From 4784870ddea6ce2d423dbf582dcdf7e8d44bd703 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Fri, 18 Sep 2026 04:43:37 +0000 Subject: [PATCH 1/7] [None][feat] Support FP8 KV cache for non-MLA PrimsTS attention Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- 3rdparty/vendor_sources.lock.yaml | 6 +- .../attention/backends/fmha/prims_ts.py | 97 +++++++--- .../attention/backends/prims_ts/decode.py | 70 +++++-- .../kernels/fmha_decode/fmha_decode_kernel.py | 31 +++- .../unittest/_torch/attention/backend_case.py | 9 +- .../test_prims_ts_attention_backend.py | 171 ++++++++++++++++++ .../_torch/attention/test_prims_ts_fmha.py | 86 +++++++-- 7 files changed, 405 insertions(+), 65 deletions(-) diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 18640486388d..210d81726c95 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: ace9ae803bdff01874bf8dd01e89ba59476dd174 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:1f4cbe2cda29673cc5c0e0086e221f61d863a7dbfc589b272a6bf7f43d0ec369 openengine-proto: url: https://github.com/ai-dynamo/openengine.git tag: v0.1.0 diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py index 89ee7d2a8126..a15509bd47a9 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py @@ -105,6 +105,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 +125,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,7 +157,12 @@ def get_attention_feature_unsupported_reason( class PrimsTSFmha(PhasedFmha): - """Blackwell task-scheduled paged context and decode FMHA library.""" + """Blackwell 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 @@ -331,11 +340,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 +385,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 +418,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 +783,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 +963,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 +1010,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 +1045,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 +1133,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 +1177,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 +1209,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 +1233,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 +1272,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 +1326,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 +1354,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 +1382,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, @@ -1398,10 +1431,26 @@ def run_generation(self, params: FmhaParams) -> None: block_tables=fixed_block_tables, bmm1_scale=self._get_bmm1_scale(attn), bmm2_scale=1.0, + bmm1_scale_device=bmm1_scale.reshape(-1)[:1] if use_fp8 else None, + bmm2_scale_device=bmm2_scale.reshape(-1)[:1] if use_fp8 else None, 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..98233168581b 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py @@ -152,8 +152,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 +192,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( @@ -1509,6 +1510,7 @@ def _make_decode_compile_spec( @functools.cache def _get_compiled_decode( compile_spec: _DecodeCompileSpec, + use_device_scales: bool = False, ): """Compile and cache one batch-dynamic TS decode topology.""" @@ -1588,8 +1590,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 +1669,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, @@ -1899,6 +1901,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 +1927,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, @@ -2247,8 +2250,9 @@ def _launch_decode( workspace.partial_o, workspace.partial_stats, workspace.attention_sinks, - runtime.bmm1_scale, - runtime.bmm2_scale, + # Producers already scale normalized partial O and log2-LSE. + 1.0, + 1.0, (), ) return runtime.out @@ -3773,6 +3777,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 +3793,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 +3873,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 ) @@ -3996,7 +4012,11 @@ def plan( kv_prefix_mode=kv_prefix_mode, kv_lengths_mode=kv_lengths_mode, ) - compiled_main, compiled_reducer = _get_compiled_decode(compile_spec) + compiled_main, compiled_reducer = ( + _get_compiled_decode(compile_spec, use_device_scales=True) + if use_device_scales + else _get_compiled_decode(compile_spec) + ) policy = spec.policy + ( ("kv_prefix_mode", kv_prefix_mode), ("kv_lengths_mode", kv_lengths_mode), @@ -4067,6 +4087,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. @@ -4084,6 +4105,8 @@ def run( qo_indptr: Optional[torch.Tensor] = None, bmm1_scale: Optional[float] = None, bmm2_scale: float = 1.0, + bmm1_scale_device: Optional[torch.Tensor] = None, + bmm2_scale_device: Optional[torch.Tensor] = None, out: Optional[torch.Tensor] = None, validate: bool = True, ) -> torch.Tensor: @@ -4132,6 +4155,10 @@ def run( bmm1_scale : float, optional QK scaling factor. Defaults to the inverse square root of ``head_dim``. + bmm1_scale_device, bmm2_scale_device : torch.Tensor, optional + Float32 CUDA tensors of shape ``(1,)`` required by a plan with + ``use_device_scales=True``. Override the host scales, and may change + in place on graph replay. BMM1 uses the regular, not log2, scale. bmm2_scale : float Value/output scaling factor. Defaults to ``1.0``. out : torch.Tensor, optional @@ -4243,6 +4270,27 @@ def run( out=out, ) + if state.use_device_scales: + for name, scale in ( + ("bmm1_scale_device", bmm1_scale_device), + ("bmm2_scale_device", bmm2_scale_device), + ): + if ( + scale is None + 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}" + ) + runtime = replace( + runtime, bmm1_scale=bmm1_scale_device, bmm2_scale=bmm2_scale_device + ) + elif bmm1_scale_device is not None or bmm2_scale_device is not None: + raise ValueError("device scales require plan(use_device_scales=True)") + return _launch_decode( runtime, seq_lens=effective_seq_lens, 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..2e5d51cb824a 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 @@ -1,3 +1,4 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -2551,8 +2552,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 +2601,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 +2903,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 +3096,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 +3342,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 +3721,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 +3735,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/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/test_prims_ts_attention_backend.py b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py index d0e61dfb5e54..d10367614960 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py +++ b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py @@ -98,17 +98,28 @@ }, 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("kv_dtype", [None, "fp8"], ids=["bf16-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, ) -> None: monkeypatch.setenv("TLLM_FMHA_LIBS", "prims_ts") case = BackendCase( **_QWEN2_7B, **phase_args, + kv_dtype=kv_dtype, use_kv_cache_manager_v2=use_kv_cache_manager_v2, ) @@ -200,8 +211,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 +225,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, @@ -1449,3 +1463,160 @@ 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: 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) + + 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_device=bmm1, + bmm2_scale_device=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) + + run() + 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) + 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..52a487503783 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"), @@ -915,15 +934,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 +965,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 +1062,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 +1076,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 +1150,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 +1161,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 +1187,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 +1289,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 +1302,11 @@ 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_device"].data_ptr() == bmm1.data_ptr() + assert run_kwargs["bmm2_scale_device"].data_ptr() == bmm2.data_ptr() 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:]) @@ -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, From 03f3b60fbf4096246a3421d5caf261e1efd70003 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Tue, 13 Oct 2026 09:17:03 +0000 Subject: [PATCH 2/7] [None][fix] Adapt FP8 decode integration to refreshed PrimTS APIs Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../_torch/attention/test_prims_ts_attention_backend.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 d10367614960..cc099b599b1d 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py +++ b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py @@ -1538,7 +1538,7 @@ def plan(output_dtype: torch.dtype, use_device_scales: bool): 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: spec) + monkeypatch.setattr(decode, "_resolve_decode_launch_spec", lambda *args, **kwargs: spec) wrapper = decode.BatchDecodePagedTSWrapper() wrapper.plan( q.device, From 44933be5b4da7a3611f064f7512e0c55220d4b5c Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Thu, 29 Oct 2026 07:07:39 +0000 Subject: [PATCH 3/7] [None][fix] Separate persistent FP8 decode scale specializations Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- 3rdparty/vendor_sources.lock.yaml | 4 ++-- .../_torch/attention/backends/prims_ts/decode.py | 7 ++++++- .../attention/test_prims_ts_attention_backend.py | 16 ++++++++++++++++ 3 files changed, 24 insertions(+), 3 deletions(-) diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 210d81726c95..4c2ed42661fc 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -15,14 +15,14 @@ vendors: flashinfer-prims-ts: url: https://github.com/yuxianq/flashinfer.git branch: prims-tsfmha-fp8-non-mla--kernels - commit: ace9ae803bdff01874bf8dd01e89ba59476dd174 + commit: ad4b6bc638237f32496b7f7a379a76e28cccb360 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:1f4cbe2cda29673cc5c0e0086e221f61d863a7dbfc589b272a6bf7f43d0ec369 + digest: sha256-tree-v1:d25b65576e4ad5dd6aeb95dee455a209fb54a05804d95aaf71b5dc4a6dafc105 openengine-proto: url: https://github.com/ai-dynamo/openengine.git tag: v0.1.0 diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py b/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py index 98233168581b..76f0eb0f4f1b 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py @@ -1968,7 +1968,12 @@ 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, + extra=(use_device_scales,), ) compiled_reducer = None if cfg.use_separate_reduction_kernel: 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 cc099b599b1d..89cb7cbf274a 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py +++ b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py @@ -1557,6 +1557,7 @@ def plan(output_dtype: torch.dtype, use_device_scales: bool): 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": @@ -1602,6 +1603,21 @@ def check_output() -> None: ) 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 + ) + run() torch.cuda.synchronize() check_output() From bc2dd7051c6f796f325de25bca329e0065ad9a26 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Fri, 30 Oct 2026 10:05:06 +0000 Subject: [PATCH 4/7] [None][fix] Enable non-MLA PrimsTS attention on Rubin Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../attention/backends/fmha/prims_ts.py | 23 +++++++++++----- .../test_prims_ts_attention_backend.py | 18 ++++++++++--- .../_torch/attention/test_prims_ts_fmha.py | 27 ++++++++++++------- 3 files changed, 48 insertions(+), 20 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py index a15509bd47a9..9cf0adabb988 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py @@ -13,7 +13,7 @@ # 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 @@ -58,6 +58,7 @@ _MIN_CUTLASS_DSL_VERSION = Version("4.7.0") +_MIN_RUBIN_CUTLASS_DSL_VERSION = Version("4.8.0") _MIN_CUTLASS_COMPILER_VERSION = "13.3" _WORKSPACE_ALIGNMENT = 32 @@ -157,7 +158,7 @@ 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 @@ -166,7 +167,7 @@ class PrimsTSFmha(PhasedFmha): supports_per_tensor_fp8 = True - SUPPORTED_SMS = frozenset({100, 103}) + SUPPORTED_SMS = frozenset({100, 103, 107}) SUPPORTED_PAGE_SIZES = {16, 32, 64, 128} SUPPORTED_CONTEXT_HEAD_DIMS = {128, 256} SUPPORTED_DECODE_HEAD_DIMS = {64, 128, 256} @@ -201,17 +202,25 @@ 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}.") + logger.debug( + f"PrimTS FMHA is unavailable: requires SM100, SM103, or SM107, got SM{sm}." + ) + return False + if sm == 107 and attn.is_mla_enable: + logger.debug("PrimTS MLA FMHA is unavailable on SM107.") return False + min_dsl_version = _MIN_RUBIN_CUTLASS_DSL_VERSION if sm == 107 else _MIN_CUTLASS_DSL_VERSION try: installed_version = Version(version("nvidia-cutlass-dsl")) except (PackageNotFoundError, InvalidVersion): - logger.debug("PrimTS FMHA is unavailable: nvidia-cutlass-dsl>=4.7 is required.") + logger.debug( + f"PrimTS FMHA is unavailable: nvidia-cutlass-dsl>={min_dsl_version} is required." + ) return False - if installed_version < _MIN_CUTLASS_DSL_VERSION: + if installed_version < min_dsl_version: logger.debug( "PrimTS FMHA is unavailable: " - f"nvidia-cutlass-dsl>={_MIN_CUTLASS_DSL_VERSION} is required, " + f"nvidia-cutlass-dsl>={min_dsl_version} is required, " f"got {installed_version}." ) return False 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 89cb7cbf274a..1a7a1e147e7b 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", ) @@ -108,16 +114,18 @@ ), ], ) -@pytest.mark.parametrize("kv_dtype", [None, "fp8"], ids=["bf16-kv", "fp8-kv"]) +@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, @@ -377,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, @@ -395,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: @@ -461,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: @@ -876,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, diff --git a/tests/unittest/_torch/attention/test_prims_ts_fmha.py b/tests/unittest/_torch/attention/test_prims_ts_fmha.py index 52a487503783..47b4a175286a 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_fmha.py +++ b/tests/unittest/_torch/attention/test_prims_ts_fmha.py @@ -722,15 +722,19 @@ def test_unsupported_matrix_falls_through(case: dict, expected_reason: str) -> N @pytest.mark.parametrize( - ("sm", "cutlass_version", "compiler_version", "expected"), + ("sm", "cutlass_version", "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, "4.7.0", "13.3", False, True), + (103, "4.7.0", "13.3", False, True), + (100, "4.7.0", "13.4", False, True), + (100, "4.7.0", "13.2", False, False), + (107, "4.7.0", "13.3", False, False), + (107, "4.8.0", "13.3", False, True), + (107, "4.8.0", "13.2", False, False), + (107, "4.8.0", "13.3", True, False), + (100, "4.7.0", "13.3", True, True), + (120, "4.7.0", "13.3", False, False), + (100, "4.6.2", "13.3", False, False), ], ) def test_static_availability_gate( @@ -738,6 +742,7 @@ def test_static_availability_gate( sm: int, cutlass_version: str, compiler_version: str, + is_mla: bool, expected: bool, ) -> None: target_version = Mock( @@ -756,8 +761,10 @@ def import_cutlass_module(module_name: str) -> object: 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) and Version(cutlass_version) >= Version("4.7.0")) or ( + sm == 107 and not is_mla and Version(cutlass_version) >= Version("4.8.0") + ): target_version.assert_called_once_with(min_version="13.3") else: target_version.assert_not_called() From 981825b7e9bc1704cdec189aed16194360900e8c Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Fri, 30 Oct 2026 10:58:30 +0000 Subject: [PATCH 5/7] [None][fix] Simplify PrimsTS FP8 scale arguments Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- 3rdparty/vendor_sources.lock.yaml | 4 +- .../attention/backends/fmha/prims_ts.py | 25 +--- .../attention/backends/prims_ts/decode.py | 108 ++++++++++-------- .../kernels/fmha_decode/fmha_decode_kernel.py | 1 - .../test_prims_ts_attention_backend.py | 24 +++- .../_torch/attention/test_prims_ts_fmha.py | 39 +++---- 6 files changed, 100 insertions(+), 101 deletions(-) diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 4c2ed42661fc..96dea4d50cd4 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -15,14 +15,14 @@ vendors: flashinfer-prims-ts: url: https://github.com/yuxianq/flashinfer.git branch: prims-tsfmha-fp8-non-mla--kernels - commit: ad4b6bc638237f32496b7f7a379a76e28cccb360 + commit: 2bbeaf3c78b237298ae18e5cbe7b32d09caff562 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:d25b65576e4ad5dd6aeb95dee455a209fb54a05804d95aaf71b5dc4a6dafc105 + digest: sha256-tree-v1:afae80127f03ecc75c57f516d564050fe7475e0f0f7a4aec52beeb1b76f50ddc openengine-proto: url: https://github.com/ai-dynamo/openengine.git tag: v0.1.0 diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py index 9cf0adabb988..2dbf1b262637 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py @@ -19,11 +19,9 @@ 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,8 +55,6 @@ ) -_MIN_CUTLASS_DSL_VERSION = Version("4.7.0") -_MIN_RUBIN_CUTLASS_DSL_VERSION = Version("4.8.0") _MIN_CUTLASS_COMPILER_VERSION = "13.3" _WORKSPACE_ALIGNMENT = 32 @@ -209,21 +205,6 @@ def _is_available(cls, attn: "TrtllmAttention") -> bool: if sm == 107 and attn.is_mla_enable: logger.debug("PrimTS MLA FMHA is unavailable on SM107.") return False - min_dsl_version = _MIN_RUBIN_CUTLASS_DSL_VERSION if sm == 107 else _MIN_CUTLASS_DSL_VERSION - try: - installed_version = Version(version("nvidia-cutlass-dsl")) - except (PackageNotFoundError, InvalidVersion): - logger.debug( - f"PrimTS FMHA is unavailable: nvidia-cutlass-dsl>={min_dsl_version} is required." - ) - return False - if installed_version < min_dsl_version: - logger.debug( - "PrimTS FMHA is unavailable: " - f"nvidia-cutlass-dsl>={min_dsl_version} is required, " - f"got {installed_version}." - ) - return False try: cutlass = import_module("cutlass") compiler_version_supported = cutlass.target_version( @@ -1438,10 +1419,8 @@ 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_device=bmm1_scale.reshape(-1)[:1] if use_fp8 else None, - bmm2_scale_device=bmm2_scale.reshape(-1)[:1] if use_fp8 else None, + 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, ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py b/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py index 76f0eb0f4f1b..f8802c532722 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py @@ -2120,9 +2120,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.""" @@ -2166,11 +2167,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, @@ -2318,9 +2323,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.""" @@ -2355,11 +2361,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), ) @@ -4108,10 +4118,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_device: Optional[torch.Tensor] = None, - bmm2_scale_device: Optional[torch.Tensor] = None, + bmm1_scale: Optional[float | torch.Tensor] = None, + bmm2_scale: float | torch.Tensor = 1.0, out: Optional[torch.Tensor] = None, validate: bool = True, ) -> torch.Tensor: @@ -4124,7 +4132,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. @@ -4157,21 +4165,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``. - bmm1_scale_device, bmm2_scale_device : torch.Tensor, optional - Float32 CUDA tensors of shape ``(1,)`` required by a plan with - ``use_device_scales=True``. Override the host scales, and may change - in place on graph replay. BMM1 uses the regular, not log2, scale. - 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 ------- @@ -4182,6 +4191,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: @@ -4250,6 +4279,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, @@ -4273,29 +4303,9 @@ def run( bmm1_scale=bmm1_scale, bmm2_scale=bmm2_scale, out=out, + use_device_scales=state.use_device_scales, ) - if state.use_device_scales: - for name, scale in ( - ("bmm1_scale_device", bmm1_scale_device), - ("bmm2_scale_device", bmm2_scale_device), - ): - if ( - scale is None - 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}" - ) - runtime = replace( - runtime, bmm1_scale=bmm1_scale_device, bmm2_scale=bmm2_scale_device - ) - elif bmm1_scale_device is not None or bmm2_scale_device is not None: - raise ValueError("device scales require plan(use_device_scales=True)") - return _launch_decode( runtime, seq_lens=effective_seq_lens, 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 2e5d51cb824a..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 @@ -1,4 +1,3 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); 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 1a7a1e147e7b..9e34f0d904bc 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py +++ b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py @@ -1580,8 +1580,8 @@ def run() -> torch.Tensor: lengths, table, out=output, - bmm1_scale_device=bmm1, - bmm2_scale_device=bmm2, + bmm1_scale=bmm1, + bmm2_scale=bmm2, validate=False, ) @@ -1630,7 +1630,18 @@ def check_output() -> None: output.float(), bf16_host_output.float(), atol=0.0005, rtol=0.008 ) - run() + 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() @@ -1639,6 +1650,13 @@ def no_host_read(*args, **kwargs): 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() diff --git a/tests/unittest/_torch/attention/test_prims_ts_fmha.py b/tests/unittest/_torch/attention/test_prims_ts_fmha.py index 47b4a175286a..06035134edb6 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_fmha.py +++ b/tests/unittest/_torch/attention/test_prims_ts_fmha.py @@ -722,25 +722,22 @@ def test_unsupported_matrix_falls_through(case: dict, expected_reason: str) -> N @pytest.mark.parametrize( - ("sm", "cutlass_version", "compiler_version", "is_mla", "expected"), + ("sm", "compiler_version", "is_mla", "expected"), [ - (100, "4.7.0", "13.3", False, True), - (103, "4.7.0", "13.3", False, True), - (100, "4.7.0", "13.4", False, True), - (100, "4.7.0", "13.2", False, False), - (107, "4.7.0", "13.3", False, False), - (107, "4.8.0", "13.3", False, True), - (107, "4.8.0", "13.2", False, False), - (107, "4.8.0", "13.3", True, False), - (100, "4.7.0", "13.3", True, True), - (120, "4.7.0", "13.3", False, False), - (100, "4.6.2", "13.3", False, 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, @@ -757,14 +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_mla=is_mla)) is expected - if (sm in (100, 103) and Version(cutlass_version) >= Version("4.7.0")) or ( - sm == 107 and not is_mla and Version(cutlass_version) >= Version("4.8.0") - ): + 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() @@ -775,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", @@ -796,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", @@ -1312,8 +1304,11 @@ def test_generation_wrapper_plans_once_and_reads_live_fixed_metadata( assert run_args[0].dtype == kernel_dtype assert generation_preprocess.call_args.args[35] is fp8 if fp8: - assert run_kwargs["bmm1_scale_device"].data_ptr() == bmm1.data_ptr() - assert run_kwargs["bmm2_scale_device"].data_ptr() == bmm2.data_ptr() + 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:]) @@ -1326,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 From 5451fb769a8b95eef35b97497c3c7e4f767473ee Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Fri, 30 Oct 2026 14:08:39 +0000 Subject: [PATCH 6/7] [None][fix] Consolidate PrimsTS decode compile identity and reducer arguments Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- 3rdparty/vendor_sources.lock.yaml | 4 +-- .../attention/backends/prims_ts/decode.py | 28 +++++-------------- .../prims_ts/kernels/fmha_decode/reduction.py | 8 ++---- .../sparse/test_prims_ts_block_sparse.py | 14 ++++++++++ 4 files changed, 26 insertions(+), 28 deletions(-) diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 96dea4d50cd4..26d111699e19 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -15,14 +15,14 @@ vendors: flashinfer-prims-ts: url: https://github.com/yuxianq/flashinfer.git branch: prims-tsfmha-fp8-non-mla--kernels - commit: 2bbeaf3c78b237298ae18e5cbe7b32d09caff562 + 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:afae80127f03ecc75c57f516d564050fe7475e0f0f7a4aec52beeb1b76f50ddc + 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/prims_ts/decode.py b/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py index f8802c532722..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 @@ -1480,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.""" @@ -1504,16 +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, - use_device_scales: bool = False, -): +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 @@ -1706,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], @@ -1754,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, @@ -1953,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, @@ -1973,7 +1968,6 @@ def compile_reducer(): compile_spec, _COMPILE_OPTIONS, compile_main, - extra=(use_device_scales,), ) compiled_reducer = None if cfg.use_separate_reduction_kernel: @@ -2260,9 +2254,6 @@ def _launch_decode( workspace.partial_o, workspace.partial_stats, workspace.attention_sinks, - # Producers already scale normalized partial O and log2-LSE. - 1.0, - 1.0, (), ) return runtime.out @@ -2600,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 @@ -4026,12 +4015,9 @@ 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, use_device_scales=True) - if use_device_scales - else _get_compiled_decode(compile_spec) - ) + compiled_main, compiled_reducer = _get_compiled_decode(compile_spec) policy = spec.policy + ( ("kv_prefix_mode", kv_prefix_mode), ("kv_lengths_mode", kv_lengths_mode), 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/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)) From 7382d0fb5920526327bde3b37ea67965463f7edb Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Fri, 30 Oct 2026 14:44:57 +0000 Subject: [PATCH 7/7] [None][fix] Preserve MLA architecture routing after rebase Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py | 2 +- tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py | 5 +++-- tests/unittest/_torch/attention/test_cute_dsl_mla.py | 1 + 3 files changed, 5 insertions(+), 3 deletions(-) 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 2dbf1b262637..02d9c6d2dc97 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py @@ -164,6 +164,7 @@ class PrimsTSFmha(PhasedFmha): supports_per_tensor_fp8 = True 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} @@ -202,8 +203,8 @@ def _is_available(cls, attn: "TrtllmAttention") -> bool: f"PrimTS FMHA is unavailable: requires SM100, SM103, or SM107, got SM{sm}." ) return False - if sm == 107 and attn.is_mla_enable: - logger.debug("PrimTS MLA FMHA is unavailable on SM107.") + 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") 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),