Skip to content

[None][perf] GVR top-K decode: enable R0 histogram-ladder admission by default#16457

Open
longcheng-nv wants to merge 50 commits into
NVIDIA:mainfrom
longcheng-nv:perf/gvr-topk-r0-histogram-ladder
Open

[None][perf] GVR top-K decode: enable R0 histogram-ladder admission by default#16457
longcheng-nv wants to merge 50 commits into
NVIDIA:mainfrom
longcheng-nv:perf/gvr-topk-r0-histogram-ladder

Conversation

@longcheng-nv

@longcheng-nv longcheng-nv commented Jul 15, 2026

Copy link
Copy Markdown
Collaborator

Note

#16424 (the semantic dependency) has merged into main (9095cc11d1); this branch is synced with current main, so the stacking is resolved — the PR's net diff vs main is exactly the 3 files listed under Scope. CI is green at the current head (dc2bdc8a9f, pipeline #49596).

Summary

Adds an R0 histogram-ladder admission fast path + a fused rank-and-scatter Phase-4 writeback to the production Blackwell cuTe-DSL GVR top-K decode kernel (GvrTopKKernel), and makes R0 the default (enable_r0=True).

Headline (real production decode captures, B200): 1.33× geomean over the shipped kernel, faster on 25/25 cells, exact on the full 2772-cell audit grid — and it repairs two correctness defects of the shipped kernel (a low-hit undershoot and a boundary-tie miss; see Correctness).

What's inside:

  • R0 admission: replaces the Phase-2 secant threshold search with a single-pass multi-threshold rung ladder seeded by a 256-bin histogram over the prev-topK gathered values; fused rank-and-scatter Phase-4 (cluster barriers ~14 → ~7).
  • Virtual seed rung (r0_vseed): P1's mean probe is folded into the count pass as a free extra rung (zero SMEM growth), adapting admission per row — fixes the cold-hint fat-admission regime.
  • p4_exact_tail (fp32 default on): bit-exact boundary-tie resolution (ambiguity-gated radix select + tiny-tie fast path); 16-bit kernels byte-identical.
  • K2048 tuning: low rung 0.85 → 0.6 recalibrated on real captures + Phase-4 histogram 2048 → 512 bins.
  • The classic secant path is retained verbatim (enable_r0=False); this PR flips the kernel default only — call-site/dispatch is untouched (guard = follow-up PR).

Performance

Methodology. nsys pure-kernel, cold-L2, single-GPU same-run A/B (R0 vs retained secant on identical inputs), 20 cold reps, B200. Both arms launched via the kernel's own launch/pick_config contract (the shapes the production runner picks). Re-measured 2026-07-20 at the current head. Inputs = per-layer indexer top-K captured from production DeepSeek-V4 Flash / Pro / V3.2 BS=1 greedy decode (9 ISL rungs 4k–1024k; N = post-compress indexer length, V4 cr=4 / V3.2 cr=1; hit = preIdx∩topK / K; V3.2 preIdx = previous step's top-K). V3.2 has 7 rungs: its 160K (163,840-token) max context truncates longer prompts — the 256k row's N=163,775 is the exact valid kv length at the benched step. A calibrated synthetic envelope is used only for exactness/audit coverage (Correctness, Known limitation).

BS=1 fp32, per seq-len — geomean 1.333×, 25/25 ≥ 1.0, exact 25/25

V4 Flash (K512) — geomean 1.285×

ISL N hit base (µs) R0 (µs) speedup
4k 1,027 0.64 8.97 7.88 1.14×
8k 2,051 0.33 8.82 8.38 1.05×
16k 4,099 0.34 11.78 9.02 1.31×
32k 8,195 0.69 11.74 9.37 1.25×
64k 16,387 0.08 14.71 12.35 1.19×
128k 32,771 0.70 16.30 12.03 1.36×
256k 65,538 0.28 17.26 13.91 1.24×
512k 131,075 0.06 43.46 20.74 2.10×
1024k 262,127 0.42 20.84 18.02 1.16×

V4 Pro (K1024) — geomean 1.299×

ISL N hit base (µs) R0 (µs) speedup
4k 1,027 1.00 19.21 10.05 1.91×
8k 2,051 0.46 9.30 9.21 1.01×
16k 4,099 0.74 15.35 10.08 1.52×
32k 8,195 0.53 18.66 11.56 1.61×
64k 16,387 0.31 12.33 11.73 1.05×
128k 32,771 0.33 16.87 13.61 1.24×
256k 65,539 0.36 17.01 15.16 1.12×
512k 131,075 0.23 21.59 17.64 1.22×
1024k 262,127 0.27 25.03 20.11 1.25×

V3.2 (K2048) — geomean 1.443×

ISL N hit base (µs) R0 (µs) speedup
4k 4,111 0.73 16.06 10.86 1.48×
8k 8,207 0.84 15.46 11.90 1.30×
16k 16,399 0.53 28.00 13.76 2.03×
32k 32,783 0.43 20.62 16.78 1.23×
64k 65,551 0.41 31.96 20.68 1.55×
128k 131,087 0.62 21.56 18.22 1.18×
256k 163,775 0.87 28.38 19.11 1.48×

All cells exact vs the captured reference. Flash 512k: the base kernel is additionally inexact there (undershoot) — R0 is exact and 2.10× faster.

BS scaling — R0/base geomean per BS (all captured ISL rungs × 3 dtypes; 825 cells, exact 825/825)

BS 1 2 4 8 16 32 64 128 256 512 1024
Flash (K512) 1.272 1.263 1.264 1.261 1.258 1.254 1.248 1.240 1.278 1.295 1.314
Pro (K1024) 1.234 1.236 1.234 1.232 1.230 1.227 1.214 1.195 1.204 1.204 1.214
V3.2 (K2048) 1.267 1.274 1.270 1.264 1.259 1.264 1.236 1.207 1.227 1.260 1.258
all 1.257 1.256 1.254 1.251 1.248 1.247 1.232 1.215 1.237 1.252 1.262
  • R0 is BS-invariant (flat 1.257× → 1.262×, floor 1.215× at BS=128): it changes Phase-2/4 arithmetic only — grid shape, smem, and cluster semantics are identical to base.
  • Launch shape, not the kernel, dominates at large BS. A config frozen at the BS=1 optimum is geomean 2.38× (max 5.8×) slower than per-(BS, N) picks. This PR therefore ships the policy as part of the kernel: pick_config (the (dtype, BS, N) → launch-shape classmethod, incl. the CUDA-graph max_seq_len contract) and launch (compiled-variant cache; **kernel_overrides for forcing knobs). cluster_size=8 (previously untested) validated: 78/78 exactness cells, beats forced cs=4 on 8/8 nsys cells.

Known limitation + follow-up

  • The residual R0 losses are confined to the cold-hint stress tail (synthetic hit ≈ 0.05), a regime not observed in any production capture (measured hit 0.27–1.00): full 2772-cell audit grid geomean 1.190×, fp32 never below 0.915×, and only 4 cells < 0.90× (min 0.887×, all bf16/fp16). An earlier revision's fat-admission regression (0.68–0.79× on Flash-1M big-BS) and 106-cell 16-bit tail are fixed by the seed rung + K2048 recalibration.
  • Follow-up PR: dispatch guard routing residual regimes to the retained secant path. Hit-rate is not observable at inference time, so the guard is an in-kernel admission escape and/or trailing-step feedbackenable_r0=False is exactly what it dispatches to.

Correctness

  • Every timing cell asserts a tie-aware value-set match vs torch.topk on the same run (plus in-range / no-duplicate / n_below=0 guards). Real 25/25, synthetic envelope 52/52, full audit grid 2772/2772 (re-verified at the current head).
  • Fixes a real boundary-tie defect: the pre-fix Phase-4 kept an arbitrary member of a sub-resolution tie set straddling the top-K boundary (real Pro-512k captures, 1-element miss, |dv| ≈ 3e-6). p4_exact_tail makes all 12 affected audit cells bit-exact; unaffected-cell cost is noise (0.998 geomean); the tiny-tie fast path returns repair-active rows to a win (Pro 512k: 1.22×).
  • Repairs the base undershoot: on real Flash-512k (hit ≈ 0.06) the shipped secant returns < K unique indices on 36 audit cells — R0 is exact on all of them.
  • New SM100 unit tests: R0/secant equivalence (both arms valid top-K + identical index sets on tie-free fp32) across dtype × K × N × BS × hint × cluster_size ∈ {1,4,8}; big-BS multi-wave grids; pick_config policy lock; launch() autoconfig incl. forced-secant arm; adversarial 5e-8 / 1-ulp tie bands (p4_exact_tail). A standalone harness independently confirms 186/186.

Scope / risk / rollback

  • Touches only gvr_topk_decode.py, gvr_topk_decode_load_balance.py (a 36-line share of the pick_config/launch plumbing), and the unit-test file; no call-site, dispatch, or config-schema changes.
  • enable_r0=False restores the pre-R0 kernel byte-for-byte (all R0 fields const-folded) — rollback = flip one default.
  • r0_vseed=False keeps the static ladder; p4_exact_tail defaults off for 16-bit (byte-identical there).

Test plan

pytest tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py

on SM100 (B200 / B300).

The suite is slimmed ~33% (2428s -> 1626s local B200, two consecutive full-green runs) so the CI bucket fits its default 60-minute budget without any test-db TIMEOUT bumps: input memoization, the seqlen_sorted axis moved from the full cross-product to a 4-cell representative test, and test_lb_vs_reference reduced to the two SMEM-layout endpoint dtypes. All protected coverage (shortrow degrade, p4_exact_tail ties, r0-equivalence incl. cs=8/big-BS, pick_config/launch policy locks) is unchanged; remaining runtime is dominated by first-instance cuteDSL compilation of the 133 kernel-variant signatures.

longcheng-nv and others added 23 commits July 15, 2026 02:14
…nert)

First of a series porting the the reference prototype histogram-ladder R0 admission fast
path into the unified GvrTopKKernel. Adds the multi-threshold GE-count
primitive block_count_ge_multi<M> (reuses block_count_ge's vectorized
memory path with M static register counters, caching each per-thread count
column) and the constructor knobs enable_r0/r0_qfracs/mt_unroll/p1b_cache/
fb_fix plus their derived rung targets.

All knobs default OFF and the helper is not yet called from any kernel
body, so a default-constructed kernel is byte-identical to the base; this
commit has no behavioral effect. The rung placement (P1b), admission
dispatch, cached-column Phase 3, log-falsi fallback and cluster-aware
reduction land in follow-up commits.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…ent (inert)

Adds phase1b_hspace_rungs: a 256-bin SMEM histogram over the prev-topK
gathered values banded by P1's pmin/pmax, with warp-0-parallel extraction
of M h-space quantile rungs (crossing qneeds[m] from the top). These rungs
are the candidate thresholds the single-pass R0 admission evaluates.

Not yet called from any kernel body (the enable_r0 dispatch lands next), so
the base kernel remains byte-identical; no behavioral effect.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…s (secant fallback)

Wires the R0 histogram-ladder admission into _run_phases as a fast path:
P1b places M h-space rungs, one block_count_ge_multi pass counts all M at
once, and the tightest rung with count in [K, kC] is accepted, seeding
Phase 3 with its cached per-thread count column (zero rescan). An R0 miss
falls back to the existing phase2_secant_search (a later commit replaces
that fallback with an inline log-falsi refine). Five R0 scratch buffers are
allocated in run_one_row (gated on enable_r0, None otherwise) and threaded
through _run_phases with the same const_expr-gate + None-propagate idiom the
kernel already uses for s_cluster_partial / smem_input.

Gated to enable_r0 and cluster_size==1 (the earlier R0 scans the full row in a
single CTA; the slice-parallel + cluster count-merge variant for the cs>1
long-row branch lands in a follow-up commit). enable_r0 defaults False, so
the compiled path is the unchanged secant search and the default build is
byte-identical to the base.

NOTE: unlike the two preceding commits, the enable_r0=True path executes new
device code and is NOT yet silicon-validated; exactness + nsys gates on a
B200 are required before enabling it by default.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
On an R0 admission miss, derive the measured [lo, hi] bracket from the
rungs and, when both ends were measured, fire ONE extra count pass at a
log-count regula-falsi aim (geometric center of the [K, kC] window). If it
lands in-window the row is resolved with smem_ptcnt fresh from that pass;
otherwise (or when an end was unmeasured) the existing secant fallback runs
from the untouched P1 bracket, so the fallback state is identical to before
this commit. Adds only the R1 fast-exit for near-boundary misses.

Entirely inside the enable_r0 && cluster_size==1 branch, so the default
build (enable_r0=False) is unchanged. Enabled path is not yet
silicon-validated; gated behind the same pending B200 exactness/nsys run as
the R0 admission itself.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…laces secant under R0)

Replaces the R0-miss secant fallback (when fb_fix=True, the default)
with the earlier bounded log-count regula-falsi refine: iterate block_count_ge
with a log2_mstar-aimed threshold update over a [pmin, pmax] bracket whose
end counts start unknown and get measured, converging into [K, kC] and
setting done=1 so upstream Phase 3 skips its retry-shrink. fb_fix=False
keeps the secant fallback. Adds fb_alpha/log2_mstar (interior aim).

Inside the enable_r0 && cluster_size==1 branch, so the default build
(enable_r0=False) is byte-identical.

CORRECTNESS-CRITICAL: the exhaustion path lands on the measured undershoot
side for value-tie plateaus; this is the case the kC >= 5K 16-bit tie
contract protects and it is NOT yet silicon-validated. This commit must be
gated behind the tie-exactness gate on a B200 before enable_r0 is defaulted
on; treat it as the highest-risk commit in the series.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
Makes the R0 admission cover the multi-CTA (cluster_size>1) path so the
load-balance kernel's long-row branch gets R0 too. block_count_ge_multi now
scans [slice_start, slice_end) and DSMEM all-reduces the M rung counts
across the cluster with a RELEASE cluster_arrive (upstream's block_count_ge
comment warns cluster_arrive_relaxed skips the release; the earlier mc used the
relaxed form). A per-cluster s_cluster_partial_m[M] scratch is allocated
(gated enable_r0 && cs>1) and the R0 / R1 / fb_fix count calls are all
slice+cluster aware. phase1b rungs stay per-CTA identical (preIdx stats are
full-row). The enable_r0 && cluster_size==1 gate is relaxed to enable_r0.

Default build (enable_r0=False) remains byte-identical: every R0 buffer and
branch is const_expr-gated on enable_r0.

Validated on B200: cluster-R0 exactness sweep 144/144 (fp32/bf16/fp16 x
cs{2,4} x K{512,1024,2048} x N up to 262144 x {real,random preIdx} + tie
plateaus), no deadlock, vs torch.topk. Cumulative exactness 330/330.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…cs + kC-diet + p1b_cache dtype-gate)

Folds the earlier host dispatch policy into the ctor, all gated on enable_r0
so an OFF kernel stays byte-identical:
 - r0_qfracs default = (0.85, 0.35) (0.85, 0.35): the reference dispatch policy ships this M=2
 ladder for every (dtype,K,N); the M=2 count pass is ~free and the R1
 falsi shot covers the 3-7% bracket misses. uh4 (M=4) was silicon-
 falsified (mc geomean 0.956 - admission != latency).
 - p1b_cache default = (dtype != fp32): the gather-cache wins +0.8-2.8% on
 16-bit; flat/negative on fp32. (The cached path itself lands in; the
 flag is inert until then.)
 - kC-diet: K512 single-CTA -> kC=3072 (saves 16KB SMEM). kC>=2560 is the
 K512 16-bit tie-safety contract, so 3072 is safe; cluster + K1024/K2048
 stay stock.

Validated on B200: default-dispatch exactness 48/48 (fp32/bf16/fp16, all K,
incl. K512 aggressive tie plateaus at kC=3072). Corrects the prior
benchmarking that used the falsified uh4 ladder.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…bit)

Implements the p1b_cache fast path that 's dispatch enables for 16-bit:
P1 (phase1_preidx_stats) optionally stashes each gathered preIdx value into
a SMEM buffer (smem_gath, top_k f32) via a new const_expr-gated store, and a
new phase1b_hspace_rungs_cached builds the rung histogram straight from that
cache -- skipping P1b's second GMEM random gather (the expensive part in
half precision). smem_gath is allocated in run_one_row gated on
enable_r0 && p1b_cache and threaded through _run_phases; the phase1b call
dispatches cached vs re-gather on the p1b_cache const_expr.

The optional smem_gath param defaults to None, so with p1b_cache OFF (all
fp32, and the whole base) the store branches are const_expr-eliminated and
the kernel is byte-identical.

Validated on B200: 16-bit default-dispatch exactness 48/48 (bf16/fp16 now
exercise the cached path, incl. K512 tie plateaus at kC=3072).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…f to Phase 3/4

The cluster (cs>1) R0 path was silently WRONG on an admission/R1 hit and
only produced correct output because it always missed and fell back to the
secant path (exactness masked the bug). ncu root-cause: forcing a hit gave
uniq 592/1024. The R0/R1 accept paths omitted the per-CTA cluster-gather
state that block_count_ge sets and Phase 3/4 consume:
 - s_iscalars[1]=1 (done): tells Phase 3 to honor the admitted threshold and
 skip its retry-shrink (missing => the cluster collect diverged).
 - s_iscalars[5]=local slice count for the chosen rung (the per-CTA
 cand_count the cluster gather sums).
 - R1 accept now also sets done=1 (block_count_ge already sets [0]/[5]).
 - R0-miss persists the measured rung bracket into s_thr[1]/s_thr[2] so the
 fb_fix refine starts tight (the reference prototype behavior) instead of P1's wide
 [pmin,pmax].

Validated on B200: cluster-R0 forced-hit now exact; full grid 60/60 exact
(fp32/bf16/fp16 x best/worst x K512/1024/2048 x cs1/cs4). Perf note: at
large N the (0.85, 0.35) (M=2) rungs straddle [K,kC] (rung0 count > kC, rung1 < K),
so cs>1 resolves via the refine (the earlier 1.14x ceiling); cs=1 small-N R0
still wins ~1.5x. enable_r0 default OFF, so base is byte-identical.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…, no re-measure)

Replaces the R0-miss 'R1 single shot + fb_fix re-measure loop' with one
bounded log-falsi loop SEEDED from the measured rung bracket + its known
counts (clo/chi) so it does regula-falsi from iter 0 with no re-measure and
no redundant R1 pass. At large N ((0.85, 0.35) rungs straddle [K,kC]) this is the
resolve path; seeding trims passes. exact 5/5 spot; cs4 ncu 49.4->48.1us
(the reference prototype mc-R0 41.6). enable_r0 default OFF -> base byte-identical.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…n converged snap

P3 stream-write (dense+compact), P4 keys/vals accesses, and the snap-loop
convergence check now go through raw integer addressing off a per-method
iterator.toint() base (_smem_ld/_smem_st helpers) instead of tensor
indexing, which re-derived the cluster SMEM window (S2R SR_CgaCtaId +
LEA<<24) around hot accesses.

P4 writeback: when the snap loop converged, the final iteration's cgt
(s_iscalars[3]) is the exact strictly-greater count at sel_thr, so gt
entries pack via counter[4] from 0 while ties pack via counter[5] from
offset cgt — ONE candidate sweep and one barrier instead of two-pass.
Same [gt | eq | pad] output partition; non-converged snap falls back to
the original two-pass (its cgt would be stale).

Measured (B200, real V4-Pro/Flash production shapes): P4 7.80->7.58us,
wall +1-2% on top of fastsync (my-pro fs2 geomean 1.032->1.043, my-flash
fs0 1.063->1.075). Full sweep 450/450 PASS. ncu follow-up: the residual
snap-top S2R stall is a post-barrier pipeline-refill accounting artifact
plus the un-hoisted scalar scratch accesses — see analysis notes.

Made-with-Claude-Code: fable-5
Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
The P4 snap loop steps one distinct candidate value per iteration
(~0.45us: full candidate re-scan + 2 barriers). Real logits concentrate
count mass at the k-th boundary, so the selected 1/1024 bin often holds
tens of values; snap stragglers of 10-20us set the wall clock at
N<=32K (P4 max 15.8us at N=8K vs 7.6us mean on synth aggregate cells).

When the selected level-1 bin holds >8 candidates, rebuild the
histogram over just that bin (+2% slop; bin width shrinks 1024x) and
re-run the 3-step k-th search before entering the snap loop. Costs one
extra candidate scan (~1us), leaves snap 0-2 steps. The snap loop
converges monotonically from any start, so exactness is untouched.

Histogram build + 3-step search factored into _hist_build /
_kth_bin_search (trace-identical to the previous inline code).

Synth sweep (B300, bs64 K=1024 aggregate cells, us):
  2K 12.24->10.48  4K 15.14->12.97  8K 20.32->13.60
  16K 24.88->17.07  32K 32.99->21.57  64K 42.35->39.99
  128K/256K flat (P2-dominated). Worst bound-gap 2.09x -> 1.48x.

Made-with-Claude-Code: fable-5
Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
Two follow-ups to the level-2 histogram refinement (eed5eb5931):

1. Refinement chain: level-2 gate lowered 8->2 (one rebuild breaks even
   vs ~2 snap steps) and a level-3 pass added (gate >8, fires only when
   level 2 fails to split the bin: heavy ties or sub-ulp window).
   Unrolled 2-iteration chain over (gate, binw_cur).

2. _kth_bin_search step 1: every lane redundantly walked a
   bins_per_warp-deep serial LDS+IADD chain (~7% of stall samples at
   N=8K). Now each lane sums bins_per_warp/32 bins + one REDUX warp
   reduce; const_expr fallback to the serial walk when bins_per_warp
   is not a multiple of 32.

Synth sweep (B300, bs64 K=1024 aggregate cells, us, vs eed5eb5931):
  2K 10.48->9.10  4K 12.97->10.71  8K 13.60->12.42  16K 17.07->15.80
  32K 21.57->21.25  64K 39.99->38.14. Worst bound-gap 1.48x -> 1.43x
  (128K, P2-straggler dominated); N<=32K all within 1.08-1.27x.
Cumulative vs pre-refinement baseline: 8K 20.3->12.4 (-39%),
16K 24.9->15.8 (-36%), 32K 33.0->21.3 (-36%).

Made-with-Claude-Code: fable-5
Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
…back

1. Histogram window from the P2 exit bracket: P4's candidate min/max
   scan is replaced (behind a runtime guard) by [vlo, vhi) from
   s_thr[1]/[2] — P3 collected v >= vlo and cnt(>=vhi) < kK puts the
   k-th inside the window. Out-of-window values clamp into edge bins
   (cumulative-from-top counts stay exact) and the window is P2's
   acceptance band, so level-1 resolution improves. Stale-bracket paths
   (degenerate fallback, probe variants) fail the guard and take the
   original scan; a plausible-but-wrong bracket costs extra snap steps
   only, never exactness.

2. Converged-path writeback drops its SMEM atomics: the converged snap
   iteration's per-warp packed(ge,gt) counts are still staged in
   smem_wcnt and cover exactly this scan's per-warp element sets, so
   warps derive deterministic [gt | eq] output bases from a register
   prefix instead of ~2*cand/32 serialized atomicAdd claims. Padding
   count becomes cge (= s_iscalars[2]) on this path.

Synth sweep (B300, bs64 K=1024): 4K 10.71->10.25, 8K 12.42->12.14,
16K 15.80->15.66, 32K 21.25->21.15; large N unchanged. Both changes
individually small (~1-2%) — the ncu stall attribution on the min/max
and writeback-atomic regions was largely pipeline-accounting; kept for
the latency-chain shortening and the documented bracket contract.

Made-with-Claude-Code: fable-5
Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
Pre-existing hole surfaced by adversarial review: a candidate span
> FLT_MAX (|v| ~ 1.7e38, fuzz-only for real logits) overflows range1
to +inf -> inv1 = +0 -> all candidates in bin 0 -> thr = NaN. NaN
breaks the snap walk's ordering (all comparisons false), so it no-ops
for snap_limit iterations and the row is emitted as pure padding.
Clamping range1 to FLT_MAX keeps the start threshold ordered; snap's
monotone walk rescues any ordered (even +/-inf) start. No perf change
(8K spot 12.17us, smoke PASS).

Made-with-Claude-Code: fable-5
Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
Adversarial review of the cs>1 path (now the recommended config for
N>=64K) surfaced a pre-existing theoretical data race from the
multi-CTA commit: block_count_ge uses a single s_cluster_partial slot
and a single cluster rendezvous per call, so a straggler CTA's
post-wait DSMEM read of a peer's slot is unfenced against that peer's
NEXT-call overwrite. If it ever fired, per-CTA totals diverge ->
thresholds diverge -> peer under-collection (wrong top-k) or a
cluster-barrier phase mismatch (hang). The schedule window is
astronomically small (peer must complete six CTA barriers + classify +
a full slice re-scan while the straggler's tid0 sits between wait and
its peer loads), but under the PTX memory model it is a race.

Fix with zero added synchronization: call k writes/reads slot k&1
(s_cluster_partial grows [1] -> [3]; slot 2 = tid0-private call
counter, zeroed in run_one_row; CTAs call in lockstep so parity stays
cluster-aligned). The peer's call-(k+2) overwrite of the same slot is
transitively ordered after my call-k reads by the call-(k+1)
rendezvous. cs2 spot 64K/128K/256K + smoke all PASS, perf unchanged.

Made-with-Claude-Code: fable-5
Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
…lag)

Add phase4_rank_scatter as an alternative Phase-4 path behind two inert
flags (enable_p4_rank_scatter / enable_p4_rank_scatter_exact, both default
False). Ported verbatim from p4_recursive_digit/gvr_topk_decode_p4.py
(PR#15709). It replaces phase4_histogram_snap's k-th-bin search +
two-pass writeback with a single fused rank-and-scatter, cutting Phase-4
barriers ~14 -> ~7; the _exact variant adds one fine-histogram recursion
on the straddling coarse bin so the result stays bit-exact vs torch.topk.

Production upstream currently ships ONLY the histogram-snap P4 (no
rank-scatter); this brings the earlier latency win into the production kernel.
All deps (llvm.bitcast, atomicAdd, _fmin_f32_inline, warp_reduce_min/max_f32,
kNumBins) already present. When both flags are False the const_expr switch
selects phase4_histogram_snap, so an OFF kernel is behaviour-identical to
the base. Exactness verified 36/36 (K512/1024/2048, cs1/4, best/worst,
N to 262144) with rank_scatter_exact ON.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…r R0

Resolve enable_p4_rank_scatter / enable_p4_rank_scatter_exact from None to
ON whenever enable_r0 is set (both remain overridable). The R0 fast path
now runs the earlier fused rank-and-scatter P4 instead of histogram-snap.

nsys A/B over the the synthetic envelope (best/worst fp32, BS=1,
seq 4k-1M, cold-L2, per-cell snap-vs-rs paired on one GPU): geomean 1.092x
(K1024 1.120 / K2048 1.115 / K512 1.045), 39/52 cells win >2%, 0/52 regress
>2% (worst -1.6%, noise). Largest wins at cs=1 small-N (up to 1.42x). When
enable_r0 is False the default resolves to OFF, so the base kernel stays
byte-identical to upstream. Exactness 36/36 (all K, cs1/4, best/worst).

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…_size>1

Make the R0 p1b_cache default cs-aware. The single-CTA path keeps the
(dtype != fp32) gate (fp32 occupancy regression at kC=6144), but the
cluster path (cluster_size>1) now defaults p1b_cache ON for all dtypes:
the SMEM gather-cache win holds and the fp32 occupancy regression does
not reproduce in the latency-bound cluster kernel (different SMEM budget).

Matches the reference prototype the reference cluster dispatch (unconditional ON). nsys A/B at cs=4 fp32
(the synthetic envelope best/worst BS=1, N 131k-524k, paired per-GPU): K1024 ~1.01x /
K2048 ~1.02x / K512 wash, 0/18 losses, exactness 18/18. Base kernel
(enable_r0=False) unaffected.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
After per-warp partials land in SMEM (one visibility barrier), EVERY
warp redundantly performs the block reduce + decision a single leader
used to do, in a fixed order — identical inputs and identical fp32
operation order make the result bit-identical across warps, so the
publish barrier, the leader serialization, and the SMEM scalar
re-reads (each costing a cluster-window S2R recompute in SASS) all
disappear.

Phase 4 (p4_warp_redundant, default on):
- k-th bin search: 3 barriers -> 1; the serial 16-slot prefix and the
  64-deep serial bin walk become lane-parallel (idx-shuffle inclusive
  scans + unique-crossing ballot); the selected threshold/count return
  in registers through level-2 refinement.
- snap loop: 2 barriers/iter -> 1; threshold + convergence state in
  registers; warp partial staging parity double-buffered in smem_hist
  (a warp one iteration ahead cannot clobber a slot a slower warp still
  reads); the converged writeback reads the winning parity bank.

Phase 2 (p2_warp_redundant, default on, cluster_size == 1):
- one barrier per secant round (the staging barrier inside
  block_count_ge); classify + secant update replayed redundantly in
  registers by every warp; smem_wcnt staging parity-banked (2x
  num_warps); canonical exit state published once for Phase 3.

Implementation notes from the falsification loop (B300, N=256K bs=64):
- cute.arch.shuffle_sync_up IGNORES its offset (probed: a lane-id
  prefix scan returned 2^k*lane) — scans use shuffle_sync with a
  computed source lane instead. The interim broadcast-LDS prefix
  (64 predicated reads per thread per call) cost +40% instructions
  and erased the barrier win (ncu: barrier stall 14.1%->6.1% but
  16.2M vs 11.5M inst).

ncu (256K, bs64): barrier-stall 14.1% -> 6.1%; regs 64 -> 64, zero
local-memory traffic both sides. A/B wall: 4K +2.6%, 32K +12%,
256K +3.3% (knobs off = leader paths, byte-identical).

Made-with-Claude-Code: fable-5
Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
- _hist_build: clamp the bin index in the float domain (fmax/fmin)
  BEFORE the int cast — fptosi is undefined for out-of-range/NaN at the
  IR level even though PTX cvt.rzi saturates; fmax first canonicalizes
  NaN to bin 0. Edge-bin semantics for in-range values are unchanged.
- run_gvr_topk.gvr_topk_decode: document p4_warp_redundant /
  p2_warp_redundant in the Args section.

Made-with-Claude-Code: fable-5
Co-Authored-By: Claude <noreply@anthropic.com>
Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
Pure formatting: run ruff-format over gvr_topk_decode.py so the R0
histogram-ladder / rank-scatter code added by the preceding commits is
ruff-format-clean and passes the pre-commit gate. No logic change' the
compiled kernel is identical (enable_r0 default stays False here; the
default flip + tests land in the next commit).

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…alence test

Flip the GvrTopKKernel default to enable_r0=True so the R0 histogram-ladder
admission + P4 rank-and-scatter become the production path, and
add a direct-drive R0-vs-secant equivalence unit test.

R0 wins on real DSv4/V3.2 decode-capture workloads: 24/25 cells >= 1.0 vs the
secant baseline, geomean 1.330x (Pro up to 2.10x), value-set-exact 25/25, and it
repairs a base-secant undershoot on one real Flash cell. The secant path is
retained verbatim as the enable_r0=False fallback; on the synthetic cold-hint
worst axis R0 can regress (min 0.917x), to be gated by a follow-up call-site
dispatch guard. enable_r0=False remains byte-identical to the pre-R0 base.

The new test drives GvrTopKKernel directly (the custom op does not expose
enable_r0) and checks top-K by index set (order-independent, per the top-K
contract), with a value-set fallback under half-precision boundary ties. It is
also the only remaining coverage of the secant path now that the op-level tests
inherit the new default.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
@longcheng-nv
longcheng-nv requested review from a team as code owners July 15, 2026 22:23
@longcheng-nv
longcheng-nv removed the request for review from zongfeijing July 15, 2026 22:25
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61207 [ run ] triggered by Bot. Commit: dc2bdc8 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61102 [ run ] completed with state ABORTED. Commit: e612fc2

Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61207 [ run ] completed with state FAILURE. Commit: dc2bdc8
/LLM/main/L0_MergeRequest_PR pipeline #49454 completed with status: 'ABORTED'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

@longcheng-nv

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61215 [ run ] triggered by Bot. Commit: dc2bdc8 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61215 [ run ] completed with state FAILURE. Commit: dc2bdc8
/LLM/main/L0_MergeRequest_PR pipeline #49460 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@longcheng-nv

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61302 [ run ] triggered by Bot. Commit: dc2bdc8 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61302 [ run ] completed with state FAILURE. Commit: dc2bdc8
/LLM/main/L0_MergeRequest_PR pipeline #49534 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@longcheng-nv

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61340 [ run ] triggered by Bot. Commit: dc2bdc8 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61340 [ run ] completed with state FAILURE. Commit: dc2bdc8
/LLM/main/L0_MergeRequest_PR pipeline #49568 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@longcheng-nv

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61367 [ run ] triggered by Bot. Commit: dc2bdc8 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #61367 [ run ] completed with state SUCCESS. Commit: dc2bdc8
/LLM/main/L0_MergeRequest_PR pipeline #49596 completed with status: 'SUCCESS'

CI Report

Link to invocation

@longcheng-nv
longcheng-nv removed request for a team, cascade812 and lori-ren July 24, 2026 00:58

@lfr-0531 lfr-0531 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since I'm not familiar with the kernel side, I didn't dive into the details of the kernel changes. But I checked that there are no interface changes, so overall LGTM.

@@ -353,6 +390,188 @@ def __init__(
self.FLT_MAX = 3.4028235e38
self.NEG_FLT_MAX = -self.FLT_MAX

# --- op#26 R0 histogram-ladder admission (default ON) ---

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These comments are a bit verbose. Could you simplify them? Also include comments in other files. (AI always tends to over-explain things!)

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed — a fair amount of the added commentary is measurement history / tuning provenance (perf numbers, audit counts, campaign refs) that belongs in the PR description rather than the code. I'll prune these down to just the invariants/contracts in the already-planned follow-up PR (the dispatch-guard one), covering all three changed files — keeping this PR's head untouched preserves the green pipeline (#49596) and your approval so it can merge as-is.

longcheng-nv added a commit to longcheng-nv/TensorRT-LLM that referenced this pull request Jul 26, 2026
…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>
return GvrTopKKernel._NUM_SMS

@staticmethod
def pick_config(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

pick_config duplicates the runner's existing policy, and already diverges from it

GvrTopKKernel.pick_config re-implements, heuristic for heuristic, what GvrTopKDecodeRunner._pick_tuning + the cluster_size is None auto-pick block in cute_dsl_custom_ops.py (~L6290 / ~L6626) already do — same thresholds, same tiers, same comments.

Why we duplicate the tune routine?

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch — the duplication is deliberate, but I agree it should collapse to a single source of truth. Context:

Why it exists: GvrTopKKernel.launch is the direct-instantiation path — the unit tests (test_cute_dsl_gvr_topk_decode.py) and standalone/op-bench snapshots import this kernel file without the custom-op layer, so the launch-shape policy was mirrored here to keep the file self-contained.

The divergences are intentional shell behavior, not drift:

  • 32B misalignment: the runner's _pick_tuning asserts (rejects), while launch silently downgrades to 128-bit loads (documented in its docstring) — the standalone path can't assume the caller pre-validated the tensor.
  • cluster cap: the runner warning_once + clamps; launch clamps silently via a guarded import so a standalone snapshot (where _query_max_cluster_size may not be importable) still works.
  • The runner additionally owns the compile-only knobs (enable_unroll_4 / enable_phase3_unroll / use_constant_hint) and the LB-mode branch, which pick_config intentionally doesn't cover.

Plan: fold the de-dup into the already-planned follow-up PR (the same dispatch-guard + comment-pruning one promised above): make the runner's _pick_tuning + cluster_size auto-pick delegate to GvrTopKKernel.pick_config (cute_dsl_custom_ops already imports this module), keeping the assert/warn and LB-mode differences in a thin runner-side shell. Keeping this PR's head untouched preserves the green pipeline (#49596) and the existing approval so it can merge as-is.

@limin2021

Copy link
Copy Markdown
Collaborator

LGTM.

# two-level histogram, and 16-bit tie plateaus are bitwise-equal,
# where arrival order is value-exact).
if p4_exact_tail is None:
p4_exact_tail = self.enable_p4_rank_scatter_exact and dtype == cutlass.Float32

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Two-level linear binning does not guarantee unique fp16/bf16 values. For fp16 range [0, 65504]
with 512 bins, 1.0 and 1.25 land in the same fine bin, so arrival-order scatter can keep the
lower value at the K boundary. Please enable exact-tail handling for 16-bit or add a proven
range guard and adversarial test.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You're right — the docstring overclaims. Binning is window-relative linear (kNumBins x 256 resolution over the Phase-2 bracket), so "16-bit keys are fully resolved" only holds when the bracket is narrow relative to the boundary-local ULP; fp16's log-dense values break it for wide brackets, as in your example. The production DSv4 path is fp32 (where p4_exact_tail handles exactly this), and the 16-bit arrival-fill predates this PR (ships with #15709's rank-scatter-exact path).

Fix will land in the already-planned follow-up PR (dispatch-guard + comment pruning + pick_config de-dup): enable exact-tail for 16-bit (2 radix levels vs fp32's 4, so cheaper), add an adversarial test, and correct this docstring.

s_iscalars[2] = cutlass.Int32(-1)
cute.arch.barrier()
rs = rs + cutlass.Int32(1)
if s_iscalars[1] != cutlass.Int32(1):

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

s_thr[2] is the undershoot side, so its measured count can be below K. Marking done=1 skips
the Phase-3 retry and makes Phase 4 pad -1; a wide boundary plateau can trigger this. Please
route this case to an exact tie-aware fallback.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Correct mechanics. Two things worth separating, though:

The plateau trigger: for fp32 (production path), a > kC - K (~3-4k) run of bitwise-equal values straddling the K boundary cannot realistically arise from continuously-valued logits — never observed across the 865-cell real production grid. It's reachable only for 16-bit dtypes and adversarial inputs.

The -1-pad terminal itself: inherited unchanged from the base kernel's non-convergence behavior (real-data non-convergence is always a small undershoot, via iteration exhaustion, not plateaus). -1 padding is a first-class encoding in this contract, not an error path: every decode step with seq_len < K necessarily pads with -1, so the DSA consumer already skips them on its hot path. Semantically an undershoot drops only the lowest-ranked boundary candidates; end-to-end accuracy has been validated with this behavior in production.

Agreed it should still be exact: in the follow-up PR I'll route the collapsed-bracket case to a tie-aware fallback — emit all strictly-greater candidates, top up the remaining slots from the [s_thr[1], s_thr[2]) band via the existing radix backstop (exact for both plateau and exhaustion cases) — plus plateau tests. This composes with the 16-bit exact-tail item from your other comment.

tj = tj + cutlass.Int32(1)
cute.arch.barrier()
else:
# Persistent scalars live above the 256 digit bins

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The exact-tail radix fallback is duplicated in the fast-path else branch and the non-fast
branch. Please extract one @cute.jit helper so future synchronization and correctness fixes
cannot drift between two copies.

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Agreed — the 4-pass radix select is textually duplicated between the p4_tail_fast else path and the non-fast branch. I'll extract a single @cute.jit helper in the follow-up PR (bundled with the items above), preserving the property that p4_tail_fast=False compiles byte-identical PTX to the original — helpers inline, and I'll re-verify the PTX diff.

@mingyangHao mingyangHao left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This PR can be merged now since production tests are passed. For the comments please address them in followup-PR. @longcheng-nv

@longcheng-nv
longcheng-nv requested a review from zongfeijing July 27, 2026 06:46
@lfr-0531
lfr-0531 enabled auto-merge (squash) July 27, 2026 09:44
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

7 participants