[None][feat] Self-sampling GVR V2 prefill indexer top-K - #18702
Conversation
|
/bot run |
|
PR_Github #71439 [ run ] triggered by Bot. Commit: |
|
PR_Github #71439 [ run ] completed with state
|
a923cc0 to
be15dca
Compare
|
/bot run |
|
PR_Github #71465 [ run ] triggered by Bot. Commit: |
|
/bot run --reuse-test --disable-fail-fast Merged |
|
PR_Github #71476 Bot args parsing error: usage: /bot [-h] |
|
/bot run --reuse-test --disable-fail-fast |
|
PR_Github #71487 [ run ] triggered by Bot. Commit: |
|
PR_Github #71465 [ run ] completed with state |
|
PR_Github #71487 [ run ] completed with state
|
|
/bot run --reuse-test --disable-fail-fast |
1 similar comment
|
/bot run --reuse-test --disable-fail-fast |
|
PR_Github #71543 [ run ] triggered by Bot. Commit: |
|
PR_Github #71543 [ run ] completed with state
|
|
/bot run --reuse-test --disable-fail-fast |
|
PR_Github #71654 [ run ] triggered by Bot. Commit: |
|
Refreshed the prefill numbers with Nsight (pure kernel GPU time, NVTX projection, cold-L2) instead of CUDA-event wall time, and swept batch as well as window — the event-based version I first posted over-smoothed a real BS dependence. Same synthetic setup (there's no captured prefill data — the §C/§D prefill streams are corrupt — so batches are built from real §G decode rows with causal / lower-triangular windows). Correctness is unchanged: GVR matches the CUDA radix and
(Pro tracks Flash.) The crossover moves with both window and batch:
So a length-gated prefill dispatch holds, but the threshold should also look at batch — near the crossover window, keep radix for very small batches. Single B200, synthetic distributions (real per-row logits and length range, no true intra-request correlation); preliminary, not a ship verdict. |
|
PR_Github #71654 [ run ] completed with state
|
|
/bot run --reuse-test --disable-fail-fast |
|
PR_Github #71674 [ run ] triggered by Bot. Commit: |
|
PR_Github #71674 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #71726 [ run ] triggered by Bot. Commit: |
|
PR_Github #71726 [ run ] completed with state |
|
PR_Github #72079 [ run ] completed with state
|
|
/bot run --reuse-test --disable-fail-fast |
|
PR_Github #72172 [ run ] triggered by Bot. Commit: |
|
PR_Github #72172 [ run ] completed with state
|
|
/bot run --reuse-test --disable-fail-fast |
|
PR_Github #72206 [ run ] triggered by Bot. Commit: |
|
PR_Github #72206 [ run ] completed with state
|
|
/bot run --reuse-test --disable-fail-fast |
|
PR_Github #72217 [ run ] triggered by Bot. Commit: |
|
PR_Github #72217 [ run ] completed with state
|
|
/bot run --reuse-test --disable-fail-fast |
|
PR_Github #72249 [ run ] triggered by Bot. Commit: |
|
PR_Github #72249 [ run ] completed with state |
reasonsolo
left a comment
There was a problem hiding this comment.
No disagg code change. Approved.
zongfeijing
left a comment
There was a problem hiding this comment.
LGTM from the kernel side
…ow engines (TRT-LLM #18702 port) Adds an optional `row_starts` argument to `top_k_varlen`: row `r` ranks the window `logits[r, row_starts[r] : row_starts[r] + seq_lens[r]]` and returns window-LOCAL indices (`column - row_starts[r]`), `-1` padded, identity for windows shorter than `top_k`. This is the DSA prefill-indexer layout (every request's keys packed along the column axis, one query row per prompt token, causal window into its own request's slice); columns outside a row's window are never read, so they may be unwritten. Before this change the only FlashInfer path for that layout was the radix `top_k_ragged_transform`. Device (`gvr2_topk_decode.py`): `GvrMainKernel(prefill=True)` const_expr variant of the streaming `main` family, port of NVIDIA/TensorRT-LLM#18702. The window is read from a float4-aligned base `col0 = ks & ~3` with `lead = ks & 3` low lanes masked (sampling, P3 histogram mask, tail loop, emit) and the emitted indices shifted by `-lead` through a smem slot, so the row base stays 16-byte aligned for any `ks`. FlashInfer ABI: the `pre_idx` slot carries `row_starts` and the `kv_lens` slot keeps its per-row LENGTH meaning (`ke = ks + len`, clamped to the logits width; upstream passes ks/ke), so a caller never materializes row ends. Compiled objects are keyed on the prefill flag (`_pf` persist name); hint-free only (`next_n == 1`, `cr == 1`, no split). Host (`gvr2_topk_host.py`): `run_varlen(..., row_starts=)` routes windowed calls to `_launch_prefill` — one engine per (row-count tier, top_k, pow2 envelope bucket), streaming `main` plan, rows slabbed at 32768 (gridDim.y), values gathered from the window; the `reference` engine copies each window into an aligned `-inf`-padded row. The tier-0 (<= 148 rows) -> tier-1 promotion for envelopes just above the BLK=512 sampling gate follows upstream main's refinement of #18702. `max_seq_len` in windowed mode is the max window length and only picks the tier / bucket: the kernel clamps every window to the logits width, never to the hint. `warmup_prefill(top_k, max_cols)` compiles the tier x bucket set (idempotent, both tier edges of every bucket) for CUDA graph capture; `prefill_ready` reports it. Main kernel, all varlen modes: short rows (`n <= k`, block-uniform) now skip the body outright and go straight to the identity epilogue. Upstream runs the body as a zero-work pass (n = 0, TGT sentinel) — every barrier, the sample scan and, in FlashInfer's hint-free form, the identity bracket gather (`gather_hint(identity=True)`, two sites x two barriers) still execute, ~1 us per CTA on B200 at BLK=256. In a causal prefill staircase the first K rows of every request are short, so this is 25% of the rows at L = 8192, K = 2048 and all of them at L <= K. The hint-free identity bracket (`gather_hint(identity=True)`) samples from the first in-window lane (`base = lead`) rather than the float4-rounded row base: the <= 3 lead lanes belong to the previous request (or are unwritten), and sampling them widened the bracket to foreign values or NaN (sentinel fallback), 50% slower on misaligned windows at K = 512 than on aligned ones — upstream's compiled-out gather has no such dependence. Same-node A/B against the upstream TRT-LLM prefill kernel (identical plans, bit-identical results on every cell; SPEEDUP = upstream time / ours, min of 3 interleaved): K R x L B200 B300 Rubin 512 1 x 8192 1.37x 1.37x 1.35x 512 4 x 2048 1.79x 1.79x 1.72x 512 4 x 2047 (misaligned) 1.75x 1.79x 1.72x 512 4 x 8192 1.39x 1.39x 1.39x 512 8 x 4096 1.82x 1.85x 1.79x 512 16 x 2048 1.72x 1.72x 1.69x 512 1 x 32768 1.05x 1.05x 1.06x 512 64 x 512 (all rows <= K) 4.35x 4.55x 4.00x 512 128 x 64 (all rows <= K) 4.35x 4.35x 4.00x 2048 1 x 8192 1.23x 1.23x 1.20x 2048 4 x 2048 (all rows <= K) 3.23x 3.12x 2.86x 2048 4 x 2047 (all rows <= K) 3.23x 3.12x 2.86x 2048 4 x 8192 1.25x 1.25x 1.23x 2048 8 x 4096 1.79x 1.79x 1.75x 2048 16 x 2048 (all rows <= K) 2.86x 2.86x 2.78x 2048 1 x 32768 1.00x 1.00x 1.01x 2048 64 x 512 (all rows <= K) 2.86x 2.94x 2.78x 2048 128 x 64 (all rows <= K) 3.12x 3.12x 2.78x API (`topk_varlen.py`): `row_starts` validated as a contiguous 1-D int32 CUDA tensor with one entry per row, `next_n == 1`, `compress_ratio == 1`, `pre_idx=None`. Only gvr_2's checker admits windowed calls, so `auto` selects gvr_2 and an explicit other backend is refused at validation. The trace template gains the optional `row_starts` input with a window-aware reference and check. Perf (SGLang-shaped staircase, R requests x L prompt tokens, causal windows, fp32, graph replay, 3 interleaved repetitions, min, same node and job for every column; `proxy` = decode engine on a column-0 layout of identical per-row work, as the control): Time in us, min of 3 interleaved repetitions, graph replay. R requests x L prompt tokens, causal windows, fp32. B200 (148 SMs) K R x L rows radix_win gvr2_win sgl_v2 vs radix vs sgl_v2 512 1 x 8192 8192 426.1 93.7 120.6 4.55x 1.28x 512 4 x 2048 8192 221.8 62.8 71.2 3.57x 1.14x 512 4 x 2047 8188 221.8 62.8 72.8 3.57x 1.16x 512 4 x 8192 32768 1690.6 358.1 473.1 4.76x 1.32x 512 8 x 4096 32768 1289.1 349.5 379.4 3.70x 1.09x 512 16 x 2048 32768 972.3 251.9 314.7 3.85x 1.25x 2048 1 x 8192 8192 375.8 124.4 109.2 3.03x 0.88x 2048 4 x 2048 8192 49.7 19.4 19.0 2.56x 0.98x 2048 4 x 2047 8188 49.5 19.2 18.8 2.56x 0.98x 2048 4 x 8192 32768 1485.5 477.4 425.5 3.12x 0.89x 2048 8 x 4096 32768 895.5 333.1 276.7 2.70x 0.83x 2048 16 x 2048 32768 239.9 86.2 107.5 2.78x 1.25x B300 (148 SMs) K R x L rows radix_win gvr2_win sgl_v2 vs radix vs sgl_v2 512 1 x 8192 8192 412.4 90.6 115.7 4.55x 1.28x 512 4 x 2048 8192 212.2 60.2 68.1 3.57x 1.14x 512 4 x 2047 8188 212.2 60.3 69.8 3.57x 1.15x 512 4 x 8192 32768 1636.8 346.7 454.2 4.76x 1.32x 512 8 x 4096 32768 1243.3 337.1 364.3 3.70x 1.08x 512 16 x 2048 32768 934.6 242.5 301.5 3.85x 1.25x 2048 1 x 8192 8192 363.8 120.4 104.9 3.03x 0.87x 2048 4 x 2048 8192 46.0 18.6 18.5 2.50x 0.99x 2048 4 x 2047 8188 45.8 18.4 18.4 2.50x 1.00x 2048 4 x 8192 32768 1437.9 462.4 408.5 3.12x 0.88x 2048 8 x 4096 32768 861.9 320.7 264.3 2.70x 0.83x 2048 16 x 2048 32768 225.1 81.7 101.3 2.78x 1.23x Rubin (208 SMs) K R x L rows radix_win gvr2_win sgl_v2 vs radix vs sgl_v2 512 1 x 8192 8192 254.6 57.4 128.5 4.35x 2.22x 512 4 x 2048 8192 130.6 38.4 72.1 3.45x 1.89x 512 4 x 2047 8188 129.6 38.4 72.9 3.33x 1.89x 512 4 x 8192 32768 1001.3 215.2 503.8 4.76x 2.33x 512 8 x 4096 32768 775.4 211.2 437.3 3.70x 2.08x 512 16 x 2048 32768 590.3 153.8 387.6 3.85x 2.50x 2048 1 x 8192 8192 226.2 77.8 119.0 2.94x 1.54x 2048 4 x 2048 8192 32.7 13.5 28.2 2.44x 2.08x 2048 4 x 2047 8188 32.3 13.1 28.1 2.44x 2.13x 2048 4 x 8192 32768 884.0 287.8 462.3 3.03x 1.61x 2048 8 x 4096 32768 547.7 202.2 339.0 2.70x 1.67x 2048 16 x 2048 32768 156.9 53.5 176.2 2.94x 3.33x The remaining K = 2048 long-row cells where SGLang's kernel is 12-21% faster on B200 / B300 are the streaming main kernel's per-row cost on non-trivial rows (shared with upstream); Rubin has no losing cell. Tests (`tests/topk_varlen/test_topk_varlen_prefill.py`): out-of-window cells poisoned with NaN/+-inf/3e38 and the <= 3 lead lanes before each misaligned `ks` set to +inf (a leak becomes top-1, a missed frame correction a negative index); causal ramp for K in {512, 1024, 2048}; leads 1..3 with K in {512, 2048}; window lengths {0, 1, K-1, K, K+1}; all-equal and two-valued ties; fewer than K finite values (`-inf` masks) at every lead; `ks + len` past the width and `ks` at the width; `return_values`; reference engine parity; CUDA-graph capture/replay with changed `row_starts` and lengths between replays; unwarmed capture raises; 70000-row slab; prefill engine key distinct from decode; API validation and backend admission. tests/topk_varlen on the final tree: 370 passed on B200, B300, B100 and DRIVE P2021, 371 on Rubin (SM107), 83 on RTX 5080, 80 on A100 / L40S / H100 (rest skipped: no gvr_2) AI-assisted. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…ow engines (TRT-LLM #18702 port) Adds an optional `row_starts` argument to `top_k_varlen`: row `r` ranks the window `logits[r, row_starts[r] : row_starts[r] + seq_lens[r]]` and returns window-LOCAL indices (`column - row_starts[r]`), `-1` padded, identity for windows shorter than `top_k`. This is the DSA prefill-indexer layout (every request's keys packed along the column axis, one query row per prompt token, causal window into its own request's slice); columns outside a row's window are never read, so they may be unwritten. Before this change the only FlashInfer path for that layout was the radix `top_k_ragged_transform`. Device (`gvr2_topk_decode.py`): `GvrMainKernel(prefill=True)` const_expr variant of the streaming `main` family, port of NVIDIA/TensorRT-LLM#18702. The window is read from a float4-aligned base `col0 = ks & ~3` with `lead = ks & 3` low lanes masked (sampling, P3 histogram mask, tail loop, emit) and the emitted indices shifted by `-lead` through a smem slot, so the row base stays 16-byte aligned for any `ks`. FlashInfer ABI: the `pre_idx` slot carries `row_starts` and the `kv_lens` slot keeps its per-row LENGTH meaning (`ke = ks + len`, clamped to the logits width; upstream passes ks/ke), so a caller never materializes row ends. Compiled objects are keyed on the prefill flag (`_pf` persist name); hint-free only (`next_n == 1`, `cr == 1`, no split). Host (`gvr2_topk_host.py`): `run_varlen(..., row_starts=)` routes windowed calls to `_launch_prefill` — one engine per (row-count tier, top_k, pow2 envelope bucket), streaming `main` plan, rows slabbed at 32768 (gridDim.y), values gathered from the window; the `reference` engine copies each window into an aligned `-inf`-padded row. The tier-0 (<= 148 rows) -> tier-1 promotion for envelopes just above the BLK=512 sampling gate follows upstream main's refinement of #18702. `max_seq_len` in windowed mode is the max window length and only picks the tier / bucket: the kernel clamps every window to the logits width, never to the hint. `warmup_prefill(top_k, max_cols)` compiles the tier x bucket set (idempotent, both tier edges of every bucket) for CUDA graph capture; `prefill_ready` reports it. Main kernel, all varlen modes: short rows (`n <= k`, block-uniform) now skip the body outright and go straight to the identity epilogue. Upstream runs the body as a zero-work pass (n = 0, TGT sentinel) — every barrier, the sample scan and, in FlashInfer's hint-free form, the identity bracket gather (`gather_hint(identity=True)`, two sites x two barriers) still execute, ~1 us per CTA on B200 at BLK=256. In a causal prefill staircase the first K rows of every request are short, so this is 25% of the rows at L = 8192, K = 2048 and all of them at L <= K. The hint-free identity bracket (`gather_hint(identity=True)`) samples from the first in-window lane (`base = lead`) rather than the float4-rounded row base: the <= 3 lead lanes belong to the previous request (or are unwritten), and sampling them widened the bracket to foreign values or NaN (sentinel fallback), 50% slower on misaligned windows at K = 512 than on aligned ones — upstream's compiled-out gather has no such dependence. Same-node A/B against the upstream TRT-LLM prefill kernel (identical plans, bit-identical results on every cell; SPEEDUP = upstream time / ours, min of 3 interleaved): K R x L B200 B300 Rubin 512 1 x 8192 1.37x 1.37x 1.35x 512 4 x 2048 1.79x 1.79x 1.72x 512 4 x 2047 (misaligned) 1.75x 1.79x 1.72x 512 4 x 8192 1.39x 1.39x 1.39x 512 8 x 4096 1.82x 1.85x 1.79x 512 16 x 2048 1.72x 1.72x 1.69x 512 1 x 32768 1.05x 1.05x 1.06x 512 64 x 512 (all rows <= K) 4.35x 4.55x 4.00x 512 128 x 64 (all rows <= K) 4.35x 4.35x 4.00x 2048 1 x 8192 1.23x 1.23x 1.20x 2048 4 x 2048 (all rows <= K) 3.23x 3.12x 2.86x 2048 4 x 2047 (all rows <= K) 3.23x 3.12x 2.86x 2048 4 x 8192 1.25x 1.25x 1.23x 2048 8 x 4096 1.79x 1.79x 1.75x 2048 16 x 2048 (all rows <= K) 2.86x 2.86x 2.78x 2048 1 x 32768 1.00x 1.00x 1.01x 2048 64 x 512 (all rows <= K) 2.86x 2.94x 2.78x 2048 128 x 64 (all rows <= K) 3.12x 3.12x 2.78x API (`topk_varlen.py`): `row_starts` validated as a contiguous 1-D int32 CUDA tensor with one entry per row, `next_n == 1`, `compress_ratio == 1`, `pre_idx=None`. Only gvr_2's checker admits windowed calls, so `auto` selects gvr_2 and an explicit other backend is refused at validation. The trace template gains the optional `row_starts` input with a window-aware reference and check. Perf (SGLang-shaped staircase, R requests x L prompt tokens, causal windows, fp32, graph replay, 3 interleaved repetitions, min, same node and job for every column; `proxy` = decode engine on a column-0 layout of identical per-row work, as the control): Time in us, min of 3 interleaved repetitions, graph replay. R requests x L prompt tokens, causal windows, fp32. B200 (148 SMs) K R x L rows radix_win gvr2_win sgl_v2 vs radix vs sgl_v2 512 1 x 8192 8192 426.1 93.7 120.6 4.55x 1.28x 512 4 x 2048 8192 221.8 62.8 71.2 3.57x 1.14x 512 4 x 2047 8188 221.8 62.8 72.8 3.57x 1.16x 512 4 x 8192 32768 1690.6 358.1 473.1 4.76x 1.32x 512 8 x 4096 32768 1289.1 349.5 379.4 3.70x 1.09x 512 16 x 2048 32768 972.3 251.9 314.7 3.85x 1.25x 2048 1 x 8192 8192 375.8 124.4 109.2 3.03x 0.88x 2048 4 x 2048 8192 49.7 19.4 19.0 2.56x 0.98x 2048 4 x 2047 8188 49.5 19.2 18.8 2.56x 0.98x 2048 4 x 8192 32768 1485.5 477.4 425.5 3.12x 0.89x 2048 8 x 4096 32768 895.5 333.1 276.7 2.70x 0.83x 2048 16 x 2048 32768 239.9 86.2 107.5 2.78x 1.25x B300 (148 SMs) K R x L rows radix_win gvr2_win sgl_v2 vs radix vs sgl_v2 512 1 x 8192 8192 412.4 90.6 115.7 4.55x 1.28x 512 4 x 2048 8192 212.2 60.2 68.1 3.57x 1.14x 512 4 x 2047 8188 212.2 60.3 69.8 3.57x 1.15x 512 4 x 8192 32768 1636.8 346.7 454.2 4.76x 1.32x 512 8 x 4096 32768 1243.3 337.1 364.3 3.70x 1.08x 512 16 x 2048 32768 934.6 242.5 301.5 3.85x 1.25x 2048 1 x 8192 8192 363.8 120.4 104.9 3.03x 0.87x 2048 4 x 2048 8192 46.0 18.6 18.5 2.50x 0.99x 2048 4 x 2047 8188 45.8 18.4 18.4 2.50x 1.00x 2048 4 x 8192 32768 1437.9 462.4 408.5 3.12x 0.88x 2048 8 x 4096 32768 861.9 320.7 264.3 2.70x 0.83x 2048 16 x 2048 32768 225.1 81.7 101.3 2.78x 1.23x Rubin (208 SMs) K R x L rows radix_win gvr2_win sgl_v2 vs radix vs sgl_v2 512 1 x 8192 8192 254.6 57.4 128.5 4.35x 2.22x 512 4 x 2048 8192 130.6 38.4 72.1 3.45x 1.89x 512 4 x 2047 8188 129.6 38.4 72.9 3.33x 1.89x 512 4 x 8192 32768 1001.3 215.2 503.8 4.76x 2.33x 512 8 x 4096 32768 775.4 211.2 437.3 3.70x 2.08x 512 16 x 2048 32768 590.3 153.8 387.6 3.85x 2.50x 2048 1 x 8192 8192 226.2 77.8 119.0 2.94x 1.54x 2048 4 x 2048 8192 32.7 13.5 28.2 2.44x 2.08x 2048 4 x 2047 8188 32.3 13.1 28.1 2.44x 2.13x 2048 4 x 8192 32768 884.0 287.8 462.3 3.03x 1.61x 2048 8 x 4096 32768 547.7 202.2 339.0 2.70x 1.67x 2048 16 x 2048 32768 156.9 53.5 176.2 2.94x 3.33x The remaining K = 2048 long-row cells where SGLang's kernel is 12-21% faster on B200 / B300 are the streaming main kernel's per-row cost on non-trivial rows (shared with upstream); Rubin has no losing cell. Tests (`tests/topk_varlen/test_topk_varlen_prefill.py`): out-of-window cells poisoned with NaN/+-inf/3e38 and the <= 3 lead lanes before each misaligned `ks` set to +inf (a leak becomes top-1, a missed frame correction a negative index); causal ramp for K in {512, 1024, 2048}; leads 1..3 with K in {512, 2048}; window lengths {0, 1, K-1, K, K+1}; all-equal and two-valued ties; fewer than K finite values (`-inf` masks) at every lead; `ks + len` past the width and `ks` at the width; `return_values`; reference engine parity; CUDA-graph capture/replay with changed `row_starts` and lengths between replays; unwarmed capture raises; 70000-row slab; prefill engine key distinct from decode; API validation and backend admission. tests/topk_varlen on the final tree: 370 passed on B200, B300, B100 and DRIVE P2021, 371 on Rubin (SM107), 83 on RTX 5080, 80 on A100 / L40S / H100 (rest skipped: no gvr_2) AI-assisted. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…ow engines (TRT-LLM #18702 port) Adds an optional `row_starts` argument to `top_k_varlen`: row `r` ranks the window `logits[r, row_starts[r] : row_starts[r] + seq_lens[r]]` and returns window-LOCAL indices (`column - row_starts[r]`), `-1` padded, identity for windows shorter than `top_k`. This is the DSA prefill-indexer layout (every request's keys packed along the column axis, one query row per prompt token, causal window into its own request's slice); columns outside a row's window are never read, so they may be unwritten. Before this change the only FlashInfer path for that layout was the radix `top_k_ragged_transform`. Device (`gvr2_topk_decode.py`): `GvrMainKernel(prefill=True)` const_expr variant of the streaming `main` family, port of NVIDIA/TensorRT-LLM#18702. The window is read from a float4-aligned base `col0 = ks & ~3` with `lead = ks & 3` low lanes masked (sampling, P3 histogram mask, tail loop, emit) and the emitted indices shifted by `-lead` through a smem slot, so the row base stays 16-byte aligned for any `ks`. FlashInfer ABI: the `pre_idx` slot carries `row_starts` and the `kv_lens` slot keeps its per-row LENGTH meaning (`ke = ks + len`, clamped to the logits width; upstream passes ks/ke), so a caller never materializes row ends. Compiled objects are keyed on the prefill flag (`_pf` persist name); hint-free only (`next_n == 1`, `cr == 1`, no split). Host (`gvr2_topk_host.py`): `run_varlen(..., row_starts=)` routes windowed calls to `_launch_prefill` — one engine per (row-count tier, top_k, pow2 envelope bucket), streaming `main` plan, rows slabbed at 32768 (gridDim.y), values gathered from the window; the `reference` engine copies each window into an aligned `-inf`-padded row. The tier-0 (<= 148 rows) -> tier-1 promotion for envelopes just above the BLK=512 sampling gate follows upstream main's refinement of #18702. `max_seq_len` in windowed mode is the max window length and only picks the tier / bucket: the kernel clamps every window to the logits width, never to the hint. `warmup_prefill(top_k, max_cols)` compiles the tier x bucket set (idempotent, both tier edges of every bucket) for CUDA graph capture; `prefill_ready` reports it. Main kernel, all varlen modes: short rows (`n <= k`, block-uniform) now skip the body outright and go straight to the identity epilogue. Upstream runs the body as a zero-work pass (n = 0, TGT sentinel) — every barrier, the sample scan and, in FlashInfer's hint-free form, the identity bracket gather (`gather_hint(identity=True)`, two sites x two barriers) still execute, ~1 us per CTA on B200 at BLK=256. In a causal prefill staircase the first K rows of every request are short, so this is 25% of the rows at L = 8192, K = 2048 and all of them at L <= K. The hint-free identity bracket (`gather_hint(identity=True)`) samples from the first in-window lane (`base = lead`) rather than the float4-rounded row base: the <= 3 lead lanes belong to the previous request (or are unwritten), and sampling them widened the bracket to foreign values or NaN (sentinel fallback), 50% slower on misaligned windows at K = 512 than on aligned ones — upstream's compiled-out gather has no such dependence. Same-node A/B against the upstream TRT-LLM prefill kernel (identical plans, bit-identical results on every cell; SPEEDUP = upstream time / ours, min of 3 interleaved): K R x L B200 B300 Rubin 512 1 x 8192 1.37x 1.37x 1.35x 512 4 x 2048 1.79x 1.79x 1.72x 512 4 x 2047 (misaligned) 1.75x 1.79x 1.72x 512 4 x 8192 1.39x 1.39x 1.39x 512 8 x 4096 1.82x 1.85x 1.79x 512 16 x 2048 1.72x 1.72x 1.69x 512 1 x 32768 1.05x 1.05x 1.06x 512 64 x 512 (all rows <= K) 4.35x 4.55x 4.00x 512 128 x 64 (all rows <= K) 4.35x 4.35x 4.00x 2048 1 x 8192 1.23x 1.23x 1.20x 2048 4 x 2048 (all rows <= K) 3.23x 3.12x 2.86x 2048 4 x 2047 (all rows <= K) 3.23x 3.12x 2.86x 2048 4 x 8192 1.25x 1.25x 1.23x 2048 8 x 4096 1.79x 1.79x 1.75x 2048 16 x 2048 (all rows <= K) 2.86x 2.86x 2.78x 2048 1 x 32768 1.00x 1.00x 1.01x 2048 64 x 512 (all rows <= K) 2.86x 2.94x 2.78x 2048 128 x 64 (all rows <= K) 3.12x 3.12x 2.78x API (`topk_varlen.py`): `row_starts` validated as a contiguous 1-D int32 CUDA tensor with one entry per row, `next_n == 1`, `compress_ratio == 1`, `pre_idx=None`. Only gvr_2's checker admits windowed calls, so `auto` selects gvr_2 and an explicit other backend is refused at validation. The trace template gains the optional `row_starts` input with a window-aware reference and check. Perf (SGLang-shaped staircase, R requests x L prompt tokens, causal windows, fp32, graph replay, 3 interleaved repetitions, min, same node and job for every column; `proxy` = decode engine on a column-0 layout of identical per-row work, as the control): Time in us, min of 3 interleaved repetitions, graph replay. R requests x L prompt tokens, causal windows, fp32. B200 (148 SMs) K R x L rows radix_win gvr2_win sgl_v2 vs radix vs sgl_v2 512 1 x 8192 8192 426.1 93.7 120.6 4.55x 1.28x 512 4 x 2048 8192 221.8 62.8 71.2 3.57x 1.14x 512 4 x 2047 8188 221.8 62.8 72.8 3.57x 1.16x 512 4 x 8192 32768 1690.6 358.1 473.1 4.76x 1.32x 512 8 x 4096 32768 1289.1 349.5 379.4 3.70x 1.09x 512 16 x 2048 32768 972.3 251.9 314.7 3.85x 1.25x 2048 1 x 8192 8192 375.8 124.4 109.2 3.03x 0.88x 2048 4 x 2048 8192 49.7 19.4 19.0 2.56x 0.98x 2048 4 x 2047 8188 49.5 19.2 18.8 2.56x 0.98x 2048 4 x 8192 32768 1485.5 477.4 425.5 3.12x 0.89x 2048 8 x 4096 32768 895.5 333.1 276.7 2.70x 0.83x 2048 16 x 2048 32768 239.9 86.2 107.5 2.78x 1.25x B300 (148 SMs) K R x L rows radix_win gvr2_win sgl_v2 vs radix vs sgl_v2 512 1 x 8192 8192 412.4 90.6 115.7 4.55x 1.28x 512 4 x 2048 8192 212.2 60.2 68.1 3.57x 1.14x 512 4 x 2047 8188 212.2 60.3 69.8 3.57x 1.15x 512 4 x 8192 32768 1636.8 346.7 454.2 4.76x 1.32x 512 8 x 4096 32768 1243.3 337.1 364.3 3.70x 1.08x 512 16 x 2048 32768 934.6 242.5 301.5 3.85x 1.25x 2048 1 x 8192 8192 363.8 120.4 104.9 3.03x 0.87x 2048 4 x 2048 8192 46.0 18.6 18.5 2.50x 0.99x 2048 4 x 2047 8188 45.8 18.4 18.4 2.50x 1.00x 2048 4 x 8192 32768 1437.9 462.4 408.5 3.12x 0.88x 2048 8 x 4096 32768 861.9 320.7 264.3 2.70x 0.83x 2048 16 x 2048 32768 225.1 81.7 101.3 2.78x 1.23x Rubin (208 SMs) K R x L rows radix_win gvr2_win sgl_v2 vs radix vs sgl_v2 512 1 x 8192 8192 254.6 57.4 128.5 4.35x 2.22x 512 4 x 2048 8192 130.6 38.4 72.1 3.45x 1.89x 512 4 x 2047 8188 129.6 38.4 72.9 3.33x 1.89x 512 4 x 8192 32768 1001.3 215.2 503.8 4.76x 2.33x 512 8 x 4096 32768 775.4 211.2 437.3 3.70x 2.08x 512 16 x 2048 32768 590.3 153.8 387.6 3.85x 2.50x 2048 1 x 8192 8192 226.2 77.8 119.0 2.94x 1.54x 2048 4 x 2048 8192 32.7 13.5 28.2 2.44x 2.08x 2048 4 x 2047 8188 32.3 13.1 28.1 2.44x 2.13x 2048 4 x 8192 32768 884.0 287.8 462.3 3.03x 1.61x 2048 8 x 4096 32768 547.7 202.2 339.0 2.70x 1.67x 2048 16 x 2048 32768 156.9 53.5 176.2 2.94x 3.33x The remaining K = 2048 long-row cells where SGLang's kernel is 12-21% faster on B200 / B300 are the streaming main kernel's per-row cost on non-trivial rows (shared with upstream); Rubin has no losing cell. Tests (`tests/topk_varlen/test_topk_varlen_prefill.py`): out-of-window cells poisoned with NaN/+-inf/3e38 and the <= 3 lead lanes before each misaligned `ks` set to +inf (a leak becomes top-1, a missed frame correction a negative index); causal ramp for K in {512, 1024, 2048}; leads 1..3 with K in {512, 2048}; window lengths {0, 1, K-1, K, K+1}; all-equal and two-valued ties; fewer than K finite values (`-inf` masks) at every lead; `ks + len` past the width and `ks` at the width; `return_values`; reference engine parity; CUDA-graph capture/replay with changed `row_starts` and lengths between replays; unwarmed capture raises; 70000-row slab; prefill engine key distinct from decode; API validation and backend admission. tests/topk_varlen on the final tree: 370 passed on B200, B300, B100 and DRIVE P2021, 371 on Rubin (SM107), 83 on RTX 5080, 80 on A100 / L40S / H100 (rest skipped: no gvr_2) AI-assisted. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
…ow engines (TRT-LLM #18702 port) Adds an optional `row_starts` argument to `top_k_varlen`: row `r` ranks the window `logits[r, row_starts[r] : row_starts[r] + seq_lens[r]]` and returns window-LOCAL indices (`column - row_starts[r]`), `-1` padded, identity for windows shorter than `top_k`. This is the DSA prefill-indexer layout (every request's keys packed along the column axis, one query row per prompt token, causal window into its own request's slice); columns outside a row's window are never read, so they may be unwritten. Before this change the only FlashInfer path for that layout was the radix `top_k_ragged_transform`. Device (`gvr2_topk_decode.py`): `GvrMainKernel(prefill=True)` const_expr variant of the streaming `main` family, port of NVIDIA/TensorRT-LLM#18702. The window is read from a float4-aligned base `col0 = ks & ~3` with `lead = ks & 3` low lanes masked (sampling, P3 histogram mask, tail loop, emit) and the emitted indices shifted by `-lead` through a smem slot, so the row base stays 16-byte aligned for any `ks`. FlashInfer ABI: the `pre_idx` slot carries `row_starts` and the `kv_lens` slot keeps its per-row LENGTH meaning (`ke = ks + len`, clamped to the logits width; upstream passes ks/ke), so a caller never materializes row ends. Compiled objects are keyed on the prefill flag (`_pf` persist name); hint-free only (`next_n == 1`, `cr == 1`, no split). Host (`gvr2_topk_host.py`): `run_varlen(..., row_starts=)` routes windowed calls to `_launch_prefill` — one engine per (row-count tier, top_k, pow2 envelope bucket), streaming `main` plan, rows slabbed at 32768 (gridDim.y), values gathered from the window; the `reference` engine copies each window into an aligned `-inf`-padded row. The tier-0 (<= 148 rows) -> tier-1 promotion for envelopes just above the BLK=512 sampling gate follows upstream main's refinement of #18702. `max_seq_len` in windowed mode is the max window length and only picks the tier / bucket: the kernel clamps every window to the logits width, never to the hint. `warmup_prefill(top_k, max_cols)` compiles the tier x bucket set (idempotent, both tier edges of every bucket) for CUDA graph capture; `prefill_ready` reports it. Main kernel, all varlen modes: short rows (`n <= k`, block-uniform) now skip the body outright and go straight to the identity epilogue. Upstream runs the body as a zero-work pass (n = 0, TGT sentinel) — every barrier, the sample scan and, in FlashInfer's hint-free form, the identity bracket gather (`gather_hint(identity=True)`, two sites x two barriers) still execute, ~1 us per CTA on B200 at BLK=256. In a causal prefill staircase the first K rows of every request are short, so this is 25% of the rows at L = 8192, K = 2048 and all of them at L <= K. The hint-free identity bracket (`gather_hint(identity=True)`) samples from the first in-window lane (`base = lead`) rather than the float4-rounded row base: the <= 3 lead lanes belong to the previous request (or are unwritten), and sampling them widened the bracket to foreign values or NaN (sentinel fallback), 50% slower on misaligned windows at K = 512 than on aligned ones — upstream's compiled-out gather has no such dependence. Same-node A/B against the upstream TRT-LLM prefill kernel (identical plans, bit-identical results on every cell; SPEEDUP = upstream time / ours, min of 3 interleaved): K R x L B200 B300 Rubin 512 1 x 8192 1.37x 1.37x 1.35x 512 4 x 2048 1.79x 1.79x 1.72x 512 4 x 2047 (misaligned) 1.75x 1.79x 1.72x 512 4 x 8192 1.39x 1.39x 1.39x 512 8 x 4096 1.82x 1.85x 1.79x 512 16 x 2048 1.72x 1.72x 1.69x 512 1 x 32768 1.05x 1.05x 1.06x 512 64 x 512 (all rows <= K) 4.35x 4.55x 4.00x 512 128 x 64 (all rows <= K) 4.35x 4.35x 4.00x 2048 1 x 8192 1.23x 1.23x 1.20x 2048 4 x 2048 (all rows <= K) 3.23x 3.12x 2.86x 2048 4 x 2047 (all rows <= K) 3.23x 3.12x 2.86x 2048 4 x 8192 1.25x 1.25x 1.23x 2048 8 x 4096 1.79x 1.79x 1.75x 2048 16 x 2048 (all rows <= K) 2.86x 2.86x 2.78x 2048 1 x 32768 1.00x 1.00x 1.01x 2048 64 x 512 (all rows <= K) 2.86x 2.94x 2.78x 2048 128 x 64 (all rows <= K) 3.12x 3.12x 2.78x API (`topk_varlen.py`): `row_starts` validated as a contiguous 1-D int32 CUDA tensor with one entry per row, `next_n == 1`, `compress_ratio == 1`, `pre_idx=None`. Only gvr_2's checker admits windowed calls, so `auto` selects gvr_2 and an explicit other backend is refused at validation. The trace template gains the optional `row_starts` input with a window-aware reference and check. Perf (SGLang-shaped staircase, R requests x L prompt tokens, causal windows, fp32, graph replay, 3 interleaved repetitions, min, same node and job for every column; `proxy` = decode engine on a column-0 layout of identical per-row work, as the control): Time in us, min of 3 interleaved repetitions, graph replay. R requests x L prompt tokens, causal windows, fp32. B200 (148 SMs) K R x L rows radix_win gvr2_win sgl_v2 vs radix vs sgl_v2 512 1 x 8192 8192 426.1 93.7 120.6 4.55x 1.28x 512 4 x 2048 8192 221.8 62.8 71.2 3.57x 1.14x 512 4 x 2047 8188 221.8 62.8 72.8 3.57x 1.16x 512 4 x 8192 32768 1690.6 358.1 473.1 4.76x 1.32x 512 8 x 4096 32768 1289.1 349.5 379.4 3.70x 1.09x 512 16 x 2048 32768 972.3 251.9 314.7 3.85x 1.25x 2048 1 x 8192 8192 375.8 124.4 109.2 3.03x 0.88x 2048 4 x 2048 8192 49.7 19.4 19.0 2.56x 0.98x 2048 4 x 2047 8188 49.5 19.2 18.8 2.56x 0.98x 2048 4 x 8192 32768 1485.5 477.4 425.5 3.12x 0.89x 2048 8 x 4096 32768 895.5 333.1 276.7 2.70x 0.83x 2048 16 x 2048 32768 239.9 86.2 107.5 2.78x 1.25x B300 (148 SMs) K R x L rows radix_win gvr2_win sgl_v2 vs radix vs sgl_v2 512 1 x 8192 8192 412.4 90.6 115.7 4.55x 1.28x 512 4 x 2048 8192 212.2 60.2 68.1 3.57x 1.14x 512 4 x 2047 8188 212.2 60.3 69.8 3.57x 1.15x 512 4 x 8192 32768 1636.8 346.7 454.2 4.76x 1.32x 512 8 x 4096 32768 1243.3 337.1 364.3 3.70x 1.08x 512 16 x 2048 32768 934.6 242.5 301.5 3.85x 1.25x 2048 1 x 8192 8192 363.8 120.4 104.9 3.03x 0.87x 2048 4 x 2048 8192 46.0 18.6 18.5 2.50x 0.99x 2048 4 x 2047 8188 45.8 18.4 18.4 2.50x 1.00x 2048 4 x 8192 32768 1437.9 462.4 408.5 3.12x 0.88x 2048 8 x 4096 32768 861.9 320.7 264.3 2.70x 0.83x 2048 16 x 2048 32768 225.1 81.7 101.3 2.78x 1.23x Rubin (208 SMs) K R x L rows radix_win gvr2_win sgl_v2 vs radix vs sgl_v2 512 1 x 8192 8192 254.6 57.4 128.5 4.35x 2.22x 512 4 x 2048 8192 130.6 38.4 72.1 3.45x 1.89x 512 4 x 2047 8188 129.6 38.4 72.9 3.33x 1.89x 512 4 x 8192 32768 1001.3 215.2 503.8 4.76x 2.33x 512 8 x 4096 32768 775.4 211.2 437.3 3.70x 2.08x 512 16 x 2048 32768 590.3 153.8 387.6 3.85x 2.50x 2048 1 x 8192 8192 226.2 77.8 119.0 2.94x 1.54x 2048 4 x 2048 8192 32.7 13.5 28.2 2.44x 2.08x 2048 4 x 2047 8188 32.3 13.1 28.1 2.44x 2.13x 2048 4 x 8192 32768 884.0 287.8 462.3 3.03x 1.61x 2048 8 x 4096 32768 547.7 202.2 339.0 2.70x 1.67x 2048 16 x 2048 32768 156.9 53.5 176.2 2.94x 3.33x The remaining K = 2048 long-row cells where SGLang's kernel is 12-21% faster on B200 / B300 are the streaming main kernel's per-row cost on non-trivial rows (shared with upstream); Rubin has no losing cell. Tests (`tests/topk_varlen/test_topk_varlen_prefill.py`): out-of-window cells poisoned with NaN/+-inf/3e38 and the <= 3 lead lanes before each misaligned `ks` set to +inf (a leak becomes top-1, a missed frame correction a negative index); causal ramp for K in {512, 1024, 2048}; leads 1..3 with K in {512, 2048}; window lengths {0, 1, K-1, K, K+1}; all-equal and two-valued ties; fewer than K finite values (`-inf` masks) at every lead; `ks + len` past the width and `ks` at the width; `return_values`; reference engine parity; CUDA-graph capture/replay with changed `row_starts` and lengths between replays; unwarmed capture raises; 70000-row slab; prefill engine key distinct from decode; API validation and backend admission. tests/topk_varlen on the final tree: 370 passed on B200, B300, B100 and DRIVE P2021, 371 on Rubin (SM107), 83 on RTX 5080, 80 on A100 / L40S / H100 (rest skipped: no gvr_2) AI-assisted. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Description
Stacked on #18446. Extends the hint-free self-sampling GVR V2 indexer top-K
(introduced for decode by #18410/#18446) to the prefill phase, so a DSA
model that selects self-sampling for decode uses the same engine for prefill
instead of the CUDA radix path — one config, one kernel family, one warmup.
Dispatch is unchanged and needs no new config/API field: prefill uses the
self-sampling engine on exactly the layers where the two-level dispatch
(
enable_heuristic_topk×use_self_sampling_topk) already picks it fordecode. The temporal-hint (V1) layers keep the exact radix prefill (that
engine has no prefill form). The
TopK.forward(..., is_prefill=True, row_starts=..., row_ends=...)interface is unchanged.Kernel (
gvr_topk_decode_self_sampling.py)A constexpr
prefillmode onGvrMainKernelonly (the streaming family, whichis correct for any row length). Per row
rit selects the top-K oflogits[r, ks:ke](ks=row_starts[r],ke=row_ends[r], compressed columnunits), writes the LOCAL frame (
column - ks) with a trailing-1pad, andemits identity
0..nv-1whennv = ke-ks <= k— matching the radix short-rowcontract.
ks/keride the otherwise-unusedpre_idx/kv_lensABI slots(hint-free never reads them), so the kernel signature is byte-identical; the
prefill compile is a distinct cache key.
The base is rounded down to a 16B boundary (
col0 = ks & ~3) so the vectorloads stay aligned; the ≤3
lead = ks & 3lanes are masked positionally (aone-bit clear of the
i0==0classification mask, a valid-lane substitution inthe sample fold, and
i >= leadguards on the two whole-row escape sweeps),never with a materialized
-inf(which would drive the sample histogram indexout of bounds and mis-rank a genuine
-inftie class). The row-extent clampstighten to the last in-window float4, so the kernel reads exactly
[r*npad + col0, r*npad + ke)with no dependence on producer slack. Outputindices subtract
leadat the six write sites. Every edit isconst_expr-gatedon the prefill flag; the decode/legacy codegen is unchanged (the full decode
suite passes bit-for-bit).
Host / module / framework
run_prefill(logits, row_starts, row_ends, indices, max_row_len=None)— samevalidation style as
run_varlenbut keyed onstride(0)for all row counts(DeepGEMM prefill rows are 1024B-aligned with ≥256-float slack, so the
single-row
shape[1]rule that would reject odd-width tiles does not apply),no device reads (envelope from
shape[1]), and it slabs launches to ≤32768rows so
gridDim.ynever overflows. ForcesR=1(single CTA per row, noSPLIT/workspace); the launcher cache is keyed on
(row tier, k, envelope bucket)so arbitrary q-tile row counts do not grow it.warmup_prefill— compiles the ≤6 engines per k before serving (placedbefore the DeepGEMM decode-stride guard in
warmup_selfsampling_topk, so anodd
msl_ccannot skip it); best-effort under the existing OOM guard.TopK._forward_prefillgains aCUTE_DSL_GVRbranch: all-short tiles(
shape[1] <= k) take the exact radix identity path silently; the hardware-format gate (fp32, unit inner stride, float4 row stride, 16B base) routes to
the engine with a one-time
info_once; a miss falls back to radix with awarning_once.CUTE_DSL_RADIXstill raisesNotImplementedError.CUTE_DSL_GVRfor prefill iff it selects it for decodeand the self-sampling prerequisites hold; no
llm_args/API change.Prefill top-K is never CUDA-graph captured (
mla_dsa_attn_inplaceis apiecewise boundary op); there are no host syncs; PDL parity with the radix
prefill is none. Disaggregated context-only servers are covered by the
unconditional warmup hook. The port is compression-ratio agnostic (
ks/keare already compressed), so DeepSeek-V3.2 (cr=1, K=2048), V4 Flash (cr=4,
K=512) and V4 Pro (cr=4, K=1024) share one prefill path.
Validation
Kernel/host exactness against a tie-aware radix-parity reference on a B200
(SM100), per-row
[ks, ke)windows on DeepGEMM-strided storage:ks % 4in{1,2,3} and
+inf/NaN/3e38/-infpoison outside the window; short rows(
nvin {0, 1, k-1, k, k+1}); all-equal / two-valued degenerate paths;odd-width single-row tiles; the
-inftie class crossed with a misalignedlead (asserts no negative index leaks); a >65535-row slab. K in
{512, 1024, 2048}.
fallthrough to radix, temporal-mode fallthrough, output-width assertion.
buckets after warmup; the prefill compile key is distinct from the decode
varlen key.
layers keep radix; the warmup prefill leg fires (including for an odd
msl_c).GSM8K accuracy — end-to-end parity (B200 TEP8, 5-shot)
End-to-end
trtllm-eval gsm8k(tp8 ep8, one served config per arm). Anarmproofgate asserts each arm actually engaged the intended top-K path —decode/prefill engagement is grepped from the served logs (
viol=0), so the"V2 prefill" numbers are verified to run the new kernel, not silently fall back
to radix. The self-sampling prefill (this PR) matches both the exact radix
baseline and the pre-PR "V2 decode + radix prefill" state within GSM8K sampling
noise (stderr ±~0.6). Scores are
exact_matchaveraged over flexible-extract +strict-match, normalized to 0–100.
All three
pf_onruns are confirmedprefill_engaged=yes— the V2 prefillkernel actually ran and the score still lands on top of both baselines, so the
prefill top-K swap is accuracy-neutral across DSV3.2, V4-Flash and V4-Pro. Each
model runs its own
index_topk(Flash 512 / Pro 1024 / DSV3.2 2048).Preliminary kernel microbenchmark
A single-B200 cold-L2 CUDA-event sweep (random Gaussian logits, 20 reps/shape) pairs
the radix
topKPerRowPrefillagainst the DSL prefill kernel across 14 prefill shapes(DSv3.2 / V4, K in {512,1024,2048}, rows 64..32768, n_kv 4096..262144). The
self-sampling kernel is faster on every shape (median ~4×, worst 0.859×), so no
N-gate to radix is added. This is a gross-regression sanity check on random data;
the in-pipeline nsys numbers below are what govern the ship decision.
Operator-level top-K speedup vs the radix CUDA path
Test env: B200 (SM100), 8 GPUs,
tp8 ep8, PyTorch backend, chunked prefill(chunk=8192), KV fp8. An nsys A/B — pf_on GVR V2 selection kernel
GvrMainKernelvs pf_off radix
topKPerRowPrefill<K>— isolated by an exact kernel-nameset-difference (~103–110 kernel types per run; the two selection kernels are the
only differentiator, verified). Numbers are rank-0 GPU kernel duration of the
prefill window (NVTX-validated), dummy-AR excluded from the DEP8 denominators.
Per-prefill-window rank-0 kernel time (radix baseline → V2), then the speedup:
topKPerRowPrefill<K>GvrMainKernel(this PR)The V2 selection kernel is 1.84–2.61× faster than the production radix kernel
(e.g. Pro DEP8/990K: radix 3071 ms → V2 1177 ms = 2.61×). The ratio grows with
context/candidate count (radix
topKPerRowPrefillscales worse with K and candidates;V2 self-sampling is ~K-invariant — Pro K=512→1024 moved it only 2.68×→2.61×).
The share is config-dependent: for the same Flash it goes 0.60%→3.43% as ISL
grows (262K→990K) and attention-DP removes the ~33% NCCL comms from the denominator
(TEP8→DEP8). Because top-K selection stays ≤3.4% of prefill, the 1.84–2.61× kernel
win translates to a modest few-% end-to-end gain at throughput scale (serving results
below), not a large e2e speedup.
Where the prefill time goes (DEP8/990016, dummy-AR-excluded), rank-0 kernel
work-share — the dominant kernel is the indexer MQA logits, not the top-K
selection, and at DEP8 there is no real TP AllReduce:
Full per-kernel breakdowns + all trace paths:
indexer_topk_op_bench/op51_gvr_v2_prefill/GVR_V2_TOPK_E2E_AND_OPERATOR_REPORT.html.No DSV3.2 (cr=1) layered breakdown: its operator share is unmeasurable under nsys
in this environment — profiling the DSA algorithm under attention-DP deadlocks CUPTI
(verified it runs cleanly without nsys; TEP8+nsys separately hits an NVRTC
cuda.hgap). Its e2e numbers (below) are unaffected. By cr=1 scaling (4× candidatesvs cr=4) its indexer-logits + top-K shares would be higher than Flash/Pro.
End-to-end serving throughput (max-config: DEP8, long-ISL, batched, 3 reps)
Because top-K selection is only 2.4–3.4% of prefill, its end-to-end effect surfaces in a
throughput config, not a BS=1 latency one (BS=1 is TTFT-neutral — full table in the
report). This config maximizes the prefill top-K share: attention-DP (DEP8, no TP
AllReduce), each model's max ISL, batched. Test env: B200 (SM100) 8 GPUs,
tp8 ep8+enable_attention_dp=true, chunked prefill, KV fp8, OSL=512, 3-rep median, all armsarmproof-verified; concurrency = the max that fits (Flash/DSV3.2 32, Pro 8). The PR'schange is radix prefill (pf_off) → V2 prefill (pf_on) — both share V2 decode, so this
isolates the prefill top-K engine. Raw values first, then the speedup:
The V2 prefill port gives a +2.9–4.5% throughput / −3.0–4.4% TTFT end-to-end gain over
radix prefill across all three models — a real (if modest) benefit at scale, consistent
with the 2.6× kernel speedup applied to top-K's ≤3.4%-of-prefill share. (Against the exact
radix decode+prefill baseline, V2 overall is ~neutral here, 0.998–1.014× throughput — at
this high concurrency the prefill-interleaved pipeline masks the decode benefit; per-arm
tables in the report.)
Test configuration (max-config e2e)
B200 (SM100) × 8, PyTorch backend,
trtllm-bench throughput --streaming, upstream-HFFP4+FP8-mixed weights, 3 reps (median). Arms: radix (
enable_heuristic_topk=false) ·pf_off (V2 decode + radix prefill,
=true+TRTLLM_GVR_PREFILL_TOPK=0, an e2e-harness-only overlay switch that is not part of this PR) · pf_on(V2 decode + V2 prefill,
=true).tp8 ep8 pp1 cp1·enable_attention_dp=true(DEP8)max_position_embeddingsceiling) / 512dtype=fp8,tokens_per_block128 (V4)/64 (V3.2),free_gpu_memory_fraction=0.8,enable_block_reuse=falsemax_num_tokens=131072; V3.2backend=TRTLLMalgorithmdeepseek_v4/dsa;index_topk512/1024/2048; cr 4/4/1;window_size=128TRTLLM_ENABLE_DUMMY_ALLREDUCE=1(a benchmark sync artifact; excluded from the nsys shares)The BS=1 latency e2e (SWE-bench-64K + random, ISL 262144/65536, TEP8) — where prefill
top-K is TTFT-neutral and the V2 decode benefit shows in TPOT — is in the standalone
report.
Status
Done: kernel + host + module + indexer + metadata (decode codegen unchanged);
kernel/host exactness (39 cases) + prefill UT (28) + CPU dispatch + warmup
no-re-JIT. End-to-end validated on B200 (8 GPU) across DSV3.2 / V4-Flash / V4-Pro:
GSM8K parity (accuracy-neutral); nsys operator profile (top-K selection 1.84–2.61×
the radix kernel, ≤3.4% of prefill, MQA logits dominant); BS=1 serving TTFT-neutral;
and a max-config (DEP8, long-ISL, batched) throughput run where the prefill port adds
+2.9–4.5% throughput / −3.0–4.4% TTFT.
Rebased onto
mainafter #18446 merged (85d9871); the diff is the prefilldelta only. Follow-ups folded in after review: the warmup hook builds the DSA
metadata itself when no warmup forward ran (context-only / draft / guided-
decoder engines), an engine missed by warmup falls back to radix instead of
raising under CUDA graph capture, and GPU tests pin the DeepGEMM logits row
stride the engine's format gate relies on.
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.Dev Engineer Review
QA Engineer Review
test_dsa_indexer.py.test_fp4_mqa_logits_pass_selfsampling_prefill_format_gate.test_top_k.py.test_gvr_selfsampling_topk.py, including alignment, padding, short rows, ties, odd widths, slab limits, engine-key separation, guards, warmup idempotence, and capture safety.test_pytorch_model_engine.py.tests/integration/test_lists/,test-db/,qa/, orwaives.txtchanges were identified.