From bb930324a17a33acdc477e1e8188fe603f89eff0 Mon Sep 17 00:00:00 2001
From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Date: Tue, 30 Jun 2026 07:02:41 +0000
Subject: [PATCH 01/13] [None][fix] enable TRTLLM MLA context backend tests

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
---
 .../_torch/attention/backend_capability.py    |  16 ---
 .../unittest/_torch/attention/backend_case.py | 124 ++++++++++++++++--
 2 files changed, 111 insertions(+), 29 deletions(-)

diff --git a/tests/unittest/_torch/attention/backend_capability.py b/tests/unittest/_torch/attention/backend_capability.py
index ab117f9ddebb..b046d784da2b 100644
--- a/tests/unittest/_torch/attention/backend_capability.py
+++ b/tests/unittest/_torch/attention/backend_capability.py
@@ -213,22 +213,6 @@ def unsupported_reason(backend: str, case) -> Optional[str]:
     ):
         return "TRTLLM Blackwell no-cache fallback is unstable for head_dim 80"
 
-    # TRTLLM MLA context is validated by test_attention_mla.py through the
-    # production MLA module path. This standalone backend harness feeds already
-    # up-projected random K/V tensors, which does not match the TRTLLM MLA
-    # context op contract and produces invalid output, while Vanilla/FlashInfer
-    # can still validate the up-projected context math and cache append.
-    if (
-        backend == "TRTLLM"
-        and getattr(case, "is_mla", False)
-        and getattr(case, "num_contexts", 0) == len(getattr(case, "seq_lens", ()))
-    ):
-        return (
-            "TRTLLM MLA context is covered by test_attention_mla.py via the "
-            "production MLA module path; standalone up-projected backend "
-            "inputs are validated with Vanilla/FlashInfer"
-        )
-
     # On Blackwell, TRTLLM-Gen is the supported MLA generation path. The current
     # kernel set explicitly lacks the DeepSeek-family decode shape at page_size
     # 32, and the generic fallback then tries to JIT an unsupported generated
diff --git a/tests/unittest/_torch/attention/backend_case.py b/tests/unittest/_torch/attention/backend_case.py
index f1a9583ab864..95368af8b98f 100644
--- a/tests/unittest/_torch/attention/backend_case.py
+++ b/tests/unittest/_torch/attention/backend_case.py
@@ -113,9 +113,9 @@ class BackendCase:
     # mla_rope_generation (feeding a pre-formed fused_q and explicit q_pe), so
     # all three backends run RoPE-free and stay aligned; TRTLLM additionally
     # pre-writes the new latent into the cache and Python-initializes the
-    # trtllm-gen scheduler buffers. MLA *context* (up-projected K/V) runs on
-    # Vanilla/FlashInfer; the TRTLLM context FMHA fuses RoPE and is validated in
-    # test_attention_mla.py.
+    # trtllm-gen scheduler buffers. MLA context uses coherent up-projected K/V
+    # and latent-cache inputs: TRTLLM fuses RoPE while Vanilla/FlashInfer receive
+    # the equivalent pre-rotated tensors.
     v_head_dim: Optional[int] = None
     q_lora_rank: Optional[int] = None
     kv_lora_rank: Optional[int] = None
@@ -562,18 +562,104 @@ def _create_metadata(AttentionCls, case, mgr):
         mgr.shutdown()
 
 
+def _mla_context_pos_embd_params(case: BackendCase) -> PositionalEmbeddingParams:
+    """Build the GPT-J-style RoPE configuration required by TRTLLM MLA context."""
+    if case.rope is None:
+        raise ValueError("TRTLLM MLA context requires RoPE parameters.")
+
+    rope_config = dict(case.rope)
+    rope_config.update(dim=case.qk_rope_head_dim, duplicate_data=True)
+    return PositionalEmbeddingParams(
+        type=PositionEmbeddingType.rope_gptj,
+        rope=_rope_params_from_dict(rope_config),
+        is_neox=False,
+    )
+
+
 def generate_mla_context_inputs(case: BackendCase, seed: int = 0) -> Dict:
-    """Random up-projected MLA context inputs (asymmetric K/V)."""
+    """Random production-layout MLA context inputs before RoPE."""
     gen = torch.Generator(device="cuda").manual_seed(seed)
     cdt = case.compute_dtype
     H, Hkv = case.num_heads, case.num_kv_heads
     qk_head = case.qk_nope_head_dim + case.qk_rope_head_dim
-    d_latent = case.kv_lora_rank + case.qk_rope_head_dim
+    compressed_kv = _randn(gen, cdt, case.nnz_q, case.kv_lora_rank)
+    k_pe = _randn(gen, cdt, case.nnz_q, case.qk_rope_head_dim)
+    packed_kv = _randn(
+        gen,
+        cdt,
+        case.nnz_q,
+        Hkv * (case.qk_nope_head_dim + case.v_head_dim),
+    )
+    k_nope, v = packed_kv.split([Hkv * case.qk_nope_head_dim, Hkv * case.v_head_dim], dim=-1)
+    k = torch.cat(
+        [
+            k_nope.view(-1, Hkv, case.qk_nope_head_dim),
+            k_pe.view(-1, 1, case.qk_rope_head_dim).expand(-1, Hkv, -1),
+        ],
+        dim=-1,
+    ).view(-1, Hkv * qk_head)
     return dict(
         q=_randn(gen, cdt, case.nnz_q, H * qk_head),
-        k=_randn(gen, cdt, case.nnz_q, Hkv * qk_head),
-        v=_randn(gen, cdt, case.nnz_q, Hkv * case.v_head_dim),
-        latent_cache=_randn(gen, cdt, case.nnz_q, d_latent),  # appended for later gen
+        k=k,
+        # Keep the split view: TRTLLM MLA context expects token stride to include
+        # the packed k_nope portion, and Vanilla/FlashInfer support that layout.
+        v=v,
+        compressed_kv=compressed_kv,
+        k_pe=k_pe,
+        latent_cache=torch.cat([compressed_kv, k_pe], dim=-1),
+    )
+
+
+def _prepare_mla_context_inputs(
+    case: BackendCase,
+    inputs: Dict,
+    pos_embd_params: PositionalEmbeddingParams,
+    *,
+    fuse_rope: bool,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Prepare backend inputs and the expected post-RoPE latent cache."""
+    position_ids = make_position_ids(case.seq_lens, case.num_cached_tokens)
+    rope_params = pos_embd_params.rope
+    assert rope_params is not None
+
+    rotated_k_pe = apply_rope(
+        inputs["k_pe"],
+        position_ids,
+        rope_params,
+        case.qk_rope_head_dim,
+        is_neox=pos_embd_params.is_neox,
+    )
+    expected_latent_cache = torch.cat([inputs["compressed_kv"], rotated_k_pe], dim=-1)
+
+    if fuse_rope:
+        return (
+            inputs["q"].clone(),
+            inputs["k"].clone(),
+            inputs["v"],
+            inputs["latent_cache"],
+            expected_latent_cache,
+        )
+
+    q = inputs["q"].clone().view(-1, case.num_heads, case.head_dim)
+    q_pe = q[..., case.qk_nope_head_dim :].reshape(
+        case.nnz_q, case.num_heads * case.qk_rope_head_dim
+    )
+    q[..., case.qk_nope_head_dim :] = apply_rope(
+        q_pe,
+        position_ids,
+        rope_params,
+        case.qk_rope_head_dim,
+        is_neox=pos_embd_params.is_neox,
+    ).view(case.nnz_q, case.num_heads, case.qk_rope_head_dim)
+
+    k = inputs["k"].clone().view(-1, case.num_kv_heads, case.head_dim)
+    k[..., case.qk_nope_head_dim :] = rotated_k_pe.view(case.nnz_q, 1, case.qk_rope_head_dim)
+    return (
+        q.view(case.nnz_q, -1),
+        k.view(case.nnz_q, -1),
+        inputs["v"],
+        expected_latent_cache,
+        expected_latent_cache,
     )
 
 
@@ -582,6 +668,7 @@ def _run_mla_context_backend(case, backend, inputs, *, kv_layout: str) -> torch.
     AttentionCls = get_attention_backend(backend)
     qk_head = case.qk_nope_head_dim + case.qk_rope_head_dim
     request_ids = list(range(case.num_seqs))
+    pos_embd_params = _mla_context_pos_embd_params(case)
     attn = create_attention(
         backend,
         layer_idx=0,
@@ -589,6 +676,7 @@ def _run_mla_context_backend(case, backend, inputs, *, kv_layout: str) -> torch.
         head_dim=qk_head,
         num_kv_heads=case.num_kv_heads,
         q_scaling=case.q_scaling,
+        pos_embd_params=pos_embd_params,
         is_mla_enable=True,
         q_lora_rank=case.q_lora_rank,
         kv_lora_rank=case.kv_lora_rank,
@@ -598,7 +686,15 @@ def _run_mla_context_backend(case, backend, inputs, *, kv_layout: str) -> torch.
     )
     mgr = _build_mla_kv_cache_manager(case, backend)
     mgr.add_dummy_requests(request_ids, case.token_nums)
-    expected_latents = _split_packed_tokens(inputs["latent_cache"], case.seq_lens)
+    fuse_rope = AttentionCls.support_fused_rope()
+    q, k, v, latent_cache, expected_latent_cache = _prepare_mla_context_inputs(
+        case,
+        inputs,
+        pos_embd_params,
+        fuse_rope=fuse_rope,
+    )
+    expected_latents = _split_packed_tokens(expected_latent_cache, case.seq_lens)
+    cache_atol, cache_rtol = _tolerances(case, case.compute_dtype) if fuse_rope else (0.0, 0.0)
     metadata = AttentionCls.Metadata(
         num_contexts=case.num_contexts,
         kv_cache_params=KVCacheParams(
@@ -615,12 +711,12 @@ def _run_mla_context_backend(case, backend, inputs, *, kv_layout: str) -> torch.
     metadata.prepare()
     try:
         out = attn.forward(
-            inputs["q"],
-            inputs["k"],
-            inputs["v"],
+            q,
+            k,
+            v,
             metadata,
             forward_args=AttentionForwardArgs(
-                latent_cache=inputs["latent_cache"],
+                latent_cache=latent_cache,
                 attention_input_type=AttentionInputType.context_only,
             ),
         )
@@ -635,6 +731,8 @@ def _run_mla_context_backend(case, backend, inputs, *, kv_layout: str) -> torch.
             expected_latents,
             kv_layout=metadata.kv_layout,
             cache_kind="mla",
+            atol=cache_atol,
+            rtol=cache_rtol,
         )
         return out[: case.nnz_q].contiguous()
     finally:

From ccf8516dc9a87aaa13a3970b12d1137647760c74 Mon Sep 17 00:00:00 2001
From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Date: Tue, 30 Jun 2026 07:43:26 +0000
Subject: [PATCH 02/13] [None][fix] enable TRTLLM Blackwell MLA generation

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
---
 .../fmha/flashinfer_trtllm_gen.py             | 13 +----------
 .../_torch/attention/backend_capability.py    | 22 -------------------
 2 files changed, 1 insertion(+), 34 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 795e1766edc1..b5fea40413d2 100644
--- a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py
+++ b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py
@@ -408,9 +408,6 @@ class FlashInferTrtllmGenFmha(PhasedFmha):
         (320, 256),
         (576, 512),
     }
-    MISSING_MLA_GENERATION_KERNELS = {
-        (576, 512, 32),
-    }
 
     def __init__(self, attn: "TrtllmAttention"):
         super().__init__(attn)
@@ -567,14 +564,6 @@ def _check_mla_generation_support(
                 f"headDimQk={head_dim_qk}, headDimV={head_dim_v}. Supported: {supported}.",
             )
 
-        if (head_dim_qk, head_dim_v, tokens_per_block) in cls.MISSING_MLA_GENERATION_KERNELS:
-            return (
-                False,
-                f"[Generation][MLA] Missing 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(
@@ -1192,7 +1181,7 @@ def run_mla_generation(
         batch_beam = params.num_requests * meta.beam_width
         if params.attention_input is None:
             raise RuntimeError("MLA generation requires attention_input.")
-        kv_cache, block_tables = thop.build_trtllm_gen_kv_cache_metadata(
+        kv_cache, block_tables, _kv_scale_pool = thop.build_trtllm_gen_kv_cache_metadata(
             meta.host_kv_cache_pool_pointers,  # host_kv_cache_pool_pointers
             meta.host_kv_cache_pool_mapping,  # host_kv_cache_pool_mapping
             meta.kv_cache_block_offsets,  # kv_cache_block_offsets
diff --git a/tests/unittest/_torch/attention/backend_capability.py b/tests/unittest/_torch/attention/backend_capability.py
index b046d784da2b..e45d4e2e9b78 100644
--- a/tests/unittest/_torch/attention/backend_capability.py
+++ b/tests/unittest/_torch/attention/backend_capability.py
@@ -213,28 +213,6 @@ def unsupported_reason(backend: str, case) -> Optional[str]:
     ):
         return "TRTLLM Blackwell no-cache fallback is unstable for head_dim 80"
 
-    # On Blackwell, TRTLLM-Gen is the supported MLA generation path. The current
-    # kernel set explicitly lacks the DeepSeek-family decode shape at page_size
-    # 32, and the generic fallback then tries to JIT an unsupported generated
-    # kernel. FlashInfer/Vanilla still validate these standalone MLA gen cases.
-    if (
-        backend == "TRTLLM"
-        and sm >= 100
-        and getattr(case, "is_mla", False)
-        and getattr(case, "num_contexts", 0) == 0
-    ):
-        head_dim_qk = getattr(case, "kv_lora_rank", 0) + getattr(case, "qk_rope_head_dim", 0)
-        head_dim_v = getattr(case, "kv_lora_rank", 0)
-        if (head_dim_qk, head_dim_v, getattr(case, "page_size", None)) == (
-            576,
-            512,
-            32,
-        ):
-            return (
-                "TRTLLM-Gen Blackwell MLA generation is missing the decode "
-                "kernel for headDimQk=576, headDimV=512, page_size=32"
-            )
-
     # TRTLLM's fp8 generation (XQA) path computes in fp8 and needs the model's
     # real KV-dequant + output scale state (fed by the Attention module's
     # projection layers). A bare standalone backend lacks it: supplying a unit

From 3e2a315f951501b4155989536e810846a63b307f Mon Sep 17 00:00:00 2001
From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Date: Tue, 30 Jun 2026 10:32:31 +0000
Subject: [PATCH 03/13] [None][fix] reduce TRTLLM attention backend exclusions

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
---
 .../_torch/attention_backend/trtllm.py        |  2 +-
 .../_torch/attention/backend_capability.py    | 70 +++++++------------
 .../unittest/_torch/attention/backend_case.py |  5 +-
 3 files changed, 28 insertions(+), 49 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py
index 420fd14a1fe0..3c1829774ac8 100644
--- a/tensorrt_llm/_torch/attention_backend/trtllm.py
+++ b/tensorrt_llm/_torch/attention_backend/trtllm.py
@@ -1495,7 +1495,7 @@ def forward(
 
         # SM90 forces `use_paged_context_fmha` on for correctness
         # (https://nvbugs/5624818).
-        if get_sm_version() == 90:
+        if get_sm_version() == 90 and metadata.num_contexts > 0:
             metadata.use_paged_context_fmha = True
 
         # Sparse mqa/gqa attention uses generation kernel which reads Q from qPtr (separate buffer).
diff --git a/tests/unittest/_torch/attention/backend_capability.py b/tests/unittest/_torch/attention/backend_capability.py
index e45d4e2e9b78..16995e4b36e8 100644
--- a/tests/unittest/_torch/attention/backend_capability.py
+++ b/tests/unittest/_torch/attention/backend_capability.py
@@ -177,21 +177,30 @@ def unsupported_reason(backend: str, case) -> Optional[str]:
     ):
         return f"TRTLLM Blackwell paged attention is unstable for head_dim {case.head_dim}"
 
-    # TRTLLM's Blackwell pure-decode sliding-window path is numerically unstable
-    # for the listed Gemma GQA shapes. Mixed batches still use a different path
-    # and match the golden.
+    # TRTLLM's Blackwell no-cache fallback mismatches the Vanilla golden for the
+    # Qwen2-VL vision tower's head_dim 80 workload; other no-cache head dims in
+    # this sweep still pass on Blackwell.
     if (
         backend == "TRTLLM"
         and sm >= 100
-        and getattr(case, "sliding_window", None) is not None
-        and getattr(case, "cache", "paged") != "none"
+        and getattr(case, "cache", "paged") == "none"
         and not getattr(case, "is_mla", False)
+        and case.head_dim == 80
+    ):
+        return "TRTLLM Blackwell no-cache fallback is unstable for head_dim 80"
+
+    # These Blackwell sliding-window decode shapes pass in isolation but become
+    # numerically unstable after earlier cases initialize trtllm-gen process
+    # state. Keep the skip exact so other sliding-window decode shapes run.
+    if (
+        backend == "TRTLLM"
+        and sm >= 100
         and case.num_contexts == 0
         and (
             case.num_heads,
             case.num_kv_heads,
             case.head_dim,
-            case.sliding_window,
+            getattr(case, "sliding_window", None),
         )
         in _TRTLLM_BLACKWELL_SLIDING_DECODE_UNSTABLE
     ):
@@ -201,48 +210,19 @@ def unsupported_reason(backend: str, case) -> Optional[str]:
             f"head_dim={case.head_dim}, window={case.sliding_window}"
         )
 
-    # TRTLLM's Blackwell no-cache fallback mismatches the Vanilla golden for the
-    # Qwen2-VL vision tower's head_dim 80 workload; other no-cache head dims in
-    # this sweep still pass on Blackwell.
-    if (
-        backend == "TRTLLM"
-        and sm >= 100
-        and getattr(case, "cache", "paged") == "none"
-        and not getattr(case, "is_mla", False)
-        and case.head_dim == 80
-    ):
-        return "TRTLLM Blackwell no-cache fallback is unstable for head_dim 80"
-
-    # TRTLLM's fp8 generation (XQA) path computes in fp8 and needs the model's
-    # real KV-dequant + output scale state (fed by the Attention module's
-    # projection layers). A bare standalone backend lacks it: supplying a unit
-    # output scale lets the kernel run, but the MHA path then produces garbage
-    # (~1e5 abs error) while only GQA happens to tolerate it. fp8 *context*
-    # (pure prefill) is exercised; FlashInfer validates fp8 decode on every arch.
-    if backend == "TRTLLM" and getattr(case, "kv_dtype", None) == "fp8":
-        has_generation = case.num_contexts < len(case.seq_lens)
-        if has_generation:
-            return (
-                "TRTLLM fp8 KV generation (XQA) needs the model's fp8 scale "
-                "state, absent in the standalone backend; fp8 context is covered "
-                "and FlashInfer validates fp8 decode"
-            )
-
-    # TRTLLM cross-attention works through the standard plumbing (prepare()
-    # derives the cross kv_lens from seq_lens_kv; the backend builds cross_kv from
-    # k/v). The aligned q_len == kv_len case is exercised on TRTLLM. The
-    # q_len != kv_len case is numerically correct in isolation (verified vs the
-    # golden to ~2e-3) but exhibits a cross-case state-dependent mismatch when run
-    # after other cross cases in the same process (a TRTLLM cross-path global not
-    # reset between fresh backend instances), so it is validated via
-    # FlashInfer/Vanilla in the suite.
+    # SM90 forces paged context FMHA whenever a mixed batch has context work.
+    # With an FP8 KV cache, that mode also reaches the MMHA generation phase and
+    # requires an attention-output scale that a standalone MHA backend does not
+    # own. Pure decode, GQA/MQA mixed batches, and Blackwell trtllm-gen all run.
     if (
         backend == "TRTLLM"
-        and getattr(case, "is_cross", False)
-        and list(case.seq_lens) != list(case.seq_lens_kv)
+        and sm in (90,)
+        and getattr(case, "kv_dtype", None) == "fp8"
+        and case.num_heads == case.num_kv_heads
+        and 0 < case.num_contexts < len(case.seq_lens)
     ):
         return (
-            "TRTLLM cross q_len != kv_len is correct standalone (~2e-3) but "
-            "flaky across cross cases in-process; validated via FlashInfer/Vanilla"
+            "TRTLLM SM90 mixed-phase MHA with FP8 KV cache requires an "
+            "attention-output scale unavailable to the standalone backend"
         )
     return None
diff --git a/tests/unittest/_torch/attention/backend_case.py b/tests/unittest/_torch/attention/backend_case.py
index 95368af8b98f..cc30f1e051e0 100644
--- a/tests/unittest/_torch/attention/backend_case.py
+++ b/tests/unittest/_torch/attention/backend_case.py
@@ -80,8 +80,7 @@ class BackendCase:
 
     # Cross-attention: new KV (encoder) tokens per request. None => self-attention
     # (KV tokens == seq_lens). When set, the case is cross-attention (must be
-    # non-causal). Same-length cross runs on TRTLLM/FlashInfer/Vanilla; unequal
-    # q/kv lengths are gated for TRTLLM in the capability matrix.
+    # non-causal).
     seq_lens_kv: Optional[List[int]] = None
 
     dtype: str = "float16"
@@ -963,7 +962,7 @@ def create_metadata(AttentionCls, case, mgr, *, num_contexts: int = 0):
                 max_num_tokens=case.max_num_tokens,
                 kv_cache_manager=mgr,
                 request_ids=request_ids,
-                prompt_lens=case.token_nums,
+                prompt_lens=case.seq_lens if case.is_cross else case.token_nums,
                 kv_layout=kv_layout,
             )
 

From ca66d0c3772e6149fbbb49bc6b2a079771b1e545 Mon Sep 17 00:00:00 2001
From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Date: Wed, 1 Jul 2026 11:38:38 +0000
Subject: [PATCH 04/13] [None][fix] unwaive stable attention backend cases

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
---
 .../_torch/attention/backend_capability.py    | 26 -------------------
 1 file changed, 26 deletions(-)

diff --git a/tests/unittest/_torch/attention/backend_capability.py b/tests/unittest/_torch/attention/backend_capability.py
index 16995e4b36e8..2137a4c22911 100644
--- a/tests/unittest/_torch/attention/backend_capability.py
+++ b/tests/unittest/_torch/attention/backend_capability.py
@@ -77,11 +77,6 @@
 
 _TRTLLM_PAGED_UNSUPPORTED_HEAD_DIMS = (512,)
 _TRTLLM_BLACKWELL_PAGED_UNSUPPORTED_HEAD_DIMS = (96,)
-_TRTLLM_BLACKWELL_SLIDING_DECODE_UNSTABLE = (
-    # Gemma3-27B local layers and Gemma4-31B sliding layers.
-    (32, 16, 128, 1024),
-    (32, 16, 256, 1024),
-)
 
 
 def required_features(case) -> set:
@@ -189,27 +184,6 @@ def unsupported_reason(backend: str, case) -> Optional[str]:
     ):
         return "TRTLLM Blackwell no-cache fallback is unstable for head_dim 80"
 
-    # These Blackwell sliding-window decode shapes pass in isolation but become
-    # numerically unstable after earlier cases initialize trtllm-gen process
-    # state. Keep the skip exact so other sliding-window decode shapes run.
-    if (
-        backend == "TRTLLM"
-        and sm >= 100
-        and case.num_contexts == 0
-        and (
-            case.num_heads,
-            case.num_kv_heads,
-            case.head_dim,
-            getattr(case, "sliding_window", None),
-        )
-        in _TRTLLM_BLACKWELL_SLIDING_DECODE_UNSTABLE
-    ):
-        return (
-            "TRTLLM Blackwell sliding-window pure decode is unstable for "
-            f"num_heads={case.num_heads}, num_kv_heads={case.num_kv_heads}, "
-            f"head_dim={case.head_dim}, window={case.sliding_window}"
-        )
-
     # SM90 forces paged context FMHA whenever a mixed batch has context work.
     # With an FP8 KV cache, that mode also reaches the MMHA generation phase and
     # requires an attention-output scale that a standalone MHA backend does not

From 4f259ca44a59555b1bfdf47b915075f0f6fd3dc2 Mon Sep 17 00:00:00 2001
From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Date: Tue, 30 Jun 2026 11:06:59 +0000
Subject: [PATCH 05/13] [None][fix] remove obsolete Hopper FMHA workaround

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
---
 tensorrt_llm/_torch/attention_backend/trtllm.py   | 10 ----------
 .../_torch/attention/backend_capability.py        | 15 ---------------
 2 files changed, 25 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py
index 3c1829774ac8..bd8f6f6d6897 100644
--- a/tensorrt_llm/_torch/attention_backend/trtllm.py
+++ b/tensorrt_llm/_torch/attention_backend/trtllm.py
@@ -1312,11 +1312,6 @@ def use_nvfp4_output(
             or metadata.runtime_features.has_speculative_draft_tokens
         ) if metadata.runtime_features else False
 
-        # This is a workaround for https://nvbugs/5624818
-        # Paged context FMHA is forced on SM90 for correctness
-        if get_sm_version() == 90:
-            use_paged_context_fmha = True
-
         return self._is_nvfp4_output_kernel_available(
             tokens_per_block=metadata.tokens_per_block,
             attention_mask=attention_mask,
@@ -1493,11 +1488,6 @@ def forward(
         if self.has_fp8_kv_cache:
             metadata.use_paged_context_fmha = True
 
-        # SM90 forces `use_paged_context_fmha` on for correctness
-        # (https://nvbugs/5624818).
-        if get_sm_version() == 90 and metadata.num_contexts > 0:
-            metadata.use_paged_context_fmha = True
-
         # Sparse mqa/gqa attention uses generation kernel which reads Q from qPtr (separate buffer).
         # Force paged context FMHA so QKV preprocessing writes Q to q_buf_2_.
         if (self.sparse_params is not None and getattr(
diff --git a/tests/unittest/_torch/attention/backend_capability.py b/tests/unittest/_torch/attention/backend_capability.py
index 2137a4c22911..03bc59195ad5 100644
--- a/tests/unittest/_torch/attention/backend_capability.py
+++ b/tests/unittest/_torch/attention/backend_capability.py
@@ -184,19 +184,4 @@ def unsupported_reason(backend: str, case) -> Optional[str]:
     ):
         return "TRTLLM Blackwell no-cache fallback is unstable for head_dim 80"
 
-    # SM90 forces paged context FMHA whenever a mixed batch has context work.
-    # With an FP8 KV cache, that mode also reaches the MMHA generation phase and
-    # requires an attention-output scale that a standalone MHA backend does not
-    # own. Pure decode, GQA/MQA mixed batches, and Blackwell trtllm-gen all run.
-    if (
-        backend == "TRTLLM"
-        and sm in (90,)
-        and getattr(case, "kv_dtype", None) == "fp8"
-        and case.num_heads == case.num_kv_heads
-        and 0 < case.num_contexts < len(case.seq_lens)
-    ):
-        return (
-            "TRTLLM SM90 mixed-phase MHA with FP8 KV cache requires an "
-            "attention-output scale unavailable to the standalone backend"
-        )
     return None

From 52d1b1824061f50ed18245697a9188c50b12ffb0 Mon Sep 17 00:00:00 2001
From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Date: Wed, 1 Jul 2026 14:58:42 +0000
Subject: [PATCH 06/13] [None][fix] support large FlashInfer KV appends

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
---
 .../_torch/attention_backend/flashinfer.py    | 76 ++++++++++++++++++-
 .../_torch/attention/backend_capability.py    | 31 ++------
 2 files changed, 82 insertions(+), 25 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py
index 4633850204b6..aac79746e665 100644
--- a/tensorrt_llm/_torch/attention_backend/flashinfer.py
+++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py
@@ -43,6 +43,80 @@
 _FORCE_RAGGED_FA2 = False
 """Used for testing."""
 
+_MAX_CUDA_THREADS_PER_BLOCK = 1024
+
+
+def _slice_paged_kv_cache_heads(
+    paged_kv_cache: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+    start: int,
+    end: int,
+    kv_layout: str,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    if kv_layout == "HND":
+        head_axis = 2
+    elif kv_layout == "NHD":
+        head_axis = 3
+    else:
+        raise ValueError(f"Unsupported kv_layout: {kv_layout}")
+
+    if isinstance(paged_kv_cache, tuple):
+        head_axis -= 1
+        index = [slice(None)] * 4
+        index[head_axis] = slice(start, end)
+        return tuple(cache[tuple(index)] for cache in paged_kv_cache)
+
+    index = [slice(None)] * 5
+    index[head_axis] = slice(start, end)
+    return paged_kv_cache[tuple(index)]
+
+
+def _append_paged_kv_cache(
+    append_key: torch.Tensor,
+    append_value: torch.Tensor,
+    batch_indices: torch.Tensor,
+    positions: torch.Tensor,
+    paged_kv_cache: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+    kv_indices: torch.Tensor,
+    kv_indptr: torch.Tensor,
+    kv_last_page_len: torch.Tensor,
+    kv_layout: str = "NHD",
+) -> None:
+    """Split FlashInfer paged-KV appends that exceed CUDA's CTA limit."""
+    head_dim = append_key.shape[-1]
+    vec_size = max(16 // append_key.element_size(), head_dim // 32)
+    threads_per_head = head_dim // vec_size
+    max_heads_per_launch = _MAX_CUDA_THREADS_PER_BLOCK // threads_per_head
+
+    num_kv_heads = append_key.shape[1]
+    if num_kv_heads <= max_heads_per_launch:
+        flashinfer.page.append_paged_kv_cache(
+            append_key=append_key,
+            append_value=append_value,
+            batch_indices=batch_indices,
+            positions=positions,
+            paged_kv_cache=paged_kv_cache,
+            kv_indices=kv_indices,
+            kv_indptr=kv_indptr,
+            kv_last_page_len=kv_last_page_len,
+            kv_layout=kv_layout,
+        )
+        return
+
+    for start in range(0, num_kv_heads, max_heads_per_launch):
+        end = min(start + max_heads_per_launch, num_kv_heads)
+        flashinfer.page.append_paged_kv_cache(
+            append_key=append_key[:, start:end],
+            append_value=append_value[:, start:end],
+            batch_indices=batch_indices,
+            positions=positions,
+            paged_kv_cache=_slice_paged_kv_cache_heads(paged_kv_cache, start,
+                                                       end, kv_layout),
+            kv_indices=kv_indices,
+            kv_indptr=kv_indptr,
+            kv_last_page_len=kv_last_page_len,
+            kv_layout=kv_layout,
+        )
+
 
 @dataclass(kw_only=True, frozen=True)
 class FlashInferMultiItemParams:
@@ -1748,7 +1822,7 @@ def forward_impl(
                 f"KV cache dtype {kv_cache.dtype} does not match k/v dtype {k.dtype}/{v.dtype}"
             )
 
-            flashinfer.page.append_paged_kv_cache(
+            _append_paged_kv_cache(
                 append_key=k,
                 append_value=v,
                 batch_indices=metadata.batch_indices,
diff --git a/tests/unittest/_torch/attention/backend_capability.py b/tests/unittest/_torch/attention/backend_capability.py
index 03bc59195ad5..91d11be455df 100644
--- a/tests/unittest/_torch/attention/backend_capability.py
+++ b/tests/unittest/_torch/attention/backend_capability.py
@@ -68,12 +68,6 @@
 }
 
 _FLASHINFER_PAGED_UNSUPPORTED_HEAD_DIMS = (96, 512)
-_FLASHINFER_PAGED_APPEND_OVER_1024_THREADS = (
-    # bf16/fp16 append launches 16 threads/head for head_dim=128. With 128 KV
-    # heads, that exceeds CUDA's 1024 threads/block launch limit.
-    (128, 128, "bfloat16"),
-    (128, 128, "float16"),
-)
 
 _TRTLLM_PAGED_UNSUPPORTED_HEAD_DIMS = (512,)
 _TRTLLM_BLACKWELL_PAGED_UNSUPPORTED_HEAD_DIMS = (96,)
@@ -128,27 +122,16 @@ def unsupported_reason(backend: str, case) -> Optional[str]:
     if kv_layout is not None and kv_layout not in caps.get("kv_layouts", ()):
         return f"{backend} does not support kv_layout '{kv_layout}'"
 
-    # FlashInfer's standard paged path has shape-specific kernel limits in this
-    # suite. Keep the skip list explicit so newly supported shapes are not
-    # hidden by a broad dispatch-set predicate.
+    # FlashInfer's standard paged kernels reject these head dimensions. Keep
+    # the skip list explicit so newly supported shapes are not hidden by a
+    # broad dispatch-set predicate.
     if (
         backend == "FLASHINFER"
         and getattr(case, "cache", "paged") != "none"
         and not getattr(case, "is_mla", False)
+        and case.head_dim in _FLASHINFER_PAGED_UNSUPPORTED_HEAD_DIMS
     ):
-        kv_dtype = getattr(case, "kv_dtype", None) or getattr(case, "dtype", None)
-        if case.head_dim in _FLASHINFER_PAGED_UNSUPPORTED_HEAD_DIMS:
-            return f"FLASHINFER paged attention is unstable for head_dim {case.head_dim}"
-        if (
-            case.head_dim,
-            case.num_kv_heads,
-            kv_dtype,
-        ) in _FLASHINFER_PAGED_APPEND_OVER_1024_THREADS:
-            return (
-                "FLASHINFER paged KV append exceeds CUDA's 1024 threads/block "
-                f"limit for head_dim={case.head_dim}, "
-                f"num_kv_heads={case.num_kv_heads}, kv_dtype={kv_dtype}"
-            )
+        return f"FLASHINFER paged kernels do not support head_dim {case.head_dim}"
 
     # TRTLLM's standard paged FMHA/MMHA kernels in this build do not cover the
     # Gemma4 head_dim 512 path.
@@ -158,7 +141,7 @@ def unsupported_reason(backend: str, case) -> Optional[str]:
         and not getattr(case, "is_mla", False)
         and case.head_dim in _TRTLLM_PAGED_UNSUPPORTED_HEAD_DIMS
     ):
-        return f"TRTLLM paged attention is unstable for head_dim {case.head_dim}"
+        return f"TRTLLM paged attention does not support head_dim {case.head_dim}"
 
     # TRTLLM's Blackwell paged fallback path aborts for the Phi-3 head_dim 96
     # shape. Hopper covers that context config, but Blackwell must skip it
@@ -170,7 +153,7 @@ def unsupported_reason(backend: str, case) -> Optional[str]:
         and not getattr(case, "is_mla", False)
         and case.head_dim in _TRTLLM_BLACKWELL_PAGED_UNSUPPORTED_HEAD_DIMS
     ):
-        return f"TRTLLM Blackwell paged attention is unstable for head_dim {case.head_dim}"
+        return f"TRTLLM Blackwell paged attention does not support head_dim {case.head_dim}"
 
     # TRTLLM's Blackwell no-cache fallback mismatches the Vanilla golden for the
     # Qwen2-VL vision tower's head_dim 80 workload; other no-cache head dims in

From 96f9813e23922ba244d6c88cd7cf205299fe8873 Mon Sep 17 00:00:00 2001
From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Date: Thu, 2 Jul 2026 01:32:09 +0000
Subject: [PATCH 07/13] [None][fix] restore missing MLA generation kernel guard

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
---
 .../fmha/flashinfer_trtllm_gen.py             | 11 ++++++++++
 .../_torch/attention/backend_capability.py    | 22 +++++++++++++++++++
 2 files changed, 33 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 b5fea40413d2..425450366d26 100644
--- a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py
+++ b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py
@@ -408,6 +408,9 @@ class FlashInferTrtllmGenFmha(PhasedFmha):
         (320, 256),
         (576, 512),
     }
+    MISSING_MLA_GENERATION_KERNELS = {
+        (576, 512, 32),
+    }
 
     def __init__(self, attn: "TrtllmAttention"):
         super().__init__(attn)
@@ -564,6 +567,14 @@ def _check_mla_generation_support(
                 f"headDimQk={head_dim_qk}, headDimV={head_dim_v}. Supported: {supported}.",
             )
 
+        if (head_dim_qk, head_dim_v, tokens_per_block) in cls.MISSING_MLA_GENERATION_KERNELS:
+            return (
+                False,
+                f"[Generation][MLA] Missing 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(
diff --git a/tests/unittest/_torch/attention/backend_capability.py b/tests/unittest/_torch/attention/backend_capability.py
index 91d11be455df..7917c878e33a 100644
--- a/tests/unittest/_torch/attention/backend_capability.py
+++ b/tests/unittest/_torch/attention/backend_capability.py
@@ -155,6 +155,28 @@ def unsupported_reason(backend: str, case) -> Optional[str]:
     ):
         return f"TRTLLM Blackwell paged attention does not support head_dim {case.head_dim}"
 
+    # On Blackwell, TRTLLM-Gen is the supported MLA generation path. The current
+    # kernel set explicitly lacks the DeepSeek-family decode shape at page_size
+    # 32, and the generic fallback then tries to JIT an unsupported generated
+    # kernel. FlashInfer/Vanilla still validate these standalone MLA gen cases.
+    if (
+        backend == "TRTLLM"
+        and sm >= 100
+        and getattr(case, "is_mla", False)
+        and getattr(case, "num_contexts", 0) == 0
+    ):
+        head_dim_qk = getattr(case, "kv_lora_rank", 0) + getattr(case, "qk_rope_head_dim", 0)
+        head_dim_v = getattr(case, "kv_lora_rank", 0)
+        if (head_dim_qk, head_dim_v, getattr(case, "page_size", None)) == (
+            576,
+            512,
+            32,
+        ):
+            return (
+                "TRTLLM-Gen Blackwell MLA generation is missing the decode "
+                "kernel for headDimQk=576, headDimV=512, page_size=32"
+            )
+
     # TRTLLM's Blackwell no-cache fallback mismatches the Vanilla golden for the
     # Qwen2-VL vision tower's head_dim 80 workload; other no-cache head dims in
     # this sweep still pass on Blackwell.

From aadd5a4e8949cc6347896c0566883a41f8e8dea3 Mon Sep 17 00:00:00 2001
From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Date: Thu, 2 Jul 2026 05:43:53 +0000
Subject: [PATCH 08/13] [None][fix] support FP8 KV MLA in TRTLLM-gen

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
---
 .../fmha/flashinfer_trtllm_gen.py             | 53 +++++++++++++------
 .../_torch/attention/backend_capability.py    | 22 --------
 2 files changed, 37 insertions(+), 38 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 425450366d26..fe4fa839ad8d 100644
--- a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py
+++ b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py
@@ -408,9 +408,6 @@ class FlashInferTrtllmGenFmha(PhasedFmha):
         (320, 256),
         (576, 512),
     }
-    MISSING_MLA_GENERATION_KERNELS = {
-        (576, 512, 32),
-    }
 
     def __init__(self, attn: "TrtllmAttention"):
         super().__init__(attn)
@@ -529,7 +526,6 @@ def _get_attention_chunk_size(attn: "TrtllmAttention") -> int:
     def _check_mla_generation_support(
         cls,
         head_size: int,
-        tokens_per_block: int,
         kv_lora_rank: Optional[int],
         qk_rope_head_dim: Optional[int],
     ) -> Tuple[bool, str]:
@@ -567,14 +563,6 @@ def _check_mla_generation_support(
                 f"headDimQk={head_dim_qk}, headDimV={head_dim_v}. Supported: {supported}.",
             )
 
-        if (head_dim_qk, head_dim_v, tokens_per_block) in cls.MISSING_MLA_GENERATION_KERNELS:
-            return (
-                False,
-                f"[Generation][MLA] Missing 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(
@@ -765,7 +753,6 @@ 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,
                     kv_lora_rank=attn.kv_lora_rank,
                     qk_rope_head_dim=attn.qk_rope_head_dim,
                 )
@@ -1222,9 +1209,43 @@ def run_mla_generation(
         mla_head_dim_qk = kv_lora_rank + qk_rope_head_dim
         q_len_per_req = params.num_tokens // batch_beam if batch_beam > 0 else 1
 
-        query = params.qkv_input.view(batch_beam, q_len_per_req, attn.num_heads, mla_head_dim_qk)
+        if QuantMode(attn.quant_mode).has_fp8_kv_cache():
+            quant_q_buffer = fwd.quant_q_buffer
+            bmm1_scale_buffer = fwd.mla_bmm1_scale
+            bmm2_scale_buffer = fwd.mla_bmm2_scale
+            if quant_q_buffer is None or bmm1_scale_buffer is None or bmm2_scale_buffer is None:
+                raise RuntimeError(
+                    "FP8 MLA generation requires quant_q_buffer, "
+                    "mla_bmm1_scale, and mla_bmm2_scale."
+                )
 
-        bmm1_scale = 1.0 / (attn.q_scaling * math.sqrt(qk_nope_head_dim + qk_rope_head_dim))
+            expected_q_elements = params.num_tokens * attn.num_heads * mla_head_dim_qk
+            if quant_q_buffer.numel() < expected_q_elements:
+                raise RuntimeError(
+                    f"FP8 MLA quant_q_buffer has {quant_q_buffer.numel()} elements; "
+                    f"expected at least {expected_q_elements}."
+                )
+            if bmm1_scale_buffer.dtype != torch.float32 or bmm1_scale_buffer.numel() < 1:
+                raise RuntimeError("FP8 MLA bmm1 scale must contain a float32 value.")
+            if bmm2_scale_buffer.dtype != torch.float32 or bmm2_scale_buffer.numel() < 1:
+                raise RuntimeError("FP8 MLA bmm2 scale must contain a float32 value.")
+
+            query = (
+                quant_q_buffer.view(torch.uint8)
+                .flatten()[:expected_q_elements]
+                .view(torch.float8_e4m3fn)
+                .view(batch_beam, q_len_per_req, attn.num_heads, mla_head_dim_qk)
+            )
+            # FlashInfer converts tensor BMM1 scales to log2 internally. The
+            # producer stores the regular scale at index 0 and log2 at index 1.
+            bmm1_scale = bmm1_scale_buffer.flatten()[:1]
+            bmm2_scale = bmm2_scale_buffer.flatten()[:1]
+        else:
+            query = params.qkv_input.view(
+                batch_beam, q_len_per_req, attn.num_heads, mla_head_dim_qk
+            )
+            bmm1_scale = 1.0 / (attn.q_scaling * math.sqrt(qk_nope_head_dim + qk_rope_head_dim))
+            bmm2_scale = 1.0
         workspace_buffer = params.workspace.view(-1, 4)
         _clear_multi_ctas_kv_counter_workspace(
             workspace_buffer, attn.num_heads, meta.max_num_requests, self._multi_processor_count
@@ -1243,7 +1264,7 @@ def run_mla_generation(
             0,  # sparse_mla_top_k
             params.context_buf.view(batch_beam, q_len_per_req, attn.num_heads, kv_lora_rank),  # out
             bmm1_scale,  # bmm1_scale
-            1.0,  # bmm2_scale
+            bmm2_scale,  # bmm2_scale
             fwd.attention_sinks,  # sinks
             None,  # skip_softmax_threshold_scale_factor
             self._enable_pdl,  # enable_pdl
diff --git a/tests/unittest/_torch/attention/backend_capability.py b/tests/unittest/_torch/attention/backend_capability.py
index 7917c878e33a..91d11be455df 100644
--- a/tests/unittest/_torch/attention/backend_capability.py
+++ b/tests/unittest/_torch/attention/backend_capability.py
@@ -155,28 +155,6 @@ def unsupported_reason(backend: str, case) -> Optional[str]:
     ):
         return f"TRTLLM Blackwell paged attention does not support head_dim {case.head_dim}"
 
-    # On Blackwell, TRTLLM-Gen is the supported MLA generation path. The current
-    # kernel set explicitly lacks the DeepSeek-family decode shape at page_size
-    # 32, and the generic fallback then tries to JIT an unsupported generated
-    # kernel. FlashInfer/Vanilla still validate these standalone MLA gen cases.
-    if (
-        backend == "TRTLLM"
-        and sm >= 100
-        and getattr(case, "is_mla", False)
-        and getattr(case, "num_contexts", 0) == 0
-    ):
-        head_dim_qk = getattr(case, "kv_lora_rank", 0) + getattr(case, "qk_rope_head_dim", 0)
-        head_dim_v = getattr(case, "kv_lora_rank", 0)
-        if (head_dim_qk, head_dim_v, getattr(case, "page_size", None)) == (
-            576,
-            512,
-            32,
-        ):
-            return (
-                "TRTLLM-Gen Blackwell MLA generation is missing the decode "
-                "kernel for headDimQk=576, headDimV=512, page_size=32"
-            )
-
     # TRTLLM's Blackwell no-cache fallback mismatches the Vanilla golden for the
     # Qwen2-VL vision tower's head_dim 80 workload; other no-cache head dims in
     # this sweep still pass on Blackwell.

From ec492107a546dc14cf5c6ee05560c17e2abc5684 Mon Sep 17 00:00:00 2001
From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Date: Thu, 2 Jul 2026 06:09:48 +0000
Subject: [PATCH 09/13] [None][fix] fall back for FP8 multi-token MLA

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
---
 .../fmha/flashinfer_trtllm_gen.py              | 18 ++++++++++++++++++
 1 file changed, 18 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 fe4fa839ad8d..7a8f8cfc8b66 100644
--- a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py
+++ b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py
@@ -758,6 +758,24 @@ def _is_supported_with_reason(
                 )
                 if not supported:
                     return False, reason
+                head_dim_qk = (attn.kv_lora_rank or 0) + (attn.qk_rope_head_dim or 0)
+                head_dim_v = attn.kv_lora_rank or 0
+                is_multi_token = meta.num_generations > 0 and q.size(0) > meta.num_generations
+                if (
+                    has_fp8_kv
+                    and is_multi_token
+                    and (
+                        head_dim_qk,
+                        head_dim_v,
+                        tokens_per_block,
+                    )
+                    == (576, 512, 32)
+                ):
+                    return (
+                        False,
+                        "[Generation][MLA] FP8 multi-token generation is not supported for "
+                        "headDimQk=576, headDimV=512, tokens_per_block=32.",
+                    )
 
         if tokens_per_block <= 0:
             return False, "tokens_per_block must be positive."

From f497c9c1f92f104fc5ed05ca37ebd42ff342d9a5 Mon Sep 17 00:00:00 2001
From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Date: Thu, 2 Jul 2026 09:58:33 +0000
Subject: [PATCH 10/13] [None][fix] gate FP8 multi-token MLA fallback by
 FlashInfer version

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
---
 .../_torch/attention_backend/fmha/flashinfer_trtllm_gen.py     | 3 +++
 1 file changed, 3 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 7a8f8cfc8b66..f860a406db47 100644
--- a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py
+++ b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py
@@ -43,6 +43,7 @@
 from typing import TYPE_CHECKING, List, Optional, Tuple
 
 import torch
+from packaging.version import Version
 
 from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE, get_env_enable_pdl
 
@@ -758,6 +759,7 @@ def _is_supported_with_reason(
                 )
                 if not supported:
                     return False, reason
+                # TODO: Remove this fallback after upgrading FlashInfer to 0.6.14 or later.
                 head_dim_qk = (attn.kv_lora_rank or 0) + (attn.qk_rope_head_dim or 0)
                 head_dim_v = attn.kv_lora_rank or 0
                 is_multi_token = meta.num_generations > 0 and q.size(0) > meta.num_generations
@@ -770,6 +772,7 @@ def _is_supported_with_reason(
                         tokens_per_block,
                     )
                     == (576, 512, 32)
+                    and Version(flashinfer.__version__) < Version("0.6.14")
                 ):
                     return (
                         False,

From a385754ca9232b316241eea297f90f3d74ae8979 Mon Sep 17 00:00:00 2001
From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Date: Mon, 6 Jul 2026 03:21:42 +0000
Subject: [PATCH 11/13] [None][fix] remove obsolete FlashInfer MLA fallback

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
---
 .../fmha/flashinfer_trtllm_gen.py             | 21 -------------------
 1 file changed, 21 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 f860a406db47..fe4fa839ad8d 100644
--- a/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py
+++ b/tensorrt_llm/_torch/attention_backend/fmha/flashinfer_trtllm_gen.py
@@ -43,7 +43,6 @@
 from typing import TYPE_CHECKING, List, Optional, Tuple
 
 import torch
-from packaging.version import Version
 
 from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE, get_env_enable_pdl
 
@@ -759,26 +758,6 @@ def _is_supported_with_reason(
                 )
                 if not supported:
                     return False, reason
-                # TODO: Remove this fallback after upgrading FlashInfer to 0.6.14 or later.
-                head_dim_qk = (attn.kv_lora_rank or 0) + (attn.qk_rope_head_dim or 0)
-                head_dim_v = attn.kv_lora_rank or 0
-                is_multi_token = meta.num_generations > 0 and q.size(0) > meta.num_generations
-                if (
-                    has_fp8_kv
-                    and is_multi_token
-                    and (
-                        head_dim_qk,
-                        head_dim_v,
-                        tokens_per_block,
-                    )
-                    == (576, 512, 32)
-                    and Version(flashinfer.__version__) < Version("0.6.14")
-                ):
-                    return (
-                        False,
-                        "[Generation][MLA] FP8 multi-token generation is not supported for "
-                        "headDimQk=576, headDimV=512, tokens_per_block=32.",
-                    )
 
         if tokens_per_block <= 0:
             return False, "tokens_per_block must be positive."

From 3a40842f694aae2eb3ac52d9409a10ded12f14c8 Mon Sep 17 00:00:00 2001
From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Date: Tue, 7 Jul 2026 14:42:22 +0000
Subject: [PATCH 12/13] [None][fix] fix FP8 KV-only attention fallback

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
---
 tensorrt_llm/_torch/attention_backend/fmha/fallback.py   | 1 +
 tensorrt_llm/_torch/attention_backend/trtllm.py          | 8 ++++----
 .../unittest/_torch/attention/test_attention_op_sync.py  | 9 +--------
 3 files changed, 6 insertions(+), 12 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py b/tensorrt_llm/_torch/attention_backend/fmha/fallback.py
index 0a0d89dd5adf..d39ecaf61865 100644
--- a/tensorrt_llm/_torch/attention_backend/fmha/fallback.py
+++ b/tensorrt_llm/_torch/attention_backend/fmha/fallback.py
@@ -36,6 +36,7 @@
         "attention_mask_data",  # custom-mask code path
         "out_scale_sf",  # promoted into ``out_scale`` in ``TrtllmAttention.forward`` for NVFP4 path
         "skip_mla_rope_generation",  # handled in ``TrtllmAttention.forward`` for the test-only MLA path
+        "timestep",  # used to populate skip-softmax params in ``TrtllmAttention.forward``
     }
 )
 
diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py
index bd8f6f6d6897..e46c13fb586f 100644
--- a/tensorrt_llm/_torch/attention_backend/trtllm.py
+++ b/tensorrt_llm/_torch/attention_backend/trtllm.py
@@ -1482,10 +1482,10 @@ def forward(
         # Cross-attention uses the THOP path; the trtllm-gen backend API does
         # not carry encoder K/V tensors yet.
 
-        # cpp/tensorrt_llm/thop/attentionOp.cpp enables mFP8ContextFMHA for an
-        # FP8 KV cache only when use_paged_context_fmha is true. Force paged
-        # context so QKV preprocessing and context FMHA use the FP8 path.
-        if self.has_fp8_kv_cache:
+        # Paged context with an FP8 KV cache selects mFP8ContextFMHA, which
+        # requires a quantized-output scale. Keep KV-only quantization on the
+        # BF16-output path when the caller did not request output quantization.
+        if self.has_fp8_kv_cache and forward_args.out_scale is not None:
             metadata.use_paged_context_fmha = True
 
         # Sparse mqa/gqa attention uses generation kernel which reads Q from qPtr (separate buffer).
diff --git a/tests/unittest/_torch/attention/test_attention_op_sync.py b/tests/unittest/_torch/attention/test_attention_op_sync.py
index d1beb769f16d..cfd56978b848 100644
--- a/tests/unittest/_torch/attention/test_attention_op_sync.py
+++ b/tests/unittest/_torch/attention/test_attention_op_sync.py
@@ -59,20 +59,13 @@
 }
 
 _THOP_KWARG_SOURCE_ALIASES: dict[str, tuple[str, tuple[str, ...]]] = {
+    "beam_width": ("metadata", ("effective_beam_width",)),
     "context_lengths": ("metadata", ("prompt_lens_cuda_runtime",)),
     "head_size": ("attn", ("head_dim",)),
     "host_context_lengths": ("metadata", ("prompt_lens_cpu_runtime",)),
     "host_past_key_value_lengths": ("metadata", ("kv_lens_runtime",)),
     "host_request_types": ("metadata", ("host_request_types_runtime",)),
     "sequence_length": ("metadata", ("kv_lens_cuda_runtime",)),
-    "skip_softmax_threshold_scale_factor_decode": (
-        "skip_softmax_kernel_params",
-        ("threshold_scale_factor_decode",),
-    ),
-    "skip_softmax_threshold_scale_factor_prefill": (
-        "skip_softmax_kernel_params",
-        ("threshold_scale_factor_prefill",),
-    ),
     "spec_decoding_target_max_draft_tokens": (
         "metadata",
         ("max_total_draft_tokens",),

From 7483843cb3b7854f5d68bb184dd78b8cf681ef42 Mon Sep 17 00:00:00 2001
From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
Date: Tue, 7 Jul 2026 15:16:43 +0000
Subject: [PATCH 13/13] [None][fix] avoid forcing paged context FMHA for FP8 KV

Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com>
---
 tensorrt_llm/_torch/attention_backend/trtllm.py | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/tensorrt_llm/_torch/attention_backend/trtllm.py b/tensorrt_llm/_torch/attention_backend/trtllm.py
index e46c13fb586f..bb3c8859b9fe 100644
--- a/tensorrt_llm/_torch/attention_backend/trtllm.py
+++ b/tensorrt_llm/_torch/attention_backend/trtllm.py
@@ -1482,12 +1482,6 @@ def forward(
         # Cross-attention uses the THOP path; the trtllm-gen backend API does
         # not carry encoder K/V tensors yet.
 
-        # Paged context with an FP8 KV cache selects mFP8ContextFMHA, which
-        # requires a quantized-output scale. Keep KV-only quantization on the
-        # BF16-output path when the caller did not request output quantization.
-        if self.has_fp8_kv_cache and forward_args.out_scale is not None:
-            metadata.use_paged_context_fmha = True
-
         # Sparse mqa/gqa attention uses generation kernel which reads Q from qPtr (separate buffer).
         # Force paged context FMHA so QKV preprocessing writes Q to q_buf_2_.
         if (self.sparse_params is not None and getattr(