[None][feat] Improve cute dsl radix top-k - #15756
Merged
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>
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>
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>
The DSA indexer pipeline (dsa.py) and CUDA indexer_topk_decode op both
only read top-K indices from the kernel output; the value buffer is
caller-allocated scratch that's never consumed. Add a kernel-level
return_output_values switch so the cuTe DSL kernel can elide all
STG.value stores when the caller doesn't need them.
Kernel (gvr_topk_decode.py):
- GvrTopKKernel gains return_output_values: bool = True.
- All 9 STG.value sites + the output_values_row slice are gated under
cutlass.const_expr(self.return_output_values), letting cute.compile
eliminate the dead writes when False.
Op (cute_dsl_custom_ops.py):
- CuteDSLGvrTopKDecodeRunner adds return_output_values to the compile
cache key + _compile signature; forward() hardcodes False, drops the
output_values arg, and passes None for the value-output slot at
launch (mirrors the optional-fake-tensor pattern at
CuteDSLTopKDecodeMultiCTARunner._compile).
- trtllm::cute_dsl_gvr_topk_decode op signature drops output_values;
mutates_args is now ("output_indices",), aligning with CUDA's
indexer_topk_decode which also only exposes indices.
Tests:
- tests/unittest/.../test_cute_dsl_gvr_topk_decode.py drops the
output_values buffer alloc + op kwarg (all 144 cases pass with the
sort+allclose strict check).
- tests/scripts/.../run_gvr_topk.py wrapper exposes
return_output_values as a knob so the standalone driver can still
capture written values; the _compile cache + cute.compile
out_values_fake placeholder are conditional on the flag.
SASS verification at bf16 K=1024 N=8K BS=384 confirms 88 STG.E.U16
writes are eliminated (kernel cubin -6KB, total SASS -416 lines).
On B200 SXM5 + synth_data, v5 (return_output_values=False) gives a
median 1.2% latency improvement over v4 with sp<1 configs nearly
halved (52 -> 27).
Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
* Modernize type hints: replace ``typing.Tuple[...]`` with the built-in ``tuple[...]`` (Python 3.10+) and drop the ``from typing import Tuple`` import in both ``_make_inputs`` / ``_tie_aware_correct`` helpers. * Fix the ``pre_idx[..., 0]`` argmax invariant for next_n > 1: argmax must come from the kernel's effective scan range ``[0, N - next_n + 1)``, not full ``[0, N)``. With the prior full-N argmax, an index landing in the ``[N_eff, N)`` tail could violate the CUDA-side ``preIdxCount == topK`` dispatch precondition (kernel still produced correct top-K because pre_idx is only a Phase-1 hint, but the test was technically exercising the kernel under invalid input). Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
The kernel derives batch_size as logits.shape[0] / next_n implicitly,
then sizes pre_idx / seq_lens / output_indices accordingly. If the
divisibility breaks, the failure modes are either an OOB write or a
ZeroDivisionError raised from deep inside the JIT-compiled kernel —
neither is actionable for callers. Add an upfront check in the op
body so the contract violation surfaces with a clear message.
Other invariants (top_k in {512,1024,2048}, compress_ratio in {1,4},
logits dtype) are already enforced by GvrTopKKernel.__init__ via
GvrParams.get and the dtype switch.
Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
* Drop the single-use N_cols local in CuteDSLGvrTopKDecodeRunner.forward; fold it directly into the N_dec ternary for less noise. * Add the input shape / dtype / knob signature to the info_once dedup key. Without the signature the first call's log message hid every subsequent shape from production diagnostics; now each new shape emits a single log line on its first run. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
Both _tie_aware_check (unittest) and _tie_aware_correct (run_gvr_topk
standalone driver) previously assumed the reference scan range was
``N - next_n + ofs + 1`` per row. That matches the kernel for cr=1
and for cr>=2 with next_n in {1, 2}, but breaks for cr>=2 with
next_n>=3 because floor-division by cr makes per-row N_eff vary
within a group in ways the simple closed form can't express.
Switch both reference helpers to mirror the kernel's exact formula:
actual_kv_len = seq_lens[row // next_n] - next_n + (row % next_n) + 1
N_eff = actual_kv_len // compress_ratio # cr=1 is identity
This requires the helpers to take ``seq_lens`` (and compress_ratio for
the standalone driver) so the reference can compute per-row N_eff
exactly as the kernel does. With this, any (next_n, cr) combo is
testable without the floor-division mismatch.
Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
- 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>
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>
New kernel GvrTopKClusterKernel + driver run_gvr_topk_cluster.py, forked
from V5 GvrTopKKernel. Adds an SM_100 thread-block-cluster path so the
dominant Phase 2 secant-search work can be parallelised across multiple
CTAs that share a single row.
Architecture (cluster_size = 2 default):
Grid = (num_rows * cluster_size, 1, 1)
Cluster = (cluster_size, 1, 1)
cta_in_cluster = cute.arch.block_idx_in_cluster()
cluster_id (= row) = bidx // cluster_size
Row slice per CTA: [cta_in_cluster*N/cs, (cta_in_cluster+1)*N/cs);
last CTA absorbs the N mod cs remainder.
Phase 1 (preIdx stats): every CTA scans pre_idx independently --
same input + deterministic math gives
identical s_thr / s_iscalars on every
CTA, no DSMEM broadcast required.
Phase 2 (secant search): every CTA scans its row slice; per-iter
block_count_ge ends with an all-reduce of
per-CTA cand_count through DSMEM
(s_cluster_partial slot at the same SMEM
offset in every CTA; mapa.shared::cluster
+ ld.shared::cluster.u32 from each peer).
The cluster total lands in s_iscalars[0]
on every CTA, and each CTA's tid==0 runs
the same secant update against it,
keeping s_thr in sync across the cluster
without an explicit broadcast.
Cluster handoff: one final cluster_arrive_relaxed +
cluster_wait. From here on only the
leader (cta_in_cluster == 0) participates;
peers fall through to the kernel epilogue.
The leader resets slice to [0, N), re-runs
block_count_ge with do_cluster_aggregation
= False on the full row to refresh
smem_ptcnt (Phase 2's last block_count_ge
cached slice-local counts, which would
give wrong Phase 3 prefix sums), then
proceeds.
Phase 3 (collect): leader-only, full row.
Phase 4 (histogram snap): leader-only.
cluster_size == 1 is a degenerate fast path -- DSMEM, cluster
barriers, and the leader-only handoff all compile out via const_expr
guards, so the kernel is byte-for-byte equivalent to V5 in that case.
DSMEM primitives (inline PTX, adapted from
single_pass_multi_cta_radix_topk_cluster.py):
mapa.shared::cluster.u32 translate a local SMEM ptr into a
peer CTA's address-space view.
ld.shared::cluster.{u32,f32} load int32 / fp32 from peer SMEM.
The MVP does not yet parallelise Phase 3/4 across the cluster (Phase
3+4 are leader-only and rescan the full row). That leaves the gather
optimisation -- "each CTA collects candidates from its slice, leader
DSMEM-merges into its smem_keys, then runs Phase 4" -- as the next
step. Estimated speedup of the current MVP at N=131K BS=64 is about
~30% (Phase 2 halves; Phase 3/4 unchanged; +2us for the leader's
full-row smem_ptcnt refresh).
Validation: 72 pytest cases pass on
dtype in {fp32, bf16, fp16}
K in {512, 1024, 2048}
N in {4096, 65536}
BS in {1, 32}
cluster_size in {1, 2}
All assertions use the tie-aware sorted-value comparator already used
by the V5 driver, so candidates that differ only at threshold ties
(common with random data) are still accepted.
Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
…input size TRUNCATE retains at most filtered_topk_smem_input_size threshold-bin candidates. When top_k exceeds that size (fp32 large_occupancy: S=8192, top_k up to 16384), a dense threshold bin can leave fewer than top_k valid candidates, and the unfilled s_indices tail is written back as undefined/duplicate indices -- violating the kernel's no-duplicate contract. This configuration is under-determined, so reject it at construction and direct callers to REREAD or GMEM_SPILL. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
…ants The CuTe DSL radix-filter decode top-k (cute_dsl_indexer_topk_decode) JIT-compiles per compile-key on first touch (~seconds). num_cols is fixed at indexer_max_seq_len for a deployment, but cluster_size = auto_cluster_size(num_cols, num_rows) steps across num_rows occupancy bands. CUDA-graph warmup only exercises cuda_graph_batch_sizes; eager decode iters (mixed prefill+decode batch, or cuda_graph disabled) whose num_rows lands in an uncovered band otherwise pay the JIT stall on a live request. Add warmup_cute_dsl_radix_topk_decode (custom_ops) that issues one decode per num_rows band so every cluster_size variant compiles during warmup, a DSA metadata method that supplies the deployment params, and a ModelEngine.warmup hook that runs it after the DG metadata prewarm. Mirrors the existing _warmup_dg_paged_mqa_logits_metadata pattern. No-op unless decode routes to the radix-filter DSL path. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
…atch modes cute_dsl_indexer_topk_decode selects among three dispatch modes via overlapping booleans whose precedence is only the if/elif ordering (radix_filter_single_pass_multi_cta > single_pass_multi_cta > 2-pass). radix_filter_single_pass_multi_cta defaults True, so a caller passing single_pass_multi_cta=True without also disabling it was silently rerouted to the radix-filter path. Rewrite the (stale, two-mode) docstring to describe all three modes and their precedence, marking radix-filter as the adaptive best-performance default and the others as legacy A/B-and-fallback overrides. Add an assert that rejects a conflicting mode-2/3 override when mode 1 is enabled, so the override fails loudly instead of being silently ignored. No behavior change for existing callers (none pass the conflicting combination). Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
…er cluster decode In the single-pass multi-CTA radix-filter decode path, a row whose effective length spans more than one chunk (need_cluster_sync) but is <= top_k has a cluster-merged histogram total (== eff_len) that never exceeds top_k, so the radix threshold search never fires and the output is undefined. The single-CTA path avoids this via its take_trivial shortcut, but the cluster path forces take_trivial off to keep every CTA cooperating on the barriers. Fix at the dispatch layer: when eff_len <= top_k there is nothing to filter, so collapse to the existing solo trivial fast path -- clear need_cluster_sync (cluster-uniform, so only cta 0 runs and no CTA waits on a barrier) and hand cta 0 the full row. No kernel change. Add a regression case (fp32 N=32768, top_k=16384, eff=12000) to test_cute_dsl_topk_decode_high_and_odd_k. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
The comment claimed no explicit CuTe DSL top-k warmup is needed; that is no longer accurate now that ModelEngine.warmup pre-compiles the radix-filter cluster_size variants. Remove it (review nit). Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
…k.py Consolidate the top-k custom-op tests in one file (review nit). Move the four overflow-policy tests and their unique helpers from test_cute_dsl_topk_overflow_policy.py into test_indexer_topk.py and delete the former. The moved tests reuse the generate_seq_lens / compare_top_k_results / create_random_logits helpers already defined in test_indexer_topk.py instead of re-importing the source-module copies, removing the duplication. 176 moved cases pass. Signed-off-by: Mindy Li <11663212+limin2021@users.noreply.github.com>
Collaborator
Author
|
/bot run --disable-fail-fast |
longcheng-nv
approved these changes
Jul 27, 2026
longcheng-nv
left a comment
longcheng-nv
left a comment
Collaborator
There was a problem hiding this comment.
All review comments are addressed on this head — approving.
Verified on 6004753:
- TRUNCATE
top_k> SMEM input size (57437ae): re-ran my duplicate-indices repro (fp32 N=32768, top_k=16384, large_occupancy, all-tie + randn, B200). The kernel now rejects the configuration at construction with a clearValueError, andREREADremains exact on the same input (0 bad rows). Docstring now states the constraint. - Warmup (
78751cd): the band-representative decode warmup covers exactly thecluster_sizecompile-key dimension the CUDA-graph path misses; the dispatch guards inwarmup_cute_dsl_radix_topkmirror the runtime routing (heuristic top-k / C++-op cases correctly skipped), and per-band failures are best-effort so a broken bucket can't abort startup. - Dispatch modes (
6981308): documented precedence + assert on conflicting mode-2/3 overrides, no behavior change for existing callers. - Also reviewed the
_eff <= top_kcluster fix (a4c3923, cluster-uniform collapse to the solo trivial path) with its regression test, and the test relocation intotest_indexer_topk.py.
One optional follow-up (non-blocking): a small negative test asserting the new ValueError for TRUNCATE with top_k > filtered_topk_smem_input_size would pin the contract.
Merge still needs a green CI run on this head.
Collaborator
|
PR_Github #61773 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #61773 [ run ] completed with state
|
mingyangHao
approved these changes
Jul 27, 2026
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #61845 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #61845 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #61894 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #61894 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #62049 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #62049 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #62182 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #62182 [ run ] completed with state |
yuxianq
approved these changes
Jul 29, 2026
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
10 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.
Summary by CodeRabbit
New Features
Bug Fixes
Description
This PR adds a radix-filter single-pass multi-CTA (cluster-DSMEM) top-k path
for the Blackwell (SM100) indexer top-k kernels, plus the overflow-policy
framework and dispatch heuristics needed to make it the default, and a large
round of correctness/perf/refactor cleanup of the existing radix top-k code.
What's new
1. Radix-filter single-pass multi-CTA kernel (cluster DSMEM).
A cluster of CTAs cooperates on one row within a single kernel launch: each CTA
histograms its chunk, then peer histograms are merged via distributed shared
memory (
mapa.shared::cluster+ld.shared::cluster) into a separate mergebuffer — no GMEM intermediate state, no second merge launch. Output offsets are
computed with a DSMEM exclusive prefix scan. Covers both decode and prefill.
2. Overflow-policy framework (
overflow_policy):REREAD(new default) — optimistic SMEM candidate buffer, re-scan the chunk onoverflow (exact).
TRUNCATE— drop overflow candidates (lossy, zero extra memory).GMEM_SPILL— spill overflow to a GMEM scratch buffer (exact).REREAD_ALWAYS— always re-scan (no SMEM candidate buffer).Default changed
GMEM_SPILL -> REREAD(exact and no scratch allocation).3. Dispatch heuristics in
cute_dsl_indexer_topk_decode:_radix_select_preferred): small-N bf16/fp16 route toradix-SELECT SP; gated on
num_rows <= num_smsso large-batch falls through tothe faster filter single-CTA.
cluster_sizeauto-tuning (auto_cluster_size/auto_cluster_size_truncate):policy-aware 2-ladder (peak-by-N ∧ occupancy-by-num_rows), REREAD variant hedges
the large-occupancy re-scan cliff.
4. Correctness fix.
DSMEM cluster-arrive publish sites use a release-fence
cluster_arrive()(notrelaxed) so a peer's
ld.shared::clustersees the published histogram; relaxedarrive risked stale reads. WAR/liveness sites stay relaxed (peer reads already
drained).
5. Perf & cleanup.
Occupancy-aware SMEM sizing (replaces hardcoded
smem_input_size), tightenedbarriers/fences in the refine loop,
cache_smem_valuesflag, device-queriednum_sms(drops hardcoded 148), and extensive dead-code/refactor removal.Notes
_radix_select_preferred,auto_cluster_size*) were tuned on aB200 randn / fixed-length-seqlen sweep; the docstrings carry that caveat.
Test Coverage
tests/unittest/_torch/thop/parallel/test_indexer_topk.py— exercisescute_dsl_indexer_topk_decode/ prefill across the single-CTA, 2-pass multi-CTA,and new single-pass multi-CTA paths (radix-filter prefix added to prefill tests).
tests/unittest/_torch/thop/parallel/test_cute_dsl_topk_overflow_policy.py(new) —correctness of REREAD / TRUNCATE / GMEM_SPILL / REREAD_ALWAYS, including the
multi-CTA short-row degrade boundary.
Performance
Measured on B200 (SM100). GPU-kernel time via CUPTI, cold L2 (8 GB flush before each launch, flush kernel excluded), randn fixed-length inputs, median of 40 iters. Repro scripts + full data attached under
ab_pr15756_vs_main/.1. Decode — this PR vs main baseline (same
cute_dsl_indexer_topk_decodeop, each version's default dispatch)The PR makes radix-filter single-pass multi-CTA the default and adds overflow policies. Speedup = baseline_us / PR_us.
Geomean: bf16 1.36x, fp32 1.11x.
Representative configs (top_k=2048):
Full decode data (60 configs)
2. Prefill — new DSL interface vs CUDA
indexer_topk_prefill(fp32,overflow_policy=REREADdefault)Prefill is new in this PR (no main baseline).
num_rows= packed query-token count (min(chunk_query_len, 2^31/num_cols), chunk capindexer_max_chunk_size=32768); the packed production workload is num_rows 8K–32K. cpp/dsl = cpp_us / dsl_us. cpp is fp32-only.Realistic regime (num_rows 8192–32768) geomean cpp/dsl = 1.25x.
Representative configs:
Full prefill fp32 data, num_rows>=8192 (realistic regime)
PR 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.