[None][perf] Emission-assisted GVR top-K decode for the DeepSeek V4 indexer - #16953
[None][perf] Emission-assisted GVR top-K decode for the DeepSeek V4 indexer#16953siyidNV wants to merge 105 commits into
Conversation
…ase-3 collect Port the block-skip consumer from the skip-finegrain development chain onto the R0 (op#26) architecture, measured-optimal configuration only (grain 32, int16 active list = 16KB SMEM, strided coalesced 3-barrier build, UN=2 software-pipelined compact scan). The active-block list is built once per row at the loosest rung (lossless for every rung count and the collect); phase3's compact stream-write replays the count pass's per-thread walk order so prefix-sum positions stay exact; a list-current flag pairs the two and is cleared on any dense fallback. Misaligned slice starts fall through to the dense path. Contract: workspace/epilogue_topk_interface.md. Correctness: REAL-data smoke (flash 16k/64k/256k/1024k + pro 64k/1024k, dense + skip arms, exactness contract) all pass incl. hit_rate=0.08. Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
The block-skip bounds tensor initially landed as a required positional in __call__, breaking every pre-existing compile path that does not pass it (16457's equivalence tests: 'Missing required argument'). Move it after stream with a None default so legacy callers are untouched; the wrapper passes it positionally last (the TVM-FFI env-stream launch takes no runtime stream arg). Verified standalone: main equivalence family 768 passed / 0 failed, launch_autoconfig 4/4, real-data smoke (flash+pro up to 1M) all green. Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
The active list previously required 32-aligned slice starts (runtime
guard falling back to dense), which silently disabled the skip on every
cluster slicing whose N/cs is not a multiple of 32 — including the
launch policy's cs=8 picks at the 512k/1024k rungs. The list now covers
FULL blocks only (first-full-block ceil in the build); the sub-block
head region of an unaligned slice is counted by all threads in a
strided scalar pass ordered BEFORE the list walk, and Phase 3's compact
write replays the same head-then-list per-thread order, so prefix-sum
positions stay exact. Boundary blocks shared with a neighbouring CTA
appear only in that CTA's head region — no double count, no gap.
Verified: skip arms exact at cs in {1,2,4,8} on flash+pro real cells
(64k/512k/1024k, unaligned N=262127 and 131075 slices).
Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
The active list was built at the loosest rung over all M columns. On
low-hit-rate rows the lowest sample-quantile rung sits far below the
K-th value (real pro 1024k: retained fraction 0.63 at that rung vs
0.08 at the final threshold), so the list barely skipped anything.
Build now iterates (cs==1): if the list exceeds CAP = 3/4 kC blocks,
DROP that rung — dropping is always correct (the rung is merely an
unmeasured probe; its partial counts are excluded from the admission
argmin and the fallback bracket seeding via a dropped-rung mask) — and
rebuild at the next tighter threshold, bounded by M-1 extra ~2-5us
builds. At cs>1 per-CTA list lengths differ (the drop decision would
diverge across the cluster), so the plain loosest-rung build stays.
Real-data 1024k, cs1, warm-L2 directional: pro 28.8 -> 16.4us (skip
ratio 1.16x -> 2.05x), flash 18.5 -> 12.8us (1.73x -> 2.49x).
Correctness: cs {1,2,4,8} x flash+pro x {64k,512k,1024k} all exact.
Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
…loads Review fixes for the compact machinery: - skip_ok now also requires nb_slice <= SKIP_MAX_BLOCKS and absolute block id < 32768 (int16 list entries); wider/higher slices fall back to the dense walk losslessly instead of silently truncating counts and the collect (or wrapping ids negative at cluster_size > 1). - Both compact walks vector-load only fully in-bounds chunks; the slice-end straddle goes through the scalar path, so an unaligned row no longer reads past the row/allocation. - Ctor rejects enable_block_skip without enable_r0 (dead 16KB SMEM). - emu_block_max defaults to records='positional' (the shipped kernel is grain 32; 'rotate' is a grain-128 fold fixture) and the wrapper asserts block_max covers every 32-position record of the row. Validated: capacity boundary (8192/8193/9375 blocks), N=1.05M at cs=4/8, unaligned exact-size rows (N=1000/65535/65529), real-data flash/pro 64k/1024k — all exact with planted tail winners. Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
The compact walk only wins on long rows (cold-L2 protocol: >= 2.18x at N=262k, 4-14% loss at N <= 131k). Gate block_max shape-based (no device sync) behind skip_min_n=200_000: below it the wrapper drops to the dense arms. Protocol after gating: flash/pro 256k/512k cells all 0.99-1.02x, 1024k wins intact (flash BS1 15.63us 2.18x, BS1024 4.18x; pro BS1024 3.15x). Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
pick_config gains has_block_max: with bounds available and N >= 200k the policy pins cluster_size = 1 - the compact list + rung tightening (cs1-only) beat the row-split configs outright once the bounds prune the scan (cold protocol, real data: BS1 1.21x, BS64 2.12x, BS1024 4.16x vs the stock picks; splitting shrinks each CTA slice below the skip break-even and disables tightening). The wrapper dispatch gains a second gate next to skip_min_n: K > 512 at num_rows < 8 keeps the stock path - the acceptance band is proportionally tighter (kC/K = 6 vs 10), the bounds prune less, and the row-split configs win (pro 262k BS1: skip 21.3-21.6us at cs1/cs8 vs stock cs8 19.7us). Cold protocol vs op26 at its own launch policy, 24 cells: zero regressions; flash 1024k 1.21/2.12/4.16x (BS 1/64/1024), pro 1024k 1.68/3.15x (BS 64/1024), everything below the gates identical to stock. Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
emu_seed_counts / emu_cand implement the A-side products per the epilogue<->topk buffer contract v2 so the consumer waterfall can be developed and tested against torch references before the fused indexer lands. Real-data coverage probe: 7/8 cells have a seed count inside [K, kC]; prev-kth drifts too loose at long context, so the L2 collect threshold should be the middle rung / xstate-adaptive. Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
use_ext_counts: rung thresholds AND their exact counts arrive from the
indexer epilogue (seed_thr/seed_counts [rows, 3], interface v2), so
P1b and the M-ary R0 count pass are skipped. The seeded refine routes
both cases: an in-band rung is re-measured once (building the
per-thread hand-off Phase 3 requires) and accepted; a full miss seeds
log-falsi with the external brackets. cs==1, requires fb_fix and a
3-slot rung config (the wrapper pins 2 qfracs + vseed; the values are
irrelevant since P1b never runs).
Real-data validation (flash/pro x 16k..1024k, thresholds {prev-kth,
q35, q85} + emu counts): 8/8 exact on stock/ext/ext+skip arms incl.
the pro-1M full-miss bracket cell. Directional: +9-10% at 1M (P1b +
M-count saved), small-N slightly negative (the waterfall routes those
to the L2 direct path instead). Next: skip P1 under ext (outer
brackets from xstate) and the L2 direct-to-P4 branch.
Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
With external epilogue counts the only surviving P1 products are the [v_lo, v_hi] outer bracket and the scalar-state init; the ext rungs provide the bracket directly (host contract: t_0 < t_2, finite, all rows valid) and tid0 initializes the scalars. A miss whose target falls outside [t_0, t_2] recovers through the refine loop's 8x bracket expansion, same as the stock fail-soft. Real data 8/8 exact unchanged; directional gains vs stock R0 improve to 1.15x at flash 256k / 1.14x at flash 1M (from 1.04x/1.09x with P1 still running). Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
v1 routed external-count admission through the dense seeded refine and forfeited the compact-walk win (flash 1M ext 34.7us vs skipR0 15.6us cold). v2 only skips P1b: the stock M-ary pass runs on the ext rungs, so list build, rung tightening, per-thread hand-off and classify compose unchanged. When an ext count is already in [K, kC] the admitted threshold is parked in ALL rung slots (v2b) — the M-ary pass degenerates to one compact single-threshold count and classify admits it; a miss keeps the distinct rungs as measured brackets. Real data 8/8 exact (stock/ext/ext+skip). Cold protocol: flash 64k 1.21x/1.25x (BS1/BS1024, the slim-admission cell); flash 1M ext+skip 16.1us vs skipR0 15.6us (composition recovered); pro-1M miss rows still pay multi-count+refine (0.7x) — routing sends those to stock skipR0 via xstate feedback (next step). Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
emit_xstate writes the per-row loop state (interface v2 layout [rows, 8]: [0] valid, [1] kth proxy, [2] accepted threshold, [3] cand_count from the pre-P4 snapshot — P4 repurposes the s_iscalars slots) at the cs==1 Phase-4 exit; degenerate identity rows write valid=0. The next step derives its seed rung group from these fields. Real-data validation: exactness unchanged; state fields exact (cand_count == count_ge(threshold): flash 991/633, pro 1854/2354); same-step reseed from the written state admits in-band with the exact count on ALL cells — including pro 1M, whose static rung group missed entirely (the temporal rung fixes the miss AND slims flash 1M admission 1290 -> 633). Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
use_ext_cand: epilogue-collected (value, index) pairs land straight in
smem_keys/vals via a cooperative sentinel-skipping SMEM-atomic load —
no P1, no counting, no Phase-3 scan. Eligibility (void == 0, claimed
<= cand_cap, collect rung count in [K, kC]) is a CTA-uniform register
predicate, so the dynamic skip of the P2/P3 slab stays convergent;
ineligible rows fall through to the ext-counts path unchanged.
Real data: 6 cells x {ext, l2, forced-void fallback} all exact. The
direct path makes top-k O(cand_count), independent of N: eligible rows
cost ~12.2us warm from 16k through 1M (flash 1M: 2.84x vs the ext
count path, below even the cold skipR0 15.6us). With the 0.89-0.97
chain in-band rates, ~90% of production rows hit this floor when the
epilogue emits cand.
Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
derive_seed_rungs places the next step's guard rungs a fixed number of count-OCTAVES from the previous accepted threshold, using the local slope of log2(count) vs threshold estimated from the previous step's own 3 rung measurements (log-linearity is the same property log-falsi exploits). Fixed spreads face a two-sided trap: too narrow misses drift, too wide puts the guard rungs themselves out of band — no single value wins both models (best fixed: pro 0.97/flash 0.92 vs flash-tuned 0.82/0.97). Real-chain kernel validation (V4-Pro/Flash multi-step captures): in-band admission pro 0.89 -> 0.99, flash 0.96, all steps exact. Combined with the L2 direct arm this routes ~97%+ of production rows to the O(cand_count) floor (cold protocol: flash 1M BS1024 30.9us = 12.1x, BS1 8.7us = 4.0x; pro 256k 1.5-1.8x). Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
The ext knobs were compile-time, so a row whose epilogue rungs all miss [K, kC] (or an xstate-invalid row, t_0 = +FLT_MAX) still paid the ext bracket-refine — measurably worse than stock (pro-1M cold: ext-miss 51us vs stock-skip 21us). Routing is now a per-row runtime predicate read from the ext counts themselves (CTA-uniform loads, so the dynamic branches with barriers inside stay convergent): in-band rows keep the ext fast path (skip P1 + P1b), miss/invalid rows run the full stock path (P1 + P1b + vseed + count) including the block-skip machinery. Warm validation: pro 1M ext+skip 43.5us (0.77x) -> 18.8us (1.80x); in-band cells unchanged (flash 1M 2.40x, pro 256k 1.07x); mixed-row chains exact with in-band 0.99/0.97 (pro/flash, adaptive rungs). Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
The BS=1 mid/long-row cells stayed on op26 because the waterfall fast paths were cs==1-only while op26's pick_config splits a single row across cs=4/8 CTAs. The pre-collected pairs are O(cand_count), so row splitting buys the direct path nothing: at cs > 1 the LEADER loads the pairs alone (take_cand is cluster-uniform - every CTA reads the same per-row control words) and peers publish zero local candidates for the DSMEM gather; ineligible/invalid rows fall through to the native stock path at op26's own cluster split. xstate writes at the leader's Phase-4 exit; the ext count pass composes with the existing cs>1 cluster merge unchanged. Validation: bl2 cells exact at cs=1/4/8 including forced-void fallback; cs1 smokes and the adaptive-rung chains unchanged (in-band 0.99). Cold protocol with the production arm (op26 launch config + ext inputs + in-kernel routing), vs op26 baseline: flash 256k BS1/64 1.24x/1.59x, flash 512k 1.33x/1.42x, flash 1M BS64 3.69x, pro 256k 1.20-1.73x - the former regression cells flip to wins; pro 512k/1M BS1 static-rung misses route to stock (adaptive xstate rungs take them direct in the closed loop, 1.4-1.6x steady-state). Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
Collect the pre-collected candidate list at the LOOSEST seed rung and admit it whenever it is complete (claimed <= K_max) and any rung counts >= K; the filter rung (count closest to K from above) is applied on the fly while loading the pairs, so P4 sees the thinnest covering set. kC leaves the admission vocabulary and remains only as the physical smem capacity guard. Correctness: C(t_lo) >= C(t_filt) >= K implies true top-K subset of list subset of filtered set; list truncation (claim order is value-blind) remains the only fatal case and falls back. - K_max = 24576, set by a four-chain search on real captures (incl. 320k/640k long decode): 16K->24K gains 8pp direct-hit rate, 24K->32K only 0.1pp (band-limited, not capacity-limited). - Loader: 4x-unrolled latency-overlapped walk with ballot-batched smem claims (loop exit must stay warp-uniform: ragged exits deadlock the warp collectives) and un-nested value loads. Device-level (nsys kern-sum, cold L2) on a 160k real chain: the naive walk ran 0.64x vs the block-skip arm; this form reaches 1.05x at full loosest-rung coverage (eligibility 1.00). - Straddle refine (cs=1): when no rung count lands in [K, kC] but the list is complete, one 256-bin histogram pass over the list finds an in-band edge and the filtered load proceeds; smem overflow demotes to the fallback. 640k chain: straddle steps 30 -> 14-16us, device mean 1.41x -> 1.73x vs block-skip. - Byte-parity routing keeps fat lists (2*claimed*cs >= N) on the fallback: measured both ways, walking them is slower at every cs. Validation (B200): four admission modes exact (direct/filter/refine/ fallback) at cs=1 and 18/18 exact at cs=1/4/8 on real V4 bundles; four real decode chains (160k/132k/320k/640k) all-step exact with wall ratios 0.97/1.00/1.04/1.21x and device-level 1.05x (160k) / 1.73x (640k) vs the block-skip arm. Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
Replace the rung-based in-list-filter admission with the count-only scheme: the candidate list is SoA (score column + position column, sentinel score -inf), collected at a single loose line, and admitted purely by entry count (K + 64 <= claimed <= K_max = 24576; the 64 is the emitter sentinel-pad bound, so the live count provably covers K). Rung admission, filter-line selection, straddle refine and the parity gate are all deleted - the seed-count columns are no longer consumed on the list path (the GEMM-side L1 pass becomes deletable, -3.2% emission tax). - THIN list (fits kC): every entry lands AT ITS LIST INDEX in the candidate buffers - no ballots, no smem atomics (128 serialized same-address atomics per trip measured ~1.1us/1k entries), no warp-uniform loop constraint. - FAT list: atomic-free copy of the score column into a dedicated 96KB smem region (sentinels sanitized to t_lo - 1), a zooming smem histogram (3 rounds, NBL^3 resolution - value-linear bins collapse on long-tailed logits) finds an edge whose exact count lands in [K, kC] (lands ~1030 for K=1024), survivors compact with one merged-ballot atomic per warp per trip. The vals slots carry LIST INDICES (no second cold gmem pass over the position column); a post-P4 repair swaps the K winners' positions with fully-parallel gathers. - Closed loop: xstate[1] publishes the exact k-th (output slot K-1 of the rank-ordered scatter), xstate[2] the ~3K-crossing anchor from the round-0 histogram. Host policy picks the anchor field per domain (tight k-th for short/stable rows, wide 3K edge for volatile long rows - the exact-k-th anchor alone shrinks the next down-guard target to 4K and slope noise then undershoots K, forcing ~26us fallbacks). GVR_P4_TAIL_DBG compiles per-phase clock64 stamps into the spare xstate slots. Validation (B200): 24/24 exact across cs=1/4/8 and the straddle- threshold suite on real V4 bundles; per-row cold device phases: thin walk 1.5-3us, fat stage+zoom+compact ~1.1us/1k entries, Phase 4 flat 5.5-6.5us. Real-chain device-level vs the block-skip arm: 640k 1.42x (fallback steps are C(t_lo) < K undershoots - a host anchor-policy matter), 160k 0.93x. Kernel-only chain means trade 5-20% vs the previous rung-based commit at B=8 in exchange for the interface collapse; the deleted L1 emission pass dominates E2E at large batch. Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
The emitter (indexer GEMM epilogue; emulated host-side in the bench
harness) now counts the two tighter lines while writing the SoA list -
two extra compares per EMITTED element only, against the full-row L1
pass this replaces - and the control words widen to {n0, void, n1,
n2}. The topk side enters with every count known and the whole list
path collapses to a scalar state machine:
- some line's count lands in the acceptance band [K, B*]: cut at the
TIGHTEST such line, ONE filtered gmem pass straight into the
candidate buffers (positions deferred as list indices; the position
column is gathered only for the K winners after Phase 4). Counts
and load predicates are the same comparison, so line cuts need no
overflow net at all.
- the band is straddled or overshot by every line: a zooming
histogram over the gmem list CLAMPED between the two known bracket
lines finds an in-band edge (narrow domain - no long-tail bin
collapse; the all-above case takes one max pass first).
- void, or n0 < K + 64 (the emitter sentinel bound, proving live
coverage of K): fallback.
The dedicated smem staging region is deleted (frees 96-128KB; the
kernel's smem drops back to the pre-list footprint), and B* / kC
become constructor knobs (accept_cap, kc_override) for the band
search. Closed loop publishes the exact k-th (rank-ordered output
slot K-1) and the loosest in-band line as the anchor.
Line placement is a searched host policy (derive_seed_lines_v4):
count targets (t0, t1, t2), grid-searched on real chains =
(4096, 3584, 1536) for short domains / (12288, 5120, 2048) for long;
physical kC stays 5120 (8192 measured no gain).
Validation (B200): five admission states each exercised exact
(hit-t2/hit-t1/bracketed-histogram/above-t2-histogram/fallback),
24/24 exact at cs=1/4/8; real-chain device-level vs the block-skip
arm: 160k 1.09x (first config of this lineage to beat the rung-based
1.05x), 640k 1.39x (residual: 2-3 volatile rows/step whose collection
count escapes any placement - a host anchor-policy iteration item).
Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
…d histogram Emitter writes the candidate list into three fixed segments (>=t2 / [t1,t2) / [t0,t1), caps B*/B*/rest, spill to the looser segment on overflow), so a line cut only ever reads the dense mapped prefix of the segments above it: the hit path becomes a pure copy (no value filter, no ballots, no atomics) and the histogram path walks mapped indices. When all three lines overshoot B*, the bracket segment's own prefix doubles as an unbiased sample: the histogram runs on it at the sample rate with 1.25x-scaled fire targets, and the exact post-load count net absorbs the sampling noise. Device-level cold-chain results vs the block-skip arm (B=1): flash 132k 1.58x, pro 160k 1.57x, pro 320k 1.54x, pro 640k 1.98x (fastest steps 8-14us). Exactness smokes pass for cs=1/4/8 including forced straddle/void routings. Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
…repair Sub-phase clock64 instrumentation (GVR_P4_SUB_DBG) showed the P4 rank-scatter core costs only ~0.6us/k candidates; the chain-observed ~1.9us/k came from the exact-tail boundary repair: the tiny-tie fast path ran an O(need x class) serial select on thread0 (~10us on real rows with need ~100 x class ~100), and bigger classes re-scanned every candidate per radix level behind ~20 block barriers. The repair is now: (1) a block-wide pure-tie check over the straddle class (bit-equal class needs no repair at all - the scatter's arrival fill is already value-set exact); (2) mixed classes are compacted IN PLACE into smem_keys/vals[0..class) with a register-buffered two-phase pass (warp-aggregated slot claims), so every later step scales with the class, never the candidate count; (3) class <= 128 takes an exact warp0 pairwise-rank rewrite, larger classes a block-parallel 4-level MSB radix over the compacted pairs with a warp0 shuffle-scan digit search (3 block barriers per level instead of 5). The full-candidate radix fallback is gone from the fast-tail variant. Device-level cold chains, 640k B=1: step mean 13.7 -> 12.6us (1.97x -> 2.13x vs block-skip; the previously slowest window improves 17.35 -> 12.6us as five 19-21us serial-repair victims drop to 9.3-10.8us); 640k B=8 19.5us (1.60x). Exactness: 18/18 microbench cells including forced tie/outlier stressors, smokes cs=1/4/8 plus forced straddle/void routings all bit-exact. Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
New self_scan mode: the kernel itself streams the row ONCE against the
three closed-loop seed lines and buckets candidates on the fly - no
external emitter, no indexer-side changes, no gmem candidate values.
One CTA per row, four phases: (0) scan-bucket - VALUES land in on-chip
segments (A/B/C at bases 0/B*/2B*, values-only 4B/entry, spill to the
looser segment, cursor totals ARE the line counts), POSITIONS stream to
a write-only gmem column reusing the cand_idx slot; (1) the v5 cut
state machine unchanged (a line cut compacts winning segment runs to
the smem prefix and fills smem_vals with segment coordinates, so P4,
the tail repair and the deferred K-gather run verbatim); ineligible
rows take the stock in-kernel fallback.
Scan-loop lessons baked in (each measured): per-element warp ballots
serialize every load (~1.8us/k); 16-wide register lists spill at 1024
threads (64 regs/thread ceiling) - values re-read from the load
fragments, positions derived arithmetically, classes recomputed;
warp-collective claim prefixes cap in-flight loads at 2/warp (ncu:
0.19% memory throughput) - final form claims passers with per-element
smem atomics, which do not synchronize the warp and hide under the
read stream (0.13us/k comp).
Exactness: 25-cell REPORT-S4 dataset x B in {1,2,4,8} = 100/100
bit-exact (flash/pro/v32 incl. K=2048, tiny-N and straddle fallbacks).
Perf vs PR16457 tip (same node, cold kernel-sum): geomean 0.64-0.71x,
short rows 0.8-1.05x, long rows 0.4-0.8x - the single-CTA read wall by
design; stage 2 (block-max skip) attacks the read itself.
Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
…gle band
Phase 0 gains a block-skip variant (enable_block_skip + self_scan):
per-32-position maxima from the GEMM tail gate whole blocks out of the
scan. Measured design pivot: skipping against the LOOSE collection
line can never pay (n0/N ~5-12% density -> ~80-99% of blocks contain a
passer; benched 0.16-0.39x), so the skip mode collects a SINGLE BAND
against the TIGHTEST line (density 0.4-0.8% -> 12-22% pass): only
segment A fills, the cursor keeps exact attempt counts, and the v5
state machine runs unchanged fed n0 == n1 == n2 - a cut lands on t2
(common), the sample-hist path absorbs over-B* rows (the A prefix
stays a value-blind sample), under-K rows take the stock fallback.
The small-batch block_max gate in the wrapper is bypassed for
self_scan (stage 2 owns its own skip economics).
Exactness: 25-cell REPORT-S4 x B in {1,2,4,8} = 100/100 bit-exact,
plus forced under-K fallback cells. Perf state (B=1 vs PR16457 tip,
same node): long rows improve markedly over the dense scan (flash
512k 33.9 -> 24.5us = 0.87x of tip; 1024k 46.7 -> 38.7; pro 1M 52 ->
44) while short/mid rows should route to the dense scan (host picks
by expected block pass rate). Known remaining work, measured and
documented: the block loop is still latency-bound on the bmax stream
(8 scalar loads/warp round); a lane-per-block + ballot variant was
tried and loses at high pass rates - loop shape per density regime is
the open optimization, along with a t2-only closed-loop line-derive
for chains.
Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
The skip scan is restructured into two passes that both run at the tuned dense-loop shape: (1) DENSE-vector-scan the block-max array itself (1/32 of the row, 128-bit vectors - the bmax row base is only 16B aligned) and compact the PASSING BLOCK IDS into the idle C segment (single-band mode never fills C; ids store exactly as floats); (2) walk the compact list, eight listed blocks per warp round issued back-to-back - every element read is useful and the loads pipeline. A list overflow (pass rate too high for skipping to ever pay) falls back to a dense full scan of the row inside the same phase. This removes both latency walls the one-pass shapes hit (8-scalar bmax rounds; serial per-block walks): flash 512k drops 25 -> 20us and BEATS the PR16457 tip (1.06x) - first cell where the fused self-contained kernel wins outright; flash 1024k 47 -> 26us (0.72x of tip), pro 1M 53 -> 36us, v32 128k+ 28us. Dense/skip best-of geomean 0.69 -> 0.73-0.75x across the 25-cell REPORT-S4 dataset, all 100 cells bit-exact. Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
…ases The pass-2 gather interleaved each block load with its smem atomic claims; atomics are memory-ordered, so the compiler could not overlap the next block load and the eight-block round degenerated into a serial latency chain (phase-0 stamp: 17.7us at flash-1024k against a ~5us budget). Loading all eight listed blocks into registers first and claiming afterwards restores the in-flight parallelism: phase 0 drops to 10.4us and the 25-cell table moves decisively - flash 512k 1.39x over the PR16457 tip, 1024k parity (19us), 256k 0.93x; v32 64k parity, 128k+ 0.90-0.92x; pro 1M 0.85x. Dense/skip best-of geomean 0.73 -> 0.84-0.86x, still 100/100 bit-exact. Remaining gap concentrates in the mid-row dense regime (64-128k, 0.63-0.72x), where the dense scan's single-CTA latency wall stands (cp.async staging is the known next lever). Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
Preload the next round's two vectors into shadow fragments before the current round's atomic claims (the pass-2 lesson applied to the dense loop). Measured neutral-to-slightly-positive (phase-0 38.7 -> 37.3us at flash-1024k): unlike pass 2 the dense loop's wall is not the cross-round atomic ordering - documented for the record; the next dense-lane lever is cp.async/smem staging. Final 25-cell state (dense/skip best-of vs PR16457 tip, B=1..8 geomean 0.84-0.86x, 100/100 bit-exact): flash 512k 1.39x / 1024k 1.00x / 256k 0.93x; v32 64k 1.00x / 128k+ 0.90x; pro 1M 0.85x; remaining gap concentrated at the 64-128k dense regime. Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
Replace the register-preload dense scan with an LDGSTS staging pipeline: each thread streams one 16B vector per step into a private slot-major smem slot (no data registers, no scoreboard stall until the wait), keeping stage_slots rounds in flight. The staging buffer aliases smem_vals - written only after phase 0, with every non-empty cp.async group drained inside the loop - so depth 2 costs zero smem; trimming cap_c to <= 16384 frees 32KB of keys for depth 4. flash-1024k phase0 37.3 -> 34.8us; exactness unchanged (fused and skip smoke 6/6, 25-cell real-data sweep 50/50 bit-exact). Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
The short-row 512-thread heuristic is tuned for the stock multi-pass
kernel; under self_scan it silently halved the warp count of every
N_dec < 65536 cell and cost ~5us/cell in the phase-0 scan (flash-128k
p0 14.4 -> 9.4us at 1024 threads). Route self_scan to 1024 threads
unconditionally.
25-cell x B{1,2,4,8} same-node sweep vs PR16457 tip: best-of geomean
0.838 -> 0.864 (B8 0.881), 100/100 bit-exact.
Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
The dense scan is instruction-issue bound (pcsamp: no_instructions +
fixed-latency wait dominate; long_scoreboard is 6%), so each pipeline
step now processes two 16B vectors per thread - loop, wait, commit and
address arithmetic amortize over 8 elements while the in-flight byte
count stays at 2 pairs x 32B across the 4 staging slots.
The pair shape needs all 4 slot rows, and the 64KB staging fits the
CTA budget only with the C segment trimmed, so self_scan now defaults
cap_c to 16384 and rejects anything larger (validated bit-exact across
the 25-cell x B{1,2,4,8} sweep).
flash p0 (warm, 1024 threads): 512k 17.8 -> 16.8us, 1024k 33.5 ->
32.4us; smoke 6/6 exact.
Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
…in-kernel) New ext_rungs mode: the host supplies only the three closed-loop rung THRESHOLDS (previous-step xstep lines); the kernel counts them itself through the stock R0 multi-count pass and admits the tightest rung with count in [K, kC], then collects and refines as usual. This is the fully self-contained two-pass shape: no emission of any kind, pass 1 = one fused 3-rung count (cluster-merge and block-skip compose unchanged), pass 2 = the stock single-line collect. Versus use_ext_counts (variant A) the only delta is where the counts come from; P1's preIdx gather and the P1b quantile rung derivation are both skipped (the seed lines carry the bracket). Smoke: 15/15 bit-exact across cs=1/4, block_max skip, and the thin (all rungs below K) and fat (all counts above kC) miss paths. Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Solid, well-tested work — the opt-in gating, runtime validity guards (exactness never rides on hint quality), and the degenerate-preIdx test battery are exactly what a change of this risk class needs. A few PR-level items:
-
PR size / splitting: +7.8k lines across kernels, host wiring, ops, and tests is past what anyone can meaningfully line-review in one pass — I spot-checked the CuTe DSL kernels and leaned on the tests rather than reading all ~5k kernel lines. Unless this is urgent, please break it into a stack of smaller PRs, each independently reviewable and CI-validated. A natural split: (a) the degenerate-preIdx gather fix on its own (it is a default-path bug fix and should be bisectable), (b) the
radix_lenschange from item 3 below, (c) the FP4 epilogue emission plumbing (block maxima → packed seed counts → bucketed candidate list, possibly one PR per tier) plus its contract tests, (d) the GVR consumption tiers plusgvr_ext.py/gvr_routing.pyhost wiring and theTRTLLM_GVR_EXTflag, (e) the CI test-list surgery. Stacking them lets CI validate each dependent layer ahead of the merge and keeps a regression attributable to one layer. If this has to land as one PR, say so and I will review it as-is — but the kernel internals will remain covered by tests rather than by line review. -
Ticket: a ~7.8k-line perf feature with new kernels, a new env var, and a default-path fix should carry a TRTLLM JIRA, not
[None]. Please file one and update the title. -
Undescribed default-path change: the
radix_lensselection change in[tensorrt_llm/_torch/attention_backend/sparse/dsa/indexer.py:1881](https://github.com/NVIDIA/TensorRT-LLM/pull/16953/files#diff-ad9a1b99228c59de01085466215ddf943ce4382c8ee8934f0c82c42f2de966a5R1881)-1893alters an existing non-ext branch (cute_dsl_indexer_topk_decodewithcompress_ratio > 1), but the description says default-path behavior is unchanged except the degenerate-preIdx fix. If this is a latent-bug fix (wrong/2-D lens previously passed on the DSL path), say so explicitly in the description — or split it into its own PR so it's bisectable. Also, the new comment claimskv_lens_cuda_2dis zero on the DSL path, butdsa/metadata.py:354fills it with the compressed lens on the non-expanded path — one of the two needs correcting. -
torch.compile hazard: both ops now write optional tensors (
seed_thr,xstate,cand_*,block_max) that stay out ofmutates_args(the code comments explain the torch.library limitation). Under functionalization those writes are invisible, so the closed-loop feedback (xstate → next step's lines) would silently read stale data if these ops ever run under torch.compile. Worth stating the eager/CUDA-graph-only constraint in the op docstrings, and re-checking whether the pinned torch version still has the Optional-mutates-args IndexError. -
TRTLLM_GVR_EXTis undocumented outside the code. Fine for an experimental flag, but at least the docstring ingvr_ext.pyshould state the env var name so it's greppable from both ends.
The host-side contracts (packed seed row layout, segment geometry, shape asserts) all check out consistent across the three files that duplicate them. The inline comments below apply regardless of how the change is split.
There was a problem hiding this comment.
Did a deep pass over the emission/consumption/wiring layers with one question in mind: the perf protocol builds batches by row replication, so what survives when rows are independent, as in real serving? The data path itself holds up well under that question: buffers, guards, admission and fallback are all row-indexed with no cross-row state, every validity guard demotes to an exact path, the counts tier re-measures its admitted line in-kernel, and block-skip consumes current-step maxima so it doesn't depend on feedback quality. The inline comments carry the specific findings (span collapse in the shipped updater, a reachable list-tier exactness gap on the line-cut branch, ext_rungs unreachable in production, static-N routing, positional-state consequences). Two review-level points:
-
Row replication measures the per-row hit path, not batch behavior. All fast/slow forks are per-row runtime decisions inside a max-over-rows launch (one CTA/cluster per row, no work stealing), and with cs=1 every B <= 148 cell is a single wave. Replication makes all forks bit-identical across rows, so a reported cell speedup S is exactly t_stock/t_hit for one row. For independent rows with per-row hit probability p and relative miss cost r, the step speedup is about 1 / (p^B/S + (1-p^B)*r): at S=1.75, r=1 that gives 1.54 at B=4/p=0.95 and 1.09 at B=16/p=0.9, and the 1M/B=256 corner needs per-row miss < 1e-4 to keep 85% of its value, while first decode steps alone contribute ~1/OSL misses per row-step, before churn or drift. The miss sources replication can't sample are structural: first ext step of every request, churned stale-valid state, MTP (next_n>1) aging, static-N routing. A cheap decisive experiment on the existing setup: force a fraction f in {0, 1/B, 2/B, 4/B} of rows cold in otherwise-replicated batches (zero their xstate) and read back the per-row xstate[:,0] / list_used the kernel already publishes. One sweep maps the collapse curve, and a chained independent-rows run through update_seed_rows then measures p directly.
-
Provenance. The 486-cell driver isn't in the PR, and the only line-derivation code here that produces healthy spans is the harness's (derive_seed_lines_v4), while the shipped update_seed_rows is span-degenerate (inline). Publishing the driver, or a committed repro script, and re-scoring a subset with the shipped closed loop on independent rows would make the headline numbers defensible as production numbers rather than per-row upper bounds.
None of this blocks the mechanism. The routed, tax-inclusive 1.296/1.387 with worst cell 1.000 is a strong result on its own; it doesn't need the unrouted corners to carry the PR.
| x = self.xstate[:num_rows] | ||
| kth = x[:, 1] | ||
| anch = torch.maximum(x[:, 2], kth + 1e-5) | ||
| span = (anch - kth).clamp_min(1e-4) |
There was a problem hiding this comment.
This span can never leave its 1e-4 floor. The kernel publish writes xstate[1]/xstate[2] with x[2] - x[1] <= 0 on every path: non-list rows at cs==1 set thr_pub = anch_pub = s_thr[0] (gvr_topk_decode.py L7561-7562), the cs>1 leader exit writes s_thr[0] into both slots (L7705-7706), and list rows upgrade thr_pub to the exact k-th value while anch_pub stays at the loose collect line below it (L7563-7571). So torch.maximum(x[:,2], kth + 1e-5) - kth <= 1e-5 always, and span == 1e-4 for every row on every step. GUARD_LO/GUARD_HI never see a real spread, and the three lines are pinned to prev_kth +/- 2e-4. For the list tier that puts t0 only 2e-4 below the exact previous k-th, so any upward k-th drift beyond 2e-4 drops the true count(>= t0) below K; see the line-cut comment for why that matters beyond hit rate.
The derivations that do produce usable spans (derive_seed_rungs / derive_seed_lines_v4 in tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py) are never called from this class, and no test chains update_seed_rows through a closed-loop step, so the perf grid's line quality may not reflect this updater. Publishing a real spread (the value at rank K-delta next to the k-th, or the rung counts, which today only go out under ext_rungs and production never enables that mode, see the comment below) and a 2-3 step chained test (kernel -> xstate -> update_seed_rows -> next-step admission) would cover it.
There was a problem hiding this comment.
You read the publish path right — thr_pub == anch_pub on every non-list exit, so the span never left the 1e-4 floor and the lines were pinned to prev-kth ± 2e-4. update_seed_rows is rewritten to slope-fit log2(count) vs threshold from the previous step's (lines, counts) — the same construction as the harness's derive_seed_lines_v4 — placing lines at K-relative target counts (8K/5K/2K); the list tier two-point-fits through (t0_prev, n0) and the published exact k-th, targeting n0 = 2.5K inside the admission band. Rows without a usable fit (cold start, tier flip, non-monotone counts) fall back to multiplicative guards ~9% of |kth| — wide, not the old 2e-4 pin. Rung counts come from the kernel's xstate[4..6] publish now that ext_rungs is reachable (next comment). New chained test ext_closed_loop runs 3 steps with the k-th drifting ±0.05 between steps through both the list and counts tiers. (fixed in 0620720)
| usable = cutlass.Int32(0) | ||
| if ( | ||
| void_c == cutlass.Int32(0) | ||
| and claimed_c >= cutlass.Int32(self.top_k + 64) |
There was a problem hiding this comment.
On the list tier the parked tight lines make n1_c = n2_c = 0, so the only line-cut that can fire is the third branch at L6455, which takes cut_n = claimed_c. claimed_c is ctl n0: exact A/B placements plus C window claims including sentinel pads. The epilogue claims hits + CAND_WIN(=8) per warp-window and sentinel-fills the tails, so the pad count scales with the number of active (CTA, warp) windows on the row, potentially hundreds on long rows. The top_k + 64 slack in this check doesn't bound that; the unit test's pads mode uses exactly 64, which matches the slack, but the epilogue contract doesn't. The LINE-cut branch then copies the prefix without filtering ("every read is a winner candidate", L6652-6654), and the exact post-load net (cnt_l < top_k or cnt_l > kC -> demote, L6777-6778) is only compiled under line_cut == 0.
So a row whose true count(>= t0) is below K, with pads keeping claimed_c inside [K+64, bstar], rank-selects over fewer than K real candidates and pads the output tail with -FLT_MAX / -1 indices (the xstate publish already anticipates -1 at slot K-1, L7568-7570). And with the shipped updater pinning t0 at 2e-4 below the previous exact k-th (see the update_seed_rows comment), count-below-K is just "the k-th rose this step", not a rare corner.
The prefix copy already touches every value, so counting non-sentinel entries during the walk and demoting on < K looks close to free; or the epilogue could publish a pad-free n0, since it knows its per-window pad counts. A 2-step chained test with the step-2 k-th drifted up by more than 2e-4 would pin this down either way. If I've missed a guard on this path, please point me to it.
There was a problem hiding this comment.
Confirmed — this was a real exactness gap, and reachable exactly as you described once the shipped updater's degenerate lines are in the picture. Fixed in the line-cut walk itself: the mapped-prefix copy now counts non-sentinel entries in flight (idx >= 0; one register add per element, one warp-reduce + smem atomic per warp at the end) and demotes to the stock path when the real count is below K — same recovery net as the histogram-cut branch. Regression test ext_list[starved] builds a row with 400 real candidates and enough sentinel pads to lift the claim into [K+64, B*]: the old kernel emits -1 tail slots, the fixed one falls back and stays exact. The chained ext_closed_loop test also crosses this path with a drifting k-th. (fixed in 0620720)
| if route.tier in ("counts", "list"): | ||
| kw["seed_thr"] = self.seed_row[:num_rows] | ||
| # rungs tier: pass no seed at all (a [rows, 3] column view of the | ||
| # packed row is non-contiguous and would trip the runner's assert) |
There was a problem hiding this comment.
topk_ext_kwargs never passes seed_thr for the rungs tier (this comment), and the op derives ext_rungs = seed_thr is not None and seed_thr.shape[1] == 3 (cute_dsl_custom_ops.py L7753), so ext_rungs is always False in the production wiring and the vb/rungs assist path is unreachable. What ships for that tier is cluster-size / num_threads selection only. If that's intentional, the vb rows in the perf report (geomean 1.032 / 1.107) were measured with a contiguous [rows, 3] seed the production path can't produce, and the report should say the shipped rungs tier runs assist-less. If it's not, update_seed_rows could materialize a contiguous [rows, 3] copy for this tier; 3 floats per row per step is noise next to the emission buffers.
There was a problem hiding this comment.
Correct on both counts. Fixed the wiring rather than the report: GvrEmissionState now keeps a persistent contiguous [rows, 3] mirror of the lines and topk_ext_kwargs passes it for the rungs tier, so the op's ext_rungs variant is what production ships; its in-kernel rung counts (xstate[4..6] publish) also feed the new slope-fit updater. The 3-floats-per-row-per-step copy rides the same update kernel. (fixed in 0620720)
|
|
||
| # indexer_max_seq_len is already the compressed length | ||
| # (get_indexer_max_seq_len divides); do not divide again. | ||
| n_comp = indexer_max_seq_len |
There was a problem hiding this comment.
n_comp here is kv_cache_manager.max_seq_len // cr, constant for the engine's lifetime, so neither per-row nor per-batch-max lengths ever reach plan_emission. Two consequences. First, the escapes fitted into plan_emission (n_comp < ASSIST_MIN_N_COMP, and the 48k-64k weak band that hands back to stock) can never fire on actual row lengths in a long-max engine, so short rows get routed onto assist machinery the planner would refuse at their true length. Second, a long-max engine at B within the list gate always plans list emission and pays its epilogue tax even when every live row is 2k. The routing constants were fitted on a grid where routing-N == row-N by construction, so these boundaries were never exercised in the mispriced regime. In eager mode, routing on the batch's actual max kv_len would recover most of this; under CUDA-graph capture the tier decision is baked into the graph, so actual-length routing would need length-bucketed graphs. If that's out of scope, a doc note that tier choice is engine-max-based would at least set expectations.
There was a problem hiding this comment.
All true, and worth stating plainly: every N the routing sees is max_seq_len // cr, engine-static. We've added that caveat to the PR description (with the mispriced-regime consequences you list), and the exactness story under mispricing stays the in-kernel guards. Actual-length routing needs either eager-only gating or length-bucketed graphs — we've scoped that as a follow-up PR (the dispatch-guard follow-up the kernel-side comment already references) rather than growing this one.
| max_seq_len=indexer_max_seq_len, | ||
| **ext_kw, | ||
| ) | ||
| st.prev_topk[:num_gen_tokens].copy_(out_slice) |
There was a problem hiding this comment.
Building on @brnguyen2's slot-churn point, two consequences of the positional keying that deserve to be spelled out:
-
Churned state is stale-valid, not cold. After a compaction/reuse event, xstate[i] col 0 is still 1.0 from the departed request, so
update_seed_rowsproduces confidently-wrong finite lines instead of the +inf cold-start lines the validity guard is designed to catch. Exactness survives on the counts tier via the in-kernel re-measurement; on the list tier, wrong-but-finite lines are exactly the corruption precondition (see the line-cut comment). -
Every request's first ext step is a guaranteed assist miss, and the baseline's warm start is bypassed. The stock path seeds
heuristic_prev_topkfrom the sequence's own last context token at the prefill->decode handoff (L1584-1599);st.prev_topk/st.xstatereceive nothing there, so step 1 runs the stock pipeline plus the p1r rescue scan, and the assist only warms at step 2. Under continuous batching that's a recurring per-request cost a fixed-composition benchmark can't sample.
Keying the ext state by the slot conventions heuristic_prev_topk already has, plus invalidation on finish and a prefill handoff seed, would address both.
There was a problem hiding this comment.
Both consequences were real. (a) is fixed at the prefill->decode handoff: the emission state is now seeded at the same point as heuristic_prev_topk (prev_topk from the last context token, xstate rows zeroed -> validity guard -> exact stock first step), so a request's first ext step is a clean cold start instead of inheriting the departed request's stale-valid state, and the warm start is no longer bypassed. (b) full slot-keying: we searched for a churn/compaction signal to key on — the heuristic path itself is purely positional and none exists — so beyond the handoff seed + zero, the positional-identity assumption and why exactness survives churn (counts re-measured in-kernel, validity/admission guards) is now stated in code where the state is consumed. Proper slot identity needs metadata/executor surgery; follow-up scope. (fixed in 0620720)
| # Emission-cost model: counts emission is a batch-only latency chain | ||
| # (hides at large batch); list emission grows with batch and context. | ||
| LIST_EMIT_MIN_N = 65536 # shorter rows: the emission outweighs the saving | ||
| LIST_EMIT_MAX_B = 4 # past four rows the list stops repaying its emission |
There was a problem hiding this comment.
LIST_EMIT_MAX_B = 4 here, but the PR description says the list tier routes at N >= 64k, B <= 16, and the wf table's B >= 32 columns (including the 12.2x corner) are outside the shipped routing either way. Could the description mark which wf columns are reachable through plan_emission, and lead with the routed, tax-inclusive grid numbers (1.296 vs reference / 1.387 vs stock, worst 1.000, from the 08-03 re-score)? Those describe what actually ships.
There was a problem hiding this comment.
Fixed the description: B <= 4 (LIST_EMIT_MAX_B) everywhere, the wf table is now explicitly labeled as the UNROUTED kernel-capability grid with the reachable columns marked, and the summary leads with the routed, tax-inclusive numbers (1.296 vs #16457 / 1.387 vs stock, worst cell 1.000) plus the per-step full-grid stats.
| self.enable_block_skip = bool(enable_block_skip) | ||
| self.SKIP_BLOCK = 32 | ||
| self.SKIP_BLOCK_LOG2 = 5 | ||
| self.SKIP_MAX_BLOCKS = 8192 # covers N up to 262144 at grain 32 |
There was a problem hiding this comment.
SKIP_MAX_BLOCKS = 8192 bounds block-skip to N_local <= 262144 (grain 32), and pick_config forces cs=1 when block_max is attached at n >= 200k, so at 512k/1M the fallback paths run dense and the block_max emission is pure tax at those shapes: the wf hit path never scans, and the miss path can't use the skip. The emission contract chains emit_cand -> emit_seed_counts -> emit_block_meta (fp4_paged_mqa_logits.py L545-559), so the tax can't simply be gated off on the consume side either. Raising SKIP_MAX_BLOCKS would let long rows skip too; if the smem active-list budget is the constraint, a comment saying so, plus a perf-report note that the 512k/1M rows owe nothing to block-skip (the wins there are pure list/counts effects), would keep the attribution straight for future readers. Separate nit: the cand_ctl signature comments at L5402/L5478 still say [numRows, 2] but the bucketed contract (and the GvrExtState allocation) is 4 ints/row.
There was a problem hiding this comment.
Comment added at the constant: SKIP_MAX_BLOCKS is the smem active-list budget (32KB of local ids), bounding block-skip to N_local <= 262144 at grain 32 — longer rows run the dense fallback and owe nothing to block-skip. The PR description now carries the same attribution note for the 512k/1M rows. The stale [numRows, 2] cand_ctl signature comments at 5402/5478 are fixed to the 4-int contract. Raising SKIP_MAX_BLOCKS trades smem against the skip range; left as is for this PR. (fixed in 0620720)
|
/bot run --disable-fail-fast |
|
/bot run |
|
PR_Github #65272 [ run ] triggered by Bot. Commit: |
Review round 2026-08-11 (brnguyen2, longcheng-nv), all inline findings: - list tier line-cut: the mapped-prefix copy re-measures the non-sentinel candidate count in flight and demotes to the stock path below K - a pad-inflated claim can no longer admit a starved list. Regression test ext_list[starved]. - update_seed_rows: slope-fit of log2(count) vs threshold from the previous step's lines+counts (harness derive_seed_lines_v4 construction) at K-relative target counts; list tier two-point-fits through the published exact k-th; no-fit rows take multiplicative guards (~9% of kth) instead of the degenerate 1e-4 span. Chained test ext_closed_loop (3 steps, kth drift 0.05, all three tiers). - rungs tier reachable: persistent contiguous [rows, 3] seed mirror passed for the rungs route; counts telemetry from the kernel's xstate[4..6] publish. - one shared per-step gate (batch <= 256, next_n == 1) for planning, emission and consumption; batch > 256 pays no emission tax and a stale route cannot be consumed. - prefill->decode handoff seeds prev_topk and zeroes xstate for new requests (same slot convention as heuristic_prev_topk); positional identity assumption documented at the consumption site. - seed width assert == 3 or == 8 with matching use_ext_counts derivation; ext tiers require enable_r0; list-capacity ctor reject; trace-time flag/tensor contract errors; block-skip single-band void contract; cand_ctl width comments; SKIP_MAX_BLOCKS budget comment; radix_lens comment reconciled with metadata (declared latent-bug fix); mutates_args limitation re-verified on the pinned torch (IndexError when a declared-mutable Optional is None at call) and documented precisely in both op docstrings. - rename per review naming feedback: TRTLLM_GVR_EXT -> TRTLLM_GVR_EMISSION, gvr_ext.py -> gvr_emission.py, GvrExtState -> GvrEmissionState. Validation on B200: full decode unit file 730 passed / 0 failed / 144 skipped / 1 xfailed; new ext_list[starved] and 3-tier ext_closed_loop chained tests pass; mutates_args declaration attempt reverted after uniform stock-path failures (None-case IndexError). Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
Per review: the driver behind the PR's performance numbers, in-tree. Per-step paired protocol (one NVTX range per arm/model/isl/B/layer, cold reps cycle all decode steps), nsys wrapper, per-step extraction from the sqlite, pairing + production-routing aggregation into the B x N mean/min tables. Capture root and output dir parameterized via GVR_CAP_ROOT / GVR_BENCH_OUT. Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
|
@brnguyen2 Thanks for the thorough pass — all five PR-level items addressed or answered:
Inline comments each carry their own reply + fix commit. New tests this revision: ext_list[starved], ext_closed_loop (3-step chained closed loop), on top of the ext tier batteries. |
|
@longcheng-nv The inline findings were all real — thank you, this was the most productive review round this PR has had. Each carries a reply + fix commit; the two review-level points:
|
|
CodeRabbit triage for this revision (18 actionable):
|
|
/bot run |
|
PR_Github #65279 [ run ] triggered by Bot. Commit: |
|
PR_Github #65272 [ run ] completed with state |
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
Two whole-PR items:
Ticket: a feature of this size (new emission pipeline, two kernel families, production routing) should carry a TRTLLM JIRA rather than [None] — the perf-grid provenance and routing thresholds will need a home when they get retuned.
Coverage of the host wiring: the unit tests cover the op contracts and GvrEmissionState well, but the closed-loop test emulates emission on host and nothing exercises the indexer.py wiring with TRTLLM_GVR_EMISSION=1 — prefill→decode handoff, batch churn, tier flips, and the ensure_block_max warmup-before-capture requirement are only validated by manual bring-up. An accuracy test (e.g., a DeepSeek Lite variant with the flag set) would be a good follow-up before any default flip; same for documenting the env var somewhere other than the source.
The description is accurate to the diff, including the two declared default-path fixes — appreciated, both are independently bisectable.
| self.cand_ctl: Optional[torch.Tensor] = None | ||
| self.cand_cur: Optional[torch.Tensor] = None | ||
| if enable_list_tier: | ||
| self.cand_vals = torch.zeros((max_rows, LIST_WIDTH), dtype=torch.float32, device=device) |
There was a problem hiding this comment.
These candidate buffers are sized (max_rows, LIST_WIDTH) = max_rows × 40960 × 8 B for vals+idx (~0.33 MB/row), allocated per indexer layer. With max_num_sequences=256 and ~61 layers that is ~5 GB for a tier that plan_emission only ever picks at batch <= LIST_EMIT_MAX_B (= 4). Two cheap fixes: (1) size the cand buffers to min(max_rows, LIST_EMIT_MAX_B) rows — nothing indexes past [:batch] on the list path, though update_seed_rows's unconditional cand_ctl[:num_rows].zero_() at line 207 would need a matching clamp (it currently runs for every tier, e.g. counts at batch 256); (2) longer term, cand_* and block_max are produced and consumed within one layer's forward (only seed_row/xstate/prev_topk carry cross-step state), so the transient buffers could be shared across layers the way heuristic_prev_topk lives on metadata.
There was a problem hiding this comment.
Done in c4de505: cand_* now allocate min(max_rows, LIST_EMIT_MAX_B) rows (~5 GB -> ~80 MB in your example) and the per-step control-word zeroing is clamped to the buffer rows (it does run for every tier). The cross-layer sharing of the transient cand_*/block_max via metadata is noted for the follow-up PR - agreed only seed_row/xstate/prev_topk are genuinely per-layer cross-step state.
| st = self._ensure_gvr_emission(metadata, q_fp8.device) | ||
| # indexer_max_seq_len is already the compressed length | ||
| # (get_indexer_max_seq_len divides); do not divide again. | ||
| emit_tier, self._gvr_route = st.plan( |
There was a problem hiding this comment.
plan() receives indexer_max_seq_len, which is engine-static (kv_cache_manager.max_seq_len // compress_ratio), not the live batch's KV lengths — so tier selection routes on engine capacity, not on the rows actually being scored. On an engine built for 256k serving 4k requests at batch ≤ 4, every step picks the list tier and pays the L2 emission tax (+9-14% of the GEMM) where live-length routing would return "none" (your own table says stock wins below 2048 comp and in the 48-64k weak band). Exactness is unaffected, and static routing is forced under CUDA-graph capture, but the "no routed cell regresses" grid presumably had engine max == ISL. Worth (a) a comment here stating the assumption, and (b) considering a live host-known bound for eager steps — or confirming the mismatch case was measured.
There was a problem hiding this comment.
Comment added at the plan call in c4de505 stating the engine-static assumption and its consequence (short rows in a long-max engine run assist machinery the planner would refuse at their true length; exactness via in-kernel guards). The mismatch case was NOT in the measured grid (engine max == ISL by construction there) - measuring it and live host-known-bound routing for eager steps are scoped to the dispatch-guard follow-up, together with the cold-fraction sweep.
Review touch-ups (approved round): the wide candidate buffers only serve the list tier, which plan_emission picks at batch <= LIST_EMIT_MAX_B - allocate that many rows instead of max_rows (~0.33 MB/row/layer, ~5 GB -> ~80 MB at max_rows=256 x 61 layers); clamp the per-step control-word zeroing accordingly. Also state the engine-static routing-N assumption at the plan call site. Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
|
Thanks for the fast re-review! Both optional touch-ups are in (c4de505: cand buffer sizing by the routing cap + the static-routing assumption comment). JIRA is being routed via longcheng-nv and the title will be updated with the id. Re-triggering CI on the new head. |
|
/bot run |
|
PR_Github #65289 [ run ] triggered by Bot. Commit: |
|
PR_Github #65279 [ run ] completed with state |
|
PR_Github #65289 [ run ] completed with state
|
|
/bot run |
|
PR_Github #65301 [ run ] triggered by Bot. Commit: |
|
PR_Github #65301 [ run ] completed with state
|
|
CI status note: two consecutive infra-category failures on this head, zero test failures in either (last run: 2219 passed / 0 failed / 62 skipped). Per the failure analysis for pipeline 53076, the GB300/NVL72 Slurm cluster now rejects job-scoped |
5 participants
Summary
Emission-assisted GVR top-K decode for the DeepSeek V4 sparse-attention indexer: the FP4 indexer GEMM epilogue now emits selection hints (per-block maxima / packed seed-count rows / a bucketed candidate list) that the GVR top-K kernel consumes through new opt-in fast paths, replacing most of its threshold-search and full-row scan work.
Shipped-path numbers (production routing, emission tax included): geomean 1.296x vs the #16457 baseline / 1.387x vs this PR's own stock kernel, worst cell 1.000 (no routed cell regresses on mean). Per-step over the full grid (B 1..1024 x raw ISL 8k..256k, all layers x all decode steps of real captures, 102k paired steps): mean 1.44-1.8x per model, zero mean-regressing cells, per-step regressions 0.14% of steps (all at raw ISL <= 64k; 128k+ has none). Unrouted kernel-capability corners reach 12.2x (table below, marked for reachability). All measured cells exact.
What's in the change
Indexer emission epilogue (
fp4_paged_mqa_logits.py, +1450)[rows, 8]— three threshold lines +count(>= line)accumulated with fp32 atomics in the epilogue (<<1% tax).{n0, void, n1, n2}control words. Measured emission tax: L1 +0.2-4.4%, L2 +9-14% of the indexer GEMM, flat in batch.GVR top-K consumption tiers (
gvr_topk_decode.py, +3603)wf: known-counts admission over the bucketed list — the tightest in-band line is a pure scalar lookup, the hit path degenerates to a filtered prefix copy straight into P4 rank selection; histogram-over-list and full fallback below.va: seed-count rows replace the preIdx gather and P2 threshold search on hit.vb: closed-loop three-line rungs (zero-emission variant, fallback tier).radix_lensfor thecute_dsl_indexer_topk_decodebranch atcompress_ratio > 1: the op takes 1-D request-level lens, but the old code passed the 2-Dkv_lens_cuda_2dslice; on the FP4-DSL path the live compressed lens aregen_indexer_kv_lens_cuda_runtime. Declared here per review so it is bisectable.[0, K)instead of computing the top-K. Found during real-model bring-up (42/231 dumped rows wrong = 21 layers x 2 sequences, first decode step each). Newphase1r_data_reseedrebuilds the refine bracket from the row itself (restores thecount(>= v_lo) >= Kinvariant), keeping the identity shortcut only where it is provably exact (all-tied row orN <= K). Non-degenerate rows pay nothing.Host routing + production wiring (new
gvr_routing.py, newgvr_emission.py,dsa.py,cute_dsl_custom_ops.py)plan_emission/pick_config: (B, N)-based tier selection (candidate-list tier only where it is net-positive: N >= 64k, B <= 4).GvrEmissionState: emission buffer lifecycle, device-side seed-row updates (CUDA-graph safe), prev-topK feedback loop.TRTLLM_GVR_EMISSION=1and composes with the existinguse_cute_dsl_topkrouting from [None][feat] top-k: route decode to CuTe DSL GVR top-k in e2e #16420 — default-path behavior is unchanged except for the two declared default-path fixes above.Tests
test_cute_dsl_fp4_paged_mqa_logits.py(+707).test_cute_dsl_gvr_topk_decode.py.xfaildocumenting a pre-existing corner inherited from the current kernel (reproduces on the unmodified upstream kernel): when the k-th tie class alone exceeds the candidate capacity, the selected value multiset is still exact but the index list can contain duplicate/unwritten slots. Requires >kC bit-identical scores at the boundary; never observed on real captures.Performance report
Protocol: real DeepSeek V4 captures (V4-Flash 21 indexer layers, V4-Pro 30 layers), all usable decode steps per layer, batch = row replication, nsys cold-L2 kernel-only timing on B200. Baseline = GVR kernel at the #16457 tip (identical to what main carries today). 486 cells, every cell exact (tie-aware score-multiset check).
wfkernel capability grid (UNROUTED: forced list tier;plan_emissionreaches only the B <= 4 columns at N >= 64k — the B >= 8 columns document kernel headroom, not shipped behavior) — speedup vs baseline:Geomeans over the full 9x9 grid (all layers x all steps):
Numbers above are kernel-only; the emission tax (L1 +0.2-4.4%, L2 +9-14% of the indexer GEMM) is charged on the indexer side and is why routing only enables the list tier at N >= 64k, B <= 4 — net accounting stays positive everywhere routed (headline geomeans above are tax-inclusive).
Routing caveats stated explicitly:
plan_emissionseeskv_cache_manager.max_seq_len // cr, an engine-lifetime constant (CUDA-graph capture bakes the tier in); per-step actual-length routing needs length-bucketed graphs and is deferred to a follow-up. Short rows in a long-max engine therefore run assist machinery the planner would refuse at their true length; the in-kernel validity/admission guards keep that exact.SKIP_MAX_BLOCKS = 8192(smem active-list budget) bounds the skip walk to N_local <= 262144; wins at 512k/1M are pure list/counts effects.Reproduction: the grid driver (per-step paired protocol, per-arm launch code, aggregation) is committed under
tests/scripts/cute_dsl_kernels/top_k/in this PR.Correctness validation
TRTLLM_GVR_EMISSION=1, the production-path selections of every indexer layer across all captured decode steps (231 rows) are score-multiset identical totorch.topk. This acceptance run is also what exposed (and now guards, via the new unit battery) the degenerate-preIdx bug fixed here.Review revision (2026-08-11)
All inline findings from the 2026-08-11 review round are addressed in the follow-up commits:
claimed_ccan no longer admit a starved list. Regression test:ext_list[starved](400 real candidates, pads lift the claim into the admission band).update_seed_rowsnow slope-fits log2(count) vs threshold from the previous step's (lines, counts) — the same construction as the harness'sderive_seed_lines_v4— and places lines at K-relative target counts; the list tier two-point-fits through the published exact k-th. No-fit rows fall back to multiplicative guards (~9% of kth, vs the degenerate 2e-4 pin). Chained test:ext_closed_loop(3 steps, k-th drift 0.05 between steps, list + counts tiers).[rows, 3]seed buffer is passed for the rungs tier (the packed-row column view is non-contiguous); count telemetry comes from the kernel's rung-count publish.heuristic_prev_topk(prev_topk from the last context token, xstate zeroed -> validity guard -> exact stock first step). Positional-identity assumption and why exactness survives churn is now stated in code; full slot-keying (no churn signal exists today) is follow-up scope.== 3 or == 8, ext-tiers-require-enable_r0and list-capacity constructor rejects, trace-time flag/tensor contract errors, block-skip single-bandvoidcontract,cand_ctlwidth comments, themutates_argslimitation re-verified on the pinned torch (IndexError fires when a declared-mutable Optional arg is None at call time - i.e. on every hint-less call), so the eager/CUDA-graph-only contract is now stated precisely in both op docstrings;TRTLLM_GVR_EMISSIONdocumented.TRTLLM_GVR_EXT->TRTLLM_GVR_EMISSION,gvr_ext.py->gvr_emission.py,GvrExtState->GvrEmissionState(review naming feedback).🤖 Generated with Claude Code