From 4ffc76b96d5334cfc6a551428d786fde7c96c407 Mon Sep 17 00:00:00 2001 From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Date: Sun, 23 Aug 2026 01:19:58 -0700 Subject: [PATCH 1/2] [https://nvbugs/6617948][fix] Restore trtllm-gen MLA decode perf gate dropped by #15300 PR #15300 ("Use public flashinfer APIs") deleted a pure performance gate from FlashInferTrtllmGenFmha as collateral of an otherwise mechanical refactor: the SLOWER_MLA_GENERATION_KERNELS constant, its membership check in _check_mla_generation_support(), and that method's tokens_per_block parameter. With the gate present, this backend declined MLA *generation* for (headDimQk=576, headDimV=512, tokens_per_block=32) and selection fell through to FallbackFmha (thop.attention). With it deleted, flashinfer's trtllm_batch_decode_with_kv_cache_mla claims the workload, so the build runs the decode kernel the deleted set itself labelled slower. That tuple is the DeepSeek-V3/R1-family and Kimi-K2/K2.5 MLA shape at the default page size, worth ~3% output token throughput on a decode-dominated case. This is not a revert of #15300 -- the public-API migration is kept in full. The restored gate is scoped to mla_backend == "trtllm-gen". The measurement behind it is of that kernel, and since the supported head dims are only (320,256) and (576,512), an unconditional gate would decline the whole MLA generation path at tokens_per_block=32 and take the cute-dsl MLA decode backend down with it. Adds four CPU-only unit tests. They call the checker rather than asserting on the constant, so a dropped parameter becomes a TypeError and a dropped constant an AttributeError -- nothing in the tree referenced either before this change, which is why the deletion was invisible to CI. Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> --- .../fmha/flashinfer_trtllm_gen.py | 23 +++++++ .../_torch/attention/test_fmha_page_index.py | 69 +++++++++++++++++++ 2 files changed, 92 insertions(+) diff --git a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py index 14840b595565..2f21e8f87950 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py @@ -507,6 +507,11 @@ class FlashInferTrtllmGenFmha(PhasedFmha): (320, 256), (576, 512), } + # (headDimQk, headDimV, tokens_per_block) whose trtllm-gen MLA decode kernel is + # slower than the thop.attention fallback, so this backend declines them. + SLOWER_MLA_GENERATION_KERNELS = { + (576, 512, 32), + } def __init__(self, attn: "TrtllmAttention") -> None: super().__init__(attn) @@ -633,6 +638,8 @@ def _get_attention_chunk_size(attn: "TrtllmAttention") -> int: def _check_mla_generation_support( cls, head_size: int, + tokens_per_block: int, + mla_backend: str, kv_lora_rank: Optional[int], qk_rope_head_dim: Optional[int], ) -> Tuple[bool, str]: @@ -670,6 +677,20 @@ def _check_mla_generation_support( f"headDimQk={head_dim_qk}, headDimV={head_dim_v}. Supported: {supported}.", ) + # Scoped to trtllm-gen: the measurement behind SLOWER_MLA_GENERATION_KERNELS + # is of that kernel, and declining here would take the cute-dsl MLA decode + # path down with it. + if ( + mla_backend == "trtllm-gen" + and (head_dim_qk, head_dim_v, tokens_per_block) in cls.SLOWER_MLA_GENERATION_KERNELS + ): + return ( + False, + f"[Generation][MLA] slower TRTLLM-GEN decode kernel for " + f"headDimQk={head_dim_qk}, headDimV={head_dim_v}, " + f"tokens_per_block={tokens_per_block}.", + ) + return True, "" def is_supported( @@ -836,6 +857,8 @@ def _is_supported_with_reason( if is_mla_enable: supported, reason = self._check_mla_generation_support( head_size=attn.head_dim, + tokens_per_block=tokens_per_block, + mla_backend=self._mla_backend, kv_lora_rank=attn.kv_lora_rank, qk_rope_head_dim=attn.qk_rope_head_dim, ) diff --git a/tests/unittest/_torch/attention/test_fmha_page_index.py b/tests/unittest/_torch/attention/test_fmha_page_index.py index 3f18a0b30e84..d06e9ceb01c0 100644 --- a/tests/unittest/_torch/attention/test_fmha_page_index.py +++ b/tests/unittest/_torch/attention/test_fmha_page_index.py @@ -249,3 +249,72 @@ def policy( assert fmha._get_effective_mla_backend(meta, 4) == "trtllm-gen" assert calls == [("cute-dsl", meta, 4)] + + +# The four tests below guard the MLA generation perf gate that #15300 removed as +# refactoring collateral, costing ~3% output token throughput on DeepSeek-V3-family +# and Kimi-K2 MLA decode at the default tokens_per_block. They deliberately call the +# checker instead of asserting on SLOWER_MLA_GENERATION_KERNELS itself: a test that +# pins the literal set would be deleted along with the constant by the next +# mechanical refactor, whereas these turn a dropped parameter into a TypeError and a +# dropped constant into an AttributeError. + + +def test_mla_generation_declines_slower_trtllm_gen_decode_kernel() -> None: + # DeepSeek-V3 / Kimi-K2 shape at the default tokens_per_block=32: the trtllm-gen + # MLA decode kernel is slower here than the thop.attention fallback, so this + # backend must decline and let selection fall through. + supported, reason = FlashInferTrtllmGenFmha._check_mla_generation_support( + head_size=576, + tokens_per_block=32, + mla_backend="trtllm-gen", + kv_lora_rank=512, + qk_rope_head_dim=64, + ) + assert not supported + assert "slower" in reason + assert "headDimQk=576" in reason + assert "headDimV=512" in reason + assert "tokens_per_block=32" in reason + + +@pytest.mark.parametrize("tokens_per_block", [16, 64]) +def test_mla_generation_gate_is_scoped_to_one_page_size(tokens_per_block: int) -> None: + # The gate must stay narrow: the same head dims at other page sizes are still + # served by this backend. Real configs run tokens_per_block=64. + supported, reason = FlashInferTrtllmGenFmha._check_mla_generation_support( + head_size=576, + tokens_per_block=tokens_per_block, + mla_backend="trtllm-gen", + kv_lora_rank=512, + qk_rope_head_dim=64, + ) + assert supported, reason + assert reason == "" + + +def test_mla_generation_gate_is_scoped_to_the_trtllm_gen_backend() -> None: + # The gated kernel is the trtllm-gen one; the cute-dsl MLA decode path shares + # this class and these head dims, and must stay selectable. + supported, reason = FlashInferTrtllmGenFmha._check_mla_generation_support( + head_size=576, + tokens_per_block=32, + mla_backend="cute-dsl", + kv_lora_rank=512, + qk_rope_head_dim=64, + ) + assert supported, reason + assert reason == "" + + +def test_mla_generation_allows_other_supported_head_dims() -> None: + # (320, 256) is unaffected at every page size. + supported, reason = FlashInferTrtllmGenFmha._check_mla_generation_support( + head_size=320, + tokens_per_block=32, + mla_backend="trtllm-gen", + kv_lora_rank=256, + qk_rope_head_dim=64, + ) + assert supported, reason + assert reason == "" From f001cb500e4fb562a193b6f003bc8504299a7104 Mon Sep 17 00:00:00 2001 From: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> Date: Sun, 23 Aug 2026 04:50:56 -0700 Subject: [PATCH 2/2] [https://nvbugs/6617948][fix] Gate on the effective MLA decode backend Address review feedback: the gate read the statically configured self._mla_backend, so a cute-dsl configuration whose per-batch policy downgrades to trtllm-gen (Kimi K3 does this for mixed batches and for speculative verification) still ran the slower kernel at (576, 512, 32) with the gate silent. is_supported() is evaluated per forward, per batch (trtllm.py), so reading the policy-resolved backend costs nothing and cannot disable a backend wholesale: declining only diverts that batch to FallbackFmha, which is the pre-#15300 selection this PR restores. MLA reaches the gate only as generation-only, so num_gen_tokens == q.size(0), matching prepare_workspace's is_gen_only branch. Two tests: the policy-downgrade composition, and a structural assertion that the call site passes the effective backend, since reverting it is a one-word change that no behavioural test here would catch. Signed-off-by: chenfeiz0326 <203214996+chenfeiz0326@users.noreply.github.com> --- .../fmha/flashinfer_trtllm_gen.py | 14 ++++-- .../_torch/attention/test_fmha_page_index.py | 50 ++++++++++++++++++- 2 files changed, 59 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py index 2f21e8f87950..52bf0a1a376c 100644 --- a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py +++ b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py @@ -677,9 +677,10 @@ def _check_mla_generation_support( f"headDimQk={head_dim_qk}, headDimV={head_dim_v}. Supported: {supported}.", ) - # Scoped to trtllm-gen: the measurement behind SLOWER_MLA_GENERATION_KERNELS - # is of that kernel, and declining here would take the cute-dsl MLA decode - # path down with it. + # Scoped to trtllm-gen: SLOWER_MLA_GENERATION_KERNELS was measured on that + # kernel. Callers pass the backend that will actually run this batch, so a + # cute-dsl batch stays selectable while a batch a policy downgraded to + # trtllm-gen is still gated. if ( mla_backend == "trtllm-gen" and (head_dim_qk, head_dim_v, tokens_per_block) in cls.SLOWER_MLA_GENERATION_KERNELS @@ -855,10 +856,15 @@ def _is_supported_with_reason( f"Q={q_dtype}, KV={kv_cache_dtype}, O={o_dtype}." ) if is_mla_enable: + # The effective backend, not the configured one: a policy may downgrade + # cute-dsl to trtllm-gen for this batch, and it is the kernel that + # actually runs that the gate is about. MLA reaches here only as + # generation-only (checked above), so num_gen_tokens == q.size(0), + # matching prepare_workspace's is_gen_only branch. supported, reason = self._check_mla_generation_support( head_size=attn.head_dim, tokens_per_block=tokens_per_block, - mla_backend=self._mla_backend, + mla_backend=self._get_effective_mla_backend(meta, q.size(0)), kv_lora_rank=attn.kv_lora_rank, qk_rope_head_dim=attn.qk_rope_head_dim, ) diff --git a/tests/unittest/_torch/attention/test_fmha_page_index.py b/tests/unittest/_torch/attention/test_fmha_page_index.py index d06e9ceb01c0..b29a78d3c715 100644 --- a/tests/unittest/_torch/attention/test_fmha_page_index.py +++ b/tests/unittest/_torch/attention/test_fmha_page_index.py @@ -1,6 +1,9 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import ast +import inspect +import textwrap from collections.abc import Callable from types import SimpleNamespace from typing import TypeAlias @@ -251,7 +254,7 @@ def policy( assert calls == [("cute-dsl", meta, 4)] -# The four tests below guard the MLA generation perf gate that #15300 removed as +# The six tests below guard the MLA generation perf gate that #15300 removed as # refactoring collateral, costing ~3% output token throughput on DeepSeek-V3-family # and Kimi-K2 MLA decode at the default tokens_per_block. They deliberately call the # checker instead of asserting on SLOWER_MLA_GENERATION_KERNELS itself: a test that @@ -307,6 +310,51 @@ def test_mla_generation_gate_is_scoped_to_the_trtllm_gen_backend() -> None: assert reason == "" +def test_mla_generation_gate_declines_a_policy_downgrade_to_trtllm_gen() -> None: + # A cute-dsl config whose per-batch policy downgrades to trtllm-gen (K3 does so + # for mixed batches and for speculative verification) still runs the gated + # kernel, so the gate must fire on the *effective* backend. Reading the static + # self._mla_backend here would let the slower kernel through. + fmha = _make_fmha("cute-dsl", mla_backend_policy=lambda *_: "trtllm-gen") + meta = SimpleNamespace(num_contexts=1, num_generations=3) + effective = fmha._get_effective_mla_backend(meta, 3) + assert effective == "trtllm-gen" + + supported, reason = FlashInferTrtllmGenFmha._check_mla_generation_support( + head_size=576, + tokens_per_block=32, + mla_backend=effective, + kv_lora_rank=512, + qk_rope_head_dim=64, + ) + assert not supported + assert "slower" in reason + + +def test_mla_generation_gate_reads_the_effective_mla_backend() -> None: + # The composition above is only load-bearing if the production call site feeds + # the gate the effective backend. Assert that structurally: reverting to the + # static self._mla_backend is a one-word change, and no behavioural test here + # would catch it because driving _is_supported_with_reason needs a full + # metadata/forward-args stub. + source = textwrap.dedent(inspect.getsource(FlashInferTrtllmGenFmha._is_supported_with_reason)) + calls = [ + node + for node in ast.walk(ast.parse(source)) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "_check_mla_generation_support" + ] + assert len(calls) == 1, "expected exactly one MLA generation gate call site" + kwargs = {kw.arg: kw.value for kw in calls[0].keywords} + passed = kwargs.get("mla_backend") + assert passed is not None, "gate call site lost its mla_backend argument" + assert ( + isinstance(passed, ast.Call) + and getattr(passed.func, "attr", None) == "_get_effective_mla_backend" + ), f"gate must receive the effective backend, got {ast.dump(passed)}" + + def test_mla_generation_allows_other_supported_head_dims() -> None: # (320, 256) is unaffected at every page size. supported, reason = FlashInferTrtllmGenFmha._check_mla_generation_support(