[None][perf] GVR top-K decode: enable R0 histogram-ladder admission by default#16457
[None][perf] GVR top-K decode: enable R0 histogram-ladder admission by default#16457longcheng-nv wants to merge 20 commits into
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>
…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>
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>
📝 WalkthroughWalkthroughChangesGVR Top-K execution paths
Estimated code review effort: 5 (Critical) | ~90+ minutes Sequence Diagram(s)sequenceDiagram
participant GvrTopKKernel
participant phase1b_hspace_rungs
participant block_count_ge_multi
participant phase4_rank_scatter
GvrTopKKernel->>phase1b_hspace_rungs: compute rung thresholds
phase1b_hspace_rungs->>block_count_ge_multi: pass thresholds for GE counting
block_count_ge_multi-->>GvrTopKKernel: return admitted rung counts
GvrTopKKernel->>phase4_rank_scatter: emit selected top-K entries
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py (1)
727-733: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd direct R0/secant equivalence coverage for
compress_ratio=4.The only baseline comparison hardcodes
compress_ratio=1, although R0 is enabled for the production DSv4 path andGvrParamsvaries by compression ratio. Parameterize the helper and add a focused DSv4 case with offset-zeropre_idxencoding.Coverage in
tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.pyis insufficient for this production path. As per path instructions, test feedback must identify the concrete test file and whether coverage is sufficient.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py` around lines 727 - 733, The helper _run_gvr_direct currently hardcodes compress_ratio=1, so it cannot compare R0/secant behavior for the production compression path. Parameterize _run_gvr_direct with compress_ratio, pass it through to GvrTopKKernel, and add a focused DSv4 test using compress_ratio=4 with offset-zero pre_idx encoding that compares direct R0 results against the secant path. Ensure coverage is added in tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py`:
- Around line 2008-2139: The final fine-histogram path in the exact rank-scatter
logic still groups distinct fp32 values into one sub-bin and can scatter them in
the wrong order at the K boundary. In the fine-bin handling around sb_star,
refine the selected sub-bin to the actual value boundary using the existing
exact snap logic before writing outputs, while preserving rank_above_fine
ordering. In
tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py lines
803-819, add deterministic fp32 values clustered within one fine-bin width
around K along with broad minimum and maximum outliers to verify exact ordering.
In `@tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py`:
- Around line 759-774: Update _make_r0_pre_idx so encoding preserves pre - 1
without clamping, allowing true index 0 to remain stored as -1 and keeping index
1 distinct. Ensure the test data remains valid for the kernel’s +1 decoding
convention and add or adjust coverage in test_cute_dsl_gvr_topk_decode.py to
exercise the index-zero boundary.
---
Nitpick comments:
In `@tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py`:
- Around line 727-733: The helper _run_gvr_direct currently hardcodes
compress_ratio=1, so it cannot compare R0/secant behavior for the production
compression path. Parameterize _run_gvr_direct with compress_ratio, pass it
through to GvrTopKKernel, and add a focused DSv4 test using compress_ratio=4
with offset-zero pre_idx encoding that compares direct R0 results against the
secant path. Ensure coverage is added in
tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: da5958ef-d00d-44a3-960e-9632c2d86656
📒 Files selected for processing (2)
tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.pytests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
39ace82 to
95a8388
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #59547 [ run ] triggered by Bot. Commit: |
…ch helper
cluster_size / num_threads / min_blocks_per_mp are compile-time ctor knobs,
so batch-size adaptation must happen at launch time by picking a different
compiled variant. Add GvrTopKKernel.pick_config — the (dtype, BS, N) ->
ctor-kwargs policy as a pure classmethod colocated with the kernel (mirrors
the production runner's cluster auto-pick + _pick_tuning, including the
CUDA-graph max_seq_len contract) — and GvrTopKKernel.launch, a thin
compiled-variant-cache wrapper so direct-drive users (tests, benchmarks)
get the same shapes production picks. No device code changes; the custom
op keeps its own inline policy until the dispatch-guard follow-up PR
unifies it onto pick_config.
Rationale (B200 nsys cold-L2, BS in {64, 256, 1024}): a config frozen at
the BS=1 optimum (cs = N>=64K ? 4 : 1, T=1024, mbpm=1) is geomean 2.27x
slower (max 6.0x) than these picks; through pick_config the R0 kernel is
0.95x of the op-bench reference (parity/better), 20/20 cells exact. The
cs=8 rung (BS<=4, N>=128K) beats a forced cs=4 on every probed cell
(geomean 0.943, up to -12% at N=262144), 78/78 standalone exactness cells
green (R0 + secant, 3 dtypes, hr in {1, mid, 0}).
Tests: cluster_size=8 added to the R0 equivalence grid (large-N gated);
new multi-wave big-BS equivalence cells (BS=256 single-CTA, BS=64 cs=4);
pick_config policy lock test; launch() autoconfig test across all four
cluster regimes incl. a forced-secant override arm.
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
|
PR_Github #59547 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59635 [ run ] triggered by Bot. Commit: |
…ung defaults
Fold P1's pmean (the secant init probe) into the M-ary R0 count pass as one
extra virtual rung column. The column reuses the secant per-thread count
buffer (zero SMEM growth, no extra memory traffic or synchronization; P1
parks the value under its own trailing barrier). Admission picks the
tightest admissible column (explicit argmin over the window), and on a full
miss the measured pmean count donates an interior bracket point to the
fallback refine.
This adapts the admission ladder to the row's value distribution and fixes
the fat-admission regime disclosed as the known limitation: on the real
Flash 1M-ISL capture at BS>=128 the coarse q.85 rung admits ~4400
candidates where pmean admits ~630 (7x P3/P4 candidate work), giving
R0/base 0.68-0.79. With the seed rung those cells return to base parity
(fp32 1.02-1.03x, up to 1.43x vs the previous R0 head).
Per-K rung defaults: with the seed rung on, pmean covers q.35's admission
region for K512/K1024, so those drop to qfracs=(0.85,) — two count columns,
zero column tax vs the previous head. K2048 keeps (0.85, 0.35): at kC/K =
2.5 a fat admit is costlier than a slim two-pass miss. Explicit
r0_vseed=False retains the previous (0.85, 0.35) ladder for all K;
enable_r0=False stays byte-identical to the pre-R0 base.
Full-envelope audit (B200 nsys cold-L2, 2772 cells: synthetic best/worst x
K{512,1024,2048} x 3 dtypes x N 4k-1M x BS 1-1024, plus real V4
Flash/Pro/V3.2 decode captures over all ISL rungs): geomean 0.9963 vs the
previous R0 head with the worst base-relative cell improving 0.678 -> 0.791,
value-set exactness unchanged. Residual disclosed: K1024 16-bit BS=128
N>=262k cells at 0.80-0.89 (pre-existing 16-bit mechanism, not admission).
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
…ng fine bin The P4 rank-scatter fine recursion resolves candidate values to range/(kNumBins*256). Distinct fp32 values spaced below that resolution which straddle the top-K boundary land in ONE fine sub-bin, and the scatter kept an arbitrary arrival-order subset of them — observed on real DSv4-Pro 512k-ISL decode captures as |miss|=1 with |dv| ~ 3e-6 (the kept and dropped values differ by less than the fine-bin width). The defect is config- and data-dependent: random logits are tie-free at that scale, so it does not reproduce on synthetic sweeps. Fix (p4_exact_tail, default ON for fp32 rank-scatter-exact kernels): gate on the only ambiguous case — the straddling-bin tie set overfilling the remaining output slots — and re-rank that tie set exactly with an MSB-first 8-bit-digit radix select over order-preserving integer keys (4 levels = bit-exact for fp32), then rewrite the tie slot range. Unambiguous rows (all rows on every natural dataset measured) pay two scalar compares; the counters, fine histogram, and per-warp scan scratch are all reused, so SMEM does not grow, and the digit scan mirrors the existing two-stage descending bin search. 16-bit inputs keep the arrival fill: their upconverted keys are fully resolved by the two-level histogram already, and 16-bit tie plateaus are bitwise-equal, where arrival order is value-exact (forcing the repair on via p4_exact_tail=True instead triggers it on every plateau row at real cost). Tie sets wider than the kC candidate budget remain outside the kernel contract (P3 collection caps them before P4 runs). Validation (B200): the 12 real Pro/512k fp32 audit failures go value-exact at BS 1/128/1024; adversarial 5e-8-spaced and 1-ulp bitwise tie bands repair across K 512/1024/2048, cr 1/4, cs 1/4/8; paired cold-L2 A/B on unaffected cells is 0.998 geomean (noise), 16-bit arms byte-identical; repair-active rows (previously WRONG) pay ~1.3x at BS=1. New unit test plants both adversarial bands and asserts value-multiset exactness against torch.topk. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
|
/bot run |
|
PR_Github #59736 [ run ] triggered by Bot. Commit: |
|
PR_Github #59635 [ run ] completed with state |
|
PR_Github #59736 [ run ] completed with state
|
Formatting-only: rewrap three multi-line conditions/expressions and normalize inline-comment spacing per ruff-format v0.9.4. No semantic change. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
|
/bot run |
|
PR_Github #59901 [ run ] triggered by Bot. Commit: |
The kernel reads logits[pre_idx + 1] (cr=1, next_n=1) behind an idx >= 0 && idx < N guard, so stored -1 is the valid encoding of true index 0. The clamp(min=0) collapsed true indices 0 and 1 onto the same stored value and left the index-zero boundary unexercised. Addresses CodeRabbit review feedback. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>
|
/bot run |
|
PR_Github #59906 [ run ] triggered by Bot. Commit: |
|
PR_Github #59901 [ run ] completed with state |
|
PR_Github #59906 [ run ] completed with state
|
Summary
Adds an R0 histogram-ladder admission fast path plus 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).R0 replaces the Phase-2 secant threshold search with a single-pass multi-threshold "rung ladder" admission, seeded by a 256-bin histogram over the prev-topK gathered values (P1b), and a fused rank-and-scatter Phase-4 writeback (barriers ~14 → ~7). Two hardening pieces are included:
r0_vseed, default on with R0). P1's mean probe (the secant init) is folded into the R0 count pass as one extra virtual rung column, reusing the existing per-thread count buffer (zero SMEM growth, no extra traffic or synchronization). Admission picks the tightest admissible column, and on a full miss the measured count donates an interior bracket point to the fallback refine. This adapts the ladder to each row's value distribution and fixes the cold-hint fat-admission regime (details under Known limitation). With the seed rung, K512/K1024 drop to a single static rung (qfracs=(0.85,), zero column tax); K2048 keeps(0.85, 0.35).p4_exact_tail, default on for fp32). The P4 fine recursion resolves values to the fine-bin width; distinct fp32 values spaced below that resolution which straddle the top-K boundary landed in one fine sub-bin and were kept in arbitrary arrival order — observed on real DSv4-Pro 512k-ISL captures as a 1-element miss with |dv| ≈ 3e-6. Now the only ambiguous case (straddling-bin tie set overfilling the remaining slots) is gated and re-ranked bit-exactly with an MSB-first 8-bit-digit radix select over order-preserving integer keys; all scratch is reused (zero SMEM growth). Unambiguous rows pay two scalar compares. 16-bit kernels are byte-identical (their upconverted keys are already fully resolved; 16-bit tie plateaus are bitwise-equal, so arrival order is value-exact).The classic secant path is retained verbatim and stays reachable via
enable_r0=False.This PR lands the kernel capability + flips the default only. It does not touch the call site / dispatch (the custom op does not plumb
enable_r0); the activation policy / dispatch guard is a follow-up PR (see Known limitation).Performance
Methodology. nsys pure-kernel, cold-L2 (512 MB L2 evict outside the timed NVTX range, per timed call), single-GPU same-run A/B (R0 vs the retained secant baseline on the identical inputs), 20 cold reps. B200 (SM100). All numbers below re-measured 2026-07-16 at the current head (seed rung + exact-tail fix included) with both arms driven through the kernel's own launch-shape contract (
GvrTopKKernel.launch→pick_config): cluster_size 8/4/2/1 by (BS, N) — incl. the cs=8 tiny-grid pick at N≥128K — T=512/1024, mbpm tiers, fp32-only 256-bit loads; i.e. the shapes the production runner picks. The two data sources are (a) a calibrated synthetic envelope (empirical-CDF + temporally-coherent preIdx) and (b) per-layer indexer top-K captured from production DeepSeek-V4 Flash / Pro and V3.2 decode runs.Data sources
Synthetic (calibrated synthetic envelope). Generated by the indexer-topk temporal-synth method rather than i.i.d. noise: each row's logit marginal is a per-layer empirical inverse-CDF with a GPD-calibrated tail fit to the real 64K production captures (this reproduces the top-K selection-boundary tail that a single Beta/Gaussian collapses); the previous-step
preIdxhint is built by a rank-conditional retention model with per-row hit-rate drawn from the real per-decode-step distribution + AR(1) temporal coherence, plus a marginal-preserving positional model (winner-position density + clustering) so the gather footprint (cache-line segments) matches production.best= warm-hint / higher-hit regime (hit-rate ≈ 0.55);worst= cold-hint tail (hit-rate ≈ 0.05).Real (production decode capture). Per-layer
(indexer logits, indexer top-K)dumped from real BS=1 greedy end-to-end decode runs of DeepSeek-V4 Flash (K=512), DeepSeek-V4 Pro (K=1024) and DeepSeek-V3.2 (K=2048), across 9 input-sequence-length rungs (ISL 4k–1024k) × 3 GVR-active layers each.N= post-compress indexer length (V4 compress_ratio=4, V3.2 cr=1);hit= |preIdx ∩ topK| / K (0.27–1.00 measured across the sweep). For V3.2 (no warm-start)preIdxis reconstructed from the previous decode step's top-K. Correctness is validated against each capture's own recorded top-K reference.Synthetic (calibrated synthetic envelope, BS=1, fp32) — µs and speedup per seq-len N
best= warm hint (hit-rate ≈ 0.55). Launch shapes perpick_config(BS=1: cs=1 below 64K, cs=4 at 64K, cs=8 from 128K). All cells value-set-exact vstorch.topk.K512 (V4 Flash)
K512 (V4 Flash) geomean speedup (N=4K–256K): best 1.276×
K1024 (V4 Pro)
K1024 (V4 Pro) geomean speedup (N=4K–256K): best 1.454×
K2048 (V3.2)
K2048 (V3.2) geomean speedup: best 1.073×
Synthetic overall (22 cells shown; best axis, V4 tables show the deployment-envelope range N=4K–256K): geomean best 1.249×; exactness 52/52 over the full measured envelope (incl. the worst-axis and V4 N≥512K cells not tabulated).
Real decode-capture (production DSv4 / V3.2 indexer top-K, BS=1, fp32) — µs and speedup per seq-len
Per-layer top-K captured from real decode runs;
N= post-compress indexer length,hit= preIdx∩topK overlap. Launch shapes perpick_config(cs=8 on the large-N rungs). Time = median over 3 captured layers.V4 Flash (K512)
V4 Flash (K512) geomean speedup: 1.276×
V4 Pro (K1024)
V4 Pro (K1024) geomean speedup: 1.271×
V3.2 (K2048)
V3.2 (K2048) geomean speedup: 1.348×
Real overall (25 cells): geomean 1.294×, 24/25 ≥ 1.0, exactness 25/25.
The single real cell < 1.0 (Pro 512k, 0.959×) is a repair-active row: the exact boundary-tie repair triggers there and pays its re-rank cost — the pre-fix kernel returned a wrong top-K set on exactly these rows (see Correctness), so the previous, faster number was for an incorrect result.
Batch-size scaling + launch-shape policy (
pick_config/launch)The BS=1 tables above raise the natural question of batch scaling. Two findings (same nsys cold-L2 protocol, BS ∈ {1…1024}, synthetic + real inputs replicated to BS independent rows):
cluster_size/num_threads/min_blocks_per_mpare compile-time constructor knobs: a compiled kernel cannot change its own grid or cluster shape, so batch adaptation must pick a different compiled variant per call. Driving the kernel with a config frozen at the BS=1 optimum (cluster_size = N≥64K ? 4 : 1, T=1024, mbpm=1) is geomean 2.38× slower (max 5.8×) than per-(BS, N) picks across BS ∈ {64, 256, 1024} — e.g. K512 bf16 N=65536 BS=1024: 419 µs frozen vs 72 µs picked. Multi-CTA splitting only pays while the grid is a single wave (num_rows × cluster_size ≤ num_SMs); past that, row parallelism already saturates the SMs and per-row splitting is pure overhead. The production custom op already makes these picks; the risk was only ever to direct-drive users.To make that policy part of the kernel's contract (and un-mismeasurable), this PR adds two host-side helpers to
GvrTopKKernel— no device-code change:pick_config(dtype, num_rows, num_candidates, max_seq_len=None)— the (dtype, BS, N) → ctor-kwargs launch-shape policy as a pure classmethod colocated with the kernel: cluster_size 8/4/2/1 by (BS, N) (8 = tiny-grid large-N: BS ≤ 4, N ≥ 128K), threads/occupancy/vec-width tiers, and the CUDA-graph contract (max_seq_len= peak runtime N so variants are picked for the replay shape, not the capture shape). Mirrors the production runner's inline policy; the follow-up dispatch-guard PR unifies the runner onto it.launch(logits, pre_idx, seq_lens, out_indices, top_k, ...)— a thin compiled-variant cache (same sym-int + tvm-ffi compile contract as the runner) so tests and benchmarks get production-equivalent shapes by default;**kernel_overrides(e.g.enable_r0=False,cluster_size=8) force any knob and participate in the cache key.cluster_size = 8 validated (previously untested anywhere): 78/78 standalone exactness cells green (R0 and secant, dtype ∈ {fp32, bf16, fp16} × K ∈ {512, 1024, 2048} × N ∈ {131072, 262144} × warm/mid/cold hint × BS ∈ {1, 4}, fp32 R0 = secant identical index sets), and the cs=8 pick beats a forced cs=4 on 8/8 nsys cells (geomean 0.943×, up to −12 % at N=262144).
Known limitation (disclosed) + follow-up
R0 is not a uniform win on the cold-hint tail, but the seed rung tightened the envelope materially:
The real decode captures do not exhibit the cold-hint distribution (measured hit-rates 0.27–1.00; 24/25 real cells ≥ 1.0), which is why the default is on. A follow-up PR adds a dispatch guard routing the residual regimes to the retained secant path. Because the hint hit-rate is not observable at inference time (it is the overlap of the current top-K with the hint), the guard is designed as an in-kernel admission escape (R0 already counts seed admissions — bail out to the secant path when the count signals a cold seed) and/or trailing-step feedback (the kernel emits its measured admission counter; the host uses the previous step's value as this step's predictor). The
enable_r0=Falsefallback in this PR is exactly what that guard dispatches to.Correctness of the results
Every timing cell above is value-set-exact vs
torch.topkand the exactness is asserted on the same run that produces the time:torch.topk's (rtol/atol 1e-5), plus in-range + no-duplicate +n_below=0 guards. This is robust to value ties (a different index with an equal value is still correct). Synthetic 52/52, real 25/25 (real vs the captured reference); over the full dtype × BS audit grid, R0 is exact on all 2772 cells.p4_exact_tailradix re-rank makes all 12 affected audit cells bit-exact (validated at BS 1/128/1024), with paired cold-L2 A/B on unaffected cells at 0.998 geomean (noise) and 16-bit arms byte-identical.test_cute_dsl_gvr_topk_decode_r0_equivalence(SM100-gated) drivesGvrTopKKerneldirectly (the custom op does not exposeenable_r0), runs bothenable_r0=Trueandenable_r0=False, and asserts (1) both arms are a valid top-K vstorch.topk(tie-aware value-multiset) and (2) identical index sets on tie-free fp32. Grid: dtype ∈ {bf16, fp16, fp32} × K ∈ {512, 1024, 2048} × N ∈ {8192, 65536} × BS ∈ {1, 16} × warm/cold hint × cluster_size ∈ {1, 4, 8} (cs=8 gated to N ≥ 65536, its production regime). Additional tests:..._r0_equivalence_bigbs(multi-wave grids: BS=256 single-CTA, BS=64 cluster_size=4),..._pick_config_policy(locks the (BS, N) → shape map and the big-BS occupancy knobs),..._launch_autoconfig(driveslaunch()across all four cluster regimes incl. a forced-secant override arm), and..._p4_exact_tail_ties(plants adversarial 5e-8-spaced and 1-ulp bitwise tie bands across the top-K boundary and asserts value-multiset exactness across K × N × cluster_size). This is also the sole coverage of the secant fallback now that op-level tests inherit the new default. A standalone exactness harness independently confirms 186/186 across dtype/K/N/BS + adversarial quantized tie plateaus + cluster.Scope / risk / rollback
tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.pyand its unit-test file. No call-site, dispatch, or config-schema changes.enable_r0=Falserestores the pre-R0 upstream kernel byte-for-byte (all R0 fields areconst_expr-folded). Full rollback = flip one default back.r0_vseed=Falseretains the previous static(0.85, 0.35)ladder for all K;p4_exact_taildefaults off for 16-bit (byte-identical there).Test plan
on SM100 (B200 / B300).