From e76b07423f61cfa23425a3ac8bd0ac663c992538 Mon Sep 17 00:00:00 2001 From: Harrison Zhang Date: Thu, 8 Oct 2026 14:53:22 -0700 Subject: [PATCH 01/11] [primsTS][fmha][visual_gen] sync vendored PrimTS context FMHA with flashinfer harrisonzhy/primsts_fmha_perf-bringup_trtllm-pin (910c3ed3) Pin the perf-bringup context kernels (P-half overlap, lazy correction, two-CTA UMMA, fp8 P in SMEM, deeper K/V ring) on the current TRT-LLM pin lineage; decode, MLA and block-sparse are unchanged. Route cache-free VisualGen attention to the contiguous PrimTS wrapper and add quant_attention_config.algorithm (primsts, sage) with the PrimTS per-tensor FP8 recipes. Move the paged context plan() call to the k_dtype/v_dtype signature. Signed-off-by: Harrison Zhang --- .../features/visualgen-quantized-attention.md | 3 + .../attention/backends/fmha/prims_ts.py | 180 +++++++++++++++++- .../_torch/attention/backends/trtllm.py | 5 + .../visual_gen/attention/backends/trtllm.py | 12 +- tensorrt_llm/visual_gen/args.py | 41 +++- .../_torch/attention/test_prims_ts_fmha.py | 50 +++++ 6 files changed, 278 insertions(+), 13 deletions(-) diff --git a/docs/source/features/visualgen-quantized-attention.md b/docs/source/features/visualgen-quantized-attention.md index 9d4f33a061e8..b08238ecad60 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 (`algorithm: primsts`, `TLLM_FMHA_LIBS=+prims_ts`) | +| `TRTLLM` | `fp8` | `fp8` | `(0, 0, 0)` | PrimTS FP8 for Blackwell (`algorithm: primsts`, `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 | @@ -62,6 +64,7 @@ After a recipe meets the quality target, benchmark its end-to-end throughput wit | `q_block_size` | int ≥ 0 | `0` | Q tokens per SageAttention quantization block. `0` outside SageAttention. | | `k_block_size` | int ≥ 0 | `0` | K tokens per SageAttention quantization block. `0` outside SageAttention. | | `v_block_size` | int ≥ 0 | `0` | V block size on the hidden dimension. `0` = one tensor-wide V scale; `1` = one scale per channel. Keep `0` if `v_dtype` defines its own scaling format (e.g., `mxfp8` or `nvfp4`) | +| `algorithm` | `"primsts" \| "sage" \| "vc_attention-qk16"` | `"sage"` | Kernel family serving the recipe on the `TRTLLM` backend. `sage` quantizes per block; `primsts` quantizes per tensor and `vc_attention-qk16` adds the VC-Attention V treatment, both served by the `prims_ts` FMHA library. | Routing (`tensorrt_llm/_torch/visual_gen/attention/backends/utils.py`) forwards the validated `quant_attention_config` into the backend constructor: `TrtllmAttention` for `TRTLLM`, `CuDNNAttention` for `CUDNN`, `FlashInferAttention` for `FLASHINFER`, and the dense `CuTeDSLAttention` FMHA backend for `CUTEDSL`. diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py index 2e67616b5e9d..8a760de3bdfb 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, @@ -173,6 +176,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 @@ -294,10 +299,6 @@ def _is_supported_with_reason( 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 +319,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, 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 +463,68 @@ 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, + attn: "TrtllmAttention", + meta: "TrtllmAttentionMetadata", + fwd: AttentionForwardArgs, + ) -> tuple[bool, str]: + """Support decision for fused-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: + return False, f"query dtype {q.dtype} is unsupported." + recipe = attn.quant_attention_config + if recipe is not None and recipe.algorithm == "vc_attention-qk16": + return False, "VC-Attention kernels are not available in this build." + if recipe is not None and recipe.qk_dtype == "bf16" 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: + return False, f"fused attention input must be rank 2, got rank {q.ndim}." + 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]}." + 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 +798,108 @@ 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: + if metadata.kv_cache_manager is None and not int(metadata.num_generations): + self._forward_without_kv_cache(q, metadata, forward_args) + return + super().forward(q, k, v, metadata, forward_args) + + def _forward_without_kv_cache( + self, + q: torch.Tensor, + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + ) -> None: + """Fused-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) + num_heads, num_kv_heads, head_dim = attn.num_heads, attn.num_kv_heads, attn.head_dim + q_width, kv_width = num_heads * head_dim, num_kv_heads * head_dim + qkv = q.view(batch_size, seq_len, q_width + 2 * kv_width) + # The fixed contiguous plan takes compact [B, S, H, D] operands. The fused-QKV + # column slices are strided, so each one is copied out. + q_bshd = qkv[..., :q_width].reshape(batch_size, seq_len, num_heads, head_dim).contiguous() + k_bshd = ( + qkv[..., q_width : q_width + kv_width] + .reshape(batch_size, seq_len, num_kv_heads, head_dim) + .contiguous() + ) + v_bshd = ( + qkv[..., q_width + kv_width :] + .reshape(batch_size, seq_len, num_kv_heads, head_dim) + .contiguous() + ) + # Per-tensor E4M3 operands as selected by the layer's attention 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. + recipe = attn.quant_attention_config + sm_scale = self._get_bmm1_scale(attn) + scale_softmax_log2 = None + if recipe is not None and recipe.qk_dtype == "fp8": + 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 * k_dequant_scale * (sm_scale * math.log2(math.e))) + .float() + .reshape(1) + ) + output_scale = None + if recipe is not None and recipe.v_dtype == "fp8": + 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/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index d29def811bc5..7e71998d3d9e 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -26,6 +26,7 @@ if TYPE_CHECKING: from tensorrt_llm.mapping import Mapping + from tensorrt_llm.visual_gen.args import QuantAttentionConfig from ...model_config import ModelConfig from ...speculative.interface import SpecMetadata @@ -1696,6 +1697,7 @@ def __init__( skip_create_weights_in_init: bool = False, attention_chunk_size: Optional[int] = None, sparse_params: Optional[SparseParams] = None, + quant_attention_config: Optional["QuantAttentionConfig"] = None, kv_cache_dtype: str = "auto", skip_correction_threshold: float = 0.0, **kwargs, @@ -1716,6 +1718,8 @@ def __init__( sparse_params (SparseParams): Optional sparse-attention backend parameters (e.g. skip-softmax). Algorithm-specific fields are documented on the corresponding ``SparseParams`` subclass. + quant_attention_config (QuantAttentionConfig): Optional attention operand + quantization recipe (Q/K and V dtypes with their block sizes). kv_cache_dtype (str): KV-cache dtype selected by ``KvCacheConfig``. Accepted values are ``auto``, ``fp8``, ``fp8_ds_mla``, ``nvfp4``, and supported torch dtype strings. ``fp8_ds_mla`` selects the packed sparse-MLA cache @@ -1726,6 +1730,7 @@ def __init__( super().__init__(layer_idx, num_heads, head_dim, num_kv_heads, quant_config, **kwargs) self.sparse_params = sparse_params + self.quant_attention_config = quant_attention_config self.kv_cache_dtype = kv_cache_dtype self.use_fp8_ds_mla = kv_cache_dtype == "fp8_ds_mla" self.is_mla_enable = mla_params is not None diff --git a/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py b/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py index 3502ca6f9a4a..514e5ef76323 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 algorithm primsts, fused QKV) - Separate-QKV forwarding for generic block-sparse attention and backends that reject fused QKV """ @@ -222,6 +223,7 @@ def __init__( head_dim=head_dim, quant_config=quant_config, sparse_params=sparse_params, + quant_attention_config=quant_attention_config, dtype=dtype, ) @@ -232,8 +234,6 @@ def __init__( attention_metadata_state=attention_metadata_state, ) - self.quant_attention_config = quant_attention_config - @property def timestep_cutoff(self) -> Optional[float]: """Normalized timestep below which the sparse algorithm is enabled, if any.""" @@ -376,7 +376,7 @@ def forward( # self-attention from fused QKV only. use_separate_qkv = ( block_sparse_inputs is not None - or self.quant_attention_config is not None + or (self.quant_attention_config is not None and not self.support_fused_qkv()) or (not self.support_fused_qkv() and self.should_use_sparse(sparse_attn_phase)) ) if use_separate_qkv and (k is None or v is None): @@ -428,5 +428,7 @@ def preferred_layout(self) -> AttentionTensorLayout: return self._preferred_layout def support_fused_qkv(self) -> bool: - """Standard path fuses QKV; SageAttention path does not.""" - return self.quant_attention_config is None + """Standard and PrimTS paths fuse QKV. SageAttention does not.""" + return ( + self.quant_attention_config is None or self.quant_attention_config.algorithm != "sage" + ) diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index dc9c0e30184f..eec6c5655de7 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -103,6 +103,14 @@ class QuantAttentionConfig(StrictBaseModel): "quantized type has to cover. SageAttention only." ), ) + algorithm: Literal["primsts", "sage", "vc_attention-qk16"] = Field( + "sage", + status="prototype", + description=( + "TRTLLM kernel family for the recipe. primsts quantizes per tensor, sage per block, " + "vc_attention-qk16 adds the VC-Attention V treatment (both via TLLM_FMHA_LIBS=+prims_ts)." + ), + ) # Discriminated union of sparse attention configs. @@ -158,6 +166,23 @@ def _validate_quant_attention_config(self) -> "AttentionConfig": ("fp8", "fp8", (1, 4, 1)), }, } + # PrimTS supports per-tensor FP8 recipes on Blackwell, with or without the VC-Attention V + # treatment. + PRIMSTS_RECIPES = { + 100: { + ("bf16", "fp8", (0, 0, 0)), + ("fp8", "fp8", (0, 0, 0)), + }, + 103: { + ("bf16", "fp8", (0, 0, 0)), + ("fp8", "fp8", (0, 0, 0)), + }, + } + TRTLLM_RECIPES = { + "primsts": PRIMSTS_RECIPES, + "sage": SAGE_RECIPES, + "vc_attention-qk16": PRIMSTS_RECIPES, + } # Other recipes verify the hardware at corresponding backend implementations. CUDNN_RECIPES = { ("fp8", "fp8", (0, 0, 0)), @@ -186,13 +211,23 @@ 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()) + recipes = TRTLLM_RECIPES[q_config.algorithm].get(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"{sorted(recipes)}." + f"for backend='TRTLLM'. Supported {q_config.algorithm} recipes on this " + f"device: {sorted(recipes)}." ) + if q_config.smooth_k and q_config.algorithm != "sage": + raise ValueError( + f"smooth_k is a SageAttention option and does not apply to " + f"algorithm='{q_config.algorithm}'." + ) + elif q_config.algorithm != "sage": + raise ValueError( + f"algorithm='{q_config.algorithm}' selects a TRTLLM kernel family and requires " + f"backend='TRTLLM', got backend='{self.backend}'." + ) 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_prims_ts_fmha.py b/tests/unittest/_torch/attention/test_prims_ts_fmha.py index 8902890021ce..16e92145755b 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_fmha.py +++ b/tests/unittest/_torch/attention/test_prims_ts_fmha.py @@ -88,6 +88,7 @@ def __init__( self.v_head_dim = 128 if is_mla else None self.predicted_tokens_per_seq = 1 self.sparse_params = None + self.quant_attention_config = None self.skip_correction_threshold = 0.0 self.position_embedding_type = 0 self.quant_mode = 0 @@ -2331,3 +2332,52 @@ 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: + 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 + 2 * attn.num_kv_heads) * attn.head_dim), + dtype=torch.bfloat16, + ) + 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, None, None, metadata, forward_args) + fmha.forward(q, None, None, 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_kwargs["out"].data_ptr() == output.data_ptr() + assert run_kwargs["scale_softmax_log2"] is None + assert run_kwargs["output_scale"] is None From 8238809514aa077abc08cbad2431d1dfcacca3b0 Mon Sep 17 00:00:00 2001 From: Harrison Zhang Date: Thu, 8 Oct 2026 16:52:07 -0700 Subject: [PATCH 02/11] [primsTS][fmha][visual_gen] require the PrimTS library for primsts recipes and cover the FP8 path Raise when a primsts recipe would be served by another FMHA library instead of running unquantized. Limit quant_attention_config.algorithm to primsts and sage. Add a cache-free FP8 recipe test (plan dtypes, folded scales, SDPA match) and the library check test. Signed-off-by: Harrison Zhang (cherry picked from commit 47bacd7dcedabb5c0e809179f34c73f94a9f7b4b) --- .../features/visualgen-quantized-attention.md | 2 +- .../attention/backends/fmha/prims_ts.py | 3 +- .../_torch/attention/backends/trtllm.py | 8 ++ tensorrt_llm/visual_gen/args.py | 15 +-- .../test_prims_ts_attention_backend.py | 120 +++++++++++++++++- .../_torch/attention/test_prims_ts_fmha.py | 1 + .../_torch/visual_gen/test_visual_gen_args.py | 14 ++ 7 files changed, 149 insertions(+), 14 deletions(-) diff --git a/docs/source/features/visualgen-quantized-attention.md b/docs/source/features/visualgen-quantized-attention.md index b08238ecad60..4698a0b232bc 100644 --- a/docs/source/features/visualgen-quantized-attention.md +++ b/docs/source/features/visualgen-quantized-attention.md @@ -64,7 +64,7 @@ After a recipe meets the quality target, benchmark its end-to-end throughput wit | `q_block_size` | int ≥ 0 | `0` | Q tokens per SageAttention quantization block. `0` outside SageAttention. | | `k_block_size` | int ≥ 0 | `0` | K tokens per SageAttention quantization block. `0` outside SageAttention. | | `v_block_size` | int ≥ 0 | `0` | V block size on the hidden dimension. `0` = one tensor-wide V scale; `1` = one scale per channel. Keep `0` if `v_dtype` defines its own scaling format (e.g., `mxfp8` or `nvfp4`) | -| `algorithm` | `"primsts" \| "sage" \| "vc_attention-qk16"` | `"sage"` | Kernel family serving the recipe on the `TRTLLM` backend. `sage` quantizes per block; `primsts` quantizes per tensor and `vc_attention-qk16` adds the VC-Attention V treatment, both served by the `prims_ts` FMHA library. | +| `algorithm` | `"primsts" \| "sage"` | `"sage"` | Kernel family serving the recipe on the `TRTLLM` backend. `sage` quantizes per block; `primsts` quantizes per tensor and is served by the `prims_ts` FMHA library. | Routing (`tensorrt_llm/_torch/visual_gen/attention/backends/utils.py`) forwards the validated `quant_attention_config` into the backend constructor: `TrtllmAttention` for `TRTLLM`, `CuDNNAttention` for `CUDNN`, `FlashInferAttention` for `FLASHINFER`, and the dense `CuTeDSLAttention` FMHA backend for `CUTEDSL`. diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py index 8a760de3bdfb..05f016c9aa86 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py @@ -495,8 +495,6 @@ def _is_supported_without_kv_cache( if q.dtype not in self.SUPPORTED_DTYPES: return False, f"query dtype {q.dtype} is unsupported." recipe = attn.quant_attention_config - if recipe is not None and recipe.algorithm == "vc_attention-qk16": - return False, "VC-Attention kernels are not available in this build." if recipe is not None and recipe.qk_dtype == "bf16" and q.dtype != torch.bfloat16: return False, f"the attention recipe keeps Q/K in BF16, got {q.dtype}." output = fwd.output @@ -806,6 +804,7 @@ def forward( 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, metadata, forward_args) return diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 7e71998d3d9e..6ce4194a9ad8 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -44,6 +44,7 @@ from ...utils import (compute_swizzled_sf_shape, get_global_attrs, get_model_extra_attrs, helix_local_len_tensor) from .fmha.manager import FmhaManager +from .fmha.prims_ts import PrimsTSFmha from .fp4_mla import can_fuse_fp4_mla_q_quant, scatter_fp4_mla_kv_cache from .fp4_mla.state import Fp4MlaState from .interface import (AttentionBackend, AttentionForwardArgs, @@ -2472,6 +2473,13 @@ def forward( if fmha is None: raise RuntimeError( "No TRT-LLM attention FMHA library supports this request.") + recipe = self.quant_attention_config + if recipe is not None and recipe.algorithm != "sage" and not isinstance( + fmha, PrimsTSFmha): + raise RuntimeError( + f"The {recipe.algorithm} attention recipe needs the PrimTS FMHA " + f"library (TLLM_FMHA_LIBS=+prims_ts), got " + f"{type(fmha).__name__}.") if metadata.is_cuda_graph or not fmha.supports_workspace_reclamation: # Conservatively disable reclamation for metadata used by graphs # or backends that can retain staged workspace state. diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index eec6c5655de7..4ad1c68becb6 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -103,12 +103,12 @@ class QuantAttentionConfig(StrictBaseModel): "quantized type has to cover. SageAttention only." ), ) - algorithm: Literal["primsts", "sage", "vc_attention-qk16"] = Field( + algorithm: Literal["primsts", "sage"] = Field( "sage", status="prototype", description=( - "TRTLLM kernel family for the recipe. primsts quantizes per tensor, sage per block, " - "vc_attention-qk16 adds the VC-Attention V treatment (both via TLLM_FMHA_LIBS=+prims_ts)." + "TRTLLM kernel family for the recipe. primsts quantizes per tensor " + "(TLLM_FMHA_LIBS=+prims_ts), sage per block." ), ) @@ -166,8 +166,7 @@ def _validate_quant_attention_config(self) -> "AttentionConfig": ("fp8", "fp8", (1, 4, 1)), }, } - # PrimTS supports per-tensor FP8 recipes on Blackwell, with or without the VC-Attention V - # treatment. + # PrimTS supports per-tensor FP8 recipes on Blackwell. PRIMSTS_RECIPES = { 100: { ("bf16", "fp8", (0, 0, 0)), @@ -178,11 +177,7 @@ def _validate_quant_attention_config(self) -> "AttentionConfig": ("fp8", "fp8", (0, 0, 0)), }, } - TRTLLM_RECIPES = { - "primsts": PRIMSTS_RECIPES, - "sage": SAGE_RECIPES, - "vc_attention-qk16": PRIMSTS_RECIPES, - } + TRTLLM_RECIPES = {"primsts": PRIMSTS_RECIPES, "sage": SAGE_RECIPES} # Other recipes verify the hardware at corresponding backend implementations. CUDNN_RECIPES = { ("fp8", "fp8", (0, 0, 0)), 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..dbba767933ed 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,116 @@ def test_prims_ts_decode_wrappers_share_workspace_across_serialized_layers() -> torch.testing.assert_close(actual, reference, atol=3e-2, rtol=3e-3) +def _build_trtllm_attention_with_recipe(monkeypatch: pytest.MonkeyPatch, recipe) -> None: + """Make the shared harness pass the attention recipe to the TRTLLM layer it builds.""" + import backend_case + + create_attention = backend_case.create_attention + + def create_attention_with_recipe(backend, **kwargs): + attn = create_attention(backend, **kwargs) + if backend == "TRTLLM": + attn.quant_attention_config = recipe + return attn + + monkeypatch.setattr(backend_case, "create_attention", create_attention_with_recipe) + + +_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(("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 +) -> None: + """Cache-free recipes plan FP8 operands, fold the per-tensor scales, and match 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, algorithm="primsts") + _build_trtllm_attention_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) + + 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 + torch.testing.assert_close(out, golden, atol=FP8_ATOL + BF16_ATOL, rtol=BF16_RTOL) + + +def test_prims_ts_recipe_requires_prims_ts_library(monkeypatch: pytest.MonkeyPatch) -> None: + """A PrimTS 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", algorithm="primsts") + _build_trtllm_attention_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 16e92145755b..e508d469d24c 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_fmha.py +++ b/tests/unittest/_torch/attention/test_prims_ts_fmha.py @@ -2335,6 +2335,7 @@ def test_phased_forward_routes_query_and_qkv_by_phase( 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) 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..9ce92edede1b 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py @@ -180,6 +180,20 @@ def test_supported_quant_config_sage( assert attention.quant_attention_config is not None + @pytest.mark.parametrize("sm_ver", [100, 103]) + @pytest.mark.parametrize("qk_dtype", ["bf16", "fp8"]) + def test_supported_quant_config_primsts(self, sm_ver, qk_dtype): + """PrimTS per-tensor recipes validate on Blackwell for both Q/K dtypes.""" + with patch("tensorrt_llm.visual_gen.args.get_sm_version", return_value=sm_ver): + attention = AttentionConfig( + backend="TRTLLM", + quant_attention_config=QuantAttentionConfig( + qk_dtype=qk_dtype, v_dtype="fp8", algorithm="primsts" + ), + ) + + assert attention.quant_attention_config.algorithm == "primsts" + @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"): From 9cb19babbbb7807f4dc20d1f868b9ab251d458b2 Mon Sep 17 00:00:00 2001 From: Harrison Zhang Date: Mon, 12 Oct 2026 11:58:10 -0700 Subject: [PATCH 03/11] [primsTS][fmha] tolerate attention backends without quant_attention_config Read the attention recipe with getattr in TrtllmAttention.forward and the PrimTS adapter so partially constructed backends and spec mocks in the attention unit tests keep working. Signed-off-by: Harrison Zhang (cherry picked from commit c89ba1c2076a973073b7339ac6fa59fb765751f2) --- tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py | 4 ++-- tensorrt_llm/_torch/attention/backends/trtllm.py | 3 ++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py index 05f016c9aa86..730ed3adc33d 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py @@ -494,7 +494,7 @@ def _is_supported_without_kv_cache( return False, f"context head dimension {attn.head_dim} is unsupported." if q.dtype not in self.SUPPORTED_DTYPES: return False, f"query dtype {q.dtype} is unsupported." - recipe = attn.quant_attention_config + recipe = getattr(attn, "quant_attention_config", None) if recipe is not None and recipe.qk_dtype == "bf16" and q.dtype != torch.bfloat16: return False, f"the attention recipe keeps Q/K in BF16, got {q.dtype}." output = fwd.output @@ -842,7 +842,7 @@ def _forward_without_kv_cache( # Per-tensor E4M3 operands as selected by the layer's attention 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. - recipe = attn.quant_attention_config + recipe = getattr(attn, "quant_attention_config", None) sm_scale = self._get_bmm1_scale(attn) scale_softmax_log2 = None if recipe is not None and recipe.qk_dtype == "fp8": diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 6ce4194a9ad8..6828db04090b 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -2473,7 +2473,8 @@ def forward( if fmha is None: raise RuntimeError( "No TRT-LLM attention FMHA library supports this request.") - recipe = self.quant_attention_config + # Backends built without __init__ (test doubles) carry no recipe. + recipe = getattr(self, "quant_attention_config", None) if recipe is not None and recipe.algorithm != "sage" and not isinstance( fmha, PrimsTSFmha): raise RuntimeError( From 89c8f0e900a2a41a714e5e1f253be0c008211a67 Mon Sep 17 00:00:00 2001 From: Harrison Zhang Date: Mon, 12 Oct 2026 15:40:14 -0700 Subject: [PATCH 04/11] [primsTS][fmha] fail early on a missing PrimTS library and re-check uniform lengths Raise in TrtllmAttention.update_quant_config when the primsts recipe is set but the PrimTS FMHA library is not enabled or available, instead of on the first forward. Re-check uniform context lengths in the cache-free PrimTS forward, since the FMHA selection cache key carries no sequence lengths. Cover a batch with samples of different magnitude in the cache-free FP8 recipe test. Signed-off-by: Harrison Zhang (cherry picked from commit 64c7e4bdd0a36117311a2fd506ae678297381921) --- .../attention/backends/fmha/prims_ts.py | 5 ++++ .../_torch/attention/backends/trtllm.py | 11 ++++++++- .../test_prims_ts_attention_backend.py | 23 ++++++++++++++++--- 3 files changed, 35 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py index 730ed3adc33d..a8847175df7b 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py @@ -823,6 +823,11 @@ def _forward_without_kv_cache( 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 q_width, kv_width = num_heads * head_dim, num_kv_heads * head_dim qkv = q.view(batch_size, seq_len, q_width + 2 * kv_width) diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 6828db04090b..54cf99c5494b 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -1829,6 +1829,15 @@ def update_quant_config(self, new_quant_config: Optional[QuantConfig]): self.has_w4a8_nvfp4_fp8 = self.quant_config.layer_quant_mode.has_w4a8_nvfp4_fp8( ) self._fmha_manager = FmhaManager(self) + recipe = getattr(self, "quant_attention_config", None) + has_prims_ts = any( + isinstance(fmha, PrimsTSFmha) + for fmha in self._fmha_manager.fmha_libs) + if recipe is not None and recipe.algorithm == "primsts" and not has_prims_ts: + raise RuntimeError( + f"The {recipe.algorithm} attention recipe needs the PrimTS FMHA " + "library (TLLM_FMHA_LIBS=+prims_ts), which is not enabled or " + "not available on this platform.") @classmethod def runtime_workspace_bytes_per_token(cls, model_config: "ModelConfig", @@ -2475,7 +2484,7 @@ def forward( "No TRT-LLM attention FMHA library supports this request.") # Backends built without __init__ (test doubles) carry no recipe. recipe = getattr(self, "quant_attention_config", None) - if recipe is not None and recipe.algorithm != "sage" and not isinstance( + if recipe is not None and recipe.algorithm == "primsts" and not isinstance( fmha, PrimsTSFmha): raise RuntimeError( f"The {recipe.algorithm} attention recipe needs the PrimTS FMHA " 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 dbba767933ed..e9688574d4fa 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py +++ b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py @@ -1311,11 +1311,17 @@ def create_attention_with_recipe(backend, **kwargs): ) +@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 + 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.""" + """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 @@ -1338,6 +1344,10 @@ def recording_run(self, *args, **kwargs): 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") @@ -1375,7 +1385,14 @@ def dequant_scale(x: torch.Tensor) -> torch.Tensor: 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 - torch.testing.assert_close(out, golden, atol=FP8_ATOL + BF16_ATOL, rtol=BF16_RTOL) + # 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: From cbf8f319f228a7154506c5c2b40e769f8f36f845 Mon Sep 17 00:00:00 2001 From: Harrison Zhang Date: Tue, 13 Oct 2026 09:55:52 -0700 Subject: [PATCH 05/11] [primsTS][fmha] point PrimTS vendor source to flashinfer harrisonzhy/trtllm-prims-ts-dev_perf-bringup-20261012 (ae312f7c) The perf-bringup FlashInfer PRs (#4879, #5280, #5475, #5900) are all in the new canonical source trtllm-prims-ts-dev-20261012 @ ae312f7cc, so the temporary source branch is that commit with nothing on top. Lock-only change, vendored tree and compat patch stay identical to main. Signed-off-by: Harrison Zhang --- 3rdparty/vendor_sources.lock.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 2e914da8f052..3045a99ba5a5 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -1,8 +1,8 @@ schema_version: 1 vendors: flashinfer-prims-ts: - url: https://github.com/yuxianq/flashinfer.git - branch: trtllm-prims-ts-dev-20261012 + url: https://github.com/harrisonzhy/flashinfer.git + branch: harrisonzhy/trtllm-prims-ts-dev_perf-bringup-20261012 commit: ae312f7cc1723b06fda5168105ade0f328ee47d8 source: flashinfer/attention/prims_ts destination: tensorrt_llm/_torch/attention/backends/prims_ts From 3c26d01ce164f34d39f0a597b3c124ac031fb4bc Mon Sep 17 00:00:00 2001 From: Harrison Zhang Date: Tue, 13 Oct 2026 22:05:00 +0000 Subject: [PATCH 06/11] [primsTS][fmha] retrigger the vendor status check Signed-off-by: Harrison Zhang From d8d1e4c4325b5c50a20c1200c76ababa9781c13d Mon Sep 17 00:00:00 2001 From: Harrison Zhang Date: Tue, 13 Oct 2026 16:56:18 -0700 Subject: [PATCH 07/11] [primsTS][fmha][visual_gen] select PrimTS from the recipe block sizes and take separate Q/K/V Drop QuantAttentionConfig.algorithm. The TRTLLM backend validates the recipe against the SAGE and PrimTS tables together, and the zero block sizes select the per-tensor PrimTS recipes. The VisualGen backend sends every recipe down the separate-QKV path and carries the per-tensor FP8 flags in the forward args like the SAGE ones, so the cache-free adapter views Q, K and V without the fused concat and split, and FMHA libraries without the supports_per_tensor_fp8 capability reject the request. The softmax scale product is formed in fp32. Signed-off-by: Harrison Zhang --- .../features/visualgen-quantized-attention.md | 7 +- .../attention/backends/fmha/interface.py | 6 ++ .../attention/backends/fmha/prims_ts.py | 87 ++++++++++--------- .../_torch/attention/backends/interface.py | 4 + .../_torch/attention/backends/trtllm.py | 29 ++----- .../visual_gen/attention/backends/trtllm.py | 26 ++++-- tensorrt_llm/visual_gen/args.py | 29 ++----- .../_torch/attention/test_fmha_interface.py | 42 +++++---- .../test_prims_ts_attention_backend.py | 34 ++++---- .../_torch/attention/test_prims_ts_fmha.py | 14 +-- .../_torch/visual_gen/test_visual_gen_args.py | 39 ++++----- 11 files changed, 159 insertions(+), 158 deletions(-) diff --git a/docs/source/features/visualgen-quantized-attention.md b/docs/source/features/visualgen-quantized-attention.md index 4698a0b232bc..7f81a62d2143 100644 --- a/docs/source/features/visualgen-quantized-attention.md +++ b/docs/source/features/visualgen-quantized-attention.md @@ -29,8 +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 (`algorithm: primsts`, `TLLM_FMHA_LIBS=+prims_ts`) | -| `TRTLLM` | `fp8` | `fp8` | `(0, 0, 0)` | PrimTS FP8 for Blackwell (`algorithm: primsts`, `TLLM_FMHA_LIBS=+prims_ts`) | +| `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 | @@ -40,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. @@ -64,7 +66,6 @@ After a recipe meets the quality target, benchmark its end-to-end throughput wit | `q_block_size` | int ≥ 0 | `0` | Q tokens per SageAttention quantization block. `0` outside SageAttention. | | `k_block_size` | int ≥ 0 | `0` | K tokens per SageAttention quantization block. `0` outside SageAttention. | | `v_block_size` | int ≥ 0 | `0` | V block size on the hidden dimension. `0` = one tensor-wide V scale; `1` = one scale per channel. Keep `0` if `v_dtype` defines its own scaling format (e.g., `mxfp8` or `nvfp4`) | -| `algorithm` | `"primsts" \| "sage"` | `"sage"` | Kernel family serving the recipe on the `TRTLLM` backend. `sage` quantizes per block; `primsts` quantizes per tensor and is served by the `prims_ts` FMHA library. | Routing (`tensorrt_llm/_torch/visual_gen/attention/backends/utils.py`) forwards the validated `quant_attention_config` into the backend constructor: `TrtllmAttention` for `TRTLLM`, `CuDNNAttention` for `CUDNN`, `FlashInferAttention` for `FLASHINFER`, and the dense `CuTeDSLAttention` FMHA backend for `CUTEDSL`. diff --git a/tensorrt_llm/_torch/attention/backends/fmha/interface.py b/tensorrt_llm/_torch/attention/backends/fmha/interface.py index 2d1bd5c86742..71edcf451796 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,11 @@ def is_supported( ): logger.debug(f"{type(self).__name__} does not support block-sparse inputs.") return False + if ( + forward_args.qk_fp8_per_tensor or forward_args.v_fp8_per_tensor + ) 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 a8847175df7b..5268e72b5bd6 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py @@ -161,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} @@ -292,10 +294,11 @@ 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." @@ -320,7 +323,7 @@ def _is_supported_with_reason( 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, attn, meta, fwd) + 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 @@ -466,11 +469,13 @@ def _is_supported_with_reason( 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 fused-QKV context attention without a KV cache.""" + """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( @@ -492,19 +497,29 @@ def _is_supported_without_kv_cache( 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: - return False, f"query dtype {q.dtype} is unsupported." - recipe = getattr(attn, "quant_attention_config", None) - if recipe is not None and recipe.qk_dtype == "bf16" and q.dtype != torch.bfloat16: + 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: - return False, f"fused attention input must be rank 2, got rank {q.ndim}." - 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]}." + 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) @@ -806,17 +821,19 @@ def forward( ) -> 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, metadata, forward_args) + 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: - """Fused-QKV context attention without a KV cache on the contiguous kernel.""" + """Separate-QKV context attention without a KV cache on the contiguous kernel.""" attn = self.attn output = forward_args.output if output is None: @@ -829,37 +846,23 @@ def _forward_without_kv_cache( 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 - q_width, kv_width = num_heads * head_dim, num_kv_heads * head_dim - qkv = q.view(batch_size, seq_len, q_width + 2 * kv_width) - # The fixed contiguous plan takes compact [B, S, H, D] operands. The fused-QKV - # column slices are strided, so each one is copied out. - q_bshd = qkv[..., :q_width].reshape(batch_size, seq_len, num_heads, head_dim).contiguous() - k_bshd = ( - qkv[..., q_width : q_width + kv_width] - .reshape(batch_size, seq_len, num_kv_heads, head_dim) - .contiguous() - ) - v_bshd = ( - qkv[..., q_width + kv_width :] - .reshape(batch_size, seq_len, num_kv_heads, head_dim) - .contiguous() - ) - # Per-tensor E4M3 operands as selected by the layer's attention 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. - recipe = getattr(attn, "quant_attention_config", None) + # 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 recipe is not None and recipe.qk_dtype == "fp8": + 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 * k_dequant_scale * (sm_scale * math.log2(math.e))) - .float() - .reshape(1) - ) + q_dequant_scale.float() * k_dequant_scale.float() * (sm_scale * math.log2(math.e)) + ).reshape(1) output_scale = None - if recipe is not None and recipe.v_dtype == "fp8": + 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) 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 54cf99c5494b..3c93681f369b 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -26,7 +26,6 @@ if TYPE_CHECKING: from tensorrt_llm.mapping import Mapping - from tensorrt_llm.visual_gen.args import QuantAttentionConfig from ...model_config import ModelConfig from ...speculative.interface import SpecMetadata @@ -44,7 +43,6 @@ from ...utils import (compute_swizzled_sf_shape, get_global_attrs, get_model_extra_attrs, helix_local_len_tensor) from .fmha.manager import FmhaManager -from .fmha.prims_ts import PrimsTSFmha from .fp4_mla import can_fuse_fp4_mla_q_quant, scatter_fp4_mla_kv_cache from .fp4_mla.state import Fp4MlaState from .interface import (AttentionBackend, AttentionForwardArgs, @@ -1698,7 +1696,6 @@ def __init__( skip_create_weights_in_init: bool = False, attention_chunk_size: Optional[int] = None, sparse_params: Optional[SparseParams] = None, - quant_attention_config: Optional["QuantAttentionConfig"] = None, kv_cache_dtype: str = "auto", skip_correction_threshold: float = 0.0, **kwargs, @@ -1719,8 +1716,6 @@ def __init__( sparse_params (SparseParams): Optional sparse-attention backend parameters (e.g. skip-softmax). Algorithm-specific fields are documented on the corresponding ``SparseParams`` subclass. - quant_attention_config (QuantAttentionConfig): Optional attention operand - quantization recipe (Q/K and V dtypes with their block sizes). kv_cache_dtype (str): KV-cache dtype selected by ``KvCacheConfig``. Accepted values are ``auto``, ``fp8``, ``fp8_ds_mla``, ``nvfp4``, and supported torch dtype strings. ``fp8_ds_mla`` selects the packed sparse-MLA cache @@ -1731,7 +1726,6 @@ def __init__( super().__init__(layer_idx, num_heads, head_dim, num_kv_heads, quant_config, **kwargs) self.sparse_params = sparse_params - self.quant_attention_config = quant_attention_config self.kv_cache_dtype = kv_cache_dtype self.use_fp8_ds_mla = kv_cache_dtype == "fp8_ds_mla" self.is_mla_enable = mla_params is not None @@ -1829,15 +1823,6 @@ def update_quant_config(self, new_quant_config: Optional[QuantConfig]): self.has_w4a8_nvfp4_fp8 = self.quant_config.layer_quant_mode.has_w4a8_nvfp4_fp8( ) self._fmha_manager = FmhaManager(self) - recipe = getattr(self, "quant_attention_config", None) - has_prims_ts = any( - isinstance(fmha, PrimsTSFmha) - for fmha in self._fmha_manager.fmha_libs) - if recipe is not None and recipe.algorithm == "primsts" and not has_prims_ts: - raise RuntimeError( - f"The {recipe.algorithm} attention recipe needs the PrimTS FMHA " - "library (TLLM_FMHA_LIBS=+prims_ts), which is not enabled or " - "not available on this platform.") @classmethod def runtime_workspace_bytes_per_token(cls, model_config: "ModelConfig", @@ -2480,16 +2465,14 @@ def forward( forward_args) if fmha is None: + if (forward_args.qk_fp8_per_tensor + or forward_args.v_fp8_per_tensor): + 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.") - # Backends built without __init__ (test doubles) carry no recipe. - recipe = getattr(self, "quant_attention_config", None) - if recipe is not None and recipe.algorithm == "primsts" and not isinstance( - fmha, PrimsTSFmha): - raise RuntimeError( - f"The {recipe.algorithm} attention recipe needs the PrimTS FMHA " - f"library (TLLM_FMHA_LIBS=+prims_ts), got " - f"{type(fmha).__name__}.") if metadata.is_cuda_graph or not fmha.supports_workspace_reclamation: # Conservatively disable reclamation for metadata used by graphs # or backends that can retain staged workspace state. diff --git a/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py b/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py index 514e5ef76323..c1f0ef0ef0fc 100644 --- a/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py @@ -191,7 +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 algorithm primsts, fused 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 """ @@ -223,7 +223,6 @@ def __init__( head_dim=head_dim, quant_config=quant_config, sparse_params=sparse_params, - quant_attention_config=quant_attention_config, dtype=dtype, ) @@ -234,6 +233,18 @@ def __init__( attention_metadata_state=attention_metadata_state, ) + 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]: """Normalized timestep below which the sparse algorithm is enabled, if any.""" @@ -376,7 +387,7 @@ def forward( # self-attention from fused QKV only. use_separate_qkv = ( block_sparse_inputs is not None - or (self.quant_attention_config is not None and not self.support_fused_qkv()) + or self.quant_attention_config is not None or (not self.support_fused_qkv() and self.should_use_sparse(sparse_attn_phase)) ) if use_separate_qkv and (k is None or v is None): @@ -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: @@ -428,7 +442,5 @@ def preferred_layout(self) -> AttentionTensorLayout: return self._preferred_layout def support_fused_qkv(self) -> bool: - """Standard and PrimTS paths fuse QKV. SageAttention does not.""" - return ( - self.quant_attention_config is None or self.quant_attention_config.algorithm != "sage" - ) + """Standard path fuses QKV; SageAttention path does not.""" + return self.quant_attention_config is None diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index 4ad1c68becb6..9f013a72d372 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -103,14 +103,6 @@ class QuantAttentionConfig(StrictBaseModel): "quantized type has to cover. SageAttention only." ), ) - algorithm: Literal["primsts", "sage"] = Field( - "sage", - status="prototype", - description=( - "TRTLLM kernel family for the recipe. primsts quantizes per tensor " - "(TLLM_FMHA_LIBS=+prims_ts), sage per block." - ), - ) # Discriminated union of sparse attention configs. @@ -166,7 +158,7 @@ def _validate_quant_attention_config(self) -> "AttentionConfig": ("fp8", "fp8", (1, 4, 1)), }, } - # PrimTS supports per-tensor FP8 recipes on Blackwell. + # PrimTS serves the per-tensor FP8 recipes on Blackwell (TLLM_FMHA_LIBS=+prims_ts). PRIMSTS_RECIPES = { 100: { ("bf16", "fp8", (0, 0, 0)), @@ -177,7 +169,6 @@ def _validate_quant_attention_config(self) -> "AttentionConfig": ("fp8", "fp8", (0, 0, 0)), }, } - TRTLLM_RECIPES = {"primsts": PRIMSTS_RECIPES, "sage": SAGE_RECIPES} # Other recipes verify the hardware at corresponding backend implementations. CUDNN_RECIPES = { ("fp8", "fp8", (0, 0, 0)), @@ -206,23 +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 = TRTLLM_RECIPES[q_config.algorithm].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 {q_config.algorithm} recipes on this " - f"device: {sorted(recipes)}." + f"for backend='TRTLLM'. Supported SAGE and PrimTS recipes on this device: " + f"{sorted(recipes)}." ) - if q_config.smooth_k and q_config.algorithm != "sage": + if q_config.smooth_k and q_config.q_block_size == 0: raise ValueError( - f"smooth_k is a SageAttention option and does not apply to " - f"algorithm='{q_config.algorithm}'." + "smooth_k is a SageAttention option and does not apply to the per-tensor " + f"recipe {recipe!r}." ) - elif q_config.algorithm != "sage": - raise ValueError( - f"algorithm='{q_config.algorithm}' selects a TRTLLM kernel family and requires " - f"backend='TRTLLM', got backend='{self.backend}'." - ) 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 e9688574d4fa..071fcc8154f8 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py +++ b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py @@ -1283,19 +1283,23 @@ def test_prims_ts_decode_wrappers_share_workspace_across_serialized_layers() -> torch.testing.assert_close(actual, reference, atol=3e-2, rtol=3e-3) -def _build_trtllm_attention_with_recipe(monkeypatch: pytest.MonkeyPatch, recipe) -> None: - """Make the shared harness pass the attention recipe to the TRTLLM layer it builds.""" +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 - create_attention = backend_case.create_attention + from tensorrt_llm._torch.attention.backends.interface import AttentionForwardArgs + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention - def create_attention_with_recipe(backend, **kwargs): - attn = create_attention(backend, **kwargs) - if backend == "TRTLLM": - attn.quant_attention_config = recipe - return attn - - monkeypatch.setattr(backend_case, "create_attention", create_attention_with_recipe) + 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( @@ -1325,8 +1329,8 @@ def test_prims_ts_cache_free_fp8_recipes( 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, algorithm="primsts") - _build_trtllm_attention_with_recipe(monkeypatch, recipe) + 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 @@ -1396,11 +1400,11 @@ def dequant_scale(x: torch.Tensor) -> torch.Tensor: def test_prims_ts_recipe_requires_prims_ts_library(monkeypatch: pytest.MonkeyPatch) -> None: - """A PrimTS recipe fails instead of running unquantized on another FMHA library.""" + """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", algorithm="primsts") - _build_trtllm_attention_with_recipe(monkeypatch, recipe) + 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) diff --git a/tests/unittest/_torch/attention/test_prims_ts_fmha.py b/tests/unittest/_torch/attention/test_prims_ts_fmha.py index e508d469d24c..8d2f771b88b2 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_fmha.py +++ b/tests/unittest/_torch/attention/test_prims_ts_fmha.py @@ -88,7 +88,6 @@ def __init__( self.v_head_dim = 128 if is_mla else None self.predicted_tokens_per_seq = 1 self.sparse_params = None - self.quant_attention_config = None self.skip_correction_threshold = 0.0 self.position_embedding_type = 0 self.quant_mode = 0 @@ -2343,10 +2342,9 @@ def test_cache_free_context_plans_contiguous_wrapper_once(monkeypatch: pytest.Mo 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 + 2 * attn.num_kv_heads) * attn.head_dim), - dtype=torch.bfloat16, - ) + 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 ) @@ -2363,8 +2361,8 @@ def test_cache_free_context_plans_contiguous_wrapper_once(monkeypatch: pytest.Mo attention_mask=PredefinedAttentionMask.FULL, ) - fmha.forward(q, None, None, metadata, forward_args) - fmha.forward(q, None, None, metadata, forward_args) + 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() @@ -2379,6 +2377,8 @@ def test_cache_free_context_plans_contiguous_wrapper_once(monkeypatch: pytest.Mo 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 9ce92edede1b..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,29 +184,16 @@ def test_supported_quant_config_sage( assert attention.quant_attention_config is not None - @pytest.mark.parametrize("sm_ver", [100, 103]) - @pytest.mark.parametrize("qk_dtype", ["bf16", "fp8"]) - def test_supported_quant_config_primsts(self, sm_ver, qk_dtype): - """PrimTS per-tensor recipes validate on Blackwell for both Q/K dtypes.""" - with patch("tensorrt_llm.visual_gen.args.get_sm_version", return_value=sm_ver): - attention = AttentionConfig( - backend="TRTLLM", - quant_attention_config=QuantAttentionConfig( - qk_dtype=qk_dtype, v_dtype="fp8", algorithm="primsts" - ), - ) - - assert attention.quant_attention_config.algorithm == "primsts" - - @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( From b3b1cde36afa29eaa0bc4d5def932fabdf3f27c4 Mon Sep 17 00:00:00 2001 From: Harrison Zhang Date: Wed, 14 Oct 2026 12:02:43 -0700 Subject: [PATCH 08/11] [primsTS][fmha] retrigger the vendor status check Signed-off-by: Harrison Zhang From 59eb9c47fb7b716527c9553c6de1f0a65e275edc Mon Sep 17 00:00:00 2001 From: Harrison Zhang Date: Wed, 14 Oct 2026 12:49:20 -0700 Subject: [PATCH 09/11] [primsTS][fmha] point PrimTS vendor source back to flashinfer harrisonzhy/trtllm-prims-ts-dev_perf-bringup (ae312f7c) Signed-off-by: Harrison Zhang --- 3rdparty/vendor_sources.lock.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 3045a99ba5a5..b44a01d29623 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -2,7 +2,7 @@ schema_version: 1 vendors: flashinfer-prims-ts: url: https://github.com/harrisonzhy/flashinfer.git - branch: harrisonzhy/trtllm-prims-ts-dev_perf-bringup-20261012 + branch: harrisonzhy/trtllm-prims-ts-dev_perf-bringup commit: ae312f7cc1723b06fda5168105ade0f328ee47d8 source: flashinfer/attention/prims_ts destination: tensorrt_llm/_torch/attention/backends/prims_ts From f87cff0998a2d86a91b15cf665b92d2f054242c7 Mon Sep 17 00:00:00 2001 From: Harrison Zhang Date: Thu, 15 Oct 2026 00:08:07 -0700 Subject: [PATCH 10/11] Update vendor source URL and branch for flashinfer Signed-off-by: Harrison Zhang --- 3rdparty/vendor_sources.lock.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index b44a01d29623..2e914da8f052 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -1,8 +1,8 @@ schema_version: 1 vendors: flashinfer-prims-ts: - url: https://github.com/harrisonzhy/flashinfer.git - branch: harrisonzhy/trtllm-prims-ts-dev_perf-bringup + url: https://github.com/yuxianq/flashinfer.git + branch: trtllm-prims-ts-dev-20261012 commit: ae312f7cc1723b06fda5168105ade0f328ee47d8 source: flashinfer/attention/prims_ts destination: tensorrt_llm/_torch/attention/backends/prims_ts From 1dd23ca5fd4877b852ea2a223b4717b1df79f0f6 Mon Sep 17 00:00:00 2001 From: Harrison Zhang Date: Thu, 15 Oct 2026 07:57:26 +0000 Subject: [PATCH 11/11] [primsTS][fmha] keep the per-tensor FP8 recipe flags out of the thop sync check and read them defensively qk_fp8_per_tensor and v_fp8_per_tensor are consumed by FMHA selection and the PrimTS FMHA, not by thop.attention, so list them in _THOP_EXCLUDED_FIELDS; the support check and the no-library error read them with getattr so forward-args stubs without the fields still dispatch. Signed-off-by: Harrison Zhang --- tensorrt_llm/_torch/attention/backends/fmha/fallback.py | 2 ++ tensorrt_llm/_torch/attention/backends/fmha/interface.py | 3 ++- tensorrt_llm/_torch/attention/backends/trtllm.py | 4 ++-- 3 files changed, 6 insertions(+), 3 deletions(-) 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 71edcf451796..e3905dd1964d 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/interface.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/interface.py @@ -121,7 +121,8 @@ def is_supported( logger.debug(f"{type(self).__name__} does not support block-sparse inputs.") return False if ( - forward_args.qk_fp8_per_tensor or forward_args.v_fp8_per_tensor + 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 diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 3c93681f369b..048138a48582 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -2465,8 +2465,8 @@ def forward( forward_args) if fmha is None: - if (forward_args.qk_fp8_per_tensor - or forward_args.v_fp8_per_tensor): + 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 "