[None][feat] Support FP8 KV cache for non-MLA PrimsTS attention

3rdparty/vendor_sources.lock.yaml

@@ -14,15 +14,15 @@ vendors:
     digest: sha256-tree-v1:d2d382335ccb68b55078523e63835c28b699607227afb81308cb7bb1d85845ac
   flashinfer-prims-ts:
     url: https://github.com/yuxianq/flashinfer.git
-    branch: trtllm-prims-ts-dev-20261012
-    commit: aa709061230a58858375c513056e7733fa74294a
+    branch: prims-tsfmha-fp8-non-mla--kernels
+    commit: dd087c6568aa144bb7ceff1306e2568279ec5068
     source: flashinfer/attention/prims_ts
     destination: tensorrt_llm/_torch/attention/backends/prims_ts
     include:
     - '**/*.py'
     patch: 3rdparty/vendor_patches/flashinfer-prims-ts.patch
     patch_digest: sha256:fc5c23ec3d65c562025a1845f211fdac66bfcd3da156dbb8385f4c4407d20d07
-    digest: sha256-tree-v1:f2f444f7e6d4ad688fbf4e1e9eed11ff7c78d835cfd0af5ef9d7660780b059e0
+    digest: sha256-tree-v1:a794ed3cc25687817ea08878b1e156ce7616ce553058c1a3a70a122ee915e9fc
   openengine-proto:
     url: https://github.com/ai-dynamo/openengine.git
     tag: v0.1.0

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

@@ -286,7 +286,7 @@ def _page_size_exclusive_to_cute_dsl(meta: "TrtllmAttentionMetadata") -> bool:
         if tokens_per_block in getattr(kv_cache_manager, "trtllm_gen_extra_tokens_per_block", ()):
             return False
         return not (
-            get_sm_version() in PrimsTSFmha.SUPPORTED_SMS
+            get_sm_version() in PrimsTSFmha.SUPPORTED_MLA_SMS
             and tokens_per_block in PrimsTSFmha.SUPPORTED_PAGE_SIZES
             and getattr(kv_cache_manager, "dtype", None) in (DataType.FP8, DataType.BF16)
         )

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

@@ -13,17 +13,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""TRT-LLM FMHA adapter for the vendored PrimTS Blackwell kernels."""
+"""TRT-LLM FMHA adapter for the vendored PrimTS kernels."""
 
 from __future__ import annotations
 
 import math
 from importlib import import_module
-from importlib.metadata import PackageNotFoundError, version
 from typing import TYPE_CHECKING, Any, Optional
 
 import torch
-from packaging.version import InvalidVersion, Version
 
 from tensorrt_llm._torch.attention.backends.interface import (
     AttentionForwardArgs,
@@ -57,7 +55,6 @@
     )
 
 
-_MIN_CUTLASS_DSL_VERSION = Version("4.7.0")
 _MIN_CUTLASS_COMPILER_VERSION = "13.3"
 _WORKSPACE_ALIGNMENT = 32
 
@@ -105,6 +102,8 @@ def get_paged_kv_storage_unsupported_reason(
 def get_paged_kv_policy_unsupported_reason(
     attn: "TrtllmAttention",
     metadata: "TrtllmAttentionMetadata",
+    *,
+    allow_fp8_kv_cache: bool = False,
 ) -> Optional[str]:
     """Return why the request's decoding policy is outside the fixed page-table envelope."""
     if metadata.beam_width != 1:
@@ -123,13 +122,15 @@ def get_paged_kv_policy_unsupported_reason(
         quant_mode = QuantMode(attn.quant_mode)
     except (TypeError, ValueError):
         return "invalid KV-cache quantization mode."
-    is_fp8_mla = (
-        attn.is_mla_enable
+    is_supported_fp8 = (
+        (attn.is_mla_enable or allow_fp8_kv_cache)
         and quant_mode.has_fp8_kv_cache()
         and not quant_mode.has_int8_kv_cache()
         and not quant_mode.has_fp4_kv_cache()
     )
-    if quant_mode.has_kv_cache_quant() and not is_fp8_mla:
+    if quant_mode.has_kv_cache_quant() and not is_supported_fp8:
+        if allow_fp8_kv_cache:
+            return "only FP8 E4M3 KV-cache quantization is supported."
         return "quantized KV cache is supported only for FP8 MLA decode."
     return None
 
@@ -153,11 +154,17 @@ def get_attention_feature_unsupported_reason(
 
 
 class PrimsTSFmha(PhasedFmha):
-    """Blackwell task-scheduled paged context and decode FMHA library."""
+    """Task-scheduled paged context and decode FMHA library.
+
+    Standard FP8 KV cache uses fused preprocessing to quantize Q with the
+    shared Q/K/V scale. Kernels consume FP8 Q/K/V and write the model's FP16
+    or BF16 output directly, using preprocessing-produced device scales.
+    """
 
     supports_per_tensor_fp8 = True
 
-    SUPPORTED_SMS = frozenset({100, 103})
+    SUPPORTED_SMS = frozenset({100, 103, 107})
+    SUPPORTED_MLA_SMS = frozenset({100, 103})
     SUPPORTED_PAGE_SIZES = {16, 32, 64, 128}
     SUPPORTED_CONTEXT_HEAD_DIMS = {128, 256}
     SUPPORTED_DECODE_HEAD_DIMS = {64, 128, 256}
@@ -192,20 +199,13 @@ def __init__(self, attn: "TrtllmAttention") -> None:
     def _is_available(cls, attn: "TrtllmAttention") -> bool:
         sm = get_sm_version()
         if sm not in cls.SUPPORTED_SMS:
-            logger.debug(f"PrimTS FMHA is unavailable: requires SM100 or SM103, got SM{sm}.")
-            return False
-        try:
-            installed_version = Version(version("nvidia-cutlass-dsl"))
-        except (PackageNotFoundError, InvalidVersion):
-            logger.debug("PrimTS FMHA is unavailable: nvidia-cutlass-dsl>=4.7 is required.")
-            return False
-        if installed_version < _MIN_CUTLASS_DSL_VERSION:
             logger.debug(
-                "PrimTS FMHA is unavailable: "
-                f"nvidia-cutlass-dsl>={_MIN_CUTLASS_DSL_VERSION} is required, "
-                f"got {installed_version}."
+                f"PrimTS FMHA is unavailable: requires SM100, SM103, or SM107, got SM{sm}."
             )
             return False
+        if attn.is_mla_enable and sm not in cls.SUPPORTED_MLA_SMS:
+            logger.debug(f"PrimTS MLA FMHA is unavailable on SM{sm}.")
+            return False
         try:
             cutlass = import_module("cutlass")
             compiler_version_supported = cutlass.target_version(
@@ -331,11 +331,12 @@ def _is_supported_with_reason(
             return False, "the attention mask is not causal or dense."
         if mask_type not in (AttentionMaskType.causal, AttentionMaskType.padding):
             return False, f"attention mask type {mask_type} is not supported."
-        policy_reason = get_paged_kv_policy_unsupported_reason(attn, meta)
+        policy_reason = get_paged_kv_policy_unsupported_reason(attn, meta, allow_fp8_kv_cache=True)
         if policy_reason is not None:
             return False, policy_reason
         is_mla = attn.is_mla_enable
-        is_fp8_mla = is_mla and QuantMode(attn.quant_mode).has_fp8_kv_cache()
+        is_fp8_kv = QuantMode(attn.quant_mode).has_fp8_kv_cache()
+        is_fp8_mla = is_mla and is_fp8_kv
 
         input_type = fwd.attention_input_type
         if input_type not in (
@@ -375,9 +376,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 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:
+        if is_fp8_kv and cache_dtype != torch.float8_e4m3fn:
+            return False, "FP8 attention requires an FP8 E4M3 KV cache."
+        if not is_fp8_kv 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}."
@@ -408,6 +409,22 @@ def _is_supported_with_reason(
                 if error is not None:
                     return False, error
         else:
+            if is_fp8_kv:
+                for name, scale in (
+                    ("kv_scale_orig_quant", fwd.kv_scale_orig_quant),
+                    ("kv_scale_quant_orig", fwd.kv_scale_quant_orig),
+                ):
+                    if (
+                        scale is None
+                        or scale.dtype != torch.float32
+                        or scale.device != q.device
+                        or not scale.is_contiguous()
+                        or scale.numel() != 1
+                    ):
+                        return (
+                            False,
+                            f"FP8 attention requires scalar float32 {name} on the query device.",
+                        )
             expected_width = (attn.num_heads + 2 * attn.num_kv_heads) * attn.head_dim
             if q.shape[1] != expected_width:
                 return False, f"fused QKV width must be {expected_width}, got {q.shape[1]}."
@@ -757,6 +774,7 @@ def _get_or_plan_decode_wrapper(
             mask_type=mask_type,
             window_left=window_left,
             workspace_buffer=workspace_buffer,
+            use_device_scales=QuantMode(self.attn.quant_mode).has_fp8_kv_cache(),
         )
         self._decode_wrappers[batch_size] = wrapper
         return wrapper
@@ -936,7 +954,7 @@ def prepare_workspace(
                 self.attn.head_dim,
                 self.attn.rope_dim,
                 True,
-                False,
+                QuantMode(self.attn.quant_mode).has_fp8_kv_cache(),
                 skip_fmha_workspace=True,
             )
             required_preprocess_bytes = max(
@@ -983,16 +1001,14 @@ def prepare_workspace(
         # sizing key aligned with the non-windowed plan selected at runtime.
         window_left = -1
 
+        kernel_dtype = (
+            torch.float8_e4m3fn if QuantMode(self.attn.quant_mode).has_fp8_kv_cache() else q.dtype
+        )
         if self.attn.is_mla_enable:
             from tensorrt_llm._torch.attention.backends.prims_ts import (
                 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,
@@ -1020,8 +1036,8 @@ def prepare_workspace(
                 int(metadata.tokens_per_block),
                 max_seq_len,
                 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,
                 window_left=window_left,
@@ -1108,8 +1124,8 @@ def run_context(self, params: FmhaParams) -> None:
             kv_pool,
             block_tables,
             _kv_scale_pool,
-            _bmm1_scale,
-            _bmm2_scale,
+            bmm1_scale,
+            bmm2_scale,
             fmha_workspace,
             cu_q_seqlens,
             _cu_kv_seqlens,
@@ -1152,7 +1168,7 @@ def run_context(self, params: FmhaParams) -> None:
             self._get_bmm1_scale(attn),
             1.0,
             attention_chunk_size,
-            False,
+            QuantMode(attn.quant_mode).has_fp8_kv_cache(),
             True,
             False,
             self._multi_processor_count,
@@ -1184,6 +1200,9 @@ def run_context(self, params: FmhaParams) -> None:
             params.sequence_lengths,
             params.batch_size,
         )
+        use_fp8 = QuantMode(attn.quant_mode).has_fp8_kv_cache()
+        if use_fp8:
+            q_processed = self._view_fp8_query(q_processed, params.num_tokens)
         mask_type = self._get_prims_mask_type(fwd)
         wrapper = self._get_or_plan_context_wrapper(
             q_processed,
@@ -1205,6 +1224,8 @@ def run_context(self, params: FmhaParams) -> None:
             cu_q_seqlens,
             block_tables=fixed_block_tables,
             seq_lens_kv=seq_lens_kv,
+            scale_softmax_log2=bmm1_scale.reshape(-1)[1:2] if use_fp8 else None,
+            output_scale=bmm2_scale.reshape(-1)[:1] if use_fp8 else None,
             out=params.output,
             validate=False,
         )
@@ -1242,7 +1263,7 @@ def run_context(self, params: FmhaParams) -> None:
             rope_params.max_positions,
             attn.position_embedding_type,
             self._get_bmm1_scale(attn),
-            False,
+            QuantMode(attn.quant_mode).has_fp8_kv_cache(),
             True,
             False,
             attention_chunk_size,
@@ -1296,7 +1317,7 @@ def _run_generation_preprocess(self, params: FmhaParams) -> tuple[Any, ...]:
             attn.position_embedding_type,
             self._get_bmm1_scale(attn),
             1.0,
-            False,
+            QuantMode(attn.quant_mode).has_fp8_kv_cache(),
             attn.predicted_tokens_per_seq,
             attention_chunk_size,
             self._multi_processor_count,
@@ -1324,8 +1345,8 @@ def run_generation(self, params: FmhaParams) -> None:
             kv_pool,
             block_tables,
             _kv_scale_pool,
-            _bmm1_scale,
-            _bmm2_scale,
+            bmm1_scale,
+            bmm2_scale,
             fmha_workspace,
             _cu_seqlens,
             _max_q_len,
@@ -1352,6 +1373,9 @@ def run_generation(self, params: FmhaParams) -> None:
         fixed_block_tables = self._get_fixed_block_tables(block_tables, batch_size)
         max_seq_len = int(block_tables.shape[-1]) * params.tokens_per_block
         seq_lens = self._get_sequence_lengths(params.sequence_lengths, batch_size)
+        use_fp8 = QuantMode(attn.quant_mode).has_fp8_kv_cache()
+        if use_fp8:
+            q_processed = self._view_fp8_query(q_processed, params.num_tokens)
         query = q_processed.view(
             batch_size,
             params.input_seq_length,
@@ -1396,12 +1420,26 @@ def run_generation(self, params: FmhaParams) -> None:
             (k_cache, v_cache),
             seq_lens,
             block_tables=fixed_block_tables,
-            bmm1_scale=self._get_bmm1_scale(attn),
-            bmm2_scale=1.0,
+            bmm1_scale=bmm1_scale.reshape(-1)[:1] if use_fp8 else self._get_bmm1_scale(attn),
+            bmm2_scale=bmm2_scale.reshape(-1)[:1] if use_fp8 else 1.0,
             out=output,
             validate=False,
         )
 
+    def _view_fp8_query(self, q_processed: torch.Tensor, num_tokens: int) -> torch.Tensor:
+        """View the packed FP8 prefix written by fused QKV preprocessing.
+
+        Generation reserves a model-dtype Q buffer, but writes one byte per
+        element when FP8 is enabled. Its unused trailing bytes are not queries.
+        """
+        numel = num_tokens * self.attn.num_heads * self.attn.head_dim
+        return (
+            q_processed.view(torch.uint8)
+            .reshape(-1)[:numel]
+            .view(torch.float8_e4m3fn)
+            .view(num_tokens, self.attn.num_heads, self.attn.head_dim)
+        )
+
     def _get_decode_workspace(
         self,
         root_workspace: torch.Tensor,