diff --git a/3rdparty/vendor_patches/flashinfer-prims-ts.patch b/3rdparty/vendor_patches/flashinfer-prims-ts.patch index 6cfb3f7bf8e3..6a6d7894f5ef 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 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 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): +@@ -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,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 b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e231c98a52 100644 +index 7774f881518b11f8e3fc0626afbcd757c77ca6cc..3c3a5d2197cfbc33fcedc7d6c2fb21ebaac2a74d 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 b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 from .decode import ( _WorkspaceSection, -@@ -1047,11 +1044,35 @@ def _make_mla_decode_compile_spec( +@@ -1155,11 +1151,35 @@ def _make_mla_decode_compile_spec( ) @@ -481,26 +476,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"]), +@@ -1798,7 +1818,7 @@ class BatchMLADecodePagedTSWrapper: + use_kv_bounds=use_kv_bounds, ) - @flashinfer_experimental_api(trace=prims_ts_decode_mla_wrapper_trace_dispatch) @@ -508,7 +485,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 def run( self, query: torch.Tensor, -@@ -1742,7 +1764,7 @@ class BatchMLADecodePagedTSWrapper: +@@ -1927,7 +1947,7 @@ class BatchMLADecodePagedTSWrapper: ) @@ -517,7 +494,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( +@@ -2306,6 +2326,7 @@ def _prepare_sparse_output_quant( ) @@ -525,7 +502,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 def _resolve_sparse_mla_plan( device, batch_size, -@@ -2119,17 +2142,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 @@ -547,7 +524,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 MlaDecodeTs, ) -@@ -2264,6 +2287,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 @@ -557,7 +534,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: +@@ -2504,6 +2528,11 @@ class BatchSparseMLADecodePagedTSWrapper: self._workspace_buffer = workspace_buffer self._state = None self._scalar_cache = {} @@ -569,7 +546,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 self.workspace_size_bytes = 0 def plan( -@@ -2386,6 +2417,7 @@ class BatchSparseMLADecodePagedTSWrapper: +@@ -2593,6 +2622,7 @@ class BatchSparseMLADecodePagedTSWrapper: kernel_signature=_mla_kernel_compile_signature(static_kernel), ) ) @@ -577,7 +554,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 self._state = dict( device=device, batch=batch_size, -@@ -2608,6 +2640,38 @@ class BatchSparseMLADecodePagedTSWrapper: +@@ -2815,6 +2845,38 @@ class BatchSparseMLADecodePagedTSWrapper: state = self._state if state is None: raise RuntimeError("plan() must be called before run()") @@ -616,7 +593,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 _validate_query( query, packed_query=state["packed"], -@@ -2781,10 +2845,124 @@ class BatchSparseMLADecodePagedTSWrapper: +@@ -2988,10 +3050,126 @@ class BatchSparseMLADecodePagedTSWrapper: if static_scales else state["compiled"], ) @@ -735,6 +712,8 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 + bmm1, + bmm2, + None, ++ None, ++ None, + ) + return out + diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 02bae0bde8cf..d2753311efdb 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: 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:5cbcaf2b60029742c55f622eadc84b32f49e574038c417972955706157715141 - digest: sha256-tree-v1:fa20dc33e52628519dba7f071f5b28a4eb7963768adb6b46eb8000e4fdb805c4 + patch_digest: sha256:011e0dab0ea60efff540c3ce80a96970f1be9c85f6f5b67522d157c31fa710fa + digest: sha256-tree-v1:a5e826af66f2cb96c5c7ce1bb785d5a413e2d2a997d7c5057fdb99080ee8c5db openengine-proto: url: https://github.com/ai-dynamo/openengine.git tag: v0.1.0 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); } 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/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..6a4b9d6c2716 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 @@ -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 ( @@ -58,10 +62,46 @@ _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 +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", @@ -105,16 +145,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): @@ -130,16 +169,18 @@ 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 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." @@ -153,11 +194,11 @@ def get_attention_feature_unsupported_reason( class PrimsTSFmha(PhasedFmha): - """Blackwell task-scheduled paged 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} @@ -170,13 +211,17 @@ 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_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" + ] = {} # 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 @@ -192,17 +237,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 @@ -293,7 +341,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." @@ -314,7 +362,13 @@ 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 = attn.is_mla_enable + is_mla_generation = ( + is_mla 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: @@ -331,10 +385,15 @@ 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 + 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 @@ -361,12 +420,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 +435,9 @@ 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 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: @@ -384,9 +445,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: @@ -428,14 +529,58 @@ 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: 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 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()) max_seq_len = int(meta.max_seq_len) @@ -776,10 +921,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) + # 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: return wrapper if torch.cuda.is_current_stream_capturing(): @@ -803,9 +953,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( @@ -913,7 +1065,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 +1079,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, @@ -984,29 +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, - 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, @@ -1090,6 +1259,141 @@ 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, 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: + 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, + 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, + ) + 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.") @@ -1472,6 +1776,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, @@ -1486,6 +1805,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, @@ -1495,6 +1816,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/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/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/tensorrt_llm/_torch/attention/backends/prims_ts/context.py b/tensorrt_llm/_torch/attention/backends/prims_ts/context.py index 1ec1080b942f..41048a80b43e 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) @@ -156,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) @@ -181,6 +183,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) @@ -192,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[ @@ -234,6 +238,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 +263,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 +285,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 +334,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 +789,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 +802,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 +1280,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 +1361,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 +1501,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 +1537,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 +1605,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 +1709,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 +1747,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 +1810,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,11 +1843,13 @@ 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], 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.""" @@ -1845,6 +1870,8 @@ 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, + stats_scale=stats_scale, ) else: fmha( @@ -1859,6 +1886,8 @@ 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, + stats_scale=stats_scale, ) def fake_compact(dtype, shape, assumed_align): @@ -1893,6 +1922,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,11 +1965,13 @@ 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, max_seq_len_q, max_seq_len_k, + scale_fake if compile_spec.store_softmax_stats else None, options=_COMPILE_OPTIONS, ) @@ -2012,6 +2048,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,10 +2080,12 @@ 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], static_max_kv_len: cutlass.Constexpr[int], + stats_scale: cute.Tensor | None, ) -> None: fmha( q, @@ -2062,6 +2101,8 @@ 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, + stats_scale=stats_scale, ) def fake_compact(dtype, shape, assumed_align): @@ -2086,6 +2127,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,10 +2158,12 @@ def compile_kernel(): qo_indptr_fake, block_tables_fake, seq_lens_fake, + softmax_stats_fake, stream_fake, max_active_clusters, max_seq_len_q, max_kv_len, + scale_fake if compile_spec.store_softmax_stats else None, options=_COMPILE_OPTIONS, ) @@ -2198,7 +2246,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 +2277,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 +2527,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 +2603,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 +2611,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 +2651,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 +2691,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) @@ -2609,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 ) @@ -2642,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, @@ -2663,6 +2763,8 @@ 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, + softmax_scale: Optional[torch.Tensor] = None, scale_softmax_log2: Optional[torch.Tensor] = None, output_scale: Optional[torch.Tensor] = None, validate: bool = True, @@ -2687,6 +2789,22 @@ 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. + 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. @@ -2705,6 +2823,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( @@ -2775,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: @@ -2817,6 +2955,8 @@ def run( runtime_window_starts, runtime_window_ends, runtime_window_cta_starts, + softmax_stats, + softmax_scale if geometry.store_softmax_stats else None, ) return out @@ -2886,6 +3026,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 +3101,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 +3131,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) @@ -2997,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) @@ -3010,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, @@ -3026,6 +3178,8 @@ def run( seq_lens_kv: torch.Tensor, *, 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, @@ -3077,6 +3231,19 @@ 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. + 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. @@ -3102,6 +3269,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( @@ -3112,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: @@ -3143,6 +3330,8 @@ def run( qo_indptr, block_tables, seq_lens_kv, + softmax_stats, + softmax_scale if geometry.store_softmax_stats else None, ) return out @@ -3163,6 +3352,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 +3392,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 +3442,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 +3457,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 +3478,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 +3527,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 +3586,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 +3596,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..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 @@ -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, @@ -691,7 +749,9 @@ 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, + 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, @@ -1124,7 +1184,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 @@ -1198,6 +1258,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 +1275,10 @@ 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, + stats_scale=stats_scale, + cum_seqlen_q=cum_seqlen_q, + q_half=0, name="tmem_vec0", ) smem_o_0 = SmemOResource( @@ -1291,6 +1356,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 +1365,10 @@ 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, + stats_scale=stats_scale, + cum_seqlen_q=cum_seqlen_q, + q_half=1, name="tmem_vec1", ) smem_o_1 = SmemOResource( @@ -2537,6 +2607,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} @@ -2581,6 +2669,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, @@ -2590,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], @@ -2674,6 +2764,8 @@ 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, + 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, @@ -3065,6 +3157,8 @@ 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, + stats_scale: cute.Tensor | None = None, ) -> None: """Set up TMA descriptors, compute grid, and launch the kernel. @@ -3383,6 +3477,7 @@ def __call__( tile_sched_params, scale_softmax_log2, output_scale, + softmax_stats, num_kv_tiles, num_seq_tiles, q_offset, @@ -3398,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], @@ -3423,6 +3519,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, @@ -3438,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. @@ -3500,6 +3598,8 @@ def kernel( variable_window_q_stride=variable_window_q_stride, 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 af69ee7d51bc..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 @@ -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 @@ -3207,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 @@ -3272,12 +3278,64 @@ 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: 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 +3396,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 +3602,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 +3655,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 +3769,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 +3827,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) @@ -3858,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) @@ -3937,6 +4044,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 +4123,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 +4270,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 +4310,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 +4654,10 @@ 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) + 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) # Precomputed per-warp TMEM vec address (once, before persistent loop). tmem_vec_addr_cached: TmemAddr | None = field(init=False, default=None) @@ -4550,6 +4677,10 @@ 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, + stats_scale: 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 +4689,10 @@ 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.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) self._smem_alloc = None if cfg.stats_via_smem: @@ -4754,6 +4889,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.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) + 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 +5709,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_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/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..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 @@ -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, @@ -369,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]], @@ -1423,6 +1425,8 @@ 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, + stats_scale=stats_scale, seqlens_kv=kv_seqlens, max_seq_len_kv=corr_max_seq_len_kv, num_heads_kv=num_heads_kv, @@ -1452,6 +1456,8 @@ 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, + stats_scale=stats_scale, seqlens_kv=kv_seqlens, max_seq_len_kv=corr_max_seq_len_kv, num_heads_kv=num_heads_kv, @@ -2589,6 +2595,8 @@ 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, + stats_scale: Float32 | None = None, ) -> None: """Run the complete decode body for one runtime-valid Q tile. @@ -2709,6 +2717,8 @@ 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, + 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, @@ -2930,6 +2940,8 @@ 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, + stats_scale: Float32 | None = None, ) -> None: """Run the general runtime split-prefix producer or retire its suffix.""" @@ -3015,6 +3027,8 @@ 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, + stats_scale=stats_scale, ) else: _run_decode_gen_inactive_cluster_rank() @@ -3069,6 +3083,8 @@ 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, + stats_scale=stats_scale, ) else: _signal_padded_pdl_producer(cfg) @@ -3117,6 +3133,8 @@ 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, + 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() @@ -3203,6 +3221,8 @@ def decode_gen_kernel( g_sparse_row_route_offsets, 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 @@ -3257,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( @@ -3307,6 +3329,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: # Packed-Q grids use a batch-wide maximum envelope. These Q CTAs own no @@ -3354,6 +3378,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 +3774,8 @@ def fmha_decode_launch( null_sparse_route_ptr, 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..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 @@ -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,20 @@ 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. + # 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 + ) + 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 +438,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 +462,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/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..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 @@ -115,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 @@ -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.stats_scale + 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. @@ -1152,7 +1238,13 @@ 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: + 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 +2329,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 +2679,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 +2722,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 +2765,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 +2860,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 +3092,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 +3391,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, @@ -3372,7 +3444,17 @@ 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: + 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. @@ -3799,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( @@ -3975,6 +4058,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 @@ -4302,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( @@ -4482,6 +4569,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/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/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/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 fe2461ad74c4..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"); @@ -104,6 +105,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 +117,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) ) @@ -270,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: @@ -343,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, @@ -362,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 ffc5c7fb609c..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"); @@ -159,6 +160,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 @@ -166,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, @@ -441,6 +444,8 @@ 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, smem_p=smem_p_arr, @@ -491,6 +496,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, @@ -831,6 +837,7 @@ def __init__( seq_len_q=1, batch_size=1, mask_type: MaskType | str = MaskType.CAUSAL, + use_kv_bounds: bool = False, ): """ Parameters @@ -913,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 @@ -977,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, @@ -1069,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, @@ -1112,6 +1122,8 @@ def __call__( output_scale: cutlass.Float32, 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): @@ -1400,6 +1412,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) @@ -1450,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, @@ -1473,6 +1490,9 @@ def _flatten_query_rows_for_tma(t): clc_tile_sched_params, output_quant_args, tma_desc_o, + softmax_stats, + softmax_stats_scale, + kv_bounds, ).launch( grid=grid, block=[cfg.threads_per_cta, 1, 1], @@ -1498,6 +1518,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 +1550,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 +1595,9 @@ 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, + kv_bounds: cute.Tensor | None = None, ) -> None: """MLA decode TS kernel: persistent tile-scheduled execution.""" if cutlass.const_expr(self.direct_sparse): @@ -1848,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), @@ -2005,10 +2034,12 @@ 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, 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, @@ -2077,6 +2108,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 +2189,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( @@ -2167,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) @@ -2183,6 +2218,8 @@ def reduction_kernel( cfg, self.reduction_split_capacity, REDUCTION_ROWS_PER_CTA, + softmax_stats, + softmax_stats_scale, ) @cute.kernel @@ -2197,6 +2234,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.""" @@ -2208,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) @@ -2236,4 +2276,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..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"); @@ -29,6 +30,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, @@ -94,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, @@ -154,6 +160,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 +175,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 +240,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 +445,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 +567,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..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"); @@ -27,6 +28,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 +56,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. @@ -137,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, @@ -217,6 +223,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..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, @@ -115,6 +117,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, @@ -477,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 @@ -2115,7 +2120,9 @@ 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 + 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 @@ -2488,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, @@ -2603,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 ) @@ -2853,17 +2880,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 +3675,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 +3752,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 +4173,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 +4192,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 +4451,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 +4472,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..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 @@ -348,7 +350,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 +361,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 +461,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 +475,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..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, @@ -480,6 +482,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, @@ -518,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, @@ -532,6 +537,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, @@ -737,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, @@ -752,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, @@ -772,6 +781,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 +791,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 +839,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 +862,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, @@ -1157,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, @@ -1171,6 +1187,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, @@ -1356,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, @@ -1375,6 +1394,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 +1404,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 +1452,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, @@ -1677,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, ): @@ -1707,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 @@ -1850,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, @@ -1935,6 +1961,8 @@ def __call__( output_scale: cutlass.Float32, 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() @@ -2173,6 +2201,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 ) @@ -2233,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, @@ -2252,6 +2285,9 @@ def __call__( tile_sched_params, block_split_kvs, output_quant_args, + softmax_stats, + softmax_stats_scale, + kv_bounds, ).launch( grid=grid, block=[cfg.threads_per_cta, 1, 1], @@ -2292,6 +2328,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 +2374,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 +2406,9 @@ def dense_kernel( tile_sched_params: object, sparse_inputs: object, 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() @@ -2493,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, @@ -2509,6 +2553,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 +2677,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 +2692,8 @@ def gmem_reduction_kernel( cu_seqlens_q, cfg, num_reduction_ctas, + softmax_stats, + softmax_stats_scale, ) @cute.kernel @@ -2657,6 +2707,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 +2730,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/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_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/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 7a305b1ec6f2..3c3a5d2197cf 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"); @@ -90,6 +90,8 @@ class _MLADecodeCompileSpec: has_kernel_workspace: bool 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) @@ -131,6 +133,8 @@ class _MLADecodePlanState: compiled: Callable[..., object] policy: tuple[tuple[str, object], ...] split_kv: int + store_softmax_stats: bool = False + use_kv_bounds: bool = False @dataclass(frozen=True) @@ -410,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. @@ -464,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( @@ -486,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) ): @@ -689,6 +714,56 @@ 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], + 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 +813,8 @@ def _resolve_mla_decode_launch_spec( output_dtype_key: str, 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.""" @@ -868,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) @@ -875,6 +953,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, @@ -926,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, @@ -950,6 +1051,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), @@ -959,6 +1061,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 +1127,8 @@ def _make_mla_decode_compile_spec( output_dtype_key: str, 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.""" @@ -1041,6 +1146,8 @@ 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, + use_kv_bounds=use_kv_bounds, ) @@ -1225,8 +1332,25 @@ 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)) + 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. @@ -1248,7 +1372,9 @@ def compile_kernel(): cutlass.Float32(1.0), cutlass.Float32(1.0), stream_fake, - output_quant_args=output_quant_args, + output_quant_args, + stats_fake, + bounds_fake, options=_COMPILE_OPTIONS, ) @@ -1279,6 +1405,8 @@ 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, + use_kv_bounds: bool = False, ) -> int: """Return caller-workspace bytes for one automatic MLA policy. @@ -1289,11 +1417,20 @@ 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. + ``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. """ + 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) @@ -1330,6 +1467,8 @@ def get_prims_ts_batch_mla_decode_workspace_size( _dtype_key(out_dtype), 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 @@ -1439,6 +1578,8 @@ def _launch_mla_decode( split_kv: int, 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.""" @@ -1480,12 +1621,11 @@ 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 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 @@ -1515,6 +1655,8 @@ def plan( kv_data_type: torch.dtype, 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: @@ -1554,6 +1696,13 @@ 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``. + 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. @@ -1564,6 +1713,10 @@ 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(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: @@ -1603,6 +1756,8 @@ def plan( _dtype_key(o_data_type), 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( @@ -1616,6 +1771,8 @@ 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, + use_kv_bounds=use_kv_bounds, ) policy = spec.policy workspace_layout = _make_mla_workspace_layout( @@ -1657,6 +1814,8 @@ def plan( compiled=compiled, policy=policy, split_kv=int(dict(policy)["split_kv"]), + store_softmax_stats=store_softmax_stats, + use_kv_bounds=use_kv_bounds, ) @flashinfer_experimental_api @@ -1668,9 +1827,11 @@ 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, + softmax_stats: Optional[torch.Tensor] = None, validate: bool = True, ) -> torch.Tensor: """Launch the most recently planned MLA decode on the current stream. @@ -1699,10 +1860,25 @@ 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 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 +1893,10 @@ 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 + ) + _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, @@ -1750,6 +1930,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, @@ -1761,6 +1942,8 @@ def run( split_kv=state.split_kv, workspace=state.workspace_views, compiled=state.compiled, + softmax_stats=softmax_stats, + kv_bounds=kv_bounds, ) @@ -1783,6 +1966,9 @@ 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, + 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. @@ -1820,6 +2006,16 @@ 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``. + 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. workspace_buffer : torch.Tensor, optional Caller-owned byte scratch, exclusive to one in-flight launch/graph. @@ -1840,8 +2036,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 @@ -1901,6 +2098,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( @@ -1927,6 +2126,8 @@ 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, + use_kv_bounds=kv_bounds is not None, ) workspace_buffer = torch.empty( workspace_bytes, dtype=torch.int8, device=query.device @@ -1949,6 +2150,8 @@ def batch_mla_decode_with_paged_kv_cache( mask_type=mask_type, workspace_buffer=workspace_buffer, validate=validate, + store_softmax_stats=store_softmax_stats, + use_kv_bounds=kv_bounds is not None, ) return wrapper.run( query, @@ -1960,6 +2163,8 @@ def batch_mla_decode_with_paged_kv_cache( bmm2_scale=bmm2_scale, out=out, validate=validate, + softmax_stats=softmax_stats, + kv_bounds=kv_bounds, ) @@ -2960,6 +3165,8 @@ def _run_prepared(self, prepared, query, indices, lengths, extra_indices, extra_ bmm1, bmm2, None, + None, + None, ) return out 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/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/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index e81ec522bc20..084c4af6e171 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 @@ -357,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 00ccf4c18de9..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 @@ -164,6 +165,53 @@ 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"]) + @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. + + 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 ( + 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]) + GSM8K(self.MODEL_NAME).evaluate(llm) + @skip_pre_blackwell @pytest.mark.skip_less_mpi_world_size(16) @pytest.mark.skip_less_device_memory(200000) @@ -412,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 @@ -421,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" @@ -451,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/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/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/_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() 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), 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_attention_backend.py b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py index d0e61dfb5e54..453684312868 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", ) @@ -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,304 @@ 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. + """ + _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 + 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.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 + + 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): + 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 + 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") + 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), + spec_config=spec_config, + ), + 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_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=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; + # 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 = [] + 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] = [] + try: + with torch.inference_mode(): + for step, q_lengths in enumerate(lengths): + if step: + for request_id, length in enumerate(q_lengths): + 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], + 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, + num_extra_kv_tokens=manager.num_extra_kv_tokens, + ), + 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, + 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( + [ + 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]) @@ -500,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( @@ -555,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() diff --git a/tests/unittest/_torch/attention/test_prims_ts_fmha.py b/tests/unittest/_torch/attention/test_prims_ts_fmha.py index c1d1fe078709..7be1629fcbdf 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 @@ -144,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, @@ -151,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, @@ -185,10 +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) - output_width = attn.num_heads * (512 if is_mla else head_dim) + 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, @@ -202,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: @@ -233,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, @@ -255,12 +279,31 @@ 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 = _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 +315,201 @@ 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("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]) +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", [ @@ -353,9 +591,11 @@ 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"), - ({"use_spec_decoding": True}, "speculative decoding"), + ({"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"), + ({"is_spec_dec_tree": True}, "tree speculative decoding"), ({"has_sparse_attention": True}, "sparse attention"), ], ) @@ -580,7 +820,7 @@ def test_is_supported_accepts_and_forwards_phase_keyword( "head_dim": 576, "is_mla": True, }, - "generation-only", + "MLA context requires QK=192", ), ( { @@ -710,6 +950,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), ], @@ -738,7 +981,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() @@ -915,6 +1159,116 @@ 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]) +@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) + 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, 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, num_heads, 128, dtype=q.dtype) + 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, + 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_input_type=AttentionInputType.context_only, + attention_mask=mask, + 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([4, 5], dtype=torch.int32), + batch_size=2, + input_seq_length=3, + max_past_kv_length=7, + tokens_per_block=128, + 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([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, 4, 9], dtype=torch.int32)) + torch.testing.assert_close( + 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, num_heads, 192) + 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 + 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" + ) + 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 + + def test_context_wrapper_plans_once_and_reads_live_fixed_metadata( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -1461,7 +1815,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) @@ -1478,6 +1832,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, @@ -1486,26 +1843,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]], @@ -1532,41 +1907,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, ) - sequence_lengths = torch.tensor([33, 64], dtype=torch.int32) - assert sequence_lengths.data_ptr() % 16 == 0 + 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 = 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, @@ -1585,24 +1975,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 @@ -1616,16 +2008,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() @@ -1636,8 +2043,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]) @@ -1761,8 +2170,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( @@ -1803,6 +2212,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, @@ -1841,6 +2295,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), @@ -1927,10 +2382,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) @@ -1943,7 +2401,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), @@ -1954,11 +2412,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) @@ -1976,6 +2437,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 @@ -1997,7 +2460,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, @@ -2016,9 +2479,10 @@ 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] is wrapper + assert fmha._mla_decode_wrappers[(2, 1)] is wrapper assert workspace.numel() == 48 workspace_size.assert_called_once() @@ -2057,15 +2521,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: @@ -2258,7 +2875,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 +2946,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/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" 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 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 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", [ 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]