From 2f2f78f77f0724f267ad5387a91fbd402462b418 Mon Sep 17 00:00:00 2001 From: yunruis <205571022+yunruis@users.noreply.github.com> Date: Tue, 13 Oct 2026 01:35:09 -0700 Subject: [PATCH 01/19] [None][chore] Update vendored PrimTS to the flashinfer sparse-MLA branch Re-pin flashinfer-prims-ts to yunruis/flashinfer@549977ef (branch yunruis/prims-ts-sparse-mla-group-major-out): FlashInfer main at ae312f7cc (the revision #20155 vendored) plus flashinfer-ai/flashinfer#5434 (native two-source sparse MLA with fused inverse-RoPE + E4M3 output, rebased head 735e1f4f3) and flashinfer-ai/flashinfer#6345 (group-major fused output and the 2CTA schedule for short lists). This brings BatchSparseMLADecodePagedTSWrapper into TensorRT-LLM with the output layout TRT-LLM's DeepSeek-V4 epilogue consumes. The compatibility patch keeps upstream's shims and additionally drops the one-shot MLA trace template (absent from the pinned flashinfer-python 0.6.18) and makes the sparse plan's kernel imports relative so the vendored package does not import flashinfer.attention.prims_ts. AI-assisted. Co-Authored-By: Claude Fable 5.1 Signed-off-by: yunruis <205571022+yunruis@users.noreply.github.com> --- .../vendor_patches/flashinfer-prims-ts.patch | 49 +- 3rdparty/vendor_sources.lock.yaml | 10 +- .../attention/backends/prims_ts/__init__.py | 8 + .../kernels/mla_decode/helpers/constants.py | 20 +- .../kernels/mla_decode/helpers/gather.py | 444 ++++++ .../kernels/mla_decode/helpers/layout.py | 4 +- .../kernels/mla_decode/helpers/math.py | 4 +- .../kernels/mla_decode/helpers/ops.py | 15 +- .../mla_decode/helpers/output_quant.py | 251 ++++ .../kernels/mla_decode/helpers/schedule.py | 12 +- .../kernels/mla_decode/helpers/tile.py | 14 +- .../kernels/mla_decode/kernel_policy.py | 2 +- .../kernels/mla_decode/sparse_policy.py | 507 +++++++ .../kernels/mla_decode/sparse_views.py | 314 +++++ .../mla_decode/throughput_2cta/config.py | 99 +- .../mla_decode/throughput_2cta/kernel.py | 559 ++++++-- .../throughput_2cta/parallel_reduction.py | 140 +- .../mla_decode/throughput_2cta/resources.py | 872 +++++++++++- .../mla_decode/throughput_2cta/tasks.py | 255 +++- .../throughput_latency_1cta/config.py | 283 +++- .../throughput_latency_1cta/kernel.py | 752 +++++++--- .../parallel_reduction.py | 77 +- .../resources/smem_p.py | 10 +- .../resources/smem_resources.py | 345 ++++- .../resources/tmem_corr.py | 433 +++++- .../resources/tmem_o.py | 4 +- .../resources/tmem_p.py | 134 +- .../resources/tmem_s.py | 227 ++- .../resources/tmem_softmax_stats.py | 76 +- .../throughput_latency_1cta/tasks.py | 475 +++++-- .../prims_ts/kernels/separate_reduction.py | 125 +- .../attention/backends/prims_ts/mla_decode.py | 1243 ++++++++++++++++- 32 files changed, 7009 insertions(+), 754 deletions(-) create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/gather.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/output_quant.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/sparse_policy.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/sparse_views.py diff --git a/3rdparty/vendor_patches/flashinfer-prims-ts.patch b/3rdparty/vendor_patches/flashinfer-prims-ts.patch index 48f3254537d8..67992c0ead8b 100644 --- a/3rdparty/vendor_patches/flashinfer-prims-ts.patch +++ b/3rdparty/vendor_patches/flashinfer-prims-ts.patch @@ -201,7 +201,7 @@ index d17584b725e24235ce0ec90caac174a29d4ab965..e6cdddf8d5cb396be5feb828535c0633 + _resolve_kv_tile_idx_context, +) diff --git a/mla_decode.py b/mla_decode.py -index 41ba3b9055d229d227ac10cfc6f7d0f403eae69f..3d5fb159a3298d1edf72935b43d9a2eef74a1a4b 100644 +index 208e495869370e5aea2693548e0ff4d186c70d5e..70c3d0e6423d9810f987e59991cbb7f5b1df8337 100644 --- a/mla_decode.py +++ b/mla_decode.py @@ -1,3 +1,4 @@ @@ -209,19 +209,21 @@ index 41ba3b9055d229d227ac10cfc6f7d0f403eae69f..3d5fb159a3298d1edf72935b43d9a2ee # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); -@@ -21,10 +22,9 @@ from typing import Any, Literal, Optional, cast +@@ -24,12 +25,7 @@ from typing import Any, Literal, NamedTuple, Optional, cast import torch -from flashinfer.api_logging import flashinfer_experimental_api -+from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api - from flashinfer.trace.templates.attention import ( - prims_ts_decode_mla_one_shot_trace_dispatch, +-from flashinfer.trace.templates.attention import ( +- prims_ts_decode_mla_one_shot_trace_dispatch, - prims_ts_decode_mla_wrapper_trace_dispatch, - ) - +-) +- ++from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api from .decode import ( -@@ -1576,7 +1576,7 @@ class BatchMLADecodePagedTSWrapper: + _WorkspaceSection, + _align_up, +@@ -1650,7 +1646,7 @@ class BatchMLADecodePagedTSWrapper: split_kv=int(dict(policy)["split_kv"]), ) @@ -230,6 +232,37 @@ index 41ba3b9055d229d227ac10cfc6f7d0f403eae69f..3d5fb159a3298d1edf72935b43d9a2ee def run( self, query: torch.Tensor, +@@ -1755,7 +1751,7 @@ class BatchMLADecodePagedTSWrapper: + ) + + +-@flashinfer_experimental_api(trace=prims_ts_decode_mla_one_shot_trace_dispatch) ++@flashinfer_experimental_api + def batch_mla_decode_with_paged_kv_cache( + query: torch.Tensor, + kv_cache: torch.Tensor, +@@ -2189,17 +2185,17 @@ def _resolve_sparse_mla_plan( + ): + """Resolve sparse geometry and scratch sizing without allocation/compilation.""" + import cutlass +- from flashinfer.attention.prims_ts.kernels.mla_decode.throughput_latency_1cta.config import ( ++ from .kernels.mla_decode.throughput_latency_1cta.config import ( + MlaProfile, + compute_workspace_size, + ) +- from flashinfer.attention.prims_ts.kernels.mla_decode.throughput_latency_1cta.kernel import ( ++ from .kernels.mla_decode.throughput_latency_1cta.kernel import ( + ThroughputLatencyMlaDecodeTs, + ) +- from flashinfer.attention.prims_ts.kernels.mla_decode.throughput_2cta.config import ( ++ from .kernels.mla_decode.throughput_2cta.config import ( + compute_workspace_size as workspace_2cta, + ) +- from flashinfer.attention.prims_ts.kernels.mla_decode.throughput_2cta.kernel import ( ++ from .kernels.mla_decode.throughput_2cta.kernel import ( + MlaDecodeTs, + ) + diff --git a/q_token_kv_block_sparse_metadata.py b/q_token_kv_block_sparse_metadata.py index 35efc60b7ff3dac6af1c5c5e46631712a171fb72..0c5cf40dbd4e4708b8f74bfa698f3c0788f9f3a4 100644 --- a/q_token_kv_block_sparse_metadata.py diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index ea34d65b41be..184250de470b 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -13,16 +13,16 @@ vendors: patch_digest: sha256:bc9b96a69873e0ddfcba66ab899c666cf52766976f24d850048d2a0bf8d967ff digest: sha256-tree-v1:d2d382335ccb68b55078523e63835c28b699607227afb81308cb7bb1d85845ac flashinfer-prims-ts: - url: https://github.com/yuxianq/flashinfer.git - branch: trtllm-prims-ts-dev-20261012 - commit: 0e0761ef5ef10cb0762cc1b890a5c1a329bb49bf + url: https://github.com/yunruis/flashinfer.git + branch: yunruis/prims-ts-sparse-mla-group-major-out + commit: 549977ef961f00a4ff2feb4c9dd75b0cbf222a3a 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:1ec340aac4fea9c77b9e9d956b1676c09cff5081053ef1e41f962ed38a6f728c - digest: sha256-tree-v1:e4129294b2b26295581cd9cbf1ac5ac4d7eec2ec65f0108626692ddf8a036f96 + patch_digest: sha256:5cf178611ac8a9db64b9465218c2ce800e0e336b80a4129d086093089898e164 + digest: sha256-tree-v1:5cfc19d5eda263af15f8e6661a44e7bfc82d4524e5b6a7734297edbb55de7dce openengine-proto: url: https://github.com/ai-dynamo/openengine.git tag: v0.1.0 diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/__init__.py b/tensorrt_llm/_torch/attention/backends/prims_ts/__init__.py index e40d8ce0d41c..e8187118271f 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/__init__.py @@ -37,6 +37,10 @@ batch_prefill_with_paged_kv_cache, ) from .mla_decode import ( + BatchSparseMLADecodePagedTSWrapper, + SparseMLAPreparedMetadata, + batch_sparse_mla_decode_with_paged_kv_cache, + get_prims_ts_sparse_mla_decode_workspace_size, BatchMLADecodePagedTSWrapper, batch_mla_decode_with_paged_kv_cache, get_prims_ts_batch_mla_decode_workspace_size, @@ -49,6 +53,10 @@ from .sage import SageAttentionConfig, SageAttentionParams __all__ = [ + "SparseMLAPreparedMetadata", + "BatchSparseMLADecodePagedTSWrapper", + "batch_sparse_mla_decode_with_paged_kv_cache", + "get_prims_ts_sparse_mla_decode_workspace_size", "BlockSparseTSWrapper", "SageAttentionConfig", "SageAttentionParams", diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/constants.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/constants.py index aa68b9efa0f1..1e47d5c5b237 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/constants.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/constants.py @@ -34,8 +34,10 @@ WARPGROUP_WARPS = 4 WARPGROUP_THREADS = WARPGROUP_WARPS * WARP_LANES -# Kernel-level TMEM lifecycle synchronization. Barrier ID 15 is reserved for -# alloc/dealloc phases in the 1CTA path; the 2CTA path uses config-owned IDs. +# Kernel-level 1CTA TMEM synchronization: 14 publishes allocation to compute +# warps; 15 handles full-CTA initialization/deallocation. The 2CTA path owns +# its barrier IDs in its configuration. +TMEM_READY_BARRIER_ID = 14 TMEM_LIFECYCLE_BARRIER_ID = 15 TMEM_DEALLOC_MBAR_THREADS = WARP_LANES @@ -69,6 +71,10 @@ TCGEN05_16X32BX2_FP8_P_STRIDE = 16 TCGEN05_SECOND_PANEL_ADDR_OFFSET = 16 << 16 +# Softmax uses base-two exponentials internally; public sparse LSE uses ln. +LN_2 = 0.6931471805599453 +LOG2_E = 1.4426950408889634 + # Softmax scratch stores the max state first and the sum state in a second # fixed-size panel. The offset is in Uint32 scratch words. SOFTMAX_SCRATCH_SUM_WORD_OFFSET = 384 @@ -87,6 +93,16 @@ PAGE_OFFSET_BYTES = 4 CP_ASYNC_CACHE_CA = "ca" +# gather4 always copies four selected KV rows. This instruction width is +# independent of the number of issuing warps or the KV tile size. +TMA_GATHER_ROWS = 4 + +# Sparse 2CTA offset caches derive their sizes from cfg token geometry: +# K owns 128/2=64 rows per CTA; V needs all 128 rows. Lane-distributed +# caches therefore hold 64/32=2 K and 128/32=4 V words per lane. A uniform +# cache instead holds the issuing warp's share (64/W K and 128/W V words +# per lane). Pending V stays lane-distributed until it becomes current. + # Dense MLA kernels specialize paged-KV addressing for this explicit ABI set. # Each listed size exactly partitions a 128-token KV tile; other page sizes are # outside the current kernel ABI. diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/gather.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/gather.py new file mode 100644 index 000000000000..c60f2070fd03 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/gather.py @@ -0,0 +1,444 @@ +# 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. + +"""Blackwell four-row TMA gather with explicit five-coordinate PTX operands. + +The CUTLASS DSL 4.7 experimental wrapper validates gather coordinates as a pair; +it cannot express four independent row coordinates through that interface. +Keep the native instruction isolated here instead of patching the dependency. +""" + +import cutlass +import cutlass.cute as cute +from cutlass import Int32, Int64 +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.experimental import primitives as prims + +from .constants import TMA_GATHER_ROWS + + +@dsl_user_op +def _address(pointer, dtype, *, loc=None, ip=None): + if hasattr(pointer, "data_ptr"): + pointer = pointer.data_ptr(loc=loc, ip=ip) + if hasattr(pointer, "toint"): + pointer = pointer.toint(loc=loc, ip=ip) + return dtype(pointer) + + +@cute.jit +def select_gather_map(primary, extra, use_extra): + """Select numeric addresses: Pointer-object rebinding is not an SSA select.""" + primary_addr = _address(primary, Int64) + extra_addr = _address(extra, Int64) + return extra_addr if use_extra else primary_addr + + +@cute.jit +def load_sparse_rows(routes, request, offset, *, fragments: cutlass.Constexpr[int]): + """Coalesced row window, distributed as one coordinate per lane/fragment.""" + lane = cute.arch.thread_idx()[0] & Int32(31) + rows = cutlass.Array(Int32, fragments, space=cutlass.AddressSpace.rmem) + for fragment in cutlass.range_constexpr(fragments): + index = offset + lane + Int32(fragment * 32) + value = Int32(0x7FFFFFFF) + if index < routes.shape[0]: + value = Int32(routes[index, request]) + rows[fragment] = value + return rows + + +@cute.jit +def load_sparse_quad(routes, request, offset): + """Keep one contiguous row quadruple in its issuing lane's registers.""" + rows = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) + if cutlass.const_expr(hasattr(routes, "routes_for_tile")): + tile_routes = routes.routes_for_tile(offset, request) + return tile_routes.mapped_quad(offset) + for j in cutlass.range_constexpr(4): + rows[j] = Int32(0x7FFFFFFF) + if offset + Int32(3) < routes.shape[0]: + if cutlass.const_expr(routes.stride[0] == 1): + values = cutlass.Pointer( + routes.iterator + offset + request * routes.stride[1], dtype=Int32 + ).load(count=4, alignment=4) + for j in cutlass.range_constexpr(4): + rows[j] = values[j] + else: + for j in cutlass.range_constexpr(4): + rows[j] = Int32(routes[offset + j, request]) + else: + for j in cutlass.range_constexpr(4): + if offset + Int32(j) < routes.shape[0]: + rows[j] = Int32(routes[offset + j, request]) + return rows + + +@cute.jit +def decode_sparse_quad(primary, extra, raw): + """Select a source and map masked row sentinels to TMA zero-fill rows.""" + rows = cutlass.Array(Int32, TMA_GATHER_ROWS, space=cutlass.AddressSpace.rmem) + for j in cutlass.range_constexpr(TMA_GATHER_ROWS): + # Native sparse pools have fewer than INT32_MAX storage rows. The + # reserved sentinel is already an out-of-bounds TMA coordinate; + # retaining it gives zero fill without a compare/select per row. + rows[j] = raw[j] & Int32(0x7FFFFFFF) + return select_gather_map(primary, extra, raw[0] < 0), rows + + +@cute.jit +def gather4_quad( + dst, primary, extra, column, raw, barrier, *, cta_group: cutlass.Constexpr[int] = 1 +): + """Issue a quad from the owning lane, without election or warp shuffles.""" + descriptor, rows = decode_sparse_quad(primary, extra, raw) + gather4( + dst, + descriptor, + column, + rows[0], + rows[1], + rows[2], + rows[3], + barrier, + cta_group=cta_group, + ) + + +@cute.jit +def gather4_cached( + dst, + primary, + extra, + column, + cached_rows, + token_offset, + barrier, + *, + cache_base: cutlass.Constexpr[int] = 0, + cta_group: cutlass.Constexpr[int] = 1, +): + """Assemble a row quadruple using only warp-local register shuffles.""" + raw0 = Int32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=cached_rows[cache_base + token_offset // 32], + offset=token_offset % 32, + mask_and_clamp=0x1F, + kind=prims.Shfl.IDX, + ) + ) + rows = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) + for j in cutlass.range_constexpr(4): + token = token_offset + Int32(j) + raw = Int32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=cached_rows[cache_base + token // 32], + offset=token % 32, + mask_and_clamp=0x1F, + kind=prims.Shfl.IDX, + ) + ) + row = raw & Int32(0x7FFFFFFF) + rows[j] = row if row != Int32(0x7FFFFFFF) else Int32(-1) + descriptor = select_gather_map(primary, extra, raw0 < 0) + if prims.elect_sync(): + gather4( + dst, + descriptor, + column, + rows[0], + rows[1], + rows[2], + rows[3], + barrier, + cta_group=cta_group, + ) + + +@cute.jit +def invalid_sparse_token( + token, length, routes, request, *, sparse: cutlass.Constexpr[bool] +): + """Mask prefix tails and tagged invalid routes before softmax. + + Bit 31 selects the source; bits 0..30 hold its TMA row coordinate. + All-one low bits denote an invalid row for either source. + """ + if cutlass.const_expr(sparse and hasattr(routes, "is_invalid")): + return (token >= length) | routes.is_invalid(token, request) + invalid = token >= length + if cutlass.const_expr(sparse): + if token < length: + invalid = (Int32(routes[token, request]) & Int32(0x7FFFFFFF)) == Int32( + 0x7FFFFFFF + ) + return invalid + + +@cute.jit +def gather4( + dst, + tensor_map, + column, + row0, + row1, + row2, + row3, + barrier, + *, + cta_group: cutlass.Constexpr[int] = 1, +): + """Issue one four-row copy from the calling thread. + + Multiple lanes may issue disjoint copies. The caller accounts for all + copies in the shared completion barrier, including masked/OOB rows. + + The descriptor is 2D with a one-row box and no interleave. Destination + alignment/swizzle and coordinate bounds follow its tensor-map contract. + Group 2 delivers completion to the even CTA of the issuing CTA pair. + """ + dst_addr = _address(dst, Int32) + bar_addr = _address(barrier, Int32) + if cutlass.const_expr(cta_group == 2): + bar_addr = bar_addr & Int32(-16777217) # Clear peer-CTA address bit 24. + instruction = ( + "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4." + "mbarrier::complete_tx::bytes.cta_group::2 " + ) + else: + assert cta_group == 1 + instruction = ( + "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4." + "mbarrier::complete_tx::bytes.cta_group::1 " + ) + cute.arch.inline_ptx( + instruction + "[{$r0}], [{$r1}, {{$r2}, {$r3}, {$r4}, {$r5}, {$r6}}], [{$r7}];", + read_only_args=[ + dst_addr, + _address(tensor_map, Int64), + Int32(column), + Int32(row0), + Int32(row1), + Int32(row2), + Int32(row3), + bar_addr, + ], + ) + + +class CachedSparseMask: + """One/two per-warp bitsets reused across all score registers and heads.""" + + def __init__(self, values, swaps): + self.values = values + self.swaps = swaps + self.low, self.high = values + + def __extract_mlir_values__(self): + return cutlass.extract_mlir_values(self.values) + + def __new_from_mlir_values__(self, values): + return CachedSparseMask( + cutlass.new_from_mlir_values(self.values, values), self.swaps + ) + + @cute.jit + def is_invalid(self, token, request): + del request + word = self.low if self.swaps or (token & Int32(32)) == Int32(0) else self.high + return ((word >> (token & Int32(31))) & Int32(1)) != Int32(0) + + +@cute.jit +def cache_sparse_mask( + routes, + tile_offset, + length, + request, + *, + swaps: cutlass.Constexpr[bool], + column_half=None, + tile_size_kv: cutlass.Constexpr[int] = 128, +): + lane = cute.arch.thread_idx()[0] % 32 + warp = cute.arch.thread_idx()[0] % 128 // 32 + if cutlass.const_expr(hasattr(routes, "prefix_mask32")): + low, high = Int32(0), Int32(0) + if not routes.tile_is_full(tile_offset, length): + if cutlass.const_expr(swaps): + low = routes.prefix_mask32(tile_offset + warp * Int32(32), length) + high = low + else: + if cutlass.const_expr(column_half is None): + column_half = lane >> Int32(4) + start = tile_offset + if cutlass.const_expr(tile_size_kv == 128): + start += column_half * Int32(64) + low = routes.prefix_mask32(start, length) + high = routes.prefix_mask32(start + Int32(32), length) + return CachedSparseMask((low, high), swaps) + if cutlass.const_expr(swaps): + token = tile_offset + warp * Int32(32) + lane + invalid = invalid_sparse_token(token, length, routes, request, sparse=True) + bits = Int32(cute.arch.vote_ballot_sync(invalid)) + return CachedSparseMask((bits, bits), True) + else: + masks = cutlass.Array( + Int32, tile_size_kv // 32, space=cutlass.AddressSpace.rmem + ) + for part in cutlass.range_constexpr(tile_size_kv // 32): + token = tile_offset + Int32(part * 32) + lane + invalid = invalid_sparse_token(token, length, routes, request, sparse=True) + masks[part] = Int32(cute.arch.vote_ballot_sync(invalid)) + if cutlass.const_expr(column_half is None): + column_half = lane >> Int32(4) + if cutlass.const_expr(tile_size_kv == 64): + low, high = masks[0], masks[1] + else: + low = masks[0] if column_half == Int32(0) else masks[2] + high = masks[1] if column_half == Int32(0) else masks[3] + return CachedSparseMask((low, high), False) + + +@dsl_user_op +def _gather4_slices( + dst, + tensor_map, + column, + rows, + barrier, + *, + num_slices, + dst_stride, + col_stride, + loc=None, + ip=None, +): + # All callers broadcast the quad before entering this collective. Keep + # feature slices in one asm block so election/uniformization is shared. + instruction = "cp.async.bulk.tensor.2d.shared::cluster.global.tile::gather4.mbarrier::complete_tx::bytes.cta_group::2 " + asm = [ + "{ .reg .b32 slice_dst, slice_col; .reg .pred issue; elect.sync _|issue, 0xffffffff;", + "mov.b32 slice_dst, {$r0};", + "mov.b32 slice_col, {$r2};", + ] + for i in range(num_slices): + asm.append( + "@issue " + + instruction + + "[slice_dst], [{$r1}, {slice_col, {$r3}, {$r4}, {$r5}, {$r6}}], [{$r7}];" + ) + if i + 1 < num_slices: + asm.append(f"add.s32 slice_dst, slice_dst, {dst_stride};") + asm.append(f"add.s32 slice_col, slice_col, {col_stride};") + asm.append("}") + cute.arch.inline_ptx( + "\n".join(asm), + read_only_args=[ + _address(dst, Int32), + _address(tensor_map, Int64), + Int32(column), + rows[0], + rows[1], + rows[2], + rows[3], + _address(barrier, Int32) & Int32(-16777217), + ], + loc=loc, + ip=ip, + ) + + +@cute.jit +def gather4_uniform_quad_slices( + dst, + primary, + extra, + column, + raw, + barrier, + *, + num_slices: cutlass.Constexpr[int], + dst_stride: cutlass.Constexpr[int], + col_stride: cutlass.Constexpr[int], +): + """Issue 2CTA feature slices from a warp-uniform cached row quadruple. + + Every lane must participate with identical arguments. One elected lane + issues all copies; dst_stride is bytes and col_stride is source elements. + The caller accounts for all copied bytes in the completion barrier. + """ + descriptor, rows = decode_sparse_quad(primary, extra, raw) + _gather4_slices( + dst, + descriptor, + column, + rows, + barrier, + num_slices=num_slices, + dst_stride=dst_stride, + col_stride=col_stride, + ) + + +@dsl_user_op +def gather4_cached_warp( + dst, + tensor_map, + column, + cached_rows, + barrier, + *, + cache_start, + num_quads, + num_slices, + quad_stride, + slice_stride, + col_stride, + loc=None, + ip=None, +): + """Issue one warp's cached quads with one election for every feature slice. + + All lanes participate. Rows are already mapped, cleaned and warp-uniform; + the caller retains them across the head-dimension stages of one K tile. + Strides for destinations are bytes, and the column stride is elements. + """ + assert num_quads <= 8 + args = [ + _address(dst, Int32), + _address(tensor_map, Int64), + Int32(column), + _address(barrier, Int32), + ] + args += [cached_rows[cache_start + i] for i in range(num_quads * 4)] + instruction = "cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4.mbarrier::complete_tx::bytes.cta_group::1 " + asm = [ + "{ .reg .b32 copy_dst, copy_col; .reg .pred issue; elect.sync _|issue, 0xffffffff;" + ] + for part in range(num_slices): + asm.append(f"add.s32 copy_col, {{$r2}}, {part * col_stride};") + for quad in range(num_quads): + asm.append( + f"add.s32 copy_dst, {{$r0}}, {part * slice_stride + quad * quad_stride};" + ) + rows = ", ".join("{$r" + str(4 + quad * 4 + j) + "}" for j in range(4)) + asm.append( + "@issue " + + instruction + + "[copy_dst], [{$r1}, {copy_col, " + + rows + + "}], [{$r3}];" + ) + asm.append("}") + cute.arch.inline_ptx("\n".join(asm), read_only_args=args, loc=loc, ip=ip) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/layout.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/layout.py index fe45148142c3..a9d28a54ecef 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/layout.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/layout.py @@ -54,7 +54,7 @@ def num_softmax_scale_groups(cfg: MlaConfig) -> int: """Return the number of independent softmax scale groups per thread.""" if cfg.kernel_variant == "keeps_mma_ab": return 1 - return max(cfg.tile_size_q // 4, 4) + return max(cfg.tile_size_q // 4, 2) def num_s_regs_per_thread(cfg: MlaConfig) -> int: @@ -73,6 +73,8 @@ def softmax_scratch_words(cfg: MlaConfig) -> int: def num_packed_p_regs(cfg: MlaConfig) -> int: """Return the number of packed P registers per softmax thread.""" + if cfg.kernel_variant == "keeps_mma_ab": + return cfg.tile_size_kv // (8 if cfg.is_fp8_qkv() else 4) q_repeats = max(cfg.tile_size_q // 8, 1) return (2 if cfg.is_fp8_qkv() else 4) * q_repeats diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/math.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/math.py index e07a06325518..b2d9933ec306 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/math.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/math.py @@ -21,9 +21,7 @@ from cutlass import Float32, Int32, Uint32 from cutlass.experimental import primitives as prims -# Softmax converts exp2 inputs with log2(e) and initializes masked scores to -# negative Float32 max. -LOG2_E = 1.4426950408889634074 +# Masked scores start at negative Float32 max. NEG_FLT_MAX = -3.4028235e38 # E4M3FN finite maximum used when clamping FP8 output conversion. diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/ops.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/ops.py index 86f1ca7cee6b..021b9ebbd79f 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/ops.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/ops.py @@ -223,6 +223,12 @@ def fp8_log2_quant_scale(): return Float32(8.8073549) +@cute.jit +def fp8_quant_scale(): + """Return the E4M3 probability scale used by sink normalization.""" + return Float32(448.0) + + @cute.jit def fp8_quant_scale_rcp(): """Return the reciprocal of the shared E4M3 probability scale.""" @@ -334,14 +340,19 @@ def store_transposed_smem8b_x4( @cute.jit -def tcgen05_store_p_16x32bx2_x16(tmem_addr, regs_p, start_idx: cutlass.Constexpr[int]): +def tcgen05_store_p_16x32bx2_x16( + tmem_addr, + regs_p, + start_idx: cutlass.Constexpr[int], + stride: cutlass.Constexpr[int] = TCGEN05_16X32BX2_BF16_P_STRIDE, +): """Store 16 packed BF16 P registers to a split TMEM-P tile.""" # The stride immediate is the BF16 lane-to-column spacing for the split # 16x32bx2 P layout; keep it named because the FP8 path uses a different # immediate with the same instruction shape. inline_ptx( "tcgen05.st.sync.aligned.16x32bx2.x16.b32 " - f"[{{$r0}}], {TCGEN05_16X32BX2_BF16_P_STRIDE}, " + f"[{{$r0}}], {stride}, " "{ {$r1}, {$r2}, {$r3}, {$r4}, {$r5}, {$r6}, {$r7}, {$r8}, " "{$r9}, {$r10}, {$r11}, {$r12}, {$r13}, {$r14}, {$r15}, {$r16} };", read_only_args=[ diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/output_quant.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/output_quant.py new file mode 100644 index 000000000000..e620e90dfb2f --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/output_quant.py @@ -0,0 +1,251 @@ +# 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. + +"""Final sparse MLA output: conjugate RoPE and block E4M3 quantization.""" + +from dataclasses import dataclass + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Int64, Uint32 +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T +from cutlass.experimental import primitives as prims + +from .math import mul_packed_f32x2, fma_packed_f32x2 +from .ops import pack_float4_to_fp8_e4m3 + + +@cutlass.dsl_user_op +def _max_abs3(x, y, z, *, loc=None, ip=None): + """SM100 max-of-magnitudes with native floating-point operand constraints.""" + return Float32( + llvm.inline_asm( + T.f32(), + [Float32(v).ir_value(loc=loc, ip=ip) for v in (x, y, z)], + "max.abs.ftz.f32 $0, $1, $2, $3;", + "=f,f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@dataclass(frozen=True) +class MlaOutputQuant: + """Static geometry; runtime args are (flat O, positions, cos/sin, scales). + + O is E4M3 stored group-major, physically [G, T, (H/G) * 512] (T is derived + from its extent); positions are INT32; cos/sin is FP32; scales are a + token-major logical view of column-major FP32 or packed INT32 storage. + Rows are token * H + head and columns are channels within the D512 head. + """ + + block_size: int + ue8m0: bool + num_heads: int + num_groups: int + + @cute.jit + def rotate(self, args, values, row, column): + """Rotate a mutable final-output fragment in place.""" + _, positions, cos_sin, _ = args + token = Int64(row) // self.num_heads + count = cutlass.const_expr(values.shape[0]) + fragment = values + + # The cache contains forward sin; conjugation changes its sign here. + if column + count - 1 >= 448: + position = positions[token] + # Load independent coefficients together before consuming them. + pairs = cutlass.const_expr(min(32, count // 2)) + for i in cutlass.range_constexpr(0, count, pairs * 2): + d = column + i + if d >= 448: + # Callers provide aligned contiguous channel vectors. + base = ( + cos_sin.iterator.raw_ptr() + + Int64(position) * 64 + + (d - 448) // 2 + ) + alignment = cutlass.const_expr(min(16, pairs * 4)) + cos = base.load(count=pairs, alignment=alignment) + sin = (base + 32).load(count=pairs, alignment=alignment) + for j in cutlass.range_constexpr(pairs): + x, y = fragment[i + 2 * j], fragment[i + 2 * j + 1] + sine = mul_packed_f32x2((y, x), (sin[j], sin[j]), ftz=True) + rotated = fma_packed_f32x2( + (x, y), (cos[j], cos[j]), (sine[0], -sine[1]), ftz=True + ) + fragment[i + 2 * j] = rotated[0] + fragment[i + 2 * j + 1] = rotated[1] + + @cute.jit + def store( + self, + args, + values, + row, + column, + lane_stride: cutlass.Constexpr[int] = 1, + rotary: cutlass.Constexpr[bool] = True, + shared_output=None, + ): + """Finalize contiguous FP32 fragments after all attention reductions. + + Blocks fit in one thread or one warp; ``lane_stride=16`` joins the + two lanes owning the halves of a keep-MMA head. Nonrotary D stages + remove the rotation at compile time, limiting epilogue code size. + ``shared_output`` is the optional 2CTA FP8 output stage: two swizzled + 64-row by 128-byte boxes, consumed by its asynchronous TMA-store task. + """ + if cutlass.const_expr(rotary): + # Some epilogues pass an immutable register vector. Two-CTA + # epilogues rotate their mutable accumulator before this store. + count = cutlass.const_expr(values.shape[0]) + fragment = cutlass.Array(Float32, count, space=cutlass.AddressSpace.rmem) + for i in cutlass.range_constexpr(count): + fragment[i] = values[i] + self.rotate(args, fragment, row, column) + else: + fragment = values + out, _, _, scales = args + token = Int64(row) // self.num_heads + head = Int64(row) % self.num_heads + # Physical [G, T, (H/G) * 512] output; the 512-channel head rows of one + # (group, token) pair are contiguous. Planning bounds T * H below 2**31, + # so the row index fits Int32. + heads_per_group = cutlass.const_expr(self.num_heads // self.num_groups) + total_tokens = Int32(Int64(cute.size(out)) // (self.num_heads * 512)) + head32 = Int32(head) + if cutlass.const_expr(heads_per_group & (heads_per_group - 1) == 0): + shift = cutlass.const_expr(heads_per_group.bit_length() - 1) + group = head32 >> shift + head_in_group = head32 & Int32(heads_per_group - 1) + else: + group = head32 // heads_per_group + head_in_group = head32 % heads_per_group + out_row = Int64( + (group * total_tokens + Int32(token)) * heads_per_group + head_in_group + ) + lane = cute.arch.thread_idx()[0] % 32 + count = cutlass.const_expr(fragment.shape[0]) + per_block = cutlass.const_expr(min(count, self.block_size)) + lanes = cutlass.const_expr(self.block_size // per_block) + # RoPE branches can leave partners at different program counters. + # Name every lane in the quant block instead of sampling activemask. + group_base = lane % lane_stride + lane // (lanes * lane_stride) * ( + lanes * lane_stride + ) + group_bits = cutlass.const_expr( + sum(1 << (i * lane_stride) for i in range(lanes)) + ) + group_mask = Uint32(group_bits) << Uint32(group_base) + for start in cutlass.range_constexpr(0, count, per_block): + # Two independent max3 chains avoid a long scalar amax dependency. + m0, m1 = Float32(1e-4), Float32(1e-4) + for i in cutlass.range_constexpr(0, per_block, 4): + m0 = _max_abs3(m0, fragment[start + i], fragment[start + i + 1]) + m1 = _max_abs3(m1, fragment[start + i + 2], fragment[start + i + 3]) + amax = _max_abs3(m0, m1, Float32(1e-4)) + for step in cutlass.range_constexpr(lanes.bit_length() - 1): + peer = prims.shfl_sync( + thread_mask=group_mask, + val=amax, + offset=(1 << step) * lane_stride, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + amax = cute.math.max(amax, peer) + scale = amax * Float32(1.0 / 448.0) + if cutlass.const_expr(self.ue8m0): + exponent = (scale.bitcast(Uint32) + Uint32(0x007FFFFF)) >> 23 + scale = (exponent << 23).bitcast(Float32) + if cutlass.const_expr(self.ue8m0): + inverse_bits = (Uint32(254) - exponent) << 23 + # UE8M0 scales are powers of two; keep reciprocal(inf) == 0. + inv_scale = ( + Float32(0) if exponent == 255 else inverse_bits.bitcast(Float32) + ) + else: + inv_scale = cute.math.rcp(scale, approx=True) + block = (column + start) // self.block_size + # The validated [G, columns, pad4(T)] backing is contiguous in + # head/block order, so group decomposition is unnecessary. + if lane // lane_stride % lanes == 0: + if cutlass.const_expr(self.ue8m0): + word = head * (512 // self.block_size // 4) + block // 4 + byte = block % 4 + offset = token + word * Int64(scales.stride[2]) + byte_ptr = cutlass.inttoptr( + (scales.iterator.raw_ptr() + offset).toint(Int64) + byte, + mem_space=1, + dtype=cutlass.Uint8, + ) + byte_ptr.store(exponent.to(cutlass.Uint8)) + else: + offset = token + (head * (512 // self.block_size) + block) * Int64( + scales.stride[2] + ) + (scales.iterator.raw_ptr() + offset).store(scale) + + words = cutlass.const_expr(min(8, per_block // 4)) + for i in cutlass.range_constexpr(0, per_block, 4 * words): + v = start + i + packed = cutlass.Array(Int32, words, space=cutlass.AddressSpace.rmem) + for w in cutlass.range_constexpr(words): + j = v + 4 * w + lo = mul_packed_f32x2( + (fragment[j], fragment[j + 1]), (inv_scale, inv_scale), ftz=True + ) + hi = mul_packed_f32x2( + (fragment[j + 2], fragment[j + 3]), + (inv_scale, inv_scale), + ftz=True, + ) + packed[w] = pack_float4_to_fp8_e4m3(lo[0], lo[1], hi[0], hi[1]) + if cutlass.const_expr(shared_output is not None): + for part in cutlass.range_constexpr(0, words, 4): + col = Int32(column + v + part * 4) & Int32(255) + smem_offset = ( + (col // 128) * 8192 + (Int32(head) & 63) * 128 + (col & 127) + ) + # TMA 128B swizzle XORs byte bits [4:7] with [7:10]. + # The 1024B-aligned base leaves these swizzle bits unchanged. + swizzled = smem_offset ^ ((smem_offset >> 3) & 0x70) + destination = cutlass.inttoptr( + shared_output + swizzled, mem_space=3, dtype=Int32 + ) + destination.store(packed.load(part, 4), alignment=16) + else: + ptr = cutlass.inttoptr( + (out.iterator.raw_ptr() + out_row * 512 + column + v).toint( + Int64 + ), + mem_space=1, + dtype=Int32, + ) + if cutlass.const_expr(words == 8): + # Native allocations support a 256-bit store. Retain the + # documented 16-byte alignment for caller-owned buffers. + if (ptr.toint(Int64) & Int64(31)) == 0: + ptr.store(packed.load(0, 8), alignment=32) + else: + ptr.store(packed.load(0, 4), alignment=16) + (ptr + 4).store(packed.load(4, 4), alignment=16) + else: + ptr.store(packed.load(0, words), alignment=4 * words) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/schedule.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/schedule.py index 133c3e15fe48..6e3626b3467c 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/schedule.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/schedule.py @@ -173,6 +173,7 @@ def staged_qk_mma( head_dim_stages=None, consumer_label=None, include_acquire=True, + release_k=True, ): """Consume staged K descriptors and produce one QK MMA score tile.""" if consumer_label is None: @@ -192,7 +193,8 @@ def staged_qk_mma( smem_kv.wait() kv_desc = getattr(smem_kv, consumer_label)(k_subtile_idx=k_subtile_idx) tmem_s.qk_mma(kv_desc=kv_desc, k_subtile_idx=k_subtile_idx) - smem_kv.release() + if release_k: + smem_kv.release() tmem_s.commit() @@ -206,6 +208,7 @@ def staged_pv_mma( consumer_label=None, producer_label=None, is_tail: bool = False, + reuse_k=False, ): """Consume staged V descriptors and produce one PV MMA output tile.""" smem_p.wait() @@ -232,7 +235,8 @@ def staged_pv_mma( producer_label is not None and producer_label.find("tail") >= 0 ) for v_subtile_idx in range(head_dim_stages): - smem_kv.wait() + if not reuse_k: + smem_kv.wait() if producer_label.endswith("_0"): v_desc_0 = getattr(smem_kv, consumer_label)(v_subtile_idx=v_subtile_idx) getattr(tmem_o, producer_label)( @@ -316,6 +320,7 @@ def staged_pv_mma_tmem_p( head_dim_stages, consumer_label, producer_label="pv_mma_loop_tmem_p", + reuse_k=False, ): """Consume TMEM P and staged V descriptors to produce one PV output tile.""" tmem_p.wait() @@ -323,7 +328,8 @@ def staged_pv_mma_tmem_p( tmem_o.acquire() is_tail = producer_label == "pv_mma_tail_tmem_p" for v_subtile_idx in range(head_dim_stages): - smem_kv.wait() + if not reuse_k: + smem_kv.wait() v_desc_0 = getattr(smem_kv, consumer_label)(v_subtile_idx=v_subtile_idx) getattr(tmem_o, producer_label)( p_stage_idx=p_stage_idx, 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 a5e7e810d770..e35517b09e93 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 @@ -45,7 +45,7 @@ def local_kv_tile_idx( return Int32(inst_id) if cutlass.const_expr(section == MlaStage.Loop): base = stage_info.loop_offset * Int32(cfg.num_insts_kv) - if cutlass.const_expr(is_v): + if cutlass.const_expr(is_v or (cfg.serial_swap_reuse and inst_id == 1)): return base + Int32(inst_id) return base + Int32(cfg.num_insts_kv + inst_id) if cutlass.const_expr(cfg.total_kv_tiles <= cfg.num_insts_kv): @@ -109,6 +109,8 @@ def runtime_split_pruning_is_profitable(configured_splits_kv: int) -> bool: def _runtime_configured_local_kv_tiles(cfg: MlaConfig, seq_len_kv): """Return the instruction-aligned local span for configured split capacity.""" + if cutlass.const_expr(cfg.fixed_sparse_kv_tiles > 0): + return Int32(cfg.fixed_sparse_kv_tiles) total_kv_tiles = runtime_total_kv_tiles(cfg, seq_len_kv) num_insts_kv = Int32(cfg.num_insts_kv) tiles_per_group = Int32(cfg.num_ctas_per_seq_kv) * num_insts_kv @@ -134,6 +136,16 @@ def runtime_num_ctas_kv(cfg: MlaConfig, seq_len_kv): @cute.jit def runtime_local_kv_tiles(cfg: MlaConfig, seq_len_kv): """Return the padded local KV tile count assigned to each KV CTA group.""" + if cutlass.const_expr(cfg.fixed_sparse_kv_tiles > 0): + return Int32(cfg.fixed_sparse_kv_tiles) + if cutlass.const_expr( + cfg.sparse_direct and cfg.seq_len_kv <= 256 and cfg.num_ctas_per_seq_kv == 1 + ): + # The short direct schedule executes its planned padded tiles. Every + # consumer must use that same domain: using the smaller live length + # here makes tail PV overwrite O with its final masked tile instead + # of accumulating into the earlier nonempty tile. + return Int32(cfg.local_kv_tiles(cfg.total_kv_tiles)) if cutlass.const_expr(cfg.use_multi_ctas_kv != 1): return runtime_total_kv_tiles(cfg, seq_len_kv) return _runtime_configured_local_kv_tiles(cfg, seq_len_kv) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/kernel_policy.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/kernel_policy.py index c0e29277c516..3949e8ca5eb8 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/kernel_policy.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/kernel_policy.py @@ -161,7 +161,7 @@ def is_throughput_2cta_mla_supported_shape( dtype in ("bf16", "e4m3") and out_dtype in ("bf16", "e4m3") and latent_dim == 512 - and rope_dim == 64 + and rope_dim in (0, 64) and page_size in SUPPORTED_MLA_PAGE_SIZES and 1 <= num_heads <= 128 and seq_len_q >= 1 diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/sparse_policy.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/sparse_policy.py new file mode 100644 index 000000000000..0b838cf4d3a3 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/sparse_policy.py @@ -0,0 +1,507 @@ +# Copyright (c) 2026 by FlashInfer team. +# Licensed under the Apache License, Version 2.0. +"""Sparse MLA policy: choose a schedule, size buffers, then fill the SM grid. + +``rows`` is planned B*Q, ``query_length`` is planned Q, and ``capacity`` counts +selected slots (each source rounded to 128, at least 256 in total), not raw KV +context. Live lengths and indices remain runtime inputs. A query/head grid has +rows*ceil(heads/tile_Q) work tiles; a 2CTA tile consumes two CTAs. + +Selection tries qualified schedules before the small-grid defaults, then +sizes buffers and fills the remaining grid with KV/V partitions. Row-to-SM +ratios estimate available work, not achieved occupancy. Head/capacity/Q and +wave cutoffs are measured GB300 crossovers, not hardware or model limits; +resource and layout eligibility is checked separately by the kernel config. +""" + +from dataclasses import dataclass, replace +from math import ceil + + +@dataclass(frozen=True) +class SparseMlaProfile: + """Kernel schedule; zero tile/split/V-partition values are resolved below.""" + + family: str + tile_size_q: int = 0 + split_kv: int = 1 + scheduler: str = "nonpersistent" + reduction: str = "gmem_separate" + head_dim_ctas: int = 1 + gather_issue_warps: int = 4 + # Separate lists help qualified small decode grids; premerged tagged + # routes remain selected where their simpler loader improves throughput. + direct_inputs: bool = True + balanced_registers: bool = False + reuse_kv: bool = False + reuse_kv_stages: int = 0 + paired_correction: bool = False + single_kv_stream: bool = False + defer_max_update: bool = False + uniform_offset_cache: bool = False + kv_pipeline_stages: int = 0 + kv_tile_size: int = 128 + page_pipeline_stages: int = 0 + + +def _base_fp8(rows, heads, capacity, sm_count): + """Small-grid tile and split selection for FP8.""" + # M64 shares KV across more heads but can leave a small grid underfilled. + # M16 exposes more tasks at the cost of repeating KV work for head tiles. + throughput_h32 = heads == 32 and rows * 2 >= sm_count and (capacity >= 512) + long_list = capacity >= 8192 + m16_work = rows * ((heads + 15) // 16) + m16_one_step_work = m16_work * ((capacity + 255) // 256) + use_keep = heads >= 64 and ( + capacity > 1024 and m16_one_step_work > sm_count or m16_work >= sm_count // 2 + ) + use_keep = use_keep or (long_list and heads >= 16 and (rows >= 8)) + use_keep = use_keep or (capacity <= 256 and heads >= 16) + use_keep = use_keep or throughput_h32 + base_clusters = rows * ((heads + 127) // 128) + keep_work = rows * ((heads + 63) // 64) + kv_tiles = (capacity + 127) // 128 + split_limit = min(kv_tiles, max(1, sm_count // 4 // keep_work)) + kv_steps = (kv_tiles + split_limit - 1) // split_limit + keep_splits = (kv_tiles + kv_steps - 1) // kv_steps + prefer_wide_q = ( + heads >= 64 + and capacity >= 512 + and (kv_steps <= 2) + and (keep_work * keep_splits * 4 >= sm_count // 2) + ) + cluster_splits = 1 << ((capacity + 255) // 256 - 1).bit_length() + m8_work = rows * ((heads + 7) // 8) + one_tile_splits = (capacity + 127) // 128 + if heads == 32 and capacity >= 512 and rows >= sm_count: + # Enough query rows amortize padded M64 work. Retaining two FP8 K + # tiles saves V reloads; uniform coordinates amortize repeated D slices. + return SparseMlaProfile( + family="keep", + direct_inputs=False, + reuse_kv=True, + reuse_kv_stages=10, + uniform_offset_cache=True, + scheduler="clc", + ) + if ( + heads >= 64 + and 512 <= capacity <= 1024 + and (sm_count // 2 <= keep_work * one_tile_splits <= sm_count) + ): + # One K tile per split fills about half to one device wave. Further + # splitting cannot shorten the tile and only increases merge traffic. + return SparseMlaProfile( + family="keep", split_kv=one_tile_splits, reuse_kv=True, reuse_kv_stages=10 + ) + if heads == 32 and 32 <= rows <= sm_count // 2 and (capacity >= 512): + # This intermediate grid favors M64 reuse, with a second split only + # for longer lists; additional splits lose to their merge overhead. + return SparseMlaProfile( + family="keep", + split_kv=1 if capacity <= 1024 else 2, + reuse_kv=True, + reuse_kv_stages=10, + ) + if heads <= 16 and rows * 4 >= sm_count and (capacity >= 512): + # Narrow Q tiles avoid padded heads; twelve stages retain two FP8 + # tiles plus four prefetched D slices (four slices per D512 tile). + return SparseMlaProfile( + family="swap", + tile_size_q=8 if heads <= 8 else 16, + split_kv=2 if capacity > 1024 and rows * 2 <= sm_count else 1, + balanced_registers=True, + reuse_kv=True, + reuse_kv_stages=12, + ) + if capacity <= 256 and heads >= 16 and (m8_work * 4 <= sm_count // 2): + # Short lists offer little KV parallelism. V partitions expose work + # despite repeating QK/softmax, while M8 limits padded arithmetic. + return SparseMlaProfile( + family="swap", tile_size_q=8, head_dim_ctas=4, balanced_registers=True + ) + # Partition V when the wide tile alone cannot fill a resident wave. + if prefer_wide_q: + return SparseMlaProfile(family="keep", split_kv=keep_splits, head_dim_ctas=4) + elif ( + 1 < cluster_splits <= 4 + and heads % 16 == 0 + and (m16_work * cluster_splits <= sm_count // 2) + ): + # A small two/four-split cluster avoids a standalone GMEM finisher. + # Its coordination/residency cost loses for larger grids or split sets. + return SparseMlaProfile( + family="swap", + split_kv=cluster_splits, + reduction="cluster", + head_dim_ctas=0, + balanced_registers=True, + ) + if ( + long_list + and heads >= 128 + and (rows >= 8) + and (base_clusters * 2 <= sm_count // 2) + ): + # Wide heads amortize the 2CTA layout; long lists provide enough KV + # work to split when the query grid alone is small. + return SparseMlaProfile(family="2cta", split_kv=0, gather_issue_warps=1) + elif heads >= 128 and capacity >= 512 and (base_clusters * 2 >= sm_count // 2): + return SparseMlaProfile(family="2cta", gather_issue_warps=1) + return SparseMlaProfile( + family="keep" if use_keep else "swap", + split_kv=1 if throughput_h32 or (use_keep and capacity <= 1024) else 0, + head_dim_ctas=1 if throughput_h32 else 0, + balanced_registers=not use_keep, + reuse_kv=capacity >= 512 and use_keep, + reuse_kv_stages=10 if capacity >= 512 and use_keep else 0, + ) + + +def _base_bf16(rows, heads, capacity, sm_count): + """Small-grid tile and split selection for BF16.""" + # BF16 BK128 cannot retain two complete D512 tiles in SMEM. A narrow + # grid favors more small-Q tasks; a large grid amortizes M64's padding + # and benefits from retaining one tile instead of reloading V. + throughput_h32 = heads == 32 and rows * 2 >= sm_count and capacity >= 512 + throughput_bf16 = ( + heads >= 32 and (rows >= sm_count or throughput_h32) and (capacity >= 512) + ) + short_list = capacity <= 256 and heads <= 64 + tile = (8 if heads <= 16 else 16) if short_list else 16 if heads <= 32 else 64 + if throughput_bf16: + tile = 64 + profile = SparseMlaProfile( + family="swap" if tile < 64 else "keep", + tile_size_q=tile, + split_kv=1 if short_list or throughput_bf16 else 0, + head_dim_ctas=1 if throughput_bf16 else 0, + balanced_registers=tile < 64, + reuse_kv=throughput_bf16, + reuse_kv_stages=4 if throughput_bf16 else 0, + uniform_offset_cache=throughput_bf16 + and rows >= sm_count + and ( + capacity >= 2048 + or (capacity >= 1024 and rows >= 2 * sm_count) + or rows >= 8 * sm_count + ), + ) + cluster_splits = 1 << ((capacity + 255) // 256 - 1).bit_length() + m16_work = rows * ((heads + 15) // 16) + if ( + not short_list + and tile == 16 + and (1 < cluster_splits <= 4) + and (heads % 16 == 0) + and (m16_work * cluster_splits <= sm_count // 2) + ): + # Same small-grid tradeoff as FP8: cluster fusion saves a launch, but + # only the qualified two/four-split M16 layouts justify its overhead. + profile = replace(profile, split_kv=cluster_splits, reduction="cluster") + return profile + + +def _choose_fp8(rows, heads, query_length, capacity, sm_count): + """Prefer resident split grids or sustained-throughput schedules when eligible.""" + # These short-Q regions favor one KV stream with retained tiles and a + # split grid near one device wave. The Q<=8 cutoff is empirical; the + # kernel also supports prefill. Powers of two are policy, not a KV ABI. + if ( + query_length <= 8 + and capacity >= 512 + and ( + heads <= 16 + and 32 <= rows <= sm_count + and (not (heads <= 8 and capacity >= 2048 and (rows * 2 >= sm_count))) + or (heads == 32 and 16 <= rows <= sm_count // 8 and (capacity >= 2048)) + ) + ): + base_ctas = rows * ((heads + 15) // 16) + wave_splits = max(1, sm_count // base_ctas) + wave_splits = 1 << (wave_splits.bit_length() - 1) + return SparseMlaProfile( + family="swap", + split_kv=wave_splits, + balanced_registers=True, + reuse_kv=True, + reuse_kv_stages=10, + single_kv_stream=True, + ) + if ( + heads == 128 + and query_length <= 8 + and (capacity >= 512) + and ((8 if capacity <= 6 * 128 else 16) <= rows <= sm_count // 4) + ): + # Two CTAs cover H128 without repeating KV loads across M64 tiles. + # Up to six KV tiles, avoiding the separate split-KV merge wins from + # eight query rows. Longer lists keep the measured sixteen-row + # crossover; these are GB300 latency crossovers, not kernel limits. + # Four splits help only the small, long-list grid measured here. + small_long_grid = capacity >= 2048 and rows <= sm_count // 8 + return SparseMlaProfile( + family="2cta", split_kv=4 if small_long_grid else 1, gather_issue_warps=1 + ) + if heads == 64 and capacity >= 512 and (rows * 2 >= sm_count): + # Sustained M64 work amortizes eight issuers, expanded coordinate + # registers and CLC setup. Paired correction reduces stats traffic. + return SparseMlaProfile( + family="keep", + gather_issue_warps=8, + direct_inputs=False, + reuse_kv=True, + reuse_kv_stages=10, + uniform_offset_cache=True, + scheduler="clc", + paired_correction=True, + ) + if heads == 128 and capacity >= 256 and (rows * 2 >= sm_count): + # The FP8 2CTA family uses its qualified static persistent scheduler; + # BF16's CLC choice below is a separate measured schedule. Short lists + # (capacity 256-384, e.g. HCA at moderate context) run this schedule + # too: one 2CTA tile covers H128 with a single KV stream, where the + # M64 keep fallback repeats it per head tile (B200: 2.4x faster). + return SparseMlaProfile(family="2cta", scheduler="static", gather_issue_warps=1) + if ( + heads <= 16 + and capacity >= 512 + and ( + capacity <= 1024 + and rows >= (sm_count if heads > 8 else 2 * sm_count) + or rows >= 8 * sm_count + ) + ): + # At high row counts, saved KV traffic outweighs M64 head padding. + # Smaller grids retain M8/M16 to preserve independent work. + return SparseMlaProfile( + family="keep", + direct_inputs=False, + reuse_kv=True, + reuse_kv_stages=10, + uniform_offset_cache=True, + scheduler="clc", + ) + return _base_fp8(rows, heads, capacity, sm_count) + + +def _choose_bf16(rows, heads, query_length, capacity, sm_count): + """Prefer resident split grids or sustained-throughput schedules when eligible.""" + # Small-head workloads can retain one BK128 tile with five/six 32-KiB + # stages: four hold D512, the rest prefetch. PV must release that tile + # before next QK; this favors a single stream in the measured regions. + if ( + heads <= 16 + and capacity >= 512 + and (rows * 2 >= sm_count) + and ( + query_length <= 8 + or ( + heads <= 8 + and capacity <= 1024 + and (query_length >= 128) + and (rows >= 8 * sm_count) + ) + ) + ): + return SparseMlaProfile( + family="swap", + tile_size_q=8 if heads <= 8 else 16, + balanced_registers=True, + reuse_kv=True, + reuse_kv_stages=6 + if heads > 8 or (query_length <= 8 and rows >= 4 * sm_count) + else 5, + single_kv_stream=True, + defer_max_update=True, + scheduler="clc", + ) + if ( + heads in (8, 16, 32, 64) + and capacity >= 512 + and (query_length >= 128) + and (rows >= 8 * sm_count) + ): + # Long-Q grids have ample tasks. Static persistence amortizes CTA + # setup here; wider heads favor M64 reuse, narrow heads avoid padding. + keep = heads >= 32 + return SparseMlaProfile( + family="keep" if keep else "swap", + tile_size_q=64 if keep else 8 if heads <= 8 else 16, + direct_inputs=False, + balanced_registers=not keep, + reuse_kv=True, + reuse_kv_stages=5 if keep else 6, + uniform_offset_cache=keep, + scheduler="static", + ) + if ( + heads in (32, 64) + and rows * 2 >= sm_count + and (query_length <= 8) + and (capacity >= 512) + ): + # Decode with enough rows favors M64 reuse. Buffer sizing below + # supplies its qualified gather/cache budget and scheduler crossover. + return SparseMlaProfile( + family="keep", + reuse_kv=True, + reuse_kv_stages=5, + uniform_offset_cache=capacity >= 2048 + and rows >= sm_count + or (capacity >= 1024 and rows >= 4 * sm_count) + or rows >= 8 * sm_count, + ) + if heads <= 16 and rows * 2 >= sm_count and (capacity >= 512): + return SparseMlaProfile( + family="swap", + tile_size_q=8 if heads <= 8 else 16, + balanced_registers=True, + reuse_kv=True, + reuse_kv_stages=6, + ) + if heads == 128 and 2 <= sm_count // (2 * rows) <= 4 and (capacity >= 512): + # Two M64 head tiles with BK64 reuse and a few KV splits fill this + # small H128 grid. Larger grids favor the wide 2CTA path below. + wave_splits = max(1, sm_count // 2 // rows) + return SparseMlaProfile( + family="keep", + split_kv=wave_splits, + gather_issue_warps=8, + uniform_offset_cache=True, + reuse_kv=True, + reuse_kv_stages=10, + kv_tile_size=64, + page_pipeline_stages=2, + defer_max_update=True, + ) + if ( + capacity >= 512 + and rows * 4 >= sm_count + and (heads == 128 or (heads in (32, 64) and rows * 2 <= sm_count)) + ): + # BF16 2CTA amortizes wide-head work once enough clusters are present. + # More K/V stages overlap gathers; coordinate caching costs registers, + # so it is enabled only for the qualified H128 configuration. + return SparseMlaProfile( + family="2cta", + kv_pipeline_stages=8, + uniform_offset_cache=heads == 128, + scheduler="clc" if heads == 128 and rows >= sm_count else "nonpersistent", + direct_inputs=False, + ) + return _base_bf16(rows, heads, capacity, sm_count) + + +def _size_buffers(profile, *, rows, heads, query_length, capacity, dtype, sm_count): + # These are coupled measured budgets, not independent universal wins: + # issuer/cache registers compete with softmax/correction, and extra KV + # stages compete with Q, P, metadata and scheduler SMEM. + if ( + dtype == "fp8" + and heads == 128 + and (capacity >= 2048) + and (query_length <= 8) + and (rows >= 4 * sm_count) + ): + profile = replace(profile, balanced_registers=True) + if ( + dtype == "bf16" + and profile.family == "keep" + and (profile.reuse_kv_stages == 5) + and (heads in (32, 64)) + ): + # Both prefill and decode wide-BF16 reuse schedules share this budget. + profile = replace( + profile, + gather_issue_warps=8, + uniform_offset_cache=True, + defer_max_update=True, + ) + if query_length <= 8 and rows >= 2 * sm_count: + profile = replace(profile, scheduler="clc", direct_inputs=True) + if ( + heads == 32 + and query_length <= 8 + and (rows >= 4 * sm_count) + and (capacity >= 1024) + ): + # BK64 halves each BF16 stage to 16 KiB: ten stages retain two + # four-slice K tiles plus two slices ahead. This restores QK/PV + # overlap, at the cost of twice as many token-tile iterations. + profile = replace( + profile, kv_tile_size=64, reuse_kv_stages=10, page_pipeline_stages=2 + ) + if ( + dtype == "bf16" + and profile.family == "keep" + and (heads == 64) + and (rows >= sm_count) + and (capacity >= 512) + and (profile.split_kv == 1) + ): + # The qualified H64 BK64 schedule serves both prefill and decode. + # Two metadata stages let preparation advance while issuers reuse + # current coordinates. The bounded maximum update is BF16-only. + profile = replace( + profile, + kv_tile_size=64, + reuse_kv=True, + reuse_kv_stages=10, + page_pipeline_stages=2, + gather_issue_warps=8, + uniform_offset_cache=True, + direct_inputs=True, + scheduler="clc", + defer_max_update=True, + ) + return profile + + +def select_sparse_mla_profile(*, rows, heads, query_length, capacity, dtype, sm_count): + """Resolve one complete schedule from the validated public plan bounds.""" + choose = _choose_fp8 if dtype == "fp8" else _choose_bf16 + profile = choose(rows, heads, query_length, capacity, sm_count) + profile = _size_buffers( + profile, + rows=rows, + heads=heads, + query_length=query_length, + capacity=capacity, + dtype=dtype, + sm_count=sm_count, + ) + tile = profile.tile_size_q or {"swap": 16, "keep": 64, "2cta": 128}[profile.family] + work = rows * ceil(heads / tile) + splits = profile.split_kv + if splits == 0: + # Fill a resident wave using instruction-aligned KV partitions. + # More splits shorten each CTA's work but increase partial-O/LSE merge + # traffic. Two-CTA work targets half as many cluster tasks as SMs. + step = 128 if tile >= 64 else 256 + steps = max(1, ceil(capacity / step)) + target = sm_count // 2 if profile.family == "2cta" else sm_count + partitions = min(steps, max(1, target // work), 128) + splits = ceil(steps / ceil(steps / partitions)) + v_ctas = profile.head_dim_ctas + if v_ctas == 0: + # V partitioning exposes more tasks but repeats QK/softmax across them. + # Prefer it only while the query/KV grid still lacks enough work. + remaining = work * splits + v_ctas = ( + 1 + if profile.family == "2cta" or remaining * 2 > sm_count + else 4 + if remaining * 4 <= sm_count + else 2 + ) + # Retaining K through PV requires a single owner of the V dimension. + reuse = profile.reuse_kv and v_ctas == 1 + return replace( + profile, + tile_size_q=tile, + split_kv=splits, + head_dim_ctas=v_ctas, + reuse_kv=reuse, + reuse_kv_stages=profile.reuse_kv_stages if reuse else 0, + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/sparse_views.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/sparse_views.py new file mode 100644 index 000000000000..ffc74f6d74af --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/sparse_views.py @@ -0,0 +1,314 @@ +# 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. + +"""CTA-local views of caller-prepared storage-row indices. + +Adapt separate primary/extra physical-row lists to the existing route/length +interface without materializing a merged GPU array. Source spans are padded +to 128 and tagged per tile; arbitrary holes remain masked. Nonpersistent CTAs +bind one request, while persistent work tiles rebind live pointers/lengths. +The common premerged-route path bypasses these views. Neither source is +required to be SWA or compressed, and logical-page conversion is external. +""" + +import cutlass +import cutlass.cute as cute +from cutlass import Int32, Int64 + + +class SparseRouteView: + def __init__(self, values, capacity, assume_valid_prefix=False): + self.values = values + self.capacity = capacity + self.assume_valid_prefix = assume_valid_prefix + self.si, self.ci, self.sl, self.cl, self.row = values + self.primary_span = (self.sl + Int32(127)) // Int32(128) * Int32(128) + self.shape = (capacity, self.si.shape[0]) + + def __extract_mlir_values__(self): + return cutlass.extract_mlir_values(self.values) + + def __new_from_mlir_values__(self, values): + return SparseRouteView( + cutlass.new_from_mlir_values(self.values, values), + self.capacity, + self.assume_valid_prefix, + ) + + @cute.jit + def __getitem__(self, coordinate): + token, _ = coordinate + result = Int32(0x7FFFFFFF) + if token < self.sl: + index = Int32(self.si[self.row, token]) + if index >= 0: + result = index + elif token >= self.primary_span: + result = Int32(-1) + position = token - self.primary_span + if position < self.cl: + index = Int32(self.ci[self.row, position]) + if index >= 0: + result = index | Int32(-2147483648) + return result + + @cute.jit + def routes_for_tile(self, token, request=None): + compressed = token >= self.primary_span + primary_ptr = self.si.iterator.toint() + self.row * Int64( + self.si.stride[0] + ) * Int64(4) + extra_ptr = self.ci.iterator.toint() + self.row * Int64( + self.ci.stride[0] + ) * Int64(4) + address = extra_ptr if compressed else primary_ptr + origin = self.primary_span if compressed else Int32(0) + length = self.cl if compressed else self.sl + tag = Int32(-2147483648) if compressed else Int32(0) + indices = cute.make_tensor( + cute.make_ptr(Int32, address, assumed_align=4), + cute.make_layout((self.capacity,), stride=(1,)), + ) + return SparseTileRoutes((indices, origin, length, tag)) + + @cute.jit + def mask_for_tile(self, token, request=None): + # Source boundaries are 128-row aligned. Select the pointer and prefix + # once per score tile, outside the unrolled register mask loop. + compressed = token >= self.primary_span + if cutlass.const_expr(self.assume_valid_prefix): + origin = self.primary_span if compressed else Int32(0) + length = self.cl if compressed else self.sl + return SparsePrefixMask((origin, length)) + routes = self.routes_for_tile(token) + return routes + + @cute.jit + def is_invalid(self, token, request): + return self.mask_for_tile(token).is_invalid(token, request) + + +class SparseLengthView: + def __init__(self, values): + self.values = values + self.length, rows = values + self.shape = (rows,) + + def __extract_mlir_values__(self): + return cutlass.extract_mlir_values(self.values) + + def __new_from_mlir_values__(self, values): + return SparseLengthView(cutlass.new_from_mlir_values(self.values, values)) + + @cute.jit + def __getitem__(self, request): + del request + return self.length + + +class SparsePrefixMask: + """Analytic validity for one source's compact, hole-free prefix.""" + + def __init__(self, values): + self.values = values + self.origin, self.length = values + + def __extract_mlir_values__(self): + return cutlass.extract_mlir_values(self.values) + + def __new_from_mlir_values__(self, values): + return SparsePrefixMask(cutlass.new_from_mlir_values(self.values, values)) + + @cute.jit + def tile_is_full(self, start, live_length): + return start + Int32(128) <= cute.math.min( + self.origin + self.length, live_length + ) + + @cute.jit + def prefix_mask32(self, start, live_length): + # Tiles and source boundaries are 128-aligned; start never precedes + # this source's origin. Include the request's live bound as well. + end = cute.math.min(self.origin + self.length, live_length) + valid = cute.math.min(cute.math.max(end - start, Int32(0)), Int32(32)) + # Mask the shift count even on the unselected side of the expression. + # A full 32-bit valid prefix must not execute an undefined shift by 32. + bits = Int32(-1) << (valid & Int32(31)) + return bits if valid < Int32(32) else Int32(0) + + @cute.jit + def is_invalid(self, token, request): + del request + return (token < self.origin) | (token >= self.origin + self.length) + + +class SparseTileRoutes: + """One source selection per tile, with bounded scalar and quad loads.""" + + def __init__(self, values): + self.values = values + self.indices, self.origin, self.length, self.tag = values + + def __extract_mlir_values__(self): + return cutlass.extract_mlir_values(self.values) + + def __new_from_mlir_values__(self, values): + return SparseTileRoutes(cutlass.new_from_mlir_values(self.values, values)) + + @cute.jit + def _bounded_index(self, token): + position = token - self.origin + # Zero-capacity sources bind a one-row dummy. Clamp even masked + # positions so both gather and softmax can use a bounded load. + safe = cute.math.min( + cute.math.max(position, Int32(0)), + cute.math.max(self.length - Int32(1), Int32(0)), + ) + index = Int32(self.indices[safe]) + valid = (position >= Int32(0)) & (position < self.length) & (index >= Int32(0)) + return index, valid + + @cute.jit + def is_invalid(self, token, request): + del request + _, valid = self._bounded_index(token) + return ~valid + + @cute.jit + def mapped_row(self, token): + index, valid = self._bounded_index(token) + return (index if valid else Int32(0x7FFFFFFF)) | self.tag + + @cute.jit + def mapped_quad(self, token): + """Load a complete active quad together and preserve partial-prefix masking.""" + rows = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) + for j in cutlass.range_constexpr(4): + rows[j] = Int32(0x7FFFFFFF) | self.tag + position = token - self.origin + if position >= Int32(0) and position + Int32(3) < self.length: + values = cutlass.Pointer( + self.indices.iterator + position, dtype=Int32 + ).load(count=4, alignment=4) + for j in cutlass.range_constexpr(4): + index = Int32(values[j]) + rows[j] = (index if index >= Int32(0) else Int32(0x7FFFFFFF)) | self.tag + elif position + Int32(3) >= Int32(0) and position < self.length: + for j in cutlass.range_constexpr(4): + rows[j] = self.mapped_row(token + Int32(j)) + return rows + + +class SparseBatchRouteView: + """Bind live source metadata to the request currently owned by a work tile.""" + + def __init__(self, values, capacity, assume_valid_prefix=False): + self.values = values + self.si, self.ci, self.sl, self.cl = values + self.capacity = capacity + self.assume_valid_prefix = assume_valid_prefix + self.shape = (capacity, self.si.shape[0]) + + def __extract_mlir_values__(self): + return cutlass.extract_mlir_values(self.values) + + def __new_from_mlir_values__(self, values): + return SparseBatchRouteView( + cutlass.new_from_mlir_values(self.values, values), + self.capacity, + self.assume_valid_prefix, + ) + + @cute.jit + def for_request(self, request, primary_length=None, extra_length=None): + # Persistent consumers may reuse lengths loaded at query-tile entry. + row = Int64(request) + return SparseRouteView( + ( + self.si, + self.ci, + Int32(self.sl[row]) + if cutlass.const_expr(primary_length is None) + else primary_length, + Int32(self.cl[row]) + if cutlass.const_expr(extra_length is None) + else extra_length, + row, + ), + self.capacity, + self.assume_valid_prefix, + ) + + @cute.jit + def __getitem__(self, coordinate): + return self.for_request(coordinate[1])[coordinate] + + @cute.jit + def routes_for_tile(self, token, request): + return self.for_request(request).routes_for_tile(token) + + @cute.jit + def mask_for_tile(self, token, request): + return self.for_request(request).mask_for_tile(token) + + @cute.jit + def is_invalid(self, token, request): + return self.mask_for_tile(token, request).is_invalid(token, request) + + +class SparseBatchLengthView: + """Live padded source extent for each virtual request.""" + + def __init__(self, values): + self.values = values + self.sl, self.cl = values + self.shape = self.sl.shape + + def __extract_mlir_values__(self): + return cutlass.extract_mlir_values(self.values) + + def __new_from_mlir_values__(self, values): + return SparseBatchLengthView(cutlass.new_from_mlir_values(self.values, values)) + + @cute.jit + def __getitem__(self, request): + return _sparse_execution_length( + Int32(self.sl[request]), Int32(self.cl[request]) + ) + + +@cute.jit +def _sparse_execution_length(primary_length, extra_length): + primary_tiles = (primary_length + Int32(127)) // Int32(128) + extra_tiles = (extra_length + Int32(127)) // Int32(128) + # Empty requests still execute a masked tile to publish O=0 and LSE=-inf. + return cute.math.max((primary_tiles + extra_tiles) * Int32(128), Int32(1)) + + +@cute.jit +def bind_sparse_views( + si, + ci, + sl, + cl, + capacity: cutlass.Constexpr[int], + *, + request=None, + assume_valid_prefix: cutlass.Constexpr[bool] = False, +): + """Bind source lists to a work queue or to one nonpersistent request.""" + routes = SparseBatchRouteView((si, ci, sl, cl), capacity, assume_valid_prefix) + if cutlass.const_expr(request is None): + return routes, SparseBatchLengthView((sl, cl)) + else: + current = routes.for_request(request) + length = _sparse_execution_length(current.sl, current.cl) + return current, SparseLengthView((length, si.shape[0])) 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 af89becc56eb..fe2461ad74c4 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 @@ -25,9 +25,6 @@ from ..helpers.mask import MaskType, normalize_mask_type -# Softmax converts natural-scale scores to exp2 with log2(e). -LOG2_E = 1.4426950408889634074 - # The separate reducer follows the public MLA output contract: partial O is # stored as BF16 while LSE and the final accumulation remain FP32. One # 512-thread CTA owns eight D512 rows, with each thread moving one 16-byte @@ -49,6 +46,10 @@ # Each transpose-TMA issue stages one 64-element slice of the latent dimension. V_TMA_LATENT_ELEMENTS = 64 +# Two N slices of the final M128xD512 FP8 output, split across two CTAs. +OUTPUT_TMA_STAGES = 2 +OUTPUT_TMA_STAGE_BYTES = 64 * 256 + def ceil_div(a: int, b: int) -> int: """Return ``ceil(a / b)`` for positive integer divisors.""" @@ -159,8 +160,12 @@ class MlaDecodeConfig: iterations_qk_stages: int = 5 iterations_pv_stages: int = 4 + # Final quantized output can use a staged TMA store on padding warp 10. + use_tma_output: bool = False + # Pipeline stage counts for the captured schedule resources. The combined # K/V stage count keeps enough delayed-V stages live for K-before-V overlap. + sparse_offset_stages: int = 4 load_q_stage: int = 1 load_k_stage: int = 3 load_v_stage: int = 2 @@ -179,6 +184,9 @@ class MlaDecodeConfig: correction_warp_ids: Tuple[int, ...] = (4, 5, 6, 7) mma_warp_id: int = 8 load_tma_warp_id: int = 9 + load_k_warp_id: int = 9 + load_v_warp_id: int = 10 + load_num_warps: int = 1 pv_mma_warp_id: int = 11 empty_warp_ids: Tuple[int, ...] = (11,) second_compute_warp_ids: Tuple[int, ...] = () @@ -258,6 +266,7 @@ class MlaDecodeConfig: # Use block_split_kvs[batch] as a per-batch cap before runtime K contracts # the useful split prefix. Grid and workspace geometry retain the maximum. is_var_split_kv: bool = False + cache_uniform_sparse_pages: bool = False # 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 @@ -334,6 +343,11 @@ def make_mla_decode_config( is_var_seq: bool = False, is_var_split_kv: bool = False, mask_type: MaskType | str = MaskType.CAUSAL, + sparse_gather_warps: int = 1, + sparse_kv_stages: int = 0, + cache_uniform_sparse_pages: bool = False, + balance_sparse_registers: bool = False, + fuse_output_quant: bool = False, ) -> MlaDecodeConfig: """Create and populate a MlaDecodeConfig from problem parameters.""" cfg = MlaDecodeConfig() @@ -373,7 +387,9 @@ def _require_divisible( if rope_dim < 0: raise ValueError(f"rope_dim must be non-negative, got {rope_dim}") _require_positive("page_size", page_size) - if page_size not in SUPPORTED_MLA_PAGE_SIZES: + if page_size not in SUPPORTED_MLA_PAGE_SIZES and not ( + page_size == 1 and rope_dim == 0 + ): raise ValueError( f"page_size must be one of {SUPPORTED_MLA_PAGE_SIZES}, got {page_size}" ) @@ -403,6 +419,81 @@ def _require_divisible( cfg.other_reg_num = 32 cfg.mma_o_stage = 2 cfg.epilogue_sync_bar_id = 4 + if page_size == 1: + # Independent QK/PV streams retain two whole K/V tiles each. A + # third V tile costs 32 KiB per CTA without improving overlap in + # the qualified sparse schedule. Keep four metadata slots so + # index prefetch can run ahead of both streams. + cfg.load_k_stage = 2 + cfg.load_v_stage = 2 + cfg.use_fp8_dual_softmax_schedule = False + cfg.second_compute_warp_ids = () + cfg.num_softmax_groups = 1 + cfg.load_k_warp_id = 12 + cfg.load_v_warp_id = 16 + cfg.load_num_warps = 4 + cfg.threads_per_cta = cfg.threads_per_warp * 20 + cfg.other_reg_num = 48 + + if sparse_gather_warps != 1: + if sparse_gather_warps not in (4, 8) or qkv_dtype != "bf16" or page_size != 1: + raise ValueError("multi-warp combined loading requires BF16 sparse 2CTA") + cfg.load_num_warps = sparse_gather_warps + cfg.load_tma_warp_id = 12 + cfg.load_k_warp_id = 12 + cfg.threads_per_cta = 384 + 32 * sparse_gather_warps + cfg.softmax_reg_num = 128 if sparse_gather_warps == 8 else 192 + cfg.correction_reg_num = 160 if sparse_gather_warps == 8 else 192 + cfg.other_reg_num = 64 + + if sparse_kv_stages: + if qkv_dtype != "bf16" or page_size != 1 or sparse_gather_warps not in (4, 8): + raise ValueError("pipeline tuning requires multi-warp BF16 sparse 2CTA") + if sparse_kv_stages < 4: + raise ValueError("combined KV pipeline requires at least four stages") + cfg.load_kv_stage = sparse_kv_stages + + if cache_uniform_sparse_pages: + if qkv_dtype != "bf16" or page_size != 1 or sparse_gather_warps not in (4, 8): + raise ValueError("uniform K cache requires multi-warp BF16 sparse loading") + cfg.cache_uniform_sparse_pages = True + cfg.softmax_reg_num = 128 + cfg.correction_reg_num = 160 if sparse_gather_warps == 8 else 192 + cfg.other_reg_num = 64 if sparse_gather_warps == 8 else 96 + + if balance_sparse_registers: + if page_size != 1 or qkv_dtype != "e4m3": + raise ValueError("balanced registers require sparse FP8") + # Sparse FP8 keeps 64 registers for the gather issuers. Split the + # remaining budget equally: 128 softmax registers spilled live loop + # state on every KV tile, while correction fits in 144 registers. + # 4*144 + 4*144 + 12*64 = 1920 warp-registers, the same CTA total + # as the previous 128/160/64 allocation (61440 registers). + cfg.softmax_reg_num = 144 + cfg.correction_reg_num = 144 + cfg.other_reg_num = 64 + + if fuse_output_quant and cfg.is_fp8_qkv() and cfg.threads_per_cta == 640: + # Quantization keeps an output block and rotary coefficients live. + # The gather tasks fit in 48 registers; use their surplus for the + # correction task while preserving the CTA's 61,440-register budget. + cfg.softmax_reg_num = 160 + cfg.correction_reg_num = 176 + cfg.other_reg_num = 48 + + # SETMAXNREG redistributes the CTA's initial allocation, rounded down + # to eight registers/thread; the remainder of the SM file is not credit. + initial_regs = 65536 // (cfg.threads_per_cta * 8) * 8 + smx_warps = cfg.num_softmax_groups * cfg.num_compute_warps + cor_warps = cfg.num_compute_warps + other_warps = cfg.threads_per_cta // 32 - smx_warps - cor_warps + claimed = ( + smx_warps * cfg.softmax_reg_num + + cor_warps * cfg.correction_reg_num + + other_warps * cfg.other_reg_num + ) * 32 + if claimed > initial_regs * cfg.threads_per_cta: + raise ValueError("2CTA task registers exceed the CTA's initial allocation") # Derived MMA tilers. FP8 latent QK uses K=128 while the separate RoPE MMA # keeps K=64. 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 5ea23ffb4238..70fa06200b13 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 @@ -32,7 +32,7 @@ import cutlass.utils as utils import cutlass.utils.blackwell_helpers as sm100_utils import cutlass.experimental.cuda as cuda -from cutlass import Int32, Int64 +from cutlass import Float32, Int32, Int64 from cutlass.cute.testing import assert_ as runtime_assert from cutlass.cute.nvgpu import OperandMajorMode, tcgen05 from cutlass.experimental import primitives as prims @@ -50,14 +50,15 @@ ) from .config import ( - LOG2_E, + OUTPUT_TMA_STAGES, + OUTPUT_TMA_STAGE_BYTES, REDUCTION_ROWS_PER_CTA, REDUCTION_THREADS_PER_ROW, V_TMA_LATENT_ELEMENTS, MlaDecodeConfig, make_mla_decode_config, ) -from ..helpers.constants import MAX_MLA_SPLITS_KV, TMEM_DEALLOC_MBAR_THREADS +from ..helpers.constants import LOG2_E, MAX_MLA_SPLITS_KV, TMEM_DEALLOC_MBAR_THREADS from ..helpers.mask import MaskType, mask_visible_k_length, normalize_mask_type from ..helpers.query import ( FlatQueryTileLayout, @@ -86,6 +87,7 @@ from ..helpers.math import qkv_dtype from ..parallel_reduction_topology import ( ParallelReductionTopology, + make_balanced_parallel_reduction_topology, make_q128_wave_limited_parallel_reduction_topology, should_use_q128_g1_parallel_reducer, validate_parallel_reduction_workspace, @@ -107,6 +109,7 @@ MlaInterleavedTask, MlaTask, create_load_k_task, + create_sparse_offsets_task, create_load_v_task, create_load_tma_task, create_mma_task, @@ -115,6 +118,7 @@ create_softmax_task, create_correction_task, create_padding_task, + create_output_tma_task, create_scheduler_task, ) @@ -132,6 +136,7 @@ def build_mla_decode_task_manager( smem_kc_arr=None, smem_vc_arr=None, smem_p_arr=None, + sparse_offsets_smem=None, # TMA descriptors (None for validate-only) tma_desc_q_latent=None, tma_desc_q_rope=None, @@ -145,6 +150,12 @@ def build_mla_decode_task_manager( tidx=None, # GMEM output tensors output=None, + atten_sinks=None, + lse_in_natural_log=False, + output_quant=None, + output_quant_args=None, + smem_o=None, + tma_o_desc=None, acc_output=None, lse=None, acc_lse=None, @@ -200,7 +211,7 @@ def build_mla_decode_task_manager( ) # Cooperative groups - tma_producer_group = pipeline.CooperativeGroup(Agent.Thread) # elect_one TMA + tma_producer_group = pipeline.CooperativeGroup(Agent.Thread) umma_hw_group = pipeline.CooperativeGroup(Agent.Thread) # UMMA hardware # Softmax warps: 4 warps * 32 threads * 2 CTAs for cluster-scoped consumer @@ -227,10 +238,15 @@ def build_mla_decode_task_manager( num_stages=cfg.load_q_stage, num_bytes=cfg.tma_copy_q_bytes, producer_group=tma_producer_group, + producer_signaling_threads=SignalingThreads.TaskWarpLeader + if cfg.load_num_warps > 1 + else SignalingThreads.All, consumer_group=umma_hw_group, cta_layout_vmnk=cluster_shape_vmnk, consumer_signaling_threads=SignalingThreads.CtaLeader, - num_bytes_per_warp_per_cta=(cfg.tma_copy_q_bytes // cfg.num_mma_ctas), + num_bytes_per_warp_per_cta=( + cfg.tma_copy_q_bytes // cfg.num_mma_ctas // cfg.load_num_warps + ), ) if cfg.use_fp8_split_mma_schedule: @@ -241,19 +257,29 @@ def build_mla_decode_task_manager( num_stages=cfg.load_k_stage, num_bytes=cfg.tma_copy_k_tile_bytes, producer_group=tma_producer_group, + producer_signaling_threads=SignalingThreads.TaskWarpLeader + if cfg.load_num_warps > 1 + else SignalingThreads.All, consumer_group=umma_hw_group, cta_layout_vmnk=cluster_shape_vmnk, consumer_signaling_threads=SignalingThreads.CtaLeader, - num_bytes_per_warp_per_cta=(cfg.tma_copy_k_tile_bytes // cfg.num_mma_ctas), + num_bytes_per_warp_per_cta=( + cfg.tma_copy_k_tile_bytes // cfg.num_mma_ctas // cfg.load_num_warps + ), ) smem_v_pipeline_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( num_stages=cfg.load_v_stage, num_bytes=cfg.tma_copy_v_tile_bytes, producer_group=tma_producer_group, + producer_signaling_threads=SignalingThreads.TaskWarpLeader + if cfg.load_num_warps > 1 + else SignalingThreads.All, consumer_group=umma_hw_group, cta_layout_vmnk=cluster_shape_vmnk, consumer_signaling_threads=SignalingThreads.CtaLeader, - num_bytes_per_warp_per_cta=(cfg.tma_copy_v_tile_bytes // cfg.num_mma_ctas), + num_bytes_per_warp_per_cta=( + cfg.tma_copy_v_tile_bytes // cfg.num_mma_ctas // cfg.load_num_warps + ), ) smem_kv_pipeline_cfg = None else: @@ -265,7 +291,12 @@ def build_mla_decode_task_manager( consumer_group=umma_hw_group, cta_layout_vmnk=cluster_shape_vmnk, consumer_signaling_threads=SignalingThreads.CtaLeader, - num_bytes_per_warp_per_cta=(cfg.tma_kc_subtile_bytes // cfg.num_mma_ctas), + num_bytes_per_warp_per_cta=( + cfg.tma_kc_subtile_bytes // cfg.num_mma_ctas // cfg.load_num_warps + ), + producer_signaling_threads=SignalingThreads.TaskWarpLeader + if cfg.load_num_warps > 1 + else SignalingThreads.All, ) smem_k_pipeline_cfg = None smem_v_pipeline_cfg = None @@ -329,10 +360,31 @@ def build_mla_decode_task_manager( # Create resource instances # ────────────────────────────────────────────────────────────── + sparse_offsets_enabled = cfg.page_size == 1 and cfg.use_fp8_split_mma_schedule + shared_mask_enabled = ( + sparse_offsets_enabled + and not cfg.use_fp8_dual_softmax_schedule + and not getattr(page_offsets, "assume_valid_prefix", False) + ) + # Prefix masks use live lengths only: softmax no longer consumes the + # offset ring. Remove its arrivals along with its task dependency; K/V + # loaders remain consumers with the same addressing and synchronization. + offsets_cfg = None + if sparse_offsets_enabled: + offsets_cfg = PipelineConfig.create_async_async_pipeline_cfg( + num_stages=cfg.sparse_offset_stages, + producer_group=pipeline.CooperativeGroup(Agent.Thread, 32), + consumer_group=pipeline.CooperativeGroup( + Agent.Thread, + (2 * cfg.load_num_warps + (4 if shared_mask_enabled else 0)) * 32, + ), + cta_layout_vmnk=cluster_shape_vmnk, + ) page_offset_window = PageOffsetWindowResource( page_offsets=page_offsets, + smem_sparse_offsets=sparse_offsets_smem, cfg=cfg, - pipeline_config=None, + pipeline_config=offsets_cfg, name="page_offset_window", ) @@ -351,6 +403,7 @@ def build_mla_decode_task_manager( if cfg.use_fp8_split_mma_schedule: smem_k = SmemKResource( + smem_sparse_offsets=sparse_offsets_smem, smem_k=smem_kc_arr, page_offsets=page_offsets, tma_desc_c_latent=tma_desc_c_latent, @@ -361,9 +414,11 @@ def build_mla_decode_task_manager( name="smem_k", ) smem_v = SmemVResource( + smem_sparse_offsets=sparse_offsets_smem, smem_v=smem_vc_arr, page_offsets=page_offsets, tma_desc_c_transpose=tma_desc_c_transpose, + tma_desc_extra_v=tma_desc_q_rope, logical_seq_len_q=logical_seq_len_q, cfg=cfg, pipeline_config=smem_v_pipeline_cfg, @@ -376,6 +431,7 @@ def build_mla_decode_task_manager( tma_desc_c_latent=tma_desc_c_latent, tma_desc_c_rope=tma_desc_c_rope, tma_desc_c_transpose=tma_desc_c_transpose, + tma_desc_extra_v=tma_desc_q_rope, cfg=cfg, pipeline_config=smem_kv_pipeline_cfg, name="smem_kv", @@ -384,6 +440,7 @@ def build_mla_decode_task_manager( smem_v = None tmem_s = TmemSResource( + page_offsets=page_offsets, smem_q_latent=smem_q_latent_arr, smem_q_rope=smem_q_rope_arr, smem_p=smem_p_arr, @@ -420,6 +477,14 @@ def build_mla_decode_task_manager( name="tmem_o", ) + output_pipeline = None + if cfg.use_tma_output: + output_pipeline = PipelineConfig.create_async_async_pipeline_cfg( + num_stages=OUTPUT_TMA_STAGES, + producer_group=compute_group_local, + consumer_group=pipeline.CooperativeGroup(Agent.Thread, 32), + cta_layout_vmnk=cluster_shape_vmnk, + ) gmem_o = GmemOResource( cfg=cfg, output=output, @@ -428,6 +493,12 @@ def build_mla_decode_task_manager( partial_lse=acc_lse, tmem_o_ref=tmem_o, tmem_corr_ref=tmem_corr, + atten_sinks=atten_sinks, + lse_in_natural_log=lse_in_natural_log, + output_quant=output_quant, + output_quant_args=output_quant_args, + smem_o=smem_o, + tma_o_desc=tma_o_desc, output_scale=None, # set at runtime softmax_scale_log2=None, # set at runtime smem_exchange=None, # set at runtime @@ -436,6 +507,7 @@ def build_mla_decode_task_manager( logical_num_heads_q=logical_num_heads_q, logical_seq_len_q=logical_seq_len_q, name="gmem_o", + pipeline_config=output_pipeline, ) # ────────────────────────────────────────────────────────────── @@ -451,12 +523,14 @@ def build_mla_decode_task_manager( load_k_task = create_load_k_task( smem_q, smem_k, + page_offsets=page_offset_window if sparse_offsets_enabled else None, work_queue=work_queue, domain=domain, num_registers=wg2_reg_count, ) load_v_task = create_load_v_task( smem_v, + page_offsets=page_offset_window if sparse_offsets_enabled else None, work_queue=work_queue, domain=domain, num_registers=wg2_reg_count, @@ -521,6 +595,7 @@ def build_mla_decode_task_manager( tmem_s, tmem_corr, smem_p, + page_offsets=page_offset_window if shared_mask_enabled else None, work_queue=work_queue, task_class=( MlaInterleavedTask @@ -580,13 +655,41 @@ def build_mla_decode_task_manager( ] task_list.extend([mma_task, softmax_task, correction_task]) + if sparse_offsets_enabled: + task_list.append( + create_sparse_offsets_task( + page_offset_window, + work_queue=work_queue, + domain=domain, + num_registers=wg2_reg_count, + ) + ) + if cfg.use_tma_output: + task_list.append( + create_output_tma_task( + gmem_o, + work_queue=work_queue, + num_registers=wg2_reg_count, + ) + ) + elif cfg.use_fp8_split_mma_schedule and cfg.load_num_warps > 1: + task_list.append( + create_padding_task( + work_queue=work_queue, + domain=domain, + num_registers=wg2_reg_count, + warp_idx=10 if sparse_offsets_enabled else 9, + num_warps=1 if sparse_offsets_enabled else 2, + ) + ) if not cfg.use_fp8_split_mma_schedule and use_clc_dynamic: padding_task = create_padding_task( work_queue=work_queue, task_class=non_interleaved_task_class, domain=task_domain, num_registers=wg2_reg_count, - warp_idx=10, + warp_idx=9 if cfg.load_num_warps > 1 else 10, + num_warps=2 if cfg.load_num_warps > 1 else 1, ) task_list.append(padding_task) scheduler_task = create_scheduler_task( @@ -604,8 +707,8 @@ def build_mla_decode_task_manager( work_queue=work_queue, domain=domain, num_registers=wg2_reg_count, - warp_idx=10, - num_warps=2, + warp_idx=9 if cfg.load_num_warps > 1 else 10, + num_warps=3 if cfg.load_num_warps > 1 else 2, ) task_list.append(padding_task) @@ -664,6 +767,12 @@ def build_mla_decode_task_manager( else: resource_dependency_graph[work_queue] = [] + if sparse_offsets_enabled: + # K/V TMA reads register snapshots; metadata has no asynchronous SMEM read edge. + resource_dependency_graph[page_offset_window] = ( + [work_queue] if work_queue is not None else [] + ) + # ────────────────────────────────────────────────────────────── # Create TaskManager # ────────────────────────────────────────────────────────────── @@ -722,6 +831,7 @@ def __init__( seq_len_q=1, batch_size=1, mask_type: MaskType | str = MaskType.CAUSAL, + device_scales: bool = False, ): """ Parameters @@ -804,6 +914,20 @@ def __init__( raise ValueError("batch_size must be positive") self.batch_size = batch_size self.mask_type = normalize_mask_type(mask_type) + self.device_scales = device_scales + self.fuse_sparse_reduction = False + self.fuse_sparse_epilogue = False + self.output_quant = None + # Separate prepared source lists, rather than a premerged tagged list. + # Both representations reuse the same metadata staging and KV loaders. + self.direct_sparse = False + self.sparse_gather_warps = 1 + self.sparse_kv_stages = 0 + self.sparse_uniform_pages = False + self.sparse_balanced_registers = False + self.assume_valid_prefix = False + self.direct_static_scales = False + self.direct_sparse_capacities = (0, 0) self.query_tile_layout = FlatQueryTileLayout.for_tile( num_heads, seq_len_q, mma_qk_tiler_mn[0] ) @@ -827,6 +951,13 @@ def _effective_reduction_shape(self) -> tuple[int, int]: return self.mma_qk_tiler_mn[0], self.num_q_tiles + def _reduction_topology(self): + if self.fuse_sparse_reduction: + return make_balanced_parallel_reduction_topology( + self.reduction_split_capacity, cluster_size=1 + ) + return self.parallel_reduction_topology + def compile_signature(self) -> tuple[object, ...]: """Return the complete batch-independent JIT identity.""" @@ -854,7 +985,18 @@ def compile_signature(self) -> tuple[object, ...]: self.tail_q_rows, self._parallel_reduction_shape_is_eligible, self.use_parallel_reduction, - self.parallel_reduction_topology, + self._reduction_topology(), + self.fuse_sparse_reduction, + self.fuse_sparse_epilogue, + self.output_quant, + self.direct_sparse, + self.assume_valid_prefix, + self.direct_static_scales, + self.direct_sparse_capacities, + self.sparse_gather_warps, + self.sparse_kv_stages, + self.sparse_uniform_pages, + self.sparse_balanced_registers, ) def _configure_parallel_reduction_topology(self) -> None: @@ -915,6 +1057,46 @@ def _configure_parallel_reduction_topology(self) -> None: self.parallel_reduction_topology = topology self.use_parallel_reduction = True + def _make_config(self): + """Share geometry and resource settings between launch and device code.""" + cfg = make_mla_decode_config( + mma_qk_tiler_mn=self.mma_qk_tiler_mn, + mma_pv_tiler_mn=self.mma_pv_tiler_mn, + rope_dim=self.rope_dim, + page_size=self.page_size, + qkv_dtype=self.qkv_dtype, + o_dtype=self.out_dtype, + max_active_clusters=self.max_active_clusters, + is_persistent=self.is_persistent, + is_var_seq=self.is_var_seq, + is_var_split_kv=self.is_var_split_kv, + mask_type=self.mask_type, + sparse_gather_warps=self.sparse_gather_warps, + sparse_kv_stages=self.sparse_kv_stages, + cache_uniform_sparse_pages=self.sparse_uniform_pages, + balance_sparse_registers=self.sparse_balanced_registers, + fuse_output_quant=self.output_quant is not None + and self.fuse_sparse_epilogue, + ) + + # One flattened sparse query supplies a complete M128 tile. Its two + # CTAs each stage 64 heads without row predicates. Warp 10 is otherwise + # padding in the sparse FP8 schedule; using it for stores preserves + # gather/MMA/softmax resources and overlaps output traffic with work. + # Split reductions and other geometries retain direct vector stores. + cfg.use_tma_output = bool( + self.output_quant is not None + and self.fuse_sparse_epilogue + and cfg.is_fp8_qkv() + and cfg.page_size == 1 + and self.num_heads == 128 + # A 64-head TMA box spans whole head groups of the group-major output. + and 64 % (self.num_heads // self.output_quant.num_groups) == 0 + and self.seq_len_q == 1 + and cfg.mma_pv_tiler_mn == (128, 256) + ) + return cfg + @cute.jit def __call__( self, @@ -929,25 +1111,40 @@ def __call__( split_kv: cutlass.Int32, cache_seqs: cute.Tensor, cu_seqlens_q: cute.Tensor | None, - block_split_kvs: cute.Tensor, + block_split_kvs: object, softmax_scale: cutlass.Float32, output_scale: cutlass.Float32, stream: object, + output_quant_args=None, ): """Execute the MLA decode TS kernel.""" - cfg = make_mla_decode_config( - mma_qk_tiler_mn=self.mma_qk_tiler_mn, - mma_pv_tiler_mn=self.mma_pv_tiler_mn, - rope_dim=self.rope_dim, - page_size=self.page_size, - qkv_dtype=self.qkv_dtype, - o_dtype=self.out_dtype, - max_active_clusters=self.max_active_clusters, - is_persistent=self.is_persistent, - is_var_seq=self.is_var_seq, - is_var_split_kv=self.is_var_split_kv, - mask_type=self.mask_type, - ) + if cutlass.const_expr(self.fuse_sparse_reduction): + if cutlass.const_expr( + self.is_persistent + or self.is_var_split_kv + or self.static_split_kv is None + or self.static_split_kv < 2 + ): + raise ValueError( + "fused sparse reduction requires fixed, nonpersistent split-KV" + ) + if cutlass.const_expr(self.fuse_sparse_epilogue): + # This CTA can finalize only an unsplit output. Split partials + # must first merge their statistics; this is not a gather4 limit. + if cutlass.const_expr(self.static_split_kv != 1): + raise ValueError("2CTA fused epilogue requires one split") + if cutlass.const_expr(self.direct_sparse): + if cutlass.const_expr( + not (self.fuse_sparse_reduction or self.fuse_sparse_epilogue) + or self.qkv_dtype not in ("e4m3", "bf16") + or self.page_size != 1 + or self.rope_dim != 0 + or self.seq_len_q != 1 + ): + raise ValueError( + "direct sparse 2CTA requires native FP8/BF16 D512 and fused output" + ) + cfg = self._make_config() physical_tile_rows = self.mma_qk_tiler_mn[0] num_query_tiles = self.num_q_tiles # Fixed public tensors retain [H,D,SQ,B]/[H,SQ,B]; variable-Q tensors @@ -1084,48 +1281,127 @@ def _flatten_query_rows_for_tma(t): c_latent.iterator, cute.select(c_latent.layout, mode=[1, 0, 2]), ) - tma_desc_c_latent = create_tensor_map_tiled_from_view( - c_latent_tma, - box_dims=(cfg.mma_qk_tiler[2], cfg.kc_page_tile_size, 1), - stride_order=(0, 1, 2), - swizzle=cuda.TensorMapSwizzle.s128b, - l2_promotion=cuda.TensorMapL2Promotion.l2_128b, - ) - - if cutlass.const_expr(cfg.rope_dim > 0): - c_rope_tma = cute.make_tensor( + if cutlass.const_expr(cfg.page_size == 1): + # D512 has no separate Q/K RoPE maps. Reuse their ABI slots for + # extra-K and extra-V descriptors, keeping dense launch ABI stable. + primary_rows = cute.make_tensor( + c_latent.iterator, + cute.make_layout( + (c_latent.shape[1], c_latent.shape[2]), + stride=(c_latent.stride[1], c_latent.stride[2]), + ), + ) + extra_rows = cute.make_tensor( c_rope.iterator, - cute.select(c_rope.layout, mode=[1, 0, 2]), + cute.make_layout( + (c_rope.shape[1], c_rope.shape[2]), + stride=(c_rope.stride[1], c_rope.stride[2]), + ), + ) + tma_desc_c_latent = create_tensor_map_tiled_from_view( + primary_rows, + box_dims=(cfg.mma_qk_tiler[2], 1), + stride_order=(0, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, ) - c_rope_swizzle = cuda.TensorMapSwizzle.s128b - if cutlass.const_expr(cfg.is_fp8_qkv() and cfg.rope_dim == 64): - c_rope_swizzle = cuda.TensorMapSwizzle.s64b tma_desc_c_rope = create_tensor_map_tiled_from_view( - c_rope_tma, - box_dims=(cfg.mma_qk_rope_tiler[2], cfg.kc_page_tile_size, 1), - stride_order=(0, 1, 2), - swizzle=c_rope_swizzle, + extra_rows, + box_dims=(cfg.mma_qk_tiler[2], 1), + stride_order=(0, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + v_swizzle = cuda.TensorMapSwizzle.s128b + sparse_v_width = 128 if cfg.is_fp8_qkv() else V_TMA_LATENT_ELEMENTS + tma_desc_c_transpose = create_tensor_map_tiled_from_view( + primary_rows, + box_dims=(sparse_v_width, 1), + stride_order=(0, 1), + swizzle=v_swizzle, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + tma_desc_q_rope = create_tensor_map_tiled_from_view( + extra_rows, + box_dims=(sparse_v_width, 1), + stride_order=(0, 1), + swizzle=v_swizzle, l2_promotion=cuda.TensorMapL2Promotion.l2_128b, ) else: - tma_desc_c_rope = tma_desc_c_latent - - c_transpose_swizzle = cuda.TensorMapSwizzle.s128b - if cutlass.const_expr(cfg.is_fp8_qkv() and cfg.mma_pv_tiler[2] == 64): - c_transpose_swizzle = cuda.TensorMapSwizzle.s64b - c_latent_transpose_layout = cute.select(c_latent.layout, mode=[1, 0, 2]) - c_latent_transpose = cute.make_tensor( - c_latent.iterator, c_latent_transpose_layout - ) - # The physical page controls only the GMEM coordinate. TMA assembles - # fixed K32 SMEM blocks for PV from one or more page-bounded copies. - tma_desc_c_transpose = create_tensor_map_tiled_from_view( - c_latent_transpose, - box_dims=(V_TMA_LATENT_ELEMENTS, cfg.v_tma_token_count, 1), - stride_order=(0, 1, 2), - swizzle=c_transpose_swizzle, - l2_promotion=cuda.TensorMapL2Promotion.l2_128b, - ) + tma_desc_c_latent = create_tensor_map_tiled_from_view( + c_latent_tma, + box_dims=(cfg.mma_qk_tiler[2], cfg.kc_page_tile_size, 1), + stride_order=(0, 1, 2), + swizzle=cuda.TensorMapSwizzle.s128b, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + + if cutlass.const_expr(cfg.rope_dim > 0): + c_rope_tma = cute.make_tensor( + c_rope.iterator, + cute.select(c_rope.layout, mode=[1, 0, 2]), + ) + c_rope_swizzle = cuda.TensorMapSwizzle.s128b + if cutlass.const_expr(cfg.is_fp8_qkv() and cfg.rope_dim == 64): + c_rope_swizzle = cuda.TensorMapSwizzle.s64b + tma_desc_c_rope = create_tensor_map_tiled_from_view( + c_rope_tma, + box_dims=(cfg.mma_qk_rope_tiler[2], cfg.kc_page_tile_size, 1), + stride_order=(0, 1, 2), + swizzle=c_rope_swizzle, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + else: + tma_desc_c_rope = tma_desc_c_latent + + c_transpose_swizzle = cuda.TensorMapSwizzle.s128b + if cutlass.const_expr(cfg.is_fp8_qkv() and cfg.mma_pv_tiler[2] == 64): + c_transpose_swizzle = cuda.TensorMapSwizzle.s64b + c_latent_transpose_layout = cute.select(c_latent.layout, mode=[1, 0, 2]) + c_latent_transpose = cute.make_tensor( + c_latent.iterator, c_latent_transpose_layout + ) + # The physical page controls only the GMEM coordinate. TMA assembles + # fixed K32 SMEM blocks for PV from one or more page-bounded copies. + tma_desc_c_transpose = create_tensor_map_tiled_from_view( + c_latent_transpose, + box_dims=(V_TMA_LATENT_ELEMENTS, cfg.v_tma_token_count, 1), + stride_order=(0, 1, 2), + swizzle=c_transpose_swizzle, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + + tma_desc_o = tma_desc_q_latent + if cutlass.const_expr(cfg.use_tma_output): + flat_o = output_quant_args[0] + output_rows = cute.size(flat_o) // (128 * 512) + # Group-major [G, T, (H/G) * 512] output viewed as (channel, head in + # group, token, group). A 64-head box covers 64 / (H/G) whole groups + # and fills 64 shared-memory rows in head order, so the staging + # layout matches the flat head-major view. + heads_per_group = ( + self.output_quant.num_heads // self.output_quant.num_groups + ) + output_view = cute.make_tensor( + flat_o.iterator, + cute.make_layout( + (512, heads_per_group, output_rows, self.output_quant.num_groups), + stride=( + 1, + 512, + heads_per_group * 512, + output_rows * heads_per_group * 512, + ), + ), + ) + tma_desc_o = create_tensor_map_tiled_from_view( + output_view, + box_dims=(128, heads_per_group, 1, 64 // heads_per_group), + stride_order=(0, 1, 2, 3), + swizzle=cuda.TensorMapSwizzle.s128b, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) softmax_scale_log2 = softmax_scale * LOG2_E @@ -1173,9 +1449,9 @@ def _flatten_query_rows_for_tma(t): kernel_split_kv, workspace, ) - # A one-wave producer can publish its dependent reducer launch while - # retiring, hiding launch latency without admitting reducer CTAs into - # an actively grid-striding persistent producer. + # Nonpersistent producers signal near retirement. The native reducer + # can overlap setup, then waits for the complete prerequisite grid + # before reading partials. Persistent schedules retain stream ordering. use_one_wave_reducer_pdl = acc_o is not None and not self.is_persistent self.split_kv_kernel( @@ -1199,6 +1475,8 @@ def _flatten_query_rows_for_tma(t): output_scale, tile_sched_params, clc_tile_sched_params, + output_quant_args, + tma_desc_o, ).launch( grid=grid, block=[cfg.threads_per_cta, 1, 1], @@ -1210,8 +1488,10 @@ def _flatten_query_rows_for_tma(t): # Reduction kernel: combine per-split results when split_kv > 1 if cutlass.const_expr(acc_o is not None): - if cutlass.const_expr(self.use_parallel_reduction): - topology = self.parallel_reduction_topology + if cutlass.const_expr( + self.use_parallel_reduction or self.fuse_sparse_reduction + ): + topology = self._reduction_topology() self.parallel_reduction_kernel( o, lse, @@ -1221,14 +1501,22 @@ def _flatten_query_rows_for_tma(t): cache_seqs, cu_seqlens_q, block_split_kvs, + output_quant_args, ).launch( grid=( - physical_tile_rows * topology.cluster_size, + ( + min(self.num_heads, physical_tile_rows) + if self.fuse_sparse_reduction + else physical_tile_rows + ) + * topology.cluster_size, num_query_tiles, batch_size, ), block=[PARALLEL_REDUCTION_THREADS, 1, 1], - cluster=[topology.cluster_size, 1, 1], + cluster=None + if topology.cluster_size == 1 + else [topology.cluster_size, 1, 1], stream=stream, min_blocks_per_mp=1, use_pdl=use_one_wave_reducer_pdl, @@ -1280,26 +1568,64 @@ def split_kv_kernel( split_kv: cutlass.Int32, cache_seqs: cute.Tensor, cu_seqlens_q: cute.Tensor | None, - block_split_kvs: cute.Tensor, + block_split_kvs: object, softmax_scale_log2: cutlass.Float32, output_scale: cutlass.Float32, tile_sched_params: MLAStaticTileSchedulerParams, clc_tile_sched_params: object, + output_quant_args=None, + tma_desc_o: cutlass.GridConstant[cuda.TensorMap] = None, ) -> None: """MLA decode TS kernel: persistent tile-scheduled execution.""" - cfg = make_mla_decode_config( - mma_qk_tiler_mn=self.mma_qk_tiler_mn, - mma_pv_tiler_mn=self.mma_pv_tiler_mn, - rope_dim=self.rope_dim, - page_size=self.page_size, - qkv_dtype=self.qkv_dtype, - o_dtype=self.out_dtype, - max_active_clusters=self.max_active_clusters, - is_persistent=self.is_persistent, - is_var_seq=self.is_var_seq, - is_var_split_kv=self.is_var_split_kv, - mask_type=self.mask_type, - ) + if cutlass.const_expr(self.direct_sparse): + from ..sparse_views import bind_sparse_views + + si, ci, sl, cl, sm, qs, ss, os, sinks = block_split_kvs + capacity = cutlass.const_expr( + max( + 256, + ( + (self.direct_sparse_capacities[0] + 127) // 128 + + (self.direct_sparse_capacities[1] + 127) // 128 + ) + * 128, + ) + ) + request = None + if cutlass.const_expr(not self.is_persistent): + _, row = divmod_constexpr_power_of_two_or_fdd( + cute.arch.block_idx()[1], + None, + tile_sched_params.problem_shape_b_fdd, + ) + request = Int64(row) + page_offsets, cache_seqs = bind_sparse_views( + si, + ci, + sl, + cl, + capacity, + request=request, + assume_valid_prefix=self.assume_valid_prefix, + ) + if cutlass.const_expr(not self.direct_static_scales): + softmax_scale_log2 = ( + Float32(sm[0]) * Float32(qs[0]) * Float32(ss[0]) * LOG2_E + ) + output_scale = Float32(os[0]) * Float32(ss[0]) + else: + if cutlass.const_expr(self.device_scales): + softmax_scale_log2 = cutlass.Float32(block_split_kvs[0]) * LOG2_E + output_scale = cutlass.Float32(block_split_kvs[1]) + if cutlass.const_expr(self.fuse_sparse_epilogue): + # A host-scale caller can bind sinks alone. Device-scale + # metadata starts with QK/PV scales, followed by sinks. + sink_offset = 2 if self.device_scales else 0 + sinks = cute.make_tensor( + block_split_kvs.iterator + sink_offset, + cute.make_layout(self.num_heads), + ) + cfg = self._make_config() num_query_tiles = self.num_q_tiles use_clc_dynamic = self.is_persistent and not cfg.is_fp8_qkv() tiled_mma_qk = None @@ -1362,6 +1688,26 @@ def split_kv_kernel( space=cutlass.AddressSpace.smem, alignment=1024, ) + # One shared metadata ring feeds both K/V issuers. Each holds its + # stage through the last coordinate read; TMA consumes register + # operands and its payload buffers have independent lifetimes. + sparse_offsets_smem = None + if cutlass.const_expr(cfg.page_size == 1 and cfg.use_fp8_split_mma_schedule): + sparse_offsets_smem = cutlass.Array( + Int32, + cfg.sparse_offset_stages * cfg.tokens_per_k_tile, + space=cutlass.AddressSpace.smem, + alignment=128, + ) + smem_o = None + if cutlass.const_expr(cfg.use_tma_output): + smem_o = cutlass.Array( + cutlass.Uint8, + OUTPUT_TMA_STAGES * OUTPUT_TMA_STAGE_BYTES, + space=cutlass.AddressSpace.smem, + alignment=1024, + ) + softmax_exchange_arr = cutlass.Array( self.acc_dtype, cfg.softmax_exchange_elems, @@ -1452,6 +1798,26 @@ def split_kv_kernel( ) blk_coord = (cluster_idx, seq_q_idx, batch_idx, split_kv_idx) tile_cluster_idx, tile_seq_q_idx, tile_batch_idx, tile_split_kv_idx = blk_coord + if cutlass.const_expr(self.fuse_sparse_reduction): + # Final output uses a fixed-S reducer: neutral LSE slots let it + # avoid reloading source lengths and dividing split ranges in + # every output CTA. Active tiles publish their own statistics. + if tidx < Int32(cfg.mma_qk_tiler[0] // cfg.num_mma_ctas): + total_tiles = (Int32(cache_seqs[tile_batch_idx]) + Int32(127)) // Int32( + 128 + ) + _, local_tiles = runtime_split_tile_range( + total_tiles, split_kv, tile_split_kv_idx + ) + if local_tiles == Int32(0): + row_in_tile = ( + tile_cluster_idx + * Int32(cfg.mma_qk_tiler[0] // cfg.num_mma_ctas) + + tidx + ) + acc_lse[ + row_in_tile, tile_split_kv_idx, tile_seq_q_idx, tile_batch_idx + ] = Float32(-Float32.inf) del tile_cluster_idx fixed_nonempty_single_split = ( @@ -1636,6 +2002,7 @@ def split_kv_kernel( smem_kc_arr=smem_kc_arr, smem_vc_arr=smem_vc_arr, smem_p_arr=smem_p_arr, + sparse_offsets_smem=sparse_offsets_smem, tma_desc_q_latent=tma_desc_q_latent.get_ptr(), tma_desc_q_rope=tma_desc_q_rope.get_ptr(), tma_desc_c_latent=tma_desc_c_latent.get_ptr(), @@ -1645,6 +2012,12 @@ def split_kv_kernel( blk_coord=blk_coord, tidx=tidx, output=o, + atten_sinks=sinks if self.fuse_sparse_epilogue else None, + lse_in_natural_log=self.fuse_sparse_epilogue, + output_quant=self.output_quant if self.fuse_sparse_epilogue else None, + output_quant_args=output_quant_args, + smem_o=smem_o, + tma_o_desc=tma_desc_o.get_ptr(), acc_output=acc_o, lse=lse, acc_lse=acc_lse, @@ -1838,7 +2211,8 @@ def parallel_reduction_kernel( split_kv: cutlass.Int32, cache_seqs: cute.Tensor, cu_seqlens_q: cute.Tensor | None, - block_split_kvs: cute.Tensor, + block_split_kvs: object, + output_quant_args=None, ): """Dispatch the high-split fixed-D512 cluster reducer.""" @@ -1853,7 +2227,16 @@ def parallel_reduction_kernel( ) if cutlass.const_expr(not self.is_persistent): prims.griddepcontrol(kind=prims.GridDepAction.WAIT) - topology = self.parallel_reduction_topology + atten_sinks = None + if cutlass.const_expr(self.fuse_sparse_reduction): + if cutlass.const_expr(self.direct_sparse): + atten_sinks = block_split_kvs[-1] + else: + atten_sinks = cute.make_tensor( + block_split_kvs.iterator + 2, cute.make_layout(self.num_heads) + ) + block_split_kvs = None + topology = self._reduction_topology() run_parallel_reduction_kernel( self.num_heads, self.seq_len_q, @@ -1870,4 +2253,8 @@ def parallel_reduction_kernel( topology.actual_splits, topology.cluster_size, topology.slots_per_rank, + atten_sinks, + self.fuse_sparse_reduction, + self.output_quant, + output_quant_args, ) 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 2512ef3dcfa1..41b61707a4c4 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 @@ -19,8 +19,12 @@ from cutlass import Float32, Int32, Int64 from cutlass.experimental import primitives as prims -from ...separate_reduction import finalize_log2_sum_exp, normalized_lse_weight -from ..helpers.constants import SPLIT_REDUCTION_SCALE_BARRIER_ID +from ...separate_reduction import ( + finalize_log2_sum_exp, + normalized_lse_weight, + attention_sink_scale, +) +from ..helpers.constants import LN_2, SPLIT_REDUCTION_SCALE_BARRIER_ID from ..helpers.mask import MaskType, mask_visible_k_length from ..helpers.math import ceil_div from ..helpers.ops import fmax_f32, warp_reduce_max_f32, warp_reduce_sum_f32 @@ -147,15 +151,30 @@ def _store_parallel_reduction_result( tidx, batch_idx, cu_seqlens_q, + finalize_attention: cutlass.Constexpr[bool] = False, + output_quant: cutlass.Constexpr = None, + output_quant_args=None, ): """Publish one normalized FP32 output fragment and final LSE.""" + if cutlass.const_expr(finalize_attention): + global_lse *= Float32(LN_2) if tidx == Int32(0) and query_is_valid: if cutlass.const_expr(cu_seqlens_q is not None): lse[logical_head_idx, storage_q_idx] = global_lse else: lse[logical_head_idx, logical_q_idx, batch_idx] = global_lse + if cutlass.const_expr(output_quant is not None): + if query_is_valid: + row = Int64(storage_q_idx) * output_quant.num_heads + logical_head_idx + if cutlass.const_expr(cu_seqlens_q is None): + row = ( + Int64(batch_idx) * output.shape[2] + logical_q_idx + ) * output_quant.num_heads + logical_head_idx + output_quant.store(output_quant_args, output_vals, row, element_idx) + return + out_element_dtype = output.element_type output_regs = cutlass.Array( out_element_dtype, @@ -193,6 +212,10 @@ def run_parallel_reduction_kernel( actual_splits: cutlass.Constexpr[int], cluster_size: cutlass.Constexpr[int], slots_per_rank: cutlass.Constexpr[int], + atten_sinks=None, + finalize_attention: cutlass.Constexpr[bool] = False, + output_quant: cutlass.Constexpr = None, + output_quant_args=None, ): """Reduce one D=512 row cooperatively across a padded CTA cluster. @@ -206,32 +229,46 @@ def run_parallel_reduction_kernel( signal before entering this body. """ + if cutlass.const_expr(finalize_attention and cluster_size != 1): + raise ValueError("final attention output uses a single-CTA reducer per row") block_idx_x, query_tile_idx, batch_idx = cute.arch.block_idx() tidx, _, _ = cute.arch.thread_idx() warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) lane_idx = tidx % Int32(cfg.threads_per_warp) - cluster_rank = cute.arch.block_idx_in_cluster() + cluster_rank = Int32(0) if cluster_size == 1 else cute.arch.block_idx_in_cluster() row_in_tile = block_idx_x // Int32(cluster_size) - ( - logical_head_idx, - logical_q_idx, - storage_q_idx, - query_is_valid, - active_split_kv, - ) = _parallel_reduction_row_state( - cfg.mma_qk_tiler[0], - num_heads, - seq_len_q, - is_var_split_kv, - cache_seqs, - cu_seqlens_q, - block_split_kvs, - split_kv, - row_in_tile, - query_tile_idx, - batch_idx, - cfg, - ) + if cutlass.const_expr(finalize_attention): + # A sparse query is already one virtual request. The producer writes + # neutral LSE for skipped splits, so the compiled slot count suffices; + # live length loads and dynamic split division would repeat per head. + if cutlass.const_expr(seq_len_q != 1 or cu_seqlens_q is not None): + raise ValueError("final sparse reduction requires virtual query rows") + logical_head_idx = query_tile_idx * Int32(cfg.mma_qk_tiler[0]) + row_in_tile + logical_q_idx = Int32(0) + storage_q_idx = batch_idx + query_is_valid = logical_head_idx < Int32(num_heads) + active_split_kv = Int32(actual_splits) + else: + ( + logical_head_idx, + logical_q_idx, + storage_q_idx, + query_is_valid, + active_split_kv, + ) = _parallel_reduction_row_state( + cfg.mma_qk_tiler[0], + num_heads, + seq_len_q, + is_var_split_kv, + cache_seqs, + cu_seqlens_q, + block_split_kvs, + split_kv, + row_in_tile, + query_tile_idx, + batch_idx, + cfg, + ) element_idx = tidx * Int32(PARALLEL_REDUCTION_ELEMENTS_PER_THREAD) neg_inf = Float32(-Float32.inf) @@ -297,16 +334,25 @@ def run_parallel_reduction_kernel( ) local_sum_lse += lane_exp[lane_slot_i] local_sum_lse = warp_reduce_sum_f32(local_sum_lse) - local_lse_value = finalize_log2_sum_exp(local_exp_frame, local_sum_lse) + # Keep this warp-local value distinct from the post-barrier SMEM load. + stats_lse_value = finalize_log2_sum_exp(local_exp_frame, local_sum_lse) + sink_scale = Float32(1) + if cutlass.const_expr(finalize_attention and atten_sinks is not None): + sink_scale = Float32(0) + if query_is_valid: + sink_scale = attention_sink_scale( + stats_lse_value, Float32(atten_sinks[logical_head_idx]) + ) if lane_idx == Int32(0): - smem_local_lse[0] = local_lse_value + smem_local_lse[0] = stats_lse_value for lane_slot_i in cutlass.range_constexpr(lse_slots_per_lane): local_slot_idx = lane_idx + Int32(lane_slot_i * cfg.threads_per_warp) if local_slot_idx < Int32(slots_per_rank): - smem_local_scale[local_slot_idx] = normalized_lse_weight( - lane_lse[lane_slot_i], local_lse_value - ) + weight = normalized_lse_weight(lane_lse[lane_slot_i], stats_lse_value) + if cutlass.const_expr(finalize_attention): + weight *= sink_scale + smem_local_scale[local_slot_idx] = weight prims.barrier_cta_sync(SPLIT_REDUCTION_SCALE_BARRIER_ID) @@ -331,23 +377,26 @@ def run_parallel_reduction_kernel( ) if active_slot: split_scale = Float32(smem_local_scale[slot_i]) - partial_offset = Int64( - row_in_tile * acc_output.stride[0] - + split_idx * acc_output.stride[1] - + element_idx * acc_output.stride[2] - + query_tile_idx * acc_output.stride[3] - + batch_idx * acc_output.stride[4] - ) - partial_output = ( - (acc_output_ptr + partial_offset) - .load( - count=PARALLEL_REDUCTION_ELEMENTS_PER_THREAD, - alignment=8, + if not finalize_attention or split_scale > Float32(0): + partial_offset = Int64( + row_in_tile * acc_output.stride[0] + + split_idx * acc_output.stride[1] + + element_idx * acc_output.stride[2] + + query_tile_idx * acc_output.stride[3] + + batch_idx * acc_output.stride[4] ) - .to(Float32) - ) - for j in cutlass.range_constexpr(PARALLEL_REDUCTION_ELEMENTS_PER_THREAD): - local_output[j] += partial_output[j] * split_scale + partial_output = ( + (acc_output_ptr + partial_offset) + .load( + count=PARALLEL_REDUCTION_ELEMENTS_PER_THREAD, + alignment=8, + ) + .to(Float32) + ) + for j in cutlass.range_constexpr( + PARALLEL_REDUCTION_ELEMENTS_PER_THREAD + ): + local_output[j] += partial_output[j] * split_scale local_lse_value = smem_local_lse[0] @@ -367,6 +416,9 @@ def run_parallel_reduction_kernel( tidx, batch_idx, cu_seqlens_q, + finalize_attention, + output_quant, + output_quant_args, ) return 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 55523137fdea..1b3950e29954 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 @@ -32,14 +32,16 @@ - TmemCorrResource : Async(2 stages), SoftmaxTask -> CorrectionTask - TmemOResource : UmmaProducerAsync(1 stage), MmaTask -> CorrectionTask -GMEM (no pipeline) ------------------- -- GmemOResource : No pipeline, Correction -> GMEM +Output +------ +- GmemOResource : Direct stores, or Async(2 stages) Correction -> TMA -> GMEM """ from dataclasses import dataclass, field from typing import Any, ClassVar +from ..helpers.constants import LN_2, LOG2_E + import cutlass import cutlass.cute as cute from cutlass.experimental import primitives as prims @@ -56,7 +58,18 @@ from ...mask import kv_tile_needs_right_mask from ...tensor_map import transform_ragged_coords +from ..helpers.gather import ( + gather4_cached, + load_sparse_rows, + invalid_sparse_token, + gather4_uniform_quad_slices, + load_sparse_quad, + cache_sparse_mask, + CachedSparseMask, +) + from .config import ( + OUTPUT_TMA_STAGE_BYTES, V_SMEM_K_BLOCK_TOKENS, V_TMA_LATENT_ELEMENTS, MlaDecodeConfig, @@ -72,6 +85,8 @@ TCGEN05_32B_REGS_PER_LOAD, TCGEN05_32B_SHAPE, WARP_LANE_SHIFT, + WARP_LANES, + TMA_GATHER_ROWS, ) from ..helpers.tile_scheduler import ( MLAStaticTileScheduler, @@ -623,10 +638,19 @@ def advance_tile(self, stage_info: StageInfo): @dataclass(kw_only=True) class PageOffsetWindowResource(HighThroughputMlaResource): - """GMEM page table cached as one register per lane of the TMA warp.""" + """Page-coordinate staging for paged and sparse KV loads. + + FP8 sparse loaders read uniform quads from the prefetched shared ring. + Other schedules retain page windows in registers. + """ + request_lengths: Any = field(init=False, default=None) page_offsets: Any = None # GMEM page-offset tensor + smem_sparse_offsets: Any = None cfg: cutlass.Constexpr = field(default_factory=MlaDecodeConfig) + cached_page_stage: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) cached_k_pages: cutlass.Constexpr[TaskLocalVariable] = ( TaskLocalVariable.uninitialized() ) @@ -639,7 +663,21 @@ class PageOffsetWindowResource(HighThroughputMlaResource): cached_window_page: cutlass.Constexpr[TaskLocalVariable] = ( TaskLocalVariable.uninitialized() ) + cached_mask_low: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + cached_mask_high: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ("cached_page_stage", Int32, Int32(0), "Held shared-offset stage."), + ("cached_mask_low", Int32, Int32(0), "Low half of the prefetched sparse mask."), + ( + "cached_mask_high", + Int32, + Int32(0), + "High half of the prefetched sparse mask.", + ), ( "cached_k_pages", cutlass.Array, @@ -666,6 +704,70 @@ class PageOffsetWindowResource(HighThroughputMlaResource): ), ) + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_sparse_state(self, stage_info: StageInfo): + del stage_info + # Allocate outside skipped-work guards; refresh values for every tile. + self.request_lengths = cutlass.Array(Int32, 2, space=cutlass.AddressSpace.rmem) + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def bind_sparse_request(self, stage_info: StageInfo): + if cutlass.const_expr(hasattr(self.page_offsets, "for_request")): + row = stage_info.work_tile.tile_idx[2] + self.request_lengths[0] = Int32(self.page_offsets.sl[row]) + self.request_lengths[1] = Int32(self.page_offsets.cl[row]) + + @producer_work + @cute.jit + def load_sparse_offsets(self, stage_info: StageInfo): + lane = cute.arch.thread_idx()[0] % 32 + k_index = stage_info.work_tile.k_index_base + Int32(stage_info.loop_offset) + token = k_index * Int32(128) + lane * Int32(4) + row = stage_info.work_tile.tile_idx[2] + if cutlass.const_expr(hasattr(self.page_offsets, "for_request")): + routes = self.page_offsets.for_request( + row, self.request_lengths[0], self.request_lengths[1] + ) + else: + routes = self.page_offsets + quad = load_sparse_quad(routes, row, token) + self.smem_sparse_offsets.data_ptr(stage_info.stage_idx * 128 + lane * 4).store( + quad.load(0, 4), alignment=16 + ) + cute.arch.fence_view_async_shared() + + @consumer_work(returns=(cached_mask_low, cached_mask_high)) + @cute.jit + def read_sparse_mask(self, stage_info: StageInfo): + lane = cute.arch.thread_idx()[0] % 32 + half = (cute.arch.thread_idx()[0] % 128) // 64 + base = stage_info.stage_idx * 128 + half * 64 + tile = stage_info.work_tile.k_index_base + Int32(stage_info.loop_offset) + token = tile * Int32(128) + half * Int32(64) + lane + length = Int32(stage_info.work_tile.k_len) + low = self.smem_sparse_offsets[base + lane] + high = self.smem_sparse_offsets[base + lane + 32] + low_bits = Int32( + cute.arch.vote_ballot_sync( + ((low & Int32(0x7FFFFFFF)) == Int32(0x7FFFFFFF)) | (token >= length) + ) + ) + high_bits = Int32( + cute.arch.vote_ballot_sync( + ((high & Int32(0x7FFFFFFF)) == Int32(0x7FFFFFFF)) + | (token + Int32(32) >= length) + ) + ) + return low_bits, high_bits + + @consumer_work(returns=cached_page_stage) + @cute.jit + def sparse_stage(self, stage_info: StageInfo): + """Return the shared-ring slot protected by the current consumer token.""" + return stage_info.stage_idx + @consumer_work( work_attrs=WorkAttr.AUXILIARY, returns=( @@ -680,20 +782,37 @@ def init_read_state(self, stage_info: StageInfo): """Create cached page-index arrays used by staged KV TMA loads.""" del stage_info cfg = self.cfg + # Uniform caches repeat each warp's K/current-V coordinates in all + # lanes, saving shuffles on later feature slices. Keep pending V in + # compact lane fragments until needed to limit live registers. return ( cutlass.Array( Int32, - cfg.pages_per_k_cta, + ( + cfg.tokens_per_k_cta // cfg.load_num_warps + if cfg.cache_uniform_sparse_pages + else cfg.tokens_per_k_cta // WARP_LANES + ) + if cfg.page_size == 1 + else cfg.pages_per_k_cta, space=cutlass.AddressSpace.rmem, ), cutlass.Array( Int32, - cfg.pages_per_v_tile, + ( + cfg.tokens_per_v_tile // cfg.load_num_warps + if cfg.cache_uniform_sparse_pages + else cfg.tokens_per_v_tile // WARP_LANES + ) + if cfg.page_size == 1 + else cfg.pages_per_v_tile, space=cutlass.AddressSpace.rmem, ), cutlass.Array( Int32, - cfg.pages_per_v_tile, + cfg.tokens_per_v_tile // WARP_LANES + if cfg.page_size == 1 + else cfg.pages_per_v_tile, space=cutlass.AddressSpace.rmem, ), Int32(0), @@ -731,6 +850,79 @@ def read_page_offset_window( blk_coord = work_tile.tile_idx local_k_index = Int32(stage_info.loop_offset) global_k_index = work_tile.k_index_base + local_k_index + if cutlass.const_expr(cfg.page_size == 1): + cta = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + new_k = load_sparse_rows( + self.page_offsets, + blk_coord[2], + global_k_index * cfg.tokens_per_k_tile + cta * cfg.tokens_per_k_cta, + fragments=cfg.tokens_per_k_cta // WARP_LANES, + ) + new_v = load_sparse_rows( + self.page_offsets, + blk_coord[2], + global_k_index * cfg.tokens_per_v_tile, + fragments=cfg.tokens_per_v_tile // WARP_LANES, + ) + # Expand only current K/V coordinates once per tile. "Uniform" + # means identical across lanes, not consecutive physical KV rows. + # Reuse across D slices trades register capacity for fewer shuffles. + if cutlass.const_expr(cfg.cache_uniform_sparse_pages): + warp = ( + cute.arch.make_warp_uniform(cute.arch.warp_idx()) + - cfg.load_tma_warp_id + ) + fragment = ( + new_k[0] if warp < Int32(cfg.load_num_warps // 2) else new_k[1] + ) + for i in cutlass.range_constexpr(64 // cfg.load_num_warps): + cached_k_pages[i] = cute.arch.make_warp_uniform( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=fragment, + offset=(warp % Int32(cfg.load_num_warps // 2)) + * Int32(64 // cfg.load_num_warps) + + Int32(i), + mask_and_clamp=0x1F, + kind=prims.Shfl.IDX, + ) + ) + else: + for i in cutlass.range_constexpr(2): + cached_k_pages[i] = new_k[i] + if cutlass.const_expr(cfg.cache_uniform_sparse_pages): + # Expand only the delayed V tile; retain the next V tile in + # lane fragments to keep the register footprint bounded. + for i in cutlass.range_constexpr(128 // cfg.load_num_warps): + if cutlass.const_expr(init_v_cache): + cached_v_pages[i] = Int32(0x7FFFFFFF) + else: + cached_v_pages[i] = cute.arch.make_warp_uniform( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=cached_next_v_pages[ + i // (32 // cfg.load_num_warps) + ], + offset=warp * Int32(32 // cfg.load_num_warps) + + Int32(i % (32 // cfg.load_num_warps)), + mask_and_clamp=0x1F, + kind=prims.Shfl.IDX, + ) + ) + for i in cutlass.range_constexpr(4): + cached_next_v_pages[i] = new_v[i] + else: + for i in cutlass.range_constexpr(4): + cached_v_pages[i] = ( + Int32(0x7FFFFFFF) if init_v_cache else cached_next_v_pages[i] + ) + cached_next_v_pages[i] = new_v[i] + return ( + cached_k_pages, + cached_v_pages, + cached_next_v_pages, + cached_window_page, + ) pages_per_k_tile = cutlass.const_expr(cfg.pages_per_k_tile) page_window_tiles = cutlass.const_expr(32 // pages_per_k_tile) page_window_mask = cutlass.const_expr(page_window_tiles - 1) @@ -864,6 +1056,15 @@ def _query_tma_coords( @producer_work @cute.jit def tma_load(self, stage_info: StageInfo) -> None: + """Let one warp issue Q while the full loader group signals readiness.""" + if cutlass.const_expr(self.cfg.load_num_warps > 1): + if cute.arch.thread_idx()[0] // 32 == self.cfg.load_k_warp_id: + self._tma_load(stage_info) + else: + self._tma_load(stage_info) + + @cute.jit + def _tma_load(self, stage_info: StageInfo) -> None: """TMA load Q latent + Q rope into SMEM (all sub-tiles in one commit).""" cfg = self.cfg work_tile = stage_info.work_tile @@ -963,6 +1164,7 @@ class SmemKVResource(HighThroughputMlaResource): tma_desc_c_latent: Any = None tma_desc_c_rope: Any = None tma_desc_c_transpose: Any = None + tma_desc_extra_v: Any = None cfg: cutlass.Constexpr = field(default_factory=MlaDecodeConfig) cta_rank: Any = field(init=False, default=None) is_leader: Any = field(init=False, default=None) @@ -1002,6 +1204,162 @@ def tma_load( cfg = self.cfg stage_idx = stage_info.stage_idx + if cutlass.const_expr(cfg.page_size == 1): + cta = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + base = stage_idx * cfg.smem_k_stage_elems + warp = ( + cute.arch.make_warp_uniform(cute.arch.warp_idx()) - cfg.load_tma_warp_id + ) + # Each issuer owns disjoint four-row quads. K covers 64 rows per + # CTA (16 quads); V copies 32-token subtiles (8 quads). BF16 V is + # blocked as 32x64 panels: two 64-feature/128-byte slices cover + # each CTA's 128 features within a 256-feature PV panel. Source + # columns are elements; gather4 destination strides are bytes. + if cutlass.const_expr(cfg.load_num_warps in (4, 8)): + if cutlass.const_expr(not is_v): + if cutlass.const_expr(not cfg.cache_uniform_sparse_pages): + fragment = ( + cached_k_pages[0] + if warp < Int32(cfg.load_num_warps // 2) + else cached_k_pages[1] + ) + for owner in cutlass.range_constexpr(16 // cfg.load_num_warps): + raw = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) + for j in cutlass.range_constexpr(4): + if cutlass.const_expr(cfg.cache_uniform_sparse_pages): + raw[j] = cached_k_pages[owner * 4 + j] + else: + lane_base = ( + warp % Int32(cfg.load_num_warps // 2) + ) * Int32(64 // cfg.load_num_warps) + Int32(owner * 4) + raw[j] = cute.arch.make_warp_uniform( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=fragment, + offset=lane_base + Int32(j), + mask_and_clamp=0x1F, + kind=prims.Shfl.IDX, + ) + ) + group = warp * Int32(16 // cfg.load_num_warps) + Int32(owner) + gather4_uniform_quad_slices( + self.smem_kv.data_ptr( + base + group * 4 * cfg.mma_qk_tiler_k + ), + self.tma_desc_c_latent, + self.tma_desc_c_rope, + Int32( + subtile_idx + * cfg.kv_subtiles_per_stage + * cfg.mma_qk_tiler_k + ), + raw, + stage_info.barrier, + num_slices=cfg.kv_subtiles_per_stage, + dst_stride=64 * cfg.mma_qk_tiler_k * cfg.qkv_dtype_bytes, + col_stride=cfg.mma_qk_tiler_k, + ) + else: + # PV consumes the delayed current-V window; the tail can + # instead request the compact pending-V coordinates. + cached = cached_next_v_pages if use_next_v_pages else cached_v_pages + pv_j = subtile_idx // cfg.kv_subtiles_per_stage + token_partition = subtile_idx % cfg.kv_subtiles_per_stage + for part in cutlass.range_constexpr(cfg.kv_subtiles_per_stage): + pv_i = token_partition * cfg.kv_subtiles_per_stage + part + for owner in cutlass.range_constexpr(8 // cfg.load_num_warps): + group = warp * Int32(8 // cfg.load_num_warps) + Int32(owner) + raw = cutlass.Array( + Int32, 4, space=cutlass.AddressSpace.rmem + ) + for j in cutlass.range_constexpr(4): + if cutlass.const_expr( + cfg.cache_uniform_sparse_pages + and not use_next_v_pages + ): + raw[j] = cached_v_pages[ + pv_i * (32 // cfg.load_num_warps) + + owner * 4 + + j + ] + else: + raw[j] = cute.arch.make_warp_uniform( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=cached[pv_i], + offset=group * Int32(4) + Int32(j), + mask_and_clamp=0x1F, + kind=prims.Shfl.IDX, + ) + ) + gather4_uniform_quad_slices( + self.smem_kv.data_ptr( + base + part * 32 * 128 + group * 4 * 64 + ), + self.tma_desc_c_transpose, + self.tma_desc_extra_v, + cta * Int32(128) + Int32(pv_j * 256), + raw, + stage_info.barrier, + num_slices=2, + dst_stride=32 * 64 * cfg.qkv_dtype_bytes, + col_stride=64, + ) + return + # Fewer issuers reconstruct quads from compact lane fragments. + if cutlass.const_expr(not is_v): + for part in cutlass.range_constexpr(cfg.kv_subtiles_per_stage): + column = Int32( + (subtile_idx * cfg.kv_subtiles_per_stage + part) + * cfg.mma_qk_tiler_k + ) + for local_group in cutlass.range( + 16 // cfg.load_num_warps, unroll_full=True + ): + group = local_group + gather4_cached( + self.smem_kv.data_ptr( + base + + part * 64 * cfg.mma_qk_tiler_k + + group * 4 * cfg.mma_qk_tiler_k + ), + self.tma_desc_c_latent, + self.tma_desc_c_rope, + column, + cached_k_pages, + group * 4, + stage_info.barrier, + cta_group=2, + ) + else: + cached = cached_next_v_pages if use_next_v_pages else cached_v_pages + pv_j = subtile_idx // cfg.kv_subtiles_per_stage + token_partition = subtile_idx % cfg.kv_subtiles_per_stage + for part in cutlass.range_constexpr(cfg.kv_subtiles_per_stage): + pv_i = token_partition * cfg.kv_subtiles_per_stage + part + for panel in cutlass.range_constexpr(2): + column = cta * 128 + Int32(pv_j * 256 + panel * 64) + for local_group in cutlass.range( + 8 // cfg.load_num_warps, unroll_full=True + ): + group = local_group + gather4_cached( + self.smem_kv.data_ptr( + base + + part * 32 * 128 + + panel * 32 * 64 + + group * 4 * 64 + ), + self.tma_desc_c_transpose, + self.tma_desc_extra_v, + column, + cached, + Int32(pv_i * 32) + group * 4, + stage_info.barrier, + cta_group=2, + ) + return + kc_page_smem_elems = cfg.kc_page_tile_size * cfg.mma_qk_tiler_k kv_mbar_arr = cutlass.Array(stage_info.barrier.data_ptr(), dtype=Int64) kv_stage_stride_elems = cutlass.const_expr(cfg.smem_k_stage_elems) @@ -1259,11 +1617,12 @@ def v_desc(self, stage_info: StageInfo, *, v_subtile_idx: cutlass.Constexpr[int] class SmemKResource(HighThroughputMlaResource): """SMEM K buffer for the FP8 split-MMA path. - Producer: LoadTma. Consumer: MmaQkTask. A single producer stage contains - all latent K sub-tiles plus the RoPE K sub-tile for one logical K tile. + Producer: LoadKTask. Consumer: MmaQkTask. A single producer stage contains + all latent K sub-tiles and optional RoPE channels for one logical K tile. """ smem_k: Any = None + smem_sparse_offsets: Any = None page_offsets: Any = None tma_desc_c_latent: Any = None tma_desc_c_rope: Any = None @@ -1284,8 +1643,42 @@ def init_load_state(self, stage_info: StageInfo) -> None: @producer_work @cute.jit - def tma_load_direct(self, stage_info: StageInfo) -> None: - """TMA load one full K tile using page offsets read directly from GMEM.""" + def tma_load_shared(self, stage_info: StageInfo, *, page_offset_stage): + """Read one warp-uniform SMEM quad and reuse it across FP8 feature slices.""" + cfg = self.cfg + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - cfg.load_k_warp_id + base = ( + cute.arch.make_warp_uniform(stage_info.stage_idx) * cfg.smem_k_stage_elems + ) + for owner in cutlass.range_constexpr( + cfg.tokens_per_k_cta // TMA_GATHER_ROWS // cfg.load_num_warps + ): + # All lanes read the same quad: shared-memory broadcast replaces + # four register shuffles. Retain the quad across feature slices. + quad = warp + Int32(owner * cfg.load_num_warps) + offset = ( + page_offset_stage * Int32(128) + + quad * Int32(4) + + cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + * Int32(64) + ) + raw = self.smem_sparse_offsets.data_ptr(offset).load(count=4, alignment=16) + gather4_uniform_quad_slices( + self.smem_k.data_ptr(base + quad * 4 * cfg.mma_qk_tiler_k), + self.tma_desc_c_latent, + self.tma_desc_c_rope, + Int32(0), + raw, + stage_info.barrier, + num_slices=cfg.iterations_qk_latent, + dst_stride=64 * cfg.mma_qk_tiler_k, + col_stride=cfg.mma_qk_tiler_k, + ) + + @producer_work + @cute.jit + def tma_load_paged(self, stage_info: StageInfo) -> None: + """Load a paged K tile; page-size-one inputs use the offset ring.""" cfg = self.cfg stage_idx = stage_info.stage_idx stage_base = stage_idx * cfg.smem_k_stage_elems @@ -1433,13 +1826,15 @@ def k_desc(self, stage_info: StageInfo, *, k_subtile_idx: cutlass.Constexpr[int] class SmemVResource(HighThroughputMlaResource): """SMEM V buffer for the FP8 split-MMA path. - Producer: LoadTma. Consumer: MmaPvTask. A single stage contains every V + Producer: LoadVTask. Consumer: MmaPvTask. A single stage contains every V sub-tile needed by the PV MMA for one logical K tile. """ smem_v: Any = None + smem_sparse_offsets: Any = None page_offsets: Any = None tma_desc_c_transpose: Any = None + tma_desc_extra_v: Any = None logical_seq_len_q: cutlass.Constexpr[int] = 1 cfg: cutlass.Constexpr = field(default_factory=MlaDecodeConfig) desc_v_base: cutlass.Constexpr[TaskLocalVariable] = ( @@ -1457,14 +1852,50 @@ def init_load_state(self, stage_info: StageInfo) -> None: @producer_work @cute.jit - def tma_load_direct(self, stage_info: StageInfo) -> None: - """TMA load one full V tile using page offsets read directly from GMEM.""" + def tma_load_shared(self, stage_info: StageInfo, *, page_offset_stage): + """Read one warp-uniform SMEM quad and reuse it across FP8 V panels.""" + cfg = self.cfg + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - cfg.load_v_warp_id + cta = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + base = ( + cute.arch.make_warp_uniform(stage_info.stage_idx) * cfg.smem_v_stage_elems + ) + for owner in cutlass.range_constexpr( + cfg.tokens_per_v_tile // TMA_GATHER_ROWS // cfg.load_num_warps + ): + # All lanes read the same quad: shared-memory broadcast replaces + # four register shuffles. Retain the quad across feature slices. + quad = warp + Int32(owner * cfg.load_num_warps) + offset = page_offset_stage * Int32(128) + quad * Int32(4) + raw = self.smem_sparse_offsets.data_ptr(offset).load(count=4, alignment=16) + call = (quad // Int32(16)) * cfg.iterations_pv_n + gather4_uniform_quad_slices( + self.smem_v.data_ptr( + base + + call * 128 * cfg.mma_pv_tiler[2] + + (quad % Int32(16)) * 4 * 128 + ), + self.tma_desc_c_transpose, + self.tma_desc_extra_v, + cta * 128, + raw, + stage_info.barrier, + num_slices=cfg.iterations_pv_n, + dst_stride=128 * cfg.mma_pv_tiler[2], + col_stride=256, + ) + + @producer_work + @cute.jit + def tma_load_paged(self, stage_info: StageInfo) -> None: + """Load a paged V tile; page-size-one inputs use the offset ring.""" cfg = self.cfg stage_idx = stage_info.stage_idx stage_base = stage_idx * cfg.smem_v_stage_elems work_tile = stage_info.work_tile blk_coord = work_tile.tile_idx k_index = work_tile.k_index_base + Int32(stage_info.loop_offset) + page_row_idx = blk_coord[2] page_offsets_batch = self.page_offsets[None, page_row_idx] @@ -1567,6 +1998,10 @@ def v_desc(self, stage_info: StageInfo, *, v_subtile_idx: cutlass.Constexpr[int] layout = cutlass.const_expr( qk_desc_layout_for_head_dim(cfg, cfg.mma_pv_tiler[2]) ) + if cutlass.const_expr(cfg.page_size == 1): + leading_byte_offset = 0 + stride_byte_offset = 1024 + layout = 2 svc_ptr = self.smem_v.data_ptr(subtile_offset) desc = Int64( prims.Tcgen05SmemDesc.build( @@ -1607,6 +2042,10 @@ def v_desc_n_major( layout = cutlass.const_expr( qk_desc_layout_for_head_dim(cfg, cfg.mma_pv_tiler[2]) ) + if cutlass.const_expr(cfg.page_size == 1): + leading_byte_offset = 0 + stride_byte_offset = 1024 + layout = 2 svc_ptr = self.smem_v.data_ptr(subtile_offset) desc = Int64( prims.Tcgen05SmemDesc.build( @@ -1639,6 +2078,7 @@ class TmemSResource(HighThroughputMlaResource): smem_p: Any = None # SMEM P array smem_exchange: Any = None # SMEM array for cross-warp max exchange softmax_scale_log2: Any = None # softmax_scale * log2(e) + page_offsets: Any = None cache_seqs: Any = None # per-batch valid K length cu_seqlens_q: Any = None # cumulative compact-Q offsets, or None for fixed Q split_kv: Any = None # per-work-tile split count @@ -1650,6 +2090,7 @@ class TmemSResource(HighThroughputMlaResource): is_leader: Any = field(init=False, default=None) row_max_state: Any = field(init=False, default=None) row_sum_state: Any = field(init=False, default=None) + request_lengths: Any = field(init=False, default=None) qk_acc_regs: cutlass.Constexpr[TaskLocalVariable] = ( TaskLocalVariable.uninitialized() ) @@ -1755,6 +2196,8 @@ def init_softmax_state(self, stage_info: StageInfo): self.is_leader = self.cta_rank == 0 self.row_max_state = Float32(-Float32.inf) self.row_sum_state = Float32(0) + # Allocate outside skipped-work guards; refresh values for every tile. + self.request_lengths = cutlass.Array(Int32, 2, space=cutlass.AddressSpace.rmem) return ( cutlass.Array( Float32, @@ -1803,6 +2246,15 @@ def init_softmax_state_odd(self, stage_info: StageInfo): Int32(0), ) + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def bind_sparse_request(self, stage_info: StageInfo): + """Cache live source lengths once per query tile, including graph replay.""" + if cutlass.const_expr(hasattr(self.page_offsets, "for_request")): + row = stage_info.work_tile.tile_idx[2] + self.request_lengths[0] = Int32(self.page_offsets.sl[row]) + self.request_lengths[1] = Int32(self.page_offsets.cl[row]) + @producer_work @cute.jit def qk_mma( @@ -1981,6 +2433,8 @@ def _load_s_impl( row_max_new, correction_factor_out, no_correction_out, + cached_mask_low=None, + cached_mask_high=None, ): """Load S from TMEM and compute the local row max.""" del row_sum_out, row_max_new, correction_factor_out, no_correction_out @@ -2040,6 +2494,46 @@ def _load_s_impl( ) qk_acc_regs.store(loaded, load_idx * TCGEN05_32B_REGS_PER_LOAD) + if cutlass.const_expr(cfg.page_size == 1): + group_needs_mask = True + if cutlass.const_expr(hasattr(self.page_offsets, "for_request")): + mask_page_offsets = self.page_offsets.for_request( + work_tile.tile_idx[2], self.request_lengths[0], self.request_lengths[1] + ) + else: + mask_page_offsets = self.page_offsets + if cutlass.const_expr(cached_mask_low is not None): + mask_routes = CachedSparseMask((cached_mask_low, cached_mask_high), False) + # The cached invalid bits already include the live K boundary. + # Whole valid halves need no per-score masking instructions. + if cutlass.const_expr(not needs_row_causal_mask): + group_needs_mask = (cached_mask_low | cached_mask_high) != Int32(0) + else: + raw_mask_routes = ( + mask_page_offsets.mask_for_tile(tile_offset_k, work_tile.tile_idx[2]) + if cutlass.const_expr(hasattr(self.page_offsets, "mask_for_tile")) + else self.page_offsets + ) + mask_routes = ( + cache_sparse_mask( + raw_mask_routes, + tile_offset_k, + K, + work_tile.tile_idx[2], + swaps=False, + column_half=local_tidx >> EPILOGUE_COLUMN_GROUP_SHIFT, + ) + if cutlass.const_expr( + cfg.page_size == 1 and self.logical_seq_len_q == 1 + ) + else raw_mask_routes + ) + if cutlass.const_expr( + cfg.page_size == 1 + and self.logical_seq_len_q == 1 + and not needs_row_causal_mask + ): + group_needs_mask = (mask_routes.low | mask_routes.high) != Int32(0) if group_needs_mask: if cutlass.const_expr(needs_row_causal_mask): # Clamp padded physical tail rows to the request's last real @@ -2070,8 +2564,21 @@ def _load_s_impl( token_idx - K + logical_seq_len_q ) token_is_visible = flat_query_row >= mask_flat_row + elif cutlass.const_expr( + cached_mask_low is not None + or getattr(self.page_offsets, "assume_valid_prefix", False) + ): + token_is_visible = not mask_routes.is_invalid( + token_idx, work_tile.tile_idx[2] + ) else: - token_is_visible = token_idx < K + token_is_visible = not invalid_sparse_token( + token_idx, + K, + mask_routes, + work_tile.tile_idx[2], + sparse=cfg.page_size == 1, + ) qk_acc_regs[i] = qk_acc_regs[i] if token_is_visible else neg_inf max0 = neg_inf @@ -2134,6 +2641,46 @@ def load_s( no_correction_out=no_correction_out, ) + @consumer_work( + returns=( + qk_acc_regs, + row_max, + row_sum, + row_sum_out, + row_max_new, + correction_factor_out, + no_correction_out, + ), + ) + @cute.jit + def load_s_cached( + self, + stage_info: StageInfo, + *, + qk_acc_regs, + row_max, + row_sum, + row_sum_out, + row_max_new, + correction_factor_out, + no_correction_out, + cached_mask_low, + cached_mask_high, + ): + """Load S for the even softmax group.""" + return self._load_s_impl( + stage_info, + qk_acc_regs=qk_acc_regs, + row_max=row_max, + row_sum=row_sum, + row_sum_out=row_sum_out, + row_max_new=row_max_new, + correction_factor_out=correction_factor_out, + no_correction_out=no_correction_out, + cached_mask_low=cached_mask_low, + cached_mask_high=cached_mask_high, + ) + @consumer_work( returns=( qk_acc_regs_odd, @@ -2264,8 +2811,8 @@ def load_peer_state(): fma_b = self.softmax_scale_log2 fma_c = Float32(0) - safe_row_max_new * self.softmax_scale_log2 if cutlass.const_expr(cfg.is_fp8_qkv()): - # Match the 448-scaled E4M3 P convention used by the reference - # output and the 1CTA implementation. + # 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( @@ -3034,8 +3581,10 @@ def rescale_o_slice( @dataclass(kw_only=True) class GmemOResource(HighThroughputMlaResource): - """GMEM output. Producer: Correction (epilogue store). No pipeline.""" + """Final output: direct stores or staged quantization with a TMA consumer.""" + smem_o: Any = None + tma_o_desc: Any = None output: Any = None partial_output: Any = None lse: Any = None @@ -3043,6 +3592,10 @@ class GmemOResource(HighThroughputMlaResource): 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 + atten_sinks: Any = None + output_quant: cutlass.Constexpr = None + output_quant_args: Any = None + lse_in_natural_log: cutlass.Constexpr[bool] = False softmax_scale_log2: Any = None smem_exchange: Any = None # SMEM for row_sum exchange (as Int32 base addr) split_kv: Any = None @@ -3051,6 +3604,66 @@ class GmemOResource(HighThroughputMlaResource): logical_seq_len_q: cutlass.Constexpr[int] = 1 cfg: cutlass.Constexpr = field(default_factory=MlaDecodeConfig) + @cute.jit + def _prefetch_output_rope(self, flat_row, batch_idx, local_tidx, head, valid): + """Fetch rotary coefficients while the nonrotary output is finalized.""" + tile_h = self.cfg.mma_pv_tiler[0] // self.cfg.num_mma_ctas + # One lane per logical query in this CTA; the first row may start + # partway through a head group. The other column group duplicates it. + if local_tidx < tile_h and valid and (head == 0 or local_tidx == 0): + row = Int64(flat_row) + if cutlass.const_expr(self.cu_seqlens_q is None): + row += ( + Int64(batch_idx) * self.logical_num_heads_q * self.logical_seq_len_q + ) + _, positions, cos_sin, _ = self.output_quant_args + position = positions[row // self.logical_num_heads_q] + cache_row = cos_sin.iterator.raw_ptr() + Int64(position) * 64 + prims.prefetch_l1(cache_row) + prims.prefetch_l1(cache_row + 32) + + @cute.jit + def _normalization_scale(self, row_sum, row_max, local_tidx, blk_coord): + safe_sum = row_sum if row_sum > Float32(0) else Float32(1) + scale = self.output_scale * cute.math.rcp(safe_sum, approx=True) + if cutlass.const_expr(self.atten_sinks is not None): + head = ( + Int32(blk_coord[1]) * self.cfg.mma_qk_tiler[0] + + Int32(blk_coord[0]) + * (self.cfg.mma_qk_tiler[0] // self.cfg.num_mma_ctas) + + (local_tidx & Int32(63)) + ) + if cutlass.const_expr(self.logical_seq_len_q == 1): + # Preserve the cheap single-query mapping, including padded heads. + head = cute.math.min(head, Int32(self.logical_num_heads_q - 1)) + else: + row = Int32(blk_coord[0]) * ( + self.cfg.mma_qk_tiler[0] // self.cfg.num_mma_ctas + ) + (local_tidx & Int32(63)) + _, head, _, _, _ = self._query_row_state( + row, Int32(blk_coord[1]), Int32(blk_coord[2]) + ) + sink = Float32(self.atten_sinks[head]) + scale = Float32(0) + if row_sum > Float32(0) and sink != Float32(Float32.inf): + max_log2 = row_max * self.softmax_scale_log2 + sink_log2 = sink * Float32(LOG2_E) + norm_max = cute.math.max(max_log2, sink_log2) + rescale = cute.math.exp2(max_log2 - norm_max, fastmath=True) + inv_quant = Float32( + fp8_quant_scale_rcp() if self.cfg.is_fp8_qkv() else 1.0 + ) + denominator = row_sum * inv_quant * rescale + cute.math.exp2( + sink_log2 - norm_max, fastmath=True + ) + scale = ( + self.output_scale + * inv_quant + * rescale + * cute.math.rcp(denominator, approx=True) + ) + return 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.""" @@ -3064,6 +3677,92 @@ def _query_row_state(self, row_in_tile, query_tile_idx, batch_idx): batch_idx, ) + @cute.jit + def _store_quantized_slice( + self, + values, + flat_row, + batch_idx, + column, + stage_idx=0, + ): + row = Int64(flat_row) + if cutlass.const_expr(self.cu_seqlens_q is None): + row += Int64(batch_idx) * self.logical_num_heads_q * self.logical_seq_len_q + shared_output = None + if cutlass.const_expr(self.cfg.use_tma_output): + shared_output = self.smem_o.data_ptr( + stage_idx * OUTPUT_TMA_STAGE_BYTES + ).toint(Int32) + self.output_quant.store( + self.output_quant_args, + values, + row, + column, + rotary=False, + shared_output=shared_output, + ) + if cutlass.const_expr(self.cfg.use_tma_output): + cute.arch.fence_view_async_shared() + + @consumer_work + @cute.jit + def tma_store_output( + self, stage_info: StageInfo, *, iter_n: cutlass.Constexpr[int] + ): + # Sparse launch flattens queries into batch rows; each CTA owns 64 heads. + head_tile_idx, _, token, _ = stage_info.work_tile.tile_idx + # Group-major output: the 64-head tile starts at this group. + heads_per_group = cutlass.const_expr( + self.output_quant.num_heads // self.output_quant.num_groups + ) + group = head_tile_idx * (64 // heads_per_group) + if prims.elect_sync(): + for half in cutlass.range_constexpr(2): + prims.cp_async_bulk_tensor_global_shared_cta( + self.tma_o_desc, + self.smem_o.subview( + stage_info.stage_idx * OUTPUT_TMA_STAGE_BYTES + half * 8192 + ), + (iter_n * 256 + half * 128, 0, token, group), + ) + prims.cp_async_bulk_commit_group() + # Release the stage only after TMA has consumed its shared source. + # The store task waits for global completion before retiring the CTA. + prims.cp_async_bulk_wait_group(0, read=True) + + @cute.jit + def _load_rotary_slice( + self, tmem_addr, norm_scale, flat_row, batch_idx, column, valid + ): + """Rotate the upper half before the other 64 accumulators become live.""" + row = Int64(flat_row) + if cutlass.const_expr(self.cu_seqlens_q is None): + row += Int64(batch_idx) * self.logical_num_heads_q * self.logical_seq_len_q + result = cutlass.Array(Float32, 128, space=cutlass.AddressSpace.rmem) + # The correction warp's register budget cannot keep 128 accumulators + # and a batch of RoPE coefficients live together without serial loads. + for half in cutlass.range_constexpr(2): + offset = cutlass.const_expr(64 if half == 0 else 0) + values = cutlass.Array(Float32, 64, space=cutlass.AddressSpace.rmem) + for chunk in cutlass.range_constexpr(2): + ptr = prims.make_tmem_ptr(tmem_addr + offset + chunk * 32, Float32) + values.store( + prims.tcgen05_ld(TCGEN05_32B_SHAPE, ptr, num=32), chunk * 32 + ) + for i in cutlass.range_constexpr(0, 64, 2): + scaled = mul_packed_f32x2( + (values[i], values[i + 1]), (norm_scale, norm_scale) + ) + values[i], values[i + 1] = scaled[0], scaled[1] + if cutlass.const_expr(half == 0): + if valid: + self.output_quant.rotate( + self.output_quant_args, values, row, column + offset + ) + result.store(values.load(0, 64), offset) + return result + @producer_work @cute.jit def epilogue_store(self, stage_info: StageInfo) -> None: @@ -3146,12 +3845,22 @@ def epilogue_store(self, stage_info: StageInfo) -> None: query_is_valid, ) = self._query_row_state(row_in_tile, seq_q_idx, batch_idx) + if cutlass.const_expr( + self.output_quant is not None and self.partial_output is None + ): + self._prefetch_output_rope( + storage_flat_query_row, + batch_idx, + local_tidx, + logical_head_idx, + row_in_tile < physical_tile_rows and query_is_valid, + ) + # Fully masked split rows can occur when physical tail rows or earlier # causal query rows have no visible K values in this split. Store zero # O and -inf LSE so split-KV reduction gives those rows zero weight. row_has_values = row_sum > Float32(0) - safe_row_sum = row_sum if row_has_values else Float32(1) - norm_scale = self.output_scale * cute.math.rcp(safe_row_sum, approx=True) + norm_scale = self._normalization_scale(row_sum, row_max, local_tidx, blk_coord) for iter_n in cutlass.range_constexpr(cfg.iterations_pv_n): # Load O from TMEM @@ -3159,22 +3868,36 @@ def epilogue_store(self, stage_info: StageInfo) -> None: tmem_raw_addr_n = tmem_raw_addr + ( iter_n * (cfg.mma_pv_tiler[1] // cfg.warps_in_n) ) - for load_idx in cutlass.range_constexpr(num_tmem_loads): - curr_addr = tmem_raw_addr_n + load_idx * TCGEN05_32B_REGS_PER_LOAD - tmem_ptr = prims.make_tmem_ptr(curr_addr, Float32) - loaded = prims.tcgen05_ld( - t2r_shape, tmem_ptr, num=TCGEN05_32B_REGS_PER_LOAD + if cutlass.const_expr( + self.output_quant is not None + and self.partial_output is None + and iter_n + 1 == cfg.iterations_pv_n + ): + qk_acc_regs = self._load_rotary_slice( + tmem_raw_addr_n, + norm_scale, + storage_flat_query_row, + batch_idx, + iter_n * tile_d + g_j, + row_in_tile < physical_tile_rows and query_is_valid, ) - qk_acc_regs.store(loaded, load_idx * TCGEN05_32B_REGS_PER_LOAD) + else: + for load_idx in cutlass.range_constexpr(num_tmem_loads): + curr_addr = tmem_raw_addr_n + load_idx * TCGEN05_32B_REGS_PER_LOAD + tmem_ptr = prims.make_tmem_ptr(curr_addr, Float32) + loaded = prims.tcgen05_ld( + t2r_shape, tmem_ptr, num=TCGEN05_32B_REGS_PER_LOAD + ) + qk_acc_regs.store(loaded, load_idx * TCGEN05_32B_REGS_PER_LOAD) - # Normalize: O = O * output_scale / row_sum - for i in cutlass.range_constexpr(0, 128, 2): - scaled = mul_packed_f32x2( - (qk_acc_regs[i], qk_acc_regs[i + 1]), - (norm_scale, norm_scale), - ) - qk_acc_regs[i] = scaled[0] - qk_acc_regs[i + 1] = scaled[1] + # Normalize: O = O * output_scale / row_sum + for i in cutlass.range_constexpr(0, 128, 2): + scaled = mul_packed_f32x2( + (qk_acc_regs[i], qk_acc_regs[i + 1]), + (norm_scale, norm_scale), + ) + qk_acc_regs[i] = scaled[0] + qk_acc_regs[i + 1] = scaled[1] # Store O to GMEM if row_in_tile < physical_tile_rows and query_is_valid: @@ -3219,6 +3942,13 @@ def epilogue_store(self, stage_info: StageInfo) -> None: vec_partial, evict="noallocate", ) + elif cutlass.const_expr(self.output_quant is not None): + self._store_quantized_slice( + qk_acc_regs, + storage_flat_query_row, + batch_idx, + iter_n * tile_d + g_j, + ) else: # 16-bit output (split_kv == 1, direct output) if cutlass.const_expr(self.cu_seqlens_q is not None): @@ -3303,6 +4033,8 @@ def epilogue_store(self, stage_info: StageInfo) -> None: else Float32(-Float32.inf) ) + if cutlass.const_expr(self.lse_in_natural_log): + lse *= Float32(LN_2) # Use local_tidx (0..127 within correction warpgroup) for LSE # indexing, not global tidx (which is 128..255 for correction warps). lse_tidx = local_tidx @@ -3386,24 +4118,8 @@ def epilogue_store_slice( t2r_shape = TCGEN05_32B_SHAPE num_tmem_loads = 4 qk_acc_regs = cutlass.Array(Float32, 128, space=cutlass.AddressSpace.rmem) - for load_idx in cutlass.range_constexpr(num_tmem_loads): - curr_addr = tmem_raw_addr + load_idx * TCGEN05_32B_REGS_PER_LOAD - tmem_ptr = prims.make_tmem_ptr(curr_addr, Float32) - loaded = prims.tcgen05_ld( - t2r_shape, tmem_ptr, num=TCGEN05_32B_REGS_PER_LOAD - ) - qk_acc_regs.store(loaded, load_idx * TCGEN05_32B_REGS_PER_LOAD) - row_has_values = row_sum > Float32(0) - safe_row_sum = row_sum if row_has_values else Float32(1) - norm_scale = self.output_scale * cute.math.rcp(safe_row_sum, approx=True) - for i in cutlass.range_constexpr(0, 128, 2): - scaled = mul_packed_f32x2( - (qk_acc_regs[i], qk_acc_regs[i + 1]), - (norm_scale, norm_scale), - ) - qk_acc_regs[i] = scaled[0] - qk_acc_regs[i + 1] = scaled[1] + norm_scale = self._normalization_scale(row_sum, row_max, local_tidx, blk_coord) tile_h = cfg.mma_pv_tiler[0] // cfg.num_mma_ctas tile_d = cfg.mma_pv_tiler[1] @@ -3427,6 +4143,48 @@ def epilogue_store_slice( query_is_valid, ) = self._query_row_state(row_in_tile, seq_q_idx, batch_idx) + if cutlass.const_expr( + self.output_quant is not None + and self.partial_output is None + and iter_n == 0 + ): + self._prefetch_output_rope( + storage_flat_query_row, + batch_idx, + local_tidx, + logical_head_idx, + row_in_tile < physical_tile_rows and query_is_valid, + ) + + if cutlass.const_expr( + self.output_quant is not None + and self.partial_output is None + and iter_n + 1 == cfg.iterations_pv_n + ): + qk_acc_regs = self._load_rotary_slice( + tmem_raw_addr, + norm_scale, + storage_flat_query_row, + batch_idx, + iter_n * tile_d + g_j, + row_in_tile < physical_tile_rows and query_is_valid, + ) + else: + for load_idx in cutlass.range_constexpr(num_tmem_loads): + curr_addr = tmem_raw_addr + load_idx * TCGEN05_32B_REGS_PER_LOAD + tmem_ptr = prims.make_tmem_ptr(curr_addr, Float32) + loaded = prims.tcgen05_ld( + t2r_shape, tmem_ptr, num=TCGEN05_32B_REGS_PER_LOAD + ) + qk_acc_regs.store(loaded, load_idx * TCGEN05_32B_REGS_PER_LOAD) + for i in cutlass.range_constexpr(0, 128, 2): + scaled = mul_packed_f32x2( + (qk_acc_regs[i], qk_acc_regs[i + 1]), + (norm_scale, norm_scale), + ) + qk_acc_regs[i] = scaled[0] + qk_acc_regs[i + 1] = scaled[1] + if row_in_tile < physical_tile_rows and query_is_valid: if cutlass.const_expr(self.partial_output is not None): S_q = ( @@ -3465,6 +4223,16 @@ def epilogue_store_slice( vec_partial, evict="noallocate", ) + elif cutlass.const_expr(self.output_quant is not None): + self._store_quantized_slice( + qk_acc_regs, + storage_flat_query_row, + batch_idx, + iter_n * tile_d + g_j, + stage_idx=stage_info.stage_idx + if cutlass.const_expr(cfg.use_tma_output) + else 0, + ) else: if cutlass.const_expr(self.cu_seqlens_q is not None): o_base_ptr = self.output.iterator.raw_ptr() + Int64( @@ -3543,6 +4311,8 @@ def epilogue_store_slice( if row_has_values else Float32(-Float32.inf) ) + if cutlass.const_expr(self.lse_in_natural_log): + lse *= Float32(LN_2) lse_tidx = local_tidx if lse_tidx < tile_h: lse_row_in_tile = head_tile_idx * tile_h + lse_tidx diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/tasks.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/tasks.py index 0ae3fa0f8534..0a98fca8ea6a 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/tasks.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/tasks.py @@ -42,6 +42,7 @@ import cutlass import cutlass.cute as cute +from cutlass.experimental import primitives as prims from cutlass.experimental.task_scheduling.memory import ResourceContext from cutlass.experimental.task_scheduling.schedule_builder import ( @@ -468,30 +469,62 @@ def load_tma_schedule(page_offset_window, smem_q, smem_kv, work_queue=None): return task_class( src_resources=src, dst_resources=[smem_q, smem_kv], - warp_idx=9, - num_warps=1, + warp_idx=smem_kv.cfg.load_tma_warp_id, + num_warps=smem_kv.cfg.load_num_warps, schedule=captured_schedule, name="LoadTmaTask", **task_kwargs, ) +def create_sparse_offsets_task( + page_offsets, work_queue=None, task_class=MlaTask, **task_kwargs +): + loop_start, loop_end, loop_step = captured_loop_bounds(task_kwargs, 0) + + @schedule + def offsets_schedule(page_offsets, work_queue=None): + page_offsets.init_sparse_state() + page_offsets.bind_sparse_request() + with domain_loop(loop_start, loop_end, loop_step): + page_offsets.acquire() + page_offsets.load_sparse_offsets() + page_offsets.commit() + work_queue_tail(work_queue, advance_label="advance_tile") + + result = ( + offsets_schedule(page_offsets) + if work_queue is None + else offsets_schedule(page_offsets, work_queue) + ) + return task_class( + src_resources=[work_queue] if work_queue is not None else [], + dst_resources=[page_offsets], + warp_idx=9, + num_warps=1, + schedule=result, + name="SparseOffsetsTask", + **task_kwargs, + ) + + def create_load_k_task( smem_q: SmemQResource, smem_k: SmemKResource, + page_offsets=None, work_queue: MlaWorkQueue = None, task_class: type = MlaTask, **task_kwargs, ) -> Task: - """Create the FP8 Q/K TMA task (warp 9). + """Create the FP8 Q/K TMA task on its configured loader warps. Q is loaded once before the loop. K uses one whole-tile pipeline stage per - logical K tile and reads page offsets directly from GMEM. + logical K tile. Sparse indices are prefetched through the shared offset + ring; ordinary paged inputs read page offsets directly from GMEM. """ loop_start, loop_end, loop_step = captured_loop_bounds(task_kwargs, 0) - @schedule - def load_k_schedule(smem_q, smem_k, work_queue=None): + def load_k_body(smem_q, smem_k, page_offsets=None, work_queue=None): """Load Q once and then publish one K stage per loop tile.""" smem_q.init_load_state() smem_k.init_load_state() @@ -500,23 +533,50 @@ def load_k_schedule(smem_q, smem_k, work_queue=None): smem_q.commit() with domain_loop(loop_start, loop_end, loop_step): - smem_k.acquire() - smem_k.tma_load_direct() - smem_k.commit() + if page_offsets is not None: + page_offsets.wait() + page_stage = page_offsets.sparse_stage() + smem_k.acquire() + smem_k.tma_load_shared(page_offset_stage=page_stage) + smem_k.commit() + # The last quad has been copied into TMA operands. Only now + # may the metadata producer overwrite this shared ring slot. + page_offsets.release() + else: + smem_k.acquire() + smem_k.tma_load_paged() + smem_k.commit() work_queue_tail(work_queue, advance_label="advance_tile") - schedule_result = ( - load_k_schedule(smem_q, smem_k) - if work_queue is None - else load_k_schedule(smem_q, smem_k, work_queue) - ) + @schedule + def load_k_schedule(smem_q, smem_k, page_offsets, work_queue=None): + load_k_body(smem_q, smem_k, page_offsets, work_queue) - src = [work_queue] if work_queue is not None else [] + @schedule + def load_k_dense_schedule(smem_q, smem_k, work_queue=None): + load_k_body(smem_q, smem_k, None, work_queue) + + if page_offsets is None: + schedule_result = ( + load_k_dense_schedule(smem_q, smem_k) + if work_queue is None + else load_k_dense_schedule(smem_q, smem_k, work_queue) + ) + else: + schedule_result = ( + load_k_schedule(smem_q, smem_k, page_offsets) + if work_queue is None + else load_k_schedule(smem_q, smem_k, page_offsets, work_queue) + ) + + src = ([page_offsets] if page_offsets is not None else []) + ( + [work_queue] if work_queue is not None else [] + ) return task_class( src_resources=src, dst_resources=[smem_q, smem_k], - warp_idx=9, - num_warps=1, + warp_idx=smem_k.cfg.load_k_warp_id, + num_warps=smem_k.cfg.load_num_warps, schedule=schedule_result, name="LoadKTask", **task_kwargs, @@ -525,35 +585,61 @@ def load_k_schedule(smem_q, smem_k, work_queue=None): def create_load_v_task( smem_v: SmemVResource, + page_offsets=None, work_queue: MlaWorkQueue = None, task_class: type = MlaTask, **task_kwargs, ) -> Task: - """Create the FP8 V TMA task (warp 10).""" + """Create the FP8 V TMA task on its configured loader warps.""" loop_start, loop_end, loop_step = captured_loop_bounds(task_kwargs, 0) - @schedule - def load_v_schedule(smem_v, work_queue=None): + def load_v_body(smem_v, page_offsets=None, work_queue=None): """Publish one V stage per logical K tile.""" smem_v.init_load_state() with domain_loop(loop_start, loop_end, loop_step): - smem_v.acquire() - smem_v.tma_load_direct() - smem_v.commit() + if page_offsets is not None: + page_offsets.wait() + page_stage = page_offsets.sparse_stage() + smem_v.acquire() + smem_v.tma_load_shared(page_offset_stage=page_stage) + smem_v.commit() + # Keep the ring slot protected through every shared read. + page_offsets.release() + else: + smem_v.acquire() + smem_v.tma_load_paged() + smem_v.commit() work_queue_tail(work_queue, advance_label="advance_tile") - schedule_result = ( - load_v_schedule(smem_v) - if work_queue is None - else load_v_schedule(smem_v, work_queue) - ) + @schedule + def load_v_schedule(smem_v, page_offsets, work_queue=None): + load_v_body(smem_v, page_offsets, work_queue) - src = [work_queue] if work_queue is not None else [] + @schedule + def load_v_dense_schedule(smem_v, work_queue=None): + load_v_body(smem_v, None, work_queue) + + if page_offsets is None: + schedule_result = ( + load_v_dense_schedule(smem_v) + if work_queue is None + else load_v_dense_schedule(smem_v, work_queue) + ) + else: + schedule_result = ( + load_v_schedule(smem_v, page_offsets) + if work_queue is None + else load_v_schedule(smem_v, page_offsets, work_queue) + ) + + src = ([page_offsets] if page_offsets is not None else []) + ( + [work_queue] if work_queue is not None else [] + ) return task_class( src_resources=src, dst_resources=[smem_v], - warp_idx=10, - num_warps=1, + warp_idx=smem_v.cfg.load_v_warp_id, + num_warps=smem_v.cfg.load_num_warps, schedule=schedule_result, name="LoadVTask", **task_kwargs, @@ -897,6 +983,7 @@ def create_softmax_task( smem_p: SmemPResource, work_queue: MlaWorkQueue = None, task_class: type = MlaTask, + page_offsets=None, warp_idx: int = 0, name: str = "SoftmaxTask", softmax_group_id: int = 0, @@ -920,9 +1007,15 @@ def softmax_prelude(tmem_s): ) return init_softmax_state() - def softmax_body(tmem_s, tmem_corr, smem_p, softmax_state): + def softmax_body(tmem_s, tmem_corr, smem_p, softmax_state, page_offsets=None): """Consume S tiles, materialize P, and publish correction factors.""" - load_s = tmem_s.load_s_odd if softmax_group_id == 1 else tmem_s.load_s + load_s = ( + tmem_s.load_s_odd + if softmax_group_id == 1 + else tmem_s.load_s_cached + if page_offsets is not None + else tmem_s.load_s + ) finish_softmax = ( tmem_s.finish_softmax_odd if softmax_group_id == 1 @@ -947,7 +1040,17 @@ def softmax_body(tmem_s, tmem_corr, smem_p, softmax_state): no_correction_out, ) = softmax_state + tmem_s.bind_sparse_request() with domain_loop(loop_start, loop_end, loop_step): + mask_kwargs = {} + if page_offsets is not None: + page_offsets.wait() + mask_low, mask_high = page_offsets.read_sparse_mask() + page_offsets.release() + mask_kwargs = { + "cached_mask_low": mask_low, + "cached_mask_high": mask_high, + } tmem_s.wait() if softmax_group_id == 1: ( @@ -984,6 +1087,7 @@ def softmax_body(tmem_s, tmem_corr, smem_p, softmax_state): row_max_new=row_max_new, correction_factor_out=correction_factor_out, no_correction_out=no_correction_out, + **mask_kwargs, ) tmem_s.release() if softmax_group_id == 1: @@ -1057,29 +1161,44 @@ def softmax_body(tmem_s, tmem_corr, smem_p, softmax_state): ) tmem_corr.commit() - @schedule - def softmax_schedule(tmem_s, tmem_corr, smem_p, work_queue=None): - """Capture one active softmax tile and unconditional queue progress.""" - + def capture_softmax(tmem_s, tmem_corr, smem_p, work_queue, page_offsets): _capture_clc_work_tile_body( work_queue, lambda softmax_state: softmax_body( - tmem_s, - tmem_corr, - smem_p, - softmax_state, + tmem_s, tmem_corr, smem_p, softmax_state, page_offsets ), lambda: softmax_prelude(tmem_s), use_clc_dynamic=use_clc_dynamic, ) - schedule_result = ( - softmax_schedule(tmem_s, tmem_corr, smem_p) - if work_queue is None - else softmax_schedule(tmem_s, tmem_corr, smem_p, work_queue) - ) + @schedule + def softmax_schedule(tmem_s, tmem_corr, smem_p, work_queue=None): + capture_softmax(tmem_s, tmem_corr, smem_p, work_queue, None) + + @schedule + def sparse_softmax_schedule( + tmem_s, tmem_corr, smem_p, page_offsets, work_queue=None + ): + capture_softmax(tmem_s, tmem_corr, smem_p, work_queue, page_offsets) + + if page_offsets is None: + schedule_result = ( + softmax_schedule(tmem_s, tmem_corr, smem_p) + if work_queue is None + else softmax_schedule(tmem_s, tmem_corr, smem_p, work_queue) + ) + else: + schedule_result = ( + sparse_softmax_schedule(tmem_s, tmem_corr, smem_p, page_offsets) + if work_queue is None + else sparse_softmax_schedule( + tmem_s, tmem_corr, smem_p, page_offsets, work_queue + ) + ) src = [tmem_s] + if page_offsets is not None: + src.append(page_offsets) if work_queue is not None: src.append(work_queue) return task_class( @@ -1117,6 +1236,7 @@ def create_correction_task( """ loop_start, loop_end, loop_step = captured_loop_bounds(task_kwargs, 1) use_clc_dynamic = bool(work_queue is not None and work_queue.use_clc_dynamic) + use_tma_output = bool(gmem_o.cfg.use_tma_output) def correction_prelude(tmem_corr): """Create correction and epilogue register state before skip guards.""" @@ -1161,12 +1281,16 @@ def correction_body(tmem_corr, tmem_o, gmem_o, correction_state): tmem_corr.prepare_epilogue_slice_store() ) for iter_n in range(iterations_pv_n): + if use_tma_output: + gmem_o.acquire() tmem_o.wait() gmem_o.epilogue_store_slice( row_sum=epilogue_row_sum, row_max=epilogue_row_max, iter_n=iter_n, ) + if use_tma_output: + gmem_o.commit() tmem_o.release() else: tmem_o.wait() @@ -1297,3 +1421,42 @@ def scheduler_schedule(work_queue, work_throttle=None): run_only_on_cta_id=0, **task_kwargs, ) + + +class MlaOutputTmaTask(MlaTask): + """Drain global TMA writes once after the persistent work loop.""" + + @cute.jit + def _drain_mla_work_tile_tails(self): + MlaTask._drain_mla_work_tile_tails(self) + prims.cp_async_bulk_wait_group(0) + + +def create_output_tma_task(output, work_queue=None, **task_kwargs): + """Consume two quantized N slices on the otherwise idle sparse FP8 warp.""" + + @schedule + def output_schedule(output, work_queue=None): + # This task only has epilogue work; declare its empty K loop. + with domain_loop(0, 0, 1): + pass + for iter_n in range(2): + output.wait() + output.tma_store_output(iter_n=iter_n) + output.release() + work_queue_tail(work_queue, advance_label="advance_tile") + + captured = ( + output_schedule(output) + if work_queue is None + else output_schedule(output, work_queue) + ) + return MlaOutputTmaTask( + src_resources=[output] + ([work_queue] if work_queue is not None else []), + dst_resources=[], + warp_idx=10, + num_warps=1, + schedule=captured, + name="OutputTmaTask", + **task_kwargs, + ) 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 dedc950131d2..73aef9cf8bc3 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 @@ -26,7 +26,7 @@ from __future__ import annotations -from dataclasses import dataclass, replace +from dataclasses import dataclass, replace as dataclass_replace, replace from math import ceil from typing import Any @@ -165,6 +165,16 @@ class MlaConfig: correction_num_warps: int = 4 mma_num_warps: int = 1 load_num_warps: int = 1 + sparse_reuse_kv: bool = False + paired_sparse_correction: bool = False + defer_sparse_max_update: bool = False + cache_uniform_sparse_quads: bool = False + fuse_sparse_epilogue: bool = False + fuse_sparse_cluster_epilogue: bool = False + fuse_sparse_reduction: bool = False + sparse_direct: bool = False + sparse_static_scales: bool = False + sparse_direct_capacities: tuple[int, int] = (0, 0) page_offsets_num_warps: int = 1 scheduler_num_warps: int = 1 clc_padding_warp_idx: int = 14 @@ -193,6 +203,47 @@ class MlaConfig: use_sliding_window_causal: int = 0 attention_window_size: int = 0 + @property + def one_insts_kv_swap(self) -> bool: + """One KV stream with swapped MMA operands (P remains in SMEM).""" + return self.num_insts_kv == 1 and self.kernel_variant == "swaps_mma_ab" + + @property + def single_kv_pipe(self) -> bool: + """Use one softmax/P stream while retaining two score stages.""" + return self.num_insts_kv == 1 + + @property + def short_merged_softmax(self) -> bool: + """Process two independent KV streams with one softmax warp group. + + Each planned split fits one two-tile group, so no online-correction + loop runs. Keep both streams' final statistics and PV accumulators; + live lengths still mask partial/empty tiles. This is num_insts_kv=2, + not the one-stream schedule selected by one_insts_kv_swap. + """ + return ( + self.sparse_direct + and not self.one_insts_kv_swap + and self.kernel_variant == "swaps_mma_ab" + and self.fixed_sparse_kv_tiles > 0 + ) + + @property + def serial_swap_reuse(self): + return ( + self.sparse_reuse_kv + and self.num_insts_kv == 2 + and not self.is_fp8_qkv() + and self.kernel_variant == "swaps_mma_ab" + ) + + @property + def sparse_cache_words(self): + if self.cache_uniform_sparse_quads: + return 6 + self.tile_size_kv // self.load_num_warps + return 6 + @property def softmax0_num_warps(self) -> int: return self.softmax_num_warps @@ -369,6 +420,21 @@ def is_fp8_qkv(self) -> bool: return self.qkv_dtype == "e4m3" + @property + def fixed_sparse_kv_tiles(self) -> int: + """Return the minimum instruction group when bounded by the plan. + + Keep-AB uses one KV tile and Swap-AB uses two independent streams. + If the maximum planned split fits that group, every shorter valid + prefix already has the same partition. Live lengths still mask data. + """ + tiles = self.local_kv_tiles(self.total_kv_tiles) + return ( + self.num_insts_kv + if self.sparse_direct and tiles == self.num_insts_kv + else 0 + ) + def local_kv_tiles(self, total_kv_tiles: int) -> int: """Return per-CTA KV tiles after multi-CTA KV splitting.""" if self.use_multi_ctas_kv != 1: @@ -1233,10 +1299,13 @@ def is_throughput_latency_mla_supported_shape( del batch_size return ( latent_dim == 512 - and rope_dim == 64 + and rope_dim in (0, 64) and 1 <= num_heads_q <= 128 and is_power_of_two(num_heads_q) - and num_tokens_per_page in SUPPORTED_MLA_PAGE_SIZES + and ( + num_tokens_per_page in SUPPORTED_MLA_PAGE_SIZES + or (num_tokens_per_page == 1 and rope_dim == 0) + ) and seq_len_q >= 1 and seq_len_kv >= 128 ) @@ -1479,7 +1548,9 @@ def make_throughput_latency_mla_config( raise ValueError("latent_dim must be positive") if rope_dim < 0: raise ValueError("rope_dim must be non-negative") - if num_tokens_per_page not in SUPPORTED_MLA_PAGE_SIZES: + if num_tokens_per_page not in SUPPORTED_MLA_PAGE_SIZES and not ( + num_tokens_per_page == 1 and rope_dim == 0 + ): raise ValueError( "num_tokens_per_page must be one of " f"{SUPPORTED_MLA_PAGE_SIZES}, got {num_tokens_per_page}" @@ -1491,7 +1562,10 @@ def make_throughput_latency_mla_config( f"num_tokens_per_page={num_tokens_per_page}" ) pages_per_kv_tile = MlaConfig.tile_size_kv // num_tokens_per_page - if pages_per_kv_tile > MlaConfig.page_offsets_entries_per_stage: + metadata_capacity = ( + 128 if num_tokens_per_page == 1 else MlaConfig.page_offsets_entries_per_stage + ) + if pages_per_kv_tile > metadata_capacity: raise ValueError( "page-offset staging capacity is smaller than one KV tile: " f"pages_per_kv_tile={pages_per_kv_tile}, " @@ -1528,6 +1602,11 @@ def make_throughput_latency_mla_config( selected_profile, explicit_persistent, ) + if selected_profile.use_persistent_scheduler and selected_profile.use_multi_ctas_kv: + # Persistent work queues enumerate Q/head-dimension tiles, not KV + # splits. Direct MlaProfile callers must obey the same restriction + # as persistent_override_profile, or only split zero gets produced. + raise ValueError("persistent scheduling is not supported with split-KV") tile_size_q = tile_size_q_for_profile( selected_profile, num_heads_q, seq_len_q, tile_size_q ) @@ -1612,6 +1691,7 @@ def make_throughput_latency_mla_config( softmax_regs=softmax_register_budget(tile_size_q), correction_regs=correction_register_budget(tile_size_q), num_tokens_per_page=num_tokens_per_page, + page_offsets_entries_per_stage=metadata_capacity, max_num_pages_per_seq_kv=max(1, ceil(seq_len_kv / num_tokens_per_page)), tmem_s_cols=tile_size_q, tmem_stats_cols=MlaConfig.tmem_stats_cols, @@ -1635,3 +1715,196 @@ def make_throughput_latency_mla_config( reduction_mode=reduction_mode, ) return cfg + + +def configure_sparse_mla( + cfg: MlaConfig, + profile, + *, + finalize_output=False, + direct_inputs=False, + static_scales=False, + source_capacities=(0, 0), +) -> MlaConfig: + """Resolve one selected sparse profile without order-dependent overrides. + + Policy selects the schedule. This builder checks its implementation + constraints, fixes topology, then assigns buffers and warp/register roles. + """ + keep = cfg.kernel_variant == "keeps_mma_ab" + issuers = profile.gather_issue_warps + reuse = profile.reuse_kv + uniform = profile.uniform_offset_cache + single = profile.single_kv_stream + if cfg.num_tokens_per_page != 1 or cfg.rope_dim != 0: + raise ValueError("sparse profiles require native page-size-one MLA") + if issuers not in (1, 2, 4, 8): + raise ValueError("sparse profiles require one, two, four or eight issuers") + if (keep and cfg.tile_size_q != 64) or ( + not keep and cfg.tile_size_q not in (8, 16, 32) + ): + raise ValueError("unsupported sparse query tile") + if single and ( + keep + or cfg.tile_size_q not in (8, 16) + or cfg.num_ctas_per_head_dim != 1 + or issuers != 4 + or cfg.use_cluster_reduction + or (cfg.num_ctas_per_seq_kv > 1 and cfg.use_persistent_scheduler) + ): + raise ValueError( + "one-instance Swap-AB requires M8/M16, V1, four issuers and no cluster reduction" + ) + if reuse and cfg.num_ctas_per_head_dim != 1: + raise ValueError("retained KV requires one V owner") + if profile.kv_tile_size not in (64, 128) or ( + profile.kv_tile_size == 64 and (not keep or cfg.is_fp8_qkv() or not reuse) + ): + raise ValueError("BK64 requires native BF16 Keep-AB reuse") + if uniform and ( + not reuse + or issuers not in (4, 8) + or (cfg.is_fp8_qkv() and not keep) + or (issuers == 8 and not keep) + ): + raise ValueError("uniform coordinate caching requires a qualified reuse layout") + if profile.balanced_registers and (keep or cfg.tile_size_q not in (8, 16)): + raise ValueError("balanced registers require Swap-AB M8/M16") + if profile.defer_max_update and cfg.is_fp8_qkv(): + raise ValueError("deferred maxima require native BF16 probabilities") + if profile.paired_correction and not keep: + raise ValueError("paired correction requires Keep-AB") + if profile.page_pipeline_stages and (profile.page_pipeline_stages < 1 or not keep): + raise ValueError("page-stage overrides require positive stages and Keep-AB") + + # Final stream count and tile geometry must precede short-loop eligibility. + splits = cfg.num_ctas_per_seq_kv + cluster_finish = finalize_output and bool(cfg.use_cluster_reduction) + direct_finish = finalize_output and splits == 1 + separate_finish = finalize_output and splits > 1 and not cluster_finish + if cluster_finish and ( + keep or cfg.tile_size_q not in (16, 32) or not cfg.use_bf16_output + ): + raise ValueError("cluster finishing requires native Swap-AB M16/M32") + if direct_finish and cfg.logical_seq_len_q != 1: + raise ValueError("fused output requires one logical query per virtual request") + if direct_inputs and (not finalize_output or cfg.qkv_dtype not in ("bf16", "e4m3")): + raise ValueError("separate sparse indices require native BF16/FP8 final output") + cfg = dataclass_replace( + cfg, + num_insts_kv=1 if keep or single else 2, + tile_size_kv=profile.kv_tile_size, + page_offsets_entries_per_stage=profile.kv_tile_size, + sparse_reuse_kv=reuse, + cache_uniform_sparse_quads=uniform, + paired_sparse_correction=profile.paired_correction, + defer_sparse_max_update=profile.defer_max_update, + sparse_direct=direct_inputs, + sparse_static_scales=static_scales, + sparse_direct_capacities=source_capacities if direct_inputs else (0, 0), + fuse_sparse_epilogue=direct_finish, + fuse_sparse_cluster_epilogue=cluster_finish, + fuse_sparse_reduction=separate_finish, + load_num_warps=issuers, + ) + short = cfg.short_merged_softmax + if ( + issuers == 8 + and not keep + and ( + cfg.tile_size_q not in (8, 16) + or cfg.local_kv_tiles(cfg.total_kv_tiles) <= cfg.num_insts_kv + ) + ): + raise ValueError("eight Swap-AB issuers require a steady-state loop") + + # Four D128 slices make one D512 tile. FP8/BK64 reuse keeps two tiles; + # BF16 BK128 retains one, with additional stages used for prefetch. + kv_stages = ( + profile.reuse_kv_stages if reuse and profile.reuse_kv_stages else cfg.kv_stages + ) + if reuse and profile.reuse_kv_stages: + retained_tiles = 2 if cfg.is_fp8_qkv() or cfg.tile_size_kv == 64 else 1 + if kv_stages < retained_tiles * cfg.qk_head_dim_stages: + raise ValueError("insufficient stages to retain KV through PV") + if cfg.tile_size_kv == 64 and profile.reuse_kv_stages < 2 * cfg.qk_head_dim_stages: + raise ValueError("BK64 requires two retained K tiles") + page_stages = cfg.page_offsets_stages + if ( + reuse + and profile.reuse_kv_stages + and keep + and (kv_stages == (12 if cfg.is_fp8_qkv() else 5) or cfg.tile_size_kv == 64) + ): + page_stages = 1 + if profile.page_pipeline_stages: + page_stages = profile.page_pipeline_stages + if single and reuse and kv_stages >= 6 and not cfg.is_fp8_qkv(): + page_stages = 1 + + # Select the final layout once. Sharing a softmax group in a short split + # preserves two numerical streams; the one-instance layout is distinct. + if single: + roles = dict( + threads_per_cta=512, + correction_warp_idx=4, + mma_warp_idx=8, + load_warp_idx=12, + page_offsets_warp_idx=10, + scheduler_warp_idx=11, + softmax_regs=128, + correction_regs=160, + mma_load_regs=64, + scheduler_regs=64, + o_stages=1, + ) + elif short: + roles = dict( + threads_per_cta=512, + correction_warp_idx=4, + mma_warp_idx=8, + load_warp_idx=12, + page_offsets_warp_idx=9, + scheduler_warp_idx=10, + softmax_regs=96, + correction_regs=96, + mma_load_regs=96, + scheduler_regs=96, + ) + elif keep: + roles = dict( + threads_per_cta=640 if issuers == 8 else 512, + load_warp_idx=12, + softmax_regs=128 if uniform or issuers == 8 else 192, + correction_regs=160 if issuers == 8 else 192, + mma_load_regs=96 if uniform and issuers == 4 else 64, + ) + else: + balanced = profile.balanced_registers + roles = dict( + threads_per_cta=768 if issuers == 8 else 640, + load_warp_idx=16, + softmax_regs=96 + if balanced + else 128 + if uniform + else 144 + if cfg.tile_size_q == 32 + else 160, + correction_regs=96 if balanced else 192 if uniform else cfg.correction_regs, + mma_load_regs=(64 if issuers == 8 else 96) + if balanced + else 96 + if uniform + else 48, + scheduler_regs=(64 if issuers == 8 else 96) + if balanced + else cfg.scheduler_regs, + ) + return dataclass_replace( + cfg, + **roles, + kv_stages=kv_stages, + page_offsets_stages=page_stages, + q_stages=1 if single or (reuse and profile.reuse_kv_stages) else cfg.q_stages, + ) 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 cab0923d5959..f7c476f7a3c9 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 @@ -55,6 +55,7 @@ from .config import ( MlaConfig, make_throughput_latency_mla_config, + configure_sparse_mla, ) from .resources import ( SmemKvResource, @@ -74,7 +75,7 @@ MlaDecodeTask, create_throughput_latency_correction_task, create_keeps_mma_ab_correction_task, - create_keeps_mma_ab_mma_task, + create_single_kv_pipe_mma_task, create_keeps_mma_ab_softmax_task, create_load_page_offsets_task, create_padding_task, @@ -94,7 +95,7 @@ supports_parallel_gmem_reduction, ) from .reduction import gmem_reduction_launch_shape, run_gmem_reduction_kernel -from ..helpers.constants import TMEM_LIFECYCLE_BARRIER_ID +from ..helpers.constants import LOG2_E, TMEM_LIFECYCLE_BARRIER_ID, TMEM_READY_BARRIER_ID from ..helpers.mask import MaskType, normalize_mask_type from ..helpers.tile import ( runtime_query_tile_is_active, @@ -109,10 +110,6 @@ ) -# Softmax uses exp2, so natural-scale scores are multiplied by log2(e). -LOG2_E = 1.4426950408889634 - - @cute.jit def _publish_neutral_standalone_partial( cfg, @@ -331,6 +328,25 @@ def __init__( self.cache_seqs = cache_seqs self.cu_seqlens_q = cu_seqlens_q self.enable_runtime_skip = cu_seqlens_q is not None + self._clc_response_alloc = SmemAllocation( + self.name + "_clc_response", dtype=cutlass.Int128, count=2, alignment=16 + ) + self._response_allocator = None + + def get_smem_requirements(self): + return [self._clc_response_alloc] + + def create_tile_scheduler(self): + response = self._response_allocator.get(self._clc_response_alloc) + response_ptr = cute.make_ptr( + cutlass.Int128, response.data_ptr(), mem_space=cutlass.AddressSpace.smem + ) + return utils.ClcDynamicPersistentTileScheduler.create( + self.tile_scheduler_config.tile_scheduler_params, + cute.arch.block_idx(), + cute.arch.grid_dim(), + response_ptr, + ) @cute.jit def skip_work_tile_if(self, work_tile: WorkTileInfo): @@ -353,6 +369,19 @@ def _default_scales(scale_softmax_log2, output_scale): return scale_softmax_log2, output_scale +def _atten_sinks_view(cfg, epilogue_params): + """Bind only sinks; packed scale/count metadata remains a separate input.""" + if epilogue_params is None or not ( + cfg.fuse_sparse_epilogue or cfg.fuse_sparse_cluster_epilogue + ): + return None + if cfg.sparse_direct: + return epilogue_params + return cute.make_tensor( + epilogue_params.iterator + 2, cute.make_layout(cfg.logical_num_heads_q) + ) + + def _check_persistent_scheduler_modes( use_clc_dynamic_scheduler: bool, use_static_persistent_scheduler: bool, @@ -378,7 +407,7 @@ def _make_static_work_queue( tile_scheduler_params=tile_sched_params, ), cfg=cfg, - batch_size=cute.size(cache_seqs), + batch_size=cute.size(cache_seqs.shape), cache_seqs=cache_seqs, cu_seqlens_q=cu_seqlens_q, name=name, @@ -457,6 +486,9 @@ def build_throughput_latency_mla_task_manager( cta_idx_head_dim_v=None, scale_softmax_log2=None, output_scale=None, + sparse_epilogue_params=None, + output_quant=None, + output_quant_args=None, o_tensor=None, lse_tensor=None, acc_o_tensor=None, @@ -485,8 +517,8 @@ def build_throughput_latency_mla_task_manager( use_clc_dynamic_scheduler, use_static_persistent_scheduler, ) - if cfg.kernel_variant == "keeps_mma_ab": - return _make_keeps_mma_ab_task_graph( + if cfg.single_kv_pipe: + return _make_single_kv_pipe_task_graph( cfg, total_kv_tiles=total_kv_tiles, use_page_offsets=use_page_offsets, @@ -506,6 +538,9 @@ def build_throughput_latency_mla_task_manager( cta_idx_head_dim_v=cta_idx_head_dim_v, scale_softmax_log2=scale_softmax_log2, output_scale=output_scale, + sparse_epilogue_params=sparse_epilogue_params, + output_quant=output_quant, + output_quant_args=output_quant_args, o_tensor=o_tensor, lse_tensor=lse_tensor, acc_o_tensor=acc_o_tensor, @@ -522,7 +557,7 @@ def build_throughput_latency_mla_task_manager( cta_layout_vmnk = (1, 1, 1, 1) tma_group = pipeline.CooperativeGroup(agent.Thread) umma_group = pipeline.CooperativeGroup(agent.Thread) - load_group = pipeline.CooperativeGroup(agent.Thread, 32) + load_group = pipeline.CooperativeGroup(agent.Thread, cfg.load_num_warps * 32) page_group = pipeline.CooperativeGroup(agent.Thread, 32) scheduler_group = pipeline.CooperativeGroup( agent.Thread, @@ -550,13 +585,27 @@ def build_throughput_latency_mla_task_manager( consumer_group=load_group, pipeline_type=PipelineType.AsyncAsync, cta_layout_vmnk=cta_layout_vmnk, - async_producer_op=pipeline.PipelineOp.AsyncLoad, + async_producer_op=( + pipeline.PipelineOp.AsyncThread + if cfg.sparse_direct + else pipeline.PipelineOp.AsyncLoad + ), advance_on_wait=True, ) smem_q_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( num_stages=cfg.q_stages, num_bytes=cfg.qk_smem_tile_bytes, + num_bytes_per_warp_per_cta=( + cfg.qk_smem_tile_bytes // cfg.load_num_warps + if cfg.load_num_warps > 1 + else None + ), producer_group=tma_group, + producer_signaling_threads=( + SignalingThreads.TaskWarpLeader + if cfg.load_num_warps > 1 + else SignalingThreads.All + ), consumer_group=umma_group, cta_layout_vmnk=cta_layout_vmnk, consumer_signaling_threads=SignalingThreads.CtaLeader, @@ -564,11 +613,23 @@ def build_throughput_latency_mla_task_manager( smem_kv_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( num_stages=cfg.kv_stages, num_bytes=cfg.kv_smem_tile_bytes, + num_bytes_per_warp_per_cta=( + cfg.kv_smem_tile_bytes // cfg.load_num_warps + if cfg.load_num_warps > 1 + else None + ), producer_group=tma_group, + producer_signaling_threads=( + SignalingThreads.TaskWarpLeader + if cfg.load_num_warps > 1 + else SignalingThreads.All + ), consumer_group=umma_group, cta_layout_vmnk=cta_layout_vmnk, consumer_signaling_threads=SignalingThreads.CtaLeader, ) + if cfg.sparse_reuse_kv: + smem_kv_cfg = dataclass_replace(smem_kv_cfg, advance_on_wait=True) tmem_s_cfg = PipelineConfig.create_umma_async_pipeline_cfg( num_stages=1, producer_group=umma_group, @@ -674,6 +735,7 @@ def build_throughput_latency_mla_task_manager( tma_desc_c_latent=tma_desc_c_latent, tma_desc_c_rope=tma_desc_c_rope, tma_desc_v=tma_desc_v, + tma_desc_bulk_extra=tma_desc_q_rope, c_rope_tensor=c_rope_tensor, page_offsets_kv=smem_page_offsets if use_page_offsets else None, page_offsets=page_offsets, @@ -699,6 +761,7 @@ def build_throughput_latency_mla_task_manager( cta_idx_kv=cta_idx_kv, sync_barrier_id=0, name="ll_mla_tmem_s0", + page_offsets=page_offsets, ) tmem_s1 = TmemSResource( cfg=cfg, @@ -713,6 +776,7 @@ def build_throughput_latency_mla_task_manager( cta_idx_kv=cta_idx_kv, sync_barrier_id=1, name="ll_mla_tmem_s1", + page_offsets=page_offsets, ) order_p01_alloc = SmemAllocation( name="ll_mla_order_p01", @@ -770,6 +834,10 @@ def build_throughput_latency_mla_task_manager( inst_id=0, scale_softmax_log2=scale_softmax_log2, output_scale=output_scale, + sparse_epilogue_params=sparse_epilogue_params, + atten_sinks=_atten_sinks_view(cfg, sparse_epilogue_params), + output_quant=output_quant, + output_quant_args=output_quant_args, o_tensor=o_tensor, lse_tensor=lse_tensor, acc_o_tensor=acc_o_tensor, @@ -788,6 +856,10 @@ def build_throughput_latency_mla_task_manager( inst_id=1, scale_softmax_log2=scale_softmax_log2, output_scale=output_scale, + sparse_epilogue_params=sparse_epilogue_params, + atten_sinks=_atten_sinks_view(cfg, sparse_epilogue_params), + output_quant=output_quant, + output_quant_args=output_quant_args, o_tensor=o_tensor, lse_tensor=lse_tensor, acc_o_tensor=acc_o_tensor, @@ -824,44 +896,51 @@ def build_throughput_latency_mla_task_manager( **task_domain_kwargs, ) ) + tasks.append( + create_throughput_latency_load_task( + smem_q, + smem_kv, + work_queue, + schedule_token_throttle, + cfg, + domain=load_domain, + smem_page_offsets=smem_page_offsets if use_page_offsets else None, + use_page_offsets=use_page_offsets, + task_class=MlaDecodeTask, + **task_domain_kwargs, + ) + ) + if not cfg.short_merged_softmax: + tasks.extend( + [ + create_throughput_latency_softmax0_task( + tmem_s0, + local0, + smem_p0, + global0, + work_queue, + cfg, + domain=softmax_domain, + task_class=MlaDecodeTask, + domain_bias=1, + **task_domain_kwargs, + ), + create_throughput_latency_softmax1_task( + tmem_s1, + local1, + smem_p1, + global1, + work_queue, + cfg, + domain=softmax_domain, + task_class=MlaDecodeTask, + domain_bias=1, + **task_domain_kwargs, + ), + ] + ) tasks.extend( [ - create_throughput_latency_load_task( - smem_q, - smem_kv, - work_queue, - schedule_token_throttle, - cfg, - domain=load_domain, - smem_page_offsets=smem_page_offsets if use_page_offsets else None, - use_page_offsets=use_page_offsets, - task_class=MlaDecodeTask, - **task_domain_kwargs, - ), - create_throughput_latency_softmax0_task( - tmem_s0, - local0, - smem_p0, - global0, - work_queue, - cfg, - domain=softmax_domain, - task_class=MlaDecodeTask, - domain_bias=1, - **task_domain_kwargs, - ), - create_throughput_latency_softmax1_task( - tmem_s1, - local1, - smem_p1, - global1, - work_queue, - cfg, - domain=softmax_domain, - task_class=MlaDecodeTask, - domain_bias=1, - **task_domain_kwargs, - ), create_throughput_latency_correction_task( local0, local1, @@ -918,6 +997,23 @@ def build_throughput_latency_mla_task_manager( ) ) + if cfg.load_warp_idx == 16: + idle = create_padding_task( + cfg, work_queue, warp_idx=15, num_warps=1, task_class=MlaDecodeTask + ) + idle.name = "FormerLoadWarpPadding" + tasks.append(idle) + if cfg.load_num_warps < 4: + idle = create_padding_task( + cfg, + work_queue, + warp_idx=16 + cfg.load_num_warps, + num_warps=4 - cfg.load_num_warps, + task_class=MlaDecodeTask, + ) + idle.name = "LoadGroupPadding" + tasks.append(idle) + deps = { smem_q: [], smem_kv: [], @@ -965,12 +1061,26 @@ def build_throughput_latency_mla_task_manager( deps[schedule_token_throttle] = [work_queue] dma_release_labels = { - (smem_kv, tmem_s0): {"k_desc_0"}, - (smem_kv, tmem_s1): {"k_desc_1"}, - (smem_kv, tmem_o): {"v_desc_0", "v_desc_1"}, + (smem_kv, tmem_s0): {"v_desc_reused" if cfg.sparse_reuse_kv else "k_desc_0"}, + (smem_kv, tmem_s1): {"v_desc_reused" if cfg.sparse_reuse_kv else "k_desc_1"}, + (smem_kv, tmem_o): ( + {"v_desc_reused"} if cfg.sparse_reuse_kv else {"v_desc_0", "v_desc_1"} + ), } + if cfg.serial_swap_reuse: + dma_release_labels = { + (smem_kv, tmem_s0): {"v_desc_reused_0"}, + (smem_kv, tmem_s1): {"v_desc_reused_1"}, + (smem_kv, tmem_o): {"v_desc_reused_0", "v_desc_reused_1"}, + } + smem_allocator = SmemAllocator() + if isinstance(work_queue, ThroughputLatencyMlaClcWorkQueue): + work_queue._response_allocator = smem_allocator + smem_allocator.add_resource(work_queue) + if schedule_token_throttle is not None: + smem_allocator.add_resource(schedule_token_throttle) smem_allocator.add_resource(smem_q) if use_page_offsets: smem_allocator.add_resource(smem_page_offsets) @@ -999,10 +1109,47 @@ def build_throughput_latency_mla_task_manager( tmem_allocator.add_resource(tmem_o) tmem_allocator.compute_layout() + if cfg.short_merged_softmax: + from .tasks import create_short_merged_softmax_task + + # Keep the two P/statistic streams, but create only their shared + # softmax worker. No steady-state correction exchanges are needed. + tasks.append( + create_short_merged_softmax_task( + tmem_s0, + tmem_s1, + local0, + local1, + smem_p0, + smem_p1, + global0, + global1, + cfg, + **task_domain_kwargs, + ) + ) + idle = create_padding_task( + cfg, work_queue, warp_idx=11, num_warps=1, task_class=MlaDecodeTask + ) + idle.name = "MergedSoftmaxPadding" + tasks.append(idle) + if cfg.load_num_warps < 4: + idle = create_padding_task( + cfg, + work_queue, + warp_idx=cfg.load_warp_idx + cfg.load_num_warps, + num_warps=4 - cfg.load_num_warps, + task_class=MlaDecodeTask, + ) + idle.name = "MergedLoadGroupPadding" + tasks.append(idle) task_manager = TaskManager( tasks=tasks, resource_dependency_graph=deps, dma_consumer_release_labels=dma_release_labels, + # DSL 4.7's bounded credit simulation can stop before delayed KV + # releases drain. Its warning mode still runs all schedule checks. + skip_validation=cfg.serial_swap_reuse, smem_allocator=smem_allocator, tmem_allocator=tmem_allocator, verbose=verbose, @@ -1011,7 +1158,7 @@ def build_throughput_latency_mla_task_manager( return task_manager, corr1 -def _make_keeps_mma_ab_task_graph( +def _make_single_kv_pipe_task_graph( cfg: MlaConfig, *, total_kv_tiles: int, @@ -1032,6 +1179,9 @@ def _make_keeps_mma_ab_task_graph( cta_idx_head_dim_v=None, scale_softmax_log2=None, output_scale=None, + sparse_epilogue_params=None, + output_quant=None, + output_quant_args=None, o_tensor=None, lse_tensor=None, acc_o_tensor=None, @@ -1043,10 +1193,10 @@ def _make_keeps_mma_ab_task_graph( verbose: bool = False, exhaustive_deadlock_race_check: bool = False, ) -> tuple[TaskManager, object]: - """Build the keeps-MMA-AB 1CTA task graph for the generic builder. + """Build the single-stream 1CTA task graph. - This variant uses one score/P pipe with P stored in TMEM. The swaps path - keeps its two SMEM-P pipes and is intentionally left untouched. + Keep stores P in TMEM; single-stream swap stores P in SMEM. Both use + one softmax group, two score stages, and one output accumulator. """ scale_softmax_log2, output_scale = _default_scales( @@ -1061,7 +1211,7 @@ def _make_keeps_mma_ab_task_graph( cta_layout_vmnk = (1, 1, 1, 1) tma_group = pipeline.CooperativeGroup(agent.Thread) umma_group = pipeline.CooperativeGroup(agent.Thread) - load_group = pipeline.CooperativeGroup(agent.Thread, 32) + load_group = pipeline.CooperativeGroup(agent.Thread, cfg.load_num_warps * 32) page_group = pipeline.CooperativeGroup(agent.Thread, 32) softmax_group = pipeline.CooperativeGroup( agent.Thread, @@ -1079,13 +1229,27 @@ def _make_keeps_mma_ab_task_graph( consumer_group=load_group, pipeline_type=PipelineType.AsyncAsync, cta_layout_vmnk=cta_layout_vmnk, - async_producer_op=pipeline.PipelineOp.AsyncLoad, + async_producer_op=( + pipeline.PipelineOp.AsyncThread + if cfg.sparse_direct + else pipeline.PipelineOp.AsyncLoad + ), advance_on_wait=True, ) smem_q_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( num_stages=cfg.q_stages, num_bytes=cfg.qk_smem_tile_bytes, + num_bytes_per_warp_per_cta=( + cfg.qk_smem_tile_bytes // cfg.load_num_warps + if cfg.load_num_warps > 1 + else None + ), producer_group=tma_group, + producer_signaling_threads=( + SignalingThreads.TaskWarpLeader + if cfg.load_num_warps > 1 + else SignalingThreads.All + ), consumer_group=umma_group, cta_layout_vmnk=cta_layout_vmnk, consumer_signaling_threads=SignalingThreads.CtaLeader, @@ -1093,11 +1257,23 @@ def _make_keeps_mma_ab_task_graph( smem_kv_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( num_stages=cfg.kv_stages, num_bytes=cfg.kv_smem_tile_bytes, + num_bytes_per_warp_per_cta=( + cfg.kv_smem_tile_bytes // cfg.load_num_warps + if cfg.load_num_warps > 1 + else None + ), producer_group=tma_group, + producer_signaling_threads=( + SignalingThreads.TaskWarpLeader + if cfg.load_num_warps > 1 + else SignalingThreads.All + ), consumer_group=umma_group, cta_layout_vmnk=cta_layout_vmnk, consumer_signaling_threads=SignalingThreads.CtaLeader, ) + if cfg.sparse_reuse_kv: + smem_kv_cfg = dataclass_replace(smem_kv_cfg, advance_on_wait=True) tmem_s_cfg = PipelineConfig.create_umma_async_pipeline_cfg( num_stages=2, producer_group=umma_group, @@ -1111,8 +1287,8 @@ def _make_keeps_mma_ab_task_graph( consumer_group=corr_group, cta_layout_vmnk=cta_layout_vmnk, ) - tmem_p_cfg = PipelineConfig.create_async_umma_pipeline_cfg( - num_stages=2, + p_pipeline_cfg = PipelineConfig.create_async_umma_pipeline_cfg( + num_stages=1 if cfg.one_insts_kv_swap else 2, producer_group=softmax_group, consumer_group=umma_group, cta_layout_vmnk=cta_layout_vmnk, @@ -1180,6 +1356,7 @@ def _make_keeps_mma_ab_task_graph( tma_desc_c_latent=tma_desc_c_latent, tma_desc_c_rope=tma_desc_c_rope, tma_desc_v=tma_desc_v, + tma_desc_bulk_extra=tma_desc_q_rope, c_rope_tensor=c_rope_tensor, page_offsets_kv=smem_page_offsets if use_page_offsets else None, page_offsets=page_offsets, @@ -1192,8 +1369,9 @@ def _make_keeps_mma_ab_task_graph( cta_idx_head_dim_v=cta_idx_head_dim_v, name="ll_mla_q64_smem_kv", ) - tmem_s = TmemSKeepsResource( + tmem_s = (TmemSResource if cfg.one_insts_kv_swap else TmemSKeepsResource)( cfg=cfg, + page_offsets=page_offsets, pipeline_config=tmem_s_cfg, inst_id=0, scale_softmax_log2=scale_softmax_log2, @@ -1206,26 +1384,47 @@ def _make_keeps_mma_ab_task_graph( sync_barrier_id=0, name="ll_mla_q64_tmem_s", ) - tmem_p = TmemPResource( - cfg=cfg, - pipeline_config=tmem_p_cfg, - inst_id=0, - scale_softmax_log2=scale_softmax_log2, - tmem_alias_ref=tmem_s, - name="ll_mla_q64_tmem_p", - ) - tmem_s.p_ref = tmem_p + # Both one- and two-instance Swap-AB read P as the SMEM B operand. + # This builder's other branch is one-instance Keep-AB: P is the TMEM A + # operand. Stream count alone does not determine probability storage. + if cfg.one_insts_kv_swap: + p_resource = SmemPResource( + cfg=cfg, + pipeline_config=p_pipeline_cfg, + inst_id=0, + scale_softmax_log2=scale_softmax_log2, + tmem_s_ref=tmem_s, + name="ll_mla_single_smem_p", + ) + else: + p_resource = TmemPResource( + cfg=cfg, + pipeline_config=p_pipeline_cfg, + inst_id=0, + scale_softmax_log2=scale_softmax_log2, + tmem_alias_ref=tmem_s, + name="ll_mla_q64_tmem_p", + ) + tmem_s.p_ref = p_resource tmem_o = TmemOResource( cfg=cfg, pipeline_config=tmem_o_cfg, - p_tmem_ref=tmem_p, + p_tmem_ref=None if cfg.one_insts_kv_swap else p_resource, + p0_ref=p_resource if cfg.one_insts_kv_swap else None, cache_seqs=cache_seqs, cu_seqlens_q=cu_seqlens_q, batch_idx=batch_idx, cta_idx_q=cta_idx_q, name="ll_mla_q64_tmem_o", ) - local = TmemSoftmaxLocalResource( + from .resources.tmem_softmax_stats import SmemSoftmaxLocalResource + + local_resource = ( + SmemSoftmaxLocalResource + if cfg.num_tokens_per_page == 1 + else TmemSoftmaxLocalResource + ) + local = local_resource( cfg=cfg, pipeline_config=local_cfg, inst_id=0, @@ -1240,6 +1439,10 @@ def _make_keeps_mma_ab_task_graph( inst_id=1, scale_softmax_log2=scale_softmax_log2, output_scale=output_scale, + sparse_epilogue_params=sparse_epilogue_params, + atten_sinks=_atten_sinks_view(cfg, sparse_epilogue_params), + output_quant=output_quant, + output_quant_args=output_quant_args, o_tensor=o_tensor, lse_tensor=lse_tensor, acc_o_tensor=acc_o_tensor, @@ -1287,10 +1490,14 @@ def _make_keeps_mma_ab_task_graph( task_class=MlaDecodeTask, **task_domain_kwargs, ), - create_keeps_mma_ab_softmax_task( + ( + create_throughput_latency_softmax0_task + if cfg.one_insts_kv_swap + else create_keeps_mma_ab_softmax_task + )( tmem_s, local, - tmem_p, + p_resource, global_softmax, work_queue, cfg, @@ -1309,11 +1516,11 @@ def _make_keeps_mma_ab_task_graph( task_class=MlaDecodeTask, **task_domain_kwargs, ), - create_keeps_mma_ab_mma_task( + create_single_kv_pipe_mma_task( smem_q, smem_kv, tmem_s, - tmem_p, + p_resource, tmem_o, work_queue, cfg, @@ -1353,14 +1560,31 @@ def _make_keeps_mma_ab_task_graph( ) ) + if cfg.load_warp_idx == 12: + idle = create_padding_task( + cfg, work_queue, warp_idx=9, num_warps=1, task_class=MlaDecodeTask + ) + idle.name = "FormerKeepLoadWarpPadding" + tasks.append(idle) + if cfg.load_num_warps < 4: + idle = create_padding_task( + cfg, + work_queue, + warp_idx=12 + cfg.load_num_warps, + num_warps=4 - cfg.load_num_warps, + task_class=MlaDecodeTask, + ) + idle.name = "KeepLoadGroupPadding" + tasks.append(idle) + deps = { smem_q: [], smem_kv: [], tmem_s: [smem_q, smem_kv], local: [tmem_s], global_softmax: [tmem_s], - tmem_p: [tmem_s], - tmem_o: [tmem_p, smem_kv], + p_resource: [tmem_s], + tmem_o: [p_resource, smem_kv], corr: [local, tmem_o], } if use_page_offsets: @@ -1373,8 +1597,8 @@ def _make_keeps_mma_ab_task_graph( tmem_s: [smem_q, smem_kv, work_queue], local: [tmem_s, work_queue], global_softmax: [tmem_s, work_queue], - tmem_p: [tmem_s, work_queue], - tmem_o: [tmem_p, smem_kv, work_queue], + p_resource: [tmem_s, work_queue], + tmem_o: [p_resource, smem_kv, work_queue], corr: [local, tmem_o, work_queue], work_queue: ( [work_queue, schedule_token_throttle] @@ -1389,15 +1613,32 @@ def _make_keeps_mma_ab_task_graph( deps[schedule_token_throttle] = [work_queue] dma_release_labels = { - (smem_kv, tmem_s): {"k_desc_0"}, - (smem_kv, tmem_o): {"v_desc_0"}, + (smem_kv, tmem_s): { + ( + ("v_desc_reused_0" if cfg.one_insts_kv_swap else "v_desc_reused") + if cfg.sparse_reuse_kv + else "k_desc_0" + ) + }, + (smem_kv, tmem_o): { + ( + ("v_desc_reused_0" if cfg.one_insts_kv_swap else "v_desc_reused") + if cfg.sparse_reuse_kv + else "v_desc_0" + ) + }, } smem_allocator = SmemAllocator() + if isinstance(work_queue, ThroughputLatencyMlaClcWorkQueue): + work_queue._response_allocator = smem_allocator + smem_allocator.add_resource(work_queue) + if schedule_token_throttle is not None: + smem_allocator.add_resource(schedule_token_throttle) smem_allocator.add_resource(smem_q) if use_page_offsets: smem_allocator.add_resource(smem_page_offsets) - for resource in (smem_kv, tmem_s, local, global_softmax, tmem_p, tmem_o, corr): + for resource in (smem_kv, tmem_s, local, global_softmax, p_resource, tmem_o, corr): smem_allocator.add_resource(resource) smem_allocator.add_tmem_ptr( SmemAllocation("ll_mla_q64_tmem_ptr_i32", dtype=Int32, alignment=4) @@ -1414,6 +1655,9 @@ def _make_keeps_mma_ab_task_graph( tasks=tasks, resource_dependency_graph=deps, dma_consumer_release_labels=dma_release_labels, + # See the delayed-release schedule above: keep the stock checks in + # warning mode for BF16 reuse instead of maintaining a second verifier. + skip_validation=cfg.sparse_reuse_kv and not cfg.is_fp8_qkv(), smem_allocator=smem_allocator, tmem_allocator=tmem_allocator, verbose=verbose, @@ -1423,7 +1667,7 @@ def _make_keeps_mma_ab_task_graph( class ThroughputLatencyMlaDecodeTs: - """Dense throughput-latency 1CTA MLA TS wrapper.""" + """Throughput-latency 1CTA MLA over dense pages or prepared sparse rows.""" def __init__( self, @@ -1449,6 +1693,9 @@ def __init__( explicit_split_kv: int | None = None, explicit_persistent: bool | None = None, mask_type: MaskType | str = MaskType.CAUSAL, + device_scales: bool = False, + sparse_profile=None, + finalize_output: bool = False, ): """Initialize one selected physical tile profile over logical flat Q rows.""" import cutlass as _cutlass @@ -1477,6 +1724,15 @@ def __init__( self.explicit_split_kv = explicit_split_kv self.explicit_persistent = explicit_persistent self.mask_type = normalize_mask_type(mask_type) + self.device_scales = device_scales + self.sparse_profile = sparse_profile + self.finalize_output = finalize_output + self.output_quant = None + # Separate prepared source lists, rather than a premerged tagged list. + # Both representations reuse the same metadata staging and KV loaders. + self.direct_sparse = False + self.direct_static_scales = False + self.direct_sparse_capacities = (0, 0) cfg = self._make_config() # Parallel standalone reduction requires a fixed split profile and a @@ -1527,18 +1783,36 @@ def __init__( self.acc_dtype = acc_dtype self.lse_dtype = lse_dtype + def _parallel_reduction_setup(self, cfg): + if cfg.fuse_sparse_reduction: + if cfg.logical_seq_len_q != 1: + raise ValueError( + "final sparse reduction requires one virtual query per request" + ) + # Partials already contain every V partition. Finalize one whole + # D512 head per CTA, sharing its statistics across four vec4 warps. + reducer_cfg = dataclass_replace( + cfg, num_ctas_per_head_dim=1, head_dim_per_cta_v=cfg.head_dim_v + ) + topology = make_balanced_parallel_reduction_topology( + cfg.num_ctas_per_seq_kv, cluster_size=1 + ) + return reducer_cfg, topology, cfg.head_dim_v + return ( + cfg, + self.parallel_reduction_topology, + self.parallel_reduction_elements_per_slice, + ) + def compile_topology_signature(self) -> tuple[object, ...]: """Describe batch-derived reducer choices without retaining batch.""" cfg = self._make_config() if cfg.use_multi_ctas_kv != 1 or cfg.use_cluster_reduction == 1: return ("no_separate_reducer",) - if self.use_parallel_reduction: - return ( - "parallel", - self.parallel_reduction_topology, - self.parallel_reduction_elements_per_slice, - ) + if self.use_parallel_reduction or cfg.fuse_sparse_reduction: + _, topology, elements = self._parallel_reduction_setup(cfg) + return ("parallel", topology, elements) grid, smem, threads, reduction_ctas = gmem_reduction_launch_shape( cfg, cfg.seq_len_q, @@ -1567,6 +1841,7 @@ def compile_signature(self) -> tuple[object, ...]: self.parallel_reduction_topology, self.parallel_reduction_elements_per_slice, self.compile_topology_signature(), + self.output_quant, ) def _make_config(self): @@ -1592,6 +1867,15 @@ def _make_config(self): explicit_persistent=self.explicit_persistent, mask_type=self.mask_type, ) + if self.sparse_profile is not None: + cfg = configure_sparse_mla( + cfg, + self.sparse_profile, + finalize_output=self.finalize_output, + direct_inputs=self.direct_sparse, + static_scales=self.direct_static_scales, + source_capacities=self.direct_sparse_capacities, + ) return cfg def validate_split_kv_launch(self, split_kv: int, workspace) -> None: @@ -1662,10 +1946,11 @@ def __call__( split_kv: cutlass.Int32, cache_seqs: cute.Tensor, cu_seqlens_q: cute.Tensor | None, - block_split_kvs: cute.Tensor, + block_split_kvs: object, softmax_scale: cutlass.Float32, output_scale: cutlass.Float32, stream: object, + output_quant_args=None, ): """Execute throughput-latency 1CTA MLA with fixed or compact ragged Q.""" cfg = self._make_config() @@ -1775,75 +2060,133 @@ def __call__( l2_promotion=cuda.TensorMapL2Promotion.l2_128b, ) - q_rope_swizzle = cuda.TensorMapSwizzle.s128b - if cutlass.const_expr(cfg.is_fp8_qkv() and cfg.rope_dim == 64): - q_rope_swizzle = cuda.TensorMapSwizzle.s64b - if cutlass.const_expr(cu_seqlens_q is not None): - q_rope_tma = cute.make_tensor( - q_rope.iterator, + if cutlass.const_expr(cfg.rope_dim > 0): + q_rope_swizzle = cuda.TensorMapSwizzle.s128b + if cutlass.const_expr(cfg.is_fp8_qkv() and cfg.rope_dim == 64): + q_rope_swizzle = cuda.TensorMapSwizzle.s64b + if cutlass.const_expr(cu_seqlens_q is not None): + q_rope_tma = cute.make_tensor( + q_rope.iterator, + cute.make_layout( + (q_rope.shape[1], q_rope.shape[0] * q_rope.shape[2]), + stride=(q_rope.stride[1], q_rope.stride[0]), + ), + ) + tma_desc_q_rope = create_tensor_map_ragged_from_tensor( + q_rope_tma, + box_dims=(min(tma_box0, cfg.rope_dim), cfg.tile_size_q), + ragged_dim=1, + stride_order=(0, 1), + swizzle=q_rope_swizzle, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + else: + q_rope_tma = cute.make_tensor( + q_rope.iterator, + cute.make_layout( + ( + q_rope.shape[1], + q_rope.shape[0] * q_rope.shape[2], + q_rope.shape[3], + ), + stride=( + q_rope.stride[1], + q_rope.stride[0], + q_rope.stride[3], + ), + ), + ) + tma_desc_q_rope = create_tensor_map_tiled_from_view( + q_rope_tma, + box_dims=(min(tma_box0, cfg.rope_dim), cfg.tile_size_q, 1), + stride_order=(0, 1, 2), + swizzle=q_rope_swizzle, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + else: + tma_desc_q_rope = tma_desc_q_latent + + c_latent_tma = cute.make_tensor( + c_latent.iterator, + cute.select(c_latent.layout, mode=[1, 0, 2]), + ) + if cutlass.const_expr(cfg.num_tokens_per_page == 1): + # Sparse routes address individual rows; the unused D512 RoPE + # descriptor slot carries the independent compressed-pool map. + c_rows = cute.make_tensor( + c_latent.iterator, cute.make_layout( - (q_rope.shape[1], q_rope.shape[0] * q_rope.shape[2]), - stride=(q_rope.stride[1], q_rope.stride[0]), + (c_latent.shape[1], c_latent.shape[2]), + stride=(c_latent.stride[1], c_latent.stride[2]), ), ) - tma_desc_q_rope = create_tensor_map_ragged_from_tensor( - q_rope_tma, - box_dims=(min(tma_box0, cfg.rope_dim), cfg.tile_size_q), - ragged_dim=1, + c_extra_rows = cute.make_tensor( + c_rope.iterator, + cute.make_layout( + (c_rope.shape[1], c_rope.shape[2]), + stride=(c_rope.stride[1], c_rope.stride[2]), + ), + ) + tma_desc_c_latent = create_tensor_map_tiled_from_view( + c_rows, + box_dims=(tma_box0, 1), stride_order=(0, 1), - swizzle=q_rope_swizzle, + swizzle=cuda.TensorMapSwizzle.s128b, l2_promotion=cuda.TensorMapL2Promotion.l2_128b, ) - else: - q_rope_tma = cute.make_tensor( - q_rope.iterator, - cute.make_layout( - ( - q_rope.shape[1], - q_rope.shape[0] * q_rope.shape[2], - q_rope.shape[3], - ), - stride=( - q_rope.stride[1], - q_rope.stride[0], - q_rope.stride[3], - ), - ), + tma_desc_c_rope = create_tensor_map_tiled_from_view( + c_extra_rows, + box_dims=(tma_box0, 1), + stride_order=(0, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, ) - tma_desc_q_rope = create_tensor_map_tiled_from_view( - q_rope_tma, - box_dims=(min(tma_box0, cfg.rope_dim), cfg.tile_size_q, 1), + else: + tma_desc_c_latent = create_tensor_map_tiled_from_view( + c_latent_tma, + box_dims=(tma_box0, tma_page_tokens, 1), stride_order=(0, 1, 2), - swizzle=q_rope_swizzle, + swizzle=cuda.TensorMapSwizzle.s128b, l2_promotion=cuda.TensorMapL2Promotion.l2_128b, ) - c_latent_tma = cute.make_tensor( - c_latent.iterator, - cute.select(c_latent.layout, mode=[1, 0, 2]), - ) - tma_desc_c_latent = create_tensor_map_tiled_from_view( - c_latent_tma, - box_dims=(tma_box0, tma_page_tokens, 1), - stride_order=(0, 1, 2), - swizzle=cuda.TensorMapSwizzle.s128b, - l2_promotion=cuda.TensorMapL2Promotion.l2_128b, - ) - - c_rope_tma = cute.make_tensor( - c_rope.iterator, - cute.select(c_rope.layout, mode=[1, 0, 2]), - ) - c_rope_swizzle = cuda.TensorMapSwizzle.s128b - if cutlass.const_expr(cfg.is_fp8_qkv() and cfg.rope_dim == 64): - c_rope_swizzle = cuda.TensorMapSwizzle.s64b - tma_desc_c_rope = create_tensor_map_tiled_from_view( - c_rope_tma, - box_dims=(min(tma_box0, cfg.rope_dim), tma_page_tokens, 1), - stride_order=(0, 1, 2), - swizzle=c_rope_swizzle, - l2_promotion=cuda.TensorMapL2Promotion.l2_128b, - ) + if cutlass.const_expr(cfg.rope_dim > 0): + c_rope_tma = cute.make_tensor( + c_rope.iterator, + cute.select(c_rope.layout, mode=[1, 0, 2]), + ) + c_rope_swizzle = cuda.TensorMapSwizzle.s128b + if cutlass.const_expr(cfg.is_fp8_qkv() and cfg.rope_dim == 64): + c_rope_swizzle = cuda.TensorMapSwizzle.s64b + tma_desc_c_rope = create_tensor_map_tiled_from_view( + c_rope_tma, + box_dims=(min(tma_box0, cfg.rope_dim), tma_page_tokens, 1), + stride_order=(0, 1, 2), + swizzle=c_rope_swizzle, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + elif cutlass.const_expr(cfg.num_tokens_per_page != 1): + tma_desc_c_rope = tma_desc_c_latent + + tma_desc_v = tma_desc_c_latent + if cutlass.const_expr(cfg.num_tokens_per_page == 1): + # D512 does not use Q-rope; reuse its descriptor slot for the + # second pool's contiguous-run map. Gather maps stay separate. + rows_per_issuer = 32 + tma_desc_v = create_tensor_map_tiled_from_view( + c_rows, + box_dims=(tma_box0, rows_per_issuer), + stride_order=(0, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + tma_desc_q_rope = create_tensor_map_tiled_from_view( + c_extra_rows, + box_dims=(tma_box0, rows_per_issuer), + stride_order=(0, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) softmax_scale_log2 = softmax_scale * LOG2_E use_gmem_reduction = cutlass.const_expr( @@ -1911,7 +2254,7 @@ def __call__( tma_desc_q_rope, tma_desc_c_latent, tma_desc_c_rope, - tma_desc_c_latent, + tma_desc_v, c_rope, page_offsets, o, @@ -1923,6 +2266,8 @@ def __call__( softmax_scale_log2, output_scale, tile_sched_params, + block_split_kvs, + output_quant_args, ).launch( grid=grid, block=[cfg.threads_per_cta, 1, 1], @@ -1932,13 +2277,24 @@ def __call__( use_pdl=acc_o is not None, ) if cutlass.const_expr(acc_o is not None): - if cutlass.const_expr(self.use_parallel_reduction): - topology = self.parallel_reduction_topology + if cutlass.const_expr( + self.use_parallel_reduction or cfg.fuse_sparse_reduction + ): + reducer_cfg, topology, elements = self._parallel_reduction_setup(cfg) + atten_sinks = None + if cutlass.const_expr(cfg.fuse_sparse_reduction): + if cutlass.const_expr(cfg.sparse_direct): + atten_sinks = block_split_kvs[-1] + else: + atten_sinks = cute.make_tensor( + block_split_kvs.iterator + 2, + cute.make_layout(cfg.logical_num_heads_q), + ) reduction_grid, reduction_cluster = ( parallel_gmem_reduction_launch_shape( - cfg, + reducer_cfg, topology, - self.parallel_reduction_elements_per_slice, + elements, ) ) reduction_grid = ( @@ -1946,9 +2302,9 @@ def __call__( reduction_grid[1], batch_size, ) - reduction_threads = parallel_gmem_reduction_threads( - self.parallel_reduction_elements_per_slice - ) + if cutlass.const_expr(cfg.fuse_sparse_reduction): + reduction_grid = (cfg.logical_num_heads_q, 1, batch_size) + reduction_threads = parallel_gmem_reduction_threads(elements) parallel_reducer = self.parallel_gmem_reduction_kernel( o, lse, @@ -1956,6 +2312,8 @@ def __call__( acc_lse, cache_seqs, cu_seqlens_q, + atten_sinks, + output_quant_args, ) if cutlass.const_expr(topology.cluster_size == 1): parallel_reducer.launch( @@ -2028,10 +2386,43 @@ def dense_kernel( softmax_scale_log2: cutlass.Float32, output_scale: cutlass.Float32, tile_sched_params: object, + sparse_scale_params: object, + output_quant_args=None, ): """Execute one flat-Q, batch, KV-split, and V head-dimension tile.""" cfg = self._make_config() + if cutlass.const_expr(cfg.sparse_direct): + from ..sparse_views import bind_sparse_views + + si, ci, sl, cl, sm, qs, ss, os, sinks = sparse_scale_params + request = None + if cutlass.const_expr(cfg.use_persistent_scheduler != 1): + request = cutlass.Int64(cute.arch.block_idx()[2]) + page_offsets, cache_seqs = bind_sparse_views( + si, + ci, + sl, + cl, + cfg.total_kv_tiles * cfg.tile_size_kv, + request=request, + ) + if cutlass.const_expr(not cfg.sparse_static_scales): + softmax_scale_log2 = ( + cutlass.Float32(sm[0]) + * cutlass.Float32(qs[0]) + * cutlass.Float32(ss[0]) + * LOG2_E + ) + output_scale = cutlass.Float32(os[0]) * cutlass.Float32(ss[0]) + sparse_epilogue_params = sinks + elif cutlass.const_expr(self.device_scales): + softmax_scale_log2 = cutlass.Float32(sparse_scale_params[0]) * LOG2_E + output_scale = cutlass.Float32(sparse_scale_params[1]) + sparse_epilogue_params = sparse_scale_params + else: + sparse_epilogue_params = None + # The grid is expressed in physical flat-Q tile coordinates. Decode it # once here; Q/O resources retain responsibility for logical-row # mapping, compact-ragged offsets, and padded-tail publication. @@ -2105,10 +2496,17 @@ def dense_kernel( prims.prefetch_tensormap(tma_desc_c_latent.get_ptr()) prims.prefetch_tensormap(tma_desc_c_rope.get_ptr()) prims.prefetch_tensormap(tma_desc_v.get_ptr()) + if cutlass.const_expr(cfg.sparse_direct): + lane = cute.arch.thread_idx()[0] % 32 + if lane < (cfg.logical_num_heads_q + 31) // 32: + address = sparse_epilogue_params.iterator.toint() + cutlass.Int64( + lane * 128 + ) + cute.arch.inline_ptx( + "prefetch.global.L2 [{$r0}];", read_only_args=[address] + ) clc_response_ptr = None - if cutlass.const_expr(use_clc_dynamic): - clc_response_ptr = cute.arch.alloc_smem(cutlass.Int128, 2) task_manager, cluster_corr_resource = build_throughput_latency_mla_task_manager( cfg, @@ -2130,6 +2528,11 @@ def dense_kernel( cta_idx_head_dim_v=cta_idx_head_dim_v, scale_softmax_log2=softmax_scale_log2, output_scale=output_scale, + sparse_epilogue_params=sparse_epilogue_params, + output_quant=self.output_quant + if (cfg.fuse_sparse_epilogue or cfg.fuse_sparse_cluster_epilogue) + else None, + output_quant_args=output_quant_args, o_tensor=o, lse_tensor=lse, acc_o_tensor=acc_o, @@ -2170,10 +2573,27 @@ def dense_kernel( ) prims.tcgen05_relinquish_alloc_permit(group=prims.CTAGroup.CTA_1) - prims.barrier_cta_sync( - barrier_id=TMEM_LIFECYCLE_BARRIER_ID, - thread_count=cfg.threads_per_cta, - ) + if cutlass.const_expr(cfg.sparse_direct): + # Loaders need the initialized pipelines, but not TMEM allocation. + # A distinct barrier keeps an early-finishing loader's final full- + # CTA rendezvous from joining the compute-only readiness phase. + if warp_idx < Int32( + cfg.correction_warp_idx + cfg.correction_num_warps + ) or warp_idx == Int32(cfg.mma_warp_idx): + prims.barrier_cta_sync( + barrier_id=TMEM_READY_BARRIER_ID, + thread_count=( + cfg.correction_warp_idx + + cfg.correction_num_warps + + cfg.mma_num_warps + ) + * 32, + ) + else: + prims.barrier_cta_sync( + barrier_id=TMEM_LIFECYCLE_BARRIER_ID, + thread_count=cfg.threads_per_cta, + ) if cutlass.const_expr(has_runtime_activity_guard): if run_task_graph: task_manager.run() @@ -2263,11 +2683,13 @@ def parallel_gmem_reduction_kernel( acc_lse: cute.Tensor, cache_seqs: cute.Tensor, cu_seqlens_q: cute.Tensor, + atten_sinks: cute.Tensor | None = None, + output_quant_args=None, ): """Dispatch the automatically selected parallel standalone reducer.""" cfg = self._make_config() - topology = self.parallel_reduction_topology + cfg, topology, elements = self._parallel_reduction_setup(cfg) run_parallel_gmem_reduction_kernel( output, lse, @@ -2279,5 +2701,9 @@ def parallel_gmem_reduction_kernel( topology.cluster_size, topology.slots_per_rank, topology.actual_splits, - self.parallel_reduction_elements_per_slice, + elements, + atten_sinks, + cfg.fuse_sparse_reduction, + self.output_quant, + output_quant_args, ) 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 11d39bdec7e6..11a40b52bfd1 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 @@ -19,8 +19,12 @@ from cutlass import Float32, Int32, Int64 from cutlass.experimental import primitives as prims -from ...separate_reduction import finalize_log2_sum_exp, normalized_lse_weight -from ..helpers.constants import SPLIT_REDUCTION_SCALE_BARRIER_ID +from ...separate_reduction import ( + finalize_log2_sum_exp, + normalized_lse_weight, + attention_sink_scale, +) +from ..helpers.constants import LN_2, SPLIT_REDUCTION_SCALE_BARRIER_ID from ..helpers.math import ceil_div from ..helpers.ops import ( fmax_f32, @@ -197,6 +201,10 @@ def _run_parallel_gmem_reduction_g1_shared_stats( cu_seqlens_q, cfg, elements_per_slice: cutlass.Constexpr[int], + atten_sinks=None, + finalize_attention: cutlass.Constexpr[bool] = False, + output_quant: cutlass.Constexpr = None, + output_quant_args=None, ): """Reduce G1 partials with the compact row-shared schedule.""" @@ -204,8 +212,15 @@ def _run_parallel_gmem_reduction_g1_shared_stats( thread_idx, _, _ = cute.arch.thread_idx() warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) lane_idx = thread_idx % Int32(GMEM_REDUCTION_WARP_LANES) - head_dim_cta_idx = block_idx_y % Int32(cfg.num_ctas_per_head_dim) - q_idx = block_idx_y // Int32(cfg.num_ctas_per_head_dim) + if cutlass.const_expr(finalize_attention): + # The final sparse grid is one CTA per logical head. Decode its + # physical producer tile without launching padded head/V slices. + q_idx = slice_idx // Int32(cfg.tile_size_q) + slice_idx = slice_idx % Int32(cfg.tile_size_q) + head_dim_cta_idx = Int32(0) + else: + head_dim_cta_idx = block_idx_y % Int32(cfg.num_ctas_per_head_dim) + q_idx = block_idx_y // Int32(cfg.num_ctas_per_head_dim) elements_per_slice = _parallel_reduction_effective_slice_elements( cfg, elements_per_slice ) @@ -271,6 +286,11 @@ def _run_parallel_gmem_reduction_g1_shared_stats( ) lse_sum = warp_reduce_sum_f32(lse_sum) global_lse = finalize_log2_sum_exp(lse_max, lse_sum) + stats_sink_scale = Float32(1) + if cutlass.const_expr(finalize_attention and atten_sinks is not None): + stats_sink_scale = attention_sink_scale( + global_lse, Float32(atten_sinks[stats_query_row[1]]) + ) if ( lane_idx == Int32(0) and head_dim_cta_idx == Int32(0) @@ -282,16 +302,23 @@ def _run_parallel_gmem_reduction_g1_shared_stats( batch_idx, cu_seqlens_q, ) - (lse.iterator.raw_ptr() + stats_output_query_row).store(global_lse) + (lse.iterator.raw_ptr() + stats_output_query_row).store( + global_lse * Float32(LN_2) if finalize_attention else global_lse + ) for lane_slot_i in cutlass.range_constexpr(lse_per_lane): stats_split_idx = lane_idx + Int32( lane_slot_i * GMEM_REDUCTION_WARP_LANES ) if stats_split_idx < Int32(cfg.num_ctas_per_seq_kv): + stats_weight = normalized_lse_weight( + lane_lse[lane_slot_i], global_lse + ) + if cutlass.const_expr(finalize_attention): + stats_weight *= stats_sink_scale smem_scale[ stats_row * Int32(cfg.num_ctas_per_seq_kv) + stats_split_idx - ] = normalized_lse_weight(lane_lse[lane_slot_i], global_lse) + ] = stats_weight prims.barrier_cta_sync( barrier_id=SPLIT_REDUCTION_SCALE_BARRIER_ID, @@ -342,16 +369,20 @@ def _run_parallel_gmem_reduction_g1_shared_stats( scale = Float32( smem_scale[row_in_slice * Int32(cfg.num_ctas_per_seq_kv) + split_idx] ) - split_ptr = acc_row_ptr + Int64(split_i * cfg.head_dim_v) - if cutlass.const_expr(output_elements_per_thread == 1): - output_acc[0] += Float32(split_ptr.load()) * scale - else: - partial_o = split_ptr.load( - count=output_elements_per_thread, - alignment=output_elements_per_thread * cfg.partial_o_dtype_bytes, - ).to(Float32) - for elem_i in cutlass.range_constexpr(output_elements_per_thread): - output_acc[elem_i] += partial_o[elem_i] * scale + # Sparse masked splits can leave their vector unused. Do not + # multiply an uninitialized/NaN partial by a zero statistical weight. + if not finalize_attention or scale > Float32(0): + split_ptr = acc_row_ptr + Int64(split_i * cfg.head_dim_v) + if cutlass.const_expr(output_elements_per_thread == 1): + output_acc[0] += Float32(split_ptr.load()) * scale + else: + partial_o = split_ptr.load( + count=output_elements_per_thread, + alignment=output_elements_per_thread + * cfg.partial_o_dtype_bytes, + ).to(Float32) + for elem_i in cutlass.range_constexpr(output_elements_per_thread): + output_acc[elem_i] += partial_o[elem_i] * scale output_query_row = public_query_flat_row( cfg, @@ -360,7 +391,9 @@ def _run_parallel_gmem_reduction_g1_shared_stats( cu_seqlens_q, ) out_elem_offset = _output_element_offset(cfg, output_query_row, dim_idx) - if cutlass.const_expr(output_elements_per_thread == 1): + if cutlass.const_expr(output_quant is not None): + output_quant.store(output_quant_args, output_acc, output_query_row, dim_idx) + elif cutlass.const_expr(output_elements_per_thread == 1): (output.iterator.raw_ptr() + out_elem_offset).store( output.element_type(output_acc[0]) ) @@ -769,9 +802,15 @@ def run_parallel_gmem_reduction_kernel( slots_per_rank: cutlass.Constexpr[int], actual_splits: cutlass.Constexpr[int], elements_per_slice: cutlass.Constexpr[int], + atten_sinks=None, + finalize_attention: cutlass.Constexpr[bool] = False, + output_quant: cutlass.Constexpr = None, + output_quant_args=None, ): """Reduce split-KV partials with row-shared local and peer statistics.""" + if cutlass.const_expr(finalize_attention and cluster_size != 1): + raise ValueError("final attention output uses a single-CTA reducer per row") prims.griddepcontrol(kind=prims.GridDepAction.WAIT) if cutlass.const_expr(cluster_size == 1): @@ -784,6 +823,10 @@ def run_parallel_gmem_reduction_kernel( cu_seqlens_q, cfg, elements_per_slice, + atten_sinks, + finalize_attention, + output_quant, + output_quant_args, ) return 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 deae8aa5fed2..14fd5758621b 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 @@ -194,7 +194,10 @@ def _init_order_p01_barriers(self): def _ordered_sequence_wait(self): """Wait for this P instance's ordered materialization turn.""" - if cutlass.const_expr(self.cfg.use_clc_dynamic_persistent_scheduler == 1): + if cutlass.const_expr( + self.cfg.use_clc_dynamic_persistent_scheduler == 1 + and self.cfg.num_insts_kv > 1 + ): cute.arch.mbarrier_wait( self._order_p01_barrier_ptr + self.inst_id, self._order_p01_phase, @@ -204,7 +207,10 @@ def _ordered_sequence_wait(self): def _ordered_sequence_arrive(self): """Signal the peer P instance after materialization completes.""" - if cutlass.const_expr(self.cfg.use_clc_dynamic_persistent_scheduler == 1): + if cutlass.const_expr( + self.cfg.use_clc_dynamic_persistent_scheduler == 1 + and self.cfg.num_insts_kv > 1 + ): signaling_id = 1 if self.inst_id == 0 else 0 cute.arch.mbarrier_arrive( self._order_p01_barrier_ptr + signaling_id, 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 8e2544d61276..e09393d91303 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 @@ -35,6 +35,8 @@ producer_work, ) +from ...helpers.gather import gather4, gather4_cached_warp, select_gather_map + from ...helpers.constants import ( CP_ASYNC_CACHE_CA, PAGE_OFFSET_BYTES, @@ -230,7 +232,11 @@ def _load_q_stage( local_flat_query_row, batch_idx, ) - if prims.elect_sync(): + if prims.elect_sync() and ( + self.cfg.load_num_warps == 1 + or cute.arch.thread_idx()[0] // 32 - self.cfg.load_warp_idx + == qk_stage_idx % self.cfg.load_num_warps + ): smem_offset = Int32(q_stage_smem_element_offset(self.cfg, qk_stage_idx)) prims.cp_async_bulk_tensor_shared_cta_global( stage_base.data_ptr(smem_offset), @@ -276,13 +282,14 @@ def load_q(self, stage_info: StageInfo): dim_offset, self.tma_desc_q_latent, ) - self._load_q_stage( - stage_info, - stage_base, - latent_stages, - Int32(0), - self.tma_desc_q_rope, - ) + if cutlass.const_expr(cfg.rope_dim > 0): + self._load_q_stage( + stage_info, + stage_base, + latent_stages, + Int32(0), + self.tma_desc_q_rope, + ) @consumer_work(returns=("q_desc", "q_desc_rope")) @cute.jit @@ -479,7 +486,15 @@ def init_read_state(self, stage_info: StageInfo): cache_slots = 2 if self.cfg.kernel_variant != "keeps_mma_ab" else 1 return cutlass.Array( Int32, - max(1, cache_slots * self.cfg.pages_per_kv_tile), + max( + 1, + cache_slots + * ( + self.cfg.sparse_cache_words + if self.cfg.num_tokens_per_page == 1 + else self.cfg.pages_per_kv_tile + ), + ), space=cutlass.AddressSpace.rmem, ) @@ -506,7 +521,7 @@ def _producer_load_page_offsets( cfg = self.cfg lane_idx = lane_idx_from_thread(cute.arch.thread_idx()[0]) pages_per_tile = Int32(cfg.pages_per_kv_tile) - if lane_idx < pages_per_tile: + if cutlass.const_expr(cfg.sparse_direct): batch_idx = batch_idx_for_stage_cfg(self.batch_idx, cfg, stage_info) cta_idx_q = cta_idx_q_for_stage(self.cta_idx_q, stage_info) cta_idx_kv = cta_idx_kv_for_stage(self.cta_idx_kv, stage_info) @@ -521,31 +536,70 @@ def _producer_load_page_offsets( cfg, stage_info, inst_id, is_v, section=section ) tile_idx = global_kv_tile_idx(cfg, local_tile_idx, seq_len_kv, cta_idx_kv) - last_valid_page = cute.math.max( - (seq_len_kv + Int32(cfg.num_tokens_per_page - 1)) - // Int32(cfg.num_tokens_per_page) - - Int32(1), - Int32(0), - ) - logical_page_idx = cute.math.min( - tile_idx * pages_per_tile + lane_idx, - last_valid_page, - ) + tile_start = tile_idx * pages_per_tile + routes = self.page_offsets.routes_for_tile(tile_start, batch_idx) + position = lane_idx * Int32(4) smem_offset = ( stage_info.stage_idx * Int32(cfg.page_offsets_entries_per_stage) - + lane_idx - ) - page_offsets_batch = self.page_offsets[None, batch_idx] - page_offsets_flat = cute.flat_divide(page_offsets_batch, (1,)) - gmem_ptr = page_offsets_flat[None, logical_page_idx].iterator.llvm_ptr - smem_ptr = cutlass.inttoptr( - self._smem_page_offsets.data_ptr(smem_offset).toint(cutlass.Int32), - 3, - cutlass.Int32, - ) - prims.cp_async_shared_global( - smem_ptr, gmem_ptr, PAGE_OFFSET_BYTES, CP_ASYNC_CACHE_CA + + position ) + if cutlass.const_expr(cfg.tile_size_kv < 128): + if position < pages_per_tile: + quad = routes.mapped_quad(tile_start + position) + self._smem_page_offsets.data_ptr(smem_offset).store( + quad.load(0, 4), alignment=16 + ) + else: + quad = routes.mapped_quad(tile_start + position) + self._smem_page_offsets.data_ptr(smem_offset).store( + quad.load(0, 4), alignment=16 + ) + cute.arch.fence_view_async_shared() + return + for fragment in cutlass.range_constexpr((cfg.pages_per_kv_tile + 31) // 32): + fragment_lane = lane_idx + Int32(fragment * 32) + if fragment_lane < pages_per_tile: + batch_idx = batch_idx_for_stage_cfg(self.batch_idx, cfg, stage_info) + cta_idx_q = cta_idx_q_for_stage(self.cta_idx_q, stage_info) + cta_idx_kv = cta_idx_kv_for_stage(self.cta_idx_kv, stage_info) + seq_len_kv = runtime_seq_len_kv_from_task_cache( + cfg, + decode_gen_task_cache(stage_info), + cta_idx_q, + self.cu_seqlens_q, + batch_idx, + ) + local_tile_idx = local_kv_tile_idx( + cfg, stage_info, inst_id, is_v, section=section + ) + tile_idx = global_kv_tile_idx( + cfg, local_tile_idx, seq_len_kv, cta_idx_kv + ) + last_valid_page = cute.math.max( + (seq_len_kv + Int32(cfg.num_tokens_per_page - 1)) + // Int32(cfg.num_tokens_per_page) + - Int32(1), + Int32(0), + ) + logical_page_idx = cute.math.min( + tile_idx * pages_per_tile + fragment_lane, + last_valid_page, + ) + smem_offset = ( + stage_info.stage_idx * Int32(cfg.page_offsets_entries_per_stage) + + fragment_lane + ) + page_offsets_batch = self.page_offsets[None, batch_idx] + page_offsets_flat = cute.flat_divide(page_offsets_batch, (1,)) + gmem_ptr = page_offsets_flat[None, logical_page_idx].iterator.llvm_ptr + smem_ptr = cutlass.inttoptr( + self._smem_page_offsets.data_ptr(smem_offset).toint(cutlass.Int32), + 3, + cutlass.Int32, + ) + prims.cp_async_shared_global( + smem_ptr, gmem_ptr, PAGE_OFFSET_BYTES, CP_ASYNC_CACHE_CA + ) @producer_work @cute.jit @@ -592,9 +646,88 @@ def read_offsets( # The load task waits on the page-offset stage, snapshots the page ids # into registers, and keeps swaps K0/K1 live until their delayed V use. del stage_info - cache_base = cache_slot * self.cfg.pages_per_kv_tile - for page_frag in cutlass.range_constexpr(self.cfg.pages_per_kv_tile): - cached_page_ids[cache_base + page_frag] = self.page_id(page_frag) + if cutlass.const_expr(self.cfg.num_tokens_per_page == 1): + lane = lane_idx_from_thread(cute.arch.thread_idx()[0]) + warp = cute.arch.thread_idx()[0] // 32 - self.cfg.load_warp_idx + groups_per_warp = self.cfg.tile_size_kv // (4 * self.cfg.load_num_warps) + for j in cutlass.range_constexpr(4): + cached_page_ids[cache_slot * self.cfg.sparse_cache_words + j] = Int32( + 0x7FFFFFFF + ) + if lane < groups_per_warp: + quad = warp * groups_per_warp + lane + offset = ( + self.consumer_work_stage * self.cfg.page_offsets_entries_per_stage + + quad * 4 + ) + values = self._smem_page_offsets.data_ptr(offset).load( + count=4, alignment=16 + ) + for j in cutlass.range_constexpr(4): + cached_page_ids[cache_slot * self.cfg.sparse_cache_words + j] = ( + values[j] + ) + base = cache_slot * self.cfg.sparse_cache_words + # Actual physical adjacency within one source is required. A + # full valid prefix can still contain random rows. Eight quads + # (32 rows) are needed for bulk; smaller issuer shares use gather4. + if cutlass.const_expr(groups_per_warp >= 8): + first = Int32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=cached_page_ids[base], + offset=lane & Int32(-8), + mask_and_clamp=0x1F, + kind=prims.Shfl.IDX, + ) + ) + contiguous = (first & Int32(0x7FFFFFFF)) != Int32(0x7FFFFFFF) + empty = cutlass.Boolean(True) + for j in cutlass.range_constexpr(4): + raw = cached_page_ids[base + j] + contiguous = contiguous and raw == first + (lane % 8) * 4 + j + empty = empty and (raw & Int32(0x7FFFFFFF)) == Int32(0x7FFFFFFF) + cached_page_ids[base + 4] = Int32( + cute.arch.vote_ballot_sync(contiguous) + ) + cached_page_ids[base + 5] = Int32(cute.arch.vote_ballot_sync(empty)) + else: + cached_page_ids[base + 4] = Int32(0) + cached_page_ids[base + 5] = Int32(0) + + else: + cache_base = cache_slot * self.cfg.pages_per_kv_tile + for page_frag in cutlass.range_constexpr(self.cfg.pages_per_kv_tile): + cached_page_ids[cache_base + page_frag] = self.page_id(page_frag) + if cutlass.const_expr(self.cfg.cache_uniform_sparse_quads): + # Stage once per token tile, then reuse these uniform coordinates + # across all four head-dimension stages. Keep the raw lane cache + # for the contiguous/empty bulk path and source selection. + base = cache_slot * self.cfg.sparse_cache_words + bulk = (cached_page_ids[base + 4] & Int32(255)) == Int32(255) + empty = (cached_page_ids[base + 5] & Int32(255)) == Int32(255) + if self.cfg.load_num_warps == 8 or not (bulk or empty): + for owner in cutlass.range_constexpr( + self.cfg.tile_size_kv // (4 * self.cfg.load_num_warps) + ): + for j in cutlass.range_constexpr(4): + raw = prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=cached_page_ids[base + j], + offset=owner, + mask_and_clamp=0x1F, + kind=prims.Shfl.IDX, + ) + row = raw & Int32(0x7FFFFFFF) + row = row if row != Int32(0x7FFFFFFF) else Int32(-1) + cached_page_ids[base + 6 + owner * 4 + j] = ( + cute.arch.make_warp_uniform(row) + ) + else: + # No gather reads these coordinates on a bulk/empty tile. + # Constants avoid carrying the previous tile's uniform cache. + for j in cutlass.range_constexpr(self.cfg.sparse_cache_words - 6): + cached_page_ids[base + 6 + j] = Int32(0) return cached_page_ids @@ -633,6 +766,7 @@ class SmemKvResource(MlaResource): tma_desc_c_latent: object = None tma_desc_c_rope: object = None tma_desc_v: object = None + tma_desc_bulk_extra: object = None c_rope_tensor: object = None page_offsets_kv: object = None page_offsets: object = None @@ -800,7 +934,11 @@ def _producer_load_kv_stage( dim_offset = head_dim_cta_offset_v(cfg, cta_idx_head_dim_v) + Int32( qk_stage_idx * cfg.head_dim_per_stage_v ) - tma_desc = self.tma_desc_v + tma_desc = ( + self.tma_desc_c_latent + if cfg.num_tokens_per_page == 1 + else self.tma_desc_v + ) else: active_width = cfg.qk_head_stage_width(qk_stage_idx) if cutlass.const_expr( @@ -839,6 +977,101 @@ def _producer_load_kv_stage( and cfg.head_dim_per_stage_kv == 128 ) + # Each lane owns four adjacent index-list entries. Bulk TMA requires + # a verified 32-row physical run within one source; irregular rows + # use gather4. The separate empty-window path only zero-fills padding. + if cutlass.const_expr(cfg.num_tokens_per_page == 1): + warp = cute.arch.thread_idx()[0] // 32 - cfg.load_warp_idx + groups_per_warp = cfg.tile_size_kv // (4 * cfg.load_num_warps) + base = page_id_slot * cfg.sparse_cache_words + for chunk in cutlass.range_constexpr( + max(1, cfg.tile_size_kv // (32 * cfg.load_num_warps)) + ): + mask = Int32(255) << Int32(chunk * 8) + contiguous = (cached_page_ids[base + 4] & mask) == mask + empty = (cached_page_ids[base + 5] & mask) == mask + first = Int32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=cached_page_ids[base], + offset=chunk * 8, + mask_and_clamp=0x1F, + kind=prims.Shfl.IDX, + ) + ) + if (contiguous or empty) and groups_per_warp >= 8: + row = Int32(-32) if empty else first & Int32(0x7FFFFFFF) + desc = select_gather_map( + self.tma_desc_v, self.tma_desc_bulk_extra, first < 0 + ) + if prims.elect_sync(): + for part in cutlass.range_constexpr( + cfg.head_dim_per_stage_kv // (128 // cfg.qkv_dtype_bytes) + ): + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.data_ptr( + (warp * groups_per_warp + chunk * 8) + * 4 + * inner_width + + part * cfg.tile_size_kv * inner_width + ), + cute.make_ptr(cutlass.Int8, desc), + (dim_offset + part * inner_width, row), + stage_info.barrier, + ) + elif cutlass.const_expr(cfg.cache_uniform_sparse_quads): + desc = select_gather_map(tma_desc, self.tma_desc_c_rope, first < 0) + gather4_cached_warp( + stage_base.data_ptr(warp * groups_per_warp * 4 * inner_width), + desc, + dim_offset, + cached_page_ids, + stage_info.barrier, + cache_start=base + 6, + num_quads=min(8, groups_per_warp), + num_slices=cfg.head_dim_per_stage_kv + // (128 // cfg.qkv_dtype_bytes), + quad_stride=4 * inner_width * cfg.qkv_dtype_bytes, + slice_stride=cfg.tile_size_kv + * inner_width + * cfg.qkv_dtype_bytes, + col_stride=inner_width, + ) + else: + # Each active lane already owns a register-cached quad. + # Issue disjoint TMA transactions directly from those + # lanes, avoiding eight rounds of shuffles/election. + lane = cute.arch.thread_idx()[0] % 32 + if lane >= chunk * 8 and lane < min(chunk * 8 + 8, groups_per_warp): + group = warp * groups_per_warp + lane + coords = cutlass.Array( + Int32, 4, space=cutlass.AddressSpace.rmem + ) + for j in cutlass.range_constexpr(4): + raw = cached_page_ids[base + j] + row = raw & Int32(0x7FFFFFFF) + coords[j] = row if row != Int32(0x7FFFFFFF) else Int32(-1) + desc = select_gather_map( + tma_desc, self.tma_desc_c_rope, first < 0 + ) + for part in cutlass.range_constexpr( + cfg.head_dim_per_stage_kv // (128 // cfg.qkv_dtype_bytes) + ): + gather4( + stage_base.data_ptr( + group * 4 * inner_width + + part * cfg.tile_size_kv * inner_width + ), + desc, + dim_offset + part * inner_width, + coords[0], + coords[1], + coords[2], + coords[3], + stage_info.barrier, + ) + return + if cutlass.const_expr(cfg.use_paged_kv == 1): if cutlass.const_expr( self.page_offsets_kv is None and self.page_offsets is None @@ -1059,6 +1292,7 @@ def _set_desc( inst_id: int = 0, *, k_subtile_idx: cutlass.Constexpr[int] = 0, + reuse_k: cutlass.Constexpr[bool] = False, ): """Build a K or V SMEM descriptor for the current producer stage.""" # Descriptor offsets mirror _init_smem_state, but are rebuilt against the @@ -1070,22 +1304,30 @@ def _set_desc( qkv_major_k_stride_bytes_for(self.cfg, self.cfg.head_dim_per_stage_kv) ) if cutlass.const_expr(is_v): + stage_base = self._stage_base(stage_info) + if cutlass.const_expr(reuse_k): + # K waits advance independently of releases. Keep its payload + # live until the corresponding delayed PV has consumed it. + stage_base = self._smem_kv.subview( + self.consumer_release_state.index + * Int32(self.cfg.kv_smem_stage_elements) + ) v_leading_byte_offset = ( Int32(0) if self.cfg.head_dim_per_stage_v == TCGEN05_BF16_K_BLOCK_WIDTH - else Int32(TCGEN05_BF16_SECOND_K_BLOCK_OFFSET_BYTES) + else Int32(self.cfg.tile_size_kv * 128) ) if cutlass.const_expr(self.cfg.is_fp8_qkv()): v_leading_byte_offset = Int32(0) desc = cprims.Tcgen05SmemDesc.build( - self._stage_base(stage_info), + stage_base, leading_byte_offset=v_leading_byte_offset, stride_byte_offset=stride_byte_offset, layout=qkv_smem_swizzle(self.cfg), ) return desc else: - k_leading_byte_offset = Int32(TCGEN05_BF16_SECOND_K_BLOCK_OFFSET_BYTES) + k_leading_byte_offset = Int32(self.cfg.tile_size_kv * 128) if cutlass.const_expr(self.cfg.is_fp8_qkv()): k_leading_byte_offset = Int32(self.cfg.kv_smem_tile_bytes) qk_stage_idx = k_subtile_idx @@ -1139,6 +1381,31 @@ def v_desc_0(self, stage_info: StageInfo, *, v_subtile_idx: cutlass.Constexpr[in del v_subtile_idx return self._set_desc(stage_info, 1, 0) + @consumer_work(returns=("v_desc_0",)) + @cute.jit + def v_desc_reused( + self, stage_info: StageInfo, *, v_subtile_idx: cutlass.Constexpr[int] + ): + """Read delayed V from the unreleased K payload in the same SMEM slot.""" + del v_subtile_idx + return self._set_desc(stage_info, 1, 0, reuse_k=True) + + @consumer_work(returns=("v_desc_0",)) + @cute.jit + def v_desc_reused_0( + self, stage_info: StageInfo, *, v_subtile_idx: cutlass.Constexpr[int] + ): + del v_subtile_idx + return self._set_desc(stage_info, 1, 0, reuse_k=True) + + @consumer_work(returns=("v_desc_1",)) + @cute.jit + def v_desc_reused_1( + self, stage_info: StageInfo, *, v_subtile_idx: cutlass.Constexpr[int] + ): + del v_subtile_idx + return self._set_desc(stage_info, 1, 1, reuse_k=True) + @consumer_work(returns=("v_desc_1",)) @cute.jit def v_desc_1(self, stage_info: StageInfo, *, v_subtile_idx: cutlass.Constexpr[int]): 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 f51036748809..c221d42f6102 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 @@ -19,6 +19,8 @@ from cutlass.experimental import primitives as prims +from ...helpers.constants import LN_2, LOG2_E + import cutlass import cutlass.cute as cute from cutlass import Float32, Int8, Int32, Int64 @@ -58,6 +60,7 @@ smem_array, ) from ...helpers.math import ( + neg_max_f32, ceil_div, fadd2, ffma2, @@ -67,6 +70,7 @@ partial_output_dtype, ) from ...helpers.ops import ( + fp8_quant_scale, fp8_quant_scale_rcp, pack_float4_to_fp8_e4m3, store_transposed_smem8b_x2, @@ -112,6 +116,10 @@ class TmemCorrResource(MlaResource): inst_id: cutlass.Constexpr[int] = 0 scale_softmax_log2: Float32 = None output_scale: Float32 = None + sparse_epilogue_params: object = None + atten_sinks: object = None + output_quant: cutlass.Constexpr = None + output_quant_args: object = None o_tensor: object = None lse_tensor: object = None acc_o_tensor: object = None @@ -133,9 +141,18 @@ class TmemCorrResource(MlaResource): _cluster_reduction_smem: object = None _cluster_reduction_barrier: object = None + def _keeps_direct_output(self): + return ( + self.cfg.kernel_variant == "keeps_mma_ab" + and self.cfg.num_tokens_per_page == 1 + and self.cfg.use_cluster_reduction != 1 + ) + def get_smem_requirements(self): """Return correction/output SMEM scratch and optional cluster buffers.""" - if self.inst_id != 1: + if self.inst_id != 1 or self._keeps_direct_output(): + # Keeps-AB normalizes through warp shuffles and stores from TMEM + # directly to GMEM. Its output and sum scratch are never used. return [] if self._alloc is None: self._alloc = SmemAllocation( @@ -170,7 +187,7 @@ def get_smem_requirements(self): @cute.jit def _init_store_state_from_context(self, context) -> None: """Create correction/output SMEM views and initialize cluster state.""" - if cutlass.const_expr(self.inst_id != 1): + if cutlass.const_expr(self.inst_id != 1 or self._keeps_direct_output()): return self._smem_o = smem_array( context, @@ -239,7 +256,7 @@ def init_store_state(self, stage_info: StageInfo) -> None: def _o_base_col(self): """Return the first TMEM column owned by the O accumulator.""" - if cutlass.const_expr(self.cfg.kernel_variant == "keeps_mma_ab"): + if cutlass.const_expr(self.cfg.single_kv_pipe): return Int32(2 * self.cfg.tmem_s_cols) return Int32(2 * self.cfg.tmem_s_cols + 2 * self.cfg.tmem_stats_cols) @@ -268,6 +285,10 @@ def _init_cluster_reduction_barrier(self): if warp_idx == Int32(cfg.correction_warp_idx): with cute.arch.elect_one(): prims.mbarrier_init(self._cluster_reduction_barrier, 1) + # Init and arrive can lower to different elected lanes. Order + # the warp's initialization before registering transactions. + prims.bar_warp_sync(cute.arch.FULL_MASK) + with cute.arch.elect_one(): prims.mbarrier_arrive_expect_tx( self._cluster_reduction_barrier, expected_bytes ) @@ -500,10 +521,13 @@ def _store_lse_to_gmem( self.cfg, col_group_idx, scale_idx ) head_idx = head_base_idx + local_row_idx + if cutlass.const_expr( + self.cfg.num_tokens_per_page == 1 and self.cfg.use_cluster_reduction != 1 + ): + should_store_lse = should_store_lse and cta_idx_head_dim_v == Int32(0) if should_store_lse: - # Swaps keeps at least four scale groups, so TileQ8 has eight - # padded rows in the correction register footprint. Do not let - # those rows publish split LSE into the next logical head tile. + # Bound physical rows before publishing statistics so a padded + # query row cannot write into the next logical head tile. if local_row_idx < Int32(self.cfg.tile_size_q) and head_idx < Int32( self.cfg.num_heads_q ): @@ -518,16 +542,45 @@ def _store_lse_to_gmem( batch_idx=batch_idx, ) ) - lse_sum = reduced_sum[scale_idx] + # Keep the four swap-M16 statistics in registers. A dynamic + # Array index here otherwise creates a local-memory round trip. + selected_sum = Float32(0) + selected_max = Float32(0) + for group in cutlass.range_constexpr( + num_softmax_scale_groups(self.cfg) + ): + if scale_idx == Int32(group): + selected_sum = reduced_sum[group] + selected_max = final_max[group] + lse_sum = selected_sum if cutlass.const_expr(self.cfg.is_fp8_qkv()): - # FP8 P is quantized as 448 * softmax(P) for BMM2. The - # online sum tracks the same scale, so undo it for the - # externally visible log-sum-exp value. + # P and the online sum use the same FP8 quantization scale. + # Undo that scale for the externally visible log-sum-exp. lse_sum = lse_sum * fp8_quant_scale_rcp() lse_val = ( cute.math.log2(lse_sum, fastmath=True) - + self.scale_softmax_log2 * final_max[scale_idx] + + self.scale_softmax_log2 * selected_max ) + if cutlass.const_expr(self.cfg.sparse_direct): + if selected_max == neg_max_f32(): + lse_val = Float32(-Float32.inf) + if cutlass.const_expr(self.cfg.fuse_sparse_epilogue): + if cutlass.const_expr(self.cfg.sparse_direct): + count = Float32(selected_max != neg_max_f32()) + else: + count = Float32( + self.sparse_epilogue_params[ + 2 + self.cfg.logical_num_heads_q + batch_idx + ] + ) + lse_val = ( + lse_val * Float32(LN_2) + if count > Float32(0) + else Float32(-Float32.inf) + ) + valid_output_row = valid_output_row and cta_idx_head_dim_v == Int32( + 0 + ) 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 @@ -796,6 +849,24 @@ def _reduce_cluster_partials_and_store_for_tile( Float32(-Float32.inf) if is_fully_masked else global_lse ) + sink_scale = Float32(1) + if cutlass.const_expr(cfg.fuse_sparse_cluster_epilogue): + head = storage_flat_query_row % Int32( + cfg.logical_num_heads_q + ) + sink = Float32(self.atten_sinks[head]) + sink_scale = Float32(0) + if has_finite_mass and sink != Float32(Float32.inf): + delta = sink * Float32(LOG2_E) - global_lse + z = cute.math.exp2( + -cute.math.abs(delta), fastmath=True + ) + inverse = cute.math.rcp(Float32(1) + z, approx=True) + sink_scale = ( + z * inverse if delta > Float32(0) else inverse + ) + published_lse *= Float32(LN_2) + if ( valid_output_row and dim_idx == Int32(0) @@ -842,6 +913,9 @@ def _reduce_cluster_partials_and_store_for_tile( ).to(Float32) acc_vec = acc_vec + partial_vec * scale + if cutlass.const_expr(cfg.fuse_sparse_cluster_epilogue): + acc_vec = acc_vec * sink_scale + if cutlass.const_expr(self.o_tensor is not None): if valid_output_row: output_query_row = public_query_flat_row( @@ -858,7 +932,16 @@ def _reduce_cluster_partials_and_store_for_tile( out_byte_offset = out_elem_offset * Int64( cfg.o_dtype_bytes ) - if cutlass.const_expr(cfg.use_fp8_output == 1): + if cutlass.const_expr( + self.output_quant is not None + ): + self.output_quant.store( + self.output_quant_args, + acc_vec, + output_query_row, + head_dim_offset + dim_idx, + ) + elif cutlass.const_expr(cfg.use_fp8_output == 1): packed_o = cutlass.Array( Int32, 2, space=cutlass.AddressSpace.rmem ) @@ -992,8 +1075,20 @@ def _apply_loop_correction_in_tmem( cfg.head_dim_per_stage_v // 2, space=cutlass.AddressSpace.rmem, ) - for reg_idx in cutlass.range_constexpr(cfg.head_dim_per_stage_v // 2): - scaled[reg_idx] = loaded[reg_idx] * scale_vals[0] + if cutlass.const_expr(cfg.paired_sparse_correction): + for reg_idx in cutlass.range_constexpr( + 0, cfg.head_dim_per_stage_v // 2, 2 + ): + pair = fmul2( + (loaded[reg_idx], loaded[reg_idx + 1]), + (scale_vals[0], scale_vals[0]), + ) + scaled[reg_idx], scaled[reg_idx + 1] = pair[0], pair[1] + else: + for reg_idx in cutlass.range_constexpr( + cfg.head_dim_per_stage_v // 2 + ): + scaled[reg_idx] = loaded[reg_idx] * scale_vals[0] scaled_vec = vector_from_scalars( tuple( scaled[reg_idx] @@ -1203,6 +1298,7 @@ def _rescale_loop_o( def _compute_tail_softmax_scales( self, task_cache, + stage_info, *, new_max_arr, sum_arr, @@ -1262,9 +1358,17 @@ def _compute_tail_softmax_scales( reduced_sum[scale_idx] = final_sum[scale_idx] else: inst0_sum = inst0_sum_arr[scale_idx] - inst1_sum = inst1_sum_arr[scale_idx] inst0_max = inst0_new_max_arr[scale_idx] - inst1_max = inst1_new_max_arr[scale_idx] + inst1_sum = ( + Float32(0.0) + if cutlass.const_expr(cfg.one_insts_kv_swap) + else inst1_sum_arr[scale_idx] + ) + inst1_max = ( + inst0_max + if cutlass.const_expr(cfg.one_insts_kv_swap) + else inst1_new_max_arr[scale_idx] + ) final_max[scale_idx] = cute.math.max(inst0_max, inst1_max, ftz=True) exp_scale0[scale_idx] = cute.math.exp2( self.scale_softmax_log2 * (inst0_max - final_max[scale_idx]), @@ -1274,6 +1378,8 @@ def _compute_tail_softmax_scales( self.scale_softmax_log2 * (inst1_max - final_max[scale_idx]), fastmath=True, ) + if cutlass.const_expr(cfg.one_insts_kv_swap): + exp_scale1[scale_idx] = Float32(0.0) final_sum[scale_idx] = ( inst0_sum * exp_scale0[scale_idx] + inst1_sum * exp_scale1[scale_idx] @@ -1306,7 +1412,7 @@ def _compute_tail_softmax_scales( ) ) - if not cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): + if cutlass.const_expr(cfg.kernel_variant != "keeps_mma_ab"): # The two softmax pipes first reduce within their owning warp, then # publish one value per column group so the warpgroup can form the # final normalization denominator before O is stored. @@ -1346,6 +1452,73 @@ def _compute_tail_softmax_scales( norm_scale = Float32(0.0) if reduced_sum[scale_idx] != Float32(0.0): norm_scale = self.output_scale / reduced_sum[scale_idx] + if cutlass.const_expr(cfg.fuse_sparse_epilogue): + # Four lanes compute the distinct row factors. Broadcast to + # their peers instead of repeating sink SFU work eight times. + norm_lane = ( + lane_idx & Int32(15) + if cfg.kernel_variant == "keeps_mma_ab" + else col_group_idx + ) + norm_lanes = 16 if cfg.kernel_variant == "keeps_mma_ab" else 4 + if lane_idx < Int32(norm_lanes): + batch_idx = batch_idx_for_stage_cfg(self.batch_idx, cfg, stage_info) + logical_head = ( + cta_idx_q_for_stage(self.cta_idx_q, stage_info) + * cfg.tile_size_q + + head_idx_for_stage(self.head_idx, cfg, stage_info) + + ( + warp_idx * Int32(16) + norm_lane + if cfg.kernel_variant == "keeps_mma_ab" + else local_q_head_idx_for_scale( + cfg, col_group_idx, scale_idx + ) + ) + ) + if cutlass.const_expr(cfg.sparse_direct): + count = Float32(final_max[scale_idx] != neg_max_f32()) + else: + count = Float32( + self.sparse_epilogue_params[ + 2 + cfg.logical_num_heads_q + batch_idx + ] + ) + norm_scale = Float32(0) + if logical_head < cfg.logical_num_heads_q and count > Float32(0): + sink = Float32(self.atten_sinks[logical_head]) + if sink == Float32(-Float32.inf): + norm_scale = self.output_scale / reduced_sum[scale_idx] + elif sink != Float32(Float32.inf): + delta = ( + sink * Float32(LOG2_E) + - self.scale_softmax_log2 * final_max[scale_idx] + ) + attention_rescale = cute.math.exp2( + -cute.math.max(delta, Float32(0)), approx=True + ) + sink_mass = cute.math.exp2( + cute.math.min(delta, Float32(0)), approx=True + ) + if cutlass.const_expr(cfg.is_fp8_qkv()): + sink_mass *= fp8_quant_scale() + norm_scale = ( + self.output_scale + * attention_rescale + * cute.math.rcp( + reduced_sum[scale_idx] * attention_rescale + + sink_mass, + approx=True, + ) + ) + norm_scale = Float32( + cprims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=norm_scale, + offset=norm_lane, + mask_and_clamp=0x1F, + kind=cprims.Shfl.IDX, + ) + ) final_scale0[scale_idx] = norm_scale * exp_scale0[scale_idx] final_scale1[scale_idx] = norm_scale * exp_scale1[scale_idx] return final_max, reduced_sum, final_scale0, final_scale1 @@ -1386,6 +1559,40 @@ def _store_keeps_tail_o_stage( cfg.head_dim_per_stage_v // 2 ) + if cutlass.const_expr(self.output_quant is not None): + # Each lane owns half a head's D stage. K128 can span the two + # lanes separated by 16; K32 blocks remain entirely lane-local. + width = min(self.output_quant.block_size, cfg.head_dim_per_stage_v // 2) + for stripe in cutlass.range_constexpr( + 0, cfg.head_dim_per_stage_v // 2, width + ): + values = cutlass.Array(Float32, width, space=cutlass.AddressSpace.rmem) + loaded = tcgen05_ld_16x32bx2_f32( + prims.make_tmem_ptr(base_addr0 + Int32(stripe), Float32), + num=width, + offset=Int32(cfg.head_dim_per_stage_v // 2), + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + cute.arch.fence_view_async_tmem_load() + for i in cutlass.range_constexpr(width): + values[i] = loaded[i] * final_scale0[0] + if global_head_idx < Int32(cfg.num_heads_q) and valid_output_row: + row = public_query_flat_row( + cfg, storage_flat_query_row, batch_idx, self.cu_seqlens_q + ) + self.output_quant.store( + self.output_quant_args, + values, + row, + head_dim_offset + + v_stage_idx * cfg.head_dim_per_stage_v + + half_warp_col_offset + + stripe, + lane_stride=16, + rotary=v_stage_idx + 1 == cfg.v_head_dim_stages, + ) + return + # ``tcgen05.ld.sync.aligned`` must remain convergent across the warp. # Padded query rows still own valid physical TMEM rows, so load and # normalize every row and predicate only the externally visible store. @@ -1695,13 +1902,16 @@ def _stage_generic_tail_o_to_smem( ), num=1, ) - o1_loaded_lo = prims.tcgen05_ld( - shape, - prims.make_tmem_ptr( - tcgen05_panel_addr(base_addr1, chunk_idx), Float32 - ), - num=1, - ) + if cutlass.const_expr(cfg.one_insts_kv_swap): + o1_loaded_lo = o0_loaded_lo + else: + o1_loaded_lo = prims.tcgen05_ld( + shape, + prims.make_tmem_ptr( + tcgen05_panel_addr(base_addr1, chunk_idx), Float32 + ), + num=1, + ) o0_loaded_hi = prims.tcgen05_ld( shape, prims.make_tmem_ptr( @@ -1712,16 +1922,19 @@ def _stage_generic_tail_o_to_smem( ), num=1, ) - o1_loaded_hi = prims.tcgen05_ld( - shape, - prims.make_tmem_ptr( - tcgen05_second_panel_addr( - tcgen05_panel_addr(base_addr1, chunk_idx) + if cutlass.const_expr(cfg.one_insts_kv_swap): + o1_loaded_hi = o0_loaded_hi + else: + o1_loaded_hi = prims.tcgen05_ld( + shape, + prims.make_tmem_ptr( + tcgen05_second_panel_addr( + tcgen05_panel_addr(base_addr1, chunk_idx) + ), + Float32, ), - Float32, - ), - num=1, - ) + num=1, + ) else: o0_loaded = prims.tcgen05_ld( shape, @@ -1730,13 +1943,16 @@ def _stage_generic_tail_o_to_smem( ), num=q_repeats, ) - o1_loaded = prims.tcgen05_ld( - shape, - prims.make_tmem_ptr( - tcgen05_panel_addr(base_addr1, chunk_idx), Float32 - ), - num=q_repeats, - ) + if cutlass.const_expr(cfg.one_insts_kv_swap): + o1_loaded = o0_loaded + else: + o1_loaded = prims.tcgen05_ld( + shape, + prims.make_tmem_ptr( + tcgen05_panel_addr(base_addr1, chunk_idx), Float32 + ), + num=q_repeats, + ) prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) cute.arch.fence_view_async_tmem_load() for pair_idx in cutlass.range_constexpr(o_reg_pair_count): @@ -1770,20 +1986,28 @@ def _stage_generic_tail_o_to_smem( o1_loaded[reg_base], o1_loaded[reg_base + 1], ) - final_pair = ffma2( - ( - final_scale0[scale_base], - final_scale0[scale_base + 1], - ), - o0_pair, - fmul2( + if cutlass.const_expr(cfg.one_insts_kv_swap): + # Only O0 exists: normalize it directly, without a second + # TMEM load or the two-stream weighted merge. + final_pair = fmul2( + (final_scale0[scale_base], final_scale0[scale_base + 1]), + o0_pair, + ) + else: + final_pair = ffma2( ( - final_scale1[scale_base], - final_scale1[scale_base + 1], + final_scale0[scale_base], + final_scale0[scale_base + 1], ), - o1_pair, - ), - ) + o0_pair, + fmul2( + ( + final_scale1[scale_base], + final_scale1[scale_base + 1], + ), + o1_pair, + ), + ) if cutlass.const_expr( cfg.use_fp8_output == 1 and self.acc_o_tensor is None @@ -1841,6 +2065,97 @@ def _stage_swaps_tail_o_to_smem( final_scale1=final_scale1, ) + @cute.jit + def _store_swap_quantized_stage( + self, + stage_info, + task_cache, + base_addr0, + base_addr1, + final_scale0, + final_scale1, + v_stage_idx, + head_dim_offset, + ): + """Transpose FP32 O in two half-Q passes using the existing O scratch. + + Half as many FP32 rows fit in the same buffer as the ordinary BF16 + stage. After the transpose, each lane owns a contiguous vector, so + rotation pairs stay local and quantization needs only warp reductions. + """ + cfg = self.cfg + tid = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] + loaded0 = prims.tcgen05_ld( + "32x32b", prims.make_tmem_ptr(base_addr0, Float32), num=cfg.tile_size_q + ) + if cutlass.const_expr(cfg.one_insts_kv_swap): + loaded1 = loaded0 + else: + loaded1 = prims.tcgen05_ld( + "32x32b", prims.make_tmem_ptr(base_addr1, Float32), num=cfg.tile_size_q + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + cute.arch.fence_view_async_tmem_load() + batch = batch_idx_for_stage_cfg(self.batch_idx, cfg, stage_info) + head_base = head_idx_for_stage(self.head_idx, cfg, stage_info) + q_tile = cta_idx_q_for_stage(self.cta_idx_q, stage_info) + rows_per_pass = cfg.tile_size_q // 2 + width = rows_per_pass # 128 channels * rows / 128 correction threads. + threads_per_row = 128 // width + for half in cutlass.range_constexpr(2): + for i in cutlass.range_constexpr(rows_per_pass): + h = half * rows_per_pass + i + scale_idx = (h // 8) * 2 + h % 2 + source_lane = (h % 8) // 2 + s0 = cprims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=final_scale0[scale_idx], + offset=source_lane, + mask_and_clamp=0x1F, + kind=cprims.Shfl.IDX, + ) + value = loaded0[h] * s0 + if cutlass.const_expr(not cfg.one_insts_kv_swap): + s1 = cprims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=final_scale1[scale_idx], + offset=source_lane, + mask_and_clamp=0x1F, + kind=cprims.Shfl.IDX, + ) + value += loaded1[h] * s1 + self._smem_o_i32[i * 128 + tid] = value.bitcast(Int32) + prims.barrier_cta_sync( + self.store_barrier_id, thread_count=WARPGROUP_THREADS + ) + h = head_base + half * rows_per_pass + tid // threads_per_row + flat_row, _, _, _, valid = flat_query_row_state( + h, + q_tile, + cfg.tile_size_q, + cfg.logical_num_heads_q, + cfg.logical_seq_len_q, + self.cu_seqlens_q, + batch, + ) + values = self._smem_o_i32.load( + tid * width, vector_size=width, alignment=width * 4 + ).bitcast(Float32) + if h < cfg.num_heads_q and valid: + row = public_query_flat_row(cfg, flat_row, batch, self.cu_seqlens_q) + self.output_quant.store( + self.output_quant_args, + values, + row, + head_dim_offset + + v_stage_idx * cfg.head_dim_per_stage_v + + (tid % threads_per_row) * width, + rotary=v_stage_idx + 1 == cfg.v_head_dim_stages, + ) + prims.barrier_cta_sync( + self.store_barrier_id, thread_count=WARPGROUP_THREADS + ) + @cute.jit def _normalize_and_store_tail_o( self, @@ -1867,6 +2182,7 @@ def _normalize_and_store_tail_o( final_max, reduced_sum, final_scale0, final_scale1 = ( self._compute_tail_softmax_scales( task_cache, + stage_info, new_max_arr=new_max_arr, sum_arr=sum_arr, inst0_new_max_arr=inst0_new_max_arr, @@ -1897,6 +2213,19 @@ def _normalize_and_store_tail_o( v_stage_idx=v_stage_idx, head_dim_offset=head_dim_offset, ) + elif cutlass.const_expr( + self.output_quant is not None and cfg.fuse_sparse_epilogue + ): + self._store_swap_quantized_stage( + stage_info, + task_cache, + base_addr0, + base_addr1, + final_scale0, + final_scale1, + v_stage_idx, + head_dim_offset, + ) else: self._stage_swaps_tail_o_to_smem( task_cache, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_o.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_o.py index eebeb725d40f..6e66442172e4 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_o.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_o.py @@ -127,7 +127,7 @@ def init_stage_state(self, stage_info: StageInfo): # Consumer aux work creates the stage-tracking variables used by the # correction task after it waits for O. self._init_tmem_state(stage_info) - return Int32(0), Int32(0), Int32(1) + return Int32(0), Int32(0), Int32(0 if self.cfg.one_insts_kv_swap else 1) @consumer_work( work_attrs=WorkAttr.AUXILIARY, @@ -138,7 +138,7 @@ def init_stage_work_tile_state(self, stage_info: StageInfo): """Initialize O-stage index variables for a persistent work tile.""" # Reset O stage bookkeeping for each persistent work tile. del stage_info - return Int32(0), Int32(0), Int32(1) + return Int32(0), Int32(0), Int32(0 if self.cfg.one_insts_kv_swap else 1) @cute.jit def _tail_has_prior_o(self, stage_info: StageInfo): 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 91dd9802cbff..e867ba4e7c90 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 @@ -43,6 +43,7 @@ ) from ...helpers.math import ( fadd2, + ffma2, fmul2, neg_max_f32, pack_float2_to_bf16, @@ -71,7 +72,7 @@ class TmemPResource(MlaResource): """TMEM probability tile exchanged from softmax to PV MMA. - The keeps-MMA-AB schedule uses one P pipe: softmax writes BF16 P directly + The keeps-MMA-AB schedule uses one P pipe: softmax writes native BF16/FP8 P directly into the score TMEM stage, and MmaTask consumes that TMEM stage with ``tcgen05.mma`` A-from-TMEM. """ @@ -147,6 +148,9 @@ def materialize_p( neg_scaled_max[idx] = -self.scale_softmax_log2 * safe_new_max local_sums[idx] = Float32(0.0) + sums = cutlass.Array(Float32, 4, space=cutlass.AddressSpace.rmem) + for i in cutlass.range_constexpr(4): + sums[i] = Float32(0) packed_p_reg_count = num_packed_p_regs(cfg) regs_p = cutlass.Array( Int32, packed_p_reg_count, space=cutlass.AddressSpace.rmem @@ -167,34 +171,44 @@ def materialize_p( p2 = Float32(0.0) p3 = Float32(0.0) if has_finite_max: - scaled01 = fadd2( - fmul2( - ( - s_arr[s_base + 0], - s_arr[s_base + 1], - ), + 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, - ) - scaled23 = fadd2( - fmul2( - ( - s_arr[s_base + 2], - s_arr[s_base + 3], + neg_scaled_pair, + ) + else: + scaled01 = fadd2( + fmul2( + (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, - ) + neg_scaled_pair, + ) + else: + scaled23 = fadd2( + fmul2( + (s_arr[s_base + 2], s_arr[s_base + 3]), log2_scale_pair + ), + neg_scaled_pair, + ) p0 = cute.math.exp2(scaled01[0], fastmath=True) p1 = cute.math.exp2(scaled01[1], fastmath=True) p2 = cute.math.exp2(scaled23[0], fastmath=True) p3 = cute.math.exp2(scaled23[1], fastmath=True) - local_sums[0] += p0 - local_sums[0] += p1 - local_sums[0] += p2 - local_sums[0] += p3 + if cutlass.const_expr(cfg.num_tokens_per_page == 1): + sums[0], sums[1] = fadd2((sums[0], sums[1]), (p0, p1)) + sums[2], sums[3] = fadd2((sums[2], sums[3]), (p2, p3)) + else: + local_sums[0] += p0 + local_sums[0] += p1 + local_sums[0] += p2 + local_sums[0] += p3 regs_p[packed_idx] = pack_float4_to_fp8_e4m3(p0, p1, p2, p3) elif cutlass.const_expr(cfg.is_fp8_qkv()): for packed_idx in cutlass.range_constexpr(packed_p_reg_count): @@ -228,16 +242,31 @@ 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(): - p0 = cute.math.exp2( - s_arr[s0] * self.scale_softmax_log2 + neg_scaled_max[0], - fastmath=True, - ) - p1 = cute.math.exp2( - s_arr[s1] * self.scale_softmax_log2 + neg_scaled_max[0], - fastmath=True, + if 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), + (neg_scaled_max[0], neg_scaled_max[0]), + ) + p0 = cute.math.exp2(scaled[0], fastmath=True) + p1 = cute.math.exp2(scaled[1], fastmath=True) + else: + p0 = cute.math.exp2( + s_arr[s0] * self.scale_softmax_log2 + neg_scaled_max[0], + fastmath=True, + ) + p1 = cute.math.exp2( + s_arr[s1] * self.scale_softmax_log2 + neg_scaled_max[0], + fastmath=True, + ) + if cutlass.const_expr(cfg.num_tokens_per_page == 1): + offset = (pair_idx % 2) * 2 + sums[offset], sums[offset + 1] = fadd2( + (sums[offset], sums[offset + 1]), (p0, p1) ) - local_sums[0] += p0 - local_sums[0] += p1 + else: + local_sums[0] += p0 + local_sums[0] += p1 else: scale0 = ((pair_idx % (2 * max(cfg.tile_size_q // 8, 1))) // 2) * 2 scale1 = scale0 + 1 @@ -259,11 +288,18 @@ def materialize_p( local_sums[scale1] += p1 regs_p[pair_idx] = pack_float2_to_bf16(p0, p1) + if cutlass.const_expr( + cfg.kernel_variant == "keeps_mma_ab" and cfg.num_tokens_per_page == 1 + ): + local_sums[0] = (sums[0] + sums[1]) + (sums[2] + sums[3]) + for scale_idx in cutlass.range_constexpr(num_scale_groups): self.tmem_alias_ref._p_local_sum_arr[scale_idx] = local_sums[scale_idx] local_sum_arr[scale_idx] = local_sums[scale_idx] - if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): + if cutlass.const_expr( + cfg.kernel_variant == "keeps_mma_ab" and cfg.num_tokens_per_page != 1 + ): state_idx = cute.arch.thread_idx()[0] state_ptr = self.tmem_alias_ref._softmax_scratch.data_ptr(state_idx) old_max = u32_bits_to_float(state_ptr.load(is_volatile=True, alignment=4)) @@ -295,15 +331,31 @@ def materialize_p( tcgen05_second_panel_addr(p_base), regs_p ) else: - tcgen05_store_p_16x32bx2_x16(p_base, regs_p, 0) - tcgen05_store_p_16x32bx2_x16(p_base + Int32(16), regs_p, 16) - if cutlass.const_expr(cfg.head_dim_per_cta_v > 256): - tcgen05_store_p_16x32bx2_x16( - tcgen05_second_panel_addr(p_base), regs_p, 0 - ) - tcgen05_store_p_16x32bx2_x16( - tcgen05_second_panel_addr(p_base) + Int32(16), regs_p, 16 - ) + if cutlass.const_expr(cfg.tile_size_kv == 128): + tcgen05_store_p_16x32bx2_x16(p_base, regs_p, 0) + tcgen05_store_p_16x32bx2_x16(p_base + Int32(16), regs_p, 16) + if cutlass.const_expr(cfg.head_dim_per_cta_v > 256): + tcgen05_store_p_16x32bx2_x16( + tcgen05_second_panel_addr(p_base), regs_p, 0 + ) + tcgen05_store_p_16x32bx2_x16( + tcgen05_second_panel_addr(p_base) + Int32(16), regs_p, 16 + ) + else: + for part in cutlass.range_constexpr(packed_p_reg_count // 16): + tcgen05_store_p_16x32bx2_x16( + p_base + Int32(part * 16), + regs_p, + part * 16, + stride=cfg.tile_size_kv // 4, + ) + if cutlass.const_expr(cfg.head_dim_per_cta_v > 256): + tcgen05_store_p_16x32bx2_x16( + tcgen05_second_panel_addr(p_base) + Int32(part * 16), + regs_p, + part * 16, + stride=cfg.tile_size_kv // 4, + ) prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) cute.arch.fence_view_async_tmem_store() 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 88c7aa259d4f..1c29fd10f8c4 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 @@ -74,6 +74,7 @@ u32_to_float_for_atomic_max, ) from ...helpers.mask import MaskType +from ...helpers.gather import invalid_sparse_token, cache_sparse_mask from ...helpers.ops import ( float_to_u32_bits, freeze_smem_descriptor, @@ -118,6 +119,7 @@ class TmemSResource(MlaResource): p_ref: Optional[MemoryResource] = None global_ref: Optional[MemoryResource] = None cache_seqs: object = None + page_offsets: object = None head_idx: object = None batch_idx: object = None cta_idx_q: object = None @@ -162,6 +164,11 @@ def _reset_softmax_state(self): def get_smem_requirements(self): """Return the softmax scratch SMEM allocation.""" + if ( + self.cfg.kernel_variant == "keeps_mma_ab" + and self.cfg.num_tokens_per_page == 1 + ): + return [] if self._scratch_alloc is None: self._scratch_alloc = SmemAllocation( name=f"{self.name}_softmaxScratch", @@ -175,8 +182,7 @@ def get_tmem_requirements(self): if self._tmem_alloc is None: num_stages = ( self.pipeline_config.num_stages - if self.cfg.kernel_variant == "keeps_mma_ab" - and self.pipeline_config is not None + if self.cfg.single_kv_pipe and self.pipeline_config is not None else 1 ) self._tmem_alloc = TmemAllocation( @@ -233,15 +239,24 @@ def init_softmax_state(self, stage_info: StageInfo): # materialize_p(), and update_softmax_sum(). context = stage_info.context self._init_tmem_state(stage_info) - self._softmax_scratch = smem_array( - context, - self._scratch_alloc, - Uint32, - softmax_scratch_words(self.cfg), - ) + if cutlass.const_expr( + not ( + self.cfg.kernel_variant == "keeps_mma_ab" + and self.cfg.num_tokens_per_page == 1 + ) + ): + self._softmax_scratch = smem_array( + context, + self._scratch_alloc, + Uint32, + softmax_scratch_words(self.cfg), + ) self.q_desc_current = Int64(0) self.q_desc_rope_current = Int64(0) - if cutlass.const_expr(self.cfg.kernel_variant == "keeps_mma_ab"): + if cutlass.const_expr( + self.cfg.kernel_variant == "keeps_mma_ab" + and self.cfg.num_tokens_per_page != 1 + ): thread_idx = cute.arch.thread_idx()[0] state_ptr = self._softmax_scratch.data_ptr(thread_idx) state_ptr.store( @@ -322,7 +337,7 @@ def qk_mma( task_cache = decode_gen_task_cache(stage_info) stage_col_offset = Int32(0) - if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): + if cutlass.const_expr(cfg.single_kv_pipe): stage_col_offset = stage_info.stage_idx * Int32(cfg.tmem_s_cols) tmem_ptr = prims.make_tmem_ptr( task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] @@ -373,7 +388,11 @@ def qk_mma( scale_d = Boolean(True) if cutlass.const_expr(k_block + 1 < k_block_count): if cutlass.const_expr(k_block == 3): - kv_desc = kv_desc + Int32(TCGEN05_DESC_WRAPPED_K_BLOCK_UNITS) + kv_desc = kv_desc + Int32( + cfg.tile_size_kv * 8 - 6 + if cfg.tile_size_kv != 128 + else TCGEN05_DESC_WRAPPED_K_BLOCK_UNITS + ) q_desc = q_desc + Int32(q_p_desc_k_block_wrap_units(cfg)) else: kv_desc = kv_desc + Int32(TCGEN05_DESC_NEXT_K_BLOCK_UNITS) @@ -420,7 +439,9 @@ def update_softmax( ) for idx in cutlass.range_constexpr(num_scale_groups): - if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): + if cutlass.const_expr( + cfg.kernel_variant == "keeps_mma_ab" and cfg.num_tokens_per_page != 1 + ): state_ptr = self._softmax_scratch.data_ptr(warp_grp_thread_idx) old_max_vals[idx] = u32_bits_to_float( state_ptr.load(is_volatile=True, alignment=4) @@ -436,11 +457,45 @@ def update_softmax( for idx in cutlass.range_constexpr(num_s_regs_per_thread(cfg)): s_vals[idx] = neg_max_f32() + batch_idx = batch_idx_for_stage_cfg(self.batch_idx, cfg, stage_info) + cta_idx_q = cta_idx_q_for_stage(self.cta_idx_q, stage_info) + cta_idx_kv = cta_idx_kv_for_stage(self.cta_idx_kv, stage_info) + seq_len_kv = runtime_seq_len_kv_from_task_cache( + cfg, + task_cache, + cta_idx_q, + self.cu_seqlens_q, + batch_idx, + ) + local_tile_idx = softmax_kv_tile_idx(cfg, stage_info, self.inst_id) + tile_idx = global_kv_tile_idx(cfg, local_tile_idx, seq_len_kv, cta_idx_kv) + tile_offset_k = tile_idx * Int32(cfg.tile_size_kv) + next_tile_offset_k = tile_offset_k + Int32(cfg.tile_size_kv) + raw_mask_routes = ( + self.page_offsets.mask_for_tile(tile_offset_k, batch_idx) + if cutlass.const_expr(cfg.sparse_direct) + else self.page_offsets + ) + mask_routes = ( + cache_sparse_mask( + raw_mask_routes, + tile_offset_k, + seq_len_kv, + batch_idx, + swaps=cfg.kernel_variant != "keeps_mma_ab", + tile_size_kv=cfg.tile_size_kv, + ) + if cutlass.const_expr( + cfg.num_tokens_per_page == 1 and cfg.logical_seq_len_q == 1 + ) + else raw_mask_routes + ) + # Normalize both score layouts into one local S register array before # masking: keeps-MMA-AB reads a single TMEM panel, swaps-MMA-AB reads # two 16x256b panels. stage_col_offset = Int32(0) - if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): + if cutlass.const_expr(cfg.single_kv_pipe): stage_col_offset = stage_info.stage_idx * Int32(cfg.tmem_s_cols) base_addr = ( task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] @@ -479,20 +534,6 @@ def update_softmax( s_vals[q_repeats * 4 + ld_base + 2] = loaded1[ld_base + 2] s_vals[q_repeats * 4 + ld_base + 3] = loaded1[ld_base + 3] - batch_idx = batch_idx_for_stage_cfg(self.batch_idx, cfg, stage_info) - cta_idx_q = cta_idx_q_for_stage(self.cta_idx_q, stage_info) - cta_idx_kv = cta_idx_kv_for_stage(self.cta_idx_kv, stage_info) - seq_len_kv = runtime_seq_len_kv_from_task_cache( - cfg, - task_cache, - cta_idx_q, - self.cu_seqlens_q, - batch_idx, - ) - local_tile_idx = softmax_kv_tile_idx(cfg, stage_info, self.inst_id) - tile_idx = global_kv_tile_idx(cfg, local_tile_idx, seq_len_kv, cta_idx_kv) - tile_offset_k = tile_idx * Int32(cfg.tile_size_kv) - next_tile_offset_k = tile_offset_k + Int32(cfg.tile_size_kv) should_apply_dense_mask = ( (seq_len_kv % Int32(cfg.tile_size_kv)) != Int32(0) ) or (next_tile_offset_k > seq_len_kv) @@ -511,6 +552,14 @@ def update_softmax( should_apply_dense_mask = should_apply_dense_mask or ( next_tile_offset_k > min_seq_len_kv ) + if cutlass.const_expr(cfg.num_tokens_per_page == 1): + if cutlass.const_expr(cfg.logical_seq_len_q == 1): + # Cached route bits include the live prefix boundary. + should_apply_dense_mask = (mask_routes.low | mask_routes.high) != Int32( + 0 + ) + 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. @@ -531,7 +580,13 @@ def update_softmax( ) for reg_idx in cutlass.range_constexpr(num_s_regs_per_thread(cfg)): token_idx = tile_offset_k + local_col_base + Int32(reg_idx) - if token_idx >= row_seq_len_kv: + if invalid_sparse_token( + token_idx, + row_seq_len_kv, + mask_routes, + batch_idx, + sparse=cfg.num_tokens_per_page == 1, + ): s_vals[reg_idx] = neg_max_f32() else: local_idx_k0 = warp_idx * Int32(WARP_LANES) + ( @@ -571,42 +626,106 @@ def update_softmax( s_base = repeat_idx * 4 s_second_panel_base = q_repeats * 4 + s_base token_idx = tile_offset_k + local_idx_k0 - if token_idx >= seq_len_kv_0: + if invalid_sparse_token( + token_idx, + seq_len_kv_0, + mask_routes, + batch_idx, + sparse=cfg.num_tokens_per_page == 1, + ): s_vals[s_base + 0] = neg_max_f32() - if token_idx >= seq_len_kv_1: + if invalid_sparse_token( + token_idx, + seq_len_kv_1, + mask_routes, + batch_idx, + sparse=cfg.num_tokens_per_page == 1, + ): s_vals[s_base + 1] = neg_max_f32() token_idx = ( tile_offset_k + local_idx_k0 + Int32(SCORE_TOKENS_PER_QK_GROUP) ) - if token_idx >= seq_len_kv_0: + if invalid_sparse_token( + token_idx, + seq_len_kv_0, + mask_routes, + batch_idx, + sparse=cfg.num_tokens_per_page == 1, + ): s_vals[s_base + 2] = neg_max_f32() - if token_idx >= seq_len_kv_1: + if invalid_sparse_token( + token_idx, + seq_len_kv_1, + mask_routes, + batch_idx, + sparse=cfg.num_tokens_per_page == 1, + ): s_vals[s_base + 3] = neg_max_f32() token_idx = ( tile_offset_k + local_idx_k0 + Int32(2 * SCORE_TOKENS_PER_QK_GROUP) ) - if token_idx >= seq_len_kv_0: + if invalid_sparse_token( + token_idx, + seq_len_kv_0, + mask_routes, + batch_idx, + sparse=cfg.num_tokens_per_page == 1, + ): s_vals[s_second_panel_base + 0] = neg_max_f32() - if token_idx >= seq_len_kv_1: + if invalid_sparse_token( + token_idx, + seq_len_kv_1, + mask_routes, + batch_idx, + sparse=cfg.num_tokens_per_page == 1, + ): s_vals[s_second_panel_base + 1] = neg_max_f32() token_idx = ( tile_offset_k + local_idx_k0 + Int32(3 * SCORE_TOKENS_PER_QK_GROUP) ) - if token_idx >= seq_len_kv_0: + if invalid_sparse_token( + token_idx, + seq_len_kv_0, + mask_routes, + batch_idx, + sparse=cfg.num_tokens_per_page == 1, + ): s_vals[s_second_panel_base + 2] = neg_max_f32() - if token_idx >= seq_len_kv_1: + if invalid_sparse_token( + token_idx, + seq_len_kv_1, + mask_routes, + batch_idx, + sparse=cfg.num_tokens_per_page == 1, + ): s_vals[s_second_panel_base + 3] = neg_max_f32() if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): # Keeps-MMA-AB keeps the softmax state in scratch words indexed by # CTA thread, so only a warp-level max is needed here. local_max = old_max_vals[0] - for reg_idx in cutlass.range_constexpr(num_s_regs_per_thread(cfg)): - local_max = cute.math.max(local_max, s_vals[reg_idx], ftz=True) + if cutlass.const_expr(cfg.num_tokens_per_page == 1): + max0 = local_max + max1 = neg_max_f32() + max2 = neg_max_f32() + max3 = neg_max_f32() + for i in cutlass.range_constexpr(0, num_s_regs_per_thread(cfg), 4): + max0 = cute.math.max(max0, s_vals[i], ftz=True) + max1 = cute.math.max(max1, s_vals[i + 1], ftz=True) + max2 = cute.math.max(max2, s_vals[i + 2], ftz=True) + max3 = cute.math.max(max3, s_vals[i + 3], ftz=True) + local_max = cute.math.max( + cute.math.max(max0, max1, ftz=True), + cute.math.max(max2, max3, ftz=True), + ftz=True, + ) + else: + for reg_idx in cutlass.range_constexpr(num_s_regs_per_thread(cfg)): + local_max = cute.math.max(local_max, s_vals[reg_idx], ftz=True) local_max = cute.math.max( local_max, Float32( @@ -718,6 +837,21 @@ def update_softmax( ) self.new_max_state[scale_idx] = new_max_vals[scale_idx] for scale_idx in cutlass.range_constexpr(num_scale_groups): + if cutlass.const_expr(cfg.defer_sparse_max_update): + old_anchor = old_max_vals[scale_idx] + candidate = new_max_vals[scale_idx] + # Any common normalization anchor gives the same attention. + # Match FlashMLA's six-log2 bound (weights up to 64). + # Native FP8 retains exact maxima and scale 448. + # Initialize normally and update on larger score increases. + if old_anchor != neg_max_f32() and self.scale_softmax_log2 > Float32( + 0.0 + ): + increase_log2 = (candidate - old_anchor) * self.scale_softmax_log2 + limit = Float32(6.0) + if increase_log2 <= limit: + new_max_vals[scale_idx] = old_anchor + self.new_max_state[scale_idx] = new_max_vals[scale_idx] old_max_arr[scale_idx] = old_max_vals[scale_idx] sum_arr[scale_idx] = sum_vals[scale_idx] new_max_arr[scale_idx] = new_max_vals[scale_idx] @@ -744,6 +878,23 @@ def update_softmax_sum( # This consumer aux step runs after SmemP/TmemP has produced P and # published local_sum_arr. It updates the running denominator for the # next score tile without consuming a new score payload. + if cutlass.const_expr( + self.cfg.kernel_variant == "keeps_mma_ab" + and self.cfg.num_tokens_per_page == 1 + ): + old_max = old_max_arr[0] + new_max = new_max_arr[0] + old_sum = sum_arr[0] + local_sum = self._p_local_sum_arr[0] + exp_scale = cute.math.exp2( + self.scale_softmax_log2 * (old_max - new_max), + fastmath=True, + ) + updated_sum = exp_scale * old_sum + local_sum + self.new_max_state[0] = new_max + self.sum_state[0] = updated_sum + sum_arr[0] = updated_sum + return old_max_arr, sum_arr, new_max_arr, local_sum_arr, s_arr if cutlass.const_expr(self.cfg.kernel_variant == "keeps_mma_ab"): state_idx = cute.arch.thread_idx()[0] state_ptr = self._softmax_scratch.data_ptr(state_idx) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_softmax_stats.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_softmax_stats.py index d9c83631fbec..f6bec6761305 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_softmax_stats.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_softmax_stats.py @@ -23,7 +23,7 @@ import cutlass.cute as cute from cutlass import Float32, Int32 from cutlass.experimental.task_scheduling.enums import WorkAttr -from cutlass.experimental.task_scheduling.memory import TmemAllocation +from cutlass.experimental.task_scheduling.memory import TmemAllocation, SmemAllocation from cutlass.experimental.task_scheduling.resources import ( StageInfo, TaskLocalVariable, @@ -39,6 +39,7 @@ _TASK_CACHE_TMEM_BASE_OFFSET, decode_gen_task_cache, num_softmax_scale_groups, + smem_array, ) from ...helpers.math import ( neg_max_f32, @@ -51,6 +52,7 @@ MlaResource, ) + # ===================================================================== # TmemSoftmaxLocalResource — Local softmax stats in TMEM # ===================================================================== @@ -561,3 +563,75 @@ def track_global(self, stage_info: StageInfo): # This resource has no payload. Its producer work gives TS a named # dependency edge between the two softmax instances and correction. del stage_info + + +@dataclass(kw_only=True) +class SmemSoftmaxLocalResource(TmemSoftmaxLocalResource): + """Keep correction statistics alive independently of recycled QK/P slots. + + Sparse loads can change task pacing enough for a future QK to overwrite + the score TMEM before correction reads aliased statistics. The stats pipe + owns these small shared-memory slots until its consumer releases them. + """ + + def get_tmem_requirements(self): + return [] + + def get_smem_requirements(self): + if self._alloc is None: + self._alloc = SmemAllocation( + name=self.name + "_stats", + size_bytes=self.pipeline_config.num_stages + * 128 + * 2 + * num_softmax_scale_groups(self.cfg) + * 4, + alignment=16, + ) + return [self._alloc] + + @cute.jit + def _stats_pointer(self, stage_info): + storage = smem_array( + stage_info.context, + self._alloc, + Float32, + self.pipeline_config.num_stages + * 128 + * 2 + * num_softmax_scale_groups(self.cfg), + ) + local_thread = cute.arch.thread_idx()[0] & Int32(127) + return storage.data_ptr( + (stage_info.stage_idx * Int32(128) + local_thread) + * Int32(2 * num_softmax_scale_groups(self.cfg)) + ) + + @producer_work + @cute.jit + def store_loop_stats( + self, + stage_info: StageInfo, + *, + old_max_arr, + new_max_arr, + sum_arr, + inst_idx: cutlass.Constexpr[int], + ): + n = num_softmax_scale_groups(self.cfg) + values = tuple( + old_max_arr[i] if inst_idx == 0 else sum_arr[i] for i in range(n) + ) + values += tuple(new_max_arr[i] for i in range(n)) + self._stats_pointer(stage_info).store( + vector_from_scalars(values, Float32), + alignment=16 if n > 1 else 8, + ) + cute.arch.fence_view_async_shared() + + @cute.jit + def _load_stats_payload(self, stage_info: StageInfo): + return self._stats_pointer(stage_info).load( + count=2 * num_softmax_scale_groups(self.cfg), + alignment=16 if num_softmax_scale_groups(self.cfg) > 1 else 8, + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/tasks.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/tasks.py index bf7715ce4ddf..1d75ee237638 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/tasks.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/tasks.py @@ -80,7 +80,7 @@ def init_variables(self, context=None): self._local_warp_idx = self._warp_grp_thread_idx >> Int32(WARP_LANE_SHIFT) self._lane_idx = self._warp_grp_thread_idx & Int32(WARP_LANE_MASK) if context is not None and context.tmem_ptr_i32 is not None: - loaded = Int32(context.tmem_ptr_i32.load()) + loaded = self._load_allocated_tmem_base(context.tmem_ptr_i32) self._tmem_base_offset = cute.arch.make_warp_uniform( prims.shfl_sync( thread_mask=0xFFFFFFFF, @@ -91,6 +91,22 @@ def init_variables(self, context=None): ) ) + @cute.jit + def _load_allocated_tmem_base(self, pointer): + # TaskManager initializes every task before dispatch. Sparse loaders + # may already be running while the MMA warp allocates TMEM, so only + # the warps participating in the readiness barrier may read its result. + if cutlass.const_expr(self.cfg is not None and self.cfg.sparse_direct): + warp = cute.arch.warp_idx() + base = Int32(0) + if warp < Int32( + self.cfg.correction_warp_idx + self.cfg.correction_num_warps + ) or warp == Int32(self.cfg.mma_warp_idx): + base = Int32(pointer.load(is_volatile=True)) + return base + else: + return Int32(pointer.load()) + @cute.jit def make_task_cache(self): """Return the compact tuple of cached per-task thread values.""" @@ -151,6 +167,8 @@ def get_domain(self, tile_coord): self.cu_seqlens_q, batch_idx, ) + if cutlass.const_expr(self.cfg.fixed_sparse_kv_tiles > 0): + return Int32(self.domain_bias) total_kv_tiles = runtime_local_kv_tiles(self.cfg, seq_len_kv) remaining_kv_tiles = cute.math.max( total_kv_tiles - Int32(self.cfg.num_insts_kv), Int32(0) @@ -309,6 +327,63 @@ def init_work_tile_state(): ) +def create_short_merged_softmax_task(s0, s1, l0, l1, p0, p1, g0, g1, cfg, **kw): + @schedule + def merged(s0, s1, l0, l1, p0, p1, g0, g1): + def process(s, local, p, global_stats): + old, total, maximum, local_sum, scores = s.init_softmax_state() + p.init_materialize_state() + s.wait() + old, total, maximum, local_sum, scores = s.update_softmax( + old_max_arr=old, + sum_arr=total, + new_max_arr=maximum, + local_sum_arr=local_sum, + s_arr=scores, + section=MlaStage.Loop, + ) + s.release() + p.acquire() + p.materialize_p(new_max_arr=maximum, s_arr=scores, local_sum_arr=local_sum) + p.commit() + global_stats.track_global() + old, total, maximum, local_sum, scores = s.update_softmax_sum( + old_max_arr=old, + sum_arr=total, + new_max_arr=maximum, + local_sum_arr=local_sum, + s_arr=scores, + ) + return old, maximum, total + + def publish_final(local, state): + old, maximum, total = state + local.acquire() + local.store_loop_stats( + old_max_arr=old, new_max_arr=maximum, sum_arr=total, inst_idx=1 + ) + local.commit() + + with domain_loop(0, 1, 1): + state0 = process(s0, l0, p0, g0) + state1 = process(s1, l1, p1, g1) + publish_final(l0, state0) + publish_final(l1, state1) + + return MlaDecodeTask( + src_resources=[s0, s1], + dst_resources=[l0, l1, p0, p1, g0, g1], + cfg=cfg, + warp_idx=0, + num_warps=4, + num_registers=cfg.softmax_regs, + schedule=merged(s0, s1, l0, l1, p0, p1, g0, g1), + name="ShortMergedSoftmaxTask", + domain_bias=1, + **kw, + ) + + def create_keeps_mma_ab_softmax_task( tmem_s, tmem_softmax_local, @@ -830,8 +905,15 @@ def init_work_tile_state(): tail_o_stage_idx_0, tail_o_stage_idx_1, ) = work_tile_state - local0_state = load_initial_local_stats(tmem_softmax_local0, local0_state) - local1_state = load_initial_local_stats(tmem_softmax_local1, local1_state) + # With only one tile per stream there is no correction loop; + # consume only final statistics, published after both P tiles. + if not cfg.short_merged_softmax: + local0_state = load_initial_local_stats( + tmem_softmax_local0, local0_state + ) + local1_state = load_initial_local_stats( + tmem_softmax_local1, local1_state + ) with domain_loop(0, domain, 1): ( @@ -1014,7 +1096,7 @@ def page_offsets_schedule(smem_page_offsets, work_queue=None): smem_page_offsets.init_load_state() with work_tile_schedule_loop(work_queue, skip_if=work_tile_skip_if): - if cfg.kernel_variant == "keeps_mma_ab": + if cfg.single_kv_pipe: page_offsets_produce( smem_page_offsets, "load_k0", section=MlaStage.Head ) @@ -1022,11 +1104,27 @@ def page_offsets_schedule(smem_page_offsets, work_queue=None): page_offsets_produce( smem_page_offsets, "load_k0", section=MlaStage.Loop ) + if not cfg.sparse_reuse_kv: + page_offsets_produce( + smem_page_offsets, "load_v0", section=MlaStage.Loop + ) + if not cfg.sparse_reuse_kv: page_offsets_produce( - smem_page_offsets, "load_v0", section=MlaStage.Loop + smem_page_offsets, "load_v0", section=MlaStage.Tail ) + elif cfg.serial_swap_reuse: page_offsets_produce( - smem_page_offsets, "load_v0", section=MlaStage.Tail + smem_page_offsets, "load_k0", section=MlaStage.Head + ) + with domain_loop(0, domain, 1): + page_offsets_produce( + smem_page_offsets, "load_k1", section=MlaStage.Loop + ) + page_offsets_produce( + smem_page_offsets, "load_k0", section=MlaStage.Loop + ) + page_offsets_produce( + smem_page_offsets, "load_k1", section=MlaStage.Tail ) else: page_offsets_produce( @@ -1132,7 +1230,7 @@ def load_schedule_body( else: cached_page_ids = None reuse_delayed_v_page_ids = ( - smem_page_offsets is not None and cfg.kernel_variant != "keeps_mma_ab" + smem_page_offsets is not None and not cfg.single_kv_pipe ) def load_kv( @@ -1171,75 +1269,106 @@ def load_kv( smem_q.acquire() smem_q.load_q() smem_q.commit() - load_kv( - head_dim_stages=cfg.qk_head_dim_stages, - producer_label="load_k0", - section=MlaStage.Head, - page_id_slot=0, - ) - if cfg.kernel_variant != "keeps_mma_ab": + if cfg.serial_swap_reuse: load_kv( head_dim_stages=cfg.qk_head_dim_stages, - producer_label="load_k1", + producer_label="load_k0", section=MlaStage.Head, - page_id_slot=1, + page_id_slot=0, ) - - with domain_loop(0, domain, 1): - if cfg.kernel_variant == "keeps_mma_ab": + with domain_loop(0, domain, 1): load_kv( head_dim_stages=cfg.qk_head_dim_stages, - producer_label="load_k0", + producer_label="load_k1", section=MlaStage.Loop, + page_id_slot=1, ) load_kv( - head_dim_stages=cfg.v_head_dim_stages, - producer_label="load_v0", + head_dim_stages=cfg.qk_head_dim_stages, + producer_label="load_k0", section=MlaStage.Loop, + page_id_slot=0, + ) + load_kv( + head_dim_stages=cfg.qk_head_dim_stages, + producer_label="load_k1", + section=MlaStage.Tail, + page_id_slot=1, + ) + else: + load_kv( + head_dim_stages=cfg.qk_head_dim_stages, + producer_label="load_k0", + section=MlaStage.Head, + page_id_slot=0, + ) + if not cfg.single_kv_pipe: + load_kv( + head_dim_stages=cfg.qk_head_dim_stages, + producer_label="load_k1", + section=MlaStage.Head, + page_id_slot=1, ) - else: + + with domain_loop(0, domain, 1): + if cfg.single_kv_pipe: + load_kv( + head_dim_stages=cfg.qk_head_dim_stages, + producer_label="load_k0", + section=MlaStage.Loop, + ) + if not cfg.sparse_reuse_kv: + load_kv( + head_dim_stages=cfg.v_head_dim_stages, + producer_label="load_v0", + section=MlaStage.Loop, + ) + else: + if not cfg.sparse_reuse_kv: + load_kv( + head_dim_stages=cfg.v_head_dim_stages, + producer_label="load_v0", + section=MlaStage.Loop, + page_id_slot=0, + reuse_cached_page_ids=True, + ) + load_kv( + head_dim_stages=cfg.qk_head_dim_stages, + producer_label="load_k0", + section=MlaStage.Loop, + page_id_slot=0, + ) + if not cfg.sparse_reuse_kv: + load_kv( + head_dim_stages=cfg.v_head_dim_stages, + producer_label="load_v1", + section=MlaStage.Loop, + page_id_slot=1, + reuse_cached_page_ids=True, + ) + load_kv( + head_dim_stages=cfg.qk_head_dim_stages, + producer_label="load_k1", + section=MlaStage.Loop, + page_id_slot=1, + ) + + if not cfg.sparse_reuse_kv: load_kv( head_dim_stages=cfg.v_head_dim_stages, producer_label="load_v0", - section=MlaStage.Loop, + section=MlaStage.Tail, page_id_slot=0, reuse_cached_page_ids=True, ) - load_kv( - head_dim_stages=cfg.qk_head_dim_stages, - producer_label="load_k0", - section=MlaStage.Loop, - page_id_slot=0, - ) + if not cfg.single_kv_pipe and not cfg.sparse_reuse_kv: load_kv( head_dim_stages=cfg.v_head_dim_stages, producer_label="load_v1", - section=MlaStage.Loop, + section=MlaStage.Tail, page_id_slot=1, reuse_cached_page_ids=True, ) - load_kv( - head_dim_stages=cfg.qk_head_dim_stages, - producer_label="load_k1", - section=MlaStage.Loop, - page_id_slot=1, - ) - - load_kv( - head_dim_stages=cfg.v_head_dim_stages, - producer_label="load_v0", - section=MlaStage.Tail, - page_id_slot=0, - reuse_cached_page_ids=True, - ) - if cfg.kernel_variant != "keeps_mma_ab": - load_kv( - head_dim_stages=cfg.v_head_dim_stages, - producer_label="load_v1", - section=MlaStage.Tail, - page_id_slot=1, - reuse_cached_page_ids=True, - ) @schedule def load_schedule(smem_q, smem_kv): @@ -1514,67 +1643,150 @@ def mma_schedule( q_desc, q_desc_rope = smem_q.q_desc() tmem_s0.set_q_desc(q_desc=q_desc, q_desc_rope=q_desc_rope) tmem_s1.set_q_desc(q_desc=q_desc, q_desc_rope=q_desc_rope) - staged_qk_mma( - smem_kv, - tmem_s0, - head_dim_stages=cfg.qk_head_dim_stages, - consumer_label="k_desc_0", - ) - staged_qk_mma( - smem_kv, - tmem_s1, - head_dim_stages=cfg.qk_head_dim_stages, - consumer_label="k_desc_1", - ) - with domain_loop(0, domain, 1): - tmem_s0.acquire() + if cfg.serial_swap_reuse: + staged_qk_mma( + smem_kv, + tmem_s0, + head_dim_stages=cfg.qk_head_dim_stages, + consumer_label="k_desc_0", + release_k=False, + ) + with domain_loop(0, domain, 1): + staged_pv_mma( + smem_kv, + smem_p0, + tmem_o, + head_dim_stages=cfg.v_head_dim_stages, + consumer_label="v_desc_reused_0", + reuse_k=True, + producer_label="pv_mma_loop_0", + ) + staged_qk_mma( + smem_kv, + tmem_s1, + head_dim_stages=cfg.qk_head_dim_stages, + consumer_label="k_desc_1", + release_k=False, + ) + staged_pv_mma( + smem_kv, + smem_p1, + tmem_o, + head_dim_stages=cfg.v_head_dim_stages, + consumer_label="v_desc_reused_1", + reuse_k=True, + producer_label="pv_mma_loop_1", + ) + staged_qk_mma( + smem_kv, + tmem_s0, + head_dim_stages=cfg.qk_head_dim_stages, + consumer_label="k_desc_0", + release_k=False, + ) staged_pv_mma( smem_kv, smem_p0, tmem_o, head_dim_stages=cfg.v_head_dim_stages, - consumer_label="v_desc_0", - producer_label="pv_mma_loop_0", + consumer_label="v_desc_reused_0", + reuse_k=True, + producer_label="pv_mma_tail_0", ) staged_qk_mma( smem_kv, - tmem_s0, + tmem_s1, head_dim_stages=cfg.qk_head_dim_stages, - consumer_label="k_desc_0", - include_acquire=False, + consumer_label="k_desc_1", + release_k=False, ) - tmem_s1.acquire() staged_pv_mma( smem_kv, smem_p1, tmem_o, head_dim_stages=cfg.v_head_dim_stages, - consumer_label="v_desc_1", - producer_label="pv_mma_loop_1", + consumer_label="v_desc_reused_1", + reuse_k=True, + producer_label="pv_mma_tail_1", + ) + else: + staged_qk_mma( + smem_kv, + tmem_s0, + head_dim_stages=cfg.qk_head_dim_stages, + consumer_label="k_desc_0", + release_k=not cfg.sparse_reuse_kv, ) staged_qk_mma( smem_kv, tmem_s1, head_dim_stages=cfg.qk_head_dim_stages, consumer_label="k_desc_1", - include_acquire=False, + release_k=not cfg.sparse_reuse_kv, + ) + with domain_loop(0, domain, 1): + tmem_s0.acquire() + staged_pv_mma( + smem_kv, + smem_p0, + tmem_o, + head_dim_stages=cfg.v_head_dim_stages, + consumer_label=( + "v_desc_reused" if cfg.sparse_reuse_kv else "v_desc_0" + ), + reuse_k=cfg.sparse_reuse_kv, + producer_label="pv_mma_loop_0", + ) + staged_qk_mma( + smem_kv, + tmem_s0, + head_dim_stages=cfg.qk_head_dim_stages, + consumer_label="k_desc_0", + release_k=not cfg.sparse_reuse_kv, + include_acquire=False, + ) + tmem_s1.acquire() + staged_pv_mma( + smem_kv, + smem_p1, + tmem_o, + head_dim_stages=cfg.v_head_dim_stages, + consumer_label=( + "v_desc_reused" if cfg.sparse_reuse_kv else "v_desc_1" + ), + reuse_k=cfg.sparse_reuse_kv, + producer_label="pv_mma_loop_1", + ) + staged_qk_mma( + smem_kv, + tmem_s1, + head_dim_stages=cfg.qk_head_dim_stages, + consumer_label="k_desc_1", + release_k=not cfg.sparse_reuse_kv, + include_acquire=False, + ) + staged_pv_mma( + smem_kv, + smem_p0, + tmem_o, + head_dim_stages=cfg.v_head_dim_stages, + consumer_label=( + "v_desc_reused" if cfg.sparse_reuse_kv else "v_desc_0" + ), + reuse_k=cfg.sparse_reuse_kv, + producer_label="pv_mma_tail_0", + ) + staged_pv_mma( + smem_kv, + smem_p1, + tmem_o, + head_dim_stages=cfg.v_head_dim_stages, + consumer_label=( + "v_desc_reused" if cfg.sparse_reuse_kv else "v_desc_1" + ), + reuse_k=cfg.sparse_reuse_kv, + producer_label="pv_mma_tail_1", ) - staged_pv_mma( - smem_kv, - smem_p0, - tmem_o, - head_dim_stages=cfg.v_head_dim_stages, - consumer_label="v_desc_0", - producer_label="pv_mma_tail_0", - ) - staged_pv_mma( - smem_kv, - smem_p1, - tmem_o, - head_dim_stages=cfg.v_head_dim_stages, - consumer_label="v_desc_1", - producer_label="pv_mma_tail_1", - ) smem_q.release() captured_schedule = ( @@ -1615,11 +1827,11 @@ def mma_schedule( ) -def create_keeps_mma_ab_mma_task( +def create_single_kv_pipe_mma_task( smem_q, smem_kv, tmem_s, - tmem_p, + p_resource, tmem_o, work_queue: WorkQueue | None, cfg, @@ -1628,33 +1840,49 @@ def create_keeps_mma_ab_mma_task( task_class=MlaDecodeTask, **kw, ) -> Task: - """Create the keeps-MMA-AB single-pipe MMA task. + """Schedule one QK/PV stream with SMEM P (swap) or TMEM P (keep). - HEAD computes QK for K[0]. LOOP computes QK for K[n] before PV for - K[n-1], and TAIL drains PV for K[last]. + HEAD computes the first QK; TAIL drains the last PV. Retaining one BF16 + BK128 tile requires PV-before-next-QK; two retained tiles permit overlap. """ work_tile_skip_if = runtime_work_tile_skip_if(work_queue) + swap = cfg.one_insts_kv_swap + pv_mma = staged_pv_mma if swap else staged_pv_mma_tmem_p + pv_loop = "pv_mma_loop_0" if swap else "pv_mma_loop_tmem_p" + pv_tail = "pv_mma_tail_0" if swap else "pv_mma_tail_tmem_p" + v_desc = "v_desc_reused_0" if swap else "v_desc_reused" + if not cfg.sparse_reuse_kv: + v_desc = "v_desc_0" + pv_before_qk = ( + cfg.sparse_reuse_kv + and not cfg.is_fp8_qkv() + and (swap or cfg.tile_size_kv == 128) + ) @schedule def mma_schedule( smem_q, smem_kv, tmem_s, - tmem_p, + p_resource, tmem_o, work_queue=None, ): - """Captured keeps-MMA-AB schedule with TMEM P.""" + """Capture identical pipeline ownership for the two P layouts.""" smem_q.init_descriptor_state() smem_kv.init_descriptor_state() tmem_o.init_mma_state() - tmem_p.init_stage_state() + if swap: + p_resource.init_descriptor_state() + else: + p_resource.init_stage_state() with work_tile_schedule_loop(work_queue, skip_if=work_tile_skip_if): if work_queue is not None: tmem_s.reset_softmax_work_tile_state() - tmem_p.init_stage_work_tile_state() + if not swap: + p_resource.init_stage_work_tile_state() smem_q.wait() q_desc, q_desc_rope = smem_q.q_desc() tmem_s.set_q_desc(q_desc=q_desc, q_desc_rope=q_desc_rope) @@ -1663,38 +1891,53 @@ def mma_schedule( tmem_s, head_dim_stages=cfg.qk_head_dim_stages, consumer_label="k_desc_0", + release_k=not cfg.sparse_reuse_kv, ) with domain_loop(0, domain, 1): + if pv_before_qk: + pv_mma( + smem_kv, + p_resource, + tmem_o, + head_dim_stages=cfg.v_head_dim_stages, + consumer_label=v_desc, + reuse_k=True, + producer_label=pv_loop, + ) staged_qk_mma( smem_kv, tmem_s, head_dim_stages=cfg.qk_head_dim_stages, consumer_label="k_desc_0", + release_k=not cfg.sparse_reuse_kv, ) - staged_pv_mma_tmem_p( - smem_kv, - tmem_p, - tmem_o, - head_dim_stages=cfg.v_head_dim_stages, - consumer_label="v_desc_0", - producer_label="pv_mma_loop_tmem_p", - ) - staged_pv_mma_tmem_p( + if not pv_before_qk: + pv_mma( + smem_kv, + p_resource, + tmem_o, + head_dim_stages=cfg.v_head_dim_stages, + consumer_label=v_desc, + reuse_k=cfg.sparse_reuse_kv, + producer_label=pv_loop, + ) + pv_mma( smem_kv, - tmem_p, + p_resource, tmem_o, head_dim_stages=cfg.v_head_dim_stages, - consumer_label="v_desc_0", - producer_label="pv_mma_tail_tmem_p", + consumer_label=v_desc, + reuse_k=cfg.sparse_reuse_kv, + producer_label=pv_tail, ) smem_q.release() schedule_result = ( - mma_schedule(smem_q, smem_kv, tmem_s, tmem_p, tmem_o) + mma_schedule(smem_q, smem_kv, tmem_s, p_resource, tmem_o) if work_queue is None - else mma_schedule(smem_q, smem_kv, tmem_s, tmem_p, tmem_o, work_queue) + else mma_schedule(smem_q, smem_kv, tmem_s, p_resource, tmem_o, work_queue) ) - src = [smem_q, smem_kv, tmem_p] + src = [smem_q, smem_kv, p_resource] if work_queue is not None: src.append(work_queue) return task_class( diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/separate_reduction.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/separate_reduction.py index 01bd9e49152d..fe8706c9e4a2 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/separate_reduction.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/separate_reduction.py @@ -26,7 +26,9 @@ import cutlass import cutlass.cute as cute -from cutlass import BFloat16, Float16, Float32, Int32 +from cutlass import BFloat16, Float16, Float32, Int32, Int64 + +from .mla_decode.helpers.constants import LN_2, LOG2_E @cute.jit @@ -98,3 +100,124 @@ def unpack_normalized_vec8( if cutlass.const_expr(use_bf16_partial): return regs_vec.bitcast(BFloat16).to(Float32) return regs_vec.bitcast(Float16).to(Float32) + + +@cute.jit +def attention_sink_scale(log2_lse, sink): + """Denominator-only sink mass; keep the KV LSE independent of the sink.""" + scale = Float32(0) + if log2_lse != Float32(-Float32.inf) and sink != Float32(Float32.inf): + delta = sink * Float32(LOG2_E) - log2_lse + z = cute.math.exp2(-cute.math.abs(delta), fastmath=True) + inverse = cute.math.rcp(Float32(1) + z, approx=True) + scale = z * inverse if delta > Float32(0) else inverse + return scale + + +class MergeMlaSourceStates: + def __init__(self, independent_sources: bool, output_quant=None): + self.independent_sources = independent_sources + self.output_quant = output_quant + + @cute.jit + def __call__( + self, + part_s, + part_c, + lse_s, + lse_c, + count_s, + count_c, + sinks, + out, + lse, + stream, + output_quant_args=None, + ): + self.finish( + part_s, + part_c, + lse_s, + lse_c, + count_s, + count_c, + sinks, + out, + lse, + output_quant_args, + ).launch( + grid=((out.shape[0] * out.shape[1] + 3) // 4, 1, 1), + block=(128, 1, 1), + stream=stream, + ) + + @cute.kernel + def finish( + self, + part_s, + part_c, + lse_s, + lse_c, + count_s, + count_c, + sinks, + out, + lse, + output_quant_args=None, + ): + block, _, _ = cute.arch.block_idx() + tid = cute.arch.thread_idx()[0] + lane = tid % 32 + # One warp handles one output head. Group four heads per CTA to + # amortize launch overhead and avoid repeating normalization 128 times. + row = Int64(block) * 4 + Int64(tid // 32) + q = row // out.shape[1] + h = row % out.shape[1] + if row < out.shape[0] * out.shape[1]: + have_s = Int32(count_s[q]) > 0 + have_c = (Int32(count_c[q]) > 0) if self.independent_sources else False + l0 = Float32(lse_s[q, h]) if have_s else Float32(-Float32.inf) + l1 = Float32(lse_c[q, h]) if have_c else Float32(-Float32.inf) + ws = Float32(0) + wc = Float32(0) + public_lse = Float32(-Float32.inf) + if have_s or have_c: + kv_lse, ws, wc = merge_log2_lse(l0, l1) + public_lse = kv_lse * Float32(LN_2) + sink_scale = attention_sink_scale(kv_lse, Float32(sinks[h])) + ws *= sink_scale + wc *= sink_scale + width = cutlass.const_expr(4 if self.output_quant is not None else 2) + for chunk in cutlass.range_constexpr(512 // (32 * width)): + column = lane * width + chunk * (32 * width) + values = cutlass.Array(Float32, width, space=cutlass.AddressSpace.rmem) + for j in cutlass.range_constexpr(width): + values[j] = Float32(0) + if have_s: + offset = q * part_s.stride[0] + h * part_s.stride[1] + column + pair = cutlass.Pointer( + part_s.iterator + offset, dtype=part_s.element_type + ).load(count=width, alignment=width * 2) + for j in cutlass.range_constexpr(width): + values[j] = Float32(pair[j]) * ws + if have_c: + offset = q * part_c.stride[0] + h * part_c.stride[1] + column + pair = cutlass.Pointer( + part_c.iterator + offset, dtype=part_c.element_type + ).load(count=width, alignment=width * 2) + for j in cutlass.range_constexpr(width): + values[j] += Float32(pair[j]) * wc + if cutlass.const_expr(self.output_quant is not None): + self.output_quant.store(output_quant_args, values, row, column) + else: + packed = cutlass.Array( + out.element_type, 2, space=cutlass.AddressSpace.rmem + ) + for j in cutlass.range_constexpr(2): + packed[j] = values[j].to(out.element_type) + offset = q * out.stride[0] + h * out.stride[1] + column + cutlass.Pointer( + out.iterator + offset, dtype=out.element_type + ).store(packed.load(0, 2), alignment=4) + if lane == 0: + lse[q, h] = public_lse 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 3d5fb159a329..70c3d0e6423d 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py @@ -16,17 +16,16 @@ """Task-scheduled paged MLA decode with a plan/run lifecycle.""" from collections.abc import Callable -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace +import copy +import math +import struct import functools -from typing import Any, Literal, Optional, cast +from typing import Any, Literal, NamedTuple, Optional, cast import torch from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api -from flashinfer.trace.templates.attention import ( - prims_ts_decode_mla_one_shot_trace_dispatch, -) - from .decode import ( _WorkspaceSection, _align_up, @@ -44,6 +43,8 @@ ) +from .kernels.mla_decode.sparse_policy import select_sparse_mla_profile + _COMPILE_OPTIONS = "--enable-tvm-ffi --opt-level 2" _MLA_LATENT_DIM = 512 _MLA_ROPE_DIM = 64 @@ -85,6 +86,7 @@ class _MLADecodeCompileSpec: has_kernel_workspace: bool split_kv: int kernel: Any = field(compare=False, hash=False, repr=False) + device_scales: bool = False @dataclass(frozen=True) @@ -136,6 +138,10 @@ class _MLARuntime: num_physical_pages: int bmm1_scale: float bmm2_scale: float + extra_cache: Optional[torch.Tensor] = None + scale_params: Optional[torch.Tensor] = None + sparse_inputs: Optional[tuple] = None + output_quant_args: Optional[tuple] = None def _make_mla_workspace_layout( @@ -533,21 +539,24 @@ def _validate_query( num_heads: Optional[int] = None, max_seq_len_q: Optional[int] = None, q_dtype: Optional[torch.dtype] = None, + query_dim: int = _MLA_QUERY_DIM, ) -> None: if not isinstance(query, torch.Tensor): raise TypeError("query must be a torch.Tensor") expected_rank = 3 if packed_query else 4 if query.ndim != expected_rank: - expected_shape = "[total_q, H, 576]" if packed_query else "[B, SQ, H, 576]" + expected_shape = ( + f"[total_q, H, {query_dim}]" if packed_query else f"[B, SQ, H, {query_dim}]" + ) raise ValueError(f"query must have shape {expected_shape}") if packed_query: if int(query.shape[1]) <= 0: raise ValueError("query head extent must be positive") elif any(int(extent) <= 0 for extent in query.shape[:-1]): raise ValueError("query row and head extents must be positive") - if query.shape[-1] != _MLA_QUERY_DIM: + if query.shape[-1] != query_dim: raise ValueError( - f"query last dimension must be {_MLA_QUERY_DIM}, got {query.shape[-1]}" + f"query last dimension must be {query_dim}, got {query.shape[-1]}" ) if query.dtype not in _SUPPORTED_INPUT_DTYPES: raise NotImplementedError( @@ -596,7 +605,9 @@ def _validate_query( f"query dtype must match the plan ({q_dtype}), got {query.dtype}" ) if not query.is_contiguous(): - layout = "[total_q, H, 576]" if packed_query else "[B, SQ, H, 576]" + layout = ( + f"[total_q, H, {query_dim}]" if packed_query else f"[B, SQ, H, {query_dim}]" + ) raise ValueError(f"query must be compact in {layout} layout") _validate_16byte_alignment(query, "query") @@ -1023,6 +1034,7 @@ def _make_mla_decode_compile_spec( output_dtype_key: str, max_seq_len_q: int, packed_query: bool, + device_scales: bool = False, ) -> _MLADecodeCompileSpec: """Keep policy resolution plan-specific and JIT identity batch-free.""" @@ -1040,6 +1052,7 @@ def _make_mla_decode_compile_spec( has_kernel_workspace=launch_spec.kernel_workspace_bytes > 0, split_kv=launch_spec.split_kv, kernel=launch_spec.kernel, + device_scales=device_scales, ) @@ -1056,6 +1069,10 @@ def _get_compiled_mla_decode( num_heads = compile_spec.num_heads kv_lora_rank = compile_spec.kv_lora_rank qk_rope_head_dim = compile_spec.qk_rope_head_dim + query_dim = kv_lora_rank + qk_rope_head_dim + # The D512 path has no separate RoPE stage. Use an unused full-width + # alias in the callable ABI rather than a zero-extent tensor map. + rope_view_dim = qk_rope_head_dim or kv_lora_rank page_size = compile_spec.page_size max_seq_len_q = compile_spec.max_seq_len_q packed_query = compile_spec.packed_query @@ -1072,21 +1089,21 @@ def _get_compiled_mla_decode( # These fake tensors pin the public ABI while allowing runtime page counts, # table widths/row strides, and batch metadata pointers to vary. - q_stride_h = _MLA_QUERY_DIM - q_stride_q = num_heads * _MLA_QUERY_DIM + q_stride_h = query_dim + q_stride_q = num_heads * query_dim q_latent_shape: tuple[int, ...] q_rope_shape: tuple[int, ...] q_stride: tuple[int, ...] if packed_query: q_latent_shape = (num_heads, kv_lora_rank, runtime_total_q) - q_rope_shape = (num_heads, qk_rope_head_dim, runtime_total_q) + q_rope_shape = (num_heads, rope_view_dim, runtime_total_q) q_stride = (q_stride_h, 1, q_stride_q) else: q_stride_batch = max_seq_len_q * q_stride_q q_latent_shape = (num_heads, kv_lora_rank, max_seq_len_q, batch_size) q_rope_shape = ( num_heads, - qk_rope_head_dim, + rope_view_dim, max_seq_len_q, batch_size, ) @@ -1097,8 +1114,8 @@ def _get_compiled_mla_decode( q_rope_fake = cute.runtime.make_fake_tensor( qkv_dtype, q_rope_shape, stride=q_stride, assumed_align=16 ) - cache_token_stride = _MLA_QUERY_DIM - cache_page_stride = page_size * _MLA_QUERY_DIM + cache_token_stride = query_dim + cache_page_stride = page_size * query_dim c_latent_fake = cute.runtime.make_fake_tensor( qkv_dtype, (page_size, kv_lora_rank, physical_pages), @@ -1107,7 +1124,11 @@ def _get_compiled_mla_decode( ) c_rope_fake = cute.runtime.make_fake_tensor( qkv_dtype, - (page_size, qk_rope_head_dim, physical_pages), + ( + page_size, + rope_view_dim, + cute.sym_int() if page_size == 1 else physical_pages, + ), stride=(cache_token_stride, 1, cache_page_stride), assumed_align=16, ) @@ -1165,8 +1186,44 @@ def _get_compiled_mla_decode( stride_order=(0,), assumed_align=4, ) + scale_params_fake = ( + cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (cute.sym_int(),), stride_order=(0,), assumed_align=4 + ) + if compile_spec.device_scales + else None + ) + if getattr(kernel, "direct_sparse", False): + + def sparse_indices(capacity): + return cute.runtime.make_fake_tensor( + cutlass.Int32, + (batch_size, max(1, capacity)), + stride=(cute.sym_int64(), 1), + assumed_align=4, + ) + + scalar = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (1,), stride_order=(0,), assumed_align=4 + ) + sink = cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (num_heads,), stride_order=(0,), assumed_align=4 + ) + scale_params_fake = ( + sparse_indices(kernel.direct_sparse_capacities[0]), + sparse_indices(kernel.direct_sparse_capacities[1]), + cache_seqs_fake, + cache_seqs_fake, + scalar, + scalar, + scalar, + scalar, + sink, + ) stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + output_quant_args = _fake_output_quant_args(getattr(kernel, "output_quant", None)) + # Task objects carry loop-local state through generated control flow, so # select the public staged frontend for this compilation. with torch.cuda.device(device_index): @@ -1183,10 +1240,11 @@ def _get_compiled_mla_decode( cutlass.Int32(compile_spec.split_kv), cache_seqs_fake, qo_indptr_fake, - None, + scale_params_fake, cutlass.Float32(1.0), cutlass.Float32(1.0), stream_fake, + output_quant_args=output_quant_args, options=_COMPILE_OPTIONS, ) return compiled @@ -1387,6 +1445,11 @@ def _launch_mla_decode( lse_kernel = workspace.lse.permute(2, 1, 0) c_latent = runtime.normalized_cache[..., :kv_lora_rank].permute(1, 2, 0) c_rope = runtime.normalized_cache[..., kv_lora_rank:].permute(1, 2, 0) + if runtime.query.shape[-1] == kv_lora_rank: + q_rope = q_latent + c_rope = c_latent + if runtime.extra_cache is not None: + c_rope = runtime.extra_cache.permute(1, 2, 0) page_offsets = block_tables.transpose(0, 1) compiled( q_latent, @@ -1400,9 +1463,16 @@ def _launch_mla_decode( split_kv, seq_lens, qo_indptr, - None, + runtime.sparse_inputs + if runtime.sparse_inputs is not None + else runtime.scale_params, runtime.bmm1_scale, runtime.bmm2_scale, + **( + {"output_quant_args": runtime.output_quant_args} + if runtime.output_quant_args is not None + else {} + ), ) return runtime.out @@ -1681,7 +1751,7 @@ def run( ) -@flashinfer_experimental_api(trace=prims_ts_decode_mla_one_shot_trace_dispatch) +@flashinfer_experimental_api def batch_mla_decode_with_paged_kv_cache( query: torch.Tensor, kv_cache: torch.Tensor, @@ -1885,3 +1955,1136 @@ def batch_mla_decode_with_paged_kv_cache( "batch_mla_decode_with_paged_kv_cache", "get_prims_ts_batch_mla_decode_workspace_size", ] + + +class SparseMLAPreparedMetadata(NamedTuple): + """Caller-owned storage-row indices and prepared attention metadata. + + Indices are int32 [rows, capacity], with -1 invalid; lengths are int32 + [rows]. Primary and extra indices address their own flattened native pool. + No block-table lookup or page-stride conversion occurs in run. + + Combined-route schedules also consume routes [passes, rows, route_capacity], + execution_lengths/valid_counts [passes, rows], and FP32 scale_params + [passes, 2 + heads + rows]. Route bit 31 selects the extra pool; bits 0:31 + encode the storage row, with 0x7fffffff invalid. Each source's execution + span is rounded to 128. Empty rows execute one masked slot. Scale rows + contain [QK scale, PV scale, sinks[heads], valid_counts[rows]]. Two passes + are required when the sources use independent KV descales. + + Preparation must refresh all dependent fields when routes, lengths, + scales or sinks change. Buffers must remain alive and unaliased with + writable attention buffers through graph replay. Applications may supply + these buffers from any preparer satisfying this contract. + """ + + indices: torch.Tensor + lengths: torch.Tensor + extra_indices: torch.Tensor | None = None + extra_lengths: torch.Tensor | None = None + routes: torch.Tensor | None = None + execution_lengths: torch.Tensor | None = None + valid_counts: torch.Tensor | None = None + scale_params: torch.Tensor | None = None + + +def _fake_sparse_tensor(dtype, shape): + """Compact, four-byte-aligned ABI shared by sparse finishing kernels.""" + import cutlass.cute as cute + + return cute.runtime.make_fake_compact_tensor( + dtype, + shape, + stride_order=tuple(reversed(range(len(shape)))), + assumed_align=4, + ) + + +@functools.cache +def _compile_finish(device_index, heads, independent, output_quant=None): + import cutlass + import cutlass.cute as cute + from .kernels.separate_reduction import MergeMlaSourceStates + + rows = cute.sym_int() + + with torch.cuda.device(device_index): + partial = _fake_sparse_tensor(cutlass.BFloat16, (rows, heads, 512)) + lse = _fake_sparse_tensor(cutlass.Float32, (rows, heads)) + lens = _fake_sparse_tensor(cutlass.Int32, (rows,)) + return cute.compile[cute.FrontendNext]( + MergeMlaSourceStates(independent, output_quant), + partial, + partial, + lse, + lse, + lens, + lens, + _fake_sparse_tensor(cutlass.Float32, (heads,)), + partial, + lse, + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + output_quant_args=_fake_output_quant_args(output_quant), + options=_COMPILE_OPTIONS, + ) + + +def _make_sparse_output_quant(heads, block_size, scale_format, groups): + if block_size is None: + if groups is not None or scale_format != "ue8m0": + raise ValueError( + "output quantization options require output_quant_block_size" + ) + return None + block_size = _validate_positive_int(block_size, "output_quant_block_size") + if block_size not in (32, 128): + raise ValueError("output_quant_block_size must be 32, 128, or None") + if scale_format not in ("fp32", "ue8m0"): + raise ValueError("out_scale_format must be 'fp32' or 'ue8m0'") + groups = _validate_positive_int(groups, "num_output_groups") + if heads % groups: + raise ValueError("num_output_groups must divide num_heads") + from .kernels.mla_decode.helpers.output_quant import MlaOutputQuant + + return MlaOutputQuant(block_size, scale_format == "ue8m0", heads, groups) + + +def _fake_output_quant_args(quant): + if quant is None: + return None + import cutlass + import cutlass.cute as cute + + rows = cute.sym_int() + padded_rows = cute.sym_int(divisibility=4) + columns = quant.num_heads // quant.num_groups * 512 // quant.block_size + if quant.ue8m0: + columns //= 4 + return ( + cute.runtime.make_fake_compact_tensor( + cutlass.Float8E4M3FN, + (cute.sym_int64(),), + stride_order=(0,), + assumed_align=16, + ), + _fake_sparse_tensor(cutlass.Int32, (rows,)), + cute.runtime.make_fake_compact_tensor( + cutlass.Float32, (cute.sym_int(), 64), stride_order=(1, 0), assumed_align=16 + ), + cute.runtime.make_fake_tensor( + cutlass.Int32 if quant.ue8m0 else cutlass.Float32, + (rows, quant.num_groups, columns), + stride=(1, columns * padded_rows, padded_rows), + assumed_align=16, + ), + ) + + +def _prepare_sparse_output_quant( + quant, query, positions, cos_sin, out, scales, validate +): + """Bind the token-major output and column-major block scales.""" + prefix = tuple(query.shape[:-2]) + rows, device = math.prod(prefix), query.device + channels = quant.num_heads // quant.num_groups * 512 + columns = channels // quant.block_size // (4 if quant.ue8m0 else 1) + padded_rows = (rows + 3) // 4 * 4 + if ( + positions is None + or positions.shape != prefix + or positions.dtype != torch.int32 + or positions.device != device + or not positions.is_contiguous() + ): + raise ValueError( + "token_positions must be contiguous CUDA INT32 matching the query prefix" + ) + if ( + cos_sin is None + or cos_sin.ndim != 2 + or cos_sin.shape[0] == 0 + or cos_sin.shape[1] != 64 + or cos_sin.dtype != torch.float32 + or cos_sin.device != device + or not cos_sin.is_contiguous() + or cos_sin.data_ptr() % 16 + ): + raise ValueError( + "cos_sin_cache must be 16-byte aligned contiguous CUDA FP32[P,64]" + ) + if ( + validate + and rows + and ((positions < 0) | (positions >= cos_sin.shape[0])).any().item() + ): + raise ValueError("token_positions must index cos_sin_cache") + shape = (rows, quant.num_groups, channels) + # Physical storage is group-major [G, T, C] (as in vLLM and TensorRT-LLM); + # the public tensor is its token-major [T, G, C] view. + if out is None: + out = torch.empty( + (quant.num_groups, rows, channels), dtype=torch.float8_e4m3fn, device=device + ).transpose(0, 1) + if ( + out.shape != shape + or out.dtype != torch.float8_e4m3fn + or out.device != device + or out.stride(2) != 1 + or out.stride(0) != channels + or (rows > 1 and out.stride(1) != rows * channels) + or out.data_ptr() % 16 + ): + raise ValueError( + f"out must be the [T, G, C] view of a contiguous group-major CUDA " + f"E4M3 [G, T, C] tensor ({shape}), aligned to 16 bytes" + ) + scale_dtype = torch.int32 if quant.ue8m0 else torch.float32 + scale_shape = (rows, quant.num_groups, columns) + scale_stride = (1, columns * padded_rows, padded_rows) + if scales is None: + scales = torch.empty( + (quant.num_groups, columns, padded_rows), dtype=scale_dtype, device=device + ).permute(2, 0, 1)[:rows] + if rows == 0: + scales = scales.as_strided(scale_shape, scale_stride) + if ( + scales.shape != scale_shape + or scales.stride() != scale_stride + or scales.dtype != scale_dtype + or scales.device != device + or scales.data_ptr() % 16 + ): + raise ValueError( + f"out_scale must be 16-byte aligned CUDA {scale_dtype}{scale_shape} " + f"with strides {scale_stride}" + ) + return ( + out, + scales, + (out.as_strided((out.numel(),), (1,)), positions.view(-1), cos_sin, scales), + ) + + +def _resolve_sparse_mla_plan( + device, + batch_size, + num_heads, + *, + max_topk, + max_extra_topk=0, + max_seq_len_q=1, + packed_query=False, + q_data_type=torch.bfloat16, + kv_layout="NHD", + has_sinks=False, + return_lse=False, + assume_valid_prefix=False, + output_quant_block_size=None, + out_scale_format="ue8m0", + num_output_groups=None, +): + """Resolve sparse geometry and scratch sizing without allocation/compilation.""" + import cutlass + from .kernels.mla_decode.throughput_latency_1cta.config import ( + MlaProfile, + compute_workspace_size, + ) + from .kernels.mla_decode.throughput_latency_1cta.kernel import ( + ThroughputLatencyMlaDecodeTs, + ) + from .kernels.mla_decode.throughput_2cta.config import ( + compute_workspace_size as workspace_2cta, + ) + from .kernels.mla_decode.throughput_2cta.kernel import ( + MlaDecodeTs, + ) + + if not isinstance(assume_valid_prefix, bool): + raise TypeError("assume_valid_prefix must be bool") + device, _ = _resolve_cuda_device(device) + if torch.cuda.get_device_capability(device) not in ((10, 0), (10, 3)): + raise NotImplementedError("sparse TS MLA requires SM100 or SM103") + for name, value in ( + ("batch_size", batch_size), + ("num_heads", num_heads), + ("max_seq_len_q", max_seq_len_q), + ): + if not isinstance(value, int) or isinstance(value, bool) or value <= 0: + raise ValueError(f"{name} must be a positive integer") + if num_heads > 128: + raise ValueError("num_heads must be at most 128") + output_quant = _make_sparse_output_quant( + num_heads, output_quant_block_size, out_scale_format, num_output_groups + ) + for value in (max_topk, max_extra_topk): + if not isinstance(value, int) or isinstance(value, bool) or value < 0: + raise ValueError("source capacities must be nonnegative integers") + if max_topk == 0: + raise ValueError("the primary selected-slot capacity must be positive") + if q_data_type not in (torch.bfloat16, torch.float8_e4m3fn): + raise ValueError("native BF16 or E4M3 is required") + if kv_layout not in ("NHD", "HND"): + raise ValueError("kv_layout must be NHD or HND") + max_rows = batch_size * max_seq_len_q + capacity = max( + 256, + ((max_topk + 127) // 128 + (max_extra_topk + 127) // 128) * 128, + ) + if max_rows * num_heads >= 2**31 or capacity >= 2**31 - 32768: + raise ValueError("query/route extent exceeds the int32 kernel domain") + sm_count = torch.cuda.get_device_properties(device).multi_processor_count + profile = select_sparse_mla_profile( + rows=max_rows, + heads=num_heads, + query_length=max_seq_len_q, + capacity=capacity, + dtype="bf16" if q_data_type == torch.bfloat16 else "fp8", + sm_count=sm_count, + ) + family, tile = profile.family, profile.tile_size_q + splits, head_dim_ctas = profile.split_kv, profile.head_dim_ctas + dtype_name = "bf16" if q_data_type == torch.bfloat16 else "e4m3" + if family == "2cta": + kernel = MlaDecodeTs( + page_size=1, + max_active_clusters=sm_count // 2, + is_persistent=profile.scheduler != "nonpersistent", + is_var_seq=False, + is_var_split_kv=False, + static_split_kv=splits, + qkv_dtype=dtype_name, + out_dtype="bf16", + rope_dim=0, + num_heads=num_heads, + seq_len_q=1, + batch_size=max_rows, + mask_type="dense", + device_scales=True, + ) + kernel.sparse_gather_warps = profile.gather_issue_warps + kernel.sparse_kv_stages = profile.kv_pipeline_stages + kernel.sparse_uniform_pages = profile.uniform_offset_cache + kernel.sparse_balanced_registers = profile.balanced_registers + core_bytes = workspace_2cta( + tile_size_q=128, + num_q_tiles=1, + latent_dim=512, + batch_size=max_rows, + split_kv=splits, + partial_o_dtype=cutlass.BFloat16, + lse_dtype=cutlass.Float32, + ) + else: + kernel_profile = MlaProfile( + name="sparse", + kernel_variant="swaps_mma_ab" if family == "swap" else "keeps_mma_ab", + tile_size_q=tile, + num_ctas_per_head_dim=head_dim_ctas, + num_ctas_per_seq_kv=splits, + use_multi_ctas_kv=int(splits > 1), + use_cluster_reduction=int(profile.reduction == "cluster"), + use_persistent_scheduler=int(profile.scheduler != "nonpersistent"), + use_clc_dynamic_persistent_scheduler=int(profile.scheduler == "clc"), + ) + kernel = ThroughputLatencyMlaDecodeTs( + batch_size=max_rows, + num_heads=tile, + seq_len_q=(num_heads + tile - 1) // tile, + seq_len_k=capacity, + rope_dim=0, + page_size=1, + max_active_clusters=sm_count, + qkv_dtype=dtype_name, + out_dtype="bf16", + profile=kernel_profile, + reduction_mode=profile.reduction, + logical_num_heads=num_heads, + logical_seq_len_q=1, + tile_size_q=tile, + sparse_profile=profile, + mask_type="dense", + device_scales=True, + ) + core_bytes = compute_workspace_size( + cfg=kernel._make_config(), + partial_o_dtype=cutlass.BFloat16, + lse_dtype=cutlass.Float32, + ) + spec = _make_mla_decode_compile_spec( + _MLADecodeLaunchSpec(kernel, (), core_bytes, splits), + device_index=device.index, + num_heads=num_heads, + kv_lora_rank=512, + qk_rope_head_dim=0, + page_size=1, + q_dtype_key=str(q_data_type).removeprefix("torch."), + output_dtype_key="bfloat16", + max_seq_len_q=1, + packed_query=False, + device_scales=True, + ) + # Only a second source can require an independent attention pass. + source_slots = 2 if max_extra_topk else 1 + sections = {} + byte_end = 0 + for name, shape, dtype in ( + ("core", (core_bytes,), torch.int8), + ("partial", (source_slots, max_rows, 1, num_heads, 512), torch.bfloat16), + # Both source views must satisfy the core's 16-byte base alignment, + # including H6/H12 and odd maximum query counts. + ("lse", (source_slots, (max_rows + 3) // 4 * 4, 1, num_heads), torch.float32), + ("public_lse", (max_rows, num_heads), torch.float32), + ): + sections[name], byte_end = _append_workspace_section(byte_end, shape, dtype) + return device, profile, spec, sections, byte_end, capacity, output_quant + + +class BatchSparseMLADecodePagedTSWrapper: + """Plan native BF16/E4M3 attention over one or two pools, then bind metadata. + + Q is [B,Sq,H,512] or packed [total_q,H,512]. Caller-prepared metadata + contains int32 [rows, capacity] storage-row indices and live lengths. + Indices already include physical page strides; -1 is masked. Both source + lists participate in one attention distribution. Packed KV is unsupported. + + ``output_quant_block_size=32`` or ``128`` fuses conjugate output RoPE on + D[448:512] and E4M3 quantization after the final attention reduction. Q/KV + remain prepared native tensors. Supply the local ``num_output_groups`` G + and ``out_scale_format="ue8m0"`` (packed exponent bytes) or ``"fp32"``. + This option does not change attention dispatch or softmax correction. + + Planning with ``assume_valid_prefix=True`` promises that every index before + each live source length is valid (no -1 holes). Direct FP8 2CTA kernels + then derive masks from lengths without reading indices or issuing ballots. + Other kernels retain their generic mask path; no gain was established. + ``run(validate=True)`` checks the promise; graph replay with validation + disabled must preserve it. Tail entries beyond the lengths are ignored. + + A wrapper/workspace permits one in-flight run. Graphs require a completed + warmup, stable addresses, preallocated outputs, and validate=False. + Input/output/workspace storage must not overlap; this is an unchecked + caller precondition. Separate wrappers/workspaces are needed per stream. + """ + + @flashinfer_experimental_api + def __init__(self, workspace_buffer=None): + self._workspace_buffer = workspace_buffer + self._state = None + self._scalar_cache = {} + self.workspace_size_bytes = 0 + + def plan( + self, + device, + batch_size, + num_heads, + *, + max_topk, + max_extra_topk=0, + max_seq_len_q=1, + packed_query=False, + q_data_type=torch.bfloat16, + kv_layout="NHD", + has_sinks=False, + return_lse=False, + assume_valid_prefix=False, + output_quant_block_size=None, + out_scale_format="ue8m0", + num_output_groups=None, + ): + """Plan D512 attention over one primary and an optional extra pool. + + ``max_topk`` describes arbitrary selected primary entries, not an SWA + window. Prepared indices already include the physical page stride, so + the core uses page size one regardless of the pools' external pages. + Planning compiles no metadata-preparation kernel. + + Output quantization is disabled by default. When enabled, G must divide + the local query-head count H. Blocks of 32 with UE8M0 scales implement + MXFP8; FP32 scales use unrounded amax/448. Both use an amax floor of + 1e-4. A block never crosses a head boundary. + """ + (device, profile, spec, sections, byte_end, capacity, output_quant) = ( + _resolve_sparse_mla_plan( + device, + batch_size, + num_heads, + max_topk=max_topk, + max_extra_topk=max_extra_topk, + max_seq_len_q=max_seq_len_q, + packed_query=packed_query, + q_data_type=q_data_type, + kv_layout=kv_layout, + has_sinks=has_sinks, + return_lse=return_lse, + assume_valid_prefix=assume_valid_prefix, + output_quant_block_size=output_quant_block_size, + out_scale_format=out_scale_format, + num_output_groups=num_output_groups, + ) + ) + kernel = spec.kernel + family = profile.family + splits = profile.split_kv + max_rows = batch_size * max_seq_len_q + # Policy resolves supported geometry before compilation. The fused + # kernel handles one split or a cluster reduction; other splits use + # the separate reducer. Independent source scales use the unfused path. + use_cluster_epilogue = ( + family == "swap" and splits > 1 and profile.reduction == "cluster" + ) + use_fused_epilogue = splits == 1 or use_cluster_epilogue + use_fused_reduction = splits > 1 and profile.reduction == "gmem_separate" + use_direct = profile.direct_inputs + workspace = self._workspace_buffer + if workspace is None: + workspace = torch.empty(byte_end, device=device, dtype=torch.uint8) + _validate_workspace_buffer(workspace, device=device, required_bytes=byte_end) + buffers = { + name: _workspace_section_view(workspace, section) + for name, section in sections.items() + } + compiled = _get_compiled_mla_decode(spec) + compiled_fused = None + compiled_static = None + if use_fused_epilogue or use_fused_reduction: + fused_kernel = copy.copy(kernel) + fused_kernel.output_quant = output_quant + if family == "2cta": + fused_kernel.fuse_sparse_epilogue = use_fused_epilogue + fused_kernel.fuse_sparse_reduction = use_fused_reduction + else: + fused_kernel.finalize_output = True + fused_kernel.direct_sparse = use_direct + if family == "2cta": + fused_kernel.assume_valid_prefix = assume_valid_prefix + fused_kernel.direct_sparse_capacities = (max_topk, max_extra_topk) + fused_spec = replace( + spec, + kernel=fused_kernel, + kernel_signature=_mla_kernel_compile_signature(fused_kernel), + ) + compiled_fused = _get_compiled_mla_decode(fused_spec) + if use_direct: + static_kernel = copy.copy(fused_kernel) + static_kernel.direct_static_scales = True + compiled_static = _get_compiled_mla_decode( + replace( + fused_spec, + kernel=static_kernel, + kernel_signature=_mla_kernel_compile_signature(static_kernel), + ) + ) + finish = tuple( + _compile_finish(device.index, num_heads, independent, output_quant) + for independent in (False, True) + ) + self._state = dict( + device=device, + batch=batch_size, + heads=num_heads, + max_q=max_seq_len_q, + max_rows=max_rows, + packed=packed_query, + dtype=q_data_type, + ks=max_topk, + kc=max_extra_topk, + capacity=capacity, + kv_layout=kv_layout, + has_sinks=has_sinks, + return_lse=return_lse, + output_quant=output_quant, + buffers=buffers, + workspace=workspace, + compiled=compiled, + compiled_fused=compiled_fused, + compiled_static=compiled_static, + direct_inputs=use_direct, + assume_valid_prefix=assume_valid_prefix, + finish=finish, + splits=splits, + default_cl=torch.full( + (max_rows,), max_extra_topk, device=device, dtype=torch.int32 + ), + dummy_indices=torch.zeros((max_rows, 1), device=device, dtype=torch.int32), + dummy_scales=torch.ones((1, 2), device=device, dtype=torch.float32), + dummy_cache=torch.zeros((1, 1, 512), device=device, dtype=q_data_type), + default_sinks=torch.full( + (num_heads,), -torch.inf, device=device, dtype=torch.float32 + ), + ) + + self.workspace_size_bytes = byte_end + self._scalar_cache.clear() + + def _scalar(self, value, name, validate): + state = self._state + if isinstance(value, torch.Tensor): + if ( + value.dtype != torch.float32 + or value.device != state["device"] + or value.numel() != 1 + or not value.is_contiguous() + ): + raise ValueError(f"{name} must be a scalar CUDA FP32 tensor") + if validate and ( + not torch.isfinite(value).all().item() or (value <= 0).any().item() + ): + raise ValueError(f"{name} must be positive and finite") + return value.reshape(1) + if ( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(value) + or value <= 0 + ): + raise ValueError(f"{name} must be positive and finite") + try: + key = struct.unpack("f", struct.pack("f", float(value)))[0] + except OverflowError as error: + raise ValueError(f"{name} must be representable in float32") from error + if not math.isfinite(key) or key <= 0: + raise ValueError(f"{name} must be positive finite float32") + if key not in self._scalar_cache: + self._scalar_cache[key] = torch.full( + (1,), key, device=state["device"], dtype=torch.float32 + ) + return self._scalar_cache[key] + + def _pool(self, cache, name): + state = self._state + if not isinstance(cache, torch.Tensor): + raise TypeError(f"{name} must be a tensor") + if cache.ndim == 4: + axis = 1 if state["kv_layout"] == "HND" else 2 + if cache.shape[axis] != 1: + raise ValueError(f"{name} requires one KV head") + cache = cache.squeeze(axis) + if ( + cache.ndim != 3 + or cache.shape[1] <= 0 + or cache.shape[2] != 512 + or cache.shape[0] <= 0 + ): + raise ValueError( + f"{name} must have shape [pages,page_size,512] with nonempty pages" + ) + page_size = cache.shape[1] + if cache.device != state["device"] or cache.dtype != state["dtype"]: + raise ValueError(f"{name} must match the planned device and dtype") + if ( + cache.stride(-1) != 1 + or cache.stride(1) != 512 + or cache.stride(0) < page_size * 512 + or cache.stride(0) % 512 + or cache.data_ptr() % 16 + ): + raise ValueError( + f"{name} requires aligned compact rows and a page stride divisible by 512 elements" + ) + row_stride = cache.stride(0) // 512 + rows = (cache.shape[0] - 1) * row_stride + page_size + if rows >= 0x7FFFFFFF: + raise ValueError(f"{name} exceeds signed gather coordinate range") + return cache.as_strided((rows, 1, 512), (512, 512, 1)) + + def _source_metadata( + self, indices, lengths, capacity, tokens, name, *, rows, validate + ): + state = self._state + if capacity == 0: + if indices is not None or lengths is not None: + raise ValueError(f"{name} metadata is not expected by the plan") + return state["dummy_indices"][:rows], state["default_cl"][:rows] + if not isinstance(indices, torch.Tensor) or ( + indices.shape != (rows, capacity) + or indices.dtype != torch.int32 + or indices.device != state["device"] + or indices.stride(-1) != 1 + ): + raise ValueError(f"invalid {name} indices shape/dtype/device/stride") + if not isinstance(lengths, torch.Tensor) or ( + lengths.shape != (rows,) + or lengths.dtype != torch.int32 + or lengths.device != state["device"] + or not lengths.is_contiguous() + ): + raise ValueError(f"invalid {name} lengths") + if validate: + if ((lengths < 0) | (lengths > capacity)).any().item(): + raise ValueError(f"{name} length exceeds capacity") + active = ( + torch.arange(capacity, device=state["device"])[None, :] + < lengths[:, None] + ) + if (active & ((indices < -1) | (indices >= tokens))).any().item(): + raise ValueError(f"{name} active index outside pool") + if state["assume_valid_prefix"] and (active & (indices < 0)).any().item(): + raise ValueError( + f"{name} active prefix contains a hole with assume_valid_prefix=True" + ) + return indices, lengths + + def _bind_metadata(self, metadata, *, rows, passes, direct, primary_lengths): + state = self._state + buffers = state["buffers"] + prepared_buffers = {} + for key, value, dtype, shape in ( + ( + "routes", + metadata.routes, + torch.int32, + (passes, rows, state["capacity"]), + ), + ( + "lengths", + metadata.execution_lengths, + torch.int32, + (passes, rows), + ), + ( + "counts", + metadata.valid_counts, + torch.int32, + (passes, rows), + ), + ( + "scales", + metadata.scale_params, + torch.float32, + (passes, 2 + state["heads"] + rows), + ), + ): + if value is None: + if not direct: + raise ValueError(f"prepared {key} are required by this schedule") + continue + if ( + value.dtype != dtype + or value.device != state["device"] + or value.shape != shape + or not value.is_contiguous() + ): + raise ValueError(f"prepared {key} must be contiguous {dtype}{shape}") + prepared_buffers[key] = value + buffers = ( + buffers + | dict( + routes=state["dummy_indices"][:rows].unsqueeze(0), + lengths=primary_lengths.unsqueeze(0), + counts=primary_lengths.unsqueeze(0), + scales=state["dummy_scales"], + ) + | prepared_buffers + ) + return buffers + + @flashinfer_experimental_api + def run( + self, + query, + kv_cache, + metadata: SparseMLAPreparedMetadata, + extra_kv_cache=None, + *, + qo_indptr=None, + softmax_scale=512**-0.5, + q_scale=1.0, + kv_scale=1.0, + extra_kv_scale=None, + output_scale=1.0, + sinks=None, + out=None, + lse=None, + token_positions=None, + cos_sin_cache=None, + out_scale=None, + validate=True, + ): + """Launch with caller-prepared metadata, without index conversion. + + Caller-owned metadata and plan workspace must outlive graph replay. Q and + native KV are used directly; flattening padded pages only creates a + tensor view. Required attention reductions/finishing remain included. + Metadata is an unchecked consistency contract: its packed routes, + counts and scales must agree with the supplied indices and scalars. + + With output quantization enabled, ``token_positions`` is contiguous + INT32 with the query's batch/token prefix; ``cos_sin_cache`` is aligned + contiguous FP32[P,64] containing forward [cos(32), sin(32)] rows. Use + absolute query positions and the same layer-specific table as Q RoPE. + Positions must be in [0,P); validate=False leaves this to the caller. + + Return (out, out_scale), or (out, out_scale, lse) with return_lse=True. + Out is the [T,G,C] view (strides (C,T*C,1)) of a 16-byte aligned + contiguous group-major E4M3[G,T,C] tensor, where T is total query + tokens and C=(H/G)*512, in natural head/channel order. Pass + ``torch.empty((G,T,C), dtype=torch.float8_e4m3fn).transpose(0,1)`` as + ``out`` to preallocate. For K=output_quant_block_size + and N=C/K, FP32 scales have logical shape [T,G,N]; UE8M0 scales have + shape [T,G,N/4], INT32 packing four block exponent bytes, first in the + least significant byte (scale=2**(byte-127)). Allocate scale backing as + 16-byte aligned [G,columns,pad4(T)] and use backing.permute(2,0,1)[:T], + with strides (1,columns*pad4(T),pad4(T)). Padding is not written. The scalar + ``output_scale`` multiplier remains applied before quantization. + + Preallocate out/out_scale/lse and warm up before graph capture. With + fusion disabled, the existing BF16 output and return contract apply. + """ + state = self._state + if state is None: + raise RuntimeError("plan() must be called before run()") + if not isinstance(metadata, SparseMLAPreparedMetadata): + raise TypeError("metadata must be SparseMLAPreparedMetadata") + _validate_query( + query, + packed_query=state["packed"], + device=state["device"], + batch_size=state["batch"], + num_heads=state["heads"], + max_seq_len_q=state["max_q"], + q_dtype=state["dtype"], + query_dim=512, + ) + prefix = tuple(query.shape[:-2]) + rows = math.prod(prefix) + if state["packed"]: + _validate_qo_indptr( + qo_indptr, device=state["device"], batch_size=state["batch"] + ) + if validate: + maximum, total, _ = _derive_max_seq_len_q( + qo_indptr, batch_size=state["batch"] + ) + if total != rows or maximum > state["max_q"]: + raise ValueError("invalid packed query offsets") + elif qo_indptr is not None: + raise ValueError("fixed queries do not accept qo_indptr") + + primary = self._pool(kv_cache, "kv_cache") + if extra_kv_cache is None: + if ( + state["kc"] + or metadata.extra_indices is not None + or metadata.extra_lengths is not None + ): + raise ValueError("extra pool and indices are required by the plan") + extra = state["dummy_cache"] + else: + extra = self._pool(extra_kv_cache, "extra_kv_cache") + + si, sl = self._source_metadata( + metadata.indices, + metadata.lengths, + state["ks"], + primary.shape[0], + "primary", + rows=rows, + validate=validate, + ) + ci, cl = self._source_metadata( + metadata.extra_indices, + metadata.extra_lengths, + state["kc"], + extra.shape[0], + "extra", + rows=rows, + validate=validate, + ) + if state["has_sinks"] != (sinks is not None): + raise ValueError("sinks presence must match the plan") + if sinks is None: + sinks = state["default_sinks"] + if ( + sinks.shape != (state["heads"],) + or sinks.dtype != torch.float32 + or sinks.device != state["device"] + or not sinks.is_contiguous() + ): + raise ValueError("sinks must be contiguous CUDA FP32[H]") + if validate and torch.isnan(sinks).any().item(): + raise ValueError("sinks must not contain NaN") + if extra_kv_scale is None: + extra_kv_scale = kv_scale + shared_scale = extra_kv_scale is kv_scale or ( + not isinstance(kv_scale, torch.Tensor) + and not isinstance(extra_kv_scale, torch.Tensor) + and kv_scale == extra_kv_scale + ) + independent = state["kc"] > 0 and not shared_scale + scale_tensors = [ + self._scalar(v, n, validate) + for v, n in ( + (softmax_scale, "softmax_scale"), + (q_scale, "q_scale"), + (kv_scale, "kv_scale"), + (extra_kv_scale, "extra_kv_scale"), + (output_scale, "output_scale"), + ) + ] + if validate and state["dtype"] == torch.bfloat16: + if any(t.item() != 1 for t in scale_tensors[1:4]): + raise ValueError("BF16 Q/KV descales must be one") + quant = state["output_quant"] + output_quant_args = None + if quant is not None: + out, out_scale, output_quant_args = _prepare_sparse_output_quant( + quant, query, token_positions, cos_sin_cache, out, out_scale, validate + ) + elif any(v is not None for v in (token_positions, cos_sin_cache, out_scale)): + raise ValueError( + "plan output quantization before passing RoPE/scale tensors" + ) + elif out is None: + out = torch.empty( + (*prefix, state["heads"], 512), + device=state["device"], + dtype=torch.bfloat16, + ) + if quant is None: + _validate_out( + out, + device=state["device"], + batch_size=state["batch"], + num_heads=state["heads"], + max_seq_len_q=state["max_q"], + packed_query=state["packed"], + total_q=rows if state["packed"] else None, + output_dtype=torch.bfloat16, + ) + if lse is None: + # A returned result must survive the next call on this wrapper. + # Use plan-owned scratch only when the LSE is not exposed. + lse = ( + torch.empty( + (*prefix, state["heads"]), + device=state["device"], + dtype=torch.float32, + ) + if state["return_lse"] + else state["buffers"]["public_lse"][:rows].view(*prefix, state["heads"]) + ) + if ( + lse.shape != (*prefix, state["heads"]) + or lse.dtype != torch.float32 + or lse.device != state["device"] + or not lse.is_contiguous() + ): + raise ValueError( + "lse must be contiguous FP32 with one value per query/head" + ) + fused = ( + state["compiled_fused"] is not None + and not independent + and lse.data_ptr() % 16 == 0 + ) + fused_main = fused + direct = fused and state["direct_inputs"] + static_scales = direct and all( + not isinstance(v, torch.Tensor) + for v in (softmax_scale, q_scale, kv_scale, output_scale) + ) + bmm1_scale, bmm2_scale = 1.0, 1.0 + if static_scales: + + def f32(value): + return struct.unpack("f", struct.pack("f", float(value)))[0] + + bmm1_scale = f32(f32(f32(softmax_scale) * f32(q_scale)) * f32(kv_scale)) + bmm2_scale = f32(f32(output_scale) * f32(kv_scale)) + sparse_inputs = ( + ( + si, + ci, + sl, + cl, + scale_tensors[0], + scale_tensors[1], + scale_tensors[2], + scale_tensors[4], + sinks, + ) + if direct + else None + ) + if rows: + buffers = self._bind_metadata( + metadata, + rows=rows, + passes=2 if independent else 1, + direct=direct, + primary_lengths=sl, + ) + query_view = query.view(rows, 1, state["heads"], 512) + # Final FP8 stores use output_quant_args. Keep the core O argument + # in its existing BF16 partial-state layout for every schedule. + for slot in range(2 if independent else 1): + _launch_mla_decode( + _MLARuntime( + query_view, + primary, + ( + out.view(rows, 1, state["heads"], 512) + if fused_main and quant is None + else buffers["partial"][slot, :rows] + ), + primary.shape[0], + bmm1_scale, + bmm2_scale, + extra_cache=extra, + scale_params=buffers["scales"][slot], + sparse_inputs=sparse_inputs, + output_quant_args=output_quant_args if fused_main else None, + ), + block_tables=buffers["routes"][slot, :rows], + seq_lens=buffers["lengths"][slot, :rows], + qo_indptr=None, + packed_query=False, + kv_lora_rank=512, + split_kv=state["splits"], + workspace=_MLAWorkspaceViews( + buffers["core"] if buffers["core"].numel() else None, + ( + lse.view(rows, 1, state["heads"]) + if fused_main + else buffers["lse"][slot, :rows] + ), + ), + compiled=( + state["compiled_static"] + if static_scales + else state["compiled_fused"] + if fused + else state["compiled"] + ), + ) + if not fused: + other_slot = 1 if independent else 0 + state["finish"][int(independent)]( + buffers["partial"][0, :rows].view(rows, state["heads"], 512), + buffers["partial"][other_slot, :rows].view( + rows, state["heads"], 512 + ), + buffers["lse"][0, :rows].view(rows, state["heads"]), + buffers["lse"][other_slot, :rows].view(rows, state["heads"]), + buffers["counts"][0, :rows], + buffers["counts"][1 if independent else 0, :rows], + sinks, + ( + out.view(rows, state["heads"], 512) + if quant is None + else buffers["partial"][0, :rows].view( + rows, state["heads"], 512 + ) + ), + lse.view(rows, state["heads"]), + **( + {"output_quant_args": output_quant_args} + if quant is not None + else {} + ), + ) + if quant is not None: + return (out, out_scale, lse) if state["return_lse"] else (out, out_scale) + return (out, lse) if state["return_lse"] else out + + +def get_prims_ts_sparse_mla_decode_workspace_size(*plan_args, **plan_kwargs): + """Return workspace bytes using the same arguments as wrapper.plan(). + + Resolves the default kernel geometry without allocating GPU scratch or + compiling a kernel. Bind the resulting byte buffer to the wrapper constructor + and call plan() before graph capture; run() is the prepared standalone launch. + """ + _, _, _, _, size_bytes, _, _ = _resolve_sparse_mla_plan(*plan_args, **plan_kwargs) + return size_bytes + + +@flashinfer_experimental_api +def batch_sparse_mla_decode_with_paged_kv_cache( + query, + kv_cache, + metadata: SparseMLAPreparedMetadata, + extra_kv_cache=None, + *, + qo_indptr=None, + max_seq_len_q=None, + kv_layout="NHD", + softmax_scale=512**-0.5, + q_scale=1.0, + kv_scale=1.0, + extra_kv_scale=None, + output_scale=1.0, + sinks=None, + out=None, + lse=None, + return_lse=False, + workspace_buffer=None, + assume_valid_prefix=False, + output_quant_block_size=None, + out_scale_format="ue8m0", + num_output_groups=None, + token_positions=None, + cos_sin_cache=None, + out_scale=None, +): + """Eager plan-and-run helper using caller-prepared metadata. + + Use a planned wrapper for CUDA Graph replay. Preparation remains external. + Output quantization arguments and return layouts follow + :class:`BatchSparseMLADecodePagedTSWrapper`. + """ + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "plan BatchSparseMLADecodePagedTSWrapper before CUDA Graph capture" + ) + if not isinstance(metadata, SparseMLAPreparedMetadata): + raise TypeError("metadata must be SparseMLAPreparedMetadata") + packed = query.ndim == 3 + if packed: + if qo_indptr is None: + raise ValueError("packed queries require qo_indptr") + batch = qo_indptr.numel() - 1 + if max_seq_len_q is None: + max_seq_len_q = max(1, int((qo_indptr[1:] - qo_indptr[:-1]).max().item())) + else: + batch = query.shape[0] + max_seq_len_q = query.shape[1] if max_seq_len_q is None else max_seq_len_q + wrapper = BatchSparseMLADecodePagedTSWrapper(workspace_buffer) + wrapper.plan( + query.device, + batch, + query.shape[-2], + max_topk=metadata.indices.shape[-1], + max_extra_topk=0 + if metadata.extra_indices is None + else metadata.extra_indices.shape[-1], + max_seq_len_q=max_seq_len_q, + packed_query=packed, + q_data_type=query.dtype, + kv_layout=kv_layout, + has_sinks=sinks is not None, + return_lse=return_lse, + assume_valid_prefix=assume_valid_prefix, + output_quant_block_size=output_quant_block_size, + out_scale_format=out_scale_format, + num_output_groups=num_output_groups, + ) + return wrapper.run( + query, + kv_cache, + metadata, + extra_kv_cache, + qo_indptr=qo_indptr, + softmax_scale=softmax_scale, + q_scale=q_scale, + kv_scale=kv_scale, + extra_kv_scale=extra_kv_scale, + output_scale=output_scale, + sinks=sinks, + out=out, + lse=lse, + token_positions=token_positions, + cos_sin_cache=cos_sin_cache, + out_scale=out_scale, + ) From 8358a0ebf7bc4f25ad1672c9d3c4b1e7f240f674 Mon Sep 17 00:00:00 2001 From: yunruis <205571022+yunruis@users.noreply.github.com> Date: Mon, 12 Oct 2026 22:45:04 -0700 Subject: [PATCH 02/19] [None][feat] Add the PrimTS native sparse MLA FMHA for DeepSeek-V4 generation Add prims_ts_sparse_mla, an opt-in FMHA library (TLLM_FMHA_LIBS=+prims_ts_sparse_mla) that runs the DeepSeek-V4 generation phase on the vendored PrimTS two-source sparse MLA kernel on SM100/SM103. It consumes the dual-pool storage-row indices produced by DeepseekV4TrtllmAttention.sparse_attn_predict (SWA pool primary, compress pool extra) without page-table conversion, supports BF16 and per-tensor FP8 KV caches (quantized Q from MLA preprocessing), attention sinks, and maps the fused DSv4 epilogue onto the kernel's block-128 FP32-scaled E4M3 output in TRT-LLM's group-major [G, T, (H/G)*512] / [G, (H/G)*4, pad4(T)] layout. Context stays on the fallback library. DeepSeek-V4 metadata now keeps absolute token positions for every KV cache mode (previously fp8_ds_mla only); the library uses them for window lengths and the fused output RoPE. Tests (B200, CuTe DSL 4.8.0.dev0): a kernel-level parity test of the fused output against trtllm.fused_inv_rope_fp8_quant_vllm_port, and the unified DeepSeek-V4 module test (decode and mixed batches, bf16 and fp8 caches) with the new library serving generation. 21 passed. AI-assisted. Co-Authored-By: Claude Fable 5.1 Signed-off-by: yunruis <205571022+yunruis@users.noreply.github.com> (cherry picked from commit 8bcce6c27af60938af8c03c6e3f9e6d1a6186610) --- .../attention/backends/fmha/__init__.py | 2 + .../backends/fmha/prims_ts_sparse_mla.py | 658 ++++++++++++++++++ .../attention/backends/fmha/registry.py | 9 +- .../backends/sparse/deepseek_v4/metadata.py | 14 +- .../test_prims_ts_sparse_mla_fmha.py | 270 +++++++ .../_torch/attention/test_fmha_registry.py | 4 + 6 files changed, 955 insertions(+), 2 deletions(-) create mode 100644 tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py create mode 100644 tests/unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py diff --git a/tensorrt_llm/_torch/attention/backends/fmha/__init__.py b/tensorrt_llm/_torch/attention/backends/fmha/__init__.py index 0fc278ce8c1c..2e1ee8497bc9 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/__init__.py @@ -24,6 +24,7 @@ from .msa_prefill import MsaPrefillFmha from .phased import FmhaParams, PhasedFmha from .prims_ts import PrimsTSFmha +from .prims_ts_sparse_mla import PrimsTSSparseMlaFmha from .registry import DEFAULT_FMHA_LIBS, FMHA_LIBS, FmhaCls, get_enabled_fmha_lib_classes from .triton_custom_mask import TritonCustomMaskFmha @@ -44,6 +45,7 @@ "MsaPrefillFmha", "PhasedFmha", "PrimsTSFmha", + "PrimsTSSparseMlaFmha", "TritonCustomMaskFmha", "get_enabled_fmha_lib_classes", ] diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py new file mode 100644 index 000000000000..ccc464e1783e --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py @@ -0,0 +1,658 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""TRT-LLM FMHA adapter for the vendored PrimTS native sparse MLA kernels. + +Serves the DeepSeek-V4 generation phase on SM100/SM103 with the Blackwell +task-scheduled two-source sparse MLA decode kernel +(``BatchSparseMLADecodePagedTSWrapper``). The sliding-window pool is the +primary source and, for compressing layers, the compressed pool is the extra +source. Both lists feed one softmax, exactly like the TRTLLM-Gen sparse MLA +kernel the fallback library runs. + +The adapter consumes the dual-pool global indices that +``DeepseekV4TrtllmAttention.sparse_attn_predict`` already produces: SWA slots +are storage rows of the SWA pool and compressed slots are storage rows of the +layer's compress pool, so no page-table conversion happens here. The fused +DSv4 epilogue (conjugate RoPE on ``D[448:512]`` + 1x128 block-scaled E4M3 +output in the group-major ``[G, T, (H/G)*512]`` layout with FP32 scales +``[G, (H/G)*4, pad4(T)]``) maps one-to-one onto the kernel's +``output_quant_block_size=128`` / ``out_scale_format="fp32"`` mode, so the +``o_a_proj`` BMM consumes the attention output without a separate quant +kernel. + +The library is opt-in (``TLLM_FMHA_LIBS=+prims_ts_sparse_mla``) and leaves +the context phase to the fallback library. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from importlib import import_module +from importlib.metadata import PackageNotFoundError, version +from typing import TYPE_CHECKING, Optional + +import torch +import triton +import triton.language as tl +from packaging.version import InvalidVersion, Version + +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + AttentionInputType, +) +from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor, get_sm_version +from tensorrt_llm.logger import logger + +from .interface import Fmha + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention.backends.prims_ts.mla_decode import ( + BatchSparseMLADecodePagedTSWrapper, + SparseMLAPreparedMetadata, + ) + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.cache_manager import ( + DeepseekV4CacheManager, + ) + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.metadata import ( + DeepseekV4TrtllmAttentionMetadata, + ) + from tensorrt_llm._torch.attention.backends.trtllm import ( + TrtllmAttention, + TrtllmAttentionMetadata, + ) + + +_MIN_CUTLASS_DSL_VERSION = Version("4.7.0") +_LATENT_DIM = 512 +_ROPE_DIM = 64 +_FUSED_OUTPUT_BLOCK = 128 +_ROUTE_INVALID = 0x7FFFFFFF +_ROUTE_EXTRA_BIT = -2147483648 + + +def get_prims_ts_toolchain_unsupported_reason() -> Optional[str]: + """Return why the CuTe DSL task-scheduling toolchain cannot run PrimTS kernels.""" + try: + installed_version = Version(version("nvidia-cutlass-dsl")) + except (PackageNotFoundError, InvalidVersion): + return "nvidia-cutlass-dsl>=4.7 is required." + if installed_version < _MIN_CUTLASS_DSL_VERSION: + return ( + f"nvidia-cutlass-dsl>={_MIN_CUTLASS_DSL_VERSION} is required, got {installed_version}." + ) + try: + import_module("cutlass.experimental.task_scheduling") + except ImportError: + return "CUTLASS task scheduling is missing." + return None + + +@triton.jit +def _pack_sparse_mla_metadata( + swa_indices, + extra_indices, + swa_lengths, + extra_lengths, + routes, + execution_lengths, + valid_counts, + scale_params, + sinks, + bmm1_scale, + bmm2_scale, + KS: tl.constexpr, + KC: tl.constexpr, + ROWS: tl.constexpr, + HEADS: tl.constexpr, + CAP: tl.constexpr, + BLOCK: tl.constexpr, + HEAD_BLOCK: tl.constexpr, +): + """Build the combined route words and scale rows for one shared-scale pass. + + Mirrors the ``SparseMLAPreparedMetadata`` contract: each source span is + rounded to 128 slots, extra-pool rows carry bit 31, ``0x7fffffff`` marks an + inactive slot, and the scale row holds ``[bmm1, bmm2, sinks[H], counts[rows]]``. + """ + row = tl.program_id(0) + ls = tl.load(swa_lengths + row) + lc = tl.load(extra_lengths + row) if KC else 0 + span_s = tl.cdiv(ls, 128) * 128 + span_c = tl.cdiv(lc, 128) * 128 + col = tl.arange(0, BLOCK) + for start in range(0, CAP, BLOCK): + p = start + col + a = tl.load(swa_indices + row * KS + p, (p < ls) & (p < CAP), 0x7FFFFFFF) + local = p - span_s + if KC: + b = tl.load( + extra_indices + row * KC + local, + (local >= 0) & (local < lc) & (p < CAP), + 0x7FFFFFFF, + ) + else: + b = tl.full((BLOCK,), 0x7FFFFFFF, tl.int32) + b = b | -2147483648 + word = tl.where(p < span_s, a, b) + tl.store(routes + row * CAP + p, word, p < CAP) + tl.store(execution_lengths + row, tl.maximum(span_s + span_c, 1)) + tl.store(valid_counts + row, ls + lc) + tl.store(scale_params + 2 + HEADS + row, (ls + lc).to(tl.float32)) + if row == 0: + tl.store(scale_params, bmm1_scale) + tl.store(scale_params + 1, bmm2_scale) + h = tl.arange(0, HEAD_BLOCK) + sink = tl.load(sinks + h, h < HEADS, 0) + tl.store(scale_params + 2 + h, sink, h < HEADS) + + +@dataclass +class _SparsePlan: + """One planned wrapper plus the caller-owned metadata buffers it binds.""" + + wrapper: "BatchSparseMLADecodePagedTSWrapper" + metadata: "SparseMLAPreparedMetadata" + default_sinks: torch.Tensor + rows: int + capacity: int + validated: bool = False + + +class PrimsTSSparseMlaFmha(Fmha): + """DeepSeek-V4 generation-phase sparse MLA on the PrimTS Blackwell kernels.""" + + # The native kernel reads one KV head of 512 channels per selected token. + MAX_HEADS = 128 + + def __init__(self, attn: "TrtllmAttention") -> None: + super().__init__(attn) + # Plans are keyed by the execution profile that changes kernel + # selection or buffer geometry: query rows, dtype, sinks, fusion. + self._plans: dict[tuple, _SparsePlan] = {} + # FP8 dequantization scales are fixed per layer; the wrapper takes host + # scalars, so read them once during eager warmup, never under capture. + self._kv_dequant_scale: Optional[float] = None + # Pool-base-rooted views of the SWA and compress pools, per manager. + self._pool_views: dict[tuple[int, int], torch.Tensor] = {} + self._sinks_cache: Optional[tuple[int, torch.Tensor]] = None + + # ------------------------------------------------------------------ availability + @classmethod + def _is_available(cls, attn: "TrtllmAttention") -> bool: + sm = get_sm_version() + if sm not in (100, 103): + logger.debug( + f"PrimTS sparse MLA FMHA is unavailable: requires SM100 or SM103, got SM{sm}." + ) + return False + if not attn.is_mla_enable: + return False + if getattr(getattr(attn, "sparse_params", None), "algorithm", None) != "deepseek_v4": + return False + if getattr(attn, "use_fp8_ds_mla", False) or getattr(attn, "kv_cache_dtype", "auto") == ( + "fp8_ds_mla" + ): + logger.debug( + "PrimTS sparse MLA FMHA is unavailable: the footer-scale fp8_ds_mla " + "cache layout is not a native E4M3 pool." + ) + return False + if getattr(attn, "_uses_nvfp4_compress", False): + logger.debug( + "PrimTS sparse MLA FMHA is unavailable: NVFP4 compressed pools are not supported." + ) + return False + if ( + int(attn.kv_lora_rank) + int(attn.qk_rope_head_dim) != _LATENT_DIM + or int(attn.head_dim) != _LATENT_DIM + ): + logger.debug( + "PrimTS sparse MLA FMHA is unavailable: requires a 512-wide latent " + f"(kv_lora_rank + qk_rope_head_dim), got {attn.kv_lora_rank} + " + f"{attn.qk_rope_head_dim} with head_dim {attn.head_dim}." + ) + return False + if int(attn.qk_rope_head_dim) != _ROPE_DIM: + return False + if not 1 <= int(attn.num_heads) <= cls.MAX_HEADS: + return False + reason = get_prims_ts_toolchain_unsupported_reason() + if reason is not None: + logger.debug(f"PrimTS sparse MLA FMHA is unavailable: {reason}") + return False + return True + + # ------------------------------------------------------------------ support + def _is_supported( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + *, + phase=None, + ) -> bool: + reason = self._unsupported_reason(q, k, v, metadata, forward_args) + if reason is not None: + logger.debug(f"PrimTS sparse MLA FMHA does not support request: {reason}") + return False + return True + + def _unsupported_reason( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + fwd: AttentionForwardArgs, + ) -> Optional[str]: + attn = self.attn + if fwd.attention_input_type != AttentionInputType.generation_only: + return "only generation-only requests are served; context stays on the fallback." + if k is not None or v is not None: + return "MLA generation takes a fused query only." + if q.device.type != "cuda" or q.ndim != 2: + return "the query must be a rank-2 CUDA tensor." + if metadata.kv_cache_manager is None or not hasattr( + metadata.kv_cache_manager, "swa_pool_ptr" + ): + return "a DeepSeek-V4 KV cache manager is required." + if metadata.beam_width != 1: + return "beam search is not supported." + if getattr(metadata, "token_positions_cuda", None) is None: + return "token positions are required for lengths and output RoPE." + if fwd.attention_mask_data is not None or fwd.relative_attention_bias is not None: + return "custom masks and relative bias are not supported." + if metadata.helix_position_offsets is not None: + return "Helix parallelism is not supported." + rows = int(q.shape[0]) + if rows == 0: + return "the request contains no generation tokens." + if int(metadata.num_tokens) - int(metadata.num_ctx_tokens) != rows: + return "the query rows do not match the generation token count." + heads = int(attn.num_heads) + if q.shape[1] != heads * _LATENT_DIM: + return f"query width must be {heads * _LATENT_DIM}, got {q.shape[1]}." + + indices = fwd.sparse_runtime_params.sparse_attn_indices + if not isinstance(indices, torch.Tensor): + return "dual-pool sparse indices are required." + if indices.ndim != 2 or indices.dtype != torch.int32 or indices.shape[0] != rows: + return "sparse indices must be int32 [rows, window + compressed]." + window = int(attn.sparse_attention_config.window_size) + extra = self._extra_capacity(metadata) + if indices.shape[1] != window + extra: + return ( + f"sparse index width {indices.shape[1]} does not match the dual-pool " + f"layout {window} + {extra}." + ) + + is_fp8 = self._is_fp8_cache(metadata) + if is_fp8: + quant_q = fwd.quant_q_buffer + if quant_q is None or quant_q.dtype not in (torch.uint8, torch.float8_e4m3fn): + return "FP8 sparse MLA requires the quantized query from MLA preprocessing." + if not quant_q.is_contiguous() or quant_q.numel() < rows * heads * _LATENT_DIM: + return "the quantized query buffer is too small or not contiguous." + if attn.kv_scale_quant_orig is None: + return "FP8 sparse MLA requires the KV-cache dequantization scale." + elif q.dtype != torch.bfloat16: + return f"BF16 KV caches require a BF16 query, got {q.dtype}." + + output = fwd.output + if output is None: + return "a preallocated output is required." + if fwd.enable_dsv4_epilogue_fusion: + groups = output.shape[0] if output.ndim == 3 else 0 + if ( + output.ndim != 3 + or output.dtype != torch.float8_e4m3fn + or groups <= 0 + or heads % groups + or output.shape[1] != rows + or output.shape[2] != heads // groups * _LATENT_DIM + or not output.is_contiguous() + ): + return "the fused epilogue output must be a contiguous E4M3 [G, T, (H/G)*512]." + sf = fwd.output_sf + if ( + sf is None + or sf.dtype != torch.float32 + or sf.ndim != 3 + or sf.shape[0] != groups + or sf.shape[1] != heads // groups * (_LATENT_DIM // _FUSED_OUTPUT_BLOCK) + or sf.shape[2] != (rows + 3) // 4 * 4 + or not sf.is_contiguous() + ): + return "the fused epilogue scales must be FP32 [G, (H/G)*4, pad4(T)]." + cos_sin = fwd.dsv4_inv_rope_cos_sin_cache + if cos_sin is None or cos_sin.dtype != torch.float32 or cos_sin.numel() % _ROPE_DIM: + return "the fused epilogue needs the FP32 inverse-RoPE cos/sin table." + else: + if output.dtype != torch.bfloat16 or output.numel() != rows * heads * _LATENT_DIM: + return "the output must be BF16 with one 512-wide row per query head." + if fwd.out_scale is not None or fwd.output_sf is not None: + return "unfused quantized attention output is not supported." + return None + + # ------------------------------------------------------------------ helpers + @staticmethod + def _is_fp8_cache(metadata: "TrtllmAttentionMetadata") -> bool: + from tensorrt_llm.bindings import DataType + + return metadata.kv_cache_manager.dtype == DataType.FP8 + + def _extra_capacity(self, metadata: "DeepseekV4TrtllmAttentionMetadata") -> int: + ratio = int(self.attn.compress_ratio) + if ratio <= 1: + return 0 + return int(metadata.max_compressed_indices[ratio]) + + def _pool_view( + self, + manager: "DeepseekV4CacheManager", + attn_type, + compress_ratio: int, + ) -> torch.Tensor: + """Return a ``[rows, 1, 512]`` view of a whole pool rooted at its base pointer. + + ``sparse_attn_predict`` emits storage rows relative to the pool base, so + the view must start there and span every layer sharing the pool. + """ + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType, + ) + + key = (id(manager), int(attn_type.value), compress_ratio) + view = self._pool_views.get(key) + if view is not None: + return view + if attn_type == DeepseekV4AttentionType.SWA: + base_ptr = manager.swa_pool_ptr + layers = list(manager.pp_layers) + else: + base_ptr = manager.compress_pool_ptrs[compress_ratio] + layers = [ + layer + for layer in manager.pp_layers + if manager._compress_ratios[layer] == compress_ratio + ] + total_rows = 0 + for layer in layers: + buffer = manager.get_buffers(layer, attn_type) + if buffer.shape[-1] != _LATENT_DIM: + raise RuntimeError( + f"PrimTS sparse MLA expects native 512-wide pool rows, got {tuple(buffer.shape)}." + ) + offset = buffer.data_ptr() - base_ptr + row_bytes = _LATENT_DIM * buffer.element_size() + if offset < 0 or offset % row_bytes: + raise RuntimeError( + f"{attn_type.name} buffer for layer {layer} is not row-aligned to its pool base." + ) + total_rows = max(total_rows, offset // row_bytes + buffer.shape[0] * buffer.shape[1]) + view = convert_to_torch_tensor( + TensorWrapper(base_ptr, manager.dtype, (total_rows, 1, _LATENT_DIM)) + ) + self._pool_views[key] = view + return view + + def _sinks(self, fwd: AttentionForwardArgs) -> Optional[torch.Tensor]: + sinks = fwd.attention_sinks + if sinks is None: + return None + cached = self._sinks_cache + if ( + cached is not None + and cached[0] == sinks.data_ptr() + and cached[1].numel() == sinks.numel() + ): + return cached[1] + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "PrimTS sparse MLA attention sinks must be prepared before CUDA graph capture." + ) + prepared = sinks.detach().to(torch.float32).reshape(-1).contiguous() + self._sinks_cache = (sinks.data_ptr(), prepared) + return prepared + + def _kv_dequant(self) -> float: + if self._kv_dequant_scale is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "PrimTS sparse MLA FP8 scales must be cached before CUDA graph capture." + ) + scale = float(self.attn.kv_scale_quant_orig.view(-1)[0].item()) + if not (math.isfinite(scale) and scale > 0): + raise RuntimeError("PrimTS sparse MLA FP8 scales must be finite and positive.") + self._kv_dequant_scale = scale + return self._kv_dequant_scale + + def _get_plan( + self, + *, + device: torch.device, + rows: int, + extra: int, + q_dtype: torch.dtype, + has_sinks: bool, + fused_groups: Optional[int], + ) -> _SparsePlan: + from tensorrt_llm._torch.attention.backends.prims_ts.mla_decode import ( + BatchSparseMLADecodePagedTSWrapper, + SparseMLAPreparedMetadata, + ) + + attn = self.attn + heads = int(attn.num_heads) + window = int(attn.sparse_attention_config.window_size) + key = (rows, extra, q_dtype, has_sinks, fused_groups) + plan = self._plans.get(key) + if plan is not None: + return plan + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "PrimTS sparse MLA plans must be created during eager warmup, not under " + f"CUDA graph capture (rows={rows}, extra={extra}, dtype={q_dtype})." + ) + wrapper = BatchSparseMLADecodePagedTSWrapper() + wrapper.plan( + device, + rows, + heads, + max_topk=window, + max_extra_topk=extra, + max_seq_len_q=1, + q_data_type=q_dtype, + kv_layout="NHD", + has_sinks=has_sinks, + return_lse=False, + # Window slots and compressed slots are valid prefixes in the + # dual-pool layout, but the ratio-4 indexer top-k may carry holes; + # keep the generic mask path. + assume_valid_prefix=False, + **( + dict( + output_quant_block_size=_FUSED_OUTPUT_BLOCK, + out_scale_format="fp32", + num_output_groups=fused_groups, + ) + if fused_groups + else {} + ), + ) + capacity = int(wrapper._state["capacity"]) + + def empty(shape, dtype=torch.int32): + return torch.empty(shape, device=device, dtype=dtype) + + metadata = SparseMLAPreparedMetadata( + indices=empty((rows, window)), + lengths=empty((rows,)), + extra_indices=empty((rows, extra)) if extra else None, + extra_lengths=empty((rows,)) if extra else None, + routes=empty((1, rows, capacity)), + execution_lengths=empty((1, rows)), + valid_counts=empty((1, rows)), + scale_params=empty((1, 2 + heads + rows), torch.float32), + ) + plan = _SparsePlan( + wrapper=wrapper, + metadata=metadata, + default_sinks=torch.full((heads,), -torch.inf, device=device, dtype=torch.float32), + rows=rows, + capacity=capacity, + ) + self._plans[key] = plan + return plan + + # ------------------------------------------------------------------ forward + def forward( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + ) -> None: + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType, + ) + + reason = self._unsupported_reason(q, k, v, metadata, forward_args) + if reason is not None: + raise RuntimeError(f"PrimTS sparse MLA FMHA cannot run this request: {reason}") + attn = self.attn + manager = metadata.kv_cache_manager + heads = int(attn.num_heads) + rows = int(q.shape[0]) + start = int(metadata.num_ctx_tokens) + end = start + rows + ratio = int(attn.compress_ratio) + window = int(attn.sparse_attention_config.window_size) + extra = self._extra_capacity(metadata) + is_fp8 = self._is_fp8_cache(metadata) + fused = bool(forward_args.enable_dsv4_epilogue_fusion) + output = forward_args.output + groups = int(output.shape[0]) if fused else None + + softmax_scale = 1.0 / (float(attn.q_scaling) * math.sqrt(float(_LATENT_DIM))) + if is_fp8: + # The MLA preprocessing quantizes Q with the KV-cache scale, so both + # descales equal the layer's kv_scale_quant_orig (see + # DeepseekV4TrtllmAttention.sparse_attn_predict). + dequant = self._kv_dequant() + q_scale = kv_scale = dequant + query = ( + forward_args.quant_q_buffer.view(-1)[: rows * heads * _LATENT_DIM] + .view(torch.float8_e4m3fn) + .view(rows, 1, heads, _LATENT_DIM) + ) + else: + q_scale = kv_scale = 1.0 + query = q.view(rows, 1, heads, _LATENT_DIM) + output_scale = 1.0 + bmm1_scale = softmax_scale * q_scale * kv_scale + bmm2_scale = output_scale * kv_scale + + sinks = self._sinks(forward_args) + plan = self._get_plan( + device=q.device, + rows=rows, + extra=extra, + q_dtype=query.dtype, + has_sinks=sinks is not None, + fused_groups=groups, + ) + prepared = plan.metadata + + # Storage-row indices straight from sparse_attn_predict: the window + # slots address the SWA pool and the compressed slots address the + # compress pool, both relative to their pool base. + indices = forward_args.sparse_runtime_params.sparse_attn_indices + positions = metadata.token_positions_cuda[start:end] + prepared.indices.copy_(indices[:, :window]) + torch.add(positions, 1, out=prepared.lengths) + prepared.lengths.clamp_(max=window) + if extra: + prepared.extra_indices.copy_(indices[:, window:]) + topk_lens = metadata.sparse_mla_topk_lens[ratio][start:end] + torch.sub(topk_lens, window, out=prepared.extra_lengths) + prepared.extra_lengths.clamp_(min=0, max=extra) + _pack_sparse_mla_metadata[(rows,)]( + prepared.indices, + prepared.extra_indices if extra else prepared.indices, + prepared.lengths, + prepared.extra_lengths if extra else prepared.lengths, + prepared.routes, + prepared.execution_lengths, + prepared.valid_counts, + prepared.scale_params, + plan.default_sinks if sinks is None else sinks, + bmm1_scale, + bmm2_scale, + window, + extra, + rows, + heads, + plan.capacity, + 256, + triton.next_power_of_2(heads), + num_warps=4, + ) + + swa_pool = self._pool_view(manager, DeepseekV4AttentionType.SWA, 1) + extra_pool = ( + self._pool_view(manager, DeepseekV4AttentionType.COMPRESS, ratio) if extra else None + ) + + run_kwargs = {} + if fused: + # TRT-LLM's fused epilogue buffers are group-major [G, T, C] with + # FP32 scales [G, (H/G)*4, pad4(T)]; the kernel takes their + # token-major logical views. + run_kwargs["out"] = output.transpose(0, 1) + run_kwargs["out_scale"] = forward_args.output_sf.permute(2, 0, 1)[:rows] + # Positions carry the query prefix shape [rows, 1]. + run_kwargs["token_positions"] = positions.contiguous().view(rows, 1) + cos_sin = forward_args.dsv4_inv_rope_cos_sin_cache + run_kwargs["cos_sin_cache"] = cos_sin.reshape(-1, _ROPE_DIM).contiguous() + else: + run_kwargs["out"] = output.view(rows, 1, heads, _LATENT_DIM) + + # Validate the metadata contract once per plan while still eager; graph + # replay must keep validation off because it synchronizes. + validate = not plan.validated and not torch.cuda.is_current_stream_capturing() + plan.wrapper.run( + query, + swa_pool, + prepared, + extra_pool, + softmax_scale=softmax_scale, + q_scale=q_scale, + kv_scale=kv_scale, + extra_kv_scale=kv_scale, + output_scale=output_scale, + sinks=sinks, + validate=validate, + **run_kwargs, + ) + plan.validated = True + + +__all__ = ["PrimsTSSparseMlaFmha", "get_prims_ts_toolchain_unsupported_reason"] diff --git a/tensorrt_llm/_torch/attention/backends/fmha/registry.py b/tensorrt_llm/_torch/attention/backends/fmha/registry.py index b7b60a0de6d1..67d4b0a99ed4 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/registry.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/registry.py @@ -38,6 +38,7 @@ def init_fmha_libs() -> dict[str, "FmhaCls"]: from .msa_decode import MsaDecodeFmha from .msa_prefill import MsaPrefillFmha from .prims_ts_block_sparse import PrimsTSBlockSparseFmha + from .prims_ts_sparse_mla import PrimsTSSparseMlaFmha return { "fp4_mla": Fp4MlaFmha, @@ -51,13 +52,19 @@ def init_fmha_libs() -> dict[str, "FmhaCls"]: "flashinfer_sparse_mla": FlashInferSparseMlaFmha, "prims_ts": PrimsTSFmha, "prims_ts_block_sparse": PrimsTSBlockSparseFmha, + # DeepSeek-V4 generation on the vendored PrimTS native sparse MLA + # kernel; opt-in like prims_ts until it graduates from validation. + "prims_ts_sparse_mla": PrimsTSSparseMlaFmha, "flashinfer_trtllm_gen": FlashInferTrtllmGenFmha, "fallback": FallbackFmha, } FMHA_LIBS: dict[str, FmhaCls] = init_fmha_libs() -DEFAULT_FMHA_LIBS: tuple[str, ...] = tuple(name for name in FMHA_LIBS if name != "prims_ts") +_OPT_IN_FMHA_LIBS: frozenset[str] = frozenset({"prims_ts", "prims_ts_sparse_mla"}) +DEFAULT_FMHA_LIBS: tuple[str, ...] = tuple( + name for name in FMHA_LIBS if name not in _OPT_IN_FMHA_LIBS +) def _parse_fmha_libs_env() -> tuple[str, ...]: diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py index 67bde65e7227..7739fc12685e 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py @@ -241,6 +241,18 @@ def __post_init__(self): capture_graph=capture_graph, ) + # Absolute token positions feed the PrimTS sparse MLA FMHA (window + # lengths and the fused output RoPE). The DSA base allocates them only + # for the fp8_ds_mla layout, so provide them for every DSv4 cache mode. + if getattr(self, "token_positions_cuda", None) is None: + self.token_positions_cuda = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_tokens,), + cache_name="token_positions_cuda", + dtype=torch.int32, + capture_graph=capture_graph, + ) + # Compute max_compressed_indices for CUDA graph compatibility. # For ratio=4, the indexer selects index_topk compressed tokens. # For ratio=128, use max_seq_len / 128 rounded up to next power of 2 @@ -504,7 +516,7 @@ def prepare_for_deepseek_v4_indices(self, token_positions=None): offsets = token_idx - cu_seq_lens[req_idx].to(torch.int32) token_positions = cached_tokens[req_idx].to(torch.int32) + offsets - if self.use_fp8_ds_mla: + if self.token_positions_cuda is not None: self.token_positions_cuda[: token_positions.shape[0]] = token_positions self._prepare_deepseek_v4_indices_compiled( 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 new file mode 100644 index 000000000000..72116b519995 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py @@ -0,0 +1,270 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Tests for the PrimTS native sparse MLA FMHA library (DeepSeek-V4, SM100/SM103).""" + +import pytest +import torch +import triton + +from tensorrt_llm._torch.attention.backends.fmha import registry +from tensorrt_llm._torch.attention.backends.fmha.prims_ts_sparse_mla import ( + PrimsTSSparseMlaFmha, + _pack_sparse_mla_metadata, + get_prims_ts_toolchain_unsupported_reason, +) +from tensorrt_llm._utils import get_sm_version + +from ..test_sparse_mla_forward import HAS_FLASH_MLA +from ..test_sparse_mla_forward import ( + test_forward_sparse_mla_unified as _run_unified_sparse_mla_forward, +) + +LIB_NAME = "prims_ts_sparse_mla" +LATENT = 512 +ROPE = 64 +BLOCK = 128 + +pytestmark = pytest.mark.skipif( + get_sm_version() not in (100, 103), reason="PrimTS sparse MLA requires SM100 or SM103" +) + + +def _skip_unless_toolchain() -> None: + reason = get_prims_ts_toolchain_unsupported_reason() + if reason is not None: + pytest.skip(reason) + + +def test_registry_entry_is_opt_in() -> None: + assert registry.FMHA_LIBS[LIB_NAME] is PrimsTSSparseMlaFmha + assert LIB_NAME not in registry.DEFAULT_FMHA_LIBS + names = tuple(registry.FMHA_LIBS) + assert names.index(LIB_NAME) < names.index("fallback") + + +def _rope_table(max_positions: int, device: torch.device) -> torch.Tensor: + """Forward [cos(32), sin(32)] rows, the layout both kernels consume.""" + inv_freq = 1.0 / ( + 10000.0 ** (torch.arange(0, ROPE, 2, dtype=torch.float32, device=device) / ROPE) + ) + angles = torch.outer(torch.arange(max_positions, dtype=torch.float32, device=device), inv_freq) + return torch.cat([angles.cos(), angles.sin()], dim=-1).contiguous() + + +def _random_indices(rows: int, capacity: int, pool_rows: int, device: torch.device): + lengths = torch.randint(1, capacity + 1, (rows,), dtype=torch.int32, device=device) + indices = torch.randint(0, pool_rows, (rows, capacity), dtype=torch.int32, device=device) + inactive = torch.arange(capacity, device=device)[None, :] >= lengths[:, None] + indices[inactive] = -1 + return indices, lengths + + +def _prepare_metadata(wrapper, indices, lengths, extra_indices, extra_lengths, heads, bmm1, bmm2): + """Build SparseMLAPreparedMetadata exactly like the FMHA adapter does.""" + from tensorrt_llm._torch.attention.backends.prims_ts.mla_decode import SparseMLAPreparedMetadata + + rows, ks = indices.shape + kc = 0 if extra_indices is None else extra_indices.shape[1] + capacity = int(wrapper._state["capacity"]) + device = indices.device + + def empty(shape, dtype=torch.int32): + return torch.empty(shape, dtype=dtype, device=device) + + prepared = SparseMLAPreparedMetadata( + indices=indices.contiguous(), + lengths=lengths.contiguous(), + extra_indices=None if kc == 0 else extra_indices.contiguous(), + extra_lengths=None if kc == 0 else extra_lengths.contiguous(), + routes=empty((1, rows, capacity)), + execution_lengths=empty((1, rows)), + valid_counts=empty((1, rows)), + scale_params=empty((1, 2 + heads + rows), torch.float32), + ) + sinks = torch.full((heads,), -torch.inf, dtype=torch.float32, device=device) + _pack_sparse_mla_metadata[(rows,)]( + prepared.indices, + prepared.extra_indices if kc else prepared.indices, + prepared.lengths, + prepared.extra_lengths if kc else prepared.lengths, + prepared.routes, + prepared.execution_lengths, + prepared.valid_counts, + prepared.scale_params, + sinks, + bmm1, + bmm2, + ks, + kc, + rows, + heads, + capacity, + 256, + triton.next_power_of_2(heads), + num_warps=4, + ) + return prepared + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn], ids=["bf16", "fp8"]) +@pytest.mark.parametrize( + "rows,heads,groups,topk,extra_topk", + [ + (1, 128, 16, 128, 0), + (5, 128, 16, 128, 512), + (64, 128, 16, 128, 512), + (16, 64, 8, 128, 256), + (3, 64, 8, 128, 0), + ], +) +def test_fused_output_matches_trtllm_inverse_rope_quant( + dtype, rows, heads, groups, topk, extra_topk +): + """The kernel's fused epilogue must reproduce TRT-LLM's inverse-RoPE + FP8 quant. + + Reference: the same attention with a BF16 output, followed by + ``trtllm.fused_inv_rope_fp8_quant_vllm_port`` (the op the unfused DSv4 + path runs). The fused kernel quantizes from FP32 accumulators instead of + the BF16 output, so compare dequantized values up to FP8 rounding. + """ + _skip_unless_toolchain() + from tensorrt_llm._torch.attention.backends.prims_ts.mla_decode import ( + BatchSparseMLADecodePagedTSWrapper, + ) + + torch.manual_seed(0) + device = torch.device("cuda") + pool_rows, extra_rows, max_positions = 4096, 2048, 4096 + + def make_pool(n): + return (torch.randn(n, 1, LATENT, dtype=torch.bfloat16, device=device)).to(dtype) + + pool = make_pool(pool_rows) + extra_pool = make_pool(extra_rows) if extra_topk else None + indices, lengths = _random_indices(rows, topk, pool_rows, device) + extra_indices, extra_lengths = ( + (None, None) if extra_topk == 0 else _random_indices(rows, extra_topk, extra_rows, device) + ) + query = (torch.randn(rows, 1, heads, LATENT, dtype=torch.bfloat16, device=device) * 0.5).to( + dtype + ) + positions = torch.randint(0, max_positions, (rows,), dtype=torch.int32, device=device) + cos_sin = _rope_table(max_positions, device) + + softmax_scale = LATENT**-0.5 + descale = 0.75 if dtype == torch.float8_e4m3fn else 1.0 + bmm1 = softmax_scale * descale * descale + bmm2 = descale + run_scales = dict( + softmax_scale=softmax_scale, q_scale=descale, kv_scale=descale, extra_kv_scale=descale + ) + plan_args = dict(max_topk=topk, max_extra_topk=extra_topk, max_seq_len_q=1, q_data_type=dtype) + + reference = BatchSparseMLADecodePagedTSWrapper() + reference.plan(device, rows, heads, **plan_args) + metadata = _prepare_metadata( + reference, indices, lengths, extra_indices, extra_lengths, heads, bmm1, bmm2 + ) + out_bf16 = reference.run(query, pool, metadata, extra_pool, **run_scales) + heads_per_group = heads // groups + ref_fp8, ref_scale = torch.ops.trtllm.fused_inv_rope_fp8_quant_vllm_port( + out_bf16.view(rows, heads, LATENT), + positions, + cos_sin, + groups, + heads_per_group, + LATENT - ROPE, + ROPE, + BLOCK, + False, + ) + + fused = BatchSparseMLADecodePagedTSWrapper() + fused.plan( + device, + rows, + heads, + output_quant_block_size=BLOCK, + out_scale_format="fp32", + num_output_groups=groups, + **plan_args, + ) + metadata = _prepare_metadata( + fused, indices, lengths, extra_indices, extra_lengths, heads, bmm1, bmm2 + ) + channels = heads_per_group * LATENT + padded_rows = (rows + 3) // 4 * 4 + # The exact buffers TRT-LLM's fused DSv4 epilogue allocates. + out_buf = torch.empty((groups, rows, channels), dtype=torch.float8_e4m3fn, device=device) + scale_buf = torch.empty( + (groups, channels // BLOCK, padded_rows), dtype=torch.float32, device=device + ) + fused_out, fused_scale = fused.run( + query, + pool, + metadata, + extra_pool, + out=out_buf.transpose(0, 1), + out_scale=scale_buf.permute(2, 0, 1)[:rows], + token_positions=positions.view(rows, 1), + cos_sin_cache=cos_sin, + **run_scales, + ) + assert fused_out.data_ptr() == out_buf.data_ptr() + assert fused_scale.data_ptr() == scale_buf.data_ptr() + assert ref_fp8.shape == out_buf.shape and ref_scale.shape == scale_buf.shape + + ref_scales = ref_scale[..., :rows].permute(0, 2, 1) # [G, T, blocks] + got_scales = scale_buf[..., :rows].permute(0, 2, 1) + torch.testing.assert_close(got_scales, ref_scales, rtol=2e-2, atol=1e-6) + + def dequant(values, scales): + blocks = values.float().view(groups, rows, channels // BLOCK, BLOCK) + return (blocks * scales[..., None]).flatten(-2) + + ref = dequant(ref_fp8, ref_scales) + got = dequant(out_buf, got_scales) + # One E4M3 ulp is 1/8 of the magnitude. The reference rotates the + # BF16-rounded output, so near-cancelling RoPE channels (D[448:512]) carry + # an absolute error of a few BF16 ulps of their O(1) inputs, which the + # kernel (rotating FP32 accumulators) does not have; that shows up as up to + # about one quantization step of the block. Allow two steps there and half + # a step elsewhere, on top of two relative ulps. + block_scales = ref_scales.repeat_interleave(BLOCK, dim=-1) + is_rope = torch.zeros(channels, dtype=torch.bool, device=device) + is_rope.view(-1, LATENT)[:, LATENT - ROPE :] = True + step_allowance = torch.where(is_rope, 2.0, 0.5) * block_scales + ulp = ref.abs() * 0.125 + block_scales * 2.0**-6 + diff = (got - ref).abs() + assert (diff <= 2.0 * ulp + step_allowance + 1e-6).all(), ( + f"max fused-vs-reference gap {diff.max().item()}" + ) + assert diff.mean().item() <= 0.02 * ref.abs().mean().item() + 1e-6 + + +@pytest.mark.skipif(not HAS_FLASH_MLA, reason="FlashMLA not available") +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +@pytest.mark.parametrize( + "batch_name", + [ + "small_decode", + "medium_decode", + "single_decode", + "small_mixed", + "mixed_heavy_decode", + ], +) +def test_forward_sparse_mla_unified_prims_ts(monkeypatch, batch_name, kv_cache_dtype): + """Run the DeepSeek-V4 module test with the PrimTS library serving generation.""" + _skip_unless_toolchain() + monkeypatch.setenv("TLLM_FMHA_LIBS", f"+{LIB_NAME}") + calls = [] + original_forward = PrimsTSSparseMlaFmha.forward + + def counting_forward(self, *args, **kwargs): + calls.append(1) + return original_forward(self, *args, **kwargs) + + monkeypatch.setattr(PrimsTSSparseMlaFmha, "forward", counting_forward) + _run_unified_sparse_mla_forward(batch_name, kv_cache_dtype, "deepseek_v4") + assert calls, "the PrimTS sparse MLA FMHA did not serve the generation phase" diff --git a/tests/unittest/_torch/attention/test_fmha_registry.py b/tests/unittest/_torch/attention/test_fmha_registry.py index bd6a65809afe..6bdb30dcb0fa 100644 --- a/tests/unittest/_torch/attention/test_fmha_registry.py +++ b/tests/unittest/_torch/attention/test_fmha_registry.py @@ -20,6 +20,7 @@ PRIMS_TS = "prims_ts" PRIMS_TS_BLOCK_SPARSE = "prims_ts_block_sparse" +PRIMS_TS_SPARSE_MLA = "prims_ts_sparse_mla" def _canonical_names() -> tuple[str, ...]: @@ -35,12 +36,15 @@ def _enabled_names() -> tuple[str, ...]: def test_prims_ts_precedes_trtllm_gen_in_canonical_order() -> None: names = _canonical_names() assert names.index(PRIMS_TS) < names.index("flashinfer_trtllm_gen") + assert names.index(PRIMS_TS_SPARSE_MLA) < names.index("flashinfer_trtllm_gen") + assert names.index(PRIMS_TS_SPARSE_MLA) < names.index("fallback") def test_default_fmha_libs_exclude_prims_ts(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.delenv("TLLM_FMHA_LIBS", raising=False) assert PRIMS_TS not in registry.DEFAULT_FMHA_LIBS + assert PRIMS_TS_SPARSE_MLA not in registry.DEFAULT_FMHA_LIBS assert PRIMS_TS_BLOCK_SPARSE in registry.DEFAULT_FMHA_LIBS assert set(registry.DEFAULT_FMHA_LIBS) <= set(registry.FMHA_LIBS) assert _enabled_names() == registry.DEFAULT_FMHA_LIBS From 659ae833b9758b81d870a726cf57d01ce455f433 Mon Sep 17 00:00:00 2001 From: yunruis <205571022+yunruis@users.noreply.github.com> Date: Wed, 14 Oct 2026 21:18:32 -0700 Subject: [PATCH 03/19] [None][feat] Enable the PrimTS sparse MLA FMHA by default for DeepSeek-V4 on Blackwell `prims_ts_sparse_mla` leaves `_OPT_IN_FMHA_LIBS`: its `is_available` already restricts it to DeepSeek-V4 layers on SM100/SM103 with a native E4M3 or BF16 pool, so every other model and architecture is unaffected, and for DeepSeek-V4 generation it now precedes the TRTLLM-Gen fallback that served that phase before. `TLLM_FMHA_LIBS=-prims_ts_sparse_mla` restores the previous library. Make the selection observable: `FmhaManager` logs, once per library per layer at info, which FMHA library serves a request shape, and the PrimTS library logs at info why it steps aside for a DeepSeek-V4 layer (fp8_ds_mla cache, NVFP4 compress, toolchain) instead of only at debug. Tests: - registry tests assert the library is in the default set and ahead of the fallback, and that the `-prims_ts_sparse_mla` delta removes it; - the module test runs on the default library set (no `TLLM_FMHA_LIBS`); - new `test_forward_sparse_mla_unified_fused_epilogue` covers the fused DSv4 epilogue (3D E4M3 output + FP32 block scales, the production attention-DP path) at module level: it wraps the harness's generation call, re-runs it with `enable_dsv4_epilogue_fusion=True` on pristine inputs and checks the PrimTS fused output against the unfused output pushed through `trtllm.fused_inv_rope_fp8_quant_vllm_port`. End to end on 8x B200 (DeepSeek-V4-Flash, TP4/EP4 attention-DP, GSM8K 1319 samples, default library set vs `TLLM_FMHA_LIBS=-prims_ts_sparse_mla`): PrimTS 95.34 vs TRTLLM-Gen fallback 94.84 (reference 95.11). The info log shows `PrimsTSSparseMlaFmha serves causal requests (generation_batch<=1, ...)` for all 43 layers on every rank and `FallbackFmha` serving context. Co-Authored-By: Claude Fable 5.1 Signed-off-by: yunruis <205571022+yunruis@users.noreply.github.com> --- .../_torch/attention/backends/fmha/manager.py | 18 ++ .../backends/fmha/prims_ts_sparse_mla.py | 56 +++---- .../attention/backends/fmha/registry.py | 7 +- .../test_prims_ts_sparse_mla_fmha.py | 158 +++++++++++++++++- .../_torch/attention/test_fmha_registry.py | 2 +- 5 files changed, 205 insertions(+), 36 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/manager.py b/tensorrt_llm/_torch/attention/backends/fmha/manager.py index c9afff66e1d8..130a22d574c3 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/manager.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/manager.py @@ -311,6 +311,23 @@ def __init__(self, attn: TrtllmAttention) -> None: for fmha_cls in get_enabled_fmha_lib_classes(): if fmha_cls.is_available(attn): self.fmha_libs.append(fmha_cls(attn)) + self._layer_idx = getattr(attn, "layer_idx", None) + # Libraries already announced in the log for this layer; one line per + # library per layer is enough to tell which kernel serves a model. + self._announced: set[type[Fmha]] = set() + + def _announce(self, fmha: Fmha, cache_key: _FmhaCacheKey) -> None: + if type(fmha) in self._announced: + return + self._announced.add(type(fmha)) + logger.info( + f"FMHA layer {self._layer_idx}: {type(fmha).__name__} serves " + f"{cache_key.attention_mask_type.name} requests " + f"(context_batch<={cache_key.context_batch_size}, " + f"generation_batch<={cache_key.generation_batch_size}, " + f"generation_seq_len_q<={cache_key.generation_seq_len_q}, " + f"output_dtype={cache_key.output_dtype})." + ) def _make_cache_key( self, @@ -422,6 +439,7 @@ def select( if fmha is None: return None self._cache[cache_key] = fmha + self._announce(fmha, cache_key) if self._cache_sanity_check_enabled: self._cache_inputs[cache_key] = _snapshot_fmha_cache_inputs( q, k, v, metadata, forward_args diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py index ccc464e1783e..5a333eaffc80 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py @@ -33,8 +33,9 @@ ``o_a_proj`` BMM consumes the attention output without a separate quant kernel. -The library is opt-in (``TLLM_FMHA_LIBS=+prims_ts_sparse_mla``) and leaves -the context phase to the fallback library. +The library is enabled by default on SM100/SM103 and precedes the TRTLLM-Gen +fallback for the DeepSeek-V4 generation phase; ``TLLM_FMHA_LIBS=-prims_ts_sparse_mla`` +restores the previous library. The context phase stays on the fallback library. """ from __future__ import annotations @@ -193,48 +194,45 @@ def __init__(self, attn: "TrtllmAttention") -> None: # ------------------------------------------------------------------ availability @classmethod def _is_available(cls, attn: "TrtllmAttention") -> bool: - sm = get_sm_version() - if sm not in (100, 103): - logger.debug( - f"PrimTS sparse MLA FMHA is unavailable: requires SM100 or SM103, got SM{sm}." - ) - return False if not attn.is_mla_enable: return False if getattr(getattr(attn, "sparse_params", None), "algorithm", None) != "deepseek_v4": return False + reason = cls._unavailable_reason(attn) + if reason is not None: + # A DeepSeek-V4 layer that does not get this library falls back to + # TRTLLM-Gen silently otherwise; say so once per layer at info. + logger.info( + f"PrimTS sparse MLA FMHA is unavailable for layer " + f"{getattr(attn, 'layer_idx', None)}: {reason}" + ) + return False + return True + + @classmethod + def _unavailable_reason(cls, attn: "TrtllmAttention") -> Optional[str]: + sm = get_sm_version() + if sm not in (100, 103): + return f"requires SM100 or SM103, got SM{sm}." if getattr(attn, "use_fp8_ds_mla", False) or getattr(attn, "kv_cache_dtype", "auto") == ( "fp8_ds_mla" ): - logger.debug( - "PrimTS sparse MLA FMHA is unavailable: the footer-scale fp8_ds_mla " - "cache layout is not a native E4M3 pool." - ) - return False + return "the footer-scale fp8_ds_mla cache layout is not a native E4M3 pool." if getattr(attn, "_uses_nvfp4_compress", False): - logger.debug( - "PrimTS sparse MLA FMHA is unavailable: NVFP4 compressed pools are not supported." - ) - return False + return "NVFP4 compressed pools are not supported." if ( int(attn.kv_lora_rank) + int(attn.qk_rope_head_dim) != _LATENT_DIM or int(attn.head_dim) != _LATENT_DIM ): - logger.debug( - "PrimTS sparse MLA FMHA is unavailable: requires a 512-wide latent " - f"(kv_lora_rank + qk_rope_head_dim), got {attn.kv_lora_rank} + " - f"{attn.qk_rope_head_dim} with head_dim {attn.head_dim}." + return ( + "requires a 512-wide latent (kv_lora_rank + qk_rope_head_dim), got " + f"{attn.kv_lora_rank} + {attn.qk_rope_head_dim} with head_dim {attn.head_dim}." ) - return False if int(attn.qk_rope_head_dim) != _ROPE_DIM: - return False + return f"requires a {_ROPE_DIM}-wide RoPE segment, got {attn.qk_rope_head_dim}." if not 1 <= int(attn.num_heads) <= cls.MAX_HEADS: - return False - reason = get_prims_ts_toolchain_unsupported_reason() - if reason is not None: - logger.debug(f"PrimTS sparse MLA FMHA is unavailable: {reason}") - return False - return True + return f"supports 1..{cls.MAX_HEADS} heads, got {attn.num_heads}." + return get_prims_ts_toolchain_unsupported_reason() # ------------------------------------------------------------------ support def _is_supported( diff --git a/tensorrt_llm/_torch/attention/backends/fmha/registry.py b/tensorrt_llm/_torch/attention/backends/fmha/registry.py index 67d4b0a99ed4..e5c8e2ec68a2 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/registry.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/registry.py @@ -53,7 +53,10 @@ def init_fmha_libs() -> dict[str, "FmhaCls"]: "prims_ts": PrimsTSFmha, "prims_ts_block_sparse": PrimsTSBlockSparseFmha, # DeepSeek-V4 generation on the vendored PrimTS native sparse MLA - # kernel; opt-in like prims_ts until it graduates from validation. + # kernel. Enabled by default: it only reports itself available for + # DeepSeek-V4 on SM100/SM103, so every other model and architecture + # never sees it, and it precedes the TRTLLM-Gen fallback that served + # that phase before. "prims_ts_sparse_mla": PrimsTSSparseMlaFmha, "flashinfer_trtllm_gen": FlashInferTrtllmGenFmha, "fallback": FallbackFmha, @@ -61,7 +64,7 @@ def init_fmha_libs() -> dict[str, "FmhaCls"]: FMHA_LIBS: dict[str, FmhaCls] = init_fmha_libs() -_OPT_IN_FMHA_LIBS: frozenset[str] = frozenset({"prims_ts", "prims_ts_sparse_mla"}) +_OPT_IN_FMHA_LIBS: frozenset[str] = frozenset({"prims_ts"}) DEFAULT_FMHA_LIBS: tuple[str, ...] = tuple( name for name in FMHA_LIBS if name not in _OPT_IN_FMHA_LIBS ) 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 72116b519995..1d14cccaff63 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 @@ -35,11 +35,21 @@ def _skip_unless_toolchain() -> None: pytest.skip(reason) -def test_registry_entry_is_opt_in() -> None: +def test_registry_entry_is_default_before_fallback(monkeypatch) -> None: + """The library ships enabled and ahead of the TRTLLM-Gen fallback. + + Its ``is_available`` restricts it to DeepSeek-V4 on SM100/SM103, so being + in the default set changes nothing for other models or architectures. + """ + monkeypatch.delenv("TLLM_FMHA_LIBS", raising=False) assert registry.FMHA_LIBS[LIB_NAME] is PrimsTSSparseMlaFmha - assert LIB_NAME not in registry.DEFAULT_FMHA_LIBS + assert LIB_NAME in registry.DEFAULT_FMHA_LIBS names = tuple(registry.FMHA_LIBS) + assert names.index(LIB_NAME) < names.index("flashinfer_trtllm_gen") assert names.index(LIB_NAME) < names.index("fallback") + assert PrimsTSSparseMlaFmha in registry.get_enabled_fmha_lib_classes() + monkeypatch.setenv("TLLM_FMHA_LIBS", f"-{LIB_NAME}") + assert PrimsTSSparseMlaFmha not in registry.get_enabled_fmha_lib_classes() def _rope_table(max_positions: int, device: torch.device) -> torch.Tensor: @@ -255,9 +265,13 @@ def dequant(values, scales): ], ) def test_forward_sparse_mla_unified_prims_ts(monkeypatch, batch_name, kv_cache_dtype): - """Run the DeepSeek-V4 module test with the PrimTS library serving generation.""" + """Run the DeepSeek-V4 module test on the default library set. + + No ``TLLM_FMHA_LIBS`` override: the PrimTS library is in the default set and + must win the generation phase on its own. + """ _skip_unless_toolchain() - monkeypatch.setenv("TLLM_FMHA_LIBS", f"+{LIB_NAME}") + monkeypatch.delenv("TLLM_FMHA_LIBS", raising=False) calls = [] original_forward = PrimsTSSparseMlaFmha.forward @@ -268,3 +282,139 @@ def counting_forward(self, *args, **kwargs): monkeypatch.setattr(PrimsTSSparseMlaFmha, "forward", counting_forward) _run_unified_sparse_mla_forward(batch_name, kv_cache_dtype, "deepseek_v4") assert calls, "the PrimTS sparse MLA FMHA did not serve the generation phase" + + +def _dequant_fused(values: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + """Dequantize a fused DSv4 epilogue buffer. + + ``values`` is ``[G, T, C]`` E4M3 and ``scales`` is ``[G, C // 128, pad4(T)]`` + FP32, exactly as ``_create_dsv4_epilogue_buffers`` allocates them. + """ + groups, rows, channels = values.shape + sc = scales[..., :rows].permute(0, 2, 1) # [G, T, blocks] + blocks = values.float().view(groups, rows, channels // BLOCK, BLOCK) + return (blocks * sc[..., None]).flatten(-2), sc + + +def _assert_fused_close(got, got_sc, ref, ref_sc, channels, device): + """Same tolerance as the kernel-level fused test: two relative E4M3 ulps plus + half a quantization step, two steps on the RoPE channels D[448:512].""" + torch.testing.assert_close(got_sc, ref_sc, rtol=2e-2, atol=1e-6) + block_scales = ref_sc.repeat_interleave(BLOCK, dim=-1) + is_rope = torch.zeros(channels, dtype=torch.bool, device=device) + is_rope.view(-1, LATENT)[:, LATENT - ROPE :] = True + step_allowance = torch.where(is_rope, 2.0, 0.5) * block_scales + diff = (got - ref).abs() + ulp = ref.abs() / 8.0 + assert (diff <= 2.0 * ulp + step_allowance + 1e-6).all(), ( + f"max diff {diff.max().item():.4g} vs allowance " + f"{(2.0 * ulp + step_allowance).max().item():.4g}" + ) + assert diff.mean().item() <= 0.02 * ref.abs().mean().item() + 1e-6 + + +@pytest.mark.skipif(not HAS_FLASH_MLA, reason="FlashMLA not available") +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +@pytest.mark.parametrize("batch_name", ["small_decode", "medium_decode", "single_decode"]) +def test_forward_sparse_mla_unified_fused_epilogue(monkeypatch, batch_name, kv_cache_dtype): + """Module-level coverage of the fused DSv4 epilogue (production ADP path). + + ``prepare_sparse_attn_outputs`` only hands out the 3D fused buffer under + attention-DP with 128 heads, an FP8 cache and an FP8 ``o_a_proj``, which the + module test harness does not configure. The generation entry point + ``forward_generation_sparse_attn`` takes ``enable_dsv4_epilogue_fusion`` + directly, so wrap the harness's call to it and, on the same prepared cache, + metadata and inputs, additionally run the fused mode on the default library + set: PrimTS must serve it with ``forward_args.enable_dsv4_epilogue_fusion`` + set, and its fused output must match the unfused BF16 output pushed through + ``trtllm.fused_inv_rope_fp8_quant_vllm_port`` (the path the module takes when + fusion is off). The harness's own reference check still runs on the unfused + call. + + The TRTLLM-Gen fallback is not run in fused mode here: the module only + enables the fused epilogue for it at 128 heads, and at the harness's 64 + heads its C++ op aborts the process rather than raising. + """ + _skip_unless_toolchain() + monkeypatch.delenv("TLLM_FMHA_LIBS", raising=False) + + from .. import test_sparse_mla_forward as harness + + original_gen = harness.forward_generation_deepseek_v4_mla + original_forward = PrimsTSSparseMlaFmha.forward + seen_fused_flags: list[bool] = [] + results: dict[str, object] = {} + + def spy_forward(self, q, k, v, metadata, forward_args): + seen_fused_flags.append(bool(forward_args.enable_dsv4_epilogue_fusion)) + return original_forward(self, q, k, v, metadata, forward_args) + + _INPLACE_INPUTS = ("q", "compressed_kv", "k_pe", "latent_cache") + + def run_fused(mla, kwargs, pristine): + # The RoPE/KV-write prologue rotates ``q`` and ``latent_cache`` in + # place, so every call needs the inputs as they were before the first. + fused_kwargs = dict(kwargs) + fused_kwargs.update({k: v.clone() for k, v in pristine.items()}) + fused_kwargs["output"] = None + fused_kwargs["enable_dsv4_epilogue_fusion"] = True + return original_gen(mla, **fused_kwargs) + + def wrapped_gen(mla, **kwargs): + pristine = {k: kwargs[k].clone() for k in _INPLACE_INPUTS if kwargs.get(k) is not None} + # 1. The harness's unfused call, untouched (it validates against the + # Python reference afterwards). + out = original_gen(mla, **kwargs) + output = kwargs["output"] + rows = output.shape[0] + heads = mla.num_heads_tp + groups = mla.n_local_groups + positions = kwargs["position_ids"].view(-1).to(torch.int32) + cos_sin = mla.inverse_rotary_emb.rotary_cos_sin + torch.cuda.synchronize() + + # Path B: the standalone inverse-RoPE + FP8 quant op on the BF16 output. + ref_fp8, ref_scale = torch.ops.trtllm.fused_inv_rope_fp8_quant_vllm_port( + output.view(rows, heads, LATENT).clone(), + positions, + cos_sin, + groups, + heads // groups, + LATENT - ROPE, + ROPE, + BLOCK, + False, + ) + + # 2. Path A on the default library set: PrimTS serves the fused epilogue. + # (The module hands out the 3D buffer only under ADP + 128 heads + FP8, + # which is what production throughput configs run.) + before = len(seen_fused_flags) + prims_fp8, prims_scale = run_fused(mla, kwargs, pristine) + torch.cuda.synchronize() + assert seen_fused_flags[before:] == [True], ( + "the PrimTS sparse MLA FMHA did not serve the fused generation call" + ) + + results.update( + rows=rows, + heads=heads, + groups=groups, + ref=(ref_fp8.clone(), ref_scale.clone()), + prims=(prims_fp8.clone(), prims_scale.clone()), + ) + return out + + monkeypatch.setattr(PrimsTSSparseMlaFmha, "forward", spy_forward) + monkeypatch.setattr(harness, "forward_generation_deepseek_v4_mla", wrapped_gen) + _run_unified_sparse_mla_forward(batch_name, kv_cache_dtype, "deepseek_v4") + + assert results, "the generation phase never ran" + assert seen_fused_flags.count(False) >= 1, "the unfused call did not reach PrimTS" + device = results["prims"][0].device + channels = results["heads"] // results["groups"] * LATENT + + ref, ref_sc = _dequant_fused(*results["ref"]) + got, got_sc = _dequant_fused(*results["prims"]) + assert got.shape == ref.shape + _assert_fused_close(got, got_sc, ref, ref_sc, channels, device) diff --git a/tests/unittest/_torch/attention/test_fmha_registry.py b/tests/unittest/_torch/attention/test_fmha_registry.py index 6bdb30dcb0fa..171feed20cb3 100644 --- a/tests/unittest/_torch/attention/test_fmha_registry.py +++ b/tests/unittest/_torch/attention/test_fmha_registry.py @@ -44,7 +44,7 @@ def test_default_fmha_libs_exclude_prims_ts(monkeypatch: pytest.MonkeyPatch) -> monkeypatch.delenv("TLLM_FMHA_LIBS", raising=False) assert PRIMS_TS not in registry.DEFAULT_FMHA_LIBS - assert PRIMS_TS_SPARSE_MLA not in registry.DEFAULT_FMHA_LIBS + assert PRIMS_TS_SPARSE_MLA in registry.DEFAULT_FMHA_LIBS assert PRIMS_TS_BLOCK_SPARSE in registry.DEFAULT_FMHA_LIBS assert set(registry.DEFAULT_FMHA_LIBS) <= set(registry.FMHA_LIBS) assert _enabled_names() == registry.DEFAULT_FMHA_LIBS From deb83e7625e3d51c32b5b287a31f25315a6ec4cf Mon Sep 17 00:00:00 2001 From: yunruis <205571022+yunruis@users.noreply.github.com> Date: Thu, 15 Oct 2026 07:40:34 -0700 Subject: [PATCH 04/19] [None][fix] Use the qk_nope + qk_rope softmax scale in the PrimTS sparse MLA FMHA The adapter scaled the generation softmax by 1 / (q_scaling * sqrt(512)), the absorbed latent width the kernel contracts over. TRT-LLM's MLA convention (attentionOp.cpp and the DeepSeek-V4 TRTLLM-Gen fallback) is 1 / (q_scaling * sqrt(qk_nope_head_dim + qk_rope_head_dim)). DeepSeek-V4 checkpoints have 448 + 64 = 512, so the released models were unaffected, but the 128 + 64 scenario of test_deepseek_v4_sparse_mla ran with a 1.63x flatter softmax: L0 #64329 reported ~0.3 BF16 output mismatches on B200 and B300 for every generation case of that test (8 failures per stage). Factor the scale into PrimsTSSparseMlaFmha._softmax_scale and pin the convention with a unit test. With the fix test_deepseek_v4_sparse_mla passes (9/9, diff mean ~1e-3 vs ~3e-2 before) and the PrimTS unit tests still pass (39/39) on B200. Co-Authored-By: Claude Fable 5.1 Signed-off-by: yunruis <205571022+yunruis@users.noreply.github.com> --- .../backends/fmha/prims_ts_sparse_mla.py | 14 +++++++++++++- .../test_prims_ts_sparse_mla_fmha.py | 17 +++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py index 5a333eaffc80..12bf8b37e985 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py @@ -429,6 +429,18 @@ def _sinks(self, fwd: AttentionForwardArgs) -> Optional[torch.Tensor]: self._sinks_cache = (sinks.data_ptr(), prepared) return prepared + @staticmethod + def _softmax_scale(attn: "TrtllmAttention") -> float: + """TRT-LLM's MLA softmax scale, ``1 / (q_scaling * sqrt(qk_nope + qk_rope))``. + + The scale follows the un-absorbed QK head width (attentionOp.cpp and the + DeepSeek-V4 fallback), not the 512-wide absorbed latent the kernel + contracts over. DeepSeek-V4 checkpoints have 448 + 64 = 512, so the two + only differ for configs such as the 128 + 64 unit-test scenario. + """ + qk_head_dim = int(attn.qk_nope_head_dim) + int(attn.qk_rope_head_dim) + return 1.0 / (float(attn.q_scaling) * math.sqrt(float(qk_head_dim))) + def _kv_dequant(self) -> float: if self._kv_dequant_scale is None: if torch.cuda.is_current_stream_capturing(): @@ -549,7 +561,7 @@ def forward( output = forward_args.output groups = int(output.shape[0]) if fused else None - softmax_scale = 1.0 / (float(attn.q_scaling) * math.sqrt(float(_LATENT_DIM))) + softmax_scale = self._softmax_scale(attn) if is_fp8: # The MLA preprocessing quantizes Q with the KV-cache scale, so both # descales equal the layer's kv_scale_quant_orig (see 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 1d14cccaff63..9fe2fc754035 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 @@ -2,6 +2,9 @@ # SPDX-License-Identifier: Apache-2.0 """Tests for the PrimTS native sparse MLA FMHA library (DeepSeek-V4, SM100/SM103).""" +import math +from types import SimpleNamespace + import pytest import torch import triton @@ -52,6 +55,20 @@ def test_registry_entry_is_default_before_fallback(monkeypatch) -> None: assert PrimsTSSparseMlaFmha not in registry.get_enabled_fmha_lib_classes() +def test_softmax_scale_follows_qk_head_dim() -> None: + """The scale uses ``qk_nope + qk_rope`` like attentionOp.cpp, not the 512 latent. + + DeepSeek-V4 checkpoints have 448 + 64 = 512, so only a narrower nope width + (the 128 + 64 scenario of ``test_deepseek_v4_sparse_mla``) tells the two apart; + L0 #64329 caught the latent-width version as a ~0.3 BF16 output mismatch. + """ + attn = SimpleNamespace(qk_nope_head_dim=128, qk_rope_head_dim=64, q_scaling=0.5336) + expected = 1.0 / (0.5336 * math.sqrt(192.0)) + assert PrimsTSSparseMlaFmha._softmax_scale(attn) == pytest.approx(expected, rel=1e-12) + attn = SimpleNamespace(qk_nope_head_dim=448, qk_rope_head_dim=64, q_scaling=1.0) + assert PrimsTSSparseMlaFmha._softmax_scale(attn) == pytest.approx(LATENT**-0.5, rel=1e-12) + + def _rope_table(max_positions: int, device: torch.device) -> torch.Tensor: """Forward [cos(32), sin(32)] rows, the layout both kernels consume.""" inv_freq = 1.0 / ( From a046511cc0b6761dba220be5c32efe8ffee1c6a9 Mon Sep 17 00:00:00 2001 From: yunruis <205571022+yunruis@users.noreply.github.com> Date: Thu, 15 Oct 2026 08:17:59 -0700 Subject: [PATCH 05/19] [None][fix] Keep SM103 on the TRTLLM-Gen fallback for the PrimTS sparse MLA FMHA The vendored PrimTS MLA kernels admit SM103 through their architecture guard, but flashinfer has only signed off accuracy and performance on SM100 (prims_ts/kernels/mla_decode/README.md: "SM103a/B300 (architecture-gated, not yet signoff-qualified)"). TRT-LLM L0 #64329 hit an illegal memory access in the GB300 DeepSeek-V4-Pro perf-sanity warmup (ctx block of gb300_deepseek-v4-pro-fp4_8k1k_con8_ctx1_dep4_gen4_tep8_eplb0_mtp3_ccb: TP4 ADP, FP8 KV, MTP 3, no CUDA graphs) right after self_attn, and that case passes on every other recent build. The same configuration at TP8 runs clean on B200 (TP4 does not fit in 178 GB), so the failure is treated as SM103-specific and B300/GB300 stay on the fallback until the kernels are qualified there. Report SM103 as unavailable, keep the unit tests that need the library to serve a batch SM100-only, and add a test pinning the SM100/SM103/other availability answers. Co-Authored-By: Claude Fable 5.1 Signed-off-by: yunruis <205571022+yunruis@users.noreply.github.com> --- .../backends/fmha/prims_ts_sparse_mla.py | 19 ++++++++--- .../attention/backends/fmha/registry.py | 2 +- .../test_prims_ts_sparse_mla_fmha.py | 34 +++++++++++++++++-- 3 files changed, 48 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py index 12bf8b37e985..71357838ce01 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py @@ -15,7 +15,7 @@ """TRT-LLM FMHA adapter for the vendored PrimTS native sparse MLA kernels. -Serves the DeepSeek-V4 generation phase on SM100/SM103 with the Blackwell +Serves the DeepSeek-V4 generation phase on SM100 with the Blackwell task-scheduled two-source sparse MLA decode kernel (``BatchSparseMLADecodePagedTSWrapper``). The sliding-window pool is the primary source and, for compressing layers, the compressed pool is the extra @@ -33,9 +33,11 @@ ``o_a_proj`` BMM consumes the attention output without a separate quant kernel. -The library is enabled by default on SM100/SM103 and precedes the TRTLLM-Gen +The library is enabled by default on SM100 and precedes the TRTLLM-Gen fallback for the DeepSeek-V4 generation phase; ``TLLM_FMHA_LIBS=-prims_ts_sparse_mla`` restores the previous library. The context phase stays on the fallback library. +SM103 (B300/GB300) is admitted by the kernels' architecture guard but not yet +signoff-qualified upstream, so it keeps the fallback for now. """ from __future__ import annotations @@ -212,8 +214,17 @@ def _is_available(cls, attn: "TrtllmAttention") -> bool: @classmethod def _unavailable_reason(cls, attn: "TrtllmAttention") -> Optional[str]: sm = get_sm_version() - if sm not in (100, 103): - return f"requires SM100 or SM103, got SM{sm}." + if sm == 103: + # The vendored kernels admit SM103 through their architecture guard + # but flashinfer has only signed off accuracy and performance on + # SM100 (prims_ts/kernels/mla_decode/README.md). TRT-LLM L0 #64329 + # hit an illegal memory access in the GB300 DeepSeek-V4-Pro perf + # sanity warmup (TP4/ADP, FP8 KV, MTP 3) that the same configuration + # at TP8 does not reproduce on B200, so B300/GB300 stay on the + # TRTLLM-Gen fallback until the kernels are qualified there. + return "SM103 is not yet qualified for the PrimTS sparse MLA kernels." + if sm != 100: + return f"requires SM100, got SM{sm}." if getattr(attn, "use_fp8_ds_mla", False) or getattr(attn, "kv_cache_dtype", "auto") == ( "fp8_ds_mla" ): diff --git a/tensorrt_llm/_torch/attention/backends/fmha/registry.py b/tensorrt_llm/_torch/attention/backends/fmha/registry.py index e5c8e2ec68a2..c0340d61b85e 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/registry.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/registry.py @@ -54,7 +54,7 @@ def init_fmha_libs() -> dict[str, "FmhaCls"]: "prims_ts_block_sparse": PrimsTSBlockSparseFmha, # DeepSeek-V4 generation on the vendored PrimTS native sparse MLA # kernel. Enabled by default: it only reports itself available for - # DeepSeek-V4 on SM100/SM103, so every other model and architecture + # DeepSeek-V4 on SM100, so every other model and architecture # never sees it, and it precedes the TRTLLM-Gen fallback that served # that phase before. "prims_ts_sparse_mla": PrimsTSSparseMlaFmha, 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 9fe2fc754035..fe4d8086e833 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 @@ -1,6 +1,6 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Tests for the PrimTS native sparse MLA FMHA library (DeepSeek-V4, SM100/SM103).""" +"""Tests for the PrimTS native sparse MLA FMHA library (DeepSeek-V4, SM100).""" import math from types import SimpleNamespace @@ -33,6 +33,9 @@ def _skip_unless_toolchain() -> None: + """Skip the tests that need the library to actually serve a batch.""" + if get_sm_version() != 100: + pytest.skip("PrimTS sparse MLA serves SM100 only (SM103 is not yet qualified)") reason = get_prims_ts_toolchain_unsupported_reason() if reason is not None: pytest.skip(reason) @@ -41,7 +44,7 @@ def _skip_unless_toolchain() -> None: def test_registry_entry_is_default_before_fallback(monkeypatch) -> None: """The library ships enabled and ahead of the TRTLLM-Gen fallback. - Its ``is_available`` restricts it to DeepSeek-V4 on SM100/SM103, so being + Its ``is_available`` restricts it to DeepSeek-V4 on SM100, so being in the default set changes nothing for other models or architectures. """ monkeypatch.delenv("TLLM_FMHA_LIBS", raising=False) @@ -55,6 +58,33 @@ def test_registry_entry_is_default_before_fallback(monkeypatch) -> None: assert PrimsTSSparseMlaFmha not in registry.get_enabled_fmha_lib_classes() +def test_sm103_stays_on_the_fallback_until_qualified(monkeypatch) -> None: + """B300/GB300 report an SM reason; SM100 proceeds to the shape and toolchain checks. + + The vendored kernels admit SM103 through their architecture guard, but the + flashinfer signoff is SM100 only and L0 #64329 hit an illegal memory access + in the GB300 DeepSeek-V4-Pro warmup, so the adapter keeps SM103 on TRTLLM-Gen. + """ + import tensorrt_llm._torch.attention.backends.fmha.prims_ts_sparse_mla as mod + + attn = SimpleNamespace( + use_fp8_ds_mla=False, + kv_cache_dtype="auto", + _uses_nvfp4_compress=False, + kv_lora_rank=448, + qk_rope_head_dim=64, + head_dim=512, + num_heads=16, + ) + monkeypatch.setattr(mod, "get_sm_version", lambda: 103) + assert "SM103" in PrimsTSSparseMlaFmha._unavailable_reason(attn) + monkeypatch.setattr(mod, "get_sm_version", lambda: 90) + assert "requires SM100" in PrimsTSSparseMlaFmha._unavailable_reason(attn) + monkeypatch.setattr(mod, "get_sm_version", lambda: 100) + reason = PrimsTSSparseMlaFmha._unavailable_reason(attn) + assert reason is None or reason == get_prims_ts_toolchain_unsupported_reason() + + def test_softmax_scale_follows_qk_head_dim() -> None: """The scale uses ``qk_nope + qk_rope`` like attentionOp.cpp, not the 512 latent. From 7b7a61d13d576dc4de10a78357d51f5c1222db85 Mon Sep 17 00:00:00 2001 From: yunruis <205571022+yunruis@users.noreply.github.com> Date: Tue, 20 Oct 2026 03:19:44 -0700 Subject: [PATCH 06/19] [None][chore] Re-vendor PrimTS from flashinfer main and adapt the sparse MLA FMHA Pin flashinfer-prims-ts to flashinfer-ai/flashinfer main 3bf6d3202. It carries the merged native two-source sparse MLA (flashinfer-ai/flashinfer#5434), including the group-major fused output and the 2CTA short-list schedule that flashinfer-ai/flashinfer#6345 had stacked on the pre-merge branch, plus the PrimTS changes main took since ae312f7cc (#6529 compile cache and FP8-to-BF16 decode, #6482, #5351, #6089, #6014). The compatibility patch is regenerated for the new tree: the flashinfer-python 0.6.18 API alias, stripped trace templates, absolute imports of the flashinfer cute_dsl/utils/jit helpers (now also jit.cute_dsl_core and cute_dsl.availability), relative kernel imports in the sparse plan resolver, and the module-scope helpers_paged import for CuTe DSL 4.8 that upstream TRT-LLM already carried. The merged kernel API takes the two index lists and their live lengths directly and combines the Q and KV descales into the softmax and output scales, so the adapter drops the SparseMLAPreparedMetadata route packing (the Triton kernel, routes, execution lengths and scale rows) and binds the sparse_attn_predict index buffer plus two int32 length buffers per plan. The fused group-major output contract is unchanged. The adapter unit test drives the wrapper through the same direct-index API. AI-assisted. Co-Authored-By: Claude Fable 5.1 Signed-off-by: yunruis <205571022+yunruis@users.noreply.github.com> --- .../vendor_patches/flashinfer-prims-ts.patch | 92 ++- 3rdparty/vendor_sources.lock.yaml | 10 +- .../backends/fmha/prims_ts_sparse_mla.py | 186 ++---- .../attention/backends/prims_ts/__init__.py | 2 - .../prims_ts/_block_sparse/compiler.py | 9 +- .../backends/prims_ts/_compile_cache.py | 108 ++++ .../attention/backends/prims_ts/_hardware.py | 41 ++ .../_q_token_kv_block_sparse_policy.py | 12 +- .../attention/backends/prims_ts/context.py | 40 +- .../attention/backends/prims_ts/decode.py | 137 +++-- .../kernels/fmha_context/fmha_kernel.py | 5 + .../kernels/fmha_context/fmha_resources.py | 1 + .../fmha_decode/block_sparse_inspect.py | 80 ++- .../kernels/fmha_decode/fmha_decode_config.py | 250 ++++++-- .../fmha_decode/fmha_decode_constants.py | 13 +- .../kernels/fmha_decode/fmha_decode_kernel.py | 166 +++-- .../fmha_decode_resources/smem_resources.py | 331 ++++++++-- .../fmha_decode_resources/tmem_o.py | 4 +- .../fmha_decode_resources/tmem_s.py | 38 +- .../tmem_transformed_kv.py | 21 +- .../kernels/fmha_decode/fmha_decode_tasks.py | 44 +- .../prims_ts/kernels/fmha_decode/reduction.py | 97 ++- .../kernels/mla_decode/helpers/gather.py | 20 +- .../mla_decode/helpers/sparse_indices.py | 308 ++++++++++ .../kernels/mla_decode/sparse_policy.py | 21 +- .../kernels/mla_decode/sparse_views.py | 314 ---------- .../mla_decode/throughput_2cta/kernel.py | 51 +- .../mla_decode/throughput_2cta/resources.py | 98 ++- .../throughput_latency_1cta/config.py | 11 +- .../throughput_latency_1cta/kernel.py | 66 +- .../resources/smem_resources.py | 23 +- .../resources/tmem_corr.py | 20 +- .../resources/tmem_s.py | 4 +- .../prims_ts/kernels/separate_reduction.py | 113 +--- .../attention/backends/prims_ts/mla_decode.py | 576 ++++++------------ .../q_token_kv_block_sparse_metadata.py | 279 +++++---- .../test_prims_ts_sparse_mla_fmha.py | 70 +-- 37 files changed, 2021 insertions(+), 1640 deletions(-) create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/_compile_cache.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/_hardware.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/sparse_indices.py delete mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/sparse_views.py diff --git a/3rdparty/vendor_patches/flashinfer-prims-ts.patch b/3rdparty/vendor_patches/flashinfer-prims-ts.patch index 67992c0ead8b..027b9b65a22f 100644 --- a/3rdparty/vendor_patches/flashinfer-prims-ts.patch +++ b/3rdparty/vendor_patches/flashinfer-prims-ts.patch @@ -1,3 +1,21 @@ +diff --git a/_compile_cache.py b/_compile_cache.py +index b374a65311402982d818d80e9145ab40ed9727a1..db8416f07e37945a8e8764ac202e446f301d54d8 100644 +--- a/_compile_cache.py ++++ b/_compile_cache.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"); +@@ -82,7 +83,7 @@ def build_and_load_attention_kernel( + extra: tuple = (), + ) -> Any: + """Build or load under the caller's selected CUDA device context.""" +- from ...jit.cute_dsl_core import ( ++ from flashinfer.jit.cute_dsl_core import ( + build_and_load_cute_dsl_kernel, + cute_dsl_cache_disabled, + ) diff --git a/block_sparse.py b/block_sparse.py index fa055618705216cddf6b7f3e9ed4c1d680c2af9a..38de4926a045f17c5b6f206f4a20ecbe2532a1c7 100644 --- a/block_sparse.py @@ -61,10 +79,10 @@ index fa055618705216cddf6b7f3e9ed4c1d680c2af9a..38de4926a045f17c5b6f206f4a20ecbe q: torch.Tensor, paged_kv_cache: PagedKVCache, diff --git a/context.py b/context.py -index 53f01af72032ae2b01ea2ac374661af6387dd33f..d297bc79c1e00610783dd54a0d8cb22d1bcc2020 100644 +index d3f104e4e34b1bcb7ada697f9e76f7013d6ca100..1ec1080b942f60fad22e3e4a1627cad9592697ad 100644 --- a/context.py +++ b/context.py -@@ -45,7 +45,7 @@ from typing import TYPE_CHECKING, Callable, Literal, Optional +@@ -48,7 +48,7 @@ from ._hardware import get_max_active_clusters import torch @@ -73,7 +91,7 @@ index 53f01af72032ae2b01ea2ac374661af6387dd33f..d297bc79c1e00610783dd54a0d8cb22d if TYPE_CHECKING: -@@ -490,7 +490,7 @@ def _validate_device(device: torch.device) -> int: +@@ -493,7 +493,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): @@ -83,7 +101,7 @@ index 53f01af72032ae2b01ea2ac374661af6387dd33f..d297bc79c1e00610783dd54a0d8cb22d require_cute_dsl_arch(device_index) return device_index diff --git a/decode.py b/decode.py -index d85dd274917cf9a5a67ef1400155bb60985f1f77..205e875fc39022ca492cb40acea5b329fc4a3c11 100644 +index 4fb096f39849a3801349758360e71b092ee200e9..f91814650c48dcf5b116273742880ac722fed527 100644 --- a/decode.py +++ b/decode.py @@ -1,3 +1,4 @@ @@ -91,9 +109,9 @@ index d85dd274917cf9a5a67ef1400155bb60985f1f77..205e875fc39022ca492cb40acea5b329 # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); -@@ -24,11 +25,7 @@ from typing import TYPE_CHECKING, Literal, NamedTuple, Optional, Union, cast +@@ -26,11 +27,7 @@ import torch - import torch + from ._compile_cache import build_and_load_attention_kernel -from flashinfer.api_logging import flashinfer_experimental_api -from flashinfer.trace.templates.attention import ( @@ -104,7 +122,7 @@ index d85dd274917cf9a5a67ef1400155bb60985f1f77..205e875fc39022ca492cb40acea5b329 from . import _q_token_kv_block_sparse_policy as _sparse_policy from ._block_sparse.common import _num_sparse_pattern_heads -@@ -789,7 +786,7 @@ def _validate_runtime_device(device: torch.device) -> None: +@@ -780,7 +777,7 @@ def _validate_runtime_device(device: torch.device) -> None: # 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): @@ -113,7 +131,7 @@ index d85dd274917cf9a5a67ef1400155bb60985f1f77..205e875fc39022ca492cb40acea5b329 require_cute_dsl_arch(device_index) -@@ -4048,7 +4045,7 @@ class BatchDecodePagedTSWrapper: +@@ -4075,7 +4072,7 @@ class BatchDecodePagedTSWrapper: # previous complete plan revision usable. self._plan_state = candidate @@ -122,7 +140,7 @@ index d85dd274917cf9a5a67ef1400155bb60985f1f77..205e875fc39022ca492cb40acea5b329 def run( self, q: torch.Tensor, -@@ -4230,7 +4227,7 @@ class BatchDecodePagedTSWrapper: +@@ -4257,7 +4254,7 @@ class BatchDecodePagedTSWrapper: ) @@ -132,10 +150,15 @@ index d85dd274917cf9a5a67ef1400155bb60985f1f77..205e875fc39022ca492cb40acea5b329 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..e6cdddf8d5cb396be5feb828535c0633428e32c3 100644 +index d17584b725e24235ce0ec90caac174a29d4ab965..af69ee7d51bc1179352d0bf2ba71fad2fe30a922 100644 --- a/kernels/fmha_context/fmha_resources.py +++ b/kernels/fmha_context/fmha_resources.py -@@ -1199,8 +1199,6 @@ class GmemQKVResource(MemoryResource): +@@ -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"); +@@ -1199,8 +1200,6 @@ class GmemQKVResource(MemoryResource): if cutlass.const_expr(self.cfg.has_uniform_varlen): seqlen_k = Int32(self.cfg.uniform_seq_len_k) else: @@ -144,7 +167,7 @@ index d17584b725e24235ce0ec90caac174a29d4ab965..e6cdddf8d5cb396be5feb828535c0633 seqlen_k = _load_runtime_seq_len_kv( self.seqlens_kv, self.max_seq_len_kv, batch_coord ) -@@ -1223,8 +1221,6 @@ class GmemQKVResource(MemoryResource): +@@ -1223,8 +1222,6 @@ class GmemQKVResource(MemoryResource): kv_request_begin = batch_coord * Int32(self.block_table_row_stride) kv_page_idx_ub = Int32(self.cfg.max_num_pages_per_seq_kv - 1) else: @@ -153,7 +176,7 @@ index d17584b725e24235ce0ec90caac174a29d4ab965..e6cdddf8d5cb396be5feb828535c0633 kv_request_begin, kv_page_idx_ub = _load_block_table_row_bounds( Int32(self.block_table_row_stride), self.cfg, -@@ -1300,13 +1296,9 @@ class GmemQKVResource(MemoryResource): +@@ -1300,13 +1297,9 @@ class GmemQKVResource(MemoryResource): kv_request_begin = batch_coord * Int32(self.block_table_row_stride) kv_page_idx_ub = Int32(self.cfg.max_num_pages_per_seq_kv - 1) else: @@ -167,7 +190,7 @@ index d17584b725e24235ce0ec90caac174a29d4ab965..e6cdddf8d5cb396be5feb828535c0633 kv_request_begin, kv_page_idx_ub = _load_block_table_row_bounds( Int32(self.block_table_row_stride), self.cfg, -@@ -1743,8 +1735,6 @@ class SmemPageOffsetsKvResource(MemoryResource): +@@ -1743,8 +1736,6 @@ class SmemPageOffsetsKvResource(MemoryResource): kv_request_begin: Int32, kv_page_idx_ub: Int32, ) -> None: @@ -176,7 +199,7 @@ index d17584b725e24235ce0ec90caac174a29d4ab965..e6cdddf8d5cb396be5feb828535c0633 cfg = self.cfg # Context's K/V tile index is kv_tile_start + loop_offset; for V the # producer runs one tile ahead in TAIL, so reuse the same expression -@@ -1835,8 +1825,6 @@ class SmemPageOffsetsKvResource(MemoryResource): +@@ -1835,8 +1826,6 @@ class SmemPageOffsetsKvResource(MemoryResource): tile_offset: cutlass.Constexpr[int] = 0, ) -> cutlass.Array: """Retain one tile's page IDs after its SMEM window is released.""" @@ -185,7 +208,7 @@ index d17584b725e24235ce0ec90caac174a29d4ab965..e6cdddf8d5cb396be5feb828535c0633 tile_idx = _resolve_kv_tile_idx_context( stage_info, kv_tile_start, tile_offset=tile_offset ) -@@ -5770,3 +5758,14 @@ class GmemOResource(MemoryResource): +@@ -5770,3 +5759,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) @@ -200,8 +223,26 @@ index d17584b725e24235ce0ec90caac174a29d4ab965..e6cdddf8d5cb396be5feb828535c0633 + _load_runtime_seq_len_kv, + _resolve_kv_tile_idx_context, +) +diff --git a/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py b/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py +index b22f0f413b7312bab9a5f8e296255fe9aeabe901..a4a87200ca5d3ace5cf899caa9672a1a32c5289e 100644 +--- a/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py ++++ b/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.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"); +@@ -36,7 +37,7 @@ from cutlass.experimental.task_scheduling.resources import ( + producer_work, + ) + +-from ......cute_dsl.availability import is_rubin_cute_dsl_available ++from flashinfer.cute_dsl.availability import is_rubin_cute_dsl_available + from ..fmha_decode_config import FmhaDecodeConfig + from .helpers_common import ( + _TASK_CACHE_TMEM_BASE_OFFSET, diff --git a/mla_decode.py b/mla_decode.py -index 208e495869370e5aea2693548e0ff4d186c70d5e..70c3d0e6423d9810f987e59991cbb7f5b1df8337 100644 +index b62bd1081fbae0e18e0950008bae87c8986677b4..e64fc538c768f08cc4d51ce12e74e00ac59cd692 100644 --- a/mla_decode.py +++ b/mla_decode.py @@ -1,3 +1,4 @@ @@ -209,7 +250,7 @@ index 208e495869370e5aea2693548e0ff4d186c70d5e..70c3d0e6423d9810f987e59991cbb7f5 # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); -@@ -24,12 +25,7 @@ from typing import Any, Literal, NamedTuple, Optional, cast +@@ -27,11 +28,7 @@ from ._hardware import get_max_active_clusters import torch @@ -218,12 +259,11 @@ index 208e495869370e5aea2693548e0ff4d186c70d5e..70c3d0e6423d9810f987e59991cbb7f5 - prims_ts_decode_mla_one_shot_trace_dispatch, - prims_ts_decode_mla_wrapper_trace_dispatch, -) -- +from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api + from .decode import ( _WorkspaceSection, - _align_up, -@@ -1650,7 +1646,7 @@ class BatchMLADecodePagedTSWrapper: +@@ -1637,7 +1634,7 @@ class BatchMLADecodePagedTSWrapper: split_kv=int(dict(policy)["split_kv"]), ) @@ -232,7 +272,7 @@ index 208e495869370e5aea2693548e0ff4d186c70d5e..70c3d0e6423d9810f987e59991cbb7f5 def run( self, query: torch.Tensor, -@@ -1755,7 +1751,7 @@ class BatchMLADecodePagedTSWrapper: +@@ -1742,7 +1739,7 @@ class BatchMLADecodePagedTSWrapper: ) @@ -241,7 +281,7 @@ index 208e495869370e5aea2693548e0ff4d186c70d5e..70c3d0e6423d9810f987e59991cbb7f5 def batch_mla_decode_with_paged_kv_cache( query: torch.Tensor, kv_cache: torch.Tensor, -@@ -2189,17 +2185,17 @@ def _resolve_sparse_mla_plan( +@@ -2119,17 +2116,17 @@ def _resolve_sparse_mla_plan( ): """Resolve sparse geometry and scratch sizing without allocation/compilation.""" import cutlass @@ -264,7 +304,7 @@ index 208e495869370e5aea2693548e0ff4d186c70d5e..70c3d0e6423d9810f987e59991cbb7f5 ) diff --git a/q_token_kv_block_sparse_metadata.py b/q_token_kv_block_sparse_metadata.py -index 35efc60b7ff3dac6af1c5c5e46631712a171fb72..0c5cf40dbd4e4708b8f74bfa698f3c0788f9f3a4 100644 +index e2bbef13c0f552030875e6b433e0fcdb20686cf6..dff16046cf7cf9ff03950043cf462ac28de6c670 100644 --- a/q_token_kv_block_sparse_metadata.py +++ b/q_token_kv_block_sparse_metadata.py @@ -1,3 +1,4 @@ @@ -272,7 +312,7 @@ index 35efc60b7ff3dac6af1c5c5e46631712a171fb72..0c5cf40dbd4e4708b8f74bfa698f3c07 # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); -@@ -40,9 +41,9 @@ from dataclasses import dataclass +@@ -41,9 +42,9 @@ from dataclasses import dataclass from typing import Any, Literal, Optional import torch @@ -284,7 +324,7 @@ index 35efc60b7ff3dac6af1c5c5e46631712a171fb72..0c5cf40dbd4e4708b8f74bfa698f3c07 from ._block_sparse.common import _num_sparse_pattern_heads _Q_TOKEN_KV_BLOCK_SPARSE_MIN_BLOCK_SIZE = 4 -@@ -62,7 +63,7 @@ _Q_TOKEN_KV_BLOCK_SPARSE_WORKSPACE_ALIGNMENT = 256 +@@ -63,7 +64,7 @@ _Q_TOKEN_KV_BLOCK_SPARSE_WORKSPACE_ALIGNMENT = 256 def _get_prims_ts_q_token_kv_block_sparse_metadata_module() -> Any: """Build and cache the CUDA C++ QToken-KvBlock-Sparse-Attention metadata implementation.""" diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 184250de470b..3d4f00be8a0a 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -13,16 +13,16 @@ vendors: patch_digest: sha256:bc9b96a69873e0ddfcba66ab899c666cf52766976f24d850048d2a0bf8d967ff digest: sha256-tree-v1:d2d382335ccb68b55078523e63835c28b699607227afb81308cb7bb1d85845ac flashinfer-prims-ts: - url: https://github.com/yunruis/flashinfer.git - branch: yunruis/prims-ts-sparse-mla-group-major-out - commit: 549977ef961f00a4ff2feb4c9dd75b0cbf222a3a + url: https://github.com/flashinfer-ai/flashinfer.git + branch: main + commit: 3bf6d3202239a107989451257c6f594cb5d2f008 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:5cf178611ac8a9db64b9465218c2ce800e0e336b80a4129d086093089898e164 - digest: sha256-tree-v1:5cfc19d5eda263af15f8e6661a44e7bfc82d4524e5b6a7734297edbb55de7dce + patch_digest: sha256:53c48256776f437b5910a58c031226c594c56902c936f21c3989d01c06fe5c48 + digest: sha256-tree-v1:1f1ab05273ece9646bfcc6b6baf27cc0dc042d3bf6e1d16f47ee45270baee78e openengine-proto: url: https://github.com/ai-dynamo/openengine.git tag: v0.1.0 diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py index 71357838ce01..80f7d47e3565 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py @@ -17,18 +17,19 @@ Serves the DeepSeek-V4 generation phase on SM100 with the Blackwell task-scheduled two-source sparse MLA decode kernel -(``BatchSparseMLADecodePagedTSWrapper``). The sliding-window pool is the -primary source and, for compressing layers, the compressed pool is the extra -source. Both lists feed one softmax, exactly like the TRTLLM-Gen sparse MLA -kernel the fallback library runs. +(``BatchSparseMLADecodePagedTSWrapper``, flashinfer-ai/flashinfer#5434). The +sliding-window pool is the primary source and, for compressing layers, the +compressed pool is the extra source. Both lists feed one softmax, exactly like +the TRTLLM-Gen sparse MLA kernel the fallback library runs. The adapter consumes the dual-pool global indices that ``DeepseekV4TrtllmAttention.sparse_attn_predict`` already produces: SWA slots are storage rows of the SWA pool and compressed slots are storage rows of the -layer's compress pool, so no page-table conversion happens here. The fused -DSv4 epilogue (conjugate RoPE on ``D[448:512]`` + 1x128 block-scaled E4M3 -output in the group-major ``[G, T, (H/G)*512]`` layout with FP32 scales -``[G, (H/G)*4, pad4(T)]``) maps one-to-one onto the kernel's +layer's compress pool, so no page-table conversion happens here. The kernel +takes the two index lists and their live lengths directly and masks ``-1`` +holes itself. The fused DSv4 epilogue (conjugate RoPE on ``D[448:512]`` + +1x128 block-scaled E4M3 output in the group-major ``[G, T, (H/G)*512]`` layout +with FP32 scales ``[G, (H/G)*4, pad4(T)]``) maps one-to-one onto the kernel's ``output_quant_block_size=128`` / ``out_scale_format="fp32"`` mode, so the ``o_a_proj`` BMM consumes the attention output without a separate quant kernel. @@ -49,8 +50,6 @@ from typing import TYPE_CHECKING, Optional import torch -import triton -import triton.language as tl from packaging.version import InvalidVersion, Version from tensorrt_llm._torch.attention.backends.interface import ( @@ -65,7 +64,6 @@ if TYPE_CHECKING: from tensorrt_llm._torch.attention.backends.prims_ts.mla_decode import ( BatchSparseMLADecodePagedTSWrapper, - SparseMLAPreparedMetadata, ) from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.cache_manager import ( DeepseekV4CacheManager, @@ -83,8 +81,6 @@ _LATENT_DIM = 512 _ROPE_DIM = 64 _FUSED_OUTPUT_BLOCK = 128 -_ROUTE_INVALID = 0x7FFFFFFF -_ROUTE_EXTRA_BIT = -2147483648 def get_prims_ts_toolchain_unsupported_reason() -> Optional[str]: @@ -104,74 +100,18 @@ def get_prims_ts_toolchain_unsupported_reason() -> Optional[str]: return None -@triton.jit -def _pack_sparse_mla_metadata( - swa_indices, - extra_indices, - swa_lengths, - extra_lengths, - routes, - execution_lengths, - valid_counts, - scale_params, - sinks, - bmm1_scale, - bmm2_scale, - KS: tl.constexpr, - KC: tl.constexpr, - ROWS: tl.constexpr, - HEADS: tl.constexpr, - CAP: tl.constexpr, - BLOCK: tl.constexpr, - HEAD_BLOCK: tl.constexpr, -): - """Build the combined route words and scale rows for one shared-scale pass. - - Mirrors the ``SparseMLAPreparedMetadata`` contract: each source span is - rounded to 128 slots, extra-pool rows carry bit 31, ``0x7fffffff`` marks an - inactive slot, and the scale row holds ``[bmm1, bmm2, sinks[H], counts[rows]]``. - """ - row = tl.program_id(0) - ls = tl.load(swa_lengths + row) - lc = tl.load(extra_lengths + row) if KC else 0 - span_s = tl.cdiv(ls, 128) * 128 - span_c = tl.cdiv(lc, 128) * 128 - col = tl.arange(0, BLOCK) - for start in range(0, CAP, BLOCK): - p = start + col - a = tl.load(swa_indices + row * KS + p, (p < ls) & (p < CAP), 0x7FFFFFFF) - local = p - span_s - if KC: - b = tl.load( - extra_indices + row * KC + local, - (local >= 0) & (local < lc) & (p < CAP), - 0x7FFFFFFF, - ) - else: - b = tl.full((BLOCK,), 0x7FFFFFFF, tl.int32) - b = b | -2147483648 - word = tl.where(p < span_s, a, b) - tl.store(routes + row * CAP + p, word, p < CAP) - tl.store(execution_lengths + row, tl.maximum(span_s + span_c, 1)) - tl.store(valid_counts + row, ls + lc) - tl.store(scale_params + 2 + HEADS + row, (ls + lc).to(tl.float32)) - if row == 0: - tl.store(scale_params, bmm1_scale) - tl.store(scale_params + 1, bmm2_scale) - h = tl.arange(0, HEAD_BLOCK) - sink = tl.load(sinks + h, h < HEADS, 0) - tl.store(scale_params + 2 + h, sink, h < HEADS) - - @dataclass class _SparsePlan: - """One planned wrapper plus the caller-owned metadata buffers it binds.""" + """One planned wrapper plus the live-length buffers it reads on every launch. + + The index lists come straight from ``sparse_attn_predict``'s metadata + buffer; only the per-row live lengths are derived here, into buffers whose + addresses stay stable for CUDA graph replay. + """ wrapper: "BatchSparseMLADecodePagedTSWrapper" - metadata: "SparseMLAPreparedMetadata" - default_sinks: torch.Tensor - rows: int - capacity: int + lengths: torch.Tensor + extra_lengths: Optional[torch.Tensor] validated: bool = False @@ -303,6 +243,8 @@ def _unsupported_reason( return "dual-pool sparse indices are required." if indices.ndim != 2 or indices.dtype != torch.int32 or indices.shape[0] != rows: return "sparse indices must be int32 [rows, window + compressed]." + if indices.stride(-1) != 1: + return "sparse index rows must be contiguous." window = int(attn.sparse_attention_config.window_size) extra = self._extra_capacity(metadata) if indices.shape[1] != window + extra: @@ -476,7 +418,6 @@ def _get_plan( ) -> _SparsePlan: from tensorrt_llm._torch.attention.backends.prims_ts.mla_decode import ( BatchSparseMLADecodePagedTSWrapper, - SparseMLAPreparedMetadata, ) attn = self.attn @@ -517,27 +458,12 @@ def _get_plan( else {} ), ) - capacity = int(wrapper._state["capacity"]) - - def empty(shape, dtype=torch.int32): - return torch.empty(shape, device=device, dtype=dtype) - - metadata = SparseMLAPreparedMetadata( - indices=empty((rows, window)), - lengths=empty((rows,)), - extra_indices=empty((rows, extra)) if extra else None, - extra_lengths=empty((rows,)) if extra else None, - routes=empty((1, rows, capacity)), - execution_lengths=empty((1, rows)), - valid_counts=empty((1, rows)), - scale_params=empty((1, 2 + heads + rows), torch.float32), - ) plan = _SparsePlan( wrapper=wrapper, - metadata=metadata, - default_sinks=torch.full((heads,), -torch.inf, device=device, dtype=torch.float32), - rows=rows, - capacity=capacity, + lengths=torch.empty((rows,), device=device, dtype=torch.int32), + extra_lengths=( + torch.empty((rows,), device=device, dtype=torch.int32) if extra else None + ), ) self._plans[key] = plan return plan @@ -572,7 +498,6 @@ def forward( output = forward_args.output groups = int(output.shape[0]) if fused else None - softmax_scale = self._softmax_scale(attn) if is_fp8: # The MLA preprocessing quantizes Q with the KV-cache scale, so both # descales equal the layer's kv_scale_quant_orig (see @@ -587,9 +512,12 @@ def forward( else: q_scale = kv_scale = 1.0 query = q.view(rows, 1, heads, _LATENT_DIM) - output_scale = 1.0 - bmm1_scale = softmax_scale * q_scale * kv_scale - bmm2_scale = output_scale * kv_scale + # The kernel takes the combined scales: the softmax scale carries the Q + # and KV descales, the output scale carries the KV descale (both pools + # share it). They are per-layer constants, so the wrapper's static-scale + # kernel serves them without a device scalar. + softmax_scale = self._softmax_scale(attn) * q_scale * kv_scale + output_scale = kv_scale sinks = self._sinks(forward_args) plan = self._get_plan( @@ -600,49 +528,27 @@ def forward( has_sinks=sinks is not None, fused_groups=groups, ) - prepared = plan.metadata # Storage-row indices straight from sparse_attn_predict: the window # slots address the SWA pool and the compressed slots address the - # compress pool, both relative to their pool base. + # compress pool, both relative to their pool base. The live window + # length is the token's position within the window; the compressed + # length is what the indexer's top-k selected beyond the window. indices = forward_args.sparse_runtime_params.sparse_attn_indices positions = metadata.token_positions_cuda[start:end] - prepared.indices.copy_(indices[:, :window]) - torch.add(positions, 1, out=prepared.lengths) - prepared.lengths.clamp_(max=window) + torch.add(positions, 1, out=plan.lengths) + plan.lengths.clamp_(max=window) + swa_pool = self._pool_view(manager, DeepseekV4AttentionType.SWA, 1) + extra_pool = None + run_kwargs = {} if extra: - prepared.extra_indices.copy_(indices[:, window:]) topk_lens = metadata.sparse_mla_topk_lens[ratio][start:end] - torch.sub(topk_lens, window, out=prepared.extra_lengths) - prepared.extra_lengths.clamp_(min=0, max=extra) - _pack_sparse_mla_metadata[(rows,)]( - prepared.indices, - prepared.extra_indices if extra else prepared.indices, - prepared.lengths, - prepared.extra_lengths if extra else prepared.lengths, - prepared.routes, - prepared.execution_lengths, - prepared.valid_counts, - prepared.scale_params, - plan.default_sinks if sinks is None else sinks, - bmm1_scale, - bmm2_scale, - window, - extra, - rows, - heads, - plan.capacity, - 256, - triton.next_power_of_2(heads), - num_warps=4, - ) + torch.sub(topk_lens, window, out=plan.extra_lengths) + plan.extra_lengths.clamp_(min=0, max=extra) + extra_pool = self._pool_view(manager, DeepseekV4AttentionType.COMPRESS, ratio) + run_kwargs["extra_indices"] = indices[:, window:] + run_kwargs["extra_lengths"] = plan.extra_lengths - swa_pool = self._pool_view(manager, DeepseekV4AttentionType.SWA, 1) - extra_pool = ( - self._pool_view(manager, DeepseekV4AttentionType.COMPRESS, ratio) if extra else None - ) - - run_kwargs = {} if fused: # TRT-LLM's fused epilogue buffers are group-major [G, T, C] with # FP32 scales [G, (H/G)*4, pad4(T)]; the kernel takes their @@ -656,18 +562,16 @@ def forward( else: run_kwargs["out"] = output.view(rows, 1, heads, _LATENT_DIM) - # Validate the metadata contract once per plan while still eager; graph + # Validate the index contract once per plan while still eager; graph # replay must keep validation off because it synchronizes. validate = not plan.validated and not torch.cuda.is_current_stream_capturing() plan.wrapper.run( query, swa_pool, - prepared, + indices[:, :window], + plan.lengths, extra_pool, softmax_scale=softmax_scale, - q_scale=q_scale, - kv_scale=kv_scale, - extra_kv_scale=kv_scale, output_scale=output_scale, sinks=sinks, validate=validate, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/__init__.py b/tensorrt_llm/_torch/attention/backends/prims_ts/__init__.py index e8187118271f..477f52de75cf 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/__init__.py @@ -38,7 +38,6 @@ ) from .mla_decode import ( BatchSparseMLADecodePagedTSWrapper, - SparseMLAPreparedMetadata, batch_sparse_mla_decode_with_paged_kv_cache, get_prims_ts_sparse_mla_decode_workspace_size, BatchMLADecodePagedTSWrapper, @@ -53,7 +52,6 @@ from .sage import SageAttentionConfig, SageAttentionParams __all__ = [ - "SparseMLAPreparedMetadata", "BatchSparseMLADecodePagedTSWrapper", "batch_sparse_mla_decode_with_paged_kv_cache", "get_prims_ts_sparse_mla_decode_workspace_size", diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/compiler.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/compiler.py index c66cf74d37d7..b0510cfc090b 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/compiler.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/compiler.py @@ -17,6 +17,8 @@ from collections.abc import Callable import functools +from .._compile_cache import build_and_load_attention_kernel + import torch from flashinfer.utils import ceil_div @@ -535,7 +537,7 @@ def fake_compact( Float32(1.0), ) - with torch.cuda.device(key.device_index): + def compile_kernel(): return cute.compile( tensor_adapter, *dynamic_args, @@ -549,6 +551,11 @@ def fake_compact( options=_COMPILE_OPTIONS, ) + with torch.cuda.device(key.device_index): + return build_and_load_attention_kernel( + "block_sparse", "main", key, _COMPILE_OPTIONS, compile_kernel, extra=() + ) + @functools.cache def _get_compiled_block_sparse( diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_compile_cache.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_compile_cache.py new file mode 100644 index 000000000000..db8416f07e37 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_compile_cache.py @@ -0,0 +1,108 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. 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. + +"""Persistent compilation shared by the PrimTS attention adapters.""" + +from collections.abc import Callable +import dataclasses +import enum +import functools +import hashlib +from pathlib import Path +from typing import Any + + +def _stable_value(value: object) -> object: + if isinstance(value, enum.Enum): + return (type(value).__module__, type(value).__qualname__, value.name) + if value is None or type(value) in (bool, int, float, str): + return value + if isinstance(value, type): + return (value.__module__, value.__qualname__) + if dataclasses.is_dataclass(value) and not isinstance(value, type): + return ( + type(value).__module__, + type(value).__qualname__, + tuple( + (field.name, _stable_value(getattr(value, field.name))) + for field in dataclasses.fields(value) + if field.compare + ), + ) + if isinstance(value, (tuple, list)): + return tuple(_stable_value(item) for item in value) + raise TypeError(f"no stable cache-key encoding for {type(value).__name__}") + + +def _compile_cache_key(spec: object, options: str, extra: tuple = ()) -> str | None: + """Hash compile identity, excluding only device placement and non-key fields. + + MLA's non-comparing kernel object is represented by its kernel_signature. + Unsupported values bypass persistence rather than risking a stale artifact. + """ + if dataclasses.is_dataclass(spec) and not isinstance(spec, type): + spec = tuple( + (field.name, getattr(spec, field.name)) + for field in dataclasses.fields(spec) + if field.compare and field.name != "device_index" + ) + try: + encoded = repr(_stable_value((spec, options, extra))) + except TypeError: + return None + return hashlib.sha256(encoded.encode()).hexdigest() + + +@functools.cache +def _kernel_source_files() -> tuple[str, ...]: + # Family adapters share masks, reductions, sparse/Sage preparation and + # compatibility helpers. Invalidate them together when any source changes. + root = Path(__file__).resolve().parent + return tuple(sorted(str(path) for path in root.rglob("*.py"))) + + +def build_and_load_attention_kernel( + family: str, + kind: str, + spec: object, + options: str, + compile_fn: Callable[[], Any], + *, + extra: tuple = (), +) -> Any: + """Build or load under the caller's selected CUDA device context.""" + from flashinfer.jit.cute_dsl_core import ( + build_and_load_cute_dsl_kernel, + cute_dsl_cache_disabled, + ) + + if cute_dsl_cache_disabled(): + return compile_fn() + + key = _compile_cache_key(spec, options, extra) + if key is None: + return compile_fn() + import torch + + # A family target such as sm100f can cover physical GPUs with different + # host-selected tuning defaults. The module metadata keys the target; + # the specialization must also distinguish the physical capability. + major, minor = torch.cuda.get_device_capability() + return build_and_load_cute_dsl_kernel( + f"prims_ts_{family}", + f"{kind}_sm{major}{minor}_{key}", + compile_fn, + extra_key_files=_kernel_source_files(), + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_hardware.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_hardware.py new file mode 100644 index 000000000000..f554a507e816 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_hardware.py @@ -0,0 +1,41 @@ +# 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. + +"""Device-local occupancy probes shared by PrimTS attention planners.""" + +import functools + + +@functools.cache +def _query_max_active_clusters(device_index: int, cluster_size: int) -> int: + import torch + import cutlass.utils as utils + from cuda.bindings import driver as cuda_drv + + with torch.cuda.device(device_index): + # HardwareInfo must first see a live primary context. Do this only on + # cache misses; new plan shapes reuse the device's occupancy result. + torch.empty(0, device="cuda") + stream = cuda_drv.CUstream(torch.cuda.current_stream().cuda_stream) + return int( + utils.HardwareInfo(device_index).get_max_active_clusters( + cluster_size, stream + ) + ) + + +def get_max_active_clusters(cluster_size: int) -> int: + import torch + + return _query_max_active_clusters(torch.cuda.current_device(), cluster_size) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_q_token_kv_block_sparse_policy.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_q_token_kv_block_sparse_policy.py index a0ea3425aab7..e1a309e07a21 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/_q_token_kv_block_sparse_policy.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_q_token_kv_block_sparse_policy.py @@ -121,7 +121,17 @@ def fanout(instances: int) -> int: tile_size_q=candidate.tile_size_q, num_insts_kv=instances, splits_kv=splits, - head_dim_per_stage_kv=128 if head_dim == 256 else 0, + # Full FP8 stages preserve the 128-byte swizzled planes while issuing + # PV at N256. BF16 and Swaps retain their existing D128 head bands. + head_dim_per_stage_kv=( + 256 + if head_dim == 256 + and q_dtype_key == "float8_e4m3fn" + and candidate.variant == "keeps_mma_ab" + else 128 + if head_dim == 256 + else 0 + ), o_stages=1 if head_dim == 256 and candidate.variant == "keeps_mma_ab" else 2, use_persistent_scheduler=( splits == 1 diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/context.py b/tensorrt_llm/_torch/attention/backends/prims_ts/context.py index d297bc79c1e0..1ec1080b942f 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/context.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/context.py @@ -43,6 +43,9 @@ import struct from typing import TYPE_CHECKING, Callable, Literal, Optional +from ._compile_cache import build_and_load_attention_kernel +from ._hardware import get_max_active_clusters + import torch from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api @@ -1461,7 +1464,6 @@ def _make_context_scheduler_probe( """Build the common static probe and logical work count for policy.""" import cutlass - import cutlass.utils as utils dtype_map = { torch.float16: cutlass.Float16, @@ -1488,7 +1490,7 @@ def _make_context_scheduler_probe( max_kv_len=max_kv_len, ) with torch.cuda.device(geometry.device_index): - max_active_clusters = int(utils.HardwareInfo().get_max_active_clusters(1)) + max_active_clusters = get_max_active_clusters(1) num_seq_tiles = ( geometry.max_seq_len_q + probe.cfg.cta_tiler[0] - 1 ) // probe.cfg.cta_tiler[0] @@ -1757,7 +1759,6 @@ def _get_compiled_context( import cutlass import cutlass.cute as cute from cuda.bindings import driver as cuda_drv - import cutlass.utils as utils dtype_map = { "float16": cutlass.Float16, @@ -1769,7 +1770,7 @@ def _get_compiled_context( output_dtype = dtype_map[output_dtype_key] has_variable_window = mask_type == "variable_window" with torch.cuda.device(device_index): - max_active_clusters = int(utils.HardwareInfo().get_max_active_clusters(1)) + max_active_clusters = get_max_active_clusters(1) fmha = _make_context_kernel( device_index=device_index, input_qk_dtype=input_qk_dtype, @@ -1916,8 +1917,8 @@ def fake_compact(dtype, shape, assumed_align): # Task objects carry loop-local state through generated control flow, so # select the public staged frontend for this compilation. - with torch.cuda.device(device_index): - compiled = cute.compile[cute.FrontendNext]( + def compile_kernel(): + return cute.compile[cute.FrontendNext]( tensor_adapter, q_fake, k_fake, @@ -1937,6 +1938,16 @@ def fake_compact(dtype, shape, assumed_align): max_seq_len_k, options=_COMPILE_OPTIONS, ) + + with torch.cuda.device(device_index): + compiled = build_and_load_attention_kernel( + "context", + "main", + compile_spec, + _COMPILE_OPTIONS, + compile_kernel, + extra=(max_active_clusters,), + ) policy = ( ("scheduler", scheduler), ("pairing", "head" if head_paired else "query"), @@ -1975,7 +1986,6 @@ def _get_compiled_paged_context( import cutlass import cutlass.cute as cute from cuda.bindings import driver as cuda_drv - import cutlass.utils as utils dtype_map = { "float16": cutlass.Float16, @@ -1986,7 +1996,7 @@ def _get_compiled_paged_context( input_pv_dtype = dtype_map[pv_dtype_key] output_dtype = dtype_map[output_dtype_key] with torch.cuda.device(device_index): - max_active_clusters = int(utils.HardwareInfo().get_max_active_clusters(1)) + max_active_clusters = get_max_active_clusters(1) fmha = _make_context_kernel( device_index=device_index, input_qk_dtype=input_qk_dtype, @@ -2090,8 +2100,8 @@ def fake_compact(dtype, shape, assumed_align): seq_lens_fake = fake_compact(cutlass.Int32, (batch_size,), 4) stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - with torch.cuda.device(device_index): - compiled = cute.compile[cute.FrontendNext]( + def compile_kernel(): + return cute.compile[cute.FrontendNext]( tensor_adapter, q_fake, k_fake, @@ -2108,6 +2118,16 @@ def fake_compact(dtype, shape, assumed_align): max_kv_len, options=_COMPILE_OPTIONS, ) + + with torch.cuda.device(device_index): + compiled = build_and_load_attention_kernel( + "context", + "paged", + compile_spec, + _COMPILE_OPTIONS, + compile_kernel, + extra=(max_active_clusters,), + ) policy = ( ("scheduler", scheduler), ("pairing", "head" if head_paired else "query"), diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py b/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py index 205e875fc390..f91814650c48 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py @@ -16,7 +16,7 @@ """Task-scheduled paged decode with a FlashInfer-style plan/run lifecycle.""" from collections.abc import Callable, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, replace import functools import math import numbers @@ -25,6 +25,8 @@ import torch +from ._compile_cache import build_and_load_attention_kernel + from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api from . import _q_token_kv_block_sparse_policy as _sparse_policy @@ -694,8 +696,6 @@ def _validate_dtype_pair( k_dtype: torch.dtype, v_dtype: torch.dtype, output_dtype: torch.dtype, - *, - allow_fp8_bf16_output: bool = False, ) -> None: _dtype_key(q_dtype) _dtype_key(k_dtype) @@ -723,25 +723,16 @@ def _validate_dtype_pair( q_dtype == torch.float8_e4m3fn and matching_kv and k_dtype in (torch.float8_e4m3fn, torch.uint8) - and output_dtype - in ( - (torch.float16, torch.bfloat16, torch.float8_e4m3fn) - if allow_fp8_bf16_output - else (torch.float16, torch.float8_e4m3fn) - ) + and output_dtype in (torch.float16, torch.bfloat16, torch.float8_e4m3fn) ) ) if not supported: raise NotImplementedError( "attention-ts decode supports FP16->FP16, BF16->BF16, " - + ( - "FP8-E4M3->FP16/BF16, and FP8-E4M3->FP8-E4M3; " - if allow_fp8_bf16_output - else "FP8-E4M3->FP16, and FP8-E4M3->FP8-E4M3; " - ) - + "BF16 Q + FP8 K/V, BF16/FP8-E4M3 Q + NVFP4 K/V, " - + "and BF16 Q/K + FP8 V; got " - + f"Q_{q_dtype}_K_{k_dtype}_V_{v_dtype}->{output_dtype}" + "FP8-E4M3->FP16/BF16/FP8-E4M3; " + "BF16 Q + FP8 K/V, BF16/FP8-E4M3 Q + NVFP4 K/V, " + "and BF16 Q/K + FP8 V; got " + f"Q_{q_dtype}_K_{k_dtype}_V_{v_dtype}->{output_dtype}" ) @@ -1357,7 +1348,7 @@ def fits_one_service_wave(config: "FmhaDecodeConfig", num_q_ctas: int) -> bool: head_dim=head_dim, page_size=page_size, max_kv_len=max_kv_len, - seq_len_q=seq_len_q, + group_size=seq_len_q, q_dtype_key=q_dtype_key, k_dtype_key=k_dtype_key, v_dtype_key=v_dtype_key, @@ -1487,12 +1478,16 @@ def _make_decode_compile_spec( kv_prefix_mode: Literal["dynamic", "planned_full"], kv_lengths_mode: Literal["dynamic", "planned_uniform_max"], direct_q1_spec: Optional[_DirectQ1CompileSpec] = None, + flat_native_kv_tma: bool = False, ) -> _DecodeCompileSpec: """Freeze the resolved topology while leaving batch in runtime tensors.""" + config_items = dict(launch_spec.config.compile_signature()) + config_items["use_flat_native_kv_tma"] = flat_native_kv_tma + config_items["flat_native_kv_num_heads"] = num_kv_heads if flat_native_kv_tma else 1 return _DecodeCompileSpec( device_index=device_index, - config_items=launch_spec.config.compile_signature(), + config_items=tuple(config_items.items()), num_qo_heads=num_qo_heads, num_kv_heads=num_kv_heads, head_dim=head_dim, @@ -1905,8 +1900,8 @@ def fake_compact(dtype, shape, assumed_align): ), ) - with torch.cuda.device(device_index): - compiled_main = cute.compile( + def compile_main(): + return cute.compile( main_tensor_adapter, q_fake, k_fake, @@ -1944,28 +1939,38 @@ def fake_compact(dtype, shape, assumed_align): static_native_uniform_kv, options=_COMPILE_OPTIONS, ) + + def compile_reducer(): + assert reduction_tensor_adapter is not None + return cute.compile( + reduction_tensor_adapter, + out_fake, + seq_lens_fake, + cu_seqlens_q_fake, + partial_o_fake, + partial_stats_fake, + attention_sinks_fake, + Float32(1.0), + Float32(1.0), + direct_q1_fake, + stream_fake, + cfg, + num_qo_heads, + num_kv_heads, + head_dim, + max_kv_len, + static_full_split_prefix, + options=_COMPILE_OPTIONS, + ) + + with torch.cuda.device(device_index): + compiled_main = build_and_load_attention_kernel( + "decode", "main", compile_spec, _COMPILE_OPTIONS, compile_main + ) compiled_reducer = None if cfg.use_separate_reduction_kernel: - assert reduction_tensor_adapter is not None - compiled_reducer = cute.compile( - reduction_tensor_adapter, - out_fake, - seq_lens_fake, - cu_seqlens_q_fake, - partial_o_fake, - partial_stats_fake, - attention_sinks_fake, - Float32(1.0), - Float32(1.0), - direct_q1_fake, - stream_fake, - cfg, - num_qo_heads, - num_kv_heads, - head_dim, - max_kv_len, - static_full_split_prefix, - options=_COMPILE_OPTIONS, + compiled_reducer = build_and_load_attention_kernel( + "decode", "reducer", compile_spec, _COMPILE_OPTIONS, compile_reducer ) return compiled_main, compiled_reducer @@ -2486,6 +2491,7 @@ class PrimsTSBatchDecodePlan: _compiled_main: Callable[..., object] _compiled_reducer: Optional[Callable[..., object]] _direct_q1_inputs: tuple[torch.Tensor, ...] = () + _flat_locator_heads: int = 1 def run( self, @@ -2854,7 +2860,7 @@ def _resolve_q_token_kv_block_sparse_decode_config( head_dim: int, page_size: int, max_kv_len: int, - seq_len_q: int, + group_size: int, q_dtype_key: str, k_dtype_key: str, v_dtype_key: str, @@ -2875,7 +2881,7 @@ def _resolve_q_token_kv_block_sparse_decode_config( heads_q_per_kv = num_qo_heads // num_kv_heads _validate_prims_ts_q_token_kv_block_sparse_group_capacity( - seq_len_q, + group_size, num_qo_heads, num_kv_heads, ) @@ -2905,7 +2911,7 @@ def _resolve_q_token_kv_block_sparse_decode_config( ) launch = _sparse_policy.select_sparse_launch( - group_size=seq_len_q, + group_size=group_size, heads_q_per_kv=heads_q_per_kv, head_dim=head_dim, q_dtype_key=q_dtype_key, @@ -2935,14 +2941,14 @@ def _resolve_q_token_kv_block_sparse_decode_config( "scheduler_warp_idx": 13, "clc_load_warp_idx": 16, "page_offsets_warp_idx": 14, - "page_offsets_num_warps": 2 if seq_len_q == 1 else 1, + "page_offsets_num_warps": 2 if group_size == 1 else 1, "load_warp_idx": 16, # FP8's predicated TMA helper constructs coordinates in every load # lane, so Q1 needs only one producer warpgroup. BF16 distributes page # fragments across two producer warpgroups. Grouped routes keep # their qualified two-warpgroup schedule for both dtypes. "load_num_warps": ( - 4 if seq_len_q == 1 and q_dtype_key == "float8_e4m3fn" else 8 + 4 if group_size == 1 and q_dtype_key == "float8_e4m3fn" else 8 ), } if launch.use_keeps_mma_ab: @@ -2956,7 +2962,7 @@ def _resolve_q_token_kv_block_sparse_decode_config( ), splits_kv=launch.splits_kv, max_splits_kv=launch.splits_kv, - min_loop_iters_per_split=(1 if seq_len_q == 1 else MIN_LOOP_ITERS_PER_SPLIT), + min_loop_iters_per_split=(1 if group_size == 1 else MIN_LOOP_ITERS_PER_SPLIT), ) @@ -3169,7 +3175,7 @@ def validate_q_token_kv_block_sparse_group_size( num_qo_heads: int, num_kv_heads: int, *, - group_size: int, + group_size: int = 1, ) -> int: """Validate a QToken-KvBlock-Sparse-Attention Q1--Q8 group. @@ -3200,7 +3206,8 @@ def validate_q_token_kv_block_sparse_group_size( Number of key/value heads. group_size : int Caller-selected number of query rows per QToken-KvBlock-Sparse-Attention route. Supported values - are Q1--Q8, subject to the TileQ128 head-capacity bound. + are Q1--Q8, subject to the TileQ128 head-capacity bound. Defaults to + one (no grouping). Returns ------- @@ -3222,16 +3229,17 @@ def make_q_token_kv_block_sparse_qo_indptr( query_start_loc_cpu: torch.Tensor, num_query_tokens: int, *, - group_size: int, + group_size: int = 1, device: Optional[Union[int, str, torch.device]] = None, ) -> torch.Tensor: """Build request-safe QToken-KvBlock-Sparse-Attention route offsets. ``query_start_loc_cpu`` contains cumulative offsets for real requests. Every request is chunked independently, so its final route may contain - fewer than ``group_size`` rows. Any inert CUDA-graph padding suffix is - partitioned separately. The returned Int32 ``qo_indptr`` can be passed to - the packed-Q wrapper together with planned ``seq_len_q=group_size``. + fewer than ``group_size`` rows. The default is one (no query grouping). + Any inert CUDA-graph padding suffix is partitioned separately. Pass the + returned Int32 ``qo_indptr`` to a packed-Q wrapper planned with the same + ``group_size``. Parameters ---------- @@ -3243,7 +3251,7 @@ def make_q_token_kv_block_sparse_qo_indptr( Total flattened query-row count, including any inert CUDA-graph padding suffix. group_size : int - Maximum number of rows in each request-safe QToken-KvBlock-Sparse-Attention route. + Maximum rows per request-safe query group, default one (no grouping). device : int, str, torch.device, or None Destination device for the returned offsets. By default, PyTorch's default tensor device is used. @@ -3350,6 +3358,7 @@ def _prepare_prims_ts_batch_decode_plan( page_size: Optional[int], kv_scale_factors: Optional[PagedKVScaleFactors] = None, q_token_kv_block_sparse_page_memberships: Optional[torch.Tensor] = None, + query_major_memberships: bool = False, use_q_token_kv_block_sparse_route: bool = False, use_pdl: bool = False, split_kv: bool = True, @@ -3357,6 +3366,7 @@ def _prepare_prims_ts_batch_decode_plan( direct_q1_max_model_len: Optional[int] = None, direct_q1_sparse_block_size: int = 4, share_pattern_across_kv_heads: bool = True, + allow_head_aware_locators: bool = False, ) -> tuple[PrimsTSBatchDecodePlan, torch.Tensor]: """Validate and freeze one dense-block-table PrimTS launch contract.""" @@ -3432,7 +3442,6 @@ def _prepare_prims_ts_batch_decode_plan( k_cache.dtype, v_cache.dtype, output_dtype, - allow_fp8_bf16_output=use_q_token_kv_block_sparse_route, ) k_sf_cache, v_sf_cache = _normalize_paged_kv_scale_factors( kv_scale_factors, @@ -3466,6 +3475,12 @@ def _prepare_prims_ts_batch_decode_plan( share_pattern_across_kv_heads, ) spec = _resolve_decode_launch_spec(*policy_args) + if query_major_memberships: + if not spec.config.supports_query_major_memberships: + raise ValueError( + "query-major memberships require one G2..G8 Keeps/page-4 CTA" + ) + spec = replace(spec, config=replace(spec.config, query_major_memberships=True)) if spec.config.uses_q_token_kv_block_sparse_page_membership: if q_token_kv_block_sparse_page_memberships is None: raise ValueError( @@ -3525,6 +3540,16 @@ def _prepare_prims_ts_batch_decode_plan( kv_block_size=direct_q1_sparse_block_size, page_capacity=max_num_pages, ) + flat_native_kv_tma = ( + use_q_token_kv_block_sparse_route + and spec.config.uses_staged_one_inst_tmem_p + and (spec.config.use_bf16_qkv or spec.config.use_fp8_qkv) + and spec.config.has_storage_subpages + and (num_kv_heads == 1 or (allow_head_aware_locators and not direct_q1_inputs)) + and k_cache.is_contiguous() + and normalized_cache.v_cache.is_contiguous() + and int(k_cache.shape[0]) * storage_page_size * num_kv_heads < (1 << 31) + ) compile_spec = _make_decode_compile_spec( spec, device_index=device_index, @@ -3542,6 +3567,7 @@ def _prepare_prims_ts_batch_decode_plan( kv_prefix_mode="dynamic", kv_lengths_mode="dynamic", direct_q1_spec=direct_q1_spec, + flat_native_kv_tma=flat_native_kv_tma, ) compiled_main, compiled_reducer = _get_compiled_decode(compile_spec) workspace = _bind_decode_workspace(workspace_buffer, layout) @@ -3572,6 +3598,7 @@ def _prepare_prims_ts_batch_decode_plan( _compiled_main=compiled_main, _compiled_reducer=compiled_reducer, _direct_q1_inputs=direct_q1_inputs, + _flat_locator_heads=num_kv_heads if flat_native_kv_tma else 1, ) return plan, out diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py index 9b23b5dd7d38..0c5a915d0bcf 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 @@ -137,6 +137,8 @@ from collections.abc import Callable from typing import Any, Tuple +import importlib + import cutlass import cutlass.experimental.cuda as cuda import cutlass.cute as cute @@ -2821,6 +2823,9 @@ def __init__( ) validate_head_paired_head_ratio(head_paired=head_paired, h_r=h_r) if use_paged_kv: + # Import before task bodies enter staged GPU control flow. The + # helper module imports FmhaConfig, so it cannot load in resources. + importlib.import_module(f"{__package__}.helpers_paged") if num_tokens_per_page not in _SUPPORTED_CONTEXT_PAGE_SIZES: raise ValueError( "paged context requires num_tokens_per_page in " 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 e6cdddf8d5cb..af69ee7d51bc 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 @@ -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"); diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_inspect.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_inspect.py index da1e558acc79..78648bd18861 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_inspect.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_inspect.py @@ -28,9 +28,12 @@ import functools from collections.abc import Callable +from dataclasses import dataclass import cutlass import cutlass.cute as cute +from ..._compile_cache import build_and_load_attention_kernel + import torch from cuda.bindings import driver as cuda_drv @@ -54,6 +57,29 @@ _ERROR_INVALID_PHYSICAL_PAGE_ID = 5 +@dataclass(frozen=True) +class _BlockSparseInspectCompileSpec: + """Static BSR geometry shared by equivalent inspection specializations.""" + + batch_size: int + num_pattern_heads: int + num_q_block_rows: int + num_kv_blocks: int + + +@dataclass(frozen=True) +class _PagedBlockSparseInspectCompileSpec: + """Static BSR geometry, paging divisors, and request-validation bounds.""" + + batch_size: int + num_pattern_heads: int + num_q_block_rows: int + minimum_seq_len_kv: int + max_seq_len_kv: int + kv_block_size: int + page_size: int + + @cute.jit def _validate_bsr_row_lane( block_indices: cute.Tensor, @@ -464,12 +490,19 @@ def compile_block_sparse_inspection( ) -> Callable[..., None]: """Compile one geometry specialization while keeping ``indices[nnz]`` dynamic. - Tensor ranks, dtypes, compact strides, and all attention geometry are part - of the specialization. Only the flat ``block_indices`` extent is symbolic; - tensor contents may of course vary between calls to the cached function. + The artifact key uses rounded Q/KV block counts: exact lengths and block + sizes that produce the same BSR tensor shape and index bounds share code. + ``block_indices[nnz]`` and tensor contents remain runtime inputs. """ num_q_block_rows = (seq_len_q + q_block_size - 1) // q_block_size + num_kv_blocks = (seq_len_kv + kv_block_size - 1) // kv_block_size + compile_spec = _BlockSparseInspectCompileSpec( + batch_size=batch_size, + num_pattern_heads=num_pattern_heads, + num_q_block_rows=num_q_block_rows, + num_kv_blocks=num_kv_blocks, + ) logical_nnz = cute.sym_int() stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) kernel = _InspectBlockSparseBsr( @@ -480,7 +513,8 @@ def compile_block_sparse_inspection( q_block_size=q_block_size, kv_block_size=kv_block_size, ) - with torch.cuda.device(device_index): + + def compile_kernel(): return cute.compile( kernel, _fake_compact( @@ -495,6 +529,16 @@ def compile_block_sparse_inspection( options=_COMPILE_OPTIONS, ) + with torch.cuda.device(device_index): + return build_and_load_attention_kernel( + "block_sparse", + "inspect", + compile_spec, + _COMPILE_OPTIONS, + compile_kernel, + extra=(), + ) + @functools.cache def compile_paged_block_sparse_metadata_inspection( @@ -509,9 +553,23 @@ def compile_paged_block_sparse_metadata_inspection( kv_block_size: int, page_size: int, ) -> Callable[..., None]: - """Compile one paged metadata entry that launches request then live-BSR.""" + """Compile request and live-BSR validation keyed by its static geometry. + + Exact Q lengths sharing a Q-block count reuse one artifact. Page/block + divisors and inclusive length limits stay in the key because they become + constants in the generated validation code; table widths remain symbolic. + """ num_q_block_rows = (seq_len_q + q_block_size - 1) // q_block_size + compile_spec = _PagedBlockSparseInspectCompileSpec( + batch_size=batch_size, + num_pattern_heads=num_pattern_heads, + num_q_block_rows=num_q_block_rows, + minimum_seq_len_kv=minimum_seq_len_kv, + max_seq_len_kv=max_seq_len_kv, + kv_block_size=kv_block_size, + page_size=page_size, + ) logical_nnz = cute.sym_int() runtime_page_columns = cute.sym_int() runtime_page_row_stride = cute.sym_int64(divisibility=1) @@ -536,7 +594,7 @@ def compile_paged_block_sparse_metadata_inspection( inspect_bsr=inspect_bsr, ) - with torch.cuda.device(device_index): + def compile_kernel(): return cute.compile( inspect_paged_block_sparse_metadata, _fake_compact( @@ -559,6 +617,16 @@ def compile_paged_block_sparse_metadata_inspection( options=_COMPILE_OPTIONS, ) + with torch.cuda.device(device_index): + return build_and_load_attention_kernel( + "block_sparse", + "inspect_paged", + compile_spec, + _COMPILE_OPTIONS, + compile_kernel, + extra=(), + ) + __all__ = [ "compile_block_sparse_inspection", 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 af54b06004ce..42ba3d0e971d 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 @@ -61,6 +61,7 @@ MMA_K_STEP_BYTES, PACKED_REGISTER_BYTES, PARALLEL_REDUCTION_BYTES_PER_SLICE, + PARALLEL_REDUCTION_SLOT_LANES, PARALLEL_REDUCTION_THREADS_PER_CTA, PARTIAL_O_ELEMENT_BYTES, PARTIAL_STATS_VALUES_PER_ROW, @@ -124,6 +125,7 @@ } _GROUPED_KEEPS_STATIC_ONLY_PROFILES = { (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, Float16, 128, 0, 2, 2), + (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, BFloat16, 128, 0, 2, 2), (BFloat16, BFloat16, BFloat16, BFloat16, 64, 0, 2, 2), (Float16, Float16, Float16, Float16, 256, 128, 1, 1), } @@ -152,7 +154,8 @@ KV_TILE_256_CORRECTION_TASK_REGISTERS = 152 # Sparse route unions reuse the dense Keeps staging: D64/D128 use two -# unstaged K/V instances; D256 uses one instance with D128 head bands. +# unstaged K/V instances; D256 uses one instance with D128 bands or a full +# D256 FP8 stage assembled from 128-byte swizzled chunks. _Q_TOKEN_KV_BLOCK_SPARSE_GROUPED_KEEPS_PROFILES = { ( dtype, @@ -172,6 +175,10 @@ (Float8E4M3FN, BFloat16), ) } +_Q_TOKEN_KV_BLOCK_SPARSE_GROUPED_KEEPS_PROFILES.update( + (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, output, 256, 256, 1, 1) + for output in (Float16, BFloat16) +) _KV_TILE_256_PHYSICAL_DEFAULTS: Mapping[str, ConfigValue] = { "tmem_s_cols": 128, @@ -221,6 +228,7 @@ def _dtype_bytes(dtype: type) -> int: (BFloat16, BFloat16, BFloat16), (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN), (Float8E4M3FN, Float8E4M3FN, Float16), + (Float8E4M3FN, Float8E4M3FN, BFloat16), } @@ -230,10 +238,29 @@ def _dtype_bytes(dtype: type) -> int: # families rather than shape exceptions: batch size, KV length, tile choice, # and legal GMEM split fanout remain unrestricted by this declaration. _GROUPED_KEEPS_PAGED_FP8_PROFILES = { + ( + Float8E4M3FN, + Float8E4M3FN, + Float8E4M3FN, + output_dtype, + head_dim, + 128 if head_dim == 256 else 0, + 1 if head_dim == 256 else 2, + 1 if head_dim == 256 else 2, + ) + for head_dim in (64, 128, 256) + for output_dtype in (Float16, BFloat16, Float8E4M3FN) +} + +# Automatic selection retains the established cost-model recipes. A working +# output epilogue alone does not qualify another MMA layout for that model; +# explicit kernel configurations can use the wider capability matrix above. +_AUTO_GROUPED_KEEPS_PAGED_FP8_PROFILES = { (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, 64, 0, 2, 2), (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, 128, 0, 2, 2), (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, 256, 128, 1, 1), (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, Float16, 256, 128, 1, 1), + (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, BFloat16, 256, 128, 1, 1), } @@ -607,6 +634,9 @@ class FmhaDecodeConfig: # sparse routes consume query membership from a separate packed-word table; # Q1 uses the same scattered page route without membership masking. use_q_token_kv_block_sparse_route: bool = False + # Private prepared-plan layout: eight query words per 32 page-4 IDs. + # Legacy public membership tensors keep the default page-byte encoding. + query_major_memberships: bool = False # None preserves each sparse family's default in raw/static configs too. share_pattern_across_kv_heads: bool | None = None # Allow the attention grid to acquire a programmatic launch dependency. @@ -834,6 +864,19 @@ def uses_q_token_kv_block_sparse_page_route(self) -> bool: and self.grouped_q_fits_tile ) + @property + def supports_query_major_memberships(self) -> bool: + """Whether this CTA can read eight-slot/page-4 query masks for G2..G8.""" + return ( + self.uses_q_token_kv_block_sparse_page_membership + and self.use_keeps_mma_ab + and 2 <= self.max_seq_len_q <= 8 + and self.has_single_q_cta + and self.num_tokens_per_page == 4 + and self.tile_size_kv == 128 + and self.tile_size_q in (64, 128) + ) + @property def uses_held_encoded_locator_window(self) -> bool: """Whether encoded paging keeps the complete CTA route in SMEM.""" @@ -872,14 +915,23 @@ def uses_task_register_reallocation(self) -> bool: @property def softmax_task_num_registers(self) -> int | None: + """Softmax's preferred grant, capped by the CTA's initial register pool. + + Dense and sparse Keeps kernels share this rule. Full 128-score rows + can use 232 registers only when the pool permits; active producer and + correction grants are unchanged. KV256 and transformed KV keep their + separate budgets. + """ if not self.uses_task_register_reallocation: return None if self.use_transform_kv: preferred = 168 if self.headdim == 64 and self.num_insts_kv != 1 else 184 else: - preferred = ( - KV_TILE_256_SOFTMAX_TASK_REGISTERS if self.tile_size_kv == 256 else 184 - ) + # Full Q128 rows keep 128 scores live through masking and P + # conversion. Give them headroom without taking registers from TMA. + preferred = 232 if self.softmax_score_fragment_regs == 128 else 184 + if self.tile_size_kv == 256: + preferred = KV_TILE_256_SOFTMAX_TASK_REGISTERS # setmaxnreg redistributes the CTA's initial register pool, not the # whole SM register file. Complete 8-register/thread quanta leave # unused SM registers when a CTA has more than 16 warps. @@ -903,6 +955,14 @@ def softmax_task_num_registers(self) -> int | None: transform_regs = self.transform_kv_task_num_registers assert transform_regs is not None available -= self.transform_kv_num_warps * (transform_regs - producer_regs) + # Only whole padding warpgroups may donate independently. Padding + # beside active tasks must execute the same setmaxnreg as its peers. + for wg_idx in range(MAX_WARP_GROUPS): + padding_warps = getattr(self, f"wg{wg_idx}_padding_num_warps") + if padding_warps == 4: + padding_regs = self.padding_task_num_registers(wg_idx * 4, 4) + assert padding_regs is not None + available += padding_warps * (producer_regs - padding_regs) return min(preferred, available // (softmax_warps * 8) * 8) @property @@ -931,6 +991,22 @@ def transform_kv_task_num_registers(self) -> int | None: return 88 return 56 + def padding_task_num_registers(self, warp_idx: int, num_warps: int) -> int | None: + """Release unused registers only from complete padding warpgroups.""" + if not self.uses_task_register_reallocation: + return None + if ( + not self.use_transform_kv + and self.tile_size_kv != 256 + and self.softmax_score_fragment_regs == 128 + and warp_idx % 4 == 0 + and num_warps == 4 + ): + # setmaxnreg's minimum legal grant; no attention work runs here. + # Smaller score fragments and KV transforms keep their budgets. + return 24 + return self.mma_load_task_num_registers + # ------------------------------------------------------------------ # SMEM allocation alignment # ------------------------------------------------------------------ @@ -952,6 +1028,13 @@ def q_tokens_per_cta(self) -> int: self.groups_tokens_heads_q, ) + @property + def q_tma_tokens_per_cta(self) -> int: + """Q tokens copied by a box, excluding spare sparse-group slots.""" + if self.use_variable_seqlens_q and self.uses_q_token_kv_block_sparse_page_route: + return min(self.max_seq_len_q, self.q_tokens_per_cta) + return self.q_tokens_per_cta + @property def num_q_ctas(self) -> int: """Return Q-CTA groups per ``(batch, KV head)`` launch tile. @@ -979,15 +1062,13 @@ def has_single_q_cta(self) -> bool: @property def q_tma_rows_per_cta(self) -> int: """Q rows represented by the configured tensor-map box.""" - return _q_tma_rows_per_cta( - self.tile_size_q, - self.heads_q_per_kv, - self.groups_tokens_heads_q, - ) + if self.groups_tokens_heads_q: + return self.q_tma_tokens_per_cta * self.heads_q_per_kv + return self.tile_size_q @property def q_manual_padding_rows(self) -> int: - """Grouped structural Q rows completed without a TMA instruction.""" + """Grouped Q rows whose transaction bytes are completed without TMA.""" if self.groups_tokens_heads_q: return self.tile_size_q - self.q_tma_rows_per_cta return 0 @@ -1076,6 +1157,23 @@ def num_head_dim_stages_kv(self) -> int: """Number of K/V head-dim stages needed to cover headdim.""" return (self.headdim + self.head_dim_kv_stage - 1) // self.head_dim_kv_stage + @property + def uses_2d_flat_kv_tma(self) -> bool: + """Whether warp-owned issuing can omit the flat map's singleton axes.""" + return ( + self.use_flat_native_kv_tma + and (self.tile_size_kv // self.num_tokens_per_page) % self.load_num_warps + == 0 + and ( + self.k_dtype == self.v_dtype == BFloat16 + or ( + self.use_fp8_qkv + and self.use_persistent_scheduler + and self.head_dim_kv_stage // 128 > 1 + ) + ) + ) + @property def tmem_o_cols_per_head_dim_stage(self) -> int: """TMEM O columns owned by one head-dim stage.""" @@ -1294,13 +1392,7 @@ def supports_reduction_dtypes(self) -> bool: and self.out_dtype in (Float16, BFloat16) ) or ( self.q_dtype == Float8E4M3FN - and ( - self.out_dtype in (Float16, Float8E4M3FN) - or ( - self.uses_q_token_kv_block_sparse_page_route - and self.out_dtype == BFloat16 - ) - ) + and self.out_dtype in (Float16, BFloat16, Float8E4M3FN) ) @property @@ -1354,6 +1446,11 @@ def use_fp8_kv(self) -> bool: # page-four table: # locator = physical_page * subpages_per_storage_page + subpage. storage_tokens_per_page: int = 0 + # Prepared compact caches can view physical pages and heads as tensor rows. + # The paired metadata plan selects the corresponding locator head span. + use_flat_native_kv_tma: bool = False + # Prepared locators count the complete physical head span between pages. + flat_native_kv_num_heads: int = 1 # Maximum number of pages per (batch, head_kv) — sizes the page index # table stride. max_num_pages_per_seq_kv: int = 1 @@ -1461,7 +1558,7 @@ def parallel_reduction_cluster_size(self) -> int: return 1 return { 8: 1, - 16: 2, + 16: 1, 32: 4, 64: 8, 128: 16, @@ -1477,12 +1574,23 @@ def parallel_reduction_splits_per_cta(self) -> int: // self.parallel_reduction_cluster_size ) + @property + def parallel_reduction_slot_lanes(self) -> int: + """Share a single-CTA fragment's split slots across adjacent lanes.""" + if ( + self.use_compact_parallel_reduction + or self.parallel_reduction_cluster_size != 1 + or self.parallel_reduction_splits_per_cta % PARALLEL_REDUCTION_SLOT_LANES + ): + return 1 + return PARALLEL_REDUCTION_SLOT_LANES + @property def parallel_reduction_threads_per_cta(self) -> int: """Return the thread count selected by the reducer schedule.""" if self.use_compact_parallel_reduction: return REDUCTION_THREADS_PER_CTA - return PARALLEL_REDUCTION_THREADS_PER_CTA + return PARALLEL_REDUCTION_THREADS_PER_CTA * self.parallel_reduction_slot_lanes @property def parallel_reduction_bytes_per_slice(self) -> int: @@ -2197,6 +2305,21 @@ def _kv256_dtypes_qualified(self) -> bool: return self._uses_contiguous_grouped_keeps_recipe return (self.q_dtype, self.k_dtype, self.out_dtype) in _KV256_DTYPE_RECIPES + @property + def supports_auto_grouped_keeps(self) -> bool: + """Whether grouped Keeps is also qualified for automatic selection.""" + if ( + self.use_paged_kv + and self.tile_size_kv == 128 + and self.q_dtype == self.k_dtype == self.v_dtype == Float8E4M3FN + ): + return ( + self._grouped_keeps_profile_key + in _AUTO_GROUPED_KEEPS_PAGED_FP8_PROFILES + and self.supports_grouped_keeps + ) + return self.supports_grouped_keeps + @property def uses_guarded_fixed_q1_grouped_keeps(self) -> bool: """Whether inactive grouped rows are excluded solely at output stores. @@ -2210,6 +2333,7 @@ def uses_guarded_fixed_q1_grouped_keeps(self) -> bool: profile = self._grouped_keeps_profile_key return ( self.use_keeps_mma_ab + and self.use_paged_kv and self.groups_tokens_heads_q and self.max_seq_len_q == 1 and not self.use_variable_seqlens_q @@ -2236,6 +2360,7 @@ def uses_guarded_grouped_keeps_output_rows(self) -> bool: profile = self._grouped_keeps_profile_key return ( self.use_keeps_mma_ab + and self.use_paged_kv and self.groups_tokens_heads_q and not self.use_variable_seqlens_q and profile in _GROUPED_KEEPS_PAGED_FP8_PROFILES @@ -2368,7 +2493,14 @@ def uses_staged_one_inst_tmem_p(self) -> bool: return ( self.use_keeps_mma_ab and self.headdim == 256 - and self.head_dim_per_stage_kv == 128 + and ( + self.head_dim_per_stage_kv == 128 + or ( + self.head_dim_per_stage_kv == 256 + and self.use_fp8_qkv + and self.uses_scattered_page_route + ) + ) and self.num_insts_kv == 1 and self.o_stages == 1 ) @@ -2873,7 +3005,7 @@ def supports_grouped_keeps(self) -> bool: if self.use_sage_attention: return False - if profile in _GROUPED_KEEPS_PAGED_FP8_PROFILES: + if self.use_paged_kv and profile in _GROUPED_KEEPS_PAGED_FP8_PROFILES: fixed_q1 = ( self.max_seq_len_q == 1 and 1 <= self.heads_q_per_kv <= self.tile_size_q ) @@ -3826,28 +3958,19 @@ def _select_auto_launch_mode( def get_max_active_clusters_for_cluster_size(cluster_size: int) -> int: - """Query cluster occupancy after establishing CUDA's primary context. + """Query device occupancy, preserving the CPU-only policy fallback.""" + from ..._hardware import get_max_active_clusters - ``HardwareInfo`` compiles and retains a tiny occupancy-query module on its - first call. If that first call precedes CUDA context creation, retrying can - reuse an invalid module handle. Establish the context before constructing - ``HardwareInfo`` so a fresh Python process is reliable as well. - """ - cuda_context_ready = False try: import torch if torch.cuda.is_available(): - torch.cuda.init() - torch.empty(0, device="cuda") - cuda_context_ready = True - except (ImportError, RuntimeError): + return get_max_active_clusters(cluster_size) + except ImportError: pass try: return utils.HardwareInfo().get_max_active_clusters(cluster_size) except RuntimeError: - if cuda_context_ready: - raise return max(FALLBACK_SM_COUNT_B200 // max(cluster_size, 1), 1) @@ -4035,6 +4158,8 @@ def _try_apply_default_wide_keeps_config( probe, explicit_fields | {"tile_size_q"}, ) + if probe.groups_tokens_heads_q and not probe.supports_auto_grouped_keeps: + continue _validate_profile_support( cfg=probe, seq_len_q=seq_len_q, @@ -4164,6 +4289,8 @@ def _resolve_grouped_q_launch_candidates( probe, explicit_fields | {"tile_size_q"}, ) + if probe.use_keeps_mma_ab and not probe.supports_auto_grouped_keeps: + return () cost_num_insts_kv = probe.num_insts_kv _validate_profile_support( cfg=probe, @@ -4744,9 +4871,9 @@ def _validate_profile_support( heads_q_per_kv=heads_q_per_kv, groups_tokens_heads_q=True, ) - # Preserve the legacy fixed FP8 Q8/D128 profile. The D64/D256 extensions - # are deliberately limited to grouped page-32 decode; broader dtypes, - # head ratios, and variable-Q layouts remain unqualified. + # Fixed FP8 Q8 producers and reducers support both 16-bit output types. + # D64/D256 retain their grouped page-32 geometry; head ratios and + # variable-Q layouts remain constrained by the same resource profile. qualified_fp8_q8_separate_reduction_supported = ( not use_keeps_mma_ab and cfg.use_split_kv @@ -4756,14 +4883,14 @@ def _validate_profile_support( and cfg.k_dtype == Float8E4M3FN and cfg.v_dtype == Float8E4M3FN and ( - (headdim == 128 and cfg.out_dtype == Float8E4M3FN) + headdim == 128 or ( use_groups_tokens_heads_q and cfg.use_paged_kv and cfg.num_tokens_per_page == 32 and ( - (headdim == 64 and cfg.out_dtype == Float8E4M3FN) - or (headdim == 256 and cfg.out_dtype == Float16) + headdim == 64 + or (headdim == 256 and cfg.out_dtype in (Float16, BFloat16)) ) ) ) @@ -4781,7 +4908,7 @@ def _validate_profile_support( padded_splits = cfg.parallel_reduction_padded_splits default_cluster_size = { 8: 1, - 16: 2, + 16: 1, 32: 4, 64: 8, 128: 16, @@ -4794,7 +4921,7 @@ def _validate_profile_support( clustered_topology_supported = ( not cfg.use_compact_parallel_reduction and default_cluster_size == cluster_size - and splits_per_cta in (2, 4, 8) + and splits_per_cta in (2, 4, 8, 16) and cluster_size * splits_per_cta == padded_splits ) if not ( @@ -4913,22 +5040,27 @@ def _validate_profile_support( effective_head_dim_stage = cfg.head_dim_per_stage_kv effective_num_insts_kv = cfg.num_insts_kv effective_o_stages = cfg.o_stages + supported_head_stage = effective_head_dim_stage == 128 or ( + effective_head_dim_stage == 256 + and cfg.use_fp8_qkv + and cfg.uses_q_token_kv_block_sparse_page_route + ) if effective_num_insts_kv == 1 and ( - headdim != 256 or effective_head_dim_stage != 128 or effective_o_stages != 1 + headdim != 256 or not supported_head_stage or effective_o_stages != 1 ): raise ValueError( "one-instance KeepsMmaAb is enabled only for the staged " - "headDim=256 profile with head_dim_per_stage_kv=128 and " + "headDim=256 profile with D128 bands (or D256 for sparse FP8) and " "o_stages=1" ) if headdim == 256 and ( - effective_head_dim_stage != 128 + not supported_head_stage or effective_num_insts_kv != 1 or effective_o_stages != 1 ): raise ValueError( "fmha_decode keepsMmaAb headDim=256 requires " - "head_dim_per_stage_kv=128, num_insts_kv=1, and o_stages=1" + "D128 bands (or D256 for sparse FP8), num_insts_kv=1, and o_stages=1" ) if headdim != 256 and effective_head_dim_stage != 0: raise ValueError( @@ -4939,20 +5071,6 @@ def _validate_profile_support( raise ValueError( "fmha_decode keepsMmaAb requires numHeadsQPerKv == tile_size_q" ) - fp8_q_token_kv_block_sparse_bf16_output = ( - cfg.out_dtype == BFloat16 - and use_groups_tokens_heads_q - and is_q_token_kv_block_sparse_grouped_keeps - ) - if ( - cfg.q_dtype == Float8E4M3FN - and cfg.out_dtype not in (Float16, Float8E4M3FN) - and not fp8_q_token_kv_block_sparse_bf16_output - and not cfg.use_sage_attention - ): - raise ValueError( - "fmha_decode keepsMmaAb fp8 qkv path supports fp16 or fp8 output" - ) use_split_kv = split_kv_mode != "disabled" or cfg.use_split_kv if use_split_kv: if cfg.q_dtype not in (Float16, BFloat16, Float8E4M3FN): @@ -4987,7 +5105,7 @@ def _validate_profile_support( and cfg.uses_q_token_kv_block_sparse_page_route ) ) - and effective_head_dim_stage == 128 + and supported_head_stage and effective_num_insts_kv == 1 and effective_o_stages == 1 ) @@ -5002,7 +5120,7 @@ def _validate_profile_support( and cfg.q_dtype == Float8E4M3FN and cfg.k_dtype == Float8E4M3FN and cfg.v_dtype == Float8E4M3FN - and cfg.out_dtype == (Float16 if headdim == 256 else Float8E4M3FN) + and (headdim == 64 or cfg.out_dtype in (Float16, BFloat16)) and cfg.splits_kv == cfg.max_splits_kv and 2 <= cfg.max_splits_kv <= 128 ) @@ -5049,7 +5167,7 @@ def _validate_profile_support( "{64,128,256}, TileSizeQ in {8,16,32}, " "either ungrouped single-token or complete-token grouped Q, " "at least two split CTAs, " - "and fp16/bf16 or fp8 qkv with fp16/fp8 output" + "and fp16/bf16 or fp8 qkv with fp16/bf16/fp8 output" ) cluster_reason = cluster_smem_reduction_unsupported_reason( max_splits_kv=cfg.max_splits_kv, @@ -5078,9 +5196,9 @@ def _validate_profile_support( raise ValueError( "separate reduction SwapsMmaAb profiles require fixed SQ=1, " "fixed grouped Q, or packed variable Q; tile_size_q in {16,32}, " - "QToken-KvBlock-Sparse-Attention TileQ8, or a fixed FP8 Q8/HqPerKv8 profile (legacy D128 with FP8 output, " - "plus grouped paged-KV/page-32 D64 with FP8 output or D256 with " - "FP16 output); valid reduction dtypes and static split-KV are required" + "QToken-KvBlock-Sparse-Attention TileQ8, or a fixed FP8 Q8/HqPerKv8 profile " + "(D128, plus grouped paged-KV/page-32 D64 or D256 with " + "FP16/BF16 output); valid reduction dtypes and static split-KV are required" ) if tile_size_q and tile_size_q not in (8, 16, 32): raise ValueError("fmha_decode SwapsMmaAb supports tile_size_q in {8,16,32}") diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_constants.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_constants.py index 909de7947f81..f120a13f02cb 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_constants.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_constants.py @@ -186,14 +186,21 @@ # Clustered standalone reducer shape. One 128-thread CTA covers a contiguous # 2 KiB partial-O slice with one 16-byte vector per thread. Each cluster rank -# owns a compile-time 2, 4, or 8 split slots. Loads are batched in groups of at -# most four; padded split slots remain neutral and never form GMEM pointers. +# owns a compile-time 2, 4, or 8 split slots. Padded and unpublished slots +# remain neutral and never form GMEM pointers. PARALLEL_REDUCTION_THREADS_PER_CTA = 128 PARALLEL_REDUCTION_BYTES_PER_SLICE = ( PARALLEL_REDUCTION_THREADS_PER_CTA * REDUCTION_BYTES_PER_THREAD ) -PARALLEL_REDUCTION_LOAD_BATCH = 4 +# Batch independent published-slot loads before the serial fold. +PARALLEL_REDUCTION_LOAD_BATCH = 16 PARALLEL_REDUCTION_FINAL_REDUCERS = 4 +# Single-CTA reducer profiles spread one output fragment's split slots +# over this many adjacent lanes and merge them with warp shuffles. The fold is +# a serial chain of dependent FP32 instructions at very low occupancy, so more +# lanes with fewer slots shorten that per-lane dependency chain. Must divide +# the slots owned by each CTA. +PARALLEL_REDUCTION_SLOT_LANES = 4 # Each reduction thread produces an 8-element O vector backed by four packed # 16-bit registers. 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 2292eb6544bb..7a2235d8bc28 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 @@ -134,10 +134,9 @@ ) _PERSISTENT_SCHEDULE_TOKEN_STAGES = 2 -# Response slots of the CLC work-queue fetch pipeline. Every role consumes -# the queue one tile behind the scheduler warp, so a second slot adds no -# lookahead, only a loop-carried stage index that spills to local memory. -_WORK_QUEUE_STAGES = 1 +# Double-buffer CLC responses independently of the schedule-token throttle. +# A single response slot triggers register spills in grouped sparse work loops. +_WORK_QUEUE_STAGES = 2 def _block_sparse_bshd_tma_strides( @@ -404,7 +403,7 @@ def _build_decode_gen_schedule( if cfg.use_keeps_mma_ab and cfg.num_insts_kv == 1 and not cfg.uses_tmem_p: raise ValueError( "one-instance KeepsMmaAb is enabled only for the staged headDim=256 " - "profile with head_dim_per_stage_kv=128 and o_stages=1" + "profile with D128 bands (or D256 for sparse FP8) and o_stages=1" ) if use_native_paged_kv and not cfg.use_paged_kv: raise ValueError("native paged-KV ABI requires cfg.use_paged_kv=True") @@ -534,7 +533,18 @@ def _build_decode_gen_schedule( cfg.k_dtype_bytes != cfg.v_dtype_bytes and cfg.tile_size_kv != 256 ) one_inst_tmem_stages = 2 if use_one_inst_qkv else 1 - one_inst_kv_stages = cfg.num_head_dim_stages_kv if use_one_inst_qkv else 1 + # Keep complete head-dimension rounds in each balanced K/V ring. FP8's + # smaller stage tiles can prefetch farther without increasing the budget. + one_inst_kv_stages = ( + max( + cfg.num_head_dim_stages_kv, + cfg.kv_stages + // (2 * cfg.num_head_dim_stages_kv) + * cfg.num_head_dim_stages_kv, + ) + if use_one_inst_qkv + else 1 + ) use_distributed_split_kv_stages = not use_one_inst_qkv if cfg.tile_size_q == 128 and use_distributed_split_kv_stages: # Q128's four instruction-local K0/K1/V0/V1 rings need equal depth. @@ -836,15 +846,14 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: advance_on_wait=True, ) - # Two-instance Keeps keeps stats outside S and orders each same-instance PV - # before the next QK, so it needs no stats-done credit. - # The staged one-instance path needs an overwrite-credit gate across its - # double-buffered S/P overlay: correction returns the stage credit before - # MMA can reissue QK into those columns. + # The single MMA issuer orders PV before QK reuses the same S/P slot. + # PTX's same-warp mma(A-read) -> mma(D-write) pipeline protects that alias. + # Separate SMEM stats have their own softmax-local handoff, just as in + # two-instance Keeps; correction need not additionally gate S reuse. stats_done0_cfg = None stats_done1_cfg = None resource_dependency_graph: dict[MemoryResource, list[MemoryResource]] - if use_one_inst_qkv: + if use_one_inst_qkv and (not cfg.keeps_stats_via_smem or cfg.mma_num_warps != 1): stats_done0_cfg = PipelineConfig.create_async_async_pipeline_cfg( num_stages=one_inst_tmem_stages, producer_group=mma_grp, @@ -3375,12 +3384,6 @@ def fmha_decode_launch( if cutlass.const_expr(use_native_paged_kv): total_pages = num_physical_kv_pages storage_tokens_per_page = Int32(cfg.effective_storage_tokens_per_page) - kv_shape = ( - kv_tma_d, - storage_tokens_per_page, - h_k, - total_pages, - ) k_tma_page_stride = k_page_stride k_tma_head_stride = k_head_stride k_tma_token_stride = k_token_stride @@ -3396,24 +3399,60 @@ def fmha_decode_launch( v_tma_page_stride = v_page_stride * Int64(2) v_tma_head_stride = v_head_stride * Int64(2) v_tma_token_stride = v_token_stride * Int64(2) - k_layout = cute.make_layout( - kv_shape, - stride=( - 1, - k_tma_token_stride, - k_tma_head_stride, - k_tma_page_stride, - ), - ) - v_layout = cute.make_layout( - kv_shape, - stride=( - 1, - v_tma_token_stride, - v_tma_head_stride, - v_tma_page_stride, - ), - ) + if cutlass.const_expr(cfg.use_flat_native_kv_tma): + # Reinterpret compact HND rows, without copying/repacking KV + # or requiring contiguous sparse selections. Planning checked + # signed-32-bit row bounds; other bindings use the general map. + # For S tokens/page, F tokens/fragment and H heads, metadata + # emits A = page*H*(S/F) + subpage. Attention addresses row + # (A + head*(S/F))*F, avoiding per-copy page/subpage divmod. + # Bytes and TMA count are unchanged; rank-2 issuing below is a + # separate optimization of this descriptor-only flat view. + flat_tokens = ( + num_physical_kv_pages + * Int64(storage_tokens_per_page) + * Int64(cfg.flat_native_kv_num_heads) + ) + kv_shape = (kv_tma_d, flat_tokens, Int32(1), Int32(1)) + k_tma_head_stride = k_tma_token_stride * flat_tokens + k_tma_page_stride = k_tma_head_stride + v_tma_head_stride = v_tma_token_stride * flat_tokens + v_tma_page_stride = v_tma_head_stride + else: + kv_shape = ( + kv_tma_d, + storage_tokens_per_page, + h_k, + total_pages, + ) + if cutlass.const_expr(cfg.uses_2d_flat_kv_tma): + # The prepared flat view has no head/page coordinates. Keep + # the same physical rows and SW128 tile without repacking KV. + k_layout = cute.make_layout( + (kv_shape[0], kv_shape[1]), stride=(1, k_tma_token_stride) + ) + v_layout = cute.make_layout( + (kv_shape[0], kv_shape[1]), stride=(1, v_tma_token_stride) + ) + else: + k_layout = cute.make_layout( + kv_shape, + stride=( + 1, + k_tma_token_stride, + k_tma_head_stride, + k_tma_page_stride, + ), + ) + v_layout = cute.make_layout( + kv_shape, + stride=( + 1, + v_tma_token_stride, + v_tma_head_stride, + v_tma_page_stride, + ), + ) k_tma = cute.make_tensor(k_iter, k_layout) v_tma = cute.make_tensor(v_iter, v_layout) else: @@ -3496,8 +3535,8 @@ def fmha_decode_launch( q_box_dims: tuple[object, ...] if cutlass.const_expr(cfg.use_variable_seqlens_q): # Packed Q is physically [sum_q_tokens, num_heads_q, head_dim]. Flatten - # the two head modes into Hq so the ragged token axis still fits in a - # rank-5 tensor map after the helper inserts its two synthetic modes. + # the head modes into Hq. Sparse Q uses this rank-3 map directly; + # other packed paths retain per-sequence ragged zero fill. q_tma = cute.make_tensor( q_iter, cute.make_layout( @@ -3509,7 +3548,7 @@ def fmha_decode_launch( q_box_dims = ( tma_box0_q, cfg.heads_q_per_kv, - cfg.q_tokens_per_cta, + cfg.q_tma_tokens_per_cta, ) q_groups = Int32( (cfg.max_seq_len_q + cfg.q_tokens_per_cta - 1) // cfg.q_tokens_per_cta @@ -3520,13 +3559,24 @@ def fmha_decode_launch( (cfg.heads_q_per_kv + cfg.tile_size_q - 1) // cfg.tile_size_q ) q_groups = Int32(cfg.max_seq_len_q) * head_ctas_per_token - tma_desc_q = create_tensor_map_ragged_from_tensor( - q_tma, - box_dims=q_box_dims, - ragged_dim=2, - stride_order=(0, 1, 2), - swizzle=tma_swizzle_q, - ) + if cutlass.const_expr(cfg.uses_q_token_kv_block_sparse_page_route): + # A partial group may read the next sequence into inactive Q rows. + # QK/softmax/PV rows are independent; row masks and guarded stores + # discard them. Physical allocation bounds still get TMA zero fill. + tma_desc_q = create_tensor_map_tiled_from_view( + q_tma, + box_dims=q_box_dims, + stride_order=(0, 1, 2), + swizzle=tma_swizzle_q, + ) + else: + tma_desc_q = create_tensor_map_ragged_from_tensor( + q_tma, + box_dims=q_box_dims, + ragged_dim=2, + stride_order=(0, 1, 2), + swizzle=tma_swizzle_q, + ) else: q_tma = cute.make_tensor( q_iter, @@ -3557,19 +3607,27 @@ def fmha_decode_launch( stride_order=(0, 1, 2, 3, 4), swizzle=tma_swizzle_q, ) + if cutlass.const_expr(cfg.uses_2d_flat_kv_tma): + k_box_dims = (tma_box0_k, tma_kv_tokens) + v_box_dims = (tma_box0_v, tma_kv_tokens) + kv_stride_order = (0, 1) + else: + k_box_dims = (tma_box0_k, tma_kv_tokens, 1, 1) + v_box_dims = (tma_box0_v, tma_kv_tokens, 1, 1) + kv_stride_order = (0, 1, 2, 3) if cutlass.const_expr(cfg.store_transformed_kv_in_tmem): tma_desc_k = create_tensor_map_tiled_from_view( k_tma, - box_dims=(tma_box0_k, tma_kv_tokens, 1, 1), - stride_order=(0, 1, 2, 3), + box_dims=k_box_dims, + stride_order=kv_stride_order, swizzle=tma_swizzle_k, dtype=cutlass.Float4E2M1FNx2, tma_format=cuda.TensorMapDataFormat.B4X16_P64, ) tma_desc_v = create_tensor_map_tiled_from_view( v_tma, - box_dims=(tma_box0_v, tma_kv_tokens, 1, 1), - stride_order=(0, 1, 2, 3), + box_dims=v_box_dims, + stride_order=kv_stride_order, swizzle=tma_swizzle_v, dtype=cutlass.Float4E2M1FNx2, tma_format=cuda.TensorMapDataFormat.B4X16_P64, @@ -3577,14 +3635,14 @@ def fmha_decode_launch( else: tma_desc_k = create_tensor_map_tiled_from_view( k_tma, - box_dims=(tma_box0_k, tma_kv_tokens, 1, 1), - stride_order=(0, 1, 2, 3), + box_dims=k_box_dims, + stride_order=kv_stride_order, swizzle=tma_swizzle_k, ) tma_desc_v = create_tensor_map_tiled_from_view( v_tma, - box_dims=(tma_box0_v, tma_kv_tokens, 1, 1), - stride_order=(0, 1, 2, 3), + box_dims=v_box_dims, + stride_order=kv_stride_order, swizzle=tma_swizzle_v, ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_resources.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_resources.py index fe49a7c67592..3e2031504ef7 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_resources.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_resources.py @@ -94,6 +94,26 @@ from .smem_block_sparse_metadata import SmemBlockSparseKvMetadataResource +@cute.jit +def _q_token_membership_keep_word16( + word0: Uint32, + word1: Uint32, + word2: Uint32, + word3: Uint32, + query_idx: Int32, +) -> Uint32: + """Extract one query's membership bit from sixteen packed page bytes.""" + keep_word = Uint32(0) + shift = query_idx & Int32(7) + words = (word0, word1, word2, word3) + for word_idx in cutlass.range_constexpr(4): + bits = (words[word_idx] >> shift) & Uint32(0x01010101) + # The disjoint partial products gather four byte bits without carries. + nibble = (bits * Uint32(0x01020408)) >> Uint32(24) + keep_word |= nibble << Uint32(word_idx * 4) + return cutlass.select_(Uint32(query_idx) < Uint32(8), keep_word, Uint32(0)) + + def _paged_sparse_kv_tma_transaction_geometry( *, tile_size_kv: int, @@ -158,8 +178,20 @@ def _kv_dtype_bytes_for_kind( def _decode_native_page_locator( cfg: Constexpr[FmhaDecodeConfig], page_locator: Int32, + kv_head: Int32 = 0, ) -> tuple[Int32, Int32]: - """Return ``(token_offset, physical_page)`` for one native table entry.""" + """Return token/page TMA coordinates for one prepared locator.""" + if cutlass.const_expr(cfg.use_flat_native_kv_tma): + # The prepared locator is the head-zero origin in the compact cache. + # Keep the entire fragment OOB when adding a head offset to a sentinel. + if cutlass.const_expr(cfg.flat_native_kv_num_heads > 1): + head_fragments = ( + cfg.effective_storage_tokens_per_page // cfg.num_tokens_per_page + ) + page_locator = (page_locator + kv_head * Int32(head_fragments)) | ( + page_locator >> Int32(31) + ) + return page_locator * Int32(cfg.num_tokens_per_page), Int32(0) token_offset = Int32(0) physical_page = page_locator if cutlass.const_expr(cfg.has_storage_subpages): @@ -224,14 +256,84 @@ def _issue_sparse_page_copies( head_dim_stage: Constexpr[int], head_dim_stage_offset: Constexpr[int], ): - """Partition (D64 chunk, physical fragment) copies across loader lanes. + """Issue physical page fragments through the selected loader schedule. The cache writer zero-fills page padding: full-fragment tail loads rely on this because score masking alone does not sanitize NaN/Inf in V. """ - chunk_hd = min(head_dim_stage, 64) + chunk_hd = min(head_dim_stage, 128 if cfg.use_fp8_qkv else 64) chunks = head_dim_stage // chunk_hd fragments = cfg.tile_size_kv // cfg.num_tokens_per_page + if cutlass.const_expr( + cfg.uses_2d_flat_kv_tma and cfg.k_dtype == cfg.v_dtype == cutlass.BFloat16 + ): + # Prefetch warp-owned locators before TMA issue and reuse each across + # BF16's head chunks. Static indexing keeps the small array in registers. + fragments_per_warp = fragments // cfg.load_num_warps + warp = _load_task_warp_rank(cfg) + if prims.elect_sync(): + locators = cutlass.Array( + Int32, (fragments_per_warp,), space=cutlass.AddressSpace.rmem + ) + for i in cutlass.range_constexpr(fragments_per_warp): + fragment = warp * Int32(fragments_per_warp) + Int32(i) + locators[i] = page_offsets.page_id(tile_idx, local_tile_idx, fragment) + for i in cutlass.range_constexpr(fragments_per_warp): + fragment = warp * Int32(fragments_per_warp) + Int32(i) + token_offset, _ = _decode_native_page_locator(cfg, locators[i], kv_head) + for chunk in cutlass.range_constexpr(chunks): + smem_offset = Int32( + chunk * chunk_hd * cfg.tile_size_kv + ) + fragment * Int32(chunk_hd * cfg.num_tokens_per_page) + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.subview(smem_offset), + tma_desc, + ( + Int32(head_dim_stage_offset + chunk * chunk_hd), + token_offset, + ), + barrier, + ) + return + if cutlass.const_expr(cfg.uses_2d_flat_kv_tma): + # A warp owns complete fragments across head planes. Its staged token + # coordinate is uniform: expose that before election so both TMA copies + # reuse uniform coordinates instead of serializing lane operands. + # General addressing retains its lane-parallel decode. One-shot + # kernels keep parallel shared loads to minimize startup latency. + uniform_fragments_per_warp = fragments // cfg.load_num_warps + uniform_warp = _load_task_warp_rank(cfg) + uniform_first_fragment = uniform_warp * Int32(uniform_fragments_per_warp) + uniform_page_ptr = page_offsets.page_id_pointer( + tile_idx, local_tile_idx, uniform_first_fragment + ) + uniform_destination_ptr = stage_base.data_ptr( + uniform_first_fragment * Int32(chunk_hd * cfg.num_tokens_per_page) + ) + # Reuse a compact set of uniform address operands across fragments; + # advance staged locators and SMEM destinations independently of the + # scattered global coordinates. Head-plane copies stay unrolled. + for _fragment_idx in cutlass.range(uniform_fragments_per_warp, unroll=1): + uniform_token_offset = cute.arch.make_warp_uniform( + Int32(uniform_page_ptr[0]) + ) + if prims.elect_sync(): + for head_chunk in cutlass.range_constexpr(chunks): + prims.cp_async_bulk_tensor_shared_cta_global( + uniform_destination_ptr + + Int32(head_chunk * chunk_hd * cfg.tile_size_kv), + tma_desc, + ( + Int32(head_dim_stage_offset + head_chunk * chunk_hd), + uniform_token_offset, + ), + barrier, + ) + uniform_page_ptr = uniform_page_ptr + Int32(1) + uniform_destination_ptr = uniform_destination_ptr + Int32( + chunk_hd * cfg.num_tokens_per_page + ) + return copies = fragments * chunks copies_per_warp = (copies + cfg.load_num_warps - 1) // cfg.load_num_warps lane = cute.arch.thread_idx()[0] & Int32(31) @@ -243,7 +345,9 @@ def _issue_sparse_page_copies( chunk = copy // Int32(fragments) fragment = copy % Int32(fragments) locator = page_offsets.page_id(tile_idx, local_tile_idx, fragment) - token_offset, physical_page = _decode_native_page_locator(cfg, locator) + token_offset, physical_page = _decode_native_page_locator( + cfg, locator, kv_head + ) smem_offset = chunk * Int32(chunk_hd * cfg.tile_size_kv) + fragment * Int32( chunk_hd * cfg.num_tokens_per_page ) @@ -253,7 +357,9 @@ def _issue_sparse_page_copies( ( Int32(head_dim_stage_offset) + chunk * Int32(chunk_hd), token_offset, - kv_head, + Int32(0) + if cutlass.const_expr(cfg.use_flat_native_kv_tma) + else kv_head, physical_page, ), barrier, @@ -338,7 +444,7 @@ def _create_initial_task_locals( leading_byte_offset = q_leading_bytes stride_byte_offset = Int32(1024) if cutlass.const_expr(self.cfg.q_dtype_bytes == 1): - q_head_dim_stage = self.cfg.head_dim_kv_stage + q_head_dim_stage = min(self.cfg.head_dim_kv_stage, 128) q_tile_bytes = Int32( self.cfg.tile_size_q * q_head_dim_stage * self.cfg.q_dtype_bytes ) @@ -389,7 +495,7 @@ def _tma_coords( logical_h_k_idx: Int32, logical_b_idx: Int32, head_dim_offset: Int32, - ) -> tuple[Int32, Int32, Int32, Int32, Int32]: + ) -> tuple[Int32, ...]: """Map a logical Q CTA to token-major tensor-map coordinates.""" cfg = self.cfg if cutlass.const_expr(cfg.use_variable_seqlens_q): @@ -407,6 +513,12 @@ def _tma_coords( cfg.heads_q_per_kv ) + head_cta_idx * Int32(cfg.tile_size_q) tokens_per_box = 1 + if cutlass.const_expr(cfg.uses_q_token_kv_block_sparse_page_route): + return ( + head_dim_offset, + global_head_idx, + self.q_token_offset + token_idx, + ) packed_coords = transform_ragged_coords( ( head_dim_offset, @@ -459,7 +571,7 @@ def _tma_coords( @cute.jit def _complete_grouped_q_padding(self, stage_info: StageInfo) -> None: - """Account structural grouped rows that have no corresponding TMA.""" + """Credit the Q tile rows not covered by the configured TMA box.""" cfg = self.cfg if cutlass.const_expr( cfg.groups_tokens_heads_q and cfg.q_manual_padding_rows > 0 @@ -486,7 +598,7 @@ def tma_load(self, stage_info: StageInfo) -> None: stage_base = self._smem_base_q.subview(stage_info.stage_idx * stage_elems) if cutlass.const_expr(cfg.q_dtype_bytes == 1): if _is_load_task_warp_leader(cfg): - if cutlass.const_expr(cfg.num_head_dim_stages_kv == 1): + if cutlass.const_expr(cfg.headdim <= 128): # FP8 with one head-dim stage is one tensor copy into the # complete staged Q tile. prims.cp_async_bulk_tensor_shared_cta_global( @@ -503,9 +615,9 @@ def tma_load(self, stage_info: StageInfo) -> None: else: # H256 FP8 stages Q by head-dim slices so each QK head-dim # stage sees a contiguous SMEM tile. - q_chunk_dim = cfg.head_dim_kv_stage + q_chunk_dim = 128 for q_chunk_idx in cutlass.range_constexpr( - cfg.num_head_dim_stages_kv + cfg.headdim // q_chunk_dim ): head_dim_offset = q_chunk_idx * q_chunk_dim smem_offset = head_dim_offset * cfg.tile_size_q @@ -734,9 +846,6 @@ def _create_initial_task_locals( # FP16, BF16, and FP8. self._desc_base = prims.Tcgen05SmemDesc(0) return {"kv_desc": cutlass.Int64(0), "v_desc": cutlass.Int64(0)} - kv_tile_bytes = Int32( - self.cfg.tile_size_kv * self.cfg.head_dim_kv_stage * inst_dtype_bytes - ) leading_byte_offset = Int32( self.cfg.tile_size_kv * min(self.cfg.head_dim_kv_stage, 64) @@ -744,13 +853,18 @@ def _create_initial_task_locals( ) stride_byte_offset = Int32(1024) if cutlass.const_expr(inst_dtype_bytes == 1): - leading_byte_offset = kv_tile_bytes + leading_byte_offset = Int32( + self.cfg.tile_size_kv * min(self.cfg.head_dim_kv_stage, 128) + ) stride_byte_offset = Int32( _major_k_stride_bytes(inst_dtype_bytes, self.cfg.head_dim_kv_stage) ) if cutlass.const_expr( self.kv_kind == KV_KIND_V - and (inst_dtype_bytes == 1 or self.cfg.headdim == 64) + and ( + (inst_dtype_bytes == 1 and self.cfg.head_dim_kv_stage <= 128) + or self.cfg.headdim == 64 + ) ): leading_byte_offset = Int32(0) inst_swizzle = prims.Tcgen05SmemSwizzle.SWIZZLE_128B @@ -1152,6 +1266,22 @@ def _producer_load( tile_idx = self._maybe_runtime_tile_idx(stage_info, local_tile_idx) inst_dtype_bytes = _kv_dtype_bytes_for_kind(self.cfg, self.kv_kind) if cutlass.const_expr(inst_dtype_bytes == 1): + if cutlass.const_expr( + cfg.use_fp8_qkv and cfg.uses_scattered_page_route + ): + _issue_sparse_page_copies( + cfg, + self.page_offsets_kv, + self._stage_base(stage_info), + tma_desc, + stage_info.barrier, + tile_idx, + local_tile_idx, + logical_h_k_idx, + head_dim_stage, + head_dim_stage_offset, + ) + return if prims.elect_sync() and _load_task_warp_rank(cfg) < Int32( min(page_fragments, cfg.load_num_warps) ): @@ -1162,7 +1292,6 @@ def _producer_load( # elected lane replay the full loop over-completes the TMA # transaction barrier when load_num_warps > 1. stage_base = self._stage_base(stage_info) - page_ids = self.page_offsets_kv.page_ids(tile_idx, local_tile_idx) sf_desc = ( self.tma_desc_v_sf if cutlass.const_expr(self.kv_kind == KV_KIND_V) @@ -1182,8 +1311,14 @@ def _producer_load( transactions_per_warp ): page_frag = first_page_frag + Int32(warp_transaction) + # A runtime warp rank cannot index a register vector + # without materializing it in local memory. Read only + # this issuer's locator from the retained SMEM window. + locator = self.page_offsets_kv.page_id( + tile_idx, local_tile_idx, page_frag + ) token_offset, physical_page = _decode_native_page_locator( - cfg, Int32(page_ids[page_frag]) + cfg, locator ) smem_page_offset = Int32( page_frag @@ -1658,10 +1793,10 @@ def page_ids(self, tile_idx: Int32, local_tile_idx: Int32) -> cutlass.Array: return self._smem_page_offsets.load(offset, vector_size=1, alignment=4) @cute.jit - def page_id( + def page_id_pointer( self, tile_idx: Int32, local_tile_idx: Int32, page_frag: Int32 - ) -> Int32: - """Load a tile's locator; offset stages are independent of K/V data stages.""" + ) -> cutlass.Pointer: + """Locate a staged ID independently of the K/V data-buffer stage.""" cfg = self.cfg pages_per_tile = cfg.tile_size_kv // cfg.num_tokens_per_page if cutlass.const_expr(self.holds_encoded_locator_window): @@ -1675,7 +1810,14 @@ def page_id( else: group_page_idx = (tile_idx * Int32(pages_per_tile)) & Int32(31) offset = self.consumer_work_stage * Int32(32) + group_page_idx - return Int32(self._smem_page_offsets[offset + page_frag]) + return self._smem_page_offsets.data_ptr(offset + page_frag) + + @cute.jit + def page_id( + self, tile_idx: Int32, local_tile_idx: Int32, page_frag: Int32 + ) -> Int32: + """Load one locator from the active page-offset stage.""" + return Int32(self.page_id_pointer(tile_idx, local_tile_idx, page_frag)[0]) @cute.jit def _q_token_kv_block_sparse_membership_word( @@ -1761,37 +1903,101 @@ def q_token_kv_block_sparse_page_membership( ) @cute.jit - def q_token_kv_block_sparse_page_memberships4( + def query_membership_keep_word( + self, stage_info: StageInfo, local_tile_idx: Int32, q_token_idx: Int32 + ) -> Uint32: + """Read a private, tail-padded query bit-plane without transposing it.""" + cfg = self.cfg + assert cfg.query_major_memberships and cfg.num_tokens_per_page == 4 + assert cfg.tile_size_kv == 128 and 2 <= cfg.max_seq_len_q <= 8 + query = q_token_idx & Int32(7) + word = Uint32(0) + if cutlass.const_expr(self.cache_memberships_in_smem): + self._create_initial_task_locals(stage_info.context) + word = Uint32( + self._smem_q_token_kv_block_sparse_memberships[ + local_tile_idx * Int32(8) + query + ] + ) + else: + task_cache = _decode_gen_task_cache(stage_info) + tile = local_tile_idx + Int32(task_cache[_TASK_CACHE_KV_RAW_TILE_BASE]) + if tile * Int32(32) <= Int32(task_cache[_TASK_CACHE_KV_PAGE_IDX_UB]): + head, batch = _logical_head_batch(stage_info, self.h_k_idx, self.b_idx) + row = cutlass.Int64(batch) + if cutlass.const_expr( + cfg.use_persistent_scheduler and not cfg.shares_sparse_pattern + ): + row = row * cutlass.Int64(self.num_heads_kv) + cutlass.Int64(head) + word = Uint32( + self.q_token_kv_block_sparse_page_memberships[ + row + * cutlass.Int64( + self.q_token_kv_block_sparse_page_membership_stride + ) + + cutlass.Int64(tile * Int32(8) + query) + ] + ) + return cutlass.select_(Uint32(q_token_idx) < Uint32(8), word, Uint32(0)) + + @cute.jit + def q_token_kv_block_sparse_page_membership_word4( self, stage_info: StageInfo, local_tile_idx: Int32, page_frag: Int32, - ) -> cutlass.Array: - """Load four aligned grouped-Q membership bytes from one packed word. - - The ``tmem_s`` Keeps Softmax path calls this helper four times to - assemble each lane-local 16-page membership word. - """ + ) -> Uint32: + """Load the packed word for four aligned grouped-Q page memberships.""" assert self.cfg.uses_q_token_kv_block_sparse_page_membership pages_per_tile = Int32(self.cfg.tile_size_kv // self.cfg.num_tokens_per_page) membership_idx = local_tile_idx * pages_per_tile + page_frag - memberships = cutlass.Array( - Uint32, - 4, - space=cutlass.AddressSpace.rmem, - ) membership_word, _ = self._q_token_kv_block_sparse_membership_word( stage_info, membership_idx ) - for elem_idx in cutlass.range_constexpr(4): - memberships[elem_idx] = ( - membership_word - >> Int32(Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIP_BITS * elem_idx) - ) & Uint32( - Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIP_MASK, + return membership_word + + @cute.jit + def q_token_kv_block_sparse_page_keep_word16( + self, + stage_info: StageInfo, + local_tile_idx: Int32, + page_frag: Int32, + query_idx: Int32, + ) -> Uint32: + """Extract sixteen memberships, widening only padded SMEM reads.""" + assert self.cfg.uses_q_token_kv_block_sparse_page_membership + assert not self.cfg.query_major_memberships + if cutlass.const_expr(self.cache_memberships_in_smem): + self._create_initial_task_locals(stage_info.context) + pages_per_tile = Int32( + self.cfg.tile_size_kv // self.cfg.num_tokens_per_page + ) + word_idx = (local_tile_idx * pages_per_tile + page_frag) // Int32(4) + # Callers select aligned sixteen-page windows; the producer owns + # and initializes the full padded SMEM tile, including its tail. + words = ( + self._smem_q_token_kv_block_sparse_memberships.subview(word_idx) + .data_ptr() + .load(count=4, alignment=16) + ) + word0, word1, word2, word3 = words[0], words[1], words[2], words[3] + else: + # Public GMEM rows need not have sixteen readable tail bytes. + # Preserve the existing live-prefix checks for each packed word. + word0 = self.q_token_kv_block_sparse_page_membership_word4( + stage_info, local_tile_idx, page_frag + ) + word1 = self.q_token_kv_block_sparse_page_membership_word4( + stage_info, local_tile_idx, page_frag + Int32(4) + ) + word2 = self.q_token_kv_block_sparse_page_membership_word4( + stage_info, local_tile_idx, page_frag + Int32(8) + ) + word3 = self.q_token_kv_block_sparse_page_membership_word4( + stage_info, local_tile_idx, page_frag + Int32(12) ) - return memberships + return _q_token_membership_keep_word16(word0, word1, word2, word3, query_idx) @consumer_work(returns=cached_page_ids) @cute.jit @@ -1840,14 +2046,19 @@ def _prefetch_page_offsets_async( lane, warp_rank, route_state, + kv_head, ): """Stage int32 IDs with cp.async and resolve direct inputs in place. Copy and mapping use the same lane ownership, so the per-thread wait suffices before in-place mapping. The existing TS full barrier then publishes all lanes' final locators to the TMA load task. - Materialized locators need no translation. Direct inputs contain only - selected block IDs; their causal tail locators are synthesized here. + Direct inputs contain only selected block IDs; their causal tail + locators are synthesized here. + The uniform FP8 issuer consumes flat token coordinates, decoded here + per staged entry rather than in the issuer. Held windows share these + across K/V; the streamed K/V staging cadence is unchanged. Other + issuers consume IDs. """ entries = self.page_ids_per_stage stage_base = stage_info.stage_idx * Int32(entries) @@ -1909,6 +2120,14 @@ def _prefetch_page_offsets_async( elif slot > last: # cp.async zero-fill is not the invalid-locator sentinel. self._smem_page_offsets[stage_base + offset] = Int32(-1) + if cutlass.const_expr( + self.cfg.use_fp8_qkv and self.cfg.uses_2d_flat_kv_tma + ): + locator = Int32(self._smem_page_offsets[stage_base + offset]) + token_offset, _ = _decode_native_page_locator( + self.cfg, locator, kv_head + ) + self._smem_page_offsets[stage_base + offset] = token_offset @cute.jit def _producer_load_page_offsets( @@ -2034,6 +2253,12 @@ def _producer_load_page_offsets( + membership_word_idx * Int32(Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIPS_PER_WORD) ) + if cutlass.const_expr(cfg.query_major_memberships): + # Every word in a query tile refers to the same 32 pages. + # The private producer initializes all eight tail words. + first_logical_page_idx = ( + raw_tile_base + membership_word_idx // Int32(8) + ) * Int32(32) if first_logical_page_idx <= page_idx_ub: membership_word = Uint32( self.q_token_kv_block_sparse_page_memberships[ @@ -2087,16 +2312,17 @@ def _producer_load_page_offsets( # The producer owns only the live byte prefix. Mask the # final word here so padding bytes remain unobservable # even when the caller reuses an uninitialized workspace. - live_memberships = ( - page_idx_ub - first_logical_page_idx + Int32(1) - ) - if live_memberships < Int32( - Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIPS_PER_WORD - ): - live_bits = live_memberships * Int32( - Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIP_BITS + if cutlass.const_expr(not cfg.query_major_memberships): + live_memberships = ( + page_idx_ub - first_logical_page_idx + Int32(1) ) - membership_word &= (Uint32(1) << live_bits) - Uint32(1) + if live_memberships < Int32( + Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIPS_PER_WORD + ): + live_bits = live_memberships * Int32( + Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIP_BITS + ) + membership_word &= (Uint32(1) << live_bits) - Uint32(1) self._smem_q_token_kv_block_sparse_memberships[ membership_word_idx ] = membership_word @@ -2117,6 +2343,7 @@ def _producer_load_page_offsets( lane_idx, page_warp_rank, direct_route, + logical_h_idx, ) elif cutlass.const_expr(self.holds_encoded_locator_window): # Stripe whole KV tiles: with two producers, warps own even/odd diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_o.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_o.py index e75003484df3..d08b770708f0 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_o.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_o.py @@ -505,7 +505,9 @@ def _vp_mma( ) else: v_operand = v_operand + Int32( - (cfg.head_dim_kv_stage * 2) if cfg.use_fp8_pv else 128 + (min(cfg.head_dim_kv_stage, 128) * 2) + if cfg.use_fp8_pv + else 128 ) if cutlass.const_expr(not cfg.uses_tmem_p): if cutlass.const_expr( diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_s.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_s.py index 715a2ec91bcc..73982aa0a912 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_s.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_s.py @@ -901,6 +901,12 @@ def _qk_mma( elif cutlass.const_expr(cfg.qk_dtype_bytes == 1): k_desc_offset = ki * Int32(2) q_desc_offset = ki * Int32(2) + if cutlass.const_expr(cfg.head_dim_kv_stage > 128): + # A full FP8 head stage consists of two SW128 planes, + # not a row-major D256 tile. Cross planes at K128. + chunk_idx = (ki * Int32(_mma_k_step_qk(cfg))) // Int32(128) + k_desc_offset += chunk_idx * Int32(8 * cfg.tile_size_kv - 8) + q_desc_offset += chunk_idx * Int32(8 * cfg.tile_size_q - 8) iter_k_desc = k_desc + k_desc_offset else: chunk_idx = (ki * Int32(_mma_k_step_qk(cfg))) // Int32(64) @@ -1072,24 +1078,42 @@ def _q_token_kv_block_sparse_keeps_kv128_membership_word( lane_idx, cfg.softmax_score_fragment_regs, ) + if cutlass.const_expr(cfg.query_major_memberships): + return self.page_offsets_ref.query_membership_keep_word( + stage_info, local_tile_idx, q_token_idx + ) >> Uint32(col_base // Int32(4)) keep_word = Uint32(0) page_span = min(cfg.num_tokens_per_page, cfg.num_s_regs_per_thread) pages_per_lane = cfg.num_s_regs_per_thread // page_span + if cutlass.const_expr(pages_per_lane >= 16): + for vector_idx in cutlass.range_constexpr(pages_per_lane // 16): + bits = self.page_offsets_ref.q_token_kv_block_sparse_page_keep_word16( + stage_info, + local_tile_idx, + col_base // Int32(cfg.num_tokens_per_page) + Int32(vector_idx * 16), + q_token_idx, + ) + keep_word |= bits << Uint32(vector_idx * 16) + return keep_word for page_vector_idx in cutlass.range_constexpr(pages_per_lane // 4): memberships = ( - self.page_offsets_ref.q_token_kv_block_sparse_page_memberships4( + self.page_offsets_ref.q_token_kv_block_sparse_page_membership_word4( stage_info, local_tile_idx, col_base // Int32(cfg.num_tokens_per_page) + Int32(page_vector_idx * 4), ) ) - for vector_elem_idx in cutlass.range_constexpr(4): - local_page_idx = page_vector_idx * 4 + vector_elem_idx - page_is_member = Uint32( - (memberships[vector_elem_idx] & membership_bit) != Uint32(0) - ) - keep_word = keep_word | (page_is_member << Uint32(local_page_idx)) + # Select query q's bit from each byte (positions 0/8/16/24), then + # pack them into bits 0..3 with a carry-free uint32 multiply, not + # a first-set-bit search. Zero q >= 8 so the bounded shift cannot + # alias another query's bit. + byte_bits = (memberships >> (q_token_idx & Int32(7))) & Uint32(0x01010101) + packed_bits = (byte_bits * Uint32(0x01020408)) >> Uint32(24) + packed_bits = cutlass.select_( + Uint32(q_token_idx) < Uint32(8), packed_bits, Uint32(0) + ) + keep_word = keep_word | (packed_bits << Uint32(page_vector_idx * 4)) if cutlass.const_expr(pages_per_lane < 4): for local_page_idx in cutlass.range_constexpr(pages_per_lane): membership = ( diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py index 18f2c272cf5c..a4a87200ca5d 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.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"); @@ -20,6 +21,8 @@ import cutlass import cutlass.cute as cute from cutlass import Int32, Int64, Uint8 +from cutlass.base_dsl.arch import Arch +from cutlass.cutlass_dsl import BaseDSL from cutlass.experimental import primitives as prims from cutlass.experimental.task_scheduling.enums import WorkAttr from cutlass.experimental.task_scheduling.memory import ( @@ -34,6 +37,7 @@ producer_work, ) +from flashinfer.cute_dsl.availability import is_rubin_cute_dsl_available from ..fmha_decode_config import FmhaDecodeConfig from .helpers_common import ( _TASK_CACHE_TMEM_BASE_OFFSET, @@ -46,14 +50,27 @@ from .smem_resources import SmemKvResource +def _has_packed_e4m3_e2m1_mul() -> bool: + """Whether the compile target has the packed E4M3 x E2M1 multiply. + + PTX 9.4 (CUDA 13.4) adds it for SM100 and SM103; ptxas rejects it for + SM107, which takes the f16x2 path instead. + """ + if not cutlass.target_version(min_version="13.4"): + return False + if not is_rubin_cute_dsl_available(): + return True + arch = BaseDSL._get_dsl().get_arch_enum() + return not (Arch.sm_107 <= arch <= Arch.sm_107f) + + @cute.jit def _mul_e2m1x4_e4m3x4(packed_fp4: Int32, packed_sf: Int32) -> Int32: """Apply four E4M3 scale factors to four unpacked E2M1 values. - PTX 9.4 (CUDA 13.4) supports the byte-padded E2M1 operand directly. Each FP4 value occupies the low nibble of one byte in the 32-bit input. """ - if cutlass.const_expr(cutlass.target_version(min_version="13.4")): + if cutlass.const_expr(_has_packed_e4m3_e2m1_mul()): return cute.arch.inline_ptx( "mul.e4m3x4.e2m1p4x4.e4m3x4.satfinite {$w0}, {$r0}, {$r1};", write_only_types=[Int32], diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_tasks.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_tasks.py index 7cea30d06cb4..7d2393db28f0 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_tasks.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_tasks.py @@ -2620,7 +2620,7 @@ def create_mma_task_one_inst_qkv( work_queue: WorkQueue | None, cfg: FmhaDecodeConfig, *, - tmem_stats_done: MemoryResource, + tmem_stats_done: MemoryResource | None, domain: int | cutlass.Int32, warp_idx: int | None = None, num_warps: int | None = None, @@ -2636,14 +2636,15 @@ def mma_schedule_body( tmem_s: MemoryResource, smem_p: MemoryResource, tmem_o: MemoryResource, - tmem_stats_done: MemoryResource, + tmem_stats_done: MemoryResource | None, q_desc: Any, ) -> None: """Schedule single-instance QK and PV waves across HEAD/LOOP/TAIL.""" def qk_mma(q_desc, qk_mma_label: str, section: FmhaStage) -> None: """Issue one scheduled single-instance QK wave.""" - tmem_stats_done.acquire() + if tmem_stats_done is not None: + tmem_stats_done.acquire() tmem_s.acquire() _issue_score_seed(tmem_s, section, cfg) for head_dim_stage_idx in range(cfg.num_head_dim_stages_kv): @@ -2662,7 +2663,8 @@ def qk_mma(q_desc, qk_mma_label: str, section: FmhaStage) -> None: ) smem_k.release() tmem_s.commit() - tmem_stats_done.commit() + if tmem_stats_done is not None: + tmem_stats_done.commit() def pv_mma(vp_mma_label: str, section: FmhaStage) -> None: """Issue one scheduled single-instance PV wave.""" @@ -2710,7 +2712,7 @@ def mma_schedule_prelude( smem_v.init_descriptor_state() smem_p.init_descriptor_state() - @schedule + @_schedule_with_optional_resources def mma_schedule( smem_q: MemoryResource, smem_k: MemoryResource, @@ -2718,8 +2720,8 @@ def mma_schedule( tmem_s: MemoryResource, smem_p: MemoryResource, tmem_o: MemoryResource, - tmem_stats_done: MemoryResource, - work_queue: WorkQueue | None = None, + tmem_stats_done: MemoryResource | None, + work_queue: WorkQueue | None, ) -> None: """Wrap one-inst MMA work in packed persistent skip handling.""" _decode_work_tile_schedule_with_invariant_bridge( @@ -2740,34 +2742,16 @@ def mma_schedule( ), ) - schedule_result = ( - mma_schedule( - smem_q, - smem_k, - smem_v, - tmem_s, - smem_p, - tmem_o, - tmem_stats_done, - ) - if work_queue is None - else mma_schedule( - smem_q, - smem_k, - smem_v, - tmem_s, - smem_p, - tmem_o, - tmem_stats_done, - work_queue, - ) + schedule_result = mma_schedule( + smem_q, smem_k, smem_v, tmem_s, smem_p, tmem_o, tmem_stats_done, work_queue ) src = [smem_q, smem_k, smem_v, smem_p] if work_queue is not None: src.append(work_queue) return task_class( src_resources=src, - dst_resources=[tmem_s, tmem_o, tmem_stats_done], + dst_resources=[tmem_s, tmem_o] + + ([tmem_stats_done] if tmem_stats_done is not None else []), cfg=cfg, warp_idx=cfg.mma_warp_idx if warp_idx is None else warp_idx, num_warps=cfg.mma_num_warps if num_warps is None else num_warps, @@ -4149,7 +4133,7 @@ def padding_body() -> None: warp_idx=warp_idx, num_warps=num_warps, schedule=captured_schedule, - num_registers=cfg.mma_load_task_num_registers, + num_registers=cfg.padding_task_num_registers(warp_idx, num_warps), name="PaddingTask", **kw, ) 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 bfd4d2229592..f28769a10c8d 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 @@ -23,9 +23,11 @@ * S2-S4 use one 512-thread CTA per 8 KiB output slice and fold the exact split count in registers. -* S5 and larger use 128-thread CTAs over 2 KiB slices. Each cluster rank folds - its split slots locally, publishes ``(LSE, O)`` through SMEM, and rank zero - merges the distributed states. G16 uses a 4x4 two-level merge. +* S5-S16 use one 512-thread CTA per 2 KiB slice, with four adjacent lanes + folding interleaved split slots and merging through warp shuffles. +* S17 and larger use 128-thread CTAs over 2 KiB slices. Each cluster rank + publishes its local ``(LSE, O)`` through SMEM; rank zero merges the states. + Sixteen-CTA clusters use a 4x4 two-level merge. A 512-thread exact-split kernel remains as the serial reference schedule. Both schedules merge the optional attention-sink denominator, pack the requested @@ -459,9 +461,12 @@ def decode_gen_parallel_separate_reduction_kernel( """Reduce split partials with a compact or clustered constexpr schedule. S2-S4 use one 512-thread CTA per 8 KiB output slice and reduce the exact - split count directly. Larger schedules use 128-thread CTAs over 2 KiB - slices. Each cluster rank owns 2, 4, or 8 split slots; G1 stores directly, - G16 uses a two-level 4x4 merge, and G2/G4/G8 finalize through rank zero. + split count directly. Larger schedules cover 2 KiB slices with 128 output + fragments per CTA. A single CTA (up to 16 slots) spreads each fragment's slots + over ``PARALLEL_REDUCTION_SLOT_LANES`` adjacent lanes, merges them with + warp shuffles, and stores directly; clustered profiles keep one lane per + fragment, each cluster rank owns 2, 4, or 8 split slots, G16 uses a + two-level 4x4 merge, and G2/G4/G8 finalize through rank zero. """ if cutlass.const_expr(cfg.use_compact_parallel_reduction): @@ -494,9 +499,15 @@ def decode_gen_parallel_separate_reduction_kernel( bytes_per_output_row = Int32(cfg.headdim * PARTIAL_O_ELEMENT_BYTES) slice_idx = block_idx_x // Int32(cfg.parallel_reduction_cluster_size) + # ``slot_lanes`` adjacent lanes share one 16-byte output fragment and each + # own an interleaved subset of its split slots (lane ``l`` folds slots + # ``l, l + slot_lanes, ...`` so a short runtime prefix stays balanced). + slot_lanes = cfg.parallel_reduction_slot_lanes + fragment_idx = thread_idx // Int32(slot_lanes) + lane_slot = thread_idx % Int32(slot_lanes) reduce_base_offset = slice_idx * Int32( PARALLEL_REDUCTION_BYTES_PER_SLICE - ) + thread_idx * Int32(REDUCTION_BYTES_PER_THREAD) + ) + fragment_idx * Int32(REDUCTION_BYTES_PER_THREAD) reduce_row_idx = reduce_base_offset // bytes_per_output_row reduce_col_idx = (reduce_base_offset % bytes_per_output_row) // Int32( PARTIAL_O_ELEMENT_BYTES @@ -531,10 +542,11 @@ def decode_gen_parallel_separate_reduction_kernel( output_vals[elem_idx] = Float32(0.0) global_lse = Float32(-Float32.inf) - # Batch up to four independent GMEM loads before folding them. The final - # batch width is compile-time. Padded split slots stay neutral and never - # form GMEM pointers, including non-power-of-two split counts. + # Issue every owned slot's independent GMEM loads before folding them. The + # final batch width is compile-time. Padded split slots stay neutral and + # never form GMEM pointers, including non-power-of-two split counts. local_splits = cfg.parallel_reduction_splits_per_cta + lane_splits = local_splits // slot_lanes partial_lse = cutlass.Array( Float32, PARALLEL_REDUCTION_LOAD_BATCH, space=cutlass.AddressSpace.rmem ) @@ -575,12 +587,18 @@ def decode_gen_parallel_separate_reduction_kernel( ) for split_base_i in cutlass.range_constexpr( - 0, local_splits, PARALLEL_REDUCTION_LOAD_BATCH + 0, lane_splits, PARALLEL_REDUCTION_LOAD_BATCH ): - batch_width = min(PARALLEL_REDUCTION_LOAD_BATCH, local_splits - split_base_i) + batch_width = min(PARALLEL_REDUCTION_LOAD_BATCH, lane_splits - split_base_i) split_base = Int32(split_base_i) + # Prefetch only published slots; inactive slots and padded Q rows + # are in bounds but may contain uninitialized caller-owned scratch. for jj in cutlass.range_constexpr(batch_width): - split_idx = cluster_rank * Int32(local_splits) + split_base + Int32(jj) + split_idx = ( + cluster_rank * Int32(local_splits) + + (split_base + Int32(jj)) * Int32(slot_lanes) + + lane_slot + ) valid_split_idx = cutlass.const_expr( cfg.parallel_reduction_padded_splits == cfg.max_splits_kv ) or split_idx < Int32(cfg.max_splits_kv) @@ -622,7 +640,11 @@ def decode_gen_parallel_separate_reduction_kernel( partial_regs[regs_base + reg_idx] = loaded_partial_regs[reg_idx] for jj in cutlass.range_constexpr(batch_width): - split_idx = cluster_rank * Int32(local_splits) + split_base + Int32(jj) + split_idx = ( + cluster_rank * Int32(local_splits) + + (split_base + Int32(jj)) * Int32(slot_lanes) + + lane_slot + ) valid_split_idx = cutlass.const_expr( cfg.parallel_reduction_padded_splits == cfg.max_splits_kv ) or split_idx < Int32(cfg.max_splits_kv) @@ -649,11 +671,50 @@ def decode_gen_parallel_separate_reduction_kernel( ) global_lse = new_lse - # G1 stores its register accumulator directly and compiles out SMEM - # publication, mapa, and cluster barriers. The branch is CTA-uniform and - # resolved at compile time. + # G1 merges the lanes sharing one fragment with a warp butterfly, then the + # first lane stores the register accumulator directly and compiles out + # SMEM publication, mapa, and cluster barriers. The branch is CTA-uniform + # and resolved at compile time. if cutlass.const_expr(cfg.parallel_reduction_cluster_size == 1): - if valid_reduce_row: + # Shuffles execute warp-uniformly outside every row/slot predicate; + # lanes of invalid fragments exchange neutral or unused state. Names + # stay distinct from the clustered finalize below so the DSL sees one + # stable type per variable across the compile-time branches. + if cutlass.const_expr(slot_lanes > 1): + for lane_level in cutlass.range_constexpr((slot_lanes - 1).bit_length()): + lane_xor = 1 << lane_level + lane_peer_lse = Float32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=global_lse, + offset=lane_xor, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + ) + lane_peer_vals = cutlass.Array( + Float32, OUTPUT_VALUES_PER_THREAD, space=cutlass.AddressSpace.rmem + ) + for elem_idx in cutlass.range_constexpr(OUTPUT_VALUES_PER_THREAD): + lane_peer_vals[elem_idx] = Float32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=output_vals[elem_idx], + offset=lane_xor, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + ) + lane_merged_lse, lane_own_weight, lane_peer_weight = merge_log2_lse( + global_lse, lane_peer_lse + ) + for elem_idx in cutlass.range_constexpr(OUTPUT_VALUES_PER_THREAD): + output_vals[elem_idx] = ( + output_vals[elem_idx] * lane_own_weight + + lane_peer_vals[elem_idx] * lane_peer_weight + ) + global_lse = lane_merged_lse + if valid_reduce_row and lane_slot == Int32(0): _store_parallel_reduction_output( output_vals, global_lse, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/gather.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/gather.py index c60f2070fd03..d36f2ad068f2 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/gather.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/gather.py @@ -23,6 +23,7 @@ from cutlass.experimental import primitives as prims from .constants import TMA_GATHER_ROWS +from .sparse_indices import SparseKVIndices, SparseSourceTile @dsl_user_op @@ -50,7 +51,9 @@ def load_sparse_rows(routes, request, offset, *, fragments: cutlass.Constexpr[in for fragment in cutlass.range_constexpr(fragments): index = offset + lane + Int32(fragment * 32) value = Int32(0x7FFFFFFF) - if index < routes.shape[0]: + if cutlass.const_expr(isinstance(routes, SparseSourceTile)): + value = routes.load_row(index) + elif index < routes.shape[0]: value = Int32(routes[index, request]) rows[fragment] = value return rows @@ -59,10 +62,10 @@ def load_sparse_rows(routes, request, offset, *, fragments: cutlass.Constexpr[in @cute.jit def load_sparse_quad(routes, request, offset): """Keep one contiguous row quadruple in its issuing lane's registers.""" + if cutlass.const_expr(isinstance(routes, SparseKVIndices)): + source = routes.select_source_tile(offset, request) + return source.load_quad(offset) rows = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) - if cutlass.const_expr(hasattr(routes, "routes_for_tile")): - tile_routes = routes.routes_for_tile(offset, request) - return tile_routes.mapped_quad(offset) for j in cutlass.range_constexpr(4): rows[j] = Int32(0x7FFFFFFF) if offset + Int32(3) < routes.shape[0]: @@ -91,6 +94,8 @@ def decode_sparse_quad(primary, extra, raw): # reserved sentinel is already an out-of-bounds TMA coordinate; # retaining it gives zero fill without a compare/select per row. rows[j] = raw[j] & Int32(0x7FFFFFFF) + # Internally generated invalid coordinates retain the source bit: the + # first coordinate selects the shared descriptor even when it is invalid. return select_gather_map(primary, extra, raw[0] < 0), rows @@ -150,6 +155,7 @@ def gather4_cached( ) row = raw & Int32(0x7FFFFFFF) rows[j] = row if row != Int32(0x7FFFFFFF) else Int32(-1) + # The same source-tag invariant applies to cached invalid anchors. descriptor = select_gather_map(primary, extra, raw0 < 0) if prims.elect_sync(): gather4( @@ -174,7 +180,9 @@ def invalid_sparse_token( Bit 31 selects the source; bits 0..30 hold its TMA row coordinate. All-one low bits denote an invalid row for either source. """ - if cutlass.const_expr(sparse and hasattr(routes, "is_invalid")): + if cutlass.const_expr( + sparse and isinstance(routes, (SparseSourceTile, CachedSparseMask)) + ): return (token >= length) | routes.is_invalid(token, request) invalid = token >= length if cutlass.const_expr(sparse): @@ -272,7 +280,7 @@ def cache_sparse_mask( ): lane = cute.arch.thread_idx()[0] % 32 warp = cute.arch.thread_idx()[0] % 128 // 32 - if cutlass.const_expr(hasattr(routes, "prefix_mask32")): + if cutlass.const_expr(isinstance(routes, SparseSourceTile) and routes.valid_prefix): low, high = Int32(0), Int32(0) if not routes.tile_is_full(tile_offset, length): if cutlass.const_expr(swaps): diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/sparse_indices.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/sparse_indices.py new file mode 100644 index 000000000000..d944aa034bd1 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/sparse_indices.py @@ -0,0 +1,308 @@ +# 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. + +"""Select a source index buffer from a tile's logical KV offset. + + row = indices.bind_query_row(query_row) + tile = row.select_source_tile(tile_start) + quad = tile.load_quad(kv_offset) + +The primary source's logical span is rounded to 128 so BK64/BK128 tiles never +cross pools, even when physical rows are random. Caller arrays/live lengths +need no such padding; bounded loads synthesize invalid tails. A 32-entry +boundary would suffice for the grouped TMA descriptors alone, but not for +this whole-tile source selection. These compiler-side objects carry existing +pointers/scalars through CuTe calls; they allocate no merged metadata buffer. +Only SparseKVLengths retains tensor indexing for the shared scheduler. +""" + +import cutlass +import cutlass.cute as cute +from cutlass import Int32, Int64 + + +class SparseKVIndices: + """Caller-owned index lists; optionally cache one query row's live lengths.""" + + def __init__(self, values, capacity, assume_valid_prefix=False, has_extra=True): + self.values = values + ( + self.primary_indices, + self.extra_indices, + self.primary_lengths, + self.extra_lengths, + self.query_row, + self.primary_length, + self.extra_length, + ) = values + self.capacity = capacity + self.assume_valid_prefix = assume_valid_prefix + self.has_extra = has_extra + self.is_bound = self.query_row is not None + if self.is_bound: + self.primary_span = ( + (self.primary_length + Int32(127)) // Int32(128) * Int32(128) + ) + + def __extract_mlir_values__(self): + return cutlass.extract_mlir_values(self.values) + + def __new_from_mlir_values__(self, values): + return SparseKVIndices( + cutlass.new_from_mlir_values(self.values, values), + self.capacity, + self.assume_valid_prefix, + self.has_extra, + ) + + @cute.jit + def bind_query_row(self, query_row, primary_length=None, extra_length=None): + if cutlass.const_expr(self.is_bound): + return self + row = Int64(query_row) + # Once bound, carry scalar lengths instead of the two length tensors. + return SparseKVIndices( + ( + self.primary_indices, + self.extra_indices, + None, + None, + row, + Int32(self.primary_lengths[row]) + if cutlass.const_expr(primary_length is None) + else primary_length, + ( + Int32(self.extra_lengths[row]) + if cutlass.const_expr(extra_length is None) + else extra_length + ) + if cutlass.const_expr(self.has_extra) + else Int32(0), + ), + self.capacity, + self.assume_valid_prefix, + self.has_extra, + ) + + @cute.jit + def select_source_tile( + self, tile_start, query_row=None, *, for_mask: cutlass.Constexpr[bool] = False + ): + row = self.bind_query_row(query_row) + if cutlass.const_expr(not self.has_extra): + # A single-pool plan needs no second length, pointer selection or tag. + if cutlass.const_expr(for_mask and self.assume_valid_prefix): + return SparseSourceTile( + (None, Int32(0), row.primary_length, None), valid_prefix=True + ) + return SparseSourceTile( + ( + row.primary_indices[row.query_row, None], + Int32(0), + row.primary_length, + Int32(0), + ), + valid_prefix=self.assume_valid_prefix, + ) + use_extra = tile_start >= row.primary_span + origin = row.primary_span if use_extra else Int32(0) + length = row.extra_length if use_extra else row.primary_length + if cutlass.const_expr(for_mask and self.assume_valid_prefix): + # Analytic masks need no index pointer or index loads. + return SparseSourceTile((None, origin, length, None), valid_prefix=True) + primary_ptr = row.primary_indices.iterator.toint() + row.query_row * Int64( + row.primary_indices.stride[0] + ) * Int64(4) + extra_ptr = row.extra_indices.iterator.toint() + row.query_row * Int64( + row.extra_indices.stride[0] + ) * Int64(4) + # Source tensors can have different layouts; select numeric addresses. + address = extra_ptr if use_extra else primary_ptr + indices = cute.make_tensor( + cute.make_ptr(Int32, address, assumed_align=4), + cute.make_layout((self.capacity,), stride=(1,)), + ) + # The first coordinate chooses a shared TMA descriptor, even when + # invalid. Preserve this source tag on every synthesized tail entry. + tag = Int32(-2147483648) if use_extra else Int32(0) + return SparseSourceTile( + (indices, origin, length, tag), valid_prefix=self.assume_valid_prefix + ) + + +class SparseSourceTile: + """Selected source and bounds; analytic prefix masks need no index pointer.""" + + def __init__(self, values, valid_prefix=False): + self.values = values + self.indices, self.origin, self.length, self.tag = values + self.valid_prefix = valid_prefix + + def __extract_mlir_values__(self): + return cutlass.extract_mlir_values(self.values) + + def __new_from_mlir_values__(self, values): + return SparseSourceTile( + cutlass.new_from_mlir_values(self.values, values), self.valid_prefix + ) + + @cute.jit + def _bounded_index(self, token): + position = token - self.origin + # A zero-capacity source binds one dummy index. Even masked scalar + # loads stay inside its backing; validity is determined separately. + safe = cute.math.min( + cute.math.max(position, Int32(0)), + cute.math.max(self.length - Int32(1), Int32(0)), + ) + index = Int32(self.indices[safe]) + valid = (position >= Int32(0)) & (position < self.length) + if cutlass.const_expr(not self.valid_prefix): + valid = valid & (index >= Int32(0)) + return index, valid + + @cute.jit + def load_row(self, token): + index, valid = self._bounded_index(token) + return (index if valid else Int32(0x7FFFFFFF)) | self.tag + + @cute.jit + def load_full_quad(self, token): + """Load an in-bounds quad; callers establish its source-local bounds.""" + values = cutlass.Pointer( + self.indices.iterator + token - self.origin, dtype=Int32 + ).load(count=4, alignment=4) + rows = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) + for j in cutlass.range_constexpr(4): + index = Int32(values[j]) + if cutlass.const_expr(self.valid_prefix): + rows[j] = index | self.tag + else: + rows[j] = (index if index >= Int32(0) else Int32(0x7FFFFFFF)) | self.tag + return rows + + @cute.jit + def load_quad(self, token): + """One vector load for an active quad; bounded scalar loads for tails.""" + rows = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) + for j in cutlass.range_constexpr(4): + rows[j] = Int32(0x7FFFFFFF) | self.tag + position = token - self.origin + if position >= Int32(0) and position + Int32(3) < self.length: + values = self.load_full_quad(token) + for j in cutlass.range_constexpr(4): + rows[j] = values[j] + elif position + Int32(3) >= Int32(0) and position < self.length: + for j in cutlass.range_constexpr(4): + rows[j] = self.load_row(token + Int32(j)) + return rows + + @cute.jit + def is_invalid(self, token, request): + del request + if cutlass.const_expr(self.valid_prefix): + return (token < self.origin) | (token >= self.origin + self.length) + _, valid = self._bounded_index(token) + return ~valid + + @cute.jit + def tile_is_full(self, start, live_length): + return start + Int32(128) <= cute.math.min( + self.origin + self.length, live_length + ) + + @cute.jit + def prefix_mask32(self, start, live_length): + # Source/tile boundaries are aligned; start is within this source. + end = cute.math.min(self.origin + self.length, live_length) + valid = cute.math.min(cute.math.max(end - start, Int32(0)), Int32(32)) + # Avoid an undefined shift by 32, including on an unselected SSA arm. + bits = Int32(-1) << (valid & Int32(31)) + return bits if valid < Int32(32) else Int32(0) + + +class SparseKVLengths: + """Adapter for the scheduler's lengths[query_row] interface only.""" + + def __init__(self, values, has_extra=True): + self.values = values + self.has_extra = has_extra + self.primary_lengths, self.extra_lengths, self.bound_length, rows = values + self.shape = ( + (rows,) if self.bound_length is not None else self.primary_lengths.shape + ) + + def __extract_mlir_values__(self): + return cutlass.extract_mlir_values(self.values) + + def __new_from_mlir_values__(self, values): + return SparseKVLengths( + cutlass.new_from_mlir_values(self.values, values), self.has_extra + ) + + @cute.jit + def __getitem__(self, query_row): + if cutlass.const_expr(self.bound_length is not None): + return self.bound_length + return _sparse_execution_length( + Int32(self.primary_lengths[query_row]), + Int32(self.extra_lengths[query_row]) + if cutlass.const_expr(self.has_extra) + else Int32(0), + ) + + +@cute.jit +def _sparse_execution_length(primary_length, extra_length): + primary_tiles = (primary_length + Int32(127)) // Int32(128) + extra_tiles = (extra_length + Int32(127)) // Int32(128) + # Empty rows still execute a masked tile to publish O=0 and LSE=-inf. + return cute.math.max((primary_tiles + extra_tiles) * Int32(128), Int32(1)) + + +@cute.jit +def bind_sparse_kv_indices( + primary_indices, + extra_indices, + primary_lengths, + extra_lengths, + capacity: cutlass.Constexpr[int], + *, + query_row=None, + assume_valid_prefix: cutlass.Constexpr[bool] = False, + has_extra: cutlass.Constexpr[bool] = True, +): + """Bind index readers and scheduler lengths without copying metadata.""" + indices = SparseKVIndices( + ( + primary_indices, + extra_indices, + primary_lengths, + extra_lengths, + None, + None, + None, + ), + capacity, + assume_valid_prefix, + has_extra, + ) + if cutlass.const_expr(query_row is not None): + indices = indices.bind_query_row(query_row) + length = _sparse_execution_length(indices.primary_length, indices.extra_length) + lengths = SparseKVLengths( + (None, None, length, primary_lengths.shape[0]), has_extra + ) + else: + lengths = SparseKVLengths( + (primary_lengths, extra_lengths, None, None), has_extra + ) + return indices, lengths diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/sparse_policy.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/sparse_policy.py index 0b838cf4d3a3..8f6100825f04 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/sparse_policy.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/sparse_policy.py @@ -29,9 +29,6 @@ class SparseMlaProfile: reduction: str = "gmem_separate" head_dim_ctas: int = 1 gather_issue_warps: int = 4 - # Separate lists help qualified small decode grids; premerged tagged - # routes remain selected where their simpler loader improves throughput. - direct_inputs: bool = True balanced_registers: bool = False reuse_kv: bool = False reuse_kv_stages: int = 0 @@ -78,7 +75,6 @@ def _base_fp8(rows, heads, capacity, sm_count): # tiles saves V reloads; uniform coordinates amortize repeated D slices. return SparseMlaProfile( family="keep", - direct_inputs=False, reuse_kv=True, reuse_kv_stages=10, uniform_offset_cache=True, @@ -249,19 +245,16 @@ def _choose_fp8(rows, heads, query_length, capacity, sm_count): return SparseMlaProfile( family="keep", gather_issue_warps=8, - direct_inputs=False, reuse_kv=True, reuse_kv_stages=10, uniform_offset_cache=True, scheduler="clc", paired_correction=True, ) - if heads == 128 and capacity >= 256 and (rows * 2 >= sm_count): - # The FP8 2CTA family uses its qualified static persistent scheduler; - # BF16's CLC choice below is a separate measured schedule. Short lists - # (capacity 256-384, e.g. HCA at moderate context) run this schedule - # too: one 2CTA tile covers H128 with a single KV stream, where the - # M64 keep fallback repeats it per head tile (B200: 2.4x faster). + if heads == 128 and rows * 2 >= sm_count: + # Planned capacity is at least 256. Once the two-CTA grid fills a + # device wave, share KV across H128 even for short selected lists. + # BF16's CLC choice below is a separately qualified schedule. return SparseMlaProfile(family="2cta", scheduler="static", gather_issue_warps=1) if ( heads <= 16 @@ -276,7 +269,6 @@ def _choose_fp8(rows, heads, query_length, capacity, sm_count): # Smaller grids retain M8/M16 to preserve independent work. return SparseMlaProfile( family="keep", - direct_inputs=False, reuse_kv=True, reuse_kv_stages=10, uniform_offset_cache=True, @@ -328,7 +320,6 @@ def _choose_bf16(rows, heads, query_length, capacity, sm_count): return SparseMlaProfile( family="keep" if keep else "swap", tile_size_q=64 if keep else 8 if heads <= 8 else 16, - direct_inputs=False, balanced_registers=not keep, reuse_kv=True, reuse_kv_stages=5 if keep else 6, @@ -388,7 +379,6 @@ def _choose_bf16(rows, heads, query_length, capacity, sm_count): kv_pipeline_stages=8, uniform_offset_cache=heads == 128, scheduler="clc" if heads == 128 and rows >= sm_count else "nonpersistent", - direct_inputs=False, ) return _base_bf16(rows, heads, capacity, sm_count) @@ -419,7 +409,7 @@ def _size_buffers(profile, *, rows, heads, query_length, capacity, dtype, sm_cou defer_max_update=True, ) if query_length <= 8 and rows >= 2 * sm_count: - profile = replace(profile, scheduler="clc", direct_inputs=True) + profile = replace(profile, scheduler="clc") if ( heads == 32 and query_length <= 8 @@ -451,7 +441,6 @@ def _size_buffers(profile, *, rows, heads, query_length, capacity, dtype, sm_cou page_pipeline_stages=2, gather_issue_warps=8, uniform_offset_cache=True, - direct_inputs=True, scheduler="clc", defer_max_update=True, ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/sparse_views.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/sparse_views.py deleted file mode 100644 index ffc74f6d74af..000000000000 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/sparse_views.py +++ /dev/null @@ -1,314 +0,0 @@ -# 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. - -"""CTA-local views of caller-prepared storage-row indices. - -Adapt separate primary/extra physical-row lists to the existing route/length -interface without materializing a merged GPU array. Source spans are padded -to 128 and tagged per tile; arbitrary holes remain masked. Nonpersistent CTAs -bind one request, while persistent work tiles rebind live pointers/lengths. -The common premerged-route path bypasses these views. Neither source is -required to be SWA or compressed, and logical-page conversion is external. -""" - -import cutlass -import cutlass.cute as cute -from cutlass import Int32, Int64 - - -class SparseRouteView: - def __init__(self, values, capacity, assume_valid_prefix=False): - self.values = values - self.capacity = capacity - self.assume_valid_prefix = assume_valid_prefix - self.si, self.ci, self.sl, self.cl, self.row = values - self.primary_span = (self.sl + Int32(127)) // Int32(128) * Int32(128) - self.shape = (capacity, self.si.shape[0]) - - def __extract_mlir_values__(self): - return cutlass.extract_mlir_values(self.values) - - def __new_from_mlir_values__(self, values): - return SparseRouteView( - cutlass.new_from_mlir_values(self.values, values), - self.capacity, - self.assume_valid_prefix, - ) - - @cute.jit - def __getitem__(self, coordinate): - token, _ = coordinate - result = Int32(0x7FFFFFFF) - if token < self.sl: - index = Int32(self.si[self.row, token]) - if index >= 0: - result = index - elif token >= self.primary_span: - result = Int32(-1) - position = token - self.primary_span - if position < self.cl: - index = Int32(self.ci[self.row, position]) - if index >= 0: - result = index | Int32(-2147483648) - return result - - @cute.jit - def routes_for_tile(self, token, request=None): - compressed = token >= self.primary_span - primary_ptr = self.si.iterator.toint() + self.row * Int64( - self.si.stride[0] - ) * Int64(4) - extra_ptr = self.ci.iterator.toint() + self.row * Int64( - self.ci.stride[0] - ) * Int64(4) - address = extra_ptr if compressed else primary_ptr - origin = self.primary_span if compressed else Int32(0) - length = self.cl if compressed else self.sl - tag = Int32(-2147483648) if compressed else Int32(0) - indices = cute.make_tensor( - cute.make_ptr(Int32, address, assumed_align=4), - cute.make_layout((self.capacity,), stride=(1,)), - ) - return SparseTileRoutes((indices, origin, length, tag)) - - @cute.jit - def mask_for_tile(self, token, request=None): - # Source boundaries are 128-row aligned. Select the pointer and prefix - # once per score tile, outside the unrolled register mask loop. - compressed = token >= self.primary_span - if cutlass.const_expr(self.assume_valid_prefix): - origin = self.primary_span if compressed else Int32(0) - length = self.cl if compressed else self.sl - return SparsePrefixMask((origin, length)) - routes = self.routes_for_tile(token) - return routes - - @cute.jit - def is_invalid(self, token, request): - return self.mask_for_tile(token).is_invalid(token, request) - - -class SparseLengthView: - def __init__(self, values): - self.values = values - self.length, rows = values - self.shape = (rows,) - - def __extract_mlir_values__(self): - return cutlass.extract_mlir_values(self.values) - - def __new_from_mlir_values__(self, values): - return SparseLengthView(cutlass.new_from_mlir_values(self.values, values)) - - @cute.jit - def __getitem__(self, request): - del request - return self.length - - -class SparsePrefixMask: - """Analytic validity for one source's compact, hole-free prefix.""" - - def __init__(self, values): - self.values = values - self.origin, self.length = values - - def __extract_mlir_values__(self): - return cutlass.extract_mlir_values(self.values) - - def __new_from_mlir_values__(self, values): - return SparsePrefixMask(cutlass.new_from_mlir_values(self.values, values)) - - @cute.jit - def tile_is_full(self, start, live_length): - return start + Int32(128) <= cute.math.min( - self.origin + self.length, live_length - ) - - @cute.jit - def prefix_mask32(self, start, live_length): - # Tiles and source boundaries are 128-aligned; start never precedes - # this source's origin. Include the request's live bound as well. - end = cute.math.min(self.origin + self.length, live_length) - valid = cute.math.min(cute.math.max(end - start, Int32(0)), Int32(32)) - # Mask the shift count even on the unselected side of the expression. - # A full 32-bit valid prefix must not execute an undefined shift by 32. - bits = Int32(-1) << (valid & Int32(31)) - return bits if valid < Int32(32) else Int32(0) - - @cute.jit - def is_invalid(self, token, request): - del request - return (token < self.origin) | (token >= self.origin + self.length) - - -class SparseTileRoutes: - """One source selection per tile, with bounded scalar and quad loads.""" - - def __init__(self, values): - self.values = values - self.indices, self.origin, self.length, self.tag = values - - def __extract_mlir_values__(self): - return cutlass.extract_mlir_values(self.values) - - def __new_from_mlir_values__(self, values): - return SparseTileRoutes(cutlass.new_from_mlir_values(self.values, values)) - - @cute.jit - def _bounded_index(self, token): - position = token - self.origin - # Zero-capacity sources bind a one-row dummy. Clamp even masked - # positions so both gather and softmax can use a bounded load. - safe = cute.math.min( - cute.math.max(position, Int32(0)), - cute.math.max(self.length - Int32(1), Int32(0)), - ) - index = Int32(self.indices[safe]) - valid = (position >= Int32(0)) & (position < self.length) & (index >= Int32(0)) - return index, valid - - @cute.jit - def is_invalid(self, token, request): - del request - _, valid = self._bounded_index(token) - return ~valid - - @cute.jit - def mapped_row(self, token): - index, valid = self._bounded_index(token) - return (index if valid else Int32(0x7FFFFFFF)) | self.tag - - @cute.jit - def mapped_quad(self, token): - """Load a complete active quad together and preserve partial-prefix masking.""" - rows = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) - for j in cutlass.range_constexpr(4): - rows[j] = Int32(0x7FFFFFFF) | self.tag - position = token - self.origin - if position >= Int32(0) and position + Int32(3) < self.length: - values = cutlass.Pointer( - self.indices.iterator + position, dtype=Int32 - ).load(count=4, alignment=4) - for j in cutlass.range_constexpr(4): - index = Int32(values[j]) - rows[j] = (index if index >= Int32(0) else Int32(0x7FFFFFFF)) | self.tag - elif position + Int32(3) >= Int32(0) and position < self.length: - for j in cutlass.range_constexpr(4): - rows[j] = self.mapped_row(token + Int32(j)) - return rows - - -class SparseBatchRouteView: - """Bind live source metadata to the request currently owned by a work tile.""" - - def __init__(self, values, capacity, assume_valid_prefix=False): - self.values = values - self.si, self.ci, self.sl, self.cl = values - self.capacity = capacity - self.assume_valid_prefix = assume_valid_prefix - self.shape = (capacity, self.si.shape[0]) - - def __extract_mlir_values__(self): - return cutlass.extract_mlir_values(self.values) - - def __new_from_mlir_values__(self, values): - return SparseBatchRouteView( - cutlass.new_from_mlir_values(self.values, values), - self.capacity, - self.assume_valid_prefix, - ) - - @cute.jit - def for_request(self, request, primary_length=None, extra_length=None): - # Persistent consumers may reuse lengths loaded at query-tile entry. - row = Int64(request) - return SparseRouteView( - ( - self.si, - self.ci, - Int32(self.sl[row]) - if cutlass.const_expr(primary_length is None) - else primary_length, - Int32(self.cl[row]) - if cutlass.const_expr(extra_length is None) - else extra_length, - row, - ), - self.capacity, - self.assume_valid_prefix, - ) - - @cute.jit - def __getitem__(self, coordinate): - return self.for_request(coordinate[1])[coordinate] - - @cute.jit - def routes_for_tile(self, token, request): - return self.for_request(request).routes_for_tile(token) - - @cute.jit - def mask_for_tile(self, token, request): - return self.for_request(request).mask_for_tile(token) - - @cute.jit - def is_invalid(self, token, request): - return self.mask_for_tile(token, request).is_invalid(token, request) - - -class SparseBatchLengthView: - """Live padded source extent for each virtual request.""" - - def __init__(self, values): - self.values = values - self.sl, self.cl = values - self.shape = self.sl.shape - - def __extract_mlir_values__(self): - return cutlass.extract_mlir_values(self.values) - - def __new_from_mlir_values__(self, values): - return SparseBatchLengthView(cutlass.new_from_mlir_values(self.values, values)) - - @cute.jit - def __getitem__(self, request): - return _sparse_execution_length( - Int32(self.sl[request]), Int32(self.cl[request]) - ) - - -@cute.jit -def _sparse_execution_length(primary_length, extra_length): - primary_tiles = (primary_length + Int32(127)) // Int32(128) - extra_tiles = (extra_length + Int32(127)) // Int32(128) - # Empty requests still execute a masked tile to publish O=0 and LSE=-inf. - return cute.math.max((primary_tiles + extra_tiles) * Int32(128), Int32(1)) - - -@cute.jit -def bind_sparse_views( - si, - ci, - sl, - cl, - capacity: cutlass.Constexpr[int], - *, - request=None, - assume_valid_prefix: cutlass.Constexpr[bool] = False, -): - """Bind source lists to a work queue or to one nonpersistent request.""" - routes = SparseBatchRouteView((si, ci, sl, cl), capacity, assume_valid_prefix) - if cutlass.const_expr(request is None): - return routes, SparseBatchLengthView((sl, cl)) - else: - current = routes.for_request(request) - length = _sparse_execution_length(current.sl, current.cl) - return current, SparseLengthView((length, si.shape[0])) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/kernel.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/kernel.py index 70fa06200b13..ffc5c7fb609c 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 @@ -831,7 +831,6 @@ def __init__( seq_len_q=1, batch_size=1, mask_type: MaskType | str = MaskType.CAUSAL, - device_scales: bool = False, ): """ Parameters @@ -914,12 +913,11 @@ def __init__( raise ValueError("batch_size must be positive") self.batch_size = batch_size self.mask_type = normalize_mask_type(mask_type) - self.device_scales = device_scales self.fuse_sparse_reduction = False self.fuse_sparse_epilogue = False self.output_quant = None - # Separate prepared source lists, rather than a premerged tagged list. - # Both representations reuse the same metadata staging and KV loaders. + # Sparse attention binds separate source lists; dense attention uses + # the ordinary page table. Both reuse the staged KV loaders. self.direct_sparse = False self.sparse_gather_warps = 1 self.sparse_kv_stages = 0 @@ -1090,8 +1088,6 @@ def _make_config(self): and cfg.is_fp8_qkv() and cfg.page_size == 1 and self.num_heads == 128 - # A 64-head TMA box spans whole head groups of the group-major output. - and 64 % (self.num_heads // self.output_quant.num_groups) == 0 and self.seq_len_q == 1 and cfg.mma_pv_tiler_mn == (128, 256) ) @@ -1377,12 +1373,12 @@ def _flatten_query_rows_for_tma(t): flat_o = output_quant_args[0] output_rows = cute.size(flat_o) // (128 * 512) # Group-major [G, T, (H/G) * 512] output viewed as (channel, head in - # group, token, group). A 64-head box covers 64 / (H/G) whole groups - # and fills 64 shared-memory rows in head order, so the staging - # layout matches the flat head-major view. + # group, token, group). A box spans whole groups, or half of the + # single H128 group, retaining the same 64 shared-memory rows. heads_per_group = ( self.output_quant.num_heads // self.output_quant.num_groups ) + box_heads = cutlass.const_expr(min(heads_per_group, 64)) output_view = cute.make_tensor( flat_o.iterator, cute.make_layout( @@ -1397,7 +1393,7 @@ def _flatten_query_rows_for_tma(t): ) tma_desc_o = create_tensor_map_tiled_from_view( output_view, - box_dims=(128, heads_per_group, 1, 64 // heads_per_group), + box_dims=(128, box_heads, 1, 64 // box_heads), stride_order=(0, 1, 2, 3), swizzle=cuda.TensorMapSwizzle.s128b, l2_promotion=cuda.TensorMapL2Promotion.l2_128b, @@ -1578,9 +1574,9 @@ def split_kv_kernel( ) -> None: """MLA decode TS kernel: persistent tile-scheduled execution.""" if cutlass.const_expr(self.direct_sparse): - from ..sparse_views import bind_sparse_views + from ..helpers.sparse_indices import bind_sparse_kv_indices - si, ci, sl, cl, sm, qs, ss, os, sinks = block_split_kvs + si, ci, sl, cl, sm, os, sinks = block_split_kvs capacity = cutlass.const_expr( max( 256, @@ -1599,32 +1595,20 @@ def split_kv_kernel( tile_sched_params.problem_shape_b_fdd, ) request = Int64(row) - page_offsets, cache_seqs = bind_sparse_views( + page_offsets, cache_seqs = bind_sparse_kv_indices( si, ci, sl, cl, capacity, - request=request, + query_row=request, assume_valid_prefix=self.assume_valid_prefix, + has_extra=self.direct_sparse_capacities[1] > 0, ) if cutlass.const_expr(not self.direct_static_scales): - softmax_scale_log2 = ( - Float32(sm[0]) * Float32(qs[0]) * Float32(ss[0]) * LOG2_E - ) - output_scale = Float32(os[0]) * Float32(ss[0]) - else: - if cutlass.const_expr(self.device_scales): - softmax_scale_log2 = cutlass.Float32(block_split_kvs[0]) * LOG2_E - output_scale = cutlass.Float32(block_split_kvs[1]) - if cutlass.const_expr(self.fuse_sparse_epilogue): - # A host-scale caller can bind sinks alone. Device-scale - # metadata starts with QK/PV scales, followed by sinks. - sink_offset = 2 if self.device_scales else 0 - sinks = cute.make_tensor( - block_split_kvs.iterator + sink_offset, - cute.make_layout(self.num_heads), - ) + # Caller scales already include Q/K/V descales. + softmax_scale_log2 = Float32(sm[0]) * LOG2_E + output_scale = Float32(os[0]) cfg = self._make_config() num_query_tiles = self.num_q_tiles use_clc_dynamic = self.is_persistent and not cfg.is_fp8_qkv() @@ -2229,12 +2213,7 @@ def parallel_reduction_kernel( prims.griddepcontrol(kind=prims.GridDepAction.WAIT) atten_sinks = None if cutlass.const_expr(self.fuse_sparse_reduction): - if cutlass.const_expr(self.direct_sparse): - atten_sinks = block_split_kvs[-1] - else: - atten_sinks = cute.make_tensor( - block_split_kvs.iterator + 2, cute.make_layout(self.num_heads) - ) + atten_sinks = block_split_kvs[-1] block_split_kvs = None topology = self._reduction_topology() run_parallel_reduction_kernel( 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 1b3950e29954..4e3d24d11fd6 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 @@ -58,6 +58,7 @@ from ...mask import kv_tile_needs_right_mask from ...tensor_map import transform_ragged_coords +from ..helpers.sparse_indices import SparseKVIndices, SparseSourceTile from ..helpers.gather import ( gather4_cached, load_sparse_rows, @@ -714,10 +715,13 @@ def init_sparse_state(self, stage_info: StageInfo): @producer_work(work_attrs=WorkAttr.AUXILIARY) @cute.jit def bind_sparse_request(self, stage_info: StageInfo): - if cutlass.const_expr(hasattr(self.page_offsets, "for_request")): + if cutlass.const_expr( + isinstance(self.page_offsets, SparseKVIndices) + and not self.page_offsets.is_bound + ): row = stage_info.work_tile.tile_idx[2] - self.request_lengths[0] = Int32(self.page_offsets.sl[row]) - self.request_lengths[1] = Int32(self.page_offsets.cl[row]) + self.request_lengths[0] = Int32(self.page_offsets.primary_lengths[row]) + self.request_lengths[1] = Int32(self.page_offsets.extra_lengths[row]) @producer_work @cute.jit @@ -726,8 +730,11 @@ def load_sparse_offsets(self, stage_info: StageInfo): k_index = stage_info.work_tile.k_index_base + Int32(stage_info.loop_offset) token = k_index * Int32(128) + lane * Int32(4) row = stage_info.work_tile.tile_idx[2] - if cutlass.const_expr(hasattr(self.page_offsets, "for_request")): - routes = self.page_offsets.for_request( + if cutlass.const_expr( + isinstance(self.page_offsets, SparseKVIndices) + and not self.page_offsets.is_bound + ): + routes = self.page_offsets.bind_query_row( row, self.request_lengths[0], self.request_lengths[1] ) else: @@ -852,18 +859,47 @@ def read_page_offset_window( global_k_index = work_tile.k_index_base + local_k_index if cutlass.const_expr(cfg.page_size == 1): cta = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) - new_k = load_sparse_rows( - self.page_offsets, - blk_coord[2], - global_k_index * cfg.tokens_per_k_tile + cta * cfg.tokens_per_k_cta, - fragments=cfg.tokens_per_k_cta // WARP_LANES, - ) - new_v = load_sparse_rows( - self.page_offsets, - blk_coord[2], - global_k_index * cfg.tokens_per_v_tile, - fragments=cfg.tokens_per_v_tile // WARP_LANES, + # K and pending V cover the same source-aligned BK128 tile. + # Bind its source once, not again for every lane/fragment. + routes = ( + self.page_offsets.select_source_tile( + global_k_index * cfg.tokens_per_k_tile, blk_coord[2] + ) + if cutlass.const_expr(isinstance(self.page_offsets, SparseKVIndices)) + else self.page_offsets ) + if cutlass.const_expr(isinstance(routes, SparseSourceTile)): + new_v = load_sparse_rows( + routes, + blk_coord[2], + global_k_index * cfg.tokens_per_v_tile, + fragments=cfg.tokens_per_v_tile // WARP_LANES, + ) + # K is one CTA half of this freshly loaded pending-V window. + # cached_v_pages instead belongs to the previous K tile. + fragments_k = cfg.tokens_per_k_cta // WARP_LANES + new_k = cutlass.Array( + Int32, fragments_k, space=cutlass.AddressSpace.rmem + ) + for fragment in cutlass.range_constexpr(fragments_k): + new_k[fragment] = ( + new_v[fragment] + if cta == Int32(0) + else new_v[fragment + fragments_k] + ) + else: + new_k = load_sparse_rows( + routes, + blk_coord[2], + global_k_index * cfg.tokens_per_k_tile + cta * cfg.tokens_per_k_cta, + fragments=cfg.tokens_per_k_cta // WARP_LANES, + ) + new_v = load_sparse_rows( + routes, + blk_coord[2], + global_k_index * cfg.tokens_per_v_tile, + fragments=cfg.tokens_per_v_tile // WARP_LANES, + ) # Expand only current K/V coordinates once per tile. "Uniform" # means identical across lanes, not consecutive physical KV rows. # Reuse across D slices trades register capacity for fewer shuffles. @@ -2250,10 +2286,13 @@ def init_softmax_state_odd(self, stage_info: StageInfo): @cute.jit def bind_sparse_request(self, stage_info: StageInfo): """Cache live source lengths once per query tile, including graph replay.""" - if cutlass.const_expr(hasattr(self.page_offsets, "for_request")): + if cutlass.const_expr( + isinstance(self.page_offsets, SparseKVIndices) + and not self.page_offsets.is_bound + ): row = stage_info.work_tile.tile_idx[2] - self.request_lengths[0] = Int32(self.page_offsets.sl[row]) - self.request_lengths[1] = Int32(self.page_offsets.cl[row]) + self.request_lengths[0] = Int32(self.page_offsets.primary_lengths[row]) + self.request_lengths[1] = Int32(self.page_offsets.extra_lengths[row]) @producer_work @cute.jit @@ -2496,8 +2535,11 @@ def _load_s_impl( if cutlass.const_expr(cfg.page_size == 1): group_needs_mask = True - if cutlass.const_expr(hasattr(self.page_offsets, "for_request")): - mask_page_offsets = self.page_offsets.for_request( + if cutlass.const_expr( + isinstance(self.page_offsets, SparseKVIndices) + and not self.page_offsets.is_bound + ): + mask_page_offsets = self.page_offsets.bind_query_row( work_tile.tile_idx[2], self.request_lengths[0], self.request_lengths[1] ) else: @@ -2510,8 +2552,10 @@ def _load_s_impl( group_needs_mask = (cached_mask_low | cached_mask_high) != Int32(0) else: raw_mask_routes = ( - mask_page_offsets.mask_for_tile(tile_offset_k, work_tile.tile_idx[2]) - if cutlass.const_expr(hasattr(self.page_offsets, "mask_for_tile")) + mask_page_offsets.select_source_tile( + tile_offset_k, work_tile.tile_idx[2], for_mask=True + ) + if cutlass.const_expr(isinstance(mask_page_offsets, SparseKVIndices)) else self.page_offsets ) mask_routes = ( @@ -3712,11 +3756,13 @@ def tma_store_output( ): # Sparse launch flattens queries into batch rows; each CTA owns 64 heads. head_tile_idx, _, token, _ = stage_info.work_tile.tile_idx - # Group-major output: the 64-head tile starts at this group. + # A tile can span groups or lie inside the single H128 group. heads_per_group = cutlass.const_expr( self.output_quant.num_heads // self.output_quant.num_groups ) - group = head_tile_idx * (64 // heads_per_group) + head_start = head_tile_idx * Int32(64) + group = head_start // heads_per_group + head_in_group = head_start % heads_per_group if prims.elect_sync(): for half in cutlass.range_constexpr(2): prims.cp_async_bulk_tensor_global_shared_cta( @@ -3724,7 +3770,7 @@ def tma_store_output( self.smem_o.subview( stage_info.stage_idx * OUTPUT_TMA_STAGE_BYTES + half * 8192 ), - (iter_n * 256 + half * 128, 0, token, group), + (iter_n * 256 + half * 128, head_in_group, token, group), ) prims.cp_async_bulk_commit_group() # Release the stage only after TMA has consumed its shared source. 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 73aef9cf8bc3..7657de52608d 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 @@ -174,6 +174,7 @@ class MlaConfig: fuse_sparse_reduction: bool = False sparse_direct: bool = False sparse_static_scales: bool = False + sparse_valid_prefix: bool = False sparse_direct_capacities: tuple[int, int] = (0, 0) page_offsets_num_warps: int = 1 scheduler_num_warps: int = 1 @@ -861,11 +862,9 @@ def resolve_runtime_cluster_reduction_mode( cfg: MlaConfig | None, *, reduction_mode: str | None, - hardware_info, - stream=None, log=None, ) -> str | None: - """Query cluster occupancy and return the final runtime reduction mode.""" + """Use cached device occupancy to resolve the runtime reduction mode.""" if ( cfg is None @@ -876,7 +875,9 @@ def resolve_runtime_cluster_reduction_mode( cluster_shape = cluster_reduction_cluster_shape(cfg) cluster_size = cluster_shape[0] * cluster_shape[1] * cluster_shape[2] - active_clusters = hardware_info.get_max_active_clusters(cluster_size, stream) + from ...._hardware import get_max_active_clusters + + active_clusters = get_max_active_clusters(cluster_size) cluster_count = cluster_reduction_cluster_count(cfg) if log is not None: log( @@ -1723,6 +1724,7 @@ def configure_sparse_mla( *, finalize_output=False, direct_inputs=False, + assume_valid_prefix=False, static_scales=False, source_capacities=(0, 0), ) -> MlaConfig: @@ -1801,6 +1803,7 @@ def configure_sparse_mla( defer_sparse_max_update=profile.defer_max_update, sparse_direct=direct_inputs, sparse_static_scales=static_scales, + sparse_valid_prefix=assume_valid_prefix, sparse_direct_capacities=source_capacities if direct_inputs else (0, 0), fuse_sparse_epilogue=direct_finish, fuse_sparse_cluster_epilogue=cluster_finish, 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 f7c476f7a3c9..130297f3b4a9 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 @@ -36,7 +36,6 @@ ) from cutlass.experimental import cuda from cutlass.experimental import primitives as prims -from cutlass.utils.static_persistent_tile_scheduler import WorkTileInfo from cutlass.experimental.task_scheduling.enums import PipelineType, SignalingThreads from cutlass.experimental.task_scheduling.memory import ( @@ -49,6 +48,7 @@ PipelineConfig, TileSchedulerConfig, WorkQueue, + WorkTileInfo, ) from cutlass.experimental.task_scheduling.task_manager import TaskManager @@ -369,19 +369,6 @@ def _default_scales(scale_softmax_log2, output_scale): return scale_softmax_log2, output_scale -def _atten_sinks_view(cfg, epilogue_params): - """Bind only sinks; packed scale/count metadata remains a separate input.""" - if epilogue_params is None or not ( - cfg.fuse_sparse_epilogue or cfg.fuse_sparse_cluster_epilogue - ): - return None - if cfg.sparse_direct: - return epilogue_params - return cute.make_tensor( - epilogue_params.iterator + 2, cute.make_layout(cfg.logical_num_heads_q) - ) - - def _check_persistent_scheduler_modes( use_clc_dynamic_scheduler: bool, use_static_persistent_scheduler: bool, @@ -834,8 +821,7 @@ def build_throughput_latency_mla_task_manager( inst_id=0, scale_softmax_log2=scale_softmax_log2, output_scale=output_scale, - sparse_epilogue_params=sparse_epilogue_params, - atten_sinks=_atten_sinks_view(cfg, sparse_epilogue_params), + atten_sinks=sparse_epilogue_params, output_quant=output_quant, output_quant_args=output_quant_args, o_tensor=o_tensor, @@ -856,8 +842,7 @@ def build_throughput_latency_mla_task_manager( inst_id=1, scale_softmax_log2=scale_softmax_log2, output_scale=output_scale, - sparse_epilogue_params=sparse_epilogue_params, - atten_sinks=_atten_sinks_view(cfg, sparse_epilogue_params), + atten_sinks=sparse_epilogue_params, output_quant=output_quant, output_quant_args=output_quant_args, o_tensor=o_tensor, @@ -1439,8 +1424,7 @@ def _make_single_kv_pipe_task_graph( inst_id=1, scale_softmax_log2=scale_softmax_log2, output_scale=output_scale, - sparse_epilogue_params=sparse_epilogue_params, - atten_sinks=_atten_sinks_view(cfg, sparse_epilogue_params), + atten_sinks=sparse_epilogue_params, output_quant=output_quant, output_quant_args=output_quant_args, o_tensor=o_tensor, @@ -1693,7 +1677,6 @@ def __init__( explicit_split_kv: int | None = None, explicit_persistent: bool | None = None, mask_type: MaskType | str = MaskType.CAUSAL, - device_scales: bool = False, sparse_profile=None, finalize_output: bool = False, ): @@ -1724,13 +1707,13 @@ def __init__( self.explicit_split_kv = explicit_split_kv self.explicit_persistent = explicit_persistent self.mask_type = normalize_mask_type(mask_type) - self.device_scales = device_scales self.sparse_profile = sparse_profile self.finalize_output = finalize_output self.output_quant = None - # Separate prepared source lists, rather than a premerged tagged list. - # Both representations reuse the same metadata staging and KV loaders. + # Sparse attention binds separate source lists; dense attention uses + # the ordinary page table. Both reuse the staged KV loaders. self.direct_sparse = False + self.assume_valid_prefix = False self.direct_static_scales = False self.direct_sparse_capacities = (0, 0) @@ -1873,6 +1856,7 @@ def _make_config(self): self.sparse_profile, finalize_output=self.finalize_output, direct_inputs=self.direct_sparse, + assume_valid_prefix=self.assume_valid_prefix, static_scales=self.direct_static_scales, source_capacities=self.direct_sparse_capacities, ) @@ -2283,13 +2267,7 @@ def __call__( reducer_cfg, topology, elements = self._parallel_reduction_setup(cfg) atten_sinks = None if cutlass.const_expr(cfg.fuse_sparse_reduction): - if cutlass.const_expr(cfg.sparse_direct): - atten_sinks = block_split_kvs[-1] - else: - atten_sinks = cute.make_tensor( - block_split_kvs.iterator + 2, - cute.make_layout(cfg.logical_num_heads_q), - ) + atten_sinks = block_split_kvs[-1] reduction_grid, reduction_cluster = ( parallel_gmem_reduction_launch_shape( reducer_cfg, @@ -2386,40 +2364,34 @@ def dense_kernel( softmax_scale_log2: cutlass.Float32, output_scale: cutlass.Float32, tile_sched_params: object, - sparse_scale_params: object, + sparse_inputs: object, output_quant_args=None, ): """Execute one flat-Q, batch, KV-split, and V head-dimension tile.""" cfg = self._make_config() if cutlass.const_expr(cfg.sparse_direct): - from ..sparse_views import bind_sparse_views + from ..helpers.sparse_indices import bind_sparse_kv_indices - si, ci, sl, cl, sm, qs, ss, os, sinks = sparse_scale_params + si, ci, sl, cl, sm, os, sinks = sparse_inputs request = None if cutlass.const_expr(cfg.use_persistent_scheduler != 1): request = cutlass.Int64(cute.arch.block_idx()[2]) - page_offsets, cache_seqs = bind_sparse_views( + page_offsets, cache_seqs = bind_sparse_kv_indices( si, ci, sl, cl, cfg.total_kv_tiles * cfg.tile_size_kv, - request=request, + query_row=request, + assume_valid_prefix=cfg.sparse_valid_prefix, + has_extra=cfg.sparse_direct_capacities[1] > 0, ) if cutlass.const_expr(not cfg.sparse_static_scales): - softmax_scale_log2 = ( - cutlass.Float32(sm[0]) - * cutlass.Float32(qs[0]) - * cutlass.Float32(ss[0]) - * LOG2_E - ) - output_scale = cutlass.Float32(os[0]) * cutlass.Float32(ss[0]) + # Caller scales already include Q/K/V descales. + softmax_scale_log2 = cutlass.Float32(sm[0]) * LOG2_E + output_scale = cutlass.Float32(os[0]) sparse_epilogue_params = sinks - elif cutlass.const_expr(self.device_scales): - softmax_scale_log2 = cutlass.Float32(sparse_scale_params[0]) * LOG2_E - output_scale = cutlass.Float32(sparse_scale_params[1]) - sparse_epilogue_params = sparse_scale_params else: sparse_epilogue_params = None 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 e09393d91303..590173465f44 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 @@ -537,23 +537,26 @@ def _producer_load_page_offsets( ) tile_idx = global_kv_tile_idx(cfg, local_tile_idx, seq_len_kv, cta_idx_kv) tile_start = tile_idx * pages_per_tile - routes = self.page_offsets.routes_for_tile(tile_start, batch_idx) + routes = self.page_offsets.select_source_tile(tile_start, batch_idx) position = lane_idx * Int32(4) smem_offset = ( stage_info.stage_idx * Int32(cfg.page_offsets_entries_per_stage) + position ) - if cutlass.const_expr(cfg.tile_size_kv < 128): - if position < pages_per_tile: - quad = routes.mapped_quad(tile_start + position) + if position < pages_per_tile: + # Source selection is tile-wide. Hoist the bounds decision out + # of individual quads so full tiles take one uniform branch; + # index values still mask holes even in an in-bounds tile. + if tile_start + pages_per_tile <= routes.origin + routes.length: + quad = routes.load_full_quad(tile_start + position) + self._smem_page_offsets.data_ptr(smem_offset).store( + quad.load(0, 4), alignment=16 + ) + else: + quad = routes.load_quad(tile_start + position) self._smem_page_offsets.data_ptr(smem_offset).store( quad.load(0, 4), alignment=16 ) - else: - quad = routes.mapped_quad(tile_start + position) - self._smem_page_offsets.data_ptr(smem_offset).store( - quad.load(0, 4), alignment=16 - ) cute.arch.fence_view_async_shared() return for fragment in cutlass.range_constexpr((cfg.pages_per_kv_tile + 31) // 32): @@ -672,6 +675,8 @@ def read_offsets( # full valid prefix can still contain random rows. Eight quads # (32 rows) are needed for bulk; smaller issuer shares use gather4. if cutlass.const_expr(groups_per_warp >= 8): + # A chunk stays within one source. Invalid coordinates retain + # its tag because the first coordinate selects the shared map. first = Int32( prims.shfl_sync( thread_mask=0xFFFFFFFF, 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 c221d42f6102..cf547165e9e1 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 @@ -116,7 +116,6 @@ class TmemCorrResource(MlaResource): inst_id: cutlass.Constexpr[int] = 0 scale_softmax_log2: Float32 = None output_scale: Float32 = None - sparse_epilogue_params: object = None atten_sinks: object = None output_quant: cutlass.Constexpr = None output_quant_args: object = None @@ -565,14 +564,7 @@ def _store_lse_to_gmem( if selected_max == neg_max_f32(): lse_val = Float32(-Float32.inf) if cutlass.const_expr(self.cfg.fuse_sparse_epilogue): - if cutlass.const_expr(self.cfg.sparse_direct): - count = Float32(selected_max != neg_max_f32()) - else: - count = Float32( - self.sparse_epilogue_params[ - 2 + self.cfg.logical_num_heads_q + batch_idx - ] - ) + count = Float32(selected_max != neg_max_f32()) lse_val = ( lse_val * Float32(LN_2) if count > Float32(0) @@ -1462,7 +1454,6 @@ def _compute_tail_softmax_scales( ) norm_lanes = 16 if cfg.kernel_variant == "keeps_mma_ab" else 4 if lane_idx < Int32(norm_lanes): - batch_idx = batch_idx_for_stage_cfg(self.batch_idx, cfg, stage_info) logical_head = ( cta_idx_q_for_stage(self.cta_idx_q, stage_info) * cfg.tile_size_q @@ -1475,14 +1466,7 @@ def _compute_tail_softmax_scales( ) ) ) - if cutlass.const_expr(cfg.sparse_direct): - count = Float32(final_max[scale_idx] != neg_max_f32()) - else: - count = Float32( - self.sparse_epilogue_params[ - 2 + cfg.logical_num_heads_q + batch_idx - ] - ) + count = Float32(final_max[scale_idx] != neg_max_f32()) norm_scale = Float32(0) if logical_head < cfg.logical_num_heads_q and count > Float32(0): sink = Float32(self.atten_sinks[logical_head]) 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 1c29fd10f8c4..9831ab71e5d3 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 @@ -472,7 +472,9 @@ def update_softmax( tile_offset_k = tile_idx * Int32(cfg.tile_size_kv) next_tile_offset_k = tile_offset_k + Int32(cfg.tile_size_kv) raw_mask_routes = ( - self.page_offsets.mask_for_tile(tile_offset_k, batch_idx) + self.page_offsets.select_source_tile( + tile_offset_k, batch_idx, for_mask=True + ) if cutlass.const_expr(cfg.sparse_direct) else self.page_offsets ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/separate_reduction.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/separate_reduction.py index fe8706c9e4a2..152e956a40a5 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/separate_reduction.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/separate_reduction.py @@ -26,9 +26,9 @@ import cutlass import cutlass.cute as cute -from cutlass import BFloat16, Float16, Float32, Int32, Int64 +from cutlass import BFloat16, Float16, Float32, Int32 -from .mla_decode.helpers.constants import LN_2, LOG2_E +from .mla_decode.helpers.constants import LOG2_E @cute.jit @@ -112,112 +112,3 @@ def attention_sink_scale(log2_lse, sink): inverse = cute.math.rcp(Float32(1) + z, approx=True) scale = z * inverse if delta > Float32(0) else inverse return scale - - -class MergeMlaSourceStates: - def __init__(self, independent_sources: bool, output_quant=None): - self.independent_sources = independent_sources - self.output_quant = output_quant - - @cute.jit - def __call__( - self, - part_s, - part_c, - lse_s, - lse_c, - count_s, - count_c, - sinks, - out, - lse, - stream, - output_quant_args=None, - ): - self.finish( - part_s, - part_c, - lse_s, - lse_c, - count_s, - count_c, - sinks, - out, - lse, - output_quant_args, - ).launch( - grid=((out.shape[0] * out.shape[1] + 3) // 4, 1, 1), - block=(128, 1, 1), - stream=stream, - ) - - @cute.kernel - def finish( - self, - part_s, - part_c, - lse_s, - lse_c, - count_s, - count_c, - sinks, - out, - lse, - output_quant_args=None, - ): - block, _, _ = cute.arch.block_idx() - tid = cute.arch.thread_idx()[0] - lane = tid % 32 - # One warp handles one output head. Group four heads per CTA to - # amortize launch overhead and avoid repeating normalization 128 times. - row = Int64(block) * 4 + Int64(tid // 32) - q = row // out.shape[1] - h = row % out.shape[1] - if row < out.shape[0] * out.shape[1]: - have_s = Int32(count_s[q]) > 0 - have_c = (Int32(count_c[q]) > 0) if self.independent_sources else False - l0 = Float32(lse_s[q, h]) if have_s else Float32(-Float32.inf) - l1 = Float32(lse_c[q, h]) if have_c else Float32(-Float32.inf) - ws = Float32(0) - wc = Float32(0) - public_lse = Float32(-Float32.inf) - if have_s or have_c: - kv_lse, ws, wc = merge_log2_lse(l0, l1) - public_lse = kv_lse * Float32(LN_2) - sink_scale = attention_sink_scale(kv_lse, Float32(sinks[h])) - ws *= sink_scale - wc *= sink_scale - width = cutlass.const_expr(4 if self.output_quant is not None else 2) - for chunk in cutlass.range_constexpr(512 // (32 * width)): - column = lane * width + chunk * (32 * width) - values = cutlass.Array(Float32, width, space=cutlass.AddressSpace.rmem) - for j in cutlass.range_constexpr(width): - values[j] = Float32(0) - if have_s: - offset = q * part_s.stride[0] + h * part_s.stride[1] + column - pair = cutlass.Pointer( - part_s.iterator + offset, dtype=part_s.element_type - ).load(count=width, alignment=width * 2) - for j in cutlass.range_constexpr(width): - values[j] = Float32(pair[j]) * ws - if have_c: - offset = q * part_c.stride[0] + h * part_c.stride[1] + column - pair = cutlass.Pointer( - part_c.iterator + offset, dtype=part_c.element_type - ).load(count=width, alignment=width * 2) - for j in cutlass.range_constexpr(width): - values[j] += Float32(pair[j]) * wc - if cutlass.const_expr(self.output_quant is not None): - self.output_quant.store(output_quant_args, values, row, column) - else: - packed = cutlass.Array( - out.element_type, 2, space=cutlass.AddressSpace.rmem - ) - for j in cutlass.range_constexpr(2): - packed[j] = values[j].to(out.element_type) - offset = q * out.stride[0] + h * out.stride[1] + column - cutlass.Pointer( - out.iterator + offset, dtype=out.element_type - ).store(packed.load(0, 2), alignment=4) - if lane == 0: - lse[q, h] = public_lse 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 70c3d0e6423d..e64fc538c768 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py @@ -21,11 +21,15 @@ import math import struct import functools -from typing import Any, Literal, NamedTuple, Optional, cast +from typing import Any, Literal, Optional, cast + +from ._compile_cache import build_and_load_attention_kernel +from ._hardware import get_max_active_clusters import torch from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api + from .decode import ( _WorkspaceSection, _align_up, @@ -86,7 +90,6 @@ class _MLADecodeCompileSpec: has_kernel_workspace: bool split_kv: int kernel: Any = field(compare=False, hash=False, repr=False) - device_scales: bool = False @dataclass(frozen=True) @@ -139,7 +142,6 @@ class _MLARuntime: bmm1_scale: float bmm2_scale: float extra_cache: Optional[torch.Tensor] = None - scale_params: Optional[torch.Tensor] = None sparse_inputs: Optional[tuple] = None output_quant_args: Optional[tuple] = None @@ -742,8 +744,6 @@ def _resolve_mla_decode_launch_spec( max_kv_len = _validate_mla_max_kv_len(max_kv_len, "max_kv_len") import cutlass - import cutlass.utils as cutlass_utils - from cuda.bindings import driver as cuda_drv from .kernels.mla_decode.kernel_policy import ( resolve_mla_kernel_policy, @@ -779,16 +779,8 @@ def _resolve_mla_decode_launch_spec( output_dtype_name = _kernel_dtype_name(output_dtype_key) with torch.cuda.device(device_index): - plan_stream = cuda_drv.CUstream( - torch.cuda.current_stream(device_index).cuda_stream - ) - hardware_info = cutlass_utils.HardwareInfo(device_index) - max_active_one_cta_clusters = hardware_info.get_max_active_clusters( - 1, plan_stream - ) - max_active_two_cta_clusters = hardware_info.get_max_active_clusters( - 2, plan_stream - ) + max_active_one_cta_clusters = get_max_active_clusters(1) + max_active_two_cta_clusters = get_max_active_clusters(2) one_cta_launch_shape = resolve_auto_flat_query_launch_shape( num_heads_q=num_heads, seq_len_q=seq_len_q, @@ -854,8 +846,6 @@ def _resolve_mla_decode_launch_spec( reduction_mode = resolve_runtime_cluster_reduction_mode( decision.config, reduction_mode=None, - hardware_info=hardware_info, - stream=plan_stream, ) kernel = ThroughputLatencyMlaDecodeTs( batch_size=batch_size, @@ -1034,7 +1024,6 @@ def _make_mla_decode_compile_spec( output_dtype_key: str, max_seq_len_q: int, packed_query: bool, - device_scales: bool = False, ) -> _MLADecodeCompileSpec: """Keep policy resolution plan-specific and JIT identity batch-free.""" @@ -1052,7 +1041,6 @@ 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, - device_scales=device_scales, ) @@ -1186,13 +1174,7 @@ def _get_compiled_mla_decode( stride_order=(0,), assumed_align=4, ) - scale_params_fake = ( - cute.runtime.make_fake_compact_tensor( - cutlass.Float32, (cute.sym_int(),), stride_order=(0,), assumed_align=4 - ) - if compile_spec.device_scales - else None - ) + sparse_inputs_fake = None if getattr(kernel, "direct_sparse", False): def sparse_indices(capacity): @@ -1209,15 +1191,13 @@ def sparse_indices(capacity): sink = cute.runtime.make_fake_compact_tensor( cutlass.Float32, (num_heads,), stride_order=(0,), assumed_align=4 ) - scale_params_fake = ( + sparse_inputs_fake = ( sparse_indices(kernel.direct_sparse_capacities[0]), sparse_indices(kernel.direct_sparse_capacities[1]), cache_seqs_fake, cache_seqs_fake, scalar, scalar, - scalar, - scalar, sink, ) stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) @@ -1226,8 +1206,8 @@ def sparse_indices(capacity): # Task objects carry loop-local state through generated control flow, so # select the public staged frontend for this compilation. - with torch.cuda.device(device_index): - compiled = cute.compile[cute.FrontendNext]( + def compile_kernel(): + return cute.compile[cute.FrontendNext]( kernel, q_latent_fake, q_rope_fake, @@ -1240,13 +1220,23 @@ def sparse_indices(capacity): cutlass.Int32(compile_spec.split_kv), cache_seqs_fake, qo_indptr_fake, - scale_params_fake, + sparse_inputs_fake, cutlass.Float32(1.0), cutlass.Float32(1.0), stream_fake, output_quant_args=output_quant_args, options=_COMPILE_OPTIONS, ) + + with torch.cuda.device(device_index): + compiled = build_and_load_attention_kernel( + "mla_decode", + "main", + compile_spec, + _COMPILE_OPTIONS, + compile_kernel, + extra=(), + ) return compiled @@ -1463,9 +1453,7 @@ def _launch_mla_decode( split_kv, seq_lens, qo_indptr, - runtime.sparse_inputs - if runtime.sparse_inputs is not None - else runtime.scale_params, + runtime.sparse_inputs, runtime.bmm1_scale, runtime.bmm2_scale, **( @@ -1957,78 +1945,6 @@ def batch_mla_decode_with_paged_kv_cache( ] -class SparseMLAPreparedMetadata(NamedTuple): - """Caller-owned storage-row indices and prepared attention metadata. - - Indices are int32 [rows, capacity], with -1 invalid; lengths are int32 - [rows]. Primary and extra indices address their own flattened native pool. - No block-table lookup or page-stride conversion occurs in run. - - Combined-route schedules also consume routes [passes, rows, route_capacity], - execution_lengths/valid_counts [passes, rows], and FP32 scale_params - [passes, 2 + heads + rows]. Route bit 31 selects the extra pool; bits 0:31 - encode the storage row, with 0x7fffffff invalid. Each source's execution - span is rounded to 128. Empty rows execute one masked slot. Scale rows - contain [QK scale, PV scale, sinks[heads], valid_counts[rows]]. Two passes - are required when the sources use independent KV descales. - - Preparation must refresh all dependent fields when routes, lengths, - scales or sinks change. Buffers must remain alive and unaliased with - writable attention buffers through graph replay. Applications may supply - these buffers from any preparer satisfying this contract. - """ - - indices: torch.Tensor - lengths: torch.Tensor - extra_indices: torch.Tensor | None = None - extra_lengths: torch.Tensor | None = None - routes: torch.Tensor | None = None - execution_lengths: torch.Tensor | None = None - valid_counts: torch.Tensor | None = None - scale_params: torch.Tensor | None = None - - -def _fake_sparse_tensor(dtype, shape): - """Compact, four-byte-aligned ABI shared by sparse finishing kernels.""" - import cutlass.cute as cute - - return cute.runtime.make_fake_compact_tensor( - dtype, - shape, - stride_order=tuple(reversed(range(len(shape)))), - assumed_align=4, - ) - - -@functools.cache -def _compile_finish(device_index, heads, independent, output_quant=None): - import cutlass - import cutlass.cute as cute - from .kernels.separate_reduction import MergeMlaSourceStates - - rows = cute.sym_int() - - with torch.cuda.device(device_index): - partial = _fake_sparse_tensor(cutlass.BFloat16, (rows, heads, 512)) - lse = _fake_sparse_tensor(cutlass.Float32, (rows, heads)) - lens = _fake_sparse_tensor(cutlass.Int32, (rows,)) - return cute.compile[cute.FrontendNext]( - MergeMlaSourceStates(independent, output_quant), - partial, - partial, - lse, - lse, - lens, - lens, - _fake_sparse_tensor(cutlass.Float32, (heads,)), - partial, - lse, - cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), - output_quant_args=_fake_output_quant_args(output_quant), - options=_COMPILE_OPTIONS, - ) - - def _make_sparse_output_quant(heads, block_size, scale_format, groups): if block_size is None: if groups is not None or scale_format != "ue8m0": @@ -2067,7 +1983,9 @@ def _fake_output_quant_args(quant): stride_order=(0,), assumed_align=16, ), - _fake_sparse_tensor(cutlass.Int32, (rows,)), + cute.runtime.make_fake_compact_tensor( + cutlass.Int32, (rows,), stride_order=(0,), assumed_align=4 + ), cute.runtime.make_fake_compact_tensor( cutlass.Float32, (cute.sym_int(), 64), stride_order=(1, 0), assumed_align=16 ), @@ -2129,15 +2047,28 @@ def _prepare_sparse_output_quant( out.shape != shape or out.dtype != torch.float8_e4m3fn or out.device != device - or out.stride(2) != 1 - or out.stride(0) != channels - or (rows > 1 and out.stride(1) != rows * channels) or out.data_ptr() % 16 ): raise ValueError( f"out must be the [T, G, C] view of a contiguous group-major CUDA " f"E4M3 [G, T, C] tensor ({shape}), aligned to 16 bytes" ) + physical_out = out.transpose(0, 1) + if rows and quant.num_groups > 1: + # A live prefix can retain a larger graph bucket's group stride. + capacity, remainder = divmod(physical_out.stride(0), channels) + if remainder or capacity < rows or capacity * quant.num_heads >= 2**31: + raise ValueError("out must have valid group-major token capacity") + try: + physical_out = physical_out.as_strided( + (quant.num_groups, capacity, channels), physical_out.stride() + ) + except RuntimeError as exc: + raise ValueError( + "out must view complete group-major backing storage" + ) from exc + if not physical_out.is_contiguous(): + raise ValueError("out must have contiguous physical group-major storage") scale_dtype = torch.int32 if quant.ue8m0 else torch.float32 scale_shape = (rows, quant.num_groups, columns) scale_stride = (1, columns * padded_rows, padded_rows) @@ -2161,7 +2092,7 @@ def _prepare_sparse_output_quant( return ( out, scales, - (out.as_strided((out.numel(),), (1,)), positions.view(-1), cos_sin, scales), + (physical_out.view(-1), positions.view(-1), cos_sin, scales), ) @@ -2259,7 +2190,6 @@ def _resolve_sparse_mla_plan( seq_len_q=1, batch_size=max_rows, mask_type="dense", - device_scales=True, ) kernel.sparse_gather_warps = profile.gather_issue_warps kernel.sparse_kv_stages = profile.kv_pipeline_stages @@ -2303,7 +2233,6 @@ def _resolve_sparse_mla_plan( tile_size_q=tile, sparse_profile=profile, mask_type="dense", - device_scales=True, ) core_bytes = compute_workspace_size( cfg=kernel._make_config(), @@ -2318,21 +2247,14 @@ def _resolve_sparse_mla_plan( qk_rope_head_dim=0, page_size=1, q_dtype_key=str(q_data_type).removeprefix("torch."), - output_dtype_key="bfloat16", + output_dtype_key="float8_e4m3fn" if output_quant is not None else "bfloat16", max_seq_len_q=1, packed_query=False, - device_scales=True, ) - # Only a second source can require an independent attention pass. - source_slots = 2 if max_extra_topk else 1 sections = {} byte_end = 0 for name, shape, dtype in ( ("core", (core_bytes,), torch.int8), - ("partial", (source_slots, max_rows, 1, num_heads, 512), torch.bfloat16), - # Both source views must satisfy the core's 16-byte base alignment, - # including H6/H12 and odd maximum query counts. - ("lse", (source_slots, (max_rows + 3) // 4 * 4, 1, num_heads), torch.float32), ("public_lse", (max_rows, num_heads), torch.float32), ): sections[name], byte_end = _append_workspace_section(byte_end, shape, dtype) @@ -2340,12 +2262,13 @@ def _resolve_sparse_mla_plan( class BatchSparseMLADecodePagedTSWrapper: - """Plan native BF16/E4M3 attention over one or two pools, then bind metadata. + """Plan native BF16/E4M3 attention, then bind one or two source index lists. - Q is [B,Sq,H,512] or packed [total_q,H,512]. Caller-prepared metadata - contains int32 [rows, capacity] storage-row indices and live lengths. - Indices already include physical page strides; -1 is masked. Both source - lists participate in one attention distribution. Packed KV is unsupported. + Q is [B,Sq,H,512] or packed [total_q,H,512]. Each source supplies int32 + [rows, capacity] storage-row indices and int32 [rows] live lengths directly + to run(). Indices include physical page strides; -1 is masked. Both source + lists participate in one attention distribution. The caller folds the + shared KV descale into the effective softmax and output scales. ``output_quant_block_size=32`` or ``128`` fuses conjugate output RoPE on D[448:512] and E4M3 quantization after the final attention reduction. Q/KV @@ -2354,9 +2277,9 @@ class BatchSparseMLADecodePagedTSWrapper: This option does not change attention dispatch or softmax correction. Planning with ``assume_valid_prefix=True`` promises that every index before - each live source length is valid (no -1 holes). Direct FP8 2CTA kernels - then derive masks from lengths without reading indices or issuing ballots. - Other kernels retain their generic mask path; no gain was established. + each live source length is valid (no -1 holes). Kernels can then derive + masks from lengths and omit per-index hole checks. This does not imply + that the physical KV rows are consecutive. ``run(validate=True)`` checks the promise; graph replay with validation disabled must preserve it. Tail entries beyond the lengths are ignored. @@ -2404,7 +2327,7 @@ def plan( MXFP8; FP32 scales use unrounded amax/448. Both use an amax floor of 1e-4. A block never crosses a head boundary. """ - (device, profile, spec, sections, byte_end, capacity, output_quant) = ( + (device, profile, spec, sections, byte_end, _, output_quant) = ( _resolve_sparse_mla_plan( device, batch_size, @@ -2427,15 +2350,6 @@ def plan( family = profile.family splits = profile.split_kv max_rows = batch_size * max_seq_len_q - # Policy resolves supported geometry before compilation. The fused - # kernel handles one split or a cluster reduction; other splits use - # the separate reducer. Independent source scales use the unfused path. - use_cluster_epilogue = ( - family == "swap" and splits > 1 and profile.reduction == "cluster" - ) - use_fused_epilogue = splits == 1 or use_cluster_epilogue - use_fused_reduction = splits > 1 and profile.reduction == "gmem_separate" - use_direct = profile.direct_inputs workspace = self._workspace_buffer if workspace is None: workspace = torch.empty(byte_end, device=device, dtype=torch.uint8) @@ -2444,40 +2358,30 @@ def plan( name: _workspace_section_view(workspace, section) for name, section in sections.items() } + # Finalize the joint distribution in the main kernel or its existing + # split-KV reducer. No per-source outputs or finishing launch are needed. + kernel = copy.copy(kernel) + kernel.output_quant = output_quant + kernel.assume_valid_prefix = assume_valid_prefix + if family == "2cta": + kernel.fuse_sparse_epilogue = splits == 1 + kernel.fuse_sparse_reduction = splits > 1 + else: + kernel.finalize_output = True + kernel.direct_sparse = True + kernel.direct_sparse_capacities = (max_topk, max_extra_topk) + spec = replace( + spec, kernel=kernel, kernel_signature=_mla_kernel_compile_signature(kernel) + ) compiled = _get_compiled_mla_decode(spec) - compiled_fused = None - compiled_static = None - if use_fused_epilogue or use_fused_reduction: - fused_kernel = copy.copy(kernel) - fused_kernel.output_quant = output_quant - if family == "2cta": - fused_kernel.fuse_sparse_epilogue = use_fused_epilogue - fused_kernel.fuse_sparse_reduction = use_fused_reduction - else: - fused_kernel.finalize_output = True - fused_kernel.direct_sparse = use_direct - if family == "2cta": - fused_kernel.assume_valid_prefix = assume_valid_prefix - fused_kernel.direct_sparse_capacities = (max_topk, max_extra_topk) - fused_spec = replace( + static_kernel = copy.copy(kernel) + static_kernel.direct_static_scales = True + compiled_static = _get_compiled_mla_decode( + replace( spec, - kernel=fused_kernel, - kernel_signature=_mla_kernel_compile_signature(fused_kernel), + kernel=static_kernel, + kernel_signature=_mla_kernel_compile_signature(static_kernel), ) - compiled_fused = _get_compiled_mla_decode(fused_spec) - if use_direct: - static_kernel = copy.copy(fused_kernel) - static_kernel.direct_static_scales = True - compiled_static = _get_compiled_mla_decode( - replace( - fused_spec, - kernel=static_kernel, - kernel_signature=_mla_kernel_compile_signature(static_kernel), - ) - ) - finish = tuple( - _compile_finish(device.index, num_heads, independent, output_quant) - for independent in (False, True) ) self._state = dict( device=device, @@ -2489,7 +2393,6 @@ def plan( dtype=q_data_type, ks=max_topk, kc=max_extra_topk, - capacity=capacity, kv_layout=kv_layout, has_sinks=has_sinks, return_lse=return_lse, @@ -2497,17 +2400,13 @@ def plan( buffers=buffers, workspace=workspace, compiled=compiled, - compiled_fused=compiled_fused, compiled_static=compiled_static, - direct_inputs=use_direct, assume_valid_prefix=assume_valid_prefix, - finish=finish, splits=splits, default_cl=torch.full( (max_rows,), max_extra_topk, device=device, dtype=torch.int32 ), dummy_indices=torch.zeros((max_rows, 1), device=device, dtype=torch.int32), - dummy_scales=torch.ones((1, 2), device=device, dtype=torch.float32), dummy_cache=torch.zeros((1, 1, 512), device=device, dtype=q_data_type), default_sinks=torch.full( (num_heads,), -torch.inf, device=device, dtype=torch.float32 @@ -2625,73 +2524,19 @@ def _source_metadata( ) return indices, lengths - def _bind_metadata(self, metadata, *, rows, passes, direct, primary_lengths): - state = self._state - buffers = state["buffers"] - prepared_buffers = {} - for key, value, dtype, shape in ( - ( - "routes", - metadata.routes, - torch.int32, - (passes, rows, state["capacity"]), - ), - ( - "lengths", - metadata.execution_lengths, - torch.int32, - (passes, rows), - ), - ( - "counts", - metadata.valid_counts, - torch.int32, - (passes, rows), - ), - ( - "scales", - metadata.scale_params, - torch.float32, - (passes, 2 + state["heads"] + rows), - ), - ): - if value is None: - if not direct: - raise ValueError(f"prepared {key} are required by this schedule") - continue - if ( - value.dtype != dtype - or value.device != state["device"] - or value.shape != shape - or not value.is_contiguous() - ): - raise ValueError(f"prepared {key} must be contiguous {dtype}{shape}") - prepared_buffers[key] = value - buffers = ( - buffers - | dict( - routes=state["dummy_indices"][:rows].unsqueeze(0), - lengths=primary_lengths.unsqueeze(0), - counts=primary_lengths.unsqueeze(0), - scales=state["dummy_scales"], - ) - | prepared_buffers - ) - return buffers - @flashinfer_experimental_api def run( self, query, kv_cache, - metadata: SparseMLAPreparedMetadata, + indices, + lengths, extra_kv_cache=None, *, + extra_indices=None, + extra_lengths=None, qo_indptr=None, softmax_scale=512**-0.5, - q_scale=1.0, - kv_scale=1.0, - extra_kv_scale=None, output_scale=1.0, sinks=None, out=None, @@ -2701,13 +2546,37 @@ def run( out_scale=None, validate=True, ): - """Launch with caller-prepared metadata, without index conversion. - - Caller-owned metadata and plan workspace must outlive graph replay. Q and - native KV are used directly; flattening padded pages only creates a - tensor view. Required attention reductions/finishing remain included. - Metadata is an unchecked consistency contract: its packed routes, - counts and scales must agree with the supplied indices and scalars. + """Attend to caller-selected storage rows in one or two native KV pools. + + ``indices`` and ``extra_indices`` are CUDA INT32 [T,K] tensors, with + T equal to the query token count and K equal to the corresponding + planned max_topk/max_extra_topk. The last dimension is contiguous; + row strides may include padding. Each index addresses its own pool's + flattened storage, including physical page strides. The caller performs + logical-token/page-table conversion and applies causality upstream. + ``lengths`` and ``extra_lengths`` are contiguous CUDA INT32 [T]. Only + entries before each length are considered; -1 entries are masked holes. + Omit both extra tensors when no extra pool was planned. Input widths + and live lengths need no 128-entry padding: the kernel aligns source + execution spans internally and bounds all source reads. + + ``softmax_scale`` and ``output_scale`` are positive finite Python + scalars or scalar CUDA FP32 tensors. They are already combined: + softmax_scale = attention_scale * q_descale * kv_descale; + output_scale = output_multiplier * kv_descale. Both pools share that + KV descale. These multiply QK scores and the weighted stored V values, + respectively; sink logits are not multiplied by softmax_scale. The + output multiplier applies before fused RoPE/quantization. ``out_scale`` + instead holds the generated per-block quantization scales. + Combine dynamic scales in upstream work before each launch/replay and + keep their tensor addresses stable for CUDA Graphs. + + Caller-owned inputs and plan workspace must outlive graph replay, + including tensor scales and sinks, and must not alias writable output + or workspace storage. Q and native KV are used directly; flattening + padded pages only creates a tensor view. Required attention reductions + remain included. Caller-owned LSE must be contiguous FP32 and 16-byte + aligned, including with validate=False. With output quantization enabled, ``token_positions`` is contiguous INT32 with the query's batch/token prefix; ``cos_sin_cache`` is aligned @@ -2716,9 +2585,11 @@ def run( Positions must be in [0,P); validate=False leaves this to the caller. Return (out, out_scale), or (out, out_scale, lse) with return_lse=True. - Out is the [T,G,C] view (strides (C,T*C,1)) of a 16-byte aligned + Out is the [T,G,C] view (usual strides (C,T*C,1)) of a 16-byte aligned contiguous group-major E4M3[G,T,C] tensor, where T is total query - tokens and C=(H/G)*512, in natural head/channel order. Pass + tokens and C=(H/G)*512, in natural head/channel order. Equivalent + singleton-axis strides are accepted. A live prefix may retain a + larger [G,T_capacity,C] allocation; only its live T rows are written. Pass ``torch.empty((G,T,C), dtype=torch.float8_e4m3fn).transpose(0,1)`` as ``out`` to preallocate. For K=output_quant_block_size and N=C/K, FP32 scales have logical shape [T,G,N]; UE8M0 scales have @@ -2734,8 +2605,6 @@ def run( state = self._state if state is None: raise RuntimeError("plan() must be called before run()") - if not isinstance(metadata, SparseMLAPreparedMetadata): - raise TypeError("metadata must be SparseMLAPreparedMetadata") _validate_query( query, packed_query=state["packed"], @@ -2763,19 +2632,15 @@ def run( primary = self._pool(kv_cache, "kv_cache") if extra_kv_cache is None: - if ( - state["kc"] - or metadata.extra_indices is not None - or metadata.extra_lengths is not None - ): + if state["kc"] or extra_indices is not None or extra_lengths is not None: raise ValueError("extra pool and indices are required by the plan") extra = state["dummy_cache"] else: extra = self._pool(extra_kv_cache, "extra_kv_cache") si, sl = self._source_metadata( - metadata.indices, - metadata.lengths, + indices, + lengths, state["ks"], primary.shape[0], "primary", @@ -2783,8 +2648,8 @@ def run( validate=validate, ) ci, cl = self._source_metadata( - metadata.extra_indices, - metadata.extra_lengths, + extra_indices, + extra_lengths, state["kc"], extra.shape[0], "extra", @@ -2804,27 +2669,13 @@ def run( raise ValueError("sinks must be contiguous CUDA FP32[H]") if validate and torch.isnan(sinks).any().item(): raise ValueError("sinks must not contain NaN") - if extra_kv_scale is None: - extra_kv_scale = kv_scale - shared_scale = extra_kv_scale is kv_scale or ( - not isinstance(kv_scale, torch.Tensor) - and not isinstance(extra_kv_scale, torch.Tensor) - and kv_scale == extra_kv_scale - ) - independent = state["kc"] > 0 and not shared_scale scale_tensors = [ self._scalar(v, n, validate) for v, n in ( (softmax_scale, "softmax_scale"), - (q_scale, "q_scale"), - (kv_scale, "kv_scale"), - (extra_kv_scale, "extra_kv_scale"), (output_scale, "output_scale"), ) ] - if validate and state["dtype"] == torch.bfloat16: - if any(t.item() != 1 for t in scale_tensors[1:4]): - raise ValueError("BF16 Q/KV descales must be one") quant = state["output_quant"] output_quant_args = None if quant is not None: @@ -2869,121 +2720,64 @@ def run( or lse.dtype != torch.float32 or lse.device != state["device"] or not lse.is_contiguous() + or lse.data_ptr() % 16 ): raise ValueError( - "lse must be contiguous FP32 with one value per query/head" + "lse must be contiguous, 16-byte-aligned FP32 with one value per query/head" ) - fused = ( - state["compiled_fused"] is not None - and not independent - and lse.data_ptr() % 16 == 0 - ) - fused_main = fused - direct = fused and state["direct_inputs"] - static_scales = direct and all( - not isinstance(v, torch.Tensor) - for v in (softmax_scale, q_scale, kv_scale, output_scale) + static_scales = all( + not isinstance(v, torch.Tensor) for v in (softmax_scale, output_scale) ) - bmm1_scale, bmm2_scale = 1.0, 1.0 - if static_scales: - - def f32(value): - return struct.unpack("f", struct.pack("f", float(value)))[0] - - bmm1_scale = f32(f32(f32(softmax_scale) * f32(q_scale)) * f32(kv_scale)) - bmm2_scale = f32(f32(output_scale) * f32(kv_scale)) + # Both paths consume the same effective scales; no Q/K/V scale + # composition or auxiliary preparation is performed here. + bmm1_scale = float(softmax_scale) if static_scales else 1.0 + bmm2_scale = float(output_scale) if static_scales else 1.0 sparse_inputs = ( - ( - si, - ci, - sl, - cl, - scale_tensors[0], - scale_tensors[1], - scale_tensors[2], - scale_tensors[4], - sinks, - ) - if direct - else None + si, + ci, + sl, + cl, + scale_tensors[0], + scale_tensors[1], + sinks, ) if rows: - buffers = self._bind_metadata( - metadata, - rows=rows, - passes=2 if independent else 1, - direct=direct, - primary_lengths=sl, - ) - query_view = query.view(rows, 1, state["heads"], 512) - # Final FP8 stores use output_quant_args. Keep the core O argument - # in its existing BF16 partial-state layout for every schedule. - for slot in range(2 if independent else 1): - _launch_mla_decode( - _MLARuntime( - query_view, - primary, - ( - out.view(rows, 1, state["heads"], 512) - if fused_main and quant is None - else buffers["partial"][slot, :rows] - ), - primary.shape[0], - bmm1_scale, - bmm2_scale, - extra_cache=extra, - scale_params=buffers["scales"][slot], - sparse_inputs=sparse_inputs, - output_quant_args=output_quant_args if fused_main else None, - ), - block_tables=buffers["routes"][slot, :rows], - seq_lens=buffers["lengths"][slot, :rows], - qo_indptr=None, - packed_query=False, - kv_lora_rank=512, - split_kv=state["splits"], - workspace=_MLAWorkspaceViews( - buffers["core"] if buffers["core"].numel() else None, - ( - lse.view(rows, 1, state["heads"]) - if fused_main - else buffers["lse"][slot, :rows] - ), - ), - compiled=( - state["compiled_static"] - if static_scales - else state["compiled_fused"] - if fused - else state["compiled"] - ), - ) - if not fused: - other_slot = 1 if independent else 0 - state["finish"][int(independent)]( - buffers["partial"][0, :rows].view(rows, state["heads"], 512), - buffers["partial"][other_slot, :rows].view( - rows, state["heads"], 512 - ), - buffers["lse"][0, :rows].view(rows, state["heads"]), - buffers["lse"][other_slot, :rows].view(rows, state["heads"]), - buffers["counts"][0, :rows], - buffers["counts"][1 if independent else 0, :rows], - sinks, - ( - out.view(rows, state["heads"], 512) - if quant is None - else buffers["partial"][0, :rows].view( - rows, state["heads"], 512 - ) - ), - lse.view(rows, state["heads"]), - **( - {"output_quant_args": output_quant_args} - if quant is not None - else {} - ), + buffers = state["buffers"] + # Quantized stores use output_quant_args, including group/bucket + # strides. The core O view supplies compact launch geometry only. + core_out = ( + out.view(rows, 1, state["heads"], 512) + if quant is None + else output_quant_args[0][: rows * state["heads"] * 512].view( + rows, 1, state["heads"], 512 ) + ) + _launch_mla_decode( + _MLARuntime( + query.view(rows, 1, state["heads"], 512), + primary, + core_out, + primary.shape[0], + bmm1_scale, + bmm2_scale, + extra_cache=extra, + sparse_inputs=sparse_inputs, + output_quant_args=output_quant_args, + ), + block_tables=si, + seq_lens=sl, + qo_indptr=None, + packed_query=False, + kv_lora_rank=512, + split_kv=state["splits"], + workspace=_MLAWorkspaceViews( + buffers["core"] if buffers["core"].numel() else None, + lse.view(rows, 1, state["heads"]), + ), + compiled=state["compiled_static"] + if static_scales + else state["compiled"], + ) if quant is not None: return (out, out_scale, lse) if state["return_lse"] else (out, out_scale) return (out, lse) if state["return_lse"] else out @@ -3004,16 +2798,16 @@ def get_prims_ts_sparse_mla_decode_workspace_size(*plan_args, **plan_kwargs): def batch_sparse_mla_decode_with_paged_kv_cache( query, kv_cache, - metadata: SparseMLAPreparedMetadata, + indices, + lengths, extra_kv_cache=None, *, + extra_indices=None, + extra_lengths=None, qo_indptr=None, max_seq_len_q=None, kv_layout="NHD", softmax_scale=512**-0.5, - q_scale=1.0, - kv_scale=1.0, - extra_kv_scale=None, output_scale=1.0, sinks=None, out=None, @@ -3028,18 +2822,18 @@ def batch_sparse_mla_decode_with_paged_kv_cache( cos_sin_cache=None, out_scale=None, ): - """Eager plan-and-run helper using caller-prepared metadata. + """Eager plan-and-run helper using separate source indices and lengths. - Use a planned wrapper for CUDA Graph replay. Preparation remains external. - Output quantization arguments and return layouts follow + Use a planned wrapper for CUDA Graph replay. Index conversion and scale + composition remain external. softmax_scale includes attention/Q/K descales; + output_scale includes the V descale and optional output multiplier. + Source contracts, output quantization arguments and return layouts follow :class:`BatchSparseMLADecodePagedTSWrapper`. """ if torch.cuda.is_current_stream_capturing(): raise RuntimeError( "plan BatchSparseMLADecodePagedTSWrapper before CUDA Graph capture" ) - if not isinstance(metadata, SparseMLAPreparedMetadata): - raise TypeError("metadata must be SparseMLAPreparedMetadata") packed = query.ndim == 3 if packed: if qo_indptr is None: @@ -3055,10 +2849,8 @@ def batch_sparse_mla_decode_with_paged_kv_cache( query.device, batch, query.shape[-2], - max_topk=metadata.indices.shape[-1], - max_extra_topk=0 - if metadata.extra_indices is None - else metadata.extra_indices.shape[-1], + max_topk=indices.shape[-1], + max_extra_topk=0 if extra_indices is None else extra_indices.shape[-1], max_seq_len_q=max_seq_len_q, packed_query=packed, q_data_type=query.dtype, @@ -3073,13 +2865,13 @@ def batch_sparse_mla_decode_with_paged_kv_cache( return wrapper.run( query, kv_cache, - metadata, + indices, + lengths, extra_kv_cache, + extra_indices=extra_indices, + extra_lengths=extra_lengths, qo_indptr=qo_indptr, softmax_scale=softmax_scale, - q_scale=q_scale, - kv_scale=kv_scale, - extra_kv_scale=extra_kv_scale, output_scale=output_scale, sinks=sinks, out=out, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/q_token_kv_block_sparse_metadata.py b/tensorrt_llm/_torch/attention/backends/prims_ts/q_token_kv_block_sparse_metadata.py index 0c5cf40dbd4e..dff16046cf7c 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/q_token_kv_block_sparse_metadata.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/q_token_kv_block_sparse_metadata.py @@ -25,11 +25,12 @@ Patterns can be shared across KV heads or supplied independently for each head. Construction uses one CUDA C++ CTA per route. Q1 maps selected and tail blocks -directly. Q2--Q8 radix-sort at most ``G * (topk + 1)`` selected/tail IDs, -segmented-OR equal-key memberships, and emit only unique pages; work and -temporary storage are independent of the model context length. On SM90 and -newer, the metadata grid releases the prepared attention grid through -programmatic dependent launch (PDL). +directly. Q2--Q8 union at most ``G * (topk + 1)`` selected/tail IDs using a +shared byte/bit map when its model-context footprint fits without adding CTA +waves, or a bounded radix-sort fallback. No global-cache-sized scratch is +required. On SM90 and newer, the metadata grid releases the attention prologue +at entry through programmatic dependent launch (PDL); the consumer wait still +covers metadata completion. """ from __future__ import annotations @@ -189,6 +190,7 @@ class _PrimsTSQTokenKvBlockSparseWorkspaceLayout: attention_workspace_byte_offset: int attention_scratch_bytes: int uses_split_kv: bool + query_major_memberships: bool max_seq_len: int total_bytes: int @@ -246,6 +248,7 @@ class _PrimsTSQTokenKvBlockSparseMetadataPlan: storage_page_size: int max_seq_len_kv: int release_attention_pdl: bool + locator_head_multiplier: int = 1 def run( self, @@ -271,6 +274,7 @@ def run( self.sparse_block_size, self.max_seq_len_kv, self.release_attention_pdl, + self.locator_head_multiplier, ) else: self.metadata_run( @@ -286,6 +290,7 @@ def run( self.sparse_block_size, self.max_seq_len_kv, self.release_attention_pdl, + self.locator_head_multiplier, ) @@ -403,7 +408,7 @@ class _QTokenKvBlockSparsePlanConfig: """Capacity and dtype contract published by the public wrapper.""" batch_size: int - seq_len_q: int + group_size: int num_qo_heads: int num_kv_heads: int head_dim: int @@ -468,7 +473,6 @@ def __init__(self) -> None: def plan( self, batch_size: int, - seq_len_q: int, num_qo_heads: int, num_kv_heads: int, head_dim: int, @@ -479,6 +483,7 @@ def plan( *, device: torch.device | str | int, workspace_buffer: torch.Tensor, + group_size: int = 1, use_packed_q: bool = False, split_kv: bool = True, share_pattern_across_kv_heads: bool = True, @@ -491,8 +496,9 @@ def plan( ``batch_size`` is the number of query routes: ``len(qo_indptr)-1`` for packed Q, or ``B * Nq`` for fixed ``[B, Nq, G, Hq, D]`` Q. The planned - ``seq_len_q`` is the maximum packed route length or fixed group size - ``G``. ``kv_block_size`` is the indexer's sparse K/V atom and currently + ``group_size`` is the maximum packed route length or fixed group size + ``G``, defaulting to one (no grouping), not the request sequence length. + ``kv_block_size`` is the indexer's sparse K/V atom and currently supports 4/8/16/32/64/128 tokens; ``page_size`` is the physical cache page. Blocks may cross cache pages. The loader uses fragments of size gcd(kv_block_size, page_size), resolving every fragment independently. @@ -516,7 +522,7 @@ def plan( ) batch_size = _validate_positive_int(batch_size, "batch_size") - seq_len_q = _validate_positive_int(seq_len_q, "seq_len_q") + group_size = _validate_positive_int(group_size, "group_size") block_topk = _validate_positive_int(block_topk, "block_topk") max_seq_len_kv = _validate_positive_int(max_seq_len_kv, "max_seq_len_kv") page_size = _validate_positive_int(page_size, "page_size") @@ -524,7 +530,7 @@ def plan( _validate_head_geometry(num_qo_heads, num_kv_heads) kv_block_size = _validate_sparse_block_size(kv_block_size) _validate_storage_page_size(page_size) - _validate_shape_parameters(batch_size * seq_len_q, block_topk, seq_len_q) + _validate_shape_parameters(batch_size * group_size, block_topk, group_size) if mask_type != "causal": raise ValueError( "QToken-KvBlock-Sparse-Attention requires mask_type='causal'" @@ -551,11 +557,11 @@ def plan( max_num_storage_pages = (max_seq_len_kv + page_size - 1) // page_size layout = _get_prims_ts_q_token_kv_block_sparse_workspace_layout( - batch_size * seq_len_q, + batch_size * group_size, block_topk, max_num_storage_pages, page_size, - seq_len_q, + group_size, max_seq_len_kv=max_seq_len_kv, num_qo_heads=num_qo_heads, num_kv_heads=num_kv_heads, @@ -581,7 +587,7 @@ def plan( ) candidate = _QTokenKvBlockSparsePlanConfig( batch_size=batch_size, - seq_len_q=seq_len_q, + group_size=group_size, num_qo_heads=num_qo_heads, num_kv_heads=num_kv_heads, head_dim=head_dim, @@ -676,7 +682,7 @@ def run( qo_indptr : torch.Tensor, optional CUDA Int32 packed-route offsets. Required for packed Q and omitted for fixed Q; each route must stay within one request and contain - at most the planned ``seq_len_q`` tokens. + at most the planned ``group_size`` tokens. sm_scale : float, optional Softmax scale, defaulting to ``head_dim**-0.5``. Include Q/K dequantization scales here when using FP8 inputs. @@ -707,7 +713,7 @@ def run( ) if qo_indptr.numel() != config.batch_size + 1: raise ValueError("qo_indptr length must match the planned route count") - if q.shape[0] <= 0 or q.shape[0] > config.batch_size * config.seq_len_q: + if q.shape[0] <= 0 or q.shape[0] > config.batch_size * config.group_size: raise ValueError("packed q exceeds the planned query-token capacity") else: if qo_indptr is not None: @@ -715,7 +721,7 @@ def run( if ( q.ndim != 5 or q.shape[0] * q.shape[1] != config.batch_size - or q.shape[2] != config.seq_len_q + or q.shape[2] != config.group_size ): raise ValueError( "fixed q must be [B,Nq,G,Hq,D] with B*Nq and G matching plan" @@ -725,7 +731,7 @@ def run( num_query_tokens = ( int(q.shape[0]) if config.use_packed_q - else (config.batch_size * config.seq_len_q) + else (config.batch_size * config.group_size) ) expected_ids_shape = ( (num_query_tokens, config.block_topk) @@ -818,9 +824,7 @@ def run( max_seq_len_kv=config.max_seq_len_kv, bmm1_scale=sm_scale, qo_indptr=qo_indptr, - max_seq_len_q=( - config.seq_len_q if config.use_packed_q else None - ), + group_size=config.group_size, sparse_block_size=config.kv_block_size, split_kv=config.split_kv, share_pattern_across_kv_heads=config.share_pattern_across_kv_heads, @@ -974,7 +978,10 @@ def _get_prims_ts_q_token_kv_block_sparse_workspace_layout( ) from .decode import ( - _resolve_decode_workspace_layout, + _dtype_key, + _make_decode_workspace_layout, + _resolve_cuda_device, + _resolve_decode_launch_spec, _validate_prims_ts_q_token_kv_block_sparse_group_capacity, ) @@ -1001,6 +1008,50 @@ def _get_prims_ts_q_token_kv_block_sparse_workspace_layout( (max_seq_len_kv + sparse_block_size - 1) // sparse_block_size, ) page_capacity = block_capacity * (sparse_block_size // fragment_size) + # Padding is an allocation property, not additional selected KV work. + max_seq_len = ( + block_topk * sparse_block_size + (sparse_block_size - 1) + if group_size == 1 + else page_capacity * fragment_size + ) + max_seq_len = min(max_seq_len, max_seq_len_kv) + if kv_dtype is None: + kv_dtype = q_dtype + if out_dtype is None: + out_dtype = q_dtype + _, device_index = _resolve_cuda_device(device) + spec = _resolve_decode_launch_spec( + device_index, + groups, + num_qo_heads, + num_kv_heads, + head_dim, + fragment_size, + max_seq_len, + group_size, + _dtype_key(q_dtype), + _dtype_key(kv_dtype), + _dtype_key(kv_dtype), + _dtype_key(out_dtype), + "HND", + "causal", + use_packed_q, + -1, + storage_page_size, + True, + True, + split_kv, + share_pattern_across_kv_heads, + ) + # Eligible FP8 G2..G8/page-4 profiles consume one bit-plane per query. + # Other profiles and the standalone metadata API retain page-byte masks. + query_major_memberships = ( + sparse_block_size == 4 + and spec.config.use_fp8_qkv + and spec.config.supports_query_major_memberships + ) + if query_major_memberships: + page_capacity = (page_capacity + 31) // 32 * 32 membership_words = ( 0 if group_size == 1 @@ -1024,38 +1075,11 @@ def _get_prims_ts_q_token_kv_block_sparse_workspace_layout( ) seq_lens_numel = metadata_rows seq_lens_bytes = seq_lens_numel * 4 - max_seq_len = ( - block_topk * sparse_block_size + (sparse_block_size - 1) - if group_size == 1 - else page_capacity * fragment_size - ) - max_seq_len = min(max_seq_len, max_seq_len_kv) - if kv_dtype is None: - kv_dtype = q_dtype - if out_dtype is None: - out_dtype = q_dtype - attention_layout = _resolve_decode_workspace_layout( - groups, - num_qo_heads, - num_kv_heads, - head_dim, - fragment_size, - max_seq_len, - group_size, - q_dtype, - kv_dtype, - kv_dtype, + attention_layout = _make_decode_workspace_layout( + spec.scratch_shapes, out_dtype, - "HND", - "causal", - use_packed_q, - -1, - storage_page_size, - device, - use_q_token_kv_block_sparse_route=True, - use_pdl=True, - split_kv=split_kv, - share_pattern_across_kv_heads=share_pattern_across_kv_heads, + use_separate_reduction_kernel=spec.config.use_separate_reduction_kernel, + use_split_kv=spec.config.use_split_kv, ) attention_scratch_bytes = attention_layout.total_bytes attention_workspace_byte_offset = _align_up_q_token_kv_block_sparse_workspace( @@ -1072,6 +1096,7 @@ def _get_prims_ts_q_token_kv_block_sparse_workspace_layout( attention_workspace_byte_offset=attention_workspace_byte_offset, attention_scratch_bytes=attention_scratch_bytes, uses_split_kv=attention_layout.uses_split_kv, + query_major_memberships=query_major_memberships, max_seq_len=max_seq_len, total_bytes=attention_workspace_byte_offset + attention_scratch_bytes, ) @@ -1087,7 +1112,7 @@ def get_q_token_kv_block_sparse_workspace_size( max_seq_len_kv: int, o_data_type: Optional[torch.dtype] = None, qo_indptr: Optional[torch.Tensor] = None, - seq_len_q: Optional[int] = None, + group_size: int = 1, kv_block_size: int = 4, split_kv: bool = True, share_pattern_across_kv_heads: bool = True, @@ -1096,7 +1121,7 @@ def get_q_token_kv_block_sparse_workspace_size( Packed queries use ``[num_query_tokens, num_qo_heads, head_dim]`` and Int32 ``qo_indptr`` partitions their rows into request-safe routes of - length at most ``seq_len_q``. Fixed queries use + length at most ``group_size``. Fixed queries use ``[batch, num_query_groups, group_size, num_qo_heads, head_dim]`` and do not require query-offset metadata. The first two fixed axes are flattened into the attention route axis without copying tensor storage. @@ -1135,9 +1160,9 @@ def get_q_token_kv_block_sparse_workspace_size( dtype. qo_indptr : torch.Tensor, optional Contiguous CUDA Int32 cumulative offsets for packed query routes. - seq_len_q : int, optional - Maximum packed-route length and QToken-KvBlock-Sparse-Attention group size. It is required when - ``qo_indptr`` is supplied. + group_size : int, optional + Maximum query tokens per group, default one (no grouping). For fixed + Q, this must match its ``G`` dimension; packed final groups may be shorter. kv_block_size : int Logical sparse-block size in tokens: 4, 8, 16, 32, 64, or 128. split_kv : bool @@ -1165,7 +1190,7 @@ def get_q_token_kv_block_sparse_workspace_size( max_seq_len_kv=max_seq_len_kv, out_dtype=o_data_type, qo_indptr=qo_indptr, - max_seq_len_q=seq_len_q, + group_size=group_size, sparse_block_size=kv_block_size, split_kv=split_kv, share_pattern_across_kv_heads=share_pattern_across_kv_heads, @@ -1181,7 +1206,7 @@ def _get_prims_ts_q_token_kv_block_sparse_workspace_layout_from_tensors( max_seq_len_kv: int, out_dtype: Optional[torch.dtype], qo_indptr: Optional[torch.Tensor] = None, - max_seq_len_q: Optional[int] = None, + group_size: int = 1, sparse_block_size: int = 4, split_kv: bool = True, share_pattern_across_kv_heads: bool = True, @@ -1196,25 +1221,20 @@ def _get_prims_ts_q_token_kv_block_sparse_workspace_layout_from_tensors( if not isinstance(block_table, torch.Tensor): raise TypeError("block_table must be a torch.Tensor") use_packed_q = qo_indptr is not None - if not use_packed_q and max_seq_len_q is not None: - raise ValueError( - "max_seq_len_q is only valid with packed QToken-KvBlock-Sparse-Attention qo_indptr" - ) if use_packed_q: if query.ndim != 3: raise ValueError( "packed QToken-KvBlock-Sparse-Attention query must have shape [total_q,Hq,D]" ) - if max_seq_len_q is None: - raise ValueError( - "max_seq_len_q is required with QToken-KvBlock-Sparse-Attention qo_indptr" - ) _validate_q_token_kv_block_sparse_qo_indptr_layout_tensor(qo_indptr) - group_size = int(max_seq_len_q) groups = int(qo_indptr.numel()) - 1 num_query_tokens, num_qo_heads, head_dim = query.shape elif query.ndim == 5: - batch_size, groups_per_request, group_size, num_qo_heads, head_dim = query.shape + batch_size, groups_per_request, query_group_size, num_qo_heads, head_dim = ( + query.shape + ) + if group_size != query_group_size: + raise ValueError("fixed group_size must match query.shape[2]") if batch_size <= 0 or groups_per_request <= 0: raise ValueError( "fixed QToken-KvBlock-Sparse-Attention batch and query-group counts must be positive" @@ -1297,11 +1317,12 @@ def _build_prims_ts_q_token_kv_block_sparse_metadata( The membership output has one Int32 word per four page slots, or zero columns for Q1. Q1 uses one CUDA C++ CTA per route to map its selected blocks and causal - tail directly. Q2--Q8 use one CUDA C++ CTA per route to radix-sort the - bounded ``group_size * (block_topk + 1)`` candidates, unique them while - OR-reducing membership bits, and map the resulting logical pages through - the dense block table. The terminal metadata grid releases a following - PDL-capable QToken-KvBlock-Sparse-Attention attention launch when requested by the combined API. + tail directly. Q2--Q8 union the bounded + ``group_size * (block_topk + 1)`` candidates with a shared map or radix + sort, OR membership bits, and map the sorted unique logical blocks through + the dense block table. A map is used only when it fits shared memory and + requires no extra CTA waves. The combined API releases attention at + metadata entry; the consumer wait still covers every output store. Packed ``qo_indptr`` must be an Int32 device copy of CPU-validated route offsets; this builder checks only its structural tensor contract and does not read route values back to the host. @@ -1385,6 +1406,7 @@ def _build_prims_ts_q_token_kv_block_sparse_metadata( sparse_block_size, max_seq_len_kv, release_attention_pdl, + 1, ) else: metadata_module.run_fixed( @@ -1400,6 +1422,7 @@ def _build_prims_ts_q_token_kv_block_sparse_metadata( sparse_block_size, max_seq_len_kv, release_attention_pdl, + 1, ) return outputs @@ -1429,7 +1452,7 @@ def _build_q_token_kv_block_sparse_metadata( physical_page * (storage_page_size/F) + subpage. No CSR conversion is needed. Q1 uses one CUDA C++ direct-mapping kernel. Q2--Q8 use one CUDA C++ - radix-sort and union kernel per route. Neither path requires + map-or-sort union kernel per route. Neither path requires caller-provided scratch. Advanced callers that capture this raw path must preallocate ``out``, warm the same metadata geometry once before capture, and retain every tensor at a stable address through replay. The three @@ -1515,8 +1538,18 @@ def _prepare_prims_ts_q_token_kv_block_sparse_metadata_plan( max_seq_len_kv: int, sparse_block_size: int = 4, qo_indptr: Optional[torch.Tensor] = None, + query_major_memberships: bool = False, + locator_head_multiplier: int = 1, ) -> _PrimsTSQTokenKvBlockSparseMetadataPlan: - """Freeze validated metadata tensors and launch constants.""" + """Freeze validated metadata tensors and launch constants. + + Query-major memberships use a private logical [row, KV tile, query slot] + layout: each uint32 masks 32 four-token union fragments for one query. + The rank-two tensor flattens eight slots per tile, with unused slots and + tail bits zero. This encoding must match the attention plan. Standalone + metadata retains per-fragment membership bytes; Q/K/V layouts, the union, + and page indices are unchanged. + """ sparse_block_size = _validate_sparse_block_size(sparse_block_size) rows = block_indices.shape[0] @@ -1532,9 +1565,16 @@ def _prepare_prims_ts_q_token_kv_block_sparse_metadata_plan( ) release_attention_pdl = device_support_pdl(block_indices.device) metadata_module = _get_prims_ts_q_token_kv_block_sparse_metadata_module() - metadata_run = ( - metadata_module.run_packed if use_packed_q else metadata_module.run_fixed - ) + if query_major_memberships: + metadata_run = ( + metadata_module.run_packed_query_major + if use_packed_q + else metadata_module.run_fixed_query_major + ) + else: + metadata_run = ( + metadata_module.run_packed if use_packed_q else metadata_module.run_fixed + ) return _PrimsTSQTokenKvBlockSparseMetadataPlan( q_token_kv_block_sparse_page_indices=q_token_kv_block_sparse_page_indices, q_token_kv_block_sparse_page_memberships=q_token_kv_block_sparse_page_memberships, @@ -1547,6 +1587,7 @@ def _prepare_prims_ts_q_token_kv_block_sparse_metadata_plan( storage_page_size=storage_page_size, max_seq_len_kv=max_seq_len_kv, release_attention_pdl=release_attention_pdl, + locator_head_multiplier=locator_head_multiplier, ) @@ -1564,7 +1605,7 @@ def _prepare_q_token_kv_block_sparse_attention( bmm1_scale: Optional[float] = None, bmm2_scale: float = 1.0, qo_indptr: Optional[torch.Tensor] = None, - max_seq_len_q: Optional[int] = None, + group_size: int = 1, sparse_block_size: int = 4, split_kv: bool = True, share_pattern_across_kv_heads: bool = True, @@ -1572,7 +1613,7 @@ def _prepare_q_token_kv_block_sparse_attention( """Prepare one graph-stable metadata-plus-attention QToken-KvBlock-Sparse-Attention launch. Packed Q/output use ``[num_query_tokens, Hq, D]`` and describe Q1--Q8 - routes with ``qo_indptr`` and ``max_seq_len_q``. A request's final route + routes with ``qo_indptr`` and ``group_size``. A request's final route may be shorter than the maximum group size. Fixed Q/output use ``[B, Nq, G, Hq, D]`` and omit ``qo_indptr``; ``B * Nq`` becomes the internal attention route count through a zero-copy view. @@ -1634,9 +1675,9 @@ def _prepare_q_token_kv_block_sparse_attention( qo_indptr : torch.Tensor, optional Contiguous CUDA Int32 cumulative route offsets selecting packed-query mode. - max_seq_len_q : int, optional - Maximum packed-route length and QToken-KvBlock-Sparse-Attention group size. It is required when - ``qo_indptr`` is supplied. + group_size : int, optional + Maximum query tokens per group, default one (no grouping). For fixed + Q, this must match its ``G`` dimension; packed final groups may be shorter. sparse_block_size : int Logical sparse-block size in tokens: 4, 8, 16, 32, 64, or 128. @@ -1666,11 +1707,6 @@ def _prepare_q_token_kv_block_sparse_attention( raise ValueError( "packed QToken-KvBlock-Sparse-Attention query must have shape [total_q,Hq,D]" ) - if max_seq_len_q is None: - raise ValueError( - "max_seq_len_q is required with QToken-KvBlock-Sparse-Attention qo_indptr" - ) - group_size = int(max_seq_len_q) num_query_tokens = int(query.shape[0]) _validate_q_token_kv_block_sparse_qo_indptr_tensor( qo_indptr, @@ -1687,7 +1723,8 @@ def _prepare_q_token_kv_block_sparse_attention( raise ValueError( "fixed QToken-KvBlock-Sparse-Attention query and output must be contiguous" ) - group_size = int(query.shape[2]) + if group_size != query.shape[2]: + raise ValueError("fixed group_size must match query.shape[2]") num_query_tokens = int(query.shape[0] * query.shape[1]) * group_size attention_query = _flatten_fixed_q_token_kv_block_sparse_groups( query, group_size @@ -1724,26 +1761,12 @@ def _prepare_q_token_kv_block_sparse_attention( max_seq_len_kv=max_seq_len_kv, out_dtype=out.dtype, qo_indptr=qo_indptr, - max_seq_len_q=max_seq_len_q, + group_size=group_size, sparse_block_size=sparse_block_size, split_kv=split_kv, share_pattern_across_kv_heads=share_pattern_across_kv_heads, ) views = layout.bind(workspace_buffer) - metadata_plan = _prepare_prims_ts_q_token_kv_block_sparse_metadata_plan( - block_indices, - block_table, - token_to_request, - query_positions, - views.q_token_kv_block_sparse_page_indices, - views.q_token_kv_block_sparse_page_memberships, - views.seq_lens, - group_size=group_size, - storage_page_size=int(k_cache.shape[2]), - max_seq_len_kv=max_seq_len_kv, - sparse_block_size=sparse_block_size, - qo_indptr=qo_indptr, - ) from .decode import _prepare_prims_ts_batch_decode_plan, _validate_scale scale_qk = _validate_scale( @@ -1751,6 +1774,8 @@ def _prepare_q_token_kv_block_sparse_attention( "bmm1_scale", ) scale_v = _validate_scale(bmm2_scale, "bmm2_scale") + # Generic FMHA sees one query group as a route, so its route-local + # sequence-length parameters receive G, not the full request length. attention_plan, _ = _prepare_prims_ts_batch_decode_plan( attention_query, paged_kv_cache, @@ -1768,6 +1793,7 @@ def _prepare_q_token_kv_block_sparse_attention( kv_layout="HND", page_size=math.gcd(sparse_block_size, int(k_cache.shape[2])), direct_q1_sparse_block_size=sparse_block_size, + query_major_memberships=layout.query_major_memberships, use_q_token_kv_block_sparse_route=True, use_pdl=True, split_kv=split_kv, @@ -1781,6 +1807,25 @@ def _prepare_q_token_kv_block_sparse_attention( else () ), direct_q1_max_model_len=max_seq_len_kv if group_size == 1 else None, + allow_head_aware_locators=group_size > 1, + ) + # Preparation only binds shapes and storage; attention does not read these + # outputs until metadata runs. Use its resolved format for the producer. + metadata_plan = _prepare_prims_ts_q_token_kv_block_sparse_metadata_plan( + block_indices, + block_table, + token_to_request, + query_positions, + views.q_token_kv_block_sparse_page_indices, + views.q_token_kv_block_sparse_page_memberships, + views.seq_lens, + group_size=group_size, + storage_page_size=int(k_cache.shape[2]), + max_seq_len_kv=max_seq_len_kv, + sparse_block_size=sparse_block_size, + qo_indptr=qo_indptr, + query_major_memberships=layout.query_major_memberships, + locator_head_multiplier=attention_plan._flat_locator_heads, ) return _PrimsTSQTokenKvBlockSparsePlan( _metadata_plan=metadata_plan, @@ -1808,7 +1853,7 @@ def q_token_kv_block_sparse_attention_with_paged_kv_cache( workspace_buffer: torch.Tensor, *, max_seq_len_kv: int, - seq_len_q: Optional[int] = None, + group_size: int = 1, kv_block_size: int = 4, mask_type: Literal["causal"] = "causal", sm_scale: Optional[float] = None, @@ -1827,9 +1872,10 @@ def q_token_kv_block_sparse_attention_with_paged_kv_cache( wrapper, call :meth:`QTokenKvBlockSparsePagedTSWrapper.plan` outside capture, run once eagerly, and capture only :meth:`~QTokenKvBlockSparsePagedTSWrapper.run`. - Packed Q is ``[total_q, Hq, D]`` with ``qo_indptr``; ``seq_len_q`` + Packed Q is ``[total_q, Hq, D]`` with ``qo_indptr``; ``group_size`` is its maximum route length. Fixed Q is ``[B, Nq, G, Hq, D]`` without - ``qo_indptr``, and ``G`` determines ``seq_len_q``. The indexer + ``qo_indptr``, and ``group_size`` must match ``G``. The default is one + (no grouping) for both layouts. The indexer output is named ``indexer_block_ids`` to distinguish it from BSR ``block_indices``. ``kv_block_size`` is the semantic indexer atom (4/8/16/32/64/128), while the physical cache extent is inferred from @@ -1862,8 +1908,9 @@ def q_token_kv_block_sparse_attention_with_paged_kv_cache( :func:`get_q_token_kv_block_sparse_workspace_size`. max_seq_len_kv : int Static maximum logical K/V length in tokens. - seq_len_q : int, optional - Maximum packed route length. Fixed Q derives it from ``G``. + group_size : int, optional + Maximum query tokens per group, default one (no grouping). Must match + ``G`` for fixed Q; packed final groups may contain fewer tokens. kv_block_size : int Semantic sparse K/V block size: 4/8/16/32/64/128, default four. mask_type : {"causal"} @@ -1918,24 +1965,19 @@ def q_token_kv_block_sparse_attention_with_paged_kv_cache( raise ValueError( "fixed q must be [B,Nq,G,Hq,D], or supply qo_indptr for packed q" ) - fixed_seq_len_q = int(q.shape[2]) - if seq_len_q is not None and seq_len_q != fixed_seq_len_q: - raise ValueError("fixed seq_len_q must match q.shape[2]") - seq_len_q = fixed_seq_len_q + if group_size != q.shape[2]: + raise ValueError("fixed group_size must match q.shape[2]") batch_size = int(q.shape[0] * q.shape[1]) use_packed_q = False else: if q.ndim != 3: raise ValueError("packed q must have shape [total_q,Hq,D]") - if seq_len_q is None: - raise ValueError("seq_len_q is required with packed qo_indptr") batch_size = int(qo_indptr.numel()) - 1 use_packed_q = True wrapper = QTokenKvBlockSparsePagedTSWrapper() wrapper.plan( batch_size, - seq_len_q, int(q.shape[-2]), int(k_cache.shape[1]), int(q.shape[-1]), @@ -1945,6 +1987,7 @@ def q_token_kv_block_sparse_attention_with_paged_kv_cache( max_seq_len_kv, device=q.device, workspace_buffer=workspace_buffer, + group_size=group_size, use_packed_q=use_packed_q, mask_type=mask_type, split_kv=split_kv, 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 fe4d8086e833..cf3fcda7f9e0 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 @@ -7,12 +7,10 @@ import pytest import torch -import triton from tensorrt_llm._torch.attention.backends.fmha import registry from tensorrt_llm._torch.attention.backends.fmha.prims_ts_sparse_mla import ( PrimsTSSparseMlaFmha, - _pack_sparse_mla_metadata, get_prims_ts_toolchain_unsupported_reason, ) from tensorrt_llm._utils import get_sm_version @@ -116,53 +114,6 @@ def _random_indices(rows: int, capacity: int, pool_rows: int, device: torch.devi return indices, lengths -def _prepare_metadata(wrapper, indices, lengths, extra_indices, extra_lengths, heads, bmm1, bmm2): - """Build SparseMLAPreparedMetadata exactly like the FMHA adapter does.""" - from tensorrt_llm._torch.attention.backends.prims_ts.mla_decode import SparseMLAPreparedMetadata - - rows, ks = indices.shape - kc = 0 if extra_indices is None else extra_indices.shape[1] - capacity = int(wrapper._state["capacity"]) - device = indices.device - - def empty(shape, dtype=torch.int32): - return torch.empty(shape, dtype=dtype, device=device) - - prepared = SparseMLAPreparedMetadata( - indices=indices.contiguous(), - lengths=lengths.contiguous(), - extra_indices=None if kc == 0 else extra_indices.contiguous(), - extra_lengths=None if kc == 0 else extra_lengths.contiguous(), - routes=empty((1, rows, capacity)), - execution_lengths=empty((1, rows)), - valid_counts=empty((1, rows)), - scale_params=empty((1, 2 + heads + rows), torch.float32), - ) - sinks = torch.full((heads,), -torch.inf, dtype=torch.float32, device=device) - _pack_sparse_mla_metadata[(rows,)]( - prepared.indices, - prepared.extra_indices if kc else prepared.indices, - prepared.lengths, - prepared.extra_lengths if kc else prepared.lengths, - prepared.routes, - prepared.execution_lengths, - prepared.valid_counts, - prepared.scale_params, - sinks, - bmm1, - bmm2, - ks, - kc, - rows, - heads, - capacity, - 256, - triton.next_power_of_2(heads), - num_warps=4, - ) - return prepared - - @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn], ids=["bf16", "fp8"]) @pytest.mark.parametrize( "rows,heads,groups,topk,extra_topk", @@ -210,19 +161,15 @@ def make_pool(n): softmax_scale = LATENT**-0.5 descale = 0.75 if dtype == torch.float8_e4m3fn else 1.0 - bmm1 = softmax_scale * descale * descale - bmm2 = descale - run_scales = dict( - softmax_scale=softmax_scale, q_scale=descale, kv_scale=descale, extra_kv_scale=descale - ) + # The kernel takes the combined scales (attention x Q descale x KV descale + # on the softmax, the KV descale on the output), as the adapter passes them. + run_scales = dict(softmax_scale=softmax_scale * descale * descale, output_scale=descale) + sources = dict(extra_indices=extra_indices, extra_lengths=extra_lengths) if extra_topk else {} plan_args = dict(max_topk=topk, max_extra_topk=extra_topk, max_seq_len_q=1, q_data_type=dtype) reference = BatchSparseMLADecodePagedTSWrapper() reference.plan(device, rows, heads, **plan_args) - metadata = _prepare_metadata( - reference, indices, lengths, extra_indices, extra_lengths, heads, bmm1, bmm2 - ) - out_bf16 = reference.run(query, pool, metadata, extra_pool, **run_scales) + out_bf16 = reference.run(query, pool, indices, lengths, extra_pool, **sources, **run_scales) heads_per_group = heads // groups ref_fp8, ref_scale = torch.ops.trtllm.fused_inv_rope_fp8_quant_vllm_port( out_bf16.view(rows, heads, LATENT), @@ -246,9 +193,6 @@ def make_pool(n): num_output_groups=groups, **plan_args, ) - metadata = _prepare_metadata( - fused, indices, lengths, extra_indices, extra_lengths, heads, bmm1, bmm2 - ) channels = heads_per_group * LATENT padded_rows = (rows + 3) // 4 * 4 # The exact buffers TRT-LLM's fused DSv4 epilogue allocates. @@ -259,8 +203,10 @@ def make_pool(n): fused_out, fused_scale = fused.run( query, pool, - metadata, + indices, + lengths, extra_pool, + **sources, out=out_buf.transpose(0, 1), out_scale=scale_buf.permute(2, 0, 1)[:rows], token_positions=positions.view(rows, 1), From f3960c117e523b330d711d058dca11379b19c334 Mon Sep 17 00:00:00 2001 From: yunruis <205571022+yunruis@users.noreply.github.com> Date: Tue, 20 Oct 2026 03:19:45 -0700 Subject: [PATCH 07/19] [None][chore] Dedup the FMHA library announcement with logger.info_once Review follow-up: the per-manager set of announced libraries duplicated the keyed once-logging the logger already provides. Co-Authored-By: Claude Fable 5.1 Signed-off-by: yunruis <205571022+yunruis@users.noreply.github.com> --- .../_torch/attention/backends/fmha/manager.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/manager.py b/tensorrt_llm/_torch/attention/backends/fmha/manager.py index 130a22d574c3..5f86c100f8c7 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/manager.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/manager.py @@ -312,21 +312,18 @@ def __init__(self, attn: TrtllmAttention) -> None: if fmha_cls.is_available(attn): self.fmha_libs.append(fmha_cls(attn)) self._layer_idx = getattr(attn, "layer_idx", None) - # Libraries already announced in the log for this layer; one line per - # library per layer is enough to tell which kernel serves a model. - self._announced: set[type[Fmha]] = set() def _announce(self, fmha: Fmha, cache_key: _FmhaCacheKey) -> None: - if type(fmha) in self._announced: - return - self._announced.add(type(fmha)) - logger.info( + # One line per library per layer is enough to tell which kernel + # serves a model; the logger dedups on the key. + logger.info_once( f"FMHA layer {self._layer_idx}: {type(fmha).__name__} serves " f"{cache_key.attention_mask_type.name} requests " f"(context_batch<={cache_key.context_batch_size}, " f"generation_batch<={cache_key.generation_batch_size}, " f"generation_seq_len_q<={cache_key.generation_seq_len_q}, " - f"output_dtype={cache_key.output_dtype})." + f"output_dtype={cache_key.output_dtype}).", + key=f"fmha_serves:{self._layer_idx}:{type(fmha).__name__}", ) def _make_cache_key( From 84bef3fd89baecb1c1e5bd85f0b298b5ce8afde7 Mon Sep 17 00:00:00 2001 From: yunruis <205571022+yunruis@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:32:49 -0700 Subject: [PATCH 08/19] [None][fix] Key the PrimTS pool-view cache by a weak manager reference The SWA/compress pool views were cached by id(manager). A freed cache manager can hand its id to a new one, in which case the stale view of released pool memory would be returned. Hold the owning manager weakly and drop the views whenever the owner changes, so a freed manager (or a new one at the same address) never resolves to an old view. Add a unit test with stand-in managers for the owner change and the freed owner. Co-Authored-By: Claude Fable 5.1 Signed-off-by: yunruis <205571022+yunruis@users.noreply.github.com> --- .../backends/fmha/prims_ts_sparse_mla.py | 13 ++++- .../test_prims_ts_sparse_mla_fmha.py | 48 +++++++++++++++++++ 2 files changed, 59 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py index 80f7d47e3565..1740bd84c717 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py @@ -44,6 +44,7 @@ from __future__ import annotations import math +import weakref from dataclasses import dataclass from importlib import import_module from importlib.metadata import PackageNotFoundError, version @@ -129,7 +130,11 @@ def __init__(self, attn: "TrtllmAttention") -> None: # FP8 dequantization scales are fixed per layer; the wrapper takes host # scalars, so read them once during eager warmup, never under capture. self._kv_dequant_scale: Optional[float] = None - # Pool-base-rooted views of the SWA and compress pools, per manager. + # Pool-base-rooted views of the SWA and compress pools, owned by the + # manager they were built from. The owner is held weakly: a freed + # manager (or a new one at the same address) never resolves to a + # stale view of released pool memory. + self._pool_views_owner: Optional[weakref.ReferenceType] = None self._pool_views: dict[tuple[int, int], torch.Tensor] = {} self._sinks_cache: Optional[tuple[int, torch.Tensor]] = None @@ -329,7 +334,11 @@ def _pool_view( DeepseekV4AttentionType, ) - key = (id(manager), int(attn_type.value), compress_ratio) + owner = self._pool_views_owner() if self._pool_views_owner is not None else None + if owner is not manager: + self._pool_views.clear() + self._pool_views_owner = weakref.ref(manager) + key = (int(attn_type.value), compress_ratio) view = self._pool_views.get(key) if view is not None: return view 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 cf3fcda7f9e0..0b9d26845a92 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 @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 """Tests for the PrimTS native sparse MLA FMHA library (DeepSeek-V4, SM100).""" +import gc import math from types import SimpleNamespace @@ -97,6 +98,53 @@ def test_softmax_scale_follows_qk_head_dim() -> None: assert PrimsTSSparseMlaFmha._softmax_scale(attn) == pytest.approx(LATENT**-0.5, rel=1e-12) +def test_pool_view_cache_follows_the_owning_manager() -> None: + """A new cache manager gets fresh pool views, even at the old manager's address. + + The views wrap the raw pool base pointer, so a view cached for a freed + manager would alias released memory. The cache is owned through a weak + reference and dropped whenever the owner changes. + """ + from tensorrt_llm._torch.attention.backends.sparse.deepseek_v4.params import ( + DeepseekV4AttentionType, + ) + + class _Owner: # weak-referenceable stand-ins for TrtllmAttention / the manager + pass + + def make_manager(rows: int) -> _Owner: + manager = _Owner() + manager.pool = torch.empty(rows, 1, LATENT, dtype=torch.bfloat16, device="cuda") + manager.swa_pool_ptr = manager.pool.data_ptr() + manager.compress_pool_ptrs = {} + manager.pp_layers = [0] + manager._compress_ratios = {0: 4} + manager.dtype = torch.bfloat16 + manager.get_buffers = lambda layer, attn_type: manager.pool + return manager + + attn = _Owner() + fmha = PrimsTSSparseMlaFmha(attn) + first = make_manager(8) + view = fmha._pool_view(first, DeepseekV4AttentionType.SWA, 1) + assert view.data_ptr() == first.swa_pool_ptr + assert tuple(view.shape) == (8, 1, LATENT) + assert fmha._pool_view(first, DeepseekV4AttentionType.SWA, 1) is view + + second = make_manager(16) + other = fmha._pool_view(second, DeepseekV4AttentionType.SWA, 1) + assert other is not view + assert other.data_ptr() == second.swa_pool_ptr + assert tuple(other.shape) == (16, 1, LATENT) + assert fmha._pool_views_owner() is second + + del second + gc.collect() + assert fmha._pool_views_owner() is None + third = make_manager(4) + assert tuple(fmha._pool_view(third, DeepseekV4AttentionType.SWA, 1).shape) == (4, 1, LATENT) + + def _rope_table(max_positions: int, device: torch.device) -> torch.Tensor: """Forward [cos(32), sin(32)] rows, the layout both kernels consume.""" inv_freq = 1.0 / ( From 7d2becf92742ff07a21b5c8ba78a32a42b7a0a54 Mon Sep 17 00:00:00 2001 From: yunruis <205571022+yunruis@users.noreply.github.com> Date: Tue, 20 Oct 2026 05:40:10 -0700 Subject: [PATCH 09/19] [None][fix] Pass output_quant_args positionally to the cached PrimTS MLA kernel The vendored PrimTS tree now persists compiled CuTe DSL kernels (flashinfer-ai/flashinfer#6529). A kernel reloaded from that cache is a positional TVM-FFI entry: it has no Python defaults and takes no keyword arguments. _launch_mla_decode omitted output_quant_args when it was None, which the in-process cute.compile object tolerated, so a cold process worked while every later process failed on each unquantized MLA decode (dense and sparse) with "TypeError: Expects 15 parameters when calling". Pass the argument positionally and unconditionally; the compile signature already contains it (None for unquantized output). Carried in the vendor compatibility patch until it lands upstream. Co-Authored-By: Claude Fable 5.1 Signed-off-by: yunruis <205571022+yunruis@users.noreply.github.com> --- .../vendor_patches/flashinfer-prims-ts.patch | 26 ++++++++++++++++--- 3rdparty/vendor_sources.lock.yaml | 4 +-- .../attention/backends/prims_ts/mla_decode.py | 11 ++++---- 3 files changed, 30 insertions(+), 11 deletions(-) diff --git a/3rdparty/vendor_patches/flashinfer-prims-ts.patch b/3rdparty/vendor_patches/flashinfer-prims-ts.patch index 027b9b65a22f..b911288a6e6f 100644 --- a/3rdparty/vendor_patches/flashinfer-prims-ts.patch +++ b/3rdparty/vendor_patches/flashinfer-prims-ts.patch @@ -242,7 +242,7 @@ index b22f0f413b7312bab9a5f8e296255fe9aeabe901..a4a87200ca5d3ace5cf899caa9672a1a from .helpers_common import ( _TASK_CACHE_TMEM_BASE_OFFSET, diff --git a/mla_decode.py b/mla_decode.py -index b62bd1081fbae0e18e0950008bae87c8986677b4..e64fc538c768f08cc4d51ce12e74e00ac59cd692 100644 +index b62bd1081fbae0e18e0950008bae87c8986677b4..0555e277288fdeef25e85de1de4415a578cfe743 100644 --- a/mla_decode.py +++ b/mla_decode.py @@ -1,3 +1,4 @@ @@ -263,7 +263,25 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..e64fc538c768f08cc4d51ce12e74e00a from .decode import ( _WorkspaceSection, -@@ -1637,7 +1634,7 @@ class BatchMLADecodePagedTSWrapper: +@@ -1459,11 +1456,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 +1635,7 @@ class BatchMLADecodePagedTSWrapper: split_kv=int(dict(policy)["split_kv"]), ) @@ -272,7 +290,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..e64fc538c768f08cc4d51ce12e74e00a def run( self, query: torch.Tensor, -@@ -1742,7 +1739,7 @@ class BatchMLADecodePagedTSWrapper: +@@ -1742,7 +1740,7 @@ class BatchMLADecodePagedTSWrapper: ) @@ -281,7 +299,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..e64fc538c768f08cc4d51ce12e74e00a def batch_mla_decode_with_paged_kv_cache( query: torch.Tensor, kv_cache: torch.Tensor, -@@ -2119,17 +2116,17 @@ def _resolve_sparse_mla_plan( +@@ -2119,17 +2117,17 @@ def _resolve_sparse_mla_plan( ): """Resolve sparse geometry and scratch sizing without allocation/compilation.""" import cutlass diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 3d4f00be8a0a..7281e21d0f6f 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -21,8 +21,8 @@ vendors: include: - '**/*.py' patch: 3rdparty/vendor_patches/flashinfer-prims-ts.patch - patch_digest: sha256:53c48256776f437b5910a58c031226c594c56902c936f21c3989d01c06fe5c48 - digest: sha256-tree-v1:1f1ab05273ece9646bfcc6b6baf27cc0dc042d3bf6e1d16f47ee45270baee78e + patch_digest: sha256:09d14793745622df73c38d33cc1101db17b7a08a800a9fa5c1be0d4c43af377c + digest: sha256-tree-v1:f2800344002d0f73d2e9dc26efcf1a100dbcf2bfe5185fa29db5f2391746515d openengine-proto: url: https://github.com/ai-dynamo/openengine.git tag: v0.1.0 diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py b/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py index e64fc538c768..0555e277288f 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py @@ -1456,11 +1456,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 From c59f3597548e07842aafa20ed08ae5a725d4a073 Mon Sep 17 00:00:00 2001 From: yunruis <205571022+yunruis@users.noreply.github.com> Date: Tue, 20 Oct 2026 22:55:34 -0700 Subject: [PATCH 10/19] [None][fix] Cache PrimTS sparse MLA plans on the attention metadata and list the test in L0 Plans were cached per layer, so every layer kept its own planned wrapper and kernel workspace for each CUDA-graph geometry. Keep them in TrtllmAttentionMetadata.fmha_plan_caches under PLAN_CACHE_KEY instead, like the block-sparse and FlashInfer wrappers: layers that share a metadata object run serially and reuse one plan per static profile (device, rows, extra capacity, heads, window, dtype, sinks, fusion groups), and each CUDA-graph batch size keeps its own metadata, so the captured plans stay stable. Add a unit test for the sharing. Add test_prims_ts_sparse_mla_fmha.py to the B200 L0 list so the default SM100 path is covered in CI (34 min cold / 3 min warm on B200). Co-Authored-By: Claude Fable 5.1 Signed-off-by: yunruis <205571022+yunruis@users.noreply.github.com> --- .../backends/fmha/prims_ts_sparse_mla.py | 28 ++++++++--- .../test_lists/test-db/l0_b200.yml | 1 + .../test_prims_ts_sparse_mla_fmha.py | 49 +++++++++++++++++++ 3 files changed, 70 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py index 1740bd84c717..7c78d19ea979 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py @@ -121,12 +121,14 @@ class PrimsTSSparseMlaFmha(Fmha): # The native kernel reads one KV head of 512 channels per selected token. MAX_HEADS = 128 + # Plans live in ``TrtllmAttentionMetadata.fmha_plan_caches`` under this key, + # like the block-sparse and FlashInfer wrappers: every layer that runs with + # one metadata object plans each static profile once and reuses its + # workspace, and layers sharing a metadata object execute serially. + PLAN_CACHE_KEY = "prims_ts_sparse_mla" def __init__(self, attn: "TrtllmAttention") -> None: super().__init__(attn) - # Plans are keyed by the execution profile that changes kernel - # selection or buffer geometry: query rows, dtype, sinks, fusion. - self._plans: dict[tuple, _SparsePlan] = {} # FP8 dequantization scales are fixed per layer; the wrapper takes host # scalars, so read them once during eager warmup, never under capture. self._kv_dequant_scale: Optional[float] = None @@ -418,6 +420,7 @@ def _kv_dequant(self) -> float: def _get_plan( self, *, + metadata: "TrtllmAttentionMetadata", device: torch.device, rows: int, extra: int, @@ -425,6 +428,12 @@ def _get_plan( has_sinks: bool, fused_groups: Optional[int], ) -> _SparsePlan: + """Return the metadata-cached plan for this static profile, planning it once. + + The key holds every value that changes kernel selection or buffer + geometry: device, query rows, extra capacity, heads, window, dtype, + sinks, fusion groups. One plan serves every layer with that profile. + """ from tensorrt_llm._torch.attention.backends.prims_ts.mla_decode import ( BatchSparseMLADecodePagedTSWrapper, ) @@ -432,8 +441,9 @@ def _get_plan( attn = self.attn heads = int(attn.num_heads) window = int(attn.sparse_attention_config.window_size) - key = (rows, extra, q_dtype, has_sinks, fused_groups) - plan = self._plans.get(key) + key = (device, rows, extra, heads, window, q_dtype, has_sinks, fused_groups) + cache = metadata.fmha_plan_caches.setdefault(self.PLAN_CACHE_KEY, {}) + plan = cache.get(key) if plan is not None: return plan if torch.cuda.is_current_stream_capturing(): @@ -474,7 +484,7 @@ def _get_plan( torch.empty((rows,), device=device, dtype=torch.int32) if extra else None ), ) - self._plans[key] = plan + cache[key] = plan return plan # ------------------------------------------------------------------ forward @@ -530,6 +540,7 @@ def forward( sinks = self._sinks(forward_args) plan = self._get_plan( + metadata=metadata, device=q.device, rows=rows, extra=extra, @@ -571,8 +582,9 @@ def forward( else: run_kwargs["out"] = output.view(rows, 1, heads, _LATENT_DIM) - # Validate the index contract once per plan while still eager; graph - # replay must keep validation off because it synchronizes. + # Validate the index contract once per plan (the first layer that uses + # it) while still eager; graph replay must keep validation off because + # it synchronizes. validate = not plan.validated and not torch.cuda.is_current_stream_capturing() plan.wrapper.run( query, diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 94cea2e5c098..e1ffa0a9a87b 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -121,6 +121,7 @@ l0_b200: - unittest/disaggregated/test_deepseek_v4_kv_transfer.py - unittest/disaggregated/test_fp4_mla_kv_transfer.py - unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py TIMEOUT (60) + - unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py TIMEOUT (90) - unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py TIMEOUT (60) - unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.py TIMEOUT (60) - unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py TIMEOUT (30) 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 0b9d26845a92..ac1ae79ca025 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 @@ -145,6 +145,55 @@ def make_manager(rows: int) -> _Owner: assert tuple(fmha._pool_view(third, DeepseekV4AttentionType.SWA, 1).shape) == (4, 1, LATENT) +def test_plans_are_shared_through_the_metadata_cache(monkeypatch) -> None: + """Layers with one static profile share one plan via ``metadata.fmha_plan_caches``. + + The plan owns the kernel workspace, so caching it per layer would keep one + workspace per layer per CUDA-graph geometry; the metadata cache (the + convention of the block-sparse and FlashInfer wrappers) keeps one per + profile, and layers sharing a metadata object run serially. + """ + import tensorrt_llm._torch.attention.backends.prims_ts.mla_decode as mla_decode + + planned = [] + + class _FakeWrapper: + def plan(self, device, rows, heads, **kwargs): + planned.append((device, rows, heads, kwargs)) + + monkeypatch.setattr(mla_decode, "BatchSparseMLADecodePagedTSWrapper", _FakeWrapper) + + class _Attn: # weak-referenceable stand-in for TrtllmAttention + num_heads = 16 + sparse_attention_config = SimpleNamespace(window_size=128) + + attn_a, attn_b = _Attn(), _Attn() + layer_a, layer_b = PrimsTSSparseMlaFmha(attn_a), PrimsTSSparseMlaFmha(attn_b) + metadata = SimpleNamespace(fmha_plan_caches={}) + profile = dict( + metadata=metadata, + device=torch.device("cuda"), + rows=4, + extra=256, + q_dtype=torch.bfloat16, + has_sinks=False, + fused_groups=None, + ) + plan_a = layer_a._get_plan(**profile) + plan_b = layer_b._get_plan(**profile) + assert plan_a is plan_b + assert len(planned) == 1 + assert planned[0][3]["max_topk"] == 128 and planned[0][3]["max_extra_topk"] == 256 + cache = metadata.fmha_plan_caches[PrimsTSSparseMlaFmha.PLAN_CACHE_KEY] + assert list(cache.values()) == [plan_a] + + other = layer_a._get_plan(**{**profile, "extra": 0}) + assert other is not plan_a + assert other.extra_lengths is None + assert len(planned) == 2 + assert len(cache) == 2 + + def _rope_table(max_positions: int, device: torch.device) -> torch.Tensor: """Forward [cos(32), sin(32)] rows, the layout both kernels consume.""" inv_freq = 1.0 / ( From f3272a496887a7aa0fc157d811e06781369e39c0 Mon Sep 17 00:00:00 2001 From: yunruis <205571022+yunruis@users.noreply.github.com> Date: Tue, 20 Oct 2026 23:48:35 -0700 Subject: [PATCH 11/19] [None][fix] Share one PrimTS sparse MLA workspace per metadata and derive live lengths per step Pre-PR review follow-ups for the PrimTS sparse MLA FMHA: - All plans of a metadata object now share one kernel workspace, grown to the largest profile, instead of each plan allocating its own: the eager metadata sees one plan per distinct generation row count, so it no longer accumulates one split-KV scratch per row count. - The live SWA and compressed lengths are derived once per step in the DeepSeek-V4 metadata (sparse_mla_swa_lens, sparse_mla_extra_lens) next to sparse_mla_topk_lens; the library slices them instead of launching 2-4 small kernels per layer, and plans no longer own length buffers. - The toolchain probe applies the same CUTLASS compiler-target gate as the dense PrimTS libraries; the unavailable reason is logged once per reason. - Tests: reloaded-kernel launch (the persistent compile cache path CI never reached), CUDA graph capture/replay with swapped indices and lengths, shared-workspace growth in the metadata plan cache, and a guard that the vendored tree does not import the installed flashinfer attention tree. - l0_b200.yml: keep the new test file out of the attention aggregate so it runs once; ATTENTION_DEVELOPER_GUIDE.md documents the default-on library and its opt-out. Comments keep the invariants and drop CI history. Co-Authored-By: Claude Fable 5.1 Signed-off-by: yunruis <205571022+yunruis@users.noreply.github.com> --- .../attention/ATTENTION_DEVELOPER_GUIDE.md | 13 +- .../backends/fmha/prims_ts_sparse_mla.py | 120 ++++++++------- .../attention/backends/fmha/registry.py | 5 +- .../backends/sparse/deepseek_v4/metadata.py | 31 +++- .../test_lists/test-db/l0_b200.yml | 2 +- .../test_prims_ts_sparse_mla_fmha.py | 145 ++++++++++++++++-- .../test_prims_ts_vendored_imports.py | 30 ++++ 7 files changed, 271 insertions(+), 75 deletions(-) create mode 100644 tests/unittest/_torch/attention/test_prims_ts_vendored_imports.py diff --git a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index f9b52aef8a9d..8cb5a8394c0c 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md @@ -479,7 +479,11 @@ the selected library. it may add host overhead; use `TLLM_FMHA_LIBS=+prims_ts` to add it to the defaults or `TLLM_FMHA_LIBS=fallback` to force the fallback path. Generic block-sparse PrimTS remains enabled by default because a dense fallback cannot -preserve its routing semantics. Delta entries update the +preserve its routing semantics. The PrimTS sparse MLA library +(`prims_ts_sparse_mla`) is also enabled by default: it serves only DeepSeek-V4 +generation on SM100 and sits ahead of the TRTLLM-Gen libraries, so DeepSeek-V4 +decode on B200/GB200 runs on it unless `TLLM_FMHA_LIBS=-prims_ts_sparse_mla` +removes it. Delta entries update the default membership and follow canonical registry order, while an exact list preserves the user-specified order. Each FMHA library exposes `is_available()` for module/static environment checks and `is_supported()` for per-forward @@ -552,6 +556,13 @@ The FMHA package is split by role: With an FP8 KV cache it passes `mla_bmm1_scale` / `mla_bmm2_scale` to the kernel as device tensors (element 0 of each), so the scales need no host read and stay valid under CUDA graph replay. +- `fmha/prims_ts_sparse_mla.py` runs DeepSeek-V4 generation on the vendored + PrimTS two-source sparse MLA kernel (SM100 only; SM103 stays on TRTLLM-Gen). + It consumes the dual-pool storage-row indices from `sparse_attn_predict` and + the per-step live lengths from the DeepSeek-V4 metadata, views the SWA and + compress pools from their base pointers, keeps its plans in + `fmha_plan_caches` with one shared kernel workspace per metadata object, and + writes the fused DSv4 epilogue (E4M3 output plus FP32 block scales) in place. - `fmha/prims_ts_block_sparse.py` adapts generic block-sparse requests to the vendored PrimTS contiguous and paged wrappers. Paged generation passes a live, zero-copy 2D K-page-table view with its TRT-LLM padded row stride; it diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py index 7c78d19ea979..192bf2d96eb5 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py @@ -79,6 +79,9 @@ _MIN_CUTLASS_DSL_VERSION = Version("4.7.0") +# Same compiler-target gate as fmha/prims_ts.py: the kernels share one CuTe DSL +# task-scheduling stack, so a toolchain the dense libraries refuse is refused here too. +_MIN_CUTLASS_COMPILER_VERSION = "13.3" _LATENT_DIM = 512 _ROPE_DIM = 64 _FUSED_OUTPUT_BLOCK = 128 @@ -94,6 +97,15 @@ def get_prims_ts_toolchain_unsupported_reason() -> Optional[str]: return ( f"nvidia-cutlass-dsl>={_MIN_CUTLASS_DSL_VERSION} is required, got {installed_version}." ) + try: + cutlass = import_module("cutlass") + compiler_version_supported = cutlass.target_version( + min_version=_MIN_CUTLASS_COMPILER_VERSION + ) + except Exception as error: # noqa: BLE001 - availability probes must fail closed + return f"could not query the active CUTLASS compiler version: {error}" + if not compiler_version_supported: + return f"the active CUTLASS compiler must target CUDA>={_MIN_CUTLASS_COMPILER_VERSION}." try: import_module("cutlass.experimental.task_scheduling") except ImportError: @@ -103,16 +115,9 @@ def get_prims_ts_toolchain_unsupported_reason() -> Optional[str]: @dataclass class _SparsePlan: - """One planned wrapper plus the live-length buffers it reads on every launch. - - The index lists come straight from ``sparse_attn_predict``'s metadata - buffer; only the per-row live lengths are derived here, into buffers whose - addresses stay stable for CUDA graph replay. - """ + """One planned wrapper for a static profile; every launch input comes from the metadata.""" wrapper: "BatchSparseMLADecodePagedTSWrapper" - lengths: torch.Tensor - extra_lengths: Optional[torch.Tensor] validated: bool = False @@ -123,9 +128,13 @@ class PrimsTSSparseMlaFmha(Fmha): MAX_HEADS = 128 # Plans live in ``TrtllmAttentionMetadata.fmha_plan_caches`` under this key, # like the block-sparse and FlashInfer wrappers: every layer that runs with - # one metadata object plans each static profile once and reuses its - # workspace, and layers sharing a metadata object execute serially. + # one metadata object plans each static profile once, and layers sharing a + # metadata object execute serially. All plans of a metadata object share + # one kernel workspace (under WORKSPACE_CACHE_KEY), grown to the largest + # profile, so the eager metadata does not accumulate one workspace per + # distinct generation row count. PLAN_CACHE_KEY = "prims_ts_sparse_mla" + WORKSPACE_CACHE_KEY = "prims_ts_sparse_mla.workspace" def __init__(self, attn: "TrtllmAttention") -> None: super().__init__(attn) @@ -150,10 +159,11 @@ def _is_available(cls, attn: "TrtllmAttention") -> bool: reason = cls._unavailable_reason(attn) if reason is not None: # A DeepSeek-V4 layer that does not get this library falls back to - # TRTLLM-Gen silently otherwise; say so once per layer at info. - logger.info( - f"PrimTS sparse MLA FMHA is unavailable for layer " - f"{getattr(attn, 'layer_idx', None)}: {reason}" + # TRTLLM-Gen silently otherwise; say so once per reason at info. + logger.info_once( + f"PrimTS sparse MLA FMHA is unavailable, DeepSeek-V4 generation stays on the " + f"TRTLLM-Gen fallback: {reason}", + key=f"prims_ts_sparse_mla_unavailable:{reason}", ) return False return True @@ -162,13 +172,9 @@ def _is_available(cls, attn: "TrtllmAttention") -> bool: def _unavailable_reason(cls, attn: "TrtllmAttention") -> Optional[str]: sm = get_sm_version() if sm == 103: - # The vendored kernels admit SM103 through their architecture guard - # but flashinfer has only signed off accuracy and performance on - # SM100 (prims_ts/kernels/mla_decode/README.md). TRT-LLM L0 #64329 - # hit an illegal memory access in the GB300 DeepSeek-V4-Pro perf - # sanity warmup (TP4/ADP, FP8 KV, MTP 3) that the same configuration - # at TP8 does not reproduce on B200, so B300/GB300 stay on the - # TRTLLM-Gen fallback until the kernels are qualified there. + # The kernels' architecture guard admits SM103, but they are only + # signed off on SM100 (prims_ts/kernels/mla_decode/README.md); keep + # B300/GB300 on the TRTLLM-Gen fallback until they are qualified. return "SM103 is not yet qualified for the PrimTS sparse MLA kernels." if sm != 100: return f"requires SM100, got SM{sm}." @@ -231,7 +237,9 @@ def _unsupported_reason( if metadata.beam_width != 1: return "beam search is not supported." if getattr(metadata, "token_positions_cuda", None) is None: - return "token positions are required for lengths and output RoPE." + return "token positions are required for the output RoPE." + if getattr(metadata, "sparse_mla_swa_lens", None) is None: + return "per-step live lengths are required (DeepSeek-V4 metadata)." if fwd.attention_mask_data is not None or fwd.relative_attention_bias is not None: return "custom masks and relative bias are not supported." if metadata.helix_position_offsets is not None: @@ -434,9 +442,7 @@ def _get_plan( geometry: device, query rows, extra capacity, heads, window, dtype, sinks, fusion groups. One plan serves every layer with that profile. """ - from tensorrt_llm._torch.attention.backends.prims_ts.mla_decode import ( - BatchSparseMLADecodePagedTSWrapper, - ) + from tensorrt_llm._torch.attention.backends.prims_ts import mla_decode attn = self.attn heads = int(attn.num_heads) @@ -451,11 +457,7 @@ def _get_plan( "PrimTS sparse MLA plans must be created during eager warmup, not under " f"CUDA graph capture (rows={rows}, extra={extra}, dtype={q_dtype})." ) - wrapper = BatchSparseMLADecodePagedTSWrapper() - wrapper.plan( - device, - rows, - heads, + plan_kwargs = dict( max_topk=window, max_extra_topk=extra, max_seq_len_q=1, @@ -467,23 +469,28 @@ def _get_plan( # dual-pool layout, but the ratio-4 indexer top-k may carry holes; # keep the generic mask path. assume_valid_prefix=False, - **( - dict( - output_quant_block_size=_FUSED_OUTPUT_BLOCK, - out_scale_format="fp32", - num_output_groups=fused_groups, - ) - if fused_groups - else {} - ), ) - plan = _SparsePlan( - wrapper=wrapper, - lengths=torch.empty((rows,), device=device, dtype=torch.int32), - extra_lengths=( - torch.empty((rows,), device=device, dtype=torch.int32) if extra else None - ), + if fused_groups: + plan_kwargs.update( + output_quant_block_size=_FUSED_OUTPUT_BLOCK, + out_scale_format="fp32", + num_output_groups=fused_groups, + ) + # One workspace per metadata object, shared by every plan and grown to + # the largest profile: plans run serially, and a plan's scratch is only + # live during its own launch. The size comes from the same resolver the + # wrapper's plan() uses (its fifth result is the byte extent). + required_bytes = int( + mla_decode._resolve_sparse_mla_plan(device, rows, heads, **plan_kwargs)[4] ) + shared = metadata.fmha_plan_caches.setdefault(self.WORKSPACE_CACHE_KEY, {}) + workspace = shared.get("buffer") + if workspace is None or workspace.numel() < required_bytes: + workspace = torch.empty((required_bytes,), device=device, dtype=torch.uint8) + shared["buffer"] = workspace + wrapper = mla_decode.BatchSparseMLADecodePagedTSWrapper(workspace_buffer=workspace) + wrapper.plan(device, rows, heads, **plan_kwargs) + plan = _SparsePlan(wrapper=wrapper) cache[key] = plan return plan @@ -518,9 +525,9 @@ def forward( groups = int(output.shape[0]) if fused else None if is_fp8: - # The MLA preprocessing quantizes Q with the KV-cache scale, so both - # descales equal the layer's kv_scale_quant_orig (see - # DeepseekV4TrtllmAttention.sparse_attn_predict). + # Mirror the fallback's scheduler prologue, which sets both + # dequant_scale_q and dequant_scale_kv to the layer's + # kv_scale_quant_orig (DeepseekV4TrtllmAttention.sparse_attn_predict). dequant = self._kv_dequant() q_scale = kv_scale = dequant query = ( @@ -551,23 +558,20 @@ def forward( # Storage-row indices straight from sparse_attn_predict: the window # slots address the SWA pool and the compressed slots address the - # compress pool, both relative to their pool base. The live window - # length is the token's position within the window; the compressed - # length is what the indexer's top-k selected beyond the window. + # compress pool, both relative to their pool base. The live lengths + # come from the metadata too, derived once per step: the token's + # position within the window, and the compressed rows selected beyond + # the window. indices = forward_args.sparse_runtime_params.sparse_attn_indices positions = metadata.token_positions_cuda[start:end] - torch.add(positions, 1, out=plan.lengths) - plan.lengths.clamp_(max=window) + lengths = metadata.sparse_mla_swa_lens[start:end] swa_pool = self._pool_view(manager, DeepseekV4AttentionType.SWA, 1) extra_pool = None run_kwargs = {} if extra: - topk_lens = metadata.sparse_mla_topk_lens[ratio][start:end] - torch.sub(topk_lens, window, out=plan.extra_lengths) - plan.extra_lengths.clamp_(min=0, max=extra) extra_pool = self._pool_view(manager, DeepseekV4AttentionType.COMPRESS, ratio) run_kwargs["extra_indices"] = indices[:, window:] - run_kwargs["extra_lengths"] = plan.extra_lengths + run_kwargs["extra_lengths"] = metadata.sparse_mla_extra_lens[ratio][start:end] if fused: # TRT-LLM's fused epilogue buffers are group-major [G, T, C] with @@ -590,7 +594,7 @@ def forward( query, swa_pool, indices[:, :window], - plan.lengths, + lengths, extra_pool, softmax_scale=softmax_scale, output_scale=output_scale, diff --git a/tensorrt_llm/_torch/attention/backends/fmha/registry.py b/tensorrt_llm/_torch/attention/backends/fmha/registry.py index c0340d61b85e..3e00148b9c6a 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/registry.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/registry.py @@ -54,9 +54,8 @@ def init_fmha_libs() -> dict[str, "FmhaCls"]: "prims_ts_block_sparse": PrimsTSBlockSparseFmha, # DeepSeek-V4 generation on the vendored PrimTS native sparse MLA # kernel. Enabled by default: it only reports itself available for - # DeepSeek-V4 on SM100, so every other model and architecture - # never sees it, and it precedes the TRTLLM-Gen fallback that served - # that phase before. + # DeepSeek-V4 on SM100, so every other model and architecture never + # sees it; everything else keeps the TRTLLM-Gen libraries below. "prims_ts_sparse_mla": PrimsTSSparseMlaFmha, "flashinfer_trtllm_gen": FlashInferTrtllmGenFmha, "fallback": FallbackFmha, diff --git a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py index 7739fc12685e..4ed9a04e968b 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/deepseek_v4/metadata.py @@ -321,6 +321,28 @@ def __post_init__(self): ) for compress_ratio in self.compress_ratio_set } + # Live per-source lengths for the dual-pool FMHA libraries, derived once + # per step next to sparse_mla_topk_lens: SWA rows live per token + # (min(pos + 1, window)) and compressed rows live per compress ratio + # (sparse_mla_topk_lens[ratio] - window). Shape: [max_num_tokens]. + self.sparse_mla_swa_lens = self.get_empty( + self.cuda_graph_buffers, + (self.max_num_tokens,), + cache_name="sparse_mla_swa_lens", + dtype=torch.int32, + capture_graph=capture_graph, + ) + self.sparse_mla_extra_lens = { + compress_ratio: self.get_empty( + self.cuda_graph_buffers, + (self.max_num_tokens,), + cache_name=f"sparse_mla_extra_lens_{compress_ratio}", + dtype=torch.int32, + capture_graph=capture_graph, + ) + for compress_ratio in self.compress_ratio_set + if compress_ratio > 1 + } # cached_token_lens_cuda: number of tokens already cached per request self.cached_token_lens_cuda = self.get_empty( @@ -528,6 +550,8 @@ def prepare_for_deepseek_v4_indices(self, token_positions=None): self.compressed_local_indices_cuda, self.sparse_mla_topk_lens, self._compress_ratios_sorted, + self.sparse_mla_swa_lens, + self.sparse_mla_extra_lens, ) @staticmethod @@ -541,8 +565,10 @@ def _prepare_deepseek_v4_indices_compiled( compressed_local_indices_buf: torch.Tensor, sparse_mla_topk_lens_bufs: Dict[int, torch.Tensor], compress_ratios: list, + swa_lens_buf: torch.Tensor, + extra_lens_bufs: Dict[int, torch.Tensor], ): - """Build SWA indices, compressed indices, and topk_lens in one fused graph.""" + """Build SWA indices, compressed indices, topk_lens and live lengths in one fused graph.""" device = token_positions.device num_tokens = token_positions.shape[0] @@ -580,6 +606,7 @@ def _prepare_deepseek_v4_indices_compiled( # For ratio>1, the SWA region always occupies exactly window_size (128) # slots. Invalid SWA positions are padded with -1 in the index buffer. kv_lens = token_positions + 1 + swa_lens_buf[:num_tokens] = kv_lens.clamp(max=window_size).to(torch.int32) for compress_ratio in compress_ratios: if compress_ratio == 1: total_count = kv_lens.clamp(max=window_size) @@ -592,6 +619,8 @@ def _prepare_deepseek_v4_indices_compiled( else: raise ValueError(f"Unsupported compress_ratio: {compress_ratio}") sparse_mla_topk_lens_bufs[compress_ratio][:num_tokens] = total_count.to(torch.int32) + if compress_ratio > 1: + extra_lens_bufs[compress_ratio][:num_tokens] = compressed_count.to(torch.int32) def _build_cache_buffer_data_pointers( self, manager: "DeepseekV4CacheManager", compress_ratios_by_layer: list[int] diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index e1ffa0a9a87b..1a30a7bd5c05 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -107,7 +107,7 @@ l0_b200: # by absolute index, and they gate on >= 2 GPUs plus Blackwell. They also # rebind mpi4py's serializer at import time, which applies process-wide on # every stage that collects them, including where the tests themselves skip. - - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py --ignore=unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py --ignore=unittest/_torch/attention/sparse/msa/test_msa_proxy_split.py TIMEOUT (90) + - unittest/_torch/attention --ignore=unittest/_torch/attention/sparse/test_cute_dsl_gvr_topk_decode.py --ignore=unittest/_torch/attention/sparse/msa/test_minimax_m3_sparse_attn_decode.py --ignore=unittest/_torch/attention/sparse/msa/test_msa_proxy_split.py --ignore=unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py TIMEOUT (90) - unittest/_torch/executor/kv_cache/test_mixed_precision_kv_cache.py # Keep cold-JIT MSA coverage out of the 60-minute attention aggregate. # Cold-JIT runs take about 25 minutes; allow a 45-minute budget. 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 ac1ae79ca025..b78a685c676f 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 @@ -60,9 +60,8 @@ def test_registry_entry_is_default_before_fallback(monkeypatch) -> None: def test_sm103_stays_on_the_fallback_until_qualified(monkeypatch) -> None: """B300/GB300 report an SM reason; SM100 proceeds to the shape and toolchain checks. - The vendored kernels admit SM103 through their architecture guard, but the - flashinfer signoff is SM100 only and L0 #64329 hit an illegal memory access - in the GB300 DeepSeek-V4-Pro warmup, so the adapter keeps SM103 on TRTLLM-Gen. + The vendored kernels admit SM103 through their architecture guard, but they + are signed off on SM100 only, so the adapter keeps SM103 on TRTLLM-Gen. """ import tensorrt_llm._torch.attention.backends.fmha.prims_ts_sparse_mla as mod @@ -88,8 +87,7 @@ def test_softmax_scale_follows_qk_head_dim() -> None: """The scale uses ``qk_nope + qk_rope`` like attentionOp.cpp, not the 512 latent. DeepSeek-V4 checkpoints have 448 + 64 = 512, so only a narrower nope width - (the 128 + 64 scenario of ``test_deepseek_v4_sparse_mla``) tells the two apart; - L0 #64329 caught the latent-width version as a ~0.3 BF16 output mismatch. + (the 128 + 64 scenario of ``test_deepseek_v4_sparse_mla``) tells the two apart. """ attn = SimpleNamespace(qk_nope_head_dim=128, qk_rope_head_dim=64, q_scaling=0.5336) expected = 1.0 / (0.5336 * math.sqrt(192.0)) @@ -158,10 +156,18 @@ def test_plans_are_shared_through_the_metadata_cache(monkeypatch) -> None: planned = [] class _FakeWrapper: + def __init__(self, workspace_buffer=None): + self.workspace_buffer = workspace_buffer + def plan(self, device, rows, heads, **kwargs): - planned.append((device, rows, heads, kwargs)) + planned.append((device, rows, heads, kwargs, self.workspace_buffer)) + + def fake_resolve(device, rows, heads, **kwargs): + # The wrapper's resolver returns the workspace byte extent fifth. + return (device, None, None, {}, 1024 * (1 + kwargs["max_extra_topk"] // 256), None, None) monkeypatch.setattr(mla_decode, "BatchSparseMLADecodePagedTSWrapper", _FakeWrapper) + monkeypatch.setattr(mla_decode, "_resolve_sparse_mla_plan", fake_resolve) class _Attn: # weak-referenceable stand-in for TrtllmAttention num_heads = 16 @@ -186,12 +192,16 @@ class _Attn: # weak-referenceable stand-in for TrtllmAttention assert planned[0][3]["max_topk"] == 128 and planned[0][3]["max_extra_topk"] == 256 cache = metadata.fmha_plan_caches[PrimsTSSparseMlaFmha.PLAN_CACHE_KEY] assert list(cache.values()) == [plan_a] + workspace = metadata.fmha_plan_caches[PrimsTSSparseMlaFmha.WORKSPACE_CACHE_KEY]["buffer"] + assert planned[0][4] is workspace and workspace.numel() == 2048 - other = layer_a._get_plan(**{**profile, "extra": 0}) - assert other is not plan_a - assert other.extra_lengths is None - assert len(planned) == 2 - assert len(cache) == 2 + # A smaller profile reuses the shared workspace; a larger one grows it. + smaller = layer_a._get_plan(**{**profile, "extra": 0}) + assert smaller is not plan_a and planned[1][4] is workspace + larger = layer_b._get_plan(**{**profile, "extra": 512}) + grown = metadata.fmha_plan_caches[PrimsTSSparseMlaFmha.WORKSPACE_CACHE_KEY]["buffer"] + assert larger is not plan_a and planned[2][4] is grown and grown.numel() == 3072 + assert len(planned) == 3 and len(cache) == 3 def _rope_table(max_positions: int, device: torch.device) -> torch.Tensor: @@ -211,6 +221,119 @@ def _random_indices(rows: int, capacity: int, pool_rows: int, device: torch.devi return indices, lengths +def _two_source_case(rows: int, heads: int, topk: int, extra_topk: int, device, dtype): + """Pools, query and two index sets for one two-source BF16/FP8 attention.""" + pool_rows, extra_rows = 2048, 1024 + pool = torch.randn(pool_rows, 1, LATENT, dtype=torch.bfloat16, device=device).to(dtype) + extra_pool = torch.randn(extra_rows, 1, LATENT, dtype=torch.bfloat16, device=device).to(dtype) + query = (torch.randn(rows, 1, heads, LATENT, dtype=torch.bfloat16, device=device) * 0.5).to( + dtype + ) + sets = [] + for _ in range(2): + indices, lengths = _random_indices(rows, topk, pool_rows, device) + extra_indices, extra_lengths = _random_indices(rows, extra_topk, extra_rows, device) + sets.append((indices, lengths, extra_indices, extra_lengths)) + return pool, extra_pool, query, sets + + +def test_reloaded_kernel_accepts_the_same_launch(): + """A kernel reloaded from the persistent compile cache runs the adapter's launch. + + CI processes always compile cold, while every later process (and + 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. + """ + _skip_unless_toolchain() + from tensorrt_llm._torch.attention.backends.prims_ts import mla_decode + + torch.manual_seed(1) + device = torch.device("cuda") + rows, heads, topk, extra_topk = 5, 64, 128, 256 + pool, extra_pool, query, sets = _two_source_case( + rows, heads, topk, extra_topk, device, torch.bfloat16 + ) + indices, lengths, extra_indices, extra_lengths = sets[0] + + def run_fresh_plan() -> torch.Tensor: + wrapper = mla_decode.BatchSparseMLADecodePagedTSWrapper() + wrapper.plan(device, rows, heads, max_topk=topk, max_extra_topk=extra_topk) + out = torch.empty(rows, 1, heads, LATENT, dtype=torch.bfloat16, device=device) + 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=out, + ) + torch.cuda.synchronize() + return out + + first = run_fresh_plan() + mla_decode._get_compiled_mla_decode.cache_clear() + second = run_fresh_plan() + torch.testing.assert_close(second, first, rtol=0, atol=0) + + +def test_graph_replay_follows_new_indices_and_lengths(): + """A captured run reads the live index and length buffers on replay. + + Warm up eagerly, capture one run with ``validate=False``, swap in a second + index set and compare the replayed output with an eager run on that set. + """ + _skip_unless_toolchain() + from tensorrt_llm._torch.attention.backends.prims_ts import mla_decode + + torch.manual_seed(2) + device = torch.device("cuda") + rows, heads, topk, extra_topk = 8, 64, 128, 256 + pool, extra_pool, query, sets = _two_source_case( + rows, heads, topk, extra_topk, device, torch.bfloat16 + ) + buffers = [tensor.clone() for tensor in sets[0]] # fixed-address launch inputs + indices, lengths, extra_indices, extra_lengths = buffers + out = torch.empty(rows, 1, heads, LATENT, dtype=torch.bfloat16, device=device) + wrapper = mla_decode.BatchSparseMLADecodePagedTSWrapper() + wrapper.plan(device, rows, heads, max_topk=topk, max_extra_topk=extra_topk) + + def launch(validate: bool) -> None: + 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=out, + validate=validate, + ) + + launch(validate=True) # eager warmup on set 0, also prepares the scalar cache + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + launch(validate=False) + for buffer, source in zip(buffers, sets[1]): + buffer.copy_(source) + graph.replay() + torch.cuda.synchronize() + replayed = out.clone() + launch(validate=True) # eager reference on set 1 through the same plan + torch.cuda.synchronize() + torch.testing.assert_close(replayed, out, rtol=0, atol=0) + assert not torch.equal(replayed, torch.zeros_like(replayed)) + + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn], ids=["bf16", "fp8"]) @pytest.mark.parametrize( "rows,heads,groups,topk,extra_topk", diff --git a/tests/unittest/_torch/attention/test_prims_ts_vendored_imports.py b/tests/unittest/_torch/attention/test_prims_ts_vendored_imports.py new file mode 100644 index 000000000000..78641108b537 --- /dev/null +++ b/tests/unittest/_torch/attention/test_prims_ts_vendored_imports.py @@ -0,0 +1,30 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The vendored PrimTS tree must stay self-contained. + +TRT-LLM pins a flashinfer-python release whose own ``flashinfer.attention`` +tree is older than the vendored one and whose ``flashinfer_api`` has no trace +templates. A vendored module that imports either would silently bind to the +installed copy, so the compatibility patch rewrites those imports; this test +keeps a future re-vendor from reintroducing them. +""" + +import pathlib +import re + +import tensorrt_llm._torch.attention.backends.prims_ts as vendored + +_ESCAPING_IMPORT = re.compile( + r"^\s*(?:from\s+flashinfer\.(?:attention|trace)\b|import\s+flashinfer\.(?:attention|trace)\b)", + re.MULTILINE, +) + + +def test_vendored_prims_ts_does_not_import_flashinfer_attention_or_trace() -> None: + root = pathlib.Path(vendored.__file__).resolve().parent + offenders = sorted( + str(path.relative_to(root)) + for path in root.rglob("*.py") + if _ESCAPING_IMPORT.search(path.read_text(encoding="utf-8")) + ) + assert not offenders, f"vendored modules import the installed flashinfer tree: {offenders}" From fee152700804d3c8e3a35a52bd3ad45a8113b6b6 Mon Sep 17 00:00:00 2001 From: yunruis <205571022+yunruis@users.noreply.github.com> Date: Wed, 21 Oct 2026 01:41:01 -0700 Subject: [PATCH 12/19] [None][perf] PrimTS sparse MLA: promise valid index prefixes unless holes are possible DeepSeek-V4's index producers emit valid prefixes: the window slots are the first min(pos + 1, window) rows, the compressed slots the first topk_lens - window rows (ratio 128 by construction, ratio 4 from the exact indexer top-k, which writes its live entries first and -1 only in the tail). Plan with assume_valid_prefix=True so the kernel derives its masks from the live lengths instead of loading and checking every index. Two configurations can leave -1 inside a live prefix and keep the generic hole-checking path: KV-cache offload (non-resident pages are masked to -1 in place) and the opt-in heuristic (GVR) top-k engines, which pad from their own candidate count. The flag is part of the kernel's compile signature (a separate variant), and the first eager run of each plan still validates the promise (run(validate=True)). Verified on B200: tests/unittest/_torch/attention/sparse/deepseek_v4 (27 passed, including the two new plan-flag unit tests). The promise also held under validate=True on every first plan use of the exclusive-node decode benchmarks (real indexer output, short and long prompts). Those benchmarks show no measurable steady-state difference from the promise (numbers in the PR description). AI-assisted. Co-Authored-By: Claude Fable 5.1 Signed-off-by: yunruis <205571022+yunruis@users.noreply.github.com> --- .../backends/fmha/prims_ts_sparse_mla.py | 26 +++++++--- .../test_prims_ts_sparse_mla_fmha.py | 47 ++++++++++++++++++- 2 files changed, 66 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py index 192bf2d96eb5..56cd7b141bb6 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py @@ -26,8 +26,10 @@ ``DeepseekV4TrtllmAttention.sparse_attn_predict`` already produces: SWA slots are storage rows of the SWA pool and compressed slots are storage rows of the layer's compress pool, so no page-table conversion happens here. The kernel -takes the two index lists and their live lengths directly and masks ``-1`` -holes itself. The fused DSv4 epilogue (conjugate RoPE on ``D[448:512]`` + +takes the two index lists and their live lengths directly; ``-1`` holes before a +live length are masked in the kernel only on the configurations that can produce +them (KV-cache offload, heuristic top-k), otherwise the plan promises valid +prefixes. The fused DSv4 epilogue (conjugate RoPE on ``D[448:512]`` + 1x128 block-scaled E4M3 output in the group-major ``[G, T, (H/G)*512]`` layout with FP32 scales ``[G, (H/G)*4, pad4(T)]``) maps one-to-one onto the kernel's ``output_quant_block_size=128`` / ``out_scale_format="fp32"`` mode, so the @@ -457,6 +459,21 @@ def _get_plan( "PrimTS sparse MLA plans must be created during eager warmup, not under " f"CUDA graph capture (rows={rows}, extra={extra}, dtype={q_dtype})." ) + # DeepSeek-V4's producers emit valid prefixes: the window slots are + # the first min(pos + 1, window) rows, the compressed slots the first + # `topk_lens - window` rows (ratio 128 by construction, ratio 4 from + # the indexer's top-k), with -1 only in the tail. Promising that lets + # the kernel bound its work by the live lengths instead of scanning + # every capacity tile for holes. KV-cache offload breaks the promise + # (non-resident pages are masked to -1 in place), and the opt-in + # heuristic (GVR) top-k engines pad from their own candidate count, + # so both keep the generic hole-checking path. The first eager run of + # every plan validates the promise (run(validate=True)); the fields + # are read directly so a renamed config cannot silently fail open. + sparse_cfg = attn.sparse_attention_config + holes_possible = bool(sparse_cfg.enable_kv_cache_offload) or bool( + sparse_cfg.enable_heuristic_topk + ) plan_kwargs = dict( max_topk=window, max_extra_topk=extra, @@ -465,10 +482,7 @@ def _get_plan( kv_layout="NHD", has_sinks=has_sinks, return_lse=False, - # Window slots and compressed slots are valid prefixes in the - # dual-pool layout, but the ratio-4 indexer top-k may carry holes; - # keep the generic mask path. - assume_valid_prefix=False, + assume_valid_prefix=not holes_possible, ) if fused_groups: plan_kwargs.update( 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 b78a685c676f..652b23d36ee7 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 @@ -171,7 +171,9 @@ def fake_resolve(device, rows, heads, **kwargs): class _Attn: # weak-referenceable stand-in for TrtllmAttention num_heads = 16 - sparse_attention_config = SimpleNamespace(window_size=128) + sparse_attention_config = SimpleNamespace( + window_size=128, enable_kv_cache_offload=False, enable_heuristic_topk=False + ) attn_a, attn_b = _Attn(), _Attn() layer_a, layer_b = PrimsTSSparseMlaFmha(attn_a), PrimsTSSparseMlaFmha(attn_b) @@ -190,6 +192,10 @@ class _Attn: # weak-referenceable stand-in for TrtllmAttention assert plan_a is plan_b assert len(planned) == 1 assert planned[0][3]["max_topk"] == 128 and planned[0][3]["max_extra_topk"] == 256 + # DeepSeek-V4's index producers emit valid prefixes, so the plan promises + # them unless KV-cache offload (which masks non-resident pages to -1 in + # place) or heuristic top-k is on; see test_plan_keeps_the_hole_checking_path. + assert planned[0][3]["assume_valid_prefix"] is True cache = metadata.fmha_plan_caches[PrimsTSSparseMlaFmha.PLAN_CACHE_KEY] assert list(cache.values()) == [plan_a] workspace = metadata.fmha_plan_caches[PrimsTSSparseMlaFmha.WORKSPACE_CACHE_KEY]["buffer"] @@ -204,6 +210,45 @@ class _Attn: # weak-referenceable stand-in for TrtllmAttention assert len(planned) == 3 and len(cache) == 3 +@pytest.mark.parametrize("offload, heuristic_topk", [(True, False), (False, True)]) +def test_plan_keeps_the_hole_checking_path(monkeypatch, offload, heuristic_topk) -> None: + """Offload (pages masked to -1 in place) and the heuristic GVR top-k (pads from its + own candidate count) can leave -1 inside a live prefix, so no prefix promise.""" + import tensorrt_llm._torch.attention.backends.prims_ts.mla_decode as mla_decode + + planned = [] + + class _FakeWrapper: + def __init__(self, workspace_buffer=None): + pass + + def plan(self, device, rows, heads, **kwargs): + planned.append(kwargs) + + monkeypatch.setattr(mla_decode, "BatchSparseMLADecodePagedTSWrapper", _FakeWrapper) + monkeypatch.setattr( + mla_decode, "_resolve_sparse_mla_plan", lambda *a, **k: (None,) * 4 + (1024, None, None) + ) + + class _Attn: + num_heads = 16 + sparse_attention_config = SimpleNamespace( + window_size=128, enable_kv_cache_offload=offload, enable_heuristic_topk=heuristic_topk + ) + + attn = _Attn() # the library only holds a weak reference to its attention module + PrimsTSSparseMlaFmha(attn)._get_plan( + metadata=SimpleNamespace(fmha_plan_caches={}), + device=torch.device("cuda"), + rows=4, + extra=256, + q_dtype=torch.bfloat16, + has_sinks=False, + fused_groups=None, + ) + assert planned[0]["assume_valid_prefix"] is False + + def _rope_table(max_positions: int, device: torch.device) -> torch.Tensor: """Forward [cos(32), sin(32)] rows, the layout both kernels consume.""" inv_freq = 1.0 / ( From 0ce1bdae86fb7789a8fa3d7227b7cdc7374af358 Mon Sep 17 00:00:00 2001 From: yunruis <205571022+yunruis@users.noreply.github.com> Date: Wed, 21 Oct 2026 21:44:40 -0700 Subject: [PATCH 13/19] [None][perf] PrimTS sparse MLA: warm kernel variants and validate each index contract once per process The generation-row count changes on almost every step of a batch's context ramp, and every mixed (context + generation) step runs eager. Each new row count planned a wrapper whose first launch validated the index contract with device-to-host checks, and the first plan of each kernel variant (the split-KV schedule changes with the row count and the extra capacity) loaded it from the persistent compile cache (~0.25 s). Both stalled the host while the GPU ran prefill: nsys on an exclusive B200 (DeepSeek-V4-Flash TP4, ~3000-token prompts, concurrency 128, CUDA graphs) showed equal GPU busy time and +0.5 s GPU idle per rank over the ramp, with +2178 cudaStreamSynchronize calls, and TRTLLM-Gen 6 % faster per decode step on that configuration while every other measured configuration was at parity. The adapter now warms a profile's kernel variants once per process on the first eager forward of the eager metadata (never on CUDA graph metadata or under capture): it scans the generation row counts a step can carry, plans one representative per distinct kernel compile signature, largest first so the shared workspace is allocated once at its final size. Each layer group (window-only, ratio-4 and ratio-128 compressed layers) has its own extra capacity and variants and warms separately. The index contract is validated once per process and contract (device, window, extra capacity, compress ratio, hole policy, query dtype, the pools' addresses and sizes) instead of once per plan, so the KV-cache estimation executor's temporary pools and the serving pools are each validated once. Plan creation, warming and validation log at debug level. Exclusive B200, graph mode, long prompts, concurrency 128: 1.06x -> 1.007x of the TRTLLM-Gen step time (41.41 vs 41.12 ms, 1.008x in a second run); steady-state decode steps within 0.5 %. The remaining ~7-15 ms per mixed step is the wrapper's per-call host cost in eager steps. AI-assisted. Co-Authored-By: Claude Fable 5.1 Signed-off-by: yunruis <205571022+yunruis@users.noreply.github.com> --- .../backends/fmha/prims_ts_sparse_mla.py | 222 ++++++++++++++---- .../test_prims_ts_sparse_mla_fmha.py | 73 ++++++ 2 files changed, 253 insertions(+), 42 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py index 56cd7b141bb6..d5dad83278cb 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py @@ -46,6 +46,7 @@ from __future__ import annotations import math +import time import weakref from dataclasses import dataclass from importlib import import_module @@ -84,6 +85,12 @@ # Same compiler-target gate as fmha/prims_ts.py: the kernels share one CuTe DSL # task-scheduling stack, so a toolchain the dense libraries refuse is refused here too. _MIN_CUTLASS_COMPILER_VERSION = "13.3" +# Process-level state shared by every library instance and metadata object: +# the static profiles whose kernel variants were warmed (compiled kernels are +# cached per process by the vendored wrapper, so one warm-up per profile is +# enough), and the index contracts that were validated with one eager launch. +_WARMED_PROFILES: set = set() +_VALIDATED_CONTRACTS: set = set() _LATENT_DIM = 512 _ROPE_DIM = 64 _FUSED_OUTPUT_BLOCK = 128 @@ -120,7 +127,6 @@ class _SparsePlan: """One planned wrapper for a static profile; every launch input comes from the metadata.""" wrapper: "BatchSparseMLADecodePagedTSWrapper" - validated: bool = False class PrimsTSSparseMlaFmha(Fmha): @@ -132,11 +138,18 @@ class PrimsTSSparseMlaFmha(Fmha): # like the block-sparse and FlashInfer wrappers: every layer that runs with # one metadata object plans each static profile once, and layers sharing a # metadata object execute serially. All plans of a metadata object share - # one kernel workspace (under WORKSPACE_CACHE_KEY), grown to the largest - # profile, so the eager metadata does not accumulate one workspace per - # distinct generation row count. + # one kernel workspace (under WORKSPACE_CACHE_KEY) that is grown to the + # largest profile planned so far (a plan keeps the buffer it was planned + # with alive until the metadata object is dropped); warming plans the + # largest row count first so the metadata that warms allocates it once. PLAN_CACHE_KEY = "prims_ts_sparse_mla" WORKSPACE_CACHE_KEY = "prims_ts_sparse_mla.workspace" + # Generation row counts scanned when warming a profile's kernel variants: + # every count up to WARM_ROW_SCAN (the split-KV schedule changes often + # at small counts), then a geometric tail up to the metadata's token + # budget. One plan is created per distinct kernel compile signature. + WARM_ROW_SCAN = 128 + WARM_ROW_TAIL = (192, 256, 384, 512, 768, 1024, 1536, 2048, 3072, 4096) def __init__(self, attn: "TrtllmAttention") -> None: super().__init__(attn) @@ -454,42 +467,13 @@ def _get_plan( plan = cache.get(key) if plan is not None: return plan + t_plan = time.perf_counter() if torch.cuda.is_current_stream_capturing(): raise RuntimeError( "PrimTS sparse MLA plans must be created during eager warmup, not under " f"CUDA graph capture (rows={rows}, extra={extra}, dtype={q_dtype})." ) - # DeepSeek-V4's producers emit valid prefixes: the window slots are - # the first min(pos + 1, window) rows, the compressed slots the first - # `topk_lens - window` rows (ratio 128 by construction, ratio 4 from - # the indexer's top-k), with -1 only in the tail. Promising that lets - # the kernel bound its work by the live lengths instead of scanning - # every capacity tile for holes. KV-cache offload breaks the promise - # (non-resident pages are masked to -1 in place), and the opt-in - # heuristic (GVR) top-k engines pad from their own candidate count, - # so both keep the generic hole-checking path. The first eager run of - # every plan validates the promise (run(validate=True)); the fields - # are read directly so a renamed config cannot silently fail open. - sparse_cfg = attn.sparse_attention_config - holes_possible = bool(sparse_cfg.enable_kv_cache_offload) or bool( - sparse_cfg.enable_heuristic_topk - ) - plan_kwargs = dict( - max_topk=window, - max_extra_topk=extra, - max_seq_len_q=1, - q_data_type=q_dtype, - kv_layout="NHD", - has_sinks=has_sinks, - return_lse=False, - assume_valid_prefix=not holes_possible, - ) - if fused_groups: - plan_kwargs.update( - output_quant_block_size=_FUSED_OUTPUT_BLOCK, - out_scale_format="fp32", - num_output_groups=fused_groups, - ) + plan_kwargs, _ = self._plan_kwargs(extra, q_dtype, has_sinks, fused_groups) # One workspace per metadata object, shared by every plan and grown to # the largest profile: plans run serially, and a plan's scratch is only # live during its own launch. The size comes from the same resolver the @@ -506,8 +490,148 @@ def _get_plan( wrapper.plan(device, rows, heads, **plan_kwargs) plan = _SparsePlan(wrapper=wrapper) cache[key] = plan + logger.debug( + f"PrimTS sparse MLA plan rows={rows} extra={extra} sinks={has_sinks} " + f"fused={fused_groups}: {(time.perf_counter() - t_plan) * 1e3:.1f} ms " + f"({len(cache)} plans on this metadata)" + ) return plan + def _holes_possible(self) -> bool: + """Whether -1 can appear inside a live index prefix with this layer's config + (KV-cache offload masks non-resident pages in place; the heuristic GVR top-k + pads from its own candidate count); read directly so a renamed field cannot + silently fail open.""" + sparse_cfg = self.attn.sparse_attention_config + return bool(sparse_cfg.enable_kv_cache_offload) or bool(sparse_cfg.enable_heuristic_topk) + + def _plan_kwargs( + self, extra: int, q_dtype: torch.dtype, has_sinks: bool, fused_groups: Optional[int] + ) -> tuple[dict, bool]: + """Return the wrapper plan() kwargs for a static profile and whether -1 holes + can appear inside a live index prefix. + + DeepSeek-V4's producers emit valid prefixes: the window slots are the + first min(pos + 1, window) rows, the compressed slots the first + `topk_lens - window` rows (ratio 128 by construction, ratio 4 from the + indexer's top-k), with -1 only in the tail. Promising that lets the + kernel derive its masks from the live lengths instead of loading and + checking every index. KV-cache offload breaks the promise (non-resident + pages are masked to -1 in place), and the opt-in heuristic (GVR) top-k + engines pad from their own candidate count, so both keep the generic + hole-checking path. The fields are read directly so a renamed config + cannot silently fail open. + """ + holes_possible = self._holes_possible() + plan_kwargs = dict( + max_topk=int(self.attn.sparse_attention_config.window_size), + max_extra_topk=extra, + max_seq_len_q=1, + q_data_type=q_dtype, + kv_layout="NHD", + has_sinks=has_sinks, + return_lse=False, + assume_valid_prefix=not holes_possible, + ) + if fused_groups: + plan_kwargs.update( + output_quant_block_size=_FUSED_OUTPUT_BLOCK, + out_scale_format="fp32", + num_output_groups=fused_groups, + ) + return plan_kwargs, holes_possible + + def _warm_kernel_variants( + self, + *, + metadata: "TrtllmAttentionMetadata", + device: torch.device, + extra: int, + q_dtype: torch.dtype, + has_sinks: bool, + fused_groups: Optional[int], + ) -> None: + """Load every kernel variant of a static profile during eager warmup. + + A batch's context ramp brings a new generation row count almost every + step, and mixed (context + generation) steps run eager. Planning a new + row count is cheap (~1 ms), but the first plan of each kernel variant + (the split-KV schedule changes with the row count and the extra + capacity) loads it from the persistent compile cache (~0.25 s), which + stalls the host while the GPU runs prefill. Compiled kernels are cached + per process by the wrapper, so this runs once per process and profile, + on the eager metadata only (CUDA graph metadata objects replay one row + count each and never need it), and never under capture. It scans the + row counts a step can carry, keeps one representative per distinct + compile signature, and plans them largest first so the shared + workspace is allocated once at its final size. + """ + if torch.cuda.is_current_stream_capturing() or getattr(metadata, "is_cuda_graph", False): + return + attn = self.attn + heads = int(attn.num_heads) + profile_key = ( + str(device), + heads, + extra, + q_dtype, + has_sinks, + fused_groups, + int(attn.sparse_attention_config.window_size), + ) + if profile_key in _WARMED_PROFILES: + return + _WARMED_PROFILES.add(profile_key) + from tensorrt_llm._torch.attention.backends.prims_ts import mla_decode + + plan_kwargs, _ = self._plan_kwargs(extra, q_dtype, has_sinks, fused_groups) + # Generation rows per step: one per request, plus its draft tokens with + # speculative decoding, within the step's token budget. + requests = int(getattr(metadata, "max_num_requests", 0) or 0) + draft = int(getattr(metadata, "max_total_draft_tokens", 0) or 0) + tokens = int(getattr(metadata, "max_num_tokens", 0) or 0) + max_rows = requests * (1 + draft) + if tokens: + max_rows = min(max_rows, tokens) if max_rows else tokens + candidates = [r for r in range(1, min(max_rows, self.WARM_ROW_SCAN) + 1)] + candidates += [r for r in self.WARM_ROW_TAIL if self.WARM_ROW_SCAN < r <= max_rows] + t_warm = time.perf_counter() + representative: dict = {} + for rows in candidates: + spec = mla_decode._resolve_sparse_mla_plan(device, rows, heads, **plan_kwargs)[2] + signature = mla_decode._mla_kernel_compile_signature(spec.kernel) + representative[signature] = rows # the largest row count of each variant + for rows in sorted(representative.values(), reverse=True): + self._get_plan( + metadata=metadata, + device=device, + rows=rows, + extra=extra, + q_dtype=q_dtype, + has_sinks=has_sinks, + fused_groups=fused_groups, + ) + logger.debug( + f"PrimTS sparse MLA warmed {len(representative)} kernel variants for extra={extra} " + f"(rows scanned up to {max_rows}): {(time.perf_counter() - t_warm) * 1e3:.0f} ms" + ) + + @staticmethod + def _claim_contract_validation(contract_key: tuple) -> bool: + """Return True exactly once per process for an index contract (never under capture). + + The contract (valid prefixes, indices inside the pools) is a property + of the index producers and the pools they address, so one validated + eager launch per process and contract is enough; validating every + plan's first launch re-synchronized the host with a GPU that was busy + with prefill on the context-ramp steps that created new plans. + """ + if torch.cuda.is_current_stream_capturing() or contract_key in _VALIDATED_CONTRACTS: + return False + _VALIDATED_CONTRACTS.add(contract_key) + logger.debug(f"PrimTS sparse MLA validating the index contract {contract_key}") + return True + # ------------------------------------------------------------------ forward def forward( self, @@ -560,15 +684,16 @@ def forward( output_scale = kv_scale sinks = self._sinks(forward_args) - plan = self._get_plan( + profile = dict( metadata=metadata, device=q.device, - rows=rows, extra=extra, q_dtype=query.dtype, has_sinks=sinks is not None, fused_groups=groups, ) + self._warm_kernel_variants(**profile) + plan = self._get_plan(rows=rows, **profile) # Storage-row indices straight from sparse_attn_predict: the window # slots address the SWA pool and the compressed slots address the @@ -600,10 +725,24 @@ def forward( else: run_kwargs["out"] = output.view(rows, 1, heads, _LATENT_DIM) - # Validate the index contract once per plan (the first layer that uses - # it) while still eager; graph replay must keep validation off because - # it synchronizes. - validate = not plan.validated and not torch.cuda.is_current_stream_capturing() + # One validated eager launch per process and contract; the key names + # the producers (window, extra capacity, compress ratio, hole policy) + # and the pools they address, so the KV-cache estimation executor's + # temporary pools and the serving pools are validated separately. + validate = self._claim_contract_validation( + ( + str(q.device), + window, + extra, + ratio, + self._holes_possible(), + str(query.dtype), + swa_pool.data_ptr(), + int(swa_pool.shape[0]), + extra_pool.data_ptr() if extra_pool is not None else 0, + int(extra_pool.shape[0]) if extra_pool is not None else 0, + ) + ) plan.wrapper.run( query, swa_pool, @@ -616,7 +755,6 @@ def forward( validate=validate, **run_kwargs, ) - plan.validated = True __all__ = ["PrimsTSSparseMlaFmha", "get_prims_ts_toolchain_unsupported_reason"] 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 652b23d36ee7..ac0497ccfc3e 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 @@ -676,3 +676,76 @@ def wrapped_gen(mla, **kwargs): got, got_sc = _dequant_fused(*results["prims"]) assert got.shape == ref.shape _assert_fused_close(got, got_sc, ref, ref_sc, channels, device) + + +def test_kernel_variants_are_warmed_once_per_profile(monkeypatch) -> None: + """The first eager forward of a static profile plans one representative row count + per distinct kernel compile signature, largest first, on the eager metadata only; + the same profile is not warmed again in the process and a second extra capacity + (another layer group) warms its own variants.""" + import tensorrt_llm._torch.attention.backends.fmha.prims_ts_sparse_mla as mod + import tensorrt_llm._torch.attention.backends.prims_ts.mla_decode as mla_decode + + planned = [] + + class _FakeWrapper: + def __init__(self, workspace_buffer=None): + pass + + def plan(self, device, rows, heads, **kwargs): + planned.append((rows, kwargs["max_extra_topk"])) + + class _FakeKernel: + def __init__(self, rows): + self.rows = rows + + def compile_signature(self): + return (self.rows // 8,) # one variant per 8 row counts + + def fake_resolve(device, rows, heads, **kwargs): + return (None, None, SimpleNamespace(kernel=_FakeKernel(rows)), None, 1024, None, None) + + monkeypatch.setattr(mla_decode, "BatchSparseMLADecodePagedTSWrapper", _FakeWrapper) + monkeypatch.setattr(mla_decode, "_resolve_sparse_mla_plan", fake_resolve) + monkeypatch.setattr(mod, "_WARMED_PROFILES", set()) + + class _Attn: + num_heads = 16 + sparse_attention_config = SimpleNamespace( + window_size=128, enable_kv_cache_offload=False, enable_heuristic_topk=False + ) + + attn = _Attn() + lib = PrimsTSSparseMlaFmha(attn) + profile = dict( + device=torch.device("cuda"), q_dtype=torch.bfloat16, has_sinks=False, fused_groups=None + ) + eager = SimpleNamespace( + fmha_plan_caches={}, max_num_tokens=20, max_num_requests=20, is_cuda_graph=False + ) + lib._warm_kernel_variants(metadata=eager, extra=512, **profile) + assert planned == [(20, 512), (15, 512), (7, 512)] # one per signature, largest first + lib._warm_kernel_variants(metadata=eager, extra=512, **profile) + assert len(planned) == 3 # same profile: already warmed in this process + graph = SimpleNamespace( + fmha_plan_caches={}, max_num_tokens=20, max_num_requests=20, is_cuda_graph=True + ) + lib._warm_kernel_variants(metadata=graph, extra=32, **profile) + assert len(planned) == 3 # CUDA graph metadata never warms + lib._warm_kernel_variants(metadata=eager, extra=32, **profile) + assert planned[3:] == [(20, 32), (15, 32), (7, 32)] # another layer group warms its own + assert lib._get_plan(metadata=eager, rows=15, extra=512, **profile) is not None + assert len(planned) == 6 # the warmed plan is reused + + +def test_index_contract_is_validated_once_per_process(monkeypatch) -> None: + """One validated eager launch per contract (producers + pools); the same contract never + re-validates, because that synchronizes the host inside the context ramp.""" + import tensorrt_llm._torch.attention.backends.fmha.prims_ts_sparse_mla as mod + + monkeypatch.setattr(mod, "_VALIDATED_CONTRACTS", set()) + key = ("cuda:0", 128, 512, 4, False, "torch.bfloat16", 1, 4096, 2, 4096) + assert PrimsTSSparseMlaFmha._claim_contract_validation(key) is True + assert PrimsTSSparseMlaFmha._claim_contract_validation(key) is False + other_pools = key[:6] + (3, 4096, 4, 4096) + assert PrimsTSSparseMlaFmha._claim_contract_validation(other_pools) is True From 37c7b7d94961045d69c2d80fe5b61a695463751e Mon Sep 17 00:00:00 2001 From: yunruis <205571022+yunruis@users.noreply.github.com> Date: Wed, 21 Oct 2026 22:45:10 -0700 Subject: [PATCH 14/19] [None][chore] Pin the vendored PrimTS tree to the TensorRT-LLM canonical branch plus the flashinfer main PrimTS commits The vendored flashinfer-prims-ts tree moves from yuxianq/flashinfer trtllm-prims-ts-dev-20261012 @ ae312f7c (TensorRT-LLM main's pin) to yunruis/flashinfer trtllm-prims-ts-dev-20261022 @ f924fd45: a linear branch on the canonical branch's current tip (ae3851ab, which adds the Blackwell ragged-Q tensor-map fixes) with the six flashinfer-ai/flashinfer main commits that touch flashinfer/attention/prims_ts since ae312f7c cherry-picked on top, in order and without conflicts: #6014, #6089, #5351, #6482, #6529 (attention kernel compile cache, FP8-to-BF16 decode) and #5434 (the native two-source sparse MLA this backend uses). The resulting prims_ts tree is byte-identical to flashinfer main @ 3bf6d3202. The compatibility patch is the one reviewed with the earlier main pin (7 files, 33 hunks: the flashinfer-python 0.6.18 import adaptations plus the positional output_quant_args launch that kernels reloaded from the compile cache need, flashinfer #6843). Nothing in tensorrt_llm/ changes. AI-assisted. Co-Authored-By: Claude Fable 5.1 Signed-off-by: yunruis <205571022+yunruis@users.noreply.github.com> --- 3rdparty/vendor_sources.lock.yaml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 7281e21d0f6f..b3af17b3d233 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -13,9 +13,9 @@ vendors: patch_digest: sha256:bc9b96a69873e0ddfcba66ab899c666cf52766976f24d850048d2a0bf8d967ff digest: sha256-tree-v1:d2d382335ccb68b55078523e63835c28b699607227afb81308cb7bb1d85845ac flashinfer-prims-ts: - url: https://github.com/flashinfer-ai/flashinfer.git - branch: main - commit: 3bf6d3202239a107989451257c6f594cb5d2f008 + url: https://github.com/yunruis/flashinfer.git + branch: trtllm-prims-ts-dev-20261022 + commit: f924fd45d1040a84a54a3f48f3978507b61064e6 source: flashinfer/attention/prims_ts destination: tensorrt_llm/_torch/attention/backends/prims_ts include: From 7cffbca9825bc754710435b84dbada27226bbc36 Mon Sep 17 00:00:00 2001 From: yunruis <205571022+yunruis@users.noreply.github.com> Date: Thu, 22 Oct 2026 08:08:12 -0700 Subject: [PATCH 15/19] [None][chore] PrimTS sparse MLA: 60-minute L0 timeout and a shorter developer-guide entry tests/unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py takes 20 min 48 s in a cold B200 process (36 cases, CuTe DSL compiles) and about 3 min from a warm compile cache; the separate L0 entry keeps the attention aggregate within its budget, and a 60-minute timeout leaves a 3x margin. The developer-guide entry for fmha/prims_ts_sparse_mla.py now states what the library serves and how to opt out, and leaves implementation details to the code (review feedback). Co-Authored-By: Claude Fable 5.1 Signed-off-by: yunruis <205571022+yunruis@users.noreply.github.com> --- .../_torch/attention/ATTENTION_DEVELOPER_GUIDE.md | 9 +++------ tests/integration/test_lists/test-db/l0_b200.yml | 2 +- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index 8cb5a8394c0c..1eff1545c3e2 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md @@ -557,12 +557,9 @@ The FMHA package is split by role: kernel as device tensors (element 0 of each), so the scales need no host read and stay valid under CUDA graph replay. - `fmha/prims_ts_sparse_mla.py` runs DeepSeek-V4 generation on the vendored - PrimTS two-source sparse MLA kernel (SM100 only; SM103 stays on TRTLLM-Gen). - It consumes the dual-pool storage-row indices from `sparse_attn_predict` and - the per-step live lengths from the DeepSeek-V4 metadata, views the SWA and - compress pools from their base pointers, keeps its plans in - `fmha_plan_caches` with one shared kernel workspace per metadata object, and - writes the fused DSv4 epilogue (E4M3 output plus FP32 block scales) in place. + PrimTS two-source sparse MLA kernel (SM100 only; SM103 stays on TRTLLM-Gen); + it is in the default library set and `TLLM_FMHA_LIBS=-prims_ts_sparse_mla` + removes it. - `fmha/prims_ts_block_sparse.py` adapts generic block-sparse requests to the vendored PrimTS contiguous and paged wrappers. Paged generation passes a live, zero-copy 2D K-page-table view with its TRT-LLM padded row stride; it diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 1a30a7bd5c05..368fc84c3750 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -121,7 +121,7 @@ l0_b200: - unittest/disaggregated/test_deepseek_v4_kv_transfer.py - unittest/disaggregated/test_fp4_mla_kv_transfer.py - unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py TIMEOUT (60) - - unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py TIMEOUT (90) + - unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py TIMEOUT (60) - unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.py TIMEOUT (60) - unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_indices_transform.py TIMEOUT (60) - unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_o_proj.py TIMEOUT (30) From 2e4eda221091d7db85b00ee7c5663ce4c35a3f9c Mon Sep 17 00:00:00 2001 From: yunruis <205571022+yunruis@users.noreply.github.com> Date: Thu, 22 Oct 2026 23:15:31 -0700 Subject: [PATCH 16/19] [None][chore] Re-pin the vendored PrimTS source branch onto the promoted canonical tip TensorRT-LLM main (#20224, promoted by #20870) moved the flashinfer-prims-ts canonical pin from ae3851ab to 0e0761ef (the #6014 SM103 block-sparse/sage tuning). The linear source branch is rebased onto that tip as yunruis/flashinfer trtllm-prims-ts-dev-20261023 @ aa709061: the duplicate #6014 cherry-pick drops out and the remaining five flashinfer main PrimTS commits (#6089, #5351, #6482, #6529, #5434) replay without conflicts. The vendored tree, patch and digests are byte-identical to the previous pin (f924fd45); only the lock's branch and commit change. AI-assisted. Co-Authored-By: Claude Fable 5.1 Signed-off-by: yunruis <205571022+yunruis@users.noreply.github.com> --- 3rdparty/vendor_sources.lock.yaml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index b3af17b3d233..57fac6d0485d 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -14,8 +14,8 @@ vendors: digest: sha256-tree-v1:d2d382335ccb68b55078523e63835c28b699607227afb81308cb7bb1d85845ac flashinfer-prims-ts: url: https://github.com/yunruis/flashinfer.git - branch: trtllm-prims-ts-dev-20261022 - commit: f924fd45d1040a84a54a3f48f3978507b61064e6 + branch: trtllm-prims-ts-dev-20261023 + commit: aa709061230a58858375c513056e7733fa74294a source: flashinfer/attention/prims_ts destination: tensorrt_llm/_torch/attention/backends/prims_ts include: From c086e387fb41420dd36f66a3a02f69ec9a442f7f Mon Sep 17 00:00:00 2001 From: yunruis <205571022+yunruis@users.noreply.github.com> Date: Fri, 23 Oct 2026 05:10:31 -0700 Subject: [PATCH 17/19] [None][perf] PrimTS sparse MLA: precompile warm-up kernel variants with worker subprocesses Without a compile cache, the first eager forward of each DeepSeek-V4 layer profile warms every kernel variant the deployment can select (Flash: 12 variants, two kernels each; Pro: 19). FlashInfer's JitSpec.build_and_load() holds the op-family file lock for the whole cute.compile, so the ranks of a launch deduplicate but serialize: the warm-up wall clock is the sum of every variant's compile time regardless of rank count (about 13 min for Flash on B100, 10 min on B200), all of it inside LLM startup. _warm_kernel_variants now probes the on-disk cache (a plan under the new _compile_cache.probe_cached() mode raises KernelNotCached instead of compiling) and, when a variant is missing, hands the deduplicated (rows, extra) list to worker subprocesses first (prims_ts_precompile.py; TLLM_PRIMTS_WARM_WORKERS, default 4, 0/1 disables; TLLM_PRIMTS_WARM_TIMEOUT bounds the wait, default 1800 s). A worker runs the same wrapper.plan() under _compile_cache.precompile_only(), which compiles through JitSpecCuteDsl.compile_and_persist() -- cute.compile outside the module lock, the artifact committed under it -- and returns a placeholder instead of loading. A per-kernel file lock in FlashInfer's tmp directory makes every kernel of a workspace compile once: worker i of N takes payload slots i, i+N, ... first, then goes through the rest skipping kernels another process holds (KernelBusy) and finishes with a blocking pass, so every worker ends with every variant on disk whether the workspace is private to the rank (TRT-LLM's default) or shared between ranks. Workers are plain interpreter subprocesses with the MPI launcher environment removed (split_mpi_env(); importing tensorrt_llm initializes MPI) and the parent's exact sys.path, because the cache key's cute_dsl_version is enumerated from it and a different meta.json makes the parent treat the module as stale and wipe it. Any failure (worker exit, timeout, cache disabled, FLASHINFER_DISABLE_JIT) falls back to the in-process compile that ran before; TLLM_PRIMTS_WARM_STRICT=1 turns it into an error for tests. Measured on dlcluster 4u4g-gen-0097 (B100), cold FLASHINFER_WORKSPACE_BASE: the 12 Flash variants take 770 s in-process versus 264 / 187 / 136 s with 4 / 8 / 12 workers, after which the serving process loads them in 2 s; the DeepSeek-V4-Flash TP4 LLM becomes ready in 827 s instead of 1369 s. A warm cache is unchanged (about 0.2 s per variant, no workers spawned). Tests: the worker entry (own shard first, busy deferral), the runner with a fake subprocess (environment, sys.path, slots, failure / cache-key mismatch / timeout paths), the probe / persist / disabled modes of build_and_load_attention_kernel, worker_count, and the adapter's probe -> fan-out -> fallback -> strict behaviour. AI-assisted. Co-Authored-By: Claude Fable 5.1 Signed-off-by: yunruis <205571022+yunruis@users.noreply.github.com> --- .../vendor_patches/flashinfer-prims-ts.patch | 232 +++++++++++- 3rdparty/vendor_sources.lock.yaml | 4 +- .../backends/fmha/prims_ts_precompile.py | 224 ++++++++++++ .../backends/fmha/prims_ts_sparse_mla.py | 61 ++++ .../backends/prims_ts/_compile_cache.py | 143 ++++++++ .../attention/backends/prims_ts/mla_decode.py | 34 +- .../test_deepseek_v4_sparse_mla.py | 6 + .../test_prims_ts_sparse_mla_fmha.py | 345 ++++++++++++++++++ 8 files changed, 1035 insertions(+), 14 deletions(-) create mode 100644 tensorrt_llm/_torch/attention/backends/fmha/prims_ts_precompile.py diff --git a/3rdparty/vendor_patches/flashinfer-prims-ts.patch b/3rdparty/vendor_patches/flashinfer-prims-ts.patch index b911288a6e6f..88315cfd30ef 100644 --- a/3rdparty/vendor_patches/flashinfer-prims-ts.patch +++ b/3rdparty/vendor_patches/flashinfer-prims-ts.patch @@ -1,5 +1,5 @@ diff --git a/_compile_cache.py b/_compile_cache.py -index b374a65311402982d818d80e9145ab40ed9727a1..db8416f07e37945a8e8764ac202e446f301d54d8 100644 +index b374a65311402982d818d80e9145ab40ed9727a1..cdfa76294f2d854117db62ad36724254006f8cb8 100644 --- a/_compile_cache.py +++ b/_compile_cache.py @@ -1,3 +1,4 @@ @@ -7,7 +7,150 @@ index b374a65311402982d818d80e9145ab40ed9727a1..db8416f07e37945a8e8764ac202e446f # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); -@@ -82,7 +83,7 @@ def build_and_load_attention_kernel( +@@ -15,10 +16,13 @@ + """Persistent compilation shared by the PrimTS attention adapters.""" + + from collections.abc import Callable ++import contextlib ++import contextvars + import dataclasses + import enum + import functools + import hashlib ++import os + from pathlib import Path + from typing import Any + +@@ -64,6 +68,128 @@ def _compile_cache_key(spec: object, options: str, extra: tuple = ()) -> str | N + return hashlib.sha256(encoded.encode()).hexdigest() + + ++# Warm-up precompile modes. Under ``precompile_only()`` the attention kernels ++# are compiled through ``JitSpecCuteDsl.compile_and_persist()`` (the ++# ``cute.compile`` runs outside the module-level file lock that ++# ``build_and_load()`` holds for a whole build, so sibling processes compile ++# different kernels concurrently), committed to disk, and not loaded: the ++# serving process loads them afterwards. Under ``probe_cached()`` nothing is ++# compiled; a kernel missing from the cache raises ``KernelNotCached``. Both ++# modes return a placeholder in place of the kernel. ++_PRECOMPILE_MODE: contextvars.ContextVar[str | None] = contextvars.ContextVar( ++ "prims_ts_precompile_mode", default=None ++) ++_SKIP_BUSY: contextvars.ContextVar[bool] = contextvars.ContextVar( ++ "prims_ts_precompile_skip_busy", default=False ++) ++ ++ ++class KernelNotCached(Exception): ++ """Raised under ``probe_cached()`` for a kernel that is not on disk.""" ++ ++ ++class KernelBusy(Exception): ++ """Raised under ``precompile_only(skip_busy=True)`` while another process compiles the kernel.""" ++ ++ ++@contextlib.contextmanager ++def precompile_only(*, skip_busy: bool = False): ++ """Compile and persist kernels without loading them (see module note). ++ ++ With ``skip_busy`` a kernel whose per-kernel lock another process holds ++ raises ``KernelBusy`` instead of waiting, so a worker can defer it and ++ compile something else first. ++ """ ++ token = _PRECOMPILE_MODE.set("persist") ++ busy_token = _SKIP_BUSY.set(skip_busy) ++ try: ++ yield ++ finally: ++ _SKIP_BUSY.reset(busy_token) ++ _PRECOMPILE_MODE.reset(token) ++ ++ ++@contextlib.contextmanager ++def probe_cached(): ++ """Report, without compiling, whether every kernel a plan needs is on disk.""" ++ token = _PRECOMPILE_MODE.set("probe") ++ try: ++ yield ++ finally: ++ _PRECOMPILE_MODE.reset(token) ++ ++ ++def expected_cache_meta() -> dict: ++ """The on-disk module metadata this process would write (arch, DSL version, source hash).""" ++ from flashinfer.jit.cute_dsl_core import ( ++ _get_compile_arch, ++ _get_cute_dsl_version, ++ _hash_source_files, ++ ) ++ ++ return { ++ "arch": _get_compile_arch(), ++ "cute_dsl_version": _get_cute_dsl_version(), ++ "source_sha256": _hash_source_files(tuple(_kernel_source_files())), ++ } ++ ++ ++def _unloaded_kernel(*args, **kwargs): ++ raise RuntimeError( ++ "PrimTS kernel was compiled in precompile-only mode and never loaded; " ++ "load it through a normal plan() outside precompile_only()." ++ ) ++ ++ ++def is_placeholder(compiled) -> bool: ++ """True for the stand-in returned under the precompile modes.""" ++ return compiled is _unloaded_kernel ++ ++ ++def _make_spec(module_name: str, kernel_name: str, compile_fn, key_files): ++ from flashinfer.jit.cute_dsl_core import JitSpecCuteDsl, _hash_source_files ++ ++ return JitSpecCuteDsl(module_name, kernel_name, compile_fn, _hash_source_files(tuple(key_files))) ++ ++ ++def _precompile_kernel(module_name: str, kernel_name: str, compile_fn, key_files) -> None: ++ """Compile one kernel into the on-disk cache unless it is already there. ++ ++ One compile per kernel per workspace: a sibling that wants the same kernel ++ waits on the per-kernel lock (or raises ``KernelBusy`` under ``skip_busy``) ++ and then finds it on disk. The lock lives in FlashInfer's tmp directory, ++ not in the module directory: creating that directory before the first ++ export would let ``compile_and_persist()`` of a sibling read a missing ++ meta.json as a stale module and wipe the exports already there. ++ """ ++ from filelock import FileLock, Timeout ++ from flashinfer.jit.core import get_tmpdir ++ ++ spec = _make_spec(module_name, kernel_name, compile_fn, key_files) ++ if spec.is_compiled: ++ return ++ lock = FileLock(get_tmpdir() / f"{spec.module_dir_name}.{spec.kernel_name}.lock", thread_local=False) ++ if _SKIP_BUSY.get(): ++ try: ++ lock.acquire(timeout=0) ++ except Timeout: ++ raise KernelBusy(spec.name) from None ++ else: ++ lock.acquire() ++ try: ++ if spec.is_compiled: ++ return ++ spec.compile_and_persist() ++ finally: ++ lock.release() ++ ++ ++def _probe_kernel(module_name: str, kernel_name: str, compile_fn, key_files) -> None: ++ spec = _make_spec(module_name, kernel_name, compile_fn, key_files) ++ if not spec.is_compiled: ++ raise KernelNotCached(spec.name) ++ ++ + @functools.cache + def _kernel_source_files() -> tuple[str, ...]: + # Family adapters share masks, reductions, sparse/Sage preparation and +@@ -82,16 +208,26 @@ def build_and_load_attention_kernel( extra: tuple = (), ) -> Any: """Build or load under the caller's selected CUDA device context.""" @@ -16,6 +159,40 @@ index b374a65311402982d818d80e9145ab40ed9727a1..db8416f07e37945a8e8764ac202e446f build_and_load_cute_dsl_kernel, cute_dsl_cache_disabled, ) + ++ mode = _PRECOMPILE_MODE.get() ++ if mode is not None and os.environ.get("FLASHINFER_DISABLE_JIT"): ++ raise RuntimeError("PrimTS warm-up precompile cannot compile with FLASHINFER_DISABLE_JIT set.") + if cute_dsl_cache_disabled(): ++ if mode is not None: ++ raise RuntimeError( ++ "PrimTS warm-up precompile needs the on-disk CuTe DSL cache " ++ "(FLASHINFER_CUTE_DSL_DISABLE_CACHE=1 is set)." ++ ) + return compile_fn() + + key = _compile_cache_key(spec, options, extra) + if key is None: ++ if mode is not None: ++ raise RuntimeError("PrimTS warm-up precompile cannot key this kernel for the on-disk cache.") + return compile_fn() + import torch + +@@ -99,6 +235,14 @@ def build_and_load_attention_kernel( + # host-selected tuning defaults. The module metadata keys the target; + # the specialization must also distinguish the physical capability. + major, minor = torch.cuda.get_device_capability() ++ if mode == "persist": ++ _precompile_kernel( ++ f"prims_ts_{family}", f"{kind}_sm{major}{minor}_{key}", compile_fn, _kernel_source_files() ++ ) ++ return _unloaded_kernel ++ if mode == "probe": ++ _probe_kernel(f"prims_ts_{family}", f"{kind}_sm{major}{minor}_{key}", compile_fn, _kernel_source_files()) ++ return _unloaded_kernel + return build_and_load_cute_dsl_kernel( + f"prims_ts_{family}", + f"{kind}_sm{major}{minor}_{key}", diff --git a/block_sparse.py b/block_sparse.py index fa055618705216cddf6b7f3e9ed4c1d680c2af9a..38de4926a045f17c5b6f206f4a20ecbe2532a1c7 100644 --- a/block_sparse.py @@ -242,7 +419,7 @@ index b22f0f413b7312bab9a5f8e296255fe9aeabe901..a4a87200ca5d3ace5cf899caa9672a1a from .helpers_common import ( _TASK_CACHE_TMEM_BASE_OFFSET, diff --git a/mla_decode.py b/mla_decode.py -index b62bd1081fbae0e18e0950008bae87c8986677b4..0555e277288fdeef25e85de1de4415a578cfe743 100644 +index b62bd1081fbae0e18e0950008bae87c8986677b4..d9cdcd13a270d4d90c6d858113047aff9c69891e 100644 --- a/mla_decode.py +++ b/mla_decode.py @@ -1,3 +1,4 @@ @@ -263,7 +440,48 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..0555e277288fdeef25e85de1de4415a5 from .decode import ( _WorkspaceSection, -@@ -1459,11 +1456,12 @@ def _launch_mla_decode( +@@ -1047,11 +1044,35 @@ def _make_mla_decode_compile_spec( + ) + + +-@functools.cache +-def _get_compiled_mla_decode( +- compile_spec: _MLADecodeCompileSpec, +-): +- """Compile and cache one batch-dynamic TS MLA topology.""" ++_COMPILED_MLA_DECODE: dict = {} ++ ++ ++def _get_compiled_mla_decode(compile_spec: _MLADecodeCompileSpec): ++ """Compile and cache one batch-dynamic TS MLA topology. ++ ++ A manual memo rather than ``functools.cache``: under the warm-up precompile ++ modes of ``_compile_cache`` the build returns a placeholder (probe / persist ++ only), which must not be remembered for the serving launches that follow. ++ """ ++ compiled = _COMPILED_MLA_DECODE.get(compile_spec) ++ if compiled is None: ++ compiled = _compile_mla_decode(compile_spec) ++ if not _compile_cache_is_placeholder(compiled): ++ _COMPILED_MLA_DECODE[compile_spec] = compiled ++ return compiled ++ ++ ++_get_compiled_mla_decode.cache_clear = _COMPILED_MLA_DECODE.clear # type: ignore[attr-defined] ++ ++ ++def _compile_cache_is_placeholder(compiled) -> bool: ++ from ._compile_cache import is_placeholder ++ ++ return is_placeholder(compiled) ++ ++ ++def _compile_mla_decode(compile_spec: _MLADecodeCompileSpec): ++ """Build (or load from the on-disk cache) one batch-dynamic TS MLA topology.""" + + import cutlass + import cutlass.cute as cute +@@ -1459,11 +1480,12 @@ def _launch_mla_decode( runtime.sparse_inputs, runtime.bmm1_scale, runtime.bmm2_scale, @@ -281,7 +499,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..0555e277288fdeef25e85de1de4415a5 ) return runtime.out -@@ -1637,7 +1635,7 @@ class BatchMLADecodePagedTSWrapper: +@@ -1637,7 +1659,7 @@ class BatchMLADecodePagedTSWrapper: split_kv=int(dict(policy)["split_kv"]), ) @@ -290,7 +508,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..0555e277288fdeef25e85de1de4415a5 def run( self, query: torch.Tensor, -@@ -1742,7 +1740,7 @@ class BatchMLADecodePagedTSWrapper: +@@ -1742,7 +1764,7 @@ class BatchMLADecodePagedTSWrapper: ) @@ -299,7 +517,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..0555e277288fdeef25e85de1de4415a5 def batch_mla_decode_with_paged_kv_cache( query: torch.Tensor, kv_cache: torch.Tensor, -@@ -2119,17 +2117,17 @@ def _resolve_sparse_mla_plan( +@@ -2119,17 +2141,17 @@ def _resolve_sparse_mla_plan( ): """Resolve sparse geometry and scratch sizing without allocation/compilation.""" import cutlass diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 57fac6d0485d..eba5560b9d96 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -21,8 +21,8 @@ vendors: include: - '**/*.py' patch: 3rdparty/vendor_patches/flashinfer-prims-ts.patch - patch_digest: sha256:09d14793745622df73c38d33cc1101db17b7a08a800a9fa5c1be0d4c43af377c - digest: sha256-tree-v1:f2800344002d0f73d2e9dc26efcf1a100dbcf2bfe5185fa29db5f2391746515d + patch_digest: sha256:ef6f887843a37314d86dbadc15db9cc1a53fb7e09fe5a7a0d4d9415cc47432ff + digest: sha256-tree-v1:ce5b8d57a8752174cc1a7d2330fdb5b6aa48463db84c7d8242f4dcec9b0e08ba openengine-proto: url: https://github.com/ai-dynamo/openengine.git tag: v0.1.0 diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_precompile.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_precompile.py new file mode 100644 index 000000000000..209ae7681b37 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_precompile.py @@ -0,0 +1,224 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Parallel precompilation of PrimTS sparse MLA kernel variants. + +``PrimsTSSparseMlaFmha._warm_kernel_variants`` plans every kernel variant a +deployment can select. Each variant is one ``cute.compile`` (tens of seconds), +and FlashInfer's on-disk CuTe DSL cache serializes the builds of one op family +behind a file lock, so N ranks warming the same variants finish no sooner than +one rank compiling them all in sequence. This module fans the variants out to +worker subprocesses that compile through ``_compile_cache.precompile_only()`` +(``cute.compile`` outside the lock, the artifact committed under it); the +serving process then loads every variant from disk. + +Workers are plain interpreter subprocesses rather than ``multiprocessing`` +children: importing ``tensorrt_llm`` initializes MPI, so a worker must not +inherit the launcher's MPI environment, and a subprocess gets a fully +controlled one. Worker ``i`` of ``N`` starts with payload slots ``i, i+N, ...`` +and then goes through the rest of the list, deferring kernels another process +is compiling (per-kernel lock) and finishing with a blocking pass, so every +worker ends with every variant on disk whether the workspace is private to the +rank (TRT-LLM's default) or shared between ranks. Failures are best-effort: +the caller falls back to compiling in-process. + +Environment: + TLLM_PRIMTS_WARM_WORKERS subprocess count, default 4; 0 or 1 disables. + TLLM_PRIMTS_WARM_TIMEOUT seconds to wait for the workers, default 1800. + TLLM_PRIMTS_WARM_STRICT 1 turns a worker failure into an error (tests). +""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +import tempfile +import time +from typing import Any, Sequence + +from tensorrt_llm.logger import logger + +_WORKER_BOOTSTRAP = """ +import json, os, sys +_path = os.environ.get("TLLM_PRIMTS_PRECOMPILE_SYS_PATH") +if _path: + sys.path[:] = json.loads(_path) +from tensorrt_llm._torch.attention.backends.fmha.prims_ts_precompile import _main +sys.exit(_main(sys.argv)) +""" + + +def worker_count(num_payloads: int) -> int: + """Subprocess count from ``TLLM_PRIMTS_WARM_WORKERS`` (default 4); ``0``/``1`` disables. + + Zero when the on-disk cache or JIT compilation is disabled: a worker could + then compile but keep nothing. + """ + if os.environ.get("FLASHINFER_CUTE_DSL_DISABLE_CACHE", "0") == "1" or os.environ.get( + "FLASHINFER_DISABLE_JIT" + ): + return 0 + raw = os.environ.get("TLLM_PRIMTS_WARM_WORKERS", "4") + try: + workers = int(raw) + except ValueError: + logger.warning(f"Ignoring TLLM_PRIMTS_WARM_WORKERS={raw!r}; expected an integer.") + workers = 4 + return max(0, min(workers, num_payloads)) + + +def strict() -> bool: + return os.environ.get("TLLM_PRIMTS_WARM_STRICT", "0") == "1" + + +def encode_payload(device_index: int, rows: int, heads: int, plan_kwargs: dict) -> dict[str, Any]: + """JSON-safe description of one ``BatchSparseMLADecodePagedTSWrapper.plan`` call.""" + kwargs = dict(plan_kwargs) + dtype = kwargs.pop("q_data_type") + kwargs["q_data_type"] = str(dtype).removeprefix("torch.") + return {"device_index": device_index, "rows": rows, "heads": heads, "plan_kwargs": kwargs} + + +def shard_order(payloads: Sequence[Any], slot: int, slots: int) -> list: + """Own shard first (``payloads[slot::slots]``), then the rest in list order.""" + mine = list(payloads[slot::slots]) + return mine + [p for p in payloads if not any(p is m for m in mine)] + + +def _worker_env() -> dict[str, str]: + from tensorrt_llm.llmapi.mpi_session import split_mpi_env + + env, _ = split_mpi_env() + env["TLLM_PRIMTS_WARM_WORKERS"] = "0" # workers never fan out again + env.setdefault("TLLM_LOG_LEVEL", "warning") + # The on-disk cache key includes the installed CuTe DSL distributions as + # enumerated along sys.path; a worker must see the parent's sys.path + # exactly (entries, order and duplicates) to write the same meta.json. + env["TLLM_PRIMTS_PRECOMPILE_SYS_PATH"] = json.dumps(sys.path) + return env + + +def run_precompile_workers(payloads: Sequence[dict[str, Any]], workers: int) -> bool: + """Compile ``payloads`` with ``workers`` subprocesses; True iff every worker succeeded. + + Worker output goes to temporary files (a pipe could fill while the worker + holds a kernel lock another worker waits on). The wait is bounded by + ``TLLM_PRIMTS_WARM_TIMEOUT``; on timeout or any error every worker is + killed and the caller compiles in-process. + """ + if workers <= 1 or not payloads: + return False + timeout_s = float(os.environ.get("TLLM_PRIMTS_WARM_TIMEOUT", "1800")) + env = _worker_env() + payload_json = json.dumps(list(payloads)) + procs: list[tuple[int, subprocess.Popen, Any]] = [] + ok = True + try: + for slot in range(workers): + out = tempfile.TemporaryFile(mode="w+", prefix="prims_ts_precompile_", suffix=".log") + proc = subprocess.Popen( + [sys.executable, "-c", _WORKER_BOOTSTRAP, payload_json, str(slot), str(workers)], + env=env, + stdout=out, + stderr=subprocess.STDOUT, + text=True, + ) + procs.append((slot, proc, out)) + deadline = time.monotonic() + timeout_s + for slot, proc, out in procs: + remaining = deadline - time.monotonic() + try: + proc.wait(timeout=max(remaining, 0.0)) + except subprocess.TimeoutExpired: + ok = False + logger.warning( + f"PrimTS precompile worker (slot {slot}) exceeded TLLM_PRIMTS_WARM_TIMEOUT=" + f"{timeout_s:g}s; the serving process compiles what is missing." + ) + continue + out.seek(0) + text = out.read() + if proc.returncode != 0: + ok = False + tail = "\n".join(text.strip().splitlines()[-15:]) + logger.warning( + f"PrimTS precompile worker (slot {slot}) exited with {proc.returncode}; " + f"the serving process compiles what is missing.\n{tail}" + ) + continue + meta_line = next( + ( + line + for line in text.splitlines() + if line.startswith("[prims-ts precompile] meta=") + ), + "", + ) + expected = str(_expected_cache_meta()) + if meta_line and meta_line[len("[prims-ts precompile] meta=") :] != expected: + ok = False + logger.warning( + f"PrimTS precompile worker (slot {slot}) wrote a different cache key than this " + f"process expects; its kernels will be recompiled.\n worker: {meta_line}\n parent: {expected}" + ) + else: + logger.debug(f"PrimTS precompile worker (slot {slot}) finished") + finally: + for _, proc, out in procs: + if proc.poll() is None: + proc.kill() + proc.wait() + out.close() + return ok + + +def _expected_cache_meta() -> dict: + from tensorrt_llm._torch.attention.backends.prims_ts import _compile_cache + + return _compile_cache.expected_cache_meta() + + +def _plan(mla_decode, payload: dict[str, Any], torch) -> None: + device = torch.device("cuda", int(payload["device_index"])) + torch.cuda.set_device(device) + kwargs = dict(payload["plan_kwargs"]) + kwargs["q_data_type"] = getattr(torch, kwargs["q_data_type"]) + mla_decode.BatchSparseMLADecodePagedTSWrapper().plan( + device, int(payload["rows"]), int(payload["heads"]), **kwargs + ) + + +def _main(argv: list[str]) -> int: + """Worker entry: compile ``payloads[slot::slots]`` first, then everything else.""" + payloads = json.loads(argv[1]) + slot, slots = (int(argv[2]), int(argv[3])) if len(argv) > 3 else (0, 1) + import torch + + from tensorrt_llm._torch.attention.backends.prims_ts import _compile_cache, mla_decode + + deferred = [] + for payload in shard_order(payloads, slot, slots): + try: + with _compile_cache.precompile_only(skip_busy=True): + _plan(mla_decode, payload, torch) + except _compile_cache.KernelBusy: + deferred.append(payload) + continue + print( + f"[prims-ts precompile] rows={payload['rows']} extra={payload['plan_kwargs'].get('max_extra_topk')} done", + flush=True, + ) + for payload in deferred: + with _compile_cache.precompile_only(): + _plan(mla_decode, payload, torch) + extra = payload["plan_kwargs"].get("max_extra_topk") + print( + f"[prims-ts precompile] rows={payload['rows']} extra={extra} done (waited)", flush=True + ) + print(f"[prims-ts precompile] meta={_compile_cache.expected_cache_meta()}", flush=True) + return 0 + + +if __name__ == "__main__": + sys.exit(_main(sys.argv)) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py index d5dad83278cb..49c7ca16534f 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py @@ -601,6 +601,9 @@ def _warm_kernel_variants( spec = mla_decode._resolve_sparse_mla_plan(device, rows, heads, **plan_kwargs)[2] signature = mla_decode._mla_kernel_compile_signature(spec.kernel) representative[signature] = rows # the largest row count of each variant + self._precompile_variants( + device, heads, plan_kwargs, sorted(representative.values(), reverse=True) + ) for rows in sorted(representative.values(), reverse=True): self._get_plan( metadata=metadata, @@ -616,6 +619,64 @@ def _warm_kernel_variants( f"(rows scanned up to {max_rows}): {(time.perf_counter() - t_warm) * 1e3:.0f} ms" ) + @staticmethod + def _precompile_variants( + device: torch.device, heads: int, plan_kwargs: dict, rows_list: list + ) -> None: + """Populate the on-disk kernel cache for ``rows_list`` with worker subprocesses. + + Skipped when every variant is already cached (a probe plans each row + count without compiling), so a warm restart spawns nothing. Otherwise + best-effort: on any failure the serial plan loop that follows compiles + in-process, exactly as without workers. See ``prims_ts_precompile``. + """ + from tensorrt_llm._torch.attention.backends.prims_ts import _compile_cache, mla_decode + + from . import prims_ts_precompile + + workers = prims_ts_precompile.worker_count(len(rows_list)) + if workers <= 1: + return + try: + with _compile_cache.probe_cached(): + for rows in rows_list: + mla_decode.BatchSparseMLADecodePagedTSWrapper().plan( + device, rows, heads, **plan_kwargs + ) + except _compile_cache.KernelNotCached: + pass + except Exception as error: # noqa: BLE001 -- a probe failure only costs the fan-out + if prims_ts_precompile.strict(): + raise + logger.warning( + f"PrimTS sparse MLA cache probe failed: {error!r}; compiling in-process." + ) + return + else: + logger.debug("PrimTS sparse MLA kernel variants are all cached; no precompile workers") + return + device_index = device.index if device.index is not None else torch.cuda.current_device() + payloads = [ + prims_ts_precompile.encode_payload(device_index, rows, heads, plan_kwargs) + for rows in rows_list + ] + t_pre = time.perf_counter() + try: + ok = prims_ts_precompile.run_precompile_workers(payloads, workers) + except Exception as error: # noqa: BLE001 -- warmup must not fail on a helper + if prims_ts_precompile.strict(): + raise + logger.warning(f"PrimTS sparse MLA parallel precompile skipped: {error!r}") + return + if not ok and prims_ts_precompile.strict(): + raise RuntimeError( + "PrimTS sparse MLA precompile workers failed (TLLM_PRIMTS_WARM_STRICT=1)" + ) + logger.info( + f"PrimTS sparse MLA precompiled {len(payloads)} kernel variants with {workers} workers " + f"in {time.perf_counter() - t_pre:.0f} s{'' if ok else ' (some workers failed)'}" + ) + @staticmethod def _claim_contract_validation(contract_key: tuple) -> bool: """Return True exactly once per process for an index contract (never under capture). diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_compile_cache.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_compile_cache.py index db8416f07e37..cdfa76294f2d 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/_compile_cache.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_compile_cache.py @@ -16,10 +16,13 @@ """Persistent compilation shared by the PrimTS attention adapters.""" from collections.abc import Callable +import contextlib +import contextvars import dataclasses import enum import functools import hashlib +import os from pathlib import Path from typing import Any @@ -65,6 +68,128 @@ def _compile_cache_key(spec: object, options: str, extra: tuple = ()) -> str | N return hashlib.sha256(encoded.encode()).hexdigest() +# Warm-up precompile modes. Under ``precompile_only()`` the attention kernels +# are compiled through ``JitSpecCuteDsl.compile_and_persist()`` (the +# ``cute.compile`` runs outside the module-level file lock that +# ``build_and_load()`` holds for a whole build, so sibling processes compile +# different kernels concurrently), committed to disk, and not loaded: the +# serving process loads them afterwards. Under ``probe_cached()`` nothing is +# compiled; a kernel missing from the cache raises ``KernelNotCached``. Both +# modes return a placeholder in place of the kernel. +_PRECOMPILE_MODE: contextvars.ContextVar[str | None] = contextvars.ContextVar( + "prims_ts_precompile_mode", default=None +) +_SKIP_BUSY: contextvars.ContextVar[bool] = contextvars.ContextVar( + "prims_ts_precompile_skip_busy", default=False +) + + +class KernelNotCached(Exception): + """Raised under ``probe_cached()`` for a kernel that is not on disk.""" + + +class KernelBusy(Exception): + """Raised under ``precompile_only(skip_busy=True)`` while another process compiles the kernel.""" + + +@contextlib.contextmanager +def precompile_only(*, skip_busy: bool = False): + """Compile and persist kernels without loading them (see module note). + + With ``skip_busy`` a kernel whose per-kernel lock another process holds + raises ``KernelBusy`` instead of waiting, so a worker can defer it and + compile something else first. + """ + token = _PRECOMPILE_MODE.set("persist") + busy_token = _SKIP_BUSY.set(skip_busy) + try: + yield + finally: + _SKIP_BUSY.reset(busy_token) + _PRECOMPILE_MODE.reset(token) + + +@contextlib.contextmanager +def probe_cached(): + """Report, without compiling, whether every kernel a plan needs is on disk.""" + token = _PRECOMPILE_MODE.set("probe") + try: + yield + finally: + _PRECOMPILE_MODE.reset(token) + + +def expected_cache_meta() -> dict: + """The on-disk module metadata this process would write (arch, DSL version, source hash).""" + from flashinfer.jit.cute_dsl_core import ( + _get_compile_arch, + _get_cute_dsl_version, + _hash_source_files, + ) + + return { + "arch": _get_compile_arch(), + "cute_dsl_version": _get_cute_dsl_version(), + "source_sha256": _hash_source_files(tuple(_kernel_source_files())), + } + + +def _unloaded_kernel(*args, **kwargs): + raise RuntimeError( + "PrimTS kernel was compiled in precompile-only mode and never loaded; " + "load it through a normal plan() outside precompile_only()." + ) + + +def is_placeholder(compiled) -> bool: + """True for the stand-in returned under the precompile modes.""" + return compiled is _unloaded_kernel + + +def _make_spec(module_name: str, kernel_name: str, compile_fn, key_files): + from flashinfer.jit.cute_dsl_core import JitSpecCuteDsl, _hash_source_files + + return JitSpecCuteDsl(module_name, kernel_name, compile_fn, _hash_source_files(tuple(key_files))) + + +def _precompile_kernel(module_name: str, kernel_name: str, compile_fn, key_files) -> None: + """Compile one kernel into the on-disk cache unless it is already there. + + One compile per kernel per workspace: a sibling that wants the same kernel + waits on the per-kernel lock (or raises ``KernelBusy`` under ``skip_busy``) + and then finds it on disk. The lock lives in FlashInfer's tmp directory, + not in the module directory: creating that directory before the first + export would let ``compile_and_persist()`` of a sibling read a missing + meta.json as a stale module and wipe the exports already there. + """ + from filelock import FileLock, Timeout + from flashinfer.jit.core import get_tmpdir + + spec = _make_spec(module_name, kernel_name, compile_fn, key_files) + if spec.is_compiled: + return + lock = FileLock(get_tmpdir() / f"{spec.module_dir_name}.{spec.kernel_name}.lock", thread_local=False) + if _SKIP_BUSY.get(): + try: + lock.acquire(timeout=0) + except Timeout: + raise KernelBusy(spec.name) from None + else: + lock.acquire() + try: + if spec.is_compiled: + return + spec.compile_and_persist() + finally: + lock.release() + + +def _probe_kernel(module_name: str, kernel_name: str, compile_fn, key_files) -> None: + spec = _make_spec(module_name, kernel_name, compile_fn, key_files) + if not spec.is_compiled: + raise KernelNotCached(spec.name) + + @functools.cache def _kernel_source_files() -> tuple[str, ...]: # Family adapters share masks, reductions, sparse/Sage preparation and @@ -88,11 +213,21 @@ def build_and_load_attention_kernel( cute_dsl_cache_disabled, ) + mode = _PRECOMPILE_MODE.get() + if mode is not None and os.environ.get("FLASHINFER_DISABLE_JIT"): + raise RuntimeError("PrimTS warm-up precompile cannot compile with FLASHINFER_DISABLE_JIT set.") if cute_dsl_cache_disabled(): + if mode is not None: + raise RuntimeError( + "PrimTS warm-up precompile needs the on-disk CuTe DSL cache " + "(FLASHINFER_CUTE_DSL_DISABLE_CACHE=1 is set)." + ) return compile_fn() key = _compile_cache_key(spec, options, extra) if key is None: + if mode is not None: + raise RuntimeError("PrimTS warm-up precompile cannot key this kernel for the on-disk cache.") return compile_fn() import torch @@ -100,6 +235,14 @@ def build_and_load_attention_kernel( # host-selected tuning defaults. The module metadata keys the target; # the specialization must also distinguish the physical capability. major, minor = torch.cuda.get_device_capability() + if mode == "persist": + _precompile_kernel( + f"prims_ts_{family}", f"{kind}_sm{major}{minor}_{key}", compile_fn, _kernel_source_files() + ) + return _unloaded_kernel + if mode == "probe": + _probe_kernel(f"prims_ts_{family}", f"{kind}_sm{major}{minor}_{key}", compile_fn, _kernel_source_files()) + return _unloaded_kernel return build_and_load_cute_dsl_kernel( f"prims_ts_{family}", f"{kind}_sm{major}{minor}_{key}", 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 0555e277288f..d9cdcd13a270 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py @@ -1044,11 +1044,35 @@ def _make_mla_decode_compile_spec( ) -@functools.cache -def _get_compiled_mla_decode( - compile_spec: _MLADecodeCompileSpec, -): - """Compile and cache one batch-dynamic TS MLA topology.""" +_COMPILED_MLA_DECODE: dict = {} + + +def _get_compiled_mla_decode(compile_spec: _MLADecodeCompileSpec): + """Compile and cache one batch-dynamic TS MLA topology. + + A manual memo rather than ``functools.cache``: under the warm-up precompile + modes of ``_compile_cache`` the build returns a placeholder (probe / persist + only), which must not be remembered for the serving launches that follow. + """ + compiled = _COMPILED_MLA_DECODE.get(compile_spec) + if compiled is None: + compiled = _compile_mla_decode(compile_spec) + if not _compile_cache_is_placeholder(compiled): + _COMPILED_MLA_DECODE[compile_spec] = compiled + return compiled + + +_get_compiled_mla_decode.cache_clear = _COMPILED_MLA_DECODE.clear # type: ignore[attr-defined] + + +def _compile_cache_is_placeholder(compiled) -> bool: + from ._compile_cache import is_placeholder + + return is_placeholder(compiled) + + +def _compile_mla_decode(compile_spec: _MLADecodeCompileSpec): + """Build (or load from the on-disk cache) one batch-dynamic TS MLA topology.""" import cutlass import cutlass.cute as cute diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py index b04787ed8944..9d6cc8cf66b6 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_sparse_mla.py @@ -48,6 +48,12 @@ from tensorrt_llm.mapping import Mapping +@pytest.fixture(autouse=True) +def _strict_precompile_workers(monkeypatch): + """A PrimTS precompile worker failure is an error here, not a logged fallback.""" + monkeypatch.setenv("TLLM_PRIMTS_WARM_STRICT", "1") + + @dataclass(kw_only=True, frozen=True) class Scenario: dtype: torch.dtype = torch.bfloat16 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 ac0497ccfc3e..43d098bb45a5 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 @@ -26,6 +26,13 @@ ROPE = 64 BLOCK = 128 + +@pytest.fixture(autouse=True) +def _strict_precompile_workers(monkeypatch): + """A precompile worker failure is an error here, not a logged fallback.""" + monkeypatch.setenv("TLLM_PRIMTS_WARM_STRICT", "1") + + pytestmark = pytest.mark.skipif( get_sm_version() not in (100, 103), reason="PrimTS sparse MLA requires SM100 or SM103" ) @@ -708,6 +715,7 @@ def fake_resolve(device, rows, heads, **kwargs): monkeypatch.setattr(mla_decode, "BatchSparseMLADecodePagedTSWrapper", _FakeWrapper) monkeypatch.setattr(mla_decode, "_resolve_sparse_mla_plan", fake_resolve) monkeypatch.setattr(mod, "_WARMED_PROFILES", set()) + monkeypatch.setenv("TLLM_PRIMTS_WARM_WORKERS", "0") # in-process planning only class _Attn: num_heads = 16 @@ -749,3 +757,340 @@ def test_index_contract_is_validated_once_per_process(monkeypatch) -> None: assert PrimsTSSparseMlaFmha._claim_contract_validation(key) is False other_pools = key[:6] + (3, 4096, 4, 4096) assert PrimsTSSparseMlaFmha._claim_contract_validation(other_pools) is True + + +def _fake_sparse_plan(monkeypatch, planned, *, signature_div=8): + """Replace the sparse wrapper and plan resolver with fakes that record plan() calls.""" + import tensorrt_llm._torch.attention.backends.prims_ts.mla_decode as mla_decode + + class _FakeWrapper: + def __init__(self, workspace_buffer=None): + pass + + def plan(self, device, rows, heads, **kwargs): + planned.append((rows, kwargs["max_extra_topk"])) + + class _FakeKernel: + def __init__(self, rows): + self.rows = rows + + def compile_signature(self): + return (self.rows // signature_div,) + + def fake_resolve(device, rows, heads, **kwargs): + return (None, None, SimpleNamespace(kernel=_FakeKernel(rows)), None, 1024, None, None) + + monkeypatch.setattr(mla_decode, "BatchSparseMLADecodePagedTSWrapper", _FakeWrapper) + monkeypatch.setattr(mla_decode, "_resolve_sparse_mla_plan", fake_resolve) + return mla_decode + + +def _attn_and_profile(): + class _Attn: + num_heads = 16 + sparse_attention_config = SimpleNamespace( + window_size=128, enable_kv_cache_offload=False, enable_heuristic_topk=False + ) + + attn = _Attn() # the library holds a weak reference to its owner + profile = dict( + device=torch.device("cuda", 0), q_dtype=torch.bfloat16, has_sinks=False, fused_groups=None + ) + eager = SimpleNamespace( + fmha_plan_caches={}, max_num_tokens=20, max_num_requests=20, is_cuda_graph=False + ) + return attn, profile, eager + + +def test_warm_up_fans_out_only_when_a_variant_is_missing(monkeypatch) -> None: + """The warm-up probes the cache first: all cached -> no workers; a miss -> every + representative row count goes to the workers as a JSON-safe payload before the + in-process plan loop; a failing helper never breaks the warm-up; strict mode raises.""" + import tensorrt_llm._torch.attention.backends.fmha.prims_ts_precompile as pp + import tensorrt_llm._torch.attention.backends.fmha.prims_ts_sparse_mla as mod + import tensorrt_llm._torch.attention.backends.prims_ts._compile_cache as cc + + planned, fanned = [], [] + _fake_sparse_plan(monkeypatch, planned) + monkeypatch.setattr(mod, "_WARMED_PROFILES", set()) + monkeypatch.setenv("TLLM_PRIMTS_WARM_WORKERS", "2") + monkeypatch.delenv("FLASHINFER_CUTE_DSL_DISABLE_CACHE", raising=False) + monkeypatch.delenv("FLASHINFER_DISABLE_JIT", raising=False) + + def fake_workers(payloads, workers): + fanned.append((list(payloads), workers)) + return True + + monkeypatch.setattr(pp, "run_precompile_workers", fake_workers) + attn, profile, eager = _attn_and_profile() + lib = PrimsTSSparseMlaFmha(attn) + + # 1. Every variant cached: the probe plans (under probe mode) and no worker runs. + lib._warm_kernel_variants(metadata=eager, extra=512, **profile) + assert fanned == [] + # probe planned the three variants, then the real loop planned them again + assert planned == [(20, 512), (15, 512), (7, 512)] * 2 + + # 2. A cache miss during the probe -> workers get every representative row count. + import tensorrt_llm._torch.attention.backends.prims_ts.mla_decode as mla_decode + + real_plan = mla_decode.BatchSparseMLADecodePagedTSWrapper.plan + + def probing_plan(self, device, rows, heads, **kwargs): + if cc._PRECOMPILE_MODE.get() == "probe": + raise cc.KernelNotCached("missing") + return real_plan(self, device, rows, heads, **kwargs) + + monkeypatch.setattr(mla_decode.BatchSparseMLADecodePagedTSWrapper, "plan", probing_plan) + planned.clear() + lib._warm_kernel_variants(metadata=eager, extra=32, **profile) + assert planned == [(20, 32), (15, 32), (7, 32)] + assert len(fanned) == 1 + payloads, workers = fanned[0] + assert workers == 2 # capped by the env value, below the three variants + assert [p["rows"] for p in payloads] == [20, 15, 7] + assert payloads[0]["plan_kwargs"]["q_data_type"] == "bfloat16" # JSON-safe dtype + assert payloads[0]["plan_kwargs"]["max_extra_topk"] == 32 + assert payloads[0]["device_index"] == 0 + import json + + json.dumps(payloads) # every payload field crosses the subprocess boundary + + # 3. A failing helper is logged and the in-process loop still plans everything... + def broken_workers(payloads, workers): + raise OSError("no subprocesses here") + + monkeypatch.setattr(pp, "run_precompile_workers", broken_workers) + monkeypatch.setattr(mod, "_WARMED_PROFILES", set()) + monkeypatch.setenv("TLLM_PRIMTS_WARM_STRICT", "0") # the module fixture sets strict mode + planned.clear() + fresh = SimpleNamespace( # a new metadata object: the earlier plans are cached on `eager` + fmha_plan_caches={}, max_num_tokens=20, max_num_requests=20, is_cuda_graph=False + ) + lib._warm_kernel_variants(metadata=fresh, extra=32, **profile) + assert planned == [(20, 32), (15, 32), (7, 32)] + # ...unless strict mode asks for the failure. + monkeypatch.setenv("TLLM_PRIMTS_WARM_STRICT", "1") + monkeypatch.setattr(mod, "_WARMED_PROFILES", set()) + with pytest.raises(OSError): + lib._warm_kernel_variants(metadata=eager, extra=32, **profile) + + +def test_precompile_worker_count_follows_the_env(monkeypatch) -> None: + import tensorrt_llm._torch.attention.backends.fmha.prims_ts_precompile as pp + + monkeypatch.delenv("TLLM_PRIMTS_WARM_WORKERS", raising=False) + monkeypatch.delenv("FLASHINFER_CUTE_DSL_DISABLE_CACHE", raising=False) + monkeypatch.delenv("FLASHINFER_DISABLE_JIT", raising=False) + assert pp.worker_count(12) == 4 # default + assert pp.worker_count(3) == 3 # never more workers than variants + monkeypatch.setenv("TLLM_PRIMTS_WARM_WORKERS", "0") + assert pp.worker_count(12) == 0 + monkeypatch.setenv("TLLM_PRIMTS_WARM_WORKERS", "8") + assert pp.worker_count(12) == 8 + monkeypatch.setenv("TLLM_PRIMTS_WARM_WORKERS", "many") + assert pp.worker_count(12) == 4 # unparsable: default with a warning + monkeypatch.setenv("FLASHINFER_CUTE_DSL_DISABLE_CACHE", "1") + assert pp.worker_count(12) == 0 # nothing to persist into + monkeypatch.delenv("FLASHINFER_CUTE_DSL_DISABLE_CACHE") + monkeypatch.setenv("FLASHINFER_DISABLE_JIT", "1") + assert pp.worker_count(12) == 0 # compiling is forbidden + + +def test_precompile_worker_entry_compiles_its_shard_first_then_the_rest(monkeypatch) -> None: + """The worker entry point runs the real plan path under precompile_only(): own + shard first, the rest afterwards, busy kernels deferred to a blocking pass.""" + import tensorrt_llm._torch.attention.backends.fmha.prims_ts_precompile as pp + import tensorrt_llm._torch.attention.backends.prims_ts._compile_cache as cc + + planned, modes = [], [] + mla_decode = _fake_sparse_plan(monkeypatch, planned) + real_plan = mla_decode.BatchSparseMLADecodePagedTSWrapper.plan + busy_once = {"rows": 15, "armed": True} + + def recording_plan(self, device, rows, heads, **kwargs): + modes.append((rows, cc._PRECOMPILE_MODE.get(), cc._SKIP_BUSY.get())) + if rows == busy_once["rows"] and busy_once["armed"] and cc._SKIP_BUSY.get(): + busy_once["armed"] = False + raise cc.KernelBusy("another worker owns it") + return real_plan(self, device, rows, heads, **kwargs) + + monkeypatch.setattr(mla_decode.BatchSparseMLADecodePagedTSWrapper, "plan", recording_plan) + monkeypatch.setattr(torch.cuda, "set_device", lambda device: None) + monkeypatch.setattr(cc, "expected_cache_meta", lambda: {"probe": "meta"}) + payloads = [ + pp.encode_payload( + 0, rows, 16, dict(max_topk=128, max_extra_topk=512, q_data_type=torch.bfloat16) + ) + for rows in (20, 15, 7, 3) + ] + import json + + assert pp._main(["worker", json.dumps(payloads), "1", "2"]) == 0 + # slot 1 of 2 owns rows 15 and 3; 15 was busy once and is retried blocking at the end + assert [m[0] for m in modes] == [15, 3, 20, 7, 15] + assert all(m[1] == "persist" for m in modes) + assert [m[2] for m in modes] == [True, True, True, True, False] + assert [r for r, _ in planned] == [3, 20, 7, 15] + assert pp.shard_order([0, 1, 2, 3, 4], 1, 2) == [1, 3, 0, 2, 4] + assert pp.shard_order([0, 1], 5, 8) == [0, 1] # a slot past the list owns nothing + + +def test_precompile_workers_get_a_clean_environment_and_bounded_wait(monkeypatch, tmp_path) -> None: + """The runner strips the MPI launcher environment, forwards sys.path and the + worker count 0, assigns one slot per worker, kills workers on timeout, and + reports a worker whose cache key differs from the parent's.""" + import tensorrt_llm._torch.attention.backends.fmha.prims_ts_precompile as pp + + spawned = [] + + class _FakeProc: + def __init__(self, argv, env, stdout, returncode=0, text="", hang=False): + self.argv, self.env, self.returncode, self._hang = argv, env, returncode, hang + stdout.write(text) + self._done = False + + def wait(self, timeout=None): + if self._hang and not getattr(self, "killed", False): + raise subprocess.TimeoutExpired(self.argv, timeout) + self._done = True # like Popen.wait() after kill() + return self.returncode + + def poll(self): + return None if (self._hang and not self._done) else self.returncode + + def kill(self): + self._done = True + self.killed = True + + import subprocess + + def make_popen(**behaviour): + def fake_popen(argv, env, stdout, stderr, text): + proc = _FakeProc(argv, env, stdout, **behaviour) + spawned.append(proc) + return proc + + return fake_popen + + monkeypatch.setenv("OMPI_COMM_WORLD_RANK", "2") + monkeypatch.setenv("SLURM_PROCID", "2") + monkeypatch.setenv("TLLM_PRIMTS_WARM_WORKERS", "4") + monkeypatch.setattr(pp, "_expected_cache_meta", lambda: {"k": 1}) + meta_ok = "[prims-ts precompile] meta={'k': 1}\n" + monkeypatch.setattr(subprocess, "Popen", make_popen(text=meta_ok)) + payloads = [{"rows": r} for r in (20, 15, 7)] + assert pp.run_precompile_workers(payloads, 2) is True + assert len(spawned) == 2 + env = spawned[0].env + assert "OMPI_COMM_WORLD_RANK" not in env and "SLURM_PROCID" not in env + assert env["TLLM_PRIMTS_WARM_WORKERS"] == "0" + import json + import sys + + assert json.loads(env["TLLM_PRIMTS_PRECOMPILE_SYS_PATH"]) == sys.path + assert [p.argv[-2:] for p in spawned] == [["0", "2"], ["1", "2"]] + assert spawned[0].argv[0] == sys.executable and spawned[0].argv[1] == "-c" + assert json.loads(spawned[0].argv[3]) == payloads + + spawned.clear() + monkeypatch.setattr( + subprocess, "Popen", make_popen(text="[prims-ts precompile] meta={'k': 2}\n") + ) + assert pp.run_precompile_workers(payloads, 1 + 1) is False # cache key mismatch is a failure + + spawned.clear() + monkeypatch.setattr(subprocess, "Popen", make_popen(returncode=3, text="boom")) + assert pp.run_precompile_workers(payloads, 2) is False + + spawned.clear() + monkeypatch.setenv("TLLM_PRIMTS_WARM_TIMEOUT", "0.01") + monkeypatch.setattr(subprocess, "Popen", make_popen(hang=True)) + assert pp.run_precompile_workers(payloads, 2) is False + assert all(getattr(p, "killed", False) for p in spawned) # nothing left running + assert pp.run_precompile_workers(payloads, 1) is False # one worker: nothing to fan out + + +def test_precompile_only_persists_without_loading(monkeypatch) -> None: + """Under precompile_only(), build_and_load_attention_kernel compiles and persists + through the unlocked path and hands back a placeholder that refuses to launch; + under probe_cached() a missing kernel raises; with the cache disabled both refuse.""" + import tensorrt_llm._torch.attention.backends.prims_ts._compile_cache as cc + + persisted, probed = [], [] + monkeypatch.setattr( + cc, + "_precompile_kernel", + lambda module, kernel, fn, files: persisted.append((module, kernel)), + ) + monkeypatch.setattr( + cc, "_probe_kernel", lambda module, kernel, fn, files: probed.append(kernel) + ) + monkeypatch.setattr(cc, "_kernel_source_files", lambda: ("a.py",)) + monkeypatch.setattr(cc, "_compile_cache_key", lambda spec, options, extra: "deadbeef") + monkeypatch.setattr( + "flashinfer.jit.cute_dsl_core.cute_dsl_cache_disabled", lambda: False, raising=False + ) + monkeypatch.delenv("FLASHINFER_DISABLE_JIT", raising=False) + compiled = [] + + def compile_fn(): + compiled.append(True) + return "kernel" + + with cc.precompile_only(): + placeholder = cc.build_and_load_attention_kernel( + "mla_decode", "main", object(), "", compile_fn + ) + assert persisted and persisted[0][0] == "prims_ts_mla_decode" + assert persisted[0][1].startswith("main_sm") and persisted[0][1].endswith("_deadbeef") + assert not compiled # the helper owns the compile; nothing was loaded in-process + with pytest.raises(RuntimeError, match="precompile-only"): + placeholder() + assert cc._PRECOMPILE_MODE.get() is None # the context manager restores the mode + with cc.probe_cached(): + cc.build_and_load_attention_kernel("mla_decode", "main", object(), "", compile_fn) + assert probed == [persisted[0][1]] + monkeypatch.setenv("FLASHINFER_DISABLE_JIT", "1") + with cc.precompile_only(), pytest.raises(RuntimeError, match="FLASHINFER_DISABLE_JIT"): + cc.build_and_load_attention_kernel("mla_decode", "main", object(), "", compile_fn) + # Outside the precompile modes FLASHINFER_DISABLE_JIT keeps the disk-cache path: + # the kernel goes through FlashInfer's builder, never a compile_fn() here. + loaded = [] + monkeypatch.setattr( + "flashinfer.jit.cute_dsl_core.build_and_load_cute_dsl_kernel", + lambda module, kernel, fn, extra_key_files=(): loaded.append((module, kernel)) or "loaded", + ) + assert ( + cc.build_and_load_attention_kernel("mla_decode", "main", object(), "", compile_fn) + == "loaded" + ) + assert loaded == [("prims_ts_mla_decode", persisted[0][1])] and not compiled + monkeypatch.delenv("FLASHINFER_DISABLE_JIT") + + +def test_precompile_placeholders_are_never_memoized(monkeypatch) -> None: + """A probe or persist build must not poison the compiled-kernel memo of the + serving process: the next real request compiles / loads the kernel.""" + import tensorrt_llm._torch.attention.backends.prims_ts._compile_cache as cc + import tensorrt_llm._torch.attention.backends.prims_ts.mla_decode as mla_decode + + calls = [] + + def fake_compile(spec): + calls.append(cc._PRECOMPILE_MODE.get()) + return cc._unloaded_kernel if cc._PRECOMPILE_MODE.get() else "real kernel" + + monkeypatch.setattr(mla_decode, "_compile_mla_decode", fake_compile) + mla_decode._get_compiled_mla_decode.cache_clear() + spec = ("fake", "compile", "spec") + with cc.probe_cached(): + assert cc.is_placeholder(mla_decode._get_compiled_mla_decode(spec)) + with cc.precompile_only(): + assert cc.is_placeholder(mla_decode._get_compiled_mla_decode(spec)) + assert ( + mla_decode._get_compiled_mla_decode(spec) == "real kernel" + ) # not the memoized placeholder + assert mla_decode._get_compiled_mla_decode(spec) == "real kernel" + assert calls == ["probe", "persist", None] # the real result is memoized + mla_decode._get_compiled_mla_decode.cache_clear() From 7281a601eb3124e9e067d011a13c20e5cd7743c7 Mon Sep 17 00:00:00 2001 From: yunruis <205571022+yunruis@users.noreply.github.com> Date: Fri, 23 Oct 2026 09:08:16 -0700 Subject: [PATCH 18/19] [None][perf] PrimTS sparse MLA: prepared per-layer launches and one request check per geometry Each eager DeepSeek-V4 decode layer paid about 140 us of host time in the PrimTS sparse MLA adapter forward on B100 (4 ranks, eager, concurrency 8), of which the wrapper's run() was about 100-120 us although the CUDA API work inside it is 13 us (nsys: one cudaLaunchKernelExC, five cuTensorMapEncodeTiled). The rest was Python: argument validation, pool view construction, scalar tensors and the kernel's descriptor arguments, rebuilt on every call although they only depend on the plan, the pools, the scales and the sinks. BatchSparseMLADecodePagedTSWrapper.run now records, at the end of a full unfused BF16 launch with a caller-owned output, everything static for the (pool storage, extra pool storage, sinks address, scales) tuple, and a later validate=False call with the same tuple goes through _run_prepared: it checks and views only the per-call query, index, length and output tensors and launches the same compiled kernel with the recorded arguments, byte-identical to the full path. One plan serves every layer of a model and each layer has its own sinks parameter, so the prepared launches are kept per key (bounded, cleared by plan()); a single slot would be evicted at every layer and never hit. The prepared launch is taken only for the shapes the adapter uses (one query token per row, no packed query, no LSE, no fused output), and it still checks the per-call tensors' device, contiguity and 16-byte alignment. The wrapper change is a generic optimization of the vendored FlashInfer tree; it ships in the vendor patch so the integration lands together and will be sent upstream to flashinfer-ai/flashinfer as a follow-up. _resolve_sparse_mla_plan is memoized: it is a pure function of its arguments, plan() copies the kernel before mutating it, and the warm-up resolves it 128 times per profile. The adapter runs the full _unsupported_reason re-check only when the request geometry (metadata, query shape and dtype, index width, output shape and dtype, fusion flag, token counts) differs from the last accepted call; the manager already selected the library through is_supported(). Attention sinks that are already contiguous FP32 per head (DeepseekV4TrtllmAttention.attn_sink) are passed through, other dtypes are converted once per source tensor. Measured on dlcluster 4x B100 (exclusive node), DeepSeek-V4 Flash, TP4, eager, before -> after: wrapper.run (1x B100 standalone, rows 32): 50 -> 20 us/call adapter forward per layer, host p50 (c8): 141 -> 91 us wrapper.run inside forward (cProfile, cum): 120 -> 48 us/call prepared-launch hits: 0% -> 96% of calls decode ms/step (c8, 960 tokens, host-noisy): 132.3 -> 128.5 The remaining per-call cost is the adapter's own bookkeeping (pool views, contract key, about 50 us) and the compiled launch (13 us of CUDA API). Tests: test_prepared_launch_matches_the_full_path_per_layer (two sinks sharing one plan both hit, bitwise equal to the full path, new pool or sinks address falls back, plan() clears), test_prepared_launch_single_pool_matches_the_full_path (SWA-only layers, extra indices rejected), test_request_check_runs_once_per_geometry and test_sinks_pass_through_or_are_prepared_once_per_layer. Co-Authored-By: Claude Fable 5.1 Signed-off-by: yunruis <205571022+yunruis@users.noreply.github.com> --- .../vendor_patches/flashinfer-prims-ts.patch | 206 ++++++++++++++++- 3rdparty/vendor_sources.lock.yaml | 4 +- .../backends/fmha/prims_ts_sparse_mla.py | 77 ++++++- .../attention/backends/prims_ts/mla_decode.py | 156 +++++++++++++ .../test_prims_ts_sparse_mla_fmha.py | 212 ++++++++++++++++++ 5 files changed, 639 insertions(+), 16 deletions(-) diff --git a/3rdparty/vendor_patches/flashinfer-prims-ts.patch b/3rdparty/vendor_patches/flashinfer-prims-ts.patch index 88315cfd30ef..f68f66608f7f 100644 --- a/3rdparty/vendor_patches/flashinfer-prims-ts.patch +++ b/3rdparty/vendor_patches/flashinfer-prims-ts.patch @@ -419,7 +419,7 @@ index b22f0f413b7312bab9a5f8e296255fe9aeabe901..a4a87200ca5d3ace5cf899caa9672a1a from .helpers_common import ( _TASK_CACHE_TMEM_BASE_OFFSET, diff --git a/mla_decode.py b/mla_decode.py -index b62bd1081fbae0e18e0950008bae87c8986677b4..d9cdcd13a270d4d90c6d858113047aff9c69891e 100644 +index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e231c98a52 100644 --- a/mla_decode.py +++ b/mla_decode.py @@ -1,3 +1,4 @@ @@ -517,7 +517,15 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..d9cdcd13a270d4d90c6d858113047aff def batch_mla_decode_with_paged_kv_cache( query: torch.Tensor, kv_cache: torch.Tensor, -@@ -2119,17 +2141,17 @@ def _resolve_sparse_mla_plan( +@@ -2099,6 +2121,7 @@ def _prepare_sparse_output_quant( + ) + + ++@functools.lru_cache(maxsize=1024) + def _resolve_sparse_mla_plan( + device, + batch_size, +@@ -2119,17 +2142,17 @@ def _resolve_sparse_mla_plan( ): """Resolve sparse geometry and scratch sizing without allocation/compilation.""" import cutlass @@ -539,6 +547,200 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..d9cdcd13a270d4d90c6d858113047aff MlaDecodeTs, ) +@@ -2264,6 +2287,9 @@ def _resolve_sparse_mla_plan( + return device, profile, spec, sections, byte_end, capacity, output_quant + + ++# Prepared launches kept per plan (one per layer sharing the plan; DeepSeek-V4 has 61). ++_PREPARED_LAUNCH_LIMIT = 256 ++ + class BatchSparseMLADecodePagedTSWrapper: + """Plan native BF16/E4M3 attention, then bind one or two source index lists. + +@@ -2297,6 +2323,11 @@ class BatchSparseMLADecodePagedTSWrapper: + self._workspace_buffer = workspace_buffer + self._state = None + self._scalar_cache = {} ++ # Prepared launches for the unfused validate=False path: everything that ++ # is static for one (plan, pools, scales, sinks) tuple, keyed by the pools' ++ # storage and the sinks address. One plan serves every layer of a model, ++ # and each layer brings its own sinks, so one entry per layer. ++ self._prepared: dict[tuple, dict] = {} + self.workspace_size_bytes = 0 + + def plan( +@@ -2386,6 +2417,7 @@ class BatchSparseMLADecodePagedTSWrapper: + kernel_signature=_mla_kernel_compile_signature(static_kernel), + ) + ) ++ self._prepared.clear() + self._state = dict( + device=device, + batch=batch_size, +@@ -2608,6 +2640,38 @@ class BatchSparseMLADecodePagedTSWrapper: + state = self._state + if state is None: + raise RuntimeError("plan() must be called before run()") ++ lse_arg = lse ++ prepared_key = ( ++ kv_cache.data_ptr(), ++ kv_cache.shape, ++ kv_cache.stride(), ++ kv_cache.dtype, ++ None ++ if extra_kv_cache is None ++ else ( ++ extra_kv_cache.data_ptr(), ++ extra_kv_cache.shape, ++ extra_kv_cache.stride(), ++ extra_kv_cache.dtype, ++ ), ++ None if sinks is None else sinks.data_ptr(), ++ softmax_scale, ++ output_scale, ++ ) ++ if ( ++ not validate ++ and out is not None ++ and lse is None ++ and qo_indptr is None ++ and token_positions is None ++ and cos_sin_cache is None ++ and out_scale is None ++ ): ++ prepared = self._prepared.get(prepared_key) ++ if prepared is not None: ++ return self._run_prepared( ++ prepared, query, indices, lengths, extra_indices, extra_lengths, out ++ ) + _validate_query( + query, + packed_query=state["packed"], +@@ -2781,10 +2845,124 @@ class BatchSparseMLADecodePagedTSWrapper: + if static_scales + else state["compiled"], + ) ++ if ( ++ rows ++ and quant is None ++ and static_scales ++ and not state["packed"] ++ and not state["return_lse"] ++ and state["max_q"] == 1 ++ and lse_arg is None ++ ): ++ # Static for this (plan, pools, scales): pool views, scalars, sinks, split, ++ # workspace. The next validate=False launch with the same pools re-forms ++ # only the query / index / output views (see _run_prepared). ++ if len(self._prepared) >= _PREPARED_LAUNCH_LIMIT: ++ self._prepared.clear() ++ self._prepared[prepared_key] = dict( ++ primary=primary, ++ extra=extra, ++ primary_rows=primary.shape[0], ++ c_latent=primary.permute(1, 2, 0), ++ c_rope=extra.permute(1, 2, 0), ++ scalars=(scale_tensors[0], scale_tensors[1], sinks), ++ bmm=(bmm1_scale, bmm2_scale), ++ ks=state["ks"], ++ kc=state["kc"], ++ dummy_indices=state["dummy_indices"], ++ default_cl=state["default_cl"], ++ ) + if quant is not None: + return (out, out_scale, lse) if state["return_lse"] else (out, out_scale) + return (out, lse) if state["return_lse"] else out + ++ def _run_prepared(self, prepared, query, indices, lengths, extra_indices, extra_lengths, out): ++ """Launch with the static arguments recorded by the last full run(). ++ ++ Only the per-call tensors are checked and viewed: query and output ++ geometry, index/length shapes, dtypes and strides. Pool geometry, ++ scales and sinks were validated by the full path that prepared this ++ launch, and the caller passed validate=False. Byte-identical to the ++ full path for the same arguments. ++ """ ++ state = self._state ++ rows, heads, device = state["batch"], state["heads"], state["device"] ++ if ( ++ query.shape != (rows, 1, heads, 512) ++ or query.dtype != state["dtype"] ++ or query.device != device ++ or not query.is_contiguous() ++ or query.data_ptr() % 16 ++ ): ++ raise ValueError( ++ "prepared launch: query must be a contiguous, 16-byte aligned " ++ "[batch, 1, heads, 512] tensor of the planned dtype on the planned device" ++ ) ++ if ( ++ out.numel() != rows * heads * 512 ++ or out.dtype != torch.bfloat16 ++ or out.device != device ++ or not out.is_contiguous() ++ or out.data_ptr() % 16 ++ ): ++ raise ValueError( ++ "prepared launch: out must be a contiguous, 16-byte aligned BF16 tensor " ++ "holding one 512-wide row per query head" ++ ) ++ if ( ++ indices.shape != (rows, prepared["ks"]) ++ or indices.dtype != torch.int32 ++ or indices.device != device ++ or indices.stride(-1) != 1 ++ or lengths.shape != (rows,) ++ or lengths.dtype != torch.int32 ++ or lengths.device != device ++ or not lengths.is_contiguous() ++ ): ++ raise ValueError("prepared launch: invalid primary indices/lengths") ++ if prepared["kc"]: ++ if ( ++ not isinstance(extra_indices, torch.Tensor) ++ or not isinstance(extra_lengths, torch.Tensor) ++ or extra_indices.shape != (rows, prepared["kc"]) ++ or extra_indices.dtype != torch.int32 ++ or extra_indices.device != device ++ or extra_indices.stride(-1) != 1 ++ or extra_lengths.shape != (rows,) ++ or extra_lengths.dtype != torch.int32 ++ or extra_lengths.device != device ++ or not extra_lengths.is_contiguous() ++ ): ++ raise ValueError("prepared launch: invalid extra indices/lengths") ++ ci, cl = extra_indices, extra_lengths ++ else: ++ if extra_indices is not None or extra_lengths is not None: ++ raise ValueError("prepared launch: extra indices without a planned extra pool") ++ ci, cl = prepared["dummy_indices"][:rows], prepared["default_cl"][:rows] ++ buffers = state["buffers"] ++ s0, s1, sinks = prepared["scalars"] ++ q_kernel = query.permute(2, 3, 1, 0) ++ lse_kernel = buffers["public_lse"][:rows].view(rows, 1, heads).permute(2, 1, 0) ++ bmm1, bmm2 = prepared["bmm"] ++ state["compiled_static"]( ++ q_kernel, ++ q_kernel, ++ prepared["c_latent"], ++ prepared["c_rope"], ++ indices.transpose(0, 1), ++ out.view(rows, 1, heads, 512).permute(2, 3, 1, 0), ++ lse_kernel, ++ buffers["core"] if buffers["core"].numel() else None, ++ state["splits"], ++ lengths, ++ None, ++ (indices, ci, lengths, cl, s0, s1, sinks), ++ bmm1, ++ bmm2, ++ None, ++ ) ++ return out ++ + + def get_prims_ts_sparse_mla_decode_workspace_size(*plan_args, **plan_kwargs): + """Return workspace bytes using the same arguments as wrapper.plan(). diff --git a/q_token_kv_block_sparse_metadata.py b/q_token_kv_block_sparse_metadata.py index e2bbef13c0f552030875e6b433e0fcdb20686cf6..dff16046cf7cf9ff03950043cf462ac28de6c670 100644 --- a/q_token_kv_block_sparse_metadata.py diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index eba5560b9d96..e19632edea75 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -21,8 +21,8 @@ vendors: include: - '**/*.py' patch: 3rdparty/vendor_patches/flashinfer-prims-ts.patch - patch_digest: sha256:ef6f887843a37314d86dbadc15db9cc1a53fb7e09fe5a7a0d4d9415cc47432ff - digest: sha256-tree-v1:ce5b8d57a8752174cc1a7d2330fdb5b6aa48463db84c7d8242f4dcec9b0e08ba + patch_digest: sha256:fc5c23ec3d65c562025a1845f211fdac66bfcd3da156dbb8385f4c4407d20d07 + digest: sha256-tree-v1:f2f444f7e6d4ad688fbf4e1e9eed11ff7c78d835cfd0af5ef9d7660780b059e0 openengine-proto: url: https://github.com/ai-dynamo/openengine.git tag: v0.1.0 diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py index 49c7ca16534f..dcb10ec3d3e5 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py @@ -92,6 +92,8 @@ _WARMED_PROFILES: set = set() _VALIDATED_CONTRACTS: set = set() _LATENT_DIM = 512 +# Prepared-sinks cache entries (one per layer; DeepSeek-V4 has 61). +_SINKS_CACHE_LIMIT = 512 _ROPE_DIM = 64 _FUSED_OUTPUT_BLOCK = 128 @@ -160,9 +162,11 @@ def __init__(self, attn: "TrtllmAttention") -> None: # manager they were built from. The owner is held weakly: a freed # manager (or a new one at the same address) never resolves to a # stale view of released pool memory. + self._checked_key = None self._pool_views_owner: Optional[weakref.ReferenceType] = None self._pool_views: dict[tuple[int, int], torch.Tensor] = {} - self._sinks_cache: Optional[tuple[int, torch.Tensor]] = None + # One prepared FP32 sinks vector per source tensor (one per layer). + self._sinks_cache: dict[tuple[int, int, torch.dtype], torch.Tensor] = {} # ------------------------------------------------------------------ availability @classmethod @@ -401,19 +405,27 @@ def _sinks(self, fwd: AttentionForwardArgs) -> Optional[torch.Tensor]: sinks = fwd.attention_sinks if sinks is None: return None - cached = self._sinks_cache if ( - cached is not None - and cached[0] == sinks.data_ptr() - and cached[1].numel() == sinks.numel() + sinks.dtype == torch.float32 + and sinks.dim() == 1 + and sinks.is_contiguous() + and not sinks.requires_grad ): - return cached[1] + # DeepseekV4TrtllmAttention.attn_sink is already a per-head FP32 + # parameter; pass it through so its address stays stable per layer. + return sinks + key = (sinks.data_ptr(), sinks.numel(), sinks.dtype) + prepared = self._sinks_cache.get(key) + if prepared is not None: + return prepared if torch.cuda.is_current_stream_capturing(): raise RuntimeError( "PrimTS sparse MLA attention sinks must be prepared before CUDA graph capture." ) prepared = sinks.detach().to(torch.float32).reshape(-1).contiguous() - self._sinks_cache = (sinks.data_ptr(), prepared) + if len(self._sinks_cache) >= _SINKS_CACHE_LIMIT: + self._sinks_cache.clear() + self._sinks_cache[key] = prepared return prepared @staticmethod @@ -694,6 +706,49 @@ def _claim_contract_validation(contract_key: tuple) -> bool: return True # ------------------------------------------------------------------ forward + def _check_request( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + fwd: AttentionForwardArgs, + ) -> None: + """Run the full request check once per request geometry. + + The manager selected this library through is_supported(); the re-check + here guards a changed request, so it runs only when the per-step + geometry (metadata object, query, indices, output, fusion mode, token + counts) differs from the last accepted call. + """ + indices = fwd.sparse_runtime_params.sparse_attn_indices + output = fwd.output + check_key = ( + id(metadata), + int(metadata.num_tokens), + int(metadata.num_ctx_tokens), + q.shape, + q.dtype, + k is None, + v is None, + None + if not isinstance(indices, torch.Tensor) + else (indices.shape, indices.dtype, indices.stride(-1)), + None if output is None else (output.shape, output.dtype, output.is_contiguous()), + fwd.enable_dsv4_epilogue_fusion, + fwd.quant_q_buffer is None, + fwd.attention_sinks is None, + fwd.output_sf is None, + fwd.out_scale is None, + fwd.dsv4_inv_rope_cos_sin_cache is None, + ) + if check_key == self._checked_key: + return + reason = self._unsupported_reason(q, k, v, metadata, fwd) + if reason is not None: + raise RuntimeError(f"PrimTS sparse MLA FMHA cannot run this request: {reason}") + self._checked_key = check_key + def forward( self, q: torch.Tensor, @@ -706,9 +761,9 @@ def forward( DeepseekV4AttentionType, ) - reason = self._unsupported_reason(q, k, v, metadata, forward_args) - if reason is not None: - raise RuntimeError(f"PrimTS sparse MLA FMHA cannot run this request: {reason}") + self._check_request(q, k, v, metadata, forward_args) + indices = forward_args.sparse_runtime_params.sparse_attn_indices + output = forward_args.output attn = self.attn manager = metadata.kv_cache_manager heads = int(attn.num_heads) @@ -720,7 +775,6 @@ def forward( extra = self._extra_capacity(metadata) is_fp8 = self._is_fp8_cache(metadata) fused = bool(forward_args.enable_dsv4_epilogue_fusion) - output = forward_args.output groups = int(output.shape[0]) if fused else None if is_fp8: @@ -762,7 +816,6 @@ def forward( # come from the metadata too, derived once per step: the token's # position within the window, and the compressed rows selected beyond # the window. - indices = forward_args.sparse_runtime_params.sparse_attn_indices positions = metadata.token_positions_cuda[start:end] lengths = metadata.sparse_mla_swa_lens[start:end] swa_pool = self._pool_view(manager, DeepseekV4AttentionType.SWA, 1) 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 d9cdcd13a270..7a305b1ec6f2 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py @@ -2121,6 +2121,7 @@ def _prepare_sparse_output_quant( ) +@functools.lru_cache(maxsize=1024) def _resolve_sparse_mla_plan( device, batch_size, @@ -2286,6 +2287,9 @@ def _resolve_sparse_mla_plan( return device, profile, spec, sections, byte_end, capacity, output_quant +# Prepared launches kept per plan (one per layer sharing the plan; DeepSeek-V4 has 61). +_PREPARED_LAUNCH_LIMIT = 256 + class BatchSparseMLADecodePagedTSWrapper: """Plan native BF16/E4M3 attention, then bind one or two source index lists. @@ -2319,6 +2323,11 @@ def __init__(self, workspace_buffer=None): self._workspace_buffer = workspace_buffer self._state = None self._scalar_cache = {} + # Prepared launches for the unfused validate=False path: everything that + # is static for one (plan, pools, scales, sinks) tuple, keyed by the pools' + # storage and the sinks address. One plan serves every layer of a model, + # and each layer brings its own sinks, so one entry per layer. + self._prepared: dict[tuple, dict] = {} self.workspace_size_bytes = 0 def plan( @@ -2408,6 +2417,7 @@ def plan( kernel_signature=_mla_kernel_compile_signature(static_kernel), ) ) + self._prepared.clear() self._state = dict( device=device, batch=batch_size, @@ -2630,6 +2640,38 @@ def run( state = self._state if state is None: raise RuntimeError("plan() must be called before run()") + lse_arg = lse + prepared_key = ( + kv_cache.data_ptr(), + kv_cache.shape, + kv_cache.stride(), + kv_cache.dtype, + None + if extra_kv_cache is None + else ( + extra_kv_cache.data_ptr(), + extra_kv_cache.shape, + extra_kv_cache.stride(), + extra_kv_cache.dtype, + ), + None if sinks is None else sinks.data_ptr(), + softmax_scale, + output_scale, + ) + if ( + not validate + and out is not None + and lse is None + and qo_indptr is None + and token_positions is None + and cos_sin_cache is None + and out_scale is None + ): + prepared = self._prepared.get(prepared_key) + if prepared is not None: + return self._run_prepared( + prepared, query, indices, lengths, extra_indices, extra_lengths, out + ) _validate_query( query, packed_query=state["packed"], @@ -2803,10 +2845,124 @@ def run( if static_scales else state["compiled"], ) + if ( + rows + and quant is None + and static_scales + and not state["packed"] + and not state["return_lse"] + and state["max_q"] == 1 + and lse_arg is None + ): + # Static for this (plan, pools, scales): pool views, scalars, sinks, split, + # workspace. The next validate=False launch with the same pools re-forms + # only the query / index / output views (see _run_prepared). + if len(self._prepared) >= _PREPARED_LAUNCH_LIMIT: + self._prepared.clear() + self._prepared[prepared_key] = dict( + primary=primary, + extra=extra, + primary_rows=primary.shape[0], + c_latent=primary.permute(1, 2, 0), + c_rope=extra.permute(1, 2, 0), + scalars=(scale_tensors[0], scale_tensors[1], sinks), + bmm=(bmm1_scale, bmm2_scale), + ks=state["ks"], + kc=state["kc"], + dummy_indices=state["dummy_indices"], + default_cl=state["default_cl"], + ) if quant is not None: return (out, out_scale, lse) if state["return_lse"] else (out, out_scale) return (out, lse) if state["return_lse"] else out + def _run_prepared(self, prepared, query, indices, lengths, extra_indices, extra_lengths, out): + """Launch with the static arguments recorded by the last full run(). + + Only the per-call tensors are checked and viewed: query and output + geometry, index/length shapes, dtypes and strides. Pool geometry, + scales and sinks were validated by the full path that prepared this + launch, and the caller passed validate=False. Byte-identical to the + full path for the same arguments. + """ + state = self._state + rows, heads, device = state["batch"], state["heads"], state["device"] + if ( + query.shape != (rows, 1, heads, 512) + or query.dtype != state["dtype"] + or query.device != device + or not query.is_contiguous() + or query.data_ptr() % 16 + ): + raise ValueError( + "prepared launch: query must be a contiguous, 16-byte aligned " + "[batch, 1, heads, 512] tensor of the planned dtype on the planned device" + ) + if ( + out.numel() != rows * heads * 512 + or out.dtype != torch.bfloat16 + or out.device != device + or not out.is_contiguous() + or out.data_ptr() % 16 + ): + raise ValueError( + "prepared launch: out must be a contiguous, 16-byte aligned BF16 tensor " + "holding one 512-wide row per query head" + ) + if ( + indices.shape != (rows, prepared["ks"]) + or indices.dtype != torch.int32 + or indices.device != device + or indices.stride(-1) != 1 + or lengths.shape != (rows,) + or lengths.dtype != torch.int32 + or lengths.device != device + or not lengths.is_contiguous() + ): + raise ValueError("prepared launch: invalid primary indices/lengths") + if prepared["kc"]: + if ( + not isinstance(extra_indices, torch.Tensor) + or not isinstance(extra_lengths, torch.Tensor) + or extra_indices.shape != (rows, prepared["kc"]) + or extra_indices.dtype != torch.int32 + or extra_indices.device != device + or extra_indices.stride(-1) != 1 + or extra_lengths.shape != (rows,) + or extra_lengths.dtype != torch.int32 + or extra_lengths.device != device + or not extra_lengths.is_contiguous() + ): + raise ValueError("prepared launch: invalid extra indices/lengths") + ci, cl = extra_indices, extra_lengths + else: + if extra_indices is not None or extra_lengths is not None: + raise ValueError("prepared launch: extra indices without a planned extra pool") + ci, cl = prepared["dummy_indices"][:rows], prepared["default_cl"][:rows] + buffers = state["buffers"] + s0, s1, sinks = prepared["scalars"] + q_kernel = query.permute(2, 3, 1, 0) + lse_kernel = buffers["public_lse"][:rows].view(rows, 1, heads).permute(2, 1, 0) + bmm1, bmm2 = prepared["bmm"] + state["compiled_static"]( + q_kernel, + q_kernel, + prepared["c_latent"], + prepared["c_rope"], + indices.transpose(0, 1), + out.view(rows, 1, heads, 512).permute(2, 3, 1, 0), + lse_kernel, + buffers["core"] if buffers["core"].numel() else None, + state["splits"], + lengths, + None, + (indices, ci, lengths, cl, s0, s1, sinks), + bmm1, + bmm2, + None, + ) + return out + def get_prims_ts_sparse_mla_decode_workspace_size(*plan_args, **plan_kwargs): """Return workspace bytes using the same arguments as wrapper.plan(). 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 43d098bb45a5..3d15a9521b59 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 @@ -386,6 +386,218 @@ def launch(validate: bool) -> None: assert not torch.equal(replayed, torch.zeros_like(replayed)) +def test_prepared_launch_matches_the_full_path_per_layer(): + """A validate=False run with the pools and sinks of an earlier full run takes + the prepared launch, byte-identical to the full path, with one prepared + entry per (pools, sinks) key so layers sharing the plan do not evict each + other; a new pool or new sinks address goes back through the full path. + """ + _skip_unless_toolchain() + from tensorrt_llm._torch.attention.backends.prims_ts import mla_decode + + torch.manual_seed(3) + device = torch.device("cuda") + rows, heads, topk, extra_topk = 6, 64, 128, 256 + pool, extra_pool, query, sets = _two_source_case( + rows, heads, topk, extra_topk, device, torch.bfloat16 + ) + indices, lengths, extra_indices, extra_lengths = sets[0] + sinks_a = torch.randn(heads, dtype=torch.float32, device=device) + sinks_b = torch.randn(heads, dtype=torch.float32, device=device) + wrapper = mla_decode.BatchSparseMLADecodePagedTSWrapper() + wrapper.plan(device, rows, heads, max_topk=topk, max_extra_topk=extra_topk, has_sinks=True) + prepared_launches = [] + original = wrapper._run_prepared + + def counting_run_prepared(*args, **kwargs): + prepared_launches.append(1) + return original(*args, **kwargs) + + wrapper._run_prepared = counting_run_prepared + + def launch(sinks, kv=pool, validate=False, idx=0): + out = torch.empty(rows, 1, heads, LATENT, dtype=torch.bfloat16, device=device) + wrapper.run( + query, + kv, + sets[idx][0], + sets[idx][1], + extra_pool, + extra_indices=sets[idx][2], + extra_lengths=sets[idx][3], + softmax_scale=LATENT**-0.5, + output_scale=1.0, + sinks=sinks, + out=out, + validate=validate, + ) + torch.cuda.synchronize() + return out + + ref_a = launch(sinks_a, validate=True) # full path records layer A + ref_b = launch(sinks_b, validate=True) # full path records layer B + assert not prepared_launches and len(wrapper._prepared) == 2 + got_a = launch(sinks_a) + got_b = launch(sinks_b) + assert len(prepared_launches) == 2, "both layers must hit their prepared launch" + torch.testing.assert_close(got_a, ref_a, rtol=0, atol=0) + torch.testing.assert_close(got_b, ref_b, rtol=0, atol=0) + assert not torch.equal(got_a, got_b) + # The prepared launch reads the live index set, not the recorded one. + got_a1 = launch(sinks_a, idx=1) + assert len(prepared_launches) == 3 + torch.testing.assert_close(got_a1, launch(sinks_a, validate=True, idx=1), rtol=0, atol=0) + # A different pool storage or sinks address is not prepared: full path. + other_pool = pool.clone() + launch(sinks_a, kv=other_pool) + launch(sinks_a.clone()) + assert len(prepared_launches) == 3 and len(wrapper._prepared) == 4 + # Re-planning drops every prepared launch. + wrapper.plan(device, rows, heads, max_topk=topk, max_extra_topk=extra_topk, has_sinks=True) + assert not wrapper._prepared + + +def test_request_check_runs_once_per_geometry(monkeypatch) -> None: + """The full request re-check runs on the first call and again only when the + request geometry changes; repeated steps of one shape skip it.""" + adapter = PrimsTSSparseMlaFmha.__new__(PrimsTSSparseMlaFmha) + adapter._checked_key = None + reasons = [] + + def fake_reason(self, q, k, v, metadata, fwd): + reasons.append(1) + return None + + monkeypatch.setattr(PrimsTSSparseMlaFmha, "_unsupported_reason", fake_reason) + + def make(rows: int, width: int = 640, fusion: bool = False, ctx: int = 0): + metadata = SimpleNamespace(num_tokens=ctx + rows, num_ctx_tokens=ctx) + indices = torch.zeros(rows, width, dtype=torch.int32) + output = torch.empty(rows, 64 * LATENT, dtype=torch.bfloat16) + fwd = SimpleNamespace( + sparse_runtime_params=SimpleNamespace(sparse_attn_indices=indices), + output=output, + enable_dsv4_epilogue_fusion=fusion, + quant_q_buffer=None, + attention_sinks=None, + output_sf=None, + out_scale=None, + dsv4_inv_rope_cos_sin_cache=None, + ) + return torch.empty(rows, 64 * LATENT, dtype=torch.bfloat16), metadata, fwd + + q, metadata, fwd = make(8) + adapter._check_request(q, None, None, metadata, fwd) + adapter._check_request(q, None, None, metadata, fwd) + adapter._check_request(q, None, None, metadata, fwd) + assert len(reasons) == 1 + q2, _, fwd2 = make(16) # new row count on the same metadata object + adapter._check_request(q2, None, None, metadata, fwd2) + assert len(reasons) == 2 + adapter._check_request(q2, None, None, metadata, fwd2) + assert len(reasons) == 2 + _, _, fwd3 = make(16, fusion=True) + adapter._check_request(q2, None, None, metadata, fwd3) + assert len(reasons) == 3 + adapter._check_request(q2, None, None, metadata, fwd3) + assert len(reasons) == 3 + _, metadata4, fwd4 = make(16, fusion=True) # only the metadata object differs + adapter._check_request(q2, None, None, metadata4, fwd4) + assert len(reasons) == 4 + + def rejecting_reason(self, q, k, v, metadata, fwd): + return "nope" + + monkeypatch.setattr(PrimsTSSparseMlaFmha, "_unsupported_reason", rejecting_reason) + adapter._check_request(q2, None, None, metadata4, fwd4) # same geometry: skipped + q5, metadata5, fwd5 = make(4) + with pytest.raises(RuntimeError, match="nope"): + adapter._check_request(q5, None, None, metadata5, fwd5) + + +def test_prepared_launch_single_pool_matches_the_full_path(): + """SWA-only layers plan no extra pool; their prepared launch must match the + full path and reject extra indices.""" + _skip_unless_toolchain() + from tensorrt_llm._torch.attention.backends.prims_ts import mla_decode + + torch.manual_seed(4) + device = torch.device("cuda") + rows, heads, topk = 6, 64, 128 + pool, _, query, sets = _two_source_case(rows, heads, topk, 256, device, torch.bfloat16) + wrapper = mla_decode.BatchSparseMLADecodePagedTSWrapper() + wrapper.plan(device, rows, heads, max_topk=topk) + prepared_launches = [] + original = wrapper._run_prepared + + def counting_run_prepared(*args, **kwargs): + prepared_launches.append(1) + return original(*args, **kwargs) + + wrapper._run_prepared = counting_run_prepared + + def launch(validate: bool, idx: int = 0) -> torch.Tensor: + out = torch.empty(rows, 1, heads, LATENT, dtype=torch.bfloat16, device=device) + wrapper.run( + query, + pool, + sets[idx][0], + sets[idx][1], + softmax_scale=LATENT**-0.5, + output_scale=1.0, + out=out, + validate=validate, + ) + torch.cuda.synchronize() + return out + + ref = launch(validate=True) + assert not prepared_launches and len(wrapper._prepared) == 1 + got = launch(validate=False) + assert len(prepared_launches) == 1 + torch.testing.assert_close(got, ref, rtol=0, atol=0) + ref1 = launch(validate=True, idx=1) + got1 = launch(validate=False, idx=1) + assert len(prepared_launches) == 2 + torch.testing.assert_close(got1, ref1, rtol=0, atol=0) + with pytest.raises(ValueError, match="extra indices"): + wrapper.run( + query, + pool, + sets[0][0], + sets[0][1], + extra_indices=sets[0][2], + extra_lengths=sets[0][3], + softmax_scale=LATENT**-0.5, + output_scale=1.0, + out=torch.empty(rows, 1, heads, LATENT, dtype=torch.bfloat16, device=device), + validate=False, + ) + + +def test_sinks_pass_through_or_are_prepared_once_per_layer() -> None: + """FP32 per-head sinks are passed through at a stable address; other + dtypes are converted once per source tensor, so the kernel key keyed by the + sinks address stays stable for every layer served by one adapter.""" + from tensorrt_llm._torch.attention.backends.fmha import prims_ts_sparse_mla as mod + + adapter = PrimsTSSparseMlaFmha.__new__(PrimsTSSparseMlaFmha) + adapter._sinks_cache = {} + fp32 = torch.nn.Parameter(torch.randn(64, device="cuda"), requires_grad=False) + assert adapter._sinks(SimpleNamespace(attention_sinks=fp32)) is fp32 + assert not adapter._sinks_cache + bf16_a = torch.randn(64, device="cuda").to(torch.bfloat16) + bf16_b = torch.randn(64, device="cuda").to(torch.bfloat16) + prep_a = adapter._sinks(SimpleNamespace(attention_sinks=bf16_a)) + prep_b = adapter._sinks(SimpleNamespace(attention_sinks=bf16_b)) + assert prep_a.dtype == torch.float32 and prep_a is not prep_b + assert adapter._sinks(SimpleNamespace(attention_sinks=bf16_a)) is prep_a + assert adapter._sinks(SimpleNamespace(attention_sinks=bf16_b)) is prep_b + torch.testing.assert_close(prep_a, bf16_a.float()) + assert adapter._sinks(SimpleNamespace(attention_sinks=None)) is None + assert mod._SINKS_CACHE_LIMIT >= 61 + + @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float8_e4m3fn], ids=["bf16", "fp8"]) @pytest.mark.parametrize( "rows,heads,groups,topk,extra_topk", From c67aa5f37cf9c7f3618401660f1ac7493381add1 Mon Sep 17 00:00:00 2001 From: yunruis <205571022+yunruis@users.noreply.github.com> Date: Fri, 23 Oct 2026 19:56:51 -0700 Subject: [PATCH 19/19] [None][fix] PrimTS sparse MLA: cap warm-up workers per host, cover the opt-out and MTP rows, add release notes Under the default launcher (`trtllm-llmapi-launch` and the `LLM()` MPI session) every rank owns a private FlashInfer workspace, so the ranks of one host do not share kernel compiles and each fans its warm-up out on its own. The per-rank default of 4 precompile subprocesses therefore started 32 `cute.compile` interpreters on a TP=8 node, each with a CUDA context. The default now follows the ranks on the host: at most 4 per rank and 16 per host together (TP=8 runs 2 per rank); an explicit `TLLM_PRIMTS_WARM_WORKERS` is never capped. The module note documents the host load. The warm-up reads `max_num_requests`, `max_total_draft_tokens` and `max_num_tokens` directly instead of through `getattr(..., 0)`, so a renamed field fails at warm-up instead of silently shrinking the warmed set and moving the MTP row counts' compiles to the first serving step. Tests: the unified DeepSeek-V4 forward test now also runs with `TLLM_FMHA_LIBS=-prims_ts_sparse_mla` and asserts that the TRTLLM-Gen fallback serves generation with PrimTS never entered (the documented opt-out, and what SM103 / fp8_ds_mla / NVFP4 deployments run); the module test's multi-token generation cases (next_n 2 and 6, ragged 1/2/4) run with a spy asserting PrimTS served them with one row per query token; a unit test asserts the warm set covers `requests * (1 + draft)` rows and is bounded by the token budget; the worker-count tests cover the per-host cap. Release notes: the default flip, the opt-out, the first-launch compile cost under per-rank workspaces, the `TLLM_PRIMTS_WARM_*` variables and the `FLASHINFER_DISABLE_JIT=1` behavior. Addresses the review by @brnguyen2 on #20209. Co-Authored-By: Claude Fable 5.1 Signed-off-by: yunruis <205571022+yunruis@users.noreply.github.com> --- docs/source/release-notes.md | 5 + .../backends/fmha/prims_ts_precompile.py | 77 ++++++- .../backends/fmha/prims_ts_sparse_mla.py | 28 ++- .../test_prims_ts_sparse_mla_fmha.py | 201 ++++++++++++++++-- 4 files changed, 278 insertions(+), 33 deletions(-) diff --git a/docs/source/release-notes.md b/docs/source/release-notes.md index 8f949a1aeb75..4aa228406737 100644 --- a/docs/source/release-notes.md +++ b/docs/source/release-notes.md @@ -8,6 +8,11 @@ All published functionality in the Release Notes has been fully tested and verif ### Key Features and Enhancements +- **DeepSeek-V4 generation on B200/GB200 (SM100) now runs on the PrimTS native sparse MLA FMHA by default.** The new `prims_ts_sparse_mla` FMHA library (a CuTe DSL kernel vendored from FlashInfer) serves the DeepSeek-V4 sparse MLA generation phase ahead of the TRTLLM-Gen path on SM100; it covers BF16 and FP8 KV caches, the fused inverse-RoPE + FP8 output epilogue, speculative decoding (MTP) rows and CUDA graph replay. SM103 (GB300), `fp8_ds_mla` and NVFP4 KV caches keep the TRTLLM-Gen path. Opt out with `TLLM_FMHA_LIBS=-prims_ts_sparse_mla`. + - **First-launch compile cost.** The kernels are JIT-compiled by the CuTe DSL and cached on disk under the FlashInfer workspace. The first start on a host (or after a FlashInfer / CuTe DSL upgrade) compiles the kernel variants the deployment's generation row counts select during `LLM()` startup, about 5 to 10 minutes on B200 for DeepSeek-V4; later starts load the cache. Under `trtllm-llmapi-launch` and the default `LLM()` MPI session every rank has its own FlashInfer workspace (`TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS=1`), so each rank compiles its own copy; set `TRTLLM_FLASHINFER_WORKSPACE_PER_PROCESS=0` to share one cache between the ranks of a host. + - **Warm-up controls.** `TLLM_PRIMTS_WARM_WORKERS` sets the compile subprocesses per rank (default 4, reduced so that the ranks on a host, counted from the model's parallel mapping, start at most 16 together; `0` compiles in-process), `TLLM_PRIMTS_WARM_TIMEOUT` bounds the wait for them (default 1800 s), and `TLLM_PRIMTS_WARM_STRICT=1` turns a worker failure into an error instead of an in-process retry. Each worker imports torch and the CuTe DSL and opens a CUDA context, so lower the worker count on nodes with little host memory. + - **`FLASHINFER_DISABLE_JIT=1`** now also refuses to compile these kernels (with the on-disk CuTe DSL cache enabled, its default): startup then fails on a variant missing from the cache instead of falling back to TRTLLM-Gen, so a deployment that sets it must pre-populate the cache (one start without the variable) or opt out of the library. + ### Infrastructure Changes ### API Changes diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_precompile.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_precompile.py index 209ae7681b37..f293388f6a1f 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_precompile.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_precompile.py @@ -21,8 +21,22 @@ rank (TRT-LLM's default) or shared between ranks. Failures are best-effort: the caller falls back to compiling in-process. +Host load. Every worker is an interpreter that imports torch and the CuTe DSL, +opens a CUDA context on its rank's GPU and runs ``cute.compile`` (one to a few +GB of host memory and one busy core each). Under the default launcher each +rank owns a private FlashInfer workspace, so the ranks of one host do not +share compiles and every rank fans out on its own. The default worker count +therefore follows the number of ranks on the host (from the model's +``Mapping``, or the launcher's environment): at most 4 per rank and at most +``_HOST_COMPILE_BUDGET`` (16) workers per host in total, so a TP=8 node runs 2 +workers per rank, not 32 subprocesses at once. Set +``TLLM_PRIMTS_WARM_WORKERS`` explicitly to trade startup time against host +memory and CPU on a given node. + Environment: - TLLM_PRIMTS_WARM_WORKERS subprocess count, default 4; 0 or 1 disables. + TLLM_PRIMTS_WARM_WORKERS subprocess count per rank; default 4 capped so + that all ranks on the host start at most 16; + 0 or 1 disables. TLLM_PRIMTS_WARM_TIMEOUT seconds to wait for the workers, default 1800. TLLM_PRIMTS_WARM_STRICT 1 turns a worker failure into an error (tests). """ @@ -31,7 +45,7 @@ import json import os -import subprocess +import subprocess # nosec B404 - runs sys.executable workers with a fixed argv, no shell import sys import tempfile import time @@ -49,9 +63,46 @@ """ -def worker_count(num_payloads: int) -> int: - """Subprocess count from ``TLLM_PRIMTS_WARM_WORKERS`` (default 4); ``0``/``1`` disables. +_MAX_DEFAULT_WORKERS = 4 +_HOST_COMPILE_BUDGET = 16 # default workers of all ranks on one host, together +_LOCAL_RANK_ENV = ("OMPI_COMM_WORLD_LOCAL_SIZE", "RAY_LOCAL_WORLD_SIZE", "SLURM_NTASKS_PER_NODE") + + +def local_rank_count(mapping) -> int: + """Ranks of this job on this host, without any MPI call. + + From the model's ``Mapping`` (``min(world_size, gpus_per_node)``) when the + caller has one, else from the launcher's environment, else 1. No + communicator query: this runs inside the first forward, where a collective + would hang if the ranks reached it at different points. + """ + if mapping is not None: + world = int(getattr(mapping, "world_size", 0) or 0) + per_node = int(getattr(mapping, "gpus_per_node", 0) or 0) + if world > 0 and per_node > 0: + return max(1, min(world, per_node)) + if world > 0: + return world + for name in _LOCAL_RANK_ENV: + raw = os.environ.get(name) + if raw: + try: + return max(1, int(raw)) + except ValueError: + continue + return 1 + + +def default_worker_count(local_ranks: int = 1) -> int: + """Per-rank default: 4, lowered so the ranks on the host start at most 16 workers.""" + return max(1, min(_MAX_DEFAULT_WORKERS, _HOST_COMPILE_BUDGET // max(1, int(local_ranks)))) + + +def worker_count(num_payloads: int, *, local_ranks: int = 1) -> int: + """Subprocess count from ``TLLM_PRIMTS_WARM_WORKERS``; ``0``/``1`` disables. + Without the variable the default comes from ``default_worker_count()`` + for ``local_ranks`` ranks on this host (see the module note on host load). Zero when the on-disk cache or JIT compilation is disabled: a worker could then compile but keep nothing. """ @@ -59,12 +110,18 @@ def worker_count(num_payloads: int) -> int: "FLASHINFER_DISABLE_JIT" ): return 0 - raw = os.environ.get("TLLM_PRIMTS_WARM_WORKERS", "4") - try: - workers = int(raw) - except ValueError: - logger.warning(f"Ignoring TLLM_PRIMTS_WARM_WORKERS={raw!r}; expected an integer.") - workers = 4 + raw = os.environ.get("TLLM_PRIMTS_WARM_WORKERS") + if raw is None: + workers = default_worker_count(local_ranks) + else: + try: + workers = int(raw) + except ValueError: + workers = default_worker_count(local_ranks) + logger.warning( + f"Ignoring TLLM_PRIMTS_WARM_WORKERS={raw!r}; expected an integer " + f"(using the default of {workers})." + ) return max(0, min(workers, num_payloads)) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py index dcb10ec3d3e5..2cb531e21c3d 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py @@ -596,12 +596,17 @@ def _warm_kernel_variants( _WARMED_PROFILES.add(profile_key) from tensorrt_llm._torch.attention.backends.prims_ts import mla_decode + from . import prims_ts_precompile + plan_kwargs, _ = self._plan_kwargs(extra, q_dtype, has_sinks, fused_groups) # Generation rows per step: one per request, plus its draft tokens with - # speculative decoding, within the step's token budget. - requests = int(getattr(metadata, "max_num_requests", 0) or 0) - draft = int(getattr(metadata, "max_total_draft_tokens", 0) or 0) - tokens = int(getattr(metadata, "max_num_tokens", 0) or 0) + # speculative decoding, within the step's token budget. The fields are + # declared on TrtllmAttentionMetadata; read them directly so a rename + # fails here instead of silently shrinking the warmed set and moving + # the MTP row counts' compiles to the first serving step. + requests = int(metadata.max_num_requests or 0) + draft = int(metadata.max_total_draft_tokens or 0) + tokens = int(metadata.max_num_tokens or 0) max_rows = requests * (1 + draft) if tokens: max_rows = min(max_rows, tokens) if max_rows else tokens @@ -614,7 +619,11 @@ def _warm_kernel_variants( signature = mla_decode._mla_kernel_compile_signature(spec.kernel) representative[signature] = rows # the largest row count of each variant self._precompile_variants( - device, heads, plan_kwargs, sorted(representative.values(), reverse=True) + device, + heads, + plan_kwargs, + sorted(representative.values(), reverse=True), + local_ranks=prims_ts_precompile.local_rank_count(metadata.mapping), ) for rows in sorted(representative.values(), reverse=True): self._get_plan( @@ -633,7 +642,12 @@ def _warm_kernel_variants( @staticmethod def _precompile_variants( - device: torch.device, heads: int, plan_kwargs: dict, rows_list: list + device: torch.device, + heads: int, + plan_kwargs: dict, + rows_list: list, + *, + local_ranks: int = 1, ) -> None: """Populate the on-disk kernel cache for ``rows_list`` with worker subprocesses. @@ -646,7 +660,7 @@ def _precompile_variants( from . import prims_ts_precompile - workers = prims_ts_precompile.worker_count(len(rows_list)) + workers = prims_ts_precompile.worker_count(len(rows_list), local_ranks=local_ranks) if workers <= 1: return try: 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 3d15a9521b59..5beda56866d2 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 @@ -730,6 +730,7 @@ def dequant(values, scales): @pytest.mark.skipif(not HAS_FLASH_MLA, reason="FlashMLA not available") +@pytest.mark.parametrize("library", ["prims_ts", "fallback"]) @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) @pytest.mark.parametrize( "batch_name", @@ -741,24 +742,89 @@ def dequant(values, scales): "mixed_heavy_decode", ], ) -def test_forward_sparse_mla_unified_prims_ts(monkeypatch, batch_name, kv_cache_dtype): - """Run the DeepSeek-V4 module test on the default library set. - - No ``TLLM_FMHA_LIBS`` override: the PrimTS library is in the default set and - must win the generation phase on its own. +def test_forward_sparse_mla_unified_prims_ts(monkeypatch, batch_name, kv_cache_dtype, library): + """Run the DeepSeek-V4 module test on the default library set and on the opt-out. + + ``prims_ts``: no ``TLLM_FMHA_LIBS`` override; the PrimTS library is in the + default set and must win the generation phase on its own. ``fallback``: + ``TLLM_FMHA_LIBS=-prims_ts_sparse_mla``, the documented opt-out, must hand + generation back to the TRTLLM-Gen fallback (``FallbackFmha``) with PrimTS + never entered, so a regression on that path (what SM103, ``fp8_ds_mla`` and + NVFP4 deployments run) still fails here. """ - _skip_unless_toolchain() - monkeypatch.delenv("TLLM_FMHA_LIBS", raising=False) - calls = [] + from tensorrt_llm._torch.attention.backends.fmha.fallback import FallbackFmha + from tensorrt_llm._torch.attention.backends.interface import AttentionInputType + + if library == "prims_ts": + _skip_unless_toolchain() + monkeypatch.delenv("TLLM_FMHA_LIBS", raising=False) + else: + monkeypatch.setenv("TLLM_FMHA_LIBS", f"-{LIB_NAME}") + prims_calls: list = [] + fallback_calls: list = [] original_forward = PrimsTSSparseMlaFmha.forward + original_fallback = FallbackFmha.forward def counting_forward(self, *args, **kwargs): - calls.append(1) + prims_calls.append(1) return original_forward(self, *args, **kwargs) + def counting_fallback(self, q, k, v, metadata, forward_args, *args, **kwargs): + if ( + metadata.num_generations > 0 + and forward_args.attention_input_type != AttentionInputType.context_only + ): + fallback_calls.append(1) + return original_fallback(self, q, k, v, metadata, forward_args, *args, **kwargs) + monkeypatch.setattr(PrimsTSSparseMlaFmha, "forward", counting_forward) + monkeypatch.setattr(FallbackFmha, "forward", counting_fallback) _run_unified_sparse_mla_forward(batch_name, kv_cache_dtype, "deepseek_v4") - assert calls, "the PrimTS sparse MLA FMHA did not serve the generation phase" + if library == "prims_ts": + assert prims_calls, "the PrimTS sparse MLA FMHA did not serve the generation phase" + else: + assert not prims_calls, "TLLM_FMHA_LIBS=-prims_ts_sparse_mla did not disable PrimTS" + assert fallback_calls, "the TRTLLM-Gen fallback did not serve the generation phase" + + +@pytest.mark.parametrize( + "generation_seq_len_q", + [2, 6, [1, 2, 4]], + ids=["mtp_next_n_2", "mtp_next_n_6", "ragged_1_2_4"], +) +def test_speculative_decoding_rows_run_on_prims_ts(monkeypatch, generation_seq_len_q): + """MTP generation (``max_draft_tokens > 0``) on the default library set. + + ``test_deepseek_v4_sparse_mla`` drives the DeepSeek-V4 attention module with + several query tokens per request (uniform ``next_n`` widths and a ragged + batch) and checks the output against its reference; here it is run with the + library default and the spy asserts that PrimTS served those steps with one + row per query token (``rows = sum(generation_lens)``, not one per request), + which is the per-token-row contract the kernel is planned for + (``max_seq_len_q = 1``). + """ + _skip_unless_toolchain() + monkeypatch.delenv("TLLM_FMHA_LIBS", raising=False) + from . import test_deepseek_v4_sparse_mla as module_test + + context_lengths = [14, 508, 3947] + generation_lens = ( + generation_seq_len_q + if isinstance(generation_seq_len_q, list) + else [generation_seq_len_q] * len(context_lengths) + ) + rows_seen: list[int] = [] + original_forward = PrimsTSSparseMlaFmha.forward + + def spy_forward(self, q, k, v, metadata, forward_args): + rows_seen.append(int(q.shape[0])) + return original_forward(self, q, k, v, metadata, forward_args) + + monkeypatch.setattr(PrimsTSSparseMlaFmha, "forward", spy_forward) + module_test.test_deepseek_v4_sparse_mla(context_lengths, generation_seq_len_q) + assert rows_seen, "PrimTS did not serve the speculative-decoding generation steps" + assert set(rows_seen) == {sum(generation_lens)}, rows_seen + assert sum(generation_lens) > len(context_lengths) # more rows than requests: draft tokens def _dequant_fused(values: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: @@ -941,14 +1007,24 @@ class _Attn: device=torch.device("cuda"), q_dtype=torch.bfloat16, has_sinks=False, fused_groups=None ) eager = SimpleNamespace( - fmha_plan_caches={}, max_num_tokens=20, max_num_requests=20, is_cuda_graph=False + fmha_plan_caches={}, + max_num_tokens=20, + max_num_requests=20, + max_total_draft_tokens=None, + mapping=None, + is_cuda_graph=False, ) lib._warm_kernel_variants(metadata=eager, extra=512, **profile) assert planned == [(20, 512), (15, 512), (7, 512)] # one per signature, largest first lib._warm_kernel_variants(metadata=eager, extra=512, **profile) assert len(planned) == 3 # same profile: already warmed in this process graph = SimpleNamespace( - fmha_plan_caches={}, max_num_tokens=20, max_num_requests=20, is_cuda_graph=True + fmha_plan_caches={}, + max_num_tokens=20, + max_num_requests=20, + max_total_draft_tokens=None, + mapping=None, + is_cuda_graph=True, ) lib._warm_kernel_variants(metadata=graph, extra=32, **profile) assert len(planned) == 3 # CUDA graph metadata never warms @@ -958,6 +1034,56 @@ class _Attn: assert len(planned) == 6 # the warmed plan is reused +def test_warm_set_covers_the_speculative_decoding_rows(monkeypatch) -> None: + """With MTP the generation step carries ``requests * (1 + draft)`` rows; the warm-up + plans the largest row count of every variant up to that ceiling (bounded by the + token budget), so the MTP row counts never compile during the first serving step.""" + import tensorrt_llm._torch.attention.backends.fmha.prims_ts_sparse_mla as mod + + planned: list = [] + _fake_sparse_plan(monkeypatch, planned, signature_div=8) + monkeypatch.setattr(mod, "_WARMED_PROFILES", set()) + monkeypatch.setenv("TLLM_PRIMTS_WARM_WORKERS", "0") + attn, profile, _ = _attn_and_profile() + lib = PrimsTSSparseMlaFmha(attn) + + mtp = SimpleNamespace( + fmha_plan_caches={}, + max_num_tokens=256, + max_num_requests=4, + max_total_draft_tokens=5, # DeepSeek-V4 MTP: next_n = 6 rows per request + mapping=None, + is_cuda_graph=False, + ) + lib._warm_kernel_variants(metadata=mtp, extra=512, **profile) + rows = [r for r, _ in planned] + assert rows[0] == 4 * (1 + 5), rows + assert rows == sorted(rows, reverse=True) + assert set(rows) == {24, 23, 15, 7} # the largest row count of each variant + + # The step's token budget bounds the rows even when the MTP ceiling is larger. + planned.clear() + monkeypatch.setattr(mod, "_WARMED_PROFILES", set()) + budget = SimpleNamespace( + fmha_plan_caches={}, + max_num_tokens=10, + max_num_requests=4, + max_total_draft_tokens=5, + mapping=None, + is_cuda_graph=False, + ) + lib._warm_kernel_variants(metadata=budget, extra=512, **profile) + assert [r for r, _ in planned][0] == 10 + + # A metadata class that does not declare the fields fails loudly. + monkeypatch.setattr(mod, "_WARMED_PROFILES", set()) + renamed = SimpleNamespace( + fmha_plan_caches={}, max_num_tokens=10, max_num_requests=4, is_cuda_graph=False + ) + with pytest.raises(AttributeError): + lib._warm_kernel_variants(metadata=renamed, extra=512, **profile) + + def test_index_contract_is_validated_once_per_process(monkeypatch) -> None: """One validated eager launch per contract (producers + pools); the same contract never re-validates, because that synchronizes the host inside the context ramp.""" @@ -1009,7 +1135,12 @@ class _Attn: device=torch.device("cuda", 0), q_dtype=torch.bfloat16, has_sinks=False, fused_groups=None ) eager = SimpleNamespace( - fmha_plan_caches={}, max_num_tokens=20, max_num_requests=20, is_cuda_graph=False + fmha_plan_caches={}, + max_num_tokens=20, + max_num_requests=20, + max_total_draft_tokens=None, + mapping=None, + is_cuda_graph=False, ) return attn, profile, eager @@ -1077,7 +1208,12 @@ def broken_workers(payloads, workers): monkeypatch.setenv("TLLM_PRIMTS_WARM_STRICT", "0") # the module fixture sets strict mode planned.clear() fresh = SimpleNamespace( # a new metadata object: the earlier plans are cached on `eager` - fmha_plan_caches={}, max_num_tokens=20, max_num_requests=20, is_cuda_graph=False + fmha_plan_caches={}, + max_num_tokens=20, + max_num_requests=20, + max_total_draft_tokens=None, + mapping=None, + is_cuda_graph=False, ) lib._warm_kernel_variants(metadata=fresh, extra=32, **profile) assert planned == [(20, 32), (15, 32), (7, 32)] @@ -1094,12 +1230,12 @@ def test_precompile_worker_count_follows_the_env(monkeypatch) -> None: monkeypatch.delenv("TLLM_PRIMTS_WARM_WORKERS", raising=False) monkeypatch.delenv("FLASHINFER_CUTE_DSL_DISABLE_CACHE", raising=False) monkeypatch.delenv("FLASHINFER_DISABLE_JIT", raising=False) - assert pp.worker_count(12) == 4 # default + assert pp.worker_count(12) == 4 # default for a single rank assert pp.worker_count(3) == 3 # never more workers than variants monkeypatch.setenv("TLLM_PRIMTS_WARM_WORKERS", "0") assert pp.worker_count(12) == 0 monkeypatch.setenv("TLLM_PRIMTS_WARM_WORKERS", "8") - assert pp.worker_count(12) == 8 + assert pp.worker_count(12) == 8 # explicit values are never capped monkeypatch.setenv("TLLM_PRIMTS_WARM_WORKERS", "many") assert pp.worker_count(12) == 4 # unparsable: default with a warning monkeypatch.setenv("FLASHINFER_CUTE_DSL_DISABLE_CACHE", "1") @@ -1109,6 +1245,39 @@ def test_precompile_worker_count_follows_the_env(monkeypatch) -> None: assert pp.worker_count(12) == 0 # compiling is forbidden +def test_default_worker_count_is_capped_by_the_ranks_on_the_host(monkeypatch) -> None: + """Under the default launcher every rank has a private FlashInfer workspace and fans + out on its own, so the per-rank default shrinks with the ranks on the host: the + host never starts more than 16 compile subprocesses at the default, and an + explicit ``TLLM_PRIMTS_WARM_WORKERS`` still wins.""" + import tensorrt_llm._torch.attention.backends.fmha.prims_ts_precompile as pp + + monkeypatch.delenv("TLLM_PRIMTS_WARM_WORKERS", raising=False) + monkeypatch.delenv("FLASHINFER_CUTE_DSL_DISABLE_CACHE", raising=False) + monkeypatch.delenv("FLASHINFER_DISABLE_JIT", raising=False) + expected = {1: 4, 2: 4, 4: 4, 8: 2, 16: 1, 32: 1} + for ranks, workers in expected.items(): + assert pp.default_worker_count(ranks) == workers, ranks + assert pp.worker_count(12, local_ranks=ranks) == workers, ranks + assert ranks * workers <= 16 or workers == 1 + monkeypatch.setenv("TLLM_PRIMTS_WARM_WORKERS", "4") + assert pp.worker_count(12, local_ranks=8) == 4 # explicit values are never capped + + # The rank count comes from the model's Mapping (no MPI call inside the + # forward), else from the launcher's environment, else 1. + for name in pp._LOCAL_RANK_ENV: + monkeypatch.delenv(name, raising=False) + assert pp.local_rank_count(SimpleNamespace(world_size=8, gpus_per_node=8)) == 8 + assert pp.local_rank_count(SimpleNamespace(world_size=16, gpus_per_node=8)) == 8 + assert pp.local_rank_count(SimpleNamespace(world_size=4, gpus_per_node=8)) == 4 + assert pp.local_rank_count(None) == 1 + monkeypatch.setenv("OMPI_COMM_WORLD_LOCAL_SIZE", "8") + assert pp.local_rank_count(None) == 8 + monkeypatch.setenv("OMPI_COMM_WORLD_LOCAL_SIZE", "many") + monkeypatch.setenv("RAY_LOCAL_WORLD_SIZE", "4") + assert pp.local_rank_count(None) == 4 + + def test_precompile_worker_entry_compiles_its_shard_first_then_the_rest(monkeypatch) -> None: """The worker entry point runs the real plan path under precompile_only(): own shard first, the rest afterwards, busy kernels deferred to a blocking pass."""