[None][feat] top-k: route decode to CuTe DSL GVR top-k in e2e

[limin2021]

Wire the existing CuTe DSL GVR (Guess-Verify-Refine) top-k decode op into the DSA indexer decode dispatch. The flag combo
use_cute_dsl_topk=True AND enable_heuristic_topk=True now routes to cute_dsl_gvr_topk_decode, covering all compress_ratio and next_n values.

• DSAMetadataParams: add use_cute_dsl_topk; propagate through all three to_sparse_metadata_params / make_deepseek_v4_sparse_metadata_params sites.
• Metadata prepare(): compute the LJF row-reorder (argsort of gen kv_lens, descending) only when num_generations >= num_sms, writing into a pre-allocated stable buffer so CUDA Graph replay reads a valid permutation instead of a freed fresh allocation.
• Indexer decode dispatch: add GVR DSL as the first branch; keep the radix DSL and C++ scheduler branches as fallbacks.
• Skip the redundant radix warmup when the GVR path is active.
• Integration test: add cute_dsl_gvr_mtp1 case to test_fp8_blockscale.

Summary by CodeRabbit

• New Features

  • Added an optional CuTE DSL TopK decoding path for supported sparse attention configurations.
  • Improved heuristic TopK ordering with CUDA-graph-friendly row reordering.
  • Propagated the CuTE DSL TopK setting consistently through DeepSeek sparse attention configuration.

• Tests

  • Expanded integration coverage for heuristic and CuTE DSL TopK combinations.
  • Added architecture-aware handling for supported hardware.

Description

Test Coverage

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.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[coderabbitai]

📝 Walkthrough

Walkthrough

The change propagates use_cute_dsl_topk through sparse-attention metadata, maintains a stable KV-length row-reorder buffer for heuristic TopK, dispatches the CuTe DSL GVR decoder, and expands DeepSeek integration coverage.

Changes

CuTe DSL heuristic TopK

Layer / File(s)	Summary
Configuration propagation and metadata contract tensorrt_llm/_torch/attention_backend/sparse/dsa.py, tensorrt_llm/llmapi/llm_args.py, tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py	Metadata parameter constructors now propagate use_cute_dsl_topk, and DSAMetadataParams declares the flag.
Stable heuristic reorder state tensorrt_llm/_torch/attention_backend/sparse/dsa.py	DSA metadata conditionally allocates a stable reorder buffer and updates it from KV-length ordering during generation updates.
GVR dispatch and integration coverage tensorrt_llm/_torch/attention_backend/sparse/dsa.py, tests/integration/defs/accuracy/test_llm_api_pytorch.py	Indexer dispatch selects cute_dsl_gvr_topk_decode for heuristic CuTe mode, and DeepSeek tests cover the new flag combinations and architecture guard.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Config
  participant DSAtrtllmAttentionMetadata
  participant Indexer
  participant cute_dsl_gvr_topk_decode
  Config->>DSAtrtllmAttentionMetadata: propagate use_cute_dsl_topk
  DSAtrtllmAttentionMetadata->>DSAtrtllmAttentionMetadata: update kv_lens_row_reorder
  DSAtrtllmAttentionMetadata->>Indexer: provide row reorder
  Indexer->>cute_dsl_gvr_topk_decode: pass order_row and TopK parameters
  cute_dsl_gvr_topk_decode-->>Indexer: produce heuristic TopK results

Loading

Suggested reviewers: pengbowang-nv, yunruis, zhaoyangwang-nvidia

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is concise, specific, and accurately reflects the main CuTe DSL GVR top-k routing change.
Description check	✅ Passed	The description covers the problem, solution, and testing at a high level, though it does not explicitly fill every template section.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.

✨ 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

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)

tensorrt_llm/_torch/attention_backend/sparse/dsa.py (2)

1791-1805: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win

Warm up the GVR kernel before capture. warmup_heuristic_topk_decode() covers the other top-k runners, but it never precompiles CuteDSLGvrTopKDecodeRunner._compile/_compile_lb. Add a GVR warmup in the same eager path so the first cute_dsl_gvr_topk_decode call can’t JIT inside CUDA graph capture.

🤖 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 `@tensorrt_llm/_torch/attention_backend/sparse/dsa.py` around lines 1791 -
1805, Extend the layer_idx == 0 eager warmup path around
warmup_heuristic_topk_decode to also precompile the GVR top-k decode runner,
covering both CuteDSLGvrTopKDecodeRunner._compile and _compile_lb. Invoke the
existing GVR warmup mechanism before CUDA Graph capture, while preserving the
current conditional behavior and top-k configuration.

---

1788-1789: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Gate GVR Top-K decode on SM100/103 only

_enable_heuristic_topk at 1788-1789 is too broad for the new cute_dsl_gvr_topk_decode branch at 2785-2803. That op only supports SM100/103, so SM120/121 can reach it and raise at runtime instead of falling back. Narrow the dispatch with is_sm_100f() (or get_sm_version() in (100, 103)).

🤖 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 `@tensorrt_llm/_torch/attention_backend/sparse/dsa.py` around lines 1788 -
1789, Narrow the _enable_heuristic_topk condition to enable heuristic Top-K only
for SM100/103 architectures, using is_sm_100f() or an equivalent
get_sm_version() check. Preserve the existing
sparse_params.enable_heuristic_topk requirement so SM120/121 fall back instead
of entering cute_dsl_gvr_topk_decode.

🧹 Nitpick comments (1)

tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py (1)

239-258: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Two parallel builders construct DeepSeekV4MetadataParams with duplicated field-mapping logic. This PR had to add use_cute_dsl_topk=... in both places to keep them in sync — a concrete DRY violation that risks future drift if only one path gets updated.

• tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py#L239-L258: consider having make_deepseek_v4_sparse_metadata_params() delegate to (or be replaced by) DeepSeekV4SparseAttentionConfig.to_sparse_metadata_params(), or extract one shared field-mapping helper both call sites use, so new fields only need to be wired once. Verify whether this function still has an active caller distinct from to_sparse_metadata_params() before merging the two.
• tensorrt_llm/llmapi/llm_args.py#L1060-L1069: keep as the canonical Pydantic-config-driven builder; if make_deepseek_v4_sparse_metadata_params() is confirmed dead/redundant, remove it, otherwise route it through this method.

🤖 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 `@tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py`
around lines 239 - 258, Unify the duplicated DeepSeek V4 metadata field mapping
by making make_deepseek_v4_sparse_metadata_params() delegate to
DeepSeekV4SparseAttentionConfig.to_sparse_metadata_params(), or extract a shared
helper used by both builders; first verify whether the standalone function has
an active distinct caller and remove it if redundant. In
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py:239-258,
route the builder through the shared canonical mapping or remove it; in
tensorrt_llm/llmapi/llm_args.py:1060-1069, retain to_sparse_metadata_params() as
the canonical Pydantic-config-driven builder and ensure both paths use it.

🤖 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/integration/defs/accuracy/test_llm_api_pytorch.py`:
- Around line 3379-3387: Update the dsa_config construction around
disable_skip_indexer, enable_heuristic_topk, and use_cute_dsl_topk to accumulate
keyword arguments before creating a single DeepSeekSparseAttentionConfig,
preserving skip_indexer_for_short_seqs=False alongside heuristic or cute-DSl
settings when flags are combined. Mirror the existing dsa_kwargs pattern used by
test_nvfp4_multi_gpus.

---

Outside diff comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/dsa.py`:
- Around line 1791-1805: Extend the layer_idx == 0 eager warmup path around
warmup_heuristic_topk_decode to also precompile the GVR top-k decode runner,
covering both CuteDSLGvrTopKDecodeRunner._compile and _compile_lb. Invoke the
existing GVR warmup mechanism before CUDA Graph capture, while preserving the
current conditional behavior and top-k configuration.
- Around line 1788-1789: Narrow the _enable_heuristic_topk condition to enable
heuristic Top-K only for SM100/103 architectures, using is_sm_100f() or an
equivalent get_sm_version() check. Preserve the existing
sparse_params.enable_heuristic_topk requirement so SM120/121 fall back instead
of entering cute_dsl_gvr_topk_decode.

---

Nitpick comments:
In `@tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py`:
- Around line 239-258: Unify the duplicated DeepSeek V4 metadata field mapping
by making make_deepseek_v4_sparse_metadata_params() delegate to
DeepSeekV4SparseAttentionConfig.to_sparse_metadata_params(), or extract a shared
helper used by both builders; first verify whether the standalone function has
an active distinct caller and remove it if redundant. In
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py:239-258,
route the builder through the shared canonical mapping or remove it; in
tensorrt_llm/llmapi/llm_args.py:1060-1069, retain to_sparse_metadata_params() as
the canonical Pydantic-config-driven builder and ensure both paths use it.

🪄 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: f9cda4a4-7d17-4a2e-bcad-b6f8d4a49d24

📥 Commits

Reviewing files that changed from the base of the PR and between 46bf21b and 50b901d.

📒 Files selected for processing (4)

• tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py
• tensorrt_llm/_torch/attention_backend/sparse/dsa.py
• tensorrt_llm/llmapi/llm_args.py
• tests/integration/defs/accuracy/test_llm_api_pytorch.py

[limin2021]

/bot run

[tensorrt-cicd]

PR_Github #59437 [ run ] triggered by Bot. Commit: 20cb47b Link to invocation

[limin2021]

/bot run --disable-fail-fast --add-multi-gpu-test

[tensorrt-cicd]

PR_Github #60663 [ run ] triggered by Bot. Commit: b782847 Link to invocation

[limin2021]

/bot run --disable-fail-fast --add-multi-gpu-test

[tensorrt-cicd]

PR_Github #60663 [ run ] completed with state FAILURE. Commit: b782847
/LLM/main/L0_MergeRequest_PR pipeline #48961 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 #60699 [ run ] triggered by Bot. Commit: 329dd6d Link to invocation

[lfr-0531]

LGTM~

[tensorrt-cicd]

PR_Github #60699 [ run ] completed with state FAILURE. Commit: 329dd6d
/LLM/main/L0_MergeRequest_PR pipeline #48990 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]

LGTM. All five review threads are resolved and verified in code: the row-count engage threshold fix (+ its new unit test), symmetric IS_CUTLASS_DSL_AVAILABLE gating, removal of the duplicate _compute_kv_lens_row_reorder call, the mtp3 GVR test case, and the exact revert of the accidental nvvm.fmin change (its 1500+ CuTe-DSL MoE failures cleared in L0 #48990).

Approving pending a green CI re-run: the 5 remaining failures on 329dd6d (A100X llmapi x2, B300 attention 'terminated unexpectedly' x2, DGX_B200 Qwen3 accuracy) are recurring infra flakes seen on unrelated PRs the same day and don't touch this PR's DSA dispatch path — please re-trigger before merge.

[limin2021]

/bot run --disable-fail-fast --add-multi-gpu-test

[mingyangHao]

LGTM

[limin2021]

/bot run --disable-fail-fast --add-multi-gpu-test

[dominicshanshan]

/bot run --disable-fail-fast --add-multi-gpu-test

[tensorrt-cicd]

PR_Github #60893 [ run ] triggered by Bot. Commit: 329dd6d Link to invocation

[tensorrt-cicd]

PR_Github #60893 [ run ] completed with state FAILURE. Commit: 329dd6d
/LLM/main/L0_MergeRequest_PR pipeline #49162 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

[limin2021]

/bot reuse-pipeline

[tensorrt-cicd]

PR_Github #60950 [ reuse-pipeline ] triggered by Bot. Commit: 329dd6d Link to invocation

[tensorrt-cicd]

PR_Github #60950 [ reuse-pipeline ] completed with state SUCCESS. Commit: 329dd6d
Can't reuse PR_Github #60893 with status: FAILED

Link to invocation

[limin2021]

/bot run --disable-fail-fast --add-multi-gpu-test

[tensorrt-cicd]

PR_Github #60967 [ run ] triggered by Bot. Commit: 329dd6d Link to invocation

[tensorrt-cicd]

PR_Github #60967 [ run ] completed with state FAILURE. Commit: 329dd6d
/LLM/main/L0_MergeRequest_PR pipeline #49229 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

[limin2021]

/bot run --disable-fail-fast --add-multi-gpu-test

[tensorrt-cicd]

PR_Github #61281 [ run ] triggered by Bot. Commit: 329dd6d Link to invocation

[tensorrt-cicd]

PR_Github #61281 [ run ] completed with state SUCCESS. Commit: 329dd6d
/LLM/main/L0_MergeRequest_PR pipeline #49515 completed with status: 'SUCCESS'

CI Report

Link to invocation