[yuxianq]

Dev Engineer Review

• The vendor pin now uses FlashInfer commit bf444275170ed1134644c5f682d179e72bb44a58. Patch and source digests were regenerated.
• The compatibility patch removes explicit trace-template dispatch objects from PrimTS and decode API decorators. It uses bare @flashinfer_api decorators instead.
• The patch adds 1-CTA MLA fallback support through the 2-CTA kernel, including short-KV bounds and optional CTA metadata.
• The patch clarifies decode sequence-length ownership and fixes contiguous-context MMA task dependencies.
• The scope excludes unrelated upstream changes, the canonical trtllm-prims-ts-dev branch, and FP8 MLA integration.
• Main regression risks are MLA kernel selection, short-KV bounds, CUDA graph replay behavior, and compatibility with callers that depended on the removed trace-template imports.

QA Engineer Review

• tests/unittest/_torch/attention/test_prims_ts_attention_backend.py was modified.
• test_prims_ts_mla_live_wrapper_cuda_graph_replay() now runs with maximum sequence lengths of 64 and 1024. The 64-token case verifies selection of the throughput_2cta kernel.
• The test covers CUDA graph replay and short-KV MLA fallback behavior. Coverage is sufficient for the stated test change.
• No changed integration test requires test-list verification from the provided changes.

Per-File QA Perspective

• 3rdparty/vendor_patches/flashinfer-prims-ts.patch: Verify 1-CTA MLA fallback, short-KV bounds, CTA metadata, decode sequence-length handling, and contiguous-context MMA dependencies. Verify that removing trace-template imports does not break API consumers.
• 3rdparty/vendor_sources.lock.yaml: Verify that the vendor commit, patch digest, and source digest match the checked-in vendor contents.
• tests/unittest/_torch/attention/test_prims_ts_attention_backend.py: Verifies MLA CUDA graph replay at 64 and 1024 maximum sequence lengths, including throughput_2cta selection for the short configuration. No test-list change is included.

Description

Backport the remaining PrimTS follow-ups from merged FlashInfer #4829 without importing unrelated changes from newer FlashInfer main.

• Update the vendor pin from main's CUTLASS-4.8-compatible bad2bdb15aac3553934e7a7a164dcb2ca4fda7f1 to b0d190d4f2420f0c2a0ab3a6efeef4439c2aa996, on the temporary branch trtllm-prims-ts-dev-4829-backport.
• Fall back from automatic 1-CTA MLA selection to the supported 2-CTA kernel when the 1-CTA configuration is unavailable, including short KV bounds.
• Make decode sequence-length ownership explicit: lengths supplied to plan() are immutable owned snapshots; otherwise each run() supplies its lengths. TRT-LLM already uses the latter mode, so its adapter is unchanged.
• Include optional caller-guaranteed zero V tails and caller-precomputed variable-window CTA start bounds.
• Fix contiguous-context MMA task dependencies so nonpaged schedules do not consume paged global metadata.
• Regenerate the existing compatibility patch and vendor digests. The patch's actual additions and removals are unchanged; only its diff context and blob IDs change.

The FlashInfer backport is one consolidated commit based on the existing dev tip. Its PrimTS Python delta is patch-equivalent to 18de70f6..c5611f76, the missing follow-up range in #4829 (squash: 3985489542af7adba88b970e634cf7490d72ff63). Tests are adapted to the older baseline without importing LDTM.STAT or newer block-sparse features.

The canonical trtllm-prims-ts-dev branch is not advanced by this PR. After merge, the maintainer can fast-forward it to the pinned commit and open the lock-only follow-up. This PR does not enable FP8 MLA in the TRT-LLM adapter and is independent of #18946.

Rebase update — 2026-09-10

• FlashInfer: bf444275 → b0d190d4, rebased onto latest trtllm-prims-ts-dev (bad2bdb1). The backport patch is unchanged (git range-diff reports equality); the CUTLASS DSL 4.8 Task API migration is inherited intact.
• TRT-LLM: a8e8a746 → 86e66439c32747439abb0427f0d629db4ac373af, rebased onto main b509ae936cd64eec2b09776568293a7170a15729. Only 3rdparty/vendor_sources.lock.yaml conflicted. The official vendor tool regenerated the pin/digest; the nine-file backport patch is unchanged outside lock metadata.
• The FMHA-decode and MLA-decode Task files are byte-identical to main's CUTLASS 4.8 migration. The existing downstream compatibility patch is preserved.
• Validation after rebase: syntax, changed-file pre-commit checks in both repositories, and vendor offline/local/pushed-upstream checks passed.
• FlashInfer host run: 1078 passed, 304 skipped, 1 failed. The failure is test_default_is_causal_leaves_other_backends_unchanged: the non-PrimTS JIT path cannot find flashinfer/data/csrc/batch_decode_customize_config.jinja in this worktree. No product-code fix was applied for this environment limitation. These host checks used the existing CUTLASS 4.7 environment; CUTLASS 4.8 GPU validation is delegated to current-main CI.
• No GPU rerun or C++ rebuild was performed for this rebase. The GPU results below are historical, before this rebase. CI is requested with /bot run --disable-fail-fast for the new head.

Test Coverage (before rebase)

• Extend test_prims_ts_mla_live_wrapper_cuda_graph_replay to cover a 64-token KV bound, assert the 2-CTA fallback, and retain the existing 1024-token case.
• FlashInfer on B200: 23 targeted checks passed, including BF16/FP8 short-KV MLA eager/one-shot/graph parity, dynamic decode lengths during graph replay, plan-owned length snapshots, V-tail handling, variable-window graph replay, and contiguous schedule construction.
• FlashInfer host/contract suite on A40: 1066 passed, 304 skipped, 13 deselected; PrimTS benchmark-adapter checks: 17 passed.
• TRT-LLM FMHA unit tests: 78 passed. Vendor lifecycle tests: 25 passed.
• TRT-LLM on B200: 11 distinct backend/wrapper cases passed across the final runs, covering Qwen2 GQA context/generation/mixed batches, DeepSeek-V3-Lite MLA generation, V-tail correctness, decode/MLA graph replay, shared workspace reuse, and fallback selection. Both short-KV and regular-KV MLA graph cases passed.
• Vendor snapshot verification passed against both the local FlashInfer commit and the pushed GitHub source (scripts/vendor_sources.py check flashinfer-prims-ts --upstream --require-access).
• pre-commit run --all-files passed in both repositories (TRT-LLM's type hook used its lightweight mode).

Local validation limitation: the reused TRT-LLM C++ build does not expose main's new reuse_match_backoff KV-cache-manager-v2 argument. Three backend tests failed at that binding boundary, before their affected attention execution; V2-dependent cases require a current-main build in CI. No C++ rebuild or full model-weight inference was performed for this source-only update.

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: 104c4650-7eb2-405d-8b27-6ce85acd7b73

📥 Commits

Reviewing files that changed from the base of the PR and between a8e8a74 and 86e6643.

📒 Files selected for processing (1)

• 3rdparty/vendor_sources.lock.yaml

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

---

Walkthrough

The vendored PrimTS source pin was updated. PrimTS API decorators and CUTE DSL imports were aligned with the updated source. MLA CUDA-graph replay coverage now tests two maximum sequence lengths and validates kernel selection.

Changes

PrimTS API updates

Layer / File(s)	Summary
Vendored source revision 3rdparty/vendor_sources.lock.yaml	The flashinfer-prims-ts source now uses a new branch and commit with updated source digests.
API decorator and import alignment 3rdparty/vendor_patches/flashinfer-prims-ts.patch	Decode, MLA, and block-sparse APIs use bare @flashinfer_api decorators. Rubin validation uses the absolute flashinfer.cute_dsl.utils import path.
MLA replay validation tests/unittest/_torch/attention/test_prims_ts_attention_backend.py	The replay test runs with maximum sequence lengths 64 and 1024. The 64-token case checks for throughput_2cta.

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

Merge Risk: ⚪ Minimal · up to 86e66

The vendored PrimTS revision and MLA replay coverage are updated, with reported validation completing successfully. No actionable merge-blocking correctness or deployment risk remains.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	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 2 functions across 1 files. (1 skipped: 1 …	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
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.
Title check	✅ Passed	The title follows the required format and clearly identifies the fix: updating vendored PrimTS with FlashInfer #4829 follow-ups.
Description check	✅ Passed	The description includes complete Description, Test Coverage, and PR Checklist sections. It explains the scope, implementation details, validation results, known limitations, and rebase status.

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 2 functions across 1 files. (1 skipped: 1 unsupported.)

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[yuxianq]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #72641 [ run ] triggered by Bot. Commit: a8e8a74 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.

[tensorrt-cicd]

PR_Github #72641 [ run ] completed with state SUCCESS. Commit: a8e8a74
/LLM/main/L0_MergeRequest_PR pipeline #59636 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation

[yuxianq]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #72707 [ run ] triggered by Bot. Commit: 86e6643 Link to invocation

[tensorrt-cicd]

PR_Github #72707 [ run ] completed with state SUCCESS. Commit: 86e6643
/LLM/main/L0_MergeRequest_PR pipeline #59696 completed with status: 'SUCCESS'

CI Report

Link to invocation