From 224c02f5f903c3684b6fe10ee30f490fe344c97d Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Tue, 13 Oct 2026 11:59:11 +0000 Subject: [PATCH 01/13] [None][chore] Pin vendored PrimTS to heyuhhh/flashinfer trtllm-prims-ts-dev Update tensorrt_llm/_torch/attention/backends/prims_ts to FlashInfer commit 0e0761ef5 of the heyuhhh/flashinfer trtllm-prims-ts-dev branch: the tip of yuxianq/flashinfer trtllm-prims-ts-dev-20261012 (ae3851ab5) plus flashinfer-ai/flashinfer#6014. The two commits added to trtllm-prims-ts-dev-20261012 since the current pin work around a Blackwell driver TMA descriptor issue in the context kernel: on SM100 and SM103 devices, ragged Q tensor maps use conservative out-of-bounds addressing unless the Q layout proves a contiguous 128 KiB prefix. flashinfer-ai/flashinfer#6014 tunes the block-sparse decode kernels for SM103 and adds a Sage plan that prepares its scales ahead of attention: K scales are read in place from the prepared buffer, V scales and the optional V mean are per batch entry ([B, Hkv, D]), and the route preparation launches as a dependent of the scale preparation. The vendor patch is unchanged apart from refreshed line offsets. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- .../vendor_patches/flashinfer-prims-ts.patch | 38 +- 3rdparty/vendor_sources.lock.yaml | 10 +- .../prims_ts/_block_sparse/compiler.py | 66 +- .../backends/prims_ts/_block_sparse/config.py | 22 +- .../prims_ts/_block_sparse/inspection.py | 4 +- .../backends/prims_ts/_block_sparse/plan.py | 33 +- .../prims_ts/_block_sparse/prepared.py | 3 +- .../prims_ts/_block_sparse/runtime.py | 12 +- .../backends/prims_ts/block_sparse.py | 74 +- .../attention/backends/prims_ts/context.py | 37 +- .../attention/backends/prims_ts/decode.py | 4 +- .../kernels/fmha_context/fmha_kernel.py | 39 + .../kernels/fmha_context/fmha_resources.py | 40 +- .../prims_ts/kernels/fmha_context/helpers.py | 23 +- .../fmha_decode/block_sparse_inspect.py | 40 +- .../fmha_decode/block_sparse_prepare.py | 1107 +++++++++++------ .../kernels/fmha_decode/fmha_decode_config.py | 207 ++- .../fmha_decode/fmha_decode_constants.py | 15 +- .../kernels/fmha_decode/fmha_decode_kernel.py | 345 ++--- .../fmha_decode_resources/__init__.py | 14 + .../helpers_kv_tile_idx.py | 43 +- .../fmha_decode_resources/helpers_softmax.py | 3 +- .../fmha_decode_resources/sage_scales.py | 969 +++++++-------- .../smem_block_sparse_metadata.py | 184 +-- .../fmha_decode_resources/smem_p.py | 218 ++-- .../fmha_decode_resources/smem_resources.py | 57 +- .../fmha_decode_resources/tmem_corr.py | 77 +- .../fmha_decode_resources/tmem_s.py | 458 +++---- .../kernels/fmha_decode/fmha_decode_tasks.py | 357 +++--- .../prims_ts/kernels/fmha_decode/reduction.py | 3 +- .../kernels/fmha_decode/sage_prepare.py | 364 ++++++ .../prims_ts/kernels/tcgen05_compat.py | 82 +- .../attention/backends/prims_ts/sage.py | 277 ++++- 33 files changed, 3071 insertions(+), 2154 deletions(-) create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/sage_prepare.py diff --git a/3rdparty/vendor_patches/flashinfer-prims-ts.patch b/3rdparty/vendor_patches/flashinfer-prims-ts.patch index 57ae33cb1935..48f3254537d8 100644 --- a/3rdparty/vendor_patches/flashinfer-prims-ts.patch +++ b/3rdparty/vendor_patches/flashinfer-prims-ts.patch @@ -1,5 +1,5 @@ diff --git a/block_sparse.py b/block_sparse.py -index d4652b8b68958ee3dab29d062a7cd082ea604c47..bb56724e986d302854722bf33962bea063f9e6e4 100644 +index fa055618705216cddf6b7f3e9ed4c1d680c2af9a..38de4926a045f17c5b6f206f4a20ecbe2532a1c7 100644 --- a/block_sparse.py +++ b/block_sparse.py @@ -1,3 +1,4 @@ @@ -7,7 +7,7 @@ index d4652b8b68958ee3dab29d062a7cd082ea604c47..bb56724e986d302854722bf33962bea0 # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); -@@ -36,13 +37,7 @@ from typing import Literal +@@ -37,13 +38,7 @@ from typing import Literal import torch @@ -22,7 +22,7 @@ index d4652b8b68958ee3dab29d062a7cd082ea604c47..bb56724e986d302854722bf33962bea0 from ._block_sparse.common import ( _num_sparse_pattern_heads, -@@ -301,7 +296,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase): +@@ -320,7 +315,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase): # previously published revision intact and runnable. self._plan_state = candidate @@ -31,7 +31,7 @@ index d4652b8b68958ee3dab29d062a7cd082ea604c47..bb56724e986d302854722bf33962bea0 def run( self, q: torch.Tensor, -@@ -433,7 +428,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase): +@@ -452,7 +447,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase): return self._launch_validated_run(state, run_args, run_stream) @@ -40,7 +40,7 @@ index d4652b8b68958ee3dab29d062a7cd082ea604c47..bb56724e986d302854722bf33962bea0 def block_sparse_attention( q: torch.Tensor, k: torch.Tensor, -@@ -726,9 +721,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase): +@@ -757,9 +752,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase): ) self._plan_state = candidate @@ -51,7 +51,7 @@ index d4652b8b68958ee3dab29d062a7cd082ea604c47..bb56724e986d302854722bf33962bea0 def run( self, q: torch.Tensor, -@@ -850,7 +843,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase): +@@ -881,7 +874,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase): return self._launch_validated_run(state, run_args, run_stream) @@ -61,10 +61,10 @@ index d4652b8b68958ee3dab29d062a7cd082ea604c47..bb56724e986d302854722bf33962bea0 q: torch.Tensor, paged_kv_cache: PagedKVCache, diff --git a/context.py b/context.py -index dd9a29c3e36111a0de9a57c487068b6b5caa4ea2..1163324ec632f2a0938f78a3636fb415450369df 100644 +index 53f01af72032ae2b01ea2ac374661af6387dd33f..d297bc79c1e00610783dd54a0d8cb22d1bcc2020 100644 --- a/context.py +++ b/context.py -@@ -46,7 +46,7 @@ from typing import TYPE_CHECKING, Callable, Literal, Optional +@@ -45,7 +45,7 @@ from typing import TYPE_CHECKING, Callable, Literal, Optional import torch @@ -73,7 +73,7 @@ index dd9a29c3e36111a0de9a57c487068b6b5caa4ea2..1163324ec632f2a0938f78a3636fb415 if TYPE_CHECKING: -@@ -488,7 +488,7 @@ def _validate_device(device: torch.device) -> int: +@@ -490,7 +490,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 +83,7 @@ index dd9a29c3e36111a0de9a57c487068b6b5caa4ea2..1163324ec632f2a0938f78a3636fb415 require_cute_dsl_arch(device_index) return device_index diff --git a/decode.py b/decode.py -index 41479a386d72c27eb0df76107318dc09256aff06..3fcec5efdc2a3825be0dc1d9a238aaaa6d76a609 100644 +index d85dd274917cf9a5a67ef1400155bb60985f1f77..205e875fc39022ca492cb40acea5b329fc4a3c11 100644 --- a/decode.py +++ b/decode.py @@ -1,3 +1,4 @@ @@ -113,7 +113,7 @@ index 41479a386d72c27eb0df76107318dc09256aff06..3fcec5efdc2a3825be0dc1d9a238aaaa require_cute_dsl_arch(device_index) -@@ -4046,7 +4043,7 @@ class BatchDecodePagedTSWrapper: +@@ -4048,7 +4045,7 @@ class BatchDecodePagedTSWrapper: # previous complete plan revision usable. self._plan_state = candidate @@ -122,7 +122,7 @@ index 41479a386d72c27eb0df76107318dc09256aff06..3fcec5efdc2a3825be0dc1d9a238aaaa def run( self, q: torch.Tensor, -@@ -4228,7 +4225,7 @@ class BatchDecodePagedTSWrapper: +@@ -4230,7 +4227,7 @@ class BatchDecodePagedTSWrapper: ) @@ -132,10 +132,10 @@ index 41479a386d72c27eb0df76107318dc09256aff06..3fcec5efdc2a3825be0dc1d9a238aaaa q: torch.Tensor, paged_kv_cache: PagedKVCache, diff --git a/kernels/fmha_context/fmha_resources.py b/kernels/fmha_context/fmha_resources.py -index c4d1b5512f9209793e42ec12bc269ac612c4703f..c21b2c9f5ab8d4f52d02024dcba15fdb2d16f080 100644 +index d17584b725e24235ce0ec90caac174a29d4ab965..e6cdddf8d5cb396be5feb828535c0633428e32c3 100644 --- a/kernels/fmha_context/fmha_resources.py +++ b/kernels/fmha_context/fmha_resources.py -@@ -1200,8 +1200,6 @@ class GmemQKVResource(MemoryResource): +@@ -1199,8 +1199,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 +144,7 @@ index c4d1b5512f9209793e42ec12bc269ac612c4703f..c21b2c9f5ab8d4f52d02024dcba15fdb seqlen_k = _load_runtime_seq_len_kv( self.seqlens_kv, self.max_seq_len_kv, batch_coord ) -@@ -1224,8 +1222,6 @@ class GmemQKVResource(MemoryResource): +@@ -1223,8 +1221,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 +153,7 @@ index c4d1b5512f9209793e42ec12bc269ac612c4703f..c21b2c9f5ab8d4f52d02024dcba15fdb kv_request_begin, kv_page_idx_ub = _load_block_table_row_bounds( Int32(self.block_table_row_stride), self.cfg, -@@ -1301,13 +1297,9 @@ class GmemQKVResource(MemoryResource): +@@ -1300,13 +1296,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 +167,7 @@ index c4d1b5512f9209793e42ec12bc269ac612c4703f..c21b2c9f5ab8d4f52d02024dcba15fdb kv_request_begin, kv_page_idx_ub = _load_block_table_row_bounds( Int32(self.block_table_row_stride), self.cfg, -@@ -1744,8 +1736,6 @@ class SmemPageOffsetsKvResource(MemoryResource): +@@ -1743,8 +1735,6 @@ class SmemPageOffsetsKvResource(MemoryResource): kv_request_begin: Int32, kv_page_idx_ub: Int32, ) -> None: @@ -176,7 +176,7 @@ index c4d1b5512f9209793e42ec12bc269ac612c4703f..c21b2c9f5ab8d4f52d02024dcba15fdb 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 -@@ -1836,8 +1826,6 @@ class SmemPageOffsetsKvResource(MemoryResource): +@@ -1835,8 +1825,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 +185,7 @@ index c4d1b5512f9209793e42ec12bc269ac612c4703f..c21b2c9f5ab8d4f52d02024dcba15fdb tile_idx = _resolve_kv_tile_idx_context( stage_info, kv_tile_start, tile_offset=tile_offset ) -@@ -5798,3 +5786,14 @@ class GmemOResource(MemoryResource): +@@ -5770,3 +5758,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) diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 617b168021c6..ba2cb10458d6 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -1,16 +1,16 @@ schema_version: 1 vendors: flashinfer-prims-ts: - url: https://github.com/yuxianq/flashinfer.git - branch: trtllm-prims-ts-dev-20261012 - commit: ae312f7cc1723b06fda5168105ade0f328ee47d8 + url: https://github.com/heyuhhh/flashinfer.git + branch: trtllm-prims-ts-dev + commit: 0e0761ef5ef10cb0762cc1b890a5c1a329bb49bf 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:a4360ca506d25545db408d22875e678871970fec94bdc1166dd6108d3b504d1c - digest: sha256-tree-v1:1bd57713d2487e94c0fc4b134bed4f7eff275282cb8ef471775e87db503ded03 + patch_digest: sha256:1ec340aac4fea9c77b9e9d956b1676c09cff5081053ef1e41f962ed38a6f728c + digest: sha256-tree-v1:e4129294b2b26295581cd9cbf1ac5ac4d7eec2ec65f0108626692ddf8a036f96 openengine-proto: url: https://github.com/ai-dynamo/openengine.git tag: v0.1.0 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 1a9d42857f94..c66cf74d37d7 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 @@ -36,11 +36,12 @@ def _compile_block_sparse(key: _BlockSparseCompileKey) -> Callable[..., object]: bitmask, compiles the same adapter signature. The plan decides at compile time which tensor slots it uses; the unused slots are ``None`` both for the compile-time fakes and for every later call, and the adapter body reads a - slot only on the branch its configuration enables. A dense plan runs no - prepare kernel and reaches the decode kernel directly; a block-sparse plan - prepares its routes and launches the prepared-route attention. Paged - block-sparse plans keep their own adapter because their K/V and request - metadata differ in kind, not only in presence. + slot only on the branch its configuration enables. A Sage plan first + prepares its scales; a dense plan then reaches the decode kernel + directly, while a block-sparse plan prepares its routes and launches the + prepared-route attention. Paged block-sparse plans keep their own adapter + because their K/V and request metadata differ in kind, not only in + presence. """ import cutlass @@ -56,6 +57,7 @@ def _compile_block_sparse(key: _BlockSparseCompileKey) -> Callable[..., object]: fmha_block_sparse_launch, fmha_decode_launch, ) + from ..kernels.fmha_decode.sage_prepare import _PrepareSageScales config = _make_block_sparse_config(key) sparse_format = key.sparse_format @@ -65,10 +67,20 @@ def _compile_block_sparse(key: _BlockSparseCompileKey) -> Callable[..., object]: pattern_heads = _num_sparse_pattern_heads( key.num_kv_heads, key.share_pattern_across_kv_heads ) + sage_scale_layout = None + prepare_sage_scales: _PrepareSageScales | None = None + if config.use_sage_attention: + sage_scale_layout = config.sage_scale_layout(key.seq_len_kv) + prepare_sage_scales = _PrepareSageScales( + cfg=config, + layout=sage_scale_layout, + batch_size=key.batch_size, + num_kv_heads=key.num_kv_heads, + ) if key.use_block_sparse: prepare_kwargs = { "batch_size": key.batch_size, - "num_kv_heads": pattern_heads, + "num_pattern_heads": pattern_heads, "seq_len_q": key.seq_len_q, "seq_len_kv": key.seq_len_kv, "q_block_size": key.q_block_size, @@ -78,6 +90,9 @@ def _compile_block_sparse(key: _BlockSparseCompileKey) -> Callable[..., object]: "use_causal_mask": key.mask_type == "causal", "apply_token_mask": key.use_kv_valid_bits, "store_score_words": config.uses_prepared_score_keep_words, + # The route prepare of a Sage plan is a programmatic dependent of + # the scale prepare launched immediately before it. + "use_pdl": prepare_sage_scales is not None, } if key.page_size is not None: if sparse_format != "bsr" or use_proxy_routes: @@ -113,6 +128,7 @@ def contiguous_adapter( kv_valid_bits: cute.Tensor | None, row_route_offsets: cute.Tensor | None, route_workspace: cute.Tensor | None, + sage_scales: cute.Tensor | None, max_blocks_per_row: cutlass.Int32, # The Sage scale slots, in ``SAGE_ADAPTER_SLOTS`` order. q_scale: cute.Tensor | None, @@ -136,26 +152,21 @@ def contiguous_adapter( Int32(static_seq_len_kv), Int32(static_head_dim), ) - # Sage scale tensors are [heads, flat slots]; the head stride is the - # flat slot count of one head. A slot the recipe does not use (the V - # mean without a mean, the summary K scales without proxy routes) is - # ``None`` and stays unbound. + # Sage scale tensors are [heads, flat slots] or [B, Hkv, D]. A slot the + # recipe does not use (the V mean without a mean, the summary K + # scales without proxy routes) is ``None`` and stays unbound. + # Attention reads every scale from the buffer the Sage prepare + # writes first, never the caller's tensors. The Sage prepare is an + # ordinary launch; the route prepare that follows it is a + # programmatic dependent launch that acquires it before exiting, so + # the stream-ordered attention launch finds both the scales and the + # routes complete. sage_kwargs = {} if cutlass.const_expr(static_config.use_sage_attention): - sage_kwargs = { - "sage_q_scale_iter": q_scale.iterator, - "sage_k_scale_iter": k_scale.iterator, - "sage_v_scale_iter": v_scale.iterator, - "sage_q_scale_head_stride": Int32(q_scale.shape[1]), - "sage_k_scale_head_stride": Int32(k_scale.shape[1]), - } - if cutlass.const_expr(static_config.sage_v_mean): - sage_kwargs["sage_v_mean_iter"] = v_mean.iterator - if cutlass.const_expr(static_config.use_block_sparse_proxy_routes): - sage_kwargs["sage_k_summary_scale_iter"] = k_summary_scale.iterator - sage_kwargs["sage_k_summary_scale_head_stride"] = Int32( - k_summary_scale.shape[1] - ) + sage_kwargs = {"sage_scales_iter": sage_scales.iterator} + prepare_sage_scales( + q_scale, k_scale, k_summary_scale, v_scale, v_mean, sage_scales, stream + ) if cutlass.const_expr(static_config.use_block_sparse): if cutlass.const_expr(sparse_format == "bsr"): prepare_routes( @@ -360,6 +371,7 @@ def fake_compact( valid_bits_fake = None row_route_offsets_fake = None route_workspace_fake = None + sage_scales_fake = None if use_proxy_routes: summary_shape = ( key.batch_size, @@ -399,6 +411,11 @@ def fake_compact( 4, ) route_workspace_fake = fake_compact(Int32, (cute.sym_int(),), 4) + if sage_scale_layout is not None: + sage_scales_fake = fake_compact( + Float32, + (sage_scale_layout.total_words(key.batch_size * key.num_kv_heads),), + ) sage_fakes: tuple[object | None, ...] = sage_adapter_slots({}) if key.sage is not None: # The plan's scale shapes fill the adapter slots through @@ -433,6 +450,7 @@ def fake_compact( valid_bits_fake, row_route_offsets_fake, route_workspace_fake, + sage_scales_fake, Int32(0), *sage_fakes, Float32(1.0), diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/config.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/config.py index 7f6fc2fb15de..e8234961ffba 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/config.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/config.py @@ -23,6 +23,7 @@ from flashinfer.utils import ceil_div from ..decode import _cutlass_dtype, _dtype_key, _validate_mask, _validate_positive_int +from ..kernels.fmha_decode.fmha_decode_constants import SOFTMAX_RESCALE_THRESHOLD_LOG2 from ..sage import SageAttentionConfig from .common import ( _PREPARED_KV_ROUTE_SIZE, @@ -30,6 +31,7 @@ _select_block_sparse_q_tile_size, _validate_sparse_kv_block_size, ) +from ..sage import _SageScaleLayout from .prepared import _BlockSparseRouteLayout if TYPE_CHECKING: @@ -92,6 +94,7 @@ class _BlockSparseCompileKey: # A dense key ignores the routing fields above. use_block_sparse: bool = True sage: SageAttentionConfig | None = None + softmax_rescale_threshold_log2: float = SOFTMAX_RESCALE_THRESHOLD_LOG2 @dataclass(frozen=True) @@ -103,6 +106,10 @@ class _BlockSparseLaunchSpec: # Whether prepared routes carry K32 score-validity words; the decode # config owns this rule and the plan sizes its route storage from it. prepares_score_words: bool + # The layout of a Sage plan's prepared scales, which the plan allocates + # and the Sage prepare writes before every attention launch; ``None`` + # otherwise. + sage_scales: _SageScaleLayout | None = None _CAPACITY_UNSET = object() @@ -136,6 +143,7 @@ class _BlockSparseStaticProfile: share_pattern_across_kv_heads: bool = False sage: SageAttentionConfig | None = None use_block_sparse: bool = True + softmax_rescale_threshold_log2: float = SOFTMAX_RESCALE_THRESHOLD_LOG2 def _select_block_sparse_kv_route_size( @@ -423,6 +431,7 @@ def _validate_block_sparse_static_profile( share_pattern_across_kv_heads: bool = False, sage: SageAttentionConfig | None = None, use_block_sparse: bool = True, + softmax_rescale_threshold_log2: float = SOFTMAX_RESCALE_THRESHOLD_LOG2, ) -> _BlockSparseStaticProfile: """Validate static policy before any device work or BSR inspection.""" @@ -517,13 +526,17 @@ def _validate_block_sparse_static_profile( share_pattern_across_kv_heads=share_pattern_across_kv_heads, sage=sage, use_block_sparse=use_block_sparse, + softmax_rescale_threshold_log2=softmax_rescale_threshold_log2, ) def _make_block_sparse_config(key: _BlockSparseCompileKey) -> "FmhaDecodeConfig": """Build one decode configuration from its exact compile cache key.""" - from ..kernels.fmha_decode.fmha_decode_config import make_decode_config + from ..kernels.fmha_decode.fmha_decode_config import ( + arch_config_args, + make_decode_config, + ) dtype = _cutlass_dtype(key.dtype_key) out_dtype = _cutlass_dtype(key.out_dtype_key) @@ -535,11 +548,13 @@ def _make_block_sparse_config(key: _BlockSparseCompileKey) -> "FmhaDecodeConfig" ) use_keeps_mma_ab = q_tile_size >= 64 config_args: dict[str, object] = { + **arch_config_args(torch.cuda.get_device_capability(key.device_index)), "use_keeps_mma_ab": use_keeps_mma_ab, "tile_size_q": q_tile_size, "tile_size_kv": key.kv_route_size, "groups_tokens_heads_q": True, "use_block_sparse": key.use_block_sparse, + "softmax_rescale_threshold_log2": key.softmax_rescale_threshold_log2, } if key.use_block_sparse: config_args.update( @@ -622,6 +637,7 @@ def _resolve_block_sparse_launch_spec( out_dtype_key: str, v_dtype_key: str, sage: SageAttentionConfig | None = None, + softmax_rescale_threshold_log2: float = SOFTMAX_RESCALE_THRESHOLD_LOG2, ) -> _BlockSparseLaunchSpec: """Resolve and cache one validated static or CLC launch. @@ -686,6 +702,7 @@ def parallel_loads(persistent: bool) -> bool: out_dtype_key=out_dtype_key, v_dtype_key=v_dtype_key, sage=sage, + softmax_rescale_threshold_log2=softmax_rescale_threshold_log2, ) try: config = _make_block_sparse_config(compile_key) @@ -735,6 +752,9 @@ def parallel_loads(persistent: bool) -> bool: policy=tuple(policy_entries), compile_key=compile_key, prepares_score_words=config.uses_prepared_score_keep_words, + sage_scales=( + config.sage_scale_layout(seq_len_kv) if config.use_sage_attention else None + ), ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/inspection.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/inspection.py index b8f0abe9820c..1b395caf8c2f 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/inspection.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/inspection.py @@ -127,7 +127,7 @@ def launch(summary: torch.Tensor, device_index: int) -> None: inspect_bsr = compile_block_sparse_inspection( device_index=device_index, batch_size=static.batch_size, - num_kv_heads=_num_sparse_pattern_heads( + num_pattern_heads=_num_sparse_pattern_heads( static.num_kv_heads, static.share_pattern_across_kv_heads ), seq_len_q=static.seq_len_q, @@ -172,7 +172,7 @@ def launch(summary: torch.Tensor, device_index: int) -> None: inspect_metadata = compile_paged_block_sparse_metadata_inspection( device_index=device_index, batch_size=static.batch_size, - num_kv_heads=_num_sparse_pattern_heads( + num_pattern_heads=_num_sparse_pattern_heads( static.num_kv_heads, static.share_pattern_across_kv_heads ), seq_len_q=static.seq_len_q, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/plan.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/plan.py index 4c721004d9d0..24b508e5b237 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/plan.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/plan.py @@ -35,6 +35,7 @@ _BlockSparseStaticProfile, _resolve_block_sparse_launch_spec, ) +from ..sage import _SageScaleLayout from .prepared import _BlockSparseRouteLayout _P = ParamSpec("_P") @@ -86,18 +87,20 @@ class _BlockSparsePlanState: Every block-sparse state executes one ``prepare -> prepared-route attention`` adapter, including a pattern whose rows select every KV block; - a dense state skips route preparation. Caller BSR and token mask tensors - belong to individual runs and are never retained here. + a dense state skips route preparation. A Sage state of either kind + prepares its scales first. Caller BSR and token mask tensors belong to + individual runs and are never retained here. Runtime geometry, dtypes, the compiled launch, and the readiness event are published together. ``run()`` therefore sees either the complete old state or the complete new one, never a mix. Plan-owned row offsets describe uniform capacity slices. Mutable route scratch, the optional unmasked-ABI - dummy mask, and the event are also state-owned. - Policy is immutable after publication; the cached compiled adapter is - shared read-only. + dummy mask, a Sage state's mutable prepared scales, and the event are also + state-owned. Policy is immutable after publication; the cached compiled + adapter is shared read-only. - One revision owns one mutable route workspace. Ordered runs on one + One revision owns one mutable route workspace and, for a Sage state, one + mutable prepared-scale buffer that every run rewrites. Ordered runs on one stream, or externally synchronized cross-stream runs, are valid. Unordered concurrent runs of the same revision are unsupported because their prepare launches would race. Different wrappers and published revisions own @@ -128,7 +131,7 @@ class _BlockSparsePlanState: page_size: int | None sage: SageAttentionConfig | None # Expected shape of every Sage scale tensor, by field name. - sage_scale_shapes: dict[str, tuple[int, int]] | None + sage_scale_shapes: dict[str, tuple[int, ...]] | None # Only an unmasked specialization needs a shape-correct ABI placeholder; # dense plans own neither it nor route storage. @@ -137,6 +140,10 @@ class _BlockSparsePlanState: # Immutable row capacities and mutable per-run route payload. row_route_offsets: torch.Tensor | None route_workspace: torch.Tensor | None + # Mutable per-run prepared scales of a Sage plan, in the layout the plan + # compiled. + sage_scales: torch.Tensor | None + sage_scales_layout: _SageScaleLayout | None # Semantic row bound; unlike route capacity, this distinguishes # multiple semantic blocks packed into one prepared route. max_blocks_per_row: int | None @@ -296,6 +303,7 @@ def _build_block_sparse_plan_state( out_dtype_key=static.out_dtype_key, v_dtype_key=static.v_dtype_key, sage=static.sage, + softmax_rescale_threshold_log2=static.softmax_rescale_threshold_log2, ) policy = ( *spec.policy, @@ -305,6 +313,13 @@ def _build_block_sparse_plan_state( dummy_kv_valid_bits = None row_route_offsets = None route_workspace = None + sage_scales = None + if spec.sage_scales is not None: + sage_scales = torch.empty( + spec.sage_scales.total_words(static.batch_size * static.num_kv_heads), + dtype=torch.float32, + device=device, + ) if static.use_block_sparse: route_layout = _BlockSparseRouteLayout.create( kv_route_size=static.kv_route_size, @@ -351,6 +366,8 @@ def _build_block_sparse_plan_state( dummy_kv_valid_bits=dummy_kv_valid_bits, row_route_offsets=row_route_offsets, route_workspace=route_workspace, + sage_scales=sage_scales, + sage_scales_layout=spec.sage_scales, max_blocks_per_row=static.max_blocks_per_row, policy=policy, compiled=compiled, @@ -377,6 +394,8 @@ def _wait_and_record_block_sparse_plan( state.row_route_offsets.record_stream(stream) if state.route_workspace is not None: state.route_workspace.record_stream(stream) + if state.sage_scales is not None: + state.sage_scales.record_stream(stream) __all__ = [ diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/prepared.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/prepared.py index 66815f890a55..8153ad8f2d53 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/prepared.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/prepared.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Shared layout of the route records produced before sparse attention.""" +"""Layout of the route records the block-sparse prepare writes before attention.""" from dataclasses import dataclass @@ -25,6 +25,7 @@ _PREPARED_ROUTE_IS_PROXY_FLAG = 1 << 1 _SUPPORTED_KV_ROUTE_SIZES = (128, 256) _SUPPORTED_PAGED_KV_PAGE_SIZES = (16, 32, 64, 128) +_FLOAT4_WORDS = 4 def _validate_int(value: object, name: str, *, allow_zero: bool) -> int: diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/runtime.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/runtime.py index 22027a9c5ef8..177145a15737 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/runtime.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/runtime.py @@ -146,7 +146,7 @@ def validate_block_sparse_metadata( batch_size: int, seq_len_q: int, seq_len_kv: int, - num_kv_heads: int, + num_pattern_heads: int, q_block_size: int, kv_block_size: int, use_kv_valid_bits: bool, @@ -167,7 +167,7 @@ def validate_block_sparse_metadata( ndim=3, dtype=torch.int32, expected_device=device, - expected_shape=(batch_size, num_kv_heads, num_q_blocks + 1), + expected_shape=(batch_size, num_pattern_heads, num_q_blocks + 1), ) _validate_metadata_tensor( block_indices, @@ -194,7 +194,7 @@ def validate_block_sparse_metadata( expected_device=device, expected_shape=( batch_size, - num_kv_heads, + num_pattern_heads, num_q_blocks, (num_kv_blocks + 31) // 32, ), @@ -362,7 +362,7 @@ def validate_block_sparse_run( batch_size=state.batch_size, seq_len_q=state.seq_len_q, seq_len_kv=state.seq_len_kv, - num_kv_heads=_num_sparse_pattern_heads( + num_pattern_heads=_num_sparse_pattern_heads( state.num_kv_heads, state.share_pattern_across_kv_heads ), q_block_size=state.q_block_size, @@ -522,7 +522,8 @@ def prepare_block_sparse_run_unchecked( Only the work every launch needs happens here: K/V view selection, the plan-owned dummy token mask when the plan disabled token bits, the default - softmax scale, and allocation of an omitted output tensor. + softmax scale, the Sage scale tensors in adapter slot order, and allocation + of an omitted output tensor. """ paged_kv: _PagedKVLaunchPayload | None = None @@ -652,6 +653,7 @@ def launch_block_sparse( run_args.kv_valid_bits, state.row_route_offsets, state.route_workspace, + state.sage_scales, 0 if state.max_blocks_per_row is None else state.max_blocks_per_row, *run_args.sage_slots, run_args.sm_scale, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/block_sparse.py b/tensorrt_llm/_torch/attention/backends/prims_ts/block_sparse.py index bb56724e986d..38de4926a045 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/block_sparse.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/block_sparse.py @@ -19,10 +19,11 @@ Q8/Q16/Q32 SWAPAB or Q64/Q128 KeepsAB specialization, and atomically publishes an immutable revision. ``run()`` validates compact BSHD Q/K/V and caller-owned routing tensors, then enqueues route preparation followed by attention in one -compiled adapter on the caller's CUDA stream. The one-shot entry point retains -synchronous canonical-BSR inspection to derive its temporary plan capacity. -A dense plan (``use_block_sparse=False``) keeps the tile selection but attends -over the whole K/V sequence without route preparation. +compiled adapter on the caller's CUDA stream; a Sage plan's run first gathers +its scales into the plan-owned prepared buffer. The one-shot entry point +retains synchronous canonical-BSR inspection to derive its temporary plan +capacity. A dense plan (``use_block_sparse=False``) keeps the tile selection +but attends over the whole K/V sequence without route preparation. ``BlockSparsePagedTSWrapper`` shares the same scheduling policy and decode kernel, but plans only fixed Q geometry and maximum K/V capacity. Page spans, @@ -75,6 +76,7 @@ _resolve_kv_dtypes, _validate_runtime_device, ) +from .kernels.fmha_decode.fmha_decode_constants import SOFTMAX_RESCALE_THRESHOLD_LOG2 from .sage import SageAttentionConfig, SageAttentionParams @@ -122,18 +124,21 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase): """Plan and reuse compact-BSHD block-sparse attention launches. Q is ``[B, Sq, Hq, D]`` and K/V are ``[B, Skv, Hkv, D]``. Sparse rows are - owned per batch, KV head, and query block, so every Q head in one grouped - KV head consumes the same sparse row. A plan fixes geometry and a per-row - capacity; every run supplies either BSR or a packed exact-block bitmask. - Proxy-enabled plans additionally consume caller-owned K/V summaries, while - an optional token mask applies only to exact routes. Callers must keep those - tensors alive and immutable until the queued run or captured graph finishes - using them. CUDA Graph capture pins plan-owned state only, so captured - routing storage remains the caller's responsibility. - - One plan revision owns one mutable route workspace. Its runs must be ordered - on one stream or externally synchronized; unordered concurrent runs require - distinct wrappers. + owned per batch, pattern head, and query block, so every Q head in one + grouped KV head consumes the same sparse row. A block-sparse plan fixes + geometry and a per-row capacity; every run supplies either BSR or a packed + exact-block bitmask. Proxy-enabled plans additionally consume caller-owned + K/V summaries, while an optional token mask applies only to exact routes. + A dense plan's run takes Q/K/V alone, and a Sage plan's run of either kind + adds its scale tensors. Callers must keep those tensors alive and immutable + until the queued run or captured graph finishes using them. CUDA Graph + capture pins plan-owned state only, so captured routing storage remains the + caller's responsibility. + + One block-sparse plan revision owns one mutable route workspace, and one + Sage plan revision owns its mutable prepared scales. Its runs must be + ordered on one stream or externally synchronized; unordered concurrent runs + require distinct wrappers. """ @_serialize_plan @@ -162,6 +167,7 @@ def plan( kv_data_type: torch.dtype | None = None, o_data_type: torch.dtype | None = None, sage_config: SageAttentionConfig | None = None, + softmax_rescale_threshold_log2: float = SOFTMAX_RESCALE_THRESHOLD_LOG2, ) -> None: """Choose a legal profile and allocate reusable routing capacity. @@ -179,7 +185,8 @@ def plan( ``use_block_sparse=False`` plans dense attention over the whole K/V sequence with the same Q-tile selection. Its route arguments must keep - their defaults, and :meth:`run` takes Q/K/V only. + their defaults, and :meth:`run` takes Q/K/V only, plus the scales of a + Sage plan. ``share_pattern_across_kv_heads=False`` retains independent per-KV-head patterns. True uses a singleton head axis in BSR/bitmask inputs and @@ -227,6 +234,17 @@ def plan( ``v_data_type``), and the output is ``torch.bfloat16`` (default) or ``torch.float16``. Every :meth:`run` then supplies the scale tensors as :class:`SageAttentionParams`. + + ``softmax_rescale_threshold_log2`` (0 to 8, default 8) applies to + block-sparse plans on the Q64/KV256 and Q128/KV128 profiles and to + dense Q64/KV256 plans: a row keeps its softmax exponent anchor while a + tile raises the row maximum by at most this many binades, which skips + the rescale of the accumulated output. Probabilities then reach + ``2**threshold``, so Sage FP8 P is quantized at ``448 * 2**-threshold`` + and probabilities far below the anchor lose E4M3 precision; rows whose + largest scores come first, such as attention sinks in the first + blocks, are the most exposed. 0 rescales on every new row maximum and + quantizes FP8 P at 448. """ if use_block_sparse: @@ -273,6 +291,7 @@ def plan( ), sage=sage_config, use_block_sparse=use_block_sparse, + softmax_rescale_threshold_log2=softmax_rescale_threshold_log2, ) if use_proxy_routes and static.mask_type != "dense": raise ValueError("block-sparse proxy routes require mask_type='dense'") @@ -451,6 +470,7 @@ def block_sparse_attention( out: torch.Tensor | None = None, sage: SageAttentionParams | None = None, sage_config: SageAttentionConfig | None = None, + softmax_rescale_threshold_log2: float = SOFTMAX_RESCALE_THRESHOLD_LOG2, ) -> torch.Tensor: """Plan and run one compact-BSHD block-sparse or dense attention launch. @@ -514,6 +534,9 @@ def block_sparse_attention( sage_config : SageAttentionConfig, optional The Sage recipe to plan. Defaults to :class:`SageAttentionConfig` with ``v_mean`` set when ``sage.v_mean`` is present. + softmax_rescale_threshold_log2 : float, optional + Binades a row maximum may rise before the output is rescaled, 0 to 8 + (default 8); see :meth:`BlockSparseTSWrapper.plan`. share_pattern_across_kv_heads : bool False (default) uses one pattern per KV head. True requires a singleton @@ -581,7 +604,7 @@ def block_sparse_attention( batch_size=static.batch_size, seq_len_q=static.seq_len_q, seq_len_kv=static.seq_len_kv, - num_kv_heads=_num_sparse_pattern_heads( + num_pattern_heads=_num_sparse_pattern_heads( static.num_kv_heads, static.share_pattern_across_kv_heads ), q_block_size=static.q_block_size, @@ -625,6 +648,7 @@ def block_sparse_attention( o_data_type=static.output_dtype, share_pattern_across_kv_heads=static.share_pattern_across_kv_heads, sage_config=sage_config, + softmax_rescale_threshold_log2=softmax_rescale_threshold_log2, ) return wrapper.run( q, @@ -666,6 +690,7 @@ def plan( q_data_type: torch.dtype = torch.float16, kv_data_type: torch.dtype | None = None, o_data_type: torch.dtype | None = None, + softmax_rescale_threshold_log2: float = SOFTMAX_RESCALE_THRESHOLD_LOG2, ) -> None: """Plan fixed-Q geometry and a maximum variable-K capacity. @@ -684,6 +709,11 @@ def plan( every physical Q tile within exactly one logical BSR row. ``kv_block_size`` remains restricted to 8, 16, 32, or a positive multiple of 64. + + ``softmax_rescale_threshold_log2`` (0 to 8, default 8) is the number of + binades a row maximum may rise on the Q64/KV256 and Q128/KV128 + profiles before the output is rescaled; see + :meth:`BlockSparseTSWrapper.plan`. """ static = _validate_block_sparse_static_profile( @@ -703,6 +733,7 @@ def plan( kv_dtype=kv_data_type, output_dtype=o_data_type, max_blocks_per_row=max_blocks_per_row, + softmax_rescale_threshold_log2=softmax_rescale_threshold_log2, ) assert static.page_size is not None device, device_index = _resolve_cuda_device(device) @@ -860,6 +891,7 @@ def block_sparse_attention_with_paged_kv_cache( share_pattern_across_kv_heads: bool = False, sm_scale: float | None = None, out: torch.Tensor | None = None, + softmax_rescale_threshold_log2: float = SOFTMAX_RESCALE_THRESHOLD_LOG2, ) -> torch.Tensor: """Plan and run one fixed-Q paged block-sparse attention launch. @@ -908,6 +940,9 @@ def block_sparse_attention_with_paged_kv_cache( out : torch.Tensor, optional Caller-owned compact output buffer ``[B, Sq, Hq, D]``. Must not overlap any live input; storage overlap is not checked. + softmax_rescale_threshold_log2 : float, optional + Binades a row maximum may rise before the output is rescaled, 0 to 8 + (default 8); see :meth:`BlockSparseTSWrapper.plan`. share_pattern_across_kv_heads : bool False (default) uses one pattern per KV head. True requires a singleton @@ -973,7 +1008,7 @@ def block_sparse_attention_with_paged_kv_cache( batch_size=static.batch_size, seq_len_q=static.seq_len_q, seq_len_kv=static.seq_len_kv, - num_kv_heads=_num_sparse_pattern_heads( + num_pattern_heads=_num_sparse_pattern_heads( static.num_kv_heads, static.share_pattern_across_kv_heads ), q_block_size=static.q_block_size, @@ -1010,6 +1045,7 @@ def block_sparse_attention_with_paged_kv_cache( kv_data_type=static.kv_dtype, o_data_type=static.output_dtype, share_pattern_across_kv_heads=static.share_pattern_across_kv_heads, + softmax_rescale_threshold_log2=softmax_rescale_threshold_log2, ) return wrapper.run( q, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/context.py b/tensorrt_llm/_torch/attention/backends/prims_ts/context.py index 1163324ec632..d297bc79c1e0 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/context.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/context.py @@ -37,7 +37,6 @@ from dataclasses import dataclass import functools -import importlib.metadata import itertools import math import numbers @@ -328,6 +327,9 @@ def _make_context_kernel( enable_skip_correction=True, uses_ldtm_stat=uses_ldtm_stat, causal_single_kv_tile=(causal_single_kv_tile and not use_paged_kv), + use_blackwell_tma_oob_workaround=( + torch.cuda.get_device_capability(device_index) in ((10, 0), (10, 3)) + ), **paged_kwargs, ) return fmha @@ -494,22 +496,6 @@ def _validate_device(device: torch.device) -> int: return device_index -@functools.cache -def _dsl_supports_ldtm_stat() -> bool: - """LDTM.STAT uses a native wrapper in DSL 4.8 or inline PTX in 4.7.""" - try: - dsl_version = importlib.metadata.version("nvidia-cutlass-dsl") - except importlib.metadata.PackageNotFoundError: - return False - from packaging import version as pkg_version - - try: - # Use .release so 4.7.0.dev* includes the inline-PTX path. - return pkg_version.Version(dsl_version).release >= (4, 7, 0) - except pkg_version.InvalidVersion: - return False - - def _default_exp2_fma_pairs(device_index: int, cfg) -> int: """FMA-pipe exp2 pairs per 16-pair softmax chunk on SM100 with 16-bit V. @@ -537,21 +523,10 @@ def _default_two_cta_umma(device_index: int) -> bool: def _default_uses_ldtm_stat(device_index: int) -> bool: - """Enable LDTM.STAT on SM103/SM107 with native or compatibility lowering.""" - if not _dsl_supports_ldtm_stat(): - return False - # tcgen05.ld.red.max (LDTM.STAT) is available on B300 (SM103) and Rubin - # (SM107), not B200 (SM100). - capability = torch.cuda.get_device_capability(device_index) - if capability == (10, 3): - return True - if capability == (10, 7): - # DSL 4.7's Rubin escape hatch targets sm_100f, which cannot emit - # LDTM.STAT. Keep Rubin on the native wrapper/target support path. - from cutlass.experimental import primitives as prims + """Enable LDTM.STAT where the shared tcgen05 gate allows it.""" + from .kernels.tcgen05_compat import ldtm_stat_supported - return hasattr(prims, "tcgen05_ld_red") - return False + return ldtm_stat_supported(torch.cuda.get_device_capability(device_index)) def _resolve_cuda_device( diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py b/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py index 3fcec5efdc2a..205e875fc390 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py @@ -1276,6 +1276,7 @@ def _resolve_decode_launch_spec( from .kernels.fmha_decode.fmha_decode_config import ( MIN_LOOP_ITERS_PER_SPLIT, + arch_config_args, get_max_active_clusters_for_cluster_size, make_decode_config, make_q_tile_geometry, @@ -1287,6 +1288,7 @@ def _resolve_decode_launch_spec( k_dtype = _cutlass_dtype(k_dtype_key) v_dtype = _cutlass_dtype(v_dtype_key) output_dtype = _cutlass_dtype(output_dtype_key) + arch_args = arch_config_args(torch.cuda.get_device_capability(device_index)) def make_config( args: object | None = None, @@ -1298,7 +1300,7 @@ def make_config( ) -> "FmhaDecodeConfig": return make_decode_config( headdim=head_dim, - args=args, + args=(arch_args, args), seq_len_q=seq_len_q, seq_len_kv=max_kv_len, batch_size=batch_size, 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 c74fbcc1287c..9b23b5dd7d38 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 @@ -200,6 +200,22 @@ from cutlass.experimental import primitives as prims +def _compact_q_prefix_min_rows(q: cute.Tensor) -> int | None: + """Rows needed for a proven 128 KiB contiguous prefix, or None if unproven.""" + if len(q.shape) != 3: + return None + h, d = q.shape[1:] + element_bits = q.element_type.width + if not all(isinstance(x, int) for x in (h, d, element_bits, *q.stride)): + return None + if h <= 0 or d <= 0 or element_bits not in (8, 16): + return None + if tuple(q.stride) != (h * d, d, 1): + return None + row_bytes = h * d * (element_bits // 8) + return (128 * 1024 + row_bytes - 1) // row_bytes + + def _as_i32(x: int | Int32) -> Int32: return Int32(x) if isinstance(x, int) else x @@ -2743,6 +2759,9 @@ class FmhaTs: Use the fixed causal one-K/V-tile task domains. The context runner enables this only for query-paired, fixed-length inputs whose K/V extent fits one 128-token tile (default: False). + use_blackwell_tma_oob_workaround : bool, optional + Use conservative OOB address generation for ragged Q tensor maps. + Enable only for a Blackwell execution device (default: False). """ def __init__( @@ -2772,6 +2791,7 @@ def __init__( causal_single_kv_tile: bool = False, exhaustive_deadlock_race_check: bool = True, d_v: int | None = None, + use_blackwell_tma_oob_workaround: bool = False, ) -> None: """Initialize mode-specific tiling, dtype, and schedule configuration.""" head_paired = resolve_head_paired_mode( @@ -2815,6 +2835,7 @@ def __init__( self.is_causal = is_causal self.is_clc_dynamic = is_clc_dynamic self.exhaustive_deadlock_race_check = exhaustive_deadlock_race_check + self.use_blackwell_tma_oob_workaround = use_blackwell_tma_oob_workaround q_dtype = in_qk_dtype or cutlass.Float16 k_dtype = in_qk_dtype or cutlass.Float16 @@ -3162,6 +3183,24 @@ def __call__( # Zero-fill partial Q rows and head-dimension fragments. oob_fill=cuda.TensorMapFloatOOBFill.none, ) + if cutlass.const_expr(self.use_blackwell_tma_oob_workaround): + # Only the real Q layout can establish a mapped prefix; + # synthetic ragged dimensions and padded head sizes cannot. + minimum_rows = _compact_q_prefix_min_rows(q_cute) + correct_oob_addressing = True + if cutlass.const_expr(minimum_rows is not None): + correct_oob_addressing = q_cute.shape[0] < minimum_rows + if correct_oob_addressing: + # Work around the Blackwell driver TMA descriptor issue documented at + # https://docs.nvidia.com/datacenter/tesla/tesla-release-notes-615-71-09/index.html#known-issues + # Clear Blackwell descriptor bit 85 on the host when a + # contiguous 128 KiB prefix has not been established. + descriptor_word1 = ( + cutlass.Pointer(tma_q_desc.get_ptr(), dtype=cutlass.Uint64) + 1 + ) + descriptor_word1.store( + descriptor_word1.load() & ~cutlass.Uint64(1 << 21) + ) else: tma_q_desc = cuda.create_tensor_map_tiled_from_view( q_cute, 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 c21b2c9f5ab8..e6cdddf8d5cb 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 @@ -87,15 +87,14 @@ from cutlass.cutlass_dsl import Boolean, Constexpr, dsl_user_op, if_generate from ..placeholder_helpers import _placeholder_smem_array, _placeholder_tmem_ptr +from ..tcgen05_compat import tcgen05_ld_32x32b_max from .helpers import ( bottom_right_window_left_bound, bottom_right_window_tile_start, freeze_smem_descriptor, - load_tmem_32x32b_max, variable_window_cta_min_start, ) from cutlass.experimental import primitives as prims -from cutlass._mlir.dialects import arith as arith_dialect SmemDescOffsets: TypeAlias = tuple[int, int] TmemAddr: TypeAlias = int | Int32 @@ -3223,38 +3222,11 @@ def _load_s_chunks_and_reduce_row_max( s_data: SoftmaxChunks = [None] * num_chunks chunk_maxima: tuple[Any, ...] = () for chunk_idx in cutlass.range_constexpr(num_chunks): - if cutlass.const_expr(hasattr(prims, "tcgen05_ld_red")): - loaded_words, red_word = prims.tcgen05_ld_red( - prims.Tcgen05LdStShape.SHAPE_32X32B, - prims.make_tmem_ptr( - tmem_s_addr + chunk_idx * tmem_x, self.cfg.qk_acc_dtype - ), - prims.ReductionKind.MAX, - num=tmem_x, - ) - s_data[chunk_idx] = cutlass.Vector.from_elements( - tuple( - loaded_words[i].bitcast(self.cfg.qk_acc_dtype) - for i in range(tmem_x) - ), - dtype=self.cfg.qk_acc_dtype, - ) - chunk_max = self.cfg.qk_acc_dtype( - arith_dialect.bitcast( - self.cfg.qk_acc_dtype.mlir_type, red_word.ir_value() - ) - ) - else: - # PTX ISA 8.8 supports LDTM.STAT on SM103 before the DSL 4.8 - # convenience wrapper is available. The geometry is fixed at - # 32 rows x 32 FP32 values for each context load fragment. - assert tmem_x == 32 - loaded = load_tmem_32x32b_max(tmem_s_addr + chunk_idx * tmem_x) - s_data[chunk_idx] = cutlass.Vector.from_elements( - loaded[:32], dtype=self.cfg.qk_acc_dtype - ) - chunk_max = loaded[32] - chunk_maxima += (chunk_max,) + loaded = tcgen05_ld_32x32b_max(tmem_s_addr + chunk_idx * tmem_x, tmem_x) + s_data[chunk_idx] = cutlass.Vector.from_elements( + loaded[:tmem_x], dtype=self.cfg.qk_acc_dtype + ) + chunk_maxima += (loaded[tmem_x],) cute.arch.fence_view_async_tmem_load() # Reduction results are asynchronous, just like the loaded scores. tile_max = cutlass.Vector.from_elements( diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/helpers.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/helpers.py index a66a6385baa7..bdaf9790be14 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/helpers.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/helpers.py @@ -17,7 +17,7 @@ import cutlass import cutlass.cute as cute -from cutlass import Float32, Int32, Int64 +from cutlass import Int32, Int64 @cute.jit @@ -94,24 +94,3 @@ def freeze_smem_descriptor(desc): write_only_types=[Int64], read_only_args=[desc], ) - - -_LDTM32_MAX_PTX = ( - "tcgen05.ld.red.sync.aligned.32x32b.x32.f32.max {" - + ", ".join(f"{{$w{i}}}" for i in range(32)) - + "}, {$w32}, [{$r0}];" -) - - -@cute.jit -def load_tmem_32x32b_max(tmem_addr: Int32) -> tuple: - """Load 32 FP32 scores and their maximum on SM103/SM107 (PTX 8.8). - - Compatibility wrapper for CUTLASS DSL 4.7. The caller gates the GPU - architecture and completes the TMEM loads before publishing/reusing data. - """ - return cute.arch.inline_ptx( - _LDTM32_MAX_PTX, - write_only_types=[Float32] * 33, - read_only_args=[tmem_addr], - ) 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 5728a7c6d373..da1e558acc79 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 @@ -15,7 +15,7 @@ """Validate paged request metadata and canonical BSR for one-shot planning. A BSR Q-block row is one ``block_indptr`` row keyed by -``(batch, kv_head, q_block)``. The maximum row width gives the one-shot API the +``(batch, pattern head, q_block)``. The maximum row width gives the one-shot API the same semantic ``max_blocks_per_row`` bound that reusable plans receive from their caller. Token-mask contents belong to the run-time prepare kernel and are not read here. @@ -92,7 +92,7 @@ def _inspect_bsr_row( block_indptr: cute.Tensor, block_indices: cute.Tensor, batch_idx: cutlass.Int32, - kv_head_idx: cutlass.Int32, + pattern_head_idx: cutlass.Int32, q_block_row_idx: cutlass.Int32, lane_idx: cutlass.Int32, bsr_row_is_valid: cutlass.Boolean, @@ -105,10 +105,10 @@ def _inspect_bsr_row( row_range_is_valid = cutlass.Int32(0) if lane_idx == cutlass.Int32(0) and bsr_row_is_valid: bsr_row_begin = cutlass.Int32( - block_indptr[batch_idx, kv_head_idx, q_block_row_idx] + block_indptr[batch_idx, pattern_head_idx, q_block_row_idx] ) bsr_row_end = cutlass.Int32( - block_indptr[batch_idx, kv_head_idx, q_block_row_idx + 1] + block_indptr[batch_idx, pattern_head_idx, q_block_row_idx + 1] ) row_range_is_valid = cutlass.Int32( bsr_row_begin >= cutlass.Int32(0) @@ -141,7 +141,7 @@ def _inspect_live_paged_bsr_row( block_indices: cute.Tensor, seq_lens_kv: cute.Tensor, batch_idx: cutlass.Int32, - kv_head_idx: cutlass.Int32, + pattern_head_idx: cutlass.Int32, q_block_row_idx: cutlass.Int32, lane_idx: cutlass.Int32, bsr_row_is_valid: cutlass.Boolean, @@ -162,7 +162,7 @@ def _inspect_live_paged_bsr_row( block_indptr, block_indices, batch_idx, - kv_head_idx, + pattern_head_idx, q_block_row_idx, lane_idx, bsr_row_is_valid, @@ -178,13 +178,13 @@ def __init__( self, *, batch_size: int, - num_kv_heads: int, + num_pattern_heads: int, seq_len_q: int, seq_len_kv: int | None, q_block_size: int, kv_block_size: int, ) -> None: - self.num_kv_heads = num_kv_heads + self.num_pattern_heads = num_pattern_heads self.num_q_block_rows = (seq_len_q + q_block_size - 1) // q_block_size self.num_kv_blocks = ( 0 @@ -192,7 +192,9 @@ def __init__( else (seq_len_kv + kv_block_size - 1) // kv_block_size ) self.kv_block_size = kv_block_size - self.total_bsr_row_count = batch_size * num_kv_heads * self.num_q_block_rows + self.total_bsr_row_count = ( + batch_size * num_pattern_heads * self.num_q_block_rows + ) @cute.jit def _launch( @@ -238,15 +240,15 @@ def kernel( ) q_block_row_idx = safe_linear_bsr_row_idx % self.num_q_block_rows linear_batch_head_idx = safe_linear_bsr_row_idx // self.num_q_block_rows - kv_head_idx = linear_batch_head_idx % self.num_kv_heads - batch_idx = linear_batch_head_idx // self.num_kv_heads + pattern_head_idx = linear_batch_head_idx % self.num_pattern_heads + batch_idx = linear_batch_head_idx // self.num_pattern_heads if cutlass.const_expr(seq_lens_kv is None): error_code, selected_kv_block_count = _inspect_bsr_row( block_indptr, block_indices, batch_idx, - kv_head_idx, + pattern_head_idx, q_block_row_idx, lane_idx, bsr_row_is_valid, @@ -258,7 +260,7 @@ def kernel( block_indices, seq_lens_kv, batch_idx, - kv_head_idx, + pattern_head_idx, q_block_row_idx, lane_idx, bsr_row_is_valid, @@ -454,7 +456,7 @@ def compile_block_sparse_inspection( *, device_index: int, batch_size: int, - num_kv_heads: int, + num_pattern_heads: int, seq_len_q: int, seq_len_kv: int, q_block_size: int, @@ -472,7 +474,7 @@ def compile_block_sparse_inspection( stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) kernel = _InspectBlockSparseBsr( batch_size=batch_size, - num_kv_heads=num_kv_heads, + num_pattern_heads=num_pattern_heads, seq_len_q=seq_len_q, seq_len_kv=seq_len_kv, q_block_size=q_block_size, @@ -483,7 +485,7 @@ def compile_block_sparse_inspection( kernel, _fake_compact( cutlass.Int32, - (batch_size, num_kv_heads, num_q_block_rows + 1), + (batch_size, num_pattern_heads, num_q_block_rows + 1), alignment=4, ), _fake_compact(cutlass.Int32, (logical_nnz,), alignment=4), @@ -499,7 +501,7 @@ def compile_paged_block_sparse_metadata_inspection( *, device_index: int, batch_size: int, - num_kv_heads: int, + num_pattern_heads: int, seq_len_q: int, minimum_seq_len_kv: int, max_seq_len_kv: int, @@ -523,7 +525,7 @@ def compile_paged_block_sparse_metadata_inspection( ) inspect_bsr = _InspectBlockSparseBsr( batch_size=batch_size, - num_kv_heads=num_kv_heads, + num_pattern_heads=num_pattern_heads, seq_len_q=seq_len_q, seq_len_kv=None, q_block_size=q_block_size, @@ -539,7 +541,7 @@ def compile_paged_block_sparse_metadata_inspection( inspect_paged_block_sparse_metadata, _fake_compact( cutlass.Int32, - (batch_size, num_kv_heads, num_q_block_rows + 1), + (batch_size, num_pattern_heads, num_q_block_rows + 1), alignment=4, ), _fake_compact(cutlass.Int32, (logical_nnz,), alignment=4), diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_prepare.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_prepare.py index 08f2872809cb..5c01660cb329 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_prepare.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_prepare.py @@ -15,18 +15,37 @@ """Prepare exact-first sparse routes for the PrimTS FMHA route consumer. The BSR frontend is shared by continuous exact/proxy and paged exact storage. -It validates each canonical row once, emits the same logical exact records, -then either resolves paged locators or appends continuous proxy records. The -bitmask frontend shares the record geometry and emitters but remains limited -to continuous storage. Proxy suffixes contain one stable record per summary +It validates each canonical row once and emits the same logical exact records; +a paged row resolves each atom's page locator in the pass that emits its +record, and a continuous row then appends its proxy records. The bitmask +frontend shares the record geometry and emitters but remains limited to +continuous storage. Proxy suffixes contain one stable record per summary group; a fully exact group remains present with zero score words. One warp owns one sparse row and four warps share a CTA. +Each lane owns one exact atom, so a warp pass emits +``32 / logical_origins_per_route`` complete exact records. A proxy record owns +an aligned group of ``max(logical_origins_per_route, token_words_per_route)`` +lanes, so a proxy pass emits 32 divided by that group size. Route masks and +flags reduce lane-group ballots. A pass stages its consecutive records in +shared memory and copies them to the workspace with coalesced stores that skip +record padding. A BSR lane reads its atom's block from ``block_indices``; the +bitmask frontend lists the selected blocks of 32 exact words at a time in a +per-warp shared-memory list, and each pass gathers its atoms' blocks from the +list, moving on to later batches until the one holding its last block. + ``row_route_offsets`` is a separate plan-owned immutable Int32 tensor. ``route_workspace`` contains only mutable row counts and route metadata described by ``_BlockSparseRouteLayout``. Payload outside each live row count is intentionally stale. ``max_blocks_per_row`` is the plan-declared semantic BSR-block limit, which remains distinct from packed-route capacity. + +A prepare launched with ``use_pdl`` is a programmatic dependent of the grid +launched immediately before it on the same stream, a Sage plan's scale +prepare, whose output the routes never read. Its CTAs may start while that +grid drains; each acquires the grid after its last store, so the prepare's +completion implies the producer's completion and memory visibility for the +attention launch that follows in stream order and reads both. """ from dataclasses import dataclass @@ -35,6 +54,8 @@ import cutlass.cute as cute from cuda.bindings import driver as cuda_drv from cutlass.cute.testing import assert_ as runtime_assert +from cutlass.experimental import primitives as prims +from cutlass.utils.smem_allocator import SmemAllocator from ..._block_sparse.prepared import ( _PREPARED_ROUTE_IS_FULL_FLAG, @@ -48,13 +69,18 @@ _WARPS_PER_CTA = 4 _WARP_SIZE = 32 _THREADS_PER_CTA = _WARPS_PER_CTA * _WARP_SIZE +# One bitmask batch loads a word per lane, so it selects at most 32 * 32 blocks. +_BITMASK_BATCH_BLOCKS = _WARP_SIZE * _WARP_SIZE @dataclass(frozen=True) class _RouteConfig: - """Compile-time route geometry shared across sparse input/storage modes.""" + """Compile-time route, warp-pass, and shared-memory geometry plus policy flags. - num_kv_heads: int + It is shared across sparse input and storage modes. + """ + + num_pattern_heads: int num_q_blocks: int num_kv_blocks: int num_exact_words: int @@ -69,6 +95,8 @@ class _RouteConfig: atom_valid_mask_word_offset: int route_flags_word_offset: int token_words_word_offset: int + # ``None`` selects the contiguous record, which has no page-ID words. + physical_page_ids_word_offset: int | None stores_score_words: bool apply_token_mask: bool use_proxy_routes: bool @@ -79,7 +107,7 @@ class _RouteConfig: def create( *, layout: _BlockSparseRouteLayout, - num_kv_heads: int, + num_pattern_heads: int, seq_len_q: int, seq_len_kv: int, q_block_size: int, @@ -96,8 +124,8 @@ def create( raise ValueError("proxy routes require prepared score words") num_q_blocks = (seq_len_q + q_block_size - 1) // q_block_size num_kv_blocks = (seq_len_kv + kv_block_size - 1) // kv_block_size - return _RouteConfig( - num_kv_heads=num_kv_heads, + cfg = _RouteConfig( + num_pattern_heads=num_pattern_heads, num_q_blocks=num_q_blocks, num_kv_blocks=num_kv_blocks, num_exact_words=(num_kv_blocks + _WARP_SIZE - 1) // _WARP_SIZE, @@ -117,12 +145,60 @@ def create( if layout.token_words_word_offset is not None else 0 ), + physical_page_ids_word_offset=( + layout.physical_page_ids_word_offset if layout.is_paged else None + ), stores_score_words=stores_score_words, apply_token_mask=apply_token_mask, use_proxy_routes=use_proxy_routes, route_metadata_stride_words=layout.route_metadata_stride_words, route_metadata_base_word_offset=layout.route_metadata_base_word_offset, ) + # Exact passes give each route origin a lane; proxy passes give each + # origin or score-word slot a lane. The layout's power-of-two route and + # atom sizes make both lane groups tile a warp. + assert _WARP_SIZE % cfg.proxy_lanes_per_route == 0 + return cfg + + @property + def exact_routes_per_pass(self) -> int: + """Exact records emitted by one warp pass holding one atom per lane.""" + + return _WARP_SIZE // self.logical_origins_per_route + + @property + def proxy_lanes_per_route(self) -> int: + """Lanes owning one proxy record's origin and score-word slots.""" + + return max(self.logical_origins_per_route, self.token_words_per_route) + + @property + def proxy_routes_per_pass(self) -> int: + """Proxy records emitted by one warp pass; never more than exact ones.""" + + return _WARP_SIZE // self.proxy_lanes_per_route + + @property + def route_record_words(self) -> int: + """Words of one record that carry data; the rest of its stride is padding.""" + + return ( + self.token_words_word_offset + self.token_words_per_route + if self.stores_score_words + else self.route_flags_word_offset + 1 + ) + + @property + def record_stage_words(self) -> int: + """Shared-memory words one warp uses to stage a pass of records.""" + + return self.exact_routes_per_pass * self.route_metadata_stride_words + + @property + def selected_block_list_words(self) -> int: + """Shared-memory words of one warp's compacted bitmask batch.""" + + return min(_BITMASK_BATCH_BLOCKS, self.num_exact_words * _WARP_SIZE) def _positive_i32_ceil_div( @@ -152,19 +228,37 @@ def _prepared_route_counts( @cute.jit -def _prepared_row_route_begin( +def _load_row_route_span( row_route_offsets: cute.Tensor, linear_row_idx: cutlass.Int32, lane_idx: cutlass.Int32, row_is_valid: cutlass.Boolean, - total_route_count: cutlass.Int32, -) -> cutlass.Int32: - """Load and validate one row's plan-owned prepared-route span.""" +) -> tuple[cutlass.Int32, cutlass.Int32]: + """Load one row's plan-owned prepared-route span into lane zero. + + The span depends only on the row, so kernels issue this load before the + row's routing metadata to overlap the two global-memory round trips. + """ row_route_begin = cutlass.Int32(0) + row_route_end = cutlass.Int32(0) if lane_idx == cutlass.Int32(0) and row_is_valid: row_route_begin = cutlass.Int32(row_route_offsets[linear_row_idx]) row_route_end = cutlass.Int32(row_route_offsets[linear_row_idx + 1]) + return row_route_begin, row_route_end + + +@cute.jit +def _prepared_row_route_begin( + row_route_begin: cutlass.Int32, + row_route_end: cutlass.Int32, + lane_idx: cutlass.Int32, + row_is_valid: cutlass.Boolean, + total_route_count: cutlass.Int32, +) -> cutlass.Int32: + """Validate one row's loaded route span and broadcast its first route.""" + + if lane_idx == cutlass.Int32(0) and row_is_valid: row_capacity = row_route_end - row_route_begin runtime_assert( row_route_begin >= cutlass.Int32(0) @@ -175,31 +269,107 @@ def _prepared_row_route_begin( return _warp_broadcast_i32(row_route_begin, 0) +@cute.jit +def _route_metadata_word_index( + row_route_begin: cutlass.Int32, + route_idx: cutlass.Int32, + cfg: cutlass.Constexpr[_RouteConfig], +) -> cutlass.Int32: + """Return the first workspace word of one row-relative route record.""" + + return cutlass.Int32(cfg.route_metadata_base_word_offset) + ( + row_route_begin + route_idx + ) * cutlass.Int32(cfg.route_metadata_stride_words) + + +@cute.jit +def _store_staged_records( + route_workspace: cute.Tensor, + record_stage: cute.Tensor, + first_record_word_index: cutlass.Int32, + live_record_count: cutlass.Int32, + lane_idx: cutlass.Int32, + records_per_pass: cutlass.Constexpr[int], + cfg: cutlass.Constexpr[_RouteConfig], +) -> None: + """Copy one pass of staged consecutive records to the workspace. + + Every store instruction covers 32 consecutive workspace words. Record + padding and records at or beyond ``live_record_count`` are skipped, so + exactly the data words of the live records are written. + """ + + cute.arch.sync_warp() + stride = cfg.route_metadata_stride_words + live_word_count = cutlass.Int32(records_per_pass * stride) + if live_record_count < cutlass.Int32(records_per_pass): + live_word_count = live_record_count * cutlass.Int32(stride) + for step in cutlass.range_constexpr( + (records_per_pass * stride + _WARP_SIZE - 1) // _WARP_SIZE + ): + word_idx = cutlass.Int32(step * _WARP_SIZE) + lane_idx + word_in_record = word_idx % cutlass.Int32(stride) + if word_idx < live_word_count and word_in_record < cutlass.Int32( + cfg.route_record_words + ): + route_workspace[first_record_word_index + word_idx] = record_stage[word_idx] + # The next pass restages only after every lane has copied this one. + cute.arch.sync_warp() + + +@cute.jit +def _lane_group_bits( + ballot: cutlass.Int32, + lane_idx: cutlass.Int32, + lanes_per_group: cutlass.Constexpr[int], +) -> cutlass.Uint32: + """Return the ballot bits of the aligned lane group containing this lane.""" + + group_bits = ballot.bitcast(cutlass.Uint32) + if cutlass.const_expr(lanes_per_group < _WARP_SIZE): + first_group_lane = lane_idx - lane_idx % cutlass.Int32(lanes_per_group) + group_bits = (group_bits >> first_group_lane) & cutlass.Uint32( + (1 << lanes_per_group) - 1 + ) + return group_bits + + +@cute.jit +def _lane_group_all( + predicate: cutlass.Boolean, + lane_idx: cutlass.Int32, + lanes_per_group: cutlass.Constexpr[int], +) -> cutlass.Boolean: + """Return whether ``predicate`` holds on every lane of this lane's group.""" + + return cutlass.Boolean( + _lane_group_bits( + cute.arch.vote_ballot_sync(predicate), lane_idx, lanes_per_group + ) + == cutlass.Uint32((1 << lanes_per_group) - 1) + ) + + @cute.jit def _resolve_route_logical_atom_origin( block_indices: cute.Tensor, row_begin: cutlass.Int32, row_end: cutlass.Int32, - route_idx: cutlass.Int32, - atom_in_route: cutlass.Int32, - kv_block_size: cutlass.Constexpr[int], - atom_size: cutlass.Constexpr[int], - logical_origins_per_route: cutlass.Constexpr[int], + flat_atom_idx: cutlass.Int32, seq_len_kv: cutlass.Int32, + cfg: cutlass.Constexpr[_RouteConfig], ) -> tuple[cutlass.Int32, cutlass.Boolean]: - """Resolve one route atom to its logical KV-token origin.""" + """Resolve a BSR row's ``flat_atom_idx``-th exact atom to its KV-token origin.""" - atoms_per_block = kv_block_size // atom_size - flat_atom_idx = route_idx * cutlass.Int32(logical_origins_per_route) + atom_in_route - bsr_entry_offset = flat_atom_idx // cutlass.Int32(atoms_per_block) - atom_in_block = flat_atom_idx % cutlass.Int32(atoms_per_block) + bsr_entry_offset = flat_atom_idx // cutlass.Int32(cfg.atoms_per_block) + atom_in_block = flat_atom_idx % cutlass.Int32(cfg.atoms_per_block) valid = cutlass.Boolean(bsr_entry_offset < row_end - row_begin) logical_origin = cutlass.Int32(-1) if valid: block_idx = cutlass.Int32(block_indices[row_begin + bsr_entry_offset]) - block_origin = block_idx * cutlass.Int32(kv_block_size) - atom_offset = atom_in_block * cutlass.Int32(atom_size) - valid = cutlass.Boolean(atom_offset < cutlass.Int32(seq_len_kv) - block_origin) + block_origin = block_idx * cutlass.Int32(cfg.kv_block_size) + atom_offset = atom_in_block * cutlass.Int32(cfg.atom_size) + valid = cutlass.Boolean(atom_offset < seq_len_kv - block_origin) if valid: logical_origin = block_origin + atom_offset return logical_origin, valid @@ -218,59 +388,38 @@ def _low_bits_mask(valid_bits: cutlass.Int32) -> cutlass.Uint32: @cute.jit -def _load_exact_score_word( - route_workspace: cute.Tensor, +def _load_atom_score_word( kv_valid_bits: cute.Tensor, - route_metadata_word_index: cutlass.Int32, - logical_word_idx: cutlass.Int32, + logical_origin: cutlass.Int32, + word_in_atom: cutlass.Constexpr[int], batch_idx: cutlass.Int32, seq_len_kv: cutlass.Int32, cfg: cutlass.Constexpr[_RouteConfig], ) -> cutlass.Uint32: - """Build one exact score word with optional caller-token masking.""" + """Build the score bits of one atom's ``word_in_atom``-th 32-token word. + An atom of at most 32 tokens has one word, returned in the low + ``atom_size`` bits. Tokens at or past ``seq_len_kv`` and, with caller-token + masking, masked tokens have no bit; an invalid atom (origin ``-1``) has + none. + """ + + word_origin = logical_origin + cutlass.Int32(word_in_atom * _WARP_SIZE) token_word = cutlass.Uint32(0) - if cutlass.const_expr(cfg.atom_size <= _WARP_SIZE): - atoms_per_word = _WARP_SIZE // cfg.atom_size - first_atom_idx = logical_word_idx * cutlass.Int32(atoms_per_word) - for atom_in_word in cutlass.range_constexpr(atoms_per_word): - atom_idx = first_atom_idx + cutlass.Int32(atom_in_word) - if atom_idx < cutlass.Int32(cfg.logical_origins_per_route): - origin = cutlass.Int32( - route_workspace[route_metadata_word_index + atom_idx] - ) - atom_word = cutlass.Uint32(0) - if origin >= cutlass.Int32(0): - if cutlass.const_expr(cfg.apply_token_mask): - source_word_idx = origin >> cutlass.Int32(5) - atom_word = cutlass.Uint32( - kv_valid_bits[batch_idx, source_word_idx] - ) - atom_word = atom_word >> (origin & cutlass.Int32(31)) - atom_word = atom_word & cutlass.Uint32((1 << cfg.atom_size) - 1) - atom_word = atom_word & _low_bits_mask(seq_len_kv - origin) - else: - atom_word = _low_bits_mask( - seq_len_kv - origin - ) & cutlass.Uint32((1 << cfg.atom_size) - 1) - token_word = token_word | ( - atom_word << cutlass.Int32(atom_in_word * cfg.atom_size) - ) - else: - words_per_atom = cfg.atom_size // _WARP_SIZE - atom_idx = logical_word_idx // cutlass.Int32(words_per_atom) - word_in_atom = logical_word_idx % cutlass.Int32(words_per_atom) - origin = cutlass.Int32(route_workspace[route_metadata_word_index + atom_idx]) - word_origin = origin + word_in_atom * cutlass.Int32(_WARP_SIZE) - if origin >= cutlass.Int32(0): - if cutlass.const_expr(cfg.apply_token_mask): - if word_origin < seq_len_kv: - source_word_idx = word_origin >> cutlass.Int32(5) - token_word = cutlass.Uint32( - kv_valid_bits[batch_idx, source_word_idx] - ) & _low_bits_mask(seq_len_kv - word_origin) - else: - token_word = _low_bits_mask(seq_len_kv - word_origin) + if logical_origin >= cutlass.Int32(0) and word_origin < seq_len_kv: + token_word = _low_bits_mask(seq_len_kv - word_origin) & cutlass.Uint32( + (1 << min(cfg.atom_size, _WARP_SIZE)) - 1 + ) + if cutlass.const_expr(cfg.apply_token_mask): + # Atoms start at multiples of their size, so only an atom narrower + # than a word sits at a nonzero bit offset of its source word. + bit_offset = word_origin & cutlass.Int32( + _WARP_SIZE - min(cfg.atom_size, _WARP_SIZE) + ) + source_word = cutlass.Uint32( + kv_valid_bits[batch_idx, word_origin >> cutlass.Int32(5)] + ) + token_word = token_word & (source_word >> bit_offset) return token_word @@ -291,11 +440,13 @@ def _resolve_prepared_bsr_row( if lane_idx == cutlass.Int32(0) and row_is_valid: q_block_row_idx = linear_row_idx % cfg.num_q_blocks linear_batch_head_idx = linear_row_idx // cfg.num_q_blocks - kv_head_idx = linear_batch_head_idx % cfg.num_kv_heads - batch_idx = linear_batch_head_idx // cfg.num_kv_heads - row_begin = cutlass.Int32(block_indptr[batch_idx, kv_head_idx, q_block_row_idx]) + pattern_head_idx = linear_batch_head_idx % cfg.num_pattern_heads + batch_idx = linear_batch_head_idx // cfg.num_pattern_heads + row_begin = cutlass.Int32( + block_indptr[batch_idx, pattern_head_idx, q_block_row_idx] + ) row_end = cutlass.Int32( - block_indptr[batch_idx, kv_head_idx, q_block_row_idx + 1] + block_indptr[batch_idx, pattern_head_idx, q_block_row_idx + 1] ) num_indices = cutlass.Int32(cute.size(block_indices)) runtime_assert( @@ -327,62 +478,94 @@ def _resolve_prepared_bsr_row( @cute.jit -def _finalize_exact_route( +def _emit_exact_route_pass( route_workspace: cute.Tensor, + record_stage: cute.Tensor, kv_valid_bits: cute.Tensor, - route_metadata_word_index: cutlass.Int32, - batch_idx: cutlass.Int32, + row_route_begin: cutlass.Int32, + first_route_idx: cutlass.Int32, + exact_route_count: cutlass.Int32, lane_idx: cutlass.Int32, - atom_is_valid: cutlass.Boolean, + logical_origin: cutlass.Int32, + physical_page_id: cutlass.Int32, + batch_idx: cutlass.Int32, seq_len_kv: cutlass.Int32, cfg: cutlass.Constexpr[_RouteConfig], ) -> None: - """Finalize an exact record after its logical origins are stored.""" - - atom_is_full = cutlass.Boolean(False) - if lane_idx < cutlass.Int32(cfg.logical_origins_per_route): - origin = cutlass.Int32(route_workspace[route_metadata_word_index + lane_idx]) - atom_is_full = cutlass.Boolean( - atom_is_valid and origin <= seq_len_kv - cutlass.Int32(cfg.atom_size) - ) - atom_valid_mask = cutlass.Int32(cute.arch.vote_ballot_sync(atom_is_valid)) - structural_full = cute.arch.vote_all_sync( - lane_idx >= cutlass.Int32(cfg.logical_origins_per_route) or atom_is_full + """Emit the exact records whose atoms the warp holds, one atom per lane. + + Lane ``i`` holds atom ``i % logical_origins_per_route`` of route + ``first_route_idx + i // logical_origins_per_route``; an invalid atom has + origin ``-1``. Records at or beyond ``exact_route_count`` are not stored. + """ + + origins_per_route = cfg.logical_origins_per_route + atom_in_route = lane_idx % cutlass.Int32(origins_per_route) + route_in_pass = lane_idx // cutlass.Int32(origins_per_route) + record_word_index = route_in_pass * cutlass.Int32(cfg.route_metadata_stride_words) + record_stage[record_word_index + atom_in_route] = logical_origin + if cutlass.const_expr(cfg.physical_page_ids_word_offset is not None): + record_stage[ + record_word_index + + cutlass.Int32(cfg.physical_page_ids_word_offset) + + atom_in_route + ] = physical_page_id + + # A record is full when every atom is structurally full and every score + # word it stores is full, so each lane folds its own words into its atom. + atom_is_valid = cutlass.Boolean(logical_origin >= cutlass.Int32(0)) + atom_is_full = cutlass.Boolean( + atom_is_valid and logical_origin <= seq_len_kv - cutlass.Int32(cfg.atom_size) ) - - score_words_are_full = cutlass.Boolean(True) if cutlass.const_expr(cfg.stores_score_words): - score_word = cutlass.Uint32(0) - if lane_idx < cutlass.Int32(cfg.token_words_per_route): - score_word = _load_exact_score_word( - route_workspace, - kv_valid_bits, - route_metadata_word_index, - lane_idx, - batch_idx, - seq_len_kv, - cfg, - ) - route_workspace[ - route_metadata_word_index - + cutlass.Int32(cfg.token_words_word_offset) - + lane_idx - ] = cutlass.Int32(score_word) - score_words_are_full = cute.arch.vote_all_sync( - lane_idx >= cutlass.Int32(cfg.token_words_per_route) - or score_word == cutlass.Uint32(0xFFFFFFFF) + # Atoms narrower than a word share it: an aligned group of adjacent + # lanes ORs its atoms' bits into one word. A wider atom fills several. + atoms_per_word = max(_WARP_SIZE // cfg.atom_size, 1) + words_per_atom = max(cfg.atom_size // _WARP_SIZE, 1) + atom_in_word = lane_idx % cutlass.Int32(atoms_per_word) + first_word_index = ( + record_word_index + + cutlass.Int32(cfg.token_words_word_offset) + + atom_in_route // cutlass.Int32(atoms_per_word) * words_per_atom ) - if lane_idx == cutlass.Int32(0): - route_workspace[ - route_metadata_word_index + cutlass.Int32(cfg.atom_valid_mask_word_offset) - ] = atom_valid_mask - route_workspace[ - route_metadata_word_index + cutlass.Int32(cfg.route_flags_word_offset) - ] = ( + for word_in_atom in cutlass.range_constexpr(words_per_atom): + score_word = _load_atom_score_word( + kv_valid_bits, logical_origin, word_in_atom, batch_idx, seq_len_kv, cfg + ) << (atom_in_word * cutlass.Int32(cfg.atom_size)) + for step in cutlass.range_constexpr(atoms_per_word.bit_length() - 1): + score_word = score_word | cute.arch.shuffle_sync_bfly( + score_word, offset=1 << step + ) + if atom_in_word == cutlass.Int32(0): + record_stage[first_word_index + word_in_atom] = cutlass.Int32( + score_word + ) + atom_is_full = cutlass.Boolean( + atom_is_full and score_word == cutlass.Uint32(0xFFFFFFFF) + ) + + atom_valid_mask = _lane_group_bits( + cute.arch.vote_ballot_sync(atom_is_valid), lane_idx, origins_per_route + ) + route_is_full = _lane_group_all(atom_is_full, lane_idx, origins_per_route) + if atom_in_route == cutlass.Int32(0): + record_stage[ + record_word_index + cutlass.Int32(cfg.atom_valid_mask_word_offset) + ] = atom_valid_mask.bitcast(cutlass.Int32) + record_stage[record_word_index + cutlass.Int32(cfg.route_flags_word_offset)] = ( cutlass.Int32(_PREPARED_ROUTE_IS_FULL_FLAG) - if structural_full and score_words_are_full + if route_is_full else cutlass.Int32(0) ) + _store_staged_records( + route_workspace, + record_stage, + _route_metadata_word_index(row_route_begin, first_route_idx, cfg), + exact_route_count - first_route_idx, + lane_idx, + cfg.exact_routes_per_pass, + cfg, + ) @cute.jit @@ -422,52 +605,28 @@ def _resolve_paged_route_atom_page_id( @cute.jit -def _exact_lane_rank( - exact_ballot: cutlass.Uint32, +def _inclusive_warp_prefix_sum( + value: cutlass.Int32, lane_idx: cutlass.Int32, - exact_prefix: cutlass.Int32, ) -> cutlass.Int32: - """Return one exact lane's global semantic-block rank.""" - - lower_lane_mask = (cutlass.Uint32(1) << lane_idx) - cutlass.Uint32(1) - return exact_prefix + cutlass.Int32(cute.arch.popc(exact_ballot & lower_lane_mask)) + """Return the inclusive warp prefix sum of one Int32 per lane.""" - -@cute.jit -def _emit_exact_block_atoms( - route_workspace: cute.Tensor, - row_route_begin: cutlass.Int32, - semantic_block_idx: cutlass.Int32, - exact_block_rank: cutlass.Int32, - cfg: cutlass.Constexpr[_RouteConfig], -) -> None: - """Expand one bitmask-selected block into fixed row-global atom slots.""" - - first_atom_rank = exact_block_rank * cutlass.Int32(cfg.atoms_per_block) - atom_in_block = cutlass.Int32(0) - while atom_in_block < cutlass.Int32(cfg.atoms_per_block): - atom_rank = first_atom_rank + atom_in_block - route_idx = atom_rank // cutlass.Int32(cfg.logical_origins_per_route) - atom_in_route = atom_rank % cutlass.Int32(cfg.logical_origins_per_route) - route_word_index = cutlass.Int32(cfg.route_metadata_base_word_offset) + ( - (row_route_begin + route_idx) - * cutlass.Int32(cfg.route_metadata_stride_words) + prefix_sum = value + for step in cutlass.range_constexpr((_WARP_SIZE - 1).bit_length()): + # A zero clamp keeps shfl.up lanes below the offset on their own value. + lower_sum = cute.arch.shuffle_sync_up( + prefix_sum, offset=1 << step, mask_and_clamp=0 ) - logical_origin = semantic_block_idx * cutlass.Int32( - cfg.kv_block_size - ) + atom_in_block * cutlass.Int32(cfg.atom_size) - stored_origin = cutlass.Int32(-1) - if logical_origin < cutlass.Int32(cfg.seq_len_kv): - stored_origin = logical_origin - route_workspace[route_word_index + atom_in_route] = stored_origin - atom_in_block += cutlass.Int32(1) + if lane_idx >= cutlass.Int32(1 << step): + prefix_sum = prefix_sum + lower_sum + return prefix_sum @cute.jit def _load_bitmask_word( exact_block_bits: cute.Tensor, batch_idx: cutlass.Int32, - kv_head_idx: cutlass.Int32, + pattern_head_idx: cutlass.Int32, q_block_idx: cutlass.Int32, logical_word_idx: cutlass.Int32, cfg: cutlass.Constexpr[_RouteConfig], @@ -479,13 +638,176 @@ def _load_bitmask_word( cutlass.Int32(cfg.num_kv_blocks) - logical_word_idx * cutlass.Int32(_WARP_SIZE) ) selected_word = cutlass.Uint32( - exact_block_bits[batch_idx, kv_head_idx, q_block_idx, logical_word_idx] + exact_block_bits[batch_idx, pattern_head_idx, q_block_idx, logical_word_idx] ) if cutlass.const_expr(for_proxy): selected_word = ~selected_word return valid_word & selected_word +@cute.jit +def _list_bitmask_batch( + exact_word: cutlass.Uint32, + batch_word_idx: cutlass.Int32, + selected_blocks: cute.Tensor, + lane_idx: cutlass.Int32, +) -> cutlass.Int32: + """List one batch's selected blocks in rank order; return their count. + + Lane ``i`` holds exact word ``batch_word_idx + i``. An inclusive warp + prefix sum of the per-word counts gives each lane its list slots, which it + fills from its highest selected block downward. The count is warp-uniform. + """ + + word_rank_end = _inclusive_warp_prefix_sum( + cutlass.Int32(cute.arch.popc(exact_word)), lane_idx + ) + first_word_block = (batch_word_idx + lane_idx) * cutlass.Int32(_WARP_SIZE) + remaining_word = exact_word + block_slot = word_rank_end + while remaining_word != cutlass.Uint32(0): + bit_idx = cutlass.Int32(cute.arch.bfind(remaining_word)) + block_slot -= cutlass.Int32(1) + selected_blocks[block_slot] = first_word_block + bit_idx + remaining_word = remaining_word ^ (cutlass.Uint32(1) << bit_idx) + # Lanes gather from the list only after every lane has filled its slots. + cute.arch.sync_warp() + return _warp_broadcast_i32(word_rank_end, _WARP_SIZE - 1) + + +@cute.jit +def _gather_listed_atom_origin( + selected_blocks: cute.Tensor, + flat_atom_idx: cutlass.Int32, + exact_atom_count: cutlass.Int32, + batch_rank_begin: cutlass.Int32, + batch_rank_end: cutlass.Int32, + logical_origin: cutlass.Int32, + cfg: cutlass.Constexpr[_RouteConfig], +) -> cutlass.Int32: + """Return the origin of an exact atom whose block is in the listed batch. + + Atoms past the row's exact atoms or outside the batch keep + ``logical_origin``. + """ + + block_rank = flat_atom_idx // cutlass.Int32(cfg.atoms_per_block) + if ( + flat_atom_idx < exact_atom_count + and block_rank >= batch_rank_begin + and block_rank < batch_rank_end + ): + semantic_block_idx = cutlass.Int32( + selected_blocks[block_rank - batch_rank_begin] + ) + candidate_origin = semantic_block_idx * cutlass.Int32( + cfg.kv_block_size + ) + flat_atom_idx % cutlass.Int32(cfg.atoms_per_block) * cutlass.Int32( + cfg.atom_size + ) + if candidate_origin < cutlass.Int32(cfg.seq_len_kv): + logical_origin = candidate_origin + return logical_origin + + +@cute.jit +def _emit_bitmask_exact_routes( + exact_block_bits: cute.Tensor, + kv_valid_bits: cute.Tensor, + route_workspace: cute.Tensor, + record_stage: cute.Tensor, + selected_blocks: cute.Tensor, + first_exact_word: cutlass.Uint32, + batch_idx: cutlass.Int32, + pattern_head_idx: cutlass.Int32, + q_block_idx: cutlass.Int32, + row_route_begin: cutlass.Int32, + exact_atom_count: cutlass.Int32, + exact_route_count: cutlass.Int32, + lane_idx: cutlass.Int32, + cfg: cutlass.Constexpr[_RouteConfig], +) -> None: + """Emit one bitmask row's exact records, one pass of atoms at a time. + + The row's selected blocks are listed in rank order one batch of 32 exact + words at a time, starting with the already loaded ``first_exact_word``. + Each pass gathers its atoms' blocks from the listed batch and moves on to + later batches, possibly empty ones, until the one holding the pass's last + block; atoms gathered from an earlier batch keep their origins. + """ + + batch_word_idx = cutlass.Int32(0) + batch_rank_begin = cutlass.Int32(0) + batch_rank_end = _list_bitmask_batch( + first_exact_word, batch_word_idx, selected_blocks, lane_idx + ) + first_route_idx = cutlass.Int32(0) + while first_route_idx < exact_route_count: + first_flat_atom_idx = first_route_idx * cutlass.Int32( + cfg.logical_origins_per_route + ) + flat_atom_idx = first_flat_atom_idx + lane_idx + last_flat_atom_idx = first_flat_atom_idx + cutlass.Int32(_WARP_SIZE - 1) + if last_flat_atom_idx >= exact_atom_count: + last_flat_atom_idx = exact_atom_count - cutlass.Int32(1) + last_block_rank = last_flat_atom_idx // cutlass.Int32(cfg.atoms_per_block) + logical_origin = _gather_listed_atom_origin( + selected_blocks, + flat_atom_idx, + exact_atom_count, + batch_rank_begin, + batch_rank_end, + cutlass.Int32(-1), + cfg, + ) + while last_block_rank >= batch_rank_end and batch_word_idx + cutlass.Int32( + _WARP_SIZE + ) < cutlass.Int32(cfg.num_exact_words): + # The next batch overwrites the list only after every lane has read it. + cute.arch.sync_warp() + batch_word_idx += cutlass.Int32(_WARP_SIZE) + word_idx = batch_word_idx + lane_idx + exact_word = cutlass.Uint32(0) + if word_idx < cutlass.Int32(cfg.num_exact_words): + exact_word = _load_bitmask_word( + exact_block_bits, + batch_idx, + pattern_head_idx, + q_block_idx, + word_idx, + cfg, + for_proxy=False, + ) + batch_rank_begin = batch_rank_end + batch_rank_end += _list_bitmask_batch( + exact_word, batch_word_idx, selected_blocks, lane_idx + ) + logical_origin = _gather_listed_atom_origin( + selected_blocks, + flat_atom_idx, + exact_atom_count, + batch_rank_begin, + batch_rank_end, + logical_origin, + cfg, + ) + _emit_exact_route_pass( + route_workspace, + record_stage, + kv_valid_bits, + row_route_begin, + first_route_idx, + exact_route_count, + lane_idx, + logical_origin, + cutlass.Int32(-1), + batch_idx, + cutlass.Int32(cfg.seq_len_kv), + cfg, + ) + first_route_idx += cutlass.Int32(cfg.exact_routes_per_pass) + + @cute.jit def _load_bsr_proxy_word( block_indices: cute.Tensor, @@ -523,66 +845,135 @@ def _load_bsr_proxy_word( @cute.jit -def _emit_proxy_route( +def _emit_proxy_routes( route_workspace: cute.Tensor, + record_stage: cute.Tensor, row_route_begin: cutlass.Int32, exact_route_count: cutlass.Int32, - group_idx: cutlass.Int32, - proxy_word: cutlass.Uint32, lane_idx: cutlass.Int32, cfg: cutlass.Constexpr[_RouteConfig], + load_proxy_word: cutlass.Constexpr, ) -> None: - """Emit one fixed summary-group proxy record, including an empty mask.""" - - route_metadata_word_index = cutlass.Int32(cfg.route_metadata_base_word_offset) + ( - row_route_begin + exact_route_count + group_idx - ) * cutlass.Int32(cfg.route_metadata_stride_words) - group_start = group_idx * cutlass.Int32(cfg.token_words_per_route * _WARP_SIZE) - group_size = cutlass.Int32(cfg.num_kv_blocks) - group_start - if group_size > cutlass.Int32(cfg.token_words_per_route * _WARP_SIZE): - group_size = cutlass.Int32(cfg.token_words_per_route * _WARP_SIZE) - origin_is_valid = cutlass.Boolean(False) - if lane_idx < cutlass.Int32(cfg.logical_origins_per_route): - summary_origin = group_start + lane_idx * cutlass.Int32(cfg.atom_size) - origin_is_valid = cutlass.Boolean(summary_origin < cfg.num_kv_blocks) - stored_origin = cutlass.Int32(-1) - if origin_is_valid: - stored_origin = summary_origin - route_workspace[route_metadata_word_index + lane_idx] = stored_origin - atom_valid_mask = cutlass.Int32(cute.arch.vote_ballot_sync(origin_is_valid)) - if lane_idx < cutlass.Int32(cfg.token_words_per_route): - route_workspace[ - route_metadata_word_index - + cutlass.Int32(cfg.token_words_word_offset) - + lane_idx - ] = cutlass.Int32(proxy_word) - score_full = cute.arch.vote_all_sync( - lane_idx >= cutlass.Int32(cfg.token_words_per_route) - or proxy_word == cutlass.Uint32(0xFFFFFFFF) - ) - if lane_idx == cutlass.Int32(0): - route_workspace[ - route_metadata_word_index + cutlass.Int32(cfg.atom_valid_mask_word_offset) - ] = atom_valid_mask - proxy_is_full = cutlass.Boolean( - group_size == cutlass.Int32(cfg.token_words_per_route * _WARP_SIZE) - and score_full + """Emit every fixed summary-group proxy record, including empty masks. + + ``load_proxy_word(logical_word_idx)`` returns one in-range word of the + row's proxy blocks, the complement of its exact blocks. Each aligned group + of ``proxy_lanes_per_route`` lanes owns one record per pass: its lane ``j`` + stores origin ``j`` and score word ``j`` when the record has them. + """ + + lanes_per_route = cfg.proxy_lanes_per_route + route_slot = lane_idx % cutlass.Int32(lanes_per_route) + route_in_pass = lane_idx // cutlass.Int32(lanes_per_route) + record_word_index = route_in_pass * cutlass.Int32(cfg.route_metadata_stride_words) + group_capacity = cfg.token_words_per_route * _WARP_SIZE + first_group_idx = cutlass.Int32(0) + while first_group_idx < cutlass.Int32(cfg.num_proxy_groups): + group_idx = first_group_idx + route_in_pass + logical_word_idx = group_idx * cutlass.Int32(cfg.token_words_per_route) + ( + route_slot + ) + proxy_word = cutlass.Uint32(0) + if route_slot < cutlass.Int32( + cfg.token_words_per_route + ) and logical_word_idx < cutlass.Int32(cfg.num_exact_words): + proxy_word = load_proxy_word(logical_word_idx) + + group_start = group_idx * cutlass.Int32(group_capacity) + origin_is_valid = cutlass.Boolean(False) + if route_slot < cutlass.Int32(cfg.logical_origins_per_route): + summary_origin = group_start + route_slot * cutlass.Int32(cfg.atom_size) + origin_is_valid = cutlass.Boolean(summary_origin < cfg.num_kv_blocks) + stored_origin = cutlass.Int32(-1) + if origin_is_valid: + stored_origin = summary_origin + record_stage[record_word_index + route_slot] = stored_origin + if route_slot < cutlass.Int32(cfg.token_words_per_route): + record_stage[ + record_word_index + + cutlass.Int32(cfg.token_words_word_offset) + + route_slot + ] = cutlass.Int32(proxy_word) + atom_valid_mask = _lane_group_bits( + cute.arch.vote_ballot_sync(origin_is_valid), lane_idx, lanes_per_route + ) + words_are_full = _lane_group_all( + route_slot >= cutlass.Int32(cfg.token_words_per_route) + or proxy_word == cutlass.Uint32(0xFFFFFFFF), + lane_idx, + lanes_per_route, ) - route_workspace[ - route_metadata_word_index + cutlass.Int32(cfg.route_flags_word_offset) - ] = cutlass.Int32(_PREPARED_ROUTE_IS_PROXY_FLAG) | ( - cutlass.Int32(proxy_is_full) * cutlass.Int32(_PREPARED_ROUTE_IS_FULL_FLAG) + if route_slot == cutlass.Int32(0): + record_stage[ + record_word_index + cutlass.Int32(cfg.atom_valid_mask_word_offset) + ] = atom_valid_mask.bitcast(cutlass.Int32) + proxy_is_full = cutlass.Boolean( + group_start + cutlass.Int32(group_capacity) + <= cutlass.Int32(cfg.num_kv_blocks) + and words_are_full + ) + record_stage[ + record_word_index + cutlass.Int32(cfg.route_flags_word_offset) + ] = cutlass.Int32(_PREPARED_ROUTE_IS_PROXY_FLAG) | ( + cutlass.Int32(proxy_is_full) + * cutlass.Int32(_PREPARED_ROUTE_IS_FULL_FLAG) + ) + _store_staged_records( + route_workspace, + record_stage, + _route_metadata_word_index( + row_route_begin, exact_route_count + first_group_idx, cfg + ), + cutlass.Int32(cfg.num_proxy_groups) - first_group_idx, + lane_idx, + cfg.proxy_routes_per_pass, + cfg, ) + first_group_idx += cutlass.Int32(cfg.proxy_routes_per_pass) + + +def _allocate_warp_words( + smem: SmemAllocator, + words_per_warp: int, + warp_idx: cutlass.Int32, +) -> cute.Tensor: + """Allocate ``words_per_warp`` shared Int32 words per warp; return this warp's.""" + + return smem.allocate_tensor( + cutlass.Int32, + cute.make_layout((words_per_warp, _WARPS_PER_CTA)), + byte_alignment=16, + )[None, warp_idx] + + +@cute.jit +def _acquire_pdl_producer(use_pdl: cutlass.Constexpr[bool]) -> None: + """Acquire the preceding grid of a programmatic dependent launch before exit. + + The routes depend on nothing that grid writes, so the prepare stores its + records first and acquires last; a prepare launched without ``use_pdl`` + has no dependency to acquire. + """ + + if cutlass.const_expr(use_pdl): + prims.griddepcontrol(kind=prims.GridDepAction.WAIT) class _PrepareRoutesBase: - """Own shared route geometry and compile-time storage/policy flags.""" + """Own shared route geometry and compile-time storage/policy flags. + + ``num_pattern_heads`` is the head axis of the BSR or bitmask inputs, one + when every KV head shares a pattern, and sizes the route rows. ``use_pdl`` + launches the prepare as a programmatic dependent of the grid launched + immediately before it on the same stream and acquires that grid before + the kernel exits. + """ def __init__( self, *, batch_size: int, - num_kv_heads: int, + num_pattern_heads: int, seq_len_q: int, seq_len_kv: int, q_block_size: int, @@ -593,6 +984,7 @@ def __init__( apply_token_mask: bool = False, store_score_words: bool = False, page_size: int | None = None, + use_pdl: bool = False, ) -> None: if not isinstance(use_proxy_routes, bool): raise TypeError("use_proxy_routes must be a bool") @@ -602,11 +994,13 @@ def __init__( raise TypeError("store_score_words must be a bool") if not isinstance(use_causal_mask, bool): raise TypeError("use_causal_mask must be a bool") + if not isinstance(use_pdl, bool): + raise TypeError("use_pdl must be a bool") if use_proxy_routes and page_size is not None: raise ValueError("paged KV does not support proxy routes") num_q_blocks = (seq_len_q + q_block_size - 1) // q_block_size - num_rows = batch_size * num_kv_heads * num_q_blocks + num_rows = batch_size * num_pattern_heads * num_q_blocks # Structural score words (sequence tail, invalid atoms) can be stored # without a caller token mask; proxy routes and token masks require them. stores_score_words = use_proxy_routes or apply_token_mask or store_score_words @@ -621,7 +1015,7 @@ def __init__( self.route_layout = layout self.cfg = _RouteConfig.create( layout=layout, - num_kv_heads=num_kv_heads, + num_pattern_heads=num_pattern_heads, seq_len_q=seq_len_q, seq_len_kv=seq_len_kv, q_block_size=q_block_size, @@ -631,11 +1025,15 @@ def __init__( ) self.page_size = page_size if page_size is not None else 1 self.minimum_seq_len_kv = seq_len_q if use_causal_mask else 1 - self.physical_page_ids_word_offset = ( - layout.physical_page_ids_word_offset if layout.is_paged else 0 - ) + self.use_pdl = use_pdl self.route_metadata_base_word_offset = layout.route_metadata_base_word_offset + @property + def grid(self) -> list[int]: + """Return the launch grid: one warp per route row, four rows per CTA.""" + + return [(self.cfg.num_rows + _WARPS_PER_CTA - 1) // _WARPS_PER_CTA, 1, 1] + class _PrepareBsrRoutes(_PrepareRoutesBase): """Prepare continuous exact/proxy or paged exact routes from one BSR flow.""" @@ -669,13 +1067,10 @@ def __call__( route_workspace, max_blocks_per_row, ).launch( - grid=[ - (self.cfg.num_rows + _WARPS_PER_CTA - 1) // _WARPS_PER_CTA, - 1, - 1, - ], + grid=self.grid, block=[_THREADS_PER_CTA, 1, 1], stream=stream, + use_pdl=self.use_pdl, ) @cute.kernel @@ -700,7 +1095,16 @@ def kernel( lane_idx = thread_idx % _WARP_SIZE linear_row_idx = block_idx * _WARPS_PER_CTA + warp_idx row_is_valid = linear_row_idx < self.cfg.num_rows + record_stage = _allocate_warp_words( + SmemAllocator(), self.cfg.record_stage_words, warp_idx + ) + row_route_span_begin, row_route_span_end = _load_row_route_span( + row_route_offsets, + linear_row_idx, + lane_idx, + row_is_valid, + ) row_begin, row_end, batch_idx = _resolve_prepared_bsr_row( block_indptr, block_indices, @@ -757,39 +1161,29 @@ def kernel( "selected BSR blocks exceed planned semantic capacity", ) row_route_begin = _prepared_row_route_begin( - row_route_offsets, - linear_row_idx, + row_route_span_begin, + row_route_span_end, lane_idx, row_is_valid, total_route_count, ) if row_is_valid: - route_idx = cutlass.Int32(0) - while route_idx < exact_route_count: - route_word_index = cutlass.Int32( - self.cfg.route_metadata_base_word_offset - ) + (row_route_begin + route_idx) * cutlass.Int32( - self.cfg.route_metadata_stride_words + first_route_idx = cutlass.Int32(0) + while first_route_idx < exact_route_count: + ( + logical_origin, + logical_origin_is_valid, + ) = _resolve_route_logical_atom_origin( + block_indices, + row_begin, + row_end, + first_route_idx * cutlass.Int32(self.cfg.logical_origins_per_route) + + lane_idx, + live_seq_len_kv, + self.cfg, ) - logical_origin = cutlass.Int32(-1) - logical_origin_is_valid = cutlass.Boolean(False) physical_page_id = cutlass.Int32(-1) - if lane_idx < cutlass.Int32(self.cfg.logical_origins_per_route): - ( - logical_origin, - logical_origin_is_valid, - ) = _resolve_route_logical_atom_origin( - block_indices, - row_begin, - row_end, - route_idx, - lane_idx, - self.cfg.kv_block_size, - self.cfg.atom_size, - self.cfg.logical_origins_per_route, - live_seq_len_kv, - ) if cutlass.const_expr(self.route_layout.is_paged): physical_page_id = _resolve_paged_route_atom_page_id( block_tables, @@ -801,58 +1195,38 @@ def kernel( self.page_size, num_physical_kv_pages, ) - if lane_idx < cutlass.Int32(self.cfg.logical_origins_per_route): - route_workspace[route_word_index + lane_idx] = logical_origin - if cutlass.const_expr(self.route_layout.is_paged): - route_workspace[ - route_word_index - + cutlass.Int32(self.physical_page_ids_word_offset) - + lane_idx - ] = physical_page_id - cute.arch.sync_warp() - - _finalize_exact_route( + _emit_exact_route_pass( route_workspace, + record_stage, kv_valid_bits, - route_word_index, - batch_idx, + row_route_begin, + first_route_idx, + exact_route_count, lane_idx, - logical_origin_is_valid, + logical_origin, + physical_page_id, + batch_idx, live_seq_len_kv, self.cfg, ) - route_idx += cutlass.Int32(1) + first_route_idx += cutlass.Int32(self.cfg.exact_routes_per_pass) if cutlass.const_expr(self.cfg.use_proxy_routes): - group_idx = cutlass.Int32(0) - while group_idx < cutlass.Int32(self.cfg.num_proxy_groups): - proxy_word = cutlass.Uint32(0) - logical_word_idx = ( - group_idx * cutlass.Int32(self.cfg.token_words_per_route) - + lane_idx - ) - if lane_idx < cutlass.Int32(self.cfg.token_words_per_route): - if logical_word_idx < cutlass.Int32(self.cfg.num_exact_words): - proxy_word = _load_bsr_proxy_word( - block_indices, - row_begin, - row_end, - logical_word_idx, - self.cfg, - ) - _emit_proxy_route( - route_workspace, - row_route_begin, - exact_route_count, - group_idx, - proxy_word, - lane_idx, - self.cfg, - ) - group_idx += cutlass.Int32(1) + _emit_proxy_routes( + route_workspace, + record_stage, + row_route_begin, + exact_route_count, + lane_idx, + self.cfg, + lambda word_idx: _load_bsr_proxy_word( + block_indices, row_begin, row_end, word_idx, self.cfg + ), + ) if lane_idx == cutlass.Int32(0) and row_is_valid: route_workspace[linear_row_idx] = total_route_count + _acquire_pdl_producer(self.use_pdl) class _PrepareBitmaskRoutes(_PrepareRoutesBase): @@ -862,7 +1236,7 @@ def __init__( self, *, batch_size: int, - num_kv_heads: int, + num_pattern_heads: int, seq_len_q: int, seq_len_kv: int, q_block_size: int, @@ -872,10 +1246,11 @@ def __init__( use_causal_mask: bool = False, apply_token_mask: bool = False, store_score_words: bool = False, + use_pdl: bool = False, ) -> None: super().__init__( batch_size=batch_size, - num_kv_heads=num_kv_heads, + num_pattern_heads=num_pattern_heads, seq_len_q=seq_len_q, seq_len_kv=seq_len_kv, q_block_size=q_block_size, @@ -885,6 +1260,7 @@ def __init__( use_causal_mask=use_causal_mask, apply_token_mask=apply_token_mask, store_score_words=store_score_words, + use_pdl=use_pdl, ) @cute.jit @@ -904,13 +1280,10 @@ def __call__( route_workspace, max_blocks_per_row, ).launch( - grid=[ - (self.cfg.num_rows + _WARPS_PER_CTA - 1) // _WARPS_PER_CTA, - 1, - 1, - ], + grid=self.grid, block=[_THREADS_PER_CTA, 1, 1], stream=stream, + use_pdl=self.use_pdl, ) @cute.kernel @@ -932,17 +1305,40 @@ def kernel( row_is_valid = linear_row_idx < self.cfg.num_rows q_block_idx = linear_row_idx % self.cfg.num_q_blocks linear_batch_head_idx = linear_row_idx // self.cfg.num_q_blocks - kv_head_idx = linear_batch_head_idx % self.cfg.num_kv_heads - batch_idx = linear_batch_head_idx // self.cfg.num_kv_heads + pattern_head_idx = linear_batch_head_idx % self.cfg.num_pattern_heads + batch_idx = linear_batch_head_idx // self.cfg.num_pattern_heads + smem = SmemAllocator() + record_stage = _allocate_warp_words(smem, self.cfg.record_stage_words, warp_idx) + selected_blocks = _allocate_warp_words( + smem, self.cfg.selected_block_list_words, warp_idx + ) + row_route_span_begin, row_route_span_end = _load_row_route_span( + row_route_offsets, + linear_row_idx, + lane_idx, + row_is_valid, + ) + # Count the first word after the loop so its load overlaps the others. + first_exact_word = cutlass.Uint32(0) + if row_is_valid and lane_idx < cutlass.Int32(self.cfg.num_exact_words): + first_exact_word = _load_bitmask_word( + exact_block_bits, + batch_idx, + pattern_head_idx, + q_block_idx, + lane_idx, + self.cfg, + for_proxy=False, + ) lane_exact_count = cutlass.Int32(0) - word_idx = lane_idx + word_idx = lane_idx + cutlass.Int32(_WARP_SIZE) while word_idx < cutlass.Int32(self.cfg.num_exact_words): if row_is_valid: exact_word = _load_bitmask_word( exact_block_bits, batch_idx, - kv_head_idx, + pattern_head_idx, q_block_idx, word_idx, self.cfg, @@ -950,6 +1346,7 @@ def kernel( ) lane_exact_count += cutlass.Int32(cute.arch.popc(exact_word)) word_idx += cutlass.Int32(_WARP_SIZE) + lane_exact_count += cutlass.Int32(cute.arch.popc(first_exact_word)) exact_block_count = cutlass.Int32( cute.arch.warp_redux_sync(lane_exact_count, "add") ) @@ -964,123 +1361,53 @@ def kernel( "selected bitmask blocks exceed planned semantic capacity", ) row_route_begin = _prepared_row_route_begin( - row_route_offsets, - linear_row_idx, + row_route_span_begin, + row_route_span_end, lane_idx, row_is_valid, total_route_count, ) if row_is_valid: - exact_prefix = cutlass.Int32(0) - word_idx = cutlass.Int32(0) - while word_idx < cutlass.Int32(self.cfg.num_exact_words): - exact_word_i32 = cutlass.Int32(0) - if lane_idx == cutlass.Int32(0): - exact_word_i32 = _load_bitmask_word( - exact_block_bits, - batch_idx, - kv_head_idx, - q_block_idx, - word_idx, - self.cfg, - for_proxy=False, - ).bitcast(cutlass.Int32) - exact_word = _warp_broadcast_i32(exact_word_i32, 0).bitcast( - cutlass.Uint32 - ) - is_exact = cutlass.Boolean( - (exact_word & (cutlass.Uint32(1) << lane_idx)) != cutlass.Uint32(0) - ) - exact_ballot = cute.arch.vote_ballot_sync(is_exact).bitcast( - cutlass.Uint32 - ) - exact_rank = _exact_lane_rank(exact_ballot, lane_idx, exact_prefix) - if is_exact: - _emit_exact_block_atoms( - route_workspace, - row_route_begin, - word_idx * cutlass.Int32(_WARP_SIZE) + lane_idx, - exact_rank, - self.cfg, - ) - exact_prefix += cutlass.Int32(cute.arch.popc(exact_ballot)) - word_idx += cutlass.Int32(1) - - final_route_atom_count = exact_atom_count % cutlass.Int32( - self.cfg.logical_origins_per_route + _emit_bitmask_exact_routes( + exact_block_bits, + kv_valid_bits, + route_workspace, + record_stage, + selected_blocks, + first_exact_word, + batch_idx, + pattern_head_idx, + q_block_idx, + row_route_begin, + exact_atom_count, + exact_route_count, + lane_idx, + self.cfg, ) - if final_route_atom_count != cutlass.Int32(0): - if lane_idx >= final_route_atom_count and lane_idx < cutlass.Int32( - self.cfg.logical_origins_per_route - ): - final_route_word_index = cutlass.Int32( - self.cfg.route_metadata_base_word_offset - ) + (row_route_begin + exact_route_count - cutlass.Int32(1)) * ( - cutlass.Int32(self.cfg.route_metadata_stride_words) - ) - route_workspace[final_route_word_index + lane_idx] = cutlass.Int32( - -1 - ) - cute.arch.sync_warp() - route_idx = cutlass.Int32(0) - while route_idx < exact_route_count: - route_word_index = cutlass.Int32( - self.cfg.route_metadata_base_word_offset - ) + (row_route_begin + route_idx) * cutlass.Int32( - self.cfg.route_metadata_stride_words - ) - atom_is_valid = cutlass.Boolean(False) - if lane_idx < cutlass.Int32(self.cfg.logical_origins_per_route): - atom_is_valid = cutlass.Boolean( - cutlass.Int32(route_workspace[route_word_index + lane_idx]) - >= cutlass.Int32(0) - ) - _finalize_exact_route( + if cutlass.const_expr(self.cfg.use_proxy_routes): + _emit_proxy_routes( route_workspace, - kv_valid_bits, - route_word_index, - batch_idx, + record_stage, + row_route_begin, + exact_route_count, lane_idx, - atom_is_valid, - cutlass.Int32(self.cfg.seq_len_kv), self.cfg, - ) - route_idx += cutlass.Int32(1) - - if cutlass.const_expr(self.cfg.use_proxy_routes): - group_idx = cutlass.Int32(0) - while group_idx < cutlass.Int32(self.cfg.num_proxy_groups): - proxy_word = cutlass.Uint32(0) - logical_word_idx = ( - group_idx * cutlass.Int32(self.cfg.token_words_per_route) - + lane_idx - ) - if lane_idx < cutlass.Int32(self.cfg.token_words_per_route): - if logical_word_idx < cutlass.Int32(self.cfg.num_exact_words): - proxy_word = _load_bitmask_word( - exact_block_bits, - batch_idx, - kv_head_idx, - q_block_idx, - logical_word_idx, - self.cfg, - for_proxy=True, - ) - _emit_proxy_route( - route_workspace, - row_route_begin, - exact_route_count, - group_idx, - proxy_word, - lane_idx, + lambda word_idx: _load_bitmask_word( + exact_block_bits, + batch_idx, + pattern_head_idx, + q_block_idx, + word_idx, self.cfg, - ) - group_idx += cutlass.Int32(1) + for_proxy=True, + ), + ) if lane_idx == cutlass.Int32(0) and row_is_valid: route_workspace[linear_row_idx] = total_route_count + _acquire_pdl_producer(self.use_pdl) __all__ = ["_PrepareBitmaskRoutes", "_PrepareBsrRoutes"] 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 b0efb7adc248..af54b06004ce 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 @@ -35,8 +35,12 @@ _select_block_sparse_q_tile_size, _validate_sparse_kv_block_size, ) +from ...sage import _SAGE_ATOM_MIN_WORDS, _SAGE_ATOM_TOKENS, _SageScaleLayout from ...split_kv_mode_policy import select_split_kv_modes +from ..tcgen05_compat import ldtm_stat_supported from .fmha_decode_constants import ( + E4M3_MAX, + SOFTMAX_RESCALE_THRESHOLD_LOG2, AUTO_LAUNCH_TILE_SIZE_KV, BITS_PER_BYTE, BYTES_PER_KIB, @@ -780,6 +784,16 @@ def uses_scattered_page_route(self) -> bool: """Whether native paged loads consume one locator per page fragment.""" return self.has_storage_subpages or self.use_q_token_kv_block_sparse_route + @property + def uses_page_offsets_task(self) -> bool: + """Whether a page-table producer task stages page IDs for the load warp. + + Every paged plan that is not block-sparse takes it, the + QToken-KvBlock-Sparse paged routes included; a block-sparse paged plan + reads each atom's physical page from its prepared route instead. + """ + return self.use_paged_kv and not self.use_block_sparse + @property def shares_sparse_pattern(self) -> bool: """Q-block routes default to per-head patterns; token routes share.""" @@ -1379,6 +1393,10 @@ def use_fp8_kv(self) -> bool: sage_k_block_size: int = 0 sage_k_summary_block_size: int = 0 sage_v_mean: bool = False + # Binades a deferred exponent anchor may lag the row maximum, at most + # ``SOFTMAX_RESCALE_THRESHOLD_LOG2``; 0 keeps exact anchors and None the + # P dtype's default. See ``softmax_anchor_headroom_log2``. + softmax_rescale_threshold_log2: float | None = None # Nonzero means each K/V stage covers only this many head-dim columns. # H256 SwapsMmaAb uses 128-column stages to keep TMA and TMEM layouts valid. head_dim_per_stage_kv: int = 0 @@ -1390,6 +1408,13 @@ def use_fp8_kv(self) -> bool: softmax_order_barrier_id: int = 8 # Both softmax task groups participate: 8 warps * 32 lanes. softmax_order_barrier_threads: int = 256 + # Score pairs per streamed K32 fragment whose exponentials run as FMA + # polynomials instead of MUFU. The best value depends on the GPU's MUFU + # rate; ``arch_config_args`` supplies it for the launch device. + ex2_emulated_pairs: int = 4 + # Reduce streamed score-fragment maxima inside the TMEM load + # (tcgen05.ld.red); ``arch_config_args`` supplies it. + uses_ldtm_stat: bool = False use_cluster_smem_reduction: bool = False use_separate_reduction_kernel: bool = False # Compile-time attention-mask selection. Public APIs normalize the string @@ -1953,6 +1978,9 @@ def validate_sage_profile(self) -> None: "Sage attention requires a two-instance Keeps profile: " "Q64/KV256 or Q128/KV128" ) + if not self.groups_tokens_heads_q: + # The prepared ``sfQ`` rows follow the grouped Q tile rows. + raise ValueError("Sage attention requires groups_tokens_heads_q=True") if self.use_paged_kv or self.headdim != 128: raise ValueError("Sage attention requires contiguous K/V with headdim=128") # Only the direct output store applies the V channel scales and means; @@ -2410,19 +2438,49 @@ def defers_softmax_anchor_updates(self) -> bool: Keeps correction skips the in-place O rescale whenever the anchor is unchanged, so keeping the prior anchor within - ``SOFTMAX_RESCALE_THRESHOLD_LOG2`` trades a bounded 16-bit P range - (2**8) for fewer TMEM rescales. The profiles listed here are the ones - where that trade was measured to pay: KV256 tiles and block-sparse - routes, whose row maximum moves often but rarely by much. FP8 P uses - the static 448 scale, which needs ``p <= 1``, so it always anchors on + ``softmax_anchor_headroom_log2`` binades trades a bounded P range for + fewer TMEM rescales. The profiles listed here are the ones where that + trade was measured to pay: KV256 tiles and block-sparse routes, whose + row maximum moves often but rarely by much. A zero headroom anchors on the exact row maximum. """ return ( self.use_keeps_mma_ab - and not self.use_fp8_pv and (self.tile_size_kv == 256 or self.use_block_sparse) + and self.softmax_anchor_headroom_log2 > 0 ) + @property + def softmax_anchor_headroom_log2(self) -> float: + """Binades a deferred exponent anchor may lag the row maximum. + + A probability is ``2**(s - anchor)`` times its scale with ``s - anchor`` + at most this headroom. ``softmax_rescale_threshold_log2`` sets it. + Otherwise 16-bit P allows ``SOFTMAX_RESCALE_THRESHOLD_LOG2``, far + inside its range, and FP8 P keeps exact anchors: FP8 headroom lowers + the E4M3 scale by the same number of binades (``fp8_p_quant_scale``). + """ + if self.softmax_rescale_threshold_log2 is not None: + return self.softmax_rescale_threshold_log2 + return 0.0 if self.use_fp8_pv else SOFTMAX_RESCALE_THRESHOLD_LOG2 + + @property + def fp8_p_quant_scale(self) -> float: + """Scale ``c`` of FP8 probabilities, which are quantized to E4M3 as ``c * p``. + + A deferred anchor lets ``p`` reach ``2**softmax_anchor_headroom_log2``, + so ``c`` is the E4M3 maximum divided by that bound; an exact anchor + keeps ``p <= 1`` and uses the whole range. + """ + if self.defers_softmax_anchor_updates: + return E4M3_MAX * 2.0**-self.softmax_anchor_headroom_log2 + return E4M3_MAX + + @property + def fp8_p_quant_log2_scale(self) -> float: + """``log2(fp8_p_quant_scale)``, the exponent addend of FP8 probabilities.""" + return math.log2(self.fp8_p_quant_scale) + @property def use_sage_attention(self) -> bool: """Whether Q/K scores carry per-block scales and V per-channel scales.""" @@ -2437,61 +2495,87 @@ def softmax_num_warps(self, inst_id: int) -> int: """Return the warps of one softmax instance.""" return self.softmax0_num_warps if inst_id == 0 else self.softmax1_num_warps - def sage_k_groups_for_block(self, k_block_size: int) -> int: - """Return the K scale groups inside one streamed K32 score fragment.""" - return max(1, self.softmax_score_fragment_regs // k_block_size) + def sage_k_scale_block_size(self, proxy: bool = False) -> int: + """Return the K block size of the ``sfK`` words one route kind reads. + + A plan's route kinds read one geometry wherever both read words in + the P pass: the finer of the K and summary block sizes, and the + prepared scales repeat the coarser kind's (``_SageScaleLayout``). A + P loop per kind would cost more than the repeated words: its second + copy of the loop spills the softmax warps' registers. One-token + summary scales are the exception: the max pass dequantizes a proxy + route's scores, so the P pass reads exact words only and each kind + keeps its own block size. + """ + if not self.use_block_sparse_proxy_routes: + return self.sage_k_block_size + if self.sage_k_summary_block_size == 1: + return 1 if proxy else self.sage_k_block_size + return min(self.sage_k_block_size, self.sage_k_summary_block_size) @property - def sage_k_groups_per_fragment(self) -> int: - """Return the scale groups per fragment of an exact route or dense tile.""" - if not self.use_sage_attention: - return 1 - return self.sage_k_groups_for_block(self.sage_k_block_size) + def sage_mixed_k_geometry(self) -> bool: + """Whether exact and proxy routes read different ``sfK`` geometries.""" + return self.use_sage_attention and ( + self.sage_k_scale_block_size(proxy=True) != self.sage_k_scale_block_size() + ) - @property - def sage_summary_k_groups_per_fragment(self) -> int: - """Return the scale groups per fragment of a proxy route's summaries.""" - if not self.use_sage_attention: - return 1 - return self.sage_k_groups_for_block(self.sage_k_summary_block_size) + def sage_scale_layout(self, seq_len_kv: int) -> _SageScaleLayout: + """Return the layout of the Sage scales a plan prepares before attention. - def sage_k_groups_per_fragment_for(self, proxy: bool) -> int: - """Return the scale groups per fragment of one route kind.""" - if proxy: - return self.sage_summary_k_groups_per_fragment - return self.sage_k_groups_per_fragment + Its K atoms are the plan's Keeps atoms, 64 K tokens of two K32 score + fragments, and its Q tiles the plan's grouped Q CTAs, so the softmax + reads a work tile's rows at ``q_group_idx * tile_size_q``. Every Sage + profile streams such fragments (``validate_sage_profile``). + """ + assert self.use_sage_attention and self.groups_tokens_heads_q + assert self.keeps_atom_tokens == _SAGE_ATOM_TOKENS + assert self.keeps_fragments_per_atom == _SAGE_ATOM_MIN_WORDS + num_summaries = 0 + if self.use_block_sparse_proxy_routes: + num_summaries, _ = _block_sparse_proxy_summary_geometry( + seq_len_kv, self.kv_block_size + ) + layout = _SageScaleLayout( + seq_len_kv=seq_len_kv, + k_block_size=self.sage_k_block_size, + num_summaries=num_summaries, + k_summary_block_size=self.sage_k_summary_block_size, + k_tokens_block_size=self.sage_k_scale_block_size(), + summaries_block_size=self.sage_k_scale_block_size(proxy=True), + seq_len_q=self.max_seq_len_q, + q_block_size=self.sage_q_block_size, + heads_q_per_kv=self.heads_q_per_kv, + tile_size_q=self.tile_size_q, + headdim=self.headdim, + v_mean=self.sage_v_mean, + ) + q_rows = layout.q_rows + assert q_rows.tokens_per_tile == self.q_tokens_per_cta + assert q_rows.tiles == self.num_q_ctas + return layout - @property - def sage_mixed_k_geometry(self) -> bool: - """Whether exact and proxy routes use different scale-group geometries.""" - return ( - self.use_sage_attention - and self.sage_summary_k_groups_per_fragment - != self.sage_k_groups_per_fragment + def sage_k_groups_per_fragment(self, proxy: bool = False) -> int: + """Return the ``sfK`` groups of one streamed K32 score fragment of a route kind.""" + if not self.use_sage_attention: + return 1 + return max( + 1, self.softmax_score_fragment_regs // self.sage_k_scale_block_size(proxy) ) - def sage_scores_dequantized_for(self, proxy: bool) -> bool: - """Whether the max pass writes one route kind's scores back dequantized. + def sage_scores_dequantized(self, proxy: bool = False) -> bool: + """Whether the max pass writes a route kind's scores back dequantized. - With one scale per score (the one-token K block), the P pass then - reads neither ``sfK`` nor the INT32 bias. Coarser geometries keep the - scores quantized to avoid a TMEM round trip per tile. + A kind with one scale per score (the one-token K block) does, so the P + pass reads neither its ``sfK`` nor the INT32 bias; other geometries + stay quantized to avoid a TMEM round trip per tile. """ return ( self.use_sage_attention - and self.sage_k_groups_per_fragment_for(proxy) + and self.sage_k_groups_per_fragment(proxy) == self.softmax_score_fragment_regs ) - def sage_k_scales_in_smem_for(self, groups: int) -> bool: - """Whether a tile with ``groups`` scale groups per fragment keeps ``sfK`` in SMEM. - - Four or more groups per K32 fragment (K blocks of 4 or 1 token) would - need 16 or more registers per lane, so those tiles read ``sfK`` from - SMEM; larger blocks keep a rotating register array. - """ - return self.use_sage_attention and groups >= 4 - @property def matches_kv256_task_topology(self) -> bool: """Whether task roles match KV256's validated 16-warp layout.""" @@ -3381,7 +3465,7 @@ def _active_warp_roles(cfg: FmhaDecodeConfig) -> list[_WarpRole]: ) ) - if cfg.use_paged_kv and not cfg.use_block_sparse: + if cfg.uses_page_offsets_task: roles.append( _WarpRole( "page_offsets", @@ -4637,7 +4721,7 @@ def _validate_profile_support( _validate_kv256_static_config(cfg) if cfg.mask_type not in (DENSE, CAUSAL): raise ValueError("mask_type must be DENSE or CAUSAL") - if cfg.use_paged_kv and not cfg.use_block_sparse: + if cfg.uses_page_offsets_task: cfg.validate_paged_kv_staging_config() if num_heads_q % num_heads_kv != 0: raise ValueError("fmha_decode requires num_heads_q divisible by num_heads_kv") @@ -4647,6 +4731,12 @@ def _validate_profile_support( "heads_q_per_kv metadata must equal num_heads_q / num_heads_kv" ) cfg.validate_sage_profile() + threshold = cfg.softmax_rescale_threshold_log2 + if threshold is not None and not 0.0 <= threshold <= SOFTMAX_RESCALE_THRESHOLD_LOG2: + raise ValueError( + "softmax_rescale_threshold_log2 must lie in " + f"[0, {SOFTMAX_RESCALE_THRESHOLD_LOG2:g}], got {threshold}" + ) cfg.validate_block_sparse_profile(heads_q_per_kv=heads_q_per_kv) if use_groups_tokens_heads_q: make_q_tile_geometry( @@ -5043,6 +5133,25 @@ def validate_storage_page_size( ) +def arch_config_args(compute_capability: tuple[int, int]) -> dict[str, object]: + """Return the config fields whose value depends on the GPU architecture. + + SM100 runs a quarter of the streamed exponentials as FMA polynomials + because the MUFU issue rate bounds its P pass while the FMA pipe is nearly + idle; larger shares grow the fragment body until the softmax warps become + instruction-fetch bound. SM103 doubles the MUFU ex2 rate (32 results per + SM clock), so the polynomials only add instructions there and every + exponential stays on MUFU. GPUs whose TMEM loads return score maxima + (LDTM.STAT) take them from the load, which replaces the max pass's + register reduction; the context kernel shares that gate. + """ + is_sm103 = tuple(compute_capability) == (10, 3) + return { + "ex2_emulated_pairs": 0 if is_sm103 else 4, + "uses_ldtm_stat": ldtm_stat_supported(compute_capability), + } + + def make_decode_config( headdim: int = 128, args: object | None = None, 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 843b7ffb44b0..909de7947f81 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 @@ -18,8 +18,6 @@ reduction code can use named values without duplicating comments. """ -import math - # B200 has 148 SMs. Use this only when the runtime SM query is unavailable, # so auto split-KV selection remains deterministic in offline/test flows. FALLBACK_SM_COUNT_B200 = 148 @@ -60,7 +58,9 @@ # FlashInfer/TRT-LLM policy, this assumes normal model logits rather than # adversarial values outside the qualified probability bound. Streamed KV256 # and block-sparse Keeps profiles apply it; see -# ``FmhaDecodeConfig.defers_softmax_anchor_updates``. +# ``FmhaDecodeConfig.defers_softmax_anchor_updates``. It bounds +# ``FmhaDecodeConfig.softmax_rescale_threshold_log2``, is its 16-bit default, +# and is the block-sparse plans' default for every P dtype. SOFTMAX_RESCALE_THRESHOLD_LOG2 = 8.0 # A launch bound makes ptxas honor warpgroup ``setmaxnreg`` allocations, but @@ -156,10 +156,11 @@ # one-byte or 16 two-byte K elements. MMA_K_STEP_BYTES = 32 -# FP8 probabilities are quantized as 448 * p (the E4M3 maximum); row sums and -# sink terms share the scale, and the output normalization divides it out. -FP8_P_QUANT_SCALE = 448.0 -FP8_P_QUANT_LOG2_SCALE = math.log2(FP8_P_QUANT_SCALE) +# FP8 probabilities are quantized as ``p * FmhaDecodeConfig.fp8_p_quant_scale``: +# the E4M3 maximum, less the binades a deferred exponent anchor may lag the +# row maximum (``FmhaDecodeConfig.softmax_anchor_headroom_log2``). Row sums +# and sink terms share the scale, and the output normalization divides it out. +E4M3_MAX = 448.0 # tcgen05 SMEM descriptor geometry: address offsets count 16-byte units and a # 128-byte swizzle atom spans one 128-byte row per K or MN index. 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 6e03c2a0c12b..2292eb6544bb 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 @@ -101,9 +101,7 @@ ) from .fmha_decode_resources.helpers_kv_tile_idx import _runtime_active_splits_kv from .fmha_decode_resources.sage_scales import ( - SAGE_K_SCALES_RING_STAGES, SageKScalesResource, - SageScaleTensors, SageVScalesResource, ) from .fmha_decode_tasks import ( @@ -371,14 +369,7 @@ def _build_decode_gen_schedule( sparse_route_metadata: cute.Pointer | None = None, sparse_row_route_begin: Int32 | None = None, sparse_route_count: Int32 | None = None, - sage_q_scale_ptr: cute.Pointer | None = None, - sage_k_scale_ptr: cute.Pointer | None = None, - sage_k_summary_scale_ptr: cute.Pointer | None = None, - sage_v_scale_ptr: cute.Pointer | None = None, - sage_v_mean_ptr: cute.Pointer | None = None, - sage_q_scale_head_stride: Int32 | None = None, - sage_k_scale_head_stride: Int32 | None = None, - sage_k_summary_scale_head_stride: Int32 | None = None, + sage_scales_ptr: cute.Pointer | None = None, ) -> tuple[ list[Task], dict[MemoryResource, list[MemoryResource]], @@ -533,9 +524,8 @@ def _build_decode_gen_schedule( # Leave barrier_ptr unset so SmemAllocator packs every pipeline barrier # into the unified block. Separate barrier arrays create an alignment gap # before the 1024-byte-aligned data block and overflow near-capacity Q128. - use_paged_kv = cfg.use_paged_kv use_one_inst_kv = cfg.num_insts_kv == 1 - use_dense_page_offsets = use_paged_kv and not cfg.use_block_sparse + uses_page_offsets_task = cfg.uses_page_offsets_task use_one_inst_qkv = cfg.use_keeps_mma_ab and cfg.num_insts_kv == 1 # K and V of different byte widths take the split-resource paths with one # K ring and one V ring, each shared by both K/V instances; equal widths @@ -636,7 +626,7 @@ def _build_decode_gen_schedule( # A held route is published once and retained through the final V issue; # only schedules that replace locators per tile need separate K/V state. use_separate_kv_page_offset_resources = ( - use_dense_page_offsets + uses_page_offsets_task and use_per_inst_kv_resources and not use_one_inst_qkv and not ( @@ -726,7 +716,7 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: smem_page_offsets_cfg = None smem_page_offsets_v_cfg = None - if use_dense_page_offsets: + if uses_page_offsets_task: page_offsets_stages = ( 3 if use_separate_kv_page_offset_resources else cfg.page_offsets_stages ) @@ -740,11 +730,25 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: smem_page_offsets_v_cfg = _make_page_offsets_cfg(page_offsets_stages) sparse_softmax_metadata0_cfg = None sparse_softmax_metadata1_cfg = None + # Load publishes a block-sparse route, and a Sage plan's K scales into the + # instance's ``sfK`` ring, right after the K tile. Softmax releases the + # route when it starts that tile and the scales after the tile's last pass + # that reads them. When the K/V ring frees a slot for an instance's next K + # tile, the oldest buffered K tile's QK has run, after the previous tile's + # P left the TMEM columns it aliases, so the instance's Softmax has + # finished every tile older than the K tiles the ring buffers per + # instance; one stage more than those K tiles never stalls the load warp + # on either pipeline. + if use_per_inst_kv_resources: + inst_k_tiles = max(split_k0_stages, split_k1_stages) + else: + # One ring cycles K and V tiles of every instance. + inst_k_tiles = math.ceil(cfg.kv_stages / (2 * cfg.num_insts_kv)) + k_tile_stages = inst_k_tiles + 1 if cfg.use_block_sparse: - # Two stages are sufficient for the split-ring cadence: one route can - # await Softmax while Load publishes the next route for the same inst. + route_metadata_stages = k_tile_stages sparse_softmax_metadata0_cfg = PipelineConfig( - num_stages=2, + num_stages=route_metadata_stages, num_bytes=0, producer_group=load_grp, consumer_group=softmax0_grp, @@ -753,7 +757,7 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: advance_on_wait=True, ) sparse_softmax_metadata1_cfg = PipelineConfig( - num_stages=2, + num_stages=route_metadata_stages, num_bytes=0, producer_group=load_grp, consumer_group=softmax1_grp, @@ -954,7 +958,7 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: kv_seqlens = seqlens_kv if use_runtime_seqlens_kv else None smem_page_offsets = None smem_page_offsets_v = None - if use_dense_page_offsets: + if uses_page_offsets_task: smem_page_offsets = SmemPageOffsetsKvResource( pipeline_config=smem_page_offsets_cfg, cfg=cfg, @@ -1002,18 +1006,6 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: sparse_kv_metadata1 = None sparse_softmax_metadata0 = None sparse_softmax_metadata1 = None - sage_scale_tensors = None - if cfg.use_sage_attention: - sage_scale_tensors = SageScaleTensors( - q_scale_ptr=sage_q_scale_ptr, - q_scale_head_stride=sage_q_scale_head_stride, - k_scale_ptr=sage_k_scale_ptr, - k_scale_head_stride=sage_k_scale_head_stride, - k_summary_scale_ptr=sage_k_summary_scale_ptr, - k_summary_scale_head_stride=sage_k_summary_scale_head_stride, - v_scale_ptr=sage_v_scale_ptr, - v_mean_ptr=sage_v_mean_ptr, - ) if cfg.use_block_sparse: # This selects the prepared-record storage ABI. Causal consumers still # intersect these column-validity words with each Q row's causal mask. @@ -1050,9 +1042,6 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: route_metadata=sparse_route_metadata, route_layout=prepared_route_layout, name="smemBlockSparseSoftmaxMetadata0", - h_k_idx=h_k_idx, - b_idx=b_idx, - scale_tensors=sage_scale_tensors, ) sparse_softmax_metadata1 = SmemBlockSparseSoftmaxMetadataResource( pipeline_config=sparse_softmax_metadata1_cfg, @@ -1061,9 +1050,6 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: route_metadata=sparse_route_metadata, route_layout=prepared_route_layout, name="smemBlockSparseSoftmaxMetadata1", - h_k_idx=h_k_idx, - b_idx=b_idx, - scale_tensors=sage_scale_tensors, ) smem_kv = None smem_k0 = None @@ -1238,85 +1224,35 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: advance_on_wait=True, ), cfg=cfg, - scale_tensors=sage_scale_tensors, + sage_scales_ptr=sage_scales_ptr, h_k_idx=h_k_idx, b_idx=b_idx, + num_heads_kv=num_heads_kv, name="sageVScales", ) - # Each softmax instance fills and reads its own ``sfK`` words; only the - # SMEM form carries a pipeline. A mixed-geometry plan adds a second - # resource for the proxy summaries. + # Each softmax instance reads its own ``sfK`` words from a ring that the + # load warp fills after each K tile, through ``cp.async`` copies that the + # ``AsyncLoad`` commit ties to the stage: the tile's atoms of the prepared + # scales, a block-sparse route's or a dense tile's. sage_k_scales0 = None sage_k_scales1 = None - sage_summary_k_scales0 = None - sage_summary_k_scales1 = None if cfg.use_sage_attention: - def _sage_k_scales_cfg(groups: int, softmax_grp) -> PipelineConfig | None: - if not cfg.sage_k_scales_in_smem_for(groups): - return None - return PipelineConfig( - num_stages=SAGE_K_SCALES_RING_STAGES, - num_bytes=0, - producer_group=softmax_grp, - consumer_group=softmax_grp, - pipeline_type=PipelineType.AsyncAsync, - cta_layout_vmnk=cta_layout, - advance_on_wait=True, - ) - - token_groups = cfg.sage_k_groups_per_fragment - sage_k_scales0 = SageKScalesResource( - pipeline_config=_sage_k_scales_cfg(token_groups, softmax0_grp), - inst_id=0, - route_metadata=sparse_softmax_metadata0, - name="sageKScales0", - cfg=cfg, - seqlens_kv=kv_seqlens, - max_seq_len_kv=max_seq_len_kv, - q_group_idx=q_group_idx, - seq_len_q=seq_len_q, - h_k_idx=h_k_idx, - b_idx=b_idx, - scale_tensors=sage_scale_tensors, - ) - sage_k_scales1 = SageKScalesResource( - pipeline_config=_sage_k_scales_cfg(token_groups, softmax1_grp), - inst_id=1, - route_metadata=sparse_softmax_metadata1, - name="sageKScales1", - cfg=cfg, - seqlens_kv=kv_seqlens, - max_seq_len_kv=max_seq_len_kv, - q_group_idx=q_group_idx, - seq_len_q=seq_len_q, - h_k_idx=h_k_idx, - b_idx=b_idx, - scale_tensors=sage_scale_tensors, - ) - if cfg.sage_mixed_k_geometry: - summary_groups = cfg.sage_summary_k_groups_per_fragment - sage_summary_k_scales0 = SageKScalesResource( - pipeline_config=_sage_k_scales_cfg(summary_groups, softmax0_grp), - inst_id=0, - summary=True, - route_metadata=sparse_softmax_metadata0, - name="sageSummaryKScales0", - cfg=cfg, - seqlens_kv=kv_seqlens, - max_seq_len_kv=max_seq_len_kv, - q_group_idx=q_group_idx, - seq_len_q=seq_len_q, - h_k_idx=h_k_idx, - b_idx=b_idx, - scale_tensors=sage_scale_tensors, - ) - sage_summary_k_scales1 = SageKScalesResource( - pipeline_config=_sage_k_scales_cfg(summary_groups, softmax1_grp), - inst_id=1, - summary=True, - route_metadata=sparse_softmax_metadata1, - name="sageSummaryKScales1", + def _sage_k_scales(inst_id: int, softmax_grp) -> SageKScalesResource: + return SageKScalesResource( + pipeline_config=PipelineConfig( + num_stages=k_tile_stages, + num_bytes=0, + producer_group=load_grp, + consumer_group=softmax_grp, + pipeline_type=PipelineType.AsyncAsync, + cta_layout_vmnk=cta_layout, + advance_on_wait=True, + async_producer_op=pipeline.PipelineOp.AsyncLoad, + ), + inst_id=inst_id, + route_layout=prepared_route_layout if cfg.use_block_sparse else None, + name=f"sageKScales{inst_id}", cfg=cfg, seqlens_kv=kv_seqlens, max_seq_len_kv=max_seq_len_kv, @@ -1324,9 +1260,13 @@ def _sage_k_scales_cfg(groups: int, softmax_grp) -> PipelineConfig | None: seq_len_q=seq_len_q, h_k_idx=h_k_idx, b_idx=b_idx, - scale_tensors=sage_scale_tensors, + num_heads_kv=num_heads_kv, + sage_scales_ptr=sage_scales_ptr, ) + sage_k_scales0 = _sage_k_scales(0, softmax0_grp) + sage_k_scales1 = _sage_k_scales(1, softmax1_grp) + tmem_s0 = TmemSResource( inst_id=0, pipeline_config=tmem_s0_cfg, @@ -1339,11 +1279,10 @@ def _sage_k_scales_cfg(groups: int, softmax_grp) -> PipelineConfig | None: seq_len_q=seq_len_q, h_k_idx=h_k_idx, b_idx=b_idx, + num_heads_kv=num_heads_kv, sync_barrier_id=0, - sage_k_scales=sage_k_scales0, - sage_summary_k_scales=sage_summary_k_scales0 or sage_k_scales0, name="tmemS0", - scale_tensors=sage_scale_tensors, + sage_scales_ptr=sage_scales_ptr, ) tmem_s1 = None if not use_one_inst_kv: @@ -1359,12 +1298,11 @@ def _sage_k_scales_cfg(groups: int, softmax_grp) -> PipelineConfig | None: seq_len_q=seq_len_q, h_k_idx=h_k_idx, b_idx=b_idx, + num_heads_kv=num_heads_kv, sync_barrier_id=1, score_seed_owner=tmem_s0, - sage_k_scales=sage_k_scales1, - sage_summary_k_scales=sage_summary_k_scales1 or sage_k_scales1, name="tmemS1", - scale_tensors=sage_scale_tensors, + sage_scales_ptr=sage_scales_ptr, ) # Packed persistent QK derives the descriptor from Q's just-waited # consumer stage, avoiding a routed HEAD-to-LOOP descriptor value across @@ -1385,8 +1323,6 @@ def _sage_k_scales_cfg(groups: int, softmax_grp) -> PipelineConfig | None: cfg=cfg, scale_softmax_log2=scale_softmax_log2, use_variable_seqlens_kv=use_runtime_seqlens_kv, - sage_k_scales=sage_k_scales0, - sage_summary_k_scales=sage_summary_k_scales0 or sage_k_scales0, name="smemP0", ) smem_p1 = None @@ -1397,8 +1333,6 @@ def _sage_k_scales_cfg(groups: int, softmax_grp) -> PipelineConfig | None: cfg=cfg, scale_softmax_log2=scale_softmax_log2, use_variable_seqlens_kv=use_runtime_seqlens_kv, - sage_k_scales=sage_k_scales1, - sage_summary_k_scales=sage_summary_k_scales1 or sage_k_scales1, name="smemP1", ) @@ -1492,7 +1426,7 @@ def _sage_k_scales_cfg(groups: int, softmax_grp) -> PipelineConfig | None: active_splits_kv=active_splits_kv, static_full_split_prefix=static_full_split_prefix, name="tmemCorr0", - sage_v_scales=sage_v_scales, + sage_v_scales_ref=sage_v_scales, ) tmem_corr0.tmem_o_ref = tmem_o tmem_corr0.softmax_local0_ref = tmem_softmax_local0 @@ -1521,7 +1455,7 @@ def _sage_k_scales_cfg(groups: int, softmax_grp) -> PipelineConfig | None: active_splits_kv=active_splits_kv, static_full_split_prefix=static_full_split_prefix, name="tmemCorr1", - sage_v_scales=sage_v_scales, + sage_v_scales_ref=sage_v_scales, ) tmem_corr1.smem_p0_ref = smem_p0 tmem_corr1.smem_p1_ref = smem_p1 @@ -1565,16 +1499,14 @@ def _sage_k_scales_cfg(groups: int, softmax_grp) -> PipelineConfig | None: smem_resources.append(smem_p0) if not use_one_inst_kv: smem_resources.append(smem_p1) - # Each instance's scale rings follow its S resource in the layout. + # An instance's ``sfK`` ring follows its S resource in the layout. smem_resources.append(tmem_s0) - smem_resources.extend( - r for r in (sage_k_scales0, sage_summary_k_scales0) if r is not None - ) + if sage_k_scales0 is not None: + smem_resources.append(sage_k_scales0) if not use_one_inst_kv: smem_resources.append(tmem_s1) - smem_resources.extend( - r for r in (sage_k_scales1, sage_summary_k_scales1) if r is not None - ) + if sage_k_scales1 is not None: + smem_resources.append(sage_k_scales1) smem_resources.append(tmem_o) smem_resources.append(tmem_softmax_local0) if not use_one_inst_kv: @@ -1682,6 +1614,12 @@ def smem_bytes(allocator: SmemAllocator) -> int: "sparse_route_count": sparse_route_count, "num_heads_kv": num_heads_kv, } + # A Sage plan's ``sfK`` rings follow the K tiles. + sage_k_scale_rings = ( + {"sage_k_scales0": sage_k_scales0, "sage_k_scales1": sage_k_scales1} + if cfg.use_sage_attention + else {} + ) if use_one_inst_qkv: load_tasks = ( create_load_task_one_inst_qkv( @@ -1718,6 +1656,7 @@ def smem_bytes(allocator: SmemAllocator) -> int: sparse_softmax_metadata1=sparse_softmax_metadata1, domain_bias=0, warp_indices=load_warp_indices, + **sage_k_scale_rings, **task_runtime_kwargs, ) else: @@ -1740,6 +1679,7 @@ def smem_bytes(allocator: SmemAllocator) -> int: sparse_softmax_metadata1=sparse_softmax_metadata1, domain_bias=0, warp_idx=cfg.clc_load_warp_idx if use_clc_dynamic else None, + **sage_k_scale_rings, **task_runtime_kwargs, ), ) @@ -1759,11 +1699,12 @@ def smem_bytes(allocator: SmemAllocator) -> int: sparse_kv_metadata1=sparse_kv_metadata1, sparse_softmax_metadata0=sparse_softmax_metadata0, sparse_softmax_metadata1=sparse_softmax_metadata1, + **sage_k_scale_rings, **task_runtime_kwargs, ), ) page_offsets_task = None - if use_dense_page_offsets: + if uses_page_offsets_task: page_offsets_warp_idx = cfg.page_offsets_warp_idx if use_separate_kv_page_offset_resources: page_offsets_task = create_page_offsets_task_split_kv( @@ -1879,7 +1820,6 @@ def smem_bytes(allocator: SmemAllocator) -> int: domain_bias=1, membership_lifetime=membership_lifetime, sage_k_scales=sage_k_scales0, - sage_summary_k_scales=sage_summary_k_scales0, **task_runtime_kwargs, ) softmax1_task = None @@ -1897,7 +1837,6 @@ def smem_bytes(allocator: SmemAllocator) -> int: domain_bias=1, membership_lifetime=membership_lifetime, sage_k_scales=sage_k_scales1, - sage_summary_k_scales=sage_summary_k_scales1, **task_runtime_kwargs, ) if use_one_inst_kv: @@ -2145,32 +2084,17 @@ def smem_bytes(allocator: SmemAllocator) -> int: assert tmem_softmax_local1 is not None resource_dependency_graph[tmem_s1].append(tmem_stats_done1) resource_dependency_graph[tmem_stats_done1] = [tmem_softmax_local1] - # A softmax instance's outputs read its K scale words; an SMEM-form - # resource fills them from the route's staged metadata, while a - # register-form one hands them straight to the outputs. The correction + # A softmax instance's outputs read its K scale words. The correction # epilogue reads the V channel scales. - for scales_pair, route_metadata, outputs in ( - ( - (sage_k_scales0, sage_summary_k_scales0), - sparse_softmax_metadata0, - (smem_p0, tmem_softmax_local0, tmem_softmax_global0), - ), - ( - (sage_k_scales1, sage_summary_k_scales1), - sparse_softmax_metadata1, - (smem_p1, tmem_softmax_local1, tmem_softmax_global1), - ), + for scales, outputs in ( + (sage_k_scales0, (smem_p0, tmem_softmax_local0, tmem_softmax_global0)), + (sage_k_scales1, (smem_p1, tmem_softmax_local1, tmem_softmax_global1)), ): - for scales in scales_pair: - if scales is None: - continue - resource_dependency_graph[scales] = ( - [route_metadata] - if route_metadata is not None and scales.in_smem - else [] - ) - for resource in outputs: - resource_dependency_graph[resource].append(scales) + if scales is None: + continue + resource_dependency_graph[scales] = [] + for resource in outputs: + resource_dependency_graph[resource].append(scales) if sage_v_scales is not None: resource_dependency_graph[sage_v_scales] = [] for resource in (tmem_corr0, tmem_corr1): @@ -2655,14 +2579,7 @@ def _run_decode_gen_active( tma_desc_v_summary: cutlass.GridConstant[cuda.TensorMap] | None = None, tma_desc_k_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, tma_desc_v_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, - g_sage_q_scale: cute.Pointer | None = None, - g_sage_k_scale: cute.Pointer | None = None, - g_sage_k_summary_scale: cute.Pointer | None = None, - g_sage_v_scale: cute.Pointer | None = None, - g_sage_v_mean: cute.Pointer | None = None, - g_sage_q_scale_head_stride: Int32 | None = None, - g_sage_k_scale_head_stride: Int32 | None = None, - g_sage_k_summary_scale_head_stride: Int32 | None = None, + g_sage_scales: cute.Pointer | None = None, ) -> None: """Run the complete decode body for one runtime-valid Q tile. @@ -2800,14 +2717,7 @@ def _run_decode_gen_active( tma_desc_v_summary=tma_desc_v_summary_ptr, tma_desc_k_summary_atom=tma_desc_k_summary_atom_ptr, tma_desc_v_summary_atom=tma_desc_v_summary_atom_ptr, - sage_q_scale_ptr=g_sage_q_scale, - sage_k_scale_ptr=g_sage_k_scale, - sage_k_summary_scale_ptr=g_sage_k_summary_scale, - sage_v_scale_ptr=g_sage_v_scale, - sage_v_mean_ptr=g_sage_v_mean, - sage_q_scale_head_stride=g_sage_q_scale_head_stride, - sage_k_scale_head_stride=g_sage_k_scale_head_stride, - sage_k_summary_scale_head_stride=g_sage_k_summary_scale_head_stride, + sage_scales_ptr=g_sage_scales, page_idx_kv=g_page_idx_kv, page_table_stride=g_page_table_stride, page_table_capacity=g_page_table_capacity, @@ -3010,14 +2920,7 @@ def _run_decode_gen_runtime_prefix( tma_desc_v_summary: cutlass.GridConstant[cuda.TensorMap] | None = None, tma_desc_k_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, tma_desc_v_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, - g_sage_q_scale: cute.Pointer | None = None, - g_sage_k_scale: cute.Pointer | None = None, - g_sage_k_summary_scale: cute.Pointer | None = None, - g_sage_v_scale: cute.Pointer | None = None, - g_sage_v_mean: cute.Pointer | None = None, - g_sage_q_scale_head_stride: Int32 | None = None, - g_sage_k_scale_head_stride: Int32 | None = None, - g_sage_k_summary_scale_head_stride: Int32 | None = None, + g_sage_scales: cute.Pointer | None = None, ) -> None: """Run the general runtime split-prefix producer or retire its suffix.""" @@ -3102,14 +3005,7 @@ def _run_decode_gen_runtime_prefix( tma_desc_v_summary=tma_desc_v_summary, tma_desc_k_summary_atom=tma_desc_k_summary_atom, tma_desc_v_summary_atom=tma_desc_v_summary_atom, - g_sage_q_scale=g_sage_q_scale, - g_sage_k_scale=g_sage_k_scale, - g_sage_k_summary_scale=g_sage_k_summary_scale, - g_sage_v_scale=g_sage_v_scale, - g_sage_v_mean=g_sage_v_mean, - g_sage_q_scale_head_stride=g_sage_q_scale_head_stride, - g_sage_k_scale_head_stride=g_sage_k_scale_head_stride, - g_sage_k_summary_scale_head_stride=g_sage_k_summary_scale_head_stride, + g_sage_scales=g_sage_scales, ) else: _run_decode_gen_inactive_cluster_rank() @@ -3163,14 +3059,7 @@ def _run_decode_gen_runtime_prefix( tma_desc_v_summary=tma_desc_v_summary, tma_desc_k_summary_atom=tma_desc_k_summary_atom, tma_desc_v_summary_atom=tma_desc_v_summary_atom, - g_sage_q_scale=g_sage_q_scale, - g_sage_k_scale=g_sage_k_scale, - g_sage_k_summary_scale=g_sage_k_summary_scale, - g_sage_v_scale=g_sage_v_scale, - g_sage_v_mean=g_sage_v_mean, - g_sage_q_scale_head_stride=g_sage_q_scale_head_stride, - g_sage_k_scale_head_stride=g_sage_k_scale_head_stride, - g_sage_k_summary_scale_head_stride=g_sage_k_summary_scale_head_stride, + g_sage_scales=g_sage_scales, ) else: _signal_padded_pdl_producer(cfg) @@ -3218,14 +3107,7 @@ def decode_gen_kernel( tma_desc_v_summary: cutlass.GridConstant[cuda.TensorMap] | None = None, tma_desc_k_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, tma_desc_v_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, - g_sage_q_scale: cute.Pointer | None = None, - g_sage_k_scale: cute.Pointer | None = None, - g_sage_k_summary_scale: cute.Pointer | None = None, - g_sage_v_scale: cute.Pointer | None = None, - g_sage_v_mean: cute.Pointer | None = None, - g_sage_q_scale_head_stride: Int32 | None = None, - g_sage_k_scale_head_stride: Int32 | None = None, - g_sage_k_summary_scale_head_stride: Int32 | None = None, + g_sage_scales: cute.Pointer | None = None, ) -> None: """Dispatch one static Q/split tile and drain padded launch slots safely.""" q_group_cta_idx, h_k_idx, b_idx = cute.arch.block_idx() @@ -3365,14 +3247,7 @@ def decode_gen_kernel( tma_desc_v_summary=tma_desc_v_summary, tma_desc_k_summary_atom=tma_desc_k_summary_atom, tma_desc_v_summary_atom=tma_desc_v_summary_atom, - g_sage_q_scale=g_sage_q_scale, - g_sage_k_scale=g_sage_k_scale, - g_sage_k_summary_scale=g_sage_k_summary_scale, - g_sage_v_scale=g_sage_v_scale, - g_sage_v_mean=g_sage_v_mean, - g_sage_q_scale_head_stride=g_sage_q_scale_head_stride, - g_sage_k_scale_head_stride=g_sage_k_scale_head_stride, - g_sage_k_summary_scale_head_stride=g_sage_k_summary_scale_head_stride, + g_sage_scales=g_sage_scales, ) else: _run_decode_gen_runtime_prefix( @@ -3422,14 +3297,7 @@ def decode_gen_kernel( tma_desc_v_summary=tma_desc_v_summary, tma_desc_k_summary_atom=tma_desc_k_summary_atom, tma_desc_v_summary_atom=tma_desc_v_summary_atom, - g_sage_q_scale=g_sage_q_scale, - g_sage_k_scale=g_sage_k_scale, - g_sage_k_summary_scale=g_sage_k_summary_scale, - g_sage_v_scale=g_sage_v_scale, - g_sage_v_mean=g_sage_v_mean, - g_sage_q_scale_head_stride=g_sage_q_scale_head_stride, - g_sage_k_scale_head_stride=g_sage_k_scale_head_stride, - g_sage_k_summary_scale_head_stride=g_sage_k_summary_scale_head_stride, + g_sage_scales=g_sage_scales, ) else: # Packed-Q grids use a batch-wide maximum envelope. These Q CTAs own no @@ -3476,16 +3344,12 @@ def fmha_decode_launch( v_token_stride: Int64 = 0, static_full_split_prefix: cutlass.Constexpr[bool] = False, use_static_native_seqlens_kv: cutlass.Constexpr[bool] = False, - sage_q_scale_iter: cute.Pointer | None = None, - sage_k_scale_iter: cute.Pointer | None = None, - sage_v_scale_iter: cute.Pointer | None = None, - sage_v_mean_iter: cute.Pointer | None = None, - sage_q_scale_head_stride: Int32 | None = None, - sage_k_scale_head_stride: Int32 | None = None, + sage_scales_iter: cute.Pointer | None = None, ) -> None: """Standalone JIT launcher for FMHA decode TS. - The ``sage_*`` scale arguments are used only by a Sage attention config. + ``sage_scales_iter`` is used only by a Sage attention config, which takes + its Q, K and V scales from the prepared scales (``_SageScaleLayout``). """ log2_e = math.log2(math.e) b, h_q, h_k, s_k, d = problem_shape @@ -3826,12 +3690,7 @@ def fmha_decode_launch( null_sparse_route_ptr, null_sparse_route_ptr, static_full_split_prefix, - g_sage_q_scale=sage_q_scale_iter, - g_sage_k_scale=sage_k_scale_iter, - g_sage_v_scale=sage_v_scale_iter, - g_sage_v_mean=sage_v_mean_iter, - g_sage_q_scale_head_stride=sage_q_scale_head_stride, - g_sage_k_scale_head_stride=sage_k_scale_head_stride, + g_sage_scales=sage_scales_iter, ).launch( grid=grid, block=[cfg.threads_per_cta, 1, 1], @@ -3866,14 +3725,7 @@ def fmha_block_sparse_launch( num_physical_kv_pages: Int64 = 0, k_page_stride: Int64 = 0, v_page_stride: Int64 = 0, - sage_q_scale_iter: cute.Pointer | None = None, - sage_k_scale_iter: cute.Pointer | None = None, - sage_k_summary_scale_iter: cute.Pointer | None = None, - sage_v_scale_iter: cute.Pointer | None = None, - sage_v_mean_iter: cute.Pointer | None = None, - sage_q_scale_head_stride: Int32 | None = None, - sage_k_scale_head_stride: Int32 | None = None, - sage_k_summary_scale_head_stride: Int32 | None = None, + sage_scales_iter: cute.Pointer | None = None, ) -> None: """Launch attention over exact and typed exact/proxy prepared KV routes. @@ -3882,8 +3734,8 @@ def fmha_block_sparse_launch( words. Exact routes address K/V; proxy routes address summary K/V. Both layouts execute the same ``decode_gen_kernel`` schedule and physical copy policy. Exact builds constexpr-elide summary TensorMaps. - The ``sage_*`` scale arguments are used only by a Sage attention config, - the summary K scales only by its proxy routes. + ``sage_scales_iter`` is used only by a Sage attention config, which takes + its Q, K and V scales from the prepared scales (``_SageScaleLayout``). """ if cutlass.const_expr(not cfg.use_block_sparse): raise ValueError("fmha_block_sparse_launch requires block-sparse config") @@ -4148,14 +4000,7 @@ def fmha_block_sparse_launch( tma_desc_v_summary=v_desc_summary_primary, tma_desc_k_summary_atom=k_desc_summary_atom, tma_desc_v_summary_atom=v_desc_summary_atom, - g_sage_q_scale=sage_q_scale_iter, - g_sage_k_scale=sage_k_scale_iter, - g_sage_k_summary_scale=sage_k_summary_scale_iter, - g_sage_v_scale=sage_v_scale_iter, - g_sage_v_mean=sage_v_mean_iter, - g_sage_q_scale_head_stride=sage_q_scale_head_stride, - g_sage_k_scale_head_stride=sage_k_scale_head_stride, - g_sage_k_summary_scale_head_stride=sage_k_summary_scale_head_stride, + g_sage_scales=sage_scales_iter, ).launch( grid=grid, block=[cfg.threads_per_cta, 1, 1], diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/__init__.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/__init__.py index 664566125c43..303992114dfb 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/__init__.py @@ -58,6 +58,17 @@ Q-side MMA input type. Consumer (MmaTask): uses the same descriptor labels as SmemKvResource. +- SageKScalesResource : One softmax instance's ring of ``sfK`` tiles. + Producer (LoadTask): ``cp.async`` copies of the + tile's atoms from the prepared scales, a dense + tile's or a block-sparse route's, after each K + tile. Consumer (Softmax): reads the waited stage + in place during the max and P passes. + +- SageVScalesResource : The work tile's per-channel ``sfV`` (and + ``v_mean``) staged in SMEM by the correction warps + before the output epilogue reads them. + - SmemPResource : P operand for BMM2. The validated one-instance staged-D256 Keeps profile places P in two TMEM views aliased with the matching S stages; other profiles @@ -125,6 +136,7 @@ """ from .helpers_common import DecodeGenResourceBase +from .sage_scales import SageKScalesResource, SageVScalesResource from .smem_resources import ( SmemKvTileResource, SmemKvResource, @@ -150,6 +162,8 @@ __all__ = [ "DecodeGenResourceBase", + "SageKScalesResource", + "SageVScalesResource", "SmemKvResource", "SmemKvTileResource", "SmemPageOffsetsKvResource", diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_kv_tile_idx.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_kv_tile_idx.py index 3a795a73a4a2..896780914d1e 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_kv_tile_idx.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_kv_tile_idx.py @@ -26,7 +26,9 @@ from cutlass.experimental.task_scheduling.resources import StageInfo from ...mask import kv_tile_is_fully_visible +from ...stage import FmhaStage from ..fmha_decode_config import CAUSAL, FmhaDecodeConfig +from ..fmha_decode_constants import KV_KIND_V from .helpers_common import ( Constexpr, _TASK_CACHE_KV_RAW_TILE_BASE, @@ -348,6 +350,34 @@ def _runtime_split_kv_global_tile_idx( ) +@cute.jit +def _local_kv_tile_idx_for_section( + cfg: Constexpr[FmhaDecodeConfig], + stage_info: StageInfo, + inst_id: Constexpr[int], + kv_kind: Constexpr[int], + section: Constexpr[FmhaStage], +) -> Int32: + """Return the local K/V tile index implied by the TS HEAD/LOOP/TAIL cadence. + + The schedule names whether a producer is K0/K1/V0/V1, while the phase + defines where that logical tile sits in the staggered decode pipeline: + + - HEAD produces only initial K tiles. + - LOOP produces V for the current MMA iteration and K for the next one. + - TAIL drains the final V tiles after the last loop iteration. + """ + num_insts_kv = Int32(cfg.num_insts_kv) + if cutlass.const_expr(section == FmhaStage.Head): + return Int32(inst_id) + if cutlass.const_expr(section == FmhaStage.Loop): + base = stage_info.loop_offset * num_insts_kv + Int32(inst_id) + if cutlass.const_expr(kv_kind == KV_KIND_V): + return base + return base + num_insts_kv + return stage_info.loop_end * num_insts_kv + Int32(inst_id) + + @cute.jit def resolve_keeps_tile_context( cfg: Constexpr[FmhaDecodeConfig], @@ -358,12 +388,16 @@ def resolve_keeps_tile_context( max_seq_len_kv: Int32, seq_len_q: Int32, q_group_idx: Int32 | None, + local_tile_idx: Int32 | None = None, ): """Resolve one score tile's logical position and boundary-mask state. - Every consumer of a softmax instance's tiles (the S resource and the - Sage K scale resource) derives the tile's token range from the same - sequence-length sources, so the resolution lives outside the S resource. + Every user of a softmax instance's dense tiles (the S resource, and the + load warp copying a dense Sage plan's K scales for the tile) derives the + tile's token range from the same sequence-length sources, so the + resolution lives outside the S resource; block-sparse routes take their + origins from the prepared record instead. ``local_tile_idx`` defaults to + the softmax's tile; a producer passes the tile it loads. """ task_cache = _decode_gen_task_cache(stage_info) if cutlass.const_expr(seqlens_kv is None): @@ -384,7 +418,8 @@ def resolve_keeps_tile_context( element_mask_end_idx = seq_len_kv - seq_len_q + q_token_base + Int32(1) use_runtime_kv_domain = seqlens_kv is not None or cfg.uses_runtime_q_kv_union - local_tile_idx = _softmax_tile_idx(cfg, stage_info, inst_id) + if cutlass.const_expr(local_tile_idx is None): + local_tile_idx = _softmax_tile_idx(cfg, stage_info, inst_id) if cutlass.const_expr(not use_runtime_kv_domain): effective_tile_idx = _static_split_kv_global_tile_idx( cfg, stage_info, local_tile_idx diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_softmax.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_softmax.py index ca7b8a662cec..bca4608eed8e 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_softmax.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_softmax.py @@ -26,7 +26,6 @@ from cutlass.experimental import primitives as prims from ..fmha_decode_config import FmhaDecodeConfig -from ..fmha_decode_constants import FP8_P_QUANT_SCALE from .helpers_common import ( Constexpr, fadd2, @@ -412,7 +411,7 @@ def _attention_sink_for_local_head( fastmath=True, ) if cutlass.const_expr(cfg.use_fp8_pv): - sink_exp = sink_exp * Float32(FP8_P_QUANT_SCALE) + sink_exp = sink_exp * Float32(cfg.fp8_p_quant_scale) return sink_exp diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/sage_scales.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/sage_scales.py index f2df9bbeafcb..0ee8f36eb909 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/sage_scales.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/sage_scales.py @@ -14,31 +14,27 @@ """Sage attention scale addressing for the decode softmax and epilogue. -Every scale load of the kernel goes through this module. ``sfQ`` and ``sfK`` -use the trtllm-gen flat layout (``flat_scale_slot`` of -:mod:`flashinfer.attention.prims_ts.sage`). Tokens past the sequence end and -Q rows past the valid row count clamp to the last valid slot, so masked -scores keep a finite scale and exponentiate to zero. - -:class:`SageKScalesResource` holds one softmax instance's ``sfK`` words for -one KV tile. The instance produces and consumes it: it fills the words ahead -of the tile's score wait, and both softmax passes read them one K32 fragment -at a time. The scale groups per fragment select the storage form -(``FmhaDecodeConfig.sage_k_scales_in_smem_for``): - -* Register form (K blocks of 16 tokens or more, at most two groups per - fragment): each lane gathers its own words into a rotating register array; - there is no pipeline. -* SMEM form (K blocks of 4 and 1 token, eight or 32 groups per fragment): - a two-slot SMEM buffer in the ``sage_k_scale_words`` layout, pipelined - between the instance's own warps. The passes read a fragment's groups with - 16-byte broadcast loads; the routed register array is a placeholder. - -A dense tile gathers ``k_scale`` over its token range, a block-sparse route -takes the words the load warp staged with its metadata, and a proxy route of -a mixed-geometry plan (``sage_mixed_k_geometry``) gathers ``k_summary_scale`` -from the route's atom origins. A mixed plan holds one resource per geometry; -the CTA-uniform route kind selects the one that serves the tile. +Every scale reaches the kernel from the buffer the Sage prepare writes +before attention (``_SageScaleLayout``), in the order the kernel reads it: +``sfQ`` as one word per row of each Q tile, ``sfK`` as one atom of words per +64 K tokens, ``sfV`` and ``v_mean`` as one row per sequence. The prepare +resolves the flat layout (``flat_scale_slot`` of +:mod:`flashinfer.attention.prims_ts.sage`, ``load_flat_scale``) and gives Q +rows past the valid row count and K tokens past the sequence end the last +valid scale, so masked scores keep a finite scale and exponentiate to zero. +Each softmax lane loads its row's ``sfQ`` word once per work tile +(``load_q_scale``). + +:class:`SageKScalesResource` is one softmax instance's ring of ``sfK`` +tiles, which the load warp produces after each K tile with ``cp.async`` +copies under the ring's ``AsyncLoad`` commit: the tile's atoms, a dense +tile's or a block-sparse route's, from the prepared scales +(``_copy_atom_scales``). After the softmax task waits a stage, +``open_stage`` publishes the stage's shared-memory offset as a task-local +value that the schedule routes into the max pass and the P pass; both read +the stage in place through ``sage_k_view`` and ``load_sage_k_fragment``, one +K32 fragment at a time with loads of at most 16 bytes, so attention never +addresses ``k_scale`` itself and no other resource holds the ring. :class:`SageVScalesResource` stages the work tile's per-channel V scales (and means) in SMEM once per tile, while no output is pending, so the @@ -46,13 +42,13 @@ waiting on global loads. """ -from collections.abc import Callable from dataclasses import dataclass -from typing import TYPE_CHECKING, ClassVar +from typing import ClassVar import cutlass import cutlass.cute as cute from cutlass import Float32, Int32, Int64 +from cutlass.experimental import primitives as prims from cutlass.experimental.task_scheduling.enums import WorkAttr from cutlass.experimental.task_scheduling.memory import ( ResourceContext, @@ -66,49 +62,31 @@ producer_work, ) -from ....sage import flat_scale_slot, log2_block_size -from ...placeholder_helpers import _placeholder_local_array +from ...._block_sparse.prepared import ( + _PREPARED_ROUTE_IS_PROXY_FLAG, + _BlockSparseRouteLayout, +) +from ....sage import _FLOAT4_WORDS, _SAGE_ATOM_TOKENS, _SageKSection, flat_scale_slot +from ...stage import FmhaStage from ..fmha_decode_config import FmhaDecodeConfig +from ..fmha_decode_constants import KV_KIND_K from .helpers_common import ( DecodeGenResourceBase, _TASK_CACHE_WARP_GRP_THREAD_IDX, _decode_gen_task_cache, - _keeps_route_atom, _keeps_spatial_half, _logical_head_batch, - _route_is_proxy, + _warp_broadcast_i32, fmul2, ) -from .helpers_kv_tile_idx import resolve_keeps_tile_context - -if TYPE_CHECKING: - from .smem_block_sparse_metadata import SmemBlockSparseSoftmaxMetadataResource +from .helpers_kv_tile_idx import ( + _local_kv_tile_idx_for_section, + resolve_keeps_tile_context, +) Constexpr = cutlass.Constexpr -@dataclass(eq=False) -class SageScaleTensors: - """The plan's Sage scale tensors, shared by every resource that reads one. - - ``sfQ``, ``sfK`` and ``k_summary_scale`` are ``[heads, slots]`` FP32 - arrays in the flat layout; the V scales and means are ``[Hkv, D]`` FP32. - ``k_summary_scale`` is ``None`` without proxy routes and ``v_mean`` - without a channel mean. The dataclass is not frozen because the DSL - replaces frozen dataclasses with proxies inside traced dynamic branches, - which changes the traced structure of the holding resource. - """ - - q_scale_ptr: cute.Pointer - q_scale_head_stride: Int32 - k_scale_ptr: cute.Pointer - k_scale_head_stride: Int32 - k_summary_scale_ptr: cute.Pointer | None - k_summary_scale_head_stride: Int32 | None - v_scale_ptr: cute.Pointer - v_mean_ptr: cute.Pointer | None - - @cute.jit def _load_scale(scale_addr: Int64, index: Int32) -> Float32: """Return one FP32 scale loaded from a global memory base address.""" @@ -121,129 +99,237 @@ def _load_scale(scale_addr: Int64, index: Int32) -> Float32: @cute.jit -def load_q_scale( - cfg: Constexpr[FmhaDecodeConfig], - q_scale_addr: Int64, - q_scale_head_stride: Int32, +def load_flat_scale( + scale_addr: Int64, + head_stride: Int32, *, - kv_head_idx: Int32, - local_head_idx: Int32, + head_idx: Int32, batch_idx: Int32, - q_token_idx: Int32, + seq_len: Int32, + token_idx: Int32, + log2_block: Int32, ) -> Float32: - """Return ``sfQ`` for one Q row (token and Q head of one KV head group).""" - heads_q_per_kv = Int32(cfg.heads_q_per_kv) - q_head_idx = kv_head_idx * heads_q_per_kv + cute.math.min( - local_head_idx, heads_q_per_kv - Int32(1) - ) - q_token_idx = cute.math.min(q_token_idx, Int32(cfg.max_seq_len_q - 1)) - slot = flat_scale_slot( - batch_idx, - q_token_idx, - Int32(cfg.max_seq_len_q), - log2_block_size(cfg.sage_q_block_size), - ) - return _load_scale(q_scale_addr, q_head_idx * q_scale_head_stride + slot) + """Return one head's scale of one token from a flat-layout scale array. + + A token past ``seq_len`` takes the last valid slot. + """ + token_idx = cute.math.min(token_idx, seq_len - Int32(1)) + slot = flat_scale_slot(batch_idx, token_idx, seq_len, log2_block) + return _load_scale(scale_addr, head_idx * head_stride + slot) @cute.jit -def load_k_scale( +def sage_sequence_words( cfg: Constexpr[FmhaDecodeConfig], - k_scale_addr: Int64, - k_scale_head_stride: Int32, *, - kv_head_idx: Int32, batch_idx: Int32, - seq_len_kv: Int32, - kv_token_idx: Int32, - log2_k_block_size: Int32 | None = None, + kv_head_idx: Int32, + num_heads_kv: Int32, +) -> Int32: + """Return the first word of sequence ``b * Hkv + h`` in the prepared scales.""" + layout = cfg.sage_scale_layout(cfg.static_seq_len_kv) + sequence_idx = batch_idx * num_heads_kv + kv_head_idx + return sequence_idx * Int32(layout.sequence_words) + + +@cute.jit +def load_q_scale( + cfg: Constexpr[FmhaDecodeConfig], + scales_ptr: cute.Pointer, + *, + sequence_words: Int32, + q_group_idx: Int32, + row_idx: Int32, ) -> Float32: - """Return ``sfK`` for one token of one KV head from a flat scale array. + """Return ``sfQ`` of row ``row_idx`` of Q tile ``q_group_idx`` from the prepared scales. - ``k_scale_addr`` is the base address of ``k_scale`` or ``k_summary_scale``, - ``seq_len_kv`` the length of the sequence it covers and - ``log2_k_block_size`` the log2 of its K block size, the recipe's token - block by default. + ``sequence_words`` is the first word of the row's sequence + (``sage_sequence_words``). The word already holds the scale of the row's + token and Q head, the last valid row's for a row past the tile's valid + rows (``_SageQSection``). """ - log2_block = log2_k_block_size - if cutlass.const_expr(log2_block is None): - log2_block = Int32(log2_block_size(cfg.sage_k_block_size)) - kv_token_idx = cute.math.min(kv_token_idx, seq_len_kv - Int32(1)) - slot = flat_scale_slot(batch_idx, kv_token_idx, seq_len_kv, log2_block) - return _load_scale(k_scale_addr, kv_head_idx * k_scale_head_stride + slot) + q_rows = cfg.sage_scale_layout(cfg.static_seq_len_kv).q_rows + word = ( + sequence_words + + Int32(q_rows.word_offset) + + q_group_idx * Int32(q_rows.tile_size_q) + + row_idx + ) + return _load_scale(scales_ptr.toint(), word) -def sage_scale_arr_size(cfg: FmhaDecodeConfig, groups: int) -> int: - """Return the number of ``sfK`` words one softmax lane holds per tile. +def sage_k_section(cfg: FmhaDecodeConfig, proxy: bool = False) -> _SageKSection: + """Return the prepared K section one route kind reads. - ``groups`` is the scale groups per fragment of the tile's route kind - (``FmhaDecodeConfig.sage_k_groups_per_fragment_for``). + A proxy route reads the summaries, every other tile the K tokens. The + section's words per atom are the softmax's scale groups of the atom's + two fragments, so an unpadded atom's words are contiguous per fragment. """ - return cfg.num_softmax_score_fragments * groups + sections = cfg.sage_scale_layout(cfg.static_seq_len_kv).k_sections + section = sections[-1] if proxy else sections[0] + assert section.words_per_atom == ( + cfg.keeps_fragments_per_atom * cfg.sage_k_groups_per_fragment(proxy) + ) + return section -def sage_k_scale_words(cfg: FmhaDecodeConfig, groups: int) -> int: - """Return the number of ``sfK`` words that cover one KV tile for every spatial half. +def sage_k_stage_words(cfg: FmhaDecodeConfig) -> int: + """Return the ``sfK`` words of one ring stage: a KV tile's atoms of the larger section. - Zero without Sage attention. Word ``half * arr_size + f * groups + g`` - holds group ``g`` of fragment ``f`` of that half's lanes. Route staging - and the SMEM form of :class:`SageKScalesResource` share this layout, so a - softmax thread reads its half with contiguous vector loads. + A route kind's atoms lie contiguously in the stage in the order of + ``FmhaDecodeConfig.keeps_route_atom_owner``: atom ``a``, held by spatial + half ``h`` at position ``p``, is stage atom ``h * atoms_per_half + p``, + so each half's words are contiguous. """ - if not cfg.use_sage_attention: - return 0 - return cfg.keeps_spatial_halves * sage_scale_arr_size(cfg, groups) - + atom_words = max(sage_k_section(cfg, proxy).atom_words for proxy in (False, True)) + return cfg.tile_size_kv // cfg.keeps_atom_tokens * atom_words -def sage_staged_k_scale_words(cfg: FmhaDecodeConfig) -> int: - """Return the ``sfK`` words a block-sparse route stages. - Exact routes stage their words; proxy routes of a mixed-geometry plan - gather theirs in the softmax warps, so the staged area covers the exact - geometry alone. +@cute.jit +def _copy_atom_scales( + cfg: Constexpr[FmhaDecodeConfig], + scales_ptr: cute.Pointer, + stage_words: cute.Pointer, + section_words: Int32, + atom_words: Constexpr[int], + atom_index: Constexpr, + num_atoms: Constexpr[int], +) -> None: + """Copy ``num_atoms`` atoms of one K section into a ring stage. + + Lane ``l`` copies float4 ``l``, ``l + 32``, ... of the atoms with + ``cp.async``, so the TMA unit stays free for K and V. Each float4 moves + from word ``section_words + atom_index(atom) * atom_words`` of the + prepared scales plus its offset in the atom, ``section_words`` being the + section's first word in the tile's sequence, to the stage word of the + atom's spatial half and position (``sage_k_stage_words``) plus the same + offset. """ - return sage_k_scale_words(cfg, cfg.sage_k_groups_per_fragment) + atom_float4s = atom_words // _FLOAT4_WORDS + num_float4s = num_atoms * atom_float4s + lane_idx = cute.arch.thread_idx()[0] & Int32(31) + for round_idx in cutlass.range_constexpr((num_float4s + 31) // 32): + float4 = lane_idx + Int32(round_idx * 32) + atom = float4 // Int32(atom_float4s) + word_in_atom = (float4 % Int32(atom_float4s)) * Int32(_FLOAT4_WORDS) + half, position = cfg.keeps_route_atom_owner(atom) + stage_atom = half * Int32(num_atoms // cfg.keeps_spatial_halves) + position + # Every lane resolves its atom's index: a route's lookup shuffles. + source_word = section_words + atom_index(atom) * Int32(atom_words) + if float4 < Int32(num_float4s): + prims.cp_async_shared_global( + stage_words + stage_atom * Int32(atom_words) + word_in_atom, + scales_ptr + source_word + word_in_atom, + _FLOAT4_WORDS * 4, + "cg", + ) @cute.jit -def sage_word_position( - cfg: Constexpr[FmhaDecodeConfig], - half: Int32, - lane_entry, - groups: Constexpr[int], -) -> tuple[Int32, Int32]: - """Return ``(atom, token offset in the atom)`` of one ``sfK`` word. - - ``lane_entry = f * groups + g`` names group ``g`` of fragment ``f`` of - spatial half ``half``. Fragment ``f`` reads the half's Keeps layout atom - ``f // fragments_per_atom`` (``_keeps_route_atom``), which block-sparse - routes share, from token ``(f % fragments_per_atom) * fragment_regs``; - group ``g`` starts ``g * group_tokens`` later. +def _route_atom_origin( + route_layout: Constexpr[_BlockSparseRouteLayout], + record_word: Int32, + origin1: Int32, + atom: Int32, +) -> Int32: + """Return the token origin of route atom ``atom`` from the resolved record. + + ``SmemBlockSparseKvMetadataResource.resolve_route`` hands the load warp + its record in one of three forms: + + * a two-atom route without one-warp transport returns both origins as + warp-uniform scalars, ``record_word`` for atom 0 and ``origin1`` for + atom 1; + * a two-atom route with one-warp transport returns lane-distributed + origins in ``record_word``, lane ``a`` holding atom ``a``'s, and atom + 1's again in ``origin1``; + * a wider route returns lane-distributed origins in ``record_word`` and + zero in ``origin1``. + + A lane-distributed origin is shuffled from its lane, one SHFL per float4 + round, also where lane ``a`` copies atom ``a`` and holds it already. An + invalid atom has origin ``-1``. """ - fragments_per_atom = cfg.keeps_fragments_per_atom - fragment_regs = cfg.softmax_score_fragment_regs - group_tokens = fragment_regs // groups - fragment_idx = lane_entry // groups - group_idx = lane_entry % groups - atom_idx = _keeps_route_atom(cfg, half, fragment_idx // fragments_per_atom) - token_offset = ( - fragment_idx % fragments_per_atom - ) * fragment_regs + group_idx * group_tokens - return atom_idx, Int32(token_offset) + if cutlass.const_expr( + route_layout.logical_origins_per_route == 2 + and not route_layout.uses_one_warp_transport + ): + origin = Int32(record_word) + if atom != Int32(0): + origin = Int32(origin1) + return origin + return cute.arch.shuffle_sync(record_word, atom) @cute.jit -def dense_k_scale_token( - cfg: Constexpr[FmhaDecodeConfig], half: Int32, lane_entry, tile_offset_k: Int32 +def _route_atom_index( + route_layout: Constexpr[_BlockSparseRouteLayout], + record_word: Int32, + origin1: Int32, + atom: Int32, ) -> Int32: - """Return the first KV token covered by one ``sfK`` word of a dense tile. + """Return the section atom of route atom ``atom``: the 64-token atom of its origin. + + An invalid atom takes atom 0, whose finite scales its masked scores + never use. + """ + origin = _route_atom_origin(route_layout, record_word, origin1, atom) + return cute.math.max(origin, Int32(0)) // Int32(_SAGE_ATOM_TOKENS) + - A dense Keeps tile lays out its atoms in order (``_keeps_score_col``). +@cute.jit +def copy_route_k_scales( + cfg: Constexpr[FmhaDecodeConfig], + route_layout: Constexpr[_BlockSparseRouteLayout], + scales_ptr: cute.Pointer, + *, + stage_words: cute.Pointer, + sequence_words: Int32, + record_word: Int32, + origin1: Int32, +) -> None: + """Copy a block-sparse route's ``sfK`` atoms from the prepared scales into a ring stage. + + ``sequence_words`` is the first word of the route's sequence. Each route + atom reads the atom of its origin (``_route_atom_index``) from the + section its route kind reads: the K tokens, or the summaries of a proxy + route. A mixed plan's sections differ in atom size, so each kind has its + own copy; otherwise one copy serves both kinds from the kind's section. """ - atom_idx, token_offset = sage_word_position( - cfg, half, lane_entry, cfg.sage_k_groups_per_fragment + k_tokens = sage_k_section(cfg) + summaries = sage_k_section(cfg, proxy=True) + num_atoms = route_layout.logical_origins_per_route + is_proxy_route = cutlass.Boolean(False) + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): + route_flags = _warp_broadcast_i32( + record_word, route_layout.route_flags_word_offset + ) + is_proxy_route = (route_flags & Int32(_PREPARED_ROUTE_IS_PROXY_FLAG)) != Int32( + 0 + ) + atom_index = lambda atom: _route_atom_index( + route_layout, record_word, origin1, atom ) - return tile_offset_k + atom_idx * Int32(cfg.keeps_atom_tokens) + token_offset + copy = lambda section, section_words: _copy_atom_scales( + cfg, + scales_ptr, + stage_words, + section_words, + section.atom_words, + atom_index, + num_atoms, + ) + if cutlass.const_expr(cfg.sage_mixed_k_geometry): + if is_proxy_route: + copy(summaries, sequence_words + Int32(summaries.word_offset)) + else: + copy(k_tokens, sequence_words + Int32(k_tokens.word_offset)) + else: + section_words = sequence_words + Int32(k_tokens.word_offset) + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): + if is_proxy_route: + section_words = sequence_words + Int32(summaries.word_offset) + copy(k_tokens, section_words) @cute.jit @@ -259,33 +345,6 @@ def scale_pairs_in_place( ) -@cute.jit -def block_sparse_k_scale_source( - cfg: Constexpr[FmhaDecodeConfig], - scales: SageScaleTensors, - *, - seq_len_kv: Int32, - route_is_proxy: cutlass.Boolean, -) -> tuple[Int64, Int32, Int32, Int32]: - """Return ``(address, head stride, sequence length, log2 block size)``. - - Exact routes read ``k_scale`` over the KV tokens; proxy routes read - ``k_summary_scale`` over the summary sequence, one summary per KV block. - The result feeds :func:`load_k_scale`. - """ - scale_addr = scales.k_scale_ptr.toint() - head_stride = scales.k_scale_head_stride - seq_len = seq_len_kv - log2_block = Int32(log2_block_size(cfg.sage_k_block_size)) - if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): - if route_is_proxy: - scale_addr = scales.k_summary_scale_ptr.toint() - head_stride = scales.k_summary_scale_head_stride - seq_len = Int32(cfg.num_proxy_summaries) - log2_block = Int32(log2_block_size(cfg.sage_k_summary_block_size)) - return scale_addr, head_stride, seq_len, log2_block - - def staged_v_channel_scale_entries(cfg: FmhaDecodeConfig) -> int: """Return the number of SMEM floats holding one KV head's V scales and means. @@ -298,28 +357,31 @@ def staged_v_channel_scale_entries(cfg: FmhaDecodeConfig) -> int: @cute.jit def stage_v_channel_scales( cfg: Constexpr[FmhaDecodeConfig], - scales: SageScaleTensors, + scales_ptr: cute.Pointer, staged: cutlass.Array, *, - kv_head_idx: Int32, + sequence_words: Int32, thread_idx: Int32, num_threads: Constexpr[int], ) -> None: - """Copy one KV head's per-channel V scales and means into SMEM. + """Copy one sequence's per-channel V scales and means into SMEM. - Each of the ``num_threads`` callers moves ``headdim / num_threads`` - channels. Callers order these writes before the epilogue reads, and the - next tile's writes after the last read. + The sequence's V rows of the prepared scales (``_SageVSection``) have the + staged form already: the scales, then the means. Each of the + ``num_threads`` callers moves ``headdim / num_threads`` channels of each + row. Callers order these writes before the epilogue reads, and the next + tile's writes after the last read. """ assert cfg.headdim % num_threads == 0 - v_scale_addr = scales.v_scale_ptr.toint() - head_base = kv_head_idx * Int32(cfg.headdim) + v_rows = cfg.sage_scale_layout(cfg.static_seq_len_kv).v_rows + scales_addr = scales_ptr.toint() + rows_base = sequence_words + Int32(v_rows.word_offset) for base in cutlass.range_constexpr(0, cfg.headdim, num_threads): channel = Int32(base) + thread_idx - staged[channel] = _load_scale(v_scale_addr, head_base + channel) + staged[channel] = _load_scale(scales_addr, rows_base + channel) if cutlass.const_expr(cfg.sage_v_mean): staged[channel + Int32(cfg.headdim)] = _load_scale( - scales.v_mean_ptr.toint(), head_base + channel + scales_addr, rows_base + Int32(v_rows.mean_word_offset) + channel ) @@ -356,222 +418,245 @@ def load_staged_v_channel_scales( return scales, means -WordSource = Callable[[Int32, Int32, object], Float32] +def sage_k_half_words(cfg: FmhaDecodeConfig, proxy: bool = False) -> int: + """Return one route kind's stage words of one spatial half's atoms.""" + tile_atoms = cfg.tile_size_kv // cfg.keeps_atom_tokens + atoms_per_half = tile_atoms // cfg.keeps_spatial_halves + return atoms_per_half * sage_k_section(cfg, proxy).atom_words + -SAGE_K_SCALES_RING_STAGES = 2 +@cute.jit +def sage_k_view( + cfg: Constexpr[FmhaDecodeConfig], + context: ResourceContext, + stage_bytes: Int32, + warp_grp_thread_idx: Int32, + factor: Float32 | None, + proxy_kind: Constexpr[bool] = False, +): + """Return a pass's view of a waited ring stage: the lane's half pointer and ``factor``. + + ``stage_bytes`` is the stage's offset in the CTA's shared memory, as + ``SageKScalesResource.open_stage`` publishes it; ``factor``, if any, is + multiplied into every word ``load_sage_k_fragment`` returns. + """ + stage_words = cutlass.Array( + context.smem_base.data_ptr() + stage_bytes, + dtype=Float32, + shape=(sage_k_stage_words(cfg),), + addrspace=3, + ).data_ptr() + half = _keeps_spatial_half(cfg, warp_grp_thread_idx) + return stage_words + half * Int32(sage_k_half_words(cfg, proxy_kind)), factor + + +@cute.jit +def load_sage_k_fragment( + cfg: Constexpr[FmhaDecodeConfig], + view, + fragment_idx: Int32, + proxy_kind: Constexpr[bool] = False, +) -> cutlass.Array: + """Return ``factor * sfK`` of one fragment's ``groups`` words of ``view``. + + Fragment ``f`` reads words ``(f % 2) * groups`` onward of atom ``f // 2`` + of the half; an unpadded atom's words are contiguous, so that is word ``f + * groups`` of the half. The loads, at most 16 bytes each, issue ahead of + the fragment's score wait in the callers. + """ + groups = cfg.sage_k_groups_per_fragment(proxy_kind) + values = cutlass.Array(Float32, groups, space=cutlass.AddressSpace.rmem) + half_ptr, factor = view + fragments_per_atom = cfg.keeps_fragments_per_atom + section = sage_k_section(cfg, proxy_kind) + if cutlass.const_expr(section.atom_words == section.words_per_atom): + words_ptr = half_ptr + fragment_idx * Int32(groups) + else: + atom_position = fragment_idx // Int32(fragments_per_atom) + words_ptr = ( + half_ptr + + atom_position * Int32(section.atom_words) + + fragment_idx % Int32(fragments_per_atom) * Int32(groups) + ) + width = min(groups, 4) + for word in cutlass.range_constexpr(0, groups, width): + loaded = (words_ptr + Int32(word)).load(count=width, alignment=width * 4) + for elem in cutlass.range_constexpr(width): + values[word + elem] = Float32(loaded[elem]) + if cutlass.const_expr(factor is not None): + if cutlass.const_expr(groups == 1): + values[0] = Float32(values[0]) * factor + else: + scale_pairs_in_place(values, factor, groups) + return values @dataclass(kw_only=True) class SageKScalesResource(DecodeGenResourceBase): - """One softmax instance's ``sfK`` words for one scale-group geometry. - - ``summary`` marks the resource of a mixed-geometry plan that serves proxy - routes from ``k_summary_scale``. The SMEM form takes the instance's - two-stage pipeline as ``pipeline_config``; the register form takes none. - ``route_metadata`` is the instance's block-sparse softmax metadata - resource, whose held stage carries the route's staged words and atom - origins; a dense plan resolves the tile's tokens from the sequence-length - fields instead. + """One softmax instance's ring of ``sfK`` tiles, produced by the load warp. + + Each stage holds one KV tile's words in the ``sage_k_stage_words`` layout. + The load warp fills it after the tile's K copy with ``cp.async`` copies of + the tile's atoms from the prepared scales, a block-sparse route's + (``copy_route``) or a dense tile's (``copy_tile``), and the pipeline's + ``AsyncLoad`` commit completes the stage once they land. The softmax waits + a stage, takes its offset from ``open_stage``, reads it in place through + ``sage_k_view`` and ``load_sage_k_fragment`` in the max pass and the P + pass, and releases it after the last pass that reads the words. """ _task_local_specs: ClassVar[tuple[tuple, ...]] = ( ( - "sage_scale_arr", - cutlass.Array, - None, - "Raw sfK words of the lane's fragments for the current tile.", + "stage_bytes", + Int32, + Int32(0), + "Shared-memory offset of the waited ``sfK`` ring stage.", ), ) + stage_bytes: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() cfg: Constexpr[FmhaDecodeConfig] = None inst_id: Constexpr[int] = 0 - summary: Constexpr[bool] = False - route_metadata: "SmemBlockSparseSoftmaxMetadataResource | None" = None + route_layout: Constexpr[_BlockSparseRouteLayout | None] = None seqlens_kv: cute.Pointer | None = None max_seq_len_kv: Int32 = None seq_len_q: Int32 = None q_group_idx: Int32 | None = None h_k_idx: Int32 | None = None b_idx: Int32 | None = None - scale_tensors: SageScaleTensors | None = None - sage_scale_arr: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + num_heads_kv: Int32 | None = None + sage_scales_ptr: cute.Pointer | None = None _alloc: Constexpr[SmemAllocation | None] = None def __post_init__(self) -> None: - assert (self.pipeline_config is not None) == self.in_smem, ( - "the SMEM form of the Sage K scales carries the instance's pipeline " - "and the register form carries none" - ) + assert self.pipeline_config is not None + assert (self.route_layout is not None) == self.cfg.use_block_sparse super().__post_init__() - @property - def groups(self) -> int: - """Return the scale groups per K32 fragment of the served geometry.""" - if self.summary: - return self.cfg.sage_summary_k_groups_per_fragment - return self.cfg.sage_k_groups_per_fragment - - @property - def in_smem(self) -> bool: - """Whether the tile's words live in the SMEM ring rather than registers.""" - return self.cfg.sage_k_scales_in_smem_for(self.groups) - - @property - def arr_size(self) -> int: - """Return the ``sfK`` words one lane holds per tile in the register form.""" - return sage_scale_arr_size(self.cfg, self.groups) - @property def tile_words(self) -> int: - """Return the ``sfK`` words of one tile slot of the SMEM ring.""" - return sage_k_scale_words(self.cfg, self.groups) + """Return the ``sfK`` words of one tile in a ring stage.""" + return sage_k_stage_words(self.cfg) @property - def routed_words(self) -> int: - """Return the length of the routed array: the lane's words, or one placeholder.""" - return 1 if self.in_smem else self.arr_size - - def _init_placeholder_state(self) -> None: - """Create the placeholder routed array for task-graph tracing.""" - self.sage_scale_arr.default = _placeholder_local_array( - Float32, self.routed_words - ) + def tile_atoms(self) -> int: + """Return the atoms of one KV tile.""" + return self.cfg.tile_size_kv // self.cfg.keeps_atom_tokens def get_smem_requirements(self) -> list[SmemAllocation]: - """Allocate the ring of the SMEM form; the register form needs none.""" - if not self.in_smem: - return [] + """Allocate the ring.""" if self._alloc is None: self._alloc = SmemAllocation( name=self.name, - size_bytes=SAGE_K_SCALES_RING_STAGES * self.tile_words * 4, + size_bytes=self.pipeline_config.num_stages * self.tile_words * 4, alignment=16, ) return [self._alloc] def get_tmem_requirements(self) -> list[TmemAllocation]: - """The words live in SMEM or registers only.""" + """The words live in SMEM only.""" return [] - # -- schedule steps ---------------------------------------------------- - - @producer_work(work_attrs=WorkAttr.AUXILIARY) - @cute.jit - def publish_tile(self, stage_info: StageInfo) -> None: - """Fill the acquired ring slot with a dense tile's words (SMEM form).""" - self._fill_ring(stage_info, cutlass.Boolean(False)) - - @producer_work(work_attrs=WorkAttr.AUXILIARY) - @cute.jit - def publish_route_tile(self, stage_info: StageInfo, *, route_flags: Int32) -> None: - """Fill the acquired ring slot with a block-sparse route's words (SMEM form). + # -- producer steps ---------------------------------------------------- - ``route_flags`` comes from the preceding ``load_route`` of the route - metadata resource, whose stage must stay held until this step and - ``take_route_tile`` have run. - """ - self._fill_ring(stage_info, self._route_is_proxy(route_flags)) - - @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=sage_scale_arr) @cute.jit - def take_tile(self, stage_info: StageInfo) -> cutlass.Array: - """Return the routed array of a dense tile.""" - return self._take(stage_info, cutlass.Boolean(False)) + def _stage_words(self, stage_info: StageInfo) -> cute.Pointer: + """Return the acquired stage's first word as an FP32 SMEM pointer.""" + return self._ring(stage_info.context).data_ptr() + Int32( + stage_info.stage_idx + ) * Int32(self.tile_words) - @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=sage_scale_arr) @cute.jit - def take_route_tile( - self, stage_info: StageInfo, *, route_flags: Int32 - ) -> cutlass.Array: - """Return the routed array of a block-sparse route.""" - return self._take(stage_info, self._route_is_proxy(route_flags)) - - @cute.jit - def _route_is_proxy(self, route_flags: Int32) -> cutlass.Boolean: - """Return the route kind, resolved only where it selects a geometry.""" - route_is_proxy = cutlass.Boolean(False) - if cutlass.const_expr(self._selects_by_route): - route_is_proxy = cute.arch.make_warp_uniform(_route_is_proxy(route_flags)) - return route_is_proxy + def _sequence_words(self, stage_info: StageInfo) -> Int32: + """Return the first word of the work tile's sequence in the prepared scales.""" + kv_head_idx, batch_idx = _logical_head_batch( + stage_info, self.h_k_idx, self.b_idx + ) + return sage_sequence_words( + self.cfg, + batch_idx=batch_idx, + kv_head_idx=kv_head_idx, + num_heads_kv=self.num_heads_kv, + ) + @producer_work @cute.jit - def _fill_ring( - self, stage_info: StageInfo, route_is_proxy: cutlass.Boolean + def copy_route( + self, + stage_info: StageInfo, + *, + resolved_record_word: Int32, + resolved_origin1: Int32, ) -> None: - """Store the tile's words into the acquired slot (SMEM form). + """Copy a block-sparse route's words into the acquired ring stage. - In a mixed-geometry plan only the resource of the route's kind fills; - the branch is CTA-uniform. + The load warp issues the ``cp.async`` copies only; the pipeline's + ``AsyncLoad`` commit completes the stage once they land. """ - assert self.in_smem - if self._serves(route_is_proxy): - self._store_tile_words(stage_info) - - @cute.jit - def _take( - self, stage_info: StageInfo, route_is_proxy: cutlass.Boolean - ) -> cutlass.Array: - """Return the routed array: the lane's words in the register form, else ones.""" - words = cutlass.Array( - Float32, self.routed_words, space=cutlass.AddressSpace.rmem + assert self.route_layout is not None + copy_route_k_scales( + self.cfg, + self.route_layout, + self.sage_scales_ptr, + stage_words=self._stage_words(stage_info), + sequence_words=self._sequence_words(stage_info), + record_word=resolved_record_word, + origin1=resolved_origin1, ) - for entry in cutlass.range_constexpr(self.routed_words): - words[entry] = Float32(1.0) - if cutlass.const_expr(not self.in_smem): - if self._serves(route_is_proxy): - self._gather_lane_words(stage_info, words) - return words - - @property - def _selects_by_route(self) -> bool: - """Whether the route kind selects between two geometries of this plan.""" - return self.cfg.sage_mixed_k_geometry and self.route_metadata is not None - - @cute.jit - def _serves(self, route_is_proxy: cutlass.Boolean) -> cutlass.Boolean: - """Whether this resource holds the route's geometry; constant for one geometry.""" - if cutlass.const_expr(not self._selects_by_route): - return True - if cutlass.const_expr(self.summary): - return route_is_proxy - return not route_is_proxy - - # -- storage fills ----------------------------------------------------- + @producer_work @cute.jit - def _gather_lane_words(self, stage_info: StageInfo, words: cutlass.Array) -> None: - """Take the lane's spatial half of the tile's words into ``words``. - - Every entry index is a constant, so no pass indexes registers at run - time. + def copy_tile( + self, stage_info: StageInfo, *, section: Constexpr[FmhaStage] + ) -> None: + """Copy a dense tile's atoms from the prepared scales into the acquired ring stage. + + The tile is the instance's K tile of ``section``, resolved as the + softmax resolves the tile it reads. Atom ``a`` of the tile reads the + K-token atom of its first token, ``tile_offset_k + a * 64``, clamped + to the sequence's last atom: a tile or atom past the sequence end has + masked scores. The pipeline's ``AsyncLoad`` commit completes the + stage once the copies land. """ - word_value = self._word_source(stage_info) - arr_size = self.arr_size - warp_grp_thread_idx = Int32( - _decode_gen_task_cache(stage_info)[_TASK_CACHE_WARP_GRP_THREAD_IDX] + assert self.route_layout is None + cfg = self.cfg + ( + _seq_len_kv, + _q_group_idx, + _element_mask_end_idx, + tile_offset_k, + _window_start_idx, + _is_valid_effective_tile, + _is_masked_final_wave, + _tile_is_unmasked, + _rows_are_active, + ) = resolve_keeps_tile_context( + cfg, + stage_info, + inst_id=self.inst_id, + seqlens_kv=self.seqlens_kv, + max_seq_len_kv=self.max_seq_len_kv, + seq_len_q=self.seq_len_q, + q_group_idx=self.q_group_idx, + local_tile_idx=_local_kv_tile_idx_for_section( + cfg, stage_info, self.inst_id, KV_KIND_K, section + ), ) - half = _keeps_spatial_half(self.cfg, warp_grp_thread_idx) - half_base = half * Int32(arr_size) - for entry in cutlass.range_constexpr(arr_size): - words[entry] = word_value(half_base + Int32(entry), half, entry) - - @cute.jit - def _store_tile_words(self, stage_info: StageInfo) -> None: - """Store the tile's words into the slot of the acquired producer stage.""" - word_value = self._word_source(stage_info) - num_words = self.tile_words - arr_size = self.arr_size - threads = 32 * self.cfg.softmax_num_warps(self.inst_id) - warp_grp_thread_idx = Int32( - _decode_gen_task_cache(stage_info)[_TASK_CACHE_WARP_GRP_THREAD_IDX] + # A dense plan prepares the K-token section alone. + k_tokens = sage_k_section(cfg) + assert cfg.sage_scale_layout(cfg.static_seq_len_kv).summaries is None + tile_atom = tile_offset_k // Int32(_SAGE_ATOM_TOKENS) + last_atom = Int32(k_tokens.atoms - 1) + _copy_atom_scales( + cfg, + self.sage_scales_ptr, + self._stage_words(stage_info), + self._sequence_words(stage_info) + Int32(k_tokens.word_offset), + k_tokens.atom_words, + lambda atom: cute.math.min(tile_atom + atom, last_atom), + self.tile_atoms, ) - buffer = self._ring(stage_info.context) - slot_base = Int32(stage_info.stage_idx) * Int32(num_words) - for round_idx in cutlass.range_constexpr((num_words + threads - 1) // threads): - word_idx = warp_grp_thread_idx + Int32(round_idx * threads) - if word_idx < Int32(num_words): - buffer[slot_base + word_idx] = word_value( - word_idx, word_idx // Int32(arr_size), word_idx % Int32(arr_size) - ) + + # -- storage helpers --------------------------------------------------- @cute.jit def _ring(self, context: ResourceContext) -> cutlass.Array: @@ -579,190 +664,35 @@ def _ring(self, context: ResourceContext) -> cutlass.Array: return cutlass.Array( context.smem_base.data_ptr() + self._alloc.offset, dtype=Float32, - shape=(SAGE_K_SCALES_RING_STAGES * self.tile_words,), + shape=(self.pipeline_config.num_stages * self.tile_words,), addrspace=3, ) - # -- word sources ------------------------------------------------------ + # -- consumer steps ---------------------------------------------------- + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=stage_bytes) @cute.jit - def _word_source(self, stage_info: StageInfo) -> WordSource: - """Return ``word_value(word_idx, half, lane_entry)`` for the current tile. - - ``word_idx = half * arr_size + lane_entry`` in the - ``sage_k_scale_words`` layout. - """ - cfg = self.cfg - if cutlass.const_expr(self.route_metadata is None): - # A dense tile gathers ``k_scale`` over its token range; the load - # clamps tokens past the sequence end. - ( - seq_len_kv, - _q_group_idx, - _element_mask_end_idx, - tile_offset_k, - _window_start_idx, - _is_valid_effective_tile, - _is_masked_final_wave, - _tile_is_unmasked, - _rows_are_active, - ) = resolve_keeps_tile_context( - cfg, - stage_info, - inst_id=self.inst_id, - seqlens_kv=self.seqlens_kv, - max_seq_len_kv=self.max_seq_len_kv, - seq_len_q=self.seq_len_q, - q_group_idx=self.q_group_idx, - ) - kv_head_idx, batch_idx = _logical_head_batch( - stage_info, self.h_k_idx, self.b_idx - ) - k_scale_addr = self.scale_tensors.k_scale_ptr.toint() - - def dense_word(word_idx: Int32, half: Int32, lane_entry) -> Float32: - _ = word_idx - return load_k_scale( - cfg, - k_scale_addr, - self.scale_tensors.k_scale_head_stride, - kv_head_idx=kv_head_idx, - batch_idx=batch_idx, - seq_len_kv=seq_len_kv, - kv_token_idx=dense_k_scale_token( - cfg, half, lane_entry, tile_offset_k - ), - ) - - return dense_word - - metadata = self.route_metadata - stage_base = metadata._consumer_stage_base() - if cutlass.const_expr(not self.summary): - # The load warp staged the route's words in the - # ``sage_k_scale_words`` layout behind the route record. - assert metadata.staging_layout.sage_scale_words_word_offset is not None - source = stage_base + Int32( - metadata.staging_layout.sage_scale_words_word_offset - ) - staged = metadata._smem_scales - - def staged_word(word_idx: Int32, half: Int32, lane_entry) -> Float32: - _ = half, lane_entry - return Float32(staged[source + word_idx]) - - return staged_word - - # A proxy route of a mixed-geometry plan stages no scale words; gather - # each summary scale from the route's staged atom origins. - groups = self.groups - num_origins = metadata.staging_layout.num_origin_words - kv_head_idx, batch_idx = _logical_head_batch( - stage_info, self.h_k_idx, self.b_idx + def open_stage(self, stage_info: StageInfo) -> Int32: + """Publish the waited stage's shared-memory offset for the passes that read it.""" + return Int32(self._alloc.offset) + Int32(self.consumer_work_stage) * Int32( + self.tile_words * 4 ) - scale_addr = self.scale_tensors.k_summary_scale_ptr.toint() - scale_head_stride = self.scale_tensors.k_summary_scale_head_stride - summary_seq_len = Int32(cfg.num_proxy_summaries) - log2_k_block_size = Int32(log2_block_size(cfg.sage_k_summary_block_size)) - origins = metadata._smem_words - - def summary_word(word_idx: Int32, half: Int32, lane_entry) -> Float32: - _ = word_idx - atom_idx, token_offset = sage_word_position(cfg, half, lane_entry, groups) - atom_idx = cute.math.min(atom_idx, Int32(num_origins - 1)) - origin = Int32(origins[stage_base + atom_idx]) - return load_k_scale( - cfg, - scale_addr, - scale_head_stride, - kv_head_idx=kv_head_idx, - batch_idx=batch_idx, - seq_len_kv=summary_seq_len, - kv_token_idx=cute.math.max(origin, Int32(0)) + token_offset, - log2_k_block_size=log2_k_block_size, - ) - - return summary_word - - # -- pass-side reads --------------------------------------------------- - - @cute.jit - def open( - self, - stage_info: StageInfo, - scale_arr: cutlass.Array, - factor: Float32 | None, - ): - """Return a pass's view of the tile's words with ``factor`` applied. - - Register form: a rotating copy of the routed array, scaled once. - SMEM form: the lane's half pointer into the waited slot and the - factor, applied per fragment. - """ - if cutlass.const_expr(self.in_smem): - warp_grp_thread_idx = Int32( - _decode_gen_task_cache(stage_info)[_TASK_CACHE_WARP_GRP_THREAD_IDX] - ) - half_base = Int32(self.consumer_work_stage) * Int32( - self.tile_words - ) + _keeps_spatial_half(self.cfg, warp_grp_thread_idx) * Int32( - self.arr_size - ) - half_ptr = self._ring(stage_info.context).data_ptr() + half_base - return half_ptr, factor - arr_size = self.arr_size - words = cutlass.Array(Float32, arr_size, space=cutlass.AddressSpace.rmem) - for entry in cutlass.range_constexpr(arr_size): - words[entry] = Float32(scale_arr[entry]) - if cutlass.const_expr(factor is not None): - scale_pairs_in_place(words, factor, arr_size) - return words - - @cute.jit - def fragment(self, view, fragment_idx: Int32) -> cutlass.Array: - """Return ``factor * sfK`` of one fragment's ``groups`` words. - - The SMEM form uses 16-byte loads, which callers issue ahead of the - fragment's score wait; the register form reads the leading entries - of the rotating array. - """ - groups = self.groups - values = cutlass.Array(Float32, groups, space=cutlass.AddressSpace.rmem) - if cutlass.const_expr(self.in_smem): - half_ptr, factor = view - assert groups % 4 == 0 - words_ptr = half_ptr + fragment_idx * Int32(groups) - for chunk in cutlass.range_constexpr(0, groups, 4): - loaded = (words_ptr + Int32(chunk)).load(count=4, alignment=16) - for elem in cutlass.range_constexpr(4): - values[chunk + elem] = Float32(loaded[elem]) - if cutlass.const_expr(factor is not None): - scale_pairs_in_place(values, factor, groups) - else: - for group_idx in cutlass.range_constexpr(groups): - values[group_idx] = Float32(view[group_idx]) - return values - - @cute.jit - def advance(self, view) -> None: - """Rotate the register array down by one fragment; the SMEM form reads in place.""" - if cutlass.const_expr(not self.in_smem): - groups = self.groups - for entry in cutlass.range_constexpr(self.arr_size - groups): - view[entry] = Float32(view[entry + groups]) @dataclass(kw_only=True) class SageVScalesResource(DecodeGenResourceBase): - """The staged ``sfV`` and ``v_mean`` of the work tile's KV head. + """The staged ``sfV`` and ``v_mean`` of the work tile's sequence. - The pipeline is a one-stage pipeline of the correction warps. + The correction warps copy the sequence's V rows of the prepared scales + into the block at work-tile start; the pipeline is a one-stage pipeline + of those warps. """ cfg: Constexpr[FmhaDecodeConfig] = None - scale_tensors: SageScaleTensors | None = None + sage_scales_ptr: cute.Pointer | None = None h_k_idx: Int32 | None = None b_idx: Int32 | None = None + num_heads_kv: Int32 | None = None _alloc: Constexpr[SmemAllocation | None] = None @property @@ -787,15 +717,22 @@ def get_tmem_requirements(self) -> list[TmemAllocation]: @producer_work(work_attrs=WorkAttr.AUXILIARY) @cute.jit def stage_tile(self, stage_info: StageInfo) -> None: - """Copy the work tile's KV head scales and means into the acquired block.""" + """Copy the work tile's V scales and means into the acquired block.""" cfg = self.cfg task_cache = _decode_gen_task_cache(stage_info) - kv_head_idx, _ = _logical_head_batch(stage_info, self.h_k_idx, self.b_idx) + kv_head_idx, batch_idx = _logical_head_batch( + stage_info, self.h_k_idx, self.b_idx + ) stage_v_channel_scales( cfg, - self.scale_tensors, + self.sage_scales_ptr, self.staged(stage_info.context), - kv_head_idx=kv_head_idx, + sequence_words=sage_sequence_words( + cfg, + batch_idx=batch_idx, + kv_head_idx=kv_head_idx, + num_heads_kv=self.num_heads_kv, + ), thread_idx=task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX], num_threads=cfg.correction_barrier_threads, ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_block_sparse_metadata.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_block_sparse_metadata.py index 22ca7169cf47..f6adced0bc95 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_block_sparse_metadata.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_block_sparse_metadata.py @@ -27,7 +27,7 @@ import cutlass import cutlass.cute as cute -from cutlass import Float32, Int32, Uint32 +from cutlass import Int32, Uint32 from cutlass.experimental.task_scheduling.enums import WorkAttr from cutlass.experimental.task_scheduling.memory import ( ResourceContext, @@ -59,28 +59,10 @@ ResourceVars, _decode_gen_task_cache, _keeps_spatial_half, - _logical_head_batch, _sparse_task_cache_route_begin, _sparse_task_cache_route_count, _warp_broadcast_i32, ) -from .sage_scales import ( - SageScaleTensors, - block_sparse_k_scale_source, - load_k_scale, - sage_scale_arr_size, - sage_word_position, - sage_k_scale_words, - sage_staged_k_scale_words, -) - - -# Keeps staging uses the low four bits for structural KV64 validity. Bit 4 -# carries the conservative prepared summary that token masking can be skipped; -# structural, tail, and causal masking remain independent. The streamed Keeps -# max pass derives its keep words from the token words directly, so the bit is -# currently staged for the consumer but not read. -_SOFTMAX_TOKEN_MASK_IS_FULL_FLAG = 1 << 4 # SWAP origins are at least eight-token aligned, so their low two bits are free @@ -159,10 +141,9 @@ def _kv_retained_route_words( class _BlockSparseSoftmaxStagingLayout: """Layout of the staged cross-warp Softmax metadata payload. - Keeps retains all route origins, a flags word, alignment padding, the - optional K32 token words and, for Sage attention, the route's ``sfK`` - words on their own aligned boundary. KV256 consumers then select the four - words owned by their spatial half. SWAP stores execution-ordered origins + Keeps retains all route origins, a flags word, alignment padding and the + optional K32 token words. KV256 consumers then select the four words + owned by their spatial half. SWAP stores execution-ordered origins followed by optional logical K32 token words, one for each consumer warp. Selected prepared route flags travel in the otherwise-zero low bits of each warp's first aligned origin. @@ -174,7 +155,6 @@ class _BlockSparseSoftmaxStagingLayout: origins_per_warp: int route_flags_word_offset: int | None token_words_word_offset: int | None - sage_scale_words_word_offset: int | None stage_stride_words: int total_words: int @@ -190,14 +170,12 @@ def create( use_keeps_mma_ab: bool, route_layout: _BlockSparseRouteLayout, num_stages: int, - sage_scale_words: int = 0, ) -> "_BlockSparseSoftmaxStagingLayout": """Build a stage-count-dependent layout for Softmax metadata.""" kv_route_size = route_layout.kv_route_size atom_size = route_layout.atom_size has_token_bits = route_layout.has_token_bits - sage_scale_words_word_offset = None if use_keeps_mma_ab: num_origin_words = kv_route_size // atom_size assert num_origin_words <= 4, ( @@ -209,11 +187,7 @@ def create( token_words_word_offset = aligned_payload_words if has_token_bits else None token_words = kv_route_size // 32 if has_token_bits else 0 stage_stride_words = ((aligned_payload_words + token_words + 3) // 4) * 4 - if sage_scale_words > 0: - sage_scale_words_word_offset = stage_stride_words - stage_stride_words += ((sage_scale_words + 3) // 4) * 4 else: - assert sage_scale_words == 0, "Sage scales are staged for Keeps only" softmax_atom_size = min(atom_size, 32) num_origin_words = kv_route_size // softmax_atom_size origins_per_warp = 32 // softmax_atom_size @@ -227,7 +201,6 @@ def create( origins_per_warp=origins_per_warp, route_flags_word_offset=route_flags_word_offset, token_words_word_offset=token_words_word_offset, - sage_scale_words_word_offset=sage_scale_words_word_offset, stage_stride_words=stage_stride_words, total_words=num_stages * stage_stride_words, ) @@ -427,12 +400,11 @@ def prefetch_route( ``target`` selects the route relative to the calling section: ``"head"`` is this instance's HEAD route, ``"first_loop"`` the route of - LOOP iteration 0 (called from HEAD), ``"current_loop"`` the route of - the calling LOOP iteration (no pipelining), and ``"next_loop"`` the - route of the following LOOP iteration. Only the lane-distributed load is issued + LOOP iteration 0 (called from HEAD), and ``"next_loop"`` the route of + the following LOOP iteration. Only the lane-distributed load is issued here; the warp broadcasts happen in ``resolve_route`` so the global - memory latency overlaps the TMA issue of the current route instead of - stalling the load warp. Layouts without one-warp transport keep their + memory latency overlaps the load warp's TMA issue in between instead of + stalling it. Layouts without one-warp transport keep their loads in ``resolve_route`` and get placeholder values here. """ @@ -442,10 +414,6 @@ def prefetch_route( route_idx = Int32(self.inst_id) elif cutlass.const_expr(target == "first_loop"): route_idx = num_insts + Int32(self.inst_id) - elif cutlass.const_expr(target == "current_loop"): - route_idx = (stage_info.loop_offset + Int32(1)) * num_insts + Int32( - self.inst_id - ) else: route_idx = (stage_info.loop_offset + Int32(2)) * num_insts + Int32( self.inst_id @@ -725,10 +693,6 @@ class SmemBlockSparseSoftmaxMetadataResource(DecodeGenResourceBase): data-dependent branch; each consumer receives at most four words through the stable task-local ABI. Runtime route flags carry the conservative FULL summary and, for proxy-capable builds, the route source kind. - - Sage attention adds the route's ``sfK`` words to the staged payload, so - the softmax threads read them from SMEM instead of issuing scattered - global loads before the score wait. """ _task_local_specs: ClassVar[tuple[tuple, ...]] = ( @@ -770,12 +734,8 @@ class SmemBlockSparseSoftmaxMetadataResource(DecodeGenResourceBase): route_metadata: cute.Pointer | None = None staging_layout: Constexpr[_BlockSparseSoftmaxStagingLayout | None] = None route_layout: Constexpr[_BlockSparseRouteLayout | None] = None - h_k_idx: Int32 | None = None - b_idx: Int32 | None = None - scale_tensors: SageScaleTensors | None = None _alloc: Constexpr[SmemAllocation | None] = None _smem_words: cutlass.Array = None - _smem_scales: cutlass.Array = None softmax_origin0_slot: Constexpr[TaskLocalVariable] = ( TaskLocalVariable.uninitialized() ) @@ -809,7 +769,6 @@ def __post_init__(self) -> None: use_keeps_mma_ab=self.cfg.use_keeps_mma_ab, route_layout=self.route_layout, num_stages=self.pipeline_config.num_stages, - sage_scale_words=sage_staged_k_scale_words(self.cfg), ) super().__post_init__() @@ -819,9 +778,6 @@ def _init_placeholder_state(self) -> None: self._smem_words = _placeholder_smem_array( Int32, self.staging_layout.total_words ) - self._smem_scales = _placeholder_smem_array( - Float32, self.staging_layout.total_words - ) def get_smem_requirements(self) -> list[SmemAllocation]: """Allocate one metadata payload per configured pipeline stage.""" @@ -852,13 +808,6 @@ def _create_initial_task_locals( shape=(self.staging_layout.total_words,), addrspace=3, ) - # Float32 view of the same allocation for the staged sfK words. - self._smem_scales = cutlass.Array( - context.smem_base.data_ptr() + self._alloc.offset, - dtype=Float32, - shape=(self.staging_layout.total_words,), - addrspace=3, - ) return {} @producer_work(work_attrs=WorkAttr.AUXILIARY) @@ -1047,25 +996,10 @@ def _store_route_keeps( ) token_word = Uint32(0) - route_token_mask_is_full = cutlass.Boolean(False) if cutlass.const_expr(self.cfg.uses_prepared_score_keep_words): assert self.route_layout.token_words_word_offset is not None assert self.route_metadata is not None if cutlass.const_expr(not uses_one_warp_transport): - gmem_route_flags = Int32(0) - if lane_idx == Int32(0) and route_record_is_valid: - gmem_route_flags = Int32( - self.route_metadata[ - route_record_word_offset - + Int32(self.route_layout.route_flags_word_offset) - ] - ) - gmem_route_flags = _warp_broadcast_i32(gmem_route_flags, 0) - # Prepared bit 0 summarizes the whole route. Staged low bits - # already hold fragment validity, so remap it above them. - route_token_mask_is_full = cutlass.Boolean( - (gmem_route_flags & Int32(_PREPARED_ROUTE_IS_FULL_FLAG)) != Int32(0) - ) if ( lane_idx < Int32(self.route_layout.token_words_per_route) and route_record_is_valid @@ -1088,15 +1022,8 @@ def _store_route_keeps( if cutlass.const_expr(uses_one_warp_transport): if lane_idx == Int32(self.route_layout.route_flags_word_offset): - prepared_route_flags = Int32(resolved_record_word) - route_flags = route_flags | ( - Int32( - (prepared_route_flags & Int32(_PREPARED_ROUTE_IS_FULL_FLAG)) - != Int32(0) - ) - * Int32(_SOFTMAX_TOKEN_MASK_IS_FULL_FLAG) - ) if cutlass.const_expr(self.cfg.use_block_sparse_proxy_routes): + prepared_route_flags = Int32(resolved_record_word) route_flags = route_flags | ( Int32( ( @@ -1111,11 +1038,6 @@ def _store_route_keeps( stage_base + Int32(self.staging_layout.route_flags_word_offset) ] = route_flags elif lane_idx == Int32(0): - if cutlass.const_expr(self.cfg.uses_prepared_score_keep_words): - route_flags = route_flags | ( - Int32(route_token_mask_is_full) - * Int32(_SOFTMAX_TOKEN_MASK_IS_FULL_FLAG) - ) self._smem_words[ stage_base + Int32(self.staging_layout.route_flags_word_offset) ] = route_flags @@ -1138,90 +1060,6 @@ def _store_route_keeps( + Int32(self.staging_layout.token_words_word_offset) + lane_idx ] = Int32(token_word) - if cutlass.const_expr(self.cfg.use_sage_attention): - # Stage the route's ``sfK`` words in the exact geometry. Lane ``l`` - # owns words ``l * words_per_lane`` onward, which lie in one - # fragment of one spatial half, so it resolves one atom origin and - # loads its words at the group stride. Invalid atoms (origin - # ``-1``) are masked; the clamp to token zero only keeps the load - # in bounds. Proxy routes of a mixed-geometry plan stage nothing: - # the softmax warps gather their summary scales behind the score - # wait. - cfg = self.cfg - assert self.staging_layout.sage_scale_words_word_offset is not None - route_is_proxy = cutlass.Boolean(False) - if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): - route_is_proxy = cutlass.Boolean( - ( - _warp_broadcast_i32( - resolved_record_word, - self.route_layout.route_flags_word_offset, - ) - & Int32(_PREPARED_ROUTE_IS_PROXY_FLAG) - ) - != Int32(0) - ) - task_cache = _decode_gen_task_cache(stage_info) - scale_addr, scale_head_stride, scale_seq_len, log2_k_block_size = ( - block_sparse_k_scale_source( - cfg, - self.scale_tensors, - seq_len_kv=Int32(task_cache[_TASK_CACHE_SEQ_LEN_KV]), - route_is_proxy=route_is_proxy, - ) - ) - kv_head_idx, batch_idx = _logical_head_batch( - stage_info, self.h_k_idx, self.b_idx - ) - groups: Constexpr[int] = cfg.sage_k_groups_per_fragment - num_words: Constexpr[int] = sage_k_scale_words(cfg, groups) - words_per_lane: Constexpr[int] = max(1, num_words // 32) - assert words_per_lane * 32 >= num_words and groups % words_per_lane == 0 - group_tokens: Constexpr[int] = cfg.softmax_score_fragment_regs // groups - arr_size: Constexpr[int] = sage_scale_arr_size(cfg, groups) - # Build the view outside the dynamic branch so the store reuses - # its base pointer. - staged = cutlass.Array( - self._smem_scales.data_ptr(), - dtype=Float32, - shape=(self.staging_layout.total_words,), - addrspace=3, - ) - if cutlass.const_expr(not cfg.sage_mixed_k_geometry) or not route_is_proxy: - first_word = lane_idx * Int32(words_per_lane) - atom_idx, token_offset = sage_word_position( - cfg, - first_word // Int32(arr_size), - first_word % Int32(arr_size), - groups, - ) - atom_idx = cute.math.min(atom_idx, Int32(num_origins - 1)) - if cutlass.const_expr(num_origins == 2 and not uses_one_warp_transport): - origin = Int32(resolved_record_word) - if atom_idx != Int32(0): - origin = Int32(resolved_origin1) - else: - origin = Int32( - cute.arch.shuffle_sync(Int32(resolved_record_word), atom_idx) - ) - if first_word < Int32(num_words): - first_token = cute.math.max(origin, Int32(0)) + token_offset - word_base = ( - stage_base - + Int32(self.staging_layout.sage_scale_words_word_offset) - + first_word - ) - for entry in cutlass.range_constexpr(words_per_lane): - staged[word_base + Int32(entry)] = load_k_scale( - cfg, - scale_addr, - scale_head_stride, - kv_head_idx=kv_head_idx, - batch_idx=batch_idx, - seq_len_kv=scale_seq_len, - kv_token_idx=first_token + Int32(entry * group_tokens), - log2_k_block_size=log2_k_block_size, - ) cute.arch.sync_warp() @producer_work @@ -1380,10 +1218,6 @@ def load_route( valid0 = (stored_route_flags >> origin0_idx) & Int32(1) valid1 = (stored_route_flags >> origin1_idx) & Int32(1) route_flags = valid0 | (valid1 << Int32(1)) - if cutlass.const_expr(self.cfg.uses_prepared_score_keep_words): - route_flags = route_flags | ( - stored_route_flags & Int32(_SOFTMAX_TOKEN_MASK_IS_FULL_FLAG) - ) if cutlass.const_expr(self.cfg.use_block_sparse_proxy_routes): route_flags = route_flags | ( stored_route_flags & Int32(_SOFTMAX_ROUTE_IS_PROXY_FLAG) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_p.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_p.py index 80254f62ea52..bb3448d40352 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_p.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_p.py @@ -43,7 +43,6 @@ from ..fmha_decode_config import FmhaDecodeConfig from ..fmha_decode_constants import ( - FP8_P_QUANT_LOG2_SCALE, INT32_SCORE_BIAS, ) from ...placeholder_helpers import _placeholder_smem_array @@ -83,20 +82,14 @@ _pack_float4_to_fp8_e4m3, _pack_float4_to_fp8_e4m3_inline, ) -from .sage_scales import SageKScalesResource +from .sage_scales import load_sage_k_fragment, sage_k_view from .tmem_s import TmemSResource -# Tunable: number of score pairs per streamed fragment whose exponentials run -# as FMA polynomials instead of MUFU. The MUFU issue rate bounds the fragment -# otherwise, while the FMA pipe is nearly idle in the softmax warps. Larger -# shares grow the fragment body and the softmax warps become instruction-fetch -# bound again, so one quarter of the 16 pairs is the measured optimum. -KV_TILE_256_EX2_EMULATED_PAIRS = 4 - -def _pair_uses_ex2_emulation(pair_idx: int, pairs_per_fragment: int) -> bool: - """Spread the emulated pairs evenly across a fragment's score pairs.""" - count = KV_TILE_256_EX2_EMULATED_PAIRS +def _pair_uses_ex2_emulation( + pair_idx: int, pairs_per_fragment: int, count: int +) -> bool: + """Spread ``count`` emulated pairs evenly across a fragment's score pairs.""" pairs = pairs_per_fragment return ((pair_idx + 1) * count) // pairs != (pair_idx * count) // pairs @@ -144,10 +137,7 @@ class SmemPResource(DecodeGenResourceBase): use_variable_seqlens_kv: Constexpr[bool] = False tmem_s_ref: Constexpr[TmemSResource] = None tmem_o_ref: Constexpr[object] = None - # The instance's ``sfK`` resources, shared with the S resource; proxy - # routes read ``sage_summary_k_scales``, everything else ``sage_k_scales``. - sage_k_scales: SageKScalesResource | None = None - sage_summary_k_scales: SageKScalesResource | None = None + # The instance's ``sfK`` words, shared with the S resource. _alloc: Constexpr[SmemAllocation | None] = None _fragment_ready_alloc: Constexpr[SmemAllocation | None] = None _tmem_alloc: Constexpr[TmemAllocation | None] = None @@ -336,7 +326,8 @@ def _exponent_addend( maximum) anchors on zero so its scores exponentiate to zero, not NaN. Callers whose masked rows are skipped or discarded downstream pass ``guards_masked_rows=False``. - - Byte-wide P adds ``log2(448)``, its static FP8 quantization scale. + - Byte-wide P adds the log2 of its quantization scale + (``fp8_p_quant_scale``). - A proxy route adds ``log2`` of its block mass, as the max pass does, so each summary probability carries the tokens it stands for. """ @@ -348,7 +339,7 @@ def _exponent_addend( # Not ``cute.math.fma``: the fused form makes ptxas spill in the callers. addend = Float32(-self.scale_softmax_log2 * safe_new_max) if cutlass.const_expr(self.cfg.use_fp8_pv): - addend += Float32(FP8_P_QUANT_LOG2_SCALE) + addend += Float32(cfg.fp8_p_quant_log2_scale) if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): if route_is_proxy: addend += Float32(cfg.proxy_log2_block_mass) @@ -377,18 +368,20 @@ def compute_sage_p_fragments( *, new_max_arr: cutlass.Array, sage_q_scale: Float32, - sage_scale_arr: cutlass.Array, - sage_summary_scale_arr: cutlass.Array, + sage_k_stage_bytes: Int32, ) -> None: - """Stream every K32 fragment with per-group Sage dequantization scales.""" + """Stream every K32 fragment with per-group Sage dequantization scales. + + ``sage_k_stage_bytes`` is the waited ``sfK`` ring stage, routed from + ``SageKScalesResource.open_stage``. + """ assert self.cfg.use_sage_attention self._compute_p_fragments_impl( stage_info, new_max_arr=new_max_arr, route_is_proxy=cutlass.Boolean(False), sage_q_scale=sage_q_scale, - sage_scale_arr=sage_scale_arr, - sage_summary_scale_arr=sage_summary_scale_arr, + sage_k_stage_bytes=sage_k_stage_bytes, ) @producer_work @@ -417,8 +410,7 @@ def compute_sage_proxy_route_p_fragments( new_max_arr: cutlass.Array, route_flags: Int32, sage_q_scale: Float32, - sage_scale_arr: cutlass.Array, - sage_summary_scale_arr: cutlass.Array, + sage_k_stage_bytes: Int32, ) -> None: """Stream exact or proxy K32 fragments with per-group Sage scales.""" assert self.cfg.use_block_sparse_proxy_routes and self.cfg.use_sage_attention @@ -427,8 +419,7 @@ def compute_sage_proxy_route_p_fragments( new_max_arr=new_max_arr, route_is_proxy=_route_is_proxy(route_flags), sage_q_scale=sage_q_scale, - sage_scale_arr=sage_scale_arr, - sage_summary_scale_arr=sage_summary_scale_arr, + sage_k_stage_bytes=sage_k_stage_bytes, ) @cute.jit @@ -439,8 +430,7 @@ def _compute_p_fragments_impl( new_max_arr: cutlass.Array, route_is_proxy: cutlass.Boolean, sage_q_scale: Float32 | None = None, - sage_scale_arr: cutlass.Array | None = None, - sage_summary_scale_arr: cutlass.Array | None = None, + sage_k_stage_bytes: Int32 | None = None, ) -> None: """Reload, exponentiate, and publish all K32 fragments in a rolled loop. @@ -467,52 +457,45 @@ def _compute_p_fragments_impl( tidx, _, _ = cute.arch.thread_idx() publishes_fragment = (tidx & Int32(31)) == Int32(0) - # ``c * sfQ`` is one factor per lane and tile: the only multiplier of a - # route kind whose max pass returns dequantized scores (no bias), - # otherwise folded into the kind's raw ``sfK`` words. - exact_dequantized: Constexpr[bool] = cfg.sage_scores_dequantized_for( - proxy=False - ) - summary_dequantized: Constexpr[bool] = cfg.sage_scores_dequantized_for( - proxy=True - ) - exact_scales_in_smem: Constexpr[bool] = cfg.sage_k_scales_in_smem_for( - cfg.sage_k_groups_per_fragment - ) - # A mixed geometry selects the resource per fragment on the CTA-uniform - # route kind, unless proxy scores are dequantized and the exact words - # sit in registers: the exact resource then serves both kinds, routing - # ones on a proxy tile. - unified_scales: Constexpr[bool] = ( - cfg.sage_mixed_k_geometry - and summary_dequantized - and not exact_scales_in_smem - ) - mixed_scales: Constexpr[bool] = cfg.sage_mixed_k_geometry and not unified_scales - exp_scale = None - scales_view = None - summary_view = None + # Every score scales by ``c * sfQ * sfK_g``. A geometry whose max pass + # returns dequantized scores (one scale per score) scales by + # ``c * sfQ`` alone and carries no bias; any other loads the tile's + # words in place, a fragment ahead like the scores. + dequantized: Constexpr[bool] = cfg.sage_scores_dequantized() + uses_words: Constexpr[bool] = cfg.use_sage_attention and not dequantized + groups: Constexpr[int] = cfg.sage_k_groups_per_fragment() + # A mixed plan's proxy tile comes back from the max pass dequantized + # (``FmhaDecodeConfig.sage_scores_dequantized``), so its multipliers + # are ``c * sfQ`` alone, with no INT32 bias, in place of the exact + # words. + replaces_proxy_words: Constexpr[bool] = uses_words and cfg.sage_mixed_k_geometry score_bias = None + exp_scale = None + words_view = None row_multipliers = None row_addends = None if cutlass.const_expr(cfg.use_sage_attention): exp_scale = self.scale_softmax_log2 * sage_q_scale - if cutlass.const_expr(exact_dequantized or summary_dequantized): + if cutlass.const_expr(uses_words): + words_view = sage_k_view( + cfg, + stage_info.context, + sage_k_stage_bytes, + Int32( + _decode_gen_task_cache(stage_info)[ + _TASK_CACHE_WARP_GRP_THREAD_IDX + ] + ), + exp_scale, + ) + else: row_multipliers = cutlass.Array( Float32, 1, space=cutlass.AddressSpace.rmem ) row_addends = cutlass.Array(Float32, 1, space=cutlass.AddressSpace.rmem) row_multipliers[0] = exp_scale row_addends[0] = exponent_addend - if cutlass.const_expr(not exact_dequantized): - scales_view = self.sage_k_scales.open( - stage_info, sage_scale_arr, exp_scale - ) - if cutlass.const_expr(mixed_scales and not summary_dequantized): - summary_view = self.sage_summary_k_scales.open( - stage_info, sage_summary_scale_arr, exp_scale - ) - if cutlass.const_expr(unified_scales and cfg.uses_int32_scores): + if cutlass.const_expr(replaces_proxy_words and cfg.uses_int32_scores): score_bias = Float32(INT32_SCORE_BIAS) if route_is_proxy: score_bias = Float32(0.0) @@ -520,7 +503,9 @@ def _compute_p_fragments_impl( # A fragment is scaled before its first exponential, and the next # fragment's TMEM load issues once the scale FFMAs have consumed the # scores, so it reuses their registers and hides behind the - # exponentials. The running sum stays a packed pair until the loop ends. + # exponentials. The next fragment's ``sfK`` words load at the same + # point, so an SMEM read never waits in front of the scale FFMAs. The + # running sum stays a packed pair until the loop ends. num_fragments = cfg.num_softmax_score_fragments last_fragment = Int32(num_fragments - 1) total_sum_pair = (Float32(0.0), Float32(0.0)) @@ -528,64 +513,44 @@ def _compute_p_fragments_impl( pending_scores = cutlass.Array( Float32, fragment_regs, space=cutlass.AddressSpace.rmem ) + pending_words = None + if cutlass.const_expr(uses_words): + pending_words = cutlass.Array( + Float32, groups, space=cutlass.AddressSpace.rmem + ) + self._take_fragment_words(words_view, Int32(0), pending_words) self._load_score_fragment(tmem_base, Int32(0), pending_scores) for fragment_idx in cutlass.range(cfg.num_softmax_score_fragments, unroll=1): fragment = Int32(fragment_idx) prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) for score_idx in cutlass.range_constexpr(fragment_regs): s_arr[score_idx] = pending_scores[score_idx] - - if cutlass.const_expr(mixed_scales): - # Each route kind scales the fragment in place with its own - # resource, so only the scores cross the branch. - if route_is_proxy: - if cutlass.const_expr(summary_dequantized): - self._scale_fragment_pairs( - s_arr, row_addends, row_multipliers, groups=1 - ) - else: - self._scale_fragment_with( - self.sage_summary_k_scales, - summary_view, - fragment, - s_arr, - exponent_addend, - ) - else: - if cutlass.const_expr(exact_dequantized): - self._scale_fragment_pairs( - s_arr, row_addends, row_multipliers, groups=1 - ) - else: - self._scale_fragment_with( - self.sage_k_scales, - scales_view, - fragment, - s_arr, - exponent_addend, - ) - elif cutlass.const_expr(exact_dequantized): + if cutlass.const_expr(dequantized): self._scale_fragment_pairs( s_arr, row_addends, row_multipliers, groups=1 ) else: - fragment_multipliers = None - if cutlass.const_expr(cfg.use_sage_attention): - fragment_multipliers = self.sage_k_scales.fragment( - scales_view, fragment + if cutlass.const_expr(replaces_proxy_words): + # The per-group select on the warp-uniform route kind + # works on a copy: one loop body serves both kinds + # because a loop per kind spills the softmax warps' + # registers, and ``pending_words`` stays read-only as + # the loop-carried prefetch array. + words = cutlass.Array( + Float32, groups, space=cutlass.AddressSpace.rmem ) - self.sage_k_scales.advance(scales_view) + for group_idx in cutlass.range_constexpr(groups): + words[group_idx] = pending_words[group_idx] + if route_is_proxy: + for group_idx in cutlass.range_constexpr(groups): + words[group_idx] = exp_scale + else: + words = pending_words group_multipliers, group_addends = self._fragment_exponent_terms( - fragment_multipliers, - exponent_addend, - groups=cfg.sage_k_groups_per_fragment, - score_bias=score_bias, + words, exponent_addend, groups=groups, score_bias=score_bias ) self._scale_fragment_pairs( - s_arr, - group_addends, - group_multipliers, - groups=cfg.sage_k_groups_per_fragment, + s_arr, group_addends, group_multipliers, groups=groups ) # The last iteration reloads its own fragment so the loop body # stays branch-free; the wait after the loop retires it. @@ -593,6 +558,8 @@ def _compute_p_fragments_impl( if next_fragment > last_fragment: next_fragment = last_fragment self._load_score_fragment(tmem_base, next_fragment, pending_scores) + if cutlass.const_expr(uses_words): + self._take_fragment_words(words_view, next_fragment, pending_words) local_sum_pair = self._exponentiate_fragment_pairs(s_arr) total_sum_pair = cute.arch.add_packed_f32x2(total_sum_pair, local_sum_pair) @@ -635,24 +602,21 @@ def _compute_p_fragments_impl( self.tmem_s_ref.store_p_local_sum(0, total_sum) @cute.jit - def _scale_fragment_with( + def _take_fragment_words( self, - scales: Constexpr[SageKScalesResource], - view, + words_view, fragment: Int32, - s_arr: cutlass.Array, - exponent_addend: Float32, + words: cutlass.Array, ) -> None: - """Turn one fragment's scores into log2 exponents with one route kind's resource.""" - groups = scales.groups - fragment_multipliers = scales.fragment(view, fragment) - scales.advance(view) - group_multipliers, group_addends = self._fragment_exponent_terms( - fragment_multipliers, exponent_addend, groups=groups - ) - self._scale_fragment_pairs( - s_arr, group_addends, group_multipliers, groups=groups - ) + """Copy one fragment's ``c * sfQ * sfK_g`` words into ``words``. + + A constant-indexed register copy, so the in-place loads can issue a + fragment ahead of the FFMAs that read them. + """ + fragment_words = load_sage_k_fragment(self.cfg, words_view, fragment) + groups: Constexpr[int] = self.cfg.sage_k_groups_per_fragment() + for group_idx in cutlass.range_constexpr(groups): + words[group_idx] = fragment_words[group_idx] @cute.jit def _load_score_fragment( @@ -716,7 +680,7 @@ def _fragment_exponent_terms( The multipliers are the resource's ``c * sfQ * sfK_g`` under Sage, else the softmax scale. Biased INT32 scores subtract ``bias * multiplier`` from the row's addend; ``score_bias`` replaces the constant bias when - given (zero on a tile whose scores come back dequantized). + given (zero on a tile whose scores came back dequantized). """ cfg = self.cfg neg_bias = Float32(-INT32_SCORE_BIAS) @@ -779,7 +743,9 @@ def _exponentiate_fragment_pairs( p0 = Float32(s_arr[value_idx]) p1 = Float32(s_arr[value_idx + 1]) if cutlass.const_expr( - _pair_uses_ex2_emulation(pair_idx, pairs_per_fragment) + _pair_uses_ex2_emulation( + pair_idx, pairs_per_fragment, self.cfg.ex2_emulated_pairs + ) ): p0, p1 = _ex2_emulation_packed_f32x2(p0, p1) else: 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 133589603fd0..fe49a7c67592 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 @@ -81,6 +81,7 @@ ) from .helpers_kv_tile_idx import ( _load_runtime_seq_len_kv, + _local_kv_tile_idx_for_section, _num_skipped_kv_tiles, _runtime_clamp_valid_tile_idx, _runtime_last_valid_page_idx, @@ -210,34 +211,6 @@ def _is_load_task_warp_leader(cfg: Constexpr[FmhaDecodeConfig]): return elected and _load_task_warp_rank(cfg) == Int32(0) -@cute.jit -def _local_kv_tile_idx_for_section( - cfg: Constexpr[FmhaDecodeConfig], - stage_info: StageInfo, - inst_id: Constexpr[int], - kv_kind: Constexpr[int], - section: Constexpr[FmhaStage], -) -> Int32: - """Return the local K/V tile index implied by the TS HEAD/LOOP/TAIL cadence. - - The schedule names whether a producer is K0/K1/V0/V1, while the phase - defines where that logical tile sits in the staggered decode pipeline: - - - HEAD produces only initial K tiles. - - LOOP produces V for the current MMA iteration and K for the next one. - - TAIL drains the final V tiles after the last loop iteration. - """ - num_insts_kv = Int32(cfg.num_insts_kv) - if cutlass.const_expr(section == FmhaStage.Head): - return Int32(inst_id) - if cutlass.const_expr(section == FmhaStage.Loop): - base = stage_info.loop_offset * num_insts_kv + Int32(inst_id) - if cutlass.const_expr(kv_kind == KV_KIND_V): - return base - return base + num_insts_kv - return stage_info.loop_end * num_insts_kv + Int32(inst_id) - - @cute.jit def _issue_sparse_page_copies( cfg: Constexpr[FmhaDecodeConfig], @@ -2946,18 +2919,30 @@ def _producer_load_kv_tile_256( route_tma_desc = ( summary_tma_desc if route_is_proxy else sparse_tma_desc ) + # Keep two retained route origins ahead of the copies. A copy may + # write SMEM, so the compiler keeps a later origin load behind it; + # loading each origin right before its copies would serialize every + # block's copies behind its own SMEM load. + route_coord = None + next_route_coord = None + if cutlass.const_expr(cfg.use_block_sparse): + assert sparse_kv_metadata is not None + route_coord = sparse_kv_metadata.route_tma_coordinate( + Int32(0), logical_b_idx + ) + next_route_coord = sparse_kv_metadata.route_tma_coordinate( + Int32(1), logical_b_idx + ) for semantic_block in cutlass.range_constexpr(4): token_coord = Int32(0) storage_coord = logical_b_idx if cutlass.const_expr(cfg.use_block_sparse): - assert sparse_kv_metadata is not None - ( - token_coord, - storage_coord, - ) = sparse_kv_metadata.route_tma_coordinate( - Int32(semantic_block), - logical_b_idx, - ) + token_coord, storage_coord = route_coord + route_coord = next_route_coord + if cutlass.const_expr(semantic_block + 2 < 4): + next_route_coord = sparse_kv_metadata.route_tma_coordinate( + Int32(semantic_block + 2), logical_b_idx + ) physical_block = semantic_block if cutlass.const_expr(kv_kind == KV_KIND_K): physical_block = KV_TILE_256_K_SLOT_FOR_SEMANTIC_ATOM[ diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_corr.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_corr.py index 1a2a04fbd976..9e46fae39bcc 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_corr.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_corr.py @@ -40,7 +40,6 @@ from ...placeholder_helpers import _placeholder_smem_array from ..fmha_decode_config import FmhaDecodeConfig -from ..fmha_decode_constants import FP8_P_QUANT_SCALE from .helpers_common import ( _TASK_CACHE_LANE_IDX, _TASK_CACHE_TMEM_BASE_OFFSET, @@ -122,9 +121,10 @@ class TmemCorrResource(DecodeGenResourceBase): partial_stats_ptr: cute.Pointer = None split_kv_counter_ptr: cute.Pointer = None attention_sinks_ptr: cute.Pointer = None - # The correction task's staged ``sfV`` / ``v_mean`` resource; ``None`` + # The ``sfV`` / ``v_mean`` block the correction task stages per work tile; + # the schedule waits and releases it, the epilogue only reads it. ``None`` # without Sage attention. - sage_v_scales: SageVScalesResource | None = None + sage_v_scales_ref: Constexpr[SageVScalesResource | None] = None seqlens_kv: cute.Pointer = None max_seq_len_kv: Constexpr[int] = 0 seq_len_q: Int32 = None @@ -217,7 +217,7 @@ def _init_placeholder_state(self) -> None: ) self._sage_v_scales = _placeholder_smem_array( Float32, - self.sage_v_scales.entries if self.sage_v_scales is not None else 1, + self.sage_v_scales_ref.entries if self.sage_v_scales_ref is not None else 1, ) def _owns_final_epilogue(self) -> bool: @@ -596,9 +596,9 @@ def _split_partial_scale(self, denominator: Float32) -> Float32: """Return the scale applied to a split-KV partial O before it is stored. The separate reduction kernel receives normalized partials. Fused GMEM - and cluster reductions receive unnormalized partials, with the static - 448 scale of E4M3 P divided out before the narrowing to 16 bits; the - reducer restores it. The 448 scale follows ``pv_dtype``, not the + and cluster reductions receive unnormalized partials, with the E4M3 P + scale (``fp8_p_quant_scale``) divided out before the narrowing to 16 + bits; the reducer restores it. The scale follows ``pv_dtype``, not the Q/K dtype. """ cfg = self.cfg @@ -606,7 +606,7 @@ def _split_partial_scale(self, denominator: Float32) -> Float32: if cutlass.const_expr(cfg.use_separate_reduction_kernel): partial_scale = self._separate_partial_norm_scale(denominator) elif cutlass.const_expr(cfg.use_fp8_pv): - partial_scale = Float32(1.0 / FP8_P_QUANT_SCALE) + partial_scale = Float32(1.0 / cfg.fp8_p_quant_scale) return partial_scale @cute.jit @@ -961,7 +961,7 @@ def _apply_sage_channel_scales( cfg = self.cfg if cutlass.const_expr(not cfg.use_sage_attention): return norm_scale - v_scales, v_means = self.sage_v_scales.channel_scales( + v_scales, v_means = self.sage_v_scales_ref.channel_scales( self._sage_v_scales, first_col=first_col, count=count ) if cutlass.const_expr(cfg.sage_v_mean): @@ -1135,8 +1135,8 @@ def _softmax_output_row_state( ) -> tuple[Int64, Float32]: """Resolve one logical row's output address and normalization once. - With E4M3 P the O accumulator and the row sum both carry the 448 P - scale, which cancels in a direct tail. + With E4M3 P the O accumulator and the row sum both carry the P scale + (``fp8_p_quant_scale``), which cancels in a direct tail. """ cfg = self.cfg attention_sink_h_r = _attention_sink_head_stride(cfg, self.h_r) @@ -1158,9 +1158,9 @@ def _softmax_output_row_state( # helper only after the cross-CTA reduction has completed. norm_scale = self.output_scale * self._safe_norm_rcp(sum_val) if cutlass.const_expr(cfg.use_fp8_pv and cfg.use_split_kv): - # Fused split partials carry O divided by 448 (see + # Fused split partials carry O divided by the FP8 P scale (see # ``_split_partial_scale``); restore it after the FP32 reduction. - norm_scale *= Float32(FP8_P_QUANT_SCALE) + norm_scale *= Float32(cfg.fp8_p_quant_scale) physical_dst_row_idx = _q_physical_output_row_from_logical( cfg, self.h_r, @@ -1333,9 +1333,9 @@ def create_function_variables( if cutlass.const_expr( context is not None and context.smem_base is not None - and self.sage_v_scales is not None + and self.sage_v_scales_ref is not None ): - self._sage_v_scales = self.sage_v_scales.staged(context) + self._sage_v_scales = self.sage_v_scales_ref.staged(context) return {} @producer_work(work_attrs=WorkAttr.AUXILIARY) @@ -3484,22 +3484,35 @@ def _keeps_tail_epilogue( ) return + # Lanes store disjoint output rows, so every store instruction + # touches one cache line per lane, and those line touches bound the + # tail. A 16-bit output chunk is one 32-byte sector (16 columns): + # whole-sector stores halve the touches of 16-byte stores, and plain + # predicated stores (no branch) let the next chunk's TMEM loads issue + # ahead of the current chunk's store. FP8 output keeps 8 columns. + chunk_cols = 8 if cfg.use_fp8_output else 16 + assert (output_pair_regs * 2) % chunk_cols == 0 + # Row strides and lane column bases are multiples of 32 bytes, so + # sector-wide stores are legal exactly when the output base is. + o_misalignment = self.o_ptr.toint() & Int64(31) + store_sectors = valid_output_row & (o_misalignment == Int64(0)) + store_halves = valid_output_row & (o_misalignment != Int64(0)) regs_o_chunk = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) - for store_idx in cutlass.range_constexpr(output_pair_regs // 4): - chunk_col = store_idx * 8 + for store_idx in cutlass.range_constexpr(output_pair_regs * 2 // chunk_cols): + chunk_col = store_idx * chunk_cols # FP16/BF16 or generic FP8 path: load one output chunk from each # tail O stage, combine the two instances, then pack to output type. o0_vals = _keeps_tcgen05_ld( cfg, prims.make_tmem_ptr(base_addr0 + Int32(chunk_col), Float32), - num=8, + num=chunk_cols, offset=keeps_o_ldst_offset, ) if cutlass.const_expr(cfg.num_insts_kv != 1): o1_vals = _keeps_tcgen05_ld( cfg, prims.make_tmem_ptr(base_addr1 + Int32(chunk_col), Float32), - num=8, + num=chunk_cols, offset=keeps_o_ldst_offset, ) prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) @@ -3548,8 +3561,10 @@ def _keeps_tail_epilogue( final_pair1[1], ) else: - final_vals = cutlass.Array(Float32, 8, space=cutlass.AddressSpace.rmem) - for chunk_idx in cutlass.range_constexpr(4): + final_vals = cutlass.Array( + Float32, chunk_cols, space=cutlass.AddressSpace.rmem + ) + for chunk_idx in cutlass.range_constexpr(chunk_cols // 2): reg_base = chunk_idx * 2 if cutlass.const_expr(cfg.num_insts_kv == 1): final_pair = fmul2( @@ -3572,15 +3587,15 @@ def _keeps_tail_epilogue( final_vals, None, first_col=col_base + Int32(chunk_col), - count=8, + count=chunk_cols, row_has_mass=reduced_sum_0 > Float32(0.0), ) - regs_o_chunk = self._pack_final_o_regs(final_vals, None, 8) + regs_o_chunk = self._pack_final_o_regs(final_vals, None, chunk_cols) dst_ptr = cutlass.inttoptr( self.o_ptr.toint() + dst_row_base + dst_col_offset - + Int64(store_idx * (8 if cfg.use_fp8_output else 16)), + + Int64(chunk_col * cfg.o_dtype_bytes), mem_space=1, dtype=Int32, ) @@ -3591,11 +3606,13 @@ def _keeps_tail_epilogue( alignment=8, ) else: - if valid_output_row: + if store_sectors: dst_ptr.store( - regs_o_chunk.data_ptr().load(count=4, alignment=4), - alignment=16, + regs_o_chunk.data_ptr().load(count=8, alignment=4), + alignment=32, ) + if store_halves: + self._store_16bit_output_chunks(dst_ptr, regs_o_chunk, chunk_cols) return @cute.jit @@ -3836,8 +3853,8 @@ def _swaps_wide_tail_epilogue( for pair_idx in cutlass.range_constexpr(output_pair_regs): # Separate reduction includes this split's reciprocal sum; # fused reduction delays normalization until the final merge. - # FP8-P fused reduction removes the 448x P quantization scale - # so the unnormalized numerator fits the 16-bit scratch. + # FP8-P fused reduction removes the P quantization scale so + # the unnormalized numerator fits the 16-bit scratch. scale_base = ((pair_idx % (2 * q_repeats)) // 2) * 2 reg_base = pair_idx * 2 partial_scale0 = ( @@ -4338,7 +4355,7 @@ def _swaps_q8_tail_epilogue( # Store normalized 16-bit O for the standalone reducer, or preserve # unnormalized 16-bit O for fused GMEM/cluster reduction. - # FP8-P fused reduction removes the 448x P quantization scale before + # FP8-P fused reduction removes the P quantization scale before # storing the unnormalized O numerator in the 16-bit scratch. regs_partial_o = cutlass.Array( Int32, 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 78317931721a..715a2ec91bcc 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 @@ -49,9 +49,8 @@ INT32_SCORE_SEED_TILE_LBO, INT32_SCORE_SEED_TILE_SBO, INT32_SCORE_SEED_TILE_WORDS, - SOFTMAX_RESCALE_THRESHOLD_LOG2, ) -from ...tcgen05_compat import tcgen05_mma_ws +from ...tcgen05_compat import tcgen05_ld_32x32b_max, tcgen05_mma_ws from ...placeholder_helpers import ( _placeholder_local_array, _placeholder_smem_array, @@ -97,7 +96,12 @@ from .smem_block_sparse_metadata import ( _swaps_forwards_packed_route_full, ) -from .sage_scales import SageKScalesResource, SageScaleTensors, load_q_scale +from .sage_scales import ( + load_q_scale, + load_sage_k_fragment, + sage_k_view, + sage_sequence_words, +) from .helpers_kv_tile_idx import ( resolve_keeps_tile_context, _kv_tile_is_fully_unmasked_for_q_group, @@ -267,7 +271,8 @@ class TmemSResource(DecodeGenResourceBase): h_k_idx: Int32 | None = None b_idx: Int32 | None = None q_group_idx: Int32 | None = None - scale_tensors: SageScaleTensors | None = None + num_heads_kv: Int32 | None = None + sage_scales_ptr: cute.Pointer | None = None q_ref: Constexpr[MemoryResource | None] = None page_offsets_ref: Constexpr[MemoryResource | None] = None _p_local_sum_arr: cutlass.Array | None = None @@ -280,11 +285,7 @@ class TmemSResource(DecodeGenResourceBase): # instance owns the tile; the other reads it through this reference. score_seed_owner: Constexpr["TmemSResource | None"] = None _seed_alloc: Constexpr[SmemAllocation | None] = None - # The instance's ``sfK`` resources. Proxy routes read - # ``sage_summary_k_scales``, which is ``sage_k_scales`` unless the summary - # scales have their own K block size. - sage_k_scales: SageKScalesResource | None = None - sage_summary_k_scales: SageKScalesResource | None = None + # The instance's ``sfK`` words of every tile and route kind. old_max_arr: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() sum_arr: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() new_max_arr: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() @@ -499,6 +500,15 @@ def _softmax_loop_stage_slot_offset(self, stage_info: StageInfo) -> Int32: return self._stage_slot_offset_from_slot(Int32(stage_info.stage_idx)) return self._stage_slot_offset_from_slot(stage_info.loop_offset % Int32(2)) + @cute.jit + def _softmax_loop_score_tmem_addr(self, stage_info: StageInfo) -> Int32: + """Return the TMEM address of the S slot that LOOP softmax reads.""" + return ( + _decode_gen_task_cache(stage_info)[_TASK_CACHE_TMEM_BASE_OFFSET] + + Int32(self._alloc.offset) + + self._softmax_loop_stage_slot_offset(stage_info) + ) + @cute.jit def _create_initial_task_locals( self, context: ResourceContext | None = None @@ -1019,17 +1029,17 @@ def _softmax_anchor(self, old_max: Float32, tile_max: Float32) -> Float32: Online softmax only requires a common finite reference for P, the running sum, and O; it does not require the exact row maximum. Profiles that defer anchor updates keep the previous reference while - the tile raises it by less than ``SOFTMAX_RESCALE_THRESHOLD_LOG2`` - log2 units, so correction can skip the in-place TMEM O rescale. The - 16-bit P path represents the bounded values above one, and the - numerator and denominator stay in the same scale frame. Larger jumps - still rebase to keep P comfortably in range. + the tile raises it by at most ``softmax_anchor_headroom_log2`` + log2 units, so correction can skip the in-place TMEM O rescale. P + represents the bounded values above one (FP8 P through the headroom of + ``fp8_p_quant_scale``), and the numerator and denominator stay in the + same scale frame. Larger jumps rebase the anchor. """ new_max = cute.math.max(old_max, tile_max, ftz=True) if cutlass.const_expr(self.cfg.defers_softmax_anchor_updates): if old_max != _neg_max_f32(): max_delta_log2 = self.scale_softmax_log2 * (old_max - new_max) - if max_delta_log2 >= Float32(-SOFTMAX_RESCALE_THRESHOLD_LOG2): + if max_delta_log2 >= Float32(-self.cfg.softmax_anchor_headroom_log2): new_max = old_max return new_max @@ -1121,11 +1131,7 @@ def _load_keeps_fragment_impl( cfg = self.cfg task_cache = _decode_gen_task_cache(stage_info) num_s_regs = cfg.num_s_regs_per_thread - base_addr = ( - task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] - + Int32(self._alloc.offset) - + self._softmax_loop_stage_slot_offset(stage_info) - ) + base_addr = self._softmax_loop_score_tmem_addr(stage_info) for load_atom_idx in cutlass.range_constexpr(num_s_regs // 32): atom_col = load_atom_idx * 32 loaded = _keeps_tcgen05_ld( @@ -1435,11 +1441,7 @@ def _compute_softmax_loop_keeps( return old_max_arr, sum_arr, new_max_arr, s_arr if tile_has_valid_scores: - base_addr = ( - task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] - + Int32(self._alloc.offset) - + self._softmax_loop_stage_slot_offset(stage_info) - ) + base_addr = self._softmax_loop_score_tmem_addr(stage_info) # Keep each intrinsic result at the native 32-register atom. TQ128 # uses four consecutive atoms and TQ64 uses two atoms with the same # half-split offset, avoiding a monolithic x64/x128 LLVM intrinsic @@ -1742,11 +1744,7 @@ def _compute_softmax_loop_swaps( # registers. Two TMEM rows cover the two K subtiles. Invalid # odd-tail waves intentionally skip the load and leave S at # -inf so the later P path contributes zero probability. - base_addr = ( - task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] - + Int32(self._alloc.offset) - + self._softmax_loop_stage_slot_offset(stage_info) - ) + base_addr = self._softmax_loop_score_tmem_addr(stage_info) shape = "16x256b" loaded0 = prims.tcgen05_ld( @@ -2233,7 +2231,12 @@ def compute_softmax_loop( @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=sage_q_scale) @cute.jit def load_sage_q_scale(self, stage_info: StageInfo) -> Float32: - """Load ``sfQ`` of the lane's Q row, fixed for the whole work tile.""" + """Load ``sfQ`` of the lane's Q row, fixed for the whole work tile. + + The word is row ``row_idx`` of the work tile's words in the prepared + ``sfQ`` image, one coalesced load per lane; the Sage prepare resolved + the flat scale slot and the last valid row's repeat. + """ cfg = self.cfg assert cfg.use_sage_attention task_cache = _decode_gen_task_cache(stage_info) @@ -2241,20 +2244,17 @@ def load_sage_q_scale(self, stage_info: StageInfo) -> Float32: kv_head_idx, batch_idx = _logical_head_batch( stage_info, self.h_k_idx, self.b_idx ) - q_token_idx, local_head_idx = _q_row_token_and_local_head( - cfg, - self.h_r, - _logical_q_group_idx(cfg, stage_info, self.q_group_idx), - _keeps_row_idx(cfg, warp_grp_thread_idx), - ) return load_q_scale( cfg, - self.scale_tensors.q_scale_ptr.toint(), - self.scale_tensors.q_scale_head_stride, - kv_head_idx=kv_head_idx, - local_head_idx=local_head_idx, - batch_idx=batch_idx, - q_token_idx=q_token_idx, + self.sage_scales_ptr, + sequence_words=sage_sequence_words( + cfg, + batch_idx=batch_idx, + kv_head_idx=kv_head_idx, + num_heads_kv=self.num_heads_kv, + ), + q_group_idx=_logical_q_group_idx(cfg, stage_info, self.q_group_idx), + row_idx=_keeps_row_idx(cfg, warp_grp_thread_idx), ) @consumer_work(returns=("old_max_arr", "sum_arr", "new_max_arr", "s_arr")) @@ -2268,10 +2268,13 @@ def compute_sage_softmax_loop( new_max_arr: cutlass.Array, s_arr: cutlass.Array, sage_q_scale: Float32, - sage_scale_arr: cutlass.Array, - sage_summary_scale_arr: cutlass.Array, + sage_k_stage_bytes: Int32, ) -> tuple[object, object, object, object]: - """Consume S with per-group Sage scales and publish the dequantized max.""" + """Consume S with per-group Sage scales and publish the dequantized max. + + ``sage_k_stage_bytes`` is the waited ``sfK`` ring stage, routed from + ``SageKScalesResource.open_stage``. + """ assert self.cfg.use_sage_attention return self._compute_softmax_loop_keeps_fragments( @@ -2282,8 +2285,7 @@ def compute_sage_softmax_loop( s_arr=s_arr, use_sparse=False, sage_q_scale=sage_q_scale, - sage_scale_arr=sage_scale_arr, - sage_summary_scale_arr=sage_summary_scale_arr, + sage_k_stage_bytes=sage_k_stage_bytes, ) @consumer_work( @@ -2396,8 +2398,7 @@ def _compute_softmax_loop_keeps_fragments( sparse_token_word2: Uint32 | None = None, sparse_token_word3: Uint32 | None = None, sage_q_scale: Float32 | None = None, - sage_scale_arr: cutlass.Array | None = None, - sage_summary_scale_arr: cutlass.Array | None = None, + sage_k_stage_bytes: Int32 | None = None, ) -> tuple[object, object, object, object]: """Mask streamed K32 score fragments in place and reduce their max. @@ -2548,14 +2549,14 @@ def _compute_softmax_loop_keeps_fragments( tail_fragment_mask = Int32(0) tail_lane: Constexpr[int] = 0 shifts_tail: Constexpr[bool] = False - route_is_proxy = cutlass.Boolean(False) - # A mixed plan quantizes proxy summaries with their own K block size, - # so exact and proxy tiles read different ``sfK`` geometries; the - # CTA-uniform route kind selects one per tile. - mixed_scales: Constexpr[bool] = use_sparse and cfg.sage_mixed_k_geometry - # The unrolled unmasked pass serves the exact geometry only. - exact_dequantized: Constexpr[bool] = cfg.sage_scores_dequantized_for( - proxy=False + dequantizes_scores: Constexpr[bool] = cfg.sage_scores_dequantized() + # With LDTM.STAT the unmasked pass takes each fragment's maximum, or + # each Sage scale group's, from the TMEM load itself. INT32 scores (the + # load reduces FP32 values) and one-scale-per-score geometries, which + # store their dequantized scores back, keep the register reduction. + groups: Constexpr[int] = cfg.sage_k_groups_per_fragment() + loads_group_maxima: Constexpr[bool] = ( + cfg.uses_ldtm_stat and not cfg.uses_int32_scores and not dequantizes_scores ) if cutlass.const_expr(use_sparse and cfg.use_block_sparse_proxy_routes): # A warp-uniform route kind keeps the fold and the load/store @@ -2601,18 +2602,17 @@ def _compute_softmax_loop_keeps_fragments( warp_scores_are_unmasked and tail_fragment_mask == Int32(0) ) ) + # A mixed plan reads each route kind's ``sfK`` geometry; the CTA-uniform + # route kind selects one per tile. Proxy tiles take the rolled pass, so + # the summary geometry's fold is emitted once rather than per unrolled + # fragment. + mixed_scales: Constexpr[bool] = use_sparse and cfg.sage_mixed_k_geometry if cutlass.const_expr(mixed_scales): - # Proxy tiles take the rolled pass, so the summary geometry's fold - # is emitted once rather than per unrolled fragment. warp_scores_are_unmasked = cutlass.Boolean( warp_scores_are_unmasked and not route_is_proxy ) - score_tmem_addr = ( - task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] - + Int32(self._alloc.offset) - + self._softmax_loop_stage_slot_offset(stage_info) - ) + score_tmem_addr = self._softmax_loop_score_tmem_addr(stage_info) max_chains = cutlass.Array(Float32, 4, space=cutlass.AddressSpace.rmem) max_identity = _neg_max_f32() if cutlass.const_expr(cfg.use_sage_attention): @@ -2622,10 +2622,20 @@ def _compute_softmax_loop_keeps_fragments( scales_view = None summary_view = None if cutlass.const_expr(cfg.use_sage_attention): - scales_view = self.sage_k_scales.open(stage_info, sage_scale_arr, None) + warp_grp_thread_idx = Int32( + _decode_gen_task_cache(stage_info)[_TASK_CACHE_WARP_GRP_THREAD_IDX] + ) + scales_view = sage_k_view( + cfg, stage_info.context, sage_k_stage_bytes, warp_grp_thread_idx, None + ) if cutlass.const_expr(mixed_scales): - summary_view = self.sage_summary_k_scales.open( - stage_info, sage_summary_scale_arr, None + summary_view = sage_k_view( + cfg, + stage_info.context, + sage_k_stage_bytes, + warp_grp_thread_idx, + None, + proxy_kind=True, ) if warp_scores_are_unmasked: @@ -2633,61 +2643,81 @@ def _compute_softmax_loop_keeps_fragments( if cutlass.const_expr(cfg.use_sage_attention): # Issued ahead of the score load so an SMEM read hides # behind the TMEM wait. - fragment_scales = self.sage_k_scales.fragment( - scales_view, Int32(fragment_idx) + fragment_scales = load_sage_k_fragment( + cfg, scales_view, Int32(fragment_idx) ) - loaded = _keeps_tcgen05_ld( - cfg, - prims.make_tmem_ptr( - score_tmem_addr + Int32(fragment_idx * fragment_regs), - Float32, - ), - num=fragment_regs, - offset=cfg.tile_size_kv // 2, - ) - prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) - if cutlass.const_expr(exact_dequantized): - scores = cutlass.Array( - Float32, fragment_regs, space=cutlass.AddressSpace.rmem - ) - for score_idx in cutlass.range_constexpr(fragment_regs): - scores[score_idx] = Float32(loaded[score_idx]) - self._dequantize_fragment_scores( - scores, fragment_scales, cfg.sage_k_groups_per_fragment - ) - self.sage_k_scales.advance(scales_view) - for score_idx in cutlass.range_constexpr(fragment_regs): - chain_idx: Constexpr[int] = score_idx % 4 + fragment_addr = score_tmem_addr + Int32(fragment_idx * fragment_regs) + if cutlass.const_expr(loads_group_maxima): + # One load per scale group returns the group's columns and + # their maximum; the P pass reloads the scores. + group_cols: Constexpr[int] = fragment_regs // groups + maxima: tuple = () + for group_idx in cutlass.range_constexpr(groups): + group_load = tcgen05_ld_32x32b_max( + fragment_addr + Int32(group_idx * group_cols), group_cols + ) + maxima += (group_load[group_cols],) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + if cutlass.const_expr(cfg.use_sage_attention): + self._fold_sage_group_maxima( + max_chains, + maxima, + fragment_scales=fragment_scales, + chain_base=fragment_idx * groups, + groups=groups, + ) + else: + chain_idx: Constexpr[int] = fragment_idx % 4 max_chains[chain_idx] = cute.math.max( - max_chains[chain_idx], Float32(scores[score_idx]), ftz=True + max_chains[chain_idx], Float32(maxima[0]), ftz=True ) - _keeps_tcgen05_st( + else: + loaded = _keeps_tcgen05_ld( cfg, - prims.make_tmem_ptr( - score_tmem_addr + Int32(fragment_idx * fragment_regs), - Float32, - ), - scores.data_ptr().load(count=fragment_regs, alignment=4), + prims.make_tmem_ptr(fragment_addr, Float32), + num=fragment_regs, offset=cfg.tile_size_kv // 2, ) - elif cutlass.const_expr(cfg.use_sage_attention): - self._fold_sage_fragment_max( - max_chains, - loaded, - fragment_scales=fragment_scales, - chain_base=fragment_idx * cfg.sage_k_groups_per_fragment, - groups=cfg.sage_k_groups_per_fragment, - ) - self.sage_k_scales.advance(scales_view) - else: - for score_idx in cutlass.range_constexpr(fragment_regs): - chain_idx: Constexpr[int] = score_idx % 4 - max_chains[chain_idx] = cute.math.max( - max_chains[chain_idx], - Float32(loaded[score_idx]), - ftz=True, + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + if cutlass.const_expr(dequantizes_scores): + scores = cutlass.Array( + Float32, fragment_regs, space=cutlass.AddressSpace.rmem + ) + for score_idx in cutlass.range_constexpr(fragment_regs): + scores[score_idx] = Float32(loaded[score_idx]) + self._dequantize_fragment_scores( + scores, fragment_scales, groups ) - if cutlass.const_expr(exact_dequantized): + for score_idx in cutlass.range_constexpr(fragment_regs): + chain_idx: Constexpr[int] = score_idx % 4 + max_chains[chain_idx] = cute.math.max( + max_chains[chain_idx], + Float32(scores[score_idx]), + ftz=True, + ) + _keeps_tcgen05_st( + cfg, + prims.make_tmem_ptr(fragment_addr, Float32), + scores.data_ptr().load(count=fragment_regs, alignment=4), + offset=cfg.tile_size_kv // 2, + ) + elif cutlass.const_expr(cfg.use_sage_attention): + self._fold_sage_fragment_max( + max_chains, + loaded, + fragment_scales=fragment_scales, + chain_base=fragment_idx * groups, + groups=groups, + ) + else: + for score_idx in cutlass.range_constexpr(fragment_regs): + chain_idx: Constexpr[int] = score_idx % 4 + max_chains[chain_idx] = cute.math.max( + max_chains[chain_idx], + Float32(loaded[score_idx]), + ftz=True, + ) + if cutlass.const_expr(dequantizes_scores): prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) cute.arch.fence_view_async_tmem_store() else: @@ -2707,45 +2737,41 @@ def _compute_softmax_loop_keeps_fragments( visible_end - fragment_token_base, fragment_regs=fragment_regs, ) - # The masked pass is one rolled loop per route kind. A mixed - # geometry selects the loop per tile: a per-fragment selection - # would be if-converted and run both folds on every tile. + mask_args = dict( + may_hold_tail=shifts_tail, + tail_fragment_mask=tail_fragment_mask, + tail_shift=tail_shift, + tail_lane=tail_lane, + ) + # A mixed plan runs one rolled loop per route kind: selecting the + # geometry per fragment would be if-converted and run both folds. if cutlass.const_expr(mixed_scales): if route_is_proxy: self._mask_score_fragments( score_tmem_addr, keep_words, max_chains, - proxy_kind=True, scales_view=summary_view, - may_hold_tail=shifts_tail, - tail_fragment_mask=tail_fragment_mask, - tail_shift=tail_shift, - tail_lane=tail_lane, + proxy_kind=True, + **mask_args, ) else: self._mask_score_fragments( score_tmem_addr, keep_words, max_chains, - proxy_kind=False, scales_view=scales_view, - may_hold_tail=shifts_tail, - tail_fragment_mask=tail_fragment_mask, - tail_shift=tail_shift, - tail_lane=tail_lane, + proxy_kind=False, + **mask_args, ) else: self._mask_score_fragments( score_tmem_addr, keep_words, max_chains, - proxy_kind=False, scales_view=scales_view, - may_hold_tail=shifts_tail, - tail_fragment_mask=tail_fragment_mask, - tail_shift=tail_shift, - tail_lane=tail_lane, + proxy_kind=False, + **mask_args, ) prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) cute.arch.fence_view_async_tmem_store() @@ -2780,43 +2806,37 @@ def _mask_score_fragments( keep_words: cutlass.Array, max_chains: cutlass.Array, *, - proxy_kind: Constexpr[bool], scales_view, + proxy_kind: Constexpr[bool], may_hold_tail: Constexpr[bool], tail_fragment_mask: Int32, tail_shift: Float32, tail_lane: Constexpr[int], ) -> None: - """Mask, fold and write back every fragment with one route kind's scales. + """Mask, fold and write back every fragment in one rolled loop. - ``scales_view`` is the opened view of the resource ``proxy_kind`` - selects; the resource is read from ``self`` because a local would be - flattened as a loop-carried value. The keep words rotate down one - entry per fragment, so the body reads ``keep_words[0]``. + ``scales_view`` is the waited ``sfK`` stage opened in the geometry of + the route kind ``proxy_kind`` names (``sage_k_view``). The keep words + rotate down one entry per fragment, so the body reads + ``keep_words[0]``. Cleared keep bits become ``-inf`` with Sage, else ``-FLT_MAX``. The fragment holding a proxy route's ragged final summary adds the mass shortfall to that score; Sage divides the shift (already divided by ``sfQ``) by the tail group's ``sfK`` with an approximate reciprocal, - ample for a shift far below the score resolution. A route kind whose + ample for a shift far below the score resolution. A geometry whose scores are dequantized writes the dequantized scores back. """ cfg = self.cfg num_fragments = cfg.num_softmax_score_fragments fragment_regs = cfg.softmax_score_fragment_regs - groups: Constexpr[int] = cfg.sage_k_groups_per_fragment_for(proxy=proxy_kind) - dequantize_scores: Constexpr[bool] = cfg.sage_scores_dequantized_for( - proxy=proxy_kind - ) + groups: Constexpr[int] = cfg.sage_k_groups_per_fragment(proxy_kind) + dequantize_scores: Constexpr[bool] = cfg.sage_scores_dequantized(proxy_kind) for fragment in cutlass.range(num_fragments, unroll=1): fragment_scales = None - if cutlass.const_expr(cfg.use_sage_attention and proxy_kind): - fragment_scales = self.sage_summary_k_scales.fragment( - scales_view, Int32(fragment) - ) - elif cutlass.const_expr(cfg.use_sage_attention): - fragment_scales = self.sage_k_scales.fragment( - scales_view, Int32(fragment) + if cutlass.const_expr(cfg.use_sage_attention): + fragment_scales = load_sage_k_fragment( + cfg, scales_view, Int32(fragment), proxy_kind ) keep_word = Uint32(keep_words[0]) fragment_addr = score_tmem_addr + Int32(fragment) * Int32(fragment_regs) @@ -2880,10 +2900,6 @@ def _mask_score_fragments( masked_scores.data_ptr().load(count=fragment_regs, alignment=4), offset=cfg.tile_size_kv // 2, ) - if cutlass.const_expr(cfg.use_sage_attention and proxy_kind): - self.sage_summary_k_scales.advance(scales_view) - elif cutlass.const_expr(cfg.use_sage_attention): - self.sage_k_scales.advance(scales_view) for entry in cutlass.range_constexpr(num_fragments - 1): keep_words[entry] = Uint32(keep_words[entry + 1]) @@ -2949,70 +2965,84 @@ def _fold_sage_fragment_max( chain_base: Constexpr[int], groups: Constexpr[int], ) -> None: - """Fold one fragment's dequantized group maxima into the max chains. + """Reduce one fragment's scale groups on the quantized scores and fold them. - Each scale group is reduced on the quantized scores, then - ``(group_max - bias) * sfK_g`` is folded; the caller applies ``sfQ`` - to the tile maximum. Groups of four or more scores reduce over four - chains; masked ``-inf`` scores need no special case. Group maxima - leave in pairs through a packed bias add (exact on the biased scores' - unit spacing) and a packed scale multiply. Group ``g`` folds into - chain ``(chain_base + g) % 4``. + Groups of four or more scores reduce over four chains; masked ``-inf`` + scores need no special case. ``_fold_sage_group_maxima`` dequantizes + and folds the group maxima. """ - cfg = self.cfg - group_regs = cfg.softmax_score_fragment_regs // groups - width: Constexpr[int] = 1 if groups == 1 else 2 - assert groups % width == 0 - for group_base in cutlass.range_constexpr(0, groups, width): - maxima: tuple = () - for elem in cutlass.range_constexpr(width): - first: Constexpr[int] = (group_base + elem) * group_regs - if cutlass.const_expr(group_regs < 4): - group_max = Float32(scores[first]) - for score_elem in cutlass.range_constexpr(1, group_regs): - group_max = cute.math.max( - group_max, Float32(scores[first + score_elem]), ftz=True - ) - else: - chains = cutlass.Array(Float32, 4, space=cutlass.AddressSpace.rmem) - for chain_idx in cutlass.range_constexpr(4): - chains[chain_idx] = Float32(scores[first + chain_idx]) - for score_elem in cutlass.range_constexpr(4, group_regs): - chain_idx: Constexpr[int] = score_elem % 4 - chains[chain_idx] = cute.math.max( - chains[chain_idx], - Float32(scores[first + score_elem]), - ftz=True, - ) + group_regs = self.cfg.softmax_score_fragment_regs // groups + maxima: tuple = () + for group_idx in cutlass.range_constexpr(groups): + first: Constexpr[int] = group_idx * group_regs + if cutlass.const_expr(group_regs < 4): + group_max = Float32(scores[first]) + for score_elem in cutlass.range_constexpr(1, group_regs): group_max = cute.math.max( - cute.math.max(chains[0], chains[1], ftz=True), - cute.math.max(chains[2], chains[3], ftz=True), - ftz=True, + group_max, Float32(scores[first + score_elem]), ftz=True ) - maxima += (group_max,) - if cutlass.const_expr(width == 1): - scaled_one = Float32(maxima[0]) - if cutlass.const_expr(cfg.uses_int32_scores): - scaled_one = scaled_one - Float32(INT32_SCORE_BIAS) - scaled = (scaled_one * Float32(fragment_scales[group_base]),) else: - scaled = maxima - if cutlass.const_expr(cfg.uses_int32_scores): - scaled = cute.arch.add_packed_f32x2( - scaled, (Float32(-INT32_SCORE_BIAS), Float32(-INT32_SCORE_BIAS)) + chains = cutlass.Array(Float32, 4, space=cutlass.AddressSpace.rmem) + for chain_idx in cutlass.range_constexpr(4): + chains[chain_idx] = Float32(scores[first + chain_idx]) + for score_elem in cutlass.range_constexpr(4, group_regs): + chain_idx: Constexpr[int] = score_elem % 4 + chains[chain_idx] = cute.math.max( + chains[chain_idx], + Float32(scores[first + score_elem]), + ftz=True, ) - scaled = fmul2( - scaled, - ( - Float32(fragment_scales[group_base]), - Float32(fragment_scales[group_base + 1]), - ), + group_max = cute.math.max( + cute.math.max(chains[0], chains[1], ftz=True), + cute.math.max(chains[2], chains[3], ftz=True), + ftz=True, ) - for elem in cutlass.range_constexpr(width): - scaled_max = Float32(scaled[elem]) + maxima += (group_max,) + self._fold_sage_group_maxima( + max_chains, + maxima, + fragment_scales=fragment_scales, + chain_base=chain_base, + groups=groups, + ) + + @cute.jit + def _fold_sage_group_maxima( + self, + max_chains: cutlass.Array, + maxima: tuple, + *, + fragment_scales: cutlass.Array, + chain_base: Constexpr[int], + groups: Constexpr[int], + ) -> None: + """Fold one fragment's dequantized group maxima into the max chains. + + Folds ``(group_max - bias) * sfK_g``; the caller applies ``sfQ`` to + the tile maximum. Group maxima leave in pairs through a packed bias + add (exact on the biased scores' unit spacing) and a packed scale + multiply; a single group pairs with itself. Group ``g`` folds into + chain ``(chain_base + g) % 4``. + """ + cfg = self.cfg + for group_base in cutlass.range_constexpr(0, groups, 2): + second: Constexpr[int] = min(group_base + 1, groups - 1) + scaled = (Float32(maxima[group_base]), Float32(maxima[second])) + if cutlass.const_expr(cfg.uses_int32_scores): + scaled = cute.arch.add_packed_f32x2( + scaled, (Float32(-INT32_SCORE_BIAS), Float32(-INT32_SCORE_BIAS)) + ) + scaled = fmul2( + scaled, + ( + Float32(fragment_scales[group_base]), + Float32(fragment_scales[second]), + ), + ) + for elem in cutlass.range_constexpr(second - group_base + 1): fold_chain: Constexpr[int] = (chain_base + group_base + elem) % 4 max_chains[fold_chain] = cute.math.max( - max_chains[fold_chain], scaled_max, ftz=True + max_chains[fold_chain], Float32(scaled[elem]), ftz=True ) @consumer_work(returns=("old_max_arr", "sum_arr", "new_max_arr", "s_arr")) @@ -3033,8 +3063,7 @@ def compute_sage_block_sparse_softmax_loop( sparse_token_word2: Uint32, sparse_token_word3: Uint32, sage_q_scale: Float32, - sage_scale_arr: cutlass.Array, - sage_summary_scale_arr: cutlass.Array, + sage_k_stage_bytes: Int32, ) -> tuple[object, object, object, object]: """Consume one routed S payload with per-group Sage scales.""" @@ -3054,8 +3083,7 @@ def compute_sage_block_sparse_softmax_loop( sparse_token_word2=sparse_token_word2, sparse_token_word3=sparse_token_word3, sage_q_scale=sage_q_scale, - sage_scale_arr=sage_scale_arr, - sage_summary_scale_arr=sage_summary_scale_arr, + sage_k_stage_bytes=sage_k_stage_bytes, ) @consumer_work(returns=("old_max_arr", "sum_arr", "new_max_arr", "s_arr")) 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 49cd91cde855..7cea30d06cb4 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 @@ -1085,37 +1085,39 @@ def get_domain(self, tile_coord: cute.Coord) -> cutlass.Int32 | int: # TAIL: V(last-1) + V(last) # Dense paged-KV consumes page offsets produced by PageTableTask. Sparse # paged-KV instead consumes physical page IDs retained with its prepared route. +# A Sage plan's load warp also fills each instance's ``sfK`` ring with +# ``cp.async`` copies right after the K tile the words belong to. # ====================================================================== +def _prefetch_sparse_route( + sparse_kv_metadata: MemoryResource | None, + target: str, +) -> tuple[Any, Any] | None: + """Issue one prepared route-record load; dense profiles get ``None``. + + Route records are loaded one resolution ahead so their global-memory + latency overlaps other load-warp work instead of stalling it. Callers + pass the result to the instance's next ``_resolve_and_store_sparse_route``, + the way ``_staged_kv_load`` threads its cached page IDs. + """ + + if sparse_kv_metadata is None: + return None + return sparse_kv_metadata.prefetch_route(target=target) + + def _resolve_and_store_sparse_route( sparse_kv_metadata: MemoryResource | None, section: FmhaStage, - prefetch: tuple[Any, Any] | None = None, - *, - pipeline: bool = True, -) -> tuple[tuple[Any, Any, Any, Any] | None, tuple[Any, Any] | None]: + prefetch: tuple[Any, Any] | None, +) -> tuple[Any, Any, Any, Any] | None: """Resolve one prepared route and retain it for the matching K/V pair. - Returns ``(route, prefetch)``. With ``pipeline`` the record load is issued - one resolution ahead: HEAD loads its own record immediately, and every - resolution issues the load for the next one (LOOP iteration 0 from HEAD, - iteration i + 1 from iteration i) before the caller's K TMA burst, so the - global-memory latency overlaps that issue instead of stalling the load - warp. Callers pass the returned ``prefetch`` back into the next resolution - of the same instance, the way ``_staged_kv_load`` threads its cached page - IDs. Without ``pipeline`` the record is loaded where it is resolved and no - state is returned; the split-ring load variants use this because the - pipelined form measured slower for them. Dense profiles pass ``None`` and - get ``(None, None)``. + ``prefetch`` is the record load issued earlier by + ``_prefetch_sparse_route``. Dense profiles pass ``None`` and get ``None``. """ if sparse_kv_metadata is None: - return None, None - if not pipeline: - prefetch = sparse_kv_metadata.prefetch_route( - target="head" if section == FmhaStage.Head else "current_loop" - ) - elif section == FmhaStage.Head: - prefetch = sparse_kv_metadata.prefetch_route(target="head") + return None assert prefetch is not None prefetched_record_word, prefetched_record_offset = prefetch ( @@ -1134,18 +1136,12 @@ def _resolve_and_store_sparse_route( resolved_atom_validity=resolved_atom_validity, route_record_word_offset=route_record_word_offset, ) - next_prefetch = None - if pipeline: - next_prefetch = sparse_kv_metadata.prefetch_route( - target="first_loop" if section == FmhaStage.Head else "next_loop" - ) - route = ( + return ( resolved_record_word, resolved_origin1, resolved_atom_validity, route_record_word_offset, ) - return route, next_prefetch def _publish_sparse_softmax_route( @@ -1173,6 +1169,32 @@ def _publish_sparse_softmax_route( sparse_softmax_metadata.commit() +def _publish_k_scales( + sage_k_scales: MemoryResource | None, + route: tuple[Any, Any, Any, Any] | None, + section: FmhaStage, +) -> None: + """Fill one ``sfK`` ring stage for the K tile just loaded. + + The words are the tile's atoms of the prepared scales, a block-sparse + route's or a dense tile's; the ring's ``AsyncLoad`` commit completes the + stage once the copies land. + """ + + if sage_k_scales is None: + return + sage_k_scales.acquire() + if route is None: + sage_k_scales.copy_tile(section=section) + else: + resolved_record_word, resolved_origin1, _, _ = route + sage_k_scales.copy_route( + resolved_record_word=resolved_record_word, + resolved_origin1=resolved_origin1, + ) + sage_k_scales.commit() + + def create_load_task( smem_q: MemoryResource, smem_kv: MemoryResource, @@ -1186,12 +1208,18 @@ def create_load_task( sparse_kv_metadata1: MemoryResource | None = None, sparse_softmax_metadata0: MemoryResource | None = None, sparse_softmax_metadata1: MemoryResource | None = None, + sage_k_scales0: MemoryResource | None = None, + sage_k_scales1: MemoryResource | None = None, warp_idx: int | None = None, num_warps: int | None = None, task_class: type[DecodeGenTask] = DecodeGenTask, **kw: TaskKwarg, ) -> Task: - """Create the shared-KV load task and optional page-offset dependency.""" + """Create the shared-KV load task and optional page-offset dependency. + + ``sage_k_scales0/1`` are a Sage plan's ``sfK`` rings, which the task + fills after each K tile. + """ hold_page_window = _can_hold_native_page_window(cfg, smem_page_offsets) head_labels = ("load_k0",) if cfg.num_insts_kv == 1 else ("load_k0", "load_k1") tail_labels = ("load_v0",) if cfg.num_insts_kv == 1 else ("load_v0", "load_v1") @@ -1205,18 +1233,20 @@ def load_schedule_body( sparse_kv_metadata1: MemoryResource | None = None, sparse_softmax_metadata0: MemoryResource | None = None, sparse_softmax_metadata1: MemoryResource | None = None, + sage_k_scales0: MemoryResource | None = None, + sage_k_scales1: MemoryResource | None = None, ) -> None: """Build the shared-KV load cadence for HEAD, LOOP, and TAIL.""" smem_q.init_load_state() smem_kv.init_load_state() - for sparse_resource in ( + for resource in ( sparse_kv_metadata0, sparse_kv_metadata1, sparse_softmax_metadata0, sparse_softmax_metadata1, ): - if sparse_resource is not None: - sparse_resource.init_load_state() + if resource is not None: + resource.init_load_state() cached_page_ids = None if smem_page_offsets is not None: if cfg.num_head_dim_stages_kv > 1 and not cfg.uses_scattered_page_route: @@ -1238,12 +1268,24 @@ def _kv_load(label: str, section: FmhaStage) -> None: cached_page_ids=cached_page_ids, ) - route_metadata_by_label = { - "load_k0": (sparse_kv_metadata0, sparse_softmax_metadata0), - "load_k1": (sparse_kv_metadata1, sparse_softmax_metadata1), + # The resources that follow each instance's K tile: its route metadata + # and Softmax route stage of a block-sparse plan, and its ``sfK`` ring + # of a Sage plan. + k_resources_by_label = { + "load_k0": (sparse_kv_metadata0, sparse_softmax_metadata0, sage_k_scales0), + "load_k1": (sparse_kv_metadata1, sparse_softmax_metadata1, sage_k_scales1), } - # HEAD: load Q once, then prefetch K for each active instance. + # HEAD: issue every instance's route-record load before Q, so their + # global-memory latency overlaps the wait for the previous work tile + # to release Q and the Q load itself. One cadence serves dense and + # block-sparse loads: the route helpers skip a dense profile's + # ``None`` metadata resources, and the K-scale publish fills a Sage + # plan's ring from the dense tile or from the route. + prefetch_by_label = { + label: _prefetch_sparse_route(k_resources_by_label[label][0], "head") + for label in head_labels + } smem_q.acquire() smem_q.tma_load() smem_q.commit() @@ -1256,22 +1298,21 @@ def _kv_load(label: str, section: FmhaStage) -> None: smem_page_offsets.wait() else: _page_offsets_consume(smem_page_offsets) - # Dense profiles have no route metadata: the resolve and publish - # helpers are no-ops for ``None`` resources, so one cadence serves - # both dense and block-sparse loads. - prefetch_by_label = {} head_routes = [] for label in head_labels: - kv_metadata, softmax_metadata = route_metadata_by_label[label] - route, prefetch_by_label[label] = _resolve_and_store_sparse_route( - kv_metadata, FmhaStage.Head + kv_metadata, softmax_metadata, sage_k_scales = k_resources_by_label[label] + route = _resolve_and_store_sparse_route( + kv_metadata, FmhaStage.Head, prefetch_by_label[label] ) + # Each resolution issues its successor's record load ahead of its + # K copies, which gives the load the longest lead time. + prefetch_by_label[label] = _prefetch_sparse_route(kv_metadata, "first_loop") _kv_load(label, FmhaStage.Head) - if route is not None: - head_routes.append((softmax_metadata, route)) + head_routes.append((softmax_metadata, route, sage_k_scales)) # Issue all K tiles before a metadata FIFO can backpressure the load warp. - for softmax_metadata, route in head_routes: + for softmax_metadata, route, sage_k_scales in head_routes: _publish_sparse_softmax_route(softmax_metadata, route) + _publish_k_scales(sage_k_scales, route, FmhaStage.Head) # LOOP: each iter prefetches the full ``num_insts_kv`` K/V pair set. # When P aliases the consumed S columns, MMA must consume each V/P pair @@ -1285,20 +1326,22 @@ def _kv_load(label: str, section: FmhaStage) -> None: loop_labels = ("load_k0", "load_v0", "load_k1", "load_v1") with domain_loop(0, domain, 1, unroll=1): # Generic V-first profiles consume their retained route before the - # matching K label replaces it. - loop_routes = [] + # matching K label replaces it. Each route goes to its Softmax + # consumer right after its own K issue instead of waiting for the + # other instance's loads. for label in loop_labels: - kv_metadata, softmax_metadata = route_metadata_by_label.get( - label, (None, None) + kv_metadata, softmax_metadata, sage_k_scales = k_resources_by_label.get( + label, (None, None, None) ) - route, prefetch_by_label[label] = _resolve_and_store_sparse_route( + route = _resolve_and_store_sparse_route( kv_metadata, FmhaStage.Loop, prefetch_by_label.get(label) ) + prefetch_by_label[label] = _prefetch_sparse_route( + kv_metadata, "next_loop" + ) _kv_load(label, FmhaStage.Loop) - if route is not None: - loop_routes.append((softmax_metadata, route)) - for sparse_softmax_metadata, route in loop_routes: - _publish_sparse_softmax_route(sparse_softmax_metadata, route) + _publish_sparse_softmax_route(softmax_metadata, route) + _publish_k_scales(sage_k_scales, route, FmhaStage.Loop) # TAIL: load the final V tile for each active instance's last BMM2. for label in tail_labels: @@ -1315,6 +1358,8 @@ def load_schedule( sparse_kv_metadata1: MemoryResource | None, sparse_softmax_metadata0: MemoryResource | None, sparse_softmax_metadata1: MemoryResource | None, + sage_k_scales0: MemoryResource | None, + sage_k_scales1: MemoryResource | None, work_queue: WorkQueue | None, schedule_token_throttle: MemoryResource | None, ) -> None: @@ -1332,6 +1377,8 @@ def load_schedule( sparse_kv_metadata1, sparse_softmax_metadata0, sparse_softmax_metadata1, + sage_k_scales0, + sage_k_scales1, ), lambda: _schedule_token_throttle_head(schedule_token_throttle), ) @@ -1359,6 +1406,8 @@ def load_schedule( sparse_kv_metadata1, sparse_softmax_metadata0, sparse_softmax_metadata1, + sage_k_scales0, + sage_k_scales1, work_queue, schedule_token_throttle, ) @@ -1371,9 +1420,9 @@ def load_schedule( if work_queue is not None: src.append(work_queue) dst = [smem_q, smem_kv] - for sparse_resource in sparse_resources: - if sparse_resource is not None and sparse_resource not in dst: - dst.append(sparse_resource) + for resource in (*sparse_resources, sage_k_scales0, sage_k_scales1): + if resource is not None and resource not in dst: + dst.append(resource) if schedule_token_throttle is not None: dst.append(schedule_token_throttle) return task_class( @@ -1670,6 +1719,8 @@ def create_load_task_split_kv( sparse_kv_metadata1: MemoryResource | None = None, sparse_softmax_metadata0: MemoryResource | None = None, sparse_softmax_metadata1: MemoryResource | None = None, + sage_k_scales0: MemoryResource | None = None, + sage_k_scales1: MemoryResource | None = None, warp_idx: int | None = None, num_warps: int | None = None, task_name: str = "LoadTask", @@ -1679,6 +1730,8 @@ def create_load_task_split_kv( """Create loads for independent K0/K1/V0/V1 resources. Here ``split`` describes those resources, not split-KV reduction. + ``sage_k_scales0/1`` are a Sage plan's ``sfK`` rings, which the task + fills after each K tile. """ if schedule_token_throttle is not None and work_queue is None: raise ValueError("schedule-token throttle requires a work queue") @@ -1705,10 +1758,13 @@ def load_schedule_body( smem_page_offsets: MemoryResource | None, smem_page_offsets_v: MemoryResource | None = None, schedule_token_throttle: MemoryResource | None = None, + sage_k_scales0: MemoryResource | None = None, + sage_k_scales1: MemoryResource | None = None, ) -> None: """Build the split-resource K/V load cadence for all schedule phases.""" if smem_q is not None: smem_q.init_load_state() + sage_k_scales_by_label = {"load_k0": sage_k_scales0, "load_k1": sage_k_scales1} active_instances = ( ( smem_k0, @@ -1768,6 +1824,14 @@ def load_tile( if not hold_page_window: _page_offsets_release(offsets) + # Issue every instance's HEAD route-record load before Q, so their + # global-memory latency overlaps the wait for the previous work tile + # to release Q and the Q load itself. + prefetch_by_label = { + load_k: _prefetch_sparse_route(sparse_kv_metadata, "head") + for smem_k, _, sparse_kv_metadata, _, load_k, _ in active_instances + if smem_k is not None + } if smem_q is not None: smem_q.acquire() smem_q.tma_load() @@ -1789,19 +1853,33 @@ def load_tile( ) in active_instances: if smem_k is None: continue - route, _ = _resolve_and_store_sparse_route( - sparse_kv_metadata, FmhaStage.Head, pipeline=False + route = _resolve_and_store_sparse_route( + sparse_kv_metadata, FmhaStage.Head, prefetch_by_label[load_k] ) load_tile(smem_k, load_k, smem_page_offsets_k, FmhaStage.Head) - head_routes.append((sparse_softmax_metadata, route)) + head_routes.append( + (sparse_softmax_metadata, route, sage_k_scales_by_label[load_k]) + ) # In the combined task, preserve both K issues ahead of Softmax # backpressure. A per-instance task naturally stages its sole route. - for sparse_softmax_metadata, route in head_routes: + for sparse_softmax_metadata, route, sage_k_scales in head_routes: _publish_sparse_softmax_route(sparse_softmax_metadata, route) + _publish_k_scales(sage_k_scales, route, FmhaStage.Head) + # Successor record loads go out after the K copies and the route + # publish, here and in the loop. Queued ahead of the copies, a record + # load delays the shared-memory reads of the route origins that feed + # their coordinates, and with them the work tile's first K tiles and + # Softmax routes. + prefetch_by_label = { + load_k: _prefetch_sparse_route(sparse_kv_metadata, "first_loop") + for smem_k, _, sparse_kv_metadata, _, load_k, _ in active_instances + if smem_k is not None + } with domain_loop(0, domain, 1, unroll=1): # V consumes the retained route before the next K overwrites it. - loop_routes = [] + # Each route goes to its Softmax consumer right after its own K + # issue instead of waiting for the other instance's loads. for ( smem_k, smem_v, @@ -1819,13 +1897,15 @@ def load_tile( smem_page_offsets_v_local, FmhaStage.Loop, ) - route, _ = _resolve_and_store_sparse_route( - sparse_kv_metadata, FmhaStage.Loop, pipeline=False + route = _resolve_and_store_sparse_route( + sparse_kv_metadata, FmhaStage.Loop, prefetch_by_label[load_k] ) load_tile(smem_k, load_k, smem_page_offsets_k, FmhaStage.Loop) - loop_routes.append((sparse_softmax_metadata, route)) - for sparse_softmax_metadata, route in loop_routes: _publish_sparse_softmax_route(sparse_softmax_metadata, route) + _publish_k_scales(sage_k_scales_by_label[load_k], route, FmhaStage.Loop) + prefetch_by_label[load_k] = _prefetch_sparse_route( + sparse_kv_metadata, "next_loop" + ) for _, smem_v, _, _, _, load_v in active_instances: if smem_v is not None: @@ -1851,6 +1931,8 @@ def load_schedule( sparse_softmax_metadata1: MemoryResource | None, smem_page_offsets: MemoryResource | None, smem_page_offsets_v: MemoryResource | None, + sage_k_scales0: MemoryResource | None, + sage_k_scales1: MemoryResource | None, work_queue: WorkQueue | None, schedule_token_throttle: MemoryResource | None, ) -> None: @@ -1871,13 +1953,15 @@ def load_schedule( smem_page_offsets, smem_page_offsets_v, schedule_token_throttle, + sage_k_scales0, + sage_k_scales1, ), lambda: _schedule_token_throttle_head(schedule_token_throttle), ) - for smem_k, sparse_kv_metadata, sparse_softmax_metadata in ( - (smem_k0, sparse_kv_metadata0, sparse_softmax_metadata0), - (smem_k1, sparse_kv_metadata1, sparse_softmax_metadata1), + for smem_k, sparse_kv_metadata, sparse_softmax_metadata, sage_k_scales in ( + (smem_k0, sparse_kv_metadata0, sparse_softmax_metadata0, sage_k_scales0), + (smem_k1, sparse_kv_metadata1, sparse_softmax_metadata1, sage_k_scales1), ): if smem_k is None and ( sparse_kv_metadata is not None or sparse_softmax_metadata is not None @@ -1885,6 +1969,8 @@ def load_schedule( raise ValueError("inactive K/V instances cannot own sparse metadata") if sparse_softmax_metadata is not None and sparse_kv_metadata is None: raise ValueError("Softmax sparse metadata requires retained KV metadata") + if sage_k_scales is not None and smem_k is None: + raise ValueError("a Sage K scale ring follows its instance's K tiles") if sparse_kv_metadata0 is not None or sparse_kv_metadata1 is not None: if smem_page_offsets is not None or smem_page_offsets_v is not None: raise ValueError( @@ -1903,6 +1989,8 @@ def load_schedule( sparse_softmax_metadata1, smem_page_offsets, smem_page_offsets_v, + sage_k_scales0, + sage_k_scales1, work_queue, schedule_token_throttle, ) @@ -1923,14 +2011,16 @@ def load_schedule( for resource in (smem_q, smem_k0, smem_k1, smem_v0, smem_v1) if resource is not None ] - for sparse_resource in ( + for resource in ( sparse_kv_metadata0, sparse_kv_metadata1, sparse_softmax_metadata0, sparse_softmax_metadata1, + sage_k_scales0, + sage_k_scales1, ): - if sparse_resource is not None: - dst.append(sparse_resource) + if resource is not None: + dst.append(resource) if schedule_token_throttle is not None: dst.append(schedule_token_throttle) q_bound_resources = tuple( @@ -1968,14 +2058,16 @@ def create_block_sparse_load_tasks_per_inst( sparse_softmax_metadata0: MemoryResource, sparse_softmax_metadata1: MemoryResource, warp_indices: tuple[int, int], + sage_k_scales0: MemoryResource | None = None, + sage_k_scales1: MemoryResource | None = None, task_class: type[DecodeGenTask] = DecodeGenTask, **kw: TaskKwarg, ) -> tuple[Task, Task]: """Assign each sparse K/V instruction stream to an independent load warp. Load0 alone owns Q and the persistent schedule-token throttle. Both tasks - consume the same logical work tile, while their K/V and sparse-metadata - pipelines remain disjoint. + consume the same logical work tile, while their K/V, sparse-metadata and + Sage K scale pipelines remain disjoint. """ if not cfg.use_block_sparse: @@ -1993,6 +2085,7 @@ def create_block_sparse_load_tasks_per_inst( domain=domain, sparse_kv_metadata0=sparse_kv_metadata0, sparse_softmax_metadata0=sparse_softmax_metadata0, + sage_k_scales0=sage_k_scales0, warp_idx=warp_indices[0], task_name="LoadTask0", task_class=task_class, @@ -2010,6 +2103,7 @@ def create_block_sparse_load_tasks_per_inst( domain=domain, sparse_kv_metadata1=sparse_kv_metadata1, sparse_softmax_metadata1=sparse_softmax_metadata1, + sage_k_scales1=sage_k_scales1, warp_idx=warp_indices[1], task_name="LoadTask1", task_class=task_class, @@ -2296,9 +2390,13 @@ def pv_mma( FmhaStage.Loop, ) + # The last BMM1 has been issued once the loop ends, and BMM2 never reads + # Q. Releasing it here lets a persistent CTA load the next work tile's + # Q while the final softmax and BMM2 waves run, as the KV256 MMA does. + smem_q.release() + pv_mma(smem_v0, smem_p0, "vp_mma_tail", KV_INST0, FmhaStage.Tail) pv_mma(smem_v1, smem_p1, "vp_mma_tail", KV_INST1, FmhaStage.Tail) - smem_q.release() def mma_schedule_prelude( smem_q: MemoryResource, @@ -2595,8 +2693,10 @@ def pv_mma(vp_mma_label: str, section: FmhaStage) -> None: qk_mma(q_desc, "qk_mma_loop", FmhaStage.Loop) pv_mma("vp_mma_loop", FmhaStage.Loop) - pv_mma("vp_mma_tail", FmhaStage.Tail) + # BMM2 never reads Q; release it before the final wave + # (see create_mma_task_split_kv). smem_q.release() + pv_mma("vp_mma_tail", FmhaStage.Tail) def mma_schedule_prelude( smem_q: MemoryResource, @@ -3072,14 +3172,11 @@ def _softmax_schedule_body( sparse_softmax_metadata: MemoryResource | None, membership_lifetime: MemoryResource | None, sage_k_scales: MemoryResource | None, - sage_summary_k_scales: MemoryResource | None, - sage_scales_in_smem: tuple[bool, ...] = (), ) -> None: """Build one softmax instance's loop, P publication, and final stats handoff. Both instances run this body and differ only in their side of the P - baton. ``sage_scales_in_smem`` flags the scale resources in SMEM form, - which carry a pipeline; schedule proxies cannot answer it themselves. + baton. """ if membership_lifetime is not None: membership_lifetime.wait() @@ -3096,16 +3193,14 @@ def _softmax_schedule_body( # The lane's Q row is fixed for the work tile. sage_q_scale = tmem_s.load_sage_q_scale() # The work framework rejects ``None`` arguments, so each Sage/sparse - # combination is its own work callable. A single-geometry plan passes - # its one routed ``sfK`` array under both names. - sage_scales = [r for r in (sage_k_scales, sage_summary_k_scales) if r is not None] - sage_smem_scales = [ - r - for r, in_smem in zip(sage_scales, sage_scales_in_smem, strict=True) - if in_smem - ] + # combination is its own work callable. use_sparse = sparse_softmax_metadata is not None use_sage = cfg.use_sage_attention + # The ``sfK`` words sit in the instance's ring, which the load warp fills + # with each tile. They are read in place, so the ring stage stays held + # through the last pass that reads them: the max pass for dequantized + # scores, else the P pass. + releases_sage_stage_after_max = cfg.sage_scores_dequantized() compute_softmax_loop = { (False, False): tmem_s.compute_softmax_loop, (False, True): tmem_s.compute_sage_softmax_loop, @@ -3135,35 +3230,12 @@ def _softmax_schedule_body( sparse_token_word2, sparse_token_word3, ) = sparse_softmax_metadata.load_route() - if cutlass.const_expr(cfg.use_sage_attention): - # Take the route's ``sfK`` words before the metadata stage - # is released. Each iteration fills and publishes this tile's - # words only; the second stage lets the fill overlap the - # previous tile's P pass on the instance's other warps. - for scales in sage_smem_scales: - scales.acquire() - scales.publish_route_tile(route_flags=sparse_route_flags) - scales.commit() - scales.wait() - sage_scale_arr = sage_k_scales.take_route_tile( - route_flags=sparse_route_flags - ) - sage_summary_scale_arr = sage_scale_arr - if sage_summary_k_scales is not None: - sage_summary_scale_arr = sage_summary_k_scales.take_route_tile( - route_flags=sparse_route_flags - ) sparse_softmax_metadata.release() - if cutlass.const_expr(cfg.use_sage_attention and not use_sparse): - # Issue the dense tile's ``sfK`` loads before the score wait so - # they hide behind the QK MMA. - for scales in sage_smem_scales: - scales.acquire() - scales.publish_tile() - scales.commit() - scales.wait() - sage_scale_arr = sage_k_scales.take_tile() - sage_summary_scale_arr = sage_scale_arr + if cutlass.const_expr(use_sage): + # ConsWait/ConsWork: the tile's ``sfK`` words; both passes read + # the waited stage through its routed offset. + sage_k_scales.wait() + sage_k_stage_bytes = sage_k_scales.open_stage() # ConsWait/ConsWork: load S from TMEM and compute the tile max. tmem_s.wait() softmax_loop_kwargs = dict( @@ -3174,8 +3246,7 @@ def _softmax_schedule_body( ) if cutlass.const_expr(cfg.use_sage_attention): softmax_loop_kwargs["sage_q_scale"] = sage_q_scale - softmax_loop_kwargs["sage_scale_arr"] = sage_scale_arr - softmax_loop_kwargs["sage_summary_scale_arr"] = sage_summary_scale_arr + softmax_loop_kwargs["sage_k_stage_bytes"] = sage_k_stage_bytes if use_sparse: softmax_loop_kwargs.update( sparse_origin0=sparse_origin0, @@ -3189,6 +3260,9 @@ def _softmax_schedule_body( old_max_arr, sum_arr, new_max_arr, s_arr = compute_softmax_loop( **softmax_loop_kwargs ) + if cutlass.const_expr(releases_sage_stage_after_max): + # Dequantized scores went back to TMEM; the P pass reads no words. + sage_k_scales.release() if cutlass.const_expr(not cfg.use_keeps_mma_ab or not cfg.uses_tmem_p): # ConsRelease: free S once the scores are in registers unless # a Keeps TMEM-P operand still aliases the consumed columns. @@ -3209,8 +3283,7 @@ def _softmax_schedule_body( p_fragments_kwargs = dict(new_max_arr=new_max_arr) if cutlass.const_expr(cfg.use_sage_attention): p_fragments_kwargs["sage_q_scale"] = sage_q_scale - p_fragments_kwargs["sage_scale_arr"] = sage_scale_arr - p_fragments_kwargs["sage_summary_scale_arr"] = sage_summary_scale_arr + p_fragments_kwargs["sage_k_stage_bytes"] = sage_k_stage_bytes if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): p_fragments_kwargs["route_flags"] = sparse_route_flags compute_p_fragments(**p_fragments_kwargs) @@ -3249,10 +3322,9 @@ def _softmax_schedule_body( # TMEM-P has consumed the aliased S columns, so the next QK # wave can now overwrite S. tmem_s.release() - if cutlass.const_expr(cfg.use_sage_attention): - # ConsRelease: both passes have read this tile's ``sfK`` slot. - for scales in sage_smem_scales: - scales.release() + if cutlass.const_expr(use_sage and not releases_sage_stage_after_max): + # ConsRelease: both passes have read this tile's ``sfK`` words. + sage_k_scales.release() # ProdWork: FP8 path applies the cross-resource sum correction # before TmemS.reduce_sums publishes the new running sums. tmem_softmax_global.global_correction( @@ -3319,15 +3391,11 @@ def create_softmax0_task( domain: int | cutlass.Int32, membership_lifetime: MemoryResource | None = None, sage_k_scales: MemoryResource | None = None, - sage_summary_k_scales: MemoryResource | None = None, task_class: type[DecodeGenTask] = DecodeGenTask, **kw: TaskKwarg, ) -> Task: """Create the first softmax task, including optional ordered publication.""" - sage_scales = [r for r in (sage_k_scales, sage_summary_k_scales) if r is not None] - sage_scales_in_smem = tuple(r.in_smem for r in sage_scales) - @_schedule_with_optional_resources def softmax0_schedule( tmem_s0: MemoryResource, @@ -3338,7 +3406,6 @@ def softmax0_schedule( sparse_softmax_metadata: MemoryResource | None, membership_lifetime: MemoryResource | None, sage_k_scales: MemoryResource | None, - sage_summary_k_scales: MemoryResource | None, work_queue: WorkQueue | None, ) -> None: """Schedule softmax0 with the supplied order and sparse resources.""" @@ -3357,8 +3424,6 @@ def softmax0_schedule( sparse_softmax_metadata, membership_lifetime, sage_k_scales, - sage_summary_k_scales, - sage_scales_in_smem, ), ) @@ -3371,7 +3436,6 @@ def softmax0_schedule( sparse_softmax_metadata, membership_lifetime, sage_k_scales, - sage_summary_k_scales, work_queue, ) src = [tmem_s0] @@ -3379,12 +3443,11 @@ def softmax0_schedule( src.append(sparse_softmax_metadata) if membership_lifetime is not None: src.append(membership_lifetime) - # The scale resources are produced and consumed by this task. - src.extend(sage_scales) + if sage_k_scales is not None: + src.append(sage_k_scales) if work_queue is not None: src.append(work_queue) dst = [tmem_softmax_local0, smem_p0, tmem_softmax_global0] - dst.extend(r for r in sage_scales if r.in_smem) if tmem_softmax_order is not None: dst.append(tmem_softmax_order) return task_class( @@ -3417,15 +3480,11 @@ def create_softmax1_task( domain: int | cutlass.Int32, membership_lifetime: MemoryResource | None = None, sage_k_scales: MemoryResource | None = None, - sage_summary_k_scales: MemoryResource | None = None, task_class: type[DecodeGenTask] = DecodeGenTask, **kw: TaskKwarg, ) -> Task: """Create the second softmax task, including optional ordered publication.""" - sage_scales = [r for r in (sage_k_scales, sage_summary_k_scales) if r is not None] - sage_scales_in_smem = tuple(r.in_smem for r in sage_scales) - @_schedule_with_optional_resources def softmax1_schedule( tmem_s1: MemoryResource, @@ -3436,7 +3495,6 @@ def softmax1_schedule( sparse_softmax_metadata: MemoryResource | None, membership_lifetime: MemoryResource | None, sage_k_scales: MemoryResource | None, - sage_summary_k_scales: MemoryResource | None, work_queue: WorkQueue | None, ) -> None: """Schedule softmax1 with the supplied order and sparse resources.""" @@ -3461,8 +3519,6 @@ def softmax1_schedule( sparse_softmax_metadata, membership_lifetime, sage_k_scales, - sage_summary_k_scales, - sage_scales_in_smem, ), ) @@ -3475,7 +3531,6 @@ def softmax1_schedule( sparse_softmax_metadata, membership_lifetime, sage_k_scales, - sage_summary_k_scales, work_queue, ) src = [tmem_s1] @@ -3485,18 +3540,14 @@ def softmax1_schedule( src.append(tmem_softmax_order) if membership_lifetime is not None: src.append(membership_lifetime) - # The scale resources are produced and consumed by this task. - src.extend(sage_scales) + if sage_k_scales is not None: + src.append(sage_k_scales) if work_queue is not None: src.append(work_queue) + dst = [tmem_softmax_local1, smem_p1, tmem_softmax_global1] return task_class( src_resources=src, - dst_resources=[ - tmem_softmax_local1, - smem_p1, - tmem_softmax_global1, - *(r for r in sage_scales if r.in_smem), - ], + dst_resources=dst, q_bound_resources=((tmem_s1, False),), cfg=cfg, warp_idx=cfg.softmax1_warp_idx, 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 ebbdc53fdfbf..bfd4d2229592 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 @@ -42,7 +42,6 @@ from .fmha_decode_constants import ( FP32_BYTES, - FP8_P_QUANT_LOG2_SCALE, FP8_PACKED_OUTPUT_REGS_PER_THREAD, FP8_VALUES_PER_REG, FP16_VALUES_PER_REG, @@ -350,7 +349,7 @@ def _attention_sink_log2_lse( ) sink_lse = sink_ptr.load() * Float32(1.4426950408889634) if cutlass.const_expr(cfg.use_fp8_pv): - sink_lse += Float32(FP8_P_QUANT_LOG2_SCALE) + sink_lse += Float32(cfg.fp8_p_quant_log2_scale) return sink_lse diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/sage_prepare.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/sage_prepare.py new file mode 100644 index 000000000000..88afcfccd3cc --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/sage_prepare.py @@ -0,0 +1,364 @@ +# 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. + +"""Prepare the Sage scales of a dense or block-sparse plan ahead of attention. + +The plan's scales are gathered into one buffer in the order the attention +kernel reads them (``_SageScaleLayout``): per sequence (batch entry and KV +head) the ``sfK`` atoms of the K tokens and of a proxy plan's block +summaries, the ``sfQ`` rows of the Q tiles, then the ``sfV`` row and, with a +V mean, the ``v_mean`` row. This kernel reads the flat ``q_scale``, +``k_scale`` and ``k_summary_scale`` tensors and the ``[B, Hkv, D]`` +``v_scale`` and ``v_mean`` tensors; one thread writes one float4 of the +buffer. +""" + +from dataclasses import dataclass + +import cutlass +import cutlass.cute as cute +from cuda.bindings import driver as cuda_drv + +from flashinfer.utils import ceil_div + +from ...sage import ( + _FLOAT4_WORDS, + _SAGE_ATOM_TOKENS, + _SageKSection, + _SageScaleLayout, + log2_block_size, +) +from .fmha_decode_config import FmhaDecodeConfig +from .fmha_decode_resources.helpers_common import ( + _q_row_token_and_local_head, + _q_tile_valid_rows_for_seq, +) +from .fmha_decode_resources.sage_scales import load_flat_scale + +_THREADS_PER_CTA = 128 + + +@dataclass +class _KSectionSource: + """One K section's atom geometry and source scales, as a float4 writer reads them. + + The dataclass is not frozen because the writer selects one of the two + sections inside a traced branch, where the DSL rebuilds a frozen + dataclass as a proxy. + """ + + # Float4s of one atom. + atom_float4s: cutlass.Int32 + # Tokens from one word's token to the next word's. + tokens_per_word: cutlass.Int32 + # Words of an atom that hold scales; the rest are zero padding. + words_per_atom: cutlass.Int32 + # log2 of the source tensor's block size. + log2_block: cutlass.Int32 + # Base address and per-head slot count of the flat source tensor. + scale_addr: cutlass.Int64 + head_stride: cutlass.Int32 + # Tokens of the section: K tokens, or block summaries. + length: cutlass.Int32 + + @staticmethod + def create(section: _SageKSection, scales: cute.Tensor) -> "_KSectionSource": + return _KSectionSource( + atom_float4s=cutlass.Int32(section.atom_float4s), + tokens_per_word=cutlass.Int32(section.tokens_per_word), + words_per_atom=cutlass.Int32(section.words_per_atom), + log2_block=cutlass.Int32(log2_block_size(section.source_block_size)), + scale_addr=scales.iterator.toint(), + head_stride=cutlass.Int32(scales.shape[1]), + length=cutlass.Int32(section.length), + ) + + +@cute.jit +def _store_float4( + scales: cute.Tensor, float4_idx: cutlass.Int32, words: cutlass.Array +) -> None: + """Store four words as float4 ``float4_idx`` of the prepared scales.""" + + float4_ptr = cutlass.inttoptr( + scales.iterator.toint() + + cutlass.Int64(float4_idx) * cutlass.Int64(_FLOAT4_WORDS * 4), + mem_space=1, + dtype=cutlass.Float32, + ) + float4_ptr.store( + words.data_ptr().load(count=_FLOAT4_WORDS, alignment=4), + alignment=_FLOAT4_WORDS * 4, + ) + + +@cute.jit +def _write_v_float4( + v_scale: cute.Tensor, + v_mean: cute.Tensor | None, + scales: cute.Tensor, + float4_idx: cutlass.Int32, + layout: cutlass.Constexpr[_SageScaleLayout], + *, + sequence_idx: cutlass.Int32, + float4_in_section: cutlass.Int32, +) -> None: + """Write float4 ``float4_idx`` of the buffer: four channels of one V row. + + ``float4_in_section`` places the float4 in the sequence's V rows: the + ``sfV`` row, copied from the sequence's row of ``v_scale``, or the + ``v_mean`` row behind it, copied from ``v_mean``. + """ + + v_rows = layout.v_rows + row_float4s = cutlass.Int32(v_rows.headdim // _FLOAT4_WORDS) + source_addr = v_scale.iterator.toint() + float4_in_row = float4_in_section + if cutlass.const_expr(v_rows.has_mean): + if float4_in_section >= row_float4s: + source_addr = v_mean.iterator.toint() + float4_in_row = float4_in_section - row_float4s + source_ptr = cutlass.inttoptr( + source_addr + + cutlass.Int64(sequence_idx * row_float4s + float4_in_row) + * cutlass.Int64(_FLOAT4_WORDS * 4), + mem_space=1, + dtype=cutlass.Float32, + ) + words = cutlass.Array( + cutlass.Float32, _FLOAT4_WORDS, space=cutlass.AddressSpace.rmem + ) + loaded = source_ptr.load(count=_FLOAT4_WORDS, alignment=_FLOAT4_WORDS * 4) + for elem in cutlass.range_constexpr(_FLOAT4_WORDS): + words[elem] = cutlass.Float32(loaded[elem]) + _store_float4(scales, float4_idx, words) + + +@cute.jit +def _write_k_float4( + k_scale: cute.Tensor, + k_summary_scale: cute.Tensor | None, + scales: cute.Tensor, + float4_idx: cutlass.Int32, + layout: cutlass.Constexpr[_SageScaleLayout], + *, + batch_idx: cutlass.Int32, + head_idx: cutlass.Int32, + float4_in_sequence: cutlass.Int32, +) -> None: + """Write float4 ``float4_idx`` of the buffer: four words of one K atom. + + ``float4_in_sequence`` places the float4 in the sequence's K sections: + the K tokens from ``k_scale``, or the block summaries from + ``k_summary_scale`` behind them. Word ``w`` of an atom is the section's + scale of the atom's token ``w * tokens_per_word``; a token past the + section's length takes the last scale, and padding words are zero. + """ + + section = _KSectionSource.create(layout.k_tokens, k_scale) + float4_in_section = float4_in_sequence + if cutlass.const_expr(layout.summaries is not None): + summaries_first_float4 = layout.summaries.word_offset // _FLOAT4_WORDS + if float4_in_sequence >= cutlass.Int32(summaries_first_float4): + section = _KSectionSource.create(layout.summaries, k_summary_scale) + float4_in_section = float4_in_sequence - cutlass.Int32( + summaries_first_float4 + ) + atom_idx = float4_in_section // section.atom_float4s + first_word = (float4_in_section - atom_idx * section.atom_float4s) * cutlass.Int32( + _FLOAT4_WORDS + ) + + words = cutlass.Array( + cutlass.Float32, _FLOAT4_WORDS, space=cutlass.AddressSpace.rmem + ) + for elem in cutlass.range_constexpr(_FLOAT4_WORDS): + word = first_word + cutlass.Int32(elem) + words[elem] = cutlass.Float32(0.0) + # Only a two-word atom pads, and it is a single float4. + if word < section.words_per_atom: + words[elem] = load_flat_scale( + section.scale_addr, + section.head_stride, + head_idx=head_idx, + batch_idx=batch_idx, + seq_len=section.length, + token_idx=atom_idx * cutlass.Int32(_SAGE_ATOM_TOKENS) + + word * section.tokens_per_word, + log2_block=section.log2_block, + ) + _store_float4(scales, float4_idx, words) + + +@cute.jit +def _write_q_float4( + q_scale: cute.Tensor, + scales: cute.Tensor, + float4_idx: cutlass.Int32, + cfg: cutlass.Constexpr[FmhaDecodeConfig], + layout: cutlass.Constexpr[_SageScaleLayout], + *, + batch_idx: cutlass.Int32, + kv_head_idx: cutlass.Int32, + float4_in_section: cutlass.Int32, +) -> None: + """Write float4 ``float4_idx`` of the buffer: four rows of one Q tile. + + ``float4_in_section`` places the float4 in the sequence's Q rows. Row + ``r`` of the tile takes the Q scale of the token and Q head the softmax + maps the row to (``_q_row_token_and_local_head``); a row past the + tile's valid rows (``_q_tile_valid_rows_for_seq``) takes the last valid + row's. + """ + + q_rows = layout.q_rows + tile_float4s = cutlass.Int32(q_rows.tile_size_q // _FLOAT4_WORDS) + tile_idx = float4_in_section // tile_float4s + first_row = (float4_in_section - tile_idx * tile_float4s) * cutlass.Int32( + _FLOAT4_WORDS + ) + # The row helpers take the plan's heads per KV head as the packed row + # count and its fixed Q length; Sage plans have no variable Q lengths. + heads_q_per_kv = cutlass.Int32(q_rows.heads_q_per_kv) + seq_len_q = cutlass.Int32(q_rows.seq_len_q) + last_row = _q_tile_valid_rows_for_seq( + cfg, heads_q_per_kv, tile_idx, seq_len_q + ) - cutlass.Int32(1) + + words = cutlass.Array( + cutlass.Float32, _FLOAT4_WORDS, space=cutlass.AddressSpace.rmem + ) + for elem in cutlass.range_constexpr(_FLOAT4_WORDS): + row = cute.math.min(first_row + cutlass.Int32(elem), last_row) + q_token_idx, local_head_idx = _q_row_token_and_local_head( + cfg, heads_q_per_kv, tile_idx, row + ) + words[elem] = load_flat_scale( + q_scale.iterator.toint(), + cutlass.Int32(q_scale.shape[1]), + head_idx=kv_head_idx * heads_q_per_kv + local_head_idx, + batch_idx=batch_idx, + seq_len=seq_len_q, + token_idx=q_token_idx, + log2_block=cutlass.Int32(log2_block_size(q_rows.q_block_size)), + ) + _store_float4(scales, float4_idx, words) + + +class _PrepareSageScales: + """Write a Sage plan's prepared scales from its flat scale tensors. + + The buffer covers ``batch_size * num_kv_heads`` sequences + (``_SageScaleLayout``); the grid gives each of its float4s one thread. The kernel is an ordinary launch on the plan's + stream. A dense plan's attention launch follows it in stream order; a + block-sparse plan's route prepare follows it as a programmatic dependent + launch that acquires this grid before exiting, and attention follows the + route prepare in stream order. + """ + + def __init__( + self, + *, + cfg: FmhaDecodeConfig, + layout: _SageScaleLayout, + batch_size: int, + num_kv_heads: int, + ) -> None: + self.cfg = cfg + self.layout = layout + self.num_kv_heads = num_kv_heads + num_sequences = batch_size * num_kv_heads + self.sequence_float4s = layout.sequence_words // _FLOAT4_WORDS + self.num_float4s = layout.total_words(num_sequences) // _FLOAT4_WORDS + + @cute.jit + def __call__( + self, + q_scale: cute.Tensor, + k_scale: cute.Tensor, + k_summary_scale: cute.Tensor | None, + v_scale: cute.Tensor, + v_mean: cute.Tensor | None, + scales: cute.Tensor, + stream: cuda_drv.CUstream, + ) -> None: + self.kernel(q_scale, k_scale, k_summary_scale, v_scale, v_mean, scales).launch( + grid=[ceil_div(self.num_float4s, _THREADS_PER_CTA), 1, 1], + block=[_THREADS_PER_CTA, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + q_scale: cute.Tensor, + k_scale: cute.Tensor, + k_summary_scale: cute.Tensor | None, + v_scale: cute.Tensor, + v_mean: cute.Tensor | None, + scales: cute.Tensor, + ) -> None: + """Write this thread's float4 of a K section, the Q rows or the V rows, if any.""" + + thread_idx, _, _ = cute.arch.thread_idx() + block_idx, _, _ = cute.arch.block_idx() + float4_idx = block_idx * cutlass.Int32(_THREADS_PER_CTA) + thread_idx + if float4_idx < cutlass.Int32(self.num_float4s): + sequence_idx = float4_idx // cutlass.Int32(self.sequence_float4s) + float4_in_sequence = float4_idx - sequence_idx * cutlass.Int32( + self.sequence_float4s + ) + batch_idx = sequence_idx // cutlass.Int32(self.num_kv_heads) + head_idx = sequence_idx - batch_idx * cutlass.Int32(self.num_kv_heads) + q_rows_first_float4 = cutlass.Int32( + self.layout.q_rows.word_offset // _FLOAT4_WORDS + ) + v_rows_first_float4 = cutlass.Int32( + self.layout.v_rows.word_offset // _FLOAT4_WORDS + ) + if float4_in_sequence >= v_rows_first_float4: + _write_v_float4( + v_scale, + v_mean, + scales, + float4_idx, + self.layout, + sequence_idx=sequence_idx, + float4_in_section=float4_in_sequence - v_rows_first_float4, + ) + elif float4_in_sequence < q_rows_first_float4: + _write_k_float4( + k_scale, + k_summary_scale, + scales, + float4_idx, + self.layout, + batch_idx=batch_idx, + head_idx=head_idx, + float4_in_sequence=float4_in_sequence, + ) + else: + _write_q_float4( + q_scale, + scales, + float4_idx, + self.cfg, + self.layout, + batch_idx=batch_idx, + kv_head_idx=head_idx, + float4_in_section=float4_in_sequence - q_rows_first_float4, + ) + + +__all__ = ["_PrepareSageScales"] diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/tcgen05_compat.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/tcgen05_compat.py index f852b1daec46..514005d3db3e 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/tcgen05_compat.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/tcgen05_compat.py @@ -14,11 +14,86 @@ """Narrow compatibility helpers for tcgen05 primitives used by PrimTS kernels.""" +import functools +import importlib.metadata + import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32 +from cutlass._mlir.dialects import arith as arith_dialect from cutlass.cutlass_dsl import dsl_user_op from cutlass.experimental import primitives as prims +def _ld_32x32b_max_ptx(num: int) -> str: + registers = ", ".join(f"{{$w{i}}}" for i in range(num)) + return ( + f"tcgen05.ld.red.sync.aligned.32x32b.x{num}.f32.max " + f"{{{registers}}}, {{$w{num}}}, [{{$r0}}];" + ) + + +@functools.cache +def dsl_supports_ldtm_stat() -> bool: + """LDTM.STAT uses a native wrapper in DSL 4.8 or inline PTX in 4.7.""" + try: + dsl_version = importlib.metadata.version("nvidia-cutlass-dsl") + except importlib.metadata.PackageNotFoundError: + return False + from packaging import version as pkg_version + + try: + # Use .release so 4.7.0.dev* includes the inline-PTX path. + return pkg_version.Version(dsl_version).release >= (4, 7, 0) + except pkg_version.InvalidVersion: + return False + + +def ldtm_stat_supported(compute_capability: tuple[int, int]) -> bool: + """Whether kernels may take score maxima from LDTM.STAT on this GPU. + + tcgen05.ld.red.max (LDTM.STAT) exists on B300 (SM103) and Rubin (SM107), + not B200 (SM100). DSL 4.7's Rubin escape hatch targets sm_100f, which + cannot emit it, so SM107 also needs the native DSL wrapper. + """ + if not dsl_supports_ldtm_stat(): + return False + capability = tuple(compute_capability) + if capability == (10, 3): + return True + if capability == (10, 7): + return hasattr(prims, "tcgen05_ld_red") + return False + + +@cute.jit +def tcgen05_ld_32x32b_max(tmem_addr: Int32, num: cutlass.Constexpr[int]) -> tuple: + """Load ``num`` FP32 TMEM columns per lane and their maximum (LDTM.STAT). + + Returns the ``num`` loaded values followed by their maximum. The native + DSL wrapper is used when present; inline PTX covers CUTLASS DSL 4.7, which + has none. The caller gates the architecture with ``ldtm_stat_supported`` + and waits for TMEM loads before reading any result. + """ + if cutlass.const_expr(hasattr(prims, "tcgen05_ld_red")): + loaded_words, red_word = prims.tcgen05_ld_red( + prims.Tcgen05LdStShape.SHAPE_32X32B, + prims.make_tmem_ptr(tmem_addr, Float32), + prims.ReductionKind.MAX, + num=num, + ) + result = tuple(loaded_words[i].bitcast(Float32) for i in range(num)) + ( + Float32(arith_dialect.bitcast(Float32.mlir_type, red_word.ir_value())), + ) + else: + result = cute.arch.inline_ptx( + _ld_32x32b_max_ptx(num), + write_only_types=[Float32] * (num + 1), + read_only_args=[tmem_addr], + ) + return result + + @dsl_user_op def tcgen05_mma_ws( mma_kind, @@ -74,4 +149,9 @@ def tcgen05_mma_ws( ) -__all__ = ["tcgen05_mma_ws"] +__all__ = [ + "dsl_supports_ldtm_stat", + "ldtm_stat_supported", + "tcgen05_ld_32x32b_max", + "tcgen05_mma_ws", +] diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/sage.py b/tensorrt_llm/_torch/attention/backends/prims_ts/sage.py index 97708c79525b..8d224328a3d9 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/sage.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/sage.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Sage attention inputs for the PrimTS decode kernels. +"""Sage attention inputs for the PrimTS decode kernels and the layout of their prepared scales. Sage attention runs ``QK^T`` on 8-bit Q/K with one dequantization scale per token block and ``PV`` on E4M3 P/V with one scale per V channel. Q/K scales use @@ -29,7 +29,7 @@ import torch -from flashinfer.utils import ceil_div +from flashinfer.utils import ceil_div, round_up SAGE_K_BLOCK_SIZES = (1, 4, 16, 32, 64, 128, 256) # Sage scale slots of the contiguous attention adapter, in ABI order. @@ -69,7 +69,8 @@ class SageAttentionParams: :class:`SageAttentionConfig`: ``q_scale`` is ``[Hq, flat_scale_numel(B, Sq, q_block_size)]``, ``k_scale`` is ``[Hkv, flat_scale_numel(B, Skv, k_block_size)]``, and ``v_scale`` and - ``v_mean`` are ``[Hkv, D]``. ``k_summary_scale`` holds the scales of the + ``v_mean`` are ``[B, Hkv, D]``, one per-channel row per batch entry and + KV head. ``k_summary_scale`` holds the scales of the K summaries of a proxy plan, ``[Hkv, flat_scale_numel(B, num_kv_blocks, k_summary_block_size)]``. ``v_mean`` and ``k_summary_scale`` are present exactly when the plan uses @@ -123,7 +124,7 @@ def _validate_scale_tensor( tensor: torch.Tensor, name: str, *, - expected_shape: tuple[int, int], + expected_shape: tuple[int, ...], device: torch.device, ) -> None: if not isinstance(tensor, torch.Tensor): @@ -152,7 +153,7 @@ def sage_scale_shapes( num_kv_heads: int, head_dim: int, summary_seq_len: int | None = None, -) -> dict[str, tuple[int, int]]: +) -> dict[str, tuple[int, ...]]: """Return the shape of every scale tensor a plan consumes, by field name. ``summary_seq_len`` is the KV block count of a proxy plan and ``None`` @@ -167,7 +168,7 @@ def sage_scale_shapes( num_kv_heads, flat_scale_numel(batch_size, seq_len_kv, config.k_block_size), ), - "v_scale": (num_kv_heads, head_dim), + "v_scale": (batch_size, num_kv_heads, head_dim), } if summary_seq_len is not None: shapes["k_summary_scale"] = ( @@ -175,13 +176,13 @@ def sage_scale_shapes( flat_scale_numel(batch_size, summary_seq_len, config.k_summary_block_size), ) if config.v_mean: - shapes["v_mean"] = (num_kv_heads, head_dim) + shapes["v_mean"] = (batch_size, num_kv_heads, head_dim) return shapes def validate_sage_params( params: SageAttentionParams, - expected_shapes: Mapping[str, tuple[int, int]], + expected_shapes: Mapping[str, tuple[int, ...]], *, device: torch.device, ) -> None: @@ -204,6 +205,266 @@ def validate_sage_params( ) +_FLOAT4_WORDS = 4 +# A Sage K-scale atom is one Keeps layout atom of 64 K tokens, read one word +# per 32-token score fragment at least, so it holds at least two words. The +# prepared scales are written and copied in float4 units, the width of one +# ``cp.async``. +_SAGE_ATOM_TOKENS = 64 +_SAGE_ATOM_MIN_WORDS = 2 + + +def _validate_i32_address(value: int, name: str) -> int: + """Reject a buffer extent device Int32 addressing cannot hold.""" + + if value > 2**31 - 1: + raise OverflowError(f"{name} must fit in signed int32") + return value + + +@dataclass(frozen=True) +class _SageKSection: + """The ``sfK`` words of one route kind: one atom per 64 tokens. + + An atom covers 64 tokens (two K32 score fragments) and holds + ``words_per_atom = max(2, 64 / block_size)`` scales, word ``w`` being the + scale of the atom's token ``w * tokens_per_word``, padded with zeros to + whole float4s. Scales come from the flat source tensor of one scale per + ``source_block_size`` tokens; a coarser source repeats its value, and a + token past ``length`` takes the last scale. + """ + + length: int + block_size: int + source_block_size: int + word_offset: int + + @property + def words_per_atom(self) -> int: + """Return the words of one atom that hold scales: one per block, at least two.""" + + return max(_SAGE_ATOM_MIN_WORDS, _SAGE_ATOM_TOKENS // self.block_size) + + @property + def tokens_per_word(self) -> int: + """Return the tokens from one word's token to the next word's.""" + + return _SAGE_ATOM_TOKENS // self.words_per_atom + + @property + def atom_words(self) -> int: + """Return the words of one atom: ``words_per_atom`` padded to whole float4s.""" + + return round_up(self.words_per_atom, _FLOAT4_WORDS) + + @property + def atom_float4s(self) -> int: + """Return the float4s of one atom.""" + + return self.atom_words // _FLOAT4_WORDS + + @property + def atoms(self) -> int: + """Return the section's 64-token atoms, the last one possibly partial.""" + + return ceil_div(self.length, _SAGE_ATOM_TOKENS) + + @property + def words(self) -> int: + """Return the words of the section's atoms.""" + + return self.atoms * self.atom_words + + +@dataclass(frozen=True) +class _SageQSection: + """The ``sfQ`` words of a sequence's Q tiles: one word per tile row. + + Tile ``t`` holds ``tile_size_q`` words from word ``word_offset + t * + tile_size_q``, in the order the softmax lanes read them: row ``r`` is + token ``t * tokens_per_tile + r // heads_q_per_kv`` of local Q head ``r % + heads_q_per_kv``, and its word is that head's scale of the token's + ``q_block_size`` block. Rows past the sequence's last token repeat the + last valid word, so a masked row keeps a finite positive scale. + """ + + seq_len_q: int + q_block_size: int + heads_q_per_kv: int + tile_size_q: int + word_offset: int + + def __post_init__(self) -> None: + if self.tile_size_q % _FLOAT4_WORDS or self.heads_q_per_kv > self.tile_size_q: + raise ValueError("sfQ tiles hold whole float4s of at least one token") + + @property + def tokens_per_tile(self) -> int: + """Return the complete Q tokens of one tile.""" + + return self.tile_size_q // self.heads_q_per_kv + + @property + def tiles(self) -> int: + """Return the Q tiles of one sequence, the last one possibly partial.""" + + return ceil_div(self.seq_len_q, self.tokens_per_tile) + + def valid_rows(self, tile_idx: int) -> int: + """Return the rows of tile ``tile_idx`` that hold a token of the sequence.""" + + remaining = self.seq_len_q - tile_idx * self.tokens_per_tile + return min(remaining, self.tokens_per_tile) * self.heads_q_per_kv + + @property + def words(self) -> int: + """Return the words of the section: ``tile_size_q`` per tile.""" + + return self.tiles * self.tile_size_q + + +@dataclass(frozen=True) +class _SageVSection: + """The ``sfV`` row and, with a V mean, the ``v_mean`` row of one sequence. + + Row ``r`` holds channel ``c`` at word ``word_offset + r * headdim + c`` + of the sequence, copied from the sequence's rows of the ``[B, Hkv, D]`` + tensors as they are. + """ + + headdim: int + has_mean: bool + word_offset: int + + def __post_init__(self) -> None: + if self.headdim % _FLOAT4_WORDS: + raise ValueError("sfV rows hold whole float4s") + + @property + def rows(self) -> int: + """Return the rows of one head: the scales, then the means.""" + + return 2 if self.has_mean else 1 + + @property + def mean_word_offset(self) -> int: + """Return the first word of the mean row within the section.""" + + return self.headdim + + @property + def words(self) -> int: + """Return the words of the section's rows.""" + + return self.rows * self.headdim + + +@dataclass(frozen=True) +class _SageScaleLayout: + """Word layout of the FP32 scale buffer a Sage plan prepares before attention. + + The buffer is one block of ``sequence_words`` words per sequence ``s = b + * Hkv + h`` (KV head ``h`` of batch entry ``b``), each holding four + sections:: + + | k_tokens | summaries (proxy plans only) | q_rows | v_rows | + + so any scale of a sequence is ``s * sequence_words + section.word_offset + + offset within the section``. + + ``k_tokens`` and ``summaries`` are ``_SageKSection`` atoms of the block + size each route kind reads (``FmhaDecodeConfig.sage_k_scale_block_size``): + 16-token K blocks give four words per atom (tokens 0, 16, 32, 48), + one-token scales 64, blocks of 64 tokens or more the two-word floor (one + per fragment) padded to a float4. ``q_rows`` is a ``_SageQSection`` and + ``v_rows`` a ``_SageVSection``. Every section starts on a float4 and K + atoms are whole float4s because the attention load warp copies them with + 16-byte ``cp.async``. + """ + + seq_len_kv: int + k_block_size: int + # Summary sequence length of a proxy plan, else zero. + num_summaries: int + k_summary_block_size: int + k_tokens_block_size: int + summaries_block_size: int + seq_len_q: int + q_block_size: int + heads_q_per_kv: int + tile_size_q: int + headdim: int + v_mean: bool + + @property + def k_tokens(self) -> _SageKSection: + """Return the K-token section, which starts a sequence's words.""" + + return _SageKSection( + length=self.seq_len_kv, + block_size=self.k_tokens_block_size, + source_block_size=self.k_block_size, + word_offset=0, + ) + + @property + def summaries(self) -> _SageKSection | None: + """Return a proxy plan's summary section, which follows the K tokens.""" + + if self.num_summaries == 0: + return None + return _SageKSection( + length=self.num_summaries, + block_size=self.summaries_block_size, + source_block_size=self.k_summary_block_size, + word_offset=self.k_tokens.words, + ) + + @property + def k_sections(self) -> tuple[_SageKSection, ...]: + """Return the K sections in order, paired with ``k_scale`` then ``k_summary_scale``.""" + + summaries = self.summaries + return (self.k_tokens,) if summaries is None else (self.k_tokens, summaries) + + @property + def q_rows(self) -> _SageQSection: + """Return the Q section, which follows the K sections.""" + + return _SageQSection( + seq_len_q=self.seq_len_q, + q_block_size=self.q_block_size, + heads_q_per_kv=self.heads_q_per_kv, + tile_size_q=self.tile_size_q, + word_offset=sum(section.words for section in self.k_sections), + ) + + @property + def v_rows(self) -> _SageVSection: + """Return the V section, which follows the Q rows.""" + + q_rows = self.q_rows + return _SageVSection( + headdim=self.headdim, + has_mean=self.v_mean, + word_offset=q_rows.word_offset + q_rows.words, + ) + + @property + def sequence_words(self) -> int: + """Return the words of one sequence: K sections, Q rows, then V rows.""" + + v_rows = self.v_rows + return v_rows.word_offset + v_rows.words + + def total_words(self, num_sequences: int) -> int: + """Return the words of the buffer for ``batch * Hkv`` sequences.""" + + return _validate_i32_address( + num_sequences * self.sequence_words, "sage_scales_words" + ) + + __all__ = [ "SAGE_ADAPTER_SLOTS", "SAGE_K_BLOCK_SIZES", From e67ac9dc4d3f80be20464d8546341892db770cb6 Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Sun, 11 Oct 2026 03:31:38 +0000 Subject: [PATCH 02/13] [None][fix] Keep VisualGen block-sparse attention working under torch.compile - Run PrimsTSBlockSparseFmha.forward outside torch.compile: planning records an external CUDA event that Dynamo rejects. - Resolve the dense-layer flag at construction in the VisualGen TRTLLM and CuTeDSL SOL backends. Testing layer_idx per call made torch.compile specialize the attention frames per layer; with SOL sparse on every layer the core forward was then traced 40 times per call on real tensors, and the traces kept each layer's Q, K, V and output alive until a full garbage collection (OOM on B200 at 720p, 81 frames). - Test that compiled callers share one graph across SOL layers. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- .../backends/fmha/prims_ts_block_sparse.py | 3 ++ .../attention/backends/sparse/sol/backend.py | 9 ++-- .../visual_gen/attention/backends/trtllm.py | 5 +- .../sol/test_sol_attention.py | 52 ++++++++++++++++--- 4 files changed, 58 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py index 6d754d7a8747..a40f828f46ac 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py @@ -551,6 +551,9 @@ def _forward_contiguous( out=forward_args.output.view_as(query), ) + # Planning records an external CUDA event that Dynamo cannot construct, so + # compiled callers run the PrimTS call eagerly between graph segments. + @torch.compiler.disable def forward( self, q: torch.Tensor, diff --git a/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/sol/backend.py b/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/sol/backend.py index f8070fb9fecd..1482ddad5948 100644 --- a/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/sol/backend.py +++ b/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/sol/backend.py @@ -96,8 +96,8 @@ class SOLCuTeDSLAttention(AttentionBackend): Inputs the kernel cannot serve (shape, dtype, architecture) are delegated to dense attention up front by ``_run_sol_attn_bthd``; a kernel error is not recovered from and propagates. This class adds the ``dense_layers`` guard - (evaluated at construction time, no external plumbing needed) and forwards - the routing knobs from config. + (resolved for this layer at construction) and forwards the routing knobs + from config. """ def __init__( @@ -136,7 +136,8 @@ def __init__( self.tau = sparse_params.tau self.thresh_type = sparse_params.thresh_type self.disabled_until_timestep = sparse_params.disabled_until_timestep - self.dense_layers = sparse_params.dense_layers + # Resolved here so compiled callers do not specialize on layer_idx. + self.is_dense_layer = layer_idx in sparse_params.dense_layers # SOL's dense steps must run the backend the user selected. Without # this they ran torch SDPA while a `backend: CUTEDSL` baseline ran @@ -284,7 +285,7 @@ def _can_serve(self, q: torch.Tensor, k: torch.Tensor, **kwargs: Any) -> bool: != PredefinedAttentionMask.FULL ): return False - if self.layer_idx in self.dense_layers: + if self.is_dense_layer: return False if self.disabled_until_timestep is not None and self._dense_by_step(kwargs.get("timestep")): return False diff --git a/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py b/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py index c1f0ef0ef0fc..a261df5e9df6 100644 --- a/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py @@ -244,6 +244,9 @@ def __init__( "(TLLM_FMHA_LIBS=+prims_ts), which is not enabled or not available on this " "platform." ) + # Compiled callers guard on this flag; testing layer_idx per call would specialize, and + # recompile, the attention frames for every layer. + self.is_dense_layer = self.layer_idx in self.dense_layers @property def timestep_cutoff(self) -> Optional[float]: @@ -283,7 +286,7 @@ def should_use_sparse(self, sparse_attn_phase: Optional[int]) -> bool: dense; a call without a phase (no cutoff or no timestep) is sparse. """ - if self.layer_idx in self.dense_layers: + if self.is_dense_layer: return False return sparse_attn_phase != 0 diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_attention.py b/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_attention.py index 31158399a46b..d12e21369212 100644 --- a/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_attention.py +++ b/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_attention.py @@ -31,6 +31,7 @@ import pytest import torch from pydantic import ValidationError +from torch._dynamo.testing import CompileCounter from tensorrt_llm._torch.attention.backends.fmha.prims_ts_block_sparse import PrimsTSBlockSparseFmha from tensorrt_llm._torch.attention.backends.interface import ( @@ -90,6 +91,7 @@ def _make_backend(params: SolParams, *, layer_idx: int = 1) -> SOLTrtllmAttentio backend.sparse_params = None backend._fmha_manager = SimpleNamespace(fmha_libs=[object.__new__(PrimsTSBlockSparseFmha)]) backend.sol_params = params + backend.is_dense_layer = layer_idx in params.dense_layers backend.metadata = TrtllmAttentionMetadata( device=torch.device("cpu"), attention_metadata_state={} ) @@ -198,6 +200,7 @@ def _stub_backend( *, seq_len: int = 64, unsupported_reason: str | None = None, + layer_idx: int = 1, ) -> tuple[SOLTrtllmAttention, SimpleNamespace]: """Backend whose predictor functions are recorded mocks.""" predictor = SimpleNamespace( @@ -206,7 +209,7 @@ def _stub_backend( ) monkeypatch.setattr(sol_predictor, "support_reason", predictor.support_reason) monkeypatch.setattr(sol_predictor, "predict", predictor.predict) - return _make_backend(params or SolParams(tau=1.0)), predictor + return _make_backend(params or SolParams(tau=1.0), layer_idx=layer_idx), predictor # --------------------------------------------------------------------------- TRTLLM backend @@ -290,6 +293,43 @@ def _base_init(self, **kwargs) -> None: assert "block_sparse_attn_predict" in SOLTrtllmAttention.__dict__ +@_CPU_ONLY +def test_sol_sparse_decision_shares_compiled_code_across_layers(monkeypatch) -> None: + """Compiled callers of the sparse decision reuse one graph for all sparse layers. + + Each backend resolves whether it is a dense layer when it is built, so a compiled + caller guards on that flag rather than on ``layer_idx``. + """ + + def _core_init(self, **kwargs) -> None: + self.layer_idx = kwargs["layer_idx"] + self.head_dim = kwargs["head_dim"] + self.q_scaling = 1.0 + + monkeypatch.setattr(CoreTrtllmAttention, "__init__", _core_init) + params = SolParams(dense_layers=frozenset({0})) + backends = [ + SOLTrtllmAttention( + layer_idx=layer_idx, + num_heads=2, + head_dim=128, + sparse_params=params, + attention_metadata_state=create_attention_metadata_state(), + ) + for layer_idx in range(4) + ] + counter = CompileCounter() + + @torch.compile(backend=counter) + def route(backend: SOLTrtllmAttention, x: torch.Tensor) -> torch.Tensor: + return x + 1 if backend.should_use_sparse(None) else x - 1 + + routed = [route(backend, torch.zeros(1)).item() for backend in backends] + + assert routed == [-1.0, 1.0, 1.0, 1.0] + assert counter.frame_count == 2 + + @_CPU_ONLY def test_sol_backend_sparse_phase_emits_proxy_bitmask_carrier(monkeypatch) -> None: batch_size, seq_len, num_heads = 1, 65, 2 @@ -405,8 +445,7 @@ def test_sol_dense_policy_returns_no_routes_without_predicting( timestep: float | None, ) -> None: q = _bshd() - backend, predictor = _stub_backend(monkeypatch, params) - backend.layer_idx = layer_idx + backend, predictor = _stub_backend(monkeypatch, params, layer_idx=layer_idx) assert _predict(backend, q, q, q, timestep=timestep) is None predictor.support_reason.assert_not_called() @@ -615,8 +654,9 @@ def _fail_if_called(*args, **kwargs): raise AssertionError("kernel must not be invoked for a dense_layers-forced layer") monkeypatch.setattr(sol_backend, "_sol_attn_run", _fail_if_called) - backend = SOLCuTeDSLAttention(layer_idx=3, num_heads=2, head_dim=16) - backend.dense_layers = frozenset({3}) + backend = SOLCuTeDSLAttention( + layer_idx=3, num_heads=2, head_dim=16, sparse_params=SolParams(dense_layers=frozenset({3})) + ) q = k = v = torch.randn(1, 4, 2, 16) out = backend.forward(q, k, v) @@ -697,7 +737,7 @@ def _spy(*args, **kwargs): assert calls["n"] == 1, "dense prefix was not delegated to the CuTeDSL dense kernel" backend, calls = _make() - backend.dense_layers = frozenset({backend.layer_idx}) + backend.is_dense_layer = True backend.forward(q, k, v) assert calls["n"] == 1, "dense_layers guard was not delegated to the CuTeDSL dense kernel" From 7792d302b05fef9cc9cb0d3562f86a83a1051d2d Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Sun, 11 Oct 2026 03:37:18 +0000 Subject: [PATCH 03/13] [None][feat] Support SageAttention in the PrimTS block-sparse FMHA The contiguous PrimsTSBlockSparseFmha path now consumes the SageAttention recipe from AttentionForwardArgs: - After route prediction, Q/K/V are quantized with FlashInfer's trtllm_sage_attention_quantize, honoring smooth_k. - Proxy K/V means are quantized by a new Triton op, trtllm::quantize_sage_summaries: one scale per K mean, while V means reuse the V channel scales so exact and proxy routes share one PV scale. - Plans carry explicit K and V dtypes and are keyed by the recipe. Sage needs BF16 input; paged block-sparse attention rejects it, and PrimTS plans check which tile profiles take it. The dense PrimsTSFmha still rejects Sage, now in its own support check instead of the shared feature gate. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- .../attention/ATTENTION_DEVELOPER_GUIDE.md | 8 + .../attention/backends/fmha/prims_ts.py | 12 +- .../backends/fmha/prims_ts_block_sparse.py | 128 +++++++++++++++- .../_torch/attention/kernels/sage_quant.py | 135 +++++++++++++++++ .../test_lists/test-db/l0_b200.yml | 1 + .../attention/kernels/test_sage_quant.py | 139 ++++++++++++++++++ 6 files changed, 412 insertions(+), 11 deletions(-) create mode 100644 tensorrt_llm/_torch/attention/kernels/sage_quant.py create mode 100644 tests/unittest/_torch/attention/kernels/test_sage_quant.py diff --git a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index 71f400ec5348..b871bc642cc3 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md @@ -233,6 +233,14 @@ routes, and optional token-validity bits mask ragged KV tails. Plans contain only static format, proxy, geometry, and capacity choices; every run receives the live routes, summaries, validity bits, page tables, and sequence lengths. +The contiguous block-sparse path also consumes the SageAttention recipe in +`AttentionForwardArgs`. Prediction sees the original BF16 Q/K/V. The adapter +then quantizes Q/K/V with FlashInfer's Sage quantizer, quantizes the summaries +with one scale per K mean and the original V channel scales, and passes all +scales to PrimTS. The recipe and quantized dtypes are part of the plan key. +Sage requires BF16 input, because the quantizer's FP16 reciprocal scale can +overflow on zero or tiny Q/K blocks. Paged block-sparse attention rejects Sage. + `PrimsTSBlockSparseFmha` keeps its planned wrappers in the attention metadata (`TrtllmAttentionMetadata.fmha_plan_caches`), like the FlashInfer wrappers and the MSA plans. Every layer that runs with one metadata object plans each static diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py index 5268e72b5bd6..29bde4a063d4 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py @@ -149,12 +149,6 @@ def get_attention_feature_unsupported_reason( return "custom attention masks are not supported." if forward_args.enable_dsv4_epilogue_fusion: return "DSv4 epilogue fusion is not supported." - if ( - forward_args.sage_attn_num_elts_per_blk_q > 0 - or forward_args.sage_attn_num_elts_per_blk_k > 0 - or forward_args.sage_attn_num_elts_per_blk_v > 0 - ): - return "SageAttention is not supported." return None @@ -322,6 +316,12 @@ def _is_supported_with_reason( feature_reason = get_attention_feature_unsupported_reason(meta, fwd) if feature_reason is not None: return False, feature_reason + if ( + fwd.sage_attn_num_elts_per_blk_q > 0 + or fwd.sage_attn_num_elts_per_blk_k > 0 + or fwd.sage_attn_num_elts_per_blk_v > 0 + ): + return False, "SageAttention is not supported." if meta.kv_cache_manager is None: return self._is_supported_without_kv_cache(q, k, v, attn, meta, fwd) storage_reason = get_paged_kv_storage_unsupported_reason(attn, meta) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py index a40f828f46ac..9a6c6d4b4a15 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py @@ -16,10 +16,11 @@ """TRT-LLM FMHA adapter for the vendored PrimTS block-sparse kernels.""" import math -from dataclasses import dataclass, fields +from dataclasses import dataclass, fields, replace from typing import TYPE_CHECKING, Literal, cast import torch +from flashinfer import trtllm_sage_attention_quantize from tensorrt_llm._torch.attention.backends.interface import ( AttentionForwardArgs, @@ -29,7 +30,12 @@ from tensorrt_llm._torch.attention.backends.prims_ts._block_sparse.config import ( _validate_block_sparse_static_profile, ) +from tensorrt_llm._torch.attention.backends.prims_ts.sage import ( + SageAttentionConfig, + SageAttentionParams, +) from tensorrt_llm._torch.attention.backends.sparse.params import BlockSparseForwardInputs +from tensorrt_llm._torch.attention.kernels.sage_quant import quantize_sage_summaries from tensorrt_llm.logger import logger from .interface import FmhaPhase @@ -82,9 +88,12 @@ class _BlockSparsePlanKey: max_blocks_per_row: int mask_type: Literal["dense", "causal"] dtype: torch.dtype + qk_dtype: torch.dtype + v_dtype: torch.dtype use_kv_valid_bits: bool sparse_format: Literal["bsr", "bitmask"] use_proxy_routes: bool + sage: SageAttentionConfig | None def unsupported_reason(self) -> str | None: try: @@ -99,11 +108,13 @@ def unsupported_reason(self) -> str | None: kv_block_size=self.kv_block_size, use_kv_valid_bits=self.use_kv_valid_bits, mask_type=self.mask_type, - q_dtype=self.dtype, - kv_dtype=self.dtype, + q_dtype=self.qk_dtype, + kv_dtype=self.qk_dtype, + v_dtype=self.v_dtype, output_dtype=self.dtype, max_blocks_per_row=self.max_blocks_per_row, page_size=self.page_size, + sage=self.sage, ) except (ValueError, NotImplementedError, OverflowError) as error: return str(error) @@ -129,16 +140,19 @@ def plan(self) -> "BlockSparseTSWrapper | BlockSparsePagedTSWrapper": "max_blocks_per_row": self.max_blocks_per_row, "use_kv_valid_bits": self.use_kv_valid_bits, "mask_type": self.mask_type, - "q_data_type": self.dtype, - "kv_data_type": self.dtype, + "q_data_type": self.qk_dtype, "o_data_type": self.dtype, } if paged: plan_args += (self.page_size,) + plan_kwargs.update(kv_data_type=self.qk_dtype) else: plan_kwargs.update( sparse_format=self.sparse_format, use_proxy_routes=self.use_proxy_routes, + k_data_type=self.qk_dtype, + v_data_type=self.v_dtype, + sage_config=self.sage, ) wrapper.plan(*plan_args, **plan_kwargs) return wrapper @@ -184,6 +198,41 @@ def _uniform_seq_len_q( return seq_len_q +def _sage_recipe( + forward_args: AttentionForwardArgs, +) -> tuple[SageAttentionConfig, torch.dtype] | None: + """Return the forward's SageAttention plan config and quantized Q/K dtype, if any.""" + q_block = forward_args.sage_attn_num_elts_per_blk_q + k_block = forward_args.sage_attn_num_elts_per_blk_k + if q_block <= 0 and k_block <= 0 and forward_args.sage_attn_num_elts_per_blk_v <= 0: + return None + # Every proxy K summary is quantized with its own scale. + config = SageAttentionConfig(q_block_size=q_block, k_block_size=k_block, k_summary_block_size=1) + return config, torch.int8 if forward_args.sage_attn_qk_int8 else torch.float8_e4m3fn + + +def _sage_unsupported_reason( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + forward_args: AttentionForwardArgs, +) -> str | None: + """Return why the forward's Sage recipe cannot quantize these contiguous inputs.""" + if _sage_recipe(forward_args) is None: + return None + if ( + forward_args.sage_attn_num_elts_per_blk_q not in (1, 4, 16) + or forward_args.sage_attn_num_elts_per_blk_k not in (1, 4, 16) + or forward_args.sage_attn_num_elts_per_blk_v != 1 + ): + return "Sage block-sparse attention requires Q/K blocks of 1, 4, or 16 and V block size 1" + if q.dtype != torch.bfloat16 or k.dtype != q.dtype or v.dtype != q.dtype: + # The quantizer computes reciprocal scales in the input dtype, which + # overflows FP16 on zero or tiny Q/K blocks such as VSA padding. + return "Sage block-sparse attention requires matching BF16 Q/K/V inputs" + return None + + class PrimsTSBlockSparseFmha(PrimsTSFmha): """Contiguous context and fixed-Q paged generation block-sparse FMHA.""" @@ -281,10 +330,15 @@ def _make_plan_key( kv_capacity: int, page_size: int | None, mask_type: Literal["dense", "causal"], + sage: tuple[SageAttentionConfig, torch.dtype] | None = None, ) -> _BlockSparsePlanKey: max_blocks_per_row = inputs.max_blocks_per_row if max_blocks_per_row is None: max_blocks_per_row = math.ceil(kv_capacity / inputs.kv_block_size) + sage_config, qk_dtype, v_dtype = None, q.dtype, q.dtype + if sage is not None: + sage_config, qk_dtype = sage + v_dtype = torch.float8_e4m3fn return _BlockSparsePlanKey( device=q.device, batch_size=batch_size, @@ -299,9 +353,12 @@ def _make_plan_key( max_blocks_per_row=max_blocks_per_row, mask_type=mask_type, dtype=q.dtype, + qk_dtype=qk_dtype, + v_dtype=v_dtype, use_kv_valid_bits=inputs.kv_valid_bits is not None, sparse_format=inputs.sparse_format, use_proxy_routes=inputs.use_proxy_routes, + sage=sage_config, ) def _get_or_plan_wrapper( @@ -336,6 +393,9 @@ def _contiguous_unsupported_reason( return "contiguous Q/K/V must have position embedding applied before attention" if forward_args.cu_q_seqlens is not None or forward_args.cu_kv_seqlens is not None: return "packed variable-length Q/KV inputs are not supported" + sage_reason = _sage_unsupported_reason(q, k, v, forward_args) + if sage_reason is not None: + return sage_reason inputs = _get_block_sparse_inputs(forward_args) assert inputs is not None mask_type = self._get_prims_mask_type(forward_args) @@ -353,6 +413,7 @@ def _contiguous_unsupported_reason( kv_capacity=int(k.shape[0]) // batch_size, page_size=None, mask_type=mask_type, + sage=_sage_recipe(forward_args), ) return key.unsupported_reason() @@ -366,6 +427,8 @@ def _paged_unsupported_reason( assert inputs is not None if inputs.sparse_format != "bsr" or inputs.use_proxy_routes: return "paged block-sparse attention only supports BSR exact routes" + if _sage_recipe(forward_args) is not None: + return "paged block-sparse attention does not support SageAttention" if not forward_args.is_fused_qkv: return "paged block-sparse attention requires fused QKV input" if ( @@ -535,8 +598,62 @@ def _forward_contiguous( kv_capacity=int(key_states.shape[1]), page_size=None, mask_type=self._get_prims_mask_type(forward_args), + sage=_sage_recipe(forward_args), ) wrapper = cast("BlockSparseTSWrapper", self._get_or_plan_wrapper(key, metadata)) + sage_params = None + if key.sage is not None: + # Prediction has consumed the unquantized tensors. Device cumulative + # lengths keep the per-sequence scale layout without a host-to-device + # copy during CUDA Graph capture. + cu_seqlens_q = torch.arange( + 0, + (batch_size + 1) * key.seq_len_q, + key.seq_len_q, + dtype=torch.int32, + device=q.device, + ) + cu_seqlens_kv = torch.arange( + 0, + (batch_size + 1) * key.kv_capacity, + key.kv_capacity, + dtype=torch.int32, + device=q.device, + ) + quantized = trtllm_sage_attention_quantize( + query.flatten(0, 1), + key_states.flatten(0, 1), + value_states.flatten(0, 1), + q_block_size=key.sage.q_block_size, + k_block_size=key.sage.k_block_size, + qk_quant_dtype=key.qk_dtype, + cum_seq_lens_q=cu_seqlens_q, + cum_seq_lens_kv=cu_seqlens_kv, + smooth_k=forward_args.sage_attn_smooth_k, + ) + q_quant, k_quant, v_quant, q_scale, k_scale, v_scale = quantized[:6] + query = q_quant.view_as(query) + key_states = k_quant.view_as(key_states) + value_states = v_quant.view_as(value_states) + k_summary_scale = None + if inputs.use_proxy_routes: + k_summary = inputs.k_summary + if forward_args.sage_attn_smooth_k: + # Proxy logits must shift by the same per-row constant as the + # smoothed exact keys, so the summaries subtract the same mean. + k_summary = (k_summary.float() - quantized[6]).to(k_summary.dtype) + k_summary, v_summary, k_summary_scale = quantize_sage_summaries( + k_summary, inputs.v_summary, v_scale, key.qk_dtype + ) + inputs = replace(inputs, k_summary=k_summary, v_summary=v_summary) + sage_params = SageAttentionParams( + q_scale=q_scale, + k_scale=k_scale, + # PrimTS takes one per-channel V scale row per batch entry; the + # quantizer collects a single row over the whole batch. + v_scale=v_scale.expand(batch_size, -1, -1).contiguous(), + k_summary_scale=k_summary_scale, + ) wrapper.run( query, key_states, @@ -549,6 +666,7 @@ def _forward_contiguous( kv_valid_bits=inputs.kv_valid_bits, sm_scale=self._get_bmm1_scale(self.attn), out=forward_args.output.view_as(query), + sage=sage_params, ) # Planning records an external CUDA event that Dynamo cannot construct, so diff --git a/tensorrt_llm/_torch/attention/kernels/sage_quant.py b/tensorrt_llm/_torch/attention/kernels/sage_quant.py new file mode 100644 index 000000000000..570f4a93fd02 --- /dev/null +++ b/tensorrt_llm/_torch/attention/kernels/sage_quant.py @@ -0,0 +1,135 @@ +# 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. +"""Quantization of the K/V block means consumed by Sage block-sparse attention.""" + +import torch +import triton +import triton.language as tl +from torch.library import triton_op, wrap_triton +from triton.language.extra.cuda import libdevice + +__all__ = ["quantize_sage_summaries"] + + +@triton.jit +def _quantize_sage_summaries_kernel( + k_ptr, + v_ptr, + v_scale_ptr, + quant_k_ptr, + quant_v_ptr, + k_scale_ptr, + num_rows, + blocks, + heads, + scale_stride, + DIM: tl.constexpr, + INT8: tl.constexpr, + ROWS: tl.constexpr, + WIDTH: tl.constexpr, +): + rows = tl.program_id(0) * ROWS + tl.arange(0, ROWS) + head = tl.program_id(1) + channels = tl.arange(0, WIDTH) + valid_rows = rows < num_rows + valid = valid_rows[:, None] & (channels < DIM)[None, :] + offsets = (rows[:, None] * heads + head) * DIM + channels[None, :] + k = tl.load(k_ptr + offsets, valid, other=0).to(tl.float32) + v = tl.load(v_ptr + offsets, valid, other=0).to(tl.float32) + + maximum = tl.maximum(tl.max(tl.abs(k), axis=1), 1e-3) + divisor: tl.constexpr = 126.9 if INT8 else 448.0 + k_scale = tl.div_rn(maximum, divisor) + normalized_k = tl.div_rn(k, k_scale[:, None]) + if INT8: + quant_k = tl.minimum(tl.maximum(libdevice.nearbyint(normalized_k), -127.0), 127.0) + else: + quant_k = tl.minimum(tl.maximum(normalized_k, -448.0), 448.0) + v_scale = tl.load(v_scale_ptr + head * DIM + channels, channels < DIM, other=1.0) + quant_v = tl.minimum(tl.maximum(tl.div_rn(v, v_scale[None, :]), -448.0), 448.0) + tl.store(quant_k_ptr + offsets, quant_k, valid) + tl.store(quant_v_ptr + offsets, quant_v, valid) + + # Sequence b starts at scale slot b * (blocks + 1): the flat layout leaves one + # unused slot between consecutive sequences. + tl.store(k_scale_ptr + head * scale_stride + rows + rows // blocks, k_scale, valid_rows) + + +@triton_op("trtllm::quantize_sage_summaries", mutates_args={}) +def quantize_sage_summaries( + k_summary: torch.Tensor, + v_summary: torch.Tensor, + v_scale: torch.Tensor, + qk_quant_dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Quantize key/value block means for the proxy routes of Sage attention. + + Every key mean gets its own scale; value means reuse the per-channel scales of + the original V, so exact and proxy routes share one PV dequantization. + + Args: + k_summary: Contiguous FP16/BF16 CUDA means, ``[batch, blocks, heads, head_dim]``. + v_summary: Value means with the same shape, dtype, and device as ``k_summary``. + v_scale: Contiguous positive FP32 ``[heads, head_dim]`` dequantization scales + from quantizing the original V. + qk_quant_dtype: ``torch.int8`` or ``torch.float8_e4m3fn`` for the key means. + + Returns: + Quantized key means, FP8 E4M3 value means, and FP32 key scales of shape + ``[heads, batch * blocks + batch - 1]``. Key scale slot ``b * (blocks + 1) + n`` + belongs to mean ``[b, n]``; the unused slot between sequences is not written. + """ + if k_summary.ndim != 4 or min(k_summary.shape) <= 0: + raise ValueError("k_summary must have nonempty [batch, blocks, heads, head_dim] shape") + if k_summary.device.type != "cuda" or k_summary.dtype not in (torch.float16, torch.bfloat16): + raise ValueError("k_summary must be a CUDA FP16 or BF16 tensor") + if ( + v_summary.shape != k_summary.shape + or v_summary.dtype != k_summary.dtype + or v_summary.device != k_summary.device + ): + raise ValueError("key and value summaries must have matching shape, dtype, and device") + if v_scale.shape != k_summary.shape[2:] or v_scale.dtype != torch.float32: + raise ValueError("v_scale must be FP32 with shape [heads, head_dim]") + if v_scale.device != k_summary.device: + raise ValueError("v_scale and summaries must be on the same CUDA device") + if not (k_summary.is_contiguous() and v_summary.is_contiguous() and v_scale.is_contiguous()): + raise ValueError("summaries and v_scale must be contiguous") + if qk_quant_dtype not in (torch.int8, torch.float8_e4m3fn): + raise ValueError("qk_quant_dtype must be torch.int8 or torch.float8_e4m3fn") + + batch, blocks, heads, dim = k_summary.shape + scale_stride = batch * blocks + batch - 1 + quant_k = torch.empty_like(k_summary, dtype=qk_quant_dtype) + quant_v = torch.empty_like(v_summary, dtype=torch.float8_e4m3fn) + k_scale = torch.empty((heads, scale_stride), device=k_summary.device, dtype=torch.float32) + wrap_triton(_quantize_sage_summaries_kernel)[(triton.cdiv(batch * blocks, 4), heads)]( + k_summary, + v_summary, + v_scale, + quant_k, + quant_v, + k_scale, + batch * blocks, + blocks, + heads, + scale_stride, + DIM=dim, + INT8=qk_quant_dtype == torch.int8, + ROWS=4, + WIDTH=triton.next_power_of_2(dim), + num_warps=4, + ) + return quant_k, quant_v, k_scale diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index e5fcb9c9de8f..e0360e1def5b 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -339,6 +339,7 @@ l0_b200: - unittest/_torch/visual_gen/sparse_attention/sol/test_sol_parity.py - unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor.py - unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor_kernels.py + - unittest/_torch/attention/kernels/test_sage_quant.py - unittest/_torch/visual_gen/test_attention_flashinfer.py - unittest/_torch/visual_gen/test_attention_trtllm_sage.py - unittest/_torch/visual_gen/test_attention_cudnn.py diff --git a/tests/unittest/_torch/attention/kernels/test_sage_quant.py b/tests/unittest/_torch/attention/kernels/test_sage_quant.py new file mode 100644 index 000000000000..979f194d4378 --- /dev/null +++ b/tests/unittest/_torch/attention/kernels/test_sage_quant.py @@ -0,0 +1,139 @@ +# 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. +"""Exact parity of the Sage K/V summary quantizer with a PyTorch reference. + +Scale slots are located with the PrimTS flat-layout helpers that the attention +kernel uses to read them. +""" + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.prims_ts.sage import flat_scale_numel, flat_scale_slot +from tensorrt_llm._torch.attention.kernels.sage_quant import quantize_sage_summaries + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def _inputs( + shape: tuple[int, ...], dtype: torch.dtype +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + generator = torch.Generator(device="cuda").manual_seed(43) + k = torch.randn(shape, device="cuda", generator=generator) + v = torch.randn(shape, device="cuda", generator=generator) + magnitudes = torch.logspace(-5, 2, shape[1], device="cuda").view(1, -1, 1, 1) + k, v = (k * magnitudes).to(dtype), (v * magnitudes).to(dtype) + # Raw V has larger extrema than its block means, and the attention epilogue + # dequantizes exact and mean V with the same channel scales. + v_scale = v.float().abs().amax(dim=(0, 1)).clamp_min(1e-3) * 2 / 448.0 + return k, v, v_scale + + +def _assert_matches_reference( + actual: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + k: torch.Tensor, + v: torch.Tensor, + v_scale: torch.Tensor, + quant_dtype: torch.dtype, +) -> None: + batch, blocks, heads, _ = k.shape + divisor = 126.9 if quant_dtype == torch.int8 else 448.0 + maximum = k.float().abs().amax(dim=-1).clamp_min(1e-3) + # Tensor division keeps FP32 round-to-nearest; dividing by a Python scalar + # multiplies by its rounded reciprocal and can move an FP8 rounding tie. + scales = maximum / torch.full_like(maximum, divisor) + normalized = k.float() / scales.unsqueeze(-1) + if quant_dtype == torch.int8: + quant_k = normalized.round().clamp(-127, 127).to(quant_dtype) + else: + quant_k = normalized.clamp(-448, 448).to(quant_dtype) + quant_v = (v.float() / v_scale).clamp(-448, 448).to(torch.float8_e4m3fn) + slots = flat_scale_slot( + torch.arange(batch, device=k.device)[:, None], + torch.arange(blocks, device=k.device)[None, :], + blocks, + 0, + ).flatten() + + actual_k, actual_v, actual_scales = actual + assert actual_scales.shape == (heads, flat_scale_numel(batch, blocks, 1)) + torch.testing.assert_close(actual_k.float(), quant_k.float(), rtol=0, atol=0) + torch.testing.assert_close(actual_v.float(), quant_v.float(), rtol=0, atol=0) + torch.testing.assert_close( + actual_scales[:, slots], scales.permute(2, 0, 1).reshape(heads, -1), rtol=0, atol=0 + ) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize("quant_dtype", [torch.int8, torch.float8_e4m3fn]) +@pytest.mark.parametrize("shape", [(2, 9, 3, 64), (3, 5, 2, 7), (1, 129, 2, 128)]) +# A 100x smaller V scale saturates the quantized V means at the E4M3 maximum. +@pytest.mark.parametrize("v_scale_divisor", [1.0, 100.0]) +def test_quantize_sage_summaries_matches_reference( + shape: tuple[int, ...], dtype: torch.dtype, quant_dtype: torch.dtype, v_scale_divisor: float +) -> None: + k, v, v_scale = _inputs(shape, dtype) + v_scale = v_scale / v_scale_divisor + actual = quantize_sage_summaries(k, v, v_scale, quant_dtype) + assert [result.dtype for result in actual] == [quant_dtype, torch.float8_e4m3fn, torch.float32] + _assert_matches_reference(actual, k, v, v_scale, quant_dtype) + + +@pytest.mark.parametrize("quant_dtype", [torch.int8, torch.float8_e4m3fn]) +def test_quantize_sage_summaries_zero_blocks_use_scale_floor(quant_dtype: torch.dtype) -> None: + k = torch.zeros(2, 9, 2, 64, device="cuda", dtype=torch.bfloat16) + v = torch.zeros_like(k) + v_scale = torch.full((2, 64), 1e-3 / 448.0, device="cuda") + actual = quantize_sage_summaries(k, v, v_scale, quant_dtype) + _assert_matches_reference(actual, k, v, v_scale, quant_dtype) + + +def test_quantize_sage_summaries_rejects_layouts_the_kernel_cannot_index() -> None: + k, v, v_scale = _inputs((2, 9, 3, 64), torch.bfloat16) + cases = ( + ((k, v[:, :1], v_scale), "matching shape, dtype, and device"), + ((k, v, v_scale[:1]), "FP32 with shape"), + ((k.transpose(0, 1), v.transpose(0, 1), v_scale), "contiguous"), + ) + for args, message in cases: + with pytest.raises(ValueError, match=message): + quantize_sage_summaries(*args, torch.int8) + + +@pytest.mark.parametrize("quant_dtype", [torch.int8, torch.float8_e4m3fn]) +def test_quantize_sage_summaries_under_compile_and_cuda_graph(quant_dtype: torch.dtype) -> None: + k, v, v_scale = _inputs((2, 9, 2, 64), torch.bfloat16) + + def quantize( + key: torch.Tensor, value: torch.Tensor, value_scale: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return quantize_sage_summaries(key, value, value_scale, quant_dtype) + + _assert_matches_reference( + torch.compile(quantize, fullgraph=True)(k, v, v_scale), k, v, v_scale, quant_dtype + ) + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + quantize(k, v, v_scale) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = quantize(k, v, v_scale) + k.mul_(2) + v.mul_(2) + graph.replay() + _assert_matches_reference(captured, k, v, v_scale, quant_dtype) From 679eb70f8f9732371c9448cce4b6205906f8d8bc Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Sun, 11 Oct 2026 03:39:18 +0000 Subject: [PATCH 04/13] [None][feat] Enable SageAttention for VisualGen VSA and SOL on TRTLLM - AttentionConfig accepts a Sage quant_attention_config together with VSA or SOL when the backend is TRTLLM; other backends still reject the combination. - The TRTLLM VisualGen backend, VSA and SOL pass the recipe to the block-sparse path instead of rejecting it, so PrimsTSBlockSparseFmha quantizes after prediction. VSA falls back to the dense fine stage when Sage meets non-BF16 inputs. - Document the composition and test Sage block-sparse attention through the VisualGen TRTLLM backend. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- .../features/visualgen-sparse-attention.md | 12 +- .../attention/backends/sparse/sol/backend.py | 2 - .../attention/backends/sparse/vsa/backend.py | 10 +- .../visual_gen/attention/backends/trtllm.py | 4 - tensorrt_llm/visual_gen/args.py | 18 +- .../test_lists/test-db/l0_b200.yml | 1 + .../sol/test_sol_attention.py | 29 +- .../test_sage_block_sparse.py | 407 ++++++++++++++++++ .../vsa/test_vsa_attention.py | 2 +- .../_torch/visual_gen/test_visual_gen_args.py | 34 +- 10 files changed, 462 insertions(+), 57 deletions(-) create mode 100644 tests/unittest/_torch/visual_gen/sparse_attention/test_sage_block_sparse.py diff --git a/docs/source/features/visualgen-sparse-attention.md b/docs/source/features/visualgen-sparse-attention.md index 9ab465921aec..095790c27520 100644 --- a/docs/source/features/visualgen-sparse-attention.md +++ b/docs/source/features/visualgen-sparse-attention.md @@ -94,7 +94,7 @@ User configuration is supplied through Python or YAML and controls how the check `threshold_scale_factor` and `target_sparsity` are alternatives: if both are present, `threshold_scale_factor` takes precedence and the calibration formula is not used. User-provided `target_sparsity` and `disabled_until_timestep` override checkpoint defaults. Checkpoint `ignore` patterns always disable Skip Softmax Attention for matching layers. -Skip Softmax Attention works with both the **TRTLLM** and **CUTEDSL** attention backends in VisualGen. Set `attention_config.backend` to either when enabling it. On CUTEDSL, Skip Softmax Attention can also be combined with `quant_attention_config`'s block-scaled Q/K recipes (MXFP8, NVFP4); VSA and SOL cannot be combined with quantized attention on any backend. +Skip Softmax Attention works with both the **TRTLLM** and **CUTEDSL** attention backends in VisualGen. Set `attention_config.backend` to either when enabling it. On CUTEDSL, Skip Softmax Attention can also be combined with `quant_attention_config`'s block-scaled Q/K recipes (MXFP8, NVFP4). VSA and SOL can be combined with quantized attention only on TRTLLM, through its SageAttention recipes. #### Mapping `disabled_until_timestep` to Actual Denoising Steps @@ -234,9 +234,11 @@ softmax. VisualGen serves it on two backends behind one `SolAttentionConfig`: attention metadata and the `timestep` in the forward arguments; dense layers and dense timestep phases return no routes. The VisualGen wrapper compacts fused projection split views once and shares those tensors between - prediction and the block-sparse FMHA. Cross-attention, context parallelism, - attention quantization and unsupported tensor envelopes raise an error - instead of silently falling back to dense attention. + prediction and the block-sparse FMHA. With a SageAttention recipe, the FMHA + quantizes Q/K/V and the proxy means after prediction; the dense timestep + phases use the dense SageAttention kernel. Cross-attention, context + parallelism and unsupported tensor envelopes raise an error instead of + silently falling back to dense attention. - **CUTEDSL** runs the fused kernel vendored from the reference implementation (`SOLCuTeDSLAttention`), which folds routing, sparse computation and the approximation correction into one online-softmax pass. It runs on datacenter @@ -318,7 +320,7 @@ Requirements: - A VSA-fine-tuned checkpoint such as [`FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers`](https://huggingface.co/FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers). Standard Wan checkpoints do not carry the learned VSA gates. - The `CUTEDSL` or `TRTLLM` attention backend. `CUTEDSL` uses the CuTe DSL fine-stage kernel; `TRTLLM` lowers the selected cubes through the generic PrimTS block-sparse FMHA contract. - A supported CUDA device and tensor shape for the selected block-sparse kernel. When that kernel is unavailable or the input is outside its envelope, the fine branch runs the compact Q/K/V through that backend's dense path (SDPA for `CUTEDSL`, TRTLLM attention for `TRTLLM`) and the VSA post-processing still applies. -- VSA cannot be combined with `quant_attention_config`. +- On `TRTLLM`, VSA can be combined with a SageAttention `quant_attention_config` for BF16 models; the block-sparse fine stage then runs with quantized Q/K/V. `CUTEDSL` rejects `quant_attention_config` with VSA. - VSA is not compatible with Ring attention or Attention2D (it does not produce per-split LSE). Ulysses is supported. Configure VSA with `VideoSparseAttentionConfig` and either backend: diff --git a/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/sol/backend.py b/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/sol/backend.py index 1482ddad5948..c2c5bef3c495 100644 --- a/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/sol/backend.py +++ b/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/sol/backend.py @@ -363,8 +363,6 @@ def block_sparse_attn_predict( if not self.should_use_sparse(forward_args.sparse_attn_phase): return None - if self.quant_attention_config is not None: - raise ValueError("SOL sparse execution does not support quant_attention_config") if not any( isinstance(fmha, PrimsTSBlockSparseFmha) for fmha in self._fmha_manager.fmha_libs ): diff --git a/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/vsa/backend.py b/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/vsa/backend.py index 386eaef42ad9..de66d8637ce2 100644 --- a/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/vsa/backend.py +++ b/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/vsa/backend.py @@ -68,6 +68,8 @@ def _get_unsupported_primts_reason( k: torch.Tensor, v: torch.Tensor, metadata: VSAMetadata, + *, + sage: bool, ) -> str | None: if q.shape != k.shape or q.shape != v.shape: return "VSA PrimTS requires matching MHA Q/K/V shapes" @@ -75,6 +77,8 @@ def _get_unsupported_primts_reason( return f"VSA PrimTS requires CUDA tensors, got {q.device}" if q.dtype not in (torch.float16, torch.bfloat16): return f"VSA PrimTS requires FP16 or BF16 tensors, got {q.dtype}" + if sage and q.dtype != torch.bfloat16: + return f"VSA PrimTS Sage attention requires BF16 tensors, got {q.dtype}" batch_size, seq_len, num_heads, head_dim = map(int, q.shape) if min(batch_size, seq_len, num_heads, metadata.num_cubes) <= 0: return "VSA PrimTS requires positive batch, sequence, head, and cube extents" @@ -142,9 +146,9 @@ def forward( use_primts = any( isinstance(fmha, PrimsTSBlockSparseFmha) for fmha in self._fmha_manager.fmha_libs ) - unsupported_reason = _get_unsupported_primts_reason(q, k, v, metadata) - if self.quant_attention_config is not None: - unsupported_reason = "VSA PrimTS does not support quant_attention_config" + unsupported_reason = _get_unsupported_primts_reason( + q, k, v, metadata, sage=self.quant_attention_config is not None + ) if not use_primts: logger.warning_once( "TRTLLM VSA cannot use PrimTS block-sparse attention because the " diff --git a/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py b/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py index a261df5e9df6..df2574bf4ae3 100644 --- a/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py @@ -395,10 +395,6 @@ def forward( ) if use_separate_qkv and (k is None or v is None): raise ValueError("This TRTLLM attention call requires separate q, k, and v tensors.") - if block_sparse_inputs is not None and self.quant_attention_config is not None: - raise ValueError( - "Generic block-sparse attention does not support quant_attention_config." - ) kv_seq_len = seq_len_kv if seq_len_kv is not None else seq_len prepared_metadata = self._prepare_metadata(batch_size, seq_len) diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index c1e28236f20b..f61b8bdf6ff0 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -276,15 +276,23 @@ def _validate_quant_sparse_mutex(self) -> "AttentionConfig": if self.quant_attention_config is None or self.sparse_attention_config is None: return self - # VSA and SOL replace the dense attention path on every backend that - # serves them and never consume quant_attention_config, so accepting a - # quantization recipe would silently ignore user configuration. + # VSA and SOL replace the dense attention path. On TRTLLM every call + # carries the SageAttention recipe to the attention kernels, which either + # apply it or reject the call, so the recipe is never silently dropped. + # The other backends never consume quant_attention_config on that path, + # so accepting a recipe there would silently ignore user configuration. # SkipSoftmax is part of the dense path itself and can compose. + if self.backend == "TRTLLM": + return self algorithm = self.sparse_attention_config.algorithm if algorithm == "vsa": - raise ValueError("VSA and quant_attention_config are mutually exclusive.") + raise ValueError( + "VSA and quant_attention_config are mutually exclusive except on backend='TRTLLM'." + ) if algorithm == "sol_attn": - raise ValueError("SOL and quant_attention_config are mutually exclusive.") + raise ValueError( + "SOL and quant_attention_config are mutually exclusive except on backend='TRTLLM'." + ) return self diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index e0360e1def5b..08b371fc7a65 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -339,6 +339,7 @@ l0_b200: - unittest/_torch/visual_gen/sparse_attention/sol/test_sol_parity.py - unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor.py - unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor_kernels.py + - unittest/_torch/visual_gen/sparse_attention/test_sage_block_sparse.py - unittest/_torch/attention/kernels/test_sage_quant.py - unittest/_torch/visual_gen/test_attention_flashinfer.py - unittest/_torch/visual_gen/test_attention_trtllm_sage.py diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_attention.py b/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_attention.py index d12e21369212..a8ec8503b307 100644 --- a/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_attention.py +++ b/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_attention.py @@ -931,27 +931,24 @@ def test_sol_public_config_requires_supported_backend() -> None: @_CPU_ONLY -@pytest.mark.parametrize( - ("backend", "quant_recipe"), - [ - ("TRTLLM", {"qk_dtype": "fp8", "q_block_size": 1, "k_block_size": 1, "v_block_size": 1}), - ("CUTEDSL", {"qk_dtype": "bf16", "v_dtype": "fp8"}), - ], - ids=["TRTLLM", "CUTEDSL"], -) -def test_sol_and_attention_quantization_are_mutually_exclusive( - backend: str, quant_recipe: dict -) -> None: - """A recipe the backend accepts on its own is still rejected next to SOL.""" +def test_sol_and_attention_quantization_compose_only_on_trtllm() -> None: + """TRTLLM runs SOL with its Sage recipe; CUTEDSL rejects a recipe it accepts on its own.""" with patch("tensorrt_llm.visual_gen.args.get_sm_version", return_value=100): - AttentionConfig( - backend=backend, quant_attention_config=QuantAttentionConfig(**quant_recipe) + sage = QuantAttentionConfig(qk_dtype="fp8", q_block_size=1, k_block_size=1, v_block_size=1) + attention = AttentionConfig( + backend="TRTLLM", + quant_attention_config=sage, + sparse_attention_config=SolAttentionConfig(), ) + assert attention.quant_attention_config is sage + + cutedsl_recipe = QuantAttentionConfig(qk_dtype="bf16", v_dtype="fp8") + AttentionConfig(backend="CUTEDSL", quant_attention_config=cutedsl_recipe) with pytest.raises(ValidationError, match="SOL and quant_attention_config"): AttentionConfig( - backend=backend, - quant_attention_config=QuantAttentionConfig(**quant_recipe), + backend="CUTEDSL", + quant_attention_config=cutedsl_recipe, sparse_attention_config=SolAttentionConfig(), ) diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/test_sage_block_sparse.py b/tests/unittest/_torch/visual_gen/sparse_attention/test_sage_block_sparse.py new file mode 100644 index 000000000000..287c65257e51 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/sparse_attention/test_sage_block_sparse.py @@ -0,0 +1,407 @@ +# 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. +"""Sage quantization of VSA and SOL block-sparse attention on the TRTLLM backend. + +Backends come from the VisualGen factory with a validated ``AttentionConfig``. +Lossy checks compare against the same sparse backend with Sage disabled on +identical BF16 inputs. CUDA Graph replay and masking checks compare against +eager or unquantized execution with near-exact tolerances. +""" + +from typing import Literal + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch.attention.backends.fmha.prims_ts_block_sparse import PrimsTSBlockSparseFmha +from tensorrt_llm._torch.attention.backends.prims_ts import BlockSparseTSWrapper +from tensorrt_llm._torch.attention.backends.sparse.params import ( + BlockSparseForwardInputs, + SparseBackendForwardArgs, +) +from tensorrt_llm._torch.visual_gen.attention.backends.sparse.sol import predictor as sol_predictor +from tensorrt_llm._torch.visual_gen.attention.backends.sparse.vsa.metadata import ( + VSAMetadataBuilder, + set_vsa_forward_context, +) +from tensorrt_llm._torch.visual_gen.attention.backends.trtllm import TrtllmAttention +from tensorrt_llm._torch.visual_gen.attention.backends.utils import create_attention +from tensorrt_llm._torch.visual_gen.config import create_attention_metadata_state +from tensorrt_llm.visual_gen.args import ( + AttentionConfig, + QuantAttentionConfig, + SolAttentionConfig, + VideoSparseAttentionConfig, +) + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() not in ((10, 0), (10, 3)), + reason="Sage block-sparse attention requires SM100 or SM103", +) + +_QuantDtype = Literal["int8", "fp8"] +_SolPhase = Literal["proxy", "exact", "dense"] +_HEADS = 2 +_HEAD_DIM = 128 +_BATCH_SIZE = 2 +# B200, BF16 inputs: INT8 worst relL2/cosine/maxAbs = .03618/.999403/.140625; +# FP8 = .04979/.998761/.167969. The bounds leave about 25-35% headroom. +_LOSSY_BOUNDS = { + "int8": { + "relative_l2": 0.045, + "cosine": 0.9991, + "max_abs": 0.18, + "p99_abs": 0.035, + "channel_mean_shift": 0.0075, + }, + "fp8": { + "relative_l2": 0.065, + "cosine": 0.9982, + "max_abs": 0.22, + "p99_abs": 0.05, + "channel_mean_shift": 0.0095, + }, +} + + +def _backend( + qk_dtype: _QuantDtype | None, + sparse_config: SolAttentionConfig | VideoSparseAttentionConfig | None = None, + *, + k_block_size: int | None = None, +) -> TrtllmAttention: + quant_config = None + if qk_dtype is not None: + if qk_dtype == "int8" and torch.cuda.get_device_capability() == (10, 3): + pytest.skip("SM103 does not have INT8 Tensor Cores") + if k_block_size is None: + k_block_size = 16 if qk_dtype == "int8" else 4 + quant_config = QuantAttentionConfig( + qk_dtype=qk_dtype, + v_dtype="fp8", + q_block_size=1, + k_block_size=k_block_size, + v_block_size=1, + ) + return create_attention( + backend="TRTLLM", + layer_idx=1, + num_heads=_HEADS, + head_dim=_HEAD_DIM, + dtype=torch.bfloat16, + attention_config=AttentionConfig( + backend="TRTLLM", + quant_attention_config=quant_config, + sparse_attention_config=sparse_config, + ), + attention_metadata_state=create_attention_metadata_state(), + ) + + +def _sol( + qk_dtype: _QuantDtype | None, *, phase: _SolPhase = "proxy", k_block_size: int | None = None +) -> TrtllmAttention: + sparse_config = SolAttentionConfig( + tau=-1.0e6 if phase == "exact" else 1.0e6, disabled_until_timestep=0.6 + ) + return _backend(qk_dtype, sparse_config, k_block_size=k_block_size) + + +def _inputs(seq_len: int, *, seed: int = 20260921) -> tuple[torch.Tensor, ...]: + """Strided projection views with batch-dependent scales and a nonzero V mean.""" + generator = torch.Generator(device="cuda").manual_seed(seed) + packed = torch.randn( + _BATCH_SIZE, + seq_len, + 3, + _HEADS, + _HEAD_DIM, + generator=generator, + device="cuda", + dtype=torch.bfloat16, + ) + batch_scales = torch.tensor([0.8, 1.2], device="cuda", dtype=packed.dtype) + channel_bias = torch.linspace(-0.2, 0.2, _HEAD_DIM, device="cuda", dtype=packed.dtype) + packed = packed * batch_scales.view(_BATCH_SIZE, 1, 1, 1, 1) + channel_bias + return packed.unbind(dim=2) + + +def _forward_sol( + backend: TrtllmAttention, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + phase: _SolPhase = "proxy", +) -> torch.Tensor: + return backend.forward( + q=q, + k=k, + v=v, + batch_size=q.shape[0], + seq_len=q.shape[1], + seq_len_kv=q.shape[1], + timestep=0.8 if phase == "dense" else 0.2, + ).view_as(q) + + +def _sparse_wrappers(backend: TrtllmAttention) -> list[BlockSparseTSWrapper]: + """Block-sparse wrappers planned in the backend's shared metadata plan caches.""" + return [ + wrapper + for entry in backend.metadata._metadata_cache.values() + for wrapper in entry["metadata"] + .fmha_plan_caches.get(PrimsTSBlockSparseFmha.PLAN_CACHE_KEY, {}) + .values() + ] + + +def _assert_sparse_sage_active(backend: TrtllmAttention, *, proxy: bool, token_mask: bool) -> None: + wrappers = _sparse_wrappers(backend) + assert wrappers, "Sage silently bypassed the PrimTS block-sparse FMHA" + quant_config = backend.quant_attention_config + qk_dtype = torch.int8 if quant_config.qk_dtype == "int8" else torch.float8_e4m3fn + for wrapper in wrappers: + state = wrapper._plan_state + assert state is not None and state.sage is not None + assert state.use_block_sparse + assert state.use_proxy_routes == proxy + assert state.use_kv_valid_bits == token_mask + assert state.q_dtype == qk_dtype + assert state.v_dtype == torch.float8_e4m3fn + assert state.output_dtype == torch.bfloat16 + assert state.sage.q_block_size == quant_config.q_block_size + assert state.sage.k_block_size == quant_config.k_block_size + + +def _parity_stats(actual: torch.Tensor, expected: torch.Tensor) -> dict[str, float]: + actual, expected = actual.float(), expected.float() + assert actual.shape == expected.shape + assert torch.isfinite(actual).all() and torch.isfinite(expected).all() + delta = actual - expected + return { + "relative_l2": ( + torch.linalg.vector_norm(delta) / torch.linalg.vector_norm(expected) + ).item(), + "cosine": F.cosine_similarity(actual.flatten(), expected.flatten(), dim=0).item(), + "max_abs": delta.abs().max().item(), + "p99_abs": torch.quantile(delta.abs().flatten(), 0.99).item(), + "channel_mean_shift": delta.mean(dim=(0, 1)).abs().max().item(), + } + + +def _assert_parity( + actual: torch.Tensor, + expected: torch.Tensor, + *, + label: str, + qk_dtype: _QuantDtype | None = None, +) -> None: + """Check lossy Sage bounds for ``qk_dtype``, or near-exact agreement without one.""" + stats = _parity_stats(actual, expected) + report = f"{label}, shape={tuple(actual.shape)}, dtype={actual.dtype}: {stats}" + print(report) + if qk_dtype is None: + torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3, msg=report) + assert stats["relative_l2"] <= 1e-2, report + assert stats["cosine"] >= 0.9999, report + return + bounds = _LOSSY_BOUNDS[qk_dtype] + assert stats["cosine"] > bounds["cosine"], report + for name in ("relative_l2", "max_abs", "p99_abs", "channel_mean_shift"): + assert stats[name] < bounds[name], report + assert stats["max_abs"] > 0.0, "Sage quantization had no observable effect" + + +@pytest.mark.parametrize("qk_dtype", ["int8", "fp8"]) +@torch.no_grad() +def test_vsa_sage_matches_unquantized_sparse_attention(qk_dtype: _QuantDtype) -> None: + """Quantize the sparse fine stage while keeping VSA padding and both gates.""" + metadata = VSAMetadataBuilder().build( + current_timestep=0, + raw_latent_shape=(3, 7, 9), + patch_size=(1, 1, 1), + vsa_sparsity=0.5, + device=torch.device("cuda"), + ) + q, k, v = _inputs(3 * 7 * 9) + assert metadata.padded_seq_length > q.shape[1] + assert not q.is_contiguous() + gate_compress = 0.25 + 0.1 * q.sigmoid() + gate_fine = 0.7 + 0.2 * k.sigmoid() + vsa_config = VideoSparseAttentionConfig(vsa_sparsity=0.5) + reference, candidate = _backend(None, vsa_config), _backend(qk_dtype, vsa_config) + with set_vsa_forward_context(metadata): + kwargs = dict( + q=q, + k=k, + v=v, + batch_size=q.shape[0], + seq_len=q.shape[1], + gate_compress=gate_compress, + gate_fine=gate_fine, + ) + expected = reference.forward(**kwargs).view_as(q) + actual = candidate.forward(**kwargs).view_as(q) + + _assert_sparse_sage_active(candidate, proxy=False, token_mask=True) + assert _sparse_wrappers(reference) + _assert_parity( + actual, + expected, + label=f"VSA Sage {qk_dtype}, padded coarse+fine", + qk_dtype=qk_dtype, + ) + + +@pytest.mark.parametrize( + ("qk_dtype", "phase", "k_block_size"), + [ + ("int8", "proxy", 16), + ("fp8", "proxy", 4), + ("fp8", "proxy", 1), + ("int8", "exact", 16), + ("fp8", "exact", 4), + ("int8", "dense", 16), + ("fp8", "dense", 4), + ], +) +@torch.no_grad() +def test_sol_sage_matches_unquantized_attention( + qk_dtype: _QuantDtype, phase: _SolPhase, k_block_size: int +) -> None: + """Identical SOL routes and means, including tail blocks and the dense timestep prefix.""" + q, k, v = _inputs(131) + reference = _sol(None, phase=phase) + candidate = _sol(qk_dtype, phase=phase, k_block_size=k_block_size) + expected = _forward_sol(reference, q, k, v, phase=phase) + actual = _forward_sol(candidate, q, k, v, phase=phase) + + if phase == "dense": + assert not _sparse_wrappers(candidate), "the dense timestep unexpectedly ran SOL" + assert not _sparse_wrappers(reference) + else: + _assert_sparse_sage_active(candidate, proxy=True, token_mask=False) + outputs = sol_predictor.predict( + q.contiguous(), + k.contiguous(), + v.contiguous(), + tau=candidate.sol_params.tau, + sm_scale=_HEAD_DIM**-0.5, + ) + exact_bits = outputs.exact_block_bits.to(torch.int64) + num_blocks = outputs.k_summary.shape[1] + selected = torch.stack( + [(exact_bits[..., idx // 32] >> (idx % 32)) & 1 for idx in range(num_blocks)] + ) + if phase == "exact": + assert selected.all(), "the all-exact SOL policy retained proxy routes" + else: + assert selected.any() and not selected.all(), "SOL did not mix exact and proxy routes" + _assert_parity( + actual, + expected, + label=f"SOL Sage {qk_dtype}, Kblock={k_block_size}, {phase}", + qk_dtype=qk_dtype, + ) + + +@pytest.mark.parametrize("qk_dtype", ["int8", "fp8"]) +@torch.no_grad() +def test_sol_sage_cuda_graph_refreshes_inputs_and_scales(qk_dtype: _QuantDtype) -> None: + """A captured mixed-proxy Sage layer consumes new Q/K/V and scales on replay.""" + backend = _sol(qk_dtype) + q, k, v = _inputs(131) + previous = _forward_sol(backend, q, k, v).clone() + _assert_sparse_sage_active(backend, proxy=True, token_mask=False) + warmup_stream = torch.cuda.Stream() + warmup_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(warmup_stream): + _forward_sol(backend, q, k, v) + torch.cuda.current_stream().wait_stream(warmup_stream) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = _forward_sol(backend, q, k, v) + for destination, source in zip((q, k, v), _inputs(131, seed=20260922)): + destination.copy_(source) + graph.replay() + replayed = captured.clone() + expected = _forward_sol(backend, q, k, v) + assert _parity_stats(expected, previous)["relative_l2"] > 0.1 + _assert_parity(replayed, expected, label=f"SOL {qk_dtype} Sage CUDA Graph replay") + + +def _masked_zero_score_inputs() -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, BlockSparseForwardInputs +]: + shape = (_BATCH_SIZE, 131, _HEADS, _HEAD_DIM) + magnitudes = torch.tensor([0.0, 1.0e-8], device="cuda", dtype=torch.bfloat16) + q = torch.ones(shape, device="cuda", dtype=torch.bfloat16) * magnitudes[:, None, None, None] + tokens = torch.arange(shape[1], device=q.device) + visible = (tokens >= 68) & (tokens < 76) + # Masked V=4 exposes leaked probability; the eight visible tokens have V=1. + v = torch.where(visible, 1.0, 4.0)[None, :, None, None].expand(shape).to(q.dtype).contiguous() + # Query blocks select a fully masked block, a partly visible block, then no block. + bits = torch.tensor([1, 2, 0], device=q.device, dtype=torch.uint32) + valid_words = torch.tensor([0, 0, 0xFF0, 0, 0], device=q.device, dtype=torch.uint32) + routes = BlockSparseForwardInputs( + q_block_size=64, + kv_block_size=64, + exact_block_bits=bits.view(1, 1, 3, 1).expand(_BATCH_SIZE, _HEADS, 3, 1).contiguous(), + kv_valid_bits=valid_words[None, :].expand(_BATCH_SIZE, -1).contiguous(), + ) + return q, -q, v, routes + + +@pytest.mark.parametrize("qk_dtype", ["int8", "fp8"]) +@torch.no_grad() +def test_sage_zero_scores_preserve_token_masks_and_empty_rows(qk_dtype: _QuantDtype) -> None: + """Zero and tiny Q/K keep token masking and empty rows exact through quantization.""" + q, k, v, routes = _masked_zero_score_inputs() + backends = (_backend(None), _backend(qk_dtype)) + expected, actual = ( + backend.forward( + q, + k, + v, + batch_size=q.shape[0], + seq_len=q.shape[1], + sparse_backend_args=SparseBackendForwardArgs(block_sparse_inputs=routes), + ).view_as(q) + for backend in backends + ) + _assert_sparse_sage_active(backends[1], proxy=False, token_mask=True) + _assert_parity(actual, expected, label=f"{qk_dtype} zero/tiny scores with masked rows") + assert torch.count_nonzero(actual[:, :64]).item() == 0, "fully masked rows must be zero" + assert torch.count_nonzero(actual[:, 128:]).item() == 0, "empty route rows must be zero" + torch.testing.assert_close(actual[:, 64:128], torch.ones_like(actual[:, 64:128])) + + +@pytest.mark.parametrize("qk_dtype", ["int8", "fp8"]) +@torch.no_grad() +def test_sol_sage_zero_values_preserve_zero_output(qk_dtype: _QuantDtype) -> None: + """Exact and proxy SOL contributions stay zero through the V-scale floor. + + The inputs and tau match the mixed-route proxy cases above, so both route + kinds contribute. + """ + q, k, v = _inputs(131) + zero_v = torch.zeros_like(v) + candidate = _sol(qk_dtype) + actual = _forward_sol(candidate, q, k, zero_v) + _assert_sparse_sage_active(candidate, proxy=True, token_mask=False) + assert torch.count_nonzero(actual).item() == 0, "zero exact values and proxy means stay zero" diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_attention.py b/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_attention.py index 24aa3b32f774..0484352b709d 100644 --- a/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_attention.py +++ b/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_attention.py @@ -295,7 +295,7 @@ def test_trtllm_vsa_hands_predicted_routes_to_core_via_sparse_backend_args( attention._fmha_manager = SimpleNamespace( fmha_libs=[object.__new__(vsa_backend.PrimsTSBlockSparseFmha)] ) - monkeypatch.setattr(vsa_backend, "_get_unsupported_primts_reason", lambda *args: None) + monkeypatch.setattr(vsa_backend, "_get_unsupported_primts_reason", lambda *args, sage: None) q = torch.randn(1, 80, 1, 8) with set_vsa_forward_context(_make_vsa_metadata()): diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py index cee4b89daebf..453f5b15ec89 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py @@ -104,32 +104,24 @@ def test_quant_config_rejected_when_unsupported(self): ), ) - @pytest.mark.parametrize( - ("backend", "quant_config"), - [ - ( - "TRTLLM", - QuantAttentionConfig( - qk_dtype="fp8", - q_block_size=1, - k_block_size=1, - v_block_size=1, - ), - ), - ( - "CUTEDSL", - QuantAttentionConfig(qk_dtype="bf16", v_dtype="fp8"), - ), - ], - ) - def test_vsa_and_quantization_are_mutually_exclusive(self, backend, quant_config): + def test_vsa_and_quantization_compose_only_on_trtllm(self): with patch("tensorrt_llm.visual_gen.args.get_sm_version", return_value=100): + sage = QuantAttentionConfig( + qk_dtype="fp8", q_block_size=1, k_block_size=1, v_block_size=1 + ) + attention = AttentionConfig( + backend="TRTLLM", + quant_attention_config=sage, + sparse_attention_config=VideoSparseAttentionConfig(vsa_sparsity=0.9), + ) + assert attention.quant_attention_config is sage + with pytest.raises( ValidationError, match="VSA and quant_attention_config are mutually exclusive" ): AttentionConfig( - backend=backend, - quant_attention_config=quant_config, + backend="CUTEDSL", + quant_attention_config=QuantAttentionConfig(qk_dtype="bf16", v_dtype="fp8"), sparse_attention_config=VideoSparseAttentionConfig(vsa_sparsity=0.9), ) From b224d07046e653158c43d55184f6685581e0379f Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Sun, 11 Oct 2026 03:40:13 +0000 Subject: [PATCH 05/13] [None][feat] Allow the FP8 SageAttention recipe (1, 16, 1) on SM100 and SM103 The trtllm-gen FMHA cubins already include dense FP8 Sage context kernels with 16-token K scale blocks for SM100 and SM103, the attention op accepts Q/K blocks of 1, 4 and 16, and the PrimTS block-sparse path supports them. Only the VisualGen recipe table left FP8 (1, 16, 1) out. Add it for SM100 and SM103, list it in the quantized-attention guide, and cover it in the recipe validation and the dense numerical test. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- docs/source/features/visualgen-quantized-attention.md | 2 +- tensorrt_llm/visual_gen/args.py | 2 ++ tests/unittest/_torch/visual_gen/test_attention_trtllm_sage.py | 1 + tests/unittest/_torch/visual_gen/test_visual_gen_args.py | 2 ++ 4 files changed, 6 insertions(+), 1 deletion(-) diff --git a/docs/source/features/visualgen-quantized-attention.md b/docs/source/features/visualgen-quantized-attention.md index 7f81a62d2143..c6d3e7aa8e23 100644 --- a/docs/source/features/visualgen-quantized-attention.md +++ b/docs/source/features/visualgen-quantized-attention.md @@ -28,7 +28,7 @@ A recipe is the tuple `(qk_dtype, v_dtype, (q_block_size, k_block_size, v_block_ |---|---|---|---|---| | `TRTLLM` | `int8` | `fp8` | `(2, 16, 1)` | SageAttention (INT8 QK) for Hopper | | `TRTLLM` | `int8` | `fp8` | `(1, 1, 1)`, `(1, 4, 1)`, `(1, 16, 1)` | SageAttention (INT8 QK) for Blackwell | -| `TRTLLM` | `fp8` | `fp8` | `(1, 1, 1)`, `(1, 4, 1)` | SageAttention (FP8 QK) for Blackwell | +| `TRTLLM` | `fp8` | `fp8` | `(1, 1, 1)`, `(1, 4, 1)`, `(1, 16, 1)` | SageAttention (FP8 QK) for Blackwell | | `TRTLLM` | `bf16` | `fp8` | `(0, 0, 0)` | PrimTS QK16PV8 for Blackwell (`TLLM_FMHA_LIBS=+prims_ts`) | | `TRTLLM` | `fp8` | `fp8` | `(0, 0, 0)` | PrimTS FP8 for Blackwell (`TLLM_FMHA_LIBS=+prims_ts`) | | `CUDNN` | `fp8` | `fp8` | `(0, 0, 0)` | cuDNN FP8 | diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index f61b8bdf6ff0..380154d92108 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -152,10 +152,12 @@ def _validate_quant_attention_config(self) -> "AttentionConfig": ("int8", "fp8", (1, 16, 1)), ("fp8", "fp8", (1, 1, 1)), ("fp8", "fp8", (1, 4, 1)), + ("fp8", "fp8", (1, 16, 1)), }, 103: { ("fp8", "fp8", (1, 1, 1)), ("fp8", "fp8", (1, 4, 1)), + ("fp8", "fp8", (1, 16, 1)), }, } # PrimTS serves the per-tensor FP8 recipes on Blackwell (TLLM_FMHA_LIBS=+prims_ts). diff --git a/tests/unittest/_torch/visual_gen/test_attention_trtllm_sage.py b/tests/unittest/_torch/visual_gen/test_attention_trtllm_sage.py index 10853ed9ffe0..09540a930c07 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_trtllm_sage.py +++ b/tests/unittest/_torch/visual_gen/test_attention_trtllm_sage.py @@ -176,6 +176,7 @@ def _test_attention_trtllm_sage( pytest.param(True, 1, 4, 1e-1, 4e-2, (100,), id="sm100-int8-q1k4"), pytest.param(True, 1, 16, 5e-1, 5e-1, (100,), id="sm100-int8-q1k16"), pytest.param(False, 1, 1, 3e-1, 2e-1, (100, 103), id="sm10x-fp8-q1k1"), + pytest.param(False, 1, 16, 3e-1, 2e-1, (100, 103), id="sm10x-fp8-q1k16"), ], ) def test_attention_trtllm_sage( diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py index 453f5b15ec89..3aead1c2c10e 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py @@ -152,8 +152,10 @@ def test_skip_softmax_and_sage_quantization_can_be_combined(self): (100, "int8", 1, 16, 1), (100, "fp8", 1, 1, 1), (100, "fp8", 1, 4, 1), + (100, "fp8", 1, 16, 1), (103, "fp8", 1, 1, 1), (103, "fp8", 1, 4, 1), + (103, "fp8", 1, 16, 1), (100, "bf16", 0, 0, 0), (100, "fp8", 0, 0, 0), (103, "bf16", 0, 0, 0), From 7f0b3de8e0a3f93a346ea8089c1214a58ac11a41 Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Wed, 14 Oct 2026 10:06:17 +0000 Subject: [PATCH 06/13] [None][feat] Quantize PrimTS Sage attention with the TensorRT-LLM Sage kernel The PrimTS block-sparse path quantized Q/K/V with FlashInfer's copy of the TensorRT-LLM Sage quantization kernel, which only offers 1, 4, or 16 tokens per Q/K scale and reduces the V scales and the K-smoothing mean over the whole batch. PrimTS now takes V scales per batch entry and accepts more K block sizes, so the path moves to the TensorRT-LLM kernel and extends it: - Per-sequence statistics: given KV cumulative lengths, the kernel reduces the V scales and the K mean over each sequence, [batch, heads_kv, head_dim], as SageAttention does. Without them it keeps the whole-batch reduction the trtllm-gen kernels consume. - Block sizes: contiguous Q/K scale blocks of 1, 4, 16, and now 32 tokens. Blocks larger than 8 tokens are read in 8-token chunks twice (maximum, then quantize), with the chunk loop kept rolled, so registers and code size stay fixed; smaller blocks keep the single pass. Other powers of two up to 256 only need a dispatch case. With 16 tokens per K scale, quantizing a real Wan VSA layer on GB300 takes 2.02 ms instead of 2.75 ms with FlashInfer's copy. - Rounding: Q/K are now scaled and converted in FP32 for every block size. The single pass converted the BF16/FP16 product to INT8, which truncates, so INT8 Q/K carried up to 1.2 grid steps of error (mean 0.44) instead of half a step (mean 0.25). This also applies to the dense INT8 Sage path. - Quantization outside the FMHA libraries: after sparse prediction and before FMHA selection, the attention core quantizes the Q/K/V of block-sparse calls and the K/V means of their proxy routes, and passes the result in AttentionForwardArgs.sage_quantized. It does so when a library of the layer declares supports_sage_quantized_inputs; PrimsTSBlockSparseFmha does, and serves Sage only with these inputs. Dense SageAttention still quantizes inside the attention op, which ignores them. AttentionForwardArgs.sage_attn_qk_dtype gives the core and the library the recipe's Q/K dtype. - Op: trtllm::sage_quantize quantizes Q/K/V, or, with reuse_stats, only K/V with the V scales and K mean of an earlier call; the Python sage_quantize takes that call's result as stats. The proxy summaries are quantized this way, one key per mean with a scale of its own, which replaces their Triton quantizer. The op has a fake implementation for torch.compile. The adapter no longer checks the quantizer's block sizes and input dtypes; the op validates its own inputs. VSA no longer limits Sage to BF16 inputs either: the scales are computed in FP32, so zero or tiny FP16 Q/K blocks no longer overflow them. - Tests compare against a PyTorch reference over all block sizes, INT8 and FP8, K smoothing, GQA, ragged sequences, and proxy summaries, check which calls the attention core quantizes, and cover zero Q/K with FP16 inputs. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- cpp/tensorrt_llm/common/attentionOp.cpp | 3 +- cpp/tensorrt_llm/common/sageQuant.cu | 224 +++++++++----- cpp/tensorrt_llm/common/sageQuant.h | 10 +- cpp/tensorrt_llm/thop/CMakeLists.txt | 10 +- cpp/tensorrt_llm/thop/sageQuantizeOp.cpp | 207 +++++++++++++ .../features/visualgen-sparse-attention.md | 12 +- .../attention/ATTENTION_DEVELOPER_GUIDE.md | 17 +- .../attention/backends/fmha/fallback.py | 1 + .../attention/backends/fmha/interface.py | 3 + .../attention/backends/fmha/prims_ts.py | 6 +- .../backends/fmha/prims_ts_block_sparse.py | 103 ++----- .../_torch/attention/backends/interface.py | 13 + .../_torch/attention/backends/trtllm.py | 73 +++++ .../_torch/attention/kernels/sage_quant.py | 262 +++++++++------- .../_torch/custom_ops/cpp_custom_ops.py | 18 ++ .../attention/backends/sparse/vsa/backend.py | 8 +- .../attention/kernels/test_sage_quant.py | 285 +++++++++++++++--- .../sparse/test_prims_ts_block_sparse.py | 80 +++++ .../test_sage_block_sparse.py | 34 ++- .../vsa/test_vsa_attention.py | 2 +- 20 files changed, 1021 insertions(+), 350 deletions(-) create mode 100644 cpp/tensorrt_llm/thop/sageQuantizeOp.cpp diff --git a/cpp/tensorrt_llm/common/attentionOp.cpp b/cpp/tensorrt_llm/common/attentionOp.cpp index 9a452d4bda52..8ac041a4c2fa 100644 --- a/cpp/tensorrt_llm/common/attentionOp.cpp +++ b/cpp/tensorrt_llm/common/attentionOp.cpp @@ -1959,7 +1959,8 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea // otherwise surface as an empty scale buffer deep inside invokeSageQuant(). TLLM_CHECK_WITH_INFO(q_max_n_blk > 0 && k_max_n_blk > 0, "No SageQuant kernel on sm_%d for block sizes (q, k, v) = (%d, %d, %d) with qk_int8=%s. SM90 " - "requires (2, 16, 1) together with INT8 Q/K; SM100 requires q/k block sizes of 1, 4 or 16.", + "requires (2, 16, 1) together with INT8 Q/K; other architectures require q/k block sizes of 1, 4, 16, " + "or 32.", mSM, mSageAttnNumEltsPerBlkQ, mSageAttnNumEltsPerBlkK, mSageAttnNumEltsPerBlkV, mSageAttnQkInt8 ? "true" : "false"); TLLM_CHECK_WITH_INFO(!params.kv_scale_quant_orig, diff --git a/cpp/tensorrt_llm/common/sageQuant.cu b/cpp/tensorrt_llm/common/sageQuant.cu index 305f39d86704..48ba3812b5cc 100644 --- a/cpp/tensorrt_llm/common/sageQuant.cu +++ b/cpp/tensorrt_llm/common/sageQuant.cu @@ -37,6 +37,33 @@ using tensorrt_llm::kernels::DATA_TYPE_E4M3; using tensorrt_llm::kernels::DATA_TYPE_FP16; using tensorrt_llm::kernels::DATA_TYPE_INT8; +// Most tokens of one Q/K scale group a thread holds in registers at once. +constexpr int MaxChunkTokens = 8; + +// The V tokens one CTA row reduces for the V scales and the K mean, and the scale row they go to. +// blockIdx.y is (headIdx * batchSize + seqIdx) when the statistics are per sequence; otherwise +// rows below numHeadsV each cover one head over the whole batch. +struct SageStatRange +{ + int headIdx; + int tokBegin; + int tokEnd; + int row; +}; + +__device__ inline SageStatRange getSageStatRange( + int const* ptrCuSeqLensV, int batchSize, int numHeadsV, int sumSeqLensV) +{ + if (ptrCuSeqLensV == nullptr) + { + int const headIdx = blockIdx.y; + return {headIdx, 0, sumSeqLensV, headIdx}; + } + int const headIdx = blockIdx.y / batchSize; + int const seqIdx = blockIdx.y % batchSize; + return {headIdx, ptrCuSeqLensV[seqIdx], ptrCuSeqLensV[seqIdx + 1], seqIdx * numHeadsV + headIdx}; +} + // SageAttention quantization kernel. Tensors are interpreted as column-major [D, H, S], which is // the same physical layout as contiguous PyTorch [S, H, D]. Each launch quantizes Q or K. It can // simultaneously collect V scales (VStage=1) or quantize V using scales collected by an earlier @@ -52,10 +79,13 @@ using tensorrt_llm::kernels::DATA_TYPE_INT8; // the head stride is scaleHeadStride(sumSeqLensQk, batchSize). Every sequence reserves a whole // spare tile, so a trailing partial tile -- whose scales the consumer still indexes in full -- // cannot run into the next sequence. +// +// V scales and the K mean are reduced per sequence when ptrCuSeqLensV is given (row +// seqIdx * numHeadsV + headIdx) and over the whole batch otherwise (row headIdx). template __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* ptrCuSeqLensQk, void const* ptrQk, void* ptrQkQuant, float* ptrQkScale, void const* ptrKForMean, float* ptrKMean, int sumSeqLensV, int numHeadsV, - void const* ptrV, void* ptrVQuant, float* ptrVScale) + void const* ptrV, void* ptrVQuant, float* ptrVScale, int const* ptrCuSeqLensV) { using namespace cute; using namespace cutlass; @@ -117,6 +147,25 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p float* ptrQkSeqScale = ptrQkScaleHead + Traits::scaleBase(seqBegin, seqIdx); int const numTiles = ceil_div(seqLen, TileTokens); + // This thread's channel slice of the K mean of this sequence (or of the whole batch). + Tensor rKMean = make_tensor(Shape{}); + if constexpr (SmoothK && VStage == 2) + { + int const meanRow = ptrCuSeqLensV != nullptr ? seqIdx * numHeads + headIdx : headIdx; + Tensor gKMean = make_tensor( + make_gmem_ptr(ptrKMean + static_cast(meanRow) * HeadDim + threadInScaleIdx * BestVL), + Shape{}); + cute::copy(AutoVectorizingCopy{}, gKMean, rKMean); + } + + // A thread keeps at most ChunkTokens tokens of a group in registers, in the input type; FP32 + // values are recomputed from them when needed. Larger groups are read in chunks twice: once + // to find the group maximum, once to quantize. + constexpr int ChunkTokens = TokensPerScale < MaxChunkTokens ? TokensPerScale : MaxChunkTokens; + constexpr int NumChunks = TokensPerScale / ChunkTokens; + static_assert(TokensPerScale % ChunkTokens == 0, "a group must split into whole chunks"); + static_assert(BestVL % 4 == 0, "quantized values are converted four at a time"); + // tileIdx/grpIdx identify a scale; isFullTile is a compile-time fast path for tiles that // lie entirely inside the sequence. In a partial tile validity has to be tested per token // rather than as a prefix count, because a group's tokens are strided in general -- for @@ -127,58 +176,100 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p constexpr bool IsFullTile = decltype(isFullTile)::value; int const tileBase = tileIdx * TileTokens; auto tokenOf = [&](int i) { return tileBase + Traits::tokenInTile(i, grpIdx); }; + auto isValid = [&](int chunk, int i) { return IsFullTile || tokenOf(chunk * ChunkTokens + i) < seqLen; }; - Tensor rQk = make_tensor(Shape>{}); - Tensor rQkQuant = make_tensor(Shape>{}); - Tensor rQkCompute = make_tensor(Shape>{}); - Tensor rQk_x2 = recast>(rQk); - Tensor rQkCompute_x2 = recast>(rQkCompute); - Tensor rQk_x4 = recast>(rQk); + Tensor rQk = make_tensor(Shape>{}); + Tensor rQkQuant = make_tensor(Shape>{}); Tensor rQkQuant_x4 = recast>(rQkQuant); - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < TokensPerScale; ++i) + auto loadChunk = [&](int chunk) { - int const tokenIdx = tokenOf(i); - if (IsFullTile || tokenIdx < seqLen) + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < ChunkTokens; ++i) { - Tensor gSrc = make_tensor( - make_gmem_ptr(ptrQkThread + static_cast(tokenIdx) * tokenStride), Shape{}); - cute::copy(AutoVectorizingCopy{}, gSrc, rQk(_, i)); + int const tokenIdx = tokenOf(chunk * ChunkTokens + i); + if (IsFullTile || tokenIdx < seqLen) + { + Tensor gSrc = make_tensor( + make_gmem_ptr(ptrQkThread + static_cast(tokenIdx) * tokenStride), Shape{}); + cute::copy(AutoVectorizingCopy{}, gSrc, rQk(_, i)); + } + else + { + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < BestVL; ++j) + { + rQk(j, i) = static_cast(0); + } + } } - else + }; + // FP32 value of channel j of token i of the loaded chunk, K-smoothed. + auto valueAt = [&](int chunk, int i, int j) + { + float x = static_cast(rQk(j, i)); + if constexpr (SmoothK && VStage == 2) { - CUTLASS_PRAGMA_UNROLL - for (int j = 0; j < BestVL; ++j) + if (isValid(chunk, i)) { - rQk(j, i) = static_cast(0); + x -= rKMean(j); } } - } - cute::transform(rQk_x2, rQkCompute_x2, NumericArrayConverter::convert); - - if constexpr (SmoothK && VStage == 2) + return x; + }; + // Scaling and conversion stay in FP32: converting a BF16/FP16 product to INT8 truncates + // instead of rounding. + auto quantizeChunk = [&](int chunk, float invScale) { - float const* ptrKMeanHead = ptrKMean + headIdx * HeadDim; CUTLASS_PRAGMA_UNROLL - for (int tokenIdx = 0; tokenIdx < TokensPerScale; ++tokenIdx) + for (int i = 0; i < ChunkTokens; ++i) { - if (IsFullTile || tokenOf(tokenIdx) < seqLen) + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < BestVL; j += 4) { + Array scaled; CUTLASS_PRAGMA_UNROLL - for (int vecIdx = 0; vecIdx < BestVL; ++vecIdx) + for (int e = 0; e < 4; ++e) { - rQkCompute(vecIdx, tokenIdx) -= ptrKMeanHead[threadInScaleIdx * BestVL + vecIdx]; + scaled[e] = valueAt(chunk, i, j + e) * invScale; } + rQkQuant_x4((i * BestVL + j) / 4) + = NumericArrayConverter::convert(scaled); } } - } + }; + auto storeChunk = [&](int chunk) + { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < ChunkTokens; ++i) + { + int const tokenIdx = tokenOf(chunk * ChunkTokens + i); + if (IsFullTile || tokenIdx < seqLen) + { + Tensor gDst = make_tensor( + make_gmem_ptr(ptrQkQuantThread + static_cast(tokenIdx) * tokenStride), + Shape{}); + cute::copy(AutoVectorizingCopy{}, rQkQuant(_, i), gDst); + } + } + }; float maxScale = 1e-3f; - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < size(rQk); ++i) + // Both chunk loops stay rolled: unrolling them multiplies the code of the large blocks by + // the chunk count. + CUTLASS_PRAGMA_NO_UNROLL + for (int chunk = 0; chunk < NumChunks; ++chunk) { - maxScale = ::fmaxf(maxScale, ::fabsf(rQkCompute(i))); + loadChunk(chunk); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < ChunkTokens; ++i) + { + CUTLASS_PRAGMA_UNROLL + for (int j = 0; j < BestVL; ++j) + { + maxScale = ::fmaxf(maxScale, ::fabsf(valueAt(chunk, i, j))); + } + } } CUTLASS_PRAGMA_UNROLL for (int delta = 1; delta < threadsPerScale; delta <<= 1) @@ -193,33 +284,16 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p // scale. The 1e-3 floor above is what guarantees that -- do not lower it to an epsilon. ptrQkSeqScale[tileIdx * ScalesPerTile + grpIdx] = maxScale; float const invScale = 1.0f / maxScale; - if constexpr (SmoothK && VStage == 2) - { - // rQkCompute holds the result value - Tensor rQkCompute_x4 = recast>(rQkCompute); - cute::transform(rQkCompute, rQkCompute, [&](float const& x) { return x * invScale; }); - cute::transform(rQkCompute_x4, rQkQuant_x4, NumericArrayConverter::convert); - } - else + CUTLASS_PRAGMA_NO_UNROLL + for (int chunk = 0; chunk < NumChunks; ++chunk) { - // rQk holds the result value - Array scaleQuant - = NumericArrayConverter::convert(Array{invScale, invScale}); - cutlass::multiplies> scaleQuantOp; - cute::transform(rQk_x2, rQk_x2, [&](auto& x) { return scaleQuantOp(x, scaleQuant); }); - cute::transform(rQk_x4, rQkQuant_x4, NumericArrayConverter::convert); - } - - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < TokensPerScale; ++i) - { - int const tokenIdx = tokenOf(i); - if (IsFullTile || tokenIdx < seqLen) + if constexpr (NumChunks > 1) { - Tensor gDst = make_tensor( - make_gmem_ptr(ptrQkQuantThread + static_cast(tokenIdx) * tokenStride), Shape{}); - cute::copy(AutoVectorizingCopy{}, rQkQuant(_, i), gDst); + // The single-chunk case still holds its tokens from the maximum pass. + loadChunk(chunk); } + quantizeChunk(chunk, invScale); + storeChunk(chunk); } }; @@ -241,7 +315,8 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p } else if (blockIdx.z == 1) { - int const headIdx = blockIdx.y; + SageStatRange const range = getSageStatRange(ptrCuSeqLensV, batchSize, numHeadsV, sumSeqLensV); + int const headIdx = range.headIdx; using ElementQuantizedV = cutlass::float_e4m3_t; constexpr int threadsPerHead = HeadDim / BestVL; static_assert(HeadDim % BestVL == 0, "VL must divide HeadDim"); @@ -250,7 +325,8 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p reinterpret_cast(ptrV), make_shape(VL{}, Int{}, numHeadsV, sumSeqLensV)); Tensor gVQuant = make_tensor(reinterpret_cast(ptrVQuant), make_shape(VL{}, Int{}, numHeadsV, sumSeqLensV)); - Tensor gVScale = make_tensor(ptrVScale, make_shape(VL{}, Int{}, numHeadsV)); + Tensor gVScale = make_tensor( + ptrVScale + static_cast(range.row) * HeadDim, make_shape(VL{}, Int{})); Tensor rV = make_tensor(Shape{}); Tensor rVMax = make_tensor(Shape{}); @@ -267,11 +343,11 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p if (headIdx < numHeadsV) { int const numToksPerWarp = 32 / threadsPerHead; - int tokIdx = warpId * numToksPerWarp + thrId / threadsPerHead; + int tokIdx = range.tokBegin + warpId * numToksPerWarp + thrId / threadsPerHead; int const threadInTokIdx = thrId % threadsPerHead; Tensor gVSeq = gV(_, threadInTokIdx, headIdx, _); Tensor gVSeqQuant = gVQuant(_, threadInTokIdx, headIdx, _); - Tensor gVSeqScale = gVScale(_, threadInTokIdx, headIdx); + Tensor gVSeqScale = gVScale(_, threadInTokIdx); if constexpr (VStage == 1) { @@ -287,7 +363,7 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p rVScale(i) = 1e-3f; } cute::transform(rVScale_x2, rVMax_x2, cutlass::NumericArrayConverter::convert); - for (; tokIdx < sumSeqLensV; tokIdx += numToksPerWave) + for (; tokIdx < range.tokEnd; tokIdx += numToksPerWave) { cute::copy(AutoVectorizingCopy{}, gVSeq(_, tokIdx), rV); cute::transform(rV_x2, rV_x2, cutlass::absolute_value_op>{}); @@ -315,7 +391,7 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p int const numToksPerWave = numWarps * numToksPerWarp; cute::copy(AutoVectorizingCopy{}, gVSeqScale, rVScale); cute::transform(rVScale_x2, rVScale_x2, cutlass::reciprocal_approximate>{}); - for (; tokIdx < sumSeqLensV; tokIdx += numToksPerWave) + for (; tokIdx < range.tokEnd; tokIdx += numToksPerWave) { cute::copy(AutoVectorizingCopy{}, gVSeq(_, tokIdx), rV); cute::transform(rV_x2, rVCompute_x2, cutlass::NumericArrayConverter::convert); @@ -331,16 +407,18 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p { if constexpr (SmoothK && VStage == 1) { - int const headIdx = blockIdx.y; + SageStatRange const range = getSageStatRange(ptrCuSeqLensV, batchSize, numHeadsV, sumSeqLensV); + int const headIdx = range.headIdx; constexpr int threadsPerHead = HeadDim / BestVL; static_assert(HeadDim % BestVL == 0, "VL must divide HeadDim"); static_assert(threadsPerHead <= 32, "One token block should never exceed warp scope"); - if (headIdx < numHeadsV) + if (headIdx < numHeadsV && range.tokBegin < range.tokEnd) { Tensor gK = make_tensor(reinterpret_cast(ptrKForMean), make_shape(VL{}, Int{}, numHeadsV, sumSeqLensV)); - Tensor gKMean = make_tensor(ptrKMean, make_shape(VL{}, Int{}, numHeadsV)); + Tensor gKMean = make_tensor( + ptrKMean + static_cast(range.row) * HeadDim, make_shape(VL{}, Int{})); Tensor rK = make_tensor(Shape{}); Tensor rKCompute = make_tensor(Shape{}); Tensor rKSum = make_tensor(Shape{}); @@ -354,13 +432,13 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p { return; } - int tokIdx = warpId * numToksPerWarp + thrId / threadsPerHead; + int tokIdx = range.tokBegin + warpId * numToksPerWarp + thrId / threadsPerHead; int const threadInTokIdx = thrId % threadsPerHead; Tensor gKSeq = gK(_, threadInTokIdx, headIdx, _); - Tensor gKMeanHead = gKMean(_, threadInTokIdx, headIdx); + Tensor gKMeanHead = gKMean(_, threadInTokIdx); clear(rKSum); - for (; tokIdx < sumSeqLensV; tokIdx += numToksPerWave) + for (; tokIdx < range.tokEnd; tokIdx += numToksPerWave) { cute::copy(AutoVectorizingCopy{}, gKSeq(_, tokIdx), rK); cute::transform(rK_x2, rKCompute_x2, NumericArrayConverter::convert); @@ -380,7 +458,7 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p } if (threadInTokIdx == thrId) { - float const invNumTokens = 1.0f / sumSeqLensV; + float const invNumTokens = 1.0f / (range.tokEnd - range.tokBegin); CUTLASS_PRAGMA_UNROLL for (int i = 0; i < BestVL; ++i) { @@ -410,6 +488,8 @@ void invokeSageQuantQkvImpl(SageQuantParams const& params) TLLM_CHECK_WITH_INFO( !params.kSmooth || (params.ptrKMean != nullptr && (params.vStage != 1 || params.ptrKForMean != nullptr)), "Invalid SageQuant K-smoothing parameters"); + TLLM_CHECK_WITH_INFO( + params.ptrCuSeqLensV == nullptr || params.vStage != 0, "SageQuant per-sequence V statistics require V staging"); auto invokeKernel = [&](auto headDimStatic, auto partitionStatic) { @@ -419,7 +499,7 @@ void invokeSageQuantQkvImpl(SageQuantParams const& params) void* kernelArgs[] = {&kernelParams.sumSeqLensQk, &kernelParams.batchSize, &kernelParams.ptrCuSeqLensQk, &kernelParams.ptrQk, &kernelParams.ptrQkQuant, &kernelParams.ptrQkScale, &kernelParams.ptrKForMean, &kernelParams.ptrKMean, &kernelParams.sumSeqLensV, &kernelParams.numHeadsV, &kernelParams.ptrV, - &kernelParams.ptrVQuant, &kernelParams.ptrVScale}; + &kernelParams.ptrVQuant, &kernelParams.ptrVScale, &kernelParams.ptrCuSeqLensV}; auto launchKernel = [&](auto smoothKStatic, auto vStageStatic) { @@ -477,6 +557,8 @@ void invokeSageQuantQkvImpl(SageQuantParams const& params) } }; + // Contiguous blocks are instantiated for the sizes in use. The kernel takes any power of two up to + // 256: add its case here and in withPartitionTraits. #define TLLM_SAGE_DISPATCH_HEAD_DIM(HEAD_DIM) \ if (params.headDim == HEAD_DIM) \ { \ @@ -490,6 +572,7 @@ void invokeSageQuantQkvImpl(SageQuantParams const& params) case 1: invokeKernel(Int{}, ContiguousPartition<1>{}); return; \ case 4: invokeKernel(Int{}, ContiguousPartition<4>{}); return; \ case 16: invokeKernel(Int{}, ContiguousPartition<16>{}); return; \ + case 32: invokeKernel(Int{}, ContiguousPartition<32>{}); return; \ default: break; \ } \ break; \ @@ -501,7 +584,7 @@ void invokeSageQuantQkvImpl(SageQuantParams const& params) #undef TLLM_SAGE_DISPATCH_HEAD_DIM TLLM_THROW( "Unsupported SageQuant dispatch config (head_dim must be 64, 128, or 256; contiguous token_block_size must " - "be 1, 4, or 16): headDim=%d partition=%d tokenBlockSize=%d", + "be 1, 4, 16, or 32): headDim=%d partition=%d tokenBlockSize=%d", params.headDim, static_cast(params.partition), params.tokenBlockSize); } @@ -538,6 +621,7 @@ int withPartitionTraits(SageScalePartition partition, int tokenBlockSize, int fa case 1: return fn(PartitionTraits>{}); case 4: return fn(PartitionTraits>{}); case 16: return fn(PartitionTraits>{}); + case 32: return fn(PartitionTraits>{}); default: break; } break; diff --git a/cpp/tensorrt_llm/common/sageQuant.h b/cpp/tensorrt_llm/common/sageQuant.h index 8be6835de483..1c0691243f04 100644 --- a/cpp/tensorrt_llm/common/sageQuant.h +++ b/cpp/tensorrt_llm/common/sageQuant.h @@ -46,7 +46,8 @@ struct SageQuantParams int batchSize{}; int numHeads{}; int headDim{}; - // Only read when partition == Contiguous; the Hopper partitions fix their own group size. + // Only read when partition == Contiguous, where it is 1, 4, 16, or 32; the Hopper + // partitions fix their own group size. int tokenBlockSize{}; SageScalePartition partition{SageScalePartition::Contiguous}; bool kSmooth{false}; @@ -57,7 +58,7 @@ struct SageQuantParams kernels::Data_type quantType{kernels::DATA_TYPE_E4M3}; float* ptrQkScale{nullptr}; // Optional source and scratch mean used to perform K-smoothing. - // (See below) collected at vStage==0, applied at vStage==1. + // (See below) collected at vStage==1, applied at vStage==2. void const* ptrKForMean{nullptr}; float* ptrKMean{nullptr}; // Optional arguments for SageQuantV: @@ -68,6 +69,11 @@ struct SageQuantParams void const* ptrV{nullptr}; void* ptrVQuant{nullptr}; float* ptrVScale{nullptr}; + // Optional cumulative KV sequence lengths, [batchSize + 1]. When set, the V scales and the K + // mean are collected per sequence, [batchSize, numHeadsV, headDim], as SageAttention and the + // PrimTS block-sparse kernel expect; otherwise one [numHeadsV, headDim] row set covers the + // whole batch, as the trtllm-gen kernels expect. + int const* ptrCuSeqLensV{nullptr}; // Hardware info. Required. int smCount{}; cudaStream_t stream{}; diff --git a/cpp/tensorrt_llm/thop/CMakeLists.txt b/cpp/tensorrt_llm/thop/CMakeLists.txt index ac9053a76e18..49d749203836 100644 --- a/cpp/tensorrt_llm/thop/CMakeLists.txt +++ b/cpp/tensorrt_llm/thop/CMakeLists.txt @@ -145,9 +145,13 @@ add_th_common_family(th_common_sequence causalConv1dOp.cpp mambaConv1dOp.cpp mamba2MTPSSMCacheOp.cpp selectiveScanOp.cpp) add_th_common_family( - th_common_visual_gen fusedAdaptiveLayerNormOp.cpp fusedDiTQKNormRopeOp.cpp - fusedDiTSplitQKNormRopeOp.cpp fusedDiTSplitNormOp.cpp - fusedDiTGateResidNormShiftScaleOp.cpp) + th_common_visual_gen + fusedAdaptiveLayerNormOp.cpp + fusedDiTQKNormRopeOp.cpp + fusedDiTSplitQKNormRopeOp.cpp + fusedDiTSplitNormOp.cpp + fusedDiTGateResidNormShiftScaleOp.cpp + sageQuantizeOp.cpp) add_th_common_family( th_common_decoding diff --git a/cpp/tensorrt_llm/thop/sageQuantizeOp.cpp b/cpp/tensorrt_llm/thop/sageQuantizeOp.cpp new file mode 100644 index 000000000000..df75ed4e5843 --- /dev/null +++ b/cpp/tensorrt_llm/thop/sageQuantizeOp.cpp @@ -0,0 +1,207 @@ +/* + * 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. + */ + +#include "tensorrt_llm/common/cudaUtils.h" +#include "tensorrt_llm/common/sageQuant.h" +#include "tensorrt_llm/thop/thUtils.h" + +#include +#include + +#include +#include + +TRTLLM_NAMESPACE_BEGIN + +namespace torch_ext +{ + +namespace +{ + +using tensorrt_llm::common::SageQuantParams; +using tensorrt_llm::common::SageScalePartition; + +void checkTokenMajor(torch::Tensor const& tensor, char const* name) +{ + TORCH_CHECK(tensor.is_cuda() && tensor.dim() == 3 && tensor.is_contiguous(), name, + " must be a contiguous CUDA [tokens, heads, head_dim] tensor"); +} + +void checkCuSeqLens(torch::Tensor const& tensor, char const* name) +{ + TORCH_CHECK(tensor.is_cuda() && tensor.dim() == 1 && tensor.is_contiguous() && tensor.scalar_type() == torch::kInt32 + && tensor.size(0) > 1, + name, " must be a contiguous CUDA int32 [batch_size + 1] tensor"); +} + +// The kernel reads the V scales and the means with 16-byte vector loads. +void checkScale(torch::Tensor const& tensor, int64_t numel, char const* name) +{ + TORCH_CHECK(tensor.is_cuda() && tensor.is_contiguous() && tensor.scalar_type() == torch::kFloat32 + && tensor.numel() == numel && reinterpret_cast(tensor.data_ptr()) % 16 == 0, + name, " must be a contiguous, 16-byte aligned CUDA float32 tensor with ", numel, " elements"); +} + +void checkKeyValue(torch::Tensor const& key, torch::Tensor const& value, torch::Tensor const& k_quant, + torch::Tensor const& v_quant, torch::Tensor const& cu_seqlens_kv) +{ + checkTokenMajor(key, "key"); + checkTokenMajor(value, "value"); + TORCH_CHECK( + key.scalar_type() == torch::kBFloat16 || key.scalar_type() == torch::kFloat16, "key must be BF16 or FP16"); + TORCH_CHECK(value.scalar_type() == key.scalar_type(), "key and value must share a dtype"); + TORCH_CHECK(key.sizes() == value.sizes(), "key and value must have the same shape"); + checkCuSeqLens(cu_seqlens_kv, "cu_seqlens_kv"); + checkTokenMajor(k_quant, "k_quant"); + checkTokenMajor(v_quant, "v_quant"); + TORCH_CHECK(k_quant.sizes() == key.sizes() && v_quant.sizes() == value.sizes(), + "quantized outputs must match the input shapes"); + TORCH_CHECK(k_quant.scalar_type() == torch::kInt8 || k_quant.scalar_type() == torch::kFloat8_e4m3fn, + "k_quant must be int8 or float8_e4m3fn"); + TORCH_CHECK(v_quant.scalar_type() == torch::kFloat8_e4m3fn, "v_quant must be float8_e4m3fn"); +} + +// Scales per head of a Q/K tensor quantized in token_block_size-token blocks. +int checkBlockScales( + torch::Tensor const& input, torch::Tensor const& scale, int64_t token_block_size, int batchSize, char const* name) +{ + int const stride = tensorrt_llm::common::getSageScaleHeadStride( + SageScalePartition::Contiguous, static_cast(token_block_size), static_cast(input.size(0)), batchSize); + TORCH_CHECK(stride > 0, "q_block_size and k_block_size must be 1, 4, 16, or 32"); + checkScale(scale, input.size(1) * stride, name); + return stride; +} + +// The V part and the hardware fields of the kernel parameters. +SageQuantParams makeParams(torch::Tensor const& value, torch::Tensor& v_quant, torch::Tensor const& v_scale, + torch::Tensor const& qk_quant, int batchSize, int const* ptrCuSeqLensV) +{ + SageQuantParams params{}; + params.headDim = static_cast(value.size(2)); + params.inputType = value.scalar_type() == torch::kBFloat16 ? tensorrt_llm::kernels::DATA_TYPE_BF16 + : tensorrt_llm::kernels::DATA_TYPE_FP16; + params.quantType = qk_quant.scalar_type() == torch::kInt8 ? tensorrt_llm::kernels::DATA_TYPE_INT8 + : tensorrt_llm::kernels::DATA_TYPE_E4M3; + params.partition = SageScalePartition::Contiguous; + params.batchSize = batchSize; + params.sumSeqLensV = static_cast(value.size(0)); + params.numHeadsV = static_cast(value.size(1)); + params.ptrV = value.data_ptr(); + params.ptrVQuant = v_quant.data_ptr(); + params.ptrVScale = v_scale.data_ptr(); + params.ptrCuSeqLensV = ptrCuSeqLensV; + params.smCount = tensorrt_llm::common::getMultiProcessorCount(); + params.stream = at::cuda::getCurrentCUDAStream(value.get_device()); + return params; +} + +// Quantizes Q or K in token_block_size-token blocks; vStage says what the launch does to V. +void launch(SageQuantParams params, torch::Tensor const& input, torch::Tensor const& cu_seqlens, torch::Tensor& quant, + torch::Tensor& scale, int64_t token_block_size, int vStage) +{ + params.sumSeqLensQk = static_cast(input.size(0)); + params.ptrCuSeqLensQk = cu_seqlens.data_ptr(); + params.numHeads = static_cast(input.size(1)); + params.tokenBlockSize = static_cast(token_block_size); + params.ptrQk = input.data_ptr(); + params.ptrQkQuant = quant.data_ptr(); + params.ptrQkScale = scale.data_ptr(); + params.vStage = vStage; + tensorrt_llm::common::invokeSageQuant(params); +} + +} // namespace + +// SageAttention Q/K/V quantization with contiguous token-block scales. +// +// Q and K are quantized to INT8 or FP8 E4M3 with one scale per q_block_size / k_block_size tokens +// of a sequence, laid out per head as tensorrt_llm::common::getSageScaleHeadStride describes. V is +// quantized to FP8 E4M3 with one scale per channel of each sequence. The V scales and the K mean +// are [batch, heads_kv, head_dim]; k_mean is the K-smoothing mean and is filled exactly when given. +// +// With reuse_stats, only K and V are quantized, with the V scales and the K mean of an earlier call +// as inputs, for tensors that must dequantize like that call's K and V, such as the K/V block means +// of PrimTS proxy routes; K still takes its own scales. The query arguments must then be absent. +void sage_quantize(std::optional const& query, torch::Tensor const& key, torch::Tensor const& value, + std::optional const& cu_seqlens_q, torch::Tensor const& cu_seqlens_kv, + std::optional q_quant, torch::Tensor& k_quant, torch::Tensor& v_quant, + std::optional q_scale, torch::Tensor& k_scale, torch::Tensor& v_scale, + std::optional k_mean, int64_t q_block_size, int64_t k_block_size, bool reuse_stats) +{ + checkKeyValue(key, value, k_quant, v_quant, cu_seqlens_kv); + int const batchSize = static_cast(cu_seqlens_kv.size(0) - 1); + checkBlockScales(key, k_scale, k_block_size, batchSize, "k_scale"); + int64_t const statRows = batchSize * value.size(1); + checkScale(v_scale, statRows * value.size(2), "v_scale"); + if (k_mean.has_value()) + { + checkScale(*k_mean, statRows * key.size(2), "k_mean"); + } + + SageQuantParams params = makeParams(value, v_quant, v_scale, k_quant, batchSize, cu_seqlens_kv.data_ptr()); + params.kSmooth = k_mean.has_value(); + params.ptrKMean = params.kSmooth ? k_mean->data_ptr() : nullptr; + if (reuse_stats) + { + TORCH_CHECK(!query.has_value() && !cu_seqlens_q.has_value() && !q_quant.has_value() && !q_scale.has_value(), + "reuse_stats quantizes K and V only: query, cu_seqlens_q, q_quant, and q_scale must be absent"); + launch(params, key, cu_seqlens_kv, k_quant, k_scale, k_block_size, 2); + return; + } + + TORCH_CHECK(query.has_value() && cu_seqlens_q.has_value() && q_quant.has_value() && q_scale.has_value(), + "query, cu_seqlens_q, q_quant, and q_scale are required unless reuse_stats is set"); + checkTokenMajor(*query, "query"); + TORCH_CHECK(query->scalar_type() == key.scalar_type(), "query, key, and value must share a dtype"); + TORCH_CHECK(query->size(2) == key.size(2), "query and key must have the same head dimension"); + checkCuSeqLens(*cu_seqlens_q, "cu_seqlens_q"); + TORCH_CHECK(cu_seqlens_q->size(0) == cu_seqlens_kv.size(0), "Q and KV must have the same batch size"); + checkTokenMajor(*q_quant, "q_quant"); + TORCH_CHECK(q_quant->sizes() == query->sizes(), "quantized outputs must match the input shapes"); + TORCH_CHECK(q_quant->scalar_type() == k_quant.scalar_type(), "q_quant and k_quant must share a dtype"); + checkBlockScales(*query, *q_scale, q_block_size, batchSize, "q_scale"); + + // The V scales and the K mean are reduced with atomics into zeroed buffers. + v_scale.zero_(); + if (k_mean.has_value()) + { + k_mean->zero_(); + } + + params.ptrKForMean = params.kSmooth ? key.data_ptr() : nullptr; + // Q, while collecting the V scales and the K mean; then K (smoothed by the K mean) and V. + launch(params, *query, *cu_seqlens_q, *q_quant, *q_scale, q_block_size, 1); + launch(params, key, cu_seqlens_kv, k_quant, k_scale, k_block_size, 2); +} + +TORCH_LIBRARY_FRAGMENT(trtllm, m) +{ + m.def( + "sage_quantize(Tensor? query, Tensor key, Tensor value, Tensor? cu_seqlens_q, Tensor cu_seqlens_kv, " + "Tensor(a!)? q_quant, Tensor(b!) k_quant, Tensor(c!) v_quant, Tensor(d!)? q_scale, Tensor(e!) k_scale, " + "Tensor(f!) v_scale, Tensor(g!)? k_mean, int q_block_size, int k_block_size, bool reuse_stats=False) -> ()"); +} + +TORCH_LIBRARY_IMPL(trtllm, CUDA, m) +{ + m.impl("sage_quantize", &sage_quantize); +} + +} // namespace torch_ext + +TRTLLM_NAMESPACE_END diff --git a/docs/source/features/visualgen-sparse-attention.md b/docs/source/features/visualgen-sparse-attention.md index 095790c27520..cac4431683d9 100644 --- a/docs/source/features/visualgen-sparse-attention.md +++ b/docs/source/features/visualgen-sparse-attention.md @@ -234,11 +234,11 @@ softmax. VisualGen serves it on two backends behind one `SolAttentionConfig`: attention metadata and the `timestep` in the forward arguments; dense layers and dense timestep phases return no routes. The VisualGen wrapper compacts fused projection split views once and shares those tensors between - prediction and the block-sparse FMHA. With a SageAttention recipe, the FMHA - quantizes Q/K/V and the proxy means after prediction; the dense timestep - phases use the dense SageAttention kernel. Cross-attention, context - parallelism and unsupported tensor envelopes raise an error instead of - silently falling back to dense attention. + prediction and the block-sparse FMHA. With a SageAttention recipe, the + attention core quantizes Q/K/V and the proxy means for the FMHA after + prediction; the dense timestep phases use the dense SageAttention kernel. + Cross-attention, context parallelism and unsupported tensor envelopes raise + an error instead of silently falling back to dense attention. - **CUTEDSL** runs the fused kernel vendored from the reference implementation (`SOLCuTeDSLAttention`), which folds routing, sparse computation and the approximation correction into one online-softmax pass. It runs on datacenter @@ -320,7 +320,7 @@ Requirements: - A VSA-fine-tuned checkpoint such as [`FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers`](https://huggingface.co/FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers). Standard Wan checkpoints do not carry the learned VSA gates. - The `CUTEDSL` or `TRTLLM` attention backend. `CUTEDSL` uses the CuTe DSL fine-stage kernel; `TRTLLM` lowers the selected cubes through the generic PrimTS block-sparse FMHA contract. - A supported CUDA device and tensor shape for the selected block-sparse kernel. When that kernel is unavailable or the input is outside its envelope, the fine branch runs the compact Q/K/V through that backend's dense path (SDPA for `CUTEDSL`, TRTLLM attention for `TRTLLM`) and the VSA post-processing still applies. -- On `TRTLLM`, VSA can be combined with a SageAttention `quant_attention_config` for BF16 models; the block-sparse fine stage then runs with quantized Q/K/V. `CUTEDSL` rejects `quant_attention_config` with VSA. +- On `TRTLLM`, VSA can be combined with a SageAttention `quant_attention_config`; the block-sparse fine stage then runs with quantized Q/K/V. `CUTEDSL` rejects `quant_attention_config` with VSA. - VSA is not compatible with Ring attention or Attention2D (it does not produce per-split LSE). Ulysses is supported. Configure VSA with `VideoSparseAttentionConfig` and either backend: diff --git a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index b871bc642cc3..2e9055d21444 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md @@ -234,12 +234,17 @@ only static format, proxy, geometry, and capacity choices; every run receives the live routes, summaries, validity bits, page tables, and sequence lengths. The contiguous block-sparse path also consumes the SageAttention recipe in -`AttentionForwardArgs`. Prediction sees the original BF16 Q/K/V. The adapter -then quantizes Q/K/V with FlashInfer's Sage quantizer, quantizes the summaries -with one scale per K mean and the original V channel scales, and passes all -scales to PrimTS. The recipe and quantized dtypes are part of the plan key. -Sage requires BF16 input, because the quantizer's FP16 reciprocal scale can -overflow on zero or tiny Q/K blocks. Paged block-sparse attention rejects Sage. +`AttentionForwardArgs`. Prediction sees the original Q/K/V. When a library of +the layer declares `supports_sage_quantized_inputs`, the attention core then +quantizes the Q/K/V of a block-sparse call before FMHA selection, with +TensorRT-LLM's Sage quantization kernel (`attention/kernels/sage_quant.py`), +and passes the result in `AttentionForwardArgs.sage_quantized`. Q/K get +per-block scales and V per-sequence channel scales; the summaries of proxy +routes get one scale per K mean and reuse the original V channel scales and K +mean. `PrimsTSBlockSparseFmha` serves Sage only with these inputs and passes +all scales to PrimTS; the attention op ignores them and quantizes dense +SageAttention inside its kernels. The recipe and quantized dtypes are part of +the plan key. Paged block-sparse attention rejects Sage. `PrimsTSBlockSparseFmha` keeps its planned wrappers in the attention metadata (`TrtllmAttentionMetadata.fmha_plan_caches`), like the FlashInfer wrappers and diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py index 448d64df9925..2c344c28680a 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py @@ -50,6 +50,7 @@ "sparse_attn_phase", # host-resolved sparse phase, consumed before FMHA dispatch "qk_fp8_per_tensor", # per-tensor FP8 recipe flag, consumed by FMHA selection and the PrimTS FMHA "v_fp8_per_tensor", # per-tensor FP8 recipe flag, consumed by FMHA selection and the PrimTS FMHA + "sage_quantized", # for libraries taking quantized Sage inputs; the op quantizes its own } ) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/interface.py b/tensorrt_llm/_torch/attention/backends/fmha/interface.py index e3905dd1964d..d65aeff14d6e 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/interface.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/interface.py @@ -43,6 +43,9 @@ class Fmha(ABC): supports_skip_correction: ClassVar[bool] = False supports_block_sparse_inputs: ClassVar[bool] = False supports_per_tensor_fp8: ClassVar[bool] = False + # Reads the SageAttention inputs the attention core quantized from + # ``AttentionForwardArgs.sage_quantized``. + supports_sage_quantized_inputs: ClassVar[bool] = False supports_workspace_reclamation: bool = False supports_fp4_mla: ClassVar[bool] = False diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py index 29bde4a063d4..2bf1554d3753 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py @@ -316,11 +316,7 @@ def _is_supported_with_reason( feature_reason = get_attention_feature_unsupported_reason(meta, fwd) if feature_reason is not None: return False, feature_reason - if ( - fwd.sage_attn_num_elts_per_blk_q > 0 - or fwd.sage_attn_num_elts_per_blk_k > 0 - or fwd.sage_attn_num_elts_per_blk_v > 0 - ): + if fwd.sage_attn_qk_dtype is not None: return False, "SageAttention is not supported." if meta.kv_cache_manager is None: return self._is_supported_without_kv_cache(q, k, v, attn, meta, fwd) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py index 9a6c6d4b4a15..6494ebc9b434 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py @@ -20,7 +20,6 @@ from typing import TYPE_CHECKING, Literal, cast import torch -from flashinfer import trtllm_sage_attention_quantize from tensorrt_llm._torch.attention.backends.interface import ( AttentionForwardArgs, @@ -35,7 +34,6 @@ SageAttentionParams, ) from tensorrt_llm._torch.attention.backends.sparse.params import BlockSparseForwardInputs -from tensorrt_llm._torch.attention.kernels.sage_quant import quantize_sage_summaries from tensorrt_llm.logger import logger from .interface import FmhaPhase @@ -202,35 +200,16 @@ def _sage_recipe( forward_args: AttentionForwardArgs, ) -> tuple[SageAttentionConfig, torch.dtype] | None: """Return the forward's SageAttention plan config and quantized Q/K dtype, if any.""" - q_block = forward_args.sage_attn_num_elts_per_blk_q - k_block = forward_args.sage_attn_num_elts_per_blk_k - if q_block <= 0 and k_block <= 0 and forward_args.sage_attn_num_elts_per_blk_v <= 0: + qk_dtype = forward_args.sage_attn_qk_dtype + if qk_dtype is None: return None # Every proxy K summary is quantized with its own scale. - config = SageAttentionConfig(q_block_size=q_block, k_block_size=k_block, k_summary_block_size=1) - return config, torch.int8 if forward_args.sage_attn_qk_int8 else torch.float8_e4m3fn - - -def _sage_unsupported_reason( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - forward_args: AttentionForwardArgs, -) -> str | None: - """Return why the forward's Sage recipe cannot quantize these contiguous inputs.""" - if _sage_recipe(forward_args) is None: - return None - if ( - forward_args.sage_attn_num_elts_per_blk_q not in (1, 4, 16) - or forward_args.sage_attn_num_elts_per_blk_k not in (1, 4, 16) - or forward_args.sage_attn_num_elts_per_blk_v != 1 - ): - return "Sage block-sparse attention requires Q/K blocks of 1, 4, or 16 and V block size 1" - if q.dtype != torch.bfloat16 or k.dtype != q.dtype or v.dtype != q.dtype: - # The quantizer computes reciprocal scales in the input dtype, which - # overflows FP16 on zero or tiny Q/K blocks such as VSA padding. - return "Sage block-sparse attention requires matching BF16 Q/K/V inputs" - return None + config = SageAttentionConfig( + q_block_size=forward_args.sage_attn_num_elts_per_blk_q, + k_block_size=forward_args.sage_attn_num_elts_per_blk_k, + k_summary_block_size=1, + ) + return config, qk_dtype class PrimsTSBlockSparseFmha(PrimsTSFmha): @@ -239,6 +218,7 @@ class PrimsTSBlockSparseFmha(PrimsTSFmha): PLAN_CACHE_KEY = "prims_ts_block_sparse" supports_block_sparse_inputs = True + supports_sage_quantized_inputs = True def _is_supported( self, @@ -389,13 +369,12 @@ def _contiguous_unsupported_reason( ) -> str | None: if forward_args.is_fused_qkv or k is None or v is None: return "contiguous block-sparse attention requires separate Q, K, and V" + if forward_args.sage_attn_qk_dtype is not None and forward_args.sage_quantized is None: + return "SageAttention requires inputs quantized by the attention core" if self.attn.position_embedding_type != 0 or forward_args.mrope_position_deltas is not None: return "contiguous Q/K/V must have position embedding applied before attention" if forward_args.cu_q_seqlens is not None or forward_args.cu_kv_seqlens is not None: return "packed variable-length Q/KV inputs are not supported" - sage_reason = _sage_unsupported_reason(q, k, v, forward_args) - if sage_reason is not None: - return sage_reason inputs = _get_block_sparse_inputs(forward_args) assert inputs is not None mask_type = self._get_prims_mask_type(forward_args) @@ -427,7 +406,7 @@ def _paged_unsupported_reason( assert inputs is not None if inputs.sparse_format != "bsr" or inputs.use_proxy_routes: return "paged block-sparse attention only supports BSR exact routes" - if _sage_recipe(forward_args) is not None: + if forward_args.sage_attn_qk_dtype is not None: return "paged block-sparse attention does not support SageAttention" if not forward_args.is_fused_qkv: return "paged block-sparse attention requires fused QKV input" @@ -603,56 +582,20 @@ def _forward_contiguous( wrapper = cast("BlockSparseTSWrapper", self._get_or_plan_wrapper(key, metadata)) sage_params = None if key.sage is not None: - # Prediction has consumed the unquantized tensors. Device cumulative - # lengths keep the per-sequence scale layout without a host-to-device - # copy during CUDA Graph capture. - cu_seqlens_q = torch.arange( - 0, - (batch_size + 1) * key.seq_len_q, - key.seq_len_q, - dtype=torch.int32, - device=q.device, - ) - cu_seqlens_kv = torch.arange( - 0, - (batch_size + 1) * key.kv_capacity, - key.kv_capacity, - dtype=torch.int32, - device=q.device, - ) - quantized = trtllm_sage_attention_quantize( - query.flatten(0, 1), - key_states.flatten(0, 1), - value_states.flatten(0, 1), - q_block_size=key.sage.q_block_size, - k_block_size=key.sage.k_block_size, - qk_quant_dtype=key.qk_dtype, - cum_seq_lens_q=cu_seqlens_q, - cum_seq_lens_kv=cu_seqlens_kv, - smooth_k=forward_args.sage_attn_smooth_k, - ) - q_quant, k_quant, v_quant, q_scale, k_scale, v_scale = quantized[:6] - query = q_quant.view_as(query) - key_states = k_quant.view_as(key_states) - value_states = v_quant.view_as(value_states) - k_summary_scale = None + quantized = forward_args.sage_quantized + assert quantized is not None, "Sage inputs are quantized before the FMHA forward" + query = quantized.q.view_as(query) + key_states = quantized.k.view_as(key_states) + value_states = quantized.v.view_as(value_states) if inputs.use_proxy_routes: - k_summary = inputs.k_summary - if forward_args.sage_attn_smooth_k: - # Proxy logits must shift by the same per-row constant as the - # smoothed exact keys, so the summaries subtract the same mean. - k_summary = (k_summary.float() - quantized[6]).to(k_summary.dtype) - k_summary, v_summary, k_summary_scale = quantize_sage_summaries( - k_summary, inputs.v_summary, v_scale, key.qk_dtype + inputs = replace( + inputs, k_summary=quantized.k_summary, v_summary=quantized.v_summary ) - inputs = replace(inputs, k_summary=k_summary, v_summary=v_summary) sage_params = SageAttentionParams( - q_scale=q_scale, - k_scale=k_scale, - # PrimTS takes one per-channel V scale row per batch entry; the - # quantizer collects a single row over the whole batch. - v_scale=v_scale.expand(batch_size, -1, -1).contiguous(), - k_summary_scale=k_summary_scale, + q_scale=quantized.q_scale, + k_scale=quantized.k_scale, + v_scale=quantized.v_scale, + k_summary_scale=quantized.k_summary_scale, ) wrapper.run( query, diff --git a/tensorrt_llm/_torch/attention/backends/interface.py b/tensorrt_llm/_torch/attention/backends/interface.py index 7ccf04220604..f9a5e631bfe4 100644 --- a/tensorrt_llm/_torch/attention/backends/interface.py +++ b/tensorrt_llm/_torch/attention/backends/interface.py @@ -17,6 +17,7 @@ from ...modules.mamba.mamba2_metadata import Mamba2Metadata from ...speculative.interface import SpecMetadata from ...speculative.spec_tree_manager import SpecTreeManager + from ..kernels.sage_quant import SageQuantized from tensorrt_llm._utils import get_hf_rope_theta, maybe_pin_memory from tensorrt_llm.functional import (AttentionMaskType, PositionEmbeddingType, @@ -968,6 +969,9 @@ class AttentionForwardArgs: sage_attn_num_elts_per_blk_v: int = 0 sage_attn_qk_int8: bool = False sage_attn_smooth_k: bool = False + # SageAttention inputs that the attention core quantized for the FMHA libraries + # taking quantized inputs; the attention op quantizes inside its kernels. + sage_quantized: Optional["SageQuantized"] = None # Per-tensor FP8 recipe of the PrimTS FMHA library. Quantize Q/K and/or V to # E4M3 with one scale each, folded into the softmax and output scales. qk_fp8_per_tensor: bool = False @@ -998,6 +1002,15 @@ def mask_type(self) -> int: raise ValueError( f"Unexpected attention mask type: {self.attention_mask!r}") + @property + def sage_attn_qk_dtype(self) -> Optional[torch.dtype]: + """Quantized Q/K dtype of the SageAttention recipe, or ``None`` without one.""" + if (self.sage_attn_num_elts_per_blk_q <= 0 + and self.sage_attn_num_elts_per_blk_k <= 0 + and self.sage_attn_num_elts_per_blk_v <= 0): + return None + return torch.int8 if self.sage_attn_qk_int8 else torch.float8_e4m3fn + _ATTENTION_FORWARD_ARGS_FIELDS = frozenset( AttentionForwardArgs.__dataclass_fields__) diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 909e208129b0..dd69011f255e 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -42,6 +42,7 @@ from ...pyexecutor.config_utils import is_mla from ...utils import (compute_swizzled_sf_shape, get_global_attrs, get_model_extra_attrs, helix_local_len_tensor) +from ..kernels.sage_quant import SageQuantized, sage_quantize from .fmha.manager import FmhaManager from .fp4_mla import can_fuse_fp4_mla_q_quant, scatter_fp4_mla_kv_cache from .fp4_mla.state import Fp4MlaState @@ -2014,6 +2015,73 @@ def _compute_flash_mla_metadata(self, self.kv_lora_rank, ) + def _quantize_sage_inputs( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: TrtllmAttentionMetadata, + forward_args: AttentionForwardArgs, + ) -> Optional[SageQuantized]: + """Quantize the SageAttention inputs of a block-sparse call. + + The attention op quantizes dense SageAttention inside its kernels, so + only block-sparse calls are quantized here, for the FMHA libraries that + take quantized inputs: separate Q, K, and V without a KV cache. + """ + qk_quant_dtype = forward_args.sage_attn_qk_dtype + if (qk_quant_dtype is None or k is None or v is None + or metadata.kv_cache_manager is not None + or not any(fmha.supports_sage_quantized_inputs + for fmha in self._fmha_manager.fmha_libs)): + return None + inputs = forward_args.sparse_runtime_params.block_sparse_inputs + if inputs is None: + return None + num_seqs = metadata.num_contexts + # Device lengths add no host-to-device copy during CUDA Graph capture. + cu_seqlens_q = torch.zeros(num_seqs + 1, + dtype=torch.int32, + device=q.device) + cu_seqlens_kv = torch.zeros_like(cu_seqlens_q) + cu_seqlens_q[1:] = metadata.seq_lens_cuda[:num_seqs].cumsum( + 0, dtype=torch.int32) + cu_seqlens_kv[1:] = metadata.seq_lens_kv_cuda[:num_seqs].cumsum( + 0, dtype=torch.int32) + quantized = sage_quantize( + q.view(-1, self.num_heads, self.head_dim), + k.view(-1, self.num_kv_heads, self.head_dim), + v.view(-1, self.num_kv_heads, self.head_dim), + cu_seqlens_q, + cu_seqlens_kv, + q_block_size=forward_args.sage_attn_num_elts_per_blk_q, + k_block_size=forward_args.sage_attn_num_elts_per_blk_k, + qk_quant_dtype=qk_quant_dtype, + smooth_k=forward_args.sage_attn_smooth_k, + ) + if not inputs.use_proxy_routes: + return quantized + # The K/V block means of proxy routes dequantize like the keys and + # values they stand for: every mean is a key of its sequence with a + # scale of its own, and reuses the V scales and K mean of this call. + batch, blocks = inputs.k_summary.shape[:2] + cu_seqlens_blocks = torch.arange(0, (batch + 1) * blocks, + blocks, + dtype=torch.int32, + device=q.device) + summary = sage_quantize(None, + inputs.k_summary.flatten(0, 1), + inputs.v_summary.flatten(0, 1), + None, + cu_seqlens_blocks, + k_block_size=1, + qk_quant_dtype=qk_quant_dtype, + stats=quantized) + return quantized._replace( + k_summary=summary.k.view(inputs.k_summary.shape), + v_summary=summary.v.view(inputs.v_summary.shape), + k_summary_scale=summary.k_scale) + def _ensure_rope_table_size(self, required_max_positions: int) -> None: floats_per_position = self.rope_params.dim * ( 2 if self.rope_params.duplicate_data else 1) @@ -2431,6 +2499,11 @@ def forward( assert metadata.kv_cache_manager is None assert metadata.num_contexts == metadata.num_seqs + # Sparse prediction has read the unquantized inputs; FMHA support checks + # see the quantized ones. + forward_args.sage_quantized = self._quantize_sage_inputs( + q, k, v, metadata, forward_args) + # On a ragged generation step the attention op is handed # one row per query token so its `num_tokens % num_seqs == 0` check # passes with seq_len == 1. Keep the presentation scoped around both diff --git a/tensorrt_llm/_torch/attention/kernels/sage_quant.py b/tensorrt_llm/_torch/attention/kernels/sage_quant.py index 570f4a93fd02..f76a80457951 100644 --- a/tensorrt_llm/_torch/attention/kernels/sage_quant.py +++ b/tensorrt_llm/_torch/attention/kernels/sage_quant.py @@ -12,124 +12,164 @@ # 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. -"""Quantization of the K/V block means consumed by Sage block-sparse attention.""" +"""SageAttention quantization with TensorRT-LLM's Sage quantization kernel. + +``sage_quantize`` quantizes Q, K, and V, or, given the statistics of an earlier call, only K and +V so that they dequantize like that call's K and V, as the K/V block means of proxy routes must. +The attention core runs it ahead of the FMHA libraries that take quantized inputs. +""" + +from typing import NamedTuple import torch -import triton -import triton.language as tl -from torch.library import triton_op, wrap_triton -from triton.language.extra.cuda import libdevice - -__all__ = ["quantize_sage_summaries"] - - -@triton.jit -def _quantize_sage_summaries_kernel( - k_ptr, - v_ptr, - v_scale_ptr, - quant_k_ptr, - quant_v_ptr, - k_scale_ptr, - num_rows, - blocks, - heads, - scale_stride, - DIM: tl.constexpr, - INT8: tl.constexpr, - ROWS: tl.constexpr, - WIDTH: tl.constexpr, -): - rows = tl.program_id(0) * ROWS + tl.arange(0, ROWS) - head = tl.program_id(1) - channels = tl.arange(0, WIDTH) - valid_rows = rows < num_rows - valid = valid_rows[:, None] & (channels < DIM)[None, :] - offsets = (rows[:, None] * heads + head) * DIM + channels[None, :] - k = tl.load(k_ptr + offsets, valid, other=0).to(tl.float32) - v = tl.load(v_ptr + offsets, valid, other=0).to(tl.float32) - - maximum = tl.maximum(tl.max(tl.abs(k), axis=1), 1e-3) - divisor: tl.constexpr = 126.9 if INT8 else 448.0 - k_scale = tl.div_rn(maximum, divisor) - normalized_k = tl.div_rn(k, k_scale[:, None]) - if INT8: - quant_k = tl.minimum(tl.maximum(libdevice.nearbyint(normalized_k), -127.0), 127.0) - else: - quant_k = tl.minimum(tl.maximum(normalized_k, -448.0), 448.0) - v_scale = tl.load(v_scale_ptr + head * DIM + channels, channels < DIM, other=1.0) - quant_v = tl.minimum(tl.maximum(tl.div_rn(v, v_scale[None, :]), -448.0), 448.0) - tl.store(quant_k_ptr + offsets, quant_k, valid) - tl.store(quant_v_ptr + offsets, quant_v, valid) - - # Sequence b starts at scale slot b * (blocks + 1): the flat layout leaves one - # unused slot between consecutive sequences. - tl.store(k_scale_ptr + head * scale_stride + rows + rows // blocks, k_scale, valid_rows) - - -@triton_op("trtllm::quantize_sage_summaries", mutates_args={}) -def quantize_sage_summaries( - k_summary: torch.Tensor, - v_summary: torch.Tensor, - v_scale: torch.Tensor, - qk_quant_dtype: torch.dtype, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Quantize key/value block means for the proxy routes of Sage attention. - Every key mean gets its own scale; value means reuse the per-channel scales of - the original V, so exact and proxy routes share one PV dequantization. +__all__ = [ + "SageQuantized", + "sage_quantize", +] - Args: - k_summary: Contiguous FP16/BF16 CUDA means, ``[batch, blocks, heads, head_dim]``. - v_summary: Value means with the same shape, dtype, and device as ``k_summary``. - v_scale: Contiguous positive FP32 ``[heads, head_dim]`` dequantization scales - from quantizing the original V. - qk_quant_dtype: ``torch.int8`` or ``torch.float8_e4m3fn`` for the key means. - - Returns: - Quantized key means, FP8 E4M3 value means, and FP32 key scales of shape - ``[heads, batch * blocks + batch - 1]``. Key scale slot ``b * (blocks + 1) + n`` - belongs to mean ``[b, n]``; the unused slot between sequences is not written. + +class SageQuantized(NamedTuple): + """Quantized Q/K/V and their FP32 dequantization scales. + + ``q_scale`` and ``k_scale`` are ``[heads, ceil(tokens / block) + batch - 1]``: sequence + ``b`` starts at slot ``cu_seqlens[b] // block + b``. ``v_scale`` and ``k_mean`` are + ``[batch, heads_kv, head_dim]``; ``k_mean`` is the mean subtracted from K before + quantization and is ``None`` without K smoothing. ``q`` and ``q_scale`` are ``None`` when + the statistics of an earlier call were reused. The summary fields hold the quantized K/V + block means of proxy routes and their per-mean key scales, for calls that have them. """ - if k_summary.ndim != 4 or min(k_summary.shape) <= 0: - raise ValueError("k_summary must have nonempty [batch, blocks, heads, head_dim] shape") - if k_summary.device.type != "cuda" or k_summary.dtype not in (torch.float16, torch.bfloat16): - raise ValueError("k_summary must be a CUDA FP16 or BF16 tensor") - if ( - v_summary.shape != k_summary.shape - or v_summary.dtype != k_summary.dtype - or v_summary.device != k_summary.device - ): - raise ValueError("key and value summaries must have matching shape, dtype, and device") - if v_scale.shape != k_summary.shape[2:] or v_scale.dtype != torch.float32: - raise ValueError("v_scale must be FP32 with shape [heads, head_dim]") - if v_scale.device != k_summary.device: - raise ValueError("v_scale and summaries must be on the same CUDA device") - if not (k_summary.is_contiguous() and v_summary.is_contiguous() and v_scale.is_contiguous()): - raise ValueError("summaries and v_scale must be contiguous") + + q: torch.Tensor | None + k: torch.Tensor + v: torch.Tensor + q_scale: torch.Tensor | None + k_scale: torch.Tensor + v_scale: torch.Tensor + k_mean: torch.Tensor | None = None + k_summary: torch.Tensor | None = None + v_summary: torch.Tensor | None = None + k_summary_scale: torch.Tensor | None = None + + +def _flat_scale_numel(tokens: int, batch: int, block_size: int) -> int: + return -(-tokens // block_size) + batch - 1 + + +def _check_qk_quant_dtype(qk_quant_dtype: torch.dtype) -> None: if qk_quant_dtype not in (torch.int8, torch.float8_e4m3fn): raise ValueError("qk_quant_dtype must be torch.int8 or torch.float8_e4m3fn") - batch, blocks, heads, dim = k_summary.shape - scale_stride = batch * blocks + batch - 1 - quant_k = torch.empty_like(k_summary, dtype=qk_quant_dtype) - quant_v = torch.empty_like(v_summary, dtype=torch.float8_e4m3fn) - k_scale = torch.empty((heads, scale_stride), device=k_summary.device, dtype=torch.float32) - wrap_triton(_quantize_sage_summaries_kernel)[(triton.cdiv(batch * blocks, 4), heads)]( - k_summary, - v_summary, - v_scale, - quant_k, - quant_v, + +def sage_quantize( + query: torch.Tensor | None, + key: torch.Tensor, + value: torch.Tensor, + cu_seqlens_q: torch.Tensor | None, + cu_seqlens_kv: torch.Tensor, + *, + q_block_size: int = 0, + k_block_size: int, + qk_quant_dtype: torch.dtype, + smooth_k: bool = False, + stats: SageQuantized | None = None, +) -> SageQuantized: + """Quantize Q/K/V for SageAttention, with V scales and the K mean per sequence. + + Q and K get one scale per ``q_block_size`` / ``k_block_size`` tokens of a sequence (1, 4, + 16, or 32); V gets one FP8 scale per channel of each sequence. With + ``smooth_k``, every sequence's per-channel K mean is subtracted before K is quantized, + which leaves the attention output unchanged. + + With ``stats``, the result of an earlier call, only K and V are quantized, with the V + scales and the K mean of ``stats``, so that they dequantize like its K and V; K still takes + its own scales. The smoothing follows ``stats``, and the result carries its statistics and + no Q. + + Args: + query: Contiguous BF16/FP16 CUDA ``[tokens_q, heads_q, head_dim]``; ``None`` with + ``stats``. + key: Contiguous ``[tokens_kv, heads_kv, head_dim]`` with the dtype of ``value``. + value: Same shape and dtype as ``key``. + cu_seqlens_q: INT32 ``[batch + 1]`` cumulative query lengths; ``None`` with ``stats``. + cu_seqlens_kv: INT32 ``[batch + 1]`` cumulative key/value lengths. + q_block_size: Query tokens per scale. + k_block_size: Key tokens per scale. + qk_quant_dtype: ``torch.int8`` or ``torch.float8_e4m3fn`` for Q and K. + smooth_k: Subtract the per-sequence K mean before quantizing K. + stats: ``None``, or an earlier result with the batch, heads, and head dimension of + ``key``, whose V scales and K mean to reuse. + """ + _check_qk_quant_dtype(qk_quant_dtype) + if stats is not None and (query is not None or cu_seqlens_q is not None or smooth_k): + raise ValueError("with stats, only K and V are quantized, with the smoothing of stats") + batch = cu_seqlens_kv.numel() - 1 + heads_kv, head_dim = key.shape[1], key.shape[2] + k = torch.empty_like(key, dtype=qk_quant_dtype) + v = torch.empty_like(value, dtype=torch.float8_e4m3fn) + k_scale = torch.empty( + (heads_kv, _flat_scale_numel(key.shape[0], batch, k_block_size)), + dtype=torch.float32, + device=key.device, + ) + if stats is not None: + torch.ops.trtllm.sage_quantize( + None, + key, + value, + None, + cu_seqlens_kv, + None, + k, + v, + None, + k_scale, + stats.v_scale, + stats.k_mean, + 0, + k_block_size, + reuse_stats=True, + ) + return SageQuantized( + q=None, + k=k, + v=v, + q_scale=None, + k_scale=k_scale, + v_scale=stats.v_scale, + k_mean=stats.k_mean, + ) + + q = torch.empty_like(query, dtype=qk_quant_dtype) + q_scale = torch.empty( + (query.shape[1], _flat_scale_numel(query.shape[0], batch, q_block_size)), + dtype=torch.float32, + device=query.device, + ) + v_scale = torch.empty((batch, heads_kv, head_dim), dtype=torch.float32, device=value.device) + k_mean = torch.empty_like(v_scale) if smooth_k else None + torch.ops.trtllm.sage_quantize( + query, + key, + value, + cu_seqlens_q, + cu_seqlens_kv, + q, + k, + v, + q_scale, k_scale, - batch * blocks, - blocks, - heads, - scale_stride, - DIM=dim, - INT8=qk_quant_dtype == torch.int8, - ROWS=4, - WIDTH=triton.next_power_of_2(dim), - num_warps=4, + v_scale, + k_mean, + q_block_size, + k_block_size, + ) + return SageQuantized( + q=q, + k=k, + v=v, + q_scale=q_scale, + k_scale=k_scale, + v_scale=v_scale, + k_mean=k_mean, ) - return quant_k, quant_v, k_scale diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index 18b68a373e21..c6fdaf92480f 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -987,6 +987,24 @@ def _(logits: torch.Tensor, draft_tokens: torch.Tensor, dtype=torch.int32), logits.new_empty( (batch_size, ), dtype=torch.int32) + @torch.library.register_fake("trtllm::sage_quantize") + def _(query: Optional[torch.Tensor], + key: torch.Tensor, + value: torch.Tensor, + cu_seqlens_q: Optional[torch.Tensor], + cu_seqlens_kv: torch.Tensor, + q_quant: Optional[torch.Tensor], + k_quant: torch.Tensor, + v_quant: torch.Tensor, + q_scale: Optional[torch.Tensor], + k_scale: torch.Tensor, + v_scale: torch.Tensor, + k_mean: Optional[torch.Tensor], + q_block_size: int, + k_block_size: int, + reuse_stats: bool = False) -> None: + return None + @torch.library.register_fake("trtllm::fp8_quantize_1x128") def _(input: torch.Tensor, use_ue8m0: bool = False): pad_m = fp4_utils.pad_up(input.shape[0], 4) diff --git a/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/vsa/backend.py b/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/vsa/backend.py index de66d8637ce2..64d792240c5e 100644 --- a/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/vsa/backend.py +++ b/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/vsa/backend.py @@ -68,8 +68,6 @@ def _get_unsupported_primts_reason( k: torch.Tensor, v: torch.Tensor, metadata: VSAMetadata, - *, - sage: bool, ) -> str | None: if q.shape != k.shape or q.shape != v.shape: return "VSA PrimTS requires matching MHA Q/K/V shapes" @@ -77,8 +75,6 @@ def _get_unsupported_primts_reason( return f"VSA PrimTS requires CUDA tensors, got {q.device}" if q.dtype not in (torch.float16, torch.bfloat16): return f"VSA PrimTS requires FP16 or BF16 tensors, got {q.dtype}" - if sage and q.dtype != torch.bfloat16: - return f"VSA PrimTS Sage attention requires BF16 tensors, got {q.dtype}" batch_size, seq_len, num_heads, head_dim = map(int, q.shape) if min(batch_size, seq_len, num_heads, metadata.num_cubes) <= 0: return "VSA PrimTS requires positive batch, sequence, head, and cube extents" @@ -146,9 +142,7 @@ def forward( use_primts = any( isinstance(fmha, PrimsTSBlockSparseFmha) for fmha in self._fmha_manager.fmha_libs ) - unsupported_reason = _get_unsupported_primts_reason( - q, k, v, metadata, sage=self.quant_attention_config is not None - ) + unsupported_reason = _get_unsupported_primts_reason(q, k, v, metadata) if not use_primts: logger.warning_once( "TRTLLM VSA cannot use PrimTS block-sparse attention because the " diff --git a/tests/unittest/_torch/attention/kernels/test_sage_quant.py b/tests/unittest/_torch/attention/kernels/test_sage_quant.py index 979f194d4378..07d9dae28dba 100644 --- a/tests/unittest/_torch/attention/kernels/test_sage_quant.py +++ b/tests/unittest/_torch/attention/kernels/test_sage_quant.py @@ -12,54 +12,89 @@ # 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. -"""Exact parity of the Sage K/V summary quantizer with a PyTorch reference. +"""Parity of the Sage quantizers used by PrimTS block-sparse attention with PyTorch references. -Scale slots are located with the PrimTS flat-layout helpers that the attention -kernel uses to read them. +Scale slots are located with the PrimTS flat-layout helpers that the attention kernel uses to +read them. """ import pytest import torch from tensorrt_llm._torch.attention.backends.prims_ts.sage import flat_scale_numel, flat_scale_slot -from tensorrt_llm._torch.attention.kernels.sage_quant import quantize_sage_summaries +from tensorrt_llm._torch.attention.kernels.sage_quant import SageQuantized, sage_quantize pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def _inputs( - shape: tuple[int, ...], dtype: torch.dtype -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: +# --------------------------------------------------------------------------- summaries +def _summary_inputs( + shape: tuple[int, ...], dtype: torch.dtype, smooth_k: bool = False +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]: generator = torch.Generator(device="cuda").manual_seed(43) k = torch.randn(shape, device="cuda", generator=generator) v = torch.randn(shape, device="cuda", generator=generator) magnitudes = torch.logspace(-5, 2, shape[1], device="cuda").view(1, -1, 1, 1) k, v = (k * magnitudes).to(dtype), (v * magnitudes).to(dtype) # Raw V has larger extrema than its block means, and the attention epilogue - # dequantizes exact and mean V with the same channel scales. - v_scale = v.float().abs().amax(dim=(0, 1)).clamp_min(1e-3) * 2 / 448.0 - return k, v, v_scale + # dequantizes exact and mean V with the same per-sequence channel scales. + v_scale = v.float().abs().amax(dim=1).clamp_min(1e-3) * 2 / 448.0 + k_mean = k.float().mean(dim=1) if smooth_k else None + return k, v, v_scale, k_mean -def _assert_matches_reference( +def _quantize_summaries( + k: torch.Tensor, + v: torch.Tensor, + v_scale: torch.Tensor, + k_mean: torch.Tensor | None, + quant_dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Quantize ``[batch, blocks, heads, head_dim]`` block means as the attention core does: + every mean is a key of its sequence with a scale of its own, reusing the given statistics.""" + batch, blocks = k.shape[:2] + cu_seqlens = torch.arange(0, (batch + 1) * blocks, blocks, dtype=torch.int32, device=k.device) + stats = SageQuantized(None, None, None, None, None, v_scale, k_mean) + actual = sage_quantize( + None, + k.flatten(0, 1), + v.flatten(0, 1), + None, + cu_seqlens, + k_block_size=1, + qk_quant_dtype=quant_dtype, + stats=stats, + ) + return actual.k.view(k.shape), actual.v.view(v.shape), actual.k_scale + + +def _assert_quantized_close( + actual: torch.Tensor, normalized: torch.Tensor, quant_dtype: torch.dtype +) -> None: + """Quantized values round the saturated FP32 reference within half a grid step.""" + limit = 127.0 if quant_dtype == torch.int8 else 448.0 + expected = normalized.clamp(-limit, limit) + if quant_dtype == torch.int8: + tolerance = torch.full_like(expected, 0.51) + else: + # E4M3 keeps 3 mantissa bits: half a step is 2^-4 of the value. + tolerance = expected.abs() * 0.0626 + 2**-9 + assert torch.all((actual.float() - expected).abs() <= tolerance) + + +def _assert_summaries_match_reference( actual: tuple[torch.Tensor, torch.Tensor, torch.Tensor], k: torch.Tensor, v: torch.Tensor, v_scale: torch.Tensor, + k_mean: torch.Tensor | None, quant_dtype: torch.dtype, ) -> None: batch, blocks, heads, _ = k.shape divisor = 126.9 if quant_dtype == torch.int8 else 448.0 - maximum = k.float().abs().amax(dim=-1).clamp_min(1e-3) - # Tensor division keeps FP32 round-to-nearest; dividing by a Python scalar - # multiplies by its rounded reciprocal and can move an FP8 rounding tie. + k_ref = k.float() if k_mean is None else k.float() - k_mean[:, None] + maximum = k_ref.abs().amax(dim=-1).clamp_min(1e-3) scales = maximum / torch.full_like(maximum, divisor) - normalized = k.float() / scales.unsqueeze(-1) - if quant_dtype == torch.int8: - quant_k = normalized.round().clamp(-127, 127).to(quant_dtype) - else: - quant_k = normalized.clamp(-448, 448).to(quant_dtype) - quant_v = (v.float() / v_scale).clamp(-448, 448).to(torch.float8_e4m3fn) slots = flat_scale_slot( torch.arange(batch, device=k.device)[:, None], torch.arange(blocks, device=k.device)[None, :], @@ -68,72 +103,230 @@ def _assert_matches_reference( ).flatten() actual_k, actual_v, actual_scales = actual + assert actual_k.shape == k.shape and actual_v.shape == v.shape assert actual_scales.shape == (heads, flat_scale_numel(batch, blocks, 1)) - torch.testing.assert_close(actual_k.float(), quant_k.float(), rtol=0, atol=0) - torch.testing.assert_close(actual_v.float(), quant_v.float(), rtol=0, atol=0) torch.testing.assert_close( - actual_scales[:, slots], scales.permute(2, 0, 1).reshape(heads, -1), rtol=0, atol=0 + actual_scales[:, slots], scales.permute(2, 0, 1).reshape(heads, -1), rtol=1e-6, atol=0 ) + _assert_quantized_close(actual_k, k_ref / scales.unsqueeze(-1), quant_dtype) + _assert_quantized_close(actual_v, v.float() / v_scale[:, None], torch.float8_e4m3fn) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("quant_dtype", [torch.int8, torch.float8_e4m3fn]) -@pytest.mark.parametrize("shape", [(2, 9, 3, 64), (3, 5, 2, 7), (1, 129, 2, 128)]) +@pytest.mark.parametrize("shape", [(2, 9, 3, 64), (3, 5, 2, 128), (1, 129, 2, 128)]) # A 100x smaller V scale saturates the quantized V means at the E4M3 maximum. @pytest.mark.parametrize("v_scale_divisor", [1.0, 100.0]) +@pytest.mark.parametrize("smooth_k", [False, True]) def test_quantize_sage_summaries_matches_reference( - shape: tuple[int, ...], dtype: torch.dtype, quant_dtype: torch.dtype, v_scale_divisor: float + shape: tuple[int, ...], + dtype: torch.dtype, + quant_dtype: torch.dtype, + v_scale_divisor: float, + smooth_k: bool, ) -> None: - k, v, v_scale = _inputs(shape, dtype) + k, v, v_scale, k_mean = _summary_inputs(shape, dtype, smooth_k) v_scale = v_scale / v_scale_divisor - actual = quantize_sage_summaries(k, v, v_scale, quant_dtype) + actual = _quantize_summaries(k, v, v_scale, k_mean, quant_dtype) assert [result.dtype for result in actual] == [quant_dtype, torch.float8_e4m3fn, torch.float32] - _assert_matches_reference(actual, k, v, v_scale, quant_dtype) + _assert_summaries_match_reference(actual, k, v, v_scale, k_mean, quant_dtype) @pytest.mark.parametrize("quant_dtype", [torch.int8, torch.float8_e4m3fn]) def test_quantize_sage_summaries_zero_blocks_use_scale_floor(quant_dtype: torch.dtype) -> None: k = torch.zeros(2, 9, 2, 64, device="cuda", dtype=torch.bfloat16) v = torch.zeros_like(k) - v_scale = torch.full((2, 64), 1e-3 / 448.0, device="cuda") - actual = quantize_sage_summaries(k, v, v_scale, quant_dtype) - _assert_matches_reference(actual, k, v, v_scale, quant_dtype) + v_scale = torch.full((2, 2, 64), 1e-3 / 448.0, device="cuda") + actual = _quantize_summaries(k, v, v_scale, None, quant_dtype) + _assert_summaries_match_reference(actual, k, v, v_scale, None, quant_dtype) -def test_quantize_sage_summaries_rejects_layouts_the_kernel_cannot_index() -> None: - k, v, v_scale = _inputs((2, 9, 3, 64), torch.bfloat16) +def test_quantize_sage_summaries_rejects_mismatched_statistics() -> None: + k, v, v_scale, _ = _summary_inputs((2, 9, 3, 64), torch.bfloat16) cases = ( - ((k, v[:, :1], v_scale), "matching shape, dtype, and device"), - ((k, v, v_scale[:1]), "FP32 with shape"), - ((k.transpose(0, 1), v.transpose(0, 1), v_scale), "contiguous"), + ((k, v[:1], v_scale, None), "key and value must have the same shape"), + ((k, v, v_scale[:1], None), "v_scale must be"), + ((k, v, v_scale, v_scale[:, :1]), "k_mean must be"), ) for args, message in cases: - with pytest.raises(ValueError, match=message): - quantize_sage_summaries(*args, torch.int8) + with pytest.raises(RuntimeError, match=message): + _quantize_summaries(*args, torch.int8) + stats = SageQuantized(None, None, None, None, None, v_scale) + cu_seqlens = _cu_seqlens([9, 9]) + keys = k.flatten(0, 1) + with pytest.raises(ValueError, match="only K and V are quantized"): + sage_quantize( + keys, + keys, + v.flatten(0, 1), + cu_seqlens, + cu_seqlens, + k_block_size=1, + qk_quant_dtype=torch.int8, + stats=stats, + ) @pytest.mark.parametrize("quant_dtype", [torch.int8, torch.float8_e4m3fn]) def test_quantize_sage_summaries_under_compile_and_cuda_graph(quant_dtype: torch.dtype) -> None: - k, v, v_scale = _inputs((2, 9, 2, 64), torch.bfloat16) + k, v, v_scale, k_mean = _summary_inputs((2, 9, 2, 64), torch.bfloat16, smooth_k=True) def quantize( - key: torch.Tensor, value: torch.Tensor, value_scale: torch.Tensor + key: torch.Tensor, value: torch.Tensor, value_scale: torch.Tensor, key_mean: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - return quantize_sage_summaries(key, value, value_scale, quant_dtype) + return _quantize_summaries(key, value, value_scale, key_mean, quant_dtype) - _assert_matches_reference( - torch.compile(quantize, fullgraph=True)(k, v, v_scale), k, v, v_scale, quant_dtype + _assert_summaries_match_reference( + torch.compile(quantize, fullgraph=True)(k, v, v_scale, k_mean), + k, + v, + v_scale, + k_mean, + quant_dtype, ) stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(stream): - quantize(k, v, v_scale) + quantize(k, v, v_scale, k_mean) torch.cuda.current_stream().wait_stream(stream) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): - captured = quantize(k, v, v_scale) + captured = quantize(k, v, v_scale, k_mean) k.mul_(2) v.mul_(2) graph.replay() - _assert_matches_reference(captured, k, v, v_scale, quant_dtype) + _assert_summaries_match_reference(captured, k, v, v_scale, k_mean, quant_dtype) + + +# --------------------------------------------------------------------------- Q/K/V +def _ragged(lengths: list[int], heads: int, head_dim: int, generator: torch.Generator): + """Token-major BF16 tensor with per-sequence channel offsets, so that K smoothing matters.""" + tokens = sum(lengths) + x = torch.randn(tokens, heads, head_dim, device="cuda", generator=generator) + offsets = 4 * torch.randn(len(lengths), heads, head_dim, device="cuda", generator=generator) + sequence = torch.repeat_interleave( + torch.arange(len(lengths), device="cuda"), torch.tensor(lengths, device="cuda") + ) + return (x + offsets[sequence]).to(torch.bfloat16) + + +def _cu_seqlens(lengths: list[int]) -> torch.Tensor: + return torch.tensor([0, *torch.tensor(lengths).cumsum(0).tolist()], device="cuda").int() + + +def _check_block_quantized( + quant: torch.Tensor, + scale: torch.Tensor, + source: torch.Tensor, + lengths: list[int], + block_size: int, + quant_dtype: torch.dtype, +) -> None: + """Scales equal the reference block maxima; values dequantize within one grid step.""" + divisor = 126.9 if quant_dtype == torch.int8 else 448.0 + batch = len(lengths) + expected_numel = flat_scale_numel(1, source.shape[0], block_size) + batch - 1 + assert scale.shape == (source.shape[1], expected_numel) + begin = 0 + for b, length in enumerate(lengths): + seq = source[begin : begin + length].float() + blocks = -(-length // block_size) + padded = torch.zeros(blocks * block_size, *seq.shape[1:], device="cuda") + padded[:length] = seq + maximum = padded.view(blocks, block_size, *seq.shape[1:]).abs().amax(dim=(1, 3)) + maximum = maximum.clamp_min(1e-3) + expected = maximum / torch.full_like(maximum, divisor) + # Sequence b's scales start at cu_seqlens[b] // block_size + b. + slots = begin // block_size + flat_scale_slot(b, torch.arange(blocks, device="cuda"), 0, 0) + torch.testing.assert_close(scale[:, slots], expected.T, rtol=1e-6, atol=0) + + token_scale = expected.repeat_interleave(block_size, dim=0)[:length].unsqueeze(-1) + _assert_quantized_close(quant[begin : begin + length], seq / token_scale, quant_dtype) + begin += length + + +@pytest.mark.parametrize("block_size", [1, 4, 16, 32]) +@pytest.mark.parametrize("quant_dtype", [torch.int8, torch.float8_e4m3fn]) +@pytest.mark.parametrize("smooth_k", [False, True]) +@pytest.mark.parametrize("head_dim", [64, 128]) +def test_sage_quantize_matches_reference( + block_size: int, quant_dtype: torch.dtype, smooth_k: bool, head_dim: int +) -> None: + generator = torch.Generator(device="cuda").manual_seed(7) + # Ragged sequences with partial final blocks; GQA with 2 query heads per KV head. + lengths = [301, 77, 520] + q = _ragged(lengths, 4, head_dim, generator) + k = _ragged(lengths, 2, head_dim, generator) + v = _ragged(lengths, 2, head_dim, generator) + cu_seqlens = _cu_seqlens(lengths) + token_lengths = torch.tensor(lengths, device="cuda") + + actual = sage_quantize( + q, + k, + v, + cu_seqlens, + cu_seqlens, + q_block_size=block_size, + k_block_size=block_size, + qk_quant_dtype=quant_dtype, + smooth_k=smooth_k, + ) + + if smooth_k: + k_mean = torch.stack([seq.float().mean(dim=0) for seq in torch.split(k, lengths)]) + torch.testing.assert_close(actual.k_mean, k_mean, rtol=1e-5, atol=1e-5) + k_ref = k.float() - torch.repeat_interleave(actual.k_mean, token_lengths, dim=0) + else: + assert actual.k_mean is None + k_ref = k.float() + _check_block_quantized(actual.q, actual.q_scale, q, lengths, block_size, quant_dtype) + _check_block_quantized(actual.k, actual.k_scale, k_ref, lengths, block_size, quant_dtype) + + # V: one FP8 scale per channel of every sequence, from the BF16 maxima. + v_max = torch.stack([seq.abs().amax(dim=0) for seq in torch.split(v, lengths)]) + floor = torch.tensor(1e-3, dtype=torch.bfloat16, device="cuda") + v_scale = torch.maximum(v_max, floor).float() * torch.tensor(1 / 448.0, device="cuda") + torch.testing.assert_close(actual.v_scale, v_scale, rtol=0, atol=0) + token_v_scale = torch.repeat_interleave(v_scale, token_lengths, dim=0) + error = (actual.v.float() * token_v_scale - v.float()).abs() + assert torch.all(error <= v.float().abs() * 0.07 + token_v_scale * 2**-6) + + +def test_sage_quantize_mixed_block_sizes_and_query_lengths() -> None: + generator = torch.Generator(device="cuda").manual_seed(11) + q_lengths, kv_lengths = [64, 200], [256, 90] + q = _ragged(q_lengths, 2, 128, generator) + k = _ragged(kv_lengths, 2, 128, generator) + v = _ragged(kv_lengths, 2, 128, generator) + actual = sage_quantize( + q, + k, + v, + _cu_seqlens(q_lengths), + _cu_seqlens(kv_lengths), + q_block_size=1, + k_block_size=32, + qk_quant_dtype=torch.float8_e4m3fn, + ) + _check_block_quantized(actual.q, actual.q_scale, q, q_lengths, 1, torch.float8_e4m3fn) + _check_block_quantized(actual.k, actual.k_scale, k, kv_lengths, 32, torch.float8_e4m3fn) + assert actual.v_scale.shape == (2, 2, 128) + + +@pytest.mark.parametrize("block_size", [2, 3, 64]) +def test_sage_quantize_rejects_unsupported_block_sizes(block_size: int) -> None: + x = torch.zeros(16, 2, 128, device="cuda", dtype=torch.bfloat16) + cu_seqlens = _cu_seqlens([16]) + with pytest.raises(RuntimeError, match="must be 1, 4, 16, or 32"): + sage_quantize( + x, + x, + x, + cu_seqlens, + cu_seqlens, + q_block_size=block_size, + k_block_size=16, + qk_quant_dtype=torch.int8, + ) diff --git a/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py b/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py index 090385c73dcb..26e387966815 100644 --- a/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py +++ b/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py @@ -23,6 +23,7 @@ from utils.util import getSMVersion from tensorrt_llm._torch.attention.backends import prims_ts +from tensorrt_llm._torch.attention.backends import trtllm as trtllm_backend from tensorrt_llm._torch.attention.backends.fmha import prims_ts_block_sparse as block_sparse_fmha from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams @@ -35,6 +36,7 @@ BlockSparseForwardInputs, SparseRuntimeParams, ) +from tensorrt_llm._torch.attention.kernels.sage_quant import SageQuantized from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm.functional import PositionEmbeddingType @@ -148,7 +150,10 @@ def _contiguous_case(): num_sparse_topk=0, helix_position_offsets=None, kv_cache_manager=None, + num_contexts=2, seq_lens=torch.tensor([64, 64], dtype=torch.int32), + seq_lens_cuda=torch.tensor([64, 64], dtype=torch.int32), + seq_lens_kv_cuda=torch.tensor([256, 256], dtype=torch.int32), fmha_plan_caches={}, ) args = AttentionForwardArgs( @@ -408,6 +413,81 @@ def test_block_sparse_plan_cache_lives_on_the_metadata(monkeypatch) -> None: assert factory.call_count == 2 +def _set_sage_recipe(args: AttentionForwardArgs) -> None: + args.sage_attn_num_elts_per_blk_q = 1 + args.sage_attn_num_elts_per_blk_k = 16 + args.sage_attn_num_elts_per_blk_v = 1 + + +def test_attention_core_quantizes_sage_inputs_ahead_of_the_fmha(monkeypatch) -> None: + attention, fmha, q, k, v, metadata, args = _contiguous_case() + attention._fmha_manager = SimpleNamespace(fmha_libs=[fmha]) + quantized = SageQuantized._make(Mock(name=name) for name in SageQuantized._fields) + sage_quantize = Mock(return_value=quantized) + monkeypatch.setattr(trtllm_backend, "sage_quantize", sage_quantize) + + def quantize(): + return trtllm_backend.TrtllmAttention._quantize_sage_inputs( + attention, q, k, v, metadata, args + ) + + assert quantize() is None + _set_sage_recipe(args) + args.sage_attn_smooth_k = True + assert quantize() is quantized + (query, key, value, cu_seqlens_q, cu_seqlens_kv), kwargs = sage_quantize.call_args + assert query.shape == (128, 2, 128) + assert key.shape == value.shape == (512, 1, 128) + assert cu_seqlens_q.tolist() == [0, 64, 128] + assert cu_seqlens_kv.tolist() == [0, 256, 512] + assert kwargs == { + "q_block_size": 1, + "k_block_size": 16, + "qk_quant_dtype": torch.float8_e4m3fn, + "smooth_k": True, + } + assert sage_quantize.call_count == 1 + + # Proxy summaries are keys of their own, with the V scales and means of the call's own K/V. + proxy = _bitmask_inputs(proxy=True) + _set_block_sparse_inputs(args, proxy) + summary = SageQuantized._make(Mock(name=name) for name in SageQuantized._fields) + sage_quantize.reset_mock() + sage_quantize.side_effect = [quantized, summary] + actual = quantize() + (query, key, value, cu_seqlens_q, cu_seqlens_blocks), kwargs = sage_quantize.call_args + assert query is None and cu_seqlens_q is None + assert key.shape == value.shape == (8, 1, 128) + assert cu_seqlens_blocks.tolist() == [0, 4, 8] + assert kwargs == {"k_block_size": 1, "qk_quant_dtype": torch.float8_e4m3fn, "stats": quantized} + assert actual.k_summary is summary.k.view.return_value + assert actual.v_summary is summary.v.view.return_value + assert actual.k_summary_scale is summary.k_scale + sage_quantize.side_effect = None + + # Dense calls and calls with a KV cache keep the attention op's quantization, and a layer + # without a library that takes quantized inputs quantizes nothing. + sage_quantize.reset_mock() + _set_block_sparse_inputs(args, None) + assert quantize() is None + _set_block_sparse_inputs(args, proxy) + attention._fmha_manager.fmha_libs = [] + assert quantize() is None + attention._fmha_manager.fmha_libs = [fmha] + metadata.kv_cache_manager = Mock() + assert quantize() is None + sage_quantize.assert_not_called() + + +def test_contiguous_sage_requires_inputs_quantized_by_the_attention_core() -> None: + _attention, fmha, q, k, v, metadata, args = _contiguous_case() + _set_sage_recipe(args) + + reason = fmha._contiguous_unsupported_reason(q, k, v, metadata, args) + + assert reason == "SageAttention requires inputs quantized by the attention core" + + def test_paged_wrapper_uses_zero_copy_padded_row_stride_block_tables(monkeypatch) -> None: attention, fmha, q, metadata, args = _paged_case() wrapper = Mock() diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/test_sage_block_sparse.py b/tests/unittest/_torch/visual_gen/sparse_attention/test_sage_block_sparse.py index 287c65257e51..78cca59740f4 100644 --- a/tests/unittest/_torch/visual_gen/sparse_attention/test_sage_block_sparse.py +++ b/tests/unittest/_torch/visual_gen/sparse_attention/test_sage_block_sparse.py @@ -82,6 +82,7 @@ def _backend( sparse_config: SolAttentionConfig | VideoSparseAttentionConfig | None = None, *, k_block_size: int | None = None, + dtype: torch.dtype = torch.bfloat16, ) -> TrtllmAttention: quant_config = None if qk_dtype is not None: @@ -101,7 +102,7 @@ def _backend( layer_idx=1, num_heads=_HEADS, head_dim=_HEAD_DIM, - dtype=torch.bfloat16, + dtype=dtype, attention_config=AttentionConfig( backend="TRTLLM", quant_attention_config=quant_config, @@ -169,7 +170,13 @@ def _sparse_wrappers(backend: TrtllmAttention) -> list[BlockSparseTSWrapper]: ] -def _assert_sparse_sage_active(backend: TrtllmAttention, *, proxy: bool, token_mask: bool) -> None: +def _assert_sparse_sage_active( + backend: TrtllmAttention, + *, + proxy: bool, + token_mask: bool, + dtype: torch.dtype = torch.bfloat16, +) -> None: wrappers = _sparse_wrappers(backend) assert wrappers, "Sage silently bypassed the PrimTS block-sparse FMHA" quant_config = backend.quant_attention_config @@ -182,7 +189,7 @@ def _assert_sparse_sage_active(backend: TrtllmAttention, *, proxy: bool, token_m assert state.use_kv_valid_bits == token_mask assert state.q_dtype == qk_dtype assert state.v_dtype == torch.float8_e4m3fn - assert state.output_dtype == torch.bfloat16 + assert state.output_dtype == dtype assert state.sage.q_block_size == quant_config.q_block_size assert state.sage.k_block_size == quant_config.k_block_size @@ -345,12 +352,12 @@ def test_sol_sage_cuda_graph_refreshes_inputs_and_scales(qk_dtype: _QuantDtype) _assert_parity(replayed, expected, label=f"SOL {qk_dtype} Sage CUDA Graph replay") -def _masked_zero_score_inputs() -> tuple[ - torch.Tensor, torch.Tensor, torch.Tensor, BlockSparseForwardInputs -]: +def _masked_zero_score_inputs( + dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, BlockSparseForwardInputs]: shape = (_BATCH_SIZE, 131, _HEADS, _HEAD_DIM) - magnitudes = torch.tensor([0.0, 1.0e-8], device="cuda", dtype=torch.bfloat16) - q = torch.ones(shape, device="cuda", dtype=torch.bfloat16) * magnitudes[:, None, None, None] + magnitudes = torch.tensor([0.0, 1.0e-8], device="cuda", dtype=dtype) + q = torch.ones(shape, device="cuda", dtype=dtype) * magnitudes[:, None, None, None] tokens = torch.arange(shape[1], device=q.device) visible = (tokens >= 68) & (tokens < 76) # Masked V=4 exposes leaked probability; the eight visible tokens have V=1. @@ -367,12 +374,15 @@ def _masked_zero_score_inputs() -> tuple[ return q, -q, v, routes +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16], ids=["bf16", "fp16"]) @pytest.mark.parametrize("qk_dtype", ["int8", "fp8"]) @torch.no_grad() -def test_sage_zero_scores_preserve_token_masks_and_empty_rows(qk_dtype: _QuantDtype) -> None: +def test_sage_zero_scores_preserve_token_masks_and_empty_rows( + qk_dtype: _QuantDtype, dtype: torch.dtype +) -> None: """Zero and tiny Q/K keep token masking and empty rows exact through quantization.""" - q, k, v, routes = _masked_zero_score_inputs() - backends = (_backend(None), _backend(qk_dtype)) + q, k, v, routes = _masked_zero_score_inputs(dtype) + backends = (_backend(None, dtype=dtype), _backend(qk_dtype, dtype=dtype)) expected, actual = ( backend.forward( q, @@ -384,7 +394,7 @@ def test_sage_zero_scores_preserve_token_masks_and_empty_rows(qk_dtype: _QuantDt ).view_as(q) for backend in backends ) - _assert_sparse_sage_active(backends[1], proxy=False, token_mask=True) + _assert_sparse_sage_active(backends[1], proxy=False, token_mask=True, dtype=dtype) _assert_parity(actual, expected, label=f"{qk_dtype} zero/tiny scores with masked rows") assert torch.count_nonzero(actual[:, :64]).item() == 0, "fully masked rows must be zero" assert torch.count_nonzero(actual[:, 128:]).item() == 0, "empty route rows must be zero" diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_attention.py b/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_attention.py index 0484352b709d..24aa3b32f774 100644 --- a/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_attention.py +++ b/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_attention.py @@ -295,7 +295,7 @@ def test_trtllm_vsa_hands_predicted_routes_to_core_via_sparse_backend_args( attention._fmha_manager = SimpleNamespace( fmha_libs=[object.__new__(vsa_backend.PrimsTSBlockSparseFmha)] ) - monkeypatch.setattr(vsa_backend, "_get_unsupported_primts_reason", lambda *args, sage: None) + monkeypatch.setattr(vsa_backend, "_get_unsupported_primts_reason", lambda *args: None) q = torch.randn(1, 80, 1, 8) with set_vsa_forward_context(_make_vsa_metadata()): From 9a06e12b074da3692396d890b8071c8d257d1397 Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Wed, 14 Oct 2026 10:07:43 +0000 Subject: [PATCH 07/13] [None][feat] Support V smoothing for SageAttention on the PrimTS kernels SageAttention can subtract the per-channel V mean before quantizing V and add it back to the output: every row's attention weights sum to one, so the output is unchanged while the range V has to cover in FP8 narrows. PrimTS takes this mean (v_mean), and the Sage path now supplies it: - Kernel: with a V-mean buffer, the statistics launch of sageQuant.cu collects each channel's V sum and signed extrema, and the quantizing launch derives scales that cover V minus its mean and subtracts the mean. V smoothing is a template parameter of the statistics launch, so the launch without it keeps its code and registers; one dispatch instantiates only the smoothing combinations that exist. trtllm::sage_quantize fills v_mean, and subtracts a given one with reuse_stats, so the V means of proxy routes subtract the mean the kernel adds back. - The attention op cannot add the mean back, so the attention core also quantizes dense contiguous calls with V smoothing, and PrimsTSBlockSparseFmha serves dense calls with core-quantized inputs through a PrimTS dense plan with 128-token Q/KV blocks, which select the faster Q128/KV128 Sage profile. With V smoothing, its Sage plans enable v_mean and take the mean. The attention op serves the remaining calls with V quantized unsmoothed and logs a warning once. - QuantAttentionConfig.smooth_v (default False) enables it for VisualGen on the TRTLLM backend; AttentionForwardArgs.sage_attn_smooth_v carries it. - Tests: quantizer references with V smoothing for Q/K/V and summaries, including -0.0 extrema; core quantization of dense calls with V smoothing; dense-plan selection; and VisualGen parity for mixed SOL routes, SOL dense timesteps, and a dense model. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- cpp/tensorrt_llm/common/sageQuant.cu | 201 +++++++++++++++--- cpp/tensorrt_llm/common/sageQuant.h | 8 + cpp/tensorrt_llm/thop/sageQuantizeOp.cpp | 42 +++- .../features/visualgen-quantized-attention.md | 3 + .../attention/ATTENTION_DEVELOPER_GUIDE.md | 8 +- .../attention/backends/fmha/fallback.py | 6 + .../backends/fmha/prims_ts_block_sparse.py | 82 ++++--- .../_torch/attention/backends/interface.py | 3 + .../_torch/attention/backends/trtllm.py | 16 +- .../_torch/attention/kernels/sage_quant.py | 32 ++- .../_torch/custom_ops/cpp_custom_ops.py | 1 + .../visual_gen/attention/backends/trtllm.py | 1 + tensorrt_llm/visual_gen/args.py | 21 +- .../attention/kernels/test_sage_quant.py | 65 ++++-- .../sparse/test_prims_ts_block_sparse.py | 44 +++- .../test_sage_block_sparse.py | 37 +++- .../_torch/visual_gen/test_visual_gen_args.py | 7 +- 17 files changed, 469 insertions(+), 108 deletions(-) diff --git a/cpp/tensorrt_llm/common/sageQuant.cu b/cpp/tensorrt_llm/common/sageQuant.cu index 48ba3812b5cc..a24a03db610f 100644 --- a/cpp/tensorrt_llm/common/sageQuant.cu +++ b/cpp/tensorrt_llm/common/sageQuant.cu @@ -64,11 +64,41 @@ __device__ inline SageStatRange getSageStatRange( return {headIdx, ptrCuSeqLensV[seqIdx], ptrCuSeqLensV[seqIdx + 1], seqIdx * numHeadsV + headIdx}; } +// Float atomic maximum and minimum through the integer orderings of IEEE-754 bit patterns: floats +// with a clear sign bit order like signed integers, floats with a set sign bit (-0.0 included) +// reversed like unsigned integers. +__device__ inline void sageAtomicMax(float* address, float value) +{ + if (__float_as_int(value) >= 0) + { + atomicMax(reinterpret_cast(address), __float_as_int(value)); + } + else + { + atomicMin(reinterpret_cast(address), __float_as_uint(value)); + } +} + +__device__ inline void sageAtomicMin(float* address, float value) +{ + if (__float_as_int(value) >= 0) + { + atomicMin(reinterpret_cast(address), __float_as_int(value)); + } + else + { + atomicMax(reinterpret_cast(address), __float_as_uint(value)); + } +} + // SageAttention quantization kernel. Tensors are interpreted as column-major [D, H, S], which is // the same physical layout as contiguous PyTorch [S, H, D]. Each launch quantizes Q or K. It can // simultaneously collect V scales (VStage=1) or quantize V using scales collected by an earlier // launch (VStage=2). When SmoothK is enabled, the VStage=1 launch also collects a per-head-per- -// channel K mean and the VStage=2 launch uses it to perform K-smoothing before quantization. +// channel K mean and the VStage=2 launch uses it to perform K-smoothing before quantization. When +// SmoothV is enabled, the VStage=1 launch collects the per-channel V mean and signed extrema instead +// of the V scales; the VStage=2 launch derives the scales from them when ptrVMax is given and +// subtracts ptrVMean from V whenever it is given. // // Which tokens share a scale is a template parameter: Partition is a CuTe layout // ((intra...), (group...)) -> token offset inside a tile (see sagePartition.h). Contiguous blocks @@ -82,10 +112,12 @@ __device__ inline SageStatRange getSageStatRange( // // V scales and the K mean are reduced per sequence when ptrCuSeqLensV is given (row // seqIdx * numHeadsV + headIdx) and over the whole batch otherwise (row headIdx). -template +template __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* ptrCuSeqLensQk, void const* ptrQk, void* ptrQkQuant, float* ptrQkScale, void const* ptrKForMean, float* ptrKMean, int sumSeqLensV, int numHeadsV, - void const* ptrV, void* ptrVQuant, float* ptrVScale, int const* ptrCuSeqLensV) + void const* ptrV, void* ptrVQuant, float* ptrVScale, int const* ptrCuSeqLensV, float* ptrVMean, float* ptrVMax, + float* ptrVMin) { using namespace cute; using namespace cutlass; @@ -94,6 +126,7 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p constexpr int ScalesPerTile = Traits::ScalesPerTile; constexpr int TileTokens = Traits::TileTokens; static_assert(!SmoothK || VStage != 0, "K smoothing requires V staging"); + static_assert(!SmoothV || VStage == 1, "V smoothing only changes the statistics launch"); static_assert(std::is_same_v || std::is_same_v, "Unrecognized target dtype for quantization"); constexpr float TypeMax = cute::is_same_v ? 448.0f : static_cast(126.9f); @@ -349,7 +382,63 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p Tensor gVSeqQuant = gVQuant(_, threadInTokIdx, headIdx, _); Tensor gVSeqScale = gVScale(_, threadInTokIdx); - if constexpr (VStage == 1) + if constexpr (VStage == 1 && SmoothV) + { + int const numWarpsToUse = cutlass::fast_min(numWarps, 256); + int const numToksPerWave = numWarpsToUse * numToksPerWarp; + if (warpId >= numWarpsToUse) + { + return; + } + // Per-channel sum and signed extrema, in FP32. + float rSum[BestVL]; + float rMax[BestVL]; + float rMin[BestVL]; + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < BestVL; ++i) + { + rSum[i] = 0.0f; + rMax[i] = -INFINITY; + rMin[i] = INFINITY; + } + for (; tokIdx < range.tokEnd; tokIdx += numToksPerWave) + { + cute::copy(AutoVectorizingCopy{}, gVSeq(_, tokIdx), rV); + cute::transform(rV_x2, rVCompute_x2, NumericArrayConverter::convert); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < BestVL; ++i) + { + rSum[i] += rVCompute(i); + rMax[i] = ::fmaxf(rMax[i], rVCompute(i)); + rMin[i] = ::fminf(rMin[i], rVCompute(i)); + } + } + for (int delta = threadsPerHead; delta < 32; delta <<= 1) + { + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < BestVL; ++i) + { + rSum[i] += __shfl_xor_sync(0xffffffffu, rSum[i], delta); + rMax[i] = ::fmaxf(rMax[i], __shfl_xor_sync(0xffffffffu, rMax[i], delta)); + rMin[i] = ::fminf(rMin[i], __shfl_xor_sync(0xffffffffu, rMin[i], delta)); + } + } + // A smoothed V's scales cover V minus its mean, which is final only after this launch, so + // the caller derives them from the mean and the extrema. + if (threadInTokIdx == thrId && range.tokBegin < range.tokEnd) + { + int64_t const statOffset = static_cast(range.row) * HeadDim + threadInTokIdx * BestVL; + float const invNumTokens = 1.0f / (range.tokEnd - range.tokBegin); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < BestVL; ++i) + { + atomicAdd(ptrVMean + statOffset + i, rSum[i] * invNumTokens); + sageAtomicMax(ptrVMax + statOffset + i, rMax[i]); + sageAtomicMin(ptrVMin + statOffset + i, rMin[i]); + } + } + } + else if constexpr (VStage == 1) { int const numWarpsToUse = cutlass::fast_min(numWarps, 256); int const numToksPerWave = numWarpsToUse * numToksPerWarp; @@ -389,13 +478,50 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p else if constexpr (VStage == 2) { int const numToksPerWave = numWarps * numToksPerWarp; - cute::copy(AutoVectorizingCopy{}, gVSeqScale, rVScale); + int64_t const statOffset = static_cast(range.row) * HeadDim + threadInTokIdx * BestVL; + // This thread's channel slice of the V mean, zero without V smoothing. + Tensor rVMean = make_tensor(Shape{}); + clear(rVMean); + if (ptrVMean != nullptr) + { + cute::copy( + AutoVectorizingCopy{}, make_tensor(make_gmem_ptr(ptrVMean + statOffset), Shape{}), rVMean); + } + if (ptrVMax != nullptr) + { + // A smoothed V's scales cover its largest distance from the mean, with the floor of the + // unsmoothed scales. One thread per channel slice writes them for the consumer. + Tensor rVLargest = make_tensor(Shape{}); + Tensor rVSmallest = make_tensor(Shape{}); + cute::copy(AutoVectorizingCopy{}, make_tensor(make_gmem_ptr(ptrVMax + statOffset), Shape{}), + rVLargest); + cute::copy(AutoVectorizingCopy{}, make_tensor(make_gmem_ptr(ptrVMin + statOffset), Shape{}), + rVSmallest); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < BestVL; ++i) + { + rVScale(i) + = ::fmaxf(::fmaxf(rVLargest(i) - rVMean(i), rVMean(i) - rVSmallest(i)), 1e-3f) / 448.0f; + } + if (warpId == 0 && threadInTokIdx == thrId) + { + cute::copy(AutoVectorizingCopy{}, rVScale, gVSeqScale); + } + } + else + { + cute::copy(AutoVectorizingCopy{}, gVSeqScale, rVScale); + } cute::transform(rVScale_x2, rVScale_x2, cutlass::reciprocal_approximate>{}); for (; tokIdx < range.tokEnd; tokIdx += numToksPerWave) { cute::copy(AutoVectorizingCopy{}, gVSeq(_, tokIdx), rV); cute::transform(rV_x2, rVCompute_x2, cutlass::NumericArrayConverter::convert); - cute::transform(rVCompute_x2, rVScale_x2, rVCompute_x2, cutlass::multiplies>{}); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < BestVL; ++i) + { + rVCompute(i) = (rVCompute(i) - rVMean(i)) * rVScale(i); + } cute::transform( rVCompute_x4, rVQuant_x4, cutlass::NumericArrayConverter::convert); cute::copy(AutoVectorizingCopy{}, rVQuant, gVSeqQuant(_, tokIdx)); @@ -490,6 +616,9 @@ void invokeSageQuantQkvImpl(SageQuantParams const& params) "Invalid SageQuant K-smoothing parameters"); TLLM_CHECK_WITH_INFO( params.ptrCuSeqLensV == nullptr || params.vStage != 0, "SageQuant per-sequence V statistics require V staging"); + TLLM_CHECK_WITH_INFO(params.ptrVMean == nullptr + || (params.vStage != 0 && (params.vStage != 1 || (params.ptrVMax != nullptr && params.ptrVMin != nullptr))), + "Invalid SageQuant V-smoothing parameters"); auto invokeKernel = [&](auto headDimStatic, auto partitionStatic) { @@ -499,22 +628,24 @@ void invokeSageQuantQkvImpl(SageQuantParams const& params) void* kernelArgs[] = {&kernelParams.sumSeqLensQk, &kernelParams.batchSize, &kernelParams.ptrCuSeqLensQk, &kernelParams.ptrQk, &kernelParams.ptrQkQuant, &kernelParams.ptrQkScale, &kernelParams.ptrKForMean, &kernelParams.ptrKMean, &kernelParams.sumSeqLensV, &kernelParams.numHeadsV, &kernelParams.ptrV, - &kernelParams.ptrVQuant, &kernelParams.ptrVScale, &kernelParams.ptrCuSeqLensV}; + &kernelParams.ptrVQuant, &kernelParams.ptrVScale, &kernelParams.ptrCuSeqLensV, &kernelParams.ptrVMean, + &kernelParams.ptrVMax, &kernelParams.ptrVMin}; - auto launchKernel = [&](auto smoothKStatic, auto vStageStatic) + auto launchKernel = [&](auto smoothKStatic, auto smoothVStatic, auto vStageStatic) { constexpr bool SmoothK_ = decltype(smoothKStatic)::value; + constexpr bool SmoothV_ = decltype(smoothVStatic)::value; constexpr int VStage_ = vStageStatic; void const* kernelFunc = nullptr; if (params.quantType == DATA_TYPE_E4M3) { - kernelFunc = reinterpret_cast( - sageQuantQkvKernel); + kernelFunc = reinterpret_cast(sageQuantQkvKernel); } else if (params.quantType == DATA_TYPE_INT8) { kernelFunc = reinterpret_cast( - sageQuantQkvKernel); + sageQuantQkvKernel); } else { @@ -530,29 +661,39 @@ void invokeSageQuantQkvImpl(SageQuantParams const& params) TLLM_CHECK_WITH_INFO(status == cudaSuccess, "%s", cudaGetErrorString(status)); }; - switch (params.vStage) + // K smoothing needs V staging, and V smoothing only changes what the VStage==1 launch + // collects, so only those combinations are instantiated. + auto launchWithSmoothKV = [&](auto vStageStatic) { - case 0: launchKernel(cute::false_type{}, Int<0>{}); return; - case 1: - if (params.kSmooth) - { - launchKernel(cute::true_type{}, Int<1>{}); - } - else - { - launchKernel(cute::false_type{}, Int<1>{}); - } - return; - case 2: - if (params.kSmooth) + constexpr int VStage_ = vStageStatic; + auto launchWithSmoothV = [&](auto smoothKStatic) { - launchKernel(cute::true_type{}, Int<2>{}); - } - else + if constexpr (VStage_ == 1) + { + if (params.ptrVMean != nullptr) + { + launchKernel(smoothKStatic, cute::true_type{}, vStageStatic); + return; + } + } + launchKernel(smoothKStatic, cute::false_type{}, vStageStatic); + }; + if constexpr (VStage_ != 0) { - launchKernel(cute::false_type{}, Int<2>{}); + if (params.kSmooth) + { + launchWithSmoothV(cute::true_type{}); + return; + } } - return; + launchWithSmoothV(cute::false_type{}); + }; + + switch (params.vStage) + { + case 0: launchWithSmoothKV(Int<0>{}); return; + case 1: launchWithSmoothKV(Int<1>{}); return; + case 2: launchWithSmoothKV(Int<2>{}); return; default: TLLM_THROW("Unsupported SageQuantV stage %d", params.vStage); } }; diff --git a/cpp/tensorrt_llm/common/sageQuant.h b/cpp/tensorrt_llm/common/sageQuant.h index 1c0691243f04..475a508c5b77 100644 --- a/cpp/tensorrt_llm/common/sageQuant.h +++ b/cpp/tensorrt_llm/common/sageQuant.h @@ -74,6 +74,14 @@ struct SageQuantParams // PrimTS block-sparse kernel expect; otherwise one [numHeadsV, headDim] row set covers the // whole batch, as the trtllm-gen kernels expect. int const* ptrCuSeqLensV{nullptr}; + // Optional V smoothing, with the layout of the V scales. ptrVMean is collected at vStage==1 and + // subtracted from V at vStage==2. A smoothed V's scales cover V minus its mean, which is final + // only after the whole vStage==1 launch, so that launch collects the signed per-channel extrema + // into ptrVMax/ptrVMin (filled with -inf/+inf beforehand) instead of writing ptrVScale, and the + // vStage==2 launch derives ptrVScale from them. Without ptrVMax, vStage==2 reads ptrVScale. + float* ptrVMean{nullptr}; + float* ptrVMax{nullptr}; + float* ptrVMin{nullptr}; // Hardware info. Required. int smCount{}; cudaStream_t stream{}; diff --git a/cpp/tensorrt_llm/thop/sageQuantizeOp.cpp b/cpp/tensorrt_llm/thop/sageQuantizeOp.cpp index df75ed4e5843..6e52efd0e02c 100644 --- a/cpp/tensorrt_llm/thop/sageQuantizeOp.cpp +++ b/cpp/tensorrt_llm/thop/sageQuantizeOp.cpp @@ -23,6 +23,7 @@ #include #include +#include #include TRTLLM_NAMESPACE_BEGIN @@ -131,17 +132,19 @@ void launch(SageQuantParams params, torch::Tensor const& input, torch::Tensor co // // Q and K are quantized to INT8 or FP8 E4M3 with one scale per q_block_size / k_block_size tokens // of a sequence, laid out per head as tensorrt_llm::common::getSageScaleHeadStride describes. V is -// quantized to FP8 E4M3 with one scale per channel of each sequence. The V scales and the K mean -// are [batch, heads_kv, head_dim]; k_mean is the K-smoothing mean and is filled exactly when given. +// quantized to FP8 E4M3 with one scale per channel of each sequence. The V scales and the means +// are [batch, heads_kv, head_dim]; k_mean and v_mean are the K- and V-smoothing means, subtracted +// before quantization and filled exactly when given; a smoothed V's scales cover V minus its mean. // -// With reuse_stats, only K and V are quantized, with the V scales and the K mean of an earlier call +// With reuse_stats, only K and V are quantized, with the V scales and the means of an earlier call // as inputs, for tensors that must dequantize like that call's K and V, such as the K/V block means // of PrimTS proxy routes; K still takes its own scales. The query arguments must then be absent. void sage_quantize(std::optional const& query, torch::Tensor const& key, torch::Tensor const& value, std::optional const& cu_seqlens_q, torch::Tensor const& cu_seqlens_kv, std::optional q_quant, torch::Tensor& k_quant, torch::Tensor& v_quant, std::optional q_scale, torch::Tensor& k_scale, torch::Tensor& v_scale, - std::optional k_mean, int64_t q_block_size, int64_t k_block_size, bool reuse_stats) + std::optional k_mean, std::optional v_mean, int64_t q_block_size, + int64_t k_block_size, bool reuse_stats) { checkKeyValue(key, value, k_quant, v_quant, cu_seqlens_kv); int const batchSize = static_cast(cu_seqlens_kv.size(0) - 1); @@ -152,10 +155,15 @@ void sage_quantize(std::optional const& query, torch::Tensor cons { checkScale(*k_mean, statRows * key.size(2), "k_mean"); } + if (v_mean.has_value()) + { + checkScale(*v_mean, statRows * value.size(2), "v_mean"); + } SageQuantParams params = makeParams(value, v_quant, v_scale, k_quant, batchSize, cu_seqlens_kv.data_ptr()); params.kSmooth = k_mean.has_value(); params.ptrKMean = params.kSmooth ? k_mean->data_ptr() : nullptr; + params.ptrVMean = v_mean.has_value() ? v_mean->data_ptr() : nullptr; if (reuse_stats) { TORCH_CHECK(!query.has_value() && !cu_seqlens_q.has_value() && !q_quant.has_value() && !q_scale.has_value(), @@ -176,15 +184,32 @@ void sage_quantize(std::optional const& query, torch::Tensor cons TORCH_CHECK(q_quant->scalar_type() == k_quant.scalar_type(), "q_quant and k_quant must share a dtype"); checkBlockScales(*query, *q_scale, q_block_size, batchSize, "q_scale"); - // The V scales and the K mean are reduced with atomics into zeroed buffers. - v_scale.zero_(); + // The V scales and the means are reduced with atomics into zeroed buffers. A smoothed V's scales + // are derived from the extrema instead. + if (!v_mean.has_value()) + { + v_scale.zero_(); + } if (k_mean.has_value()) { k_mean->zero_(); } + if (v_mean.has_value()) + { + v_mean->zero_(); + } params.ptrKForMean = params.kSmooth ? key.data_ptr() : nullptr; - // Q, while collecting the V scales and the K mean; then K (smoothed by the K mean) and V. + std::optional v_max; + std::optional v_min; + if (v_mean.has_value()) + { + v_max = torch::full_like(v_scale, -std::numeric_limits::infinity()); + v_min = torch::full_like(v_scale, std::numeric_limits::infinity()); + params.ptrVMax = v_max->data_ptr(); + params.ptrVMin = v_min->data_ptr(); + } + // Q, while collecting the V statistics and the K mean; then K (smoothed by the K mean) and V. launch(params, *query, *cu_seqlens_q, *q_quant, *q_scale, q_block_size, 1); launch(params, key, cu_seqlens_kv, k_quant, k_scale, k_block_size, 2); } @@ -194,7 +219,8 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m) m.def( "sage_quantize(Tensor? query, Tensor key, Tensor value, Tensor? cu_seqlens_q, Tensor cu_seqlens_kv, " "Tensor(a!)? q_quant, Tensor(b!) k_quant, Tensor(c!) v_quant, Tensor(d!)? q_scale, Tensor(e!) k_scale, " - "Tensor(f!) v_scale, Tensor(g!)? k_mean, int q_block_size, int k_block_size, bool reuse_stats=False) -> ()"); + "Tensor(f!) v_scale, Tensor(g!)? k_mean, Tensor(h!)? v_mean, int q_block_size, int k_block_size, " + "bool reuse_stats=False) -> ()"); } TORCH_LIBRARY_IMPL(trtllm, CUDA, m) diff --git a/docs/source/features/visualgen-quantized-attention.md b/docs/source/features/visualgen-quantized-attention.md index c6d3e7aa8e23..78663f54bf4e 100644 --- a/docs/source/features/visualgen-quantized-attention.md +++ b/docs/source/features/visualgen-quantized-attention.md @@ -66,6 +66,8 @@ After a recipe meets the quality target, benchmark its end-to-end throughput wit | `q_block_size` | int ≥ 0 | `0` | Q tokens per SageAttention quantization block. `0` outside SageAttention. | | `k_block_size` | int ≥ 0 | `0` | K tokens per SageAttention quantization block. `0` outside SageAttention. | | `v_block_size` | int ≥ 0 | `0` | V block size on the hidden dimension. `0` = one tensor-wide V scale; `1` = one scale per channel. Keep `0` if `v_dtype` defines its own scaling format (e.g., `mxfp8` or `nvfp4`) | +| `smooth_k` | bool | `False` | SageAttention only. Subtract the per-channel K mean before quantizing K. | +| `smooth_v` | bool | `False` | SageAttention only. Subtract the per-channel V mean before quantizing V and add it back to the output. Applied where attention runs on the PrimTS kernels; see [SageAttention in the `TRTLLM` backend](#sageattention-in-the-trtllm-backend). | Routing (`tensorrt_llm/_torch/visual_gen/attention/backends/utils.py`) forwards the validated `quant_attention_config` into the backend constructor: `TrtllmAttention` for `TRTLLM`, `CuDNNAttention` for `CUDNN`, `FlashInferAttention` for `FLASHINFER`, and the dense `CuTeDSLAttention` FMHA backend for `CUTEDSL`. @@ -118,6 +120,7 @@ attention_config: - On H100/H200/H800/H20, only one recipe with fixed `qk_dtype: "int8"` is supported. - On B200/GB200, it is recommended to use the `qk_dtype: "int8"` recipes. - On B300/GB300, use `qk_dtype: "fp8"` and evaluate output quality because it can be less accurate than the INT8 Q/K recipe on B200/GB200. +- `smooth_v` needs a kernel that adds the V mean back to the output. The PrimTS kernels behind VSA and SOL do; with `smooth_v`, dense attention calls (dense layers and timesteps, or a dense model) also run on PrimTS where their shape allows. The trtllm-gen kernels quantize V without the mean and log a warning once. **Configuration.** diff --git a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index 2e9055d21444..03f8c537456d 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md @@ -244,7 +244,13 @@ routes get one scale per K mean and reuse the original V channel scales and K mean. `PrimsTSBlockSparseFmha` serves Sage only with these inputs and passes all scales to PrimTS; the attention op ignores them and quantizes dense SageAttention inside its kernels. The recipe and quantized dtypes are part of -the plan key. Paged block-sparse attention rejects Sage. +the plan key. Paged block-sparse attention rejects Sage. With +`sage_attn_smooth_v`, V and its summaries also subtract the per-sequence V +mean, which PrimTS adds back to the output. The attention op cannot, so the +core also quantizes dense contiguous calls with V smoothing, and +`PrimsTSBlockSparseFmha` serves dense calls with core-quantized inputs through +a dense plan; the attention op serves the rest without V smoothing and logs a +warning once. `PrimsTSBlockSparseFmha` keeps its planned wrappers in the attention metadata (`TrtllmAttentionMetadata.fmha_plan_caches`), like the FlashInfer wrappers and diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py index 2c344c28680a..0831ace484f8 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py @@ -297,6 +297,12 @@ def forward( forward_args: AttentionForwardArgs, ) -> None: attn = self.attn + if forward_args.sage_attn_smooth_v: + logger.warning_once( + "The attention op quantizes V without the V smoothing this SageAttention " + "recipe requests.", + key="fallback_fmha_sage_attn_smooth_v", + ) # Every kwarg sources from ``attn`` / ``metadata`` / ``forward_args`` # (with ``forward_args.sparse_runtime_params`` for sparse inputs), diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py index 6494ebc9b434..4aa2760ef406 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py @@ -83,7 +83,7 @@ class _BlockSparsePlanKey: page_size: int | None q_block_size: int kv_block_size: int - max_blocks_per_row: int + max_blocks_per_row: int | None mask_type: Literal["dense", "causal"] dtype: torch.dtype qk_dtype: torch.dtype @@ -92,8 +92,11 @@ class _BlockSparsePlanKey: sparse_format: Literal["bsr", "bitmask"] use_proxy_routes: bool sage: SageAttentionConfig | None + use_block_sparse: bool = True def unsupported_reason(self) -> str | None: + # A dense plan has no route capacity. + capacity = {"max_blocks_per_row": self.max_blocks_per_row} if self.use_block_sparse else {} try: _validate_block_sparse_static_profile( batch_size=self.batch_size, @@ -110,9 +113,10 @@ def unsupported_reason(self) -> str | None: kv_dtype=self.qk_dtype, v_dtype=self.v_dtype, output_dtype=self.dtype, - max_blocks_per_row=self.max_blocks_per_row, page_size=self.page_size, sage=self.sage, + use_block_sparse=self.use_block_sparse, + **capacity, ) except (ValueError, NotImplementedError, OverflowError) as error: return str(error) @@ -146,6 +150,7 @@ def plan(self) -> "BlockSparseTSWrapper | BlockSparsePagedTSWrapper": plan_kwargs.update(kv_data_type=self.qk_dtype) else: plan_kwargs.update( + use_block_sparse=self.use_block_sparse, sparse_format=self.sparse_format, use_proxy_routes=self.use_proxy_routes, k_data_type=self.qk_dtype, @@ -174,7 +179,12 @@ def _has_other_sparse_runtime(forward_args: AttentionForwardArgs) -> bool: return False -def _route_batch_size(inputs: BlockSparseForwardInputs) -> int: +def _contiguous_batch_size( + inputs: BlockSparseForwardInputs | None, metadata: "TrtllmAttentionMetadata" +) -> int: + """Return the route batch of block-sparse inputs, or the context batch of a dense call.""" + if inputs is None: + return int(metadata.num_contexts) routes = inputs.block_indptr if inputs.sparse_format == "bsr" else inputs.exact_block_bits return int(routes.shape[0]) @@ -207,13 +217,18 @@ def _sage_recipe( config = SageAttentionConfig( q_block_size=forward_args.sage_attn_num_elts_per_blk_q, k_block_size=forward_args.sage_attn_num_elts_per_blk_k, + v_mean=forward_args.sage_attn_smooth_v, k_summary_block_size=1, ) return config, qk_dtype class PrimsTSBlockSparseFmha(PrimsTSFmha): - """Contiguous context and fixed-Q paged generation block-sparse FMHA.""" + """Contiguous context and fixed-Q paged generation block-sparse FMHA. + + Contiguous dense calls whose SageAttention inputs the attention core quantized also run + here, on a dense plan. + """ PLAN_CACHE_KEY = "prims_ts_block_sparse" @@ -274,7 +289,7 @@ def _common_unsupported_reason( forward_args: AttentionForwardArgs, ) -> str | None: """Gates shared by the contiguous and paged block-sparse paths.""" - if _get_block_sparse_inputs(forward_args) is None: + if _get_block_sparse_inputs(forward_args) is None and forward_args.sage_quantized is None: return "block-sparse forward inputs are required" if metadata.is_cross: return "cross attention is not supported" @@ -303,7 +318,7 @@ def _common_unsupported_reason( def _make_plan_key( self, q: torch.Tensor, - inputs: BlockSparseForwardInputs, + inputs: BlockSparseForwardInputs | None, *, batch_size: int, seq_len_q: int, @@ -312,9 +327,16 @@ def _make_plan_key( mask_type: Literal["dense", "causal"], sage: tuple[SageAttentionConfig, torch.dtype] | None = None, ) -> _BlockSparsePlanKey: - max_blocks_per_row = inputs.max_blocks_per_row - if max_blocks_per_row is None: - max_blocks_per_row = math.ceil(kv_capacity / inputs.kv_block_size) + if inputs is None: + # Dense plans use 128-token blocks, which select the Q128/KV128 Sage profile, faster + # for dense attention than Q64/KV256. + q_block_size, kv_block_size = 128, 128 + max_blocks_per_row = None + else: + q_block_size, kv_block_size = inputs.q_block_size, inputs.kv_block_size + max_blocks_per_row = inputs.max_blocks_per_row + if max_blocks_per_row is None: + max_blocks_per_row = math.ceil(kv_capacity / kv_block_size) sage_config, qk_dtype, v_dtype = None, q.dtype, q.dtype if sage is not None: sage_config, qk_dtype = sage @@ -328,17 +350,18 @@ def _make_plan_key( num_kv_heads=self.attn.num_kv_heads, head_dim=self.attn.head_dim, page_size=page_size, - q_block_size=inputs.q_block_size, - kv_block_size=inputs.kv_block_size, + q_block_size=q_block_size, + kv_block_size=kv_block_size, max_blocks_per_row=max_blocks_per_row, mask_type=mask_type, dtype=q.dtype, qk_dtype=qk_dtype, v_dtype=v_dtype, - use_kv_valid_bits=inputs.kv_valid_bits is not None, - sparse_format=inputs.sparse_format, - use_proxy_routes=inputs.use_proxy_routes, + use_kv_valid_bits=inputs is not None and inputs.kv_valid_bits is not None, + sparse_format="bsr" if inputs is None else inputs.sparse_format, + use_proxy_routes=inputs is not None and inputs.use_proxy_routes, sage=sage_config, + use_block_sparse=inputs is not None, ) def _get_or_plan_wrapper( @@ -376,14 +399,13 @@ def _contiguous_unsupported_reason( if forward_args.cu_q_seqlens is not None or forward_args.cu_kv_seqlens is not None: return "packed variable-length Q/KV inputs are not supported" inputs = _get_block_sparse_inputs(forward_args) - assert inputs is not None mask_type = self._get_prims_mask_type(forward_args) - if inputs.use_proxy_routes and mask_type != "dense": + if inputs is not None and inputs.use_proxy_routes and mask_type != "dense": return "block-sparse proxy routes require mask_type='dense'" - batch_size = _route_batch_size(inputs) + batch_size = _contiguous_batch_size(inputs, metadata) seq_len_q = _uniform_seq_len_q(q, metadata, batch_size) if seq_len_q is None or k.shape[0] % batch_size: - return "query and KV token counts must be batch-uniform over the route batch size" + return "query and KV token counts must be batch-uniform over the batch size" key = self._make_plan_key( q, inputs, @@ -563,9 +585,8 @@ def _forward_contiguous( forward_args: AttentionForwardArgs, ) -> None: inputs = _get_block_sparse_inputs(forward_args) - assert inputs is not None assert forward_args.output is not None - batch_size = _route_batch_size(inputs) + batch_size = _contiguous_batch_size(inputs, metadata) query = q.view(batch_size, -1, self.attn.num_heads, self.attn.head_dim) key_states = k.view(batch_size, -1, self.attn.num_kv_heads, self.attn.head_dim) value_states = v.view_as(key_states) @@ -587,7 +608,7 @@ def _forward_contiguous( query = quantized.q.view_as(query) key_states = quantized.k.view_as(key_states) value_states = quantized.v.view_as(value_states) - if inputs.use_proxy_routes: + if inputs is not None and inputs.use_proxy_routes: inputs = replace( inputs, k_summary=quantized.k_summary, v_summary=quantized.v_summary ) @@ -596,20 +617,27 @@ def _forward_contiguous( k_scale=quantized.k_scale, v_scale=quantized.v_scale, k_summary_scale=quantized.k_summary_scale, + v_mean=quantized.v_mean, ) + # A dense plan takes no routes. + routes = {} + if inputs is not None: + routes = { + "block_indptr": inputs.block_indptr, + "block_indices": inputs.block_indices, + "exact_block_bits": inputs.exact_block_bits, + "k_summary": inputs.k_summary, + "v_summary": inputs.v_summary, + "kv_valid_bits": inputs.kv_valid_bits, + } wrapper.run( query, key_states, value_states, - block_indptr=inputs.block_indptr, - block_indices=inputs.block_indices, - exact_block_bits=inputs.exact_block_bits, - k_summary=inputs.k_summary, - v_summary=inputs.v_summary, - kv_valid_bits=inputs.kv_valid_bits, sm_scale=self._get_bmm1_scale(self.attn), out=forward_args.output.view_as(query), sage=sage_params, + **routes, ) # Planning records an external CUDA event that Dynamo cannot construct, so diff --git a/tensorrt_llm/_torch/attention/backends/interface.py b/tensorrt_llm/_torch/attention/backends/interface.py index f9a5e631bfe4..7558793e0ef0 100644 --- a/tensorrt_llm/_torch/attention/backends/interface.py +++ b/tensorrt_llm/_torch/attention/backends/interface.py @@ -969,6 +969,9 @@ class AttentionForwardArgs: sage_attn_num_elts_per_blk_v: int = 0 sage_attn_qk_int8: bool = False sage_attn_smooth_k: bool = False + # Subtract the per-channel V mean before quantizing V. Applied by FMHA libraries that add + # the mean back to the output; the others quantize V unsmoothed. + sage_attn_smooth_v: bool = False # SageAttention inputs that the attention core quantized for the FMHA libraries # taking quantized inputs; the attention op quantizes inside its kernels. sage_quantized: Optional["SageQuantized"] = None diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index dd69011f255e..79dd50b31676 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -2023,11 +2023,12 @@ def _quantize_sage_inputs( metadata: TrtllmAttentionMetadata, forward_args: AttentionForwardArgs, ) -> Optional[SageQuantized]: - """Quantize the SageAttention inputs of a block-sparse call. + """Quantize the Sage inputs of block-sparse and V-smoothing calls. - The attention op quantizes dense SageAttention inside its kernels, so - only block-sparse calls are quantized here, for the FMHA libraries that - take quantized inputs: separate Q, K, and V without a KV cache. + The attention op quantizes dense SageAttention inside its kernels and + cannot add a V mean back, so only block-sparse calls and calls with V + smoothing are quantized here, for the FMHA libraries that take + quantized inputs: separate Q, K, and V without a KV cache. """ qk_quant_dtype = forward_args.sage_attn_qk_dtype if (qk_quant_dtype is None or k is None or v is None @@ -2036,7 +2037,7 @@ def _quantize_sage_inputs( for fmha in self._fmha_manager.fmha_libs)): return None inputs = forward_args.sparse_runtime_params.block_sparse_inputs - if inputs is None: + if inputs is None and not forward_args.sage_attn_smooth_v: return None num_seqs = metadata.num_contexts # Device lengths add no host-to-device copy during CUDA Graph capture. @@ -2058,12 +2059,13 @@ def _quantize_sage_inputs( k_block_size=forward_args.sage_attn_num_elts_per_blk_k, qk_quant_dtype=qk_quant_dtype, smooth_k=forward_args.sage_attn_smooth_k, + smooth_v=forward_args.sage_attn_smooth_v, ) - if not inputs.use_proxy_routes: + if inputs is None or not inputs.use_proxy_routes: return quantized # The K/V block means of proxy routes dequantize like the keys and # values they stand for: every mean is a key of its sequence with a - # scale of its own, and reuses the V scales and K mean of this call. + # scale of its own, and reuses the V scales and means of this call. batch, blocks = inputs.k_summary.shape[:2] cu_seqlens_blocks = torch.arange(0, (batch + 1) * blocks, blocks, diff --git a/tensorrt_llm/_torch/attention/kernels/sage_quant.py b/tensorrt_llm/_torch/attention/kernels/sage_quant.py index f76a80457951..c3fb9f6a9033 100644 --- a/tensorrt_llm/_torch/attention/kernels/sage_quant.py +++ b/tensorrt_llm/_torch/attention/kernels/sage_quant.py @@ -33,11 +33,12 @@ class SageQuantized(NamedTuple): """Quantized Q/K/V and their FP32 dequantization scales. ``q_scale`` and ``k_scale`` are ``[heads, ceil(tokens / block) + batch - 1]``: sequence - ``b`` starts at slot ``cu_seqlens[b] // block + b``. ``v_scale`` and ``k_mean`` are - ``[batch, heads_kv, head_dim]``; ``k_mean`` is the mean subtracted from K before - quantization and is ``None`` without K smoothing. ``q`` and ``q_scale`` are ``None`` when - the statistics of an earlier call were reused. The summary fields hold the quantized K/V - block means of proxy routes and their per-mean key scales, for calls that have them. + ``b`` starts at slot ``cu_seqlens[b] // block + b``. ``v_scale``, ``k_mean``, and + ``v_mean`` are ``[batch, heads_kv, head_dim]``; the means are subtracted from K and V + before quantization and are ``None`` without K or V smoothing. ``q`` and ``q_scale`` are + ``None`` when the statistics of an earlier call were reused. The summary fields hold the + quantized K/V block means of proxy routes and their per-mean key scales, for calls that + have them. """ q: torch.Tensor | None @@ -47,6 +48,7 @@ class SageQuantized(NamedTuple): k_scale: torch.Tensor v_scale: torch.Tensor k_mean: torch.Tensor | None = None + v_mean: torch.Tensor | None = None k_summary: torch.Tensor | None = None v_summary: torch.Tensor | None = None k_summary_scale: torch.Tensor | None = None @@ -72,17 +74,19 @@ def sage_quantize( k_block_size: int, qk_quant_dtype: torch.dtype, smooth_k: bool = False, + smooth_v: bool = False, stats: SageQuantized | None = None, ) -> SageQuantized: - """Quantize Q/K/V for SageAttention, with V scales and the K mean per sequence. + """Quantize Q/K/V for SageAttention, with V scales and the means per sequence. Q and K get one scale per ``q_block_size`` / ``k_block_size`` tokens of a sequence (1, 4, 16, or 32); V gets one FP8 scale per channel of each sequence. With ``smooth_k``, every sequence's per-channel K mean is subtracted before K is quantized, - which leaves the attention output unchanged. + which leaves the attention output unchanged. With ``smooth_v``, the per-channel V mean + is subtracted before V is quantized; the attention kernel must add it back. With ``stats``, the result of an earlier call, only K and V are quantized, with the V - scales and the K mean of ``stats``, so that they dequantize like its K and V; K still takes + scales and the means of ``stats``, so that they dequantize like its K and V; K still takes its own scales. The smoothing follows ``stats``, and the result carries its statistics and no Q. @@ -97,11 +101,14 @@ def sage_quantize( k_block_size: Key tokens per scale. qk_quant_dtype: ``torch.int8`` or ``torch.float8_e4m3fn`` for Q and K. smooth_k: Subtract the per-sequence K mean before quantizing K. + smooth_v: Subtract the per-sequence V mean before quantizing V. stats: ``None``, or an earlier result with the batch, heads, and head dimension of - ``key``, whose V scales and K mean to reuse. + ``key``, whose V scales and means to reuse. """ _check_qk_quant_dtype(qk_quant_dtype) - if stats is not None and (query is not None or cu_seqlens_q is not None or smooth_k): + if stats is not None and ( + query is not None or cu_seqlens_q is not None or smooth_k or smooth_v + ): raise ValueError("with stats, only K and V are quantized, with the smoothing of stats") batch = cu_seqlens_kv.numel() - 1 heads_kv, head_dim = key.shape[1], key.shape[2] @@ -126,6 +133,7 @@ def sage_quantize( k_scale, stats.v_scale, stats.k_mean, + stats.v_mean, 0, k_block_size, reuse_stats=True, @@ -138,6 +146,7 @@ def sage_quantize( k_scale=k_scale, v_scale=stats.v_scale, k_mean=stats.k_mean, + v_mean=stats.v_mean, ) q = torch.empty_like(query, dtype=qk_quant_dtype) @@ -148,6 +157,7 @@ def sage_quantize( ) v_scale = torch.empty((batch, heads_kv, head_dim), dtype=torch.float32, device=value.device) k_mean = torch.empty_like(v_scale) if smooth_k else None + v_mean = torch.empty_like(v_scale) if smooth_v else None torch.ops.trtllm.sage_quantize( query, key, @@ -161,6 +171,7 @@ def sage_quantize( k_scale, v_scale, k_mean, + v_mean, q_block_size, k_block_size, ) @@ -172,4 +183,5 @@ def sage_quantize( k_scale=k_scale, v_scale=v_scale, k_mean=k_mean, + v_mean=v_mean, ) diff --git a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py index c6fdaf92480f..93744c2ad66f 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -1000,6 +1000,7 @@ def _(query: Optional[torch.Tensor], k_scale: torch.Tensor, v_scale: torch.Tensor, k_mean: Optional[torch.Tensor], + v_mean: Optional[torch.Tensor], q_block_size: int, k_block_size: int, reuse_stats: bool = False) -> None: diff --git a/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py b/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py index df2574bf4ae3..90b564070728 100644 --- a/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/visual_gen/attention/backends/trtllm.py @@ -409,6 +409,7 @@ def forward( "sage_attn_num_elts_per_blk_v": quant_cfg.v_block_size, "sage_attn_qk_int8": quant_cfg.qk_dtype == "int8", "sage_attn_smooth_k": quant_cfg.smooth_k, + "sage_attn_smooth_v": quant_cfg.smooth_v, "qk_fp8_per_tensor": quant_cfg.q_block_size == 0 and quant_cfg.qk_dtype == "fp8", "v_fp8_per_tensor": quant_cfg.v_block_size == 0 and quant_cfg.v_dtype == "fp8", diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index 380154d92108..9ce1f0472fd5 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -103,6 +103,15 @@ class QuantAttentionConfig(StrictBaseModel): "quantized type has to cover. SageAttention only." ), ) + smooth_v: bool = Field( + False, + status="prototype", + description=( + "Subtract the per-channel V mean before quantizing V and add it back to the output. " + "SageAttention only; applied where the attention runs on the PrimTS kernels, while " + "the trtllm-gen kernels quantize V without it." + ), + ) # Discriminated union of sparse attention configs. @@ -207,15 +216,15 @@ def _validate_quant_attention_config(self) -> "AttentionConfig": f"for backend='TRTLLM'. Supported SAGE and PrimTS recipes on this device: " f"{sorted(recipes)}." ) - if q_config.smooth_k and q_config.q_block_size == 0: + if (q_config.smooth_k or q_config.smooth_v) and q_config.q_block_size == 0: raise ValueError( - "smooth_k is a SageAttention option and does not apply to the per-tensor " - f"recipe {recipe!r}." + "smooth_k and smooth_v are SageAttention options and do not apply to the " + f"per-tensor recipe {recipe!r}." ) - elif q_config.smooth_k: + elif q_config.smooth_k or q_config.smooth_v: raise ValueError( - f"smooth_k is a SageAttention option and requires backend='TRTLLM', got " - f"backend='{self.backend}'." + f"smooth_k and smooth_v are SageAttention options and require backend='TRTLLM', " + f"got backend='{self.backend}'." ) elif self.backend == "CUTEDSL": if recipe not in CUTEDSL_RECIPES: diff --git a/tests/unittest/_torch/attention/kernels/test_sage_quant.py b/tests/unittest/_torch/attention/kernels/test_sage_quant.py index 07d9dae28dba..7ae17e78ea5b 100644 --- a/tests/unittest/_torch/attention/kernels/test_sage_quant.py +++ b/tests/unittest/_torch/attention/kernels/test_sage_quant.py @@ -49,12 +49,13 @@ def _quantize_summaries( v_scale: torch.Tensor, k_mean: torch.Tensor | None, quant_dtype: torch.dtype, + v_mean: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Quantize ``[batch, blocks, heads, head_dim]`` block means as the attention core does: every mean is a key of its sequence with a scale of its own, reusing the given statistics.""" batch, blocks = k.shape[:2] cu_seqlens = torch.arange(0, (batch + 1) * blocks, blocks, dtype=torch.int32, device=k.device) - stats = SageQuantized(None, None, None, None, None, v_scale, k_mean) + stats = SageQuantized(None, None, None, None, None, v_scale, k_mean, v_mean) actual = sage_quantize( None, k.flatten(0, 1), @@ -89,6 +90,7 @@ def _assert_summaries_match_reference( v_scale: torch.Tensor, k_mean: torch.Tensor | None, quant_dtype: torch.dtype, + v_mean: torch.Tensor | None = None, ) -> None: batch, blocks, heads, _ = k.shape divisor = 126.9 if quant_dtype == torch.int8 else 448.0 @@ -109,7 +111,8 @@ def _assert_summaries_match_reference( actual_scales[:, slots], scales.permute(2, 0, 1).reshape(heads, -1), rtol=1e-6, atol=0 ) _assert_quantized_close(actual_k, k_ref / scales.unsqueeze(-1), quant_dtype) - _assert_quantized_close(actual_v, v.float() / v_scale[:, None], torch.float8_e4m3fn) + v_ref = v.float() if v_mean is None else v.float() - v_mean[:, None] + _assert_quantized_close(actual_v, v_ref / v_scale[:, None], torch.float8_e4m3fn) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @@ -118,18 +121,21 @@ def _assert_summaries_match_reference( # A 100x smaller V scale saturates the quantized V means at the E4M3 maximum. @pytest.mark.parametrize("v_scale_divisor", [1.0, 100.0]) @pytest.mark.parametrize("smooth_k", [False, True]) +@pytest.mark.parametrize("smooth_v", [False, True]) def test_quantize_sage_summaries_matches_reference( shape: tuple[int, ...], dtype: torch.dtype, quant_dtype: torch.dtype, v_scale_divisor: float, smooth_k: bool, + smooth_v: bool, ) -> None: k, v, v_scale, k_mean = _summary_inputs(shape, dtype, smooth_k) v_scale = v_scale / v_scale_divisor - actual = _quantize_summaries(k, v, v_scale, k_mean, quant_dtype) + v_mean = v.float().mean(dim=1) if smooth_v else None + actual = _quantize_summaries(k, v, v_scale, k_mean, quant_dtype, v_mean) assert [result.dtype for result in actual] == [quant_dtype, torch.float8_e4m3fn, torch.float32] - _assert_summaries_match_reference(actual, k, v, v_scale, k_mean, quant_dtype) + _assert_summaries_match_reference(actual, k, v, v_scale, k_mean, quant_dtype, v_mean) @pytest.mark.parametrize("quant_dtype", [torch.int8, torch.float8_e4m3fn]) @@ -249,9 +255,10 @@ def _check_block_quantized( @pytest.mark.parametrize("block_size", [1, 4, 16, 32]) @pytest.mark.parametrize("quant_dtype", [torch.int8, torch.float8_e4m3fn]) @pytest.mark.parametrize("smooth_k", [False, True]) +@pytest.mark.parametrize("smooth_v", [False, True]) @pytest.mark.parametrize("head_dim", [64, 128]) def test_sage_quantize_matches_reference( - block_size: int, quant_dtype: torch.dtype, smooth_k: bool, head_dim: int + block_size: int, quant_dtype: torch.dtype, smooth_k: bool, smooth_v: bool, head_dim: int ) -> None: generator = torch.Generator(device="cuda").manual_seed(7) # Ragged sequences with partial final blocks; GQA with 2 query heads per KV head. @@ -272,6 +279,7 @@ def test_sage_quantize_matches_reference( k_block_size=block_size, qk_quant_dtype=quant_dtype, smooth_k=smooth_k, + smooth_v=smooth_v, ) if smooth_k: @@ -284,14 +292,24 @@ def test_sage_quantize_matches_reference( _check_block_quantized(actual.q, actual.q_scale, q, lengths, block_size, quant_dtype) _check_block_quantized(actual.k, actual.k_scale, k_ref, lengths, block_size, quant_dtype) - # V: one FP8 scale per channel of every sequence, from the BF16 maxima. - v_max = torch.stack([seq.abs().amax(dim=0) for seq in torch.split(v, lengths)]) - floor = torch.tensor(1e-3, dtype=torch.bfloat16, device="cuda") - v_scale = torch.maximum(v_max, floor).float() * torch.tensor(1 / 448.0, device="cuda") - torch.testing.assert_close(actual.v_scale, v_scale, rtol=0, atol=0) + # V: one FP8 scale per channel of every sequence, from the maxima of V or of V minus its mean. + if smooth_v: + v_mean = torch.stack([seq.float().mean(dim=0) for seq in torch.split(v, lengths)]) + torch.testing.assert_close(actual.v_mean, v_mean, rtol=1e-5, atol=1e-5) + v_ref = v.float() - torch.repeat_interleave(actual.v_mean, token_lengths, dim=0) + v_max = torch.stack([seq.abs().amax(dim=0) for seq in torch.split(v_ref, lengths)]) + v_scale = v_max.clamp_min(1e-3) / 448.0 + torch.testing.assert_close(actual.v_scale, v_scale, rtol=1e-6, atol=0) + else: + assert actual.v_mean is None + v_ref = v.float() + v_max = torch.stack([seq.abs().amax(dim=0) for seq in torch.split(v, lengths)]) + floor = torch.tensor(1e-3, dtype=torch.bfloat16, device="cuda") + v_scale = torch.maximum(v_max, floor).float() * torch.tensor(1 / 448.0, device="cuda") + torch.testing.assert_close(actual.v_scale, v_scale, rtol=0, atol=0) token_v_scale = torch.repeat_interleave(v_scale, token_lengths, dim=0) - error = (actual.v.float() * token_v_scale - v.float()).abs() - assert torch.all(error <= v.float().abs() * 0.07 + token_v_scale * 2**-6) + error = (actual.v.float() * token_v_scale - v_ref).abs() + assert torch.all(error <= v_ref.abs() * 0.07 + token_v_scale * 2**-6) def test_sage_quantize_mixed_block_sizes_and_query_lengths() -> None: @@ -315,6 +333,29 @@ def test_sage_quantize_mixed_block_sizes_and_query_lengths() -> None: assert actual.v_scale.shape == (2, 2, 128) +def test_sage_quantize_v_extrema_order_negative_zero() -> None: + # Every channel is -0.0 except for a -5 in the first token: across the warps that merge the + # V extrema, -0.0 must win the maximum and lose the minimum to -5. + x = torch.zeros(4096, 2, 128, device="cuda", dtype=torch.bfloat16) + v = torch.full_like(x, -0.0) + v[0] = -5.0 + cu_seqlens = _cu_seqlens([4096]) + actual = sage_quantize( + x, + x, + v, + cu_seqlens, + cu_seqlens, + q_block_size=1, + k_block_size=16, + qk_quant_dtype=torch.float8_e4m3fn, + smooth_v=True, + ) + v_ref = v.float() - v.float().mean(dim=0) + expected = v_ref.abs().amax(dim=0).clamp_min(1e-3) / 448.0 + torch.testing.assert_close(actual.v_scale[0], expected, rtol=1e-6, atol=0) + + @pytest.mark.parametrize("block_size", [2, 3, 64]) def test_sage_quantize_rejects_unsupported_block_sizes(block_size: int) -> None: x = torch.zeros(16, 2, 128, device="cuda", dtype=torch.bfloat16) diff --git a/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py b/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py index 26e387966815..f3f05a4e699a 100644 --- a/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py +++ b/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py @@ -445,6 +445,7 @@ def quantize(): "k_block_size": 16, "qk_quant_dtype": torch.float8_e4m3fn, "smooth_k": True, + "smooth_v": False, } assert sage_quantize.call_count == 1 @@ -465,11 +466,18 @@ def quantize(): assert actual.k_summary_scale is summary.k_scale sage_quantize.side_effect = None - # Dense calls and calls with a KV cache keep the attention op's quantization, and a layer - # without a library that takes quantized inputs quantizes nothing. + # Dense calls keep the attention op's quantization unless they need V smoothing, which + # the attention op cannot apply. sage_quantize.reset_mock() _set_block_sparse_inputs(args, None) assert quantize() is None + args.sage_attn_smooth_v = True + assert quantize() is quantized + assert sage_quantize.call_args.kwargs["smooth_v"] + + # Calls with a KV cache, and layers without a library that takes quantized inputs, + # quantize nothing. + sage_quantize.reset_mock() _set_block_sparse_inputs(args, proxy) attention._fmha_manager.fmha_libs = [] assert quantize() is None @@ -488,6 +496,38 @@ def test_contiguous_sage_requires_inputs_quantized_by_the_attention_core() -> No assert reason == "SageAttention requires inputs quantized by the attention core" +def test_dense_sage_with_v_smoothing_runs_on_a_dense_plan() -> None: + _attention, fmha, q, k, v, metadata, args = _contiguous_case() + _set_block_sparse_inputs(args, None) + _set_sage_recipe(args) + supported, reason = fmha._is_supported_with_reason( + q, k, v, metadata, args, phase=FmhaPhase.CONTEXT + ) + assert not supported + assert reason == "block-sparse forward inputs are required" + + args.sage_attn_smooth_v = True + args.sage_quantized = Mock() + supported, reason = fmha._is_supported_with_reason( + q, k, v, metadata, args, phase=FmhaPhase.CONTEXT + ) + assert supported, reason + key = fmha._make_plan_key( + q.view(2, 64, 2, 128), + None, + batch_size=2, + seq_len_q=64, + kv_capacity=256, + page_size=None, + mask_type="dense", + sage=block_sparse_fmha._sage_recipe(args), + ) + assert not key.use_block_sparse + assert (key.q_block_size, key.kv_block_size, key.max_blocks_per_row) == (128, 128, None) + assert key.sage.v_mean + assert key.unsupported_reason() is None + + def test_paged_wrapper_uses_zero_copy_padded_row_stride_block_tables(monkeypatch) -> None: attention, fmha, q, metadata, args = _paged_case() wrapper = Mock() diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/test_sage_block_sparse.py b/tests/unittest/_torch/visual_gen/sparse_attention/test_sage_block_sparse.py index 78cca59740f4..b8414ec1b265 100644 --- a/tests/unittest/_torch/visual_gen/sparse_attention/test_sage_block_sparse.py +++ b/tests/unittest/_torch/visual_gen/sparse_attention/test_sage_block_sparse.py @@ -82,6 +82,7 @@ def _backend( sparse_config: SolAttentionConfig | VideoSparseAttentionConfig | None = None, *, k_block_size: int | None = None, + smooth_v: bool = False, dtype: torch.dtype = torch.bfloat16, ) -> TrtllmAttention: quant_config = None @@ -96,6 +97,7 @@ def _backend( q_block_size=1, k_block_size=k_block_size, v_block_size=1, + smooth_v=smooth_v, ) return create_attention( backend="TRTLLM", @@ -113,12 +115,16 @@ def _backend( def _sol( - qk_dtype: _QuantDtype | None, *, phase: _SolPhase = "proxy", k_block_size: int | None = None + qk_dtype: _QuantDtype | None, + *, + phase: _SolPhase = "proxy", + k_block_size: int | None = None, + smooth_v: bool = False, ) -> TrtllmAttention: sparse_config = SolAttentionConfig( tau=-1.0e6 if phase == "exact" else 1.0e6, disabled_until_timestep=0.6 ) - return _backend(qk_dtype, sparse_config, k_block_size=k_block_size) + return _backend(qk_dtype, sparse_config, k_block_size=k_block_size, smooth_v=smooth_v) def _inputs(seq_len: int, *, seed: int = 20260921) -> tuple[torch.Tensor, ...]: @@ -192,6 +198,7 @@ def _assert_sparse_sage_active( assert state.output_dtype == dtype assert state.sage.q_block_size == quant_config.q_block_size assert state.sage.k_block_size == quant_config.k_block_size + assert state.sage.v_mean == quant_config.smooth_v def _parity_stats(actual: torch.Tensor, expected: torch.Tensor) -> dict[str, float]: @@ -326,6 +333,32 @@ def test_sol_sage_matches_unquantized_attention( ) +@pytest.mark.parametrize("case", ["sol_proxy", "sol_dense_timestep", "dense_model"]) +@torch.no_grad() +def test_sage_smooth_v_matches_unquantized_attention(case: str) -> None: + """V smoothing on mixed SOL routes, and on dense calls, which then run on a PrimTS dense plan.""" + q, k, v = _inputs(131) + if case == "dense_model": + reference, candidate = _backend(None), _backend("fp8", smooth_v=True) + expected, actual = ( + backend.forward(q=q, k=k, v=v, batch_size=q.shape[0], seq_len=q.shape[1]).view_as(q) + for backend in (reference, candidate) + ) + else: + phase = "dense" if case == "sol_dense_timestep" else "proxy" + reference, candidate = _sol(None, phase=phase), _sol("fp8", phase=phase, smooth_v=True) + expected = _forward_sol(reference, q, k, v, phase=phase) + actual = _forward_sol(candidate, q, k, v, phase=phase) + + wrappers = _sparse_wrappers(candidate) + assert wrappers, "V smoothing bypassed the PrimTS kernels" + for wrapper in wrappers: + state = wrapper._plan_state + assert state.sage is not None and state.sage.v_mean + assert state.use_block_sparse == (case == "sol_proxy") + _assert_parity(actual, expected, label=f"Sage fp8 smooth_v, {case}", qk_dtype="fp8") + + @pytest.mark.parametrize("qk_dtype", ["int8", "fp8"]) @torch.no_grad() def test_sol_sage_cuda_graph_refreshes_inputs_and_scales(qk_dtype: _QuantDtype) -> None: diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py index 3aead1c2c10e..e6118f5ac277 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py @@ -178,14 +178,15 @@ def test_supported_quant_config_trtllm( assert attention.quant_attention_config is not None + @pytest.mark.parametrize("smoothing", ["smooth_k", "smooth_v"]) @pytest.mark.parametrize("backend", ["TRTLLM", "CUTEDSL", "CUDNN", "FLASHINFER", "VANILLA"]) - def test_smooth_k_rejected_without_sage_recipe(self, backend): + def test_smoothing_rejected_without_sage_recipe(self, backend, smoothing): with patch("tensorrt_llm.visual_gen.args.get_sm_version", return_value=100): - with pytest.raises(ValidationError, match="smooth_k is a SageAttention option"): + with pytest.raises(ValidationError, match="are SageAttention options"): AttentionConfig( backend=backend, quant_attention_config=QuantAttentionConfig( - qk_dtype="bf16", v_dtype="fp8", smooth_k=True + qk_dtype="bf16", v_dtype="fp8", **{smoothing: True} ), ) From 27d9997fd8b9ea27481c8a10278069f799d561e8 Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Wed, 14 Oct 2026 10:08:19 +0000 Subject: [PATCH 08/13] [None][fix] Make SageAttention smoothing means bitwise reproducible The Sage quantization kernel reduced the K-smoothing and V-smoothing means with float atomics: every warp adds its partial mean, and the order the atomics land in changes from run to run, so the last bits of the means did too. Those bits flip some quantized K/V values, and a diffusion model amplifies the flips over its steps: two runs of the same request with smoothing produced visibly different videos. The statistics launch now sums the means in fixed point (32 fraction bits) in int64. Every warp's partial mean is rounded once and integer additions are exact, so the sum no longer depends on the order. Partial means are clamped to 2^22 so that the sums cannot overflow. Any mean keeps smoothing exact as long as every consumer uses the same one. - A small kernel after the statistics launch converts the sums into the float means that the quantizing launch and the consumers read. With V smoothing it also turns the V maximum, now collected in the V scale buffer itself, and the V minimum into the scales of V minus its mean, so the quantizing launch reads the V scales as without smoothing. - Another small kernel before the statistics launch starts every reduction from its identity, so callers no longer clear buffers: the attention op drops its memsets, and trtllm::sage_quantize allocates its sums and V minima uninitialized. - The attention op's K-smoothing workspace gains the int64 sum, which also makes dense SageAttention with smooth_k reproducible; the workspace test checks the order of the Sage slices. - Test that repeated sage_quantize calls with K and V smoothing give bitwise identical means and quantized K and V. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- cpp/tensorrt_llm/common/attentionOp.cpp | 11 +- cpp/tensorrt_llm/common/attentionWorkspace.h | 5 + cpp/tensorrt_llm/common/sageQuant.cu | 178 +++++++++++++----- cpp/tensorrt_llm/common/sageQuant.h | 21 ++- cpp/tensorrt_llm/thop/sageQuantizeOp.cpp | 31 ++- .../common/attentionWorkspaceTest.cpp | 6 +- .../attention/kernels/test_sage_quant.py | 35 ++++ 7 files changed, 200 insertions(+), 87 deletions(-) diff --git a/cpp/tensorrt_llm/common/attentionOp.cpp b/cpp/tensorrt_llm/common/attentionOp.cpp index 8ac041a4c2fa..ee7fd2fe5b5a 100644 --- a/cpp/tensorrt_llm/common/attentionOp.cpp +++ b/cpp/tensorrt_llm/common/attentionOp.cpp @@ -903,6 +903,7 @@ size_t AttentionOp::getWorkspaceSizeForContext(tensorrt_llm::DataType type, int3 size_t const sage_k_sfs_buffer_size = sizeof(float) * mNumAttnKVHeads * static_cast(k_max_n_blk); size_t const sage_v_sfs_buffer_size = sizeof(float) * static_cast(v_max_n_blk); size_t const sage_k_mean_buffer_size = mSageAttnSmoothK ? sizeof(float) * local_hidden_units_kv : 0; + size_t const sage_k_mean_sum_buffer_size = mSageAttnSmoothK ? sizeof(int64_t) * local_hidden_units_kv : 0; size_t const padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * padded_num_tokens; size_t const encoder_padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * padded_kv_tokens; @@ -946,6 +947,7 @@ size_t AttentionOp::getWorkspaceSizeForContext(tensorrt_llm::DataType type, int3 workspaceSizes.sageKScale = sage_k_sfs_buffer_size; workspaceSizes.sageVScale = sage_v_sfs_buffer_size; workspaceSizes.sageKMean = sage_k_mean_buffer_size; + workspaceSizes.sageKMeanSum = sage_k_mean_sum_buffer_size; workspaceSizes.cpWorkspace = cpWorkspaceSize; workspaceSizes.fmhaMultiCtasKvScratch = fmha_multi_ctas_kv_scratch_size; context_workspace_size = AttentionWorkspaceManager::buildContextLayout(workspaceSizes).totalSize; @@ -1629,6 +1631,7 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea size_t const sage_k_sfs_buffer_size = sizeof(float) * mNumAttnKVHeads * static_cast(k_max_n_blk); size_t const sage_v_sfs_buffer_size = sizeof(float) * static_cast(v_max_n_blk); size_t const sage_k_mean_buffer_size = mSageAttnSmoothK ? sizeof(float) * local_hidden_units_kv : 0; + size_t const sage_k_mean_sum_buffer_size = mSageAttnSmoothK ? sizeof(int64_t) * local_hidden_units_kv : 0; size_t const padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * params.batch_size * params.input_seq_length; @@ -1675,6 +1678,7 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea workspaceSizes.sageKScale = sage_k_sfs_buffer_size; workspaceSizes.sageVScale = sage_v_sfs_buffer_size; workspaceSizes.sageKMean = sage_k_mean_buffer_size; + workspaceSizes.sageKMeanSum = sage_k_mean_sum_buffer_size; workspaceSizes.cpWorkspace = cpWorkspaceSize; workspaceSizes.fmhaMultiCtasKvScratch = fmha_multi_ctas_kv_scratch_size; auto const workspaceLayout = AttentionWorkspaceManager::buildContextLayout(workspaceSizes); @@ -1965,12 +1969,6 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea mSageAttnQkInt8 ? "true" : "false"); TLLM_CHECK_WITH_INFO(!params.kv_scale_quant_orig, "SageAttention disregards the configured params.kv_scale_quant_orig, invalidating the result."); - // Reduction buffers must be initialized prior to invokeSageQuant(). - check_cuda_error(cudaMemsetAsync(workspaceViews.sageVScale, 0, sage_v_sfs_buffer_size, stream)); - if (mSageAttnSmoothK) - { - check_cuda_error(cudaMemsetAsync(workspaceViews.sageKMean, 0, sage_k_mean_buffer_size, stream)); - } // Common params for sageQuant tc::SageQuantParams sageQuantParams{}; @@ -1985,6 +1983,7 @@ int AttentionOp::enqueueContext(EnqueueContextParams const& params, cudaStrea sageQuantParams.ptrVQuant = workspaceViews.fp8VBuf; sageQuantParams.ptrVScale = workspaceViews.sageVScale; sageQuantParams.ptrKForMean = params.k_ptr; + sageQuantParams.ptrKMeanSum = workspaceViews.sageKMeanSum; sageQuantParams.ptrKMean = workspaceViews.sageKMean; sageQuantParams.smCount = mMultiProcessorCount; sageQuantParams.stream = stream; diff --git a/cpp/tensorrt_llm/common/attentionWorkspace.h b/cpp/tensorrt_llm/common/attentionWorkspace.h index 774235881aa0..99f71dc144d3 100644 --- a/cpp/tensorrt_llm/common/attentionWorkspace.h +++ b/cpp/tensorrt_llm/common/attentionWorkspace.h @@ -66,6 +66,7 @@ struct AttentionContextWorkspaceSizes size_t sageKScale{}; size_t sageVScale{}; size_t sageKMean{}; + size_t sageKMeanSum{}; size_t cpWorkspace{}; size_t fmhaMultiCtasKvScratch{}; }; @@ -98,6 +99,7 @@ struct AttentionContextWorkspaceLayout WorkspaceSlice sageKScale{}; WorkspaceSlice sageVScale{}; WorkspaceSlice sageKMean{}; + WorkspaceSlice sageKMeanSum{}; WorkspaceSlice cpWorkspace{}; WorkspaceSlice fmhaMultiCtasKvScratch{}; size_t totalSize{}; @@ -132,6 +134,7 @@ struct AttentionContextWorkspaceViews float* sageKScale{}; float* sageVScale{}; float* sageKMean{}; + int64_t* sageKMeanSum{}; T* gatherInBuffer{}; T* gatherOutBuffer{}; int* cuCpPartialSeqlens{}; @@ -258,6 +261,7 @@ class AttentionWorkspaceManager layout.sageKScale = nextSlice(offset, sizes.sageKScale, alignment); layout.sageVScale = nextSlice(offset, sizes.sageVScale, alignment); layout.sageKMean = nextSlice(offset, sizes.sageKMean, alignment); + layout.sageKMeanSum = nextSlice(offset, sizes.sageKMeanSum, alignment); layout.cpWorkspace = nextSlice(offset, sizes.cpWorkspace, alignment); layout.fmhaMultiCtasKvScratch = nextSlice(offset, sizes.fmhaMultiCtasKvScratch, alignment); layout.totalSize = offset; @@ -296,6 +300,7 @@ class AttentionWorkspaceManager views.sageKScale = ptr(workspace, layout.sageKScale); views.sageVScale = ptr(workspace, layout.sageVScale); views.sageKMean = ptr(workspace, layout.sageKMean); + views.sageKMeanSum = ptr(workspace, layout.sageKMeanSum); views.gatherInBuffer = ptr(workspace, layout.cpWorkspace); if (views.gatherInBuffer != nullptr) diff --git a/cpp/tensorrt_llm/common/sageQuant.cu b/cpp/tensorrt_llm/common/sageQuant.cu index a24a03db610f..614666b64deb 100644 --- a/cpp/tensorrt_llm/common/sageQuant.cu +++ b/cpp/tensorrt_llm/common/sageQuant.cu @@ -91,14 +91,93 @@ __device__ inline void sageAtomicMin(float* address, float value) } } +// Warps that reduce one row of the V statistics or the K mean. +constexpr int MaxStatWarps = 256; + +// The smoothing means are summed in fixed point with 32 fraction bits: each warp adds its partial +// mean, rounded once, and integer additions are exact, so the mean does not depend on the order the +// atomics land in, as it would with float atomics. A partial mean is bounded by the largest +// |input|; clamping it to MaxPartialMean keeps MaxStatWarps of them within int64, and keeps the +// power-of-two scalings below exact. Any mean is exact for smoothing as long as every consumer uses +// the same one: softmax ignores a constant subtracted from all keys, and V's mean is added back. +constexpr float MeanFixedScale = static_cast(1LL << 32); +constexpr float MaxPartialMean = static_cast(1LL << 62) / MeanFixedScale / MaxStatWarps; + +__device__ inline void sageMeanAdd(int64_t* sum, float partialMean) +{ + float const clamped = ::fminf(::fmaxf(partialMean, -MaxPartialMean), MaxPartialMean); + long long const fixed = __float2ll_rn(clamped * MeanFixedScale); + atomicAdd(reinterpret_cast(sum), static_cast(fixed)); +} + +// Starts the reductions of a VStage=1 launch from their identities: zero for the mean sums and for +// the V scales, which collect absolute maxima; with V smoothing, -inf for the signed maxima the V +// scale buffer collects instead and +inf for the minima. A null sum is skipped. +__global__ void sageInitStatsKernel(int64_t* kSum, int64_t* vSum, float* vScale, float* vMin, bool smoothV, int numel) +{ + int const i = blockIdx.x * blockDim.x + threadIdx.x; + if (i >= numel) + { + return; + } + if (kSum != nullptr) + { + kSum[i] = 0; + } + if (vSum != nullptr) + { + vSum[i] = 0; + } + vScale[i] = smoothV ? -INFINITY : 0.0f; + if (smoothV) + { + vMin[i] = INFINITY; + } +} + +// Converts the summed means of a VStage=1 launch into the float means the VStage=2 launch and the +// consumers read; a null sum is skipped. With V smoothing, vScale holds the per-channel maximum of V +// and vMin the minimum; the maximum becomes the scale of V minus its mean, which covers the larger +// distance from the mean, with the floor of the unsmoothed scales. +__global__ void sageMeanFromSumKernel( + int64_t const* kSum, float* kMean, int64_t const* vSum, float* vMean, float* vScale, float const* vMin, int numel) +{ + int const i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < numel && kSum != nullptr) + { + kMean[i] = __ll2float_rn(kSum[i]) * (1.0f / MeanFixedScale); + } + if (i < numel && vSum != nullptr) + { + float const mean = __ll2float_rn(vSum[i]) * (1.0f / MeanFixedScale); + vMean[i] = mean; + vScale[i] = ::fmaxf(::fmaxf(vScale[i] - mean, mean - vMin[i]), 1e-3f) / 448.0f; + } +} + +namespace +{ + +constexpr int StatThreads = 256; + +// The V scales and the means hold headDim values per sequence and head, or per head without +// ptrCuSeqLensV. +int getSageStatNumel(SageQuantParams const& params) +{ + return (params.ptrCuSeqLensV != nullptr ? params.batchSize : 1) * params.numHeadsV * params.headDim; +} + +} // namespace + // SageAttention quantization kernel. Tensors are interpreted as column-major [D, H, S], which is // the same physical layout as contiguous PyTorch [S, H, D]. Each launch quantizes Q or K. It can // simultaneously collect V scales (VStage=1) or quantize V using scales collected by an earlier // launch (VStage=2). When SmoothK is enabled, the VStage=1 launch also collects a per-head-per- // channel K mean and the VStage=2 launch uses it to perform K-smoothing before quantization. When -// SmoothV is enabled, the VStage=1 launch collects the per-channel V mean and signed extrema instead -// of the V scales; the VStage=2 launch derives the scales from them when ptrVMax is given and -// subtracts ptrVMean from V whenever it is given. +// SmoothV is enabled, the VStage=1 launch collects the per-channel V mean and signed extrema, the +// maximum in ptrVScale itself, and sageMeanFromSumKernel turns them into the V scales; the VStage=2 +// launch subtracts ptrVMean from V whenever it is given. The VStage=1 launch sums the means into +// ptrKMeanSum and ptrVMeanSum (see sageMeanAdd). // // Which tokens share a scale is a template parameter: Partition is a CuTe layout // ((intra...), (group...)) -> token offset inside a tile (see sagePartition.h). Contiguous blocks @@ -115,9 +194,9 @@ __device__ inline void sageAtomicMin(float* address, float value) template __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* ptrCuSeqLensQk, void const* ptrQk, - void* ptrQkQuant, float* ptrQkScale, void const* ptrKForMean, float* ptrKMean, int sumSeqLensV, int numHeadsV, - void const* ptrV, void* ptrVQuant, float* ptrVScale, int const* ptrCuSeqLensV, float* ptrVMean, float* ptrVMax, - float* ptrVMin) + void* ptrQkQuant, float* ptrQkScale, void const* ptrKForMean, int64_t* ptrKMeanSum, float* ptrKMean, + int sumSeqLensV, int numHeadsV, void const* ptrV, void* ptrVQuant, float* ptrVScale, int const* ptrCuSeqLensV, + int64_t* ptrVMeanSum, float* ptrVMean, float* ptrVMin) { using namespace cute; using namespace cutlass; @@ -384,7 +463,7 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p if constexpr (VStage == 1 && SmoothV) { - int const numWarpsToUse = cutlass::fast_min(numWarps, 256); + int const numWarpsToUse = cutlass::fast_min(numWarps, MaxStatWarps); int const numToksPerWave = numWarpsToUse * numToksPerWarp; if (warpId >= numWarpsToUse) { @@ -423,8 +502,8 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p rMin[i] = ::fminf(rMin[i], __shfl_xor_sync(0xffffffffu, rMin[i], delta)); } } - // A smoothed V's scales cover V minus its mean, which is final only after this launch, so - // the caller derives them from the mean and the extrema. + // A smoothed V's scales cover V minus its mean, which is final only after this launch: the + // maximum goes to the scale buffer, and sageMeanFromSumKernel turns it into the scale. if (threadInTokIdx == thrId && range.tokBegin < range.tokEnd) { int64_t const statOffset = static_cast(range.row) * HeadDim + threadInTokIdx * BestVL; @@ -432,15 +511,15 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p CUTLASS_PRAGMA_UNROLL for (int i = 0; i < BestVL; ++i) { - atomicAdd(ptrVMean + statOffset + i, rSum[i] * invNumTokens); - sageAtomicMax(ptrVMax + statOffset + i, rMax[i]); + sageMeanAdd(ptrVMeanSum + statOffset + i, rSum[i] * invNumTokens); + sageAtomicMax(ptrVScale + statOffset + i, rMax[i]); sageAtomicMin(ptrVMin + statOffset + i, rMin[i]); } } } else if constexpr (VStage == 1) { - int const numWarpsToUse = cutlass::fast_min(numWarps, 256); + int const numWarpsToUse = cutlass::fast_min(numWarps, MaxStatWarps); int const numToksPerWave = numWarpsToUse * numToksPerWarp; if (warpId >= numWarpsToUse) { @@ -487,31 +566,7 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p cute::copy( AutoVectorizingCopy{}, make_tensor(make_gmem_ptr(ptrVMean + statOffset), Shape{}), rVMean); } - if (ptrVMax != nullptr) - { - // A smoothed V's scales cover its largest distance from the mean, with the floor of the - // unsmoothed scales. One thread per channel slice writes them for the consumer. - Tensor rVLargest = make_tensor(Shape{}); - Tensor rVSmallest = make_tensor(Shape{}); - cute::copy(AutoVectorizingCopy{}, make_tensor(make_gmem_ptr(ptrVMax + statOffset), Shape{}), - rVLargest); - cute::copy(AutoVectorizingCopy{}, make_tensor(make_gmem_ptr(ptrVMin + statOffset), Shape{}), - rVSmallest); - CUTLASS_PRAGMA_UNROLL - for (int i = 0; i < BestVL; ++i) - { - rVScale(i) - = ::fmaxf(::fmaxf(rVLargest(i) - rVMean(i), rVMean(i) - rVSmallest(i)), 1e-3f) / 448.0f; - } - if (warpId == 0 && threadInTokIdx == thrId) - { - cute::copy(AutoVectorizingCopy{}, rVScale, gVSeqScale); - } - } - else - { - cute::copy(AutoVectorizingCopy{}, gVSeqScale, rVScale); - } + cute::copy(AutoVectorizingCopy{}, gVSeqScale, rVScale); cute::transform(rVScale_x2, rVScale_x2, cutlass::reciprocal_approximate>{}); for (; tokIdx < range.tokEnd; tokIdx += numToksPerWave) { @@ -543,8 +598,6 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p { Tensor gK = make_tensor(reinterpret_cast(ptrKForMean), make_shape(VL{}, Int{}, numHeadsV, sumSeqLensV)); - Tensor gKMean = make_tensor( - ptrKMean + static_cast(range.row) * HeadDim, make_shape(VL{}, Int{})); Tensor rK = make_tensor(Shape{}); Tensor rKCompute = make_tensor(Shape{}); Tensor rKSum = make_tensor(Shape{}); @@ -552,7 +605,7 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p Tensor rKCompute_x2 = recast>(rKCompute); int const numToksPerWarp = 32 / threadsPerHead; - int const numWarpsToUse = cutlass::fast_min(numWarps, 256); + int const numWarpsToUse = cutlass::fast_min(numWarps, MaxStatWarps); int const numToksPerWave = numWarpsToUse * numToksPerWarp; if (warpId >= numWarpsToUse) { @@ -561,7 +614,6 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p int tokIdx = range.tokBegin + warpId * numToksPerWarp + thrId / threadsPerHead; int const threadInTokIdx = thrId % threadsPerHead; Tensor gKSeq = gK(_, threadInTokIdx, headIdx, _); - Tensor gKMeanHead = gKMean(_, threadInTokIdx); clear(rKSum); for (; tokIdx < range.tokEnd; tokIdx += numToksPerWave) @@ -584,11 +636,12 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p } if (threadInTokIdx == thrId) { + int64_t const statOffset = static_cast(range.row) * HeadDim + threadInTokIdx * BestVL; float const invNumTokens = 1.0f / (range.tokEnd - range.tokBegin); CUTLASS_PRAGMA_UNROLL for (int i = 0; i < BestVL; ++i) { - atomicAdd(&gKMeanHead(i), rKSum(i) * invNumTokens); + sageMeanAdd(ptrKMeanSum + statOffset + i, rKSum(i) * invNumTokens); } } } @@ -611,15 +664,27 @@ void invokeSageQuantQkvImpl(SageQuantParams const& params) && params.ptrVScale != nullptr), "Invalid SageQuantV parameters"); TLLM_CHECK_WITH_INFO(!params.kSmooth || params.vStage != 0, "SageQuant K smoothing requires V staging"); - TLLM_CHECK_WITH_INFO( - !params.kSmooth || (params.ptrKMean != nullptr && (params.vStage != 1 || params.ptrKForMean != nullptr)), + TLLM_CHECK_WITH_INFO(!params.kSmooth + || (params.ptrKMean != nullptr + && (params.vStage != 1 || (params.ptrKForMean != nullptr && params.ptrKMeanSum != nullptr))), "Invalid SageQuant K-smoothing parameters"); TLLM_CHECK_WITH_INFO( params.ptrCuSeqLensV == nullptr || params.vStage != 0, "SageQuant per-sequence V statistics require V staging"); TLLM_CHECK_WITH_INFO(params.ptrVMean == nullptr - || (params.vStage != 0 && (params.vStage != 1 || (params.ptrVMax != nullptr && params.ptrVMin != nullptr))), + || (params.vStage != 0 + && (params.vStage != 1 || (params.ptrVMeanSum != nullptr && params.ptrVMin != nullptr))), "Invalid SageQuant V-smoothing parameters"); + if (params.vStage == 1) + { + // The VStage=1 launch reduces into the mean sums and the V statistics with atomics. + int const numel = getSageStatNumel(params); + sageInitStatsKernel<<(divUp(numel, StatThreads)), StatThreads, 0, params.stream>>>( + params.ptrKMeanSum, params.ptrVMeanSum, params.ptrVScale, params.ptrVMin, params.ptrVMean != nullptr, + numel); + TLLM_CUDA_CHECK(cudaPeekAtLastError()); + } + auto invokeKernel = [&](auto headDimStatic, auto partitionStatic) { constexpr int HeadDim_ = headDimStatic; @@ -627,9 +692,9 @@ void invokeSageQuantQkvImpl(SageQuantParams const& params) SageQuantParams kernelParams = params; void* kernelArgs[] = {&kernelParams.sumSeqLensQk, &kernelParams.batchSize, &kernelParams.ptrCuSeqLensQk, &kernelParams.ptrQk, &kernelParams.ptrQkQuant, &kernelParams.ptrQkScale, &kernelParams.ptrKForMean, - &kernelParams.ptrKMean, &kernelParams.sumSeqLensV, &kernelParams.numHeadsV, &kernelParams.ptrV, - &kernelParams.ptrVQuant, &kernelParams.ptrVScale, &kernelParams.ptrCuSeqLensV, &kernelParams.ptrVMean, - &kernelParams.ptrVMax, &kernelParams.ptrVMin}; + &kernelParams.ptrKMeanSum, &kernelParams.ptrKMean, &kernelParams.sumSeqLensV, &kernelParams.numHeadsV, + &kernelParams.ptrV, &kernelParams.ptrVQuant, &kernelParams.ptrVScale, &kernelParams.ptrCuSeqLensV, + &kernelParams.ptrVMeanSum, &kernelParams.ptrVMean, &kernelParams.ptrVMin}; auto launchKernel = [&](auto smoothKStatic, auto smoothVStatic, auto vStageStatic) { @@ -734,14 +799,23 @@ void invokeSageQuant(SageQuantParams const& params) if (params.inputType == DATA_TYPE_FP16) { invokeSageQuantQkvImpl(params); - return; } - if (params.inputType == DATA_TYPE_BF16) + else if (params.inputType == DATA_TYPE_BF16) { invokeSageQuantQkvImpl(params); - return; } - TLLM_THROW("SageQuant input must be FP16 or BF16"); + else + { + TLLM_THROW("SageQuant input must be FP16 or BF16"); + } + if (params.vStage == 1 && (params.ptrKMeanSum != nullptr || params.ptrVMeanSum != nullptr)) + { + int const numel = getSageStatNumel(params); + sageMeanFromSumKernel<<(divUp(numel, StatThreads)), StatThreads, 0, params.stream>>>( + params.ptrKMeanSum, params.ptrKMean, params.ptrVMeanSum, params.ptrVMean, params.ptrVScale, params.ptrVMin, + numel); + TLLM_CUDA_CHECK(cudaPeekAtLastError()); + } } namespace diff --git a/cpp/tensorrt_llm/common/sageQuant.h b/cpp/tensorrt_llm/common/sageQuant.h index 475a508c5b77..12a710a19b8a 100644 --- a/cpp/tensorrt_llm/common/sageQuant.h +++ b/cpp/tensorrt_llm/common/sageQuant.h @@ -19,6 +19,7 @@ #include "tensorrt_llm/common/cudaBf16Wrapper.h" #include "tensorrt_llm/kernels/multiHeadAttentionCommon.h" +#include #include #include #ifdef ENABLE_FP8 @@ -57,12 +58,14 @@ struct SageQuantParams kernels::Data_type inputType{kernels::DATA_TYPE_FP16}; kernels::Data_type quantType{kernels::DATA_TYPE_E4M3}; float* ptrQkScale{nullptr}; - // Optional source and scratch mean used to perform K-smoothing. - // (See below) collected at vStage==1, applied at vStage==2. + // Optional K-smoothing. The vStage==1 launch sums the mean of ptrKForMean into the int64 buffer + // ptrKMeanSum and converts it into ptrKMean, which the vStage==2 launch subtracts from K. void const* ptrKForMean{nullptr}; + int64_t* ptrKMeanSum{nullptr}; float* ptrKMean{nullptr}; // Optional arguments for SageQuantV: - // vStage: 0: disabled, 1: collect scales, 2: quantize + // vStage: 0: disabled, 1: collect scales, 2: quantize. The vStage==1 launch initializes the + // buffers it reduces into, ptrVScale and the sums and minima of the smoothing means, itself. int vStage{}; int sumSeqLensV{}; int numHeadsV{}; @@ -74,13 +77,13 @@ struct SageQuantParams // PrimTS block-sparse kernel expect; otherwise one [numHeadsV, headDim] row set covers the // whole batch, as the trtllm-gen kernels expect. int const* ptrCuSeqLensV{nullptr}; - // Optional V smoothing, with the layout of the V scales. ptrVMean is collected at vStage==1 and - // subtracted from V at vStage==2. A smoothed V's scales cover V minus its mean, which is final - // only after the whole vStage==1 launch, so that launch collects the signed per-channel extrema - // into ptrVMax/ptrVMin (filled with -inf/+inf beforehand) instead of writing ptrVScale, and the - // vStage==2 launch derives ptrVScale from them. Without ptrVMax, vStage==2 reads ptrVScale. + // Optional V smoothing, with the layout of the V scales. The vStage==1 launch sums the V mean + // into the int64 buffer ptrVMeanSum and converts it into ptrVMean, which the vStage==2 launch + // subtracts from V. A smoothed V's scales cover V minus its mean, which is final only after the + // whole vStage==1 launch, so that launch collects the signed per-channel maximum into ptrVScale + // and the minimum into ptrVMin, and turns them into the scales once the mean is known. + int64_t* ptrVMeanSum{nullptr}; float* ptrVMean{nullptr}; - float* ptrVMax{nullptr}; float* ptrVMin{nullptr}; // Hardware info. Required. int smCount{}; diff --git a/cpp/tensorrt_llm/thop/sageQuantizeOp.cpp b/cpp/tensorrt_llm/thop/sageQuantizeOp.cpp index 6e52efd0e02c..0974bc1c1110 100644 --- a/cpp/tensorrt_llm/thop/sageQuantizeOp.cpp +++ b/cpp/tensorrt_llm/thop/sageQuantizeOp.cpp @@ -23,7 +23,6 @@ #include #include -#include #include TRTLLM_NAMESPACE_BEGIN @@ -135,6 +134,7 @@ void launch(SageQuantParams params, torch::Tensor const& input, torch::Tensor co // quantized to FP8 E4M3 with one scale per channel of each sequence. The V scales and the means // are [batch, heads_kv, head_dim]; k_mean and v_mean are the K- and V-smoothing means, subtracted // before quantization and filled exactly when given; a smoothed V's scales cover V minus its mean. +// Repeated calls on the same GPU write bitwise identical values, the means included. // // With reuse_stats, only K and V are quantized, with the V scales and the means of an earlier call // as inputs, for tensors that must dequantize like that call's K and V, such as the K/V block means @@ -184,29 +184,22 @@ void sage_quantize(std::optional const& query, torch::Tensor cons TORCH_CHECK(q_quant->scalar_type() == k_quant.scalar_type(), "q_quant and k_quant must share a dtype"); checkBlockScales(*query, *q_scale, q_block_size, batchSize, "q_scale"); - // The V scales and the means are reduced with atomics into zeroed buffers. A smoothed V's scales - // are derived from the extrema instead. - if (!v_mean.has_value()) + // The launch initializes the buffers it reduces into: the V scales, and the sums and minima of + // the smoothing means below. + std::optional k_mean_sum; + if (params.kSmooth) { - v_scale.zero_(); + k_mean_sum = torch::empty_like(*k_mean, torch::kInt64); + params.ptrKForMean = key.data_ptr(); + params.ptrKMeanSum = k_mean_sum->data_ptr(); } - if (k_mean.has_value()) - { - k_mean->zero_(); - } - if (v_mean.has_value()) - { - v_mean->zero_(); - } - - params.ptrKForMean = params.kSmooth ? key.data_ptr() : nullptr; - std::optional v_max; + std::optional v_mean_sum; std::optional v_min; if (v_mean.has_value()) { - v_max = torch::full_like(v_scale, -std::numeric_limits::infinity()); - v_min = torch::full_like(v_scale, std::numeric_limits::infinity()); - params.ptrVMax = v_max->data_ptr(); + v_mean_sum = torch::empty_like(*v_mean, torch::kInt64); + v_min = torch::empty_like(v_scale); + params.ptrVMeanSum = v_mean_sum->data_ptr(); params.ptrVMin = v_min->data_ptr(); } // Q, while collecting the V statistics and the K mean; then K (smoothed by the K mean) and V. diff --git a/cpp/tests/unit_tests/common/attentionWorkspaceTest.cpp b/cpp/tests/unit_tests/common/attentionWorkspaceTest.cpp index b0f2dc0dcd1a..bb74d6fc3a34 100644 --- a/cpp/tests/unit_tests/common/attentionWorkspaceTest.cpp +++ b/cpp/tests/unit_tests/common/attentionWorkspaceTest.cpp @@ -152,7 +152,9 @@ TEST(AttentionWorkspaceManagerTest, ContextLayoutMatchesAttentionOpOrdering) sizes.sageQScale = 101; sizes.sageKScale = 103; sizes.sageVScale = 107; - sizes.cpWorkspace = 109; + sizes.sageKMean = 109; + sizes.sageKMeanSum = 113; + sizes.cpWorkspace = 137; auto const layout = tcop::AttentionWorkspaceManager::buildContextLayout(sizes); @@ -182,6 +184,8 @@ TEST(AttentionWorkspaceManagerTest, ContextLayoutMatchesAttentionOpOrdering) expectNextSlice("sageQScale", layout.sageQScale, sizes.sageQScale, expectedOffset); expectNextSlice("sageKScale", layout.sageKScale, sizes.sageKScale, expectedOffset); expectNextSlice("sageVScale", layout.sageVScale, sizes.sageVScale, expectedOffset); + expectNextSlice("sageKMean", layout.sageKMean, sizes.sageKMean, expectedOffset); + expectNextSlice("sageKMeanSum", layout.sageKMeanSum, sizes.sageKMeanSum, expectedOffset); expectNextSlice("cpWorkspace", layout.cpWorkspace, sizes.cpWorkspace, expectedOffset); EXPECT_EQ(layout.totalSize, expectedOffset); } diff --git a/tests/unittest/_torch/attention/kernels/test_sage_quant.py b/tests/unittest/_torch/attention/kernels/test_sage_quant.py index 7ae17e78ea5b..3dc2d0bd7f06 100644 --- a/tests/unittest/_torch/attention/kernels/test_sage_quant.py +++ b/tests/unittest/_torch/attention/kernels/test_sage_quant.py @@ -356,6 +356,41 @@ def test_sage_quantize_v_extrema_order_negative_zero() -> None: torch.testing.assert_close(actual.v_scale[0], expected, rtol=1e-6, atol=0) +@pytest.mark.parametrize("quant_dtype", [torch.int8, torch.float8_e4m3fn]) +def test_sage_quantize_smoothing_is_bitwise_reproducible(quant_dtype: torch.dtype) -> None: + # Long sequences spread every mean over the most warps; summed with float atomics, the means + # and with them the quantized K and V would change from call to call. + generator = torch.Generator(device="cuda").manual_seed(5) + lengths = [9000, 4097, 7001] + q = _ragged(lengths, 4, 128, generator) + k = _ragged(lengths, 4, 128, generator) + v = _ragged(lengths, 4, 128, generator) + cu_seqlens = _cu_seqlens(lengths) + + def quantize(): + return sage_quantize( + q, + k, + v, + cu_seqlens, + cu_seqlens, + q_block_size=1, + k_block_size=16, + qk_quant_dtype=quant_dtype, + smooth_k=True, + smooth_v=True, + ) + + # Q does not use the means; the K scales are checked through K, as their buffer also holds + # slots that are never written. + first = quantize() + for _ in range(4): + again = quantize() + for name in ("k", "v", "v_scale", "k_mean", "v_mean"): + expected, actual = getattr(first, name), getattr(again, name) + assert torch.equal(expected.view(torch.uint8), actual.view(torch.uint8)), name + + @pytest.mark.parametrize("block_size", [2, 3, 64]) def test_sage_quantize_rejects_unsupported_block_sizes(block_size: int) -> None: x = torch.zeros(16, 2, 128, device="cuda", dtype=torch.bfloat16) From 2a4a797943370fcba45c85365bce798df783c928 Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Wed, 14 Oct 2026 10:53:32 +0000 Subject: [PATCH 09/13] [None][test] Trim the Sage quantization tests to their distinct cases The quantizer and VisualGen Sage tests multiplied parameters whose combinations exercise no further code: every smoothing pair with every head dimension for each block size, 96 summary cases that include V scales no real call produces, and both Q/K dtypes where the path under test does not depend on them. Every kernel instantiation, block size, Q/K dtype, and SOL route kind stays covered, as do the torch.compile and CUDA Graph paths, the reproducibility check, and the edge cases; 192 cases become 71. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- .../attention/kernels/test_sage_quant.py | 32 +++++++------------ .../test_sage_block_sparse.py | 22 ++++++------- 2 files changed, 22 insertions(+), 32 deletions(-) diff --git a/tests/unittest/_torch/attention/kernels/test_sage_quant.py b/tests/unittest/_torch/attention/kernels/test_sage_quant.py index 3dc2d0bd7f06..7db0a5da2473 100644 --- a/tests/unittest/_torch/attention/kernels/test_sage_quant.py +++ b/tests/unittest/_torch/attention/kernels/test_sage_quant.py @@ -117,22 +117,13 @@ def _assert_summaries_match_reference( @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize("quant_dtype", [torch.int8, torch.float8_e4m3fn]) -@pytest.mark.parametrize("shape", [(2, 9, 3, 64), (3, 5, 2, 128), (1, 129, 2, 128)]) -# A 100x smaller V scale saturates the quantized V means at the E4M3 maximum. -@pytest.mark.parametrize("v_scale_divisor", [1.0, 100.0]) -@pytest.mark.parametrize("smooth_k", [False, True]) -@pytest.mark.parametrize("smooth_v", [False, True]) +@pytest.mark.parametrize("shape", [(2, 9, 3, 64), (1, 129, 2, 128)]) +@pytest.mark.parametrize("smooth", [False, True]) def test_quantize_sage_summaries_matches_reference( - shape: tuple[int, ...], - dtype: torch.dtype, - quant_dtype: torch.dtype, - v_scale_divisor: float, - smooth_k: bool, - smooth_v: bool, + shape: tuple[int, ...], dtype: torch.dtype, quant_dtype: torch.dtype, smooth: bool ) -> None: - k, v, v_scale, k_mean = _summary_inputs(shape, dtype, smooth_k) - v_scale = v_scale / v_scale_divisor - v_mean = v.float().mean(dim=1) if smooth_v else None + k, v, v_scale, k_mean = _summary_inputs(shape, dtype, smooth_k=smooth) + v_mean = v.float().mean(dim=1) if smooth else None actual = _quantize_summaries(k, v, v_scale, k_mean, quant_dtype, v_mean) assert [result.dtype for result in actual] == [quant_dtype, torch.float8_e4m3fn, torch.float32] _assert_summaries_match_reference(actual, k, v, v_scale, k_mean, quant_dtype, v_mean) @@ -254,9 +245,11 @@ def _check_block_quantized( @pytest.mark.parametrize("block_size", [1, 4, 16, 32]) @pytest.mark.parametrize("quant_dtype", [torch.int8, torch.float8_e4m3fn]) -@pytest.mark.parametrize("smooth_k", [False, True]) -@pytest.mark.parametrize("smooth_v", [False, True]) -@pytest.mark.parametrize("head_dim", [64, 128]) +# Every smoothing combination is its own kernel; each also takes one of the head dimensions. +@pytest.mark.parametrize( + ("smooth_k", "smooth_v", "head_dim"), + [(False, False, 64), (True, False, 128), (False, True, 64), (True, True, 128)], +) def test_sage_quantize_matches_reference( block_size: int, quant_dtype: torch.dtype, smooth_k: bool, smooth_v: bool, head_dim: int ) -> None: @@ -391,8 +384,7 @@ def quantize(): assert torch.equal(expected.view(torch.uint8), actual.view(torch.uint8)), name -@pytest.mark.parametrize("block_size", [2, 3, 64]) -def test_sage_quantize_rejects_unsupported_block_sizes(block_size: int) -> None: +def test_sage_quantize_rejects_unsupported_block_sizes() -> None: x = torch.zeros(16, 2, 128, device="cuda", dtype=torch.bfloat16) cu_seqlens = _cu_seqlens([16]) with pytest.raises(RuntimeError, match="must be 1, 4, 16, or 32"): @@ -402,7 +394,7 @@ def test_sage_quantize_rejects_unsupported_block_sizes(block_size: int) -> None: x, cu_seqlens, cu_seqlens, - q_block_size=block_size, + q_block_size=3, k_block_size=16, qk_quant_dtype=torch.int8, ) diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/test_sage_block_sparse.py b/tests/unittest/_torch/visual_gen/sparse_attention/test_sage_block_sparse.py index b8414ec1b265..74dd6445f7f5 100644 --- a/tests/unittest/_torch/visual_gen/sparse_attention/test_sage_block_sparse.py +++ b/tests/unittest/_torch/visual_gen/sparse_attention/test_sage_block_sparse.py @@ -285,11 +285,8 @@ def test_vsa_sage_matches_unquantized_sparse_attention(qk_dtype: _QuantDtype) -> ("qk_dtype", "phase", "k_block_size"), [ ("int8", "proxy", 16), - ("fp8", "proxy", 4), ("fp8", "proxy", 1), - ("int8", "exact", 16), ("fp8", "exact", 4), - ("int8", "dense", 16), ("fp8", "dense", 4), ], ) @@ -359,11 +356,10 @@ def test_sage_smooth_v_matches_unquantized_attention(case: str) -> None: _assert_parity(actual, expected, label=f"Sage fp8 smooth_v, {case}", qk_dtype="fp8") -@pytest.mark.parametrize("qk_dtype", ["int8", "fp8"]) @torch.no_grad() -def test_sol_sage_cuda_graph_refreshes_inputs_and_scales(qk_dtype: _QuantDtype) -> None: +def test_sol_sage_cuda_graph_refreshes_inputs_and_scales() -> None: """A captured mixed-proxy Sage layer consumes new Q/K/V and scales on replay.""" - backend = _sol(qk_dtype) + backend = _sol("fp8") q, k, v = _inputs(131) previous = _forward_sol(backend, q, k, v).clone() _assert_sparse_sage_active(backend, proxy=True, token_mask=False) @@ -382,7 +378,7 @@ def test_sol_sage_cuda_graph_refreshes_inputs_and_scales(qk_dtype: _QuantDtype) replayed = captured.clone() expected = _forward_sol(backend, q, k, v) assert _parity_stats(expected, previous)["relative_l2"] > 0.1 - _assert_parity(replayed, expected, label=f"SOL {qk_dtype} Sage CUDA Graph replay") + _assert_parity(replayed, expected, label="SOL Sage CUDA Graph replay") def _masked_zero_score_inputs( @@ -407,8 +403,11 @@ def _masked_zero_score_inputs( return q, -q, v, routes -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16], ids=["bf16", "fp16"]) -@pytest.mark.parametrize("qk_dtype", ["int8", "fp8"]) +@pytest.mark.parametrize( + ("dtype", "qk_dtype"), + [(torch.bfloat16, "int8"), (torch.float16, "fp8")], + ids=["bf16-int8", "fp16-fp8"], +) @torch.no_grad() def test_sage_zero_scores_preserve_token_masks_and_empty_rows( qk_dtype: _QuantDtype, dtype: torch.dtype @@ -434,9 +433,8 @@ def test_sage_zero_scores_preserve_token_masks_and_empty_rows( torch.testing.assert_close(actual[:, 64:128], torch.ones_like(actual[:, 64:128])) -@pytest.mark.parametrize("qk_dtype", ["int8", "fp8"]) @torch.no_grad() -def test_sol_sage_zero_values_preserve_zero_output(qk_dtype: _QuantDtype) -> None: +def test_sol_sage_zero_values_preserve_zero_output() -> None: """Exact and proxy SOL contributions stay zero through the V-scale floor. The inputs and tau match the mixed-route proxy cases above, so both route @@ -444,7 +442,7 @@ def test_sol_sage_zero_values_preserve_zero_output(qk_dtype: _QuantDtype) -> Non """ q, k, v = _inputs(131) zero_v = torch.zeros_like(v) - candidate = _sol(qk_dtype) + candidate = _sol("fp8") actual = _forward_sol(candidate, q, k, zero_v) _assert_sparse_sage_active(candidate, proxy=True, token_mask=False) assert torch.count_nonzero(actual).item() == 0, "zero exact values and proxy means stay zero" From a139886c2ec5530f5481da9497967b0eb09726bc Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Thu, 15 Oct 2026 01:05:14 +0000 Subject: [PATCH 10/13] [None][test] Give the VisualGen attention test double its dense-layer flag TrtllmAttention resolves at construction whether its layer is dense, and should_use_sparse reads that flag. The helper that builds the attention with object.__new__ did not set it, so test_forward_hands_separate_qkv_to_core_when_backend_rejects_fused_qkv failed with an AttributeError. The B200 test list also drops its separate test_sage_quant.py entry, which the unittest/_torch/attention entry already runs. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- tests/integration/test_lists/test-db/l0_b200.yml | 1 - .../unittest/_torch/visual_gen/test_trtllm_attention_metadata.py | 1 + 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 08b371fc7a65..88dbd05b0221 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -340,7 +340,6 @@ l0_b200: - unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor.py - unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor_kernels.py - unittest/_torch/visual_gen/sparse_attention/test_sage_block_sparse.py - - unittest/_torch/attention/kernels/test_sage_quant.py - unittest/_torch/visual_gen/test_attention_flashinfer.py - unittest/_torch/visual_gen/test_attention_trtllm_sage.py - unittest/_torch/visual_gen/test_attention_cudnn.py diff --git a/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py b/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py index 056cfb5d39d8..a8dad92f8c06 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py @@ -81,6 +81,7 @@ def _make_wrapper(cls=visual_trtllm.TrtllmAttention, *, quant_attention_config=N attention.quant_attention_config = quant_attention_config attention.layer_idx = 0 attention.sparse_params = None + attention.is_dense_layer = False attention.metadata = visual_trtllm.TrtllmAttentionMetadata( device=torch.device("cpu"), attention_metadata_state={} ) From 4651f0c53b41ab46f7362bfb3fe527eda60aa038 Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Thu, 15 Oct 2026 15:52:06 +0000 Subject: [PATCH 11/13] [None][fix] Check SageQuant mean sums against their outputs and name the V-scale constants - The mean conversion after a statistics launch writes the mean of every sum it is given and, for V, reads the minimum. invokeSageQuant now rejects a K or V mean sum without them, as it already rejected smoothing without the sums. - The K mean of a sequence is read at row seqIdx * numHeadsV + headIdx, the row the statistics launch writes, instead of relying on the K launch having as many heads as V. - The FP8 E4M3 maximum and the V-scale floor are named constants shared by the smoothed and unsmoothed V scales; the FP8 Q/K scales use the same maximum. The quantized outputs are bitwise unchanged. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- cpp/tensorrt_llm/common/sageQuant.cu | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/cpp/tensorrt_llm/common/sageQuant.cu b/cpp/tensorrt_llm/common/sageQuant.cu index 614666b64deb..b32ef9f75a86 100644 --- a/cpp/tensorrt_llm/common/sageQuant.cu +++ b/cpp/tensorrt_llm/common/sageQuant.cu @@ -40,6 +40,12 @@ using tensorrt_llm::kernels::DATA_TYPE_INT8; // Most tokens of one Q/K scale group a thread holds in registers at once. constexpr int MaxChunkTokens = 8; +// Largest FP8 E4M3 magnitude. A V scale maps the largest magnitude of its channel onto it, with a +// floor that keeps the reciprocal of a constant channel's scale finite; the smoothed and unsmoothed +// V scales share both. +constexpr float E4M3Max = 448.0f; +constexpr float VScaleFloor = 1e-3f; + // The V tokens one CTA row reduces for the V scales and the K mean, and the scale row they go to. // blockIdx.y is (headIdx * batchSize + seqIdx) when the statistics are per sequence; otherwise // rows below numHeadsV each cover one head over the whole batch. @@ -151,7 +157,7 @@ __global__ void sageMeanFromSumKernel( { float const mean = __ll2float_rn(vSum[i]) * (1.0f / MeanFixedScale); vMean[i] = mean; - vScale[i] = ::fmaxf(::fmaxf(vScale[i] - mean, mean - vMin[i]), 1e-3f) / 448.0f; + vScale[i] = ::fmaxf(::fmaxf(vScale[i] - mean, mean - vMin[i]), VScaleFloor) / E4M3Max; } } @@ -208,7 +214,7 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p static_assert(!SmoothV || VStage == 1, "V smoothing only changes the statistics launch"); static_assert(std::is_same_v || std::is_same_v, "Unrecognized target dtype for quantization"); - constexpr float TypeMax = cute::is_same_v ? 448.0f : static_cast(126.9f); + constexpr float TypeMax = cute::is_same_v ? E4M3Max : static_cast(126.9f); constexpr int BestVL = 128 / sizeof_bits_v; using VL = Int; @@ -263,7 +269,7 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p Tensor rKMean = make_tensor(Shape{}); if constexpr (SmoothK && VStage == 2) { - int const meanRow = ptrCuSeqLensV != nullptr ? seqIdx * numHeads + headIdx : headIdx; + int const meanRow = ptrCuSeqLensV != nullptr ? seqIdx * numHeadsV + headIdx : headIdx; Tensor gKMean = make_tensor( make_gmem_ptr(ptrKMean + static_cast(meanRow) * HeadDim + threadInScaleIdx * BestVL), Shape{}); @@ -528,7 +534,7 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p CUTLASS_PRAGMA_UNROLL for (int i = 0; i < size(rVScale); ++i) { - rVScale(i) = 1e-3f; + rVScale(i) = VScaleFloor; } cute::transform(rVScale_x2, rVMax_x2, cutlass::NumericArrayConverter::convert); for (; tokIdx < range.tokEnd; tokIdx += numToksPerWave) @@ -538,7 +544,7 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p cute::transform(rV_x2, rVMax_x2, rVMax_x2, cutlass::maximum>{}); } cute::transform(rVMax_x2, rVScale_x2, cutlass::NumericArrayConverter::convert); - cute::transform(rVScale_x2, rVScale_x2, cutlass::scale>{1 / 448.0f}); + cute::transform(rVScale_x2, rVScale_x2, cutlass::scale>{1.0f / E4M3Max}); for (int delta = threadsPerHead; delta < 32; delta <<= 1) { cute::transform(rVScale, rVScale, @@ -674,6 +680,12 @@ void invokeSageQuantQkvImpl(SageQuantParams const& params) || (params.vStage != 0 && (params.vStage != 1 || (params.ptrVMeanSum != nullptr && params.ptrVMin != nullptr))), "Invalid SageQuant V-smoothing parameters"); + // The mean conversion after a VStage=1 launch writes the mean of every sum it is given, and for V + // also reads the minimum. + TLLM_CHECK_WITH_INFO( + params.ptrKMeanSum == nullptr || params.ptrKMean != nullptr, "Invalid SageQuant K-smoothing parameters"); + TLLM_CHECK_WITH_INFO(params.ptrVMeanSum == nullptr || (params.ptrVMean != nullptr && params.ptrVMin != nullptr), + "Invalid SageQuant V-smoothing parameters"); if (params.vStage == 1) { From ee1eaaf62edfcba2d324e7e20e69f1d531ef87b1 Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Sat, 17 Oct 2026 02:18:43 +0000 Subject: [PATCH 12/13] [None][chore] List smooth_v in the trtllm-serve VisualGen config schema The checked-in schema for trtllm-serve --visual_gen_args now includes QuantAttentionConfig.smooth_v. It is regenerated with scripts/generate_trtllm_serve_schemas.py; the other two schemas are left unchanged, since their differences come from LLM arguments this change does not touch. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- .../schemas/trtllm-serve-visual-gen-config.schema.json | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tensorrt_llm/schemas/trtllm-serve-visual-gen-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-visual-gen-config.schema.json index cffa1974bcdc..1ac421acdac7 100644 --- a/tensorrt_llm/schemas/trtllm-serve-visual-gen-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-visual-gen-config.schema.json @@ -510,6 +510,13 @@ "title": "Smooth K", "type": "boolean" }, + "smooth_v": { + "default": false, + "description": "Subtract the per-channel V mean before quantizing V and add it back to the output. SageAttention only; applied where the attention runs on the PrimTS kernels, while the trtllm-gen kernels quantize V without it.", + "status": "prototype", + "title": "Smooth V", + "type": "boolean" + }, "v_block_size": { "default": 0, "description": "V quantization block size on the hidden dimension; 0 uses one tensor-wide V scale.", From 08d8cfa731414f27a20c31bbe9967f254a711cd3 Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Mon, 19 Oct 2026 14:23:00 +0000 Subject: [PATCH 13/13] [None][fix] Quantize Sage inputs for every FMHA library and reject smooth_v on SM90 The attention core quantizes the Q/K/V of block-sparse and V-smoothing calls without asking which FMHA libraries of the layer take quantized inputs; each library decides in its support check whether it uses them, and the attention op ignores them. Fmha.supports_sage_quantized_inputs had no other reader and is removed. V smoothing needs the PrimTS kernels, which add the V mean back to the attention output and run only on SM100 and SM103, and the core quantizer has no layout for the SM90 recipe. AttentionConfig now rejects smooth_v on other architectures instead of quantizing inputs that no kernel there can use. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- .../features/visualgen-quantized-attention.md | 2 +- .../attention/ATTENTION_DEVELOPER_GUIDE.md | 5 ++-- .../attention/backends/fmha/interface.py | 3 --- .../backends/fmha/prims_ts_block_sparse.py | 1 - .../_torch/attention/backends/trtllm.py | 9 ++++--- tensorrt_llm/visual_gen/args.py | 5 ++++ .../sparse/test_prims_ts_block_sparse.py | 12 ++++------ .../_torch/visual_gen/test_visual_gen_args.py | 24 +++++++++++++++++++ 8 files changed, 41 insertions(+), 20 deletions(-) diff --git a/docs/source/features/visualgen-quantized-attention.md b/docs/source/features/visualgen-quantized-attention.md index 78663f54bf4e..4ecc17b732f1 100644 --- a/docs/source/features/visualgen-quantized-attention.md +++ b/docs/source/features/visualgen-quantized-attention.md @@ -120,7 +120,7 @@ attention_config: - On H100/H200/H800/H20, only one recipe with fixed `qk_dtype: "int8"` is supported. - On B200/GB200, it is recommended to use the `qk_dtype: "int8"` recipes. - On B300/GB300, use `qk_dtype: "fp8"` and evaluate output quality because it can be less accurate than the INT8 Q/K recipe on B200/GB200. -- `smooth_v` needs a kernel that adds the V mean back to the output. The PrimTS kernels behind VSA and SOL do; with `smooth_v`, dense attention calls (dense layers and timesteps, or a dense model) also run on PrimTS where their shape allows. The trtllm-gen kernels quantize V without the mean and log a warning once. +- `smooth_v` needs a kernel that adds the V mean back to the output. The PrimTS kernels behind VSA and SOL do; with `smooth_v`, dense attention calls (dense layers and timesteps, or a dense model) also run on PrimTS where their shape allows. The trtllm-gen kernels quantize V without the mean and log a warning once. PrimTS runs on B200/GB200 and B300/GB300 only, so H100/H200/H800/H20 reject `smooth_v`. **Configuration.** diff --git a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index 03f8c537456d..9da5468d2b2e 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md @@ -234,9 +234,8 @@ only static format, proxy, geometry, and capacity choices; every run receives the live routes, summaries, validity bits, page tables, and sequence lengths. The contiguous block-sparse path also consumes the SageAttention recipe in -`AttentionForwardArgs`. Prediction sees the original Q/K/V. When a library of -the layer declares `supports_sage_quantized_inputs`, the attention core then -quantizes the Q/K/V of a block-sparse call before FMHA selection, with +`AttentionForwardArgs`. Prediction sees the original Q/K/V. The attention core +then quantizes the Q/K/V of a block-sparse call before FMHA selection, with TensorRT-LLM's Sage quantization kernel (`attention/kernels/sage_quant.py`), and passes the result in `AttentionForwardArgs.sage_quantized`. Q/K get per-block scales and V per-sequence channel scales; the summaries of proxy diff --git a/tensorrt_llm/_torch/attention/backends/fmha/interface.py b/tensorrt_llm/_torch/attention/backends/fmha/interface.py index d65aeff14d6e..e3905dd1964d 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/interface.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/interface.py @@ -43,9 +43,6 @@ class Fmha(ABC): supports_skip_correction: ClassVar[bool] = False supports_block_sparse_inputs: ClassVar[bool] = False supports_per_tensor_fp8: ClassVar[bool] = False - # Reads the SageAttention inputs the attention core quantized from - # ``AttentionForwardArgs.sage_quantized``. - supports_sage_quantized_inputs: ClassVar[bool] = False supports_workspace_reclamation: bool = False supports_fp4_mla: ClassVar[bool] = False diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py index 4aa2760ef406..7dab70862935 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py @@ -233,7 +233,6 @@ class PrimsTSBlockSparseFmha(PrimsTSFmha): PLAN_CACHE_KEY = "prims_ts_block_sparse" supports_block_sparse_inputs = True - supports_sage_quantized_inputs = True def _is_supported( self, diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 79dd50b31676..02f16c1484b3 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -2027,14 +2027,13 @@ def _quantize_sage_inputs( The attention op quantizes dense SageAttention inside its kernels and cannot add a V mean back, so only block-sparse calls and calls with V - smoothing are quantized here, for the FMHA libraries that take - quantized inputs: separate Q, K, and V without a KV cache. + smoothing are quantized here, given separate Q, K, and V without a KV + cache. Each FMHA library decides whether it takes the quantized inputs; + the attention op ignores them. """ qk_quant_dtype = forward_args.sage_attn_qk_dtype if (qk_quant_dtype is None or k is None or v is None - or metadata.kv_cache_manager is not None - or not any(fmha.supports_sage_quantized_inputs - for fmha in self._fmha_manager.fmha_libs)): + or metadata.kv_cache_manager is not None): return None inputs = forward_args.sparse_runtime_params.block_sparse_inputs if inputs is None and not forward_args.sage_attn_smooth_v: diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index 9ce1f0472fd5..7228a9b4a79c 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -221,6 +221,11 @@ def _validate_quant_attention_config(self) -> "AttentionConfig": "smooth_k and smooth_v are SageAttention options and do not apply to the " f"per-tensor recipe {recipe!r}." ) + if q_config.smooth_v and sm_version not in (100, 103): + raise ValueError( + "smooth_v needs the PrimTS kernels, which add the V mean back to the " + f"attention output and run on SM100 and SM103 only; got SM{sm_version}." + ) elif q_config.smooth_k or q_config.smooth_v: raise ValueError( f"smooth_k and smooth_v are SageAttention options and require backend='TRTLLM', " diff --git a/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py b/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py index f3f05a4e699a..f8c682f65a6f 100644 --- a/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py +++ b/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py @@ -420,8 +420,10 @@ def _set_sage_recipe(args: AttentionForwardArgs) -> None: def test_attention_core_quantizes_sage_inputs_ahead_of_the_fmha(monkeypatch) -> None: - attention, fmha, q, k, v, metadata, args = _contiguous_case() - attention._fmha_manager = SimpleNamespace(fmha_libs=[fmha]) + attention, _fmha, q, k, v, metadata, args = _contiguous_case() + # The core quantizes whichever libraries the layer has; each library decides whether it + # takes the quantized inputs. + attention._fmha_manager = SimpleNamespace(fmha_libs=[]) quantized = SageQuantized._make(Mock(name=name) for name in SageQuantized._fields) sage_quantize = Mock(return_value=quantized) monkeypatch.setattr(trtllm_backend, "sage_quantize", sage_quantize) @@ -475,13 +477,9 @@ def quantize(): assert quantize() is quantized assert sage_quantize.call_args.kwargs["smooth_v"] - # Calls with a KV cache, and layers without a library that takes quantized inputs, - # quantize nothing. + # Calls with a KV cache quantize nothing. sage_quantize.reset_mock() _set_block_sparse_inputs(args, proxy) - attention._fmha_manager.fmha_libs = [] - assert quantize() is None - attention._fmha_manager.fmha_libs = [fmha] metadata.kv_cache_manager = Mock() assert quantize() is None sage_quantize.assert_not_called() diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py index e6118f5ac277..85d2c29d4478 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py @@ -190,6 +190,30 @@ def test_smoothing_rejected_without_sage_recipe(self, backend, smoothing): ), ) + @pytest.mark.parametrize( + ("sm_ver", "qk_dtype", "q_block_size", "accepted"), + [(90, "int8", 2, False), (100, "int8", 1, True), (103, "fp8", 1, True)], + ) + def test_smooth_v_requires_primts_architecture(self, sm_ver, qk_dtype, q_block_size, accepted): + def build(): + return AttentionConfig( + backend="TRTLLM", + quant_attention_config=QuantAttentionConfig( + qk_dtype=qk_dtype, + q_block_size=q_block_size, + k_block_size=16, + v_block_size=1, + smooth_v=True, + ), + ) + + with patch("tensorrt_llm.visual_gen.args.get_sm_version", return_value=sm_ver): + if accepted: + assert build().quant_attention_config.smooth_v + else: + with pytest.raises(ValidationError, match="smooth_v needs the PrimTS kernels"): + build() + def test_supported_quant_config_cute(self): attention = AttentionConfig( backend="CUTEDSL",