[MoE] Gather the cutlass MoE activation and its scales in one launch

[yuan-luo]

Motivation

The fp8 blockwise CUTLASS MoE quantizes its activation once and then replicates rows per routed expert, so the gather it performs walks a dst2src map of m * topk entries. Until now it walked that map twice: one shuffle_rows launch for the fp8 values and a second one for the fp32 group scales.

The second walk moves a thirty-second of the bytes the first one does — k // 128 fp32 against k fp8 — so as a launch of its own it is almost entirely latency. That is exactly the cost that matters where this path spends its time: at low concurrency the whole gather is a few tens of KB, and the launch is the work.

Where the gather sits, between quantizing the activation once and the first grouped GEMM:

Modifications

A triton kernel that walks the map once and writes both tensors, registered as a TRITON KernelSpec (op="moe.shuffle_rows_with_scales") per RFC #29630. The grid is (dst_row, k_tile); the tile that owns column 0 also carries that row's scales, which are 1/32 the width and do not deserve a launch of their own.

The non-mxfp8 path calls it instead of the two shuffle_rows. The SM100 mxfp8 ES branch gathers a bf16 activation through a different kernel and is untouched.

The two blue arrows are the point: token 0 is read twice because it routes to two experts, which is why the gather moves topk times the rows the quantizer ever saw. The narrow cell trailing each row is that row's scales, and it travels under the same a_map — that shared map is the whole reason one kernel can produce both.

Rows are copied as bytes rather than as fp8 — the values are never decoded — so the kernel works for any 1-byte dtype and cannot perturb them. Row bases are computed in int64, matching the CUDA shuffle_rows it replaces, since rows * k overflows int32 well inside the shapes this path serves.

The kernel also masks its tails. The CUDA shuffle_rows 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 // 128) % 4 == 0; every k that is a multiple of 512 is unaffected either way, which covers the shapes this path serves today.

mxfp8 MoE is not affected by this PR. It will be addressed in the next PR.

Accuracy Tests

The result is bit-identical to the pair of calls it replaces, by construction rather than by tolerance.

test/registered/kernels/ops/moe/test_shuffle_rows_with_scales.py (new) checks that three ways, on B200: against plain torch advanced indexing as the oracle, against the two shuffle_rows calls for drop-in equivalence, and on the zero-row short circuit. Comparisons are made on integer views of both tensors — a random fp8 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 what the kernel promises.

The shape list includes k = 896 (7 scale groups), which the CUDA reference cannot express because of the remainder above, so that case is checked against the torch oracle only. Maps are drawn with randint and therefore contain duplicate source rows, which is the normal case here: a token is replicated once per expert it routes to.

22 passed, 6 warnings in 16.02s

Speed Tests and Profiling

test/registered/kernels/benchmark/moe/bench_shuffle_rows_with_scales.py (new) benchmarks the fused kernel against the two launches it replaces. On B200 (sm100, torch 2.11.0+cu130, triton 3.6.0), topk 8:

k	tokens	rows	two launches	fused	speedup
7168	1	8	3.06 us	1.31 us	2.33x
2048	1	8	2.47 us	1.42 us	1.74x
7168	8	64	3.25 us	1.60 us	2.03x
7168	64	512	4.30 us	2.10 us	2.05x
7168	1024	8192	23.78 us	11.14 us	2.14x

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

---

CI States

Latest PR Test (Base): ❌ Run #32585586997
Latest PR Test (Extra): ✅ Run #32585586727
Latest PR Test (AMD ROCm 7.2): ❌ Run #32585586801

[BBuf]

LGTM.