[None][perf] GVR top-K decode: enable R0 histogram-ladder admission by default - #16457
Merged
lfr-0531 merged 50 commits intoJul 28, 2026
Merged
Conversation
…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>
7 tasks
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>
limin2021
reviewed
Jul 27, 2026
Collaborator
|
LGTM. |
limin2021
approved these changes
Jul 27, 2026
mingyangHao
reviewed
Jul 27, 2026
mingyangHao
reviewed
Jul 27, 2026
mingyangHao
approved these changes
Jul 27, 2026
mingyangHao
left a comment
mingyangHao
left a comment
Collaborator
There was a problem hiding this comment.
This PR can be merged now since production tests are passed. For the comments please address them in followup-PR. @longcheng-nv
zongfeijing
approved these changes
Jul 27, 2026
lfr-0531
enabled auto-merge (squash)
July 27, 2026 09:44
pengbowang-nv
approved these changes
Jul 28, 2026
pengbowang-nv
left a comment
pengbowang-nv
left a comment
Collaborator
There was a problem hiding this comment.
Approve for attention as it is an internal kernel change that won't affect current attention interface or computation.
longcheng-nv
added a commit
to longcheng-nv/TensorRT-LLM
that referenced
this pull request
Jul 28, 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>
longcheng-nv
added a commit
to longcheng-nv/TensorRT-LLM
that referenced
this pull request
Jul 29, 2026
…6457 Four of the five follow-up items committed to reviewers on NVIDIA#16457 (the fifth, the dispatch guard itself, is this PR's dispatcher): 1. Comment pruning (requested by @lfr-0531): measurement-history / tuning-provenance commentary reduced to invariants and contracts across the kernel and custom-op files. 2. Launch-shape policy single source (requested by @limin2021): pick_config is split into pick_cluster_size + pick_tuning on the kernel class (the single source of truth); the production runner's _pick_tuning becomes a thin adapter and its cluster auto-pick delegates to pick_cluster_size. The intentional shell divergence is kept and documented: the runner ASSERTS on a 32B-misaligned logits pointer (contract violation) while GvrTopKKernel.launch silently downgrades to 128-bit loads. New test test_..._pick_policy_single_source sweeps dtype x BS x N x graph-capture and pins runner == kernel policy. 3. 16-bit exact-tail (requested by @mingyangHao): p4_exact_tail now defaults ON for fp16/bf16 as well - candidate keys are ALWAYS fp32 (16-bit inputs are upcast injectively at collect), so the tail radix re-rank on the full fp32 order key is exact for every dtype; the overclaiming 'fully resolved' docstring is corrected. New adversarial test: two distinct 16-bit values (1.0 vs 1.25) in one fine bin straddling the K boundary under a wide Phase-2 bracket, fp16 + bf16. 4. P4 exact-tail radix de-duplication (requested by @mingyangHao): the two verbatim copies (tiny-tie fast path's large-class fallback and the plain exact-tail path; token-identical, 1162 tokens) collapse into one @cute.jit helper _p4_exact_tail_radix_select. Verified: the p4_tail_fast=False variant compiles to BYTE-IDENTICAL PTX before and after (465,875 bytes, CUTE_DSL_KEEP=ptx). Remaining item (plateau undershoot terminal routing to an exact tie-aware fallback) follows as its own commit: the audit found the rank-scatter path currently has no cand_count < K branch at all, so the fix is wider than the review comment assumed and deserves isolated review. Gates: full sparse-attention suite 674 passed / 144 skipped (includes the two new tests); PTX identity proof above. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
longcheng-nv
added a commit
to longcheng-nv/TensorRT-LLM
that referenced
this pull request
Jul 29, 2026
Resolves the last NVIDIA#16457 review follow-up: a bitwise-equal plateau wider than the candidate buffer (kC) straddling the K boundary has no admissible threshold - every threshold either overflows the buffer or undershoots K - so Phase 2 previously gave up on the undershoot side and Phase 4 padded the tail with -1, dropping entries that belong in the answer. Fix: on that terminal the driver first collapses the bracket by pure bisection until [lo, hi] are ADJACENT floats. Every value in [lo, hi) is then bitwise-equal, i.e. a genuine tie class, so the row can be completed exactly: Phase 4 emits the cnt(>= hi) sure winners and a ticketed fill takes any (K - count)-subset of the tie class - a valid tie-aware completion by definition. A count that lands in [K, kC] mid-collapse converges normally. Both terminals (the base secant driver and the admission path's tie-plateau fail-soft) take this route; the guard requires a coherent undershoot-overflow bracket with both counts current, so the admission retry's widened brackets - whose adjacency carries no tie-class meaning - are excluded. Non-plateau undershoot keeps the documented -1-pad encoding. Note for future work in Phase 4: the terminal is captured into a dedicated SMEM slot BEFORE Phase 4 runs, because Phase 4 reuses s_iscalars[1] as radix scratch - reading the terminal back from it afterwards yields a mid-radix value. Adversarial test: a plateau wider than kC straddling K, fp32 + fp16 x {rank-scatter cs=1, cs=4, histogram-snap}, 6/6. Full sparse-attention suite unchanged at 674 passed / 144 skipped. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
longcheng-nv
added a commit
to longcheng-nv/TensorRT-LLM
that referenced
this pull request
Aug 5, 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>
longcheng-nv
added a commit
to longcheng-nv/TensorRT-LLM
that referenced
this pull request
Aug 5, 2026
…6457 Four of the five follow-up items committed to reviewers on NVIDIA#16457 (the fifth, the dispatch guard itself, is this PR's dispatcher): 1. Comment pruning (requested by @lfr-0531): measurement-history / tuning-provenance commentary reduced to invariants and contracts across the kernel and custom-op files. 2. Launch-shape policy single source (requested by @limin2021): pick_config is split into pick_cluster_size + pick_tuning on the kernel class (the single source of truth); the production runner's _pick_tuning becomes a thin adapter and its cluster auto-pick delegates to pick_cluster_size. The intentional shell divergence is kept and documented: the runner ASSERTS on a 32B-misaligned logits pointer (contract violation) while GvrTopKKernel.launch silently downgrades to 128-bit loads. New test test_..._pick_policy_single_source sweeps dtype x BS x N x graph-capture and pins runner == kernel policy. 3. 16-bit exact-tail (requested by @mingyangHao): p4_exact_tail now defaults ON for fp16/bf16 as well - candidate keys are ALWAYS fp32 (16-bit inputs are upcast injectively at collect), so the tail radix re-rank on the full fp32 order key is exact for every dtype; the overclaiming 'fully resolved' docstring is corrected. New adversarial test: two distinct 16-bit values (1.0 vs 1.25) in one fine bin straddling the K boundary under a wide Phase-2 bracket, fp16 + bf16. 4. P4 exact-tail radix de-duplication (requested by @mingyangHao): the two verbatim copies (tiny-tie fast path's large-class fallback and the plain exact-tail path; token-identical, 1162 tokens) collapse into one @cute.jit helper _p4_exact_tail_radix_select. Verified: the p4_tail_fast=False variant compiles to BYTE-IDENTICAL PTX before and after (465,875 bytes, CUTE_DSL_KEEP=ptx). Remaining item (plateau undershoot terminal routing to an exact tie-aware fallback) follows as its own commit: the audit found the rank-scatter path currently has no cand_count < K branch at all, so the fix is wider than the review comment assumed and deserves isolated review. Gates: full sparse-attention suite 674 passed / 144 skipped (includes the two new tests); PTX identity proof above. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
longcheng-nv
added a commit
to longcheng-nv/TensorRT-LLM
that referenced
this pull request
Aug 5, 2026
Resolves the last NVIDIA#16457 review follow-up: a bitwise-equal plateau wider than the candidate buffer (kC) straddling the K boundary has no admissible threshold - every threshold either overflows the buffer or undershoots K - so Phase 2 previously gave up on the undershoot side and Phase 4 padded the tail with -1, dropping entries that belong in the answer. Fix: on that terminal the driver first collapses the bracket by pure bisection until [lo, hi] are ADJACENT floats. Every value in [lo, hi) is then bitwise-equal, i.e. a genuine tie class, so the row can be completed exactly: Phase 4 emits the cnt(>= hi) sure winners and a ticketed fill takes any (K - count)-subset of the tie class - a valid tie-aware completion by definition. A count that lands in [K, kC] mid-collapse converges normally. Both terminals (the base secant driver and the admission path's tie-plateau fail-soft) take this route; the guard requires a coherent undershoot-overflow bracket with both counts current, so the admission retry's widened brackets - whose adjacency carries no tie-class meaning - are excluded. Non-plateau undershoot keeps the documented -1-pad encoding. Note for future work in Phase 4: the terminal is captured into a dedicated SMEM slot BEFORE Phase 4 runs, because Phase 4 reuses s_iscalars[1] as radix scratch - reading the terminal back from it afterwards yields a mid-radix value. Adversarial test: a plateau wider than kC straddling K, fp32 + fp16 x {rank-scatter cs=1, cs=4, histogram-snap}, 6/6. Full sparse-attention suite unchanged at 674 passed / 144 skipped. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
dhiraj113
added a commit
to dhiraj113/flashinfer
that referenced
this pull request
Sep 1, 2026
…LM #18094) FlashInfer's GVR top_k_varlen backend was ported from TRT-LLM's V1 kernel at its ~Jul 22 2026 state and missed two upstream hardening layers. The Phase-2 secant threshold search can terminate without a threshold whose count lands in the [K, kC] acceptance window; the old code then shipped a silently wrong top-K: identity indices row[0:K] on a degenerate hint bracket, or an underfilled row whose untouched output slots keep stale / -1 garbage. Reproducible triggers: hostile or degenerate pre_idx hints, tie plateaus wider than the candidate buffer (ReLU-sparse rows), and batches where most rows have N_eff = K + 1 (found by the gvr_2 comparison sweep: short-scenario B>=64, N=8192, K=1024 returned out-of-range indices on N_eff=1025 rows, identically for LB and non-LB). This ports the correctness-relevant subset of three upstream commits, adapted to FlashInfer's diverged copy of the kernel (extra launch knobs, no R0/tiers machinery): - NVIDIA/TensorRT-LLM#16457 (Jul 28) + #16877 (Aug 13): the tie-plateau layer — adjacent-float bracket terminal (done=3, threshold = the sure-winner side) and budget-exhausted plateau-collapse bisection in BOTH phase2_secant_search copies (redundant-warp and leader); s_iscalars grows (6,)->(8,) with [6] = plateau flag captured before Phase 4 and [7] = fill ticket; Branch-C's -1 pad is gated off under the plateau terminal and a post-Phase-4 fill completes the row from the bitwise-equal tie class (any (K - count)-subset is a valid tie-aware completion). - NVIDIA/TensorRT-LLM#18094 (Aug 25, the upstream tip for this kernel): the two-sided repair — Phase 3's retry loop (previously overflow-only, 10 arithmetic-midpoint iters that stall on adjacent floats) becomes an anchored bisection on the signed fp32 order-key image (provable collapse in <= 32 steps, budget 48), handling undershoot by anchoring the untested bracket end at a float extreme, restoring to val_lo when the collapse ends under K, and handing collapsed tie plateaus to the done=3 machinery; the degenerate-hint identity emit is replaced by a synthetic-bracket fall-through (correctness no longer depends on the hint at all; cnt_hi is seeded with top_k so the collapse guard cannot fire on the unmeasured bracket). Both LB paths are covered automatically (GvrTopKLBKernel reuses GvrTopKKernel.run_one_row). Converging rows — the common case — pay only a flag capture and one extra barrier before Phase 4. Tests: tests/topk_varlen/test_gvr_threshold_repair.py ports upstream #18094's regression patterns onto the FlashInfer API (hostile bottom-k / uniform / random hints x K x LB modes, ReLU-sparse plateaus, MTP hostile hints with a mod-cr boundary) plus the FlashInfer-found N_eff=K+1 batch case; 31/31 pass, and the full topk_varlen suite passes 146/146 on B200 (SM100). Off-Blackwell suites unaffected (gvr is sm_100/103-gated). AI-assisted (Claude Code): ported with a 3-lens adversarial review against the upstream reference (port fidelity, DSL/barrier execution, FlashInfer-divergence interactions); all lenses clean. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
dhiraj113
added a commit
to dhiraj113/flashinfer
that referenced
this pull request
Sep 1, 2026
…LM #18094) FlashInfer's GVR top_k_varlen backend was ported from TRT-LLM's V1 kernel at its ~Jul 22 2026 state and missed two upstream hardening layers. The Phase-2 secant threshold search can terminate without a threshold whose count lands in the [K, kC] acceptance window; the old code then shipped a silently wrong top-K: identity indices row[0:K] on a degenerate hint bracket, or an underfilled row whose untouched output slots keep stale / -1 garbage. Reproducible triggers: hostile or degenerate pre_idx hints, tie plateaus wider than the candidate buffer (ReLU-sparse rows), and batches where most rows have N_eff = K + 1 (found by the gvr_2 comparison sweep: short-scenario B>=64, N=8192, K=1024 returned out-of-range indices on N_eff=1025 rows, identically for LB and non-LB). This ports the correctness-relevant subset of three upstream commits, adapted to FlashInfer's diverged copy of the kernel (extra launch knobs, no R0/tiers machinery): - NVIDIA/TensorRT-LLM#16457 (Jul 28) + #16877 (Aug 13): the tie-plateau layer — adjacent-float bracket terminal (done=3, threshold = the sure-winner side) and budget-exhausted plateau-collapse bisection in BOTH phase2_secant_search copies (redundant-warp and leader); s_iscalars grows (6,)->(8,) with [6] = plateau flag captured before Phase 4 and [7] = fill ticket; Branch-C's -1 pad is gated off under the plateau terminal and a post-Phase-4 fill completes the row from the bitwise-equal tie class (any (K - count)-subset is a valid tie-aware completion). - NVIDIA/TensorRT-LLM#18094 (Aug 25, the upstream tip for this kernel): the two-sided repair — Phase 3's retry loop (previously overflow-only, 10 arithmetic-midpoint iters that stall on adjacent floats) becomes an anchored bisection on the signed fp32 order-key image (provable collapse in <= 32 steps, budget 48), handling undershoot by anchoring the untested bracket end at a float extreme, restoring to val_lo when the collapse ends under K, and handing collapsed tie plateaus to the done=3 machinery; the degenerate-hint identity emit is replaced by a synthetic-bracket fall-through (correctness no longer depends on the hint at all; cnt_hi is seeded with top_k so the collapse guard cannot fire on the unmeasured bracket). Both LB paths are covered automatically (GvrTopKLBKernel reuses GvrTopKKernel.run_one_row). Converging rows — the common case — pay only a flag capture and one extra barrier before Phase 4. Tests: tests/topk_varlen/test_gvr_threshold_repair.py ports upstream #18094's regression patterns onto the FlashInfer API (hostile bottom-k / uniform / random hints x K x LB modes, ReLU-sparse plateaus, MTP hostile hints with a mod-cr boundary) plus the FlashInfer-found N_eff=K+1 batch case; 31/31 pass, and the full topk_varlen suite passes 146/146 on B200 (SM100). Off-Blackwell suites unaffected (gvr is sm_100/103-gated). AI-assisted (Claude Code): ported with a 3-lens adversarial review against the upstream reference (port fidelity, DSL/barrier execution, FlashInfer-divergence interactions); all lenses clean. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
dhiraj113
added a commit
to dhiraj113/flashinfer
that referenced
this pull request
Sep 2, 2026
…LM #18094) FlashInfer's GVR top_k_varlen backend was ported from TRT-LLM's V1 kernel at its ~Jul 22 2026 state and missed two upstream hardening layers. The Phase-2 secant threshold search can terminate without a threshold whose count lands in the [K, kC] acceptance window; the old code then shipped a silently wrong top-K: identity indices row[0:K] on a degenerate hint bracket, or an underfilled row whose untouched output slots keep stale / -1 garbage. Reproducible triggers: hostile or degenerate pre_idx hints, tie plateaus wider than the candidate buffer (ReLU-sparse rows), and batches where most rows have N_eff = K + 1 (found by the gvr_2 comparison sweep: short-scenario B>=64, N=8192, K=1024 returned out-of-range indices on N_eff=1025 rows, identically for LB and non-LB). This ports the correctness-relevant subset of three upstream commits, adapted to FlashInfer's diverged copy of the kernel (extra launch knobs, no R0/tiers machinery): - NVIDIA/TensorRT-LLM#16457 (Jul 28) + #16877 (Aug 13): the tie-plateau layer — adjacent-float bracket terminal (done=3, threshold = the sure-winner side) and budget-exhausted plateau-collapse bisection in BOTH phase2_secant_search copies (redundant-warp and leader); s_iscalars grows (6,)->(8,) with [6] = plateau flag captured before Phase 4 and [7] = fill ticket; Branch-C's -1 pad is gated off under the plateau terminal and a post-Phase-4 fill completes the row from the bitwise-equal tie class (any (K - count)-subset is a valid tie-aware completion). - NVIDIA/TensorRT-LLM#18094 (Aug 25, the upstream tip for this kernel): the two-sided repair — Phase 3's retry loop (previously overflow-only, 10 arithmetic-midpoint iters that stall on adjacent floats) becomes an anchored bisection on the signed fp32 order-key image (provable collapse in <= 32 steps, budget 48), handling undershoot by anchoring the untested bracket end at a float extreme, restoring to val_lo when the collapse ends under K, and handing collapsed tie plateaus to the done=3 machinery; the degenerate-hint identity emit is replaced by a synthetic-bracket fall-through (correctness no longer depends on the hint at all; cnt_hi is seeded with top_k so the collapse guard cannot fire on the unmeasured bracket). Both LB paths are covered automatically (GvrTopKLBKernel reuses GvrTopKKernel.run_one_row). Converging rows — the common case — pay only a flag capture and one extra barrier before Phase 4. Tests: tests/topk_varlen/test_gvr_threshold_repair.py ports upstream #18094's regression patterns onto the FlashInfer API (hostile bottom-k / uniform / random hints x K x LB modes, ReLU-sparse plateaus, MTP hostile hints with a mod-cr boundary) plus the FlashInfer-found N_eff=K+1 batch case; 31/31 pass, and the full topk_varlen suite passes 146/146 on B200 (SM100). Off-Blackwell suites unaffected (gvr is sm_100/103-gated). AI-assisted (Claude Code): ported with a 3-lens adversarial review against the upstream reference (port fidelity, DSL/barrier execution, FlashInfer-divergence interactions); all lenses clean. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
dhiraj113
added a commit
to flashinfer-ai/flashinfer
that referenced
this pull request
Sep 4, 2026
…king auto, and the V1 threshold repair (#4811) # feat(topk): self-sampling GVR V2 backend, oracle-tracking `auto`, and the V1 threshold repair Three commits, one arc: port TRT-LLM's self-sampling GVR V2 top-K decode as a new `top_k_varlen` backend, make `backend="auto"` track the measured per-config winner, and fix the correctness bug in the existing V1 `gvr` backend that the new benchmark sweep uncovered. ## 1. `gvr_2` — self-sampling GVR V2 port (NVIDIA/TensorRT-LLM#17821, commit `ed94d4cfbf`) **Algorithm** — V1 GVR guesses its selection threshold from the previous step's `pre_idx` hint and pays full-row refinement rescans when the hint is stale. V2 instead derives a *bracketed ladder* of candidate thresholds from an in-kernel sample of the row itself (Floyd–Rivest-style) and resolves exact, tie-interchangeable top-K in a **single streaming pass**; exactness is guaranteed by count-crossing invariants, never by the estimate, and the hint survives only as a degenerate-case anchor. Four kernel families (streaming `main` with multi-CTA SPLIT, register-resident `reg`/`regimg`, clustered `clus`/`reg_clus`) are chosen by a pure host dispatch `route(b, n, npad, k)`; per-request lengths are read **on device**, so one launch serves the whole ragged batch — no prepare kernel, no LJF sort, no host reads (sync-free and CUDA-graph safe; the length envelope comes from the logits row width). **Port notes** - `flashinfer/topk_varlen/kernels/gvr2_topk_decode.py` (device) and `gvr2_topk_host.py` (dispatch/workspace/entry) are near-verbatim upstream drops, excluded from ruff/mypy like the repo's other verbatim kernel ports so future syncs stay mechanical. Local changes: module rename, provenance notes, a `_persist()` hook routing the four `get_compiled*` builders through the persistent CuTe-DSL kernel cache (the upstream compile closures already use `--enable-tvm-ffi` + fully symbolic shapes, matching the cache's TVM-FFI reload convention), and a fixed missing kv-arg in the `regclus_topk` debug entry. - Backend contract: fp32 logits only (bf16/fp16 are an upstream follow-up), `top_k ∈ {512, 1024, 2048}`, `compress_ratio ∈ {1, 4}`, `pre_idx` required with width == `top_k`, datacenter Blackwell (sm_100/103). The ~21 MB per-device workspace slab is zero-initialized once and self-restoring; multi-stream callers can pass `workspace={"gvr2_workspace": ...}`. - Upstream caveat found by this PR's adversarial tests and reproduced **bit-identically by the TRT-LLM implementation**: literal `+inf` logits may not be selected (at least in the clustered-register family); all finite values — including 3.1e38, above the kernel's 3e38 pad sentinel — and `-inf` are tie-aware exact. Documented in the API docs + a non-strict `xfail` (`test_gvr2_plus_inf_upstream_caveat`) that flips visible if a future upstream sync fixes it. Same class as upstream's implementation-specific NaN ordering. **Performance** (B200, fp32, CUDA-graph timing, tie-aware output validation before every timing; 213 configs across uniform/mixed/short lengths, K ∈ {512, 1024, 2048}, B ∈ [1, 256], N ∈ [1K, 128K], cr ∈ {1, 4}, next_n ∈ {1, 2}): | vs backend | geomean | gvr_2 faster | range | |---|---|---|---| | `gvr` (existing, LB) | **2.58×** | 210/211 | 0.47–6.23× | | `gvr` (non-LB) | 3.80× | 210/211 | 0.76–10.94× | | `radix` (CuTe DSL) | 2.13× | 211/213 | 0.30–10.05× | | `radix_cutlass` | 3.29× | 211/213 | 0.56–5.92× | | TRT-LLM upstream twin (port parity) | **1.00×** | — | 0.95–1.05× | The only losing regime is batches whose rows are mostly barely longer than K (upstream's own documented short-row trade-off). ## 2. Shape/dtype-aware `auto` (oracle-tracking) The previous static gvr-first `auto` never reached `gvr_2`, always picked `gvr` for hinted bf16/fp16 where `radix` wins by up to 2.8×, and never reached `radix_cutlass` in its fp32 big corner (up to 2.8×). The new ranking uses only capture-stable host facts (dtype, N, B — never `seq_lens` contents, which would cost a D2H sync), derived from a 500+-cell sweep: 1. hinted fp32 → **`gvr_2`**; 2. `gvr` outranks `radix` only when `B·N ≥ 2²²` (fp32) / `2²³` (bf16/fp16); 3. `radix_cutlass` outranks `radix` only in the fp32 corner `N ≥ 65536` and `B·N ≥ 2²³`; in half precision `radix` always leads. **Validated on 100 off-grid configs** (shapes never used to fit the thresholds), regret = t(auto's pick) / t(fastest backend): | policy | geomean | max | oracle hits | |---|---|---|---| | new | **1.004×** | 1.16× | 97/100 | | old | 1.317× (hinted 1.694×) | 3.05× | 55/100 | Documented static blind spot: mostly-short rows inside a wide N favor `radix` but are indistinguishable without reading `seq_lens`; such callers should pass `backend="radix"` explicitly. `benchmarks/bench_topk_varlen_gvr2.py` (new comparison benchmark, with an upstream-twin mode) gained `--dtype`; the bf16/fp16 sweeps drove rules 2–3. ## 3. V1 `gvr` threshold-search repair (folded from #4813) This PR's sweep found the existing `gvr` backend shipping silently wrong top-K when its Phase-2 threshold search terminates without converging: identity indices `row[0:K]` on degenerate hint brackets, or underfilled rows (stale/−1 output slots) on hostile hints, tie plateaus wider than the candidate buffer, and `N_eff = K+1` batches. Fixed by porting the correctness-relevant subset of three upstream commits FlashInfer's V1 snapshot (~Jul 22) predates: - NVIDIA/TensorRT-LLM#16457 + #16877: the tie-plateau layer — adjacent-float bracket terminal (`done=3`) in both `phase2_secant_search` copies, `s_iscalars (6,)→(8,)`, plateau flag/ticket, gated Branch-C pad, and a post-Phase-4 fill completing the row from the bitwise-equal tie class; - NVIDIA/TensorRT-LLM#18094 (the upstream tip for this kernel, verified via live GitHub): the two-sided repair — Phase 3's overflow-only retry becomes an anchored bisection on the signed fp32 order-key image (provable collapse in ≤32 steps), handling undershoot, restoring to `val_lo` when a collapse ends under K, and handing collapsed plateaus to the `done=3` machinery; the degenerate-hint identity emit becomes a synthetic-bracket fall-through, so correctness no longer depends on the hint at all. Both LB paths are covered automatically (`GvrTopKLBKernel` reuses `GvrTopKKernel.run_one_row`). Before/after: the fix is ~4% *faster* geomean in the converging common case (the repair shrinks oversized candidate sets) and turns the previously-wrong regimes correct at 1.1–7× cost (their old timings were the price of not computing the answer). Full A/B table and root-cause discussion preserved in the #4813 thread. Ported with a 3-lens adversarial review against the upstream reference (port fidelity, DSL/barrier execution, FlashInfer-divergence interactions) — all clean. ## Testing - `tests/topk_varlen/test_topk_varlen_gvr2.py` (79 tests): tie-aware exactness with poisoned pads, short-row identity+pad contract, degenerate hints, varlen/MTP/compress-ratio grids, zero-kv-slot rows, per-family admission parity, CUDA-graph capture/replay with in-place growing `seq_lens`, warmup-then-capture, non-contiguous arena views, workspace override, adversarial tie/huge-value/denormal patterns, a >10k-point `route_split == route` dispatch fuzz, and the `+inf` caveat xfail. - `tests/topk_varlen/test_gvr_threshold_repair.py` (31 tests): upstream #18094's regression patterns on the FlashInfer API (hostile/degenerate hints × K × LB modes, ReLU-sparse plateaus, MTP hostile hints) plus the FlashInfer-found `N_eff = K+1` case — all fail pre-fix. - Cross-backend value-multiset consistency and the shape-aware heuristic-priority unit test (meta tensors, runs off-GPU) extended in `test_topk_varlen.py`. - Results: see section 4 for the post-rebase sweep on SM80/89/90/100 (B200 and B100)/107 (Rubin)/120. ## 4. Rubin (SM107) enablement and rebase onto #4621 Rebased onto current `main`, which now includes #4621 (Rubin support for `top_k_varlen` plus the `radix_filter` backend). Conflicts were confined to `topk_varlen.py` (docstrings, the backend `Literal`/registration, the `auto` heuristic body, and the `next_n` validation), `pyproject.toml`, and `.pre-commit-config.yaml` (both sides added vendored-kernel exclusions; union kept). Resolutions of note: - The `next_n`/`seq_lens` grouped-row validation lands as #4621's `ValueError` raises (they hold under `python -O`); this branch's equivalent asserts were dropped and its up-front-validation test now expects `ValueError`. - `auto` keeps this branch's shape-aware ranking; `radix_filter` stays explicit-only, with #4621's rationale comment retained. `gvr_2` on Rubin (commit `feat(topk): enable gvr_2 on Rubin (SM107)`): - Admission comes through the shared `_GVR_CCS = [100, 103, 107]` list from #4621; the gvr_2 checker now uses #4621's per-device `_cute_dsl_ready()` probe instead of the process-wide DSL flag, so a DSL predating the device degrades `auto` cleanly. - No kernel changes: the device module uses only family-portable ops (no `tcgen05` / block-scaled MMA), so the same source compiles for `sm_107a`; the compile target follows the current device and the persistent CuTe-DSL cache is arch-namespaced. - The host router's `148` (B200 SM count, mirrored from the upstream CUDA dispatch) is documented as an occupancy heuristic only; on Rubin's 208 SMs the same constants are conservative, never incorrect. Per-arch retuning is a perf follow-up. Verification after the rebase (full `tests/topk_varlen/`, which now includes #4621's `test_radix_filter.py`): | Arch | GPU | `tests/topk_varlen/` at `9f518813` (identical to the swept tree except one unused loop index in a test) | |---|---|---| | SM80 | A100-PCIE-40GB | 54 passed, 203 skipped | | SM89 | L40S | 54 passed, 203 skipped | | SM90 | H100 NVL | 54 passed, 203 skipped | | SM100 | B200 | 252 passed, 3 skipped, 2 xfailed | | SM100 | B100 | 252 passed, 3 skipped, 2 xfailed | | **SM107** | **Rubin GR100** | **253 passed, 2 skipped, 2 xfailed** (7 min 19 s cold, including the first `sm_107a` compiles of every gvr_2 family) | | SM120 | RTX 5080 | 54 passed, 203 skipped | The skips on SM100/107 are all in `test_radix_filter.py` (two need a second visible GPU; one wants an async-TMA default that is active on SM107 but not SM100, hence Rubin's extra pass). The two xfails are the known upstream caveats (radix_filter's padded-merge tie, gvr_2's `+inf`). Every gvr_2 test ran on SM100 and SM107. Basic perf after the rebase (`benchmarks/bench_topk_varlen_gvr2.py`, fp32, K = 1024, B ∈ {1, 16, 64, 256}, N ∈ {8K, 32K, 128K}, uniform and mixed lengths, CUDA-graph timing, tie-aware validation before every timing): | GPU | vs `gvr` | vs `radix` | vs `radix_cutlass` | gvr_2 wins | |---|---|---|---|---| | B100 (SM100) | 2.40× | 2.49× | 3.40× | 24/24 vs each | | Rubin GR100 (SM107) | 2.31× | 2.55× | 3.37× | 24/24 vs each | Consistent with the B200 sweep in section 1, and Rubin lands there with the router's B200 occupancy constants untouched. ## Follow-ups (out of scope) - Consider an `N_eff`-aware escape for gvr_2's short-row regime if it is ever promoted for length-skewed workloads. - bf16/fp16 gvr_2 tracks the upstream roadmap; once ported, the V1 `gvr` backend's last winning corner (hinted bf16/fp16 at large B·N) disappears and it can be deprecated. - Report the `+inf` selection caveat upstream (reproducer in the xfail test). AI-assisted (Claude Code). 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Added the GVR V2 backend for variable-length top-k operations on supported Blackwell hardware. - Added support for variable sequence lengths, returned values, workspace controls, warmup, and CUDA Graph workflows. - Added automatic backend selection based on input shape and data type. - Added a configurable benchmark for comparing top-k implementations. - **Bug Fixes** - Improved handling of ties, plateaus, sparse results, short rows, and challenging threshold-search cases. - Added validation and clearer handling for unsupported configurations. - **Tests** - Expanded correctness, consistency, validation, and CUDA Graph coverage. <!-- end of auto-generated comment: release notes by coderabbit.ai --> ## 5. Review follow-up (`ea069804`) Addresses the open review threads: * **Row-length clamp on arena views.** `gvr_main`'s varlen prologue clamped each row's kv-derived length to the row stride; the other three families clamp to the envelope. It now clamps to the envelope too (carried in the previously dead `n` launch slot), so an oversized `seq_lens` value on a wider-stride view can never classify arena columns in `[width, stride)`. Contiguous inputs are unchanged. New test `test_gvr2_arena_view_oversized_seq_lens` covers all families; it fails on the previous kernel for both `gvr_main` parametrisations. * **Heterogeneous multi-GPU processes.** `run` / `run_ws` / `run_varlen` re-enter under `torch.cuda.device(logits.device)` when the logits are not on the current device, so the launcher cache key, the DSL compile target and the launch agree by construction. Every gvr_2 `cute.compile` also passes an explicit `--gpu-arch` for the current device: the DSL otherwise detects its target once per process from device 0, so a kernel compiled while another device is current was built for device 0's architecture and failed to load (`cudaErrorNoKernelImageForDevice`). New test `test_gvr2_logits_on_non_current_device` (exercised with device 0 = L40S, device 1 = B100; failed before both fixes). * Benchmark: `mixed`/`short` rows drawn from the documented domain, GVR-family backends report `n/a` outside `top_k ∈ {512, 1024, 2048}`, summary no longer requires `fi_gvr2`; `benchmarks/routines/topk_varlen.py` drops `gvr_2` when `max_seq_len % 4 != 0`; raw-string regex patterns in tests. * The `+inf` caveat stays an `xfail` with the upstream tracking issue; see the thread for the rationale. **Public API note.** The "potential breaking change" advisory is the `backend` `Literal` gaining the `"gvr_2"` member. That is additive: every previously valid call, including positional ones, is unchanged. ## 6. Known kernel defect: DKG issue #58 (`+inf` / NaN dropped) and the TRT-LLM #18501 port (`8e76f056`) While validating this port I filed DKG issue #58 (internal GitLab, `dlarch-fastkernels/dynamic-kernel-generator`) against the upstream self-sampling GVR V2 kernel: a row containing a literal `+inf` returns a top-k without it, and NaN is not ranked on top the way `torch.topk` does. The result does not depend on hint quality (oracle and `-1` hints fail identically). The `+inf` case is inside the kernel's documented "finite + inf exact" contract; the NaN case is a torch-parity gap (upstream documents NaN ordering as implementation-specific). The kernel owner pointed at TensorRT-LLM PR #18501 as the latest fix. **Reproduced on this branch and ported here** (register family `GvrTopkRegKernel.kern`, verbatim apart from our module rename): count-crossing enforcement when the hint-derived bracket makes the histogram total fall short of k, a radix-descent escape replacing the 32-step key-space bisection, and the `-inf` fill-lane bound. Before the port, the two upstream regressions were severe on this branch: a cold-start hint buffer (all zeros plus an argmax anchor) left 130,304 of 131,072 output slots unwritten across 256 rows, and an in-window `-inf` in a row's tail column left every row entirely unwritten. After the port both are exact with every slot written, on B100, B200 and Rubin (`test_gvr2_high_anchor_hint_completeness`, `test_gvr2_neginf_tail_completeness`). **`+inf` closed (`48572ef2` + `274fad3e`):** the kernel owner's follow-up TensorRT-LLM PR #18625 is ported in full. Its first commit fixes `GvrTopkRegKernel.kern` (an infinite bracket width fails the collapse guard and the row takes the key-space escape); its second commit, written after the repro posted on DKG issue #58, fixes `GvrRegClusKernel.kern` the same way (bounded guard, and an infinite-width bracket forces the whole-row `degen` fallback instead of the collapsed histogram, which used to turn the `+inf` into NaN in the trash bin and return the true top-K with the `+inf` replaced by the (K+1)-th value). Verified on B200, B100 and Rubin against upstream at `9f920331` and through `top_k_varlen(backend="gvr_2")`: `+inf` inside or outside the bracket sample, K in {512, 1024, 2048}, more `+inf` than K, the 4 x 32K and 4 x 64K `reg_clus` shapes, with oracle, `-1` and random hints; `main` and `clus` were already exact. `test_gvr2_posinf_completeness` (register family) and `test_gvr2_posinf_reg_clus` (clustered-register family, replaces the former strict `xfail`) pin it; the API docstring caveat is gone. NaN ordering stays implementation-specific upstream, although NaN now ranks on top in every shape tested. Perf of the reg_clus hunks (B100, graph replay, 5 alternating passes, deterministic to the nanosecond): 4 x 32K +2.5%, 8 x 32K +2.8% (about 0.15 us each), 16 x 32K and the register family unchanged. ## 7. `auto` heuristic refinement (`993f6c9e`): radix_filter admission and the half-precision gvr rule Measured `auto` against an oracle (fastest explicit backend per cell, CUDA-graph replay, tie-aware validation) over 4 batch sizes x 5 lengths (8K..1M) x uniform/mixed/short lengths x hinted/hint-free, fp32 and bf16, K=1024, on B100/B200 (SM100) and Rubin (SM107). Efficiency = best-of-all time / auto time (1.00 = auto matched the oracle). Findings and the resulting rules (one rule set for every arch; the SM100 crossovers are conservative on SM107, where radix_filter's margins are larger): * Hinted fp32 (the DSA decode path) was already at 0.95-1.00: gvr_2 first stays. The remaining sub-0.95 cells are rows whose valid length is within 1-4 of K, a register-family fallback path (perf follow-up). * Hint-free fp32 sat at 0.83 (SM100) / 0.72 (SM107) and bf16 at 0.86 / 0.82 purely because `radix_filter` was excluded from `auto` while being the fastest hint-free backend in most cells from 32K columns up. It is now admitted where it measured fastest: fp32 for N >= 32K (except the single-row case at N >= 512K) and at every N once B >= 256; half precision for 32K <= N <= 128K with B <= 16, N >= 128K with B >= 64, and N <= 8K with B >= 256. Its checker no longer rejects `pre_idx`: the hint is optional steering for the GVR family, so a hinted caller may use any hint-free backend, explicitly or via `auto`, exactly as radix and radix_cutlass already allowed; the hint is accepted and ignored. * One genuine ranking error: for half precision the old `B*N >= 2^23` rule chose gvr (V1) at B <= 64 for N >= 512K, where it loses 1.5-7x to radix (16 rows x 2M: 1518 us vs 210 us). gvr is now chosen for half precision only when B >= 256 and 32K <= N <= 512K, where it measured fastest. The fp32 gvr rule is unchanged (it only matters when gvr_2 is unsuitable; radix_filter ranks ahead of it there). Re-measured with the refined heuristic (cells >= 0.95 in parentheses; worst remaining cell listed): | dtype, K | hint | GPU | before | after | worst cell after | |---|---|---|---|---|---| | fp32, K=1024 | yes | SM100 (B100) | 0.95 (46/48) | **0.99** (57/60) | short B=64 N=8K | | fp32, K=1024 | yes | SM107 (Rubin) | 0.95 (45/48) | not re-measured (GPU shared with another job at the time; same rule set, radix_filter's margins are larger on SM107 so the SM100 crossovers are conservative there) | | | fp32, K=1024 | no | SM100 (B100) | 0.83 (19/48) | **0.99** (58/60) | short B=256 N=128K | | bf16, K=1024 | yes | SM100 (B100) | 0.92 (29/48) | **0.97** (53/60) | short B=64 N=1024K | | bf16, K=1024 | no | SM100 (B100) | 0.86 (26/48) | **1.00** (59/60) | mixed B=16 N=128K | `test_backend_heuristic_priority` pins the new boundaries off-GPU; the public docstring's `backend` section now documents `radix_filter` and the refined `auto` order. --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
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
8 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Note
#16424 (the semantic dependency) has merged into
main(9095cc11d1); this branch is synced with currentmain, so the stacking is resolved — the PR's net diff vsmainis 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_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.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_configcontract (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.2preIdx= 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×
V4 Pro (K1024) — geomean 1.299×
V3.2 (K2048) — geomean 1.443×
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)
pick_config(the (dtype, BS, N) → launch-shape classmethod, incl. the CUDA-graphmax_seq_lencontract) andlaunch(compiled-variant cache;**kernel_overridesfor 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
enable_r0=Falseis exactly what it dispatches to.Correctness
torch.topkon 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).p4_exact_tailmakes 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×).pick_configpolicy 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
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=Falserestores the pre-R0 kernel byte-for-byte (all R0 fields const-folded) — rollback = flip one default.r0_vseed=Falsekeeps the static ladder;p4_exact_taildefaults off for 16-bit (byte-identical there).Test plan
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_sortedaxis moved from the full cross-product to a 4-cell representative test, andtest_lb_vs_referencereduced to the two SMEM-layout endpoint dtypes. All protected coverage (shortrow degrade,p4_exact_tailties, r0-equivalence incl. cs=8/big-BS,pick_config/launchpolicy locks) is unchanged; remaining runtime is dominated by first-instance cuteDSL compilation of the 133 kernel-variant signatures.