[None][perf] Add tiered GVR CuTe DSL top-k decode kernels (stacked on #16457)

[longcheng-nv]

What

Adds a tiered GVR CuTe DSL top-K family (three tiers behind one host-side dispatcher) as a guarded fp32 fast path inside the existing trtllm::cute_dsl_gvr_topk_decode op — signature unchanged, no call-site change; everything outside the guard (bf16/fp16, load-balance mode, npad > 256K, banded shapes) takes the in-tree #16457 kernel. Rebased onto main. Net result on 9,515 real decode-capture cases: gm 1.3996× vs the in-tree kernel, worst case 0.9516× — no case regresses more than 10%.

Operational kill switch: TRTLLM_GVR_TIERS_DISABLE=1 disables the tiered fast path entirely — every call takes the in-tree kernel path. Read once per process at first use (TRTLLM_GVR_FALLBACK_BANDS / TRTLLM_GVR_TP_BS / TRTLLM_GVR_DENSE_BS tune routing; see the guard details in Design).

This PR also changes the in-tree kernel itself (the #16457 review follow-ups; full detail in the follow-ups section below). What a bisect landing here should know:

• Output-visible: on rows whose boundary tie class is wider than the candidate buffer, the Phase-2 plateau terminal (done=3) plus the Phase-4 plateau fill now complete the row tie-aware — previously the give-up path could emit -1 pads or arrival-order picks. Both old and new outputs are "some K indices"; the new one is a valid tie-aware top-K, and the emitted index set on such rows changes with this PR.
• Capability, default unchanged: p4_exact_tail gains 16-bit support but the default stays fp32-only (forcing it on bf16 measured gm 1.29–1.36× slower, worst 2.27×); bf16/fp16 production behavior is unchanged.
• Refactors, no output change: the two token-identical exact-tail radix selects collapse into _p4_exact_tail_radix_select (byte-identical PTX verified for the p4_tail_fast=False variant); the launch-tuning policy moves into GvrTopKKernel.pick_tuning / pick_cluster_size with the runner as a thin adapter (sweep test pins runner == kernel policy).

Naming (review follow-up): the family's development codename "BSX" is gone — the tiers are optimizations of the same GVR algorithm, not a different one, so files are gvr_topk_decode_{dispatch,direct,reg,tp}.py, the op-facing symbols are tiered_topk / is_tiered_topk_supported, and the env knobs are TRTLLM_GVR_*.

Design

The problem being solved. Each decode step, the DSA indexer hands this op a batch of rows; each row is up to 262,144 fp32 scores, and the op must return the indices of the K largest (K = 512/1024/2048). The op runs every layer, every step, so it needs to be fast at every batch size from 1 to 1024.

The one idea everything builds on: if you somehow knew the value of the K-th largest score, selection would be a single cheap pass — keep everything above that cutoff. So the whole game is getting a good cutoff estimate cheaply, and then proving it safe before trusting it. That is GVR (Guess–Verify–Refine):

• Guess a cutoff. The best free source is the previous decode step: the model ran this exact selection on almost the same scores one step ago, and its top-K values bracket where today's cutoff lies. (Where per-row history isn't practical, sample a few hundred elements of the row and estimate the cutoff from the sample.)
• Verify by exactly counting how many scores pass the guessed cutoff. If at least K pass, the true top-K is guaranteed to be inside the kept set — the guess only decided how much work remains, never which answer comes out. If the count comes back bad (the guess was too tight or too loose), tighten or loosen and re-count; a deterministic fallback (secant step on the two nearest counts, then a plateau walk) always terminates.
• Refine: run an exact radix select inside the kept set — now a few thousand elements instead of 262K — to produce the final K indices, including correct handling of ties at the boundary.

The guess quality only affects speed, never correctness: every emitted index is justified by exact counts on the current row.

Kernel selection. One host-side chain of four questions picks the implementation per call (pure function of (BS, npad, K, dtype) — no device sync, CUDA-graph safe):

flowchart TD
    OP["cute_dsl_gvr_topk_decode(...)"] --> G{"① dtype and shape<br/>supported by the tiers?"}
    G -- no --> IT["in-tree #16457 kernel"]
    G -- yes --> BAND{"② is (npad, BS) a bucket where<br/>the in-tree kernel measured faster?"}
    BAND -- yes --> IT
    BAND -- no --> TPQ{"③ batch big enough<br/>to stream (tp)?"}
    TPQ -- yes --> TP(["tp tier"])
    TPQ -- no --> DQ{"④ row short enough<br/>to collect whole?"}
    DQ -- yes --> DIR(["direct tier"])
    DQ -- no --> REG(["reg tier"])

Loading

The same chain, evaluated over the whole plane — where any (npad, BS) call lands (identical for K = 512/1024/2048; generated from the dispatch code):

npad \ BS	1–4	8	16–32	64	128	≥256
≤ 6143	direct	direct	direct	direct	direct	tp
6144 – 12288	direct	direct	direct	direct	direct	in-tree
12289 – 20480	reg-L	reg-L	reg-L	reg-D	reg-D	in-tree
20481 – 24575	reg-L	reg-L	reg-L	reg-D	tp	in-tree
24576 – 98303	reg-L	reg-L	in-tree	in-tree	in-tree	in-tree
98304 – 147456	reg-L	in-tree	in-tree	in-tree	in-tree	tp
147457 – 196607	reg-L	reg-L	in-tree	in-tree	in-tree	tp
196608 – 262144	reg-L	reg-D	in-tree	in-tree	tp	tp

The three GVR tiers run the same Guess–Verify–Refine loop; they differ in where the row's data lives while it runs, which is what actually decides speed at each shape:

• direct (short rows, npad ≤ 12288): the row is small enough to skip guessing entirely. One CTA loads the whole row and runs one exact radix select over it. There is nothing to estimate when you can afford to look at everything.
• reg (longer rows, small batch): the row is too big to skip estimation but the batch is small, so latency is what matters — and the enemy of latency is touching DRAM twice. The CTA(s) assigned to a row load it into registers once; the Guess (from the previous step's top-K), every Verify count, and the final Refine all run against those registers. DRAM is read exactly one time per element, period. Two variants exist (reg-L with 512-thread blocks for the smallest batches, reg-D with 1024-thread blocks when there are enough rows to keep the machine busy); the dispatch table picks per shape.
• tp (large batch): with hundreds of rows in flight the machine is throughput-bound, and holding every row in registers no longer pays. Each row is streamed: first a small sample of the row estimates the cutoff (no per-row history needed at this scale), then one pass over the row keeps only the scores above it. A statistical guard rail ("lean-pivot admission") kicks in when the sample says the kept set came out much fatter than K — it tightens the cutoff before the expensive part rather than after. Rows short enough to fit the candidate buffer (npad ≤ kC) skip all of this and just collect everything.
• in-tree #16457 is the dtype-generic GVR kernel already on main, unchanged. It serves everything the tiers decline (bf16/fp16, load-balance mode, npad > 256K, cluster shapes over the device limit) — and the map's mid-table cells, where it is genuinely faster: those rows fit in L2 cache, so re-scanning a cached row is nearly free, and its strategy of repeatedly re-counting to shrink the kept set beats the tiers' one-shot estimates there (measured, >1.10× on at least one production layer per routed bucket). Routing those buckets back is what caps this PR's worst case at 0.95× by construction.

Reading the map's other edges: at small batch, direct gives way to reg exactly where the row stops fitting one CTA's buffer (12288); at large batch, tp takes over at a BS threshold that drops as rows get longer (256 → 128 → 16) — a longer row is more work per row, so fewer rows are needed before streaming keeps the whole GPU busy. Band boundaries sit at values like 24576 and 98304 because the fallback table buckets npad by nearest power of two.

Guard details (①): the tiers require fp32 logits, K ∈ {512, 1024, 2048}, cr ∈ {1, 4}, next_n ≥ 1 with num_rows divisible by next_n, npad ≤ 262144 and a multiple of 64, contiguous 16B-aligned tensors, and the routed tier's cluster size within the device limit. order_row (the row-scheduling hint dsa.py sends for every batch with num_rows ≥ 2×num_sms) is accepted and ignored — the tiers launch per-row CTAs and never consumes the permutation. counters (load-balance mode) always takes the in-tree path. Env controls: TRTLLM_GVR_TIERS_DISABLE=1 turns the fast path off entirely, TRTLLM_GVR_FALLBACK_BANDS=0 disables the band table (②), TRTLLM_GVR_TP_BS / TRTLLM_GVR_DENSE_BS override the tier thresholds (③); malformed values log a warning and fall back to the baked defaults instead of failing the decode step.

Performance (B200, fp32, real SWE-bench decode captures, paired same-rep cold-L2 nsys)

Full-mesh re-measure of this PR's head — 865 real decode-capture cells x 11 batch sizes = 9,515 paired cases, 8 GPUs, 0 harness failures, and a tie-aware exactness check inside every case (see Correctness).

vs the #16457 kernel now on main, shipped operator (band table on):

overall	floor	cases below 0.909	win rate
gm 1.3996x	0.9516x	0 (no case regresses more than 10%)	67.2%

By model: K=2048/cr=1 1.4056x · K=1024/cr=4 1.4063x · K=512/cr=4 1.3772x

By sequence length (npad = post-compression row width; token length ~ npad x cr):

npad	1-2K	4K	8K	16K	32K	64K	128K	256K
gm	2.04 / 1.73	1.67	1.34	1.22	1.17	1.19	1.46	1.44

By batch size:

model \ BS	1	2	4	8	16	32	64	128	256	512	1024
K=2048, cr=1	1.80	1.78	1.76	1.53	1.25	1.24	1.25	1.24	1.26	1.27	1.28
K=512, cr=4	1.67	1.64	1.63	1.48	1.29	1.29	1.26	1.28	1.23	1.24	1.25
K=1024, cr=4	1.72	1.68	1.67	1.51	1.30	1.30	1.28	1.31	1.26	1.27	1.29

By layer: all 109/109 production layers win — worst 1.273x, medians 1.37-1.41x, best 1.540x. (The loss tail is a per-(layer, shape) phenomenon; the band table absorbs it.)

The same run also isolates the cost of everything added on top of the reviewed tier commits: measured against the previous full grid, the in-tree arm drifts 1.0051x and the tier arm 1.0060x — i.e. the follow-up work (comment prune, launch-policy de-duplication, 16-bit exact-tail capability, P4 helper extraction, plateau terminal) costs nothing measurable on either side.

Development A/B additionally falsified (measured, component-isolated): exact-count admission as a wholesale replacement (gain and harm co-sourced), a K-scaled candidate-budget diet (pure harm), ladder-quantile re-placement under the shipped admission, cluster-size occupancy cuts, multi-pass-straggler hypotheses (the residual band is single-pass; the gap is candidate-set fatness), and a conditional lean-pivot for the register tier (short kernels cannot amortise an extra cluster round trip).

Review follow-ups from #16457 (resolved here)

The follow-up items committed to reviewers on #16457 all land in this PR (f47dda5c, 3959d327):

item (reviewer)	resolution
comment pruning (@lfr-0531)	provenance commentary reduced to invariants/contracts across kernel + custom-op files
launch-shape policy duplication (@limin2021)	pick_config split into pick_cluster_size + pick_tuning (kernel = single source); runner _pick_tuning is now a thin adapter, cluster auto-pick delegates; the intentional divergence is documented (runner asserts on 32B misalignment, launch downgrades); new sweep test pins runner == kernel policy
16-bit exact-tail (@mingyangHao)	capability landed, default kept fp32-only (review-round measurement: on 16-bit the ambiguity gate fires on virtually every input — bf16 quantization plateaus — costing gm 1.29–1.36×, worst 2.27× across the envelope, while typical 16-bit inputs are value-exact without it, 48/48 paired runs); candidate keys are always fp32 (injective upcast at collect), so p4_exact_tail=True is exact for every dtype; new adversarial test (two distinct 16-bit values in one fine bin straddling K, fp16 + bf16) opts in explicitly
P4 radix duplication (@mingyangHao)	the two token-identical copies collapse into one @cute.jit helper; byte-identical PTX verified for the p4_tail_fast=False variant (465,875 B)
dispatch guard	this PR's dispatcher
plateau undershoot terminal (@mingyangHao)	resolved (3959d327). Confirmed site: the admission path's tie-plateau fail-soft landed done=1 on the undershoot side, so Phase 4 padded the tail with -1. Both terminals now first collapse the bracket by bounded bisection to ADJACENT floats — every in-bracket value is then bitwise-equal, a genuine tie class — so Phase 4 emits the sure winners and a ticketed fill completes the row from that class (any (K-count)-subset of a tie class is a valid tie-aware completion). The guard requires a coherent undershoot-overflow bracket with both counts current, so the retry path's widened brackets are excluded; non-plateau undershoot keeps the documented -1 encoding. New adversarial test: a plateau wider than the candidate buffer straddling K, fp32 + fp16 x {rank-scatter cs=1, cs=4, histogram-snap}, 6/6. Implementation note now in the code: the terminal is captured into a dedicated SMEM slot before Phase 4, because Phase 4 reuses that scalar slot as radix scratch. Follow-up (9eadcf20): Phase 2 has two secant drivers — the SMEM/leader one and a register-resident redundant-warp one used at cluster_size == 1 — and the terminal initially landed only in the former, so the same plateau still padded with -1 on the classic admission path. Both terminals are now mirrored into the register-resident driver (warp-uniform, so the counting barrier cadence is unchanged), and the adversarial matrix grew to 5 variants × 2 dtypes (adding enable_r0=False at cluster size 1 and 4, the route that exposed the gap), 10/10.

Gate for the above: GVR top-K suite 684 passed / 144 skipped, tiers suite 94 passed / 8 skipped (order_row acceptance, kill switch, env soft-fail tests added in the review round), plus the full-mesh re-measure in the Performance section.

Test Coverage

• New test_cute_dsl_gvr_topk_tiers.py (103 cases, shared tie-aware checker in conftest.py as a fixture, registered in l0_b300.yml; CI wall-clock managed by pinning cases onto a minimal covering set of JIT variants: 570 s → 306 s, no code-path loss).
• Existing test_cute_dsl_gvr_topk_decode.py full suite green.

Draft checklist (before ready-for-review)

• Direct A/B perf table vs [None][perf] GVR top-K decode: enable R0 histogram-ladder admission by default #16457 kernel (full 865×11 grid); ragged-mix perf axis deferred (correctness-covered by varlen/MTP tests)
• CUDA-graph capture/replay test
• Rebase onto main

PR Checklist

• PR title follows the required format
• Description explains what and why
• Test cases added and passing locally
• DCO sign-off on all commits

🤖 Generated with Claude Code

Dev Engineer Review

• Adds a guarded fp32 BSX fast path to trtllm::cute_dsl_gvr_topk_decode.
• Routes supported cases to direct, register-resident, or throughput CuTe DSL tiers.
• Preserves the existing API and falls back to the in-tree #16457 kernel for unsupported cases.
• Adds guards for dtype, shape, alignment, K, next_n, compression ratio, hardware limits, and calibrated fallback bands.
• Adds TRTLLM_BSX_DISABLE and soft-failing environment configuration.
• Refactors launch and tuning policy through GvrTopKKernel.pick_cluster_size and GvrTopKKernel.pick_tuning.
• Extends plateau-terminal handling and exact-tail tie repair.
• Adds Blackwell BSX direct, register-resident, and throughput kernels with compilation caching and clustered launch support.
• Review should verify CUDA synchronization, atomic operations, cache invalidation, fallback behavior, and API consistency with CODING_GUIDELINES.md.

QA Engineer Review

Test changes

• Added test_cute_dsl_bsx_topk_decode.py with coverage for:
  • CUDA graph replay.
  • Routing, environment controls, fallback bands, and launch policies.
  • Direct, register-resident, and throughput paths.
  • Degenerate rows, invalid pre_idx, poisoned tails, ties, and MTP behavior.
• Updated test_cute_dsl_gvr_topk_decode.py with:
  • Shared tie-aware checking.
  • 16-bit exact-tail coverage.
  • Launch-policy consistency checks.
  • Plateau-terminal coverage across variants and cluster sizes.
• Added the shared tie_aware_check fixture in conftest.py.
• Updated tests/integration/test_lists/test-db/l0_b300.yml to exclude the BSX and GVR test modules from the pre_merge PyTorch attention list.

Coverage verdict

• Both modified test modules are listed in l0_b300.yml, but the entries exclude them from that test list.
• The provided context reports successful local BSX and GVR suites, but it does not identify positive CI or manual-QA coverage.
• Verdict: needs follow-up.

[longcheng-nv]

Pushed ef1c5e08 — band-table recalibration (131072 band lower bound 16 -> 8) + updated the PR-body Evidence section accordingly.

Context: while building a unified-dispatch framework on top of this branch, a full-grid re-measure of the exact shipped head revealed that the original calibration harness measured the reg tier with a faster experimental streaming phase1 inherited from GvrTpKernel (the measurement arm rebound only the tp entry point — GvrRegKernel inherits phase1/streaming-load helpers, so reg-routed shapes silently rode the experimental base, ~20% fast at 32K-128K x BS1-8). On the shipped code the 128K x BS8 bucket dips to 0.82-0.91x vs the in-tree kernel (5 production layers), below the 0.909 floor the band table guarantees — hence the one-bucket recalibration. tp/direct tiers and the headline gm reproduce within 0.1-0.3% (1.3966 claimed vs 1.3965 re-measured; 1.3899 after routing the extra bucket). Unit test extended (test_bsx_fallback_band_table).

[longcheng-nv]

/bot run

[tensorrt-cicd]

PR_Github #62363 [ run ] triggered by Bot. Commit: ef1c5e0 Link to invocation

[longcheng-nv]

Pushed a55ed714 — bsx test-suite CI slimming: 570s -> 306s measured on B200 (-46%), zero code-path coverage lost. The suite's cost is DSL JIT compiles (15-34s per constexpr variant), so cases are now pinned onto a minimal covering set of variants: MTP matrix 18 -> 8 variant compiles (full {2,3}x{1,4} cross on tp, diagonals on direct/reg; next_n=4 dropped as structurally redundant with 2), reg launch table keeps all 14 route asserts but live-launches only the 6 codegen-distinct instances, hardening/admission cells re-pinned onto already-compiled variants, bf16 fallback check is guard-only (its in-tree execution is covered by the sibling gvr suite). PR-body Test Coverage section updated.

[longcheng-nv]

/bot run

[longcheng-nv]

/bot run

[tensorrt-cicd]

PR_Github #62390 [ run ] triggered by Bot. Commit: a55ed71 Link to invocation

[tensorrt-cicd]

PR_Github #62363 [ run ] completed with state ABORTED. Commit: ef1c5e0

Link to invocation

[tensorrt-cicd]

PR_Github #62390 [ run ] completed with state SUCCESS. Commit: a55ed71
/LLM/main/L0_MergeRequest_PR pipeline #50551 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]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #62444 [ run ] triggered by Bot. Commit: a55ed71 Link to invocation

[longcheng-nv]

/bot run

[tensorrt-cicd]

PR_Github #62475 [ run ] triggered by Bot. Commit: d6d960d Link to invocation

[tensorrt-cicd]

PR_Github #62444 [ run ] completed with state ABORTED. Commit: a55ed71

Link to invocation

[tensorrt-cicd]

PR_Github #62475 [ run ] completed with state SUCCESS. Commit: d6d960d
/LLM/main/L0_MergeRequest_PR pipeline #50625 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]

Pushed f47dda5c — resolves the follow-up items committed to reviewers on #16457 (@lfr-0531 comment pruning; @limin2021 launch-shape policy single-source with a runner==kernel sweep test; @mingyangHao 16-bit exact-tail enablement with a new fp16/bf16 adversarial tie test, and the P4 exact-tail radix de-duplication with byte-identical-PTX verification for the p4_tail_fast=False variant). The plateau-undershoot terminal follows as its own commit — the audit found the rank-scatter path has no cand_count<K branch at all, so the fix is wider than the review comment assumed. PR body has the resolution table. Gates: sparse-attention suite 674 passed / 144 skipped; BSX suite 91 passed / 8 skipped.

[longcheng-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #62499 [ run ] triggered by Bot. Commit: f47dda5 Link to invocation

[tensorrt-cicd]

PR_Github #62499 [ run ] completed with state SUCCESS. Commit: f47dda5
/LLM/main/L0_MergeRequest_PR pipeline #50645 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]

Pushed 72436e98 — fixes the Release-Check failure from the previous run (ruff-format on three hand-wrapped expressions the follow-up commit introduced, plus the appended tests). Formatting only; ruff format --check is clean on every touched file and the affected test families re-run green (95 passed / 16 skipped).

The other 4 failures in that run are unrelated to this PR (A100 llmapi/test_llm_pytorch.py disagg-streaming / part0 and B200 _torch/sampler beam-search speculative).

[longcheng-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #62528 [ run ] triggered by Bot. Commit: 72436e9 Link to invocation

[tensorrt-cicd]

PR_Github #62528 [ run ] completed with state SUCCESS. Commit: 72436e9
/LLM/main/L0_MergeRequest_PR pipeline #50671 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]

Pushed 3959d327 (plateau terminal — the last #16457 follow-up) and 828e42ae (band-table scoping), and re-measured the full 865x11 grid on this head across 8 GPUs.

Correctness: 9,515/9,515 cases pass the in-measurement tie-aware exactness check on real decode captures; 0 harness failures.

Performance vs the #16457 kernel: gm 1.3996x, floor 0.9516x, zero cases below 0.909 (no case more than 10% slower), win rate 67.2%. Per-model 1.4056 / 1.4063 / 1.3772; all 109 production layers win (worst 1.273x).

The run also isolated the cost of the follow-up work itself: in-tree arm drift 1.0051x, bsx arm drift 1.0060x vs the previous full grid — the comment prune, policy de-duplication, 16-bit exact-tail enablement, P4 helper extraction and plateau terminal are all performance-neutral.

It additionally caught a mistake in the earlier band recalibration: the nearest-power-of-two bucket for 131072 mixes npad=131136 (which holds the layers that force the routing) with npad~163776 (which only rounds into that bucket and runs 1.36-1.60x ahead), so the bs=8 extension was giving up 58 winning cells to protect 5. 828e42ae scopes the extension to true 128K shapes, verified by re-measuring all 58 affected cells x 6 batch sizes. Full Evidence section updated.

[longcheng-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #62620 [ run ] triggered by Bot. Commit: 828e42a Link to invocation

[longcheng-nv]

Pushed d29db18566 addressing all five review comments (details in the per-thread replies):

1. order_row accepted-and-ignored — the guard no longer turns bsx off for num_rows >= 2*num_sms batches; new TRTLLM_BSX_DISABLE kill switch (also the tests' force-in-tree mechanism now) + test_bsx_accepts_order_row.
2. 296 // bs — derived as 2 * multi_processor_count (bit-identical on B200) + explicit frozen-B200-calibration note for everything else.
3. 16-bit p4_exact_tail default REVERTED to fp32-only, with data — on 16-bit the ambiguity gate fires on virtually every input (bf16 quantization plateaus): measured envelope cost gm 1.29–1.36x, worst 2.27x, while typical bf16 inputs are value-exact without it (48/48 paired runs). Capability stays behind the knob; the adversarial test opts in explicitly. With the revert, the PR-body "half-precision path unchanged" claim is accurate again.
4. env knobs fail soft — warn once + baked default instead of ValueError on the decode path; new test.
5. _tie_aware_check — single definition in conftest.py, injected as the tie_aware_check fixture; both modules' copies removed.

Both suites green locally on B200 at the new head (bsx 94 passed / 8 skipped; gvr 684 passed / 144 skipped). Note on the last CI round (pipeline 51940): all 5 failures were infrastructure — 4x H100 CPP unit tests died in cmake configure on a nanobind FetchContent clone failure, and the prefix-aware-scheduling smoke failed on a broken synthetic-multi-round-qa dependency in the CI image (ModuleNotFoundError: utils) — none touch this PR's code paths.

/bot run --disable-fail-fast

[longcheng-nv]

/bot run --disable-fail-fast

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@tests/unittest/_torch/attention/sparse/conftest.py`:
- Around line 28-36: Update the _tie_aware_check_impl signature so
ref_vals_cache uses the explicit type dict[tuple[int, int, int, int, int],
torch.Tensor] | None instead of dict = None, preserving its default value and
existing cache behavior.

🪄 Autofix

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: af540bb7-a420-40cc-9959-ac1ae0c3c679

📥 Commits

Reviewing files that changed from the base of the PR and between e382f98 and d29db18.

📒 Files selected for processing (7)

• tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py
• tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
• tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_dispatch.py
• tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py
• tests/unittest/_torch/attention/sparse/conftest.py
• tests/unittest/_torch/attention/sparse/test_cute_dsl_bsx_topk_decode.py
• tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py

🚧 Files skipped from review as they are similar to previous changes (2)

• tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py
• tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_bsx_tp.py

[tensorrt-cicd]

PR_Github #65208 [ run ] triggered by Bot. Commit: d29db18 Link to invocation

[longcheng-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #65212 [ run ] triggered by Bot. Commit: fd97411 Link to invocation

[tensorrt-cicd]

PR_Github #65208 [ run ] completed with state ABORTED. Commit: d29db18

Link to invocation

[longcheng-nv]

/bot run --disable-fail-fast

[longcheng-nv]

CI update (pipeline 52997, commit fd97411): 56,988 passed / 3 failed / 15,910 skipped. All previous-round failures are gone (the cmake/nanobind clones and the prefix-aware smoke both passed this time, confirming they were infra).

The 3 remaining failures are unrelated to this PR: unittest/auto_deploy/singlegpu/models/test_glm4_moe_modeling — test_glm4_moe_moe_equivalence[dtype0-2-6] and test_glm4_moe_decoder_layer_equivalence[1-dtype0-2-6] fail with rmse_ref=inf (the test's own reference model produces inf on that runner), plus the suite aggregator. Both cases were SKIPPED in the previous run of this same branch base (51940) — the skip gate flapped between runners — and this PR's diff (cute_dsl top-k + its tests) has no intersection with AutoDeploy GLM4-MoE modeling.

Re-triggering to get a clean run.

[longcheng-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

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

PR_Github #65258 [ run ] triggered by Bot. Commit: fd97411 Link to invocation

[tensorrt-cicd]

PR_Github #65258 [ run ] completed with state SUCCESS. Commit: fd97411
/LLM/main/L0_MergeRequest_PR pipeline #53036 completed with status: 'SUCCESS'

CI Report

Link to invocation

[limin2021]

Why ignore the order_row ? It's better to make the behavior keep consistent with the paramter.

[longcheng-nv]

order_row only tells the in-tree persistent kernel which rows to start first (longest-job-first, so a long row does not start last and become the straggler). It never changes results: output row i always holds row i's top-K, whatever the order. The tiers launch one CTA (or cluster) per row and the hardware scheduler decides execution order, so the hint has nothing left to control.

Being "consistent" the other way, i.e. rejecting calls that carry it, would switch the tiers off for every batch with num_rows >= 2*num_sms, because dsa.py sends the hint for all such batches, and those are the shapes the tp tier exists for.

brnguyen2 raised the same question in the thread above. The accept-and-ignore contract is documented in the dispatcher docstring, and test_tiers_accepts_order_row checks both halves: the index sets do not change under the permutation, and a large batch with order_row stays on the tiers.

[limin2021]

I think the name of bsx is not very good, whose meaning is not easy to follow. The new updates in this MR mainly related to how to read/resident the logits. Current naming and switch method seems like a new kinds of algo. How about the following naming?

gvr_topk_decode_reg/smem.py: read/resident logits into registers/smem.
gvr_topk_decode_fallback.py 【direct】: don't need gvr, fallback to radix impl.
gvr_topk_decode_streaming.py: not resident, every time we load from gmem.

Please don't use direct/tp like naming. It's not very clear.

Note: I just give an example.

[longcheng-nv]

Agreed, and done in 4e4abbd — the working name is gone. You are right that these are optimizations of the same GVR algorithm (the tiers differ only in where the row's logits are resident while Guess–Verify–Refine runs), so the naming now says exactly that: files are gvr_topk_decode_{dispatch,direct,reg,tp}.py, the op-facing symbols are tiered_topk / is_tiered_topk_supported, the env knobs are TRTLLM_GVR_* (kill switch: TRTLLM_GVR_TIERS_DISABLE), and the test module is test_cute_dsl_gvr_topk_tiers.py. On your reg/smem suggestion: reg (register-resident) and direct (whole-row collect) match your read/resident framing; the third tier streams the row rather than holding it in SMEM, so it keeps its tp (throughput) name. Pure rename commit, zero logic change; both suites green post-rename (tiers 95, GVR 684). PR title and description updated to "tiered GVR".

[limin2021]

I would suggest remove the bsx concept. And switch kernels directly from the previous existing gvr kernel impl and the newly add kenrel variants.

[longcheng-nv]

The naming half is done in 4e4abbd (no more "bsx" anywhere — see the reply on the naming thread). On folding the switch into the previous kernel instead of a dispatcher module: the routing is deliberately a separate host-pure module rather than branches inside the existing kernel's launcher, for three reasons. (1) It is not one switch — it is a guard (dtype/shape/alignment/hw-cluster-cap), a measured per-(npad,BS) fallback-band table, a tier route, and an env-knob family with caching; inlining that into the existing kernel file would couple its correctness-critical code to routing policy that will be recalibrated. (2) The dispatcher is the documented operational surface (kill switch + knobs in one docstring) and is what the unit tests pin (route table, band table, env soft-fail, cluster-cap memoization). (3) The previous kernel stays a plain callable with zero new coupling — it is also the fallback arm the differential tests compare against, which requires it to be reachable without the new code. So the concept that remains is just "the GVR tiers + their dispatcher", named as such.

[limin2021]

What's the difference between this kernel and the previous exiting kernel?

[longcheng-nv]

Both run the same Guess-Verify-Refine loop and return the same contract. The differences are in how it executes, and they matter specifically at large batch:

1. The existing kernel finds the cutoff by re-counting: several passes over the row, each with a tighter threshold. That is cheap while the row sits in L2. At BS 16-1024 the batch working set is hundreds of MB, rows are not L2-resident, and every extra pass is a DRAM pass. So tp samples the row (every 32nd float4) to estimate the cutoff once, then makes a single streaming pass that counts and collects candidates at the same time. A bad estimate costs one re-stream, and the hint ladder from the previous decode step makes bad estimates rare.

2. The existing kernel is persistent: one cluster cooperates on one row, split into slices. tp launches BS x CS CTAs, one cluster per row, with CS picked so two waves of rows fill the SMs (min_blocks_per_mp=2). Its parallelism scales with the batch rather than the row.

3. Candidates are packed (key<<32 | idx) u64 into CTA0's SMEM (capped at kC), and CTA0 alone runs the final 4x8-bit radix select with tie-aware emit. The existing kernel instead keeps per-slice candidates and selects cooperatively across the row slices.

4. Scope: tp is fp32-only, no LB mode, npad <= 262144. The existing kernel keeps everything else, plus the mid-N shapes where L2-resident re-counting genuinely wins; the fallback-band table routes those back to it.

The module docstring maps each phase (P1/P2a/P2b/P2c/P3/P4) to the CUDA original.

[limin2021]

How about abstact the commons into a util file, and the use them in the kernel file.

[longcheng-nv]

They are shared already, through inheritance rather than a util file: GvrRegKernel subclasses GvrTpKernel (this import block), so Phase 1, _row_n_eff, exchange_counts and the DSMEM idioms come from the tp module, along with the 11 imported PTX wrappers (_f2u_bits, _shfl_*, _mapa_shared_cluster, ...).

The few helpers defined locally differ on purpose rather than being copies: the atomicAdd here is CTA-scope relaxed because a wider scope stops ptxas from warp-aggregating it (docstring note 1), and the cluster sync is the aligned variant that cg::cluster.sync() parity needs. The direct tier has nothing worth extracting: no GVR loop, just a whole-row radix select.

If you would rather have the shared helpers in a neutral _ptx_common.py so the tp file is not their owner, that is a mechanical move and I can do it in this PR. I held off only to avoid touching all four files again right after the rename.

[limin2021]

How about the perf on varlen inputs?

[brnguyen2]

Remaining items are description-level; the code itself looks ready.

1. Ticket: a new kernel family of this size (+5.9k lines, production dispatch change) should carry a TRTLLM JIRA ticket in the title rather than [None].

2. Description under-reports the in-tree kernel changes. The description says everything outside the guard "falls back to the #16457 kernel unchanged", but this PR changes that kernel too: the Phase-2 plateau terminal (done=3) plus the Phase-4 plateau fill change output on rows whose boundary tie class is wider than kC (previously the give-up path could emit -1 pads / arrival-order picks; now they're tie-aware filled — an improvement, but a behavior change a bisect will land on), the exact-tail radix select was deduplicated into _p4_exact_tail_radix_select, and the launch-tuning policy moved into GvrTopKKernel.pick_tuning/pick_cluster_size with the runner as an adapter. Please list these in the description so in-tree output changes are attributable to this PR.

3. Kill switch visibility: TRTLLM_BSX_DISABLE is the operational escape hatch for the new path; it deserves a line in the PR description so on-call can find it without reading the dispatcher docstring.

[brnguyen2]

route_cluster_size re-runs route() + the _parse_reg string parse on every call, and is_bsx_supported runs on every eager forward() — so the ~2.5-3µs host cost that the _DISPATCH_CACHE comment (line 234) says motivated caching the dispatch is still paid per call by the guard, on kernels the same comment says run <20µs. Memoize the cluster-cap verdict keyed on (bs, npad, top_k) (cleared in _reset_env_cache, like _DISPATCH_CACHE), or fold the cap check into _bind so one cache covers both. Irrelevant under CUDA-graph replay, but the eager decode path pays it.

[longcheng-nv]

Fixed in abee8b8 — the cluster-cap verdict is memoized exactly as you framed it: _CAP_OK_CACHE keyed on (bs, npad, top_k), populated on first miss in is_bsx_supported, and cleared by _reset_env_cache() alongside _DISPATCH_CACHE (the verdict embeds the routed tier, so it inherits the same env dependence). Steady-state eager forward() now pays a dict hit instead of route() + _parse_reg + tp_cluster_size. Covered by test_bsx_cluster_cap_verdict_memoized (verdict correctness vs a fresh recompute, memo population, and reset clearing).

[brnguyen2]

The env-knob list here omits TRTLLM_BSX_FALLBACK_BANDS (introduced at line 131, tested via the bands fixture). Since this docstring is the reference for the knob family, add it — one line: 0 disables the fallback-band table so bsx serves every guarded shape.

[longcheng-nv]

Fixed in abee8b8 — TRTLLM_BSX_FALLBACK_BANDS added to the docstring knob list (0 disables the measured fallback-band table so bsx serves every guarded shape; unset/other keeps it active). It parses through _env_threshold, so the existing malformed-values-fail-soft note already covers it.

[brnguyen2]

Internal development codenames (op9/op17 here, op#26 at lines 38 and 1201) survived the codename-scrub commit (2a1950c), which removed the same references from gvr_topk_decode.py. Rephrase to match — e.g. "in earlier cross-arch A/Bs" — so the scrub is complete.

[longcheng-nv]

Fixed in abee8b8 — all three survivors scrubbed to match the codename-scrub commit conventions: lines 38 and 1201 keep the load-bearing "in-tree R0 parity" anchor and drop the campaign reference; the _two_waves_rows comment now reads "measured HW-invariant in earlier B200/B300 cross-arch A/Bs" (your phrasing, with the two architectures kept since that is the fact the device-independent bands rest on). Also re-swept every file this PR touches with the codename patterns — zero remaining hits.

[longcheng-nv]

/bot run

[tensorrt-cicd]

PR_Github #65313 [ run ] triggered by Bot. Commit: abee8b8 Link to invocation

[tensorrt-cicd]

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

/bot run

[tensorrt-cicd]

PR_Github #65321 [ run ] triggered by Bot. Commit: 4e4abbd Link to invocation

[tensorrt-cicd]

PR_Github #65321 [ run ] completed with state FAILURE. Commit: 4e4abbd
/LLM/main/L0_MergeRequest_PR pipeline #53095 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