fix(layout): validate SF tensors are CUDA; add MXFP4 (1, 32) SF transform test - #75
Open
qqtang-code wants to merge 2 commits into
Open
fix(layout): validate SF tensors are CUDA; add MXFP4 (1, 32) SF transform test#75qqtang-code wants to merge 2 commits into
qqtang-code wants to merge 2 commits into
Conversation
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.
Author
|
Validation summary (4× RTX 6000D, SM120):
|
qqtang-code
marked this pull request as ready for review
August 14, 2026 09:33
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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:
preprocess_sfdevice 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, sopreprocess_sfnow assertssf.is_cuda()up front (coversget_mn_major_tma_aligned_tensorandget_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 intorch::from_blob); this guard covers the pybind/direct path and is defense-in-depth rather than the fix for that boundary crash.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_recipepassed on SM120 (RTX 6000D, fresh build from this branch): shapes (1, 3072, 3584), (4, 1024, 3584), (8, 3072, 7168) all match the torch reference.--moe-runner-backend deep_gemm; decode TPOT p50 76.0 ms vs marlin baseline 74.9 ms.Critical: the published
sgl-deep-gemm0.1.5.post2 wheel is brokenField crashes ("Unknown SF transformation" / dtype assertion / segfault during MXFP4 weight prep on SM120) trace to the published PyPI wheel, not the source:
TypeMeta, wrong dims) —transform_sf_into_required_layout,preprocess_sf,get_mn_major_tma_aligned_packed_ue8m0_tensorall fail for any dtype/shape, withAssertion error (layout.hpp:93)or a segfault inTypeMeta::error_unsupported_typemeta._C.so.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 oftransform_sf_into_required_layoutcoversarch_major == 10 or 12since #56; the test locks in the SM120 path that SGLang's MXFP4 MoE runner relies on once enabled there (sgl-project/sglang#34827).