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/cpp/tensorrt_llm/common/attentionOp.cpp b/cpp/tensorrt_llm/common/attentionOp.cpp index 9a452d4bda52..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); @@ -1959,17 +1963,12 @@ 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, "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{}; @@ -1984,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 305f39d86704..b32ef9f75a86 100644 --- a/cpp/tensorrt_llm/common/sageQuant.cu +++ b/cpp/tensorrt_llm/common/sageQuant.cu @@ -37,11 +37,153 @@ 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; + +// 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. +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}; +} + +// 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)); + } +} + +// 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]), VScaleFloor) / E4M3Max; + } +} + +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. +// 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, 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 @@ -52,10 +194,15 @@ 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. -template +// +// 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* 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; @@ -64,9 +211,10 @@ __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); + constexpr float TypeMax = cute::is_same_v ? E4M3Max : static_cast(126.9f); constexpr int BestVL = 128 / sizeof_bits_v; using VL = Int; @@ -117,6 +265,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 * numHeadsV + 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 +294,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 +402,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) + CUTLASS_PRAGMA_NO_UNROLL + for (int chunk = 0; chunk < NumChunks; ++chunk) { - // 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 - { - // 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 +433,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 +443,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,15 +461,71 @@ __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) + if constexpr (VStage == 1 && SmoothV) + { + int const numWarpsToUse = cutlass::fast_min(numWarps, MaxStatWarps); + 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: 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; + float const invNumTokens = 1.0f / (range.tokEnd - range.tokBegin); + CUTLASS_PRAGMA_UNROLL + for (int i = 0; i < BestVL; ++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) { @@ -284,17 +534,17 @@ __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 < 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>{}); 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, @@ -313,13 +563,26 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p else if constexpr (VStage == 2) { int const numToksPerWave = numWarps * numToksPerWarp; + 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); + } 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); - 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)); @@ -331,16 +594,16 @@ __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 rK = make_tensor(Shape{}); Tensor rKCompute = make_tensor(Shape{}); Tensor rKSum = make_tensor(Shape{}); @@ -348,19 +611,18 @@ __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) { 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); 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,11 +642,12 @@ __global__ void sageQuantQkvKernel(int sumSeqLensQk, int batchSize, int const* p } if (threadInTokIdx == thrId) { - float const invNumTokens = 1.0f / sumSeqLensV; + 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); } } } @@ -407,9 +670,32 @@ 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.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) + { + // 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) { @@ -418,23 +704,25 @@ 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.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 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 { @@ -450,33 +738,45 @@ 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); } }; + // 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 +790,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 +802,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); } @@ -510,14 +811,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 @@ -538,6 +848,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..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 @@ -46,7 +47,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}; @@ -56,18 +58,33 @@ 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==0, applied at vStage==1. + // 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{}; 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}; + // 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* ptrVMin{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..0974bc1c1110 --- /dev/null +++ b/cpp/tensorrt_llm/thop/sageQuantizeOp.cpp @@ -0,0 +1,226 @@ +/* + * 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 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 +// 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, 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); + 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"); + } + 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(), + "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 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) + { + k_mean_sum = torch::empty_like(*k_mean, torch::kInt64); + params.ptrKForMean = key.data_ptr(); + params.ptrKMeanSum = k_mean_sum->data_ptr(); + } + std::optional v_mean_sum; + std::optional v_min; + if (v_mean.has_value()) + { + 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. + 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, Tensor(h!)? v_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/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/docs/source/features/visualgen-quantized-attention.md b/docs/source/features/visualgen-quantized-attention.md index 7f81a62d2143..4ecc17b732f1 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 | @@ -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. PrimTS runs on B200/GB200 and B300/GB300 only, so H100/H200/H800/H20 reject `smooth_v`. **Configuration.** diff --git a/docs/source/features/visualgen-sparse-attention.md b/docs/source/features/visualgen-sparse-attention.md index 9ab465921aec..cac4431683d9 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 + 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 @@ -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`; 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 71f400ec5348..9da5468d2b2e 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md @@ -233,6 +233,24 @@ 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 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 +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. 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 the MSA plans. Every layer that runs with one metadata object plans each static diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py index 448d64df9925..0831ace484f8 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 } ) @@ -296,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.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py index 5268e72b5bd6..2bf1554d3753 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,8 @@ 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_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) 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 6d754d7a8747..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 @@ -16,7 +16,7 @@ """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 @@ -29,6 +29,10 @@ 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.logger import logger @@ -79,14 +83,20 @@ 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 + v_dtype: torch.dtype use_kv_valid_bits: bool 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, @@ -99,11 +109,14 @@ 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, + use_block_sparse=self.use_block_sparse, + **capacity, ) except (ValueError, NotImplementedError, OverflowError) as error: return str(error) @@ -129,16 +142,20 @@ 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( + 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, + v_data_type=self.v_dtype, + sage_config=self.sage, ) wrapper.plan(*plan_args, **plan_kwargs) return wrapper @@ -162,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]) @@ -184,8 +206,29 @@ 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.""" + 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=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" @@ -245,7 +288,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" @@ -274,17 +317,29 @@ def _common_unsupported_reason( def _make_plan_key( self, q: torch.Tensor, - inputs: BlockSparseForwardInputs, + inputs: BlockSparseForwardInputs | None, *, batch_size: int, seq_len_q: int, 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) + 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 + v_dtype = torch.float8_e4m3fn return _BlockSparsePlanKey( device=q.device, batch_size=batch_size, @@ -294,14 +349,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, - use_kv_valid_bits=inputs.kv_valid_bits is not None, - sparse_format=inputs.sparse_format, - use_proxy_routes=inputs.use_proxy_routes, + qk_dtype=qk_dtype, + v_dtype=v_dtype, + 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( @@ -332,19 +391,20 @@ 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" 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, @@ -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 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" if ( @@ -521,9 +584,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) @@ -535,22 +597,51 @@ 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: + 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 is not None and inputs.use_proxy_routes: + inputs = replace( + inputs, k_summary=quantized.k_summary, v_summary=quantized.v_summary + ) + sage_params = SageAttentionParams( + q_scale=quantized.q_scale, + 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 + # 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/attention/backends/interface.py b/tensorrt_llm/_torch/attention/backends/interface.py index 7ccf04220604..7558793e0ef0 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,12 @@ 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 # 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 +1005,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/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", diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 909e208129b0..02f16c1484b3 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,74 @@ 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 Sage inputs of block-sparse and V-smoothing calls. + + 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, 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): + return None + inputs = forward_args.sparse_runtime_params.block_sparse_inputs + 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. + 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, + smooth_v=forward_args.sage_attn_smooth_v, + ) + 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 means 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 +2500,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 new file mode 100644 index 000000000000..c3fb9f6a9033 --- /dev/null +++ b/tensorrt_llm/_torch/attention/kernels/sage_quant.py @@ -0,0 +1,187 @@ +# 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. +"""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 + +__all__ = [ + "SageQuantized", + "sage_quantize", +] + + +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``, ``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 + 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 + 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 + + +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") + + +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, + smooth_v: bool = False, + stats: SageQuantized | None = None, +) -> SageQuantized: + """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. 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 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. + + 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. + 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 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 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] + 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, + stats.v_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, + v_mean=stats.v_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 + v_mean = torch.empty_like(v_scale) if smooth_v else None + torch.ops.trtllm.sage_quantize( + query, + key, + value, + cu_seqlens_q, + cu_seqlens_kv, + q, + k, + v, + q_scale, + k_scale, + v_scale, + k_mean, + v_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, + 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 18b68a373e21..93744c2ad66f 100644 --- a/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py @@ -987,6 +987,25 @@ 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], + v_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/sol/backend.py b/tensorrt_llm/_torch/visual_gen/attention/backends/sparse/sol/backend.py index f8070fb9fecd..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 @@ -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 @@ -362,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..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 @@ -143,8 +143,6 @@ def forward( 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" 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 c1f0ef0ef0fc..90b564070728 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 @@ -392,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) @@ -410,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/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.", diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index c1e28236f20b..7228a9b4a79c 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. @@ -152,10 +161,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). @@ -205,15 +216,20 @@ 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 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_k is a SageAttention option and does not apply to the per-tensor " - f"recipe {recipe!r}." + "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: + 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: @@ -276,15 +292,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 e5fcb9c9de8f..88dbd05b0221 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/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..7db0a5da2473 --- /dev/null +++ b/tests/unittest/_torch/attention/kernels/test_sage_quant.py @@ -0,0 +1,400 @@ +# 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. +"""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. +""" + +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 SageQuantized, sage_quantize + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +# --------------------------------------------------------------------------- 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 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 _quantize_summaries( + k: torch.Tensor, + v: torch.Tensor, + 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, v_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, + v_mean: torch.Tensor | None = None, +) -> None: + batch, blocks, heads, _ = k.shape + divisor = 126.9 if quant_dtype == torch.int8 else 448.0 + 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) + 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_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_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) + 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]) +@pytest.mark.parametrize("quant_dtype", [torch.int8, torch.float8_e4m3fn]) +@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, smooth: bool +) -> 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) + + +@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, 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_mismatched_statistics() -> None: + k, v, v_scale, _ = _summary_inputs((2, 9, 3, 64), torch.bfloat16) + cases = ( + ((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(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, 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_mean: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return _quantize_summaries(key, value, value_scale, key_mean, 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, k_mean) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = quantize(k, v, v_scale, k_mean) + k.mul_(2) + v.mul_(2) + graph.replay() + _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]) +# 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: + 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, + smooth_v=smooth_v, + ) + + 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 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_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: + 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) + + +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("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 + + +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"): + sage_quantize( + x, + x, + x, + cu_seqlens, + cu_seqlens, + q_block_size=3, + 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..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 @@ -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,119 @@ 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() + # 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) + + 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, + "smooth_v": False, + } + 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 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 quantize nothing. + sage_quantize.reset_mock() + _set_block_sparse_inputs(args, proxy) + 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_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/sol/test_sol_attention.py b/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_attention.py index 31158399a46b..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 @@ -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" @@ -891,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..74dd6445f7f5 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/sparse_attention/test_sage_block_sparse.py @@ -0,0 +1,448 @@ +# 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, + smooth_v: bool = False, + dtype: torch.dtype = torch.bfloat16, +) -> 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, + smooth_v=smooth_v, + ) + return create_attention( + backend="TRTLLM", + layer_idx=1, + num_heads=_HEADS, + head_dim=_HEAD_DIM, + dtype=dtype, + 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, + 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, smooth_v=smooth_v) + + +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, + 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 + 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 == 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]: + 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", 1), + ("fp8", "exact", 4), + ("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("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") + + +@torch.no_grad() +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("fp8") + 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="SOL Sage CUDA Graph replay") + + +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=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. + 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( + ("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 +) -> None: + """Zero and tiny Q/K keep token masking and empty rows exact through quantization.""" + q, k, v, routes = _masked_zero_score_inputs(dtype) + backends = (_backend(None, dtype=dtype), _backend(qk_dtype, dtype=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, 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" + torch.testing.assert_close(actual[:, 64:128], torch.ones_like(actual[:, 64:128])) + + +@torch.no_grad() +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 + kinds contribute. + """ + q, k, v = _inputs(131) + zero_v = torch.zeros_like(v) + 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" 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_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={} ) 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..85d2c29d4478 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), ) @@ -160,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), @@ -184,17 +178,42 @@ 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} ), ) + @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",