Skip to content

fix(layout): validate SF tensors are CUDA; add MXFP4 (1, 32) SF transform test - #75

Open
qqtang-code wants to merge 2 commits into
sgl-project:devfrom
qqtang-code:feat/sm120-sf-layout-guard
Open

fix(layout): validate SF tensors are CUDA; add MXFP4 (1, 32) SF transform test#75
qqtang-code wants to merge 2 commits into
sgl-project:devfrom
qqtang-code:feat/sm120-sf-layout-guard

Conversation

@qqtang-code

@qqtang-code qqtang-code commented Aug 14, 2026

Copy link
Copy Markdown

Summary

Two small hardening changes for the scale-factor (SF) layout APIs, motivated by SM120 (RTX PRO 6000 / RTX 6000D) usage of the MXFP4 weight-prep path from SGLang:

  1. preprocess_sf device validation — the layout-transform functions feed .data_ptr() straight into JIT-launched device kernels. A CPU tensor reaching this point would crash the process instead of failing cleanly, so preprocess_sf now asserts sf.is_cuda() up front (covers get_mn_major_tma_aligned_tensor and get_mn_major_tma_aligned_packed_ue8m0_tensor).

    Note: through the tvm-ffi boundary a CPU tensor actually dies earlier, inside convert_to_torch_tensor (host pointer + CUDA device options in torch::from_blob); this guard covers the pybind/direct path and is defense-in-depth rather than the fix for that boundary crash.

  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 where the (1, 32) UE8M0 branch does not exist.

Validation

  • test_transform_sf_into_required_layout_mxfp4_recipe passed on SM120 (RTX 6000D, fresh build from this branch): shapes (1, 3072, 3584), (4, 1024, 3584), (8, 3072, 7168) all match the torch reference.
  • SGLang end-to-end (feat(moe): enable DeepGEMM MoE runner on SM120 (RTX PRO 6000 / RTX 6000D) sglang#34827): Kimi-K3 TP32×EP32 on 32× RTX 6000D loads and generates with --moe-runner-backend deep_gemm; decode TPOT p50 76.0 ms vs marlin baseline 74.9 ms.

Critical: the published sgl-deep-gemm 0.1.5.post2 wheel is broken

Field crashes ("Unknown SF transformation" / dtype assertion / segfault during MXFP4 weight prep on SM120) trace to the published PyPI wheel, not the source:

  • Through the wheel's tvm-ffi boundary, every tensor argument arrives corrupted on the C++ side (undefined TypeMeta, wrong dims) — transform_sf_into_required_layout, preprocess_sf, get_mn_major_tma_aligned_packed_ue8m0_tensor all fail for any dtype/shape, with Assertion error (layout.hpp:93) or a segfault in TypeMeta::error_unsupported_typemeta.
  • A minimal tvm-ffi module compiled in the same container receives the same tensors correctly — the boundary/ABI is fine; the breakage is specific to the wheel's _C.so.
  • Building the wheel from dev (this branch, version 0.1.5.post3 in the test) fixes everything on SM120.

Please rebuild/re-release the wheel — the current PyPI artifact corrupts tensor arguments at the tvm-ffi boundary and blocks MXFP4 (K3/DSV4) weight prep on any architecture that reaches transform_sf_into_required_layout.

Context

The (1, 32) UE8M0 branch of transform_sf_into_required_layout covers arch_major == 10 or 12 since #56; the test locks in the SM120 path that SGLang's MXFP4 MoE runner relies on once enabled there (sgl-project/sglang#34827).

preprocess_sf feeds .data_ptr() straight into device kernels; a CPU
tensor segfaults the process at launch (illegal device address) instead
of failing cleanly, so assert the device up front.

Also add a regression test covering transform_sf_into_required_layout
with the MXFP4 (1, 32) recipe used by sglang's Kimi-K3 / DeepSeek-V4
weight prep (packed-UE8M0 output), gated off SM90 where that UE8M0
branch does not exist.
@qqtang-code

Copy link
Copy Markdown
Author

Validation summary (4× RTX 6000D, SM120):

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant