[TRTLLM-35882][feat] Add cute dsl gvr top-k decode kernel#14602
Merged
Conversation
Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
…invariant - Replace from_dlpack/static-shape compile with make_fake_compact_tensor + sym_int for batch/num_tokens dims, keeping (dtype, top_k, next_n) as the cache key. Reduces unique compile entries from 810 to 27 across the bench sweep; correctness verified (no OOB writes from cache reuse with wrong shape) via 288-config pytest + cross-impl A/B match. - Fix test_gvr_topk_decode: (1) pre_idx_count now uses top_k (matches CUDA dispatch precondition preIdxCount == topK at heuristic_topk.cuh:810); (2) tie-aware reference now masks logits to per-row effective_len = seq_len - next_n + 1, avoiding false negatives when next_n > 1 makes the kernel skip the last next_n-1 columns. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
…imit / mask guard) Four small mechanical alignments — each isolated, removes only redundant work the CUDA reference does not do. Correctness verified: 288/288 in test_gvr_topk_decode.py. Perf delta within measurement noise (~0.15us estimated, 21us baseline DSL — under the ~0.5us spread floor) but the changes match heuristic_topk.cuh semantics 1:1 and pave the way for later batches. - block_count_ge: drop the trailing barrier (gvr_topk_decode.py:422 -> removed). CUDA blockCountGE (heuristic_topk.cuh:441) returns without a sync because callers already insert their own __syncthreads after their tid==0 post-processing. The previous DSL trailing barrier was redundant (tid==0 reads its own write in-thread, no sync needed). - Phase 4 snap_limit: change from cand_count>128 ? cand_count/4 : 32 to cand_count (matches heuristic_topk.cuh:985). The older bound silently accepted a non-converged threshold in ~0.09 % of adversarial distributions; correctness improvement only, common case still converges in 1-3 iters. - Phase 4 block_min/max: every thread now recomputes block_min/max from the warp-staged smem slots into local registers (matches heuristic_topk.cuh:891-898). Replaces the prior `tid==0 writes s_thr[1]/s_thr[2] then broadcast via __syncthreads` pattern, saving one block barrier in Phase 4. - Phase 4 Pass 1/Pass 2 writeback: wrap popc + atomicAdd + shuffle in `if mask != 0` warp-uniform guard (mirrors heuristic_topk.cuh:1020, 1045). Skips the atomic round-trip when no lane in the warp emits, most impactful for Pass 2 where only K-th-rank ties emit. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
…le + early break) Convert the Phase 2 secant refinement loop and the Phase 3 retry-shrink loop from Python-unrolled `for in range(N)` (every body wrapped in an `if not done:` guard) to runtime `while` with the convergence condition in the loop predicate. This matches CUDA's pattern at heuristic_topk.cuh: 683 (Phase 2) and :769 (Phase 3 retry). Previously, after the kernel converged at iteration k, the remaining N-k unrolled bodies still each issued an LDS+ICMP+branch guard. With secant typically converging at iter 3 of 15 and retry-shrink usually 0 of 10, this saved ~12 + ~10 = ~22 wasted guard sites per kernel call. Tradeoff: lose Python-time const-fold of `if it == 0: f = min(f, 0.5)`, which now becomes a runtime compare. CUDA does the same runtime compare (heuristic_topk.cuh:698-699), so this is alignment not regression. Measured impact (median config bf16 K=1024 N=32768 BS=1 next_n=2, same-process A/B vs CUDA GVR, 5 reps alternating order): DSL_us 21.01 -> 20.28 (-0.73 us, -3.5%) C/G 0.869 -> 0.903 (+3.4 percentage points) Above the ~1.5% bench_kineto spread floor. 288/288 tests pass. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
Two SASS-alignment changes verified against the CUDA reference at the
median config (bf16 K=1024 N=32768 BS=1 next_n=2):
1. cute.make_ptr(..., cute.AddressSpace.gmem, ...) at the two 128-bit
vec-load sites in block_count_ge and phase3_collect_candidates.
Default AddressSpace.generic lowered to SASS LD.E.128; explicit gmem
hint flips to LDG.E.128 (matches CUDA __ldg path, minus .CONSTANT
which still requires CopyG2ROp+invariant).
2. phase1_preidx_stats: replace the runtime `while i < pre_idx_count`
strided loop with `range_constexpr(pre_idx_count // num_threads)`.
pre_idx.shape[1] is a compile-time constant (top_k baked into JIT
cache key); supported top_k in {512, 1024, 2048} are all multiples
of num_threads (512), so n_iters ∈ {1, 2, 4} unrolls cleanly. cute
emits straight-line code (no BRA / ISETP / counter update) and
issues both preIdx LDG.E and input LDG.E.U16 back-to-back, enabling
LSU pipelining (in flight ILP). Mirrors what nvcc/ptxas does for
the equivalent CUDA loop via auto-partial-unroll.
Bench (same-process A/B, 5 repeats × 100 iters, kineto + L2 flush):
Before: C/G = 0.903 (DSL 10.7% slow) -- post-Batch 2 baseline
After: C/G = 0.922 (DSL 8.5% slow)
Δ = +1.9pp
Resource use after changes:
regs/thread: 34 -> 39 (still 3 blocks/SM, occupancy unchanged 75%)
dynamic smem: unchanged (~44 KB)
total SASS instructions: 2935 -> 2944 (codegen ripple, mostly
FMNMX3 +6; loop overhead ISETP/BRA -6/-2/-3 offset by +25 IMAD)
Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
…tail) Replaces the runtime `while i + (vec_w - 1) < N` vec loop in block_count_ge with a 4-way unrolled fast path + 1-way tail. The fast path issues 4 independent LDG.E.128 per round (separate fragments so cute schedules them concurrently), mirroring what nvcc/ptxas does for the equivalent CUDA loop via auto-partial-unroll. SASS verification at median config (bf16 K=1024 N=32768 BS=1 nn=2): - 4 LDG.E.128 per inline at addresses base / base+0x2000 / +0x4000 / +0x6000 — exact match to CUDA's LDG.E.128.CONSTANT pattern (minus the CONSTANT cache hint, which still requires CopyG2ROp+invariant). - Total LDG.E.128 count: 5 -> 21 (4 inlines * 4 + 4 tails + 1 phase3). - Cute software-pipelines: 3 LDGs issued back-to-back, then consume of iter 0 starts while iter 3's LDG is issued in parallel. All 4 are in flight before HBM responds (latency ~600 cy >> 23 inst slots). Resource impact: - Regs/thread: 39 -> 39 (cute reuses fragment regs across loop body; Phase 4 likely remains the kernel-wide peak) - Dynamic smem: unchanged (~44 KB) - Static SASS size: 2944 -> 3672 inst (+25%) -- code bloat acceptable, well within icache; Block Limit Reg = 3 unchanged at occupancy=75%. Bench results (kineto, L2 flush, n_iters=30): Median config (bf16 K=1024 N=32768 BS=1 nn=2), same-process A/B: Before this commit: C/G = 0.922 (DSL 8.5% slow) After this commit: C/G = 0.976 (DSL 2.4% slow) Delta: +5.4pp Full sweep (804 configs = 3 dtype * 3 top_k * 6 N * 5 BS * 3 next_n): Median C/G: 0.860 (baseline post-Batch-2) -> 0.988 (now) Geomean C/G: 0.869 -> 0.999 (parity with CUDA) DSL faster: 17% -> 46% Within 5%: 13% -> 34% Within 10%: 28% -> 55% By dtype: bf16 1.000, fp16 1.022, fp32 0.951 (fp32 has slightly less runway since vec_w=4 vs 8 for bf16/fp16). By N: gap remains at large N (>=64K: median ~0.87-0.90), where the LSU-pipelining win is already saturated and other phases dominate. The single-config worst slowdowns observed (C/G ~0.4) are concentrated in nn=3 + small-mid N (4-32K) + BS>=64 configs whose CUDA-side numbers also moved 5-15x between runs -- short-runtime measurement noise, not real regressions. This commit completes the SASS-alignment campaign objective (gap < 5% on median config). Remaining ~10-13% at very large N is deferred. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
Three new switches gate the block_count_ge vec-load fast path:
enable_unroll_4 (default True): 4-way unrolled fast path
enable_unroll_2 (default by dtype): 2-way cascade between fast and tail
use_strided_layout (default by dtype): True → single make_ptr +
(UNROLL, vec_w) strided layout (cute emits 4 LDG.E.128 sharing
base reg with +0x2000/+0x4000/+0x6000 imm offsets, matching the
CUDA SASS pattern). False → 4 separate make_ptr calls (matches
the prior b459a8f commit style with 4 independent base regs).
Dtype-aware defaults (validated via per-config A/B testing on B200):
bf16 / fp16: enable_unroll_2=True, use_strided_layout=True
Strided cascade gives clean wins: cascade flips DSL from CUDA
parity to consistently faster on small-N where the 4-way fast
path doesn't fully cover N, and the medium 2-way path keeps two
LDG.E.128 in flight. Strided layout keeps the SASS shared-base
pattern that nvcc/ptxas auto-partial-unroll also produces.
fp32: enable_unroll_2=False, use_strided_layout=False
For fp32 (vec_w=4) the strided layout pushes regs 38 → 40 and
regresses fp32 large-grid configs by 30-60pp (worst observed:
K=1024 BS=128 nn=2 → 0.753 vs 1.364 with separate-ptrs). The
cascade similarly hurts in 12% of fp32 configs. Separate-ptrs
4-way unroll alone is the sweet spot.
Cache key includes the three switches so different settings produce
separate compiled kernels.
Full sweep results (804 configs, n_iters=30 kineto, L2 flush):
baseline cascade-all dtype-policy
Median C/G: 0.988 1.011 1.006
Geomean C/G: 0.999 1.047 1.038
DSL faster %: 46% 54% 52%
Within 10%: 55% 65% 65%
By dtype:
bf16: 1.000 -> 1.043 (cascade wins preserved)
fp16: 1.022 -> 1.038
fp32: 0.951 -> 0.960 (anom K=1024 BS=128 fixed: 0.610 -> 1.038)
By N (the original "large-N gap"):
N=8192: 1.097 -> 1.172 (+7pp, cascade hides medium-path remainder)
N=65536: 0.897 -> 0.928 (+3pp)
N=131072: 0.866 -> 0.923 (+6pp)
Remaining slow configs (fp32 K=2048 + BS>=64) were already <0.7 in
the baseline -- this commit doesn't introduce new regressions there.
Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
Adds two new switches to the DSL GVR kernel:
enable_phase3_unroll (default True): master gate for phase3_collect
unrolling. When ON, the inner enable_unroll_4 / enable_unroll_2
switches independently control 4-way fast and 2-way medium paths
in phase3 (same semantics as block_count_ge). When OFF, only the
tail 1-way loop runs.
use_constant_hint (default False): True → CopyG2ROp(invariant=True)
→ SASS LDG.E.*.CONSTANT (read-only data cache, matches CUDA
__ldg). Default False because cute's invariant lowering triggers
aggressive rematerialization in LLVM/NVPTX (+272 inst, 4 spills,
net -7pp geomean), outweighing the cache hint benefit.
Phase3_collect is now a 3-tier cascade (4-way fast + 2-way medium +
1-way tail) mirroring block_count_ge. The cascade gives:
N>=65K: +5-7% (large-N main path, LSU pipelining wins)
N<=32K: -1-3% (unroll setup overhead exceeds benefit at small N)
Median geomean: +2.2pp from phase3 unroll alone
Resource analysis (bf16/fp16/fp32 x phase3 ON/OFF):
REG/thread:
bf16: 39 -> 39 (no change, cute reuses fragments)
fp16: 39 -> 39 (no change)
fp32: 38 -> 40 (+2, separate-ptrs path)
Static SASS:
bf16: 3936 -> 4368 (+11%)
fp16: 4096 -> 4512 (+10%)
fp32: 3368 -> 3480 (+3%)
Theoretical Occupancy: 75% all configs (smem-limited to 3 blocks/SM,
binding limit unaffected by phase3 unroll). Phase3 unroll has
*zero* occupancy cost.
Wrapper signature gains both switches; cache key includes them so
different settings produce separate compiled kernels. A small helper
method _make_load_copy_atom() factors out the CopyG2ROp/Universal
selection to avoid Python if-else NameError inside @cute.jit scope.
Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
…uristic Phase 2 (block_count_ge) and Phase 3 (phase3_collect_candidates) replace the manual `while + range_constexpr(UNROLL)` fast/medium-cascade unrolling with a single `for k in cutlass.range(big_iters, unroll=4)` loop. LLVM's loop unroll pass + GVN/CSE folds the 4 derived vec loads into the CUDA-style shared base + immediate offsets pattern, emitting 4 back-to-back LDG.E.128 [base+0x2000/0x4000/0x6000] instructions. Add `min_blocks_per_mp` field on `GvrTopKKernel` and a 3-tier shape-aware heuristic in the host wrapper: * n_vec_iters < 4 -> 0 (no launch_bounds, natural ptxas allocation) * num_rows <= 148 (B200 SMs) -> 1 (allow many regs, 4xLDG fold survives) * else -> 3 (keep 3 CTA/SM occupancy, ~42 reg cap) The heuristic lifts fp32 K=512 large-N out of its regression zone (worst case C/G 0.62 -> 1.07 at K=512 N=131072 BS=16 nn=2). Cache key extended so each min_blocks value gets its own compiled kernel. Random sweep vs phase3_unroll baseline (804 configs): geomean 1.060 -> 1.149, faster%-than-CUDA 62% -> 92%, losses 304 -> 63. CUDA Graph: heuristic reads `logits.shape` (host int, no GPU sync) so capture is safe; per-graph capture selects the right kernel per shape. For dynamic-shape single-graph use, caller can pin `min_blocks_per_mp=3` to disable the heuristic. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
Three new kernel knobs on GvrTopKKernel + gvr_topk_decode host wrapper:
* use_256bit_load (default False): emit LDG.E.256 (8 fp32 / 16 bf16-fp16
elements per LDG) instead of LDG.E.128. Address alignment hint is
raised from 16 to 32 bytes. Phase 2/3 unroll factor is dtype-aware:
fp32 keeps unroll=4 (no cvt-to-fp32 overhead); bf16/fp16 drops to
unroll=2 to limit the cvt register pressure that otherwise spills
under min_blocks=3.
* num_threads_per_block (default 512): configurable per-instance.
BLOCK_SIZE / WARP_SIZE / NUM_WARPS are moved from module-level to
GvrTopKKernel instance attrs (self.WARP_SIZE, self.num_threads,
self.num_warps). Phase 1 preIdx loop gains an else branch for the
K < num_threads case (e.g. num_threads=1024 with K=512): only the
first K threads load a preIdx, others keep reduction-identity
values which the warp/block reduces naturally absorb.
* vec_bits / vec_align_bytes derived from use_256bit_load; cache key
extended with use_256bit_load + num_threads_per_block.
Heuristic uses the resolved num_threads_per_block (not a hardcoded 512)
when computing n_vec_iters.
Tests parametrize use_256bit_load and num_threads_per_block; pytest
runs 288/288 PASS at use_256bit_load=True and at num_threads=1024.
Synth bench on BS<=128:
- 128-bit + heuristic baseline: gm=1.131, faster%=99%, lose=9
- 256-bit + heuristic : gm=1.121 (fp32 wins +3pp; bf16/fp16
flat-to-negative due to cvt-to-fp32 reg pressure spills under mb=3)
Random sweep on BS up to 128: 256-bit shows niche win on
(fp32, num_rows<=148, large N); should be opt-in.
Synth data generator (multi-BS) and bench script env vars
(DSL_USE_256BIT/DSL_MIN_BLOCKS/DSL_NUM_THREADS) live in the gvr-topk-opt
workspace and are not part of this commit.
Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
…nobs
Add a fourth perf knob `enable_warp_parallel_reduce` to GvrTopKKernel +
gvr_topk_decode and replace the four `tid==0` serial loops over
num_warps slots with warp-parallel reduce/scan in warp 0:
* Phase 1 block aggregate (4-way reduce):
min/max/sum_f32/sum_i32 -> 4x warp_reduce in warp 0.
* Phase 2 / blockCountGE total (1-way reduce):
sum_i32 -> warp_reduce_sum_i32.
* Phase 3 collect block prefix sum (exclusive scan):
Hillis-Steele inclusive scan via block_scan.warp_scan, then
exclusive = inclusive - val; total = inclusive at last lane.
* Phase 2 secant aggregate (3-way reduce):
packed sum_i32 + min_f32 + max_f32, with bound update on lane 0.
Default is False since at num_threads=512 (num_warps=16) the per-warp
ILP loss exceeds the serial-loop savings (~2pp regression on synth).
At num_threads=1024 (num_warps=32) the switch is essential -- without
it 1024 regresses vs baseline (gm 1.131 -> 1.123); with it 1024 wins
(gm -> 1.154 on synth BS<=128). Pair as
`enable_warp_parallel_reduce = (num_threads_per_block >= 1024)`.
Phase 1 also gains an `active_preidx_warps` optimization: when
`pre_idx_count < num_threads` (e.g. K=512 with num_threads=1024) only
the first ceil(K/32) warps have real data, so the warp_reduce + smem
write step is now gated to those warps. Saves ~30 cy/dummy-warp; the
full barrier afterwards still keeps all 1024 threads aligned for
Phase 2. The constexpr is clamped to num_warps so the K>num_threads
case (K=2048 with num_threads=512) doesn't index past the smem
buffers, and the same value drives both the warp_reduce gate and the
Site-1 block aggregate's smem read range.
Remove two now-dead switches:
* `enable_unroll_2` -- only referenced in the commented-out manual
2-way medium path that the `cutlass.range(unroll=4)` rewrite
replaced.
* `use_strided_layout` -- only referenced in the commented-out manual
4-way strided-layout path, also replaced.
Cache key drops the two dead entries and gains
`enable_warp_parallel_reduce`. The cleanup is a no-op functionally
(the dead values were ignored by the active code paths) but removes
two cache-bucket dimensions.
Test parametrize expanded to 4-way matrix:
next_n in {1, 2} (was {1, 2, 3, 4} -- trimmed to keep walltime)
use_256bit_load in {False, True}
num_threads_per_block in {512, 1024}
enable_warp_parallel_reduce in {False, True}
1152 / 1152 PASS in 20:22.
Synth bench (BS<=128, threads=512 baseline -> threads=1024+wpON):
geomean 1.131 -> 1.154 (+2.3pp)
fp32 geomean 1.127 -> 1.177 (+5.0pp; up to +28pp at fp32 K=2048
N=131072 -- 1.50x vs CUDA)
Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
Contributor
📝 WalkthroughWalkthroughThis PR adds test coverage and public API exposure for the Blackwell CuTE ChangesGVR Top-K Kernel Export and Testing
🎯 2 (Simple) | ⏱️ ~10 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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/test_gvr_topk_decode.py`:
- Around line 29-45: The test picks argmax_idx over the full N, which breaks the
pre_idx[...,0] invariant when the kernel only scans the first (N - next_n + 1)
columns; update _make_inputs to compute an effective_scan = N - next_n + 1 (or
accept next_n as a parameter) and compute argmax_idx = int(logits[0,
:effective_scan].argmax().item()) so pre_idx_list[0] is in-range, then build
pre_idx as before; apply the same change to the other occurrence (the second
_make_inputs usage around the later test).
- Line 21: Remove the unused "from typing import Tuple" import and replace the
two occurrences of typing.Tuple[...] return annotations in this test module with
the native Python 3.10+ generic syntax tuple[...] (e.g., change "Tuple[int,
str]" to "tuple[int, str]") in the two functions in
tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/test_gvr_topk_decode.py so
the file uses built-in tuple typing and no longer imports typing.Tuple.
🪄 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: 031d6484-acc6-4f3e-b7b3-e6defb77c52c
📒 Files selected for processing (3)
tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/__init__.pytensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.pytensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/test_gvr_topk_decode.py
Mirrors the kFTarget=kK alignment for K=512/1024 (all dtypes) from CUDA PR NVIDIA#14413 on the DSL GVR Top-K kernel so the DSL Phase-2 secant behavior matches the new CUDA reference. Old pre-NVIDIA#14413 values kept as inline comments for easy rollback. Verified: 768/768 pytest configs pass for K=512/1024 across all dtypes, N, next_n, use_256bit_load, num_threads_per_block, and enable_warp_parallel_reduce. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
Move module-level MAX_REFINE_ITERS / FLT_MAX / NEG_FLT_MAX into instance attributes so all kernel-wide knobs live in one place. Inline NUM_BINS_DEFAULT (2048) directly into the GvrParams table since it was only used in three K=2048 entries. Drop dead MAX_CANDIDATES. Pure refactor — values, control flow, and DSL IR are unchanged. Also removes the previously-commented-out A/B layout/unroll dead code in block_count_ge. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
Mirrors heuristicTopKDecode.cu PR NVIDIA#14219 cr-aware branch in the DSL GVR Top-K kernel. compress_ratio=1 (default) preserves DSv3.2 behavior exactly; compress_ratio=4 enables the DSv4 (overlap-compressor) indexer path: * pre_idx_offset = 0 (vs (row % next_n) + 1 for cr=1) — in compressed- index space, new entries append at the end so prev-step indices remain valid as-is. * N = actual_kv_len / cr — logits/preIdx live in compressed-token- index space when cr > 1. GvrParams TABLE is also keyed by (dtype, K, cr) so V3.2 and V4 use their respectively tuned kFTarget values: cr=1 (V3.2): kFTarget = 384 (K=512) / 2560 (K=1024), pre-NVIDIA#14413. cr=4 (V4): kFTarget = kK = 512 (K=512) / 1024 (K=1024), PR NVIDIA#14413. K=2048: identical across cr (V4 doesn't natively use K=2048). Cache key includes compress_ratio so different cr settings compile separate kernels. assert restricts compress_ratio in {1, 4}. Verified: 1152/1152 pytest configs pass on cr=1 default path. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
…3 mb
Two paired host-wrapper heuristic refinements:
1. enable_warp_parallel_reduce: bool → Optional[bool] = None, default
auto-coupled to num_threads_per_block: enabled iff threads == 1024 (32
warps, where serial tid==0 cost dominates). At threads == 512 (16
warps) the warp-parallel path measured a ~2pp synth regression so it
stays off. Cache key sees the concrete bool. Explicit True/False still
overrides for A/B testing.
2. tier-3 (large grid + large N) min_blocks_per_mp hardcoded "= 3"
replaced by a (T, dtype) lookup:
T == 1024 or dtype == fp32 → mb=2
T == 512 and dtype in (bf16, fp16) → mb=3
Derived from BS{256,384,512} × N{16K,32K,65K} × all 9 (dtype, K) sweep
(gvr-topk-opt/sweep_tv_mb_kineto/mb_sweep.png). Old mb=3 default
regressed by 25-37% on (T=512 + fp32 + large N/BS) configs because
cap=42 starves the 4-LDG-inflight ILP (fp32 vec_w=4 × unroll-4 needs
50+ regs). bf16/fp16 keep mb=3 since cvt-to-fp32 ILP fits in 40 regs
and the extra CTA/SM (3 vs 2) hides cvt latency.
Pure default-policy change — no behavioral effect when caller passes
explicit values. Verified: pytest smoke 4/4 on cr=1 default path.
Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
Two more host-wrapper Optional[*]=None defaults so callers no longer
need to pick threads/vec-bits per shape:
num_threads_per_block (None default):
1024 iff num_rows <= num_sms (1 CTA/SM bound) AND N >= 65536
(so each of the 1024 threads has meaningful vec-loop work).
Otherwise 512.
use_256bit_load (None default):
True iff dtype == fp32 AND N >= 16384.
Half-prec (bf16/fp16) cvt-to-fp32 doubles fragment reg footprint
and regresses 5-11% at K=512/1024; LDG already saturates at 128b
anyway. fp32 N=8K dips 5-8% with 256b at small grid so the N
threshold excludes that single tier.
Cache key sees concrete values; (None, X) and (None, Y) hash apart.
Explicit values still override for A/B testing.
Derivation: sweep BS{1,4,16,64,128,256,384,512} x N{4K..131K} x all 9
(dtype, K), gvr-topk-opt/sweep_tv_kineto/auto_speedup.csv. Net vs
baseline (T=512, V=128):
- median speedup vs CUDA 1.09x -> 1.10x
- mean speedup vs CUDA 1.11x -> 1.15x (+3.8pp)
- max speedup vs CUDA 1.45x -> 1.52x
- 21 of 22 sp<1 configs were already sp<1 in baseline (BS=384 grid
quirk, unrelated to this change). 1 new config introduces a 0.8pp
sp<1 dip (within bench noise).
Pure default-policy change. Verified: 4-case auto-path smoke + pytest
smoke 4/4 on cr=1 default.
Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
Decode runs under CUDA graph, where the (T, V) heuristic baked in at
capture time is reused across all replays. The capture-time
logits.shape[1] is typically much smaller than peak runtime N, so the
captured kernel misses the large-N (T=1024, V=256) path. Add an
optional max_seq_len hint so the caller (e.g. dsa.py) can pass the
peak compressed-N for the model; the heuristic then tunes the captured
kernel for the peak.
Usage guidance baked into the docstring + inline comment:
* CUDA Graph mode: CALLER MUST PASS max_seq_len.
* Eager mode: leave max_seq_len=None (heuristic adapts per call).
Rules with max_seq_len:
* T=1024 threshold becomes dtype-aware to avoid half-prec K=512/1024
small-N replay regression (14-16% when forced T=1024 at small N):
fp32 -> 65536 (small-N replay 1-9% loss, net win)
half -> 131072 (only forced at very large peak)
* V=256 still gated by fp32 + N >= 16384.
Without max_seq_len, dtype-split is NOT applied because per-call
adaptive decisions never force T=1024 onto small N — heuristic only
fires for N >= 65536 by definition — so the half-prec N=65K-128K
+4-6% T=1024 win is preserved.
Cache key sees concrete (T, V), so different max_seq_len hints compile
distinct kernels. Pure default-policy extension. Verified with 4-case
auto smoke (no hint / fp32+131K / bf16+131K / bf16+200K).
Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
The min_blocks_per_mp tier heuristic was still computing n_vec_iters from logits.shape[1] (capture-time N). In graph mode with max_seq_len hint, this would stick small-capture-N calls in tier-0 (mb=0) and miss the tier-3 occupancy choice for large-N replays — same pitfall the (T, V) heuristic was fixed against in the previous commit. Switch to N_dec (= max_seq_len if provided, else logits.shape[1]) so the tier classification is consistent with how T/V are picked. Smoke 4/4. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
Add a wave-fit branch in the fp32 tier-3 path: when num_rows ∈ (296, 444]
(i.e. fits 1 wave at 3 CTAs/SM but needs partial 2nd wave at 2 CTAs/SM
with num_sms=148), pick mb=3 instead of mb=2. This recovers ~15% perf
on fp32 BS=384 across (K, N) — verified against CUDA which already uses
__launch_bounds__(BS, 3) for this exact reason.
Math:
mb=2 cap → 2 CTAs/SM × 148 SMs = 296 CTAs in 1 wave.
mb=3 cap → 3 CTAs/SM × 148 SMs = 444 CTAs in 1 wave.
For BS=384 (× next_n=1):
mb=2: 384 / 296 = 1.30 waves → tail wave wastes ~70% SMs.
mb=3: 384 / 444 = 0.86 waves → 1 wave fits, max SM utilization.
Verified perf gains (fp32 T=512, both V=128 and V=256 default paths):
fp32 K=512 N=4K-32K BS=384: +11-23%
fp32 K=1024 N=4K-32K BS=384: +16-19%
fp32 K=2048 N=8K-32K BS=384: +5-9%
Other BS unaffected:
BS ≤ 296 (192, 256): mb=2 already fits 1 wave → rule keeps mb=2 (no change)
BS > 444 (512): both need >1 wave → rule keeps mb=2 (ILP > occupancy)
Half-prec heuristic unchanged (already uses mb=3 in tier-3 via the
dtype-split path from a prior commit).
Bench artifacts: gvr-topk-opt/auto_full_bench/fp32_bs384_cluster/
(mb sweep CSV + NCU reports + drivers). Smoke: pytest 4/4 + spot tests
across BS={256, 384, 512}.
Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
Functional change: limit wave-fit mb=3 branch to N <= 32768. Beyond
that threshold the kernel becomes bandwidth-bound and mb=3's 3-way L2
sharing causes contention; mb=2's lower occupancy gives each CTA more
bandwidth and wins +21-30% at fp32 K=512 N=65K BS=384.
The full wave-fit rule for fp32 tier-3 is now:
if 2*num_sms < num_rows <= 3*num_sms and N_dec <= 32768:
mb = 3
else:
mb = 2
Also cleans up file comments: remove obsolete TODO list, trim refs to
specific CUDA line numbers, simplify class docstring.
Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
limin2021
changed the title
limin2021
changed the title
Registers torch.ops.trtllm.cute_dsl_gvr_topk_decode as the production entry point for the cuTe DSL GVR Top-K decode kernel (Blackwell SM100). Op writes values + indices into caller-allocated buffers (mutates_args style), matching the existing cute_dsl_indexer_topk_decode pattern so the DSA indexer pipeline can drop it in. CuteDSLGvrTopKDecodeRunner takes ownership of the JIT compile cache and the auto-heuristic for T (threads/block), V (vec-load width), min_blocks_per_mp and enable_warp_parallel_reduce. The previous module-level wrapper in gvr_topk_decode.py is removed; standalone bench / A-B testing with the full tuning knob set lives in tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py. Tests: - tests/unittest/.../test_cute_dsl_gvr_topk_decode.py: production correctness sweep via the op (dtype x K x N x next_n x batch_size x compress_ratio) with vectorized tie-aware + strict sort+allclose check. - tests/scripts/.../run_gvr_topk.py: dual-mode driver -- pytest sweep over T/V/wp knobs and standalone CLI for single-case verification. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
Collaborator
Author
|
@longcheng-nv Thanks for the thorough review. Replies per item below. #1 — Perf table (replace PNG)Added 3 tables to PR description:
Median 0.964× → 1.152× across 408 configs; sp<1 dropped 230 → 7. #2 — Alignment check for
|
- Document 32B alignment caller contract when use_256bit_load=True (no runtime assert; matches existing DSL op convention — see PR reply for risk-model rationale) - Clarify return_output_values policy: op hardcodes False (matches CUDA indexer_topk_decode); kernel retains True branch for future caller flexibility - Add preidx_hit_rate parametrize axis to op unittest (0.0 worst-case + 0.5 realistic, matching V3.2/V4 Pro production preIdx∩topK overlap); test matrix 144 → 288 - Add 'Not in CI' header docstring to standalone driver run_gvr_topk.py explaining the trtllm-runtime-free design - Add --num_sms CLI to standalone driver for heuristic-edge debug Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #51301 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #51301 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #51310 [ run ] triggered by Bot. Commit: |
yuxianq
reviewed
Jun 1, 2026
longcheng-nv
reviewed
Jun 1, 2026
longcheng-nv
left a comment
longcheng-nv
left a comment
Collaborator
There was a problem hiding this comment.
Follow-up review — please address the barrier race + the still-open alignment guard
Thanks @limin2021 for the round-2 updates. Re-reviewing at head e219a000… against my previous review (2026-05-29) and aligning with @yuxianq's inline note from earlier today.
📋 Status of prior asks
| # | Ask | Status | Notes |
|---|---|---|---|
| 1 | Replace perf PNG with numerical table | ✅ Addressed | PR body now has a 408-row table across (dtype × K × N × BS) + per-optimization cumulative table (v1 → v5). Strong evidence. |
| 2 | Runtime assert when use_256bit_load=True |
❌ Open | No data_ptr() % 32 check anywhere in the new code; still a real footgun for downstream callers (views, padded layouts). |
| 3 | Clarify return_output_values policy |
✅ Addressed | Op-level docstring + the comment "kernel keeps both branches to support enabling value writeback in the future" makes the intent explicit. |
| 4 | Realistic preIdx hit-rate test |
✅ Addressed | Unit test now parametrizes preidx_hit_rate ∈ {0.0, 0.5}; covers both worst-case and realistic Guess-phase short-circuit. |
🐛 New / outstanding correctness asks
A. Confirming @yuxianq's barrier race in phase4_histogram_snap Branch B (gvr_topk_decode.py L1497–L1521)
Yuxianq's concern is correct — there is a missing cute.arch.barrier() between the "every thread recomputes block_min/block_max from staged warp slots" read and the "zero histogram" write that follows. Reproducing the sequence:
# L1488–L1492: lane==0 stages cmin/cmax bits → smem_wcnt[w] / smem_hist[w]
cute.arch.barrier() # L1493 ✓
# L1498–L1511: every thread (re)reads smem_hist[w] for w ∈ [0, num_warps)
for w in cutlass.range_constexpr(self.num_warps):
vmin_bits = smem_wcnt[w]
vmax_bits = smem_hist[w] # READ smem_hist[0..NW)
...
# L1517–L1520: zero histogram across [0, kBins)
i6 = tidx
while i6 < cutlass.Int32(kBins):
smem_hist[i6] = cutlass.Int32(0) # WRITE smem_hist[0..kBins)
i6 = i6 + cutlass.Int32(num_threads)
cute.arch.barrier() # L1521The unrolled read loop runs in-warp-lockstep but not in-block-lockstep — warp 0's thread 0 can finish the entire read loop and start zeroing smem_hist[0] while warp 1 is still reading smem_hist[0]. Result: vmax_bits (and therefore bmax_r) is read as 0, the patched bmax_r ≤ bmin_r → bmax_r = bmin_r + 1e-6 may fire spuriously, and the histogram bin width gets squashed → all candidates land in bin 0 → wrong K-th threshold.
This is hit-rate-dependent and won't show on most synth distributions (which is why the 144-cell unit sweep passes), but I'd expect it to surface under adversarial inputs.
Fix: insert a cute.arch.barrier() after the read loop. The cleanest spot is between the bmax_r ≤ bmin_r patch and the zero loop:
if bmax_r <= bmin_r:
bmax_r = bmin_r + cutlass.Float32(1e-6)
cute.arch.barrier() # ← ADD: ensure all reads of smem_hist[0..NW) complete before zeroing
# Zero histogram
i6 = tidx
while i6 < cutlass.Int32(kBins):
...B. Ask #2 from previous review — still open
Re-quoting for visibility; should land in this PR rather than as a follow-up since it protects the op's public contract:
if use_256bit_load:
assert logits.data_ptr() % 32 == 0, (
f"256-bit vec load requires 32B-aligned logits.data_ptr(), "
f"got {logits.data_ptr()} % 32 = {logits.data_ptr() % 32}"
)🔍 Minor (nice-to-fix in this PR, won't block)
-
Stale comment in
phase4_histogram_snapBranch B (gvr_topk_decode.py around thebmin/bmaxrecompute loop):# Note: unrolled for 64 times.The constexpr unrolls
num_warpstimes, which is 16 or 32, never 64. Looks like leftover from an earliernum_threads=2048variant. -
Docstring/allocation mismatch in
phase1_preidx_stats:s_thr, # cute.Tensor [3] float32: [threshold, val_lo, val_hi, pmax_saved]
Says 4-wide but allocated as
(3,);pmax_savedlives in the separates_thr_extraallocation. Cosmetic but misleading. -
packed = (lge << 16) | lgtinblock_fused_snap_itersilently assumeslge, lgt < 2^16 = 65536. The current bound iscand_count ≤ kC ≤ 6144, so OK — but worth astatic_assert-style const check / comment so a future bump ofkCdoesn't quietly corrupt the packing. -
num_sms: int = 148hardcoded default intests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py(gvr_topk_decodesignature). The standalone driver's__main__and pytest path both override it viatorch.cuda.get_device_properties(0).multi_processor_count, so this is only a footgun for someone callinggvr_topk_decode()directly on B300 (132 SMs/die) without overriding. Either passnum_smsalways or remove the default. -
Perf regression cells — the v5 table lists
sp<1: 7/408, but specific cells worth flagging in the PR body so future maintainers know which corners the heuristic loses:fp32 K=512 N=4096 BS=384: 0.777× (largest regression)bf16/fp16 K=1024–2048 N=131072 BS=512: 0.912× / 0.933× / 0.957×
The pattern (high-BS / large-N / fp32-small-K corner) suggests the
min_blocks_per_mptier-3 rule may need one more sub-bucket, but acceptable as a follow-up if the per-step impact on DSv4 indexer is bounded.
Status
Holding at COMMENT until A (barrier race) and B (alignment guard) land. Happy to flip to APPROVE once both are in. cc @Kefeng-Duan / @yuxianq for visibility.
yuxianq
reviewed
Jun 1, 2026
yuxianq
reviewed
Jun 1, 2026
yuxianq
reviewed
Jun 1, 2026
Bug fixes: - phase4_histogram_snap: add missing cute.arch.barrier() between the smem_hist[0..NW-1] read loop (recomputing block_min/block_max from warp-staged cmax slots) and the histogram-zeroing write loop. Without this, warp-0 thread-0 could finish the unrolled read and start zeroing smem_hist[0] while warp-N was still reading the staged cmax → squashed bmax_r → all candidates land in bin 0 → wrong K-th threshold. Hit-rate-dependent; covered by existing tests once they span realistic preIdx hit rates (NVIDIA#4 from prior review). - Add runtime 32B alignment assert on logits.data_ptr() when the use_256bit_load heuristic fires (LDG.E.256 path). Catches view-with- unaligned-offset misuse with a clear error rather than silent miscompiled addresses / faults. Code hygiene: - Fix stale "unrolled for 64 times" comment in the block_min/max recompute loop; it unrolls num_warps times (16 or 32). - Fix phase1_preidx_stats s_thr docstring: [3] [threshold, val_lo, val_hi]; pmax_saved lives in the separate s_thr_extra allocation. - Document the (lge << 16) | lgt packing's <2^16 per-warp count precondition (currently holds via cand_count ≤ kC ≤ 6144 in GvrParams; future kC bump past 65536 would silently corrupt). Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
s_thr_extra was a [1]-wide fp32 smem buffer written twice in phase1_preidx_stats (parallel + serial paths) to "save pmax", but never read anywhere in the kernel — leftover from an earlier design that was refactored. The same value is already in s_thr[2] at the point of the redundant write; subsequent secant updates of s_thr[2] don't need the original pmax. Drops: smem allocation, phase1_preidx_stats parameter + call-site arg, and 2 write statements. Reported by @yuxianq in PR review. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
Collaborator
Author
|
@longcheng-nv fixed the 2nd round review. For Minor 4th question, please ignore this one. This script is only used for local debug/test in no-trtllm env. These are not production path. |
Collaborator
Author
|
/bot run |
longcheng-nv
approved these changes
Jun 1, 2026
longcheng-nv
left a comment
longcheng-nv
left a comment
Collaborator
There was a problem hiding this comment.
✅ Approved — all follow-up asks landed cleanly
Verified against head 97c47dc8 (commits d23ba74a + 97c47dc8 since my COMMENT review).
| Ask | Status | Notes |
|---|---|---|
A. Barrier race fix in phase4_histogram_snap |
✅ | cute.arch.barrier() inserted between the smem_hist[0..NW-1] read and the zero loop, with reasoning in the comment. |
B. 32-byte alignment assert for use_256bit_load=True |
✅ | Guards logits.data_ptr() % 32 == 0 with a clear error pointing at column-slice / stride-padded layouts as the cause. |
| Minor 1 — stale "unrolled for 64 times" comment | ✅ | Updated to reflect the actual num_warps-times unroll (16 or 32). |
Minor 2 — s_thr docstring vs allocation |
✅ | Went beyond the doc fix: removed the dead s_thr_extra allocation entirely (was written but never read). |
| Minor 3 — 16-bit packing assumption | ✅ | Comment block documents the lge/lgt < 2^16 invariant tied to kC ≤ 6144 and warns about future kC bumps. |
Delta is 5 hunks across the two files I previously reviewed, no collateral substantive changes. Race fix is the right scope (per-row CTA, smem-only).
Thanks @limin2021 for the fast turnaround, and thanks @yuxianq for catching the barrier race. Approving.
yuxianq
approved these changes
Jun 1, 2026
Collaborator
|
PR_Github #51321 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #51310 [ run ] completed with state |
hyukn
approved these changes
Jun 1, 2026
Collaborator
|
PR_Github #51321 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #51357 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #51357 [ run ] completed with state |
mingyangHao
reviewed
Jun 2, 2026
Collaborator
Author
|
TODO: test more N, including [K, 4096] |
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
7 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.
Description
Add cute dsl gvr top-k decode kernel.
(1) port cuda gvr kernel to cute dsl. Thx for Long's help, who gives the 1st version porting code.
(2) add some extra optimizations, e.g., unroll, 256bits vectorization, num_threads tune, to further improve the perf.
for more visualization：
Cheng, Long, et al. "Guess-Verify-Refine: Data-Aware Top-K for Sparse-Attention Decoding on Blackwell via Temporal Correlation." arXiv preprint arXiv:2604.22312 (2026). https://doi.org/10.48550/arXiv.2604.22312
Test Coverage
Performance
Hardware: B200 SXM5 1000W + Intel Xeon Platinum 8570
Data:
swebench-temporal-synthbeta_deepbundlesBench:
bench_kineton_iters=30, warmup=1, L2 flushEnv:
TRTLLM_FORCE_HEURISTIC_TOPK=1TRTLLM_HEURISTIC_NMIN=1024DSL_AUTO=1Baseline: CUDA
heuristicTopKMultiRowKernel(production GVR Top-K)Per-optimization class gain (cumulative)
min_blocks_per_mpautoreturn_output_values=FalsePR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.