diff --git a/csrc/jit_kernels/impls/smxx_layout.hpp b/csrc/jit_kernels/impls/smxx_layout.hpp index 9beaab8dfa..ce82e8a39f 100644 --- a/csrc/jit_kernels/impls/smxx_layout.hpp +++ b/csrc/jit_kernels/impls/smxx_layout.hpp @@ -110,6 +110,11 @@ static std::tuple preprocess_sf(const to const auto dim = sf.dim(); DG_HOST_ASSERT(dim == 2 or dim == 3); DG_HOST_ASSERT(sf.scalar_type() == torch::kFloat); + // The downstream paths launch device kernels over `sf` (or read + // `.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()); const auto batched_sf = dim == 2 ? sf.unsqueeze(0) : sf; const auto [num_sf_batches, mn, sf_k] = get_shape<3>(batched_sf); diff --git a/sgl_deep_gemm/tests/test_layout.py b/sgl_deep_gemm/tests/test_layout.py index 14c3a7dab4..de3a93edc2 100644 --- a/sgl_deep_gemm/tests/test_layout.py +++ b/sgl_deep_gemm/tests/test_layout.py @@ -1,5 +1,6 @@ import torch import random +from deep_gemm import transform_sf_into_required_layout from deep_gemm.testing import bench_kineto, count_bytes, get_arch_major from deep_gemm.utils import ( align, ceil_div, @@ -80,6 +81,31 @@ def test_sf_layout_kernels() -> None: print() +def test_transform_sf_into_required_layout_mxfp4_recipe() -> None: + # Mirror the MXFP4 weight-prep call made by sglang for Kimi-K3 / + # DeepSeek-V4: FP32 per-tile (1, 32) scales -> packed-UE8M0 int32 output. + # The (1, 32) UE8M0 branch of transform_sf_into_required_layout only + # exists on SM100/SM120 (SM90 uses the (1, 128) FP32 TMA layout instead). + print('Testing transform_sf_into_required_layout (MXFP4, recipe (1, 32)):') + if get_arch_major() == 9: + print(' > Skipped ((1, 32) UE8M0 layout transform is SM100/SM120-only)') + return + + for num_groups, mn, k in [(1, 3072, 3584), (4, 1024, 3584), (8, 3072, 7168)]: + # Exponent-only FP32 values (0.5 -> 0x3f000000) are valid UE8M0 + # payloads; the pack kernels assert zero sign/mantissa bits. + sf = torch.full((num_groups, mn, k // 32), 0.5, dtype=torch.float, device='cuda') + packed_sf = transform_sf_into_required_layout( + sf, mn=mn, k=k, recipe=(1, 32), num_groups=num_groups, + disable_ue8m0_cast=False, + ) + ref_packed_sf = get_mn_major_tma_aligned_packed_ue8m0_tensor_torch_impl(sf) + assert torch.equal(packed_sf, ref_packed_sf), f'{num_groups=}, {mn=}, {k=}' + assert packed_sf.shape == ref_packed_sf.shape + assert all([packed_sf.stride(i) == ref_packed_sf.stride(i) for i in range(packed_sf.dim())]) + print() + + def test_k_grouped_sf_layout_kernels() -> None: print('Testing k-grouped SF layout kernels:') for mn, ks_cpu, num_groups, gran_k in enumerate_k_grouped_sf_layout(): @@ -156,5 +182,6 @@ def test_k_grouped_psum_sf_layout_kernels() -> None: random.seed(1) test_sf_layout_kernels() + test_transform_sf_into_required_layout_mxfp4_recipe() test_k_grouped_sf_layout_kernels() test_k_grouped_psum_sf_layout_kernels()