Skip to content

[None][perf] Emission-assisted GVR top-K decode for the DeepSeek V4 indexer - #16953

Open
siyidNV wants to merge 82 commits into
NVIDIA:mainfrom
siyidNV:perf/gvr-emission-topk
Open

[None][perf] Emission-assisted GVR top-K decode for the DeepSeek V4 indexer#16953
siyidNV wants to merge 82 commits into
NVIDIA:mainfrom
siyidNV:perf/gvr-emission-topk

Conversation

@siyidNV

@siyidNV siyidNV commented Jul 28, 2026

Copy link
Copy Markdown
Collaborator

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. Kernel-only geomean speedup vs the GVR kernel currently in main (PR #16457 tip) is 1.75x (wf path) across all layers x all decode steps of real DeepSeek V4 captures, with up to 12.2x in the long-context / large-batch corner. All 486 measured cells remain exact.

What's in the change

Indexer emission epilogue (fp4_paged_mqa_logits.py, +1450)

  • P0: per-256-token-block maxima carried out of the GEMM epilogue (enables block skipping in top-K).
  • L1: packed seed rows [rows, 8] — three threshold lines + count(>= line) accumulated with fp32 atomics in the epilogue (<<1% tax).
  • L2: bucketed candidate list — three fixed SoA segments classified by the tightest passed line, exact ballot claiming (pad-free prefixes), overflow spill chain, {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).
  • Block skipping over emitted block maxima; P4 rank-scatter exact-tail repairs (three-tier boundary-class handling); 512-thread small-K configs; runtime line-validity guards so exactness never rides on host hint quality.
  • Fix (default-path behavior): a degenerate preIdx gather (duplicate/invalid indices — e.g. the zero-initialized first-step feedback buffer, or a reused batch slot whose stale indices all fall past the new row's length) used to hit a shortcut that emitted identity indices [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). New phase1r_data_reseed rebuilds the refine bracket from the row itself (restores the count(>= v_lo) >= K invariant), keeping the identity shortcut only where it is provably exact (all-tied row or N <= K). Non-degenerate rows pay nothing.

Host routing + production wiring (new gvr_routing.py, new gvr_ext.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 <= 16).
  • GvrExtState: emission buffer lifecycle, device-side seed-row updates (CUDA-graph safe), prev-topK feedback loop.
  • op faces extended (modes derived from tensor presence); the whole pipeline is gated behind TRTLLM_GVR_EXT=1 and composes with the existing use_cute_dsl_topk routing from [None][feat] top-k: route decode to CuTe DSL GVR top-k in e2e #16420default-path behavior is unchanged except for the degenerate-preIdx fix above.

Tests

  • Emission contract unit tests (packed/bucketed, segment invariants) — test_cute_dsl_fp4_paged_mqa_logits.py (+707).
  • Degenerate-preIdx battery (37 cells: zero/dup/out-of-range pre x random/all-tied/tie-flood data x cr x K, plus a cs=4 cell) — test_cute_dsl_gvr_topk_decode.py.
  • One xfail documenting 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).

wf (bucketed-list path, production default where routed) — speedup vs baseline:

N \ B 1 2 4 8 16 32 64 128 256
V4-Flash 4k 1.10 1.07 1.08 1.35 1.21 1.19 1.16 1.13 1.05
32k 1.48 1.29 1.27 1.28 1.36 1.30 1.30 1.30 1.36
128k 1.74 1.90 1.81 1.67 1.78 1.79 1.83 1.86 2.21
512k 2.24 2.18 2.24 2.19 2.25 2.36 3.01 4.12 7.05
1M 3.15 2.97 2.80 2.79 2.85 3.21 4.62 7.99 12.20
V4-Pro 512k 1.91 1.78 1.88 2.11 2.14 2.19 2.58 3.48 6.25
1M 2.43 2.21 2.30 2.36 2.47 2.79 3.91 6.79 10.51

Geomeans over the full 9x9 grid (all layers x all steps):

path V4-Flash V4-Pro
wf (bucketed list) 1.749 1.581
va (seed counts) 1.229 1.250
vb (rungs, fallback) 1.032 1.107

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 <= 16 — net accounting stays positive everywhere routed.

Correctness validation

  • 486-cell campaign: exact in every cell (tie-aware multiset).
  • Unit suites: emission contracts, degenerate-preIdx battery, CUDA-graph capture/replay 11/11 identical to eager.
  • Real-model E2E (DeepSeek V4-Flash, TP2, 2x B200): with TRTLLM_GVR_EXT=1, the production-path selections of every indexer layer across all captured decode steps (231 rows) are score-multiset identical to torch.topk. This acceptance run is also what exposed (and now guards, via the new unit battery) the degenerate-preIdx bug fixed here.

🤖 Generated with Claude Code

Dev Engineer Review

  • Adds an opt-in, env-gated emission-assisted GVR top-K decode path for DeepSeek V4 sparse attention (TRTLLM_GVR_EXT=1), focused on FP4 CuTE DSL paged-MQA logits + CuTE GVR top-K flow.
  • Updates Indexer to lazily construct per-indexer GvrExtState and a host-side emission plan/route; forwards emitted seed/candidate hints and (optionally) block_max into the subsequent CuTE FP4 paged-MQA logits and GVR top-K stages. The path is restricted to safe conditions (notably next_n==1 and no DSL atom-split path).
  • Introduces GvrExtState (gvr_ext.py) to manage persistent device buffers (seed/xstate, optional list-tier candidates, and lazily allocated block_max) with CUDA-graph-safe lifecycle (explicit error on reallocation during capture). Includes tier warm-start via prev_topk/emitted_tier and builds the correct kwargs for both emission and top-K decode.
  • Extends CuTe custom ops and runners:
    • trtllm::cute_dsl_fp4_paged_mqa_logits now optionally emits block-meta/seed-count/candidate artifacts used by the GVR extension, while the public op wrapper remains logits-only.
    • trtllm::cute_dsl_gvr_topk_decode now accepts optional emission-side inputs (seed_thr, seed_counts, xstate, candidate SoA buffers, and block_max) and tuning/throughput knobs (accept_cap, kc_override, num_threads).
  • Adds routing/planning helpers (gvr_routing.py) selecting emission tiers (list/counts/rungs) and concrete launch knobs via TopkRoute, including whether block_max must be attached.
  • Extends the GVR top-K kernel (gvr_topk_decode.py) with emission-assisted admission/candidate filtering, block-skip pipelines, self-scan support, and CUDA-graph-safe state updates; includes a phase1r_data_reseed recovery to handle degenerate preIdx-derived brackets.
  • Improves radix-top-K dispatch for compress_ratio>1 by preferring metadata.gen_indexer_kv_lens_cuda_runtime when available (instead of defaulting to context_lens).
  • Notes the earlier fused indexer-to-top-K “handshake” behavior was reverted/removed from this extension (no arrival/handshake-style state in the new GvrExtState), matching the provided commit-message intent.

QA Engineer Review

Unit tests added/modified (in tests/)

  • tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py
    • Added: test_cute_dsl_gvr_topk_decode_degenerate_preidx (parametrized pre_mode/data_mode for next_n==1)
    • Added: test_cute_dsl_gvr_topk_decode_degenerate_preidx_cs4
    • Added (xfail): test_cute_dsl_gvr_topk_decode_tie_flood_beyond_capacity
  • tests/unittest/_torch/attention/sparse/test_cute_dsl_fp4_paged_mqa_logits.py
    • Added: test_cute_dsl_fp4_paged_mqa_logits_block_meta
    • Added: test_cute_dsl_fp4_paged_mqa_logits_seed_counts
    • Added: test_cute_dsl_fp4_paged_mqa_logits_cand
    • Added: test_cute_dsl_fp4_paged_mqa_logits_cand_bucketed
  • tests/scripts/cute_dsl_kernels/top_k/run_gvr_topk.py
    • Updated: test harness/wrapper to wire new optional GVR/emission inputs and modes

Integration test-list coverage

  • Not verified/provided for new or modified unit tests.
  • Verdict: needs follow-up.

Stock-path isolation, and the one deliberate exception

Everything this PR adds is behind a compile-time gate, so a build that
enables none of it is byte-identical to main. That is verified, not
asserted: compiling the stock configuration (K=1024, cr=4, 512 threads,
cs=1, only arguments main's ctor also has) and diffing the emitted
PTX gives 7316 instructions on main and 7316 on this branch once
p1r_rescue is off. The knobs are p4_tail_v3, p4_fine_rangetest,
p4_scat_rangetest, pdl_wait_late, enable_block_skip,
use_ext_counts, use_ext_cand, ext_rungs, emit_xstate,
self_scan; all default to the stock behaviour, and p4_tail_v3 is
derived from the others so the assist tiers keep the faster repair.

The single exception is p1r_rescue, default on, worth +216
instructions (+3.0%):

The kernel seeds its threshold search from the previous step's top-K
indices. When those produce an unusable bracket it takes a shortcut and
emits identity output - indices [0, K). That shortcut is only correct
for the case it was written for (a row whose values are all equal).

heuristic_prev_topk is zero-initialised on main as well, with the
comment that the kernel's +1 offset makes index 1 "a valid but benign
candidate". That holds for the C++ path, but not here: all K hints land
on the same position, so the gathered values are identical, the
bracket collapses, and the shortcut fires - on the first decode step of
every request. Measured on a V4-Flash TP2 capture: 42 of 231 rows
wrong before the fix (21 layers x 2 sequences = every first step),
231/231 exact after.

So this is a pre-existing defect in the kernel, unreachable in a default
deployment because use_cute_dsl_topk and enable_heuristic_topk both
default to False, and unavoidable once they are on. The rescue rebuilds
the bracket from the row's own min/max and keeps the identity shortcut
only where it is provably exact (all values equal, or N <= K). Its
runtime cost when it does not fire is below noise (-0.26 us median over
40 interleaved paired reps, 17/40 slower); the +3% is code size, on a
branch that does not execute.

siyidNV added 30 commits July 28, 2026 03:47
…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>
The previous two commits set the list-tier batch cap from a tax model,
first flat and then proportional to the indexer time. Both were wrong.
Measured on the FP4 indexer itself - same kernel, emission outputs
attached, wall delta - at ctx 256k and 1M:

  batch      1       2       4       8      32     128
  list   +13.8%  +64.3%  +60.7%  +62.0% +122.9% +150.7%
  counts  +9.4%   +3.5%   +0.9%   +1.1%   +0.9%   +2.6%

The list tier clears its own tax only at batch 1. From batch 2 the tax
is 20-60us against a top-k saving of a few, and it keeps growing. The
proportional model underestimated batch 128 by an order of magnitude
(45us modelled, 589us measured).

Re-scoring the 126-cell grid with the measured tax:

  rule                        geomean  worst  losing
  list B<=16, N>=65536          0.969  0.159   43
  list B<=16, N>=32768          0.926  0.159   41
  list B==1,  N>=32768          1.155  0.758   17
  per-cell optimum              1.207  0.763   14

So the cap goes to 1. Note the first row: the tier was already a net
loss as shipped before this series - the 3.4us flat tax it was sized
against only holds at batch 1.

The counts tax stays inside ~2.6us at every batch measured, so the
counts/rungs split from the previous commit is unaffected.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
@siyidNV

siyidNV commented Jul 30, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #62748 [ run ] triggered by Bot. Commit: 52973e6 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #62745 [ run ] completed with state ABORTED. Commit: cf87bf0

Link to invocation

siyidNV added 4 commits July 30, 2026 02:04
Measuring the indexer across batch 1..1024 and ctx 32k..1M shows the
list tier's tax tracks B*N - the volume it emits - rather than batch:

  B*N (raw tokens)   <=0.5M     1M       >=2M
  list                +11-14%  +55-71%  +150-270%
  counts               +1-3%    +1-3%     +0-3%

Volume is the main term but not the only one: the same 1M tokens costs
+21% as 1M x B1 and +71% as 128k x B8, so there is a per-row component
too. Both cheap regions are covered by "a single row, or under ~0.75M
tokens", which lets 128k rows carry the list to batch 4 and 256k to
batch 2 - where the previous commit's flat B==1 cap gave them up.

Scored on the 126-cell grid with the measured tax:

  rule                       geomean  worst  losing
  pre-series (B<=16, N>=64k)   0.894  0.078   46
  previous (B==1)              1.145  0.769   19
  this (B==1 or <=0.75M)       1.162  0.786   17
  per-cell optimum             1.182  0.786   14

The counts tax stays inside 3% at every batch measured, so it remains
the default and the rungs handover at B>=32 is unchanged.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
Scored on the full mean-over-layers-and-steps grid (length 8k-1024k,
batch 1-1024, all powers of two, measured emission tax):

  rule                          geomean  worst  losing
  list N>=32768, rungs B>=32      1.133  0.802  20/154
  list N>=16384                   1.140  0.802  16/154
  rungs B>=16                     1.133  0.802  19/154
  both                            1.140  0.802  15/154
  per-cell optimum                1.147  0.802  13/154

64k rows (n_comp 16387) were below the length gate but the list wins
there by ~2us at batch 2-8 once the volume gate lets it through, and at
batch 16 several cells prefer the zero-emission tier to counts.

The remaining band - 64k-512k at batch 2-16 - is not a routing problem:
the list tier is far worse there (0.65 down to 0.09; the emitted volume
crosses a megatoken and the tax takes the kernel from 17us to 100-200),
and counts is already the best of the three.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
The counts tier attached block_max from n_comp 65536, but the prefix
does not earn its setup there. Attaching it to the same shape and
diffing, 30 interleaved cold reps:

  n_comp    batch 1   batch 4   batch 16
   65538      +0.0      -0.1      +0.6   (9-16 of 30 reps faster)
  131075      -4.6      -5.0      -4.5   (30 of 30)
  262127     -13.8     -14.2     -14.0   (30 of 30)

At 65538 it is a wash and a small loss at batch 16; the win starts at
131072. That threshold was where the mean-metric grid lost worst:
flash 256k (n_comp 65538) measured 0.802 / 0.834 / 0.859 at batch
4 / 8 / 16 against the baseline. With the prefix off those cells become
1.086 / 1.150 / 1.199 - the loss was the prefix's setup, not the tier.

The time model says the same thing: the counts tier's slope is already
better than the baseline's (28-35 vs 32-52 us per million compressed
elements) and its fixed cost is the same, so a 3-4us step at one
length was the whole gap.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
Default unchanged. Added to test whether the P4 head's bin search sets
the short-row floor - it does not: at pro n=4099 and n=8195, batch 1,
widths 128/256/512/1024 land within 0.5us of each other and wider is
marginally faster, so the search is not the bottleneck. Kept because it
is the only way to ask that question again without patching the file.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #62748 [ run ] completed with state SUCCESS. Commit: 52973e6
/LLM/main/L0_MergeRequest_PR pipeline #50880 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

siyidNV added 2 commits July 30, 2026 23:55
The previous commit raised the counts tier's block_max threshold to
131072 on a synthetic-Gaussian A/B. On captured V4 rows the prefix
already pays at 65536 (flash n_comp 65537, nsys kernel-only over 5
layers x all decode steps: 16.1us with vs 17.3us without at batch
4/8/16), so put it back and note that this threshold must be set from
captures.

Raising it also silently enabled cluster split for the rungs tier at
65537, because the cs=1 clamp it used to hit was conditioned on
block_max being attached. Cluster split is a loss for the assist tiers
at that length whether or not the prefix is on (rungs at batch 4:
17.9 / 19.0 / 20.4us at cs 1 / 4 / 8; cs8 spills to 31.5us at batch
16), so gate the split on its own threshold instead.

Worst mean-metric cell on the 154-cell grid: 0.794 -> 0.833.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
p1r_rescue defaulted to True, so the single-operator path compiled it
too: 14016 vs 13570 PTX instructions against upstream for the same
launch (flash 256k B4, both trees pick cs=4 / 512 threads), and 3.5-6%
slower on the flash 128k/256k cells of the captured grid.

Derive it from the assist tiers instead, the way p4_tail_v3 already is.
The stock build is then instruction-identical to upstream (13570 / 14498
lines on both trees) and the tiers - which are the ones seeding from a
possibly zero-init prev_topk - keep the rescue. Exactness unchanged on
counts / rungs / list at flash 256k.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
@siyidNV

siyidNV commented Jul 31, 2026

Copy link
Copy Markdown
Collaborator Author

Two fixes from the captured-grid loop.

Routing gates. The previous commit raised the counts tier's block_max threshold to 131072 on a synthetic-Gaussian A/B. On captured V4 rows the prefix already pays at 65536 (flash n_comp 65537, nsys kernel-only over 5 layers x all decode steps: 16.1us with vs 17.3us without, at batch 4/8/16), so it goes back. Raising it had also silently enabled cluster split for the rungs tier at that length, because the cs=1 clamp it used to hit was conditioned on block_max being attached; the split is a loss there either way (rungs at batch 4: 17.9 / 19.0 / 20.4us at cs 1 / 4 / 8, and cs8 spills to 31.5us at batch 16), so it now has its own gate. Worst mean-metric cell on the 154-cell grid: 0.794 -> 0.833.

Stock-path isolation. p1r_rescue defaulted to True, so the single-operator path compiled the refine rescue too - 14016 vs 13570 PTX instructions against upstream for the same launch (both trees pick cs=4 / 512 threads at flash 256k B4), and 3.5-6% slower on the flash 128k/256k cells. It is now derived from the assist tiers the way p4_tail_v3 already is; the stock build is instruction-identical to upstream (13570 / 14498 lines on both trees) and the tiers keep the rescue. Exactness unchanged on counts / rungs / list.

/bot run

There is a band of row lengths where the stock kernel splits the row
across a cluster and its split grid still fits one wave. The assist
tiers cannot follow there - splitting a row costs them more than the
scan it saves, with or without the block-skip prefix - so the stock
kernel wins the cell outright and the epilogue should emit nothing.

Measured against the stock kernel over the 154-cell captured grid
(2 models x 7 context lengths x 11 batches, nsys kernel-only, emission
tax included): without the band 15 cells run below stock, worst 0.863
at flash 512k batch 16; with it, 2 cells at 0.99. Against the reference
baseline the grid geomean goes 1.139 -> 1.145 and the worst cell
0.833 -> 0.912. Batch 1-2 keep the list tier inside the band, where it
still beats stock by 1.5x, so the list check runs first.

The larger K is, the less row the tiers need before they pay, hence the
two upper bounds.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
@siyidNV

siyidNV commented Jul 31, 2026

Copy link
Copy Markdown
Collaborator Author

Third fix in this batch: a mid-row band now routes back to the stock kernel.

There is a range of row lengths where the stock kernel splits the row across a cluster and its split grid still fits one wave. The assist tiers cannot follow - splitting a row costs them more than the scan it saves, with or without the block-skip prefix (at n_comp 65537, batch 4: cs1 15.8us / cs4 18.6us with rung tightening forced off on both, so the split itself is what loses). In that band the stock kernel wins the cell outright, so the epilogue now emits nothing there and the top-k falls through to the stock branch.

Scored over the full 154-cell captured grid (2 models x 7 context lengths x 11 batches, nsys kernel-only, emission tax measured on the FP4 indexer and included), by the shipped plan_emission rather than an offline copy of it:

vs the stock kernel this PR changes vs the reference baseline
before geomean 1.196, 15 cells below parity, worst 0.863 geomean 1.139, worst 0.833
after geomean 1.203, 2 cells at 0.99, worst 0.992 geomean 1.145, worst 0.912, best 2.76

So the PR no longer regresses any grid cell against the code it is changing. Batch 1-2 keep the list tier inside the band - it still beats stock by ~1.5x there - so the list check runs first.

/bot run

@siyidNV

siyidNV commented Jul 31, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #63090 [ run ] triggered by Bot. Commit: 1e70d74 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #63090 [ run ] completed with state FAILURE. Commit: 1e70d74
/LLM/main/L0_MergeRequest_PR pipeline #51184 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@siyidNV

siyidNV commented Aug 1, 2026

Copy link
Copy Markdown
Collaborator Author

The failure in #63090 is infrastructure, not this PR: L0_Test-x86_64-Single-GPU #244 died after ~18s loading trtllm-jenkins-shared-lib@main with git@gitlab-master.nvidia.com: Permission denied (publickey), before any test ran. Build-x86_64, Build-SBSA and L0_Test-SBSA-Single-GPU all passed. The bot analysis reaches the same conclusion. Re-running.

/bot run

@siyidNV

siyidNV commented Aug 1, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #63184 [ run ] triggered by Bot. Commit: 1e70d74 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #63184 [ run ] completed with state SUCCESS. Commit: 1e70d74
/LLM/main/L0_MergeRequest_PR pipeline #51268 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

This reverts commit 78227d5.

The reverted commit read the rescue's presence on the stock path as an
isolation leak. It is not: `143453c294` added it deliberately as a
correctness fix for that exact path - the first decode step of a request
feeds a zero-init prev_topk, and the old degenerate shortcut answers it
with identity indices [0, K) instead of the top-K. The tests that pin
this (test_cute_dsl_gvr_topk_decode_degenerate_preidx and its cs4
variant) call the stock custom op directly and do not exist on main, so
deriving the knob from the assist tiers turned the fix off exactly where
it is needed and broke them.

The +446 PTX instructions and the 3.5-6% it costs the stock path on the
flash 128k/256k cells are the price of that fix, not leakage from the
emission tiers. The tiers stay isolated the way they already were -
every emission knob is still compile-time gated off the stock build.

Signed-off-by: siyidNV <297196620+siyidNV@users.noreply.github.com>
@siyidNV

siyidNV commented Aug 1, 2026

Copy link
Copy Markdown
Collaborator Author

Correction, and a revert. My earlier comment in this batch claimed p1r_rescue defaulting to True was an isolation leak into the stock path. That was wrong: 143453c294 put it there deliberately as a correctness fix for that exact path, and test_cute_dsl_gvr_topk_decode_degenerate_preidx (added by this PR, not present on main) calls the stock custom op directly to pin it. Deriving the knob from the assist tiers turned the fix off where it is needed.

Reproduced at kernel level over the same parameter matrix the test uses (dtype x top_k x compress_ratio x pre_mode x data_mode, plus the cs4 variant):

degenerate-preIdx cases passing
rescue on (reverted state) 36/36, cs4 variant passes
rescue derived from tiers 12/36, cs4 variant fails

So 78227d5397 is reverted. The +446 PTX instructions and the 3.5-6% it costs the stock path on the flash 128k/256k cells are the price of the correctness fix, not leakage from the emission tiers - the tiers themselves remain compile-time gated off the stock build.

Re-scored the whole 154-cell grid against a stock arm re-measured with the rescue on, through the shipped plan_emission. The mid-row escape band's constants come out unchanged (still exactly optimal on the corrected data):

vs the stock kernel this PR changes vs the reference baseline
without the band geomean 1.219, 13 cells below parity, worst 0.867 geomean 1.139, worst 0.833
with the band geomean 1.224, 1 cell at 0.997 geomean 1.143, worst 0.877, best 2.76

On the previous CI run: the gvr top-k degenerate tests were genuinely broken by my commit, so the DGX_B200-PyTorch-9 unittest/_torch/attention shard was ours. The other failing shards (test_llm_pytorch.py part0, and the _torch/attention shard that explicitly --ignores both files this PR touches) cannot be, and match the harness Unexpected triton_kernels path error the bot analysis flagged.

/bot run

@siyidNV

siyidNV commented Aug 1, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #63199 [ run ] triggered by Bot. Commit: 0982dc1 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #63199 [ run ] completed with state DISABLED
Pipeline is freezed and top-1 instance is under maintenance. For urgent request, contact Yiteng Niu

Link to invocation

@siyidNV

siyidNV commented Aug 1, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #63212 [ run ] triggered by Bot. Commit: 0982dc1 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #63212 [ run ] completed with state DISABLED
Pipeline is freezed and top-1 instance is under maintenance. For urgent request, contact Yiteng Niu

Link to invocation

@siyidNV

siyidNV commented Aug 2, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants