[yuxianq]

Dev Engineer Review

• The inspected working-tree diff changes only import formatting in mla_decode_fp8.py. The page-table bounds checks described in the PR objectives are present in the file, but they are not part of the inspected diff.
• No current review findings or test results were supplied.

QA Engineer Review

• No test files appear in the inspected working-tree diff. The test file described in the supplied change summary exists in the repository, but is not modified in this diff.
• The available evidence does not establish test-list changes or test execution results. Coverage verdict: needs follow-up.

Per-File QA Perspective

• tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py: The diff changes import formatting only. The file contains bounds checks that default out-of-range page indices to page 0, but those checks are not changes in the inspected diff.

Description

Fixes https://nvbugs/6862662: test_attention_mla (V1 KV cache, BF16 activations, FP8 KV cache, 64 requests × 10 context tokens, page size 32) intermittently hits a CUDA illegal memory access on B200.

Root cause. In the CuTe DSL FP8 MLA decode kernel (mla_decode_fp8.py), the K and V TMA load paths read page_per_tile page IDs for every KV tile. A 128-token tile can need up to four page IDs at page size 32, even when a request owns only one page. The kernel read those entries without checking the page-table width. For the last request, those reads run past the end of the page-table allocation. Here that allocation is 64 requests × 2 interleaved K/V slots × 1 int32 page ID = 512 bytes. Whether the access faults depends on what the caching allocator has placed after the table, so the failure is order-dependent. The original traceback surfaced later, in the test's reference computation, because CUDA reports the fault asynchronously.

Fix. Bound each page-table read by mPT.shape[0]. Lanes past the table width use physical page 0. Those lanes are masked out by the sequence length anyway, and still issuing the TMA load keeps each tile's expected byte count unchanged. This mirrors the existing guard in the FP16 kernel's page-table loader (mla_decode_fp16.py, load_page_table), which zero-fills out-of-range entries.

Test Coverage

• New tests/unittest/_torch/attention/test_cute_dsl_mla_page_table.py (8 variants): page sizes 32/64, persistent/non-persistent, split-KV 1/2, with one compiled kernel reused across table widths 4, 1, 2, 3, 5. Uses V1-style interleaved K/V page-table strides. Run under compute-sanitizer --tool memcheck with PYTORCH_NO_CUDA_MEMORY_CACHING=1, it detects the out-of-bounds reads before the fix. It is collected by the existing unittest/_torch/attention entry in l0_b200.yml.
• Existing tests/unittest/_torch/attention/test_attention_mla.py.

Local results on B200, release/1.3 at 7c2f9389a7:

Check	Before fix	After fix
Reported case, run after its V2-cache variant, under memcheck	64 invalid global reads in FP8 MLA split_kv_kernel	2 passed, 0 errors
Reported case ×10	—	10 passed
test_attention_mla.py (full file)	152 passed (fault is intermittent without sanitizer)	152 passed
New regression under memcheck	64 invalid reads (pre-fix, historical revision)	8 passed, 0 errors
Full unittest/_torch/attention suite at the historical failing revision	illegal access after 4,412 passed	4,786 passed

The same out-of-bounds reads reproduce on current main (c9b18eefa7); a separate PR will carry the fix and waiver removal there.

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.

🤖 Generated with Claude Code

[yuxianq]

/bot run --disable-fail-fast

[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: 64c7f333-a9a1-43f2-81ed-d79da2d1ce6c

📥 Commits

Reviewing files that changed from the base of the PR and between 42ca6ba and 6fd4a54.

📒 Files selected for processing: 2

• tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py
• tests/unittest/_torch/attention/test_cute_dsl_mla_page_table.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.

Walkthrough

The MLA decode kernel now bounds Q/K and V page-table lookups and uses page 0 when an index is out of range. A new CUDA test covers partial page-table lengths across page sizes, persistence modes, and split_kv settings.

Changes

MLA page-table bounds

Layer / File(s)	Summary
Guard Q/K and V page lookups tensorrt_llm/_torch/cute_dsl_kernels/blackwell/attention/mla/mla_decode_fp8.py	Q/K and V page-table reads now occur only when the computed index is in range. Each lookup defaults to page 0.
Test partial page tables tests/unittest/_torch/attention/test_cute_dsl_mla_page_table.py	A parameterized CUDA test runs decode with partial page tables and compares the output with the expected cache mean.

Priority: ⬆️ High

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

Change: Bug fix

Suggested reviewers: reasonsolo

Merge Risk: ⚪ Minimal · up to 6fd4a

No merge-blocking issue is established; the partial-page-table regression test is included in the B200 test list.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

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

Passed checks: 4 passed

Check name	Status	Explanation
Title check		The title clearly identifies the NVBugs issue, marks the change as a fix, and summarizes the main change: bounding FP8 MLA page-table loads for partial tiles.
Description check		The description is complete and relevant. It explains the root cause and fix, documents detailed test coverage and results, and includes the required checklist sections.
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.

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

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

Assignment explanation

Components	kernels/cute-dsl-mla (75%); integration/attention-tests (25%)
Tags	algorithm:dense (100%); representation:mla (100%); phase:decode (100%); features:quantization (50%); features:paged-kv (50%)
PR effort	4 points — The implementation change is localized, but review must validate distinct QK/V indexing, page-0 fallback safety, and preservation of TMA completion behavior. Regression review also requires understanding table strides and compiled-kernel reuse across changing widths.
Importance	Important — The guards address GPU memory safety in both FP8 MLA page-table load paths. The page-0 fallback requires correctness reasoning about masked lanes and continued TMA loads.
Required reviewer coverage	NVIDIA/trt-llm-kernels-devs
Point allocation	4 points per assigned reviewer
Affinity (0–1)	0.745 component/tag match; 0.559 after PIC preference
Outstanding points at selection	49 → 49 (+0)
Points already reserved on this PR	4
Effective weight	0.5
Completed review points (preceding 168 hours)	49
Combined points after assignment	98
Demand-derived target	43.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	required existing reviewers for other-group coverage

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

[tensorrt-cicd]

PR_Github #78117 [ run ] triggered by Bot. Commit: 6fd4a54 Link to invocation

[tensorrt-cicd]

PR_Github #78117 [ run ] completed with state SUCCESS. Commit: 6fd4a54
/LLM/release-1.3/L0_MergeRequest_PR pipeline #253 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.