From 6549db793df05732aec44ee87f9e4a128bafafaf Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Wed, 28 Oct 2026 04:26:33 +0000 Subject: [PATCH 01/19] [None][feat] Vendor merge-compatible PrimTS softmax statistics Import FlashInfer 75961f74dac2c438cd7bce3dd135d527e78e49ec from prims-ts-mla-prefill--flashinfer. This consolidates the functional vendor refresh and its rebase/ABI fixups without changing the pinned source or generated bytes. The runtime exposes store_softmax_stats for continuous/paged context, standard decode, and MLA decode. These statistics enable cached/chunked MLA partial-output merging. Dense-Q scheduling optimization is deliberately excluded and reviewed separately. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../vendor_patches/flashinfer-prims-ts.patch | 57 ++--- 3rdparty/vendor_sources.lock.yaml | 8 +- .../attention/backends/prims_ts/context.py | 158 ++++++++++++- .../attention/backends/prims_ts/decode.py | 105 ++++++++- .../kernels/fmha_context/fmha_kernel.py | 15 ++ .../kernels/fmha_context/fmha_resources.py | 217 ++++++++++++++---- .../kernels/fmha_decode/fmha_decode_config.py | 9 +- .../kernels/fmha_decode/fmha_decode_kernel.py | 13 ++ .../fmha_decode_resources/smem_p.py | 28 +++ .../fmha_decode_resources/tmem_corr.py | 142 +++++++++--- .../prims_ts/kernels/fmha_decode/reduction.py | 134 ++++++++++- .../mla_decode/helpers/softmax_stats.py | 79 +++++++ .../mla_decode/throughput_2cta/config.py | 5 +- .../mla_decode/throughput_2cta/kernel.py | 23 ++ .../throughput_2cta/parallel_reduction.py | 31 +++ .../mla_decode/throughput_2cta/reduction.py | 11 + .../mla_decode/throughput_2cta/resources.py | 51 +++- .../throughput_latency_1cta/config.py | 13 +- .../throughput_latency_1cta/kernel.py | 36 +++ .../parallel_reduction.py | 27 +++ .../throughput_latency_1cta/reduction.py | 36 +++ .../resources/smem_p.py | 51 ++-- .../resources/tmem_corr.py | 140 ++++++++++- .../resources/tmem_p.py | 105 +++++++-- .../attention/backends/prims_ts/mla_decode.py | 96 +++++++- 25 files changed, 1389 insertions(+), 201 deletions(-) create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/softmax_stats.py diff --git a/3rdparty/vendor_patches/flashinfer-prims-ts.patch b/3rdparty/vendor_patches/flashinfer-prims-ts.patch index 6cfb3f7bf8e3..b47dd844bfeb 100644 --- a/3rdparty/vendor_patches/flashinfer-prims-ts.patch +++ b/3rdparty/vendor_patches/flashinfer-prims-ts.patch @@ -256,7 +256,7 @@ index fa055618705216cddf6b7f3e9ed4c1d680c2af9a..38de4926a045f17c5b6f206f4a20ecbe q: torch.Tensor, paged_kv_cache: PagedKVCache, diff --git a/context.py b/context.py -index d3f104e4e34b1bcb7ada697f9e76f7013d6ca100..1ec1080b942f60fad22e3e4a1627cad9592697ad 100644 +index 1cab933986621c42e01ae1f42fe8c23c99e94b04..b41e7b47cca1ed11eba11bddf1b0426c51071681 100644 --- a/context.py +++ b/context.py @@ -48,7 +48,7 @@ from ._hardware import get_max_active_clusters @@ -268,7 +268,7 @@ index d3f104e4e34b1bcb7ada697f9e76f7013d6ca100..1ec1080b942f60fad22e3e4a1627cad9 if TYPE_CHECKING: -@@ -493,7 +493,7 @@ def _validate_device(device: torch.device) -> int: +@@ -500,7 +500,7 @@ def _validate_device(device: torch.device) -> int: # Rubin runs through the sm_100f family target; a CuTe DSL older than 4.8 # cannot emit for it unless CUTE_DSL_ARCH=sm_100f is set before import. if capability == (10, 7): @@ -278,7 +278,7 @@ index d3f104e4e34b1bcb7ada697f9e76f7013d6ca100..1ec1080b942f60fad22e3e4a1627cad9 require_cute_dsl_arch(device_index) return device_index diff --git a/decode.py b/decode.py -index 4fb096f39849a3801349758360e71b092ee200e9..f91814650c48dcf5b116273742880ac722fed527 100644 +index 8c331d05a5b55717df3ba4dd1ff55b6ca44c831e..9c0a5e7243ee2d4ad33eed1afd45f9b8aa87f9e7 100644 --- a/decode.py +++ b/decode.py @@ -1,3 +1,4 @@ @@ -308,7 +308,7 @@ index 4fb096f39849a3801349758360e71b092ee200e9..f91814650c48dcf5b116273742880ac7 require_cute_dsl_arch(device_index) -@@ -4075,7 +4072,7 @@ class BatchDecodePagedTSWrapper: +@@ -4143,7 +4140,7 @@ class BatchDecodePagedTSWrapper: # previous complete plan revision usable. self._plan_state = candidate @@ -317,7 +317,7 @@ index 4fb096f39849a3801349758360e71b092ee200e9..f91814650c48dcf5b116273742880ac7 def run( self, q: torch.Tensor, -@@ -4257,7 +4254,7 @@ class BatchDecodePagedTSWrapper: +@@ -4340,7 +4337,7 @@ class BatchDecodePagedTSWrapper: ) @@ -327,7 +327,7 @@ index 4fb096f39849a3801349758360e71b092ee200e9..f91814650c48dcf5b116273742880ac7 q: torch.Tensor, paged_kv_cache: PagedKVCache, diff --git a/kernels/fmha_context/fmha_resources.py b/kernels/fmha_context/fmha_resources.py -index d17584b725e24235ce0ec90caac174a29d4ab965..af69ee7d51bc1179352d0bf2ba71fad2fe30a922 100644 +index f4e67b642efb6d8f8ad5e5519c1679c47e776f54..9ab920d4ba7017db8d26f884ae2e420b67ddb920 100644 --- a/kernels/fmha_context/fmha_resources.py +++ b/kernels/fmha_context/fmha_resources.py @@ -1,3 +1,4 @@ @@ -385,7 +385,7 @@ index d17584b725e24235ce0ec90caac174a29d4ab965..af69ee7d51bc1179352d0bf2ba71fad2 tile_idx = _resolve_kv_tile_idx_context( stage_info, kv_tile_start, tile_offset=tile_offset ) -@@ -5770,3 +5759,14 @@ class GmemOResource(MemoryResource): +@@ -5903,3 +5892,14 @@ class GmemOResource(MemoryResource): prims.cp_async_bulk_commit_group() if cutlass.const_expr(self.cfg.gmem_o_store_wait_after_write): prims.cp_async_bulk_wait_group(0, read=True) @@ -419,7 +419,7 @@ index b22f0f413b7312bab9a5f8e296255fe9aeabe901..a4a87200ca5d3ace5cf899caa9672a1a from .helpers_common import ( _TASK_CACHE_TMEM_BASE_OFFSET, diff --git a/mla_decode.py b/mla_decode.py -index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e231c98a52 100644 +index 6ca10c5f7195531890abd3f0bf643f6ceab36e5e..594d160b83aaa1b7bde74aa691e5e4ab92e630db 100644 --- a/mla_decode.py +++ b/mla_decode.py @@ -1,3 +1,4 @@ @@ -440,7 +440,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 from .decode import ( _WorkspaceSection, -@@ -1047,11 +1044,35 @@ def _make_mla_decode_compile_spec( +@@ -1080,11 +1077,35 @@ def _make_mla_decode_compile_spec( ) @@ -481,26 +481,8 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 import cutlass import cutlass.cute as cute -@@ -1459,11 +1480,12 @@ def _launch_mla_decode( - runtime.sparse_inputs, - runtime.bmm1_scale, - runtime.bmm2_scale, -- **( -- {"output_quant_args": runtime.output_quant_args} -- if runtime.output_quant_args is not None -- else {} -- ), -+ # Passed positionally and also when None. The kernel was compiled with -+ # ``output_quant_args`` in its signature (``None`` for unquantized -+ # output), and a kernel reloaded from the persistent compile cache is -+ # a positional TVM-FFI entry with no Python defaults: it rejects both -+ # a missing parameter and a keyword argument. -+ runtime.output_quant_args, - ) - return runtime.out - -@@ -1637,7 +1659,7 @@ class BatchMLADecodePagedTSWrapper: - split_kv=int(dict(policy)["split_kv"]), +@@ -1697,7 +1718,7 @@ class BatchMLADecodePagedTSWrapper: + store_softmax_stats=store_softmax_stats, ) - @flashinfer_experimental_api(trace=prims_ts_decode_mla_wrapper_trace_dispatch) @@ -508,7 +490,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 def run( self, query: torch.Tensor, -@@ -1742,7 +1764,7 @@ class BatchMLADecodePagedTSWrapper: +@@ -1814,7 +1835,7 @@ class BatchMLADecodePagedTSWrapper: ) @@ -517,7 +499,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 def batch_mla_decode_with_paged_kv_cache( query: torch.Tensor, kv_cache: torch.Tensor, -@@ -2099,6 +2121,7 @@ def _prepare_sparse_output_quant( +@@ -2181,6 +2202,7 @@ def _prepare_sparse_output_quant( ) @@ -525,7 +507,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 def _resolve_sparse_mla_plan( device, batch_size, -@@ -2119,17 +2142,17 @@ def _resolve_sparse_mla_plan( +@@ -2201,17 +2223,17 @@ def _resolve_sparse_mla_plan( ): """Resolve sparse geometry and scratch sizing without allocation/compilation.""" import cutlass @@ -547,7 +529,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 MlaDecodeTs, ) -@@ -2264,6 +2287,9 @@ def _resolve_sparse_mla_plan( +@@ -2346,6 +2368,9 @@ def _resolve_sparse_mla_plan( return device, profile, spec, sections, byte_end, capacity, output_quant @@ -557,7 +539,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 class BatchSparseMLADecodePagedTSWrapper: """Plan native BF16/E4M3 attention, then bind one or two source index lists. -@@ -2297,6 +2323,11 @@ class BatchSparseMLADecodePagedTSWrapper: +@@ -2379,6 +2404,11 @@ class BatchSparseMLADecodePagedTSWrapper: self._workspace_buffer = workspace_buffer self._state = None self._scalar_cache = {} @@ -569,7 +551,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 self.workspace_size_bytes = 0 def plan( -@@ -2386,6 +2417,7 @@ class BatchSparseMLADecodePagedTSWrapper: +@@ -2468,6 +2498,7 @@ class BatchSparseMLADecodePagedTSWrapper: kernel_signature=_mla_kernel_compile_signature(static_kernel), ) ) @@ -577,7 +559,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 self._state = dict( device=device, batch=batch_size, -@@ -2608,6 +2640,38 @@ class BatchSparseMLADecodePagedTSWrapper: +@@ -2690,6 +2721,38 @@ class BatchSparseMLADecodePagedTSWrapper: state = self._state if state is None: raise RuntimeError("plan() must be called before run()") @@ -616,7 +598,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 _validate_query( query, packed_query=state["packed"], -@@ -2781,10 +2845,124 @@ class BatchSparseMLADecodePagedTSWrapper: +@@ -2863,10 +2926,125 @@ class BatchSparseMLADecodePagedTSWrapper: if static_scales else state["compiled"], ) @@ -735,6 +717,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 + bmm1, + bmm2, + None, ++ None, + ) + return out + diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 02bae0bde8cf..4037e23240be 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-ts-mla-prefill--flashinfer + commit: 75961f74dac2c438cd7bce3dd135d527e78e49ec 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:5cbcaf2b60029742c55f622eadc84b32f49e574038c417972955706157715141 - digest: sha256-tree-v1:fa20dc33e52628519dba7f071f5b28a4eb7963768adb6b46eb8000e4fdb805c4 + patch_digest: sha256:d4557cc5072cab63378dd76c251f0718582578ae38f95f0303656854e88a8d69 + digest: sha256-tree-v1:15e6ad32e7aeec1a152d49cdc43c371adae4a072b8de0d362562ecad3c937666 openengine-proto: url: https://github.com/ai-dynamo/openengine.git tag: v0.1.0 diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/context.py b/tensorrt_llm/_torch/attention/backends/prims_ts/context.py index 1ec1080b942f..b41e7b47cca1 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/context.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/context.py @@ -142,6 +142,7 @@ class _ContextPlanGeometry: causal_single_kv_tile: bool packed_dense_k_mask: bool head_dim_vo: int | None = None + store_softmax_stats: bool = False @dataclass(frozen=True) @@ -181,6 +182,7 @@ class _PagedContextPlanGeometry: has_q_offset: bool paged_v_tail_is_zero: bool packed_dense_k_mask: bool + store_softmax_stats: bool = False @dataclass(frozen=True) @@ -234,6 +236,7 @@ class _ContextCompileSpec: scheduler: _ContextScheduler head_dim_vo: int | None = None two_cta_umma: bool = False + store_softmax_stats: bool = False @dataclass(frozen=True) @@ -258,6 +261,7 @@ class _PagedContextCompileSpec: paged_v_tail_is_zero: bool packed_dense_k_mask: bool scheduler: _ContextScheduler + store_softmax_stats: bool = False def _make_context_kernel( @@ -279,6 +283,7 @@ def _make_context_kernel( uses_ldtm_stat: bool, fp8_psmem_early_token: bool = False, two_cta_umma: bool = False, + store_softmax_stats: bool = False, page_size: int | None = None, max_kv_len: int | None = None, ): @@ -327,7 +332,9 @@ def _make_context_kernel( head_paired=head_paired, window_size_left=window_left if window_left > 0 else 0, h_r=num_qo_heads // num_kv_heads, - enable_skip_correction=True, + # Export actual maxima, not the deferred normalization anchors that + # are sufficient when only the normalized attention output is needed. + enable_skip_correction=not store_softmax_stats, uses_ldtm_stat=uses_ldtm_stat, causal_single_kv_tile=(causal_single_kv_tile and not use_paged_kv), use_blackwell_tma_oob_workaround=( @@ -780,8 +787,9 @@ def _read_indptr( name: str, *, expected_total: int, + allow_empty: bool = False, ) -> tuple[tuple[int, ...], tuple[int, ...]]: - """Read cumulative offsets and validate strictly positive row lengths.""" + """Read cumulative offsets and validate nonempty (or explicitly empty) rows.""" values = tuple(int(value) for value in indptr.tolist()) if values[0] != 0: @@ -792,8 +800,9 @@ def _read_indptr( f"expected {expected_total}, got {values[-1]}" ) lengths = tuple(curr - prev for prev, curr in itertools.pairwise(values)) - if any(length <= 0 for length in lengths): - raise ValueError(f"{name} offsets must be strictly increasing") + if any(length < (0 if allow_empty else 1) for length in lengths): + ordering = "nondecreasing" if allow_empty else "strictly increasing" + raise ValueError(f"{name} offsets must be {ordering}") return values, lengths @@ -1269,10 +1278,13 @@ def _resolve_context_plan_geometry( mask_type: str, window_left: int, output_dtype: torch.dtype, + store_softmax_stats: bool = False, ) -> _ContextPlanGeometry: """Validate explicit static bounds for a reusable contiguous plan.""" _validate_context_dtype_pair(qk_dtype, qk_dtype, pv_dtype, output_dtype) + if not isinstance(store_softmax_stats, bool): + raise TypeError("store_softmax_stats must be a bool") _validate_mask(mask_type) window_left = _validate_window_left(window_left, mask_type) if not isinstance(packed, bool): @@ -1347,6 +1359,7 @@ def _resolve_context_plan_geometry( and max_kv_len <= _CONTEXT_KV_TILE_N ), packed_dense_k_mask=packed and mask_type == "dense", + store_softmax_stats=store_softmax_stats, ) @@ -1486,6 +1499,7 @@ def _make_context_scheduler_probe( causal_single_kv_tile=causal_single_kv_tile, scheduler="static_persistent", uses_ldtm_stat=_default_uses_ldtm_stat(geometry.device_index), + store_softmax_stats=geometry.store_softmax_stats, page_size=page_size, max_kv_len=max_kv_len, ) @@ -1521,10 +1535,13 @@ def _resolve_paged_plan_geometry( uniform_packed_lengths: bool = False, has_q_offset: bool = True, paged_v_tail_is_zero: bool = False, + store_softmax_stats: bool = False, ) -> _PagedContextPlanGeometry: """Validate explicit static bounds for a reusable paged specialization.""" _validate_paged_dtype_pair(qk_dtype, qk_dtype, pv_dtype, output_dtype) + if not isinstance(store_softmax_stats, bool): + raise TypeError("store_softmax_stats must be a bool") _validate_mask(mask_type) if not isinstance(uniform_packed_lengths, bool): raise TypeError("uniform_packed_lengths must be a bool") @@ -1586,6 +1603,7 @@ def _resolve_paged_plan_geometry( mask_type == "dense" and (not uniform_packed_lengths or max_kv_len % _CONTEXT_KV_TILE_N != 0) ), + store_softmax_stats=store_softmax_stats, ) @@ -1689,6 +1707,7 @@ def _context_compile_spec(geometry: _ContextPlanGeometry) -> _ContextCompileSpec head_dim_vo=geometry.head_dim_vo, qk_dtype_key=_dtype_key(geometry.qk_dtype), pv_dtype_key=_dtype_key(geometry.pv_dtype), + store_softmax_stats=geometry.store_softmax_stats, output_dtype_key=_dtype_key(geometry.output_dtype), mask_type=geometry.mask_type, window_left=geometry.window_left, @@ -1726,6 +1745,7 @@ def _paged_context_compile_spec( paged_v_tail_is_zero=geometry.paged_v_tail_is_zero, packed_dense_k_mask=geometry.packed_dense_k_mask, scheduler=_resolve_paged_context_scheduler(geometry), + store_softmax_stats=geometry.store_softmax_stats, ) @@ -1788,6 +1808,7 @@ def _get_compiled_context( scheduler=scheduler, two_cta_umma=two_cta_umma, uses_ldtm_stat=_default_uses_ldtm_stat(device_index), + store_softmax_stats=compile_spec.store_softmax_stats, fp8_psmem_early_token=_default_fp8_psmem_early_token(device_index), ) fmha.cfg.has_varlen = packed @@ -1820,6 +1841,7 @@ def tensor_adapter( variable_window_token_starts: cute.Tensor, variable_window_token_ends: cute.Tensor, variable_window_cta_starts: cute.Tensor, + softmax_stats: cute.Tensor | None, stream: cuda_drv.CUstream, static_max_active_clusters: cutlass.Constexpr[int], static_packed: cutlass.Constexpr[bool], @@ -1845,6 +1867,7 @@ def tensor_adapter( variable_window_token_starts=variable_window_token_starts, variable_window_token_ends=variable_window_token_ends, variable_window_cta_starts=variable_window_cta_starts, + softmax_stats=softmax_stats, ) else: fmha( @@ -1859,6 +1882,7 @@ def tensor_adapter( variable_window_token_starts=variable_window_token_starts, variable_window_token_ends=variable_window_token_ends, variable_window_cta_starts=variable_window_cta_starts, + softmax_stats=softmax_stats, ) def fake_compact(dtype, shape, assumed_align): @@ -1893,6 +1917,11 @@ def fake_compact(dtype, shape, assumed_align): k_fake = fake_compact(input_qk_dtype, kv_shape, 16) v_fake = fake_compact(input_pv_dtype, (*kv_shape[:-1], head_dim_vo), 16) out_fake = fake_compact(output_dtype, out_shape, 16) + softmax_stats_fake = ( + fake_compact(cutlass.Float32, (*out_shape[:-1], 2), 4) + if compile_spec.store_softmax_stats + else None + ) scale_fake = fake_compact(cutlass.Float32, (1,), 4) output_scale_fake = fake_compact(cutlass.Float32, (1,), 4) qo_indptr_fake = fake_compact(cutlass.Int32, qo_indptr_shape, 4) @@ -1931,6 +1960,7 @@ def compile_kernel(): variable_window_starts_fake, variable_window_ends_fake, variable_window_cta_starts_fake, + softmax_stats_fake, stream_fake, max_active_clusters, packed, @@ -2012,6 +2042,7 @@ def _get_compiled_paged_context( causal_single_kv_tile=False, scheduler=scheduler, uses_ldtm_stat=_default_uses_ldtm_stat(device_index), + store_softmax_stats=compile_spec.store_softmax_stats, fp8_psmem_early_token=_default_fp8_psmem_early_token(device_index), page_size=page_size, max_kv_len=max_kv_len, @@ -2043,6 +2074,7 @@ def tensor_adapter( qo_indptr: cute.Tensor, block_tables: cute.Tensor, seq_lens_kv: cute.Tensor, + softmax_stats: cute.Tensor | None, stream: cuda_drv.CUstream, static_max_active_clusters: cutlass.Constexpr[int], static_max_seq_len_q: cutlass.Constexpr[int], @@ -2062,6 +2094,7 @@ def tensor_adapter( max_seqlen_k=cutlass.Int32(static_max_kv_len), block_tables=block_tables, seq_lens_kv=seq_lens_kv, + softmax_stats=softmax_stats, ) def fake_compact(dtype, shape, assumed_align): @@ -2086,6 +2119,11 @@ def fake_compact(dtype, shape, assumed_align): k_fake = fake_compact(input_qk_dtype, kv_shape, 16) v_fake = fake_compact(input_pv_dtype, kv_shape, 16) out_fake = fake_compact(output_dtype, (runtime_total_q, num_qo_heads, head_dim), 16) + softmax_stats_fake = ( + fake_compact(cutlass.Float32, (runtime_total_q, num_qo_heads, 2), 4) + if compile_spec.store_softmax_stats + else None + ) scale_fake = fake_compact(cutlass.Float32, (1,), 4) output_scale_fake = fake_compact(cutlass.Float32, (1,), 4) qo_indptr_fake = fake_compact(cutlass.Int32, (runtime_num_q_offsets,), 4) @@ -2112,6 +2150,7 @@ def compile_kernel(): qo_indptr_fake, block_tables_fake, seq_lens_fake, + softmax_stats_fake, stream_fake, max_active_clusters, max_seq_len_q, @@ -2198,7 +2237,12 @@ def _validate_runtime_inputs( f"{total_q} for batch_size={geometry.batch_size} and " f"max_seq_len_q={geometry.max_seq_len_q}" ) - if total_k < geometry.batch_size or total_k > ( + min_total_k = ( + 1 + if geometry.store_softmax_stats and geometry.mask_type == "dense" + else geometry.batch_size + ) + if total_k < min_total_k or total_k > ( geometry.batch_size * geometry.max_seq_len_k ): raise ValueError( @@ -2224,6 +2268,11 @@ def _validate_runtime_inputs( indptr, name, expected_total=expected_total, + allow_empty=( + name == "kv_indptr" + and geometry.store_softmax_stats + and geometry.mask_type == "dense" + ), ) request_lengths[name] = lengths capacity = ( @@ -2469,6 +2518,32 @@ def _prepare_out( return out +def _validate_softmax_stats( + softmax_stats: Optional[torch.Tensor], + q: torch.Tensor, + *, + store_softmax_stats: bool, + validate: bool, +) -> None: + """Check the compile-time output contract without reading device values.""" + if store_softmax_stats != (softmax_stats is not None): + raise ValueError( + "softmax_stats must be supplied exactly when " + "store_softmax_stats=True was planned" + ) + if validate and softmax_stats is not None: + if not isinstance(softmax_stats, torch.Tensor): + raise TypeError("softmax_stats must be a torch.Tensor") + if softmax_stats.dtype != torch.float32: + raise ValueError("softmax_stats must have dtype float32") + if softmax_stats.device != q.device: + raise ValueError("softmax_stats must be on the same device as q") + if tuple(softmax_stats.shape) != (*q.shape[:-1], 2): + raise ValueError("softmax_stats must have shape [*q.shape[:-1], 2]") + _validate_compact(softmax_stats, "softmax_stats", "[*q.shape[:-1], 2]") + _validate_alignment(softmax_stats, "softmax_stats", 4) + + class BatchPrefillTSWrapper: """Compile and reuse fixed or packed-ragged contiguous context attention. @@ -2519,6 +2594,7 @@ def plan( v_dtype: Optional[torch.dtype] = None, out_dtype: Optional[torch.dtype] = None, packed: bool = False, + store_softmax_stats: bool = False, mask_type: Literal["dense", "causal", "variable_window"] = "dense", window_left: int = -1, sm_scale: Optional[float] = None, @@ -2526,8 +2602,9 @@ def plan( ) -> None: """Compile one reusable specialization from explicit static geometry. - Packed plans conservatively support any positive per-request Q and K/V - lengths bounded by ``max_seq_len_q`` and ``max_kv_len``. Fixed plans + Packed plans support positive per-request Q and K/V lengths bounded by + ``max_seq_len_q`` and ``max_kv_len`` (empty dense K/V chunks are also + supported when statistics are enabled; see below). Fixed plans use those bounds as their exact tensor extents. ``batch_size`` is exact. Calling ``plan`` again replaces all plan-owned tensors and invalidates CUDA graphs captured from the previous plan. Complete every launch and @@ -2565,6 +2642,13 @@ def plan( Output dtype; defaults to ``q_dtype``. packed : bool Whether run tensors use packed ``[total_tokens, H, D]`` storage. + store_softmax_stats : bool + Compile an additional output for merging independently computed + attention chunks. When enabled, every run must supply a float32 + ``softmax_stats`` tensor with shape ``[*q.shape[:-1], 2]``. + Packed dense plans also permit empty K/V rows for absent chunks, + provided the packed K/V tensors contain at least one token. + Empty rows produce zero output and statistics ``(-inf, 0)``. mask_type : {"dense", "causal", "variable_window"} Attention mask mode. ``variable_window`` is supported only for fixed-shape inputs. @@ -2598,6 +2682,7 @@ def plan( mask_type=mask_type, window_left=window_left, output_dtype=resolved_out_dtype, + store_softmax_stats=store_softmax_stats, ) if sm_scale is None: sm_scale = 1.0 / math.sqrt(geometry.head_dim) @@ -2663,6 +2748,7 @@ def run( variable_window_token_ends: Optional[torch.Tensor] = None, variable_window_cta_starts: Optional[torch.Tensor] = None, out: Optional[torch.Tensor] = None, + softmax_stats: Optional[torch.Tensor] = None, scale_softmax_log2: Optional[torch.Tensor] = None, output_scale: Optional[torch.Tensor] = None, validate: bool = True, @@ -2687,6 +2773,16 @@ def run( wrapper derives these values from ``variable_window_token_starts``. out : torch.Tensor, optional Caller-owned output tensor. A new tensor is allocated when omitted. + softmax_stats : torch.Tensor, optional + Caller-owned compact float32 tensor shaped ``[*q.shape[:-1], 2]``. + Required exactly when ``store_softmax_stats=True`` was planned. + Each pair is ``(m, s)``, where ``m`` is the maximum scaled logit + in natural-log units and ``s = sum(exp(logit - m))``. Together + with the normalized output, these values allow stable merging of + attention over disjoint K/V chunks. They are independent of + ``output_scale``; merged outputs must use a common output scale. + Must not overlap ``out``, any live input, or a plan-owned buffer; + storage overlap is a caller precondition and is not checked. scale_softmax_log2 : torch.Tensor, optional Per-run one-element float32 softmax scale in base-2 form. Defaults to the scale retained by the plan. @@ -2705,6 +2801,12 @@ def run( if not isinstance(validate, bool): raise TypeError("validate must be a bool") geometry = state.geometry + _validate_softmax_stats( + softmax_stats, + q, + store_softmax_stats=geometry.store_softmax_stats, + validate=validate, + ) if geometry.packed: if qo_indptr is None or kv_indptr is None: raise ValueError( @@ -2817,6 +2919,7 @@ def run( runtime_window_starts, runtime_window_ends, runtime_window_cta_starts, + softmax_stats, ) return out @@ -2886,6 +2989,7 @@ def plan( uniform_packed_lengths: bool = False, has_q_offset: bool = True, paged_v_tail_is_zero: bool = False, + store_softmax_stats: bool = False, ) -> None: """Compile one reusable specialization from explicit static geometry. @@ -2960,6 +3064,10 @@ def plan( Whether unused rows in every active final V-cache page are guaranteed to contain zero. Setting this to ``True`` compiles out the consumer-side V-tail clear. Defaults to ``False``. + store_softmax_stats : bool + Compile a final maximum/denominator output. Each run must supply + a float32 ``softmax_stats`` tensor shaped ``[total_q, Hq, 2]``. + Defaults to ``False``. """ resolved_out_dtype = q_dtype if out_dtype is None else out_dtype @@ -2986,6 +3094,7 @@ def plan( uniform_packed_lengths=uniform_packed_lengths, has_q_offset=has_q_offset, paged_v_tail_is_zero=paged_v_tail_is_zero, + store_softmax_stats=store_softmax_stats, ) if sm_scale is None: sm_scale = 1.0 / math.sqrt(geometry.head_dim) @@ -3026,6 +3135,7 @@ def run( seq_lens_kv: torch.Tensor, *, out: Optional[torch.Tensor] = None, + softmax_stats: Optional[torch.Tensor] = None, scale_softmax_log2: Optional[torch.Tensor] = None, output_scale: Optional[torch.Tensor] = None, validate: bool = True, @@ -3077,6 +3187,13 @@ def run( Logical K/V lengths with shape ``[B]``. out : torch.Tensor, optional Caller-owned output tensor. A new tensor is allocated when omitted. + softmax_stats : torch.Tensor, optional + Caller-owned contiguous float32 tensor shaped ``[total_q, Hq, 2]``. + Required exactly when ``store_softmax_stats=True`` was planned. + Stores the scaled logit maximum ``m`` in natural-log units and + ``sum(exp(logit - m))``, without internal probability scaling. + Must not overlap output, inputs, or plan buffers; overlap is not + checked. CUDA graphs retain this buffer's address. scale_softmax_log2 : torch.Tensor, optional Per-run one-element float32 softmax scale in base-2 form. Defaults to the scale retained by the plan. @@ -3102,6 +3219,12 @@ def run( if not isinstance(validate, bool): raise TypeError("validate must be a bool") geometry = state.geometry + _validate_softmax_stats( + softmax_stats, + q, + store_softmax_stats=geometry.store_softmax_stats, + validate=validate, + ) if validate: _validate_paged_runtime_inputs(q, k_cache, v_cache, geometry) _validate_paged_runtime_metadata( @@ -3143,6 +3266,7 @@ def run( qo_indptr, block_tables, seq_lens_kv, + softmax_stats, ) return out @@ -3163,6 +3287,8 @@ def batch_prefill( output_scale: float = 1.0, out_dtype: Optional[torch.dtype] = None, out: Optional[torch.Tensor] = None, + store_softmax_stats: bool = False, + softmax_stats: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Run one-shot fixed or packed-ragged task-scheduled context attention. @@ -3201,6 +3327,12 @@ def batch_prefill( otherwise the query dtype. out : torch.Tensor, optional Caller-owned output tensor. + store_softmax_stats : bool + Compile an additional final maximum/denominator output; defaults to False. + softmax_stats : torch.Tensor, optional + Caller-owned contiguous float32 tensor shaped ``[*q.shape[:-1], 2]``. + Required exactly when ``store_softmax_stats=True``. See + :meth:`BatchPrefillTSWrapper.run` for the numerical contract. Returns ------- @@ -3245,6 +3377,7 @@ def batch_prefill( v_dtype=geometry.pv_dtype, out_dtype=geometry.output_dtype, packed=geometry.packed, + store_softmax_stats=store_softmax_stats, mask_type=mask_type, window_left=window_left, sm_scale=sm_scale, @@ -3259,6 +3392,7 @@ def batch_prefill( variable_window_token_starts=variable_window_token_starts, variable_window_token_ends=variable_window_token_ends, out=out, + softmax_stats=softmax_stats, ) @@ -3279,6 +3413,8 @@ def batch_prefill_with_paged_kv_cache( output_scale: float = 1.0, out_dtype: Optional[torch.dtype] = None, out: Optional[torch.Tensor] = None, + store_softmax_stats: bool = False, + softmax_stats: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Run one-shot packed-Q context attention over separate HND page pools. @@ -3326,6 +3462,12 @@ def batch_prefill_with_paged_kv_cache( otherwise the query dtype. out : torch.Tensor, optional Caller-owned output tensor. + store_softmax_stats : bool + Compile an additional final maximum/denominator output; defaults to False. + softmax_stats : torch.Tensor, optional + Caller-owned contiguous float32 tensor shaped ``[total_q, Hq, 2]``. + Required exactly when ``store_softmax_stats=True``. See + :meth:`BatchPrefillPagedTSWrapper.run` for the numerical contract. Returns ------- @@ -3379,6 +3521,7 @@ def batch_prefill_with_paged_kv_cache( uniform_packed_lengths=geometry.uniform_packed_lengths, has_q_offset=geometry.has_q_offset, paged_v_tail_is_zero=geometry.paged_v_tail_is_zero, + store_softmax_stats=store_softmax_stats, ) return wrapper.run( q, @@ -3388,6 +3531,7 @@ def batch_prefill_with_paged_kv_cache( block_tables, seq_lens_kv, out=out, + softmax_stats=softmax_stats, ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py b/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py index f91814650c48..9c0a5e7243ee 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py @@ -1187,11 +1187,14 @@ def _decode_scratch_shapes( head_ratio * seq_len_q, head_dim, ) - partial_stats_shape = ( - partial_o_shape[:-1] - if cfg.use_separate_reduction_kernel - else partial_o_shape[:-1] + (2,) - ) + if cfg.use_separate_reduction_kernel: + partial_stats_shape = ( + partial_o_shape[:-1] + (3,) + if cfg.store_softmax_stats + else partial_o_shape[:-1] + ) + else: + partial_stats_shape = partial_o_shape[:-1] + (2,) counter_shape = (batch_size, num_kv_heads, num_q_groups) return partial_o_shape, partial_stats_shape, counter_shape @@ -1248,6 +1251,7 @@ def _resolve_decode_launch_spec( use_pdl: bool = False, split_kv: bool = True, share_pattern_across_kv_heads: bool = True, + store_softmax_stats: bool = False, ) -> _DecodeLaunchSpec: """Resolve automatic policy and workspace geometry without compiling.""" @@ -1448,6 +1452,7 @@ def fits_one_service_wave(config: "FmhaDecodeConfig", num_q_ctas: int) -> bool: max_splits_kv=safe_splits, ) + cfg.store_softmax_stats = store_softmax_stats _validate_decode_policy_kv_tile_size(cfg) return _decode_launch_spec_from_config( cfg, @@ -1591,6 +1596,7 @@ def main_tensor_adapter( bmm1_scale: cutlass.Float32, bmm2_scale: cutlass.Float32, direct_q1_inputs: tuple, + softmax_stats: cute.Tensor | None, stream: cuda_drv.CUstream, static_cfg: cutlass.Constexpr[FmhaDecodeConfig], static_seq_len_q: cutlass.Constexpr[int], @@ -1688,6 +1694,9 @@ def main_tensor_adapter( v_token_stride, static_full_split_prefix, static_native_uniform_kv, + softmax_stats_iter=( + softmax_stats.iterator if softmax_stats is not None else None + ), ) reduction_tensor_adapter = None @@ -1707,6 +1716,7 @@ def reduction_tensor_adapter( bmm1_scale: cutlass.Float32, bmm2_scale: cutlass.Float32, direct_q1_inputs: tuple, + softmax_stats: cute.Tensor | None, stream: cuda_drv.CUstream, static_cfg: cutlass.Constexpr[FmhaDecodeConfig], static_num_qo_heads: cutlass.Constexpr[int], @@ -1757,6 +1767,9 @@ def reduction_tensor_adapter( stream, static_cfg, static_full_split_prefix, + softmax_stats_iter=( + softmax_stats.iterator if softmax_stats is not None else None + ), ) physical_pages = cute.sym_int() @@ -1867,6 +1880,11 @@ def fake_compact(dtype, shape, assumed_align): else: k_sf_fake = fake_compact(cutlass.Uint8, (1,), 1) v_sf_fake = fake_compact(cutlass.Uint8, (1,), 1) + softmax_stats_fake = ( + fake_compact(Float32, out_shape[:-1] + (2,), 4) + if cfg.store_softmax_stats + else None + ) stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) direct_q1_fake: tuple = () if direct_q1_spec is not None: @@ -1927,6 +1945,7 @@ def compile_main(): Float32(1.0), Float32(1.0), direct_q1_fake, + softmax_stats_fake, stream_fake, cfg, seq_len_q, @@ -1953,6 +1972,7 @@ def compile_reducer(): Float32(1.0), Float32(1.0), direct_q1_fake, + softmax_stats_fake, stream_fake, cfg, num_qo_heads, @@ -1998,6 +2018,7 @@ def get_prims_ts_batch_decode_workspace_size( storage_page_size: Optional[int] = None, device: Optional[Union[int, str, torch.device]] = None, split_kv: bool = True, + store_softmax_stats: bool = False, ) -> int: """Return caller-workspace bytes for one automatic FMHA policy. @@ -2025,8 +2046,13 @@ def get_prims_ts_batch_decode_workspace_size( ``split_kv`` permits automatic split fanout when True (default), or forces nonsplit execution when False. It is independent of Q layout and must match the value supplied to attention planning/launch. + ``store_softmax_stats`` must also match the plan: enabled standalone + reducers reserve an additional FP32 maximum and denominator per split and + output row. """ + if not isinstance(store_softmax_stats, bool): + raise TypeError("store_softmax_stats must be a bool") batch_size = _validate_positive_int(batch_size, "batch_size") use_packed_q, resolved_seq_len_q = _resolve_q_mode( seq_len_q=seq_len_q, @@ -2092,6 +2118,7 @@ def get_prims_ts_batch_decode_workspace_size( storage_page_size, resolved_device, split_kv=split_kv, + store_softmax_stats=store_softmax_stats, ).total_bytes @@ -2200,6 +2227,32 @@ def _prepare_decode_runtime( ) +def _validate_decode_softmax_stats( + store_softmax_stats: bool, + softmax_stats: Optional[torch.Tensor], + q: torch.Tensor, + *, + validate: bool, +) -> None: + """Enforce the compile-time binding contract without device synchronization.""" + if store_softmax_stats != (softmax_stats is not None): + raise ValueError( + "softmax_stats must be supplied exactly when plan() sets " + "store_softmax_stats=True" + ) + if validate and softmax_stats is not None: + if not isinstance(softmax_stats, torch.Tensor): + raise TypeError("softmax_stats must be a torch.Tensor") + if softmax_stats.dtype != torch.float32: + raise TypeError("softmax_stats must have dtype torch.float32") + if softmax_stats.device != q.device: + raise ValueError("softmax_stats must be on the query device") + if tuple(softmax_stats.shape) != tuple(q.shape[:-1]) + (2,): + raise ValueError("softmax_stats must have shape [*out.shape[:-1], 2]") + if not softmax_stats.is_contiguous(): + raise ValueError("softmax_stats must be contiguous") + + def _launch_decode( runtime: _DecodeRuntime, *, @@ -2209,6 +2262,7 @@ def _launch_decode( workspace: _DecodeWorkspaceViews, compiled_main: Callable[..., object], compiled_reducer: Optional[Callable[..., object]], + softmax_stats: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Launch the compiled main kernel and its optional standalone reducer.""" @@ -2238,6 +2292,7 @@ def _launch_decode( runtime.bmm1_scale, runtime.bmm2_scale, (), + softmax_stats, ) if compiled_reducer is not None: compiled_reducer( @@ -2250,6 +2305,7 @@ def _launch_decode( runtime.bmm1_scale, runtime.bmm2_scale, (), + softmax_stats, ) return runtime.out @@ -2572,6 +2628,7 @@ def _run_unchecked( scale_qk, scale_v, direct_inputs, + None, ) if self._compiled_reducer is not None: self._compiled_reducer( @@ -2584,6 +2641,7 @@ def _run_unchecked( scale_qk, scale_v, direct_inputs, + None, ) return out @@ -2610,6 +2668,7 @@ def _resolve_decode_workspace_layout( use_pdl: bool = False, split_kv: bool = True, share_pattern_across_kv_heads: bool = True, + store_softmax_stats: bool = False, ) -> _DecodeWorkspaceLayout: """Resolve the byte layout for one already-validated semantic key.""" @@ -2637,6 +2696,7 @@ def _resolve_decode_workspace_layout( use_pdl, split_kv, share_pattern_across_kv_heads, + store_softmax_stats=store_softmax_stats, ) return _make_decode_workspace_layout( spec.scratch_shapes, @@ -3773,6 +3833,7 @@ def plan( split_kv: bool = True, validate: bool = True, initialize_workspace: bool = True, + store_softmax_stats: bool = False, ) -> None: """Compile one static-capacity plan and optionally own sequence lengths. @@ -3851,6 +3912,9 @@ def plan( must be 32-byte aligned and large enough for the selected plan. When omitted, planning allocates the buffer. The retained buffer is exclusive to one in-flight launch or graph replay. + store_softmax_stats : bool + Compile export of merge-compatible softmax statistics. Defaults + to ``False``; the corresponding buffer belongs to each ``run``. split_kv : bool Permit automatic useful split fanout (True, default), or force unsplit execution (False), independently of ``packed_query``. @@ -3878,6 +3942,8 @@ def plan( raise TypeError("validate must be a bool") if not isinstance(initialize_workspace, bool): raise TypeError("initialize_workspace must be a bool") + if not isinstance(store_softmax_stats, bool): + raise TypeError("store_softmax_stats must be a bool") if validate: if not isinstance(packed_query, bool): raise TypeError("packed_query must be a bool") @@ -3946,7 +4012,9 @@ def plan( window_left, storage_page_size, ) - spec = _resolve_decode_launch_spec(*policy_args, split_kv=split_kv) + spec = _resolve_decode_launch_spec( + *policy_args, split_kv=split_kv, store_softmax_stats=store_softmax_stats + ) if specialization_seq_lens is None: kv_prefix_mode: Literal["dynamic", "planned_full"] = "dynamic" kv_lengths_mode: Literal["dynamic", "planned_uniform_max"] = "dynamic" @@ -4086,6 +4154,7 @@ def run( bmm2_scale: float = 1.0, out: Optional[torch.Tensor] = None, validate: bool = True, + softmax_stats: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Launch the current plan with one sequence-length owner. @@ -4106,6 +4175,9 @@ def run( In either validation mode, output and workspace must be disjoint from each other and from all inputs. Storage overlap is not checked. + ``softmax_stats`` must not overlap inputs, output, metadata, or workspace; + retain its address during graph replay. Its presence must agree with + ``store_softmax_stats`` in either validation mode. Parameters ---------- @@ -4136,6 +4208,13 @@ def run( Value/output scaling factor. Defaults to ``1.0``. out : torch.Tensor, optional Caller-owned output tensor. A new tensor is allocated when omitted. + softmax_stats : torch.Tensor, optional + Contiguous FP32 output ``[*out.shape[:-1], 2]`` on the query device. + Stores the scaled token-logit maximum in natural-log units and + ``sum(exp(logit - maximum))``, with internal FP8 probability + scaling removed. Where supported, attention sinks contribute once + to the denominator, not the maximum. Required exactly when the + plan enables ``store_softmax_stats``. Empty rows store ``(-inf, 0)``. validate : bool Run explicit structural and value validation. Disable only when the caller guarantees the complete runtime contract. Sequence @@ -4148,6 +4227,9 @@ def run( """ state = self._require_plan_state() + _validate_decode_softmax_stats( + state.config.store_softmax_stats, softmax_stats, q, validate=validate + ) if not isinstance(validate, bool): raise TypeError("validate must be a bool") planned_seq_lens = state.planned_seq_lens_device @@ -4251,6 +4333,7 @@ def run( workspace=state.workspace, compiled_main=state.compiled_main, compiled_reducer=state.compiled_reducer, + softmax_stats=softmax_stats, ) @@ -4277,6 +4360,8 @@ def batch_decode_with_paged_kv_cache( workspace_buffer: Optional[torch.Tensor] = None, max_kv_len: Optional[int] = None, validate: bool = True, + store_softmax_stats: bool = False, + softmax_stats: Optional[torch.Tensor] = None, ) -> torch.Tensor: """One-shot fixed or packed-Q paged decode from fixed page tables. @@ -4345,6 +4430,12 @@ def batch_decode_with_paged_kv_cache( Validate tensors and metadata values (may synchronize), default True. False trusts the caller and requires workspace and explicit bounds; skips tensor and metadata validation. Invalid inputs have undefined behavior. + store_softmax_stats : bool + Compile statistics export, disabled by default. Match workspace sizing. + softmax_stats : torch.Tensor, optional + Caller-owned statistics buffer with the same contract as + :meth:`BatchDecodePagedTSWrapper.run`. Required exactly when + ``store_softmax_stats=True`` and disjoint from all other buffers. Returns ------- @@ -4523,6 +4614,7 @@ def batch_decode_with_paged_kv_cache( workspace_buffer=workspace_buffer, validate=validate, initialize_workspace=allocate_workspace, + store_softmax_stats=store_softmax_stats, ) return wrapper.run( q, @@ -4535,6 +4627,7 @@ def batch_decode_with_paged_kv_cache( bmm2_scale=bmm2_scale, out=out, validate=validate, + softmax_stats=softmax_stats, ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py index 0c5a915d0bcf..f0e4ce0cd3d2 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py @@ -691,6 +691,7 @@ def build_context_task_manager( variable_window_q_stride: int | Int32 = 0, scale_softmax_log2: cute.Tensor | None = None, output_scale: cute.Tensor | None = None, + softmax_stats: cute.Tensor | None = None, g_block_tables: cute.Pointer | None = None, block_table_row_stride: int | Int32 = 0, g_seq_lens_kv: cute.Pointer | None = None, @@ -1198,6 +1199,7 @@ def make_page_offsets_pipeline_cfg(num_stages: int) -> PipelineConfig: variable_window_q_stride=variable_window_q_stride, scale_softmax_log2=scale_softmax_log2, smem_p=smem_p0, + store_softmax_stats=softmax_stats is not None, name="tmem_sp0", ) tmem_p0: TmemPResource | None = None @@ -1214,6 +1216,9 @@ def make_page_offsets_pipeline_cfg(num_stages: int) -> PipelineConfig: tmem_vec_offset=cfg.tmem_vec0_offset, scale_softmax_log2=scale_softmax_log2, output_scale=output_scale, + softmax_stats=softmax_stats, + cum_seqlen_q=cum_seqlen_q, + q_half=0, name="tmem_vec0", ) smem_o_0 = SmemOResource( @@ -1291,6 +1296,7 @@ def make_page_offsets_pipeline_cfg(num_stages: int) -> PipelineConfig: variable_window_q_stride=variable_window_q_stride, scale_softmax_log2=scale_softmax_log2, smem_p=smem_p1, + store_softmax_stats=softmax_stats is not None, name="tmem_sp1", ) tmem_vec1 = TmemStatsResource( @@ -1299,6 +1305,9 @@ def make_page_offsets_pipeline_cfg(num_stages: int) -> PipelineConfig: tmem_vec_offset=cfg.tmem_vec1_offset, scale_softmax_log2=scale_softmax_log2, output_scale=output_scale, + softmax_stats=softmax_stats, + cum_seqlen_q=cum_seqlen_q, + q_half=1, name="tmem_vec1", ) smem_o_1 = SmemOResource( @@ -2581,6 +2590,7 @@ def build_fmha_task_manager( variable_window_q_stride: int | Int32 = 0, scale_softmax_log2: cute.Tensor | None = None, output_scale: cute.Tensor | None = None, + softmax_stats: cute.Tensor | None = None, q_offset: int | Int32 = 0, g_block_tables: cute.Pointer | None = None, block_table_row_stride: int | Int32 = 0, @@ -2674,6 +2684,7 @@ def build_fmha_task_manager( variable_window_q_stride=variable_window_q_stride, scale_softmax_log2=scale_softmax_log2, output_scale=output_scale, + softmax_stats=softmax_stats, g_block_tables=g_block_tables, block_table_row_stride=block_table_row_stride, g_seq_lens_kv=g_seq_lens_kv, @@ -3065,6 +3076,7 @@ def __call__( variable_window_token_starts: cute.Tensor | None = None, variable_window_token_ends: cute.Tensor | None = None, variable_window_cta_starts: cute.Tensor | None = None, + softmax_stats: cute.Tensor | None = None, ) -> None: """Set up TMA descriptors, compute grid, and launch the kernel. @@ -3383,6 +3395,7 @@ def __call__( tile_sched_params, scale_softmax_log2, output_scale, + softmax_stats, num_kv_tiles, num_seq_tiles, q_offset, @@ -3423,6 +3436,7 @@ def kernel( ), scale_softmax_log2: cute.Tensor, output_scale: cute.Tensor, + softmax_stats: cute.Tensor | None, num_kv_tiles: Int32, num_seq_tiles: Int32, q_offset: Int32, @@ -3500,6 +3514,7 @@ def kernel( variable_window_q_stride=variable_window_q_stride, scale_softmax_log2=scale_softmax_log2, output_scale=output_scale, + softmax_stats=softmax_stats, q_offset=q_offset, is_persistent=is_persistent, is_clc_dynamic=is_clc_dynamic, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py index af69ee7d51bc..9ab920d4ba70 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py @@ -2507,6 +2507,7 @@ class TmemSPResource(MemoryResource): # 0 for SP0 (uses Q0), 1 for SP1 (uses Q1). q_half: Constexpr[int] = 0 enable_early_tile_sum: Constexpr[bool] = False + store_softmax_stats: Constexpr[bool] = False q_offset_default: int | Int32 = field(init=False, default=0) cum_seqlen_q: cute.Tensor | None = field(init=False, default=None) cum_seqlen_k: cute.Tensor | None = field(init=False, default=None) @@ -2559,6 +2560,7 @@ def __init__( variable_window_q_stride: int | Int32 = 0, scale_softmax_log2: cute.Tensor | None = None, smem_p: Optional[SmemPResource] = None, + store_softmax_stats: bool = False, **kwargs: Any, ) -> None: """Bind S/P TMEM offsets, Q peer index, and optional varlen metadata.""" @@ -2568,6 +2570,7 @@ def __init__( self.tmem_p_offset = tmem_p_offset self.q_half = q_half self.enable_early_tile_sum = cfg.enable_early_tile_sum + self.store_softmax_stats = store_softmax_stats self.q_offset_default = q_offset self.cum_seqlen_q = cum_seqlen_q self.cum_seqlen_k = cum_seqlen_k @@ -3278,6 +3281,30 @@ def _load_s_chunks_and_reduce_row_max( row_max_safe = Float32(0.0) return old_row_max, row_max_safe + @cute.jit + def _centered_scaled_score_pair( + self, + score_0: SoftmaxScalar, + score_1: SoftmaxScalar, + row_max: SoftmaxScalar, + scale: SoftmaxScalar, + ) -> tuple[SoftmaxScalar, SoftmaxScalar]: + """Center raw scores before scaling to avoid a large-logit FMA residual.""" + centered = cute.arch.add_packed_f32x2( + (score_0, score_1), + (-row_max, -row_max), + rnd="rn", + ftz=False, + ) + p_scale_log2 = Float32(self.cfg.pv_p_scale_log2) + return cute.arch.fma_packed_f32x2( + centered, + (scale, scale), + (p_scale_log2, p_scale_log2), + rnd="rn", + ftz=False, + ) + @cute.jit def _exp2_p_store( self, @@ -3338,18 +3365,27 @@ def _exp2_p_store( for chunk_idx in cutlass.range_constexpr(chunk_lo, chunk_hi): p_vals = () for elem_idx in cutlass.range_constexpr(0, tmem_x, 2): - fma_pair = cute.arch.fma_packed_f32x2( - ( + if cutlass.const_expr(self.store_softmax_stats): + fma_pair = self._centered_scaled_score_pair( s_data[chunk_idx][elem_idx], s_data[chunk_idx][elem_idx + 1], - ), - (scale, scale), - (minus_row_max_scale, minus_row_max_scale), - rnd="rn", - ftz=False, - ) + row_max, + scale, + ) + else: + fma_pair = cute.arch.fma_packed_f32x2( + ( + s_data[chunk_idx][elem_idx], + s_data[chunk_idx][elem_idx + 1], + ), + (scale, scale), + (minus_row_max_scale, minus_row_max_scale), + rnd="rn", + ftz=False, + ) if cutlass.const_expr( - elem_idx // 2 >= tmem_x // 2 - self.cfg.exp2_fma_pairs + not self.store_softmax_stats + and elem_idx // 2 >= tmem_x // 2 - self.cfg.exp2_fma_pairs ): p0, p1 = _exp2_fma_packed(fma_pair[0], fma_pair[1]) else: @@ -3535,16 +3571,24 @@ def _exp2_p_store_d128_fp8_cadence( num_chunks = num_values // tmem_x for chunk_idx in cutlass.range_constexpr(num_chunks): for local_idx in cutlass.range_constexpr(0, fma_lookahead, 2): - fma_pair = cute.arch.fma_packed_f32x2( - ( + if cutlass.const_expr(self.store_softmax_stats): + fma_pair = self._centered_scaled_score_pair( s_data[chunk_idx][local_idx], s_data[chunk_idx][local_idx + 1], - ), - (scale, scale), - (minus_row_max_scale, minus_row_max_scale), - rnd="rn", - ftz=False, - ) + row_max, + scale, + ) + else: + fma_pair = cute.arch.fma_packed_f32x2( + ( + s_data[chunk_idx][local_idx], + s_data[chunk_idx][local_idx + 1], + ), + (scale, scale), + (minus_row_max_scale, minus_row_max_scale), + rnd="rn", + ftz=False, + ) fma_ring[local_idx] = fma_pair[0] fma_ring[local_idx + 1] = fma_pair[1] @@ -3580,16 +3624,24 @@ def _exp2_p_store_d128_fp8_cadence( fma_1 = fma_ring[fma_idx + 1] if cutlass.const_expr(local_idx + fma_lookahead < tmem_x): future_idx = local_idx + fma_lookahead - fma_pair = cute.arch.fma_packed_f32x2( - ( + if cutlass.const_expr(self.store_softmax_stats): + fma_pair = self._centered_scaled_score_pair( s_data[chunk_idx][future_idx], s_data[chunk_idx][future_idx + 1], - ), - (scale, scale), - (minus_row_max_scale, minus_row_max_scale), - rnd="rn", - ftz=False, - ) + row_max, + scale, + ) + else: + fma_pair = cute.arch.fma_packed_f32x2( + ( + s_data[chunk_idx][future_idx], + s_data[chunk_idx][future_idx + 1], + ), + (scale, scale), + (minus_row_max_scale, minus_row_max_scale), + rnd="rn", + ftz=False, + ) fma_ring[fma_idx] = fma_pair[0] fma_ring[fma_idx + 1] = fma_pair[1] p_1 = cute.math.exp2(fma_1, fastmath=True) @@ -3686,16 +3738,24 @@ def _exp2_p_store_d256_fp8_cadence( for flat_idx in cutlass.range_constexpr(0, 8, 2): chunk_idx = flat_idx // tmem_x elem_idx = flat_idx % tmem_x - fma_pair = cute.arch.fma_packed_f32x2( - ( + if cutlass.const_expr(self.store_softmax_stats): + fma_pair = self._centered_scaled_score_pair( s_data[chunk_idx][elem_idx], s_data[chunk_idx][elem_idx + 1], - ), - (scale, scale), - (minus_row_max_scale, minus_row_max_scale), - rnd="rn", - ftz=False, - ) + row_max, + scale, + ) + else: + fma_pair = cute.arch.fma_packed_f32x2( + ( + s_data[chunk_idx][elem_idx], + s_data[chunk_idx][elem_idx + 1], + ), + (scale, scale), + (minus_row_max_scale, minus_row_max_scale), + rnd="rn", + ftz=False, + ) fma_values += (fma_pair[0], fma_pair[1]) p_values: tuple[Any, ...] = () @@ -3736,16 +3796,24 @@ def _exp2_p_store_d256_fp8_cadence( future_idx = flat_idx + 8 chunk_idx = future_idx // tmem_x elem_idx = future_idx % tmem_x - fma_pair = cute.arch.fma_packed_f32x2( - ( + if cutlass.const_expr(self.store_softmax_stats): + fma_pair = self._centered_scaled_score_pair( s_data[chunk_idx][elem_idx], s_data[chunk_idx][elem_idx + 1], - ), - (scale, scale), - (minus_row_max_scale, minus_row_max_scale), - rnd="rn", - ftz=False, - ) + row_max, + scale, + ) + else: + fma_pair = cute.arch.fma_packed_f32x2( + ( + s_data[chunk_idx][elem_idx], + s_data[chunk_idx][elem_idx + 1], + ), + (scale, scale), + (minus_row_max_scale, minus_row_max_scale), + rnd="rn", + ftz=False, + ) fma_values += (fma_pair[0], fma_pair[1]) p1 = cute.math.exp2(fma_values[flat_idx + 1], fastmath=True) p_values += (p0, p1) @@ -3937,6 +4005,12 @@ def variable_window_row_max( s_data[chunk_idx] = cutlass.Vector.from_elements( tuple(masked_scores), self.cfg.qk_acc_dtype ) + if cutlass.const_expr(self.store_softmax_stats): + # Fully masked earlier tiles use a safe zero exponent anchor. + # Discard it before the first visible tile so an all-negative + # valid row exports its actual maximum, not the safe anchor. + if window_start >= kv_tile_base: + row_max = Float32(-Float32.inf) return self._reduce_row_max(s_data, row_max) @consumer_work(returns=(old_row_max, row_max)) @@ -4010,6 +4084,9 @@ def left_masked_row_max( mask = mask & right_mask s_data[chunk_idx] = cutlass.vector.where(mask, s_data[chunk_idx], neg_inf) + if cutlass.const_expr(self.store_softmax_stats): + if window_bound_left >= kv_tile_abs * self.cfg.kv_tile_n: + row_max = Float32(-Float32.inf) return self._reduce_row_max(s_data, row_max) @consumer_work(returns=(old_row_max, row_max)) @@ -4154,7 +4231,9 @@ def right_masked_row_max( (seq_tile_coord * self.cfg.q_tile_m - self.cfg.window_size_left) // self.cfg.kv_tile_n, ) - if cutlass.const_expr(self.needs_window_tail_left_mask): + if cutlass.const_expr( + self.needs_window_tail_left_mask or self.store_softmax_stats + ): window_bound_left = bottom_right_window_left_bound( index_q, q_offset, @@ -4192,6 +4271,11 @@ def right_masked_row_max( s_data[chunk_idx] = cutlass.vector.where( mask, s_data[chunk_idx], neg_inf ) + if cutlass.const_expr( + self.store_softmax_stats and self.cfg.window_size_left > 0 + ): + if window_bound_left >= base_k: + row_max = Float32(-Float32.inf) return self._reduce_row_max(s_data, row_max) @cute.jit @@ -4531,6 +4615,9 @@ class TmemStatsResource(MemoryResource): tmem_vec_offset: Constexpr[int] = field(init=False, default=None) scale_softmax_log2: cute.Tensor | None = field(init=False, default=None) output_scale: cute.Tensor | None = field(init=False, default=None) + softmax_stats: cute.Tensor | None = field(init=False, default=None) + cum_seqlen_q: cute.Tensor | None = field(init=False, default=None) + q_half: Constexpr[int] = 0 tmem_addr_cached: TmemAddr | None = field(init=False, default=None) # Precomputed per-warp TMEM vec address (once, before persistent loop). tmem_vec_addr_cached: TmemAddr | None = field(init=False, default=None) @@ -4550,6 +4637,9 @@ def __init__( tmem_vec_offset: int, scale_softmax_log2: cute.Tensor | None = None, output_scale: cute.Tensor | None = None, + softmax_stats: cute.Tensor | None = None, + cum_seqlen_q: cute.Tensor | None = None, + q_half: int = 0, **kwargs: Any, ) -> None: """Bind the correction-stat TMEM vector offset and allocation.""" @@ -4558,6 +4648,9 @@ def __init__( self.tmem_vec_offset = tmem_vec_offset self.scale_softmax_log2 = scale_softmax_log2 self.output_scale = output_scale + self.softmax_stats = softmax_stats + self.cum_seqlen_q = cum_seqlen_q + self.q_half = q_half self._alloc = TmemAllocation(f"tmem_vec_{tmem_vec_offset}", cfg.tmem_stats_cols) self._smem_alloc = None if cfg.stats_via_smem: @@ -4754,6 +4847,40 @@ def store_vec( scale = exp2(scale_log2 * (old_max - new_max)) and to forward row_sum to SmemO for the final normalization. """ + if cutlass.const_expr(final_stats and self.softmax_stats is not None): + seq_coord, head_coord, batch_coord = _resolve_work_tile_coords( + self.cfg, stage_info.work_tile.tile_idx + ) + row_in_tile = ( + cute.arch.warp_idx() % 4 + ) * cute.arch.WARP_SIZE + cute.arch.lane_idx() + row = ( + seq_coord * self.cfg.q_tile_m * self.cfg.work_tile_q_seq_tiles + + self.q_half * self.cfg.peer_q_seq_tile_stride * self.cfg.q_tile_m + + row_in_tile + ) + head = ( + head_coord * self.cfg.work_tile_q_heads + + self.q_half * self.cfg.peer_q_head_stride + ) + # Softmax uses exp2(raw_score * scale_log2). The chunk merger + # expects maxima in natural-log units, not raw QK scores or LSE. + maximum = row_max * self.scale_softmax_log2[0] * Float32(0.6931471805599453) + # FP8 PV stores exp(logit - maximum) scaled by 448. Export the + # unscaled denominator; keep the internal correction stats intact. + denominator = row_sum / Float32(self.cfg.pv_p_scale) + if row_sum == Float32(0.0): + maximum = -Float32.inf + if cutlass.const_expr(self.cfg.has_varlen): + start = Int32(self.cum_seqlen_q[batch_coord]) + end = Int32(self.cum_seqlen_q[batch_coord + Int32(1)]) + if start + row < end: + self.softmax_stats[start + row, head, 0] = maximum + self.softmax_stats[start + row, head, 1] = denominator + else: + if row < self.softmax_stats.shape[1]: + self.softmax_stats[batch_coord, row, head, 0] = maximum + self.softmax_stats[batch_coord, row, head, 1] = denominator if cutlass.const_expr(self.cfg.stats_via_smem): stat0 = old_row_max if cutlass.const_expr(final_stats): @@ -5540,6 +5667,12 @@ def _store_o( # in TmemStats.consumer_work. correction_scale = vec_scale scale = output_scale * correction_scale / vec_row_sum + if cutlass.const_expr(self.tmem_vec_resource is not None): + if cutlass.const_expr(self.tmem_vec_resource.softmax_stats is not None): + # Absent chunks have zero accumulated O. Avoid 0 * inf so + # their outputs can be merged using the exported (m, s). + if vec_row_sum == Float32(0.0): + scale = Float32(0.0) num_correction_warps = 4 tidx, _, _ = cute.arch.thread_idx() diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_config.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_config.py index 42ba3d0e971d..24f58601043c 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_config.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_config.py @@ -1480,6 +1480,9 @@ def use_fp8_kv(self) -> bool: # denominator (extra logit that absorbs probability mass). Requires the # `attention_sinks` tensor argument at launch. use_attention_sinks: bool = False + # Export merge-compatible (natural-log maximum, unscaled denominator). + # Only separate reducers require an additional per-split maximum scalar. + store_softmax_stats: bool = False # Optional profile and reduction knobs selected by launcher policy or tests. use_keeps_mma_ab: bool = False # Sage attention: a nonzero ``sage_k_block_size`` enables per-block Q/K @@ -2576,8 +2579,12 @@ def defers_softmax_anchor_updates(self) -> bool: row maximum moves often but rarely by much. A zero headroom anchors on the exact row maximum. """ + # The public statistics contract requires the actual maximum, not a + # merge-equivalent delayed anchor. Keep the fast default policy intact; + # statistics-enabled kernels maintain the exact online maximum. return ( - self.use_keeps_mma_ab + not self.store_softmax_stats + and self.use_keeps_mma_ab and (self.tile_size_kv == 256 or self.use_block_sparse) and self.softmax_anchor_headroom_log2 > 0 ) 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..5dc25c60ebf5 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 @@ -330,6 +330,7 @@ def _build_decode_gen_schedule( partial_stats_ptr: cute.Pointer | None = None, split_kv_counter_ptr: cute.Pointer | None = None, attention_sinks_ptr: cute.Pointer | None = None, + softmax_stats_ptr: cute.Pointer | None = None, seqlens_kv: cute.Pointer | None = None, cu_seqlens_q: cute.Pointer | None = None, max_seq_len_kv: int | Int32 = 0, @@ -1423,6 +1424,7 @@ def _sage_k_scales(inst_id: int, softmax_grp) -> SageKScalesResource: partial_stats_ptr=partial_stats_ptr, split_kv_counter_ptr=split_kv_counter_ptr, attention_sinks_ptr=attention_sinks_ptr, + softmax_stats_ptr=softmax_stats_ptr, seqlens_kv=kv_seqlens, max_seq_len_kv=corr_max_seq_len_kv, num_heads_kv=num_heads_kv, @@ -1452,6 +1454,7 @@ def _sage_k_scales(inst_id: int, softmax_grp) -> SageKScalesResource: partial_stats_ptr=partial_stats_ptr, split_kv_counter_ptr=split_kv_counter_ptr, attention_sinks_ptr=attention_sinks_ptr, + softmax_stats_ptr=softmax_stats_ptr, seqlens_kv=kv_seqlens, max_seq_len_kv=corr_max_seq_len_kv, num_heads_kv=num_heads_kv, @@ -2589,6 +2592,7 @@ def _run_decode_gen_active( tma_desc_k_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, tma_desc_v_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, g_sage_scales: cute.Pointer | None = None, + g_softmax_stats: cute.Pointer | None = None, ) -> None: """Run the complete decode body for one runtime-valid Q tile. @@ -2709,6 +2713,7 @@ def _run_decode_gen_active( partial_stats_ptr=g_partial_stats, split_kv_counter_ptr=g_split_kv_counter, attention_sinks_ptr=g_attention_sinks, + softmax_stats_ptr=g_softmax_stats, seqlens_kv=runtime_seqlens_kv, cu_seqlens_q=g_cu_seqlens_q, max_seq_len_kv=runtime_max_seq_len_kv, @@ -2930,6 +2935,7 @@ def _run_decode_gen_runtime_prefix( tma_desc_k_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, tma_desc_v_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, g_sage_scales: cute.Pointer | None = None, + g_softmax_stats: cute.Pointer | None = None, ) -> None: """Run the general runtime split-prefix producer or retire its suffix.""" @@ -3015,6 +3021,7 @@ def _run_decode_gen_runtime_prefix( tma_desc_k_summary_atom=tma_desc_k_summary_atom, tma_desc_v_summary_atom=tma_desc_v_summary_atom, g_sage_scales=g_sage_scales, + g_softmax_stats=g_softmax_stats, ) else: _run_decode_gen_inactive_cluster_rank() @@ -3069,6 +3076,7 @@ def _run_decode_gen_runtime_prefix( tma_desc_k_summary_atom=tma_desc_k_summary_atom, tma_desc_v_summary_atom=tma_desc_v_summary_atom, g_sage_scales=g_sage_scales, + g_softmax_stats=g_softmax_stats, ) else: _signal_padded_pdl_producer(cfg) @@ -3117,6 +3125,7 @@ def decode_gen_kernel( tma_desc_k_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, tma_desc_v_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, g_sage_scales: cute.Pointer | None = None, + g_softmax_stats: cute.Pointer | None = None, ) -> None: """Dispatch one static Q/split tile and drain padded launch slots safely.""" q_group_cta_idx, h_k_idx, b_idx = cute.arch.block_idx() @@ -3203,6 +3212,7 @@ def decode_gen_kernel( g_sparse_row_route_offsets, g_sparse_row_route_counts, g_sparse_route_metadata, + g_softmax_stats=g_softmax_stats, ) elif cutlass.const_expr(cfg.use_split_kv and cfg.use_pdl): # The immediately preceding grid may produce seq_lens. Send every @@ -3307,6 +3317,7 @@ def decode_gen_kernel( tma_desc_k_summary_atom=tma_desc_k_summary_atom, tma_desc_v_summary_atom=tma_desc_v_summary_atom, g_sage_scales=g_sage_scales, + g_softmax_stats=g_softmax_stats, ) else: # Packed-Q grids use a batch-wide maximum envelope. These Q CTAs own no @@ -3354,6 +3365,7 @@ def fmha_decode_launch( static_full_split_prefix: cutlass.Constexpr[bool] = False, use_static_native_seqlens_kv: cutlass.Constexpr[bool] = False, sage_scales_iter: cute.Pointer | None = None, + softmax_stats_iter: cute.Pointer | None = None, ) -> None: """Standalone JIT launcher for FMHA decode TS. @@ -3749,6 +3761,7 @@ def fmha_decode_launch( null_sparse_route_ptr, static_full_split_prefix, g_sage_scales=sage_scales_iter, + g_softmax_stats=softmax_stats_iter, ).launch( grid=grid, block=[cfg.threads_per_cta, 1, 1], diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_p.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_p.py index bb3448d40352..4e254ca60ef7 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_p.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_p.py @@ -336,6 +336,10 @@ def _exponent_addend( if cutlass.const_expr(guards_masked_rows): if safe_new_max == _neg_max_f32(): safe_new_max = Float32(0.0) + if cutlass.const_expr(cfg.store_softmax_stats): + # Statistics-enabled paths center the raw scores before scaling, + # avoiding cancellation against a rounded absolute exponent offset. + safe_new_max = Float32(0.0) # Not ``cute.math.fma``: the fused form makes ptxas spill in the callers. addend = Float32(-self.scale_softmax_log2 * safe_new_max) if cutlass.const_expr(self.cfg.use_fp8_pv): @@ -525,6 +529,16 @@ def _compute_p_fragments_impl( prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) for score_idx in cutlass.range_constexpr(fragment_regs): s_arr[score_idx] = pending_scores[score_idx] + if cutlass.const_expr(cfg.store_softmax_stats): + safe_new_max = new_max_arr[0] + if safe_new_max == _neg_max_f32(): + safe_new_max = Float32(0.0) + for pair_idx in cutlass.range_constexpr(fragment_regs // 2): + idx = pair_idx * 2 + s_arr[idx], s_arr[idx + 1] = fadd2( + (s_arr[idx], s_arr[idx + 1]), (-safe_new_max, -safe_new_max) + ) + if cutlass.const_expr(dequantized): self._scale_fragment_pairs( s_arr, row_addends, row_multipliers, groups=1 @@ -979,6 +993,20 @@ def _compute_p_impl( ) -> None: """Compute P from S, stage its BMM2 operand, and publish local sums.""" cfg = self.cfg + if cutlass.const_expr(cfg.store_softmax_stats): + # Keeps lanes own a row; the small-Q layouts interleave pairs of + # rows and repeat their scale groups across the K fragments. + for score_idx in cutlass.range_constexpr(cfg.num_s_regs_per_thread): + if cutlass.const_expr(cfg.use_keeps_mma_ab): + scale_idx = 0 + else: + scale_idx = (score_idx // 4) % max( + cfg.tile_size_q // 8, 1 + ) * 2 + score_idx % 2 + safe_new_max = new_max_arr[scale_idx] + if safe_new_max == _neg_max_f32(): + safe_new_max = Float32(0.0) + s_arr[score_idx] = s_arr[score_idx] - safe_new_max if cutlass.const_expr(cfg.use_keeps_mma_ab): self._compute_keeps_p( stage_info, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_corr.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_corr.py index 9e46fae39bcc..762572c12280 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_corr.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_corr.py @@ -21,6 +21,7 @@ """ from dataclasses import dataclass +import math import cutlass import cutlass.cute as cute @@ -125,6 +126,7 @@ class TmemCorrResource(DecodeGenResourceBase): # the schedule waits and releases it, the epilogue only reads it. ``None`` # without Sage attention. sage_v_scales_ref: Constexpr[SageVScalesResource | None] = None + softmax_stats_ptr: cute.Pointer | None = None seqlens_kv: cute.Pointer = None max_seq_len_kv: Constexpr[int] = 0 seq_len_q: Int32 = None @@ -620,6 +622,88 @@ def _separate_partial_lse(self, max_val: Float32, sum_val: Float32) -> Float32: ) return lse_val + @cute.jit + def _store_separate_stats( + self, workspace_row: Int64, max_val: Float32, sum_val: Float32 + ) -> None: + """Publish LSE and optional raw QK maximum / unscaled denominator.""" + stride = 3 if self.cfg.store_softmax_stats else 1 + ptr = self.partial_stats_ptr + workspace_row * Int64(stride) + ptr.store(self._separate_partial_lse(max_val, sum_val)) + if cutlass.const_expr(self.cfg.store_softmax_stats): + raw_max = Float32(-Float32.inf) + if sum_val > Float32(0.0): + raw_max = max_val + denominator = sum_val + if cutlass.const_expr(self.cfg.use_fp8_pv): + denominator *= Float32(1.0 / 448.0) + (ptr + Int32(1)).store(raw_max) + (ptr + Int32(2)).store(denominator) + + @cute.jit + def _store_final_stats( + self, + logical_h_k_idx: Int32, + logical_b_idx: Int32, + logical_row: Int32, + max_val: Float32, + sum_val: Float32, + ) -> None: + """Write one owned row after its final reduction and sink contribution.""" + if cutlass.const_expr(self.cfg.store_softmax_stats): + row = _q_physical_output_row_from_logical( + self.cfg, + self.h_r, + self.num_heads_kv, + logical_b_idx, + logical_h_k_idx, + logical_row, + self.q_token_offset, + ) + scaled_max = Float32(-Float32.inf) + if sum_val > Float32(0.0): + scaled_max = max_val * self.scale_softmax_log2 * Float32(math.log(2.0)) + denominator = sum_val + if cutlass.const_expr(self.cfg.use_fp8_pv): + denominator *= Float32(1.0 / 448.0) + (self.softmax_stats_ptr + Int64(row) * Int64(2)).store(scaled_max) + (self.softmax_stats_ptr + Int64(row) * Int64(2) + Int64(1)).store( + denominator + ) + + @cute.jit + def _store_direct_swaps_stats( + self, + stage_info: StageInfo, + warp_idx: Int32, + lane_idx: Int32, + final_max: cutlass.Array, + reduced_sum: cutlass.Array, + num_scale_groups: Constexpr[int], + ) -> None: + if cutlass.const_expr(self.cfg.store_softmax_stats): + if warp_idx == Int32(0) and lane_idx < Int32(4 * num_scale_groups): + stats_idx = lane_idx >> Int32(2) + row = ( + (lane_idx & Int32(3)) * Int32(2) + + (stats_idx >> Int32(1)) * Int32(8) + + (stats_idx & Int32(1)) + ) + q_group = _logical_q_group_idx(self.cfg, stage_info, self.q_group_idx) + if _q_row_is_valid_for_seq( + self.cfg, self.h_r, q_group, row, self.seq_len_q + ): + head, batch = _logical_head_batch( + stage_info, self.h_k_idx, self.b_idx + ) + self._store_final_stats( + head, + batch, + _q_tile_output_row_base(self.cfg, q_group) + row, + final_max[stats_idx], + reduced_sum[stats_idx], + ) + @cute.jit def _pack_separate_partial_o_pair(self, val0: Float32, val1: Float32) -> Int32: """Pack normalized separate partial O in the selected 16-bit type.""" @@ -1109,6 +1193,7 @@ def _store_softmax_normalized_o_vec8( reduce_row_idx, sum_val, max_val, + store_stats=reduce_col_idx == Int32(0), ) dst_offset = dst_row_base + reduce_col_idx * Int32(self.cfg.o_dtype_bytes) final_o_dst = cutlass.inttoptr( @@ -1132,6 +1217,7 @@ def _softmax_output_row_state( reduce_row_idx: Int32, sum_val: Float32, max_val: Float32, + store_stats: cutlass.Boolean = False, ) -> tuple[Int64, Float32]: """Resolve one logical row's output address and normalization once. @@ -1153,6 +1239,11 @@ def _softmax_output_row_state( self.num_heads_kv, attention_sink_head_idx, ) + if cutlass.const_expr(cfg.store_softmax_stats): + if store_stats: + self._store_final_stats( + logical_h_k_idx, logical_b_idx, reduce_row_idx, max_val, sum_val + ) # ``output_scale`` is the public bmm2_scale. Fold it into the final # normalization for every output dtype; split partials reach this # helper only after the cross-CTA reduction has completed. @@ -2237,16 +2328,8 @@ def _store_split_stats_from_scale_arrays( cta_idx_kv, stats_row_idx, ) - lse_ptr = cutlass.inttoptr( - self.partial_stats_ptr.toint() + lse_row * Int64(4), - mem_space=1, - dtype=Float32, - ) - lse_ptr.store( - self._separate_partial_lse( - final_max[stats_idx], reduced_sum[stats_idx] - ), - alignment=4, + self._store_separate_stats( + lse_row, final_max[stats_idx], reduced_sum[stats_idx] ) elif cutlass.const_expr(enable_cluster): cluster_owner = self._cluster_reduction_cta_for_row( @@ -2595,6 +2678,7 @@ def _kv_tile_256_output_row_state( self, *, resolves: cutlass.Boolean, + stats_owner: cutlass.Boolean, logical_h_k_idx: Int32, logical_b_idx: Int32, logical_kv_idx: Int32, @@ -2637,6 +2721,7 @@ def _kv_tile_256_output_row_state( logical_output_row_idx, denominator, global_max, + store_stats=valid_output_row & stats_owner, ) return valid_output_row, row_base, row_scale @@ -2679,6 +2764,7 @@ def _kv_tile_256_merge_spatial_output( logical_output_row_idx = q_row_offset + exchange_row_idx valid_output_row, row_base, row_scale = self._kv_tile_256_output_row_state( resolves=output_lane, + stats_owner=output_lane, logical_h_k_idx=logical_h_k_idx, logical_b_idx=logical_b_idx, logical_kv_idx=logical_kv_idx, @@ -2773,15 +2859,7 @@ def _store_kv_tile_256_partial_stats( logical_output_row_idx, ) if cutlass.const_expr(self.cfg.use_separate_reduction_kernel): - stats_ptr = cutlass.inttoptr( - self.partial_stats_ptr.toint() + stats_row * Int64(4), - mem_space=1, - dtype=Float32, - ) - stats_ptr.store( - self._separate_partial_lse(global_max, denominator), - alignment=4, - ) + self._store_separate_stats(stats_row, global_max, denominator) else: stats_ptr = cutlass.inttoptr( self.partial_stats_ptr.toint() + stats_row * Int64(2 * 4), @@ -3013,6 +3091,7 @@ def _kv_tile_256_merge_owned_columns( logical_output_row_idx = q_row_offset + exchange_row_idx valid_output_row, row_base, row_scale = self._kv_tile_256_output_row_state( resolves=cutlass.Boolean(True), + stats_owner=owned_col_base == Int32(0), logical_h_k_idx=logical_h_k_idx, logical_b_idx=logical_b_idx, logical_kv_idx=logical_kv_idx, @@ -3311,15 +3390,7 @@ def _keeps_tail_epilogue( cta_idx_kv, logical_output_row_idx, ) - stats_ptr = cutlass.inttoptr( - self.partial_stats_ptr.toint() + stats_row * Int64(4), - mem_space=1, - dtype=Float32, - ) - stats_ptr.store( - self._separate_partial_lse(final_max_0, reduced_sum_0), - alignment=4, - ) + self._store_separate_stats(stats_row, final_max_0, reduced_sum_0) else: stats_row = self._gmem_partial_row_offset( logical_kv_idx, @@ -3373,6 +3444,15 @@ def _keeps_tail_epilogue( self.num_heads_kv, attention_sink_head_idx, ) + if cutlass.const_expr(cfg.store_softmax_stats): + if col_base == Int32(0) and valid_output_row: + self._store_final_stats( + logical_h_k_idx, + logical_b_idx, + logical_output_row_idx, + final_max_0, + reduced_sum_0, + ) # Apply public bmm2_scale in the final direct-output normalization. # The split-KV branch above returns before reaching this point, so its # partial O remains unscaled for the selected reducer. @@ -3975,6 +4055,9 @@ def _swaps_wide_tail_epilogue( ) return + self._store_direct_swaps_stats( + stage_info, warp_idx, lane_idx, final_max, reduced_sum, num_scale_groups + ) for scale_idx in cutlass.range_constexpr(num_scale_groups): # Direct final-output scales include both the reciprocal softmax # denominator and public bmm2_scale. Split partials returned above @@ -4482,6 +4565,9 @@ def _swaps_q8_tail_epilogue( ) return + self._store_direct_swaps_stats( + stage_info, warp_idx, lane_idx, final_max, reduced_sum, num_scale_groups + ) for scale_idx in cutlass.range_constexpr(num_scale_groups): # Direct final-output scales include both the reciprocal softmax # denominator and public bmm2_scale. Split partials returned above 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..5623287f7cd6 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 @@ -15,9 +15,12 @@ """Standalone GMEM reducers for FMHA decode TS split-KV profiles. The decode kernel publishes one normalized 16-bit O vector and one FP32 -log2-LSE scalar per ``(batch, kv_head, split_kv, output_row)``. Reducer threads -own 16-byte output fragments and combine those normalized states with the -shared log2-LSE recurrence. +log2-LSE scalar per ``(batch, kv_head, split_kv, output_row)``. Statistics-enabled +specializations append the raw QK maximum and actual unscaled denominator. +Reducer threads own 16-byte output fragments and combine normalized states with +the shared log2-LSE +recurrence; only the final row owner additionally merges the optional max/sum +states independently of the rounded absolute LSE. The production schedule is selected from the split count: @@ -35,6 +38,8 @@ only split-KV is reduced here. """ +import math + import cutlass import cutlass.cute as cute from .direct_sparse_metadata import DirectSparseMetadataView, HeadIndexedMetadataView @@ -149,11 +154,14 @@ def _reduce_exact_splits_body( cfg: cutlass.Constexpr[FmhaDecodeConfig], static_full_split_prefix: cutlass.Constexpr[bool] = False, wait_for_pdl_producer: cutlass.Constexpr[bool] = False, + g_softmax_stats: cute.Pointer | None = None, + stats_scale: Float32 = 1.0, ) -> None: """Fold every split in one 512-thread CTA over one 8 KiB output slice. ``g_partial_stats`` stores one log2-LSE scalar for each split and output - row. ``g_partial_o`` stores the corresponding normalized 16-bit O fragment. + row, or an interleaved (log2-LSE, raw QK maximum, denominator) triple when + exporting statistics. ``g_partial_o`` stores the normalized 16-bit O fragment. This body is shared by the serial reference kernel and the compact S2-S4 production schedule. The production caller requests its PDL acquire after this body's local register state is initialized and before partial reads. @@ -249,7 +257,12 @@ def _reduce_exact_splits_body( cfg, ) stats_offset = workspace_row * Int64( - SEPARATE_REDUCTION_LSE_VALUES_PER_ROW * FP32_BYTES + ( + 3 + if cfg.store_softmax_stats + else SEPARATE_REDUCTION_LSE_VALUES_PER_ROW + ) + * FP32_BYTES ) stats_src = cutlass.inttoptr( g_partial_stats.toint() + stats_offset, @@ -297,6 +310,13 @@ def _reduce_exact_splits_body( grid_h_k, attention_sink_head_idx, cfg, + g_softmax_stats, + g_partial_stats, + logical_kv_idx, + reduce_row_idx, + g_q_output_rows, + active_splits_kv, + stats_scale, ) @@ -311,6 +331,8 @@ def decode_gen_separate_reduction_kernel( g_q_output_rows: Int32, cfg: cutlass.Constexpr[FmhaDecodeConfig], static_full_split_prefix: cutlass.Constexpr[bool] = False, + g_softmax_stats: cute.Pointer | None = None, + stats_scale: Float32 = 1.0, ) -> None: """Run the 512-thread exact-split reference schedule.""" @@ -324,6 +346,8 @@ def decode_gen_separate_reduction_kernel( g_q_output_rows, cfg, static_full_split_prefix, + g_softmax_stats=g_softmax_stats, + stats_scale=stats_scale, ) @@ -368,12 +392,20 @@ def _store_parallel_reduction_output( grid_h_k: Int32, attention_sink_head_idx: Int32, cfg: cutlass.Constexpr[FmhaDecodeConfig], + g_softmax_stats: cute.Pointer | None = None, + g_partial_stats: cute.Pointer | None = None, + logical_kv_idx: Int64 = 0, + logical_output_row_idx: Int32 = 0, + rows_per_split: Int32 = 0, + active_splits_kv: Int32 = 0, + stats_scale: Float32 = 1.0, ) -> None: """Merge the optional sink and store one normalized output fragment. Producer CTAs already fold the softmax scale into log2-LSE and, for FP8 output, fold the output quantization scale into normalized partial O. The - standalone reducer therefore only merges normalized states here. + output reduction therefore only merges normalized states here. Statistics + retain raw maxima and apply the runtime scale only after subtraction. """ sink_lse = _attention_sink_log2_lse( @@ -385,6 +417,60 @@ def _store_parallel_reduction_output( attention_sink_head_idx, ) _, split_weight, _ = merge_log2_lse(global_lse, sink_lse) + if cutlass.const_expr(cfg.store_softmax_stats): + if reduce_col_idx == Int32(0): + # Only the final row owner merges the optional max/sum states. + # The original LSE/O recurrence and cluster communication stay unchanged; + # statistics-disabled kernels allocate no additional state. + raw_maximum = Float32(-Float32.inf) + for split_idx in cutlass.range(active_splits_kv): + workspace_row = _separate_workspace_row_offset( + logical_kv_idx, + split_idx, + logical_output_row_idx, + rows_per_split, + cfg, + ) + partial_max = ( + g_partial_stats + workspace_row * Int64(3) + Int64(1) + ).load() + raw_maximum = cute.math.max(raw_maximum, partial_max, ftz=True) + denominator = Float32(0.0) + maximum = Float32(-Float32.inf) + if raw_maximum != Float32(-Float32.inf): + maximum = raw_maximum * stats_scale + scale_log2 = stats_scale * Float32(math.log2(math.e)) + for split_idx in cutlass.range(active_splits_kv): + workspace_row = _separate_workspace_row_offset( + logical_kv_idx, + split_idx, + logical_output_row_idx, + rows_per_split, + cfg, + ) + partial_ptr = g_partial_stats + workspace_row * Int64(3) + partial_max = (partial_ptr + Int64(1)).load() + partial_sum = (partial_ptr + Int64(2)).load() + # Preserve small differences between large finite raw + # maxima before conversion to base-2 exponent units. + denominator += partial_sum * cute.math.exp2( + (partial_max - raw_maximum) * scale_log2, fastmath=True + ) + if cutlass.const_expr(cfg.use_attention_sinks): + # Load the unscaled sink directly: adding/removing the FP8 + # probability scale in absolute LSE can itself lose bits. + head_idx = cute.math.min( + h_k_idx * attention_sink_h_r + attention_sink_head_idx, + attention_sink_h_r * grid_h_k - Int32(1), + ) + sink_delta = g_attention_sinks[head_idx] - maximum + denominator += cute.math.exp2( + sink_delta * Float32(math.log2(math.e)), fastmath=True + ) + (g_softmax_stats + Int64(output_row_idx) * Int64(2)).store(maximum) + (g_softmax_stats + Int64(output_row_idx) * Int64(2) + Int64(1)).store( + denominator + ) final_regs = cutlass.Array( Int32, @@ -457,6 +543,8 @@ def decode_gen_parallel_separate_reduction_kernel( g_q_output_rows: Int32, cfg: cutlass.Constexpr[FmhaDecodeConfig], static_full_split_prefix: cutlass.Constexpr[bool] = False, + g_softmax_stats: cute.Pointer | None = None, + stats_scale: Float32 = 1.0, ) -> None: """Reduce split partials with a compact or clustered constexpr schedule. @@ -481,6 +569,8 @@ def decode_gen_parallel_separate_reduction_kernel( cfg, static_full_split_prefix, cfg.use_parallel_separate_reduction_pdl, + g_softmax_stats, + stats_scale, ) return @@ -614,7 +704,12 @@ def decode_gen_parallel_separate_reduction_kernel( cfg, ) stats_offset = workspace_row * Int64( - SEPARATE_REDUCTION_LSE_VALUES_PER_ROW * FP32_BYTES + ( + 3 + if cfg.store_softmax_stats + else SEPARATE_REDUCTION_LSE_VALUES_PER_ROW + ) + * FP32_BYTES ) stats_src = cutlass.inttoptr( g_partial_stats.toint() + stats_offset, @@ -727,6 +822,13 @@ def decode_gen_parallel_separate_reduction_kernel( grid_h_k, attention_sink_head_idx, cfg, + g_softmax_stats, + g_partial_stats, + logical_kv_idx, + reduce_row_idx, + g_q_output_rows, + active_splits_kv, + stats_scale, ) return @@ -883,6 +985,13 @@ def decode_gen_parallel_separate_reduction_kernel( grid_h_k, attention_sink_head_idx, cfg, + g_softmax_stats, + g_partial_stats, + logical_kv_idx, + reduce_row_idx, + g_q_output_rows, + active_splits_kv, + stats_scale, ) # Keep every peer CTA alive until rank zero has finished its DSMEM reads. @@ -904,14 +1013,15 @@ def fmha_decode_separate_reduction_launch( stream: cuda_drv.CUstream, cfg: cutlass.Constexpr[FmhaDecodeConfig], static_full_split_prefix: cutlass.Constexpr[bool] = False, + softmax_stats_iter: cute.Pointer | None = None, ) -> None: """Launch the standalone reducer after the main split-KV decode kernel. 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. Producers already apply + ``scale_s`` and ``output_scale`` to log2-LSE and normalized partial O. + The statistics-only max/sum merge uses ``scale_s`` on raw maximum differences. ``seqlens_kv_iter`` is internal reducer metadata used only to bound the runtime split prefix; the decode launch ABI remains unchanged. """ @@ -960,6 +1070,8 @@ def fmha_decode_separate_reduction_launch( q_output_rows, cfg, static_full_split_prefix, + softmax_stats_iter, + scale_s, ).launch( grid=(num_parallel_slices * cluster_size, h_k, b), block=[cfg.parallel_reduction_threads_per_cta, 1, 1], @@ -978,6 +1090,8 @@ def fmha_decode_separate_reduction_launch( q_output_rows, cfg, static_full_split_prefix, + softmax_stats_iter, + scale_s, ).launch( grid=(num_reduction_slices, h_k, b), # Reduction parallelism is entirely in block.x; y/z are singleton diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/softmax_stats.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/softmax_stats.py new file mode 100644 index 000000000000..224f154229ee --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/softmax_stats.py @@ -0,0 +1,79 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Optional merge-compatible MLA statistics without changing output reduction.""" + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int64 + +from .ops import fmax_f32 + + +@cute.jit +def partial_max_tensor(acc_lse, softmax_stats): + """Bind optional raw QK maxima; the matching sums follow this tensor.""" + result = None + if cutlass.const_expr(softmax_stats is not None and acc_lse is not None): + result = cute.make_tensor( + acc_lse.iterator + Int64(cute.cosize(acc_lse.layout)), acc_lse.layout + ) + return result + + +@cute.jit +def store_softmax_stats(softmax_stats, row, raw_maximum, denominator, softmax_scale): + """Publish actual natural-log max and unscaled exp-sum from one owner.""" + if cutlass.const_expr(softmax_stats is not None): + maximum = raw_maximum * softmax_scale + if denominator == Float32(0.0): + maximum = Float32(-Float32.inf) + ptr = softmax_stats.iterator.raw_ptr() + Int64(row) * Int64(2) + ptr.store(maximum) + (ptr + Int64(1)).store(denominator) + + +@cute.jit +def store_reduced_softmax_stats( + softmax_stats, row, acc_lse, row_lse, active_splits, softmax_scale +): + """Merge actual max/sum states on the final statistics writer lane. + + Existing output reducers use LSE to rescale partial O. Their LSE maximum + is not the logit maximum. Absolute FP32 LSE also loses small log-sums at + large finite maxima, so retain both maximum and denominator separately. + Store raw QK maxima and scale only their differences: scaling absolute + maxima first would round away small gaps between large finite scores. + Inactive split slots are deliberately never read. + """ + if cutlass.const_expr(softmax_stats is not None): + scale_log2 = softmax_scale * Float32(1.4426950408889634) + partial_rows = Int64(cute.cosize(acc_lse.layout)) + max_ptr = row_lse.iterator.raw_ptr() + partial_rows + sum_ptr = max_ptr + partial_rows + maximum = Float32(-Float32.inf) + for split_idx in cutlass.range(active_splits): + if row_lse[split_idx] != Float32(-Float32.inf): + value = (max_ptr + Int64(split_idx)).load() + maximum = fmax_f32(maximum, value) + denominator = Float32(0.0) + for split_idx in cutlass.range(active_splits): + if row_lse[split_idx] != Float32(-Float32.inf): + partial_max = (max_ptr + Int64(split_idx)).load() + partial_sum = (sum_ptr + Int64(split_idx)).load() + denominator += partial_sum * cute.math.exp2( + (partial_max - maximum) * scale_log2, fastmath=True + ) + store_softmax_stats(softmax_stats, row, maximum, denominator, softmax_scale) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/config.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/config.py index fe2461ad74c4..c144ac7babe1 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/config.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/config.py @@ -104,6 +104,7 @@ def compute_workspace_size( split_kv: int, partial_o_dtype, lse_dtype, + store_softmax_stats: bool = False, ) -> int: """Return the physical flat-tile split-KV workspace size in bytes.""" @@ -115,7 +116,9 @@ def compute_workspace_size( raise ValueError(f"num_q_tiles must be positive, got {num_q_tiles}") partial_rows = batch_size * tile_size_q * num_q_tiles * split_kv return partial_rows * ( - latent_dim * partial_o_dtype.width // 8 + lse_dtype.width // 8 + latent_dim * partial_o_dtype.width // 8 + + lse_dtype.width // 8 + + (8 if store_softmax_stats else 0) ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/kernel.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/kernel.py index ffc5c7fb609c..36652e5496de 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/kernel.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/kernel.py @@ -159,6 +159,7 @@ def build_mla_decode_task_manager( acc_output=None, lse=None, acc_lse=None, + softmax_stats=None, # Domain (k_tile_count) domain=4, # Persistent loop support @@ -441,6 +442,7 @@ def build_mla_decode_task_manager( tmem_s = TmemSResource( page_offsets=page_offsets, + store_softmax_stats=softmax_stats is not None, smem_q_latent=smem_q_latent_arr, smem_q_rope=smem_q_rope_arr, smem_p=smem_p_arr, @@ -491,6 +493,7 @@ def build_mla_decode_task_manager( partial_output=acc_output, lse=lse, partial_lse=acc_lse, + softmax_stats=softmax_stats, tmem_o_ref=tmem_o, tmem_corr_ref=tmem_corr, atten_sinks=atten_sinks, @@ -1112,6 +1115,7 @@ def __call__( output_scale: cutlass.Float32, stream: object, output_quant_args=None, + softmax_stats: cute.Tensor | None = None, ): """Execute the MLA decode TS kernel.""" if cutlass.const_expr(self.fuse_sparse_reduction): @@ -1400,6 +1404,7 @@ def _flatten_query_rows_for_tma(t): ) softmax_scale_log2 = softmax_scale * LOG2_E + softmax_stats_scale = softmax_scale if softmax_stats is not None else None kernel_split_kv = ( Int32(self.static_split_kv) @@ -1473,6 +1478,8 @@ def _flatten_query_rows_for_tma(t): clc_tile_sched_params, output_quant_args, tma_desc_o, + softmax_stats, + softmax_stats_scale, ).launch( grid=grid, block=[cfg.threads_per_cta, 1, 1], @@ -1498,6 +1505,8 @@ def _flatten_query_rows_for_tma(t): cu_seqlens_q, block_split_kvs, output_quant_args, + softmax_stats, + softmax_stats_scale, ).launch( grid=( ( @@ -1528,6 +1537,8 @@ def _flatten_query_rows_for_tma(t): cache_seqs, cu_seqlens_q, block_split_kvs, + softmax_stats, + softmax_stats_scale, ).launch( grid=( ceil_div(logical_query_rows, REDUCTION_ROWS_PER_CTA), @@ -1571,6 +1582,8 @@ def split_kv_kernel( clc_tile_sched_params: object, output_quant_args=None, tma_desc_o: cutlass.GridConstant[cuda.TensorMap] = None, + softmax_stats: cute.Tensor | None = None, + softmax_stats_scale=None, ) -> None: """MLA decode TS kernel: persistent tile-scheduled execution.""" if cutlass.const_expr(self.direct_sparse): @@ -2005,6 +2018,7 @@ def split_kv_kernel( acc_output=acc_o, lse=lse, acc_lse=acc_lse, + softmax_stats=softmax_stats, domain=Int32(1), # dummy; MlaTask recomputes per-task to avoid spills work_queue=mla_work_queue, cache_seqs=cache_seqs, @@ -2077,6 +2091,7 @@ def split_kv_kernel( # Set runtime params on GmemO (epilogue) named_res["gmem_o"].output_scale = output_scale named_res["gmem_o"].softmax_scale_log2 = softmax_scale_log2 + named_res["gmem_o"].softmax_stats_scale = softmax_stats_scale named_res["gmem_o"].smem_exchange = epilogue_exchange_arr.data_ptr().toint( Int32 ) @@ -2157,6 +2172,8 @@ def reduction_kernel( cache_seqs: cute.Tensor, cu_seqlens_q: cute.Tensor | None, block_split_kvs: cute.Tensor, + softmax_stats: cute.Tensor | None = None, + softmax_stats_scale=None, ): """Dispatch the throughput 2CTA split-KV reduction body.""" cfg = make_mla_decode_config( @@ -2183,6 +2200,8 @@ def reduction_kernel( cfg, self.reduction_split_capacity, REDUCTION_ROWS_PER_CTA, + softmax_stats, + softmax_stats_scale, ) @cute.kernel @@ -2197,6 +2216,8 @@ def parallel_reduction_kernel( cu_seqlens_q: cute.Tensor | None, block_split_kvs: object, output_quant_args=None, + softmax_stats: cute.Tensor | None = None, + softmax_stats_scale=None, ): """Dispatch the high-split fixed-D512 cluster reducer.""" @@ -2236,4 +2257,6 @@ def parallel_reduction_kernel( self.fuse_sparse_reduction, self.output_quant, output_quant_args, + softmax_stats, + softmax_stats_scale, ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/parallel_reduction.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/parallel_reduction.py index 41b61707a4c4..eb0e39ba4efb 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/parallel_reduction.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/parallel_reduction.py @@ -29,6 +29,7 @@ from ..helpers.math import ceil_div from ..helpers.ops import fmax_f32, warp_reduce_max_f32, warp_reduce_sum_f32 from ..helpers.query import flat_query_row_state, query_batch_bounds +from ..helpers.softmax_stats import store_reduced_softmax_stats from .work_partition import ( runtime_row_prefix_active_split_count, runtime_split_kv_cap, @@ -154,6 +155,11 @@ def _store_parallel_reduction_result( finalize_attention: cutlass.Constexpr[bool] = False, output_quant: cutlass.Constexpr = None, output_quant_args=None, + softmax_stats=None, + acc_lse=None, + row_lse=None, + active_split_kv=None, + softmax_stats_scale=None, ): """Publish one normalized FP32 output fragment and final LSE.""" @@ -164,6 +170,19 @@ def _store_parallel_reduction_result( lse[logical_head_idx, storage_q_idx] = global_lse else: lse[logical_head_idx, logical_q_idx, batch_idx] = global_lse + public_row = logical_head_idx + output.shape[0] * storage_q_idx + if cutlass.const_expr(cu_seqlens_q is None): + public_row = logical_head_idx + output.shape[0] * ( + logical_q_idx + output.shape[2] * batch_idx + ) + store_reduced_softmax_stats( + softmax_stats, + public_row, + acc_lse, + row_lse, + active_split_kv, + softmax_stats_scale, + ) if cutlass.const_expr(output_quant is not None): if query_is_valid: @@ -216,6 +235,8 @@ def run_parallel_reduction_kernel( finalize_attention: cutlass.Constexpr[bool] = False, output_quant: cutlass.Constexpr = None, output_quant_args=None, + softmax_stats=None, + softmax_stats_scale=None, ): """Reduce one D=512 row cooperatively across a padded CTA cluster. @@ -419,6 +440,11 @@ def run_parallel_reduction_kernel( finalize_attention, output_quant, output_quant_args, + softmax_stats, + acc_lse, + acc_lse[row_in_tile, None, query_tile_idx, batch_idx], + active_split_kv, + softmax_stats_scale, ) return @@ -536,6 +562,11 @@ def run_parallel_reduction_kernel( tidx, batch_idx, cu_seqlens_q, + softmax_stats=softmax_stats, + acc_lse=acc_lse, + row_lse=acc_lse[row_in_tile, None, query_tile_idx, batch_idx], + active_split_kv=active_split_kv, + softmax_stats_scale=softmax_stats_scale, ) prims.barrier_cluster_arrive_relaxed() diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/reduction.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/reduction.py index 5de512dea6d0..8306baaa66d1 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/reduction.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/reduction.py @@ -27,6 +27,7 @@ from ..helpers.constants import SPLIT_REDUCTION_SCALE_BARRIER_ID from ..helpers.math import ceil_div from ..helpers.mask import MaskType, mask_visible_k_length +from ..helpers.softmax_stats import store_reduced_softmax_stats from ..helpers.ops import ( fmax_f32, warp_reduce_max_f32, @@ -54,6 +55,8 @@ def run_reduction_kernel( cfg, max_splits: cutlass.Constexpr[int], rows_per_cta: cutlass.Constexpr[int], + softmax_stats=None, + softmax_stats_scale=None, ): """Combine consecutive logical split rows in one reference-reducer CTA. @@ -217,6 +220,14 @@ def run_reduction_kernel( ) if lane_idx == 0 and query_is_valid: (lse.iterator.raw_ptr() + Int64(public_flat_query_row)).store(global_lse) + store_reduced_softmax_stats( + softmax_stats, + public_flat_query_row, + acc_lse, + acc_lse_tile, + local_split_kv, + softmax_stats_scale, + ) for i in cutlass.range_constexpr(lse_per_thread): split_kv_idx = lane_idx + i * cfg.threads_per_warp if cute.elem_less(split_kv_idx, local_split_kv): diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/resources.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/resources.py index 4e3d24d11fd6..a9b9c9f62496 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/resources.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/resources.py @@ -115,6 +115,7 @@ qkv_dtype, qkv_major_k_stride_bytes_for, ) +from ..helpers.softmax_stats import store_softmax_stats from ..helpers.ops import ( fp8_log2_quant_scale, fp8_quant_scale_rcp, @@ -2115,6 +2116,7 @@ class TmemSResource(HighThroughputMlaResource): smem_exchange: Any = None # SMEM array for cross-warp max exchange softmax_scale_log2: Any = None # softmax_scale * log2(e) page_offsets: Any = None + store_softmax_stats: cutlass.Constexpr[bool] = False cache_seqs: Any = None # per-batch valid K length cu_seqlens_q: Any = None # cumulative compact-Q offsets, or None for fixed Q split_kv: Any = None # per-work-tile split count @@ -2853,17 +2855,32 @@ def load_peer_state(): no_correction = Int32(not max_changed) fma_b = self.softmax_scale_log2 - fma_c = Float32(0) - safe_row_max_new * self.softmax_scale_log2 + if cutlass.const_expr(self.store_softmax_stats): + fma_c = Float32(0) + else: + fma_c = Float32(0) - safe_row_max_new * self.softmax_scale_log2 if cutlass.const_expr(cfg.is_fp8_qkv()): # Use the same selected FP8 probability scale as output and LSE. # Native FP8 uses exact maxima and scale 448. fma_c = fma_c + fp8_log2_quant_scale() for i in cutlass.range_constexpr(0, 64, 2): - fma_result = fma_packed_f32x2( - (qk_acc_regs[i], qk_acc_regs[i + 1]), - (fma_b, fma_b), - (fma_c, fma_c), - ) + if cutlass.const_expr(self.store_softmax_stats): + # Fusing absolute scaled scores with a rounded scaled maximum + # can make exp(score - max) differ from one when score == max. + fma_result = fma_packed_f32x2( + ( + qk_acc_regs[i] - safe_row_max_new, + qk_acc_regs[i + 1] - safe_row_max_new, + ), + (fma_b, fma_b), + (fma_c, fma_c), + ) + else: + fma_result = fma_packed_f32x2( + (qk_acc_regs[i], qk_acc_regs[i + 1]), + (fma_b, fma_b), + (fma_c, fma_c), + ) qk_acc_regs[i] = cute.math.exp2(fma_result[0], fastmath=True) qk_acc_regs[i + 1] = cute.math.exp2(fma_result[1], fastmath=True) @@ -3633,6 +3650,8 @@ class GmemOResource(HighThroughputMlaResource): partial_output: Any = None lse: Any = None partial_lse: Any = None + softmax_stats: Any = None + softmax_stats_scale: Any = None tmem_o_ref: Any = None # Reference to TmemOResource for tmem_base_addr tmem_corr_ref: Any = None # Reference to TmemCorrResource for correction data output_scale: Any = None @@ -3708,6 +3727,22 @@ def _normalization_scale(self, row_sum, row_max, local_tidx, blk_coord): ) return scale + @cute.jit + def _store_softmax_stats(self, lse_ptr, row_max, row_sum): + if cutlass.const_expr(self.softmax_stats is not None): + maximum = row_max + if cutlass.const_expr(self.partial_lse is not None): + partial_rows = Int64(cute.cosize(self.partial_lse.layout)) + (lse_ptr + partial_rows).store(maximum) + (lse_ptr + partial_rows + partial_rows).store(row_sum) + else: + row = ( + lse_ptr.toint(Int64) - self.lse.iterator.raw_ptr().toint(Int64) + ) // Int64(4) + store_softmax_stats( + self.softmax_stats, row, maximum, row_sum, self.softmax_stats_scale + ) + @cute.jit def _query_row_state(self, row_in_tile, query_tile_idx, batch_idx): """Map one physical flat-tile row to public storage.""" @@ -4113,6 +4148,7 @@ def epilogue_store(self, stage_info: StageInfo) -> None: * Int64(S_q) ) lse_base_ptr.store(lse) + self._store_softmax_stats(lse_base_ptr, row_max, lse_row_sum) elif cutlass.const_expr(self.lse is not None): if cutlass.const_expr(self.cu_seqlens_q is not None): lse_base_ptr = ( @@ -4131,6 +4167,7 @@ def epilogue_store(self, stage_info: StageInfo) -> None: + Int64(batch_idx) * Int64(logical_num_heads_q) * Int64(S_q) ) lse_base_ptr.store(lse) + self._store_softmax_stats(lse_base_ptr, row_max, lse_row_sum) prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) cute.arch.fence_view_async_tmem_load() @@ -4389,6 +4426,7 @@ def epilogue_store_slice( * Int64(S_q) ) lse_base_ptr.store(lse) + self._store_softmax_stats(lse_base_ptr, row_max, lse_row_sum) elif cutlass.const_expr(self.lse is not None): if cutlass.const_expr(self.cu_seqlens_q is not None): lse_base_ptr = ( @@ -4409,6 +4447,7 @@ def epilogue_store_slice( * Int64(S_q) ) lse_base_ptr.store(lse) + self._store_softmax_stats(lse_base_ptr, row_max, lse_row_sum) prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) cute.arch.fence_view_async_tmem_load() diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/config.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/config.py index 7657de52608d..330430d7d98b 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/config.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/config.py @@ -348,7 +348,9 @@ def cluster_reduction_rows_per_slice(self) -> int: def cluster_reduction_slices(self) -> int: return ceil(self.tile_size_q / self.cluster_reduction_rows_per_slice) - def cluster_reduction_smem_bytes_for(self, num_ctas_kv: int) -> int: + def cluster_reduction_smem_bytes_for( + self, num_ctas_kv: int, *, store_softmax_stats: bool = False + ) -> int: """Return multi-CTA KV cluster reduction SMEM footprint.""" if num_ctas_kv <= 1: return 0 @@ -357,7 +359,9 @@ def cluster_reduction_smem_bytes_for(self, num_ctas_kv: int) -> int: num_slices_per_cta = ceil(num_slices / num_ctas_kv) num_rows_per_cta = num_slices_per_cta * rows_per_slice num_bytes_per_row_o = self.head_dim_per_cta_v * self.partial_o_dtype_bytes - num_bytes_per_row_stats = 2 * self.acc_dtype_bytes + num_bytes_per_row_stats = ( + 3 if store_softmax_stats else 2 + ) * self.acc_dtype_bytes return ( num_ctas_kv * num_rows_per_cta @@ -455,6 +459,7 @@ def compute_workspace_size( cfg: MlaConfig, partial_o_dtype, lse_dtype, + store_softmax_stats: bool = False, ) -> int: """Return the exact 1CTA split-KV GMEM workspace size in bytes. @@ -468,7 +473,9 @@ def compute_workspace_size( return 0 partial_rows = cfg.batch_size * cfg.seq_len_q * cfg.num_heads_q * split_kv return partial_rows * ( - cfg.head_dim_v * partial_o_dtype.width // 8 + lse_dtype.width // 8 + cfg.head_dim_v * partial_o_dtype.width // 8 + + lse_dtype.width // 8 + + (8 if store_softmax_stats else 0) ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/kernel.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/kernel.py index 130297f3b4a9..d515cb596f38 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/kernel.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/kernel.py @@ -480,6 +480,8 @@ def build_throughput_latency_mla_task_manager( lse_tensor=None, acc_o_tensor=None, acc_lse_tensor=None, + softmax_stats=None, + softmax_stats_scale=None, tile_sched_params=None, clc_response_ptr=None, use_clc_dynamic_scheduler: bool = False, @@ -532,6 +534,8 @@ def build_throughput_latency_mla_task_manager( lse_tensor=lse_tensor, acc_o_tensor=acc_o_tensor, acc_lse_tensor=acc_lse_tensor, + softmax_stats=softmax_stats, + softmax_stats_scale=softmax_stats_scale, tile_sched_params=tile_sched_params, clc_response_ptr=clc_response_ptr, use_clc_dynamic_scheduler=use_clc_dynamic_scheduler, @@ -772,6 +776,7 @@ def build_throughput_latency_mla_task_manager( ) smem_p0 = SmemPResource( cfg=cfg, + store_softmax_stats=softmax_stats is not None, pipeline_config=smem_p0_cfg, inst_id=0, scale_softmax_log2=scale_softmax_log2, @@ -781,6 +786,7 @@ def build_throughput_latency_mla_task_manager( ) smem_p1 = SmemPResource( cfg=cfg, + store_softmax_stats=softmax_stats is not None, pipeline_config=smem_p1_cfg, inst_id=1, scale_softmax_log2=scale_softmax_log2, @@ -828,6 +834,8 @@ def build_throughput_latency_mla_task_manager( lse_tensor=lse_tensor, acc_o_tensor=acc_o_tensor, acc_lse_tensor=acc_lse_tensor, + softmax_stats=softmax_stats, + softmax_stats_scale=softmax_stats_scale, cache_seqs=cache_seqs, cu_seqlens_q=cu_seqlens_q, head_idx=head_idx, @@ -849,6 +857,8 @@ def build_throughput_latency_mla_task_manager( lse_tensor=lse_tensor, acc_o_tensor=acc_o_tensor, acc_lse_tensor=acc_lse_tensor, + softmax_stats=softmax_stats, + softmax_stats_scale=softmax_stats_scale, cache_seqs=cache_seqs, cu_seqlens_q=cu_seqlens_q, head_idx=head_idx, @@ -1171,6 +1181,8 @@ def _make_single_kv_pipe_task_graph( lse_tensor=None, acc_o_tensor=None, acc_lse_tensor=None, + softmax_stats=None, + softmax_stats_scale=None, tile_sched_params=None, clc_response_ptr=None, use_clc_dynamic_scheduler: bool = False, @@ -1375,6 +1387,7 @@ def _make_single_kv_pipe_task_graph( if cfg.one_insts_kv_swap: p_resource = SmemPResource( cfg=cfg, + store_softmax_stats=softmax_stats is not None, pipeline_config=p_pipeline_cfg, inst_id=0, scale_softmax_log2=scale_softmax_log2, @@ -1384,6 +1397,7 @@ def _make_single_kv_pipe_task_graph( else: p_resource = TmemPResource( cfg=cfg, + store_softmax_stats=softmax_stats is not None, pipeline_config=p_pipeline_cfg, inst_id=0, scale_softmax_log2=scale_softmax_log2, @@ -1431,6 +1445,8 @@ def _make_single_kv_pipe_task_graph( lse_tensor=lse_tensor, acc_o_tensor=acc_o_tensor, acc_lse_tensor=acc_lse_tensor, + softmax_stats=softmax_stats, + softmax_stats_scale=softmax_stats_scale, cache_seqs=cache_seqs, cu_seqlens_q=cu_seqlens_q, head_idx=head_idx, @@ -1935,6 +1951,7 @@ def __call__( output_scale: cutlass.Float32, stream: object, output_quant_args=None, + softmax_stats: cute.Tensor | None = None, ): """Execute throughput-latency 1CTA MLA with fixed or compact ragged Q.""" cfg = self._make_config() @@ -2173,6 +2190,7 @@ def __call__( ) softmax_scale_log2 = softmax_scale * LOG2_E + softmax_stats_scale = softmax_scale if softmax_stats is not None else None use_gmem_reduction = cutlass.const_expr( cfg.use_multi_ctas_kv == 1 and cfg.use_cluster_reduction != 1 ) @@ -2252,6 +2270,8 @@ def __call__( tile_sched_params, block_split_kvs, output_quant_args, + softmax_stats, + softmax_stats_scale, ).launch( grid=grid, block=[cfg.threads_per_cta, 1, 1], @@ -2292,6 +2312,8 @@ def __call__( cu_seqlens_q, atten_sinks, output_quant_args, + softmax_stats, + softmax_stats_scale, ) if cutlass.const_expr(topology.cluster_size == 1): parallel_reducer.launch( @@ -2336,6 +2358,8 @@ def __call__( cache_seqs, cu_seqlens_q, cutlass.Int32(reduction_ctas), + softmax_stats, + softmax_stats_scale, ).launch( grid=reduction_grid, block=[reduction_threads, 1, 1], @@ -2366,6 +2390,8 @@ def dense_kernel( tile_sched_params: object, sparse_inputs: object, output_quant_args=None, + softmax_stats: cute.Tensor | None = None, + softmax_stats_scale=None, ): """Execute one flat-Q, batch, KV-split, and V head-dimension tile.""" cfg = self._make_config() @@ -2509,6 +2535,8 @@ def dense_kernel( lse_tensor=lse, acc_o_tensor=acc_o, acc_lse_tensor=acc_lse, + softmax_stats=softmax_stats, + softmax_stats_scale=softmax_stats_scale, tile_sched_params=tile_sched_params, clc_response_ptr=clc_response_ptr, use_clc_dynamic_scheduler=use_clc_dynamic, @@ -2631,6 +2659,8 @@ def gmem_reduction_kernel( cache_seqs: cute.Tensor, cu_seqlens_q: cute.Tensor, num_reduction_ctas: cutlass.Int32, + softmax_stats: cute.Tensor | None = None, + softmax_stats_scale=None, ): """Dispatch the throughput-latency split-KV reduction body.""" cfg = self._make_config() @@ -2644,6 +2674,8 @@ def gmem_reduction_kernel( cu_seqlens_q, cfg, num_reduction_ctas, + softmax_stats, + softmax_stats_scale, ) @cute.kernel @@ -2657,6 +2689,8 @@ def parallel_gmem_reduction_kernel( cu_seqlens_q: cute.Tensor, atten_sinks: cute.Tensor | None = None, output_quant_args=None, + softmax_stats: cute.Tensor | None = None, + softmax_stats_scale=None, ): """Dispatch the automatically selected parallel standalone reducer.""" @@ -2678,4 +2712,6 @@ def parallel_gmem_reduction_kernel( cfg.fuse_sparse_reduction, self.output_quant, output_quant_args, + softmax_stats, + softmax_stats_scale, ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/parallel_reduction.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/parallel_reduction.py index 11a40b52bfd1..7b55be63c99b 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/parallel_reduction.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/parallel_reduction.py @@ -32,6 +32,7 @@ warp_reduce_sum_f32, ) from ..helpers.query import flat_query_row_state, public_query_flat_row +from ..helpers.softmax_stats import store_reduced_softmax_stats from ..parallel_reduction_topology import ParallelReductionTopology from .config import MlaConfig @@ -205,6 +206,8 @@ def _run_parallel_gmem_reduction_g1_shared_stats( finalize_attention: cutlass.Constexpr[bool] = False, output_quant: cutlass.Constexpr = None, output_quant_args=None, + softmax_stats=None, + softmax_stats_scale=None, ): """Reduce G1 partials with the compact row-shared schedule.""" @@ -305,6 +308,14 @@ def _run_parallel_gmem_reduction_g1_shared_stats( (lse.iterator.raw_ptr() + stats_output_query_row).store( global_lse * Float32(LN_2) if finalize_attention else global_lse ) + store_reduced_softmax_stats( + softmax_stats, + stats_output_query_row, + acc_lse, + row_lse, + cfg.num_ctas_per_seq_kv, + softmax_stats_scale, + ) for lane_slot_i in cutlass.range_constexpr(lse_per_lane): stats_split_idx = lane_idx + Int32( @@ -427,6 +438,8 @@ def _run_parallel_gmem_reduction_shared_stats( slots_per_rank: cutlass.Constexpr[int], actual_splits: cutlass.Constexpr[int], elements_per_slice: cutlass.Constexpr[int], + softmax_stats=None, + softmax_stats_scale=None, ): """Reduce split-KV with one cooperative statistics warp per output row. @@ -763,6 +776,14 @@ def _run_parallel_gmem_reduction_shared_stats( (lse.iterator.raw_ptr() + output_query_row).store( smem_global_lse[row_in_slice] ) + store_reduced_softmax_stats( + softmax_stats, + output_query_row, + acc_lse, + acc_lse[head_idx, None, q_idx, batch_idx], + cfg.num_ctas_per_seq_kv, + softmax_stats_scale, + ) out_elem_offset = _output_element_offset(cfg, output_query_row, dim_idx) if cutlass.const_expr(output_elements_per_thread == 1): (output.iterator.raw_ptr() + out_elem_offset).store( @@ -806,6 +827,8 @@ def run_parallel_gmem_reduction_kernel( finalize_attention: cutlass.Constexpr[bool] = False, output_quant: cutlass.Constexpr = None, output_quant_args=None, + softmax_stats=None, + softmax_stats_scale=None, ): """Reduce split-KV partials with row-shared local and peer statistics.""" @@ -827,6 +850,8 @@ def run_parallel_gmem_reduction_kernel( finalize_attention, output_quant, output_quant_args, + softmax_stats, + softmax_stats_scale, ) return @@ -842,4 +867,6 @@ def run_parallel_gmem_reduction_kernel( slots_per_rank, actual_splits, elements_per_slice, + softmax_stats, + softmax_stats_scale, ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/reduction.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/reduction.py index c5a694dbdba1..b5967d5248e0 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/reduction.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/reduction.py @@ -33,6 +33,7 @@ split_o_element_offset, ) from ..helpers.tile import runtime_seq_len_kv_for_q +from ..helpers.softmax_stats import store_reduced_softmax_stats # Baseline split-KV reduction uses one 128-thread CTA. Each thread owns eight @@ -235,6 +236,8 @@ def run_gmem_reduction_kernel( cu_seqlens_q, cfg, num_reduction_ctas, + softmax_stats=None, + softmax_stats_scale=None, ): """Reduce split-KV partial O/LSE rows written by the 1CTA main kernel. @@ -342,6 +345,15 @@ def run_gmem_reduction_kernel( cu_seqlens_q, ) (lse.iterator.raw_ptr() + output_query_row).store(global_lse) + if head_dim_cta_idx == Int32(0): + store_reduced_softmax_stats( + softmax_stats, + output_query_row, + acc_lse, + row_lse, + cfg.num_ctas_per_seq_kv, + softmax_stats_scale, + ) acc_vec = vector_from_scalars( ( @@ -487,6 +499,14 @@ def run_gmem_reduction_kernel( cu_seqlens_q, ) (lse.iterator.raw_ptr() + output_query_row).store(global_lse) + store_reduced_softmax_stats( + softmax_stats, + output_query_row, + acc_lse, + row_lse, + cfg.num_ctas_per_seq_kv, + softmax_stats_scale, + ) for i in cutlass.range_constexpr(lse_per_thread): split_idx = lane_idx + Int32(i * GMEM_REDUCTION_WARP_LANES) if cute.elem_less(split_idx, Int32(cfg.num_ctas_per_seq_kv)): @@ -654,6 +674,14 @@ def run_gmem_reduction_kernel( cu_seqlens_q, ) (lse.iterator.raw_ptr() + output_query_row).store(global_lse) + store_reduced_softmax_stats( + softmax_stats, + output_query_row, + acc_lse, + row_lse, + cfg.num_ctas_per_seq_kv, + softmax_stats_scale, + ) for i in cutlass.range_constexpr(lse_per_thread): split_idx = lane_idx + Int32(i * GMEM_REDUCTION_WARP_LANES) if cute.elem_less(split_idx, Int32(cfg.num_ctas_per_seq_kv)): @@ -811,6 +839,14 @@ def run_gmem_reduction_kernel( cu_seqlens_q, ) (lse.iterator.raw_ptr() + output_query_row).store(global_lse) + store_reduced_softmax_stats( + softmax_stats, + output_query_row, + acc_lse, + g_lse, + cfg.num_ctas_per_seq_kv, + softmax_stats_scale, + ) for i in cutlass.range_constexpr(lse_per_thread): split_idx = lane_idx + Int32(i * GMEM_REDUCTION_WARP_LANES) if cute.elem_less(split_idx, Int32(cfg.num_ctas_per_seq_kv)): diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_p.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_p.py index 14fd5758621b..c02df54a683f 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_p.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_p.py @@ -81,6 +81,7 @@ class SmemPResource(MlaResource): """ inst_id: cutlass.Constexpr[int] = 0 + store_softmax_stats: cutlass.Constexpr[bool] = False scale_softmax_log2: Float32 = None tmem_s_ref: Optional[MemoryResource] = None order_p01_alloc: cutlass.Constexpr[Optional[SmemAllocation]] = None @@ -270,11 +271,21 @@ def materialize_p( scale_idx = scale_base + (s_idx % 2) p_val = Float32(0.0) if new_max_arr[scale_idx] != neg_max_f32(): - p_val = cute.math.exp2( - s_arr[s_idx] * self.scale_softmax_log2 - + neg_scaled_max[scale_idx], - fastmath=True, - ) + if cutlass.const_expr(self.store_softmax_stats): + # Subtract raw maxima before scaling: a fused + # multiply/add of large absolutes loses the origin. + p_val = cute.math.exp2( + (s_arr[s_idx] - new_max_arr[scale_idx]) + * self.scale_softmax_log2 + + fp8_log2_quant_scale(), + fastmath=True, + ) + else: + p_val = cute.math.exp2( + s_arr[s_idx] * self.scale_softmax_log2 + + neg_scaled_max[scale_idx], + fastmath=True, + ) p_vals[elem_idx] = p_val local_sums[scale_idx] += p_val regs_p[packed_idx] = pack_float4_to_fp8_e4m3( @@ -293,15 +304,29 @@ def materialize_p( new_max0 = new_max_arr[scale0] new_max1 = new_max_arr[scale1] if new_max0 != neg_max_f32(): - p0 = cute.math.exp2( - s_arr[s0] * self.scale_softmax_log2 + neg_scaled_max[scale0], - fastmath=True, - ) + if cutlass.const_expr(self.store_softmax_stats): + p0 = cute.math.exp2( + (s_arr[s0] - new_max0) * self.scale_softmax_log2, + fastmath=True, + ) + else: + p0 = cute.math.exp2( + s_arr[s0] * self.scale_softmax_log2 + + neg_scaled_max[scale0], + fastmath=True, + ) if new_max1 != neg_max_f32(): - p1 = cute.math.exp2( - s_arr[s1] * self.scale_softmax_log2 + neg_scaled_max[scale1], - fastmath=True, - ) + if cutlass.const_expr(self.store_softmax_stats): + p1 = cute.math.exp2( + (s_arr[s1] - new_max1) * self.scale_softmax_log2, + fastmath=True, + ) + else: + p1 = cute.math.exp2( + s_arr[s1] * self.scale_softmax_log2 + + neg_scaled_max[scale1], + fastmath=True, + ) local_sums[scale0] += p0 local_sums[scale1] += p1 regs_p[pair_idx] = pack_float2_to_bf16(p0, p1) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_corr.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_corr.py index cf547165e9e1..68f66289de8b 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_corr.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_corr.py @@ -97,6 +97,7 @@ from .common import ( MlaResource, ) +from ...helpers.softmax_stats import partial_max_tensor, store_softmax_stats # ===================================================================== # TmemCorrResource — Correction and output store to GMEM @@ -123,6 +124,8 @@ class TmemCorrResource(MlaResource): lse_tensor: object = None acc_o_tensor: object = None acc_lse_tensor: object = None + softmax_stats: object = None + softmax_stats_scale: object = None cache_seqs: object = None head_idx: object = None batch_idx: object = None @@ -168,9 +171,15 @@ def get_smem_requirements(self): allocs = [self._alloc, self._sum_alloc] if self.cfg.cluster_reduction_smem_bytes and self.acc_o_tensor is None: if self._cluster_reduction_alloc is None: + cluster_bytes = max( + self.cfg.cluster_reduction_smem_bytes_for( + splits, store_softmax_stats=self.softmax_stats is not None + ) + for splits in range(2, self.cfg.num_ctas_per_seq_kv + 1) + ) self._cluster_reduction_alloc = SmemAllocation( name=f"{self.name}_clusterReduction", - size_bytes=self.cfg.cluster_reduction_smem_bytes, + size_bytes=cluster_bytes, alignment=64, ) if self._cluster_reduction_barrier_alloc is None: @@ -211,7 +220,7 @@ def _init_store_state_from_context(self, context) -> None: context, self._cluster_reduction_alloc, Int8, - self.cfg.cluster_reduction_smem_bytes, + self._cluster_reduction_alloc.size_bytes, ) self._cluster_reduction_barrier = smem_array( context, @@ -278,7 +287,8 @@ def _init_cluster_reduction_barrier(self): cute.math.min(Int32(rows_per_cta), remaining_rows), ) bytes_per_row = Int32( - cfg.head_dim_per_cta_v * cfg.partial_o_dtype_bytes + cfg.acc_dtype_bytes + cfg.head_dim_per_cta_v * cfg.partial_o_dtype_bytes + + cfg.acc_dtype_bytes * (3 if self.softmax_stats is not None else 1) ) expected_bytes = valid_rows * Int32(cfg.num_ctas_per_seq_kv) * bytes_per_row if warp_idx == Int32(cfg.correction_warp_idx): @@ -368,14 +378,15 @@ def _store_partial_o_to_cluster_smem( ) @cute.jit - def _store_partial_lse_to_cluster_smem(self, local_row_idx, cta_idx_kv, lse_val): + def _store_partial_lse_to_cluster_smem( + self, local_row_idx, cta_idx_kv, lse_val, raw_maximum=None, row_sum=None + ): """Send one partial log-sum-exp value to the matching reduction owner.""" rows_per_cta = self._cluster_rows_per_cta() owner_rank = local_row_idx // rows_per_cta owner_row_idx = local_row_idx - owner_rank * rows_per_cta - stats_elem_offset = cta_idx_kv * rows_per_cta * Int32( - 2 - ) + owner_row_idx * Int32(2) + stats_width = Int32(3 if self.softmax_stats is not None else 2) + stats_elem_offset = (cta_idx_kv * rows_per_cta + owner_row_idx) * stats_width byte_offset = self._cluster_o_bytes() + Int64( stats_elem_offset * Int32(self.cfg.acc_dtype_bytes) ) @@ -390,6 +401,31 @@ def _store_partial_lse_to_cluster_smem(self, local_row_idx, cta_idx_kv, lse_val) remote_barrier.ir_value(), ], ) + if cutlass.const_expr(self.softmax_stats is not None): + if cutlass.const_expr(raw_maximum is None): + raw_maximum = Float32(-Float32.inf) + if cutlass.const_expr(row_sum is None): + row_sum = Float32(0.0) + max_ptr = remote_ptr + Int32(1) + cute.arch.inline_ptx( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 " + "[{$r0}], {$r1}, [{$r2}];", + read_only_args=[ + max_ptr.ir_value(), + raw_maximum.bitcast(Int32), + remote_barrier.ir_value(), + ], + ) + sum_ptr = remote_ptr + Int32(2) + cute.arch.inline_ptx( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.b32 " + "[{$r0}], {$r1}, [{$r2}];", + read_only_args=[ + sum_ptr.ir_value(), + row_sum.bitcast(Int32), + remote_barrier.ir_value(), + ], + ) @cute.jit def _store_o_slice_to_gmem( @@ -575,7 +611,11 @@ def _store_lse_to_gmem( ) if cutlass.const_expr(self.cfg.use_cluster_reduction == 1): self._store_partial_lse_to_cluster_smem( - local_row_idx, cta_idx_kv, lse_val + local_row_idx, + cta_idx_kv, + lse_val, + final_max[scale_idx], + lse_sum, ) elif cutlass.const_expr(self.acc_lse_tensor is not None): if valid_output_row and cta_idx_head_dim_v == Int32(0): @@ -592,6 +632,18 @@ def _store_lse_to_gmem( (self.acc_lse_tensor.iterator.raw_ptr() + elem_offset).store( lse_val ) + if cutlass.const_expr(self.softmax_stats is not None): + acc_max = partial_max_tensor( + self.acc_lse_tensor, self.softmax_stats + ) + (acc_max.iterator.raw_ptr() + elem_offset).store( + final_max[scale_idx] + ) + ( + acc_max.iterator.raw_ptr() + + Int64(cute.cosize(self.acc_lse_tensor.layout)) + + elem_offset + ).store(lse_sum) else: if valid_output_row: elem_offset = public_query_flat_row( @@ -603,13 +655,20 @@ def _store_lse_to_gmem( (self.lse_tensor.iterator.raw_ptr() + elem_offset).store( lse_val ) + if cta_idx_head_dim_v == Int32(0): + store_softmax_stats( + self.softmax_stats, + elem_offset, + final_max[scale_idx], + lse_sum, + self.softmax_stats_scale, + ) @cute.jit def _cluster_lse_byte_offset(self, split_idx, owner_row_idx): rows_per_cta = self._cluster_rows_per_cta() - stats_elem_offset = split_idx * rows_per_cta * Int32(2) + owner_row_idx * Int32( - 2 - ) + stats_width = Int32(3 if self.softmax_stats is not None else 2) + stats_elem_offset = (split_idx * rows_per_cta + owner_row_idx) * stats_width return self._cluster_o_bytes() + Int64( stats_elem_offset * Int32(self.cfg.acc_dtype_bytes) ) @@ -874,6 +933,65 @@ def _reduce_cluster_partials_and_store_for_tile( ( self.lse_tensor.iterator.raw_ptr() + lse_offset ).store(published_lse) + if cutlass.const_expr( + self.softmax_stats is not None + ): + true_max = Float32(-Float32.inf) + for split_idx in cutlass.range_constexpr( + cfg.num_ctas_per_seq_kv + ): + if local_lse[split_idx] != Float32( + -Float32.inf + ): + max_ptr = self._cluster_local_smem_ptr( + self._cluster_lse_byte_offset( + Int32(split_idx), owner_row_idx + ) + + Int64(cfg.acc_dtype_bytes), + Int32, + ) + true_max = cute.math.max( + true_max, + max_ptr.load().bitcast(Float32), + ftz=True, + ) + denominator = Float32(0.0) + for split_idx in cutlass.range_constexpr( + cfg.num_ctas_per_seq_kv + ): + if local_lse[split_idx] != Float32( + -Float32.inf + ): + max_ptr = self._cluster_local_smem_ptr( + self._cluster_lse_byte_offset( + Int32(split_idx), owner_row_idx + ) + + Int64(cfg.acc_dtype_bytes), + Int32, + ) + partial_max = max_ptr.load().bitcast( + Float32 + ) + partial_sum = ( + (max_ptr + Int32(1)) + .load() + .bitcast(Float32) + ) + denominator += ( + partial_sum + * cute.math.exp2( + (partial_max - true_max) + * self.scale_softmax_log2, + fastmath=True, + ) + ) + store_softmax_stats( + self.softmax_stats, + lse_offset, + true_max, + denominator, + self.softmax_stats_scale, + ) acc_vec = vector_from_scalars( ( diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_p.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_p.py index e867ba4e7c90..cf9348747a25 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_p.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_p.py @@ -81,6 +81,7 @@ class TmemPResource(MlaResource): ("p_stage_idx", Int32, Int32(0), "Current TMEM P pipeline stage."), ) inst_id: cutlass.Constexpr[int] = 0 + store_softmax_stats: cutlass.Constexpr[bool] = False scale_softmax_log2: Float32 = None tmem_alias_ref: Optional[MlaResource] = None p_stage_idx: cutlass.Constexpr[TaskLocalVariable] = ( @@ -171,26 +172,51 @@ def materialize_p( p2 = Float32(0.0) p3 = Float32(0.0) if has_finite_max: - if cutlass.const_expr(cfg.num_tokens_per_page == 1): - scaled01 = ffma2( - (s_arr[s_base + 0], s_arr[s_base + 1]), - log2_scale_pair, - neg_scaled_pair, + if cutlass.const_expr(self.store_softmax_stats): + # Keep the probability origin exact before log2 scaling. + neg_raw_pair = (-new_max_arr[0], -new_max_arr[0]) + quant_pair = ( + fp8_log2_quant_scale(), + fp8_log2_quant_scale(), ) - else: scaled01 = fadd2( fmul2( - (s_arr[s_base + 0], s_arr[s_base + 1]), log2_scale_pair + fadd2( + (s_arr[s_base + 0], s_arr[s_base + 1]), + neg_raw_pair, + ), + log2_scale_pair, ), + quant_pair, + ) + scaled23 = fadd2( + fmul2( + fadd2( + (s_arr[s_base + 2], s_arr[s_base + 3]), + neg_raw_pair, + ), + log2_scale_pair, + ), + quant_pair, + ) + elif cutlass.const_expr(cfg.num_tokens_per_page == 1): + scaled01 = ffma2( + (s_arr[s_base + 0], s_arr[s_base + 1]), + log2_scale_pair, neg_scaled_pair, ) - if cutlass.const_expr(cfg.num_tokens_per_page == 1): scaled23 = ffma2( (s_arr[s_base + 2], s_arr[s_base + 3]), log2_scale_pair, neg_scaled_pair, ) else: + scaled01 = fadd2( + fmul2( + (s_arr[s_base + 0], s_arr[s_base + 1]), log2_scale_pair + ), + neg_scaled_pair, + ) scaled23 = fadd2( fmul2( (s_arr[s_base + 2], s_arr[s_base + 3]), log2_scale_pair @@ -222,12 +248,20 @@ def materialize_p( scale_idx = scale_base + (s_idx % 2) p_val = Float32(0.0) if new_max_arr[scale_idx] != neg_max_f32(): - p_val = cute.math.exp2( - s_arr[s_idx] * self.scale_softmax_log2 - + neg_scaled_max[scale_idx] - + fp8_log2_quant_scale(), - fastmath=True, - ) + if cutlass.const_expr(self.store_softmax_stats): + p_val = cute.math.exp2( + (s_arr[s_idx] - new_max_arr[scale_idx]) + * self.scale_softmax_log2 + + fp8_log2_quant_scale(), + fastmath=True, + ) + else: + p_val = cute.math.exp2( + s_arr[s_idx] * self.scale_softmax_log2 + + neg_scaled_max[scale_idx] + + fp8_log2_quant_scale(), + fastmath=True, + ) p_vals[elem_idx] = p_val local_sums[scale_idx] += p_val regs_p[packed_idx] = pack_float4_to_fp8_e4m3( @@ -242,7 +276,16 @@ def materialize_p( if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): new_max = new_max_arr[0] if new_max != neg_max_f32(): - if cutlass.const_expr(cfg.num_tokens_per_page == 1): + if cutlass.const_expr(self.store_softmax_stats): + p0 = cute.math.exp2( + (s_arr[s0] - new_max) * self.scale_softmax_log2, + fastmath=True, + ) + p1 = cute.math.exp2( + (s_arr[s1] - new_max) * self.scale_softmax_log2, + fastmath=True, + ) + elif cutlass.const_expr(cfg.num_tokens_per_page == 1): scaled = ffma2( (s_arr[s0], s_arr[s1]), (self.scale_softmax_log2, self.scale_softmax_log2), @@ -273,17 +316,29 @@ def materialize_p( new_max0 = new_max_arr[scale0] new_max1 = new_max_arr[scale1] if new_max0 != neg_max_f32(): - p0 = cute.math.exp2( - s_arr[s0] * self.scale_softmax_log2 - + neg_scaled_max[scale0], - fastmath=True, - ) + if cutlass.const_expr(self.store_softmax_stats): + p0 = cute.math.exp2( + (s_arr[s0] - new_max0) * self.scale_softmax_log2, + fastmath=True, + ) + else: + p0 = cute.math.exp2( + s_arr[s0] * self.scale_softmax_log2 + + neg_scaled_max[scale0], + fastmath=True, + ) if new_max1 != neg_max_f32(): - p1 = cute.math.exp2( - s_arr[s1] * self.scale_softmax_log2 - + neg_scaled_max[scale1], - fastmath=True, - ) + if cutlass.const_expr(self.store_softmax_stats): + p1 = cute.math.exp2( + (s_arr[s1] - new_max1) * self.scale_softmax_log2, + fastmath=True, + ) + else: + p1 = cute.math.exp2( + s_arr[s1] * self.scale_softmax_log2 + + neg_scaled_max[scale1], + fastmath=True, + ) local_sums[scale0] += p0 local_sums[scale1] += p1 regs_p[pair_idx] = pack_float2_to_bf16(p0, p1) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py b/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py index 7a305b1ec6f2..594d160b83aa 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py @@ -90,6 +90,7 @@ class _MLADecodeCompileSpec: has_kernel_workspace: bool split_kv: int kernel: Any = field(compare=False, hash=False, repr=False) + store_softmax_stats: bool = False @dataclass(frozen=True) @@ -131,6 +132,7 @@ class _MLADecodePlanState: compiled: Callable[..., object] policy: tuple[tuple[str, object], ...] split_kv: int + store_softmax_stats: bool = False @dataclass(frozen=True) @@ -689,6 +691,32 @@ def _validate_out( _validate_16byte_alignment(out, "out") +def _validate_mla_softmax_stats( + query: torch.Tensor, + softmax_stats: Optional[torch.Tensor], + store_softmax_stats: bool, + validate: bool, +) -> None: + """Always enforce plan/buffer agreement; validate storage without CUDA sync.""" + if (softmax_stats is not None) != store_softmax_stats: + raise ValueError( + "softmax_stats must be supplied exactly when store_softmax_stats=True" + ) + if validate and softmax_stats is not None: + if not isinstance(softmax_stats, torch.Tensor): + raise TypeError("softmax_stats must be a torch.Tensor") + if softmax_stats.dtype != torch.float32: + raise TypeError("softmax_stats must have dtype torch.float32") + if softmax_stats.device != query.device: + raise ValueError("softmax_stats must be on the query device") + if tuple(softmax_stats.shape) != (*query.shape[:-1], 2): + raise ValueError("softmax_stats must have shape [*query.shape[:-1], 2]") + if not softmax_stats.is_contiguous(): + raise ValueError("softmax_stats must be contiguous") + if softmax_stats.data_ptr() % 4: + raise ValueError("softmax_stats must be 4-byte aligned") + + def _kernel_dtype_name(dtype_key: str) -> str: names = { "bfloat16": "bf16", @@ -738,6 +766,7 @@ def _resolve_mla_decode_launch_spec( output_dtype_key: str, mask_type: str, seq_len_q: int = 1, + store_softmax_stats: bool = False, ): """Resolve and cache MLA policy/workspace without compiling.""" @@ -875,6 +904,7 @@ def _resolve_mla_decode_launch_spec( cfg=final_cfg, partial_o_dtype=cutlass.BFloat16, lse_dtype=cutlass.Float32, + store_softmax_stats=store_softmax_stats, ) separate_reducer_impl, reducer_cluster_size = _separate_reducer_provenance( kernel, @@ -959,6 +989,7 @@ def _resolve_mla_decode_launch_spec( split_kv=split_kv, partial_o_dtype=cutlass.BFloat16, lse_dtype=cutlass.Float32, + store_softmax_stats=store_softmax_stats, ) separate_reducer_impl, reducer_cluster_size = _separate_reducer_provenance( kernel, @@ -1024,6 +1055,7 @@ def _make_mla_decode_compile_spec( output_dtype_key: str, max_seq_len_q: int, packed_query: bool, + store_softmax_stats: bool = False, ) -> _MLADecodeCompileSpec: """Keep policy resolution plan-specific and JIT identity batch-free.""" @@ -1041,6 +1073,7 @@ def _make_mla_decode_compile_spec( has_kernel_workspace=launch_spec.kernel_workspace_bytes > 0, split_kv=launch_spec.split_kv, kernel=launch_spec.kernel, + store_softmax_stats=store_softmax_stats, ) @@ -1225,6 +1258,16 @@ def sparse_indices(capacity): sink, ) stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + stats_fake = None + if compile_spec.store_softmax_stats: + stats_rows = ( + runtime_total_q * num_heads + if packed_query + else batch_size * max_seq_len_q * num_heads + ) + stats_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (stats_rows, 2), stride_order=(1, 0), assumed_align=4 + ) output_quant_args = _fake_output_quant_args(getattr(kernel, "output_quant", None)) @@ -1248,7 +1291,8 @@ def compile_kernel(): cutlass.Float32(1.0), cutlass.Float32(1.0), stream_fake, - output_quant_args=output_quant_args, + output_quant_args, + stats_fake, options=_COMPILE_OPTIONS, ) @@ -1279,6 +1323,7 @@ def get_prims_ts_batch_mla_decode_workspace_size( out_dtype: torch.dtype = torch.bfloat16, mask_type: Literal["dense", "causal"] = "causal", device=None, + store_softmax_stats: bool = False, ) -> int: """Return caller-workspace bytes for one automatic MLA policy. @@ -1289,11 +1334,16 @@ def get_prims_ts_batch_mla_decode_workspace_size( packed-query launches; ``seq_len_q`` remains a backward-compatible fixed-Q alias. If neither is supplied, the bound is one. The returned byte count includes both split-KV - scratch and the internal FP32 LSE tensor. Allocate a contiguous + scratch and the internal FP32 LSE tensor. With ``store_softmax_stats=True``, + it also includes split maxima and sums to export actual maximum/denominator + statistics; the public statistics buffer is caller-owned and excluded. + Allocate a contiguous ``torch.int8`` or ``torch.uint8`` CUDA buffer; MLA does not require its contents to be initialized before first use. """ + if not isinstance(store_softmax_stats, bool): + raise TypeError("store_softmax_stats must be a bool") batch_size = _validate_positive_int(batch_size, "batch_size") num_heads = _validate_positive_int(num_heads, "num_heads") _validate_mla_dims(kv_lora_rank, qk_rope_head_dim) @@ -1330,6 +1380,7 @@ def get_prims_ts_batch_mla_decode_workspace_size( _dtype_key(out_dtype), mask_type, max_seq_len_q, + store_softmax_stats, ) return _make_mla_workspace_layout( spec.kernel_workspace_bytes, batch_size, num_heads, max_seq_len_q @@ -1439,6 +1490,7 @@ def _launch_mla_decode( split_kv: int, workspace: _MLAWorkspaceViews, compiled: Callable[..., object], + softmax_stats: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Form the dimension-first views and launch one compiled MLA kernel.""" @@ -1480,12 +1532,10 @@ def _launch_mla_decode( runtime.sparse_inputs, runtime.bmm1_scale, runtime.bmm2_scale, - # Passed positionally and also when None. The kernel was compiled with - # ``output_quant_args`` in its signature (``None`` for unquantized - # output), and a kernel reloaded from the persistent compile cache is - # a positional TVM-FFI entry with no Python defaults: it rejects both - # a missing parameter and a keyword argument. + # Reloaded persistent-cache kernels expose a positional TVM-FFI entry + # without Python defaults. Keep both optional arguments in the ABI. runtime.output_quant_args, + None if softmax_stats is None else softmax_stats.view(-1, 2), ) return runtime.out @@ -1515,6 +1565,7 @@ def plan( kv_data_type: torch.dtype, o_data_type: torch.dtype, mask_type: Literal["dense", "causal"] = "causal", + store_softmax_stats: bool = False, workspace_buffer: Optional[torch.Tensor] = None, validate: bool = True, ) -> None: @@ -1554,6 +1605,9 @@ def plan( Query, K/V, and output dtypes used to compile the plan. mask_type : {"dense", "causal"} Attention mask mode. + store_softmax_stats : bool + Compile final FP32 maximum/denominator export. Supply its caller-owned + buffer to every run when enabled. Defaults to ``False``. workspace_buffer : torch.Tensor, optional Caller-owned contiguous int8 or uint8 scratch on ``device``. It must be 32-byte aligned and large enough for the selected plan. @@ -1564,6 +1618,8 @@ def plan( Disable only for previously validated inputs and warmed topology. """ + if not isinstance(store_softmax_stats, bool): + raise TypeError("store_softmax_stats must be a bool") if not isinstance(validate, bool): raise TypeError("validate must be a bool") if validate: @@ -1603,6 +1659,7 @@ def plan( _dtype_key(o_data_type), mask_type, max_seq_len_q, + store_softmax_stats, ) spec = _resolve_mla_decode_launch_spec(*spec_key) compile_spec = _make_mla_decode_compile_spec( @@ -1616,6 +1673,7 @@ def plan( output_dtype_key=_dtype_key(o_data_type), max_seq_len_q=max_seq_len_q, packed_query=packed_query, + store_softmax_stats=store_softmax_stats, ) policy = spec.policy workspace_layout = _make_mla_workspace_layout( @@ -1657,6 +1715,7 @@ def plan( compiled=compiled, policy=policy, split_kv=int(dict(policy)["split_kv"]), + store_softmax_stats=store_softmax_stats, ) @flashinfer_experimental_api @@ -1671,6 +1730,7 @@ def run( bmm1_scale: float = 1.0, bmm2_scale: float = 1.0, out: Optional[torch.Tensor] = None, + softmax_stats: Optional[torch.Tensor] = None, validate: bool = True, ) -> torch.Tensor: """Launch the most recently planned MLA decode on the current stream. @@ -1703,6 +1763,13 @@ def run( QK and value/output scaling factors. out : torch.Tensor, optional Caller-owned output tensor. A new tensor is allocated when omitted. + softmax_stats : torch.Tensor, optional + Contiguous FP32 buffer ``[*query.shape[:-1], 2]`` on the query device. + The last dimension stores the actual maximum scaled logit in natural-log + units and ``sum(exp(logit - max))``, without internal FP8 scaling. + Required exactly when the plan has ``store_softmax_stats=True``, even + with ``validate=False``. Must not alias inputs, output, or workspace; + overlap is a caller precondition and is not checked. validate : bool Enable explicit runtime validation. Defaults to ``True``. @@ -1717,6 +1784,9 @@ def run( raise RuntimeError("plan() must be called before run()") if not isinstance(validate, bool): raise TypeError("validate must be a bool") + _validate_mla_softmax_stats( + query, softmax_stats, state.store_softmax_stats, validate + ) runtime_qo_indptr = qo_indptr if state.packed_query else None runtime = _prepare_mla_runtime( query, @@ -1761,6 +1831,7 @@ def run( split_kv=state.split_kv, workspace=state.workspace_views, compiled=state.compiled, + softmax_stats=softmax_stats, ) @@ -1783,6 +1854,8 @@ def batch_mla_decode_with_paged_kv_cache( out_dtype: torch.dtype = torch.bfloat16, workspace_buffer: Optional[torch.Tensor] = None, validate: bool = True, + store_softmax_stats: bool = False, + softmax_stats: Optional[torch.Tensor] = None, ) -> torch.Tensor: """One-shot convenience wrapper for fixed or packed-query MLA decode. @@ -1820,6 +1893,11 @@ def batch_mla_decode_with_paged_kv_cache( Caller-owned output tensor. out_dtype : torch.dtype Output dtype. + store_softmax_stats : bool + Compile statistics export; defaults to ``False``. + softmax_stats : torch.Tensor, optional + Caller-owned FP32 maximum/denominator buffer with shape + ``[*query.shape[:-1], 2]``; required exactly when statistics are enabled. workspace_buffer : torch.Tensor, optional Caller-owned byte scratch, exclusive to one in-flight launch/graph. @@ -1927,6 +2005,7 @@ def batch_mla_decode_with_paged_kv_cache( out_dtype=out_dtype, mask_type=mask_type, device=query.device, + store_softmax_stats=store_softmax_stats, ) workspace_buffer = torch.empty( workspace_bytes, dtype=torch.int8, device=query.device @@ -1949,6 +2028,7 @@ def batch_mla_decode_with_paged_kv_cache( mask_type=mask_type, workspace_buffer=workspace_buffer, validate=validate, + store_softmax_stats=store_softmax_stats, ) return wrapper.run( query, @@ -1960,6 +2040,7 @@ def batch_mla_decode_with_paged_kv_cache( bmm2_scale=bmm2_scale, out=out, validate=validate, + softmax_stats=softmax_stats, ) @@ -2960,6 +3041,7 @@ def _run_prepared(self, prepared, query, indices, lengths, extra_indices, extra_ bmm1, bmm2, None, + None, ) return out From 9ee50623afdf0ce2e0ddfa578f6f6dafcffac38f Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Wed, 28 Oct 2026 04:27:21 +0000 Subject: [PATCH 02/19] [None][feat] Expose shared MLA context preprocessing without FMHA dispatch Add mla_context_preprocess with caller-owned buffers. Reuse RoPE/cache append and FP8 QKV quantization, allow preprocessed cached/chunked inputs, and pack the strided V projection. This commit adds the prologue contract only; the PrimTS adapter follows separately. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp | 134 ++++++++++++++++++++++ 1 file changed, 134 insertions(+) diff --git a/cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp b/cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp index e9af3e5ca586..796b0a019373 100644 --- a/cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp +++ b/cpp/tensorrt_llm/thop/mlaPreprocessOp.cpp @@ -105,6 +105,132 @@ void mergeChunkedAttentionForMLAHelper(torch::Tensor& merged_attn, torch::Tensor } // namespace +// Shared dense-MLA context prologue, without dispatching an attention kernel. +// Cached/chunked callers pass no latent_cache: their Q/K RoPE and cache append +// have already run. Output buffers are owned and accounted for by the caller. +void mlaContextPreprocess(torch::Tensor q, torch::Tensor k, torch::Tensor const& v, + torch::optional latent_cache, torch::Tensor const& cu_q_seqlens, + torch::Tensor const& sequence_lengths, torch::Tensor const& cos_sin_cache, + torch::Tensor const& kv_cache_block_offsets, torch::Tensor const& host_kv_cache_pool_pointers, + torch::Tensor const& host_kv_cache_pool_mapping, torch::optional kv_scale_orig_quant, + torch::optional kv_scale_quant_orig, int64_t num_heads, int64_t max_input_seq_len, int64_t layer_idx, + int64_t tokens_per_block, int64_t attention_window_size, int64_t quant_mode, double sm_scale, torch::Tensor q_out, + torch::Tensor k_out, torch::Tensor v_out, torch::Tensor bmm1_scale, torch::Tensor bmm2_scale) +{ + TORCH_CHECK(q.scalar_type() == torch::kBFloat16 || q.scalar_type() == torch::kFloat16); + CHECK_INPUT(q, q.scalar_type()); + CHECK_INPUT(k, q.scalar_type()); + TORCH_CHECK(q.dim() == 2 && k.dim() == 2 && v.dim() == 2 && num_heads > 0); + TORCH_CHECK(q.size(1) == num_heads * 192 && k.size(1) == num_heads * 192); + TORCH_CHECK(v.device() == q.device() && k.device() == q.device()); + TORCH_CHECK(v.scalar_type() == q.scalar_type() && v.size(0) == k.size(0)); + TORCH_CHECK(v.size(1) == num_heads * 128 && v.stride(1) == 1); + CHECK_INPUT(cu_q_seqlens, torch::kInt32); + CHECK_INPUT(sequence_lengths, torch::kInt32); + CHECK_INPUT(cos_sin_cache, torch::kFloat32); + TORCH_CHECK(cu_q_seqlens.device() == q.device() && sequence_lengths.device() == q.device() + && cos_sin_cache.device() == q.device()); + TORCH_CHECK(cu_q_seqlens.dim() == 1 && cu_q_seqlens.numel() > 1); + auto const batch_size = cu_q_seqlens.numel() - 1; + TORCH_CHECK(sequence_lengths.numel() >= batch_size && max_input_seq_len > 0); + TORCH_CHECK(cos_sin_cache.numel() >= attention_window_size * 64); + auto const mode = tc::QuantMode(static_cast(quant_mode)); + TORCH_CHECK(!mode.hasInt8KvCache() && !mode.hasFp4KvCache()); + bool const fp8 = mode.hasFp8KvCache(); + auto const output_dtype = fp8 ? torch::kFloat8_e4m3fn : q.scalar_type(); + for (auto const& tensor : {q_out, k_out, v_out}) + { + CHECK_INPUT(tensor, output_dtype); + TORCH_CHECK(tensor.device() == q.device()); + } + TORCH_CHECK(q_out.sizes() == q.sizes() && k_out.sizes() == k.sizes() && v_out.sizes() == v.sizes()); + CHECK_INPUT(bmm1_scale, torch::kFloat32); + CHECK_INPUT(bmm2_scale, torch::kFloat32); + TORCH_CHECK(bmm1_scale.numel() == 2 && bmm2_scale.numel() == 1); + TORCH_CHECK(bmm1_scale.device() == q.device() && bmm2_scale.device() == q.device()); + if (fp8) + { + // The shared quantizer consumes the kv_b_proj split view of V. + TORCH_CHECK(v.stride(0) == num_heads * 256); + } + else + { + TORCH_CHECK(q_out.data_ptr() == q.data_ptr() && k_out.data_ptr() == k.data_ptr()); + } + for (auto const& scale : {kv_scale_orig_quant, kv_scale_quant_orig}) + { + if (scale.has_value()) + { + CHECK_INPUT(scale.value(), torch::kFloat32); + TORCH_CHECK(scale->device() == q.device() && scale->numel() >= 1); + } + } + if (latent_cache.has_value()) + { + CHECK_INPUT(latent_cache.value(), q.scalar_type()); + TORCH_CHECK(latent_cache->device() == q.device()); + TORCH_CHECK(latent_cache->numel() == q.size(0) * 576 && k.size(0) == q.size(0)); + } + auto cache + = buildPagedKvCacheBuffers(std::optional(kv_cache_block_offsets), std::optional(host_kv_cache_pool_pointers), + std::optional(host_kv_cache_pool_mapping), mode, layer_idx, batch_size, tokens_per_block, 1, 576, + attention_window_size, attention_window_size, 1, 0, true, q.element_size()) + .kvCacheBuffer; + auto stream = at::cuda::getCurrentCUDAStream(q.get_device()); + auto launch = [&](auto type_tag) + { + using T = decltype(type_tag); + tk::MlaParams params{}; + params.q_buf = static_cast(q.data_ptr()); + params.k_buf = static_cast(k.data_ptr()); + params.v_buf = static_cast(v.data_ptr()); + params.latent_cache = latent_cache ? static_cast(latent_cache->data_ptr()) : nullptr; + params.cu_q_seqlens = cu_q_seqlens.data_ptr(); + params.cache_seq_lens = sequence_lengths.data_ptr(); + params.cos_sin_cache = static_cast(cos_sin_cache.data_ptr()); + params.batch_size = batch_size; + params.acc_q_len = q.size(0); + params.head_num = num_heads; + params.max_input_seq_len = max_input_seq_len; + params.meta.kv_lora_rank = 512; + params.meta.qk_nope_head_dim = 128; + params.meta.qk_rope_head_dim = 64; + params.meta.v_head_dim = 128; + params.cache_type = tk::cacheTypeFromQuantMode(mode); + params.quant_scale_kv = kv_scale_orig_quant ? kv_scale_orig_quant->data_ptr() : nullptr; + if (latent_cache) + { + tk::invokeMLARopeContext(params, cache, stream); + } + if (fp8) + { + params.quant_q_buf = q_out.data_ptr(); + params.quant_k_buf = k_out.data_ptr(); + params.quant_v_buf = v_out.data_ptr(); + params.quant_scale_qkv = params.quant_scale_kv; + params.dequant_scale_q = kv_scale_quant_orig ? kv_scale_quant_orig->data_ptr() : nullptr; + params.dequant_scale_kv = params.dequant_scale_q; + params.host_bmm1_scale = sm_scale; + params.bmm1_scale = bmm1_scale.data_ptr(); + params.bmm2_scale = bmm2_scale.data_ptr(); + tk::invokeMLAContextFp8Quantize(params, k.size(0), stream); + } + }; + if (q.scalar_type() == torch::kBFloat16) + { + launch(__nv_bfloat16{}); + } + else + { + launch(half{}); + } + if (!fp8) + { + // Q/K alias the caller's inputs; only V's split projection needs packing. + v_out.copy_(v); + } +} + std::vector loadPagedKVCacheForMLA(torch::ScalarType out_dtype, int64_t const num_contexts, int64_t const num_ctx_cached_tokens, int64_t const max_ctx_cached_kv_len, torch::Tensor& cu_ctx_cached_kv_lens, torch::Tensor const& kv_cache_block_offsets, torch::Tensor const& host_kv_cache_pool_pointers, @@ -419,6 +545,13 @@ TRTLLM_NAMESPACE_END TORCH_LIBRARY_FRAGMENT(trtllm, m) { + m.def( + "mla_context_preprocess(Tensor(a!) q, Tensor(b!) k, Tensor v, Tensor? latent_cache, " + "Tensor cu_q_seqlens, Tensor sequence_lengths, Tensor cos_sin_cache, Tensor kv_cache_block_offsets, " + "Tensor host_kv_cache_pool_pointers, Tensor host_kv_cache_pool_mapping, Tensor? kv_scale_orig_quant, " + "Tensor? kv_scale_quant_orig, int num_heads, int max_input_seq_len, int layer_idx, int tokens_per_block, " + "int attention_window_size, int quant_mode, float sm_scale, Tensor(a!) q_out, Tensor(b!) k_out, " + "Tensor(c!) v_out, Tensor(d!) bmm1_scale, Tensor(e!) bmm2_scale) -> ()"); m.def( "load_paged_kv_cache_for_mla(" "ScalarType out_dtype" @@ -442,6 +575,7 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) TORCH_LIBRARY_IMPL(trtllm, CUDA, m) { + m.impl("mla_context_preprocess", &tensorrt_llm::torch_ext::mlaContextPreprocess); m.impl("load_paged_kv_cache_for_mla", &tensorrt_llm::torch_ext::loadPagedKVCacheForMLA); } From 27c68b9db5d823c37862e5ef4db4f2e157cdc98c Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Wed, 28 Oct 2026 04:28:11 +0000 Subject: [PATCH 03/19] [None][feat] Integrate PrimTS MLA context and cached-prefix merging Route fresh, cached and chunked MLA prefill through the planned PrimTS context wrapper, including FP8 scales and merge-compatible softmax statistics. Preserve full cached K/V spans and account for BF16 V-packing workspace. Fold the FMHA-field and statistics-API rebase adaptations into this implementation. Correct MLA generation to require query-only input (is_fused_qkv=False), matching PhasedFmha. Include selector/contract and workspace tests; Rubin enablement and performance scheduling are separate commits. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../_torch/attention/backends/fmha/phased.py | 9 +- .../attention/backends/fmha/prims_ts.py | 196 +++++++++++++++++- .../_torch/attention/backends/trtllm.py | 9 + .../_torch/attention/test_prims_ts_fmha.py | 123 ++++++++++- .../pyexecutor/test_mla_workspace_reserve.py | 14 ++ 5 files changed, 332 insertions(+), 19 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/phased.py b/tensorrt_llm/_torch/attention/backends/fmha/phased.py index 22ed69809be3..261a18abe013 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/phased.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/phased.py @@ -309,8 +309,13 @@ def context_params( phase_input = _leading_rows(q, num_ctx_tokens) params.qkv_input = phase_input if forward_args.is_fused_qkv else None params.query_input = None if forward_args.is_fused_qkv else phase_input - params.key_input = _leading_rows(k, num_ctx_tokens) if k is not None else None - params.value_input = _leading_rows(v, num_ctx_tokens) if v is not None else None + if self.attn.is_mla_enable: + # An expanded MLA prefix or KV chunk can be longer than its Q span. + params.key_input = k + params.value_input = v + else: + params.key_input = _leading_rows(k, num_ctx_tokens) if k is not None else None + params.value_input = _leading_rows(v, num_ctx_tokens) if v is not None else None params.output = out_tensor if fused_epilogue else _leading_rows(out_tensor, num_ctx_tokens) params.sequence_lengths = metadata.kv_lens_cuda_runtime params.context_lengths = metadata.prompt_lens_cuda_runtime diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py index 89ee7d2a8126..9caaf5bafcce 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py @@ -130,7 +130,7 @@ def get_paged_kv_policy_unsupported_reason( and not quant_mode.has_fp4_kv_cache() ) if quant_mode.has_kv_cache_quant() and not is_fp8_mla: - return "quantized KV cache is supported only for FP8 MLA decode." + return "quantized KV cache is supported only for FP8 MLA." return None @@ -153,7 +153,7 @@ def get_attention_feature_unsupported_reason( class PrimsTSFmha(PhasedFmha): - """Blackwell task-scheduled paged context and decode FMHA library.""" + """Blackwell task-scheduled standard/MLA context and decode FMHA library.""" supports_per_tensor_fp8 = True @@ -177,6 +177,9 @@ def __init__(self, attn: "TrtllmAttention") -> None: self._contiguous_context_wrappers: dict[tuple[Any, ...], "BatchPrefillTSWrapper"] = {} self._decode_wrappers: dict[int, "BatchDecodePagedTSWrapper"] = {} self._mla_decode_wrappers: dict[int, "BatchMLADecodePagedTSWrapper"] = {} + self._mla_context_wrappers: dict[ + tuple[int, int, int, str, bool], "BatchPrefillTSWrapper" + ] = {} # Dense MLA's quantization scales are fixed for this layer/model. PrimTS # takes host scalars, so read them once during eager warmup, not capture. self._mla_fp8_scales: tuple[float, float] | None = None @@ -293,7 +296,7 @@ def _is_supported_with_reason( if meta.kv_cache_manager is None: if k is None or v is None or fwd.is_fused_qkv: return False, "cache-free attention takes separate Q, K and V." - elif k is not None or v is not None or (not attn.is_mla_enable and not fwd.is_fused_qkv): + elif not attn.is_mla_enable and (k is not None or v is not None or not fwd.is_fused_qkv): return False, "only fused QKV input is supported." if meta.is_cross: return False, "cross attention is not supported." @@ -361,12 +364,12 @@ def _is_supported_with_reason( ) if attn.num_heads <= 0 or attn.num_kv_heads <= 0: return False, "query and KV head counts must be positive." - if is_mla: + if is_mla and has_generation: if attn.num_kv_heads != 1: return False, "MLA decode requires one logical KV head." if attn.num_heads > 128: return False, "MLA decode supports at most 128 local query heads." - else: + elif not is_mla: if attn.num_heads % attn.num_kv_heads != 0: return False, "the query head count must be divisible by the KV head count." if has_generation and attn.num_heads // attn.num_kv_heads > self.MAX_DECODE_GQA_RATIO: @@ -376,7 +379,7 @@ def _is_supported_with_reason( 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." + return False, "FP8 MLA requires an FP8 E4M3 KV cache." if not is_fp8_mla 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: @@ -384,9 +387,49 @@ def _is_supported_with_reason( if q.ndim != 2: return False, f"fused attention input must be rank 2, got rank {q.ndim}." - if is_mla: + if is_mla and has_context: + if input_type != AttentionInputType.context_only: + return False, "MLA context must be dispatched separately from generation." + if ( + attn.head_dim, + attn.qk_nope_head_dim, + attn.qk_rope_head_dim, + attn.v_head_dim, + attn.kv_lora_rank, + ) != (192, 128, 64, 128, 512): + return False, "MLA context requires QK=192, V=128 and latent=512." + if attn.num_kv_heads != attn.num_heads: + return False, "MLA context requires expanded K/V for every query head." + if k is None or v is None or fwd.is_fused_qkv: + return False, "MLA context requires separate Q/K/V." + if q.shape[1] != attn.num_heads * 192: + return False, "MLA context query has an incompatible extent." + if ( + k.ndim != 2 + or k.shape[1] != attn.num_heads * 192 + or v.ndim != 2 + or v.shape != (k.shape[0], attn.num_heads * 128) + ): + return False, "MLA context K/V have incompatible extents." + if ( + k.device != q.device + or v.device != q.device + or k.dtype != q.dtype + or v.dtype != q.dtype + or not k.is_contiguous() + ): + return False, "MLA context requires contiguous K and matching Q/K/V dtype/device." + if v.stride(1) != 1 or (is_fp8_mla and v.stride(0) != attn.num_heads * 256): + return False, "FP8 MLA context requires the strided kv_b_proj V view." + if output.numel() != q.shape[0] * attn.num_heads * 128: + return False, "MLA context output has an incompatible extent." + elif is_mla: if has_context or input_type != AttentionInputType.generation_only: return False, "MLA is supported only for generation-only requests." + # The latent and RoPE components form one query, not fused QKV. + # PhasedFmha routes this through query_input with is_fused_qkv=False. + if k is not None or v is not None or fwd.is_fused_qkv: + return False, "MLA decode requires query-only input and no separate K/V." if q.dtype != torch.bfloat16 or output.dtype != torch.bfloat16: return False, "MLA decode requires BF16 query input and output." if attn.kv_lora_rank != 512 or attn.qk_rope_head_dim != 64: @@ -435,7 +478,14 @@ def _is_supported_with_reason( if host_kv_lens is None or host_kv_lens.numel() < num_contexts + num_generations: return False, "host KV lengths are required for safe policy selection." active_kv_lens = host_kv_lens[: num_contexts + num_generations] - if active_kv_lens.numel() == 0 or int(active_kv_lens.min()) <= 0: + allows_empty_chunk = ( + is_mla + and has_context + and mask_type == AttentionMaskType.padding + and fwd.softmax_stats_tensor is not None + ) + min_kv_length = 0 if allows_empty_chunk else 1 + if active_kv_lens.numel() == 0 or int(active_kv_lens.min()) < min_kv_length: return False, "every active request must contain at least one KV token." max_kv_length = int(active_kv_lens.max()) max_seq_len = int(meta.max_seq_len) @@ -913,7 +963,7 @@ def prepare_workspace( forward_args: AttentionForwardArgs, workspace: torch.Tensor, ) -> None: - del k, v + del v block_offsets = metadata.kv_cache_block_offsets if block_offsets is None: raise RuntimeError("PrimTS requires paged KV-cache block offsets.") @@ -927,6 +977,11 @@ def prepare_workspace( self._multi_processor_count = get_multi_processor_count_for_device(q.device.index) required_preprocess_bytes = 0 + if has_context and self.attn.is_mla_enable: + assert k is not None + required_preprocess_bytes = self._mla_context_workspace_layout( + q, k, int(metadata.num_contexts) + )[1] if has_context and not self.attn.is_mla_enable: context_layout = thop.get_trtllm_gen_context_workspace_layout( q.dtype, @@ -1090,6 +1145,129 @@ def _standard_kv_views( kv_pool.narrow(0, kv_page_offset, usable_pages), ) + def _mla_context_workspace_layout( + self, q: torch.Tensor, k: torch.Tensor, batch_size: int + ) -> tuple[dict[str, tuple[int, int, torch.dtype, tuple[int, ...]]], int]: + """Pack context scratch into the shared, memory-accounted FMHA allocation.""" + fp8 = QuantMode(self.attn.quant_mode).has_fp8_kv_cache() + tensors = [ + ("kv_indptr", (batch_size + 1,), torch.int32), + ("bmm1_scale", (2,), torch.float32), + ("bmm2_scale", (1,), torch.float32), + ("v", (k.shape[0], self.attn.num_heads * 128), torch.float8_e4m3fn if fp8 else q.dtype), + ] + if fp8: + tensors.extend( + (name, tuple(tensor.shape), torch.float8_e4m3fn) + for name, tensor in (("q", q), ("k", k)) + ) + layout = {} + offset = 0 + for name, shape, dtype in tensors: + size = math.prod(shape) * dtype.itemsize + layout[name] = (offset, size, dtype, shape) + offset = pad_up(offset + size, _WORKSPACE_ALIGNMENT) + return layout, offset + + def run_mla_context(self, params: FmhaParams) -> None: + """Run expanded MLA, including cached-prefix and chunk-merge calls.""" + attn, meta, fwd = params.attn, params.meta, params.fwd + q, k, v = params.query_input, params.key_input, params.value_input + assert q is not None and k is not None and v is not None + assert params.output is not None + layout, _ = self._mla_context_workspace_layout(q, k, params.batch_size) + workspace = params.workspace.view(torch.uint8).view(-1) + buffers = { + name: workspace.narrow(0, offset, size).view(dtype).view(shape) + for name, (offset, size, dtype, shape) in layout.items() + } + qo_indptr = meta.mla_prepare_ctx_cu_seqlens() + assert qo_indptr is not None + kv_indptr = buffers["kv_indptr"] + kv_indptr[0].zero_() + # KV lengths change for each cached-prefill chunk, unlike the Q offsets. + torch.cumsum( + params.sequence_lengths[: params.batch_size], 0, dtype=torch.int32, out=kv_indptr[1:] + ) + attn._ensure_rope_table_size(meta.max_seq_len) + fp8 = QuantMode(attn.quant_mode).has_fp8_kv_cache() + q_processed, k_processed = (buffers["q"], buffers["k"]) if fp8 else (q, k) + v_processed = buffers["v"] + sm_scale = self._get_bmm1_scale(attn) + torch.ops.trtllm.mla_context_preprocess( + q, + k, + v, + fwd.latent_cache, + qo_indptr, + params.sequence_lengths[: params.batch_size], + attn.rotary_cos_sin, + meta.kv_cache_block_offsets, + meta.host_kv_cache_pool_pointers, + meta.host_kv_cache_pool_mapping, + fwd.kv_scale_orig_quant, + fwd.kv_scale_quant_orig, + attn.num_heads, + params.input_seq_length, + attn.get_local_layer_idx(meta), + params.tokens_per_block, + params.cyclic_attention_window_size, + attn.quant_mode, + sm_scale, + q_processed, + k_processed, + v_processed, + buffers["bmm1_scale"], + buffers["bmm2_scale"], + ) + mask = self._get_prims_mask_type(fwd) + stats = fwd.softmax_stats_tensor + # Dense cached chunks traverse the planned KV domain. Bucket current + # capacities instead of traversing max_seq_len for every small chunk. + q_capacity = 1 << (max(params.input_seq_length, 1) - 1).bit_length() + kv_capacity = 1 << (max(params.max_past_kv_length, 1) - 1).bit_length() + key = (params.batch_size, q_capacity, kv_capacity, mask, stats is not None) + wrapper = self._mla_context_wrappers.get(key) + if wrapper is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("PrimTS MLA context must be planned before CUDA graph capture.") + from tensorrt_llm._torch.attention.backends.prims_ts.context import ( + BatchPrefillTSWrapper, + ) + + wrapper = BatchPrefillTSWrapper() + wrapper.plan( + device=q.device, + batch_size=params.batch_size, + max_seq_len_q=q_capacity, + max_kv_len=kv_capacity, + num_qo_heads=attn.num_heads, + num_kv_heads=attn.num_heads, + head_dim=192, + head_dim_vo=128, + q_dtype=q_processed.dtype, + k_dtype=k_processed.dtype, + v_dtype=v_processed.dtype, + out_dtype=params.output.dtype, + packed=True, + mask_type=mask, + sm_scale=sm_scale, + store_softmax_stats=stats is not None, + ) + self._mla_context_wrappers[key] = wrapper + wrapper.run( + q_processed.view(-1, attn.num_heads, 192), + k_processed.view(-1, attn.num_heads, 192), + v_processed.view(-1, attn.num_heads, 128), + qo_indptr, + kv_indptr, + out=params.output, + softmax_stats=stats, + scale_softmax_log2=buffers["bmm1_scale"][1:] if fp8 else None, + output_scale=buffers["bmm2_scale"] if fp8 else None, + validate=False, + ) + def run_context(self, params: FmhaParams) -> None: if params.qkv_input is None or params.output is None: raise RuntimeError("PrimTS context requires QKV input and an output buffer.") diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 4e95b945e8ae..291fe49829c1 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -2074,6 +2074,15 @@ def runtime_workspace_bytes_per_token(cls, model_config: "ModelConfig", and (sm_version == 90 or is_sm_100f(sm_version) or is_sm_120f(sm_version))) if not fp8_context_mla: + # PrimTS packs the strided kv_b_proj V view even with BF16 KV. + # Cached prefixes can exceed the fresh-prefill profiling extent. + from .fmha.prims_ts import PrimsTSFmha + from .fmha.registry import get_enabled_fmha_lib_classes + + if (sparse_algorithm is None and get_sm_version() in (100, 103) + and PrimsTSFmha in get_enabled_fmha_lib_classes()): + attn_tp = 1 if mapping.enable_attention_dp else mapping.tp_size + return config.num_attention_heads // attn_tp * config.v_head_dim * 2 return 0 # Attention-DP runs the full head set per rank; otherwise heads shard across TP (mirror # mNumAttnHeads). diff --git a/tests/unittest/_torch/attention/test_prims_ts_fmha.py b/tests/unittest/_torch/attention/test_prims_ts_fmha.py index c1d1fe078709..9bccd4d3797a 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_fmha.py +++ b/tests/unittest/_torch/attention/test_prims_ts_fmha.py @@ -69,6 +69,11 @@ def is_contiguous(self) -> bool: def numel(self) -> int: return math.prod(self.shape) + def stride(self, dim: int) -> int: + if dim == 1: + return 1 + return self.shape[1] if self._contiguous else self.shape[1] * 2 + class _Attention: out_head_size = TrtllmAttention.out_head_size @@ -186,7 +191,8 @@ def _support_result( if is_fused_qkv and not is_mla: q_width += 2 * attn.num_kv_heads * head_dim q = _TensorSpec((4, q_width), dtype) - output_width = attn.num_heads * (512 if is_mla else head_dim) + mla_context = is_mla and attention_input_type == AttentionInputType.context_only + output_width = attn.num_heads * (128 if mla_context else 512 if is_mla else head_dim) forward_args = AttentionForwardArgs( output=_TensorSpec((4, output_width), output_dtype) if has_output else None, attention_input_type=attention_input_type, @@ -260,7 +266,15 @@ def _support_result( ) fmha = PrimsTSFmha(attn) k = _TensorSpec((4, attn.num_kv_heads * head_dim), dtype) if has_separate_kv else None - v = _TensorSpec((4, attn.num_kv_heads * head_dim), dtype) if has_separate_kv else None + v = ( + _TensorSpec( + (4, attn.num_kv_heads * (128 if mla_context else head_dim)), + dtype, + contiguous=not mla_context, + ) + if has_separate_kv + else None + ) return fmha._is_supported_with_reason( q, k, @@ -272,6 +286,23 @@ def _support_result( ) +@pytest.mark.parametrize("num_heads", [6, 12, 96]) +@pytest.mark.parametrize("fp8", [False, True]) +def test_supports_mla_context(num_heads: int, fp8: bool) -> None: + supported, reason = _support_result( + attention_input_type=AttentionInputType.context_only, + head_dim=192, + is_mla=True, + num_heads=num_heads, + num_kv_heads=num_heads, + has_separate_kv=True, + is_fused_qkv=False, + kv_dtype=DataType.FP8 if fp8 else DataType.BF16, + quant_mode=QuantMode.from_description(use_fp8_kv_cache=fp8), + ) + assert supported, reason + + @pytest.mark.parametrize( "case", [ @@ -353,8 +384,10 @@ def test_fp8_mla_supported( ({"quant_mode": QuantMode(0)}, "query and KV-cache dtypes"), ({"dtype": torch.float16}, "BF16 query input and output"), ({"output_dtype": torch.float8_e4m3fn}, "output dtype must match"), - ({"attention_input_type": AttentionInputType.context_only}, "generation-only"), - ({"attention_input_type": AttentionInputType.mixed}, "generation-only"), + ({"has_separate_kv": True}, "MLA decode requires query-only input"), + ({"is_fused_qkv": True}, "MLA decode requires query-only input"), + ({"attention_input_type": AttentionInputType.context_only}, "MLA context requires QK=192"), + ({"attention_input_type": AttentionInputType.mixed}, "dispatched separately"), ({"use_spec_decoding": True}, "speculative decoding"), ({"has_sparse_attention": True}, "sparse attention"), ], @@ -580,7 +613,7 @@ def test_is_supported_accepts_and_forwards_phase_keyword( "head_dim": 576, "is_mla": True, }, - "generation-only", + "MLA context requires QK=192", ), ( { @@ -915,6 +948,75 @@ def test_sequence_lengths_are_live_source_view() -> None: torch.testing.assert_close(actual, torch.tensor([33, 65], dtype=torch.int32)) +def test_mla_context_wrapper_reuses_plan_and_updates_chunk_offsets( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + attn = _Attention(head_dim=192, is_mla=True, num_heads=6, num_kv_heads=6) + attn._ensure_rope_table_size = Mock() + attn.get_local_layer_idx = lambda _meta: 0 + attn.rotary_cos_sin = torch.empty(1) + fmha = PrimsTSFmha(attn) + q = torch.empty(5, 6 * 192, dtype=torch.bfloat16) + k = torch.empty(9, 6 * 192, dtype=q.dtype) + v = torch.empty(9, 6 * 128, dtype=q.dtype) + qo_indptr = torch.tensor([0, 2, 5], dtype=torch.int32) + output = torch.empty(5, 6, 128, dtype=q.dtype) + stats = torch.empty(5, 6, 2) + meta = SimpleNamespace( + mla_prepare_ctx_cu_seqlens=lambda: qo_indptr, + max_seq_len=4096, + kv_cache_block_offsets=torch.empty(2, 2, 64, dtype=torch.int32), + host_kv_cache_pool_pointers=torch.empty(1), + host_kv_cache_pool_mapping=torch.zeros(1, 2, dtype=torch.int32), + ) + workspace = torch.empty(fmha._mla_context_workspace_layout(q, k, 2)[1], dtype=torch.uint8) + params = FmhaParams( + attn=attn, + meta=meta, + fwd=AttentionForwardArgs( + attention_mask=PredefinedAttentionMask.FULL, + softmax_stats_tensor=stats, + is_fused_qkv=False, + ), + workspace=workspace, + query_input=q, + key_input=k, + value_input=v, + output=output, + sequence_lengths=torch.tensor([7, 2], dtype=torch.int32), + batch_size=2, + input_seq_length=3, + max_past_kv_length=7, + tokens_per_block=64, + cyclic_attention_window_size=4096, + ) + native_preprocess = Mock() + monkeypatch.setattr(torch.ops.trtllm, "mla_context_preprocess", native_preprocess) + wrapper = Mock() + factory = Mock(return_value=wrapper) + monkeypatch.setattr(prims_context_module, "BatchPrefillTSWrapper", factory) + fmha.run_mla_context(params) + first_kv_offsets = wrapper.run.call_args.args[4].clone() + params.sequence_lengths.copy_(torch.tensor([2, 7], dtype=torch.int32)) + qo_indptr.copy_(torch.tensor([0, 3, 5], dtype=torch.int32)) + fmha.run_mla_context(params) + assert factory.call_count == wrapper.plan.call_count == 1 + assert native_preprocess.call_count == wrapper.run.call_count == 2 + torch.testing.assert_close(first_kv_offsets, torch.tensor([0, 7, 9], dtype=torch.int32)) + torch.testing.assert_close( + wrapper.run.call_args.args[4], torch.tensor([0, 2, 9], dtype=torch.int32) + ) + assert wrapper.run.call_args.args[3] is qo_indptr + assert wrapper.run.call_args.kwargs["softmax_stats"] is stats + assert wrapper.run.call_args.kwargs["out"] is output + assert wrapper.run.call_args.args[1].shape == (9, 6, 192) + assert wrapper.plan.call_args.kwargs["store_softmax_stats"] + assert wrapper.plan.call_args.kwargs["k_dtype"] == k.dtype + assert wrapper.plan.call_args.kwargs["v_dtype"] == v.dtype + assert wrapper.run.call_args.kwargs["validate"] is False + + def test_context_wrapper_plans_once_and_reads_live_fixed_metadata( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -2258,7 +2360,8 @@ def test_phased_forward_routes_query_and_qkv_by_phase( else: q_size = (attn.num_heads + (2 * attn.num_kv_heads if is_fused_qkv else 0)) * attn.head_dim q = torch.arange(num_tokens * q_size, dtype=torch.float32).view(num_tokens, q_size) - k = torch.randn((num_tokens, attn.num_kv_heads * attn.head_dim)) if has_kv else None + num_kv_tokens = num_tokens + 7 if is_mla and has_kv else num_tokens + k = torch.randn((num_kv_tokens, attn.num_kv_heads * attn.head_dim)) if has_kv else None v = torch.randn_like(k) if k is not None else None out_head_size = ( fmha.generation_out_head_size @@ -2328,8 +2431,12 @@ def test_phased_forward_routes_query_and_qkv_by_phase( torch.testing.assert_close(phase_input, q[token_slice]) assert phase_input.data_ptr() == q[token_slice].data_ptr() if has_kv: - torch.testing.assert_close(params.key_input, k[token_slice]) - torch.testing.assert_close(params.value_input, v[token_slice]) + if is_mla: + assert params.key_input is k + assert params.value_input is v + else: + torch.testing.assert_close(params.key_input, k[token_slice]) + torch.testing.assert_close(params.value_input, v[token_slice]) else: assert params.key_input is None assert params.value_input is None diff --git a/tests/unittest/_torch/pyexecutor/test_mla_workspace_reserve.py b/tests/unittest/_torch/pyexecutor/test_mla_workspace_reserve.py index 19ebb0ef5b60..98bf74f2321e 100644 --- a/tests/unittest/_torch/pyexecutor/test_mla_workspace_reserve.py +++ b/tests/unittest/_torch/pyexecutor/test_mla_workspace_reserve.py @@ -213,6 +213,20 @@ def _fp8_mla_model_config(sparse_algorithm): ) +@pytest.mark.parametrize("library", ["fallback", "prims_ts"]) +@pytest.mark.parametrize("attention_dp", [False, True]) +def test_prims_ts_bf16_mla_v_packing_reserve(monkeypatch, library, attention_dp): + monkeypatch.setenv("TLLM_FMHA_LIBS", library) + monkeypatch.setattr(trtllm_backend, "get_sm_version", lambda: 100) + config = _fp8_mla_model_config(None) + config.quant_config.quant_mode.has_fp8_kv_cache = lambda: False + mapping = SimpleNamespace(enable_attention_dp=attention_dp, tp_size=16) + local_heads = 128 if attention_dp else 8 + assert get_attention_workspace_bytes_per_token(config, mapping) == ( + local_heads * 128 * 2 if library == "prims_ts" else 0 + ) + + @pytest.mark.parametrize( "sparse_algorithm,sm,short_seq_mha,expect_zero", [ From 29d3ab388f4879307773d54516d87741d7cfb6a6 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Wed, 28 Oct 2026 04:28:58 +0000 Subject: [PATCH 04/19] [None][test] Cover Kimi K3 PrimTS prefill and decode end to end Add K3 full-attention coverage for fresh/cached/chunked prefill, BF16/FP8 KV, TP/attention-DP heads, and CUDA graph decode. Check fresh prefill against independent FP64 math after preprocessing; retain direct backend comparisons for cached/chunk-merged prefill and decode. Add forced-PrimTS GSM8K CI cases, supported KvCacheConfig dtype selection, and the missing FP8-KV reference without relaxing accuracy thresholds. Hecate qualification on the equivalent full functional snapshot passed 1319 examples per dtype (BF16 96.058%, FP8 96.740%). Scheduling optimization is not part of this commit. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../defs/accuracy/references/gsm8k.yaml | 3 + tests/integration/defs/accuracy/test_kimi3.py | 42 +++ .../l0_gb300_multi_nodes_node4_gpu16.yml | 2 + .../test_prims_ts_attention_backend.py | 243 +++++++++++++++++- 4 files changed, 284 insertions(+), 6 deletions(-) diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index e81ec522bc20..fc34253a66a7 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -344,6 +344,9 @@ moonshotai/Kimi-K2.5: # layer does not surface — the reference matcher sees quant_algo=None. moonshotai/Kimi-K3: - accuracy: 96.5 + # Apply the existing full-model accuracy gate to FP8 KV-cache qualification. + - kv_cache_quant_algo: FP8 + accuracy: 96.5 # SA spec dec is lossless; scores match the baseline within noise. - spec_dec_algo: SA accuracy: 96.5 diff --git a/tests/integration/defs/accuracy/test_kimi3.py b/tests/integration/defs/accuracy/test_kimi3.py index 00ccf4c18de9..7c8bedbfe799 100644 --- a/tests/integration/defs/accuracy/test_kimi3.py +++ b/tests/integration/defs/accuracy/test_kimi3.py @@ -164,6 +164,48 @@ def test_w4a16_mxfp4(self, mode: str, monkeypatch: pytest.MonkeyPatch) -> None: llm, ) + @skip_pre_blackwell + @pytest.mark.skip_less_mpi_world_size(16) + @pytest.mark.skip_less_device_memory(200000) + @pytest.mark.parametrize("fp8_kv", [False, True], ids=["bf16-kv", "fp8-kv"]) + def test_prims_ts(self, fp8_kv: bool, monkeypatch: pytest.MonkeyPatch) -> None: + """Full K3: strict PrimTS MLA context/decode, chunking, reuse and graphs. + + An exact FMHA list makes an unsupported full-attention phase fail on + the workers rather than silently falling back. KDA is unchanged. + This shares the existing DEP16 GB300 NVL72 qualification envelope. + """ + monkeypatch.setenv("TLLM_FMHA_LIBS", "prims_ts") + self._assert_checkpoint_routing() + with LLM( + self.MODEL_PATH, + tensor_parallel_size=16, + moe_expert_parallel_size=16, + enable_attention_dp=True, + max_batch_size=32, + # The explicit 4096-token reuse probe must span context chunks. + max_num_tokens=2048, + max_seq_len=8192, + trust_remote_code=True, + enable_chunked_prefill=True, + cuda_graph_config=CudaGraphConfig(enable_padding=True, max_batch_size=32), + moe_config=MoeConfig(max_num_tokens=33024, use_low_precision_moe_combine=True), + kv_cache_config=KvCacheConfig( + dtype="fp8" if fp8_kv else "auto", + enable_block_reuse=True, + free_gpu_memory_fraction=0.25, + tokens_per_block=64, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=256), + ), + enable_iter_perf_stats=True, + max_stats_len=256, + return_perf_metrics=True, + ) as llm: + self._assert_resolved_args(llm, "reuse") + assert llm.args.quant_config.kv_cache_quant_algo == (QuantAlgo.FP8 if fp8_kv else None) + assert_kv_cache_reuse_for_llm(llm, [1] + [42] * 4094 + [43]) + GSM8K(self.MODEL_NAME).evaluate(llm) + @skip_pre_blackwell @pytest.mark.skip_less_mpi_world_size(16) @pytest.mark.skip_less_device_memory(200000) diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node4_gpu16.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node4_gpu16.yml index 1e01a37b6b2a..37dff661066b 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node4_gpu16.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node4_gpu16.yml @@ -21,3 +21,5 @@ l0_gb300_multi_nodes_node4_gpu16: - accuracy/test_kimi3.py::TestKimiK3::test_w4a16_mxfp4[baseline] TIMEOUT (90) - accuracy/test_kimi3.py::TestKimiK3::test_w4a16_mxfp4[sa] TIMEOUT (90) - test_e2e.py::test_multi_nodes_eval[MiniMax-M3-tp16-mmlu] + - accuracy/test_kimi3.py::TestKimiK3::test_prims_ts[bf16-kv] TIMEOUT (90) + - accuracy/test_kimi3.py::TestKimiK3::test_prims_ts[fp8-kv] TIMEOUT (90) 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..370cfcb0ea7c 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py +++ b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py @@ -379,24 +379,37 @@ def test_prims_ts_deepseek_v3_lite_mla_generation( run_case(case) +@pytest.mark.parametrize("kv_cache_dtype", [torch.bfloat16, torch.float8_e4m3fn]) @pytest.mark.parametrize("num_heads", [6, 12, 96]) @pytest.mark.parametrize("use_kv_cache_manager_v2", [False, True], ids=["v1", "v2"]) -def test_prims_ts_fp8_mla_preprocessing( - monkeypatch: pytest.MonkeyPatch, num_heads: int, use_kv_cache_manager_v2: bool +def test_prims_ts_mla_preprocessing( + monkeypatch: pytest.MonkeyPatch, + num_heads: int, + use_kv_cache_manager_v2: bool, + kv_cache_dtype: torch.dtype, ) -> None: from test_attention_mla import RopeConfig, _run_test_for_backend from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams from tensorrt_llm._torch.attention.backends.fmha.prims_ts import PrimsTSFmha - # Context remains on the regular backend; every decode must select PrimTS. - monkeypatch.setenv("TLLM_FMHA_LIBS", "prims_ts,fallback") + # No fallback: both context and every decode must select PrimTS. + monkeypatch.setenv("TLLM_FMHA_LIBS", "prims_ts") original_run = PrimsTSFmha.run_mla_generation + original_context = PrimsTSFmha.run_mla_context calls = [] + context_calls = [] + + def run_context(fmha: PrimsTSFmha, params: FmhaParams) -> None: + original_context(fmha, params) + context_calls.append(params.attn.local_layer_idx) + + monkeypatch.setattr(PrimsTSFmha, "run_mla_context", run_context) def run_and_capture(fmha: PrimsTSFmha, params: FmhaParams) -> None: assert params.query_input.dtype == torch.bfloat16 - assert params.fwd.quant_q_buffer.dtype == torch.uint8 + if kv_cache_dtype == torch.float8_e4m3fn: + assert params.fwd.quant_q_buffer.dtype == torch.uint8 original_run(fmha, params) eager = params.output.clone() graph = torch.cuda.CUDAGraph() @@ -435,13 +448,231 @@ def run_and_capture(fmha: PrimsTSFmha, params: FmhaParams) -> None: kv_cache_tokens_per_block=32, device=torch.device("cuda"), dtype=torch.bfloat16, - kv_cache_dtype=torch.float8_e4m3fn, + kv_cache_dtype=kv_cache_dtype, context_sequence_lengths=[129, 191], generation_seq_len_q=1, num_generation_steps=2, v2_kv_cache=use_kv_cache_manager_v2, ) assert calls == [0, 1, 0, 1] + assert context_calls == [0, 1] + + +@pytest.mark.parametrize("num_heads", [6, 96], ids=["tp16-local-heads", "dp-heads"]) +@pytest.mark.parametrize("fp8_kv", [False, True], ids=["bf16-kv", "fp8-kv"]) +@pytest.mark.parametrize("chunked", [False, True], ids=["cached", "chunked"]) +def test_prims_ts_kimi_k3_full_attention( + monkeypatch: pytest.MonkeyPatch, num_heads: int, fp8_kv: bool, chunked: bool +) -> None: + """Exercise K3's real NoPE/gated MLA module, without full-model weights. + + Head counts cover the TP16-local and attention-DP geometries, not distributed + collectives. The full-model GB300 test covers those separately. + """ + import tensorrt_llm + from tensorrt_llm._torch.attention.backends import TrtllmAttention + from tensorrt_llm._torch.attention.backends.fmha.prims_ts import PrimsTSFmha + from tensorrt_llm._torch.attention.backends.interface import AttentionRuntimeFeatures + from tensorrt_llm._torch.metadata import KVCacheParams + from tensorrt_llm._torch.model_config import ModelConfig + from tensorrt_llm._torch.modules.kimi_k3_mla.kimi_k3_mla_attention import KimiK3MLAAttention + from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, SamplingConfig + from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager + from tensorrt_llm._utils import torch_dtype_to_binding + from tensorrt_llm.llmapi import KvCacheConfig, SamplingParams + from tensorrt_llm.mapping import Mapping + from tensorrt_llm.models.modeling_utils import QuantConfig + from tensorrt_llm.quantization import QuantAlgo + + context_calls = [] + generation_calls = [] + prefill_reference = None + run_context = PrimsTSFmha.run_mla_context + run_generation = PrimsTSFmha.run_mla_generation + + def record_context(fmha, params): + nonlocal prefill_reference + context_calls.append( + (params.fwd.softmax_stats_tensor is not None, params.key_input.shape[0]) + ) + run_context(fmha, params) + if len(context_calls) == 1: + # Preprocessing has now populated K's initially uninitialized RoPE + # slice. Compare both FP8 kernels to an independent reference, not + # to each other: their softmax rounding errors can add rather than + # cancel even when each meets the same accuracy bound. + operands = (params.query_input, params.key_input, params.value_input) + scale = fmha._get_bmm1_scale(params.attn) + output_scale = 1.0 + if fp8_kv: + operands = tuple( + (tensor.float() * params.fwd.kv_scale_orig_quant).to(torch.float8_e4m3fn) + for tensor in operands + ) + output_scale = params.fwd.kv_scale_quant_orig.item() + scale *= output_scale**2 + q_ref, k_ref, v_ref = ( + tensor.reshape(tensor.shape[0], num_heads, -1).double() for tensor in operands + ) + references = [] + offset = 0 + for length in lengths[0]: + q_seq, k_seq, v_seq = ( + tensor[offset : offset + length].transpose(0, 1) + for tensor in (q_ref, k_ref, v_ref) + ) + scores = (q_seq @ k_seq.transpose(-1, -2)) * scale + causal_mask = torch.ones( + (length, length), device=scores.device, dtype=torch.bool + ).tril() + scores.masked_fill_(~causal_mask, -torch.inf) + references.append(((scores.softmax(-1) @ v_seq) * output_scale).transpose(0, 1)) + offset += length + prefill_reference = torch.cat(references).flatten(1).to(params.output.dtype) + + def record_generation(fmha, params): + generation_calls.append(params.batch_size) + run_generation(fmha, params) + + monkeypatch.setattr(PrimsTSFmha, "run_mla_context", record_context) + monkeypatch.setattr(PrimsTSFmha, "run_mla_generation", record_generation) + kv_dtype = torch.float8_e4m3fn if fp8_kv else torch.bfloat16 + lengths = [[129, 65], [17, 193], [1, 1]] + torch.manual_seed(2718) + hidden_steps = [ + torch.randn(sum(q_lengths), 256, dtype=torch.bfloat16, device="cuda") + for q_lengths in lengths + ] + results = {} + for library in ("fallback", "prims_ts"): + monkeypatch.setenv("TLLM_FMHA_LIBS", library) + torch.manual_seed(1729) + mapping = Mapping(world_size=1, tp_size=1, rank=0) + module = ( + KimiK3MLAAttention( + hidden_size=256, + num_heads=num_heads, + q_lora_rank=128, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + dtype=torch.bfloat16, + model_config=ModelConfig( + mapping=mapping, + quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8 if fp8_kv else None), + ), + aux_stream_dict={}, + ) + .cuda() + .eval() + ) + # Use the module's eager entry, retaining the real projection, cache, + # context/chunk-merge, absorption, gate and output-projection paths. + module.register_to_config = False + with torch.no_grad(): + for name, parameter in module.named_parameters(): + if "layernorm" in name: + parameter.fill_(1.0) + else: + parameter.normal_(std=0.02) + k_weight, v_weight = module.kv_b_proj.weight.split(num_heads * 128) + module.k_b_proj_trans.copy_(k_weight.view(num_heads, 128, 512).transpose(1, 2)) + module.v_b_proj.copy_(v_weight.view(num_heads, 128, 512)) + manager = KVCacheManager( + KvCacheConfig(max_tokens=1024, enable_block_reuse=False), + tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY, + num_layers=1, + num_kv_heads=1, + head_dim=576, + tokens_per_block=64, + max_seq_len=512, + max_batch_size=2, + mapping=mapping, + dtype=torch_dtype_to_binding(kv_dtype), + ) + # Both decode implementations require finite values in unused cache + # slots. Reused BF16 allocations can contain FP8 NaN byte patterns; + # keep padding deterministic while real context writes every live slot. + manager.get_buffers(0).fill_(1.0) + requests = [] + for request_id, length in enumerate(lengths[0]): + request = LlmRequest( + request_id=request_id, + input_tokens=[1] * length, + max_new_tokens=256, + sampling_config=SamplingConfig(SamplingParams()._get_sampling_config()), + is_streaming=False, + ) + request.paged_kv_block_ids = [] + manager.impl.add_sequence_batch([(request_id, length, 1)], [request]) + requests.append(request) + cached = [0, 0] + results[library] = [] + try: + with torch.inference_mode(): + for step, q_lengths in enumerate(lengths): + if step: + for request_id, length in enumerate(q_lengths): + for _ in range(length): + manager.impl.add_token(request_id) + metadata = TrtllmAttention.Metadata( + seq_lens=torch.tensor(q_lengths, dtype=torch.int32), + request_ids=[0, 1], + max_num_requests=2, + num_contexts=2 if step < 2 else 0, + prompt_lens=q_lengths, + max_num_tokens=512, + kv_cache_manager=manager, + kv_cache_params=KVCacheParams( + use_cache=True, num_cached_tokens_per_seq=cached + ), + mapping=mapping, + enable_context_mla_with_cached_kv=True, + runtime_features=AttentionRuntimeFeatures( + cache_reuse=True, + chunked_prefill=chunked, + chunk_size=128, + chunked_prefill_buffer_batch_size=1, + ), + ) + metadata.prepare() + hidden = hidden_steps[step] + positions = torch.cat( + [ + torch.arange(start, start + length, device="cuda") + for start, length in zip(cached, q_lengths) + ] + ) + result = module(positions, hidden, metadata) + assert torch.isfinite(result).all(), (library, step, "eager") + results[library].append(result.clone()) + if step == 0 and library == "prims_ts": + assert prefill_reference is not None + expected = module.o_proj( + module._apply_output_gate(hidden, prefill_reference) + ) + for candidate in (results["fallback"][0], result): + torch.testing.assert_close( + candidate, expected, atol=0.02 if fp8_kv else 0.003, rtol=0.03 + ) + if step == 2 and library == "prims_ts": + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = module(positions, hidden, metadata) + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(captured, result, atol=0, rtol=0) + cached = [start + length for start, length in zip(cached, q_lengths)] + finally: + manager.shutdown() + # Fresh prefill was checked against FP64 math above; retain the direct + # backend comparison for cached/chunk-merged prefill and decode. + for actual, expected in zip(results["prims_ts"][1:], results["fallback"][1:]): + torch.testing.assert_close(actual, expected, atol=0.02 if fp8_kv else 0.003, rtol=0.03) + assert len(context_calls) == (4 if chunked else 2) + assert any(has_stats for has_stats, _ in context_calls) == chunked + assert generation_calls == [2, 2] # Eager and graph capture; replay is device-only. @pytest.mark.parametrize("num_heads", [6, 12, 96]) From 7c77d54076f883d61e8153e5ebd762a06ef425da Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Wed, 28 Oct 2026 04:29:55 +0000 Subject: [PATCH 05/19] [None][feat] Enable PrimTS attention on Rubin with CUTLASS DSL 4.8 Keep Blackwell on the existing DSL 4.7 minimum and require DSL 4.8 for SM107, retaining the compiler 13.3 gate. Extend availability tests and full-attention test collection to Rubin. This is architecture enablement only, not the dense-Q scheduling optimization. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../_torch/attention/backends/fmha/prims_ts.py | 18 +++++++++++------- .../test_prims_ts_attention_backend.py | 4 ++-- .../_torch/attention/test_prims_ts_fmha.py | 6 +++++- 3 files changed, 18 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py index 9caaf5bafcce..dfdb58fe5c04 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 Blackwell/Rubin 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 @@ -153,11 +154,11 @@ def get_attention_feature_unsupported_reason( class PrimsTSFmha(PhasedFmha): - """Blackwell task-scheduled standard/MLA context and decode FMHA library.""" + """Blackwell/Rubin task-scheduled standard/MLA context and decode FMHA library.""" 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} @@ -195,17 +196,20 @@ 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 + 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 370cfcb0ea7c..5ba32f99335b 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,8 @@ from utils.util import getSMVersion pytestmark = pytest.mark.skipif( - getSMVersion() not in (100, 103), - reason="PrimsTS attention kernels require SM100 or SM103", + getSMVersion() not in (100, 103, 107), + reason="PrimsTS attention kernels require SM100, SM103 or SM107", ) diff --git a/tests/unittest/_torch/attention/test_prims_ts_fmha.py b/tests/unittest/_torch/attention/test_prims_ts_fmha.py index 9bccd4d3797a..008ae7bb8f22 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_fmha.py +++ b/tests/unittest/_torch/attention/test_prims_ts_fmha.py @@ -743,6 +743,9 @@ def test_unsupported_matrix_falls_through(case: dict, expected_reason: str) -> N (100, "4.7.0", "13.4", True), (100, "4.7.0", "13.2", False), (107, "4.7.0", "13.3", False), + (107, "4.8.0", "13.4", True), + (107, "4.8.0", "13.2", False), + (107, "4.8.0.dev0", "13.4", False), (120, "4.7.0", "13.3", False), (100, "4.6.2", "13.3", False), ], @@ -771,7 +774,8 @@ def import_cutlass_module(module_name: str) -> object: 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"): + min_dsl_version = Version("4.8.0" if sm == 107 else "4.7.0") + if sm in (100, 103, 107) and Version(cutlass_version) >= min_dsl_version: target_version.assert_called_once_with(min_version="13.3") else: target_version.assert_not_called() From 05fe1a81f68aeed86972206374b6e94854e9cb5a Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Thu, 29 Oct 2026 01:58:53 +0000 Subject: [PATCH 06/19] [None][feat] Vendor PrimTS per-query MLA bounds for Helix Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../vendor_patches/flashinfer-prims-ts.patch | 31 ++--- 3rdparty/vendor_sources.lock.yaml | 6 +- .../kernels/mla_decode/helpers/kv_bounds.py | 71 +++++++++++ .../kernels/mla_decode/helpers/tile.py | 22 +++- .../mla_decode/throughput_2cta/config.py | 4 + .../mla_decode/throughput_2cta/kernel.py | 21 +++- .../throughput_2cta/parallel_reduction.py | 7 +- .../mla_decode/throughput_2cta/reduction.py | 5 +- .../mla_decode/throughput_2cta/resources.py | 33 ++++- .../throughput_latency_1cta/config.py | 2 + .../throughput_latency_1cta/kernel.py | 18 +++ .../resources/smem_resources.py | 26 +++- .../resources/tmem_s.py | 21 ++-- .../attention/backends/prims_ts/mla_decode.py | 114 +++++++++++++++++- 14 files changed, 331 insertions(+), 50 deletions(-) create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/kv_bounds.py diff --git a/3rdparty/vendor_patches/flashinfer-prims-ts.patch b/3rdparty/vendor_patches/flashinfer-prims-ts.patch index b47dd844bfeb..d749481966f5 100644 --- a/3rdparty/vendor_patches/flashinfer-prims-ts.patch +++ b/3rdparty/vendor_patches/flashinfer-prims-ts.patch @@ -419,15 +419,10 @@ index b22f0f413b7312bab9a5f8e296255fe9aeabe901..a4a87200ca5d3ace5cf899caa9672a1a from .helpers_common import ( _TASK_CACHE_TMEM_BASE_OFFSET, diff --git a/mla_decode.py b/mla_decode.py -index 6ca10c5f7195531890abd3f0bf643f6ceab36e5e..594d160b83aaa1b7bde74aa691e5e4ab92e630db 100644 +index 7ea49a00d3ed6a236bd9da2854b45a985e23d4a5..f0a6e9ddf697fc80748b859fd2818f683ab31eed 100644 --- a/mla_decode.py +++ b/mla_decode.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"); -@@ -27,11 +28,7 @@ from ._hardware import get_max_active_clusters +@@ -28,11 +28,7 @@ from ._hardware import get_max_active_clusters import torch @@ -440,7 +435,7 @@ index 6ca10c5f7195531890abd3f0bf643f6ceab36e5e..594d160b83aaa1b7bde74aa691e5e4ab from .decode import ( _WorkspaceSection, -@@ -1080,11 +1077,35 @@ def _make_mla_decode_compile_spec( +@@ -1133,11 +1129,35 @@ def _make_mla_decode_compile_spec( ) @@ -481,8 +476,8 @@ index 6ca10c5f7195531890abd3f0bf643f6ceab36e5e..594d160b83aaa1b7bde74aa691e5e4ab import cutlass import cutlass.cute as cute -@@ -1697,7 +1718,7 @@ class BatchMLADecodePagedTSWrapper: - store_softmax_stats=store_softmax_stats, +@@ -1776,7 +1796,7 @@ class BatchMLADecodePagedTSWrapper: + use_kv_bounds=use_kv_bounds, ) - @flashinfer_experimental_api(trace=prims_ts_decode_mla_wrapper_trace_dispatch) @@ -490,7 +485,7 @@ index 6ca10c5f7195531890abd3f0bf643f6ceab36e5e..594d160b83aaa1b7bde74aa691e5e4ab def run( self, query: torch.Tensor, -@@ -1814,7 +1835,7 @@ class BatchMLADecodePagedTSWrapper: +@@ -1905,7 +1925,7 @@ class BatchMLADecodePagedTSWrapper: ) @@ -499,7 +494,7 @@ index 6ca10c5f7195531890abd3f0bf643f6ceab36e5e..594d160b83aaa1b7bde74aa691e5e4ab def batch_mla_decode_with_paged_kv_cache( query: torch.Tensor, kv_cache: torch.Tensor, -@@ -2181,6 +2202,7 @@ def _prepare_sparse_output_quant( +@@ -2284,6 +2304,7 @@ def _prepare_sparse_output_quant( ) @@ -507,7 +502,7 @@ index 6ca10c5f7195531890abd3f0bf643f6ceab36e5e..594d160b83aaa1b7bde74aa691e5e4ab def _resolve_sparse_mla_plan( device, batch_size, -@@ -2201,17 +2223,17 @@ def _resolve_sparse_mla_plan( +@@ -2304,17 +2325,17 @@ def _resolve_sparse_mla_plan( ): """Resolve sparse geometry and scratch sizing without allocation/compilation.""" import cutlass @@ -529,7 +524,7 @@ index 6ca10c5f7195531890abd3f0bf643f6ceab36e5e..594d160b83aaa1b7bde74aa691e5e4ab MlaDecodeTs, ) -@@ -2346,6 +2368,9 @@ def _resolve_sparse_mla_plan( +@@ -2449,6 +2470,9 @@ def _resolve_sparse_mla_plan( return device, profile, spec, sections, byte_end, capacity, output_quant @@ -539,7 +534,7 @@ index 6ca10c5f7195531890abd3f0bf643f6ceab36e5e..594d160b83aaa1b7bde74aa691e5e4ab class BatchSparseMLADecodePagedTSWrapper: """Plan native BF16/E4M3 attention, then bind one or two source index lists. -@@ -2379,6 +2404,11 @@ class BatchSparseMLADecodePagedTSWrapper: +@@ -2482,6 +2506,11 @@ class BatchSparseMLADecodePagedTSWrapper: self._workspace_buffer = workspace_buffer self._state = None self._scalar_cache = {} @@ -551,7 +546,7 @@ index 6ca10c5f7195531890abd3f0bf643f6ceab36e5e..594d160b83aaa1b7bde74aa691e5e4ab self.workspace_size_bytes = 0 def plan( -@@ -2468,6 +2498,7 @@ class BatchSparseMLADecodePagedTSWrapper: +@@ -2571,6 +2600,7 @@ class BatchSparseMLADecodePagedTSWrapper: kernel_signature=_mla_kernel_compile_signature(static_kernel), ) ) @@ -559,7 +554,7 @@ index 6ca10c5f7195531890abd3f0bf643f6ceab36e5e..594d160b83aaa1b7bde74aa691e5e4ab self._state = dict( device=device, batch=batch_size, -@@ -2690,6 +2721,38 @@ class BatchSparseMLADecodePagedTSWrapper: +@@ -2793,6 +2823,38 @@ class BatchSparseMLADecodePagedTSWrapper: state = self._state if state is None: raise RuntimeError("plan() must be called before run()") @@ -598,7 +593,7 @@ index 6ca10c5f7195531890abd3f0bf643f6ceab36e5e..594d160b83aaa1b7bde74aa691e5e4ab _validate_query( query, packed_query=state["packed"], -@@ -2863,10 +2926,125 @@ class BatchSparseMLADecodePagedTSWrapper: +@@ -2966,10 +3028,125 @@ class BatchSparseMLADecodePagedTSWrapper: if static_scales else state["compiled"], ) diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 4037e23240be..995b256547fd 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-ts-mla-prefill--flashinfer - commit: 75961f74dac2c438cd7bce3dd135d527e78e49ec + commit: df3434b1dc23df96527e202aa43ae233ed8641ee 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:d4557cc5072cab63378dd76c251f0718582578ae38f95f0303656854e88a8d69 - digest: sha256-tree-v1:15e6ad32e7aeec1a152d49cdc43c371adae4a072b8de0d362562ecad3c937666 + patch_digest: sha256:27cc330c1eff406250b2ea7212d7a3200773a4c2d9eaacd57e57033344d39f3c + digest: sha256-tree-v1:0ea9593edb2c167fcd478c30d6ac423d31e6b2eb51c8f50e25a463b00d045158 openengine-proto: url: https://github.com/ai-dynamo/openengine.git tag: v0.1.0 diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/kv_bounds.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/kv_bounds.py new file mode 100644 index 000000000000..52bd8d04d38d --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/kv_bounds.py @@ -0,0 +1,71 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Runtime per-query prefixes and neutral publication for empty MLA rows.""" + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Int64 + +from .query import query_batch_bounds + + +@cute.jit +def query_kv_bound(kv_bounds, cu_seqlens_q, batch_idx, logical_q_idx, seq_len_q): + """Load a logical-token bound without indexing packed or fixed padding.""" + q_start, q_length = query_batch_bounds(cu_seqlens_q, batch_idx, seq_len_q) + if cutlass.const_expr(cu_seqlens_q is None): + q_start = Int32(batch_idx) * Int32(seq_len_q) + bound = Int32(0) + if logical_q_idx >= Int32(0) and logical_q_idx < q_length: + bound = Int32(kv_bounds[q_start + logical_q_idx]) + return bound + + +@cute.kernel +def initialize_empty_query_rows(output, lse, softmax_stats, kv_bounds): + """Publish identity rows even when the main schedule skips an empty cache. + + The public tensors have contiguous token/head storage before their + dimension-first kernel views. Bound values remain graph-live device data. + """ + row = Int64(cute.arch.block_idx()[0]) + thread = Int32(cute.arch.thread_idx()[0]) + heads = Int64(output.shape[0]) + dim = cutlass.const_expr(output.shape[1]) + if kv_bounds[row // heads] == Int32(0): + for column in cutlass.range(thread, dim, 128): + (output.iterator.raw_ptr() + row * Int64(dim) + Int64(column)).store( + output.element_type(0) + ) + if thread == Int32(0): + (lse.iterator.raw_ptr() + row).store(Float32(-Float32.inf)) + if cutlass.const_expr(softmax_stats is not None): + (softmax_stats.iterator.raw_ptr() + row * Int64(2)).store( + Float32(-Float32.inf) + ) + (softmax_stats.iterator.raw_ptr() + row * Int64(2) + Int64(1)).store( + Float32(0) + ) + + +@cute.jit +def initialize_empty_rows(output, lse, softmax_stats, kv_bounds, stream): + """Launch the bounded-mode neutral writer on the caller's stream.""" + initialize_empty_query_rows(output, lse, softmax_stats, kv_bounds).launch( + grid=(cute.size(kv_bounds) * output.shape[0], 1, 1), + block=(128, 1, 1), + stream=stream, + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/tile.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/tile.py index e35517b09e93..b322af4c9e2f 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/tile.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/tile.py @@ -1,3 +1,4 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -21,6 +22,7 @@ from cutlass.experimental.task_scheduling.resources import StageInfo from .layout import _TASK_CACHE_SEQ_LEN_KV +from .kv_bounds import query_kv_bound from .mask import MaskType, mask_visible_k_length from .query import ( flat_query_row_state, @@ -166,12 +168,18 @@ def runtime_seq_len_kv_for_logical_q( batch_idx, logical_q_idx, cu_seqlens_q=None, + kv_bounds=None, ): """Return the mask-visible KV length for one logical Q row. Dense decode exposes the full runtime cache. Bottom-right causal decode removes the speculative K positions following ``logical_q_idx``. + Explicit bounds replace either policy with a per-query prefix count. """ + if cutlass.const_expr(kv_bounds is not None): + return query_kv_bound( + kv_bounds, cu_seqlens_q, batch_idx, logical_q_idx, cfg.logical_seq_len_q + ) _, logical_seq_len_q = query_batch_bounds( cu_seqlens_q, batch_idx, @@ -195,11 +203,13 @@ def runtime_seq_len_kv_for_q( ): """Return the KV domain shared by flat query rows in a CTA. - The last physical row has the largest causal domain. Row-causal softmax - applies the remaining row-specific mask; all other modes can use - this CTA-visible length with the ordinary dense tail predicate. + The last physical row has the largest causal domain. Explicit per-query + bounds instead retain the full raw cache domain. Softmax applies each + row's remaining prefix mask within this shared producer/reducer domain. """ - if cutlass.const_expr(cfg.mask_type == MaskType.DENSE.value): + if cutlass.const_expr(cfg.use_kv_bounds or cfg.mask_type == MaskType.DENSE.value): + # Explicit bounds replace causality. The raw cache domain conservatively + # covers every row's prefix and is shared with split reduction. return runtime_base_seq_len_kv(cfg, cache_seqs, batch_idx) _, _, logical_q_idx, _, _ = flat_query_row_state( Int32(cfg.tile_size_q - 1), @@ -320,6 +330,7 @@ def runtime_seq_len_kv_for_query_row( cta_idx_q, row_in_tile, cu_seqlens_q=None, + kv_bounds=None, ): """Return the KV length visible to one physical flat-query row.""" _, _, logical_q_idx, _, _ = flat_query_row_state( @@ -337,6 +348,7 @@ def runtime_seq_len_kv_for_query_row( batch_idx, logical_q_idx, cu_seqlens_q, + kv_bounds, ) @@ -350,7 +362,7 @@ def runtime_seq_len_kv_from_task_cache( ): """Return the CTA-visible KV domain from the task-cached raw length.""" seq_len_kv = Int32(task_cache[_TASK_CACHE_SEQ_LEN_KV]) - if cutlass.const_expr(cfg.mask_type == MaskType.DENSE.value): + if cutlass.const_expr(cfg.use_kv_bounds or cfg.mask_type == MaskType.DENSE.value): return seq_len_kv _, logical_seq_len_q = query_batch_bounds( cu_seqlens_q, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/config.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/config.py index c144ac7babe1..ed44508836e3 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/config.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/config.py @@ -1,3 +1,4 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -273,6 +274,7 @@ class MlaDecodeConfig: # Causal is bottom-right aligned for speculative decode. Dense still masks # the ordinary per-batch KV tail at ``cache_seqs[batch]``. mask_type: str = MaskType.CAUSAL.value + use_kv_bounds: bool = False @property def tokens_per_k_tile(self) -> int: @@ -346,6 +348,7 @@ def make_mla_decode_config( is_var_seq: bool = False, is_var_split_kv: bool = False, mask_type: MaskType | str = MaskType.CAUSAL, + use_kv_bounds: bool = False, sparse_gather_warps: int = 1, sparse_kv_stages: int = 0, cache_uniform_sparse_pages: bool = False, @@ -365,6 +368,7 @@ def make_mla_decode_config( cfg.is_var_seq = is_var_seq cfg.is_var_split_kv = is_var_split_kv cfg.mask_type = normalize_mask_type(mask_type) + cfg.use_kv_bounds = use_kv_bounds def _require_positive(name: str, value: int) -> None: """Validate that a named configuration value is positive.""" diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/kernel.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/kernel.py index 36652e5496de..99137802b5a7 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/kernel.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/kernel.py @@ -1,3 +1,4 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -167,6 +168,7 @@ def build_mla_decode_task_manager( # For k_index_base computation (split_kv support) cache_seqs=None, cu_seqlens_q=None, + kv_bounds=None, split_kv=None, logical_num_heads_q=128, logical_seq_len_q=1, @@ -442,6 +444,7 @@ def build_mla_decode_task_manager( tmem_s = TmemSResource( page_offsets=page_offsets, + kv_bounds=kv_bounds, store_softmax_stats=softmax_stats is not None, smem_q_latent=smem_q_latent_arr, smem_q_rope=smem_q_rope_arr, @@ -834,6 +837,7 @@ def __init__( seq_len_q=1, batch_size=1, mask_type: MaskType | str = MaskType.CAUSAL, + use_kv_bounds: bool = False, ): """ Parameters @@ -916,6 +920,7 @@ def __init__( raise ValueError("batch_size must be positive") self.batch_size = batch_size self.mask_type = normalize_mask_type(mask_type) + self.use_kv_bounds = use_kv_bounds self.fuse_sparse_reduction = False self.fuse_sparse_epilogue = False self.output_quant = None @@ -980,6 +985,7 @@ def compile_signature(self) -> tuple[object, ...]: self.num_heads, self.seq_len_q, self.mask_type, + self.use_kv_bounds, self.reduction_split_capacity, self.query_tile_layout, self.num_q_tiles, @@ -1072,6 +1078,7 @@ def _make_config(self): is_var_seq=self.is_var_seq, is_var_split_kv=self.is_var_split_kv, mask_type=self.mask_type, + use_kv_bounds=self.use_kv_bounds, sparse_gather_warps=self.sparse_gather_warps, sparse_kv_stages=self.sparse_kv_stages, cache_uniform_sparse_pages=self.sparse_uniform_pages, @@ -1116,6 +1123,7 @@ def __call__( stream: object, output_quant_args=None, softmax_stats: cute.Tensor | None = None, + kv_bounds: cute.Tensor | None = None, ): """Execute the MLA decode TS kernel.""" if cutlass.const_expr(self.fuse_sparse_reduction): @@ -1455,6 +1463,10 @@ def _flatten_query_rows_for_tma(t): # before reading partials. Persistent schedules retain stream ordering. use_one_wave_reducer_pdl = acc_o is not None and not self.is_persistent + if cutlass.const_expr(self.use_kv_bounds): + from ..helpers.kv_bounds import initialize_empty_rows + + initialize_empty_rows(o, lse, softmax_stats, kv_bounds, stream) self.split_kv_kernel( tma_desc_q_latent, tma_desc_q_rope, @@ -1480,6 +1492,7 @@ def _flatten_query_rows_for_tma(t): tma_desc_o, softmax_stats, softmax_stats_scale, + kv_bounds, ).launch( grid=grid, block=[cfg.threads_per_cta, 1, 1], @@ -1584,6 +1597,7 @@ def split_kv_kernel( tma_desc_o: cutlass.GridConstant[cuda.TensorMap] = None, softmax_stats: cute.Tensor | None = None, softmax_stats_scale=None, + kv_bounds: cute.Tensor | None = None, ) -> None: """MLA decode TS kernel: persistent tile-scheduled execution.""" if cutlass.const_expr(self.direct_sparse): @@ -1861,7 +1875,9 @@ def split_kv_kernel( tile_batch_idx, ) if cutlass.const_expr( - cfg.mask_type == MaskType.CAUSAL.value and self.seq_len_q > 1 + not cfg.use_kv_bounds + and cfg.mask_type == MaskType.CAUSAL.value + and self.seq_len_q > 1 ): _, _, logical_q_idx, _, _ = flat_query_row_state( Int32(self.mma_qk_tiler_mn[0] - 1), @@ -2023,6 +2039,7 @@ def split_kv_kernel( work_queue=mla_work_queue, cache_seqs=cache_seqs, cu_seqlens_q=cu_seqlens_q, + kv_bounds=kv_bounds, split_kv=max_split_kv, logical_num_heads_q=self.num_heads, logical_seq_len_q=self.seq_len_q, @@ -2184,6 +2201,7 @@ def reduction_kernel( qkv_dtype=self.qkv_dtype, o_dtype=self.out_dtype, mask_type=self.mask_type, + use_kv_bounds=self.use_kv_bounds, ) if cutlass.const_expr(not self.is_persistent): prims.griddepcontrol(kind=prims.GridDepAction.WAIT) @@ -2229,6 +2247,7 @@ def parallel_reduction_kernel( qkv_dtype=self.qkv_dtype, o_dtype=self.out_dtype, mask_type=self.mask_type, + use_kv_bounds=self.use_kv_bounds, ) if cutlass.const_expr(not self.is_persistent): prims.griddepcontrol(kind=prims.GridDepAction.WAIT) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/parallel_reduction.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/parallel_reduction.py index eb0e39ba4efb..392486521653 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/parallel_reduction.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/parallel_reduction.py @@ -1,3 +1,4 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -95,7 +96,11 @@ def _parallel_reduction_row_state( ) tile_k = cache_seqs[batch_idx] row_k = tile_k - if cutlass.const_expr(cfg.mask_type == MaskType.CAUSAL.value and seq_len_q > 1): + if cutlass.const_expr( + not cfg.use_kv_bounds + and cfg.mask_type == MaskType.CAUSAL.value + and seq_len_q > 1 + ): _, logical_seq_len_q = query_batch_bounds( cu_seqlens_q, batch_idx, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/reduction.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/reduction.py index 8306baaa66d1..8a36ba147773 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/reduction.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/reduction.py @@ -1,3 +1,4 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -140,7 +141,9 @@ def run_reduction_kernel( tile_k = cache_seqs[batch_idx] row_k = tile_k if cutlass.const_expr( - cfg.mask_type == MaskType.CAUSAL.value and kernel.seq_len_q > 1 + not cfg.use_kv_bounds + and cfg.mask_type == MaskType.CAUSAL.value + and kernel.seq_len_q > 1 ): _, logical_seq_len_q = query_batch_bounds( cu_seqlens_q, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/resources.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/resources.py index a9b9c9f62496..dd5c83c596e5 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/resources.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/resources.py @@ -1,3 +1,4 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -59,6 +60,7 @@ from ...tensor_map import transform_ragged_coords from ..helpers.sparse_indices import SparseKVIndices, SparseSourceTile +from ..helpers.kv_bounds import query_kv_bound from ..helpers.gather import ( gather4_cached, load_sparse_rows, @@ -478,7 +480,9 @@ def _make_work_tile_info( safe_b_idx, ) if cutlass.const_expr( - cfg.mask_type == MaskType.CAUSAL.value and self.logical_seq_len_q > 1 + not cfg.use_kv_bounds + and cfg.mask_type == MaskType.CAUSAL.value + and self.logical_seq_len_q > 1 ): # Split partitioning owns the mask-visible CTA domain. The final # physical row resolves to the tile's latest valid Q @@ -2118,6 +2122,7 @@ class TmemSResource(HighThroughputMlaResource): page_offsets: Any = None store_softmax_stats: cutlass.Constexpr[bool] = False cache_seqs: Any = None # per-batch valid K length + kv_bounds: Any = None # per-logical-query prefix length cu_seqlens_q: Any = None # cumulative compact-Q offsets, or None for fixed Q split_kv: Any = None # per-work-tile split count logical_num_heads_q: cutlass.Constexpr[int] = 128 @@ -2490,9 +2495,27 @@ def _load_s_impl( k_index = work_tile.k_index_base + Int32(stage_info.loop_offset) tile_offset_k = k_index * Int32(cfg.mma_qk_tiler[1]) needs_row_causal_mask = cutlass.const_expr( - cfg.mask_type == MaskType.CAUSAL.value and self.logical_seq_len_q > 1 + not cfg.use_kv_bounds + and cfg.mask_type == MaskType.CAUSAL.value + and self.logical_seq_len_q > 1 ) - if cutlass.const_expr(needs_row_causal_mask): + row_kv_bound = Int32(0) + if cutlass.const_expr(cfg.use_kv_bounds): + row_in_tile = self.cta_rank * Int32( + cfg.mma_qk_tiler[0] // cfg.num_mma_ctas + ) + (local_tidx & Int32(EPILOGUE_ROW_MASK)) + flat_query_row = ( + Int32(work_tile.tile_idx[1]) * Int32(cfg.mma_qk_tiler[0]) + row_in_tile + ) + row_kv_bound = query_kv_bound( + self.kv_bounds, + self.cu_seqlens_q, + Int32(work_tile.tile_idx[2]), + flat_query_row // Int32(self.logical_num_heads_q), + self.logical_seq_len_q, + ) + group_needs_mask = True + elif cutlass.const_expr(needs_row_causal_mask): batch_idx = Int32(work_tile.tile_idx[2]) _, logical_seq_len_q = query_batch_bounds( self.cu_seqlens_q, @@ -2605,7 +2628,9 @@ def _load_s_impl( ) << EPILOGUE_COLUMN_GROUP_SHIFT for i in cutlass.range_constexpr(64): token_idx = tile_offset_k + tidx_col + Int32(i) - if cutlass.const_expr(needs_row_causal_mask): + if cutlass.const_expr(cfg.use_kv_bounds): + token_is_visible = token_idx < row_kv_bound + elif cutlass.const_expr(needs_row_causal_mask): mask_flat_row = Int32(self.logical_num_heads_q) * ( token_idx - K + logical_seq_len_q ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/config.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/config.py index 330430d7d98b..ad5d052233cd 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/config.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/config.py @@ -1,3 +1,4 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -98,6 +99,7 @@ class MlaConfig: # Causal is bottom-right aligned for speculative decode. Dense keeps only # the ordinary per-batch KV-tail predicate. mask_type: str = MaskType.CAUSAL.value + use_kv_bounds: bool = False head_dim_qk: int = 576 head_dim_v: int = 512 latent_dim: int = 512 diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/kernel.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/kernel.py index d515cb596f38..e93bd0e08b5d 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/kernel.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/kernel.py @@ -1,3 +1,4 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -466,6 +467,7 @@ def build_throughput_latency_mla_task_manager( page_offsets=None, cache_seqs=None, cu_seqlens_q=None, + kv_bounds=None, head_idx=None, batch_idx=None, cta_idx_q=None, @@ -520,6 +522,7 @@ def build_throughput_latency_mla_task_manager( page_offsets=page_offsets, cache_seqs=cache_seqs, cu_seqlens_q=cu_seqlens_q, + kv_bounds=kv_bounds, head_idx=head_idx, batch_idx=batch_idx, cta_idx_q=cta_idx_q, @@ -741,6 +744,7 @@ def build_throughput_latency_mla_task_manager( ) tmem_s0 = TmemSResource( cfg=cfg, + kv_bounds=kv_bounds, pipeline_config=tmem_s_cfg, inst_id=0, scale_softmax_log2=scale_softmax_log2, @@ -756,6 +760,7 @@ def build_throughput_latency_mla_task_manager( ) tmem_s1 = TmemSResource( cfg=cfg, + kv_bounds=kv_bounds, pipeline_config=tmem_s1_cfg, inst_id=1, scale_softmax_log2=scale_softmax_log2, @@ -1167,6 +1172,7 @@ def _make_single_kv_pipe_task_graph( page_offsets=None, cache_seqs=None, cu_seqlens_q=None, + kv_bounds=None, head_idx=None, batch_idx=None, cta_idx_q=None, @@ -1368,6 +1374,7 @@ def _make_single_kv_pipe_task_graph( ) tmem_s = (TmemSResource if cfg.one_insts_kv_swap else TmemSKeepsResource)( cfg=cfg, + kv_bounds=kv_bounds, page_offsets=page_offsets, pipeline_config=tmem_s_cfg, inst_id=0, @@ -1693,6 +1700,7 @@ def __init__( explicit_split_kv: int | None = None, explicit_persistent: bool | None = None, mask_type: MaskType | str = MaskType.CAUSAL, + use_kv_bounds: bool = False, sparse_profile=None, finalize_output: bool = False, ): @@ -1723,6 +1731,7 @@ def __init__( self.explicit_split_kv = explicit_split_kv self.explicit_persistent = explicit_persistent self.mask_type = normalize_mask_type(mask_type) + self.use_kv_bounds = use_kv_bounds self.sparse_profile = sparse_profile self.finalize_output = finalize_output self.output_quant = None @@ -1866,6 +1875,7 @@ def _make_config(self): explicit_persistent=self.explicit_persistent, mask_type=self.mask_type, ) + cfg = dataclass_replace(cfg, use_kv_bounds=self.use_kv_bounds) if self.sparse_profile is not None: cfg = configure_sparse_mla( cfg, @@ -1952,6 +1962,7 @@ def __call__( stream: object, output_quant_args=None, softmax_stats: cute.Tensor | None = None, + kv_bounds: cute.Tensor | None = None, ): """Execute throughput-latency 1CTA MLA with fixed or compact ragged Q.""" cfg = self._make_config() @@ -2251,6 +2262,10 @@ def __call__( cluster_shape = None if cutlass.const_expr(cfg.use_cluster_reduction == 1): cluster_shape = (cfg.num_ctas_per_seq_kv, 1, 1) + if cutlass.const_expr(self.use_kv_bounds): + from ..helpers.kv_bounds import initialize_empty_rows + + initialize_empty_rows(o, lse, softmax_stats, kv_bounds, stream) self.dense_kernel( tma_desc_q_latent, tma_desc_q_rope, @@ -2272,6 +2287,7 @@ def __call__( output_quant_args, softmax_stats, softmax_stats_scale, + kv_bounds, ).launch( grid=grid, block=[cfg.threads_per_cta, 1, 1], @@ -2392,6 +2408,7 @@ def dense_kernel( output_quant_args=None, softmax_stats: cute.Tensor | None = None, softmax_stats_scale=None, + kv_bounds: cute.Tensor | None = None, ): """Execute one flat-Q, batch, KV-split, and V head-dimension tile.""" cfg = self._make_config() @@ -2519,6 +2536,7 @@ def dense_kernel( page_offsets=page_offsets, cache_seqs=cache_seqs, cu_seqlens_q=cu_seqlens_q, + kv_bounds=kv_bounds, head_idx=head_idx, batch_idx=batch_idx, cta_idx_q=cta_idx_q, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_resources.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_resources.py index 590173465f44..26e7bd5408a8 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_resources.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_resources.py @@ -1,3 +1,4 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -600,9 +601,19 @@ def _producer_load_page_offsets( 3, cutlass.Int32, ) - prims.cp_async_shared_global( - smem_ptr, gmem_ptr, PAGE_OFFSET_BYTES, CP_ASYNC_CACHE_CA - ) + if cutlass.const_expr(cfg.use_kv_bounds): + if seq_len_kv == Int32(0): + # An empty rank need not own any valid physical pages. + # A negative page selects TMA's out-of-bounds zero fill. + self._smem_page_offsets.data_ptr(smem_offset).store(Int32(-1)) + else: + prims.cp_async_shared_global( + smem_ptr, gmem_ptr, PAGE_OFFSET_BYTES, CP_ASYNC_CACHE_CA + ) + else: + prims.cp_async_shared_global( + smem_ptr, gmem_ptr, PAGE_OFFSET_BYTES, CP_ASYNC_CACHE_CA + ) @producer_work @cute.jit @@ -1106,7 +1117,14 @@ def _producer_load_kv_stage( tile_idx * Int32(pages_per_tile) + Int32(page_frag), last_valid_page, ) - page_id = Int32(self.page_offsets[logical_page_idx, batch_idx]) + if cutlass.const_expr(cfg.use_kv_bounds): + page_id = Int32(-1) + if seq_len_kv > Int32(0): + page_id = Int32( + self.page_offsets[logical_page_idx, batch_idx] + ) + else: + page_id = Int32(self.page_offsets[logical_page_idx, batch_idx]) page_base = Int32(page_frag * first_page_elems) smem_page_offset = page_base if prims.elect_sync(): diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_s.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_s.py index 9831ab71e5d3..577f90d26ab6 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_s.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_s.py @@ -1,3 +1,4 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -119,6 +120,7 @@ class TmemSResource(MlaResource): p_ref: Optional[MemoryResource] = None global_ref: Optional[MemoryResource] = None cache_seqs: object = None + kv_bounds: object = None page_offsets: object = None head_idx: object = None batch_idx: object = None @@ -539,10 +541,11 @@ def update_softmax( should_apply_dense_mask = ( (seq_len_kv % Int32(cfg.tile_size_kv)) != Int32(0) ) or (next_tile_offset_k > seq_len_kv) - needs_row_causal_mask = cutlass.const_expr( - cfg.mask_type == MaskType.CAUSAL.value and cfg.logical_seq_len_q > 1 + needs_row_prefix_mask = cutlass.const_expr( + cfg.use_kv_bounds + or (cfg.mask_type == MaskType.CAUSAL.value and cfg.logical_seq_len_q > 1) ) - if cutlass.const_expr(needs_row_causal_mask): + if cutlass.const_expr(needs_row_prefix_mask): min_seq_len_kv = runtime_seq_len_kv_for_query_row( cfg, self.cache_seqs, @@ -550,6 +553,7 @@ def update_softmax( cta_idx_q, Int32(0), self.cu_seqlens_q, + self.kv_bounds, ) should_apply_dense_mask = should_apply_dense_mask or ( next_tile_offset_k > min_seq_len_kv @@ -563,15 +567,15 @@ def update_softmax( else: should_apply_dense_mask = True if should_apply_dense_mask: - # The CTA domain follows the latest logical query row in the flat - # tile. Earlier rows can have a narrower bottom-right causal limit. + # Each row's causal or explicit prefix can be narrower than the + # shared CTA domain used by the producer and split reduction. warp_idx = task_cache[_TASK_CACHE_WARP_IDX] lane_idx = task_cache[_TASK_CACHE_LANE_IDX] if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): local_col_base = (lane_idx >> Int32(4)) * Int32(cfg.tile_size_kv // 2) local_row_idx = warp_idx * Int32(16) + (lane_idx & Int32(0xF)) row_seq_len_kv = seq_len_kv - if cutlass.const_expr(needs_row_causal_mask): + if cutlass.const_expr(needs_row_prefix_mask): row_seq_len_kv = runtime_seq_len_kv_for_query_row( cfg, self.cache_seqs, @@ -579,6 +583,7 @@ def update_softmax( cta_idx_q, local_row_idx, self.cu_seqlens_q, + self.kv_bounds, ) for reg_idx in cutlass.range_constexpr(num_s_regs_per_thread(cfg)): token_idx = tile_offset_k + local_col_base + Int32(reg_idx) @@ -608,7 +613,7 @@ def update_softmax( row_in_tile_1 = row_in_tile_0 + Int32(1) seq_len_kv_0 = seq_len_kv seq_len_kv_1 = seq_len_kv - if cutlass.const_expr(needs_row_causal_mask): + if cutlass.const_expr(needs_row_prefix_mask): seq_len_kv_0 = runtime_seq_len_kv_for_query_row( cfg, self.cache_seqs, @@ -616,6 +621,7 @@ def update_softmax( cta_idx_q, row_in_tile_0, self.cu_seqlens_q, + self.kv_bounds, ) seq_len_kv_1 = runtime_seq_len_kv_for_query_row( cfg, @@ -624,6 +630,7 @@ def update_softmax( cta_idx_q, row_in_tile_1, self.cu_seqlens_q, + self.kv_bounds, ) s_base = repeat_idx * 4 s_second_panel_base = q_repeats * 4 + s_base diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py b/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py index 594d160b83aa..f0a6e9ddf697 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py @@ -1,4 +1,4 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -91,6 +91,7 @@ class _MLADecodeCompileSpec: split_kv: int kernel: Any = field(compare=False, hash=False, repr=False) store_softmax_stats: bool = False + use_kv_bounds: bool = False @dataclass(frozen=True) @@ -133,6 +134,7 @@ class _MLADecodePlanState: policy: tuple[tuple[str, object], ...] split_kv: int store_softmax_stats: bool = False + use_kv_bounds: bool = False @dataclass(frozen=True) @@ -412,6 +414,7 @@ def _validate_mla_run_metadata( max_kv_len: int, page_size: int, mask_type: str, + kv_bounds: Optional[torch.Tensor] = None, ) -> None: """Validate per-run request metadata against one static MLA plan. @@ -466,8 +469,10 @@ def _validate_mla_run_metadata( q_lengths = (max_seq_len_q,) * batch_size seq_lens_host = tuple(int(value) for value in seq_lens.tolist()) - if any(seq_len <= 0 for seq_len in seq_lens_host): + if kv_bounds is None and any(seq_len <= 0 for seq_len in seq_lens_host): raise ValueError("every runtime request must contain at least one KV token") + if kv_bounds is not None and any(seq_len < 0 for seq_len in seq_lens_host): + raise ValueError("runtime KV lengths must be nonnegative with kv_bounds") runtime_max_kv_len = max(seq_lens_host) if runtime_max_kv_len > max_kv_len: raise ValueError( @@ -488,7 +493,25 @@ def _validate_mla_run_metadata( f"K/V cache in [0, {runtime.num_physical_pages}); request " f"{request_idx} contains an invalid page ID" ) - if mask_type == "causal": + if kv_bounds is not None: + bounds_host = kv_bounds.tolist() + offset = 0 + for request_idx, (q_len, kv_len) in enumerate( + zip(q_lengths, seq_lens_host, strict=True) + ): + row_bounds = bounds_host[offset : offset + q_len] + if any(bound < 0 or bound > kv_len for bound in row_bounds): + raise ValueError( + "kv_bounds values must be in [0, seq_lens[request]]; " + f"request {request_idx} has K/V={kv_len}" + ) + if any( + left > right + for left, right in zip(row_bounds[:-1], row_bounds[1:], strict=True) + ): + raise ValueError("kv_bounds must be nondecreasing within each request") + offset += q_len + elif mask_type == "causal": for request_idx, (q_len, kv_len) in enumerate( zip(q_lengths, seq_lens_host, strict=True) ): @@ -691,6 +714,30 @@ def _validate_out( _validate_16byte_alignment(out, "out") +def _validate_mla_kv_bounds( + query: torch.Tensor, + kv_bounds: Optional[torch.Tensor], + use_kv_bounds: bool, + validate: bool, +) -> None: + """Enforce mode agreement and optionally check per-logical-query storage.""" + if (kv_bounds is not None) != use_kv_bounds: + raise ValueError("kv_bounds must be supplied exactly when use_kv_bounds=True") + if validate and kv_bounds is not None: + _validate_int32_cuda_tensor( + kv_bounds, "kv_bounds", ndim=1, require_16byte_alignment=False + ) + if kv_bounds.device != query.device: + raise ValueError("kv_bounds must be on the query device") + expected = math.prod(query.shape[:-2]) + if kv_bounds.numel() != expected: + raise ValueError( + f"kv_bounds must have shape [{expected}], one bound per query token" + ) + if kv_bounds.data_ptr() % 4: + raise ValueError("kv_bounds must be 4-byte aligned") + + def _validate_mla_softmax_stats( query: torch.Tensor, softmax_stats: Optional[torch.Tensor], @@ -767,6 +814,7 @@ def _resolve_mla_decode_launch_spec( mask_type: str, seq_len_q: int = 1, store_softmax_stats: bool = False, + use_kv_bounds: bool = False, ): """Resolve and cache MLA policy/workspace without compiling.""" @@ -897,6 +945,7 @@ def _resolve_mla_decode_launch_spec( explicit_split_kv=None, explicit_persistent=None, mask_type=mask_type, + use_kv_bounds=use_kv_bounds, ) final_cfg = kernel._make_config() split_kv = int(final_cfg.num_ctas_per_seq_kv) @@ -980,6 +1029,7 @@ def _resolve_mla_decode_launch_spec( seq_len_q=seq_len_q, batch_size=batch_size, mask_type=mask_type, + use_kv_bounds=use_kv_bounds, ) workspace_size = compute_2cta_workspace_size( tile_size_q=int(launch_shape.tile_size_q), @@ -1056,6 +1106,7 @@ def _make_mla_decode_compile_spec( max_seq_len_q: int, packed_query: bool, store_softmax_stats: bool = False, + use_kv_bounds: bool = False, ) -> _MLADecodeCompileSpec: """Keep policy resolution plan-specific and JIT identity batch-free.""" @@ -1074,6 +1125,7 @@ def _make_mla_decode_compile_spec( split_kv=launch_spec.split_kv, kernel=launch_spec.kernel, store_softmax_stats=store_softmax_stats, + use_kv_bounds=use_kv_bounds, ) @@ -1270,6 +1322,13 @@ def sparse_indices(capacity): ) output_quant_args = _fake_output_quant_args(getattr(kernel, "output_quant", None)) + bounds_fake = None + if compile_spec.use_kv_bounds: + bounds_fake = cute.runtime.make_fake_compact_tensor( + cutlass.Int32, + (runtime_total_q if packed_query else batch_size * max_seq_len_q,), + assumed_align=4, + ) # Task objects carry loop-local state through generated control flow, so # select the public staged frontend for this compilation. @@ -1293,6 +1352,7 @@ def compile_kernel(): stream_fake, output_quant_args, stats_fake, + bounds_fake, options=_COMPILE_OPTIONS, ) @@ -1324,6 +1384,7 @@ def get_prims_ts_batch_mla_decode_workspace_size( mask_type: Literal["dense", "causal"] = "causal", device=None, store_softmax_stats: bool = False, + use_kv_bounds: bool = False, ) -> int: """Return caller-workspace bytes for one automatic MLA policy. @@ -1337,6 +1398,8 @@ def get_prims_ts_batch_mla_decode_workspace_size( scratch and the internal FP32 LSE tensor. With ``store_softmax_stats=True``, it also includes split maxima and sums to export actual maximum/denominator statistics; the public statistics buffer is caller-owned and excluded. + ``use_kv_bounds=True`` selects per-query prefix masking without adding + scratch; supply the same flag to the plan that consumes this workspace. Allocate a contiguous ``torch.int8`` or ``torch.uint8`` CUDA buffer; MLA does not require its contents to be initialized before first use. @@ -1344,6 +1407,8 @@ def get_prims_ts_batch_mla_decode_workspace_size( if not isinstance(store_softmax_stats, bool): raise TypeError("store_softmax_stats must be a bool") + if not isinstance(use_kv_bounds, bool): + raise TypeError("use_kv_bounds must be a bool") batch_size = _validate_positive_int(batch_size, "batch_size") num_heads = _validate_positive_int(num_heads, "num_heads") _validate_mla_dims(kv_lora_rank, qk_rope_head_dim) @@ -1381,6 +1446,7 @@ def get_prims_ts_batch_mla_decode_workspace_size( mask_type, max_seq_len_q, store_softmax_stats, + use_kv_bounds, ) return _make_mla_workspace_layout( spec.kernel_workspace_bytes, batch_size, num_heads, max_seq_len_q @@ -1491,6 +1557,7 @@ def _launch_mla_decode( workspace: _MLAWorkspaceViews, compiled: Callable[..., object], softmax_stats: Optional[torch.Tensor] = None, + kv_bounds: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Form the dimension-first views and launch one compiled MLA kernel.""" @@ -1533,9 +1600,10 @@ def _launch_mla_decode( runtime.bmm1_scale, runtime.bmm2_scale, # Reloaded persistent-cache kernels expose a positional TVM-FFI entry - # without Python defaults. Keep both optional arguments in the ABI. + # without Python defaults. Keep all optional arguments in the ABI. runtime.output_quant_args, None if softmax_stats is None else softmax_stats.view(-1, 2), + kv_bounds, ) return runtime.out @@ -1566,6 +1634,7 @@ def plan( o_data_type: torch.dtype, mask_type: Literal["dense", "causal"] = "causal", store_softmax_stats: bool = False, + use_kv_bounds: bool = False, workspace_buffer: Optional[torch.Tensor] = None, validate: bool = True, ) -> None: @@ -1608,6 +1677,10 @@ def plan( store_softmax_stats : bool Compile final FP32 maximum/denominator export. Supply its caller-owned buffer to every run when enabled. Defaults to ``False``. + use_kv_bounds : bool + Compile explicit per-query KV prefix masking. When enabled, every + run requires ``kv_bounds`` and may use zero runtime K/V lengths. + Bounds override ``mask_type``; the static K/V capacity stays positive. workspace_buffer : torch.Tensor, optional Caller-owned contiguous int8 or uint8 scratch on ``device``. It must be 32-byte aligned and large enough for the selected plan. @@ -1620,6 +1693,8 @@ def plan( if not isinstance(store_softmax_stats, bool): raise TypeError("store_softmax_stats must be a bool") + if not isinstance(use_kv_bounds, bool): + raise TypeError("use_kv_bounds must be a bool") if not isinstance(validate, bool): raise TypeError("validate must be a bool") if validate: @@ -1660,6 +1735,7 @@ def plan( mask_type, max_seq_len_q, store_softmax_stats, + use_kv_bounds, ) spec = _resolve_mla_decode_launch_spec(*spec_key) compile_spec = _make_mla_decode_compile_spec( @@ -1674,6 +1750,7 @@ def plan( max_seq_len_q=max_seq_len_q, packed_query=packed_query, store_softmax_stats=store_softmax_stats, + use_kv_bounds=use_kv_bounds, ) policy = spec.policy workspace_layout = _make_mla_workspace_layout( @@ -1716,6 +1793,7 @@ def plan( policy=policy, split_kv=int(dict(policy)["split_kv"]), store_softmax_stats=store_softmax_stats, + use_kv_bounds=use_kv_bounds, ) @flashinfer_experimental_api @@ -1727,6 +1805,7 @@ def run( seq_lens: torch.Tensor, *, qo_indptr: Optional[torch.Tensor] = None, + kv_bounds: Optional[torch.Tensor] = None, bmm1_scale: float = 1.0, bmm2_scale: float = 1.0, out: Optional[torch.Tensor] = None, @@ -1759,6 +1838,14 @@ def run( Runtime K/V lengths with one element per request. qo_indptr : torch.Tensor, optional Runtime cumulative query offsets for a packed-query plan. + kv_bounds : torch.Tensor, optional + Contiguous CUDA int32 prefix counts, flat ``[B * SQ]`` for fixed Q + or ``[total_q]`` for packed Q, shared by every head of a query token. + Required exactly when ``plan(use_kv_bounds=True)``, even with + ``validate=False``. Counts must be nondecreasing within each request + and lie in ``[0, seq_lens[b]]``. They replace dense/causal visibility. + A zero prefix returns zero O and FP32 statistics ``(-inf, 0)``. + Structure and values are caller preconditions with ``validate=False``. bmm1_scale, bmm2_scale : float QK and value/output scaling factors. out : torch.Tensor, optional @@ -1787,6 +1874,7 @@ def run( _validate_mla_softmax_stats( query, softmax_stats, state.store_softmax_stats, validate ) + _validate_mla_kv_bounds(query, kv_bounds, state.use_kv_bounds, validate) runtime_qo_indptr = qo_indptr if state.packed_query else None runtime = _prepare_mla_runtime( query, @@ -1820,6 +1908,7 @@ def run( max_kv_len=state.max_kv_len, page_size=state.page_size, mask_type=state.mask_type, + kv_bounds=kv_bounds, ) return _launch_mla_decode( runtime, @@ -1832,6 +1921,7 @@ def run( workspace=state.workspace_views, compiled=state.compiled, softmax_stats=softmax_stats, + kv_bounds=kv_bounds, ) @@ -1855,6 +1945,7 @@ def batch_mla_decode_with_paged_kv_cache( workspace_buffer: Optional[torch.Tensor] = None, validate: bool = True, store_softmax_stats: bool = False, + kv_bounds: Optional[torch.Tensor] = None, softmax_stats: Optional[torch.Tensor] = None, ) -> torch.Tensor: """One-shot convenience wrapper for fixed or packed-query MLA decode. @@ -1895,6 +1986,11 @@ def batch_mla_decode_with_paged_kv_cache( Output dtype. store_softmax_stats : bool Compile statistics export; defaults to ``False``. + kv_bounds : torch.Tensor, optional + Per-query CUDA int32 prefix counts with the same contract as + :meth:`BatchMLADecodePagedTSWrapper.run`. Supplying this tensor enables + bounded mode and permits zero runtime K/V lengths. If all lengths are + zero, an omitted static ``max_kv_len`` is conservatively set to one. softmax_stats : torch.Tensor, optional Caller-owned FP32 maximum/denominator buffer with shape ``[*query.shape[:-1], 2]``; required exactly when statistics are enabled. @@ -1918,8 +2014,9 @@ def batch_mla_decode_with_paged_kv_cache( ``max_seq_len_q``. Warm this exact topology outside capture first. Retain stable tensor storage and mutate metadata only between completed launches/replays. All live K/V - lengths must be positive and within the static bound, active page IDs must - index the cache, and causal per-request Q lengths must not exceed K/V lengths. + lengths must be positive (nonnegative with ``kv_bounds``) and within the + static bound, and active page IDs must index the cache. Without ``kv_bounds``, + causal per-request Q lengths must not exceed K/V lengths. Packed offsets must start at zero, end at the query token count, and have nonnegative deltas within the Q bound. Scratch/output must not alias any inputs or each other. No metadata is copied to the host on the trusted @@ -1979,6 +2076,8 @@ def batch_mla_decode_with_paged_kv_cache( max_seq_len_q = int(query.shape[1]) if max_kv_len is None: max_kv_len = int(seq_lens.max().item()) + if kv_bounds is not None: + max_kv_len = max(max_kv_len, 1) assert max_seq_len_q is not None if validate and out is not None: _validate_out( @@ -2006,6 +2105,7 @@ def batch_mla_decode_with_paged_kv_cache( mask_type=mask_type, device=query.device, store_softmax_stats=store_softmax_stats, + use_kv_bounds=kv_bounds is not None, ) workspace_buffer = torch.empty( workspace_bytes, dtype=torch.int8, device=query.device @@ -2029,6 +2129,7 @@ def batch_mla_decode_with_paged_kv_cache( workspace_buffer=workspace_buffer, validate=validate, store_softmax_stats=store_softmax_stats, + use_kv_bounds=kv_bounds is not None, ) return wrapper.run( query, @@ -2041,6 +2142,7 @@ def batch_mla_decode_with_paged_kv_cache( out=out, validate=validate, softmax_stats=softmax_stats, + kv_bounds=kv_bounds, ) From a835a737de47f4ae16ce13706dbbc7426594c68e Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Thu, 29 Oct 2026 02:05:24 +0000 Subject: [PATCH 07/19] [None][feat] Integrate linear MTP and Helix bounds in PrimTS MLA Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../_torch/attention/backends/fmha/manager.py | 13 +- .../attention/backends/fmha/prims_ts.py | 111 +++++-- .../_torch/attention/backends/fmha/utils.py | 21 ++ .../_torch/attention/test_fmha_manager.py | 102 ++++++- .../_torch/attention/test_prims_ts_fmha.py | 289 ++++++++++++++++-- 5 files changed, 485 insertions(+), 51 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/manager.py b/tensorrt_llm/_torch/attention/backends/fmha/manager.py index e17a087c7e54..10e5bf3cbb6e 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/manager.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/manager.py @@ -39,6 +39,7 @@ from .interface import Fmha, FmhaPhase from .phased import PhasedFmha from .registry import get_enabled_fmha_lib_classes +from .utils import has_uniform_generation_query_lengths if TYPE_CHECKING: from tensorrt_llm._torch.attention.backends.trtllm import ( @@ -132,6 +133,7 @@ class _FmhaCacheKey(NamedTuple): context_batch_size: int generation_batch_size: int generation_seq_len_q: int + generation_query_is_uniform: bool attention_mask_type: AttentionMaskType use_spec_decoding: bool has_block_sparse_inputs: bool @@ -371,8 +373,10 @@ def _make_cache_key( Batch size and generation Q length use ceiling grid buckets based on TRTLLM-Gen JIT warmup candidates. Each candidate and its predecessor are separate grid points so a support boundary at the candidate cannot - alias with the interval below it. Generation Q lengths are uniform in - the executor, including padded speculative-decoding batches. + alias with the interval below it. Eager speculative batches may have + ragged generation rows with the same total Q extent as uniform rows. + Keep their support decisions separate using the host query lengths; + CUDA-graph batches retain their padded uniform widths. The batch fields describe the complete source scheduler batch, even when ``q`` contains only one compacted phase. Ordinary attention @@ -403,10 +407,14 @@ def _make_cache_key( metadata.num_generations, _FMHA_CACHE_BATCH_SIZE_GRID ) generation_seq_len_q = 0 + generation_query_is_uniform = True if attention_input_type != AttentionInputType.context_only and metadata.num_generations > 0: generation_num_tokens = q.shape[0] if attention_input_type == AttentionInputType.mixed: generation_num_tokens -= metadata.num_ctx_tokens + generation_query_is_uniform = has_uniform_generation_query_lengths( + metadata, generation_num_tokens + ) generation_seq_len_q = generation_num_tokens // metadata.num_generations generation_seq_len_q = _normalize_fmha_cache_grid_value( generation_seq_len_q, _FMHA_CACHE_SEQ_LEN_Q_GRID @@ -417,6 +425,7 @@ def _make_cache_key( context_batch_size=context_batch_size, generation_batch_size=generation_batch_size, generation_seq_len_q=generation_seq_len_q, + generation_query_is_uniform=generation_query_is_uniform, attention_mask_type=attention_mask_type, use_spec_decoding=metadata.use_spec_decoding, has_block_sparse_inputs=block_sparse_inputs is not None, diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py index dfdb58fe5c04..e91d6e9a5adc 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py @@ -40,7 +40,11 @@ from .interface import FmhaPhase from .phased import FmhaParams, PhasedFmha -from .utils import get_kv_page_offset, get_multi_processor_count_for_device +from .utils import ( + get_kv_page_offset, + get_multi_processor_count_for_device, + has_uniform_generation_query_lengths, +) if TYPE_CHECKING: from tensorrt_llm._torch.attention.backends.prims_ts.context import ( @@ -106,16 +110,15 @@ def get_paged_kv_storage_unsupported_reason( def get_paged_kv_policy_unsupported_reason( attn: "TrtllmAttention", metadata: "TrtllmAttentionMetadata", + *, + allow_linear_mla: bool = False, ) -> Optional[str]: """Return why the request's decoding policy is outside the fixed page-table envelope.""" if metadata.beam_width != 1: return "beam search is not supported." - if ( - metadata.is_spec_decoding_enabled - or metadata.use_spec_decoding - or metadata.is_spec_dec_tree - or metadata.is_spec_dec_dynamic_tree - ): + if metadata.is_spec_dec_tree or metadata.is_spec_dec_dynamic_tree: + return "tree speculative decoding is not supported." + if not allow_linear_mla and (metadata.is_spec_decoding_enabled or metadata.use_spec_decoding): return "speculative decoding is not supported by the initial adapter." position_embedding_type = int(attn.position_embedding_type) if position_embedding_type in (4, 5, 6, 7, 10): @@ -138,9 +141,11 @@ def get_paged_kv_policy_unsupported_reason( def get_attention_feature_unsupported_reason( metadata: "TrtllmAttentionMetadata", forward_args: "AttentionForwardArgs", + *, + allow_helix: bool = False, ) -> Optional[str]: """Return which optional attention feature the fused-kernel adapters do not implement.""" - if metadata.helix_position_offsets is not None: + if metadata.helix_position_offsets is not None and not allow_helix: return "Helix parallelism is not supported." if forward_args.relative_attention_bias is not None: return "relative attention bias is not supported." @@ -171,13 +176,12 @@ def __init__(self, attn: "TrtllmAttention") -> None: # the manager itself. self._kv_page_offset_cache: dict[tuple[int, int], int] = {} self._multi_processor_count: Optional[int] = None - # Every other plan attribute is fixed by this layer/model instance. - # Batch size is the only execution profile that needs its own wrapper. + # Standard-attention decode admits only single-token generation. self._context_wrappers: dict[int, "BatchPrefillPagedTSWrapper"] = {} # Cache-free context plans, keyed by (batch, seq_len, q/v/out dtype, mask). self._contiguous_context_wrappers: dict[tuple[Any, ...], "BatchPrefillTSWrapper"] = {} self._decode_wrappers: dict[int, "BatchDecodePagedTSWrapper"] = {} - self._mla_decode_wrappers: dict[int, "BatchMLADecodePagedTSWrapper"] = {} + self._mla_decode_wrappers: dict[tuple[int, int], "BatchMLADecodePagedTSWrapper"] = {} self._mla_context_wrappers: dict[ tuple[int, int, int, str, bool], "BatchPrefillTSWrapper" ] = {} @@ -321,7 +325,12 @@ def _is_supported_with_reason( return False, "sparse attention metadata is not supported." if meta.num_sparse_topk > 0: return False, "sparse attention metadata is not supported." - feature_reason = get_attention_feature_unsupported_reason(meta, fwd) + is_mla_generation = ( + attn.is_mla_enable and fwd.attention_input_type == AttentionInputType.generation_only + ) + feature_reason = get_attention_feature_unsupported_reason( + meta, fwd, allow_helix=is_mla_generation + ) if feature_reason is not None: return False, feature_reason if fwd.sage_attn_qk_dtype is not None: @@ -338,7 +347,9 @@ 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_linear_mla=is_mla_generation + ) if policy_reason is not None: return False, policy_reason is_mla = attn.is_mla_enable @@ -384,6 +395,8 @@ def _is_supported_with_reason( 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 requires an FP8 E4M3 KV cache." + if is_fp8_mla and getattr(meta.kv_cache_manager, "mla_kv_cache_residual_dim", 0): + return False, "FP8 MLA requires a full FP8 KV cache without a high-precision residual." if not is_fp8_mla 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: @@ -475,8 +488,45 @@ def _is_supported_with_reason( num_gen_tokens <= 0 or num_generations <= 0 or num_gen_tokens % num_generations != 0 ): return False, "generation tokens must be uniformly divisible across requests." - if has_generation and num_gen_tokens != num_generations: + if has_generation and not is_mla and num_gen_tokens != num_generations: return False, "only single-token generation is supported by the initial adapter." + is_helix = meta.helix_position_offsets is not None + if is_mla_generation: + seq_len_q = num_gen_tokens // num_generations + if seq_len_q > 1: + if mask_type != AttentionMaskType.causal: + return False, "MLA multi-token generation requires a causal mask." + if not has_uniform_generation_query_lengths(meta, num_gen_tokens): + return False, "MLA generation requires uniform per-request query lengths." + stats = fwd.softmax_stats_tensor + if is_helix and stats is None: + return False, "Helix MLA generation requires softmax statistics." + if stats is not None and ( + stats.dtype != torch.float32 + or stats.device != q.device + or not stats.is_contiguous() + or stats.shape != (num_gen_tokens, attn.num_heads, 2) + ): + return False, "MLA softmax statistics must be contiguous FP32 [tokens, heads, 2]." + if is_helix: + if mask_type != AttentionMaskType.causal: + return False, "Helix MLA generation requires a causal mask." + if meta._helix_spec_tokens_valid: + bounds = meta.helix_kv_bounds + if ( + bounds is None + or bounds.dtype != torch.int32 + or bounds.device != q.device + or bounds.ndim != 1 + or not bounds.is_contiguous() + or bounds.numel() < num_gen_tokens + ): + return False, "Helix MLA requires contiguous int32 per-token KV bounds." + elif seq_len_q != 1: + return ( + False, + "Helix MLA multi-token generation requires valid per-token KV bounds.", + ) host_kv_lens = meta.kv_lens_runtime if host_kv_lens is None or host_kv_lens.numel() < num_contexts + num_generations: @@ -488,7 +538,7 @@ def _is_supported_with_reason( and mask_type == AttentionMaskType.padding and fwd.softmax_stats_tensor is not None ) - min_kv_length = 0 if allows_empty_chunk else 1 + min_kv_length = 0 if allows_empty_chunk or (is_mla_generation and is_helix) else 1 if active_kv_lens.numel() == 0 or int(active_kv_lens.min()) < min_kv_length: return False, "every active request must contain at least one KV token." max_kv_length = int(active_kv_lens.max()) @@ -830,10 +880,15 @@ def _get_or_plan_mla_decode_wrapper( kv_dtype: torch.dtype, output_dtype: torch.dtype, mask_type: str, + store_softmax_stats: bool = False, + use_kv_bounds: bool = False, ) -> "BatchMLADecodePagedTSWrapper": if batch_size <= 0: raise RuntimeError("PrimTS MLA decode requires a positive batch size.") - wrapper = self._mla_decode_wrappers.get(batch_size) + # Dense MLA instances fix the cache capacity, mask, dtypes and Helix + # stats/bounds modes. Only batch size and verify width vary per request. + key = (batch_size, max_seq_len_q) + wrapper = self._mla_decode_wrappers.get(key) if wrapper is not None: return wrapper if torch.cuda.is_current_stream_capturing(): @@ -857,9 +912,11 @@ def _get_or_plan_mla_decode_wrapper( kv_data_type=kv_dtype, o_data_type=output_dtype, mask_type=mask_type, + store_softmax_stats=store_softmax_stats, + use_kv_bounds=use_kv_bounds, workspace_buffer=workspace_buffer, ) - self._mla_decode_wrappers[batch_size] = wrapper + self._mla_decode_wrappers[key] = wrapper return wrapper def forward( @@ -1064,6 +1121,7 @@ def prepare_workspace( kv_dtype=kernel_dtype, out_dtype=forward_args.output.dtype, mask_type=mask_type, + store_softmax_stats=forward_args.softmax_stats_tensor is not None, device=q.device, ) else: @@ -1654,6 +1712,21 @@ def run_mla_generation(self, params: FmhaParams) -> None: params.sequence_lengths, batch_size, ) + kv_bounds = None + if meta.helix_position_offsets is not None: + if meta._helix_spec_tokens_valid: + # Helix token buffers are generation-relative, even when the + # batch's request metadata has a context prefix. + kv_bounds = meta.helix_kv_bounds[: batch_size * seq_len_q] + elif seq_len_q == 1: + # The ordinary Helix length is already rank-local and corrected + # on device. Reuse its view without allocating a graph buffer. + kv_bounds = seq_lens + else: + raise RuntimeError("Helix MLA multi-token generation requires per-token KV bounds.") + stats = params.fwd.softmax_stats_tensor + if stats is not None: + stats = stats.view(batch_size, seq_len_q, attn.num_heads, 2) caller_workspace = params.workspace.reshape(-1).view(torch.uint8) wrapper = self._get_or_plan_mla_decode_wrapper( caller_workspace, @@ -1668,6 +1741,8 @@ def run_mla_generation(self, params: FmhaParams) -> None: kv_dtype=kv_cache.dtype, output_dtype=output.dtype, mask_type=mask_type, + store_softmax_stats=stats is not None, + use_kv_bounds=kv_bounds is not None, ) wrapper.run( query, @@ -1677,6 +1752,8 @@ def run_mla_generation(self, params: FmhaParams) -> None: out=output, bmm1_scale=bmm1_scale, bmm2_scale=bmm2_scale, + softmax_stats=stats, + kv_bounds=kv_bounds, validate=False, ) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/utils.py b/tensorrt_llm/_torch/attention/backends/fmha/utils.py index 7cce8c3db6da..33b4a8b849a1 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/utils.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/utils.py @@ -31,6 +31,27 @@ ) +def has_uniform_generation_query_lengths( + metadata: "TrtllmAttentionMetadata", num_gen_tokens: int +) -> bool: + """Check active generation query widths using only host sequence metadata.""" + num_generations = metadata.num_generations + query_lens = metadata.seq_lens + if ( + num_generations <= 0 + or num_gen_tokens <= 0 + or num_gen_tokens % num_generations + or query_lens is None + or query_lens.device.type != "cpu" + or query_lens.ndim != 1 + or query_lens.numel() < metadata.num_contexts + num_generations + ): + return False + # Original prompt lengths do not describe this step's verification rows. + query_lens = query_lens[metadata.num_contexts : metadata.num_contexts + num_generations] + return bool((query_lens == num_gen_tokens // num_generations).all()) + + def get_kv_page_offset( attn: "TrtllmAttention", metadata: "TrtllmAttentionMetadata", diff --git a/tests/unittest/_torch/attention/test_fmha_manager.py b/tests/unittest/_torch/attention/test_fmha_manager.py index 500b5c5e337d..56743ebf5bf0 100644 --- a/tests/unittest/_torch/attention/test_fmha_manager.py +++ b/tests/unittest/_torch/attention/test_fmha_manager.py @@ -16,7 +16,7 @@ import gc import weakref from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import Mock, patch import pytest import torch @@ -26,6 +26,7 @@ from tensorrt_llm._torch.attention.backends.fmha.combined import CombinedFmha from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase from tensorrt_llm._torch.attention.backends.fmha.manager import FmhaManager +from tensorrt_llm._torch.attention.backends.fmha.utils import has_uniform_generation_query_lengths from tensorrt_llm._torch.attention.backends.interface import ( AttentionForwardArgs, AttentionInputType, @@ -55,6 +56,7 @@ def _make_metadata( num_ctx_tokens: int = 0, use_spec_decoding: bool = False, num_sparse_topk: int = 0, + seq_lens: torch.Tensor | None = None, ) -> SimpleNamespace: return SimpleNamespace( num_contexts=num_contexts, @@ -62,6 +64,7 @@ def _make_metadata( num_ctx_tokens=num_ctx_tokens, use_spec_decoding=use_spec_decoding, num_sparse_topk=num_sparse_topk, + seq_lens=seq_lens, ) @@ -232,6 +235,103 @@ def test_select_non_mla_fmha_preserves_registry_order() -> None: assert events == [("support", "non-phased", None)] +@pytest.mark.parametrize("uniform_first", [False, True]) +@pytest.mark.parametrize( + "input_type", [AttentionInputType.generation_only, AttentionInputType.mixed] +) +def test_fmha_cache_separates_uniform_and_ragged_generation( + uniform_first: bool, input_type: AttentionInputType +) -> None: + events: list[tuple] = [] + attn, manager = _make_manager() + attn.is_mla_enable = input_type == AttentionInputType.generation_only + num_tokens = 12 if attn.is_mla_enable else 15 + uniform_only = FakeFmha( + attn, + "uniform-only", + events, + request_support_predicate=lambda q, metadata: has_uniform_generation_query_lengths( + metadata, q.shape[0] if attn.is_mla_enable else q.shape[0] - metadata.num_ctx_tokens + ), + ) + fallback = FakeFmha(attn, "fallback", events) + manager.fmha_libs = [uniform_only, fallback] + metadata = _make_metadata( + num_contexts=1, + num_generations=2, + num_ctx_tokens=3, + seq_lens=torch.tensor([3, 6, 6], dtype=torch.int32), + ) + q = torch.empty((num_tokens, 4)) + forward_args = AttentionForwardArgs(attention_input_type=input_type) + order = (True, False) if uniform_first else (False, True) + selected = {} + keys = {} + with patch.object(fmha_manager, "_is_fmha_cache_enabled", return_value=True): + for uniform in order + order: + metadata.seq_lens[1:].copy_(torch.tensor([6, 6] if uniform else [5, 7])) + keys[uniform] = manager._make_cache_key(q, metadata, forward_args) + selected[uniform] = manager.select(attn, q, None, None, metadata, forward_args) + assert selected == {True: uniform_only, False: fallback} + assert keys[True]._replace(generation_query_is_uniform=False) == keys[False] + assert len(manager._cache) == 2 + assert events.count(("support", "uniform-only", None)) == 2 + assert events.count(("support", "fallback", None)) == 1 + + +def test_context_cache_key_does_not_inspect_generation_query_lengths() -> None: + attn, manager = _make_manager() + attn.is_mla_enable = True + metadata = _make_metadata(num_contexts=1, num_generations=2, num_ctx_tokens=3) + q = torch.empty((3, 4)) + forward_args = AttentionForwardArgs(attention_input_type=AttentionInputType.context_only) + with patch.object( + fmha_manager, + "has_uniform_generation_query_lengths", + side_effect=AssertionError("context-only must not inspect generation rows"), + ): + metadata.seq_lens = torch.tensor([3, 6, 6]) + uniform = manager._make_cache_key(q, metadata, forward_args) + metadata.seq_lens = torch.tensor([3, 5, 7]) + ragged = manager._make_cache_key(q, metadata, forward_args) + assert uniform == ragged + assert uniform.generation_seq_len_q == 0 + + +@pytest.mark.parametrize( + "query_lens,num_gen_tokens,expected", + [ + (torch.tensor([3, 6, 6]), 12, True), + (torch.tensor([3, 5, 7]), 12, False), + (torch.tensor([3, 6, 6]), 11, False), + (torch.tensor([3, 6, 6]), 0, False), + (torch.tensor([3, 6]), 12, False), + (torch.tensor([[3, 6, 6]]), 12, False), + (torch.empty(3, device="meta", dtype=torch.int32), 12, False), + (None, 12, False), + ], +) +def test_generation_query_uniformity_uses_only_cpu_metadata( + query_lens: torch.Tensor | None, + num_gen_tokens: int, + expected: bool, +) -> None: + metadata = _make_metadata( + num_contexts=1, + num_generations=2, + num_ctx_tokens=3, + seq_lens=query_lens, + ) + device_read = Mock(side_effect=AssertionError("unexpected device read")) + with ( + patch.object(torch.Tensor, "cpu", device_read), + patch.object(torch.Tensor, "item", device_read), + patch.object(torch.Tensor, "to", device_read), + ): + assert has_uniform_generation_query_lengths(metadata, num_gen_tokens) is expected + device_read.assert_not_called() + + @pytest.mark.parametrize(("lower_q_length", "upper_q_length"), [(3, 4), (7, 8)]) @pytest.mark.parametrize("upper_first", [False, True]) def test_fmha_cache_separates_generation_q_boundaries( diff --git a/tests/unittest/_torch/attention/test_prims_ts_fmha.py b/tests/unittest/_torch/attention/test_prims_ts_fmha.py index 008ae7bb8f22..78dac7347741 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_fmha.py +++ b/tests/unittest/_torch/attention/test_prims_ts_fmha.py @@ -149,6 +149,13 @@ def _support_result( beam_width: int = 1, use_spec_decoding: bool = False, is_spec_dec_tree: bool = False, + query_width: int = 1, + query_lens: list[int] | None = None, + helix: bool = False, + helix_bounds_valid: bool = False, + kv_lens_override: list[int] | None = None, + metadata_overrides: dict | None = None, + forward_overrides: dict | None = None, has_attention_sinks: bool = False, has_relative_attention_bias: bool = False, has_sparse_attention: bool = False, @@ -156,6 +163,7 @@ def _support_result( position_embedding_type: int = 0, kv_lora_rank: int | None = None, qk_rope_head_dim: int | None = None, + mla_kv_cache_residual_dim: int = 0, has_output: bool = True, attention_window_size: int = 128, attention_chunk_size: int = 0, @@ -190,11 +198,12 @@ def _support_result( q_width = attn.num_heads * head_dim if is_fused_qkv and not is_mla: q_width += 2 * attn.num_kv_heads * head_dim - q = _TensorSpec((4, q_width), dtype) + num_tokens = 4 * query_width + q = _TensorSpec((num_tokens, q_width), dtype) mla_context = is_mla and attention_input_type == AttentionInputType.context_only output_width = attn.num_heads * (128 if mla_context else 512 if is_mla else head_dim) forward_args = AttentionForwardArgs( - output=_TensorSpec((4, output_width), output_dtype) if has_output else None, + output=_TensorSpec((num_tokens, output_width), output_dtype) if has_output else None, attention_input_type=attention_input_type, attention_mask=PredefinedAttentionMask.CAUSAL, attention_window_size=attention_window_size, @@ -208,6 +217,12 @@ def _support_result( forward_args.mla_bmm2_scale = _TensorSpec((1,), torch.float32) if mla_fp8_inputs is not None: forward_args = replace(forward_args, **mla_fp8_inputs) + if helix: + forward_args.softmax_stats_tensor = _TensorSpec( + (num_tokens, attn.num_heads, 2), torch.float32 + ) + if forward_overrides is not None: + forward_args = replace(forward_args, **forward_overrides) if has_sparse_runtime_metadata: forward_args.sparse_runtime_params = SparseRuntimeParams(sparse_kv_indices=torch.empty(1)) if attention_input_type == AttentionInputType.context_only: @@ -239,8 +254,11 @@ def _support_result( dtype=torch.int32, ), ) + kv_cache_manager.mla_kv_cache_residual_dim = mla_kv_cache_residual_dim metadata = SimpleNamespace( - helix_position_offsets=None, + helix_position_offsets=_TensorSpec((num_tokens,), torch.int32) if helix else None, + helix_kv_bounds=_TensorSpec((num_tokens,), torch.int32) if helix else None, + _helix_spec_tokens_valid=helix_bounds_valid, num_sparse_topk=0, use_spec_decoding=use_spec_decoding, is_spec_dec_tree=is_spec_dec_tree, @@ -261,9 +279,20 @@ def _support_result( num_contexts=num_contexts, num_generations=num_generations, num_ctx_tokens=num_ctx_tokens, - kv_lens_runtime=torch.tensor(kv_lens, dtype=torch.int32), + kv_lens_runtime=torch.tensor( + kv_lens_override if kv_lens_override is not None else kv_lens, dtype=torch.int32 + ), + seq_lens=torch.tensor( + query_lens + if query_lens is not None + else [query_width] * (num_contexts + num_generations), + dtype=torch.int32, + ), max_seq_len=max_seq_len, ) + if metadata_overrides is not None: + for name, value in metadata_overrides.items(): + setattr(metadata, name, value) fmha = PrimsTSFmha(attn) k = _TensorSpec((4, attn.num_kv_heads * head_dim), dtype) if has_separate_kv else None v = ( @@ -303,6 +332,95 @@ def test_supports_mla_context(num_heads: int, fp8: bool) -> None: assert supported, reason +@pytest.mark.parametrize("query_width", [1, 6, 8]) +@pytest.mark.parametrize("helix", [False, True]) +@pytest.mark.parametrize("page_size", [32, 128]) +def test_supports_uniform_linear_fp8_mla(query_width: int, helix: bool, page_size: int) -> None: + supported, reason = _support_result( + attention_input_type=AttentionInputType.generation_only, + head_dim=576, + is_mla=True, + quant_mode=QuantMode.FP8_KV_CACHE, + kv_dtype=DataType.FP8, + use_spec_decoding=True, + query_width=query_width, + tokens_per_block=page_size, + helix=helix, + helix_bounds_valid=query_width > 1, + kv_lens_override=[0, 0, 32, 64] if helix else None, + ) + assert supported, reason + + +@pytest.mark.parametrize( + "overrides,expected_reason", + [ + ({"query_lens": [1, 11, 6, 6]}, "uniform per-request"), + ({"is_spec_dec_tree": True}, "tree speculative"), + ({"metadata_overrides": {"is_spec_dec_dynamic_tree": True}}, "tree speculative"), + ({"forward_overrides": {"attention_mask_data": torch.empty(1)}}, "custom attention"), + ({"helix_bounds_valid": False}, "valid per-token KV bounds"), + ({"forward_overrides": {"softmax_stats_tensor": None}}, "requires softmax statistics"), + ( + { + "forward_overrides": { + "softmax_stats_tensor": _TensorSpec((24, 8, 2), torch.bfloat16) + } + }, + "contiguous FP32", + ), + ( + {"metadata_overrides": {"helix_kv_bounds": _TensorSpec((24,), torch.int64)}}, + "int32 per-token", + ), + ( + {"metadata_overrides": {"helix_kv_bounds": _TensorSpec((23,), torch.int32)}}, + "int32 per-token", + ), + ( + { + "metadata_overrides": { + "helix_kv_bounds": _TensorSpec((24,), torch.int32, device="cpu") + } + }, + "int32 per-token", + ), + ( + { + "metadata_overrides": { + "helix_kv_bounds": _TensorSpec((24,), torch.int32, contiguous=False) + } + }, + "int32 per-token", + ), + ( + {"forward_overrides": {"attention_mask": PredefinedAttentionMask.FULL}}, + "requires a causal mask", + ), + ({"helix": False, "kv_lens_override": [0, 32, 64, 96]}, "at least one KV"), + ({"mla_kv_cache_residual_dim": 64}, "full FP8 KV cache"), + ], +) +def test_rejects_unsupported_mla_verification_layouts( + overrides: dict, expected_reason: str +) -> None: + args = dict( + attention_input_type=AttentionInputType.generation_only, + head_dim=576, + is_mla=True, + quant_mode=QuantMode.FP8_KV_CACHE, + kv_dtype=DataType.FP8, + use_spec_decoding=True, + query_width=6, + helix=True, + helix_bounds_valid=True, + ) + args.update(overrides) + supported, reason = _support_result(**args) + assert not supported + assert expected_reason in reason + + @pytest.mark.parametrize( "case", [ @@ -388,7 +506,7 @@ def test_fp8_mla_supported( ({"is_fused_qkv": True}, "MLA decode requires query-only input"), ({"attention_input_type": AttentionInputType.context_only}, "MLA context requires QK=192"), ({"attention_input_type": AttentionInputType.mixed}, "dispatched separately"), - ({"use_spec_decoding": True}, "speculative decoding"), + ({"is_spec_dec_tree": True}, "tree speculative decoding"), ({"has_sparse_attention": True}, "sparse attention"), ], ) @@ -1567,7 +1685,7 @@ def test_workspace_allocation_change_invalidates_only_workspace_bound_wrappers() fmha._update_workspace_allocation(first_workspace) fmha._context_wrappers[1] = Mock() fmha._decode_wrappers[1] = Mock() - fmha._mla_decode_wrappers[1] = Mock() + fmha._mla_decode_wrappers[(1, 1)] = Mock() second_workspace = torch.empty(64, dtype=torch.uint8) fmha._update_workspace_allocation(second_workspace) @@ -1584,6 +1702,9 @@ def _get_test_mla_wrapper( workspace_buffer: torch.Tensor, *, mask_type: str = "causal", + query_width: int = 1, + store_softmax_stats: bool = False, + use_kv_bounds: bool = False, ) -> object: return fmha._get_or_plan_mla_decode_wrapper( workspace_buffer, @@ -1592,26 +1713,44 @@ def _get_test_mla_wrapper( kv_lora_rank=512, qk_rope_head_dim=64, page_size=32, - max_seq_len_q=1, - max_kv_len=96, + max_seq_len_q=query_width, + max_kv_len=block_tables.shape[-1] * 32, q_dtype=torch.bfloat16, kv_dtype=torch.bfloat16, output_dtype=torch.bfloat16, mask_type=mask_type, + store_softmax_stats=store_softmax_stats, + use_kv_bounds=use_kv_bounds, ) @pytest.mark.parametrize("fp8", [False, True], ids=["bf16", "fp8"]) +@pytest.mark.parametrize("page_size", [32, 128]) +@pytest.mark.parametrize( + "query_width,helix_mode", + [ + (1, "none"), + (6, "none"), + (8, "none"), + (1, "single"), + (1, "token"), + (6, "token"), + (8, "token"), + ], +) def test_mla_eager_wrapper_plans_once_and_reads_live_fixed_metadata( monkeypatch: pytest.MonkeyPatch, fp8: bool, + page_size: int, + query_width: int, + helix_mode: str, ) -> None: monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) attn = _Attention(head_dim=576, is_mla=True, num_heads=4) attn.quant_mode = QuantMode.FP8_KV_CACHE if fp8 else QuantMode(0) fmha = PrimsTSFmha(attn) kernel_dtype = torch.float8_e4m3fn if fp8 else torch.bfloat16 - kv_cache = torch.empty((20, 1, 32, 576), dtype=kernel_dtype) + kv_cache = torch.empty((20, 1, page_size, 576), dtype=kernel_dtype) block_tables = torch.tensor( [ [[0, 1, 2], [10, 11, 12]], @@ -1638,41 +1777,56 @@ def test_mla_eager_wrapper_plans_once_and_reads_live_fixed_metadata( wrapper_factory, ) + num_tokens = 2 * query_width + helix = helix_mode != "none" + # Include unused capacity to catch accidental context/padded-row offsets. + kv_bounds = torch.tensor( + [0] * query_width + [page_size - 1] + [page_size] * (query_width - 1) + [999] * 4, + dtype=torch.int32, + ) + stats = torch.empty((num_tokens, attn.num_heads, 2), dtype=torch.float32) if helix else None metadata = SimpleNamespace( is_cuda_graph=False, + helix_position_offsets=torch.empty(num_tokens, dtype=torch.int32) if helix else None, + _helix_spec_tokens_valid=helix_mode == "token", + helix_kv_bounds=kv_bounds, beam_width=1, kv_cache_block_offsets=torch.empty((2, 2, 3), dtype=torch.int32), host_kv_cache_pool_pointers=torch.tensor([1234], dtype=torch.int64), host_kv_cache_pool_mapping=torch.tensor([[0, 0]], dtype=torch.int32), ) - output = torch.empty((2, attn.num_heads, 512), dtype=torch.bfloat16) + output = torch.empty((num_tokens, attn.num_heads, 512), dtype=torch.bfloat16) forward_args = AttentionForwardArgs( output=output, attention_input_type=AttentionInputType.generation_only, attention_window_size=64, - quant_q_buffer=torch.zeros(2 * attn.num_heads * 576 + 16, dtype=torch.uint8) + quant_q_buffer=torch.zeros(num_tokens * attn.num_heads * 576 + 16, dtype=torch.uint8) if fp8 else None, mla_bmm1_scale=torch.tensor([0.125, 0.125 * math.log2(math.e)]) if fp8 else None, mla_bmm2_scale=torch.tensor([2.0]) if fp8 else None, + softmax_stats_tensor=stats, + ) + all_sequence_lengths = torch.tensor( + [999, 999, 0, page_size] if helix else [999, 999, page_size + 1, page_size * 2], + dtype=torch.int32, ) - sequence_lengths = torch.tensor([33, 64], dtype=torch.int32) - assert sequence_lengths.data_ptr() % 16 == 0 + sequence_lengths = all_sequence_lengths[2:] params = FmhaParams( attn=attn, meta=metadata, fwd=forward_args, workspace=torch.empty(64, dtype=torch.uint8), - query_input=torch.empty((2, attn.num_heads * 576), dtype=torch.bfloat16), + query_input=torch.empty((num_tokens, attn.num_heads * 576), dtype=torch.bfloat16), output=output, sequence_lengths=sequence_lengths, - input_seq_length=1, + input_seq_length=query_width, max_past_kv_length=64, max_attention_window_size=64, cyclic_attention_window_size=64, - num_tokens=2, + num_tokens=num_tokens, seq_offset=2, - tokens_per_block=32, + tokens_per_block=page_size, kv_factor=1, total_num_blocks=20, batch_size=2, @@ -1691,24 +1845,26 @@ def test_mla_eager_wrapper_plans_once_and_reads_live_fixed_metadata( attn.num_heads, 512, 64, - 32, - 96, + page_size, + 3 * page_size, ) workspace_buffer = plan_kwargs["workspace_buffer"] assert workspace_buffer.data_ptr() == params.workspace.data_ptr() assert workspace_buffer.numel() == 64 assert {key: value for key, value in plan_kwargs.items() if key != "workspace_buffer"} == { - "max_seq_len_q": 1, + "max_seq_len_q": query_width, "packed_query": False, "q_data_type": kernel_dtype, "kv_data_type": kernel_dtype, "o_data_type": torch.bfloat16, "mask_type": "causal", + "store_softmax_stats": helix, + "use_kv_bounds": helix, } wrapper.run.assert_called_once() run_args = wrapper.run.call_args.args run_kwargs = wrapper.run.call_args.kwargs - assert run_args[0].shape == (2, 1, attn.num_heads, 576) + assert run_args[0].shape == (2, query_width, attn.num_heads, 576) query_buffer = forward_args.quant_q_buffer if fp8 else params.query_input assert run_args[0].data_ptr() == query_buffer.data_ptr() assert run_args[0].dtype == kernel_dtype @@ -1722,16 +1878,31 @@ def test_mla_eager_wrapper_plans_once_and_reads_live_fixed_metadata( ) assert run_kwargs["seq_lens"].data_ptr() == sequence_lengths.data_ptr() torch.testing.assert_close(run_kwargs["seq_lens"], sequence_lengths) - assert run_kwargs["out"].shape == (2, 1, attn.num_heads, 512) + assert run_kwargs["out"].shape == (2, query_width, attn.num_heads, 512) assert run_kwargs["out"].data_ptr() == output.data_ptr() assert run_kwargs["bmm1_scale"] == pytest.approx(0.125 if fp8 else 1.0 / math.sqrt(128 + 64)) assert run_kwargs["bmm2_scale"] == (2.0 if fp8 else 1.0) assert run_kwargs["validate"] is False + if helix: + expected_bounds = kv_bounds[:num_tokens] if helix_mode == "token" else sequence_lengths + assert run_kwargs["kv_bounds"].data_ptr() == expected_bounds.data_ptr() + assert run_kwargs["kv_bounds"].shape == (num_tokens,) + torch.testing.assert_close(run_kwargs["kv_bounds"], expected_bounds) + assert run_kwargs["softmax_stats"].shape == (2, query_width, attn.num_heads, 2) + assert run_kwargs["softmax_stats"].data_ptr() == stats.data_ptr() + else: + assert run_kwargs["kv_bounds"] is None + assert run_kwargs["softmax_stats"] is None block_tables[:, 0].add_(20) block_tables[:, 1].add_(200) sequence_lengths.add_(1) - fmha.run_mla_generation(params) + kv_bounds[:num_tokens].add_(1) + with monkeypatch.context() as capture_patch: + capture_patch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + capture_patch.setattr(torch.Tensor, "item", Mock(side_effect=AssertionError("device read"))) + capture_patch.setattr(torch.Tensor, "cpu", Mock(side_effect=AssertionError("device read"))) + fmha.run_mla_generation(params) wrapper_factory.assert_called_once_with() wrapper.plan.assert_called_once() @@ -1742,8 +1913,10 @@ def test_mla_eager_wrapper_plans_once_and_reads_live_fixed_metadata( run_kwargs["block_tables"], torch.tensor([[20, 21, 22], [23, 24, 25]], dtype=torch.int32), ) - torch.testing.assert_close(run_kwargs["seq_lens"], torch.tensor([34, 65], dtype=torch.int32)) - assert fmha._mla_decode_wrappers[2] is wrapper + torch.testing.assert_close(run_kwargs["seq_lens"], sequence_lengths) + if helix: + torch.testing.assert_close(wrapper.run.call_args.kwargs["kv_bounds"], expected_bounds) + assert fmha._mla_decode_wrappers[(2, query_width)] is wrapper @pytest.mark.parametrize("buffer_dtype", [torch.uint8, torch.float8_e4m3fn]) @@ -1867,8 +2040,8 @@ def test_mla_wrapper_cache_plans_each_batch_once_and_reuses_a_b_a( for wrapper in wrappers: wrapper.plan.assert_called_once() assert wrapper.plan.call_args.kwargs["workspace_buffer"] is workspace - assert fmha._mla_decode_wrappers[2] is wrappers[0] - assert fmha._mla_decode_wrappers[1] is wrappers[1] + assert fmha._mla_decode_wrappers[(2, 1)] is wrappers[0] + assert fmha._mla_decode_wrappers[(1, 1)] is wrappers[1] def test_mla_wrapper_capture_uses_cached_plan_and_rejects_plan_miss( @@ -1909,6 +2082,51 @@ def test_mla_wrapper_capture_uses_cached_plan_and_rejects_plan_miss( wrapper.plan.assert_called_once() +@pytest.mark.parametrize("helix", [False, True]) +def test_mla_wrapper_cache_separates_query_widths( + monkeypatch: pytest.MonkeyPatch, + helix: bool, +) -> None: + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + wrapper_factory = Mock(side_effect=lambda: Mock()) + monkeypatch.setattr(prims_mla_module, "BatchMLADecodePagedTSWrapper", wrapper_factory) + fmha = PrimsTSFmha(_Attention(head_dim=576, is_mla=True, num_heads=4)) + block_tables = torch.zeros((2, 3), dtype=torch.int32) + seq_lens = torch.tensor([33, 64], dtype=torch.int32) + workspace = torch.empty(64, dtype=torch.uint8) + widths = (1, 6, 8) + planned = [] + for query_width in widths: + planned.append( + _get_test_mla_wrapper( + fmha, + block_tables, + seq_lens, + workspace, + query_width=query_width, + store_softmax_stats=helix, + use_kv_bounds=helix, + ) + ) + assert wrapper_factory.call_count == len(widths) + assert len({id(wrapper) for wrapper in planned}) == len(widths) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + for query_width, wrapper in zip(reversed(widths), reversed(planned)): + assert ( + _get_test_mla_wrapper( + fmha, + block_tables, + seq_lens, + workspace, + query_width=query_width, + store_softmax_stats=helix, + use_kv_bounds=helix, + ) + is wrapper + ) + wrapper.plan.assert_called_once() + + @pytest.mark.parametrize("is_cuda_graph", [False, True], ids=["eager", "cuda-graph"]) def test_mla_wrapper_receives_v2_bound_and_shared_workspace( monkeypatch: pytest.MonkeyPatch, @@ -1947,6 +2165,7 @@ def test_mla_wrapper_receives_v2_bound_and_shared_workspace( get_page_index_upper_bound = Mock(return_value=20) metadata = SimpleNamespace( is_cuda_graph=is_cuda_graph, + helix_position_offsets=None, beam_width=1, kv_cache_block_offsets=torch.empty((2, 2, 3), dtype=torch.int32), host_kv_cache_pool_pointers=torch.tensor([1234], dtype=torch.int64), @@ -2033,10 +2252,13 @@ def test_mla_wrapper_receives_v2_bound_and_shared_workspace( @pytest.mark.parametrize("is_cuda_graph", [False, True], ids=["eager", "cuda-graph"]) @pytest.mark.parametrize("fp8", [False, True], ids=["bf16", "fp8"]) +@pytest.mark.parametrize("query_width,with_stats", [(1, False), (6, True), (8, True)]) def test_mla_prepare_workspace_sizes_caller_owned_workspace( monkeypatch: pytest.MonkeyPatch, is_cuda_graph: bool, fp8: bool, + query_width: int, + with_stats: bool, ) -> None: monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) workspace_size = Mock(return_value=48) @@ -2049,7 +2271,7 @@ def test_mla_prepare_workspace_sizes_caller_owned_workspace( attn.quant_mode = QuantMode.FP8_KV_CACHE if fp8 else QuantMode(0) fmha = PrimsTSFmha(attn) fmha._multi_processor_count = 1 - q = torch.empty((2, attn.num_heads * 576), dtype=torch.bfloat16) + q = torch.empty((2 * query_width, attn.num_heads * 576), dtype=torch.bfloat16) metadata = SimpleNamespace( is_cuda_graph=is_cuda_graph, kv_cache_block_offsets=torch.empty((2, 2, 3), dtype=torch.int32), @@ -2060,11 +2282,14 @@ def test_mla_prepare_workspace_sizes_caller_owned_workspace( tokens_per_block=32, kv_lens_runtime=torch.tensor([33, 64], dtype=torch.int32), ) - output = torch.empty((2, attn.num_heads * 512), dtype=torch.bfloat16) + output = torch.empty((2 * query_width, attn.num_heads * 512), dtype=torch.bfloat16) forward_args = AttentionForwardArgs( output=output, attention_input_type=AttentionInputType.generation_only, attention_window_size=96, + softmax_stats_tensor=torch.empty((2 * query_width, attn.num_heads, 2), dtype=torch.float32) + if with_stats + else None, ) workspace = torch.empty(0, dtype=torch.uint8) @@ -2082,6 +2307,8 @@ def test_mla_prepare_workspace_sizes_caller_owned_workspace( assert workspace_size.call_args.kwargs["q_dtype"] == kernel_dtype assert workspace_size.call_args.kwargs["kv_dtype"] == kernel_dtype assert workspace_size.call_args.kwargs["out_dtype"] == torch.bfloat16 + assert workspace_size.call_args.kwargs["max_seq_len_q"] == query_width + assert workspace_size.call_args.kwargs["store_softmax_stats"] == with_stats assert workspace.numel() == 48 @@ -2103,7 +2330,7 @@ def test_mla_prepare_workspace_preserves_cached_wrappers_with_stable_allocation( workspace = torch.empty(48, dtype=torch.uint8) fmha._update_workspace_allocation(workspace) wrapper = Mock() - fmha._mla_decode_wrappers[2] = wrapper + fmha._mla_decode_wrappers[(2, 1)] = wrapper q = torch.empty((2, attn.num_heads * 576), dtype=torch.bfloat16) metadata = SimpleNamespace( is_cuda_graph=True, @@ -2124,7 +2351,7 @@ def test_mla_prepare_workspace_preserves_cached_wrappers_with_stable_allocation( fmha.prepare_workspace(q, None, None, metadata, forward_args, workspace) stream.synchronize.assert_not_called() - assert fmha._mla_decode_wrappers[2] is wrapper + assert fmha._mla_decode_wrappers[(2, 1)] is wrapper assert workspace.numel() == 48 workspace_size.assert_called_once() From dcc5edc79218d196e0f8979a4e32d4c413e3785d Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Thu, 29 Oct 2026 04:07:52 +0000 Subject: [PATCH 08/19] [None][feat] Allow linear speculative MLA context in PrimTS Permit explicit MLA context calls under linear speculative metadata, retaining tree, custom-mask, mixed-phase and Helix-context exclusions. Cover causal fresh/final chunks, FULL cached chunks, FP8 statistics and live metadata reuse. Validation: 363 adapter tests and 8 Kimi K3 module tests passed with no skips in the matching Rubin runtime; two independent review passes and pre-commit passed. Full-model MTP qualification remains separate. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../attention/backends/fmha/prims_ts.py | 10 +- .../_torch/attention/test_prims_ts_fmha.py | 137 ++++++++++++++++-- 2 files changed, 129 insertions(+), 18 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py index e91d6e9a5adc..6426ec084a57 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py @@ -325,8 +325,9 @@ def _is_supported_with_reason( return False, "sparse attention metadata is not supported." if meta.num_sparse_topk > 0: return False, "sparse attention metadata is not supported." + is_mla = attn.is_mla_enable is_mla_generation = ( - attn.is_mla_enable and fwd.attention_input_type == AttentionInputType.generation_only + is_mla and fwd.attention_input_type == AttentionInputType.generation_only ) feature_reason = get_attention_feature_unsupported_reason( meta, fwd, allow_helix=is_mla_generation @@ -348,11 +349,14 @@ def _is_supported_with_reason( 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, allow_linear_mla=is_mla_generation + attn, + meta, + allow_linear_mla=is_mla + and fwd.attention_input_type + in (AttentionInputType.context_only, AttentionInputType.generation_only), ) 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() input_type = fwd.attention_input_type diff --git a/tests/unittest/_torch/attention/test_prims_ts_fmha.py b/tests/unittest/_torch/attention/test_prims_ts_fmha.py index 78dac7347741..1a42c15c48c9 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_fmha.py +++ b/tests/unittest/_torch/attention/test_prims_ts_fmha.py @@ -332,6 +332,95 @@ def test_supports_mla_context(num_heads: int, fp8: bool) -> None: assert supported, reason +@pytest.mark.parametrize("fp8", [False, True]) +@pytest.mark.parametrize("spec_flags", [(True, False), (False, True), (True, True)]) +@pytest.mark.parametrize( + "mask,store_stats,kv_len", + [ + (PredefinedAttentionMask.CAUSAL, False, 4), + (PredefinedAttentionMask.CAUSAL, True, 4), + (PredefinedAttentionMask.FULL, True, 4), + (PredefinedAttentionMask.FULL, True, 0), + ], + ids=["fresh", "final-chunk", "cached-chunk", "empty-cached-chunk"], +) +def test_supports_linear_mla_context( + fp8: bool, + spec_flags: tuple[bool, bool], + mask: PredefinedAttentionMask, + store_stats: bool, + kv_len: int, +) -> None: + supported, reason = _support_result( + attention_input_type=AttentionInputType.context_only, + head_dim=192, + is_mla=True, + num_heads=96, + num_kv_heads=96, + has_separate_kv=True, + tokens_per_block=128, + kv_dtype=DataType.FP8 if fp8 else DataType.BF16, + quant_mode=QuantMode.from_description(use_fp8_kv_cache=fp8), + metadata_overrides={ + "is_spec_decoding_enabled": spec_flags[0], + "use_spec_decoding": spec_flags[1], + }, + forward_overrides={ + "attention_mask": mask, + "softmax_stats_tensor": _TensorSpec((4, 96, 2), torch.float32) if store_stats else None, + }, + kv_lens_override=[kv_len], + ) + assert supported, reason + + +@pytest.mark.parametrize( + "overrides,expected_reason", + [ + ({"is_spec_dec_tree": True}, "tree speculative"), + ({"metadata_overrides": {"is_spec_dec_dynamic_tree": True}}, "tree speculative"), + ({"forward_overrides": {"attention_mask_data": torch.empty(1)}}, "custom attention"), + ({"helix": True}, "Helix parallelism"), + ({"attention_input_type": AttentionInputType.mixed}, "speculative decoding"), + ( + { + "is_mla": False, + "head_dim": 128, + "num_kv_heads": 96, + "is_fused_qkv": True, + "has_separate_kv": False, + }, + "speculative decoding", + ), + ( + { + "kv_lens_override": [0], + "forward_overrides": { + "attention_mask": PredefinedAttentionMask.FULL, + "softmax_stats_tensor": None, + }, + }, + "at least one KV", + ), + ], +) +def test_rejects_unsupported_linear_mla_context(overrides: dict, expected_reason: str) -> None: + args = dict( + attention_input_type=AttentionInputType.context_only, + head_dim=192, + is_mla=True, + num_heads=96, + num_kv_heads=96, + has_separate_kv=True, + tokens_per_block=128, + use_spec_decoding=True, + ) + args.update(overrides) + supported, reason = _support_result(**args) + assert not supported + assert expected_reason in reason + + @pytest.mark.parametrize("query_width", [1, 6, 8]) @pytest.mark.parametrize("helix", [False, True]) @pytest.mark.parametrize("page_size", [32, 128]) @@ -1070,34 +1159,44 @@ def test_sequence_lengths_are_live_source_view() -> None: torch.testing.assert_close(actual, torch.tensor([33, 65], dtype=torch.int32)) +@pytest.mark.parametrize("num_heads", [6, 96]) +@pytest.mark.parametrize("fp8", [False, True]) +@pytest.mark.parametrize("mask", [PredefinedAttentionMask.CAUSAL, PredefinedAttentionMask.FULL]) def test_mla_context_wrapper_reuses_plan_and_updates_chunk_offsets( monkeypatch: pytest.MonkeyPatch, + num_heads: int, + fp8: bool, + mask: PredefinedAttentionMask, ) -> None: monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) - attn = _Attention(head_dim=192, is_mla=True, num_heads=6, num_kv_heads=6) + attn = _Attention(head_dim=192, is_mla=True, num_heads=num_heads, num_kv_heads=num_heads) + attn.quant_mode = QuantMode.from_description(use_fp8_kv_cache=fp8) attn._ensure_rope_table_size = Mock() attn.get_local_layer_idx = lambda _meta: 0 attn.rotary_cos_sin = torch.empty(1) fmha = PrimsTSFmha(attn) - q = torch.empty(5, 6 * 192, dtype=torch.bfloat16) - k = torch.empty(9, 6 * 192, dtype=q.dtype) - v = torch.empty(9, 6 * 128, dtype=q.dtype) + q = torch.empty(5, num_heads * 192, dtype=torch.bfloat16) + k = torch.empty(9, num_heads * 192, dtype=q.dtype) + v = torch.empty(9, num_heads * 256, dtype=q.dtype)[:, num_heads * 128 :] qo_indptr = torch.tensor([0, 2, 5], dtype=torch.int32) - output = torch.empty(5, 6, 128, dtype=q.dtype) - stats = torch.empty(5, 6, 2) + output = torch.empty(5, num_heads, 128, dtype=q.dtype) + stats = torch.empty(5, num_heads, 2) meta = SimpleNamespace( mla_prepare_ctx_cu_seqlens=lambda: qo_indptr, max_seq_len=4096, kv_cache_block_offsets=torch.empty(2, 2, 64, dtype=torch.int32), host_kv_cache_pool_pointers=torch.empty(1), host_kv_cache_pool_mapping=torch.zeros(1, 2, dtype=torch.int32), + is_spec_decoding_enabled=True, + use_spec_decoding=True, ) workspace = torch.empty(fmha._mla_context_workspace_layout(q, k, 2)[1], dtype=torch.uint8) params = FmhaParams( attn=attn, meta=meta, fwd=AttentionForwardArgs( - attention_mask=PredefinedAttentionMask.FULL, + attention_input_type=AttentionInputType.context_only, + attention_mask=mask, softmax_stats_tensor=stats, is_fused_qkv=False, ), @@ -1106,11 +1205,11 @@ def test_mla_context_wrapper_reuses_plan_and_updates_chunk_offsets( key_input=k, value_input=v, output=output, - sequence_lengths=torch.tensor([7, 2], dtype=torch.int32), + sequence_lengths=torch.tensor([4, 5], dtype=torch.int32), batch_size=2, input_seq_length=3, max_past_kv_length=7, - tokens_per_block=64, + tokens_per_block=128, cyclic_attention_window_size=4096, ) native_preprocess = Mock() @@ -1120,22 +1219,30 @@ def test_mla_context_wrapper_reuses_plan_and_updates_chunk_offsets( monkeypatch.setattr(prims_context_module, "BatchPrefillTSWrapper", factory) fmha.run_mla_context(params) first_kv_offsets = wrapper.run.call_args.args[4].clone() - params.sequence_lengths.copy_(torch.tensor([2, 7], dtype=torch.int32)) + params.sequence_lengths.copy_(torch.tensor([3, 6], dtype=torch.int32)) qo_indptr.copy_(torch.tensor([0, 3, 5], dtype=torch.int32)) fmha.run_mla_context(params) assert factory.call_count == wrapper.plan.call_count == 1 assert native_preprocess.call_count == wrapper.run.call_count == 2 - torch.testing.assert_close(first_kv_offsets, torch.tensor([0, 7, 9], dtype=torch.int32)) + torch.testing.assert_close(first_kv_offsets, torch.tensor([0, 4, 9], dtype=torch.int32)) torch.testing.assert_close( - wrapper.run.call_args.args[4], torch.tensor([0, 2, 9], dtype=torch.int32) + wrapper.run.call_args.args[4], torch.tensor([0, 3, 9], dtype=torch.int32) ) assert wrapper.run.call_args.args[3] is qo_indptr assert wrapper.run.call_args.kwargs["softmax_stats"] is stats assert wrapper.run.call_args.kwargs["out"] is output - assert wrapper.run.call_args.args[1].shape == (9, 6, 192) + assert wrapper.run.call_args.args[1].shape == (9, num_heads, 192) assert wrapper.plan.call_args.kwargs["store_softmax_stats"] - assert wrapper.plan.call_args.kwargs["k_dtype"] == k.dtype - assert wrapper.plan.call_args.kwargs["v_dtype"] == v.dtype + processed_dtype = torch.float8_e4m3fn if fp8 else q.dtype + for name in ("q_dtype", "k_dtype", "v_dtype"): + assert wrapper.plan.call_args.kwargs[name] == processed_dtype + for tensor in wrapper.run.call_args.args[:3]: + assert tensor.dtype == processed_dtype + assert wrapper.plan.call_args.kwargs["mask_type"] == ( + "causal" if mask == PredefinedAttentionMask.CAUSAL else "dense" + ) + assert (wrapper.run.call_args.kwargs["scale_softmax_log2"] is not None) == fp8 + assert (wrapper.run.call_args.kwargs["output_scale"] is not None) == fp8 assert wrapper.run.call_args.kwargs["validate"] is False From 82db3c19c1a5dc809f635836892ae3b2891a9747 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Thu, 29 Oct 2026 07:39:25 +0000 Subject: [PATCH 09/19] [None][test] Cover FP8 MLA linear MTP with real K3 modules Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../test_prims_ts_attention_backend.py | 90 +++++++++++++++++-- 1 file changed, 82 insertions(+), 8 deletions(-) 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 5ba32f99335b..453684312868 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py +++ b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py @@ -469,6 +469,27 @@ def test_prims_ts_kimi_k3_full_attention( Head counts cover the TP16-local and attention-DP geometries, not distributed collectives. The full-model GB300 test covers those separately. """ + _run_prims_ts_kimi_k3_full_attention(monkeypatch, num_heads, fp8_kv, chunked) + + +@pytest.mark.parametrize("query_width", [6, 8], ids=["mtp5", "mtp7"]) +def test_prims_ts_kimi_k3_fp8_mla_linear_mtp( + monkeypatch: pytest.MonkeyPatch, query_width: int +) -> None: + """Verify uniform MTP with the real K3 module, V2 FP8/page128 cache and graph.""" + _run_prims_ts_kimi_k3_full_attention( + monkeypatch, 96, True, False, generation_query_width=query_width + ) + + +def _run_prims_ts_kimi_k3_full_attention( + monkeypatch: pytest.MonkeyPatch, + num_heads: int, + fp8_kv: bool, + chunked: bool, + *, + generation_query_width: int = 1, +) -> None: import tensorrt_llm from tensorrt_llm._torch.attention.backends import TrtllmAttention from tensorrt_llm._torch.attention.backends.fmha.prims_ts import PrimsTSFmha @@ -476,10 +497,12 @@ def test_prims_ts_kimi_k3_full_attention( from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.modules.kimi_k3_mla.kimi_k3_mla_attention import KimiK3MLAAttention + from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, SamplingConfig from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._utils import torch_dtype_to_binding from tensorrt_llm.llmapi import KvCacheConfig, SamplingParams + from tensorrt_llm.llmapi.llm_args import DSparkDecodingConfig from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization import QuantAlgo @@ -531,13 +554,32 @@ def record_context(fmha, params): prefill_reference = torch.cat(references).flatten(1).to(params.output.dtype) def record_generation(fmha, params): + assert params.input_seq_length == generation_query_width + assert params.attn.predicted_tokens_per_seq == 1 + if generation_query_width > 1: + assert params.tokens_per_block == 128 + assert params.fwd.quant_q_buffer.dtype == torch.uint8 + assert isinstance(params.meta.kv_cache_manager, KVCacheManagerV2) generation_calls.append(params.batch_size) run_generation(fmha, params) monkeypatch.setattr(PrimsTSFmha, "run_mla_context", record_context) monkeypatch.setattr(PrimsTSFmha, "run_mla_generation", record_generation) kv_dtype = torch.float8_e4m3fn if fp8_kv else torch.bfloat16 - lengths = [[129, 65], [17, 193], [1, 1]] + use_mtp = generation_query_width > 1 + # The MTP verification group crosses a page boundary in both requests. + lengths = ( + [[65, 129], [62, 126], [generation_query_width] * 2] + if use_mtp + else [[129, 65], [17, 193], [1, 1]] + ) + spec_config = ( + DSparkDecodingConfig( + max_draft_len=generation_query_width - 1, block_size=generation_query_width - 1 + ) + if use_mtp + else None + ) torch.manual_seed(2718) hidden_steps = [ torch.randn(sum(q_lengths), 256, dtype=torch.bfloat16, device="cuda") @@ -561,6 +603,7 @@ def record_generation(fmha, params): model_config=ModelConfig( mapping=mapping, quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8 if fp8_kv else None), + spec_config=spec_config, ), aux_stream_dict={}, ) @@ -579,17 +622,23 @@ def record_generation(fmha, params): k_weight, v_weight = module.kv_b_proj.weight.split(num_heads * 128) module.k_b_proj_trans.copy_(k_weight.view(num_heads, 128, 512).transpose(1, 2)) module.v_b_proj.copy_(v_weight.view(num_heads, 128, 512)) - manager = KVCacheManager( + manager_cls = KVCacheManagerV2 if use_mtp else KVCacheManager + manager = manager_cls( KvCacheConfig(max_tokens=1024, enable_block_reuse=False), tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY, num_layers=1, num_kv_heads=1, head_dim=576, - tokens_per_block=64, + tokens_per_block=128 if use_mtp else 64, max_seq_len=512, max_batch_size=2, mapping=mapping, dtype=torch_dtype_to_binding(kv_dtype), + **( + {"spec_config": spec_config, "layer_mask": [True], "vocab_size": 128} + if use_mtp + else {} + ), ) # Both decode implementations require finite values in unused cache # slots. Reused BF16 allocations can contain FP8 NaN byte patterns; @@ -605,7 +654,13 @@ def record_generation(fmha, params): is_streaming=False, ) request.paged_kv_block_ids = [] - manager.impl.add_sequence_batch([(request_id, length, 1)], [request]) + if use_mtp: + cache = manager._create_kv_cache(request_id, None, [1] * length) + assert cache is not None + assert cache.resume(torch.cuda.current_stream().cuda_stream) + assert cache.resize(length + manager.num_extra_kv_tokens) + else: + manager.impl.add_sequence_batch([(request_id, length, 1)], [request]) requests.append(request) cached = [0, 0] results[library] = [] @@ -614,8 +669,14 @@ def record_generation(fmha, params): for step, q_lengths in enumerate(lengths): if step: for request_id, length in enumerate(q_lengths): - for _ in range(length): - manager.impl.add_token(request_id) + if use_mtp: + cache = manager.kv_cache_map[request_id] + assert cache.resize( + cached[request_id] + length + manager.num_extra_kv_tokens + ) + else: + for _ in range(length): + manager.impl.add_token(request_id) metadata = TrtllmAttention.Metadata( seq_lens=torch.tensor(q_lengths, dtype=torch.int32), request_ids=[0, 1], @@ -625,7 +686,9 @@ def record_generation(fmha, params): max_num_tokens=512, kv_cache_manager=manager, kv_cache_params=KVCacheParams( - use_cache=True, num_cached_tokens_per_seq=cached + use_cache=True, + num_cached_tokens_per_seq=cached, + num_extra_kv_tokens=manager.num_extra_kv_tokens, ), mapping=mapping, enable_context_mla_with_cached_kv=True, @@ -634,9 +697,19 @@ def record_generation(fmha, params): chunked_prefill=chunked, chunk_size=128, chunked_prefill_buffer_batch_size=1, + has_speculative_draft_tokens=use_mtp, ), ) metadata.prepare() + if use_mtp: + metadata.update_spec_dec_param( + batch_size=2, + is_spec_decoding_enabled=True, + is_spec_dec_tree=False, + is_spec_dec_dynamic_tree=False, + max_draft_len=generation_query_width - 1, + max_total_draft_tokens=generation_query_width - 1, + ) hidden = hidden_steps[step] positions = torch.cat( [ @@ -731,6 +804,7 @@ def test_prims_ts_fp8_mla_scales_and_graph_replay( num_ctx_tokens=0, num_generations=batch_size, tokens_per_block=32, + helix_position_offsets=None, ) # Pool ownership/indexing is covered by the real v1/v2 preprocessing test. monkeypatch.setattr( @@ -786,7 +860,7 @@ def check_output() -> None: fmha.run_mla_generation(params) if max_seq_len < 128: # Auto policy must fall back when the 1CTA profile cannot cover this capacity. - wrapper = fmha._mla_decode_wrappers[batch_size] + wrapper = fmha._mla_decode_wrappers[(batch_size, 1)] assert dict(wrapper._plan_state.policy)["kernel"] == "throughput_2cta" check_output() graph = torch.cuda.CUDAGraph() From b7766e64e7615ef46e888ea9298662feeeec73a3 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Sat, 31 Oct 2026 03:38:02 +0000 Subject: [PATCH 10/19] [None][test] Cover Kimi K3 FP8 MLA MTP with PrimTS Add the FP8 MLA DSpark7 integration case to the existing eight-GPU CI stage. Propagate strict target FMHA selection to pre-existing workers, restore each worker environment, and cover success and failure cleanup with CPU regressions. Preserve existing model resources, GSM8K criteria, acceptance length, and FP32 KDA state. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../defs/accuracy/references/gsm8k.yaml | 4 + tests/integration/defs/accuracy/test_kimi3.py | 77 +++++--- .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../l0_gb300_multi_nodes_node2_gpu8.yml | 3 + tests/test_common/grouped_test_utils.py | 44 ++++- .../others/test_kimi3_ci_selection.py | 186 ++++++++++++++++++ 6 files changed, 288 insertions(+), 27 deletions(-) create mode 100644 tests/unittest/others/test_kimi3_ci_selection.py diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index fc34253a66a7..084c4af6e171 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -360,6 +360,10 @@ moonshotai/Kimi-K3: - quant_algo: MIXED_PRECISION spec_dec_algo: DSpark accuracy: 96.5 + - quant_algo: MIXED_PRECISION + kv_cache_quant_algo: FP8 + spec_dec_algo: DSpark + accuracy: 96.5 # Same requant, non-speculative, on the aggregate DEP8 recipe with # TRTLLM-Gen MoE (TestKimiK3::test_gsm8k_nvfp4_dep8). Measured 96.361 on all # 1,319 samples on SM107; 96.0 keeps margin for run-to-run variation. diff --git a/tests/integration/defs/accuracy/test_kimi3.py b/tests/integration/defs/accuracy/test_kimi3.py index 7c8bedbfe799..0425d89ed766 100644 --- a/tests/integration/defs/accuracy/test_kimi3.py +++ b/tests/integration/defs/accuracy/test_kimi3.py @@ -14,6 +14,7 @@ # limitations under the License. import pytest +from test_common.grouped_test_utils import preserve_worker_env from tensorrt_llm import LLM from tensorrt_llm._torch.configs import KimiK3Config, KimiLinearConfig @@ -168,6 +169,7 @@ def test_w4a16_mxfp4(self, mode: str, monkeypatch: pytest.MonkeyPatch) -> None: @pytest.mark.skip_less_mpi_world_size(16) @pytest.mark.skip_less_device_memory(200000) @pytest.mark.parametrize("fp8_kv", [False, True], ids=["bf16-kv", "fp8-kv"]) + @pytest.mark.private_mpi_session def test_prims_ts(self, fp8_kv: bool, monkeypatch: pytest.MonkeyPatch) -> None: """Full K3: strict PrimTS MLA context/decode, chunking, reuse and graphs. @@ -177,30 +179,34 @@ def test_prims_ts(self, fp8_kv: bool, monkeypatch: pytest.MonkeyPatch) -> None: """ monkeypatch.setenv("TLLM_FMHA_LIBS", "prims_ts") self._assert_checkpoint_routing() - with LLM( - self.MODEL_PATH, - tensor_parallel_size=16, - moe_expert_parallel_size=16, - enable_attention_dp=True, - max_batch_size=32, - # The explicit 4096-token reuse probe must span context chunks. - max_num_tokens=2048, - max_seq_len=8192, - trust_remote_code=True, - enable_chunked_prefill=True, - cuda_graph_config=CudaGraphConfig(enable_padding=True, max_batch_size=32), - moe_config=MoeConfig(max_num_tokens=33024, use_low_precision_moe_combine=True), - kv_cache_config=KvCacheConfig( - dtype="fp8" if fp8_kv else "auto", - enable_block_reuse=True, - free_gpu_memory_fraction=0.25, - tokens_per_block=64, - mamba_state_config=MambaStateConfig(periodic_snapshot_interval=256), - ), - enable_iter_perf_stats=True, - max_stats_len=256, - return_perf_metrics=True, - ) as llm: + with ( + preserve_worker_env(16, "TLLM_FMHA_LIBS"), + LLM( + self.MODEL_PATH, + tensor_parallel_size=16, + moe_expert_parallel_size=16, + enable_attention_dp=True, + max_batch_size=32, + # The explicit 4096-token reuse probe must span context chunks. + max_num_tokens=2048, + max_seq_len=8192, + trust_remote_code=True, + enable_chunked_prefill=True, + cuda_graph_config=CudaGraphConfig(enable_padding=True, max_batch_size=32), + moe_config=MoeConfig(max_num_tokens=33024, use_low_precision_moe_combine=True), + kv_cache_config=KvCacheConfig( + dtype="fp8" if fp8_kv else "auto", + enable_block_reuse=True, + free_gpu_memory_fraction=0.25, + tokens_per_block=64, + mamba_state_config=MambaStateConfig(periodic_snapshot_interval=256), + ), + enable_iter_perf_stats=True, + max_stats_len=256, + return_perf_metrics=True, + env_overrides={"TLLM_FMHA_LIBS": "prims_ts"}, + ) as llm, + ): self._assert_resolved_args(llm, "reuse") assert llm.args.quant_config.kv_cache_quant_algo == (QuantAlgo.FP8 if fp8_kv else None) assert_kv_cache_reuse_for_llm(llm, [1] + [42] * 4094 + [43]) @@ -454,7 +460,7 @@ class TestKimiK3DSpark(LlmapiAccuracyTestHarness): DSPARK_MODEL_PATH = f"{llm_models_root()}/RadixArk-Kimi-K3-DSpark" @pytest.mark.skip_less_mpi_world_size(8) - def test_gsm8k_tep8(self): + def test_gsm8k_tep8(self) -> None: """GSM8K plus an acceptance-length gate on the TEP8 DSpark recipe. TEP, not TestKimiK3's DEP16 recipe: with attention DP off the drafter gets @@ -463,8 +469,24 @@ def test_gsm8k_tep8(self): a drafter producing garbage scores the same and only runs slower -- hence the acceptance-length assertion. """ + self._run_gsm8k_tep8(fp8_kv=False) + + @pytest.mark.skip_less_mpi_world_size(8) + @pytest.mark.private_mpi_session + def test_prims_ts_fp8_tep8(self, monkeypatch: pytest.MonkeyPatch) -> None: + """Qualify strict PrimTS FP8 MLA prefill/MTP decode with FP32 KDA state.""" + monkeypatch.setenv("TLLM_FMHA_LIBS", "prims_ts") + with preserve_worker_env(8, "TLLM_FMHA_LIBS"): + self._run_gsm8k_tep8(fp8_kv=True, env_overrides={"TLLM_FMHA_LIBS": "prims_ts"}) + + def _run_gsm8k_tep8(self, *, fp8_kv: bool, env_overrides: dict[str, str] | None = None) -> None: + """Run the shared DSpark7 workload and its baseline acceptance gate.""" kv_cache_config = KvCacheConfig( - enable_block_reuse=False, free_gpu_memory_fraction=0.20, tokens_per_block=64 + dtype="fp8" if fp8_kv else "auto", + enable_block_reuse=False, + free_gpu_memory_fraction=0.20, + tokens_per_block=64, + mamba_ssm_cache_dtype="float32", ) spec_config = DSparkDecodingConfig( max_draft_len=7, speculative_model=self.DSPARK_MODEL_PATH, attention_backend="TRTLLM" @@ -493,11 +515,14 @@ def test_gsm8k_tep8(self): max_stats_len=-1, enable_iter_perf_stats=True, speculative_config=spec_config, + env_overrides=env_overrides, ) as llm: # Unlike TestKimiK3's MXFP4 weights, the requant ships # hf_quant_config.json, so the matcher sees MIXED_PRECISION -- # only the experts are 4-bit. assert llm.args.quant_config.quant_algo == QuantAlgo.MIXED_PRECISION + assert llm.args.quant_config.kv_cache_quant_algo == (QuantAlgo.FP8 if fp8_kv else None) + assert llm.args.kv_cache_config.mamba_ssm_cache_dtype == "float32" task = GSM8K(self.MODEL_NAME) task.evaluate(llm) assert_acceptance_length_for_llm( diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 550796e3fa0c..3299cdcb1459 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -38,6 +38,7 @@ l0_cpu: - model_express/test_model_express_accuracy.py::test_mx_accuracy_cases_have_references - unittest/test_s3_output.py - unittest/others/test_http_utils_fail_fast.py + - unittest/others/test_kimi3_ci_selection.py - unittest/_torch/visual_gen/multi_gpu/test_worker_pool.py # Covers the cpu_only cases of tests/unittest/_torch/attention/. On this CPU-only list, # pytest_ignore_collect in tests/unittest/conftest.py drops every file without a diff --git a/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml b/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml index 522774da2c26..b9d4843e0026 100644 --- a/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml +++ b/tests/integration/test_lists/test-db/l0_gb300_multi_nodes_node2_gpu8.yml @@ -1,3 +1,5 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 version: 0.0.1 l0_gb300_multi_nodes_node2_gpu8: - condition: @@ -17,6 +19,7 @@ l0_gb300_multi_nodes_node2_gpu8: # Kimi K3 GSM8K accuracy plus an acceptance-length gate on the TEP8 DSpark # recipe. Measured 25m45s on 2x4 GB300. - accuracy/test_kimi3.py::TestKimiK3DSpark::test_gsm8k_tep8 TIMEOUT (60) + - accuracy/test_kimi3.py::TestKimiK3DSpark::test_prims_ts_fp8_tep8 TIMEOUT (60) # The MiniMax-M3 modeling_v2 tp8 target (two trays): GSM8K, and with Eagle3 GSM8K plus a chat acceptance gate. - accuracy/modeling_v2/test_minimax_m3.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp8::test_gsm8k TIMEOUT (60) - accuracy/modeling_v2/test_minimax_m3.py::TestModelingV2MiniMaxM3Nvfp4Sm103Tp8::test_eagle3_gsm8k_and_acceptance TIMEOUT (60) diff --git a/tests/test_common/grouped_test_utils.py b/tests/test_common/grouped_test_utils.py index bfca9d0e8d2d..5fc1960c5a8d 100644 --- a/tests/test_common/grouped_test_utils.py +++ b/tests/test_common/grouped_test_utils.py @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Worker-side helpers for tests that reuse one ``MpiPoolSession``. +"""Worker-side helpers for tests that reuse MPI workers. Used by the automatic session-reuse layer (``test_common/session_reuse.py``). This is test infrastructure only and must not be used on production inference @@ -8,6 +8,48 @@ circular import during package initialization. """ +from collections.abc import Iterator +from contextlib import contextmanager + + +@contextmanager +def preserve_worker_env(n_workers: int, name: str) -> Iterator[None]: + """Restore a variable on launcher-owned ranks after the executor shuts down. + + Callers must use ``private_mpi_session`` for locally spawned workers. + """ + from tensorrt_llm.executor.utils import create_mpi_comm_session, get_spawn_proxy_process_env + + session = create_mpi_comm_session(n_workers) if get_spawn_proxy_process_env() else None + previous = dict(session.submit_sync(get_worker_env, name)) if session else {} + try: + yield + finally: + if session: + session.submit_sync(restore_worker_env, name, previous) + + +def get_worker_env(name: str) -> tuple[int, str | None]: + """Read an environment variable with its rank for order-independent restoration.""" + import os + + from mpi4py import MPI + + return MPI.COMM_WORLD.Get_rank(), os.environ.get(name) + + +def restore_worker_env(name: str, previous: dict[int, str | None]) -> None: + """Restore this rank's environment variable, including its original absence.""" + import os + + from mpi4py import MPI + + value = previous[MPI.COMM_WORLD.Get_rank()] + if value is None: + os.environ.pop(name, None) + else: + os.environ[name] = value + def _ranked_run(fn): """Runs inside a worker; module-level so it is picklable.""" diff --git a/tests/unittest/others/test_kimi3_ci_selection.py b/tests/unittest/others/test_kimi3_ci_selection.py new file mode 100644 index 000000000000..c566d054cb88 --- /dev/null +++ b/tests/unittest/others/test_kimi3_ci_selection.py @@ -0,0 +1,186 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU checks for the Kimi K3 strict FMHA test's remote-worker configuration.""" + +import ast +import os +import pickle +import sys +from collections.abc import Callable +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +from test_common.grouped_test_utils import preserve_worker_env + +pytestmark = pytest.mark.cpu_only + +_FMHA_ENV = "TLLM_FMHA_LIBS" +_KIMI_TEST = Path(__file__).resolve().parents[2] / "integration/defs/accuracy/test_kimi3.py" + + +def _load_test_case(class_name: str, namespace: dict[str, object]) -> type: + """Load the real test methods without importing the GPU integration harness.""" + tree = ast.parse(_KIMI_TEST.read_text()) + case = next( + node for node in tree.body if isinstance(node, ast.ClassDef) and node.name == class_name + ) + case.decorator_list = [] + for method in case.body: + if isinstance(method, ast.FunctionDef): + if method.name in ("test_prims_ts", "test_prims_ts_fp8_tep8"): + assert "pytest.mark.private_mpi_session" in ( + ast.unparse(decorator) for decorator in method.decorator_list + ) + method.decorator_list = [] + namespace.update( + pytest=pytest, + LlmapiAccuracyTestHarness=object, + llm_models_root=lambda: "/unused-models", + KvCacheConfig=SimpleNamespace, + DSparkDecodingConfig=SimpleNamespace, + CudaGraphConfig=SimpleNamespace, + MambaStateConfig=SimpleNamespace, + MoeConfig=SimpleNamespace, + QuantAlgo=SimpleNamespace(MIXED_PRECISION="mixed", FP8="fp8"), + ) + exec(compile(ast.Module(body=[case], type_ignores=[]), str(_KIMI_TEST), "exec"), namespace) + return namespace[class_name] + + +class _RemoteWorkers: + """Pre-existing rank environments that do not inherit subsequent parent changes.""" + + def __init__(self, n_workers: int) -> None: + self.environments = [ + {} if value is None else {_FMHA_ENV: value} + for value in (None, "", "fallback", "+prims_ts") * (n_workers // 4) + ] + self.rank = 0 + + def submit_sync(self, callback: Callable[..., object], *args: object) -> list[object]: + # RemoteTask and the MPI transport use standard pickle, not cloudpickle. + callback, args = pickle.loads(pickle.dumps((callback, args))) + results = [] + for self.rank, environment in enumerate(self.environments): + with patch.dict(os.environ, environment, clear=True): + try: + results.append(callback(*args)) + finally: + environment.clear() + environment.update(os.environ) + return list(reversed(results)) + + def selections(self) -> list[str | None]: + return [environment.get(_FMHA_ENV) for environment in self.environments] + + +@pytest.mark.parametrize("failure", [None, "constructor", "evaluation"]) +@pytest.mark.parametrize("variant", ["dspark", "bf16-kv", "fp8-kv"]) +def test_strict_selection_reaches_existing_workers_and_is_restored( + monkeypatch: pytest.MonkeyPatch, failure: str | None, variant: str +) -> None: + n_workers = 8 if variant == "dspark" else 16 + workers = _RemoteWorkers(n_workers) + original = workers.selections() + monkeypatch.setitem( + sys.modules, + "mpi4py", + SimpleNamespace( + MPI=SimpleNamespace(COMM_WORLD=SimpleNamespace(Get_rank=lambda: workers.rank)) + ), + ) + monkeypatch.setenv(_FMHA_ENV, "fallback") + observed = [] + strict_run = False + + class FakeLLM: + def __init__(self, _model: str, **kwargs: object) -> None: + self.strict = strict_run + assert kwargs["tensor_parallel_size"] == n_workers + kv_dtype = getattr(kwargs["kv_cache_config"], "dtype", "auto") + if self.strict: + assert kv_dtype == ("auto" if variant == "bf16-kv" else "fp8") + overrides = kwargs.get("env_overrides") or {} + # worker_main receives only the serialized LLM args, not the + # pytest process's current environment. + for environment in workers.environments: + environment.update(overrides) + observed.append(workers.selections()) + assert observed[-1] == (["prims_ts"] * len(original) if self.strict else original) + os.environ.update(overrides) + if self.strict and failure == "constructor": + raise RuntimeError("constructor failed") + self.args = SimpleNamespace( + kv_cache_config=kwargs["kv_cache_config"], + quant_config=SimpleNamespace( + quant_algo="mixed", kv_cache_quant_algo="fp8" if kv_dtype == "fp8" else None + ), + ) + + def __enter__(self) -> "FakeLLM": + return self + + def __exit__(self, *_args: object) -> None: + pass + + def evaluate(llm: FakeLLM) -> None: + if llm.strict and failure == "evaluation": + raise RuntimeError("evaluation failed") + + def create_session(size: int) -> _RemoteWorkers: + assert size == n_workers + return workers + + monkeypatch.setitem( + sys.modules, + "tensorrt_llm.executor.utils", + SimpleNamespace( + create_mpi_comm_session=create_session, + get_spawn_proxy_process_env=lambda: True, + ), + ) + case_type = _load_test_case( + "TestKimiK3DSpark" if variant == "dspark" else "TestKimiK3", + dict( + LLM=FakeLLM, + GSM8K=lambda _model: SimpleNamespace(evaluate=evaluate), + assert_acceptance_length_for_llm=lambda *_args: None, + assert_guided_decoding_regex=lambda *_args: None, + assert_kv_cache_reuse_for_llm=lambda *_args: None, + preserve_worker_env=preserve_worker_env, + ), + ) + case = case_type() + case._assert_checkpoint_routing = lambda: None + case._assert_resolved_args = lambda *_args: None + case._assert_attention_dp = lambda *_args: None + case._assert_logits_processor = lambda *_args: None + + def run_baseline() -> None: + if variant == "dspark": + case.test_gsm8k_tep8() + else: + with monkeypatch.context() as baseline_patch: + case.test_w4a16_mxfp4("baseline", baseline_patch) + + def run_strict(strict_patch: pytest.MonkeyPatch) -> None: + if variant == "dspark": + case.test_prims_ts_fp8_tep8(strict_patch) + else: + case.test_prims_ts(variant == "fp8-kv", strict_patch) + + run_baseline() + strict_run = True + with monkeypatch.context() as strict_patch: + if failure: + with pytest.raises(RuntimeError, match=f"{failure} failed"): + run_strict(strict_patch) + else: + run_strict(strict_patch) + assert os.environ[_FMHA_ENV] == "fallback" + assert workers.selections() == original + strict_run = False + run_baseline() + assert observed == [original, ["prims_ts"] * len(original), original] From f5f0ba48b3f1d15beac0ffb5e19eac7533d7b4ea Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Sun, 1 Nov 2026 10:52:52 +0000 Subject: [PATCH 11/19] [None][test] Align MLA page-size routing coverage with Rubin support Preserve the static PrimTS capability contract after rebasing onto the shared MLA routing helper. SM107 FP8 pages are not exclusive to CuTe DSL once PrimTS supports them. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- tests/unittest/_torch/attention/test_cute_dsl_mla.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unittest/_torch/attention/test_cute_dsl_mla.py b/tests/unittest/_torch/attention/test_cute_dsl_mla.py index 63a9dc77ad06..33119bb4a8f1 100644 --- a/tests/unittest/_torch/attention/test_cute_dsl_mla.py +++ b/tests/unittest/_torch/attention/test_cute_dsl_mla.py @@ -387,10 +387,10 @@ def test_mixed_batch_gate_admits_only_page_sizes_exclusive_to_cute_dsl( # FlashInfer TRTLLM-Gen reads the manager's extra page sizes. (107, 128, DataType.FP8, (128,), False), (100, 128, DataType.NVFP4, (128,), False), - # PrimTS reads 128-token pages of FP8 or BF16 managers on SM100/SM103. + # PrimTS reads 128-token pages of FP8 or BF16 managers on SM100/SM103/SM107. (100, 128, DataType.FP8, (), False), (103, 128, DataType.BF16, (), False), - (107, 128, DataType.FP8, (), True), + (107, 128, DataType.FP8, (), False), # An NVFP4 manager's FP8 fallback layers on the FP4 MLA pages. (100, 128, DataType.NVFP4, (), True), (103, 128, DataType.NVFP4, (), True), From e486a1789799ce532aedcda11c7fca518846b863 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Fri, 30 Oct 2026 05:23:49 +0000 Subject: [PATCH 12/19] [None][perf] Normalize dense MLA context capacity buckets Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../attention/backends/fmha/prims_ts.py | 50 +++++++- .../_torch/attention/test_prims_ts_fmha.py | 2 + .../attention/test_prims_ts_mla_capacity.py | 120 ++++++++++++++++++ 3 files changed, 170 insertions(+), 2 deletions(-) create mode 100644 tests/unittest/_torch/attention/test_prims_ts_mla_capacity.py diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py index 6426ec084a57..db8d111089af 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py @@ -67,6 +67,41 @@ _WORKSPACE_ALIGNMENT = 32 +def _mla_context_capacities( + query_length: int, + kv_length: int, + *, + head_dim: int, + head_dim_vo: int | None, + qk_dtype: torch.dtype, + pv_dtype: torch.dtype, + output_dtype: torch.dtype, + packed: bool, + uniform_packed_lengths: bool, + mask_type: str, + window_left: int, +) -> tuple[int, int]: + """Bucket expanded MLA capacities without changing its live packed lengths.""" + query_capacity = 1 << (max(query_length, 1) - 1).bit_length() + kv_capacity = 1 << (max(kv_length, 1) - 1).bit_length() + if ( + (head_dim, head_dim_vo) == (192, 128) + and packed + and not uniform_packed_lengths + and mask_type in ("causal", "dense") + and window_left == -1 + and qk_dtype == pv_dtype + and qk_dtype in (torch.float16, torch.bfloat16, torch.float8_e4m3fn) + and output_dtype in (torch.float16, torch.bfloat16) + ): + # This query-paired kernel has two M128 peers and N128 KV tiles. + # Raising sub-tile capacities preserves the work domain; indptrs + # continue to supply the actual ragged lengths and causal offsets. + query_capacity = max(query_capacity, 256) + kv_capacity = max(kv_capacity, 128) + return query_capacity, kv_capacity + + def get_paged_kv_storage_unsupported_reason( attn: "TrtllmAttention", metadata: "TrtllmAttentionMetadata", @@ -1290,8 +1325,19 @@ def run_mla_context(self, params: FmhaParams) -> None: stats = fwd.softmax_stats_tensor # Dense cached chunks traverse the planned KV domain. Bucket current # capacities instead of traversing max_seq_len for every small chunk. - q_capacity = 1 << (max(params.input_seq_length, 1) - 1).bit_length() - kv_capacity = 1 << (max(params.max_past_kv_length, 1) - 1).bit_length() + q_capacity, kv_capacity = _mla_context_capacities( + params.input_seq_length, + params.max_past_kv_length, + head_dim=192, + head_dim_vo=128, + qk_dtype=q_processed.dtype, + pv_dtype=v_processed.dtype, + output_dtype=params.output.dtype, + packed=True, + uniform_packed_lengths=False, + mask_type=mask, + window_left=-1, + ) key = (params.batch_size, q_capacity, kv_capacity, mask, stats is not None) wrapper = self._mla_context_wrappers.get(key) if wrapper is None: diff --git a/tests/unittest/_torch/attention/test_prims_ts_fmha.py b/tests/unittest/_torch/attention/test_prims_ts_fmha.py index 1a42c15c48c9..62cbbae7e7cc 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_fmha.py +++ b/tests/unittest/_torch/attention/test_prims_ts_fmha.py @@ -1233,6 +1233,8 @@ def test_mla_context_wrapper_reuses_plan_and_updates_chunk_offsets( assert wrapper.run.call_args.kwargs["out"] is output assert wrapper.run.call_args.args[1].shape == (9, num_heads, 192) assert wrapper.plan.call_args.kwargs["store_softmax_stats"] + assert wrapper.plan.call_args.kwargs["max_seq_len_q"] == 256 + assert wrapper.plan.call_args.kwargs["max_kv_len"] == 128 processed_dtype = torch.float8_e4m3fn if fp8 else q.dtype for name in ("q_dtype", "k_dtype", "v_dtype"): assert wrapper.plan.call_args.kwargs[name] == processed_dtype diff --git a/tests/unittest/_torch/attention/test_prims_ts_mla_capacity.py b/tests/unittest/_torch/attention/test_prims_ts_mla_capacity.py new file mode 100644 index 000000000000..7432370c11e7 --- /dev/null +++ b/tests/unittest/_torch/attention/test_prims_ts_mla_capacity.py @@ -0,0 +1,120 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only capacity contracts, runnable without the TensorRT-LLM extension.""" + +import ast +from collections.abc import Callable +from pathlib import Path + +import pytest +import torch + +_ROOT = Path(__file__).resolve().parents[4] +_ADAPTER = _ROOT / "tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py" + + +def _load_capacities() -> Callable[..., tuple[int, int]]: + tree = ast.parse(_ADAPTER.read_text()) + function = next( + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == "_mla_context_capacities" + ) + namespace = {"torch": torch} + exec(compile(ast.Module(body=[function], type_ignores=[]), str(_ADAPTER), "exec"), namespace) + return namespace[function.name] + + +_capacities = _load_capacities() +_GEOMETRY = dict( + head_dim=192, + head_dim_vo=128, + qk_dtype=torch.float8_e4m3fn, + pv_dtype=torch.float8_e4m3fn, + output_dtype=torch.bfloat16, + packed=True, + uniform_packed_lengths=False, + mask_type="causal", + window_left=-1, +) + + +@pytest.mark.parametrize("query_length", [1, 127, 128, 129, 255, 256, 257, 2048]) +@pytest.mark.parametrize("kv_length", [0, 1, 127, 128, 129, 255, 256, 257, 4607]) +@pytest.mark.parametrize("mask_type", ["causal", "dense"]) +def test_capacity_floor_preserves_tile_counts( + query_length: int, kv_length: int, mask_type: str +) -> None: + actual = _capacities(query_length, kv_length, **(_GEOMETRY | {"mask_type": mask_type})) + original = tuple(1 << (max(length, 1) - 1).bit_length() for length in (query_length, kv_length)) + assert actual == (max(original[0], 256), max(original[1], 128)) + assert (actual[0] + 255) // 256 == (original[0] + 255) // 256 + assert (actual[1] + 127) // 128 == (original[1] + 127) // 128 + assert actual[0] >= query_length and actual[1] >= kv_length + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float8_e4m3fn]) +@pytest.mark.parametrize("output_dtype", [torch.float16, torch.bfloat16]) +def test_capacity_floor_supported_precisions(dtype: torch.dtype, output_dtype: torch.dtype) -> None: + assert _capacities( + 1, + 1, + **(_GEOMETRY | {"qk_dtype": dtype, "pv_dtype": dtype, "output_dtype": output_dtype}), + ) == (256, 128) + + +@pytest.mark.parametrize( + "override", + [ + {"head_dim": 128}, + {"head_dim_vo": 192}, + {"packed": False}, + {"uniform_packed_lengths": True}, + {"mask_type": "variable_window"}, + {"window_left": 128}, + {"qk_dtype": torch.bfloat16}, + {"pv_dtype": torch.bfloat16}, + {"qk_dtype": torch.float32, "pv_dtype": torch.float32}, + {"output_dtype": torch.float8_e4m3fn}, + ], +) +def test_unproved_geometry_keeps_original_buckets(override: dict) -> None: + assert _capacities(3, 5, **(_GEOMETRY | override)) == (4, 8) + + +def test_runtime_key_and_plan_use_the_same_normalized_capacities() -> None: + tree = ast.parse(_ADAPTER.read_text()) + method = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.FunctionDef) and node.name == "run_mla_context" + ) + assignment = next( + node + for node in method.body + if isinstance(node, ast.Assign) + and isinstance(node.value, ast.Call) + and isinstance(node.value.func, ast.Name) + and node.value.func.id == "_mla_context_capacities" + ) + assert ast.unparse(assignment.targets[0]) == "(q_capacity, kv_capacity)" + key = next( + node + for node in method.body + if isinstance(node, ast.Assign) and ast.unparse(node.targets[0]) == "key" + ) + assert ( + ast.unparse(key.value) + == "(params.batch_size, q_capacity, kv_capacity, mask, stats is not None)" + ) + plan = next( + node + for node in ast.walk(method) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "plan" + ) + keywords = {keyword.arg: ast.unparse(keyword.value) for keyword in plan.keywords} + assert keywords["max_seq_len_q"] == "q_capacity" + assert keywords["max_kv_len"] == "kv_capacity" + assert keywords["packed"] == "True" From c14e0abc430a153cb8d8391a7f2b5dd914a5eeeb Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Fri, 30 Oct 2026 05:27:37 +0000 Subject: [PATCH 13/19] [None][perf] Prewarm selected dense MLA profiles before readiness Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../backends/fmha/prims_ts_mla_warmup.py | 283 ++++++++++++++ .../engine/runners/decoder/warmup.py | 10 + .../attention/test_prims_ts_mla_warmup.py | 365 ++++++++++++++++++ 3 files changed, 658 insertions(+) create mode 100644 tensorrt_llm/_torch/attention/backends/fmha/prims_ts_mla_warmup.py create mode 100644 tests/unittest/_torch/attention/test_prims_ts_mla_warmup.py diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_mla_warmup.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_mla_warmup.py new file mode 100644 index 000000000000..f66999360c13 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_mla_warmup.py @@ -0,0 +1,283 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Startup-only planning for dense MLA profiles already selected by PrimTS.""" + +from __future__ import annotations + +from collections.abc import Callable, Hashable +from functools import partial +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import torch + + from tensorrt_llm._torch.attention.backends.prims_ts.context import _ContextPlanGeometry + from tensorrt_llm._torch.attention.backends.prims_ts.mla_decode import _MLADecodePlanState + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata + + +def _bucket_minima(limit: int) -> list[int]: + """Smallest positive length in each power-of-two bucket up to limit.""" + if limit < 1: + return [] + return [1] + [(1 << exponent) + 1 for exponent in range((limit - 1).bit_length())] + + +def _context_profiles( + *, + max_num_tokens: int, + max_seq_len: int, + cached_context: bool, + chunked_prefill: bool, + chunk_size: int, + chunk_batch_size: int, + capacities: Callable[[int, int], tuple[int, int]], +) -> list[tuple[int, int, str, bool]]: + """Enumerate the normalized fresh, cached-prefix and current-chunk envelope.""" + queries = _bucket_minima(min(max_num_tokens, max_seq_len)) + profiles = {(*capacities(query, query), "causal", False) for query in queries} + if not cached_context: + return sorted(profiles) + if chunked_prefill: + if chunk_size <= 0 or chunk_batch_size <= 0: + raise ValueError("MLA chunked-prefill warmup requires a positive staging capacity.") + max_cached_kv = min(max_seq_len - 1, chunk_size * chunk_batch_size) + for query in queries: + if query < max_seq_len: + profiles.add((*capacities(query, query), "causal", True)) + # Runtime takes these maxima independently across requests, so + # their sum need not fit within one request's sequence limit. + for cached_kv in _bucket_minima(max_cached_kv): + profiles.add((*capacities(query, cached_kv), "dense", True)) + else: + for query in queries: + for kv_length in _bucket_minima(max_seq_len): + if query <= kv_length: + profiles.add((*capacities(query, kv_length), "causal", False)) + return sorted(profiles) + + +class PrimsTSDenseMlaWarmup: + """Model-lifetime compile deduplication, with no retained runtime workspace. + + Context covers the configured normalized envelope. Decode expands batch + counts only for widths and modes observed in this model's startup plans; + arbitrary future proposal widths from variable drafters are not inferred. + """ + + def __init__(self) -> None: + self._completed: set[Hashable] = set() + + def run( + self, + model: torch.nn.Module, + metadata: TrtllmAttentionMetadata | None, + *, + max_batch_size: int, + max_num_tokens: int, + max_seq_len: int, + ) -> tuple[int, int]: + """Load selected providers' missing profiles outside CUDA graph capture.""" + import torch + + from tensorrt_llm._torch.attention.backends.trtllm import ( + TrtllmAttention, + TrtllmAttentionMetadata, + ) + from tensorrt_llm._torch.attention.mla import MLA + from tensorrt_llm.logger import logger + + from .prims_ts import PrimsTSFmha + + if not isinstance(metadata, TrtllmAttentionMetadata): + return 0, 0 + context_geometries: list[_ContextPlanGeometry] = [] + decode_states: list[_MLADecodePlanState] = [] + for module in model.modules(): + if not isinstance(module, MLA) or module.sparse_attn_hooks is not None: + continue + for backend in (module.mha, module.mqa): + if not isinstance(backend, TrtllmAttention): + continue + for library in backend._fmha_manager.fmha_libs: + if not isinstance(library, PrimsTSFmha): + continue + for context_wrapper in library._mla_context_wrappers.values(): + context_state = context_wrapper._plan_state + if ( + context_state is not None + and context_state.geometry.device.type == "cuda" + ): + context_geometries.append(context_state.geometry) + for decode_wrapper in library._mla_decode_wrappers.values(): + decode_state = decode_wrapper._plan_state + if decode_state is not None and decode_state.device.type == "cuda": + decode_states.append(decode_state) + if not context_geometries and not decode_states: + return 0, 0 + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("PrimTS dense MLA startup planning cannot run during capture.") + + contexts = decodes = 0 + widths = set() + for geometry in context_geometries: + contexts += self._prewarm_context( + geometry, + metadata, + max_num_tokens=max_num_tokens, + max_seq_len=max_seq_len, + ) + for state in decode_states: + widths.add(state.max_seq_len_q) + decodes += self._prewarm_decode( + state, max_batch_size=max_batch_size, max_num_tokens=max_num_tokens + ) + logger.info( + f"PrimTS dense MLA prewarm loaded {contexts} context profiles and {decodes} " + f"decode topologies; observed decode query widths: {sorted(widths)}." + ) + return contexts, decodes + + def _prewarm_context( + self, + geometry: _ContextPlanGeometry, + metadata: TrtllmAttentionMetadata, + *, + max_num_tokens: int, + max_seq_len: int, + ) -> int: + from tensorrt_llm._torch.attention.backends.prims_ts import context + + from .prims_ts import _mla_context_capacities + + capacities = partial( + _mla_context_capacities, + head_dim=geometry.head_dim, + head_dim_vo=geometry.head_dim_vo, + qk_dtype=geometry.qk_dtype, + pv_dtype=geometry.pv_dtype, + output_dtype=geometry.output_dtype, + packed=geometry.packed, + uniform_packed_lengths=geometry.uniform_packed_lengths, + mask_type=geometry.mask_type, + window_left=geometry.window_left, + ) + # Only the proven query-paired layout has a bounded normalized mesh. + if geometry.num_qo_heads != geometry.num_kv_heads or capacities(1, 1) != (256, 128): + return 0 + features = metadata.runtime_features + profiles = _context_profiles( + max_num_tokens=max_num_tokens, + max_seq_len=max_seq_len, + cached_context=metadata.enable_context_mla_with_cached_kv, + chunked_prefill=features.chunked_prefill, + chunk_size=features.chunk_size, + chunk_batch_size=features.chunked_prefill_buffer_batch_size, + capacities=capacities, + ) + loaded = 0 + for query_capacity, kv_capacity, mask, stats in profiles: + kwargs = dict( + device=geometry.device, + batch_size=1, + max_seq_len_q=query_capacity, + max_kv_len=kv_capacity, + num_qo_heads=geometry.num_qo_heads, + num_kv_heads=geometry.num_kv_heads, + head_dim=geometry.head_dim, + head_dim_vo=geometry.head_dim_vo, + packed=True, + mask_type=mask, + window_left=-1, + store_softmax_stats=stats, + ) + candidate = context._resolve_context_plan_geometry( + **kwargs, + qk_dtype=geometry.qk_dtype, + pv_dtype=geometry.pv_dtype, + output_dtype=geometry.output_dtype, + ) + signature = context._context_compile_spec(candidate) + if signature in self._completed: + continue + # Normal plan() loads real callables. Probe/precompile-only modes + # must not populate the context compiler's functools cache. + context.BatchPrefillTSWrapper().plan( + **kwargs, + q_dtype=geometry.qk_dtype, + k_dtype=geometry.qk_dtype, + v_dtype=geometry.pv_dtype, + out_dtype=geometry.output_dtype, + ) + self._completed.add(signature) + loaded += 1 + return loaded + + def _prewarm_decode( + self, + state: _MLADecodePlanState, + *, + max_batch_size: int, + max_num_tokens: int, + ) -> int: + from tensorrt_llm._torch.attention.backends.prims_ts import mla_decode + + query_width = state.max_seq_len_q + max_batch = min(max_batch_size, max_num_tokens // query_width) + representatives = {} + q_dtype_key = mla_decode._dtype_key(state.q_dtype) + kv_dtype_key = mla_decode._dtype_key(state.kv_dtype) + output_dtype_key = mla_decode._dtype_key(state.output_dtype) + for batch_size in range(1, max_batch + 1): + launch = mla_decode._resolve_mla_decode_launch_spec( + state.device.index, + batch_size, + state.num_heads, + state.kv_lora_rank, + state.qk_rope_head_dim, + state.page_size, + state.max_kv_len, + q_dtype_key, + kv_dtype_key, + output_dtype_key, + state.mask_type, + query_width, + state.store_softmax_stats, + state.use_kv_bounds, + ) + signature = mla_decode._make_mla_decode_compile_spec( + launch, + device_index=state.device.index, + num_heads=state.num_heads, + kv_lora_rank=state.kv_lora_rank, + qk_rope_head_dim=state.qk_rope_head_dim, + page_size=state.page_size, + q_dtype_key=q_dtype_key, + output_dtype_key=output_dtype_key, + max_seq_len_q=query_width, + packed_query=state.packed_query, + store_softmax_stats=state.store_softmax_stats, + use_kv_bounds=state.use_kv_bounds, + ) + if signature not in self._completed: + representatives[signature] = batch_size + for signature, batch_size in representatives.items(): + mla_decode.BatchMLADecodePagedTSWrapper().plan( + state.device, + batch_size, + state.num_heads, + state.kv_lora_rank, + state.qk_rope_head_dim, + state.page_size, + state.max_kv_len, + max_seq_len_q=query_width, + packed_query=state.packed_query, + q_data_type=state.q_dtype, + kv_data_type=state.kv_dtype, + o_data_type=state.output_dtype, + mask_type=state.mask_type, + store_softmax_stats=state.store_softmax_stats, + use_kv_bounds=state.use_kv_bounds, + ) + self._completed.add(signature) + return len(representatives) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py index 0fc2fc5f62f2..1a3bf1a532e7 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py @@ -13,6 +13,7 @@ import torch +from tensorrt_llm._torch.attention.backends.fmha.prims_ts_mla_warmup import PrimsTSDenseMlaWarmup from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.autotuner import AutoTuner, autotune from tensorrt_llm._torch.distributed import Distributed @@ -275,6 +276,7 @@ def __init__( self._eager_workspace_reclaimable = eager_workspace_reclaimable # Token shape of the last autotuner warmup that ran a forward pass. self._completed_autotuner_warmup_num_tokens: int | None = None + self._prims_ts_dense_mla_warmup = PrimsTSDenseMlaWarmup() @property def guided_decoder(self) -> CapturableGuidedDecoder | None: @@ -474,6 +476,14 @@ def _warmup_scheduled( with self._warmup_timer.phase("cuda_graph_capture"): self._graphs.warmup_and_capture(resource_manager) log_mem_snapshot("warmup/after_cuda_graph_capture") + with self._warmup_timer.phase("prims_ts_dense_mla_prewarm"): + self._prims_ts_dense_mla_warmup.run( + self.model, + self._metadata.attn_metadata, + max_batch_size=self._config.max_batch_size, + max_num_tokens=self._config.max_num_tokens, + max_seq_len=self._config.max_seq_len, + ) # Pre-compile DeepGEMM paged_mqa_logits_metadata for every 32-aligned # batch bucket the runtime can produce (max_batch_size scaled by the # MTP / DSL expansion factor when applicable). CUDA-graph warmup only diff --git a/tests/unittest/_torch/attention/test_prims_ts_mla_warmup.py b/tests/unittest/_torch/attention/test_prims_ts_mla_warmup.py new file mode 100644 index 000000000000..83e4f168e519 --- /dev/null +++ b/tests/unittest/_torch/attention/test_prims_ts_mla_warmup.py @@ -0,0 +1,365 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Startup planning contracts without importing the TensorRT-LLM extension.""" + +import ast +import importlib.util +import sys +from collections.abc import Callable +from functools import partial +from pathlib import Path +from types import ModuleType, SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +_ROOT = Path(__file__).resolve().parents[4] +_FMHA = _ROOT / "tensorrt_llm/_torch/attention/backends/fmha" +_MODULE_NAME = "tensorrt_llm._torch.attention.backends.fmha.prims_ts_mla_warmup" +_SPEC = importlib.util.spec_from_file_location(_MODULE_NAME, _FMHA / "prims_ts_mla_warmup.py") +_warmup = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(_warmup) + + +def _capacity_function() -> Callable[..., tuple[int, int]]: + source = _FMHA / "prims_ts.py" + function = next( + node + for node in ast.parse(source.read_text()).body + if isinstance(node, ast.FunctionDef) and node.name == "_mla_context_capacities" + ) + namespace = {"torch": torch} + exec(compile(ast.Module(body=[function], type_ignores=[]), str(source), "exec"), namespace) + return namespace[function.name] + + +def _geometry(**overrides: object) -> SimpleNamespace: + fields = dict( + device=torch.device("cuda:0"), + batch_size=4, + max_seq_len_q=2048, + max_seq_len_k=2048, + num_qo_heads=96, + num_kv_heads=96, + head_dim=192, + head_dim_vo=128, + qk_dtype=torch.float8_e4m3fn, + pv_dtype=torch.float8_e4m3fn, + output_dtype=torch.bfloat16, + packed=True, + uniform_packed_lengths=False, + mask_type="causal", + window_left=-1, + store_softmax_stats=False, + ) + return SimpleNamespace(**(fields | overrides)) + + +def _capacities() -> Callable[[int, int], tuple[int, int]]: + geometry = vars(_geometry()) + fields = ( + "head_dim", + "head_dim_vo", + "qk_dtype", + "pv_dtype", + "output_dtype", + "packed", + "uniform_packed_lengths", + "mask_type", + "window_left", + ) + return partial(_capacity_function(), **{name: geometry[name] for name in fields}) + + +def _profiles(**overrides: object) -> list[tuple[int, int, str, bool]]: + fields = dict( + max_num_tokens=2048, + max_seq_len=4608, + cached_context=True, + chunked_prefill=True, + chunk_size=2048, + chunk_batch_size=4, + capacities=_capacities(), + ) + return _warmup._context_profiles(**(fields | overrides)) + + +def test_normalized_context_envelope_includes_tails_and_three_families() -> None: + profiles = _profiles() + assert len(profiles) == 38 + assert (256, 128, "causal", False) in profiles + assert (256, 128, "causal", True) in profiles + assert (256, 4096, "dense", True) in profiles + assert (512, 8192, "dense", True) in profiles + assert (1024, 8192, "dense", True) in profiles + assert len(set(profiles)) == len(profiles) + + +def test_context_envelope_covers_independent_mixed_batch_maxima() -> None: + max_num_tokens, max_seq_len = 2048, 4608 + chunk_size, chunk_batch_size = 2048, 4 + query_lengths = torch.tensor([1024, 128], device="cpu") + cached_kv_lengths = torch.tensor([0, 4224], device="cpu") + assert int(query_lengths.sum()) <= max_num_tokens + assert bool(torch.all(query_lengths + cached_kv_lengths <= max_seq_len)) + assert int(cached_kv_lengths.sum()) <= chunk_size * chunk_batch_size + + max_query = int(query_lengths.max()) + max_cached_kv = int(cached_kv_lengths.max()) + assert max_query + max_cached_kv > max_seq_len + profile = (*_capacities()(max_query, max_cached_kv), "dense", True) + assert profile == (1024, 8192, "dense", True) + assert profile in _profiles( + max_num_tokens=max_num_tokens, + max_seq_len=max_seq_len, + chunk_size=chunk_size, + chunk_batch_size=chunk_batch_size, + ) + + +def test_context_envelope_covers_every_small_reachable_length() -> None: + capacities = _capacities() + profiles = set( + _profiles(max_num_tokens=257, max_seq_len=513, chunk_size=129, chunk_batch_size=1) + ) + for query in range(1, 258): + assert (*capacities(query, query), "causal", False) in profiles + assert (*capacities(query, query), "causal", True) in profiles + for kv_length in range(130): + assert (*capacities(query, kv_length), "dense", True) in profiles + + +def test_context_profiles_respect_staging_and_cache_modes() -> None: + assert len(_profiles(cached_context=False)) == 5 + limited = _profiles(chunk_size=128, chunk_batch_size=1) + assert all(kv == 128 for _, kv, mask, _ in limited if mask == "dense") + unchunked = _profiles(chunked_prefill=False) + assert (256, 8192, "causal", False) in unchunked + assert all(mask == "causal" and not stats for _, _, mask, stats in unchunked) + assert _profiles(max_seq_len=1) == [(256, 128, "causal", False)] + with pytest.raises(ValueError, match="positive staging capacity"): + _profiles(chunk_size=0) + + +@pytest.fixture +def planners(monkeypatch: pytest.MonkeyPatch) -> SimpleNamespace: + for name in ( + "tensorrt_llm", + "tensorrt_llm._torch", + "tensorrt_llm._torch.attention", + "tensorrt_llm._torch.attention.backends", + "tensorrt_llm._torch.attention.backends.fmha", + "tensorrt_llm._torch.attention.backends.prims_ts", + ): + module = ModuleType(name) + module.__path__ = [] + monkeypatch.setitem(sys.modules, name, module) + context = SimpleNamespace() + context._resolve_context_plan_geometry = Mock(side_effect=lambda **kwargs: kwargs) + context._context_compile_spec = Mock( + side_effect=lambda geometry: ("context", tuple(sorted(geometry.items()))) + ) + context.plan = Mock() + context.BatchPrefillTSWrapper = Mock(side_effect=lambda: SimpleNamespace(plan=context.plan)) + decode = SimpleNamespace() + decode._dtype_key = str + decode._resolve_mla_decode_launch_spec = Mock(side_effect=lambda *args: (args[1] % 3, args[0])) + decode._make_mla_decode_compile_spec = Mock( + side_effect=lambda launch, **kwargs: ("decode", launch, tuple(sorted(kwargs.items()))) + ) + decode.plan = Mock() + decode.BatchMLADecodePagedTSWrapper = Mock( + side_effect=lambda: SimpleNamespace(plan=decode.plan) + ) + package = sys.modules["tensorrt_llm._torch.attention.backends.prims_ts"] + package.context, package.mla_decode = context, decode + + class MLA: + sparse_attn_hooks = None + + class TrtllmAttention: + pass + + class Metadata: + enable_context_mla_with_cached_kv = True + runtime_features = SimpleNamespace( + chunked_prefill=True, chunk_size=2048, chunked_prefill_buffer_batch_size=4 + ) + + class PrimsTSFmha: + pass + + prims = ModuleType("tensorrt_llm._torch.attention.backends.fmha.prims_ts") + prims._mla_context_capacities = _capacity_function() + prims.PrimsTSFmha = PrimsTSFmha + monkeypatch.setitem(sys.modules, prims.__name__, prims) + trtllm = ModuleType("tensorrt_llm._torch.attention.backends.trtllm") + trtllm.TrtllmAttention, trtllm.TrtllmAttentionMetadata = TrtllmAttention, Metadata + monkeypatch.setitem(sys.modules, trtllm.__name__, trtllm) + mla = ModuleType("tensorrt_llm._torch.attention.mla") + mla.MLA = MLA + monkeypatch.setitem(sys.modules, mla.__name__, mla) + log = ModuleType("tensorrt_llm.logger") + log.logger = SimpleNamespace(info=Mock()) + monkeypatch.setitem(sys.modules, log.__name__, log) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + return SimpleNamespace( + context=context, + decode=decode, + MLA=MLA, + backend=TrtllmAttention, + metadata=Metadata(), + provider=PrimsTSFmha, + ) + + +def _decode_state(**overrides: object) -> SimpleNamespace: + fields = dict( + device=torch.device("cuda:0"), + batch_size=1, + num_heads=96, + kv_lora_rank=512, + qk_rope_head_dim=64, + page_size=128, + max_kv_len=4608, + max_seq_len_q=6, + packed_query=False, + q_dtype=torch.float8_e4m3fn, + kv_dtype=torch.float8_e4m3fn, + output_dtype=torch.bfloat16, + mask_type="causal", + store_softmax_stats=False, + use_kv_bounds=False, + workspace_buffer=object(), + ) + return SimpleNamespace(**(fields | overrides)) + + +def test_context_normal_planning_and_full_spec_dedup(planners: SimpleNamespace) -> None: + warmer = _warmup.PrimsTSDenseMlaWarmup() + args = dict(max_num_tokens=2048, max_seq_len=4608) + assert warmer._prewarm_context(_geometry(), planners.metadata, **args) == 38 + assert warmer._prewarm_context(_geometry(batch_size=16), planners.metadata, **args) == 0 + assert planners.context.plan.call_count == 38 + for call in planners.context.plan.call_args_list: + assert call.kwargs["batch_size"] == 1 + assert call.kwargs["q_dtype"] == torch.float8_e4m3fn + assert ( + warmer._prewarm_context(_geometry(device=torch.device("cuda:1")), planners.metadata, **args) + == 38 + ) + assert warmer._prewarm_context(_geometry(packed=False), planners.metadata, **args) == 0 + + +def test_failed_plan_is_retried(planners: SimpleNamespace) -> None: + warmer = _warmup.PrimsTSDenseMlaWarmup() + planners.context.plan.side_effect = RuntimeError("compile failure") + args = dict(max_num_tokens=1, max_seq_len=1) + with pytest.raises(RuntimeError, match="compile failure"): + warmer._prewarm_context(_geometry(), planners.metadata, **args) + assert not warmer._completed + planners.context.plan.side_effect = None + assert warmer._prewarm_context(_geometry(), planners.metadata, **args) == 1 + + +def test_decode_scans_batches_not_unobserved_widths(planners: SimpleNamespace) -> None: + warmer = _warmup.PrimsTSDenseMlaWarmup() + state = _decode_state() + original_workspace = state.workspace_buffer + assert warmer._prewarm_decode(state, max_batch_size=16, max_num_tokens=2048) == 3 + resolve_calls = planners.decode._resolve_mla_decode_launch_spec.call_args_list + assert [call.args[1] for call in resolve_calls] == list(range(1, 17)) + assert {call.args[11] for call in resolve_calls} == {6} + assert {call.args[6] for call in resolve_calls} == {4608} + assert {call.args[1] for call in planners.decode.plan.call_args_list} == {14, 15, 16} + assert all( + "workspace_buffer" not in call.kwargs for call in planners.decode.plan.call_args_list + ) + assert state.workspace_buffer is original_workspace + assert warmer._prewarm_decode(state, max_batch_size=16, max_num_tokens=2048) == 0 + assert ( + warmer._prewarm_decode( + _decode_state(store_softmax_stats=True, use_kv_bounds=True), + max_batch_size=16, + max_num_tokens=2048, + ) + == 3 + ) + assert ( + warmer._prewarm_decode( + _decode_state(device=torch.device("cuda:1")), max_batch_size=16, max_num_tokens=2048 + ) + == 3 + ) + + +def test_decode_respects_token_budget_and_failure_retry(planners: SimpleNamespace) -> None: + warmer = _warmup.PrimsTSDenseMlaWarmup() + planners.decode.plan.side_effect = RuntimeError("compile failure") + with pytest.raises(RuntimeError, match="compile failure"): + warmer._prewarm_decode(_decode_state(), max_batch_size=16, max_num_tokens=13) + assert not warmer._completed + assert [ + call.args[1] for call in planners.decode._resolve_mla_decode_launch_spec.call_args_list + ] == [1, 2] + planners.decode.plan.side_effect = None + assert warmer._prewarm_decode(_decode_state(), max_batch_size=16, max_num_tokens=13) == 2 + + +def _model( + planners: SimpleNamespace, *, selected: bool = True, device: str = "cuda:0" +) -> tuple[SimpleNamespace, object, SimpleNamespace]: + library = planners.provider() + wrapper = SimpleNamespace(_plan_state=_decode_state(device=torch.device(device))) + library._mla_context_wrappers = {} + library._mla_decode_wrappers = {(1, 6): wrapper} if selected else {} + backend = planners.backend() + backend._fmha_manager = SimpleNamespace(fmha_libs=[library]) + module = planners.MLA() + module.mha, module.mqa = None, backend + return SimpleNamespace(modules=lambda: [module]), library, wrapper + + +def test_selected_provider_and_capture_gates( + planners: SimpleNamespace, monkeypatch: pytest.MonkeyPatch +) -> None: + warmer = _warmup.PrimsTSDenseMlaWarmup() + args = dict(max_batch_size=16, max_num_tokens=2048, max_seq_len=4608) + model, library, wrapper = _model(planners, selected=False) + assert warmer.run(model, planners.metadata, **args) == (0, 0) + assert not planners.decode.plan.called + model, library, wrapper = _model(planners) + assert warmer.run(model, None, **args) == (0, 0) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + with pytest.raises(RuntimeError, match="during capture"): + warmer.run(model, planners.metadata, **args) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + assert warmer.run(model, planners.metadata, **args) == (0, 3) + assert library._mla_decode_wrappers == {(1, 6): wrapper} + assert warmer.run(model, planners.metadata, **args) == (0, 0) + cpu_model, _, _ = _model(planners, device="cpu") + monkeypatch.setattr( + torch.cuda, "is_current_stream_capturing", Mock(side_effect=AssertionError("CUDA touched")) + ) + assert warmer.run(cpu_model, planners.metadata, **args) == (0, 0) + + +def test_startup_hook_is_after_graphs_and_not_gated_on_graph_enablement() -> None: + source = _ROOT / "tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/warmup.py" + tree = ast.parse(source.read_text()) + scheduled = next( + node + for node in ast.walk(tree) + if isinstance(node, ast.FunctionDef) and node.name == "_warmup_scheduled" + ) + phases = [ + (index, ast.unparse(node)) + for index, node in enumerate(scheduled.body) + if isinstance(node, ast.With) + ] + graph_index = next(index for index, text in phases if "self._graphs.warmup_and_capture" in text) + prewarm_index = next( + index for index, text in phases if "self._prims_ts_dense_mla_warmup.run" in text + ) + assert prewarm_index > graph_index From dd49775236864c55a889aad97ea7d2cc0f41145c Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Thu, 29 Oct 2026 12:51:58 +0000 Subject: [PATCH 14/19] [None][perf] Cache PrimTS MLA workspace sizes per request geometry Avoid repeated workspace policy and sizing work on warmed eager MLA calls. Scope scalar sizes to the prepared KV capacity, invalidate dependent plans when that capacity changes, and preserve workspace allocation and CUDA graph capture checks. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../attention/backends/fmha/prims_ts.py | 63 ++++--- .../_torch/attention/test_prims_ts_fmha.py | 164 +++++++++++++++++- 2 files changed, 197 insertions(+), 30 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py index db8d111089af..7465e0d9e0ec 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py @@ -217,6 +217,8 @@ def __init__(self, attn: "TrtllmAttention") -> None: self._contiguous_context_wrappers: dict[tuple[Any, ...], "BatchPrefillTSWrapper"] = {} self._decode_wrappers: dict[int, "BatchDecodePagedTSWrapper"] = {} self._mla_decode_wrappers: dict[tuple[int, int], "BatchMLADecodePagedTSWrapper"] = {} + self._mla_decode_workspace_sizes: dict[tuple[int, int], int] = {} + self._mla_decode_workspace_capacity: int | None = None self._mla_context_wrappers: dict[ tuple[int, int, int, str, bool], "BatchPrefillTSWrapper" ] = {} @@ -924,8 +926,8 @@ def _get_or_plan_mla_decode_wrapper( ) -> "BatchMLADecodePagedTSWrapper": if batch_size <= 0: raise RuntimeError("PrimTS MLA decode requires a positive batch size.") - # Dense MLA instances fix the cache capacity, mask, dtypes and Helix - # stats/bounds modes. Only batch size and verify width vary per request. + # Within a prepared KV capacity, dense MLA instances fix the mask, + # dtypes and Helix stats/bounds modes. Only batch and verify width vary. key = (batch_size, max_seq_len_q) wrapper = self._mla_decode_wrappers.get(key) if wrapper is not None: @@ -1139,30 +1141,41 @@ def prepare_workspace( window_left = -1 if self.attn.is_mla_enable: - from tensorrt_llm._torch.attention.backends.prims_ts import ( - get_prims_ts_batch_mla_decode_workspace_size, - ) + # Like MLA decode plans, scalar sizes fix the layer's mask, dtypes + # and statistics mode. Warmup can replace the page-table capacity. + if self._mla_decode_workspace_capacity != max_seq_len: + if self._mla_decode_workspace_capacity is not None: + self._mla_decode_wrappers.clear() + self._mla_decode_workspace_capacity = max_seq_len + self._mla_decode_workspace_sizes.clear() + size_key = (batch_size, seq_len_q) + required_bytes = self._mla_decode_workspace_sizes.get(size_key) + if required_bytes is None: + 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, - int(self.attn.kv_lora_rank), - int(self.attn.qk_rope_head_dim), - int(metadata.tokens_per_block), - max_seq_len, - max_seq_len_q=seq_len_q, - q_dtype=kernel_dtype, - kv_dtype=kernel_dtype, - out_dtype=forward_args.output.dtype, - mask_type=mask_type, - store_softmax_stats=forward_args.softmax_stats_tensor is not None, - device=q.device, - ) + 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, + int(self.attn.kv_lora_rank), + int(self.attn.qk_rope_head_dim), + int(metadata.tokens_per_block), + max_seq_len, + max_seq_len_q=seq_len_q, + q_dtype=kernel_dtype, + kv_dtype=kernel_dtype, + out_dtype=forward_args.output.dtype, + mask_type=mask_type, + store_softmax_stats=forward_args.softmax_stats_tensor is not None, + device=q.device, + ) + self._mla_decode_workspace_sizes[size_key] = required_bytes else: from tensorrt_llm._torch.attention.backends.prims_ts import ( get_prims_ts_batch_decode_workspace_size, diff --git a/tests/unittest/_torch/attention/test_prims_ts_fmha.py b/tests/unittest/_torch/attention/test_prims_ts_fmha.py index 62cbbae7e7cc..d7831b62237f 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_fmha.py +++ b/tests/unittest/_torch/attention/test_prims_ts_fmha.py @@ -2458,6 +2458,7 @@ def test_mla_prepare_workspace_preserves_cached_wrappers_with_stable_allocation( attention_window_size=96, ) fmha.prepare_workspace(q, None, None, metadata, forward_args, workspace) + fmha.prepare_workspace(q.clone(), None, None, metadata, forward_args, workspace) stream.synchronize.assert_not_called() assert fmha._mla_decode_wrappers[(2, 1)] is wrapper @@ -2499,15 +2500,168 @@ def test_mla_caller_workspace_grows_across_plan_profiles( fmha.prepare_workspace(q, None, None, metadata, forward_args, workspace) assert workspace.numel() == 32 - dense_forward_args = replace( - forward_args, - attention_mask=PredefinedAttentionMask.FULL, - ) - fmha.prepare_workspace(q, None, None, metadata, dense_forward_args, workspace) + metadata.kv_cache_block_offsets = torch.empty((2, 2, 6), dtype=torch.int32) + fmha.prepare_workspace(q, None, None, metadata, forward_args, workspace) assert workspace.numel() == 64 +@pytest.mark.parametrize("fp8,with_stats", [(False, False), (True, False), (True, True)]) +def test_mla_workspace_size_cache_tracks_request_geometry( + monkeypatch: pytest.MonkeyPatch, fp8: bool, with_stats: bool +) -> None: + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + workspace_size = Mock( + side_effect=lambda batch, heads, latent, rope, page, capacity, **kwargs: ( + 32 * batch * kwargs["max_seq_len_q"] + capacity + ) + ) + monkeypatch.setattr( + prims_ts_package, "get_prims_ts_batch_mla_decode_workspace_size", workspace_size + ) + attn = _Attention(head_dim=576, is_mla=True, num_heads=4) + attn.quant_mode = QuantMode.FP8_KV_CACHE if fp8 else QuantMode(0) + fmha = PrimsTSFmha(attn) + fmha._multi_processor_count = 1 + metadata = SimpleNamespace( + is_cuda_graph=False, + kv_cache_block_offsets=torch.empty((4, 2, 3), dtype=torch.int32), + max_num_requests=4, + max_num_tokens=32, + num_contexts=0, + num_generations=2, + num_ctx_tokens=0, + tokens_per_block=32, + ) + workspace = torch.empty(0, dtype=torch.uint8) + for index, (batch, query_width) in enumerate(((2, 1), (2, 6), (4, 8), (2, 1))): + metadata.num_generations = batch + metadata.max_num_requests = 4 * (index + 1) + metadata.max_num_tokens = 32 * (index + 1) + metadata.is_cuda_graph = bool(index % 2) + q = torch.empty((batch * query_width, attn.num_heads * 576), dtype=torch.bfloat16) + forward_args = AttentionForwardArgs( + output=torch.empty((batch * query_width, attn.num_heads * 512), dtype=q.dtype), + attention_input_type=AttentionInputType.generation_only, + softmax_stats_tensor=torch.empty((batch * query_width, attn.num_heads, 2)) + if with_stats + else None, + ) + fmha.prepare_workspace(q, None, None, metadata, forward_args, workspace) + assert workspace.numel() >= 32 * batch * query_width + 96 + + assert workspace_size.call_count == 3 + assert set(fmha._mla_decode_workspace_sizes) == {(2, 1), (2, 6), (4, 8)} + assert all(isinstance(size, int) for size in fmha._mla_decode_workspace_sizes.values()) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + fmha.prepare_workspace(q, None, None, metadata, forward_args, workspace) + assert workspace_size.call_count == 3 + + # A larger table invalidates capacity-bound plans even when existing scratch + # already fits, so pointer stability does not conceal a changed plan bound. + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + previous_pointer = workspace.data_ptr() + fmha._mla_decode_wrappers[(2, 1)] = Mock() + metadata.kv_cache_block_offsets = torch.empty((4, 2, 6), dtype=torch.int32) + fmha.prepare_workspace(q, None, None, metadata, forward_args, workspace) + assert workspace_size.call_count == 4 + assert workspace_size.call_args.args[5] == 192 + assert set(fmha._mla_decode_workspace_sizes) == {(2, 1)} + assert fmha._mla_decode_wrappers == {} + assert workspace.data_ptr() == previous_pointer + + metadata.kv_cache_block_offsets = torch.empty((4, 2, 64), dtype=torch.int32) + fmha.prepare_workspace(q, None, None, metadata, forward_args, workspace) + assert workspace_size.call_count == 5 + assert workspace_size.call_args.args[5] == 2048 + assert workspace.numel() == 32 * 2 + 2048 + + +def test_mla_workspace_size_cache_preserves_allocation_checks( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + workspace_size = Mock(return_value=48) + monkeypatch.setattr( + prims_ts_package, "get_prims_ts_batch_mla_decode_workspace_size", workspace_size + ) + attn = _Attention(head_dim=576, is_mla=True, num_heads=4) + fmha = PrimsTSFmha(attn) + fmha._multi_processor_count = 1 + metadata = SimpleNamespace( + kv_cache_block_offsets=torch.empty((2, 2, 3), dtype=torch.int32), + num_contexts=0, + num_generations=2, + num_ctx_tokens=0, + tokens_per_block=32, + ) + q = torch.empty((12, attn.num_heads * 576), dtype=torch.bfloat16) + forward_args = AttentionForwardArgs( + output=torch.empty((12, attn.num_heads * 512), dtype=q.dtype), + attention_input_type=AttentionInputType.generation_only, + ) + first_workspace = torch.empty(0, dtype=torch.uint8) + fmha.prepare_workspace(q, None, None, metadata, forward_args, first_workspace) + fmha._mla_decode_wrappers[(2, 6)] = Mock() + context_wrapper = Mock() + fmha._context_wrappers[2] = context_wrapper + + second_workspace = torch.empty(48, dtype=torch.uint8) + assert second_workspace.data_ptr() != first_workspace.data_ptr() + fmha.prepare_workspace(q, None, None, metadata, forward_args, second_workspace) + workspace_size.assert_called_once() + assert fmha._mla_decode_wrappers == {} + assert fmha._context_wrappers[2] is context_wrapper + + too_small = torch.empty(16, dtype=torch.uint8) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + with pytest.raises(RuntimeError, match="must be sized before CUDA graph capture"): + fmha.prepare_workspace(q, None, None, metadata, forward_args, too_small) + workspace_size.assert_called_once() + assert too_small.numel() == 16 + + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + fmha.prepare_workspace(q, None, None, metadata, forward_args, too_small) + workspace_size.assert_called_once() + assert too_small.numel() == 48 + + +@pytest.mark.parametrize("fp8", [False, True], ids=["bf16", "fp8"]) +def test_mla_context_workspace_tracks_each_chunk_without_decode_size_cache( + monkeypatch: pytest.MonkeyPatch, fp8: bool +) -> None: + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + workspace_size = Mock(side_effect=AssertionError("context must not use decode sizing")) + monkeypatch.setattr( + prims_ts_package, "get_prims_ts_batch_mla_decode_workspace_size", workspace_size + ) + attn = _Attention(head_dim=192, is_mla=True, num_heads=4) + attn.quant_mode = QuantMode.FP8_KV_CACHE if fp8 else QuantMode(0) + fmha = PrimsTSFmha(attn) + fmha._multi_processor_count = 1 + metadata = SimpleNamespace( + kv_cache_block_offsets=torch.empty((2, 2, 8), dtype=torch.int32), + num_contexts=2, + num_generations=0, + num_ctx_tokens=8, + tokens_per_block=32, + ) + workspace = torch.empty(0, dtype=torch.uint8) + for q_tokens, kv_tokens in ((8, 8), (16, 32)): + metadata.num_ctx_tokens = q_tokens + q = torch.empty((q_tokens, attn.num_heads * 192), dtype=torch.bfloat16) + k = torch.empty((kv_tokens, attn.num_heads * 192), dtype=q.dtype) + forward_args = AttentionForwardArgs( + output=torch.empty((q_tokens, attn.num_heads * 128), dtype=q.dtype), + attention_input_type=AttentionInputType.context_only, + ) + expected_bytes = fmha._mla_context_workspace_layout(q, k, 2)[1] + fmha.prepare_workspace(q, k, None, metadata, forward_args, workspace) + assert workspace.numel() == expected_bytes + workspace_size.assert_not_called() + assert fmha._mla_decode_workspace_sizes == {} + + def test_decode_prepare_workspace_reserves_tail_after_compact_preprocessing( monkeypatch: pytest.MonkeyPatch, ) -> None: From 0c80a0cfc64417813965f808502cc51c41e2e639 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Sun, 1 Nov 2026 10:54:14 +0000 Subject: [PATCH 15/19] [None][perf] Vendor validated PrimTS MLA scheduling optimizations Import the separately reviewed FlashInfer single-split, packed-query pruning, LDRED and runtime-KV domain optimizations from the published temporary source branch. Preserve main compatibility adaptations and the functionality/performance split. The materialized production tree matches the Rubin-validated source, except for the independently audited block-sparse metadata compatibility fix from main. Keep the benchmark results scoped to their frozen pre-rebase runtime. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- 3rdparty/vendor_sources.lock.yaml | 4 +- .../kernels/fmha_context/fmha_kernel.py | 80 ++++++++++++++++++- .../kernels/fmha_context/fmha_resources.py | 45 ++++++++++- .../kernels/fmha_context/fmha_tasks.py | 6 +- .../attention/backends/prims_ts/mla_decode.py | 22 +++++ 5 files changed, 149 insertions(+), 8 deletions(-) diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 995b256547fd..1d489aa27c0e 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-ts-mla-prefill--flashinfer - commit: df3434b1dc23df96527e202aa43ae233ed8641ee + commit: 76e8229748606d0e51f27f84fe19335cbe605ae9 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:27cc330c1eff406250b2ea7212d7a3200773a4c2d9eaacd57e57033344d39f3c - digest: sha256-tree-v1:0ea9593edb2c167fcd478c30d6ac423d31e6b2eb51c8f50e25a463b00d045158 + digest: sha256-tree-v1:ba233486a857302bc90ca13de3448c2f78c437a83afb253e118bb36403133306 openengine-proto: url: https://github.com/ai-dynamo/openengine.git tag: v0.1.0 diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py index f0e4ce0cd3d2..2c4a64d77a92 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py @@ -246,7 +246,7 @@ def _init_task_with_domain( """Initialize a captured Task with its static validation loop domain.""" schedule = kwargs.get("schedule") if schedule is None: - raise ValueError("Causal domain tasks require a captured schedule.") + raise ValueError("Domain tasks require a captured schedule.") if domain is not None: schedule.loop_end = domain task_init(task, **kwargs) @@ -325,6 +325,40 @@ def _init_causal_domain_state( task._packed_window = packed_window +class PackedDenseDomainTask(Task): + """Request-local dense KV extent with one masked tile for empty rows.""" + + def __init__( + self, + num_kv_tiles: int | Int32, + tile_size_kv: int, + batch_idx: int, + cum_seqlen_k: cute.Tensor, + offset: int = 0, + **kwargs: Any, + ) -> None: + """Retain the planned domain for validation and bound runtime KV work.""" + static_domain = ( + max(num_kv_tiles - offset, 0) if isinstance(num_kv_tiles, int) else None + ) + _init_task_with_domain(self, kwargs, static_domain) + self._num_kv_tiles = num_kv_tiles + self._tile_size_kv = tile_size_kv + self._batch_idx = batch_idx + self._cum_seqlen_k = cum_seqlen_k + self._offset = offset + + def get_domain(self, tile_coord: cute.Coord) -> int | Int32: + """Use the same live KV bound for each task's N or N-1 domain.""" + batch_coord = Int32(tile_coord[self._batch_idx]) + k_begin = Int32(self._cum_seqlen_k[batch_coord]) + seqlen_k = Int32(self._cum_seqlen_k[batch_coord + Int32(1)]) - k_begin + runtime_tiles = cute.ceil_div(seqlen_k, self._tile_size_kv) + # Empty KV still traverses HEAD/TAIL to write zero O and (-inf, 0). + domain = _domain_min(self._num_kv_tiles, _domain_max(1, runtime_tiles)) + return domain - self._offset + + class CausalDomainTask(Task): """Task with causal and optional left-window masking domain. @@ -671,6 +705,30 @@ class FmhaDomainPolicy: softmax1_domain_kwargs: DomainKwargs +def _uses_sm107_dense_mla_q_pruning(cfg: FmhaConfig) -> bool: + """Prune unused packed-Q pairs in native SM107 FP8 expanded MLA.""" + if not ( + not cfg.is_causal + and not cfg.has_variable_window + and cfg.window_size_left == 0 + and not cfg.head_paired + and not cfg.use_paged_kv + and cfg.has_varlen + and not cfg.has_uniform_varlen + and cfg.num_qkv_instances == 2 + and not cfg.two_cta_umma + and cfg.logical_head_dim_qk == 192 + and cfg.pv_mma_tiler[1] == 128 + and cfg.q_dtype == cfg.k_dtype == cfg.v_dtype == cutlass.Float8E4M3FN + and cfg.o_dtype == cutlass.BFloat16 + ): + return False + from cutlass.cutlass_dsl import BaseDSL + + arch = BaseDSL._get_dsl().get_arch_enum() + return (arch.major, arch.minor) == (10, 7) + + def build_context_task_manager( *, cfg: FmhaConfig, @@ -1125,7 +1183,7 @@ def make_page_offsets_pipeline_cfg(num_stages: int) -> PipelineConfig: "name": "work_queue", } if ( - cfg.is_causal + (cfg.is_causal or _uses_sm107_dense_mla_q_pruning(cfg)) and cfg.has_varlen and not cfg.has_uniform_varlen and cum_seqlen_q is not None @@ -2546,6 +2604,24 @@ def _select_fmha_domain_policy( softmax0_domain_kwargs=softmax0_domain, softmax1_domain_kwargs=softmax1_domain, ) + if _uses_sm107_dense_mla_q_pruning(cfg): + if cum_seqlen_k is None: + raise ValueError("Packed dense MLA domains require runtime KV offsets") + dense_kwargs: DomainKwargs = { + "task_class": PackedDenseDomainTask, + "num_kv_tiles": num_kv_tiles, + "tile_size_kv": cfg.kv_tile_n, + "batch_idx": batch_idx, + "cum_seqlen_k": cum_seqlen_k, + } + domain_n_kwargs = {**dense_kwargs, "offset": 0} + domain_n_minus_1_kwargs = {**dense_kwargs, "offset": 1} + return FmhaDomainPolicy( + domain_n_kwargs=domain_n_kwargs, + domain_n_minus_1_kwargs=domain_n_minus_1_kwargs, + softmax0_domain_kwargs=domain_n_kwargs, + softmax1_domain_kwargs=domain_n_kwargs, + ) # Dense query-paired fallback: no head-pairing or causal mask, so both # peers traverse the full static K/V domain. domain_n_kwargs = {"domain": num_kv_tiles} diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py index 9ab920d4ba70..c4b11379ab30 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py @@ -3210,14 +3210,17 @@ def _load_s_chunks_and_reduce_row_max( *, causal_loop_mask: cutlass.Constexpr[bool] = False, q_offset: Int32 = 0, + dense_k_mask: cutlass.Constexpr[bool] = False, + seqlen_k: Int32 = 0, + section: cutlass.Constexpr[FmhaStage] = FmhaStage.Loop, ) -> tuple[SoftmaxScalar, SoftmaxScalar]: """LDTM.STAT: fuse S load and per-chunk row_max via tcgen05.ld.red.max. Instead of loading each S chunk with ``tcgen05.ld`` and then folding a per-lane max in registers, use ``tcgen05.ld.red.max`` so the - reduction happens as part of the TMEM load. For causal loop masking, - only tiles crossing the runtime right bound need a software reduction - of masked scores; wholly visible tiles retain the hardware maximum. + reduction happens as part of the TMEM load. For causal or packed dense + masking, only tiles crossing the runtime right bound need a software + reduction of masked scores; full tiles retain the hardware maximum. """ tmem_s_addr = self.tmem_s_addr_cached + self._stage_col_offset(stage_info) tmem_x = self.cfg.tmem_x_load_s @@ -3275,6 +3278,34 @@ def _load_s_chunks_and_reduce_row_max( masked_row_max, s_data[chunk_idx].reduce("max") ) row_max = masked_row_max + elif cutlass.const_expr(dense_k_mask): + if cutlass.const_expr(section == FmhaStage.Loop): + kv_tile_idx = stage_info.loop_offset + else: + kv_tile_idx = stage_info.loop_end + kv_base = kv_tile_idx * self.cfg.kv_tile_n + kv_end = kv_base + Int32(self.cfg.kv_tile_n) + if seqlen_k < kv_end: + # Invalid lanes can contain another request's scores or TMA + # padding. Neither may contribute to the maximum or softmax. + masked_row_max = old_row_max + for chunk_idx in cutlass.range_constexpr(num_chunks): + chunk_base = kv_base + Int32(chunk_idx * tmem_x) + valid_in_chunk = cute.math.min( + cute.math.max(seqlen_k - chunk_base, Int32(0)), + Int32(tmem_x), + ) + mask = cutlass.vector.create_mask([tmem_x], [valid_in_chunk]) + neg_inf = cutlass.vector.full_like( + s_data[chunk_idx], Float32(-Float32.inf) + ) + s_data[chunk_idx] = cutlass.vector.where( + mask, s_data[chunk_idx], neg_inf + ) + masked_row_max = cute.math.max( + masked_row_max, s_data[chunk_idx].reduce("max") + ) + row_max = masked_row_max _tmem_sp_sdata[id(self)] = s_data row_max_safe = row_max if row_max == -Float32.inf: @@ -3926,6 +3957,14 @@ def packed_dense_k_masked_row_max( section: cutlass.Constexpr[FmhaStage], ) -> tuple[SoftmaxScalar, SoftmaxScalar]: """Mask scores beyond one packed request's K/V right edge.""" + if cutlass.const_expr(self.cfg.uses_ldtm_stat): + return self._load_s_chunks_and_reduce_row_max( + stage_info, + row_max, + dense_k_mask=True, + seqlen_k=seqlen_k, + section=section, + ) tmem_x = self.cfg.tmem_x_load_s num_chunks = self.cfg.qk_mma_tiler[1] // tmem_x s_data = self._load_s_chunks(stage_info) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_tasks.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_tasks.py index 2e2d8581378e..bc2c3a4d942d 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_tasks.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_tasks.py @@ -63,7 +63,11 @@ @dataclass(kw_only=True) class PackedContextWorkQueue(WorkQueue): - """Persistent queue that skips Q tiles outside a live packed request.""" + """Persistent queue that skips Q tiles outside a live packed request. + + The predicate depends only on Q validity, not mask type or K/V length. + Valid Q rows with no K/V still execute the output/statistics epilogue. + """ cfg: cutlass.Constexpr[FmhaConfig] = field(init=False, default=None) cum_seqlen_q: Any = field(init=False, default=None) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py b/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py index f0a6e9ddf697..5b1cb61c1317 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py @@ -1005,6 +1005,28 @@ def _resolve_mla_decode_launch_spec( if not decision.implementation_ready: raise NotImplementedError(decision.reason) split_kv = two_cta_split_kv + # These short, bounded SM107 plans avoid split scratch and a + # separate reduction. The rule describes local geometry, including + # packed queries, rather than a context-parallel rank count. + if ( + policy_source == "auto" + and batch_size == 1 + and num_heads == 96 + and seq_len_q == 8 + and kv_lora_rank == 512 + and qk_rope_head_dim == 64 + and page_size == 128 + and max_kv_len == 1024 + and q_dtype_key == kv_dtype_key == "float8_e4m3fn" + and output_dtype_key == "bfloat16" + and mask_type == "causal" + and store_softmax_stats + and use_kv_bounds + and torch.cuda.get_device_capability(device_index) == (10, 7) + and torch.cuda.get_device_properties(device_index).multi_processor_count + == 212 + ): + split_kv = 1 work_clusters = batch_size * launch_shape.num_tiles * max(split_kv, 1) # Dynamic cluster stealing only helps once logical work exceeds a # resident wave. Within one wave every cluster already launches, From 6f075ff46486ebe3dc6a427317fa6247f906cbf8 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Sun, 1 Nov 2026 14:30:28 +0000 Subject: [PATCH 16/19] [None][fix] Preserve the prepared sparse MLA launch ABI Forward the trailing disabled KV-bounds operand through the downstream prepared launcher and regenerate the vendor compatibility patch with the vendor toolkit. Cover host argument parity for both cache layouts and workspace modes, and exercise prepared execution after persistent-kernel reload. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../vendor_patches/flashinfer-prims-ts.patch | 27 ++-- 3rdparty/vendor_sources.lock.yaml | 4 +- .../attention/backends/prims_ts/mla_decode.py | 1 + .../test_prims_ts_sparse_mla_prepared.py | 150 ++++++++++++++++++ .../test_prims_ts_sparse_mla_fmha.py | 17 +- 5 files changed, 183 insertions(+), 16 deletions(-) create mode 100644 tests/unittest/_torch/attention/hw_agnostic/test_prims_ts_sparse_mla_prepared.py diff --git a/3rdparty/vendor_patches/flashinfer-prims-ts.patch b/3rdparty/vendor_patches/flashinfer-prims-ts.patch index d749481966f5..6a6d7894f5ef 100644 --- a/3rdparty/vendor_patches/flashinfer-prims-ts.patch +++ b/3rdparty/vendor_patches/flashinfer-prims-ts.patch @@ -327,7 +327,7 @@ index 8c331d05a5b55717df3ba4dd1ff55b6ca44c831e..9c0a5e7243ee2d4ad33eed1afd45f9b8 q: torch.Tensor, paged_kv_cache: PagedKVCache, diff --git a/kernels/fmha_context/fmha_resources.py b/kernels/fmha_context/fmha_resources.py -index f4e67b642efb6d8f8ad5e5519c1679c47e776f54..9ab920d4ba7017db8d26f884ae2e420b67ddb920 100644 +index d29e7145f1da027ec86be7247a64a0a7d730c102..c4b11379ab3076644f4e7bdadf9948049bb23556 100644 --- a/kernels/fmha_context/fmha_resources.py +++ b/kernels/fmha_context/fmha_resources.py @@ -1,3 +1,4 @@ @@ -385,7 +385,7 @@ index f4e67b642efb6d8f8ad5e5519c1679c47e776f54..9ab920d4ba7017db8d26f884ae2e420b tile_idx = _resolve_kv_tile_idx_context( stage_info, kv_tile_start, tile_offset=tile_offset ) -@@ -5903,3 +5892,14 @@ class GmemOResource(MemoryResource): +@@ -5942,3 +5931,14 @@ class GmemOResource(MemoryResource): prims.cp_async_bulk_commit_group() if cutlass.const_expr(self.cfg.gmem_o_store_wait_after_write): prims.cp_async_bulk_wait_group(0, read=True) @@ -419,7 +419,7 @@ index b22f0f413b7312bab9a5f8e296255fe9aeabe901..a4a87200ca5d3ace5cf899caa9672a1a from .helpers_common import ( _TASK_CACHE_TMEM_BASE_OFFSET, diff --git a/mla_decode.py b/mla_decode.py -index 7ea49a00d3ed6a236bd9da2854b45a985e23d4a5..f0a6e9ddf697fc80748b859fd2818f683ab31eed 100644 +index 7774f881518b11f8e3fc0626afbcd757c77ca6cc..3c3a5d2197cfbc33fcedc7d6c2fb21ebaac2a74d 100644 --- a/mla_decode.py +++ b/mla_decode.py @@ -28,11 +28,7 @@ from ._hardware import get_max_active_clusters @@ -435,7 +435,7 @@ index 7ea49a00d3ed6a236bd9da2854b45a985e23d4a5..f0a6e9ddf697fc80748b859fd2818f68 from .decode import ( _WorkspaceSection, -@@ -1133,11 +1129,35 @@ def _make_mla_decode_compile_spec( +@@ -1155,11 +1151,35 @@ def _make_mla_decode_compile_spec( ) @@ -476,7 +476,7 @@ index 7ea49a00d3ed6a236bd9da2854b45a985e23d4a5..f0a6e9ddf697fc80748b859fd2818f68 import cutlass import cutlass.cute as cute -@@ -1776,7 +1796,7 @@ class BatchMLADecodePagedTSWrapper: +@@ -1798,7 +1818,7 @@ class BatchMLADecodePagedTSWrapper: use_kv_bounds=use_kv_bounds, ) @@ -485,7 +485,7 @@ index 7ea49a00d3ed6a236bd9da2854b45a985e23d4a5..f0a6e9ddf697fc80748b859fd2818f68 def run( self, query: torch.Tensor, -@@ -1905,7 +1925,7 @@ class BatchMLADecodePagedTSWrapper: +@@ -1927,7 +1947,7 @@ class BatchMLADecodePagedTSWrapper: ) @@ -494,7 +494,7 @@ index 7ea49a00d3ed6a236bd9da2854b45a985e23d4a5..f0a6e9ddf697fc80748b859fd2818f68 def batch_mla_decode_with_paged_kv_cache( query: torch.Tensor, kv_cache: torch.Tensor, -@@ -2284,6 +2304,7 @@ def _prepare_sparse_output_quant( +@@ -2306,6 +2326,7 @@ def _prepare_sparse_output_quant( ) @@ -502,7 +502,7 @@ index 7ea49a00d3ed6a236bd9da2854b45a985e23d4a5..f0a6e9ddf697fc80748b859fd2818f68 def _resolve_sparse_mla_plan( device, batch_size, -@@ -2304,17 +2325,17 @@ def _resolve_sparse_mla_plan( +@@ -2326,17 +2347,17 @@ def _resolve_sparse_mla_plan( ): """Resolve sparse geometry and scratch sizing without allocation/compilation.""" import cutlass @@ -524,7 +524,7 @@ index 7ea49a00d3ed6a236bd9da2854b45a985e23d4a5..f0a6e9ddf697fc80748b859fd2818f68 MlaDecodeTs, ) -@@ -2449,6 +2470,9 @@ def _resolve_sparse_mla_plan( +@@ -2471,6 +2492,9 @@ def _resolve_sparse_mla_plan( return device, profile, spec, sections, byte_end, capacity, output_quant @@ -534,7 +534,7 @@ index 7ea49a00d3ed6a236bd9da2854b45a985e23d4a5..f0a6e9ddf697fc80748b859fd2818f68 class BatchSparseMLADecodePagedTSWrapper: """Plan native BF16/E4M3 attention, then bind one or two source index lists. -@@ -2482,6 +2506,11 @@ class BatchSparseMLADecodePagedTSWrapper: +@@ -2504,6 +2528,11 @@ class BatchSparseMLADecodePagedTSWrapper: self._workspace_buffer = workspace_buffer self._state = None self._scalar_cache = {} @@ -546,7 +546,7 @@ index 7ea49a00d3ed6a236bd9da2854b45a985e23d4a5..f0a6e9ddf697fc80748b859fd2818f68 self.workspace_size_bytes = 0 def plan( -@@ -2571,6 +2600,7 @@ class BatchSparseMLADecodePagedTSWrapper: +@@ -2593,6 +2622,7 @@ class BatchSparseMLADecodePagedTSWrapper: kernel_signature=_mla_kernel_compile_signature(static_kernel), ) ) @@ -554,7 +554,7 @@ index 7ea49a00d3ed6a236bd9da2854b45a985e23d4a5..f0a6e9ddf697fc80748b859fd2818f68 self._state = dict( device=device, batch=batch_size, -@@ -2793,6 +2823,38 @@ class BatchSparseMLADecodePagedTSWrapper: +@@ -2815,6 +2845,38 @@ class BatchSparseMLADecodePagedTSWrapper: state = self._state if state is None: raise RuntimeError("plan() must be called before run()") @@ -593,7 +593,7 @@ index 7ea49a00d3ed6a236bd9da2854b45a985e23d4a5..f0a6e9ddf697fc80748b859fd2818f68 _validate_query( query, packed_query=state["packed"], -@@ -2966,10 +3028,125 @@ class BatchSparseMLADecodePagedTSWrapper: +@@ -2988,10 +3050,126 @@ class BatchSparseMLADecodePagedTSWrapper: if static_scales else state["compiled"], ) @@ -713,6 +713,7 @@ index 7ea49a00d3ed6a236bd9da2854b45a985e23d4a5..f0a6e9ddf697fc80748b859fd2818f68 + bmm2, + None, + None, ++ None, + ) + return out + diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 1d489aa27c0e..1c532591164a 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -21,8 +21,8 @@ vendors: include: - '**/*.py' patch: 3rdparty/vendor_patches/flashinfer-prims-ts.patch - patch_digest: sha256:27cc330c1eff406250b2ea7212d7a3200773a4c2d9eaacd57e57033344d39f3c - digest: sha256-tree-v1:ba233486a857302bc90ca13de3448c2f78c437a83afb253e118bb36403133306 + patch_digest: sha256:011e0dab0ea60efff540c3ce80a96970f1be9c85f6f5b67522d157c31fa710fa + digest: sha256-tree-v1:5e79472af6592b6030d60e51d8498cb0f45fd55c5b974625cafa966ad2e58a39 openengine-proto: url: https://github.com/ai-dynamo/openengine.git tag: v0.1.0 diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py b/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py index 5b1cb61c1317..3c3a5d2197cf 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py @@ -3166,6 +3166,7 @@ def _run_prepared(self, prepared, query, indices, lengths, extra_indices, extra_ bmm2, None, None, + None, ) return out diff --git a/tests/unittest/_torch/attention/hw_agnostic/test_prims_ts_sparse_mla_prepared.py b/tests/unittest/_torch/attention/hw_agnostic/test_prims_ts_sparse_mla_prepared.py new file mode 100644 index 000000000000..3871361028a8 --- /dev/null +++ b/tests/unittest/_torch/attention/hw_agnostic/test_prims_ts_sparse_mla_prepared.py @@ -0,0 +1,150 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Host-only checks for the prepared sparse MLA launch contract.""" + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.prims_ts import mla_decode + +pytestmark = pytest.mark.cpu_only + + +def _assert_same_argument(actual: object, expected: object) -> None: + if isinstance(expected, torch.Tensor): + assert isinstance(actual, torch.Tensor) + assert actual.data_ptr() == expected.data_ptr() + assert actual.shape == expected.shape + assert actual.stride() == expected.stride() + assert actual.dtype == expected.dtype + elif isinstance(expected, tuple): + assert isinstance(actual, tuple) + assert len(actual) == len(expected) + for actual_item, expected_item in zip(actual, expected, strict=True): + _assert_same_argument(actual_item, expected_item) + else: + assert actual == expected + + +@pytest.mark.parametrize("extra_pool", [False, True]) +@pytest.mark.parametrize("has_workspace", [False, True]) +def test_prepared_sparse_mla_matches_full_launch_abi(extra_pool: bool, has_workspace: bool) -> None: + rows, heads, topk = 2, 64, 4 + extra_topk = 3 if extra_pool else 0 + query = torch.zeros(rows, 1, heads, 512, dtype=torch.bfloat16) + out = torch.empty_like(query) + indices = torch.zeros(rows, topk, dtype=torch.int32) + lengths = torch.full((rows,), topk, dtype=torch.int32) + extra_indices = torch.zeros(rows, max(extra_topk, 1), dtype=torch.int32) + extra_lengths = torch.full((rows,), extra_topk, dtype=torch.int32) + primary = torch.ones(8, 1, 512, dtype=torch.bfloat16) + secondary = torch.full_like(primary, 2) if extra_pool else primary + core = torch.empty(1024 if has_workspace else 0, dtype=torch.int8) + lse = torch.empty(rows, 1, heads, dtype=torch.float32) + scales_and_sinks = (torch.tensor([0.5]), torch.tensor([1.0]), torch.zeros(heads)) + split_kv = 2 if has_workspace else 1 + calls = [] + + def compiled( + q_latent: torch.Tensor, + q_rope: torch.Tensor, + c_latent: torch.Tensor, + c_rope: torch.Tensor, + page_offsets: torch.Tensor, + output: torch.Tensor, + softmax_lse: torch.Tensor, + workspace: torch.Tensor | None, + splits: int, + cache_seqs: torch.Tensor, + cu_seqlens_q: torch.Tensor | None, + sparse_inputs: tuple, + softmax_scale: float, + output_scale: float, + output_quant_args: tuple | None, + softmax_stats: torch.Tensor | None, + kv_bounds: torch.Tensor | None, + ) -> None: + # The exported TVM-FFI callable has no Python defaults for these slots. + calls.append( + ( + q_latent, + q_rope, + c_latent, + c_rope, + page_offsets, + output, + softmax_lse, + workspace, + splits, + cache_seqs, + cu_seqlens_q, + sparse_inputs, + softmax_scale, + output_scale, + output_quant_args, + softmax_stats, + kv_bounds, + ) + ) + + wrapper = object.__new__(mla_decode.BatchSparseMLADecodePagedTSWrapper) + wrapper._state = dict( + batch=rows, + heads=heads, + device=query.device, + dtype=query.dtype, + splits=split_kv, + buffers={"core": core, "public_lse": lse}, + compiled_static=compiled, + ) + prepared = dict( + ks=topk, + kc=extra_topk, + dummy_indices=extra_indices, + default_cl=extra_lengths, + c_latent=primary.permute(1, 2, 0), + c_rope=secondary.permute(1, 2, 0), + scalars=scales_and_sinks, + bmm=(0.5, 1.0), + ) + assert ( + wrapper._run_prepared( + prepared, + query, + indices, + lengths, + extra_indices if extra_pool else None, + extra_lengths if extra_pool else None, + out, + ) + is out + ) + assert len(calls) == 1 + assert calls[0][-3:] == (None, None, None) + + runtime = mla_decode._MLARuntime( + query=query, + normalized_cache=primary, + out=out, + num_physical_pages=primary.shape[0], + bmm1_scale=0.5, + bmm2_scale=1.0, + extra_cache=secondary, + sparse_inputs=(indices, extra_indices, lengths, extra_lengths, *scales_and_sinks), + ) + assert ( + mla_decode._launch_mla_decode( + runtime, + block_tables=indices, + seq_lens=lengths, + qo_indptr=None, + packed_query=False, + kv_lora_rank=512, + split_kv=split_kv, + workspace=mla_decode._MLAWorkspaceViews(core if has_workspace else None, lse), + compiled=compiled, + ) + is out + ) + assert len(calls) == 2 + _assert_same_argument(calls[0], calls[1]) diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py index 5beda56866d2..60e2e1b026e1 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py @@ -296,7 +296,7 @@ def test_reloaded_kernel_accepts_the_same_launch(): production restarts) reload the exported TVM-FFI kernels, whose positional entry has no Python defaults. Drop the in-process kernel memo after the first run so the second plan takes the reload path, and require identical - output. + output from both the full and prepared launches. """ _skip_unless_toolchain() from tensorrt_llm._torch.attention.backends.prims_ts import mla_decode @@ -325,7 +325,22 @@ def run_fresh_plan() -> torch.Tensor: output_scale=1.0, out=out, ) + prepared_out = torch.empty_like(out) + wrapper.run( + query, + pool, + indices, + lengths, + extra_pool, + extra_indices=extra_indices, + extra_lengths=extra_lengths, + softmax_scale=LATENT**-0.5, + output_scale=1.0, + out=prepared_out, + validate=False, + ) torch.cuda.synchronize() + torch.testing.assert_close(prepared_out, out, rtol=0, atol=0) return out first = run_fresh_plan() From 0299c41c78b1e2444d0282f04318c756a6b4f0de Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Sun, 1 Nov 2026 14:39:35 +0000 Subject: [PATCH 17/19] [None][test] Initialize PrimTS warmup collaborators in phase fixtures Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../compilation/test_phase_selective_forward.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/tests/unittest/_torch/compilation/test_phase_selective_forward.py b/tests/unittest/_torch/compilation/test_phase_selective_forward.py index b28f11134e29..bcf30fce56b5 100644 --- a/tests/unittest/_torch/compilation/test_phase_selective_forward.py +++ b/tests/unittest/_torch/compilation/test_phase_selective_forward.py @@ -237,14 +237,20 @@ def test_decoder_runner_warmup_bypasses_only_auxiliary_phases(monkeypatch, enabl ) graphs.preallocate_padding_dummies = Mock() warmup = object.__new__(DecoderWarmup) + warmup.model = model warmup.mapping = SimpleNamespace(cp_size=1, has_cp_helix=lambda: False) warmup._inputs = SimpleNamespace(guided_decoder=None) warmup._metrics = {} - warmup._config = SimpleNamespace(max_beam_width=1) + warmup._config = SimpleNamespace( + max_beam_width=1, max_batch_size=8, max_num_tokens=64, max_seq_len=128 + ) warmup._warmup_timer = _WarmupTimer(rank=0) warmup._agreement = SimpleNamespace(agree_flag=lambda value: value, agree_shapes=lambda s: s) warmup._graphs = graphs - warmup._metadata = SimpleNamespace(_warmup_dspark_ragged_compressor_metadata=Mock()) + warmup._metadata = SimpleNamespace( + attn_metadata=object(), _warmup_dspark_ragged_compressor_metadata=Mock() + ) + warmup._prims_ts_dense_mla_warmup = SimpleNamespace(run=Mock()) warmup._context_warmups_supported = True warmup._torch_compile_bypass = runner._bypass_torch_compile warmup._get_full_general_warmup_requests = Mock(return_value=[(4, 0)]) @@ -303,6 +309,13 @@ def forward(*args): ("_general_warmup", not enabled), ] assert not model._bypass_active + warmup._prims_ts_dense_mla_warmup.run.assert_called_once_with( + model, + warmup._metadata.attn_metadata, + max_batch_size=8, + max_num_tokens=64, + max_seq_len=128, + ) @skip_pre_hopper From 06b382b78c3a06c982e6ab0bbb3c5b429ed66c3f Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Sun, 1 Nov 2026 14:47:52 +0000 Subject: [PATCH 18/19] [None][fix] Preserve original softmax scales for PrimTS partition merging Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- 3rdparty/vendor_sources.lock.yaml | 4 +- .../attention/backends/fmha/prims_ts.py | 1 + .../attention/backends/prims_ts/context.py | 67 ++++++++++++++++++- .../kernels/fmha_context/fmha_kernel.py | 9 +++ .../kernels/fmha_context/fmha_resources.py | 5 +- .../kernels/fmha_decode/fmha_decode_kernel.py | 14 ++++ .../fmha_decode_resources/helpers_softmax.py | 16 +++-- .../fmha_decode_resources/tmem_corr.py | 8 ++- .../_torch/attention/test_prims_ts_fmha.py | 29 ++++++-- 9 files changed, 139 insertions(+), 14 deletions(-) diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 1c532591164a..51aed6027efc 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-ts-mla-prefill--flashinfer - commit: 76e8229748606d0e51f27f84fe19335cbe605ae9 + commit: eef263b632a86a1003f8c3824c29a8bc8d3957ff 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:011e0dab0ea60efff540c3ce80a96970f1be9c85f6f5b67522d157c31fa710fa - digest: sha256-tree-v1:5e79472af6592b6030d60e51d8498cb0f45fd55c5b974625cafa966ad2e58a39 + digest: sha256-tree-v1:60780d99f92f1b186c4d3dc2e0600480a3d93ddd6fde94cdd61d3dc7adcb4baf 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 7465e0d9e0ec..6a4b9d6c2716 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py @@ -1388,6 +1388,7 @@ def run_mla_context(self, params: FmhaParams) -> None: kv_indptr, out=params.output, softmax_stats=stats, + softmax_scale=buffers["bmm1_scale"][:1] if fp8 and stats is not None else None, scale_softmax_log2=buffers["bmm1_scale"][1:] if fp8 else None, output_scale=buffers["bmm2_scale"] if fp8 else None, validate=False, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/context.py b/tensorrt_llm/_torch/attention/backends/prims_ts/context.py index b41e7b47cca1..41048a80b43e 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/context.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/context.py @@ -157,6 +157,7 @@ class _ContextPlanState: variable_window_cta_starts: torch.Tensor compiled: Callable[..., None] policy: tuple[tuple[str, object], ...] + softmax_scale: Optional[torch.Tensor] = None @dataclass(frozen=True) @@ -194,6 +195,7 @@ class _PagedContextPlanState: output_scale: torch.Tensor compiled: Callable[..., None] policy: tuple[tuple[str, object], ...] + softmax_scale: Optional[torch.Tensor] = None _ContextScheduler = Literal[ @@ -1847,6 +1849,7 @@ def tensor_adapter( static_packed: cutlass.Constexpr[bool], static_max_seq_len_q: cutlass.Constexpr[int], static_max_seq_len_k: cutlass.Constexpr[int], + stats_scale: cute.Tensor | None, ) -> None: """Adapt torch TVM-FFI tensors to the FmhaTs host entry point.""" @@ -1868,6 +1871,7 @@ def tensor_adapter( variable_window_token_ends=variable_window_token_ends, variable_window_cta_starts=variable_window_cta_starts, softmax_stats=softmax_stats, + stats_scale=stats_scale, ) else: fmha( @@ -1883,6 +1887,7 @@ def tensor_adapter( variable_window_token_ends=variable_window_token_ends, variable_window_cta_starts=variable_window_cta_starts, softmax_stats=softmax_stats, + stats_scale=stats_scale, ) def fake_compact(dtype, shape, assumed_align): @@ -1966,6 +1971,7 @@ def compile_kernel(): packed, max_seq_len_q, max_seq_len_k, + scale_fake if compile_spec.store_softmax_stats else None, options=_COMPILE_OPTIONS, ) @@ -2079,6 +2085,7 @@ def tensor_adapter( static_max_active_clusters: cutlass.Constexpr[int], static_max_seq_len_q: cutlass.Constexpr[int], static_max_kv_len: cutlass.Constexpr[int], + stats_scale: cute.Tensor | None, ) -> None: fmha( q, @@ -2095,6 +2102,7 @@ def tensor_adapter( block_tables=block_tables, seq_lens_kv=seq_lens_kv, softmax_stats=softmax_stats, + stats_scale=stats_scale, ) def fake_compact(dtype, shape, assumed_align): @@ -2155,6 +2163,7 @@ def compile_kernel(): max_active_clusters, max_seq_len_q, max_kv_len, + scale_fake if compile_spec.store_softmax_stats else None, options=_COMPILE_OPTIONS, ) @@ -2694,6 +2703,11 @@ def plan( scale_tensor = torch.tensor( [scale_softmax_log2], dtype=torch.float32, device=geometry.device ) + softmax_scale_tensor = ( + torch.tensor([sm_scale], dtype=torch.float32, device=geometry.device) + if geometry.store_softmax_stats + else None + ) output_scale_tensor = torch.tensor( [output_scale], dtype=torch.float32, device=geometry.device ) @@ -2727,6 +2741,7 @@ def plan( self._plan_state = _ContextPlanState( geometry=geometry, scale_softmax_log2=scale_tensor, + softmax_scale=softmax_scale_tensor, output_scale=output_scale_tensor, empty_i32=empty_i32, variable_window_padded_starts=variable_window_padded_starts, @@ -2749,6 +2764,7 @@ def run( variable_window_cta_starts: Optional[torch.Tensor] = None, out: Optional[torch.Tensor] = None, softmax_stats: Optional[torch.Tensor] = None, + softmax_scale: Optional[torch.Tensor] = None, scale_softmax_log2: Optional[torch.Tensor] = None, output_scale: Optional[torch.Tensor] = None, validate: bool = True, @@ -2783,6 +2799,12 @@ def run( ``output_scale``; merged outputs must use a common output scale. Must not overlap ``out``, any live input, or a plan-owned buffer; storage overlap is a caller precondition and is not checked. + softmax_scale : torch.Tensor, optional + Natural-log one-element float32 scale for exported statistics. + Defaults to the exact plan scale. With statistics enabled, a + runtime overrides must supply both this and ``scale_softmax_log2``; + callers must keep their values consistent before graph replay. + With ``validate=False``, no tensor value is read back to the host. scale_softmax_log2 : torch.Tensor, optional Per-run one-element float32 softmax scale in base-2 form. Defaults to the scale retained by the plan. @@ -2877,6 +2899,20 @@ def run( runtime_window_starts = state.empty_i32 runtime_window_ends = state.empty_i32 runtime_window_cta_starts = state.empty_i32 + if geometry.store_softmax_stats: + if (softmax_scale is None) != (scale_softmax_log2 is None): + raise ValueError( + "softmax_scale and scale_softmax_log2 must be overridden together " + "when store_softmax_stats=True" + ) + if softmax_scale is None: + softmax_scale = state.softmax_scale + elif validate: + _validate_runtime_scale_tensor( + softmax_scale, "softmax_scale", device=geometry.device + ) + elif softmax_scale is not None: + raise ValueError("softmax_scale requires store_softmax_stats=True") if scale_softmax_log2 is None: scale_softmax_log2 = state.scale_softmax_log2 elif validate: @@ -2920,6 +2956,7 @@ def run( runtime_window_ends, runtime_window_cta_starts, softmax_stats, + softmax_scale if geometry.store_softmax_stats else None, ) return out @@ -3106,11 +3143,16 @@ def plan( scale_tensor = torch.tensor( [scale_softmax_log2], dtype=torch.float32, device=geometry.device ) + softmax_scale_tensor = ( + torch.tensor([sm_scale], dtype=torch.float32, device=geometry.device) + if geometry.store_softmax_stats + else None + ) output_scale_tensor = torch.tensor( [output_scale], dtype=torch.float32, device=geometry.device ) - # Keep the two plan-owned scalar allocations ahead of CUTLASS JIT, + # Keep plan-owned scalar allocations ahead of CUTLASS JIT, # matching the compute-sanitizer-safe ordering of the other wrapper. compiled, policy = _get_compiled_paged_context( _paged_context_compile_spec(geometry) @@ -3119,6 +3161,7 @@ def plan( self._plan_state = _PagedContextPlanState( geometry=geometry, scale_softmax_log2=scale_tensor, + softmax_scale=softmax_scale_tensor, output_scale=output_scale_tensor, compiled=compiled, policy=policy, @@ -3136,6 +3179,7 @@ def run( *, out: Optional[torch.Tensor] = None, softmax_stats: Optional[torch.Tensor] = None, + softmax_scale: Optional[torch.Tensor] = None, scale_softmax_log2: Optional[torch.Tensor] = None, output_scale: Optional[torch.Tensor] = None, validate: bool = True, @@ -3194,6 +3238,12 @@ def run( ``sum(exp(logit - m))``, without internal probability scaling. Must not overlap output, inputs, or plan buffers; overlap is not checked. CUDA graphs retain this buffer's address. + softmax_scale : torch.Tensor, optional + Natural-log one-element float32 scale for exported statistics. + Defaults to the exact plan scale. With statistics enabled, a + runtime overrides must supply both this and ``scale_softmax_log2``; + callers must keep their values consistent before graph replay. + With ``validate=False``, no tensor value is read back to the host. scale_softmax_log2 : torch.Tensor, optional Per-run one-element float32 softmax scale in base-2 form. Defaults to the scale retained by the plan. @@ -3235,6 +3285,20 @@ def run( total_q=int(q.shape[0]), num_physical_pages=int(k_cache.shape[0]), ) + if geometry.store_softmax_stats: + if (softmax_scale is None) != (scale_softmax_log2 is None): + raise ValueError( + "softmax_scale and scale_softmax_log2 must be overridden together " + "when store_softmax_stats=True" + ) + if softmax_scale is None: + softmax_scale = state.softmax_scale + elif validate: + _validate_runtime_scale_tensor( + softmax_scale, "softmax_scale", device=geometry.device + ) + elif softmax_scale is not None: + raise ValueError("softmax_scale requires store_softmax_stats=True") if scale_softmax_log2 is None: scale_softmax_log2 = state.scale_softmax_log2 elif validate: @@ -3267,6 +3331,7 @@ def run( block_tables, seq_lens_kv, softmax_stats, + softmax_scale if geometry.store_softmax_stats else None, ) return out diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py index 2c4a64d77a92..93481959aa08 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py @@ -751,6 +751,7 @@ def build_context_task_manager( output_scale: cute.Tensor | None = None, softmax_stats: cute.Tensor | None = None, g_block_tables: cute.Pointer | None = None, + stats_scale: cute.Tensor | None = None, block_table_row_stride: int | Int32 = 0, g_seq_lens_kv: cute.Pointer | None = None, max_seq_len_kv: int | Int32 | None = None, @@ -1275,6 +1276,7 @@ def make_page_offsets_pipeline_cfg(num_stages: int) -> PipelineConfig: scale_softmax_log2=scale_softmax_log2, output_scale=output_scale, softmax_stats=softmax_stats, + stats_scale=stats_scale, cum_seqlen_q=cum_seqlen_q, q_half=0, name="tmem_vec0", @@ -1364,6 +1366,7 @@ def make_page_offsets_pipeline_cfg(num_stages: int) -> PipelineConfig: scale_softmax_log2=scale_softmax_log2, output_scale=output_scale, softmax_stats=softmax_stats, + stats_scale=stats_scale, cum_seqlen_q=cum_seqlen_q, q_half=1, name="tmem_vec1", @@ -2676,6 +2679,7 @@ def build_fmha_task_manager( is_clc_dynamic: bool = False, clc_response_ptr: cute.Pointer | None = None, exhaustive_deadlock_race_check: bool = True, + stats_scale: cute.Tensor | None = None, ) -> Tuple[ TaskManager, list[MemoryResource], @@ -2761,6 +2765,7 @@ def build_fmha_task_manager( scale_softmax_log2=scale_softmax_log2, output_scale=output_scale, softmax_stats=softmax_stats, + stats_scale=stats_scale, g_block_tables=g_block_tables, block_table_row_stride=block_table_row_stride, g_seq_lens_kv=g_seq_lens_kv, @@ -3153,6 +3158,7 @@ def __call__( variable_window_token_ends: cute.Tensor | None = None, variable_window_cta_starts: cute.Tensor | None = None, softmax_stats: cute.Tensor | None = None, + stats_scale: cute.Tensor | None = None, ) -> None: """Set up TMA descriptors, compute grid, and launch the kernel. @@ -3487,6 +3493,7 @@ def __call__( Int32(s_q), self.is_persistent, self.is_clc_dynamic, + stats_scale=stats_scale, ).launch( grid=grid, block=[block_size, 1, 1], @@ -3528,6 +3535,7 @@ def kernel( variable_window_q_stride: Int32, is_persistent: cutlass.Constexpr[bool] = True, is_clc_dynamic: cutlass.Constexpr[bool] = False, + stats_scale: cute.Tensor | None = None, ) -> None: """Warp-specialised persistent FMHA TS kernel. @@ -3591,6 +3599,7 @@ def kernel( scale_softmax_log2=scale_softmax_log2, output_scale=output_scale, softmax_stats=softmax_stats, + stats_scale=stats_scale, q_offset=q_offset, is_persistent=is_persistent, is_clc_dynamic=is_clc_dynamic, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py index c4b11379ab30..e9f0e8f093ec 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py @@ -4655,6 +4655,7 @@ class TmemStatsResource(MemoryResource): scale_softmax_log2: cute.Tensor | None = field(init=False, default=None) output_scale: cute.Tensor | None = field(init=False, default=None) softmax_stats: cute.Tensor | None = field(init=False, default=None) + stats_scale: cute.Tensor | None = field(init=False, default=None) cum_seqlen_q: cute.Tensor | None = field(init=False, default=None) q_half: Constexpr[int] = 0 tmem_addr_cached: TmemAddr | None = field(init=False, default=None) @@ -4677,6 +4678,7 @@ def __init__( scale_softmax_log2: cute.Tensor | None = None, output_scale: cute.Tensor | None = None, softmax_stats: cute.Tensor | None = None, + stats_scale: cute.Tensor | None = None, cum_seqlen_q: cute.Tensor | None = None, q_half: int = 0, **kwargs: Any, @@ -4688,6 +4690,7 @@ def __init__( self.scale_softmax_log2 = scale_softmax_log2 self.output_scale = output_scale self.softmax_stats = softmax_stats + self.stats_scale = stats_scale self.cum_seqlen_q = cum_seqlen_q self.q_half = q_half self._alloc = TmemAllocation(f"tmem_vec_{tmem_vec_offset}", cfg.tmem_stats_cols) @@ -4904,7 +4907,7 @@ def store_vec( ) # Softmax uses exp2(raw_score * scale_log2). The chunk merger # expects maxima in natural-log units, not raw QK scores or LSE. - maximum = row_max * self.scale_softmax_log2[0] * Float32(0.6931471805599453) + maximum = row_max * self.stats_scale[0] # FP8 PV stores exp(logit - maximum) scaled by 448. Export the # unscaled denominator; keep the internal correction stats intact. denominator = row_sum / Float32(self.cfg.pv_p_scale) 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 5dc25c60ebf5..8bad4a884e1b 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 @@ -370,6 +370,7 @@ def _build_decode_gen_schedule( sparse_row_route_begin: Int32 | None = None, sparse_route_count: Int32 | None = None, sage_scales_ptr: cute.Pointer | None = None, + stats_scale: Float32 | None = None, ) -> tuple[ list[Task], dict[MemoryResource, list[MemoryResource]], @@ -1425,6 +1426,7 @@ def _sage_k_scales(inst_id: int, softmax_grp) -> SageKScalesResource: split_kv_counter_ptr=split_kv_counter_ptr, attention_sinks_ptr=attention_sinks_ptr, softmax_stats_ptr=softmax_stats_ptr, + stats_scale=stats_scale, seqlens_kv=kv_seqlens, max_seq_len_kv=corr_max_seq_len_kv, num_heads_kv=num_heads_kv, @@ -1455,6 +1457,7 @@ def _sage_k_scales(inst_id: int, softmax_grp) -> SageKScalesResource: split_kv_counter_ptr=split_kv_counter_ptr, attention_sinks_ptr=attention_sinks_ptr, softmax_stats_ptr=softmax_stats_ptr, + stats_scale=stats_scale, seqlens_kv=kv_seqlens, max_seq_len_kv=corr_max_seq_len_kv, num_heads_kv=num_heads_kv, @@ -2593,6 +2596,7 @@ def _run_decode_gen_active( tma_desc_v_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, g_sage_scales: cute.Pointer | None = None, g_softmax_stats: cute.Pointer | None = None, + stats_scale: Float32 | None = None, ) -> None: """Run the complete decode body for one runtime-valid Q tile. @@ -2714,6 +2718,7 @@ def _run_decode_gen_active( split_kv_counter_ptr=g_split_kv_counter, attention_sinks_ptr=g_attention_sinks, softmax_stats_ptr=g_softmax_stats, + stats_scale=stats_scale, seqlens_kv=runtime_seqlens_kv, cu_seqlens_q=g_cu_seqlens_q, max_seq_len_kv=runtime_max_seq_len_kv, @@ -2936,6 +2941,7 @@ def _run_decode_gen_runtime_prefix( tma_desc_v_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, g_sage_scales: cute.Pointer | None = None, g_softmax_stats: cute.Pointer | None = None, + stats_scale: Float32 | None = None, ) -> None: """Run the general runtime split-prefix producer or retire its suffix.""" @@ -3022,6 +3028,7 @@ def _run_decode_gen_runtime_prefix( tma_desc_v_summary_atom=tma_desc_v_summary_atom, g_sage_scales=g_sage_scales, g_softmax_stats=g_softmax_stats, + stats_scale=stats_scale, ) else: _run_decode_gen_inactive_cluster_rank() @@ -3077,6 +3084,7 @@ def _run_decode_gen_runtime_prefix( tma_desc_v_summary_atom=tma_desc_v_summary_atom, g_sage_scales=g_sage_scales, g_softmax_stats=g_softmax_stats, + stats_scale=stats_scale, ) else: _signal_padded_pdl_producer(cfg) @@ -3126,6 +3134,7 @@ def decode_gen_kernel( tma_desc_v_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, g_sage_scales: cute.Pointer | None = None, g_softmax_stats: cute.Pointer | None = None, + stats_scale: Float32 | None = None, ) -> None: """Dispatch one static Q/split tile and drain padded launch slots safely.""" q_group_cta_idx, h_k_idx, b_idx = cute.arch.block_idx() @@ -3213,6 +3222,7 @@ def decode_gen_kernel( g_sparse_row_route_counts, g_sparse_route_metadata, g_softmax_stats=g_softmax_stats, + stats_scale=stats_scale, ) elif cutlass.const_expr(cfg.use_split_kv and cfg.use_pdl): # The immediately preceding grid may produce seq_lens. Send every @@ -3267,6 +3277,8 @@ def decode_gen_kernel( tma_desc_k_summary_atom=tma_desc_k_summary_atom, tma_desc_v_summary_atom=tma_desc_v_summary_atom, g_sage_scales=g_sage_scales, + g_softmax_stats=g_softmax_stats, + stats_scale=stats_scale, ) else: _run_decode_gen_runtime_prefix( @@ -3318,6 +3330,7 @@ def decode_gen_kernel( tma_desc_v_summary_atom=tma_desc_v_summary_atom, g_sage_scales=g_sage_scales, g_softmax_stats=g_softmax_stats, + stats_scale=stats_scale, ) else: # Packed-Q grids use a batch-wide maximum envelope. These Q CTAs own no @@ -3762,6 +3775,7 @@ def fmha_decode_launch( static_full_split_prefix, g_sage_scales=sage_scales_iter, g_softmax_stats=softmax_stats_iter, + stats_scale=scale_s, ).launch( grid=grid, block=[cfg.threads_per_cta, 1, 1], diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_softmax.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_softmax.py index bca4608eed8e..cd661b194921 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_softmax.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_softmax.py @@ -389,6 +389,7 @@ def _attention_sink_for_local_head( h_r: Int32, num_heads_kv: Int32, local_head_idx: Int32, + stats_scale: Float32 | None = None, ) -> Float32: """Return the attention-sink denominator contribution for one local head.""" # Attention sinks add a synthetic key/value entry to the normalization @@ -406,10 +407,15 @@ def _attention_sink_for_local_head( dtype=Float32, ) sink_val = sink_ptr.load(count=1, alignment=4)[0] - sink_exp = cute.math.exp2( - sink_val * Float32(1.4426950408889634) - max_val * scale_softmax_log2, - fastmath=True, - ) + if cutlass.const_expr(cfg.store_softmax_stats): + # Center in natural-log units before the base conversion so exported + # statistics and normalization share the exact original FP32 anchor. + sink_delta = (sink_val - max_val * stats_scale) * Float32(1.4426950408889634) + else: + sink_delta = ( + sink_val * Float32(1.4426950408889634) - max_val * scale_softmax_log2 + ) + sink_exp = cute.math.exp2(sink_delta, fastmath=True) if cutlass.const_expr(cfg.use_fp8_pv): sink_exp = sink_exp * Float32(cfg.fp8_p_quant_scale) return sink_exp @@ -427,6 +433,7 @@ def _attention_sink_for_scale_idx( logical_q_group_idx: Int32, col_group_idx: Int32, scale_idx: Constexpr[int], + stats_scale: Float32 | None = None, ) -> Float32: """Map a softmax scale group to its attention-sink contribution.""" tile_row_idx = ( @@ -450,4 +457,5 @@ def _attention_sink_for_scale_idx( head_stride, num_heads_kv, local_head_idx, + stats_scale=stats_scale, ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_corr.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_corr.py index 762572c12280..8ed84da0873b 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_corr.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_corr.py @@ -21,7 +21,6 @@ """ from dataclasses import dataclass -import math import cutlass import cutlass.cute as cute @@ -116,6 +115,7 @@ class TmemCorrResource(DecodeGenResourceBase): softmax_local0_ref: Constexpr[TmemSoftmaxLocalResource] = None softmax_local1_ref: Constexpr[TmemSoftmaxLocalResource] = None scale_softmax_log2: Float32 = None + stats_scale: Float32 = None output_scale: Float32 = None o_ptr: cute.Pointer = None partial_o_ptr: cute.Pointer = None @@ -662,7 +662,7 @@ def _store_final_stats( ) scaled_max = Float32(-Float32.inf) if sum_val > Float32(0.0): - scaled_max = max_val * self.scale_softmax_log2 * Float32(math.log(2.0)) + scaled_max = max_val * self.stats_scale denominator = sum_val if cutlass.const_expr(self.cfg.use_fp8_pv): denominator *= Float32(1.0 / 448.0) @@ -1238,6 +1238,7 @@ def _softmax_output_row_state( attention_sink_h_r, self.num_heads_kv, attention_sink_head_idx, + stats_scale=self.stats_scale, ) if cutlass.const_expr(cfg.store_softmax_stats): if store_stats: @@ -3443,6 +3444,7 @@ def _keeps_tail_epilogue( attention_sink_h_r, self.num_heads_kv, attention_sink_head_idx, + stats_scale=self.stats_scale, ) if cutlass.const_expr(cfg.store_softmax_stats): if col_base == Int32(0) and valid_output_row: @@ -3879,6 +3881,7 @@ def _swaps_wide_tail_epilogue( logical_q_group_idx, col_group_idx, scale_idx, + stats_scale=self.stats_scale, ) base_addr0 = self._swaps_o_stage_base_addr( @@ -4385,6 +4388,7 @@ def _swaps_q8_tail_epilogue( logical_q_group_idx, col_group_idx, scale_idx, + stats_scale=self.stats_scale, ) base_addr0 = self._swaps_o_stage_base_addr( diff --git a/tests/unittest/_torch/attention/test_prims_ts_fmha.py b/tests/unittest/_torch/attention/test_prims_ts_fmha.py index d7831b62237f..7be1629fcbdf 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_fmha.py +++ b/tests/unittest/_torch/attention/test_prims_ts_fmha.py @@ -1162,11 +1162,13 @@ def test_sequence_lengths_are_live_source_view() -> None: @pytest.mark.parametrize("num_heads", [6, 96]) @pytest.mark.parametrize("fp8", [False, True]) @pytest.mark.parametrize("mask", [PredefinedAttentionMask.CAUSAL, PredefinedAttentionMask.FULL]) +@pytest.mark.parametrize("store_softmax_stats", [False, True]) def test_mla_context_wrapper_reuses_plan_and_updates_chunk_offsets( monkeypatch: pytest.MonkeyPatch, num_heads: int, fp8: bool, mask: PredefinedAttentionMask, + store_softmax_stats: bool, ) -> None: monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) attn = _Attention(head_dim=192, is_mla=True, num_heads=num_heads, num_kv_heads=num_heads) @@ -1180,7 +1182,7 @@ def test_mla_context_wrapper_reuses_plan_and_updates_chunk_offsets( v = torch.empty(9, num_heads * 256, dtype=q.dtype)[:, num_heads * 128 :] qo_indptr = torch.tensor([0, 2, 5], dtype=torch.int32) output = torch.empty(5, num_heads, 128, dtype=q.dtype) - stats = torch.empty(5, num_heads, 2) + stats = torch.empty(5, num_heads, 2) if store_softmax_stats else None meta = SimpleNamespace( mla_prepare_ctx_cu_seqlens=lambda: qo_indptr, max_seq_len=4096, @@ -1232,7 +1234,7 @@ def test_mla_context_wrapper_reuses_plan_and_updates_chunk_offsets( assert wrapper.run.call_args.kwargs["softmax_stats"] is stats assert wrapper.run.call_args.kwargs["out"] is output assert wrapper.run.call_args.args[1].shape == (9, num_heads, 192) - assert wrapper.plan.call_args.kwargs["store_softmax_stats"] + assert wrapper.plan.call_args.kwargs["store_softmax_stats"] is store_softmax_stats assert wrapper.plan.call_args.kwargs["max_seq_len_q"] == 256 assert wrapper.plan.call_args.kwargs["max_kv_len"] == 128 processed_dtype = torch.float8_e4m3fn if fp8 else q.dtype @@ -1243,8 +1245,27 @@ def test_mla_context_wrapper_reuses_plan_and_updates_chunk_offsets( assert wrapper.plan.call_args.kwargs["mask_type"] == ( "causal" if mask == PredefinedAttentionMask.CAUSAL else "dense" ) - assert (wrapper.run.call_args.kwargs["scale_softmax_log2"] is not None) == fp8 - assert (wrapper.run.call_args.kwargs["output_scale"] is not None) == fp8 + for preprocess_call, run_call in zip( + native_preprocess.call_args_list, wrapper.run.call_args_list, strict=True + ): + bmm1_scale, bmm2_scale = preprocess_call.args[-2:] + softmax_scale = run_call.kwargs["softmax_scale"] + scale_softmax_log2 = run_call.kwargs["scale_softmax_log2"] + output_scale = run_call.kwargs["output_scale"] + if fp8 and store_softmax_stats: + assert softmax_scale.shape == (1,) + assert softmax_scale.data_ptr() == bmm1_scale.data_ptr() + else: + assert softmax_scale is None + if fp8: + assert scale_softmax_log2.shape == (1,) + assert scale_softmax_log2.data_ptr() == ( + bmm1_scale.data_ptr() + bmm1_scale.element_size() + ) + assert output_scale.data_ptr() == bmm2_scale.data_ptr() + else: + assert scale_softmax_log2 is None + assert output_scale is None assert wrapper.run.call_args.kwargs["validate"] is False From a7da965b611830fbc1347e84646b687adc177f67 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Sun, 1 Nov 2026 15:04:50 +0000 Subject: [PATCH 19/19] [None][fix] Preserve rounded PrimTS attention-sink anchors Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- 3rdparty/vendor_sources.lock.yaml | 4 ++-- .../fmha_decode/fmha_decode_resources/helpers_softmax.py | 7 ++++++- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 51aed6027efc..d2753311efdb 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-ts-mla-prefill--flashinfer - commit: eef263b632a86a1003f8c3824c29a8bc8d3957ff + commit: 069139a65aad98f3e5c61f049c761e28eff0b14e 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:011e0dab0ea60efff540c3ce80a96970f1be9c85f6f5b67522d157c31fa710fa - digest: sha256-tree-v1:60780d99f92f1b186c4d3dc2e0600480a3d93ddd6fde94cdd61d3dc7adcb4baf + digest: sha256-tree-v1:a5e826af66f2cb96c5c7ce1bb785d5a413e2d2a997d7c5057fdb99080ee8c5db openengine-proto: url: https://github.com/ai-dynamo/openengine.git tag: v0.1.0 diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_softmax.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_softmax.py index cd661b194921..1b71a958ff97 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_softmax.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_softmax.py @@ -410,7 +410,12 @@ def _attention_sink_for_local_head( if cutlass.const_expr(cfg.store_softmax_stats): # Center in natural-log units before the base conversion so exported # statistics and normalization share the exact original FP32 anchor. - sink_delta = (sink_val - max_val * stats_scale) * Float32(1.4426950408889634) + # Explicit rounding prevents contraction with the sink subtraction: + # that must use the same rounded maximum published in the statistics. + scaled_max = cute.math.mul( + max_val, stats_scale, rounding=cute.math.RoundingMode.NEAREST_EVEN + ) + sink_delta = (sink_val - scaled_max) * Float32(1.4426950408889634) else: sink_delta = ( sink_val * Float32(1.4426950408889634) - max_val * scale_softmax_log2