diff --git a/docs/source/features/visualgen-quantized-attention.md b/docs/source/features/visualgen-quantized-attention.md index 9d4f33a061e8..7f81a62d2143 100644 --- a/docs/source/features/visualgen-quantized-attention.md +++ b/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. diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py index 3cd442b39341..18a5effd1931 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py +++ b/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 } ) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/interface.py b/tensorrt_llm/_torch/attention/backends/fmha/interface.py index 2d1bd5c86742..e3905dd1964d 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/interface.py +++ b/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( diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py index 2e67616b5e9d..5268e72b5bd6 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/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, diff --git a/tensorrt_llm/_torch/attention/backends/interface.py b/tensorrt_llm/_torch/attention/backends/interface.py index 6a52657739df..7ccf04220604 100644 --- a/tensorrt_llm/_torch/attention/backends/interface.py +++ b/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 diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index d29def811bc5..048138a48582 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/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: diff --git a/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py b/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py index 3502ca6f9a4a..c1f0ef0ef0fc 100644 --- a/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py +++ b/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: diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index dc9c0e30184f..9f013a72d372 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/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 " diff --git a/tests/unittest/_torch/attention/test_fmha_interface.py b/tests/unittest/_torch/attention/test_fmha_interface.py index 015b7497e734..6c4f74997307 100644 --- a/tests/unittest/_torch/attention/test_fmha_interface.py +++ b/tests/unittest/_torch/attention/test_fmha_interface.py @@ -182,29 +182,39 @@ def _is_available(cls, attn: TrtllmAttention) -> bool: assert _SkipCorrectionFmha.is_available(attn) -@pytest.mark.parametrize("has_block_sparse_inputs", [False, True]) +@pytest.mark.parametrize( + ("capability", "request_args"), + [ + ( + "supports_block_sparse_inputs", + { + "sparse_runtime_params": SparseRuntimeParams( + block_sparse_inputs=BlockSparseForwardInputs( + q_block_size=64, + kv_block_size=64, + exact_block_bits=torch.zeros((1, 1, 1, 1), dtype=torch.int32), + ) + ) + }, + ), + ("supports_per_tensor_fp8", {"qk_fp8_per_tensor": True}), + ("supports_per_tensor_fp8", {"v_fp8_per_tensor": True}), + ], +) +@pytest.mark.parametrize("requested", [False, True]) @pytest.mark.parametrize("supported", [False, True]) -def test_support_checks_block_sparse_capability_before_implementation( - has_block_sparse_inputs: bool, supported: bool +def test_support_checks_capabilities_before_implementation( + capability: str, request_args: dict, requested: bool, supported: bool ) -> None: - class _BlockSparseFmha(_MinimalFmha): - supports_block_sparse_inputs = True + _CapableFmha = type("_CapableFmha", (_MinimalFmha,), {capability: True}) attn = cast(TrtllmAttention, FakeAttention()) q, k, v = (torch.empty((2, 4)) for _ in range(3)) metadata = Mock(spec=TrtllmAttentionMetadata) - forward_args = AttentionForwardArgs() - if has_block_sparse_inputs: - forward_args.sparse_runtime_params = SparseRuntimeParams( - block_sparse_inputs=BlockSparseForwardInputs( - q_block_size=64, - kv_block_size=64, - exact_block_bits=torch.zeros((1, 1, 1, 1), dtype=torch.int32), - ) - ) + forward_args = AttentionForwardArgs(**request_args) if requested else AttentionForwardArgs() - for fmha in (_MinimalFmha(attn), _BlockSparseFmha(attn)): - capability_supported = not has_block_sparse_inputs or fmha.supports_block_sparse_inputs + for fmha in (_MinimalFmha(attn), _CapableFmha(attn)): + capability_supported = not requested or getattr(fmha, capability) with patch.object(fmha, "_is_supported", return_value=supported) as hook: assert fmha.is_supported(q, k, v, metadata, forward_args) is ( capability_supported and supported diff --git a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py index 2adfac26c11c..071fcc8154f8 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py +++ b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py @@ -21,7 +21,15 @@ import pytest import torch -from backend_case import BackendCase, generate_inputs, run_backend, run_case +from backend_case import ( + BF16_ATOL, + BF16_RTOL, + FP8_ATOL, + BackendCase, + generate_inputs, + run_backend, + run_case, +) from utils.util import getSMVersion pytestmark = pytest.mark.skipif( @@ -1275,6 +1283,137 @@ def test_prims_ts_decode_wrappers_share_workspace_across_serialized_layers() -> torch.testing.assert_close(actual, reference, atol=3e-2, rtol=3e-3) +def _forward_with_recipe(monkeypatch: pytest.MonkeyPatch, recipe) -> None: + """Make the shared harness send separate Q/K/V with the per-tensor recipe flags.""" + import backend_case + + from tensorrt_llm._torch.attention.backends.interface import AttentionForwardArgs + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention + + monkeypatch.setattr(TrtllmAttention, "support_fused_qkv", classmethod(lambda cls: False)) + monkeypatch.setattr( + backend_case, + "AttentionForwardArgs", + functools.partial( + AttentionForwardArgs, + qk_fp8_per_tensor=recipe.qk_dtype == "fp8", + v_fp8_per_tensor=recipe.v_dtype == "fp8", + ), + ) + + +_CACHE_FREE_CASE = dict( + num_heads=8, + num_kv_heads=8, + head_dim=128, + seq_lens=[256, 256], + num_cached_tokens=[0, 0], + num_contexts=2, + dtype="bfloat16", + cache="none", + causal=False, +) + + +@pytest.mark.parametrize("outlier", [1.0, 10.0]) +@pytest.mark.parametrize(("qk_dtype", "v_dtype"), [("bf16", "fp8"), ("fp8", "fp8")]) +def test_prims_ts_cache_free_fp8_recipes( + monkeypatch: pytest.MonkeyPatch, qk_dtype: str, v_dtype: str, outlier: float +) -> None: + """Cache-free recipes plan FP8 operands, fold the per-tensor scales, and match SDPA. + + ``outlier`` multiplies the second sample's Q and V and divides its K, so the two + samples differ in magnitude and set different per-tensor scales while their QK + logits are unchanged. Each sample is checked against SDPA. + """ + from tensorrt_llm._torch.attention.backends.prims_ts.context import BatchPrefillTSWrapper + from tensorrt_llm.visual_gen.args import QuantAttentionConfig + + recipe = QuantAttentionConfig(qk_dtype=qk_dtype, v_dtype=v_dtype) + _forward_with_recipe(monkeypatch, recipe) + plans: list[dict] = [] + runs: list[dict] = [] + plan, run = BatchPrefillTSWrapper.plan, BatchPrefillTSWrapper.run + + def recording_plan(self, *args, **kwargs): + plans.append(kwargs) + return plan(self, *args, **kwargs) + + def recording_run(self, *args, **kwargs): + runs.append(kwargs) + return run(self, *args, **kwargs) + + monkeypatch.setattr(BatchPrefillTSWrapper, "plan", recording_plan) + monkeypatch.setattr(BatchPrefillTSWrapper, "run", recording_run) + monkeypatch.setenv("TLLM_FMHA_LIBS", "+prims_ts") + case = BackendCase(**_CACHE_FREE_CASE) + inputs = generate_inputs(case, seed=0) + second = slice(case.seq_lens[0], None) + inputs["q"][second] *= outlier + inputs["new_k"][second] /= outlier + inputs["new_v"][second] *= outlier + + out = run_backend(case, "TRTLLM", inputs, kv_dtype=torch.bfloat16, kv_layout="HND") + + def bhsd(x: torch.Tensor) -> torch.Tensor: + return x.view(case.num_seqs, -1, case.num_heads, case.head_dim).transpose(1, 2).float() + + golden = torch.nn.functional.scaled_dot_product_attention( + bhsd(inputs["q"]), bhsd(inputs["new_k"]), bhsd(inputs["new_v"]) + ) + golden = golden.transpose(1, 2).reshape_as(out).to(out.dtype) + + fp8, bf16 = torch.float8_e4m3fn, torch.bfloat16 + (plan_kwargs,) = plans + (run_kwargs,) = runs + expected_dtypes = (fp8, fp8, fp8) if qk_dtype == "fp8" else (bf16, bf16, fp8) + assert ( + plan_kwargs["q_dtype"], + plan_kwargs["k_dtype"], + plan_kwargs["v_dtype"], + ) == expected_dtypes + + def dequant_scale(x: torch.Tensor) -> torch.Tensor: + return (x.float().abs().amax() / 448.0).reshape(1) + + # The quantize op returns BF16-rounded dequant scales. + torch.testing.assert_close( + run_kwargs["output_scale"], dequant_scale(inputs["new_v"]), rtol=1e-2, atol=0.0 + ) + if qk_dtype == "fp8": + expected = ( + dequant_scale(inputs["q"]) + * dequant_scale(inputs["new_k"]) + * (math.log2(math.e) / math.sqrt(case.head_dim)) + ) + torch.testing.assert_close(run_kwargs["scale_softmax_log2"], expected, rtol=1e-2, atol=0.0) + else: + assert run_kwargs["scale_softmax_log2"] is None + # The second sample's output is ``outlier`` times larger. Normalize it before comparing. + sample_gain = torch.tensor([1.0, outlier], device=out.device).view(-1, 1) + torch.testing.assert_close( + out.view(case.num_seqs, -1) / sample_gain, + golden.view(case.num_seqs, -1) / sample_gain, + atol=FP8_ATOL + BF16_ATOL, + rtol=BF16_RTOL, + ) + + +def test_prims_ts_recipe_requires_prims_ts_library(monkeypatch: pytest.MonkeyPatch) -> None: + """A per-tensor recipe fails instead of running unquantized on another FMHA library.""" + from tensorrt_llm.visual_gen.args import QuantAttentionConfig + + recipe = QuantAttentionConfig(qk_dtype="bf16", v_dtype="fp8") + _forward_with_recipe(monkeypatch, recipe) + monkeypatch.delenv("TLLM_FMHA_LIBS", raising=False) + case = BackendCase(**_CACHE_FREE_CASE) + + with pytest.raises(RuntimeError, match="needs the PrimTS FMHA library"): + run_backend( + case, "TRTLLM", generate_inputs(case, seed=0), kv_dtype=torch.bfloat16, kv_layout="HND" + ) + + def test_prims_ts_unsupported_context_falls_back(monkeypatch: pytest.MonkeyPatch) -> None: from tensorrt_llm._torch.attention.backends.fmha.fallback import FallbackFmha from tensorrt_llm._torch.attention.backends.fmha.prims_ts import PrimsTSFmha diff --git a/tests/unittest/_torch/attention/test_prims_ts_fmha.py b/tests/unittest/_torch/attention/test_prims_ts_fmha.py index 8902890021ce..8d2f771b88b2 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_fmha.py +++ b/tests/unittest/_torch/attention/test_prims_ts_fmha.py @@ -2331,3 +2331,54 @@ def test_phased_forward_routes_query_and_qkv_by_phase( assert params.value_input is None assert params.output.shape == (params.num_tokens, attn.num_heads, out_head_size) assert params.output.data_ptr() == output[token_slice].data_ptr() + + +def test_cache_free_context_plans_contiguous_wrapper_once(monkeypatch: pytest.MonkeyPatch) -> None: + """The cache-free context path plans one contiguous wrapper per shape and reuses it.""" + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + attn = _Attention() + fmha = PrimsTSFmha(attn) + wrapper = Mock() + wrapper_factory = Mock(return_value=wrapper) + monkeypatch.setattr(prims_context_module, "BatchPrefillTSWrapper", wrapper_factory) + batch_size, seq_len = 2, 16 + q = torch.empty((batch_size * seq_len, attn.num_heads * attn.head_dim), dtype=torch.bfloat16) + k = torch.empty((batch_size * seq_len, attn.num_kv_heads * attn.head_dim), dtype=torch.bfloat16) + v = torch.empty_like(k) + output = torch.empty( + (batch_size * seq_len, attn.num_heads * attn.head_dim), dtype=torch.bfloat16 + ) + metadata = SimpleNamespace( + kv_cache_manager=None, + num_contexts=batch_size, + num_generations=0, + max_seq_len=seq_len, + prompt_lens_cpu_runtime=torch.full((batch_size,), seq_len, dtype=torch.int32), + ) + forward_args = AttentionForwardArgs( + output=output, + attention_input_type=AttentionInputType.context_only, + attention_mask=PredefinedAttentionMask.FULL, + ) + + fmha.forward(q, k, v, metadata, forward_args) + fmha.forward(q, k, v, metadata, forward_args) + + wrapper_factory.assert_called_once_with() + wrapper.plan.assert_called_once() + plan_kwargs = wrapper.plan.call_args.kwargs + assert plan_kwargs["batch_size"] == batch_size + assert plan_kwargs["max_seq_len_q"] == plan_kwargs["max_kv_len"] == seq_len + assert ( + plan_kwargs["q_dtype"] == plan_kwargs["k_dtype"] == plan_kwargs["v_dtype"] == torch.bfloat16 + ) + assert plan_kwargs["mask_type"] == "dense" + assert wrapper.run.call_count == 2 + run_args, run_kwargs = wrapper.run.call_args + assert tuple(run_args[0].shape) == (batch_size, seq_len, attn.num_heads, attn.head_dim) + assert tuple(run_args[2].shape) == (batch_size, seq_len, attn.num_kv_heads, attn.head_dim) + assert run_args[0].data_ptr() == q.data_ptr() + assert run_args[2].data_ptr() == v.data_ptr() + assert run_kwargs["out"].data_ptr() == output.data_ptr() + assert run_kwargs["scale_softmax_log2"] is None + assert run_kwargs["output_scale"] is None diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py index e90434e29bce..4c46a064dc47 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py @@ -162,9 +162,13 @@ def test_skip_softmax_and_sage_quantization_can_be_combined(self): (100, "fp8", 1, 4, 1), (103, "fp8", 1, 1, 1), (103, "fp8", 1, 4, 1), + (100, "bf16", 0, 0, 0), + (100, "fp8", 0, 0, 0), + (103, "bf16", 0, 0, 0), + (103, "fp8", 0, 0, 0), ], ) - def test_supported_quant_config_sage( + def test_supported_quant_config_trtllm( self, sm_ver, qk_dtype, q_block_size, k_block_size, v_block_size ): with patch("tensorrt_llm.visual_gen.args.get_sm_version", return_value=sm_ver): @@ -180,15 +184,16 @@ def test_supported_quant_config_sage( assert attention.quant_attention_config is not None - @pytest.mark.parametrize("backend", ["CUTEDSL", "CUDNN", "FLASHINFER", "VANILLA"]) - def test_smooth_k_rejected_on_non_trtllm_backend(self, backend): - with pytest.raises(ValidationError, match="smooth_k is a SageAttention option"): - AttentionConfig( - backend=backend, - quant_attention_config=QuantAttentionConfig( - qk_dtype="bf16", v_dtype="fp8", smooth_k=True - ), - ) + @pytest.mark.parametrize("backend", ["TRTLLM", "CUTEDSL", "CUDNN", "FLASHINFER", "VANILLA"]) + def test_smooth_k_rejected_without_sage_recipe(self, backend): + with patch("tensorrt_llm.visual_gen.args.get_sm_version", return_value=100): + with pytest.raises(ValidationError, match="smooth_k is a SageAttention option"): + AttentionConfig( + backend=backend, + quant_attention_config=QuantAttentionConfig( + qk_dtype="bf16", v_dtype="fp8", smooth_k=True + ), + ) def test_supported_quant_config_cute(self): attention = AttentionConfig(