[Common] Optimize fused router forward/backward kernels

[harryzhou2000]

Summary

Optimizes the fused router CUDA kernels introduced in #2821 (fused_topk_with_score_function and fused_score_for_moe_aux_loss). Achieves significant bandwidth improvements for large expert counts and topk values while preserving identical performance for smaller configurations (e.g., E=256, topk=4).

Key results (B300, float32, 8192 tokens):

• Forward (E=2304, K=36, softmax): 673 → 964 GB/s (+43%)
• Backward (E=2304, K=36, softmax): 543 → 2766 GB/s (+410%)
• Forward (E=512, K=4): no regression (±0.3%)

Changes

Forward kernels

• Persistent grid with async double-buffered prefetch: RawAsyncLoader<T> uses cp.async (sm_80+) for non-blocking global→shmem loads. Occupancy-aware grid sizing (compute_persistent_grid) keeps all SMs saturated across multiple rounds.
• Packed 8-bit radix histogram: Reduces radix topk register usage from 32 to 4 registers by packing 16 bucket counts into 4×u32 with 8-bit fields. Eliminates local memory spill at large E.
• Compile-time score function dispatch: ScoreFunc template parameter with if constexpr removes runtime branches from the hot loop.
• Simple kernel path for small topk: When topk < NVTE_RADIX_TOPK_THRESHOLD (default 8), dispatches to a lightweight kernel matching the original structure — no async loader, no persistent grid — avoiding scheduling overhead that dominates at small K.

Backward kernels

• Two-pass fused design: Pass 1 accumulates warp-level sums via register reduction + warp_allreduce_sum. Pass 2 computes per-element gradients using scalar helpers. Eliminates the comp_buf shared memory buffer (saves E × warps × 4 bytes per block).
• Double-buffered async loading: All backward inputs (grad, activation, mask) loaded through RawAsyncLoader with always-on double buffering.

Infrastructure

• async_loader.h: RawAsyncLoader<T>, compute_persistent_grid(), choose_num_buffers(), vectorized global store/fill helpers.
• NVTE_RADIX_TOPK_THRESHOLD env var (default 8): configurable naive↔radix crossover.
• Templated warp_reduce_on_shmem<T, ReduceFuncType> eliminates function-pointer overhead.

Hardening

• Host-side: num_tokens * num_experts <= INT_MAX, topk ∈ [1, E], topk % group_topk == 0
• Device-side: assert(data_size <= kMaxExpertsRadixTopk) in radix path
• Correct cudaDevAttrMaxSharedMemoryPerMultiprocessor for buffer-count decision
• Fix: single-buffer prefetch clobber when shmem is too tight for double buffering

Compatibility

• No regression for small configs: The simple forward kernel path is an exact replica of the original kernel structure, ensuring E=256/topk=4 (common in standard MoE) performs identically.
• All existing tests pass: 891/891 test_fused_router.py tests pass, 117 skipped (fp8/multi-node).
• No API changes: Same Python/C++ interface, same output semantics.
• Tunable: Set NVTE_RADIX_TOPK_THRESHOLD=0 to force radix everywhere, or =16 to use naive for topk<16.

Performance (B300 SXM6, sm_103, float32, 8192 tokens)

Effective bandwidth (GB/s) is computed as the minimum bytes that must be transferred to/from global memory for one kernel invocation, divided by the measured wall time. For example, the topk forward kernel reads logits (T×E×dtype) and writes probs (T×E×dtype), routing_map (T×E×1), and intermediate_output (T×E×4). This metric captures how well the kernel utilizes memory bandwidth — higher is better, with the device peak around 8 TB/s on B300. Config format is num_experts/topk.

Full benchmark table (softmax)

kernel	pass	config	before	after
topk	fprop	512/4	1779	1784 (+0.3%)
topk	fprop	512/8	798	904 (+13%)
topk	fprop	512/22	514	924 (+80%)
topk	fprop	512/36	499	908 (+82%)
topk	fprop	2304/4	1803	1802 (0%)
topk	fprop	2304/8	660	993 (+51%)
topk	fprop	2304/22	602	972 (+61%)
topk	fprop	2304/36	673	964 (+43%)
topk	bprop	512/22	3391	5362 (+58%)
topk	bprop	2304/36	543	2766 (+410%)
aux_loss	fprop	512/22	519	896 (+73%)
aux_loss	fprop	2304/36	645	891 (+38%)
aux_loss	bprop	512/22	5289	6155 (+16%)
aux_loss	bprop	2304/36	2272	4201 (+85%)

Full benchmark table (sigmoid)

kernel	pass	config	before	after
topk	fprop	512/4	1728	1736 (+0.5%)
topk	fprop	512/22	470	891 (+90%)
topk	fprop	2304/36	639	798 (+25%)
topk	bprop	512/22	3169	4398 (+39%)
topk	bprop	2304/36	533	2274 (+327%)
aux_loss	fprop	512/22	475	912 (+92%)
aux_loss	fprop	2304/36	598	867 (+45%)
aux_loss	bprop	2304/36	1965	2757 (+40%)

[greptile-apps]

Greptile Summary

This PR significantly optimizes the fused MoE router CUDA kernels for large expert counts and topk values by introducing persistent grids with double-buffered cp.async prefetch, packed 8-bit radix histogram counters, compile-time score-function dispatch, and a two-pass backward that eliminates the shared-memory comp_buf. A lightweight simple-kernel path is retained for small topk configurations to avoid overhead that would otherwise regress common use-cases.

• Forward: split dispatch routes small topk (below NVTE_RADIX_TOPK_THRESHOLD, default 16) to a replica of the original kernel and larger topk to the optimized async+radix path; choose_num_buffers selects single vs. double buffering based on device shared-memory headroom.
• Backward: all three inputs (grad, activations, mask) loaded through RawAsyncLoader; a warp-level two-pass reduction replaces the previous comp_buf shared-memory accumulation, yielding up to 4× bandwidth improvement at E=2304/K=36.
• Infrastructure: new async_loader.h centralizes cp.async wrappers, occupancy-aware grid sizing, and vectorized store/fill helpers; host-side input validation (INT_MAX overflow, topk ∈ [1,E], divisibility) added throughout.

Confidence Score: 5/5

Safe to merge — the correctness-critical changes (shmem check split, buffer-count selection, two-pass backward math) are all verified sound, and 891/891 existing tests pass.

The forward dispatch correctly applies the shmem capacity check against the actual kernel's allocation. The backward two-pass reduction accurately implements normalization, softmax, sigmoid, and sqrtsoftplus gradients. The packed 8-bit radix histogram is guarded by kMaxExpertsRadixTopk on both host and device. Only minor documentation mismatches and a dead parameter were found.

No files require special attention; the style notes are in utils.h and fused_score_for_moe_aux_loss.cu.

Important Files Changed

Filename	Overview
transformer_engine/common/fused_router/async_loader.h	New header introducing RawAsyncLoader (cp.async double-buffering), compute_persistent_grid, choose_num_buffers, and vectorized store/fill helpers — well-structured and correct.
transformer_engine/common/fused_router/utils.h	Packed 8-bit radix histogram, compile-time warp_reduce_on_shmem, and new scalar score helpers — logic is correct; minor comment inaccuracy about inline static TU behavior.
transformer_engine/common/fused_router/fused_topk_with_score_function.cu	Adds simple/optimized forward dispatch and two-pass fused backward; shmem checks are now correctly split per code path; all previous review concerns resolved.
transformer_engine/common/fused_router/fused_score_for_moe_aux_loss.cu	Parallel simple/optimized forward kernels and new two-pass backward; dead topk parameter in backward launcher, shmem dispatch logic is correct.
transformer_engine/common/fused_router/fused_moe_aux_loss.cu	Minimal change: updates warp_reduce_on_shmem call to new compile-time template signature — correct.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Launcher called] --> B{topk >= threshold\nAND num_experts <= 8160?}
    B -- No --> C[Simple kernel path\nother_shmem only\nNaive topk, runtime score_function]
    B -- Yes --> D[choose_num_buffers\nlogits_buf + other_shmem]
    D --> E{num_buffers == 2?}
    E -- Yes --> F[Optimized kernel\nDouble-buffered cp.async\nPersistent grid\nRadix topk\nCompile-time ScoreFunc]
    E -- No --> G[Optimized kernel\nSingle-buffer cp.async\nPersistent grid\nRadix topk\nCompile-time ScoreFunc]
    subgraph Backward
        H[Launcher] --> I[choose_num_buffers\nall 3 loaders combined]
        I --> J[Two-pass kernel\nPass 1: warp allreduce sums\nPass 2: per-element grad\nDouble-buffered grad+act+mask]
    end

Loading

_{Reviews (6): Last reviewed commit: "[Common] Match radix topk threshold to u..." | Re-trigger Greptile}

[denera]

LGMT. I'll fire off the CI once you resolve merge conflicts with main and then we can merge as soon as it clears through. Thanks!

[denera]

/te-ci L0

[denera]

LGTM, CI errors are unrelated to this branch, merging.