[None][feat] Enable block reuse for flashinfer

[2ez4bz]

Dev Engineer Review

• Regular FLASHINFER now allows block reuse.
• FLASHINFER_STAR_ATTENTION still disables block reuse.
• Unit tests cover both backend behaviors.
• KV-cache statistics are enabled for test_chunked_prefill.
• The reusedBlocks > 0.5 assertion can affect TRTLLM pre-merge tests.
• The reported changes do not confirm removal or explanation of the unrelated Gemma4 configuration change.
• Both L0 pipeline runs failed. Investigate and fix the failures before rerunning CI.
• No test-list changes are reported.

QA Engineer Review

Test functions added or modified:

• _latest_kv_cache_stats
• test_chunked_prefill
• test_flashinfer_preserves_cache_reuse
• test_flashinfer_star_attention_disables_cache_reuse
• test_llm_prefix_cache_reuse

Coverage includes:

• FlashInfer block reuse.
• FlashInfer Star Attention block-reuse disabling.
• TRTLLM and FlashInfer prefix-cache reuse.
• Chunked-prefill cache reuse.

The summary does not show corresponding test-db/ or qa/ entries. FlashInfer chunked-prefill MMLU coverage remains post-merge. Existing pre-merge prefix-cache coverage does not include partial non-page-aligned reuse, chunked prefill, or CUDA graphs.

Verdict: needs follow-up.

Description

Removes outdated guard for block reuse with the FLASHINFER
attention backend.

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]

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: 8fc3f5e2-d38a-4e20-bd76-1a16a3b0c1de

📥 Commits

Reviewing files that changed from the base of the PR and between 7485ba9 and 9640690.

📒 Files selected for processing (4)

• tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
• tests/integration/defs/accuracy/test_llm_api_pytorch.py
• tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py
• tests/unittest/llmapi/test_llm_pytorch.py

🚧 Files skipped from review as they are similar to previous changes (4)

• tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
• tests/unittest/llmapi/test_llm_pytorch.py
• tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py
• tests/integration/defs/accuracy/test_llm_api_pytorch.py

---

Walkthrough

The executor now disables block reuse only for FlashInfer Star Attention. Tests cover FlashInfer and TRTLLM prefix-cache reuse, FlashInfer Star Attention metadata, and chunked-prefill cache statistics.

Changes

KV-cache reuse behavior

Layer / File(s)	Summary
Backend-specific reuse policy tensorrt_llm/_torch/pyexecutor/py_executor_creator.py, tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py	The executor preserves block reuse for FlashInfer and disables it for FlashInfer Star Attention. Unit tests verify backend selection, KV-cache configuration, and runtime metadata.
Chunked-prefill reuse validation tests/integration/defs/accuracy/test_llm_api_pytorch.py	The integration test reads cumulative KV-cache statistics and requires more than 0.5 reused blocks after MMLU evaluation.
End-to-end prefix-cache coverage tests/unittest/llmapi/test_llm_pytorch.py	The test validates cold and warm prefix-cache reuse for TRTLLM and FlashInfer and checks deterministic output. It also adds the SPDX and Apache-2.0 license headers.

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

Suggested reviewers: brnguyen2, cascade812, allisonlim-nv

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the feature change: enabling block reuse for the FLASHINFER attention backend.
Description check	✅ Passed	The description explains the main change and includes the checklist, but the Test Coverage section does not list the relevant tests.
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]

🧹 Nitpick comments (1)

tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py (1)

213-244: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Add annotations to the new and modified Python function interfaces.

The coding guidelines require annotations for every function.

• tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py#L213-L244: Annotate _run_create_py_executor parameters and return type.
• tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py#L392-L417: Annotate both added test function interfaces.
• tests/integration/defs/accuracy/test_llm_api_pytorch.py#L85-L96: Annotate _latest_kv_cache_stats parameter and return type.
• tests/unittest/llmapi/test_llm_pytorch.py#L255-L280: Annotate test_llm_prefix_cache_reuse parameter and return type.

As per coding guidelines, “Annotate every function, use None for non-returning functions.”

🤖 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
`@tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py`
around lines 213 - 244, Annotate every affected function interface: add
parameter and return annotations to _run_create_py_executor in
tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py:213-244,
both added test functions in
tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py:392-417,
_latest_kv_cache_stats in
tests/integration/defs/accuracy/test_llm_api_pytorch.py:85-96, and
test_llm_prefix_cache_reuse in
tests/unittest/llmapi/test_llm_pytorch.py:255-280; use None for functions that
do not return a value.

Source: Coding guidelines

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

Nitpick comments:
In
`@tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py`:
- Around line 213-244: Annotate every affected function interface: add parameter
and return annotations to _run_create_py_executor in
tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py:213-244,
both added test functions in
tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py:392-417,
_latest_kv_cache_stats in
tests/integration/defs/accuracy/test_llm_api_pytorch.py:85-96, and
test_llm_prefix_cache_reuse in
tests/unittest/llmapi/test_llm_pytorch.py:255-280; use None for functions that
do not return a value.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: a0b6452a-e518-421d-9e57-e5d4172ac7d9

📥 Commits

Reviewing files that changed from the base of the PR and between 8e602fa and 453c4ab.

📒 Files selected for processing (5)

• tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
• tests/integration/defs/accuracy/test_llm_api_pytorch.py
• tests/integration/defs/accuracy/test_llm_api_pytorch_multimodal.py
• tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py
• tests/unittest/llmapi/test_llm_pytorch.py

[BowenFu]

Reviewed 453c4abc8. I think the premise is right — the guard does look genuinely outdated — but I'm not approving yet, for the reasons below.

Why I agree the guard is stale. I checked the metadata builder rather than taking it on faith. flashinfer.py:1141-1143 derives kv_lens_host from num_cached_tokens_per_seq + seq_lens_kv and num_blocks from that, and :1159-1160 carries an explicit note — "do not use len(block_ids) — that will give you a number that can be too big if using chunked prefill/kv cache reuse". So the builder was written reuse-aware on purpose, and the append path derives new-token positions from seq_len - append_len, which handles a non-page-aligned cached prefix. py_executor_creator.py is also the only place in the tree that couples FLASHINFER to enable_block_reuse, and no doc claims the limitation. Your new star-attention comment matches the actual distinction. Worth noting for anyone reading: attn_backend defaults to TRTLLM, so this affects only users who opt into FlashInfer.

1. No L0 pipeline has run. All 11 checks on this head are GitHub Actions gates (title format, DCO, pre-commit, base freshness). For a PR whose entire content is a behavior flip on the KV-cache path, the pipeline result is the evidence — please /bot run.

2. Your strongest regression guard doesn't run pre-merge. TestLlama3_1_8BInstruct::test_chunked_prefill[...attn_backend=FLASHINFER] is scheduled only in tests/integration/test_lists/qa/llm_function_core.txt:579-582; there is no test-db/l0_*.yml entry for it. So the MMLU reusedBlocks > 0 assertion lands post-merge only.

What does run pre-merge is test_llm_prefix_cache_reuse — part3 is under stage: pre_merge in l0_a100.yml:19 and l0_h100.yml:268 — and that's the right instinct. But it's TinyLlama, one repeated prompt, 5 output tokens, cuda_graph_config=None, no chunked prefill. The combinations where reuse actually gets interesting on this backend — partial reuse yielding a non-page-aligned num_cached_tokens, reuse under chunked prefill, and reuse under CUDA graphs (note :1245 asserts mixed decode/prefill is impossible under graphs) — have no pre-merge coverage. Adding the FLASHINFER chunked-prefill case to a test-db list would make this self-evidencing.

3. The Gemma4 change is unrelated to this PR and it lands on a pre-merge gate. TestGemma4_26B_A4B builds LLM(...) with no attn_backend, so it runs on the default TRTLLM backend — the guard you're removing never applied to it. You added that enable_block_reuse=False, enable_partial_reuse=False yourself in #16108 (fd9166c0a7), and TestGemma4_26B_A4B::test_nvfp4 is scheduled pre-merge at l0_b200.yml:325 with enable_chunked_prefill=True and fp8 KV. If the original disable was working around something real, this re-enables it inside an accuracy gate; if it was just conservative, please say so in the description. Either way it's a separate concern from the FlashInfer guard and would read better as its own PR.

4. Minor. assert stats["reusedBlocks"] > 0 also now applies to the attn_backend=TRTLLM parametrization, so a change in MMLU's prompt structure or in eviction behavior can turn an accuracy test red for a reason unrelated to accuracy. Fine if that's intended, just noting the new failure mode.

[BowenFu]

Correction to my point 2 above — I got the evidence wrong, and the conclusion needs restating more precisely.

TestLlama3_1_8BInstruct::test_chunked_prefill[...attn_backend=FLASHINFER] is in a test-db list: l0_h100.yml:377-378. My grep missed it. But that entry sits in the condition block starting at l0_h100.yml:305, whose terms.stage is post_merge — whereas the TRTLLM variants at l0_h100.yml:139-140 are in the pre_merge block starting at :3.

So, concretely, for the assertion you added:

• attn_backend=TRTLLM → runs pre-merge (l0_h100.yml:139-140).
• attn_backend=FLASHINFER → runs post-merge (l0_h100.yml:377-378) and in the QA list.

The ask is therefore not "wire it into test-db" — ignore that part. It's that the FlashInfer half of your MMLU guard, which is the half that covers the behavior this PR changes, won't report in this PR's own pipeline. Pre-merge coverage of FlashInfer+reuse comes from test_llm_prefix_cache_reuse instead (part3, l0_a100.yml:19 / l0_h100.yml:268), which is real but thin as described above. Promoting the two FLASHINFER test_chunked_prefill lines into the pre-merge block would close that, if the ~90s timeout budget allows.

One consequence worth flagging that I understated: because the TRTLLM lanes are pre-merge, enable_iter_perf_stats=True and assert stats["reusedBlocks"] > 0 land in pre-merge CI immediately. That's the new failure mode from my point 4, and it arrives on the parametrization this PR isn't about.

Everything else in my earlier comment stands — in particular the Gemma4 config flip (point 3), which is the one I'd most like addressed before this merges.

[BowenFu]

Approving — rechecked at e415c8e1.

Gemma4 is gone; that was my blocker. Both sides of the flip are now pinned pre-merge: test_flashinfer_preserves_cache_reuse / ..._star_attention_disables_cache_reuse via unittest/_torch/executor (l0_h100.yml:22), and test_llm_prefix_cache_reuse[TRTLLM|FLASHINFER] via part3 (l0_a100.yml:19, l0_h100.yml:269).

Also re-checked that py_executor_creator.py is still the only place coupling FLASHINFER to enable_block_reuse — _util.py:2839 only asserts the star-attention backend name — so narrowing the guard leaves no stale assumption behind it.

Post-merge is fine for the FLASHINFER test_chunked_prefill case, and no need to rework the test changes.

One to note, not act on: assert reusedBlocks > 0 and enable_iter_perf_stats=True also hit the TRTLLM parametrization, which runs pre-merge — an eviction or prompt-structure change can redden an accuracy gate for a non-accuracy reason.

[2ez4bz]

/bot run

[tensorrt-cicd]

PR_Github #63502 [ run ] triggered by Bot. Commit: 7485ba9 Link to invocation

[tensorrt-cicd]

PR_Github #63502 [ run ] completed with state SUCCESS. Commit: 7485ba9
/LLM/main/L0_MergeRequest_PR pipeline #51471 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

[2ez4bz]

/bot run

[tensorrt-cicd]

PR_Github #63554 [ run ] triggered by Bot. Commit: 7485ba9 Link to invocation

[tensorrt-cicd]

PR_Github #63554 [ run ] completed with state FAILURE. Commit: 7485ba9
/LLM/main/L0_MergeRequest_PR pipeline #51520 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

[2ez4bz]

/bot run

[tensorrt-cicd]

PR_Github #63575 [ run ] triggered by Bot. Commit: 9640690 Link to invocation

[tensorrt-cicd]

PR_Github #63575 [ run ] completed with state FAILURE. Commit: 9640690
/LLM/main/L0_MergeRequest_PR pipeline #51539 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

[2ez4bz]

/bot run

[tensorrt-cicd]

PR_Github #63658 [ run ] triggered by Bot. Commit: 9640690 Link to invocation

[tensorrt-cicd]

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

CI Report

Link to invocation