[None][fix] Guard the CuTe DSL MLA decode batch and add a regression test

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

@@ -466,7 +466,7 @@ def _run_mla_decode(
         local_layer_idx = attn.get_local_layer_idx(meta)
         pool_idx = int(pool_mapping[local_layer_idx, 0])
         page_table_layer = block_offsets[pool_idx, :, 0, :]
-        cache_seqs_base = params.sequence_lengths.to(torch.int32)
+        cache_seqs_base = params.sequence_lengths[:batch_size].to(torch.int32)
         # The page-table batch dimension determines the kernel launch batch.
         page_table = (
             page_table_layer[meta.num_contexts : meta.num_contexts + batch_size]

tensorrt_llm/_torch/custom_ops/cute_dsl_custom_ops.py

@@ -12366,6 +12366,17 @@ def forward(
 
             # workspace = lse + split_kv_workspace
             batch_size = cache_seqs.shape[0]
+            # The kernel launches one request per page-table column, so a
+            # wider table would read q, o and cache_seqs past their ends.
+            if (page_table.shape[1] != batch_size
+                    or q_latent.shape[3] != batch_size
+                    or o.shape[3] != batch_size):
+                raise RuntimeError(
+                    "CuteDSLNVMlaDecodeBlackwellRunner requires page_table "
+                    "(max_blocks, B), q_latent and o (..., B) to agree with "
+                    f"cache_seqs (B,), got B={batch_size}, page_table="
+                    f"{tuple(page_table.shape)}, q_latent="
+                    f"{tuple(q_latent.shape)}, o={tuple(o.shape)}.")
             d_latent = q_latent.shape[1]
             softmax_stats_kernel = None
             if softmax_stats is not None:

tests/unittest/_torch/attention/test_cute_dsl_mla_helix.py

@@ -146,3 +146,118 @@ def test_cute_dsl_mla_helix_stats_and_empty_local_kv(
     actual_log_partition = softmax_stats[1, :, 0] + softmax_stats[1, :, 1].log()
     expected_log_partition = torch.logsumexp(scores, dim=-1)
     torch.testing.assert_close(actual_log_partition, expected_log_partition, rtol=1e-5, atol=1e-5)
+
+
+class _FakeKvCacheManager:
+    def __init__(self, kv_pool: torch.Tensor):
+        self._kv_pool = kv_pool
+
+    def get_buffers(self, layer_idx: int) -> torch.Tensor:
+        return self._kv_pool
+
+
+@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float8_e4m3fn])
+@pytest.mark.parametrize("batch_size", [1, 3])
+def test_cute_dsl_mla_decode_batch_below_max_num_sequences(
+    input_dtype: torch.dtype, batch_size: int
+) -> None:
+    """A decode batch smaller than the metadata's sequence capacity.
+
+    kv_cache_block_offsets holds one row per sequence slot (max_num_sequences),
+    while the kernel derives its batch size from the page table it receives.
+    The FMHA wrapper must hand it only this batch's generation rows; any extra
+    column launches CTAs for requests past the end of q, the output and the
+    sequence lengths. Runs the K3 helix verify-group shape (H=96, 1 + 6 draft
+    tokens, per-token kv_bounds) through CuteDslMlaFmha._run_mla_decode.
+    """
+    from types import SimpleNamespace
+
+    from tensorrt_llm._torch.attention.backends.fmha.cute_dsl_mla import CuteDslMlaFmha
+
+    torch.manual_seed(23)
+    device = torch.device("cuda")
+    max_num_sequences, seq_len_q, num_heads = 32, 7, 96
+    latent_dim, rope_dim, qk_nope_head_dim, page_size = 512, 64, 128, 64
+    max_blocks = 8
+    num_pages = max_num_sequences * max_blocks
+    head_dim = latent_dim + rope_dim
+    kv_lens = [200, 70, 135][:batch_size]
+    softmax_scale = 1.0 / math.sqrt(qk_nope_head_dim + rope_dim)
+
+    kv_pool = (
+        torch.randn(num_pages, 1, page_size, 1, head_dim, device=device, dtype=torch.bfloat16) * 0.1
+    ).to(input_dtype)
+    # One disjoint page range per sequence slot. Slots past the batch are
+    # allocated but must never reach the kernel.
+    block_offsets = (
+        torch.arange(num_pages, dtype=torch.int32, device=device)
+        .view(1, max_num_sequences, 1, max_blocks)
+        .expand(1, max_num_sequences, 2, max_blocks)
+        .contiguous()
+    )
+    kv_bounds = torch.cat(
+        [
+            torch.arange(kv_len - seq_len_q + 1, kv_len + 1, dtype=torch.int32, device=device)
+            for kv_len in kv_lens
+        ]
+    )
+    meta = SimpleNamespace(
+        kv_cache_manager=_FakeKvCacheManager(kv_pool),
+        kv_cache_block_offsets=block_offsets,
+        host_kv_cache_pool_mapping=torch.zeros(1, 2, dtype=torch.int32),
+        num_contexts=0,
+        max_num_requests=max_num_sequences,
+        helix_position_offsets=torch.zeros(
+            batch_size * seq_len_q, dtype=torch.int32, device=device
+        ),
+        _helix_spec_tokens_valid=True,
+        helix_kv_bounds=kv_bounds,
+    )
+    attn = SimpleNamespace(
+        kv_lora_rank=latent_dim,
+        qk_rope_head_dim=rope_dim,
+        qk_nope_head_dim=qk_nope_head_dim,
+        num_heads=num_heads,
+        q_scaling=1.0,
+        layer_idx=0,
+        get_local_layer_idx=lambda _meta: 0,
+    )
+    num_tokens = batch_size * seq_len_q
+    q = torch.randn(num_tokens, num_heads * head_dim, device=device, dtype=torch.bfloat16) * 0.1
+    output = torch.empty(num_tokens, num_heads * latent_dim, device=device, dtype=torch.bfloat16)
+    softmax_stats = torch.empty(num_tokens, num_heads, 2, device=device, dtype=torch.float32)
+    # As in production: one sequence length per generation request.
+    sequence_lengths = torch.tensor(kv_lens, dtype=torch.int32, device=device)
+
+    fmha = object.__new__(CuteDslMlaFmha)
+    if input_dtype == torch.float8_e4m3fn:
+        fmha._cute_dsl_fp8_scale = (softmax_scale, 1.0)
+    workspace = torch.empty(
+        fmha._required_workspace_size(num_heads, seq_len_q, latent_dim, max_num_sequences),
+        device=device,
+        dtype=torch.uint8,
+    )
+    params = SimpleNamespace(
+        attn=attn,
+        meta=meta,
+        num_requests=batch_size,
+        tokens_per_block=page_size,
+        sequence_lengths=sequence_lengths,
+        workspace=workspace,
+        fwd=SimpleNamespace(quant_q_buffer=None, softmax_stats_tensor=softmax_stats),
+    )
+
+    fmha._run_mla_decode(q, output, params, input_dtype)
+    torch.cuda.synchronize()
+
+    q_rows = q.to(input_dtype).float().view(batch_size, seq_len_q, num_heads, head_dim)
+    out_rows = output.float().view(batch_size, seq_len_q, num_heads, latent_dim)
+    tol = 5e-2 if input_dtype == torch.float8_e4m3fn else 3e-2
+    for b, kv_len in enumerate(kv_lens):
+        pages = block_offsets[0, b, 0, : -(-kv_len // page_size)].long()
+        keys = kv_pool[pages, 0, :, 0].float().reshape(-1, head_dim)[:kv_len]
+        for t in range(seq_len_q):
+            bound = int(kv_bounds[b * seq_len_q + t])
+            scores = q_rows[b, t] @ keys[:bound].transpose(0, 1) * softmax_scale
+            expected = torch.softmax(scores, dim=-1) @ keys[:bound, :latent_dim]
+            torch.testing.assert_close(out_rows[b, t], expected, rtol=tol, atol=tol)