[MoE] Gather the cutlass MoE activation and its scales in one launch (#34915)

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>

test/registered/kernels/benchmark/moe/bench_shuffle_rows_with_scales.py

@@ -0,0 +1,69 @@
+import torch
+import triton
+import triton.testing
+
+from sglang.kernels.jit.benchmark.utils import (
+    DEFAULT_DEVICE,
+    get_benchmark_range,
+    run_benchmark,
+)
+from sglang.kernels.ops.moe.shuffle_rows_with_scales import shuffle_rows_with_scales
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(
+    est_time=15, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)
+
+GROUP_SIZE = 128
+
+# (hidden, tokens). The cutlass fp8 blockwise MoE gathers tokens * topk rows out
+# of a tokens-row source, so the small end is bs=1 decode and the large end is a
+# prefill-sized batch. 7168 is a DeepSeek-class hidden size.
+SHAPES = get_benchmark_range(
+    full_range=[(7168, 1), (2048, 1), (7168, 8), (7168, 64), (7168, 1024)],
+    ci_range=[(7168, 1), (7168, 1024)],
+)
+
+
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["k", "tokens"],
+        x_vals=SHAPES,
+        line_arg="provider",
+        line_vals=["fused", "two_launches"],
+        line_names=["Fused gather", "Two shuffle_rows"],
+        styles=[("blue", "-"), ("red", "--")],
+        ylabel="us",
+        plot_name="shuffle-rows-with-scales-performance",
+        args={"topk": 8},
+    )
+)
+def benchmark(k: int, tokens: int, topk: int, provider: str):
+    from sgl_kernel import shuffle_rows
+
+    rows = tokens * topk
+    q = torch.randint(
+        0, 256, (tokens, k), dtype=torch.uint8, device=DEFAULT_DEVICE
+    ).view(torch.float8_e4m3fn)
+    scale = torch.randn(
+        (tokens, k // GROUP_SIZE), dtype=torch.float32, device=DEFAULT_DEVICE
+    )
+    # Duplicate source rows are the normal case: a token is replicated once per
+    # expert it routes to.
+    dst2src = torch.randint(
+        0, tokens, (rows,), dtype=torch.int32, device=DEFAULT_DEVICE
+    )
+
+    if provider == "fused":
+        fn = lambda: shuffle_rows_with_scales(q, scale, dst2src, rows)
+    else:
+        fn = lambda: (
+            shuffle_rows(q, dst2src, (rows, k)),
+            shuffle_rows(scale, dst2src, (rows, k // GROUP_SIZE)),
+        )
+
+    return run_benchmark(fn)
+
+
+if __name__ == "__main__":
+    benchmark.run(print_data=True)

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

@@ -0,0 +1,123 @@
+"""Bit-exactness of the fused value+scale row gather.
+
+The oracle is plain torch advanced indexing, which is the whole contract:
+``out[i] = src[dst2src_map[i]]`` for both tensors. A second test cross-checks the
+pair of `shuffle_rows` calls this replaces, to back the drop-in claim.
+
+Comparisons are made on integer views. The values are fp8 and a random byte
+pattern is a NaN often enough that `torch.equal` on the float view would report a
+difference where the bytes agree -- and bytes are exactly what this kernel
+promises to preserve.
+"""
+
+import itertools
+import sys
+
+import pytest
+import torch
+
+from sglang.kernels.jit.utils import get_ci_test_range
+from sglang.kernels.ops.moe.shuffle_rows_with_scales import shuffle_rows_with_scales
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+GROUP_SIZE = 128
+
+# The CUDA shuffle_rows this replaces loads 128 bits per thread and takes its
+# element count as num_cols / elems_per_thread with no remainder handling, so it
+# only moves a whole fp32 scale row when (k // GROUP_SIZE) % 4 == 0. Shapes below
+# that bound are checked against the torch oracle only -- the reference itself
+# drops the tail there.
+CUDA_XCHECK_SCALE_ALIGN = 4
+
+# k = 896 gives 7 scale groups, the shape whose scale tail the CUDA reference
+# cannot express; 7168 is the DeepSeek-class hidden size. Destination rows exceed
+# source rows because the map replicates each token once per routed expert.
+CASES = get_ci_test_range(
+    [
+        (k, src, dst)
+        for k, (src, dst) in itertools.product(
+            [512, 896, 2560, 7168], [(1, 8), (17, 136), (64, 512)]
+        )
+    ],
+    [
+        (896, 1, 8),
+        (7168, 1, 8),
+        (2560, 17, 136),
+        (7168, 64, 512),
+    ],
+)
+
+
+def _inputs(k, num_src_rows, num_dst_rows, seed):
+    torch.manual_seed(seed)
+    q = torch.randint(0, 256, (num_src_rows, k), dtype=torch.uint8, device="cuda").view(
+        torch.float8_e4m3fn
+    )
+    scale = torch.randn(
+        (num_src_rows, k // GROUP_SIZE), dtype=torch.float32, device="cuda"
+    )
+    # Duplicate source rows are the normal case here: a token is replicated once
+    # per expert it routes to.
+    dst2src = torch.randint(
+        0, num_src_rows, (num_dst_rows,), dtype=torch.int32, device="cuda"
+    )
+    return q, scale, dst2src
+
+
+def _assert_same_bytes(got, ref, what):
+    assert torch.equal(
+        got.view(torch.int8), ref.view(torch.int8)
+    ), f"{what} bytes differ"
+
+
+@pytest.mark.parametrize("k,num_src_rows,num_dst_rows", CASES)
+def test_matches_torch_gather(k, num_src_rows, num_dst_rows):
+    q, scale, dst2src = _inputs(k, num_src_rows, num_dst_rows, seed=0)
+
+    got_q, got_scale = shuffle_rows_with_scales(q, scale, dst2src, num_dst_rows)
+
+    idx = dst2src.long()
+    _assert_same_bytes(got_q, q[idx], "values")
+    _assert_same_bytes(got_scale, scale[idx], "scales")
+
+
+CUDA_XCHECK_CASES = [
+    c for c in CASES if (c[0] // GROUP_SIZE) % CUDA_XCHECK_SCALE_ALIGN == 0
+]
+# An empty parametrize list collects zero tests and reports success, which would
+# retire the drop-in check without saying so. Fail at collection instead.
+assert CUDA_XCHECK_CASES, "no case survives the CUDA cross-check shape filter"
+
+
+@pytest.mark.parametrize("k,num_src_rows,num_dst_rows", CUDA_XCHECK_CASES)
+def test_matches_shuffle_rows_pair(k, num_src_rows, num_dst_rows):
+    """Drop-in equivalence with the two launches this replaces."""
+    from sgl_kernel import shuffle_rows
+
+    q, scale, dst2src = _inputs(k, num_src_rows, num_dst_rows, seed=1)
+
+    got_q, got_scale = shuffle_rows_with_scales(q, scale, dst2src, num_dst_rows)
+
+    _assert_same_bytes(got_q, shuffle_rows(q, dst2src, (num_dst_rows, k)), "values")
+    _assert_same_bytes(
+        got_scale,
+        shuffle_rows(scale, dst2src, (num_dst_rows, k // GROUP_SIZE)),
+        "scales",
+    )
+
+
+def test_empty_destination_does_not_launch():
+    """Zero rows takes the short circuit instead of a zero-sized grid."""
+    q, scale, dst2src = _inputs(512, 4, 0, seed=2)
+
+    got_q, got_scale = shuffle_rows_with_scales(q, scale, dst2src, 0)
+
+    assert got_q.shape == (0, 512)
+    assert got_scale.shape == (0, 512 // GROUP_SIZE)
+    assert got_q.dtype == q.dtype and got_scale.dtype == scale.dtype
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))