[None][feat] support FP8 KV cache in PrimTS MLA decode

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

@@ -79,6 +79,9 @@ def __init__(self, attn: "TrtllmAttention") -> None:
         self._context_wrappers: dict[int, "BatchPrefillPagedTSWrapper"] = {}
         self._decode_wrappers: dict[int, "BatchDecodePagedTSWrapper"] = {}
         self._mla_decode_wrappers: dict[int, "BatchMLADecodePagedTSWrapper"] = {}
+        # Dense MLA's quantization scales are fixed for this layer/model. PrimTS
+        # takes host scalars, so read them once during eager warmup, not capture.
+        self._mla_fp8_scales: tuple[float, float] | None = None
         # Decode plans retain views into the shared workspace and are invalidated
         # whenever its underlying allocation changes.
         self._workspace_allocation: Optional[tuple[object, ...]] = None
@@ -288,8 +291,12 @@ def _is_supported_with_reason(
             quant_mode = QuantMode(attn.quant_mode)
         except (TypeError, ValueError):
             return False, "invalid KV-cache quantization mode."
-        if quant_mode.has_kv_cache_quant():
-            return False, "quantized KV cache is not supported by the initial adapter."
+        is_mla = attn.is_mla_enable
+        is_fp8_mla = is_mla and quant_mode.has_fp8_kv_cache()
+        if quant_mode.has_kv_cache_quant() and (
+            not is_fp8_mla or quant_mode.has_int8_kv_cache() or quant_mode.has_fp4_kv_cache()
+        ):
+            return False, "quantized KV cache is supported only for FP8 MLA decode."
 
         input_type = fwd.attention_input_type
         if input_type not in (
@@ -315,7 +322,6 @@ def _is_supported_with_reason(
             )
         if attn.num_heads <= 0 or attn.num_kv_heads <= 0:
             return False, "query and KV head counts must be positive."
-        is_mla = attn.is_mla_enable
         if is_mla:
             if attn.num_kv_heads != 1:
                 return False, "MLA decode requires one logical KV head."
@@ -330,7 +336,9 @@ def _is_supported_with_reason(
         if q.dtype not in self.SUPPORTED_DTYPES:
             return False, f"query dtype {q.dtype} is unsupported."
         cache_dtype = binding_to_torch_dtype(meta.kv_cache_manager.dtype)
-        if cache_dtype != q.dtype:
+        if is_fp8_mla and cache_dtype != torch.float8_e4m3fn:
+            return False, "FP8 MLA decode requires an FP8 E4M3 KV cache."
+        if not is_fp8_mla and cache_dtype != q.dtype:
             return False, f"query and KV-cache dtypes must match, got {q.dtype} and {cache_dtype}."
         if output.dtype != q.dtype:
             return False, f"output dtype must match query dtype, got {output.dtype} and {q.dtype}."
@@ -341,7 +349,7 @@ def _is_supported_with_reason(
             if has_context or input_type != AttentionInputType.generation_only:
                 return False, "MLA is supported only for generation-only requests."
             if q.dtype != torch.bfloat16 or output.dtype != torch.bfloat16:
-                return False, "MLA decode requires BF16 query, cache, and output."
+                return False, "MLA decode requires BF16 query input and output."
             if attn.kv_lora_rank != 512 or attn.qk_rope_head_dim != 64:
                 return False, (
                     "MLA decode requires kv_lora_rank=512 and qk_rope_head_dim=64, got "
@@ -356,6 +364,17 @@ def _is_supported_with_reason(
                 return False, f"MLA query width must be {expected_width}, got {q.shape[1]}."
             if output.numel() != q.shape[0] * attn.num_heads * attn.kv_lora_rank:
                 return False, "MLA output has an incompatible extent."
+            if is_fp8_mla:
+                # The automatic 1CTA policy for <=64 query rows has no profile
+                # below 128 KV tokens. This is the plan's page-table capacity,
+                # not the live request length; short requests remain supported.
+                if attn.num_heads <= 64 and (
+                    meta.kv_cache_block_offsets.shape[-1] * tokens_per_block < 128
+                ):
+                    return False, "FP8 MLA with <=64 heads requires a paged KV capacity >=128."
+                error = self._mla_fp8_input_error(q, fwd)
+                if error is not None:
+                    return False, error
         else:
             expected_width = (attn.num_heads + 2 * attn.num_kv_heads) * attn.head_dim
             if q.shape[1] != expected_width:
@@ -412,6 +431,60 @@ def _is_supported_with_reason(
             return False, "the K-to-V page displacement could not be resolved."
         return True, ""
 
+    @staticmethod
+    def _mla_fp8_input_error(q: torch.Tensor, fwd: AttentionForwardArgs) -> str | None:
+        """Check the preprocessing buffers without reading device values."""
+        quant_q = fwd.quant_q_buffer
+        if quant_q is None:
+            return "FP8 MLA decode requires quant_q_buffer from MLA preprocessing."
+        if (
+            quant_q.dtype not in (torch.uint8, torch.float8_e4m3fn)
+            or quant_q.device != q.device
+            or not quant_q.is_contiguous()
+            or quant_q.numel() < q.numel()
+        ):
+            return (
+                "FP8 MLA quant_q_buffer must be contiguous uint8 or FP8 E4M3 on the "
+                "query device, with at least one element per query element."
+            )
+        for name, scale in (
+            ("mla_bmm1_scale", fwd.mla_bmm1_scale),
+            ("mla_bmm2_scale", fwd.mla_bmm2_scale),
+        ):
+            if (
+                scale is None
+                or scale.dtype != torch.float32
+                or scale.device != q.device
+                or not scale.is_contiguous()
+                or scale.numel() < 1
+            ):
+                return (
+                    f"FP8 MLA {name} must contain a contiguous float32 value on the query device."
+                )
+        return None
+
+    def _get_mla_fp8_inputs(
+        self, q: torch.Tensor, fwd: AttentionForwardArgs
+    ) -> tuple[torch.Tensor, float, float]:
+        error = self._mla_fp8_input_error(q, fwd)
+        if error is not None:
+            raise RuntimeError(error)
+        if self._mla_fp8_scales is None:
+            if torch.cuda.is_current_stream_capturing():
+                raise RuntimeError(
+                    "PrimTS FP8 MLA scales must be cached before CUDA graph capture."
+                )
+            # The producer stores the regular BMM1 scale at index 0, and its
+            # log2 version at index 1. PrimTS expects the regular scale.
+            bmm1_scale = float(fwd.mla_bmm1_scale.view(-1)[0].item())
+            bmm2_scale = float(fwd.mla_bmm2_scale.view(-1)[0].item())
+            if not all(math.isfinite(scale) and scale > 0 for scale in (bmm1_scale, bmm2_scale)):
+                raise RuntimeError("PrimTS FP8 MLA scales must be finite and positive.")
+            self._mla_fp8_scales = (bmm1_scale, bmm2_scale)
+        # Reinterpret the producer's bytes; do not quantize the BF16 query again.
+        query = fwd.quant_q_buffer.view(-1)[: q.numel()].view(torch.float8_e4m3fn)
+        return query, *self._mla_fp8_scales
+
     @staticmethod
     def _get_fixed_block_tables(
         block_tables: torch.Tensor,
@@ -699,6 +772,11 @@ def prepare_workspace(
                 get_prims_ts_batch_mla_decode_workspace_size,
             )
 
+            kernel_dtype = (
+                torch.float8_e4m3fn
+                if QuantMode(self.attn.quant_mode).has_fp8_kv_cache()
+                else q.dtype
+            )
             required_bytes = get_prims_ts_batch_mla_decode_workspace_size(
                 batch_size,
                 self.attn.num_heads,
@@ -707,8 +785,8 @@ def prepare_workspace(
                 int(metadata.tokens_per_block),
                 max_seq_len,
                 max_seq_len_q=seq_len_q,
-                q_dtype=q.dtype,
-                kv_dtype=q.dtype,
+                q_dtype=kernel_dtype,
+                kv_dtype=kernel_dtype,
                 out_dtype=forward_args.output.dtype,
                 mask_type=mask_type,
                 device=q.device,
@@ -1142,7 +1220,15 @@ def run_mla_generation(self, params: FmhaParams) -> None:
             raise RuntimeError("TRT-LLM did not return PrimTS MLA KV metadata.")
         fixed_block_tables = self._get_fixed_block_tables(block_tables, batch_size)
         seq_len_q = params.input_seq_length
-        query = params.qkv_input.view(
+        if QuantMode(attn.quant_mode).has_fp8_kv_cache():
+            query, bmm1_scale, bmm2_scale = self._get_mla_fp8_inputs(params.qkv_input, params.fwd)
+        else:
+            query = params.qkv_input
+            bmm1_scale = 1.0 / (
+                attn.q_scaling * math.sqrt(int(attn.qk_nope_head_dim) + int(attn.qk_rope_head_dim))
+            )
+            bmm2_scale = 1.0
+        query = query.view(
             batch_size,
             seq_len_q,
             attn.num_heads,
@@ -1155,9 +1241,6 @@ def run_mla_generation(self, params: FmhaParams) -> None:
             int(attn.kv_lora_rank),
         )
         max_seq_len = int(block_tables.shape[-1]) * params.tokens_per_block
-        bmm1_scale = 1.0 / (
-            attn.q_scaling * math.sqrt(int(attn.qk_nope_head_dim) + int(attn.qk_rope_head_dim))
-        )
         mask_type = self._get_prims_mask_type(params.fwd)
         seq_lens = self._get_sequence_lengths(
             params.sequence_lengths,
@@ -1185,7 +1268,7 @@ def run_mla_generation(self, params: FmhaParams) -> None:
             seq_lens=seq_lens,
             out=output,
             bmm1_scale=bmm1_scale,
-            bmm2_scale=1.0,
+            bmm2_scale=bmm2_scale,
             validate=False,
         )
 

tests/unittest/_torch/attention/test_prims_ts_attention_backend.py

@@ -15,6 +15,9 @@
 
 import functools
 import inspect
+import math
+from types import SimpleNamespace
+from unittest.mock import Mock
 
 import pytest
 import torch
@@ -344,6 +347,193 @@ def test_prims_ts_deepseek_v3_lite_mla_generation(
     run_case(case)
 
 
+@pytest.mark.parametrize("num_heads", [6, 12, 96])
+@pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], ids=["v1", "v2"])
+def test_prims_ts_fp8_mla_preprocessing(
+    monkeypatch: pytest.MonkeyPatch, num_heads: int, use_kv_cache_manager_v2: bool
+) -> None:
+    from test_attention_mla import RopeConfig, _run_test_for_backend
+
+    from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams
+    from tensorrt_llm._torch.attention.backends.fmha.prims_ts import PrimsTSFmha
+
+    # Context remains on the regular backend; every decode must select PrimTS.
+    monkeypatch.setenv("TLLM_FMHA_LIBS", "prims_ts,fallback")
+    original_run = PrimsTSFmha.run_mla_generation
+    calls = []
+
+    def run_and_capture(fmha: PrimsTSFmha, params: FmhaParams) -> None:
+        assert params.qkv_input.dtype == torch.bfloat16
+        assert params.fwd.quant_q_buffer.dtype == torch.uint8
+        original_run(fmha, params)
+        eager = params.context_buf.clone()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            original_run(fmha, params)
+        params.context_buf.zero_()
+        graph.replay()
+        torch.cuda.synchronize()
+        torch.testing.assert_close(params.context_buf, eager, atol=0, rtol=0)
+        calls.append(params.attn.local_layer_idx)
+
+    monkeypatch.setattr(PrimsTSFmha, "run_mla_generation", run_and_capture)
+    _run_test_for_backend(
+        backend_name="TRTLLM",
+        num_heads=num_heads,
+        num_kv_heads=num_heads,
+        num_layers=2,
+        q_lora_rank=1536,
+        kv_lora_rank=512,
+        qk_nope_head_dim=128,
+        qk_rope_head_dim=64,
+        v_head_dim=128,
+        rope_config=RopeConfig(
+            num_attention_heads=num_heads,
+            max_position_embeddings=4096,
+            rope_scaling={
+                "beta_fast": 32,
+                "beta_slow": 1,
+                "factor": 40.0,
+                "mscale": 1.0,
+                "mscale_all_dim": 1.0,
+                "original_max_position_embeddings": 4096,
+                "type": "yarn",
+            },
+        ),
+        kv_cache_tokens_per_block=32,
+        device=torch.device("cuda"),
+        dtype=torch.bfloat16,
+        kv_cache_dtype=torch.float8_e4m3fn,
+        context_sequence_lengths=[129, 191],
+        generation_seq_len_q=1,
+        num_generation_steps=2,
+        v2_kv_cache=use_kv_cache_manager_v2,
+    )
+    assert calls == [0, 1, 0, 1]
+
+
+@pytest.mark.parametrize("num_heads", [6, 12, 96])
+@pytest.mark.parametrize("batch_size", [2, 65])
+def test_prims_ts_fp8_mla_scales_and_graph_replay(
+    monkeypatch: pytest.MonkeyPatch, num_heads: int, batch_size: int
+) -> None:
+    import tensorrt_llm._torch.attention.backends.fmha.prims_ts as prims_ts_module
+    from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams
+    from tensorrt_llm._torch.attention.backends.fmha.prims_ts import PrimsTSFmha
+    from tensorrt_llm._torch.attention.backends.interface import (
+        AttentionForwardArgs,
+        AttentionInputType,
+    )
+    from tensorrt_llm.quantization.mode import QuantMode
+
+    torch.manual_seed(0)
+    device = torch.device("cuda")
+    attn = Mock(
+        is_mla_enable=True,
+        num_heads=num_heads,
+        num_kv_heads=1,
+        kv_lora_rank=512,
+        qk_rope_head_dim=64,
+        qk_nope_head_dim=128,
+        head_dim=576,
+        v_head_dim=128,
+        local_layer_idx=0,
+        quant_mode=QuantMode.FP8_KV_CACHE,
+        q_scaling=1.0,
+    )
+    fmha = PrimsTSFmha(attn)
+    query = torch.randn(batch_size, num_heads, 576, device=device).to(torch.float8_e4m3fn)
+    kv_cache = torch.randn(batch_size * 2, 1, 32, 576, device=device).to(torch.float8_e4m3fn)
+    block_tables = torch.zeros((batch_size, 2, 4), dtype=torch.int32, device=device)
+    block_tables[:, 0, :2] = torch.arange(batch_size * 2, device=device).view(batch_size, 2)
+    seq_lens = torch.full((batch_size,), 33, dtype=torch.int32, device=device)
+    output = torch.empty((batch_size, num_heads * 512), dtype=torch.bfloat16, device=device)
+    # Deliberately use different, non-unit scales. BMM1[1] is log2-scaled and
+    # must not be passed to PrimTS; BMM2 must scale V independently of QK.
+    bmm1_scale, bmm2_scale = 0.025, 1.75
+    fwd = AttentionForwardArgs(
+        attention_input_type=AttentionInputType.generation_only,
+        attention_window_size=128,
+        output=output,
+        quant_q_buffer=query.view(torch.uint8),
+        mla_bmm1_scale=torch.tensor([bmm1_scale, bmm1_scale * math.log2(math.e)], device=device),
+        mla_bmm2_scale=torch.tensor([bmm2_scale], device=device),
+    )
+    metadata = SimpleNamespace(
+        host_kv_cache_pool_pointers=None,
+        host_kv_cache_pool_mapping=None,
+        kv_cache_block_offsets=block_tables,
+        num_contexts=0,
+        num_ctx_tokens=0,
+        num_generations=batch_size,
+        tokens_per_block=32,
+    )
+    # Pool ownership/indexing is covered by the real v1/v2 preprocessing test.
+    monkeypatch.setattr(
+        prims_ts_module.thop,
+        "build_trtllm_gen_kv_cache_metadata",
+        lambda *args: (kv_cache, block_tables, None),
+    )
+    # Poison the BF16 input: the adapter must consume the preprocessing bytes.
+    q = torch.full((batch_size, num_heads * 576), float("nan"), dtype=torch.bfloat16, device=device)
+    workspace = torch.empty(0, dtype=torch.uint8, device=device)
+    fmha.prepare_workspace(q, None, None, metadata, fwd, workspace)
+    params = FmhaParams(
+        attn=attn,
+        meta=metadata,
+        fwd=fwd,
+        workspace=workspace,
+        qkv_input=q,
+        context_buf=output,
+        sequence_lengths=seq_lens,
+        input_seq_length=1,
+        num_tokens=batch_size,
+        batch_size=batch_size,
+        num_requests=batch_size,
+        tokens_per_block=32,
+        kv_factor=1,
+        total_num_blocks=batch_size * 2,
+    )
+
+    def check_output() -> None:
+        pages = kv_cache[:, 0].float()[block_tables[:, 0, :2].long()].reshape(batch_size, 64, 576)
+        scores = torch.einsum("bhd,bkd->bhk", query.float(), pages) * bmm1_scale
+        invalid = torch.arange(64, device=device)[None, :] >= seq_lens[:, None]
+        scores.masked_fill_(invalid[:, None, :], float("-inf"))
+        ideal = torch.einsum("bhk,bkd->bhd", scores.softmax(-1), pages[..., :512]) * bmm2_scale
+        # This case fits one KV tile. FP8 MLA rounds the unnormalized P tile
+        # (scaled to E4M3's maximum 448) before PV, but keeps its sum in FP32.
+        # Model that intermediate rounding for the elementwise comparison.
+        probabilities = (scores - scores.amax(dim=-1, keepdim=True)).exp() * 448.0
+        rounded_p = probabilities.to(torch.float8_e4m3fn).float()
+        expected = (
+            torch.einsum("bhk,bkd->bhd", rounded_p, pages[..., :512])
+            / probabilities.sum(dim=-1, keepdim=True)
+            * bmm2_scale
+        )
+        torch.testing.assert_close(
+            output, expected.reshape_as(output).to(output.dtype), atol=2e-2, rtol=2e-2
+        )
+        # Also bound aggregate error against full-precision attention so the
+        # low-precision oracle does not hide a large numerical regression.
+        relative_error = (output.float().view_as(ideal) - ideal).norm() / ideal.norm()
+        assert relative_error.item() < 0.04
+
+    fmha.run_mla_generation(params)
+    check_output()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        fmha.run_mla_generation(params)
+    # New query bytes and page-table/length contents must be read on replay.
+    query.view(torch.uint8).copy_(query.view(torch.uint8).flip(0))
+    block_tables[:, 0, :2].copy_(block_tables[:, 0, :2].flip(0))
+    seq_lens.add_(7)
+    output.zero_()
+    graph.replay()
+    torch.cuda.synchronize()
+    check_output()
+
+
 def test_prims_ts_context_wrapper_cuda_graph_replay_with_updated_metadata(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None: