perf: add small-batch route MMA for MoE

[Religious-J]

perf(moe): add a route-direct WGMMA path for small-batch FP8 MoE

Summary

This PR adds an SM90a route-direct WGMMA path for small-batch FP8 MoE decode, covering both per-tensor and 128x128 blockwise quantization.

For the production dispatch range (batch 1-4, top-k 8), the new path reduces single-op CUDA Graph latency by 9.25%-38.57% (1.10x-1.63x) on the three H20 workloads measured below.

The fast path:

• reads the expert directly from each token/top-k route;
• bypasses expert sorting, input gathering, prefix sums, and task-map construction;
• uses route-direct FP8 WGMMA for both Gate/Up and Down projections;
• optionally splits the long Gate/Up K dimension to expose enough parallel work;
• combines split reduction, SwiGLU, and Down-input quantization;
• preserves EP routing, shared output, preallocated output, and CUDA Graph capture.

Shapes outside the selected range continue to use the existing implementation.

Motivation

The existing MoE path is designed around grouping rows by expert:

count/sort -> gather -> prefix/task map -> Gate/Up grouped GEMM -> activation/quantization -> Down grouped GEMM -> scatter/top-k reduction

This is efficient once each expert receives enough rows. During decode, however, B * top_k can be much smaller than the expert count. For example, batch 1 with top-k 8 creates only eight routes across 128 or 192 experts. Most experts receive zero rows, so sorting, padding, gathering, metadata construction, and intermediate traffic dominate the useful matrix multiplication.

The new path keeps data in route order:

topk routes -> route Gate/Up WGMMA -> split reduction + SwiGLU + FP8 quantization -> route Down WGMMA -> contiguous top-k reduction

Each logical route is assigned to one SM90 M=8 WGMMA tile. Row 0 carries the route and rows 1-7 are zero-filled. Although this intentionally under-utilizes the M dimension, it avoids the larger fixed cost of expert grouping for very small batches.

Code changes

• Added src/fuse_moe/small_batch_route_mma.{h,cu} to orchestrate the per-tensor and blockwise route pipelines.
• Added route-direct and split-K launchers to src/group_gemm/cp_async/group_gemm.{h,cu}.
• Added small-batch dispatch, split selection, contiguous workspace allocation, output validation, and preallocated-output support in src/fuse_moe/entry.cc.
• Tightened the existing per-tensor cp.async task-map bound to the actual non-empty workload and added missing input/output shape, dtype, device, and contiguity validation.
• Updated src/fuse_moe/reduce.cu to reduce contiguous route-order outputs without a topk_pos map.

Dispatch

The production dispatch remains deliberately conservative.

Per-tensor FP8

• 1 <= batch <= 4
• top_k == 8
• 0 < I <= 512
• H % 64 == 0
• I % 64 == 0

Blockwise FP8

• 1 <= batch <= 4
• top_k == 8
• H <= 4096
• 128 <= I <= 768
• H % 128 == 0
• I % 64 == 0

All other shapes use the existing path. The forced B=5-8 measurements below characterize the profitability boundary only; this PR does not enable those batches in production dispatch.

Correctness

Regression and targeted validation

Check	Result
Full pytest regression	452 passed in 183.76 s
Additional targeted cases	14/14 passed
CUDA Graph replay	Bitwise equal to eager output

Reference-comparison metrics:

Quantization	Worst relative L2	Minimum cosine similarity
Blockwise	1.61%	0.999873
Per-tensor	1.47%	0.999892

Performance

Methodology

• GPU: NVIDIA H20 96 GB (SM90)
• Driver: 535.161.08
• CUDA: 13.0
• PyTorch: 2.11.0+cu130
• top_k=8
• Single-op CUDA Graph capture with preallocated output

To isolate the fast path, both binaries were built from the same clean bc97d55 source:

• Same-source fallback: route-MMA dispatch disabled.
• Forced route-MMA: batch 1-8 forced through route-MMA while retaining all other shape constraints.

Speedup = fallback latency / route-MMA latency.

Per-tensor Hunyuan-V3 TP8: E=192, H=4096, I=192

Batch	Same-source fallback (us)	Forced route-MMA (us)	Speedup	Latency reduction
1	20.767	13.243	1.568x	36.23%
2	27.400	16.833	1.628x	38.57%
3	33.864	22.773	1.487x	32.75%
4	41.668	30.663	1.359x	26.41%
5	45.830	35.983	1.274x	21.49%
6	49.833	41.554	1.199x	16.61%
7	53.052	47.259	1.123x	10.92%
8	56.848	54.124	1.050x	4.79%

This shape remains profitable through B=8. A future shape-aware dispatch could extend its range independently of the blockwise paths.

Blockwise Qwen3: E=128, H=2048, I=768

Batch	Same-source fallback (us)	Forced route-MMA (us)	Speedup	Latency reduction
1	24.725	17.349	1.425x	29.83%
2	41.528	29.863	1.391x	28.09%
3	50.437	40.708	1.239x	19.29%
4	58.109	52.732	1.102x	9.25%
5	67.782	62.510	1.084x	7.78%
6	75.988	73.040	1.040x	3.88%
7	83.356	84.419	0.987x	-1.28%
8	89.050	94.352	0.944x	-5.95%

Route-MMA remains profitable through B=6 for this measured shape and crosses over between B=6 and B=7.

Blockwise: E=128, H=4096, I=512

Batch	Same-source fallback (us)	Forced route-MMA (us)	Speedup	Latency reduction
1	29.845	19.779	1.509x	33.73%
2	50.389	37.585	1.341x	25.41%
3	64.916	52.295	1.241x	19.44%
4	73.862	63.833	1.157x	13.58%
5	86.072	85.218	1.010x	0.99%
6	98.718	98.465	1.003x	0.26%
7	107.532	111.362	0.966x	-3.56%
8	117.451	121.922	0.963x	-3.81%

B=5-6 are effectively at the noise boundary, while B=7-8 regress. The conservative B<=4 production threshold is appropriate for this shape.

Conclusions

Within the current B=1-4 production range, route-MMA improves all three measured workloads by 1.10x-1.63x.