[PR #407] fix(layout): validate SF tensors are CUDA; add MXFP4 (1, 32) SF transform test #424

New issue

Open

opened 2026-10-26 21:17:37 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:17:37 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/407
Author: @qqtang-code
Created: 8/14/2026
Status: 🔄 Open

Base: main ← Head: feat/sf-layout-device-guard

📝 Commits (2)

08562eb fix(layout): validate SF tensors are CUDA before JIT launch
73ccae7 review: address ds-review-bot feedback

📊 Changes

2 files changed (+33 additions, -0 deletions)

View changed files

📝 csrc/jit_kernels/impls/smxx_layout.hpp (+5 -0)
📝 tests/test_layout.py (+28 -0)

📄 Description

Summary

Two small hardening changes for the scale-factor (SF) layout APIs, shared with the SGLang fork (sgl-project/DeepGEMM#75):

preprocess_sf device validation — the layout-transform functions feed .data_ptr() straight into JIT-launched device kernels (or into the torch fallback). A CPU tensor reaches the kernel launch with host pointers and segfaults the process (illegal device address) instead of failing cleanly. preprocess_sf now asserts sf.is_cuda() up front, covering both get_mn_major_tma_aligned_tensor and get_mn_major_tma_aligned_packed_ue8m0_tensor.
Regression test for the MXFP4 (1, 32) recipe — mirrors the exact call SGLang makes during Kimi-K3 / DeepSeek-V4 weight loading (transform_sf_into_required_layout(sf, mn=..., k=..., recipe=(1, 32), num_groups=..., disable_ue8m0_cast=False)), asserting the packed-UE8M0 output matches the torch reference (values, shape, strides). Skipped on SM90 (which uses the (1, 128) FP32 TMA layout) and on SM120 (upstream main has no arch-12 branch; SM120 support lives on the sgl-project fork).

Context

The device assert turns a silent process crash (CPU tensor into a device kernel) into a clean assertion error. The regression test locks in the MXFP4 weight-prep path that SGLang's MoE runner relies on (validated end-to-end on 4× RTX 6000D / SM120 with the sgl fork: sgl-project/sglang#34827).

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepGEMM/pull/407 **Author:** [@qqtang-code](https://github.com/qqtang-code) **Created:** 8/14/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `feat/sf-layout-device-guard` --- ### 📝 Commits (2) - [`08562eb`](https://github.com/deepseek-ai/DeepGEMM/commit/08562eb18b67137b0c0405f3c263e66e74beb99d) fix(layout): validate SF tensors are CUDA before JIT launch - [`73ccae7`](https://github.com/deepseek-ai/DeepGEMM/commit/73ccae73395a67fce3830db26189367ad3d3587f) review: address ds-review-bot feedback ### 📊 Changes **2 files changed** (+33 additions, -0 deletions) <details> <summary>View changed files</summary> 📝 `csrc/jit_kernels/impls/smxx_layout.hpp` (+5 -0) 📝 `tests/test_layout.py` (+28 -0) </details> ### 📄 Description ## Summary Two small hardening changes for the scale-factor (SF) layout APIs, shared with the SGLang fork (sgl-project/DeepGEMM#75): 1. **`preprocess_sf` device validation** — the layout-transform functions feed `.data_ptr()` straight into JIT-launched device kernels (or into the `torch` fallback). A CPU tensor reaches the kernel launch with host pointers and segfaults the process (illegal device address) instead of failing cleanly. `preprocess_sf` now asserts `sf.is_cuda()` up front, covering both `get_mn_major_tma_aligned_tensor` and `get_mn_major_tma_aligned_packed_ue8m0_tensor`. 2. **Regression test for the MXFP4 `(1, 32)` recipe** — mirrors the exact call SGLang makes during Kimi-K3 / DeepSeek-V4 weight loading (`transform_sf_into_required_layout(sf, mn=..., k=..., recipe=(1, 32), num_groups=..., disable_ue8m0_cast=False)`), asserting the packed-UE8M0 output matches the torch reference (values, shape, strides). Skipped on SM90 (which uses the `(1, 128)` FP32 TMA layout) and on SM120 (upstream main has no arch-12 branch; SM120 support lives on the sgl-project fork). ## Context The device assert turns a silent process crash (CPU tensor into a device kernel) into a clean assertion error. The regression test locks in the MXFP4 weight-prep path that SGLang's MoE runner relies on (validated end-to-end on 4× RTX 6000D / SM120 with the sgl fork: sgl-project/sglang#34827). --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-26 21:17:37 +00:00

bsakaguchi referenced this issue

2026-10-26 21:17:45 +00:00

[PR #424] docs: note Mega MoE requires PyTorch >= 2.9 in Requirements #436