[None][perf] primsTS dense bf16, qk-bf16/pv-fp8, fp8 fmha context bringup

docs/source/features/visualgen-quantized-attention.md

@@ -29,6 +29,8 @@ A recipe is the tuple `(qk_dtype, v_dtype, (q_block_size, k_block_size, v_block_
 | `TRTLLM` | `int8` | `fp8` | `(2, 16, 1)` | SageAttention (INT8 QK) for Hopper |
 | `TRTLLM` | `int8` | `fp8` | `(1, 1, 1)`, `(1, 4, 1)`, `(1, 16, 1)` | SageAttention (INT8 QK) for Blackwell |
 | `TRTLLM` | `fp8` | `fp8` | `(1, 1, 1)`, `(1, 4, 1)` | SageAttention (FP8 QK) for Blackwell |
+| `TRTLLM` | `bf16` | `fp8` | `(0, 0, 0)` | PrimTS QK16PV8 for Blackwell (`TLLM_FMHA_LIBS=+prims_ts`) |
+| `TRTLLM` | `fp8` | `fp8` | `(0, 0, 0)` | PrimTS FP8 for Blackwell (`TLLM_FMHA_LIBS=+prims_ts`) |
 | `CUDNN` | `fp8` | `fp8` | `(0, 0, 0)` | cuDNN FP8 |
 | `CUDNN` | `mxfp8` | `mxfp8` | `(0, 0, 0)` | cuDNN MXFP8 |
 | `CUTEDSL` | `bf16` | `fp8` | `(0, 0, 0)` | QK16PV8 |
@@ -38,6 +40,8 @@ A recipe is the tuple `(qk_dtype, v_dtype, (q_block_size, k_block_size, v_block_
 | `FLASHINFER` | `nvfp4` | `fp8` | `(0, 0, 0)` | NVFP4 Q/K |
 | `FLASHINFER` | `nvfp4` | `nvfp4` | `(0, 0, 0)` | NVFP4 Attention |
 
+The `TRTLLM` backend picks the kernel from the block sizes: SageAttention quantizes per block, and the `(0, 0, 0)` per-tensor recipes run on the PrimTS FMHA library, which `TLLM_FMHA_LIBS=+prims_ts` enables.
+
 ### Choosing and Tuning a Recipe
 
 Choose a quantized-attention recipe by output quality first. Establish an unquantized quality baseline, then evaluate each compatible recipe with representative prompts, input media, resolutions, and fixed seeds. Quantization sensitivity varies by model, so do not assume that a recipe validated for one model will preserve quality for another.

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

@@ -48,6 +48,8 @@
         "skip_mla_rope_generation",  # handled in ``TrtllmAttention.forward`` for the test-only MLA path
         "timestep",  # consumed by sparse prediction before FMHA dispatch
         "sparse_attn_phase",  # host-resolved sparse phase, consumed before FMHA dispatch
+        "qk_fp8_per_tensor",  # per-tensor FP8 recipe flag, consumed by FMHA selection and the PrimTS FMHA
+        "v_fp8_per_tensor",  # per-tensor FP8 recipe flag, consumed by FMHA selection and the PrimTS FMHA
     }
 )
 

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

@@ -42,6 +42,7 @@ class Fmha(ABC):
 
     supports_skip_correction: ClassVar[bool] = False
     supports_block_sparse_inputs: ClassVar[bool] = False
+    supports_per_tensor_fp8: ClassVar[bool] = False
     supports_workspace_reclamation: bool = False
     supports_fp4_mla: ClassVar[bool] = False
 
@@ -119,6 +120,12 @@ def is_supported(
         ):
             logger.debug(f"{type(self).__name__} does not support block-sparse inputs.")
             return False
+        if (
+            getattr(forward_args, "qk_fp8_per_tensor", False)
+            or getattr(forward_args, "v_fp8_per_tensor", False)
+        ) and not self.supports_per_tensor_fp8:
+            logger.debug(f"{type(self).__name__} does not support per-tensor FP8 recipes.")
+            return False
         return self._is_supported(q, k, v, metadata, forward_args, phase=phase)
 
     def _is_supported(

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

@@ -43,7 +43,10 @@
 from .utils import get_kv_page_offset, get_multi_processor_count_for_device
 
 if TYPE_CHECKING:
-    from tensorrt_llm._torch.attention.backends.prims_ts.context import BatchPrefillPagedTSWrapper
+    from tensorrt_llm._torch.attention.backends.prims_ts.context import (
+        BatchPrefillPagedTSWrapper,
+        BatchPrefillTSWrapper,
+    )
     from tensorrt_llm._torch.attention.backends.prims_ts.decode import BatchDecodePagedTSWrapper
     from tensorrt_llm._torch.attention.backends.prims_ts.mla_decode import (
         BatchMLADecodePagedTSWrapper,
@@ -158,6 +161,8 @@ def get_attention_feature_unsupported_reason(
 class PrimsTSFmha(PhasedFmha):
     """Blackwell task-scheduled paged context and decode FMHA library."""
 
+    supports_per_tensor_fp8 = True
+
     SUPPORTED_PAGE_SIZES = {16, 32, 64, 128}
     SUPPORTED_CONTEXT_HEAD_DIMS = {128, 256}
     SUPPORTED_DECODE_HEAD_DIMS = {64, 128, 256}
@@ -173,6 +178,8 @@ def __init__(self, attn: "TrtllmAttention") -> None:
         # Every other plan attribute is fixed by this layer/model instance.
         # Batch size is the only execution profile that needs its own wrapper.
         self._context_wrappers: dict[int, "BatchPrefillPagedTSWrapper"] = {}
+        # Cache-free context plans, keyed by (batch, seq_len, q/v/out dtype, mask).
+        self._contiguous_context_wrappers: dict[tuple[Any, ...], "BatchPrefillTSWrapper"] = {}
         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
@@ -287,17 +294,14 @@ def _is_supported_with_reason(
         if q.device.type != "cuda":
             return False, "CUDA tensors are required."
         if not q.is_contiguous():
-            return False, "the fused attention input must be contiguous."
-        if k is not None or v is not None:
-            return False, "only fused QKV input is supported."
-        if not attn.is_mla_enable and not fwd.is_fused_qkv:
+            return False, "the attention input must be contiguous."
+        if meta.kv_cache_manager is None:
+            if k is None or v is None or fwd.is_fused_qkv:
+                return False, "cache-free attention takes separate Q, K and V."
+        elif k is not None or v is not None or (not attn.is_mla_enable and not fwd.is_fused_qkv):
             return False, "only fused QKV input is supported."
         if meta.is_cross:
             return False, "cross attention is not supported."
-        storage_reason = get_paged_kv_storage_unsupported_reason(attn, meta)
-        if storage_reason is not None:
-            return False, storage_reason
-
         output = fwd.output
         if output is None:
             return False, "an output tensor is required."
@@ -318,6 +322,11 @@ def _is_supported_with_reason(
         feature_reason = get_attention_feature_unsupported_reason(meta, fwd)
         if feature_reason is not None:
             return False, feature_reason
+        if meta.kv_cache_manager is None:
+            return self._is_supported_without_kv_cache(q, k, v, attn, meta, fwd)
+        storage_reason = get_paged_kv_storage_unsupported_reason(attn, meta)
+        if storage_reason is not None:
+            return False, storage_reason
 
         try:
             mask_type = AttentionMaskType(fwd.mask_type)
@@ -457,6 +466,78 @@ def _is_supported_with_reason(
             return False, "the K-to-V page displacement could not be resolved."
         return True, ""
 
+    def _is_supported_without_kv_cache(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        attn: "TrtllmAttention",
+        meta: "TrtllmAttentionMetadata",
+        fwd: AttentionForwardArgs,
+    ) -> tuple[bool, str]:
+        """Support decision for separate-QKV context attention without a KV cache."""
+        if attn.is_mla_enable:
+            return False, "cache-free attention does not support MLA."
+        if fwd.attention_input_type == AttentionInputType.generation_only or int(
+            meta.num_generations
+        ):
+            return False, "cache-free attention serves the context phase only."
+        num_contexts = int(meta.num_contexts)
+        if num_contexts <= 0:
+            return False, "the request contains no active context requests."
+        if (attn.attention_chunk_size or 0) != 0:
+            return False, "chunked context attention is not supported."
+        try:
+            mask_type = AttentionMaskType(fwd.mask_type)
+        except (AttributeError, TypeError, ValueError):
+            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."
+        if attn.num_heads <= 0 or attn.num_kv_heads <= 0 or attn.num_heads % attn.num_kv_heads:
+            return False, "the query head count must be a positive multiple of the KV head count."
+        if attn.head_dim not in self.SUPPORTED_CONTEXT_HEAD_DIMS:
+            return False, f"context head dimension {attn.head_dim} is unsupported."
+        if q.dtype not in self.SUPPORTED_DTYPES or k.dtype != q.dtype or v.dtype != q.dtype:
+            return (
+                False,
+                f"Q, K and V must share a supported dtype, got {q.dtype}/{k.dtype}/{v.dtype}.",
+            )
+        if fwd.v_fp8_per_tensor and not fwd.qk_fp8_per_tensor and q.dtype != torch.bfloat16:
+            return False, f"the attention recipe keeps Q/K in BF16, got {q.dtype}."
+        output = fwd.output
+        if output.dtype != q.dtype:
+            return False, f"output dtype must match query dtype, got {output.dtype} and {q.dtype}."
+        if q.ndim != 2 or k.ndim != 2 or v.ndim != 2:
+            return False, "attention inputs must be rank 2."
+        if not k.is_contiguous() or not v.is_contiguous():
+            return False, "the attention input must be contiguous."
+        kv_width = attn.num_kv_heads * attn.head_dim
+        if (
+            q.shape[1] != attn.num_heads * attn.head_dim
+            or k.shape[1] != kv_width
+            or v.shape[1] != kv_width
+        ):
+            return False, "Q, K and V widths do not match the layer's heads and head dimension."
+        if k.shape[0] != q.shape[0] or v.shape[0] != q.shape[0]:
+            return False, "cache-free attention requires equal Q and K/V lengths."
+        if output.numel() != q.shape[0] * attn.num_heads * attn.head_dim:
+            return False, "attention output has an incompatible extent."
+        seq_len = int(meta.max_seq_len)
+        host_lens = meta.prompt_lens_cpu_runtime
+        if host_lens is None or host_lens.numel() < num_contexts:
+            return False, "host context lengths are required."
+        active_lens = host_lens[:num_contexts]
+        if (
+            seq_len <= 0
+            or q.shape[0] != num_contexts * seq_len
+            or not bool((active_lens == seq_len).all())
+        ):
+            return False, "cache-free attention requires uniform sequence lengths."
+        window = fwd.attention_window_size
+        if isinstance(window, int) and 0 < window < seq_len:
+            return False, "sliding-window attention is not supported."
+        return True, ""
+
     @staticmethod
     def _mla_fp8_input_error(q: torch.Tensor, fwd: AttentionForwardArgs) -> str | None:
         """Check the preprocessing buffers without reading device values."""
@@ -730,6 +811,102 @@ def _get_or_plan_mla_decode_wrapper(
         self._mla_decode_wrappers[batch_size] = wrapper
         return wrapper
 
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: Optional[torch.Tensor],
+        v: Optional[torch.Tensor],
+        metadata: "TrtllmAttentionMetadata",
+        forward_args: AttentionForwardArgs,
+    ) -> None:
+        """Run the request on the cache-free context path or the phased paged path."""
+        if metadata.kv_cache_manager is None and not int(metadata.num_generations):
+            self._forward_without_kv_cache(q, k, v, metadata, forward_args)
+            return
+        super().forward(q, k, v, metadata, forward_args)
+
+    def _forward_without_kv_cache(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        metadata: "TrtllmAttentionMetadata",
+        forward_args: AttentionForwardArgs,
+    ) -> None:
+        """Separate-QKV context attention without a KV cache on the contiguous kernel."""
+        attn = self.attn
+        output = forward_args.output
+        if output is None:
+            raise RuntimeError("PrimTS context requires an output buffer.")
+        batch_size = int(metadata.num_contexts)
+        seq_len = int(metadata.max_seq_len)
+        # Library selection is cached without sequence lengths, so a ragged batch can
+        # reach this path on a cache hit. Re-check before the view below.
+        host_lens = metadata.prompt_lens_cpu_runtime[:batch_size]
+        if q.shape[0] != batch_size * seq_len or not bool((host_lens == seq_len).all()):
+            raise RuntimeError("PrimTS cache-free context requires uniform sequence lengths.")
+        num_heads, num_kv_heads, head_dim = attn.num_heads, attn.num_kv_heads, attn.head_dim
+        # The fixed contiguous plan takes compact [B, S, H, D] operands.
+        q_bshd = q.view(batch_size, seq_len, num_heads, head_dim)
+        k_bshd = k.view(batch_size, seq_len, num_kv_heads, head_dim)
+        v_bshd = v.view(batch_size, seq_len, num_kv_heads, head_dim)
+        # Per-tensor E4M3 operands as selected by the recipe. The Q and K dequant scales
+        # fold into the per-run base-2 softmax scale and the V dequant scale into the
+        # kernel output scale.
+        sm_scale = self._get_bmm1_scale(attn)
+        scale_softmax_log2 = None
+        if forward_args.qk_fp8_per_tensor:
+            q_bshd, q_dequant_scale = torch.ops.tensorrt_llm.quantize_e4m3_per_tensor(q_bshd)
+            k_bshd, k_dequant_scale = torch.ops.tensorrt_llm.quantize_e4m3_per_tensor(k_bshd)
+            scale_softmax_log2 = (
+                q_dequant_scale.float() * k_dequant_scale.float() * (sm_scale * math.log2(math.e))
+            ).reshape(1)
+        output_scale = None
+        if forward_args.v_fp8_per_tensor:
+            v_bshd, v_dequant_scale = torch.ops.tensorrt_llm.quantize_e4m3_per_tensor(v_bshd)
+            output_scale = v_dequant_scale.float().reshape(1)
+        out = output.view(batch_size, seq_len, num_heads, head_dim)
+        mask_type = self._get_prims_mask_type(forward_args)
+        key = (batch_size, seq_len, q_bshd.dtype, v_bshd.dtype, out.dtype, mask_type)
+        wrapper = self._contiguous_context_wrappers.get(key)
+        if wrapper is None:
+            if torch.cuda.is_current_stream_capturing():
+                raise RuntimeError("PrimTS context must be planned before CUDA graph capture.")
+            from tensorrt_llm._torch.attention.backends.prims_ts.context import (
+                BatchPrefillTSWrapper,
+            )
+
+            logger.info(
+                f"PrimTS cache-free context: q/k_dtype={q_bshd.dtype} v_dtype={v_bshd.dtype} "
+                f"batch={batch_size} seq_len={seq_len} heads={num_heads}/{num_kv_heads} head_dim={head_dim}"
+            )
+            wrapper = BatchPrefillTSWrapper()
+            wrapper.plan(
+                device=q.device,
+                batch_size=batch_size,
+                max_seq_len_q=seq_len,
+                max_kv_len=seq_len,
+                num_qo_heads=num_heads,
+                num_kv_heads=num_kv_heads,
+                head_dim=head_dim,
+                q_dtype=q_bshd.dtype,
+                k_dtype=k_bshd.dtype,
+                v_dtype=v_bshd.dtype,
+                out_dtype=out.dtype,
+                mask_type=mask_type,
+                sm_scale=sm_scale,
+            )
+            self._contiguous_context_wrappers[key] = wrapper
+        wrapper.run(
+            q_bshd,
+            k_bshd,
+            v_bshd,
+            out=out,
+            scale_softmax_log2=scale_softmax_log2,
+            output_scale=output_scale,
+            validate=False,
+        )
+
     def prepare_workspace(
         self,
         q: torch.Tensor,

tensorrt_llm/_torch/attention/backends/interface.py

@@ -968,6 +968,10 @@ class AttentionForwardArgs:
     sage_attn_num_elts_per_blk_v: int = 0
     sage_attn_qk_int8: bool = False
     sage_attn_smooth_k: bool = False
+    # Per-tensor FP8 recipe of the PrimTS FMHA library. Quantize Q/K and/or V to
+    # E4M3 with one scale each, folded into the softmax and output scales.
+    qk_fp8_per_tensor: bool = False
+    v_fp8_per_tensor: bool = False
 
     # Packed QKV for non-MLA attention. MLA always passes a separate query.
     is_fused_qkv: bool = False

tensorrt_llm/_torch/attention/backends/trtllm.py

@@ -2465,6 +2465,12 @@ def forward(
                                              forward_args)
 
             if fmha is None:
+                if (getattr(forward_args, "qk_fp8_per_tensor", False)
+                        or getattr(forward_args, "v_fp8_per_tensor", False)):
+                    raise RuntimeError(
+                        "The per-tensor FP8 attention recipe needs the PrimTS "
+                        "FMHA library (TLLM_FMHA_LIBS=+prims_ts), which is not "
+                        "enabled or does not support this request.")
                 raise RuntimeError(
                     "No TRT-LLM attention FMHA library supports this request.")
             if metadata.is_cuda_graph or not fmha.supports_workspace_reclamation:

tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py

@@ -191,6 +191,7 @@ class TrtllmAttention(BaseTrtllmAttention, AttentionBackend):
     - Metadata creation and preparation
     - No KV cache operation
     - SageAttention per-block QKV quantization (when a quant_attention_config is provided. requires unfused QKV)
+    - PrimTS per-tensor FP8 recipes (quant_attention_config with zero block sizes, requires unfused QKV)
     - Separate-QKV forwarding for generic block-sparse attention and backends that reject fused QKV
     """
 
@@ -233,6 +234,16 @@ def __init__(
         )
 
         self.quant_attention_config = quant_attention_config
+        if (
+            quant_attention_config is not None
+            and quant_attention_config.q_block_size == 0
+            and not any(fmha.supports_per_tensor_fp8 for fmha in self._fmha_manager.fmha_libs)
+        ):
+            raise RuntimeError(
+                "The per-tensor FP8 attention recipe needs the PrimTS FMHA library "
+                "(TLLM_FMHA_LIBS=+prims_ts), which is not enabled or not available on this "
+                "platform."
+            )
 
     @property
     def timestep_cutoff(self) -> Optional[float]:
@@ -399,6 +410,9 @@ def forward(
                     "sage_attn_num_elts_per_blk_v": quant_cfg.v_block_size,
                     "sage_attn_qk_int8": quant_cfg.qk_dtype == "int8",
                     "sage_attn_smooth_k": quant_cfg.smooth_k,
+                    "qk_fp8_per_tensor": quant_cfg.q_block_size == 0
+                    and quant_cfg.qk_dtype == "fp8",
+                    "v_fp8_per_tensor": quant_cfg.v_block_size == 0 and quant_cfg.v_dtype == "fp8",
                 }
         else:
             if k is None and v is None:

tensorrt_llm/visual_gen/args.py

@@ -158,6 +158,17 @@ def _validate_quant_attention_config(self) -> "AttentionConfig":
                 ("fp8", "fp8", (1, 4, 1)),
             },
         }
+        # PrimTS serves the per-tensor FP8 recipes on Blackwell (TLLM_FMHA_LIBS=+prims_ts).
+        PRIMSTS_RECIPES = {
+            100: {
+                ("bf16", "fp8", (0, 0, 0)),
+                ("fp8", "fp8", (0, 0, 0)),
+            },
+            103: {
+                ("bf16", "fp8", (0, 0, 0)),
+                ("fp8", "fp8", (0, 0, 0)),
+            },
+        }
         # Other recipes verify the hardware at corresponding backend implementations.
         CUDNN_RECIPES = {
             ("fp8", "fp8", (0, 0, 0)),
@@ -186,13 +197,19 @@ def _validate_quant_attention_config(self) -> "AttentionConfig":
             (q_config.q_block_size, q_config.k_block_size, q_config.v_block_size),
         )
         if self.backend == "TRTLLM":
-            recipes = SAGE_RECIPES.get(get_sm_version(), set())
+            sm_version = get_sm_version()
+            recipes = SAGE_RECIPES.get(sm_version, set()) | PRIMSTS_RECIPES.get(sm_version, set())
             if recipe not in recipes:
                 raise ValueError(
                     f"Unsupported quant_attention_config={self.quant_attention_config!r} "
-                    f"for backend='TRTLLM'. Supported SAGE recipes on this device: "
+                    f"for backend='TRTLLM'. Supported SAGE and PrimTS recipes on this device: "
                     f"{sorted(recipes)}."
                 )
+            if q_config.smooth_k and q_config.q_block_size == 0:
+                raise ValueError(
+                    "smooth_k is a SageAttention option and does not apply to the per-tensor "
+                    f"recipe {recipe!r}."
+                )
         elif q_config.smooth_k:
             raise ValueError(
                 f"smooth_k is a SageAttention option and requires backend='TRTLLM', got "