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

csrc/jit_kernels/impls/smxx_layout.hpp

@@ -110,6 +110,11 @@ static std::tuple<int, int, int, int, int, torch::Tensor> 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);

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()