[TRTLLM-15293][perf] Add self-sampling (GVR V2) top-K decode kernels by longcheng-nv · Pull Request #17821 · NVIDIA/TensorRT-LLM

longcheng-nv · 2026-08-17T10:35:33Z

Summary

Self-sampling GVR (guess-verify-refine) top-K decode for the DSA indexer, written in CuTeDSL for Blackwell. This is the second-generation GVR operator, translated from an optimized CUDA implementation. The kernel estimates its threshold from a sparse sample of the current logits row, then verifies and refines. The previous step's top-K indices serve only as hints; the output is an exact top-K with ties interchangeable.

Two modules under tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/:

gvr_topk_decode_self_sampling.py — device code, four kernel families (main, reg, clus, reg_clus), JIT-compiled and cached per compile-time configuration.
gvr_topk_decode_self_sampling_host.py — dispatch (route(b, n, npad, k)), per-device workspace, and the run/run_ws entries. Standalone contract: batch-uniform host n_valid, fp32, K in {512, 1024, 2048}, row stride a multiple of 64.

run_varlen covers the production per-row contract: it reads kv_lens on device, derives each row's valid length from its MTP offset and compression ratio, handles short rows in-kernel, and runs the whole batch in one launch. With max_seq_len supplied, dispatch freezes before CUDA Graph capture and replay needs no host reads.

The DSA path turns on with TRTLLM_GVR_SELF_SAMPLING=1 when heuristic top-K is enabled, CUTLASS DSL is present on SM100+, and the K/compression/layout gates pass; otherwise dispatch falls through to the existing paths. Enabled means self-sampling everywhere: every decode batch size, envelopes to 1M kv tokens, no rows-based fall-through. Each row derives its own length and hints (heterogeneous-kv_lens oracle tests pin this), the dispatcher picks main or the per-row reg_clus port from the capture-stable launch key, and warmup compiles one representative row per engine key. TSH-floor staging now also covers small batches; the old num_rows > 15 gate left them a ~6x tail on deep-layer captures.

Kernel-level results on real captures (cold-L2, tie-aware exact on every case): 2.9-4.7x over the in-tree production dispatch at N=262144 for 64-1024 rows, 3.1-3.7x on the former v4_pro_512k deep-layer pathology band, 1.6-3.5x on DSv3.2 163k. Short rows (N <= 2k) at high row counts run 0.67-0.94x at microsecond absolute scale — an accepted trade-off. Hints are consumed without the +1 offset the CUDA path applies at compressRatio == 1; raw hints overlap the current top-K better on real DSv3.2 captures (0.773 vs 0.536).

How to enable (GVR V2) — minimal setting per model

Enabling is exactly one environment variable + one YAML field. Set the environment variable before process start (with mpirun or multi-node launchers, export it to all ranks, e.g. -x TRTLLM_GVR_SELF_SAMPLING):

export TRTLLM_GVR_SELF_SAMPLING=1

With the YAML alone (environment variable unset), dispatch stays on the in-tree GVR path that ships in main today: the CUDA heuristic indexer_topk_decode by default, or CuTeDSL GVR (#16877) when use_cute_dsl_topk: true. Setting TRTLLM_GVR_SELF_SAMPLING=1 promotes the self-sampling engine to the front of that dispatch chain; everything else (hint buffers, YAML, scheduling) is unchanged. This is also why the accuracy A/B below is a strict V2-vs-V1 comparison: both arms ran the identical YAML, and the baseline arm is the production GVR path.

DeepSeek-V3.2 (`deepseek-ai/DeepSeek-V3.2-Exp`)

Minimal extra.yaml:

sparse_attention_config:
  algorithm: dsa
  enable_heuristic_topk: true   # required: allocates the prev-top-K hint buffers V2 consumes
speculative_config:             # optional; canonical MTP for V3.2
  decoding_type: MTP
  num_nextn_predict_layers: 1

TRTLLM_GVR_SELF_SAMPLING=1 trtllm-serve deepseek-ai/DeepSeek-V3.2-Exp \
  --config extra.yaml --tp_size 8 --ep_size 8 --port 8000

DeepSeek-V4 Flash (`deepseek-ai/DeepSeek-V4-Flash`)

Minimal extra.yaml:

sparse_attention_config:
  algorithm: deepseek_v4
  enable_heuristic_topk: true
speculative_config:             # optional; canonical MTP for V4
  decoding_type: MTP
  num_nextn_predict_layers: 3

TRTLLM_GVR_SELF_SAMPLING=1 trtllm-serve deepseek-ai/DeepSeek-V4-Flash \
  --config extra.yaml --tp_size 8 --ep_size 8 --port 8000

DeepSeek-V4 Pro (`deepseek-ai/DeepSeek-V4-Pro`)

Same YAML as Flash (algorithm: deepseek_v4, canonical MTP 3). Pro needs TP8 on 8xB200 (~806 GB weights):

TRTLLM_GVR_SELF_SAMPLING=1 trtllm-serve deepseek-ai/DeepSeek-V4-Pro \
  --config extra.yaml --tp_size 8 --ep_size 8 --port 8000

Auto-checked prerequisites, verification, notes

Everything else is checked by the dispatch gate and falls through to the in-tree top-K path with a one-time warning when unmet — none of it is a user setting:

Datacenter Blackwell (sm_100/sm_103: B200/B300 class) with CUTLASS DSL available (nvidia-cutlass-dsl).
index_topk in {512, 1024, 2048} and compression ratio in {1, 4} — both come from the model config; all three models above satisfy them.
fp32 row-major decode logits with a float4-aligned row stride and 16-byte base (the paged-MQA arena layout satisfies this).

Verifying engagement: the log prints self-sampling GVR top-K engaged (K=..., cr=..., next_n=...) once on first decode — or a one-time falling through to the in-tree top-K path warning when the layout gate rejects — and Nsight Systems shows the self_sampling kernel symbol during decode.

Notes: trtllm-bench works the same way (TRTLLM_GVR_SELF_SAMPLING=1 trtllm-bench --model <id> throughput --backend pytorch --dataset <jsonl> --tp 8 --ep 8 --extra_llm_api_options extra.yaml); the same holds for trtllm-eval, which is what the accuracy runs below used. The path only affects generation-phase indexer top-K; prefill is unchanged. CUDA graphs and MTP are fully compatible (accuracy-validated at MTP 0/1/3, graphs and eager); kernels are pre-compiled during engine warmup, so serving pays no first-touch JIT. use_cute_dsl_topk does not need to be set.

End-to-end decode TPOT (serving A/B)

trtllm-bench serving pairs on 8xB200 (TEP8, batch size = concurrency = 1, SWE-bench 64K prompt, ISL 68,656, OSL 1024, streaming). Both arms run the identical engine config; the only difference is TRTLLM_GVR_SELF_SAMPLING. Baseline is the production default top-K (radix CUDA). Arms run back-to-back within each repetition; reductions are paired per repetition (mean +/- sample stdev).

model	MTP	baseline TPOT (ms)	self-sampling TPOT (ms)	TPOT reduction	pairs
DSv4 Flash	0	6.254	5.622	-10.1% ± 0.7	10
DSv4 Flash	1	3.684	3.364	-8.7% ± 0.9	10
DSv4 Flash	3	2.934	2.709	-7.6% ± 1.1	9
DSv4 Pro	0	10.481	9.537	-9.0% ± 0.4	10
DSv4 Pro	1	6.171	5.667	-8.2% ± 0.7	9
DSv4 Pro	3	4.988	4.682	-6.1% ± 2.7	10
DSv3.2	0	11.908	9.728	-18.3% ± 0.4	7
DSv3.2	1	7.219	6.069	-15.9% ± 0.7	10
DSv3.2	3	5.660	4.907	-13.3% ± 1.0	7

The 10-repetition campaign (MTP {0, 1, 3} x {baseline, self-sampling}) is complete for both DSv4 models. Two DSv4 cells report 9 pairs: in each, one repetition lost its self-sampling run to a transient MPI worker exit during engine setup (unrelated to the kernel; the paired baseline run is dropped from the stats). The DSv3.2 campaign is complete. Its MTP-1 rows aggregate two hosts whose per-rep values agree within ~2% (baseline 7.12-7.33 ms across both); MTP-0/-3 report 7 pairs because one host repeatedly showed 40%+ TPOT inflation with heavy jitter from co-located load, and all of its pairs were discarded wholesale, both arms alike, leaving the single-host clean set. Engagement is proven per run via the self-sampling GVR top-K engaged log marker.

Correctness and performance summary

On B200, the CuTeDSL implementation was compared with its self-sampling CUDA source using the full real-decode-capture grid: 886 shapes at 11 batch sizes, or 9,746 cases. All cases were exact. The geometric mean of CuTeDSL time / CUDA time was 0.974.

The per-row engine was also checked against the batch-uniform reference on heterogeneous, MTP, and short-row batches, and under CUDA Graph capture/replay. The offset-free hint contract has a separate 9,746-case A/B comparison between raw and +1-shifted hints; the V3.2 segment was performance-neutral, with a geometric mean ratio of 0.9998.

Performance comparison

The table below covers the 886-shape x 11-batch-size real-decode-capture grid (9,746 cases). Inputs come from DSv3.2, DSv4 Flash, and DSv4 Pro captures and are byte-identical across arms. Measurements were collected on B200 with FP32 inputs, cold L2, kernel-only NVTX sums, and at least 10 cold repetitions per case.

Ratios are time(comparison arm) / time(self-sampling CuTeDSL), so a value above 1 means self-sampling CuTeDSL is faster.

Comparison arm	Cases	Geometric mean	Min	P5	P95	Max	Self-sampling faster
TRT-LLM radix CUDA (`indexer_topk_decode`, production)	9,746	5.49x	1.58	2.12	11.36	31.5	100%
TRT-LLM GVR CuTeDSL (#16877)	9,746	1.62x	0.70	1.13	2.62	7.0	99.6%
SGLang v2 (plan + transform, complete operator)	9,746	1.91x	0.95	1.39	3.69	11.1	99.9%
FlashInfer 0.6.14 `top_k`	9,746	2.35x	1.24	1.58	4.46	9.9	100%
TRT-LLM radix CuTeDSL (`single_pass_multi_cta`)	9,746 (same-run)¹	2.64x	1.24	1.35	5.69	6.9	100%
Self-sampling CUDA (translation source, same-run A/B)	9,746	1.03x	0.87	0.96	1.11	1.22	74.1%²

Per-segment geometric means:

Radix CUDA: Flash 5.24 / Pro 5.18 / V3.2 5.82
GVR [TRTLLM-15293][perf] Add tiered GVR CuTe DSL top-k decode kernels (stacked on #16457) #16877: Flash 1.71 / Pro 1.70 / V3.2 1.53
SGLang v2: Flash 1.97 / Pro 1.93 / V3.2 1.88
FlashInfer: Flash 2.33 / Pro 2.34 / V3.2 2.36
Radix CuTeDSL: Flash 2.96 / Pro 2.84 / V3.2 2.40

For radix CuTeDSL, the geometric-mean speedup grows with batch size over the full 886-shape grid: 2.19 at BS=1, 2.16 at BS=8, 2.60 at BS=32, 2.79 at BS=64, 3.30 at BS=256, 3.75 at BS=512, and 3.96 at BS=1024. The local value at BS=128 is 2.54.

Ratio distribution by case count:

Compared with	<0.7	0.7-0.9	0.9-1.0	1.0-1.2	1.2-1.5	1.5-2.0	2.0-3.0	>=3.0
Radix CUDA (9,746)						266	834	8,646
GVR #16877 (9,746)	1	17	19	882	3,090	3,940	1,589	208
SGLang v2 (9,746)			10	51	1,218	5,516	2,010	941
FlashInfer (9,746)					135	2,989	5,091	1,531
Radix CuTeDSL (9,746, same-run)					610	1,994	3,972	3,170

Measurement notes

Exactness: every case for every arm is tie-aware exact, comparing the value multiset with torch.topk. The self-sampling grid has 0 inexact results out of 9,746 cases.
Operator scope: the table measures the standalone batch-uniform entry (run) — the full four-family dispatch (main/reg/clus/reg_clus) that the CUDA source implements, which is also the contract every comparison arm uses. The DSA-wired run_varlen engine currently restricts per-row dispatch to the streaming main family (the one family that is correct for any row length). On this grid, route picks a specialist family on 54% of cases, amounting to 21% of total kernel time — concentrated in short rows; within the N ≥ 64K deployment band the specialist share is 6.9% of kernel time. On those shapes the wired path runs R-split main instead, so serving-path times there sit somewhat above this table; the clustered-family port below recovers it.
The self-sampling CUDA-to-CuTeDSL row is a same-run paired A/B measurement from 2026-08-17 on umbriel-b200-027.
The other baseline times come from earlier campaigns on the same captured inputs: radix CUDA and [TRTLLM-15293][perf] Add tiered GVR CuTe DSL top-k decode kernels (stacked on #16457) #16877 were measured on 2026-07-30 on umbriel-b200-027 in a quiet single lane; SGLang v2 and FlashInfer used the V3.2 segment from 2026-07-30 on umbriel-b200-027 and the V4 segment from 2026-07-31 on umbriel-b200-048, an equivalent B200 host. Those ratios join absolute microsecond measurements by (shape, batch size) and are therefore cross-run comparisons, not same-session pairs.
The cross-run numbers are internally consistent. The same-run [TRTLLM-15293][perf] Add tiered GVR CuTe DSL top-k decode kernels (stacked on #16457) #16877-to-radix ratio from 2026-07-30 is 3.394x, and 3.394 x 1.62 ~= 5.49, matching the direct join. The radix CuTeDSL row similarly implies a radix-CUDA-to-radix-CuTeDSL ratio of approximately 5.49 / 2.64 = 2.08x.

CUTLASS DSL 4.6.1 compatibility and production-varlen certification

main moved to CUTLASS DSL 4.6.1 (#17274) after the table above was measured on 4.5.0. This PR is ported to 4.6.1 (make_rmem_tensor API) and re-certified on it. Two DSL 4.6.1 performance regressions were found on the way and are worked around inside the kernel module (both reported to the DSL team with PTX/ncu evidence):

v2.b64 pair-load emission: 4.6.1 rewrites adjacent 128-bit f32 loads into .b64 register-pair loads; the even-aligned pairs fragment register allocation at the 64-register wall and spill (clustered family +20% kernel time). Worked around by pinning the row loads to ld.global.nc.v4.f32 via inline assembly.
Compile-time smem carveout ignoring dynamic smem: with min_blocks_per_mp > 1 the 4.6.1 compiler derives a shared-memory carveout from static smem only, selecting a 16 KiB config that pins the register family to one resident CTA per SM (achieved occupancy 67.6% → 11.6%, up to 3.25x slower at saturating grids). Worked around by extending the existing _no_carveout() compile scope to drop the new attribute; the other families launch with min_blocks_per_mp == 1 and were verified neutral.

With both workarounds, the full 886-shape x 11-batch grid was re-run on the PR head under DSL 4.6.1 through the production varlen entry (run_varlen, per-row lengths from device kv_lens), same-process paired against the standalone entry, on 3x8 B200:

Exactness: 9,746/9,746 tie-aware exact.
DSL 4.6.1 vs 4.5.0: cross-run geometric mean 1.00 (0.9975) on identical shapes — the 4.6.1 port is performance-neutral end to end.
Production varlen vs SGLang v2 (plan + transform, cross-run join by shape and batch size, 9,515 joinable cases): geometric mean 1.64x, self-sampling faster in 99.89% (9,505/9,515) of cases; without the carveout workaround this was 98.07% with the losses concentrated at short rows x large batch.
Varlen integration tax (production varlen vs standalone, same process): geometric mean 1.12; within the N >= 64K deployment band 1.10.

End-to-end accuracy A/B — COMPLETE: 12/12 full-dataset pairs, all passing

Protocol: one serving stack per model, identical config for both arms; the arm is toggled only by TRTLLM_GVR_SELF_SAMPLING. Arm identity proven per model AND per host by nsys kernel tables in real serving (5 host-model probes confirmed; e.g. 1,440 self-sampling launches, median 7.3 us, during V4 Pro TP8 serving). Runs execute on a validated rc21 serving base with this PR's kernels + dispatch seam grafted verbatim; a PR-head build pass follows CI.

Full-dataset pairs (strict-match / weighted; each pair same host + same TP)

Model	Task (samples)	MTP (kernel `next_n`)	base (production CUDA GVR)	self-sampling (this PR)	delta	verdict
V4 Flash (TP4)	GSM8K (1,319)	MTP3 (4)	96.44 ± 0.51	96.29 ± 0.52	−0.15pp	pass
V4 Flash (TP4)	GSM8K (1,319)	off (1)	96.59 ± 0.50	96.74 ± 0.49	+0.15pp	pass
V4 Flash (TP4)	MMLU (14,042)	MTP3 (4)	88.51	88.51	0	pass
V4 Flash (TP4)	GPQA-Diamond (198)	MTP3 (4)	73.74 ± 3.14	72.22	−1.5pp (within SE)	pass
DSv3.2 (TP4)	GSM8K (1,319)	MTP1 (2)	96.36 ± 0.52	96.06 ± 0.54	−0.30pp	pass
DSv3.2 (TP4)	GSM8K (1,319)	off (1)	96.36 ± 0.52	96.06 ± 0.54	−0.30pp	pass
DSv3.2 (TP8)	MMLU (14,042)	MTP1 (2)	87.77	87.74	−0.03pp	pass
DSv3.2 (TP8)	GPQA-Diamond (198)	MTP1 (2)	73.23	76.77	+3.5pp (within SE)	pass
V4 Pro (TP8)	GSM8K (1,319)	MTP3 (4)	96.82 ± 0.48	96.66 ± 0.49	−0.15pp	pass
V4 Pro (TP8)	GSM8K (1,319)	off (1)	96.44 ± 0.51	96.66 ± 0.49	+0.23pp	pass
V4 Pro (TP8)	MMLU (14,042)	MTP3 (4)	89.67	89.67	0	pass
V4 Pro (TP8)	GPQA-Diamond (198)	MTP3 (4)	74.75	75.25	+0.5pp	pass

Deltas are symmetric around zero (6 negative, 4 positive, 2 exact ties) and all within standard error — the statistical signature of pure scheduling noise. MMLU ties at 14k samples on two models (and −0.03pp on the third) are the strongest single pieces of evidence: with exact top-K indices, greedy decoding has nowhere to diverge.

MTP coverage

Speculative decoding (MTP) is a first-class axis — the kernel's per-row MTP window (n_r = (kv_len - next_n + (r % next_n) + 1) // compress_ratio, request-level shared hints) is exercised end to end:

Every model ran both MTP-off and its canonical MTP at full-dataset scale, collectively covering the kernel's full supported next_n set {1, 2, 4} in real serving.
Greedy MTP-invariance holds on both arms: DSv3.2 per-arm scores at MTP0 and MTP1 agree to 4 decimals; V4 Pro self-sampling scores at MTP0 and MTP3 agree to 4 decimals (96.6641); Flash within 0.15pp per arm.
Unit tests additionally cover next_n in {1, 2, 3, 4} x compress_ratio in {1, 4} with in-request row-length variation, mixed batches, and CUDA-graph capture/replay under MTP (next_n=3/MTP2 has no production config; verified to generalize by a dedicated unit-test pair).

Notes

Each (base, ss) pair ran on one host at one TP (annotated per row). DSv3.2 MMLU/GPQA pairs use TP8: the 643 GB checkpoint at TP4 leaves ~161 GiB/GPU for weights and MMLU's request packing tripped a VRAM OOM — an environment constraint, not a kernel issue.
V4 Pro rows ran eager: the baseline (production) arm hits a CUDA illegal-memory-access during MTP-3 CUDA-graph warmup on the rc21 base (Flash MTP-3 + graphs is fine). Tracked as an independent serving issue, unrelated to this PR; accuracy is launch-mode-independent.
¹ The radix CuTeDSL row is a same-run paired A/B comparison over the full grid, measured on 2026-08-18 on umbriel-b200-038 with eight idle-host lanes. Both arms were exact in all 9,746 cases. On the earlier 594-case V4 subset, this run reproduces the retired cross-run result (2.88 versus 2.89 geometric mean; for the 459 BS=1 cases, 2.42 versus 2.45). Self-sampling absolute times on hosts 038 and 027 also agree by (shape, batch size), with a geometric mean ratio of 1.002 (P5 0.956 / P95 1.051). The earlier partial-coverage results (2.89 overall and 2.45 at BS=1) are superseded.
² For the near-parity CUDA row, the geometric mean is more meaningful than the per-case win rate. The harness's A/A noise floor is approximately 1.05 at P99, and 81% of cases in which CUDA appears faster fall in (0.96, 1.0]. The isolated-confirmed CUDA advantage (era-4 funnel FINAL, 2026-08-19: grid -> bulk x3 -> serial confirm) spans 835 of 9,746 cases across 47 (family, BS) groups — dominated by the Pro-64k / Flash-64k / Pro-16k / V3.2-4k small-batch micro-kernel bands — with CuTeDSL at most 12.1% slower (grid-wide GM 0.973, i.e. CuTeDSL faster overall); the 106 rep-spread SUSPECT cases were re-arbitrated on quiet lanes (2026-08-19, temp-guard clean): 17 washed, leaving a retest-confirmed envelope of 819 cases / 47 (family, BS) groups, max 1.121. NCU PC sampling attributes this to barrier-arrival imbalance in the generated count-round code: CuTeDSL executes about 4% fewer instructions but spends about 20% more time stalled at the count-round BAR.SYNC. This remains a code-generation optimization item.
The self-sampling grid was measured under concurrent multi-GPU fleet load, which can only increase its own measured times. The cross-run speedup ratios are therefore conservative for self-sampling.
All V3.2 rows use captures regenerated on 2026-07-29, so the inputs match exactly.

Implementation status and follow-up work

Landed in this PR (reviewed incrementally while it was a draft):

Test coverage

tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py is included in the existing unittest/_torch/thop/parallel sweep in l0_b200.yml. The tests require SM100 and skip when CUTLASS DSL is unavailable. CI cost is negligible, measured from the L0 #56075 testReport: the file runs as 84 parameterized cases on each of three Blackwell shards (DGX-B200, B300, GB300-4GPU) at ~32 s per shard — 252 case instances, 97 s aggregate, all passing. The slowest case is 3.2 s; the rest are sub-second (kernels JIT once, later cases hit the compile cache). Non-SM100 shards skip at collection time, adding no cost.

Coverage includes:

tie-aware exactness using signed-zero-normalized value-multiset comparison against torch.topk;
9 (K, N) shapes at batch sizes 1 and 4, including the K=2048 dispatch boundaries at N=131075/131076 and the deployment maximum N=262144;
padding poisoned with 3e38, so any read beyond n_valid fails correctness;
uniform and per-row variable lengths, heterogeneous and short-row batches, MTP widths 2 and 4, and compression ratios 1 and 4;
differential checks between the in-kernel and reference variable-length engines;
CUDA Graph capture/replay with changing kv_lens;
run_ws with caller-owned workspace;
optional value output and production-compatible short-row padding;
degenerate previous-index hints; and
guard predicates for dtype, negative n_valid, batch-size mismatch, invalid kv_lens, non-contiguous inputs, and unsupported compression ratios, plus dispatch totality over the supported envelope.

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-compatible or api-breaking. For api-breaking, include BREAKING in 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.

Summary

Added opt-in Blackwell CuTeDSL GVR self-sampling top-K decode kernels.
Added reg_clus variable-length dispatch, ABI handling, and CUDA Graph support.
Added routing, workspace management, validation, launch caching, warmup, and optional value output.
Added run, run_ws, and run_varlen APIs.
Added support for device-side row lengths, MTP windows, compression ratios, short rows, zero-window rows, and packed buffers.
Added a float4-aligned width gate for single-row decode logits.
Warmed all admissible row-count and batch-size combinations.
Enabled TSH-floor staging for split launches with k <= 1024.
Removed a redundant compile-key slot.
Integrated CUDA graph batch sizes into ModelEngine warmup.
Integrated the path through TRTLLM_GVR_SELF_SAMPLING=1.
Restricted dispatch and warmup to validated sm_100 and sm_103 platforms.
Reported 9,746 exact kernel cases and 12 passing end-to-end accuracy cases.
Reported a CuTeDSL/CUDA geometric-mean time ratio of 0.974.

Dev Engineer Review

Host dispatch validates supported configurations before selecting the self-sampling path.
reg_clus variable-length launches receive dedicated compilation parameters and ABI dispatch.
Warmup filters unsupported widths, row counts, hardware configurations, and row strides.
Warmup catches CUDA OOM and permits lazy JIT compilation.
run_varlen supports device-side kv_lens, MTP offsets, compression ratios, short rows, optional values, explicit workspaces, and CUDA Graph-safe execution.
Single-row dispatch rejects unsupported non-float4-aligned widths.
Existing CuTeDSL and in-tree top-K implementations remain available as fallbacks.
Follow-up scope includes BF16/FP16 support, the clustered non-register family, preIdxCount != K, eager launcher bucketing, and PDL attribute parity.
No configuration or test-list changes are reported.

QA Engineer Review

Test code changed in tests/unittest/_torch/thop/parallel/test_gvr_selfsampling_topk.py.

Added tests cover:

Zero-window variable-length rows.
Warmup row-stride compatibility.
The supported row-count envelope.
Warmup filtering above the supported envelope.
Clustered-register variable-length parity and oracle results.
Clustered-register CUDA Graph replay.

The test file also covers exactness, short rows, optional values, varlen and MTP handling, engine parity, launch modes, packed buffers, validation, CUDA Graph replay, explicit workspaces, dtype enforcement, routing, and warmup.

No corresponding tests/integration/test_lists/, test-db/, or qa/ coverage entries are provided. Verdict: needs follow-up.

…ndalone) Add a self-sampling variant of the GVR (Guess-Verify-Refine) heuristic top-K decode for the DSA indexer, translated to CuTeDSL from the optimized CUDA line (fork branch GVR-selfsampling-CuTeDSL), as two standalone modules under cute_dsl_kernels/blackwell/top_k: - gvr_topk_decode_self_sampling.py: merged device module — four kernel families (sampling-ladder main / register-resident reg / cluster clus / cluster-register reg_clus), lazily JIT-compiled per constexpr tuple. - gvr_topk_decode_self_sampling_host.py: host companion — pure-function dispatch route(b, n, npad, k) (bit-exact transcription of the CUDA host dispatch, cross-checked by a 1,159,168-case boundary+fuzz sweep plus a 300k-case parity fuzz of this merged form), per-device workspace slab, and run/run_ws DPS entries with the CUDA binding's hardening battery. Contract (documented in the host module): batch-uniform host-int n_valid in compressed index space; fp32; K in {512, 1024, 2048}; 64-element-multiple row stride. Exact (tie-interchangeable) top-K. NOT wired into the decode path: the production engine reads per-request seq_lens on-device with per-row MTP offsets (heuristicTopKDecode.cu); adopting that per-row contract inside these kernels is follow-up work, so this module must not substitute for the tiered path under continuous batching, MTP, or CUDA-graph capture. Evidence (B200): 886-cell x 11-BS real-decode-capture grid = 9,746 cases vs the production CUDA arm: 0 INEXACT, geomean ratio 0.974. Unit tests: tie-aware exactness (signed-zero normalized, poisoned-pad immunity) across gate-edge and envelope shapes, run_ws with caller workspace, guard predicates, dispatch totality; sm_100-gated, picked up by the existing unittest/_torch/thop/parallel sweep. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>

… raw indices, uniform across DSv3.2/Flash/Pro The contract docstring previously instructed callers to apply the cr==1 +1 temporal shift to pre_idx (mirroring heuristicTopKDecode.cu). Drop it: hints only steer the sampling ladder — exactness never depends on them — and on real V3.2 decode captures raw prev-step hints overlap the current top-K at 0.773 vs 0.536 when +1-shifted (15 cells x 14 consecutive step-pairs, gap widening with ISL). One offset-free hint convention now serves all three models; the kernels already consume pre_idx as-is, so this is a contract-documentation fix only. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>

…ction pad convention) When the row has no more than topK valid entries every valid position is in the top-K: emit identity indices and pad the tail with -1, mirroring heuristicTopKDecode.cu:72-84. Host-level torch-op branch for the standalone module (the CUDA-graph-safe per-row rewrite will move it in-kernel, where a per-row fallback is impossible inside a graph). Closes the 'n <= topK unproven' gap from the integration audit. Tests: 8 boundary shapes (n in {64..2048}, k in {512,1024,2048}, n < k / n == k-1 / n == k) x bs {1,4} with poisoned padding, plus kernel-path regression just above the boundary (k512_n4099, k2048_n4111) verified exact on B200. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>

…lt off) Optional `values` DPS output on run()/run_ws(), production parity for the heuristicTopKDecode values writeback. Default None = OFF, matching dsa.py, which allocates the values scratch only for the non-CuTeDSL path. The indices are exact, so a gather epilogue reproduces the in-kernel writeback bit-for-bit at zero cost when disabled; the constexpr in-kernel form rides the CUDA-graph per-row rewrite. Short path pads values with -FLT_MAX (production convention). Tests: kernel path (values == gathered top-K == torch.topk multiset), short path (head copies logits, -FLT_MAX tail), wide-buffer packed re-view and dtype guard, verified on B200. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>

…ests Production can hand the kernel degenerate hint buffers: dsa.py initializes heuristic_prev_topk with zero_() (all-zero cold start; the prefill->decode seeding covers the common path but zero-init corners remain), and nothing forbids duplicated hints. Exactness must never depend on hint quality (hint-robustness bug class of PR NVIDIA#17550). Adds all-zero / all-same / all-max / half-duplicated hint cases on (k512, n8192) and the k2048 gate-edge n131075 — all verified exact on B200 (also probed on n131072/k1024 x bs{1,4} pre-commit, 24 configs, zero failures). Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>

…act entry (reference engine) Adds run_varlen(logits, pre_idx, kv_lens, indices, next_n=, compress_ratio=, values=None) — the heuristicTopKDecode contract: per-request device kv_lens (TOTAL cache length, uncompressed token space), per-row n = (kv_len - next_n + row%next_n + 1)/compressRatio (the MTP window formula, cr 1 = DSv3.2 / 4 = DSv4 Flash+Pro), request-level raw pre_idx shared by a request's next_n rows, per-row n <= k short path. REFERENCE engine: one documented host read of kv_lens (raises under CUDA-graph capture), rows driven as b=1 launches through the batch-uniform engine. This pins the varlen/MTP contract and its test battery; the per-row in-kernel engine (device kv_lens reads, fixed-R thin-slicing, n-band parameter table) replaces the loop next without changing either. Tests: heterogeneous lengths with per-row poisoned padding, cr in {1,4} x next_n in {1,2,4} incl. in-request n variation and compressed-boundary-crossing rows, mixed short rows, values output, contract guards — all verified exact on B200. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>

…it (route_static/route_dynamic) Factors route(b, n, npad, k) into route_static — family, compile tuple, grid/cluster/block and the rt scalars that change only at discrete n-thresholds (freezable at CUDA-graph capture time, mirroring the in-tree runner's pick_tuning(graph_capture=...) pattern) — and route_dynamic — the n-continuous scalars a per-row kernel recomputes from its own row length (CMP and the reg smem footprint; the SMP/TGT/SS2/TGT2/Q sampling ladder for the streaming families). The device-side per-row engine will mirror exactly the route_dynamic formulas. Lossless by construction and by fuzz: recombining the halves reproduces route() bit-exactly on 163,755 (b, n, k) points (threshold windows, R-boundaries, prime-stride sweep, LCG random; full result-dict equality). route_bands() enumerates the static-constant n-intervals — the (b=8, k=1024, 262144) envelope collapses to 10 contiguous bands, so the eventual in-kernel band table is tiny. Host-side groundwork only; no kernel behavior change. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>

…gine (gvr_main port) The gvr_main family gains a per-row varlen mode (production heuristicTopKDecode contract): each CTA reads its row's kv_len from a device kv_lens tensor and re-derives n plus the whole sampling ladder (SMP/SS2/TGT/TGT2/Q) via an exact device transcription of the route_dynamic() host formulas — integer round(sqrt(6n)) with an isqrt fixup (bit-parity with the host double form), Int64 TGT products (the CUDA host math is 64-bit), and constexpr next_n / cr_shift / r_const so every division strength-reduces. All values are pure functions of the row index, so the R split CTAs of a row stay grid-uniform by construction (workspace handshake unchanged). Design points: - No runtime return in CuTe DSL (in-tree gvr_topk_decode.py precedent): n <= k rows run the body as a zero-work pass (n=0, TGT=INT_MAX so no rung accepts, Q=0 empties the split slices) and an epilogue emits the production identity/-1-pad short path. - The TSH-floor staging gate becomes per-row runtime (tsh_en && per-row n4 <= 32768), exactly the CUDA original's grid-uniform runtime gate; legacy compiles keep bit-identical behavior (tsh_run == 1). - pre_idx is REQUEST-level in varlen mode (row // next_n mapping). - run_varlen(engine="auto") launches the batch in ONE kernel; with max_seq_len given (capture-stable engine constant) the call performs no host reads. engine="reference" keeps the b=1 loop as the differential oracle. Capture-time tuning comes from route_streaming() (the streaming half of route(), 110,003-point fuzz agreement). - Legacy batch-uniform ABI extended with dead trailing args (dummy kv_lens + five zeros) — one kern body, no duplication. Validated on B200: legacy regression 5/5 exact (extended ABI), engine vs reference differential 9/9 mixed-batch configs row-for-row equal — deep SPLIT (b=1, n=200k), 8-row cr=4 mix to n=225k, tsh band (b=16), MTP next_n {2,4}, all-short batches, 200-row small_dense k=2048, and the no-host-read max_seq_len path. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>

…est for the varlen engine Warm up, capture one run_varlen(engine=auto, max_seq_len=...) launch — no host reads, no JIT inside capture — then replay while kv_lens grows in place: a row crossing the n <= topK short-path boundary INSIDE the graph, a row walking the 131072 band edge, and a 200k deep row. Every replay verified tie-aware exact on B200 (8-replay standalone run all green). This closes the CUDA-graph-safety design goal for the gvr_main varlen port: geometry and tuple are frozen from capture-stable quantities (route_streaming at max_seq_len), all N-dependence is per-row device arithmetic. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>

…TRTLLM_GVR_SELF_SAMPLING=1) Wires the varlen engine into the DSA indexer decode top-K dispatch as the highest-priority branch, env-gated (TRTLLM_GVR_SELF_SAMPLING=1) and contract-gated at init (cutlass DSL present, sm100+, index_topk in {512,1024,2048}, compress_ratio in {1,4}) — covering DSv3.2 (K=2048, cr=1), DSv4 Flash (K=512, cr=4) and DSv4 Pro (K=1024, cr=4). The call reuses the exact buffers of the existing cute_dsl_gvr_topk_decode branch (request-level heuristic_prev_topk, kv_lens_cuda_runtime, topk_indices_buffer, indexer_max_seq_len as the capture-stable tuning constant) — no new metadata plumbing. Lazy import behind the gate; contract violations raise loudly (explicit experiment flag, not a silent fallback). First warmup call pays the one-time DSL JIT. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>

…ive the TGT*2 scan target The zero-work short pass used TGT = 0x7FFFFFFF, whose TGT*2 third scan target overflows Int32 to -2, flipping every tot0 >= TGT*2 gate on the all-zero histogram (benign downstream today — empty candidate sets emit nothing — but an unnecessary cliff). Use 2^30-1 so the doubled target stays positive; 'never accepts' semantics unchanged. Full differential battery (9/9 mixed-batch configs) + graph capture/replay (8 replays) re-verified green on B200. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>

… of the varlen engine and seam Fixes from a 4-dimension adversarial review of the varlen stack: - CRITICAL seam unit bug: indexer_max_seq_len is COMPRESSED index space (metadata divides by compress_ratio) but run_varlen's max_seq_len is kv-token space — the seam shifted twice, freezing a 4x-too-small tuning envelope for cr=4. Seam now multiplies back (exact identity: n_env == indexer_max_seq_len). - Seam crash on legal configs: the DSL paged-MQA logits arena is a 256-aligned buffer column-sliced to max_seq_len — a NON-contiguous view the engine used to reject. The engine now accepts row-major views (inner stride 1) and widens them back to a compact [rows, row_stride] view over the same storage (as_strided, zero copy; the tail columns are never classified — per-row n gates all reads). A dispatch-site hardware-format gate (stride %4, 16B base) falls through to the existing branches for layouts the kernel cannot address (odd-npad DeepGEMM). - OOB-write guard (three reviewers converged): the engine never validated indices/values batch dims — a request-level-shaped buffer under MTP would be silently written past its end (grid comes from logits rows). Full B1-style battery now runs on the engine path (CUDA/dtype/2-D/contiguity/batch/width/alignment), kv_lens contiguity included. - Engine/reference convention alignment: wider-than-k buffers now follow the flat-packed contract identically on both engines; the reference clamps kv_len < next_n to the empty row (all -1) exactly like the kernel — padded/evicted graph slots are a legal input, and the differential oracle can now cover them. - Eager-mode compile churn: without max_seq_len the data-dependent envelope is quantized up to the next power of two (bounded plan set and _VARLEN_CACHE; a growing decode no longer recompiles at every R increment). - Multi-stream escape hatch: run_varlen(workspace=...) (run_ws parity) so concurrent streams do not share the SPLIT publish slab. - Legacy hot path: per-device cached dummy kv_lens (no per-call allocation for the dead ABI slot). New tests: b=16 SPLIT + per-row TSH runtime gate, b=200 BLK=512 non-split, zero-kv slot mixed with live MTP rows (both engines), wide-buffer flat-packed convention (both engines), num_rows / strided-kv_lens guards. Full battery re-verified on B200 (10/10 incl. a non-contiguous arena-view differential). Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>

longcheng-nv · 2026-08-19T01:17:40Z

/bot run

tensorrt-cicd · 2026-08-19T01:28:55Z

PR_Github #67203 [ run ] triggered by Bot. Commit: 2efd66e Link to invocation

tensorrt-cicd · 2026-08-19T01:42:33Z

PR_Github #67203 [ run ] completed with state FAILURE. Commit: 2efd66e
/LLM/main/L0_MergeRequest_PR pipeline #54734 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

Release-Check (PR_Github #67203) failed on pre-commit: ruff-format reflow on 3 files, one ruff F841 (unused next_n unpack in the varlen differential test), one codespell hit (statics -> static fields). Formatting-only plus the two mechanical fixes; kernel exactness re-verified on GPU after the reflow: full unit file 65/65 passed (sm100, standalone overlay stack). Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>

longcheng-nv · 2026-08-19T01:57:23Z

/bot run

tensorrt-cicd · 2026-08-19T02:03:40Z

PR_Github #67217 [ run ] triggered by Bot. Commit: 06af855 Link to invocation

tensorrt-cicd · 2026-08-19T04:21:30Z

PR_Github #67217 [ run ] completed with state SUCCESS. Commit: 06af855
/LLM/main/L0_MergeRequest_PR pipeline #54745 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

longcheng-nv · 2026-08-19T04:31:52Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-19T04:38:43Z

PR_Github #67274 [ run ] triggered by Bot. Commit: 06af855 Link to invocation

…ng-only surface Public package exports now carry ONLY the production contract (selfsampling_topk_run_varlen: per-request device kv_lens, no batch-uniformity assumption — mirrors the single-op shape of the CUDA indexer_topk_decode integration). run/run_ws keep serving as the bench DPS contract and the reference-oracle plumbing (_run_impl is what engine="reference" walks row by row), but are no longer package-exported and carry TESTING/BENCH ONLY docstring warnings. Tests already import the host module directly — zero test churn. Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>

tensorrt-cicd · 2026-08-19T08:32:08Z

PR_Github #67274 [ run ] completed with state SUCCESS. Commit: 06af855
/LLM/main/L0_MergeRequest_PR pipeline #54799 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

longcheng-nv · 2026-08-19T08:35:20Z

/bot run

tensorrt-cicd · 2026-08-19T08:43:54Z

PR_Github #67371 [ run ] triggered by Bot. Commit: fa12c67 Link to invocation

…y, hint/domain tests Framework-integration audit hardening (all GPU-verified, 70/70 unit file): - dispatch gate now also requires fp32 logits (falls through loudly instead of feeding a non-fp32 tensor into the fp32-typed DSL engine; production DSA logits are always fp32 today, this is belt-and-braces for future paths) - logger.info_once on first engagement + logger.warning_once on first hardware-format fall-through: operators can tell which arm served without profiling - new tests: hints containing -1 (the production short-row pad tail that flows back through heuristic_prev_topk; engine's unsigned-compare guard verified exact) and route() domain up to 8192 rows (max_batch x next_n exceeds the bench grid's b<=1024 envelope) Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>

…elper Mirrors warmup_heuristic_topk_decode / warmup_cute_dsl_radix_topk: one tiny real launch per requested num_rows compiles the varlen engine's envelope tuples so no live request pays the first-touch DSL JIT (measured: 9.1 s cold compile for two tuples on a fresh cache; idempotent re-call 0 ms; post-warmup first real call 0.1 ms). Exposed as a module-level helper — CUDA-graph capture warmup already compiles the captured batch sizes, and wiring an automatic init hook needs max_seq_len plumbing that is not available at Indexer.__init__ time (follow-up). Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>

…e.warmup Completes the warmup story: DSAtrtllmAttentionMetadata.warmup_selfsampling_topk mirrors warmup_cute_dsl_radix_topk (same ModelEngine hook, which is where max_seq_len is actually available — Indexer.__init__ is not). Gated on the env flag + the same init-contract conditions; compiles the eager first-touch (num_rows=next_n) tuple via warmup_varlen so no live request pays the DSL JIT. Also: seam init comment now states the two-guard reality (format gate falls through with a one-time warning; in-engine contract violations raise), and run_varlen documents the inherited NaN-ordering limitation (finite inputs incl. +/-inf are tie-aware exact). Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>

tensorrt-cicd · 2026-08-24T03:37:23Z

PR_Github #68664 [ run ] completed with state ABORTED. Commit: ff08b1d

Link to invocation

tensorrt-cicd · 2026-08-24T09:37:32Z

PR_Github #68667 [ run ] completed with state FAILURE. Commit: ff08b1d
/LLM/main/L0_MergeRequest_PR pipeline #56075 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

longcheng-nv · 2026-08-24T21:57:52Z

/bot run

longcheng-nv · 2026-08-24T22:01:45Z

/bot run --reuse-test 56075

tensorrt-cicd · 2026-08-24T22:03:04Z

PR_Github #68907 [ run ] triggered by Bot. Commit: ff08b1d Link to invocation

tensorrt-cicd · 2026-08-24T22:07:10Z

PR_Github #68908 [ run ] triggered by Bot. Commit: ff08b1d Link to invocation

tensorrt-cicd · 2026-08-24T22:07:24Z

PR_Github #68907 [ run ] completed with state ABORTED. Commit: ff08b1d

Link to invocation

tensorrt-cicd · 2026-08-25T02:32:57Z

PR_Github #68908 [ run ] completed with state SUCCESS. Commit: ff08b1d
/LLM/main/L0_MergeRequest_PR pipeline #56296 completed with status: 'SUCCESS'

CI Report

Link to invocation

…e contract comments Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>

longcheng-nv · 2026-08-25T03:25:25Z

/bot run --reuse-test 56296

tensorrt-cicd · 2026-08-25T03:33:28Z

PR_Github #68980 [ run ] triggered by Bot. Commit: 62cfe5d Link to invocation

longcheng-nv · 2026-08-25T07:09:29Z

/bot run --reuse-test 56296

tensorrt-cicd · 2026-08-25T07:10:10Z

PR_Github #68980 [ run ] completed with state FAILURE. Commit: 62cfe5d
/LLM/main/L0_MergeRequest_PR pipeline #56364 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

tensorrt-cicd · 2026-08-25T07:14:50Z

PR_Github #69050 [ run ] triggered by Bot. Commit: 62cfe5d Link to invocation

tensorrt-cicd · 2026-08-25T12:13:56Z

PR_Github #69050 [ run ] completed with state SUCCESS. Commit: 62cfe5d
/LLM/main/L0_MergeRequest_PR pipeline #56424 completed with status: 'SUCCESS'

CI Report

Link to invocation

nv-xtf

LGTM from the disagg side.

yuxianq

attention part LGTM

zongfeijing

LGTM from kernel side

Port TRT-LLM's self-sampling GVR V2 top-K decode kernels (NVIDIA/TensorRT-LLM#17821, commit ed94d4cfbf) as a new opt-in backend "gvr_2" of top_k_varlen, kept separate from the existing "gvr" backend until the perf comparison settles the default (auto order: gvr > gvr_2 > radix > radix_cutlass, so auto never reaches gvr_2 while gvr is suitable). Unlike V1's guess-verify-refine (threshold guessed from the previous step's pre_idx hint, refined by rescans when the hint is stale), V2 derives a bracketed ladder of candidate thresholds from an in-kernel sample of the row itself and resolves exact top-K in a single streaming pass; exactness is guaranteed by count-crossing invariants, never by the estimate, and pre_idx survives only as a degenerate-case anchor. Four kernel families (streaming main w/ multi-CTA SPLIT, register-resident reg/regimg, clustered clus/reg_clus) are picked by a pure host dispatch; per-row lengths are read on device, so one launch serves the whole ragged batch with no prepare kernel, no LJF sort, and no host reads (CUDA-graph safe; the length envelope comes from the logits row width). Implementation notes: - kernels/gvr2_topk_decode.py and kernels/gvr2_topk_host.py are near-verbatim upstream drops (ruff/mypy-excluded like the other verbatim kernel ports) so future syncs stay mechanical. Local changes are limited to the module rename, provenance notes, a _persist() hook that routes the four get_compiled* builders through the persistent CuTe-DSL kernel cache (the upstream compile closures already use --enable-tvm-ffi + symbolic shapes, matching the cache's reload convention), and a fixed kv-arg slot in the regclus debug entry. - Backend contract: fp32 logits only (bf16/fp16 are an upstream follow-up), top_k in {512, 1024, 2048}, compress_ratio in {1, 4}, pre_idx required with width == top_k, sm_100/103 only. The ~21 MB per-device workspace slab is zero-initialized once and self-restoring; multi-stream callers can pass workspace={"gvr2_workspace": ...}. - Tests: tests/topk_varlen/test_topk_varlen_gvr2.py ports the upstream suite onto the FlashInfer API (tie-aware exactness with poisoned pads, short rows, degenerate hints, varlen/MTP/compress-ratio grids, zero-kv rows, family-admission parity, CUDA-graph capture/replay with growing seq_lens, warmup-then-capture, non-contiguous arena views, workspace override, adversarial tie/inf/denormal patterns, and a >10k-point route_split == route dispatch fuzz). Cross-backend consistency tests in test_topk_varlen.py now include gvr_2. - Validation: gvr_2 suite 78/78 and full topk_varlen regression 116/116 on B200 (SM100); SM80/89/90/120 sweeps pass with gvr_2 correctly filtered (51 passed / 155 skipped each). - Perf (B200, fp32, K=1024, uniform lengths, hint quality 0.6, CUDA-graph timing via benchmarks/bench_topk_varlen_gvr2.py): geomean speedups of 2.52x vs gvr (42/42 configs), 3.63x vs gvr non-LB, 2.29x vs radix, 3.41x vs radix_cutlass; 1.00x (0.95-1.05x) vs the upstream TRT-LLM implementation on identical inputs, i.e. the port adds no overhead. AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

github-actions Bot assigned longcheng-nv Aug 17, 2026

longcheng-nv changed the title ~~[None][perf] Add self-sampling GVR top-K decode kernels (CuTeDSL, standalone)~~ [TRTLLM-15293][perf] Add self-sampling GVR top-K decode kernels (CuTeDSL, standalone) Aug 18, 2026

longcheng-nv added 11 commits August 18, 2026 13:02

longcheng-nv added 3 commits August 19, 2026 08:48

BowenFu approved these changes Aug 24, 2026

View reviewed changes

longcheng-nv mentioned this pull request Aug 24, 2026

[None][perf] Block-skip fast path for self-sampling GVR top-K decode (stacked on #17821) #18122

Draft

4 tasks

[TRTLLM-15293][chore] self-sampling top-K: dedupe the per-row prologu…

62cfe5d

…e contract comments Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>

limin2021 approved these changes Aug 26, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_self_sampling.py

lori-ren approved these changes Aug 26, 2026

View reviewed changes

nv-xtf approved these changes Aug 26, 2026

View reviewed changes

yweng0828 approved these changes Aug 26, 2026

View reviewed changes

yuxianq approved these changes Aug 26, 2026

View reviewed changes

longcheng-nv requested a review from zongfeijing August 26, 2026 08:29

BowenFu approved these changes Aug 26, 2026

View reviewed changes

zongfeijing approved these changes Aug 26, 2026

View reviewed changes

juney-nvidia approved these changes Aug 26, 2026

View reviewed changes

juney-nvidia merged commit ed94d4c into NVIDIA:main Aug 26, 2026
10 checks passed

dhiraj113 mentioned this pull request Aug 28, 2026

feat(topk): add self-sampling GVR V2 (gvr_2) backend for top_k_varlen flashinfer-ai/flashinfer#4811

Open

Uh oh!

Conversation

longcheng-nv commented Aug 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

How to enable (GVR V2) — minimal setting per model

DeepSeek-V3.2 (deepseek-ai/DeepSeek-V3.2-Exp)

DeepSeek-V4 Flash (deepseek-ai/DeepSeek-V4-Flash)

DeepSeek-V4 Pro (deepseek-ai/DeepSeek-V4-Pro)

Auto-checked prerequisites, verification, notes

End-to-end decode TPOT (serving A/B)

Correctness and performance summary

Performance comparison

Measurement notes

CUTLASS DSL 4.6.1 compatibility and production-varlen certification

End-to-end accuracy A/B — COMPLETE: 12/12 full-dataset pairs, all passing

Full-dataset pairs (strict-match / weighted; each pair same host + same TP)

MTP coverage

Notes

Implementation status and follow-up work

Test coverage

PR Checklist

Summary

Dev Engineer Review

QA Engineer Review

Uh oh!

longcheng-nv commented Aug 19, 2026

Uh oh!

tensorrt-cicd commented Aug 19, 2026

Uh oh!

tensorrt-cicd commented Aug 19, 2026

Uh oh!

longcheng-nv commented Aug 19, 2026

Uh oh!

tensorrt-cicd commented Aug 19, 2026

Uh oh!

tensorrt-cicd commented Aug 19, 2026

Uh oh!

longcheng-nv commented Aug 19, 2026

Uh oh!

tensorrt-cicd commented Aug 19, 2026

Uh oh!

tensorrt-cicd commented Aug 19, 2026

Uh oh!

longcheng-nv commented Aug 19, 2026

Uh oh!

tensorrt-cicd commented Aug 19, 2026

Uh oh!

tensorrt-cicd commented Aug 24, 2026

Uh oh!

tensorrt-cicd commented Aug 24, 2026

Uh oh!

longcheng-nv commented Aug 24, 2026

Uh oh!

longcheng-nv commented Aug 24, 2026

Uh oh!

tensorrt-cicd commented Aug 24, 2026

Uh oh!

tensorrt-cicd commented Aug 24, 2026

Uh oh!

tensorrt-cicd commented Aug 24, 2026

Uh oh!

tensorrt-cicd commented Aug 25, 2026

Uh oh!

longcheng-nv commented Aug 25, 2026

Uh oh!

tensorrt-cicd commented Aug 25, 2026

Uh oh!

longcheng-nv commented Aug 25, 2026

Uh oh!

tensorrt-cicd commented Aug 25, 2026

Uh oh!

tensorrt-cicd commented Aug 25, 2026

Uh oh!

tensorrt-cicd commented Aug 25, 2026

Uh oh!

Uh oh!

nv-xtf left a comment

Choose a reason for hiding this comment

Uh oh!

longcheng-nv commented Aug 17, 2026 •

edited

Loading

DeepSeek-V3.2 (`deepseek-ai/DeepSeek-V3.2-Exp`)

DeepSeek-V4 Flash (`deepseek-ai/DeepSeek-V4-Flash`)

DeepSeek-V4 Pro (`deepseek-ai/DeepSeek-V4-Pro`)