diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 18640486388d..26d111699e19 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -14,15 +14,15 @@ vendors: digest: sha256-tree-v1:d2d382335ccb68b55078523e63835c28b699607227afb81308cb7bb1d85845ac flashinfer-prims-ts: url: https://github.com/yuxianq/flashinfer.git - branch: trtllm-prims-ts-dev-20261012 - commit: aa709061230a58858375c513056e7733fa74294a + branch: prims-tsfmha-fp8-non-mla--kernels + commit: dd087c6568aa144bb7ceff1306e2568279ec5068 source: flashinfer/attention/prims_ts destination: tensorrt_llm/_torch/attention/backends/prims_ts include: - '**/*.py' patch: 3rdparty/vendor_patches/flashinfer-prims-ts.patch patch_digest: sha256:fc5c23ec3d65c562025a1845f211fdac66bfcd3da156dbb8385f4c4407d20d07 - digest: sha256-tree-v1:f2f444f7e6d4ad688fbf4e1e9eed11ff7c78d835cfd0af5ef9d7660780b059e0 + digest: sha256-tree-v1:a794ed3cc25687817ea08878b1e156ce7616ce553058c1a3a70a122ee915e9fc openengine-proto: url: https://github.com/ai-dynamo/openengine.git tag: v0.1.0 diff --git a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py index c42f48995f46..27398ecfeeda 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py @@ -286,7 +286,7 @@ def _page_size_exclusive_to_cute_dsl(meta: "TrtllmAttentionMetadata") -> bool: if tokens_per_block in getattr(kv_cache_manager, "trtllm_gen_extra_tokens_per_block", ()): return False return not ( - get_sm_version() in PrimsTSFmha.SUPPORTED_SMS + get_sm_version() in PrimsTSFmha.SUPPORTED_MLA_SMS and tokens_per_block in PrimsTSFmha.SUPPORTED_PAGE_SIZES and getattr(kv_cache_manager, "dtype", None) in (DataType.FP8, DataType.BF16) ) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py index 89ee7d2a8126..02d9c6d2dc97 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py @@ -13,17 +13,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""TRT-LLM FMHA adapter for the vendored PrimTS Blackwell kernels.""" +"""TRT-LLM FMHA adapter for the vendored PrimTS kernels.""" from __future__ import annotations import math from importlib import import_module -from importlib.metadata import PackageNotFoundError, version from typing import TYPE_CHECKING, Any, Optional import torch -from packaging.version import InvalidVersion, Version from tensorrt_llm._torch.attention.backends.interface import ( AttentionForwardArgs, @@ -57,7 +55,6 @@ ) -_MIN_CUTLASS_DSL_VERSION = Version("4.7.0") _MIN_CUTLASS_COMPILER_VERSION = "13.3" _WORKSPACE_ALIGNMENT = 32 @@ -105,6 +102,8 @@ def get_paged_kv_storage_unsupported_reason( def get_paged_kv_policy_unsupported_reason( attn: "TrtllmAttention", metadata: "TrtllmAttentionMetadata", + *, + allow_fp8_kv_cache: bool = False, ) -> Optional[str]: """Return why the request's decoding policy is outside the fixed page-table envelope.""" if metadata.beam_width != 1: @@ -123,13 +122,15 @@ def get_paged_kv_policy_unsupported_reason( quant_mode = QuantMode(attn.quant_mode) except (TypeError, ValueError): return "invalid KV-cache quantization mode." - is_fp8_mla = ( - attn.is_mla_enable + is_supported_fp8 = ( + (attn.is_mla_enable or allow_fp8_kv_cache) and quant_mode.has_fp8_kv_cache() and not quant_mode.has_int8_kv_cache() and not quant_mode.has_fp4_kv_cache() ) - if quant_mode.has_kv_cache_quant() and not is_fp8_mla: + if quant_mode.has_kv_cache_quant() and not is_supported_fp8: + if allow_fp8_kv_cache: + return "only FP8 E4M3 KV-cache quantization is supported." return "quantized KV cache is supported only for FP8 MLA decode." return None @@ -153,11 +154,17 @@ def get_attention_feature_unsupported_reason( class PrimsTSFmha(PhasedFmha): - """Blackwell task-scheduled paged context and decode FMHA library.""" + """Task-scheduled paged context and decode FMHA library. + + Standard FP8 KV cache uses fused preprocessing to quantize Q with the + shared Q/K/V scale. Kernels consume FP8 Q/K/V and write the model's FP16 + or BF16 output directly, using preprocessing-produced device scales. + """ supports_per_tensor_fp8 = True - SUPPORTED_SMS = frozenset({100, 103}) + SUPPORTED_SMS = frozenset({100, 103, 107}) + SUPPORTED_MLA_SMS = frozenset({100, 103}) SUPPORTED_PAGE_SIZES = {16, 32, 64, 128} SUPPORTED_CONTEXT_HEAD_DIMS = {128, 256} SUPPORTED_DECODE_HEAD_DIMS = {64, 128, 256} @@ -192,20 +199,13 @@ def __init__(self, attn: "TrtllmAttention") -> None: def _is_available(cls, attn: "TrtllmAttention") -> bool: sm = get_sm_version() if sm not in cls.SUPPORTED_SMS: - logger.debug(f"PrimTS FMHA is unavailable: requires SM100 or SM103, got SM{sm}.") - return False - try: - installed_version = Version(version("nvidia-cutlass-dsl")) - except (PackageNotFoundError, InvalidVersion): - logger.debug("PrimTS FMHA is unavailable: nvidia-cutlass-dsl>=4.7 is required.") - return False - if installed_version < _MIN_CUTLASS_DSL_VERSION: logger.debug( - "PrimTS FMHA is unavailable: " - f"nvidia-cutlass-dsl>={_MIN_CUTLASS_DSL_VERSION} is required, " - f"got {installed_version}." + f"PrimTS FMHA is unavailable: requires SM100, SM103, or SM107, got SM{sm}." ) return False + if attn.is_mla_enable and sm not in cls.SUPPORTED_MLA_SMS: + logger.debug(f"PrimTS MLA FMHA is unavailable on SM{sm}.") + return False try: cutlass = import_module("cutlass") compiler_version_supported = cutlass.target_version( @@ -331,11 +331,12 @@ def _is_supported_with_reason( return False, "the attention mask is not causal or dense." if mask_type not in (AttentionMaskType.causal, AttentionMaskType.padding): return False, f"attention mask type {mask_type} is not supported." - policy_reason = get_paged_kv_policy_unsupported_reason(attn, meta) + policy_reason = get_paged_kv_policy_unsupported_reason(attn, meta, allow_fp8_kv_cache=True) if policy_reason is not None: return False, policy_reason is_mla = attn.is_mla_enable - is_fp8_mla = is_mla and QuantMode(attn.quant_mode).has_fp8_kv_cache() + is_fp8_kv = QuantMode(attn.quant_mode).has_fp8_kv_cache() + is_fp8_mla = is_mla and is_fp8_kv input_type = fwd.attention_input_type if input_type not in ( @@ -375,9 +376,9 @@ def _is_supported_with_reason( if q.dtype not in self.SUPPORTED_DTYPES: return False, f"query dtype {q.dtype} is unsupported." cache_dtype = binding_to_torch_dtype(meta.kv_cache_manager.dtype) - if is_fp8_mla and cache_dtype != torch.float8_e4m3fn: - return False, "FP8 MLA decode requires an FP8 E4M3 KV cache." - if not is_fp8_mla and cache_dtype != q.dtype: + if is_fp8_kv and cache_dtype != torch.float8_e4m3fn: + return False, "FP8 attention requires an FP8 E4M3 KV cache." + if not is_fp8_kv and cache_dtype != q.dtype: return False, f"query and KV-cache dtypes must match, got {q.dtype} and {cache_dtype}." if output.dtype != q.dtype: return False, f"output dtype must match query dtype, got {output.dtype} and {q.dtype}." @@ -408,6 +409,22 @@ def _is_supported_with_reason( if error is not None: return False, error else: + if is_fp8_kv: + for name, scale in ( + ("kv_scale_orig_quant", fwd.kv_scale_orig_quant), + ("kv_scale_quant_orig", fwd.kv_scale_quant_orig), + ): + if ( + scale is None + or scale.dtype != torch.float32 + or scale.device != q.device + or not scale.is_contiguous() + or scale.numel() != 1 + ): + return ( + False, + f"FP8 attention requires scalar float32 {name} on the query device.", + ) expected_width = (attn.num_heads + 2 * attn.num_kv_heads) * attn.head_dim if q.shape[1] != expected_width: return False, f"fused QKV width must be {expected_width}, got {q.shape[1]}." @@ -757,6 +774,7 @@ def _get_or_plan_decode_wrapper( mask_type=mask_type, window_left=window_left, workspace_buffer=workspace_buffer, + use_device_scales=QuantMode(self.attn.quant_mode).has_fp8_kv_cache(), ) self._decode_wrappers[batch_size] = wrapper return wrapper @@ -936,7 +954,7 @@ def prepare_workspace( self.attn.head_dim, self.attn.rope_dim, True, - False, + QuantMode(self.attn.quant_mode).has_fp8_kv_cache(), skip_fmha_workspace=True, ) required_preprocess_bytes = max( @@ -983,16 +1001,14 @@ def prepare_workspace( # sizing key aligned with the non-windowed plan selected at runtime. window_left = -1 + kernel_dtype = ( + torch.float8_e4m3fn if QuantMode(self.attn.quant_mode).has_fp8_kv_cache() else q.dtype + ) if self.attn.is_mla_enable: from tensorrt_llm._torch.attention.backends.prims_ts import ( get_prims_ts_batch_mla_decode_workspace_size, ) - kernel_dtype = ( - torch.float8_e4m3fn - if QuantMode(self.attn.quant_mode).has_fp8_kv_cache() - else q.dtype - ) required_bytes = get_prims_ts_batch_mla_decode_workspace_size( batch_size, self.attn.num_heads, @@ -1020,8 +1036,8 @@ def prepare_workspace( int(metadata.tokens_per_block), max_seq_len, seq_len_q=seq_len_q, - q_dtype=q.dtype, - kv_dtype=q.dtype, + q_dtype=kernel_dtype, + kv_dtype=kernel_dtype, out_dtype=forward_args.output.dtype, mask_type=mask_type, window_left=window_left, @@ -1108,8 +1124,8 @@ def run_context(self, params: FmhaParams) -> None: kv_pool, block_tables, _kv_scale_pool, - _bmm1_scale, - _bmm2_scale, + bmm1_scale, + bmm2_scale, fmha_workspace, cu_q_seqlens, _cu_kv_seqlens, @@ -1152,7 +1168,7 @@ def run_context(self, params: FmhaParams) -> None: self._get_bmm1_scale(attn), 1.0, attention_chunk_size, - False, + QuantMode(attn.quant_mode).has_fp8_kv_cache(), True, False, self._multi_processor_count, @@ -1184,6 +1200,9 @@ def run_context(self, params: FmhaParams) -> None: params.sequence_lengths, params.batch_size, ) + use_fp8 = QuantMode(attn.quant_mode).has_fp8_kv_cache() + if use_fp8: + q_processed = self._view_fp8_query(q_processed, params.num_tokens) mask_type = self._get_prims_mask_type(fwd) wrapper = self._get_or_plan_context_wrapper( q_processed, @@ -1205,6 +1224,8 @@ def run_context(self, params: FmhaParams) -> None: cu_q_seqlens, block_tables=fixed_block_tables, seq_lens_kv=seq_lens_kv, + scale_softmax_log2=bmm1_scale.reshape(-1)[1:2] if use_fp8 else None, + output_scale=bmm2_scale.reshape(-1)[:1] if use_fp8 else None, out=params.output, validate=False, ) @@ -1242,7 +1263,7 @@ def run_context(self, params: FmhaParams) -> None: rope_params.max_positions, attn.position_embedding_type, self._get_bmm1_scale(attn), - False, + QuantMode(attn.quant_mode).has_fp8_kv_cache(), True, False, attention_chunk_size, @@ -1296,7 +1317,7 @@ def _run_generation_preprocess(self, params: FmhaParams) -> tuple[Any, ...]: attn.position_embedding_type, self._get_bmm1_scale(attn), 1.0, - False, + QuantMode(attn.quant_mode).has_fp8_kv_cache(), attn.predicted_tokens_per_seq, attention_chunk_size, self._multi_processor_count, @@ -1324,8 +1345,8 @@ def run_generation(self, params: FmhaParams) -> None: kv_pool, block_tables, _kv_scale_pool, - _bmm1_scale, - _bmm2_scale, + bmm1_scale, + bmm2_scale, fmha_workspace, _cu_seqlens, _max_q_len, @@ -1352,6 +1373,9 @@ def run_generation(self, params: FmhaParams) -> None: fixed_block_tables = self._get_fixed_block_tables(block_tables, batch_size) max_seq_len = int(block_tables.shape[-1]) * params.tokens_per_block seq_lens = self._get_sequence_lengths(params.sequence_lengths, batch_size) + use_fp8 = QuantMode(attn.quant_mode).has_fp8_kv_cache() + if use_fp8: + q_processed = self._view_fp8_query(q_processed, params.num_tokens) query = q_processed.view( batch_size, params.input_seq_length, @@ -1396,12 +1420,26 @@ def run_generation(self, params: FmhaParams) -> None: (k_cache, v_cache), seq_lens, block_tables=fixed_block_tables, - bmm1_scale=self._get_bmm1_scale(attn), - bmm2_scale=1.0, + bmm1_scale=bmm1_scale.reshape(-1)[:1] if use_fp8 else self._get_bmm1_scale(attn), + bmm2_scale=bmm2_scale.reshape(-1)[:1] if use_fp8 else 1.0, out=output, validate=False, ) + def _view_fp8_query(self, q_processed: torch.Tensor, num_tokens: int) -> torch.Tensor: + """View the packed FP8 prefix written by fused QKV preprocessing. + + Generation reserves a model-dtype Q buffer, but writes one byte per + element when FP8 is enabled. Its unused trailing bytes are not queries. + """ + numel = num_tokens * self.attn.num_heads * self.attn.head_dim + return ( + q_processed.view(torch.uint8) + .reshape(-1)[:numel] + .view(torch.float8_e4m3fn) + .view(num_tokens, self.attn.num_heads, self.attn.head_dim) + ) + def _get_decode_workspace( self, root_workspace: torch.Tensor, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py b/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py index f91814650c48..47c7b9a33480 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py @@ -132,6 +132,7 @@ class _DecodeCompileSpec: max_active_clusters: int kv_prefix_mode: Literal["dynamic", "planned_full"] kv_lengths_mode: Literal["dynamic", "planned_uniform_max"] + use_device_scales: bool = False direct_q1_spec: Optional[_DirectQ1CompileSpec] = None @@ -152,8 +153,8 @@ class _DecodeRuntime: v_page_stride: int v_head_stride: int v_token_stride: int - bmm1_scale: float - bmm2_scale: float + bmm1_scale: float | torch.Tensor + bmm2_scale: float | torch.Tensor @dataclass(frozen=True) @@ -192,6 +193,7 @@ class _DecodePlanState: planned_seq_lens_device: Optional[torch.Tensor] policy: tuple[tuple[str, object], ...] storage_page_size: Optional[int] = None + use_device_scales: bool = False def _decode_policy_from_config( @@ -1479,6 +1481,7 @@ def _make_decode_compile_spec( kv_lengths_mode: Literal["dynamic", "planned_uniform_max"], direct_q1_spec: Optional[_DirectQ1CompileSpec] = None, flat_native_kv_tma: bool = False, + use_device_scales: bool = False, ) -> _DecodeCompileSpec: """Freeze the resolved topology while leaving batch in runtime tensors.""" @@ -1503,15 +1506,15 @@ def _make_decode_compile_spec( kv_prefix_mode=kv_prefix_mode, kv_lengths_mode=kv_lengths_mode, direct_q1_spec=direct_q1_spec, + use_device_scales=use_device_scales, ) @functools.cache -def _get_compiled_decode( - compile_spec: _DecodeCompileSpec, -): +def _get_compiled_decode(compile_spec: _DecodeCompileSpec): """Compile and cache one batch-dynamic TS decode topology.""" + use_device_scales = compile_spec.use_device_scales device_index = compile_spec.device_index num_qo_heads = compile_spec.num_qo_heads num_kv_heads = compile_spec.num_kv_heads @@ -1588,8 +1591,8 @@ def main_tensor_adapter( v_page_stride: cutlass.Int64, v_head_stride: cutlass.Int64, v_token_stride: cutlass.Int64, - bmm1_scale: cutlass.Float32, - bmm2_scale: cutlass.Float32, + bmm1_scale: cutlass.Float32 | cute.Tensor, + bmm2_scale: cutlass.Float32 | cute.Tensor, direct_q1_inputs: tuple, stream: cuda_drv.CUstream, static_cfg: cutlass.Constexpr[FmhaDecodeConfig], @@ -1667,8 +1670,8 @@ def main_tensor_adapter( partial_stats.iterator, split_kv_counter.iterator, attention_sinks.iterator, - bmm1_scale, - bmm2_scale, + bmm1_scale.iterator if use_device_scales else bmm1_scale, + bmm2_scale.iterator if use_device_scales else bmm2_scale, Int32(0), Int32(static_max_active_clusters), stream, @@ -1704,8 +1707,6 @@ def reduction_tensor_adapter( partial_o: cute.Tensor, partial_stats: cute.Tensor, attention_sinks: cute.Tensor, - bmm1_scale: cutlass.Float32, - bmm2_scale: cutlass.Float32, direct_q1_inputs: tuple, stream: cuda_drv.CUstream, static_cfg: cutlass.Constexpr[FmhaDecodeConfig], @@ -1752,8 +1753,6 @@ def reduction_tensor_adapter( partial_o.iterator, partial_stats.iterator, attention_sinks.iterator, - bmm1_scale, - bmm2_scale, stream, static_cfg, static_full_split_prefix, @@ -1899,6 +1898,7 @@ def fake_compact(dtype, shape, assumed_align): 4, ), ) + scale_fake = fake_compact(Float32, (1,), 4) if use_device_scales else Float32(1.0) def compile_main(): return cute.compile( @@ -1924,8 +1924,8 @@ def compile_main(): Int64(1), Int64(1), Int64(1), - Float32(1.0), - Float32(1.0), + scale_fake, + scale_fake, direct_q1_fake, stream_fake, cfg, @@ -1950,8 +1950,6 @@ def compile_reducer(): partial_o_fake, partial_stats_fake, attention_sinks_fake, - Float32(1.0), - Float32(1.0), direct_q1_fake, stream_fake, cfg, @@ -1965,7 +1963,11 @@ def compile_reducer(): with torch.cuda.device(device_index): compiled_main = build_and_load_attention_kernel( - "decode", "main", compile_spec, _COMPILE_OPTIONS, compile_main + "decode", + "main", + compile_spec, + _COMPILE_OPTIONS, + compile_main, ) compiled_reducer = None if cfg.use_separate_reduction_kernel: @@ -2112,9 +2114,10 @@ def _prepare_decode_runtime( k_dtype: torch.dtype, v_dtype: torch.dtype, output_dtype: torch.dtype, - bmm1_scale: Optional[float], - bmm2_scale: float, + bmm1_scale: Optional[float | torch.Tensor], + bmm2_scale: float | torch.Tensor, out: Optional[torch.Tensor], + use_device_scales: bool = False, ) -> _DecodeRuntime: """Validate runtime tensors and normalize zero-copy K/V views.""" @@ -2158,11 +2161,15 @@ def _prepare_decode_runtime( k_cache=k_cache, logical_head_dim=head_dim, ) - effective_bmm1_scale = _validate_scale( - 1.0 / math.sqrt(head_dim) if bmm1_scale is None else bmm1_scale, - "bmm1_scale", - ) - effective_bmm2_scale = _validate_scale(bmm2_scale, "bmm2_scale") + if use_device_scales: + effective_bmm1_scale = bmm1_scale + effective_bmm2_scale = bmm2_scale + else: + effective_bmm1_scale = _validate_scale( + 1.0 / math.sqrt(head_dim) if bmm1_scale is None else bmm1_scale, + "bmm1_scale", + ) + effective_bmm2_scale = _validate_scale(bmm2_scale, "bmm2_scale") output_shape = _decode_output_shape( batch_size=batch_size, num_qo_heads=num_qo_heads, @@ -2247,8 +2254,6 @@ def _launch_decode( workspace.partial_o, workspace.partial_stats, workspace.attention_sinks, - runtime.bmm1_scale, - runtime.bmm2_scale, (), ) return runtime.out @@ -2309,9 +2314,10 @@ def _prepare_decode_runtime_unchecked( kv_scale_factors: Optional[PagedKVScaleFactors], *, output_dtype: torch.dtype, - bmm1_scale: Optional[float], - bmm2_scale: float, + bmm1_scale: Optional[float | torch.Tensor], + bmm2_scale: float | torch.Tensor, out: Optional[torch.Tensor], + use_device_scales: bool = False, ) -> _DecodeRuntime: """Canonicalize one trusted run without invoking explicit validators.""" @@ -2346,11 +2352,15 @@ def _prepare_decode_runtime_unchecked( v_head_stride=int(v_cache.stride(1)), v_token_stride=int(v_cache.stride(2)), bmm1_scale=( - 1.0 / math.sqrt(int(q.shape[-1])) - if bmm1_scale is None - else float(bmm1_scale) + bmm1_scale + if use_device_scales + else ( + 1.0 / math.sqrt(int(q.shape[-1])) + if bmm1_scale is None + else float(bmm1_scale) + ) ), - bmm2_scale=float(bmm2_scale), + bmm2_scale=bmm2_scale if use_device_scales else float(bmm2_scale), ) @@ -2581,8 +2591,6 @@ def _run_unchecked( self._workspace.partial_o, self._workspace.partial_stats, self._workspace.attention_sinks, - scale_qk, - scale_v, direct_inputs, ) return out @@ -3773,6 +3781,7 @@ def plan( split_kv: bool = True, validate: bool = True, initialize_workspace: bool = True, + use_device_scales: bool = False, ) -> None: """Compile one static-capacity plan and optionally own sequence lengths. @@ -3788,6 +3797,11 @@ def plan( omitted, both K/V prefix and length handling compile dynamically and every run must supply a CUDA length tensor. + ``use_device_scales=True`` compiles a plan that reads float32 BMM + scales from CUDA tensors supplied to every run. Values must be finite + and positive; they are not read back to the host for validation. This + supports preprocessing-produced scales and CUDA graph replay. + ``workspace_buffer`` is caller-owned scratch for this plan. It is allocated when omitted, retained by the frozen plan state, and never reset by ``run``. By default planning initializes its control sections; @@ -3863,8 +3877,14 @@ def plan( Initialize control sections during planning, default ``True``. Set to ``False`` only for caller-initialized scratch with unchanged layout. Newly allocated scratch is always initialized. + use_device_scales : bool + Read BMM scales from per-run CUDA tensors instead of host scalars. + Defaults to ``False``. """ + if not isinstance(use_device_scales, bool): + raise TypeError("use_device_scales must be a bool") + k_data_type, v_data_type = _resolve_kv_dtypes( q_data_type, k_data_type, v_data_type, kv_data_type ) @@ -3995,6 +4015,7 @@ def plan( use_packed_q=packed_query, kv_prefix_mode=kv_prefix_mode, kv_lengths_mode=kv_lengths_mode, + use_device_scales=use_device_scales, ) compiled_main, compiled_reducer = _get_compiled_decode(compile_spec) policy = spec.policy + ( @@ -4067,6 +4088,7 @@ def plan( planned_seq_lens_device=planned_seq_lens_device, policy=policy, storage_page_size=storage_page_size, + use_device_scales=use_device_scales, ) # This is the only wrapper mutation. Any failure above leaves the # previous complete plan revision usable. @@ -4082,8 +4104,8 @@ def run( *, kv_scale_factors: Optional[PagedKVScaleFactors] = None, qo_indptr: Optional[torch.Tensor] = None, - bmm1_scale: Optional[float] = None, - bmm2_scale: float = 1.0, + bmm1_scale: Optional[float | torch.Tensor] = None, + bmm2_scale: float | torch.Tensor = 1.0, out: Optional[torch.Tensor] = None, validate: bool = True, ) -> torch.Tensor: @@ -4096,7 +4118,7 @@ def run( ``validate=True`` performs structural and value validation. It reads per-run metadata values back to the host. ``validate=False`` treats every run argument as a trusted binding, - performs no explicit wrapper validation, and remains free of metadata + skips metadata validation and remains free of metadata device-to-host synchronization. K/V view selection, scale forwarding, and optional output allocation are unavoidable in both modes. @@ -4129,17 +4151,22 @@ def run( qo_indptr : torch.Tensor, optional Per-run cumulative query offsets with shape ``[B + 1]``. Required for a packed-query plan and rejected for a fixed-query plan. - bmm1_scale : float, optional - QK scaling factor. Defaults to the inverse square root of - ``head_dim``. - bmm2_scale : float - Value/output scaling factor. Defaults to ``1.0``. + bmm1_scale : float or torch.Tensor, optional + QK scaling factor in regular, not log2, units. For host-scale plans, + defaults to the inverse square root of ``head_dim``. + bmm2_scale : float or torch.Tensor + Value/output scaling factor. Defaults to ``1.0`` for host-scale plans. + A plan with ``use_device_scales=True`` requires both scales to be + contiguous float32 CUDA tensors of shape ``(1,)`` on the plan device. + Their values must be finite and positive; values are not read back + to the host for validation and may change in place on graph replay. out : torch.Tensor, optional Caller-owned output tensor. A new tensor is allocated when omitted. validate : bool Run explicit structural and value validation. Disable only when the caller guarantees the complete runtime contract. Sequence - length ownership is enforced in either mode. Defaults to ``True``. + length ownership and scale argument types are enforced in either mode. + Defaults to ``True``. Returns ------- @@ -4150,6 +4177,26 @@ def run( state = self._require_plan_state() if not isinstance(validate, bool): raise TypeError("validate must be a bool") + if state.use_device_scales: + for name, scale in ( + ("bmm1_scale", bmm1_scale), + ("bmm2_scale", bmm2_scale), + ): + if ( + not isinstance(scale, torch.Tensor) + or scale.dtype != torch.float32 + or scale.device != state.device + or scale.shape != (1,) + or not scale.is_contiguous() + ): + raise ValueError( + f"{name} must be a float32 tensor of shape (1,) on {state.device}" + ) + elif isinstance(bmm1_scale, torch.Tensor) or isinstance( + bmm2_scale, torch.Tensor + ): + raise ValueError("device scales require plan(use_device_scales=True)") + planned_seq_lens = state.planned_seq_lens_device plan_owns_seq_lens = planned_seq_lens is not None if plan_owns_seq_lens: @@ -4218,6 +4265,7 @@ def run( bmm1_scale=bmm1_scale, bmm2_scale=bmm2_scale, out=out, + use_device_scales=state.use_device_scales, ) _validate_decode_run_metadata_values( runtime, @@ -4241,6 +4289,7 @@ def run( bmm1_scale=bmm1_scale, bmm2_scale=bmm2_scale, out=out, + use_device_scales=state.use_device_scales, ) return _launch_decode( diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_kernel.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_kernel.py index 7a2235d8bc28..4cacee3615ac 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_kernel.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_kernel.py @@ -2551,8 +2551,8 @@ def _run_decode_gen_active( o_iter: cute.Pointer, g_s_k: Int32, g_h_k: Int32, - g_scale_s_log2_e: Float32, - g_output_scale: Float32, + g_scale_s_log2_e: Float32 | cute.Pointer, + g_output_scale: Float32 | cute.Pointer, g_seqlens_kv: cute.Pointer | DirectSparseMetadataView, g_cu_seqlens_q: cute.Pointer, g_page_idx_kv: cute.Pointer | DirectSparseMetadataView, @@ -2600,6 +2600,15 @@ def _run_decode_gen_active( dynamic branch. """ + # Host scales arrive in log2 form; device scales are produced in regular + # form by QKV preprocessing and must be loaded by the GPU on every launch. + if cutlass.const_expr(isinstance(g_scale_s_log2_e, cute.Pointer)): + g_scale_s_log2_e = cute.make_tensor(g_scale_s_log2_e, cute.make_layout(1))[ + 0 + ] * Float32(math.log2(math.e)) + if cutlass.const_expr(isinstance(g_output_scale, cute.Pointer)): + g_output_scale = cute.make_tensor(g_output_scale, cute.make_layout(1))[0] + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) bias_static_sliding_kv_tma = ( @@ -2893,8 +2902,8 @@ def _run_decode_gen_runtime_prefix( o_iter: cute.Pointer, g_s_k: Int32, g_h_k: Int32, - g_scale_s_log2_e: Float32, - g_output_scale: Float32, + g_scale_s_log2_e: Float32 | cute.Pointer, + g_output_scale: Float32 | cute.Pointer, g_seqlens_kv: cute.Pointer | DirectSparseMetadataView, g_cu_seqlens_q: cute.Pointer, g_page_idx_kv: cute.Pointer | DirectSparseMetadataView, @@ -3086,8 +3095,8 @@ def decode_gen_kernel( o_iter: cute.Pointer, g_s_k: Int32, g_h_k: Int32, - g_scale_s_log2_e: Float32, - g_output_scale: Float32, + g_scale_s_log2_e: Float32 | cute.Pointer, + g_output_scale: Float32 | cute.Pointer, g_seqlens_kv: cute.Pointer | DirectSparseMetadataView, g_cu_seqlens_q: cute.Pointer, g_page_idx_kv: cute.Pointer | DirectSparseMetadataView, @@ -3332,8 +3341,8 @@ def fmha_decode_launch( partial_stats_iter: cute.Pointer, split_kv_counter_iter: cute.Pointer, attention_sinks_iter: cute.Pointer, - scale_s: Float32, - output_scale: Float32, + scale_s: Float32 | cute.Pointer, + output_scale: Float32 | cute.Pointer, kv_b_stride: Int32, max_active_clusters: Int32, stream: cuda_drv.CUstream, @@ -3711,6 +3720,9 @@ def fmha_decode_launch( 0, mem_space=cutlass.AddressSpace.gmem, ) + scale_log2 = scale_s + if cutlass.const_expr(not isinstance(scale_s, cute.Pointer)): + scale_log2 = Float32(scale_s * log2_e) decode_gen_kernel( tma_desc_q, tma_desc_k, @@ -3722,7 +3734,7 @@ def fmha_decode_launch( o_iter, s_k, h_k, - Float32(scale_s * log2_e), + scale_log2, output_scale, seqlens_kv_iter, cu_seqlens_q_iter, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/reduction.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/reduction.py index f28769a10c8d..f27ff6a4ee6b 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/reduction.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/reduction.py @@ -899,8 +899,6 @@ def fmha_decode_separate_reduction_launch( partial_o_iter: cute.Pointer, partial_stats_iter: cute.Pointer, attention_sinks_iter: cute.Pointer, - scale_s: Float32, - output_scale: Float32, stream: cuda_drv.CUstream, cfg: cutlass.Constexpr[FmhaDecodeConfig], static_full_split_prefix: cutlass.Constexpr[bool] = False, @@ -909,9 +907,9 @@ def fmha_decode_separate_reduction_launch( The serial reference grid is ``(slice, kv_head, batch)`` with 8 KiB slices. The production reducer keeps that geometry for compact S2-S4 and otherwise - maps cluster ranks into grid.x for each 2 KiB slice. ``scale_s`` and - ``output_scale`` remain parameters for launch-ABI compatibility; producer - CTAs already applied them to log2-LSE and normalized partial O, respectively. + maps cluster ranks into grid.x for each 2 KiB slice. Producer CTAs already + apply the QK scale to log2-LSE and the output scale to normalized partial O; + the reducer only merges these scaled states. ``seqlens_kv_iter`` is internal reducer metadata used only to bound the runtime split prefix; the decode launch ABI remains unchanged. """ diff --git a/tests/unittest/_torch/attention/backend_case.py b/tests/unittest/_torch/attention/backend_case.py index 4c1b1cdfa02e..beaad315ee56 100644 --- a/tests/unittest/_torch/attention/backend_case.py +++ b/tests/unittest/_torch/attention/backend_case.py @@ -1651,6 +1651,13 @@ def run_backend( def create_metadata(AttentionCls, case, mgr, *, num_contexts: int = 0): seq_lens_kv = torch.tensor(case.seq_lens_kv, dtype=torch.int) if case.is_cross else None + # Context preprocessing needs this chunk's Q length; the cached + # prefix is already included separately in kv_lens by prepare(). + prompt_lens = ( + case.seq_lens + if case.is_cross + else case.seq_lens[:num_contexts] + case.token_nums[num_contexts:] + ) return AttentionCls.Metadata( num_contexts=num_contexts, kv_cache_params=KVCacheParams( @@ -1662,7 +1669,7 @@ def create_metadata(AttentionCls, case, mgr, *, num_contexts: int = 0): max_num_tokens=case.max_num_tokens, kv_cache_manager=mgr, request_ids=request_ids, - prompt_lens=case.seq_lens if case.is_cross else case.token_nums, + prompt_lens=prompt_lens, kv_layout=kv_layout, ) diff --git a/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py b/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py index 30d37453f1f7..47f1eaea946b 100644 --- a/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py +++ b/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py @@ -40,6 +40,7 @@ from tensorrt_llm._torch.attention.kernels.sage_quant import SageQuantized from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.quantization.mode import QuantMode pytestmark = pytest.mark.cpu_only @@ -270,6 +271,19 @@ def test_block_sparse_support_is_phase_specific_and_paged_proxy_is_rejected( assert reason == "paged block-sparse attention only supports BSR exact routes" +def test_paged_block_sparse_rejects_fp8_kv_cache() -> None: + attention, fmha, q, metadata, args = _paged_case() + attention.quant_mode = QuantMode.FP8_KV_CACHE + metadata.kv_cache_manager.dtype = torch.float8_e4m3fn + + supported, reason = fmha._is_supported_with_reason( + q, None, None, metadata, args, phase=FmhaPhase.GENERATION + ) + + assert not supported + assert "quantized KV cache" in reason + + def test_contiguous_proxy_routes_reject_causal_mask_before_planning(monkeypatch) -> None: _attention, fmha, q, k, v, metadata, args = _contiguous_case() _set_block_sparse_inputs(args, _bitmask_inputs(proxy=True)) diff --git a/tests/unittest/_torch/attention/test_cute_dsl_mla.py b/tests/unittest/_torch/attention/test_cute_dsl_mla.py index 63a9dc77ad06..14e7cb621a10 100644 --- a/tests/unittest/_torch/attention/test_cute_dsl_mla.py +++ b/tests/unittest/_torch/attention/test_cute_dsl_mla.py @@ -391,6 +391,7 @@ def test_mixed_batch_gate_admits_only_page_sizes_exclusive_to_cute_dsl( (100, 128, DataType.FP8, (), False), (103, 128, DataType.BF16, (), False), (107, 128, DataType.FP8, (), True), + (107, 128, DataType.BF16, (), True), # An NVFP4 manager's FP8 fallback layers on the FP4 MLA pages. (100, 128, DataType.NVFP4, (), True), (103, 128, DataType.NVFP4, (), True), 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 d0e61dfb5e54..9e34f0d904bc 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py +++ b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py @@ -33,8 +33,14 @@ from utils.util import getSMVersion pytestmark = pytest.mark.skipif( + getSMVersion() not in (100, 103, 107), + reason="PrimsTS attention kernels require SM100, SM103, or SM107", +) + + +_REQUIRES_BLACKWELL_MLA = pytest.mark.skipif( getSMVersion() not in (100, 103), - reason="PrimsTS attention kernels require SM100 or SM103", + reason="MLA PrimsTS integration requires SM100 or SM103", ) @@ -98,17 +104,30 @@ }, id="mixed", ), + pytest.param( + {"seq_lens": [41, 17], "num_cached_tokens": [63, 32], "num_contexts": 2}, + id="cached-context", + ), + pytest.param( + {"seq_lens": [1, 1], "num_cached_tokens": [8191, 4096], "num_contexts": 0}, + id="long-generation", + ), ], ) +@pytest.mark.parametrize("dtype", ["bfloat16", "float16"], ids=["bf16", "fp16"]) +@pytest.mark.parametrize("kv_dtype", [None, "fp8"], ids=["native-kv", "fp8-kv"]) def test_prims_ts_qwen2_gqa( monkeypatch: pytest.MonkeyPatch, use_kv_cache_manager_v2: bool, phase_args: dict, + kv_dtype: str | None, + dtype: str, ) -> None: monkeypatch.setenv("TLLM_FMHA_LIBS", "prims_ts") case = BackendCase( - **_QWEN2_7B, + **{**_QWEN2_7B, "dtype": dtype}, **phase_args, + kv_dtype=kv_dtype, use_kv_cache_manager_v2=use_kv_cache_manager_v2, ) @@ -200,8 +219,10 @@ def run_with_zeroed_v_tail_check( assert zeroed_tails == [(0, 1), (1, 5)] +@pytest.mark.parametrize("kv_dtype", [None, "fp8"], ids=["fp16-kv", "fp8-kv"]) def test_prims_ts_fp16_dense_context_with_alternate_shape( monkeypatch: pytest.MonkeyPatch, + kv_dtype: str | None, ) -> None: monkeypatch.setenv("TLLM_FMHA_LIBS", "prims_ts") case = BackendCase( @@ -212,6 +233,7 @@ def test_prims_ts_fp16_dense_context_with_alternate_shape( num_cached_tokens=[0, 0], num_contexts=2, dtype="float16", + kv_dtype=kv_dtype, causal=False, kv_layout="HND", page_size=64, @@ -363,6 +385,7 @@ def record_decode_workspace( @pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], ids=["v1", "v2"]) +@_REQUIRES_BLACKWELL_MLA def test_prims_ts_deepseek_v3_lite_mla_generation( monkeypatch: pytest.MonkeyPatch, use_kv_cache_manager_v2: bool, @@ -381,6 +404,7 @@ def test_prims_ts_deepseek_v3_lite_mla_generation( @pytest.mark.parametrize("num_heads", [6, 12, 96]) @pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], ids=["v1", "v2"]) +@_REQUIRES_BLACKWELL_MLA def test_prims_ts_fp8_mla_preprocessing( monkeypatch: pytest.MonkeyPatch, num_heads: int, use_kv_cache_manager_v2: bool ) -> None: @@ -447,6 +471,7 @@ def run_and_capture(fmha: PrimsTSFmha, params: FmhaParams) -> None: @pytest.mark.parametrize("num_heads", [6, 12, 96]) @pytest.mark.parametrize("batch_size", [2, 65]) @pytest.mark.parametrize("max_seq_len", [64, 128], ids=["short-kv", "regular-kv"]) +@_REQUIRES_BLACKWELL_MLA def test_prims_ts_fp8_mla_scales_and_graph_replay( monkeypatch: pytest.MonkeyPatch, num_heads: int, batch_size: int, max_seq_len: int ) -> None: @@ -862,6 +887,7 @@ def test_prims_ts_decode_live_wrapper_cuda_graph_replay() -> None: @pytest.mark.parametrize("max_seq_len", [64, 1024], ids=["short-kv", "regular-kv"]) +@_REQUIRES_BLACKWELL_MLA def test_prims_ts_mla_live_wrapper_cuda_graph_replay(max_seq_len: int) -> None: from tensorrt_llm._torch.attention.backends.prims_ts import ( BatchMLADecodePagedTSWrapper, @@ -1449,3 +1475,194 @@ def counted_prims_context(self, *args, **kwargs): assert calls["fallback"] > 0 assert calls["prims_context"] == 0 + + +@pytest.mark.parametrize("head_dim", [64, 128, 256]) +@pytest.mark.parametrize( + ("mode", "page_size", "head_ratio"), + [ + ("disabled", 16, 1), + ("gmem_reduction", 32, 7), + ("gmem_reduction_with_separate_kernel", 64, 8), + ("cluster_smem_reduction", 128, 16), + ], +) +def test_prims_ts_fp8_decode_bf16_device_scales( + monkeypatch: pytest.MonkeyPatch, head_dim: int, mode: str, page_size: int, head_ratio: int +) -> None: + """Qualify BF16 stores/reducers and live device scales across graph replay.""" + from dataclasses import replace + + import cutlass + + from tensorrt_llm._torch.attention.backends.prims_ts import decode + from tensorrt_llm._torch.attention.backends.prims_ts.kernels.fmha_decode.fmha_decode_config import ( + make_decode_config, + ) + + torch.manual_seed(123) + batch_size, kv_heads = 2, 2 + qo_heads = kv_heads * head_ratio + max_kv_len = 2048 + pages_per_seq = max_kv_len // page_size + num_pages = batch_size * pages_per_seq + q = torch.randn(batch_size, qo_heads, head_dim, device="cuda").to(torch.float8_e4m3fn) + k = torch.randn(num_pages, kv_heads, page_size, head_dim, device="cuda").to(torch.float8_e4m3fn) + v = torch.randn_like(k, dtype=torch.float32).to(torch.float8_e4m3fn) + # Permuted pages and padded table rows catch assumptions about pool layout. + table = torch.full((batch_size, pages_per_seq + 3), -1, dtype=torch.int32, device="cuda") + table[:, :pages_per_seq] = torch.randperm(num_pages, device="cuda").reshape(batch_size, -1) + lengths_host = [max_kv_len - 13, max_kv_len // 2 + 7] + lengths = torch.tensor(lengths_host, dtype=torch.int32, device="cuda") + bmm1 = torch.tensor([0.25 / math.sqrt(head_dim)], device="cuda") + bmm2 = torch.tensor([1.75], device="cuda") + output = torch.empty_like(q, dtype=torch.bfloat16) + splits = 1 if mode == "disabled" else 4 + cfg = make_decode_config( + headdim=head_dim, + args={"tile_size_q": 16, "groups_tokens_heads_q": True, "use_keeps_mma_ab": False}, + seq_len_q=1, + seq_len_kv=max_kv_len, + batch_size=batch_size, + num_heads_q=qo_heads, + num_heads_kv=kv_heads, + qkv_dtype=cutlass.Float8E4M3FN, + o_dtype=cutlass.BFloat16, + qkv_layout="pagedKv", + num_tokens_per_page=page_size, + split_kv_mode=mode, + splits_kv=splits, + max_splits_kv=splits, + mask_type="causal", + auto_tuner=False, + ) + + def plan(output_dtype: torch.dtype, use_device_scales: bool): + spec = decode._decode_launch_spec_from_config( + replace( + cfg, + out_dtype=cutlass.BFloat16 if output_dtype == torch.bfloat16 else cutlass.Float16, + ), + batch_size=batch_size, + num_qo_heads=qo_heads, + num_kv_heads=kv_heads, + head_dim=head_dim, + seq_len_q=1, + max_active_clusters=torch.cuda.get_device_properties(0).multi_processor_count, + ) + monkeypatch.setattr(decode, "_resolve_decode_launch_spec", lambda *args, **kwargs: spec) + wrapper = decode.BatchDecodePagedTSWrapper() + wrapper.plan( + q.device, + batch_size, + qo_heads, + kv_heads, + head_dim, + page_size, + max_kv_len, + q_data_type=q.dtype, + o_data_type=output_dtype, + mask_type="causal", + use_device_scales=use_device_scales, + ) + return wrapper + + wrapper = plan(torch.bfloat16, True) + fp16_wrapper = plan(torch.float16, False) + bf16_host_wrapper = plan(torch.bfloat16, False) + + def run() -> torch.Tensor: + if mode == "gmem_reduction": + wrapper._plan_state.workspace.split_kv_counter.zero_() + return wrapper.run( + q, + (k, v), + lengths, + table, + out=output, + bmm1_scale=bmm1, + bmm2_scale=bmm2, + validate=False, + ) + + def reference() -> torch.Tensor: + results = [] + for batch_idx, length in enumerate(lengths_host): + ids = table[batch_idx, :pages_per_seq].long() + keys = k.float()[ids].permute(1, 0, 2, 3).reshape(kv_heads, -1, head_dim)[:, :length] + values = v.float()[ids].permute(1, 0, 2, 3).reshape(kv_heads, -1, head_dim)[:, :length] + keys = keys.repeat_interleave(head_ratio, dim=0) + values = values.repeat_interleave(head_ratio, dim=0) + scores = torch.einsum("hd,htd->ht", q[batch_idx].float(), keys) * bmm1 + results.append(torch.einsum("ht,htd->hd", scores.softmax(-1), values) * bmm2) + return torch.stack(results) + + def check_output() -> None: + # The FP8 kernel rounds the intermediate P operand to E4M3. Compare + # math with that quantization tolerance, and the existing FP16 output + # path tightly to isolate BF16 stores and device-scale handling. + torch.testing.assert_close(output.float(), reference(), atol=0.01, rtol=0.05) + if mode == "gmem_reduction": + fp16_wrapper._plan_state.workspace.split_kv_counter.zero_() + fp16_output = fp16_wrapper.run( + q, + (k, v), + lengths, + table, + bmm1_scale=bmm1.item(), + bmm2_scale=bmm2.item(), + validate=False, + ) + torch.testing.assert_close(output.float(), fp16_output.float(), atol=0.0005, rtol=0.008) + + if mode == "gmem_reduction": + bf16_host_wrapper._plan_state.workspace.split_kv_counter.zero_() + bf16_host_output = bf16_host_wrapper.run( + q, + (k, v), + lengths, + table, + bmm1_scale=bmm1.item(), + bmm2_scale=bmm2.item(), + validate=False, + ) + torch.testing.assert_close( + output.float(), bf16_host_output.float(), atol=0.0005, rtol=0.008 + ) + + if mode == "gmem_reduction": + wrapper._plan_state.workspace.split_kv_counter.zero_() + wrapper.run( + q, + (k, v), + lengths, + table, + out=output, + bmm1_scale=bmm1, + bmm2_scale=bmm2, + validate=True, + ) + torch.cuda.synchronize() + check_output() + + def no_host_read(*args, **kwargs): + raise AssertionError("decode must not read scales back to the host") + + with monkeypatch.context() as guard: + guard.setattr(torch.Tensor, "item", no_host_read) + guard.setattr(torch.Tensor, "__float__", no_host_read) + sync_debug_mode = torch.cuda.get_sync_debug_mode() + try: + torch.cuda.set_sync_debug_mode("error") + run() + finally: + torch.cuda.set_sync_debug_mode(sync_debug_mode) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + # Both attention scores and final output must consume updated scales. + bmm1.mul_(2) + bmm2.mul_(0.5) + graph.replay() + torch.cuda.synchronize() + check_output() diff --git a/tests/unittest/_torch/attention/test_prims_ts_fmha.py b/tests/unittest/_torch/attention/test_prims_ts_fmha.py index c1d1fe078709..06035134edb6 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_fmha.py +++ b/tests/unittest/_torch/attention/test_prims_ts_fmha.py @@ -196,6 +196,9 @@ def _support_result( relative_attention_bias=torch.empty(1) if has_relative_attention_bias else None, is_fused_qkv=is_fused_qkv, ) + if quant_mode.has_fp8_kv_cache(): + forward_args.kv_scale_orig_quant = _TensorSpec((1,), torch.float32) + forward_args.kv_scale_quant_orig = _TensorSpec((1,), torch.float32) if is_mla and quant_mode.has_fp8_kv_cache(): forward_args.quant_q_buffer = _TensorSpec(q.shape, torch.uint8) forward_args.mla_bmm1_scale = _TensorSpec((2,), torch.float32) @@ -321,6 +324,20 @@ def test_supported_matrix(case: dict) -> None: assert supported, reason +@pytest.mark.parametrize("phase", list(AttentionInputType)) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("use_v2", [False, True]) +def test_fp8_standard_supported(phase, dtype, use_v2) -> None: + supported, reason = _support_result( + attention_input_type=phase, + dtype=dtype, + kv_dtype=DataType.FP8, + quant_mode=QuantMode.FP8_KV_CACHE, + use_kv_cache_v2=use_v2, + ) + assert supported, reason + + @pytest.mark.parametrize("num_heads", [6, 12, 64, 96, 128]) @pytest.mark.parametrize("use_kv_cache_v2", [False, True]) @pytest.mark.parametrize("tokens_per_block", [16, 32, 64, 128]) @@ -345,10 +362,12 @@ def test_fp8_mla_supported( @pytest.mark.parametrize( ("overrides", "expected_reason"), [ - ({"is_mla": False, "head_dim": 128}, "only for FP8 MLA"), - ({"quant_mode": QuantMode.INT8_KV_CACHE}, "only for FP8 MLA"), - ({"quant_mode": QuantMode.NVFP4_KV_CACHE}, "only for FP8 MLA"), - ({"quant_mode": QuantMode.FP8_KV_CACHE | QuantMode.INT8_KV_CACHE}, "only for FP8 MLA"), + ({"quant_mode": QuantMode.INT8_KV_CACHE}, "only FP8 E4M3 KV-cache quantization"), + ({"quant_mode": QuantMode.NVFP4_KV_CACHE}, "only FP8 E4M3 KV-cache quantization"), + ( + {"quant_mode": QuantMode.FP8_KV_CACHE | QuantMode.INT8_KV_CACHE}, + "only FP8 E4M3 KV-cache quantization", + ), ({"kv_dtype": DataType.BF16}, "FP8 E4M3 KV cache"), ({"quant_mode": QuantMode(0)}, "query and KV-cache dtypes"), ({"dtype": torch.float16}, "BF16 query input and output"), @@ -703,22 +722,24 @@ def test_unsupported_matrix_falls_through(case: dict, expected_reason: str) -> N @pytest.mark.parametrize( - ("sm", "cutlass_version", "compiler_version", "expected"), + ("sm", "compiler_version", "is_mla", "expected"), [ - (100, "4.7.0", "13.3", True), - (103, "4.7.0", "13.3", True), - (100, "4.7.0", "13.4", True), - (100, "4.7.0", "13.2", False), - (107, "4.7.0", "13.3", False), - (120, "4.7.0", "13.3", False), - (100, "4.6.2", "13.3", False), + (100, "13.3", False, True), + (103, "13.3", False, True), + (100, "13.4", False, True), + (100, "13.2", False, False), + (107, "13.3", False, True), + (107, "13.2", False, False), + (107, "13.3", True, False), + (100, "13.3", True, True), + (120, "13.3", False, False), ], ) def test_static_availability_gate( monkeypatch: pytest.MonkeyPatch, sm: int, - cutlass_version: str, compiler_version: str, + is_mla: bool, expected: bool, ) -> None: target_version = Mock( @@ -733,12 +754,11 @@ def import_cutlass_module(module_name: str) -> object: return object() monkeypatch.setattr(prims_ts_module, "get_sm_version", lambda: sm) - monkeypatch.setattr(prims_ts_module, "version", lambda _: cutlass_version) monkeypatch.setattr(prims_ts_module, "import_module", import_cutlass_module) monkeypatch.setattr(PrimsTSFmha, "_missing_fused_nanobind_ops", staticmethod(lambda: [])) - assert PrimsTSFmha.is_available(_Attention()) is expected - if sm in (100, 103) and Version(cutlass_version) >= Version("4.7.0"): + assert PrimsTSFmha.is_available(_Attention(is_mla=is_mla)) is expected + if sm in (100, 103) or (sm == 107 and not is_mla): target_version.assert_called_once_with(min_version="13.3") else: target_version.assert_not_called() @@ -749,7 +769,6 @@ def test_static_availability_gate_fails_closed_when_compiler_query_fails( ) -> None: cutlass = SimpleNamespace(target_version=Mock(side_effect=RuntimeError("query failed"))) monkeypatch.setattr(prims_ts_module, "get_sm_version", lambda: 100) - monkeypatch.setattr(prims_ts_module, "version", lambda _: "4.7.0") monkeypatch.setattr( prims_ts_module, "import_module", @@ -770,7 +789,6 @@ def test_unsupported_cutlass_compiler_excludes_prims_ts_from_fmha_selection( cutlass = SimpleNamespace(target_version=target_version) monkeypatch.setenv("TLLM_FMHA_LIBS", "prims_ts,fallback") monkeypatch.setattr(prims_ts_module, "get_sm_version", lambda: 100) - monkeypatch.setattr(prims_ts_module, "version", lambda _: "4.7.0") monkeypatch.setattr( prims_ts_module, "import_module", @@ -915,15 +933,21 @@ def test_sequence_lengths_are_live_source_view() -> None: torch.testing.assert_close(actual, torch.tensor([33, 65], dtype=torch.int32)) +@pytest.mark.parametrize("fp8", [False, True], ids=["bf16", "fp8"]) def test_context_wrapper_plans_once_and_reads_live_fixed_metadata( monkeypatch: pytest.MonkeyPatch, + fp8: bool, ) -> None: monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) attn = _Attention() + attn.quant_mode = QuantMode.FP8_KV_CACHE if fp8 else QuantMode(0) + kernel_dtype = torch.float8_e4m3fn if fp8 else torch.bfloat16 + bmm1 = torch.tensor([0.125, 0.125 * math.log2(math.e)]) + bmm2 = torch.tensor([0.75]) fmha = PrimsTSFmha(attn) fmha._multi_processor_count = 120 q_processed = torch.empty((3, attn.num_heads, attn.head_dim), dtype=torch.bfloat16) - kv_pool = torch.empty((12, attn.num_kv_heads, 32, attn.head_dim), dtype=torch.bfloat16) + kv_pool = torch.empty((12, attn.num_kv_heads, 32, attn.head_dim), dtype=kernel_dtype) block_tables = torch.tensor( [ [[0, 1, 2, 3], [6, 7, 8, 9]], @@ -940,8 +964,8 @@ def test_context_wrapper_plans_once_and_reads_live_fixed_metadata( kv_pool, block_tables, None, - 1.0, - 1.0, + bmm1, + bmm2, fmha_workspace, cu_q_seqlens, cu_kv_seqlens, @@ -1037,9 +1061,9 @@ def test_context_wrapper_plans_once_and_reads_live_fixed_metadata( "num_qo_heads": attn.num_heads, "num_kv_heads": attn.num_kv_heads, "head_dim": attn.head_dim, - "q_dtype": torch.bfloat16, - "k_dtype": torch.bfloat16, - "v_dtype": torch.bfloat16, + "q_dtype": kernel_dtype, + "k_dtype": kernel_dtype, + "v_dtype": kernel_dtype, "out_dtype": torch.bfloat16, "page_size": 32, "mask_type": "causal", @@ -1051,7 +1075,13 @@ def test_context_wrapper_plans_once_and_reads_live_fixed_metadata( wrapper.run.assert_called_once() run_args = wrapper.run.call_args.args run_kwargs = wrapper.run.call_args.kwargs - assert run_args[0] is q_processed + assert run_args[0].data_ptr() == q_processed.data_ptr() + assert run_args[0].dtype == kernel_dtype + assert context_preprocess.call_args.args[35] is fp8 + assert context_postprocess.call_args.args[32] is fp8 + if fp8: + assert run_kwargs["scale_softmax_log2"].data_ptr() == bmm1[1:].data_ptr() + assert run_kwargs["output_scale"].data_ptr() == bmm2.data_ptr() k_cache, v_cache = run_args[1], run_args[2] assert k_cache.shape == v_cache.shape == (6, attn.num_kv_heads, 32, attn.head_dim) assert v_cache.storage_offset() - k_cache.storage_offset() == 6 * math.prod(kv_pool.shape[1:]) @@ -1119,8 +1149,10 @@ def test_context_wrapper_plans_once_and_reads_live_fixed_metadata( pytest.param(True, False, True, False, id="cluster-smem-reduction"), ), ) +@pytest.mark.parametrize("fp8", [False, True], ids=["bf16", "fp8"]) def test_generation_wrapper_plans_once_and_reads_live_fixed_metadata( monkeypatch: pytest.MonkeyPatch, + fp8: bool, use_split_kv: bool, use_separate_reduction_kernel: bool, use_cluster_smem_reduction: bool, @@ -1128,13 +1160,17 @@ def test_generation_wrapper_plans_once_and_reads_live_fixed_metadata( ) -> None: monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) attn = _Attention() + attn.quant_mode = QuantMode.FP8_KV_CACHE if fp8 else QuantMode(0) + kernel_dtype = torch.float8_e4m3fn if fp8 else torch.bfloat16 + bmm1 = torch.tensor([0.125, 0.125 * math.log2(math.e)]) + bmm2 = torch.tensor([0.75]) fmha = PrimsTSFmha(attn) fmha._multi_processor_count = 120 fmha._decode_workspace_offset_bytes = 0 fmha._decode_workspace_required_bytes = 64 fmha_workspace = torch.empty(0, dtype=torch.uint8) q_processed = torch.empty((2, attn.num_heads, attn.head_dim), dtype=torch.bfloat16) - kv_pool = torch.empty((12, attn.num_kv_heads, 32, attn.head_dim), dtype=torch.bfloat16) + kv_pool = torch.empty((12, attn.num_kv_heads, 32, attn.head_dim), dtype=kernel_dtype) block_tables = torch.tensor( [ [[0, 1, 2, 3], [6, 7, 8, 9]], @@ -1150,8 +1186,8 @@ def test_generation_wrapper_plans_once_and_reads_live_fixed_metadata( kv_pool, block_tables, None, - 1.0, - 1.0, + bmm1, + bmm2, fmha_workspace, None, 1, @@ -1252,8 +1288,9 @@ def test_generation_wrapper_plans_once_and_reads_live_fixed_metadata( assert {key: value for key, value in plan_kwargs.items() if key != "workspace_buffer"} == { "max_seq_len_q": 1, "packed_query": False, - "q_data_type": torch.bfloat16, - "kv_data_type": torch.bfloat16, + "use_device_scales": fp8, + "q_data_type": kernel_dtype, + "kv_data_type": kernel_dtype, "o_data_type": torch.bfloat16, "mask_type": "causal", "window_left": 15, @@ -1264,6 +1301,14 @@ def test_generation_wrapper_plans_once_and_reads_live_fixed_metadata( run_kwargs = wrapper.run.call_args.kwargs assert run_args[0].shape == (2, attn.num_heads, attn.head_dim) assert run_args[0].data_ptr() == q_processed.data_ptr() + assert run_args[0].dtype == kernel_dtype + assert generation_preprocess.call_args.args[35] is fp8 + if fp8: + assert run_kwargs["bmm1_scale"].data_ptr() == bmm1.data_ptr() + assert run_kwargs["bmm2_scale"].data_ptr() == bmm2.data_ptr() + else: + assert run_kwargs["bmm1_scale"] == pytest.approx(1.0 / math.sqrt(attn.head_dim)) + assert run_kwargs["bmm2_scale"] == 1.0 k_cache, v_cache = run_args[1] assert k_cache.shape == v_cache.shape == (6, attn.num_kv_heads, 32, attn.head_dim) assert v_cache.storage_offset() - k_cache.storage_offset() == 6 * math.prod(kv_pool.shape[1:]) @@ -1276,8 +1321,6 @@ def test_generation_wrapper_plans_once_and_reads_live_fixed_metadata( fixed_block_tables, torch.tensor([[0, 1, 2, 3], [2, 3, 4, 5]], dtype=torch.int32), ) - assert run_kwargs["bmm1_scale"] == pytest.approx(1.0 / math.sqrt(attn.head_dim)) - assert run_kwargs["bmm2_scale"] == 1.0 assert run_kwargs["out"].shape == (2, attn.num_heads, attn.head_dim) assert run_kwargs["out"].data_ptr() == output.data_ptr() assert run_kwargs["validate"] is False @@ -2066,8 +2109,10 @@ def test_mla_caller_workspace_grows_across_plan_profiles( assert workspace.numel() == 64 +@pytest.mark.parametrize("fp8", [False, True]) def test_decode_prepare_workspace_reserves_tail_after_compact_preprocessing( monkeypatch: pytest.MonkeyPatch, + fp8: bool, ) -> None: monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) monkeypatch.setattr( @@ -2075,12 +2120,14 @@ def test_decode_prepare_workspace_reserves_tail_after_compact_preprocessing( "get_trtllm_gen_generation_workspace_layout", lambda *args, **kwargs: {"total_size": 64}, ) + workspace_size = Mock(return_value=48) monkeypatch.setattr( prims_ts_package, "get_prims_ts_batch_decode_workspace_size", - lambda *args, **kwargs: 48, + workspace_size, ) attn = _Attention() + attn.quant_mode = QuantMode.FP8_KV_CACHE if fp8 else QuantMode(0) fmha = PrimsTSFmha(attn) fmha._multi_processor_count = 1 metadata = SimpleNamespace( @@ -2114,6 +2161,11 @@ def test_decode_prepare_workspace_reserves_tail_after_compact_preprocessing( assert decode_workspace.data_ptr() == workspace.data_ptr() + 64 assert decode_workspace.numel() == 48 + expected_dtype = torch.float8_e4m3fn if fp8 else torch.bfloat16 + assert workspace_size.call_args.kwargs["q_dtype"] == expected_dtype + assert workspace_size.call_args.kwargs["kv_dtype"] == expected_dtype + assert workspace_size.call_args.kwargs["out_dtype"] == torch.bfloat16 + def test_decode_workspace_tail_is_stable_across_mixed_context_layouts( monkeypatch: pytest.MonkeyPatch,