[None][feat] Rubin fp4 mla refactor(P1)

tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md

@@ -258,6 +258,7 @@ The core contract is:
   - `support_fused_rope()`
   - `support_fused_qkv()`
   - `support_mla()`
+  - `support_fp4_kv_cache()`
 - `runtime_workspace_bytes_per_token(model_config, mapping)` — the memory-accounting
   contract (default `0`); see below
 - `runtime_workspace_is_chunked_prefill_bounded(model_config)` — whether
@@ -399,8 +400,8 @@ independently with `is_supported(..., phase=...)`; a phased library accepts only
 phases backed by its corresponding `run_*()` entry point.
 
 `Fmha` owns both entry points. Libraries declare shared capabilities through
-class attributes, such as `supports_skip_correction` and
-`supports_block_sparse_inputs`, and override only
+class attributes, such as `supports_skip_correction`, `supports_block_sparse_inputs`,
+and `supports_fp4_mla`, and override only
 `_is_available()` and `_is_supported()` for implementation-specific checks.
 `is_available()` rejects unsupported static capabilities before calling
 `_is_available()`. `is_supported()` provides the same boundary for shared
@@ -427,6 +428,9 @@ The FMHA package is split by role:
   MLA uses `query_input` with `is_fused_qkv=False`.
 - `fmha/combined.py` composes different context and generation implementations
   for non-MLA mixed batches.
+- `fmha/fp4_mla.py` implements FP4 MLA using FP8 context attention with FP4
+  cache updates and FP4 no-dequant decode. It uses KV Cache Manager V2;
+  batch state, cache storage, and kernels live in `fp4_mla/`.
 - `fmha/triton_custom_mask.py` implements the Triton custom-mask context phase.
   Custom-mask data applies to context requests; for mixed batches,
   `TrtllmAttention` can pair it with a later causal-generation provider through

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

@@ -18,6 +18,7 @@
 from .fallback import FallbackFmha
 from .flashinfer_sparse_mla import FlashInferSparseMlaFmha
 from .flashinfer_trtllm_gen import FlashInferTrtllmGenFmha
+from .fp4_mla import Fp4MlaFmha
 from .interface import Fmha, FmhaPhase
 from .msa_decode import MsaDecodeFmha
 from .msa_prefill import MsaPrefillFmha
@@ -34,6 +35,7 @@
     "FallbackFmha",
     "FlashInferSparseMlaFmha",
     "FlashInferTrtllmGenFmha",
+    "Fp4MlaFmha",
     "Fmha",
     "FmhaCls",
     "FmhaParams",

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

@@ -14,7 +14,7 @@
 # limitations under the License.
 
 from dataclasses import replace
-from typing import TYPE_CHECKING, Optional, Tuple
+from typing import TYPE_CHECKING, Optional
 
 import torch
 
@@ -23,7 +23,6 @@
     scatter_fp4_mla_kv_cache,
 )
 from tensorrt_llm._torch.attention.backends.fp4_mla.fp4_mla_context import (
-    _FP8_CONTEXT_ATTN_ATTR,
     _FP8_CONTEXT_SCRATCH_ATTR,
     _build_fp8_mla_context_attn,
     _build_fp8_mla_context_metadata,
@@ -39,7 +38,7 @@
 )
 from tensorrt_llm.bindings import DataType
 
-from .fallback import FallbackFmha
+from .interface import FmhaPhase
 from .phased import FmhaParams, PhasedFmha
 
 if TYPE_CHECKING:
@@ -52,49 +51,37 @@
 class Fp4MlaFmha(PhasedFmha):
     """TRTLLM FMHA library for FP4 MLA context and no-dequant decode."""
 
+    supports_fp4_mla = True
+
+    def __init__(self, attn: "TrtllmAttention") -> None:
+        super().__init__(attn)
+        self._fp8_attention: Optional["TrtllmAttention"] = None
+
     @classmethod
     def _is_available(cls, attn: "TrtllmAttention") -> bool:
-        return attn.uses_fp4_mla_attention
+        return bool(getattr(attn, "uses_fp4_mla_attention", False))
 
-    def forward(
+    def _is_supported(
         self,
         q: torch.Tensor,
         k: Optional[torch.Tensor],
         v: Optional[torch.Tensor],
         metadata: "TrtllmAttentionMetadata",
         forward_args: AttentionForwardArgs,
-    ) -> None:
-        self._validate_request(k, v, metadata, forward_args)
-        super().forward(q, k, v, metadata, forward_args)
-
-    def _validate_request(
-        self,
-        k: Optional[torch.Tensor],
-        v: Optional[torch.Tensor],
-        metadata: "TrtllmAttentionMetadata",
-        forward_args: AttentionForwardArgs,
-    ) -> None:
+        *,
+        phase: Optional[FmhaPhase] = None,
+    ) -> bool:
+        # Mask/output formats and the attention phase are represented in the cache key.
+        # K/V presence is fixed per phase; optional features and pools are instance invariants.
+        del q, phase
         if forward_args.output_sf is not None:
             raise NotImplementedError("FP4 MLA does not support quantized attention output.")
         if forward_args.attention_mask != PredefinedAttentionMask.CAUSAL:
             raise NotImplementedError("FP4 MLA requires a causal attention mask.")
         if forward_args.attention_mask_data is not None:
             raise NotImplementedError("FP4 MLA does not support custom attention masks.")
-        if forward_args.attention_sinks is not None:
-            raise NotImplementedError("FP4 MLA does not support attention sinks.")
-
-        sparse_runtime_params = forward_args.sparse_runtime_params
-        if (
-            (
-                sparse_runtime_params.sparse_kv_indices is not None
-                and sparse_runtime_params.sparse_kv_indices.numel() > 0
-            )
-            or (
-                sparse_runtime_params.sparse_attn_indices is not None
-                and sparse_runtime_params.sparse_attn_indices.numel() > 0
-            )
-            or metadata.num_sparse_topk > 0
-        ):
+        sparse_params = self.attn.sparse_params
+        if getattr(sparse_params, "uses_spcompress", False):
             raise NotImplementedError("FP4 MLA does not support sparse attention.")
 
         kv_cache_manager = metadata.kv_cache_manager
@@ -104,39 +91,89 @@ def _validate_request(
             raise RuntimeError("FP4 MLA requires NVFP4 KV cache storage.")
         if kv_cache_manager.kv_factor != 1:
             raise RuntimeError("FP4 MLA requires a SELF-K-only KV cache.")
-        if metadata.high_precision_kv_pool is None:
-            raise RuntimeError("FP4 MLA requires the high-precision KV pool.")
-        if metadata.fp4_mla_v_scale_pool is None:
-            raise RuntimeError("FP4 MLA requires the V-scale pool.")
         if metadata.beam_width != 1:
             raise NotImplementedError("FP4 MLA does not support beam search.")
+        if forward_args.attention_sinks is not None:
+            raise NotImplementedError("FP4 MLA does not support attention sinks.")
+        sparse_inputs = forward_args.sparse_runtime_params
+        if (
+            metadata.num_sparse_topk > 0
+            or sparse_inputs.block_sparse_inputs is not None
+            or (
+                sparse_inputs.sparse_kv_indices is not None
+                and sparse_inputs.sparse_kv_indices.numel() > 0
+            )
+            or (
+                sparse_inputs.sparse_attn_indices is not None
+                and sparse_inputs.sparse_attn_indices.numel() > 0
+            )
+        ):
+            raise NotImplementedError("FP4 MLA does not support sparse attention.")
+        state = metadata.fp4_mla_state
+        if state is None or state.hp_pool is None:
+            raise RuntimeError("FP4 MLA requires prepared high-precision KV state.")
+        if state.v_scale_pool is None:
+            raise RuntimeError("FP4 MLA requires the V-scale pool.")
 
         attention_input_type = forward_args.attention_input_type
         if attention_input_type == AttentionInputType.context_only:
             if k is None or v is None:
                 raise RuntimeError("FP4 MLA context requires expanded K and V tensors.")
-            return
-        if attention_input_type == AttentionInputType.generation_only:
+        elif attention_input_type == AttentionInputType.generation_only:
             if k is not None or v is not None:
                 raise RuntimeError("FP4 MLA generation expects a fused query input.")
-            return
-        raise NotImplementedError("FP4 MLA requires a context-only or generation-only call.")
+        else:
+            raise NotImplementedError("FP4 MLA requires a context-only or generation-only call.")
+        return True
+
+    def _get_fp8_context_resources(
+        self,
+        metadata: "TrtllmAttentionMetadata",
+        q: torch.Tensor,
+    ) -> tuple["TrtllmAttention", _Fp8MlaContextScratch]:
+        """Reuse one manager-owned scratch and one layer-local FP8 attention view."""
+        attn = self.attn
+        kv_cache_manager = metadata.kv_cache_manager
+        if kv_cache_manager is None:
+            raise RuntimeError("FP8 MLA context scratch requires a KV cache manager.")
+        scratch = getattr(kv_cache_manager, _FP8_CONTEXT_SCRATCH_ATTR, None)
+        scratch_head_dim = (attn.kv_lora_rank or 0) + (attn.qk_rope_head_dim or 0)
+        if scratch is None:
+            scratch = _Fp8MlaContextScratch.create(
+                metadata,
+                device=q.device,
+                head_dim=scratch_head_dim,
+            )
+            setattr(kv_cache_manager, _FP8_CONTEXT_SCRATCH_ATTR, scratch)
+        else:
+            assert isinstance(scratch, _Fp8MlaContextScratch)
+            assert scratch.matches(metadata, device=q.device, head_dim=scratch_head_dim), (
+                "FP8 MLA context scratch geometry must be shared by all layers of its KV manager."
+            )
+
+        fp8_attention = self._fp8_attention
+        if fp8_attention is None:
+            fp8_attention = _build_fp8_mla_context_attn(attn)
+            self._fp8_attention = fp8_attention
+        fp8_attention.rotary_inv_freq = attn.rotary_inv_freq
+        fp8_attention.rotary_cos_sin = attn.rotary_cos_sin
+        return fp8_attention, scratch
 
     def run_mla_context(self, params: FmhaParams) -> None:
         attn = params.attn
         metadata = params.meta
         forward_args = params.fwd
-        q = params.qkv_input
-        k = params.k_input
-        v = params.v_input
-        output = params.context_buf
+        q = params.query_input
+        k = params.key_input
+        v = params.value_input
+        output = params.output
         if q is None or k is None or v is None:
             raise RuntimeError("FP4 MLA context requires expanded Q, K, and V tensors.")
         if output is None:
             raise RuntimeError("FP4 MLA context requires an output buffer.")
         if forward_args.latent_cache is None:
             raise RuntimeError("FP4 MLA context requires latent_cache.")
-        if metadata.positions is None:
+        if metadata.fp4_mla_state.positions is None:
             raise RuntimeError("FP4 MLA context requires token positions.")
         if metadata.num_contexts <= 0:
             raise RuntimeError("FP4 MLA context requires context requests.")
@@ -148,18 +185,12 @@ def run_mla_context(self, params: FmhaParams) -> None:
             raise RuntimeError("FP4 MLA context Q/K/V token counts do not match.")
 
         require_fp4_mla_fp8_context_support()
-        if metadata.is_cuda_graph:
-            raise NotImplementedError(
-                "FP4 MLA context does not support CUDA graphs with TRT-LLM FP8 FMHA."
-            )
 
         num_tokens = q.shape[0]
         output = output.view(num_tokens, -1)
         local_layer = attn.get_fp4_mla_local_layer_idx(metadata)
         kv_lora_rank = attn.kv_lora_rank or 0
-        qk_rope_head_dim = attn.qk_rope_head_dim or 0
 
-        attn._ensure_rope_table_size(metadata.max_seq_len)
         latent_cache = forward_args.latent_cache[:num_tokens]
 
         def update_fp4_cache() -> None:
@@ -176,30 +207,7 @@ def update_fp4_cache() -> None:
             if not hp_pool_updated:
                 raise RuntimeError("Fused FP4 MLA context scatter did not update the HP pool.")
 
-        kv_cache_manager = metadata.kv_cache_manager
-        if kv_cache_manager is None:
-            raise RuntimeError("FP8 MLA context scratch requires a KV cache manager.")
-        scratch = getattr(kv_cache_manager, _FP8_CONTEXT_SCRATCH_ATTR, None)
-        scratch_head_dim = kv_lora_rank + qk_rope_head_dim
-        if not isinstance(scratch, _Fp8MlaContextScratch) or not scratch.matches(
-            metadata,
-            device=q.device,
-            head_dim=scratch_head_dim,
-        ):
-            scratch = _Fp8MlaContextScratch.create(
-                metadata,
-                device=q.device,
-                head_dim=scratch_head_dim,
-            )
-            setattr(kv_cache_manager, _FP8_CONTEXT_SCRATCH_ATTR, scratch)
-
-        fp8_attention = getattr(attn, _FP8_CONTEXT_ATTN_ATTR, None)
-        if fp8_attention is None:
-            fp8_attention = _build_fp8_mla_context_attn(attn)
-            fp8_attention.fmha_libs = [FallbackFmha(fp8_attention)]
-            setattr(attn, _FP8_CONTEXT_ATTN_ATTR, fp8_attention)
-        fp8_attention.rotary_inv_freq = attn.rotary_inv_freq
-        fp8_attention.rotary_cos_sin = attn.rotary_cos_sin
+        fp8_attention, scratch = self._get_fp8_context_resources(metadata, q)
         fp8_metadata = _get_fp8_mla_context_metadata(metadata, scratch)
         fp8_forward_args = replace(
             forward_args,
@@ -230,7 +238,7 @@ def forward_context_partition(
         *,
         kv_lens_cuda: torch.Tensor,
         kv_lens_cpu: torch.Tensor,
-    ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
         """Run one explicit-KV partition through the FP8 MLA context kernel.
 
         Returns whatever ``TrtllmAttention.forward`` returns for the FP8
@@ -281,33 +289,7 @@ def forward_context_partition(
             raise NotImplementedError("FP4 MLA chunked prefill does not support sparse attention.")
 
         require_fp4_mla_fp8_context_support()
-        kv_cache_manager = metadata.kv_cache_manager
-        if kv_cache_manager is None:
-            raise RuntimeError("FP8 MLA context scratch requires a KV cache manager.")
-        kv_lora_rank = self.attn.kv_lora_rank or 0
-        qk_rope_head_dim = self.attn.qk_rope_head_dim or 0
-        scratch_head_dim = kv_lora_rank + qk_rope_head_dim
-        scratch = getattr(kv_cache_manager, _FP8_CONTEXT_SCRATCH_ATTR, None)
-        if not isinstance(scratch, _Fp8MlaContextScratch) or not scratch.matches(
-            metadata,
-            device=q.device,
-            head_dim=scratch_head_dim,
-        ):
-            scratch = _Fp8MlaContextScratch.create(
-                metadata,
-                device=q.device,
-                head_dim=scratch_head_dim,
-            )
-            setattr(kv_cache_manager, _FP8_CONTEXT_SCRATCH_ATTR, scratch)
-
-        attn = self.attn
-        fp8_attention = getattr(attn, _FP8_CONTEXT_ATTN_ATTR, None)
-        if fp8_attention is None:
-            fp8_attention = _build_fp8_mla_context_attn(attn)
-            fp8_attention.fmha_libs = [FallbackFmha(fp8_attention)]
-            setattr(attn, _FP8_CONTEXT_ATTN_ATTR, fp8_attention)
-        fp8_attention.rotary_inv_freq = attn.rotary_inv_freq
-        fp8_attention.rotary_cos_sin = attn.rotary_cos_sin
+        fp8_attention, scratch = self._get_fp8_context_resources(metadata, q)
 
         expected_kv_tokens = int(kv_lens_cpu[: metadata.num_contexts].sum().item())
         if k.shape[0] != expected_kv_tokens:
@@ -339,8 +321,8 @@ def run_mla_generation(self, params: FmhaParams) -> None:
         attn = params.attn
         metadata = params.meta
         forward_args = params.fwd
-        q = params.qkv_input
-        output = params.context_buf
+        q = params.query_input
+        output = params.output
         if q is None:
             raise RuntimeError("FP4 MLA generation requires a fused query input.")
         if output is None:
@@ -355,18 +337,18 @@ def run_mla_generation(self, params: FmhaParams) -> None:
         qk_rope_head_dim = attn.qk_rope_head_dim or 0
         fused_head_dim = kv_lora_rank + qk_rope_head_dim
 
-        if not bool(getattr(metadata, "_fp4_mla_generation_cache_scattered", False)):
+        if not bool(getattr(metadata.fp4_mla_state, "generation_cache_scattered", False)):
             raise RuntimeError(
                 "FP4 MLA generation requires fused RoPE/Q quantization and "
                 "cache/HP-pool update before attention."
             )
-        metadata._fp4_mla_generation_cache_scattered = False
+        metadata.fp4_mla_state.generation_cache_scattered = False
         query = q.view(q.shape[0], attn.num_heads, fused_head_dim)
         output_view = output.view(q.shape[0], attn.num_heads, kv_lora_rank)
         sm_scale = 1.0 / (attn.q_scaling * ((attn.qk_nope_head_dim or 0) + qk_rope_head_dim) ** 0.5)
-        prequantized_q = getattr(metadata, "_fp4_mla_prequantized_q", None)
-        prequantized_q_sf = getattr(metadata, "_fp4_mla_prequantized_q_sf", None)
-        q_batch_capacity = getattr(metadata, "_fp4_mla_q_batch_capacity", None)
+        prequantized_q = getattr(metadata.fp4_mla_state, "prequantized_q", None)
+        prequantized_q_sf = getattr(metadata.fp4_mla_state, "prequantized_q_sf", None)
+        q_batch_capacity = getattr(metadata.fp4_mla_state, "q_batch_capacity", None)
         try:
             run_fp4_mla_attention_decode(
                 metadata,
@@ -383,9 +365,9 @@ def run_mla_generation(self, params: FmhaParams) -> None:
                 softmax_stats_tensor=params.fwd.softmax_stats_tensor,
             )
         finally:
-            metadata._fp4_mla_prequantized_q = None
-            metadata._fp4_mla_prequantized_q_sf = None
-            metadata._fp4_mla_q_batch_capacity = None
+            metadata.fp4_mla_state.prequantized_q = None
+            metadata.fp4_mla_state.prequantized_q_sf = None
+            metadata.fp4_mla_state.q_batch_capacity = None
 
 
 __all__ = ["Fp4MlaFmha"]

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

@@ -43,6 +43,7 @@ class Fmha(ABC):
     supports_skip_correction: ClassVar[bool] = False
     supports_block_sparse_inputs: ClassVar[bool] = False
     supports_workspace_reclamation: bool = False
+    supports_fp4_mla: ClassVar[bool] = False
 
     def __init__(self, attn: "TrtllmAttention"):
         self._attn_ref: weakref.ReferenceType["TrtllmAttention"] = weakref.ref(attn)
@@ -76,6 +77,8 @@ def is_available(cls, attn: "TrtllmAttention") -> bool:
                 f"{cls.__name__} is unavailable: skip-correction is enabled and unsupported."
             )
             return False
+        if getattr(attn, "uses_fp4_mla_attention", False) and not cls.supports_fp4_mla:
+            return False
         return cls._is_available(attn)
 
     @classmethod