[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, writes one FP32 cutoff per token to honor
the common API output contract, and directly issues one global int32 atomic per token/expert
pair in the router epilogue. It avoids the 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, and 1,000 timed calls
split across 20 CUDA-event samples. The table reports milliseconds and averages the median
latency from two cutoff-writing runs that bracketed a no-write control.

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.460507	0.030977	0.041409	0.036838	0.080	1.189	0.890
BYTEMAP	1,024	0.530411	0.032497	0.050846	0.039474	0.074	1.215	0.776
BYTEMAP	4,096	0.559368	0.063485	0.085612	0.085905	0.154	1.353	1.003
BYTEMAP	8,192	0.713347	0.106315	0.141176	0.134075	0.188	1.261	0.950
BYTEMAP	16,384	1.198404	0.195790	0.280938	0.248273	0.207	1.268	0.884
dense int16	256	0.441163	0.031119	0.040860	0.037275	0.084	1.198	0.912
dense int16	1,024	0.510570	0.032113	0.051104	0.039511	0.077	1.230	0.773
dense int16	4,096	0.534869	0.062886	0.085992	0.086992	0.163	1.383	1.012
dense int16	8,192	0.699211	0.104915	0.137759	0.132246	0.189	1.261	0.960
dense int16	16,384	1.177213	0.195510	0.278774	0.245874	0.209	1.258	0.882

The fused-atomic implementation takes 7.4% to 20.9% of the PyTorch QB latency
(4.8x to 13.5x faster). Relative to TE's existing no-QB router, the QB feature costs
18.9% to 38.3%; this includes Top-(k+1) selection, the per-token cutoff store, and histogram
atomics. Relative to the two-kernel QB path, fused atomic is 4.0% to 22.7% faster in eight
cases and 0.3% to 1.2% slower in the two 4,096-token cases. The two modes are retained
because larger expert/bin counts or different score distributions can change global-atomic
contention.

An A/B/A comparison against the otherwise identical fused-atomic kernel without the cutoff
store measured a mean 0.68% overhead across the 1,024-to-16,384-token cases, with individual
results from 0.44% to 1.03%. The 256-token cases were within sub-percent measurement noise.

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 histogram collection to the fused sigmoid MoE router while preserving the existing non-QB path.

• Adds Top-(k+1) QB routing with deterministic cutoff compaction and two histogram accumulation modes.
• Exposes the feature through the common C API, PyTorch binding, and autograd wrapper.
• Adds validation, multi-device, CUDA graph, routing-layout, histogram, and backward-parity coverage.

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 cutoff extraction, histogram kernels, recoverable bounds validation, and checked common entry points.
transformer_engine/common/include/transformer_engine/fused_router.h	Declares the QB histogram modes and checked or prevalidated common API variants.
transformer_engine/pytorch/csrc/extensions/router.cpp	Adds the guarded QB PyTorch binding with consistent-device validation, output allocation, and common-core dispatch.
transformer_engine/pytorch/router.py	Adds the opt-in QB autograd path, argument validation, bounds-version caching, and histogram-mode selection.
tests/pytorch/test_fused_router.py	Covers QB numerics, gradients, layouts, accumulation, bounds errors, multi-device execution, ties, clamping, and CUDA graph capture.

Sequence Diagram

sequenceDiagram
  participant Caller
  participant Python as PyTorch router wrapper
  participant Binding as C++ binding
  participant Kernel as Fused CUDA router
  participant Hist as Caller-owned histogram
  Caller->>Python: logits, bias, histogram, bounds, mode
  Python->>Python: Validate QB arguments and bounds
  Python->>Binding: QB forward
  Binding->>Binding: Guard logits device and validate tensors
  Binding->>Kernel: Top-(k+1) selection
  Kernel->>Kernel: Drop cutoff route and normalize Top-k
  Kernel->>Hist: Accumulate required-bias bins
  Kernel-->>Caller: probabilities and routing output

Loading

_{Reviews (5): Last reviewed commit: "[Common][PyTorch] Write QB cutoff in bot..." | Re-trigger Greptile}

[denera]

LGTM besides needed clarification in one area. Please see comment.