fix(layout): validate SF tensors are CUDA; add MXFP4 (1, 32) SF transform test - #407
fix(layout): validate SF tensors are CUDA; add MXFP4 (1, 32) SF transform test#407qqtang-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). SM90 uses a different (1, 128) FP32 layout and upstream main has no SM120 branch (SM120 support lives on the sgl-project fork), so the test skips both.
| # exists on SM100 (SM90 uses the (1, 128) FP32 TMA layout; SM120 support | ||
| # lives on the sgl-project fork, see sgl-project/DeepGEMM). | ||
| if get_arch_major() in (9, 12): | ||
| print(' > Skipped ((1, 32) UE8M0 layout transform is SM100-only upstream)') |
There was a problem hiding this comment.
🔵 suggestion: 跳过分支在打印 'Testing transform_sf_into_required_layout ...' 标题之前就输出 ' > Skipped ...'，与同文件 test_k_grouped_psum_sf_layout_kernels（先打印标题、再输出跳过信息）的风格不一致。建议将 print('Testing ...') 移到跳过判断之前以保持日志风格统一；纯输出格式问题，不阻塞合入。
🤖 v5
| // `.data_ptr()` on the JIT launch path); a CPU tensor would crash the | ||
| // process at launch time (illegal device address) instead of failing | ||
| // cleanly. Validate the device up front. | ||
| DG_HOST_ASSERT(sf.is_cuda()); |
There was a problem hiding this comment.
🔵 suggestion: 新增的 DG_HOST_ASSERT(sf.is_cuda()) 与下一行 get_shape<3>(batched_sf) 中已有的 DG_HOST_ASSERT(t.is_cuda())（csrc/utils/layout.hpp）重复：preprocess_sf 的两个调用方在把 data_ptr 传给 JIT 内核之前都会先经过 get_shape<3>，CPU 张量在那里就已经抛出 DGException，并不会真的以主机指针进入内核导致进程 segfault。因此注释中 CPU 张量会崩溃进程的说法不准确，该断言属于冗余，仅让失败更早、报错更明确。建议修正注释表述，或直接删除此冗余断言。
🤖 v4p
🤖 ds-review-bot Code Reviewv6未发现会破坏现有功能或测试的缺陷。CUDA 设备校验位置合理，新增测试与当前 SM100 实现及关联仓库契约一致。 v5本 MR 的两处变更均已正确落地且与对端 MR（sgl-project#75）契约一致，可以合入。1) csrc/jit_kernels/impls/smxx_layout.hpp 在 preprocess_sf 中新增 DG_HOST_ASSERT(sf.is_cuda())：位置正确（在 dim/dtype 检查之后、任何 .data_ptr()/设备 kernel 使用之前），且 preprocess_sf 是 get_mn_major_tma_aligned_tensor 与 get_mn_major_tma_aligned_packed_ue8m0_tensor 的唯一入口，torch 回退实现 get_mn_major_tma_aligned_packed_ue8m0_tensor_torch 只从后者内部（L206/L231）可达，因此也被断言覆盖，把 CPU 张量导致的进程级 segfault（illegal device address）变为干净的断言错误。2) tests/test_layout.py 新增 MXFP4 (1, 32) recipe 回归测试：调用契约与 pybind 签名（csrc/apis/layout.hpp L126-131）匹配，传 2 元组 recipe 且不传 is_sfa 满足 DG_HOST_ASSERT(not is_sfa.has_value())（L31）；0.5 (0x3f000000) 是符号/尾数为零的合法 UE8M0 载荷，pack kernel 断言不会误触发；对值/shape/stride 均与 torch 参考实现 get_mn_major_tma_aligned_packed_ue8m0_tensor_torch_impl 比对。架构跳过逻辑与上游 layout.hpp 门控一致：(FP32, gran_k==32) 分支上游仅 arch_major==10（L49），SM90 走 (1,128) FP32 TMA 路径、SM120 会落到 DG_HOST_UNREACHABLE，故跳过 (9, 12) 必要且正确。跨仓一致性核对：C++ hunk 在两仓逐字节相同；测试差异为有意为之且与各自代码匹配——对端 fork 的 layout.hpp L49/L57 为 'arch_major == 10 or arch_major == 12' 故其测试只跳过 SM90，上游无 arch-12 分支故本 MR 额外跳过 SM120，测试注释也如实说明 SM120 支持位于 sgl fork；两侧测试的调用参数、参考实现、断言完全一致，未发现接口/协议/数据格式/调用约定不一致。仅有一处非阻塞的输出格式小建议（见 comment）。 v4p本 MR 为 SF 布局 API 增加设备校验，并新增 MXFP4 (1, 32) recipe 的回归测试，覆盖 SGLang 权重加载时调用 transform_sf_into_required_layout 的路径。整体上变更安全、测试逻辑正确（形状、stride 与值与 torch 参考实现一致，SM90/SM120 的跳过条件合理），跨仓与 sgl-project fork 的契约一致。唯一需要注意的是新增的 CUDA 断言与既有 get_shape 中的设备校验重复，注释中对崩溃原因的描述并不准确。 Files reviewed: 2 |
- preprocess_sf: clarify that the device check is an early explicit validation at the SF-layout entry (get_shape re-validates downstream), rather than claiming CPU tensors would otherwise reach kernel launch. - test: print the section title before the arch skip for log-style consistency with test_k_grouped_psum_sf_layout_kernels.
|
Thanks for the thorough review — both suggestions addressed in 73ccae7:
|
2 participants
Summary
Two small hardening changes for the scale-factor (SF) layout APIs, shared with the SGLang fork (sgl-project#75):
preprocess_sfdevice validation — the layout-transform functions feed.data_ptr()straight into JIT-launched device kernels (or into thetorchfallback). A CPU tensor reaches the kernel launch with host pointers and segfaults the process (illegal device address) instead of failing cleanly.preprocess_sfnow assertssf.is_cuda()up front, covering bothget_mn_major_tma_aligned_tensorandget_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).