[TRTLLM-13580][test] Add model-derived PyTorch attention backend test suite

tensorrt_llm/_torch/attention_backend/flashinfer.py

@@ -5,7 +5,7 @@
 import weakref
 from dataclasses import dataclass, field
 from itertools import chain
-from typing import Any, Dict, Literal, NewType, Optional, TypeAlias, cast
+from typing import Any, Dict, NewType, Optional, TypeAlias, cast
 
 if sys.version_info[:2] >= (3, 12):
     from typing import override
@@ -127,11 +127,6 @@ class FlashInferWrappers:
 class FlashInferAttentionMetadata(AttentionMetadata):
     workspace_buffer: Optional[torch.Tensor] = None
 
-    # cache concat/split kernels when using PD disaggregation
-    # expects KV cache in [max_num_pages, 2, num_kv_heads, page_size, head_dim] layout,
-    # so set kv_layout as "HND" here
-    kv_layout: Literal["NHD", "HND"] = "HND"
-
     paged_kv_indptr_decode: torch.Tensor = field(init=False)
     paged_kv_indptr_prefill: torch.Tensor = field(init=False)
     _paged_kv_indices: torch.Tensor = field(init=False, repr=False)
@@ -1390,33 +1385,6 @@ def update_quant_config(self, new_quant_config: Optional[QuantConfig]):
             self.has_fp8_kv_cache = self.quant_config.layer_quant_mode.has_fp8_kv_cache(
             )
 
-    def mla_rope_generation(
-        self,
-        fused_q: torch.Tensor,
-        q_pe: torch.Tensor,
-        latent_cache: torch.Tensor,
-        metadata,
-        cu_q_seqlens: torch.Tensor,
-        cu_kv_seqlens: torch.Tensor,
-        fmha_scheduler_counter: torch.Tensor,
-        mla_bmm1_scale,
-        mla_bmm2_scale,
-        quant_q_buffer,
-        out_scale=None,
-    ) -> None:
-        """Stub for MLA generation rope step used when FlashInfer is the mqa backend.
-
-        FlashInferAttention does not fuse RoPE (support_fused_rope returns False),
-        so RoPE is applied externally in MLA.forward_impl before this point.
-        q_pe already has RoPE applied; we just copy it into the rope slot of
-        fused_q so that forward_absorption_generation can pass fused_q directly
-        to _mla_forward_generation.  The latent_cache KV-cache append is handled
-        inside _mla_forward_generation when forward() is called.
-        """
-        # fused_q shape: [num_tokens, num_heads, kv_lora_rank + qk_rope_head_dim]
-        # q_pe shape:    [num_tokens, num_heads, qk_rope_head_dim]
-        fused_q[..., self.kv_lora_rank:] = q_pe
-
     def _get_mla_caches(
         self,
         metadata: "FlashInferAttentionMetadata",
@@ -1929,8 +1897,8 @@ def forward(self,
             else:
                 output = torch.empty_like(q)
 
-        # FlashInfer's sliding window attention is inclusive, while the attention window size defined in TRTLLM is exclusive.
-        # So we need to subtract 1 from the attention window size for a consistent behavior.
+        # FlashInfer's sliding window attention is inclusive, while the attention window size defined
+        # in TRTLLM is exclusive. Subtract 1 from the attention window size for consistent behavior.
         attention_window_size = forward_args.attention_window_size
         if attention_window_size is not None:
             attention_window_size = attention_window_size - 1

tensorrt_llm/_torch/attention_backend/interface.py

@@ -3,8 +3,8 @@
 from collections import namedtuple
 from dataclasses import dataclass, field
 from enum import Enum, IntEnum
-from typing import (TYPE_CHECKING, Any, Dict, Generic, List, Optional, Protocol,
-                    Tuple, Type, TypeVar, Union)
+from typing import (TYPE_CHECKING, Any, Dict, Generic, List, Literal, Optional,
+                    Protocol, Tuple, Type, TypeVar, Union)
 
 import torch
 from typing_extensions import Self
@@ -42,8 +42,10 @@ class AttentionRuntimeFeatures:
     chunked_prefill: bool = False
     cache_reuse: bool = False
     has_speculative_draft_tokens: bool = False
-    chunk_size: int = 0  # this is the chunk size for MLA chunked prefill, it will split kv cache into chunks to save global memory.
-    chunked_prefill_buffer_batch_size: int = 4  # real chunk size for MLA chunked prefill is chunked_prefill_buffer_batch_size * chunk_size.
+    # This is the chunk size for MLA chunked prefill, which splits KV cache into chunks.
+    chunk_size: int = 0
+    # The real chunk size for MLA chunked prefill is this value * chunk_size.
+    chunked_prefill_buffer_batch_size: int = 4
 
 
 # The type of requests in qkv passed to attention
@@ -73,6 +75,10 @@ class AttentionMetadata:
     mapping: Optional[Mapping] = None
     # Sparse settings for metadata allocation/update; dense metadata leaves it None.
     sparse_metadata_params: Optional[SparseMetadataParams] = None
+    # Paged KV-cache block layout:
+    # NHD: [max_num_pages, 2, page_size, num_kv_heads, head_dim]
+    # HND: [max_num_pages, 2, num_kv_heads, page_size, head_dim]
+    kv_layout: Literal["NHD", "HND"] = "HND"
 
     enable_flash_mla: bool = False
     enable_context_mla_with_cached_kv: bool = False
@@ -853,6 +859,10 @@ class AttentionForwardArgs:
     cu_q_seqlens: Optional[torch.Tensor] = None
     cu_kv_seqlens: Optional[torch.Tensor] = None
     fmha_scheduler_counter: Optional[torch.Tensor] = None
+    # Testing only: skip the RoPE step of MLA generation (the standalone harness
+    # feeds a pre-RoPE'd fused_q). The TRTLLM backend then appends the new latent
+    # and inits the trtllm-gen scheduler buffers itself.
+    skip_mla_rope_generation: bool = False
 
     mla_bmm1_scale: Optional[torch.Tensor] = None
     mla_bmm2_scale: Optional[torch.Tensor] = None

tensorrt_llm/_torch/attention_backend/trtllm.py

@@ -1473,6 +1473,49 @@ def forward(
         if forward_args.cu_kv_seqlens is None:
             forward_args.cu_kv_seqlens = metadata.cu_kv_seqlens
 
+        # Testing only: ``mla_rope_generation`` normally rotates q_pe, appends the
+        # new latent to the paged cache, and fills the trtllm-gen scheduler
+        # buffers (cumulative q/kv seqlens + the FMHA scheduler counter). When the
+        # harness sets ``skip_mla_rope_generation`` it feeds a pre-RoPE'd fused_q,
+        # so we skip only the RoPE and do the append + scheduler init here: the
+        # generation FMHA only reads the cache, and the fallback path needs the
+        # scheduler buffers (the flashinfer trtllm-gen decode kernel ignores them).
+        if (self.is_mla_enable and forward_args.skip_mla_rope_generation
+                and forward_args.attention_input_type
+                == AttentionInputType.generation_only):
+            num_ctx = metadata.num_contexts
+            n_gen = metadata.num_generations
+            # Use the GPU-resident length tensors (no host->device copy) so this
+            # stays CUDA-graph-capturable.
+            gen_q_lens = metadata.seq_lens_cuda[num_ctx:num_ctx + n_gen].to(
+                torch.int32)
+            gen_kv_lens = metadata.kv_lens_cuda_runtime[num_ctx:num_ctx +
+                                                        n_gen].to(torch.int32)
+            cu_q = torch.zeros(n_gen + 1, dtype=torch.int32, device=q.device)
+            cu_kv = torch.zeros(n_gen + 1, dtype=torch.int32, device=q.device)
+            cu_q[1:] = torch.cumsum(gen_q_lens, dim=0).to(
+                torch.int32) * self.num_heads
+            cu_kv[1:] = torch.cumsum(gen_kv_lens, dim=0).to(torch.int32)
+            forward_args.cu_q_seqlens = cu_q
+            forward_args.cu_kv_seqlens = cu_kv
+            if forward_args.fmha_scheduler_counter is None:
+                forward_args.fmha_scheduler_counter = torch.zeros(
+                    1, dtype=torch.uint32, device=q.device)
+            else:
+                forward_args.fmha_scheduler_counter.zero_()
+            assert forward_args.latent_cache is not None
+            from .utils import append_mla_latent_cache
+            append_mla_latent_cache(
+                metadata.kv_cache_manager,
+                self.get_local_layer_idx(metadata),
+                metadata.request_ids,
+                metadata.seq_lens.tolist(),
+                metadata.kv_cache_params.num_cached_tokens_per_seq,
+                forward_args.latent_cache,
+                kv_layout=metadata.kv_layout,
+                seq_start=num_ctx,
+            )
+
         # RocketKV and DSA predict which blocks to keep, so build their sparse
         # index tensors here. Skip-softmax needs no prediction.
         sparse_params = self.sparse_params

tensorrt_llm/_torch/attention_backend/utils.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Optional, Type, Union
+from typing import TYPE_CHECKING, Optional, Sequence, Type, Union
 
 import torch
 
@@ -125,3 +125,55 @@ def create_attention(
         num_kv_heads,
         **kwargs,
     )
+
+
+def append_mla_latent_cache(
+    kv_cache_manager,
+    layer_idx: int,
+    request_ids: Sequence[int],
+    seq_lens: Sequence[int],
+    num_cached_tokens: Sequence[int],
+    latent_cache: torch.Tensor,
+    *,
+    kv_layout: str = "NHD",
+    seq_start: int = 0,
+) -> torch.Tensor:
+    """Append packed MLA latent tokens into a paged KV cache.
+
+    The MLA cache has one latent head with ``kv_factor=1`` and stores
+    ``[compressed_kv | k_pe]``. ``latent_cache`` is packed by request for the
+    sequences in ``seq_lens[seq_start:]``.
+    """
+    kv_cache = kv_cache_manager.get_buffers(layer_idx, kv_layout=kv_layout)
+    if kv_layout == "NHD":
+        tokens_per_block = kv_cache.shape[2]
+    elif kv_layout == "HND":
+        tokens_per_block = kv_cache.shape[3]
+    else:
+        raise ValueError(f"Unsupported kv_layout: {kv_layout}")
+
+    blocks_per_seq = kv_cache_manager.get_batch_cache_indices(
+        list(request_ids), layer_idx)
+
+    offset = 0
+    for i in range(seq_start, len(seq_lens)):
+        q_len = int(seq_lens[i])
+        new = latent_cache[offset:offset + q_len].to(kv_cache.dtype)
+        start = int(num_cached_tokens[i])
+        blocks = [b for b in blocks_per_seq[i] if b != -1]
+        written = 0
+        while written < q_len:
+            pos = start + written
+            block = blocks[pos // tokens_per_block]
+            block_offset = pos % tokens_per_block
+            n = min(tokens_per_block - block_offset, q_len - written)
+            if kv_layout == "NHD":
+                kv_cache[block, 0, block_offset:block_offset + n,
+                         0, :].copy_(new[written:written + n])
+            else:
+                kv_cache[block, 0, 0, block_offset:block_offset + n, :].copy_(
+                    new[written:written + n])
+            written += n
+        offset += q_len
+
+    return kv_cache