diff --git a/python/sglang/kernels/jit/csrc/deepseek_v4/mega_moe_pre_dispatch.cuh b/python/sglang/kernels/jit/csrc/deepseek_v4/mega_moe_pre_dispatch.cuh index 666c6b178..03507c7d5 100644 --- a/python/sglang/kernels/jit/csrc/deepseek_v4/mega_moe_pre_dispatch.cuh +++ b/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(token_id) * params.buf_x_sf_stride_bytes + group_id; reinterpret_cast(params.buf_x_sf)[byte_off] = static_cast(ue8m0_exp); } }; @@ -188,11 +188,10 @@ struct MegaMoEPreDispatchKernel { .with_dtype() .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() .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(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(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(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(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(buf_x_sf.data_ptr()), .buf_topk_idx = static_cast(buf_topk_idx.data_ptr()), .buf_topk_weights = static_cast(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, diff --git a/test/registered/kernels/ops/moe/test_mega_moe_pre_dispatch.py b/test/registered/kernels/ops/moe/test_mega_moe_pre_dispatch.py new file mode 100644 index 000000000..6f9537c6b --- /dev/null +++ b/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"]))