[None][feat] DSv4 prep: attention op plumbing

[lfr-0531]

Summary by CodeRabbit

• New Features
  • Added support for DeepSeek-V4 sparse-MLA (Multi-Head Latent Attention) algorithm with optimized attention handling for sparse patterns.
  • Enhanced speculative decoding integration with improved metadata tracking.

Description

This PR adds the remaining Python-side attention-op plumbing needed for the DSv4 split:

• exposes optional sparse-MLA handoff fields on AttentionSparseArgs
• forwards sparse_mla_topk_lens and compressed_kv_cache_pool_ptr into thop.attention
• populates those fields for sparse configs with algorithm == "deepseek_v4" using duck-typed metadata/config fields, without importing DeepSeekV4SparseAttentionConfig or any DSv4 model/backend/cache-manager modules
• explicitly declares TrtllmAttentionMetadata.max_total_draft_tokens so the attention-op sync test can validate the thop call site without a dynamic getattr

This is intentionally limited to shared TRTLLM attention backend/interface plumbing. The DSv4 sparse backend and cache-manager implementation are left for the later split PR that consumes these fields.

Test Coverage

• Build/install:
  • python3 ./scripts/build_wheel.py --trt_root /usr/local/tensorrt --benchmarks --use_ccache --cuda_architectures "90-real;100-real" --configure_cmake
  • .venv-3.12/bin/python -m pip install --force-reinstall --no-deps build/tensorrt_llm-1.3.0rc18-cp312-cp312-linux_x86_64.whl
• Import/scope:
  • import smoke confirmed tensorrt_llm, tensorrt_llm._torch.attention_backend.trtllm, and tensorrt_llm.bindings.internal.thop load from this worktree / installed wheel
  • negative import assertion confirmed no modeling_deepseekv4, DSv4 sparse backend, or DSv4 sparse cache-manager module is loaded by importing trtllm.py
  • scope grep confirmed only tensorrt_llm/_torch/attention_backend/interface.py and tensorrt_llm/_torch/attention_backend/trtllm.py are changed, with no compressor/mHC, IndexerTopK/TopK, MoE, DSv4 model/tokenizer, fusion/custom-op, or sparse-backend implementation files
• Tests:
  • CUDA_VISIBLE_DEVICES=0 timeout 1800 .venv-3.12/bin/python -m pytest -q --tb=short -ra tests/unittest/_torch/attention/sparse/test_cpp_custom_ops.py
    • 43 passed, 2 warnings in 2.33s
  • CUDA_VISIBLE_DEVICES=0 timeout 600 .venv-3.12/bin/python -m pytest -q --tb=short -ra tests/unittest/_torch/attention_backend/test_attention_op_sync.py
    • 7 passed, 2 warnings in 1.25s

Also attempted tests/unittest/_torch/attention/test_attention_mla.py on idle GPU 0. It failed in existing TRTLLM-Gen FMHA NVRTC preprocessing at step 1 with could not open source file "cuda.h" (no directories in search list). Retrying the failing case with CPATH=/usr/local/cuda-13.1/targets/x86_64-linux/include produced the same NVRTC include-path error, so this is recorded as an environment/runtime include issue rather than a PR-4 plumbing assertion failure.

Build note: attribution generation warned ninja -t inputs returned no results for wheel targets, but the wheel build exited 0.

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]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 565c8220-1b9a-49bc-9a15-ca9e319b0d6d

📥 Commits

Reviewing files that changed from the base of the PR and between 0d4bab9 and 82536f0.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/attention_backend/interface.py
• tensorrt_llm/_torch/attention_backend/trtllm.py

---

📝 Walkthrough

Walkthrough

Two new optional fields (sparse_mla_topk_lens, compressed_kv_cache_pool_ptr) are added to AttentionSparseArgs. TrtllmAttentionMetadata gains max_total_draft_tokens. A new "deepseek_v4" branch in TrtllmAttention.forward computes per-layer compression ratios and populates those sparse args from metadata. The thop.attention call in _run is rewired to forward those values instead of hardcoded Nones.

Changes

DeepSeek-V4 Sparse-MLA Support

Layer / File(s)	Summary
Data contracts: AttentionSparseArgs fields, metadata, and THOP literals tensorrt_llm/_torch/attention_backend/interface.py, tensorrt_llm/_torch/attention_backend/trtllm.py	AttentionSparseArgs adds sparse_mla_topk_lens: Optional[torch.Tensor] and compressed_kv_cache_pool_ptr: Optional[int]; TrtllmAttentionMetadata adds max_total_draft_tokens: Optional[int]; _THOP_LITERALS is changed to an empty dict, removing hardcoded None stubs for those keys.
DeepSeek-V4 forward branch and THOP call rewiring tensorrt_llm/_torch/attention_backend/trtllm.py	TrtllmAttention.forward adds a "deepseek_v4" sparse attention branch that selects a per-layer ratio, slices metadata.sparse_mla_topk_lens[ratio] by input layout, updates metadata.num_sparse_topk, and conditionally sets compressed_kv_cache_pool_ptr. _run then passes those values to thop.attention and reads spec_decoding_target_max_draft_tokens directly from metadata.max_total_draft_tokens.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.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 clearly summarizes the main change: attention op plumbing preparation for DeepSeek-V4 sparse-MLA support.
Description check	✅ Passed	The PR description is comprehensive with clear sections: description of changes, detailed test coverage results, and completed PR checklist addressing all repository requirements.
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.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[lfr-0531]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54390 [ run ] triggered by Bot. Commit: 82536f0 Link to invocation

[heyuhhh]

It totally looks good to me with little comments: it would be better if we can just keep forward_args process in this PR, many other operations can be lived in V4 metadata/attention.

[lfr-0531]

/bot kill

[tensorrt-cicd]

PR_Github #54427 [ kill ] triggered by Bot. Commit: 82536f0 Link to invocation

[tensorrt-cicd]

PR_Github #54390 [ run ] completed with state ABORTED. Commit: 82536f0

Link to invocation

[tensorrt-cicd]

PR_Github #54427 [ kill ] completed with state SUCCESS. Commit: 82536f0
Successfully killed previous jobs for commit 82536f0

Link to invocation

[lfr-0531]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54453 [ run ] triggered by Bot. Commit: 9e86d65 Link to invocation

[tensorrt-cicd]

PR_Github #54453 [ run ] completed with state FAILURE. Commit: 9e86d65
/LLM/main/L0_MergeRequest_PR pipeline #43517 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

[lfr-0531]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54491 [ run ] triggered by Bot. Commit: 9e86d65 Link to invocation

[lfr-0531]

/bot kill

[tensorrt-cicd]

PR_Github #54544 [ kill ] triggered by Bot. Commit: 9e86d65 Link to invocation

[tensorrt-cicd]

PR_Github #54491 [ run ] completed with state ABORTED. Commit: 9e86d65

Link to invocation

[tensorrt-cicd]

PR_Github #54544 [ kill ] completed with state SUCCESS. Commit: 9e86d65
Successfully killed previous jobs for commit 9e86d65

Link to invocation

[lfr-0531]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54549 [ run ] triggered by Bot. Commit: bfacdce Link to invocation

[lfr-0531]

/bot kill

[tensorrt-cicd]

PR_Github #54551 [ kill ] triggered by Bot. Commit: 1a5e6d6 Link to invocation

[tensorrt-cicd]

PR_Github #54549 [ run ] completed with state ABORTED. Commit: bfacdce

Link to invocation

[tensorrt-cicd]

PR_Github #54551 [ kill ] completed with state SUCCESS. Commit: 1a5e6d6
Successfully killed previous jobs for commit 1a5e6d6

Link to invocation

[lfr-0531]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54565 [ run ] triggered by Bot. Commit: edcd002 Link to invocation

[tensorrt-cicd]

PR_Github #54565 [ run ] completed with state SUCCESS. Commit: edcd002
/LLM/main/L0_MergeRequest_PR pipeline #43607 completed with status: 'SUCCESS'

CI Report

Link to invocation