[longcheng-nv]

Description

Fix excessive overhead in deterministic GVR V2 top-K, especially at K=1024 and K=2048. Follow-up to #20100.

With clustered logits, a few warps receive most of the candidates. Their local buffers can overflow even when the shared buffer has enough total space, triggering an expensive scan of the whole row.

This fix:

• Recovers overflowing candidates in parallel using exact per-warp counts.
• Adds candidate-buffer headroom without changing sampling thresholds.
• Removes unused shared-memory allocations and adjusts buffer sizes to limit memory pressure.

Measured improvement

At BS=512 and N≈131,000, added cost drops from 208–221 µs to 1.05–2.76 µs.

Each row compares identical input tensors on the same B200, using nsys + cold L2. Overhead is relative to the default-path control in each run.

Model / K	Shape (BS × N)	Kernel time: before → after	Overhead: before → after
V4 Pro · K=1024	512 × 32,771	53.01 → 30.23 µs	+83.63% → +6.02%
V4 Pro · K=1024	512 × 131,075	300.57 → 94.25 µs	+225.04% → +1.13%
V3.2 · K=2048	512 × 32,783	63.21 → 33.80 µs	+106.67% → +11.31%
V3.2 · K=2048	512 × 131,087	295.88 → 77.19 µs	+296.63% → +3.71%

Final performance across the full grid

B200 · nsys + cold L2 · 275 captured-logit shapes · all 11 power-of-two batch sizes from 1 to 1024.

Overhead below is relative to the default path, using deterministic=True with per-route output order.

Model	K	Geometric mean overhead	Maximum added time
DeepSeek V4 Flash	512	+8.94%	+2.63 µs
DeepSeek V4 Pro	1024	+9.13%	+5.47 µs
DeepSeek V3.2	2048	+14.12%	+7.68 µs

Inputs are model logits captured from random-token prompts, with stored rows reused to form batches.

Test Coverage

• Correctness and repeatability: 275 captured-logit and 84 synthetic shapes passed; 48 shapes were also checked on another GPU.
• Regression checks: 54 kernel tests, 24 boundary cases and 45 default-path comparisons passed.
• Memory and style: Compute Sanitizer reported 0 device-memory errors; all applicable pre-commit checks passed.

Measurement and validation details

• Before/after provenance: before values are from the archived pre-fix snapshot 2807acd470e6; after values use the final source matching PR commit 358db83774. Input hashes, GPU identity and timing settings were checked for every comparison.
• Timing: CUDA graphs, 5 warmups and 24 repetitions per arm. Two default-path controls and the deterministic arm run in rotated order on the same GPU. The baseline is the mean of the two controls.
• Cold L2: a 512 MiB eviction runs before every replay, outside the measured range.
• Scope: measurements cover per-route deterministic output; optional canonical ascending-index sorting is excluded. Geometric means weight shapes equally.
• Regression tests: the actual kernel sources run through a standalone package shim, without importing the full TensorRT-LLM runtime or repository conftest. Coverage includes clustered scores, large batches, odd valid lengths and misaligned prefill windows.
• Sanitizer: device-memory checks remain enabled. API-error reporting is disabled for version-negotiation diagnostics between CUDA 13.3 bindings and the 13.2 driver.

Focused regression command in an installed TensorRT-LLM environment on B200:

pytest tests/unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_deterministic.py -k clustered

PR Checklist

• Regression coverage is included and local checks passed.

• API, dependency and ownership changes: none.

• CI results and reviewer approvals are tracked in GitHub checks and reviews.

• Please check this after reviewing the above items as appropriate for this PR.

[longcheng-nv]

/bot run

[tensorrt-cicd]

PR_Github #77809 [ run ] triggered by Bot. Commit: 358db83 Link to invocation

[tensorrt-cicd]

PR_Github #77809 [ run ] completed with state SUCCESS. Commit: 358db83
/LLM/main/L0_MergeRequest_PR pipeline #64144 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 --reuse-test

[tensorrt-cicd]

PR_Github #77958 [ run ] triggered by Bot. Commit: 358db83 Link to invocation

[tensorrt-cicd]

PR_Github #77958 [ run ] completed with state FAILURE. Commit: 358db83
/LLM/main/L0_MergeRequest_PR pipeline #64266 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 --reuse-test --disable-fail-fast

[tensorrt-cicd]

PR_Github #78022 [ run ] triggered by Bot. Commit: 358db83 Link to invocation

[tensorrt-cicd]

PR_Github #78022 [ run ] completed with state FAILURE. Commit: 358db83
/LLM/main/L0_MergeRequest_PR pipeline #64314 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 --reuse-test --disable-fail-fast

[tensorrt-cicd]

PR_Github #78270 [ run ] triggered by Bot. Commit: 358db83 Link to invocation

[tensorrt-cicd]

PR_Github #78270 [ run ] completed with state SUCCESS. Commit: 358db83
/LLM/main/L0_MergeRequest_PR pipeline #64517 completed with status: 'SUCCESS'

CI Report

Link to invocation

[github-actions]

Automatically added "ci: full pre-merge approved" because this PR has satisfied the required GitHub review approvals. Unresolved review conversations and other required checks remain independent merge requirements.