Skip to content

[feat] Enable MLA chunked prefill and KV cache reuse on SM121 #15347

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
CodersAcademy006 wants to merge 2 commits into from
+4 −4
Closed

[feat] Enable MLA chunked prefill and KV cache reuse on SM121 #15347

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
156e45b
feat: Enable SM121 support for MLA chunked prefill and KV block reuse
CodersAcademy006 Jun 14, 2026
2c586ac
Merge branch 'main' into feat/sm121-mla
CodersAcademy006 Jun 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -673,10 +673,10 @@ def drafting_loop_wrapper(model):

sm_version = get_sm_version()
if kv_cache_config.enable_block_reuse and sm_version not in [
90, 100, 103, 120
90, 100, 103, 120, 121
]:
logger.warning(
f"KV cache reuse for MLA can only be enabled on SM90/SM100/SM103/SM120, "
f"KV cache reuse for MLA can only be enabled on SM90/SM100/SM103/SM120/SM121, "
f"disable enable_block_reuse for SM{sm_version}")
kv_cache_config.enable_block_reuse = False
_set_model_engines_cache_reuse([model_engine, draft_model_engine],
Expand All @@ -693,9 +693,9 @@ def drafting_loop_wrapper(model):
kv_cache_config.enable_block_reuse = False
_set_model_engines_cache_reuse([model_engine, draft_model_engine],
False)
if enable_chunked_context and sm_version not in [90, 100, 103, 120]:
if enable_chunked_context and sm_version not in [90, 100, 103, 120, 121]:
logger.warning(
"Chunked Prefill for MLA can only be enabled on SM90/SM100/SM103/SM120, "
"Chunked Prefill for MLA can only be enabled on SM90/SM100/SM103/SM120/SM121, "
f"disable enable_chunked_context for SM{sm_version}")
enable_chunked_context = False
model_engine.attn_runtime_features.chunked_prefill = False
Expand Down