[TRTLLM-16624][fix] Fix Helix MLA decode kernel selection

tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py

@@ -60,6 +60,9 @@
 # ``assumed_align`` the decode runner compiles the page table and the KV
 # lengths with; their runtime pointers must keep it.
 _KERNEL_INPUT_ALIGN_BYTES = 16
+# Page size of the fallback thop op's precompiled TRTLLM-Gen MLA generation
+# kernels.
+_FALLBACK_MLA_GENERATION_TOKENS_PER_BLOCK = 32
 
 
 @triton.jit(
@@ -476,13 +479,27 @@ def _is_supported_with_reason(
     ) -> tuple[bool, str]:
         if fwd.attention_input_type != AttentionInputType.generation_only:
             return False, "CuTe DSL MLA FMHA only supports generation-only attention."
-        # Disable mixed context/generation batches until the CuTe DSL host
-        # overhead is reduced, except when no other library reads this
-        # layer's FP4 MLA 128-token KV pages (for example an NVFP4 manager's
-        # FP8 fallback layer); this path slices the generation inputs and
-        # page tables past num_contexts.
+        if meta.num_generations <= 0:
+            return False, "CuTe DSL MLA FMHA only supports decode-only batches."
+        seq_len_q = q.shape[0] // meta.num_generations
+        batch_size = meta.num_generations
+        helix_verify_group = seq_len_q > 1 and meta.helix_position_offsets is not None
+        # Under Helix the only other MLA generation library is the fallback
+        # thop op, since FlashInfer TRTLLM-Gen and PrimTS do not support Helix.
+        # Its TRTLLM-Gen generation pads 33 to 63 heads to a 64-row Q tile whose
+        # kernels exist only precompiled, for 32-token pages.
+        helix_decode_without_fallback_kernel = (
+            meta.helix_position_offsets is not None
+            and 32 < attn.num_heads < 64
+            and meta.tokens_per_block != _FALLBACK_MLA_GENERATION_TOKENS_PER_BLOCK
+        )
+        # Keep the mixed-batch performance restriction unless the generation
+        # sub-batch has no supported fallback. PhasedFmha slices generation
+        # lengths, and _run_mla_decode offsets the page table past contexts.
         exclusive_page_size = self._page_size_exclusive_to_cute_dsl(meta)
-        if meta.num_generations <= 0 or (meta.num_contexts != 0 and not exclusive_page_size):
+        if meta.num_contexts != 0 and not (
+            exclusive_page_size or helix_verify_group or helix_decode_without_fallback_kernel
+        ):
             return False, "CuTe DSL MLA FMHA only supports decode-only batches."
         if meta.beam_width != 1:
             return False, f"Beam search is not supported, got beam_width={meta.beam_width}."
@@ -498,8 +515,6 @@ def _is_supported_with_reason(
             or getattr(meta, "is_spec_dec_dynamic_tree", False)
         ):
             return False, "CuTe DSL MLA FMHA does not support custom/tree speculative masks."
-        seq_len_q = q.shape[0] // meta.num_generations
-        batch_size = meta.num_generations
         if meta.helix_position_offsets is not None:
             if seq_len_q != 1 and not meta._helix_spec_tokens_valid:
                 # Multi-token decode under helix needs the per-token bound /
@@ -522,19 +537,11 @@ def _is_supported_with_reason(
 
         from tensorrt_llm._torch.autotuner import AutoTuner
 
-        # A multi-token verify group at H=96 has nowhere else to go: TRTLLM-Gen
-        # rejects 64 < num_heads_q < 128 outright, so falling through the perf
-        # gate does not reach a faster kernel, it reaches an executor-init
-        # failure. That makes this a correctness carve-out rather than a perf
-        # tradeoff, so it is decided here instead of being folded into
-        # _PERF_MIN_BATCH_FP8, which stays a pure measured-win table.
-        # Single-token H=96 is already a measured table entry and still goes
-        # through the gate. Helix with seq_len_q > 1 implies
-        # _helix_spec_tokens_valid -- the helix block above returns False
-        # otherwise.
-        helix_h96_verify_group = (
-            attn.num_heads == 96 and seq_len_q > 1 and meta.helix_position_offsets is not None
-        )
+        # No fallback library supports a multi-token Helix verify group.
+        # Keep its correctness admission separate from the measured-win table;
+        # the metadata checks above and kernel can_implement below still apply.
+        # Other single-token Helix shapes retain the performance gate unless
+        # the fallback has no kernel for them.
         # The same failure without Helix: an FP8-KV multi-token verify group
         # (e.g. DSpark under attention DP, where every rank keeps all 96 heads)
         # is not in the FP8 perf table, and on SM100/SM103 the next library is
@@ -551,7 +558,15 @@ def _is_supported_with_reason(
             and self._get_kernel_dtype(attn, q) == torch.float8_e4m3fn
             and get_sm_version() in (100, 103)
         )
-        if not (helix_h96_verify_group or fp8_wide_head_verify_group or exclusive_page_size):
+        # Single-token Helix decode without a fallback kernel is another
+        # coverage case kept out of the measured-win table; can_implement
+        # remains authoritative.
+        if not (
+            helix_verify_group
+            or fp8_wide_head_verify_group
+            or helix_decode_without_fallback_kernel
+            or exclusive_page_size
+        ):
             # Perf gate (NOT a correctness limit).
             favorable, reason = self._is_perf_favorable(
                 attn.num_heads,

tests/unittest/_torch/attention/test_cute_dsl_mla.py

@@ -136,6 +136,9 @@ def _select_mla_generation_fmha(
     fp8_kv: bool,
     helix: bool,
     batch_size: int,
+    tokens_per_block: int = 64,
+    kernel_supported: bool = True,
+    num_contexts: int = 0,
 ):
     """Run FmhaManager._select_mla over [CuteDslMlaFmha, TRTLLM-Gen stub] on CPU."""
     from tensorrt_llm._torch.attention.backends.fmha import cute_dsl_mla
@@ -154,7 +157,9 @@ def _select_mla_generation_fmha(
     # The kernel's can_implement needs nvidia-cutlass-dsl; this test covers only
     # the host-side routing in front of it.
     monkeypatch.setattr(
-        CuteDslMlaFmha, "_kernel_can_implement", staticmethod(lambda *args, **kwargs: (True, ""))
+        CuteDslMlaFmha,
+        "_kernel_can_implement",
+        staticmethod(lambda *args, **kwargs: (kernel_supported, "")),
     )
 
     num_tokens = batch_size * seq_len_q
@@ -166,17 +171,18 @@ def _select_mla_generation_fmha(
         layer_idx=0,
     )
     meta = SimpleNamespace(
-        num_contexts=0,
+        num_contexts=num_contexts,
         num_generations=batch_size,
         beam_width=1,
         is_spec_dec_tree=False,
         is_spec_dec_dynamic_tree=False,
         helix_position_offsets=torch.zeros(num_tokens, dtype=torch.int32) if helix else None,
         _helix_spec_tokens_valid=helix,
         kv_cache_manager=SimpleNamespace(
-            get_buffers=lambda layer_idx: torch.empty(0, dtype=torch.bfloat16)
+            dtype=DataType.FP8 if fp8_kv else DataType.BF16,
+            get_buffers=lambda layer_idx: torch.empty(0, dtype=torch.bfloat16),
         ),
-        tokens_per_block=64,
+        tokens_per_block=tokens_per_block,
     )
     fwd = AttentionForwardArgs(
         attention_input_type=AttentionInputType.generation_only,
@@ -196,6 +202,7 @@ def _select_mla_generation_fmha(
     return "cute_dsl_mla" if selected is cute else "trtllm_gen"
 
 
+@pytest.mark.cpu_only
 @pytest.mark.parametrize(
     "sm, num_heads, seq_len_q, fp8_kv, helix, expected",
     [
@@ -209,6 +216,19 @@ def _select_mla_generation_fmha(
         # Helix H=96 verify groups are unchanged.
         (103, 96, 4, True, True, "cute_dsl_mla"),
         (107, 96, 4, True, True, "cute_dsl_mla"),
+        # Helix verify groups also require CuTe DSL with fewer TP-local heads.
+        (103, 48, 6, True, True, "cute_dsl_mla"),
+        (100, 48, 6, True, True, "cute_dsl_mla"),
+        (103, 48, 6, False, True, "cute_dsl_mla"),
+        (103, 24, 6, True, True, "cute_dsl_mla"),
+        (100, 24, 6, True, True, "cute_dsl_mla"),
+        (103, 24, 6, False, True, "cute_dsl_mla"),
+        # The same verify shape without Helix keeps the perf gate.
+        (103, 48, 6, True, False, "trtllm_gen"),
+        # Single-token Helix decode at 48 heads on 64-token pages has no
+        # fallback kernel on either architecture.
+        (103, 48, 1, True, True, "cute_dsl_mla"),
+        (100, 48, 1, True, True, "cute_dsl_mla"),
         # Single-token H=96 is unchanged (measured perf-table entry).
         (103, 96, 1, True, False, "cute_dsl_mla"),
         # Head counts TRTLLM-Gen serves keep the perf gate.
@@ -218,7 +238,7 @@ def _select_mla_generation_fmha(
         (103, 96, 8, False, False, "cute_dsl_mla"),
     ],
 )
-def test_mla_generation_routes_wide_head_fp8_verify_away_from_trtllm_gen(
+def test_mla_generation_routes_verify_groups_to_supported_fmha(
     monkeypatch, sm, num_heads, seq_len_q, fp8_kv, helix, expected
 ) -> None:
     # The perf gate admits single-token H=96 FP8 decode only from batch_size 8.
@@ -258,6 +278,150 @@ def test_mla_generation_single_token_wide_head_fp8_below_perf_floor_falls_throug
     )
 
 
+@pytest.mark.cpu_only
+@pytest.mark.parametrize(
+    "page_size, fp8_kv, helix, sm, num_heads, kernel_supported, expected",
+    [
+        # The fallback pads 33-63 heads to a 64-row Q tile that has kernels
+        # for 32-token pages only.
+        (64, True, True, 103, 48, True, "cute_dsl_mla"),
+        (64, True, True, 100, 48, True, "cute_dsl_mla"),
+        (64, False, True, 103, 48, True, "cute_dsl_mla"),
+        (128, True, True, 103, 48, True, "cute_dsl_mla"),
+        (64, True, True, 103, 40, True, "cute_dsl_mla"),
+        # The fallback serves these, so the perf gate still decides.
+        (32, True, True, 103, 48, True, "trtllm_gen"),
+        (64, True, True, 103, 24, True, "trtllm_gen"),
+        (64, True, True, 103, 64, True, "trtllm_gen"),
+        (64, True, True, 103, 96, True, "trtllm_gen"),
+        (64, True, False, 103, 48, True, "trtllm_gen"),
+        (64, True, True, 103, 48, False, "trtllm_gen"),
+    ],
+)
+def test_mla_generation_helix_decode_without_fallback_kernel(
+    monkeypatch, page_size, fp8_kv, helix, sm, num_heads, kernel_supported, expected
+) -> None:
+    assert (
+        _select_mla_generation_fmha(
+            monkeypatch,
+            sm=sm,
+            num_heads=num_heads,
+            seq_len_q=1,
+            fp8_kv=fp8_kv,
+            helix=helix,
+            batch_size=1,
+            tokens_per_block=page_size,
+            kernel_supported=kernel_supported,
+        )
+        == expected
+    )
+
+
+@pytest.mark.cpu_only
+@pytest.mark.parametrize(
+    "page, seq_len_q, fp8_kv, helix, supported, expected",
+    [
+        (64, 1, True, True, True, "cute_dsl_mla"),
+        (64, 1, True, True, False, "trtllm_gen"),
+        (32, 1, True, True, True, "trtllm_gen"),
+        (64, 1, False, True, True, "cute_dsl_mla"),
+        (64, 1, True, False, True, "trtllm_gen"),
+        (64, 6, True, True, True, "cute_dsl_mla"),
+        (32, 6, True, True, True, "cute_dsl_mla"),
+    ],
+)
+def test_mla_generation_mixed_helix_coverage(
+    monkeypatch, page, seq_len_q, fp8_kv, helix, supported, expected
+) -> None:
+    assert (
+        _select_mla_generation_fmha(
+            monkeypatch,
+            sm=103,
+            num_heads=48,
+            seq_len_q=seq_len_q,
+            fp8_kv=fp8_kv,
+            helix=helix,
+            batch_size=2,
+            tokens_per_block=page,
+            kernel_supported=supported,
+            num_contexts=2,
+        )
+        == expected
+    )
+
+
+@pytest.mark.cpu_only
+@pytest.mark.parametrize("seq_len_q", [1, 6])
+def test_mla_generation_mixed_page_table_offset(monkeypatch, seq_len_q) -> None:
+    page, heads, latent, rope = 64, 48, 512, 64
+    num_tokens = 2 * seq_len_q
+    cache = torch.zeros((12, 1, page, 1, latent + rope), dtype=torch.float8_e4m3fn)
+    blocks = torch.arange(12).view(1, 4, 1, 3)
+    lengths = torch.tensor([999, 888, 5, 7], dtype=torch.int32)
+    # The Helix buffers are generation-relative; the tail stands for entries
+    # past this step's generation tokens.
+    kv_bounds = torch.arange(100, 100 + num_tokens + 4, dtype=torch.int32)
+    meta = SimpleNamespace(
+        num_contexts=2,
+        max_num_requests=4,
+        helix_position_offsets=torch.zeros(num_tokens, dtype=torch.int32),
+        helix_kv_bounds=kv_bounds,
+        _helix_spec_tokens_valid=seq_len_q > 1,
+        kv_cache_manager=SimpleNamespace(get_buffers=lambda _: cache),
+        kv_cache_block_offsets=blocks,
+        host_kv_cache_pool_mapping=torch.tensor([[0, 0]]),
+    )
+    attn = SimpleNamespace(
+        kv_lora_rank=latent,
+        qk_rope_head_dim=rope,
+        qk_nope_head_dim=128,
+        num_heads=heads,
+        q_scaling=1,
+        layer_idx=0,
+        get_local_layer_idx=lambda _: 0,
+    )
+    params = SimpleNamespace(
+        attn=attn,
+        meta=meta,
+        num_requests=2,
+        tokens_per_block=page,
+        sequence_lengths=lengths[2:],
+        workspace=torch.empty(0, dtype=torch.int8),
+        fwd=SimpleNamespace(
+            quant_q_buffer=None,
+            mla_bmm1_scale=torch.ones(2, dtype=torch.float32),
+            mla_bmm2_scale=torch.ones(1, dtype=torch.float32),
+            softmax_stats_tensor=torch.empty((num_tokens, heads, 2), dtype=torch.float32),
+        ),
+    )
+    library = CuteDslMlaFmha.__new__(CuteDslMlaFmha)
+    monkeypatch.setattr(library, "_required_workspace_size", lambda *args: 0)
+    monkeypatch.setattr(library, "_split_kv_plan", lambda *args: None)
+    calls = []
+    monkeypatch.setattr(
+        torch.ops.trtllm,
+        "cute_dsl_mla_decode_fp8_blackwell",
+        lambda *args: calls.append(args),
+        raising=False,
+    )
+    library._run_mla_decode(
+        torch.zeros((num_tokens, heads, latent + rope), dtype=torch.bfloat16),
+        torch.empty((num_tokens, heads, latent), dtype=torch.bfloat16),
+        params,
+        torch.float8_e4m3fn,
+    )
+    assert len(calls) == 1
+    torch.testing.assert_close(calls[0][4], blocks[0, 2:, 0].T.to(torch.int32))
+    torch.testing.assert_close(calls[0][5], lengths[2:])
+    assert calls[0][9] == seq_len_q
+    assert calls[0][14] is params.fwd.softmax_stats_tensor
+    # A verify group passes the bounds of exactly this step's generation tokens.
+    if seq_len_q > 1:
+        torch.testing.assert_close(calls[0][15], kv_bounds[:num_tokens])
+    else:
+        assert calls[0][15] is None
+
+
 # kv_cache_block_offsets geometry with an odd page count per sequence, e.g. a
 # max_seq_len of 8320 over 128-token pages: each [num_seqs, 2, max_blocks] row
 # is 8 * 65 bytes, so an odd row offset leaves an 8-byte aligned slice.