[yuxianq]

@coderabbitai summary

Description

Follow-up to #20301, which merged into release/1.3 on 2026-10-16 as f8033a06da456453fd7cc6aff339adf0ef0ea6e5. This PR is now rebased onto that release tip and contains only the six-file test/documentation follow-up in 5b0228a954. The original cherry-pick PR remains unchanged.

Addresses these review items without changing the vendored kernels, vendor lock, patch, runtime adapters, or dependencies:

• Item 1: GPU test selection: replace the block-sparse module's blanket cpu_only marker with per-test markers. CPU/mock cases remain in CPU CI; real GPU cases are eligible for the existing B200/B300 attention stages and retain their SM100/SM103 guards.
• Item 2: partial proxy blocks: weight proxy means by the number of tokens represented, clamp partial exact blocks, and add a 160-token replay case alongside 192 tokens. A CPU reference regression covers the final block as either proxy or exact. Also addresses the matching CodeRabbit finding.
• Item 4: lazy rescale coverage: add deterministic FP16/BF16 paged-context inputs with a high-score final key page after earlier KV tiles. Scaled log2 score jumps of 7 and 9 exercise both sides of the kernel's threshold of 8. Eager and CUDA-graph replay outputs are compared with an independent FP32 softmax reference.
• Item 5: wrapper API contracts: autospec decode and MLA wrappers in plan/run, cache, shared-workspace, and capture tests while retaining explicit argument assertions.
• Item 6: summary documentation: state in all three attention guides that both K and V summaries are means over present tokens, excluding final-block padding; the kernel supplies proxy token-count weighting.

Item 3 (the release scope of the original refresh) is intentionally left to the maintainer and is not addressed here.

Test Coverage

• All applicable pre-commit hooks passed, including lint, formatting, test-list validation, vendor-sync verification, and DCO checks. git diff --check passed.
• 15 focused local checks passed using FlashInfer 0.6.18 and the worktree's real wrapper classes. These load selected definitions without TRT-LLM's native bootstrap: four existing adapter cache/workspace/capture tests, four new proxy-reference cases, marker placement, two autospec rejection checks, and four rescale-input math checks. This is not a full repository pytest run or GPU validation.
• Confirmed the new 160-token reference cases fail against the parent helper: a partial proxy produces 0.3333 rather than 0.2, and a partial exact block raises IndexError. Both now produce 0.2.
• Full repository pytest collection was attempted but is blocked locally by the fresh worktree's missing tensorrt_llm.bindings. The local GPU is SM86/A40, so SM100/SM103 execution is pending. This PR remains draft pending Blackwell validation.
• Existing l0_cpu.yml, l0_b200.yml, and l0_b300.yml already include the attention test directory; no test-list expansion is required.

Relevant CI commands in a built TRT-LLM environment:

python3 -m pytest tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py -m cpu_only
python3 -m pytest tests/unittest/_torch/attention/test_prims_ts_fmha.py
python3 -m pytest tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py -m 'not cpu_only'
python3 -m pytest tests/unittest/_torch/attention/test_prims_ts_attention_backend.py -k lazy_rescale

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).

• No public API, dependency, ownership, or architecture changes in this follow-up.

• Documentation updated as needed.

• Required review and Blackwell execution remain pending.

• 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.

[yuxianq]

/bot run

[tensorrt-cicd]

PR_Github #78017 [ run ] triggered by Bot. Commit: 3ad9117 Link to invocation

[yuxianq]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #78030 [ run ] triggered by Bot. Commit: 3ad9117 Link to invocation

[tensorrt-cicd]

PR_Github #78017 [ run ] completed with state ABORTED. Commit: 3ad9117

Link to invocation

[tensorrt-cicd]

PR_Github #78030 [ run ] completed with state FAILURE. Commit: 3ad9117
/LLM/release-1.3/L0_MergeRequest_PR pipeline #243 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

[yuxianq]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #78334 [ run ] triggered by Bot. Commit: 5b0228a Link to invocation

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

• Configuration used: Repository: NVIDIA/TensorRT-LLM/.coderabbit.yaml
• Review profile: CHILL
• Plan: Enterprise
• Run ID: fa424e7c-03f2-4ae0-b9e6-e2cb67c794c7

📥 Commits

Reviewing files that changed from the base of the PR and between f8033a0 and 5b0228a.

📒 Files selected for processing: 6

• docs/source/developer-guide/sparse-attention-development-guide.md
• docs/source/features/sparse-attention.md
• tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md
• tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py
• tests/unittest/_torch/attention/test_prims_ts_attention_backend.py
• tests/unittest/_torch/attention/test_prims_ts_fmha.py

Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.

⚠️ A high-level summary could not be generated for this review. CodeRabbit will regenerate it on the next update, or you can request a refresh with @coderabbitai summary.

Walkthrough

The pull request clarifies block-sparse proxy-summary requirements and updates tests for partial KV blocks. It also adds coverage for context-attention lazy rescaling and uses autospecced mocks in decode and MLA wrapper tests.

Changes

Block-sparse proxy summaries

Layer / File(s)	Summary
Proxy summary contract and coverage docs/source/developer-guide/sparse-attention-development-guide.md, docs/source/features/sparse-attention.md, tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md, tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py	The documentation specifies mean summaries over present tokens and token-count weighting by the kernel. The reference calculation and tests cover partial final blocks; CPU-only markers are applied to individual tests. The GPU adapter test covers KV lengths 160 and 192.

Context attention lazy rescaling

Layer / File(s)	Summary
Late-key-block rescaling test tests/unittest/_torch/attention/test_prims_ts_attention_backend.py	A new test checks FP16 and BF16 results for log2 jumps of 7 and 9 against a softmax reference. It also checks CUDA graph replay.

FMHA wrapper mock specifications

Layer / File(s)	Summary
Autospecced decode and MLA wrapper tests tests/unittest/_torch/attention/test_prims_ts_fmha.py	Decode and MLA wrapper tests use autospecced mocks for wrapper reuse, capture-plan, and shared-workspace cases.

Priority: ⬇️ Low

Estimated code review effort: 2 (Simple) | ~12 minutes

Change: Other

Merge Risk: ⚪ Minimal · up to 5b022

No actionable merge-blocking issue is established by this review. Blackwell validation remains pending as reported.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 29 functions across 3 files. (3 skipped: 3…	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Title check		The title identifies this as a fix follow-up for the PrimTS refresh and accurately highlights the test and documentation updates.
Description check		The description follows the repository template. It explains the purpose, lists the affected review items, documents test coverage and limitations, and includes the checklist.
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

Full details: Docstring Coverage

Explanation

Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 29 functions across 3 files. (3 skipped: 3 unsupported.)

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

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

[yuxianq]

@PerkzZheng, could you help to review this PR as trt-llm-torch-attention-devs? Thanks~

Assignment explanation

Components	integration/attention-tests (90%); integration/documentation (10%)
Tags	algorithm:sparse (55%); algorithm:dense (45%); representation:mha (40%); representation:gqa (25%); representation:mla (35%); phase:prefill (65%); phase:decode (35%); features:cuda-graphs (50%); features:paged-kv (50%)
PR effort	4 points — Review requires checking proxy multiplicities and partial exact blocks, whether the constructed score jumps exercise lazy rescaling, and wrapper signature enforcement. This is bounded test/documentation work with more reasoning than a mechanical test update.
Affinity (0–1)	0.504 component/tag match
Outstanding points at selection	32 → 36 (+4)
Points already reserved on this PR	0
Effective weight	0.5
Completed review points (preceding 168 hours)	8
Combined points after assignment	44
Demand-derived target	40.00 points (soft share)
Additional fair share before assignment	0.00 points
Share explanation	Whole PR / qualification or coverage constraint exceeds a soft share; assignment is allowed.
Selection rule	balance-first maintenance

Affinity and workload are snapshots at selection, not live counters. Existing reservations on this PR are counted once.

[tensorrt-cicd]

PR_Github #78334 [ run ] completed with state SUCCESS. Commit: 5b0228a
/LLM/release-1.3/L0_MergeRequest_PR pipeline #285 completed with status: 'SUCCESS'

CI Report

Link to invocation

[nv-guomingz]

LGTM for doc part

[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.

[farazkh80]

LGTM, thanks!

[yuxianq]

@farazkh80, thanks for picking up this PR! I’ve handed off the trt-llm-torch-attention-devs review from @PerkzZheng to you.

The 4 attention-review points now belong to @farazkh80; @PerkzZheng’s attention reservation has been released. Your review is already submitted, so the points count as completed work; no new review request was added.

@PerkzZheng, no further attention review is requested from you. Thanks~