[None][feat] Add Phase-4 rank-scatter variant to GVR Top-K decode kernels by longcheng-nv · Pull Request #15709 · NVIDIA/TensorRT-LLM

longcheng-nv · 2026-06-29T07:37:05Z

Summary

Update (rev 2)

Two follow-up commits since initial submission:

Default flipped to exact rank-scatter. enable_p4_rank_scatter and enable_p4_rank_scatter_exact now default True in both GvrTopKKernel and GvrTopKLBKernel. Exact rank-scatter is a drop-in for the snap Phase-4 (both exact, vdiff=0) and the faster path, so it is now the default. The cute_dsl_custom_ops construction sites do not pass these flags, so the production decode path now uses exact rank-scatter for Phase-4 (i.e. this PR is no longer kernel-only — the framework default is switched here). Pass enable_p4_rank_scatter_exact=False to opt back to the cheaper approximate-tie variant, or enable_p4_rank_scatter=False for the original snap.
Correctness fix in the exact fine-search (flagged by CodeRabbit). The fine-recursion staged its prefix/target-warp metadata in smem_hist[0]/[1], which are real fine-histogram bins; when the straddling value lands in the lowest fine warp (twf2 == num_warps-1) the reverse scan reads that metadata back as counts and returns the wrong top-K set. Fixed by staging the two scalars in spare s_iscalars[4]/[1] (outside the fine-histogram range, dead in that window). Pure scratch relocation, no extra barriers, no perf impact.

The published nsys cold-L2 ratios were re-confirmed on a fresh B200 host after both commits (same seed=42 inputs, same methodology): K=2048 fp32 N=8K-256K measured 1.013-1.148 (report 1.017-1.163); K=1024 fp16 N=16K peak 1.328 (report 1.313); K=512 fp32 N=16K peak 1.321 (report 1.469). All spot-checked cells reproduce rank-scatter >= snap.

Adds an opt-in Phase-4 (snap) replacement to the cuTe DSL GVR Top-K decode kernels — a fused rank-and-scatter — to both the single-CTA-per-row path (GvrTopKKernel) and the hybrid multi-CTA load-balance path (GvrTopKLBKernel, the kernel introduced for long-context decode). It mainly improves single-CTA and multi-CTA GVR at small/medium sequence lengths, where Phase-4 is the dominant phase.

Both flags now default on (exact rank-scatter); see the Update section above. Pass enable_p4_rank_scatter=False to restore the original snap Phase-4.

As of rev 2 this PR flips the framework default (the production decode path now uses exact rank-scatter, see Update above). There is still no DeepSeek-V3.2/V4-specific end-to-end wiring beyond the default switch; model-level enablement/validation lands in a follow-up PR.

What changes

GvrTopKKernel.phase4_rank_scatter (new), gated by two constructor flags:

flag	effect
`enable_p4_rank_scatter`	dispatch rank-scatter instead of the iterative histogram-snap
`enable_p4_rank_scatter_exact`	add a fixed-256-bin fine recursion on the straddling bin for an exact tie-break

The default Phase-4 (phase4_histogram_snap) finds the K-th value by an iterative histogram bracket + a 2-pass writeback (~14 block barriers). The rank-scatter path instead:

builds one coarse kNumBins histogram over the candidate values,
does a 3-step high→low bin search to the straddling bin b* and the rank of everything strictly above it (rank_above),
(exact only) runs one fixed 256-sub-bin fine-histogram recursion on b* (256 ≤ kNumBins for every supported dtype/K, so it reuses the existing smem_hist allocation),
scatters each survivor directly to its output rank in a single pass.

Net: the Phase-4 barrier count drops from ~14 → ~7, exact (vdiff=0) on the realistic temporally-coherent indexer distribution.

Because GvrTopKKernel.run_one_row is shared by the single-CTA path and the cluster leader-only Phase-4 path, the lever applies to both GVR branches (single-CTA and multi-CTA cluster). After the DSMEM gather the candidate set is fully resident in the leader CTA's SMEM, so the cluster Phase-4 is structurally identical to the single-CTA case. GvrTopKLBKernel forwards the two flags to both its underlying instances (_cluster_kernel for long rows, _single_kernel for short rows).

Why this helps (and where)

GVR's Phase-4 share of kernel time is largest at small/medium N — it operates only on the ~3×K SMEM candidate set (N-independent), so as N shrinks it dominates (≈42–55% of the kernel at N ≤ 16K, falling to ≈15–20% at N ≥ 128K where the full-N Phase-2/3 scans dominate). Cutting the Phase-4 barrier count therefore yields the biggest relative gain exactly in the small/medium-N decode regime, and the gain is dtype- and K-dependent (fp32/fp16 and larger K benefit most; bf16 least, because its half-precision ties collapse the snap into few histogram buckets so the baseline is already cheap).

Performance — single-CTA rank-scatter vs the snap baseline (same kernel, only Phase-4 swapped)

Metric: nsys pure-kernel GPU time, cold-L2 (512 MB-flush, CUDA-graph replay, per-cell NVTX range) — the canonical metric from the op-bench's full 11-op nsys re-test. Ratios are rank-scatter speedup over snap = t(snap) / t(rank-scatter) (>1 = rank-scatter faster). Synthetic inputs reproduce real DeepSeek-V4 indexer statistics (temporally-coherent, preIdx hit-rate 0.6). B200 and B300 use byte-identical inputs, so the comparison is pure hardware.

B300 (sm_103)

(A.1) BS=1, by seq-len, per (dtype, K):

dtype	K	N=4K	N=8K	N=16K	N=32K	N=64K	N=128K	N=256K
fp32	512	1.110	1.138	1.481	1.127	1.052	1.014	1.058
fp32	1024	1.099	1.214	1.077	1.277	1.240	1.054	1.131
fp32	2048	–	1.124	1.078	1.149	1.163	1.137	1.015
bf16	512	1.042	1.117	1.066	1.073	1.003	1.007	1.028
bf16	1024	1.039	1.144	1.245	1.050	1.068	1.053	1.014
bf16	2048	–	1.076	1.091	1.085	1.059	1.017	1.004
fp16	512	1.044	1.089	1.140	1.237	1.104	1.066	1.066
fp16	1024	1.161	1.147	1.321	1.195	1.142	1.010	1.049
fp16	2048	–	1.172	1.125	1.115	1.151	1.062	1.029

(A.2) BS-scaling, per (dtype, K) (median across N — isolates the BS axis):

dtype	K	BS=1	BS=2	BS=4	BS=8	BS=16	BS=32	BS=64	BS=128	BS=256	BS=512	BS=1024	BS=2048
fp32	512	1.117	1.121	1.113	1.114	1.117	1.123	1.119	1.108	1.080	1.075	1.068	1.077
fp32	1024	1.124	1.128	1.128	1.126	1.125	1.120	1.121	1.096	1.101	1.050	1.060	1.054
fp32	2048	1.131	1.141	1.136	1.132	1.139	1.135	1.124	1.106	1.069	1.056	1.056	1.041
bf16	512	1.052	1.045	1.046	1.044	1.047	1.038	1.050	1.039	1.044	1.042	1.059	1.062
bf16	1024	1.055	1.052	1.053	1.057	1.053	1.053	1.051	1.049	1.047	1.038	1.037	1.064
bf16	2048	1.070	1.068	1.070	1.064	1.067	1.067	1.066	1.062	1.044	1.056	1.063	1.054
fp16	512	1.092	1.083	1.089	1.087	1.083	1.090	1.075	1.087	1.062	1.058	1.074	1.074
fp16	1024	1.159	1.144	1.145	1.145	1.143	1.148	1.146	1.147	1.171	1.186	1.185	1.190
fp16	2048	1.121	1.118	1.116	1.125	1.118	1.115	1.116	1.106	1.099	1.115	1.122	1.126

Overall B300: median 1.077×, 704/720 cells faster.

B200 (sm_100) — essentially identical (hardware-invariant)

(B.1) BS=1, by seq-len, per (dtype, K):

dtype	K	N=4K	N=8K	N=16K	N=32K	N=64K	N=128K	N=256K
fp32	512	1.110	1.126	1.469	1.126	1.049	1.015	1.065
fp32	1024	1.085	1.208	1.074	1.275	1.234	1.052	1.130
fp32	2048	–	1.132	1.071	1.090	1.163	1.135	1.017
bf16	512	1.048	1.116	1.069	1.085	1.004	1.007	1.031
bf16	1024	1.038	1.146	1.239	1.048	1.065	1.051	1.015
bf16	2048	–	1.078	1.084	1.086	1.056	1.014	1.005
fp16	512	1.035	1.084	1.141	1.236	1.100	1.065	1.073
fp16	1024	1.159	1.151	1.313	1.192	1.138	1.010	1.050
fp16	2048	–	1.174	1.126	1.110	1.144	1.062	1.026

(B.2) BS-scaling, per (dtype, K) (median across N):

dtype	K	BS=1	BS=2	BS=4	BS=8	BS=16	BS=32	BS=64	BS=128	BS=256	BS=512	BS=1024	BS=2048
fp32	512	1.114	1.113	1.118	1.125	1.124	1.114	1.109	1.104	1.079	1.078	1.074	1.080
fp32	1024	1.129	1.127	1.127	1.122	1.124	1.121	1.123	1.092	1.100	1.061	1.071	1.051
fp32	2048	1.112	1.119	1.124	1.114	1.114	1.113	1.104	1.093	1.068	1.062	1.053	1.043
bf16	512	1.045	1.045	1.046	1.055	1.044	1.043	1.049	1.049	1.054	1.043	1.059	1.063
bf16	1024	1.054	1.052	1.052	1.054	1.055	1.052	1.049	1.048	1.044	1.035	1.036	1.060
bf16	2048	1.067	1.067	1.069	1.065	1.070	1.067	1.066	1.065	1.046	1.058	1.062	1.053
fp16	512	1.080	1.085	1.090	1.093	1.088	1.094	1.079	1.088	1.066	1.059	1.114	1.072
fp16	1024	1.144	1.154	1.154	1.153	1.149	1.152	1.147	1.149	1.164	1.172	1.185	1.192
fp16	2048	1.117	1.119	1.117	1.119	1.118	1.116	1.116	1.110	1.102	1.113	1.121	1.124

Overall B200: median 1.078×, 707/720 cells faster.

Takeaways:

Hardware-invariant. B200 and B300 ratios match cell-for-cell (overall 1.078× vs 1.077×), so no per-SKU re-tuning — the rank-scatter win reproduces identically on sm_100 and sm_103.
BS-invariant. Flat ~1.04–1.19× from BS=1 to BS=2048 across every (dtype, K): the Phase-4 barrier saving is per-row, independent of how many rows fill the GPU.
dtype/K/N shape. Peaks in the medium-N band (8K–32K, single-row hot spot at N=16K) and for K≥1024 / fp16·fp32; bf16 gains least but stays ≥1.0; never regresses materially (worst cells ≈ 0.996×, all at BS=1).

Performance — cluster + rank-scatter vs the multi-CTA cluster (PR #15198)

The two levers are orthogonal: multi-CTA-per-row (PR #15198) closes the SM-count axis on the full-N Phase-2/3 scans; rank-scatter cuts the leader's Phase-4. Stacked (synthetic grid, nsys cold-L2, full K×dtype×N×BS):

comparison	B300	B200
op8 (cluster + rank-scatter) vs GVR multi-CTA cluster (PR #15198)	1.090× (703/720)	1.088× (701/720)
op8 vs single-CTA snap baseline	1.138× (652/720)	1.132× (650/720)

On real DeepSeek-V3.2 K=2048 decode (N≈70.7K, B300) this is corroborated by a dedicated nsys run: op8 16.88 µs/call vs cluster 18.80 (1.114×) vs single-CTA snap 21.60 (1.28×), reaching parity with the fastest in-tree radix-cuteDSL (16.95 µs/call) while keeping GVR's exact-with-temporal-seed property.

All numbers are extracted from the op-bench full nsys re-test (indexer_topk_op_bench/report/{seqlen_data,bs_data}.csv, sourced from results_b200_nsys / results_b300_nsys), not produced by this PR's CI.

Correctness / testing

tests/unittest/_torch/test_gvr_topk_rank_scatter.py (Blackwell-gated, pure cuTe DSL — no engine build / model weights):

single-CTA rank-scatter (exact) vs masked torch.topk, tie-aware, across fp32/bf16/fp16 × K∈{512,1024,2048};
rank-scatter vs the default snap Phase-4 cross-check (both exact);
hybrid GvrTopKLBKernel with a long+short batch so prepare classifies 1 long (cluster branch) + 2 short (single-CTA branch) — exercising both Phase-4 paths in one launch.

All cases verified exact on B200 (sm_100); the correctness criterion is value-multiset equality, so boundary ties are accepted regardless of which tied index is selected.

Scope

✅ kernel implementation (single-CTA + cluster, default exact rank-scatter as of rev 2)
✅ basic precision tests
✅ performance characterization by dtype × top-K (BS=1 seq-len + BS-scaling), nsys cold-L2 on B200 + B300
✅ framework default switched to exact rank-scatter (production decode path)
✅ correctness fix for the exact fine-search (smem scratch aliasing)
⬜ DeepSeek-V3.2/V4 end-to-end model-level validation — follow-up PR

coderabbitai · 2026-06-29T07:42:45Z

📝 Walkthrough

Walkthrough

Adds a new phase4_rank_scatter Phase-4 implementation to GvrTopKKernel as an alternative to phase4_histogram_snap. Two new boolean flags (enable_p4_rank_scatter, enable_p4_rank_scatter_exact) control selection and tie-break exactness. The flags propagate through GvrTopKLBKernel. A new test module validates single-CTA and load-balance correctness using a tie-aware oracle.

Changes

GVR Top-K Phase-4 Rank-Scatter

Layer / File(s)	Summary
`phase4_rank_scatter` implementation and dispatch `tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py`	Adds `enable_p4_rank_scatter`/`enable_p4_rank_scatter_exact` constructor params and stores them on the instance. Adds the `phase4_rank_scatter` method (coarse histogram + bin-search + optional 256-bin fine-histogram recursion or approximate scatter). Wires conditional dispatch in `run_one_row`.
Load-balance kernel flag forwarding `tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode_load_balance.py`	`GvrTopKLBKernel.__init__` accepts the two new flags, forwards them into `common_kwargs` for both internal `GvrTopKKernel` instances, and stores them on `self`.
Correctness tests `tests/unittest/_torch/test_gvr_topk_rank_scatter.py`	New test module with SM-version gating, `_make_varlen_inputs`, `_assert_tie_aware_correct` oracle, `_run_single_cta`/`_run_lb` runners, and three test functions covering exact single-CTA, baseline comparison, and load-balance hybrid paths.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 60.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is concise, matches the main kernel change, and follows the repository's required ticket/type format.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Description check	✅ Passed	The PR description is detailed and covers the change, rationale, tests, performance, and scope, even though it doesn’t strictly follow the template headings.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

coderabbitai

Actionable comments posted: 1

🧹 Nitpick comments (1)

tests/unittest/_torch/test_gvr_topk_rank_scatter.py (1)
148-235: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Coverage is still insufficient around the fine-search boundary.

tests/unittest/_torch/test_gvr_topk_rank_scatter.py only uses random logits, so it does not reliably drive the exact fine search into the lowest fine-bin slice. Please add one deterministic _run_single_cta(...) case and one deterministic _run_lb(...) case whose K-boundary lands in that slice, so this PR covers the exact-path edge case directly. As per path instructions, coverage in tests/unittest/_torch/test_gvr_topk_rank_scatter.py is insufficient and should be expanded here.
🤖 Prompt for 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.

In `@tests/unittest/_torch/test_gvr_topk_rank_scatter.py` around lines 148 - 235,
Coverage for the exact fine-search boundary is missing, so add deterministic
tests that force the K-boundary into the lowest fine-bin slice instead of
relying on random logits. In test_gvr_topk_rank_scatter.py, extend
test_single_cta_rank_scatter_exact using _run_single_cta(...) with fixed inputs
that hit that slice, and add a matching _run_lb(...) case in
test_lb_hybrid_rank_scatter_exact so both the single-CTA path and the hybrid
path exercise the exact fine-search edge case directly. Use the existing helpers
_make_varlen_inputs, _run_single_cta, _run_lb, and _assert_tie_aware_correct to
keep the new cases targeted and deterministic.
Source: Path instructions

🤖 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 `@tensorrt_llm/_torch/cute_dsl_kernels/blackwell/top_k/gvr_topk_decode.py`:
- Around line 1930-1948: The fine-search logic in gvr_topk_decode is reusing
smem_hist bins 0 and 1 for metadata, which corrupts the histogram needed by the
reverse scan when twf points to the last fine warp. Update the fine-warp
selection path around the smem_hist writes and the subsequent pre_f/twf2 reads
so that prefix/target-warp metadata is stored in separate scratch space, not in
histogram bins used later by the exact top-K scan. Keep the rank-above, twf, and
pre computation in the same control flow, but preserve the original histogram
counts for sb_star and ra_fine.

---

Nitpick comments:
In `@tests/unittest/_torch/test_gvr_topk_rank_scatter.py`:
- Around line 148-235: Coverage for the exact fine-search boundary is missing,
so add deterministic tests that force the K-boundary into the lowest fine-bin
slice instead of relying on random logits. In test_gvr_topk_rank_scatter.py,
extend test_single_cta_rank_scatter_exact using _run_single_cta(...) with fixed
inputs that hit that slice, and add a matching _run_lb(...) case in
test_lb_hybrid_rank_scatter_exact so both the single-CTA path and the hybrid
path exercise the exact fine-search edge case directly. Use the existing helpers
_make_varlen_inputs, _run_single_cta, _run_lb, and _assert_tie_aware_correct to
keep the new cases targeted and deterministic.

🪄 Autofix (Beta)

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: 3e9d7015-202a-4626-a2fa-5fae4ab4a5a4

📥 Commits

Reviewing files that changed from the base of the PR and between 552f462 and 2e591f8.

📒 Files selected for processing (3)

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_load_balance.py
tests/unittest/_torch/test_gvr_topk_rank_scatter.py

…ch scratch In phase4_rank_scatter's exact fine-histogram recursion, the prefix-count and target-fine-warp index were staged in smem_hist[0]/[1]. When the kth-largest candidate inside the straddling coarse bin lands in the lowest fine warp (twf2 == num_warps-1), that warp's reverse bin scan walks fine bins down to 0 and 1 — reading the metadata back as if it were histogram counts and corrupting sb_star / ra_fine, so the exact path can return the wrong top-K set. Stage the two scalars in spare s_iscalars[4]/[1] instead; both are dead in this window (re-zeroed as cnt_above/cnt_strad before the scatter pass). smem_hist[2]/ [3] are only written after the scan loop, so they stay valid and need no change. Flagged by CodeRabbit on PR NVIDIA#15709. Made-with: Claude Code Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>

longcheng-nv · 2026-06-30T03:25:22Z

/bot run --disable-fail-fast

longcheng-nv · 2026-06-30T05:46:24Z

/bot run --disable-fail-fast

longcheng-nv · 2026-06-30T05:49:49Z

/bot run --disable-fail-fast

longcheng-nv · 2026-06-30T13:24:36Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-30T13:31:28Z

PR_Github #56612 [ run ] triggered by Bot. Commit: 994364f Link to invocation

tensorrt-cicd · 2026-06-30T17:29:00Z

PR_Github #56612 [ run ] completed with state SUCCESS. Commit: 994364f
/LLM/main/L0_MergeRequest_PR pipeline #45439 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-07-01T01:44:31Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-01T01:52:06Z

PR_Github #56782 [ run ] triggered by Bot. Commit: 994364f Link to invocation

tensorrt-cicd · 2026-07-01T03:54:15Z

PR_Github #56782 [ run ] completed with state SUCCESS. Commit: 994364f
/LLM/main/L0_MergeRequest_PR pipeline #45600 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-07-01T06:50:19Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-01T06:56:04Z

PR_Github #56869 [ run ] triggered by Bot. Commit: 994364f Link to invocation

longcheng-nv · 2026-07-01T08:54:39Z

Pushed 80281c2f0c: default Phase-4 back to the exact histogram-snap; rank-scatter is now opt-in.

Why: the 30 test_cute_dsl_gvr_topk_decode failures in L0 (build 45600) were the rank-scatter _exact path being non-exact on adversarial multi-bucket batches (large BS / cluster / varlen with uniform-random logits). A fixed-depth histogram cannot separate two distinct values that fall in the same sub-bin, so the straddling bin can emit a value below the true K-th rank. The single 256-bin fine recursion resolves the easy single-row cases (which is why test_gvr_topk_rank_scatter passed) but not the harder ones.

Why not just make rank-scatter exact: I prototyped a bounded value-edge fine-recursion loop that is exact on all 768 grid cases, but nsys cold-L2 (B200, median-timed) shows it is slower than the existing snap on every config — 1.06–1.08x on random, 1.04–1.24x on production temporal-coherent synth (V4 Pro K1024 N64K). Since snap is already exact, an exact rank-scatter that must iterate is strictly dominated by snap and buys nothing as a default.

This commit: enable_p4_rank_scatter now defaults to False (exact snap) in both GvrTopKKernel and GvrTopKLBKernel; docstrings corrected. rank-scatter stays available as an explicit opt-in fast path (~1.35x vs snap, exact on the production temporal-coherent distribution — verified on V4 Pro/Flash synth — just not on adversarial random). No kernel-algorithm change; test_gvr_topk_rank_scatter is unaffected.

Note: the other ~3900 failures in build 45600 were infrastructure (/root/.triton/cache missing, DeepGEMM JIT), unrelated to this PR.

longcheng-nv · 2026-07-01T08:54:44Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-01T09:00:50Z

PR_Github #56891 [ run ] triggered by Bot. Commit: 80281c2 Link to invocation

tensorrt-cicd · 2026-07-01T09:00:55Z

PR_Github #56869 [ run ] completed with state ABORTED. Commit: 994364f
/LLM/main/L0_MergeRequest_PR pipeline #45680 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-07-01T12:39:58Z

PR_Github #56891 [ run ] completed with state SUCCESS. Commit: 80281c2
/LLM/main/L0_MergeRequest_PR pipeline #45699 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

…nels Add an opt-in Phase-4 (snap) replacement to the cuTe DSL GVR Top-K decode kernels: a fused rank-and-scatter that resolves the K-th rank with one coarse histogram + an optional fixed-256-bin fine recursion, then scatters survivors directly to their output rank in a single pass. This collapses the iterative histogram-snap loop's block-barrier count (~14 -> ~7) while remaining exact (vdiff=0 on the realistic temporally-coherent indexer distribution). The new code lives in GvrTopKKernel.phase4_rank_scatter and is gated by two constructor flags (both default off): - enable_p4_rank_scatter : use rank-scatter instead of snap - enable_p4_rank_scatter_exact : add the fine recursion (exact tie-break) Because GvrTopKKernel.run_one_row is shared by the single-CTA path and the cluster leader-only path, the lever applies to both, and GvrTopKLBKernel forwards the two flags to its underlying cluster (long-row) and single-CTA (short-row) instances. The candidate set is fully resident in the leader CTA's SMEM after the DSMEM gather, so the cluster path is structurally identical to the single-CTA path. Measured benefit (standalone op-bench, cold-L2, CUDA-graph replay): - Single-CTA (vs snap baseline): B300 1.067x median (win ~98%), peaking at small/medium seq-len (N=8K-32K ~1.10x) where Phase-4 dominates the kernel; tapering to ~1.03x at N>=128K. Mixed on B200 at low BS (opt-in). - Cluster + rank-scatter (vs PR NVIDIA#15198 multi-CTA cluster): B300 1.066x cold-L2 median (real V3.2 K=2048), 1.114x by nsys; 1.17-1.28x vs the single-CTA snap baseline. No framework integration or DeepSeek V3.2/V4 end-to-end wiring in this PR; that follows separately. Adds tests/unittest/_torch/test_gvr_topk_rank_scatter.py covering the single-CTA and hybrid load-balance paths (exact vs masked torch.topk, tie-aware) across fp32/bf16/fp16 and K=512/1024/2048. Made-with: Claude Code Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>

…ch scratch In phase4_rank_scatter's exact fine-histogram recursion, the prefix-count and target-fine-warp index were staged in smem_hist[0]/[1]. When the kth-largest candidate inside the straddling coarse bin lands in the lowest fine warp (twf2 == num_warps-1), that warp's reverse bin scan walks fine bins down to 0 and 1 — reading the metadata back as if it were histogram counts and corrupting sb_star / ra_fine, so the exact path can return the wrong top-K set. Stage the two scalars in spare s_iscalars[4]/[1] instead; both are dead in this window (re-zeroed as cnt_above/cnt_strad before the scatter pass). smem_hist[2]/ [3] are only written after the scan loop, so they stay valid and need no change. Flagged by CodeRabbit on PR NVIDIA#15709. Made-with: Claude Code Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>

…am-snap Flip enable_p4_rank_scatter / enable_p4_rank_scatter_exact to True by default in GvrTopKKernel and GvrTopKLBKernel. Exact rank-scatter is a drop-in for the snap Phase-4 (both exact, vdiff=0) and the faster path (nsys cold-L2 median ~1.077x B300 / ~1.078x B200, 704-707/720 cells faster, HW- and BS-invariant). _exact is defaulted on so the default stays exact like snap; the non-exact rank-scatter emits the straddling bin in arbitrary order and is opt-in via enable_p4_rank_scatter_exact=False. The cute_dsl_custom_ops construction sites do not pass these flags, so the production decode path now uses exact rank-scatter for Phase-4. Made-with: Claude Code Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>

CI Pre-commit Check failed on the PR files: - codespell flagged the fake-output tensor local `fo` in test_gvr_topk_rank_scatter.py (4 sites) -> renamed to `fout`. - ruff (1 auto-fix) + ruff-format normalized gvr_topk_decode.py and the test (the base commit's additions were not ruff-formatted). No semantic change to the kernels; the Phase-4 rank-scatter scratch-aliasing fix (s_iscalars[4]/[1]) is unchanged. All three hooks pass locally. Made-with: Claude Code Co-Authored-By: Claude <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>

…catter opt-in The Phase-4 rank-scatter variant is not exact on arbitrary continuous logits: a fixed-depth histogram cannot separate two distinct values that land in the same (sub-)bin, so the straddling bin can emit a value below the true K-th rank. This caused 30 test_cute_dsl_gvr_topk_decode failures (adversarial multi-bucket batches: large BS / cluster / varlen with uniform-random logits) where the single 256-bin fine recursion is insufficient. Making rank-scatter exact for all continuous inputs requires iterating to a real-value threshold (a bounded value-edge fine-recursion loop), which nsys cold-L2 (B200) measures as slower than the existing exact histogram-snap on every config (1.04-1.24x on production temporal-coherent synth, 1.06-1.08x on random) — so an exact rank-scatter is strictly dominated by snap and buys nothing as a default. Therefore flip the default to the exact histogram-snap and keep rank-scatter as an explicit opt-in fast path (enable_p4_rank_scatter=True), which is ~1.35x faster than snap and exact on the production temporally-coherent decode distribution (verified on V4 Pro/Flash synth), just not on adversarial random inputs. Correct the docstrings that claimed the rank-scatter default was a drop-in exact replacement. No kernel-algorithm change; only the default variant selection and docs. The opt-in rank-scatter tests (test_gvr_topk_rank_scatter) are unaffected. Made-with: Claude Code Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: longcheng-nv <243710427+longcheng-nv@users.noreply.github.com>

tburt-nv · 2026-07-01T18:09:14Z

/bot run --disable-fail-fast

tburt-nv · 2026-07-01T18:09:32Z

Rebased to pick up #15713

tensorrt-cicd · 2026-07-01T18:16:35Z

PR_Github #56996 [ run ] triggered by Bot. Commit: ec04147 Link to invocation

tensorrt-cicd · 2026-07-01T21:57:22Z

PR_Github #56996 [ run ] completed with state SUCCESS. Commit: ec04147
/LLM/main/L0_MergeRequest_PR pipeline #45796 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 requested a review from a team as a code owner June 29, 2026 07:37

longcheng-nv requested review from byshiue and syuoni June 29, 2026 07:37

github-actions Bot assigned longcheng-nv Jun 29, 2026

coderabbitai Bot reviewed Jun 29, 2026

View reviewed changes

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

longcheng-nv requested review from lfr-0531, limin2021 and mingyangHao and removed request for byshiue and syuoni June 30, 2026 01:52

longcheng-nv added the deepseek-v4 label Jun 30, 2026

longcheng-nv and others added 4 commits July 1, 2026 14:09

tburt-nv force-pushed the feat/gvr-rank-scatter-p4 branch from 80281c2 to ec04147 Compare July 1, 2026 18:09

Uh oh!

Conversation

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

Uh oh!

Summary

Update (rev 2)

What changes

Why this helps (and where)

Performance — single-CTA rank-scatter vs the snap baseline (same kernel, only Phase-4 swapped)

B300 (sm_103)

B200 (sm_100) — essentially identical (hardware-invariant)

Performance — cluster + rank-scatter vs the multi-CTA cluster (PR #15198)

Correctness / testing

Scope

Uh oh!

coderabbitai Bot commented Jun 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Estimated code review effort

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

longcheng-nv commented Jun 30, 2026

Uh oh!

longcheng-nv commented Jun 30, 2026

Uh oh!

longcheng-nv commented Jun 30, 2026

Uh oh!

longcheng-nv commented Jun 30, 2026

Uh oh!

tensorrt-cicd commented Jun 30, 2026

Uh oh!

tensorrt-cicd commented Jun 30, 2026

Uh oh!

longcheng-nv commented Jul 1, 2026

Uh oh!

tensorrt-cicd commented Jul 1, 2026

Uh oh!

tensorrt-cicd commented Jul 1, 2026

Uh oh!

longcheng-nv commented Jul 1, 2026

Uh oh!

tensorrt-cicd commented Jul 1, 2026

Uh oh!

longcheng-nv commented Jul 1, 2026

Uh oh!

longcheng-nv commented Jul 1, 2026

Uh oh!

tensorrt-cicd commented Jul 1, 2026

Uh oh!

tensorrt-cicd commented Jul 1, 2026

Uh oh!

tensorrt-cicd commented Jul 1, 2026

Uh oh!

tburt-nv commented Jul 1, 2026

Uh oh!

tburt-nv commented Jul 1, 2026

Uh oh!

tensorrt-cicd commented Jul 1, 2026

Uh oh!

tensorrt-cicd commented Jul 1, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

longcheng-nv commented Jun 29, 2026 •

edited

Loading

coderabbitai Bot commented Jun 29, 2026 •

edited

Loading