[#15344][fix] Add SM121 to MLA allowlists

[peter941221]

Description

py_executor_creator.py treats SM121 as unsupported for both MLA KV cache reuse and MLA chunked prefill. That fallback disables both features.

This patch adds SM121 to both MLA allowlists so SM121 follows the same MLA path as SM120.

The unit coverage in tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py now checks both SM121 gates:

• test_mla_sm121_supported_configuration_preserves_cache_reuse
• test_mla_sm121_supported_configuration_preserves_chunked_prefill

Test Coverage

• python -m py_compile tensorrt_llm/_torch/pyexecutor/py_executor_creator.py tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py
• Added mock-based regression coverage for both SM121 allowlist checks in tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py

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.

Summary by CodeRabbit

• New Features

  • Added support for SM121 GPU architecture for KV cache block reuse and chunked prefill features, enabling optimized performance on newer NVIDIA hardware.

• Tests

  • Enhanced test coverage to verify KV cache reuse and chunked prefill behavior on SM121, ensuring proper feature support across hardware configurations.

[peter941221]

/bot run

[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: e975424c-60be-45aa-8d56-71ece7399628

📥 Commits

Reviewing files that changed from the base of the PR and between 18cb08e and cfb27eb.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
• tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py

---

📝 Walkthrough

Walkthrough

SM121 is added to the allowed SM version sets in two MLA compatibility guards inside py_executor_creator.py: one for KV cache block reuse and one for chunked prefill. The test helper _run_create_py_executor is extended with an enable_chunked_prefill parameter and now returns a 3-tuple including the runtime chunked prefill flag. Two SM121 regression tests are added.

Changes

SM121 MLA Feature Support

Layer / File(s)	Summary
SM121 added to MLA KV cache reuse and chunked prefill guards tensorrt_llm/_torch/pyexecutor/py_executor_creator.py	SM121 is appended to the allowed SM version tuple in both the KV cache reuse guard (lines 676–679) and the chunked prefill guard (lines 696–698), and the warning messages in both paths are updated to list the new supported set.
Test helper and SM121 regression tests tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py	_run_create_py_executor gains an enable_chunked_prefill kwarg, propagates it to llm_args, and returns a 3-tuple including attn_runtime_features.chunked_prefill. Existing fallback tests are updated to unpack the 3-tuple. Two new SM121-specific regression tests assert that KV cache reuse and runtime chunked prefill remain enabled on SM121.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Suggested reviewers

• QiJune
• mikeiovine

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and specifically describes the main change: adding SM121 to MLA allowlists. It directly aligns with the primary purpose of the PR.
Description check	✅ Passed	The description covers all essential sections from the template: it explains what was changed (adding SM121 to allowlists), why (SM121 was incorrectly unsupported), test coverage details, and confirms checklist completion.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
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.}

[VALLIS-NERIA]

/bot run

[tensorrt-cicd]

PR_Github #54790 [ run ] triggered by Bot. Commit: cfb27eb Link to invocation

[tensorrt-cicd]

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

Link to invocation

[peter941221]

Fixed in c026457.

The failure on cfb27eb was the pre-commit gate, not a behavior regression:

• yapf rewrote tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
• yapf and ruff-format rewrote tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py

I pushed the formatter output only. No semantic change.

An NVIDIA team member will need to rerun CI.