[None][perf] Add BSX multi-tier CuTe DSL top-k decode kernels (stacked on #16457) - #16877
Draft
longcheng-nv wants to merge 6 commits into
Draft
[None][perf] Add BSX multi-tier CuTe DSL top-k decode kernels (stacked on #16457)#16877longcheng-nv wants to merge 6 commits into
longcheng-nv wants to merge 6 commits into
Conversation
…GVR decode op
Adds the op43 bsx CuTe DSL kernel family as a guarded fast path inside
trtllm::cute_dsl_gvr_topk_decode (op signature unchanged):
- gvr_topk_decode_bsx_tp.py: throughput GVR tier (gvr_topk_tp port) with
ragged-N lane masking, pre_idx clamping, in-kernel degenerate emit and
the mandatory cluster exit rendezvous.
- gvr_topk_decode_bsx_reg.py: register-resident tier (gvr_topk_reg port);
ragged-N handled by extending the OOR-lane -FLT_MAX idiom at the single
register load; packed u64 candidate pushes and aligned cluster barriers
preserved from the op43 convergence rounds.
- gvr_topk_decode_bsx_direct.py: short-row (npad <= 12288) exact radix
tier with ragged-N key substitution and degenerate emit.
- gvr_topk_decode_bsx_dispatch.py: route table transcribed from op42
gvr_bsx.cu; env knobs renamed GVR_BSX_* -> TRTLLM_BSX_*; v1 guard
(fp32, next_n=1, cr=4, no order_row/counters, K in {512,1024,2048},
npad <= 262144, npad % 64 == 0) plus a hardware cluster-size cap check
that falls back to the in-tree kernel instead of degrading silently.
- CI-sized exactness test covering every reg launch-table instance once,
direct/tp tiers, ragged rows with poisoned tails, quantized ties,
degenerate rows, pre_idx hardening, route-table asserts and dispatcher
fallback; registered in l0_b300.yml.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…gh-hit-rate rows) Port the in-tree R0 histogram-ladder ADMISSION concept (op#26, PR NVIDIA#16457) into the bsx throughput (tp) tier's pivot selection, closing the pr1 full-grid losses against the in-tree GVR kernel on real-capture rows at BS >= 16 while keeping the fused one-pass structure and every existing exactness invariant. Mechanism (measured on pr1 real-capture cells): the old 3-stage pivot pick targets ~3K sampled candidates inside a narrow [1.5K, 0.6kC] band. On real rows this either (a) picks a FAT rung (2-4x more P3 pushes and P4 candidates than needed - the v32/pro BS>=16 band), or (b) trusts a clustering-inflated sampled estimate and undershoots K, which fails the fused-pass reuse check and re-streams / secant-loops the full row (the flash_512k 1.6-1.8x losses; spatially clustered rows inflate a float4-sampled estimate up to 2.5x over the true count). Changes (gvr_topk_decode_bsx_tp.py only): * P2a stage-0 ADMISSION pick (R0 parity): accept the TIGHTEST ladder rung whose sampled-count confidence interval sits inside the [K, kC] acceptance window - the same "smallest exact count in [K, kC]" rule as the in-tree R0 admission, applied to the pre-pass estimates. The legacy 3-stage pick is unchanged as the fallback when no rung qualifies (cold-start / degenerate ladders take exactly the old path). * Clustering-aware sigma: sample_count now also tracks per-rung float4 OCCUPANCY, packed into the same per-thread accumulator (occ << 16 | cnt), giving the compound-Poisson sigma cnt/sqrt(occ) (equals the classic Poisson sigma on IID rows). Packing keeps registers, SMEM and the exchange at their pre-change sizes - an unpacked occ implementation measured a 14-25% whole-kernel regression (register spill in the streaming loops; A/B/bisect on v32_64k_L20). No field overflow: bsx guards npad <= 262144 => cluster-total cnt <= 8192. * Lower margin 2-sigma; K2048 uses 1.5-sigma (its [K, 4K] window is too narrow for 2-sigma to fire tight) backed by a RESCUE rung: the fused pass's second count column is now the next-fatter ladder rung instead of hmin (identical cost), so a pivot undershoot is caught with ONE collect re-stream instead of the multi-pass secant loop. Exactness machinery (accept window, P4 + tie tickets, plateau descent, ragged-N masking, degenerate emit, exit rendezvous) is untouched; admission never drops a top-K element because acceptance still requires the EXACT fused count in [K, kC]. Perf (nsys cold-L2 paired, same-rep A/B vs in-tree GVR kernel, real pr1 capture cells, B200; ratio = bsx/in-tree, lower is better): target cells before after flash_512k_L34 BS16 1.57 1.13 flash_512k_L34 BS64 1.72 1.11 flash_512k_L34 BS128 1.83 1.04 flash_512k_L34 BS256 1.54 0.73 flash_512k_L34 BS512 1.64 0.74 flash_512k_L34 BS1024 1.58 0.74 v32_32k_L50 BS16 1.40 1.21 v32_32k_L50 BS64 1.54 1.30 v32_32k_L50 BS256-1024 1.29-1.32 1.10-1.11 pro_128k_L54 BS16 1.27 1.19 pro_128k_L54 BS64 1.35 1.26 pro_128k_L54 BS128-1024 1.21-1.32 1.12-1.25 no-regression cells before after flash_512k_L36 BS1-128 0.55-0.94 0.56-0.92 flash_512k_L36 BS256-1024 0.67-0.68 0.66-0.67 v32_64k_L20 BS1-8 (reg) 0.63-0.65 0.63 v32_64k_L20 BS16-128 1.09-1.19 1.00-1.10 v32_64k_L20 BS256-1024 1.00-1.05 0.93-0.97 v32_32k_L04 BS16-1024 0.88-1.06 0.90-1.10 pro_256k_L30 BS16-1024 0.89-1.07 0.91-1.12 flash_16k_L26 BS1-1024 0.61-0.86 0.62-0.86 (reg/direct tiers untouched; the three +3-5% readings - pro_256k_L30 BS16/64, v32_32k_L04 BS64 - are inside the +-4.5% session noise floor measured on the fully untouched direct tier, e.g. flash_16k_L26 BS1 0.607 -> 0.634 with zero code change) The remaining v32_32k/pro_128k BS16-128 gap (1.19-1.30) is NOT the admission/pass-count mechanism: with admission the reuse check fires (one fused pass, admitted set slimmed 6678->2917 on v32_32k_L50) and a CS in {1,2,4} launch-shape sweep moves <= 7%. On those L2-resident shapes the in-tree kernel's whole-row-per-CTA structure is simply faster than the tp cluster split; closing it needs a tier-structure change, out of scope for this admission port. Exactness: bsx suite 56/56 (incl. new admission cases: hit-rate extremes on all three production shapes, tie plateau AT the admission threshold, count>kC overflow fallback, mixed admit/fallback ragged batch), in-tree gvr suite 671 passed / 144 skipped, 11/11 targeted screen (clustered/ties/zeros/oor/uniform/trivial), all pr1 A/B runs value-set-exact. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
… target
Fix round for the hint-ladder admission fast path (previous commit): the
full-grid pr2-vs-pr1 A/B (9515 cases) showed the admission closed the
deep tail (flash_512k_L34 1.83 -> 1.05) but regressed 396 cases by >5pp
and flipped 47 former wins, netting the overall gm to a wash.
Verified root cause (in-kernel decision probe + host-side P2a emulation
on the regression rows) - TWO distinct failure modes, both in the PICK:
1. Fat admit: the [K, kC] acceptance window let the tightest CI-passing
rung admit a candidate set 2.5x the legacy ~3K pivot-band target
(pro_64k_L06: admitted 4134 vs legacy 1671; pro_32k_L28 4047 vs
1554; flash_128k_L10 1970 vs 788), inflating P3 pushes + P4 scan.
2. Undershoot admit: the K2048 1.5-sigma lower margin accepted rungs
whose true count lands under K on clustered rows (v32_32k_L23 rung3
est 2720, sigma 374, TRUE count 1959 < 2048) - the rescue re-stream
then costs more than the fat-but-valid legacy pick it displaced
(22.6 -> 28.2us at BS64). The improved twin v32_32k_L50 has nearly
identical rung3 stats (est 2784, sigma 382) with true count 2917:
the pivot margin (+1.80 vs +1.92 sigma) is the only discriminant.
Changes (pick logic only; fused pass, rescue and all exactness
machinery are untouched from the admission commit):
* Admission upper bound tightened from kC to the legacy pivot-band hi
(0.6*kC), so admission only fires when the rung is genuinely tight.
* K2048 lower margin raised 1.5 -> 1.85 sigma: rejects the L23-class
undershoots (+1.80 sigma) while keeping the genuine tight admits
(L50 +1.92 sigma). K512/K1024 stay at 2.0 sigma.
* Stage 0b: the legacy band pick (min |est-tgt| in [1.5K, 0.6kC], the
exact pr1 stage-1 rule) now OVERRIDES the admitted rung when it is
strictly leaner and safe by its own clustering-aware 1.5-sigma lower
CI; with no admission it is taken as-is (pr1 parity).
* The pick's margin tests are now sqrt/div-free (squared comparisons):
the pick is a thread-0 serial section between two CTA barriers, and
the previous sqrt+fdiv chain measured 2-5% whole-kernel on
L2-resident accept-path rows.
Perf (nsys cold-L2 paired vs in-tree GVR kernel, real capture cells,
B200; ratio = bsx/in-tree, lower is better; bar = pr1 ratio + 3% for
the regression cells, absolute for the retention cells):
regression cells pr1 pr2 fixed bar
v32_32k_L23 BS16 1.202 1.521 1.198 <=1.238 PASS
v32_32k_L23 BS64 1.334 1.662 1.325-1.352 <=1.374 PASS
v32_32k_L23 BS1024 1.131 1.413 1.123 <=1.165 PASS
pro_64k_L06 BS256 1.062 1.352 1.061 <=1.094 PASS
pro_64k_L06 BS1024 1.099 1.414 1.121 <=1.132 PASS
pro_32k_L28 BS1024 1.059 1.295 1.084 <=1.091 PASS
flash_128k_L10 BS64 1.100 1.279 1.094 <=1.133 PASS
pro_32k_L02 BS64 0.652 0.659 0.652 <=0.672 PASS
retention cells pr1 pr2 fixed bar
flash_512k_L34 BS128 1.831 1.052 1.020 <=1.08 PASS
flash_512k_L34 BS256 1.537 0.730 0.725 <=0.80 PASS
v32_32k_L50 BS64 1.537 1.290 1.267 <=1.35 PASS
v32_32k_L50 BS16/1024 1.401/1.316 1.199/1.130 1.177/1.112
no-regression spots pr1 pr2 fixed
v32_32k_L04 BS64 1.057 1.086 1.032-1.068 (3 reps)
flash_16k_L26 BS1 (direct) 0.607 0.622 0.600
v32_64k_L20 BS8 (reg) 0.633 0.630 0.624
flash_512k_L36 BS64/256 0.944/0.674 0.898/0.656 0.883/0.655
pro_128k_L54 BS64/256 1.351/1.209 1.245/1.122 1.234/1.093
Known give-back: v32_64k_L20 BS16-128 (rung margin +1.53 sigma, below
the new 1.85 K2048 cut) reverts from the pr2 admission pick to the
legacy band pick; BS256-1024 keep the tight rung via stage 0b.
Falsified along the way (kept out): stashing the [rescue, pivot) band
in the candidate-buffer tail during the fused pass (per-element atomic,
per-warp stripes, and warp-aggregated ballot variants) - the fused
streaming loop cannot absorb ANY extra per-element code within the
+-3% bar (measured +4-10% whole-kernel on accept-path rows for all
three variants).
Exactness: bsx suite 56/56 (admission cases included; acceptance-rule
comment updated), in-tree gvr suite K2048 subset green, all A/B runs
value-set-exact (exact=True on every measured cell x BS x arm).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…guard narrowed
Extend the op43 bsx CuTe DSL tier family (direct/reg/tp) to the full
in-tree run_one_row contract: next_n >= 1 (MTP speculative-decode rows)
and compress_ratio in {1, 4}, and drop the next_n==1 / cr==4 conditions
from the dispatcher guard.
Semantics mirror the in-tree kernel line-for-line:
- request-level hint sharing: pre_idx row = row // next_n (pre_idx and
seq_lens are [num_rows // next_n, ...], validated by the guard);
- cr==1 temporal hint shift (row % next_n) + 1, with shifted-OOR hints
falling into the tiers' existing clamp hardening;
- per-row N_eff = (seq_lens[req] - next_n + row % next_n + 1) // cr
(already general in _row_n_eff).
next_n / cr are ctor constexpr: the next_n==1 / cr==4 hot path traces
identically to the v1 port (const_expr branches; no offset value is even
computed on cr>1 builds). All compile / dispatch / runner caches key on
(next_n, cr).
Validation (umbriel-b200-074, torch nv26.05):
- bsx suite extended with the MTP axis (next_n in {2,3,4} x cr in {1,4}
x three tiers x {random-garbage, noised-hint, cold-start-zeros, ties}),
each case checked against torch.topk host N_eff/offset simulation AND
a differential in-tree arm (order_row-forced, per-row value-multiset
equality): 140/140 passed.
- full in-tree GVR suite (fp32 MTP/cr=1 cells now auto-route to bsx):
671 passed, 144 skipped.
- next_n=1 perf anchors (ab_pr smoke, paired same-GPU back-to-back vs
stashed baseline): bsxd/gvrpr ratio deltas within +/-3% on
flash_16k_L26 BS1, v32_64k_L20 BS8, flash_512k_L34 BS128/256,
v32_32k_L23 BS64, pro_256k_L30 BS256; all exact.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
longcheng-nv
force-pushed
the
perf/gvr-topk-bsx-cutedsl-tiers
branch
from
July 28, 2026 14:18
1f4bdff to
4f95451
Compare
…-K interpolated tpush) and mask-hoisted streaming loops Ships the op43-pr6 campaign candidate (measured head 1a88ede77d) as a clean patch on the pr4 head: - stage 0c LEAN-PIVOT OVERRIDE (pick-only): on stage-0 CI-admission miss with npad in [16384, 262144], a strictly leaner ladder rung whose clustering-aware 1.5-sigma CI lands in [K, kC] becomes the pivot; undershoot costs one rescue re-stream (pr4 economics preserved). - per-K interpolated lean tpush (log2-count interpolation between pick and next-tighter rung; target 1.5K, 2.0K for K=512 sampling noise; npad <= 98304 gate) with original-pick rescue. - mask-hoisted streaming loops in count/max-below/fused-count-collect/ collect (unmasked main loop + masked vec-tail) - uniform micro-gain, independently exonerated of the pr5 regression by the C5 fingerprint ablation. - candidate budget pinned at the flat pr4 kC (8192 for K>=2048, else 6144); wide pr4 ladder quantiles pinned. Full-grid evidence (865 real decode cells x 11 BS, same-rep cold-L2 paired vs the in-tree kernel): gm 1.4170 (pr4 head 1.4120), win 88.0%, <0.909 cases 517 -> from 636. Exactness: 140-test production suite + adversarial + synthetic gates all PASS; measured module and this file are AST-identical modulo dead scaffolding removal. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…e kernel Routes the (npad, bs) buckets where the op43-pr6 full-grid verdict found at least one production layer >1.10x slower on every bsx tier than the in-tree kernel (L2-resident mid-N regime: the in-tree exact-count ladder admits a leaner candidate set and its row-slice cluster split keeps all CTAs busy through P4). With the table the operator is a strict Pareto improvement over the in-tree kernel: routed buckets run at parity, everything else keeps the bsx win (full-grid gm 1.40 vs the in-tree head, worst case capped at 1.10x). npad keys resolve by nearest power of two; TRTLLM_BSX_FALLBACK_BANDS=0 disables the table. Band data: 865 real decode cells x 11 BS same-rep cold-L2 nsys pairs, 2026-07-28. Tests: kernel-contract cases pin the table off via an autouse fixture (they exist to exercise the bsx tiers); the table itself is covered by test_bsx_fallback_band_table (bucket membership, nearest-pow2 resolution, neighbour non-routing, kill-switch), and CUDA-graph capture/replay by test_bsx_cuda_graph_capture_replay on both sides of the table (host-side dispatch bakes at capture; replay stays exact over in-place rewritten inputs). Suite: 143 passed; sibling test_cute_dsl_gvr_topk_decode suite: 671 passed / 144 skipped. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
longcheng-nv
force-pushed
the
perf/gvr-topk-bsx-cutedsl-tiers
branch
from
July 28, 2026 14:34
4f95451 to
78df9cb
Compare
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
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.
Description
Rebased onto
main(#16457 merged asf3d4c85461). Six commits: the four reviewed bsx-tier commits (e92519d226tier port,a196684b8atp admission fast path,a8f4899bcfacceptance-bound fix,7c9cd2a4a3MTP/compress_ratio=1) plus the op43-pr6 campaign result:59a019bb70lean-pivot admission + mask-hoisted streaming loops,78df9cbb90measured fallback band table + tests.Adds the BSX multi-tier CuTe DSL top-K decode kernel family as a guarded fast path inside the existing
trtllm::cute_dsl_gvr_topk_decodecustom op (op signature unchanged, no production call-site change — DSL-path activation remains the deferred follow-up, same as #16457):gvr_topk_decode_bsx_direct.py— short-row (npad <= 12288) single-CTA exact radix tiergvr_topk_decode_bsx_reg.py— register-resident tier (one-time row load, rank-scatter P3, packed-u64 candidates)gvr_topk_decode_bsx_tp.py— throughput tier (sampled pivot + multi-rung secant + fused count/collect)gvr_topk_decode_bsx_dispatch.py— route table + v1 guard; anything outside the guard falls back to the existing [None][perf] GVR top-K decode: enable R0 histogram-ladder admission by default #16457 kernel unchangedDispatcher guard: fp32, no load-balance (
order_row/counters),K in {512, 1024, 2048},npad <= 262144, cluster size within_query_max_cluster_size(). MTP (next_n >= 1) andcompress_ratio in {1, 4}are fully supported (4th commit): hint-row sharingrow // next_n, cr=1 temporal hint offset, per-rowN_eff— all mirroring the in-tree kernel's semantics line-by-line;next_n==1/cr==4hot paths are constexpr-guarded to compile to exactly the previous code. MTP correctness is fully unit-tested (differential vs the in-tree kernel per row); MTP performance validation is deferred to e2e integration — the full-grid re-measure below confirms the next_n=1 envelope is unchanged (9490/9515 cases flat within 5pp vs the pre-MTP head).Production-contract adaptations (vs the op-bench origin): ragged per-row
N_efflane masking (poisoned-tail safe),pre_idxout-of-range clamping (all-zero cold start safe), in-kernel degenerate-row (N_eff <= K) identity emit with-1padding, cluster-size clamp with dispatcher fallback.Evidence (local, B200; text tables only — no profiler artifacts)
Benchmark: 865 real DSv4/V3.2 decode-capture cells x 11 batch sizes = 9515 paired same-rep cold-L2 nsys cases, 0 skips (every cell falls inside the v1 guard), 0 exactness mismatches.
Direct A/B: this PR's head (pr6: 4 tier commits + lean-pivot + band table) vs the #16457 kernel now on
main— full 865x11 real decode-capture grid (9515 same-rep cold-L2 nsys pairs), re-measured 2026-07-28:Bare-kernel loss tail vs the in-tree head: <0.909 cases 636 (pr4) -> 517; the shipped operator routes the 27 (npad, bs) buckets containing them (30.9% of grid cases, the L2-resident mid-N regime) back to the in-tree kernel — the operator is a strict Pareto improvement: routed buckets run at parity, everything else keeps the bsx win, floor 0.9125. The pr6 campaign additionally falsified (measured: 16-subset ablation + component fingerprints): exact-count admission as a wholesale replacement (gain and harm co-sourced), K-scaled kC diet (pure harm), ladder-quantile re-spread, CS occupancy cuts, and multi-pass straggler hypotheses (the residual band is 1-pass; the gap is candidate-set fatness — hence the lean pivot).
fp32 gm by model x batch size (bare kernels; the shipped-operator numbers differ only in the routed mid-N columns by 2-5pp):
Layer dimension (per-layer gm across that layer's own 9515-grid slice — the loss tail is a per-(layer, shape) phenomenon, not a per-layer one): all 109 production layers win at layer granularity, none below 1.28x — v32: 58 layers, worst L26 = 1.286 / median 1.444 / best L52 = 1.544; flash: 21 layers, worst L06 = 1.297 / median 1.380 / best L24 = 1.490; pro: 30 layers, worst L16 = 1.303 / median 1.388 / best L10 = 1.594.
dtype scope. All numbers above are fp32 — the production DSA indexer logits path, and the bsx guard intentionally restricts the fast path to fp32: bf16/fp16 calls flow to the in-tree kernel unchanged (functional support only in this PR, verified by the dispatcher-fallback test). Non-fp32 bsx tiers are future work if a production 16-bit indexer materializes.
Test Coverage
tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py(143 cases): all 14 reg launch-table instances, fallback band-table routing (bucket membership, nearest-pow2 resolution, kill-switch), CUDA-graph capture/replay on both sides of the band table, admission fast path (high/low hit-rate, tie plateau at the admission threshold, overflow fallback, mixed admit/fallback ragged batch), MTP axes (next_n {2,3,4} x cr {1,4} x all tiers x 4 input classes, each with dual oracles: host-emulated torch.topk AND per-row differential against the in-tree kernel), MTP hardening (all-zero/garbage hints) and degenerate-row boundaries, direct/tp incl. trivial path, ragged varlen with poisoned tails (+1e30), quantized-tie plateaus, degenerate mixed batches, all-zero/OORpre_idx, route-table + env-knob asserts, bf16 /next_n=2fallback-to-in-tree checks. Registered inl0_b300.yml.test_cute_dsl_gvr_topk_decode.py: full suite green on this branch (671 passed / 144 skipped) — fp32/K2048/cr4 cells of that suite now exercise the bsx route.Draft checklist (before ready-for-review)
PR Checklist
🤖 Generated with Claude Code