[Common][PyTorch] Add QB router histogram paths

[harryzhou2000]

Description

Add opt-in Quantile Balancing (QB) support to the fused sigmoid router used by
Kimi-K3-style MoE models. The router now has statically dispatched QB specializations that
select Top-(k+1), retain the actual Top-k routes, and accumulate the per-expert histogram
needed by the QB bias update without materializing a [num_tokens, num_experts] bin-index
tensor.

The existing non-QB specialization and API behavior are unchanged when the three QB
arguments are omitted.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Quantile Balancing math

This implementation follows the
Kimi K3 Technical Report, Section 2.3.3, Eqs. (13)-(14), and Appendices C-D.
Let m be the number of tokens in the global training step, E the number of experts, and
k the number of experts routed per token. Perfect balance gives every expert the target
load

q = m * k / E.

The exact derivation assumes q is integral and that cutoff ties do not occur; the practical
histogram recovery below uses ceil(q), and TE supplies deterministic cutoff tie handling.

Routing and the token-side cutoff

For token i and expert j, the raw router score and biased selection score at step t are

s_ij = sigmoid(z_ij)
u_ij^(t) = s_ij + b_j^(t).

The bias is used only for expert selection. Mixture weights omit it, as in report Eq. (13):

T_i = Top-k_j(u_ij^(t))
p_ij = s_ij / sum_{l in T_i} s_il,  j in T_i.

To derive the next bias, QB selects Top-(k+1) under u^(t). The first k experts are the
actual routes; the (k+1)-th score is the token-side cutoff

alpha_i^(t) = (k+1)-th largest_j(s_ij + b_j^(t)).

An expert must exceed alpha_i^(t) to enter token i's Top-k. The report assumes no ties.
This implementation makes ties deterministic by discarding the largest expert ID among equal
cutoff candidates, leaving exactly k output routes.

Expert-side quantile update

Hold the cutoffs from the current forward pass fixed and consider a candidate next-step bias
bhat_j^(t+1). Its implied load for expert j is

load_j(bhat_j^(t+1))
  = sum_i 1[s_ij + bhat_j^(t+1) > alpha_i^(t)].

Setting this count to q means exactly q values of the margin
g_ij = s_ij - alpha_i^(t) exceed -bhat_j^(t+1). Therefore report Eq. (14) gives

bhat_j^(t+1) = -Quantile_(1-k/E)(s_:,j - alpha^(t)).

Appendix D expresses the same update using the required bias

r_ij = alpha_i^(t) - s_ij.

Indeed, s_ij + bhat_j^(t+1) > alpha_i^(t) iff bhat_j^(t+1) > r_ij.
Negating the margin reverses its ordering, so the update is equivalently

bhat_j^(t+1) = Quantile_(k/E)(r_:,j).

This explains both QB-specific router operations. Top-(k+1) exposes the token's current
competition boundary, while subtracting the raw score converts that boundary into the exact
bias threshold for each token/expert pair. A histogram of raw scores alone would discard the
token-local boundary created by the competing experts.

In particular, the histogram quantity is alpha_i - s_ij, not
alpha_i - (s_ij + b_j): alpha_i is a biased cutoff, but s_ij is the raw sigmoid score.
An underloaded expert consequently receives a relatively larger recovered bias, while an
overloaded expert receives a smaller one.

The report mean-centers the recovered biases,

b^(t+1) = bhat^(t+1) - mean(bhat^(t+1)) * 1,

because adding one common constant to every bias shifts every biased score and cutoff equally
and does not change Top-k. The update takes effect only in the next training step, so the
batch is never routed with a bias derived from itself.

Histogram approximation

The exact update would retain m * E required biases. Instead, for B uniform bins with
bounds [L, U], this PR accumulates

bin_ij = clamp(floor((r_ij - L) * B / (U - L)), 0, B - 1)
H[j, bin_ij] += 1.

Appendix D proves a natural per-step range. Since sigmoid scores are in (0, 1) and the
cutoff is one current biased score,

L = min(b^(t)) - 1
U = max(b^(t)) + 1.

The report uses B=1000, all-reduces the per-expert H[E, B] counts once per step, selects
the first bin whose cumulative count reaches ceil(q), and interpolates within it. If
beta_j is that bin's zero-based index, c_j is its preceding cumulative count, h_j is its
own count, and w = (U-L)/B, Appendix D recovers

bhat_j = L + (beta_j + clip((q-c_j)/h_j, 0, 1)) * w.

The result is then mean-centered. Its quantile error is bounded by one bin width. Counts
accumulate exactly across ranks and microbatches, so this recovers the pooled-global-batch
quantile rather than an average of per-rank quantiles.

TE receives [L, U] as a caller-owned CUDA tensor and only performs Top-(k+1), exact Top-k
output, and local histogram accumulation. Histogram all-reduce, within-bin interpolation,
bias update/centering, and next-step bounds update remain caller responsibilities.

Changes

• Add NVTEQBHistogramMode and C/PyTorch entry points for QB routing.
• Add a QB autograd wrapper selected only when qb_histogram, qb_bin_bounds, and
  qb_histogram_mode are all provided.
• Support BYTEMAP, BITMAP_U8, and caller-owned dense int16/int32/int64 Top-k indices.
• Support both the simple and radix Top-k kernels, including Top-8 and Top-16 cases.
• Reuse the existing FP32 sigmoid intermediate tensor for histogram inputs.
• Keep the existing fused-router backward kernel: the histogram and discrete route selection
  are non-differentiable, while selected sigmoid probabilities use the existing gradient.
• Add a pure-PyTorch QB reference and tests for both implementation modes, routing layouts,
  forward/backward parity, deterministic cutoff ties, bin clamping, accumulation across
  microbatches, and argument validation.

Histogram implementation choices

two_kernel writes one FP32 cutoff per token. A second kernel rereads the existing FP32 raw
scores, accumulates eight experts' bins in shared memory per CTA, and issues global atomics
only for nonzero shared bins. With B=1000, its dynamic shared-memory footprint is about
32 KB per CTA. This mode adds a launch, a 4 * num_tokens-byte cutoff buffer, and a read of
the [T, E] score tensor, but substantially combines contended updates before global memory.

fused_atomic keeps alpha_i in the router warp and directly issues one global int32 atomic
per token/expert pair in the router epilogue. It avoids the cutoff store, second launch, and
score reread, but global contention can become configuration-dependent. The caller-owned
histogram is only 4 * E * B bytes (3.584 MB for 896 experts and 1000 bins); neither mode
creates a 4 * T * E-byte bin-index buffer.

Both are compile-time QB specializations, so ordinary routing does not execute QB conditionals,
Top-(k+1), or histogram atomics.

Performance

Measured on an NVIDIA B300 SXM6 AC with the NVIDIA PyTorch 26.06 container,
nvidia-cutlass-dsl==4.5.0, and nvidia-cudnn-frontend==1.26.0. Each case uses
896 experts, Top-16, 1,000 histogram bins, FP32 logits, 100 warmups, 500 timed calls per
sample, and ten CUDA-event samples. The table reports median latency in milliseconds.

Routing output	Tokens	PyTorch QB	TE no QB	QB two-kernel	QB fused atomic	Fused / PyTorch	Fused / TE no QB	Fused / two-kernel
BYTEMAP	256	0.463196	0.031181	0.041287	0.036797	0.079	1.180	0.891
BYTEMAP	1,024	0.534397	0.032685	0.051062	0.039163	0.073	1.198	0.767
BYTEMAP	4,096	0.560477	0.063509	0.086649	0.084004	0.150	1.323	0.969
BYTEMAP	8,192	0.715650	0.105313	0.139695	0.134476	0.188	1.277	0.963
BYTEMAP	16,384	1.197801	0.194862	0.281411	0.250105	0.209	1.283	0.889
dense int16	256	0.442827	0.031265	0.040977	0.037445	0.085	1.198	0.914
dense int16	1,024	0.514043	0.032203	0.050953	0.039062	0.076	1.213	0.767
dense int16	4,096	0.538624	0.062701	0.085145	0.084522	0.157	1.348	0.993
dense int16	8,192	0.700264	0.103942	0.135977	0.132874	0.190	1.278	0.977
dense int16	16,384	1.181452	0.192812	0.276613	0.246737	0.209	1.280	0.892

The fused-atomic implementation takes 7.3% to 20.9% of the PyTorch QB latency
(4.8x to 13.6x faster). Relative to TE's existing no-QB router, the QB feature costs
18.0% to 34.8%; this includes both Top-(k+1) selection and histogram atomics. Relative to
the two-kernel QB path, fused atomic is 0.7% to 23.3% faster in every measured case. The
two modes are retained because larger expert/bin counts or different score distributions can
change global-atomic contention.

The performance comparison includes four implementations: the pure-PyTorch QB reference,
the existing TE router without QB, QB two_kernel, and QB fused_atomic. Every case checks
route, probability, and histogram correctness before timing. The TE-no-QB comparison isolates
the feature cost; the PyTorch-QB comparison shows the value of avoiding framework-level
intermediates and launches.

Validation

The focused QB matrix passes:

25 passed, 3648 deselected, 4 warnings in 23.34s

The entire fused-router test file passes:

3229 passed, 444 skipped, 4 warnings in 34.83s

The same B300 run also passed an actual CUDA tensor smoke test and confirmed that the loaded
extension exports both QB modes and all three opt-in Python arguments. Targeted pre-commit
checks (Python formatting, clang-format, whitespace, EOF, merge-conflict, large-file, and
Python-version checks) pass on the seven changed files.

Build configuration:

NVTE_BUILD_THREADS_PER_JOB=4
NVTE_CUDA_ARCHS="100;103a;"
NVTE_USE_CCACHE=1
/usr/bin/python3 -m pip install --no-build-isolation -e ".[test]" --verbose

Checklist

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding API docstring/header changes
• My changes generate no new warnings
• I have added tests that prove the feature works
• New and existing fused-router unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

The PR adds opt-in Quantile Balancing to the fused sigmoid MoE router, including Top-(k+1) cutoff selection and caller-owned histogram accumulation.

• Adds two-kernel and fused-atomic histogram modes across the common CUDA API and PyTorch binding.
• Supports routing maps and caller-owned dense index outputs while preserving the existing non-QB path.
• Adds bounds validation, CUDA-graph prevalidation, cross-device handling, and forward/backward correctness tests.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/common/fused_router/fused_topk_with_score_function.cu	Adds statically dispatched QB routing and histogram kernels, with checked and explicitly unchecked bounds paths; the previously reported bounds issues are addressed.
transformer_engine/common/include/transformer_engine/fused_router.h	Exposes QB modes and checked/unchecked common APIs with documented bounds-validation contracts.
transformer_engine/pytorch/csrc/extensions/router.cpp	Adds the QB binding with logits-device guarding, cross-device validation, output allocation, and common-kernel dispatch.
transformer_engine/pytorch/router.py	Adds the QB autograd path, argument validation, mutation-version-aware bounds validation, and CUDA-graph prevalidation behavior.
tests/pytorch/test_fused_router.py	Adds QB routing, histogram, accumulation, device, bounds, graph-capture, tie-handling, and gradient coverage.

Sequence Diagram

sequenceDiagram
    participant User
    participant Python as PyTorch router
    participant Binding as C++ binding
    participant Core as CUDA router
    participant Histogram as Caller histogram
    User->>Python: fused_topk(..., QB arguments)
    Python->>Python: validate/cache bin bounds
    Python->>Binding: QB forward
    Binding->>Core: Top-(k+1) sigmoid routing
    Core->>Core: retain Top-k and derive cutoff
    Core->>Histogram: accumulate expert/bin counts
    Core-->>Binding: probabilities and routes
    Binding-->>User: probabilities and routing output

Loading

_{Reviews (4): Last reviewed commit: "[Common][PyTorch] Validate QB bounds in ..." | Re-trigger Greptile}