[MegaMoE] Respect padded MXFP8 scale row strides in pre-dispatch (#36237)

Co-authored-by: yangliu991 <yangliu991@fb.com>

python/sglang/kernels/jit/csrc/deepseek_v4/mega_moe_pre_dispatch.cuh

@@ -28,10 +28,11 @@ struct MegaMoEPreDispatchParams {
   const float* __restrict__ topk_weights;  // [num_tokens, top_k]
 
   fp8_e4m3_t* __restrict__ buf_x;        // [padded_max, hidden]
-  int32_t* __restrict__ buf_x_sf;        // contiguous int32 [P, G/4]; see layout comment
+  int32_t* __restrict__ buf_x_sf;        // row-major int32 [P, G/4]
   int64_t* __restrict__ buf_topk_idx;    // [padded_max, top_k]
   float* __restrict__ buf_topk_weights;  // [padded_max, top_k]
 
+  uint64_t buf_x_sf_stride_bytes;
   uint32_t num_tokens;
   uint32_t padded_max;
   uint32_t hidden;
@@ -96,13 +97,12 @@ __global__ __launch_bounds__(kMaxBlockThreads, 2) void  //
       }
       out_vec.store(token_out, chunk);
 
-      // One thread per group writes its UE8M0 byte into the contiguous
-      // row-major int32-packed layout: byte address = t*num_groups + g
-      // (see layout comment at the top of the file).
+      // One thread per group writes its UE8M0 byte into the row-major int32-
+      // packed layout. DeepGEMM may pad each row to preserve TMA alignment.
       const uint32_t group_id = chunk / kThreadsPerGroup;
       const uint32_t within_group_id = chunk % kThreadsPerGroup;
       if (within_group_id == 0 && group_id < params.num_groups) {
-        const uint32_t byte_off = token_id * params.num_groups + group_id;
+        const uint64_t byte_off = static_cast<uint64_t>(token_id) * params.buf_x_sf_stride_bytes + group_id;
         reinterpret_cast<uint8_t*>(params.buf_x_sf)[byte_off] = static_cast<uint8_t>(ue8m0_exp);
       }
     };
@@ -188,11 +188,10 @@ struct MegaMoEPreDispatchKernel {
         .with_dtype<int8_t, fp8_e4m3_t>()
         .with_device(device)
         .verify(buf_x);
-    // buf.x_sf is the contiguous row-major int32 view from DeepGEMM's mega
-    // symm buffer (DeepGEMM/csrc/apis/mega.hpp): shape (P, G/4), strides
-    // (G/4, 1). No explicit strides required -> TensorMatcher enforces
-    // is_contiguous().
+    // DeepGEMM exposes only the logical G scale bytes but may pad its physical
+    // row stride to 16 bytes for TMA (for example, H=2304 is 72B -> 80B).
     TensorMatcher({P, G4})  // buf_x_sf
+        .with_strides({-1, 1})
         .with_dtype<int32_t>()
         .with_device(device)
         .verify(buf_x_sf);
@@ -215,6 +214,13 @@ struct MegaMoEPreDispatchKernel {
     RuntimeCheck(hidden % kGroupSize == 0, "hidden must be a multiple of group_size");
     const auto num_groups = hidden / static_cast<uint32_t>(kGroupSize);
     RuntimeCheck(num_groups == num_groups_div_4 * 4u, "num_groups must be a multiple of 4");
+    RuntimeCheck(
+        buf_x_sf.stride(0) >= static_cast<int64_t>(num_groups_div_4),
+        "buf_x_sf row stride is smaller than its logical row width");
+    const auto buf_x_sf_stride_bytes = static_cast<uint64_t>(buf_x_sf.stride(0)) * sizeof(int32_t);
+    RuntimeCheck(buf_x_sf_stride_bytes % 16u == 0, "buf_x_sf row stride must be 16B-aligned");
+    RuntimeCheck(
+        reinterpret_cast<uintptr_t>(buf_x_sf.data_ptr()) % 16u == 0, "buf_x_sf base address must be 16B-aligned");
     RuntimeCheck(hidden % 8u == 0, "hidden must be a multiple of 8 (16B bf16 loads)");
     const auto num_chunks = hidden / 8u;
     const auto block_size = std::min(num_chunks, kMaxBlockThreads);
@@ -238,6 +244,7 @@ struct MegaMoEPreDispatchKernel {
         .buf_x_sf = static_cast<int32_t*>(buf_x_sf.data_ptr()),
         .buf_topk_idx = static_cast<int64_t*>(buf_topk_idx.data_ptr()),
         .buf_topk_weights = static_cast<float*>(buf_topk_weights.data_ptr()),
+        .buf_x_sf_stride_bytes = buf_x_sf_stride_bytes,
         .num_tokens = num_tokens,
         .padded_max = padded_max,
         .hidden = hidden,

test/registered/kernels/ops/moe/test_mega_moe_pre_dispatch.py

@@ -0,0 +1,70 @@
+import sys
+
+import pytest
+import torch
+
+from sglang.kernels.ops.attention.dsv4 import mega_moe_pre_dispatch
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10,
+    reason="Requires Blackwell GPU (sm_100+)",
+)
+def test_mxfp8_scale_output_uses_padded_row_stride() -> None:
+    """Write each MXFP8 scale row without overwriting its physical padding."""
+    torch.manual_seed(42)
+    num_tokens, padded_max, hidden, top_k = 5, 8, 2304, 8
+    num_groups = hidden // 32
+    logical_scale_int32 = num_groups // 4
+    scale_stride_int32 = 20
+    marker = 0xA5
+
+    x = torch.randn(num_tokens, hidden, device="cuda", dtype=torch.bfloat16)
+    topk_idx = (
+        torch.arange(num_tokens * top_k, device="cuda", dtype=torch.int32)
+        .reshape(num_tokens, top_k)
+        .remainder(256)
+    )
+    topk_weights = torch.rand(num_tokens, top_k, device="cuda", dtype=torch.float32)
+
+    buf_x = torch.empty(padded_max, hidden, device="cuda", dtype=torch.float8_e4m3fn)
+    scale_bytes = torch.full(
+        (padded_max, scale_stride_int32 * 4),
+        marker,
+        device="cuda",
+        dtype=torch.uint8,
+    )
+    buf_x_sf = scale_bytes.view(torch.int32)[:, :logical_scale_int32]
+    buf_topk_idx = torch.empty(padded_max, top_k, device="cuda", dtype=torch.int64)
+    buf_topk_weights = torch.empty(
+        padded_max, top_k, device="cuda", dtype=torch.float32
+    )
+
+    assert buf_x_sf.shape == (padded_max, logical_scale_int32)
+    assert buf_x_sf.stride() == (scale_stride_int32, 1)
+    mega_moe_pre_dispatch(
+        x,
+        topk_idx,
+        topk_weights,
+        buf_x,
+        buf_x_sf,
+        buf_topk_idx,
+        buf_topk_weights,
+    )
+    torch.cuda.synchronize()
+
+    logical_scale_bytes = logical_scale_int32 * 4
+    assert torch.all(scale_bytes[:num_tokens, :logical_scale_bytes] != marker)
+    assert torch.all(scale_bytes[:num_tokens, logical_scale_bytes:] == marker)
+    assert torch.all(scale_bytes[num_tokens:] == marker)
+    torch.testing.assert_close(buf_topk_idx[:num_tokens], topk_idx.to(torch.int64))
+    torch.testing.assert_close(buf_topk_weights[:num_tokens], topk_weights)
+    assert torch.all(buf_topk_idx[num_tokens:] == -1)
+    assert torch.all(buf_topk_weights[num_tokens:] == 0)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))