[Common] Optimize fused router forward/backward kernels#3012
Open
harryzhou2000 wants to merge 15 commits into
Open
[Common] Optimize fused router forward/backward kernels#3012harryzhou2000 wants to merge 15 commits into
harryzhou2000 wants to merge 15 commits into
Conversation
harryzhou2000
force-pushed
the
hhanyu/router_fix_p3R
branch
2 times, most recently
from
14a302c to
a805f38
Compare
Contributor
Greptile SummaryThis PR optimizes the fused router CUDA kernels by introducing persistent grids, double-buffered cp.async loads, packed 8-bit radix histograms, compile-time score-function dispatch, and a two-pass backward design that eliminates the comp_buf shared-memory buffer. Previous review concerns (premature shmem check, hardcoded backward buffer count) are addressed.
Confidence Score: 5/5Safe to merge; all 891 tests pass and the kernel logic is correct throughout. The implementation is carefully crafted: premature shmem-check and hardcoded backward buffer count from the prior review are both fixed. The async-loader pipeline, packed radix histogram, and two-pass backward gradient math are all correct. The behavioral change in the topk pre-softmax backward (non-selected experts now receive the mathematically correct softmax-coupling gradient rather than zero) does not affect test outcomes. fused_topk_with_score_function.cu - the pre-softmax backward behavioral change and the choose_num_buffers argument pattern are worth a second look before landing. Important Files Changed
Reviews (2): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Replace multi-loop preprocess (separate clear/load/score/save/bias loops) with single fused loops per score function in all 4 kernel paths (topk forward, topk backward, aux_loss forward, aux_loss backward). Replace multi-pass backward (array-based helpers + comp_buf shmem) with a two-pass approach using scalar helpers: Pass 1: reduction — warp-level sums via warp_allreduce_sum() Pass 2: element-wise — scalar gradient computation → write to global Add scalar helpers to utils.h: sigmoid_scalar, sqrtsoftplus_scalar, sigmoid_bwd_scalar, sqrtsoftplus_bwd_scalar, normalize_bwd_scalar, softmax_bwd_scalar. Remove dead array helpers from utils.h: apply_sigmoid_on_float, apply_sigmoid_bwd_on_float, apply_sqrtsoftplus_on_float, apply_sqrtsoftplus_bwd_on_float, apply_softmax_bwd_on_float, masked_warp_reduce_on_shmem. Backward shmem reduced by E×W×sizeof(float) per kernel (comp_buf eliminated). Net -226 lines across 3 files. Signed-off-by: Harry Zhou <hhanyu@nvidia.com>
Add async_loader.h with:
- RawAsyncLoader<T>: cp.async on sm_80+, int4 fallback on sm_70,
stores data in original type (no conversion during copy)
- compute_persistent_grid(): occupancy-based grid sizing
- choose_num_buffers(): shmem-aware 1-vs-2 buffer decision
- vec_fill_global(), vec_store_global(): vectorized output helpers
Forward kernels (topk + aux_loss):
- Logits loaded via RawAsyncLoader with double-buffered prefetch
- Persistent grid replaces 1-shot grid launch
- DataType→CompType conversion during compute, not during load
- vec_fill_global for clearing probs/routing_map
Backward kernels (topk + aux_loss):
- All inputs loaded via RawAsyncLoader (topk: 3 loaders for
grad/act/mask; aux_loss: 2 loaders for grad/act)
- Always double-buffered (kBwdNumBuffers=2, kAuxBwdNumBuffers=2)
- Persistent grid with occupancy-based sizing
Signed-off-by: Harry Zhou <hhanyu@nvidia.com>
Replace counts[16] + total_counts[16] (32 registers) with 4 packed u32 registers using 8-bit fields (4 counters per register). Eliminates massive register spill to local memory on large kernels (81% of L1 traffic on E=2304, K=36). Add kMaxExpertsRadixTopk constant (8160 = 255 * 32) and runtime checks in both forward launchers to guard against 8-bit overflow. All current MoE configurations (max E=2304) are well within this limit. Signed-off-by: Harry Zhou <hhanyu@nvidia.com>
…dispatch Replace runtime score_function parameter in all 4 kernel __global__ functions with template int ScoreFunc (0=sigmoid, 1=softmax, 2=sqrtsoftplus). All score_function branches now use if constexpr, eliminating dead-code register pressure and branch overhead. Forward launchers dispatch on TopkFunc × ScoreFunc = 6 instantiations per DataType. Backward launchers dispatch on ScoreFunc = 3 instantiations per DataType. Signed-off-by: Harry Zhou <hhanyu@nvidia.com>
Fix broken topk < 0 threshold (radix was always selected, naive unreachable). Replace with configurable NVTE_RADIX_TOPK_THRESHOLD env var (default 0, i.e. always use radix). Set to 16 to restore the old naive-for-small-K behavior. Uses the standard TE pattern: static local + getenv (read once, cached for process lifetime). Signed-off-by: Harry Zhou <hhanyu@nvidia.com>
When choose_num_buffers() returns 1 (shmem too tight for double buffering, e.g. E=1024 with group_topk scratch), buf_[0] and buf_[1] alias the same memory. The prefetch via start_load(next_buf()) then overwrites the current buffer while compute is still reading it. Fix: guard the prefetch on num_buffers > 1. When single-buffered, load the current round's data at the top of each iteration instead. The first round's load_current is still issued before the loop. Backward kernels are unaffected (always kBwdNumBuffers=2). Signed-off-by: Harry Zhou <hhanyu@nvidia.com>
Signed-off-by: Harry Zhou <hhanyu@nvidia.com>
Code review fixes: - C1: choose_num_buffers() now queries cudaDevAttrMaxSharedMemoryPerMultiprocessor (per-SM budget) instead of cudaDevAttrMaxSharedMemoryPerBlockOptin (per-block max). These coincide on Hopper/Blackwell but differ on Ampere. - H3: Remove dead fallback branch in choose_num_buffers() — since total_double >= total_single always, blocks_single >= blocks_double, so the old ternary always returned 1 anyway. - H4/M8: Add host-side NVTE_CHECK in all 4 launchers: - num_experts > 0 - topk in [1, num_experts] - (int64_t)num_tokens * num_experts <= INT_MAX (kernel uses int offsets) - M9: Assert topk % group_topk == 0 when group_topk > 0. - H6: Add device-side assert(data_size <= kMaxExpertsRadixTopk) in radix_topk_and_mask() — zero cost in release (NDEBUG), catches 8-bit histogram overflow in debug builds. - L1: Fix stale comments claiming default threshold is 16 (it is 0). - L4: Fix typo 'hanlded' -> 'handled'. - L8: Remove unused topk parameter from aux loss backward kernel. Signed-off-by: Harry Zhou <hhanyu@nvidia.com>
Move the duplicated static function from both .cu files into utils.h as an inline function. Each TU gets its own static local (read-once per TU), which is safe since environment variables are immutable during process lifetime. Documented this in a NOTE comment. Signed-off-by: Harry Zhou <hhanyu@nvidia.com>
Replace runtime function-pointer dispatch with compile-time if constexpr. Eliminates indirect call overhead in the reduction loop and warp shuffle butterfly, allowing the compiler to emit straight-line arithmetic. Removes the now-unused max<T>() and sum<T>() helper functions. Signed-off-by: Harry Zhou <hhanyu@nvidia.com>
When topk < NVTE_RADIX_TOPK_THRESHOLD (default 8), use a lightweight forward kernel that avoids the async loader and persistent grid overhead. The simple kernel loads logits directly from global memory to shmem and uses Naive iterative-argmax topk — matching the baseline structure that was faster for small K due to lower launch/scheduling overhead. The optimized path (async loader + persistent grid + radix topk) remains the default for topk >= 8 where the compute savings dominate. Both topk and aux_loss forward kernels get the simple variant. Backward kernels are unchanged (always use the optimized path). Signed-off-by: Harry Zhou <hhanyu@nvidia.com>
Signed-off-by: Harry Zhou <hhanyu@nvidia.com>
Use 0.0f instead of 0 to avoid ambiguity between __nv_bfloat16(float) and __nv_bfloat16(double) constructors on older CUDA toolkits. Signed-off-by: Harry Zhou <hhanyu@nvidia.com>
Signed-off-by: Harry Zhou <hhanyu@nvidia.com>
harryzhou2000
force-pushed
the
hhanyu/router_fix_p3R
branch
from
9a7cb7e to
3bab7cb
Compare
for more information, see https://pre-commit.ci
denera
requested changes
May 22, 2026
| logits, num_tokens, num_experts, topk, use_pre_softmax, num_groups, group_topk, | ||
| scaling_factor, score_function, expert_bias, probs, routing_map, intermediate_output); | ||
| // Optimized path: async loader + persistent grid + radix topk. | ||
| NVTE_CHECK(num_experts <= kMaxExpertsRadixTopk, |
Collaborator
There was a problem hiding this comment.
Unless I'm misreading the code, this is going to hard-fail for num_experts > 8160.
@ptrendx Is this a restriction we can accept?
If not, I would recommend augmenting the algorithm selection (naive topk vs. radix-sort), and just simply fall back on the naive option whenever num_experts > kMaxExpertsRadixTopk. It may not be performant but it's probably better than a hard failure.
| logits, num_tokens, num_experts, topk, score_function, scores, routing_map, | ||
| intermediate_output); | ||
| // Optimized path: async loader + persistent grid + radix topk. | ||
| NVTE_CHECK(num_experts <= kMaxExpertsRadixTopk, |
Collaborator
There was a problem hiding this comment.
Same issue as fused_topk_with_score_function. We should probably fall back on the naive path here instead of hard-failing when num_experts > 8160.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Optimizes the fused router CUDA kernels introduced in #2821 (
fused_topk_with_score_functionandfused_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):
Changes
Forward kernels
RawAsyncLoader<T>usescp.async(sm_80+) for non-blocking global→shmem loads. Occupancy-aware grid sizing (compute_persistent_grid) keeps all SMs saturated across multiple rounds.ScoreFunctemplate parameter withif constexprremoves runtime branches from the hot loop.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
warp_allreduce_sum. Pass 2 computes per-element gradients using scalar helpers. Eliminates thecomp_bufshared memory buffer (savesE × warps × 4bytes per block).RawAsyncLoaderwith always-on double buffering.Infrastructure
async_loader.h:RawAsyncLoader<T>,compute_persistent_grid(),choose_num_buffers(), vectorized global store/fill helpers.NVTE_RADIX_TOPK_THRESHOLDenv var (default 8): configurable naive↔radix crossover.warp_reduce_on_shmem<T, ReduceFuncType>eliminates function-pointer overhead.Hardening
num_tokens * num_experts <= INT_MAX,topk ∈ [1, E],topk % group_topk == 0assert(data_size <= kMaxExpertsRadixTopk)in radix pathcudaDevAttrMaxSharedMemoryPerMultiprocessorfor buffer-count decisionCompatibility
test_fused_router.pytests pass, 117 skipped (fp8/multi-node).NVTE_RADIX_TOPK_THRESHOLD=0to force radix everywhere, or=16to 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 isnum_experts/topk.Full benchmark table (softmax)
Full benchmark table (sigmoid)