From d5451b00608c9c6020e0f120b2189619b505e1d8 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Fri, 30 Oct 2026 18:14:06 +0000 Subject: [PATCH] [None][chore] Refresh PrimTS source onto FlashInfer main Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../vendor_patches/flashinfer-prims-ts.patch | 449 +++- 3rdparty/vendor_sources.lock.yaml | 8 +- .../attention/backends/fmha/prims_ts.py | 57 +- .../backends/fmha/prims_ts_precompile.py | 6 +- .../backends/fmha/prims_ts_sparse_mla.py | 8 +- .../attention/backends/prims_ts/__init__.py | 12 + .../backends/prims_ts/_block_sparse/common.py | 44 +- .../prims_ts/_block_sparse/compiler.py | 11 +- .../backends/prims_ts/_block_sparse/config.py | 27 +- .../backends/prims_ts/_block_sparse/plan.py | 25 +- .../prims_ts/_block_sparse/runtime.py | 3 +- .../_q_token_kv_block_sparse_policy.py | 64 +- .../backends/prims_ts/block_sparse.py | 44 +- .../attention/backends/prims_ts/context.py | 2076 +++++++++------ .../attention/backends/prims_ts/decode.py | 1958 +++++++------- .../kernels/fmha_context/fmha_kernel.py | 593 ++++- .../kernels/fmha_context/fmha_resources.py | 2293 +++++++++++++++-- .../kernels/fmha_context/fmha_tasks.py | 712 ++++- .../prims_ts/kernels/fmha_context/helpers.py | 144 +- .../kernels/fmha_context/helpers_paged.py | 15 - .../kernels/fmha_context/sage_scales.py | 233 ++ .../fmha_context/vc_attention/__init__.py | 7 + .../fmha_context/vc_attention/vc_resources.py | 258 ++ .../fmha_context/vc_attention/vc_scales.py | 129 + .../fmha_decode/direct_sparse_metadata.py | 106 +- .../kernels/fmha_decode/fmha_decode_config.py | 827 ++++-- .../fmha_decode/fmha_decode_constants.py | 27 +- .../kernels/fmha_decode/fmha_decode_kernel.py | 449 ++-- .../fmha_decode_resources/helpers_common.py | 50 +- .../helpers_kv_tile_idx.py | 75 +- .../fmha_decode_resources/helpers_output.py | 6 +- .../fmha_decode_resources/helpers_softmax.py | 22 + .../fmha_decode_resources/sage_scales.py | 112 +- .../fmha_decode_resources/smem_p.py | 49 +- .../fmha_decode_resources/smem_resources.py | 100 +- .../fmha_decode_resources/tmem_corr.py | 312 ++- .../fmha_decode_resources/tmem_o.py | 4 + .../fmha_decode_resources/tmem_s.py | 413 ++- .../kernels/fmha_decode/fmha_decode_tasks.py | 205 +- .../prims_ts/kernels/fmha_decode/reduction.py | 250 +- .../kernels/mla_decode/helpers/constants.py | 6 +- .../kernels/mla_decode/helpers/diagnostics.py | 55 + .../kernels/mla_decode/helpers/gather.py | 21 + .../kernels/mla_decode/helpers/layout.py | 4 + .../kernels/mla_decode/helpers/math.py | 67 + .../kernels/mla_decode/helpers/nvfp4_kv.py | 759 ++++++ .../kernels/mla_decode/helpers/ops.py | 34 + .../kernels/mla_decode/helpers/query.py | 88 +- .../kernels/mla_decode/helpers/schedule.py | 42 +- .../mla_decode/helpers/sparse_indices.py | 115 +- .../kernels/mla_decode/helpers/tile.py | 107 +- .../kernels/mla_decode/kernel_policy.py | 7 +- .../mla_decode/parallel_reduction_topology.py | 7 +- .../kernels/mla_decode/sparse_policy.py | 361 ++- .../mla_decode/throughput_2cta/config.py | 104 +- .../mla_decode/throughput_2cta/kernel.py | 645 ++++- .../throughput_2cta/nvfp4_kv_load.py | 392 +++ .../throughput_2cta/parallel_reduction.py | 245 +- .../mla_decode/throughput_2cta/reduction.py | 87 +- .../mla_decode/throughput_2cta/resources.py | 645 ++++- .../mla_decode/throughput_2cta/tasks.py | 116 +- .../throughput_2cta/work_partition.py | 9 +- .../throughput_latency_1cta/config.py | 274 +- .../throughput_latency_1cta/kernel.py | 479 +++- .../parallel_reduction.py | 279 +- .../throughput_latency_1cta/reduction.py | 355 ++- .../resources/common.py | 1 + .../resources/smem_p.py | 4 +- .../resources/smem_resources.py | 840 +++++- .../resources/tmem_corr.py | 455 +++- .../resources/tmem_o.py | 54 +- .../resources/tmem_p.py | 51 +- .../resources/tmem_s.py | 205 +- .../resources/tmem_softmax_stats.py | 46 +- .../throughput_latency_1cta/tasks.py | 67 +- .../backends/prims_ts/kernels/tensor_map.py | 11 +- .../attention/backends/prims_ts/mla_decode.py | 1816 ++++++++----- .../q_token_kv_block_sparse_metadata.py | 356 ++- .../attention/backends/prims_ts/sage.py | 480 ---- .../backends/prims_ts/sage/__init__.py | 59 + .../backends/prims_ts/sage/config.py | 267 ++ .../attention/backends/prims_ts/sage/index.py | 134 + .../backends/prims_ts/sage/int32_scores.py | 114 + .../backends/prims_ts/sage/layout.py | 288 +++ .../attention/backends/prims_ts/sage/modes.py | 52 + .../attention/backends/prims_ts/sage/plan.py | 140 + .../sage_prepare.py => sage/prepare.py} | 65 +- .../backends/prims_ts/split_kv_mode_policy.py | 2 +- .../prims_ts/vc_attention/__init__.py | 311 +++ .../backends/prims_ts/vc_attention/prepare.py | 1498 +++++++++++ .../prims_ts/vc_attention/quantization.py | 503 ++++ .../attention/kernels/test_sage_quant.py | 15 +- .../test_prims_ts_sparse_mla_fmha.py | 18 +- .../test_prims_ts_attention_backend.py | 52 +- .../test_prims_ts_context_cuda_scales.py | 71 + .../_torch/attention/test_prims_ts_fmha.py | 50 +- 96 files changed, 19793 insertions(+), 5718 deletions(-) create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/sage_scales.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/vc_attention/__init__.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/vc_attention/vc_resources.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/vc_attention/vc_scales.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/diagnostics.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/nvfp4_kv.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/nvfp4_kv_load.py delete mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/sage.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/sage/__init__.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/sage/config.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/sage/index.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/sage/int32_scores.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/sage/layout.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/sage/modes.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/sage/plan.py rename tensorrt_llm/_torch/attention/backends/prims_ts/{kernels/fmha_decode/sage_prepare.py => sage/prepare.py} (86%) create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/vc_attention/__init__.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/vc_attention/prepare.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/vc_attention/quantization.py create mode 100644 tests/unittest/_torch/attention/test_prims_ts_context_cuda_scales.py diff --git a/3rdparty/vendor_patches/flashinfer-prims-ts.patch b/3rdparty/vendor_patches/flashinfer-prims-ts.patch index f68f66608f7f..67322017b8b2 100644 --- a/3rdparty/vendor_patches/flashinfer-prims-ts.patch +++ b/3rdparty/vendor_patches/flashinfer-prims-ts.patch @@ -194,7 +194,7 @@ index b374a65311402982d818d80e9145ab40ed9727a1..cdfa76294f2d854117db62ad36724254 f"prims_ts_{family}", f"{kind}_sm{major}{minor}_{key}", diff --git a/block_sparse.py b/block_sparse.py -index fa055618705216cddf6b7f3e9ed4c1d680c2af9a..38de4926a045f17c5b6f206f4a20ecbe2532a1c7 100644 +index ea05f2b3e10190fac869bbef3802adcf35418b3e..dc0ffd1d3219c873a1a57e3fe72aa642053c7a99 100644 --- a/block_sparse.py +++ b/block_sparse.py @@ -1,3 +1,4 @@ @@ -217,7 +217,7 @@ index fa055618705216cddf6b7f3e9ed4c1d680c2af9a..38de4926a045f17c5b6f206f4a20ecbe from ._block_sparse.common import ( _num_sparse_pattern_heads, -@@ -320,7 +315,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase): +@@ -322,7 +317,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase): # previously published revision intact and runnable. self._plan_state = candidate @@ -226,7 +226,7 @@ index fa055618705216cddf6b7f3e9ed4c1d680c2af9a..38de4926a045f17c5b6f206f4a20ecbe def run( self, q: torch.Tensor, -@@ -452,7 +447,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase): +@@ -454,7 +449,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase): return self._launch_validated_run(state, run_args, run_stream) @@ -235,7 +235,7 @@ index fa055618705216cddf6b7f3e9ed4c1d680c2af9a..38de4926a045f17c5b6f206f4a20ecbe def block_sparse_attention( q: torch.Tensor, k: torch.Tensor, -@@ -757,9 +752,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase): +@@ -761,9 +756,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase): ) self._plan_state = candidate @@ -246,7 +246,7 @@ index fa055618705216cddf6b7f3e9ed4c1d680c2af9a..38de4926a045f17c5b6f206f4a20ecbe def run( self, q: torch.Tensor, -@@ -881,7 +874,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase): +@@ -885,7 +878,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase): return self._launch_validated_run(state, run_args, run_stream) @@ -256,19 +256,46 @@ index fa055618705216cddf6b7f3e9ed4c1d680c2af9a..38de4926a045f17c5b6f206f4a20ecbe q: torch.Tensor, paged_kv_cache: PagedKVCache, diff --git a/context.py b/context.py -index d3f104e4e34b1bcb7ada697f9e76f7013d6ca100..1ec1080b942f60fad22e3e4a1627cad9592697ad 100644 +index f209da5fc9d515bf29fb7ff4b9aba459cd6165b5..2cdb3a0be100f63fb79fa651342f553bff9e0ea9 100644 --- a/context.py +++ b/context.py -@@ -48,7 +48,7 @@ from ._hardware import get_max_active_clusters - - import torch +@@ -54,7 +54,7 @@ from .vc_attention import ( + vc_scale_shapes, + ) -from flashinfer.api_logging import flashinfer_experimental_api +from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api - - - if TYPE_CHECKING: -@@ -493,7 +493,7 @@ def _validate_device(device: torch.device) -> int: + from .decode import ( + _validate_native_hnd_tma_strides, + prepare_lse, +@@ -184,8 +184,8 @@ class _ContextPlanState: + """Atomically published state for one compiled contiguous context plan.""" + + geometry: _ContextPlanGeometry +- scale_softmax_log2: float +- output_scale: float ++ scale_softmax_log2: float | torch.Tensor ++ output_scale: float | torch.Tensor + empty_i32: torch.Tensor + variable_window_padded_starts: Optional[torch.Tensor] + variable_window_cta_starts: torch.Tensor +@@ -197,6 +197,7 @@ class _ContextPlanState: + lse_mode: str = "none" + use_attention_sinks: bool = False + sage: Optional[SagePreparedScales] = None ++ cuda_scales: bool = False + + + @dataclass(frozen=True) +@@ -284,6 +285,7 @@ class _ContextCompileSpec: + shape_signature: tuple[int, ...] = () + # A Sage plan's prepared scale layout; the kernel reads the scales. + sage_scale_layout: _SageScaleLayout | None = None ++ cuda_scales: bool = False + + + @dataclass(frozen=True) +@@ -564,7 +566,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): @@ -277,8 +304,185 @@ index d3f104e4e34b1bcb7ada697f9e76f7013d6ca100..1ec1080b942f60fad22e3e4a1627cad9 require_cute_dsl_arch(device_index) return device_index +@@ -1892,6 +1894,7 @@ def _context_compile_spec( + runtime_lengths: bool | None = None, + lse_mode: str = "none", + use_attention_sinks: bool = False, ++ cuda_scales: bool = False, + ) -> _ContextCompileSpec: + """Keep resource/precision choices static and ordinary lengths runtime.""" + fixed_seqlen = geometry.sage is not None or geometry.vc is not None +@@ -1926,6 +1929,7 @@ def _context_compile_spec( + runtime_lengths=runtime_lengths, + lse_mode=lse_mode, + use_attention_sinks=use_attention_sinks, ++ cuda_scales=cuda_scales, + scheduler=_resolve_context_scheduler(geometry), + two_cta_umma=_two_cta_umma_geometry_eligible(geometry), + vc=geometry.vc, +@@ -2094,7 +2098,7 @@ def _get_compiled_context( + out: cute.Tensor, + lse: cute.Tensor | None, + attention_sinks: cute.Tensor | None, +- scale_softmax_log2: cutlass.Float32, ++ scale_softmax_log2: cutlass.Float32 | cute.Tensor, + output_scale: cutlass.Float32 | cute.Tensor, + qo_indptr: cute.Tensor, + kv_indptr: cute.Tensor, +@@ -2206,6 +2210,9 @@ def _get_compiled_context( + ) + scale_fake = cutlass.Float32(1.0) + output_scale_fake = cutlass.Float32(1.0) ++ if compile_spec.cuda_scales: ++ scale_fake = fake_compact(cutlass.Float32, (1,), 4) ++ output_scale_fake = fake_compact(cutlass.Float32, (1,), 4) + if vc_attention: + output_scale_fake = fake_compact( + cutlass.Float32, (batch_size, num_qo_heads, head_dim_vo), 64 +@@ -2868,6 +2875,15 @@ def _variable_window_buffers(geometry, empty_i32): + return padded_starts, cta_starts + + ++def _cuda_scale(value: torch.Tensor, name: str, device: torch.device) -> torch.Tensor: ++ """Validate a live device scalar without reading it back to the host.""" ++ if not isinstance(value, torch.Tensor): ++ raise TypeError(f"{name} must be a CUDA FP32 tensor") ++ if value.device != device or value.dtype != torch.float32 or value.numel() != 1: ++ raise ValueError(f"{name} must contain one FP32 value on {device}") ++ return value.reshape(1) ++ ++ + class BatchPrefillTSWrapper: + """Reusable contiguous context attention with eager planning. + +@@ -2909,6 +2925,7 @@ class BatchPrefillTSWrapper: + lse_mode: Literal["none", "base2"] = "none", + use_attention_sinks: bool = False, + sage: Optional[SageAttentionConfig] = None, ++ cuda_scales: bool = False, + ) -> None: + """Select/cache one execution policy, obtain compiled kernels, and bind resources. + +@@ -2928,6 +2945,8 @@ class BatchPrefillTSWrapper: + a per-run sink binding. Context scales are positive host scalar defaults. + Sage plans use BF16 output by default and require per-run Sage scales. + Sage V-mean restoration cannot be combined with attention sinks. ++ cuda_scales binds one-element CUDA FP32 scales for graph-replayed ++ per-tensor quantization; it cannot be combined with Sage or VC plans. + Equivalent plans reuse policy and compiled-kernel caches. + VC runs sharing a plan must not overlap. + """ +@@ -2981,6 +3000,10 @@ class BatchPrefillTSWrapper: + raise NotImplementedError( + "VC-Attention does not support LSE or attention sinks" + ) ++ if not isinstance(cuda_scales, bool): ++ raise TypeError("cuda_scales must be a bool") ++ if cuda_scales and (vc_config is not None or sage is not None): ++ raise ValueError("cuda_scales requires a plain context plan") + geometry = replace(geometry, vc=vc_config) + if softmax_scale is None: + softmax_scale = geometry.head_dim**-0.5 +@@ -2989,6 +3012,13 @@ class BatchPrefillTSWrapper: + "softmax_scale * log2(e)", + ) + output_scale = _validate_scale(output_scale, "output_scale") ++ if cuda_scales: ++ scale_softmax_log2 = torch.tensor( ++ [scale_softmax_log2], dtype=torch.float32, device=geometry.device ++ ) ++ output_scale = torch.tensor( ++ [output_scale], dtype=torch.float32, device=geometry.device ++ ) + vc_output_scale_buffer = ( + torch.empty( + ( +@@ -3031,6 +3061,7 @@ class BatchPrefillTSWrapper: + runtime_lengths=vc_config is None, + lse_mode=lse_mode, + use_attention_sinks=use_attention_sinks, ++ cuda_scales=cuda_scales, + ) + ) + self._plan_state = _ContextPlanState( +@@ -3044,6 +3075,7 @@ class BatchPrefillTSWrapper: + policy=policy, + vc_scale_shapes=vc_shapes, + vc_output_scale=vc_output_scale_buffer, ++ cuda_scales=cuda_scales, + lse_mode=lse_mode, + use_attention_sinks=use_attention_sinks, + sage=sage_state, +@@ -3068,7 +3100,8 @@ class BatchPrefillTSWrapper: + return_lse: bool = False, + attention_sinks: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, +- output_scale: Optional[float] = None, ++ output_scale: float | torch.Tensor | None = None, ++ scale_softmax_log2: torch.Tensor | None = None, + vc: Optional[VCAttentionParams] = None, + sage: Optional[SageAttentionParams] = None, + validate: bool = True, +@@ -3095,6 +3128,9 @@ class BatchPrefillTSWrapper: + wrapper-owned fallback scratch. Scale overrides are finite positive host + scalars: softmax_scale is the effective QK scale, and + output_scale scales V/output. They are runtime kernel operands, not JIT keys. ++ With cuda_scales, supply scale_softmax_log2 and output_scale as ++ one-element CUDA FP32 tensors instead. Their values remain on device ++ and must be positive and finite, including after graph updates. + + Enabled attention_sinks supplies compact CUDA FP32 [Hq] natural-log sink + logits after QK scaling. Values must be finite or -inf; validation checks only +@@ -3234,19 +3270,35 @@ class BatchPrefillTSWrapper: + runtime_window_starts = state.empty_i32 + runtime_window_ends = state.empty_i32 + runtime_window_cta_starts = state.empty_i32 +- scale_softmax_log2 = ( +- state.scale_softmax_log2 +- if softmax_scale is None +- else _validate_scale( +- _validate_scale(softmax_scale, "softmax_scale") * math.log2(math.e), +- "softmax_scale * log2(e)", ++ if state.cuda_scales: ++ if softmax_scale is not None: ++ raise ValueError("cuda_scales uses the device scale_softmax_log2 operand") ++ scale_softmax_log2 = _cuda_scale( ++ state.scale_softmax_log2 if scale_softmax_log2 is None else scale_softmax_log2, ++ "scale_softmax_log2", ++ geometry.device, ++ ) ++ output_scale = _cuda_scale( ++ state.output_scale if output_scale is None else output_scale, ++ "output_scale", ++ geometry.device, ++ ) ++ else: ++ if scale_softmax_log2 is not None: ++ raise ValueError("device scale_softmax_log2 requires cuda_scales=True") ++ scale_softmax_log2 = ( ++ state.scale_softmax_log2 ++ if softmax_scale is None ++ else _validate_scale( ++ _validate_scale(softmax_scale, "softmax_scale") * math.log2(math.e), ++ "softmax_scale * log2(e)", ++ ) ++ ) ++ output_scale = ( ++ state.output_scale ++ if output_scale is None ++ else _validate_scale(output_scale, "output_scale") + ) +- ) +- output_scale = ( +- state.output_scale +- if output_scale is None +- else _validate_scale(output_scale, "output_scale") +- ) + if vc_output_scale is not None: + output_scale = torch.mul( + vc_output_scale, output_scale, out=state.vc_output_scale diff --git a/decode.py b/decode.py -index 4fb096f39849a3801349758360e71b092ee200e9..f91814650c48dcf5b116273742880ac722fed527 100644 +index cd2472da75685f073a78177c6f0e6ee8728f7274..0609b55f408f2d62e072fd11463ff01c6f1f6d75 100644 --- a/decode.py +++ b/decode.py @@ -1,3 +1,4 @@ @@ -286,20 +490,26 @@ index 4fb096f39849a3801349758360e71b092ee200e9..f91814650c48dcf5b116273742880ac7 # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); -@@ -26,11 +27,7 @@ import torch +@@ -26,17 +27,13 @@ import torch from ._compile_cache import build_and_load_attention_kernel -from flashinfer.api_logging import flashinfer_experimental_api ++from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api + from flashinfer.utils import ( + get_device_sm_count, + _check_workspace_buffer_alignment, + check_shape_dtype_device, + ) + -from flashinfer.trace.templates.attention import ( - attention_ts_decode_trace_dispatch, - prims_ts_decode_wrapper_trace_dispatch, -) -+from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api from . import _q_token_kv_block_sparse_policy as _sparse_policy from ._block_sparse.common import _num_sparse_pattern_heads -@@ -780,7 +777,7 @@ def _validate_runtime_device(device: torch.device) -> None: +@@ -770,7 +767,7 @@ def _validate_runtime_device(device: torch.device) -> None: # Rubin runs through the sm_100f family target; a CuTe DSL older than 4.8 # cannot emit for it unless CUTE_DSL_ARCH=sm_100f is set before import. if capability == (10, 7): @@ -308,7 +518,7 @@ index 4fb096f39849a3801349758360e71b092ee200e9..f91814650c48dcf5b116273742880ac7 require_cute_dsl_arch(device_index) -@@ -4075,7 +4072,7 @@ class BatchDecodePagedTSWrapper: +@@ -4281,7 +4278,7 @@ class BatchDecodePagedTSWrapper: # previous complete plan revision usable. self._plan_state = candidate @@ -317,8 +527,8 @@ index 4fb096f39849a3801349758360e71b092ee200e9..f91814650c48dcf5b116273742880ac7 def run( self, q: torch.Tensor, -@@ -4257,7 +4254,7 @@ class BatchDecodePagedTSWrapper: - ) +@@ -4481,7 +4478,7 @@ class BatchDecodePagedTSWrapper: + return (result, lse) if return_lse else result -@flashinfer_experimental_api(trace=attention_ts_decode_trace_dispatch) @@ -326,8 +536,46 @@ index 4fb096f39849a3801349758360e71b092ee200e9..f91814650c48dcf5b116273742880ac7 def batch_decode_with_paged_kv_cache( q: torch.Tensor, paged_kv_cache: PagedKVCache, +diff --git a/kernels/fmha_context/fmha_kernel.py b/kernels/fmha_context/fmha_kernel.py +index effc3d218a66309e48e3bbd9b31923c4cd02765a..d447b52846fa6ad8aaebf635c72cceb8973d6fda 100644 +--- a/kernels/fmha_context/fmha_kernel.py ++++ b/kernels/fmha_context/fmha_kernel.py +@@ -1,3 +1,4 @@ ++# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + # Copyright (c) 2026 by FlashInfer team. + # + # Licensed under the Apache License, Version 2.0 (the "License"); +@@ -3317,7 +3318,7 @@ class FmhaTs: + k_cute: cute.Tensor, + v_cute: cute.Tensor, + o_cute: cute.Tensor, +- scale_softmax_log2: Float32, ++ scale_softmax_log2: Float32 | cute.Tensor, + output_scale: Float32 | cute.Tensor, + max_active_clusters: int, + stream: cuda_drv.CUstream, +@@ -3751,7 +3752,7 @@ class FmhaTs: + | utils.ClcDynamicPersistentTileSchedulerParams + | None + ), +- scale_softmax_log2: Float32, ++ scale_softmax_log2: Float32 | cute.Tensor, + output_scale: Float32 | cute.Tensor, + num_kv_tiles: Int32, + num_seq_tiles: Int32, +@@ -3797,6 +3798,10 @@ class FmhaTs: + + cfg = copy(cfg) + cfg.fixed_dense_k_tail = dense_tail_capacity ++ if cutlass.const_expr(isinstance(scale_softmax_log2, cute.Tensor)): ++ scale_softmax_log2 = scale_softmax_log2[0] ++ if cutlass.const_expr(isinstance(output_scale, cute.Tensor) and not cfg.vc_attention): ++ output_scale = output_scale[0] + cfg.num_seq_tiles = num_seq_tiles + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + diff --git a/kernels/fmha_context/fmha_resources.py b/kernels/fmha_context/fmha_resources.py -index d17584b725e24235ce0ec90caac174a29d4ab965..af69ee7d51bc1179352d0bf2ba71fad2fe30a922 100644 +index 203cc25fb3814fa0fc793894a26d717515f31f76..ba7273fcea22713387a6234e9b935c028416fd16 100644 --- a/kernels/fmha_context/fmha_resources.py +++ b/kernels/fmha_context/fmha_resources.py @@ -1,3 +1,4 @@ @@ -335,7 +583,7 @@ index d17584b725e24235ce0ec90caac174a29d4ab965..af69ee7d51bc1179352d0bf2ba71fad2 # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); -@@ -1199,8 +1200,6 @@ class GmemQKVResource(MemoryResource): +@@ -1424,8 +1425,6 @@ class GmemQKVResource(MemoryResource): if cutlass.const_expr(self.cfg.has_uniform_varlen): seqlen_k = Int32(self.cfg.uniform_seq_len_k) else: @@ -344,7 +592,7 @@ index d17584b725e24235ce0ec90caac174a29d4ab965..af69ee7d51bc1179352d0bf2ba71fad2 seqlen_k = _load_runtime_seq_len_kv( self.seqlens_kv, self.max_seq_len_kv, batch_coord ) -@@ -1223,8 +1222,6 @@ class GmemQKVResource(MemoryResource): +@@ -1448,8 +1447,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: @@ -353,7 +601,7 @@ index d17584b725e24235ce0ec90caac174a29d4ab965..af69ee7d51bc1179352d0bf2ba71fad2 kv_request_begin, kv_page_idx_ub = _load_block_table_row_bounds( Int32(self.block_table_row_stride), self.cfg, -@@ -1300,13 +1297,9 @@ class GmemQKVResource(MemoryResource): +@@ -1536,13 +1533,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: @@ -367,25 +615,7 @@ index d17584b725e24235ce0ec90caac174a29d4ab965..af69ee7d51bc1179352d0bf2ba71fad2 kv_request_begin, kv_page_idx_ub = _load_block_table_row_bounds( Int32(self.block_table_row_stride), self.cfg, -@@ -1743,8 +1736,6 @@ class SmemPageOffsetsKvResource(MemoryResource): - kv_request_begin: Int32, - kv_page_idx_ub: Int32, - ) -> None: -- from .helpers_paged import _resolve_kv_tile_idx_context -- - cfg = self.cfg - # Context's K/V tile index is kv_tile_start + loop_offset; for V the - # producer runs one tile ahead in TAIL, so reuse the same expression -@@ -1835,8 +1826,6 @@ class SmemPageOffsetsKvResource(MemoryResource): - tile_offset: cutlass.Constexpr[int] = 0, - ) -> cutlass.Array: - """Retain one tile's page IDs after its SMEM window is released.""" -- from .helpers_paged import _resolve_kv_tile_idx_context -- - tile_idx = _resolve_kv_tile_idx_context( - stage_info, kv_tile_start, tile_offset=tile_offset - ) -@@ -5770,3 +5759,14 @@ class GmemOResource(MemoryResource): +@@ -7616,3 +7609,13 @@ 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) @@ -398,7 +628,6 @@ index d17584b725e24235ce0ec90caac174a29d4ab965..af69ee7d51bc1179352d0bf2ba71fad2 +from .helpers_paged import ( + _load_block_table_row_bounds, + _load_runtime_seq_len_kv, -+ _resolve_kv_tile_idx_context, +) diff --git a/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py b/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py index b22f0f413b7312bab9a5f8e296255fe9aeabe901..a4a87200ca5d3ace5cf899caa9672a1a32c5289e 100644 @@ -419,7 +648,7 @@ index b22f0f413b7312bab9a5f8e296255fe9aeabe901..a4a87200ca5d3ace5cf899caa9672a1a from .helpers_common import ( _TASK_CACHE_TMEM_BASE_OFFSET, diff --git a/mla_decode.py b/mla_decode.py -index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e231c98a52 100644 +index a42843f6ab47c0b1cefa0a0b51e27b355e0a4974..e60454d40ec7caf98e2bd5e70ebc951bd1e94fa2 100644 --- a/mla_decode.py +++ b/mla_decode.py @@ -1,3 +1,4 @@ @@ -427,28 +656,26 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); -@@ -27,11 +28,7 @@ from ._hardware import get_max_active_clusters - - import torch +@@ -33,12 +34,8 @@ from .kernels.mla_decode.helpers.nvfp4_kv import ( + validate_nvfp4_kv_formats, + ) -from flashinfer.api_logging import flashinfer_experimental_api ++from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api + from flashinfer.utils import get_device_sm_count -from flashinfer.trace.templates.attention import ( - prims_ts_decode_mla_one_shot_trace_dispatch, - prims_ts_decode_mla_wrapper_trace_dispatch, -) -+from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api + from .decode import ( - _WorkspaceSection, -@@ -1047,11 +1044,35 @@ def _make_mla_decode_compile_spec( - ) +@@ -1185,8 +1182,34 @@ def _make_mla_decode_compile_spec( + _cached_mla_decode_compile_spec = functools.cache(_make_mla_decode_compile_spec) -@functools.cache -def _get_compiled_mla_decode( -- compile_spec: _MLADecodeCompileSpec, --): -- """Compile and cache one batch-dynamic TS MLA topology.""" +_COMPILED_MLA_DECODE: dict = {} + + @@ -476,31 +703,12 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 + return is_placeholder(compiled) + + -+def _compile_mla_decode(compile_spec: _MLADecodeCompileSpec): -+ """Build (or load from the on-disk cache) one batch-dynamic TS MLA topology.""" - - import cutlass - import cutlass.cute as cute -@@ -1459,11 +1480,12 @@ def _launch_mla_decode( - runtime.sparse_inputs, - runtime.bmm1_scale, - runtime.bmm2_scale, -- **( -- {"output_quant_args": runtime.output_quant_args} -- if runtime.output_quant_args is not None -- else {} -- ), -+ # Passed positionally and also when None. The kernel was compiled with -+ # ``output_quant_args`` in its signature (``None`` for unquantized -+ # output), and a kernel reloaded from the persistent compile cache is -+ # a positional TVM-FFI entry with no Python defaults: it rejects both -+ # a missing parameter and a keyword argument. -+ runtime.output_quant_args, - ) - return runtime.out - -@@ -1637,7 +1659,7 @@ class BatchMLADecodePagedTSWrapper: - split_kv=int(dict(policy)["split_kv"]), ++def _compile_mla_decode( + compile_spec: _MLADecodeCompileSpec, + ): + """Compile and cache one batch-dynamic TS MLA topology.""" +@@ -1872,7 +1895,7 @@ class BatchMLADecodePagedTSWrapper: + softmax_rescale_threshold_log2=softmax_rescale_threshold_log2, ) - @flashinfer_experimental_api(trace=prims_ts_decode_mla_wrapper_trace_dispatch) @@ -508,8 +716,8 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 def run( self, query: torch.Tensor, -@@ -1742,7 +1764,7 @@ class BatchMLADecodePagedTSWrapper: - ) +@@ -1979,7 +2002,7 @@ class BatchMLADecodePagedTSWrapper: + return (result, lse) if return_lse else result -@flashinfer_experimental_api(trace=prims_ts_decode_mla_one_shot_trace_dispatch) @@ -517,7 +725,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 def batch_mla_decode_with_paged_kv_cache( query: torch.Tensor, kv_cache: torch.Tensor, -@@ -2099,6 +2121,7 @@ def _prepare_sparse_output_quant( +@@ -2343,6 +2366,7 @@ def _prepare_sparse_output_quant( ) @@ -525,7 +733,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 def _resolve_sparse_mla_plan( device, batch_size, -@@ -2119,17 +2142,17 @@ def _resolve_sparse_mla_plan( +@@ -2368,17 +2392,17 @@ def _resolve_sparse_mla_plan( ): """Resolve sparse geometry and scratch sizing without allocation/compilation.""" import cutlass @@ -547,40 +755,36 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 MlaDecodeTs, ) -@@ -2264,6 +2287,9 @@ def _resolve_sparse_mla_plan( - return device, profile, spec, sections, byte_end, capacity, output_quant +@@ -2754,6 +2778,9 @@ def resolve_sparse_d576_reload_plan( + return device, profile, spec, sections, byte_end, capacity, None -+# Prepared launches kept per plan (one per layer sharing the plan; DeepSeek-V4 has 61). +_PREPARED_LAUNCH_LIMIT = 256 ++ + class BatchSparseMLADecodePagedTSWrapper: """Plan native BF16/E4M3 attention, then bind one or two source index lists. -@@ -2297,6 +2323,11 @@ class BatchSparseMLADecodePagedTSWrapper: +@@ -2803,6 +2830,7 @@ class BatchSparseMLADecodePagedTSWrapper: + def __init__(self, workspace_buffer=None): self._workspace_buffer = workspace_buffer self._state = None - self._scalar_cache = {} -+ # Prepared launches for the unfused validate=False path: everything that -+ # is static for one (plan, pools, scales, sinks) tuple, keyed by the pools' -+ # storage and the sinks address. One plan serves every layer of a model, -+ # and each layer brings its own sinks, so one entry per layer. + self._prepared: dict[tuple, dict] = {} self.workspace_size_bytes = 0 def plan( -@@ -2386,6 +2417,7 @@ class BatchSparseMLADecodePagedTSWrapper: - kernel_signature=_mla_kernel_compile_signature(static_kernel), +@@ -2964,6 +2992,7 @@ class BatchSparseMLADecodePagedTSWrapper: + kernel_signature=_mla_kernel_compile_signature(kernel), ) - ) + compiled = _get_compiled_mla_decode(spec) + self._prepared.clear() self._state = dict( + reload_d576=reload_d576, device=device, - batch=batch_size, -@@ -2608,6 +2640,38 @@ class BatchSparseMLADecodePagedTSWrapper: - state = self._state - if state is None: +@@ -3210,6 +3239,40 @@ class BatchSparseMLADecodePagedTSWrapper: raise RuntimeError("plan() must be called before run()") + if softmax_scale is None: + softmax_scale = state["query_dim"] ** -0.5 + lse_arg = lse + prepared_key = ( + kv_cache.data_ptr(), @@ -595,7 +799,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 + extra_kv_cache.stride(), + extra_kv_cache.dtype, + ), -+ None if sinks is None else sinks.data_ptr(), ++ None if attention_sinks is None else attention_sinks.data_ptr(), + softmax_scale, + output_scale, + ) @@ -607,6 +811,8 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 + and token_positions is None + and cos_sin_cache is None + and out_scale is None ++ and type(kv_cache_scale) in (int, float) ++ and kv_cache_scale == 1.0 + ): + prepared = self._prepared.get(prepared_key) + if prepared is not None: @@ -615,18 +821,20 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 + ) _validate_query( query, - packed_query=state["packed"], -@@ -2781,10 +2845,124 @@ class BatchSparseMLADecodePagedTSWrapper: - if static_scales - else state["compiled"], + packed_query=state["packed_query"], +@@ -3405,10 +3468,127 @@ class BatchSparseMLADecodePagedTSWrapper: + compiled=state["compiled"], + max_q_len=1, ) + if ( + rows + and quant is None -+ and static_scales -+ and not state["packed"] ++ and not state["packed_query"] + and not state["return_lse"] -+ and state["max_q"] == 1 ++ and state["max_q_len"] == 1 ++ and state["query_dim"] == 512 ++ and state["kv_cache_format"] == "native" ++ and state["extra_kv_cache_format"] == "native" + and lse_arg is None + ): + # Static for this (plan, pools, scales): pool views, scalars, sinks, split, @@ -640,10 +848,10 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 + primary_rows=primary.shape[0], + c_latent=primary.permute(1, 2, 0), + c_rope=extra.permute(1, 2, 0), -+ scalars=(scale_tensors[0], scale_tensors[1], sinks), -+ bmm=(bmm1_scale, bmm2_scale), -+ ks=state["ks"], -+ kc=state["kc"], ++ sinks=attention_sinks, ++ bmm=(softmax_scale, output_scale), ++ ks=state["max_topk"], ++ kc=state["max_extra_topk"], + dummy_indices=state["dummy_indices"], + default_cl=state["default_cl"], + ) @@ -661,10 +869,10 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 + full path for the same arguments. + """ + state = self._state -+ rows, heads, device = state["batch"], state["heads"], state["device"] ++ rows, heads, device = state["batch_size"], state["num_qo_heads"], state["device"] + if ( + query.shape != (rows, 1, heads, 512) -+ or query.dtype != state["dtype"] ++ or query.dtype != state["q_dtype"] + or query.device != device + or not query.is_contiguous() + or query.data_ptr() % 16 @@ -715,11 +923,11 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 + raise ValueError("prepared launch: extra indices without a planned extra pool") + ci, cl = prepared["dummy_indices"][:rows], prepared["default_cl"][:rows] + buffers = state["buffers"] -+ s0, s1, sinks = prepared["scalars"] ++ sinks = prepared["sinks"] + q_kernel = query.permute(2, 3, 1, 0) + lse_kernel = buffers["public_lse"][:rows].view(rows, 1, heads).permute(2, 1, 0) + bmm1, bmm2 = prepared["bmm"] -+ state["compiled_static"]( ++ state["compiled"]( + q_kernel, + q_kernel, + prepared["c_latent"], @@ -728,13 +936,14 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 + out.view(rows, 1, heads, 512).permute(2, 3, 1, 0), + lse_kernel, + buffers["core"] if buffers["core"].numel() else None, -+ state["splits"], ++ state["requested_splits"], + lengths, + None, -+ (indices, ci, lengths, cl, s0, s1, sinks), ++ (indices, ci, lengths, cl, sinks, 1.0), + bmm1, + bmm2, + None, ++ 1, + ) + return out + @@ -742,7 +951,7 @@ index b62bd1081fbae0e18e0950008bae87c8986677b4..7a305b1ec6f2e860b3627e6d71c562e2 def get_prims_ts_sparse_mla_decode_workspace_size(*plan_args, **plan_kwargs): """Return workspace bytes using the same arguments as wrapper.plan(). diff --git a/q_token_kv_block_sparse_metadata.py b/q_token_kv_block_sparse_metadata.py -index e2bbef13c0f552030875e6b433e0fcdb20686cf6..dff16046cf7cf9ff03950043cf462ac28de6c670 100644 +index 146921278514aa664590b95362a9533a287b145d..142cdb1ec83a9471fce26398fedba94312c1a76c 100644 --- a/q_token_kv_block_sparse_metadata.py +++ b/q_token_kv_block_sparse_metadata.py @@ -1,3 +1,4 @@ @@ -750,7 +959,7 @@ index e2bbef13c0f552030875e6b433e0fcdb20686cf6..dff16046cf7cf9ff03950043cf462ac2 # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); -@@ -41,9 +42,9 @@ from dataclasses import dataclass +@@ -41,9 +42,9 @@ from dataclasses import dataclass, replace from typing import Any, Literal, Optional import torch diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 18640486388d..71bc9d030900 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -14,15 +14,15 @@ vendors: digest: sha256-tree-v1:d2d382335ccb68b55078523e63835c28b699607227afb81308cb7bb1d85845ac flashinfer-prims-ts: url: https://github.com/yuxianq/flashinfer.git - branch: trtllm-prims-ts-dev-20261012 - commit: aa709061230a58858375c513056e7733fa74294a + branch: trtllm-prims-ts-dev-20261030 + commit: d1bed93d69d172fd8f9af96cf14f4f869ac488ef 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:fc5c23ec3d65c562025a1845f211fdac66bfcd3da156dbb8385f4c4407d20d07 - digest: sha256-tree-v1:f2f444f7e6d4ad688fbf4e1e9eed11ff7c78d835cfd0af5ef9d7660780b059e0 + patch_digest: sha256:b41e7a8c0f5eef13fd7b3284402ce8c489a544e590883500735a9b757a29e403 + digest: sha256-tree-v1:741bf149276a18236e5b91cde228626a001e8c66a71a275d0db4128cd254e0df openengine-proto: url: https://github.com/ai-dynamo/openengine.git tag: v0.1.0 diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py index 89ee7d2a8126..159f5ac1b96d 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py @@ -672,12 +672,12 @@ def _get_or_plan_context_wrapper( v_cache: torch.Tensor, *, batch_size: int, - max_seq_len_q: int, + max_q_len: int, max_kv_len: int, page_size: int, mask_type: str, window_left: int, - sm_scale: float, + softmax_scale: float, output_dtype: torch.dtype, ) -> "BatchPrefillPagedTSWrapper": wrapper = self._context_wrappers.get(batch_size) @@ -693,7 +693,7 @@ def _get_or_plan_context_wrapper( wrapper.plan( device=q.device, batch_size=batch_size, - max_seq_len_q=max_seq_len_q, + max_q_len=max_q_len, max_kv_len=max_kv_len, num_qo_heads=int(q.shape[-2]), num_kv_heads=int(k_cache.shape[1]), @@ -705,7 +705,7 @@ def _get_or_plan_context_wrapper( page_size=page_size, mask_type=mask_type, window_left=window_left, - sm_scale=sm_scale, + softmax_scale=softmax_scale, output_scale=1.0, # Preprocessing establishes this invariant for every context # request admitted by this adapter. @@ -749,11 +749,11 @@ def _get_or_plan_decode_wrapper( head_dim, page_size, max_kv_len, - max_seq_len_q=seq_len_q, + max_q_len=seq_len_q, packed_query=False, - q_data_type=q_dtype, - kv_data_type=kv_dtype, - o_data_type=output_dtype, + q_dtype=q_dtype, + kv_dtype=kv_dtype, + out_dtype=output_dtype, mask_type=mask_type, window_left=window_left, workspace_buffer=workspace_buffer, @@ -770,7 +770,7 @@ def _get_or_plan_mla_decode_wrapper( kv_lora_rank: int, qk_rope_head_dim: int, page_size: int, - max_seq_len_q: int, + max_q_len: int, max_kv_len: int, q_dtype: torch.dtype, kv_dtype: torch.dtype, @@ -797,11 +797,11 @@ def _get_or_plan_mla_decode_wrapper( qk_rope_head_dim, page_size, max_kv_len, - max_seq_len_q=max_seq_len_q, + max_q_len=max_q_len, packed_query=False, - q_data_type=q_dtype, - kv_data_type=kv_dtype, - o_data_type=output_dtype, + q_dtype=q_dtype, + kv_dtype=kv_dtype, + out_dtype=output_dtype, mask_type=mask_type, workspace_buffer=workspace_buffer, ) @@ -850,13 +850,15 @@ def _forward_without_kv_cache( # Per-tensor E4M3 operands as selected by the recipe. The Q and K dequant scales # fold into the per-run base-2 softmax scale and the V dequant scale into the # kernel output scale. - sm_scale = self._get_bmm1_scale(attn) + softmax_scale = self._get_bmm1_scale(attn) scale_softmax_log2 = None if forward_args.qk_fp8_per_tensor: q_bshd, q_dequant_scale = torch.ops.tensorrt_llm.quantize_e4m3_per_tensor(q_bshd) k_bshd, k_dequant_scale = torch.ops.tensorrt_llm.quantize_e4m3_per_tensor(k_bshd) scale_softmax_log2 = ( - q_dequant_scale.float() * k_dequant_scale.float() * (sm_scale * math.log2(math.e)) + q_dequant_scale.float() + * k_dequant_scale.float() + * (softmax_scale * math.log2(math.e)) ).reshape(1) output_scale = None if forward_args.v_fp8_per_tensor: @@ -881,7 +883,7 @@ def _forward_without_kv_cache( wrapper.plan( device=q.device, batch_size=batch_size, - max_seq_len_q=seq_len, + max_q_len=seq_len, max_kv_len=seq_len, num_qo_heads=num_heads, num_kv_heads=num_kv_heads, @@ -891,7 +893,8 @@ def _forward_without_kv_cache( v_dtype=v_bshd.dtype, out_dtype=out.dtype, mask_type=mask_type, - sm_scale=sm_scale, + softmax_scale=softmax_scale, + cuda_scales=forward_args.qk_fp8_per_tensor or forward_args.v_fp8_per_tensor, ) self._contiguous_context_wrappers[key] = wrapper wrapper.run( @@ -1000,7 +1003,7 @@ def prepare_workspace( int(self.attn.qk_rope_head_dim), int(metadata.tokens_per_block), max_seq_len, - max_seq_len_q=seq_len_q, + max_q_len=seq_len_q, q_dtype=kernel_dtype, kv_dtype=kernel_dtype, out_dtype=forward_args.output.dtype, @@ -1177,7 +1180,7 @@ def run_context(self, params: FmhaParams) -> None: if kv_page_offset is None: raise RuntimeError("PrimTS could not resolve the K-to-V page displacement.") k_cache, v_cache = self._standard_kv_views(kv_pool, kv_page_offset) - max_seq_len_q = int(meta.max_context_length) + max_q_len = int(meta.max_context_length) max_kv_len = int(meta.max_seq_len) fixed_block_tables = self._get_fixed_block_tables(block_tables, params.batch_size) seq_lens_kv = self._get_sequence_lengths( @@ -1190,12 +1193,12 @@ def run_context(self, params: FmhaParams) -> None: k_cache, v_cache, batch_size=params.batch_size, - max_seq_len_q=max_seq_len_q, + max_q_len=max_q_len, max_kv_len=max_kv_len, page_size=params.tokens_per_block, mask_type=mask_type, window_left=window_left, - sm_scale=self._get_bmm1_scale(attn), + softmax_scale=self._get_bmm1_scale(attn), output_dtype=params.output.dtype, ) wrapper.run( @@ -1396,8 +1399,8 @@ def run_generation(self, params: FmhaParams) -> None: (k_cache, v_cache), seq_lens, block_tables=fixed_block_tables, - bmm1_scale=self._get_bmm1_scale(attn), - bmm2_scale=1.0, + softmax_scale=self._get_bmm1_scale(attn), + output_scale=1.0, out=output, validate=False, ) @@ -1480,7 +1483,7 @@ def run_mla_generation(self, params: FmhaParams) -> None: kv_lora_rank=int(attn.kv_lora_rank), qk_rope_head_dim=int(attn.qk_rope_head_dim), page_size=params.tokens_per_block, - max_seq_len_q=seq_len_q, + max_q_len=seq_len_q, max_kv_len=max_seq_len, q_dtype=query.dtype, kv_dtype=kv_cache.dtype, @@ -1491,10 +1494,10 @@ def run_mla_generation(self, params: FmhaParams) -> None: query, kv_cache, block_tables=fixed_block_tables, - seq_lens=seq_lens, + seq_lens_kv=seq_lens, out=output, - bmm1_scale=bmm1_scale, - bmm2_scale=bmm2_scale, + softmax_scale=bmm1_scale, + output_scale=bmm2_scale, validate=False, ) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_precompile.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_precompile.py index f293388f6a1f..770758751a1a 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_precompile.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_precompile.py @@ -132,8 +132,8 @@ def strict() -> bool: def encode_payload(device_index: int, rows: int, heads: int, plan_kwargs: dict) -> dict[str, Any]: """JSON-safe description of one ``BatchSparseMLADecodePagedTSWrapper.plan`` call.""" kwargs = dict(plan_kwargs) - dtype = kwargs.pop("q_data_type") - kwargs["q_data_type"] = str(dtype).removeprefix("torch.") + dtype = kwargs.pop("q_dtype") + kwargs["q_dtype"] = str(dtype).removeprefix("torch.") return {"device_index": device_index, "rows": rows, "heads": heads, "plan_kwargs": kwargs} @@ -240,7 +240,7 @@ def _plan(mla_decode, payload: dict[str, Any], torch) -> None: device = torch.device("cuda", int(payload["device_index"])) torch.cuda.set_device(device) kwargs = dict(payload["plan_kwargs"]) - kwargs["q_data_type"] = getattr(torch, kwargs["q_data_type"]) + kwargs["q_dtype"] = getattr(torch, kwargs["q_dtype"]) mla_decode.BatchSparseMLADecodePagedTSWrapper().plan( device, int(payload["rows"]), int(payload["heads"]), **kwargs ) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py index 2cb531e21c3d..299df0c01577 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_sparse_mla.py @@ -538,10 +538,10 @@ def _plan_kwargs( plan_kwargs = dict( max_topk=int(self.attn.sparse_attention_config.window_size), max_extra_topk=extra, - max_seq_len_q=1, - q_data_type=q_dtype, + max_q_len=1, + q_dtype=q_dtype, kv_layout="NHD", - has_sinks=has_sinks, + use_attention_sinks=has_sinks, return_lse=False, assume_valid_prefix=not holes_possible, ) @@ -879,7 +879,7 @@ def forward( extra_pool, softmax_scale=softmax_scale, output_scale=output_scale, - sinks=sinks, + attention_sinks=sinks, validate=validate, **run_kwargs, ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/__init__.py b/tensorrt_llm/_torch/attention/backends/prims_ts/__init__.py index 477f52de75cf..429c9c88f7a6 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/__init__.py @@ -41,6 +41,7 @@ batch_sparse_mla_decode_with_paged_kv_cache, get_prims_ts_sparse_mla_decode_workspace_size, BatchMLADecodePagedTSWrapper, + MLAAttentionTuning, batch_mla_decode_with_paged_kv_cache, get_prims_ts_batch_mla_decode_workspace_size, ) @@ -50,6 +51,13 @@ q_token_kv_block_sparse_attention_with_paged_kv_cache, ) from .sage import SageAttentionConfig, SageAttentionParams +from .vc_attention import ( + VCAttentionConfig, + VCAttentionParams, +) +from .vc_attention.quantization import ( + VCAttentionPreprocessor, +) __all__ = [ "BatchSparseMLADecodePagedTSWrapper", @@ -58,6 +66,9 @@ "BlockSparseTSWrapper", "SageAttentionConfig", "SageAttentionParams", + "VCAttentionConfig", + "VCAttentionParams", + "VCAttentionPreprocessor", "BlockSparsePagedTSWrapper", "block_sparse_attention", "block_sparse_attention_with_paged_kv_cache", @@ -77,6 +88,7 @@ "get_q_token_kv_block_sparse_workspace_size", "q_token_kv_block_sparse_attention_with_paged_kv_cache", "BatchMLADecodePagedTSWrapper", + "MLAAttentionTuning", "batch_mla_decode_with_paged_kv_cache", "get_prims_ts_batch_mla_decode_workspace_size", ] diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/common.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/common.py index f67fdba1602b..60f880d9de8e 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/common.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/common.py @@ -14,7 +14,8 @@ """Dependency-neutral semantic contract for PrimTS block-sparse attention.""" -from collections.abc import Iterable +from collections.abc import Iterable, Mapping +from typing import TypeVar _FINE_KV_BLOCK_SIZES = (8, 16, 32) _PREPARED_KV_ROUTE_SIZE = 128 @@ -22,6 +23,47 @@ _SIGNED_INT32_MAX = (1 << 31) - 1 _BLOCK_SPARSE_Q_TILE_SIZES = (8, 16, 32, 64, 128) _BLOCK_SPARSE_MAX_HEADS_Q_PER_KV = 32 +# Q and KV block sizes of a dense plan that names none; they select the +# Q128/KV128 Keeps tile. +_DENSE_Q_BLOCK_SIZE = 128 +_DENSE_KV_BLOCK_SIZE = 128 +# Sage scale slots of the block-sparse attention adapter, in ABI order. +SAGE_ADAPTER_SLOTS = ("q_scale", "k_scale", "k_summary_scale", "v_scale", "v_mean") + +_T = TypeVar("_T") + + +def sage_adapter_slots(values: Mapping[str, _T | None]) -> tuple[_T | None, ...]: + """Arrange Sage values named by scale into the adapter slots. + + An absent name leaves its slot ``None``, which the adapter binds as a null + pointer. + """ + return tuple(values.get(name) for name in SAGE_ADAPTER_SLOTS) + + +def _resolve_plan_block_sizes( + q_block_size: int | None, + kv_block_size: int | None, + *, + use_block_sparse: bool, +) -> tuple[int, int]: + """Return a plan's Q and KV block sizes. + + A block-sparse plan names both, since they give its routes their meaning; + a dense plan may omit either, which then takes the dense default. + """ + + if use_block_sparse: + if q_block_size is None or kv_block_size is None: + raise ValueError( + "q_block_size and kv_block_size are required by a block-sparse plan" + ) + return q_block_size, kv_block_size + return ( + _DENSE_Q_BLOCK_SIZE if q_block_size is None else q_block_size, + _DENSE_KV_BLOCK_SIZE if kv_block_size is None else kv_block_size, + ) def _num_sparse_pattern_heads( 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 b0510cfc090b..4da185c88739 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 @@ -23,8 +23,8 @@ from flashinfer.utils import ceil_div -from .common import _num_sparse_pattern_heads -from ..sage import sage_adapter_slots, sage_scale_shapes +from .common import _num_sparse_pattern_heads, sage_adapter_slots +from ..sage.config import sage_scale_shapes from .config import _BlockSparseCompileKey, _make_block_sparse_config @@ -59,7 +59,7 @@ def _compile_block_sparse(key: _BlockSparseCompileKey) -> Callable[..., object]: fmha_block_sparse_launch, fmha_decode_launch, ) - from ..kernels.fmha_decode.sage_prepare import _PrepareSageScales + from ..sage.prepare import _PrepareSageScales config = _make_block_sparse_config(key) sparse_format = key.sparse_format @@ -74,10 +74,7 @@ def _compile_block_sparse(key: _BlockSparseCompileKey) -> Callable[..., object]: 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, + layout=sage_scale_layout, num_kv_heads=key.num_kv_heads ) if key.use_block_sparse: prepare_kwargs = { 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 e8234961ffba..4293ceebd7a5 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 @@ -24,14 +24,15 @@ 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 ..sage.config import SageAttentionConfig, validate_sage_recipe from .common import ( _PREPARED_KV_ROUTE_SIZE, _SIGNED_INT32_MAX, + _resolve_plan_block_sizes, _select_block_sparse_q_tile_size, _validate_sparse_kv_block_size, ) -from ..sage import _SageScaleLayout +from ..sage.layout import _SageScaleLayout from .prepared import _BlockSparseRouteLayout if TYPE_CHECKING: @@ -385,7 +386,7 @@ def _resolve_plan_dtypes( """Return ``(kv_dtype, v_dtype, output_dtype, dtype_key)`` of one plan. K defaults to Q and V to K. The output defaults to Q, or to bfloat16 with - Sage; the decode configuration validates the Sage dtype recipe. + Sage, whose recipe ``validate_sage_recipe`` checks. """ if kv_dtype is None: @@ -401,12 +402,15 @@ def _resolve_plan_dtypes( output_dtype, _validate_matching_dtypes(q_dtype, kv_dtype, v_dtype, output_dtype), ) - if not isinstance(sage, SageAttentionConfig): - raise TypeError("sage must be a SageAttentionConfig instance") - if q_dtype != kv_dtype: - raise ValueError("Sage attention requires Q and K in one dtype") if output_dtype is None: output_dtype = torch.bfloat16 + validate_sage_recipe( + sage, + q_dtype=q_dtype, + k_dtype=kv_dtype, + v_dtype=v_dtype, + output_dtype=output_dtype, + ) return kv_dtype, v_dtype, output_dtype, _plan_dtype_key(q_dtype) @@ -418,8 +422,8 @@ def _validate_block_sparse_static_profile( num_qo_heads: int, num_kv_heads: int, head_dim: int, - q_block_size: int, - kv_block_size: int, + q_block_size: int | None, + kv_block_size: int | None, use_kv_valid_bits: bool, mask_type: Literal["dense", "causal"], q_dtype: torch.dtype, @@ -461,6 +465,9 @@ def _validate_block_sparse_static_profile( ) if not isinstance(use_kv_valid_bits, bool): raise TypeError("use_kv_valid_bits must be a bool") + q_block_size, kv_block_size = _resolve_plan_block_sizes( + q_block_size, kv_block_size, use_block_sparse=use_block_sparse + ) kv_block_size = _validate_sparse_kv_block_size(kv_block_size) if num_qo_heads % num_kv_heads != 0: raise ValueError("num_qo_heads must be divisible by num_kv_heads") @@ -555,6 +562,8 @@ def _make_block_sparse_config(key: _BlockSparseCompileKey) -> "FmhaDecodeConfig" "groups_tokens_heads_q": True, "use_block_sparse": key.use_block_sparse, "softmax_rescale_threshold_log2": key.softmax_rescale_threshold_log2, + # Block-sparse and Sage plans compile for their exact lengths. + "fixed_seqlen": True, } if key.use_block_sparse: config_args.update( 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 24b508e5b237..3342375287c2 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 @@ -24,7 +24,11 @@ from flashinfer.utils import ceil_div -from ..sage import SageAttentionConfig, sage_scale_shapes +from ..sage.config import ( + SageAttentionConfig, + sage_scale_shapes, + _validate_int8_qk_device, +) from .common import ( _SIGNED_INT32_MAX, _block_sparse_proxy_summary_geometry, @@ -35,29 +39,12 @@ _BlockSparseStaticProfile, _resolve_block_sparse_launch_spec, ) -from ..sage import _SageScaleLayout +from ..sage.layout import _SageScaleLayout from .prepared import _BlockSparseRouteLayout _P = ParamSpec("_P") _R = TypeVar("_R") -# Only SM100 has the INT8 tcgen05 MMA kind; later GPUs use E4M3 Q/K. -_INT8_QK_COMPUTE_CAPABILITIES = ((10, 0),) - - -def _validate_int8_qk_device(q_dtype: torch.dtype, device_index: int) -> None: - """Reject Int8 Q/K on GPUs without the INT8 tcgen05 MMA.""" - - if q_dtype != torch.int8: - return - capability = torch.cuda.get_device_capability(device_index) - if capability not in _INT8_QK_COMPUTE_CAPABILITIES: - raise NotImplementedError( - "Int8 Q/K Sage attention requires an SM100a/B200 GPU, whose tcgen05 " - "MMA supports INT8 operands; use the torch.float8_e4m3fn Q/K recipe " - f"on device cuda:{device_index} with compute capability {capability}" - ) - class _PlanLockOwner(Protocol): """Structural type required by the plan-serialization decorator.""" 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 177145a15737..0754cad8950f 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 @@ -28,8 +28,9 @@ _validate_exact_compact_strides, _validate_scale, ) -from ..sage import SageAttentionParams, sage_adapter_slots, validate_sage_params +from ..sage.config import SageAttentionParams, validate_sage_params from .common import ( + sage_adapter_slots, _SIGNED_INT32_MAX, _num_sparse_pattern_heads, _validate_dense_contiguous_plan_inputs, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_q_token_kv_block_sparse_policy.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_q_token_kv_block_sparse_policy.py index e1a309e07a21..eaec54c0e79d 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/_q_token_kv_block_sparse_policy.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_q_token_kv_block_sparse_policy.py @@ -69,17 +69,18 @@ def select_sparse_launch( # capacity bound here; runtime masks the inactive split prefix. The G # recommendation below scores useful work with a shared-KV estimate. work_tokens = max(1, route_kv_tokens) - # D256 FP8 G1 requires the staged Keeps recipe. D64/D128 also have + # Wide-head FP8 G1 uses the qualified staged Keeps recipe. D64/D128 also have # qualified unstaged Swaps recipes. minimum_q = ( 64 - if head_dim == 256 and q_dtype_key == "float8_e4m3fn" and group_size == 1 + if head_dim in (192, 256) and q_dtype_key == "float8_e4m3fn" and group_size == 1 else 8 ) candidates = enumerate_grouped_q_mma_candidates( heads_q_per_kv=heads_q_per_kv, seq_len_q=group_size ) + # Choose the first Q tile covering the route group and minimum MMA rows. def recipe(minimum: int): return next( c @@ -87,11 +88,23 @@ def recipe(minimum: int): if c.tile_size_q >= minimum and c.q_tokens_per_cta >= group_size ) + # Choose split fanout using KV instances per loop and a useful-work floor. def fanout(instances: int) -> int: maximum = MAX_SPLITS_KV + minimum_iterations = MIN_LOOP_ITERS_PER_SPLIT if group_size == 1: + # Staged FP8 Keeps has one KV tile per iteration, whereas Swaps + # has two. Preserve a two-tile work floor for FP8 in either recipe: + # one-tile splits spend more on partial-output reduction than they + # save in attention. Scale the cap with work, not a fixed split count. + minimum_tiles = ( + max(2, instances) if q_dtype_key == "float8_e4m3fn" else instances + ) + minimum_iterations = minimum_tiles // instances # A partial final iteration must not buy an otherwise empty split. - maximum = min(maximum, max(1, work_tokens // (TILE_SIZE_KV * instances))) + maximum = min( + maximum, max(1, work_tokens // (TILE_SIZE_KV * minimum_tiles)) + ) return select_splits_kv( seq_len_kv=work_tokens, batch_size=num_routes, @@ -100,16 +113,16 @@ def fanout(instances: int) -> int: num_insts_kv=instances, service_capacity=multi_processor_count, max_splits_kv=maximum, - min_loop_iters_per_split=1 if group_size == 1 else MIN_LOOP_ITERS_PER_SPLIT, + min_loop_iters_per_split=minimum_iterations, split_kv=split_kv, ) candidate = recipe(minimum_q) - instances = 1 if head_dim == 256 and candidate.variant == "keeps_mma_ab" else 2 + instances = 1 if head_dim > 128 and candidate.variant == "keeps_mma_ab" else 2 splits = fanout(instances) - # Other D256 FP8 groups may use Swaps only when split-KV is active. + # Other wide-head FP8 groups may use Swaps only when split-KV is active. if ( - head_dim == 256 + head_dim in (192, 256) and q_dtype_key == "float8_e4m3fn" and splits == 1 and candidate.tile_size_q < 64 @@ -121,29 +134,42 @@ def fanout(instances: int) -> int: tile_size_q=candidate.tile_size_q, num_insts_kv=instances, splits_kv=splits, - # Full FP8 stages preserve the 128-byte swizzled planes while issuing + # Q128 amortizes a full D192 stage; narrower Q tiles retain D128 + # bands. Full FP8 stages preserve the 128-byte swizzled planes while issuing # PV at N256. BF16 and Swaps retain their existing D128 head bands. head_dim_per_stage_kv=( - 256 + 192 + if head_dim == 192 + and candidate.tile_size_q == 128 + and q_dtype_key == "bfloat16" + else 256 if head_dim == 256 and q_dtype_key == "float8_e4m3fn" and candidate.variant == "keeps_mma_ab" else 128 - if head_dim == 256 + if head_dim > 128 else 0 ), - o_stages=1 if head_dim == 256 and candidate.variant == "keeps_mma_ab" else 2, + o_stages=1 if head_dim > 128 and candidate.variant == "keeps_mma_ab" else 2, + # A caller that disables split-KV (packed prefill) plans each step's own + # route count. Taking the persistent scheduler for every count keeps one + # program per group size instead of a static/persistent pair keyed by + # whether a step happens to exceed one SM wave. Within one wave the two + # schedulers measure within about 1 us of each other. use_persistent_scheduler=( splits == 1 - and _select_auto_launch_mode( - batch_size=num_routes, - num_heads_kv=num_kv_heads, - seq_len_kv=work_tokens, - tile_size_kv=TILE_SIZE_KV, - split_kv=False, - service_capacity=multi_processor_count, + and ( + not split_kv + or _select_auto_launch_mode( + batch_size=num_routes, + num_heads_kv=num_kv_heads, + seq_len_kv=work_tokens, + tile_size_kv=TILE_SIZE_KV, + split_kv=False, + service_capacity=multi_processor_count, + ) + == "persistent" ) - == "persistent" ), ) 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 38de4926a045..dc0ffd1d3219 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/block_sparse.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/block_sparse.py @@ -77,7 +77,7 @@ _validate_runtime_device, ) from .kernels.fmha_decode.fmha_decode_constants import SOFTMAX_RESCALE_THRESHOLD_LOG2 -from .sage import SageAttentionConfig, SageAttentionParams +from .sage.config import SageAttentionConfig, SageAttentionParams class _BlockSparseWrapperBase: @@ -150,8 +150,8 @@ def plan( num_qo_heads: int, num_kv_heads: int, head_dim: int, - q_block_size: int, - kv_block_size: int, + q_block_size: int | None = None, + kv_block_size: int | None = None, *, device: torch.device | str | int, use_block_sparse: bool = True, @@ -186,7 +186,9 @@ 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, plus the scales of a - Sage plan. + Sage plan. A dense plan may omit ``q_block_size`` and + ``kv_block_size``; each then defaults to 128, which selects the + Q128/KV128 Keeps tile. A block-sparse plan requires both. ``share_pattern_across_kv_heads=False`` retains independent per-KV-head patterns. True uses a singleton head axis in BSR/bitmask inputs and @@ -235,16 +237,16 @@ def plan( ``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. + ``softmax_rescale_threshold_log2`` (0 to 8, default 8) applies to the + Q64/KV256 and Q128/KV128 profiles of dense and block-sparse 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: @@ -454,8 +456,8 @@ def block_sparse_attention( v: torch.Tensor, block_indptr: torch.Tensor | None, block_indices: torch.Tensor | None, - q_block_size: int, - kv_block_size: int, + q_block_size: int | None = None, + kv_block_size: int | None = None, *, exact_block_bits: torch.Tensor | None = None, k_summary: torch.Tensor | None = None, @@ -497,13 +499,15 @@ def block_sparse_attention( block_indices : torch.Tensor, optional Contiguous Int32 semantic KV-block IDs referenced by ``block_indptr``. Required in BSR mode. - q_block_size : int + q_block_size : int, optional Positive number of logical query tokens represented by one BSR row. The product with ``Hq / Hkv`` must be divisible by 8 so a physical Q - tile does not cross row boundaries. - kv_block_size : int + tile does not cross row boundaries. Required by a block-sparse call; + a dense call defaults it to 128. + kv_block_size : int, optional Number of logical KV tokens represented by one BSR block ID; it must - be 8, 16, 32, or a positive multiple of 64. + be 8, 16, 32, or a positive multiple of 64. Required by a block-sparse + call; a dense call defaults it to 128. exact_block_bits : torch.Tensor, optional Compact UInt32 exact-block bitmap required in bitmask mode. k_summary : torch.Tensor, optional diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/context.py b/tensorrt_llm/_torch/attention/backends/prims_ts/context.py index 1ec1080b942f..2cdb3a0be100 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/context.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/context.py @@ -17,9 +17,9 @@ The public surface intentionally exposes attention semantics and optional compile-time metadata contracts, not scheduler choices. Reusable contiguous -and paged plans compile from static capacities; packed request lengths and -metadata remain per-run inputs. Both wrappers select scheduling from static -topology while consuming the current request extents at launch. Paged plans +and paged plans select resource topology from capacities, while logical Q/KV +lengths and request metadata remain runtime inputs. Both wrappers consume +the current request extents at launch. Paged plans conservatively treat all request lengths as variable unless the caller opts into an exact-uniform or zero-causal-offset contract. Single-instance dense paged plans use a static persistent launch when their @@ -35,20 +35,50 @@ register trace templates. """ -from dataclasses import dataclass +from dataclasses import dataclass, field, replace import functools import itertools import math -import numbers -import struct from typing import TYPE_CHECKING, Callable, Literal, Optional +import torch + from ._compile_cache import build_and_load_attention_kernel from ._hardware import get_max_active_clusters -import torch +from .vc_attention import ( + VCAttentionConfig, + VCAttentionParams, + validate_vc_params, + vc_repair_kv_len, + vc_scale_shapes, +) from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api +from .decode import ( + _validate_native_hnd_tma_strides, + prepare_lse, + validate_lse_mode, + validate_attention_sinks, + _validate_scale, + _validate_layout as _validate_kv_layout, +) + +from .sage.config import ( + SageAttentionConfig, + SageAttentionParams, + validate_sage_params, + validate_sage_recipe, + _validate_int8_qk_device, +) +from .sage.layout import ( + _SAGE_ATOM_TOKENS, + _SageKSection, + _SageQSection, + _SageScaleLayout, + _SageVSection, +) +from .sage.plan import SagePreparedScales, build_sage_prepared_scales if TYPE_CHECKING: @@ -56,9 +86,9 @@ _COMPILE_OPTIONS = "--enable-tvm-ffi --opt-level 3" -_SUPPORTED_HEAD_DIMS = (128, 256) +_SUPPORTED_HEAD_DIMS = (128, 192, 256) _DEFAULT_PAGED_KV_PAGE_SIZE = 32 -_SUPPORTED_PAGED_KV_PAGE_SIZES = (16, 32, 64, 128) +_SUPPORTED_PAGED_KV_PAGE_SIZES = (4, 8, 16, 32, 64, 128) _SUPPORTED_DTYPES = ( torch.float16, torch.bfloat16, @@ -94,18 +124,18 @@ class _ContextGeometry: device: torch.device device_index: int - packed: bool + packed_query: bool batch_size: int total_q: int total_k: int - max_seq_len_q: int + max_q_len: int max_seq_len_k: int num_qo_heads: int num_kv_heads: int head_dim: int qk_dtype: torch.dtype pv_dtype: torch.dtype - output_dtype: torch.dtype + out_dtype: torch.dtype mask_type: str window_left: int head_paired: bool @@ -124,16 +154,16 @@ class _ContextPlanGeometry: device: torch.device device_index: int - packed: bool + packed_query: bool batch_size: int - max_seq_len_q: int + max_q_len: int max_seq_len_k: int num_qo_heads: int num_kv_heads: int head_dim: int qk_dtype: torch.dtype pv_dtype: torch.dtype - output_dtype: torch.dtype + out_dtype: torch.dtype mask_type: str window_left: int head_paired: bool @@ -142,6 +172,11 @@ class _ContextPlanGeometry: causal_single_kv_tile: bool packed_dense_k_mask: bool head_dim_vo: int | None = None + # VC-Attention recipe, or ``None`` for a plain plan. + vc: Optional[VCAttentionConfig] = None + # The plan's Sage recipe, part of the compile identity; ``None`` without + # Sage attention. + sage: SageAttentionConfig | None = None @dataclass(frozen=True) @@ -149,13 +184,20 @@ class _ContextPlanState: """Atomically published state for one compiled contiguous context plan.""" geometry: _ContextPlanGeometry - scale_softmax_log2: torch.Tensor - output_scale: torch.Tensor + scale_softmax_log2: float | torch.Tensor + output_scale: float | torch.Tensor empty_i32: torch.Tensor variable_window_padded_starts: Optional[torch.Tensor] variable_window_cta_starts: torch.Tensor compiled: Callable[..., None] policy: tuple[tuple[str, object], ...] + vc_scale_shapes: Optional[dict[str, tuple[int, ...]]] = None + # Output scale of VC runs: the per-channel V scale times the planned or passed scale. + vc_output_scale: Optional[torch.Tensor] = None + lse_mode: str = "none" + use_attention_sinks: bool = False + sage: Optional[SagePreparedScales] = None + cuda_scales: bool = False @dataclass(frozen=True) @@ -165,7 +207,7 @@ class _PagedContextPlanGeometry: device: torch.device device_index: int batch_size: int - max_seq_len_q: int + max_q_len: int max_kv_len: int page_size: int num_qo_heads: int @@ -173,7 +215,7 @@ class _PagedContextPlanGeometry: head_dim: int qk_dtype: torch.dtype pv_dtype: torch.dtype - output_dtype: torch.dtype + out_dtype: torch.dtype mask_type: str window_left: int head_paired: bool @@ -188,10 +230,12 @@ class _PagedContextPlanState: """Atomically published state for one compiled paged-context plan.""" geometry: _PagedContextPlanGeometry - scale_softmax_log2: torch.Tensor - output_scale: torch.Tensor + scale_softmax_log2: float + output_scale: float compiled: Callable[..., None] policy: tuple[tuple[str, object], ...] + lse_mode: str = "none" + use_attention_sinks: bool = False _ContextScheduler = Literal[ @@ -215,17 +259,17 @@ class _ContextCompileSpec: """Batch-independent identity for one contiguous context compile.""" device_index: int - max_seq_len_q: int - max_seq_len_k: int + max_q_len: int | None = field(compare=False, hash=False) + max_seq_len_k: int | None = field(compare=False, hash=False) num_qo_heads: int num_kv_heads: int head_dim: int qk_dtype_key: str pv_dtype_key: str - output_dtype_key: str + out_dtype_key: str mask_type: str window_left: int - packed: bool + packed_query: bool head_paired: bool uniform_packed_lengths: bool has_q_offset: bool @@ -234,6 +278,14 @@ class _ContextCompileSpec: scheduler: _ContextScheduler head_dim_vo: int | None = None two_cta_umma: bool = False + vc: Optional[VCAttentionConfig] = None + lse_mode: str = "none" + use_attention_sinks: bool = False + runtime_lengths: bool = False + shape_signature: tuple[int, ...] = () + # A Sage plan's prepared scale layout; the kernel reads the scales. + sage_scale_layout: _SageScaleLayout | None = None + cuda_scales: bool = False @dataclass(frozen=True) @@ -241,15 +293,15 @@ class _PagedContextCompileSpec: """Batch-independent identity for one paged context compile.""" device_index: int - max_seq_len_q: int - max_kv_len: int + max_q_len: int | None = field(compare=False, hash=False) + max_kv_len: int | None = field(compare=False, hash=False) page_size: int num_qo_heads: int num_kv_heads: int head_dim: int qk_dtype_key: str pv_dtype_key: str - output_dtype_key: str + out_dtype_key: str mask_type: str window_left: int head_paired: bool @@ -258,6 +310,10 @@ class _PagedContextCompileSpec: paged_v_tail_is_zero: bool packed_dense_k_mask: bool scheduler: _ContextScheduler + lse_mode: str = "none" + use_attention_sinks: bool = False + runtime_lengths: bool = False + shape_signature: tuple[int, ...] = () def _make_context_kernel( @@ -265,7 +321,7 @@ def _make_context_kernel( device_index: int, input_qk_dtype, input_pv_dtype, - output_dtype, + out_dtype, head_dim: int, head_dim_vo: int | None = None, mask_type: str, @@ -278,9 +334,13 @@ def _make_context_kernel( scheduler: _ContextScheduler, uses_ldtm_stat: bool, fp8_psmem_early_token: bool = False, + kv_ring_stages: int = 0, + umma_fp8_k64: bool = False, two_cta_umma: bool = False, page_size: int | None = None, max_kv_len: int | None = None, + vc_attention: str | None = None, + sage_scale_layout: _SageScaleLayout | None = None, ): """Build one context kernel from its batch-independent static topology.""" @@ -289,29 +349,27 @@ def _make_context_kernel( from .kernels.fmha_context.fmha_kernel import FmhaTs use_paged_kv = page_size is not None - if use_paged_kv != (max_kv_len is not None): - raise ValueError("paged context requires both page geometry values") + # VC QK8 also fixes the K-scale extent through max_kv_len. Paged runtime + # context only needs page geometry; its live KV bound is a launch operand. is_persistent = scheduler != "nonpersistent" is_clc_dynamic = scheduler == "clc_dynamic_persistent" paged_kwargs = ( - { - "use_paged_kv": True, - "num_tokens_per_page": page_size, - "max_kv_len": max_kv_len, - } - if use_paged_kv - else {} + {"use_paged_kv": True, "num_tokens_per_page": page_size} if use_paged_kv else {} ) + if max_kv_len is not None: + paged_kwargs["max_kv_len"] = max_kv_len fmha = FmhaTs( qk_acc_dtype=cutlass.Float32, pv_acc_dtype=cutlass.Float32, in_qk_dtype=input_qk_dtype, in_pv_dtype=input_pv_dtype, - out_dtype=output_dtype, + out_dtype=out_dtype, d=head_dim, d_v=head_dim_vo, is_persistent=is_persistent, fp8_psmem_early_token=fp8_psmem_early_token, + kv_ring_stages=kv_ring_stages, + umma_fp8_k64=umma_fp8_k64, two_cta_umma=two_cta_umma, is_causal=mask_type == "causal", has_variable_window=mask_type == "variable_window", @@ -330,9 +388,11 @@ 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), + vc_attention=vc_attention, use_blackwell_tma_oob_workaround=( torch.cuda.get_device_capability(device_index) in ((10, 0), (10, 3)) ), + sage_scale_layout=sage_scale_layout, **paged_kwargs, ) return fmha @@ -373,20 +433,33 @@ def _dtype_key(dtype: torch.dtype) -> str: torch.float16: "float16", torch.bfloat16: "bfloat16", torch.float8_e4m3fn: "float8_e4m3fn", + torch.int8: "int8", } try: return keys[dtype] except KeyError as error: raise NotImplementedError( "attention-ts context supports torch.float16, torch.bfloat16, " - f"and torch.float8_e4m3fn; got {dtype}" + f"torch.float8_e4m3fn, and torch.int8 Q/K; got {dtype}" ) from error +def _v_dtype_follows_qk(q_dtype: torch.dtype, v_dtype: torch.dtype) -> bool: + """Whether V may take ``v_dtype`` next to ``q_dtype`` Q/K. + + V matches Q/K except on the QK-BF16/PV-FP8 path and under a Sage recipe's + Int8 Q/K, which keep E4M3 V. + """ + return v_dtype == q_dtype or ( + q_dtype in (torch.bfloat16, torch.int8) and v_dtype == torch.float8_e4m3fn + ) + + def _validate_qkv_dtype(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> None: """Validate run-time Q/K/V tensor dtypes. - Q and K must match; V may differ only for the QK-BF16/PV-FP8 path. + Q and K must match; V may differ only for the QK-BF16/PV-FP8 path and + for Int8 Q/K with E4M3 V. """ _dtype_key(q.dtype) if q.dtype != k.dtype: @@ -394,27 +467,27 @@ def _validate_qkv_dtype(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> No "attention-ts context requires Q and K to use the same dtype; " f"got Q {q.dtype} and K {k.dtype}" ) - # V may keep its own dtype only for the supported QK-BF16/PV-FP8 path. - qk_bf16_pv_fp8 = q.dtype == torch.bfloat16 and v.dtype == torch.float8_e4m3fn - if v.dtype != q.dtype and not qk_bf16_pv_fp8: + if not _v_dtype_follows_qk(q.dtype, v.dtype): raise NotImplementedError( "attention-ts context requires Q, K, and V to use the same dtype " - "except for QK-BF16/PV-FP8; " + "except for QK-BF16/PV-FP8 and Int8 Q/K with E4M3 V; " f"got Q {q.dtype}, K {k.dtype}, and V {v.dtype}" ) -def _validate_output_dtype(output_dtype: torch.dtype) -> None: - if not isinstance(output_dtype, torch.dtype): +def _validate_output_dtype(out_dtype: torch.dtype) -> None: + if not isinstance(out_dtype, torch.dtype): raise TypeError("out_dtype must be a torch.dtype") - _dtype_key(output_dtype) + _dtype_key(out_dtype) + if out_dtype == torch.int8: + raise NotImplementedError("attention-ts context does not write an int8 output") def _validate_context_dtype_pair( q_dtype: torch.dtype, k_dtype: torch.dtype, v_dtype: torch.dtype, - output_dtype: torch.dtype, + out_dtype: torch.dtype, ) -> None: """Validate the explicit dtypes of a contiguous-context specialization.""" @@ -422,22 +495,20 @@ def _validate_context_dtype_pair( (q_dtype, "q_dtype"), (k_dtype, "k_dtype"), (v_dtype, "v_dtype"), - (output_dtype, "out_dtype"), ): if not isinstance(dtype, torch.dtype): raise TypeError(f"{name} must be a torch.dtype") _dtype_key(dtype) + _validate_output_dtype(out_dtype) if k_dtype != q_dtype: raise NotImplementedError( "attention-ts context requires Q and K to use the same dtype; " f"got q_dtype={q_dtype} and k_dtype={k_dtype}" ) - if v_dtype != q_dtype and not ( - q_dtype == torch.bfloat16 and v_dtype == torch.float8_e4m3fn - ): + if not _v_dtype_follows_qk(q_dtype, v_dtype): raise NotImplementedError( "attention-ts context requires Q, K, and V to use the same dtype " - "except for QK-BF16/PV-FP8; " + "except for QK-BF16/PV-FP8 and Int8 Q/K with E4M3 V; " f"got q_dtype={q_dtype} and v_dtype={v_dtype}" ) @@ -446,7 +517,7 @@ def _validate_paged_dtype_pair( q_dtype: torch.dtype, k_dtype: torch.dtype, v_dtype: torch.dtype, - output_dtype: torch.dtype, + out_dtype: torch.dtype, ) -> None: """Validate the explicit dtypes of a paged-context specialization.""" @@ -454,19 +525,21 @@ def _validate_paged_dtype_pair( (q_dtype, "q_dtype"), (k_dtype, "k_dtype"), (v_dtype, "v_dtype"), - (output_dtype, "out_dtype"), + (out_dtype, "out_dtype"), ): if not isinstance(dtype, torch.dtype): raise TypeError(f"{name} must be a torch.dtype") _dtype_key(dtype) + if dtype == torch.int8: + raise NotImplementedError( + "attention-ts paged context does not support int8 tensors" + ) if k_dtype != q_dtype: raise NotImplementedError( "attention-ts paged context requires Q and K to use the same " f"dtype; got q_dtype={q_dtype} and k_dtype={k_dtype}" ) - if v_dtype != q_dtype and not ( - q_dtype == torch.bfloat16 and v_dtype == torch.float8_e4m3fn - ): + if not _v_dtype_follows_qk(q_dtype, v_dtype): raise NotImplementedError( "attention-ts paged context requires Q, K, and V to use the same " "dtype except for QK-BF16/PV-FP8; " @@ -500,36 +573,76 @@ def _validate_device(device: torch.device) -> int: 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. + """FMA-pipe exp2 pairs per 16-pair softmax chunk on SM100. - Measured on B200. The paired D128 dense path is MUFU bound and takes 4. + Measured on B200. The paired D128 dense path is MUFU bound: with 16-bit V + it takes 4, and with 8-bit V and P in SMEM, fixed-length plans take 3. + Packed and paged 8-bit plans spend the FMA pipe on their masks and take 0. Single-QKV dense takes 2. Causal is issue bound and takes 0. Callers can set cfg.exp2_fma_pairs to any value on any path.""" - if torch.cuda.get_device_capability(device_index) != (10, 0): + if torch.cuda.get_device_capability(device_index) != (10, 0) or cfg.is_causal: return 0 - if cfg.v_dtype.width != 16 or cfg.is_causal: + if cfg.v_dtype.width == 8: + return 0 if cfg.has_varlen or cfg.single_qkv_instance else 3 + if cfg.v_dtype.width != 16: return 0 if cfg.pv_half_overlap: return 4 return 2 if cfg.single_qkv_instance else 0 +def _default_exp2_packed_f16x2(device_index: int, cfg) -> bool: + """Softmax exp2 as one packed f16x2 MUFU op per score pair, on for the SM107 + paths with 8-bit V. SM100 and SM103 split the packed op into two MUFU ops and + keep the scalar form.""" + return ( + torch.cuda.get_device_capability(device_index) == (10, 7) + and cfg.v_dtype.width == 8 + ) + + +def _default_kv_ring_stages(device_index: int, qk_bits: int, pv_bits: int) -> int: + """K/V ring depth of the SM107 dense paths with 8-bit V, 0 elsewhere for the + budget-derived depth. Deeper rings stall the softmax P stores behind the TMA + fills, so the fp8 K ring stops at 6 stages, and the bf16 K ring takes the 4 + stages the larger SM107 shared memory fits.""" + if torch.cuda.get_device_capability(device_index) == (10, 7) and pv_bits == 8: + return 6 if qk_bits == 8 else 4 + return 0 + + +def _default_umma_fp8_k64(device_index: int, qk_bits: int) -> bool: + """K=64 fp8 UMMA instructions for the PV of the SM107 bf16-QK paths. The + fp8-QK paths are shared-memory-port bound on operand fetch and keep K=32.""" + return torch.cuda.get_device_capability(device_index) == (10, 7) and qk_bits == 16 + + def _default_fp8_psmem_early_token(device_index: int) -> bool: """Release the S0/S1 token early on the fp8 SMEM-P path, on for SM103.""" return torch.cuda.get_device_capability(device_index) == (10, 3) -def _default_two_cta_umma(device_index: int) -> bool: - """Two-CTA UMMA for the dense paired D128 context kernel: on for SM103, where the - kernel is tensor-pipe bound and halving UMMA issues and K/V traffic per SM pays.""" - return torch.cuda.get_device_capability(device_index) == (10, 3) +def _default_two_cta_umma(device_index: int, qk_bits: int) -> bool: + """Two-CTA UMMA for the dense paired D128 context kernel. On for SM103, where + the kernel is tensor-pipe bound and halving the UMMA issues and the K/V + traffic per SM pays. On SM107 only for bf16 QK, which is shared-memory-port + bound on operand fetch, since fp8 QK is slower with it.""" + capability = torch.cuda.get_device_capability(device_index) + if capability == (10, 3): + return True + return capability == (10, 7) and qk_bits == 16 -def _default_uses_ldtm_stat(device_index: int) -> bool: - """Enable LDTM.STAT where the shared tcgen05 gate allows it.""" +def _default_uses_ldtm_stat(device_index: int, qk_bits: int, pv_bits: int) -> bool: + """Enable LDTM.STAT where the shared tcgen05 gate allows it. SM107 keeps it + off for the dense QK-BF16 PV-FP8 path, which is faster with the register + row max.""" from .kernels.tcgen05_compat import ldtm_stat_supported - return ldtm_stat_supported(torch.cuda.get_device_capability(device_index)) + capability = torch.cuda.get_device_capability(device_index) + if capability == (10, 7) and qk_bits == 16 and pv_bits == 8: + return False + return ldtm_stat_supported(capability) def _resolve_cuda_device( @@ -589,11 +702,11 @@ def _validate_variable_window_bounds( "variable_window_token_starts and variable_window_token_ends are " "required for mask_type='variable_window'" ) - if geometry.packed: + if geometry.packed_query: raise NotImplementedError( "variable_window currently requires fixed [B, S, H, D] tensors" ) - expected_shape = (geometry.batch_size, geometry.max_seq_len_q) + expected_shape = (geometry.batch_size, geometry.max_q_len) for tensor, name in ( (starts, "variable_window_token_starts"), (ends, "variable_window_token_ends"), @@ -621,6 +734,13 @@ def _validate_variable_window_bounds( f"{offset} has start={start}, end={end}, and " f"max_kv_len={geometry.max_seq_len_k}" ) + if offset % geometry.max_q_len and end < end_values[offset - 1]: + batch_idx, query_idx = divmod(offset, geometry.max_q_len) + raise ValueError( + "variable_window_token_ends must be nondecreasing within each " + f"batch item; batch {batch_idx}, query {query_idx} has " + f"end={end} after end={end_values[offset - 1]}" + ) return starts.flatten(), ends.flatten(), start_values @@ -631,15 +751,15 @@ def _refresh_variable_window_cta_starts( geometry = state.geometry tile_size_q = _variable_window_tile_size_q(geometry.head_dim) - num_seq_tiles = (geometry.max_seq_len_q + tile_size_q - 1) // tile_size_q + num_seq_tiles = (geometry.max_q_len + tile_size_q - 1) // tile_size_q padded_rows = num_seq_tiles * tile_size_q - starts_2d = starts.view(geometry.batch_size, geometry.max_seq_len_q) + starts_2d = starts.view(geometry.batch_size, geometry.max_q_len) padded = state.variable_window_padded_starts if padded is not None: padded.fill_(_INT32_MAX) - padded[:, : geometry.max_seq_len_q].copy_(starts_2d) + padded[:, : geometry.max_q_len].copy_(starts_2d) starts_2d = padded - elif padded_rows != geometry.max_seq_len_q: + elif padded_rows != geometry.max_q_len: raise RuntimeError("variable-window padded-start scratch is missing") torch.amin( starts_2d.view(geometry.batch_size, num_seq_tiles, tile_size_q), @@ -669,7 +789,7 @@ def _validate_variable_window_cta_starts( f"got {cta_starts.dtype}" ) tile_size_q = _variable_window_tile_size_q(geometry.head_dim) - num_seq_tiles = (geometry.max_seq_len_q + tile_size_q - 1) // tile_size_q + num_seq_tiles = (geometry.max_q_len + tile_size_q - 1) // tile_size_q expected_shape = (geometry.batch_size, num_seq_tiles) if tuple(cta_starts.shape) != expected_shape: raise ValueError( @@ -681,12 +801,12 @@ def _validate_variable_window_cta_starts( actual_values = tuple(int(value) for value in cta_starts.flatten().tolist()) for batch_idx in range(geometry.batch_size): - row_offset = batch_idx * geometry.max_seq_len_q + row_offset = batch_idx * geometry.max_q_len for tile_idx in range(num_seq_tiles): tile_begin = row_offset + tile_idx * tile_size_q tile_end = min( tile_begin + tile_size_q, - row_offset + geometry.max_seq_len_q, + row_offset + geometry.max_q_len, ) expected = min(token_start_values[tile_begin:tile_end]) actual = actual_values[batch_idx * num_seq_tiles + tile_idx] @@ -699,26 +819,6 @@ def _validate_variable_window_cta_starts( return cta_starts.flatten() -def _validate_scale(value: object, name: str) -> float: - if isinstance(value, bool) or not isinstance(value, numbers.Real): - raise TypeError(f"{name} must be a positive Python scalar") - try: - as_float = float(value) - except (OverflowError, TypeError, ValueError) as error: - raise TypeError(f"{name} must be a positive Python scalar") from error - if not math.isfinite(as_float) or as_float <= 0.0: - raise ValueError(f"{name} must be finite and positive") - try: - as_float32 = struct.unpack("=f", struct.pack("=f", as_float))[0] - except (OverflowError, struct.error) as error: - raise ValueError( - f"{name} must be representable as a positive float32" - ) from error - if not math.isfinite(as_float32) or as_float32 <= 0.0: - raise ValueError(f"{name} must be representable as a positive float32") - return as_float32 - - def _validate_extent(value: int, name: str) -> int: if value <= 0: raise ValueError(f"{name} must be positive") @@ -781,7 +881,7 @@ def _read_indptr( *, expected_total: int, ) -> tuple[tuple[int, ...], tuple[int, ...]]: - """Read cumulative offsets and validate strictly positive row lengths.""" + """Read cumulative offsets and validate nonnegative row lengths.""" values = tuple(int(value) for value in indptr.tolist()) if values[0] != 0: @@ -792,8 +892,8 @@ def _read_indptr( f"expected {expected_total}, got {values[-1]}" ) lengths = tuple(curr - prev for prev, curr in itertools.pairwise(values)) - if any(length <= 0 for length in lengths): - raise ValueError(f"{name} offsets must be strictly increasing") + if any(length < 0 for length in lengths): + raise ValueError(f"{name} offsets must be nondecreasing") return values, lengths @@ -880,15 +980,6 @@ def _validate_page_size(page_size: int) -> int: return page_size -def _validate_kv_layout(kv_layout: str) -> None: - if not isinstance(kv_layout, str): - raise TypeError("kv_layout must be a string") - if kv_layout != "HND": - raise NotImplementedError( - f"attention-ts paged context supports only kv_layout='HND'; got {kv_layout!r}" - ) - - def _validate_base_tensors( q: torch.Tensor, k: torch.Tensor, @@ -1002,7 +1093,7 @@ def _contiguous_context_uses_clc_scheduler( *, single_qkv_instance: bool, head_paired: bool, - packed: bool, + packed_query: bool, uniform_packed_lengths: bool, is_causal: bool, has_q_offset: bool, @@ -1020,7 +1111,7 @@ def _contiguous_context_uses_clc_scheduler( return ( head_paired or not single_qkv_instance - or (packed and not uniform_packed_lengths) + or (packed_query and not uniform_packed_lengths) or (is_causal and not has_q_offset) ) @@ -1029,7 +1120,7 @@ def _contiguous_context_uses_persistent_scheduler( *, single_qkv_instance: bool, head_paired: bool, - packed: bool, + packed_query: bool, uniform_packed_lengths: bool, logical_work_tiles: int, max_active_clusters: int, @@ -1050,7 +1141,7 @@ def _contiguous_context_uses_persistent_scheduler( return ( head_paired or not single_qkv_instance - or (packed and not uniform_packed_lengths) + or (packed_query and not uniform_packed_lengths) or (is_causal and not has_q_offset) or logical_work_tiles > max_active_clusters or batch_size > _CUDA_GRID_YZ_MAX @@ -1060,7 +1151,7 @@ def _contiguous_context_uses_persistent_scheduler( def _needs_packed_dense_k_mask( *, - packed: bool, + packed_query: bool, mask_type: str, k_lengths: tuple[int, ...], ) -> bool: @@ -1072,7 +1163,7 @@ def _needs_packed_dense_k_mask( """ return ( - packed + packed_query and mask_type == "dense" and ( any(length != k_lengths[0] for length in k_lengths[1:]) @@ -1117,21 +1208,21 @@ def _resolve_geometry( kv_indptr: Optional[torch.Tensor], mask_type: str, window_left: int, - output_dtype: torch.dtype, + out_dtype: torch.dtype, ) -> _ContextGeometry: """Validate one-shot inputs and derive their static wrapper bounds.""" _validate_base_tensors(q, k, v, allow_mla=True) - _validate_output_dtype(output_dtype) + _validate_output_dtype(out_dtype) _validate_mask(mask_type) window_left = _validate_window_left(window_left, mask_type) device_index = _validate_device(q.device) device = torch.device("cuda", device_index) if (qo_indptr is None) != (kv_indptr is None): raise ValueError("qo_indptr and kv_indptr must be provided together") - packed = qo_indptr is not None + packed_query = qo_indptr is not None - if packed: + if packed_query: if q.ndim != 3 or k.ndim != 3: raise ValueError( "packed Q/K/V must use [total_tokens, H, D] storage; " @@ -1156,7 +1247,7 @@ def _resolve_geometry( _, k_lengths = _read_indptr(kv_indptr, "kv_indptr", expected_total=total_k) batch_size = len(q_lengths) _validate_extent(batch_size, "batch_size") - max_seq_len_q = max(q_lengths) + max_q_len = max(q_lengths) max_seq_len_k = max(k_lengths) q_shape = tuple(q.shape) kv_shape = tuple(k.shape) @@ -1166,23 +1257,23 @@ def _resolve_geometry( "fixed Q/K/V must use [B, S, H, D] storage; " f"got q rank {q.ndim} and k rank {k.ndim}" ) - batch_size, max_seq_len_q, num_qo_heads, q_head_dim = map(int, q.shape) + batch_size, max_q_len, num_qo_heads, q_head_dim = map(int, q.shape) k_batch, max_seq_len_k, num_kv_heads, kv_head_dim = map(int, k.shape) if batch_size != k_batch: raise ValueError( f"q and k batch dimensions must match; got {batch_size} and {k_batch}" ) _validate_extent(batch_size, "batch_size") - _validate_extent(max_seq_len_q, "seq_len_q") + _validate_extent(max_q_len, "seq_len_q") _validate_extent(max_seq_len_k, "seq_len_k") - total_q = batch_size * max_seq_len_q + total_q = batch_size * max_q_len total_k = batch_size * max_seq_len_k _validate_padded_data_extent(total_q, "B * seq_len_q") _validate_padded_data_extent(total_k, "B * seq_len_k") _validate_compact(q, "q", "[B, Sq, Hq, D]") _validate_compact(k, "k", "[B, Sk, Hkv, D]") _validate_compact(v, "v", "[B, Sk, Hkv, D]") - q_lengths = (max_seq_len_q,) * batch_size + q_lengths = (max_q_len,) * batch_size k_lengths = (max_seq_len_k,) * batch_size q_shape = tuple(q.shape) kv_shape = tuple(k.shape) @@ -1193,13 +1284,13 @@ def _resolve_geometry( for batch_idx, (q_length, k_length) in enumerate( zip(q_lengths, k_lengths, strict=True) ): - if q_length > k_length: + if k_length > 0 and q_length > k_length: raise ValueError( "bottom-right causal context requires Sq <= Sk for each " f"request; got batch {batch_idx}: Sq={q_length}, Sk={k_length}" ) - if q_head_dim == 192 and window_left > 0: + if (q_head_dim, int(v.shape[-1])) == (192, 128) and window_left > 0: raise NotImplementedError("192/128 context does not support a left window") head_paired = window_left > 0 if head_paired and (head_ratio <= 1 or head_ratio % 2 != 0): @@ -1209,30 +1300,30 @@ def _resolve_geometry( ) has_q_offset = _derive_has_q_offset(q_lengths, k_lengths, mask_type) uniform_packed_lengths = ( - packed + packed_query and all(length == q_lengths[0] for length in q_lengths[1:]) and all(length == k_lengths[0] for length in k_lengths[1:]) ) causal_single_kv_tile = ( mask_type == "causal" - and not packed + and not packed_query and not head_paired - and max_seq_len_q == max_seq_len_k + and max_q_len == max_seq_len_k and max_seq_len_k <= 128 ) packed_dense_k_mask = _needs_packed_dense_k_mask( - packed=packed, + packed_query=packed_query, mask_type=mask_type, k_lengths=k_lengths, ) return _ContextGeometry( device=device, device_index=device_index, - packed=packed, + packed_query=packed_query, batch_size=batch_size, total_q=total_q, total_k=total_k, - max_seq_len_q=max_seq_len_q, + max_q_len=max_q_len, max_seq_len_k=max_seq_len_k, num_qo_heads=num_qo_heads, num_kv_heads=num_kv_heads, @@ -1240,7 +1331,7 @@ def _resolve_geometry( head_dim_vo=int(v.shape[-1]), qk_dtype=q.dtype, pv_dtype=v.dtype, - output_dtype=output_dtype, + out_dtype=out_dtype, mask_type=mask_type, window_left=window_left, head_paired=head_paired, @@ -1257,7 +1348,7 @@ def _resolve_context_plan_geometry( *, device: int | str | torch.device, batch_size: int, - max_seq_len_q: int, + max_q_len: int, max_kv_len: int, num_qo_heads: int, num_kv_heads: int, @@ -1265,25 +1356,26 @@ def _resolve_context_plan_geometry( head_dim_vo: int | None = None, qk_dtype: torch.dtype, pv_dtype: torch.dtype, - packed: bool, + packed_query: bool, mask_type: str, window_left: int, - output_dtype: torch.dtype, + out_dtype: torch.dtype, + sage: SageAttentionConfig | None = None, ) -> _ContextPlanGeometry: """Validate explicit static bounds for a reusable contiguous plan.""" - _validate_context_dtype_pair(qk_dtype, qk_dtype, pv_dtype, output_dtype) + _validate_context_dtype_pair(qk_dtype, qk_dtype, pv_dtype, out_dtype) _validate_mask(mask_type) window_left = _validate_window_left(window_left, mask_type) - if not isinstance(packed, bool): - raise TypeError("packed must be a bool") - if packed and mask_type == "variable_window": + if not isinstance(packed_query, bool): + raise TypeError("packed_query must be a bool") + if packed_query and mask_type == "variable_window": raise NotImplementedError( "mask_type='variable_window' requires fixed [B, S, H, D] tensors" ) batch_size = _validate_static_extent(batch_size, "batch_size") - max_seq_len_q = _validate_static_extent(max_seq_len_q, "max_seq_len_q") + max_q_len = _validate_static_extent(max_q_len, "max_q_len") max_kv_len = _validate_static_extent(max_kv_len, "max_kv_len") num_qo_heads = _validate_static_extent(num_qo_heads, "num_qo_heads") num_kv_heads = _validate_static_extent(num_kv_heads, "num_kv_heads") @@ -1295,20 +1387,18 @@ def _resolve_context_plan_geometry( ) device, device_index = _resolve_cuda_device(device) - _validate_padded_data_extent( - batch_size * max_seq_len_q, "batch_size * max_seq_len_q" - ) + _validate_padded_data_extent(batch_size * max_q_len, "batch_size * max_q_len") _validate_padded_data_extent(batch_size * max_kv_len, "batch_size * max_kv_len") head_ratio = _validate_head_geometry(num_qo_heads, num_kv_heads) _validate_head_dim(head_dim, head_dim, head_dim_vo) - if head_dim == 192 and window_left > 0: + if (head_dim, head_dim_vo) == (192, 128) and window_left > 0: raise NotImplementedError( "QK/V=(192, 128) does not support a positive left window" ) - if not packed and mask_type == "causal" and max_seq_len_q > max_kv_len: + if not packed_query and mask_type == "causal" and max_q_len > max_kv_len: raise ValueError( - "bottom-right causal context requires max_seq_len_q <= max_kv_len; " - f"got {max_seq_len_q} and {max_kv_len}" + "bottom-right causal context requires max_q_len <= max_kv_len; " + f"got {max_q_len} and {max_kv_len}" ) head_paired = window_left > 0 if head_paired and (head_ratio <= 1 or head_ratio % 2 != 0): @@ -1319,13 +1409,13 @@ def _resolve_context_plan_geometry( # Packed request lengths are run inputs, so select conservative compile # facts that remain valid for every request within the plan capacities. - has_q_offset = mask_type == "causal" and (packed or max_seq_len_q != max_kv_len) - return _ContextPlanGeometry( + has_q_offset = mask_type == "causal" and (packed_query or max_q_len != max_kv_len) + geometry = _ContextPlanGeometry( device=device, device_index=device_index, - packed=packed, + packed_query=packed_query, batch_size=batch_size, - max_seq_len_q=max_seq_len_q, + max_q_len=max_q_len, max_seq_len_k=max_kv_len, num_qo_heads=num_qo_heads, num_kv_heads=num_kv_heads, @@ -1333,7 +1423,7 @@ def _resolve_context_plan_geometry( head_dim_vo=head_dim_vo, qk_dtype=qk_dtype, pv_dtype=pv_dtype, - output_dtype=output_dtype, + out_dtype=out_dtype, mask_type=mask_type, window_left=window_left, head_paired=head_paired, @@ -1341,13 +1431,24 @@ def _resolve_context_plan_geometry( has_q_offset=has_q_offset, causal_single_kv_tile=( mask_type == "causal" - and not packed + and not packed_query and not head_paired - and max_seq_len_q == max_kv_len + and max_q_len == max_kv_len and max_kv_len <= _CONTEXT_KV_TILE_N ), - packed_dense_k_mask=packed and mask_type == "dense", + packed_dense_k_mask=packed_query and mask_type == "dense", + sage=sage, ) + if sage is not None: + _validate_context_sage_config(geometry, sage) + elif qk_dtype == torch.int8: + raise NotImplementedError( + "int8 Q/K require a Sage recipe (sage=SageAttentionConfig(...)) that " + "dequantizes them" + ) + # The recipe is rejected before the device, as in block-sparse plans. + _validate_int8_qk_device(qk_dtype, device_index) + return geometry def _resolve_paged_geometry( @@ -1361,7 +1462,7 @@ def _resolve_paged_geometry( page_size: int, mask_type: str, window_left: int, - output_dtype: torch.dtype, + out_dtype: torch.dtype, ) -> _PagedContextPlanGeometry: """Validate one fixed-table paged launch and derive its plan bounds.""" @@ -1399,8 +1500,8 @@ def _resolve_paged_geometry( ) _validate_head_dim(q_head_dim, kv_head_dim) _validate_compact(q, "q", "[total_q, Hq, D]") - _validate_compact(k_cache, "k_cache", "[num_pages, Hkv, page_size, D]") - _validate_compact(v_cache, "v_cache", "[num_pages, Hkv, page_size, D]") + _validate_native_hnd_tma_strides(k_cache, "k_cache") + _validate_native_hnd_tma_strides(v_cache, "v_cache") _validate_indptr_tensor(qo_indptr, "qo_indptr", device=device) batch_size = int(qo_indptr.numel()) - 1 @@ -1413,22 +1514,22 @@ def _resolve_paged_geometry( expected_count=batch_size, ) for batch_idx, k_length in enumerate(k_lengths): - if k_length <= 0: + if k_length < 0: raise ValueError( - "seq_lens_kv entries must be positive; " + "seq_lens_kv entries must be nonnegative; " f"batch {batch_idx} has {k_length}" ) - max_seq_len_q = max(q_lengths) + max_q_len = max(q_lengths) max_kv_len = max(k_lengths) - uniform_packed_lengths = all( - length == max_seq_len_q for length in q_lengths - ) and all(length == max_kv_len for length in k_lengths) + uniform_packed_lengths = all(length == max_q_len for length in q_lengths) and all( + length == max_kv_len for length in k_lengths + ) has_q_offset = _derive_has_q_offset(q_lengths, k_lengths, mask_type) geometry = _resolve_paged_plan_geometry( device=device, batch_size=batch_size, - max_seq_len_q=max_seq_len_q, + max_q_len=max_q_len, max_kv_len=max_kv_len, num_qo_heads=num_qo_heads, num_kv_heads=num_kv_heads, @@ -1438,7 +1539,7 @@ def _resolve_paged_geometry( page_size=page_size, mask_type=mask_type, window_left=window_left, - output_dtype=output_dtype, + out_dtype=out_dtype, uniform_packed_lengths=uniform_packed_lengths, has_q_offset=has_q_offset, ) @@ -1469,12 +1570,13 @@ def _make_context_scheduler_probe( torch.float16: cutlass.Float16, torch.bfloat16: cutlass.BFloat16, torch.float8_e4m3fn: cutlass.Float8E4M3FN, + torch.int8: cutlass.Int8, } probe = _make_context_kernel( device_index=geometry.device_index, input_qk_dtype=dtype_map[geometry.qk_dtype], input_pv_dtype=dtype_map[geometry.pv_dtype], - output_dtype=dtype_map[geometry.output_dtype], + out_dtype=dtype_map[geometry.out_dtype], head_dim=geometry.head_dim, head_dim_vo=getattr(geometry, "head_dim_vo", None), mask_type=geometry.mask_type, @@ -1485,14 +1587,18 @@ def _make_context_scheduler_probe( has_q_offset=geometry.has_q_offset, causal_single_kv_tile=causal_single_kv_tile, scheduler="static_persistent", - uses_ldtm_stat=_default_uses_ldtm_stat(geometry.device_index), + uses_ldtm_stat=_default_uses_ldtm_stat( + geometry.device_index, + geometry.qk_dtype.itemsize * 8, + geometry.pv_dtype.itemsize * 8, + ), page_size=page_size, max_kv_len=max_kv_len, ) with torch.cuda.device(geometry.device_index): max_active_clusters = get_max_active_clusters(1) num_seq_tiles = ( - geometry.max_seq_len_q + probe.cfg.cta_tiler[0] - 1 + geometry.max_q_len + probe.cfg.cta_tiler[0] - 1 ) // probe.cfg.cta_tiler[0] num_head_tiles = geometry.num_qo_heads // probe.cfg.work_tile_q_heads logical_work_tiles = geometry.batch_size * num_seq_tiles * num_head_tiles @@ -1507,7 +1613,7 @@ def _resolve_paged_plan_geometry( *, device: int | str | torch.device, batch_size: int, - max_seq_len_q: int, + max_q_len: int, max_kv_len: int, num_qo_heads: int, num_kv_heads: int, @@ -1517,14 +1623,14 @@ def _resolve_paged_plan_geometry( page_size: int, mask_type: str, window_left: int, - output_dtype: torch.dtype, + out_dtype: torch.dtype, uniform_packed_lengths: bool = False, has_q_offset: bool = True, paged_v_tail_is_zero: bool = False, ) -> _PagedContextPlanGeometry: """Validate explicit static bounds for a reusable paged specialization.""" - _validate_paged_dtype_pair(qk_dtype, qk_dtype, pv_dtype, output_dtype) + _validate_paged_dtype_pair(qk_dtype, qk_dtype, pv_dtype, out_dtype) _validate_mask(mask_type) if not isinstance(uniform_packed_lengths, bool): raise TypeError("uniform_packed_lengths must be a bool") @@ -1539,16 +1645,14 @@ def _resolve_paged_plan_geometry( window_left = _validate_window_left(window_left, mask_type) page_size = _validate_page_size(page_size) batch_size = _validate_static_extent(batch_size, "batch_size") - max_seq_len_q = _validate_static_extent(max_seq_len_q, "max_seq_len_q") + max_q_len = _validate_static_extent(max_q_len, "max_q_len") max_kv_len = _validate_static_extent(max_kv_len, "max_kv_len") num_qo_heads = _validate_static_extent(num_qo_heads, "num_qo_heads") num_kv_heads = _validate_static_extent(num_kv_heads, "num_kv_heads") head_dim = _validate_static_extent(head_dim, "head_dim") device, device_index = _resolve_cuda_device(device) - _validate_padded_data_extent( - batch_size * max_seq_len_q, "batch_size * max_seq_len_q" - ) + _validate_padded_data_extent(batch_size * max_q_len, "batch_size * max_q_len") head_ratio = _validate_head_geometry(num_qo_heads, num_kv_heads) _validate_head_dim(head_dim, head_dim) head_paired = window_left > 0 @@ -1567,7 +1671,7 @@ def _resolve_paged_plan_geometry( device=device, device_index=device_index, batch_size=batch_size, - max_seq_len_q=max_seq_len_q, + max_q_len=max_q_len, max_kv_len=max_kv_len, page_size=page_size, num_qo_heads=num_qo_heads, @@ -1575,7 +1679,7 @@ def _resolve_paged_plan_geometry( head_dim=head_dim, qk_dtype=qk_dtype, pv_dtype=pv_dtype, - output_dtype=output_dtype, + out_dtype=out_dtype, mask_type=mask_type, window_left=window_left, head_paired=head_paired, @@ -1595,14 +1699,35 @@ def _two_cta_umma_geometry_eligible(geometry: _ContextPlanGeometry) -> bool: The two-CTA launch is non-persistent with heads on grid Y and batch on grid Z. CUDA limits grid Y and Z to 65,535. Oversized geometries keep the persistent flattened grid so the valid int32 plans stay launchable.""" + # VC-Attention-QK16 always runs two-CTA. With P in SMEM, the bf16 K ring, + # tile-mean ring and row-sum operands leave room for one K/V stage per CTA + # at one CTA and three in the halved form. The VC tail needs two. + # VC-Attention-QK8 keeps the E4M3 K ring and follows the device default. + vc_qk16 = geometry.vc is not None and geometry.qk_dtype != torch.float8_e4m3fn + two_cta = _default_two_cta_umma( + geometry.device_index, geometry.qk_dtype.itemsize * 8 + ) + if two_cta and torch.cuda.get_device_capability(geometry.device_index) == (10, 7): + # Pairs are coarser work units. Keep them only where they fill their + # waves at least as well as single CTAs fill theirs. + sms = torch.cuda.get_device_properties( + geometry.device_index + ).multi_processor_count + tiles = math.ceil(geometry.max_q_len / (2 * _CONTEXT_TILE_SIZE_Q)) + heads = geometry.batch_size * geometry.num_qo_heads + single_waves = heads * tiles / sms + pair_waves = heads * math.ceil(tiles / 2) / (sms // 2) + two_cta = pair_waves / math.ceil(pair_waves) >= single_waves / math.ceil( + single_waves + ) return ( - _default_two_cta_umma(geometry.device_index) + (two_cta or vc_qk16) and geometry.head_dim == 128 and geometry.head_dim_vo in (None, 128) and geometry.mask_type == "dense" - and not geometry.packed + and not geometry.packed_query and not geometry.head_paired - and torch.finfo(geometry.qk_dtype).bits in (8, 16) + and geometry.qk_dtype in (torch.float16, torch.bfloat16, torch.float8_e4m3fn) and not ( geometry.batch_size > _CUDA_GRID_YZ_MAX or geometry.num_qo_heads > _CUDA_GRID_YZ_MAX @@ -1623,7 +1748,7 @@ def _resolve_context_scheduler(geometry: _ContextPlanGeometry) -> _ContextSchedu is_persistent = _contiguous_context_uses_persistent_scheduler( single_qkv_instance=probe.single_qkv_instance, head_paired=geometry.head_paired, - packed=geometry.packed, + packed_query=geometry.packed_query, uniform_packed_lengths=geometry.uniform_packed_lengths, logical_work_tiles=probe.logical_work_tiles, max_active_clusters=probe.max_active_clusters, @@ -1634,10 +1759,21 @@ def _resolve_context_scheduler(geometry: _ContextPlanGeometry) -> _ContextSchedu ) if not is_persistent: return "nonpersistent" + # Fixed causal requests with one work item per head all traverse the same + # KV domain. Static persistence avoids CLC traffic without a load-balance + # penalty; live packed offsets and multiple Q work items still use CLC. + if ( + geometry.mask_type == "causal" + and not geometry.packed_query + and not geometry.head_paired + and geometry.window_left <= 0 + and probe.logical_work_tiles == geometry.batch_size * geometry.num_qo_heads + ): + return "static_persistent" if _contiguous_context_uses_clc_scheduler( single_qkv_instance=probe.single_qkv_instance, head_paired=geometry.head_paired, - packed=geometry.packed, + packed_query=geometry.packed_query, uniform_packed_lengths=geometry.uniform_packed_lengths, is_causal=geometry.mask_type == "causal", has_q_offset=geometry.has_q_offset, @@ -1677,54 +1813,163 @@ def _resolve_paged_context_scheduler( return "static_persistent" -def _context_compile_spec(geometry: _ContextPlanGeometry) -> _ContextCompileSpec: - """Build the contiguous compile key, including separate QK and PV dtype keys.""" +def _fixed_kv_rows(max_seq_len_k: int, vc: Optional[VCAttentionConfig]) -> int: + """K/V rows of a fixed plan. V repair adds its tiles to the sequence.""" + if vc is not None and vc.repair_tiles: + return vc_repair_kv_len(max_seq_len_k, vc.repair_tiles) + return max_seq_len_k + + +def _validate_context_sage_config( + geometry: _ContextPlanGeometry, sage: SageAttentionConfig +) -> None: + """Reject a Sage recipe the contiguous context kernel does not run. + + Beyond the recipe every wrapper checks (``validate_sage_recipe``), Sage + attention covers the dense, fixed-length, query-paired D128 path. + """ + + validate_sage_recipe( + sage, + q_dtype=geometry.qk_dtype, + k_dtype=geometry.qk_dtype, + v_dtype=geometry.pv_dtype, + output_dtype=geometry.out_dtype, + ) + if geometry.mask_type != "dense" or geometry.packed_query or geometry.head_paired: + raise NotImplementedError( + "Sage attention requires a dense, fixed-length, query-paired context plan" + ) + if geometry.head_dim != 128 or geometry.head_dim_vo not in (None, 128): + raise NotImplementedError("Sage attention requires head_dim 128") + + +def _context_sage_layout(geometry: _ContextPlanGeometry) -> _SageScaleLayout: + """Return the prepared scale layout of a contiguous context Sage plan. + + The Q rows hold one Q tile row per token of one Q head (``_SageQSection`` + without interleaved heads): the query-paired D128 tile. + """ + + sage = geometry.sage + return _SageScaleLayout( + k_tokens=_SageKSection( + length=geometry.max_seq_len_k, + block_size=sage.k_block_size, + source_block_size=sage.k_block_size, + tile_atoms=_CONTEXT_KV_TILE_N // _SAGE_ATOM_TOKENS, + ), + q_rows=_SageQSection( + seq_len_q=geometry.max_q_len, + q_block_size=sage.q_block_size, + heads_q_per_kv=geometry.num_qo_heads // geometry.num_kv_heads, + rows_per_tile=_CONTEXT_TILE_SIZE_Q, + heads_interleaved=False, + ), + v_rows=_SageVSection(headdim=geometry.head_dim, has_mean=sage.v_mean), + ) + + +def _build_context_sage_state(geometry: _ContextPlanGeometry) -> SagePreparedScales: + """Allocate a plan's prepared scale buffer and compile its prepare launch.""" + + return build_sage_prepared_scales( + _context_sage_layout(geometry), + geometry.sage, + batch_size=geometry.batch_size, + seq_len_q=geometry.max_q_len, + seq_len_kv=geometry.max_seq_len_k, + num_qo_heads=geometry.num_qo_heads, + num_kv_heads=geometry.num_kv_heads, + head_dim=geometry.head_dim, + device=geometry.device, + device_index=geometry.device_index, + ) + + +@functools.cache +def _context_compile_spec( + geometry: _ContextPlanGeometry, + *, + runtime_lengths: bool | None = None, + lse_mode: str = "none", + use_attention_sinks: bool = False, + cuda_scales: bool = False, +) -> _ContextCompileSpec: + """Keep resource/precision choices static and ordinary lengths runtime.""" + fixed_seqlen = geometry.sage is not None or geometry.vc is not None + runtime_lengths = not fixed_seqlen + shape_signature = ( + () if runtime_lengths else (geometry.max_q_len, geometry.max_seq_len_k) + ) return _ContextCompileSpec( device_index=geometry.device_index, - max_seq_len_q=geometry.max_seq_len_q, - max_seq_len_k=geometry.max_seq_len_k, + max_q_len=geometry.max_q_len if fixed_seqlen else None, + max_seq_len_k=geometry.max_seq_len_k if fixed_seqlen else None, num_qo_heads=geometry.num_qo_heads, num_kv_heads=geometry.num_kv_heads, head_dim=geometry.head_dim, head_dim_vo=geometry.head_dim_vo, qk_dtype_key=_dtype_key(geometry.qk_dtype), pv_dtype_key=_dtype_key(geometry.pv_dtype), - output_dtype_key=_dtype_key(geometry.output_dtype), + out_dtype_key=_dtype_key(geometry.out_dtype), mask_type=geometry.mask_type, window_left=geometry.window_left, - packed=geometry.packed, + packed_query=geometry.packed_query, head_paired=geometry.head_paired, - uniform_packed_lengths=geometry.uniform_packed_lengths, + uniform_packed_lengths=geometry.uniform_packed_lengths and fixed_seqlen, has_q_offset=geometry.has_q_offset, causal_single_kv_tile=geometry.causal_single_kv_tile, - packed_dense_k_mask=geometry.packed_dense_k_mask, + packed_dense_k_mask=( + geometry.packed_dense_k_mask + if fixed_seqlen + else geometry.mask_type != "causal" + ), + shape_signature=shape_signature, + runtime_lengths=runtime_lengths, + lse_mode=lse_mode, + use_attention_sinks=use_attention_sinks, + cuda_scales=cuda_scales, scheduler=_resolve_context_scheduler(geometry), two_cta_umma=_two_cta_umma_geometry_eligible(geometry), + vc=geometry.vc, + sage_scale_layout=( + _context_sage_layout(geometry) if geometry.sage is not None else None + ), ) +@functools.cache def _paged_context_compile_spec( geometry: _PagedContextPlanGeometry, + *, + runtime_lengths: bool = False, + lse_mode: str = "none", + use_attention_sinks: bool = False, ) -> _PagedContextCompileSpec: """Build the paged compile key, including separate QK and PV dtype keys.""" return _PagedContextCompileSpec( device_index=geometry.device_index, - max_seq_len_q=geometry.max_seq_len_q, - max_kv_len=geometry.max_kv_len, + max_q_len=None, + max_kv_len=None, page_size=geometry.page_size, num_qo_heads=geometry.num_qo_heads, num_kv_heads=geometry.num_kv_heads, head_dim=geometry.head_dim, qk_dtype_key=_dtype_key(geometry.qk_dtype), pv_dtype_key=_dtype_key(geometry.pv_dtype), - output_dtype_key=_dtype_key(geometry.output_dtype), + out_dtype_key=_dtype_key(geometry.out_dtype), mask_type=geometry.mask_type, window_left=geometry.window_left, head_paired=geometry.head_paired, - uniform_packed_lengths=geometry.uniform_packed_lengths, + uniform_packed_lengths=False, has_q_offset=geometry.has_q_offset, paged_v_tail_is_zero=geometry.paged_v_tail_is_zero, - packed_dense_k_mask=geometry.packed_dense_k_mask, + packed_dense_k_mask=(geometry.mask_type != "causal"), + shape_signature=(), + runtime_lengths=True, + lse_mode=lse_mode, + use_attention_sinks=use_attention_sinks, scheduler=_resolve_paged_context_scheduler(geometry), ) @@ -1736,18 +1981,20 @@ def _get_compiled_context( """Compile one batch-dynamic context specialization and cache by topology.""" device_index = compile_spec.device_index - max_seq_len_q = compile_spec.max_seq_len_q - max_seq_len_k = compile_spec.max_seq_len_k + # Only Sage and VC plans keep exact lengths in the key; they compile static shapes. + fixed_seqlen = compile_spec.max_q_len is not None + max_q_len = compile_spec.max_q_len or 0 + max_seq_len_k = compile_spec.max_seq_len_k or 0 num_qo_heads = compile_spec.num_qo_heads num_kv_heads = compile_spec.num_kv_heads head_dim = compile_spec.head_dim head_dim_vo = compile_spec.head_dim_vo or head_dim qk_dtype_key = compile_spec.qk_dtype_key pv_dtype_key = compile_spec.pv_dtype_key - output_dtype_key = compile_spec.output_dtype_key + out_dtype_key = compile_spec.out_dtype_key mask_type = compile_spec.mask_type window_left = compile_spec.window_left - packed = compile_spec.packed + packed_query = compile_spec.packed_query head_paired = compile_spec.head_paired uniform_packed_lengths = compile_spec.uniform_packed_lengths has_q_offset = compile_spec.has_q_offset @@ -1755,6 +2002,9 @@ def _get_compiled_context( packed_dense_k_mask = compile_spec.packed_dense_k_mask scheduler = compile_spec.scheduler two_cta_umma = compile_spec.two_cta_umma + vc_attention = compile_spec.vc is not None + sage_scale_layout = compile_spec.sage_scale_layout + vc_qk8 = vc_attention and qk_dtype_key == "float8_e4m3fn" import cutlass import cutlass.cute as cute @@ -1764,10 +2014,11 @@ def _get_compiled_context( "float16": cutlass.Float16, "bfloat16": cutlass.BFloat16, "float8_e4m3fn": cutlass.Float8E4M3FN, + "int8": cutlass.Int8, } input_qk_dtype = dtype_map[qk_dtype_key] input_pv_dtype = dtype_map[pv_dtype_key] - output_dtype = dtype_map[output_dtype_key] + out_dtype = dtype_map[out_dtype_key] has_variable_window = mask_type == "variable_window" with torch.cuda.device(device_index): max_active_clusters = get_max_active_clusters(1) @@ -1775,7 +2026,7 @@ def _get_compiled_context( device_index=device_index, input_qk_dtype=input_qk_dtype, input_pv_dtype=input_pv_dtype, - output_dtype=output_dtype, + out_dtype=out_dtype, head_dim=head_dim, head_dim_vo=head_dim_vo, mask_type=mask_type, @@ -1787,16 +2038,41 @@ def _get_compiled_context( causal_single_kv_tile=causal_single_kv_tile, scheduler=scheduler, two_cta_umma=two_cta_umma, - uses_ldtm_stat=_default_uses_ldtm_stat(device_index), + uses_ldtm_stat=_default_uses_ldtm_stat( + device_index, input_qk_dtype.width, input_pv_dtype.width + ), fp8_psmem_early_token=_default_fp8_psmem_early_token(device_index), + kv_ring_stages=_default_kv_ring_stages( + device_index, input_qk_dtype.width, input_pv_dtype.width + ), + umma_fp8_k64=_default_umma_fp8_k64(device_index, input_qk_dtype.width), + vc_attention=("means" if compile_spec.vc.demean else "outliers") + if vc_attention + else None, + sage_scale_layout=sage_scale_layout, + # The QK8 K-scale table is sized from the K/V rows of the fixed plan. + max_kv_len=_fixed_kv_rows(max_seq_len_k, compile_spec.vc) if vc_qk8 else None, ) - fmha.cfg.has_varlen = packed + if vc_attention and torch.cuda.get_device_capability(device_index) == (10, 7): + raise ValueError( + "vc_config is not supported on SM107, where the dense fp8 kernel runs " + "the softmax exp2 natively" + ) + fmha.cfg.lse_mode = compile_spec.lse_mode + fmha.cfg.use_attention_sinks = compile_spec.use_attention_sinks + fmha.cfg.has_varlen = packed_query fmha.cfg.has_uniform_varlen = uniform_packed_lengths + if vc_qk8: + # The QK8 scale tables are addressed with the fixed plan lengths, the K/V + # rows including the repair tiles. + fmha.cfg.uniform_seq_len_q = max_q_len + fmha.cfg.uniform_seq_len_k = _fixed_kv_rows(max_seq_len_k, compile_spec.vc) if uniform_packed_lengths: - fmha.cfg.uniform_seq_len_q = max_seq_len_q + fmha.cfg.uniform_seq_len_q = max_q_len fmha.cfg.uniform_seq_len_k = max_seq_len_k fmha.cfg.has_q_offset = has_q_offset fmha.cfg.exp2_fma_pairs = _default_exp2_fma_pairs(device_index, fmha.cfg) + fmha.cfg.exp2_packed_f16x2 = _default_exp2_packed_f16x2(device_index, fmha.cfg) if fmha.cfg.kv_tile_n != _CONTEXT_KV_TILE_N: raise RuntimeError( "context packed-K specialization assumes kv_tile_n=" @@ -1804,8 +2080,15 @@ def _get_compiled_context( ) _validate_query_work_tile_span(fmha.cfg) fmha.cfg.packed_dense_k_mask = packed_dense_k_mask - if mask_type != "causal" and not packed: - fmha.cfg.fixed_dense_k_tail = max_seq_len_k % fmha.cfg.kv_tile_n + if mask_type != "causal" and not packed_query: + fmha.cfg.fixed_dense_k_tail = ( + max_seq_len_k % fmha.cfg.kv_tile_n if fixed_seqlen else fmha.cfg.kv_tile_n + ) + + # Compile both bounded and general masks once; choose from runtime K. + fmha.runtime_dense_tail = ( + not fixed_seqlen and not packed_query and mask_type == "dense" + ) @cute.jit def tensor_adapter( @@ -1813,22 +2096,29 @@ def tensor_adapter( k: cute.Tensor, v: cute.Tensor, out: cute.Tensor, - scale_softmax_log2: cute.Tensor, - output_scale: cute.Tensor, + lse: cute.Tensor | None, + attention_sinks: cute.Tensor | None, + scale_softmax_log2: cutlass.Float32 | cute.Tensor, + output_scale: cutlass.Float32 | cute.Tensor, qo_indptr: cute.Tensor, kv_indptr: cute.Tensor, variable_window_token_starts: cute.Tensor, variable_window_token_ends: cute.Tensor, variable_window_cta_starts: cute.Tensor, + vc_mu: cute.Tensor | None, + vc_q_scale: cute.Tensor | None, + vc_k_scale: cute.Tensor | None, + sage_scales: cute.Tensor | None, stream: cuda_drv.CUstream, static_max_active_clusters: cutlass.Constexpr[int], - static_packed: cutlass.Constexpr[bool], - static_max_seq_len_q: cutlass.Constexpr[int], - static_max_seq_len_k: cutlass.Constexpr[int], + static_packed_query: cutlass.Constexpr[bool], + max_q_len: cutlass.Int32, + runtime_max_seq_len_k: cutlass.Int32, + static_vc_repair_tiles: cutlass.Constexpr[int], ) -> None: """Adapt torch TVM-FFI tensors to the FmhaTs host entry point.""" - if cutlass.const_expr(static_packed): + if cutlass.const_expr(static_packed_query): fmha( q, k, @@ -1840,11 +2130,17 @@ def tensor_adapter( stream, qo_indptr, kv_indptr, - cutlass.Int32(static_max_seq_len_q), - cutlass.Int32(static_max_seq_len_k), + max_q_len, + runtime_max_seq_len_k, variable_window_token_starts=variable_window_token_starts, variable_window_token_ends=variable_window_token_ends, variable_window_cta_starts=variable_window_cta_starts, + vc_mu=vc_mu, + vc_q_scale=vc_q_scale, + vc_k_scale=vc_k_scale, + vc_repair_tiles=static_vc_repair_tiles, + lse_output=lse, + attention_sinks=attention_sinks, ) else: fmha( @@ -1859,8 +2155,16 @@ def tensor_adapter( variable_window_token_starts=variable_window_token_starts, variable_window_token_ends=variable_window_token_ends, variable_window_cta_starts=variable_window_cta_starts, + vc_mu=vc_mu, + vc_q_scale=vc_q_scale, + vc_k_scale=vc_k_scale, + vc_repair_tiles=static_vc_repair_tiles, + lse_output=lse, + attention_sinks=attention_sinks, + sage_scales=sage_scales, ) + # Describe row-major symbolic extents; assumed_align is in bytes. def fake_compact(dtype, shape, assumed_align): return cute.runtime.make_fake_compact_tensor( dtype, @@ -1872,7 +2176,7 @@ def fake_compact(dtype, shape, assumed_align): q_shape: tuple[object, ...] kv_shape: tuple[object, ...] out_shape: tuple[object, ...] - if packed: + if packed_query: runtime_total_q = cute.sym_int() runtime_total_k = cute.sym_int() runtime_num_q_offsets = cute.sym_int() @@ -1884,26 +2188,46 @@ def fake_compact(dtype, shape, assumed_align): kv_indptr_shape = (runtime_num_k_offsets,) else: batch_size = cute.sym_int() - q_shape = (batch_size, max_seq_len_q, num_qo_heads, head_dim) - kv_shape = (batch_size, max_seq_len_k, num_kv_heads, head_dim) + physical_sq = cute.sym_int() if compile_spec.runtime_lengths else max_q_len + physical_sk = ( + cute.sym_int() + if compile_spec.runtime_lengths + else _fixed_kv_rows(max_seq_len_k, compile_spec.vc) + ) + q_shape = (batch_size, physical_sq, num_qo_heads, head_dim) + kv_shape = (batch_size, physical_sk, num_kv_heads, head_dim) out_shape = (*q_shape[:-1], head_dim_vo) qo_indptr_shape = (1,) kv_indptr_shape = (1,) q_fake = fake_compact(input_qk_dtype, q_shape, 16) k_fake = fake_compact(input_qk_dtype, kv_shape, 16) v_fake = fake_compact(input_pv_dtype, (*kv_shape[:-1], head_dim_vo), 16) - out_fake = fake_compact(output_dtype, out_shape, 16) - scale_fake = fake_compact(cutlass.Float32, (1,), 4) - output_scale_fake = fake_compact(cutlass.Float32, (1,), 4) + out_fake = fake_compact(out_dtype, out_shape, 16) + lse_fake = ( + fake_compact(cutlass.Float32, out_fake.shape[:-1], 4) + if compile_spec.lse_mode != "none" + else None + ) + scale_fake = cutlass.Float32(1.0) + output_scale_fake = cutlass.Float32(1.0) + if compile_spec.cuda_scales: + scale_fake = fake_compact(cutlass.Float32, (1,), 4) + output_scale_fake = fake_compact(cutlass.Float32, (1,), 4) + if vc_attention: + output_scale_fake = fake_compact( + cutlass.Float32, (batch_size, num_qo_heads, head_dim_vo), 64 + ) qo_indptr_fake = fake_compact(cutlass.Int32, qo_indptr_shape, 4) kv_indptr_fake = fake_compact(cutlass.Int32, kv_indptr_shape, 4) if has_variable_window: - if packed: + if packed_query: raise RuntimeError("variable-window context requires fixed tensors") - variable_window_shape = (batch_size * max_seq_len_q,) + variable_window_shape = (batch_size * q_shape[1],) variable_window_tile_size_q = _variable_window_tile_size_q(head_dim) variable_window_cta_shape = ( - batch_size * cute.ceil_div(max_seq_len_q, variable_window_tile_size_q), + (cute.sym_int(),) + if not fixed_seqlen + else (batch_size * cute.ceil_div(max_q_len, variable_window_tile_size_q),) ) else: variable_window_shape = (1,) @@ -1913,39 +2237,88 @@ def fake_compact(dtype, shape, assumed_align): variable_window_cta_starts_fake = fake_compact( cutlass.Int32, variable_window_cta_shape, 4 ) + vc_mu_fake = None + vc_q_scale_fake = None + vc_k_scale_fake = None + if vc_attention: + if packed_query: + raise RuntimeError("VC-Attention context requires fixed tensors") + # The plan fixes every VC operand shape, so the kernel addresses the + # tile means and the Q and K scale tables with compile-time extents. + vc_shapes = vc_scale_shapes( + compile_spec.vc, + batch_size=batch_size, + seq_len_kv=max_seq_len_k, + num_kv_heads=num_kv_heads, + head_dim=head_dim_vo or head_dim, + seq_len_q=max_q_len, + num_qo_heads=num_qo_heads, + qk_fp8=vc_qk8, + ) + if compile_spec.vc.demean: + vc_mu_fake = fake_compact(cutlass.BFloat16, vc_shapes["tile_means"], 16) + if vc_qk8: + # The kernel reads the Q and K scale tables as one row of blocks per head. + vc_q_scale_fake = fake_compact( + cutlass.Float32, + (num_qo_heads, batch_size * vc_shapes["q_scale"][2]), + 16, + ) + vc_k_scale_fake = fake_compact( + cutlass.Float32, + (num_kv_heads, batch_size * vc_shapes["k_scale"][2]), + 16, + ) + sinks_fake = ( + fake_compact(cutlass.Float32, (num_qo_heads,), 4) + if compile_spec.use_attention_sinks + else None + ) + # The prepared Sage scales span ``batch * Hkv`` sequences, so their length + # stays a runtime extent like the batch size; a plan without Sage has no + # scales. + sage_scales_fake = ( + fake_compact(cutlass.Float32, (cute.sym_int(),), 16) + if sage_scale_layout is not None + else None + ) stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) # Task objects carry loop-local state through generated control flow, so # select the public staged frontend for this compilation. - def compile_kernel(): - return cute.compile[cute.FrontendNext]( - tensor_adapter, - q_fake, - k_fake, - v_fake, - out_fake, - scale_fake, - output_scale_fake, - qo_indptr_fake, - kv_indptr_fake, - variable_window_starts_fake, - variable_window_ends_fake, - variable_window_cta_starts_fake, - stream_fake, - max_active_clusters, - packed, - max_seq_len_q, - max_seq_len_k, - options=_COMPILE_OPTIONS, - ) - with torch.cuda.device(device_index): compiled = build_and_load_attention_kernel( "context", "main", compile_spec, _COMPILE_OPTIONS, - compile_kernel, + lambda: cute.compile[cute.FrontendNext]( + tensor_adapter, + q_fake, + k_fake, + v_fake, + out_fake, + lse_fake, + sinks_fake, + scale_fake, + output_scale_fake, + qo_indptr_fake, + kv_indptr_fake, + variable_window_starts_fake, + variable_window_ends_fake, + variable_window_cta_starts_fake, + vc_mu_fake, + vc_q_scale_fake, + vc_k_scale_fake, + sage_scales_fake, + stream_fake, + max_active_clusters, + packed_query, + cutlass.Int32(0), + cutlass.Int32(0), + compile_spec.vc.repair_tiles if vc_attention else 0, + options=_COMPILE_OPTIONS, + ), extra=(max_active_clusters,), ) policy = ( @@ -1954,6 +2327,7 @@ def compile_kernel(): ("uniform_packed_lengths", uniform_packed_lengths), ("causal_single_kv_tile", causal_single_kv_tile), ("packed_dense_k_mask", packed_dense_k_mask), + ("vc_attention", vc_attention), ) return compiled, policy @@ -1965,15 +2339,15 @@ def _get_compiled_paged_context( """Compile one batch-dynamic packed-Q, paged-K/V specialization.""" device_index = compile_spec.device_index - max_seq_len_q = compile_spec.max_seq_len_q - max_kv_len = compile_spec.max_kv_len + max_q_len = compile_spec.max_q_len or 0 + max_kv_len = compile_spec.max_kv_len or 0 page_size = compile_spec.page_size num_qo_heads = compile_spec.num_qo_heads num_kv_heads = compile_spec.num_kv_heads head_dim = compile_spec.head_dim qk_dtype_key = compile_spec.qk_dtype_key pv_dtype_key = compile_spec.pv_dtype_key - output_dtype_key = compile_spec.output_dtype_key + out_dtype_key = compile_spec.out_dtype_key mask_type = compile_spec.mask_type window_left = compile_spec.window_left head_paired = compile_spec.head_paired @@ -1994,14 +2368,14 @@ def _get_compiled_paged_context( } input_qk_dtype = dtype_map[qk_dtype_key] input_pv_dtype = dtype_map[pv_dtype_key] - output_dtype = dtype_map[output_dtype_key] + out_dtype = dtype_map[out_dtype_key] with torch.cuda.device(device_index): max_active_clusters = get_max_active_clusters(1) fmha = _make_context_kernel( device_index=device_index, input_qk_dtype=input_qk_dtype, input_pv_dtype=input_pv_dtype, - output_dtype=output_dtype, + out_dtype=out_dtype, head_dim=head_dim, mask_type=mask_type, window_left=window_left, @@ -2011,18 +2385,27 @@ def _get_compiled_paged_context( has_q_offset=has_q_offset, causal_single_kv_tile=False, scheduler=scheduler, - uses_ldtm_stat=_default_uses_ldtm_stat(device_index), + uses_ldtm_stat=_default_uses_ldtm_stat( + device_index, input_qk_dtype.width, input_pv_dtype.width + ), fp8_psmem_early_token=_default_fp8_psmem_early_token(device_index), + kv_ring_stages=_default_kv_ring_stages( + device_index, input_qk_dtype.width, input_pv_dtype.width + ), + umma_fp8_k64=_default_umma_fp8_k64(device_index, input_qk_dtype.width), page_size=page_size, - max_kv_len=max_kv_len, + max_kv_len=compile_spec.max_kv_len, ) + fmha.cfg.lse_mode = compile_spec.lse_mode + fmha.cfg.use_attention_sinks = compile_spec.use_attention_sinks fmha.cfg.has_varlen = True fmha.cfg.has_uniform_varlen = uniform_packed_lengths if uniform_packed_lengths: - fmha.cfg.uniform_seq_len_q = max_seq_len_q + fmha.cfg.uniform_seq_len_q = max_q_len fmha.cfg.uniform_seq_len_k = max_kv_len fmha.cfg.has_q_offset = has_q_offset fmha.cfg.exp2_fma_pairs = _default_exp2_fma_pairs(device_index, fmha.cfg) + fmha.cfg.exp2_packed_f16x2 = _default_exp2_packed_f16x2(device_index, fmha.cfg) fmha.cfg.paged_v_tail_is_zero = paged_v_tail_is_zero if fmha.cfg.kv_tile_n != _CONTEXT_KV_TILE_N: raise RuntimeError( @@ -2032,21 +2415,24 @@ def _get_compiled_paged_context( _validate_query_work_tile_span(fmha.cfg) fmha.cfg.packed_dense_k_mask = packed_dense_k_mask + # Bind packed Q, native pages and live token bounds to the context launcher. @cute.jit def tensor_adapter( q: cute.Tensor, k_cache: cute.Tensor, v_cache: cute.Tensor, out: cute.Tensor, - scale_softmax_log2: cute.Tensor, - output_scale: cute.Tensor, + lse: cute.Tensor | None, + attention_sinks: cute.Tensor | None, + scale_softmax_log2: cutlass.Float32, + output_scale: cutlass.Float32, qo_indptr: cute.Tensor, block_tables: cute.Tensor, seq_lens_kv: cute.Tensor, stream: cuda_drv.CUstream, static_max_active_clusters: cutlass.Constexpr[int], - static_max_seq_len_q: cutlass.Constexpr[int], - static_max_kv_len: cutlass.Constexpr[int], + max_q_len: cutlass.Int32, + runtime_max_kv_len: cutlass.Int32, ) -> None: fmha( q, @@ -2058,12 +2444,15 @@ def tensor_adapter( static_max_active_clusters, stream, cum_seqlen_q=qo_indptr, - max_seqlen_q=cutlass.Int32(static_max_seq_len_q), - max_seqlen_k=cutlass.Int32(static_max_kv_len), + max_seqlen_q=max_q_len, + max_seqlen_k=runtime_max_kv_len, block_tables=block_tables, seq_lens_kv=seq_lens_kv, + lse_output=lse, + attention_sinks=attention_sinks, ) + # Describe compact metadata or Q/O storage with byte-aligned pointers. def fake_compact(dtype, shape, assumed_align): return cute.runtime.make_fake_compact_tensor( dtype, @@ -2083,11 +2472,26 @@ def fake_compact(dtype, shape, assumed_align): page_size, head_dim, ) - k_fake = fake_compact(input_qk_dtype, kv_shape, 16) - v_fake = fake_compact(input_pv_dtype, kv_shape, 16) - out_fake = fake_compact(output_dtype, (runtime_total_q, num_qo_heads, head_dim), 16) - scale_fake = fake_compact(cutlass.Float32, (1,), 4) - output_scale_fake = fake_compact(cutlass.Float32, (1,), 4) + + # Keep page/head/token strides dynamic while each head dimension is contiguous. + def fake_paged_cache(dtype): + return cute.runtime.make_fake_tensor( + dtype, + kv_shape, + stride=(cute.sym_int64(), cute.sym_int64(), cute.sym_int64(), 1), + assumed_align=16, + ) + + k_fake = fake_paged_cache(input_qk_dtype) + v_fake = fake_paged_cache(input_pv_dtype) + out_fake = fake_compact(out_dtype, (runtime_total_q, num_qo_heads, head_dim), 16) + lse_fake = ( + fake_compact(cutlass.Float32, out_fake.shape[:-1], 4) + if compile_spec.lse_mode != "none" + else None + ) + scale_fake = cutlass.Float32(1.0) + output_scale_fake = cutlass.Float32(1.0) qo_indptr_fake = fake_compact(cutlass.Int32, (runtime_num_q_offsets,), 4) runtime_page_columns = cute.sym_int() runtime_block_table_row_stride = cute.sym_int() @@ -2098,34 +2502,38 @@ def fake_compact(dtype, shape, assumed_align): assumed_align=4, ) seq_lens_fake = fake_compact(cutlass.Int32, (batch_size,), 4) + sinks_fake = ( + fake_compact(cutlass.Float32, (num_qo_heads,), 4) + if compile_spec.use_attention_sinks + else None + ) stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) - def compile_kernel(): - return cute.compile[cute.FrontendNext]( - tensor_adapter, - q_fake, - k_fake, - v_fake, - out_fake, - scale_fake, - output_scale_fake, - qo_indptr_fake, - block_tables_fake, - seq_lens_fake, - stream_fake, - max_active_clusters, - max_seq_len_q, - max_kv_len, - options=_COMPILE_OPTIONS, - ) - with torch.cuda.device(device_index): compiled = build_and_load_attention_kernel( "context", "paged", compile_spec, _COMPILE_OPTIONS, - compile_kernel, + lambda: cute.compile[cute.FrontendNext]( + tensor_adapter, + q_fake, + k_fake, + v_fake, + out_fake, + lse_fake, + sinks_fake, + scale_fake, + output_scale_fake, + qo_indptr_fake, + block_tables_fake, + seq_lens_fake, + stream_fake, + max_active_clusters, + cutlass.Int32(0), + cutlass.Int32(0), + options=_COMPILE_OPTIONS, + ), extra=(max_active_clusters,), ) policy = ( @@ -2150,7 +2558,7 @@ def _validate_runtime_inputs( *, qo_indptr: Optional[torch.Tensor], kv_indptr: Optional[torch.Tensor], -) -> None: +) -> tuple[int, int]: """Validate contiguous data and request metadata against one plan.""" _validate_base_tensors(q, k, v, allow_mla=True) @@ -2166,7 +2574,7 @@ def _validate_runtime_inputs( if tuple(v.shape) != v_shape: raise ValueError(f"v must have shape {v_shape}, got {tuple(v.shape)}") - if geometry.packed: + if geometry.packed_query: if q.ndim != 3 or tuple(q.shape[1:]) != ( geometry.num_qo_heads, geometry.head_dim, @@ -2189,20 +2597,16 @@ def _validate_runtime_inputs( total_k = int(k.shape[0]) _validate_padded_data_extent(total_q, "total_q") _validate_padded_data_extent(total_k, "total_k") - if total_q < geometry.batch_size or total_q > ( - geometry.batch_size * geometry.max_seq_len_q - ): + if total_q > (geometry.batch_size * geometry.max_q_len): raise ValueError( - "q must contain between batch_size and " - "batch_size * max_seq_len_q rows; got " + "q must contain no more than " + "batch_size * max_q_len rows; got " f"{total_q} for batch_size={geometry.batch_size} and " - f"max_seq_len_q={geometry.max_seq_len_q}" + f"max_q_len={geometry.max_q_len}" ) - if total_k < geometry.batch_size or total_k > ( - geometry.batch_size * geometry.max_seq_len_k - ): + if total_k > (geometry.batch_size * geometry.max_seq_len_k): raise ValueError( - "k and v must contain between batch_size and " + "k and v must contain no more than " "batch_size * max_kv_len rows; got " f"{total_k} for batch_size={geometry.batch_size} and " f"max_kv_len={geometry.max_seq_len_k}" @@ -2227,9 +2631,7 @@ def _validate_runtime_inputs( ) request_lengths[name] = lengths capacity = ( - geometry.max_seq_len_q - if name == "qo_indptr" - else geometry.max_seq_len_k + geometry.max_q_len if name == "qo_indptr" else geometry.max_seq_len_k ) for batch_idx, length in enumerate(lengths): if length > capacity: @@ -2243,7 +2645,7 @@ def _validate_runtime_inputs( for batch_idx, (q_length, k_length) in enumerate( zip(q_lengths, k_lengths, strict=True) ): - if q_length > k_length: + if k_length > 0 and q_length > k_length: raise ValueError( "bottom-right causal context requires Sq <= Sk for each " f"request; got batch {batch_idx}: Sq={q_length}, " @@ -2252,18 +2654,21 @@ def _validate_runtime_inputs( _validate_compact(q, "q", "[total_q, Hq, D]") _validate_compact(k, "k", "[total_k, Hkv, D]") _validate_compact(v, "v", "[total_k, Hkv, D]") + return max(q_lengths), max(k_lengths) else: if qo_indptr is not None or kv_indptr is not None: - raise ValueError("qo_indptr and kv_indptr require packed=True in plan()") + raise ValueError( + "qo_indptr and kv_indptr require packed_query=True in plan()" + ) q_shape = ( geometry.batch_size, - geometry.max_seq_len_q, + geometry.max_q_len, geometry.num_qo_heads, geometry.head_dim, ) kv_shape = ( geometry.batch_size, - geometry.max_seq_len_k, + _fixed_kv_rows(geometry.max_seq_len_k, geometry.vc), geometry.num_kv_heads, geometry.head_dim, ) @@ -2276,6 +2681,7 @@ def _validate_runtime_inputs( _validate_compact(q, "q", "[B, Sq, Hq, D]") _validate_compact(k, "k", "[B, Sk, Hkv, D]") _validate_compact(v, "v", "[B, Sk, Hkv, D]") + return int(q.shape[1]), geometry.max_seq_len_k def _validate_paged_runtime_inputs( @@ -2319,18 +2725,16 @@ def _validate_paged_runtime_inputs( ) total_q = int(q.shape[0]) _validate_padded_data_extent(total_q, "total_q") - if total_q < geometry.batch_size or total_q > ( - geometry.batch_size * geometry.max_seq_len_q - ): + if total_q > (geometry.batch_size * geometry.max_q_len): raise ValueError( - "q must contain between batch_size and batch_size * max_seq_len_q " + "q must contain no more than batch_size * max_q_len " f"rows; got {total_q} rows for batch_size={geometry.batch_size} and " - f"max_seq_len_q={geometry.max_seq_len_q}" + f"max_q_len={geometry.max_q_len}" ) _validate_extent(int(k_cache.shape[0]), "num_physical_pages") _validate_compact(q, "q", "[total_q, Hq, D]") - _validate_compact(k_cache, "k_cache", "[num_pages, Hkv, page_size, D]") - _validate_compact(v_cache, "v_cache", "[num_pages, Hkv, page_size, D]") + _validate_native_hnd_tma_strides(k_cache, "k_cache") + _validate_native_hnd_tma_strides(v_cache, "v_cache") def _validate_paged_runtime_metadata( @@ -2341,7 +2745,7 @@ def _validate_paged_runtime_metadata( *, total_q: int, num_physical_pages: int, -) -> None: +) -> tuple[int, int]: """Validate per-run fixed-table storage and active metadata values.""" _validate_indptr_tensor(qo_indptr, "qo_indptr", device=geometry.device) @@ -2369,30 +2773,30 @@ def _validate_paged_runtime_metadata( ) _, q_lengths = _read_indptr(qo_indptr, "qo_indptr", expected_total=total_q) - seq_lens = _read_int32_values( + seq_lens_kv = _read_int32_values( seq_lens_kv, "seq_lens_kv", expected_count=geometry.batch_size, ) block_table_rows = block_tables.tolist() for batch_idx, (q_length, kv_length, block_table_row) in enumerate( - zip(q_lengths, seq_lens, block_table_rows, strict=True) + zip(q_lengths, seq_lens_kv, block_table_rows, strict=True) ): - if q_length > geometry.max_seq_len_q: + if q_length > geometry.max_q_len: raise ValueError( - "qo_indptr deltas must not exceed max_seq_len_q; " + "qo_indptr deltas must not exceed max_q_len; " f"batch {batch_idx} has {q_length}, maximum is " - f"{geometry.max_seq_len_q}" + f"{geometry.max_q_len}" ) - if geometry.uniform_packed_lengths and q_length != geometry.max_seq_len_q: + if geometry.uniform_packed_lengths and q_length != geometry.max_q_len: raise ValueError( "uniform_packed_lengths=True requires every qo_indptr delta " - f"to equal max_seq_len_q={geometry.max_seq_len_q}; batch " + f"to equal max_q_len={geometry.max_q_len}; batch " f"{batch_idx} has {q_length}" ) - if kv_length <= 0 or kv_length > geometry.max_kv_len: + if kv_length < 0 or kv_length > geometry.max_kv_len: raise ValueError( - "seq_lens_kv entries must be positive and not exceed " + "seq_lens_kv entries must be nonnegative and not exceed " f"max_kv_len={geometry.max_kv_len}; batch {batch_idx} has " f"{kv_length}" ) @@ -2403,7 +2807,7 @@ def _validate_paged_runtime_metadata( f"{batch_idx} has {kv_length}" ) required_pages = (kv_length + geometry.page_size - 1) // geometry.page_size - if geometry.mask_type == "causal" and q_length > kv_length: + if geometry.mask_type == "causal" and kv_length > 0 and q_length > kv_length: raise ValueError( "bottom-right causal context requires Sq <= Sk for each " f"request; got batch {batch_idx}: Sq={q_length}, Sk={kv_length}" @@ -2424,43 +2828,24 @@ def _validate_paged_runtime_metadata( f"page pool; batch {batch_idx}, row entry {row_offset} is " f"{page_idx}, pool has {num_physical_pages} pages" ) - - -def _validate_runtime_scale_tensor( - scale: torch.Tensor, - name: str, - *, - device: torch.device, -) -> None: - """Validate one allocation-free per-run scale input.""" - - _validate_tensor(scale, name) - if scale.device != device: - raise ValueError(f"{name} must be on {device}, got {scale.device}") - if scale.dtype != torch.float32: - raise TypeError(f"{name} must have dtype torch.float32") - if tuple(scale.shape) != (1,): - raise ValueError(f"{name} must have shape [1], got {tuple(scale.shape)}") - _validate_compact(scale, name, "[1]") - _validate_alignment(scale, name, 4) - _validate_scale(scale.item(), name) + return max(q_lengths), max(seq_lens_kv) def _prepare_out( out: Optional[torch.Tensor], *, q: torch.Tensor, - output_dtype: torch.dtype, + out_dtype: torch.dtype, head_dim_vo: int | None = None, ) -> torch.Tensor: out_shape = (*q.shape[:-1], head_dim_vo or q.shape[-1]) if out is None: - return torch.empty(out_shape, dtype=output_dtype, device=q.device) + return torch.empty(out_shape, dtype=out_dtype, device=q.device) _validate_tensor(out, "out") if tuple(out.shape) != out_shape: raise ValueError(f"out must have shape {out_shape}, got {tuple(out.shape)}") - if out.dtype != output_dtype: - raise ValueError(f"out must have dtype {output_dtype}, got {out.dtype}") + if out.dtype != out_dtype: + raise ValueError(f"out must have dtype {out_dtype}, got {out.dtype}") if out.device != q.device: raise ValueError(f"out must be on {q.device}, got {out.device}") layout = "[total_q, Hq, D]" if q.ndim == 3 else "[B, Sq, Hq, D]" @@ -2469,33 +2854,46 @@ def _prepare_out( return out +def _variable_window_buffers(geometry, empty_i32): + if geometry.mask_type != "variable_window": + return None, empty_i32 + tile_size_q = _variable_window_tile_size_q(geometry.head_dim) + num_seq_tiles = (geometry.max_q_len + tile_size_q - 1) // tile_size_q + padded_rows = num_seq_tiles * tile_size_q + cta_starts = torch.empty( + geometry.batch_size * num_seq_tiles, dtype=torch.int32, device=geometry.device + ) + padded_starts = ( + torch.empty( + (geometry.batch_size, padded_rows), + dtype=torch.int32, + device=geometry.device, + ) + if padded_rows != geometry.max_q_len + else None + ) + return padded_starts, cta_starts + + +def _cuda_scale(value: torch.Tensor, name: str, device: torch.device) -> torch.Tensor: + """Validate a live device scalar without reading it back to the host.""" + if not isinstance(value, torch.Tensor): + raise TypeError(f"{name} must be a CUDA FP32 tensor") + if value.device != device or value.dtype != torch.float32 or value.numel() != 1: + raise ValueError(f"{name} must contain one FP32 value on {device}") + return value.reshape(1) + + class BatchPrefillTSWrapper: - """Compile and reuse fixed or packed-ragged contiguous context attention. - - Supports equal QK/V dimensions 128 and 256, and separate QK=192/V=128 - for non-absorbed MLA. Output has Q's leading dimensions and V's head - dimension. The 192/128 geometry does not support a positive left window. - - ``plan`` accepts only static compilation geometry. Q/K/V tensors, packed - cumulative offsets, variable-window metadata, and optional scale overrides - are supplied to ``run``. Variable-window metadata may include precomputed - CTA minima. Packed offset values may change between launches while their - request lengths stay within the planned capacities. - - Plan-time scalar defaults are stored as one-element device tensors. - Variable-window plans also own private fallback reduction scratch refreshed - when a run omits precomputed CTA minima; only launches using that scratch - must not overlap across streams or captured graphs. With caller-owned - output, ``validate=False`` performs no allocation or metadata readback and - is CUDA-graph-capturable when the caller guarantees the complete runtime - contract and stable addresses. Replanning replaces plan-owned tensors and - invalidates graphs captured from the previous plan; all prior launches and - replays must finish before ``plan`` is called again. Keep the wrapper and - all captured runtime tensors alive until every graph using the current - plan is destroyed. Before running on a CUDA stream that is not already - ordered after the planning stream, the caller must establish that - dependency. - """ + """Reusable contiguous context attention with eager planning. + + run() accepts fixed BSHD or packed THD Q/K/V. Plan owns model defaults and + variable-window fallback metadata, while every run supplies request tensors. + Supports equal D128/D192/D256 or contiguous non-absorbed MLA QK192/V128 geometry. + Dense fixed D128 Sage plans bind prepared per-block/channel scale tensors. + Replanning invalidates graphs; complete prior launches and replays first. + VC plans also own a refreshed per-channel output-scale buffer; their runs + must not overlap.""" @flashinfer_experimental_api def __init__(self) -> None: @@ -2508,7 +2906,7 @@ def plan( *, device: int | str | torch.device, batch_size: int, - max_seq_len_q: int, + max_q_len: int, max_kv_len: int, num_qo_heads: int, num_kv_heads: int, @@ -2518,65 +2916,51 @@ def plan( k_dtype: torch.dtype, v_dtype: Optional[torch.dtype] = None, out_dtype: Optional[torch.dtype] = None, - packed: bool = False, + packed_query: bool = False, mask_type: Literal["dense", "causal", "variable_window"] = "dense", window_left: int = -1, - sm_scale: Optional[float] = None, + softmax_scale: Optional[float] = None, output_scale: float = 1.0, + vc_config: Optional[VCAttentionConfig] = None, + lse_mode: Literal["none", "base2"] = "none", + use_attention_sinks: bool = False, + sage: Optional[SageAttentionConfig] = None, + cuda_scales: bool = False, ) -> None: - """Compile one reusable specialization from explicit static geometry. - - Packed plans conservatively support any positive per-request Q and K/V - lengths bounded by ``max_seq_len_q`` and ``max_kv_len``. Fixed plans - use those bounds as their exact tensor extents. ``batch_size`` is exact. - Calling ``plan`` again replaces all plan-owned tensors and invalidates - CUDA graphs captured from the previous plan. Complete every launch and - replay that uses the previous plan before replanning. Planning - allocates and compiles; complete it before CUDA Graph capture. - - Parameters - ---------- - device : int, str, or torch.device - CUDA device on which the specialization is compiled and run. - batch_size : int - Exact number of requests in every run. - max_seq_len_q : int - Maximum query length of any request. - max_kv_len : int - Maximum K/V length of any request. - num_qo_heads : int - Number of query/output heads. - num_kv_heads : int - Number of key/value heads. - head_dim : int - Query/key head dimension. - head_dim_vo : int, optional - Value/output head dimension; defaults to ``head_dim``. Separate - dimensions are supported for non-absorbed MLA with QK=192/V=128. - q_dtype : torch.dtype - Query dtype. - k_dtype : torch.dtype - Key dtype, and the value dtype when ``v_dtype`` is omitted. It - must currently equal ``q_dtype``. - v_dtype : torch.dtype, optional - Value dtype; defaults to ``k_dtype``. May differ from - ``k_dtype`` only for the QK-BF16/PV-FP8 combination. - out_dtype : torch.dtype, optional - Output dtype; defaults to ``q_dtype``. - packed : bool - Whether run tensors use packed ``[total_tokens, H, D]`` storage. - mask_type : {"dense", "causal", "variable_window"} - Attention mask mode. ``variable_window`` is supported only for - fixed-shape inputs. - window_left : int - Left sliding-window extent, or ``-1`` to disable the window. - sm_scale : float, optional - Softmax scale; defaults to the inverse square root of head size. - output_scale : float - Scale applied to the attention output. + """Select/cache one execution policy, obtain compiled kernels, and bind resources. + + ``vc_config`` selects a VC-Attention recipe and the planned Q/K dtype + picks its Q/K side. bf16 Q/K runs VC-Attention-QK16 and E4M3 + Q/K runs VC-Attention-QK8, whose Hadamard-rotated, per-block scaled Q + and K need ``vc.q_scale`` and ``vc.k_scale`` on every :meth:`run`. V is + E4M3 in both, stored as per-128-token-tile residuals whose bf16 tile + means the kernel restores. Every :meth:`run` then supplies the permuted + K, the E4M3 V and ``vc=VCAttentionParams(...)``. + + batch_size is exact. Fixed contiguous inputs use exact Q/KV tensor extents; + packed/paged request lengths may vary within max_q_len/max_kv_len. + Complete planning outside CUDA graph capture. Replanning replaces owned + resources and invalidates previous graphs; finish their launches first. + lse_mode='base2' enables final LSE stores, and use_attention_sinks enables + a per-run sink binding. Context scales are positive host scalar defaults. + Sage plans use BF16 output by default and require per-run Sage scales. + Sage V-mean restoration cannot be combined with attention sinks. + cuda_scales binds one-element CUDA FP32 scales for graph-replayed + per-tensor quantization; it cannot be combined with Sage or VC plans. + Equivalent plans reuse policy and compiled-kernel caches. + VC runs sharing a plan must not overlap. """ - resolved_out_dtype = q_dtype if out_dtype is None else out_dtype + lse_mode = validate_lse_mode(lse_mode) + if not isinstance(use_attention_sinks, bool): + raise TypeError("use_attention_sinks must be a bool") + if out_dtype is not None: + resolved_out_dtype = out_dtype + elif sage is not None: + # As in block-sparse Sage plans; an 8-bit output is not a recipe. + resolved_out_dtype = torch.bfloat16 + else: + resolved_out_dtype = q_dtype resolved_v_dtype = k_dtype if v_dtype is None else v_dtype if k_dtype != q_dtype: raise NotImplementedError( @@ -2586,7 +2970,7 @@ def plan( geometry = _resolve_context_plan_geometry( device=device, batch_size=batch_size, - max_seq_len_q=max_seq_len_q, + max_q_len=max_q_len, max_kv_len=max_kv_len, num_qo_heads=num_qo_heads, num_kv_heads=num_kv_heads, @@ -2594,60 +2978,107 @@ def plan( head_dim_vo=head_dim_vo, qk_dtype=k_dtype, pv_dtype=resolved_v_dtype, - packed=packed, + packed_query=packed_query, mask_type=mask_type, window_left=window_left, - output_dtype=resolved_out_dtype, + out_dtype=resolved_out_dtype, + sage=sage, ) - if sm_scale is None: - sm_scale = 1.0 / math.sqrt(geometry.head_dim) - sm_scale = _validate_scale(sm_scale, "sm_scale") - output_scale = _validate_scale(output_scale, "output_scale") + if sage is not None and sage.v_mean and use_attention_sinks: + raise NotImplementedError( + "Sage V-mean restoration does not support attention sinks" + ) + if vc_config is not None and not isinstance(vc_config, VCAttentionConfig): + raise TypeError("vc_config must be a VCAttentionConfig instance") + if vc_config is not None and (packed_query or mask_type != "dense"): + raise ValueError( + "vc_config requires a dense contiguous fixed-length plan; VC-Attention " + "has no packed, causal or variable-window kernel" + ) + if vc_config is not None and (lse_mode != "none" or use_attention_sinks): + # VC uses quantized P sums; ordinary context LSE/sink math does not apply. + raise NotImplementedError( + "VC-Attention does not support LSE or attention sinks" + ) + if not isinstance(cuda_scales, bool): + raise TypeError("cuda_scales must be a bool") + if cuda_scales and (vc_config is not None or sage is not None): + raise ValueError("cuda_scales requires a plain context plan") + geometry = replace(geometry, vc=vc_config) + if softmax_scale is None: + softmax_scale = geometry.head_dim**-0.5 scale_softmax_log2 = _validate_scale( - sm_scale * math.log2(math.e), "sm_scale * log2(e)" - ) - scale_tensor = torch.tensor( - [scale_softmax_log2], dtype=torch.float32, device=geometry.device - ) - output_scale_tensor = torch.tensor( - [output_scale], dtype=torch.float32, device=geometry.device + _validate_scale(softmax_scale, "softmax_scale") * math.log2(math.e), + "softmax_scale * log2(e)", ) - empty_i32 = torch.empty(1, dtype=torch.int32, device=geometry.device) - if geometry.mask_type == "variable_window": - tile_size_q = _variable_window_tile_size_q(geometry.head_dim) - num_seq_tiles = (geometry.max_seq_len_q + tile_size_q - 1) // tile_size_q - padded_rows = num_seq_tiles * tile_size_q - variable_window_cta_starts = torch.empty( - geometry.batch_size * num_seq_tiles, - dtype=torch.int32, - device=geometry.device, + output_scale = _validate_scale(output_scale, "output_scale") + if cuda_scales: + scale_softmax_log2 = torch.tensor( + [scale_softmax_log2], dtype=torch.float32, device=geometry.device ) - variable_window_padded_starts = ( - torch.empty( - (geometry.batch_size, padded_rows), - dtype=torch.int32, - device=geometry.device, - ) - if padded_rows != geometry.max_seq_len_q - else None + output_scale = torch.tensor( + [output_scale], dtype=torch.float32, device=geometry.device ) - else: - variable_window_cta_starts = empty_i32 - variable_window_padded_starts = None + vc_output_scale_buffer = ( + torch.empty( + ( + geometry.batch_size, + geometry.num_qo_heads, + geometry.head_dim_vo or geometry.head_dim, + ), + dtype=torch.float32, + device=geometry.device, + ) + if vc_config is not None + else None + ) + empty_i32 = torch.empty(1, dtype=torch.int32, device=geometry.device) + variable_window_padded_starts, variable_window_cta_starts = ( + _variable_window_buffers(geometry, empty_i32) + ) # Keep all runtime tensor allocation ahead of CUTLASS JIT. Besides # making plan publication atomic, this lets compute-sanitizer patch the # generated attention kernel without interleaving later PyTorch setup # launches with the DSL compiler/runtime callbacks. - compiled, policy = _get_compiled_context(_context_compile_spec(geometry)) + vc_shapes = None + if vc_config is not None: + vc_shapes = vc_scale_shapes( + vc_config, + batch_size=geometry.batch_size, + seq_len_kv=geometry.max_seq_len_k, + num_kv_heads=geometry.num_kv_heads, + head_dim=geometry.head_dim_vo or geometry.head_dim, + seq_len_q=geometry.max_q_len, + num_qo_heads=geometry.num_qo_heads, + qk_fp8=geometry.qk_dtype == torch.float8_e4m3fn, + ) + sage_state = ( + _build_context_sage_state(geometry) if geometry.sage is not None else None + ) + compiled, policy = _get_compiled_context( + _context_compile_spec( + geometry, + runtime_lengths=vc_config is None, + lse_mode=lse_mode, + use_attention_sinks=use_attention_sinks, + cuda_scales=cuda_scales, + ) + ) self._plan_state = _ContextPlanState( geometry=geometry, - scale_softmax_log2=scale_tensor, - output_scale=output_scale_tensor, + scale_softmax_log2=scale_softmax_log2, + output_scale=output_scale, empty_i32=empty_i32, variable_window_padded_starts=variable_window_padded_starts, variable_window_cta_starts=variable_window_cta_starts, compiled=compiled, policy=policy, + vc_scale_shapes=vc_shapes, + vc_output_scale=vc_output_scale_buffer, + cuda_scales=cuda_scales, + lse_mode=lse_mode, + use_attention_sinks=use_attention_sinks, + sage=sage_state, ) @flashinfer_experimental_api @@ -2662,50 +3093,74 @@ def run( variable_window_token_starts: Optional[torch.Tensor] = None, variable_window_token_ends: Optional[torch.Tensor] = None, variable_window_cta_starts: Optional[torch.Tensor] = None, + max_q_len: int | None = None, + max_kv_len: int | None = None, out: Optional[torch.Tensor] = None, - scale_softmax_log2: Optional[torch.Tensor] = None, - output_scale: Optional[torch.Tensor] = None, + lse: Optional[torch.Tensor] = None, + return_lse: bool = False, + attention_sinks: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + output_scale: float | torch.Tensor | None = None, + scale_softmax_log2: torch.Tensor | None = None, + vc: Optional[VCAttentionParams] = None, + sage: Optional[SageAttentionParams] = None, validate: bool = True, - ) -> torch.Tensor: - """Launch the compiled specialization with per-request tensors. - - Parameters - ---------- - q, k, v : torch.Tensor - Runtime query, key, and value tensors matching the plan. - qo_indptr, kv_indptr : torch.Tensor, optional - Cumulative packed Q and K/V offsets. Required by packed plans and - forbidden by fixed plans. - variable_window_token_starts, variable_window_token_ends : torch.Tensor, optional - Inclusive K-token bounds for every fixed Q row. Required only for - ``mask_type='variable_window'`` and shaped ``[B, max_seq_len_q]``. - variable_window_cta_starts : torch.Tensor, optional - Precomputed exact minimum start bound for each Q work tile, shaped - ``[B, ceil(max_seq_len_q / tile_size_q)]`` where ``tile_size_q`` - is 256 for head dimension 128 and 128 for head dimension 256. - Valid only for ``mask_type='variable_window'``. When omitted, the - wrapper derives these values from ``variable_window_token_starts``. - out : torch.Tensor, optional - Caller-owned output tensor. A new tensor is allocated when omitted. - scale_softmax_log2 : torch.Tensor, optional - Per-run one-element float32 softmax scale in base-2 form. Defaults - to the scale retained by the plan. - output_scale : torch.Tensor, optional - Per-run one-element float32 output scale. Defaults to the plan. - validate : bool - Validate tensor structure and metadata values. Disable only when - the caller guarantees these contracts. Output must not overlap - any live input or plan-owned buffer in either mode; storage - overlap and lifetime are caller preconditions. Defaults to ``True``. - """ + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Launch the prepared contiguous attention with live request tensors. + + Uses the policy and compiled kernels prepared by plan(). Batch and query + shape must match that plan; live K/V lengths stay within its capacity. + Supply output/LSE buffers and validate=False during graph capture, and + retain the wrapper and captured buffers for every replay. + + An enabled lse_mode="base2" plan writes compact FP32 log2(sum(exp(scores))) + with the query shape minus its final dimension. return_lse controls tuple + return only; an omitted enabled buffer is allocated. Output/LSE/workspace must + be disjoint from inputs and live buffers. One wrapper serves one ordered lane. + + max_q_len/max_kv_len are launch bounds within plan capacities. + Validation infers omitted values from existing metadata; trusted packed + calls use plan bounds. Exact runtime lengths do not specialize kernels. + + Packed inputs require qo_indptr and kv_indptr. Variable-window fixed inputs + require inclusive start/end bounds; ends must be nondecreasing within each + request, while starts may be unordered. Optional CTA minima avoid refreshing + wrapper-owned fallback scratch. Scale overrides are finite positive host + scalars: softmax_scale is the effective QK scale, and + output_scale scales V/output. They are runtime kernel operands, not JIT keys. + With cuda_scales, supply scale_softmax_log2 and output_scale as + one-element CUDA FP32 tensors instead. Their values remain on device + and must be positive and finite, including after graph updates. + + Enabled attention_sinks supplies compact CUDA FP32 [Hq] natural-log sink + logits after QK scaling. Values must be finite or -inf; validation checks only + metadata. Each sink adds denominator mass with zero V contribution. + Sage plans require matching SageAttentionParams; their prepared scales + are refreshed on the current stream before attention. + VC plans require vc=VCAttentionParams(...) with permuted K and residual V; + its per-channel V scale multiplies output_scale in the plan-owned buffer.""" + + if not isinstance(validate, bool): + raise TypeError("validate must be a bool") state = self._plan_state if state is None: raise RuntimeError("plan() must be called before run()") - if not isinstance(validate, bool): - raise TypeError("validate must be a bool") geometry = state.geometry - if geometry.packed: + if geometry.mask_type != "variable_window" and ( + variable_window_token_starts is not None + or variable_window_token_ends is not None + or variable_window_cta_starts is not None + ): + raise ValueError( + "variable-window metadata requires mask_type='variable_window'" + ) + if not validate and not geometry.packed_query: + max_q_len = int(q.shape[1]) + max_kv_len = ( + geometry.max_seq_len_k if geometry.vc is not None else int(k.shape[1]) + ) + if geometry.packed_query: if qo_indptr is None or kv_indptr is None: raise ValueError( "packed plans require qo_indptr and kv_indptr in run()" @@ -2715,13 +3170,41 @@ def run( else: if qo_indptr is not None or kv_indptr is not None: raise ValueError( - "qo_indptr and kv_indptr require packed=True in plan()" + "qo_indptr and kv_indptr require packed_query=True in plan()" ) runtime_qo_indptr = state.empty_i32 runtime_kv_indptr = state.empty_i32 + if (vc is None) == (geometry.vc is not None): + raise ValueError( + "vc=VCAttentionParams(...) is required by a VC plan and rejected " + "by a plan without vc_config" + ) + vc_output_scale: Optional[torch.Tensor] = None + vc_mu: Optional[torch.Tensor] = None + vc_q_scale: Optional[torch.Tensor] = None + vc_k_scale: Optional[torch.Tensor] = None + if vc is not None: + if validate: + validate_vc_params(vc, state.vc_scale_shapes, device=geometry.device) + if geometry.qk_dtype == torch.float8_e4m3fn: + if vc.q_scale is None or vc.k_scale is None: + raise ValueError( + "VC-Attention-QK8 plans require vc.q_scale and vc.k_scale" + ) + vc_q_scale = vc.q_scale.view(geometry.num_qo_heads, -1) + vc_k_scale = vc.k_scale.view(geometry.num_kv_heads, -1) + elif vc.q_scale is not None or vc.k_scale is not None: + raise ValueError( + "vc.q_scale and vc.k_scale belong to VC-Attention-QK8 plans with " + "E4M3 Q/K" + ) + vc_mu = vc.tile_means + # VC plans have equal Q and K/V head counts, so the per-channel V + # scale is already in output order. + vc_output_scale = vc.v_scale if validate: - _validate_runtime_inputs( + actual_q, actual_kv = _validate_runtime_inputs( q, k, v, @@ -2729,6 +3212,26 @@ def run( qo_indptr=qo_indptr, kv_indptr=kv_indptr, ) + from .decode import resolve_max_sequence_length + + max_q_len = resolve_max_sequence_length( + max_q_len, actual_q, "max_q_len", geometry.max_q_len + ) + max_kv_len = resolve_max_sequence_length( + max_kv_len, actual_kv, "max_kv_len", geometry.max_seq_len_k + ) + else: + if max_q_len is None: + max_q_len = geometry.max_q_len + if max_kv_len is None: + max_kv_len = geometry.max_seq_len_k + if (sage is None) != (state.sage is None): + raise ValueError( + "sage=SageAttentionParams(...) is required by a Sage plan and " + "rejected by a plan without Sage attention" + ) + if sage is not None and validate: + validate_sage_params(sage, state.sage.scale_shapes, device=geometry.device) if geometry.mask_type == "variable_window": if ( variable_window_token_starts is None @@ -2764,52 +3267,83 @@ def run( else: runtime_window_cta_starts = variable_window_cta_starts.flatten() else: - if ( - variable_window_token_starts is not None - or variable_window_token_ends is not None - or variable_window_cta_starts is not None - ): - raise ValueError( - "variable-window metadata requires mask_type='variable_window'" - ) runtime_window_starts = state.empty_i32 runtime_window_ends = state.empty_i32 runtime_window_cta_starts = state.empty_i32 - if scale_softmax_log2 is None: - scale_softmax_log2 = state.scale_softmax_log2 - elif validate: - _validate_runtime_scale_tensor( - scale_softmax_log2, + if state.cuda_scales: + if softmax_scale is not None: + raise ValueError("cuda_scales uses the device scale_softmax_log2 operand") + scale_softmax_log2 = _cuda_scale( + state.scale_softmax_log2 if scale_softmax_log2 is None else scale_softmax_log2, "scale_softmax_log2", - device=geometry.device, + geometry.device, ) - if output_scale is None: - output_scale = state.output_scale - elif validate: - _validate_runtime_scale_tensor( - output_scale, + output_scale = _cuda_scale( + state.output_scale if output_scale is None else output_scale, "output_scale", - device=geometry.device, + geometry.device, + ) + else: + if scale_softmax_log2 is not None: + raise ValueError("device scale_softmax_log2 requires cuda_scales=True") + scale_softmax_log2 = ( + state.scale_softmax_log2 + if softmax_scale is None + else _validate_scale( + _validate_scale(softmax_scale, "softmax_scale") * math.log2(math.e), + "softmax_scale * log2(e)", + ) + ) + output_scale = ( + state.output_scale + if output_scale is None + else _validate_scale(output_scale, "output_scale") + ) + if vc_output_scale is not None: + output_scale = torch.mul( + vc_output_scale, output_scale, out=state.vc_output_scale ) if out is None: out = torch.empty( (*q.shape[:-1], geometry.head_dim_vo or geometry.head_dim), - dtype=geometry.output_dtype, + dtype=geometry.out_dtype, device=q.device, ) elif validate: out = _prepare_out( out, q=q, - output_dtype=geometry.output_dtype, + out_dtype=geometry.out_dtype, head_dim_vo=geometry.head_dim_vo, ) + lse = prepare_lse( + lse, + q=q, + mode=state.lse_mode, + return_lse=return_lse, + validate=validate, + ) + attention_sinks = validate_attention_sinks( + attention_sinks, + q=q, + enabled=state.use_attention_sinks, + validate=validate, + ) + if sage is not None: + state.sage.prepare( + sage.q_scale, sage.k_scale, sage.v_scale, sage.v_mean, state.sage.scales + ) + runtime_sage_scales = state.sage.scales + else: + runtime_sage_scales = None state.compiled( q, k, v, out, + lse, + attention_sinks, scale_softmax_log2, output_scale, runtime_qo_indptr, @@ -2817,48 +3351,27 @@ def run( runtime_window_starts, runtime_window_ends, runtime_window_cta_starts, + vc_mu, + vc_q_scale, + vc_k_scale, + runtime_sage_scales, + max_q_len, + max_kv_len, ) - return out + return (out, lse) if return_lse else out class BatchPrefillPagedTSWrapper: - """Compile and reuse packed-Q context attention over HND paged K/V. - - ``plan`` accepts only static compilation geometry. All request metadata is - supplied to ``run``: cumulative Q offsets, a fixed row-strided K/V page - table, and logical K/V lengths. Metadata values may change between launches - while staying within the planned capacities. The wrapper owns no context - workspace and does not transform the fixed page table. - - Plan-time scalar defaults are stored as one-element device tensors. - ``run`` may replace either scale tensor without changing the compiled - specialization. Plans default to a conservative dynamic-length contract; - callers may instead promise exact-uniform packed lengths or no causal Q - offset to compile one narrower specialization. Callers whose cache - preparation zeroes every unused row in an active final V page may also - compile out the defensive consumer-side V-tail clear. Validation reads - request metadata back to the host, checks the length promises, and may - synchronize; it does not inspect V-cache values. With caller-owned output, - ``validate=False`` performs no allocation, metadata readback, or - synchronization and is suitable for CUDA graph capture only when the caller - already enforces the selected contract. - Replanning invalidates captured graphs; prior launches and replays must - finish before ``plan`` is called again. Keep the wrapper and all captured - runtime tensors alive until every graph using the current plan is destroyed. - Before running on a CUDA stream that is not already ordered after the - planning stream, the caller must establish that dependency. - """ + """Reusable packed-Q context attention over HND or NHD paged K/V. - @flashinfer_experimental_api - def __init__(self, kv_layout: Literal["HND"] = "HND") -> None: - """Create an unplanned paged-context wrapper. - - Parameters - ---------- - kv_layout : {"HND"} - K/V layout. Only ``"HND"`` is supported. - """ + plan() retains model semantics and optional exact-uniform/zero-offset/zero-V-tail + promises. run() supplies Q offsets, a row-strided page table, KV lengths, and + host execution bounds. Cache normalization changes metadata without copying. + Replanning invalidates graphs; complete prior launches and replays first.""" + @flashinfer_experimental_api + def __init__(self, kv_layout: Literal["HND", "NHD"] = "HND") -> None: + """Initialize an unplanned paged-context wrapper with a K/V layout.""" _validate_kv_layout(kv_layout) self._kv_layout = kv_layout self._plan_state: Optional[_PagedContextPlanState] = None @@ -2869,7 +3382,7 @@ def plan( *, device: int | str | torch.device, batch_size: int, - max_seq_len_q: int, + max_q_len: int, max_kv_len: int, num_qo_heads: int, num_kv_heads: int, @@ -2881,87 +3394,29 @@ def plan( page_size: int = _DEFAULT_PAGED_KV_PAGE_SIZE, mask_type: Literal["dense", "causal"] = "dense", window_left: int = -1, - sm_scale: Optional[float] = None, + softmax_scale: Optional[float] = None, output_scale: float = 1.0, + lse_mode: Literal["none", "base2"] = "none", + use_attention_sinks: bool = False, uniform_packed_lengths: bool = False, has_q_offset: bool = True, paged_v_tail_is_zero: bool = False, ) -> None: - """Compile one reusable specialization from explicit static geometry. - - Sequence lengths and page indices are deliberately absent. By default, - the plan supports dynamic positive lengths bounded by - ``max_seq_len_q`` and ``max_kv_len``. ``uniform_packed_lengths=True`` - instead promises that every runtime Q length equals ``max_seq_len_q`` - and every runtime K/V length equals ``max_kv_len``. - ``has_q_offset=False`` promises that every causal request has - ``Sq == Sk``. Dense attention ignores and canonicalizes the latter - flag. ``paged_v_tail_is_zero=True`` separately promises that the unused - rows following each request's logical K/V length in its active final V - page are zero. The selected contract compiles exactly one - specialization. ``run(validate=True)`` checks the length promises but - cannot cheaply inspect V-cache contents; callers always own the V-tail - promise. Violating a selected promise can produce incorrect results or - invalid memory accesses. - - ``batch_size`` is exact. Runtime page-table rows must expose at least - ``ceil(max_kv_len / page_size)`` columns, including inactive padding - columns. Calling ``plan`` again replaces plan-owned tensors and - invalidates CUDA graphs captured from the previous plan. Complete every - prior launch and replay before replanning. Planning allocates and - compiles; complete it before CUDA Graph capture. - - Parameters - ---------- - device : int, str, or torch.device - CUDA device on which the specialization is compiled and run. - batch_size : int - Exact number of requests in every run. - max_seq_len_q : int - Maximum query length of any request. - max_kv_len : int - Maximum logical K/V length of any request. - num_qo_heads : int - Number of query/output heads. - num_kv_heads : int - Number of key/value heads. ``num_qo_heads`` must be divisible by - this value. - head_dim : int - Query, key, value, and output head dimension. - q_dtype : torch.dtype - Query dtype. - k_dtype : torch.dtype - Key cache dtype, and the value cache dtype when ``v_dtype`` is - omitted. It must currently equal ``q_dtype``. - v_dtype : torch.dtype, optional - Value cache dtype; defaults to ``k_dtype``. May differ from - ``k_dtype`` only for the QK-BF16/PV-FP8 combination. - out_dtype : torch.dtype, optional - Output dtype; defaults to ``q_dtype``. - page_size : int - Number of K/V tokens stored in each page. - mask_type : {"dense", "causal"} - Attention mask mode. - window_left : int - Left sliding-window extent, or ``-1`` to disable the window. - sm_scale : float, optional - Softmax scale; defaults to the inverse square root of head size. - output_scale : float - Scale applied to the attention output. - uniform_packed_lengths : bool - Whether every run has Q lengths exactly ``max_seq_len_q`` and K/V - lengths exactly ``max_kv_len``. Defaults to ``False``. - has_q_offset : bool - Whether causal runs may have a nonzero bottom-right Q offset - ``Sk - Sq``. Setting this to ``False`` promises ``Sq == Sk`` for - every request. Dense attention ignores this flag. Defaults to - ``True``. - paged_v_tail_is_zero : bool - Whether unused rows in every active final V-cache page are - guaranteed to contain zero. Setting this to ``True`` compiles out - the consumer-side V-tail clear. Defaults to ``False``. - """ - + """Select/cache one execution policy, obtain compiled kernels, and bind resources. + + Page-table padding exposes ceil(max_kv_len / page_size) columns. + Uniform-length and zero-V-tail promises remain caller-owned. + batch_size is exact. Fixed contiguous inputs use exact Q/KV tensor extents; + packed/paged request lengths may vary within max_q_len/max_kv_len. + Complete planning outside CUDA graph capture. Replanning replaces owned + resources and invalidates previous graphs; finish their launches first. + lse_mode='base2' enables final LSE stores, and use_attention_sinks enables + a per-run sink binding. Context scales are positive host scalar defaults. + Equivalent plans reuse policy and compiled-kernel caches.""" + + lse_mode = validate_lse_mode(lse_mode) + if not isinstance(use_attention_sinks, bool): + raise TypeError("use_attention_sinks must be a bool") resolved_out_dtype = q_dtype if out_dtype is None else out_dtype resolved_v_dtype = k_dtype if v_dtype is None else v_dtype if k_dtype != q_dtype: @@ -2972,7 +3427,7 @@ def plan( geometry = _resolve_paged_plan_geometry( device=device, batch_size=batch_size, - max_seq_len_q=max_seq_len_q, + max_q_len=max_q_len, max_kv_len=max_kv_len, num_qo_heads=num_qo_heads, num_kv_heads=num_kv_heads, @@ -2982,37 +3437,39 @@ def plan( page_size=page_size, mask_type=mask_type, window_left=window_left, - output_dtype=resolved_out_dtype, + out_dtype=resolved_out_dtype, uniform_packed_lengths=uniform_packed_lengths, has_q_offset=has_q_offset, paged_v_tail_is_zero=paged_v_tail_is_zero, ) - if sm_scale is None: - sm_scale = 1.0 / math.sqrt(geometry.head_dim) - sm_scale = _validate_scale(sm_scale, "sm_scale") - output_scale = _validate_scale(output_scale, "output_scale") + if softmax_scale is None: + softmax_scale = geometry.head_dim**-0.5 scale_softmax_log2 = _validate_scale( - sm_scale * math.log2(math.e), "sm_scale * log2(e)" - ) - scale_tensor = torch.tensor( - [scale_softmax_log2], dtype=torch.float32, device=geometry.device + _validate_scale(softmax_scale, "softmax_scale") * math.log2(math.e), + "softmax_scale * log2(e)", ) - output_scale_tensor = torch.tensor( - [output_scale], dtype=torch.float32, device=geometry.device - ) - - # Keep the two plan-owned scalar allocations ahead of CUTLASS JIT, - # matching the compute-sanitizer-safe ordering of the other wrapper. + output_scale = _validate_scale(output_scale, "output_scale") compiled, policy = _get_compiled_paged_context( - _paged_context_compile_spec(geometry) + _paged_context_compile_spec( + geometry, + runtime_lengths=True, + lse_mode=lse_mode, + use_attention_sinks=use_attention_sinks, + ) + ) + policy = tuple( + (key, f"paged_{self._kv_layout.lower()}" if key == "kv_layout" else value) + for key, value in policy ) # Publish one immutable state object only after every fallible step. self._plan_state = _PagedContextPlanState( geometry=geometry, - scale_softmax_log2=scale_tensor, - output_scale=output_scale_tensor, + scale_softmax_log2=scale_softmax_log2, + output_scale=output_scale, compiled=compiled, policy=policy, + lse_mode=lse_mode, + use_attention_sinks=use_attention_sinks, ) @flashinfer_experimental_api @@ -3025,86 +3482,67 @@ def run( block_tables: torch.Tensor, seq_lens_kv: torch.Tensor, *, + max_q_len: int | None = None, + max_kv_len: int | None = None, out: Optional[torch.Tensor] = None, - scale_softmax_log2: Optional[torch.Tensor] = None, - output_scale: Optional[torch.Tensor] = None, + lse: Optional[torch.Tensor] = None, + return_lse: bool = False, + attention_sinks: Optional[torch.Tensor] = None, + softmax_scale: Optional[float] = None, + output_scale: Optional[float] = None, validate: bool = True, - ) -> torch.Tensor: - """Launch the compiled specialization with per-run request metadata. + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Launch packed-Q context with live offsets, page tables and KV lengths. - Metadata must use int32 CUDA storage. ``qo_indptr`` is compact with - shape ``[B + 1]``; ``block_tables`` has shape ``[B, C]``, unit column - stride, and row stride at least ``C``; and ``seq_lens_kv`` is compact - with shape ``[B]``. Their values must satisfy the capacities documented - by ``plan``: + Uses the policy and compiled kernels prepared by plan(). Batch and query + shape must match that plan; live K/V lengths stay within its capacity. + Supply output/LSE buffers and validate=False during graph capture, and + retain the wrapper and captured buffers for every replay. - * ``qo_indptr`` starts at zero, increases strictly, ends at the packed - Q extent, and has deltas no greater than ``max_seq_len_q``; + An enabled lse_mode="base2" plan writes compact FP32 log2(sum(exp(scores))) + with the query shape minus its final dimension. return_lse controls tuple + return only; an omitted enabled buffer is allocated. Output/LSE/workspace must + be disjoint from inputs and live buffers. One wrapper serves one ordered lane. + + * ``qo_indptr`` starts at zero, is nondecreasing, ends at the packed + Q extent, and has deltas no greater than ``max_q_len``; * ``C`` is at least ``ceil(max_kv_len / page_size)``; - * ``seq_lens_kv`` entries are positive and do not exceed + * ``seq_lens_kv`` entries are nonnegative and do not exceed ``max_kv_len``; - * causal launches additionally satisfy ``Sq[b] <= Sk[b]``; and + * active causal requests additionally satisfy ``Sq[b] <= Sk[b]``; and * every active page-table entry is a valid physical page ID in both separate, isomorphic K and V pools. Padding entries beyond the logical K/V length are not dereferenced. + Metadata uses CUDA int32 on Q's device. qo_indptr starts at zero, is + nondecreasing, and ends at the packed Q extent. Page-table columns are contiguous; + rows may have non-overlapping padding. Active page IDs must index both caches. + Active lengths must respect the supplied bounds and selected model promises, with + Sq[b] <= Sk[b] for causal attention. + When ``plan`` selected ``uniform_packed_lengths=True``, every Q delta - must equal ``max_seq_len_q`` and every K/V length must equal + must equal ``max_q_len`` and every K/V length must equal ``max_kv_len``. For a causal plan with ``has_q_offset=False``, every request must additionally satisfy ``Sq[b] == Sk[b]``. - With ``validate=True``, the host reads these values and may synchronize. - Metadata may change only between completed launches or graph replays; - CUDA graph capture additionally requires stable tensor shapes, strides, - and addresses plus ``validate=False``. That mode does not verify the - plan-time metadata promises; the caller is responsible for preserving - them across every replay. - - Parameters - ---------- - q : torch.Tensor - Runtime packed query tensor matching the plan. - k_cache : torch.Tensor - Runtime HND key page pool matching the plan. - v_cache : torch.Tensor - Runtime HND value page pool matching the plan. - qo_indptr : torch.Tensor - Cumulative packed-query offsets with shape ``[B + 1]``. - block_tables : torch.Tensor - Physical page IDs with shape ``[B, C]``. Column stride must be one - and row stride may exceed ``C``. - seq_lens_kv : torch.Tensor - Logical K/V lengths with shape ``[B]``. - out : torch.Tensor, optional - Caller-owned output tensor. A new tensor is allocated when omitted. - scale_softmax_log2 : torch.Tensor, optional - Per-run one-element float32 softmax scale in base-2 form. Defaults - to the scale retained by the plan. - output_scale : torch.Tensor, optional - Per-run one-element float32 output scale. Defaults to the scale - retained by the plan. - validate : bool - Run explicit storage, shape, dtype, device, scale, output, and - metadata-value validators. Disable only when the caller guarantees - these contracts. Output must not overlap any live input or - plan-owned buffer in either mode; storage overlap is not checked. - Defaults to ``True``. - - Returns - ------- - torch.Tensor - The packed attention output. - """ + Scale overrides are finite positive host scalars in the same units as the + contiguous wrapper. Enabled attention_sinks is a compact CUDA FP32 [Hq] + vector of finite or -inf natural-log sink logits after QK scaling, with zero + V contribution. Sink values and cache zero-tail promises are caller-owned.""" + + if not isinstance(validate, bool): + raise TypeError("validate must be a bool") state = self._plan_state if state is None: raise RuntimeError("plan() must be called before run()") - if not isinstance(validate, bool): - raise TypeError("validate must be a bool") geometry = state.geometry + if self._kv_layout == "NHD": + k_cache = k_cache.transpose(1, 2) + v_cache = v_cache.transpose(1, 2) if validate: _validate_paged_runtime_inputs(q, k_cache, v_cache, geometry) - _validate_paged_runtime_metadata( + actual_q, actual_kv = _validate_paged_runtime_metadata( qo_indptr, block_tables, seq_lens_kv, @@ -3112,39 +3550,69 @@ def run( total_q=int(q.shape[0]), num_physical_pages=int(k_cache.shape[0]), ) - if scale_softmax_log2 is None: - scale_softmax_log2 = state.scale_softmax_log2 - elif validate: - _validate_runtime_scale_tensor( - scale_softmax_log2, - "scale_softmax_log2", - device=geometry.device, + from .decode import resolve_max_sequence_length + + max_q_len = resolve_max_sequence_length( + max_q_len, + actual_q, + "max_q_len", + geometry.max_q_len, ) - if output_scale is None: - output_scale = state.output_scale - elif validate: - _validate_runtime_scale_tensor( - output_scale, "output_scale", device=geometry.device + max_kv_len = resolve_max_sequence_length( + max_kv_len, actual_kv, "max_kv_len", geometry.max_kv_len + ) + else: + if max_q_len is None: + max_q_len = geometry.max_q_len + if max_kv_len is None: + max_kv_len = geometry.max_kv_len + scale_softmax_log2 = ( + state.scale_softmax_log2 + if softmax_scale is None + else _validate_scale( + _validate_scale(softmax_scale, "softmax_scale") * math.log2(math.e), + "softmax_scale * log2(e)", ) + ) + output_scale = ( + state.output_scale + if output_scale is None + else _validate_scale(output_scale, "output_scale") + ) if out is None: - out = torch.empty( - tuple(q.shape), dtype=geometry.output_dtype, device=q.device - ) + out = torch.empty(tuple(q.shape), dtype=geometry.out_dtype, device=q.device) elif validate: - out = _prepare_out(out, q=q, output_dtype=geometry.output_dtype) + out = _prepare_out(out, q=q, out_dtype=geometry.out_dtype) + lse = prepare_lse( + lse, + q=q, + mode=state.lse_mode, + return_lse=return_lse, + validate=validate, + ) + attention_sinks = validate_attention_sinks( + attention_sinks, + q=q, + enabled=state.use_attention_sinks, + validate=validate, + ) state.compiled( q, k_cache, v_cache, out, + lse, + attention_sinks, scale_softmax_log2, output_scale, qo_indptr, block_tables, seq_lens_kv, + max_q_len, + max_kv_len, ) - return out + return (out, lse) if return_lse else out @flashinfer_experimental_api @@ -3159,16 +3627,23 @@ def batch_prefill( window_left: int = -1, variable_window_token_starts: Optional[torch.Tensor] = None, variable_window_token_ends: Optional[torch.Tensor] = None, - sm_scale: Optional[float] = None, + softmax_scale: Optional[float] = None, output_scale: float = 1.0, out_dtype: Optional[torch.dtype] = None, out: Optional[torch.Tensor] = None, -) -> torch.Tensor: + vc: Optional[VCAttentionParams] = None, + vc_config: Optional[VCAttentionConfig] = None, + lse: Optional[torch.Tensor] = None, + return_lse: bool = False, + attention_sinks: Optional[torch.Tensor] = None, + lse_mode: Literal["base2"] = "base2", +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """Run one-shot fixed or packed-ragged task-scheduled context attention. Fixed tensors use ``[B, S, H, D]`` storage. Providing both cumulative int32 offset tensors selects packed ``[total_tokens, H, D]`` storage. - Equal Q/K/V head dimensions may be 128 or 256. Separate Q/K=192 and + Equal Q/K/V head dimensions may be 128, 192, or 256. The dtype contract + is shared by all three dimensions. Separate Q/K=192 and V=128 supports non-absorbed MLA; O uses V's head dimension. Causal masking is bottom-right aligned. ``window_left=-1`` disables the left window; a positive value selects the private head-paired GQA policy @@ -3192,7 +3667,9 @@ def batch_prefill( Inclusive per-query K bounds required for ``variable_window``. Both must be CUDA int32 tensors shaped ``[B, Sq]`` and satisfy ``0 <= starts[b, q] <= ends[b, q] < Sk``. - sm_scale : float, optional + Ends must be nondecreasing along Q within each batch item; starts + need not be ordered. + softmax_scale : float, optional Softmax scale; defaults to the inverse square root of head size. output_scale : float Scale applied to the attention output. @@ -3206,6 +3683,25 @@ def batch_prefill( ------- torch.Tensor The fixed or packed attention output. + + ``return_lse=True`` returns ``(output, lse)``. ``lse_mode`` selects + ``"base2"`` (the only supported base); it is used only when returning LSE + or writing a caller-provided ``lse`` buffer. + + ``lse`` is an optional independent, compact float32 output with shape + ``query.shape[:-1]`` (``q.shape[:-1]`` for FMHA). An enabled LSE plan + allocates it when omitted, including when ``return_lse=False``. A supplied + buffer is written regardless of ``return_lse``; the flag only controls + the Python return type. LSE includes any sink mass and excludes output + scaling. Keep LSE disjoint from output, inputs, and workspace. Supply + both output buffers and disable validation for allocation-free capture. + + ``attention_sinks`` is a compact CUDA float32 vector ``[Hq]`` shared + across batch/query positions. Values are natural-log sink logits after + QK scaling, with zero V contribution. Finite values and ``-inf`` are + accepted. Validation rejects NaN/positive infinity and may synchronize. + With validation disabled, the caller must preserve these contracts. + Sink values may change between completed CUDA graph replays. """ resolved_out_dtype = ( @@ -3220,6 +3716,21 @@ def batch_prefill( "batch_prefill cannot derive host plan bounds during CUDA graph " "capture; plan BatchPrefillTSWrapper before capture" ) + if vc_config is not None and vc is None: + raise ValueError("vc_config requires the vc operands of this call") + if (vc_config is not None and vc_config.repair_tiles) or ( + vc_config is None and vc is not None and vc.tile_means is None + ): + # The K/V rows of V repair operands include the repair tiles, so the + # plan bounds cannot be derived from the operands alone. Without a + # config, operands with no tile means may be repair operands. + raise NotImplementedError( + "batch_prefill does not run V repair operands; plan " + "BatchPrefillTSWrapper with VCAttentionConfig(repair_tiles=...), or " + "pass vc_config=VCAttentionConfig(demean=False) for plain operands" + ) + if vc is not None and vc_config is None: + vc_config = VCAttentionConfig(demean=vc.tile_means is not None) geometry = _resolve_geometry( q, k, @@ -3228,13 +3739,13 @@ def batch_prefill( kv_indptr=kv_indptr, mask_type=mask_type, window_left=window_left, - output_dtype=resolved_out_dtype, + out_dtype=resolved_out_dtype, ) wrapper = BatchPrefillTSWrapper() wrapper.plan( device=geometry.device, batch_size=geometry.batch_size, - max_seq_len_q=geometry.max_seq_len_q, + max_q_len=geometry.max_q_len, max_kv_len=geometry.max_seq_len_k, num_qo_heads=geometry.num_qo_heads, num_kv_heads=geometry.num_kv_heads, @@ -3243,12 +3754,17 @@ def batch_prefill( q_dtype=geometry.qk_dtype, k_dtype=geometry.qk_dtype, v_dtype=geometry.pv_dtype, - out_dtype=geometry.output_dtype, - packed=geometry.packed, + out_dtype=geometry.out_dtype, + packed_query=geometry.packed_query, mask_type=mask_type, window_left=window_left, - sm_scale=sm_scale, + softmax_scale=softmax_scale, output_scale=output_scale, + vc_config=vc_config, + use_attention_sinks=attention_sinks is not None, + lse_mode=validate_lse_mode(lse_mode) + if return_lse or lse is not None + else "none", ) return wrapper.run( q, @@ -3258,7 +3774,13 @@ def batch_prefill( kv_indptr, variable_window_token_starts=variable_window_token_starts, variable_window_token_ends=variable_window_token_ends, + max_q_len=geometry.max_q_len, + max_kv_len=geometry.max_seq_len_k, out=out, + vc=vc, + lse=lse, + return_lse=return_lse, + attention_sinks=attention_sinks, ) @@ -3272,18 +3794,23 @@ def batch_prefill_with_paged_kv_cache( seq_lens_kv: torch.Tensor, *, page_size: int = _DEFAULT_PAGED_KV_PAGE_SIZE, - kv_layout: Literal["HND"] = "HND", + kv_layout: Literal["HND", "NHD"] = "HND", mask_type: Literal["dense", "causal"] = "dense", window_left: int = -1, - sm_scale: Optional[float] = None, + softmax_scale: Optional[float] = None, output_scale: float = 1.0, out_dtype: Optional[torch.dtype] = None, out: Optional[torch.Tensor] = None, -) -> torch.Tensor: - """Run one-shot packed-Q context attention over separate HND page pools. + lse: Optional[torch.Tensor] = None, + return_lse: bool = False, + attention_sinks: Optional[torch.Tensor] = None, + lse_mode: Literal["base2"] = "base2", +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Run one-shot packed-Q context attention over HND or NHD page pools. Q/output use ``[total_q, Hq, D]`` storage and K/V each use - ``[num_pages, Hkv, page_size, D]`` with page size 16, 32, 64, or 128. + ``[num_pages, Hkv, page_size, D]`` for HND or + ``[num_pages, page_size, Hkv, D]`` for NHD, with page size 4/8/16/32/64/128. ``qo_indptr`` describes Q rows, ``block_tables`` supplies one fixed row-strided page table, and ``seq_lens_kv`` supplies logical K/V lengths. Physical page indices need not be identity ordered. ``D`` may be 128 or @@ -3311,13 +3838,13 @@ def batch_prefill_with_paged_kv_cache( Logical K/V lengths with shape ``[B]``. page_size : int Number of K/V tokens stored in each page. - kv_layout : {"HND"} + kv_layout : {"HND", "NHD"} Layout of the separate K and V page pools. mask_type : {"dense", "causal"} Attention mask mode. window_left : int Left sliding-window extent, or ``-1`` to disable the window. - sm_scale : float, optional + softmax_scale : float, optional Softmax scale; defaults to the inverse square root of head size. output_scale : float Scale applied to the attention output. @@ -3331,6 +3858,25 @@ def batch_prefill_with_paged_kv_cache( ------- torch.Tensor The packed attention output. + + ``return_lse=True`` returns ``(output, lse)``. ``lse_mode`` selects + ``"base2"`` (the only supported base); it is used only when returning LSE + or writing a caller-provided ``lse`` buffer. + + ``lse`` is an optional independent, compact float32 output with shape + ``query.shape[:-1]`` (``q.shape[:-1]`` for FMHA). An enabled LSE plan + allocates it when omitted, including when ``return_lse=False``. A supplied + buffer is written regardless of ``return_lse``; the flag only controls + the Python return type. LSE includes any sink mass and excludes output + scaling. Keep LSE disjoint from output, inputs, and workspace. Supply + both output buffers and disable validation for allocation-free capture. + + ``attention_sinks`` is a compact CUDA float32 vector ``[Hq]`` shared + across batch/query positions. Values are natural-log sink logits after + QK scaling, with zero V contribution. Finite values and ``-inf`` are + accepted. Validation rejects NaN/positive infinity and may synchronize. + With validation disabled, the caller must preserve these contracts. + Sink values may change between completed CUDA graph replays. """ resolved_out_dtype = ( @@ -3346,23 +3892,26 @@ def batch_prefill_with_paged_kv_cache( "during CUDA graph capture; plan BatchPrefillPagedTSWrapper before " "capture" ) + _validate_kv_layout(kv_layout) + logical_k = k_cache.transpose(1, 2) if kv_layout == "NHD" else k_cache + logical_v = v_cache.transpose(1, 2) if kv_layout == "NHD" else v_cache geometry = _resolve_paged_geometry( q, - k_cache, - v_cache, + logical_k, + logical_v, qo_indptr=qo_indptr, block_tables=block_tables, seq_lens_kv=seq_lens_kv, page_size=page_size, mask_type=mask_type, window_left=window_left, - output_dtype=resolved_out_dtype, + out_dtype=resolved_out_dtype, ) wrapper = BatchPrefillPagedTSWrapper(kv_layout=kv_layout) wrapper.plan( device=geometry.device, batch_size=geometry.batch_size, - max_seq_len_q=geometry.max_seq_len_q, + max_q_len=geometry.max_q_len, max_kv_len=geometry.max_kv_len, num_qo_heads=geometry.num_qo_heads, num_kv_heads=geometry.num_kv_heads, @@ -3370,13 +3919,17 @@ def batch_prefill_with_paged_kv_cache( q_dtype=geometry.qk_dtype, k_dtype=geometry.qk_dtype, v_dtype=geometry.pv_dtype, - out_dtype=geometry.output_dtype, + out_dtype=geometry.out_dtype, page_size=page_size, mask_type=mask_type, window_left=window_left, - sm_scale=sm_scale, + softmax_scale=softmax_scale, output_scale=output_scale, - uniform_packed_lengths=geometry.uniform_packed_lengths, + use_attention_sinks=attention_sinks is not None, + lse_mode=validate_lse_mode(lse_mode) + if return_lse or lse is not None + else "none", + uniform_packed_lengths=False, has_q_offset=geometry.has_q_offset, paged_v_tail_is_zero=geometry.paged_v_tail_is_zero, ) @@ -3387,7 +3940,12 @@ def batch_prefill_with_paged_kv_cache( qo_indptr, block_tables, seq_lens_kv, + max_q_len=geometry.max_q_len, + max_kv_len=geometry.max_kv_len, out=out, + lse=lse, + return_lse=return_lse, + attention_sinks=attention_sinks, ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py b/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py index f91814650c48..0609b55f408f 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py @@ -16,18 +16,24 @@ """Task-scheduled paged decode with a FlashInfer-style plan/run lifecycle.""" from collections.abc import Callable, Sequence -from dataclasses import dataclass, replace +from dataclasses import dataclass, field, replace import functools import math import numbers import struct -from typing import TYPE_CHECKING, Literal, NamedTuple, Optional, Union, cast +from typing import TYPE_CHECKING, Literal, Optional, Union, cast import torch from ._compile_cache import build_and_load_attention_kernel from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api +from flashinfer.utils import ( + get_device_sm_count, + _check_workspace_buffer_alignment, + check_shape_dtype_device, +) + from . import _q_token_kv_block_sparse_policy as _sparse_policy from ._block_sparse.common import _num_sparse_pattern_heads @@ -38,7 +44,7 @@ if TYPE_CHECKING: from .kernels.fmha_decode.fmha_decode_config import FmhaDecodeConfig -_SUPPORTED_HEAD_DIMS = (64, 128, 256) +_SUPPORTED_HEAD_DIMS = (64, 128, 192, 256) _SUPPORTED_PAGE_SIZES = (4, 8, 16, 32, 64, 128) _MAX_INT32 = 2**31 - 1 # Decode K/V masks form an exclusive tile endpoint as @@ -100,16 +106,25 @@ class _DecodeLaunchSpec: max_active_clusters: int policy: tuple[tuple[str, object], ...] scratch_shapes: tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]] + requested_splits: Optional[int] = None -class _DirectQ1CompileSpec(NamedTuple): - """Static raw-index geometry; tensor values and strides remain runtime inputs.""" +@dataclass(frozen=True) +class _DirectQ1CompileSpec: + """Sparse format identity; index and storage extents are runtime operands.""" position_dtype_key: str - max_seq_len_kv: int - block_topk: int + max_seq_len_kv: int = field(compare=False, hash=False) + block_topk: int = field(compare=False, hash=False) kv_block_size: int - page_capacity: int + page_capacity: int = field(compare=False, hash=False) + count_current_block_in_topk: bool = True + fixed_query_len: int = field(default=0, compare=False, hash=False) + fixed_query_layout: bool = field(init=False) + + # Key fixed versus packed metadata layout, without keying the Q token count. + def __post_init__(self): + object.__setattr__(self, "fixed_query_layout", self.fixed_query_len > 0) @dataclass(frozen=True) @@ -122,17 +137,16 @@ class _DecodeCompileSpec: num_kv_heads: int head_dim: int page_size: int - max_kv_len: int - seq_len_q: int q_dtype_key: str k_dtype_key: str v_dtype_key: str - output_dtype_key: str - use_packed_q: bool + out_dtype_key: str + packed_query: bool max_active_clusters: int kv_prefix_mode: Literal["dynamic", "planned_full"] kv_lengths_mode: Literal["dynamic", "planned_uniform_max"] direct_q1_spec: Optional[_DirectQ1CompileSpec] = None + runtime_kv_lengths: bool = False @dataclass(frozen=True) @@ -152,8 +166,8 @@ class _DecodeRuntime: v_page_stride: int v_head_stride: int v_token_stride: int - bmm1_scale: float - bmm2_scale: float + softmax_scale: float + output_scale: float @dataclass(frozen=True) @@ -168,7 +182,7 @@ class _DecodePlanState: device_index: int batch_size: int seq_len_q: int - use_packed_q: bool + packed_query: bool num_qo_heads: int num_kv_heads: int head_dim: int @@ -177,7 +191,7 @@ class _DecodePlanState: q_dtype: torch.dtype k_dtype: torch.dtype v_dtype: torch.dtype - output_dtype: torch.dtype + out_dtype: torch.dtype mask_type: str window_left: int config: "FmhaDecodeConfig" @@ -188,10 +202,14 @@ class _DecodePlanState: compiled_reducer: Optional[Callable[..., object]] kv_prefix_mode: Literal["dynamic", "planned_full"] kv_lengths_mode: Literal["dynamic", "planned_uniform_max"] - planned_seq_lens_host: Optional[tuple[int, ...]] - planned_seq_lens_device: Optional[torch.Tensor] policy: tuple[tuple[str, object], ...] storage_page_size: Optional[int] = None + lse_mode: str = "none" + use_attention_sinks: bool = False + requested_splits: int = 1 + planned_seq_lens_kv_host: Optional[tuple[int, ...]] = None + planned_seq_lens_kv_device: Optional[torch.Tensor] = None + fixed_q_offsets: Optional[torch.Tensor] = None def _decode_policy_from_config( @@ -208,8 +226,8 @@ def _decode_policy_from_config( ) return ( ("seq_len_q", seq_len_q), - ("max_seq_len_q", seq_len_q), - ("use_packed_q", uses_packed_q), + ("max_q_len", seq_len_q), + ("packed_query", uses_packed_q), ("query_layout", query_layout), ("output_layout", query_layout), ( @@ -238,95 +256,6 @@ def _decode_policy_from_config( ) -def _planned_full_split_prefix( - config: "FmhaDecodeConfig", - seq_lens: tuple[int, ...], - *, - seq_len_q: int, - max_kv_len: int, - mask_type: str, -) -> bool: - """Prove that host length evidence uses every configured split CTA. - - A successful proof permits a private JIT specialization that removes only - the no-op runtime split-prefix branch. The plan retains the evidenced - lengths and supplies an owned device copy to every launch. - - Q groups are enumerated with the same token-base/union rule as the device - helper. Every batch/group pair must prove the configured fanout; otherwise - the general runtime-pruning kernel is retained. - """ - - if ( - not bool(config.use_split_kv) - or int(config.splits_kv) <= 1 - or bool(config.use_variable_seqlens_q) - or bool(config.use_sliding_window_causal) - ): - return False - from .kernels.fmha_decode.fmha_decode_config import ( - compute_runtime_active_splits_kv, - ) - - configured_splits = int(config.splits_kv) - if bool(config.uses_nontrivial_grouped_q_layout): - q_group_token_bases = range(0, seq_len_q, int(config.q_tokens_per_cta)) - q_tokens_per_group = int(config.q_tokens_per_cta) - else: - q_group_token_bases = range(seq_len_q) - q_tokens_per_group = 1 - for seq_len_kv in seq_lens: - if seq_len_kv <= 0 or seq_len_kv > max_kv_len: - return False - for q_token_base in q_group_token_bases: - valid_k = seq_len_kv - if mask_type == "causal": - q_token_end = min(q_token_base + q_tokens_per_group, seq_len_q) - valid_k = max(seq_len_kv - seq_len_q + q_token_end, 0) - if ( - compute_runtime_active_splits_kv( - valid_k=valid_k, - tile_size_kv=int(config.tile_size_kv), - num_insts_kv=int(config.num_insts_kv), - configured_splits_kv=configured_splits, - ) - != configured_splits - ): - return False - return True - - -def _planned_kv_lengths_mode( - seq_lens: tuple[int, ...], - *, - max_kv_len: int, -) -> Literal["dynamic", "planned_uniform_max"]: - """Classify host length evidence for fixed-length kernel scheduling. - - When every evidenced request is exactly the compiled maximum, native page - addressing still uses the runtime block table while task domains and - masks can use the compile-time length. The plan retains the evidenced - lengths and supplies an owned device copy to every launch. - """ - - if not seq_lens or max_kv_len <= 0: - return "dynamic" - if all(seq_len == max_kv_len for seq_len in seq_lens): - return "planned_uniform_max" - return "dynamic" - - -def _planned_kv_domain_has_unpaired_tail( - config: "FmhaDecodeConfig", max_kv_len: int -) -> bool: - """Return whether the planned K domain ends with one inactive KV instance.""" - - tile_size_kv = int(config.tile_size_kv) - num_insts_kv = int(config.num_insts_kv) - total_kv_tiles = (max_kv_len + tile_size_kv - 1) // tile_size_kv - return total_kv_tiles % num_insts_kv != 0 - - def _align_up(value: int, alignment: int = _WORKSPACE_ALIGNMENT) -> int: return (value + alignment - 1) // alignment * alignment @@ -361,19 +290,19 @@ def _append_workspace_section( def _make_decode_workspace_layout( scratch_shapes: tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]], - output_dtype: torch.dtype, + out_dtype: torch.dtype, *, use_separate_reduction_kernel: bool, use_split_kv: bool = False, ) -> _DecodeWorkspaceLayout: partial_o_shape, partial_stats_shape, counter_shape = scratch_shapes - partial_o_dtype = output_dtype - if use_separate_reduction_kernel and output_dtype in ( + partial_o_dtype = out_dtype + if use_separate_reduction_kernel and out_dtype in ( torch.bfloat16, torch.float8_e4m3fn, ): partial_o_dtype = torch.bfloat16 - elif output_dtype == torch.float8_e4m3fn or partial_o_shape == (1, 1, 1, 1, 1): + elif out_dtype == torch.float8_e4m3fn or partial_o_shape == (1, 1, 1, 1, 1): partial_o_dtype = torch.float16 byte_end = 0 @@ -489,39 +418,39 @@ def _resolve_q_mode( *, seq_len_q: int, qo_indptr: Optional[torch.Tensor], - max_seq_len_q: Optional[int], + max_q_len: Optional[int], require_packed_max: bool, ) -> tuple[bool, Optional[int]]: """Resolve fixed versus packed Q without exposing an internal mode knob.""" seq_len_q = _validate_seq_len_q(seq_len_q) if qo_indptr is None: - if max_seq_len_q is not None: - max_seq_len_q = _validate_seq_len_q(max_seq_len_q) - if max_seq_len_q != seq_len_q: + if max_q_len is not None: + max_q_len = _validate_seq_len_q(max_q_len) + if max_q_len != seq_len_q: raise ValueError( - "fixed seq_len_q and max_seq_len_q must agree: " - f"got {seq_len_q} and {max_seq_len_q}" + "fixed seq_len_q and max_q_len must agree: " + f"got {seq_len_q} and {max_q_len}" ) return False, seq_len_q - if max_seq_len_q is None and seq_len_q != 1: + if max_q_len is None and seq_len_q != 1: # Preserve the legacy name as a packed static-bound alias. The # nullable qo_indptr alone still selects fixed versus packed storage. return True, seq_len_q - if max_seq_len_q is None: + if max_q_len is None: if require_packed_max: raise ValueError( - "max_seq_len_q is required with qo_indptr for the standalone " + "max_q_len is required with qo_indptr for the standalone " "workspace/JIT interface" ) return True, None - max_seq_len_q = _validate_seq_len_q(max_seq_len_q) - if seq_len_q != 1 and seq_len_q != max_seq_len_q: + max_q_len = _validate_seq_len_q(max_q_len) + if seq_len_q != 1 and seq_len_q != max_q_len: raise ValueError( - "seq_len_q and max_seq_len_q must agree when both provide the " - f"packed static bound: got {seq_len_q} and {max_seq_len_q}" + "seq_len_q and max_q_len must agree when both provide the " + f"packed static bound: got {seq_len_q} and {max_q_len}" ) - return True, max_seq_len_q + return True, max_q_len def _validate_max_kv_len(value: int, name: str) -> int: @@ -536,6 +465,24 @@ def _validate_max_kv_len(value: int, name: str) -> int: return value +def resolve_max_sequence_length( + value: int | None, actual: int, name: str, capacity: int +) -> int: + """Infer an omitted launch bound, or check coverage within plan capacity.""" + + if value is not None: + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be an integer") + if value < actual: + raise ValueError(f"{name} must cover the maximum input length ({actual})") + if value > capacity: + raise ValueError( + f"{name} must not exceed the planned capacity ({capacity})" + ) + return value + return actual + + def _validate_decode_policy_kv_tile_size(config: "FmhaDecodeConfig") -> None: """Keep the public K/V bound coupled to generated decode policies.""" @@ -551,17 +498,17 @@ def _validate_decode_query_head_extent( *, batch_size: int, num_qo_heads: int, - max_seq_len_q: int, + max_q_len: int, ) -> None: """Keep every fixed-capacity or packed Q/head coordinate in Int32.""" batch_size = _validate_positive_int(batch_size, "batch_size") num_qo_heads = _validate_positive_int(num_qo_heads, "num_qo_heads") - max_seq_len_q = _validate_seq_len_q(max_seq_len_q) - extent = batch_size * max_seq_len_q * num_qo_heads + max_q_len = _validate_seq_len_q(max_q_len) + extent = batch_size * max_q_len * num_qo_heads if extent > _MAX_INT32: raise NotImplementedError( - "batch_size * max_seq_len_q * num_qo_heads must fit in a signed int32" + "batch_size * max_q_len * num_qo_heads must fit in a signed int32" ) @@ -592,15 +539,78 @@ def _validate_16byte_alignment(tensor: torch.Tensor, name: str) -> None: raise ValueError(f"{name} data pointer must be 16-byte aligned") +def _validate_scale(value: object, name: str) -> float: + if isinstance(value, bool) or not isinstance(value, numbers.Real): + raise TypeError(f"{name} must be a positive Python scalar") + try: + value_as_float = float(value) + except (OverflowError, TypeError, ValueError) as error: + raise TypeError(f"{name} must be a positive Python scalar") from error + if not math.isfinite(value_as_float) or value_as_float <= 0.0: + raise ValueError(f"{name} must be finite and positive") + try: + value_as_float32 = struct.unpack("=f", struct.pack("=f", value_as_float))[0] + except (OverflowError, struct.error) as error: + raise ValueError( + f"{name} must be representable as a positive float32" + ) from error + if not math.isfinite(value_as_float32) or value_as_float32 <= 0.0: + raise ValueError(f"{name} must be representable as a positive float32") + return value_as_float32 + + def _validate_layout(kv_layout: str) -> None: if not isinstance(kv_layout, str): raise TypeError("kv_layout must be a string") - if kv_layout == "NHD": - raise NotImplementedError( - "attention-ts decode currently supports kv_layout='HND' only" + if kv_layout not in ("HND", "NHD"): + raise ValueError(f"kv_layout must be 'HND' or 'NHD', got {kv_layout!r}") + + +def validate_lse_mode(mode: str) -> Literal["none", "base2"]: + if mode not in ("none", "base2"): + raise ValueError("lse_mode must be 'none' or 'base2'") + return cast(Literal["none", "base2"], mode) + + +def prepare_lse( + lse: torch.Tensor | None, + *, + q: torch.Tensor, + mode: str, + return_lse: bool, + validate: bool, +) -> torch.Tensor | None: + """Bind an independent caller buffer, or allocate the enabled output.""" + if mode == "none": + if return_lse or lse is not None: + raise ValueError("plan() must enable lse_mode='base2'") + return None + if lse is None: + return torch.empty(q.shape[:-1], dtype=torch.float32, device=q.device) + if validate: + check_shape_dtype_device(lse, q.shape[:-1], torch.float32, q.device, "lse") + if not lse.is_contiguous(): + raise ValueError("lse must be contiguous") + _check_workspace_buffer_alignment(lse, "lse", alignment=4) + return lse + + +def validate_attention_sinks(attention_sinks, *, q, enabled, validate): + """Bind sink logits; callers guarantee finite values or negative infinity.""" + if enabled != (attention_sinks is not None): + raise ValueError( + "attention_sinks must be supplied exactly when plan(use_attention_sinks=True)" ) - if kv_layout != "HND": - raise ValueError(f"kv_layout must be exactly 'HND', got {kv_layout!r}") + if attention_sinks is not None and validate: + check_shape_dtype_device( + attention_sinks, (q.shape[-2],), torch.float32, q.device, "attention_sinks" + ) + if not attention_sinks.is_contiguous(): + raise ValueError("attention_sinks must be contiguous") + _check_workspace_buffer_alignment( + attention_sinks, "attention_sinks", alignment=4 + ) + return attention_sinks def _validate_mask(mask_type: str) -> None: @@ -638,26 +648,6 @@ def _validate_head_geometry(num_qo_heads: int, num_kv_heads: int) -> None: ) -def _validate_scale(value: object, name: str) -> float: - if isinstance(value, bool) or not isinstance(value, numbers.Real): - raise TypeError(f"{name} must be a positive Python scalar") - try: - value_as_float = float(value) - except (OverflowError, TypeError, ValueError) as error: - raise TypeError(f"{name} must be a positive Python scalar") from error - if not math.isfinite(value_as_float) or value_as_float <= 0.0: - raise ValueError(f"{name} must be finite and positive") - try: - value_as_float32 = struct.unpack("=f", struct.pack("=f", value_as_float))[0] - except (OverflowError, struct.error) as error: - raise ValueError( - f"{name} must be representable as a positive float32" - ) from error - if not math.isfinite(value_as_float32) or value_as_float32 <= 0.0: - raise ValueError(f"{name} must be representable as a positive float32") - return value_as_float32 - - def _dtype_key(dtype: torch.dtype) -> str: if not isinstance(dtype, torch.dtype): raise TypeError("attention-ts dtypes must be torch.dtype values") @@ -695,22 +685,22 @@ def _validate_dtype_pair( q_dtype: torch.dtype, k_dtype: torch.dtype, v_dtype: torch.dtype, - output_dtype: torch.dtype, + out_dtype: torch.dtype, ) -> None: _dtype_key(q_dtype) _dtype_key(k_dtype) _dtype_key(v_dtype) - _dtype_key(output_dtype) + _dtype_key(out_dtype) matching_kv = k_dtype == v_dtype supported = ( ( q_dtype == torch.float16 and k_dtype == v_dtype == torch.float16 - and output_dtype == torch.float16 + and out_dtype == torch.float16 ) or ( q_dtype == torch.bfloat16 - and output_dtype == torch.bfloat16 + and out_dtype == torch.bfloat16 and ( ( matching_kv @@ -723,7 +713,7 @@ def _validate_dtype_pair( q_dtype == torch.float8_e4m3fn and matching_kv and k_dtype in (torch.float8_e4m3fn, torch.uint8) - and output_dtype in (torch.float16, torch.bfloat16, torch.float8_e4m3fn) + and out_dtype in (torch.float16, torch.bfloat16, torch.float8_e4m3fn) ) ) if not supported: @@ -732,7 +722,7 @@ def _validate_dtype_pair( "FP8-E4M3->FP16/BF16/FP8-E4M3; " "BF16 Q + FP8 K/V, BF16/FP8-E4M3 Q + NVFP4 K/V, " "and BF16 Q/K + FP8 V; got " - f"Q_{q_dtype}_K_{k_dtype}_V_{v_dtype}->{output_dtype}" + f"Q_{q_dtype}_K_{k_dtype}_V_{v_dtype}->{out_dtype}" ) @@ -809,7 +799,7 @@ def _validate_q( q: torch.Tensor, *, seq_len_q: int = 1, - use_packed_q: bool = False, + packed_query: bool = False, device: Optional[torch.device] = None, batch_size: Optional[int] = None, num_qo_heads: Optional[int] = None, @@ -819,21 +809,21 @@ def _validate_q( seq_len_q = _validate_seq_len_q(seq_len_q) if not isinstance(q, torch.Tensor): raise TypeError("q must be a torch.Tensor") - expected_rank = 3 if use_packed_q or seq_len_q == 1 else 4 + expected_rank = 3 if packed_query or seq_len_q == 1 else 4 if q.ndim != expected_rank: expected_layout = ( "[total_q, Hq, D]" - if use_packed_q + if packed_query else ("[B, Hq, D]" if seq_len_q == 1 else "[B, SQ, Hq, D]") ) raise ValueError(f"q must have shape {expected_layout} for this Q layout") - if not use_packed_q and seq_len_q > 1 and q.shape[1] != seq_len_q: + if not packed_query and seq_len_q > 1 and q.shape[1] != seq_len_q: raise ValueError( f"q sequence length must match seq_len_q ({seq_len_q}), got {q.shape[1]}" ) num_heads = int(q.shape[-2]) if q.shape[0] <= 0 or num_heads <= 0: - leading_name = "total Q token count" if use_packed_q else "batch size" + leading_name = "total Q token count" if packed_query else "batch size" raise ValueError(f"q {leading_name} and head count must be positive") _validate_head_dim(int(q.shape[-1])) if q.dtype not in _SUPPORTED_INPUT_DTYPES: @@ -842,11 +832,11 @@ def _validate_q( raise ValueError("q must be a CUDA tensor") if device is not None and q.device != device: raise ValueError(f"q must be on planned device {device}, got {q.device}") - if not use_packed_q and batch_size is not None and q.shape[0] != batch_size: + if not packed_query and batch_size is not None and q.shape[0] != batch_size: raise ValueError( f"q batch size must match the plan ({batch_size}), got {q.shape[0]}" ) - if use_packed_q and batch_size is not None: + if packed_query and batch_size is not None: total_q = int(q.shape[0]) max_total_q = batch_size * seq_len_q if total_q < batch_size or total_q > max_total_q: @@ -866,7 +856,7 @@ def _validate_q( raise ValueError(f"q dtype must match the plan ({q_dtype}), got {q.dtype}") layout = ( "[total_q, Hq, D]" - if use_packed_q + if packed_query else ("[B, Hq, D]" if seq_len_q == 1 else "[B, SQ, Hq, D]") ) _validate_exact_compact_strides(q, "q", layout) @@ -924,16 +914,16 @@ def _read_packed_q_plan_metadata( def _validate_packed_q_plan_values( qo_indptr: torch.Tensor, *, - max_seq_len_q: int, + max_q_len: int, expected_total_q: Optional[int] = None, ) -> tuple[int, int]: """Synchronize once to validate the packed-Q values against a static bound.""" derived_max_seq_len_q, total_q, _ = _read_packed_q_plan_metadata(qo_indptr) - if derived_max_seq_len_q > max_seq_len_q: + if derived_max_seq_len_q > max_q_len: raise ValueError( "qo_indptr contains a per-request Q length larger than " - f"max_seq_len_q ({max_seq_len_q}): got {derived_max_seq_len_q}" + f"max_q_len ({max_q_len}): got {derived_max_seq_len_q}" ) if expected_total_q is not None and total_q != expected_total_q: raise ValueError( @@ -1076,7 +1066,7 @@ def _validate_block_tables( def _validate_block_table_metadata( block_tables: torch.Tensor, - seq_lens: torch.Tensor, + seq_lens_kv: torch.Tensor, ) -> tuple[torch.device, int, int]: """Validate fixed page-table and sequence-length tensor structure.""" @@ -1084,30 +1074,30 @@ def _validate_block_table_metadata( block_tables ) - if not isinstance(seq_lens, torch.Tensor): - raise TypeError("seq_lens must be a torch.Tensor") - if seq_lens.ndim != 1: - raise ValueError("seq_lens must be one-dimensional") - if seq_lens.dtype != torch.int32: - raise TypeError("seq_lens must have dtype torch.int32") - if seq_lens.device.type != "cuda": - raise ValueError("seq_lens must be a CUDA tensor") - if not seq_lens.is_contiguous(): - raise ValueError("seq_lens must be contiguous") - if seq_lens.data_ptr() % 4 != 0: - raise ValueError("seq_lens data pointer must be 4-byte aligned") - if block_table_device != seq_lens.device: - raise ValueError("block_tables and seq_lens must be on the same device") - - batch_size = int(seq_lens.numel()) + if not isinstance(seq_lens_kv, torch.Tensor): + raise TypeError("seq_lens_kv must be a torch.Tensor") + if seq_lens_kv.ndim != 1: + raise ValueError("seq_lens_kv must be one-dimensional") + if seq_lens_kv.dtype != torch.int32: + raise TypeError("seq_lens_kv must have dtype torch.int32") + if seq_lens_kv.device.type != "cuda": + raise ValueError("seq_lens_kv must be a CUDA tensor") + if not seq_lens_kv.is_contiguous(): + raise ValueError("seq_lens_kv must be contiguous") + if seq_lens_kv.data_ptr() % 4 != 0: + raise ValueError("seq_lens_kv data pointer must be 4-byte aligned") + if block_table_device != seq_lens_kv.device: + raise ValueError("block_tables and seq_lens_kv must be on the same device") + + batch_size = int(seq_lens_kv.numel()) if batch_size <= 0: - raise ValueError("seq_lens must contain at least one request") + raise ValueError("seq_lens_kv must contain at least one request") if block_table_batch_size != batch_size: raise ValueError( "block_tables must have one row per request: expected " f"{batch_size}, got {block_table_batch_size}" ) - return seq_lens.device, batch_size, table_capacity + return seq_lens_kv.device, batch_size, table_capacity def _decode_output_shape( @@ -1131,8 +1121,8 @@ def _validate_out( q: torch.Tensor, expected_shape: tuple[int, ...], seq_len_q: int, - use_packed_q: bool, - output_dtype: torch.dtype, + packed_query: bool, + out_dtype: torch.dtype, ) -> None: if not isinstance(out, torch.Tensor): raise TypeError("out must be a torch.Tensor") @@ -1140,13 +1130,13 @@ def _validate_out( raise ValueError( f"out must have shape {expected_shape}, got {tuple(out.shape)}" ) - if out.dtype != output_dtype: - raise ValueError(f"out must have dtype {output_dtype}, got {out.dtype}") + if out.dtype != out_dtype: + raise ValueError(f"out must have dtype {out_dtype}, got {out.dtype}") if out.device != q.device: raise ValueError(f"out must be on {q.device}, got {out.device}") layout = ( "[total_q, Hq, D]" - if use_packed_q + if packed_query else ("[B, Hq, D]" if seq_len_q == 1 else "[B, SQ, Hq, D]") ) _validate_exact_compact_strides(out, "out", layout) @@ -1179,7 +1169,14 @@ def _decode_scratch_shapes( heads_q_per_kv=head_ratio, groups_tokens_heads_q=cfg.groups_tokens_heads_q, ) - num_q_groups = max(int(geometry.num_q_ctas(seq_len_q)), 1) + if isinstance(seq_len_q, int): + num_q_groups = max(int(geometry.num_q_ctas(seq_len_q)), 1) + elif cfg.groups_tokens_heads_q: + import cutlass.cute as cute + + num_q_groups = cute.sym_int() + else: + num_q_groups = seq_len_q * geometry.head_ctas_per_token partial_o_shape = ( batch_size, num_kv_heads, @@ -1225,6 +1222,7 @@ def _decode_launch_spec_from_config( ) +# Share CPU policy selection across equivalent plans/layers. @functools.cache def _resolve_decode_launch_spec( device_index: int, @@ -1238,16 +1236,17 @@ def _resolve_decode_launch_spec( q_dtype_key: str, k_dtype_key: str, v_dtype_key: str, - output_dtype_key: str, + out_dtype_key: str, kv_layout: str, mask_type: str, - use_packed_q: bool, + packed_query: bool, window_left: int, storage_page_size: Optional[int] = None, use_q_token_kv_block_sparse_route: bool = False, use_pdl: bool = False, split_kv: bool = True, share_pattern_across_kv_heads: bool = True, + use_attention_sinks: bool = False, ) -> _DecodeLaunchSpec: """Resolve automatic policy and workspace geometry without compiling.""" @@ -1256,7 +1255,7 @@ def _resolve_decode_launch_spec( _validate_decode_query_head_extent( batch_size=batch_size, num_qo_heads=num_qo_heads, - max_seq_len_q=seq_len_q, + max_q_len=seq_len_q, ) max_kv_len = _validate_max_kv_len(max_kv_len, "max_kv_len") window_left = _validate_window_left(window_left, mask_type) @@ -1273,14 +1272,14 @@ def _resolve_decode_launch_spec( make_q_tile_geometry, ) - if kv_layout != "HND": - raise ValueError("the cached TS decode compiler accepts HND only") + _validate_layout(kv_layout) q_dtype = _cutlass_dtype(q_dtype_key) k_dtype = _cutlass_dtype(k_dtype_key) v_dtype = _cutlass_dtype(v_dtype_key) - output_dtype = _cutlass_dtype(output_dtype_key) + out_dtype = _cutlass_dtype(out_dtype_key) arch_args = arch_config_args(torch.cuda.get_device_capability(device_index)) + # Resolve a candidate against this request geometry and device policy. def make_config( args: object | None = None, *, @@ -1300,7 +1299,7 @@ def make_config( q_dtype=q_dtype, k_dtype=k_dtype, v_dtype=v_dtype, - o_dtype=output_dtype, + o_dtype=out_dtype, qkv_layout="pagedKv", num_tokens_per_page=page_size, storage_tokens_per_page=storage_page_size, @@ -1313,8 +1312,10 @@ def make_config( attention_window_size=window_left + 1 if window_left >= 0 else 0, auto_tuner=True, split_kv=split_kv, + use_attention_sinks=use_attention_sinks, ) + # Count logical Q CTAs per KV head, before split-KV fanout. def q_ctas(config: "FmhaDecodeConfig") -> int: geometry = make_q_tile_geometry( rows_per_cta=config.tile_size_q, @@ -1323,6 +1324,7 @@ def q_ctas(config: "FmhaDecodeConfig") -> int: ) return max(int(geometry.num_q_ctas(seq_len_q)), 1) + # Check CTA or cluster occupancy for the complete nonpersistent grid. def fits_one_service_wave(config: "FmhaDecodeConfig", num_q_ctas: int) -> bool: logical_grid = batch_size * num_kv_heads * num_q_ctas if config.use_persistent_scheduler: @@ -1352,9 +1354,9 @@ def fits_one_service_wave(config: "FmhaDecodeConfig", num_q_ctas: int) -> bool: q_dtype_key=q_dtype_key, k_dtype_key=k_dtype_key, v_dtype_key=v_dtype_key, - output_dtype_key=output_dtype_key, + out_dtype_key=out_dtype_key, mask_type=mask_type, - use_packed_q=use_packed_q, + packed_query=packed_query, window_left=window_left, max_active_clusters=max_active_clusters, use_pdl=use_pdl, @@ -1363,7 +1365,7 @@ def fits_one_service_wave(config: "FmhaDecodeConfig", num_q_ctas: int) -> bool: ) else: config_overrides = {"use_pdl": use_pdl} - if use_packed_q: + if packed_query: config_overrides["use_variable_seqlens_q"] = True cfg = make_config(config_overrides) @@ -1373,7 +1375,7 @@ def fits_one_service_wave(config: "FmhaDecodeConfig", num_q_ctas: int) -> bool: # resident wave without reducing KV fanout or changing topology. if ( seq_len_q == 1 - and not use_packed_q + and not packed_query and cfg.groups_tokens_heads_q and not cfg.use_keeps_mma_ab ): @@ -1440,7 +1442,7 @@ def fits_one_service_wave(config: "FmhaDecodeConfig", num_q_ctas: int) -> bool: { "groups_tokens_heads_q": cfg.groups_tokens_heads_q, "tile_size_q": safe_tile_size_q, - "use_variable_seqlens_q": use_packed_q, + "use_variable_seqlens_q": packed_query, "use_pdl": use_pdl, }, split_kv_mode="gmem_reduction_with_separate_kernel", @@ -1448,8 +1450,27 @@ def fits_one_service_wave(config: "FmhaDecodeConfig", num_q_ctas: int) -> bool: max_splits_kv=safe_splits, ) + requested_splits = int(cfg.splits_kv) + if cfg.use_split_kv and not use_q_token_kv_block_sparse_route: + sm_count = get_device_sm_count(torch.device("cuda", device_index)) + requested_splits = min(requested_splits, sm_count) + # GMEM capacity is a device bound, independent of the runtime fanout. + # Keep Q/batch fusion within CUDA's grid-Z limit; larger grids retain + # the original exact-fanout launch. Cluster sizes remain resident topology. + compact_grid = batch_size * q_ctas(cfg) <= 65535 + capacity = ( + requested_splits + if cfg.use_cluster_smem_reduction or not compact_grid + else sm_count + ) + cfg = replace( + cfg, + splits_kv=capacity, + max_splits_kv=capacity, + runtime_split_kv=compact_grid, + ) _validate_decode_policy_kv_tile_size(cfg) - return _decode_launch_spec_from_config( + resolved = _decode_launch_spec_from_config( cfg, batch_size=batch_size, num_qo_heads=num_qo_heads, @@ -1458,6 +1479,12 @@ def fits_one_service_wave(config: "FmhaDecodeConfig", num_q_ctas: int) -> bool: seq_len_q=seq_len_q, max_active_clusters=int(max_active_clusters), ) + # Report the runtime fanout separately from the SM-bound workspace capacity. + policy = dict(resolved.policy) + policy["splits_kv"] = requested_splits + return replace( + resolved, policy=tuple(policy.items()), requested_splits=requested_splits + ) def _make_decode_compile_spec( @@ -1468,21 +1495,45 @@ def _make_decode_compile_spec( num_kv_heads: int, head_dim: int, page_size: int, - max_kv_len: int, - seq_len_q: int, q_dtype_key: str, k_dtype_key: str, v_dtype_key: str, - output_dtype_key: str, - use_packed_q: bool, - kv_prefix_mode: Literal["dynamic", "planned_full"], - kv_lengths_mode: Literal["dynamic", "planned_uniform_max"], + out_dtype_key: str, direct_q1_spec: Optional[_DirectQ1CompileSpec] = None, flat_native_kv_tma: bool = False, + lse_mode: str = "none", + runtime_kv_lengths: bool = False, + packed_query: Optional[bool] = None, + kv_prefix_mode: Literal["dynamic", "planned_full"] = "dynamic", + kv_lengths_mode: Literal["dynamic", "planned_uniform_max"] = "dynamic", ) -> _DecodeCompileSpec: - """Freeze the resolved topology while leaving batch in runtime tensors.""" + """Freeze the resolved topology while leaving batch and lengths in runtime tensors. + The config's compile signature already drops every length-derived field, and + all requests launch through the flattened-Q ABI. Uniform inputs derive + their offsets from the runtime tensor extent; packed inputs read metadata. + """ + + if launch_spec.config.fixed_seqlen: + raise ValueError("cached TS decode compiles runtime sequence lengths only") config_items = dict(launch_spec.config.compile_signature()) + # Public wrapper plans use runtime KV metadata in compiled kernels. False is retained only + # for internal prepared/sparse recipes with a fixed native layout. + if runtime_kv_lengths: + # Odd, non-aligned geometry keeps static full-tile shortcuts disabled. + # Native paged requests obtain every real bound from live metadata. + cfg = launch_spec.config + max_kv_len = int(cfg.tile_size_kv * cfg.num_insts_kv * 2 + 1) + config_items["total_kv_tiles"] = ( + max_kv_len + cfg.tile_size_kv - 1 + ) // cfg.tile_size_kv + config_items["static_seq_len_kv"] = max_kv_len + config_items["max_num_pages_per_seq_kv"] = ( + max_kv_len + cfg.num_tokens_per_page - 1 + ) // cfg.num_tokens_per_page + config_items["static_num_skipped_kv_tiles"] = 0 + config_items["static_window_start_idx"] = 0 + config_items["lse_mode"] = lse_mode config_items["use_flat_native_kv_tma"] = flat_native_kv_tma config_items["flat_native_kv_num_heads"] = num_kv_heads if flat_native_kv_tma else 1 return _DecodeCompileSpec( @@ -1492,17 +1543,21 @@ def _make_decode_compile_spec( num_kv_heads=num_kv_heads, head_dim=head_dim, page_size=page_size, - max_kv_len=max_kv_len, - seq_len_q=seq_len_q, q_dtype_key=q_dtype_key, k_dtype_key=k_dtype_key, v_dtype_key=v_dtype_key, - output_dtype_key=output_dtype_key, - use_packed_q=use_packed_q, + out_dtype_key=out_dtype_key, + packed_query=( + launch_spec.config.use_variable_seqlens_q + if packed_query is None + else packed_query + ) + or direct_q1_spec is not None, max_active_clusters=launch_spec.max_active_clusters, kv_prefix_mode=kv_prefix_mode, kv_lengths_mode=kv_lengths_mode, direct_q1_spec=direct_q1_spec, + runtime_kv_lengths=runtime_kv_lengths, ) @@ -1516,13 +1571,13 @@ def _get_compiled_decode( num_qo_heads = compile_spec.num_qo_heads num_kv_heads = compile_spec.num_kv_heads head_dim = compile_spec.head_dim - max_kv_len = compile_spec.max_kv_len - seq_len_q = compile_spec.seq_len_q q_dtype_key = compile_spec.q_dtype_key k_dtype_key = compile_spec.k_dtype_key v_dtype_key = compile_spec.v_dtype_key - output_dtype_key = compile_spec.output_dtype_key - use_packed_q = compile_spec.use_packed_q + out_dtype_key = compile_spec.out_dtype_key + # Both layouts retain the flattened tensor ABI; only the offset source differs. + packed_query = True + uniform_q = not compile_spec.packed_query max_active_clusters = compile_spec.max_active_clusters kv_prefix_mode = compile_spec.kv_prefix_mode kv_lengths_mode = compile_spec.kv_lengths_mode @@ -1541,24 +1596,28 @@ def _get_compiled_decode( from .kernels.fmha_decode.fmha_decode_config import FmhaDecodeConfig from .kernels.fmha_decode.fmha_decode_kernel import fmha_decode_launch from .kernels.fmha_decode.direct_sparse_metadata import DirectSparseMetadataView + from .kernels.fmha_decode.fmha_decode_resources.helpers_common import ( + UniformQueryOffsets, + ) + seq_len_q = cute.sym_int() direct_q1_spec = compile_spec.direct_q1_spec q_dtype = _cutlass_dtype(q_dtype_key) k_dtype = _cutlass_dtype(k_dtype_key) v_dtype = _cutlass_dtype(v_dtype_key) - output_dtype = _cutlass_dtype(output_dtype_key) + out_dtype = _cutlass_dtype(out_dtype_key) cfg = FmhaDecodeConfig(**dict(compile_spec.config_items)) k_storage_dtype = cutlass.Uint8 if k_dtype == cutlass.Float4E2M1FN else k_dtype v_storage_dtype = cutlass.Uint8 if v_dtype == cutlass.Float4E2M1FN else v_dtype storage_page_size = int(cfg.effective_storage_tokens_per_page) - partial_dtype = output_dtype - if cfg.use_separate_reduction_kernel and output_dtype in ( + partial_dtype = out_dtype + if cfg.use_separate_reduction_kernel and out_dtype in ( cutlass.BFloat16, cutlass.Float8E4M3FN, ): partial_dtype = cutlass.BFloat16 - elif output_dtype == cutlass.Float8E4M3FN or not cfg.use_split_kv: + elif out_dtype == cutlass.Float8E4M3FN or not cfg.use_split_kv: partial_dtype = cutlass.Float16 Int32 = cutlass.Int32 @@ -1573,7 +1632,8 @@ def main_tensor_adapter( k_sf_cache: cute.Tensor, v_sf_cache: cute.Tensor, out: cute.Tensor, - seq_lens: cute.Tensor, + lse: cute.Tensor | None, + seq_lens_kv: cute.Tensor, cu_seqlens_q: cute.Tensor, block_table: cute.Tensor, q_token_kv_block_sparse_page_memberships: cute.Tensor, @@ -1588,16 +1648,17 @@ def main_tensor_adapter( v_page_stride: cutlass.Int64, v_head_stride: cutlass.Int64, v_token_stride: cutlass.Int64, - bmm1_scale: cutlass.Float32, - bmm2_scale: cutlass.Float32, + softmax_scale: cutlass.Float32, + output_scale: cutlass.Float32, + requested_splits: cutlass.Int32, + max_q_len: cutlass.Int32, direct_q1_inputs: tuple, + runtime_max_kv_len: cutlass.Int32, stream: cuda_drv.CUstream, static_cfg: cutlass.Constexpr[FmhaDecodeConfig], - static_seq_len_q: cutlass.Constexpr[int], static_num_qo_heads: cutlass.Constexpr[int], static_num_kv_heads: cutlass.Constexpr[int], static_head_dim: cutlass.Constexpr[int], - static_max_kv_len: cutlass.Constexpr[int], static_max_active_clusters: cutlass.Constexpr[int], static_full_split_prefix: cutlass.Constexpr[bool], static_native_uniform_kv: cutlass.Constexpr[bool], @@ -1607,9 +1668,11 @@ def main_tensor_adapter( pattern_heads = _num_sparse_pattern_heads( static_num_kv_heads, static_cfg.shares_sparse_pattern ) - batch_size = Int32(cute.size(seq_lens)) // Int32(pattern_heads) + batch_size = Int32(cute.size(seq_lens_kv)) // Int32(pattern_heads) q_offsets_iter = cu_seqlens_q.iterator - total_q_tokens = batch_size * Int32(static_seq_len_q) + if cutlass.const_expr(uniform_q): + q_offsets_iter = UniformQueryOffsets(Int32(q.shape[0]) // batch_size) + total_q_tokens = Int32(q.shape[0]) if cutlass.const_expr(not static_cfg.use_variable_seqlens_q): # Fixed-Q is a distinct specialization. Keep a uniform TVM-FFI # wrapper signature, but pass a real null pointer to the kernel so @@ -1618,17 +1681,20 @@ def main_tensor_adapter( else: total_q_tokens = Int32(q.shape[0]) - lengths_iter = seq_lens.iterator + lengths_iter = seq_lens_kv.iterator pages_iter = block_table.iterator if cutlass.const_expr(direct_q1_spec is not None): lengths_iter = DirectSparseMetadataView( direct_q1_inputs, cu_seqlens_q, packed=static_cfg.use_variable_seqlens_q, - model_len=direct_q1_spec.max_seq_len_kv, + model_len=direct_q1_inputs[4], page_size=storage_page_size, - page_capacity=direct_q1_spec.page_capacity, + page_capacity=direct_q1_inputs[5], sparse_block_size=direct_q1_spec.kv_block_size, + count_current_block_in_topk=direct_q1_spec.count_current_block_in_topk, + fixed_query_len=direct_q1_inputs[6], + fixed_query_layout=direct_q1_spec.fixed_query_layout, fragment_size=static_cfg.num_tokens_per_page, lengths=True, ) @@ -1636,10 +1702,13 @@ def main_tensor_adapter( direct_q1_inputs, cu_seqlens_q, packed=static_cfg.use_variable_seqlens_q, - model_len=direct_q1_spec.max_seq_len_kv, + model_len=direct_q1_inputs[4], page_size=storage_page_size, - page_capacity=direct_q1_spec.page_capacity, + page_capacity=direct_q1_inputs[5], sparse_block_size=direct_q1_spec.kv_block_size, + count_current_block_in_topk=direct_q1_spec.count_current_block_in_topk, + fixed_query_len=direct_q1_inputs[6], + fixed_query_layout=direct_q1_spec.fixed_query_layout, fragment_size=static_cfg.num_tokens_per_page, lengths=False, ) @@ -1649,7 +1718,7 @@ def main_tensor_adapter( batch_size, Int32(static_num_qo_heads), Int32(static_num_kv_heads), - Int32(static_max_kv_len), + runtime_max_kv_len, Int32(static_head_dim), ), q.iterator, @@ -1667,13 +1736,13 @@ def main_tensor_adapter( partial_stats.iterator, split_kv_counter.iterator, attention_sinks.iterator, - bmm1_scale, - bmm2_scale, + softmax_scale, + output_scale, Int32(0), Int32(static_max_active_clusters), stream, static_cfg, - static_max_kv_len, + None, False, True, Int64(block_table.stride[0]), @@ -1688,6 +1757,14 @@ def main_tensor_adapter( v_token_stride, static_full_split_prefix, static_native_uniform_kv, + lse_ptr=lse.iterator if lse is not None else None, + requested_splits=requested_splits if static_cfg.runtime_split_kv else None, + max_q_len=max_q_len, + runtime_q_capacity_rows=( + Int32(partial_o.shape[-2]) + if static_cfg.use_split_kv + else max_q_len * Int32(static_num_qo_heads // static_num_kv_heads) + ), ) reduction_tensor_adapter = None @@ -1699,20 +1776,23 @@ def main_tensor_adapter( @cute.jit def reduction_tensor_adapter( out: cute.Tensor, - seq_lens: cute.Tensor, + lse: cute.Tensor | None, + seq_lens_kv: cute.Tensor, cu_seqlens_q: cute.Tensor, partial_o: cute.Tensor, partial_stats: cute.Tensor, attention_sinks: cute.Tensor, - bmm1_scale: cutlass.Float32, - bmm2_scale: cutlass.Float32, + softmax_scale: cutlass.Float32, + output_scale: cutlass.Float32, + requested_splits: cutlass.Int32, + max_q_len: cutlass.Int32, direct_q1_inputs: tuple, + runtime_max_kv_len: cutlass.Int32, stream: cuda_drv.CUstream, static_cfg: cutlass.Constexpr[FmhaDecodeConfig], static_num_qo_heads: cutlass.Constexpr[int], static_num_kv_heads: cutlass.Constexpr[int], static_head_dim: cutlass.Constexpr[int], - static_max_kv_len: cutlass.Constexpr[int], static_full_split_prefix: cutlass.Constexpr[bool], ) -> None: """Adapt TVM-FFI tensors to the raw standalone split reducer.""" @@ -1720,21 +1800,26 @@ def reduction_tensor_adapter( pattern_heads = _num_sparse_pattern_heads( static_num_kv_heads, static_cfg.shares_sparse_pattern ) - batch_size = Int32(cute.size(seq_lens)) // Int32(pattern_heads) + batch_size = Int32(cute.size(seq_lens_kv)) // Int32(pattern_heads) q_offsets_iter = cu_seqlens_q.iterator + if cutlass.const_expr(uniform_q): + q_offsets_iter = UniformQueryOffsets(Int32(out.shape[0]) // batch_size) if cutlass.const_expr(not static_cfg.use_variable_seqlens_q): q_offsets_iter = cute.make_ptr(Int32, 0) - lengths_iter = seq_lens.iterator + lengths_iter = seq_lens_kv.iterator if cutlass.const_expr(direct_q1_spec is not None): lengths_iter = DirectSparseMetadataView( direct_q1_inputs, cu_seqlens_q, packed=static_cfg.use_variable_seqlens_q, - model_len=direct_q1_spec.max_seq_len_kv, + model_len=direct_q1_inputs[4], page_size=storage_page_size, - page_capacity=direct_q1_spec.page_capacity, + page_capacity=direct_q1_inputs[5], sparse_block_size=direct_q1_spec.kv_block_size, + count_current_block_in_topk=direct_q1_spec.count_current_block_in_topk, + fixed_query_len=direct_q1_inputs[6], + fixed_query_layout=direct_q1_spec.fixed_query_layout, fragment_size=static_cfg.num_tokens_per_page, lengths=True, ) @@ -1743,7 +1828,7 @@ def reduction_tensor_adapter( batch_size, Int32(static_num_qo_heads), Int32(static_num_kv_heads), - Int32(static_max_kv_len), + runtime_max_kv_len, Int32(static_head_dim), ), out.iterator, @@ -1752,11 +1837,16 @@ def reduction_tensor_adapter( partial_o.iterator, partial_stats.iterator, attention_sinks.iterator, - bmm1_scale, - bmm2_scale, + softmax_scale, + output_scale, stream, static_cfg, static_full_split_prefix, + lse_ptr=lse.iterator if lse is not None else None, + requested_splits=requested_splits + if static_cfg.runtime_split_kv + else None, + runtime_q_capacity_rows=Int32(partial_o.shape[-2]), ) physical_pages = cute.sym_int() @@ -1774,7 +1864,7 @@ def reduction_tensor_adapter( runtime_num_q_offsets = cute.sym_int() q_shape = ( (total_q_tokens, num_qo_heads, head_dim) - if use_packed_q + if packed_query else ( (batch_size, num_qo_heads, head_dim) if seq_len_q == 1 @@ -1816,15 +1906,16 @@ def reduction_tensor_adapter( num_qo_heads=num_qo_heads, seq_len_q=seq_len_q, head_dim=head_dim, - total_q_tokens=total_q_tokens if use_packed_q else None, + total_q_tokens=total_q_tokens if packed_query else None, ) out_fake = cute.runtime.make_fake_compact_tensor( - output_dtype, + out_dtype, out_shape, stride_order=tuple(reversed(range(len(out_shape)))), assumed_align=16, ) + # Describe row-major symbolic storage; assumed_align is in bytes. def fake_compact(dtype, shape, assumed_align): return cute.runtime.make_fake_compact_tensor( dtype, @@ -1833,11 +1924,14 @@ def fake_compact(dtype, shape, assumed_align): assumed_align=assumed_align, ) + lse_fake = ( + fake_compact(Float32, q_shape[:-1], 4) if cfg.lse_mode != "none" else None + ) pattern_heads = _num_sparse_pattern_heads(num_kv_heads, cfg.shares_sparse_pattern) metadata_rows = batch_size * pattern_heads - seq_lens_fake = fake_compact(Int32, (metadata_rows,), 4) + seq_lens_kv_fake = fake_compact(Int32, (metadata_rows,), 4) cu_seqlens_q_fake = fake_compact( - Int32, (runtime_num_q_offsets,) if use_packed_q else (1,), 4 + Int32, (runtime_num_q_offsets,) if packed_query else (1,), 4 ) block_table_fake = cute.runtime.make_fake_tensor( Int32, @@ -1859,7 +1953,9 @@ def fake_compact(dtype, shape, assumed_align): partial_o_fake = fake_compact(partial_dtype, partial_o_shape, 16) partial_stats_fake = fake_compact(Float32, partial_stats_shape, 16) counter_fake = fake_compact(Int32, counter_shape, 4) - attention_sinks_fake = fake_compact(Float32, (1,), 4) + attention_sinks_fake = fake_compact( + Float32, (num_qo_heads,) if cfg.use_attention_sinks else (1,), 4 + ) if cfg.use_nvfp4_kv: sf_shape = (physical_pages, num_kv_heads, storage_page_size, head_dim // 16) k_sf_fake = fake_compact(cutlass.Float8E4M3FN, sf_shape, 16) @@ -1875,9 +1971,9 @@ def fake_compact(dtype, shape, assumed_align): cute.runtime.make_fake_tensor( Int32, ( - (rows, direct_q1_spec.block_topk) + (rows, cute.sym_int()) if cfg.shares_sparse_pattern - else (rows, num_kv_heads, direct_q1_spec.block_topk) + else (rows, num_kv_heads, cute.sym_int()) ), stride=( (cute.sym_int64(), 1) @@ -1892,85 +1988,99 @@ def fake_compact(dtype, shape, assumed_align): stride=(cute.sym_int64(), 1), assumed_align=4, ), - fake_compact(Int32, (rows,), 4), + ( + None + if direct_q1_spec.fixed_query_layout + else fake_compact(Int32, (rows,), 4) + ), fake_compact( Int64 if direct_q1_spec.position_dtype_key == "int64" else Int32, - (rows,), + (cute.sym_int(),) if direct_q1_spec.fixed_query_layout else (rows,), 4, ), - ) - - def compile_main(): - return cute.compile( - main_tensor_adapter, - q_fake, - k_fake, - v_fake, - k_sf_fake, - v_sf_fake, - out_fake, - seq_lens_fake, - cu_seqlens_q_fake, - block_table_fake, - q_token_kv_block_sparse_page_memberships_fake, - partial_o_fake, - partial_stats_fake, - counter_fake, - attention_sinks_fake, - Int64(1), - Int64(1), - Int64(1), - Int64(1), - Int64(1), - Int64(1), - Int64(1), - Float32(1.0), - Float32(1.0), - direct_q1_fake, - stream_fake, - cfg, - seq_len_q, - num_qo_heads, - num_kv_heads, - head_dim, - max_kv_len, - max_active_clusters, - static_full_split_prefix, - static_native_uniform_kv, - options=_COMPILE_OPTIONS, - ) - - def compile_reducer(): - assert reduction_tensor_adapter is not None - return cute.compile( - reduction_tensor_adapter, - out_fake, - seq_lens_fake, - cu_seqlens_q_fake, - partial_o_fake, - partial_stats_fake, - attention_sinks_fake, - Float32(1.0), - Float32(1.0), - direct_q1_fake, - stream_fake, - cfg, - num_qo_heads, - num_kv_heads, - head_dim, - max_kv_len, - static_full_split_prefix, - options=_COMPILE_OPTIONS, + Int32(0), + Int32(0), + Int32(0), ) with torch.cuda.device(device_index): compiled_main = build_and_load_attention_kernel( - "decode", "main", compile_spec, _COMPILE_OPTIONS, compile_main + "decode", + "main", + compile_spec, + _COMPILE_OPTIONS, + lambda: cute.compile( + main_tensor_adapter, + q_fake, + k_fake, + v_fake, + k_sf_fake, + v_sf_fake, + out_fake, + lse_fake, + seq_lens_kv_fake, + cu_seqlens_q_fake, + block_table_fake, + q_token_kv_block_sparse_page_memberships_fake, + partial_o_fake, + partial_stats_fake, + counter_fake, + attention_sinks_fake, + Int64(1), + Int64(1), + Int64(1), + Int64(1), + Int64(1), + Int64(1), + Int64(1), + Float32(1.0), + Float32(1.0), + Int32(1), + Int32(0), + direct_q1_fake, + Int32(0), + stream_fake, + cfg, + num_qo_heads, + num_kv_heads, + head_dim, + max_active_clusters, + static_full_split_prefix, + static_native_uniform_kv, + options=_COMPILE_OPTIONS, + ), ) compiled_reducer = None if cfg.use_separate_reduction_kernel: + assert reduction_tensor_adapter is not None compiled_reducer = build_and_load_attention_kernel( - "decode", "reducer", compile_spec, _COMPILE_OPTIONS, compile_reducer + "decode", + "reducer", + compile_spec, + _COMPILE_OPTIONS, + lambda: cute.compile( + reduction_tensor_adapter, + out_fake, + lse_fake, + seq_lens_kv_fake, + cu_seqlens_q_fake, + partial_o_fake, + partial_stats_fake, + attention_sinks_fake, + Float32(1.0), + Float32(1.0), + Int32(1), + Int32(0), + direct_q1_fake, + Int32(0), + stream_fake, + cfg, + num_qo_heads, + num_kv_heads, + head_dim, + static_full_split_prefix, + options=_COMPILE_OPTIONS, + ), ) return compiled_main, compiled_reducer @@ -1982,11 +2092,11 @@ def get_prims_ts_batch_decode_workspace_size( num_kv_heads: int, head_dim: int, page_size: int, - max_seq_len: int, + max_kv_len: int, *, seq_len_q: int = 1, qo_indptr: Optional[torch.Tensor] = None, - max_seq_len_q: Optional[int] = None, + max_q_len: Optional[int] = None, q_dtype: torch.dtype = torch.float16, k_dtype: Optional[torch.dtype] = None, v_dtype: Optional[torch.dtype] = None, @@ -1994,10 +2104,11 @@ def get_prims_ts_batch_decode_workspace_size( out_dtype: Optional[torch.dtype] = None, mask_type: Literal["dense", "causal"] = "dense", window_left: int = -1, - kv_layout: Literal["HND"] = "HND", + kv_layout: Literal["HND", "NHD"] = "HND", storage_page_size: Optional[int] = None, device: Optional[Union[int, str, torch.device]] = None, split_kv: bool = True, + use_attention_sinks: bool = False, ) -> int: """Return caller-workspace bytes for one automatic FMHA policy. @@ -2009,8 +2120,8 @@ def get_prims_ts_batch_decode_workspace_size( including ``batch_size``, changes because the internal section offsets can move even when the compiled callable is reused. Fixed-Q launches use ``seq_len_q``. Packed-Q launches provide ``qo_indptr`` and the explicit - static ``max_seq_len_q`` bound used for workspace geometry and JIT policy. - ``max_seq_len`` must be no larger than ``2,147,483,392`` so the padded + static ``max_q_len`` bound used for workspace geometry and JIT policy. + ``max_kv_len`` must be no larger than ``2,147,483,392`` so the padded 256-token K/V tile endpoint remains representable as signed Int32. ``storage_page_size`` defaults to the semantic ``page_size``. A larger physical page is supported only for semantic page size four and activates @@ -2025,13 +2136,15 @@ def get_prims_ts_batch_decode_workspace_size( ``split_kv`` permits automatic split fanout when True (default), or forces nonsplit execution when False. It is independent of Q layout and must match the value supplied to attention planning/launch. + ``use_attention_sinks`` must also match planning because it can affect + the selected policy and scratch layout. """ batch_size = _validate_positive_int(batch_size, "batch_size") - use_packed_q, resolved_seq_len_q = _resolve_q_mode( + packed_query, resolved_seq_len_q = _resolve_q_mode( seq_len_q=seq_len_q, qo_indptr=qo_indptr, - max_seq_len_q=max_seq_len_q, + max_q_len=max_q_len, require_packed_max=True, ) assert resolved_seq_len_q is not None @@ -2040,7 +2153,7 @@ def get_prims_ts_batch_decode_workspace_size( _validate_decode_query_head_extent( batch_size=batch_size, num_qo_heads=num_qo_heads, - max_seq_len_q=seq_len_q, + max_q_len=seq_len_q, ) head_dim = _validate_head_dim(head_dim) page_size = _validate_page_size(page_size) @@ -2048,7 +2161,7 @@ def get_prims_ts_batch_decode_workspace_size( page_size, page_size if storage_page_size is None else storage_page_size, ) - max_seq_len = _validate_max_kv_len(max_seq_len, "max_seq_len") + max_kv_len = _validate_max_kv_len(max_kv_len, "max_kv_len") _validate_layout(kv_layout) _validate_mask(mask_type) window_left = _validate_window_left(window_left, mask_type) @@ -2070,7 +2183,7 @@ def get_prims_ts_batch_decode_workspace_size( ) _validate_packed_q_plan_values( qo_indptr, - max_seq_len_q=seq_len_q, + max_q_len=seq_len_q, ) return _resolve_decode_workspace_layout( @@ -2079,7 +2192,7 @@ def get_prims_ts_batch_decode_workspace_size( num_kv_heads, head_dim, page_size, - max_seq_len, + max_kv_len, seq_len_q, q_dtype, k_dtype, @@ -2087,11 +2200,12 @@ def get_prims_ts_batch_decode_workspace_size( out_dtype, kv_layout, mask_type, - use_packed_q, + packed_query, window_left, storage_page_size, resolved_device, split_kv=split_kv, + use_attention_sinks=use_attention_sinks, ).total_bytes @@ -2103,7 +2217,7 @@ def _prepare_decode_runtime( device: torch.device, batch_size: int, seq_len_q: int, - use_packed_q: bool, + packed_query: bool, num_qo_heads: int, num_kv_heads: int, head_dim: int, @@ -2111,10 +2225,11 @@ def _prepare_decode_runtime( q_dtype: torch.dtype, k_dtype: torch.dtype, v_dtype: torch.dtype, - output_dtype: torch.dtype, - bmm1_scale: Optional[float], - bmm2_scale: float, + out_dtype: torch.dtype, + softmax_scale: Optional[float], + output_scale: float, out: Optional[torch.Tensor], + kv_layout: str = "HND", ) -> _DecodeRuntime: """Validate runtime tensors and normalize zero-copy K/V views.""" @@ -2123,7 +2238,7 @@ def _prepare_decode_runtime( _validate_q( q, seq_len_q=seq_len_q, - use_packed_q=use_packed_q, + packed_query=packed_query, device=device, batch_size=batch_size, num_qo_heads=num_qo_heads, @@ -2133,6 +2248,7 @@ def _prepare_decode_runtime( normalized_cache = _normalize_native_paged_kv_cache( paged_kv_cache, expected_device=device, + kv_layout=kv_layout, ) k_cache = normalized_cache.k_cache v_cache = normalized_cache.v_cache @@ -2158,28 +2274,28 @@ def _prepare_decode_runtime( k_cache=k_cache, logical_head_dim=head_dim, ) - effective_bmm1_scale = _validate_scale( - 1.0 / math.sqrt(head_dim) if bmm1_scale is None else bmm1_scale, - "bmm1_scale", + effective_softmax_scale = _validate_scale( + 1.0 / math.sqrt(head_dim) if softmax_scale is None else softmax_scale, + "softmax_scale", ) - effective_bmm2_scale = _validate_scale(bmm2_scale, "bmm2_scale") + effective_output_scale = _validate_scale(output_scale, "output_scale") output_shape = _decode_output_shape( batch_size=batch_size, num_qo_heads=num_qo_heads, seq_len_q=seq_len_q, head_dim=head_dim, - total_q_tokens=int(q.shape[0]) if use_packed_q else None, + total_q_tokens=int(q.shape[0]) if packed_query else None, ) if out is None: - out = torch.empty(output_shape, device=device, dtype=output_dtype) + out = torch.empty(output_shape, device=device, dtype=out_dtype) else: _validate_out( out, q=q, expected_shape=output_shape, seq_len_q=seq_len_q, - use_packed_q=use_packed_q, - output_dtype=output_dtype, + packed_query=packed_query, + out_dtype=out_dtype, ) return _DecodeRuntime( q=q, @@ -2195,39 +2311,46 @@ def _prepare_decode_runtime( v_page_stride=normalized_cache.v_page_stride, v_head_stride=normalized_cache.v_head_stride, v_token_stride=normalized_cache.v_token_stride, - bmm1_scale=effective_bmm1_scale, - bmm2_scale=effective_bmm2_scale, + softmax_scale=effective_softmax_scale, + output_scale=effective_output_scale, ) def _launch_decode( runtime: _DecodeRuntime, *, - seq_lens: torch.Tensor, + seq_lens_kv: torch.Tensor, qo_indptr: Optional[torch.Tensor], block_tables: torch.Tensor, workspace: _DecodeWorkspaceViews, compiled_main: Callable[..., object], compiled_reducer: Optional[Callable[..., object]], + lse: Optional[torch.Tensor] = None, + attention_sinks: Optional[torch.Tensor] = None, + requested_splits: int = 1, + max_kv_len: int, + max_q_len: int, ) -> torch.Tensor: """Launch the compiled main kernel and its optional standalone reducer.""" + sinks = workspace.attention_sinks if attention_sinks is None else attention_sinks q_offsets = workspace.cu_seqlens_q if qo_indptr is None else qo_indptr compiled_main( - runtime.q, + runtime.q.view(-1, runtime.q.shape[-2], runtime.q.shape[-1]), runtime.k_cache, runtime.v_cache, runtime.k_sf_cache, runtime.v_sf_cache, - runtime.out, - seq_lens, + runtime.out.view(-1, runtime.out.shape[-2], runtime.out.shape[-1]), + None if lse is None else lse.view(-1, lse.shape[-1]), + seq_lens_kv, q_offsets, block_tables, block_tables, # Dense decode constexpr-elides the membership table. workspace.partial_o, workspace.partial_stats, workspace.split_kv_counter, - workspace.attention_sinks, + sinks, runtime.num_physical_pages, runtime.k_page_stride, runtime.k_head_stride, @@ -2235,69 +2358,76 @@ def _launch_decode( runtime.v_page_stride, runtime.v_head_stride, runtime.v_token_stride, - runtime.bmm1_scale, - runtime.bmm2_scale, + runtime.softmax_scale, + runtime.output_scale, + requested_splits, + max_q_len, (), + max_kv_len, ) if compiled_reducer is not None: compiled_reducer( - runtime.out, - seq_lens, + runtime.out.view(-1, runtime.out.shape[-2], runtime.out.shape[-1]), + None if lse is None else lse.view(-1, lse.shape[-1]), + seq_lens_kv, q_offsets, workspace.partial_o, workspace.partial_stats, - workspace.attention_sinks, - runtime.bmm1_scale, - runtime.bmm2_scale, + sinks, + runtime.softmax_scale, + runtime.output_scale, + requested_splits, + max_q_len, (), + max_kv_len, ) return runtime.out -def _normalize_plan_seq_lens( - seq_lens: Optional[Union[Sequence[int], torch.Tensor]], +def _normalize_plan_seq_lens_kv( + seq_lens_kv: Optional[Union[Sequence[int], torch.Tensor]], *, batch_size: int, max_kv_len: int, ) -> Optional[tuple[int, ...]]: """Validate optional host-only lengths retained by the plan.""" - if seq_lens is None: + if seq_lens_kv is None: return None - if isinstance(seq_lens, torch.Tensor): - if seq_lens.device.type != "cpu": - raise ValueError("plan seq_lens must be on CPU") - if seq_lens.ndim != 1: - raise ValueError("plan seq_lens must be 1D") - if seq_lens.dtype not in (torch.int32, torch.int64): - raise TypeError("plan seq_lens must have int32 or int64 dtype") - raw_values: Sequence[object] = seq_lens.tolist() - elif isinstance(seq_lens, Sequence) and not isinstance( - seq_lens, (str, bytes, bytearray) + if isinstance(seq_lens_kv, torch.Tensor): + if seq_lens_kv.device.type != "cpu": + raise ValueError("plan seq_lens_kv must be on CPU") + if seq_lens_kv.ndim != 1: + raise ValueError("plan seq_lens_kv must be 1D") + if seq_lens_kv.dtype not in (torch.int32, torch.int64): + raise TypeError("plan seq_lens_kv must have int32 or int64 dtype") + raw_values: Sequence[object] = seq_lens_kv.tolist() + elif isinstance(seq_lens_kv, Sequence) and not isinstance( + seq_lens_kv, (str, bytes, bytearray) ): - raw_values = seq_lens + raw_values = seq_lens_kv else: raise TypeError( - "plan seq_lens must be a host sequence of integers or a CPU tensor" + "plan seq_lens_kv must be a host sequence of integers or a CPU tensor" ) if len(raw_values) != batch_size: raise ValueError( - "plan seq_lens must contain exactly " + "plan seq_lens_kv must contain exactly " f"batch_size ({batch_size}) values, got {len(raw_values)}" ) normalized = [] for request_idx, value in enumerate(raw_values): if isinstance(value, bool) or not isinstance(value, numbers.Integral): raise TypeError( - "plan seq_lens must contain integers; " + "plan seq_lens_kv must contain integers; " f"request {request_idx} has {value!r}" ) - seq_len = int(value) - if seq_len <= 0 or seq_len > max_kv_len: + seq_len = max(int(value), 0) + if seq_len > max_kv_len: raise ValueError( - "plan seq_lens values must be within " - f"[1, {max_kv_len}]; request {request_idx} has {seq_len}" + "plan seq_lens_kv values must be within " + f"[0, {max_kv_len}]; request {request_idx} has {seq_len}" ) normalized.append(seq_len) return tuple(normalized) @@ -2308,10 +2438,11 @@ def _prepare_decode_runtime_unchecked( paged_kv_cache: PagedKVCache, kv_scale_factors: Optional[PagedKVScaleFactors], *, - output_dtype: torch.dtype, - bmm1_scale: Optional[float], - bmm2_scale: float, + out_dtype: torch.dtype, + softmax_scale: Optional[float], + output_scale: float, out: Optional[torch.Tensor], + kv_layout: str = "HND", ) -> _DecodeRuntime: """Canonicalize one trusted run without invoking explicit validators.""" @@ -2320,6 +2451,9 @@ def _prepare_decode_runtime_unchecked( v_cache = paged_kv_cache[:, 1] else: k_cache, v_cache = paged_kv_cache + if kv_layout == "NHD": + k_cache = k_cache.transpose(1, 2) + v_cache = v_cache.transpose(1, 2) if k_cache.dtype == torch.uint8: k_sf_cache, v_sf_cache = cast(PagedKVScaleFactors, kv_scale_factors) else: @@ -2329,7 +2463,7 @@ def _prepare_decode_runtime_unchecked( out = torch.empty( q.shape, device=q.device, - dtype=output_dtype, + dtype=out_dtype, ) return _DecodeRuntime( q=q, @@ -2345,46 +2479,48 @@ def _prepare_decode_runtime_unchecked( v_page_stride=int(v_cache.stride(0)), v_head_stride=int(v_cache.stride(1)), v_token_stride=int(v_cache.stride(2)), - bmm1_scale=( + softmax_scale=( 1.0 / math.sqrt(int(q.shape[-1])) - if bmm1_scale is None - else float(bmm1_scale) + if softmax_scale is None + else float(softmax_scale) ), - bmm2_scale=float(bmm2_scale), + output_scale=float(output_scale), ) def _validate_decode_run_metadata_values( runtime: _DecodeRuntime, *, - seq_lens: torch.Tensor, + seq_lens_kv: torch.Tensor, block_tables: torch.Tensor, qo_indptr: Optional[torch.Tensor], - planned_seq_lens_host: Optional[tuple[int, ...]], max_kv_len: int, page_size: int, - use_packed_q: bool, + packed_query: bool, seq_len_q: int, batch_size: int, mask_type: str, -) -> None: - """Synchronously validate per-run metadata values.""" + planned_seq_lens_kv_host: Optional[tuple[int, ...]] = None, +) -> tuple[int, int]: + """Validate metadata and return live maxima from the same readback.""" - runtime_seq_lens = planned_seq_lens_host - if runtime_seq_lens is None: - runtime_seq_lens = tuple(int(value) for value in seq_lens.tolist()) + runtime_seq_lens_kv = planned_seq_lens_kv_host + if runtime_seq_lens_kv is None: + runtime_seq_lens_kv = tuple(int(value) for value in seq_lens_kv.tolist()) table_capacity = int(block_tables.shape[1]) - for request_idx, seq_len in enumerate(runtime_seq_lens): - if seq_len <= 0 or seq_len > max_kv_len: + for request_idx, seq_len in enumerate(runtime_seq_lens_kv): + if seq_len <= 0: + continue + if seq_len > max_kv_len: raise ValueError( - f"seq_lens values must be within [1, {max_kv_len}]; " + f"seq_lens_kv values must be within [1, {max_kv_len}]; " f"request {request_idx} has {seq_len}" ) required_pages = (seq_len + page_size - 1) // page_size if table_capacity < required_pages: raise ValueError( "block_tables does not have enough columns for " - f"seq_lens[{request_idx}]={seq_len}: requires {required_pages}, " + f"seq_lens_kv[{request_idx}]={seq_len}: requires {required_pages}, " f"got {table_capacity}" ) @@ -2393,8 +2529,10 @@ def _validate_decode_run_metadata_values( int(runtime.k_cache.shape[2]) // page_size ) for request_idx, (row, seq_len) in enumerate( - zip(block_table_rows, runtime_seq_lens, strict=True) + zip(block_table_rows, runtime_seq_lens_kv, strict=True) ): + if seq_len <= 0: + continue required_pages = (seq_len + page_size - 1) // page_size if any( int(page_id) < 0 or int(page_id) >= locator_capacity @@ -2406,15 +2544,14 @@ def _validate_decode_run_metadata_values( f"{request_idx} contains an invalid page ID" ) - if use_packed_q: + if packed_query: assert qo_indptr is not None - _, _, q_lengths = _read_packed_q_plan_metadata(qo_indptr) - if max(q_lengths) > seq_len_q: + runtime_max_q_len, total_q, q_lengths = _read_packed_q_plan_metadata(qo_indptr) + if runtime_max_q_len > seq_len_q: raise ValueError( "qo_indptr contains a per-request Q length larger than " - f"max_seq_len_q ({seq_len_q}): got {max(q_lengths)}" + f"max_q_len ({seq_len_q}): got {runtime_max_q_len}" ) - total_q = sum(q_lengths) if total_q != int(runtime.q.shape[0]): raise ValueError( "the final qo_indptr offset must equal the packed q token count: " @@ -2422,17 +2559,21 @@ def _validate_decode_run_metadata_values( ) else: q_lengths = (seq_len_q,) * batch_size + runtime_max_q_len = seq_len_q if mask_type == "causal": for request_idx, (q_len, kv_len) in enumerate( - zip(q_lengths, runtime_seq_lens, strict=True) + zip(q_lengths, runtime_seq_lens_kv, strict=True) ): + if kv_len <= 0: + continue if q_len > kv_len: raise ValueError( "causal decode requires every per-request Q length to be " "no greater than its K/V length; request " f"{request_idx} has Q={q_len} and K/V={kv_len}" ) + return runtime_max_q_len, max(0, max(runtime_seq_lens_kv)) @dataclass(frozen=True) @@ -2478,30 +2619,53 @@ class PrimsTSBatchDecodePlan: _output_stride: tuple[int, ...] _device: torch.device _q_dtype: torch.dtype - _output_dtype: torch.dtype + _out_dtype: torch.dtype _head_dim: int _cache: _NativePagedKVCache _k_sf_cache: torch.Tensor _v_sf_cache: torch.Tensor - _seq_lens: torch.Tensor + _seq_lens_kv: torch.Tensor _qo_indptr: Optional[torch.Tensor] _block_table: torch.Tensor _q_token_kv_block_sparse_page_memberships: torch.Tensor _workspace: _DecodeWorkspaceViews _compiled_main: Callable[..., object] _compiled_reducer: Optional[Callable[..., object]] - _direct_q1_inputs: tuple[torch.Tensor, ...] = () + _direct_q1_inputs: tuple[Optional[torch.Tensor], ...] = () _flat_locator_heads: int = 1 + _lse_mode: str = "none" + _use_attention_sinks: bool = False + _requested_splits: int = 1 + q_capacity: int = 1 + kv_capacity: int = 1 + direct_model_capacity: int = 0 + direct_page_capacity: int = 0 + direct_fixed_query_len: int = 0 def run( self, query: torch.Tensor, *, out: torch.Tensor, - bmm1_scale: Optional[float] = None, - bmm2_scale: float = 1.0, - ) -> torch.Tensor: - """Launch with lightweight checks against the validated plan.""" + softmax_scale: Optional[float] = None, + output_scale: float = 1.0, + lse: Optional[torch.Tensor] = None, + return_lse: bool = False, + attention_sinks: Optional[torch.Tensor] = None, + validate: bool = True, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Launch with lightweight checks against the validated plan. + ``lse`` is a compact float32 output of shape ``query.shape[:-1]`` + (``q.shape[:-1]`` for FMHA). Enable its base with ``plan(lse_mode=...)``. + It is allocated if omitted; ``return_lse`` controls only tuple return. + Supply output/LSE buffers and ``validate=False`` for allocation-free + capture. See the module guide for scale and buffer ownership contracts. + + ``attention_sinks`` is a compact CUDA float32 ``[Hq]`` vector of + natural-log sink logits after QK scaling, with zero V contribution. + Values must be finite or ``-inf``; this is a caller precondition. + Validation checks tensor metadata without reading the sink values. + """ if ( query.shape != self._query_shape @@ -2517,27 +2681,46 @@ def run( out.shape != self._output_shape or out.stride() != self._output_stride or out.device != self._device - or out.dtype != self._output_dtype + or out.dtype != self._out_dtype ): raise ValueError( "out must preserve the shape, strides, device, and dtype " "validated by the PrimTS plan" ) - scale_qk = _validate_scale( - 1.0 / math.sqrt(self._head_dim) if bmm1_scale is None else bmm1_scale, - "bmm1_scale", + softmax_scale = _validate_scale( + 1.0 / math.sqrt(self._head_dim) if softmax_scale is None else softmax_scale, + "softmax_scale", + ) + output_scale = _validate_scale(output_scale, "output_scale") + lse = prepare_lse( + lse, q=query, mode=self._lse_mode, return_lse=return_lse, validate=validate + ) + attention_sinks = validate_attention_sinks( + attention_sinks, + q=query, + enabled=self._use_attention_sinks, + validate=validate, + ) + result = self._run_unchecked( + query, + out, + softmax_scale, + output_scale, + lse=lse, + attention_sinks=attention_sinks, ) - scale_v = _validate_scale(bmm2_scale, "bmm2_scale") - return self._run_unchecked(query, out, scale_qk, scale_v) + return (result, lse) if return_lse else result def _run_unchecked( self, query: torch.Tensor, out: torch.Tensor, - scale_qk: float, - scale_v: float, - direct_q1_inputs: Optional[tuple[torch.Tensor, ...]] = None, + softmax_scale: float, + output_scale: float, + direct_q1_inputs: Optional[tuple[Optional[torch.Tensor], ...]] = None, + lse: Optional[torch.Tensor] = None, + attention_sinks: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Launch state already proven by a framework-owned outer plan.""" @@ -2547,21 +2730,33 @@ def _run_unchecked( direct_inputs = ( self._direct_q1_inputs if direct_q1_inputs is None else direct_q1_inputs ) + if direct_inputs: + direct_inputs = ( + *direct_inputs, + self.direct_model_capacity, + self.direct_page_capacity, + self.direct_fixed_query_len, + ) + query_view = query.view(-1, query.shape[-2], query.shape[-1]) + out_view = out.view(-1, out.shape[-2], out.shape[-1]) self._compiled_main( - query, + query_view, self._cache.k_cache, self._cache.v_cache, self._k_sf_cache, self._v_sf_cache, - out, - self._seq_lens, + out_view, + None if lse is None else lse.view(-1, lse.shape[-1]), + self._seq_lens_kv, q_offsets, self._block_table, self._q_token_kv_block_sparse_page_memberships, self._workspace.partial_o, self._workspace.partial_stats, self._workspace.split_kv_counter, - self._workspace.attention_sinks, + self._workspace.attention_sinks + if attention_sinks is None + else attention_sinks, self._cache.num_physical_pages, self._cache.k_page_stride, self._cache.k_head_stride, @@ -2569,21 +2764,30 @@ def _run_unchecked( self._cache.v_page_stride, self._cache.v_head_stride, self._cache.v_token_stride, - scale_qk, - scale_v, + softmax_scale, + output_scale, + self._requested_splits, + self.q_capacity, direct_inputs, + self.kv_capacity, ) if self._compiled_reducer is not None: self._compiled_reducer( - out, - self._seq_lens, + out_view, + None if lse is None else lse.view(-1, lse.shape[-1]), + self._seq_lens_kv, q_offsets, self._workspace.partial_o, self._workspace.partial_stats, - self._workspace.attention_sinks, - scale_qk, - scale_v, + self._workspace.attention_sinks + if attention_sinks is None + else attention_sinks, + softmax_scale, + output_scale, + self._requested_splits, + self.q_capacity, direct_inputs, + self.kv_capacity, ) return out @@ -2594,7 +2798,7 @@ def _resolve_decode_workspace_layout( num_kv_heads: int, head_dim: int, page_size: int, - max_seq_len: int, + max_kv_len: int, seq_len_q: int, q_dtype: torch.dtype, k_dtype: torch.dtype, @@ -2602,7 +2806,7 @@ def _resolve_decode_workspace_layout( out_dtype: torch.dtype, kv_layout: str, mask_type: str, - use_packed_q: bool, + packed_query: bool, window_left: int, storage_page_size: int, device: Optional[Union[int, str, torch.device]], @@ -2610,6 +2814,7 @@ def _resolve_decode_workspace_layout( use_pdl: bool = False, split_kv: bool = True, share_pattern_across_kv_heads: bool = True, + use_attention_sinks: bool = False, ) -> _DecodeWorkspaceLayout: """Resolve the byte layout for one already-validated semantic key.""" @@ -2622,7 +2827,7 @@ def _resolve_decode_workspace_layout( num_kv_heads, head_dim, page_size, - max_seq_len, + max_kv_len, seq_len_q, _dtype_key(q_dtype), _dtype_key(k_dtype), @@ -2630,13 +2835,14 @@ def _resolve_decode_workspace_layout( _dtype_key(out_dtype), kv_layout, mask_type, - use_packed_q, + packed_query, window_left, storage_page_size, use_q_token_kv_block_sparse_route, use_pdl, split_kv, share_pattern_across_kv_heads, + use_attention_sinks, ) return _make_decode_workspace_layout( spec.scratch_shapes, @@ -2658,6 +2864,7 @@ def _normalize_paged_kv_cache_views( paged_kv_cache: PagedKVCache, *, expected_device: torch.device, + kv_layout: str = "HND", ) -> tuple[torch.Tensor, torch.Tensor, int, int, int, int]: """Return validated zero-copy HND views and the logical head dimension.""" @@ -2695,7 +2902,10 @@ def _normalize_paged_kv_cache_views( ) if k_cache.ndim != 4 or v_cache.ndim != 4: - raise ValueError("K/V cache views must be rank-4 HND tensors") + raise ValueError("K/V cache views must be rank-4 tensors") + if kv_layout == "NHD": + k_cache = k_cache.transpose(1, 2) + v_cache = v_cache.transpose(1, 2) if k_cache.shape != v_cache.shape: raise ValueError("K and V cache views must have identical logical shapes") if k_cache.device != v_cache.device: @@ -2775,6 +2985,7 @@ def _normalize_native_paged_kv_cache( paged_kv_cache: PagedKVCache, *, expected_device: torch.device, + kv_layout: str = "HND", ) -> _NativePagedKVCache: """Return native decode views supporting HND- or NHD-physical storage.""" @@ -2788,6 +2999,7 @@ def _normalize_native_paged_kv_cache( ) = _normalize_paged_kv_cache_views( paged_kv_cache, expected_device=expected_device, + kv_layout=kv_layout, ) k_page_stride, k_head_stride, k_token_stride = _validate_native_hnd_tma_strides( k_cache, "K cache" @@ -2864,9 +3076,9 @@ def _resolve_q_token_kv_block_sparse_decode_config( q_dtype_key: str, k_dtype_key: str, v_dtype_key: str, - output_dtype_key: str, + out_dtype_key: str, mask_type: str, - use_packed_q: bool, + packed_query: bool, window_left: int, max_active_clusters: int, use_pdl: bool, @@ -2889,11 +3101,11 @@ def _resolve_q_token_kv_block_sparse_decode_config( q_token_kv_block_sparse_dtype_supported = ( q_dtype_key == k_dtype_key == v_dtype_key and ( - q_dtype_key == output_dtype_key + q_dtype_key == out_dtype_key and q_dtype_key in ("float16", "bfloat16") or ( q_dtype_key == "float8_e4m3fn" - and output_dtype_key in ("float16", "bfloat16") + and out_dtype_key in ("float16", "bfloat16") ) ) ) @@ -2905,7 +3117,7 @@ def _resolve_q_token_kv_block_sparse_decode_config( and window_left < 0 ): raise ValueError( - "PrimTS QToken-KvBlock-Sparse-Attention requires a supported sparse-fragment size, head_dim in {64,128,256}, " + "PrimTS QToken-KvBlock-Sparse-Attention requires a supported sparse-fragment size, head_dim in {64,128,192,256}, " "FP16/BF16 Q/K/V with matching output or FP8 Q/K/V with FP16/BF16 output, " "and a causal non-windowed mask" ) @@ -2923,7 +3135,7 @@ def _resolve_q_token_kv_block_sparse_decode_config( ) q_token_kv_block_sparse_profile = { - "use_variable_seqlens_q": use_packed_q, + "use_variable_seqlens_q": packed_query, "use_q_token_kv_block_sparse_route": True, "share_pattern_across_kv_heads": share_pattern_across_kv_heads, "use_pdl": use_pdl, @@ -3015,7 +3227,7 @@ def suggest_q_token_kv_block_sparse_group_size( multi_processor_count: int, *, head_dim: int = 256, - q_data_type: torch.dtype = torch.bfloat16, + q_dtype: torch.dtype = torch.bfloat16, split_kv: bool = True, kv_block_size: int = 4, ) -> int: @@ -3060,7 +3272,7 @@ def suggest_q_token_kv_block_sparse_group_size( and pass the cached value; this function performs no device query. head_dim : int Common Q/K/V head dimension: 64, 128, or 256. - q_data_type : torch.dtype + q_dtype : torch.dtype Q/K/V dtype used to select the matching dense MMA recipe. split_kv : bool Allow automatic split-KV fanout when True (default); False forces @@ -3094,8 +3306,8 @@ def suggest_q_token_kv_block_sparse_group_size( heads_q_per_kv = num_qo_heads // num_kv_heads head_dim = _validate_head_dim(head_dim) - if q_data_type not in _SUPPORTED_INPUT_DTYPES: - raise ValueError("q_data_type must be float16, bfloat16, or float8_e4m3fn") + if q_dtype not in _SUPPORTED_INPUT_DTYPES: + raise ValueError("q_dtype must be float16, bfloat16, or float8_e4m3fn") if not isinstance(split_kv, bool): raise TypeError("split_kv must be a bool") return _sparse_policy.suggest_sparse_group( @@ -3106,7 +3318,7 @@ def suggest_q_token_kv_block_sparse_group_size( num_kv_heads=num_kv_heads, multi_processor_count=multi_processor_count, head_dim=head_dim, - q_dtype_key=_dtype_key(q_data_type), + q_dtype_key=_dtype_key(q_dtype), split_kv=split_kv, block_size=kv_block_size, ) @@ -3344,17 +3556,17 @@ def _prepare_prims_ts_batch_decode_plan( kv_cache: PagedKVCache, workspace_buffer: torch.Tensor, block_table: torch.Tensor, - seq_lens: torch.Tensor, - max_seq_len: int, + seq_lens_kv: torch.Tensor, + max_kv_len: int, *, seq_len_q: int, qo_indptr: Optional[torch.Tensor], - max_seq_len_q: Optional[int], + max_q_len: Optional[int], out: Optional[torch.Tensor], out_dtype: Optional[torch.dtype], mask_type: Literal["dense", "causal"], window_left: int, - kv_layout: Literal["HND"], + kv_layout: Literal["HND", "NHD"], page_size: Optional[int], kv_scale_factors: Optional[PagedKVScaleFactors] = None, q_token_kv_block_sparse_page_memberships: Optional[torch.Tensor] = None, @@ -3362,31 +3574,36 @@ def _prepare_prims_ts_batch_decode_plan( use_q_token_kv_block_sparse_route: bool = False, use_pdl: bool = False, split_kv: bool = True, - direct_q1_inputs: tuple[torch.Tensor, ...] = (), + direct_q1_inputs: tuple[Optional[torch.Tensor], ...] = (), direct_q1_max_model_len: Optional[int] = None, direct_q1_sparse_block_size: int = 4, + direct_q1_fixed_query_len: int = 0, + count_current_block_in_topk: bool = True, share_pattern_across_kv_heads: bool = True, allow_head_aware_locators: bool = False, + use_attention_sinks: bool = False, + lse_mode: Literal["none", "base2"] = "none", ) -> tuple[PrimsTSBatchDecodePlan, torch.Tensor]: """Validate and freeze one dense-block-table PrimTS launch contract.""" + lse_mode = validate_lse_mode(lse_mode) _validate_layout(kv_layout) _validate_mask(mask_type) window_left = _validate_window_left(window_left, mask_type) - use_packed_q, resolved_seq_len_q = _resolve_q_mode( + packed_query, resolved_seq_len_q = _resolve_q_mode( seq_len_q=seq_len_q, qo_indptr=qo_indptr, - max_seq_len_q=max_seq_len_q, + max_q_len=max_q_len, require_packed_max=True, ) assert resolved_seq_len_q is not None seq_len_q = resolved_seq_len_q metadata_device, metadata_rows, max_num_pages = _validate_block_table_metadata( block_table, - seq_lens, + seq_lens_kv, ) normalized_cache = _normalize_native_paged_kv_cache( - kv_cache, expected_device=query.device + kv_cache, expected_device=query.device, kv_layout=kv_layout ) pattern_heads = _num_sparse_pattern_heads( normalized_cache.num_kv_heads, share_pattern_across_kv_heads @@ -3397,7 +3614,7 @@ def _prepare_prims_ts_batch_decode_plan( _validate_q( query, seq_len_q=seq_len_q, - use_packed_q=use_packed_q, + packed_query=packed_query, device=metadata_device, batch_size=batch_size, ) @@ -3423,25 +3640,25 @@ def _prepare_prims_ts_batch_decode_plan( storage_page_size if page_size is None else page_size ) storage_page_size = _validate_storage_page_size(page_size, storage_page_size) - max_seq_len = _validate_max_kv_len(max_seq_len, "max_seq_len") - required_page_columns = (max_seq_len + page_size - 1) // page_size + max_kv_len = _validate_max_kv_len(max_kv_len, "max_kv_len") + required_page_columns = (max_kv_len + page_size - 1) // page_size if max_num_pages < required_page_columns: raise ValueError( - "block_table must have at least ceil(max_seq_len / page_size) " + "block_table must have at least ceil(max_kv_len / page_size) " f"columns ({required_page_columns}), got {max_num_pages}" ) - output_dtype = out_dtype - if output_dtype is None: + out_dtype = out_dtype + if out_dtype is None: if out is not None and not isinstance(out, torch.Tensor): raise TypeError("out must be a torch.Tensor") - output_dtype = out.dtype if out is not None else query.dtype - elif not isinstance(output_dtype, torch.dtype): + out_dtype = out.dtype if out is not None else query.dtype + elif not isinstance(out_dtype, torch.dtype): raise TypeError("out_dtype must be a torch.dtype") _validate_dtype_pair( query.dtype, k_cache.dtype, v_cache.dtype, - output_dtype, + out_dtype, ) k_sf_cache, v_sf_cache = _normalize_paged_kv_scale_factors( kv_scale_factors, @@ -3458,15 +3675,15 @@ def _prepare_prims_ts_batch_decode_plan( num_kv_heads, head_dim, page_size, - max_seq_len, + max_kv_len, seq_len_q, _dtype_key(query.dtype), _dtype_key(k_cache.dtype), _dtype_key(v_cache.dtype), - _dtype_key(output_dtype), + _dtype_key(out_dtype), kv_layout, mask_type, - use_packed_q, + packed_query, window_left, storage_page_size, use_q_token_kv_block_sparse_route, @@ -3474,14 +3691,18 @@ def _prepare_prims_ts_batch_decode_plan( split_kv, share_pattern_across_kv_heads, ) - spec = _resolve_decode_launch_spec(*policy_args) + spec = _resolve_decode_launch_spec( + *policy_args, + use_attention_sinks=use_attention_sinks, + ) + policy_config = replace(spec.config, fixed_seqlen=True) if query_major_memberships: - if not spec.config.supports_query_major_memberships: + if not policy_config.supports_query_major_memberships: raise ValueError( "query-major memberships require one G2..G8 Keeps/page-4 CTA" ) spec = replace(spec, config=replace(spec.config, query_major_memberships=True)) - if spec.config.uses_q_token_kv_block_sparse_page_membership: + if policy_config.uses_q_token_kv_block_sparse_page_membership: if q_token_kv_block_sparse_page_memberships is None: raise ValueError( "grouped QToken-KvBlock-Sparse-Attention requires a separate q_token_kv_block_sparse_page_memberships tensor" @@ -3502,7 +3723,7 @@ def _prepare_prims_ts_batch_decode_plan( q_token_kv_block_sparse_page_memberships = block_table layout = _make_decode_workspace_layout( spec.scratch_shapes, - output_dtype, + out_dtype, use_separate_reduction_kernel=spec.config.use_separate_reduction_kernel, use_split_kv=spec.config.use_split_kv, ) @@ -3516,29 +3737,32 @@ def _prepare_prims_ts_batch_decode_plan( num_qo_heads=num_qo_heads, seq_len_q=seq_len_q, head_dim=head_dim, - total_q_tokens=int(query.shape[0]) if use_packed_q else None, + total_q_tokens=int(query.shape[0]) if packed_query else None, ) if out is None: - out = torch.empty(output_shape, device=query.device, dtype=output_dtype) + out = torch.empty(output_shape, device=query.device, dtype=out_dtype) else: _validate_out( out, q=query, expected_shape=output_shape, seq_len_q=seq_len_q, - use_packed_q=use_packed_q, - output_dtype=output_dtype, + packed_query=packed_query, + out_dtype=out_dtype, ) direct_q1_spec = None if direct_q1_inputs: assert use_q_token_kv_block_sparse_route and seq_len_q == 1 assert direct_q1_max_model_len is not None + assert direct_q1_inputs[0] is not None and direct_q1_inputs[3] is not None direct_q1_spec = _DirectQ1CompileSpec( position_dtype_key=str(direct_q1_inputs[3].dtype).removeprefix("torch."), max_seq_len_kv=direct_q1_max_model_len, block_topk=int(direct_q1_inputs[0].shape[-1]), kv_block_size=direct_q1_sparse_block_size, + count_current_block_in_topk=count_current_block_in_topk, page_capacity=max_num_pages, + fixed_query_len=direct_q1_fixed_query_len, ) flat_native_kv_tma = ( use_q_token_kv_block_sparse_route @@ -3557,17 +3781,13 @@ def _prepare_prims_ts_batch_decode_plan( num_kv_heads=num_kv_heads, head_dim=head_dim, page_size=page_size, - max_kv_len=max_seq_len, - seq_len_q=seq_len_q, q_dtype_key=_dtype_key(query.dtype), k_dtype_key=_dtype_key(k_cache.dtype), v_dtype_key=_dtype_key(v_cache.dtype), - output_dtype_key=_dtype_key(output_dtype), - use_packed_q=use_packed_q, - kv_prefix_mode="dynamic", - kv_lengths_mode="dynamic", + out_dtype_key=_dtype_key(out_dtype), direct_q1_spec=direct_q1_spec, flat_native_kv_tma=flat_native_kv_tma, + lse_mode=lse_mode, ) compiled_main, compiled_reducer = _get_compiled_decode(compile_spec) workspace = _bind_decode_workspace(workspace_buffer, layout) @@ -3585,13 +3805,18 @@ def _prepare_prims_ts_batch_decode_plan( _output_stride=tuple(out.stride()), _device=query.device, _q_dtype=query.dtype, - _output_dtype=output_dtype, + _out_dtype=out_dtype, _head_dim=head_dim, _cache=normalized_cache, _k_sf_cache=k_sf_cache, _v_sf_cache=v_sf_cache, - _seq_lens=seq_lens, - _qo_indptr=qo_indptr, + _seq_lens_kv=seq_lens_kv, + _qo_indptr=( + qo_indptr + if qo_indptr is not None + else torch.arange(batch_size + 1, dtype=torch.int32, device=query.device) + * seq_len_q + ), _block_table=block_table, _q_token_kv_block_sparse_page_memberships=q_token_kv_block_sparse_page_memberships, _workspace=workspace, @@ -3599,6 +3824,14 @@ def _prepare_prims_ts_batch_decode_plan( _compiled_reducer=compiled_reducer, _direct_q1_inputs=direct_q1_inputs, _flat_locator_heads=num_kv_heads if flat_native_kv_tma else 1, + _lse_mode=lse_mode, + _use_attention_sinks=use_attention_sinks, + _requested_splits=spec.requested_splits or int(spec.config.splits_kv), + q_capacity=seq_len_q, + kv_capacity=max_kv_len, + direct_model_capacity=direct_q1_max_model_len or 0, + direct_page_capacity=max_num_pages, + direct_fixed_query_len=direct_q1_fixed_query_len, ) return plan, out @@ -3609,20 +3842,22 @@ def prepare_prims_ts_batch_decode_with_kv_cache( kv_cache: PagedKVCache, workspace_buffer: torch.Tensor, block_table: torch.Tensor, - seq_lens: torch.Tensor, - max_seq_len: int, + seq_lens_kv: torch.Tensor, + max_kv_len: int, *, out: torch.Tensor, kv_scale_factors: Optional[PagedKVScaleFactors] = None, seq_len_q: int = 1, qo_indptr: Optional[torch.Tensor] = None, - max_seq_len_q: Optional[int] = None, + max_q_len: Optional[int] = None, out_dtype: Optional[torch.dtype] = None, mask_type: Literal["dense", "causal"] = "dense", window_left: int = -1, - kv_layout: Literal["HND"] = "HND", + kv_layout: Literal["HND", "NHD"] = "HND", page_size: Optional[int] = None, split_kv: bool = True, + use_attention_sinks: bool = False, + lse_mode: Literal["none", "base2"] = "none", ) -> PrimsTSBatchDecodePlan: """Validate and prepare a reusable dense-block-table PrimTS launch. @@ -3646,9 +3881,9 @@ def prepare_prims_ts_batch_decode_with_kv_cache( for this semantic configuration. block_table : torch.Tensor Contiguous CUDA Int32 dense page table shaped ``[B, max_pages]``. - seq_lens : torch.Tensor + seq_lens_kv : torch.Tensor Positive live K/V sequence lengths, one per request. - max_seq_len : int + max_kv_len : int Static maximum K/V length used for policy selection and JIT caching. out : torch.Tensor Caller-owned output tensor whose storage remains stable across replays. @@ -3660,7 +3895,7 @@ def prepare_prims_ts_batch_decode_with_kv_cache( Fixed query length when ``qo_indptr`` is omitted. qo_indptr : torch.Tensor, optional Cumulative Int32 query offsets selecting packed-query mode. - max_seq_len_q : int, optional + max_q_len : int, optional Static maximum packed-query route length. out_dtype : torch.dtype, optional Output dtype; defaults to ``out.dtype``. @@ -3668,7 +3903,7 @@ def prepare_prims_ts_batch_decode_with_kv_cache( Attention mask mode. window_left : int Left sliding-window extent, or ``-1`` to disable the window. - kv_layout : {"HND"} + kv_layout : {"HND", "NHD"} Layout of the paged K/V cache. page_size : int, optional Semantic page size; defaults to the physical cache-page extent. @@ -3681,6 +3916,14 @@ def prepare_prims_ts_batch_decode_with_kv_cache( PrimsTSBatchDecodePlan Reusable prepared launch plan. Call :meth:`PrimsTSBatchDecodePlan.run` after updating values in the prepared tensor storage. + + ``lse_mode="none"`` (default) keeps the output-only contract. + Enable ``"base2"`` while planning to write final base-2 LSE. + Writing LSE is a compilation feature; the output pointer stays runtime. + + ``use_attention_sinks=True`` enables the per-run ``attention_sinks`` + binding; supply it on every run. Defaults to ``False``. Sizing caller + workspace must use the same sink capability as planning. """ if not isinstance(out, torch.Tensor): @@ -3690,11 +3933,11 @@ def prepare_prims_ts_batch_decode_with_kv_cache( kv_cache, workspace_buffer, block_table, - seq_lens, - max_seq_len, + seq_lens_kv, + max_kv_len, seq_len_q=seq_len_q, qo_indptr=qo_indptr, - max_seq_len_q=max_seq_len_q, + max_q_len=max_q_len, out=out, out_dtype=out_dtype, mask_type=mask_type, @@ -3704,31 +3947,28 @@ def prepare_prims_ts_batch_decode_with_kv_cache( kv_scale_factors=kv_scale_factors, split_kv=split_kv, use_q_token_kv_block_sparse_route=False, + lse_mode=lse_mode, + use_attention_sinks=use_attention_sinks, ) return plan class BatchDecodePagedTSWrapper: - """Plan static paged-decode capacity and bind request metadata. + """Reusable paged decode with eager planning and shared scratch. - A plan fixes device, batch size, geometry, dtypes, query storage mode, and - maximum K/V capacity. Sequence lengths are supplied either once to - :meth:`plan`, which retains an owned CUDA copy, or to every :meth:`run`. - Every run supplies a row-strided page table; packed-query plans also consume - per-run Q offsets. - - One wrapper supports one ordered execution lane. Concurrent streams or - graph replays require separate wrappers and workspace buffers. - """ + Plan selects a cached execution from problem bounds and binds scratch. + Each run supplies query, lengths, page tables and optional packed-Q offsets. + Policies/compiled kernels are cached across layers; resources remain local. + Concurrent streams or graph replays require separate wrappers/workspaces.""" @flashinfer_experimental_api - def __init__(self, kv_layout: Literal["HND"] = "HND") -> None: + def __init__(self, kv_layout: Literal["HND", "NHD"] = "HND") -> None: """Initialize an unplanned wrapper with one static K/V layout. Parameters ---------- - kv_layout : {"HND"} - Layout of the paged K/V cache. Only ``"HND"`` is supported. + kv_layout : {"HND", "NHD"} + Physical K/V layout, ``"HND"`` or ``"NHD"``. """ _validate_layout(kv_layout) @@ -3758,122 +3998,52 @@ def plan( page_size: int, max_kv_len: int, *, - max_seq_len_q: int = 1, + max_q_len: int | None = None, packed_query: bool = False, - q_data_type: torch.dtype = torch.float16, - k_data_type: Optional[torch.dtype] = None, - v_data_type: Optional[torch.dtype] = None, - kv_data_type: Optional[torch.dtype] = None, - o_data_type: Optional[torch.dtype] = None, + q_dtype: torch.dtype = torch.float16, + k_dtype: Optional[torch.dtype] = None, + v_dtype: Optional[torch.dtype] = None, + kv_dtype: Optional[torch.dtype] = None, + out_dtype: Optional[torch.dtype] = None, mask_type: Literal["dense", "causal"] = "dense", + lse_mode: Literal["none", "base2"] = "none", + use_attention_sinks: bool = False, window_left: int = -1, - seq_lens: Optional[Union[Sequence[int], torch.Tensor]] = None, + seq_lens_kv: Optional[Union[Sequence[int], torch.Tensor]] = None, workspace_buffer: Optional[torch.Tensor] = None, storage_page_size: Optional[int] = None, split_kv: bool = True, validate: bool = True, initialize_workspace: bool = True, ) -> None: - """Compile one static-capacity plan and optionally own sequence lengths. - - ``max_seq_len_q`` is the exact fixed Q length when ``packed_query`` is - false and the per-request capacity when ``packed_query`` is true. A - packed run supplies ``qo_indptr`` and compact ``[total_q, Hq, D]`` - storage. - - ``seq_lens`` optionally fixes the per-request lengths for the lifetime - of this plan. Planning validates the host values, retains them, and - creates an owned CUDA copy used by every launch. In this mode, - :meth:`run` requires its ``seq_lens`` argument to be ``None``. When - omitted, both K/V prefix and length handling compile dynamically and - every run must supply a CUDA length tensor. - - ``workspace_buffer`` is caller-owned scratch for this plan. It is - allocated when omitted, retained by the frozen plan state, and never - reset by ``run``. By default planning initializes its control sections; - ``initialize_workspace=False`` preserves a caller-initialized buffer. - - Parameters - ---------- - device : int, str, or torch.device - CUDA device on which the specialization is compiled and run. - batch_size : int - Exact number of requests in every run. - num_qo_heads : int - Number of query/output heads. - num_kv_heads : int - Number of K/V heads. - head_dim : int - Query, key, value, and output head dimension. - page_size : int - Number of K/V tokens addressed by each page-table entry. - storage_page_size : int, optional - Physical cache-page extent, defaulting to ``page_size``. Larger - physical pages must be divisible by the logical ``page_size``. - max_kv_len : int - Per-request K/V length capacity used for policy selection, - compilation, and workspace sizing. - max_seq_len_q : int - Exact fixed Q length, or per-request capacity for packed Q. - Defaults to ``1``. - packed_query : bool - Select compact ``[total_q, Hq, D]`` query/output storage instead - of fixed ``[B, SQ, Hq, D]`` storage. Fixed SQ1 storage is - ``[B, Hq, D]``. Defaults to ``False``. - q_data_type : torch.dtype - Query dtype used to compile the plan. Defaults to - ``torch.float16``. - k_data_type : torch.dtype, optional - K dtype used to compile the plan. Defaults to ``q_data_type``. - v_data_type : torch.dtype, optional - V dtype used to compile the plan. Defaults to ``k_data_type``; - ``torch.float8_e4m3fn`` with BF16 Q/K selects the QK-BF16/PV-FP8 - path and requires separate ``(K, V)`` cache tensors. - kv_data_type : torch.dtype, optional - Compatibility alias setting both K and V storage dtypes. Explicit - K/V dtypes must agree with it. Use ``torch.uint8`` for packed - NVFP4; ``head_dim`` stays logical and runs supply scale tensors. - o_data_type : torch.dtype, optional - Output dtype used to compile the plan. Defaults to - ``q_data_type``. - mask_type : {"dense", "causal"} - Attention mask mode. Defaults to ``"dense"``. - window_left : int - Left sliding-window extent, or ``-1`` to disable the window. A - non-negative value requires causal masking. - seq_lens : Sequence[int] or torch.Tensor, optional - Host-only per-request K/V lengths owned by the resulting plan and - used for specialization. A tensor must be a one-dimensional CPU - int32 or int64 tensor. The sequence must contain exactly - ``batch_size`` positive values no larger than ``max_kv_len``. - workspace_buffer : torch.Tensor, optional - Caller-owned contiguous int8 or uint8 scratch on ``device``. It - must be 32-byte aligned and large enough for the selected plan. - When omitted, planning allocates the buffer. The retained buffer - is exclusive to one in-flight launch or graph replay. - split_kv : bool - Permit automatic useful split fanout (True, default), or force - unsplit execution (False), independently of ``packed_query``. - Typically False for prefill and True for decode. This choice is - frozen for the lifetime of the plan. - validate : bool - Validate static geometry and caller scratch. Defaults to ``True``. - Disable only for previously validated inputs and warmed topology. - initialize_workspace : bool - Initialize control sections during planning, default ``True``. - Set to ``False`` only for caller-initialized scratch with unchanged - layout. Newly allocated scratch is always initialized. - """ - - k_data_type, v_data_type = _resolve_kv_dtypes( - q_data_type, k_data_type, v_data_type, kv_data_type - ) - if o_data_type is None: - o_data_type = q_data_type - seq_len_q = max_seq_len_q + """Select/cache a bounded paged-decode execution and bind reusable scratch. + + batch_size is exact. max_q_len is the fixed query length, or an + explicit per-request bound for packed queries. Live K/V lengths remain + within max_kv_len. Exact Q/KV bounds remain runtime operands when topology is unchanged. seq_lens_kv optionally supplies host lengths copied into + plan-owned CUDA storage; otherwise every run supplies its own lengths. + Kernel selection and compilation/loading finish here, before capture. + Replanning replaces state and invalidates its prior captured graphs. + Scratch is sized for the selected policy and problem geometry. Caller + scratch may be passed without a size query. Planning initializes control + storage unless initialize_workspace=False transfers that to the caller. + Active split counts are automatic runtime operands for main/reduction + kernels; users never supply the count.""" + + if packed_query and max_q_len is None: + raise ValueError("packed Q requires max_q_len in plan()") + if max_q_len is None: + max_q_len = 1 + k_dtype, v_dtype = _resolve_kv_dtypes(q_dtype, k_dtype, v_dtype, kv_dtype) + if out_dtype is None: + out_dtype = q_dtype + seq_len_q = max_q_len storage_page_size = ( page_size if storage_page_size is None else storage_page_size ) + lse_mode = validate_lse_mode(lse_mode) + if not isinstance(use_attention_sinks, bool): + raise TypeError("use_attention_sinks must be a bool") if not isinstance(validate, bool): raise TypeError("validate must be a bool") if not isinstance(initialize_workspace, bool): @@ -3888,30 +4058,35 @@ def plan( page_size, storage_page_size ) max_kv_len = _validate_max_kv_len(max_kv_len, "max_kv_len") - seq_len_q = _validate_seq_len_q(max_seq_len_q) + seq_len_q = _validate_seq_len_q(max_q_len) _validate_head_geometry(num_qo_heads, num_kv_heads) _validate_decode_query_head_extent( batch_size=batch_size, num_qo_heads=num_qo_heads, - max_seq_len_q=seq_len_q, + max_q_len=seq_len_q, ) _validate_mask(mask_type) window_left = _validate_window_left(window_left, mask_type) - _validate_dtype_pair(q_data_type, k_data_type, v_data_type, o_data_type) + _validate_dtype_pair( + q_dtype, + k_dtype, + v_dtype, + out_dtype, + ) - specialization_seq_lens = _normalize_plan_seq_lens( - seq_lens, + planned_seq_lens = _normalize_plan_seq_lens_kv( + seq_lens_kv, batch_size=batch_size, max_kv_len=max_kv_len, ) if ( - specialization_seq_lens is not None + planned_seq_lens is not None and mask_type == "causal" and not packed_query ): - for request_idx, kv_len in enumerate(specialization_seq_lens): - if seq_len_q > kv_len: + for request_idx, kv_len in enumerate(planned_seq_lens): + if kv_len > 0 and seq_len_q > kv_len: raise ValueError( "causal decode requires every per-request Q length to be " "no greater than its K/V length; request " @@ -3919,8 +4094,8 @@ def plan( ) else: - specialization_seq_lens = _normalize_plan_seq_lens( - seq_lens, batch_size=batch_size, max_kv_len=max_kv_len + planned_seq_lens = _normalize_plan_seq_lens_kv( + seq_lens_kv, batch_size=batch_size, max_kv_len=max_kv_len ) device, device_index = _resolve_cuda_device(device) @@ -3936,49 +4111,70 @@ def plan( page_size, max_kv_len, seq_len_q, - _dtype_key(q_data_type), - _dtype_key(k_data_type), - _dtype_key(v_data_type), - _dtype_key(o_data_type), + _dtype_key(q_dtype), + _dtype_key(k_dtype), + _dtype_key(v_dtype), + _dtype_key(out_dtype), self._kv_layout, mask_type, packed_query, window_left, storage_page_size, ) - spec = _resolve_decode_launch_spec(*policy_args, split_kv=split_kv) - if specialization_seq_lens is None: - kv_prefix_mode: Literal["dynamic", "planned_full"] = "dynamic" - kv_lengths_mode: Literal["dynamic", "planned_uniform_max"] = "dynamic" - else: - static_full_split_prefix = _planned_full_split_prefix( - spec.config, - specialization_seq_lens, - seq_len_q=seq_len_q, - max_kv_len=max_kv_len, - mask_type=mask_type, - ) - kv_prefix_mode = "planned_full" if static_full_split_prefix else "dynamic" - has_unpaired_kv_tail = _planned_kv_domain_has_unpaired_tail( - spec.config, - max_kv_len, + spec = _resolve_decode_launch_spec( + *policy_args, split_kv=split_kv, use_attention_sinks=use_attention_sinks + ) + kv_prefix_mode: Literal["dynamic", "planned_full"] = "dynamic" + kv_lengths_mode: Literal["dynamic", "planned_uniform_max"] = "dynamic" + cfg = spec.config + # Reuse declared uniform bounds. Single-Q GMEM keeps its runtime page window. + uniform_kv = ( + not packed_query + and not cfg.use_sliding_window_causal + # Keep FP8 Keeps multi-Q arithmetic identical to standalone calls, + # whose KV bounds are runtime metadata rather than static lengths. + and not ( + cfg.use_keeps_mma_ab + and cfg.use_fp8_qkv + and head_dim == 128 + and seq_len_q > 1 + and cfg.o_dtype_bytes == 2 ) - requires_runtime_kv_lengths = ( - has_unpaired_kv_tail - or spec.config.use_sliding_window_causal - or ( - spec.config.use_persistent_scheduler - and spec.config.uses_runtime_q_kv_union - ) + and planned_seq_lens is not None + and all(length == max_kv_len for length in planned_seq_lens) + and ((max_kv_len + cfg.tile_size_kv - 1) // cfg.tile_size_kv) + % cfg.num_insts_kv + == 0 + ) + fixed_uniform_kv = uniform_kv and ( + cfg.use_cluster_smem_reduction or not cfg.has_single_q_cta + ) + if uniform_kv: + from .kernels.fmha_decode.fmha_decode_config import ( + compute_runtime_active_splits_kv, ) - kv_lengths_mode = ( - "dynamic" - if requires_runtime_kv_lengths - else _planned_kv_lengths_mode( - specialization_seq_lens, - max_kv_len=max_kv_len, + + q_step = cfg.q_tokens_per_cta if cfg.uses_nontrivial_grouped_q_layout else 1 + if cfg.use_split_kv and all( + compute_runtime_active_splits_kv( + valid_k=(max_kv_len - seq_len_q + min(q_base + q_step, seq_len_q)) + if mask_type == "causal" + else max_kv_len, + tile_size_kv=cfg.tile_size_kv, + num_insts_kv=cfg.num_insts_kv, + configured_splits_kv=spec.requested_splits, ) - ) + == spec.requested_splits + for q_base in range(0, seq_len_q, q_step) + ): + kv_prefix_mode = "planned_full" + fixed_uniform_kv = True + if fixed_uniform_kv: + kv_lengths_mode = "planned_uniform_max" + if fixed_uniform_kv and spec.requested_splits == cfg.splits_kv: + # The live bound equals this declared topology; partial activity + # is still derived from the owned lengths, as before runtime buckets. + spec = replace(spec, config=replace(cfg, runtime_split_kv=False)) compile_spec = _make_decode_compile_spec( spec, device_index=device_index, @@ -3986,25 +4182,26 @@ def plan( num_kv_heads=num_kv_heads, head_dim=head_dim, page_size=page_size, - max_kv_len=max_kv_len, - seq_len_q=seq_len_q, - q_dtype_key=_dtype_key(q_data_type), - k_dtype_key=_dtype_key(k_data_type), - v_dtype_key=_dtype_key(v_data_type), - output_dtype_key=_dtype_key(o_data_type), - use_packed_q=packed_query, + q_dtype_key=_dtype_key(q_dtype), + k_dtype_key=_dtype_key(k_dtype), + v_dtype_key=_dtype_key(v_dtype), + out_dtype_key=_dtype_key(out_dtype), + packed_query=packed_query, kv_prefix_mode=kv_prefix_mode, kv_lengths_mode=kv_lengths_mode, + lse_mode=lse_mode, + runtime_kv_lengths=not fixed_uniform_kv + and (cfg.runtime_split_kv or not cfg.use_split_kv), ) compiled_main, compiled_reducer = _get_compiled_decode(compile_spec) policy = spec.policy + ( - ("kv_prefix_mode", kv_prefix_mode), - ("kv_lengths_mode", kv_lengths_mode), + ("kv_prefix_mode", compile_spec.kv_prefix_mode), + ("kv_lengths_mode", compile_spec.kv_lengths_mode), ("split_kv_allowed", split_kv), ) workspace_layout = _make_decode_workspace_layout( spec.scratch_shapes, - o_data_type, + out_dtype, use_separate_reduction_kernel=(spec.config.use_separate_reduction_kernel), use_split_kv=spec.config.use_split_kv, ) @@ -4028,11 +4225,11 @@ def plan( workspace.attention_sinks.zero_() # Materialize from the normalized tuple so the plan never aliases # caller-owned host or device storage. - planned_seq_lens_device = ( + planned_seq_lens_kv_device = ( None - if specialization_seq_lens is None + if planned_seq_lens is None else torch.tensor( - specialization_seq_lens, + planned_seq_lens, dtype=torch.int32, device=device, ) @@ -4043,16 +4240,16 @@ def plan( device_index=device_index, batch_size=batch_size, seq_len_q=seq_len_q, - use_packed_q=packed_query, + packed_query=packed_query, num_qo_heads=num_qo_heads, num_kv_heads=num_kv_heads, head_dim=head_dim, page_size=page_size, max_kv_len=max_kv_len, - q_dtype=q_data_type, - k_dtype=k_data_type, - v_dtype=v_data_type, - output_dtype=o_data_type, + q_dtype=q_dtype, + k_dtype=k_dtype, + v_dtype=v_dtype, + out_dtype=out_dtype, mask_type=mask_type, window_left=window_left, config=spec.config, @@ -4061,12 +4258,21 @@ def plan( workspace=workspace, compiled_main=compiled_main, compiled_reducer=compiled_reducer, - kv_prefix_mode=kv_prefix_mode, - kv_lengths_mode=kv_lengths_mode, - planned_seq_lens_host=specialization_seq_lens, - planned_seq_lens_device=planned_seq_lens_device, + kv_prefix_mode=compile_spec.kv_prefix_mode, + kv_lengths_mode=compile_spec.kv_lengths_mode, + lse_mode=lse_mode, + planned_seq_lens_kv_host=planned_seq_lens, + planned_seq_lens_kv_device=planned_seq_lens_kv_device, policy=policy, storage_page_size=storage_page_size, + use_attention_sinks=use_attention_sinks, + requested_splits=spec.requested_splits or int(spec.config.splits_kv), + fixed_q_offsets=( + None + if packed_query + else torch.arange(batch_size + 1, device=device, dtype=torch.int32) + * seq_len_q + ), ) # This is the only wrapper mutation. Any failure above leaves the # previous complete plan revision usable. @@ -4077,21 +4283,47 @@ def run( self, q: torch.Tensor, paged_kv_cache: PagedKVCache, - seq_lens: Optional[torch.Tensor], + seq_lens_kv: Optional[torch.Tensor], block_tables: torch.Tensor, *, kv_scale_factors: Optional[PagedKVScaleFactors] = None, qo_indptr: Optional[torch.Tensor] = None, - bmm1_scale: Optional[float] = None, - bmm2_scale: float = 1.0, + softmax_scale: Optional[float] = None, + output_scale: float = 1.0, + max_q_len: int | None = None, + max_kv_len: int | None = None, out: Optional[torch.Tensor] = None, + lse: Optional[torch.Tensor] = None, + return_lse: bool = False, + attention_sinks: Optional[torch.Tensor] = None, validate: bool = True, - ) -> torch.Tensor: - """Launch the current plan with one sequence-length owner. - - When :meth:`plan` received sequence lengths, ``seq_lens`` must be - ``None`` and the plan's owned CUDA copy is used. Otherwise, ``seq_lens`` - must be supplied on every run. This ownership check is unconditional. + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Launch paged FMHA using live lengths and request metadata. + + Uses the policy and compiled kernels prepared by plan(). Batch and query + shape must match that plan; live K/V lengths stay within its capacity. + Supply output/LSE buffers and validate=False during graph capture, and + retain the wrapper and captured buffers for every replay. + + An enabled lse_mode="base2" plan writes compact FP32 log2(sum(exp(scores))) + with the query shape minus its final dimension. return_lse controls tuple + return only; an omitted enabled buffer is allocated. Output/LSE/workspace must + be disjoint from inputs and live buffers. One wrapper serves one ordered lane. + ``max_q_len`` and ``max_kv_len`` are runtime launch operands. + Bounds must cover the live request lengths without exceeding plan + capacity. Validation infers omitted bounds from the metadata already + read and checks supplied bounds. With ``validate=False``, bounds are + trusted; omitted packed-Q and KV bounds use plan capacities, and + fixed-Q bounds come from the query shape. Captured + launches can therefore replay with new metadata while the kernel reads + each request's live lengths from that metadata. + + Nonpositive per-run K/V lengths mark inactive requests in every + profile. Their page-table rows and data are not read; caller-owned + output is left untouched (newly allocated inactive output is unspecified). + Producers and reducers retain scheduler/PDL progress. Plan-owned lengths + may also mark inactive requests but remain immutable; use per-run lengths + for empty/live graph transitions. ``validate=True`` performs structural and value validation. It reads per-run metadata values back to the host. @@ -4100,85 +4332,39 @@ def run( device-to-host synchronization. K/V view selection, scale forwarding, and optional output allocation are unavoidable in both modes. - Packed plans require ``qo_indptr`` with ``B + 1`` int32 offsets. Fixed - plans require ``qo_indptr`` to be omitted. Per-run metadata tensors may - change identity between ordered runs. + Supply seq_lens_kv on every run unless plan() owns its CUDA copy. Packed queries require qo_indptr; fixed Q + rejects it under validation. Q/K/V dtypes, cache layout and model semantics + must match plan(). softmax_scale is the finite positive host effective QK scale + (default 1/sqrt(D)); output_scale is the host V/output scale (default 1). - In either validation mode, output and workspace must be disjoint from - each other and from all inputs. Storage overlap is not checked. - - Parameters - ---------- - q : torch.Tensor - Runtime fixed or packed query tensor matching the plan. - paged_kv_cache : torch.Tensor or tuple[torch.Tensor, torch.Tensor] - Runtime combined or separate paged K/V storage matching the plan. - seq_lens : torch.Tensor, optional - Per-run contiguous int32 CUDA K/V lengths with shape ``[B]``. - Required when omitted from :meth:`plan` and otherwise required to - be ``None``. - block_tables : torch.Tensor - Per-run int32 CUDA physical page IDs with shape ``[B, C]``. Entries - must be contiguous within each row; the row stride may be any value - at least ``C``. Inactive tail entries are ignored. - kv_scale_factors : tuple[torch.Tensor, torch.Tensor], optional - Required for packed NVFP4 K/V. Compact FP8 ``(K_SF, V_SF)`` tensors - have shape ``[pages, Hkv, storage_page_size, D/16]``; K scales are token-major - and V scales use the 4-token interleaved layout. Omit for other - dtypes. Scale tensors must not overlap output or workspace storage. - qo_indptr : torch.Tensor, optional - Per-run cumulative query offsets with shape ``[B + 1]``. Required for - a packed-query plan and rejected for a fixed-query plan. - bmm1_scale : float, optional - QK scaling factor. Defaults to the inverse square root of - ``head_dim``. - bmm2_scale : float - Value/output scaling factor. Defaults to ``1.0``. - out : torch.Tensor, optional - Caller-owned output tensor. A new tensor is allocated when omitted. - validate : bool - Run explicit structural and value validation. Disable only - when the caller guarantees the complete runtime contract. Sequence - length ownership is enforced in either mode. Defaults to ``True``. - - Returns - ------- - torch.Tensor - The fixed or packed attention output. - """ + Enabled attention_sinks requires compact CUDA FP32 [Hq] natural-log logits + after QK scaling. Values must be finite or -inf and add zero V contribution; + validation checks metadata without reading sink values.""" - state = self._require_plan_state() if not isinstance(validate, bool): raise TypeError("validate must be a bool") - planned_seq_lens = state.planned_seq_lens_device - plan_owns_seq_lens = planned_seq_lens is not None - if plan_owns_seq_lens: - if seq_lens is not None: + + state = self._require_plan_state() + if state.planned_seq_lens_kv_device is not None: + if seq_lens_kv is not None: raise ValueError( - "seq_lens must be None when plan() owns sequence lengths" + "run seq_lens_kv must be None when plan() owns sequence lengths" ) - effective_seq_lens = planned_seq_lens + effective_seq_lens = state.planned_seq_lens_kv_device else: - if seq_lens is None: + if seq_lens_kv is None: raise ValueError( - "seq_lens is required when plan() does not own sequence lengths" + "seq_lens_kv is required in run() when plan() does not own lengths" ) - effective_seq_lens = seq_lens + effective_seq_lens = seq_lens_kv - runtime_qo_indptr = qo_indptr if state.use_packed_q else None + runtime_qo_indptr = qo_indptr if state.packed_query else None if validate: ( metadata_device, metadata_batch_size, _, - ) = ( - _validate_block_tables(block_tables) - if plan_owns_seq_lens - else _validate_block_table_metadata( - block_tables, - effective_seq_lens, - ) - ) + ) = _validate_block_table_metadata(block_tables, effective_seq_lens) if metadata_device != state.device: raise ValueError( f"per-run metadata must be on {state.device}, got {metadata_device}" @@ -4188,7 +4374,7 @@ def run( "per-run metadata batch size must match the plan " f"({state.batch_size}), got {metadata_batch_size}" ) - if state.use_packed_q: + if state.packed_query: if qo_indptr is None: raise ValueError("qo_indptr is required for a packed-Q plan") _validate_qo_indptr( @@ -4206,7 +4392,7 @@ def run( device=state.device, batch_size=state.batch_size, seq_len_q=state.seq_len_q, - use_packed_q=state.use_packed_q, + packed_query=state.packed_query, num_qo_heads=state.num_qo_heads, num_kv_heads=state.num_kv_heads, head_dim=state.head_dim, @@ -4214,44 +4400,82 @@ def run( q_dtype=state.q_dtype, k_dtype=state.k_dtype, v_dtype=state.v_dtype, - output_dtype=state.output_dtype, - bmm1_scale=bmm1_scale, - bmm2_scale=bmm2_scale, + out_dtype=state.out_dtype, + softmax_scale=softmax_scale, + output_scale=output_scale, out=out, + kv_layout=self._kv_layout, ) - _validate_decode_run_metadata_values( + actual_q, actual_kv = _validate_decode_run_metadata_values( runtime, - planned_seq_lens_host=state.planned_seq_lens_host, + planned_seq_lens_kv_host=state.planned_seq_lens_kv_host, max_kv_len=state.max_kv_len, page_size=state.page_size, - use_packed_q=state.use_packed_q, + packed_query=state.packed_query, seq_len_q=state.seq_len_q, batch_size=state.batch_size, mask_type=state.mask_type, - seq_lens=effective_seq_lens, + seq_lens_kv=effective_seq_lens, block_tables=block_tables, qo_indptr=runtime_qo_indptr, ) + max_q_len = resolve_max_sequence_length( + max_q_len, actual_q, "max_q_len", state.seq_len_q + ) + max_kv_len = resolve_max_sequence_length( + max_kv_len, actual_kv, "max_kv_len", state.max_kv_len + ) else: runtime = _prepare_decode_runtime_unchecked( q, paged_kv_cache, kv_scale_factors, - output_dtype=state.output_dtype, - bmm1_scale=bmm1_scale, - bmm2_scale=bmm2_scale, + out_dtype=state.out_dtype, + softmax_scale=softmax_scale, + output_scale=output_scale, out=out, + kv_layout=self._kv_layout, ) - return _launch_decode( + if max_q_len is None: + max_q_len = ( + state.seq_len_q + if state.packed_query + else (1 if q.ndim == 3 else int(q.shape[1])) + ) + if max_kv_len is None: + max_kv_len = state.max_kv_len + + lse = prepare_lse( + lse, + q=q, + mode=state.lse_mode, + return_lse=return_lse, + validate=validate, + ) + attention_sinks = validate_attention_sinks( + attention_sinks, + q=q, + enabled=state.use_attention_sinks, + validate=validate, + ) + result = _launch_decode( runtime, - seq_lens=effective_seq_lens, - qo_indptr=runtime_qo_indptr, + seq_lens_kv=effective_seq_lens, + qo_indptr=( + runtime_qo_indptr if state.packed_query else state.fixed_q_offsets + ), block_tables=block_tables, workspace=state.workspace, compiled_main=state.compiled_main, compiled_reducer=state.compiled_reducer, + lse=lse, + attention_sinks=attention_sinks, + requested_splits=state.requested_splits, + max_kv_len=max_kv_len, + max_q_len=max_q_len, ) + return (result, lse) if return_lse else result @flashinfer_experimental_api @@ -4264,30 +4488,34 @@ def batch_decode_with_paged_kv_cache( kv_scale_factors: Optional[PagedKVScaleFactors] = None, seq_len_q: int = 1, qo_indptr: Optional[torch.Tensor] = None, - max_seq_len_q: Optional[int] = None, + max_q_len: Optional[int] = None, mask_type: Literal["dense", "causal"] = "dense", window_left: int = -1, - kv_layout: Literal["HND"] = "HND", - bmm1_scale: Optional[float] = None, - bmm2_scale: float = 1.0, + kv_layout: Literal["HND", "NHD"] = "HND", + softmax_scale: Optional[float] = None, + output_scale: float = 1.0, out: Optional[torch.Tensor] = None, + lse: Optional[torch.Tensor] = None, + return_lse: bool = False, + attention_sinks: Optional[torch.Tensor] = None, + lse_mode: Literal["base2"] = "base2", out_dtype: Optional[torch.dtype] = None, page_size: Optional[int] = None, split_kv: bool = True, workspace_buffer: Optional[torch.Tensor] = None, max_kv_len: Optional[int] = None, validate: bool = True, -) -> torch.Tensor: +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """One-shot fixed or packed-Q paged decode from fixed page tables. SQ1 preserves the ``[B, Hq, D]`` query/output contract. For fixed ``seq_len_q>1``, query and output are both token-major ``[B, SQ, Hq, D]``. Providing cumulative ``qo_indptr`` selects packed - ``[total_q, Hq, D]`` query/output; the wrapper derives ``max_seq_len_q`` - once when it is omitted. No transpose is hidden here. + ``[total_q, Hq, D]`` query/output; packed Q requires ``max_q_len`` + or a non-default ``seq_len_q`` alias. No transpose is hidden here. Without caller scratch, the convenience planner reads ``seq_lens_kv`` on - the host and retains the existing length-specialized kernel path. This mode + the host to size its plan and retain an owned metadata copy. This mode is not CUDA-graph-capturable. See Notes for the explicit capture-safe mode. Parameters @@ -4310,20 +4538,21 @@ def batch_decode_with_paged_kv_cache( seq_len_q : int Fixed query length when ``qo_indptr`` is omitted. In packed-query mode, a non-default value is a backward-compatible alias for - ``max_seq_len_q`` and must agree with it when both are provided. + ``max_q_len`` and must agree with it when both are provided. qo_indptr : torch.Tensor, optional Cumulative query offsets selecting packed-query mode. - max_seq_len_q : int, optional - Per-request packed-query length capacity. When omitted for packed Q, - it is derived from ``qo_indptr`` unless a non-default ``seq_len_q`` - supplies the bound. In fixed-query mode, it must equal ``seq_len_q``. + max_q_len : int, optional + Per-request packed-query length capacity. Required for packed Q + unless a non-default ``seq_len_q`` supplies the bound. It is never + inferred from ``qo_indptr``. In fixed-query mode, it must equal + ``seq_len_q``. mask_type : {"dense", "causal"} Attention mask mode. window_left : int Left sliding-window extent, or ``-1`` to disable the window. - kv_layout : {"HND"} + kv_layout : {"HND", "NHD"} Layout of the paged K/V cache. - bmm1_scale, bmm2_scale : float, optional + softmax_scale, output_scale : float, optional QK and value/output scaling factors. out : torch.Tensor, optional Caller-owned output tensor. @@ -4355,10 +4584,11 @@ def batch_decode_with_paged_kv_cache( Both owned and caller-provided scratch use the same wrapper plan/run path. For CUDA Graph capture, supply ``workspace_buffer``, ``max_kv_len``, ``out``, and ``validate=False``; packed Q also needs an explicit - ``max_seq_len_q`` (or its non-default ``seq_len_q`` alias). + ``max_q_len`` (or its non-default ``seq_len_q`` alias). Warm this exact topology outside capture first. Retain stable tensor storage and mutate metadata only between completed launches/replays. All live K/V - lengths must be positive and within the static bound, active page IDs must + positive lengths must be within the static bound; nonpositive lengths mark + inactive requests. Active page IDs must index the cache, and causal per-request Q lengths must not exceed K/V lengths. Packed offsets must start at zero, end at the query token count, and have strictly positive deltas within the Q bound. Scratch/output must not alias @@ -4366,6 +4596,25 @@ def batch_decode_with_paged_kv_cache( explicit path. Kernel policy, necessary control resets, and output layout are unchanged. Missing output may be allocated only outside capture. + + ``return_lse=True`` returns ``(output, lse)``. ``lse_mode`` selects + ``"base2"`` (the only supported base); it is used only when returning LSE + or writing a caller-provided ``lse`` buffer. + + ``lse`` is an optional independent, compact float32 output with shape + ``query.shape[:-1]`` (``q.shape[:-1]`` for FMHA). An enabled LSE plan + allocates it when omitted, including when ``return_lse=False``. A supplied + buffer is written regardless of ``return_lse``; the flag only controls + the Python return type. LSE includes any sink mass and excludes output + scaling. Keep LSE disjoint from output, inputs, and workspace. Supply + both output buffers and disable validation for allocation-free capture. + + ``attention_sinks`` is a compact CUDA float32 vector ``[Hq]`` shared + across batch/query positions. Values are natural-log sink logits after + QK scaling, with zero V contribution. Finite values and ``-inf`` are + accepted. Validation rejects NaN/positive infinity and may synchronize. + With validation disabled, the caller must preserve these contracts. + Sink values may change between completed CUDA graph replays. """ if not isinstance(validate, bool): @@ -4381,16 +4630,17 @@ def batch_decode_with_paged_kv_cache( raise ValueError("validate=False requires workspace_buffer") if max_kv_len is None: raise ValueError("validate=False requires max_kv_len") - if qo_indptr is not None and max_seq_len_q is None and seq_len_q == 1: - raise ValueError("validate=False requires max_seq_len_q for packed Q") + if qo_indptr is not None and max_q_len is None and seq_len_q == 1: + raise ValueError("validate=False requires max_q_len for packed Q") + _validate_layout(kv_layout) allocate_workspace = workspace_buffer is None if validate: - use_packed_q, resolved_seq_len_q = _resolve_q_mode( + packed_query, resolved_seq_len_q = _resolve_q_mode( seq_len_q=seq_len_q, qo_indptr=qo_indptr, - max_seq_len_q=max_seq_len_q, - require_packed_max=False, + max_q_len=max_q_len, + require_packed_max=True, ) metadata_device, batch_size, _ = _validate_block_table_metadata( block_tables, seq_lens_kv @@ -4403,23 +4653,23 @@ def batch_decode_with_paged_kv_cache( _validate_qo_indptr( qo_indptr, expected_device=q.device, batch_size=batch_size ) - if resolved_seq_len_q is None: - resolved_seq_len_q, _, _ = _read_packed_q_plan_metadata(qo_indptr) assert resolved_seq_len_q is not None _validate_q( q, seq_len_q=resolved_seq_len_q, - use_packed_q=use_packed_q, + packed_query=packed_query, device=q.device, batch_size=batch_size, ) k_cache, v_cache, _, num_kv_heads, storage_page_size, head_dim = ( - _normalize_paged_kv_cache_views(paged_kv_cache, expected_device=q.device) + _normalize_paged_kv_cache_views( + paged_kv_cache, expected_device=q.device, kv_layout=kv_layout + ) ) else: - use_packed_q = qo_indptr is not None + packed_query = qo_indptr is not None resolved_seq_len_q = ( - max_seq_len_q if use_packed_q and max_seq_len_q is not None else seq_len_q + max_q_len if packed_query and max_q_len is not None else seq_len_q ) batch_size = int(seq_lens_kv.shape[0]) k_cache = ( @@ -4432,6 +4682,9 @@ def batch_decode_with_paged_kv_cache( if isinstance(paged_kv_cache, torch.Tensor) else paged_kv_cache[1] ) + if kv_layout == "NHD": + k_cache = k_cache.transpose(1, 2) + v_cache = v_cache.transpose(1, 2) num_kv_heads, storage_page_size, storage_head_dim = map(int, k_cache.shape[1:]) head_dim = ( storage_head_dim * 2 if k_cache.dtype == torch.uint8 else storage_head_dim @@ -4441,24 +4694,13 @@ def batch_decode_with_paged_kv_cache( num_qo_heads = int(q.shape[-2]) if validate and out is not None and not isinstance(out, torch.Tensor): raise TypeError("out must be a torch.Tensor") - output_dtype = ( + out_dtype = ( out_dtype if out_dtype is not None else (out.dtype if out is not None else q.dtype) ) - # Only the allocating convenience call freezes lengths. Caller scratch - # always uses dynamic metadata, including when validation is enabled. - seq_lens_host = ( - tuple(int(value) for value in seq_lens_kv.tolist()) - if allocate_workspace - else None - ) if max_kv_len is None: - max_kv_len = ( - max(seq_lens_host) - if seq_lens_host is not None - else int(seq_lens_kv.max().item()) - ) + max_kv_len = int(seq_lens_kv.max().item()) if validate and out is not None: _validate_out( out, @@ -4468,58 +4710,35 @@ def batch_decode_with_paged_kv_cache( num_qo_heads=num_qo_heads, seq_len_q=resolved_seq_len_q, head_dim=head_dim, - total_q_tokens=int(q.shape[0]) if use_packed_q else None, + total_q_tokens=int(q.shape[0]) if packed_query else None, ), seq_len_q=resolved_seq_len_q, - use_packed_q=use_packed_q, - output_dtype=output_dtype, - ) - if allocate_workspace: - workspace_bytes = get_prims_ts_batch_decode_workspace_size( - batch_size, - num_qo_heads, - num_kv_heads, - head_dim, - page_size, - max_kv_len, - seq_len_q=seq_len_q, - qo_indptr=qo_indptr, - max_seq_len_q=resolved_seq_len_q, - q_dtype=q.dtype, - k_dtype=k_cache.dtype, - v_dtype=v_cache.dtype, - out_dtype=output_dtype, - mask_type=mask_type, - window_left=window_left, - kv_layout=kv_layout, - storage_page_size=storage_page_size, - device=q.device, - split_kv=split_kv, - ) - workspace_buffer = torch.empty( - workspace_bytes, dtype=torch.int8, device=q.device + packed_query=packed_query, + out_dtype=out_dtype, ) - wrapper = BatchDecodePagedTSWrapper(kv_layout=kv_layout) wrapper.plan( q.device, - batch_size, - num_qo_heads, - num_kv_heads, - head_dim, - page_size, - max_kv_len, - max_seq_len_q=resolved_seq_len_q, + batch_size=batch_size, + max_q_len=resolved_seq_len_q, + max_kv_len=max_kv_len, + num_qo_heads=num_qo_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + page_size=page_size, storage_page_size=storage_page_size, split_kv=split_kv, - packed_query=use_packed_q, - q_data_type=q.dtype, - k_data_type=k_cache.dtype, - v_data_type=v_cache.dtype, - o_data_type=output_dtype, + packed_query=packed_query, + q_dtype=q.dtype, + k_dtype=k_cache.dtype, + v_dtype=v_cache.dtype, + out_dtype=out_dtype, mask_type=mask_type, window_left=window_left, - seq_lens=seq_lens_host, + use_attention_sinks=attention_sinks is not None, + lse_mode=validate_lse_mode(lse_mode) + if return_lse or lse is not None + else "none", workspace_buffer=workspace_buffer, validate=validate, initialize_workspace=allocate_workspace, @@ -4527,13 +4746,16 @@ def batch_decode_with_paged_kv_cache( return wrapper.run( q, paged_kv_cache, - None if seq_lens_host is not None else seq_lens_kv, + seq_lens_kv, block_tables, kv_scale_factors=kv_scale_factors, qo_indptr=qo_indptr, - bmm1_scale=bmm1_scale, - bmm2_scale=bmm2_scale, + softmax_scale=softmax_scale, + output_scale=output_scale, out=out, + lse=lse, + return_lse=return_lse, + attention_sinks=attention_sinks, validate=validate, ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_kernel.py index 0c5a915d0bcf..d447b52846fa 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 @@ -1,3 +1,4 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -33,6 +34,13 @@ schedule before the repeated K/V tile loop, LOOP is the repeated K/V tile body, and TAIL is the one-time cleanup and drain after LOOP exits. +VC-Attention is the low-bit attention recipe of Li et al. 2026, arXiv +2609.15810, value smoothing with a softmax cast. Names with the ``vc`` prefix +belong to its recipes: QK16 keeps bf16 Q and K, QK8 adds centred, Hadamard +rotated E4M3 Q and K with block scales, and both store E4M3 V as residuals +around each 128-token tile mean, which the kernel restores inside the online +softmax. + Feature Support Matrix: | Feature | Support | @@ -145,7 +153,7 @@ import cutlass.pipeline as pipeline import cutlass.utils as utils from cuda.bindings import driver as cuda_drv -from cutlass import Int32 +from cutlass import Float32, Int32 from ..tensor_map import create_tensor_map_ragged_from_tensor from cutlass.experimental.task_scheduling.memory import ( @@ -172,6 +180,7 @@ SmemKVResource, SmemOResource, SmemPageOffsetsKvResource, + SmemSageKScalesResource, SmemQResource, TmemOResource, TmemPResource, @@ -181,6 +190,11 @@ TmemStatsDoneResource, TmemPPrefixReadyResource, ) +from ...vc_attention import vc_mean_operand_shape +from .vc_attention.vc_resources import SmemMuResource +from .sage_scales import validate_context_sage_layout +from ...sage.layout import _SageScaleLayout +from ...sage.modes import SageKTransport from .fmha_tasks import ( PackedContextWorkQueue, create_correction_task, @@ -683,18 +697,26 @@ def build_context_task_manager( tma_k_desc: cutlass.Pointer | None, tma_v_desc: cutlass.Pointer | None, tma_o_desc: cutlass.Pointer | None, + tma_mu_desc: cutlass.Pointer | None = None, + vc_q_scale: cute.Tensor | None = None, + vc_k_scale: cute.Tensor | None = None, + vc_repair_tiles: int = 0, cum_seqlen_q: cute.Tensor | None, cum_seqlen_k: cute.Tensor | None, variable_window_token_starts: cute.Tensor | None = None, variable_window_token_ends: cute.Tensor | None = None, variable_window_cta_starts: cute.Tensor | None = None, variable_window_q_stride: int | Int32 = 0, - scale_softmax_log2: cute.Tensor | None = None, - output_scale: cute.Tensor | None = None, + scale_softmax_log2: Float32 | None = None, + output_scale: Float32 | cute.Tensor | None = None, + lse_output: cute.Tensor | None = None, + attention_sinks: cute.Tensor | None = None, g_block_tables: cute.Pointer | None = None, block_table_row_stride: int | Int32 = 0, g_seq_lens_kv: cute.Pointer | None = None, max_seq_len_kv: int | Int32 | None = None, + sage_scales: cute.Pointer | None = None, + num_kv_heads: int | Int32 = 0, num_kv_tiles: int | Int32, q_offset: int | Int32, domain_n_kwargs: DomainKwargs, @@ -705,7 +727,7 @@ def build_context_task_manager( is_persistent: bool = True, is_clc_dynamic: bool = False, clc_response_ptr: cute.Pointer | None = None, - exhaustive_deadlock_race_check: bool = True, + exhaustive_deadlock_race_check: bool | None = None, ) -> Tuple[ TaskManager, list[MemoryResource], @@ -717,9 +739,13 @@ def build_context_task_manager( ]: """Build the FMHA TaskManager with all resources, tasks, and dependency graph. - ``scale_softmax_log2`` and ``output_scale`` are optional one-element - float32 device tensors. Resource setup loads element 0 once before the K/V - loop and reuses the value through task-local dataflow. ``scale_softmax_log2`` + ``sage_scales`` is a Sage plan's prepared FP32 scale buffer + (``cfg.sage_scale_layout``), whose sequences are strided by + ``num_kv_heads``; the softmax groups and the correction warps read it in + place. + + ``scale_softmax_log2`` and ``output_scale`` are runtime float32 scalars. + Resources carry these operands through task-local dataflow. ``scale_softmax_log2`` is the base-2 softmax multiplier, normally ``softmax_scale * log2(e)``; FP8 callers can fold Q/K dequant scales into it. ``output_scale`` multiplies the final O store; FP8 output can fold the @@ -741,6 +767,16 @@ def build_context_task_manager( TMA descriptor pointers for Q, K, V, O. cum_seqlen_q, cum_seqlen_k : cute.Tensor, optional Cumulative sequence-length metadata for varlen launches. + tma_mu_desc : cutlass.Pointer, optional + TMA descriptor of the VC-Attention tile-mean operands, the B operand of + the mean UMMA step. None unless the kernel restores means. + vc_q_scale, vc_k_scale : cute.Tensor, optional + VC-Attention-QK8 dequantization scales, one per 128-token Q block and + per K/V tile, in the flat per-sequence slot order of ``vc_scales``. + None under QK16. + vc_repair_tiles : int + VC-Attention repair tiles appended to the K/V sequence, kept out of + the softmax denominator. Zero without V repair. num_kv_tiles : int or Int32 Number of KV tiles in the loop domain. q_offset : int or Int32 @@ -762,9 +798,10 @@ def build_context_task_manager( clc_response_ptr : cute.Pointer, optional SMEM pointer for the CLC response buffer, required when ``is_clc_dynamic`` is true. - exhaustive_deadlock_race_check : bool - Whether TaskManager runs the exhaustive schedule checker during - construction. + exhaustive_deadlock_race_check : bool, optional + Whether TaskManager runs the exhaustive schedule checker. By default, + concrete offline domains enable it; symbolic JIT domains retain the + structural checks. An explicit Boolean overrides that default. Returns ------- tuple @@ -881,6 +918,7 @@ def build_context_task_manager( # The load warp (1 warp = 32 threads) is the only consumer; signal-All # requires the consumer group to size match num_warps × warp_size. + # Connect the auxiliary page-ID producer to the load warp for num_stages slots. def make_page_offsets_pipeline_cfg(num_stages: int) -> PipelineConfig: return PipelineConfig.create_async_async_pipeline_cfg( num_stages=num_stages, @@ -1086,6 +1124,35 @@ def make_page_offsets_pipeline_cfg(num_stages: int) -> PipelineConfig: name="smem_v", **kv_resource_kwargs, ) + # Mu is the per-tile mean of V. V tile t is stored as E4M3 residuals + # V8_t = V_t - mu_t, so the PV product splits into + # O = sum_t P_t V8_t + sum_t rowsum(P_t) mu_t + # and the second term is a bf16 K=16 UMMA per group of sixteen tiles with + # the row sums as A and the means as B. Mu is therefore a UMMA operand + # like K and V and must sit in shared memory. The load warp stages it + # beside each V tile and the UMMA warp consumes it. + smem_mu: SmemMuResource | None = None + if cfg.vc_restores_means: + smem_mu_pipeline_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( + num_stages=cfg.vc_mean_stages, + # Under two-CTA UMMA each CTA stages the N/2 channels of the mean + # operand its half of the M=256 step reads (like V columns). + num_bytes=cfg.vc_mean_tile_bytes, + producer_group=tma_producer_group, + consumer_group=pipeline.CooperativeGroup(Agent.Thread), + cta_layout_vmnk=cluster_shape_vmnk, + advance_on_wait=True, + num_bytes_per_warp_per_cta=cfg.vc_mean_tile_bytes // cfg.cta_group_size + if two_cta + else None, + **tma_umma_leader_kwargs, + ) + smem_mu = SmemMuResource( + tma_mu_desc=tma_mu_desc, + pipeline_config=smem_mu_pipeline_cfg, + cfg=cfg, + name="smem_mu", + ) # WorkQueue: persistent tile scheduler state (static or CLC dynamic), not # Q/K/V/O dataflow. Non-persistent launches omit it so each CTA executes @@ -1123,15 +1190,12 @@ def make_page_offsets_pipeline_cfg(num_stages: int) -> PipelineConfig: "pipeline_config": work_queue_pipeline_cfg, "name": "work_queue", } - if ( - cfg.is_causal - and cfg.has_varlen - and not cfg.has_uniform_varlen - and cum_seqlen_q is not None - ): + if (cfg.has_varlen and not cfg.has_uniform_varlen) or g_seq_lens_kv is not None: work_queue = PackedContextWorkQueue( cfg=cfg, cum_seqlen_q=cum_seqlen_q, + cum_seqlen_k=cum_seqlen_k, + seq_lens_kv=g_seq_lens_kv, **work_queue_kwargs, ) else: @@ -1146,15 +1210,12 @@ def make_page_offsets_pipeline_cfg(num_stages: int) -> PipelineConfig: "tile_scheduler_config": tile_scheduler_config, "name": "work_queue", } - if ( - cfg.is_causal - and cfg.has_varlen - and not cfg.has_uniform_varlen - and cum_seqlen_q is not None - ): + if (cfg.has_varlen and not cfg.has_uniform_varlen) or g_seq_lens_kv is not None: work_queue = PackedContextWorkQueue( cfg=cfg, cum_seqlen_q=cum_seqlen_q, + cum_seqlen_k=cum_seqlen_k, + seq_lens_kv=g_seq_lens_kv, **work_queue_kwargs, ) else: @@ -1182,6 +1243,26 @@ def make_page_offsets_pipeline_cfg(num_stages: int) -> PipelineConfig: for index in range(cfg.num_qkv_instances) ] smem_p0, smem_p1 = smem_p_resources + # Sage ring transport (``SageKTransport.SMEM``): the load warp stages each + # K tile's sfK words in a ring both softmax groups read. + smem_sage_k: SmemSageKScalesResource | None = None + if cfg.sage_k_transport is SageKTransport.SMEM: + smem_sage_k = SmemSageKScalesResource( + pipeline_config=PipelineConfig.create_async_async_pipeline_cfg( + num_stages=cfg.sage_k_scale_stages, + producer_group=pipeline.CooperativeGroup(Agent.Thread, warp_size), + consumer_group=pipeline.CooperativeGroup( + Agent.Thread, + cfg.num_qkv_instances * num_softmax_warps * warp_size, + ), + cta_layout_vmnk=cluster_shape_vmnk, + producer_op=pipeline.PipelineOp.AsyncLoad, + ), + cfg=cfg, + sage_scales=sage_scales, + num_kv_heads=num_kv_heads, + name="smem_sage_k", + ) tmem_sp0 = TmemSPResource( pipeline_config=tmem_sp0_pipeline_cfg, cfg=cfg, @@ -1192,12 +1273,18 @@ def make_page_offsets_pipeline_cfg(num_stages: int) -> PipelineConfig: cum_seqlen_q=cum_seqlen_q, cum_seqlen_k=cum_seqlen_k, seq_lens_kv=g_seq_lens_kv, + max_seq_len_kv=max_seq_len_kv, variable_window_token_starts=variable_window_token_starts, variable_window_token_ends=variable_window_token_ends, variable_window_cta_starts=variable_window_cta_starts, variable_window_q_stride=variable_window_q_stride, scale_softmax_log2=scale_softmax_log2, smem_p=smem_p0, + vc_q_scale=vc_q_scale, + vc_k_scale=vc_k_scale, + vc_repair_tiles=vc_repair_tiles, + sage_scales=sage_scales, + num_kv_heads=num_kv_heads, name="tmem_sp0", ) tmem_p0: TmemPResource | None = None @@ -1214,6 +1301,11 @@ def make_page_offsets_pipeline_cfg(num_stages: int) -> PipelineConfig: tmem_vec_offset=cfg.tmem_vec0_offset, scale_softmax_log2=scale_softmax_log2, output_scale=output_scale, + lse_output=lse_output, + attention_sinks=attention_sinks, + cum_seqlen_q=cum_seqlen_q, + query_seq_len=variable_window_q_stride, + peer_idx=0, name="tmem_vec0", ) smem_o_0 = SmemOResource( @@ -1221,6 +1313,8 @@ def make_page_offsets_pipeline_cfg(num_stages: int) -> PipelineConfig: cfg=cfg, stage_idx=0, tmem_vec_resource=tmem_vec0, + sage_scales=sage_scales, + num_kv_heads=num_kv_heads, name="smem_o_0", ) gmem_o_0 = GmemOResource( @@ -1285,12 +1379,19 @@ def make_page_offsets_pipeline_cfg(num_stages: int) -> PipelineConfig: cum_seqlen_q=cum_seqlen_q, cum_seqlen_k=cum_seqlen_k, seq_lens_kv=g_seq_lens_kv, + max_seq_len_kv=max_seq_len_kv, variable_window_token_starts=variable_window_token_starts, variable_window_token_ends=variable_window_token_ends, variable_window_cta_starts=variable_window_cta_starts, variable_window_q_stride=variable_window_q_stride, scale_softmax_log2=scale_softmax_log2, smem_p=smem_p1, + vc_q_scale=vc_q_scale, + vc_k_scale=vc_k_scale, + vc_repair_tiles=vc_repair_tiles, + sage_scales=sage_scales, + num_kv_heads=num_kv_heads, + score_seed_owner=tmem_sp0, name="tmem_sp1", ) tmem_vec1 = TmemStatsResource( @@ -1299,6 +1400,11 @@ def make_page_offsets_pipeline_cfg(num_stages: int) -> PipelineConfig: tmem_vec_offset=cfg.tmem_vec1_offset, scale_softmax_log2=scale_softmax_log2, output_scale=output_scale, + lse_output=lse_output, + attention_sinks=attention_sinks, + cum_seqlen_q=cum_seqlen_q, + query_seq_len=variable_window_q_stride, + peer_idx=1, name="tmem_vec1", ) smem_o_1 = SmemOResource( @@ -1306,6 +1412,8 @@ def make_page_offsets_pipeline_cfg(num_stages: int) -> PipelineConfig: cfg=cfg, stage_idx=1, tmem_vec_resource=tmem_vec1, + sage_scales=sage_scales, + num_kv_heads=num_kv_heads, name="smem_o_1", ) gmem_o_1 = GmemOResource( @@ -1373,6 +1481,8 @@ def scheduler_deps(*resources: MemoryResource) -> list[MemoryResource]: work_queue, smem_page_offsets_kv=smem_page_offsets_kv, smem_page_offsets_v=smem_page_offsets_v, + smem_mu=smem_mu, + smem_sage_k=smem_sage_k, **domain_n_kwargs, ) mma_task = create_mma_task( @@ -1391,6 +1501,7 @@ def scheduler_deps(*resources: MemoryResource) -> list[MemoryResource]: tmem_p_prefix_ready_1=tmem_p_prefix_ready_1, smem_p0=smem_p0, smem_p1=smem_p1, + smem_mu=smem_mu, **mma_domain_kwargs, ) @@ -1405,6 +1516,7 @@ def scheduler_deps(*resources: MemoryResource) -> list[MemoryResource]: work_queue, tmem_p_prefix_ready=tmem_p_prefix_ready_0, smem_p=smem_p0, + smem_sage_k=smem_sage_k, **softmax0_domain_kwargs, ) softmax1_task: Task | None = None @@ -1420,6 +1532,7 @@ def scheduler_deps(*resources: MemoryResource) -> list[MemoryResource]: work_queue, tmem_p_prefix_ready=tmem_p_prefix_ready_1, smem_p=smem_p1, + smem_sage_k=smem_sage_k, **softmax1_domain_kwargs, ) correction_task = create_correction_task( @@ -1544,6 +1657,8 @@ def scheduler_deps(*resources: MemoryResource) -> list[MemoryResource]: if smem_page_offsets_v is not None: smem_v_deps.append(smem_page_offsets_v) resource_dependency_graph[smem_v] = scheduler_deps(*smem_v_deps) + if smem_mu is not None: + resource_dependency_graph[smem_mu] = scheduler_deps(gmem_qkv) resource_dependency_graph.update( { tmem_sp0: scheduler_deps( @@ -1588,6 +1703,8 @@ def scheduler_deps(*resources: MemoryResource) -> list[MemoryResource]: resource_dependency_graph[work_queue] = [work_queue] if is_clc_dynamic else [] if smem_page_offsets_kv is not None: resource_dependency_graph[smem_page_offsets_kv] = scheduler_deps(gmem_qkv) + if smem_sage_k is not None: + resource_dependency_graph[smem_sage_k] = scheduler_deps(gmem_qkv) if smem_page_offsets_v is not None: resource_dependency_graph[smem_page_offsets_v] = scheduler_deps(gmem_qkv) @@ -1605,10 +1722,13 @@ def add_smem_resource(resource: MemoryResource | None) -> None: add_smem_resource(smem_kv) if cfg.split_kv_pipelines: add_smem_resource(smem_v) + if smem_mu is not None: + add_smem_resource(smem_mu) if smem_page_offsets_kv is not None: add_smem_resource(smem_page_offsets_kv) if smem_page_offsets_v is not None: add_smem_resource(smem_page_offsets_v) + add_smem_resource(smem_sage_k) # Register every pipeline resource, including pipeline-only TMEM handoffs, # with the unified allocator. Otherwise CUTLASS materializes those # barriers as separate dynamic-SMEM arrays that the capacity selector @@ -1738,6 +1858,10 @@ def add_smem_resource(resource: MemoryResource | None) -> None: tmem_allocator.compute_layout() skip = not isinstance(num_kv_tiles, int) + # skip_validation only downgrades errors to warnings. Exhaustive searches + # over a symbolic JIT domain add cold-start cost without an offline proof. + if exhaustive_deadlock_race_check is None: + exhaustive_deadlock_race_check = not skip task_manager = TaskManager( tasks=task_list, resource_dependency_graph=resource_dependency_graph, @@ -1806,12 +1930,16 @@ def _context_pipeline_stage_counts( ), "tmem_sp": cfg.mma_softmax_stage * cfg.num_qkv_instances, "tmem_p": cfg.mma_softmax_stage if cfg.has_tmem_p_pipeline else 0, + "smem_sage_k": cfg.sage_k_scale_stages + if cfg.sage_k_transport is SageKTransport.SMEM + else 0, "tmem_vec": cfg.softmax_corr_stage * cfg.num_qkv_instances, "tmem_o": cfg.mma_corr_stage, "smem_o": cfg.num_qkv_instances, "s0s1_seq": 0 if cfg.single_qkv_instance else 1, "tmem_p_prefix_ready": 2 if cfg.pv_half_overlap else 0, "smem_p": cfg.num_qkv_instances if cfg.p_in_smem else 0, + "smem_mu": cfg.vc_mean_stages if cfg.vc_restores_means else 0, "tmem_stats_done": 0 if cfg.stats_via_smem else cfg.num_qkv_instances, "work_queue": 1 if is_clc_dynamic else 0, } @@ -1826,8 +1954,9 @@ def _kv_ring_smem_budget_bytes( ) -> int: """Return the SMEM left for the K/V ring(s) after every fixed allocation. - Fixed terms are Q, O, correction statistics, page-ID rings, control - records, and one 16-byte barrier per stage. The task manager still checks. + Fixed terms are Q, O, correction statistics, page-ID rings, the Sage sfK + ring, control records, and one 16-byte barrier per stage. The task + manager still checks. """ q_row_bytes = (cfg.q_dtype.width * cfg.smem_q_head_dim + 7) // 8 q_row_bytes = ( @@ -1872,9 +2001,20 @@ def _kv_ring_smem_budget_bytes( align = cfg.buffer_align_bytes smem_p_bytes = 0 if cfg.p_in_smem: + p_block_bytes = cfg.smem_p_bytes + if cfg.vc_restores_means: + # The row-sum operand lives behind each P tile. + p_block_bytes += cfg.vc_rowsum_tile_bytes + if cfg.vc_qk8: + # VC-Attention-QK8 keeps its K-scale table behind the P tile. + p_block_bytes += cfg.vc_kscale_table_bytes smem_p_bytes = ( - (cfg.smem_p_bytes + align - 1) // align * align * cfg.num_qkv_instances + (p_block_bytes + align - 1) // align * align * cfg.num_qkv_instances ) + sage_k_bytes = 0 + if cfg.sage_k_transport is SageKTransport.SMEM: + sage_k_bytes = cfg.sage_k_scale_stages * cfg.sage_k_tile_bytes + score_seed_bytes = cfg.sage_score_seed_tile_bytes if cfg.uses_int32_scores else 0 control_bytes = (2 * cutlass.Int32.width + cutlass.Int64.width) // 8 if is_clc_dynamic: control_bytes = (control_bytes + 15) // 16 * 16 + cutlass.Int128.width // 8 @@ -1885,12 +2025,17 @@ def _kv_ring_smem_budget_bytes( is_clc_dynamic=is_clc_dynamic, ).values() ) + if cfg.vc_restores_means: + # The callers charge the tile-mean ring per K/V stage. + fixed_barrier_stages -= cfg.vc_mean_stages fixed_smem_bytes = ( (q_tile_bytes * cfg.q_stage + align - 1) // align * align + (o_stage_bytes + align - 1) // align * align * cfg.num_qkv_instances + smem_p_bytes + (stats_bytes + 15) // 16 * 16 + (page_offsets_bytes + 15) // 16 * 16 + + (sage_k_bytes + 15) // 16 * 16 + + (score_seed_bytes + 127) // 128 * 128 + (control_bytes + 7) // 8 * 8 + fixed_barrier_stages * _PIPELINE_BARRIER_BYTES_PER_STAGE ) @@ -1924,6 +2069,10 @@ def _infer_single_instance_kv_stages( cfg.qk_mma_tiler[1] // cfg.cta_group_size * kv_head_dim * kv_dtype_width // 8 ) kv_stage_footprint_bytes = kv_stage_bytes + _PIPELINE_BARRIER_BYTES_PER_STAGE + if cfg.vc_restores_means: + kv_stage_footprint_bytes += ( + cfg.vc_mean_tile_bytes + _PIPELINE_BARRIER_BYTES_PER_STAGE + ) memory_fit_stages = kv_budget_bytes // kv_stage_footprint_bytes cadence_stages = cfg.num_head_dim_stages_k + cfg.num_head_dim_stages_v if require_cadence and memory_fit_stages < cadence_stages: @@ -1949,6 +2098,16 @@ def _configure_kv_ring_depths(cfg: FmhaConfig, *, is_clc_dynamic: bool) -> None: cfg, is_clc_dynamic=is_clc_dynamic, require_cadence=False ), ) + if cfg.kv_ring_stages: + n_stages = max( + cfg.num_head_dim_stages_k + cfg.num_head_dim_stages_v, + cfg.kv_ring_stages, + ) + if cfg.vc_restores_means and n_stages < 2: + # The tail holds the last two tile-mean operands at once. + raise ValueError( + f"VC-Attention needs at least 2 K/V stages, kv_ring_stages gives {n_stages}" + ) cfg.kv_stage = n_stages cfg.kv_stage_k = n_stages cfg.kv_stage_v = n_stages @@ -1968,6 +2127,8 @@ def _configure_kv_ring_depths(cfg: FmhaConfig, *, is_clc_dynamic: bool) -> None: kv_rows_per_cta * kv_head_dim * cfg.v_dtype.width // 8 + _PIPELINE_BARRIER_BYTES_PER_STAGE ) + if cfg.vc_restores_means: + v_stage_footprint += cfg.vc_mean_tile_bytes + _PIPELINE_BARRIER_BYTES_PER_STAGE # K and V share the same head_dim and head_dim_per_stage_kv, so their # minimum ring depths (num_head_dim_stages) are equal and both rings are sized identically. cadence = cfg.num_head_dim_stages_k @@ -1978,6 +2139,14 @@ def _configure_kv_ring_depths(cfg: FmhaConfig, *, is_clc_dynamic: bool) -> None: f"({cadence * (k_stage_footprint + v_stage_footprint)} bytes), but the " f"shared-memory budget fits only {budget_bytes} bytes" ) + if cfg.vc_restores_means and n_stages < 2: + # The tail holds the last two tile-mean operands at once. + raise ValueError( + f"VC-Attention needs at least 2 stages per K/V ring, but the " + f"shared-memory budget of {budget_bytes} bytes fits only {n_stages}" + ) + if cfg.kv_ring_stages: + n_stages = max(cadence, cfg.kv_ring_stages) cfg.kv_stage_k = n_stages cfg.kv_stage_v = n_stages @@ -2145,6 +2314,36 @@ def _configure_single_instance_warp_layout(cfg: FmhaConfig) -> None: cfg.num_regs_other = 112 +def _check_sage_attention(cfg: FmhaConfig) -> None: + """Check a Sage plan's scale layout against the dense query-paired FP8 schedule. + + Sage attention runs on the dense query-paired D128 schedule with FP8 P + staged in SMEM, in its single-CTA or two-CTA UMMA form. Its row max + reduces 16-column score blocks, with LDTM.STAT where the hardware has it. + """ + layout = cfg.sage_scale_layout + validate_context_sage_layout(layout, cfg.kv_tile_n) + if ( + cfg.single_qkv_instance + or cfg.head_paired + or cfg.is_causal + or cfg.has_variable_window + or cfg.use_paged_kv + or cfg.split_kv_pipelines + or cfg.kv_tile_start_window_size_left > 0 + ): + raise ValueError("Sage attention requires the dense query-paired D128 schedule") + if not cfg.p_in_smem: + raise ValueError( + "Sage attention requires E4M3 V with a 16-bit output (P in SMEM)" + ) + if ( + layout.v_rows.headdim != cfg.pv_mma_tiler[1] + or layout.q_rows.rows_per_tile != cfg.q_tile_m + ): + raise ValueError("the Sage scale layout does not match the kernel's tiles") + + def _uses_smem_p(cfg: FmhaConfig, *, has_variable_window: bool) -> bool: """Stage fp8 P in SMEM for the dense query-paired D128 schedule with a 16-bit output. It needs a 128-byte P row, one SW128 atom, and an O tile stored in @@ -2163,8 +2362,10 @@ def _uses_smem_p(cfg: FmhaConfig, *, has_variable_window: bool) -> bool: ) -def _configure_head_dim_staging(cfg: FmhaConfig) -> None: - """Stage K/V in 128-wide slices and size O staging to the paired footprint.""" +def _configure_head_dim_staging( + cfg: FmhaConfig, *, head_dim_per_stage: int = 128 +) -> None: + """Stage K/V by head slices and keep O staging within the SMEM budget.""" cfg.head_dim_per_stage_kv = 0 cfg.num_head_dim_stages_k = 1 cfg.num_head_dim_stages_v = 1 @@ -2179,25 +2380,43 @@ def _configure_head_dim_staging(cfg: FmhaConfig) -> None: return if cfg.num_qkv_instances != 1 and cfg.logical_head_dim_qk != 192: return - cfg.head_dim_per_stage_kv = 128 + if head_dim_per_stage not in (128, 192): + raise ValueError("head staging must use 128 or 192 channels") + if head_dim_per_stage == 192: + if not ( + cfg.logical_head_dim_qk == 192 + and cfg.single_qkv_instance + and cfg.q_dtype.width == cfg.v_dtype.width == 16 + ): + raise ValueError( + "Full D192 staging requires single-instance D192 with 16-bit Q/K/V" + ) + cfg.qk_mma_tiler = (*cfg.qk_mma_tiler[:2], 192) + cfg.pv_mma_tiler = (cfg.pv_mma_tiler[0], 192, cfg.pv_mma_tiler[2]) + cfg.epi_tile = cfg.pv_mma_tiler[:2] + cfg.head_dim_per_stage_kv = head_dim_per_stage cfg.num_head_dim_stages_k = cfg.qk_mma_tiler[2] // cfg.head_dim_per_stage_kv - cfg.num_head_dim_stages_v = cfg.pv_mma_tiler[1] // cfg.head_dim_per_stage_kv - # K/V transfers remain 128-wide. Unlike the single-query schedule, paired + cfg.num_head_dim_stages_v = ( + cfg.pv_mma_tiler[1] + cfg.head_dim_per_stage_kv - 1 + ) // cfg.head_dim_per_stage_kv + # Output staging can use narrower head bands. Unlike the single-query schedule, paired # 16-bit Q/O needs two 64-wide O buffers to leave room for both Q tiles # and a complete K/K/V cadence. FP8 can retain 128-wide O stores. cfg.head_dim_per_stage_o = ( 64 if ( - not cfg.single_qkv_instance + (not cfg.single_qkv_instance or head_dim_per_stage == 192) and cfg.q_dtype.width == 16 and cfg.o_dtype.width == 16 ) else 128 ) - cfg.num_o_head_dim_stages = cfg.epi_tile[1] // cfg.head_dim_per_stage_o + cfg.num_o_head_dim_stages = ( + cfg.epi_tile[1] + cfg.head_dim_per_stage_o - 1 + ) // cfg.head_dim_per_stage_o cfg.stage_kv_by_head_dim = True cfg.stage_o_by_head_dim = cfg.single_qkv_instance or cfg.num_o_head_dim_stages > 1 - # K/V stays 128-wide. Smaller O stages let paired 16-bit geometry + # Default K/V staging stays 128-wide. Smaller O stages let paired geometry # retain its K/V ring without exceeding Blackwell's SMEM budget. @@ -2226,13 +2445,13 @@ def _configure_head_paired_tma_copy_metadata( o_dtype: type, ) -> None: """Derive TMA copy granularities for head-paired Q/K/V/O tensors.""" - inner_dim_size = cfg.qk_mma_tiler[2] * q_dtype.width // 8 + inner_dim_size = cfg.smem_q_head_dim * q_dtype.width // 8 cfg.tma_copy_qkv_iters = 1 if inner_dim_size % 128 == 0: cfg.tma_copy_qkv_iters = inner_dim_size // 128 elif inner_dim_size != 64 and inner_dim_size != 32: raise RuntimeError(f"Unsupported inner dimension size: {inner_dim_size}") - tma_copy_qkv_granu_inner = cfg.qk_mma_tiler[2] // cfg.tma_copy_qkv_iters + tma_copy_qkv_granu_inner = cfg.smem_q_head_dim // cfg.tma_copy_qkv_iters cfg.q_tile_m = cfg.qk_mma_tiler[0] cfg.tma_copy_q_elements = cfg.sQ_shape[1] @@ -2253,18 +2472,18 @@ def _configure_head_paired_tma_copy_metadata( ) cfg.tma_copy_kv_bytes = cfg.tma_copy_kv_elements * k_dtype.width // 8 - v_inner_dim_size = cfg.pv_mma_tiler[1] * v_dtype.width // 8 + v_inner_dim_size = cfg.smem_v_head_dim * v_dtype.width // 8 cfg.tma_copy_v_iters = 1 if v_inner_dim_size % 128 == 0: cfg.tma_copy_v_iters = v_inner_dim_size // 128 elif v_inner_dim_size != 64 and v_inner_dim_size != 32: raise RuntimeError(f"Unsupported inner dimension size: {v_inner_dim_size}") - cfg.tma_copy_v_granu_inner = cfg.pv_mma_tiler[1] // cfg.tma_copy_v_iters + cfg.tma_copy_v_granu_inner = cfg.smem_v_head_dim // cfg.tma_copy_v_iters cfg.tma_copy_v_stage_iters = kv_head_dim // cfg.tma_copy_v_granu_inner cfg.tma_copy_v_granu_elems = cfg.tma_copy_kv_elements // cfg.tma_copy_v_stage_iters cfg.tma_copy_v_bytes = cfg.tma_copy_kv_elements * v_dtype.width // 8 - output_inner_dim_size = cfg.epi_tile[1] * o_dtype.width // 8 + output_inner_dim_size = cfg.smem_o_head_dim * o_dtype.width // 8 cfg.tma_copy_o_iters = 1 if output_inner_dim_size % 128 == 0: cfg.tma_copy_o_iters = output_inner_dim_size // 128 @@ -2272,7 +2491,7 @@ def _configure_head_paired_tma_copy_metadata( raise RuntimeError( f"Unsupported output inner dimension size: {output_inner_dim_size}" ) - cfg.tma_copy_o_granu_inner = cfg.epi_tile[1] // cfg.tma_copy_o_iters + cfg.tma_copy_o_granu_inner = cfg.smem_o_head_dim // cfg.tma_copy_o_iters o_head_dim = ( cfg.head_dim_per_stage_o if cfg.stage_o_by_head_dim else cfg.epi_tile[1] ) @@ -2579,17 +2798,25 @@ def build_fmha_task_manager( variable_window_token_ends: cute.Tensor | None = None, variable_window_cta_starts: cute.Tensor | None = None, variable_window_q_stride: int | Int32 = 0, - scale_softmax_log2: cute.Tensor | None = None, - output_scale: cute.Tensor | None = None, + scale_softmax_log2: Float32 | None = None, + output_scale: Float32 | cute.Tensor | None = None, + lse_output: cute.Tensor | None = None, + attention_sinks: cute.Tensor | None = None, q_offset: int | Int32 = 0, g_block_tables: cute.Pointer | None = None, block_table_row_stride: int | Int32 = 0, g_seq_lens_kv: cute.Pointer | None = None, max_seq_len_kv: int | Int32 | None = None, + tma_mu_desc: cutlass.Pointer | None = None, + vc_q_scale: cute.Tensor | None = None, + vc_k_scale: cute.Tensor | None = None, + vc_repair_tiles: int = 0, + sage_scales: cute.Pointer | None = None, + num_kv_heads: int | Int32 = 0, is_persistent: bool = True, is_clc_dynamic: bool = False, clc_response_ptr: cute.Pointer | None = None, - exhaustive_deadlock_race_check: bool = True, + exhaustive_deadlock_race_check: bool | None = None, ) -> Tuple[ TaskManager, list[MemoryResource], @@ -2602,8 +2829,8 @@ def build_fmha_task_manager( """Build the FMHA TaskManager using the shared context TS graph. Runtime scale arguments follow the same contract as - :func:`build_context_task_manager`: one-element float32 device tensors, - indexed at element 0 and cached by resource auxiliary work. + :func:`build_context_task_manager`: runtime float32 scalar operands + carried through task-local auxiliary dataflow. """ if ( cfg.causal_single_kv_tile @@ -2666,6 +2893,10 @@ def build_fmha_task_manager( tma_k_desc=tma_k_desc, tma_v_desc=tma_v_desc, tma_o_desc=tma_o_desc, + tma_mu_desc=tma_mu_desc, + vc_q_scale=vc_q_scale, + vc_k_scale=vc_k_scale, + vc_repair_tiles=vc_repair_tiles, cum_seqlen_q=cum_seqlen_q, cum_seqlen_k=cum_seqlen_k, variable_window_token_starts=variable_window_token_starts, @@ -2674,6 +2905,10 @@ def build_fmha_task_manager( variable_window_q_stride=variable_window_q_stride, scale_softmax_log2=scale_softmax_log2, output_scale=output_scale, + sage_scales=sage_scales, + num_kv_heads=num_kv_heads, + lse_output=lse_output, + attention_sinks=attention_sinks, g_block_tables=g_block_tables, block_table_row_stride=block_table_row_stride, g_seq_lens_kv=g_seq_lens_kv, @@ -2746,6 +2981,10 @@ class FmhaTs: into the TMEM S load on the non-masked path (default: False). The context runner enables this by default on SM103 (B300) and SM107 (Rubin). + kv_ring_stages : int, optional + K/V ring depth. Zero derives the depth from the shared-memory budget. + umma_fp8_k64 : bool, optional + Issue the fp8 PV UMMAs with K=64 per instruction. two_cta_umma : bool, optional Issue the QK and PV UMMAs in ``cta_group::2`` across a 2-CTA cluster, each CTA staging half of every K/V tile. Dense contiguous query-paired D128 with @@ -2786,14 +3025,18 @@ def __init__( enable_skip_correction: bool = True, uses_ldtm_stat: bool = False, fp8_psmem_early_token: bool = False, + kv_ring_stages: int = 0, + umma_fp8_k64: bool = False, two_cta_umma: bool = False, use_paged_kv: bool = False, num_tokens_per_page: int = 32, - max_kv_len: int = 1, + max_kv_len: int | None = 1, causal_single_kv_tile: bool = False, - exhaustive_deadlock_race_check: bool = True, + exhaustive_deadlock_race_check: bool | None = None, d_v: int | None = None, + vc_attention: str | None = None, use_blackwell_tma_oob_workaround: bool = False, + sage_scale_layout: _SageScaleLayout | None = None, ) -> None: """Initialize mode-specific tiling, dtype, and schedule configuration.""" head_paired = resolve_head_paired_mode( @@ -2801,6 +3044,9 @@ def __init__( is_causal=is_causal, window_size_left=window_size_left, ) + # Load helpers before staged Sage copies enter GPU control flow. + importlib.import_module(f"{__package__}.helpers_paged") + if causal_single_kv_tile and (not is_causal or head_paired): raise ValueError( "causal_single_kv_tile requires query-paired causal attention" @@ -2823,16 +3069,13 @@ def __init__( ) validate_head_paired_head_ratio(head_paired=head_paired, h_r=h_r) if use_paged_kv: - # Import before task bodies enter staged GPU control flow. The - # helper module imports FmhaConfig, so it cannot load in resources. - importlib.import_module(f"{__package__}.helpers_paged") if num_tokens_per_page not in _SUPPORTED_CONTEXT_PAGE_SIZES: raise ValueError( "paged context requires num_tokens_per_page in " f"{_SUPPORTED_CONTEXT_PAGE_SIZES}; got " f"{num_tokens_per_page}" ) - if max_kv_len < 1: + if max_kv_len is not None and max_kv_len < 1: raise ValueError( f"max_kv_len must be >= 1 for paged context, got {max_kv_len}" ) @@ -2850,9 +3093,9 @@ def __init__( # Validate the logical Q/K and V head dimensions before choosing tiles. d_v = d if d_v is None else d_v - if (d_v != d or d == 192) and (d, d_v) != (192, 128): + if d_v != d and (d, d_v) != (192, 128): raise ValueError("asymmetric context requires QK=192 and V=128") - if d == 192 and (head_paired or use_paged_kv): + if (d, d_v) == (192, 128) and (head_paired or use_paged_kv): raise NotImplementedError( "192/128 requires contiguous query-paired context" ) @@ -2870,6 +3113,7 @@ def __init__( or d != 128 or d_v != 128 or q_dtype.width not in (8, 16) + or q_dtype == cutlass.Int8 ): raise ValueError( "two-CTA UMMA requires the non-persistent dense contiguous " @@ -2878,26 +3122,27 @@ def __init__( cfg.two_cta_umma = two_cta_umma if two_cta_umma: cfg.cluster_shape_mn = (2, 1) + cfg.vc_attention = vc_attention + cfg.fp8_psmem_early_token = fp8_psmem_early_token self.cfg = cfg + self.runtime_dense_tail = False # Compact Q and staged O fit two resident query tiles plus the K/V # ring. Pairing overlaps the two softmax groups with MMA work. - paired_mla = d == 192 and not has_variable_window + paired_mla = (d, d_v) == (192, 128) and not has_variable_window if d > 128 and not paired_mla: cfg.num_qkv_instances = 1 cfg.use_paged_kv = use_paged_kv - single_instance_persistent = ( - is_persistent and cfg.single_qkv_instance and not head_paired - ) - # Paired Q2 and persistent D256 schedules route correction statistics - # through a compact SMEM ring and omit the per-K StatsDone - # serialization. This keeps the stats payload disjoint from the S/P - # TMEM columns while producer latency varies across storage layouts. - cfg.stats_via_smem = single_instance_persistent or (not cfg.single_qkv_instance) + # Keep correction statistics disjoint from the S/P TMEM columns for + # every scheduler. Non-persistent single-instance schedules can also + # overwrite aliased statistics before correction consumes them. + cfg.stats_via_smem = True cfg.fuse_epilogue_into_correction = cfg.single_qkv_instance cfg.num_tokens_per_page = num_tokens_per_page cfg.max_num_pages_per_seq_kv = ( - max_kv_len + num_tokens_per_page - 1 - ) // num_tokens_per_page + 0 + if max_kv_len is None + else (max_kv_len + num_tokens_per_page - 1) // num_tokens_per_page + ) cfg.causal_single_kv_tile = causal_single_kv_tile # FP16/BF16 causal attention retains the default 192/96/32 # softmax/correction/auxiliary split. Other topologies start from @@ -2908,10 +3153,15 @@ def __init__( cfg.num_regs_correction = 88 cfg.num_regs_other = 56 cfg.enable_skip_correction = enable_skip_correction - if enable_skip_correction: + # VC-Attention skips no correction. The skip keeps a stale row max, + # under which a larger score codes past VC_EXPCAST_CODE_BIAS and out of + # the E4M3 range. The exact per-tile max keeps every code at or below it. + if enable_skip_correction and not cfg.vc_attention: cfg.corr_skip_threshold_log2 = _CORR_SKIP_THRESHOLD_LOG2 cfg.uses_ldtm_stat = uses_ldtm_stat cfg.fp8_psmem_early_token = fp8_psmem_early_token + cfg.kv_ring_stages = kv_ring_stages + cfg.umma_fp8_k64 = umma_fp8_k64 cfg.qk_acc_dtype = qk_acc_dtype or cutlass.Float32 cfg.pv_acc_dtype = pv_acc_dtype or cutlass.Float32 @@ -2922,12 +3172,21 @@ def __init__( cfg.o_dtype = o_dtype cfg.head_paired = head_paired cfg.is_causal = is_causal + cfg.h_r = h_r + cfg.has_variable_window = has_variable_window balance_causal_workload = balance_causal_workload or ( is_causal and is_clc_dynamic and not cfg.single_qkv_instance ) + if head_paired and cfg.vc_attention: + raise ValueError( + "VC-Attention requires the query-paired context kernel; a positive " + "window_left selects the head-paired kernel" + ) if head_paired: - _configure_head_paired_tilers(cfg, mma_tiler_mn=mma_tiler_mn, d=d) + _configure_head_paired_tilers(cfg, mma_tiler_mn=mma_tiler_mn, d=padded_d) + cfg.pv_mma_tiler = (cfg.pv_mma_tiler[0], d_v, cfg.pv_mma_tiler[2]) + cfg.epi_tile = cfg.pv_mma_tiler[:2] _configure_head_dim_staging(cfg) _configure_pipeline_stages( cfg, is_clc_dynamic=is_clc_dynamic, is_persistent=is_persistent @@ -2961,7 +3220,12 @@ def __init__( cfg.pv_mma_tiler = (mma_tiler[0], d_v, mma_tiler[1]) cfg.epi_tile = cfg.pv_mma_tiler[:2] cfg.p_in_smem = _uses_smem_p(cfg, has_variable_window=has_variable_window) + cfg.validate_vc_profile() _configure_head_dim_staging(cfg) + # The Sage layout decides whether an sfK ring takes shared memory, so + # it is bound before the K/V ring is sized; its schedule checks follow + # the pipeline configuration below. + cfg.sage_scale_layout = sage_scale_layout _configure_pipeline_stages( cfg, is_clc_dynamic=is_clc_dynamic, is_persistent=is_persistent ) @@ -3003,13 +3267,13 @@ def __init__( cfg.tma_copy_kv_bytes = cfg.tma_copy_kv_elements * k_dtype.width // 8 # TMA copy granularity for V - v_inner_dim_size = cfg.pv_mma_tiler[1] * v_dtype.width // 8 + v_inner_dim_size = cfg.smem_v_head_dim * v_dtype.width // 8 cfg.tma_copy_v_iters = 1 if v_inner_dim_size % 128 == 0: cfg.tma_copy_v_iters = v_inner_dim_size // 128 elif v_inner_dim_size != 64 and v_inner_dim_size != 32: raise RuntimeError(f"Unsupported inner dimension size: {v_inner_dim_size}") - cfg.tma_copy_v_granu_inner = cfg.pv_mma_tiler[1] // cfg.tma_copy_v_iters + cfg.tma_copy_v_granu_inner = cfg.smem_v_head_dim // cfg.tma_copy_v_iters cfg.tma_copy_v_stage_iters = kv_head_dim // cfg.tma_copy_v_granu_inner if cfg.two_cta_umma: cfg.tma_copy_v_iters = 1 @@ -3021,8 +3285,8 @@ def __init__( cfg.tma_copy_v_bytes = cfg.tma_copy_kv_elements * v_dtype.width // 8 # TMA copy granularity for O - cfg.tma_copy_o_iters = (cfg.epi_tile[1] * o_dtype.width) // 1024 - cfg.tma_copy_o_granu_inner = cfg.epi_tile[1] // cfg.tma_copy_o_iters + cfg.tma_copy_o_iters = (cfg.smem_o_head_dim * o_dtype.width) // 1024 + cfg.tma_copy_o_granu_inner = cfg.smem_o_head_dim // cfg.tma_copy_o_iters o_head_dim = ( cfg.head_dim_per_stage_o if cfg.stage_o_by_head_dim else cfg.epi_tile[1] ) @@ -3041,6 +3305,8 @@ def __init__( window_size_left=window_size_left, has_variable_window=has_variable_window, ) + if cfg.sage_scale_layout is not None: + _check_sage_attention(cfg) # --------------------------------------------------------------------------- # Host entry point @@ -3052,8 +3318,8 @@ def __call__( k_cute: cute.Tensor, v_cute: cute.Tensor, o_cute: cute.Tensor, - scale_softmax_log2: cute.Tensor, - output_scale: cute.Tensor, + scale_softmax_log2: Float32 | cute.Tensor, + output_scale: Float32 | cute.Tensor, max_active_clusters: int, stream: cuda_drv.CUstream, cum_seqlen_q: cute.Tensor | None = None, @@ -3065,15 +3331,31 @@ def __call__( variable_window_token_starts: cute.Tensor | None = None, variable_window_token_ends: cute.Tensor | None = None, variable_window_cta_starts: cute.Tensor | None = None, + vc_mu: cute.Tensor | None = None, + vc_q_scale: cute.Tensor | None = None, + vc_k_scale: cute.Tensor | None = None, + vc_repair_tiles: int = 0, + sage_scales: cute.Tensor | None = None, + lse_output: cute.Tensor | None = None, + attention_sinks: cute.Tensor | None = None, ) -> None: """Set up TMA descriptors, compute grid, and launch the kernel. - ``scale_softmax_log2`` and ``output_scale`` must be one-element float32 - device tensors. Example host values are - ``[math.log2(math.e) / math.sqrt(d)]`` for the softmax scale and - ``[1.0]`` for FP16 output. FP8 callers may fold Q/K/V and output - quantization scales into these runtime tensors. + VC-Attention adds the tile means, the bf16 UMMA operands packed by the + host with sixteen tiles per group, and the repair tile count, a + compile-time constant giving how many outlier tiles end the K/V loop of + an outlier plan. VC-Attention-QK8 also takes the Q and K dequant scales + as ``[H, B, blocks]`` tables with one scale per 128-token block. + + ``scale_softmax_log2`` and ``output_scale`` are runtime float32 scalar + operands, normally ``log2(e) / sqrt(d)`` and ``1.0``. FP8 callers fold + scalar Q/K/V dequantization and output quantization scales into them. + + VC-Attention additionally binds tile means and a static repair-tile count. + Its output_scale is a per-channel V scale table. """ + if cutlass.const_expr(self.cfg.use_sage_attention != (sage_scales is not None)): + raise ValueError("sage_scales is required by, and only by, a Sage kernel") cfg = self.cfg if cutlass.const_expr( cfg.use_paged_kv @@ -3168,7 +3450,7 @@ def __call__( v_box_dims = ( cfg.pv_mma_tiler[2], 1, - cfg.pv_mma_tiler[1] // cfg.tma_copy_v_iters, + cfg.tma_copy_v_granu_inner, ) kv_stride_order = stride_order if cutlass.const_expr(cum_seqlen_q is not None): @@ -3288,6 +3570,25 @@ def __call__( l2_promotion=cuda.TensorMapL2Promotion.none, ) + # VC-Attention tile-mean operands, one packed tile per K/V tile. + tma_mu_desc = tma_v_desc + if cutlass.const_expr( + cfg.vc_qk8 and (vc_q_scale is None or vc_k_scale is None) + ): + raise ValueError("VC-Attention-QK8 requires vc_q_scale and vc_k_scale") + if cutlass.const_expr(cfg.vc_restores_means): + if cutlass.const_expr(vc_mu is None): + raise ValueError("VC-Attention requires vc_mu") + mu_rows, mu_cols = vc_mean_operand_shape(cfg.pv_mma_tiler[1]) + tma_mu_desc = cuda.create_tensor_map_tiled_from_view( + vc_mu, + # Two-CTA loads the half of the rows its channels need. + box_dims=(1, 1, 1, mu_rows // cfg.cta_group_size, mu_cols), + stride_order=(4, 3, 2, 1, 0), + swizzle=cuda.TensorMapSwizzle.none, + l2_promotion=cuda.TensorMapL2Promotion.none, + ) + # Compute tile scheduler and grid if cutlass.const_expr(cum_seqlen_q is None): b = o_cute.shape[0] @@ -3297,13 +3598,16 @@ def __call__( # k_cute is the page pool — its sequence axis is per-page, # not the logical max_seq_len_kv. Caller must pass it in. s_k = max_seqlen_k + h_kv = k_cute.shape[1] else: s_k = k_cute.shape[1] + h_kv = k_cute.shape[2] else: b = cute.size(cum_seqlen_q) - 1 s_q = max_seqlen_q h_q = o_cute.shape[1] s_k = max_seqlen_k + h_kv = k_cute.shape[1] num_seq_tiles = cute.ceil_div(s_q, cfg.cta_tiler[0]) num_kv_tiles = cute.ceil_div(s_k, cfg.kv_tile_n) @@ -3374,38 +3678,62 @@ def __call__( if cutlass.const_expr(seq_lens_kv is not None) else None ) - - self.kernel( - tma_q_desc, - tma_k_desc, - tma_v_desc, - tma_o_desc, - tile_sched_params, - scale_softmax_log2, - output_scale, - num_kv_tiles, - num_seq_tiles, - q_offset, - cum_seqlen_q, - cum_seqlen_k, - block_tables_iter, - block_table_row_stride, - seq_lens_kv_iter, - Int32(s_k), - variable_window_token_starts, - variable_window_token_ends, - variable_window_cta_starts, - Int32(s_q), - self.is_persistent, - self.is_clc_dynamic, - ).launch( - grid=grid, - block=[block_size, 1, 1], - cluster=cfg.cluster_shape_mn + (1,), - stream=stream, - min_blocks_per_mp=1, + sage_scales_iter = ( + sage_scales.iterator + if cutlass.const_expr(sage_scales is not None) + else None ) + def launch_kernel(capacity: cutlass.Constexpr[int]): + self.kernel( + tma_q_desc, + tma_k_desc, + tma_v_desc, + tma_o_desc, + tma_mu_desc, + vc_q_scale, + vc_k_scale, + tile_sched_params, + scale_softmax_log2, + output_scale, + num_kv_tiles, + num_seq_tiles, + q_offset, + cum_seqlen_q, + cum_seqlen_k, + block_tables_iter, + block_table_row_stride, + seq_lens_kv_iter, + Int32(s_k), + variable_window_token_starts, + variable_window_token_ends, + variable_window_cta_starts, + Int32(s_q), + sage_scales_iter, + Int32(h_kv), + vc_repair_tiles, + self.is_persistent, + self.is_clc_dynamic, + lse_output, + attention_sinks, + dense_tail_capacity=capacity, + ).launch( + grid=grid, + block=[block_size, 1, 1], + cluster=cfg.cluster_shape_mn + (1,), + stream=stream, + min_blocks_per_mp=1, + ) + + if cutlass.const_expr(self.runtime_dense_tail): + tail_rows = (Int32(s_k) - Int32(1)) % Int32(cfg.kv_tile_n) + Int32(1) + if tail_rows <= Int32(4): + launch_kernel(4) + else: + launch_kernel(0) + else: + launch_kernel(0) + # --------------------------------------------------------------------------- # Device kernel # --------------------------------------------------------------------------- @@ -3416,13 +3744,16 @@ def kernel( tma_k_desc: cutlass.GridConstant[cuda.TensorMap], tma_v_desc: cutlass.GridConstant[cuda.TensorMap], tma_o_desc: cutlass.GridConstant[cuda.TensorMap], + tma_mu_desc: cutlass.GridConstant[cuda.TensorMap], + vc_q_scale: cute.Tensor | None, + vc_k_scale: cute.Tensor | None, tile_sched_params: ( utils.PersistentTileSchedulerParams | utils.ClcDynamicPersistentTileSchedulerParams | None ), - scale_softmax_log2: cute.Tensor, - output_scale: cute.Tensor, + scale_softmax_log2: Float32 | cute.Tensor, + output_scale: Float32 | cute.Tensor, num_kv_tiles: Int32, num_seq_tiles: Int32, q_offset: Int32, @@ -3436,15 +3767,20 @@ def kernel( variable_window_token_ends: cute.Tensor | None, variable_window_cta_starts: cute.Tensor | None, variable_window_q_stride: Int32, + sage_scales: cute.Pointer | None, + num_kv_heads: Int32, + vc_repair_tiles: cutlass.Constexpr[int] = 0, is_persistent: cutlass.Constexpr[bool] = True, is_clc_dynamic: cutlass.Constexpr[bool] = False, + lse_output: cute.Tensor | None = None, + attention_sinks: cute.Tensor | None = None, + dense_tail_capacity: cutlass.Constexpr[int] = 0, ) -> None: """Warp-specialised persistent FMHA TS kernel. - ``scale_softmax_log2`` and ``output_scale`` are one-element float32 - device tensors passed through to the TS resources. The resources load - element 0 in auxiliary setup before the K/V loop, so the hot loop uses - the cached task-local values rather than reloading global memory. + ``scale_softmax_log2`` and ``output_scale`` are runtime float32 scalar + operands passed through task-local resource dataflow. + ``sage_scales`` is the prepared per-block/channel scale buffer. Structure: 1. Prefetch TMA descriptors @@ -3457,6 +3793,15 @@ def kernel( 8. TMEM deallocation """ cfg = self.cfg + if cutlass.const_expr(dense_tail_capacity > 0): + from copy import copy + + cfg = copy(cfg) + cfg.fixed_dense_k_tail = dense_tail_capacity + if cutlass.const_expr(isinstance(scale_softmax_log2, cute.Tensor)): + scale_softmax_log2 = scale_softmax_log2[0] + if cutlass.const_expr(isinstance(output_scale, cute.Tensor) and not cfg.vc_attention): + output_scale = output_scale[0] cfg.num_seq_tiles = num_seq_tiles warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) @@ -3466,6 +3811,8 @@ def kernel( prims.prefetch_tensormap(tma_k_desc.get_ptr()) prims.prefetch_tensormap(tma_v_desc.get_ptr()) prims.prefetch_tensormap(tma_o_desc.get_ptr()) + if cutlass.const_expr(cfg.vc_restores_means): + prims.prefetch_tensormap(tma_mu_desc.get_ptr()) # 2. CLC dynamic: the response buffer is declared by the builder and # bound from the unified task-manager SMEM allocation below. @@ -3500,7 +3847,15 @@ def kernel( variable_window_q_stride=variable_window_q_stride, scale_softmax_log2=scale_softmax_log2, output_scale=output_scale, + sage_scales=sage_scales, + num_kv_heads=num_kv_heads, + lse_output=lse_output, + attention_sinks=attention_sinks, q_offset=q_offset, + tma_mu_desc=tma_mu_desc.get_ptr() if cfg.vc_restores_means else None, + vc_q_scale=vc_q_scale if cfg.vc_qk8 else None, + vc_k_scale=vc_k_scale if cfg.vc_qk8 else None, + vc_repair_tiles=vc_repair_tiles, is_persistent=is_persistent, is_clc_dynamic=is_clc_dynamic, clc_response_ptr=clc_response_ptr, @@ -3590,7 +3945,17 @@ def kernel( r.tmem_addr_cached = tmem_addr_cached # 7. Execute all tasks via the TS scheduler - task_manager.run() + active_request = cutlass.Boolean(True) + if cutlass.const_expr(not is_persistent): + batch_coord = cute.arch.block_idx()[cfg.work_tile_coord_indices[2]] + if cutlass.const_expr(seq_lens_kv is not None): + active_request = Int32(seq_lens_kv[batch_coord]) > Int32(0) + elif cutlass.const_expr(cfg.has_varlen and not cfg.has_uniform_varlen): + active_request = ( + cum_seqlen_k[batch_coord + Int32(1)] > cum_seqlen_k[batch_coord] + ) + if active_request: + task_manager.run() # 8. TMEM deallocation is_softmax_warp = ( diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_resources.py index af69ee7d51bc..ba7273fcea22 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 @@ -25,6 +25,15 @@ schedule before the repeated K/V tile loop, LOOP is the repeated K/V tile body, and TAIL is the one-time cleanup and drain after LOOP exits. +VC-Attention is the low-bit attention recipe of Li et al. 2026, arXiv +2609.15810, value smoothing with a softmax cast. Names with the ``vc`` prefix +belong to its recipes, QK16 with bf16 Q and K and QK8 with centred, Hadamard +rotated E4M3 Q and K under block scales, both with E4M3 V stored as residuals +around each 128-token tile mean. The kernel restores the means inside the +online softmax as ``O += sum_t rowsum(P_t) * mu_t``, one bf16 K=16 UMMA step +per group of sixteen tiles, and casts the probabilities straight to E4M3 +codes. ``vc_resources`` holds the mean operand ring and the cast helpers. + SMEM resources (TMA pipelines) ------------------------------ - SmemQResource : SMEM Q buffer, TmaUmma pipeline. Load -> MMA. @@ -66,9 +75,10 @@ from dataclasses import dataclass, field, replace from typing import Any, Optional, TypeAlias +from ..separate_reduction import merge_log2_lse import cutlass import cutlass.cute as cute -from cutlass import Float32, Int16, Int32 +from cutlass import Float32, Int16, Int32, Int64 from ..tensor_map import transform_ragged_coords from cutlass.experimental.task_scheduling.enums import WorkAttr @@ -83,24 +93,67 @@ from cutlass.experimental.task_scheduling.resources import consumer_work, producer_work -_SUPPORTED_CONTEXT_PAGE_SIZES = (16, 32, 64, 128) +_SUPPORTED_CONTEXT_PAGE_SIZES = (4, 8, 16, 32, 64, 128) from cutlass.pipeline import PipelineAsync, PipelineState from cutlass.cutlass_dsl import Boolean, Constexpr, dsl_user_op, if_generate from ..placeholder_helpers import _placeholder_smem_array, _placeholder_tmem_ptr +from ...vc_attention import VC_MEAN_GROUP_TILES, VC_MEAN_MMA_K, VC_MEAN_OPERANDS +from ..fmha_decode.fmha_decode_resources.helpers_common import _placeholder_local_array +from .vc_attention.vc_resources import ( + E4M3_CODES_PER_OCTAVE, + VC_EXPCAST_CODE_BIAS, + VC_MEAN_TILE_LBO, + VC_MEAN_TILE_SBO, +) +from .vc_attention.vc_scales import VCKScaleTable, vc_flat_k_base, vc_flat_q_slot from ..tcgen05_compat import tcgen05_ld_32x32b_max +from ...sage.int32_scores import ( + INT32_SCORE_BIAS, + fill_sage_score_seed_tile, + issue_sage_score_seed_mma, + sage_score_seed_tile_bytes, + sage_score_seed_words, +) +from ...sage.index import load_sage_q_scale, sage_section_word, sage_sequence_words +from ...sage.layout import _FLOAT4_WORDS, _SageScaleLayout +from ...sage.modes import SageKTransport, SageScoreScaling from .helpers import ( + add_f16x2_to_f32x2, + _resolve_kv_tile_idx_context, bottom_right_window_left_bound, bottom_right_window_tile_start, + exp2_f16x2_word, + expcast_e4m3_quad_relu, + f16x2_sum4, + f16x2_to_f32x2, freeze_smem_descriptor, + pack_f16x2x2_to_fp8_e4m3, + split_bf16_hi_lo_word, variable_window_cta_min_start, ) +from .sage_scales import ( + SAGE_SCORE_BLOCK_COLS, + load_sage_k_tile, + load_sage_v_channels, + prefetch_sage_k_tile, + sage_k_block_word, + sage_k_section_address, + sage_k_tile_bytes, + sage_k_tile_words, + sage_k_transport, + sage_score_scaling, + sage_v_section_address, +) from cutlass.experimental import primitives as prims SmemDescOffsets: TypeAlias = tuple[int, int] TmemAddr: TypeAlias = int | Int32 TmemPtr: TypeAlias = cutlass.Array +# Keep masked scores finite: an all-masked tile must not form inf - inf. +_SOFTMAX_MASK_VALUE = -3.4028234663852886e38 # -FLT_MAX, bits 0xff7fffff + SoftmaxScalar: TypeAlias = Float32 SoftmaxChunk: TypeAlias = cutlass.Vector SoftmaxChunks: TypeAlias = list[SoftmaxChunk] @@ -112,6 +165,22 @@ _tmem_sp_sdata: dict[int, list] = {} # Packed fp8 P words carried from exp2_p to store_p when P is staged in SMEM. _tmem_sp_pwords: dict[int, tuple] = {} +# VC-Attention: the log2 softmax scale carried from the row-max pass to exp2_p +# and the fp32 tile row sum carried from exp2_p to store_p. +_tmem_sp_tile_scale: dict[int, Any] = {} +_tmem_sp_tile_sum: dict[int, Any] = {} + + +@cute.jit +def _softmax_exp2_bias(row_max: Float32, scale: Float32, p_scale: Float32) -> Float32: + """Use a zero exponent anchor for masked rows without changing their max.""" + return Float32( + cutlass.select_( + row_max == Float32(_SOFTMAX_MASK_VALUE), + p_scale, + (Float32(0.0) - row_max) * scale + p_scale, + ) + ) @cute.jit @@ -150,10 +219,10 @@ def _mask_score_quad( setp.ne.u32 valid2, bit, 0; and.b32 bit, {$r0}, 0x8; setp.ne.u32 valid3, bit, 0; - @!valid0 mov.b32 {$w0}, 0xff800000; - @!valid1 mov.b32 {$w1}, 0xff800000; - @!valid2 mov.b32 {$w2}, 0xff800000; - @!valid3 mov.b32 {$w3}, 0xff800000; + @!valid0 mov.b32 {$w0}, 0xff7fffff; + @!valid1 mov.b32 {$w1}, 0xff7fffff; + @!valid2 mov.b32 {$w2}, 0xff7fffff; + @!valid3 mov.b32 {$w3}, 0xff7fffff; } """, write_only_types=[Float32, Float32, Float32, Float32], @@ -222,46 +291,6 @@ def _exp2_fma_packed(x0: Float32, x1: Float32) -> tuple[Float32, Float32]: ) -@cute.jit -def _f32_bits(x: Float32) -> Int32: - return cutlass.Vector.from_elements((x,), Float32).bitcast(Int32)[0] - - -@cute.jit -def _f32_from_bits(x: Int32) -> Float32: - return cutlass.Vector.from_elements((x,), Int32).bitcast(Float32)[0] - - -@cute.jit -def _exp2_fma_packed(x0: Float32, x1: Float32) -> tuple[Float32, Float32]: - """exp2 of two fp32 values on the FMA pipe, 2^x = 2^floor(x) * poly(x - floor(x)). - Same routine as ex2_emulation_f32x2_value in the SM110 GQA decode kernel.""" - # Below -127 the result flushes to zero in e4m3 anyway. - x0 = cute.math.max(x0, Float32(-127.0), ftz=True) - x1 = cute.math.max(x1, Float32(-127.0), ftz=True) - # Adding 1.5 * 2^23 with round-down leaves floor(x) in the low mantissa bits. - bias = Float32(1.5 * 2.0**23) - t = cute.arch.add_packed_f32x2((x0, x1), (bias, bias), rnd="rm", ftz=False) - n = cute.arch.add_packed_f32x2(t, (-bias, -bias), rnd="rn", ftz=False) - f = cute.arch.fma_packed_f32x2( - n, (Float32(-1.0), Float32(-1.0)), (x0, x1), rnd="rn", ftz=False - ) - # Degree-3 fit of 2^f on [0, 1). - c1 = Float32(0.695146143436431884765625) - c2 = Float32(0.227564394474029541015625) - c3 = Float32(0.077119089663028717041015625) - p = cute.arch.fma_packed_f32x2(f, (c3, c3), (c2, c2), rnd="rn", ftz=False) - p = cute.arch.fma_packed_f32x2(p, f, (c1, c1), rnd="rn", ftz=False) - p = cute.arch.fma_packed_f32x2( - p, f, (Float32(1.0), Float32(1.0)), rnd="rn", ftz=False - ) - # The bias has zero low bits, so bits(t) << 23 is floor(x) << 23. - return ( - _f32_from_bits(_f32_bits(p[0]) + (_f32_bits(t[0]) << 23)), - _f32_from_bits(_f32_bits(p[1]) + (_f32_bits(t[1]) << 23)), - ) - - def _placeholder_softmax_chunks(cfg: Any) -> SoftmaxChunks: """Build zero P chunks with the same structure as runtime softmax chunks.""" try: @@ -276,6 +305,16 @@ def _placeholder_softmax_chunks(cfg: Any) -> SoftmaxChunks: return [] +def _placeholder_unit_vector(length: int) -> cutlass.Vector | None: + """Build a vector of ``length`` unit scales with the runtime vector's structure.""" + try: + return cutlass.Vector.from_elements( + tuple(Float32(1.0) for _ in range(length)), Float32 + ) + except RuntimeError: + return None + + # --------------------------------------------------------------------------- # FmhaConfig -- kernel-wide configuration # --------------------------------------------------------------------------- @@ -297,6 +336,8 @@ class FmhaConfig: k_dtype: type | None = None v_dtype: type | None = None o_dtype: type | None = None + # The dtype the softmax reads S as. INT8 Q/K accumulate INT32 scores onto + # ``INT32_SCORE_BIAS``, whose words are the same FP32 values. qk_acc_dtype: type | None = None pv_acc_dtype: type | None = None @@ -317,6 +358,11 @@ class FmhaConfig: # (as the TMEM-P fp8 cadence does), so the peer softmax group starts its row max # while this group computes P. fp8_psmem_early_token: bool = False + # K/V ring depth. 0 derives it from the shared-memory budget. + kv_ring_stages: int = 0 + # K=64 per fp8 PV UMMA instruction through bit 29 of the instruction + # descriptor, halving the PV issue count. + umma_fp8_k64: bool = False # One TMA pipeline has one expected-transaction byte count per stage, so K # and V share a ring of kv_stage stages only while their dtype widths # match. Mixed widths set split_kv_pipelines and size one ring per side. @@ -411,8 +457,28 @@ class FmhaConfig: # Keep the running row max while a tile raises it by at most this many log2 units. corr_skip_threshold_log2: float = 0.0 two_cta_umma: bool = False + lse_mode: str = "none" + use_attention_sinks: bool = False # exp2 pairs per 16-pair softmax chunk computed on the FMA pipe. exp2_fma_pairs: int = 0 + # Softmax exp2 as one f16x2 MUFU op per score pair. + exp2_packed_f16x2: bool = False + # VC-Attention mode on the E4M3 V kernel with 16-bit or E4M3 Q/K. V stores + # residuals around a per-tile mean and the output scale is a per-channel + # table. In "means" mode the kernel adds each tile mean back into the + # output with one UMMA step. In "outliers" mode the kernel skips the means + # and instead consumes extra K/V tiles of the largest residuals that + # follow the sequence. None turns VC-Attention off. + vc_attention: str | None = None + # Sage attention: the prepared scale buffer's layout, else None. The + # softmax dequantizes each 16-column score block by its ``sfK`` word and + # the row's ``sfQ`` word, and the epilogue scales each output column by + # ``sfV`` and adds ``v_mean``. + sage_scale_layout: _SageScaleLayout | None = None + # Stages of the load warp's ``sfK`` ring (``SageKTransport.SMEM``): the K + # tiles the shared K/V ring keeps in flight, so the ring never throttles + # the K loads. + sage_k_scale_stages: int = 4 # Variable sequence length mode stores Q/K/V/O as flattened # [sum_seqlen, head, dim] tensors and uses cum_seqlen_* for per-batch @@ -421,6 +487,11 @@ class FmhaConfig: # Uniform packed plans retain the ragged tensor-map ABI but derive their # cumulative offsets arithmetically, avoiding dependent indptr loads on # every persistent work tile. + # VC-Attention-QK8 fixed plans also set the two lengths while + # has_uniform_varlen stays False. There they are the Q tokens and the K/V + # rows of the fixed plan, the rows including the repair tiles, and they fold + # the flat scale table slots at compile time so the P store addresses stay + # out of the loop. has_uniform_varlen: bool = False uniform_seq_len_q: int = 0 uniform_seq_len_k: int = 0 @@ -449,7 +520,7 @@ class FmhaConfig: # synthetic peer0 tail slot used by the general query-paired schedule. causal_single_kv_tile: bool = False window_size_left: int = 0 - # Number of valid K/V rows in the final fixed-length dense tile. Zero + # Upper bound on valid K/V rows in the final dense tile. Zero # means that the K/V extent is tile-aligned (or that this specialization # does not use the fixed dense-tail mask). fixed_dense_k_tail: int = 0 @@ -487,9 +558,83 @@ class FmhaConfig: # a public tuning input. page_table_window_entries: int = 32 + @property + def vc_qk8(self) -> bool: + """VC-Attention-QK8, the E4M3 Q/K recipe with per-block dequant scales.""" + return self.vc_attention is not None and self.q_dtype.width == 8 + + @property + def vc_restores_means(self) -> bool: + """Whether the kernel adds the tile means back into O.""" + return self.vc_attention == "means" + + def validate_vc_profile(self) -> None: + """Validate the VC-Attention recipe against the configured kernel.""" + if self.vc_attention is None: + return + if self.vc_attention not in ("means", "outliers"): + raise ValueError(f"unknown vc_attention mode {self.vc_attention!r}") + if ( + self.is_causal + or self.has_variable_window + or self.head_paired + or self.use_paged_kv + or self.h_r != 1 + ): + raise ValueError( + "VC-Attention requires the dense contiguous query-paired context " + "kernel with equal Q and K/V head counts" + ) + # The packed mean operand and the row-sum operand are laid out for 128 + # channels and 128 query rows. + if self.logical_head_dim_qk != 128 or self.pv_mma_tiler[1] != 128: + raise ValueError("VC-Attention requires head_dim 128") + if self.q_dtype.width not in (8, 16) or self.v_dtype.width != 8: + raise ValueError("VC-Attention requires 16-bit or E4M3 Q/K and E4M3 V") + if not self.p_in_smem: + raise ValueError( + "VC-Attention requires P staged in SMEM, the dense paired D128 " + "schedule with a 16-bit output" + ) + + @property + def vc_mean_operand_bytes(self) -> int: + """Bytes of one bf16 [D x 16] K-major mean operand (B of one mean UMMA step).""" + return self.pv_mma_tiler[1] * VC_MEAN_MMA_K * 2 + + @property + def vc_mean_tile_bytes(self) -> int: + """Bytes of the mean operands of one tile group (one ring stage).""" + return VC_MEAN_OPERANDS * self.vc_mean_operand_bytes + + @property + def vc_mean_stages(self) -> int: + """Tile-mean ring depth: one operand per staged V tile.""" + return self.kv_stage_v if self.split_kv_pipelines else self.kv_stage_k + + @property + def vc_rowsum_operand_bytes(self) -> int: + """Bytes of one bf16 [128 x 16] K-major row-sum operand (A of one mean UMMA step).""" + return self.qk_mma_tiler[0] * VC_MEAN_MMA_K * 2 + + @property + def vc_rowsum_tile_bytes(self) -> int: + """Bytes of the row-sum operands of one tile group.""" + return VC_MEAN_OPERANDS * self.vc_rowsum_operand_bytes + # Work-tile mapping for the two peer Q/O tiles handled by each CTA: # query-paired maps peers to two sequence tiles in one Q head, while # head-paired mode maps peers to two Q heads at one sequence tile. + @property + def vc_kscale_table_bytes(self) -> int: + """Bytes of one softmax group's fp32 K-scale table, one word per K/V + tile of the K/V length the kernel was built with, 16-byte padded. The + table is sized at construction from ``max_kv_len``, before the fixed + plan lengths reach ``cfg``.""" + max_kv_len = self.max_num_pages_per_seq_kv * self.num_tokens_per_page + entries = max((max_kv_len + self.kv_tile_n - 1) // self.kv_tile_n, 1) + return (entries * 4 + 15) // 16 * 16 + @property def smem_q_head_dim(self) -> int: """Q storage rounded to one 128-byte TMA fragment.""" @@ -502,6 +647,18 @@ def smem_q_head_dim(self) -> int: ) return self.qk_mma_tiler[2] + @property + def smem_v_head_dim(self) -> int: + """Physical V extent in complete 128-byte TMA fragments.""" + fragment = min(1024 // self.v_dtype.width, self.pv_mma_tiler[1]) + return (self.pv_mma_tiler[1] + fragment - 1) // fragment * fragment + + @property + def smem_o_head_dim(self) -> int: + """Physical O extent; the final TMA store masks channels beyond D.""" + fragment = min(1024 // self.o_dtype.width, self.epi_tile[1]) + return (self.epi_tile[1] + fragment - 1) // fragment * fragment + @property def single_qkv_instance(self) -> bool: """Return whether one work tile carries a single Q/KV/O instance.""" @@ -605,7 +762,8 @@ def reuses_page_table_windows(self) -> bool: window_entries = self.page_table_window_entries if ( pages_per_tile <= 0 - or pages_per_tile > window_entries + # A one-tile window cannot retain distinct K/V page IDs. + or pages_per_tile >= window_entries or window_entries % pages_per_tile != 0 ): return False @@ -751,6 +909,64 @@ def work_tile_coord_indices(self) -> tuple[int, int, int]: return 2, 0, 1 return 0, 1, 2 + @property + def use_sage_attention(self) -> bool: + """Return whether the softmax and epilogue apply Sage scales.""" + return self.sage_scale_layout is not None + + @property + def sage_score_blocks(self) -> int: + """Return the 16-column score blocks of one K tile, one ``sfK`` word each.""" + return self.qk_mma_tiler[1] // SAGE_SCORE_BLOCK_COLS + + @property + def sage_k_transport(self) -> SageKTransport | None: + """Return how the softmax obtains a K tile's ``sfK`` words (``sage_k_transport``).""" + if not self.use_sage_attention: + return None + return sage_k_transport(self.sage_scale_layout) + + @property + def sage_score_scaling(self) -> SageScoreScaling | None: + """Return how the softmax applies ``sfK`` to the scores (``sage_score_scaling``).""" + if not self.use_sage_attention: + return None + return sage_score_scaling(self.sage_scale_layout) + + @property + def sage_uses_ldtm_stat(self) -> bool: + """Return whether the Sage max pass takes the block maxima from LDTM.STAT pieces.""" + return ( + self.uses_ldtm_stat + and self.sage_score_scaling is SageScoreScaling.BLOCK + and not self.uses_int32_scores + ) + + @property + def uses_int32_scores(self) -> bool: + """Return whether QK accumulates INT8 Q/K into INT32 scores on a float bias. + + A Sage recipe's INT8 Q/K accumulate onto the bit pattern of + ``INT32_SCORE_BIAS``, so the softmax reads the scores as FP32 and + removes the bias where it applies ``sfK``. + """ + return self.q_dtype == cutlass.Int8 + + @property + def sage_score_seed_tile_bytes(self) -> int: + """Return the bytes of the BF16 operand tile that seeds INT32 scores.""" + return sage_score_seed_tile_bytes(self.qk_mma_tiler[0], self.qk_mma_tiler[1]) + + @property + def sage_k_tile_words(self) -> int: + """Return the ``sfK`` words of one K tile.""" + return sage_k_tile_words(self.sage_scale_layout) + + @property + def sage_k_tile_bytes(self) -> int: + """Return the bytes of one K tile's ``sfK`` words, one SMEM ring stage.""" + return sage_k_tile_bytes(self.sage_scale_layout) + @property def pv_p_scale(self) -> float: """Return the P scale applied before PV MMA.""" @@ -1029,6 +1245,7 @@ class GmemQKVResource(MemoryResource): seqlen_k: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() kv_tile_start: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() kv_request_begin: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + kv_page_idx_lb: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() kv_page_idx_ub: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() def __init__( @@ -1123,6 +1340,11 @@ def __init__( default=Int32(0), docs="Element offset of the request's block-table row.", ) + self.kv_page_idx_lb = TaskLocalVariable( + dtype=Int32, + default=Int32(0), + docs="First logical page containing a visible sliding-window key.", + ) self.kv_page_idx_ub = TaskLocalVariable( dtype=Int32, default=Int32(0), @@ -1143,6 +1365,7 @@ def __init__( seqlen_k, kv_tile_start, kv_request_begin, + kv_page_idx_lb, kv_page_idx_ub, ) ) @@ -1163,6 +1386,7 @@ def compute_coords( Int32, Int32, Int32, + Int32, ]: """Resolve per-tile coordinates from work_tile for downstream use. @@ -1185,6 +1409,7 @@ def compute_coords( window_q_offset = Int32(self.q_offset_default) kv_tile_start = Int32(0) kv_request_begin = Int32(0) + kv_page_idx_lb = Int32(0) kv_page_idx_ub = Int32(0) if cutlass.const_expr(self.cfg.has_varlen): if cutlass.const_expr(self.cfg.has_uniform_varlen): @@ -1248,6 +1473,16 @@ def compute_coords( ) // self.cfg.seq_tile_n, ) + if cutlass.const_expr( + self.cfg.use_paged_kv and self.cfg.kv_tile_start_window_size_left > 0 + ): + first_key = cute.math.max( + Int32(0), + seq_coord * self.cfg.q_tile_m + + window_q_offset + - self.cfg.kv_tile_start_window_size_left, + ) + kv_page_idx_lb = first_key // Int32(self.cfg.num_tokens_per_page) if cutlass.const_expr(self.cfg.has_variable_window): min_window_start = variable_window_cta_min_start( self.variable_window_cta_starts, @@ -1270,6 +1505,7 @@ def compute_coords( seqlen_k, kv_tile_start, kv_request_begin, + kv_page_idx_lb, kv_page_idx_ub, ) @@ -1349,12 +1585,12 @@ def _qk_inner_dim_size_bytes(cfg: FmhaConfig) -> int: def _pv_inner_dim_size_bytes(cfg: FmhaConfig) -> int: """Return the byte width of one P/V tile inner dimension.""" - return cfg.pv_n_per_cta * cfg.v_dtype.width // 8 + return cfg.smem_v_head_dim // cfg.cta_group_size * cfg.v_dtype.width // 8 def _o_inner_dim_size_bytes(cfg: FmhaConfig) -> int: """Return the byte width of one O tile inner dimension.""" - return cfg.qk_mma_tiler[2] * cfg.o_dtype.width // 8 + return cfg.smem_o_head_dim * cfg.o_dtype.width // 8 def _smem_layout_for_inner_bytes(inner_dim_size: int) -> int: @@ -1398,7 +1634,7 @@ def _pv_smem_desc_offsets(cfg: FmhaConfig) -> SmemDescOffsets: else cfg.tma_copy_v_iters ) leading_byte_offset = cfg.tma_copy_v_bytes // tma_copy_v_iters - stride_byte_offset = cfg.pv_mma_tiler[1] * cfg.v_dtype.width // cfg.tma_copy_v_iters + stride_byte_offset = cfg.tma_copy_v_granu_inner * cfg.v_dtype.width return leading_byte_offset, stride_byte_offset @@ -1656,10 +1892,8 @@ def __init__( self._alloc = SmemAllocation( "smem_page_offsets_v" if page_table_is_v else "smem_page_offsets_k", size_bytes=total_entries * 4, - # Page-size 16 consumes eight page IDs per 128-token K/V tile. - # Align only that specialization for one 32-byte vector load; - # preserve the established layout for larger page sizes. - alignment=(32 if cfg.kv_tile_n // cfg.num_tokens_per_page == 8 else 16), + # Eight-ID vector loads also serve smaller pages in chunks. + alignment=(32 if cfg.kv_tile_n // cfg.num_tokens_per_page >= 8 else 16), ) self._smem_page_offsets = _placeholder_smem_array(Int32, total_entries) self.cached_page_ids = TaskLocalVariable( @@ -1718,14 +1952,22 @@ def page_ids(self, tile_idx: Int32) -> cutlass.Array: stage_idx = self.state_src.consumer_work_stage group_page_idx = (tile_idx * Int32(pages_per_tile)) & Int32(window_entries - 1) offset = stage_idx * Int32(window_entries) + group_page_idx - if cutlass.const_expr(pages_per_tile == 8): - return self._smem_page_offsets.load(offset, vector_size=8, alignment=32) - if cutlass.const_expr(pages_per_tile == 4): - return self._smem_page_offsets.load(offset, vector_size=4, alignment=16) - if cutlass.const_expr(pages_per_tile == 2): - return self._smem_page_offsets.load(offset, vector_size=2, alignment=8) - return self._smem_page_offsets.load(offset, vector_size=1, alignment=4) + if cutlass.const_expr(pages_per_tile <= 8): + return self._smem_page_offsets.load( + offset, vector_size=pages_per_tile, alignment=pages_per_tile * 4 + ) + tile_page_ids = cutlass.Array( + Int32, pages_per_tile, space=cutlass.AddressSpace.rmem + ) + for chunk in cutlass.range_constexpr(pages_per_tile // 8): + values = self._smem_page_offsets.load( + offset + Int32(chunk * 8), vector_size=8, alignment=32 + ) + for i in cutlass.range_constexpr(8): + tile_page_ids[chunk * 8 + i] = Int32(values[i]) + return tile_page_ids + # Stage the current K/V tile page IDs into its acquired SMEM ring slot. @cute.jit def _producer_load_page_offsets( self, @@ -1853,6 +2095,126 @@ def dma_consumer_release_labels_for( return None +# --------------------------------------------------------------------------- +# SmemSageKScalesResource -- SMEM ring of one K tile's Sage sfK words per stage +# --------------------------------------------------------------------------- + + +@dataclass(kw_only=True) +class SmemSageKScalesResource(MemoryResource): + """Ring of one K tile's ``sfK`` words per stage, staged by the load warp. + + A tile of one-token K blocks has too many words to keep in the softmax + lanes' registers or to broadcast-load from global memory + (``SageKTransport.SMEM``). After each K tile's TMA load the load warp + copies the tile's atoms from the prepared scales into the acquired stage + with lane-parallel 16-byte ``cp.async`` copies; the pipeline's + ``AsyncLoad`` commit completes the stage once they land. Both softmax + groups wait the stage before their score wait, take its shared-memory + offset from ``open_stage``, read the words after the score wait + (``TmemSPResource.read_staged_sage_k_words``) and release the stage after + their max pass: a release orders every read before it, so an earlier one + would keep all the words in registers while the scores load instead of + letting the compiler place each read next to its use. + """ + + cfg: Constexpr[FmhaConfig] = field(init=False, default=None) + sage_scales: cute.Pointer | None = field(init=False, default=None) + num_kv_heads: Int32 | None = field(init=False, default=None) + _alloc: Constexpr[SmemAllocation | None] = field(init=False, default=None) + stage_words: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + + def __init__( + self, + *, + pipeline_config: PipelineConfig, + cfg: FmhaConfig, + sage_scales: cute.Pointer | None, + num_kv_heads: Int32, + **kwargs: Any, + ) -> None: + pipeline_config = replace(pipeline_config, advance_on_wait=True) + super().__init__(pipeline_config=pipeline_config, **kwargs) + self.cfg = cfg + self.sage_scales = sage_scales + self.num_kv_heads = num_kv_heads + self._alloc = SmemAllocation( + "smem_sage_k", + size_bytes=pipeline_config.num_stages * cfg.sage_k_tile_bytes, + alignment=16, + ) + self.stage_words = TaskLocalVariable( + dtype=Int32, + default=Int32(0), + docs="Shared-memory byte offset of the waited ``sfK`` ring stage.", + ) + + def get_smem_requirements(self) -> list[SmemAllocation]: + return [self._alloc] + + @producer_work + @cute.jit + def copy_tile( + self, + stage_info: StageInfo, + *, + kv_head_coord: Int32, + batch_coord: Int32, + kv_tile_start: Int32, + ) -> None: + """Copy the current K tile's ``sfK`` words into the acquired stage. + + The tile is the one the K load of this iteration staged + (``_resolve_kv_tile_idx_context``); the K-token section is padded to + whole tiles, so the copy never runs past the sequence. Lane ``l`` + copies float4 ``l``, ``l + 32``, ... of the tile. + """ + cfg = self.cfg + layout = cfg.sage_scale_layout + tile_idx = _resolve_kv_tile_idx_context(stage_info, kv_tile_start) + sequence_words = sage_sequence_words( + layout, + batch_idx=batch_coord, + kv_head_idx=kv_head_coord, + num_kv_heads=self.num_kv_heads, + ) + tile_words = cfg.sage_k_tile_words + source_word = sage_section_word( + sequence_words, layout.k_tokens + ) + tile_idx * Int32(tile_words) + stage_bytes = Int32(self._alloc.offset) + stage_info.stage_idx * Int32( + cfg.sage_k_tile_bytes + ) + stage_ptr = cutlass.Array( + stage_info.context.smem_base.data_ptr() + stage_bytes, + dtype=Float32, + shape=(tile_words,), + addrspace=3, + ).data_ptr() + lane_idx = cute.arch.thread_idx()[0] & Int32(cute.arch.WARP_SIZE - 1) + num_float4s = tile_words // _FLOAT4_WORDS + for round_idx in cutlass.range_constexpr( + (num_float4s + cute.arch.WARP_SIZE - 1) // cute.arch.WARP_SIZE + ): + float4_idx = lane_idx + Int32(round_idx * cute.arch.WARP_SIZE) + word = float4_idx * Int32(_FLOAT4_WORDS) + if float4_idx < Int32(num_float4s): + prims.cp_async_shared_global( + stage_ptr + word, + self.sage_scales + source_word + word, + _FLOAT4_WORDS * 4, + "cg", + ) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=stage_words) + @cute.jit + 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.cfg.sage_k_tile_bytes + ) + + # --------------------------------------------------------------------------- # SmemKVResource -- SMEM K/V tile buffer with TMA pipeline # --------------------------------------------------------------------------- @@ -2005,6 +2367,7 @@ def _tma_load( seqlen_k: Int32, kv_tile_start: Int32, kv_request_begin: Int32, + kv_page_idx_lb: Int32, kv_page_idx_ub: Int32, ) -> None: """Issue TMA bulk-copy for one K or V tile.""" @@ -2070,6 +2433,14 @@ def _tma_load( ) for frag in cutlass.range_constexpr(pages_per_tile): page_id = Int32(page_ids[frag]) + if cutlass.const_expr(self.cfg.kv_tile_start_window_size_left > 0): + # Old pages may be recycled while their table IDs persist. + # Zero-fill their K/V rather than multiplying masked NaNs. + if ( + tile_idx * Int32(pages_per_tile) + Int32(frag) + < kv_page_idx_lb + ): + page_id = Int32(-1) for i in cutlass.range_constexpr(stage_iters): d_offset = Int32( head_dim_stage_idx * self.cfg.head_dim_per_stage_kv @@ -2143,6 +2514,7 @@ def k_load( seqlen_k: Int32, kv_tile_start: Int32, kv_request_begin: Int32, + kv_page_idx_lb: Int32, kv_page_idx_ub: Int32, ) -> None: """TMA load K tile from GMEM to SMEM.""" @@ -2157,6 +2529,7 @@ def k_load( seqlen_k=seqlen_k, kv_tile_start=kv_tile_start, kv_request_begin=kv_request_begin, + kv_page_idx_lb=kv_page_idx_lb, kv_page_idx_ub=kv_page_idx_ub, ) @@ -2174,6 +2547,7 @@ def k_load_stage( seqlen_k: Int32, kv_tile_start: Int32, kv_request_begin: Int32, + kv_page_idx_lb: Int32, kv_page_idx_ub: Int32, ) -> None: """TMA load one K head-dim stage for split D scheduling.""" @@ -2189,6 +2563,7 @@ def k_load_stage( seqlen_k=seqlen_k, kv_tile_start=kv_tile_start, kv_request_begin=kv_request_begin, + kv_page_idx_lb=kv_page_idx_lb, kv_page_idx_ub=kv_page_idx_ub, ) @@ -2206,6 +2581,7 @@ def v_load( seqlen_k: Int32, kv_tile_start: Int32, kv_request_begin: Int32, + kv_page_idx_lb: Int32, kv_page_idx_ub: Int32, ) -> None: """TMA load V tile from GMEM to SMEM.""" @@ -2221,6 +2597,7 @@ def v_load( seqlen_k=seqlen_k, kv_tile_start=kv_tile_start, kv_request_begin=kv_request_begin, + kv_page_idx_lb=kv_page_idx_lb, kv_page_idx_ub=kv_page_idx_ub, ) @@ -2239,6 +2616,7 @@ def v_load_stage( seqlen_k: Int32, kv_tile_start: Int32, kv_request_begin: Int32, + kv_page_idx_lb: Int32, kv_page_idx_ub: Int32, ) -> None: """TMA load one current or previous V head-dim stage.""" @@ -2254,6 +2632,7 @@ def v_load_stage( seqlen_k=seqlen_k, kv_tile_start=kv_tile_start, kv_request_begin=kv_request_begin, + kv_page_idx_lb=kv_page_idx_lb, kv_page_idx_ub=kv_page_idx_ub, ) @@ -2272,6 +2651,7 @@ def v_load_stage_cached( seqlen_k: Int32, kv_tile_start: Int32, kv_request_begin: Int32, + kv_page_idx_lb: Int32, kv_page_idx_ub: Int32, ) -> None: """Load one V head-dimension stage using register-cached page IDs.""" @@ -2288,6 +2668,7 @@ def v_load_stage_cached( seqlen_k=seqlen_k, kv_tile_start=kv_tile_start, kv_request_begin=kv_request_begin, + kv_page_idx_lb=kv_page_idx_lb, kv_page_idx_ub=kv_page_idx_ub, ) @@ -2464,13 +2845,35 @@ def __init__( ) -> None: super().__init__(pipeline_config=pipeline_config, **kwargs) self.cfg = cfg + # VC-Attention appends the group's bf16 [128 x 16] row-sum operands behind + # the P tile when it restores the tile means, and VC-Attention-QK8 its + # K-scale table behind those. + extra_bytes = 0 + if cfg.vc_restores_means: + extra_bytes += cfg.vc_rowsum_tile_bytes + if cfg.vc_qk8: + extra_bytes += cfg.vc_kscale_table_bytes self._alloc = SmemAllocation( - f"smem_p{group_idx}", cfg.smem_p_bytes, alignment=cfg.buffer_align_bytes + f"smem_p{group_idx}", + cfg.smem_p_bytes + extra_bytes, + alignment=cfg.buffer_align_bytes, ) def get_smem_requirements(self) -> list[SmemAllocation]: return [self._alloc] + @property + def rowsum_tile_offset(self) -> int: + """Byte offset of the VC-Attention row-sum operands inside this allocation.""" + return self.cfg.smem_p_bytes + + @property + def kscale_table_offset(self) -> int: + """Byte offset of the VC-Attention-QK8 K-scale table inside this allocation, + behind the row-sum operand when the tile means are restored.""" + rowsum = self.cfg.vc_rowsum_tile_bytes if self.cfg.vc_restores_means else 0 + return self.cfg.smem_p_bytes + rowsum + @property def row_bytes(self) -> int: return self.cfg.qk_mma_tiler[1] * self.cfg.v_dtype.width // 8 @@ -2508,6 +2911,9 @@ class TmemSPResource(MemoryResource): q_half: Constexpr[int] = 0 enable_early_tile_sum: Constexpr[bool] = False q_offset_default: int | Int32 = field(init=False, default=0) + max_seq_len_kv: int | Int32 = field(init=False, default=0) + # Valid K/V rows of the final tile, computed once when the lengths bind. + k_tail_rows: int | Int32 = field(init=False, default=0) cum_seqlen_q: cute.Tensor | None = field(init=False, default=None) cum_seqlen_k: cute.Tensor | None = field(init=False, default=None) seq_lens_kv: cute.Pointer | None = field(init=False, default=None) @@ -2515,7 +2921,21 @@ class TmemSPResource(MemoryResource): variable_window_token_ends: cute.Tensor | None = field(init=False, default=None) variable_window_cta_starts: cute.Tensor | None = field(init=False, default=None) variable_window_q_stride: int | Int32 = field(init=False, default=0) - scale_softmax_log2: cute.Tensor | None = field(init=False, default=None) + scale_softmax_log2: Float32 | None = field(init=False, default=None) + # Sage attention: the prepared scales and the KV head count that strides + # their sequences. The lane's ``sfQ`` word and the address of the + # sequence's first ``sfK`` atom are cached per work tile; each tile's + # ``sfK`` words pass from their load to the max pass and, when the + # softmax scales blocks, to the P pass. + sage_scales: cute.Pointer | None = field(init=False, default=None) + num_kv_heads: int | Int32 = field(init=False, default=0) + sage_q_scale: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + sage_k_first_atom: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + sage_k_words: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + # INT32 scores: the first instance owns the seed MMA's constant operand + # tile and the second reads it. + score_seed_owner: Optional["TmemSPResource"] = field(init=False, default=None) + _seed_alloc: Constexpr[Optional[SmemAllocation]] = field(init=False, default=None) tmem_addr_cached: TmemAddr | None = field(init=False, default=None) # Precomputed TMEM pointers/addresses (set by auxiliary work). Avoids # per-iteration inttoptr + address math. @@ -2541,6 +2961,25 @@ class TmemSPResource(MemoryResource): variable_window_end: Constexpr[TaskLocalVariable] = ( TaskLocalVariable.uninitialized() ) + # VC-Attention-QK8 dequant scales in the flat layout: Q [Hq, numel] and K/V + # tile [Hkv, numel], both fp32. + vc_q_scale: cute.Tensor | None = field(init=False, default=None) + vc_k_scale: cute.Tensor | None = field(init=False, default=None) + # VC-Attention: the row's softmax scale in log2 units (per work tile, times + # sfQ under QK8) and the previous tile's fp32 row sum awaiting its deferred + # mean step. + vc_row_scale: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + vc_prev_tile_sum: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + vc_kept_row_sum: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + # VC-Attention-QK8: sfK of the next K/V tile, read from the SMEM table one + # tile ahead so no shared load sits on the row-max critical chain. + vc_k_scale_next: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + # Row sums of the sixteen K/V tiles of the current mean group, in the running max's + # units. The mean restore is one bf16 K=16 UMMA per group, O += sum_t rowsum(P_t) * mu_t + # with these sums as the A operand and the tile means as B, so each tile's sum stays live + # until its group closes and is rescaled by every later max correction. One register + # array of VC_MEAN_GROUP_TILES values. + vc_meansums: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() def __init__( self, @@ -2553,17 +2992,32 @@ def __init__( cum_seqlen_q: cute.Tensor | None = None, cum_seqlen_k: cute.Tensor | None = None, seq_lens_kv: cute.Pointer | None = None, + max_seq_len_kv: int | Int32 = 0, variable_window_token_starts: cute.Tensor | None = None, variable_window_token_ends: cute.Tensor | None = None, variable_window_cta_starts: cute.Tensor | None = None, variable_window_q_stride: int | Int32 = 0, - scale_softmax_log2: cute.Tensor | None = None, + scale_softmax_log2: Float32 | None = None, smem_p: Optional[SmemPResource] = None, + vc_q_scale: cute.Tensor | None = None, + vc_k_scale: cute.Tensor | None = None, + vc_repair_tiles: int = 0, + sage_scales: cute.Pointer | None = None, + num_kv_heads: int | Int32 = 0, + score_seed_owner: Optional["TmemSPResource"] = None, **kwargs: Any, ) -> None: - """Bind S/P TMEM offsets, Q peer index, and optional varlen metadata.""" + """Bind S/P TMEM offsets, Q peer index, and optional varlen metadata. + + The repair tile count is a compile-time constant. It says how many + K/V tiles at the end of the loop are VC-Attention outlier tiles. It is + zero on plans that restore the tile means. VC-Attention-QK8 plans carry + the fixed Q and K/V lengths on ``cfg``, which address the scale tables. + """ super().__init__(pipeline_config=pipeline_config, **kwargs) self.cfg = cfg + self.vc_repair_tiles = vc_repair_tiles + self.score_seed_owner = score_seed_owner self.tmem_s_offset = tmem_s_offset self.tmem_p_offset = tmem_p_offset self.q_half = q_half @@ -2572,11 +3026,34 @@ def __init__( self.cum_seqlen_q = cum_seqlen_q self.cum_seqlen_k = cum_seqlen_k self.seq_lens_kv = seq_lens_kv + self.max_seq_len_kv = max_seq_len_kv + self.k_tail_rows = Int32(0) self.variable_window_token_starts = variable_window_token_starts self.variable_window_token_ends = variable_window_token_ends self.variable_window_cta_starts = variable_window_cta_starts self.variable_window_q_stride = variable_window_q_stride self.scale_softmax_log2 = scale_softmax_log2 + self.sage_scales = sage_scales + self.num_kv_heads = num_kv_heads + self.sage_q_scale = TaskLocalVariable( + dtype=Float32, + default=Float32(1.0), + docs="Sage ``sfQ`` of this lane's Q row, cached per work tile.", + ) + self.sage_k_first_atom = TaskLocalVariable( + dtype=Int64, + default=Int64(0), + docs="Address of the work tile's first Sage ``sfK`` atom.", + ) + self.sage_k_words = TaskLocalVariable( + dtype=cutlass.Vector, + default_factory=lambda: _placeholder_unit_vector( + cfg.sage_k_tile_words if cfg.use_sage_attention else 1 + ), + docs="Sage ``sfK`` of each word of the current K tile.", + ) + self.vc_q_scale = vc_q_scale + self.vc_k_scale = vc_k_scale self._alloc = TmemAllocation( f"tmem_sp_q{q_half}", cfg.qk_mma_tiler[1] * cfg.mma_softmax_stage, @@ -2588,12 +3065,12 @@ def __init__( self.smem_p = smem_p self.old_row_max = TaskLocalVariable( dtype=Float32, - default=Float32(-Float32.inf), + default=Float32(_SOFTMAX_MASK_VALUE), docs="Softmax row maximum from the previous K/V tile.", ) self.row_max = TaskLocalVariable( dtype=Float32, - default=Float32(-Float32.inf), + default=Float32(_SOFTMAX_MASK_VALUE), docs="Softmax row maximum for the current K/V tile.", ) self.row_sum = TaskLocalVariable( @@ -2644,11 +3121,70 @@ def __init__( default=Float32(0.0), docs="Softmax scale cached from the runtime scale tensor.", ) + self.vc_row_scale = TaskLocalVariable( + dtype=Float32, + default=Float32(0.0), + docs="VC-Attention: the softmax scale in log2 units.", + ) + self.vc_prev_tile_sum = TaskLocalVariable( + dtype=Float32, + default=Float32(0.0), + docs="VC-Attention: previous tile's row sum, stored one tile late.", + ) + self.vc_kept_row_sum = TaskLocalVariable( + dtype=Float32, + default=Float32(0.0), + docs="VC-Attention V repair, row sum over the original tokens.", + ) + self.vc_k_scale_next = TaskLocalVariable( + dtype=Float32, + default=Float32(1.0), + docs="VC-Attention-QK8: K dequant scale of the next K/V tile.", + ) + self.vc_meansums = TaskLocalVariable( + dtype=cutlass.Array, + default=_placeholder_local_array(Float32, VC_MEAN_GROUP_TILES), + docs="VC-Attention: row sums of the tiles of the current mean group, " + "in the running max's units, lanes of the K=16 mean UMMA A operand.", + ) def get_tmem_requirements(self) -> list[TmemAllocation]: """Return the TMEM allocation for this S/P ping-pong resource.""" return [self._alloc] + def get_smem_requirements(self) -> list[SmemAllocation]: + """Allocate the INT32 score seed tile; the second instance shares the first's.""" + if not self.cfg.uses_int32_scores or self.score_seed_owner is not None: + return [] + if self._seed_alloc is None: + self._seed_alloc = SmemAllocation( + f"{self.name}_score_seed", + self.cfg.sage_score_seed_tile_bytes, + alignment=128, + ) + return [self._seed_alloc] + + @cute.jit + def _score_seed_words(self, context: Any) -> cutlass.Array: + """Return the seed MMA's operand tile; the second instance reads the first's.""" + owner = self if self.score_seed_owner is None else self.score_seed_owner + return sage_score_seed_words( + context.smem_base, + owner._seed_alloc.offset, + self.cfg.sage_score_seed_tile_bytes, + ) + + @producer_work + @cute.jit + def seed_scores(self, stage_info: StageInfo) -> None: + """Write ``INT32_SCORE_BIAS`` into the acquired S stage before its K wait.""" + issue_sage_score_seed_mma( + self.tmem_ptr_s_cached.subview(self._stage_col_offset(stage_info)), + self._score_seed_words(stage_info.context), + self.cfg.qk_mma_tiler[0], + self.cfg.qk_mma_tiler[1], + ) + @property def loop_offset_sensitive(self) -> bool: """Return true because MMA and masking decisions use loop_offset.""" @@ -2827,7 +3363,13 @@ def _qk_mma_impl( self._stage_col_offset(stage_info) ) - if cutlass.const_expr(self.cfg.q_dtype.width == 8): + acc_format = ( + cutlass.Int32 if self.cfg.uses_int32_scores else cutlass.Float32 + ) + if cutlass.const_expr(self.cfg.uses_int32_scores): + mma_kind = prims.Tcgen05MMAKind.INT8 + ab_format = cutlass.Int8 + elif cutlass.const_expr(self.cfg.q_dtype.width == 8): mma_kind = prims.Tcgen05MMAKind.F8F6F4 # E4M3 operands use the Float16 encoding handle. ab_format = cutlass.Float16 @@ -2839,7 +3381,7 @@ def _qk_mma_impl( ab_format = cutlass.Float16 idesc_qk = prims.Tcgen05InstrDesc.build( - c_dtype=cutlass.Float32, + c_dtype=acc_format, a_dtype=ab_format, b_dtype=ab_format, n_dim=self.cfg.qk_mma_tiler[1], @@ -2885,9 +3427,11 @@ def _qk_mma_impl( desc_q_base_ = freeze_smem_descriptor(desc_q_base) desc_k_base_ = freeze_smem_descriptor(desc_k_base) - scale_d = False + # INT32 scores accumulate onto the seeded bias; FP32 scores + # overwrite S with the first K phase. + scale_d = self.cfg.uses_int32_scores if cutlass.const_expr(self.cfg.stage_kv_by_head_dim): - scale_d = head_dim_stage_idx != 0 + scale_d = scale_d or head_dim_stage_idx != 0 for tma_iter in cutlass.range_constexpr(num_tma_iters_qk): q_tma_iter = head_dim_stage_idx * num_tma_iters_qk + tma_iter q_tma_iter_offset = chunk_bytes_q * q_tma_iter @@ -2980,12 +3524,27 @@ def _default_p_chunk(self) -> SoftmaxRowSumContribution: @cute.jit def init_mma_state(self, stage_info: StageInfo) -> None: self._init_function_state(stage_info) + if cutlass.const_expr( + self.cfg.uses_int32_scores and self.score_seed_owner is None + ): + # The MMA task is one warp; it fills the tile it will read. + tidx, _, _ = cute.arch.thread_idx() + fill_sage_score_seed_tile( + self._score_seed_words(stage_info.context), + self.cfg.sage_score_seed_tile_bytes // 4, + tidx % Int32(32), + 32, + ) @consumer_work(work_attrs=WorkAttr.AUXILIARY) @cute.jit def init_softmax_state_early(self, stage_info: StageInfo) -> None: """Initialize softmax TMEM state without a function-lifetime P value.""" self._init_function_state(stage_info) + if cutlass.const_expr(self.cfg.vc_restores_means): + # The row-sum operand exists behind the P tile only when the tile + # means are restored; V repair plans have no such buffer. + self._zero_vc_rowsum_row(stage_info) @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=p_chunk) @cute.jit @@ -3001,57 +3560,248 @@ def load_scale_softmax_log2(self, stage_info: StageInfo) -> Float32: if cutlass.const_expr(self.scale_softmax_log2 is None): # Safe fallback for validation-only resource construction. return Float32(0.0) - return self.scale_softmax_log2[0] - - @cute.jit - def _init_work_tile_state(self, stage_info: StageInfo) -> None: - """Reset softmax state and recompute per-warp TMEM addresses each tile. - - Softmax-warp addresses are computed here (inside the persistent loop, - after setmaxnreg) to avoid spilling them across the register-budget - boundary in ungated HEAD. The returned q_offset defaults to the - uniform kernel argument; varlen causal masks overwrite it once per - work tile via cache_q_offset(). - """ - num_softmax_warps = 4 - warp_id_in_sg = cute.arch.warp_idx() % num_softmax_warps - tmem_raw_addr = self.tmem_addr_cached - tmem_base_row = tmem_raw_addr >> 16 - tmem_base_col = tmem_raw_addr & Int32(0xFFFF) - row_id = tmem_base_row + warp_id_in_sg * cute.arch.WARP_SIZE - self.tmem_s_addr_cached = (row_id << 16) | (tmem_base_col + self.tmem_s_offset) - self.tmem_p_addr_cached = (row_id << 16) | (tmem_base_col + self.tmem_p_offset) - _ = stage_info - - @producer_work(work_attrs=WorkAttr.AUXILIARY) - @cute.jit - def init_mma_work_tile_state(self, stage_info: StageInfo) -> None: - self._init_work_tile_state(stage_info) + return self.scale_softmax_log2 @consumer_work( work_attrs=WorkAttr.AUXILIARY, - returns=(old_row_max, row_max, row_sum, q_offset), + returns=( + vc_row_scale, + vc_prev_tile_sum, + vc_meansums, + vc_kept_row_sum, + vc_k_scale_next, + ), ) @cute.jit - def init_softmax_work_tile_state( - self, stage_info: StageInfo - ) -> tuple[Float32, Float32, Float32, Int32]: - self._init_work_tile_state(stage_info) - return ( - Float32(-Float32.inf), - Float32(-Float32.inf), - Float32(0.0), - Int32(self.q_offset_default), + def vc_init_row_scale(self, stage_info: StageInfo) -> tuple[Float32, ...]: + """VC-Attention work-tile setup. Returns the softmax scale in log2 units, + a zero previous-tile row sum, the zeroed group row sum array, a zero kept row sum + (V repair) and tile 0's sfK. + + Under VC-Attention-QK8 the row scale also carries this row's sfQ, and the + group's SMEM sfK table is filled with one coalesced pass. Both named + barrier syncs of the fill cover the group's four warps: the first keeps + a fast warp from refilling while a peer still reads the previous work + tile's entries, the second publishes the new table. + """ + row_scale = self.scale_softmax_log2 + k_scale_next = Float32(1.0) + if cutlass.const_expr(self.cfg.vc_qk8): + batch_coord, head, row = self._vc_coords(stage_info) + kv_head = head // self.cfg.h_r + q_slot = vc_flat_q_slot( + batch_coord, + row, + self.cfg.uniform_seq_len_q, + self.cfg.kv_tile_n.bit_length() - 1, + ) + row_scale = row_scale * Float32(self.vc_q_scale[(head, q_slot)]) + k_base = vc_flat_k_base( + batch_coord, self.cfg.uniform_seq_len_k, self.cfg.kv_tile_n + ) + num_tiles = ( + self.cfg.uniform_seq_len_k + self.cfg.kv_tile_n - 1 + ) // self.cfg.kv_tile_n + k_scale_next = self._vc_k_scales.fill( + stage_info, self.vc_k_scale, kv_head, k_base, num_tiles + ) + sums = cutlass.Array( + Float32, VC_MEAN_GROUP_TILES, space=cutlass.AddressSpace.rmem ) + for i in cutlass.range_constexpr(VC_MEAN_GROUP_TILES): + sums[i] = Float32(0.0) + return (row_scale, Float32(0.0), sums, Float32(0.0), k_scale_next) @cute.jit - def _varlen_batch_coord(self, stage_info: StageInfo) -> Int32: - """Return the batch coordinate for the active tile-order policy.""" - _, _, batch_coord = _resolve_work_tile_coords( + def _vc_coords(self, stage_info: StageInfo) -> tuple[Int32, Int32, Int32]: + """Return ``(batch, head, q_row)`` of this softmax thread's row.""" + seq_coord, head_coord, batch_coord = _resolve_work_tile_coords( self.cfg, stage_info.work_tile.tile_idx ) - return batch_coord - + warp_id_in_sg = cute.arch.warp_idx() % len(self.cfg.softmax0_warp_ids) + row_in_tile = warp_id_in_sg * cute.arch.WARP_SIZE + cute.arch.lane_idx() + row = ( + seq_coord * self.cfg.q_tile_m * self.cfg.work_tile_q_seq_tiles + + self.q_half * self.cfg.peer_q_seq_tile_stride * self.cfg.q_tile_m + + row_in_tile + ) + head = head_coord * self.cfg.work_tile_q_heads + return batch_coord, head, row + + @property + def _vc_k_scales(self) -> VCKScaleTable: + """This softmax group's staged ``sfK`` table (see ``vc_scales``).""" + return VCKScaleTable(self.cfg, self.smem_p, self.q_half) + + @consumer_work(returns=(old_row_max, row_max, vc_k_scale_next)) + @cute.jit + def vc_compute_row_max( + self, + stage_info: StageInfo, + *, + row_max: SoftmaxScalar, + vc_row_scale: SoftmaxScalar, + vc_k_scale_next: SoftmaxScalar, + ) -> tuple[SoftmaxScalar, SoftmaxScalar, Float32]: + """VC-Attention K-loop row max in log2 units. Under QK8 the tile maximum is + scaled by ``vc_row_scale * sfK[tile]`` (prefetched) and the next tile's + ``sfK`` is read.""" + tile_scale = vc_row_scale + k_scale_after = vc_k_scale_next + if cutlass.const_expr(self.cfg.vc_qk8): + k_scale_after = self._vc_k_scales.next_scale( + stage_info, Int32(stage_info.loop_offset) + ) + tile_scale = vc_row_scale * vc_k_scale_next + if cutlass.const_expr(self.cfg.uses_ldtm_stat): + old_row_max, new_row_max = self._load_s_chunks_and_reduce_row_max( + stage_info, row_max, tile_scale=tile_scale + ) + else: + s_data = self._load_s_chunks(stage_info) + old_row_max, new_row_max = self._reduce_row_max(s_data, row_max, tile_scale) + return old_row_max, new_row_max, k_scale_after + + @consumer_work(returns=(old_row_max, row_max)) + @cute.jit + def vc_fixed_dense_k_tail_masked_row_max( + self, + stage_info: StageInfo, + *, + row_max: SoftmaxScalar, + vc_row_scale: SoftmaxScalar, + vc_k_scale_next: SoftmaxScalar, + ) -> tuple[SoftmaxScalar, SoftmaxScalar]: + """VC-Attention dense tail: mask the zero-filled lanes of the partial last + K/V tile before the row max; under QK8 the tile's prefetched ``sfK`` scales it.""" + tmem_x = self.cfg.tmem_x_load_s + num_chunks = self.cfg.qk_mma_tiler[1] // tmem_x + tile_scale = vc_row_scale + if cutlass.const_expr(self.cfg.vc_qk8): + tile_scale = vc_row_scale * vc_k_scale_next + s_data = self._load_s_chunks(stage_info) + neg_inf = cutlass.vector.full( + [tmem_x], + self.cfg.qk_acc_dtype(-Float32.inf), + dtype=self.cfg.qk_acc_dtype, + ) + for chunk_idx in cutlass.range_constexpr(num_chunks): + valid_in_chunk = cute.math.min( + cute.math.max( + self.k_tail_rows - Int32(chunk_idx * tmem_x), + Int32(0), + ), + Int32(tmem_x), + ) + mask = cutlass.vector.create_mask([tmem_x], [valid_in_chunk]) + s_data[chunk_idx] = cutlass.vector.where(mask, s_data[chunk_idx], neg_inf) + return self._reduce_row_max(s_data, row_max, tile_scale) + + @cute.jit + def _init_work_tile_state(self, stage_info: StageInfo) -> None: + """Reset softmax state and recompute per-warp TMEM addresses each tile. + + Softmax-warp addresses are computed here (inside the persistent loop, + after setmaxnreg) to avoid spilling them across the register-budget + boundary in ungated HEAD. The returned q_offset defaults to the + uniform kernel argument; varlen causal masks overwrite it once per + work tile via cache_q_offset(). + """ + num_softmax_warps = 4 + warp_id_in_sg = cute.arch.warp_idx() % num_softmax_warps + tmem_raw_addr = self.tmem_addr_cached + tmem_base_row = tmem_raw_addr >> 16 + tmem_base_col = tmem_raw_addr & Int32(0xFFFF) + row_id = tmem_base_row + warp_id_in_sg * cute.arch.WARP_SIZE + self.tmem_s_addr_cached = (row_id << 16) | (tmem_base_col + self.tmem_s_offset) + self.tmem_p_addr_cached = (row_id << 16) | (tmem_base_col + self.tmem_p_offset) + _ = stage_info + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_mma_work_tile_state(self, stage_info: StageInfo) -> None: + self._init_work_tile_state(stage_info) + + @consumer_work( + work_attrs=WorkAttr.AUXILIARY, + returns=(old_row_max, row_max, row_sum, q_offset), + ) + @cute.jit + def init_softmax_work_tile_state( + self, stage_info: StageInfo + ) -> tuple[Float32, Float32, Float32, Int32]: + self._init_work_tile_state(stage_info) + if cutlass.const_expr(self.cfg.fixed_dense_k_tail > 0): + # Bind after the softmax register allocation, like its TMEM addresses. + self.k_tail_rows = (Int32(self.max_seq_len_kv) - Int32(1)) % Int32( + self.cfg.kv_tile_n + ) + Int32(1) + return ( + Float32(_SOFTMAX_MASK_VALUE), + Float32(_SOFTMAX_MASK_VALUE), + Float32(0.0), + Int32(self.q_offset_default), + ) + + @cute.jit + def _varlen_batch_coord(self, stage_info: StageInfo) -> Int32: + """Return the batch coordinate for the active tile-order policy.""" + _, _, batch_coord = _resolve_work_tile_coords( + self.cfg, stage_info.work_tile.tile_idx + ) + return batch_coord + + @consumer_work( + work_attrs=WorkAttr.AUXILIARY, + returns=(sage_q_scale, sage_k_first_atom), + ) + @cute.jit + def cache_sage_scales(self, stage_info: StageInfo) -> tuple[Float32, Int64]: + """Cache this lane's Sage ``sfQ`` word and the work tile's first ``sfK`` atom. + + The lane's Q row is row ``(warp % 4) * 32 + lane`` of token range + ``2 * seq + q_half``, at most the head's last token range, of head + ``head``, which is local head ``head % h_r`` of KV head ``head // + h_r``; ``_SageQSection.tile_index`` names + the prepared Q tile holding it. The ``sfK`` atoms are the + sequence's K-token section (``_SageKSection``), padded to whole K + tiles; the max pass addresses each tile's atoms from the first one. + """ + cfg = self.cfg + layout = cfg.sage_scale_layout + seq_coord, head_coord, batch_coord = _resolve_work_tile_coords( + cfg, stage_info.work_tile.tile_idx + ) + kv_head_coord = head_coord // Int32(cfg.h_r) + local_head_coord = head_coord - kv_head_coord * Int32(cfg.h_r) + row_idx = ( + cute.arch.warp_idx() % 4 + ) * cute.arch.WARP_SIZE + cute.arch.lane_idx() + sequence_words = sage_sequence_words( + layout, + batch_idx=batch_coord, + kv_head_idx=kv_head_coord, + num_kv_heads=self.num_kv_heads, + ) + # The second Q half of a sequence's last work tile can start past its + # last Q tile; its rows are past the sequence and their output is + # dropped, so it reads the head's last tile instead of the next one's. + token_tile = cute.math.min( + seq_coord * Int32(cfg.work_tile_q_seq_tiles) + Int32(self.q_half), + Int32(layout.q_rows.token_tiles - 1), + ) + sage_q_scale = load_sage_q_scale( + layout, + self.sage_scales, + sequence_words=sequence_words, + q_tile_idx=layout.q_rows.tile_index(local_head_coord, token_tile), + row_idx=row_idx, + ) + first_atom = sage_k_section_address( + layout, self.sage_scales, sequence_words=sequence_words + ) + return sage_q_scale, first_atom + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=q_offset) @cute.jit def cache_q_offset(self, stage_info: StageInfo) -> Int32: @@ -3152,17 +3902,170 @@ def _load_s_chunks(self, stage_info: StageInfo) -> SoftmaxChunks: s_data[chunk_idx] = s_data[chunk_idx][0:tmem_x] return s_data + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=sage_k_words) + @cute.jit + def load_sage_k_words( + self, + stage_info: StageInfo, + *, + sage_k_first_atom: Int64, + section: cutlass.Constexpr[FmhaStage] = FmhaStage.Loop, + ) -> cutlass.Vector: + """Load the current K tile's ``sfK`` words and warm L1 for the next tile. + + The global transport (``SageKTransport.GMEM``), scheduled before + the wait for the tile's scores: the wait and the TMEM load then hide + the load latency, which the compiler would otherwise expose by + sinking the loads next to their first use. ``section=Tail`` loads the + tile the schedule materializes after the loop, which has no next tile. + """ + layout = self.cfg.sage_scale_layout + if cutlass.const_expr(section == FmhaStage.Loop): + kv_tile_idx = stage_info.loop_offset + else: + kv_tile_idx = stage_info.loop_end + tile_bytes = Int64(sage_k_tile_bytes(layout)) + tile_address = sage_k_first_atom + Int64(kv_tile_idx) * tile_bytes + words = load_sage_k_tile(layout, tile_address) + if cutlass.const_expr(section == FmhaStage.Loop): + prefetch_sage_k_tile(layout, tile_address + tile_bytes) + values: tuple[Any, ...] = () + for word_idx in cutlass.range_constexpr(self.cfg.sage_k_tile_words): + values += (Float32(words[word_idx]),) + return cutlass.Vector.from_elements(values, Float32) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=sage_k_words) + @cute.jit + def read_staged_sage_k_words( + self, stage_info: StageInfo, *, sage_k_stage_words: Int32 + ) -> cutlass.Vector: + """Read the current K tile's ``sfK`` words from its waited ring stage. + + The ring transport (``SageKTransport.SMEM``, protocol on + ``SmemSageKScalesResource``): ``sage_k_stage_words`` is the stage's + shared-memory byte offset; one float4 load per four words. + """ + tile_words = self.cfg.sage_k_tile_words + stage = cutlass.Array( + stage_info.context.smem_base.data_ptr() + sage_k_stage_words, + dtype=Float32, + shape=(tile_words,), + addrspace=3, + ) + values: tuple[Any, ...] = () + for base in cutlass.range_constexpr(0, tile_words, _FLOAT4_WORDS): + loaded = stage.load(base, vector_size=_FLOAT4_WORDS, alignment=16) + for elem in cutlass.range_constexpr(_FLOAT4_WORDS): + values += (Float32(loaded[elem]),) + return cutlass.Vector.from_elements(values, Float32) + + @cute.jit + def _sage_block_scales(self, sage_k_words: cutlass.Vector) -> cutlass.Vector: + """Return ``sfK`` of each 16-column score block of the tile, from its words.""" + layout = self.cfg.sage_scale_layout + values: tuple[Any, ...] = () + for block_idx in cutlass.range_constexpr(self.cfg.sage_score_blocks): + values += (Float32(sage_k_words[sage_k_block_word(layout, block_idx)]),) + return cutlass.Vector.from_elements(values, Float32) + + @cute.jit + def _sage_dequantized_tile_max( + self, + s_data: SoftmaxChunks, + *, + sage_q_scale: Float32, + sage_k_block_scales: cutlass.Vector, + ) -> Float32: + """Return the row's maximum dequantized score of the tile. + + Each 16-column block's raw maximum is scaled by its ``sfK`` word, the + scaled block maxima are reduced, and the result is scaled by ``sfQ``: + one multiply per block, as positive scales commute with the maximum. + The row max therefore stays in the dequantized units + ``scale_softmax_log2`` expects, so the lazy correction and the + statistics are unchanged by Sage. Both reductions are vector + reductions, which lower to plain FMNMX chains; a scalar ``max`` per + block would add NaN-handling compare-and-select pairs. + """ + tmem_x = self.cfg.tmem_x_load_s + blocks_per_chunk = tmem_x // SAGE_SCORE_BLOCK_COLS + block_maxima: tuple[Any, ...] = () + for block_idx in cutlass.range_constexpr(self.cfg.sage_score_blocks): + chunk_idx = block_idx // blocks_per_chunk + first_elem = (block_idx % blocks_per_chunk) * SAGE_SCORE_BLOCK_COLS + block_values: tuple[Any, ...] = () + for elem_idx in cutlass.range_constexpr( + first_elem, first_elem + SAGE_SCORE_BLOCK_COLS + ): + block_values += (s_data[chunk_idx][elem_idx],) + block_maxima += ( + cutlass.Vector.from_elements( + block_values, self.cfg.qk_acc_dtype + ).reduce("max"), + ) + return self._sage_scale_block_maxima( + block_maxima, + sage_q_scale=sage_q_scale, + sage_k_block_scales=sage_k_block_scales, + ) + + @cute.jit + def _sage_scale_block_maxima( + self, + block_maxima: tuple, + *, + sage_q_scale: Float32, + sage_k_block_scales: cutlass.Vector, + ) -> Float32: + """Return the tile's dequantized maximum from its raw 16-column block maxima.""" + scaled_block_maxima: tuple[Any, ...] = () + for block_idx in cutlass.range_constexpr(self.cfg.sage_score_blocks): + block_max = block_maxima[block_idx] + if cutlass.const_expr(self.cfg.uses_int32_scores): + # Exact: biased scores have unit spacing. + block_max = block_max - Float32(INT32_SCORE_BIAS) + scaled_block_maxima += ( + block_max * Float32(sage_k_block_scales[block_idx]), + ) + tile_max = cutlass.Vector.from_elements( + scaled_block_maxima, self.cfg.qk_acc_dtype + ).reduce("max") + return tile_max * sage_q_scale + @cute.jit def _reduce_row_max( self, s_data: SoftmaxChunks, row_max: SoftmaxScalar, + tile_scale: SoftmaxScalar | None = None, + *, + sage_q_scale: Float32 | None = None, + sage_k_block_scales: cutlass.Vector | None = None, ) -> tuple[SoftmaxScalar, SoftmaxScalar]: - """Reduce per-chunk maximums into row_max, stash s_data.""" + """Reduce per-chunk maximums into row_max, stash s_data. + + VC-Attention passes ``tile_scale``: the raw tile maximum is scaled into + log2 units (the scale is positive, so max commutes with it) before the + running maximum merge, and the scale is stashed for ``exp2_p``. + + With block-scaled Sage attention the maximum is of the dequantized + scores (``_sage_dequantized_tile_max``) and the stashed scores stay + raw. A Sage plan that multiplied its scores by ``sfK`` in place + passes ``sage_q_scale`` alone: the tile maximum takes ``sfQ`` once, + so the row max stays in dequantized units while the stashed scores + carry ``sfK`` only. + """ tmem_x = self.cfg.tmem_x_load_s num_chunks = self.cfg.qk_mma_tiler[1] // tmem_x old_row_max = row_max - if cutlass.const_expr( + if cutlass.const_expr(sage_k_block_scales is not None): + tile_row_max = self._sage_dequantized_tile_max( + s_data, + sage_q_scale=sage_q_scale, + sage_k_block_scales=sage_k_block_scales, + ) + row_max = cute.math.max(row_max, tile_row_max) + elif cutlass.const_expr( self.cfg.uses_d128_fp8_softmax_cadence or self.cfg.uses_d256_fp8_softmax_cadence ): @@ -3192,12 +4095,14 @@ def _reduce_row_max( row_values += (s_data[chunk_idx][elem_idx],) row_vector = cutlass.Vector.from_elements(row_values, self.cfg.qk_acc_dtype) tile_row_max = row_vector.reduce("max") + if cutlass.const_expr(self.cfg.vc_attention): + tile_row_max = tile_row_max * tile_scale + _tmem_sp_tile_scale[id(self)] = tile_scale + if cutlass.const_expr(sage_q_scale is not None): + tile_row_max = tile_row_max * sage_q_scale row_max = cute.math.max(row_max, tile_row_max) _tmem_sp_sdata[id(self)] = s_data - row_max_safe = row_max - if row_max == -Float32.inf: - row_max_safe = Float32(0.0) - return old_row_max, row_max_safe + return old_row_max, row_max @cute.jit def _load_s_chunks_and_reduce_row_max( @@ -3207,6 +4112,7 @@ def _load_s_chunks_and_reduce_row_max( *, causal_loop_mask: cutlass.Constexpr[bool] = False, q_offset: Int32 = 0, + tile_scale: SoftmaxScalar | None = None, ) -> tuple[SoftmaxScalar, SoftmaxScalar]: """LDTM.STAT: fuse S load and per-chunk row_max via tcgen05.ld.red.max. @@ -3215,6 +4121,10 @@ def _load_s_chunks_and_reduce_row_max( reduction happens as part of the TMEM load. For causal loop masking, only tiles crossing the runtime right bound need a software reduction of masked scores; wholly visible tiles retain the hardware maximum. + + VC-Attention keeps the row max in log2 units. ``tile_scale`` is + the softmax scale, applied to the hardware tile maximum before the + merge and carried to ``exp2_p`` like the loaded S chunks. """ tmem_s_addr = self.tmem_s_addr_cached + self._stage_col_offset(stage_info) tmem_x = self.cfg.tmem_x_load_s @@ -3233,6 +4143,9 @@ def _load_s_chunks_and_reduce_row_max( tile_max = cutlass.Vector.from_elements( chunk_maxima, self.cfg.qk_acc_dtype ).reduce("max") + if cutlass.const_expr(self.cfg.vc_attention): + tile_max = tile_max * tile_scale + _tmem_sp_tile_scale[id(self)] = tile_scale row_max = cute.math.max(row_max, tile_max) if cutlass.const_expr(causal_loop_mask): seq_coord, _, _ = _resolve_work_tile_coords( @@ -3262,21 +4175,18 @@ def _load_s_chunks_and_reduce_row_max( Int32(tmem_x), ) mask = cutlass.vector.create_mask([tmem_x], [num_valid]) - neg_inf = cutlass.vector.full_like( - s_data[chunk_idx], Float32(-Float32.inf) + masked_values = cutlass.vector.full_like( + s_data[chunk_idx], Float32(_SOFTMAX_MASK_VALUE) ) s_data[chunk_idx] = cutlass.vector.where( - mask, s_data[chunk_idx], neg_inf + mask, s_data[chunk_idx], masked_values ) masked_row_max = cute.math.max( masked_row_max, s_data[chunk_idx].reduce("max") ) row_max = masked_row_max _tmem_sp_sdata[id(self)] = s_data - row_max_safe = row_max - if row_max == -Float32.inf: - row_max_safe = Float32(0.0) - return old_row_max, row_max_safe + return old_row_max, row_max @cute.jit def _exp2_p_store( @@ -3287,12 +4197,24 @@ def _exp2_p_store( chunk_lo: cutlass.Constexpr[int] = 0, chunk_hi: cutlass.Constexpr[int | None] = None, sum_in: Float32 | None = None, + sage_q_scale: Float32 | None = None, + sage_k_block_scales: cutlass.Vector | None = None, ) -> SoftmaxRowSumContribution: """Apply exp2 softmax P, fold the PV P scale, and store P to TMEM. ``chunk_lo``/``chunk_hi`` select the key chunks to process; the whole tile by default. A partial range fences its stores and returns ``sum_in`` plus its row sum, for the two-half publish. + + With block-scaled Sage attention each 16-column block's raw scores + take the fused multiplier ``sfQ * sfK_block * scale_softmax_log2`` in + place of the scalar scale, so the exp2 argument stays one FMA per + score; the bias is unchanged because the row max is already + dequantized. Each block's multiplier is formed where it is first + used, so no per-block array stays live across the pass. A plan that + multiplied its scores by ``sfK`` in place passes ``sage_q_scale`` + alone: the exp2 scale takes ``sfQ`` while the bias keeps the + dequantized row max. """ tmem_p_addr = self.tmem_p_addr_cached + stage_col_offset tmem_shape = "32x32b" @@ -3303,6 +4225,13 @@ def _exp2_p_store( partial = chunk_lo > 0 or chunk_hi < num_chunks p_packing_ratio = self.cfg.qk_acc_dtype.width // self.cfg.v_dtype.width scale = scale_softmax_log2 + if cutlass.const_expr(self.cfg.vc_attention): + # The row-max pass kept row_max in log2 units of the scaled scores. + scale = _tmem_sp_tile_scale.pop(id(self)) + block_scaled = sage_k_block_scales is not None + blocks_per_chunk = tmem_x // SAGE_SCORE_BLOCK_COLS + if cutlass.const_expr(block_scaled): + sage_q_scale_log2 = sage_q_scale * scale_softmax_log2 if cutlass.const_expr(not partial and self.cfg.uses_d256_fp8_softmax_cadence): return self._exp2_p_store_d256_fp8_cadence( tmem_p_addr, @@ -3322,12 +4251,35 @@ def _exp2_p_store( dtype=self.cfg.v_dtype, ) p_scale_log2 = Float32(self.cfg.pv_p_scale_log2) - minus_row_max_scale = (Float32(0.0) - row_max) * scale + p_scale_log2 + if cutlass.const_expr(self.cfg.vc_attention): + # The paper's 2^8 P scaling is part of VC_EXPCAST_CODE_BIAS, so + # p_scale_log2 is not added here. + return self._exp2_p_store_expcast(scale, Float32(0.0) - row_max) + minus_row_max_scale = _softmax_exp2_bias(row_max, scale, p_scale_log2) + # The exp2 addend of a score pair; a block-scaled INT32 plan folds + # the score bias into it per block. + addend_pair = (minus_row_max_scale, minus_row_max_scale) + if cutlass.const_expr(sage_q_scale is not None and not block_scaled): + scale = scale_softmax_log2 * sage_q_scale if cutlass.const_expr(chunk_hi < num_chunks): # A later range reads the same S data. s_data = _tmem_sp_sdata[id(self)] else: s_data = _tmem_sp_sdata.pop(id(self)) + if cutlass.const_expr( + self.cfg.exp2_packed_f16x2 + and not block_scaled + and (self.enable_early_tile_sum or partial) + ): + return self._exp2_p_store_packed_words( + tmem_p_addr, + scale, + minus_row_max_scale, + s_data, + chunk_lo, + chunk_hi, + sum_in, + ) if cutlass.const_expr(self.enable_early_tile_sum or partial): # Four packed float2 accumulators keep eight independent FADD2 # chains, so no add waits on the MUFU result it consumes. @@ -3338,13 +4290,30 @@ def _exp2_p_store( for chunk_idx in cutlass.range_constexpr(chunk_lo, chunk_hi): p_vals = () for elem_idx in cutlass.range_constexpr(0, tmem_x, 2): + if cutlass.const_expr( + block_scaled and elem_idx % SAGE_SCORE_BLOCK_COLS == 0 + ): + block_idx = ( + chunk_idx * blocks_per_chunk + elem_idx // SAGE_SCORE_BLOCK_COLS + ) + scale = Float32(sage_k_block_scales[block_idx]) * sage_q_scale_log2 + if cutlass.const_expr(self.cfg.uses_int32_scores): + # One rounding per block instead of one per score; it + # stays below half a quantized score step. + addend_pair = cute.arch.fma_packed_f32x2( + (Float32(-INT32_SCORE_BIAS), Float32(-INT32_SCORE_BIAS)), + (scale, scale), + (minus_row_max_scale, minus_row_max_scale), + rnd="rn", + ftz=False, + ) fma_pair = cute.arch.fma_packed_f32x2( ( s_data[chunk_idx][elem_idx], s_data[chunk_idx][elem_idx + 1], ), (scale, scale), - (minus_row_max_scale, minus_row_max_scale), + addend_pair, rnd="rn", ftz=False, ) @@ -3469,6 +4438,145 @@ def _exp2_p_store( result.append(s_data[chunk_idx]) return result + @cute.jit + def _exp2_p_store_packed_words( + self, + tmem_p_addr: TmemAddr, + scale: SoftmaxScalar, + minus_row_max_scale: Float32, + s_data: list, + chunk_lo: cutlass.Constexpr[int], + chunk_hi: cutlass.Constexpr[int], + sum_in: Float32 | None, + ) -> Float32: + """exp2 softmax P in f16x2 words with one MUFU op per score pair. The row + sum adds the packed halves and the fp8 P is packed from the words. Returns + the tile sum.""" + assert self.cfg.v_dtype.width == 8 + tmem_x = self.cfg.tmem_x_load_s + num_chunks = self.cfg.qk_mma_tiler[1] // tmem_x + partial = chunk_lo > 0 or chunk_hi < num_chunks + # Four accumulator pairs keep eight independent sum chains, each seeded + # by its first pair. + sums: list[Any] = [None] * 4 + p_words: dict[int, tuple[Any, ...]] = {} + for chunk_idx in cutlass.range_constexpr(chunk_lo, chunk_hi): + words: tuple[Any, ...] = () + for elem_idx in cutlass.range_constexpr(0, tmem_x, 2): + fma_pair = cute.arch.fma_packed_f32x2( + ( + s_data[chunk_idx][elem_idx], + s_data[chunk_idx][elem_idx + 1], + ), + (scale, scale), + (minus_row_max_scale, minus_row_max_scale), + rnd="rn", + ftz=False, + ) + word = exp2_f16x2_word(fma_pair[0], fma_pair[1]) + pair_idx = (chunk_idx - chunk_lo) * (tmem_x // 2) + elem_idx // 2 + if cutlass.const_expr(pair_idx < 4): + sums[pair_idx] = f16x2_to_f32x2(word) + else: + sums[pair_idx % 4] = add_f16x2_to_f32x2(sums[pair_idx % 4], word) + words += (word,) + p_words[chunk_idx] = words + p_packing_ratio = self.cfg.qk_acc_dtype.width // self.cfg.v_dtype.width + for pair_idx in cutlass.range_constexpr( + chunk_lo // p_packing_ratio, + (chunk_hi + p_packing_ratio - 1) // p_packing_ratio, + ): + word_lo = max(chunk_lo - pair_idx * p_packing_ratio, 0) * tmem_x // 4 + word_hi = ( + min(chunk_hi - pair_idx * p_packing_ratio, p_packing_ratio) + * tmem_x + // 4 + ) + frag_words: tuple[Any, ...] = () + for word_idx in cutlass.range_constexpr(word_lo, word_hi): + flat_idx = word_idx * 4 + chunk_words = p_words[pair_idx * p_packing_ratio + flat_idx // tmem_x] + half_idx = (flat_idx % tmem_x) // 2 + frag_words += ( + pack_f16x2x2_to_fp8_e4m3( + chunk_words[half_idx], chunk_words[half_idx + 1] + ), + ) + if cutlass.const_expr(self.cfg.p_in_smem): + _tmem_sp_pwords[id(self)] = frag_words + else: + prims.tcgen05_st( + "32x32b", + prims.make_tmem_ptr( + tmem_p_addr + pair_idx * tmem_x + word_lo, cutlass.Int8 + ), + cutlass.Vector.from_elements(frag_words, Int32), + ) + sum_02 = cute.arch.add_packed_f32x2(sums[0], sums[2], rnd="rn", ftz=False) + sum_13 = cute.arch.add_packed_f32x2(sums[1], sums[3], rnd="rn", ftz=False) + local_sum_pair = cute.arch.add_packed_f32x2(sum_02, sum_13, rnd="rn", ftz=False) + tile_sum = local_sum_pair[0] + local_sum_pair[1] + if cutlass.const_expr(partial): + # Make this range's P visible before the barrier arrive that follows. + cute.arch.fence_view_async_tmem_store() + return sum_in + tile_sum + if cutlass.const_expr( + self.enable_early_tile_sum or self.cfg.has_tmem_p_pipeline + ): + cute.arch.fence_view_async_tmem_store() + else: + prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) + return tile_sum + + @cute.jit + def _exp2_p_store_expcast( + self, scale: SoftmaxScalar, minus_row_max_scale: SoftmaxScalar + ) -> Float32: + """VC-Attention ExpCast-FP8 P with one packed FFMA per element. + + ``code = 8 * (s * scale - m) + VC_EXPCAST_CODE_BIAS`` + + Packed P words go to ``store_p``. Returns the tile row sum of the + quantized P. + """ + assert self.cfg.p_in_smem and self.enable_early_tile_sum + tmem_x = self.cfg.tmem_x_load_s + num_chunks = self.cfg.qk_mma_tiler[1] // tmem_x + # The exact per-tile row max keeps u <= 0, so codes stay at or below + # the code bias. + code_scale = scale * Float32(E4M3_CODES_PER_OCTAVE) + code_bias = minus_row_max_scale * Float32(E4M3_CODES_PER_OCTAVE) + Float32( + VC_EXPCAST_CODE_BIAS + ) + s_data = _tmem_sp_sdata.pop(id(self)) + packed_words: tuple[Any, ...] = () + acc: list[Any] = [Int32(0)] * 4 + for chunk_idx in cutlass.range_constexpr(num_chunks): + for elem_idx in cutlass.range_constexpr(0, tmem_x, 4): + quad_idx = (chunk_idx * tmem_x + elem_idx) // 4 + chain = 2 * (quad_idx % 2) + c0, c1 = cute.arch.fma_packed_f32x2( + (s_data[chunk_idx][elem_idx], s_data[chunk_idx][elem_idx + 1]), + (code_scale, code_scale), + (code_bias, code_bias), + ) + c2, c3 = cute.arch.fma_packed_f32x2( + ( + s_data[chunk_idx][elem_idx + 2], + s_data[chunk_idx][elem_idx + 3], + ), + (code_scale, code_scale), + (code_bias, code_bias), + ) + word, acc_lo, acc_hi = expcast_e4m3_quad_relu( + c0, c1, c2, c3, acc[chain], acc[chain + 1], quad_idx < 2 + ) + acc[chain] = acc_lo + acc[chain + 1] = acc_hi + packed_words += (word,) + _tmem_sp_pwords[id(self)] = packed_words + return f16x2_sum4(acc[0], acc[1], acc[2], acc[3]) + @cute.jit def _store_p_row_smem( self, stage_info: StageInfo, packed_words: tuple[Any, ...] @@ -3498,6 +4606,71 @@ def _store_p_row_smem( chunk_words, alignment=16 ) + @cute.jit + def _vc_rowsum_view(self, stage_info: StageInfo) -> cutlass.Array: + """This group's bf16 row-sum operand: K-major 32-byte rows in 8-row core + matrices (``k`` 0-7 in the first 128 bytes of a group, 8-15 in the next).""" + assert self.smem_p is not None + context = stage_info.context + return cutlass.Array( + context.smem_base.data_ptr() + + self.smem_p._alloc.offset + + self.smem_p.rowsum_tile_offset, + dtype=cutlass.BFloat16, + shape=(self.cfg.vc_rowsum_tile_bytes // 2,), + addrspace=3, + ) + + @cute.jit + def _vc_row_base(self) -> Int32: + """Element offset of this thread's row inside the row-sum operand.""" + warp_id_in_sg = cute.arch.warp_idx() % len(self.cfg.softmax0_warp_ids) + row = Int32(warp_id_in_sg * cute.arch.WARP_SIZE + cute.arch.lane_idx()) + return (row >> Int32(3)) * Int32(128) + (row & Int32(7)) * Int32(8) + + @cute.jit + def _zero_vc_rowsum_row(self, stage_info: StageInfo) -> None: + """Clear this thread's 32-byte row of every row-sum operand once.""" + view = self._vc_rowsum_view(stage_info) + base = self._vc_row_base() + zero = cutlass.BFloat16(0.0) + zeros = cutlass.Vector.from_elements((zero,) * 8, cutlass.BFloat16) + for operand in cutlass.range_constexpr(VC_MEAN_OPERANDS): + row = base + Int32(operand * self.cfg.vc_rowsum_operand_bytes // 2) + view.subview(row).data_ptr().store(zeros, alignment=16) + view.subview(row + Int32(64)).data_ptr().store(zeros, alignment=16) + + @cute.jit + def _store_vc_rowsum_group( + self, stage_info: StageInfo, sums: cutlass.Array + ) -> None: + """Write the group's row sums as bf16 hi/lo pairs: tile ``8o+i`` in K slots + ``2i``, ``2i+1`` of operand ``o``, as one 16-byte store per core matrix + (slots 0-7 in the first, 8-15 in the next). + + The split truncates and the low half carries the remainder, so each + pair is exact to 2^-16 and runs on the integer pipe. + """ + view = cutlass.Array( + self._vc_rowsum_view(stage_info).data_ptr(), + dtype=Int32, + shape=(self.cfg.vc_rowsum_tile_bytes // 4,), + addrspace=3, + ) + base = self._vc_row_base() >> Int32(1) + tiles_per_operand = VC_MEAN_MMA_K // 2 + for operand in cutlass.range_constexpr(VC_MEAN_OPERANDS): + words = [ + split_bf16_hi_lo_word(sums[operand * tiles_per_operand + i]) + for i in range(tiles_per_operand) + ] + row = base + Int32(operand * self.cfg.vc_rowsum_operand_bytes // 4) + for half in cutlass.range_constexpr(2): + vec = cutlass.Vector.from_elements( + tuple(words[half * 4 : half * 4 + 4]), Int32 + ) + view.subview(row + Int32(half * 32)).data_ptr().store(vec, alignment=16) + @cute.jit def _exp2_p_store_d128_fp8_cadence( self, @@ -3518,7 +4691,7 @@ def _exp2_p_store_d128_fp8_cadence( store_group_words = 4 words_per_chunk = 2 * store_group_words p_scale_log2 = Float32(self.cfg.pv_p_scale_log2) - minus_row_max_scale = (Float32(0.0) - row_max) * scale + p_scale_log2 + minus_row_max_scale = _softmax_exp2_bias(row_max, scale, p_scale_log2) s_data = _tmem_sp_sdata.pop(id(self)) local_sum_chains = cutlass.Array( Float32, @@ -3679,7 +4852,7 @@ def _exp2_p_store_d256_fp8_cadence( tmem_x = self.cfg.tmem_x_load_s num_values = self.cfg.qk_mma_tiler[1] p_scale_log2 = Float32(self.cfg.pv_p_scale_log2) - minus_row_max_scale = (Float32(0.0) - row_max) * scale + p_scale_log2 + minus_row_max_scale = _softmax_exp2_bias(row_max, scale, p_scale_log2) s_data = _tmem_sp_sdata.pop(id(self)) fma_values: tuple[Any, ...] = () @@ -3773,78 +4946,278 @@ def _exp2_p_store_d256_fp8_cadence( rnd="rn", ftz=False, ) - - if cutlass.const_expr(use_two_sum_pairs): - local_sum_pair_0 = cute.arch.add_packed_f32x2( - local_sum_pair_0, - local_sum_pair_1, - rnd="rn", - ftz=False, - ) - - store_fragment = cutlass.Vector.from_elements(packed_words, Int32) - prims.tcgen05_st( - "32x32b", - prims.make_tmem_ptr(tmem_p_addr, cutlass.Int8), - store_fragment, - ) - cute.arch.fence_view_async_tmem_store() - return local_sum_pair_0[0] + local_sum_pair_0[1] + + if cutlass.const_expr(use_two_sum_pairs): + local_sum_pair_0 = cute.arch.add_packed_f32x2( + local_sum_pair_0, + local_sum_pair_1, + rnd="rn", + ftz=False, + ) + + store_fragment = cutlass.Vector.from_elements(packed_words, Int32) + prims.tcgen05_st( + "32x32b", + prims.make_tmem_ptr(tmem_p_addr, cutlass.Int8), + store_fragment, + ) + cute.arch.fence_view_async_tmem_store() + return local_sum_pair_0[0] + local_sum_pair_0[1] + + @consumer_work(returns=(old_row_max, row_max)) + @cute.jit + def compute_row_max( + self, + stage_info: StageInfo, + *, + row_max: SoftmaxScalar, + ) -> tuple[SoftmaxScalar, SoftmaxScalar]: + """Main K-loop stage: load S from TMEM and compute unmasked row_max.""" + if cutlass.const_expr(self.cfg.uses_ldtm_stat): + return self._load_s_chunks_and_reduce_row_max(stage_info, row_max) + s_data = self._load_s_chunks(stage_info) + return self._reduce_row_max(s_data, row_max) + + @consumer_work(returns=(old_row_max, row_max)) + @cute.jit + def sage_row_max( + self, + stage_info: StageInfo, + *, + row_max: SoftmaxScalar, + sage_q_scale: Float32, + sage_k_words: cutlass.Vector, + section: cutlass.Constexpr[FmhaStage] = FmhaStage.Loop, + ) -> tuple[SoftmaxScalar, SoftmaxScalar]: + """Sage row max of block-scaled scores (``SageScoreScaling.BLOCK``). + + ``sage_q_scale`` is the work tile's cached ``sfQ`` (``cache_sage_scales``) + and ``sage_k_words`` the tile's ``sfK`` words, which the P pass reuses; + each 16-column block takes its word's scale. The maximum is of the + dequantized scores; the stashed scores stay raw. ``section=Tail`` + masks the zero-filled lanes of the dense tail tile. + """ + sage_k_block_scales = self._sage_block_scales(sage_k_words) + if cutlass.const_expr( + section == FmhaStage.Loop and self.cfg.sage_uses_ldtm_stat + ): + return self._sage_load_s_chunks_and_reduce_row_max( + stage_info, + row_max, + sage_q_scale=sage_q_scale, + sage_k_block_scales=sage_k_block_scales, + ) + s_data = self._load_s_chunks(stage_info) + if cutlass.const_expr(section == FmhaStage.Tail): + s_data = self._mask_fixed_dense_k_tail(stage_info, s_data, section=section) + return self._reduce_row_max( + s_data, + row_max, + sage_q_scale=sage_q_scale, + sage_k_block_scales=sage_k_block_scales, + ) + + @cute.jit + def _sage_load_s_chunks_and_reduce_row_max( + self, + stage_info: StageInfo, + row_max: SoftmaxScalar, + *, + sage_q_scale: Float32, + sage_k_block_scales: cutlass.Vector, + ) -> tuple[SoftmaxScalar, SoftmaxScalar]: + """LDTM.STAT for Sage: load S in 16-column pieces with their maxima. + + ``tcgen05.ld.red.max`` over each 16-column block returns the block's + raw maximum with its scores, so the per-block scaling of + ``_sage_dequantized_tile_max`` costs no software reduction. The + pieces of one 32-column chunk are joined into the chunk vector the P + pass expects. + """ + tmem_s_addr = self.tmem_s_addr_cached + self._stage_col_offset(stage_info) + tmem_x = self.cfg.tmem_x_load_s + num_chunks = self.cfg.qk_mma_tiler[1] // tmem_x + blocks_per_chunk = tmem_x // SAGE_SCORE_BLOCK_COLS + old_row_max = row_max + s_data: SoftmaxChunks = [None] * num_chunks + block_maxima: tuple[Any, ...] = () + for chunk_idx in cutlass.range_constexpr(num_chunks): + chunk_values: tuple[Any, ...] = () + for block_in_chunk in cutlass.range_constexpr(blocks_per_chunk): + loaded = tcgen05_ld_32x32b_max( + tmem_s_addr + + chunk_idx * tmem_x + + block_in_chunk * SAGE_SCORE_BLOCK_COLS, + SAGE_SCORE_BLOCK_COLS, + ) + chunk_values += loaded[:SAGE_SCORE_BLOCK_COLS] + block_maxima += (loaded[SAGE_SCORE_BLOCK_COLS],) + s_data[chunk_idx] = cutlass.Vector.from_elements( + chunk_values, dtype=self.cfg.qk_acc_dtype + ) + cute.arch.fence_view_async_tmem_load() + tile_max = self._sage_scale_block_maxima( + block_maxima, + sage_q_scale=sage_q_scale, + sage_k_block_scales=sage_k_block_scales, + ) + row_max = cute.math.max(row_max, tile_max) + _tmem_sp_sdata[id(self)] = s_data + row_max_safe = row_max + if row_max == -Float32.inf: + row_max_safe = Float32(0.0) + return old_row_max, row_max_safe + + @cute.jit + def _dequantize_s_chunks( + self, s_data: SoftmaxChunks, sage_k_words: cutlass.Vector + ) -> SoftmaxChunks: + """Multiply the loaded scores by ``sfK`` in place (``SageScoreScaling.DEQUANTIZED``). + + ``sage_k_words`` holds the tile's ``sfK`` words; an atom's words are + unpadded at these block sizes, so word ``w`` scales the columns + ``w * tokens_per_word`` onward. The scores take the words in packed + pairs. ``sfQ`` is not applied here: the max pass folds it into the + tile maximum and the P pass into the exp2 scale, which saves one + multiply per word; the lazy correction sees dequantized row maxima. + INT32 scores drop their bias here: words that cover several tokens + fold ``-bias * sfK`` into one FMA per score pair (one rounding per + word, below half a quantized score step), one-token words subtract + it first, exactly, which is cheaper than a product per word. + """ + cfg = self.cfg + tmem_x = cfg.tmem_x_load_s + num_chunks = cfg.qk_mma_tiler[1] // tmem_x + tokens_per_word = cfg.sage_scale_layout.k_tokens.tokens_per_word + folds_bias = cfg.uses_int32_scores and tokens_per_word > 1 + subtracts_bias = cfg.uses_int32_scores and not folds_bias + if cutlass.const_expr(cfg.uses_int32_scores): + minus_bias_pair = (Float32(-INT32_SCORE_BIAS), Float32(-INT32_SCORE_BIAS)) + if cutlass.const_expr(folds_bias): + bias_words: tuple[Any, ...] = () + for word_idx in cutlass.range_constexpr(0, cfg.sage_k_tile_words, 2): + bias_words += cute.arch.mul_packed_f32x2( + ( + Float32(sage_k_words[word_idx]), + Float32(sage_k_words[word_idx + 1]), + ), + minus_bias_pair, + rnd="rn", + ftz=False, + ) + for chunk_idx in cutlass.range_constexpr(num_chunks): + scaled: tuple[Any, ...] = () + for elem_idx in cutlass.range_constexpr(0, tmem_x, 2): + col = chunk_idx * tmem_x + elem_idx + words = (col // tokens_per_word, (col + 1) // tokens_per_word) + score_pair = ( + s_data[chunk_idx][elem_idx], + s_data[chunk_idx][elem_idx + 1], + ) + word_pair = ( + Float32(sage_k_words[words[0]]), + Float32(sage_k_words[words[1]]), + ) + if cutlass.const_expr(folds_bias): + scaled += cute.arch.fma_packed_f32x2( + score_pair, + word_pair, + (bias_words[words[0]], bias_words[words[1]]), + rnd="rn", + ftz=False, + ) + else: + if cutlass.const_expr(subtracts_bias): + score_pair = cute.arch.add_packed_f32x2( + score_pair, minus_bias_pair, rnd="rn", ftz=False + ) + scaled += cute.arch.mul_packed_f32x2( + score_pair, word_pair, rnd="rn", ftz=False + ) + s_data[chunk_idx] = cutlass.Vector.from_elements(scaled, cfg.qk_acc_dtype) + return s_data @consumer_work(returns=(old_row_max, row_max)) @cute.jit - def compute_row_max( + def sage_dequantized_row_max( self, stage_info: StageInfo, *, row_max: SoftmaxScalar, + sage_q_scale: Float32, + sage_k_words: cutlass.Vector, + section: cutlass.Constexpr[FmhaStage] = FmhaStage.Loop, ) -> tuple[SoftmaxScalar, SoftmaxScalar]: - """Main K-loop stage: load S from TMEM and compute unmasked row_max.""" - if cutlass.const_expr(self.cfg.uses_ldtm_stat): - return self._load_s_chunks_and_reduce_row_max(stage_info, row_max) + """Row max of scores multiplied by ``sfK`` in place (``SageScoreScaling.DEQUANTIZED``). + + ``sage_k_words`` is the tile's ``sfK`` per word, from either + transport; ``sfQ`` scales the tile maximum. ``section=Tail`` masks + the zero-filled lanes of the dense tail tile. + """ s_data = self._load_s_chunks(stage_info) - return self._reduce_row_max(s_data, row_max) + s_data = self._dequantize_s_chunks(s_data, sage_k_words) + if cutlass.const_expr(section == FmhaStage.Tail): + s_data = self._mask_fixed_dense_k_tail(stage_info, s_data, section=section) + return self._reduce_row_max(s_data, row_max, sage_q_scale=sage_q_scale) - @consumer_work(returns=(old_row_max, row_max)) @cute.jit - def fixed_dense_k_tail_masked_row_max( + def _mask_fixed_dense_k_tail( self, stage_info: StageInfo, + s_data: SoftmaxChunks, *, - row_max: SoftmaxScalar, - section: cutlass.Constexpr[FmhaStage] = FmhaStage.Loop, - ) -> tuple[SoftmaxScalar, SoftmaxScalar]: - """Exclude TMA zero-fill lanes in a partial fixed dense K/V tile. + section: cutlass.Constexpr[FmhaStage], + ) -> SoftmaxChunks: + """Replace the TMA zero-fill lanes of the partial last K/V tile by -inf. - ``section=Tail``: called once for the last tile, mask always applied. - ``section=Loop``: used in loop, masks only on the last iteration. + ``section=Tail``: the last tile, masked unconditionally. + ``section=Loop``: masked only on the last iteration. """ tmem_x = self.cfg.tmem_x_load_s num_chunks = self.cfg.qk_mma_tiler[1] // tmem_x - s_data = self._load_s_chunks(stage_info) - if cutlass.const_expr(section == FmhaStage.Tail): is_tail_tile = True else: is_tail_tile = stage_info.loop_offset == stage_info.loop_end - Int32(1) if is_tail_tile: - neg_inf = cutlass.vector.full( + masked_values = cutlass.vector.full( [tmem_x], - self.cfg.qk_acc_dtype(-Float32.inf), + self.cfg.qk_acc_dtype(_SOFTMAX_MASK_VALUE), dtype=self.cfg.qk_acc_dtype, ) for chunk_idx in cutlass.range_constexpr(num_chunks): - valid_in_chunk = cute.math.min( - cute.math.max( - Int32(self.cfg.fixed_dense_k_tail) - Int32(chunk_idx * tmem_x), - Int32(0), - ), - Int32(tmem_x), - ) - mask = cutlass.vector.create_mask([tmem_x], [valid_in_chunk]) - s_data[chunk_idx] = cutlass.vector.where( - mask, s_data[chunk_idx], neg_inf + capacity = min( + tmem_x, max(0, self.cfg.fixed_dense_k_tail - chunk_idx * tmem_x) ) + if cutlass.const_expr(capacity == 0): + s_data[chunk_idx] = masked_values + else: + bits = () + for column in cutlass.range_constexpr(tmem_x): + if cutlass.const_expr(column < capacity): + bits += ( + self.k_tail_rows > Int32(chunk_idx * tmem_x + column), + ) + else: + bits += (cutlass.Boolean(False),) + mask = cutlass.Vector.from_elements(bits, cutlass.Boolean) + s_data[chunk_idx] = cutlass.vector.where( + mask, s_data[chunk_idx], masked_values + ) + return s_data + + @consumer_work(returns=(old_row_max, row_max)) + @cute.jit + def fixed_dense_k_tail_masked_row_max( + self, + stage_info: StageInfo, + *, + row_max: SoftmaxScalar, + section: cutlass.Constexpr[FmhaStage] = FmhaStage.Loop, + ) -> tuple[SoftmaxScalar, SoftmaxScalar]: + """Exclude TMA zero-fill lanes in a partial fixed dense K/V tile.""" + s_data = self._load_s_chunks(stage_info) + s_data = self._mask_fixed_dense_k_tail(stage_info, s_data, section=section) return self._reduce_row_max(s_data, row_max) @consumer_work(returns=(old_row_max, row_max)) @@ -3869,9 +5242,9 @@ def packed_dense_k_masked_row_max( kv_base = kv_tile_idx * self.cfg.kv_tile_n kv_end = kv_base + Int32(self.cfg.kv_tile_n) if seqlen_k < kv_end: - neg_inf = cutlass.vector.full( + masked_values = cutlass.vector.full( [tmem_x], - self.cfg.qk_acc_dtype(-Float32.inf), + self.cfg.qk_acc_dtype(_SOFTMAX_MASK_VALUE), dtype=self.cfg.qk_acc_dtype, ) for chunk_idx in cutlass.range_constexpr(num_chunks): @@ -3882,7 +5255,7 @@ def packed_dense_k_masked_row_max( ) mask = cutlass.vector.create_mask([tmem_x], [valid_in_chunk]) s_data[chunk_idx] = cutlass.vector.where( - mask, s_data[chunk_idx], neg_inf + mask, s_data[chunk_idx], masked_values ) return self._reduce_row_max(s_data, row_max) @@ -3985,8 +5358,10 @@ def left_masked_row_max( ) kv_tile_abs = kv_tile_start + stage_info.loop_offset - neg_inf = cutlass.vector.full( - [tmem_x], self.cfg.qk_acc_dtype(-Float32.inf), dtype=self.cfg.qk_acc_dtype + masked_values = cutlass.vector.full( + [tmem_x], + self.cfg.qk_acc_dtype(_SOFTMAX_MASK_VALUE), + dtype=self.cfg.qk_acc_dtype, ) all_true_mask = cutlass.vector.create_mask([tmem_x], [tmem_x]) @@ -4008,7 +5383,9 @@ def left_masked_row_max( ) right_mask = cutlass.vector.create_mask([tmem_x], [right_oob_start_idx]) mask = mask & right_mask - s_data[chunk_idx] = cutlass.vector.where(mask, s_data[chunk_idx], neg_inf) + s_data[chunk_idx] = cutlass.vector.where( + mask, s_data[chunk_idx], masked_values + ) return self._reduce_row_max(s_data, row_max) @@ -4044,6 +5421,114 @@ def store_p(self, stage_info: StageInfo) -> None: space=prims.SharedSpace.shared_cta, ) + @consumer_work( + returns=( + vc_prev_tile_sum, + row_sum, + vc_meansums, + ) + ) + @cute.jit + def vc_store_p( + self, + stage_info: StageInfo, + *, + old_row_max: SoftmaxScalar, + row_max: SoftmaxScalar, + row_sum: SoftmaxScalar, + vc_prev_tile_sum: SoftmaxScalar, + p_chunk: SoftmaxRowSumContribution, + vc_meansums: cutlass.Array, + is_tail: cutlass.Constexpr[bool] = False, + ) -> tuple[Float32, ...]: + """VC-Attention P store. Writes the P row, rescales the group row sums + by this tile's max correction and slots in the previous tile's sum. Every + 8th tile writes the completed group operand under the same P-ready + handoff. Returns this tile's sum, the updated row sum and the group sums. + """ + self._store_p_row_smem(stage_info, _tmem_sp_pwords.pop(id(self))) + acc_scale = cute.math.exp2(old_row_max - row_max, fastmath=True) + # The peeled tail entry carries the loop offset of the previous tile. + if cutlass.const_expr(is_tail): + tile = Int32(stage_info.loop_end) + else: + tile = Int32(stage_info.loop_offset) + slot = (tile - Int32(1)) & Int32(VC_MEAN_GROUP_TILES - 1) + new_sum = vc_prev_tile_sum * acc_scale + for i in cutlass.range_constexpr(VC_MEAN_GROUP_TILES): + scaled = vc_meansums[i] * acc_scale + if slot == Int32(i): + scaled = new_sum + vc_meansums[i] = scaled + keep = Float32(1.0) + if ((tile & Int32(VC_MEAN_GROUP_TILES - 1)) == Int32(0)) & (tile > Int32(0)): + self._store_vc_rowsum_group(stage_info, vc_meansums) + keep = Float32(0.0) + prims.fence_proxy( + kind=prims.Proxy.ASYNC_SHARED, + space=prims.SharedSpace.shared_cta, + ) + for i in cutlass.range_constexpr(VC_MEAN_GROUP_TILES): + vc_meansums[i] = vc_meansums[i] * keep + return (p_chunk, row_sum * acc_scale + p_chunk, vc_meansums) + + @consumer_work(returns=(vc_kept_row_sum, row_sum)) + @cute.jit + def vc_store_p_repair( + self, + stage_info: StageInfo, + *, + old_row_max: SoftmaxScalar, + row_max: SoftmaxScalar, + row_sum: SoftmaxScalar, + vc_kept_row_sum: SoftmaxScalar, + p_chunk: SoftmaxRowSumContribution, + is_tail: cutlass.Constexpr[bool] = False, + ) -> tuple[Float32, Float32]: + """VC-Attention V repair P store. The last ``vc_repair_tiles`` of the + loop domain stay out of the kept row sum that normalizes the output.""" + self._store_p_row_smem(stage_info, _tmem_sp_pwords.pop(id(self))) + prims.fence_proxy( + kind=prims.Proxy.ASYNC_SHARED, + space=prims.SharedSpace.shared_cta, + ) + acc_scale = cute.math.exp2(old_row_max - row_max, fastmath=True) + kept_chunk = p_chunk + if cutlass.const_expr(not is_tail): + first_repair = Int32(stage_info.loop_end) - Int32(self.vc_repair_tiles) + if Int32(stage_info.loop_offset) >= first_repair: + kept_chunk = Float32(0.0) + return ( + vc_kept_row_sum * acc_scale + kept_chunk, + row_sum * acc_scale + p_chunk, + ) + + @consumer_work + @cute.jit + def vc_store_rowsum_final( + self, + stage_info: StageInfo, + *, + vc_prev_tile_sum: SoftmaxScalar, + vc_meansums: cutlass.Array, + ) -> None: + """VC-Attention: slot the last tile's row sum into its group and + store the group operands for the tail mean step.""" + last_tile = Int32(stage_info.loop_end) + if cutlass.const_expr(not self.uses_fixed_dense_k_tail_mask): + last_tile = last_tile - Int32(1) + slot = last_tile & Int32(VC_MEAN_GROUP_TILES - 1) + for i in cutlass.range_constexpr(VC_MEAN_GROUP_TILES): + value = vc_meansums[i] + if slot == Int32(i): + value = vc_prev_tile_sum + vc_meansums[i] = value + self._store_vc_rowsum_group(stage_info, vc_meansums) + prims.fence_proxy( + kind=prims.Proxy.ASYNC_SHARED, + space=prims.SharedSpace.shared_cta, + ) + @consumer_work(returns=p_chunk) @cute.jit def exp2_p( @@ -4072,6 +5557,45 @@ def exp2_p_smem( self._stage_col_offset(stage_info), row_max, scale_softmax_log2 ) + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=p_chunk) + @cute.jit + def sage_dequantized_exp2_p_smem( + self, + stage_info: StageInfo, + *, + row_max: SoftmaxScalar, + scale_softmax_log2: SoftmaxScalar, + sage_q_scale: Float32, + ) -> SoftmaxRowSumContribution: + """``exp2_p_smem`` on scores that carry ``sfK`` only: ``sfQ`` joins the exp2 scale.""" + return self._exp2_p_store( + self._stage_col_offset(stage_info), + row_max, + scale_softmax_log2, + sage_q_scale=sage_q_scale, + ) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=p_chunk) + @cute.jit + def sage_exp2_p_smem( + self, + stage_info: StageInfo, + *, + row_max: SoftmaxScalar, + scale_softmax_log2: SoftmaxScalar, + sage_q_scale: Float32, + sage_k_words: cutlass.Vector, + ) -> SoftmaxRowSumContribution: + """``exp2_p_smem`` of block-scaled scores: the tile's ``sfK`` words of the + max pass (``sage_row_max``) give the fused per-block multipliers.""" + return self._exp2_p_store( + self._stage_col_offset(stage_info), + row_max, + scale_softmax_log2, + sage_q_scale=sage_q_scale, + sage_k_block_scales=self._sage_block_scales(sage_k_words), + ) + @consumer_work(returns=p_lo) @cute.jit def exp2_p_lo( @@ -4184,13 +5708,13 @@ def right_masked_row_max( all_true_mask = cutlass.vector.create_mask([tmem_x], [tmem_x]) left_mask = left_mask_inverted ^ all_true_mask mask = mask & left_mask - neg_inf = cutlass.vector.full( + masked_values = cutlass.vector.full( [tmem_x], - self.cfg.qk_acc_dtype(-Float32.inf), + self.cfg.qk_acc_dtype(_SOFTMAX_MASK_VALUE), dtype=self.cfg.qk_acc_dtype, ) s_data[chunk_idx] = cutlass.vector.where( - mask, s_data[chunk_idx], neg_inf + mask, s_data[chunk_idx], masked_values ) return self._reduce_row_max(s_data, row_max) @@ -4233,11 +5757,11 @@ def _apply_causal_mask_for_kv_tile( Int32(tmem_x), ) causal_mask = cutlass.vector.create_mask([tmem_x], [num_valid]) - neg_inf_vec = cutlass.vector.full_like( - s_data[chunk_idx], Float32(-Float32.inf) + masked_values_vec = cutlass.vector.full_like( + s_data[chunk_idx], Float32(_SOFTMAX_MASK_VALUE) ) s_data[chunk_idx] = cutlass.vector.where( - causal_mask, s_data[chunk_idx], neg_inf_vec + causal_mask, s_data[chunk_idx], masked_values_vec ) return s_data @@ -4305,11 +5829,8 @@ def invalid_row_max( row_max: SoftmaxScalar, ) -> tuple[SoftmaxScalar, SoftmaxScalar]: """Tail stage for softmax group 0: identity row_max, no S load.""" - row_max_safe = row_max - if row_max == -Float32.inf: - row_max_safe = Float32(0.0) _ = stage_info - return row_max, row_max_safe + return row_max, row_max @consumer_work @cute.jit @@ -4336,6 +5857,9 @@ def softmax_aux_reduce( ) -> SoftmaxScalar: """Accumulate row_sum from vector P fragments or their scalar sum.""" _ = stage_info + if cutlass.const_expr(self.cfg.vc_attention): + # VC-Attention keeps the row max in log2 units. + scale_softmax_log2 = Float32(1.0) if cutlass.const_expr(self.enable_early_tile_sum): acc_scale = cute.math.exp2( scale_softmax_log2 * (old_row_max - row_max), @@ -4376,7 +5900,12 @@ def freeze_row_max( _ = stage_info threshold = Float32(self.cfg.corr_skip_threshold_log2) frozen = row_max - if (row_max - old_row_max) * scale_softmax_log2 <= threshold: + if cutlass.const_expr(self.cfg.vc_attention): + # VC-Attention keeps the row max in log2 units. + scale_softmax_log2 = Float32(1.0) + if (old_row_max != Float32(_SOFTMAX_MASK_VALUE)) & ( + (row_max - old_row_max) * scale_softmax_log2 <= threshold + ): frozen = old_row_max return frozen @@ -4522,15 +6051,20 @@ class TmemStatsResource(MemoryResource): """Correction statistics with an AsyncAsync pipeline. Producer: Softmax writes old_max/row_max/row_sum stats. Consumer: - Correction reads them for O rescaling. Persistent D256 keeps the payload - in a small staged SMEM ring so the stats no longer alias S/P TMEM columns; - other schedules retain the original TMEM storage. + Correction reads them for O rescaling. The staged SMEM ring keeps the + payload disjoint from S/P TMEM columns, including while QK and correction + overlap in a non-persistent single-instance schedule. """ cfg: Constexpr[FmhaConfig] = field(init=False, default=None) tmem_vec_offset: Constexpr[int] = field(init=False, default=None) - scale_softmax_log2: cute.Tensor | None = field(init=False, default=None) - output_scale: cute.Tensor | None = field(init=False, default=None) + scale_softmax_log2: Float32 | None = field(init=False, default=None) + output_scale: Float32 | cute.Tensor | None = field(init=False, default=None) + lse_output: cute.Tensor | None = field(init=False, default=None) + attention_sinks: cute.Tensor | None = field(init=False, default=None) + cum_seqlen_q: cute.Tensor | None = field(init=False, default=None) + query_seq_len: Int32 | int = field(init=False, default=0) + peer_idx: Constexpr[int] = field(init=False, default=0) tmem_addr_cached: TmemAddr | None = field(init=False, default=None) # Precomputed per-warp TMEM vec address (once, before persistent loop). tmem_vec_addr_cached: TmemAddr | None = field(init=False, default=None) @@ -4548,8 +6082,13 @@ def __init__( pipeline_config: PipelineConfig, cfg: FmhaConfig, tmem_vec_offset: int, - scale_softmax_log2: cute.Tensor | None = None, - output_scale: cute.Tensor | None = None, + scale_softmax_log2: Float32 | None = None, + output_scale: Float32 | cute.Tensor | None = None, + lse_output: cute.Tensor | None = None, + attention_sinks: cute.Tensor | None = None, + cum_seqlen_q: cute.Tensor | None = None, + query_seq_len: Int32 | int = 0, + peer_idx: int = 0, **kwargs: Any, ) -> None: """Bind the correction-stat TMEM vector offset and allocation.""" @@ -4558,6 +6097,11 @@ def __init__( self.tmem_vec_offset = tmem_vec_offset self.scale_softmax_log2 = scale_softmax_log2 self.output_scale = output_scale + self.lse_output = lse_output + self.attention_sinks = attention_sinks + self.cum_seqlen_q = cum_seqlen_q + self.query_seq_len = query_seq_len + self.peer_idx = peer_idx self._alloc = TmemAllocation(f"tmem_vec_{tmem_vec_offset}", cfg.tmem_stats_cols) self._smem_alloc = None if cfg.stats_via_smem: @@ -4650,10 +6194,13 @@ def init_read_state(self, stage_info: StageInfo) -> None: def load_scale_softmax_log2(self, stage_info: StageInfo) -> Float32: """Load the runtime softmax scale once before the correction loop.""" _ = stage_info + if cutlass.const_expr(self.cfg.vc_attention): + # VC-Attention softmax publishes row maxima already in log2 units. + return Float32(1.0) if cutlass.const_expr(self.scale_softmax_log2 is None): # Safe fallback for validation-only resource construction. return Float32(0.0) - return self.scale_softmax_log2[0] + return self.scale_softmax_log2 @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns="output_scale_value") @cute.jit @@ -4663,7 +6210,11 @@ def load_output_scale(self, stage_info: StageInfo) -> Float32: if cutlass.const_expr(self.output_scale is None): # Identity fallback for validation-only resource construction. return Float32(1.0) - return self.output_scale[0] + if cutlass.const_expr(self.cfg.vc_attention): + # VC-Attention: output_scale is the per-(batch, head, channel) V + # residual scale table, applied per column in SmemOResource._store_o. + return Float32(1.0) + return self.output_scale @cute.jit def _init_work_tile_state(self, stage_info: StageInfo) -> None: @@ -4831,9 +6382,63 @@ def _read_vec( else: vec_row_sum = vec_rmem[0] vec_old_max = vec_new_max - _ = stage_info + if cutlass.const_expr( + final_stats + and (self.lse_output is not None or self.attention_sinks is not None) + ): + scale *= self._store_final_lse( + stage_info, vec_new_max, vec_row_sum, scale_softmax_log2 + ) return vec_old_max, vec_new_max, vec_row_sum, scale + @cute.jit + def _store_final_lse(self, stage_info, row_max, row_sum, scale_softmax_log2): + # Match GmemO's logical row mapping, including paired heads/queries and + # cluster tiles. Correction has one thread for each physical Q row. + seq_tile, head_tile, batch = _resolve_work_tile_coords( + self.cfg, stage_info.work_tile.tile_idx + ) + head = ( + head_tile * self.cfg.work_tile_q_heads + + self.peer_idx * self.cfg.peer_q_head_stride + ) + row = cute.arch.thread_idx()[0] % Int32(128) + query = ( + seq_tile * self.cfg.work_tile_q_seq_tiles + + self.peer_idx * self.cfg.peer_q_seq_tile_stride + ) * self.cfg.q_tile_m + row + heads = ( + Int32(self.lse_output.shape[-1]) + if self.lse_output is not None + else Int32(self.attention_sinks.shape[0]) + ) + weight = Float32(1.0) + start = Int32(0) + length = Int32(self.query_seq_len) + if cutlass.const_expr(self.cum_seqlen_q is not None): + start = Int32(self.cum_seqlen_q[batch]) + length = Int32(self.cum_seqlen_q[batch + Int32(1)]) - start + if query < length and head < heads: + flat_query = batch * length + query + if cutlass.const_expr(self.cum_seqlen_q is not None): + flat_query = start + query + value = Float32(-Float32.inf) + if row_sum > Float32(0.0): + value = ( + scale_softmax_log2 * row_max + + cute.math.log2(row_sum, fastmath=True) + - Float32(self.cfg.pv_p_scale_log2) + ) + if cutlass.const_expr(self.attention_sinks is not None): + sink = Float32(self.attention_sinks[head]) * Float32(1.4426950408889634) + value, weight, _ = merge_log2_lse(value, sink) + if cutlass.const_expr(self.lse_output is not None): + offset = cutlass.Int64(flat_query) * cutlass.Int64( + heads + ) + cutlass.Int64(head) + (self.lse_output.iterator + offset).store(value) + return weight + @consumer_work(returns=(vec_old_max, vec_new_max, vec_row_sum, vec_scale)) @cute.jit def read_vec( @@ -4870,6 +6475,16 @@ class TmemOResource(MemoryResource): tmem_addr_cached: TmemAddr | None = field(init=False, default=None) # Precomputed TMEM raw pointer (inttoptr of tmem_addr_cached). tmem_ptr_raw_cached: TmemPtr | None = field(init=False, default=None) + # VC-Attention mean step state, set once per kernel in the prologue. + # Whether this CTA issues the mean UMMA. Under two-CTA UMMA only the leader + # CTA of the pair issues, so this is block_idx_in_cluster() == 0 made warp + # uniform; otherwise it is True. + vc_mean_issue_cached: Any | None = field(init=False, default=None) + # SMEM descriptors of the two row-sum A operands of the mean step, one per + # UMMA step of a tile group. + vc_desc_l_cached: Any | None = field(init=False, default=None) + # UMMA instruction descriptor of the bf16 K=16 mean step. + vc_idesc_mu_cached: Any | None = field(init=False, default=None) # Precomputed per-warp TMEM O address base: (row_id << 16) | tmem_base_col. # consumer_work adds tmem_o_offset to get the final O0/O1 address. tmem_o_addr_base_cached: TmemAddr | None = field(init=False, default=None) @@ -4909,15 +6524,19 @@ def __init__( self.tmem_o1_offset = tmem_o1_offset self.tmem_vec0_resource = tmem_vec0_resource self.tmem_vec1_resource = tmem_vec1_resource - self._alloc_o0 = TmemAllocation("tmem_o0", 128) - self._alloc_o1 = TmemAllocation("tmem_o1", 128) + self._alloc_o0 = TmemAllocation("tmem_o0", cfg.epi_tile[1]) + self._alloc_o1 = TmemAllocation("tmem_o1", cfg.epi_tile[1]) self.tmem_addr_cached = Int32(0) self.tmem_ptr_raw_cached = _placeholder_tmem_ptr() self.tmem_o_addr_base_cached = Int32(0) def get_tmem_requirements(self) -> list[TmemAllocation]: """Return the TMEM allocations for the double-buffered O accumulators.""" - return [self._alloc_o0, self._alloc_o1] + return ( + [self._alloc_o0] + if self.cfg.single_qkv_instance + else [self._alloc_o0, self._alloc_o1] + ) @property def loop_offset_sensitive(self) -> bool: @@ -4943,6 +6562,24 @@ def _init_function_state(self, stage_info: StageInfo) -> None: self.tmem_addr_cached, cutlass.Int8 ) self.tmem_o_addr_base_cached = Int32(0) + if cutlass.const_expr(self.cfg.vc_restores_means): + issue = cutlass.Boolean(True) + if cutlass.const_expr(self.cfg.two_cta_umma): + issue = ( + cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) == 0 + ) + self.vc_mean_issue_cached = issue + self.vc_desc_l_cached = tuple( + self._vc_rowsum_desc(stage_info, smem_p) + for smem_p in (self.smem_p0_resource, self.smem_p1_resource) + ) + self.vc_idesc_mu_cached = prims.Tcgen05InstrDesc.build( + c_dtype=cutlass.Float32, + a_dtype=cutlass.BFloat16, + b_dtype=cutlass.BFloat16, + n_dim=self.cfg.pv_mma_tiler[1], + m_dim=self.cfg.pv_mma_tiler[0] * self.cfg.cta_group_size, + ) self.tmem_p_base_cached = Int32(0) _ = stage_info @@ -4997,6 +6634,68 @@ def pv_mma( inst_idx: cutlass.Constexpr[int] = 0, is_tail: cutlass.Constexpr[bool] = False, k_half: cutlass.Constexpr[int | None] = None, + ) -> None: + """PV MMA: P*V -> O (double-buffered O0/O1). See ``_pv_mma_impl``.""" + self._pv_mma_impl( + stage_info, + desc_v_base=desc_v_base, + section=section, + head_dim_stage_idx=head_dim_stage_idx, + inst_idx=inst_idx, + is_tail=is_tail, + k_half=k_half, + ) + + @producer_work + @cute.jit + def vc_mean_mma( + self, + stage_info: StageInfo, + *, + desc_mu_base: prims.Tcgen05SmemDesc, + inst_idx: cutlass.Constexpr[int] = 0, + is_tail: cutlass.Constexpr[bool] = False, + ) -> None: + """VC-Attention: ``O += sum_i rowsum(P_i) * mean_i`` over the 16 tiles + before this one for query group ``inst_idx``, two bf16 K=16 UMMA steps + after that group's PV steps every 16th tile. At tile i this adds the + group that ended at tile i - 1, so the loop never adds the last group. + The tail call adds it with the operand loaded after the loop.""" + if cutlass.const_expr(is_tail): + tile = Int32(stage_info.loop_end) + else: + tile = Int32(stage_info.loop_offset) + group_done = ((tile & Int32(VC_MEAN_GROUP_TILES - 1)) == Int32(0)) & ( + tile > Int32(0) + ) + self._vc_mean_mma( + inst_idx, desc_mu_base, self.vc_mean_issue_cached & group_done + ) + + @producer_work + @cute.jit + def vc_mean_mma_last( + self, + stage_info: StageInfo, + *, + desc_mu_last: prims.Tcgen05SmemDesc, + inst_idx: cutlass.Constexpr[int] = 0, + ) -> None: + """Tail variant of ``vc_mean_mma``. Adds the last group, the only one with + no loop tile after it, from the operand the loader fetched after the loop.""" + self._vc_mean_mma(inst_idx, desc_mu_last, self.vc_mean_issue_cached) + + @cute.jit + def _pv_mma_impl( + self, + stage_info: StageInfo, + *, + desc_v_base: prims.Tcgen05SmemDesc, + section: cutlass.Constexpr[FmhaStage], + head_dim_stage_idx: cutlass.Constexpr[int] = 0, + inst_idx: cutlass.Constexpr[int] = 0, + is_tail: cutlass.Constexpr[bool] = False, + k_half: cutlass.Constexpr[int | None] = None, ) -> None: """PV MMA: P*V -> O (double-buffered O0/O1). @@ -5058,18 +6757,24 @@ def pv_mma( else: ab_format = cutlass.Float16 + pv_fp8_k64 = self.cfg.v_dtype.width == 8 and self.cfg.umma_fp8_k64 idesc_pv = prims.Tcgen05InstrDesc.build( c_dtype=cutlass.Float32, a_dtype=ab_format, b_dtype=ab_format, n_dim=( - self.cfg.head_dim_per_stage_kv - if self.cfg.single_qkv_instance and self.cfg.pv_mma_tiler[1] == 256 + min( + self.cfg.head_dim_per_stage_kv, + self.cfg.pv_mma_tiler[1] + - head_dim_stage_idx * self.cfg.head_dim_per_stage_kv, + ) + if self.cfg.stage_kv_by_head_dim else self.cfg.pv_mma_tiler[1] ), m_dim=self.cfg.pv_mma_tiler[0] * self.cfg.cta_group_size, # V is row-major / MN-major. b_major=1, + k_dim=1 if pv_fp8_k64 else 0, ) cta_group = ( prims.CTAGroup.CTA_2 @@ -5084,11 +6789,9 @@ def pv_mma( pv_n_dim = self.cfg.pv_mma_tiler[1] num_head_dim_stages = 1 - if cutlass.const_expr( - self.cfg.single_qkv_instance and self.cfg.pv_mma_tiler[1] == 256 - ): + if cutlass.const_expr(self.cfg.stage_kv_by_head_dim): pv_n_dim = self.cfg.head_dim_per_stage_kv - num_head_dim_stages = self.cfg.pv_mma_tiler[1] // pv_n_dim + num_head_dim_stages = self.cfg.num_head_dim_stages_v head_dim_stage_start = 0 num_head_dim_stages_to_issue = num_head_dim_stages if cutlass.const_expr(self.cfg.stage_kv_by_head_dim): @@ -5097,7 +6800,7 @@ def pv_mma( k_dim_per_mma = 16 if cutlass.const_expr(self.cfg.v_dtype.width != 16): - k_dim_per_mma = 32 + k_dim_per_mma = 64 if pv_fp8_k64 else 32 num_kphases_pv = self.cfg.pv_mma_tiler[2] // k_dim_per_mma inc_tmem_p = ( k_dim_per_mma * self.cfg.v_dtype.width // self.cfg.qk_acc_dtype.width @@ -5251,6 +6954,74 @@ def pv_mma( ) scale_d_stage = True + def _vc_mean_cta_group(self): + """cta_group::2 under two-CTA UMMA (M=256 over both CTAs' row sums).""" + if self.cfg.two_cta_umma: + return prims.CTAGroup.CTA_2 + return prims.CTAGroup.CTA_1 + + @cute.jit + def _vc_rowsum_desc( + self, stage_info: StageInfo, smem_p: SmemPResource + ) -> prims.Tcgen05SmemDesc: + """K-major descriptor of a query group's row-sum operands (A of the mean steps).""" + sL = cutlass.Array( + stage_info.context.smem_base.data_ptr() + + smem_p._alloc.offset + + smem_p.rowsum_tile_offset, + dtype=cutlass.Int8, + shape=(self.cfg.vc_rowsum_tile_bytes,), + addrspace=3, + ) + return freeze_smem_descriptor( + prims.Tcgen05SmemDesc.build( + sL, + leading_byte_offset=VC_MEAN_TILE_LBO, + stride_byte_offset=VC_MEAN_TILE_SBO, + layout=prims.Tcgen05SmemSwizzle.NONE, + ) + ) + + @cute.jit + def _vc_mean_mma( + self, + inst_idx: cutlass.Constexpr[int], + desc_mu_base: prims.Tcgen05SmemDesc, + issue_mma: Any, + ) -> None: + """Issue ``O += sum_i rowsum(P_i) * mean_i`` for query group ``inst_idx`` as + ``VC_MEAN_OPERANDS`` bf16 K=16 UMMA steps back to back. + + A is the query group's row-sum operands behind its P tile, B the + TMA-staged mean operands of the tile group. + """ + if issue_mma: + if prims.elect_sync(): + if cutlass.const_expr(inst_idx == 0): + tmem_ptr_o = self.tmem_ptr_raw_cached.subview(self.tmem_o0_offset) + else: + tmem_ptr_o = self.tmem_ptr_raw_cached.subview(self.tmem_o1_offset) + desc_l = self.vc_desc_l_cached[inst_idx] + desc_mu = freeze_smem_descriptor(desc_mu_base) + for operand in cutlass.range_constexpr(VC_MEAN_OPERANDS): + prims.tcgen05_mma( + prims.Tcgen05MMAKind.F16, + self._vc_mean_cta_group(), + tmem_ptr_o, + desc_l + ((operand * self.cfg.vc_rowsum_operand_bytes) >> 4), + desc_mu + + ( + ( + operand + * self.cfg.vc_mean_operand_bytes + // self.cfg.cta_group_size + ) + >> 4 + ), + self.vc_idesc_mu_cached, + True, + ) + @consumer_work @cute.jit def correct( @@ -5299,9 +7070,9 @@ def _correct_impl( Uses inst_idx to select O0/O1 and forwards row_sum to SmemO.producer_vars for the tail epilog. - With skip_correction enabled: uses vote_ballot_sync to check if - old_max == new_max across all threads. If true, scale=1.0 and - we skip the expensive TMEM load+rescale+store. + With skip_correction enabled, rescale only when a row with prior + visible keys changes its maximum. Empty histories and unchanged + maxima skip the TMEM load+rescale+store. """ # Select O offset from captured correction-call position. if cutlass.const_expr(self.cfg.single_qkv_instance or inst_idx == 0): @@ -5325,7 +7096,13 @@ def _correct_impl( # Check if we should skip correction (when old_max == new_max) should_rescale = cutlass.Boolean(True) if cutlass.const_expr(self.cfg.enable_skip_correction): - vote_ballot_cnt = cute.arch.vote_ballot_sync(vec_old_max != vec_new_max) + # A neutral old maximum means every prior P tile was zero. + # The first visible tile establishes an anchor but has no old O + # contribution to rescale, including a late first tile in one Q peer. + needs_rescale = (vec_old_max != vec_new_max) & ( + vec_old_max != Float32(_SOFTMAX_MASK_VALUE) + ) + vote_ballot_cnt = cute.arch.vote_ballot_sync(needs_rescale) should_rescale = vote_ballot_cnt != Int32(0) scale = Float32(1.0) @@ -5340,9 +7117,7 @@ def _correct_impl( # MMA (Thread 0) does tcgen05.mma → tcgen05.commit on O_full. # Correction (Thread 1) does mbarrier.try_wait on O_full → tcgen05.ld. # The fence orders the prior tcgen05.commit's completion with our tcgen05.ld. - from cutlass.experimental import primitives as _prims - - _prims.tcgen05_fence("after") + prims.tcgen05_fence("after") # Only rescale if old_max != new_max AND not in invalid O0 iteration if should_rescale and not skip_o0_invalid: @@ -5397,6 +7172,10 @@ class SmemOResource(MemoryResource): # Reference to the TmemStats resource for this stage's correction stats. tmem_vec_resource: TmemStatsResource | None = field(init=False, default=None) stage_idx: Constexpr[int] = field(init=False, default=0) + # Sage attention: the prepared scales and the KV head count that strides + # their sequences; the epilogue reads the work tile's ``sfV`` / ``v_mean``. + sage_scales: cute.Pointer | None = field(init=False, default=None) + num_kv_heads: int | Int32 = field(init=False, default=0) _alloc: Constexpr[Optional[SmemAllocation]] = field(init=False, default=None) head_coord: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() batch_coord: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() @@ -5408,6 +7187,8 @@ def __init__( cfg: FmhaConfig, stage_idx: int = 0, tmem_vec_resource: TmemStatsResource | None = None, + sage_scales: cute.Pointer | None = None, + num_kv_heads: int | Int32 = 0, **kwargs: Any, ) -> None: """Bind one output subtile stage and reserve its SMEM staging buffer.""" @@ -5415,6 +7196,8 @@ def __init__( self.cfg = cfg self.stage_idx = stage_idx self.tmem_vec_resource = tmem_vec_resource + self.sage_scales = sage_scales + self.num_kv_heads = num_kv_heads stage_elements = cfg.sO_stage_elements size_bytes = stage_elements * cfg.o_dtype.width // 8 self._alloc = SmemAllocation( @@ -5528,7 +7311,13 @@ def _store_o( output_scale: SoftmaxScalar, head_dim_stage_idx: cutlass.Constexpr[int] = 0, ) -> None: - """Correct and normalize O in one TMEM pass, then stage it in SMEM.""" + """Correct and normalize O in one TMEM pass, then stage it in SMEM. + + With Sage attention column ``c`` becomes ``o * scale * sfV[c] + + output_scale * v_mean[c]``: the per-channel scale joins the row's + normalization and the mean is added after it. The scales are read + from the prepared buffer once per output slice, as warp broadcasts. + """ if cutlass.const_expr(self.stage_idx == 0): tmem_o_offset = self.cfg.tmem_o0_offset else: @@ -5545,6 +7334,23 @@ def _store_o( tidx, _, _ = cute.arch.thread_idx() tid_in_wg = tidx % (cute.arch.WARP_SIZE * num_correction_warps) + vc_scale_row = Int64(0) + if cutlass.const_expr(self.cfg.vc_attention): + # Start of this batch and head row in the per-channel V scale table. + assert self.tmem_vec_resource is not None + assert self.tmem_vec_resource.output_scale is not None + assert self.cfg.work_tile_q_heads == 1 + vc_scale = self.tmem_vec_resource.output_scale + _, head_coord_wt, batch_coord = _resolve_work_tile_coords( + self.cfg, stage_info.work_tile.tile_idx + ) + vc_scale_row = vc_scale.iterator.toint() + Int64( + cute.crd2idx( + (batch_coord, head_coord_wt * self.cfg.work_tile_q_heads, 0), + vc_scale.layout, + ) + ) * Int64(4) + o_head_dim = self.cfg.epi_tile[1] tma_copy_o_iters = self.cfg.tma_copy_o_iters if cutlass.const_expr(self.cfg.stage_o_by_head_dim): @@ -5559,13 +7365,35 @@ def _store_o( tmem_shape = "32x32b" tmem_x = 16 - num_iters = o_head_dim // tmem_x + # Do not load unallocated TMEM columns in a partial final O band. + valid_head_dim = min( + o_head_dim, self.cfg.epi_tile[1] - head_dim_stage_idx * o_head_dim + ) + num_iters = valid_head_dim // tmem_x smem_o_swizzle = _smem_o_swizzle(self.cfg) d_block_size = o_head_dim // tma_copy_o_iters row_offset = tid_in_wg * d_block_size + if cutlass.const_expr(self.cfg.use_sage_attention): + # The sequence's first sfV word plus this stage's first channel; + # the per-iteration channels are constant offsets from it. + _, head_coord, batch_coord = _resolve_work_tile_coords( + self.cfg, stage_info.work_tile.tile_idx + ) + sage_sequence = sage_sequence_words( + self.cfg.sage_scale_layout, + batch_idx=batch_coord, + kv_head_idx=head_coord // Int32(self.cfg.h_r), + num_kv_heads=self.num_kv_heads, + ) + sage_v_address = sage_v_section_address( + self.cfg.sage_scale_layout, + self.sage_scales, + sequence_words=sage_sequence, + ) + Int64(head_dim_stage_idx * o_head_dim) * Int64(4) + for i in cutlass.range_constexpr(num_iters): tmem_offset_tile = tmem_offset_o + i * tmem_x @@ -5579,7 +7407,29 @@ def _store_o( cute.arch.fence_view_async_tmem_load() scale_vec = cutlass.vector.full_like(o_rmem, scale) - o_rmem = o_rmem * scale_vec + if cutlass.const_expr(self.cfg.use_sage_attention): + v_scales, v_means = load_sage_v_channels( + self.cfg.sage_scale_layout, + sage_v_address, + offset=i * tmem_x, + count=tmem_x, + ) + o_rmem = o_rmem * (scale_vec * v_scales) + if cutlass.const_expr(v_means is not None): + o_rmem = o_rmem + v_means * cutlass.vector.full_like( + o_rmem, output_scale + ) + elif cutlass.const_expr(self.cfg.vc_attention): + o_rmem = o_rmem * scale_vec + # The tile means were stored divided by the per-channel V scale, + # so the whole accumulator is rescaled by it here. + col0 = head_dim_stage_idx * o_head_dim + i * tmem_x + chan_scale = cutlass.inttoptr( + vc_scale_row + Int64(col0) * Int64(4), mem_space=1, dtype=Float32 + ).load(count=tmem_x, alignment=64) + o_rmem = o_rmem * chan_scale + else: + o_rmem = o_rmem * scale_vec o_rmem_dtype = o_rmem.to(self.cfg.o_dtype) @@ -5768,5 +7618,4 @@ def tma_store( from .helpers_paged import ( _load_block_table_row_bounds, _load_runtime_seq_len_kv, - _resolve_kv_tile_idx_context, ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_tasks.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_tasks.py index 2e2d8581378e..cd2fc67f6018 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_tasks.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/fmha_tasks.py @@ -34,6 +34,7 @@ from cutlass import Int32 from ..stage import FmhaStage +from ...sage.modes import SageKTransport, SageScoreScaling from cutlass.experimental.task_scheduling.schedule_builder import ( domain_loop, schedule, @@ -50,6 +51,7 @@ SmemKVResource, SmemOResource, SmemPageOffsetsKvResource, + SmemSageKScalesResource, SmemQResource, TmemOResource, TmemPResource, @@ -59,25 +61,32 @@ TmemStatsDoneResource, TmemPPrefixReadyResource, ) +from .vc_attention.vc_resources import SmemMuResource @dataclass(kw_only=True) class PackedContextWorkQueue(WorkQueue): - """Persistent queue that skips Q tiles outside a live packed request.""" + """Persistent queue that skips empty Q/KV requests and padded Q tiles.""" cfg: cutlass.Constexpr[FmhaConfig] = field(init=False, default=None) cum_seqlen_q: Any = field(init=False, default=None) + cum_seqlen_k: Any = field(init=False, default=None) + seq_lens_kv: Any = field(init=False, default=None) def __init__( self, cfg: FmhaConfig, - cum_seqlen_q: cute.Tensor, + cum_seqlen_q: cute.Tensor | None, + cum_seqlen_k: cute.Tensor | None = None, + seq_lens_kv: cute.Pointer | None = None, **kwargs: Any, ) -> None: """Attach the live packed-Q metadata used by the skip predicate.""" super().__init__(**kwargs) self.cfg = cfg self.cum_seqlen_q = cum_seqlen_q + self.cum_seqlen_k = cum_seqlen_k + self.seq_lens_kv = seq_lens_kv @cute.jit def skip_work_tile_if(self, work_tile: Any) -> cutlass.Boolean: @@ -87,10 +96,24 @@ def skip_work_tile_if(self, work_tile: Any) -> cutlass.Boolean: if cutlass.const_expr(self.cfg.uses_causal_reversed_head_batch_seq_tile_order): seq_coord = Int32(self.cfg.num_seq_tiles) - seq_coord - Int32(1) batch_coord = Int32(work_tile.tile_idx[batch_idx]) - q_begin = Int32(self.cum_seqlen_q[batch_coord]) - q_end = Int32(self.cum_seqlen_q[batch_coord + Int32(1)]) - seqlen_q = q_end - q_begin - return seq_coord * Int32(self.cfg.cta_tiler[0]) >= seqlen_q + skip = cutlass.Boolean(False) + if cutlass.const_expr( + self.cum_seqlen_q is not None and not self.cfg.has_uniform_varlen + ): + q_begin = Int32(self.cum_seqlen_q[batch_coord]) + q_end = Int32(self.cum_seqlen_q[batch_coord + Int32(1)]) + skip = seq_coord * Int32(self.cfg.cta_tiler[0]) >= q_end - q_begin + if cutlass.const_expr(self.seq_lens_kv is not None): + skip = skip or Int32(self.seq_lens_kv[batch_coord]) <= Int32(0) + elif cutlass.const_expr( + self.cum_seqlen_k is not None and not self.cfg.has_uniform_varlen + ): + skip = ( + skip + or self.cum_seqlen_k[batch_coord + Int32(1)] + <= self.cum_seqlen_k[batch_coord] + ) + return skip def _persistent_tail(work_queue: WorkQueue) -> None: @@ -188,13 +211,19 @@ def create_load_task( task_class: type[Task] = Task, smem_page_offsets_kv: SmemPageOffsetsKvResource | None = None, smem_page_offsets_v: SmemPageOffsetsKvResource | None = None, + smem_mu: SmemMuResource | None = None, + smem_sage_k: SmemSageKScalesResource | None = None, **task_kwargs: Any, ) -> Task: """Create the one-warp TMA load task. + ``smem_mu`` (VC-Attention) stages one bf16 tile-mean operand per K/V tile + beside V. + When ``smem_page_offsets_kv`` is provided, each K/V TMA load consumes page IDs prefetched by the auxiliary warp through the ordinary asynchronous - page-offset pipeline. + page-offset pipeline. When ``smem_sage_k`` is provided, the warp also + stages each K tile's Sage ``sfK`` words after the tile's K load. """ loop_start, loop_end, loop_step = _captured_loop_bounds(task_class, task_kwargs) skip_work_tile_if = _packed_context_skip_predicate(work_queue) @@ -205,10 +234,14 @@ def create_load_task( raise ValueError("split K/V staging requires a separate V buffer") kv_resources = (smem_k_or_kv, smem_v) if split_kv else (smem_k_or_kv,) dst = [smem_q, *kv_resources] + if smem_mu is not None: + dst.append(smem_mu) if smem_page_offsets_kv is not None: src.append(smem_page_offsets_kv) if smem_page_offsets_v is not None: src.append(smem_page_offsets_v) + if smem_sage_k is not None: + dst.append(smem_sage_k) if smem_q.cfg.single_qkv_instance and smem_q.cfg.has_tmem_p_pipeline: num_head_dim_stages_k = smem_k_or_kv.cfg.num_head_dim_stages_k num_head_dim_stages_v = smem_k_or_kv.cfg.num_head_dim_stages_v @@ -266,6 +299,7 @@ def load_reused_page_windows_schedule_body( seqlen_k, kv_tile_start, kv_request_begin, + kv_page_idx_lb, kv_page_idx_ub, ) = gqkv.compute_coords() sq.acquire() @@ -293,6 +327,7 @@ def load_k_tile(*, tile_offset: int) -> None: seqlen_k=seqlen_k, kv_tile_start=kv_tile_start, kv_request_begin=kv_request_begin, + kv_page_idx_lb=kv_page_idx_lb, kv_page_idx_ub=kv_page_idx_ub, ) sk.commit() @@ -327,6 +362,7 @@ def load_v_tile( seqlen_k=seqlen_k, kv_tile_start=kv_tile_start, kv_request_begin=kv_request_begin, + kv_page_idx_lb=kv_page_idx_lb, kv_page_idx_ub=kv_page_idx_ub, ) else: @@ -339,6 +375,7 @@ def load_v_tile( seqlen_k=seqlen_k, kv_tile_start=kv_tile_start, kv_request_begin=kv_request_begin, + kv_page_idx_lb=kv_page_idx_lb, kv_page_idx_ub=kv_page_idx_ub, ) sv.commit() @@ -471,6 +508,7 @@ def load_schedule_body( seqlen_k, kv_tile_start, kv_request_begin, + kv_page_idx_lb, kv_page_idx_ub, ) = coords sq.acquire() @@ -498,6 +536,7 @@ def load_schedule_body( seqlen_k=seqlen_k, kv_tile_start=kv_tile_start, kv_request_begin=kv_request_begin, + kv_page_idx_lb=kv_page_idx_lb, kv_page_idx_ub=kv_page_idx_ub, ) sk.commit() @@ -519,6 +558,7 @@ def load_schedule_body( seqlen_k=seqlen_k, kv_tile_start=kv_tile_start, kv_request_begin=kv_request_begin, + kv_page_idx_lb=kv_page_idx_lb, kv_page_idx_ub=kv_page_idx_ub, ) sk.commit() @@ -540,6 +580,7 @@ def load_schedule_body( seqlen_k=seqlen_k, kv_tile_start=kv_tile_start, kv_request_begin=kv_request_begin, + kv_page_idx_lb=kv_page_idx_lb, kv_page_idx_ub=kv_page_idx_ub, ) sv.commit() @@ -561,6 +602,7 @@ def load_schedule_body( seqlen_k=seqlen_k, kv_tile_start=kv_tile_start, kv_request_begin=kv_request_begin, + kv_page_idx_lb=kv_page_idx_lb, kv_page_idx_ub=kv_page_idx_ub, ) sv.commit() @@ -651,12 +693,32 @@ def load_schedule_body( sq: SmemQResource, sk: SmemKVResource, sv: SmemKVResource, + skr: SmemSageKScalesResource | None, wq: WorkQueue | None, + smu: SmemMuResource | None = None, ) -> None: - """Load paired Q instances and their directly addressed K/V tiles.""" + """Load paired Q instances and their directly addressed K/V tiles. + + With ``skr`` each K tile's Sage ``sfK`` words follow the tile's K + load into the scale ring. + """ sq.init_load_state() sk.init_load_state() sv.init_load_state() + if smu is not None: + smu.init_load_state() + + def stage_sage_k_scales() -> None: + if skr is None: + return + skr.acquire() + skr.copy_tile( + kv_head_coord=kv_head_coord, + batch_coord=batch_coord, + kv_tile_start=kv_tile_start, + ) + skr.commit() + with _work_tile_schedule_loop(wq, skip_if=skip_work_tile_if): # noqa: SIM117 # The first K-loop iteration also loads Q0/Q1. Later iterations # only stream the next K/V tiles through their respective pipelines. @@ -675,6 +737,7 @@ def load_schedule_body( seqlen_k, kv_tile_start, kv_request_begin, + kv_page_idx_lb, kv_page_idx_ub, ) = gqkv.compute_coords() # Load Q0 for the first Q tile in this work tile. @@ -714,9 +777,11 @@ def load_schedule_body( kv_tile_start=kv_tile_start, seqlen_k=seqlen_k, kv_request_begin=kv_request_begin, + kv_page_idx_lb=kv_page_idx_lb, kv_page_idx_ub=kv_page_idx_ub, ) sk.commit() + stage_sage_k_scales() else: # Throttle TMA before reserving a KV stage. sk.try_acquire() @@ -729,9 +794,11 @@ def load_schedule_body( kv_tile_start=kv_tile_start, seqlen_k=seqlen_k, kv_request_begin=kv_request_begin, + kv_page_idx_lb=kv_page_idx_lb, kv_page_idx_ub=kv_page_idx_ub, ) sk.commit() + stage_sage_k_scales() with d.first_iter(): # Load Q1 for the second Q tile in this work tile. sq.acquire() @@ -755,9 +822,73 @@ def load_schedule_body( seqlen_k=seqlen_k, kv_tile_start=kv_tile_start, kv_request_begin=kv_request_begin, + kv_page_idx_lb=kv_page_idx_lb, kv_page_idx_ub=kv_page_idx_ub, ) sv.commit() + if smu is not None: + # VC-Attention: the mean operand of tile i-1 rides beside V_i. + smu.try_acquire() + smu.acquire() + smu.mu_load( + kv_head_coord=kv_head_coord, + batch_coord=batch_coord, + kv_tile_start=kv_tile_start, + ) + smu.commit() + if smu is not None: + # The loop step at tile i adds the means of the group that ended + # at tile i - 1, so the loop covers groups 0 .. N - 2 and the + # group of the last tile never meets a following loop tile. Its + # operand is fetched once more here for the tail mean step. + ( + _seq_coord, + _head_coord, + kv_head_coord, + _head_coord_kv, + batch_coord, + _seq_coord_q, + _cuseqlen_q, + _cuseqlen_k, + _seqlen_q, + _seqlen_k, + kv_tile_start, + _kv_request_begin, + _kv_page_idx_lb, + _kv_page_idx_ub, + ) = gqkv.compute_coords() + smu.try_acquire() + smu.acquire() + smu.mu_load( + kv_head_coord=kv_head_coord, + batch_coord=batch_coord, + kv_tile_start=kv_tile_start, + is_tail=True, + ) + smu.commit() + + @schedule + def load_schedule_mu( + gqkv: GmemQKVResource, + sq: SmemQResource, + skv: SmemKVResource, + smu: SmemMuResource, + wq: WorkQueue | None = None, + ) -> None: + """Shared K/V buffer plus the VC-Attention tile-mean ring.""" + load_schedule_body(gqkv, sq, skv, skv, None, wq, smu) # type: ignore[call-arg] + + @schedule + def load_split_schedule_mu( + gqkv: GmemQKVResource, + sq: SmemQResource, + sk: SmemKVResource, + sv: SmemKVResource, + smu: SmemMuResource, + wq: WorkQueue | None = None, + ) -> None: + """Split K and V buffers plus the VC-Attention tile-mean ring.""" + load_schedule_body(gqkv, sq, sk, sv, None, wq, smu) # type: ignore[call-arg] @schedule def load_schedule( @@ -768,7 +899,18 @@ def load_schedule( ) -> None: """Contiguous-KV captured schedule over a shared K/V buffer.""" # Mypy retains the earlier branch's five-argument closure signature. - load_schedule_body(gqkv, sq, skv, skv, wq) # type: ignore[call-arg] + load_schedule_body(gqkv, sq, skv, skv, None, wq) # type: ignore[call-arg] + + @schedule + def load_sage_schedule( + gqkv: GmemQKVResource, + sq: SmemQResource, + skv: SmemKVResource, + skr: SmemSageKScalesResource, + wq: WorkQueue | None = None, + ) -> None: + """Contiguous-KV captured schedule that also stages Sage ``sfK`` tiles.""" + load_schedule_body(gqkv, sq, skv, skv, skr, wq) # type: ignore[call-arg] @schedule def load_split_schedule( @@ -779,15 +921,36 @@ def load_split_schedule( wq: WorkQueue | None = None, ) -> None: """Contiguous-KV captured schedule over split K and V buffers.""" - load_schedule_body(gqkv, sq, sk, sv, wq) # type: ignore[call-arg] + load_schedule_body(gqkv, sq, sk, sv, None, wq) # type: ignore[call-arg] - captured_schedule = _schedule_with_work_queue( - load_split_schedule if split_kv else load_schedule, - gmem_qkv, - smem_q, - *kv_resources, - work_queue=work_queue, - ) + if smem_mu is not None: + captured_schedule = _schedule_with_work_queue( + load_split_schedule_mu if split_kv else load_schedule_mu, + gmem_qkv, + smem_q, + *kv_resources, + smem_mu, + work_queue=work_queue, + ) + elif smem_sage_k is not None: + if split_kv: + raise ValueError("the Sage scale ring requires a shared K/V buffer") + captured_schedule = _schedule_with_work_queue( + load_sage_schedule, + gmem_qkv, + smem_q, + *kv_resources, + smem_sage_k, + work_queue=work_queue, + ) + else: + captured_schedule = _schedule_with_work_queue( + load_split_schedule if split_kv else load_schedule, + gmem_qkv, + smem_q, + *kv_resources, + work_queue=work_queue, + ) return task_class( src_resources=src, dst_resources=dst, @@ -817,9 +980,14 @@ def create_mma_task( tmem_p_prefix_ready_1: TmemPPrefixReadyResource | None = None, smem_p0: SmemPResource | None = None, smem_p1: SmemPResource | None = None, + smem_mu: SmemMuResource | None = None, **task_kwargs: Any, ) -> Task: - """Create the one-warp MMA compute task.""" + """Create the one-warp MMA compute task. + + ``smem_mu`` (VC-Attention) supplies the bf16 tile-mean operand consumed by + one extra ``kind::f16`` step after each tile's PV steps. + """ loop_start, loop_end, loop_step = _captured_loop_bounds(task_class, task_kwargs) skip_work_tile_if = _packed_context_skip_predicate(work_queue) # Only paged MMA schedules read request coordinates from global metadata. @@ -840,12 +1008,14 @@ def create_mma_task( if p_in_smem and (smem_p0 is None or smem_p1 is None): raise ValueError("p_in_smem requires both SMEM P resources") smem_p_resources = (smem_p0, smem_p1) if p_in_smem else () + mu_resources = (smem_mu,) if smem_mu is not None else () src = _src_resources( *qkv_resources, smem_q, *kv_resources, *p_prefix_resources, *smem_p_resources, + *mu_resources, work_queue=work_queue, ) num_head_dim_stages_k = smem_k_or_kv.cfg.num_head_dim_stages_k @@ -861,12 +1031,10 @@ def create_mma_task( ) num_head_dim_stages_k = smem_k_or_kv.cfg.num_head_dim_stages_k num_head_dim_stages_v = smem_k_or_kv.cfg.num_head_dim_stages_v - loop_carried_head_dim_stages = 2 - if ( - num_head_dim_stages_k != loop_carried_head_dim_stages - or num_head_dim_stages_v != loop_carried_head_dim_stages - ): - raise ValueError("loop-carried split S/P scheduling expects two K/V stages") + if num_head_dim_stages_k not in (1, 2) or num_head_dim_stages_v not in (1, 2): + raise ValueError( + "loop-carried split S/P scheduling expects one or two K/V stages" + ) def mma_schedule_body( gqkv: GmemQKVResource, @@ -908,6 +1076,7 @@ def mma_schedule_body( v_seqlen_k, v_kv_tile_start, _kv_request_begin, + _kv_page_idx_lb, _kv_page_idx_ub, ) = gqkv.compute_coords() @@ -947,15 +1116,16 @@ def mma_schedule_body( ) sk.release() - sk.wait() - desc_k_base = sk.k_desc() - sp0.qk_mma( - desc_q_base=desc_q0_base, - desc_k_base=desc_k_base, - section=FmhaStage.Loop, - head_dim_stage_idx=1, - ) - sk.release() + if num_head_dim_stages_k > 1: + sk.wait() + desc_k_base = sk.k_desc() + sp0.qk_mma( + desc_q_base=desc_q0_base, + desc_k_base=desc_k_base, + section=FmhaStage.Loop, + head_dim_stage_idx=1, + ) + sk.release() sp0.commit() if not smem_q.cfg.stats_via_smem: vd0.commit() @@ -986,21 +1156,22 @@ def mma_schedule_body( ) sv.release() - sv.wait() - if cutlass.const_expr(smem_q.cfg.needs_paged_v_tail_clear): - desc_v_base = sv.v_desc_paged( + if num_head_dim_stages_v > 1: + sv.wait() + if cutlass.const_expr(smem_q.cfg.needs_paged_v_tail_clear): + desc_v_base = sv.v_desc_paged( + section=FmhaStage.Loop, + seqlen_k=v_seqlen_k, + kv_tile_start=v_kv_tile_start, + ) + else: + desc_v_base = sv.v_desc() + to.pv_mma( + desc_v_base=desc_v_base, section=FmhaStage.Loop, - seqlen_k=v_seqlen_k, - kv_tile_start=v_kv_tile_start, + head_dim_stage_idx=1, ) - else: - desc_v_base = sv.v_desc() - to.pv_mma( - desc_v_base=desc_v_base, - section=FmhaStage.Loop, - head_dim_stage_idx=1, - ) - sv.release() + sv.release() to.commit() tp0.release() @@ -1127,6 +1298,7 @@ def mma_schedule_body( v_seqlen_k, v_kv_tile_start, _kv_request_begin, + _kv_page_idx_lb, _kv_page_idx_ub, ) = gqkv.compute_coords() if wq is not None: @@ -1257,6 +1429,11 @@ def load_k_stage(sk, head_dim_stage_idx): return sk.k_desc() return sk.k_stage_desc() + def seed_scores(sp): + """Seed an acquired S stage with the INT32 score bias before its K wait.""" + if smem_q.cfg.uses_int32_scores: + sp.seed_scores() + def qk_mma_stages(sk, sp, desc_q_base, section, *, descriptors=None): """Issue all K slices, loading descriptors once and reusing them for Q1.""" stages = [] @@ -1285,8 +1462,12 @@ def mma_schedule_body( pr1: TmemPPrefixReadyResource | None = None, pb0: SmemPResource | None = None, pb1: SmemPResource | None = None, + smu: SmemMuResource | None = None, ) -> None: """Interleave paired QK/PV while retaining each K slice for both Qs.""" + # VC-Attention: the current tile's mean-operand descriptor, waited and + # released in lockstep with V. + mu_state: dict[str, Any] = {"desc": None} def pv( sp: TmemSPResource, @@ -1301,6 +1482,14 @@ def pv( if p_in_smem: pb.wait() to.pv_mma(inst_idx=group, **kw) + if smu is not None: + # VC-Attention: a group's row sums travel with the P of the + # tile after it, so its mean step follows that tile's PV steps. + to.vc_mean_mma( + inst_idx=group, + desc_mu_base=mu_state["desc"], + is_tail=kw.get("is_tail", False), + ) pb.release() return if pv_half_overlap: @@ -1314,6 +1503,8 @@ def pv( sq.init_descriptor_state() sk.init_descriptor_state() sv.init_descriptor_state() + if smu is not None: + smu.init_descriptor_state() sp0.init_mma_state() sp1.init_mma_state() to.init_mma_state() @@ -1334,6 +1525,7 @@ def pv( v_seqlen_k, v_kv_tile_start, _kv_request_begin, + _kv_page_idx_lb, _kv_page_idx_ub, ) = gqkv.compute_coords() # HEAD: consume Q0, K0, Q1, and V0. TmemStatsDone starts empty, so @@ -1349,6 +1541,7 @@ def pv( if not smem_q.cfg.stats_via_smem: vd0.acquire() sp0.acquire() + seed_scores(sp0) qk_mma_slice( sp0, desc_q0_base, desc_k, FmhaStage.Head, head_dim_stage_idx ) @@ -1362,6 +1555,7 @@ def pv( if not smem_q.cfg.stats_via_smem: vd1.acquire() sp1.acquire() + seed_scores(sp1) qk_mma_slice( sp1, desc_q1_base, desc_k, FmhaStage.Head, head_dim_stage_idx ) @@ -1414,6 +1608,17 @@ def wait_next_v() -> Any: ) return sv.v_desc() + def release_mu() -> None: + """Release the mean operand of the tile whose PV steps were issued.""" + if smu is not None: + smu.release() + + def wait_next_mu() -> None: + """Wait for the next tile's mean operand and cache its descriptor.""" + if smu is not None: + smu.wait() + mu_state["desc"] = smu.mu_desc() + if p_in_smem: # With P in SMEM, softmax releases the S stage right after loading # S, so QK(i+1) overlaps exp2(i). The MMA issues QK0(i+1), PV0(i), @@ -1421,8 +1626,11 @@ def wait_next_v() -> Any: # waits on its group's P tile instead. The tail drains no SP slot. with domain_loop(loop_start, loop_end, loop_step): sp0.acquire() + seed_scores(sp0) k_descriptors = qk_mma_stages(sk, sp0, desc_q0_base, FmhaStage.Loop) sp0.commit() + # Wait for the mean operand here, off the QK issue path. + wait_next_mu() to.acquire() pv( sp0, @@ -1434,6 +1642,7 @@ def wait_next_v() -> Any: ) to.commit() sp1.acquire() + seed_scores(sp1) qk_mma_stages( sk, sp1, desc_q1_base, FmhaStage.Loop, descriptors=k_descriptors ) @@ -1449,9 +1658,27 @@ def wait_next_v() -> Any: ) to.commit() sv.release() + release_mu() desc_v_base = wait_next_v() sq.release() sq.release() + if smu is not None: + # The loop issued the mean step of group g at the first tile + # after it, so with N groups it covered 0 .. N - 2. The tail + # O window of each query group adds group N - 1 from the + # operand the loader fetched after the loop. Both ring slots + # are held until then. + wait_next_mu() + smu.wait() + desc_mu_last = smu.mu_desc_last() + + def tail_mean(group: int, pb: Any) -> None: + """Issue the final mean step of ``group`` inside its O window.""" + if smu is not None: + pb.wait() + to.vc_mean_mma_last(inst_idx=group, desc_mu_last=desc_mu_last) + pb.release() + to.acquire() pv( sp0, @@ -1462,6 +1689,7 @@ def wait_next_v() -> Any: section=FmhaStage.Tail, is_tail=True, ) + tail_mean(0, pb0) to.commit() to.acquire() pv( @@ -1473,8 +1701,11 @@ def wait_next_v() -> Any: section=FmhaStage.Tail, is_tail=True, ) + tail_mean(1, pb1) to.commit() sv.release() + release_mu() + release_mu() else: # QK0 -> PV1 -> QK1 -> PV0: retain each K slice and the previous V # until both peer query tiles have consumed the corresponding data. @@ -1635,7 +1866,52 @@ def mma_schedule_pb( gqkv, sq, skv, skv, sp0, sp1, to, vd0, vd1, wq, None, None, pb0, pb1 ) - if p_in_smem: + @schedule + def mma_split_schedule_pb_mu( + gqkv: GmemQKVResource, + sq: SmemQResource, + sk: SmemKVResource, + sv: SmemKVResource, + sp0: TmemSPResource, + sp1: TmemSPResource, + to: TmemOResource, + vd0: TmemStatsDoneResource, + vd1: TmemStatsDoneResource, + pb0: SmemPResource, + pb1: SmemPResource, + smu: SmemMuResource, + wq: WorkQueue | None = None, + ) -> None: + """Split K/V captured schedule with P in SMEM and VC tile means.""" + mma_schedule_body( + gqkv, sq, sk, sv, sp0, sp1, to, vd0, vd1, wq, None, None, pb0, pb1, smu + ) + + @schedule + def mma_schedule_pb_mu( + gqkv: GmemQKVResource, + sq: SmemQResource, + skv: SmemKVResource, + sp0: TmemSPResource, + sp1: TmemSPResource, + to: TmemOResource, + vd0: TmemStatsDoneResource, + vd1: TmemStatsDoneResource, + pb0: SmemPResource, + pb1: SmemPResource, + smu: SmemMuResource, + wq: WorkQueue | None = None, + ) -> None: + """Shared-buffer captured schedule with P in SMEM and VC tile means.""" + mma_schedule_body( + gqkv, sq, skv, skv, sp0, sp1, to, vd0, vd1, wq, None, None, pb0, pb1, smu + ) + + if smem_mu is not None: + selected_mma_schedule = ( + mma_split_schedule_pb_mu if split_kv else mma_schedule_pb_mu + ) + elif p_in_smem: selected_mma_schedule = mma_split_schedule_pb if split_kv else mma_schedule_pb elif pv_half_overlap: selected_mma_schedule = mma_split_schedule_pr if split_kv else mma_schedule_pr @@ -1653,6 +1929,7 @@ def mma_schedule_pb( tmem_vec_done_1, *p_prefix_resources, *smem_p_resources, + *mu_resources, work_queue=work_queue, ) return task_class( @@ -1678,6 +1955,7 @@ def create_softmax_task( task_class: type[Task] = Task, tmem_p_prefix_ready: TmemPPrefixReadyResource | None = None, smem_p: SmemPResource | None = None, + smem_sage_k: SmemSageKScalesResource | None = None, **task_kwargs: Any, ) -> Task: """Create a four-warp Softmax task. @@ -2190,6 +2468,26 @@ def softmax_schedule( if smem_p is None: raise ValueError("p_in_smem requires the group's SMEM P resource") dst.append(smem_p) + vc_attention = tmem_sp.cfg.vc_attention + # VC-Attention V treatment: tile means restored in-kernel, or V repair tiles. + vc_restores_means = tmem_sp.cfg.vc_restores_means + vc_repairs_v = vc_attention and not vc_restores_means + # Sage attention takes each K tile's sfK words from the prepared buffer + # or from the load warp's ring (``SageKTransport``), and scales blocks of + # raw scores or dequantizes them (``SageScoreScaling``). + use_sage = tmem_sp.cfg.use_sage_attention + sage_transport = tmem_sp.cfg.sage_k_transport + sage_scaling = tmem_sp.cfg.sage_score_scaling + if use_sage and not p_in_smem: + raise ValueError("Sage attention requires P in SMEM") + if sage_transport is SageKTransport.SMEM and sage_scaling is SageScoreScaling.BLOCK: + # The P pass of block-scaled scores reuses the max pass's words, and + # a ring stage is released after the max pass. + raise ValueError("block-scaled Sage scores load their sfK words directly") + if sage_transport is SageKTransport.SMEM: + if smem_sage_k is None: + raise ValueError("this Sage K block size requires the sfK ring") + src.append(smem_sage_k) def softmax_schedule_body( sp: TmemSPResource, @@ -2198,6 +2496,7 @@ def softmax_schedule_body( wq: WorkQueue | None = None, pr: TmemPPrefixReadyResource | None = None, pb: SmemPResource | None = None, + skr: SmemSageKScalesResource | None = None, ) -> None: """Captured schedule for one softmax warp group.""" if tmem_sp.enable_early_tile_sum: @@ -2208,10 +2507,54 @@ def softmax_schedule_body( p_chunk = sp.init_softmax_state() scale_softmax_log2 = sp.load_scale_softmax_log2() - def exp2_p(sp: TmemSPResource, *, row_max: Any, scale_softmax_log2: Any) -> Any: + # VC-Attention sums the E4M3 probability codes inside the P store, + # so vc_store_p returns the row sum and softmax_aux_reduce is skipped. + vc_state: dict[str, Any] = { + "prev_sum": None, + "row_sum": None, + "sums": None, + "k_next": None, + "kept": None, + } + + def exp2_p( + sp: TmemSPResource, + *, + row_max: Any, + scale_softmax_log2: Any, + old_row_max: Any = None, + is_tail: bool = False, + sage_q_scale: Any = None, + sage_k_words: Any = None, + ) -> Any: """Softmax and P store. With P in SMEM the store goes to the SMEM P tile, with the half overlap the leading half is published behind ``pr``, else P goes - to the TMEM S/P stage.""" + to the TMEM S/P stage. VC-Attention then carries the group row sums + (rescaled to this tile's max) and stores the completed group's operand + for the deferred mean step. A Sage plan passes the work tile's ``sfQ`` and, when it + scales blocks, the tile's ``sfK`` words from its max pass.""" + if sage_scaling is SageScoreScaling.BLOCK: + p_chunk = sp.sage_exp2_p_smem( + row_max=row_max, + scale_softmax_log2=scale_softmax_log2, + sage_q_scale=sage_q_scale, + sage_k_words=sage_k_words, + ) + pb.acquire() + sp.store_p() + pb.commit() + return p_chunk + if use_sage: + # The scores carry sfK already; sfQ joins the exp2 scale. + p_chunk = sp.sage_dequantized_exp2_p_smem( + row_max=row_max, + scale_softmax_log2=scale_softmax_log2, + sage_q_scale=sage_q_scale, + ) + pb.acquire() + sp.store_p() + pb.commit() + return p_chunk if p_in_smem: # The S stage is released already, so exp2 and the P store are # auxiliary work on the loaded S and the SMEM P tile. @@ -2219,7 +2562,30 @@ def exp2_p(sp: TmemSPResource, *, row_max: Any, scale_softmax_log2: Any) -> Any: row_max=row_max, scale_softmax_log2=scale_softmax_log2 ) pb.acquire() - sp.store_p() + if vc_repairs_v: + out = sp.vc_store_p_repair( + old_row_max=old_row_max, + row_max=row_max, + row_sum=vc_state["row_sum"], + vc_kept_row_sum=vc_state["kept"], + p_chunk=p_chunk, + is_tail=is_tail, + ) + vc_state["kept"], vc_state["row_sum"] = out[0], out[1] + elif vc_restores_means: + out = sp.vc_store_p( + old_row_max=old_row_max, + row_max=row_max, + row_sum=vc_state["row_sum"], + vc_prev_tile_sum=vc_state["prev_sum"], + p_chunk=p_chunk, + is_tail=is_tail, + vc_meansums=vc_state["sums"], + ) + vc_state["prev_sum"], vc_state["row_sum"] = out[0], out[1] + vc_state["sums"] = out[2] + else: + sp.store_p() pb.commit() return p_chunk if not pv_half_overlap: @@ -2236,6 +2602,13 @@ def exp2_p(sp: TmemSPResource, *, row_max: Any, scale_softmax_log2: Any) -> Any: # Recompute per-tile SP/Vec TMEM state. old_row_max, row_max, row_sum, q_offset = sp.init_softmax_work_tile_state() vec.init_store_work_tile_state() + vc_state["row_sum"] = row_sum + if vc_attention: + init = sp.vc_init_row_scale() + vc_row_scale, vc_state["prev_sum"] = init[0], init[1] + vc_state["sums"] = init[2] + vc_state["kept"] = init[3] + vc_state["k_next"] = init[4] if tmem_sp.uses_varlen_q_offset_cache: q_offset = sp.cache_q_offset() if tmem_sp.uses_packed_dense_k_mask: @@ -2244,12 +2617,60 @@ def exp2_p(sp: TmemSPResource, *, row_max: Any, scale_softmax_log2: Any) -> Any: window_end = Int32(0) if tmem_sp.uses_variable_window: window_start, window_end = sp.cache_variable_window_bounds() + sage_q_scale = sage_k_words = None + if use_sage: + sage_q_scale, sage_k_first_atom = sp.cache_sage_scales() + + def start_sage_k_words(section: FmhaStage) -> Any: + """Start obtaining the K tile's ``sfK`` words before the wait for its scores. + + The global transport loads the words, so the wait hides their + latency; the ring transport waits for the load warp's stage and + returns its offset, read after the score wait.""" + if sage_transport is SageKTransport.GMEM: + return sp.load_sage_k_words( + sage_k_first_atom=sage_k_first_atom, section=section + ) + skr.wait() + return skr.open_stage() + + def sage_row_max(row_max: Any, k_words: Any, section: FmhaStage) -> Any: + """Row max of the dequantized scores. + + A ring stage's words are read after the score wait and the + stage is released after the max pass (``SmemSageKScalesResource``).""" + if sage_transport is SageKTransport.SMEM: + k_words = sp.read_staged_sage_k_words(sage_k_stage_words=k_words) + if sage_scaling is SageScoreScaling.BLOCK: + maxima = sp.sage_row_max( + row_max=row_max, + sage_q_scale=sage_q_scale, + sage_k_words=k_words, + section=section, + ) + else: + maxima = sp.sage_dequantized_row_max( + row_max=row_max, + sage_q_scale=sage_q_scale, + sage_k_words=k_words, + section=section, + ) + if sage_transport is SageKTransport.SMEM: + skr.release() + return maxima + # Reserve a stats slot before the first softmax result is published. vec.acquire() with domain_loop(loop_start, loop_end, loop_step): + if use_sage: + sage_k_words = start_sage_k_words(FmhaStage.Loop) sp.wait() # Compute row max and publish vec. - if tmem_sp.uses_variable_window: + if use_sage: + old_row_max, row_max = sage_row_max( + row_max, sage_k_words, FmhaStage.Loop + ) + elif tmem_sp.uses_variable_window: old_row_max, row_max = sp.variable_window_row_max( row_max=row_max, window_start=window_start, @@ -2277,6 +2698,12 @@ def exp2_p(sp: TmemSPResource, *, row_max: Any, scale_softmax_log2: Any) -> Any: seqlen_k=seqlen_k, section=FmhaStage.Loop, ) + elif vc_attention: + old_row_max, row_max, vc_state["k_next"] = sp.vc_compute_row_max( + row_max=row_max, + vc_row_scale=vc_row_scale, + vc_k_scale_next=vc_state["k_next"], + ) else: old_row_max, row_max = sp.compute_row_max(row_max=row_max) if tmem_sp.cfg.corr_skip_threshold_log2 > 0: @@ -2317,6 +2744,9 @@ def exp2_p(sp: TmemSPResource, *, row_max: Any, scale_softmax_log2: Any) -> Any: sp, row_max=row_max, scale_softmax_log2=scale_softmax_log2, + old_row_max=old_row_max, + sage_q_scale=sage_q_scale, + sage_k_words=sage_k_words, ) if s0s1_seq is None or early_token: pass @@ -2327,15 +2757,44 @@ def exp2_p(sp: TmemSPResource, *, row_max: Any, scale_softmax_log2: Any) -> Any: if not p_in_smem: sp.release() # Reduction. - row_sum = sp.softmax_aux_reduce( + if vc_attention: + row_sum = vc_state["row_sum"] + else: + row_sum = sp.softmax_aux_reduce( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + p_chunk=p_chunk, + scale_softmax_log2=scale_softmax_log2, + ) + # Acquire vec for next iter. + vec.acquire() + + def publish_final_stats(row_max: Any, row_sum: Any) -> None: + """Store the last mean-group row sums or take the repair denominator, drain + the final SP slot and publish the identity stats of the work tile.""" + if vc_restores_means: + # The last row sums ride a second P handoff to the tail step. + pb.acquire() + sp.vc_store_rowsum_final( + vc_prev_tile_sum=vc_state["prev_sum"], + vc_meansums=vc_state["sums"], + ) + pb.commit() + elif vc_repairs_v: + # The repair tiles stay out of the denominator. + row_sum = vc_state["kept"] + if not p_in_smem: + sp.wait() + sp.release() + old_row_max = sp.softmax_aux_identity(row_max=row_max) + vec.store_vec( old_row_max=old_row_max, row_max=row_max, row_sum=row_sum, - p_chunk=p_chunk, - scale_softmax_log2=scale_softmax_log2, + final_stats=True, ) - # Acquire vec for next iter. - vec.acquire() + vec.commit() if tmem_sp.uses_head_paired_causal_tail_mask: # Head-paired maps Q0/Q1 to adjacent Hq slices at the same S @@ -2370,13 +2829,16 @@ def exp2_p(sp: TmemSPResource, *, row_max: Any, scale_softmax_log2: Any) -> Any: else: seq.release() sp.release() - row_sum = sp.softmax_aux_reduce( - old_row_max=old_row_max, - row_max=row_max, - row_sum=row_sum, - p_chunk=p_chunk, - scale_softmax_log2=scale_softmax_log2, - ) + if vc_attention: + row_sum = vc_state["row_sum"] + else: + row_sum = sp.softmax_aux_reduce( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + p_chunk=p_chunk, + scale_softmax_log2=scale_softmax_log2, + ) vec.acquire() sp.wait() sp.release() @@ -2424,13 +2886,16 @@ def exp2_p(sp: TmemSPResource, *, row_max: Any, scale_softmax_log2: Any) -> Any: seq.commit() if not p_in_smem: sp.release() - row_sum = sp.softmax_aux_reduce( - old_row_max=old_row_max, - row_max=row_max, - row_sum=row_sum, - p_chunk=p_chunk, - scale_softmax_log2=scale_softmax_log2, - ) + if vc_attention: + row_sum = vc_state["row_sum"] + else: + row_sum = sp.softmax_aux_reduce( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + p_chunk=p_chunk, + scale_softmax_log2=scale_softmax_log2, + ) if tmem_sp.uses_query_paired_invalid_tail: sp.wait() old_row_max, row_max = sp.invalid_row_max(row_max=row_max) @@ -2497,13 +2962,16 @@ def exp2_p(sp: TmemSPResource, *, row_max: Any, scale_softmax_log2: Any) -> Any: seq.release() if not p_in_smem: sp.release() - row_sum = sp.softmax_aux_reduce( - old_row_max=old_row_max, - row_max=row_max, - row_sum=row_sum, - p_chunk=p_chunk, - scale_softmax_log2=scale_softmax_log2, - ) + if vc_attention: + row_sum = vc_state["row_sum"] + else: + row_sum = sp.softmax_aux_reduce( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + p_chunk=p_chunk, + scale_softmax_log2=scale_softmax_log2, + ) if not p_in_smem: sp.wait() sp.release() @@ -2518,11 +2986,24 @@ def exp2_p(sp: TmemSPResource, *, row_max: Any, scale_softmax_log2: Any) -> Any: vec.commit() elif tmem_sp.uses_fixed_dense_k_tail_mask: # Dense tail: one more loop step with the zero-filled lanes masked out. + if use_sage: + sage_k_words = start_sage_k_words(FmhaStage.Tail) sp.wait() - old_row_max, row_max = sp.fixed_dense_k_tail_masked_row_max( - row_max=row_max, - section=FmhaStage.Tail, - ) + if vc_attention: + old_row_max, row_max = sp.vc_fixed_dense_k_tail_masked_row_max( + row_max=row_max, + vc_row_scale=vc_row_scale, + vc_k_scale_next=vc_state["k_next"], + ) + elif use_sage: + old_row_max, row_max = sage_row_max( + row_max, sage_k_words, FmhaStage.Tail + ) + else: + old_row_max, row_max = sp.fixed_dense_k_tail_masked_row_max( + row_max=row_max, + section=FmhaStage.Tail, + ) vec.store_vec( old_row_max=old_row_max, row_max=row_max, @@ -2541,6 +3022,10 @@ def exp2_p(sp: TmemSPResource, *, row_max: Any, scale_softmax_log2: Any) -> Any: sp, row_max=row_max, scale_softmax_log2=scale_softmax_log2, + old_row_max=old_row_max, + is_tail=True, + sage_q_scale=sage_q_scale, + sage_k_words=sage_k_words, ) if s0s1_seq is None: pass @@ -2550,39 +3035,21 @@ def exp2_p(sp: TmemSPResource, *, row_max: Any, scale_softmax_log2: Any) -> Any: seq.release() if not p_in_smem: sp.release() - row_sum = sp.softmax_aux_reduce( - old_row_max=old_row_max, - row_max=row_max, - row_sum=row_sum, - p_chunk=p_chunk, - scale_softmax_log2=scale_softmax_log2, - ) + if vc_attention: + row_sum = vc_state["row_sum"] + else: + row_sum = sp.softmax_aux_reduce( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + p_chunk=p_chunk, + scale_softmax_log2=scale_softmax_log2, + ) vec.acquire() - # Cleanup: drain the final SP slot and publish identity stats. - if not p_in_smem: - sp.wait() - sp.release() - old_row_max = sp.softmax_aux_identity(row_max=row_max) - vec.store_vec( - old_row_max=old_row_max, - row_max=row_max, - row_sum=row_sum, - final_stats=True, - ) - vec.commit() + publish_final_stats(row_max, row_sum) else: # Non-causal tail: no more tiles, just publish the final stats. - if not p_in_smem: - sp.wait() - sp.release() - old_row_max = sp.softmax_aux_identity(row_max=row_max) - vec.store_vec( - old_row_max=old_row_max, - row_max=row_max, - row_sum=row_sum, - final_stats=True, - ) - vec.commit() + publish_final_stats(row_max, row_sum) @schedule def softmax_schedule( @@ -2613,7 +3080,28 @@ def softmax_schedule_pb( ) -> None: softmax_schedule_body(sp, vec, seq, wq, None, pb) - if p_in_smem: + @schedule + def softmax_schedule_pb_skr( + sp: TmemSPResource, + vec: TmemStatsResource, + seq: S0S1SequenceResource, + pb: SmemPResource, + skr: SmemSageKScalesResource, + wq: WorkQueue | None = None, + ) -> None: + softmax_schedule_body(sp, vec, seq, wq, None, pb, skr) + + if p_in_smem and sage_transport is SageKTransport.SMEM: + captured_schedule = _schedule_with_work_queue( + softmax_schedule_pb_skr, + tmem_sp, + tmem_vec, + s0s1_seq, + smem_p, + smem_sage_k, + work_queue=work_queue, + ) + elif p_in_smem: captured_schedule = _schedule_with_work_queue( softmax_schedule_pb, tmem_sp, 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 bdaf9790be14..77aa2049afe7 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,23 @@ import cutlass import cutlass.cute as cute -from cutlass import Int32, Int64 +from cutlass import Float16, Float32, Int32, Int64, Uint32 +from cutlass.cutlass_dsl import Constexpr +from cutlass.experimental.task_scheduling.resources import StageInfo + + +@cute.jit +def _resolve_kv_tile_idx_context( + stage_info: StageInfo, + kv_tile_start: Int32, + tile_offset: cutlass.Constexpr[int] = 0, +) -> Int32: + """Resolve the context KV tile, including a compile-time pipeline offset. + + Sliding-window prefix skipping is already included in ``kv_tile_start``; + context does not apply a multi-CTA-KV split transform. + """ + return kv_tile_start + Int32(stage_info.loop_offset) + Int32(tile_offset) @cute.jit @@ -94,3 +110,129 @@ def freeze_smem_descriptor(desc): write_only_types=[Int64], read_only_args=[desc], ) + + +@cute.jit +def split_bf16_hi_lo_word(value: Float32) -> Int32: + """Pack ``value`` as two bf16 halves: low 16 bits = truncated high part, + high 16 bits = truncated remainder ``value - hi``.""" + bits = cutlass.Vector.from_elements((value,), Float32).bitcast(Uint32)[0] + hi_bits = bits >> Uint32(16) + hi = cutlass.Vector.from_elements((hi_bits << Uint32(16),), Uint32).bitcast( + Float32 + )[0] + lo_bits = cutlass.Vector.from_elements((value - hi,), Float32).bitcast(Uint32)[0] + word = hi_bits | (lo_bits & Uint32(0xFFFF0000)) + return cutlass.Vector.from_elements((word,), Uint32).bitcast(Int32)[0] + + +@cute.jit +def expcast_e4m3_quad_relu( + c0: Float32, + c1: Float32, + c2: Float32, + c3: Float32, + acc01: Int32, + acc23: Int32, + init: Constexpr[bool], +) -> tuple[Int32, Int32, Int32]: + """ExpCast four unclamped codes into one packed E4M3 word and accumulate + their values into two f16x2 row-sum registers. + + With ``init`` the quad starts a chain and the accumulators are ignored. The + returned sums are then the quantized values themselves. + + ``c = 8 * (s * scale - m) + VC_EXPCAST_CODE_BIAS`` + + ``cvt.rn.relu.f16x2.f32`` packs a pair and clamps underflow to code 0, an + f16 FMA adds 1024 so the low byte of each half is the rounded code, one + ``prmt`` gathers the four bytes, and the E4M3-pair to f16x2 conversion + with one f16 FMA per pair accumulates the quantized values. + """ + sum_pair = ( + "mov.b32 {$w1}, v01;\n mov.b32 {$w2}, v23;" + if init + else "fma.rn.f16x2 {$w1}, v01, one2, {$r4};\n" + " fma.rn.f16x2 {$w2}, v23, one2, {$r5};" + ) + return cute.arch.inline_ptx( + """ + { + .reg .b32 h01, h23, y01, y23, v01, v23, one2, k1024; + .reg .b16 c01, c23; + mov.b32 one2, 0x3C003C00; + mov.b32 k1024, 0x64006400; + cvt.rn.relu.f16x2.f32 h01, {$r1}, {$r0}; + cvt.rn.relu.f16x2.f32 h23, {$r3}, {$r2}; + fma.rn.f16x2 y01, h01, one2, k1024; + fma.rn.f16x2 y23, h23, one2, k1024; + prmt.b32 {$w0}, y01, y23, 0x6420; + mov.b32 {c01, c23}, {$w0}; + cvt.rn.f16x2.e4m3x2 v01, c01; + cvt.rn.f16x2.e4m3x2 v23, c23; + """ + + sum_pair + + """ + }""", + write_only_types=[Int32, Int32, Int32], + read_only_args=[c0, c1, c2, c3, acc01, acc23], + ) + + +@cute.jit +def f16x2_sum4(a0: Int32, a1: Int32, a2: Int32, a3: Int32) -> Float32: + """fp32 sum of the eight fp16 halves of four f16x2 registers.""" + f = [] + for word in (a0, a1, a2, a3): + halves = cutlass.Vector.from_elements((word,), Int32).bitcast(Float16) + f.append(halves[0].to(Float32)) + f.append(halves[1].to(Float32)) + return ((f[0] + f[1]) + (f[2] + f[3])) + ((f[4] + f[5]) + (f[6] + f[7])) + + +@cute.jit +def exp2_f16x2_word(x0: Float32, x1: Float32) -> Int32: + """exp2 of two fp32 values as one f16x2 word, x0 in the low half.""" + x_f16x2 = cutlass.Vector.from_elements((x0, x1), Float32).to(cutlass.Float16) + return cute.arch.exp2_packed_f16x2(x_f16x2.bitcast(Int32)[0]) + + +@cute.jit +def add_f16x2_to_f32x2( + acc: tuple[Float32, Float32], word: Int32 +) -> tuple[Float32, Float32]: + """Add the two f16 halves of a word to two fp32 accumulators with one + add.f32x2.f16x2.f32x2 instruction.""" + acc_i64 = cutlass.Vector.from_elements((acc[0], acc[1]), Float32).bitcast( + cutlass.Int64 + )[0] + res_i64 = cute.arch.inline_ptx( + "add.f32x2.f16x2.f32x2 {$w0}, {$r0}, {$r1};", + write_only_types=[cutlass.Int64], + read_only_args=[word, acc_i64], + ) + res = cutlass.Vector.from_elements((res_i64,), cutlass.Int64).bitcast(Float32) + return (res[0], res[1]) + + +@cute.jit +def f16x2_to_f32x2(word: Int32) -> tuple[Float32, Float32]: + """The two f16 halves of a word as fp32.""" + halves = cutlass.Vector.from_elements((word,), Int32).bitcast(cutlass.Float16) + halves_f32 = halves.to(Float32) + return (halves_f32[0], halves_f32[1]) + + +@cute.jit +def pack_f16x2x2_to_fp8_e4m3(w0: Int32, w1: Int32) -> Int32: + """Pack two f16x2 words into one e4m3x4 word, w0 in the low half.""" + return cute.arch.inline_ptx( + """{ + .reg .b16 lo, hi; + cvt.rn.satfinite.e4m3x2.f16x2 lo, {$r0}; + cvt.rn.satfinite.e4m3x2.f16x2 hi, {$r1}; + mov.b32 {$w0}, {lo, hi}; + }""", + write_only_types=[Int32], + read_only_args=[w0, w1], + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/helpers_paged.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/helpers_paged.py index 195d7b50c9ee..897822ec726b 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/helpers_paged.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/helpers_paged.py @@ -23,8 +23,6 @@ import cutlass.cute as cute from cutlass import Int32 -from cutlass.experimental.task_scheduling.resources import StageInfo - from .fmha_resources import FmhaConfig @@ -59,16 +57,3 @@ def _load_block_table_row_bounds( row_begin = batch_coord * block_table_row_stride return row_begin, _runtime_last_valid_page_idx(cfg, seq_len_kv) - - -@cute.jit -def _resolve_kv_tile_idx_context( - stage_info: StageInfo, - kv_tile_start: Int32, - tile_offset: cutlass.Constexpr[int] = 0, -) -> Int32: - # Context's current K/V tile index is kv_tile_start + loop_offset, with an - # optional compile-time shift for the staged D>128 schedule's previous V. - # No multi-CTA-KV transform; sliding-window prefix skipping is already - # folded into kv_tile_start by GmemQKVResource. - return kv_tile_start + Int32(stage_info.loop_offset) + Int32(tile_offset) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/sage_scales.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/sage_scales.py new file mode 100644 index 000000000000..75fbfa35c45d --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/sage_scales.py @@ -0,0 +1,233 @@ +# 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. + +"""Sage attention scale addressing for the context softmax and epilogue. + +The Sage prepare writes one FP32 buffer per plan before attention, laid out +as ``_SageScaleLayout`` describes; a context Q tile holds 128 tokens of one Q +head. A softmax lane reads its ``sfQ`` word once per work tile with +``sage.index.load_sage_q_scale``, which the decode kernel shares. The helpers +here turn work-tile coordinates into words of that buffer and move the other +words the kernel reads: + +* ``load_sage_k_tile`` reads one K tile's ``sfK`` atoms from the tile's + address (``sage_k_section_address`` gives the sequence's first atom) for + the softmax lanes; ``sage_k_transport`` says whether the plan's K block + size takes that global-memory path or the load warp's shared-memory ring + (``SmemSageKScalesResource``), and ``sage_score_scaling`` how the softmax + applies the words; + ``prefetch_sage_k_tile`` warms L1 with the next tile's atoms one tile ahead + of direct loads; +* ``load_sage_v_channels`` reads the sequence's ``sfV`` / ``v_mean`` words per + output column in the epilogue (``sage_v_section_address`` gives the + sequence's first ``sfV`` word). +""" + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Int64 + +from ...sage.index import sage_section_word +from ...sage.layout import _FLOAT4_WORDS, _SAGE_ATOM_TOKENS, _SageScaleLayout +from ...sage.modes import SageKTransport, SageScoreScaling + +Constexpr = cutlass.Constexpr + +# Score columns that share one ``sfK`` word when the softmax scales blocks +# (``SageScoreScaling.BLOCK``), the width of one LDTM.STAT piece. +SAGE_SCORE_BLOCK_COLS = 16 +# Bytes of one L1 line, the granularity of ``prefetch_sage_k_tile``. +_L1_LINE_BYTES = 128 +# Most ``sfK`` words of one K tile the softmax lanes load from global memory +# and carry across the S wait as a task-local vector +# (``SageKTransport.GMEM``); more words go through the load warp's ring. +_GMEM_MAX_TILE_WORDS = 32 + + +def validate_context_sage_layout(layout: _SageScaleLayout, kv_tile_n: int) -> None: + """Reject a layout the context softmax cannot read tile by tile. + + The K-token atoms must be padded to whole K tiles, since the softmax + addresses each tile's atoms without clamping. + """ + if layout.k_tokens.tile_atoms * _SAGE_ATOM_TOKENS != kv_tile_n: + raise ValueError( + "the Sage K-token atoms must be padded to the kernel's K tiles" + ) + if layout.q_rows.heads_interleaved: + raise ValueError("the context kernel reads one Q head per 128 Q rows") + + +def sage_score_scaling(layout: _SageScaleLayout) -> SageScoreScaling: + """Return how the context softmax applies the layout's ``sfK`` words. + + K blocks of 16 tokens or more give one word per 16-column score block: + the scores stay raw and each block's maximum and exp2 multiplier carry + ``sfQ * sfK``. Finer blocks need one multiplier per word of fewer + columns, so the max pass multiplies the scores in place and the P pass + runs the per-tensor path with ``sfQ`` in its exp2 scale. + """ + if layout.k_tokens.tokens_per_word >= SAGE_SCORE_BLOCK_COLS: + return SageScoreScaling.BLOCK + return SageScoreScaling.DEQUANTIZED + + +def sage_k_transport(layout: _SageScaleLayout) -> SageKTransport: + """Return how the context softmax obtains a K tile's ``sfK`` words. + + The lanes load a tile of at most ``_GMEM_MAX_TILE_WORDS`` words from + global memory before the S wait; a tile with more words (one-token + blocks) is too many to carry or to broadcast-load per lane, so the load + warp stages it in the ``sfK`` ring (``SmemSageKScalesResource``). + """ + if sage_k_tile_words(layout) <= _GMEM_MAX_TILE_WORDS: + return SageKTransport.GMEM + return SageKTransport.SMEM + + +def sage_k_tile_words(layout: _SageScaleLayout) -> int: + """Return the ``sfK`` words of one K tile: its atoms, each padded to float4s.""" + return layout.k_tokens.tile_atoms * layout.k_tokens.atom_words + + +def sage_k_tile_bytes(layout: _SageScaleLayout) -> int: + """Return the bytes of one K tile's ``sfK`` atoms.""" + return sage_k_tile_words(layout) * 4 + + +def sage_k_block_word(layout: _SageScaleLayout, block_idx: int) -> int: + """Return the tile word holding ``sfK`` of score block ``block_idx``. + + Block ``b`` starts at token ``16 b`` of the tile: word + ``token_in_atom // tokens_per_word`` of atom ``token // 64``. + """ + k_tokens = layout.k_tokens + atom, token_in_atom = divmod(block_idx * SAGE_SCORE_BLOCK_COLS, _SAGE_ATOM_TOKENS) + return atom * k_tokens.atom_words + token_in_atom // k_tokens.tokens_per_word + + +@cute.jit +def _section_address( + scales: cute.Pointer, sequence_words: Int32, section: Constexpr +) -> Int64: + """Return the address of the first word of a sequence's section.""" + return scales.toint() + Int64(sage_section_word(sequence_words, section)) * Int64(4) + + +@cute.jit +def sage_k_section_address( + layout: Constexpr[_SageScaleLayout], + scales: cute.Pointer, + *, + sequence_words: Int32, +) -> Int64: + """Return the address of a sequence's first ``sfK`` atom. + + K tile ``t`` starts ``t * sage_k_tile_bytes`` after it. The section is + padded to whole tiles, so every tile's atoms lie inside it; a tile past + the sequence end reads the pad atoms, which repeat the last scale. + """ + return _section_address(scales, sequence_words, layout.k_tokens) + + +@cute.jit +def sage_v_section_address( + layout: Constexpr[_SageScaleLayout], + scales: cute.Pointer, + *, + sequence_words: Int32, +) -> Int64: + """Return the address of a sequence's first ``sfV`` word (``_SageVSection``).""" + return _section_address(scales, sequence_words, layout.v_rows) + + +@cute.jit +def _load_float4_at(address: Int64): + """Return the four words at a 16-byte aligned global address.""" + value_ptr = cutlass.inttoptr(address, mem_space=1, dtype=Float32) + return value_ptr.load(count=_FLOAT4_WORDS, alignment=16) + + +@cute.jit +def load_sage_k_tile( + layout: Constexpr[_SageScaleLayout], tile_address: Int64 +) -> cutlass.Array: + """Return the ``sfK`` words of the K tile at ``tile_address``. + + One float4 load per four words; every lane of a warp reads the same + words, so the loads are warp broadcasts that hit L1 once + ``prefetch_sage_k_tile`` has warmed the tile's lines. + """ + tile_words = sage_k_tile_words(layout) + words = cutlass.Array(Float32, tile_words, space=cutlass.AddressSpace.rmem) + for word_base in cutlass.range_constexpr(0, tile_words, _FLOAT4_WORDS): + loaded = _load_float4_at(tile_address + Int64(word_base * 4)) + for elem in cutlass.range_constexpr(_FLOAT4_WORDS): + words[word_base + elem] = Float32(loaded[elem]) + return words + + +@cute.jit +def prefetch_sage_k_tile( + layout: Constexpr[_SageScaleLayout], tile_address: Int64 +) -> None: + """Bring the K tile at ``tile_address`` into L1 for a later ``load_sage_k_tile``. + + The tile after a sequence's last one lies in the sequence's Q rows, so + the prefetch of a loop's final iteration stays inside the buffer. + """ + for offset in cutlass.range_constexpr(0, sage_k_tile_bytes(layout), _L1_LINE_BYTES): + # A tile is a whole number of 32-byte units, so it never straddles a + # line boundary that this stride misses. + cute.arch.inline_ptx( + "prefetch.global.L1 [{$r0}];", + read_only_args=[tile_address + Int64(offset)], + ) + + +@cute.jit +def load_sage_v_channels( + layout: Constexpr[_SageScaleLayout], + channel_address: Int64, + *, + offset: Constexpr[int], + count: Constexpr[int], +) -> tuple[cutlass.Vector, cutlass.Vector | None]: + """Return ``count`` ``sfV`` and ``v_mean`` values of consecutive channels. + + ``channel_address`` is the ``sfV`` word of the first channel the caller + handles (``sage_v_section_address`` plus its channel); ``offset`` is this + call's constant channel offset from it, so the loads take immediate + offsets from one address. One float4 per load; the correction warps + read the same words, so the loads are warp broadcasts. The means are + ``None`` for a recipe without a V mean. + """ + assert count % _FLOAT4_WORDS == 0 + mean_bytes = layout.v_rows.mean_word_offset * 4 + scale_values: tuple = () + mean_values: tuple = () + for channel in cutlass.range_constexpr(offset, offset + count, _FLOAT4_WORDS): + scale_vec = _load_float4_at(channel_address + Int64(channel * 4)) + for elem in cutlass.range_constexpr(_FLOAT4_WORDS): + scale_values += (Float32(scale_vec[elem]),) + if cutlass.const_expr(layout.v_rows.has_mean): + mean_vec = _load_float4_at( + channel_address + Int64(mean_bytes + channel * 4) + ) + for elem in cutlass.range_constexpr(_FLOAT4_WORDS): + mean_values += (Float32(mean_vec[elem]),) + scales = cutlass.Vector.from_elements(scale_values, Float32) + if cutlass.const_expr(layout.v_rows.has_mean): + return scales, cutlass.Vector.from_elements(mean_values, Float32) + return scales, None diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/vc_attention/__init__.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/vc_attention/__init__.py new file mode 100644 index 000000000000..76a82805fe5a --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/vc_attention/__init__.py @@ -0,0 +1,7 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""VC-Attention pieces of the context kernel. + +``vc_resources`` holds the tile-mean operand ring and the ExpCast helpers, +``vc_scales`` the per-block Q/K scale tables of VC-Attention-QK8. +""" diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/vc_attention/vc_resources.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/vc_attention/vc_resources.py new file mode 100644 index 000000000000..7753b1b30b1d --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/vc_attention/vc_resources.py @@ -0,0 +1,258 @@ +# 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. + +"""VC-Attention operands and helpers of the context kernel. + +VC-Attention stores E4M3 V as residuals around the bf16 mean of each +128-token K/V tile and restores the mean inside the online softmax with one +bf16 K=16 UMMA step per tile, ``O += rowsum(P_tile) * mean_tile``. +:class:`SmemMuResource` is the TMA-fed ring of the ``mean_tile`` operands. The +ExpCast helpers in ``helpers`` code the probabilities directly as E4M3 bytes +and accumulate the quantized row sums that become the step's A operand. +The resources in ``fmha_resources`` own the row-sum operand (behind the +SMEM P tile), the exp2 and store paths and the mean UMMA issue. +""" + +from ....vc_attention import VC_MEAN_GROUP_TILES, VC_MEAN_OPERANDS +from dataclasses import dataclass, field +from typing import Any, Optional + +import cutlass +import cutlass.cute as cute +from cutlass import Int16, Int32 +from cutlass.cutlass_dsl import Constexpr +from cutlass.experimental import primitives as prims +from cutlass.experimental.task_scheduling.enums import WorkAttr +from cutlass.experimental.task_scheduling.memory import SmemAllocation +from cutlass.experimental.task_scheduling.resources import ( + MemoryResource, + PipelineConfig, + StageInfo, + TaskLocalVariable, + consumer_work, + producer_work, +) + +from ...placeholder_helpers import _placeholder_smem_array + +# VC-Attention mean-restore operands of the bf16 K=16 UMMA step that adds +# ``rowsum(P_tile) * mean_tile`` into O. Unswizzled K-major 32-byte rows with +# LBO 128 between the two K core matrices and SBO 256 between 8-row groups. +VC_MEAN_TILE_LBO = 128 +VC_MEAN_TILE_SBO = 256 + +# ExpCast-FP8 (VC-Attention). An E4M3 byte read as an integer equals +# ``8*log2(v) + 56 + eps``. With ``u`` the log2-domain probability and the +# paper's 2^8 P scaling, ``code = round(8*u + VC_EXPCAST_CODE_BIAS)`` replaces +# exp2 and the E4M3 cast by one FMA. beta is the paper's centering of +# Mitchell's error. +_E4M3_EXP_BIAS = 7 +E4M3_CODES_PER_OCTAVE = 8 +_VC_EXPCAST_BETA = -0.35 +_VC_EXPCAST_P_SCALE_LOG2 = 8 +VC_EXPCAST_CODE_BIAS = ( + E4M3_CODES_PER_OCTAVE * (_E4M3_EXP_BIAS + _VC_EXPCAST_P_SCALE_LOG2) + + _VC_EXPCAST_BETA +) + + +# --------------------------------------------------------------------------- +# SmemMuResource -- VC-Attention bf16 tile-mean operand ring +# --------------------------------------------------------------------------- +@dataclass(kw_only=True) +class SmemMuResource(MemoryResource): + """Shared memory buffer for mu, the paper's name for the per-tile mean of V. + + One buffer holds the bf16 means of a group of sixteen K/V tiles, packed as + K-major UMMA operands. The means are a tensor core B operand and the MMA + instruction reads its operands through shared memory descriptors. They + therefore follow the same path as K and V in SmemKVResource and are staged + into shared memory by TMA with a TmaUmma pipeline. + + Producer: the load warp TMA-loads the operands of the group holding K/V tile + ``i-1`` beside V tile ``i``. Consumer: the UMMA warp issues two + ``kind::f16`` K=16 steps ``O += sum_j rowsum(P_j) * mean_j`` over the group's + tiles after the PV steps of the tile that follows the group, so the means + restored into O are rescaled by the same correction as the residual + product. Rows ``k=2i`` and ``k=2i+1`` of operand ``o`` hold the mean of tile + ``8o+i`` of the group (divided by the V dequant scale). + """ + + cfg: Constexpr[Any] = field(init=False, default=None) + tma_mu_desc: cutlass.Pointer | None = field(init=False, default=None) + sMu_array: cutlass.Array = field(init=False, default=None) + _alloc: Constexpr[Optional[SmemAllocation]] = field(init=False, default=None) + desc_mu_base: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + # Second live handle for the tail, where the last tile's operand is waited + # while the previous tile's is still in use. + desc_mu_last: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + + def __init__( + self, + tma_mu_desc: cutlass.Pointer | None, + pipeline_config: PipelineConfig, + cfg: Any, + **kwargs: Any, + ) -> None: + super().__init__(pipeline_config=pipeline_config, **kwargs) + self.cfg = cfg + self.tma_mu_desc = tma_mu_desc + self._alloc = SmemAllocation( + "smem_mu", + pipeline_config.num_stages * cfg.vc_mean_tile_bytes, + alignment=cfg.buffer_align_bytes, + ) + self.sMu_array = _placeholder_smem_array(cutlass.BFloat16) + self.desc_mu_base = TaskLocalVariable( + dtype=cutlass.Int64, + default=cutlass.Int64(0), + docs="SMEM descriptor base for the current tile-mean operand.", + ) + self.desc_mu_last = TaskLocalVariable( + dtype=cutlass.Int64, + default=cutlass.Int64(0), + docs="SMEM descriptor base for the last tile's mean operand.", + ) + + def get_smem_requirements(self) -> list[SmemAllocation]: + return [self._alloc] + + @property + def stage_elements(self) -> int: + return self.cfg.vc_mean_tile_bytes // 2 + + @cute.jit + def _init_smem_state(self, stage_info: StageInfo) -> None: + smem_base = stage_info.context.smem_base + self.sMu_array = cutlass.Array( + smem_base.data_ptr() + self._alloc.offset, + dtype=cutlass.BFloat16, + shape=(self.pipeline_config.num_stages * self.stage_elements,), + addrspace=3, + ) + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_load_state(self, stage_info: StageInfo) -> None: + self._init_smem_state(stage_info) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_descriptor_state(self, stage_info: StageInfo) -> None: + self._init_smem_state(stage_info) + + @property + def loop_offset_sensitive(self) -> bool: + return True + + @cute.jit + def _tma_load_tile( + self, + stage_info: StageInfo, + tile_idx: Int32, + kv_head_coord: Int32, + batch_coord: Int32, + ) -> None: + if prims.elect_sync(): + operand_rows = self.cfg.vc_mean_operand_bytes // 512 + operand_elems = ( + self.cfg.vc_mean_operand_bytes // self.cfg.cta_group_size // 2 + ) + for operand in cutlass.range_constexpr(VC_MEAN_OPERANDS): + sMu_curr = self.sMu_array.subview( + stage_info.stage_idx * self.stage_elements + operand * operand_elems + ) + row = Int32(operand * operand_rows) + if cutlass.const_expr(self.cfg.two_cta_umma): + # Each CTA stages the 4 packed rows (64 channels) its half of + # the M=256 mean step reads, like K/V. + cta_rank = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + prims.cp_async_bulk_tensor_shared_cluster_global( + sMu_curr, + self.tma_mu_desc, + ( + Int32(0), + row + cta_rank * Int32(operand_rows // 2), + tile_idx, + kv_head_coord, + batch_coord, + ), + cutlass.Array( + stage_info.barrier.data_ptr(), dtype=cutlass.Int64 + ), + [], + multicast_mask=Int16(Int32(1) << cta_rank), + group=prims.CTAGroup.CTA_2, + ) + else: + prims.cp_async_bulk_tensor_shared_cta_global( + sMu_curr, + self.tma_mu_desc, + (Int32(0), row, tile_idx, kv_head_coord, batch_coord), + stage_info.barrier, + ) + + @producer_work + @cute.jit + def mu_load( + self, + stage_info: StageInfo, + *, + kv_head_coord: Int32, + batch_coord: Int32, + kv_tile_start: Int32, + is_tail: Constexpr[bool] = False, + ) -> None: + """TMA-load the mean operand of the group holding K/V tile ``loop_offset - 1``. + + A group's mean step is issued one tile late, beside PV of the tile after + the group, so the ring runs one tile behind V. The tail call loads the + group of the last tile for the tail mean step. + """ + if cutlass.const_expr(is_tail): + tile_idx = kv_tile_start + Int32(stage_info.loop_end) - Int32(1) + else: + tile_idx = cute.math.max( + kv_tile_start + stage_info.loop_offset - Int32(1), Int32(0) + ) + self._tma_load_tile( + stage_info, + tile_idx // Int32(VC_MEAN_GROUP_TILES), + kv_head_coord, + batch_coord, + ) + + @cute.jit + def _build_desc(self, stage_info: StageInfo) -> prims.Tcgen05SmemDesc: + sMu_curr = self.sMu_array.subview(stage_info.stage_idx * self.stage_elements) + return prims.Tcgen05SmemDesc.build( + sMu_curr, + leading_byte_offset=VC_MEAN_TILE_LBO, + stride_byte_offset=VC_MEAN_TILE_SBO, + layout=prims.Tcgen05SmemSwizzle.NONE, + ) + + @consumer_work(returns=desc_mu_base) + @cute.jit + def mu_desc(self, stage_info: StageInfo) -> prims.Tcgen05SmemDesc: + """Build the unswizzled K-major descriptor of the current mean operand.""" + return self._build_desc(stage_info) + + @consumer_work(returns=desc_mu_last) + @cute.jit + def mu_desc_last(self, stage_info: StageInfo) -> prims.Tcgen05SmemDesc: + """Descriptor of the last tile's mean operand, held beside the previous one.""" + return self._build_desc(stage_info) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/vc_attention/vc_scales.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/vc_attention/vc_scales.py new file mode 100644 index 000000000000..7cc0097ba5cb --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/vc_attention/vc_scales.py @@ -0,0 +1,129 @@ +# 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. +"""VC-Attention scale addressing for the context softmax. + +Every dequant-scale load of the VC-Attention context kernel goes through this +module, the context counterpart of the decode kernels' +:mod:`..fmha_decode.fmha_decode_resources.sage_scales`. ``sfQ`` and ``sfK`` use +the trtllm-gen flat layout (``flat_scale_slot`` of +:mod:`flashinfer.attention.prims_ts.sage`): per head, sequence ``b`` of length +``S`` starts at slot ``(b * S >> log2(blk)) + b`` and token ``t`` uses slot +``t >> log2(blk)`` inside it. Q rows past the valid row count clamp to the +last valid slot, so masked rows keep a finite scale. + +:class:`VCKScaleTable` is one softmax group's ``sfK`` words for the whole K/V +loop of a work tile: the group fills it with one coalesced pass before the +loop (published by a named barrier that also keeps a fast warp from refilling +while a peer still reads the previous work tile's words) and every softmax +step reads the next tile's word one iteration ahead. The words live behind the +group's SMEM P tile (``SmemPResource.kscale_table_offset``). +""" + +from __future__ import annotations + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32 +from cutlass.experimental.task_scheduling.resources import StageInfo + +# First of the two named barrier ids, one per softmax group, that publish a +# refilled K-scale table. +VC_KSCALE_BARRIER_ID = 8 + + +@cute.jit +def vc_flat_q_slot( + batch_coord: Int32, + row: Int32, + seq_len_q: cutlass.Constexpr[int], + q_block_log2: cutlass.Constexpr[int], +) -> Int32: + """Slot of query ``row`` of sequence ``batch_coord`` in the Q scale table, which + holds one scale per block of every sequence in turn.""" + lbq = Int32(q_block_log2) + row_c = cute.math.min(row, Int32(seq_len_q - 1)) + blocks = (seq_len_q + (1 << q_block_log2) - 1) >> q_block_log2 + return batch_coord * Int32(blocks) + (row_c >> lbq) + + +@cute.jit +def vc_flat_k_base( + batch_coord: Int32, + seq_len_k: cutlass.Constexpr[int], + kv_tile_n: cutlass.Constexpr[int], +) -> Int32: + """Slot of K/V tile 0 of sequence ``batch_coord`` in the K scale table, which + holds one scale per tile of every sequence in turn. Tile ``t`` is ``base + t``.""" + return batch_coord * Int32((seq_len_k + kv_tile_n - 1) // kv_tile_n) + + +class VCKScaleTable: + """One softmax group's ``sfK`` table for the work tile, staged in SMEM.""" + + def __init__(self, cfg, smem_p, q_half: int) -> None: + self.cfg = cfg + self.smem_p = smem_p + self.q_half = q_half + + @property + def words(self) -> int: + return self.cfg.vc_kscale_table_bytes // 4 + + @cute.jit + def view(self, stage_info: StageInfo) -> cutlass.Array: + """The table as an fp32 SMEM array.""" + context = stage_info.context + return cutlass.Array( + context.smem_base.data_ptr() + + self.smem_p._alloc.offset + + self.smem_p.kscale_table_offset, + dtype=Float32, + shape=(self.words,), + addrspace=3, + ) + + @cute.jit + def fill( + self, + stage_info: StageInfo, + vc_k_scale: cute.Tensor, + kv_head: Int32, + k_base: Int32, + num_tiles: cutlass.Constexpr[int], + ) -> Float32: + """Gather this work tile's ``sfK`` words from the flat layout and publish + the table to the group's four warps; returns tile 0's scale.""" + cfg = self.cfg + num_threads = len(cfg.softmax0_warp_ids) * cute.arch.WARP_SIZE + barrier_id = VC_KSCALE_BARRIER_ID + self.q_half + table = self.view(stage_info) + num_tiles = Int32(num_tiles) + warp_id_in_sg = cute.arch.warp_idx() % len(cfg.softmax0_warp_ids) + thread = warp_id_in_sg * cute.arch.WARP_SIZE + cute.arch.lane_idx() + cute.arch.barrier(barrier_id=barrier_id, number_of_threads=num_threads) + for pass_idx in cutlass.range_constexpr( + (self.words + num_threads - 1) // num_threads + ): + tile_idx = Int32(pass_idx * num_threads) + thread + if tile_idx < num_tiles: + table[tile_idx] = Float32(vc_k_scale[(kv_head, k_base + tile_idx)]) + cute.arch.barrier(barrier_id=barrier_id, number_of_threads=num_threads) + return Float32(table[Int32(0)]) + + @cute.jit + def next_scale(self, stage_info: StageInfo, tile_idx: Int32) -> Float32: + """``sfK`` of tile ``tile_idx + 1`` (clamped to the table), read one tile ahead.""" + table = self.view(stage_info) + next_idx = cute.math.min(tile_idx + Int32(1), Int32(self.words) - Int32(1)) + return Float32(table[next_idx]) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/direct_sparse_metadata.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/direct_sparse_metadata.py index 8bffb076c4a3..ec2ee534739e 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/direct_sparse_metadata.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/direct_sparse_metadata.py @@ -21,17 +21,21 @@ class HeadIndexedMetadataView: """Read one KV head's values from interleaved [route, pattern_head] rows.""" + # Bind interleaved route/head metadata and the selected KV-head index. def __init__(self, values, heads, head): self.values = (values, heads, head) + # Flatten runtime tensors and indices for the CuTe value protocol. def __extract_mlir_values__(self): return cutlass.extract_mlir_values(self.values) + # Rebind runtime metadata after a CuTe call or control-flow boundary. def __new_from_mlir_values__(self, values): return HeadIndexedMetadataView( *cutlass.new_from_mlir_values(self.values, values) ) + # Read the selected head at a logical route index. @cute.jit def __getitem__(self, route): values, heads, head = self.values @@ -47,6 +51,11 @@ class DirectSparseMetadataView: requests may each contain several routes. Split reduction can still run. """ + # Bind block IDs, page table, request IDs and token positions without copying. + # model_len and fixed_query_len count tokens; page_capacity counts fragment slots. + # packed selects route offsets; fixed_query_layout derives request/position from + # uniform Q rows instead of request IDs. lengths selects scheduler lengths + # rather than physical fragment locators; head selects an unshared pattern. def __init__( self, inputs, @@ -60,35 +69,66 @@ def __init__( sparse_block_size=4, fragment_size=4, head=0, + count_current_block_in_topk=True, + fixed_query_len=0, + fixed_query_layout=False, ): - self.values = (*inputs, qo_indptr, head) + self.values = ( + *inputs[:4], + qo_indptr, + head, + model_len, + page_capacity, + fixed_query_len, + ) self.packed = packed - self.model_len = model_len self.page_size = page_size - self.page_capacity = page_capacity self.lengths = lengths self.sparse_block_size = sparse_block_size + self.count_current_block_in_topk = count_current_block_in_topk + self.fixed_query_layout = fixed_query_layout self.fragment_size = fragment_size self.fragments_per_block = sparse_block_size // fragment_size + # Return the runtime exclusive token-position bound. + @property + def model_len(self): + return self.values[6] + + # Return fragment slots per route, the flattened locator row stride. + @property + def page_capacity(self): + return self.values[7] + + # Return tokens per request for the fixed-query metadata layout. + @property + def fixed_query_len(self): + return self.values[8] + + # Carry tensors and live bounds through the CuTe value protocol. def __extract_mlir_values__(self): return cutlass.extract_mlir_values(self.values) + # Rebind runtime values while preserving the static sparse metadata format. def __new_from_mlir_values__(self, values): new = cutlass.new_from_mlir_values(self.values, values) return DirectSparseMetadataView( new[:4], new[4], packed=self.packed, - model_len=self.model_len, + model_len=new[6], page_size=self.page_size, - page_capacity=self.page_capacity, + page_capacity=new[7], lengths=self.lengths, sparse_block_size=self.sparse_block_size, fragment_size=self.fragment_size, + count_current_block_in_topk=self.count_current_block_in_topk, + fixed_query_len=new[8], + fixed_query_layout=self.fixed_query_layout, head=new[5], ) + # Share the metadata and live bounds while selecting another KV-head pattern. def with_head(self, head): return DirectSparseMetadataView( self.values[:4], @@ -100,12 +140,17 @@ def with_head(self, head): lengths=self.lengths, sparse_block_size=self.sparse_block_size, fragment_size=self.fragment_size, + count_current_block_in_topk=self.count_current_block_in_topk, + fixed_query_len=self.fixed_query_len, + fixed_query_layout=self.fixed_query_layout, head=head, ) + # Resolve a route to (query row, request, visible KV tokens); invalid routes have + # zero visibility. @cute.jit def resolve_route(self, route): - blocks, table, requests, positions, offsets, _ = self.values + blocks, table, requests, positions, offsets, _ = self.values[:6] row = route valid = (row >= Int32(0)) & (row < Int32(blocks.shape[0])) if cutlass.const_expr(self.packed): @@ -119,8 +164,20 @@ def resolve_route(self, route): request = Int32(-1) visible = Int32(0) if valid: - request = Int32(requests[row]) - position = Int64(positions[row]) + if cutlass.const_expr(self.fixed_query_layout): + request = row // Int32(self.fixed_query_len) + # Fixed Q ends at the per-request KV length. + length = Int64(positions[request]) + position = Int64(-1) + if length > Int64(0): + position = ( + length + - Int64(self.fixed_query_len) + + Int64(row % Int32(self.fixed_query_len)) + ) + else: + request = Int32(requests[row]) + position = Int64(positions[row]) valid = ( (request >= Int32(0)) & (request < Int32(table.shape[0])) @@ -131,12 +188,21 @@ def resolve_route(self, route): visible = Int32(position + Int64(1)) return row, request, visible + @cute.jit + def selected_block_bound(self, visible): + """Return the exclusive top-k block bound under the current-block policy.""" + eligible = visible + if cutlass.const_expr(not self.count_current_block_in_topk): + eligible = cute.math.max(visible - Int32(1), Int32(0)) + return eligible // Int32(self.sparse_block_size) + + # Read a selected block ID and map one fragment slot to its physical locator. @cute.jit def load_page(self, route_state, slot): blocks = self.values[0] row, _, visible = route_state complete = cute.math.min( - visible // Int32(self.sparse_block_size), Int32(blocks.shape[-1]) + self.selected_block_bound(visible), Int32(blocks.shape[-1]) ) logical = Int32(-1) candidate = slot // Int32(self.fragments_per_block) @@ -147,25 +213,26 @@ def load_page(self, route_state, slot): logical = Int32(blocks[row, candidate]) return self.map_page(route_state, slot, logical) + # Map a block ID and fragment slot through the page table; invalid slots return -1. @cute.jit def map_page(self, route_state, slot, logical): - blocks, table, _, _, _, _ = self.values + blocks, table, _, _, _, _ = self.values[:6] _, request, visible = route_state complete = cute.math.min( - visible // Int32(self.sparse_block_size), Int32(blocks.shape[-1]) + self.selected_block_bound(visible), Int32(blocks.shape[-1]) + ) + tail = visible - self.selected_block_bound(visible) * Int32( + self.sparse_block_size ) - tail = visible % Int32(self.sparse_block_size) candidate = slot // Int32(self.fragments_per_block) result = Int32(-1) if candidate < complete: - if (logical < Int32(0)) | ( - logical >= visible // Int32(self.sparse_block_size) - ): + if (logical < Int32(0)) | (logical >= self.selected_block_bound(visible)): logical = Int32(-1) elif (candidate == complete) & ( slot % Int32(self.fragments_per_block) * Int32(self.fragment_size) < tail ): - logical = visible // Int32(self.sparse_block_size) + logical = self.selected_block_bound(visible) else: logical = Int32(-1) if logical >= Int32(0): @@ -181,15 +248,18 @@ def map_page(self, route_state, slot, logical): result = physical * fragments_per_page + subpage return result + # Expose scheduler token lengths or flattened physical fragment locators. @cute.jit def __getitem__(self, index): if cutlass.const_expr(self.lengths): _, _, visible = self.resolve_route(Int32(index)) complete = cute.math.min( - visible // Int32(self.sparse_block_size), + self.selected_block_bound(visible), Int32(self.values[0].shape[-1]), ) - tail = visible % Int32(self.sparse_block_size) + tail = visible - self.selected_block_bound(visible) * Int32( + self.sparse_block_size + ) result = cute.math.max( complete * Int32(self.sparse_block_size) + tail, Int32(1) ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_config.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_config.py index 42ba3d0e971d..e492e92a92e1 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 @@ -28,20 +28,29 @@ import cutlass.utils as utils from cutlass import BFloat16, Float16, Float32, Float4E2M1FN, Float8E4M3FN, Int8 -from ...sage import SAGE_K_BLOCK_SIZES, is_power_of_two +from ...sage.config import SAGE_K_BLOCK_SIZES, is_power_of_two from ..._block_sparse.common import ( _block_sparse_kv_atom_size, _block_sparse_proxy_summary_geometry, _select_block_sparse_q_tile_size, _validate_sparse_kv_block_size, ) -from ...sage import _SAGE_ATOM_MIN_WORDS, _SAGE_ATOM_TOKENS, _SageScaleLayout +from ...sage.layout import ( + _SAGE_ATOM_MIN_WORDS, + _SAGE_ATOM_TOKENS, + _SageKSection, + _SageQSection, + _SageScaleLayout, + _SageVSection, +) +from ...sage.modes import SageScoreScaling 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, + KV_TILE_256_MAX_EX2_EMULATED_PAIRS, BITS_PER_BYTE, BYTES_PER_KIB, FALLBACK_SM_COUNT_B200, @@ -80,6 +89,15 @@ WARP_THREADS, ) +LENGTH_DERIVED_FIELDS = ( + "max_seq_len_q", + "total_kv_tiles", + "static_seq_len_kv", + "max_num_pages_per_seq_kv", + "static_num_skipped_kv_tiles", + "static_window_start_idx", +) + ConfigValue = int | float | bool | str | type | None _GroupedKeepsProfileKey = tuple[type, type, type, type, int, int, int, int] @@ -111,6 +129,17 @@ _GROUPED_KEEPS_MAIN_PROFILE, (BFloat16, BFloat16, BFloat16, BFloat16, 128, 0, 2, 2), } +_D192_GROUPED_KEEPS_PROFILES = { + (dtype, dtype, dtype, output, 192, stage, 1, 1) + for dtype, output in ( + (Float16, Float16), + (BFloat16, BFloat16), + (Float8E4M3FN, Float16), + (Float8E4M3FN, BFloat16), + (Float8E4M3FN, Float8E4M3FN), + ) + for stage in ((128, 192) if dtype in (Float16, BFloat16) else (128,)) +} # Sage recipes: 8-bit Q/K, E4M3 V and a 16-bit output. ``use_sage_attention`` # selects this table because the same dtypes without scales name the # static-only FP8 profile below. @@ -175,6 +204,9 @@ (Float8E4M3FN, BFloat16), ) } +_Q_TOKEN_KV_BLOCK_SPARSE_GROUPED_KEEPS_PROFILES.update( + profile for profile in _D192_GROUPED_KEEPS_PROFILES if profile[3] != Float8E4M3FN +) _Q_TOKEN_KV_BLOCK_SPARSE_GROUPED_KEEPS_PROFILES.update( (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, output, 256, 256, 1, 1) for output in (Float16, BFloat16) @@ -262,6 +294,17 @@ def _dtype_bytes(dtype: type) -> int: (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, Float16, 256, 128, 1, 1), (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, BFloat16, 256, 128, 1, 1), } +_AUTO_GROUPED_KEEPS_PAGED_FP8_PROFILES.update( + profile for profile in _D192_GROUPED_KEEPS_PROFILES if profile[0] == Float8E4M3FN +) + +# FP16/BF16 grouped Keeps Q128 recipes on paged K/V: the contiguous 16-bit Q128 +# resources with the paged FP8 Q128 K/V loads. Paged 16-bit Q64 keeps its own +# profile below. +_GROUPED_KEEPS_PAGED_16BIT_Q128_PROFILES = { + _GROUPED_KEEPS_MAIN_PROFILE, + (BFloat16, BFloat16, BFloat16, BFloat16, 128, 0, 2, 2), +} def normalize_mask_type( @@ -413,19 +456,20 @@ class GroupedQLaunchCandidate: # The selector reuses the measured relative costs but resolves its own legal # profiles, Q geometry, split fanout, and cluster promotion instead of copying # a final shape policy. +# Both tables are keyed by a K/V width (``_grouped_q_cost_kv_width``), then by +# TileQ. Width 16 scores 16-bit K and V at head dim 128, the one head dim whose +# TileQ64 candidate is promoted to the Q64/KV256 CTA; every other profile, +# 16-bit D192 included, takes width 8. The widths share every factor but one: +# a 16-bit D128 Keeps Q128 CTA runs its MMAs at half the 8-bit rate, so +# TileQ128 is scored against the Q64/KV256 CTA that the Q64 factor stands for +# after the KV256 promotion rather than with the shared TileQ128 factor. _GROUPED_Q_MAINLOOP_COST = { - 8: 1.0, - 16: 1.2, - 32: 1.48, - 64: 1.68, - 128: 2.2, + 8: {8: 1.0, 16: 1.2, 32: 1.48, 64: 1.68, 128: 2.2}, + 16: {8: 1.0, 16: 1.2, 32: 1.48, 64: 1.68, 128: 3.2}, } _GROUPED_Q_REDUCTION_COST = { - 8: 1.0, - 16: 1.03, - 32: 1.08, - 64: 1.2, - 128: 1.32, + 8: {8: 1.0, 16: 1.03, 32: 1.08, 64: 1.2, 128: 1.32}, + 16: {8: 1.0, 16: 1.03, 32: 1.08, 64: 1.2, 128: 1.32}, } _GROUPED_Q_REDUCTION_SEQ_LEN_FACTOR = 128.0 @@ -475,6 +519,7 @@ def make_grouped_q_launch_candidate( batch_size: int, num_heads_kv: int, service_capacity: int, + cost_kv_width: int, ) -> GroupedQLaunchCandidate: """Build one empirical grouped-Q launch-cost record. @@ -483,9 +528,10 @@ def make_grouped_q_launch_candidate( ``(mainloop_factor[TileQ] * seq_len_per_cta_kv +`` `` reduction_factor[TileQ] * 128 * splits_kv) * waves``. - This is a small deterministic default heuristic. It deliberately uses the - same factors for every dtype; the production profile validator still - determines which actual-dtype recipes are legal. + This is a small deterministic default heuristic. Both factors come from + the tables of ``cost_kv_width`` (``_grouped_q_cost_kv_width``); the + production profile validator still determines which actual-dtype recipes + are legal. """ if splits_kv <= 0: raise ValueError("splits_kv must be positive") @@ -507,8 +553,9 @@ def make_grouped_q_launch_candidate( launched_ctas = base_ctas * splits_kv waves = (launched_ctas + service_capacity - 1) // service_capacity modeled_time = ( - _GROUPED_Q_MAINLOOP_COST[candidate.tile_size_q] * seq_len_per_cta_kv - + _GROUPED_Q_REDUCTION_COST[candidate.tile_size_q] + _GROUPED_Q_MAINLOOP_COST[cost_kv_width][candidate.tile_size_q] + * seq_len_per_cta_kv + + _GROUPED_Q_REDUCTION_COST[cost_kv_width][candidate.tile_size_q] * _GROUPED_Q_REDUCTION_SEQ_LEN_FACTOR * splits_kv ) * waves @@ -620,6 +667,14 @@ class FmhaDecodeConfig: # Number of KV tiles in the launch (= ceil(seq_len_kv / tile_size_kv)). # Populated by the launcher once seq_len_kv is known. total_kv_tiles: int = 0 + # Compile for the exact Q/KV lengths recorded in this config. By default the + # logical lengths are runtime launch operands and stay out of the compiled + # program and its cache identity (see ``compile_signature``). + fixed_seqlen: bool = False + # Whether a grouped sparse route carries a packed-membership table. Runtime Q + # lengths cannot imply it, and a single-token route keeps the lean layout, so + # the launch policy records it before the lengths leave the identity. + query_memberships: bool | None = None # Q rows placed on the small MMA N dimension. Shape-aware selection derives # this tile from the Q-heads-per-KV-head ratio and grouping policy. tile_size_q: int = 8 @@ -629,6 +684,8 @@ class FmhaDecodeConfig: # Select the packed Q/O ABI. Q and O are laid out as # [sum_q_tokens, num_heads_q, head_dim] and indexed by cu_seqlens_q. use_variable_seqlens_q: bool = False + # Derived from the affine metadata view; lengths themselves stay runtime. + uniform_query_offsets: bool = False # Select the private sparse-block route. This explicit discriminator keeps # ordinary storage-subpage callers on the generic paged-KV path. Grouped # sparse routes consume query membership from a separate packed-word table; @@ -811,8 +868,8 @@ def has_storage_subpages(self) -> bool: @property 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 + """Whether a prepared sparse route supplies page-fragment locators.""" + return self.use_q_token_kv_block_sparse_route @property def uses_page_offsets_task(self) -> bool: @@ -834,12 +891,23 @@ def shares_sparse_pattern(self) -> bool: @property def uses_q_token_kv_block_sparse_page_membership(self) -> bool: """Whether grouped sparse attention uses a packed-membership table.""" - return self.uses_q_token_kv_block_sparse_page_route and self.max_seq_len_q > 1 + return self.uses_q_token_kv_block_sparse_page_route and ( + self.is_multi_token_q + if self.query_memberships is None + else self.query_memberships + ) + + @property + def is_multi_token_q(self) -> bool: + """Runtime Q domains retain general token indexing and masking.""" + return not self.fixed_seqlen or self.max_seq_len_q > 1 @property def grouped_q_rows(self) -> int: """Return logical query/head rows represented by one sparse-route group.""" - return self.heads_q_per_kv * self.max_seq_len_q + return self.heads_q_per_kv * ( + self.q_tokens_per_cta if not self.fixed_seqlen else self.max_seq_len_q + ) @property def grouped_q_fits_tile(self) -> bool: @@ -860,7 +928,12 @@ def uses_q_token_kv_block_sparse_page_route(self) -> bool: and self.num_tokens_per_page in (4, 8, 16, 32, 64, 128) and self.groups_tokens_heads_q and self.use_q_token_kv_block_sparse_route - and 0 < self.max_seq_len_q <= Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIP_BITS + and ( + not self.fixed_seqlen + or 0 + < self.max_seq_len_q + <= Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIP_BITS + ) and self.grouped_q_fits_tile ) @@ -870,7 +943,7 @@ def supports_query_major_memberships(self) -> bool: return ( self.uses_q_token_kv_block_sparse_page_membership and self.use_keeps_mma_ab - and 2 <= self.max_seq_len_q <= 8 + and (self.fixed_seqlen and 2 <= self.max_seq_len_q <= 8) and self.has_single_q_cta and self.num_tokens_per_page == 4 and self.tile_size_kv == 128 @@ -882,14 +955,18 @@ def uses_held_encoded_locator_window(self) -> bool: """Whether encoded paging keeps the complete CTA route in SMEM.""" return ( - self.use_paged_kv + not self.runtime_split_kv + and self.use_paged_kv and self.uses_scattered_page_route and self.tile_size_kv == 128 + and self.fixed_seqlen and not self.use_sliding_window_causal and self.static_local_kv_tiles <= Q_TOKEN_KV_BLOCK_SPARSE_HELD_LOCATOR_MAX_TILES ) + # Select profiles whose task roles hand registers between producer and compute + # warps. @property def uses_task_register_reallocation(self) -> bool: if ( @@ -906,7 +983,11 @@ def uses_task_register_reallocation(self) -> bool: self.use_keeps_mma_ab and self.tile_size_q == 64 and self.tile_size_kv == 256 - and self.total_kv_tiles >= KV_TILE_256_REGISTER_REALLOCATION_MIN_TILES + and ( + not self.fixed_seqlen + or self.total_kv_tiles + >= KV_TILE_256_REGISTER_REALLOCATION_MIN_TILES + ) ) ): return True @@ -971,6 +1052,17 @@ def correction_task_num_registers(self) -> int | None: return None if self.use_transform_kv: return 88 + # Runtime partition metadata extends the full score tile's live ranges. + # Grant one eight-register quantum per Softmax thread from Correction; + # the two-instance Q128 graph returns 16 registers per correction lane. + if self.runtime_split_kv and self.use_keeps_mma_ab and self.tile_size_q == 128: + softmax_warps = self.softmax0_num_warps + ( + self.softmax1_num_warps if self.num_insts_kv != 1 else 0 + ) + quanta = ( + softmax_warps + self.correction_num_warps - 1 + ) // self.correction_num_warps + return 88 - quanta * 8 return KV_TILE_256_CORRECTION_TASK_REGISTERS if self.tile_size_kv == 256 else 88 @property @@ -1014,10 +1106,16 @@ def padding_task_num_registers(self, warp_idx: int, num_warps: int) -> int | Non # the swizzled TMA descriptors used here. stensor_align: int = 1024 + @property + def smem_q_head_dim(self) -> int: + """Q storage rounded to a complete TMA swizzle fragment.""" + fragment = min(128 // self.q_dtype_bytes, self.headdim) + return (self.headdim + fragment - 1) // fragment * fragment + @property def smem_q_tile_bytes(self) -> int: - """SMEM bytes for one Q tile (tile_size_q × headdim elements).""" - return self.tile_size_q * self.headdim * self.q_dtype_bytes + """SMEM bytes for one Q tile, including TMA-filled head padding.""" + return self.tile_size_q * self.smem_q_head_dim * self.q_dtype_bytes @property def q_tokens_per_cta(self) -> int: @@ -1032,7 +1130,11 @@ def q_tokens_per_cta(self) -> int: def q_tma_tokens_per_cta(self) -> int: """Q tokens copied by a box, excluding spare sparse-group slots.""" if self.use_variable_seqlens_q and self.uses_q_token_kv_block_sparse_page_route: - return min(self.max_seq_len_q, self.q_tokens_per_cta) + return ( + self.q_tokens_per_cta + if not self.fixed_seqlen + else min(self.max_seq_len_q, self.q_tokens_per_cta) + ) return self.q_tokens_per_cta @property @@ -1043,6 +1145,10 @@ def num_q_ctas(self) -> int: raw shape-less configs conservative so this scheduling predicate can never specialize geometry that has not been resolved yet. """ + if not self.fixed_seqlen: + if hasattr(self, "runtime_q_ctas"): + return self.runtime_q_ctas + raise ValueError("runtime Q CTA count is a launch operand") if self.heads_q_per_kv <= 0 or self.max_seq_len_q <= 0: return max(self.max_seq_len_q, 1) q_geometry = make_q_tile_geometry( @@ -1052,11 +1158,24 @@ def num_q_ctas(self) -> int: ) return max(q_geometry.num_q_ctas(self.max_seq_len_q), 1) + @property + def uses_compact_split_grid(self) -> bool: + """Launch actual split work in X and fuse Q/batch in Z.""" + return ( + self.runtime_split_kv + and self.use_split_kv + and not self.use_cluster_smem_reduction + and not self.use_persistent_scheduler + ) + @property def has_single_q_cta(self) -> bool: """Whether every physical split maps to one logical Q CTA.""" return ( - self.heads_q_per_kv > 0 and self.max_seq_len_q > 0 and self.num_q_ctas == 1 + self.fixed_seqlen + and self.heads_q_per_kv > 0 + and self.max_seq_len_q > 0 + and self.num_q_ctas == 1 ) @property @@ -1431,6 +1550,9 @@ def use_fp8_kv(self) -> bool: # Split-KV fanout: number of CTAs that cooperate on one (batch, head_kv) # when use_split_kv is enabled. The launcher picks this based on SM count. splits_kv: int = 1 + runtime_split_kv: bool = False + # Internal launch profile for a bounded CTA-relative page-ID span. + runtime_page_window: bool = False # Upper bound of splits_kv across all (batch, head_kv) groups, # used to size the per-launch partial-O / partial-stats / counter # scratch GMEM buffers. @@ -1480,6 +1602,7 @@ def use_fp8_kv(self) -> bool: # denominator (extra logit that absorbs probability mass). Requires the # `attention_sinks` tensor argument at launch. use_attention_sinks: bool = False + lse_mode: str = "none" # Optional profile and reduction knobs selected by launcher policy or tests. use_keeps_mma_ab: bool = False # Sage attention: a nonzero ``sage_k_block_size`` enables per-block Q/K @@ -1507,8 +1630,9 @@ def use_fp8_kv(self) -> bool: 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 + # rate; ``arch_config_args`` supplies it for the launch device, and + # ``fragment_ex2_emulated_pairs`` caps it on KV256 tiles. + ex2_emulated_pairs: int = 6 # Reduce streamed score-fragment maxima inside the TMEM load # (tcgen05.ld.red); ``arch_config_args`` supplies it. uses_ldtm_stat: bool = False @@ -1554,6 +1678,7 @@ def parallel_reduction_cluster_size(self) -> int: if ( not self.use_separate_reduction_kernel or self.use_compact_parallel_reduction + or self.runtime_split_kv ): return 1 return { @@ -1640,20 +1765,43 @@ def num_s_regs_per_thread(self) -> int: return self.tile_size_kv // 2 return self.num_softmax_scale_groups * 4 + @property + def keeps_score_chunk_regs(self) -> int: + """Return the scores of one Keeps K32 chunk, a native 32-register LDTM atom. + + Streamed profiles publish P one chunk at a time + (``softmax_score_fragment_regs``). Keeps layout atoms and Sage ``sfK`` + groups count in chunks whether the softmax streams the row or holds + it complete, so both forms read one scale layout. + """ + return 32 + @property def softmax_score_fragment_regs(self) -> int: """Return the maximum score fragment kept live in registers. Streamed profiles own 128 score values per lane but process them as - four native 32-register LDTM atoms. Other profiles retain their - complete score fragment, so this property is intentionally distinct - from ``num_s_regs_per_thread`` (the total logical ownership). The - selection of streamed profiles lives in ``streams_tmem_p_fragments``. + four K32 chunks (``keeps_score_chunk_regs``). Other profiles retain + their complete score fragment, so this property is intentionally + distinct from ``num_s_regs_per_thread`` (the total logical ownership). + The selection of streamed profiles lives in + ``streams_tmem_p_fragments``. """ if self.streams_tmem_p_fragments: - return 32 + return self.keeps_score_chunk_regs return self.num_s_regs_per_thread + @property + def fragment_ex2_emulated_pairs(self) -> int: + """Return a streamed fragment's score pairs exponentiated on the FMA pipe. + + KV128 tiles take ``ex2_emulated_pairs``; KV256 tiles take at most + ``KV_TILE_256_MAX_EX2_EMULATED_PAIRS``. + """ + if self.tile_size_kv == 256: + return min(self.ex2_emulated_pairs, KV_TILE_256_MAX_EX2_EMULATED_PAIRS) + return self.ex2_emulated_pairs + @property def num_softmax_score_fragments(self) -> int: """Return score fragments used to cover one logical KV tile.""" @@ -1729,7 +1877,7 @@ def keeps_route_atom_owner(self, atom: int) -> tuple[int, int]: @property def keeps_fragments_per_atom(self) -> int: - """Return the K32 score fragments of one Keeps layout atom. + """Return the K32 score chunks of one Keeps layout atom. Fragment ``f`` of a thread is the low or high half of its spatial half's KV64 atom ``f // 2``. Block-sparse route atoms span the same @@ -1740,7 +1888,7 @@ def keeps_fragments_per_atom(self) -> int: @property def keeps_atom_tokens(self) -> int: """Return the K tokens of one Keeps layout atom.""" - return self.keeps_fragments_per_atom * self.softmax_score_fragment_regs + return self.keeps_fragments_per_atom * self.keeps_score_chunk_regs @property def proxy_summary_geometry(self) -> tuple[int, int]: @@ -1775,6 +1923,12 @@ def proxy_tail_summary(self) -> tuple[int, float]: block boundary. """ num_summaries, tail_mass = self.proxy_summary_geometry + if not self.fixed_seqlen: + import cutlass.cute as cute + + return num_summaries - 1, cute.math.log2(Float32(tail_mass)) - math.log2( + self.kv_block_size + ) return num_summaries - 1, math.log2(tail_mass) - math.log2(self.kv_block_size) @property @@ -1787,10 +1941,10 @@ def num_proxy_groups(self) -> int: return (self.num_proxy_summaries + self.tile_size_kv - 1) // self.tile_size_kv @property - def proxy_static_tail_fragment(self) -> tuple[int, int]: + def proxy_tail_fragment(self) -> tuple[int, int]: """Return ``(spatial half, fragment)`` holding the ragged final summary. - Proxy groups are fixed, so the final summary sits at the compile-time + Proxy groups cover consecutive summaries, so the final summary sits at position ``tail_summary_idx % tile_size_kv`` of the last group. Only meaningful when ``proxy_tail_summary`` reports a non-zero shortfall. """ @@ -1840,15 +1994,24 @@ def tmem_p_cols_per_inst(self) -> int: def num_fp8_output_regs(self) -> int: """Return packed FP8 output registers owned by each correction lane.""" return max( - (self.tile_size_q * self.headdim) // FP8_OUTPUT_ELEMENTS_PER_REG_GROUP, + (self.tile_size_q * self.swaps_output_head_dim) + // FP8_OUTPUT_ELEMENTS_PER_REG_GROUP, 1, ) + @property + def swaps_output_head_dim(self) -> int: + """STSM layout covers complete head bands; GMEM stores remain logical.""" + if self.use_keeps_mma_ab: + return self.headdim + return self.num_head_dim_stages_kv * self.head_dim_kv_stage + @property def num_fp16_output_regs(self) -> int: """Return packed 16-bit output registers owned by each correction lane.""" return max( - (self.tile_size_q * self.headdim) // FP16_OUTPUT_ELEMENTS_PER_REG_GROUP, + (self.tile_size_q * self.swaps_output_head_dim) + // FP16_OUTPUT_ELEMENTS_PER_REG_GROUP, 1, ) @@ -1885,7 +2048,11 @@ def inferred_kv_stages(self) -> int: """Return the deepest K/V ring that fits the shared-memory budget.""" q_dtype_bits = self.q_dtype_bytes * BITS_PER_BYTE q_row_bytes = ( - (q_dtype_bits * self.headdim // BITS_PER_BYTE + Q_ROW_ALIGNMENT_BYTES - 1) + ( + q_dtype_bits * self.smem_q_head_dim // BITS_PER_BYTE + + Q_ROW_ALIGNMENT_BYTES + - 1 + ) // Q_ROW_ALIGNMENT_BYTES ) * Q_ROW_ALIGNMENT_BYTES fixed_bytes = ( @@ -1895,7 +2062,11 @@ def inferred_kv_stages(self) -> int: if not self.uses_tmem_p: fixed_bytes += self.num_insts_kv * self.smem_p_tile_bytes if not self.use_keeps_mma_ab: - fixed_bytes += self.tile_size_q * self.headdim * max(self.o_dtype_bytes, 2) + fixed_bytes += ( + self.tile_size_q + * self.swaps_output_head_dim + * max(self.o_dtype_bytes, 2) + ) if self.use_paged_kv: page_ids_per_stage = WARP_THREADS if self.uses_held_encoded_locator_window: @@ -2076,12 +2247,11 @@ def validate_sage_profile(self) -> None: f"got {self.sage_k_summary_block_size}" ) q_block_size = self.sage_q_block_size - if not is_power_of_two(q_block_size) or q_block_size > self.tile_size_q: + if not is_power_of_two(q_block_size): raise ValueError( - "sage_q_block_size must be a power of two no larger than " - f"tile_size_q ({self.tile_size_q}), got {q_block_size}" + f"sage_q_block_size must be a power of two, got {q_block_size}" ) - if not self.streams_tmem_p_fragments: + if not (self.streams_tmem_p_fragments or self.stages_q128_p_in_smem): raise ValueError( "Sage attention requires a two-instance Keeps profile: " "Q64/KV256 or Q128/KV128" @@ -2201,13 +2371,24 @@ def compile_signature(self) -> tuple[tuple[str, object], ...]: The public planner can reuse one compiled topology across batch sizes. Keeping the complete static config in the cache key prevents a callable compiled for one scheduler, reduction, or tile layout from being reused - by another. + by another. Unless ``fixed_seqlen`` asks for exact lengths, every + length-derived field is dropped from the key: Q lengths come from packed + offsets and K/V lengths from metadata at launch. """ - return tuple( - (config_field.name, getattr(self, config_field.name)) + items = { + config_field.name: getattr(self, config_field.name) for config_field in fields(self) - ) + } + if not self.fixed_seqlen: + policy = replace(self, fixed_seqlen=True) + for name in LENGTH_DERIVED_FIELDS: + del items[name] + items.update( + use_variable_seqlens_q=True, + query_memberships=policy.uses_q_token_kv_block_sparse_page_membership, + ) + return tuple(items.items()) @property def uses_prepared_score_keep_words(self) -> bool: @@ -2253,7 +2434,7 @@ def uses_q_desc_ref(self) -> bool: @property def has_odd_kv_tail(self) -> bool: """Whether the K/V tile count leaves an unpaired tail instance.""" - return (self.total_kv_tiles % self.num_insts_kv) != 0 + return not self.fixed_seqlen or (self.total_kv_tiles % self.num_insts_kv) != 0 @property def uses_uniform_causal_mask(self) -> bool: @@ -2298,6 +2479,23 @@ def _uses_contiguous_grouped_keeps_recipe(self) -> bool: ) return self._grouped_keeps_profile_key in profiles + @property + def _in_grouped_q_selection_scope(self) -> bool: + """Whether the grouped-Q cost model may select this launch's Q tile. + + It covers fixed multi-token grouped launches over paged K/V with 4- to + 128-token pages and no attention sinks; the selector also requires the + caller to leave the MMA and launch fields unpinned. + """ + return ( + self.groups_tokens_heads_q + and self.use_paged_kv + and self.num_tokens_per_page in (4, 8, 16, 32, 64, 128) + and not self.use_variable_seqlens_q + and self.max_seq_len_q > 1 + and not self.use_attention_sinks + ) + @property def _kv256_dtypes_qualified(self) -> bool: """Whether the Q64/KV256 profile is qualified for this dtype combination.""" @@ -2307,18 +2505,38 @@ def _kv256_dtypes_qualified(self) -> bool: @property def supports_auto_grouped_keeps(self) -> bool: - """Whether grouped Keeps is also qualified for automatic selection.""" + """Whether automatic selection may choose this grouped Keeps config. + + ``supports_grouped_keeps`` decides whether the kernel runs a config; + automatic selection, with or without a cost comparison, takes a recipe + only on the launches it is qualified for. Paged FP8 qualifies the + established cost-model recipes. Paged 16-bit Q128 qualifies the fixed + multi-token launches that the grouped-Q cost model scores; single-token + and variable-length launches keep their previous tiles. + """ + if not self.supports_grouped_keeps: + return False + profile = self._grouped_keeps_profile_key + if ( + self.use_paged_kv + and self.tile_size_q == 128 + and profile in _GROUPED_KEEPS_PAGED_16BIT_Q128_PROFILES + ): + return self._in_grouped_q_selection_scope if ( self.use_paged_kv and self.tile_size_kv == 128 and self.q_dtype == self.k_dtype == self.v_dtype == Float8E4M3FN ): - return ( - self._grouped_keeps_profile_key - in _AUTO_GROUPED_KEEPS_PAGED_FP8_PROFILES - and self.supports_grouped_keeps + # DSL 4.8 measurements also qualify D128 FP8 -> 16-bit fixed + # multi-Q recipes. Let the existing cost model compare them; + # retaining only Swaps Q32 leaves a substantial runtime-split gap. + return profile in _AUTO_GROUPED_KEEPS_PAGED_FP8_PROFILES or ( + self._in_grouped_q_selection_scope + and self.headdim == 128 + and self.out_dtype in (Float16, BFloat16) ) - return self.supports_grouped_keeps + return True @property def uses_guarded_fixed_q1_grouped_keeps(self) -> bool: @@ -2335,6 +2553,7 @@ def uses_guarded_fixed_q1_grouped_keeps(self) -> bool: self.use_keeps_mma_ab and self.use_paged_kv and self.groups_tokens_heads_q + and self.fixed_seqlen and self.max_seq_len_q == 1 and not self.use_variable_seqlens_q and self.q_manual_padding_rows == 0 @@ -2362,7 +2581,10 @@ def uses_guarded_grouped_keeps_output_rows(self) -> bool: self.use_keeps_mma_ab and self.use_paged_kv and self.groups_tokens_heads_q - and not self.use_variable_seqlens_q + and ( + not self.use_variable_seqlens_q + or (self.uniform_query_offsets and not self.use_sage_attention) + ) and profile in _GROUPED_KEEPS_PAGED_FP8_PROFILES and self.supports_grouped_keeps and not (self.uses_staged_one_inst_tmem_p and self.tile_size_q == 64) @@ -2372,7 +2594,8 @@ def uses_guarded_grouped_keeps_output_rows(self) -> bool: def has_static_dense_full_kv_tiles(self) -> bool: """Whether static dense KV avoids masking and runtime tile remapping.""" return ( - not self.use_split_kv + self.fixed_seqlen + and not self.use_split_kv and not self.use_block_sparse and self.mask_type == DENSE and not self.use_sliding_window_causal @@ -2383,10 +2606,11 @@ def has_static_dense_full_kv_tiles(self) -> bool: @property def uses_ordered_softmax_barrier(self) -> bool: """Whether this profile selects the ordered P0/P1 softmax barrier.""" - if self.streams_tmem_p_fragments: - # Streamed profiles use independent per-fragment P pipelines. - # Ordering the two softmax groups would serialize fragment - # production and defeat the intended P/PV overlap. + if self.streams_tmem_p_fragments or self.stages_q128_p_in_smem: + # Streamed profiles use independent per-fragment P pipelines, and + # SMEM-P Q128 tiles overlap each softmax with its own next QK. + # Ordering the two softmax groups would make one wait for the other + # and defeat that overlap. return False if self.ordered_softmax_barrier_mode == 2: return True @@ -2431,15 +2655,24 @@ def uses_nontrivial_grouped_q_layout(self) -> bool: @property def q_tiles_are_full(self) -> bool: - """Whether every launched Q CTA owns all of its MMA rows.""" - if self.use_variable_seqlens_q: + """Whether every active Q CTA owns all of its MMA rows.""" + # Empty/overlaunched Q tiles are discarded before the task graph. + # One-token tiles with complete head bands cannot have partial rows, + # even when request lengths are packed or change at runtime. + if ( + self.heads_q_per_kv > 0 + and self.q_tokens_per_cta == 1 + and self.heads_q_per_kv % self.tile_size_q == 0 + ): + return True + if self.use_variable_seqlens_q or not self.fixed_seqlen: return False if self.uses_nontrivial_grouped_q_layout: return ( self.q_manual_padding_rows == 0 and self.max_seq_len_q % self.q_tokens_per_cta == 0 ) - if self.max_seq_len_q > 1: + if self.is_multi_token_q: return self.heads_q_per_kv % self.tile_size_q == 0 return self.heads_q_per_kv == self.tile_size_q @@ -2447,8 +2680,8 @@ def q_tiles_are_full(self) -> bool: def q_tiles_need_row_mask(self) -> bool: """Whether softmax must suppress inactive Q rows.""" if self.use_variable_seqlens_q: - return True - if self.uses_nontrivial_grouped_q_layout or self.max_seq_len_q > 1: + return not self.q_tiles_are_full + if self.uses_nontrivial_grouped_q_layout or self.is_multi_token_q: return not self.q_tiles_are_full return False @@ -2465,7 +2698,7 @@ def uses_q_cta_sliding_union(self) -> bool: """Whether the causal/window KV union depends on the logical Q CTA.""" return ( self.mask_type == CAUSAL - and self.max_seq_len_q > 1 + and self.is_multi_token_q and not self.has_single_q_cta ) @@ -2473,13 +2706,14 @@ def uses_q_cta_sliding_union(self) -> bool: def uses_runtime_q_kv_union(self) -> bool: """Whether task/resource KV geometry must retain runtime Q metadata. - Multiple causal Q CTAs have distinct right bounds. A multi-token - sliding-window launch also retains the runtime path even when it fits - in one Q CTA, because the existing static window-prefix metadata is - specialized only for SQ1. + This is a geometry property, not the ``fixed_seqlen`` switch: it holds + even when every length is static. Multiple causal Q CTAs have distinct + right bounds. A multi-token sliding-window launch also retains the + runtime path even when it fits in one Q CTA, because the existing static + window-prefix metadata is specialized only for SQ1. """ return self.uses_q_cta_sliding_union or ( - self.use_sliding_window_causal and self.max_seq_len_q > 1 + self.use_sliding_window_causal and self.is_multi_token_q ) @property @@ -2489,12 +2723,13 @@ def uses_tmem_p(self) -> bool: @property def uses_staged_one_inst_tmem_p(self) -> bool: - """Whether P uses the double-buffered D256 TMEM overlay.""" + """Whether staged wide-head P uses the double-buffered TMEM overlay.""" return ( self.use_keeps_mma_ab - and self.headdim == 256 + and self.headdim in (192, 256) and ( self.head_dim_per_stage_kv == 128 + or (self.headdim == 192 and self.head_dim_per_stage_kv == 192) or ( self.head_dim_per_stage_kv == 256 and self.use_fp8_qkv @@ -2509,11 +2744,11 @@ def uses_staged_one_inst_tmem_p(self) -> bool: def uses_two_inst_tmem_p(self) -> bool: """Whether a two-instance Keeps profile uses the TMEM-P overlay. - Dense 16-bit Q128/KV128 and the other unstreamed profiles publish a - complete packed-P row per pipeline token. Streamed profiles use the - same S-to-P aliasing contract but stream four independently ready K32 - fragments (see ``streams_tmem_p_fragments``). Q64/KV128 keeps the base - kernel's faster SMEM-P cadence. + Unstreamed profiles publish a complete packed-P row per pipeline + token. Streamed profiles use the same S-to-P aliasing contract but + stream four independently ready K32 fragments (see + ``streams_tmem_p_fragments``). Q64/KV128 and the 8-bit Q128/KV128 + tiles of ``stages_q128_p_in_smem`` stage P in SMEM instead. """ # Two-instance Keeps keeps stats outside S, so both static and persistent # work tiles can overlay P on the consumed S instance. The split K/V @@ -2529,6 +2764,33 @@ def uses_two_inst_tmem_p(self) -> bool: and self.num_insts_kv == 2 and self.o_stages == 2 and self.tmem_total_cols <= 512 + and not self.stages_q128_p_in_smem + ) + + @property + def stages_q128_p_in_smem(self) -> bool: + """Whether a two-instance Keeps Q128/KV128 tile stages its P in SMEM. + + P aliasing S in TMEM ties each instance's next QK to its PV: the QK + may only overwrite S once PV has read the whole P. With P in SMEM, + Softmax holds the complete score row in registers and releases S after + the max pass, so the MMA warp issues the next QK before this tile's PV + and the exponentials overlap that QK. Dense 8-bit Q/K/V tiles take it, + Sage included; an 8-bit P tile fits beside the K/V ring. Sage's max + pass dequantizes the row in registers one K32 chunk of ``sfK`` groups + at a time, so its P pass reads no scale words. Block-sparse routes + keep the streamed TMEM P. + """ + return ( + self.use_keeps_mma_ab + and self.tile_size_q == 128 + and self.tile_size_kv == 128 + and self.head_dim_per_stage_kv == 0 + and self.num_insts_kv == 2 + and self.o_stages == 2 + and self.k_dtype.width == 8 + and self.v_dtype.width == 8 + and not self.use_block_sparse ) @property @@ -2544,12 +2806,15 @@ def streams_tmem_p_fragments(self) -> bool: A two-instance TMEM-P profile streams for KV256 tiles and block-sparse routes, whose loops wait on K/V loads that the earlier PV start hides. - It also streams with 8-bit K and V and SMEM softmax stats (the 8-bit - Q/K/V and Sage recipes) unless Q-token page membership applies: a - dense fragment's keep word is one visible token range and cannot - express per-page membership. Dense 16-bit Q128/KV128 (not load-bound), - FP8 D64 (softmax stats in TMEM) and QK-BF16/PV-FP8 keep the complete - row. + Dense KV128 tiles stream too when K and V share a width and the + softmax stats travel through SMEM (the 16-bit recipes; dense 8-bit + tiles, Sage included, stage P in SMEM through + ``stages_q128_p_in_smem``): the PV MMA starts on the first fragment + instead of the complete row, and the exponentials share the FMA pipe + (``fragment_ex2_emulated_pairs``). Q-token page membership keeps the + complete row, since a dense fragment's keep word is one visible token + range and cannot express per-page membership; so do FP8 D64 (softmax + stats in TMEM) and QK-BF16/PV-FP8. """ if not self.uses_two_inst_tmem_p: return False @@ -2557,8 +2822,7 @@ def streams_tmem_p_fragments(self) -> bool: self.tile_size_kv == 256 or self.use_block_sparse or ( - self.k_dtype.width == 8 - and self.v_dtype.width == 8 + self.k_dtype.width == self.v_dtype.width and self.keeps_stats_via_smem and not self.uses_q_token_kv_block_sparse_page_membership ) @@ -2571,16 +2835,11 @@ 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_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. + fewer TMEM rescales. A row maximum moves often but rarely by much, so + every Keeps profile takes the trade. A zero headroom anchors on the + exact row maximum. """ - return ( - self.use_keeps_mma_ab - and (self.tile_size_kv == 256 or self.use_block_sparse) - and self.softmax_anchor_headroom_log2 > 0 - ) + return self.use_keeps_mma_ab and self.softmax_anchor_headroom_log2 > 0 @property def softmax_anchor_headroom_log2(self) -> float: @@ -2656,57 +2915,74 @@ def sage_scale_layout(self, seq_len_kv: int) -> _SageScaleLayout: """Return the layout of the Sage scales a plan prepares before attention. 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``). + chunks, 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``. Streamed + and complete-row softmax passes read the words per chunk alike + (``keeps_score_chunk_regs``). """ 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 + summaries = None if self.use_block_sparse_proxy_routes: num_summaries, _ = _block_sparse_proxy_summary_geometry( seq_len_kv, self.kv_block_size ) + summaries = _SageKSection( + length=num_summaries, + block_size=self.sage_k_scale_block_size(proxy=True), + source_block_size=self.sage_k_summary_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, + k_tokens=_SageKSection( + length=seq_len_kv, + block_size=self.sage_k_scale_block_size(), + source_block_size=self.sage_k_block_size, + ), + q_rows=_SageQSection( + 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, + rows_per_tile=self.tile_size_q, + heads_interleaved=True, + ), + v_rows=_SageVSection(headdim=self.headdim, has_mean=self.sage_v_mean), + summaries=summaries, ) q_rows = layout.q_rows assert q_rows.tokens_per_tile == self.q_tokens_per_cta - assert q_rows.tiles == self.num_q_ctas + if self.fixed_seqlen: + assert q_rows.tiles == self.num_q_ctas return layout 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.""" + """Return the ``sfK`` groups of one K32 score chunk 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) + 1, self.keeps_score_chunk_regs // self.sage_k_scale_block_size(proxy) ) - def sage_scores_dequantized(self, proxy: bool = False) -> bool: - """Whether the max pass writes a route kind's scores back dequantized. + def sage_score_scaling(self, proxy: bool = False) -> SageScoreScaling | None: + """Return how the softmax applies a route kind's ``sfK`` to its scores. - 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. + A complete-row SMEM-P tile (``stages_q128_p_in_smem``) dequantizes its + register-resident row in every geometry. A streamed kind is dequantized + only with one scale per score (the one-token K block): the max pass + writes its scores back multiplied. Either way the P pass reads neither + ``sfK`` nor the INT32 bias. Other streamed geometries scale blocks of + ``sage_k_scale_block_size`` columns, avoiding a TMEM round trip per + tile. Every kind's words come through the ``SageKScalesResource`` ring + (``SageKTransport.SMEM``). """ - return ( - self.use_sage_attention - and self.sage_k_groups_per_fragment(proxy) - == self.softmax_score_fragment_regs - ) + if not self.use_sage_attention: + return None + if ( + self.stages_q128_p_in_smem + or self.sage_k_groups_per_fragment(proxy) == self.keeps_score_chunk_regs + ): + return SageScoreScaling.DEQUANTIZED + return SageScoreScaling.BLOCK @property def matches_kv256_task_topology(self) -> bool: @@ -2805,6 +3081,7 @@ def can_use_cluster_smem_reduction(self) -> bool: ungrouped_q_layout = ( not self.groups_tokens_heads_q and not self.use_variable_seqlens_q + and self.fixed_seqlen and self.max_seq_len_q == 1 ) grouped_q_layout = ( @@ -3005,32 +3282,44 @@ def supports_grouped_keeps(self) -> bool: if self.use_sage_attention: return False - if self.use_paged_kv and profile in _GROUPED_KEEPS_PAGED_FP8_PROFILES: - fixed_q1 = ( - self.max_seq_len_q == 1 and 1 <= self.heads_q_per_kv <= self.tile_size_q - ) - fixed_grouped_q = self.max_seq_len_q > 1 + if profile in _D192_GROUPED_KEEPS_PROFILES: + # Fixed-Q staged resources support dense/paged K/V, CLC and + # non-windowed GMEM splits. Common validators enforce the + # remaining reduction geometry and exclude cluster-SMEM Keeps. return ( - (fixed_q1 or fixed_grouped_q) - and not self.use_variable_seqlens_q - and self.use_paged_kv - and self.num_tokens_per_page == 32 - and self.mask_type in (DENSE, CAUSAL) - and not any( - ( - self.use_cluster_smem_reduction, - self.use_sliding_window_causal, - self.use_attention_sinks, - ) - ) - and ( + ( direct or ( - self.use_split_kv - and self.splits_kv > 1 - and self.max_splits_kv >= self.splits_kv + self.has_valid_split_kv_fanout + and not self.use_sliding_window_causal ) ) + and not self.uses_q_token_kv_block_sparse_page_route + and (not self.use_variable_seqlens_q or not self.fixed_seqlen) + and not self.use_attention_sinks + and not self.use_cluster_smem_reduction + ) + + paged_16bit_q128 = ( + self.tile_size_q == 128 + and profile in _GROUPED_KEEPS_PAGED_16BIT_Q128_PROFILES + ) + if self.use_paged_kv and ( + paged_16bit_q128 or profile in _GROUPED_KEEPS_PAGED_FP8_PROFILES + ): + # Paged grouped Keeps recipes exclude sliding windows and attention + # sinks. The 16-bit Q128 recipe, which reuses the contiguous 16-bit + # resources, runs fixed or packed Q in direct nonpersistent + # launches; the FP8 recipes run fixed Q, direct or split. + if self.use_sliding_window_causal or self.use_attention_sinks: + return False + if paged_16bit_q128: + return direct and not self.use_persistent_scheduler + return ( + (self.is_multi_token_q or self.grouped_q_fits_tile) + and (not self.use_variable_seqlens_q or not self.fixed_seqlen) + and self.mask_type in (DENSE, CAUSAL) + and (direct or self.has_valid_split_kv_fanout) ) if profile in _GROUPED_KEEPS_STATIC_ONLY_PROFILES: @@ -3643,6 +3932,10 @@ def _make_static_decode_config( """ cfg = FmhaDecodeConfig(headdim=headdim) explicit_fields = _apply_config_source(cfg, args) + # Policy selection uses the planned lengths; only the compiled program + # replaces them with runtime operands. + fixed_seqlen = cfg.fixed_seqlen + cfg.fixed_seqlen = True _apply_mask_type_config( cfg, source=args, @@ -3654,6 +3947,7 @@ def _make_static_decode_config( _validate_mixed_kv_dtype_profile(cfg) _validate_kv256_static_config(cfg) _finalize_warp_roles(cfg) + cfg.fixed_seqlen = fixed_seqlen return cfg @@ -3807,12 +4101,15 @@ def enumerate_auto_splits_kv( num_q_tiles: int, service_capacity: int, split_kv: bool = True, + requested_splits_kv: int = -1, + max_splits_kv: int | None = None, ) -> tuple[int, ...]: """Enumerate direct and useful split fanouts for an under-filled Q grid. All one-wave fanouts and the first capacity-crossing fanout participate in the empirical score. The latter is important when a partially filled wave cannot be completed by any uniform integer fanout. + A positive requested fanout is pinned, subject to the same KV-work cap. """ if not split_kv: return (1,) @@ -3824,7 +4121,10 @@ def enumerate_auto_splits_kv( seq_len_kv=seq_len_kv, tile_size_kv=tile_size_kv, num_insts_kv=num_insts_kv, + max_splits_kv=max_splits_kv, ) + if requested_splits_kv > 0: + return (max(1, min(requested_splits_kv, max_by_work)),) base_grid = max(1, batch_size * num_heads_kv * num_q_tiles) if base_grid >= service_capacity or max_by_work <= 1: return (1,) @@ -4106,8 +4406,6 @@ def _apply_swaps_tile_config( _LAUNCH_SELECTION_FIELDS = { "use_split_kv", - "splits_kv", - "max_splits_kv", "use_separate_reduction_kernel", "use_cluster_smem_reduction", "use_persistent_scheduler", @@ -4242,6 +4540,18 @@ def _apply_grouped_q_mma_candidate( cfg.mma_tile_n_bmm2 = candidate.tile_size_q +def _grouped_q_cost_kv_width(cfg: FmhaDecodeConfig) -> int: + """Return the cost-table width that scores ``cfg``'s grouped-Q candidates. + + Only 16-bit K and V at head dim 128 take width 16, since only that head + dim promotes its TileQ64 candidate to KV256. Every other profile takes + width 8, under which D192 picks its faster Q128 tile. + """ + if cfg.headdim == 128 and cfg.k_dtype.width == 16 and cfg.v_dtype.width == 16: + return 16 + return 8 + + def _resolve_grouped_q_launch_candidates( cfg: FmhaDecodeConfig, candidate: GroupedQMmaCandidate, @@ -4253,6 +4563,8 @@ def _resolve_grouped_q_launch_candidates( num_heads_q: int, num_heads_kv: int, service_capacity: int, + splits_kv: int, + max_splits_kv: int | None, split_kv: bool = True, ) -> tuple[GroupedQLaunchCandidate, ...]: """Resolve one TileQ with Q-grid-aware direct and split recipes. @@ -4280,6 +4592,19 @@ def _resolve_grouped_q_launch_candidates( # its Q geometry and modeled cost are identical to the qualified FP16 # logical candidate. The KV selector materializes and validates the # actual BF16 profile only after the Q winner is known. + if cfg.tile_size_kv != 256 and not _try_apply_auto_kv256_profile( + probe, + q_candidate=candidate, + explicit_fields=explicit_fields, + auto_tuner=True, + split_kv_mode="disabled", + num_heads_q=num_heads_q, + num_heads_kv=num_heads_kv, + splits_kv=splits_kv, + max_splits_kv=max_splits_kv, + ): + return () + probe.tile_size_kv = 128 probe.q_dtype = Float16 probe.k_dtype = Float16 probe.v_dtype = Float16 @@ -4289,6 +4614,13 @@ def _resolve_grouped_q_launch_candidates( probe, explicit_fields | {"tile_size_q"}, ) + # Explicit MMA fields constrain candidates instead of disabling the + # search over the remaining geometry and launch choices. + if any( + getattr(probe, field) != getattr(cfg, field) + for field in _MMA_SELECTION_FIELDS & explicit_fields + ): + return () if probe.use_keeps_mma_ab and not probe.supports_auto_grouped_keeps: return () cost_num_insts_kv = probe.num_insts_kv @@ -4301,16 +4633,23 @@ def _resolve_grouped_q_launch_candidates( ) except ValueError: return () + cost_kv_width = _grouped_q_cost_kv_width(probe) - split_candidates = enumerate_auto_splits_kv( - seq_len_kv=seq_len_kv, - batch_size=batch_size, - num_heads_kv=num_heads_kv, - tile_size_kv=cost_tile_size_kv, - num_insts_kv=cost_num_insts_kv, - num_q_tiles=candidate.q_tiles, - service_capacity=service_capacity, - split_kv=split_kv, + split_candidates = ( + (cfg.splits_kv if cfg.use_split_kv else 1,) + if _LAUNCH_SELECTION_FIELDS & explicit_fields + else enumerate_auto_splits_kv( + seq_len_kv=seq_len_kv, + batch_size=batch_size, + num_heads_kv=num_heads_kv, + tile_size_kv=cost_tile_size_kv, + num_insts_kv=cost_num_insts_kv, + num_q_tiles=candidate.q_tiles, + service_capacity=service_capacity, + split_kv=split_kv, + requested_splits_kv=splits_kv, + max_splits_kv=max_splits_kv, + ) ) recipes = [] for candidate_splits_kv in split_candidates: @@ -4339,6 +4678,7 @@ def _resolve_grouped_q_launch_candidates( batch_size=batch_size, num_heads_kv=num_heads_kv, service_capacity=service_capacity, + cost_kv_width=cost_kv_width, ) ) return tuple(recipes) @@ -4359,28 +4699,26 @@ def _apply_auto_grouped_q_mma_config( max_splits_kv: int | None, split_kv: bool = True, ) -> GroupedQLaunchCandidate | None: - """Select a legal fixed multi-Q MMA and KV-split recipe when unpinned. + """Select a legal multi-Q MMA and KV-split recipe within caller constraints. - SQ1, packed/variable Q, explicit launch modes, ungrouped layouts, and - caller-provided MMA fields retain the existing path. Explicit fanout - controls bypass this selector together with explicit launch and MMA fields. + SQ1, packed sliding windows, explicit launch modes, and ungrouped layouts + retain the existing path. Explicit MMA fields filter candidates; explicit launch + fields and fanout controls constrain their launch recipes. Mixed-precision candidates use SwapsMmaAb since Keeps is not supported yet. + + Fanout arguments here are caller overrides, before automatic selection. + Each Q candidate gets its own automatic split search; final KV/reduction + choices are applied after the joint recipe has been selected. """ + # TODO: Refine short-KV setup/epilogue costs. Wider dense Q tiles can + # regress even when their current mainloop/reduction proxy scores lower. if ( not auto_tuner or split_kv_mode != "disabled" or seq_len_q <= 1 - or cfg.use_variable_seqlens_q or not cfg.groups_tokens_heads_q or not cfg.use_paged_kv - or cfg.num_tokens_per_page != 32 - or cfg.mask_type != CAUSAL - or cfg.use_sliding_window_causal - or cfg.use_attention_sinks - or bool(_MMA_SELECTION_FIELDS & explicit_fields) - or bool(_LAUNCH_SELECTION_FIELDS & explicit_fields) - or splits_kv != -1 - or max_splits_kv is not None + or (cfg.use_variable_seqlens_q and cfg.use_sliding_window_causal) ): return None @@ -4392,6 +4730,15 @@ def _apply_auto_grouped_q_mma_config( if not candidates: return None service_capacity = get_max_active_clusters_for_cluster_size(1) + if cfg.use_sliding_window_causal: + # Wider Q tiles reuse overlapping windows, but only win when fewer + # Q-grid waves repay their per-tile cost. Score the visible union, + # not the model's maximum context. Windowed split-KV recipes remain + # outside this selector; qualified direct Keeps tiles can reuse a + # query group with fewer KV loads than narrower Swaps tiles. Partial + # final groups retain their ceil-divided CTA count in the cost model. + seq_len_kv = min(seq_len_kv, cfg.attention_window_size + seq_len_q - 1) + split_kv = False supported = tuple( recipe for candidate in candidates @@ -4405,6 +4752,8 @@ def _apply_auto_grouped_q_mma_config( num_heads_q=num_heads_q, num_heads_kv=num_heads_kv, service_capacity=service_capacity, + splits_kv=splits_kv, + max_splits_kv=max_splits_kv, split_kv=split_kv, ) ) @@ -4423,11 +4772,17 @@ def _apply_auto_grouped_q_mma_config( # mainloop-aware score instead of rewarding narrow tiles that reread # the complete KV sequence for each additional Q tile. direct = tuple(recipe for recipe in supported if recipe.splits_kv == 1) - if not direct: - return None - selected = select_grouped_q_direct_wave_candidate(direct) + selected = ( + select_grouped_q_direct_wave_candidate(direct) + if direct + else select_grouped_q_launch_candidate(supported) + ) _apply_grouped_q_mma_candidate(cfg, selected.mma) - if selected.splits_kv == 1 and selected.base_ctas > service_capacity: + if ( + selected.splits_kv == 1 + and selected.base_ctas > service_capacity + and not (_LAUNCH_SELECTION_FIELDS & explicit_fields) + ): # CLC persistence pays for work discovery by reusing one resident CTA # wave. Select it only when the direct launch has more than one wave; # a single-wave grid has no launch work for persistence to eliminate. @@ -4630,8 +4985,8 @@ def _apply_auto_launch_mode( if mode not in ("gmem_reduction", "persistent"): return split_kv_mode - # An explicitly selected MMA profile can bypass the joint selector while - # still leaving launch mode automatic. Probe the complete derived launch + # Profiles outside the joint selector can still leave launch mode + # automatic. Probe the complete derived launch # before committing it, so an unsupported Keeps split/persistent recipe # falls back to the caller's valid direct profile. probe = deepcopy(cfg) @@ -4872,7 +5227,7 @@ def _validate_profile_support( groups_tokens_heads_q=True, ) # Fixed FP8 Q8 producers and reducers support both 16-bit output types. - # D64/D256 retain their grouped page-32 geometry; head ratios and + # D64/D256 retain their grouped paged-KV geometry; head ratios and # variable-Q layouts remain constrained by the same resource profile. qualified_fp8_q8_separate_reduction_supported = ( not use_keeps_mma_ab @@ -4887,7 +5242,6 @@ def _validate_profile_support( or ( use_groups_tokens_heads_q and cfg.use_paged_kv - and cfg.num_tokens_per_page == 32 and ( headdim == 64 or (headdim == 256 and cfg.out_dtype in (Float16, BFloat16)) @@ -4998,11 +5352,15 @@ def _validate_profile_support( "packed ungrouped SwapsMmaAb does not support cluster SMEM reduction; " "use grouped Q or a GMEM reduction mode" ) - if cfg.use_transform_kv and cfg.use_paged_kv and cfg.num_tokens_per_page == 16: - raise ValueError( - "mixed Q/KV dtypes do not support paged-KV page size 16; " - "use 32, 64, or 128 tokens per page" - ) + # Page16 currently misaligns D64 NVFP4 staging; other transformed profiles + # use the ordinary page-size and staging checks. + if ( + cfg.use_nvfp4_kv + and cfg.headdim == 64 + and cfg.use_paged_kv + and cfg.num_tokens_per_page == 16 + ): + raise ValueError("NVFP4 KV with head_dim=64 does not support page_size=16") if cfg.use_transform_kv and not cfg.store_transformed_kv_in_tmem: # The SMEM transform assigns eight elements per thread per iteration. total_elems = cfg.tile_size_kv * cfg.head_dim_kv_stage @@ -5033,39 +5391,48 @@ def _validate_profile_support( if use_keeps_mma_ab: if cfg.use_transform_kv: raise ValueError("transform-KV is not supported with KeepsMmaAb for now") - if headdim not in (64, 128, 256): - raise ValueError("fmha_decode keepsMmaAb supports headdim=64, 128, or 256") + if headdim not in (64, 128, 192, 256): + raise ValueError( + "fmha_decode keepsMmaAb supports headdim=64, 128, 192, or 256" + ) if tile_size_q not in (64, 128): raise ValueError("fmha_decode keepsMmaAb supports tile_size_q=64 or 128") effective_head_dim_stage = cfg.head_dim_per_stage_kv effective_num_insts_kv = cfg.num_insts_kv effective_o_stages = cfg.o_stages - supported_head_stage = effective_head_dim_stage == 128 or ( - effective_head_dim_stage == 256 - and cfg.use_fp8_qkv - and cfg.uses_q_token_kv_block_sparse_page_route + supported_head_stage = ( + effective_head_dim_stage == 128 + or (headdim == 192 and effective_head_dim_stage == 192) + or ( + effective_head_dim_stage == 256 + and cfg.use_fp8_qkv + and cfg.uses_q_token_kv_block_sparse_page_route + ) ) if effective_num_insts_kv == 1 and ( - headdim != 256 or not supported_head_stage or effective_o_stages != 1 + headdim not in (192, 256) + or not supported_head_stage + or effective_o_stages != 1 ): raise ValueError( "one-instance KeepsMmaAb is enabled only for the staged " - "headDim=256 profile with D128 bands (or D256 for sparse FP8) and " + "headDim=192/256 profiles with D128 bands, full 16-bit D192, " + "or full sparse FP8 D256, and " "o_stages=1" ) - if headdim == 256 and ( + if headdim in (192, 256) and ( not supported_head_stage or effective_num_insts_kv != 1 or effective_o_stages != 1 ): raise ValueError( - "fmha_decode keepsMmaAb headDim=256 requires " - "D128 bands (or D256 for sparse FP8), num_insts_kv=1, and o_stages=1" + "fmha_decode keepsMmaAb headDim=192/256 requires supported " + "head staging, num_insts_kv=1, and o_stages=1" ) - if headdim != 256 and effective_head_dim_stage != 0: + if headdim not in (192, 256) and effective_head_dim_stage != 0: raise ValueError( "split head_dim_per_stage_kv keepsMmaAb profiles are enabled " - "only for headDim=256" + "only for headDim=192 or 256" ) if not use_groups_tokens_heads_q and heads_q_per_kv != tile_size_q: raise ValueError( @@ -5095,14 +5462,14 @@ def _validate_profile_support( and effective_num_insts_kv == 2 and effective_o_stages == 2 ) - separate_reduction_h256_supported = ( - headdim == 256 + separate_reduction_staged_supported = ( + headdim in (192, 256) and ( (tile_size_q == 128 and cfg.tile_size_kv == 128) or ( tile_size_q == 64 and cfg.tile_size_kv == 128 - and cfg.uses_q_token_kv_block_sparse_page_route + and (headdim == 192 or cfg.uses_q_token_kv_block_sparse_page_route) ) ) and supported_head_stage @@ -5116,7 +5483,6 @@ def _validate_profile_support( and cfg.max_seq_len_q == 1 and not cfg.use_variable_seqlens_q and cfg.use_paged_kv - and cfg.num_tokens_per_page == 32 and cfg.q_dtype == Float8E4M3FN and cfg.k_dtype == Float8E4M3FN and cfg.v_dtype == Float8E4M3FN @@ -5126,7 +5492,7 @@ def _validate_profile_support( ) separate_reduction_profile_supported = ( separate_reduction_unstaged_supported - or separate_reduction_h256_supported + or separate_reduction_staged_supported or fixed_fp8_new_keeps_separate_reduction_supported ) if use_separate_reduction_kernel and not ( @@ -5137,9 +5503,10 @@ def _validate_profile_support( ): raise ValueError( "separate reduction keepsMmaAb profiles require " - "the established D128/Q64-Q128, D256/Q128, or " - "QToken-KvBlock-Sparse-Attention D64/D128/D256 Q64/Q128-KV128 profiles, or a " - "fixed FP8/page-32 D64/Q64-Q128 or D256/Q64 profile; valid " + "the established D128/Q64-Q128 or D192-D256/Q128 profiles, " + "D192/Q64-KV128, QToken-KvBlock-Sparse-Attention " + "D64/D128/D192/D256 Q64/Q128-KV128 profiles, or a " + "fixed FP8/paged D64/Q64-Q128 or D256/Q64 profile; valid " "reduction dtypes, Q layout, and static split-KV are required" ) return @@ -5148,15 +5515,15 @@ def _validate_profile_support( # Keep the ungrouped one-token-per-CTA control available at the same tile Q # as grouped profiles so explicit ungrouped launches retain the MMA shape. effective_head_dim_stage = cfg.head_dim_per_stage_kv - if headdim == 256: + if headdim in (192, 256): if effective_head_dim_stage != 128: raise ValueError( - "fmha_decode SwapsMmaAb headDim=256 requires head_dim_per_stage_kv=128" + "fmha_decode SwapsMmaAb headDim=192/256 requires head_dim_per_stage_kv=128" ) elif effective_head_dim_stage != 0: raise ValueError( "split head_dim_per_stage_kv SwapsMmaAb profiles are enabled only " - "for headDim=256" + "for headDim=192 or 256" ) if cfg.use_cluster_smem_reduction: # Single source of truth for structural eligibility plus dtype and @@ -5197,15 +5564,15 @@ def _validate_profile_support( "separate reduction SwapsMmaAb profiles require fixed SQ=1, " "fixed grouped Q, or packed variable Q; tile_size_q in {16,32}, " "QToken-KvBlock-Sparse-Attention TileQ8, or a fixed FP8 Q8/HqPerKv8 profile " - "(D128, plus grouped paged-KV/page-32 D64 or D256 with " + "(D128, plus grouped paged-KV D64 or D256 with " "FP16/BF16 output); valid reduction dtypes and static split-KV are required" ) if tile_size_q and tile_size_q not in (8, 16, 32): raise ValueError("fmha_decode SwapsMmaAb supports tile_size_q in {8,16,32}") - if headdim not in (64, 128, 256): + if headdim not in (64, 128, 192, 256): raise ValueError( "fmha_decode SwapsMmaAb supports headDim in " - "{64,128,256}. headDim=256 uses the staged profile with " + "{64,128,192,256}. headDim=192/256 uses the staged profile with " "head_dim_per_stage_kv=128." ) # Validate constraints imposed by mixed precision KV. @@ -5254,10 +5621,11 @@ 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 + SM100 runs three eighths of the streamed exponentials of KV128 tiles (a + quarter on KV256 tiles) as FMA polynomials because the MUFU issue rate + bounds its P pass while the FMA pipe has room; 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 @@ -5265,7 +5633,7 @@ def arch_config_args(compute_capability: tuple[int, int]) -> dict[str, object]: """ is_sm103 = tuple(compute_capability) == (10, 3) return { - "ex2_emulated_pairs": 0 if is_sm103 else 4, + "ex2_emulated_pairs": 0 if is_sm103 else 6, "uses_ldtm_stat": ldtm_stat_supported(compute_capability), } @@ -5317,8 +5685,8 @@ def make_decode_config( 2. Fill profile-derived defaults: dtype fields, paged-KV metadata, the dense/causal mask, sliding-window flags, attention-sink flags, default grouped-Q metadata for fixed or packed launches, and the SMEM-derived - KV stage count. For an unpinned fixed multi-Q paged-causal page-32 - launch, the Q selector first enumerates Swaps8/16/32 and Keeps64/128 + KV stage count. For a multi-Q paged launch, the Q selector first + enumerates Swaps8/16/32 and Keeps64/128 tiles and every useful direct/split recipe through the first capacity-crossing fanout. Production-valid recipes minimize the empirical TileQ @@ -5347,8 +5715,8 @@ def make_decode_config( downward for a one-wave cluster split. 6. Validate the final profile combination before returning the config. - Attention-sink paths skip automatic launch-mode selection. Explicit launch - modes are still validated. + Outside joint grouped-Q selection, attention-sink paths retain their + existing launch mode. Explicit launch modes are still validated. ``split_kv`` is a common planning control for every FMHA kernel family. True permits supported automatic split recipes; False restricts selection @@ -5401,6 +5769,10 @@ def make_decode_config( validate_sliding_window_args(sliding_window_causal, attention_window_size) cfg = FmhaDecodeConfig(headdim=headdim) explicit_fields = _apply_config_source(cfg, args) + # Policy selection uses the planned lengths; only the compiled program + # replaces them with runtime operands. + fixed_seqlen = cfg.fixed_seqlen + cfg.fixed_seqlen = True if not split_kv and (cfg.use_split_kv or cfg.splits_kv > 1): raise ValueError("split_kv=False conflicts with an explicit split config") splits_kv, max_splits_kv = _resolve_explicit_split_controls( @@ -5461,8 +5833,8 @@ def make_decode_config( "fixed-Q max_seq_len_q must equal seq_len_q; use " "use_variable_seqlens_q for a runtime-varying Q length" ) - # Resolve Q first. The KV selector consumes the multi-Q cost-model result - # instead of forcing a Q64 profile ahead of the established Q policy. + # Score Q geometry with candidate KV fanouts on the common KV128 basis. + # Finalize the KV tile and reduction mode after selecting that joint recipe. auto_selection_explicit_fields = set(explicit_fields) selected_grouped_q_recipe = _apply_auto_grouped_q_mma_config( cfg, @@ -5535,7 +5907,9 @@ def make_decode_config( # Auto launch-mode selection. Only kicks in if the caller has not already # opted into a specific mode. The caller's split permission is independent # of Q storage; unsupported automatic recipes retain their direct fallback. - if selected_grouped_q_recipe is not None: + if selected_grouped_q_recipe is not None and not ( + _LAUNCH_SELECTION_FIELDS & explicit_fields + ): if selected_grouped_q_recipe.splits_kv > 1: split_kv_mode = selected_grouped_q_recipe.split_kv_mode splits_kv = selected_grouped_q_recipe.splits_kv @@ -5605,4 +5979,5 @@ def make_decode_config( num_heads_kv=num_heads_kv, split_kv_mode=split_kv_mode, ) + cfg.fixed_seqlen = fixed_seqlen return cfg 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 f120a13f02cb..7e7be50d43fb 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 @@ -56,13 +56,18 @@ # letting an intermediate probability exceed 2**8; the softmax identity is # unchanged apart from normal finite-precision rounding. As in the # 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``. It bounds +# adversarial values outside the qualified probability bound. Keeps profiles +# apply it; see ``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 +# KV256 tiles run at most this many score pairs of a streamed K32 fragment as +# FMA polynomials, below the share that KV128 tiles take on SM100: their +# 16-bit fragment loops slow down with a larger share. See +# ``FmhaDecodeConfig.fragment_ex2_emulated_pairs``. +KV_TILE_256_MAX_EX2_EMULATED_PAIRS = 4 + # A launch bound makes ptxas honor warpgroup ``setmaxnreg`` allocations, but # the resulting register hand-off has a fixed cost. Paired B200 measurements # show that it is amortized once a Q64/KV256 CTA processes at least 32 tiles @@ -207,19 +212,3 @@ OUTPUT_VALUES_PER_THREAD = 8 FP8_PACKED_OUTPUT_REGS_PER_THREAD = 2 PACKED_OUTPUT_REGS_PER_THREAD = 4 - -# INT8 Q/K scores accumulate onto the bit pattern of ``1.5 * 2**23``. Integers -# in ``[2**23, 2**24)`` are FP32 numbers with unit spacing and every D = 128 -# INT8 dot product has ``|score| <= 2**21``, so the INT32 accumulator is the -# exact FP32 number ``12582912.0 + score`` and the softmax reads it without a -# conversion, removing the bias once per scale group. -INT32_SCORE_BIAS = 12582912.0 -# One ``kind::f16`` MMA step writes the bias before the INT8 K steps: both -# operands read one BF16 tile whose K = 16 rows repeat -# ``[1024, 1024, 1024, 0]``, so the twelve products ``2**20`` sum exactly to -# ``1.5 * 2**23``. The tile is unswizzled K-major 32-byte rows with 128-byte -# LBO and 256-byte SBO core-matrix strides. -INT32_SCORE_SEED_MMA_K = 16 -INT32_SCORE_SEED_TILE_WORDS = (0x44804480, 0x00004480) -INT32_SCORE_SEED_TILE_LBO = 128 -INT32_SCORE_SEED_TILE_SBO = 256 diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_kernel.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_kernel.py index 7a2235d8bc28..41372840bf31 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 @@ -30,6 +30,7 @@ """ import math +from dataclasses import replace import cutlass import cutlass.experimental.cuda as cuda @@ -51,7 +52,6 @@ MemoryResource, PipelineConfig, TileSchedulerConfig, - WorkQueue, ) from cutlass.experimental.task_scheduling.task import Task from cutlass.experimental.task_scheduling.task_manager import TaskManager @@ -64,7 +64,6 @@ ) from ..._block_sparse.prepared import _BlockSparseRouteLayout from ..tensor_map import ( - create_tensor_map_ragged_from_tensor, create_tensor_map_tiled, create_tensor_map_tiled_from_view, ) @@ -96,10 +95,15 @@ TmemTransformedKvResource, ) from .fmha_decode_resources.helpers_common import ( + split_kv_grid_coord, + UniformQueryOffsets, _q_group_token_base, _q_seq_bounds, ) -from .fmha_decode_resources.helpers_kv_tile_idx import _runtime_active_splits_kv +from .fmha_decode_resources.helpers_kv_tile_idx import ( + _runtime_active_splits_kv, + runtime_split_kv_is_active, +) from .fmha_decode_resources.sage_scales import ( SageKScalesResource, SageVScalesResource, @@ -249,7 +253,7 @@ def _compute_total_kv_tiles(seq_len_kv: int, tile_size_kv: int) -> int: return (seq_len_kv + tile_size_kv - 1) // tile_size_kv -def _decode_min_blocks_per_mp(cfg: FmhaDecodeConfig, seq_len_kv: int) -> int: +def _decode_min_blocks_per_mp(cfg: FmhaDecodeConfig, seq_len_kv: int | None) -> int: """Return the launch bound needed for dynamic register reallocation. ``setmaxnreg`` needs a kernel-entry occupancy bound before ptxas can infer @@ -257,9 +261,9 @@ def _decode_min_blocks_per_mp(cfg: FmhaDecodeConfig, seq_len_kv: int) -> int: hand-off cost for a long enough mainloop; the established profiles below retain their existing unconditional launch bounds. """ - kv256_reallocation = ( - cfg.tile_size_kv == 256 - and _compute_total_kv_tiles(seq_len_kv, cfg.tile_size_kv) + kv256_reallocation = cfg.tile_size_kv == 256 and ( + seq_len_kv is None + or _compute_total_kv_tiles(seq_len_kv, cfg.tile_size_kv) >= KV_TILE_256_REGISTER_REALLOCATION_MIN_TILES ) return int( @@ -273,22 +277,27 @@ def _decode_min_blocks_per_mp(cfg: FmhaDecodeConfig, seq_len_kv: int) -> int: def _compute_static_num_skipped_kv_tiles(cfg: FmhaDecodeConfig, seq_len_kv: int) -> int: """Return full leading KV tiles skipped by a static sliding window.""" - if not cfg.use_sliding_window_causal or cfg.max_seq_len_q > 1: + if not cfg.use_sliding_window_causal or cfg.is_multi_token_q: return 0 return max(seq_len_kv - cfg.attention_window_size, 0) // cfg.tile_size_kv def _compute_static_window_start_idx(cfg: FmhaDecodeConfig, seq_len_kv: int) -> int: """Return the token index where a static sliding window begins.""" - if not cfg.use_sliding_window_causal or cfg.max_seq_len_q > 1: + if not cfg.use_sliding_window_causal or cfg.is_multi_token_q: return 0 return max(seq_len_kv - cfg.attention_window_size, 0) def _configure_static_sliding_window( - cfg: FmhaDecodeConfig, seq_len_kv: int, bias_kv_tma: bool = False + cfg: FmhaDecodeConfig, seq_len_kv: int | None, bias_kv_tma: bool = False ) -> int: """Populate fixed-length sliding metadata and return effective seqLenKv.""" + if seq_len_kv is None: + cfg.use_static_sliding_kv_tma_bias = False + cfg.static_num_skipped_kv_tiles = 0 + cfg.static_window_start_idx = 0 + return 0 skipped_tiles = _compute_static_num_skipped_kv_tiles(cfg, seq_len_kv) skipped_tokens = skipped_tiles * cfg.tile_size_kv window_start_idx = _compute_static_window_start_idx(cfg, seq_len_kv) @@ -331,11 +340,12 @@ def _build_decode_gen_schedule( split_kv_counter_ptr: cute.Pointer | None = None, attention_sinks_ptr: cute.Pointer | None = None, seqlens_kv: cute.Pointer | None = None, - cu_seqlens_q: cute.Pointer | None = None, + cu_seqlens_q: cute.Pointer | UniformQueryOffsets | None = None, max_seq_len_kv: int | Int32 = 0, corr_max_seq_len_kv: int | Int32 | None = None, num_heads_kv: Int32 | None = None, h_r: Int32 | None = None, + max_q_len: int | Int32 = 1, tma_desc_q: cutlass.Pointer | None = None, tma_desc_k: cutlass.Pointer | None = None, tma_desc_v: cutlass.Pointer | None = None, @@ -369,6 +379,8 @@ def _build_decode_gen_schedule( sparse_row_route_begin: Int32 | None = None, sparse_route_count: Int32 | None = None, sage_scales_ptr: cute.Pointer | None = None, + lse_ptr: cute.Pointer | None = None, + requested_splits=None, ) -> tuple[ list[Task], dict[MemoryResource, list[MemoryResource]], @@ -639,6 +651,7 @@ def _build_decode_gen_schedule( uses_page_offsets_task and use_per_inst_kv_resources and not use_one_inst_qkv + and not cfg.runtime_page_window and not ( use_native_paged_kv and cfg.uses_q_token_kv_block_sparse_page_route @@ -745,10 +758,12 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: # 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. + # P left the TMEM columns it aliases or, with SMEM P, after the previous + # tile's max pass released S. A route or ``sfK`` stage is read no later + # than that pass (an SMEM-P Sage row is dequantized there), so the + # instance's Softmax has released every stage 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: @@ -905,14 +920,13 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: "pipeline_config": wq_pipeline_config, "name": "work_queue", } - if cfg.use_variable_seqlens_q: - work_queue = PackedDecodeWorkQueue( - cfg=cfg, - cu_seqlens_q=cu_seqlens_q, - **work_queue_kwargs, - ) - else: - work_queue = WorkQueue(**work_queue_kwargs) + work_queue = PackedDecodeWorkQueue( + cfg=cfg, + cu_seqlens_q=cu_seqlens_q, + seqlens_kv=seqlens_kv, + num_heads_kv=num_heads_kv, + **work_queue_kwargs, + ) schedule_token_throttle = ScheduleTokenThrottleResource( pipeline_config=PipelineConfig.create_async_async_pipeline_cfg( @@ -972,6 +986,9 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: pipeline_config=smem_page_offsets_cfg, cfg=cfg, stage_page_ids_per_tile=stage_page_ids_per_tile, + # The full membership row has no static allocation extent when + # KV lengths are runtime. Read its immutable GMEM words on demand. + cache_memberships_in_smem=cfg.fixed_seqlen, page_idx_kv=page_idx_kv, page_table_stride=page_table_stride, num_heads_kv=num_heads_kv, @@ -1246,7 +1263,7 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: sage_k_scales0 = None sage_k_scales1 = None if cfg.use_sage_attention: - + # Connect one K-scale SMEM ring from the load warp to its softmax instance. def _sage_k_scales(inst_id: int, softmax_grp) -> SageKScalesResource: return SageKScalesResource( pipeline_config=PipelineConfig( @@ -1423,8 +1440,10 @@ def _sage_k_scales(inst_id: int, softmax_grp) -> SageKScalesResource: partial_stats_ptr=partial_stats_ptr, split_kv_counter_ptr=split_kv_counter_ptr, attention_sinks_ptr=attention_sinks_ptr, + lse_ptr=lse_ptr, seqlens_kv=kv_seqlens, max_seq_len_kv=corr_max_seq_len_kv, + max_q_len=max_q_len, num_heads_kv=num_heads_kv, h_r=h_r, h_k_idx=h_k_idx, @@ -1452,8 +1471,10 @@ def _sage_k_scales(inst_id: int, softmax_grp) -> SageKScalesResource: partial_stats_ptr=partial_stats_ptr, split_kv_counter_ptr=split_kv_counter_ptr, attention_sinks_ptr=attention_sinks_ptr, + lse_ptr=lse_ptr, seqlens_kv=kv_seqlens, max_seq_len_kv=corr_max_seq_len_kv, + max_q_len=max_q_len, num_heads_kv=num_heads_kv, h_r=h_r, h_k_idx=h_k_idx, @@ -1529,6 +1550,7 @@ def _sage_k_scales(inst_id: int, softmax_grp) -> SageKScalesResource: if sage_v_scales is not None: smem_resources.append(sage_v_scales) + # Lay out all resource storage, CLC responses and the shared TMEM pointer. def allocate_smem() -> tuple[SmemAllocator, SmemAllocation | None]: allocator = SmemAllocator() clc_response_alloc = None @@ -2265,7 +2287,7 @@ def smem_bytes(allocator: SmemAllocator) -> int: tmem_allocator.add_resource(transformed_kv) tmem_allocator.compute_layout() if cfg.use_keeps_mma_ab and not use_one_inst_qkv and not cfg.uses_tmem_p: - # Two-inst Keeps with SMEM P (currently Q64) must retain the historical + # Two-inst Keeps with SMEM P (Q64/KV128 and 8-bit Q128/KV128) keeps the # standalone-first S/O layout after unused stats aliases are removed. s0_alloc = tmem_s0.get_tmem_requirements()[0] s1_alloc = tmem_s1.get_tmem_requirements()[0] @@ -2360,6 +2382,8 @@ def smem_bytes(allocator: SmemAllocator) -> int: # The score-seed operand tile is written once by its owning instance. eager_init_resources.append(tmem_s0) + for task in task_list: + task.requested_splits = requested_splits return ( task_list, resource_dependency_graph, @@ -2435,6 +2459,7 @@ def build_decode_task_manager( total_kv_tiles, tile_sched_params=None, num_heads_kv=Int32(num_heads_kv), + lse_ptr=None, ) tm = TaskManager( @@ -2551,10 +2576,10 @@ def _run_decode_gen_active( o_iter: cute.Pointer, g_s_k: Int32, g_h_k: Int32, - g_scale_s_log2_e: Float32, - g_output_scale: Float32, + scale_softmax_log2: Float32, + output_scale: Float32, g_seqlens_kv: cute.Pointer | DirectSparseMetadataView, - g_cu_seqlens_q: cute.Pointer, + g_cu_seqlens_q: cute.Pointer | UniformQueryOffsets, g_page_idx_kv: cute.Pointer | DirectSparseMetadataView, g_k_sf: cute.Pointer, g_v_sf: cute.Pointer, @@ -2577,7 +2602,7 @@ def _run_decode_gen_active( defer_runtime_split_pruning: cutlass.Constexpr[bool], tile_sched_params: utils.ClcDynamicPersistentTileSchedulerParams | None, cfg: cutlass.Constexpr[FmhaDecodeConfig], - seq_len_kv: cutlass.Constexpr[int] = 2048, + seq_len_kv: cutlass.Constexpr[int | None] = 2048, use_variable_seqlens_kv: cutlass.Constexpr[bool] = False, use_native_paged_kv: cutlass.Constexpr[bool] = False, use_static_native_seqlens_kv: cutlass.Constexpr[bool] = False, @@ -2589,6 +2614,9 @@ def _run_decode_gen_active( tma_desc_k_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, tma_desc_v_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, g_sage_scales: cute.Pointer | None = None, + lse_ptr: cute.Pointer | None = None, + requested_splits=None, + g_max_q_len: Int32 = 1, ) -> None: """Run the complete decode body for one runtime-valid Q tile. @@ -2604,6 +2632,7 @@ def _run_decode_gen_active( bias_static_sliding_kv_tma = ( cfg.use_sliding_window_causal + and cfg.fixed_seqlen and cfg.max_seq_len_q == 1 and not cfg.use_paged_kv and not cfg.use_split_kv @@ -2622,7 +2651,7 @@ def _run_decode_gen_active( ) runtime_max_seq_len_kv = ( g_s_k - if cutlass.const_expr(use_runtime_seqlens_kv) + if cutlass.const_expr(use_runtime_seqlens_kv or not cfg.fixed_seqlen) else Int32(cfg.static_seq_len_kv) ) tma_desc_k_summary_ptr = None @@ -2668,7 +2697,7 @@ def _run_decode_gen_active( init_warp += 1 q_output_rows = g_h_r - if cutlass.const_expr(cfg.max_seq_len_q > 1): + if cutlass.const_expr(cfg.fixed_seqlen and cfg.is_multi_token_q): q_output_rows = g_h_r * Int32(cfg.max_seq_len_q) # Static block-sparse tiles read their prepared row header here, before @@ -2702,9 +2731,9 @@ def _run_decode_gen_active( ) = _build_decode_gen_schedule( cfg, total_kv_tiles, - scale_softmax_log2=g_scale_s_log2_e, + scale_softmax_log2=scale_softmax_log2, o_ptr=o_iter, - output_scale=g_output_scale, + output_scale=output_scale, partial_o_ptr=g_partial_o, partial_stats_ptr=g_partial_stats, split_kv_counter_ptr=g_split_kv_counter, @@ -2712,9 +2741,10 @@ def _run_decode_gen_active( seqlens_kv=runtime_seqlens_kv, cu_seqlens_q=g_cu_seqlens_q, max_seq_len_kv=runtime_max_seq_len_kv, - corr_max_seq_len_kv=seq_len_kv, + corr_max_seq_len_kv=runtime_max_seq_len_kv, num_heads_kv=g_h_k, h_r=q_output_rows, + max_q_len=g_max_q_len, tma_desc_q=tma_desc_q.get_ptr(), tma_desc_k=tma_desc_k.get_ptr(), tma_desc_v=tma_desc_v.get_ptr(), @@ -2748,6 +2778,8 @@ def _run_decode_gen_active( sparse_route_metadata=g_sparse_route_metadata, sparse_row_route_begin=sparse_row_route_begin, sparse_route_count=sparse_route_count, + lse_ptr=lse_ptr, + requested_splits=requested_splits, ) smem_allocator.allocate() @@ -2817,16 +2849,28 @@ def _run_decode_gen_active( or (use_native_paged_kv and not use_static_native_seqlens_kv) ): runtime_seq_len_kv = Int32(g_seqlens_kv[b_idx]) - deferred_active_splits_kv = _runtime_active_splits_kv( - cfg, - runtime_seq_len_kv, - seq_len_q, - _q_group_token_base(cfg, q_group_idx), - ) - split_idx = cute.arch.block_idx()[0] % Int32(cfg.splits_kv) - run_task_graph = cute.arch.make_warp_uniform( - split_idx < deferred_active_splits_kv - ) + split_idx = Int32(cute.arch.block_idx()[0]) + if cutlass.const_expr(not cfg.uses_compact_split_grid): + split_idx = split_idx % Int32(cfg.splits_kv) + if cutlass.const_expr(cfg.runtime_split_kv and requested_splits is not None): + active_splits = _runtime_active_splits_kv( + cfg, + runtime_seq_len_kv, + seq_len_q, + _q_group_token_base(cfg, q_group_idx), + requested_splits=requested_splits, + ) + run_task_graph = cute.arch.make_warp_uniform(split_idx < active_splits) + else: + run_task_graph = cute.arch.make_warp_uniform( + runtime_split_kv_is_active( + cfg, + runtime_seq_len_kv, + seq_len_q, + _q_group_token_base(cfg, q_group_idx), + split_idx, + ) + ) _retire_inactive_pdl_split( run_task_graph, tmem_ptr_i32, @@ -2893,10 +2937,10 @@ def _run_decode_gen_runtime_prefix( o_iter: cute.Pointer, g_s_k: Int32, g_h_k: Int32, - g_scale_s_log2_e: Float32, - g_output_scale: Float32, + scale_softmax_log2: Float32, + output_scale: Float32, g_seqlens_kv: cute.Pointer | DirectSparseMetadataView, - g_cu_seqlens_q: cute.Pointer, + g_cu_seqlens_q: cute.Pointer | UniformQueryOffsets, g_page_idx_kv: cute.Pointer | DirectSparseMetadataView, g_k_sf: cute.Pointer, g_v_sf: cute.Pointer, @@ -2918,7 +2962,7 @@ def _run_decode_gen_runtime_prefix( q_token_base: Int32, tile_sched_params: utils.ClcDynamicPersistentTileSchedulerParams | None, cfg: cutlass.Constexpr[FmhaDecodeConfig], - seq_len_kv: cutlass.Constexpr[int], + seq_len_kv: cutlass.Constexpr[int | None], use_variable_seqlens_kv: cutlass.Constexpr[bool], use_native_paged_kv: cutlass.Constexpr[bool], use_static_native_seqlens_kv: cutlass.Constexpr[bool], @@ -2930,6 +2974,9 @@ def _run_decode_gen_runtime_prefix( tma_desc_k_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, tma_desc_v_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, g_sage_scales: cute.Pointer | None = None, + lse_ptr: cute.Pointer | None = None, + requested_splits=None, + g_max_q_len: Int32 = 1, ) -> None: """Run the general runtime split-prefix producer or retire its suffix.""" @@ -2947,6 +2994,7 @@ def _run_decode_gen_runtime_prefix( runtime_seq_len_kv, seq_len_q, q_token_base, + requested_splits=requested_splits, ) if cutlass.const_expr(not cfg.use_separate_reduction_kernel): # Preserve one neutral producer for fused empty-K semantics. @@ -2961,7 +3009,9 @@ def _run_decode_gen_runtime_prefix( # This also avoids changing the grid or adding a reducer launch. if active_splits_kv == Int32(cfg.splits_kv - 1): active_splits_kv = Int32(cfg.splits_kv) - split_idx = q_group_cta_idx % Int32(cfg.splits_kv) + split_idx = Int32(cute.arch.block_idx()[0]) + if cutlass.const_expr(not cfg.uses_compact_split_grid): + split_idx = q_group_cta_idx % Int32(cfg.splits_kv) split_is_active = split_idx < active_splits_kv if cutlass.const_expr(cfg.supports_cluster_smem_reduction): @@ -2977,8 +3027,8 @@ def _run_decode_gen_runtime_prefix( o_iter, g_s_k, g_h_k, - g_scale_s_log2_e, - g_output_scale, + scale_softmax_log2, + output_scale, g_seqlens_kv, g_cu_seqlens_q, g_page_idx_kv, @@ -3015,6 +3065,9 @@ def _run_decode_gen_runtime_prefix( tma_desc_k_summary_atom=tma_desc_k_summary_atom, tma_desc_v_summary_atom=tma_desc_v_summary_atom, g_sage_scales=g_sage_scales, + lse_ptr=lse_ptr, + requested_splits=requested_splits, + g_max_q_len=g_max_q_len, ) else: _run_decode_gen_inactive_cluster_rank() @@ -3031,8 +3084,8 @@ def _run_decode_gen_runtime_prefix( o_iter, g_s_k, g_h_k, - g_scale_s_log2_e, - g_output_scale, + scale_softmax_log2, + output_scale, g_seqlens_kv, g_cu_seqlens_q, g_page_idx_kv, @@ -3069,6 +3122,9 @@ def _run_decode_gen_runtime_prefix( tma_desc_k_summary_atom=tma_desc_k_summary_atom, tma_desc_v_summary_atom=tma_desc_v_summary_atom, g_sage_scales=g_sage_scales, + lse_ptr=lse_ptr, + requested_splits=requested_splits, + g_max_q_len=g_max_q_len, ) else: _signal_padded_pdl_producer(cfg) @@ -3086,10 +3142,10 @@ def decode_gen_kernel( o_iter: cute.Pointer, g_s_k: Int32, g_h_k: Int32, - g_scale_s_log2_e: Float32, - g_output_scale: Float32, + scale_softmax_log2: Float32, + output_scale: Float32, g_seqlens_kv: cute.Pointer | DirectSparseMetadataView, - g_cu_seqlens_q: cute.Pointer, + g_cu_seqlens_q: cute.Pointer | UniformQueryOffsets, g_page_idx_kv: cute.Pointer | DirectSparseMetadataView, g_k_sf: cute.Pointer, g_v_sf: cute.Pointer, @@ -3104,7 +3160,7 @@ def decode_gen_kernel( g_h_r: Int32, tile_sched_params: utils.ClcDynamicPersistentTileSchedulerParams | None, cfg: cutlass.Constexpr[FmhaDecodeConfig], - seq_len_kv: cutlass.Constexpr[int] = 2048, + seq_len_kv: cutlass.Constexpr[int | None] = 2048, use_variable_seqlens_kv: cutlass.Constexpr[bool] = False, use_native_paged_kv: cutlass.Constexpr[bool] = False, use_static_native_seqlens_kv: cutlass.Constexpr[bool] = False, @@ -3117,9 +3173,27 @@ def decode_gen_kernel( tma_desc_k_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, tma_desc_v_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, g_sage_scales: cute.Pointer | None = None, + lse_ptr: cute.Pointer | None = None, + requested_splits=None, + g_max_q_len: Int32 = 1, ) -> 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() + if cutlass.const_expr(not cfg.fixed_seqlen): + from copy import copy + + cfg = copy(cfg) + cfg.static_seq_len_kv = g_s_k + cfg.uniform_query_offsets = isinstance(g_cu_seqlens_q, UniformQueryOffsets) + if cutlass.const_expr(cfg.groups_tokens_heads_q): + cfg.runtime_q_ctas = Int32( + (cutlass.Uint32(g_max_q_len) + cfg.q_tokens_per_cta - 1) + // cfg.q_tokens_per_cta + ) + else: + cfg.runtime_q_ctas = g_max_q_len * ( + (cfg.heads_q_per_kv + cfg.tile_size_q - 1) // cfg.tile_size_q + ) + q_group_cta_idx, h_k_idx, b_idx = split_kv_grid_coord(cfg, cute.arch.block_idx()) if cutlass.const_expr( cfg.use_q_token_kv_block_sparse_route and not cfg.shares_sparse_pattern @@ -3140,12 +3214,15 @@ def decode_gen_kernel( g_q_token_kv_block_sparse_page_membership_stride * g_h_k ) q_group_idx = q_group_cta_idx - if cutlass.const_expr(cfg.use_split_kv): - # Grid coordinates and scratch linearization retain configured fanout. + if cutlass.const_expr(cfg.uses_compact_split_grid): + q_group_idx = Int32(cute.arch.block_idx()[2]) % Int32(cfg.num_q_ctas) + elif cutlass.const_expr(cfg.use_split_kv): q_group_idx = q_group_cta_idx // Int32(cfg.splits_kv) q_token_offset = Int32(0) - seq_len_q = Int32(cfg.max_seq_len_q) + seq_len_q = Int32(0) + if cutlass.const_expr(not cfg.use_variable_seqlens_q): + seq_len_q = Int32(cfg.max_seq_len_q) q_token_base = Int32(0) q_tile_is_active = cutlass.Boolean(True) if cutlass.const_expr(not cfg.use_persistent_scheduler): @@ -3153,6 +3230,15 @@ def decode_gen_kernel( q_token_offset, seq_len_q = _q_seq_bounds(cfg, g_cu_seqlens_q, b_idx) q_token_base = _q_group_token_base(cfg, q_group_idx) q_tile_is_active = q_token_base < seq_len_q + if cutlass.const_expr( + use_variable_seqlens_kv + or (use_native_paged_kv and not use_static_native_seqlens_kv) + ): + if cutlass.const_expr(cfg.use_pdl): + prims.griddepcontrol(kind=prims.GridDepAction.WAIT) + q_tile_is_active = q_tile_is_active & ( + Int32(g_seqlens_kv[b_idx]) > Int32(0) + ) # Persistent block coordinates identify physical workers rather than # logical tiles; their WorkQueue owns the equivalent Q predicate. Split-KV @@ -3170,8 +3256,8 @@ def decode_gen_kernel( o_iter, g_s_k, g_h_k, - g_scale_s_log2_e, - g_output_scale, + scale_softmax_log2, + output_scale, g_seqlens_kv, g_cu_seqlens_q, g_page_idx_kv, @@ -3191,8 +3277,10 @@ def decode_gen_kernel( b_idx, q_token_offset, seq_len_q, - Int32(cfg.splits_kv), - True, + Int32(cfg.splits_kv) + if requested_splits is None + else requested_splits[0], + requested_splits is None, False, tile_sched_params, cfg, @@ -3203,6 +3291,9 @@ def decode_gen_kernel( g_sparse_row_route_offsets, g_sparse_row_route_counts, g_sparse_route_metadata, + lse_ptr=lse_ptr, + requested_splits=requested_splits, + g_max_q_len=g_max_q_len, ) elif cutlass.const_expr(cfg.use_split_kv and cfg.use_pdl): # The immediately preceding grid may produce seq_lens. Send every @@ -3219,8 +3310,8 @@ def decode_gen_kernel( o_iter, g_s_k, g_h_k, - g_scale_s_log2_e, - g_output_scale, + scale_softmax_log2, + output_scale, g_seqlens_kv, g_cu_seqlens_q, g_page_idx_kv, @@ -3257,6 +3348,9 @@ def decode_gen_kernel( tma_desc_k_summary_atom=tma_desc_k_summary_atom, tma_desc_v_summary_atom=tma_desc_v_summary_atom, g_sage_scales=g_sage_scales, + lse_ptr=lse_ptr, + requested_splits=requested_splits, + g_max_q_len=g_max_q_len, ) else: _run_decode_gen_runtime_prefix( @@ -3270,8 +3364,8 @@ def decode_gen_kernel( o_iter, g_s_k, g_h_k, - g_scale_s_log2_e, - g_output_scale, + scale_softmax_log2, + output_scale, g_seqlens_kv, g_cu_seqlens_q, g_page_idx_kv, @@ -3307,10 +3401,13 @@ def decode_gen_kernel( tma_desc_k_summary_atom=tma_desc_k_summary_atom, tma_desc_v_summary_atom=tma_desc_v_summary_atom, g_sage_scales=g_sage_scales, + lse_ptr=lse_ptr, + requested_splits=requested_splits, + g_max_q_len=g_max_q_len, ) else: - # Packed-Q grids use a batch-wide maximum envelope. These Q CTAs own no - # producer state, unlike a valid Q tile whose K domain is empty. + # Inactive requests publish no output or partials; dependent reducers + # use the same live-length predicate. Every physical CTA still signals PDL. _signal_padded_pdl_producer(cfg) @@ -3324,7 +3421,7 @@ def fmha_decode_launch( v_sf_iter: cute.Pointer, o_iter: cute.Pointer, seqlens_kv_iter: cute.Pointer | DirectSparseMetadataView, - cu_seqlens_q_iter: cute.Pointer, + cu_seqlens_q_iter: cute.Pointer | UniformQueryOffsets, total_q_tokens: Int32, page_idx_kv_iter: cute.Pointer | DirectSparseMetadataView, q_token_kv_block_sparse_page_memberships_iter: cute.Pointer, @@ -3332,13 +3429,13 @@ def fmha_decode_launch( partial_stats_iter: cute.Pointer, split_kv_counter_iter: cute.Pointer, attention_sinks_iter: cute.Pointer, - scale_s: Float32, + softmax_scale: Float32, output_scale: Float32, kv_b_stride: Int32, max_active_clusters: Int32, stream: cuda_drv.CUstream, cfg: cutlass.Constexpr[FmhaDecodeConfig], - seq_len_kv: cutlass.Constexpr[int] = 2048, + seq_len_kv: cutlass.Constexpr[int | None] = 2048, use_variable_seqlens_kv: cutlass.Constexpr[bool] = False, use_native_paged_kv: cutlass.Constexpr[bool] = False, page_table_stride: Int64 = 0, @@ -3354,17 +3451,31 @@ def fmha_decode_launch( static_full_split_prefix: cutlass.Constexpr[bool] = False, use_static_native_seqlens_kv: cutlass.Constexpr[bool] = False, sage_scales_iter: cute.Pointer | None = None, + lse_ptr: cute.Pointer | None = None, + requested_splits=None, + max_q_len: Int32 = 0, + runtime_q_capacity_rows: Int32 = 0, ) -> None: """Standalone JIT launcher for FMHA decode TS. + ``max_q_len`` bounds the launched Q-token grid. ``runtime_q_capacity_rows`` + is the partial-output row stride per request/KV head and can exceed that + grid's live rows; both operands remain runtime values. ``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) + if cutlass.const_expr(requested_splits is not None): + # Prepare the reciprocal once on the host; counts stay runtime operands. + requested_splits = ( + requested_splits, + cute.fast_divmod_create_divisor(requested_splits * Int32(cfg.num_insts_kv)), + ) b, h_q, h_k, s_k, d = problem_shape h_r = h_q // h_k bias_static_sliding_kv_tma = ( cfg.use_sliding_window_causal + and cfg.fixed_seqlen and cfg.max_seq_len_q == 1 and not cfg.use_paged_kv and not cfg.use_split_kv @@ -3374,7 +3485,7 @@ def fmha_decode_launch( cfg, seq_len_kv, bias_static_sliding_kv_tma ) - q_seq = Int32(cfg.max_seq_len_q) + q_seq = max_q_len if not cfg.fixed_seqlen else Int32(cfg.max_seq_len_q) if cutlass.const_expr(cfg.use_paged_kv): kv_tma_d = d if cutlass.const_expr( @@ -3534,9 +3645,8 @@ def fmha_decode_launch( ) q_box_dims: tuple[object, ...] if cutlass.const_expr(cfg.use_variable_seqlens_q): - # Packed Q is physically [sum_q_tokens, num_heads_q, head_dim]. Flatten - # the head modes into Hq. Sparse Q uses this rank-3 map directly; - # other packed paths retain per-sequence ragged zero fill. + # Bound packed Q by its physical [D, Hq, total_q_tokens] allocation. + # Per-route Q masks handle tiles that cross into the next request. q_tma = cute.make_tensor( q_iter, cute.make_layout( @@ -3550,33 +3660,21 @@ def fmha_decode_launch( cfg.heads_q_per_kv, cfg.q_tma_tokens_per_cta, ) - q_groups = Int32( - (cfg.max_seq_len_q + cfg.q_tokens_per_cta - 1) // cfg.q_tokens_per_cta + q_groups = (q_seq + Int32(cfg.q_tokens_per_cta - 1)) // Int32( + cfg.q_tokens_per_cta ) else: q_box_dims = (tma_box0_q, cfg.tile_size_q, 1) head_ctas_per_token = Int32( (cfg.heads_q_per_kv + cfg.tile_size_q - 1) // cfg.tile_size_q ) - q_groups = Int32(cfg.max_seq_len_q) * head_ctas_per_token - if cutlass.const_expr(cfg.uses_q_token_kv_block_sparse_page_route): - # A partial group may read the next sequence into inactive Q rows. - # QK/softmax/PV rows are independent; row masks and guarded stores - # discard them. Physical allocation bounds still get TMA zero fill. - tma_desc_q = create_tensor_map_tiled_from_view( - q_tma, - box_dims=q_box_dims, - stride_order=(0, 1, 2), - swizzle=tma_swizzle_q, - ) - else: - tma_desc_q = create_tensor_map_ragged_from_tensor( - q_tma, - box_dims=q_box_dims, - ragged_dim=2, - stride_order=(0, 1, 2), - swizzle=tma_swizzle_q, - ) + q_groups = q_seq * head_ctas_per_token + tma_desc_q = create_tensor_map_tiled_from_view( + q_tma, + box_dims=q_box_dims, + stride_order=(0, 1, 2), + swizzle=tma_swizzle_q, + ) else: q_tma = cute.make_tensor( q_iter, @@ -3593,8 +3691,8 @@ def fmha_decode_launch( cfg.q_tokens_per_cta, 1, ) - q_groups = Int32( - (cfg.max_seq_len_q + cfg.q_tokens_per_cta - 1) // cfg.q_tokens_per_cta + q_groups = (q_seq + Int32(cfg.q_tokens_per_cta - 1)) // Int32( + cfg.q_tokens_per_cta ) else: q_box_dims = (tma_box0_q, cfg.tile_size_q, 1, 1, 1) @@ -3688,8 +3786,18 @@ def fmha_decode_launch( ) grid_x = q_groups + grid_batch = b if cutlass.const_expr(cfg.use_split_kv): grid_x = q_groups * Int32(cfg.splits_kv) + if cutlass.const_expr(cfg.uses_compact_split_grid): + # X contains only live policy splits. Q/batch share Z with a + # compile-time Q stride, avoiding runtime-count coordinate division. + grid_x = ( + Int32(cfg.splits_kv) + if requested_splits is None + else requested_splits[0] + ) + grid_batch = b * q_groups if cutlass.const_expr(cfg.use_persistent_scheduler): tile_sched_params = utils.ClcDynamicPersistentTileSchedulerParams( @@ -3699,7 +3807,7 @@ def fmha_decode_launch( grid = tile_sched_params.get_grid_shape() else: tile_sched_params = None - grid = (grid_x, h_k, b) + grid = (grid_x, h_k, grid_batch) # cluster distributed-SMEM reduction groups the splits_kv split CTAs # of each sequence (contiguous in grid-X) into one cluster so they can write # partials into each other's shared memory via prims.mapa. @@ -3711,55 +3819,84 @@ def fmha_decode_launch( 0, mem_space=cutlass.AddressSpace.gmem, ) - decode_gen_kernel( - tma_desc_q, - tma_desc_k, - tma_desc_v, - tma_desc_k_sf, - tma_desc_v_sf, - tma_desc_k, - tma_desc_v, - o_iter, - s_k, - h_k, - Float32(scale_s * log2_e), - output_scale, - seqlens_kv_iter, - cu_seqlens_q_iter, - page_idx_kv_iter, - k_sf_iter, - v_sf_iter, - page_table_stride, - page_table_capacity, - q_token_kv_block_sparse_page_memberships_iter, - q_token_kv_block_sparse_page_membership_stride, - partial_o_iter, - partial_stats_iter, - split_kv_counter_iter, - attention_sinks_iter, - h_r, - tile_sched_params, - cfg, - seq_len_kv, - use_variable_seqlens_kv, - use_native_paged_kv, - use_static_native_seqlens_kv, - null_sparse_route_ptr, - null_sparse_route_ptr, - null_sparse_route_ptr, - static_full_split_prefix, - g_sage_scales=sage_scales_iter, - ).launch( - grid=grid, - block=[cfg.threads_per_cta, 1, 1], - cluster=cluster_shape, - stream=stream, - min_blocks_per_mp=_decode_min_blocks_per_mp(cfg, effective_seq_len_kv), - # PDL-capable attention initializes all local task resources before - # acquiring its immediately preceding producer. Split attention also - # defers runtime prefix pruning until after that acquire. - use_pdl=cfg.use_pdl, + # Both schedules are compiled together and selected by runtime bounds. + # A CTA-relative 32-ID stage covers any instruction-aligned short span; + # longer spans retain the ordinary staged pipeline. Counts never enter + # the compile identity or require another JIT when this choice changes. + can_hold_runtime_window = ( + cfg.runtime_split_kv + and requested_splits is not None + and use_native_paged_kv + and cfg.use_split_kv + and not cfg.use_sliding_window_causal + and not cfg.uses_scattered_page_route + and not cfg.use_block_sparse + and cfg.num_insts_kv * (cfg.tile_size_kv // cfg.num_tokens_per_page) <= 32 ) + use_runtime_window = cutlass.Boolean(False) + if cutlass.const_expr(can_hold_runtime_window): + pages_per_tile = cfg.tile_size_kv // cfg.num_tokens_per_page + window_tiles = (32 // pages_per_tile // cfg.num_insts_kv) * cfg.num_insts_kv + use_runtime_window = s_k <= requested_splits[0] * Int32( + window_tiles * cfg.tile_size_kv + ) + for profile in cutlass.range_constexpr(2 if can_hold_runtime_window else 1): + launch_cfg = replace(cfg, runtime_page_window=(profile == 1)) + if use_runtime_window == cutlass.Boolean(profile == 1): + decode_gen_kernel( + tma_desc_q, + tma_desc_k, + tma_desc_v, + tma_desc_k_sf, + tma_desc_v_sf, + tma_desc_k, + tma_desc_v, + o_iter, + s_k, + h_k, + Float32(softmax_scale * log2_e), + output_scale, + seqlens_kv_iter, + cu_seqlens_q_iter, + page_idx_kv_iter, + k_sf_iter, + v_sf_iter, + page_table_stride, + page_table_capacity, + q_token_kv_block_sparse_page_memberships_iter, + q_token_kv_block_sparse_page_membership_stride, + partial_o_iter, + partial_stats_iter, + split_kv_counter_iter, + attention_sinks_iter, + (runtime_q_capacity_rows if not cfg.fixed_seqlen else h_r), + tile_sched_params, + launch_cfg, + seq_len_kv, + use_variable_seqlens_kv, + use_native_paged_kv, + use_static_native_seqlens_kv, + null_sparse_route_ptr, + null_sparse_route_ptr, + null_sparse_route_ptr, + static_full_split_prefix, + g_sage_scales=sage_scales_iter, + lse_ptr=lse_ptr, + requested_splits=requested_splits, + g_max_q_len=max_q_len, + ).launch( + grid=grid, + block=[cfg.threads_per_cta, 1, 1], + cluster=cluster_shape, + stream=stream, + min_blocks_per_mp=_decode_min_blocks_per_mp( + cfg, None if seq_len_kv is None else effective_seq_len_kv + ), + # PDL-capable attention initializes all local task resources before + # acquiring its immediately preceding producer. Split attention also + # defers runtime prefix pruning until after that acquire. + use_pdl=cfg.use_pdl, + ) @cute.jit @@ -3774,10 +3911,10 @@ def fmha_block_sparse_launch( row_route_offsets_iter: cute.Pointer, row_route_counts_iter: cute.Pointer, route_metadata_iter: cute.Pointer, - scale_s: Float32, + softmax_scale: Float32, stream: cuda_drv.CUstream, cfg: cutlass.Constexpr[FmhaDecodeConfig], - seq_len_kv: cutlass.Constexpr[int], + seq_len_kv: cutlass.Constexpr[int | None], g_seqlens_kv: cute.Pointer | None = None, use_variable_seqlens_kv: cutlass.Constexpr[bool] = False, num_physical_kv_pages: Int64 = 0, @@ -3943,7 +4080,7 @@ def fmha_block_sparse_launch( if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): num_kv_blocks, _ = _block_sparse_proxy_summary_geometry( - seq_len_kv, + s_k, cfg.kv_block_size, ) _, summary_kv_strides = _block_sparse_bshd_tma_strides( @@ -3992,7 +4129,7 @@ def fmha_block_sparse_launch( ) q_groups = Int32( - (cfg.max_seq_len_q + cfg.q_tokens_per_cta - 1) // cfg.q_tokens_per_cta + (q_seq + Int32(cfg.q_tokens_per_cta - 1)) // Int32(cfg.q_tokens_per_cta) ) if cutlass.const_expr(cfg.use_persistent_scheduler): tile_sched_params = utils.ClcDynamicPersistentTileSchedulerParams( @@ -4028,7 +4165,7 @@ def fmha_block_sparse_launch( o_iter, s_k, h_k, - Float32(scale_s * log2_e), + Float32(softmax_scale * log2_e), Float32(1.0), seqlens_kv_iter, null_i32_ptr, @@ -4059,6 +4196,8 @@ def fmha_block_sparse_launch( tma_desc_k_summary_atom=k_desc_summary_atom, tma_desc_v_summary_atom=v_desc_summary_atom, g_sage_scales=sage_scales_iter, + lse_ptr=None, + g_max_q_len=q_seq, ).launch( grid=grid, block=[cfg.threads_per_cta, 1, 1], diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_common.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_common.py index 23fc82d712d5..7bef39eff3cf 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_common.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_common.py @@ -554,6 +554,18 @@ def _logical_head_batch( return fallback_h_k_idx, fallback_b_idx +@cute.jit +def split_kv_grid_coord( + cfg: Constexpr[FmhaDecodeConfig], coord: cute.Coord +) -> cute.Coord: + """Map compact split/Q/batch launches to capacity-strided task coordinates.""" + if cutlass.const_expr(cfg.uses_compact_split_grid): + q_group = Int32(coord[2]) % Int32(cfg.num_q_ctas) + batch = Int32(coord[2]) // Int32(cfg.num_q_ctas) + return (q_group * Int32(cfg.splits_kv) + Int32(coord[0]), coord[1], batch) + return coord + + @cute.jit def _logical_q_group_idx( cfg: Constexpr[FmhaDecodeConfig], @@ -561,6 +573,8 @@ def _logical_q_group_idx( fallback_q_group_idx: Int32, ) -> Int32: """Resolve the q-group from the persistent work tile or static launch.""" + if cutlass.const_expr(cfg.uses_compact_split_grid): + return Int32(cute.arch.block_idx()[2]) % Int32(cfg.num_q_ctas) if cutlass.const_expr(cfg.has_single_q_cta): # A split coordinate may still vary in grid X, but every physical CTA # maps to logical Q group zero. State this explicitly because the @@ -581,7 +595,7 @@ def _q_tile_output_row_base( """Return the first packed output row owned by a logical Q CTA.""" if cutlass.const_expr(cfg.uses_nontrivial_grouped_q_layout): return q_group_idx * Int32(cfg.q_tma_rows_per_cta) - if cutlass.const_expr(cfg.max_seq_len_q > 1): + if cutlass.const_expr(cfg.is_multi_token_q): head_ctas_per_token = Int32( (cfg.heads_q_per_kv + cfg.tile_size_q - 1) // cfg.tile_size_q ) @@ -593,13 +607,37 @@ def _q_tile_output_row_base( return q_group_idx * Int32(cfg.tile_size_q) +class UniformQueryOffsets: + """Affine offsets for fixed-layout Q, with a runtime token count.""" + + def __init__(self, seq_len_q): + self.values = (seq_len_q,) + + def __extract_mlir_values__(self): + return cutlass.extract_mlir_values(self.values) + + def __new_from_mlir_values__(self, values): + return UniformQueryOffsets(cutlass.new_from_mlir_values(self.values, values)[0]) + + @property + def seq_len_q(self): + return self.values[0] + + @cute.jit + def __getitem__(self, index): + return Int32(index) * Int32(self.seq_len_q) + + @cute.jit def _q_seq_bounds( cfg: Constexpr[FmhaDecodeConfig], - cu_seqlens_q: cute.Pointer | None, + cu_seqlens_q: cute.Pointer | UniformQueryOffsets | None, batch_idx: Int32, ) -> tuple[Int32, Int32]: """Return packed Q token offset/length, or fixed-SQ neutral bounds.""" + if cutlass.const_expr(isinstance(cu_seqlens_q, UniformQueryOffsets)): + length = Int32(cu_seqlens_q.seq_len_q) + return Int32(batch_idx) * length, length if cutlass.const_expr(cfg.use_variable_seqlens_q): q_begin = Int32(cu_seqlens_q[batch_idx]) q_end = Int32(cu_seqlens_q[batch_idx + Int32(1)]) @@ -612,7 +650,7 @@ def _q_group_token_base(cfg: Constexpr[FmhaDecodeConfig], q_group_idx: Int32) -> """Return the first Q token owned by a logical Q CTA.""" if cutlass.const_expr(cfg.uses_nontrivial_grouped_q_layout): return q_group_idx * Int32(cfg.q_tokens_per_cta) - if cutlass.const_expr(cfg.max_seq_len_q > 1): + if cutlass.const_expr(cfg.is_multi_token_q): head_ctas_per_token = Int32( (cfg.heads_q_per_kv + cfg.tile_size_q - 1) // cfg.tile_size_q ) @@ -638,7 +676,7 @@ def _q_tile_valid_rows_for_seq( remaining_tokens = cute.math.max(seq_len_q - token_base, Int32(0)) valid_tokens = cute.math.min(remaining_tokens, Int32(cfg.q_tokens_per_cta)) return valid_tokens * Int32(cfg.heads_q_per_kv) - if cutlass.const_expr(cfg.max_seq_len_q > 1): + if cutlass.const_expr(cfg.is_multi_token_q): head_ctas_per_token = Int32( (cfg.heads_q_per_kv + cfg.tile_size_q - 1) // cfg.tile_size_q ) @@ -707,7 +745,7 @@ def _q_physical_output_row_from_token_and_local_head( num_heads_q = num_heads_kv * heads_q_per_kv global_head_idx = kv_head_idx * heads_q_per_kv + local_head_idx return (q_token_offset + q_token_idx) * num_heads_q + global_head_idx - if cutlass.const_expr(cfg.max_seq_len_q > 1): + if cutlass.const_expr(cfg.is_multi_token_q): heads_q_per_kv = Int32(cfg.heads_q_per_kv) num_heads_q = num_heads_kv * heads_q_per_kv global_head_idx = kv_head_idx * heads_q_per_kv + local_head_idx @@ -815,7 +853,7 @@ def _q_row_token_and_local_head( q_group_idx * Int32(cfg.q_tokens_per_cta) + token_offset, local_head, ) - if cutlass.const_expr(cfg.max_seq_len_q > 1): + if cutlass.const_expr(cfg.is_multi_token_q): head_ctas_per_token = Int32( (cfg.heads_q_per_kv + cfg.tile_size_q - 1) // cfg.tile_size_q ) 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 896780914d1e..6f5e84ce9f86 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 @@ -82,6 +82,7 @@ def _runtime_configured_local_kv_tiles( seq_len_kv: Int32, seq_len_q: Int32, q_token_base: Int32, + requested_splits=None, ) -> Int32: """Return the instruction-aligned local span for configured split capacity.""" total_kv_tiles = _runtime_total_kv_tiles( @@ -91,9 +92,20 @@ def _runtime_configured_local_kv_tiles( q_token_base, ) num_insts_kv = Int32(cfg.num_insts_kv) - configured_splits = Int32(cfg.splits_kv if cfg.use_split_kv else 1) + configured_splits = ( + Int32(cfg.splits_kv if cfg.use_split_kv else 1) + if requested_splits is None + else Int32( + requested_splits[0] + if isinstance(requested_splits, tuple) + else requested_splits + ) + ) tiles_per_group = configured_splits * num_insts_kv - num_groups = (total_kv_tiles + tiles_per_group - Int32(1)) // tiles_per_group + divisor = ( + requested_splits[1] if isinstance(requested_splits, tuple) else tiles_per_group + ) + num_groups = (total_kv_tiles + tiles_per_group - Int32(1)) // divisor return cute.math.max(num_groups * num_insts_kv, num_insts_kv) @@ -219,6 +231,7 @@ def _runtime_active_splits_kv( seq_len_kv: Int32, seq_len_q: Int32, q_token_base: Int32, + requested_splits=None, ) -> Int32: """Return the useful split prefix for one runtime Q/KV work item.""" if cutlass.const_expr(not cfg.use_split_kv): @@ -230,26 +243,46 @@ def _runtime_active_splits_kv( q_token_base, ) local_kv_tiles = _runtime_configured_local_kv_tiles( - cfg, - seq_len_kv, - seq_len_q, - q_token_base, + cfg, seq_len_kv, seq_len_q, q_token_base, requested_splits=requested_splits ) # Split ranges are instruction-group aligned. Only ranges intersecting the # valid K domain participate; remaining configured grid slots are padding. + if cutlass.const_expr(cfg.splits_kv == 2): + # The configured span covers at most two nonempty ranges. Comparisons + # recover their count without a divide by the runtime local span. + return Int32(total_kv_tiles > Int32(0)) + Int32(total_kv_tiles > local_kv_tiles) return (total_kv_tiles + local_kv_tiles - Int32(1)) // local_kv_tiles +@cute.jit +def runtime_split_kv_is_active( + cfg: FmhaDecodeConfig, + seq_len_kv: Int32, + seq_len_q: Int32, + q_token_base: Int32, + split_idx: Int32, +) -> cutlass.Boolean: + """Test a split's first tile without dividing by its runtime local span.""" + total_kv_tiles = _runtime_total_kv_tiles(cfg, seq_len_kv, seq_len_q, q_token_base) + local_kv_tiles = _runtime_configured_local_kv_tiles( + cfg, seq_len_kv, seq_len_q, q_token_base + ) + return split_idx * local_kv_tiles < total_kv_tiles + + @cute.jit def _runtime_execution_splits_kv( cfg: FmhaDecodeConfig, seq_len_kv: Int32, seq_len_q: Int32, q_token_base: Int32, + requested_splits=None, ) -> Int32: """Return a nonzero producer/reduction fanout for a runtime-valid Q tile.""" return cute.math.max( - _runtime_active_splits_kv(cfg, seq_len_kv, seq_len_q, q_token_base), + _runtime_active_splits_kv( + cfg, seq_len_kv, seq_len_q, q_token_base, requested_splits=requested_splits + ), Int32(1), ) @@ -306,6 +339,8 @@ def _static_split_kv_global_tile_idx( @cute.jit def _logical_cta_kv_idx(cfg: FmhaDecodeConfig, stage_info: StageInfo) -> Int32: """Resolve the split index from the Q-group-major launch coordinate.""" + if cutlass.const_expr(cfg.uses_compact_split_grid): + return Int32(cute.arch.block_idx()[0]) if cutlass.const_expr(stage_info.work_tile is not None): q_group_cta_idx = Int32(stage_info.work_tile.tile_idx[0]) else: @@ -321,13 +356,11 @@ def _runtime_local_kv_tiles( seq_len_kv: Int32, seq_len_q: Int32, q_token_base: Int32, + requested_splits=None, ) -> Int32: """Return local KV tiles assigned to one split CTA at runtime.""" return _runtime_configured_local_kv_tiles( - cfg, - seq_len_kv, - seq_len_q, - q_token_base, + cfg, seq_len_kv, seq_len_q, q_token_base, requested_splits=requested_splits ) @@ -343,9 +376,23 @@ def _runtime_split_kv_global_tile_idx( """Map a runtime local tile index to a global split-KV tile index.""" if cutlass.const_expr(not cfg.use_split_kv): return local_tile_idx + if cutlass.const_expr(cfg.runtime_split_kv): + # get_domain has already resolved the split's absolute window origin. + # Consumers need only add their local tile, without dividing by the + # launch count again for every K/V or score tile. + return ( + Int32(_decode_gen_task_cache(stage_info)[_TASK_CACHE_KV_RAW_TILE_BASE]) + - _num_skipped_kv_tiles(cfg, seq_len_kv, seq_len_q, q_token_base) + + local_tile_idx + ) return ( _logical_cta_kv_idx(cfg, stage_info) - * _runtime_local_kv_tiles(cfg, seq_len_kv, seq_len_q, q_token_base) + * _runtime_local_kv_tiles( + cfg, + seq_len_kv, + seq_len_q, + q_token_base, + ) + local_tile_idx ) @@ -417,7 +464,9 @@ def resolve_keeps_tile_context( if cutlass.const_expr(cfg.uses_uniform_causal_mask): 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 + use_runtime_kv_domain = seqlens_kv is not None or ( + not cfg.fixed_seqlen or cfg.uses_runtime_q_kv_union + ) 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): diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_output.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_output.py index 24a7a4c8b29e..8979076fcaea 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_output.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_output.py @@ -265,7 +265,8 @@ def _copy_transposed_smem8b_to_gmem( headdim = cfg.headdim tile_size_q = cfg.tile_size_q num_bytes_per_smem_row = Int32(headdim if headdim <= 128 else 128) - num_copy_segments = max((tile_size_q * headdim + 2047) // 2048, 1) + # STSM covers full physical bands, including a partial logical head tail. + num_copy_segments = max((tile_size_q * cfg.swaps_output_head_dim + 2047) // 2048, 1) for copy_segment_idx in cutlass.range_constexpr(num_copy_segments): base_offset = warp_grp_thread_idx * Int32(16) + Int32(copy_segment_idx * 2048) smem_row_idx = base_offset // num_bytes_per_smem_row @@ -296,7 +297,8 @@ def _copy_transposed_smem8b_to_gmem( dst_row_idx, seq_len_q, ) - if cutlass.const_expr(full_tile_rows): + valid_output_row = valid_output_row & (dst_col_offset < headdim) + if cutlass.const_expr(full_tile_rows and cfg.swaps_output_head_dim == headdim): smem_src_full = ( smem_base_i32.subview((load_smem_offset >> Int32(2))) ).data_ptr() diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_softmax.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_softmax.py index bca4608eed8e..f24d7cb5ef31 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 @@ -25,6 +25,7 @@ from cutlass.experimental import primitives as prims +from ...separate_reduction import merge_log2_lse from ..fmha_decode_config import FmhaDecodeConfig from .helpers_common import ( Constexpr, @@ -451,3 +452,24 @@ def _attention_sink_for_scale_idx( num_heads_kv, local_head_idx, ) + + +@cute.jit +def _final_output_lse(cfg, sink_ptr, lse_ptr, mass_lse, output_row, head, writer): + """Merge a sink once in log space and publish unamplified final LSE. + + Input mass includes the internal FP8 P amplification. Partial statistics + remain in that domain; only the externally visible final store removes it. + """ + weight = Float32(1.0) + if cutlass.const_expr(cfg.use_attention_sinks or cfg.lse_mode != "none"): + value = mass_lse + if cutlass.const_expr(cfg.use_fp8_pv): + value -= Float32(cfg.fp8_p_quant_log2_scale) + if cutlass.const_expr(cfg.use_attention_sinks): + sink = (sink_ptr + cutlass.Int64(head)).load() * Float32(1.4426950408889634) + value, weight, _ = merge_log2_lse(value, sink) + if cutlass.const_expr(cfg.lse_mode != "none"): + if writer: + (lse_ptr + cutlass.Int64(output_row)).store(value) + return weight 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 0ee8f36eb909..ea835ed902d6 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 @@ -18,12 +18,12 @@ 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 +resolves the flat layout (``sage.config.sage_flat_scale_slot`` and +``sage.index.load_sage_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``). +Each softmax lane loads its row's ``sfQ`` word once per work tile with +``sage.index.load_sage_q_scale``, which the context kernel shares. :class:`SageKScalesResource` is one softmax instance's ring of ``sfK`` tiles, which the load warp produces after each K tile with ``cp.async`` @@ -47,7 +47,7 @@ import cutlass import cutlass.cute as cute -from cutlass import Float32, Int32, Int64 +from cutlass import Float32, Int32 from cutlass.experimental import primitives as prims from cutlass.experimental.task_scheduling.enums import WorkAttr from cutlass.experimental.task_scheduling.memory import ( @@ -66,7 +66,12 @@ _PREPARED_ROUTE_IS_PROXY_FLAG, _BlockSparseRouteLayout, ) -from ....sage import _FLOAT4_WORDS, _SAGE_ATOM_TOKENS, _SageKSection, flat_scale_slot +from ....sage.index import ( + _load_sage_scale, + sage_section_word, + sage_sequence_words, +) +from ....sage.layout import _FLOAT4_WORDS, _SAGE_ATOM_TOKENS, _SageKSection from ...stage import FmhaStage from ..fmha_decode_config import FmhaDecodeConfig from ..fmha_decode_constants import KV_KIND_K @@ -87,77 +92,6 @@ Constexpr = cutlass.Constexpr -@cute.jit -def _load_scale(scale_addr: Int64, index: Int32) -> Float32: - """Return one FP32 scale loaded from a global memory base address.""" - value_ptr = cutlass.inttoptr( - scale_addr + Int64(index) * Int64(4), - mem_space=1, - dtype=Float32, - ) - return Float32(value_ptr.load(count=1, alignment=4)[0]) - - -@cute.jit -def load_flat_scale( - scale_addr: Int64, - head_stride: Int32, - *, - head_idx: Int32, - batch_idx: Int32, - seq_len: Int32, - token_idx: Int32, - log2_block: Int32, -) -> Float32: - """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 sage_sequence_words( - cfg: Constexpr[FmhaDecodeConfig], - *, - batch_idx: Int32, - 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 ``sfQ`` of row ``row_idx`` of Q tile ``q_group_idx`` from the prepared scales. - - ``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``). - """ - 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_k_section(cfg: FmhaDecodeConfig, proxy: bool = False) -> _SageKSection: """Return the prepared K section one route kind reads. @@ -321,14 +255,14 @@ def copy_route_k_scales( ) if cutlass.const_expr(cfg.sage_mixed_k_geometry): if is_proxy_route: - copy(summaries, sequence_words + Int32(summaries.word_offset)) + copy(summaries, sage_section_word(sequence_words, summaries)) else: - copy(k_tokens, sequence_words + Int32(k_tokens.word_offset)) + copy(k_tokens, sage_section_word(sequence_words, k_tokens)) else: - section_words = sequence_words + Int32(k_tokens.word_offset) + section_words = sage_section_word(sequence_words, k_tokens) if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): if is_proxy_route: - section_words = sequence_words + Int32(summaries.word_offset) + section_words = sage_section_word(sequence_words, summaries) copy(k_tokens, section_words) @@ -375,12 +309,12 @@ def stage_v_channel_scales( assert cfg.headdim % num_threads == 0 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) + rows_base = sage_section_word(sequence_words, v_rows) for base in cutlass.range_constexpr(0, cfg.headdim, num_threads): channel = Int32(base) + thread_idx - staged[channel] = _load_scale(scales_addr, rows_base + channel) + staged[channel] = _load_sage_scale(scales_addr, rows_base + channel) if cutlass.const_expr(cfg.sage_v_mean): - staged[channel + Int32(cfg.headdim)] = _load_scale( + staged[channel + Int32(cfg.headdim)] = _load_sage_scale( scales_addr, rows_base + Int32(v_rows.mean_word_offset) + channel ) @@ -572,10 +506,10 @@ def _sequence_words(self, stage_info: StageInfo) -> Int32: stage_info, self.h_k_idx, self.b_idx ) return sage_sequence_words( - self.cfg, + self.cfg.sage_scale_layout(self.cfg.static_seq_len_kv), batch_idx=batch_idx, kv_head_idx=kv_head_idx, - num_heads_kv=self.num_heads_kv, + num_kv_heads=self.num_heads_kv, ) @producer_work @@ -650,7 +584,7 @@ def copy_tile( cfg, self.sage_scales_ptr, self._stage_words(stage_info), - self._sequence_words(stage_info) + Int32(k_tokens.word_offset), + sage_section_word(self._sequence_words(stage_info), k_tokens), k_tokens.atom_words, lambda atom: cute.math.min(tile_atom + atom, last_atom), self.tile_atoms, @@ -728,10 +662,10 @@ def stage_tile(self, stage_info: StageInfo) -> None: self.sage_scales_ptr, self.staged(stage_info.context), sequence_words=sage_sequence_words( - cfg, + cfg.sage_scale_layout(cfg.static_seq_len_kv), batch_idx=batch_idx, kv_head_idx=kv_head_idx, - num_heads_kv=self.num_heads_kv, + num_kv_heads=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_p.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_p.py index bb3448d40352..0ff00f6ee1dc 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 @@ -42,9 +42,8 @@ ) from ..fmha_decode_config import FmhaDecodeConfig -from ..fmha_decode_constants import ( - INT32_SCORE_BIAS, -) +from ....sage.int32_scores import INT32_SCORE_BIAS +from ....sage.modes import SageScoreScaling from ...placeholder_helpers import _placeholder_smem_array from .helpers_common import ( _route_is_proxy, @@ -461,11 +460,13 @@ def _compute_p_fragments_impl( # 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() + dequantized: Constexpr[bool] = ( + cfg.sage_score_scaling() is SageScoreScaling.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 + # (``FmhaDecodeConfig.sage_score_scaling``), 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 @@ -744,7 +745,7 @@ def _exponentiate_fragment_pairs( p1 = Float32(s_arr[value_idx + 1]) if cutlass.const_expr( _pair_uses_ex2_emulation( - pair_idx, pairs_per_fragment, self.cfg.ex2_emulated_pairs + pair_idx, pairs_per_fragment, self.cfg.fragment_ex2_emulated_pairs ) ): p0, p1 = _ex2_emulation_packed_f32x2(p0, p1) @@ -782,6 +783,7 @@ def _compute_keeps_p( new_max_arr: cutlass.Array, s_arr: cutlass.Array, route_is_proxy: cutlass.Boolean, + sage_q_scale: Float32 | None = None, ) -> None: """Materialize one complete-row Keeps probability tile. @@ -789,7 +791,9 @@ def _compute_keeps_p( gives paired lanes the low/high 64-column halves of one row. Each lane writes disjoint packed blocks into the TMEM or SMEM layout consumed by BMM2. Streamed profiles, including every block-sparse Keeps profile, - produce P through the rolled fragment loop instead. + produce P through the rolled fragment loop instead. A Sage row arrives + dequantized without ``sfQ`` (``TmemSResource._dequantize_sage_keeps_row``), + so its exponent multiplier is ``c * sfQ``. """ cfg = self.cfg # Every block-sparse Keeps profile streams P; only dense complete rows @@ -812,6 +816,9 @@ def _compute_keeps_p( ) exponent_addend = self._exponent_addend(new_max_arr[0], route_is_proxy) + exp_scale = self.scale_softmax_log2 + if cutlass.const_expr(cfg.use_sage_attention): + exp_scale = self.scale_softmax_log2 * sage_q_scale # Preserve four independent modulo-4 sum chains as two packed pairs, # without keeping a second 16-value P array live beside the S row. @@ -832,7 +839,7 @@ def _compute_keeps_p( s_arr[s_base + val_base], s_arr[s_base + val_base + 1], ), - (self.scale_softmax_log2, self.scale_softmax_log2), + (exp_scale, exp_scale), (exponent_addend, exponent_addend), ) p_pair_01 = ( @@ -845,7 +852,7 @@ def _compute_keeps_p( s_arr[s_base + val_base + 2], s_arr[s_base + val_base + 3], ), - (self.scale_softmax_log2, self.scale_softmax_log2), + (exp_scale, exp_scale), (exponent_addend, exponent_addend), ) p_pair_23 = ( @@ -869,7 +876,7 @@ def _compute_keeps_p( s_arr[s_base + val_base], s_arr[s_base + val_base + 1], ), - (self.scale_softmax_log2, self.scale_softmax_log2), + (exp_scale, exp_scale), (exponent_addend, exponent_addend), ) p_pair = ( @@ -1276,6 +1283,7 @@ def _compute_p_impl( ) if cutlass.const_expr( not self.use_variable_seqlens_kv + and cfg.fixed_seqlen and not cfg.uses_runtime_q_kv_union and not cfg.use_split_kv and cfg.has_odd_kv_tail @@ -1444,6 +1452,27 @@ def compute_p( route_is_proxy=cutlass.Boolean(False), ) + @producer_work + @cute.jit + def compute_sage_p( + self, + stage_info: StageInfo, + *, + new_max_arr: cutlass.Array, + s_arr: cutlass.Array, + sage_q_scale: Float32, + ) -> None: + """Compute a complete-row Sage P tile from the max pass's dequantized row.""" + + assert self.cfg.use_sage_attention and self.cfg.stages_q128_p_in_smem + self._compute_keeps_p( + stage_info, + new_max_arr=new_max_arr, + s_arr=s_arr, + route_is_proxy=cutlass.Boolean(False), + sage_q_scale=sage_q_scale, + ) + @producer_work @cute.jit def compute_proxy_route_p( 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 3e2031504ef7..919189399c6d 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 @@ -58,7 +58,6 @@ SWIZZLE_128B_ROW_BYTES, ) from ...stage import FmhaStage -from ...tensor_map import transform_ragged_coords from ...placeholder_helpers import ( _placeholder_local_array, _placeholder_smem_array, @@ -66,6 +65,7 @@ from .helpers_common import ( _TASK_CACHE_KV_PAGE_IDX_UB, _TASK_CACHE_KV_RAW_TILE_BASE, + _TASK_CACHE_KV_WINDOW_START, _TASK_CACHE_SEQ_LEN_KV, Constexpr, DecodeGenResourceBase, @@ -450,7 +450,9 @@ def _create_initial_task_locals( ) leading_byte_offset = q_tile_bytes stride_byte_offset = Int32( - _major_k_stride_bytes(self.cfg.q_dtype_bytes, self.cfg.headdim) + _major_k_stride_bytes( + self.cfg.q_dtype_bytes, self.cfg.smem_q_head_dim + ) ) for stage_idx in cutlass.range_constexpr(self.cfg.q_stages): # Advance the base address for each stage while preserving the @@ -502,7 +504,6 @@ def _tma_coords( if cutlass.const_expr(cfg.groups_tokens_heads_q): token_idx = logical_q_group_idx * Int32(cfg.q_tokens_per_cta) global_head_idx = logical_h_k_idx * Int32(cfg.heads_q_per_kv) - tokens_per_box = cfg.q_tokens_per_cta else: head_ctas_per_token = Int32( (cfg.heads_q_per_kv + cfg.tile_size_q - 1) // cfg.tile_size_q @@ -512,29 +513,10 @@ def _tma_coords( global_head_idx = logical_h_k_idx * Int32( cfg.heads_q_per_kv ) + head_cta_idx * Int32(cfg.tile_size_q) - tokens_per_box = 1 - if cutlass.const_expr(cfg.uses_q_token_kv_block_sparse_page_route): - return ( - head_dim_offset, - global_head_idx, - self.q_token_offset + token_idx, - ) - packed_coords = transform_ragged_coords( - ( - head_dim_offset, - global_head_idx, - self.q_token_offset + token_idx, - ), - ragged_dim_idx=2, - ragged_box_size=tokens_per_box, - ragged_extent=self.seq_len_q - token_idx, - ) return ( - packed_coords[0], - packed_coords[1], - packed_coords[2], - packed_coords[3], - packed_coords[4], + head_dim_offset, + global_head_idx, + self.q_token_offset + token_idx, ) if cutlass.const_expr(cfg.uses_nontrivial_grouped_q_layout): @@ -547,7 +529,7 @@ def _tma_coords( logical_b_idx, ) - if cutlass.const_expr(cfg.max_seq_len_q == 1): + if cutlass.const_expr(cfg.fixed_seqlen and cfg.max_seq_len_q == 1): return ( head_dim_offset, logical_q_group_idx * Int32(cfg.tile_size_q), @@ -576,7 +558,9 @@ def _complete_grouped_q_padding(self, stage_info: StageInfo) -> None: if cutlass.const_expr( cfg.groups_tokens_heads_q and cfg.q_manual_padding_rows > 0 ): - padding_bytes = cfg.q_manual_padding_rows * cfg.headdim * cfg.q_dtype_bytes + padding_bytes = ( + cfg.q_manual_padding_rows * cfg.smem_q_head_dim * cfg.q_dtype_bytes + ) prims.mbarrier_complete_tx(stage_info.barrier, Int32(padding_bytes)) @producer_work @@ -613,11 +597,11 @@ def tma_load(self, stage_info: StageInfo) -> None: stage_info.barrier, ) else: - # H256 FP8 stages Q by head-dim slices so each QK head-dim - # stage sees a contiguous SMEM tile. + # FP8 stages Q by head slices; TMA zero-fills a partial + # final slice while completing the full transaction. q_chunk_dim = 128 for q_chunk_idx in cutlass.range_constexpr( - cfg.headdim // q_chunk_dim + cfg.smem_q_head_dim // q_chunk_dim ): head_dim_offset = q_chunk_idx * q_chunk_dim smem_offset = head_dim_offset * cfg.tile_size_q @@ -647,7 +631,7 @@ def tma_load(self, stage_info: StageInfo) -> None: q_chunk_count = num_chunks issues_q_chunk = cutlass.Boolean(True) if cutlass.const_expr(cfg.load_num_warps > 1): - assert cfg.headdim in (64, 128, 256) + assert cfg.headdim in (64, 128, 192, 256) q_issuer_warps = min(num_chunks, cfg.load_num_warps) assert num_chunks % q_issuer_warps == 0 load_warp_rank = _load_task_warp_rank(cfg) @@ -939,7 +923,9 @@ def _maybe_runtime_tile_idx(self, stage_info: StageInfo, tile_idx: Int32) -> Int + tile_idx ) if cutlass.const_expr( - self.seqlens_kv is None and not self.cfg.uses_runtime_q_kv_union + self.seqlens_kv is None + and self.cfg.fixed_seqlen + and not self.cfg.uses_runtime_q_kv_union ): if cutlass.const_expr(self.cfg.use_split_kv): tile_idx = _static_split_kv_global_tile_idx( @@ -1763,6 +1749,10 @@ def page_ids(self, tile_idx: Int32, local_tile_idx: Int32) -> cutlass.Array: offset = self.consumer_work_stage * Int32( self.page_ids_per_stage ) + local_tile_idx * Int32(pages_per_tile) + elif cutlass.const_expr(cfg.runtime_page_window): + offset = self.consumer_work_stage * Int32(32) + local_tile_idx * Int32( + pages_per_tile + ) elif cutlass.const_expr( self.stage_page_ids_per_tile or self.stages_encoded_locators_per_tile ): @@ -1803,6 +1793,10 @@ def page_id_pointer( offset = self.consumer_work_stage * Int32( self.page_ids_per_stage ) + local_tile_idx * Int32(pages_per_tile) + elif cutlass.const_expr(cfg.runtime_page_window): + offset = self.consumer_work_stage * Int32(32) + local_tile_idx * Int32( + pages_per_tile + ) elif cutlass.const_expr( self.stage_page_ids_per_tile or self.stages_encoded_locators_per_tile ): @@ -1909,7 +1903,7 @@ def query_membership_keep_word( """Read a private, tail-padded query bit-plane without transposing it.""" cfg = self.cfg assert cfg.query_major_memberships and cfg.num_tokens_per_page == 4 - assert cfg.tile_size_kv == 128 and 2 <= cfg.max_seq_len_q <= 8 + assert cfg.tile_size_kv == 128 query = q_token_idx & Int32(7) word = Uint32(0) if cutlass.const_expr(self.cache_memberships_in_smem): @@ -2162,6 +2156,11 @@ def _producer_load_page_offsets( self.num_heads_kv ) + cutlass.Int64(logical_h_idx) pages_per_tile = Int32(cfg.tile_size_kv // cfg.num_tokens_per_page) + page_idx_lb = Int32(0) + if cutlass.const_expr(cfg.use_sliding_window_causal): + page_idx_lb = Int32( + _decode_gen_task_cache(stage_info)[_TASK_CACHE_KV_WINDOW_START] + ) // Int32(cfg.num_tokens_per_page) if cutlass.const_expr(self.use_native_paged_kv): task_cache = _decode_gen_task_cache(stage_info) page_idx_ub = Int32(task_cache[_TASK_CACHE_KV_PAGE_IDX_UB]) @@ -2207,6 +2206,9 @@ def _producer_load_page_offsets( Int32(task_cache[_TASK_CACHE_SEQ_LEN_KV]), self.seq_len_q, _q_group_token_base(cfg, runtime_q_group), + requested_splits=_decode_gen_task_cache(stage_info)[10] + if cfg.runtime_split_kv + else None, ) local_page_end = ( Int32(task_cache[_TASK_CACHE_KV_RAW_TILE_BASE]) + local_tiles @@ -2423,23 +2425,39 @@ def _producer_load_page_offsets( tile_idx * pages_per_tile + lane_idx, page_idx_ub ) smem_offset = stage_info.stage_idx * pages_per_tile + lane_idx - smem_page_offsets[smem_offset] = Int32( - page_idx_kv[page_table_offset + logical_page_idx] - ) + page_is_live = cutlass.Boolean(True) + if cutlass.const_expr(cfg.use_sliding_window_causal): + page_is_live = logical_page_idx >= page_idx_lb + page_id = Int32(-1) + if page_is_live: + page_id = Int32(page_idx_kv[page_table_offset + logical_page_idx]) + smem_page_offsets[smem_offset] = page_id else: # Shared-offset schedules use one coalesced warp load for an # aligned 32-ID window; consumers select their tile within it. grouped_base_page_idx = ((tile_idx * pages_per_tile) >> Int32(5)) << Int32( 5 ) + if cutlass.const_expr(cfg.runtime_page_window): + grouped_base_page_idx = ( + Int32(task_cache[_TASK_CACHE_KV_RAW_TILE_BASE]) * pages_per_tile + ) if page_warp_rank == Int32(0): grouped_logical_page_idx = cute.math.min( grouped_base_page_idx + lane_idx, page_idx_ub ) grouped_smem_offset = stage_info.stage_idx * Int32(32) + lane_idx - smem_page_offsets[grouped_smem_offset] = Int32( - page_idx_kv[page_table_offset + grouped_logical_page_idx] - ) + # Evicted pages can be reused under a different storage dtype. + # TensorMap zero-fill keeps masked NaNs out of the PV MMA. + page_is_live = cutlass.Boolean(True) + if cutlass.const_expr(cfg.use_sliding_window_causal): + page_is_live = grouped_logical_page_idx >= page_idx_lb + page_id = Int32(-1) + if page_is_live: + page_id = Int32( + page_idx_kv[page_table_offset + grouped_logical_page_idx] + ) + smem_page_offsets[grouped_smem_offset] = page_id @producer_work @cute.jit @@ -2778,7 +2796,9 @@ def _maybe_runtime_tile_idx(self, stage_info: StageInfo, tile_idx: Int32) -> Int # Apply runtime sequence-length, split-KV, and sliding-window # transforms to a local prefetch tile. if cutlass.const_expr( - self.seqlens_kv is None and not self.cfg.uses_runtime_q_kv_union + self.seqlens_kv is None + and self.cfg.fixed_seqlen + and not self.cfg.uses_runtime_q_kv_union ): if cutlass.const_expr(self.cfg.use_split_kv): tile_idx = _static_split_kv_global_tile_idx( diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_corr.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_corr.py index 9e46fae39bcc..9cfacbe37862 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 @@ -61,6 +61,7 @@ _q_physical_output_row, _q_physical_output_row_from_logical, _q_row_is_valid_for_seq, + _q_logical_output_row_is_valid_for_seq, _q_tile_output_row_base, _q_tile_valid_rows_for_seq, _masked_weighted_sum, @@ -83,8 +84,7 @@ _store_transposed_smem8b, ) from .helpers_softmax import ( - _attention_sink_for_local_head, - _attention_sink_for_scale_idx, + _final_output_lse, _pack_float4_to_fp8_e4m3, ) from .sage_scales import SageVScalesResource @@ -121,12 +121,15 @@ class TmemCorrResource(DecodeGenResourceBase): partial_stats_ptr: cute.Pointer = None split_kv_counter_ptr: cute.Pointer = None attention_sinks_ptr: cute.Pointer = None + lse_ptr: cute.Pointer | None = 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_ref: Constexpr[SageVScalesResource | None] = None seqlens_kv: cute.Pointer = None - max_seq_len_kv: Constexpr[int] = 0 + max_seq_len_kv: Int32 = 0 + # Launch bound in Q tokens; h_r separately carries partial-output row capacity. + max_q_len: Int32 = 1 seq_len_q: Int32 = None q_token_offset: Int32 = None num_heads_kv: Int32 = None @@ -187,7 +190,10 @@ def _init_placeholder_state(self) -> None: """Create placeholder SMEM views for correction and split reduction.""" o_stage_dtype_bytes = self.get_o_stage_dtype_bytes() o_entries = max( - self.cfg.tile_size_q * self.cfg.headdim * o_stage_dtype_bytes // 4, + self.cfg.tile_size_q + * self.cfg.swaps_output_head_dim + * o_stage_dtype_bytes + // 4, 1, ) # Keeps reduces each row's denominator in registers (plus the q64 @@ -251,6 +257,9 @@ def _runtime_splits_kv(self, stage_info: StageInfo) -> Int32: seq_len_kv, self.seq_len_q, _q_group_token_base(cfg, logical_q_group_idx), + requested_splits=_decode_gen_task_cache(stage_info)[10] + if cfg.runtime_split_kv + else None, ) def _can_alias_o_smem(self) -> bool: @@ -259,7 +268,9 @@ def _can_alias_o_smem(self) -> bool: # draining the final O staging path. Keep O staging separate from SmemP # there so the next tile's P producer cannot overwrite the copy source. o_stage_dtype_bytes = self.get_o_stage_dtype_bytes() - required_o_bytes = self.cfg.tile_size_q * self.cfg.headdim * o_stage_dtype_bytes + required_o_bytes = ( + self.cfg.tile_size_q * self.cfg.swaps_output_head_dim * o_stage_dtype_bytes + ) available_p_bytes = 2 * self.cfg.smem_p_tile_bytes return ( self.smem_p0_ref is not None @@ -280,7 +291,7 @@ def get_smem_requirements(self) -> list[SmemAllocation]: self._alloc = SmemAllocation( name=f"{self.name}", size_bytes=self.cfg.tile_size_q - * self.cfg.headdim + * self.cfg.swaps_output_head_dim * o_stage_dtype_bytes, alignment=self.cfg.stensor_align, ) @@ -381,20 +392,19 @@ def _split_reduction_row_is_owned( @cute.jit def _multi_cta_counter_q_groups(self, h_r: Int32) -> Int32: - """Return the number of split-KV counter groups in the output tile.""" + """Index counters by the launch Q bound, independent of partial-O capacity.""" cfg = self.cfg + seq_len_q = self.max_q_len if not cfg.fixed_seqlen else Int32(cfg.max_seq_len_q) if cutlass.const_expr(cfg.groups_tokens_heads_q): - return Int32( - (cfg.max_seq_len_q + cfg.q_tokens_per_cta - 1) // cfg.q_tokens_per_cta - ) + return Int32((seq_len_q + cfg.q_tokens_per_cta - 1) // cfg.q_tokens_per_cta) if cutlass.const_expr( - cfg.max_seq_len_q > 1 + cfg.is_multi_token_q and not cfg.groups_tokens_heads_q and cfg.heads_q_per_kv != 0 ): return Int32( ((cfg.heads_q_per_kv + cfg.tile_size_q - 1) // cfg.tile_size_q) - * cfg.max_seq_len_q + * seq_len_q ) return (h_r + Int32(cfg.tile_size_q - 1)) // Int32(cfg.tile_size_q) @@ -1109,6 +1119,7 @@ def _store_softmax_normalized_o_vec8( reduce_row_idx, sum_val, max_val, + lse_writer=reduce_col_idx == Int32(0), ) dst_offset = dst_row_base + reduce_col_idx * Int32(self.cfg.o_dtype_bytes) final_o_dst = cutlass.inttoptr( @@ -1124,6 +1135,43 @@ def _store_softmax_normalized_o_vec8( sector_aligned=cutlass.Boolean(False), ) + @cute.jit + def _final_row_weight(self, head_kv, batch, row, row_sum, row_max, writer): + cfg = self.cfg + weight = Float32(1.0) + if cutlass.const_expr(cfg.use_attention_sinks or cfg.lse_mode != "none"): + local_head = _local_head_from_q_output_row(cfg, self.h_r, row) + head_stride = _attention_sink_head_stride(cfg, self.h_r) + physical_row = _q_physical_output_row_from_logical( + cfg, + self.h_r, + self.num_heads_kv, + batch, + head_kv, + row, + self.q_token_offset, + ) + # This compact-row check excludes tail queries. Direct scale-group + # callers also mask CTA-local padding, which can alias the next tile. + valid = _q_logical_output_row_is_valid_for_seq( + cfg, self.h_r, row, self.seq_len_q + ) + mass_lse = Float32(-Float32.inf) + if row_sum > Float32(0.0): + mass_lse = self.scale_softmax_log2 * row_max + cute.math.log2( + row_sum, fastmath=True + ) + weight = _final_output_lse( + cfg, + self.attention_sinks_ptr, + self.lse_ptr, + mass_lse, + physical_row, + head_kv * head_stride + local_head, + writer and valid, + ) + return weight + @cute.jit def _softmax_output_row_state( self, @@ -1132,6 +1180,7 @@ def _softmax_output_row_state( reduce_row_idx: Int32, sum_val: Float32, max_val: Float32, + lse_writer: cutlass.Boolean = False, ) -> tuple[Int64, Float32]: """Resolve one logical row's output address and normalization once. @@ -1139,24 +1188,13 @@ def _softmax_output_row_state( (``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) - attention_sink_head_idx = _local_head_from_q_output_row( - cfg, self.h_r, reduce_row_idx - ) - sum_val += _attention_sink_for_local_head( - cfg, - self.attention_sinks_ptr, - self.scale_softmax_log2, - max_val, - logical_h_k_idx, - attention_sink_h_r, - self.num_heads_kv, - attention_sink_head_idx, + final_weight = self._final_row_weight( + logical_h_k_idx, logical_b_idx, reduce_row_idx, sum_val, max_val, lse_writer ) # ``output_scale`` is the public bmm2_scale. Fold it into the final # normalization for every output dtype; split partials reach this # helper only after the cross-CTA reduction has completed. - norm_scale = self.output_scale * self._safe_norm_rcp(sum_val) + norm_scale = self.output_scale * self._safe_norm_rcp(sum_val) * final_weight if cutlass.const_expr(cfg.use_fp8_pv and cfg.use_split_kv): # Fused split partials carry O divided by the FP8 P scale (see # ``_split_partial_scale``); restore it after the FP32 reduction. @@ -1247,7 +1285,10 @@ def create_function_variables( + (alias_offset if alias_offset is not None else self._alloc.offset), dtype=Int32, shape=( - self.cfg.tile_size_q * self.cfg.headdim * o_stage_dtype_bytes // 4, + self.cfg.tile_size_q + * self.cfg.swaps_output_head_dim + * o_stage_dtype_bytes + // 4, ), addrspace=3, ) @@ -1428,16 +1469,16 @@ def _reduce_and_store_gmem_split_o_vec8( partial_max = cutlass.Array(Float32, 4, space=cutlass.AddressSpace.rmem) partial_sum = cutlass.Array(Float32, 4, space=cutlass.AddressSpace.rmem) partial_regs = cutlass.Array(Int32, 16, space=cutlass.AddressSpace.rmem) - for split_base_i in cutlass.range_constexpr(0, cfg.max_splits_kv, 4): + for split_base_i in cutlass.range( + 0, reduction_splits_kv, 4, unroll_full=full_prefix + ): split_base = Int32(split_base_i) # Load up to four split partials first. Keeping the stats and O # loads grouped lets the fold loop below operate on registers only. for jj in cutlass.range_constexpr(4): if cutlass.const_expr(full_prefix): split_idx = Int32(split_base_i + jj) - valid_split_idx = cutlass.const_expr( - split_base_i + jj < cfg.splits_kv - ) + valid_split_idx = split_idx < Int32(cfg.splits_kv) else: split_idx = split_base + Int32(jj) valid_split_idx = split_idx < reduction_splits_kv @@ -1478,8 +1519,8 @@ def _reduce_and_store_gmem_split_o_vec8( for jj in cutlass.range_constexpr(4): valid_split_for_apply = split_base + Int32(jj) < reduction_splits_kv if cutlass.const_expr(full_prefix): - valid_split_for_apply = cutlass.const_expr( - split_base_i + jj < cfg.splits_kv + valid_split_for_apply = split_base + Int32(jj) < Int32( + cfg.splits_kv ) if valid_split_for_apply and valid_reduce_row: regs_base = jj * 4 @@ -1991,7 +2032,11 @@ def _copy_staged_fp16_o_segments( dst_row_idx, self.seq_len_q, ) - if valid_partial_row: + # Only padded STSM head bands need a runtime column check. + if valid_partial_row and ( + cfg.swaps_output_head_dim == cfg.headdim + or dst_col_offset < Int32(cfg.headdim * 2) + ): smem_src = self._smem_base_o_i32.subview( load_smem_offset >> Int32(2) ).data_ptr() @@ -2049,6 +2094,9 @@ def _copy_staged_fp16_o_segments( logical_q_group_idx, dst_row_idx, self.seq_len_q, + ) and ( + cfg.swaps_output_head_dim == cfg.headdim + or dst_col_offset < Int32(cfg.headdim * cfg.o_dtype_bytes) ): smem_src = self._smem_base_o_i32.subview( load_smem_offset >> Int32(2) @@ -2604,6 +2652,7 @@ def _kv_tile_256_output_row_state( logical_output_row_idx: Int32, denominator: Float32, global_max: Float32, + lse_writer: cutlass.Boolean, ) -> tuple[cutlass.Boolean, Int64, Float32]: """Return ``(valid_output_row, row_base, row_scale)`` of a lane's output row. @@ -2637,6 +2686,7 @@ def _kv_tile_256_output_row_state( logical_output_row_idx, denominator, global_max, + lse_writer=lse_writer and valid_output_row, ) return valid_output_row, row_base, row_scale @@ -2679,6 +2729,7 @@ def _kv_tile_256_merge_spatial_output( logical_output_row_idx = q_row_offset + exchange_row_idx valid_output_row, row_base, row_scale = self._kv_tile_256_output_row_state( resolves=output_lane, + lse_writer=output_lane, logical_h_k_idx=logical_h_k_idx, logical_b_idx=logical_b_idx, logical_kv_idx=logical_kv_idx, @@ -3013,6 +3064,7 @@ def _kv_tile_256_merge_owned_columns( logical_output_row_idx = q_row_offset + exchange_row_idx valid_output_row, row_base, row_scale = self._kv_tile_256_output_row_state( resolves=cutlass.Boolean(True), + lse_writer=owned_col_base == Int32(0), logical_h_k_idx=logical_h_k_idx, logical_b_idx=logical_b_idx, logical_kv_idx=logical_kv_idx, @@ -3153,6 +3205,7 @@ def _keeps_tail_epilogue( and cfg.tile_size_q == 128 and cfg.headdim == 128 and cfg.num_insts_kv == 2 + and cfg.fixed_seqlen and cfg.max_seq_len_q == 1 and cfg.has_static_dense_full_kv_tiles ): @@ -3223,32 +3276,42 @@ def _keeps_tail_epilogue( base_addr1 = base_addr0 partial_dst_col_offset = col_base * Int32(2) partial_norm_scale = self._split_partial_scale(reduced_sum_0) - regs_o_chunk = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) + # Retain the narrow fallback for low-level pointers whose base + # does not satisfy the public workspace's 32-byte alignment. + partial_misalignment = self.partial_o_ptr.toint() & Int64(31) + store_sectors = valid_output_row & (partial_misalignment == Int64(0)) + store_halves = valid_output_row & (partial_misalignment != Int64(0)) + chunk_cols = min(16, output_pair_regs * 2) + regs_o_chunk = cutlass.Array( + Int32, chunk_cols // 2, space=cutlass.AddressSpace.rmem + ) partial_o_row_base = self._gmem_partial_row_offset( logical_kv_idx, cta_idx_kv, logical_output_row_idx, ) * Int64(cfg.headdim * 2) - 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 # Load O0/O1 from the tail TMEM slots, combine with # max-correction scales, then store normalized 16-bit O for # standalone reduction or unnormalized O for fused reduction. 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) - for chunk_idx in cutlass.range_constexpr(4): + for chunk_idx in cutlass.range_constexpr(chunk_cols // 2): reg_base = chunk_idx * 2 if cutlass.const_expr(cfg.num_insts_kv == 1): partial_pair = fmul2( @@ -3292,15 +3355,25 @@ def _keeps_tail_epilogue( self.partial_o_ptr.toint() + partial_o_row_base + Int64(partial_dst_col_offset) - + Int64(store_idx * 16), + + Int64(store_idx * chunk_cols * 2), mem_space=1, dtype=Int32, ) - if valid_output_row: - partial_o_dst.store( - regs_o_chunk.data_ptr().load(count=4, alignment=4), - alignment=16, - ) + if cutlass.const_expr(chunk_cols == 16): + if store_sectors: + partial_o_dst.store( + regs_o_chunk.data_ptr().load(count=8, alignment=4), + alignment=32, + ) + if store_halves: + self._store_16bit_output_chunks( + partial_o_dst, regs_o_chunk, chunk_cols + ) + else: + if valid_output_row: + self._store_16bit_output_chunks( + partial_o_dst, regs_o_chunk, chunk_cols + ) if col_base == Int32(0) and valid_output_row: # One state record per output row accompanies all partial-O @@ -3359,24 +3432,20 @@ def _keeps_tail_epilogue( ) return - attention_sink_h_r = _attention_sink_head_stride(cfg, self.h_r) - attention_sink_head_idx = _local_head_from_q_output_row( - cfg, self.h_r, logical_output_row_idx - ) - reduced_sum_0 += _attention_sink_for_local_head( - cfg, - self.attention_sinks_ptr, - self.scale_softmax_log2, - final_max_0, + final_weight = self._final_row_weight( logical_h_k_idx, - attention_sink_h_r, - self.num_heads_kv, - attention_sink_head_idx, + logical_b_idx, + logical_output_row_idx, + reduced_sum_0, + final_max_0, + col_base == Int32(0) and valid_output_row, ) # Apply public bmm2_scale in the final direct-output normalization. # The split-KV branch above returns before reaching this point, so its # partial O remains unscaled for the selected reducer. - norm_scale_0 = self.output_scale * self._safe_norm_rcp(reduced_sum_0) + norm_scale_0 = ( + self.output_scale * self._safe_norm_rcp(reduced_sum_0) * final_weight + ) # Direct-output path: fold attention sinks into the denominator, then # bake normalization and per-instance max correction into the O scales. final_scale0_0 = norm_scale_0 * exp_scale0_0 @@ -3778,28 +3847,6 @@ def _swaps_wide_tail_epilogue( ) reduced_sum[scale_base] = reduced_pair[0] reduced_sum[scale_base + 1] = reduced_pair[1] - if cutlass.const_expr(not cfg.use_split_kv): - logical_h_k_idx, _ = _logical_head_batch( - stage_info, self.h_k_idx, self.b_idx - ) - logical_q_group_idx = _logical_q_group_idx( - cfg, stage_info, self.q_group_idx - ) - for scale_idx in cutlass.range_constexpr(num_scale_groups): - # Attention sinks add to the denominator before the - # final normalization scale is computed. - reduced_sum[scale_idx] += _attention_sink_for_scale_idx( - cfg, - self.attention_sinks_ptr, - self.scale_softmax_log2, - final_max[scale_idx], - logical_h_k_idx, - self.h_r, - self.num_heads_kv, - logical_q_group_idx, - col_group_idx, - scale_idx, - ) base_addr0 = self._swaps_o_stage_base_addr( tmem_row_base, o_base_col, tail_o_stage_idx_0 @@ -3975,11 +4022,43 @@ def _swaps_wide_tail_epilogue( ) return + logical_h_k_idx, logical_b_idx = _logical_head_batch( + stage_info, self.h_k_idx, self.b_idx + ) + logical_q_group_idx = _logical_q_group_idx(cfg, stage_info, self.q_group_idx) for scale_idx in cutlass.range_constexpr(num_scale_groups): # Direct final-output scales include both the reciprocal softmax # denominator and public bmm2_scale. Split partials returned above # use only the unnormalized ``exp_scale*`` arrays. - norm_scale = self.output_scale * self._safe_norm_rcp(reduced_sum[scale_idx]) + row_in_tile = ( + Int32((scale_idx // 2) * 8) + + col_group_idx * Int32(2) + + Int32(scale_idx % 2) + ) + output_row = _q_tile_output_row_base(cfg, logical_q_group_idx) + row_in_tile + final_weight = self._final_row_weight( + logical_h_k_idx, + logical_b_idx, + output_row, + reduced_sum[scale_idx], + final_max[scale_idx], + ( + warp_idx == Int32(0) + and lane_idx < Int32(4) + and _q_row_is_valid_for_seq( + cfg, + self.h_r, + logical_q_group_idx, + row_in_tile, + self.seq_len_q, + ) + ), + ) + norm_scale = ( + self.output_scale + * self._safe_norm_rcp(reduced_sum[scale_idx]) + * final_weight + ) final_scale0[scale_idx] = norm_scale * exp_scale0[scale_idx] final_scale1[scale_idx] = norm_scale * exp_scale1[scale_idx] @@ -4087,7 +4166,7 @@ def _swaps_wide_tail_epilogue( ), warp_grp_thread_idx, cfg.tile_size_q, - cfg.headdim, + cfg.swaps_output_head_dim, cfg.num_fp8_output_regs, ) cute.arch.fence_view_async_shared() @@ -4283,27 +4362,6 @@ def _swaps_q8_tail_epilogue( for scale_idx in cutlass.range_constexpr(num_scale_groups): reduced_sum[scale_idx] = reduced_pair[scale_idx] - if cutlass.const_expr(not cfg.use_split_kv): - logical_h_k_idx, _ = _logical_head_batch( - stage_info, self.h_k_idx, self.b_idx - ) - logical_q_group_idx = _logical_q_group_idx( - cfg, stage_info, self.q_group_idx - ) - for scale_idx in cutlass.range_constexpr(num_scale_groups): - reduced_sum[scale_idx] += _attention_sink_for_scale_idx( - cfg, - self.attention_sinks_ptr, - self.scale_softmax_log2, - final_max[scale_idx], - logical_h_k_idx, - self.h_r, - self.num_heads_kv, - logical_q_group_idx, - col_group_idx, - scale_idx, - ) - base_addr0 = self._swaps_o_stage_base_addr( tmem_row_base, o_base_col, tail_o_stage_idx_0 ) @@ -4314,7 +4372,7 @@ def _swaps_q8_tail_epilogue( ) output_pair_regs = cfg.num_fp16_output_regs output_f32_regs = output_pair_regs * 2 - num_o_chunks = cfg.headdim // 64 + num_o_chunks = cfg.swaps_output_head_dim // 64 if cutlass.const_expr(cfg.num_insts_kv == 1): o0_vals = self._swaps_load_o_stage_chunks( base_addr0, @@ -4482,11 +4540,43 @@ def _swaps_q8_tail_epilogue( ) return + logical_h_k_idx, logical_b_idx = _logical_head_batch( + stage_info, self.h_k_idx, self.b_idx + ) + logical_q_group_idx = _logical_q_group_idx(cfg, stage_info, self.q_group_idx) for scale_idx in cutlass.range_constexpr(num_scale_groups): # Direct final-output scales include both the reciprocal softmax # denominator and public bmm2_scale. Split partials returned above # use only the unnormalized ``exp_scale*`` arrays. - norm_scale = self.output_scale * self._safe_norm_rcp(reduced_sum[scale_idx]) + row_in_tile = ( + Int32((scale_idx // 2) * 8) + + col_group_idx * Int32(2) + + Int32(scale_idx % 2) + ) + output_row = _q_tile_output_row_base(cfg, logical_q_group_idx) + row_in_tile + final_weight = self._final_row_weight( + logical_h_k_idx, + logical_b_idx, + output_row, + reduced_sum[scale_idx], + final_max[scale_idx], + ( + warp_idx == Int32(0) + and lane_idx < Int32(4) + and _q_row_is_valid_for_seq( + cfg, + self.h_r, + logical_q_group_idx, + row_in_tile, + self.seq_len_q, + ) + ), + ) + norm_scale = ( + self.output_scale + * self._safe_norm_rcp(reduced_sum[scale_idx]) + * final_weight + ) final_scale0[scale_idx] = norm_scale * exp_scale0[scale_idx] final_scale1[scale_idx] = norm_scale * exp_scale1[scale_idx] @@ -4572,7 +4662,7 @@ def _swaps_q8_tail_epilogue( ), warp_grp_thread_idx, cfg.tile_size_q, - cfg.headdim, + cfg.swaps_output_head_dim, cfg.num_fp8_output_regs, ) cute.arch.fence_view_async_shared() @@ -4739,7 +4829,7 @@ def correction_loop_epilogue( num_scale_groups = cfg.num_softmax_scale_groups output_pair_regs = cfg.num_fp16_output_regs output_f32_regs = output_pair_regs * 2 - num_o_chunks = cfg.headdim // 64 + num_o_chunks = cfg.swaps_output_head_dim // 64 scale_vals = cutlass.Array( Float32, num_scale_groups, space=cutlass.AddressSpace.rmem ) @@ -4795,7 +4885,7 @@ def correction_loop_epilogue( num_scale_groups = cfg.num_softmax_scale_groups output_pair_regs = cfg.num_fp16_output_regs output_f32_regs = output_pair_regs * 2 - num_o_chunks = cfg.headdim // 64 + num_o_chunks = cfg.swaps_output_head_dim // 64 scale_vals = cutlass.Array( Float32, num_scale_groups, space=cutlass.AddressSpace.rmem ) @@ -4857,7 +4947,7 @@ def correction_loop_epilogue( ) output_pair_regs = cfg.num_fp16_output_regs output_f32_regs = output_pair_regs * 2 - num_o_chunks = cfg.headdim // 64 + num_o_chunks = cfg.swaps_output_head_dim // 64 self._swaps_rescale_o_stage_in_tmem( base_addr, scale_vals, @@ -4960,7 +5050,7 @@ def _correction_tail_epilogue_impl( num_scale_groups = cfg.num_softmax_scale_groups output_pair_regs = cfg.num_fp16_output_regs output_f32_regs = output_pair_regs * 2 - num_o_chunks = cfg.headdim // 64 + num_o_chunks = cfg.swaps_output_head_dim // 64 self._swaps_wide_tail_epilogue( stage_info, tail_o_stage_idx_0=tail_o_stage_idx_0, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_o.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_o.py index d08b770708f0..f2d546bddca2 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_o.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_o.py @@ -359,6 +359,10 @@ def _vp_mma( cfg ) + if cutlass.const_expr(cfg.use_keeps_mma_ab): + # SMEM keeps a complete staging band; MMA/TMEM own only live columns. + mma_n = min(mma_n, cfg.headdim - head_dim_stage_idx * cfg.head_dim_kv_stage) + if cutlass.const_expr(cfg.head_dim_per_stage_kv == 0): # Full-headDim path: all V columns for this O stage live in one # contiguous TMEM stage selected by the TS pipeline stage index. 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 73982aa0a912..61ac46bfd161 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 @@ -24,7 +24,7 @@ import cutlass import cutlass.cute as cute -from cutlass import BFloat16, Boolean, Float32, Int32, Uint32 +from cutlass import Boolean, Float32, Int32, Uint32 from cutlass.experimental import primitives as prims from cutlass.experimental.task_scheduling.enums import WorkAttr @@ -43,12 +43,12 @@ from ...._block_sparse.prepared import _PREPARED_ROUTE_IS_FULL_FLAG from ..fmha_decode_config import CAUSAL, FmhaDecodeConfig -from ..fmha_decode_constants import ( +from ....sage.int32_scores import ( INT32_SCORE_BIAS, - INT32_SCORE_SEED_MMA_K, - INT32_SCORE_SEED_TILE_LBO, - INT32_SCORE_SEED_TILE_SBO, - INT32_SCORE_SEED_TILE_WORDS, + fill_sage_score_seed_tile, + issue_sage_score_seed_mma, + sage_score_seed_tile_bytes, + sage_score_seed_words, ) from ...tcgen05_compat import tcgen05_ld_32x32b_max, tcgen05_mma_ws from ...placeholder_helpers import ( @@ -96,11 +96,11 @@ from .smem_block_sparse_metadata import ( _swaps_forwards_packed_route_full, ) +from ....sage.index import load_sage_q_scale, sage_sequence_words +from ....sage.modes import SageScoreScaling 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, @@ -352,22 +352,16 @@ def get_smem_requirements(self) -> list[SmemAllocation]: return allocations def _seed_mma_tile_bytes(self) -> int: - """Return the byte size of the seeding MMA's shared operand tile. - - A reads the first M rows and B the first N rows of the same tile. - """ + """Return the byte size of the seeding MMA's operand tile.""" _, mma_m, mma_n = _qk_mma_operand_contract_for_config(self.cfg) - return max(mma_m, mma_n) * INT32_SCORE_SEED_MMA_K * 2 + return sage_score_seed_tile_bytes(mma_m, mma_n) @cute.jit def _seed_mma_tile_words(self, context: ResourceContext) -> cutlass.Array: - """Return the seeding MMA's operand tile as a 32-bit SMEM array.""" + """Return the seeding MMA's operand tile; the second instance reads the first's.""" owner = self if self.score_seed_owner is None else self.score_seed_owner - return cutlass.Array( - context.smem_base.data_ptr() + owner._seed_alloc.offset, - dtype=Uint32, - shape=(owner._seed_mma_tile_bytes() // 4,), - addrspace=3, + return sage_score_seed_words( + context.smem_base, owner._seed_alloc.offset, owner._seed_mma_tile_bytes() ) @cute.jit @@ -385,26 +379,15 @@ def create_function_variables( @cute.jit def fill_score_seed_tiles(self, context: ResourceContext) -> None: - """Fill the seeding MMA's constant operand tile with the whole CTA. - - The fence makes the generic-proxy stores visible to the tensor core. - """ + """Fill the seeding MMA's constant operand tile with the whole CTA.""" assert self.score_seed_owner is None, "only the owning instance fills the tile" - words = self._seed_mma_tile_words(context) - num_words = self._seed_mma_tile_bytes() // 4 - num_threads = self.cfg.threads_per_cta tidx, _, _ = cute.arch.thread_idx() - even_word, odd_word = INT32_SCORE_SEED_TILE_WORDS - for round_idx in cutlass.range_constexpr( - (num_words + num_threads - 1) // num_threads - ): - word_idx = tidx + Int32(round_idx * num_threads) - if word_idx < Int32(num_words): - word = Uint32(even_word) - if (word_idx & Int32(1)) != Int32(0): - word = Uint32(odd_word) - words[word_idx] = word - cute.arch.fence_view_async_shared() + fill_sage_score_seed_tile( + self._seed_mma_tile_words(context), + self._seed_mma_tile_bytes() // 4, + tidx, + self.cfg.threads_per_cta, + ) def get_tmem_requirements(self) -> list[TmemAllocation]: """Allocate TMEM S score columns for QK MMA output.""" @@ -643,51 +626,20 @@ def _qk_tmem_col(self, stage_info: StageInfo, stage_slot_offset: Int32): @cute.jit def _seed_scores(self, stage_info: StageInfo, stage_slot_offset: Int32) -> None: - """Write ``INT32_SCORE_BIAS`` into the acquired S slot with one BF16 MMA step. - - The step reads the constant operand tile and no accumulator; the INT8 - K steps accumulate onto the seed's bit pattern in issue order. - """ + """Write ``INT32_SCORE_BIAS`` into the acquired S slot before its K steps.""" cfg = self.cfg assert cfg.uses_int32_scores # The bias binade relies on ``|score| <= 2**21``, which holds for # 128-wide INT8 dot products. assert cfg.headdim == 128 - if prims.elect_sync(): - tmem_col = self._qk_tmem_col(stage_info, stage_slot_offset) - _, mma_m, mma_n = _qk_mma_operand_contract_for_config(cfg) - tile_desc = prims.Tcgen05SmemDesc.build( - self._seed_mma_tile_words(stage_info.context), - leading_byte_offset=INT32_SCORE_SEED_TILE_LBO, - stride_byte_offset=INT32_SCORE_SEED_TILE_SBO, - layout=prims.Tcgen05SmemSwizzle.NONE, - ) - idesc = prims.Tcgen05InstrDesc.build( - c_dtype=Float32, - a_dtype=BFloat16, - b_dtype=BFloat16, - n_dim=mma_n, - m_dim=mma_m, - ) - if cutlass.const_expr(cfg.uses_ws_2x2_datapath): - tcgen05_mma_ws( - prims.Tcgen05MMAKind.F16, - tmem_col, - tile_desc, - tile_desc, - idesc, - False, - ) - else: - prims.tcgen05_mma( - prims.Tcgen05MMAKind.F16, - prims.CTAGroup.CTA_1, - tmem_col, - tile_desc, - tile_desc, - idesc, - False, - ) + _, mma_m, mma_n = _qk_mma_operand_contract_for_config(cfg) + issue_sage_score_seed_mma( + self._qk_tmem_col(stage_info, stage_slot_offset), + self._seed_mma_tile_words(stage_info.context), + mma_m, + mma_n, + ws=cfg.uses_ws_2x2_datapath, + ) @producer_work @cute.jit @@ -869,7 +821,10 @@ def _qk_mma( assert not cfg.uses_int32_scores, ( "staged head-dim BMM1 does not seed scores" ) - mma_k_steps = cfg.head_dim_kv_stage // _mma_k_step_qk(cfg) + mma_k_steps = min( + cfg.head_dim_kv_stage, + cfg.headdim - head_dim_stage_idx * cfg.head_dim_kv_stage, + ) // _mma_k_step_qk(cfg) if prims.elect_sync(): # Peel the first MMA so overwrite-vs-accumulate remains a # compile-time value rather than loop-carried state. @@ -947,6 +902,22 @@ def _resolve_keeps_tile_context(self, stage_info: StageInfo): q_group_idx=self.q_group_idx, ) + def _masked_score(self) -> Float32: + """Return the score a complete-row mask writes. + + A fully masked row must reach the P pass with the ``-FLT_MAX`` + maximum, which ``SmemPResource._exponent_addend`` anchors on zero so + that its masked scores exponentiate to zero. Most profiles write that + sentinel directly. Sage writes ``-inf``: its max pass dequantizes each + score with a positive ``sfK`` and scales the maximum by ``sfQ``, which + keeps ``-inf`` but can move ``-FLT_MAX`` off the sentinel, and it maps + an all ``-inf`` maximum back to the sentinel + (``_dequantize_sage_keeps_row``). + """ + if self.cfg.use_sage_attention: + return Float32(float("-inf")) + return _neg_max_f32() + @cute.jit def _reduce_keeps_row_max(self, s_vals: cutlass.Array) -> Float32: """Reduce one Keeps score row while preserving its lane ownership.""" @@ -985,6 +956,76 @@ def _reduce_keeps_row_max(self, s_vals: cutlass.Array) -> Float32: ) return tile_max + @cute.jit + def _dequantize_sage_keeps_row( + self, + stage_info: StageInfo, + s_vals: cutlass.Array, + sage_q_scale: Float32, + sage_k_stage_bytes: Int32, + ) -> Float32: + """Dequantize a complete Sage score row in place and return its tile max. + + Each K32 chunk reads its ``sfK`` words from the waited stage and + becomes ``(s - bias) * sfK_g`` (``_dequantize_fragment_scores``), so + the row reaches the P pass without scale words or bias. ``sfQ`` + applies once to the tile maximum; masked ``-inf`` scores stay + ``-inf``, and only the ``-inf`` identity maps to the finite sentinel. + """ + cfg = self.cfg + assert cfg.use_sage_attention and cfg.tile_size_q == 128 + chunk_regs = cfg.keeps_score_chunk_regs + groups: Constexpr[int] = cfg.sage_k_groups_per_fragment() + 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 + ) + max_chains = cutlass.Array(Float32, 4, space=cutlass.AddressSpace.rmem) + for chain_idx in cutlass.range_constexpr(4): + max_chains[chain_idx] = Float32(float("-inf")) + chunk = cutlass.Array(Float32, chunk_regs, space=cutlass.AddressSpace.rmem) + for chunk_idx in cutlass.range_constexpr( + cfg.num_s_regs_per_thread // chunk_regs + ): + first: Constexpr[int] = chunk_idx * chunk_regs + chunk_scales = load_sage_k_fragment(cfg, scales_view, Int32(chunk_idx)) + for score_idx in cutlass.range_constexpr(chunk_regs): + chunk[score_idx] = s_vals[first + score_idx] + self._dequantize_fragment_scores(chunk, chunk_scales, groups) + for score_idx in cutlass.range_constexpr(chunk_regs): + s_vals[first + score_idx] = chunk[score_idx] + chain_idx: Constexpr[int] = score_idx % 4 + max_chains[chain_idx] = cute.math.max( + max_chains[chain_idx], Float32(chunk[score_idx]), ftz=True + ) + tile_max = cute.math.max( + cute.math.max(max_chains[0], max_chains[1], ftz=True), + cute.math.max(max_chains[2], max_chains[3], ftz=True), + ftz=True, + ) + if tile_max == Float32(float("-inf")): + tile_max = _neg_max_f32() + else: + tile_max = tile_max * sage_q_scale + return tile_max + + @cute.jit + def _reduce_keeps_tile_max( + self, + stage_info: StageInfo, + s_vals: cutlass.Array, + sage_q_scale: Float32 | None, + sage_k_stage_bytes: Int32 | None, + ) -> Float32: + """Return a complete row's tile maximum, dequantizing a Sage row first.""" + if cutlass.const_expr(self.cfg.use_sage_attention): + return self._dequantize_sage_keeps_row( + stage_info, s_vals, sage_q_scale, sage_k_stage_bytes + ) + return self._reduce_keeps_row_max(s_vals) + @cute.jit def _publish_keeps_softmax_state( self, @@ -1194,7 +1235,7 @@ def _load_keeps_fragment_impl( and not is_masked_final_wave ): for reg_idx in cutlass.range_constexpr(num_s_regs): - s_vals[reg_idx] = _neg_max_f32() + s_vals[reg_idx] = self._masked_score() # A per-row causal endpoint is always <= seq_len_kv and is # applied by the loop below. Avoid emitting a second, @@ -1208,10 +1249,10 @@ def _load_keeps_fragment_impl( col_base, ) if token_idx >= element_mask_end_idx: - s_vals[reg_idx] = _neg_max_f32() + s_vals[reg_idx] = self._masked_score() if cutlass.const_expr(cfg.use_sliding_window_causal): if token_idx < window_start_idx: - s_vals[reg_idx] = _neg_max_f32() + s_vals[reg_idx] = self._masked_score() if cutlass.const_expr(cfg.uses_per_row_causal_mask): q_token_idx, _ = _q_row_token_and_local_head( @@ -1265,7 +1306,7 @@ def _load_keeps_fragment_impl( ) reg = local_page_idx * page_span + token_in_span s_vals[reg] = cutlass.select_( - valid, s_vals[reg], _neg_max_f32() + valid, s_vals[reg], self._masked_score() ) else: for reg_idx in cutlass.range_constexpr(num_s_regs): @@ -1277,9 +1318,9 @@ def _load_keeps_fragment_impl( ) if cutlass.const_expr(cfg.use_sliding_window_causal): if score_col < causal_start_rel: - s_vals[reg_idx] = _neg_max_f32() + s_vals[reg_idx] = self._masked_score() if score_col >= causal_end_rel: - s_vals[reg_idx] = _neg_max_f32() + s_vals[reg_idx] = self._masked_score() if cutlass.const_expr(cfg.uses_q_token_kv_block_sparse_page_membership): # Q64 lanes own a contiguous 64-column half; Q128 lanes own the @@ -1297,7 +1338,7 @@ def _load_keeps_fragment_impl( s_vals[membership_reg_idx] = cutlass.select_( page_is_member, s_vals[membership_reg_idx], - _neg_max_f32(), + self._masked_score(), ) if cutlass.const_expr(cfg.q_score_rows_need_mask): @@ -1309,7 +1350,7 @@ def _load_keeps_fragment_impl( self.seq_len_q, ): for invalid_q_reg_idx in cutlass.range_constexpr(num_s_regs): - s_vals[invalid_q_reg_idx] = _neg_max_f32() + s_vals[invalid_q_reg_idx] = self._masked_score() @cute.jit def _load_keeps_fragment( @@ -1383,6 +1424,8 @@ def _compute_softmax_loop_keeps( sum_arr: cutlass.Array, new_max_arr: cutlass.Array, s_arr: cutlass.Array, + sage_q_scale: Float32 | None = None, + sage_k_stage_bytes: Int32 | None = None, ) -> tuple[object, object, object, object]: """Load and reduce the row-major Keeps S fragment. @@ -1390,6 +1433,8 @@ def _compute_softmax_loop_keeps( assigns one row to a lane pair: lanes ``xor 16`` own the low/high 64-column halves. This path deliberately avoids the Swaps scratch reduction, whose 16x256b register mapping is unrelated to Keeps. + A Sage row is dequantized in registers before the maximum + (``_dequantize_sage_keeps_row``). """ cfg = self.cfg if cutlass.const_expr(cfg.streams_tmem_p_fragments): @@ -1402,25 +1447,13 @@ def _compute_softmax_loop_keeps( new_max_arr=new_max_arr, s_arr=s_arr, use_sparse=False, + sage_q_scale=sage_q_scale, + sage_k_stage_bytes=sage_k_stage_bytes, ) - task_cache = _decode_gen_task_cache(stage_info) num_s_regs = cfg.num_s_regs_per_thread old_max = new_max_arr[0] running_sum = sum_arr[0] s_vals = cutlass.Array(Float32, num_s_regs, space=cutlass.AddressSpace.rmem) - use_runtime_paged_dense_load = ( - self.seqlens_kv is not None - and cfg.use_paged_kv - and not cfg.use_sliding_window_causal - and cfg.tile_size_q in (64, 128) - ) - use_preload_mask_split = ( - use_runtime_paged_dense_load or cfg.uses_per_row_causal_mask - ) - if cutlass.const_expr(not use_preload_mask_split): - for reg_idx in cutlass.range_constexpr(num_s_regs): - s_vals[reg_idx] = _neg_max_f32() - ( seq_len_kv, logical_q_group_idx, @@ -1430,110 +1463,29 @@ def _compute_softmax_loop_keeps( is_valid_effective_tile, is_masked_final_wave, tile_is_unmasked, - tile_has_valid_scores, + _tile_has_valid_scores, ) = self._resolve_keeps_tile_context(stage_info) - if cutlass.const_expr(use_preload_mask_split): - # Select the complete unmasked/masked TMEM load+max path before any S - # registers are materialized. The shared predicate covers the - # intersection of all active grouped-Q causal/window intervals. - tile_max = _neg_max_f32() - self._load_keeps_fragment( - stage_info, - s_vals, - tile_offset_k, - element_mask_end_idx, - window_start_idx, - seq_len_kv, - logical_q_group_idx, - is_valid_effective_tile, - is_masked_final_wave, - tile_is_unmasked, - ) - tile_max = self._reduce_keeps_row_max(s_vals) - - self._publish_keeps_softmax_state( - s_vals, - tile_max, - old_max, - running_sum, - old_max_arr, - sum_arr, - new_max_arr, - s_arr, - ) - return old_max_arr, sum_arr, new_max_arr, s_arr - - if tile_has_valid_scores: - 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 - # result. - load_atom_regs = 32 - for load_atom_idx in cutlass.range_constexpr(num_s_regs // load_atom_regs): - atom_col = load_atom_idx * load_atom_regs - loaded = _keeps_tcgen05_ld( - cfg, - prims.make_tmem_ptr(base_addr + Int32(atom_col), Float32), - num=load_atom_regs, - offset=cfg.tile_size_kv // 2, - ) - prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) - for atom_reg_idx in cutlass.range_constexpr(load_atom_regs): - s_vals[atom_col + atom_reg_idx] = loaded[atom_reg_idx] - - warp_grp_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] - lane_idx = task_cache[_TASK_CACHE_LANE_IDX] - tile_row_idx = _keeps_row_idx(cfg, warp_grp_thread_idx) - col_base = _keeps_col_base(cfg, lane_idx, num_s_regs) - - # Tail/uniform-causal/window masks use each register's true logical K - # column. For q64, the paired lane owns the complementary 64-column half. - for reg_idx in cutlass.range_constexpr(num_s_regs): - token_idx = tile_offset_k + _keeps_score_col( - cfg, warp_grp_thread_idx, reg_idx, col_base - ) - # The per-row causal pass below subsumes seq_len_kv's upper bound. - if cutlass.const_expr(not cfg.uses_per_row_causal_mask): - if token_idx >= element_mask_end_idx: - s_vals[reg_idx] = _neg_max_f32() - if cutlass.const_expr(cfg.use_sliding_window_causal): - if token_idx < window_start_idx: - s_vals[reg_idx] = _neg_max_f32() - - if cutlass.const_expr(cfg.uses_per_row_causal_mask): - q_token_idx, _ = _q_row_token_and_local_head( - cfg, - self.h_r, - logical_q_group_idx, - tile_row_idx, - ) - causal_end = seq_len_kv - self.seq_len_q + q_token_idx + Int32(1) - causal_start = Int32(0) - if cutlass.const_expr(cfg.use_sliding_window_causal): - causal_start = cute.math.max( - causal_end - Int32(cfg.attention_window_size), Int32(0) - ) - for reg_idx in cutlass.range_constexpr(num_s_regs): - token_idx = tile_offset_k + _keeps_score_col( - cfg, warp_grp_thread_idx, reg_idx, col_base - ) - if token_idx < causal_start or token_idx >= causal_end: - s_vals[reg_idx] = _neg_max_f32() - - if cutlass.const_expr(cfg.q_score_rows_need_mask): - if not _q_row_is_valid_for_seq( - cfg, - self.h_r, - logical_q_group_idx, - tile_row_idx, - self.seq_len_q, - ): - for reg_idx in cutlass.range_constexpr(num_s_regs): - s_vals[reg_idx] = _neg_max_f32() - - tile_max = self._reduce_keeps_row_max(s_vals) + # Select the complete unmasked/masked TMEM load+max path before any S + # registers are materialized: fully visible tiles then skip the mask + # pass, and no loaded score crosses a post-load branch. The shared + # predicate covers the intersection of all active grouped-Q + # causal/window intervals. + self._load_keeps_fragment( + stage_info, + s_vals, + tile_offset_k, + element_mask_end_idx, + window_start_idx, + seq_len_kv, + logical_q_group_idx, + is_valid_effective_tile, + is_masked_final_wave, + tile_is_unmasked, + ) + tile_max = self._reduce_keeps_tile_max( + stage_info, s_vals, sage_q_scale, sage_k_stage_bytes + ) self._publish_keeps_softmax_state( s_vals, tile_max, @@ -1654,6 +1606,7 @@ def _compute_softmax_loop_swaps( self.seqlens_kv is not None and cfg.use_paged_kv and not cfg.use_split_kv + and cfg.fixed_seqlen and cfg.max_seq_len_q == 1 and not cfg.use_sliding_window_causal ) @@ -1685,8 +1638,8 @@ def _compute_softmax_loop_swaps( element_mask_end_idx = seq_len_kv - self.seq_len_q + q_token_base + Int32(1) should_load_s = True if cutlass.const_expr(not use_sparse): - use_runtime_kv_domain = ( - self.seqlens_kv is not None or cfg.uses_runtime_q_kv_union + use_runtime_kv_domain = self.seqlens_kv is not None or ( + not cfg.fixed_seqlen or cfg.uses_runtime_q_kv_union ) local_tile_idx = _softmax_tile_idx(cfg, stage_info, self.inst_id) if cutlass.const_expr(not use_runtime_kv_domain): @@ -2067,7 +2020,7 @@ def _compute_softmax_loop_swaps( proxy_max_shift, tail_shift = self._proxy_score_shifts( route_is_proxy_uniform ) - if cutlass.const_expr(tail_log2_delta != 0.0): + if cutlass.const_expr(not cfg.fixed_seqlen or tail_log2_delta != 0.0): if route_is_proxy_uniform: for token_group_idx in cutlass.range_constexpr(4): atom_origin, logical_summary = _swaps_routed_coordinate( @@ -2268,16 +2221,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 ) - return load_q_scale( - cfg, + layout = cfg.sage_scale_layout(cfg.static_seq_len_kv) + return load_sage_q_scale( + layout, self.sage_scales_ptr, sequence_words=sage_sequence_words( - cfg, + layout, batch_idx=batch_idx, kv_head_idx=kv_head_idx, - num_heads_kv=self.num_heads_kv, + num_kv_heads=self.num_heads_kv, ), - q_group_idx=_logical_q_group_idx(cfg, stage_info, self.q_group_idx), + q_tile_idx=_logical_q_group_idx(cfg, stage_info, self.q_group_idx), row_idx=_keeps_row_idx(cfg, warp_grp_thread_idx), ) @@ -2297,17 +2251,18 @@ def compute_sage_softmax_loop( """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``. + ``SageKScalesResource.open_stage``. Streamed profiles fold the + fragments in TMEM; a complete-row SMEM-P tile returns its row + dequantized in ``s_arr``. """ assert self.cfg.use_sage_attention - return self._compute_softmax_loop_keeps_fragments( + return self._compute_softmax_loop_keeps( stage_info, old_max_arr=old_max_arr, sum_arr=sum_arr, new_max_arr=new_max_arr, s_arr=s_arr, - use_sparse=False, sage_q_scale=sage_q_scale, sage_k_stage_bytes=sage_k_stage_bytes, ) @@ -2571,9 +2526,11 @@ def _compute_softmax_loop_keeps_fragments( proxy_max_shift = Float32(0.0) tail_shift = Float32(0.0) tail_fragment_mask = Int32(0) - tail_lane: Constexpr[int] = 0 + tail_lane = Int32(0) shifts_tail: Constexpr[bool] = False - dequantizes_scores: Constexpr[bool] = cfg.sage_scores_dequantized() + dequantizes_scores: Constexpr[bool] = ( + cfg.sage_score_scaling() is SageScoreScaling.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 @@ -2590,7 +2547,7 @@ def _compute_softmax_loop_keeps_fragments( ) tail_summary_idx, tail_log2_delta = cfg.proxy_tail_summary tail_lane = tail_summary_idx % fragment_regs - shifts_tail = tail_log2_delta != 0.0 + shifts_tail = not cfg.fixed_seqlen or tail_log2_delta != 0.0 proxy_max_shift, tail_shift = self._proxy_score_shifts(route_is_proxy) if cutlass.const_expr(shifts_tail and cfg.use_sage_attention): # Sage scores are quantized: the shift is divided by the row's @@ -2598,17 +2555,17 @@ def _compute_softmax_loop_keeps_fragments( tail_shift = tail_shift * cute.math.rcp(sage_q_scale, approx=True) if cutlass.const_expr(shifts_tail): # One mask bit per fragment marks the one holding the ragged - # final summary, whose half and fragment are compile-time - # (``proxy_static_tail_fragment``). A thread holds it on a + # final summary. Its half and fragment follow the live KV + # length. A thread holds it on a # proxy route of the last summary group, in the tail's half; # with several groups, ``origin0`` identifies the last one. - tail_half, tail_fragment = cfg.proxy_static_tail_fragment + tail_half, tail_fragment = cfg.proxy_tail_fragment holds_tail = cutlass.Boolean( route_is_proxy and _keeps_spatial_half(cfg, warp_grp_thread_idx) == Int32(tail_half) ) - if cutlass.const_expr(cfg.num_proxy_groups > 1): + if cfg.num_proxy_groups > 1: # The half's first atom is atom ``tail_half`` of the group. tail_group_origin0 = ( cfg.num_proxy_groups - 1 @@ -2835,7 +2792,7 @@ def _mask_score_fragments( may_hold_tail: Constexpr[bool], tail_fragment_mask: Int32, tail_shift: Float32, - tail_lane: Constexpr[int], + tail_lane: Int32, ) -> None: """Mask, fold and write back every fragment in one rolled loop. @@ -2855,7 +2812,9 @@ def _mask_score_fragments( num_fragments = cfg.num_softmax_score_fragments fragment_regs = cfg.softmax_score_fragment_regs groups: Constexpr[int] = cfg.sage_k_groups_per_fragment(proxy_kind) - dequantize_scores: Constexpr[bool] = cfg.sage_scores_dequantized(proxy_kind) + dequantize_scores: Constexpr[bool] = ( + cfg.sage_score_scaling(proxy_kind) is SageScoreScaling.DEQUANTIZED + ) for fragment in cutlass.range(num_fragments, unroll=1): fragment_scales = None if cutlass.const_expr(cfg.use_sage_attention): @@ -2884,11 +2843,13 @@ def _mask_score_fragments( score = Float32(float("-inf")) else: score = _neg_max_f32() - if cutlass.const_expr(may_hold_tail and score_idx == tail_lane): - if ((tail_fragment_mask >> Int32(fragment)) & Int32(1)) != Int32(0): + if cutlass.const_expr(may_hold_tail): + if score_idx == tail_lane and ( + (tail_fragment_mask >> Int32(fragment)) & Int32(1) + ) != Int32(0): lane_shift = tail_shift if cutlass.const_expr(cfg.use_sage_attention): - tail_group: Constexpr[int] = tail_lane // ( + tail_group: Constexpr[int] = score_idx // ( fragment_regs // groups ) lane_shift = tail_shift * cute.math.rcp( @@ -2940,7 +2901,7 @@ def _dequantize_fragment_scores( masked ``-inf`` score stays ``-inf`` under any positive scale. """ cfg = self.cfg - fragment_regs = cfg.softmax_score_fragment_regs + fragment_regs = cfg.keeps_score_chunk_regs group_regs = fragment_regs // groups group_addends = None if cutlass.const_expr(cfg.uses_int32_scores): 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 7d2393db28f0..2edea7357a2b 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 @@ -51,9 +51,11 @@ ) from cutlass.experimental.task_scheduling.task import Task +from ...sage.modes import SageScoreScaling from ..stage import FmhaStage from ..task_compat import task_context_kwargs from .direct_sparse_metadata import DirectSparseMetadataView +from .direct_sparse_metadata import HeadIndexedMetadataView from .fmha_decode_config import FmhaDecodeConfig from .fmha_decode_constants import ( KV_INST0, @@ -61,7 +63,11 @@ KV_KIND_K, KV_KIND_V, ) +from cutlass.utils.static_persistent_tile_scheduler import WorkTileInfo + from .fmha_decode_resources.helpers_common import ( + split_kv_grid_coord, + UniformQueryOffsets, _assume_nonnegative_i32, _q_group_token_base, _q_seq_bounds, @@ -69,6 +75,7 @@ ) from .fmha_decode_resources.helpers_kv_tile_idx import ( _runtime_last_valid_page_idx, + _runtime_configured_local_kv_tiles, _runtime_total_kv_tiles, _sliding_window_start_idx, ) @@ -169,35 +176,55 @@ def consume_schedule_token(self, stage_info: StageInfo) -> None: @dataclass(kw_only=True) class PackedDecodeWorkQueue(WorkQueue): - """CLC work queue that drops packed-Q tiles beyond a batch's Q length.""" + """CLC work queue that drops Q padding and inactive KV requests.""" cfg: cutlass.Constexpr[FmhaDecodeConfig] = field(init=False, default=None) cu_seqlens_q: Any = field(init=False, default=None) + seqlens_kv: Any = field(init=False, default=None) + num_heads_kv: Any = field(init=False, default=None) def __init__( self, cfg: FmhaDecodeConfig, - cu_seqlens_q: cute.Pointer, + cu_seqlens_q: cute.Pointer | UniformQueryOffsets, + seqlens_kv: cute.Pointer | None = None, + num_heads_kv: cutlass.Int32 | None = None, **kwargs: Any, ) -> None: - """Attach packed-Q metadata to the shared persistent work queue.""" + """Attach request metadata to the shared persistent work queue.""" super().__init__(**kwargs) self.cfg = cfg self.cu_seqlens_q = cu_seqlens_q + self.seqlens_kv = seqlens_kv + self.num_heads_kv = num_heads_kv @cute.jit def skip_work_tile_if(self, work_tile: Any) -> cutlass.Boolean: - """Skip a fetched Q group when its first token is outside this batch.""" - q_group_cta_idx, _, b_idx = work_tile.tile_idx + """Skip empty KV requests and Q groups outside their live sequence.""" + q_group_cta_idx, h_idx, b_idx = work_tile.tile_idx + empty_kv = cutlass.Boolean(False) + if cutlass.const_expr(self.seqlens_kv is not None): + lengths = self.seqlens_kv + if cutlass.const_expr( + self.cfg.use_q_token_kv_block_sparse_route + and not self.cfg.shares_sparse_pattern + ): + if cutlass.const_expr(isinstance(lengths, DirectSparseMetadataView)): + lengths = lengths.with_head(h_idx) + else: + lengths = HeadIndexedMetadataView(lengths, self.num_heads_kv, h_idx) + empty_kv = cutlass.Int32(lengths[b_idx]) <= cutlass.Int32(0) q_group_idx = cutlass.Int32(q_group_cta_idx) - if cutlass.const_expr(self.cfg.use_split_kv): + if cutlass.const_expr(self.cfg.uses_compact_split_grid): + q_group_idx = cutlass.Int32(cute.arch.block_idx()[2]) % self.cfg.num_q_ctas + elif cutlass.const_expr(self.cfg.use_split_kv): q_group_idx = q_group_idx // cutlass.Int32(self.cfg.splits_kv) _, seq_len_q = _q_seq_bounds( self.cfg, self.cu_seqlens_q, cutlass.Int32(b_idx), ) - return _q_group_token_base(self.cfg, q_group_idx) >= seq_len_q + return empty_kv or _q_group_token_base(self.cfg, q_group_idx) >= seq_len_q def _page_offsets_consume( @@ -236,13 +263,21 @@ def _can_hold_native_page_window( cfg: FmhaDecodeConfig, smem_page_offsets: MemoryResource | None, ) -> bool: - """Return whether one native page-ID stage covers every runtime KV range.""" + """Return whether one native page-ID stage covers every runtime KV range. + + A runtime launch selects the CTA-relative window only when its host KV + bound and count guarantee the entire instruction-aligned span fits. + Longer spans retain the ordinary staged pipeline. + """ if ( smem_page_offsets is None or not smem_page_offsets.use_native_paged_kv or cfg.use_sliding_window_causal + or not cfg.fixed_seqlen ): return False + if cfg.runtime_split_kv: + return cfg.runtime_page_window pages_per_tile = cfg.tile_size_kv // cfg.num_tokens_per_page if cfg.uses_scattered_page_route: # Each sparse route gets a page-offset stage sized for its complete @@ -514,16 +549,10 @@ def _decode_work_tile_schedule( body: Callable[[], None], non_skippable_head: Callable[[], None] | None = None, ) -> None: - """Trace one worker schedule, making only packed persistent data skippable.""" + """Trace data work separately from persistent queue bookkeeping.""" if cfg.use_persistent_scheduler: assert work_queue is not None - if not cfg.use_variable_seqlens_q: - with work_tile_loop(work_queue): - if non_skippable_head is not None: - non_skippable_head() - body() - _work_queue_tail(work_queue) - return + # Schedule tracing passes resource proxies, so use the static config. with work_tile_loop( work_queue, skip_if=PackedDecodeWorkQueue.skip_work_tile_if, @@ -562,13 +591,6 @@ def _decode_work_tile_schedule_with_invariant_bridge( # Descriptor backing storage depends only on the task's SMEM # allocation, so initialize it once before the persistent loop. invariant_setup() - if not cfg.use_variable_seqlens_q: - with work_tile_loop(work_queue): - active_prelude() - invariant = bridge() - body(invariant) - _work_queue_tail(work_queue) - return with work_tile_loop( work_queue, skip_if=PackedDecodeWorkQueue.skip_work_tile_if, @@ -577,7 +599,11 @@ def _decode_work_tile_schedule_with_invariant_bridge( active_prelude() # Packed persistent QK reads Q's consumer-work stage directly, # so it needs no descriptor task local crossing this guard. - body(None) + if cfg.use_variable_seqlens_q: + body(None) + else: + invariant = bridge() + body(invariant) _work_queue_tail(work_queue) return @@ -662,6 +688,7 @@ class DecodeGenTask(Task): def __init__(self, **kwargs: TaskKwarg) -> None: """Capture decode-specific task config and initialize cache slots.""" self.cfg = kwargs.pop("cfg", None) + self.requested_splits = kwargs.pop("requested_splits", None) self.seqlens_kv = kwargs.pop("seqlens_kv", None) self.block_table_capacity = kwargs.pop("block_table_capacity", None) self.sparse_row_route_offsets = kwargs.pop("sparse_row_route_offsets", None) @@ -709,25 +736,12 @@ def init_variables(self, context: cute.Pointer | None = None) -> None: self._tmem_base_offset = _warp_broadcast_i32(loaded, 0) @cute.jit - def make_task_cache( - self, - ) -> tuple[ - cutlass.Int32, - cutlass.Int32, - cutlass.Int32, - cutlass.Int32, - cutlass.Int32, - cutlass.Int32, - cutlass.Int32, - cutlass.Int32, - cutlass.Int32, - cutlass.Int32, - ]: + def make_task_cache(self): """Return cached task-local values passed through StageInfo.""" # The cache is threaded through StageInfo so resource methods can read # task-local constants without recomputing them or adding more # resource variables to every schedule edge. - return ( + values = ( self._tmem_base_offset, self._warp_grp_thread_idx, self._local_warp_idx, @@ -740,6 +754,15 @@ def make_task_cache( self._kv_window_start, ) + if cutlass.const_expr(self.cfg.runtime_split_kv): + count = ( + cutlass.Int32(self.cfg.splits_kv) + if self.requested_splits is None + else self.requested_splits + ) + values = (*values, count) + return values + @cute.jit def _refresh_packed_q_bounds(self, work_tile: Any) -> None: """Attach one active CLC tile's packed Q bounds to its data resources.""" @@ -796,6 +819,11 @@ def _run_task_body_impl( context: ResourceContext | None = None, ) -> None: """Run one ordinary task tile and synchronize attention-sink tails.""" + if cutlass.const_expr(self.cfg is not None and self.cfg.runtime_split_kv): + work_tile = WorkTileInfo( + split_kv_grid_coord(self.cfg, work_tile.tile_idx), + work_tile.is_valid_tile, + ) Task._run_task_body_impl( self, work_tile, @@ -832,10 +860,9 @@ def _run_task_body_persistent( self, context: ResourceContext | None = None, ) -> None: - """Drain inactive packed tiles before each unconditional active body.""" + """Drain inactive tiles before each unconditional active body.""" use_packed_early_stop = ( self.cfg is not None - and self.cfg.use_variable_seqlens_q and self.cfg.use_persistent_scheduler and self._has_skip_if ) @@ -866,7 +893,8 @@ def _run_task_body_persistent( self.dummy = cutlass.Boolean(True) while work_tile.is_valid_tile: - self._refresh_packed_q_bounds(work_tile) + if cutlass.const_expr(self.cfg.use_variable_seqlens_q): + self._refresh_packed_q_bounds(work_tile) # The tile is known active here. Running the complete schedule # without a dynamic skip guard keeps HEAD-produced pipeline state # in scope for LOOP and TAIL. @@ -906,6 +934,9 @@ def get_domain(self, tile_coord: cute.Coord) -> cutlass.Int32 | int: """Return this task's loop domain for one static or persistent tile.""" if self.cfg is None: return self.domain + if cutlass.const_expr(not self.cfg.fixed_seqlen and self.seq_len_q is None): + # Padding and scheduler tasks do not consume request dimensions. + return self.domain # Sparse rows have a runtime-dependent number of prepared KV routes # even when sequence lengths are static. Load their compact header @@ -993,6 +1024,7 @@ def get_domain(self, tile_coord: cute.Coord) -> cutlass.Int32 | int: if ( self.seqlens_kv is None and not self.cfg.use_split_kv + and self.cfg.fixed_seqlen and not self.cfg.uses_runtime_q_kv_union ): return self.domain @@ -1003,6 +1035,7 @@ def get_domain(self, tile_coord: cute.Coord) -> cutlass.Int32 | int: # causal/window/split coordinate construction in every task warp. if cutlass.const_expr( not self.cfg.use_split_kv + and self.cfg.fixed_seqlen and not self.cfg.uses_runtime_q_kv_union and not self.cfg.use_sliding_window_causal ): @@ -1018,11 +1051,12 @@ def get_domain(self, tile_coord: cute.Coord) -> cutlass.Int32 | int: ) return loop_domain + cutlass.Int32(self.domain_bias) - # Decode the logical Q tile with the configured physical split fanout, - # then derive its causal/window K union and useful runtime split prefix. + # Work-tile coordinates retain the configured capacity stride. q_group_cta_idx = cutlass.Int32(tile_coord[0]) q_group_idx = q_group_cta_idx - if self.cfg.use_split_kv: + if self.cfg.uses_compact_split_grid: + q_group_idx = cutlass.Int32(cute.arch.block_idx()[2]) % self.cfg.num_q_ctas + elif self.cfg.use_split_kv: q_group_idx = q_group_cta_idx // cutlass.Int32(self.cfg.splits_kv) q_token_base = _q_group_token_base(self.cfg, q_group_idx) seq_len_q = ( @@ -1048,24 +1082,20 @@ def get_domain(self, tile_coord: cute.Coord) -> cutlass.Int32 | int: # Split-KV groups CTAs by K/V split. The loop domain is rounded so # each CTA in the group executes the same number of instruction pairs. if self.cfg.use_split_kv: - # The useful active prefix is derived from this configured-fanout - # local span at kernel entry. Recomputing it in every scheduled - # task produces the same span but adds a runtime integer division - # to each task warp. - splits_kv = cutlass.Int32(self.cfg.splits_kv) - num_insts_kv = cutlass.Int32(self.cfg.num_insts_kv) - tiles_per_cta_group = splits_kv * num_insts_kv - num_groups = ( - total_kv_tiles + tiles_per_cta_group - cutlass.Int32(1) - ) // tiles_per_cta_group - total_kv_tiles = cute.math.max( - num_groups * num_insts_kv, - num_insts_kv, + # The domain and tile consumers share the instruction-aligned + # partition. Runtime launches carry its host-prepared reciprocal. + total_kv_tiles = _runtime_configured_local_kv_tiles( + self.cfg, + seq_len_kv, + seq_len_q, + q_token_base, + requested_splits=self.requested_splits, ) - # Physical split coordinates and page-cache strides are laid out with - # the configured launch fanout. Runtime pruning only shortens the - # useful prefix; it must not renumber the remaining CTAs. - split_idx = cutlass.Int32(tile_coord[0]) % cutlass.Int32(self.cfg.splits_kv) + split_idx = cutlass.Int32(cute.arch.block_idx()[0]) + if not self.cfg.uses_compact_split_grid: + split_idx = cutlass.Int32(tile_coord[0]) % cutlass.Int32( + self.cfg.splits_kv + ) self._kv_raw_tile_base = skipped_tiles + split_idx * total_kv_tiles else: self._kv_raw_tile_base = skipped_tiles @@ -2370,25 +2400,28 @@ def pv_mma( FmhaStage.Head, ) + # TMEM P aliases S, so an instance's next QK must follow the PV that + # reads P. A Q128 tile with SMEM P leaves S once Softmax has the scores + # in registers: the next QK goes first and overlaps the exponentials + # of this tile (``stages_q128_p_in_smem``). + qk_before_pv = cfg.stages_q128_p_in_smem with domain_loop(0, domain, 1, unroll=1): - pv_mma(smem_v0, smem_p0, "vp_mma_loop", KV_INST0, FmhaStage.Loop) - qk_mma( - smem_k0, - tmem_s0, - tmem_stats_done0, - q_desc, - "qk_mma_loop", - FmhaStage.Loop, - ) - pv_mma(smem_v1, smem_p1, "vp_mma_loop", KV_INST1, FmhaStage.Loop) - qk_mma( - smem_k1, - tmem_s1, - tmem_stats_done1, - q_desc, - "qk_mma_loop", - FmhaStage.Loop, - ) + for inst_idx, smem_k, smem_v, tmem_s, tmem_stats_done, smem_p in ( + (KV_INST0, smem_k0, smem_v0, tmem_s0, tmem_stats_done0, smem_p0), + (KV_INST1, smem_k1, smem_v1, tmem_s1, tmem_stats_done1, smem_p1), + ): + if not qk_before_pv: + pv_mma(smem_v, smem_p, "vp_mma_loop", inst_idx, FmhaStage.Loop) + qk_mma( + smem_k, + tmem_s, + tmem_stats_done, + q_desc, + "qk_mma_loop", + FmhaStage.Loop, + ) + if qk_before_pv: + pv_mma(smem_v, smem_p, "vp_mma_loop", inst_idx, 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 @@ -3184,7 +3217,9 @@ def _softmax_schedule_body( # 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() + releases_sage_stage_after_max = ( + cfg.sage_score_scaling() is SageScoreScaling.DEQUANTIZED + ) compute_softmax_loop = { (False, False): tmem_s.compute_softmax_loop, (False, True): tmem_s.compute_sage_softmax_loop, @@ -3245,7 +3280,7 @@ def _softmax_schedule_body( **softmax_loop_kwargs ) if cutlass.const_expr(releases_sage_stage_after_max): - # Dequantized scores went back to TMEM; the P pass reads no words. + # The scores came back dequantized; 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 @@ -3291,6 +3326,12 @@ def _softmax_schedule_body( keeps_route_flags=sparse_route_flags, swaps_route_flags=sparse_token_word2, ) + elif cutlass.const_expr(use_sage): + smem_p.compute_sage_p( + new_max_arr=new_max_arr, + s_arr=s_arr, + sage_q_scale=sage_q_scale, + ) else: smem_p.compute_p( new_max_arr=new_max_arr, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/reduction.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/reduction.py index f28769a10c8d..7a518ab32716 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 @@ -62,6 +62,7 @@ ) from .fmha_decode_config import FmhaDecodeConfig from .fmha_decode_resources.helpers_common import ( + UniformQueryOffsets, _attention_sink_head_stride, _local_head_from_q_output_row, _pack_float2_to_bf16, @@ -74,7 +75,10 @@ fmul2, ) from .fmha_decode_resources.helpers_kv_tile_idx import _runtime_active_splits_kv -from .fmha_decode_resources.helpers_softmax import _pack_float4_to_fp8_e4m3 +from .fmha_decode_resources.helpers_softmax import ( + _pack_float4_to_fp8_e4m3, + _final_output_lse, +) from ..separate_reduction import ( merge_log2_lse, unpack_normalized_vec8, @@ -108,7 +112,7 @@ def _reduction_q_group_idx( ) if cutlass.const_expr(cfg.uses_nontrivial_grouped_q_layout): return token_idx // Int32(cfg.q_tokens_per_cta) - if cutlass.const_expr(cfg.max_seq_len_q > 1): + if cutlass.const_expr(cfg.is_multi_token_q): head_ctas_per_token = Int32( (cfg.heads_q_per_kv + cfg.tile_size_q - 1) // cfg.tile_size_q ) @@ -126,14 +130,17 @@ def _reduction_active_splits_kv( h_r: Int32, logical_output_row_idx: Int32, seq_len_q: Int32, + requested_splits=None, ) -> Int32: """Recompute the producer's useful split prefix for one output row.""" q_group_idx = _reduction_q_group_idx(cfg, h_r, logical_output_row_idx) + seq_len_kv = cute.math.max(Int32(g_seqlens_kv[b_idx]), Int32(0)) return _runtime_active_splits_kv( cfg, - Int32(g_seqlens_kv[b_idx]), + seq_len_kv, seq_len_q, _q_group_token_base(cfg, q_group_idx), + requested_splits=requested_splits, ) @@ -141,7 +148,7 @@ def _reduction_active_splits_kv( def _reduce_exact_splits_body( o_iter: cute.Pointer, g_seqlens_kv: cute.Pointer | DirectSparseMetadataView, - g_cu_seqlens_q: cute.Pointer, + g_cu_seqlens_q: cute.Pointer | UniformQueryOffsets, g_partial_o: cute.Pointer, g_partial_stats: cute.Pointer, g_attention_sinks: cute.Pointer, @@ -149,6 +156,8 @@ def _reduce_exact_splits_body( cfg: cutlass.Constexpr[FmhaDecodeConfig], static_full_split_prefix: cutlass.Constexpr[bool] = False, wait_for_pdl_producer: cutlass.Constexpr[bool] = False, + lse_ptr: cute.Pointer | None = None, + requested_splits=None, ) -> None: """Fold every split in one 512-thread CTA over one 8 KiB output slice. @@ -185,7 +194,9 @@ def _reduce_exact_splits_body( PARTIAL_O_ELEMENT_BYTES ) q_token_offset, seq_len_q = _q_seq_bounds(cfg, g_cu_seqlens_q, b_idx) - active_splits_kv = Int32(cfg.splits_kv) + active_splits_kv = ( + Int32(cfg.splits_kv) if requested_splits is None else Int32(requested_splits) + ) valid_reduce_row = _q_logical_output_row_is_valid_for_seq( cfg, g_q_output_rows, @@ -229,16 +240,22 @@ def _reduce_exact_splits_body( g_q_output_rows, reduce_row_idx, seq_len_q, + requested_splits=requested_splits, ) + # A zero prefix publishes no partials; reuse it without reloading lengths. + valid_reduce_row = valid_reduce_row & (active_splits_kv > Int32(0)) if valid_reduce_row: # The workspace retains configured-max strides, but only the runtime # active prefix was published by producer CTAs. - for split_idx_i in cutlass.range_constexpr(cfg.max_splits_kv): + for split_idx_i in cutlass.range( + active_splits_kv if cfg.runtime_split_kv else cfg.max_splits_kv, + unroll_full=not cfg.runtime_split_kv, + ): split_idx = Int32(split_idx_i) split_is_active = split_idx < active_splits_kv if cutlass.const_expr(static_full_split_prefix): - split_is_active = cutlass.const_expr(split_idx_i < cfg.splits_kv) + split_is_active = split_idx < Int32(cfg.splits_kv) if split_is_active: # LSE layout: [logical_kv][configured split][row]. workspace_row = _separate_workspace_row_offset( @@ -297,6 +314,7 @@ def _reduce_exact_splits_body( grid_h_k, attention_sink_head_idx, cfg, + lse_ptr=lse_ptr, ) @@ -304,13 +322,15 @@ def _reduce_exact_splits_body( def decode_gen_separate_reduction_kernel( o_iter: cute.Pointer, g_seqlens_kv: cute.Pointer | DirectSparseMetadataView, - g_cu_seqlens_q: cute.Pointer, + g_cu_seqlens_q: cute.Pointer | UniformQueryOffsets, g_partial_o: cute.Pointer, g_partial_stats: cute.Pointer, g_attention_sinks: cute.Pointer, g_q_output_rows: Int32, cfg: cutlass.Constexpr[FmhaDecodeConfig], static_full_split_prefix: cutlass.Constexpr[bool] = False, + lse_ptr: cute.Pointer | None = None, + requested_splits=None, ) -> None: """Run the 512-thread exact-split reference schedule.""" @@ -324,6 +344,8 @@ def decode_gen_separate_reduction_kernel( g_q_output_rows, cfg, static_full_split_prefix, + lse_ptr=lse_ptr, + requested_splits=requested_splits, ) @@ -368,6 +390,7 @@ def _store_parallel_reduction_output( grid_h_k: Int32, attention_sink_head_idx: Int32, cfg: cutlass.Constexpr[FmhaDecodeConfig], + lse_ptr: cute.Pointer | None = None, ) -> None: """Merge the optional sink and store one normalized output fragment. @@ -376,15 +399,15 @@ def _store_parallel_reduction_output( standalone reducer therefore only merges normalized states here. """ - sink_lse = _attention_sink_log2_lse( + split_weight = _final_output_lse( cfg, g_attention_sinks, - h_k_idx, - attention_sink_h_r, - grid_h_k, - attention_sink_head_idx, + lse_ptr, + global_lse, + output_row_idx, + h_k_idx * attention_sink_h_r + attention_sink_head_idx, + reduce_col_idx == Int32(0), ) - _, split_weight, _ = merge_log2_lse(global_lse, sink_lse) final_regs = cutlass.Array( Int32, @@ -450,13 +473,16 @@ def _store_parallel_reduction_output( def decode_gen_parallel_separate_reduction_kernel( o_iter: cute.Pointer, g_seqlens_kv: cute.Pointer | DirectSparseMetadataView, - g_cu_seqlens_q: cute.Pointer, + g_cu_seqlens_q: cute.Pointer | UniformQueryOffsets, g_partial_o: cute.Pointer, g_partial_stats: cute.Pointer, g_attention_sinks: cute.Pointer, g_q_output_rows: Int32, cfg: cutlass.Constexpr[FmhaDecodeConfig], static_full_split_prefix: cutlass.Constexpr[bool] = False, + lse_ptr: cute.Pointer | None = None, + requested_splits=None, + compact_schedule: cutlass.Constexpr[bool] = False, ) -> None: """Reduce split partials with a compact or clustered constexpr schedule. @@ -469,7 +495,7 @@ def decode_gen_parallel_separate_reduction_kernel( two-level 4x4 merge, and G2/G4/G8 finalize through rank zero. """ - if cutlass.const_expr(cfg.use_compact_parallel_reduction): + if cutlass.const_expr(cfg.use_compact_parallel_reduction or compact_schedule): _reduce_exact_splits_body( o_iter, g_seqlens_kv, @@ -481,6 +507,8 @@ def decode_gen_parallel_separate_reduction_kernel( cfg, static_full_split_prefix, cfg.use_parallel_separate_reduction_pdl, + lse_ptr=lse_ptr, + requested_splits=requested_splits, ) return @@ -513,7 +541,9 @@ def decode_gen_parallel_separate_reduction_kernel( PARTIAL_O_ELEMENT_BYTES ) q_token_offset, seq_len_q = _q_seq_bounds(cfg, g_cu_seqlens_q, b_idx) - active_splits_kv = Int32(cfg.splits_kv) + active_splits_kv = ( + Int32(cfg.splits_kv) if requested_splits is None else Int32(requested_splits) + ) valid_reduce_row = _q_logical_output_row_is_valid_for_seq( cfg, g_q_output_rows, @@ -542,19 +572,11 @@ def decode_gen_parallel_separate_reduction_kernel( output_vals[elem_idx] = Float32(0.0) global_lse = Float32(-Float32.inf) - # Issue every owned slot's independent GMEM loads before folding them. The - # final batch width is compile-time. Padded split slots stay neutral and - # never form GMEM pointers, including non-power-of-two split counts. + # Batch independent GMEM loads before folding them. Storage capacity fixes + # the layout; runtime profiles only traverse their published split prefix. + load_batch = 2 if cfg.runtime_split_kv else PARALLEL_REDUCTION_LOAD_BATCH local_splits = cfg.parallel_reduction_splits_per_cta lane_splits = local_splits // slot_lanes - partial_lse = cutlass.Array( - Float32, PARALLEL_REDUCTION_LOAD_BATCH, space=cutlass.AddressSpace.rmem - ) - partial_regs = cutlass.Array( - Int32, - PARALLEL_REDUCTION_LOAD_BATCH * PACKED_OUTPUT_REGS_PER_THREAD, - space=cutlass.AddressSpace.rmem, - ) # Allocate the clustered path's local SMEM before acquiring the producer. # Cluster-size-one profiles compile out every later use of these arrays. @@ -584,12 +606,33 @@ def decode_gen_parallel_separate_reduction_kernel( g_q_output_rows, reduce_row_idx, seq_len_q, + requested_splits=requested_splits, ) - - for split_base_i in cutlass.range_constexpr( - 0, lane_splits, PARALLEL_REDUCTION_LOAD_BATCH + # A zero prefix publishes no partials; reuse it without reloading lengths. + valid_reduce_row = valid_reduce_row & (active_splits_kv > Int32(0)) + + split_rounds = lane_splits + if cutlass.const_expr(cfg.runtime_split_kv): + # Capacity only determines scratch strides. Distribute the live prefix + # cyclically across the fixed reducer topology and stop at its runtime end. + split_rounds = cute.ceil_div( + active_splits_kv, slot_lanes * cfg.parallel_reduction_cluster_size + ) + for split_base_i in cutlass.range( + 0, + split_rounds, + load_batch, + unroll_full=not cfg.runtime_split_kv, ): - batch_width = min(PARALLEL_REDUCTION_LOAD_BATCH, lane_splits - split_base_i) + partial_lse = cutlass.Array( + Float32, load_batch, space=cutlass.AddressSpace.rmem + ) + partial_regs = cutlass.Array( + Int32, + load_batch * PACKED_OUTPUT_REGS_PER_THREAD, + space=cutlass.AddressSpace.rmem, + ) + batch_width = load_batch split_base = Int32(split_base_i) # Prefetch only published slots; inactive slots and padded Q rows # are in bounds but may contain uninitialized caller-owned scratch. @@ -599,11 +642,18 @@ def decode_gen_parallel_separate_reduction_kernel( + (split_base + Int32(jj)) * Int32(slot_lanes) + lane_slot ) - valid_split_idx = cutlass.const_expr( - cfg.parallel_reduction_padded_splits == cfg.max_splits_kv - ) or split_idx < Int32(cfg.max_splits_kv) + if cutlass.const_expr(cfg.runtime_split_kv): + split_idx = ( + (split_base + Int32(jj)) * Int32(slot_lanes) + lane_slot + ) * Int32(cfg.parallel_reduction_cluster_size) + cluster_rank + valid_split_idx = cutlass.const_expr(cfg.runtime_split_kv) or ( + split_base + Int32(jj) < Int32(lane_splits) + and split_idx < Int32(cfg.max_splits_kv) + ) active_split_idx = split_idx < active_splits_kv - if cutlass.const_expr(static_full_split_prefix): + if cutlass.const_expr( + static_full_split_prefix and not cfg.runtime_split_kv + ): active_split_idx = cutlass.Boolean(True) if valid_split_idx and active_split_idx and valid_reduce_row: workspace_row = _separate_workspace_row_offset( @@ -645,11 +695,18 @@ def decode_gen_parallel_separate_reduction_kernel( + (split_base + Int32(jj)) * Int32(slot_lanes) + lane_slot ) - valid_split_idx = cutlass.const_expr( - cfg.parallel_reduction_padded_splits == cfg.max_splits_kv - ) or split_idx < Int32(cfg.max_splits_kv) + if cutlass.const_expr(cfg.runtime_split_kv): + split_idx = ( + (split_base + Int32(jj)) * Int32(slot_lanes) + lane_slot + ) * Int32(cfg.parallel_reduction_cluster_size) + cluster_rank + valid_split_idx = cutlass.const_expr(cfg.runtime_split_kv) or ( + split_base + Int32(jj) < Int32(lane_splits) + and split_idx < Int32(cfg.max_splits_kv) + ) active_split_idx = split_idx < active_splits_kv - if cutlass.const_expr(static_full_split_prefix): + if cutlass.const_expr( + static_full_split_prefix and not cfg.runtime_split_kv + ): active_split_idx = cutlass.Boolean(True) if valid_split_idx and active_split_idx and valid_reduce_row: regs_base = jj * PACKED_OUTPUT_REGS_PER_THREAD @@ -727,6 +784,7 @@ def decode_gen_parallel_separate_reduction_kernel( grid_h_k, attention_sink_head_idx, cfg, + lse_ptr=lse_ptr, ) return @@ -883,6 +941,7 @@ def decode_gen_parallel_separate_reduction_kernel( grid_h_k, attention_sink_head_idx, cfg, + lse_ptr=lse_ptr, ) # Keep every peer CTA alive until rank zero has finished its DSMEM reads. @@ -895,21 +954,24 @@ def fmha_decode_separate_reduction_launch( problem_shape: tuple[Int32, Int32, Int32, Int32, Int32], o_iter: cute.Pointer, seqlens_kv_iter: cute.Pointer | DirectSparseMetadataView, - cu_seqlens_q_iter: cute.Pointer, + cu_seqlens_q_iter: cute.Pointer | UniformQueryOffsets, partial_o_iter: cute.Pointer, partial_stats_iter: cute.Pointer, attention_sinks_iter: cute.Pointer, - scale_s: Float32, + softmax_scale: Float32, output_scale: Float32, stream: cuda_drv.CUstream, cfg: cutlass.Constexpr[FmhaDecodeConfig], static_full_split_prefix: cutlass.Constexpr[bool] = False, + lse_ptr: cute.Pointer | None = None, + requested_splits=None, + runtime_q_capacity_rows: Int32 = 0, ) -> None: """Launch the standalone reducer after the main split-KV decode kernel. The serial reference grid is ``(slice, kv_head, batch)`` with 8 KiB slices. The production reducer keeps that geometry for compact S2-S4 and otherwise - maps cluster ranks into grid.x for each 2 KiB slice. ``scale_s`` and + maps cluster ranks into grid.x for each 2 KiB slice. ``softmax_scale`` and ``output_scale`` remain parameters for launch-ABI compatibility; producer CTAs already applied them to log2-LSE and normalized partial O, respectively. ``seqlens_kv_iter`` is internal reducer metadata used only to @@ -917,19 +979,21 @@ def fmha_decode_separate_reduction_launch( """ b, h_q, h_k, _, _ = problem_shape h_r = h_q // h_k - q_output_rows = h_r - if cutlass.const_expr(cfg.max_seq_len_q > 1): + q_output_rows = runtime_q_capacity_rows if not cfg.fixed_seqlen else h_r + if cutlass.const_expr(cfg.fixed_seqlen and cfg.is_multi_token_q): q_output_rows = h_r * Int32(cfg.max_seq_len_q) bytes_per_slice = REDUCTION_BYTES_PER_SLICE grid_q_output_rows = q_output_rows - if cutlass.const_expr(cfg.max_seq_len_q > 1 and cfg.heads_q_per_kv != 0): + if cutlass.const_expr( + cfg.fixed_seqlen and cfg.is_multi_token_q and cfg.heads_q_per_kv != 0 + ): # Grouped multi-token Q profiles lay rows out by the configured # heads-per-KV group instead of the launch-time h_q / h_k ratio. grid_q_output_rows = cfg.heads_q_per_kv * cfg.max_seq_len_q # The reference schedule uses 512-thread CTAs over contiguous 8 KiB slices. # Multiple CTAs cover wide dimensions or grouped SQ rows without changing # the producer workspace layout. - num_reduction_slices = max( + num_reduction_slices = cute.math.max( ( grid_q_output_rows * cfg.headdim * PARTIAL_O_ELEMENT_BYTES + bytes_per_slice @@ -941,7 +1005,7 @@ def fmha_decode_separate_reduction_launch( if cutlass.const_expr(cfg.use_parallel_separate_reduction): cluster_size = cfg.parallel_reduction_cluster_size parallel_bytes_per_slice = cfg.parallel_reduction_bytes_per_slice - num_parallel_slices = max( + num_parallel_slices = cute.math.max( ( grid_q_output_rows * cfg.headdim * PARTIAL_O_ELEMENT_BYTES + parallel_bytes_per_slice @@ -950,23 +1014,77 @@ def fmha_decode_separate_reduction_launch( // parallel_bytes_per_slice, 1, ) - decode_gen_parallel_separate_reduction_kernel( - o_iter, - seqlens_kv_iter, - cu_seqlens_q_iter, - partial_o_iter, - partial_stats_iter, - attention_sinks_iter, - q_output_rows, - cfg, - static_full_split_prefix, - ).launch( - grid=(num_parallel_slices * cluster_size, h_k, b), - block=[cfg.parallel_reduction_threads_per_cta, 1, 1], - cluster=[cluster_size, 1, 1], - stream=stream, - use_pdl=cfg.use_parallel_separate_reduction_pdl, - ) + if cutlass.const_expr( + cfg.runtime_split_kv + and requested_splits is not None + and not cfg.use_compact_parallel_reduction + ): + # Capacity fixes scratch strides, not the useful reducer grid. + # S1-S4 retain the wide 8KiB schedule; the generic runtime reducer + # assigns four lanes per fragment and otherwise launches 4x as many + # 2KiB CTAs even when only two split partials were published. + # Both entry points share one cached module across runtime counts. + if requested_splits <= Int32(4): + decode_gen_parallel_separate_reduction_kernel( + o_iter, + seqlens_kv_iter, + cu_seqlens_q_iter, + partial_o_iter, + partial_stats_iter, + attention_sinks_iter, + q_output_rows, + cfg, + static_full_split_prefix, + lse_ptr=lse_ptr, + requested_splits=requested_splits, + compact_schedule=True, + ).launch( + grid=(num_reduction_slices, h_k, b), + block=[REDUCTION_THREADS_PER_CTA, 1, 1], + cluster=[1, 1, 1], + stream=stream, + use_pdl=cfg.use_parallel_separate_reduction_pdl, + ) + else: + decode_gen_parallel_separate_reduction_kernel( + o_iter, + seqlens_kv_iter, + cu_seqlens_q_iter, + partial_o_iter, + partial_stats_iter, + attention_sinks_iter, + q_output_rows, + cfg, + static_full_split_prefix, + lse_ptr=lse_ptr, + requested_splits=requested_splits, + ).launch( + grid=(num_parallel_slices * cluster_size, h_k, b), + block=[cfg.parallel_reduction_threads_per_cta, 1, 1], + cluster=[cluster_size, 1, 1], + stream=stream, + use_pdl=cfg.use_parallel_separate_reduction_pdl, + ) + else: + decode_gen_parallel_separate_reduction_kernel( + o_iter, + seqlens_kv_iter, + cu_seqlens_q_iter, + partial_o_iter, + partial_stats_iter, + attention_sinks_iter, + q_output_rows, + cfg, + static_full_split_prefix, + lse_ptr=lse_ptr, + requested_splits=requested_splits, + ).launch( + grid=(num_parallel_slices * cluster_size, h_k, b), + block=[cfg.parallel_reduction_threads_per_cta, 1, 1], + cluster=[cluster_size, 1, 1], + stream=stream, + use_pdl=cfg.use_parallel_separate_reduction_pdl, + ) return decode_gen_separate_reduction_kernel( o_iter, @@ -978,6 +1096,8 @@ def fmha_decode_separate_reduction_launch( q_output_rows, cfg, static_full_split_prefix, + lse_ptr=lse_ptr, + requested_splits=requested_splits, ).launch( grid=(num_reduction_slices, h_k, b), # Reduction parallelism is entirely in block.x; y/z are singleton diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/constants.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/constants.py index 1e47d5c5b237..c903df7684c2 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/constants.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/constants.py @@ -106,11 +106,7 @@ # Dense MLA kernels specialize paged-KV addressing for this explicit ABI set. # Each listed size exactly partitions a 128-token KV tile; other page sizes are # outside the current kernel ABI. -SUPPORTED_MLA_PAGE_SIZES = (16, 32, 64, 128) - -# Both MLA launch families and their standalone reducers share one fixed -# workspace/scheduler split capacity. -MAX_MLA_SPLITS_KV = 128 +SUPPORTED_MLA_PAGE_SIZES = (4, 8, 16, 32, 64, 128) # Throughput 2CTA epilogue maps 128 local threads onto two 64-row groups and # 128-column output halves for vectorized GMEM publication. diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/diagnostics.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/diagnostics.py new file mode 100644 index 000000000000..0162f03def17 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/diagnostics.py @@ -0,0 +1,55 @@ +# Copyright (c) 2026 by FlashInfer team. +# Licensed under the Apache License, Version 2.0. +"""Opt-in advisory diagnostics for MLA task-schedule construction.""" + +from contextlib import contextmanager +import os +import warnings + +from cutlass.experimental.task_scheduling.task_manager import ( + TaskManager as BaseTaskManager, +) + + +def mla_debug_checks_enabled() -> bool: + """Read the shared Prims-TS debug flag at build time, independent of run validation.""" + return os.environ.get("FLASHINFER_PRIMS_TS_DEBUG_CHECKS", "0").lower() in ( + "1", + "true", + "yes", + "on", + ) + + +class MlaTaskManager(BaseTaskManager): + """Retain schedule verification while making advisory overlap notes opt-in.""" + + def _verify_head_loop_producer_overlap(self) -> None: + """Show the upstream advisory for explicit verbose/debug builds. + + MLA intentionally produces in HEAD and LOOP with separate stage/offset + handling. The upstream check only prints a possible-ambiguity note; + all dependency, resource, barrier and memory verification still runs. + The override name follows the upstream TaskManager interface. + """ + if self._verbose or mla_debug_checks_enabled(): + super()._verify_head_loop_producer_overlap() + + +@contextmanager +def mla_build_diagnostics(): + """Hide unused-routing advisories during normal builds, preserving other warnings. + + MLA carries some task-local state directly rather than through returned + routes. Suppress only that known schedule-builder advisory; compiler and + schedule verification failures retain their normal reporting. + """ + with warnings.catch_warnings(): + if not mla_debug_checks_enabled(): + warnings.filterwarnings( + "ignore", + message=r".* produces '.*' but it is never routed to a downstream work call\.", + category=UserWarning, + module=r"cutlass\.experimental\.task_scheduling\.schedule_builder", + ) + yield diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/gather.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/gather.py index d36f2ad068f2..367f659a951b 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/gather.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/gather.py @@ -21,11 +21,32 @@ from cutlass import Int32, Int64 from cutlass.cutlass_dsl import dsl_user_op from cutlass.experimental import primitives as prims +from cutlass.experimental import cuda + +from ...tensor_map import create_tensor_map_tiled_from_view from .constants import TMA_GATHER_ROWS from .sparse_indices import SparseKVIndices, SparseSourceTile +@cute.jit +def make_sparse_rope_map(cache: cute.Tensor, dtype_bytes: cutlass.Constexpr[int]): + """Map the QK-only tail of compact D576 storage for gather4 loads.""" + rope = cute.make_tensor( + cache.iterator + 512, + cute.make_layout((64, cache.shape[2]), stride=(1, cache.stride[2])), + ) + return create_tensor_map_tiled_from_view( + rope, + box_dims=(64, 1), + stride_order=(0, 1), + swizzle=cuda.TensorMapSwizzle.s64b + if dtype_bytes == 1 + else cuda.TensorMapSwizzle.s128b, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + + @dsl_user_op def _address(pointer, dtype, *, loc=None, ip=None): if hasattr(pointer, "data_ptr"): diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/layout.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/layout.py index a9d28a54ecef..feec2d33c3fc 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/layout.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/layout.py @@ -48,6 +48,7 @@ _TASK_CACHE_WARP_IDX = 2 _TASK_CACHE_LANE_IDX = 3 _TASK_CACHE_SEQ_LEN_KV = 4 +_TASK_CACHE_LOCAL_KV_TILES = 5 def num_softmax_scale_groups(cfg: MlaConfig) -> int: @@ -291,6 +292,9 @@ def o_stage_tmem_col_offset( v_stage_idx: int, ): """Return the TMEM column offset for one O pipeline and V head-dim stage.""" + if getattr(cfg, "ws_pv", False): + return o_stage_idx * Int32(256) + Int32((v_stage_idx // 2) * 128) + if cfg.kernel_variant == "keeps_mma_ab" and cfg.head_dim_per_cta_v > 256: return ( o_stage_idx * Int32(2 * cfg.tmem_o_buffer_cols) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/math.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/math.py index b2d9933ec306..e4e18432c766 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/math.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/math.py @@ -15,6 +15,9 @@ """Math, dtype, and atomic helper functions for MLA decode TS examples.""" from functools import partial +import math +import struct +import numbers import cutlass import cutlass.cute as cute @@ -41,6 +44,61 @@ def ceil_div(a, b): return (a + b - 1) // b +def max_softmax_rescale_threshold_log2(qkv_dtype: str) -> float: + """Return the largest legal ``softmax_rescale_threshold_log2`` per P dtype. + + P reaches ``2**T``. ``"e4m3"`` caps at 8 binades: P is quantized at + ``448 * 2**-T``, and a larger T would push more of the row into E4M3 + subnormals. ``"bf16"`` caps at 32 binades (P <= ``2**32``). Both bounds + match the trtllm-gen FMHA ``skipCorrThreshold`` checks for E4M3 and BF16 + BMM2. + """ + + if qkv_dtype == "e4m3": + return 8.0 + if qkv_dtype == "bf16": + return 32.0 + raise ValueError(f"unsupported qkv_dtype={qkv_dtype!r}") + + +def canonicalize_softmax_rescale_threshold_log2( + threshold: object, *, qkv_dtype: str +) -> float: + """Range-check ``softmax_rescale_threshold_log2``; return it float32-rounded. + + Zero keeps the exact rescale. A positive value is compiled into the + kernel, which compares it with the log2-scaled row-max increase. Any zero + (including ``-0.0``) returns ``0.0`` so it selects the default kernel. + """ + + if isinstance(threshold, bool) or not isinstance(threshold, numbers.Real): + raise TypeError("softmax_rescale_threshold_log2 must be a real number") + try: + value = float(threshold) + except OverflowError: + value = math.inf + if value == 0.0: + # The default needs no per-dtype bound. + return 0.0 + bound = max_softmax_rescale_threshold_log2(qkv_dtype) + if math.isfinite(value) and 0.0 < value <= bound: + # The kernel compares in float32: canonicalize so values equal in + # float32 share one compiled kernel, and reject positive values that + # round to zero (they would change the FP8 P scale without skipping). + value = struct.unpack(" int: + """Elements per token; quantized formats use uint8 elements.""" + return { + "native": 512 + rope_dim, + "nvfp4": 288, + "nvfp4.fp8rope": 352, + "fp8": 512, + }[cache_format] + + +def validate_nvfp4_kv_formats( + primary_format, extra_format, rope_dim, qkv_dtype, page_size +): + """Reject unsupported storage/compute pairs before constructing tensor maps. + + Page size one is the internal physical-row view, not an external cache-page + restriction. Native/native retains the caller's ordinary paged-cache path. + A mixed pair has FP8 only in the primary SWA pool and NVFP4 in the extra. + """ + formats = (primary_format, extra_format) + allowed = ( + ("native", "native"), + ("nvfp4", "nvfp4"), + ("nvfp4.fp8rope", "nvfp4.fp8rope"), + ("fp8", "nvfp4"), + ) + if formats not in allowed: + raise ValueError( + "cache formats must be native/native, nvfp4/nvfp4, nvfp4.fp8rope/nvfp4.fp8rope, or fp8 primary/nvfp4 extra" + ) + if primary_format == "native": + return + if qkv_dtype != "e4m3": + raise ValueError("NVFP4 KV formats require E4M3 Q") + if page_size != 1: + raise ValueError("NVFP4 KV requires an internal page-size-one storage-row view") + if primary_format == "nvfp4.fp8rope": + if rope_dim != 64: + raise ValueError("nvfp4.fp8rope requires 512 NoPE and 64 RoPE channels") + elif rope_dim != 0: + raise ValueError("nvfp4/fp8 require D512 QK/V without an extra RoPE tail") + + +class Nvfp4KVConfig: + """Common format predicates for both attention families.""" + + @property + def has_nvfp4_kv(self) -> bool: + """Select row-record loaders for NVFP4 pools or the FP8-primary/NVFP4-extra pair.""" + return self.kv_cache_format != "native" + + @property + def has_fp8_primary(self) -> bool: + """Return whether primary K/V bytes need a separate per-tensor descale.""" + return self.kv_cache_format == "fp8" + + @property + def max_cache_record_size(self) -> int: + """Return the largest storage-row extent, in bytes for the NVFP4 loader family.""" + return max( + cache_record_size(self.kv_cache_format, self.rope_dim), + cache_record_size(self.extra_kv_cache_format, self.rope_dim), + ) + + @property + def raw_record_stride(self) -> int: + # gather4 requires every four-row destination to be 128-byte aligned. + """Return the raw SMEM row pitch in bytes, aligned for four-row TMA gathers.""" + return (self.max_cache_record_size + 31) // 32 * 32 + + +# =========================================================================== +# Shared-memory addressing and publication +# =========================================================================== + + +@cute.jit +def publish_nvfp4_kv(barrier, count, *, cta_group: cutlass.Constexpr[int] = 1): + """Publish ordinary shared stores before signaling the MMA pipeline.""" + cute.arch.fence_view_async_shared() + # The elected publisher must observe every lane's converted stores. + # Proxy fencing alone does not make a relaxed manual complete_tx a + # release operation. Join the warp, then publish through a release fence. + cute.arch.sync_warp() + if prims.elect_sync(): + if cutlass.const_expr(cta_group == 2): + prims.fence_sync_restrict(prims.MemOrder.RELEASE) + # Use the cluster address mapping API, not encoded CTA-rank bits. + leader = prims.mapa(barrier, Int32(0)) + prims.mbarrier_complete_tx( + leader, Int32(count), scope=prims.MemScope.CLUSTER + ) + else: + cute.arch.inline_ptx("fence.acq_rel.cta;") + prims.mbarrier_complete_tx(barrier, Int32(count)) + + +@cute.jit +def swizzle_byte_offset(offset, width: cutlass.Constexpr[int]): + """Map an unswizzled byte offset into the matching 32/64/128-byte TMA swizzle. + + The selected XOR mask changes the 16-byte bank group using row bits; + callers must use the same swizzle width as the tensor map and MMA view.""" + mask = {32: 0x10, 64: 0x30, 128: 0x70}[width] + return offset ^ ((offset >> Int32(3)) & Int32(mask)) + + +@cute.jit +def store_shared_u32x4(dst, offset, a, b, c, d): + """Store four packed 32-bit words at dst + byte offset with one 16-byte SMEM store. + + The destination must be 16-byte aligned. Publication is a separate + operation so a loader can finish its complete tile before signaling MMA.""" + cute.arch.inline_ptx( + "st.shared.v4.b32 [{$r0}], {{$r1}, {$r2}, {$r3}, {$r4}};", + read_only_args=[_address(dst, Int32) + offset, a, b, c, d], + ) + + +# =========================================================================== +# Primary-pool FP8 selection +# =========================================================================== + + +@cute.jit +def source_is_fp8(cfg, use_extra): + """Select unchanged FP8 primary bytes from a source-tagged gather group. + + The extra pool is always NVFP4 in a mixed pair; the tag selects a pool, + not an element dtype within a record. Broadcast from lane zero because + only metadata-issuing lanes have valid source coordinates.""" + if cutlass.const_expr(cfg.kv_cache_format != "fp8"): + return cutlass.Boolean(False) + # Only the primary pool can be per-tensor FP8. The extra-source tag means + # NVFP4, even when the first selected row is a masked hole. Only issuing + # lanes own metadata, so broadcast that tag to every conversion lane. + extra = Int32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=Int32(use_extra), + offset=0, + mask_and_clamp=0x1F, + kind=prims.Shfl.IDX, + ) + ) + return extra == Int32(0) + + +# =========================================================================== +# NVFP4 conversion +# =========================================================================== + +NVFP4_DATA_BYTES = 256 + + +NVFP4_SCALE_BYTES = 32 + + +@cute.jit +def nvfp4_scale_byte_offset(block, rope_dim: cutlass.Constexpr[int] = 64): + """Offset within the scale region; only the mixed format permutes it.""" + if cutlass.const_expr(rope_dim == 0): + return block + return 8 * (block % 4) + block // 4 + + +def nvfp4_row_bytes(rope_dim: int) -> int: + """Packed D512, optional unscaled FP8 tail, and block16 E4M3 scales.""" + return NVFP4_DATA_BYTES + rope_dim + NVFP4_SCALE_BYTES + + +@cute.jit +def convert_nvfp4x8(word: Int32, scale: Int32): + """Convert eight packed E2M1 values with one E4M3 block-scale byte. + + ``word`` contains eight consecutive four-bit values; ``scale`` stores the + E4M3 byte in its low eight bits. Return two words containing eight E4M3 + bytes, rounded once with saturation for ordinary FP8 MMA. No tensor-wide + NVFP4 scale is applied. + """ + if cutlass.const_expr(not _has_packed_e4m3_e2m1_mul()): + # Older DSL builds lack packed multiply even with a newer system + # toolkit. Convert dense nibbles directly instead of unpacking to + # bytes only to repack them inside the multiply fallback. + return _cvt_e2m1_word_to_e4m3_words(word, scale) + scales = (scale & Int32(255)) * Int32(0x01010101) + # Each 16-bit half holds four contiguous E2M1 nibbles. CUDA 13.4's + # dense packed operand avoids padded-nibble masks and output permutations. + values = cute.arch.inline_ptx( + "{ .reg .b16 lo, hi;\n" + " mov.b32 {lo, hi}, {$r0};\n" + " mul.rn.e4m3x4.e2m1x4.e4m3x4.satfinite {$w0}, lo, {$r1};\n" + " mul.rn.e4m3x4.e2m1x4.e4m3x4.satfinite {$w1}, hi, {$r1};\n" + "}", + write_only_types=[Int32, Int32], + read_only_args=[word, scales], + ) + return values[0], values[1] + + +# =========================================================================== +# TMA descriptors, gathering and staging +# =========================================================================== + + +@cute.jit +def make_whole_record_map(pool, record_bytes: cutlass.Constexpr[int]): + """Describe a complete byte record as Uint16 values for unswizzled gather4. + + All supported records fit the 256-element TMA box limit using Uint16. + A 32-byte raw row pitch makes each four-row destination 128-byte aligned; + global rows retain their original record stride and are not repacked. + """ + word_bytes = 2 + raw_bytes = (record_bytes + 31) // 32 * 32 + packed = cute.make_tensor( + cute.make_ptr( + cutlass.Uint16, + pool.iterator.toint(), + cute.AddressSpace.gmem, + assumed_align=16, + ), + cute.make_layout( + (record_bytes // word_bytes, pool.shape[2]), + stride=(1, pool.stride[2] // word_bytes), + ), + ) + return create_tensor_map_tiled_from_view( + packed, + box_dims=(raw_bytes // word_bytes, 1), + stride_order=(0, 1), + swizzle=cuda.TensorMapSwizzle.none, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + + +@cute.jit +def make_nvfp4_kv_maps( + primary, + extra, + rope_dim: cutlass.Constexpr[int], + whole_record: cutlass.Constexpr[bool] = False, + primary_format: cutlass.Constexpr[str] = "nvfp4", + extra_format: cutlass.Constexpr[str] = "nvfp4", +): + """Build paired payload, block-scale and optional RoPE tensor maps. + + Each pool is the internal [1, record_bytes, physical_rows] byte view. + Return primary/extra maps in payload, scales, RoPE order. Whole-record + maps preserve the token-interleaved representation; panel maps place + 128 logical values in the MMA swizzle. FP8 primary data has no trailer, + so its unused scale/RoPE descriptors are safe aliases inside that pool.""" + nope = 512 + maps = () + for pool, fmt in ((primary, primary_format), (extra, extra_format)): + fp8 = fmt == "fp8" + data = cute.make_tensor( + pool.iterator, + cute.make_layout( + (((nope + 127) // 128) * 128, pool.shape[2]), + stride=(1, pool.stride[2] * (1 if fp8 else 2)), + ), + ) + scales = cute.make_tensor( + pool.iterator + (0 if fp8 else nope // 2 + rope_dim), + cute.make_layout( + (16 if fp8 else 32, pool.shape[2]), stride=(1, pool.stride[2]) + ), + ) + rope = cute.make_tensor( + pool.iterator + (nope // 2 if rope_dim else 0), + cute.make_layout((64, pool.shape[2]), stride=(1, pool.stride[2])), + ) + if cutlass.const_expr(whole_record): + data_map = make_whole_record_map(pool, cache_record_size(fmt, rope_dim)) + elif cutlass.const_expr(fp8): + data_map = create_tensor_map_tiled_from_view( + data, + box_dims=(128, 1), + stride_order=(0, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + else: + data_map = create_tensor_map_tiled_from_view( + data, + box_dims=(128, 1), + stride_order=(0, 1), + swizzle=cuda.TensorMapSwizzle.s128b, + dtype=cutlass.Float4E2M1FNx2, + tma_format=cuda.TensorMapDataFormat.B4X16_P64, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ) + maps += ( + data_map, + create_tensor_map_tiled_from_view( + scales, + box_dims=(32, 1), + stride_order=(0, 1), + swizzle=cuda.TensorMapSwizzle.s32b, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ), + create_tensor_map_tiled_from_view( + rope, + box_dims=(64, 1), + stride_order=(0, 1), + swizzle=cuda.TensorMapSwizzle.s64b, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ), + ) + return maps[0], maps[3], maps[1], maps[4], maps[2], maps[5] + + +@cute.jit +def convert_nvfp4_panel( + dst, + scales, + column, + *, + rows: cutlass.Constexpr[int], + rope_dim: cutlass.Constexpr[int], +): + """Expand a gathered 128-column E2M1 panel in place into E4M3 MMA bytes. + + TMA B4X16_P64 leaves each eight packed bytes in a sixteen-byte slot. + One lane reads sixteen values and their block16 scale, rounds once to + E4M3, then fills that slot; D576 uses its permuted scale-byte order.""" + lane = cute.arch.thread_idx()[0] % 32 + for chunk in cutlass.range(lane, rows * 8, 32, unroll=4): + row, col = chunk // 8, chunk % 8 * 16 + logical_col = column + col + if logical_col < 512: + byte = swizzle_byte_offset(row * 128 + col, 128) + packed = cutlass.Pointer(dst + byte, dtype=Int32).load(count=2, alignment=8) + block = logical_col // 16 + scale_index = nvfp4_scale_byte_offset(block, rope_dim) + scale_byte = cutlass.Array(scales, dtype=cutlass.Uint8)[ + swizzle_byte_offset(row * 32 + scale_index, 32) + ] + a, b = convert_nvfp4x8(packed[0], Int32(scale_byte)) + c, d = convert_nvfp4x8(packed[1], Int32(scale_byte)) + store_shared_u32x4(dst, byte, a, b, c, d) + + +@cute.jit +def begin_raw_load(barrier, count): + """Have lane zero arm a raw TMA barrier for count bytes and return its epoch token. + + The warp joins before issuing gathers. Only lane zero later consumes + the returned token; other lanes wait through the matching warp join.""" + state = Int64(0) + if cute.arch.thread_idx()[0] % 32 == 0: + state = prims.mbarrier_arrive_expect_tx(barrier, Int32(count)) + cute.arch.sync_warp() + return state + + +@cute.jit +def wait_raw_load(barrier, state): + """Wait in lane zero for the armed raw TMA epoch, then release the whole warp.""" + if cute.arch.thread_idx()[0] % 32 == 0: + while not prims.mbarrier_try_wait(barrier, state): + pass + cute.arch.sync_warp() + + +@cute.jit +def gather_cached( + dst, + primary_map, + extra_map, + column, + barrier, + cached, + base, + *, + rows: cutlass.Constexpr[int], + width, +): + """Issue four-row TMA gathers from register-cached, source-tagged storage indices. + + Each issuing lane owns one quad; width is its destination bytes per row. + The high source bit chooses the pool map and is stripped from row IDs. + Invalid rows become the TMA out-of-bounds sentinel and are zero-filled.""" + lane = cute.arch.thread_idx()[0] % 32 + if lane < rows // 4: + raw = cached[base] + coords = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) + for j in cutlass.range_constexpr(4): + row = cached[base + j] & Int32(0x7FFFFFFF) + coords[j] = row if row != Int32(0x7FFFFFFF) else Int32(-1) + descriptor = select_gather_map(primary_map, extra_map, raw < 0) + gather4( + dst + lane * 4 * width, + descriptor, + column, + coords[0], + coords[1], + coords[2], + coords[3], + barrier, + ) + + +@cute.jit +def load_nvfp4_panel( + dst, + scratch, + barrier, + cached, + cache_base, + column, + data_map, + extra_data_map, + scale_map, + extra_scale_map, + rope_map, + extra_rope_map, + *, + rows: cutlass.Constexpr[int], + rope_dim: cutlass.Constexpr[int], + rope_only: cutlass.Constexpr[bool], + load_scales: cutlass.Constexpr[bool], + fp8=False, +): + """Gather one feature panel, applying NVFP4 block scales only for packed storage. + + The caller owns the raw barrier and scratch. Payload and scale byte + credits match their tensor-map transfers; FP8 primary and unscaled + RoPE bypass conversion. MMA publication remains the caller's job.""" + if cutlass.const_expr(rope_only): + state = begin_raw_load(barrier, rows * 64) + gather_cached( + dst, + rope_map, + extra_rope_map, + Int32(0), + barrier, + cached, + cache_base, + rows=rows, + width=64, + ) + wait_raw_load(barrier, state) + else: + # B4X16_P64 writes 8 packed bytes followed by 8 pad bytes per 16 + # values. Its completion count is the packed bytes, not padded SMEM. + state = begin_raw_load( + barrier, rows * (128 if fp8 else 64 + (32 if load_scales else 0)) + ) + gather_cached( + dst, + data_map, + extra_data_map, + column, + barrier, + cached, + cache_base, + rows=rows, + width=128, + ) + if load_scales and not fp8: + gather_cached( + scratch, + scale_map, + extra_scale_map, + Int32(0), + barrier, + cached, + cache_base, + rows=rows, + width=32, + ) + wait_raw_load(barrier, state) + if not fp8: + convert_nvfp4_panel(dst, scratch, column, rows=rows, rope_dim=rope_dim) + + +@cute.jit +def load_nvfp4_panels( + destinations, + columns, + scratch, + barrier, + cached, + cache_base, + data_map, + extra_data_map, + scale_map, + extra_scale_map, + rope_map, + extra_rope_map, + *, + rows: cutlass.Constexpr[int], + rope_dim: cutlass.Constexpr[int], + separate_rope: cutlass.Constexpr[bool], + fp8=False, +): + """Overlap all feature-panel gathers before one raw-data wait.""" + panels = len(columns) + state = begin_raw_load( + barrier, + rows + * ( + panels * (128 if fp8 else 64) + + (0 if fp8 else 32) + + (64 if separate_rope else 0) + ), + ) + for part in cutlass.range_constexpr(panels): + gather_cached( + destinations[part], + data_map, + extra_data_map, + columns[part], + barrier, + cached, + cache_base, + rows=rows, + width=128, + ) + if not fp8: + gather_cached( + scratch, + scale_map, + extra_scale_map, + Int32(0), + barrier, + cached, + cache_base, + rows=rows, + width=32, + ) + if cutlass.const_expr(separate_rope): + gather_cached( + destinations[panels], + rope_map, + extra_rope_map, + Int32(0), + barrier, + cached, + cache_base, + rows=rows, + width=64, + ) + wait_raw_load(barrier, state) + for part in cutlass.range_constexpr(panels): + if not fp8: + convert_nvfp4_panel( + destinations[part], scratch, columns[part], rows=rows, rope_dim=rope_dim + ) + + +@cute.jit +def issue_whole_record( + dst, barrier, cached, data_map, extra_data_map, *, issuer, row_bytes +): + """One gather per quad for a complete NVFP4 or per-tensor FP8 storage record.""" + begin_raw_load(barrier, Int32(64) * row_bytes) + gather_cached( + dst + issuer * Int32(64) * row_bytes, + data_map, + extra_data_map, + Int32(0), + barrier, + cached, + 0, + rows=64, + width=row_bytes, + ) + + +@cute.jit +def convert_nvfp4_record( + dst, + thread_idx, + *, + stage_bytes: cutlass.Constexpr[int], + rope_dim: cutlass.Constexpr[int], + raw=None, + raw_empty=None, +): + """Snapshot two16-value rows per lane, then write conflict-free panels.""" + row_bytes = nvfp4_row_bytes(rope_dim) + if cutlass.const_expr(raw is None): + raw = dst + packed = cutlass.Array(Int32, 32, space=cutlass.AddressSpace.rmem) + scales = cutlass.Array(Int32, 8, space=cutlass.AddressSpace.rmem) + rope = cutlass.Array(Int32, 8, space=cutlass.AddressSpace.rmem) + lane = thread_idx & Int32(31) + first_row = thread_idx >> Int32(5) + column = lane * Int32(16) + sf_column = nvfp4_scale_byte_offset(lane, rope_dim) + for chunk in cutlass.range_constexpr(8): + row0 = first_row + Int32(chunk * 16) + row1 = row0 + Int32(8) + words0 = cutlass.Pointer( + raw + row0 * Int32(row_bytes) + column // Int32(2), dtype=Int32 + ).load(count=2, alignment=8) + words1 = cutlass.Pointer( + raw + row1 * Int32(row_bytes) + column // Int32(2), dtype=Int32 + ).load(count=2, alignment=8) + packed[chunk * 4], packed[chunk * 4 + 1] = words0[0], words0[1] + packed[chunk * 4 + 2], packed[chunk * 4 + 3] = words1[0], words1[1] + sf0 = cutlass.Array(raw, dtype=cutlass.Uint8)[ + row0 * Int32(row_bytes) + Int32(NVFP4_DATA_BYTES + rope_dim) + sf_column + ] + sf1 = cutlass.Array(raw, dtype=cutlass.Uint8)[ + row1 * Int32(row_bytes) + Int32(NVFP4_DATA_BYTES + rope_dim) + sf_column + ] + scales[chunk] = Int32(sf0) | (Int32(sf1) << Int32(8)) + if cutlass.const_expr(rope_dim): + rope0 = cutlass.Pointer( + raw + + row0 * Int32(row_bytes) + + Int32(NVFP4_DATA_BYTES) + + lane * Int32(2), + dtype=cutlass.Uint16, + ).load() + rope1 = cutlass.Pointer( + raw + + row1 * Int32(row_bytes) + + Int32(NVFP4_DATA_BYTES) + + lane * Int32(2), + dtype=cutlass.Uint16, + ).load() + rope[chunk] = Int32(rope0) | (Int32(rope1) << Int32(16)) + if cutlass.const_expr(raw_empty is not None): + # Raw storage is disjoint from every expanded destination. Each + # warp snapshots all of its bytes before its release-arrival. The + # raw-empty barrier has eight arrivals; the metadata TMA producer + # acquires it and proxy-fences before any overwrite. + cute.arch.sync_warp() + if prims.elect_sync(): + prims.mbarrier_arrive(raw_empty, scope=prims.MemScope.CTA) + else: + # In-place expansion must join all readers before any output store. + prims.barrier_cta_sync(barrier_id=12, thread_count=256) + panel = lane >> Int32(3) + local_column = (lane & Int32(7)) * Int32(16) + for chunk in cutlass.range_constexpr(8): + for pair in cutlass.range_constexpr(2): + row = first_row + Int32(chunk * 16 + pair * 8) + sf = (scales[chunk] >> Int32(pair * 8)) & Int32(255) + a, b = convert_nvfp4x8(packed[chunk * 4 + pair * 2], sf) + c, d = convert_nvfp4x8(packed[chunk * 4 + pair * 2 + 1], sf) + store_shared_u32x4( + dst + panel * Int32(stage_bytes), + swizzle_byte_offset(row * Int32(128) + local_column, 128), + a, + b, + c, + d, + ) + if cutlass.const_expr(rope_dim): + packed_rope = cutlass.Uint16( + (rope[chunk] >> Int32(pair * 16)) & Int32(65535) + ) + cutlass.Pointer( + dst + + Int32(4 * stage_bytes) + + swizzle_byte_offset(row * Int32(64) + lane * Int32(2), 64), + dtype=cutlass.Uint16, + ).store(packed_rope) + + +@cute.jit +def copy_fp8_record(dst, raw, thread_idx, raw_empty): + """Reorder unmodified FP8 bytes into four retained MMA panels.""" + lane = thread_idx % Int32(32) + warp = thread_idx // Int32(32) + for i in cutlass.range(16, unroll=2): + row = warp + i * Int32(8) + offset = (lane // Int32(8)) * Int32(16384) + swizzle_byte_offset( + row * Int32(128) + (lane % Int32(8)) * Int32(16), 128 + ) + words = cutlass.Pointer( + raw + row * Int32(512) + lane * Int32(16), dtype=Int32 + ).load(count=4, alignment=16) + store_shared_u32x4(dst, offset, words[0], words[1], words[2], words[3]) + cute.arch.sync_warp() + if prims.elect_sync(): + prims.mbarrier_arrive(raw_empty, scope=prims.MemScope.CTA) + + +@cute.jit +def prefetch_nvfp4_kv_row(primary, extra, route, cfg): + """Warm the complete selected record while metadata runs ahead of conversion.""" + row = route & Int32(0x7FFFFFFF) + address = primary.iterator.toint() + capacity = Int32(primary.shape[2]) + stride = Int32(primary.stride[2]) + width = Int32(cache_record_size(cfg.kv_cache_format, cfg.rope_dim)) + if route < Int32(0): + address = extra.iterator.toint() + capacity = Int32(extra.shape[2]) + stride = Int32(extra.stride[2]) + width = Int32(cache_record_size(cfg.extra_kv_cache_format, cfg.rope_dim)) + if row < capacity: + address += Int64(row) * Int64(stride) + for line in cutlass.range_constexpr(3): + cute.arch.inline_ptx( + "prefetch.global.L2 [{$r0}];", + read_only_args=[address + Int64(line * 128)], + ) + if cutlass.const_expr(cfg.rope_dim): + cute.arch.inline_ptx( + "prefetch.global.L2 [{$r0}];", read_only_args=[address + Int64(351)] + ) + if cutlass.const_expr(cfg.has_fp8_primary): + if width == Int32(512): + for line in cutlass.range_constexpr(3, 4): + cute.arch.inline_ptx( + "prefetch.global.L2 [{$r0}];", + read_only_args=[address + Int64(line * 128)], + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/ops.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/ops.py index 021b9ebbd79f..ff2881d89831 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/ops.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/ops.py @@ -14,6 +14,8 @@ """Reusable low-level operations for MLA decode TS kernels.""" +import math + import cutlass import cutlass.cute as cute from cutlass import Float32, Int32, Int64, Uint32 @@ -235,6 +237,38 @@ def fp8_quant_scale_rcp(): return Float32(1.0 / 448.0) +# A positive ``softmax_rescale_threshold_log2`` T keeps a row's max anchor while +# a K tile raises it by at most T binades, so P reaches 2**T. E4M3 P then uses +# the scale 448 * 2**-T, as Prims-TS FMHA decode does (``fp8_p_quant_scale``), +# so the scaled P stays within 448. This equals the trtllm-gen FMHA +# ``getSoftmaxPScale`` at T = 0 (448) and at T = 8 (1.75). T is a compile-time +# config value, so these are plain Python helpers rather than ``cute.jit`` +# functions. +def _fp8_rescale_threshold_p_scale(cfg) -> float: + return 448.0 * 2.0**-cfg.softmax_rescale_threshold_log2 + + +def fp8_log2_p_scale(cfg): + """Return log2 of the E4M3 P scale selected by the config.""" + if cfg.softmax_rescale_threshold_log2 > 0.0: + return Float32(math.log2(_fp8_rescale_threshold_p_scale(cfg))) + return fp8_log2_quant_scale() + + +def fp8_p_scale(cfg): + """Return the E4M3 P scale selected by the config (``448 * 2**-T``).""" + if cfg.softmax_rescale_threshold_log2 > 0.0: + return Float32(_fp8_rescale_threshold_p_scale(cfg)) + return fp8_quant_scale() + + +def fp8_p_scale_rcp(cfg): + """Return the reciprocal of ``fp8_p_scale``.""" + if cfg.softmax_rescale_threshold_log2 > 0.0: + return Float32(1.0 / _fp8_rescale_threshold_p_scale(cfg)) + return fp8_quant_scale_rcp() + + @cute.jit def fp8_stsm_smem_dst( smem_base_i32, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/query.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/query.py index b1269f976be5..91fef29c24bb 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/query.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/query.py @@ -19,6 +19,7 @@ import cutlass import cutlass.cute as cute from cutlass import Int32, Int64 +from cutlass.experimental import primitives as prims @dataclass(frozen=True) @@ -75,7 +76,7 @@ def for_tile( def query_batch_bounds( cu_seqlens_q, batch_idx, - logical_seq_len_q: cutlass.Constexpr[int], + logical_seq_len_q: Int32, ): """Return the compact-storage offset and logical Q length for a batch. @@ -97,7 +98,7 @@ def runtime_flat_query_tile_valid_rows( query_tile_idx, tile_size_q: cutlass.Constexpr[int], logical_num_heads_q: cutlass.Constexpr[int], - logical_seq_len_q: cutlass.Constexpr[int], + logical_seq_len_q: Int32, cu_seqlens_q=None, batch_idx=None, ): @@ -122,7 +123,7 @@ def runtime_flat_query_tile_has_rows( query_tile_idx, tile_size_q: cutlass.Constexpr[int], logical_num_heads_q: cutlass.Constexpr[int], - logical_seq_len_q: cutlass.Constexpr[int], + logical_seq_len_q: Int32, cu_seqlens_q=None, batch_idx=None, ): @@ -144,7 +145,7 @@ def flat_query_row_state( query_tile_idx, tile_size_q: cutlass.Constexpr[int], logical_num_heads_q: cutlass.Constexpr[int], - logical_seq_len_q: cutlass.Constexpr[int], + logical_seq_len_q: Int32, cu_seqlens_q=None, batch_idx=None, ): @@ -158,8 +159,24 @@ def flat_query_row_state( local_flat_query_row = Int32(query_tile_idx) * Int32(tile_size_q) + Int32( row_in_tile ) - logical_q_idx = local_flat_query_row // Int32(logical_num_heads_q) - logical_head_idx = local_flat_query_row - logical_q_idx * Int32(logical_num_heads_q) + if cutlass.const_expr( + not isinstance(logical_seq_len_q, int) + and tile_size_q % logical_num_heads_q == 0 + ): + # Runtime Q benefits from factoring the tile/head ratio. Static Q keeps + # the shared flat-coordinate quotient/remainder to limit live values + # in the fused RoPE/quantization epilogue. + row_q_idx = Int32(row_in_tile) // Int32(logical_num_heads_q) + logical_q_idx = ( + Int32(query_tile_idx) * Int32(tile_size_q // logical_num_heads_q) + + row_q_idx + ) + logical_head_idx = Int32(row_in_tile) - row_q_idx * Int32(logical_num_heads_q) + else: + logical_q_idx = local_flat_query_row // Int32(logical_num_heads_q) + logical_head_idx = local_flat_query_row - logical_q_idx * Int32( + logical_num_heads_q + ) q_start, q_len = query_batch_bounds( cu_seqlens_q, batch_idx, @@ -208,31 +225,50 @@ def split_o_element_offset( ): """Return a batch-dynamic partial-O offset with 64-bit-safe products. - Batch size is intentionally absent from the compile signature, so the - complete workspace extent cannot prove that every offset fits in Int32. - The within-batch layout is still compile-time constant, however. Keep that - common bounded part in 32-bit arithmetic and widen only the batch-stride - product. Use fully widened arithmetic when a profile's per-batch layout - alone exceeds Int32. + Batch and Q extents are runtime values. Widen their tile index, but keep + the bounded head/split offset in 32-bit arithmetic to avoid carrying wide + intermediates through the output epilogue. """ - elements_per_batch = ( - cfg.seq_len_q * cfg.num_heads_q * cfg.num_ctas_per_seq_kv * cfg.head_dim_v - ) - if cutlass.const_expr(elements_per_batch <= (1 << 31) - 1): - within_batch_offset = ( - (Int32(q_idx) * Int32(cfg.num_heads_q) + Int32(head_idx)) - * Int32(cfg.num_ctas_per_seq_kv) - + Int32(split_idx) + elements_per_tile = cfg.num_heads_q * cfg.num_ctas_per_seq_kv * cfg.head_dim_v + tile_idx = Int64(batch_idx) * Int64(cfg.seq_len_q) + Int64(q_idx) + if cutlass.const_expr(elements_per_tile <= (1 << 31) - 1): + within_tile_offset = ( + Int32(head_idx) * Int32(cfg.num_ctas_per_seq_kv) + Int32(split_idx) ) * Int32(cfg.head_dim_v) + Int32(dim_idx) - return Int64(batch_idx) * Int64(elements_per_batch) + Int64(within_batch_offset) + return tile_idx * Int64(elements_per_tile) + Int64(within_tile_offset) return ( - ( - (Int64(batch_idx) * Int64(cfg.seq_len_q) + Int64(q_idx)) - * Int64(cfg.num_heads_q) - + Int64(head_idx) - ) + (tile_idx * Int64(cfg.num_heads_q) + Int64(head_idx)) * Int64(cfg.num_ctas_per_seq_kv) + Int64(split_idx) ) * Int64(cfg.head_dim_v) + Int64(dim_idx) + + +@cute.jit +def copy_q_fp8_to_tmem(source, tmem_base, *, cta_group: cutlass.Constexpr[int]): + """Copy64x512 E4M3 Q values to Layout E(1CTA) or Layout B(2CTA). + + Source consists of four SW128 panels. Callers own source lifetime and + copy-to-consumer ordering; this helper adds no synchronization. + """ + assert cta_group in (1, 2) + group = prims.CTAGroup.CTA_1 if cta_group == 1 else prims.CTAGroup.CTA_2 + base_desc = Int64( + prims.Tcgen05SmemDesc.build( + source, + leading_byte_offset=16, + stride_byte_offset=1024, + layout=2, + ) + ) + for part in cutlass.range_constexpr(32): + # The aligned start-address field stays inside the256KiB SMEM range. + descriptor = base_desc + Int64((part // 8) * 512 + part % 8) + prims.tcgen05_cp( + prims.Tcgen05CpShape.SHAPE_64X128B, + prims.make_tmem_ptr(tmem_base + Int32(part * 4), Int32), + descriptor, + group=group, + multicast=prims.Tcgen05CpMulticast.WARPX2_02_13, + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/schedule.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/schedule.py index 6e3626b3467c..6b13e64d8776 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/schedule.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/schedule.py @@ -150,10 +150,11 @@ def staged_kv_load( section: MlaStage, smem_page_offsets=None, cached_page_ids=None, + whole_kv=False, ): """Issue all K/V producer stages for one logical K or V tile.""" cached_page_ids = page_offsets_consume(smem_page_offsets, cached_page_ids) - for stage_idx in range(head_dim_stages): + for stage_idx in range(1 if whole_kv else head_dim_stages): smem_kv.acquire() getattr(smem_kv, producer_label)( stage_idx=stage_idx, @@ -174,6 +175,7 @@ def staged_qk_mma( consumer_label=None, include_acquire=True, release_k=True, + whole_kv=False, ): """Consume staged K descriptors and produce one QK MMA score tile.""" if consumer_label is None: @@ -189,13 +191,18 @@ def staged_qk_mma( if include_acquire: tmem_s.acquire() - for k_subtile_idx in range(head_dim_stages): + if whole_kv: smem_kv.wait() + for k_subtile_idx in range(head_dim_stages): + if not whole_kv: + smem_kv.wait() kv_desc = getattr(smem_kv, consumer_label)(k_subtile_idx=k_subtile_idx) tmem_s.qk_mma(kv_desc=kv_desc, k_subtile_idx=k_subtile_idx) - if release_k: + if release_k and not whole_kv: smem_kv.release() tmem_s.commit() + if whole_kv and release_k: + smem_kv.release() def staged_pv_mma( @@ -209,6 +216,8 @@ def staged_pv_mma( producer_label=None, is_tail: bool = False, reuse_k=False, + retained_k_stages=0, + whole_kv=False, ): """Consume staged V descriptors and produce one PV MMA output tile.""" smem_p.wait() @@ -251,7 +260,16 @@ def staged_pv_mma( v_subtile_idx=v_subtile_idx, is_tail=is_tail, ) + if not whole_kv: + smem_kv.release() + if whole_kv: + # QK and every PV feature slice consume the same complete latent tile. smem_kv.release() + elif reuse_k: + # QK-only RoPE stages follow the retained V channels. They were + # consumed by QK and must retire before the next tile's V slices. + for _ in range(retained_k_stages - head_dim_stages): + smem_kv.release() tmem_o.commit() smem_p.release() @@ -321,13 +339,17 @@ def staged_pv_mma_tmem_p( consumer_label, producer_label="pv_mma_loop_tmem_p", reuse_k=False, + retained_k_stages=0, + ws_pair=False, + whole_kv=False, ): """Consume TMEM P and staged V descriptors to produce one PV output tile.""" tmem_p.wait() p_stage_idx = tmem_p.p_stage() tmem_o.acquire() is_tail = producer_label == "pv_mma_tail_tmem_p" - for v_subtile_idx in range(head_dim_stages): + assert not ws_pair or reuse_k + for v_subtile_idx in range(0, head_dim_stages, 2 if ws_pair else 1): if not reuse_k: smem_kv.wait() v_desc_0 = getattr(smem_kv, consumer_label)(v_subtile_idx=v_subtile_idx) @@ -337,6 +359,18 @@ def staged_pv_mma_tmem_p( v_subtile_idx=v_subtile_idx, is_tail=is_tail, ) + # Both retained128-column panels were acquired by QK and remain live. + # Issue the combined PV before returning either panel to the producer. + if not whole_kv: + smem_kv.release() + if ws_pair: + smem_kv.release() + if whole_kv: smem_kv.release() + elif reuse_k: + # QK-only RoPE stages follow the retained V channels. They were + # consumed by QK and must retire before the next tile's V slices. + for _ in range(retained_k_stages - head_dim_stages): + smem_kv.release() tmem_o.commit() tmem_p.release() diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/sparse_indices.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/sparse_indices.py index d944aa034bd1..03df371cba68 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/sparse_indices.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/sparse_indices.py @@ -32,8 +32,13 @@ class SparseKVIndices: """Caller-owned index lists; optionally cache one query row's live lengths.""" + # Bind source tensors and optional cached row lengths; capacity counts padded KV + # slots. + # assume_valid_prefix promises nonnegative indices within each live prefix; + # has_extra enables the second source, whose tile domain follows the primary span. def __init__(self, values, capacity, assume_valid_prefix=False, has_extra=True): - self.values = values + """Carry caller-owned source metadata and an optional FP8-primary scale through JIT calls.""" + self.values = (*values, capacity) ( self.primary_indices, self.extra_indices, @@ -42,6 +47,7 @@ def __init__(self, values, capacity, assume_valid_prefix=False, has_extra=True): self.query_row, self.primary_length, self.extra_length, + self.primary_tensor_scale, ) = values self.capacity = capacity self.assume_valid_prefix = assume_valid_prefix @@ -52,19 +58,27 @@ def __init__(self, values, capacity, assume_valid_prefix=False, has_extra=True): (self.primary_length + Int32(127)) // Int32(128) * Int32(128) ) + # Flatten source tensors, row state and runtime capacity for CuTe calls. def __extract_mlir_values__(self): return cutlass.extract_mlir_values(self.values) + # Rebind source values and live capacity while preserving format promises. def __new_from_mlir_values__(self, values): + rebound = cutlass.new_from_mlir_values(self.values, values) return SparseKVIndices( - cutlass.new_from_mlir_values(self.values, values), - self.capacity, + rebound[:-1], + rebound[-1], self.assume_valid_prefix, self.has_extra, ) + # Load one query row's lengths once, or use the supplied token-count overrides. @cute.jit def bind_query_row(self, query_row, primary_length=None, extra_length=None): + """Bind one query to live source lengths, optionally reusing register-cached lengths. + + Round the primary span to 128 slots so a BK64/BK128 tile never mixes + pools. Padding is logical: caller index arrays and cache rows are not copied.""" if cutlass.const_expr(self.is_bound): return self row = Int64(query_row) @@ -86,12 +100,62 @@ def bind_query_row(self, query_row, primary_length=None, extra_length=None): ) if cutlass.const_expr(self.has_extra) else Int32(0), + self.primary_tensor_scale, ), self.capacity, self.assume_valid_prefix, self.has_extra, ) + # Select a source at a logical KV-token offset; for_mask may omit index loads. + @cute.jit + def tensor_scale_is_extra(self, tile_start, query_row=None): + """Return whether a logical KV tile starts in the extra pool. + + The primary list occupies ceil(primary_length/128)*128 logical slots. + Later slots select extra NVFP4 rows and therefore have tensor scale one. + Use the current graph-live length, not the planned list capacity.""" + if cutlass.const_expr(not self.has_extra): + return cutlass.Boolean(False) + length = ( + self.primary_length + if cutlass.const_expr(self.is_bound) + else Int32(self.primary_lengths[Int64(query_row)]) + ) + return tile_start >= (length + Int32(127)) // Int32(128) * Int32(128) + + @cute.jit + def tensor_scale(self, tile_start, query_row=None): + """Return K/V's descale for this tile: primary FP8 scale or one for NVFP4. + + NVFP4 block scales have already been applied in the loader. The FP8 + primary bytes remain unchanged, so score/output arithmetic consumes + the host-runtime FP32 scalar instead of rounding scaled K/V back into E4M3. + """ + scale = cutlass.Float32(1) + if not self.tensor_scale_is_extra(tile_start, query_row): + scale = cutlass.Float32(self.primary_tensor_scale) + return scale + + @cute.jit + def tensor_scale_ratio(self, current, previous, query_row=None): + """Convert accumulated O from the previous tile's V units to the current's. + + PV uses raw FP8 primary values and already-decoded NVFP4 extra values. + Therefore O_new = alpha * (s_previous / s_current) * O_old + P @ V_raw, + where alpha is the ordinary online-softmax correction. Primary tiles + precede extra tiles, so the only transition multiplies O by the FP8 + primary scale (the extra scale is one). Stay exactly one within a pool. + This factor changes O only; the softmax denominator and LSE use alpha. + """ + ratio = cutlass.Float32(1) + if cutlass.const_expr(self.has_extra): + if self.tensor_scale_is_extra( + current, query_row + ) and not self.tensor_scale_is_extra(previous, query_row): + ratio = cutlass.Float32(self.primary_tensor_scale) + return ratio + @cute.jit def select_source_tile( self, tile_start, query_row=None, *, for_mask: cutlass.Constexpr[bool] = False @@ -141,19 +205,25 @@ def select_source_tile( class SparseSourceTile: """Selected source and bounds; analytic prefix masks need no index pointer.""" + # Bind one source's indices, token origin/length and descriptor tag; valid_prefix + # skips index validity checks. def __init__(self, values, valid_prefix=False): self.values = values self.indices, self.origin, self.length, self.tag = values self.valid_prefix = valid_prefix + # Flatten the selected source pointer and token bounds for CuTe calls. def __extract_mlir_values__(self): return cutlass.extract_mlir_values(self.values) + # Rebind a source tile while retaining its valid-prefix promise. def __new_from_mlir_values__(self, values): return SparseSourceTile( cutlass.new_from_mlir_values(self.values, values), self.valid_prefix ) + # Clamp an index load inside source storage and return its separate live-token + # validity. @cute.jit def _bounded_index(self, token): position = token - self.origin @@ -169,6 +239,7 @@ def _bounded_index(self, token): valid = valid & (index >= Int32(0)) return index, valid + # Return one physical row with its source tag, or the tagged invalid-row sentinel. @cute.jit def load_row(self, token): index, valid = self._bounded_index(token) @@ -205,6 +276,7 @@ def load_quad(self, token): rows[j] = self.load_row(token + Int32(j)) return rows + # Test one logical KV token; request is unused because the source is row-bound. @cute.jit def is_invalid(self, token, request): del request @@ -213,12 +285,14 @@ def is_invalid(self, token, request): _, valid = self._bounded_index(token) return ~valid + # Check that a 128-token tile fits both the selected source and live execution span. @cute.jit def tile_is_full(self, start, live_length): return start + Int32(128) <= cute.math.min( self.origin + self.length, live_length ) + # Return one invalid-bit mask for 32 tokens starting within the selected source. @cute.jit def prefix_mask32(self, start, live_length): # Source/tile boundaries are aligned; start is within this source. @@ -232,6 +306,8 @@ def prefix_mask32(self, start, live_length): class SparseKVLengths: """Adapter for the scheduler's lengths[query_row] interface only.""" + # Bind per-query lengths or a cached scalar execution span; has_extra selects the + # second pool. def __init__(self, values, has_extra=True): self.values = values self.has_extra = has_extra @@ -240,14 +316,17 @@ def __init__(self, values, has_extra=True): (rows,) if self.bound_length is not None else self.primary_lengths.shape ) + # Flatten scheduler length tensors or the bound scalar for CuTe calls. def __extract_mlir_values__(self): return cutlass.extract_mlir_values(self.values) + # Rebind scheduler lengths while retaining the source-count contract. def __new_from_mlir_values__(self, values): return SparseKVLengths( cutlass.new_from_mlir_values(self.values, values), self.has_extra ) + # Return the padded execution span in tokens for one virtual query request. @cute.jit def __getitem__(self, query_row): if cutlass.const_expr(self.bound_length is not None): @@ -260,6 +339,7 @@ def __getitem__(self, query_row): ) +# Round each source to 128-token tiles; empty rows still execute one masked token. @cute.jit def _sparse_execution_length(primary_length, extra_length): primary_tiles = (primary_length + Int32(127)) // Int32(128) @@ -274,11 +354,12 @@ def bind_sparse_kv_indices( extra_indices, primary_lengths, extra_lengths, - capacity: cutlass.Constexpr[int], + capacity: Int32, *, query_row=None, assume_valid_prefix: cutlass.Constexpr[bool] = False, has_extra: cutlass.Constexpr[bool] = True, + primary_tensor_scale=None, ): """Bind index readers and scheduler lengths without copying metadata.""" indices = SparseKVIndices( @@ -290,6 +371,7 @@ def bind_sparse_kv_indices( None, None, None, + primary_tensor_scale, ), capacity, assume_valid_prefix, @@ -306,3 +388,28 @@ def bind_sparse_kv_indices( (primary_lengths, extra_lengths, None, None), has_extra ) return indices, lengths + + +@cute.jit +def cache_sparse_source_lengths( + page_offsets, + query_row, + lengths, + *, + include_primary_tiles: cutlass.Constexpr[bool] = False, +): + """Snapshot graph-live source lengths once per scheduled query work tile. + + Bound readers already carry scalar lengths. Unbound persistent readers + cache primary/extra lengths in registers; loaders can additionally retain + ceil(primary_length/128), the source boundary in full KV tiles. + """ + if cutlass.const_expr(not page_offsets.is_bound): + lengths[0] = Int32(page_offsets.primary_lengths[query_row]) + lengths[1] = ( + Int32(page_offsets.extra_lengths[query_row]) + if page_offsets.has_extra + else Int32(0) + ) + if cutlass.const_expr(include_primary_tiles): + lengths[2] = (lengths[0] + Int32(127)) // Int32(128) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/tile.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/tile.py index e35517b09e93..3216accc6635 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/tile.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/helpers/tile.py @@ -20,7 +20,7 @@ from cutlass.experimental.task_scheduling.resources import StageInfo -from .layout import _TASK_CACHE_SEQ_LEN_KV +from .layout import _TASK_CACHE_SEQ_LEN_KV, _TASK_CACHE_LOCAL_KV_TILES from .mask import MaskType, mask_visible_k_length from .query import ( flat_query_row_state, @@ -31,6 +31,14 @@ from ..throughput_latency_1cta.config import MlaConfig +@cute.jit +def runtime_split_count(requested_splits): + """Read the scalar count from optional host-prepared division metadata.""" + return Int32( + requested_splits[0] if isinstance(requested_splits, tuple) else requested_splits + ) + + @cute.jit def local_kv_tile_idx( cfg: MlaConfig, @@ -39,6 +47,7 @@ def local_kv_tile_idx( is_v: int, *, section: cutlass.Constexpr[MlaStage], + requested_splits=None, ): """Return the local K/V tile id for one staged pipeline instance.""" if cutlass.const_expr(section == MlaStage.Head): @@ -48,7 +57,11 @@ def local_kv_tile_idx( if cutlass.const_expr(is_v or (cfg.serial_swap_reuse and inst_id == 1)): return base + Int32(inst_id) return base + Int32(cfg.num_insts_kv + inst_id) - if cutlass.const_expr(cfg.total_kv_tiles <= cfg.num_insts_kv): + if cutlass.const_expr( + requested_splits is None + and cfg.seq_len_kv is not None + and cfg.total_kv_tiles <= cfg.num_insts_kv + ): return Int32(inst_id) return stage_info.loop_end * Int32(cfg.num_insts_kv) + Int32(inst_id) @@ -106,20 +119,36 @@ def runtime_split_pruning_is_profitable(configured_splits_kv: int) -> bool: @cute.jit -def _runtime_configured_local_kv_tiles(cfg: MlaConfig, seq_len_kv): +def _runtime_configured_local_kv_tiles( + cfg: MlaConfig, + seq_len_kv, + requested_splits=None, +): """Return the instruction-aligned local span for configured split capacity.""" if cutlass.const_expr(cfg.fixed_sparse_kv_tiles > 0): return Int32(cfg.fixed_sparse_kv_tiles) total_kv_tiles = runtime_total_kv_tiles(cfg, seq_len_kv) num_insts_kv = Int32(cfg.num_insts_kv) - tiles_per_group = Int32(cfg.num_ctas_per_seq_kv) * num_insts_kv - num_groups = (total_kv_tiles + tiles_per_group - Int32(1)) // tiles_per_group + split_count = ( + Int32(cfg.num_ctas_per_seq_kv) + if requested_splits is None + else runtime_split_count(requested_splits) + ) + tiles_per_group = split_count * num_insts_kv + divisor = ( + requested_splits[1] if isinstance(requested_splits, tuple) else tiles_per_group + ) + num_groups = (total_kv_tiles + tiles_per_group - Int32(1)) // divisor return cute.math.max(num_groups * num_insts_kv, num_insts_kv) @cute.jit -def runtime_num_ctas_kv(cfg: MlaConfig, seq_len_kv): +def runtime_num_ctas_kv( + cfg: MlaConfig, + seq_len_kv, + requested_splits=None, +): """Return the active prefix of configured split-KV CTAs. One CTA's minimum useful K unit is ``tile_size_kv * num_insts_kv``. @@ -129,17 +158,26 @@ def runtime_num_ctas_kv(cfg: MlaConfig, seq_len_kv): if cutlass.const_expr(cfg.use_multi_ctas_kv != 1): return Int32(1) total_kv_tiles = runtime_total_kv_tiles(cfg, seq_len_kv) - local_kv_tiles = _runtime_configured_local_kv_tiles(cfg, seq_len_kv) + local_kv_tiles = _runtime_configured_local_kv_tiles( + cfg, seq_len_kv, requested_splits=requested_splits + ) return (total_kv_tiles + local_kv_tiles - Int32(1)) // local_kv_tiles @cute.jit -def runtime_local_kv_tiles(cfg: MlaConfig, seq_len_kv): +def runtime_local_kv_tiles( + cfg: MlaConfig, + seq_len_kv, + requested_splits=None, +): """Return the padded local KV tile count assigned to each KV CTA group.""" if cutlass.const_expr(cfg.fixed_sparse_kv_tiles > 0): return Int32(cfg.fixed_sparse_kv_tiles) if cutlass.const_expr( - cfg.sparse_direct and cfg.seq_len_kv <= 256 and cfg.num_ctas_per_seq_kv == 1 + cfg.seq_len_kv is not None + and cfg.sparse_direct + and cfg.seq_len_kv <= 256 + and cfg.num_ctas_per_seq_kv == 1 ): # The short direct schedule executes its planned padded tiles. Every # consumer must use that same domain: using the smaller live length @@ -148,7 +186,9 @@ def runtime_local_kv_tiles(cfg: MlaConfig, seq_len_kv): return Int32(cfg.local_kv_tiles(cfg.total_kv_tiles)) if cutlass.const_expr(cfg.use_multi_ctas_kv != 1): return runtime_total_kv_tiles(cfg, seq_len_kv) - return _runtime_configured_local_kv_tiles(cfg, seq_len_kv) + return _runtime_configured_local_kv_tiles( + cfg, seq_len_kv, requested_splits=requested_splits + ) @cute.jit @@ -159,6 +199,19 @@ def runtime_base_seq_len_kv(cfg: MlaConfig, cache_seqs, batch_idx): return Int32(cache_seqs[batch_idx]) +@cute.jit +def runtime_request_publishes(cfg: MlaConfig, cache_seqs, batch_idx): + """Return whether a batch row publishes O / LSE (and split scratch). + + An empty request (zero KV length) publishes nothing, which preserves its + output. D576 sparse rows keep their contract instead: a token whose + top-k length is zero still writes O = 0 and LSE = -inf. + """ + if cutlass.const_expr(cfg.sparse_d576_reload): + return cutlass.Boolean(True) + return runtime_base_seq_len_kv(cfg, cache_seqs, batch_idx) > Int32(0) + + @cute.jit def runtime_seq_len_kv_for_logical_q( cfg: MlaConfig, @@ -249,9 +302,12 @@ def runtime_split_tile_is_active( batch_idx, cta_idx_q, cta_idx_kv, + requested_splits=None, ): """Return whether a configured split rank owns runtime KV work.""" + if cutlass.const_expr(cfg.use_multi_ctas_kv != 1): + return Int32(cta_idx_kv) < Int32(1) seq_len_kv = runtime_seq_len_kv_for_q( cfg, cache_seqs, @@ -259,7 +315,16 @@ def runtime_split_tile_is_active( cta_idx_q, cu_seqlens_q, ) - return Int32(cta_idx_kv) < runtime_num_ctas_kv(cfg, seq_len_kv) + total_kv_tiles = runtime_total_kv_tiles(cfg, seq_len_kv) + local_kv_tiles = runtime_local_kv_tiles( + cfg, seq_len_kv, requested_splits=requested_splits + ) + # A split owns work exactly when its first tile intersects the valid K + # prefix. Avoid dividing by the local span just to compare with its ceil. + # The physical suffix can be much larger than the requested count. + return cutlass.Int64(cta_idx_kv) * cutlass.Int64(local_kv_tiles) < cutlass.Int64( + total_kv_tiles + ) @cute.jit @@ -270,6 +335,7 @@ def runtime_work_tile_activity( batch_idx, cta_idx_q, cta_idx_kv, + requested_splits=None, ): """Return the independent runtime Q and split-KV activity predicates.""" @@ -286,6 +352,10 @@ def runtime_work_tile_activity( batch_idx, cta_idx_q, cta_idx_kv, + requested_splits=requested_splits, + ) + query_is_active = query_is_active & ( + runtime_request_publishes(cfg, cache_seqs, batch_idx) ) return query_is_active, split_is_active @@ -298,6 +368,7 @@ def runtime_work_tile_is_active( batch_idx, cta_idx_q, cta_idx_kv, + requested_splits=None, ): """Return whether a Q/split work tile owns any runtime data.""" @@ -308,6 +379,7 @@ def runtime_work_tile_is_active( batch_idx, cta_idx_q, cta_idx_kv, + requested_splits=requested_splits, ) return query_is_active and split_is_active @@ -380,11 +452,22 @@ def global_kv_tile_idx( local_tile_idx, seq_len_kv, cta_idx_kv, + requested_splits=None, + task_cache=None, ): """Map a local KV tile id to the global KV tile id for split-KV mode.""" if cutlass.const_expr(cfg.use_multi_ctas_kv != 1): return local_tile_idx - return Int32(cta_idx_kv) * runtime_local_kv_tiles(cfg, seq_len_kv) + local_tile_idx + if cutlass.const_expr(requested_splits is not None and task_cache is not None): + return ( + Int32(cta_idx_kv) * Int32(task_cache[_TASK_CACHE_LOCAL_KV_TILES]) + + local_tile_idx + ) + return ( + Int32(cta_idx_kv) + * runtime_local_kv_tiles(cfg, seq_len_kv, requested_splits=requested_splits) + + local_tile_idx + ) @cute.jit diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/kernel_policy.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/kernel_policy.py index 3949e8ca5eb8..28397af587ed 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/kernel_policy.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/kernel_policy.py @@ -152,8 +152,9 @@ def is_throughput_2cta_mla_supported_shape( ) -> bool: """Return whether the 2CTA M128 TS MLA path is eligible. - The predicate is intentionally side-effect free: it only checks the dense - MLA shape constraints that can be known before kernel construction. + Physical M128 tiles cover consecutive logical (query, head) rows, including + partial final tiles. The tile width does not limit the logical head count; + the public planner validates flattened coordinate and workspace bounds. """ del batch_size @@ -163,7 +164,7 @@ def is_throughput_2cta_mla_supported_shape( and latent_dim == 512 and rope_dim in (0, 64) and page_size in SUPPORTED_MLA_PAGE_SIZES - and 1 <= num_heads <= 128 + and num_heads >= 1 and seq_len_q >= 1 and seq_len_k >= 1 ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/parallel_reduction_topology.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/parallel_reduction_topology.py index 7ee4779fac6e..4f5108c94adc 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/parallel_reduction_topology.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/parallel_reduction_topology.py @@ -18,7 +18,6 @@ from dataclasses import dataclass -from .helpers.constants import MAX_MLA_SPLITS_KV _TARGET_SPLITS_PER_RANK = 8 _PARALLEL_REDUCER_MIN_SPLITS_PER_RANK = 32 @@ -75,8 +74,8 @@ def _next_power_of_two(value: int) -> int: def _validate_splits_kv(splits_kv: int) -> None: if isinstance(splits_kv, bool) or not isinstance(splits_kv, int): raise TypeError("splits_kv must be an integer") - if not 1 <= splits_kv <= MAX_MLA_SPLITS_KV: - raise ValueError(f"splits_kv must be in [1, {MAX_MLA_SPLITS_KV}]") + if splits_kv < 1: + raise ValueError("splits_kv must be positive") def _validate_positive_int(value: int, name: str) -> None: @@ -145,7 +144,7 @@ def validate_parallel_reduction_workspace( _validate_splits_kv(splits_kv) _validate_positive_int(head_dim, "head_dim") if splits_kv == 1: - raise ValueError("parallel reduction requires splits_kv in [2, 128]") + raise ValueError("parallel reduction requires at least two splits") workspace_elements = batch_size * num_heads_q * seq_len_q * splits_kv * head_dim if workspace_elements > _MAX_PARTIAL_O_WORKSPACE_ELEMENTS: diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/sparse_policy.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/sparse_policy.py index 8f6100825f04..df5c5524d2bc 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/sparse_policy.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/sparse_policy.py @@ -39,10 +39,16 @@ class SparseMlaProfile: kv_pipeline_stages: int = 0 kv_tile_size: int = 128 page_pipeline_stages: int = 0 + whole_kv_tile: bool = False -def _base_fp8(rows, heads, capacity, sm_count): - """Small-grid tile and split selection for FP8.""" +def base_fp8_profile(rows, heads, capacity, sm_count, rope_dim): + """Small-grid tile and split selection for FP8. + + ``rope_dim`` is the QK-only RoPE width (64 for D576). Rules gated on + ``rope_dim == 64`` are D576 crossovers measured on GB300 with top-k 2048; + D512 plans use the remaining rules. + """ # M64 shares KV across more heads but can leave a small grid underfilled. # M16 exposes more tasks at the cost of repeating KV work for head tiles. throughput_h32 = heads == 32 and rows * 2 >= sm_count and (capacity >= 512) @@ -73,11 +79,16 @@ def _base_fp8(rows, heads, capacity, sm_count): if heads == 32 and capacity >= 512 and rows >= sm_count: # Enough query rows amortize padded M64 work. Retaining two FP8 K # tiles saves V reloads; uniform coordinates amortize repeated D slices. + # D576 prefill-sized grids over long lists use the H64 issuer budget + # (GB300 top-k 2048, H24 prefill 8k: 1669 -> 1511 us). + prefill = rope_dim == 64 and capacity >= 2048 and rows >= 8 * sm_count return SparseMlaProfile( family="keep", + gather_issue_warps=8 if prefill else 4, reuse_kv=True, reuse_kv_stages=10, uniform_offset_cache=True, + paired_correction=prefill, scheduler="clc", ) if ( @@ -93,11 +104,35 @@ def _base_fp8(rows, heads, capacity, sm_count): if heads == 32 and 32 <= rows <= sm_count // 2 and (capacity >= 512): # This intermediate grid favors M64 reuse, with a second split only # for longer lists; additional splits lose to their merge overhead. + if rope_dim != 64 or capacity <= 1024: + return SparseMlaProfile( + family="keep", + split_kv=1 if capacity <= 1024 else 2, + reuse_kv=True, + reuse_kv_stages=10, + ) + # D576 long lists split KV up to one resident wave: four splits at + # 32-37 rows (GB300 top-k 2048, 29.6 -> 21.8 us), otherwise two. They + # use eight issuers; four splits finish their short K loops faster by + # streaming V than by retaining K (22.9 -> 20.7 us), two-split grids + # retain K with uniform coordinates (37.8 -> 36.3 us). + wave_splits = 1 << ((sm_count // rows).bit_length() - 1) + splits = min(4, max(2, wave_splits)) + retain = splits < 4 return SparseMlaProfile( family="keep", - split_kv=1 if capacity <= 1024 else 2, - reuse_kv=True, - reuse_kv_stages=10, + split_kv=splits, + gather_issue_warps=8, + reuse_kv=retain, + reuse_kv_stages=10 if retain else 0, + uniform_offset_cache=retain, + ) + if rope_dim == 64 and heads == 32 and 8 <= rows < 16 and capacity >= 2048: + # Few D576 rows: M16 Swap-AB with resident-wave KV splits; retaining K + # through PV beats reloading V (GB300 top-k 2048, B1 Sq8: 13.1 -> + # 12.2 us). + return SparseMlaProfile( + family="swap", split_kv=0, reuse_kv=True, reuse_kv_stages=10 ) if heads <= 16 and rows * 4 >= sm_count and (capacity >= 512): # Narrow Q tiles avoid padded heads; twelve stages retain two FP8 @@ -154,7 +189,7 @@ def _base_fp8(rows, heads, capacity, sm_count): ) -def _base_bf16(rows, heads, capacity, sm_count): +def base_bf16_profile(rows, heads, capacity, sm_count): """Small-grid tile and split selection for BF16.""" # BF16 BK128 cannot retain two complete D512 tiles in SMEM. A narrow # grid favors more small-Q tasks; a large grid amortizes M64's padding @@ -198,8 +233,30 @@ def _base_bf16(rows, heads, capacity, sm_count): return profile -def _choose_fp8(rows, heads, query_length, capacity, sm_count): - """Prefer resident split grids or sustained-throughput schedules when eligible.""" +def choose_fp8_profile(rows, heads, query_length, capacity, sm_count, rope_dim): + """Prefer resident split grids or sustained-throughput schedules when eligible. + + Falls back to :func:`base_fp8_profile`; ``rope_dim`` has the same meaning. + """ + if ( + rope_dim == 64 + and heads == 32 + and query_length <= 8 + and capacity >= 2048 + and 16 <= rows <= sm_count // 8 + ): + # D576: one M64 tile per row with a resident wave of KV splits beats + # the one-stream M16 split below (GB300 top-k 2048, B4 Sq4: 17.2 -> + # 15.6 us). + wave_splits = 1 << ((sm_count // rows).bit_length() - 1) + return SparseMlaProfile( + family="keep", + split_kv=min(wave_splits, capacity // 256), + reuse_kv=True, + reuse_kv_stages=10, + paired_correction=True, + page_pipeline_stages=2, + ) # These short-Q regions favor one KV stream with retained tiles and a # split grid near one device wave. The Q<=8 cutoff is empirical; the # kernel also supports prefill. Powers of two are policy, not a KV ABI. @@ -256,6 +313,20 @@ def _choose_fp8(rows, heads, query_length, capacity, sm_count): # device wave, share KV across H128 even for short selected lists. # BF16's CLC choice below is a separately qualified schedule. return SparseMlaProfile(family="2cta", scheduler="static", gather_issue_warps=1) + if rope_dim == 64 and 8 < heads <= 16 and capacity >= 2048 and rows >= 8 * sm_count: + # Long D576 lists keep M16 Swap-AB on prefill-sized grids: eight + # issuers, twelve retained stages and CLC persistence gather fast + # enough that M64 head padding no longer pays (GB300 top-k 2048, H12 + # prefill 8k: 1637 -> 1374 us). + return SparseMlaProfile( + family="swap", + split_kv=1, + scheduler="clc", + gather_issue_warps=8, + balanced_registers=True, + reuse_kv=True, + reuse_kv_stages=12, + ) if ( heads <= 16 and capacity >= 512 @@ -274,10 +345,10 @@ def _choose_fp8(rows, heads, query_length, capacity, sm_count): uniform_offset_cache=True, scheduler="clc", ) - return _base_fp8(rows, heads, capacity, sm_count) + return base_fp8_profile(rows, heads, capacity, sm_count, rope_dim) -def _choose_bf16(rows, heads, query_length, capacity, sm_count): +def choose_bf16_profile(rows, heads, query_length, capacity, sm_count): """Prefer resident split grids or sustained-throughput schedules when eligible.""" # Small-head workloads can retain one BK128 tile with five/six 32-KiB # stages: four hold D512, the rest prefetch. PV must release that tile @@ -380,10 +451,13 @@ def _choose_bf16(rows, heads, query_length, capacity, sm_count): uniform_offset_cache=heads == 128, scheduler="clc" if heads == 128 and rows >= sm_count else "nonpersistent", ) - return _base_bf16(rows, heads, capacity, sm_count) + return base_bf16_profile(rows, heads, capacity, sm_count) -def _size_buffers(profile, *, rows, heads, query_length, capacity, dtype, sm_count): +def size_profile_buffers( + profile, *, rows, heads, query_length, capacity, dtype, sm_count +): + """Set register-budget and KV/page stage counts for a chosen schedule.""" # These are coupled measured budgets, not independent universal wins: # issuer/cache registers compete with softmax/correction, and extra KV # stages compete with Q, P, metadata and scheduler SMEM. @@ -447,11 +521,153 @@ def _size_buffers(profile, *, rows, heads, query_length, capacity, dtype, sm_cou return profile -def select_sparse_mla_profile(*, rows, heads, query_length, capacity, dtype, sm_count): - """Resolve one complete schedule from the validated public plan bounds.""" - choose = _choose_fp8 if dtype == "fp8" else _choose_bf16 - profile = choose(rows, heads, query_length, capacity, sm_count) - profile = _size_buffers( +QUALIFIED_POLICY_HEADS = (8, 16, 32, 64, 128) + + +def sparse_policy_heads(heads, *, dtype, rope_dim): + """Return the head count whose measured schedule serves ``heads``. + + The crossovers are measured for H8/16/32/64/128. For D576 FP8-MMA plans, + other counts (e.g. 12/24/48/96) use the next qualified count's schedule: + its Q tile already pads that count and the kernel masks the extra rows. + Matching them on their raw count instead falls through to a default M16 + Swap-AB grid that repeats every KV tile once per 16-head tile (three times + for H48, six for H96). + BF16 and D512 plans keep their raw count until measured the same way. + """ + if dtype != "fp8" or rope_dim != 64: + return heads + return next((h for h in QUALIFIED_POLICY_HEADS if heads <= h), heads) + + +def select_sparse_mla_profile( + *, + rows, + heads, + query_length, + capacity, + dtype, + sm_count, + rope_dim=0, + kv_cache_format="native", +): + """Resolve one complete schedule from the validated public plan bounds. + + ``heads`` is first mapped by :func:`sparse_policy_heads`. A mapped + schedule with cluster reduction is replaced by the raw-count schedule when + ``heads`` is not a multiple of 16: the cluster finisher needs whole M16 + head tiles, a kernel constraint rather than a measured crossover. + """ + kwargs = dict( + rows=rows, + query_length=query_length, + capacity=capacity, + dtype=dtype, + sm_count=sm_count, + rope_dim=rope_dim, + kv_cache_format=kv_cache_format, + ) + policy_heads = sparse_policy_heads(heads, dtype=dtype, rope_dim=rope_dim) + profile = resolve_sparse_mla_profile(heads=policy_heads, **kwargs) + if policy_heads != heads and profile.reduction == "cluster" and heads % 16: + profile = resolve_sparse_mla_profile(heads=heads, **kwargs) + return profile + + +def resolve_sparse_mla_profile( + *, + rows, + heads, + query_length, + capacity, + dtype, + sm_count, + rope_dim=0, + kv_cache_format="native", +): + """Choose, size and grid-fill the schedule for exactly ``heads`` heads. + + :func:`select_sparse_mla_profile` is the planning entry point; it maps + unmeasured head counts before calling this. + """ + if ( + kv_cache_format in ("nvfp4", "nvfp4.fp8rope", "fp8") + and dtype == "fp8" + and rope_dim in (0, 64) + and heads <= 32 + and capacity >= 512 + and (rows * 2 >= sm_count or heads >= 16) + ): + # Retain complete converted tiles even when the query grid is small. + # Fill one resident wave with power-of-two KV partitions, preserving + # at least two tiles per split for the eight-issuer steady-state loop. + # Persistent 1CTA queues do not enumerate KV splits, so small grids + # use the same reuse resources with a nonpersistent split schedule. + persistent = rows * 2 >= sm_count + split_limit = min(capacity // 128 // 2, max(1, sm_count // rows)) + splits = 1 if persistent else 1 << (split_limit.bit_length() - 1) + # Above sixteen heads, WS M64 wins on GB300: its N256 PV consumes + # Dv512 in two passes, whereas the narrow Swap-AB path needs four. + keep = heads > 16 + return SparseMlaProfile( + family="keep" if keep else "swap", + tile_size_q=64 if keep else 8 if heads <= 8 else 16, + split_kv=splits, + gather_issue_warps=8, + reuse_kv=True, + whole_kv_tile=True, + single_kv_stream=not keep, + uniform_offset_cache=keep, + paired_correction=keep, + scheduler="clc" if persistent else "nonpersistent", + ) + if dtype == "fp8": + profile = choose_fp8_profile( + rows, heads, query_length, capacity, sm_count, rope_dim + ) + else: + profile = choose_bf16_profile(rows, heads, query_length, capacity, sm_count) + # D576 NVFP4 keeps this schedule down to one row: the general small-grid + # choice partitions V over four CTAs that each repeat the K conversion + # (GB300 H64 top-k 2048, B1 Sq4: 32.3 -> 13.1 us; Sq8: 18.0 -> 14.2 us). + if ( + kv_cache_format != "native" + and dtype == "fp8" + and query_length <= 8 + and capacity >= 2048 + and rows >= (1 if rope_dim == 64 else 16) + and rows * 2 < sm_count + ): + if ( + heads == 64 + or (heads == 128 and kv_cache_format == "fp8" and rows * 4 < sm_count) + # D576 NVFP4 H128 below 16 rows: two whole-tile M64 heads beat the + # 2CTA split (GB300 top-k 2048, H96 B1 Sq8: 27.3 -> 16.8 us). + or (heads == 128 and rope_dim == 64 and rows < 16) + ): + # Small packed grids still benefit from whole-record gathers and + # retained WS tiles. At low row counts, two M64 head tiles also + # beat group-2 mixed-pool startup despite repeating KV conversion. + # Power-of-two splits fill most of a resident wave without adding + # a short split for a source-boundary tile. + work = rows * ceil(heads / 64) + split_limit = min(capacity // 128, max(1, sm_count // work)) + wave_splits = 1 << (split_limit.bit_length() - 1) + profile = SparseMlaProfile( + family="keep", + split_kv=wave_splits, + gather_issue_warps=8, + reuse_kv=True, + reuse_kv_stages=10, + whole_kv_tile=True, + uniform_offset_cache=True, + paired_correction=True, + ) + elif profile.family == "2cta": + # Conversion makes longer per-cluster loops costly. Use the same + # resident-wave split sizing below, with two SMs per cluster. + profile = replace(profile, split_kv=0) + profile = size_profile_buffers( profile, rows=rows, heads=heads, @@ -460,6 +676,32 @@ def select_sparse_mla_profile(*, rows, heads, query_length, capacity, dtype, sm_ dtype=dtype, sm_count=sm_count, ) + if ( + rope_dim + and dtype == "bf16" + and profile.family == "2cta" + and profile.kv_pipeline_stages + ): + # Each CTA's Q RoPE tail needs 64*rope_dim BF16 elements. A combined + # K/V stage holds 64*128, so trade whole prefetch stages for that storage. + profile = replace( + profile, + kv_pipeline_stages=max( + 4, profile.kv_pipeline_stages - ceil(rope_dim / 128) + ), + ) + if rope_dim and profile.reuse_kv: + if dtype == "bf16" and profile.family == "keep": + # A D576 M64 query plus a retained BK128 BF16 tile exceeds SMEM. + # BK64 retains one five-slice K tile plus one prefetched slice; + # its pipeline serializes PV before the next complete QK tile. + profile = replace(profile, kv_tile_size=64, reuse_kv_stages=6) + else: + retained_tiles = 2 if dtype == "fp8" else 1 + profile = replace( + profile, + reuse_kv_stages=max(profile.reuse_kv_stages, retained_tiles * 5), + ) tile = profile.tile_size_q or {"swap": 16, "keep": 64, "2cta": 128}[profile.family] work = rows * ceil(heads / tile) splits = profile.split_kv @@ -486,6 +728,10 @@ def select_sparse_mla_profile(*, rows, heads, query_length, capacity, dtype, sm_ ) # Retaining K through PV requires a single owner of the V dimension. reuse = profile.reuse_kv and v_ctas == 1 + if profile.family == "2cta" and kv_cache_format != "native": + # Convert each CTA's64 K rows once, then exchange the peer's V + # features through DSMEM. Buffering depends on the QK width. + profile = replace(profile, balanced_registers=True) return replace( profile, tile_size_q=tile, @@ -494,3 +740,84 @@ def select_sparse_mla_profile(*, rows, heads, query_length, capacity, dtype, sm_ reuse_kv=reuse, reuse_kv_stages=profile.reuse_kv_stages if reuse else 0, ) + + +# D576 (DSA) sparse MLA: one CTA per (query token, 64 heads) on the 1CTA +# keeps-MMA-AB q64 schedule with four Gather4 issue warps and the uniform quad +# cache (fixed by the D576 1CTA config). Only the split varies with the plan: +# below D576_SPLIT_KV_MIN_CTAS CTAs most SMs would idle, so every token's top-k +# list is split into S contiguous tile ranges (one CTA each) and the 1CTA GMEM +# reducer merges the partials, following trtllm-gen's MLA generation rule. +D576_SPLIT_KV_MIN_CTAS = 96 +D576_SPLIT_KV_CHOICES = (16, 8, 4, 2) + + +def d576_v_split(total_ctas: int, sm_count: int) -> int: + """V head-dim split of a split launch of ``total_ctas`` CTAs: 4 / 2 while the + grid fills at most a quarter / half of the SMs, else 1.""" + if total_ctas * 4 <= sm_count: + return 4 + if total_ctas * 2 <= sm_count: + return 2 + return 1 + + +def select_sparse_mla_d576_split( + num_ctas: int, sparse_mla_top_k: int, sm_count: int +) -> tuple[int, int]: + """Return ``(split_kv, head_dim_split)`` for an E4M3 launch of ``num_ctas`` CTAs. + + trtllm-gen's MLA generation rule: split the top-k list so every CTA keeps + at least two 128-key tiles within one wave (a top-k of one tile pair: + one tile, if the V split then applies), then, while fewer than half the + SMs have work, split the 512 V columns over 4 (or 2) CTAs. Each V + slice repeats the QK / softmax of its tiles, but a one- or two-tile CTA is + bound by its fixed latency (Q load, first gather, the single QK-softmax-PV + chain, the partial-O store), which the slice cuts. B200, K=2048: B1 is + 8 x 4 = 32 CTAs instead of 16 x 1. + """ + if num_ctas >= D576_SPLIT_KV_MIN_CTAS: + return 1, 1 + pairs = -(-sparse_mla_top_k // 256) + limit = min(pairs, sm_count // num_ctas) + split = next((s for s in D576_SPLIT_KV_CHOICES if s <= limit), 1) + if split == 1 and sparse_mla_top_k > 128 and num_ctas * 4 <= sm_count: + # A single tile pair (top-k <= 256): one tile per CTA, so the V split + # below can apply. B200, top-k 256: B1 9.77 -> 7.52 us (trtllm-gen + # 8.13), B32 9.78 -> 9.35; without the V split (B64) one-tile CTAs + # are slower (10.6 -> 12.0), hence the quarter-wave bound. + split = 2 + if split == 1: + return 1, 1 + return split, d576_v_split(num_ctas * split, sm_count) + + +def prefer_sparse_d576_reload( + *, rows, heads, query_length, capacity, sm_count, general_profile +): + """Choose native FP8 reload versus the general retained-KV/2CTA pipeline. + + ``capacity`` is the primary top-k slot count and ``rows`` is B*Q. H128 + benefits from sharing KV in 2CTA beyond the measured four-or-more native + KV-partition region. For H64, small grids and short lists favor early K + release. Longer queries increase within-request reuse; measured GB300 + crossovers keep q<=16 on reload, q<=32 through four SM-sized row groups, + and larger query grids on retained KV. These are schedule crossovers, not + supported-shape limits. ``general_profile`` is the existing alternative: + its narrow-Q/V partitions and cluster finish beat a separate native merge + on short lists. The native reload launch alone uses grid-z per row. + """ + if rows > 65535: + return False + if capacity <= 512 and general_profile.family == "swap": + return False + if heads == 128: + return select_sparse_mla_d576_split(rows * 2, capacity, sm_count)[0] >= 4 + if heads != 64: + return False + return ( + capacity <= 256 + or rows <= sm_count + or query_length <= 16 + or (query_length <= 32 and rows <= 4 * sm_count) + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/config.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/config.py index fe2461ad74c4..3482f2398a1e 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/config.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/config.py @@ -21,7 +21,11 @@ from dataclasses import dataclass from typing import Tuple -from ..helpers.constants import MAX_MLA_SPLITS_KV, SUPPORTED_MLA_PAGE_SIZES +from ..helpers.constants import ( + SUPPORTED_MLA_PAGE_SIZES, + WARP_LANES, +) +from ..helpers.nvfp4_kv import Nvfp4KVConfig, validate_nvfp4_kv_formats from ..helpers.mask import MaskType, normalize_mask_type @@ -92,7 +96,7 @@ def compute_split_kv( split_heur = min(max_splits, blocks_per_batch) k_waves = ceil_div(max_splits, split_heur) split_wave_aware = ceil_div(max_splits, k_waves) - return min(split_wave_aware, MAX_MLA_SPLITS_KV) + return split_wave_aware def compute_workspace_size( @@ -120,7 +124,7 @@ def compute_workspace_size( @dataclass -class MlaDecodeConfig: +class MlaDecodeConfig(Nvfp4KVConfig): """MLA decode kernel configuration. All values are plain Python ints/tuples so they can be used as @@ -267,9 +271,24 @@ class MlaDecodeConfig: # the useful split prefix. Grid and workspace geometry retain the maximum. is_var_split_kv: bool = False cache_uniform_sparse_pages: bool = False + # Storage pairs: native/native; D512 nvfp4/nvfp4 or fp8/nvfp4; + # D576 nvfp4.fp8rope/nvfp4.fp8rope. NVFP4 paths require E4M3 Q/MMA. + # The FP8 primary has a separate tensor scale; NVFP4 uses embedded block + # scales. Both pools retain their physical row strides and share stages. + kv_cache_format: str = "native" + extra_kv_cache_format: str = "native" # Causal is bottom-right aligned for speculative decode. Dense still masks # the ordinary per-batch KV tail at ``cache_seqs[batch]``. mask_type: str = MaskType.CAUSAL.value + # Threshold skip correction in log2 units (binades); 0 disables it. A positive value + # freezes the running row max while a K tile raises it by at most this + # amount. It is a compile-time constant: a runtime scalar stays live in + # every softmax iteration and costs registers under the 128-register cap. + # With E4M3 P it also lowers the P scale from 448 to ``448 * 2**-T``; + # BF16 P stays unscaled. See ``max_softmax_rescale_threshold_log2`` for + # the bounds. + softmax_rescale_threshold_log2: float = 0.0 + lse_base_on_e: bool = False @property def tokens_per_k_tile(self) -> int: @@ -348,9 +367,33 @@ def make_mla_decode_config( cache_uniform_sparse_pages: bool = False, balance_sparse_registers: bool = False, fuse_output_quant: bool = False, + softmax_rescale_threshold_log2: float = 0.0, + kv_cache_format: str = "native", + extra_kv_cache_format: str = "native", + lse_base_on_e: bool = False, ) -> MlaDecodeConfig: """Create and populate a MlaDecodeConfig from problem parameters.""" + validate_nvfp4_kv_formats( + kv_cache_format, extra_kv_cache_format, rope_dim, qkv_dtype, page_size + ) + has_nvfp4_kv = kv_cache_format != "native" + if has_nvfp4_kv and not ( + qkv_dtype == "e4m3" + and page_size == 1 + and rope_dim in (0, 64) + and sparse_gather_warps == 1 + ): + raise ValueError("NVFP4 KV needs sparse E4M3 D512/D576 (internal page size 1)") + if has_nvfp4_kv and ( + mma_qk_tiler_mn != (128, 128) or mma_pv_tiler_mn != (128, 256) + ): + # The loader converts 64 token rows per CTA and exchanges 256-feature + # V halves in fixed SMEM panels. Other tiles need a matching producer + # layout, not just different pipeline byte counts. + raise ValueError("NVFP4 2CTA requires QK (128, 128) and PV (128, 256) tiles") cfg = MlaDecodeConfig() + cfg.kv_cache_format = kv_cache_format + cfg.extra_kv_cache_format = extra_kv_cache_format cfg.mma_qk_tiler_mn = mma_qk_tiler_mn cfg.mma_pv_tiler_mn = mma_pv_tiler_mn cfg.rope_dim = rope_dim @@ -362,6 +405,9 @@ def make_mla_decode_config( cfg.is_var_seq = is_var_seq cfg.is_var_split_kv = is_var_split_kv cfg.mask_type = normalize_mask_type(mask_type) + # Range-checked by MlaDecodeTs and plan(). + cfg.softmax_rescale_threshold_log2 = softmax_rescale_threshold_log2 + cfg.lse_base_on_e = lse_base_on_e def _require_positive(name: str, value: int) -> None: """Validate that a named configuration value is positive.""" @@ -387,8 +433,10 @@ def _require_divisible( if rope_dim < 0: raise ValueError(f"rope_dim must be non-negative, got {rope_dim}") _require_positive("page_size", page_size) + # Page size one is the sparse gather path: D512 (no RoPE) or D576 (RoPE is the + # last 64 channels of each gathered row). if page_size not in SUPPORTED_MLA_PAGE_SIZES and not ( - page_size == 1 and rope_dim == 0 + page_size == 1 and rope_dim in (0, 64) ): raise ValueError( f"page_size must be one of {SUPPORTED_MLA_PAGE_SIZES}, got {page_size}" @@ -424,8 +472,9 @@ def _require_divisible( # third V tile costs 32 KiB per CTA without improving overlap in # the qualified sparse schedule. Keep four metadata slots so # index prefetch can run ahead of both streams. - cfg.load_k_stage = 2 - cfg.load_v_stage = 2 + # Packed rows already have independent prefetch scratch. Two K + # slots let QK advance while the previous V tile is consumed. + cfg.load_k_stage, cfg.load_v_stage = (2, 1) if has_nvfp4_kv else (2, 2) cfg.use_fp8_dual_softmax_schedule = False cfg.second_compute_warp_ids = () cfg.num_softmax_groups = 1 @@ -435,6 +484,22 @@ def _require_divisible( cfg.threads_per_cta = cfg.threads_per_warp * 20 cfg.other_reg_num = 48 + if has_nvfp4_kv: + # Convert each CTA's 64 rows once. D512 sends V directly from registers, + # allowing three V stages against two K stages; D576 bulk exchange + # retains the K source, so K and V lifetimes use paired two-stage slots. + # One S stage leaves the TMEM columns needed to retain FP8 NoPE Q. + # Fused FP8 output needs a 32-KiB shared output stage. Pair K/V at + # two stages for that path to fit the SMEM budget; equal stage counts + # select bulk DSMEM exchange and retain K until both V readers finish. + cfg.load_k_stage, cfg.load_v_stage = ( + 2, + 2 if rope_dim or fuse_output_quant else 3, + ) + cfg.load_num_warps = 8 + cfg.mma_s_stage = 1 + cfg.load_v_warp_id = cfg.load_k_warp_id + if sparse_gather_warps != 1: if sparse_gather_warps not in (4, 8) or qkv_dtype != "bf16" or page_size != 1: raise ValueError("multi-warp combined loading requires BF16 sparse 2CTA") @@ -464,14 +529,19 @@ def _require_divisible( if balance_sparse_registers: if page_size != 1 or qkv_dtype != "e4m3": raise ValueError("balanced registers require sparse FP8") - # Sparse FP8 keeps 64 registers for the gather issuers. Split the - # remaining budget equally: 128 softmax registers spilled live loop - # state on every KV tile, while correction fits in 144 registers. - # 4*144 + 4*144 + 12*64 = 1920 warp-registers, the same CTA total - # as the previous 128/160/64 allocation (61440 registers). + # D512's register DSMEM stores benefit from a larger loader budget; + # D576's bulk copies retain more registers for softmax/correction. + # Base splits claim 61,440 registers in the 640-thread CTA. cfg.softmax_reg_num = 144 cfg.correction_reg_num = 144 cfg.other_reg_num = 64 + if has_nvfp4_kv and not rope_dim: + # Tensor-scale score/correction state needs more softmax registers. + # Borrow eight from the other warps; stay within the CTA's initial + # allocation (60,416 registers versus its 61,440-register credit). + cfg.softmax_reg_num = 112 if cfg.has_fp8_primary else 96 + cfg.correction_reg_num = 96 + cfg.other_reg_num = 88 if cfg.has_fp8_primary else 96 if fuse_output_quant and cfg.is_fp8_qkv() and cfg.threads_per_cta == 640: # Quantization keeps an output block and rotary coefficients live. @@ -481,6 +551,10 @@ def _require_divisible( cfg.correction_reg_num = 176 cfg.other_reg_num = 48 + if cfg.has_nvfp4_kv: + cfg.p_mma_stage = 2 if rope_dim else 1 + assert cfg.sparse_offset_stages >= 2 + # SETMAXNREG redistributes the CTA's initial allocation, rounded down # to eight registers/thread; the remainder of the SM file is not credit. initial_regs = 65536 // (cfg.threads_per_cta * 8) * 8 @@ -570,6 +644,13 @@ def _require_divisible( "one full-tile page: " f"tokens_per_k_cta={cfg.tokens_per_k_cta}, page_size={cfg.page_size}" ) + if cfg.page_size > 1 and WARP_LANES % cfg.pages_per_k_tile != 0: + # The BF16 TMA warp and the FP8 K/V loaders read dense page IDs from a + # one-entry-per-lane window that must hold whole K tiles. + raise ValueError( + f"dense MLA needs pages_per_k_tile to divide {WARP_LANES}: " + f"pages_per_k_tile={cfg.pages_per_k_tile}" + ) _require_divisible( "mma_pv_tiler_mn[0]", cfg.mma_pv_tiler[0], "num_mma_ctas", num_mma_ctas ) @@ -679,7 +760,6 @@ def _require_divisible( cfg.tma_copy_v_tile_bytes = ( cfg.tma_copy_vc_bytes * cfg.iterations_pv_k * cfg.iterations_pv_n ) - # TMEM sync barrier thread count: QK MMA + softmax + correction, plus # the FP8 PV-MMA warp when it participates in TMEM O production. cfg.tmem_sync_bar_threads = ( diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/kernel.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/kernel.py index ffc5c7fb609c..de83adb98a80 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/kernel.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/kernel.py @@ -43,7 +43,7 @@ TileSchedulerConfig, ) from cutlass.experimental.task_scheduling.enums import SignalingThreads -from cutlass.experimental.task_scheduling.task_manager import TaskManager +from ..helpers.diagnostics import MlaTaskManager as TaskManager from ...tensor_map import ( create_tensor_map_ragged_from_tensor, create_tensor_map_tiled_from_view, @@ -58,7 +58,7 @@ MlaDecodeConfig, make_mla_decode_config, ) -from ..helpers.constants import LOG2_E, MAX_MLA_SPLITS_KV, TMEM_DEALLOC_MBAR_THREADS +from ..helpers.constants import LOG2_E, TMEM_DEALLOC_MBAR_THREADS from ..helpers.mask import MaskType, mask_visible_k_length, normalize_mask_type from ..helpers.query import ( FlatQueryTileLayout, @@ -84,7 +84,14 @@ MlaWorkQueue, WorkThrottleBarrierResource, ) -from ..helpers.math import qkv_dtype +from ..helpers.math import qkv_dtype, canonicalize_softmax_rescale_threshold_log2 +from .nvfp4_kv_load import nvfp4_kv_scratch_bytes +from ..helpers.gather import make_sparse_rope_map +from ..helpers.nvfp4_kv import ( + nvfp4_row_bytes, + Nvfp4KVConfig, + make_whole_record_map, +) from ..parallel_reduction_topology import ( ParallelReductionTopology, make_balanced_parallel_reduction_topology, @@ -109,6 +116,7 @@ MlaInterleavedTask, MlaTask, create_load_k_task, + create_load_combined_task, create_sparse_offsets_task, create_load_v_task, create_load_tma_task, @@ -137,20 +145,27 @@ def build_mla_decode_task_manager( smem_vc_arr=None, smem_p_arr=None, sparse_offsets_smem=None, + smem_raw_arr=None, + raw_mbar_arr=None, # TMA descriptors (None for validate-only) tma_desc_q_latent=None, tma_desc_q_rope=None, tma_desc_c_latent=None, tma_desc_c_rope=None, tma_desc_c_transpose=None, + tma_desc_extra_v=None, + tma_desc_sparse_rope=None, + tma_desc_sparse_extra_rope=None, # Page-offset tensor + primary_cache=None, + extra_cache=None, page_offsets=None, # Runtime coordinates blk_coord=None, tidx=None, # GMEM output tensors output=None, - atten_sinks=None, + attention_sinks=None, lse_in_natural_log=False, output_quant=None, output_quant_args=None, @@ -171,14 +186,19 @@ def build_mla_decode_task_manager( logical_seq_len_q=1, tiled_mma_qk=None, verbose=False, - exhaustive_deadlock_race_check=True, + exhaustive_deadlock_race_check: bool | None = None, ) -> "tuple[TaskManager, list[MemoryResource], dict[str, MemoryResource]]": """Build the MLA decode TaskManager with all resources, tasks, and dependency graph. - The throughput 2CTA graph uses Q/K/V TMA loads, a register-held BF16 page-ID - window, QK/PV MMA, softmax, and correction/store tasks. Validation-only + The throughput 2CTA graph uses Q/K/V TMA loads, register-held page-ID + windows (BF16: refreshed at each window boundary; dense FP8 K/V loaders: + the next window is also prefetched), QK/PV MMA, softmax, and + correction/store tasks. Validation-only calls pass a concrete integer ``domain`` and no runtime tensors; JIT calls - pass symbolic runtime state and skip schedule validation. The dependency + pass symbolic runtime state and retain structural checks as warnings. + Exhaustive interleaving checks default to validation-only calls, whose + concrete domains permit meaningful simulation; an explicit Boolean + overrides that default. The dependency graph relies on TaskManager DMA-order validation to keep SMEM producers alive until async TMEM consumers have launched. @@ -376,11 +396,17 @@ def build_mla_decode_task_manager( producer_group=pipeline.CooperativeGroup(Agent.Thread, 32), consumer_group=pipeline.CooperativeGroup( Agent.Thread, - (2 * cfg.load_num_warps + (4 if shared_mask_enabled else 0)) * 32, + ( + (1 if cfg.has_nvfp4_kv else 2) * cfg.load_num_warps + + (4 if shared_mask_enabled else 0) + ) + * 32, ), cta_layout_vmnk=cluster_shape_vmnk, ) page_offset_window = PageOffsetWindowResource( + primary_cache=primary_cache, + extra_cache=extra_cache, page_offsets=page_offsets, smem_sparse_offsets=sparse_offsets_smem, cfg=cfg, @@ -403,11 +429,17 @@ def build_mla_decode_task_manager( if cfg.use_fp8_split_mma_schedule: smem_k = SmemKResource( + # NVFP4 loaders cache graph-live source lengths; dense page IDs + # are supplied by the independent prefetch window instead. + page_offsets=page_offsets if cfg.has_nvfp4_kv else None, smem_sparse_offsets=sparse_offsets_smem, smem_k=smem_kc_arr, - page_offsets=page_offsets, tma_desc_c_latent=tma_desc_c_latent, tma_desc_c_rope=tma_desc_c_rope, + tma_desc_sparse_rope=tma_desc_sparse_rope, + tma_desc_sparse_extra_rope=tma_desc_sparse_extra_rope, + smem_raw=smem_raw_arr, + raw_mbar=raw_mbar_arr, logical_seq_len_q=logical_seq_len_q, cfg=cfg, pipeline_config=smem_k_pipeline_cfg, @@ -416,9 +448,12 @@ def build_mla_decode_task_manager( smem_v = SmemVResource( smem_sparse_offsets=sparse_offsets_smem, smem_v=smem_vc_arr, - page_offsets=page_offsets, tma_desc_c_transpose=tma_desc_c_transpose, - tma_desc_extra_v=tma_desc_q_rope, + tma_desc_extra_v=tma_desc_extra_v, + tma_desc_sparse_rope=tma_desc_sparse_rope, + tma_desc_sparse_extra_rope=tma_desc_sparse_extra_rope, + smem_raw=smem_raw_arr, + raw_mbar=raw_mbar_arr, logical_seq_len_q=logical_seq_len_q, cfg=cfg, pipeline_config=smem_v_pipeline_cfg, @@ -430,8 +465,10 @@ def build_mla_decode_task_manager( smem_kv=smem_kc_arr, tma_desc_c_latent=tma_desc_c_latent, tma_desc_c_rope=tma_desc_c_rope, + tma_desc_sparse_rope=tma_desc_sparse_rope, + tma_desc_sparse_extra_rope=tma_desc_sparse_extra_rope, tma_desc_c_transpose=tma_desc_c_transpose, - tma_desc_extra_v=tma_desc_q_rope, + tma_desc_extra_v=tma_desc_extra_v, cfg=cfg, pipeline_config=smem_kv_pipeline_cfg, name="smem_kv", @@ -445,7 +482,7 @@ def build_mla_decode_task_manager( smem_q_rope=smem_q_rope_arr, smem_p=smem_p_arr, smem_exchange=None, # Set at runtime - softmax_scale_log2=None, # Set at runtime + scale_softmax_log2=None, # Set at runtime cache_seqs=cache_seqs, cu_seqlens_q=cu_seqlens_q, split_kv=split_kv, @@ -493,14 +530,14 @@ def build_mla_decode_task_manager( partial_lse=acc_lse, tmem_o_ref=tmem_o, tmem_corr_ref=tmem_corr, - atten_sinks=atten_sinks, + attention_sinks=attention_sinks, lse_in_natural_log=lse_in_natural_log, output_quant=output_quant, output_quant_args=output_quant_args, smem_o=smem_o, tma_o_desc=tma_o_desc, output_scale=None, # set at runtime - softmax_scale_log2=None, # set at runtime + scale_softmax_log2=None, # set at runtime smem_exchange=None, # set at runtime split_kv=None, # set at runtime cu_seqlens_q=cu_seqlens_q, @@ -519,22 +556,46 @@ def build_mla_decode_task_manager( # register-budget check accounts for the real MLA warpgroup layout. wg2_reg_count = cfg.other_reg_num + # Dense FP8 K/V loaders take page IDs from a prefetched register window + # (one table entry per lane, 32 per warp) instead of reading GMEM after + # each slot wait; make_mla_decode_config checks the window geometry. + dense_page_window = ( + page_offset_window + if cfg.use_fp8_split_mma_schedule and not sparse_offsets_enabled + else None + ) if cfg.use_fp8_split_mma_schedule: - load_k_task = create_load_k_task( - smem_q, - smem_k, - page_offsets=page_offset_window if sparse_offsets_enabled else None, - work_queue=work_queue, - domain=domain, - num_registers=wg2_reg_count, - ) - load_v_task = create_load_v_task( - smem_v, - page_offsets=page_offset_window if sparse_offsets_enabled else None, - work_queue=work_queue, - domain=domain, - num_registers=wg2_reg_count, - ) + if cfg.has_nvfp4_kv: + smem_k.paired_v = smem_v + smem_k.page_window = page_offset_window + load_k_task = create_load_combined_task( + smem_q, + smem_k, + smem_v, + page_offset_window, + work_queue=work_queue, + domain=domain, + num_registers=wg2_reg_count, + ) + load_v_task = None + else: + load_k_task = create_load_k_task( + smem_q, + smem_k, + page_offsets=page_offset_window if sparse_offsets_enabled else None, + page_window=dense_page_window, + work_queue=work_queue, + domain=domain, + num_registers=wg2_reg_count, + ) + load_v_task = create_load_v_task( + smem_v, + page_offsets=page_offset_window if sparse_offsets_enabled else None, + page_window=dense_page_window, + work_queue=work_queue, + domain=domain, + num_registers=wg2_reg_count, + ) mma_task = None mma_qk_task = create_mma_qk_direct_task( smem_q, @@ -641,10 +702,9 @@ def build_mla_decode_task_manager( ) if cfg.use_fp8_split_mma_schedule: - task_list = [ - load_k_task, - load_v_task, - ] + task_list = [load_k_task] + if load_v_task is not None: + task_list.append(load_v_task) task_list.extend([mma_qk_task, softmax_task]) if second_softmax_task is not None: task_list.append(second_softmax_task) @@ -718,8 +778,8 @@ def build_mla_decode_task_manager( if cfg.use_fp8_split_mma_schedule: resource_dependency_graph = { smem_q: [], # Q loads (independent) - smem_k: [], # K loads read page offsets directly - smem_v: [], # V loads read page offsets directly + smem_k: [], # K loads take page IDs from the window or offset ring + smem_v: [], # V loads take page IDs from the window or offset ring tmem_s: [smem_k, smem_q], # QK MMA needs K and Q smem_p: [tmem_s], # softmax reads S -> writes P tmem_corr: [tmem_s], # softmax produces correction @@ -772,11 +832,18 @@ def build_mla_decode_task_manager( resource_dependency_graph[page_offset_window] = ( [work_queue] if work_queue is not None else [] ) + elif dense_page_window is not None: + resource_dependency_graph[page_offset_window] = [] # register-held window # ────────────────────────────────────────────────────────────── # Create TaskManager # ────────────────────────────────────────────────────────────── skip = not isinstance(domain, int) + # skip_validation only downgrades failures to warnings in the DSL. Avoid + # repeating the expensive interleaving search for every JIT specialization + # with a symbolic domain, which the checker approximates as one iteration. + if exhaustive_deadlock_race_check is None: + exhaustive_deadlock_race_check = not skip task_manager = TaskManager( tasks=task_list, resource_dependency_graph=resource_dependency_graph, @@ -788,6 +855,7 @@ def build_mla_decode_task_manager( tmem_resources = [tmem_s, tmem_o, tmem_corr] named_resources = { + "smem_q": smem_q, "tmem_s": tmem_s, "tmem_o": tmem_o, "tmem_corr": tmem_corr, @@ -800,7 +868,7 @@ def build_mla_decode_task_manager( # ===================================================================== -class MlaDecodeTs: +class MlaDecodeTs(Nvfp4KVConfig): """Warp-specialised MLA decode kernel using the TS framework. Usage:: @@ -811,6 +879,8 @@ class MlaDecodeTs: block_split_kvs, softmax_scale, output_scale, stream) """ + lse_base_on_e = False + def __init__( self, acc_dtype=None, @@ -823,6 +893,7 @@ def __init__( is_var_seq=False, is_var_split_kv=False, static_split_kv=None, + split_kv_capacity: int | None = None, static_seq_len_k=None, qkv_dtype="bf16", out_dtype="bf16", @@ -831,6 +902,8 @@ def __init__( seq_len_q=1, batch_size=1, mask_type: MaskType | str = MaskType.CAUSAL, + softmax_rescale_threshold_log2: float = 0.0, + fixed_seqlen: bool = False, ): """ Parameters @@ -878,6 +951,12 @@ def __init__( mask_type : MaskType or str, optional ``causal`` (default) for bottom-right speculative decoding or ``dense`` for full per-batch KV visibility. + softmax_rescale_threshold_log2 : float, optional + Compile-time skip-correction threshold in log2 units. ``0.0`` + (default) keeps the exact online-softmax rescale. + fixed_seqlen : bool, optional + ``False`` (default) makes the query lengths runtime launch + operands; ``True`` compiles for the exact static ``seq_len_q``. """ import cutlass as _cutlass @@ -897,22 +976,39 @@ def __init__( self.is_var_split_kv = is_var_split_kv self.static_split_kv = static_split_kv self.reduction_split_capacity = ( - static_split_kv if static_split_kv is not None else MAX_MLA_SPLITS_KV + split_kv_capacity + if split_kv_capacity is not None + else static_split_kv + if static_split_kv is not None + else max_active_clusters * 2 ) - if not 1 <= self.reduction_split_capacity <= MAX_MLA_SPLITS_KV: - raise ValueError(f"static_split_kv must be in [1, {MAX_MLA_SPLITS_KV}]") + if self.reduction_split_capacity < 1: + raise ValueError("split_kv_capacity must be positive") + if ( + static_split_kv is not None + and not 1 <= static_split_kv <= self.reduction_split_capacity + ): + raise ValueError("static_split_kv exceeds split_kv_capacity") self.static_seq_len_k = static_seq_len_k self.qkv_dtype = qkv_dtype self.out_dtype = out_dtype self.rope_dim = rope_dim self.num_heads = num_heads self.seq_len_q = seq_len_q + self.fixed_seqlen = fixed_seqlen if isinstance(batch_size, bool) or not isinstance(batch_size, int): raise TypeError("batch_size must be an integer") if batch_size <= 0: raise ValueError("batch_size must be positive") self.batch_size = batch_size self.mask_type = normalize_mask_type(mask_type) + # Validate eagerly: the config (and its own check) is only built when + # the kernel is traced, after plan() has bound a workspace. + self.softmax_rescale_threshold_log2 = ( + canonicalize_softmax_rescale_threshold_log2( + softmax_rescale_threshold_log2, qkv_dtype=qkv_dtype + ) + ) self.fuse_sparse_reduction = False self.fuse_sparse_epilogue = False self.output_quant = None @@ -923,8 +1019,9 @@ def __init__( self.sparse_kv_stages = 0 self.sparse_uniform_pages = False self.sparse_balanced_registers = False + self.kv_cache_format = "native" + self.extra_kv_cache_format = "native" self.assume_valid_prefix = False - self.direct_static_scales = False self.direct_sparse_capacities = (0, 0) self.query_tile_layout = FlatQueryTileLayout.for_tile( num_heads, seq_len_q, mma_qk_tiler_mn[0] @@ -935,8 +1032,7 @@ def __init__( self.use_parallel_reduction = False self._parallel_reduction_shape_is_eligible = ( not is_var_split_kv - and static_split_kv is not None - and 2 <= static_split_kv <= 128 + and self.reduction_split_capacity >= 2 and mma_qk_tiler_mn[0] == 128 and acc_dtype == _cutlass.Float32 and lse_dtype == _cutlass.Float32 @@ -962,6 +1058,7 @@ def compile_signature(self) -> tuple[object, ...]: return ( self.acc_dtype, self.lse_dtype, + self.lse_base_on_e, self.mma_qk_tiler_mn, self.mma_pv_tiler_mn, self.max_active_clusters, @@ -975,29 +1072,34 @@ def compile_signature(self) -> tuple[object, ...]: self.out_dtype, self.rope_dim, self.num_heads, - self.seq_len_q, + self.seq_len_q if self.fixed_seqlen else None, self.mask_type, + self.softmax_rescale_threshold_log2, self.reduction_split_capacity, - self.query_tile_layout, - self.num_q_tiles, - self.tail_q_rows, + self.query_tile_layout if self.fixed_seqlen else None, + self.num_q_tiles if self.fixed_seqlen else None, + self.tail_q_rows if self.fixed_seqlen else None, self._parallel_reduction_shape_is_eligible, self.use_parallel_reduction, self._reduction_topology(), self.fuse_sparse_reduction, self.fuse_sparse_epilogue, self.output_quant, + self.fixed_seqlen, self.direct_sparse, self.assume_valid_prefix, - self.direct_static_scales, - self.direct_sparse_capacities, + tuple(int(x > 0) for x in self.direct_sparse_capacities), self.sparse_gather_warps, self.sparse_kv_stages, self.sparse_uniform_pages, self.sparse_balanced_registers, + self.kv_cache_format, + self.extra_kv_cache_format, ) - def _configure_parallel_reduction_topology(self) -> None: + def _configure_parallel_reduction_topology( + self, policy_splits: int | None = None + ) -> None: """Refresh reducer topology after any host-side launch-shape update.""" self.parallel_reduction_topology = None @@ -1017,9 +1119,11 @@ def _configure_parallel_reduction_topology(self) -> None: if not self._parallel_reduction_shape_is_eligible: return - assert self.static_split_kv is not None + policy_splits = ( + self.reduction_split_capacity if policy_splits is None else policy_splits + ) topology = make_q128_wave_limited_parallel_reduction_topology( - self.static_split_kv, + policy_splits, logical_rows=(self.num_heads * self.seq_len_q * self.batch_size), physical_sm_count=self.max_active_clusters * 2, max_cluster_size=8, @@ -1032,27 +1136,27 @@ def _configure_parallel_reduction_topology(self) -> None: # physical row. Padded M128 tails and intermediate split counts retain # the compact reference grid; high split counts may use clusters. use_small_split_g1 = ( - self.static_split_kv <= 16 + policy_splits <= 16 and topology is not None and topology.cluster_size == 1 and parallel_g1_grid_has_no_padded_rows and should_use_q128_g1_parallel_reducer( batch_size=self.batch_size, physical_rows_per_batch=(physical_tile_rows * num_query_tiles), - producer_ctas=( - self.batch_size * num_query_tiles * self.static_split_kv * 2 - ), + producer_ctas=(self.batch_size * num_query_tiles * policy_splits * 2), reference_rows_per_cta=REDUCTION_ROWS_PER_CTA, physical_sm_count=self.max_active_clusters * 2, ) ) use_high_split_cluster = ( - self.static_split_kv > 32 - and topology is not None - and topology.cluster_size > 1 + policy_splits > 32 and topology is not None and topology.cluster_size > 1 ) - if use_small_split_g1 or use_high_split_cluster: - self.parallel_reduction_topology = topology + if topology is not None and (use_small_split_g1 or use_high_split_cluster): + self.parallel_reduction_topology = ( + make_balanced_parallel_reduction_topology( + self.reduction_split_capacity, cluster_size=topology.cluster_size + ) + ) self.use_parallel_reduction = True def _make_config(self): @@ -1069,12 +1173,16 @@ def _make_config(self): is_var_seq=self.is_var_seq, is_var_split_kv=self.is_var_split_kv, mask_type=self.mask_type, + lse_base_on_e=self.lse_base_on_e, sparse_gather_warps=self.sparse_gather_warps, sparse_kv_stages=self.sparse_kv_stages, cache_uniform_sparse_pages=self.sparse_uniform_pages, balance_sparse_registers=self.sparse_balanced_registers, fuse_output_quant=self.output_quant is not None and self.fuse_sparse_epilogue, + softmax_rescale_threshold_log2=self.softmax_rescale_threshold_log2, + kv_cache_format=self.kv_cache_format, + extra_kv_cache_format=self.extra_kv_cache_format, ) # One flattened sparse query supplies a complete M128 tile. Its two @@ -1112,8 +1220,10 @@ def __call__( output_scale: cutlass.Float32, stream: object, output_quant_args=None, + max_q_len: cutlass.Int32 = 0, ): """Execute the MLA decode TS kernel.""" + seq_len_q = self.seq_len_q if self.fixed_seqlen else max_q_len if cutlass.const_expr(self.fuse_sparse_reduction): if cutlass.const_expr( self.is_persistent @@ -1134,15 +1244,17 @@ def __call__( not (self.fuse_sparse_reduction or self.fuse_sparse_epilogue) or self.qkv_dtype not in ("e4m3", "bf16") or self.page_size != 1 - or self.rope_dim != 0 + or self.rope_dim not in (0, 64) or self.seq_len_q != 1 ): raise ValueError( - "direct sparse 2CTA requires native FP8/BF16 D512 and fused output" + "direct sparse 2CTA requires native FP8/BF16 D512 or D576 and fused output" ) cfg = self._make_config() physical_tile_rows = self.mma_qk_tiler_mn[0] - num_query_tiles = self.num_q_tiles + num_query_tiles = ( + seq_len_q * self.num_heads + self.mma_qk_tiler_mn[0] - 1 + ) // self.mma_qk_tiler_mn[0] # Fixed public tensors retain [H,D,SQ,B]/[H,SQ,B]; variable-Q tensors # compact batches into [H,D,totalQ]/[H,totalQ]. The scheduler/workspace # use physical flat-query tile coordinates, while resources map @@ -1185,19 +1297,20 @@ def __call__( o.stride[3] == o.shape[2] * o.stride[2], "o must be compact from the query axis into the batch axis", ) - runtime_assert( - lse.stride[0] == 1, - "lse must have a contiguous head axis", - ) - runtime_assert( - lse.stride[1] == lse.shape[0] * lse.stride[0], - "lse must be compact from the head axis into the query axis", - ) - if cutlass.const_expr(cu_seqlens_q is None): + if cutlass.const_expr(lse is not None): runtime_assert( - lse.stride[2] == lse.shape[1] * lse.stride[1], - "lse must be compact from the query axis into the batch axis", + lse.stride[0] == 1, + "lse must have a contiguous head axis", ) + runtime_assert( + lse.stride[1] == lse.shape[0] * lse.stride[0], + "lse must be compact from the head axis into the query axis", + ) + if cutlass.const_expr(cu_seqlens_q is None): + runtime_assert( + lse.stride[2] == lse.shape[1] * lse.stride[1], + "lse must be compact from the query axis into the batch axis", + ) def _flatten_query_rows_for_tma(t): """Expose logical ``(SQ, H)`` as one bounded TMA row dimension.""" @@ -1228,6 +1341,7 @@ def _flatten_query_rows_for_tma(t): q_latent_tma, box_dims=(cfg.mma_qk_tiler[2], cfg.mma_qk_tiler[0] // 2), ragged_dim=1, + use_blackwell_tma_oob_workaround=True, stride_order=(0, 1), swizzle=cuda.TensorMapSwizzle.s128b, l2_promotion=cuda.TensorMapL2Promotion.l2_128b, @@ -1254,6 +1368,7 @@ def _flatten_query_rows_for_tma(t): cfg.mma_qk_rope_tiler[0] // 2, ), ragged_dim=1, + use_blackwell_tma_oob_workaround=True, stride_order=(0, 1), swizzle=q_rope_swizzle, l2_promotion=cuda.TensorMapL2Promotion.l2_128b, @@ -1277,7 +1392,52 @@ def _flatten_query_rows_for_tma(t): c_latent.iterator, cute.select(c_latent.layout, mode=[1, 0, 2]), ) - if cutlass.const_expr(cfg.page_size == 1): + if cutlass.const_expr(cfg.has_nvfp4_kv): + # Byte views are [1, record_bytes, physical_rows], preserving each + # pool's page stride. Each logical token is contiguous: + # D512 NVFP4: [256 B nibbles | 32 B block16 scales] + # D576 NVFP4: [256 B nibbles | 64 B FP8 RoPE | 32 B scales] + # FP8 SWA: [512 B E4M3], primary only; scale is a run input. + # NVFP4 whole-record gathers feed conversion scratch. FP8 panels + # land directly in K's MMA swizzle. Source tags select primary or + # extra maps through the existing K/tail/V descriptor ABI slots. + # Full offsets and nibble/scale ordering: helpers/nvfp4_kv.py. + maps = () + for pool, fmt in ( + (c_latent, cfg.kv_cache_format), + (c_rope, cfg.extra_kv_cache_format), + ): + if cutlass.const_expr(fmt != "fp8"): + whole_map = make_whole_record_map( + pool, nvfp4_row_bytes(cfg.rope_dim) + ) + maps += (whole_map, whole_map, whole_map) + else: + pieces = ( + (0, 512, 128, cuda.TensorMapSwizzle.s128b), + (0, 16, 32, cuda.TensorMapSwizzle.s32b), + (0, 512, 128, cuda.TensorMapSwizzle.s128b), + ) + for offset, width, box, swizzle in pieces: + view = cute.make_tensor( + pool.iterator + offset, + cute.make_layout( + (width, pool.shape[2]), stride=(1, pool.stride[2]) + ), + ) + maps += ( + create_tensor_map_tiled_from_view( + view, + box_dims=(box, 1), + stride_order=(0, 1), + swizzle=swizzle, + l2_promotion=cuda.TensorMapL2Promotion.l2_128b, + ), + ) + tma_desc_c_latent, tma_desc_c_rope = maps[0], maps[3] + tma_desc_sparse_rope, tma_desc_sparse_extra_rope = maps[1], maps[4] + tma_desc_c_transpose, tma_desc_extra_v = maps[2], maps[5] + elif cutlass.const_expr(cfg.page_size == 1): # D512 has no separate Q/K RoPE maps. Reuse their ABI slots for # extra-K and extra-V descriptors, keeping dense launch ABI stable. primary_rows = cute.make_tensor( @@ -1317,7 +1477,7 @@ def _flatten_query_rows_for_tma(t): swizzle=v_swizzle, l2_promotion=cuda.TensorMapL2Promotion.l2_128b, ) - tma_desc_q_rope = create_tensor_map_tiled_from_view( + tma_desc_extra_v = create_tensor_map_tiled_from_view( extra_rows, box_dims=(sparse_v_width, 1), stride_order=(0, 1), @@ -1368,6 +1528,19 @@ def _flatten_query_rows_for_tma(t): l2_promotion=cuda.TensorMapL2Promotion.l2_128b, ) + if cutlass.const_expr(cfg.page_size != 1): + tma_desc_extra_v = tma_desc_c_transpose + if cutlass.const_expr(not cfg.has_nvfp4_kv): + tma_desc_sparse_rope = tma_desc_c_latent + tma_desc_sparse_extra_rope = tma_desc_c_rope + if cutlass.const_expr( + cfg.page_size == 1 and cfg.rope_dim and not cfg.has_nvfp4_kv + ): + tma_desc_sparse_rope = make_sparse_rope_map(c_latent, cfg.qkv_dtype_bytes) + tma_desc_sparse_extra_rope = make_sparse_rope_map( + c_rope, cfg.qkv_dtype_bytes + ) + tma_desc_o = tma_desc_q_latent if cutlass.const_expr(cfg.use_tma_output): flat_o = output_quant_args[0] @@ -1399,7 +1572,7 @@ def _flatten_query_rows_for_tma(t): l2_promotion=cuda.TensorMapL2Promotion.l2_128b, ) - softmax_scale_log2 = softmax_scale * LOG2_E + scale_softmax_log2 = softmax_scale * LOG2_E kernel_split_kv = ( Int32(self.static_split_kv) @@ -1407,6 +1580,12 @@ def _flatten_query_rows_for_tma(t): else split_kv ) + runtime_assert(kernel_split_kv >= Int32(1), "split count must be positive") + runtime_assert( + kernel_split_kv <= Int32(self.reduction_split_capacity), + "split count exceeds prepared capacity", + ) + # Compute grid tile_sched_params = create_mla_static_tile_scheduler_params( self.is_persistent, @@ -1432,8 +1611,19 @@ def _flatten_query_rows_for_tma(t): ) grid = clc_tile_sched_params.get_grid_shape() else: + launch_clusters = Int32(self.max_active_clusters) + if cutlass.const_expr( + self.direct_sparse and self.has_nvfp4_kv and self.is_persistent + ): + # Equalize the persistent waves using the live host-side + # tensor extent. Baking this row-dependent cap into self + # creates a new JIT key for many otherwise identical graph + # buckets. Only the physical SM limit remains static. + work_clusters = batch_size * Int32(num_query_tiles) * kernel_split_kv + waves = (work_clusters + launch_clusters - Int32(1)) // launch_clusters + launch_clusters = (work_clusters + waves - Int32(1)) // waves grid = MLAStaticTileScheduler.get_grid_shape( - tile_sched_params, self.max_active_clusters + tile_sched_params, launch_clusters ) # Initialize workspace for split_kv > 1 @@ -1467,12 +1657,16 @@ def _flatten_query_rows_for_tma(t): cache_seqs, cu_seqlens_q, block_split_kvs, - softmax_scale_log2, + scale_softmax_log2, output_scale, tile_sched_params, clc_tile_sched_params, output_quant_args, tma_desc_o, + tma_desc_extra_v, + tma_desc_sparse_rope, + tma_desc_sparse_extra_rope, + max_q_len, ).launch( grid=grid, block=[cfg.threads_per_cta, 1, 1], @@ -1498,6 +1692,10 @@ def _flatten_query_rows_for_tma(t): cu_seqlens_q, block_split_kvs, output_quant_args, + split_divisor=tile_sched_params.split_kv_fdd + if self.static_split_kv is None and not self.is_var_split_kv + else None, + max_q_len=max_q_len, ).launch( grid=( ( @@ -1518,33 +1716,152 @@ def _flatten_query_rows_for_tma(t): use_pdl=use_one_wave_reducer_pdl, ) else: - logical_query_rows = self.num_heads * self.seq_len_q - self.reduction_kernel( - o, - lse, - acc_o, - acc_lse, - kernel_split_kv, - cache_seqs, - cu_seqlens_q, - block_split_kvs, - ).launch( - grid=( - ceil_div(logical_query_rows, REDUCTION_ROWS_PER_CTA), - 1, - batch_size, - ), - block=[REDUCTION_ROWS_PER_CTA * REDUCTION_THREADS_PER_ROW, 1, 1], - smem=( - REDUCTION_ROWS_PER_CTA - * self.reduction_split_capacity - * self.lse_dtype.width - // 8 - ), - stream=stream, - min_blocks_per_mp=2, - use_pdl=use_one_wave_reducer_pdl, - ) + logical_query_rows = self.num_heads * seq_len_q + if cutlass.const_expr( + self.reduction_split_capacity > 4 + and self.static_split_kv is None + and not self.is_var_split_kv + ): + # Allocation capacity is independent of the active split + # prefix. Keep the existing S2/S4 reducer schedules and their + # small scale rings; all three entrypoints live + # in one cached module and retain runtime workspace strides. + if kernel_split_kv <= Int32(2): + self.reduction_kernel( + o, + lse, + acc_o, + acc_lse, + kernel_split_kv, + cache_seqs, + cu_seqlens_q, + block_split_kvs, + split_divisor=tile_sched_params.split_kv_fdd + if self.static_split_kv is None and not self.is_var_split_kv + else None, + max_q_len=max_q_len, + local_split_capacity=2, + ).launch( + grid=( + ceil_div(logical_query_rows, REDUCTION_ROWS_PER_CTA), + 1, + batch_size, + ), + block=[ + REDUCTION_ROWS_PER_CTA * REDUCTION_THREADS_PER_ROW, + 1, + 1, + ], + smem=( + REDUCTION_ROWS_PER_CTA * 2 * self.lse_dtype.width // 8 + ), + stream=stream, + min_blocks_per_mp=2, + use_pdl=use_one_wave_reducer_pdl, + ) + elif kernel_split_kv <= Int32(4): + self.reduction_kernel( + o, + lse, + acc_o, + acc_lse, + kernel_split_kv, + cache_seqs, + cu_seqlens_q, + block_split_kvs, + split_divisor=tile_sched_params.split_kv_fdd + if self.static_split_kv is None and not self.is_var_split_kv + else None, + max_q_len=max_q_len, + local_split_capacity=4, + ).launch( + grid=( + ceil_div(logical_query_rows, REDUCTION_ROWS_PER_CTA), + 1, + batch_size, + ), + block=[ + REDUCTION_ROWS_PER_CTA * REDUCTION_THREADS_PER_ROW, + 1, + 1, + ], + smem=( + REDUCTION_ROWS_PER_CTA * 4 * self.lse_dtype.width // 8 + ), + stream=stream, + min_blocks_per_mp=2, + use_pdl=use_one_wave_reducer_pdl, + ) + else: + self.reduction_kernel( + o, + lse, + acc_o, + acc_lse, + kernel_split_kv, + cache_seqs, + cu_seqlens_q, + block_split_kvs, + split_divisor=tile_sched_params.split_kv_fdd + if self.static_split_kv is None and not self.is_var_split_kv + else None, + max_q_len=max_q_len, + ).launch( + grid=( + ceil_div(logical_query_rows, REDUCTION_ROWS_PER_CTA), + 1, + batch_size, + ), + block=[ + REDUCTION_ROWS_PER_CTA * REDUCTION_THREADS_PER_ROW, + 1, + 1, + ], + smem=( + REDUCTION_ROWS_PER_CTA + * self.reduction_split_capacity + * self.lse_dtype.width + // 8 + ), + stream=stream, + min_blocks_per_mp=2, + use_pdl=use_one_wave_reducer_pdl, + ) + else: + self.reduction_kernel( + o, + lse, + acc_o, + acc_lse, + kernel_split_kv, + cache_seqs, + cu_seqlens_q, + block_split_kvs, + split_divisor=tile_sched_params.split_kv_fdd + if self.static_split_kv is None and not self.is_var_split_kv + else None, + max_q_len=max_q_len, + ).launch( + grid=( + ceil_div(logical_query_rows, REDUCTION_ROWS_PER_CTA), + 1, + batch_size, + ), + block=[ + REDUCTION_ROWS_PER_CTA * REDUCTION_THREADS_PER_ROW, + 1, + 1, + ], + smem=( + REDUCTION_ROWS_PER_CTA + * self.reduction_split_capacity + * self.lse_dtype.width + // 8 + ), + stream=stream, + min_blocks_per_mp=2, + use_pdl=use_one_wave_reducer_pdl, + ) @cute.kernel def split_kv_kernel( @@ -1565,28 +1882,27 @@ def split_kv_kernel( cache_seqs: cute.Tensor, cu_seqlens_q: cute.Tensor | None, block_split_kvs: object, - softmax_scale_log2: cutlass.Float32, + scale_softmax_log2: cutlass.Float32, output_scale: cutlass.Float32, tile_sched_params: MLAStaticTileSchedulerParams, clc_tile_sched_params: object, output_quant_args=None, tma_desc_o: cutlass.GridConstant[cuda.TensorMap] = None, + tma_desc_extra_v: cutlass.GridConstant[cuda.TensorMap] = None, + tma_desc_sparse_rope: cutlass.GridConstant[cuda.TensorMap] = None, + tma_desc_sparse_extra_rope: cutlass.GridConstant[cuda.TensorMap] = None, + max_q_len: cutlass.Int32 = 0, ) -> None: """MLA decode TS kernel: persistent tile-scheduled execution.""" + seq_len_q = self.seq_len_q if self.fixed_seqlen else max_q_len if cutlass.const_expr(self.direct_sparse): from ..helpers.sparse_indices import bind_sparse_kv_indices - si, ci, sl, cl, sm, os, sinks = block_split_kvs - capacity = cutlass.const_expr( - max( - 256, - ( - (self.direct_sparse_capacities[0] + 127) // 128 - + (self.direct_sparse_capacities[1] + 127) // 128 - ) - * 128, - ) - ) + si, ci, sl, cl, sinks, ks = block_split_kvs + capacity = ((Int32(si.shape[-1]) + 127) // 128) * 128 + if cutlass.const_expr(self.direct_sparse_capacities[1] > 0): + capacity = capacity + ((Int32(ci.shape[-1]) + 127) // 128) * 128 + capacity = cute.math.max(Int32(256), capacity) request = None if cutlass.const_expr(not self.is_persistent): _, row = divmod_constexpr_power_of_two_or_fdd( @@ -1604,13 +1920,12 @@ def split_kv_kernel( query_row=request, assume_valid_prefix=self.assume_valid_prefix, has_extra=self.direct_sparse_capacities[1] > 0, + primary_tensor_scale=ks if self.kv_cache_format == "fp8" else None, ) - if cutlass.const_expr(not self.direct_static_scales): - # Caller scales already include Q/K/V descales. - softmax_scale_log2 = Float32(sm[0]) * LOG2_E - output_scale = Float32(os[0]) cfg = self._make_config() - num_query_tiles = self.num_q_tiles + num_query_tiles = ( + seq_len_q * self.num_heads + self.mma_qk_tiler_mn[0] - 1 + ) // self.mma_qk_tiler_mn[0] use_clc_dynamic = self.is_persistent and not cfg.is_fp8_qkv() tiled_mma_qk = None if cutlass.const_expr(cfg.is_fp8_qkv()): @@ -1683,6 +1998,22 @@ def split_kv_kernel( space=cutlass.AddressSpace.smem, alignment=128, ) + # NVFP4 pool: per-loader-warp raw-row scratch buffers and their re-armed mbarriers. + smem_raw_arr = None + raw_mbar_arr = None + if cutlass.const_expr(cfg.has_nvfp4_kv): + smem_raw_arr = cutlass.Array( + cutlass.Uint8, + nvfp4_kv_scratch_bytes(cfg.rope_dim), + space=cutlass.AddressSpace.smem, + alignment=128, + ) + raw_mbar_arr = cutlass.Array( + Int64, + cfg.load_num_warps + cfg.load_v_stage, + space=cutlass.AddressSpace.smem, + alignment=16, + ) smem_o = None if cutlass.const_expr(cfg.use_tma_output): smem_o = cutlass.Array( @@ -1837,25 +2168,26 @@ def split_kv_kernel( _, q_len = query_batch_bounds( cu_seqlens_q, tile_batch_idx, - self.seq_len_q, + seq_len_q, ) query_tile_has_rows = runtime_flat_query_tile_has_rows( tile_seq_q_idx, self.mma_qk_tiler_mn[0], self.num_heads, - self.seq_len_q, + seq_len_q, cu_seqlens_q, tile_batch_idx, ) if cutlass.const_expr( - cfg.mask_type == MaskType.CAUSAL.value and self.seq_len_q > 1 + cfg.mask_type == MaskType.CAUSAL.value + and (not self.fixed_seqlen or self.seq_len_q > 1) ): _, _, logical_q_idx, _, _ = flat_query_row_state( Int32(self.mma_qk_tiler_mn[0] - 1), tile_seq_q_idx, self.mma_qk_tiler_mn[0], self.num_heads, - self.seq_len_q, + seq_len_q, cu_seqlens_q, tile_batch_idx, ) @@ -1884,6 +2216,9 @@ def split_kv_kernel( k_tile_total, split_kv_cap, tile_split_kv_idx, + split_divisor=tile_sched_params.split_kv_fdd + if self.static_split_kv is None and not self.is_var_split_kv + else None, ) exit_early = k_tile_count <= Int32(0) if cutlass.const_expr(self.is_persistent): @@ -1970,7 +2305,7 @@ def split_kv_kernel( static_seq_len_k=self.static_seq_len_k, cu_seqlens_q=cu_seqlens_q, logical_num_heads_q=self.num_heads, - logical_seq_len_q=self.seq_len_q, + logical_seq_len_q=seq_len_q, static_problem_shape_b=None, static_problem_shape_s=num_query_tiles, use_clc_dynamic=use_clc_dynamic, @@ -1987,16 +2322,23 @@ def split_kv_kernel( smem_vc_arr=smem_vc_arr, smem_p_arr=smem_p_arr, sparse_offsets_smem=sparse_offsets_smem, + smem_raw_arr=smem_raw_arr, + raw_mbar_arr=raw_mbar_arr, tma_desc_q_latent=tma_desc_q_latent.get_ptr(), tma_desc_q_rope=tma_desc_q_rope.get_ptr(), tma_desc_c_latent=tma_desc_c_latent.get_ptr(), tma_desc_c_rope=tma_desc_c_rope.get_ptr(), tma_desc_c_transpose=tma_desc_c_transpose.get_ptr(), + tma_desc_extra_v=tma_desc_extra_v.get_ptr(), + tma_desc_sparse_rope=tma_desc_sparse_rope.get_ptr(), + tma_desc_sparse_extra_rope=tma_desc_sparse_extra_rope.get_ptr(), + primary_cache=c_latent, + extra_cache=c_rope, page_offsets=page_offsets, blk_coord=blk_coord, tidx=tidx, output=o, - atten_sinks=sinks if self.fuse_sparse_epilogue else None, + attention_sinks=sinks if self.fuse_sparse_epilogue else None, lse_in_natural_log=self.fuse_sparse_epilogue, output_quant=self.output_quant if self.fuse_sparse_epilogue else None, output_quant_args=output_quant_args, @@ -2011,12 +2353,18 @@ def split_kv_kernel( cu_seqlens_q=cu_seqlens_q, split_kv=max_split_kv, logical_num_heads_q=self.num_heads, - logical_seq_len_q=self.seq_len_q, + logical_seq_len_q=seq_len_q, tiled_mma_qk=tiled_mma_qk, ) # Initialize pipelines (creates mbarriers in SMEM) task_manager.setup_resources_and_tasks() + if cutlass.const_expr(cfg.has_nvfp4_kv): + if warp_idx == 0 and prims.elect_sync(): + for stage in cutlass.range_constexpr(cfg.load_v_stage): + prims.mbarrier_init( + raw_mbar_arr.data_ptr(cfg.load_num_warps + stage), 1 + ) # Fence mbarrier init then cluster sync so both CTAs see # initialized barriers before tcgen05_alloc. The alloc with @@ -2063,11 +2411,15 @@ def split_kv_kernel( # Set TMEM base address on resources named_res["tmem_s"].tmem_base_addr = tmem_base_addr + named_res["smem_q"].tmem_base_addr = tmem_base_addr named_res["tmem_o"].tmem_base_addr = tmem_base_addr named_res["tmem_corr"].tmem_base_addr = tmem_base_addr # Set runtime params on TmemS (softmax) - named_res["tmem_s"].softmax_scale_log2 = softmax_scale_log2 + named_res["tmem_s"].scale_softmax_log2 = scale_softmax_log2 + named_res["tmem_corr"].page_offsets = page_offsets + named_res["tmem_corr"].score_resource = named_res["tmem_s"] + named_res["gmem_o"].page_offsets = page_offsets named_res["tmem_s"].smem_exchange = softmax_exchange_arr named_res[ @@ -2076,7 +2428,7 @@ def split_kv_kernel( # Set runtime params on GmemO (epilogue) named_res["gmem_o"].output_scale = output_scale - named_res["gmem_o"].softmax_scale_log2 = softmax_scale_log2 + named_res["gmem_o"].scale_softmax_log2 = scale_softmax_log2 named_res["gmem_o"].smem_exchange = epilogue_exchange_arr.data_ptr().toint( Int32 ) @@ -2117,6 +2469,7 @@ def initialize_workspace( workspace: cute.Tensor, ): """Construct acc_o and acc_lse tensors from the workspace buffer.""" + # Allocation is SM-bounded; live tensors use the runtime split stride. acc_o, acc_lse = None, None if cutlass.const_expr(workspace is not None): # Workspace strides are aligned to 256 bits, expressed in Float16 @@ -2157,8 +2510,12 @@ def reduction_kernel( cache_seqs: cute.Tensor, cu_seqlens_q: cute.Tensor | None, block_split_kvs: cute.Tensor, + split_divisor=None, + max_q_len: cutlass.Int32 = 0, + local_split_capacity: cutlass.Constexpr[int] = 0, ): """Dispatch the throughput 2CTA split-KV reduction body.""" + seq_len_q = self.seq_len_q if self.fixed_seqlen else max_q_len cfg = make_mla_decode_config( mma_qk_tiler_mn=self.mma_qk_tiler_mn, mma_pv_tiler_mn=self.mma_pv_tiler_mn, @@ -2167,6 +2524,7 @@ def reduction_kernel( qkv_dtype=self.qkv_dtype, o_dtype=self.out_dtype, mask_type=self.mask_type, + lse_base_on_e=self.lse_base_on_e, ) if cutlass.const_expr(not self.is_persistent): prims.griddepcontrol(kind=prims.GridDepAction.WAIT) @@ -2181,8 +2539,12 @@ def reduction_kernel( cu_seqlens_q, block_split_kvs, cfg, - self.reduction_split_capacity, + local_split_capacity + if local_split_capacity > 0 + else self.reduction_split_capacity, REDUCTION_ROWS_PER_CTA, + split_divisor=split_divisor, + runtime_seq_len_q=seq_len_q, ) @cute.kernel @@ -2197,8 +2559,11 @@ def parallel_reduction_kernel( cu_seqlens_q: cute.Tensor | None, block_split_kvs: object, output_quant_args=None, + split_divisor=None, + max_q_len: cutlass.Int32 = 0, ): """Dispatch the high-split fixed-D512 cluster reducer.""" + seq_len_q = self.seq_len_q if self.fixed_seqlen else max_q_len cfg = make_mla_decode_config( mma_qk_tiler_mn=self.mma_qk_tiler_mn, @@ -2208,17 +2573,18 @@ def parallel_reduction_kernel( qkv_dtype=self.qkv_dtype, o_dtype=self.out_dtype, mask_type=self.mask_type, + lse_base_on_e=self.lse_base_on_e, ) if cutlass.const_expr(not self.is_persistent): prims.griddepcontrol(kind=prims.GridDepAction.WAIT) - atten_sinks = None + attention_sinks = None if cutlass.const_expr(self.fuse_sparse_reduction): - atten_sinks = block_split_kvs[-1] + attention_sinks = block_split_kvs[4] block_split_kvs = None topology = self._reduction_topology() run_parallel_reduction_kernel( self.num_heads, - self.seq_len_q, + seq_len_q, self.is_var_split_kv, output, lse, @@ -2232,8 +2598,9 @@ def parallel_reduction_kernel( topology.actual_splits, topology.cluster_size, topology.slots_per_rank, - atten_sinks, + attention_sinks, self.fuse_sparse_reduction, self.output_quant, output_quant_args, + split_divisor=split_divisor, ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/nvfp4_kv_load.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/nvfp4_kv_load.py new file mode 100644 index 000000000000..358df5051ade --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/nvfp4_kv_load.py @@ -0,0 +1,392 @@ +# Copyright (c) 2026 by FlashInfer team. +# Licensed under the Apache License, Version 2.0. +"""Convert disjoint K rows once, then exchange peer V features through DSMEM. + +Each CTA owns 64 of 128 token rows. It writes its local V half and sends the +other 256 features per row to its peer. Register stores permit unequal K/V +stage counts; bulk copies retain their K source through coupled K/V reuse. +""" + +import cutlass +import cutlass.cute as cute +from cutlass import Int32 +from cutlass.base_dsl.common import DSLRuntimeError +from cutlass.cutlass_dsl import dsl_user_op +from cutlass.experimental import primitives as prims +from cutlass.pipeline import PipelineState + +from ..helpers.gather import gather4_quad +from ..helpers.nvfp4_kv import ( + convert_nvfp4x8, + nvfp4_row_bytes, + nvfp4_scale_byte_offset, + publish_nvfp4_kv, + store_shared_u32x4, + swizzle_byte_offset, +) + + +@cute.jit +def shared_byte_view(ptr, count: cutlass.Constexpr[int]): + """Expose count bytes of CTA-local raw scratch for packed scale/payload reads.""" + return cutlass.Array(ptr, dtype=cutlass.Uint8, shape=(count,), addrspace=3) + + +@cute.jit +def arm_raw_load(barrier, tx_bytes: cutlass.Constexpr[int], reset=False): + """Initialize once per work tile, then advance the raw barrier every KV tile.""" + if reset: + prims.mbarrier_init(barrier, 1) + prims.fence_mbarrier_init() + prims.mbarrier_arrive_expect_tx(barrier, tx_bytes) + + +@cute.jit +def wait_raw_load(barrier, phase): + """Acquire completion of the current per-warp TMA epoch before reading raw bytes.""" + while not prims.mbarrier_wait_parity(barrier, phase, prims.MBarrierWait.TRY): + pass + + +@cute.jit +def load_shared_row_quad(smem_sparse_offsets, offset): + """Four storage rows of one gather quad from the offsets ring; masked rows keep the OOB sentinel.""" + raw = smem_sparse_offsets.data_ptr(offset).load(count=4, alignment=16) + rows = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) + for j in cutlass.range_constexpr(4): + rows[j] = raw[j] + return rows + + +@cute.jit +def load_shared_words(addr, count: cutlass.Constexpr[int]): + """Vector shared load of ``count`` words (2 or 4) from a suitably aligned address.""" + return cutlass.inttoptr(addr, 3, Int32).load(count=count, alignment=4 * count) + + +def nvfp4_kv_scratch_bytes(rope_dim): + """Return bytes for the 64 NVFP4 rows owned by one CTA; FP8 gathers directly into K.""" + return 64 * nvfp4_row_bytes(rope_dim) + + +@cute.jit +def issue_k_gathers(k, raw, dst, barrier, rows, q: cutlass.Constexpr[int], warp, fp8): + """Gather one warp-owned four-row group into raw NVFP4 scratch or FP8 K panels. + + FP8 primary emits four 128-byte panel gathers. NVFP4 emits one whole + record gather and converts later; both routes use the same source tag.""" + if fp8: + quad = warp * Int32(2) + Int32(q) + for part in cutlass.range_constexpr(4): + gather4_quad( + dst + part * 8192 + quad * Int32(512), + k.tma_desc_c_latent, + k.tma_desc_c_rope, + Int32(part * 128), + rows, + barrier, + ) + else: + gather4_quad( + raw + q * 4 * nvfp4_row_bytes(k.cfg.rope_dim), + k.tma_desc_c_latent, + k.tma_desc_c_rope, + Int32(0), + rows, + barrier, + ) + + +@dsl_user_op +def store_v_remote(dst, offset, done, owner, a, b, c, d, *, loc=None, ip=None): + """Send sixteen converted V bytes and credit the peer completion barrier. + + Offset is in bytes within an acquired peer V stage. Prefer the native + vector store for compiler scheduling. DSL 4.7 validates tuples but then + converts them to a scalar; use inline PTX only for that wrapper failure. + """ + destination = cute.make_ptr( + cutlass.Int32, + dst.toint(Int32) + offset, + cute.AddressSpace.smem, + assumed_align=16, + loc=loc, + ip=ip, + ) + completion = cute.make_ptr( + cutlass.Int64, + done.toint(Int32), + cute.AddressSpace.smem, + assumed_align=8, + loc=loc, + ip=ip, + ) + try: + cute.arch.store_async_dsmem( + destination, (a, b, c, d), completion, owner, loc=loc, ip=ip + ) + return + except DSLRuntimeError as error: + if "to integer conversion is not supported" not in str(error): + raise + + remote = prims.mapa(dst + offset, owner, loc=loc, ip=ip) + remote_completion = prims.mapa(done, owner, loc=loc, ip=ip) + cute.arch.inline_ptx( + "st.async.shared::cluster.mbarrier::complete_tx::bytes.v4.b32 " + "[{$r0}], {{$r1}, {$r2}, {$r3}, {$r4}}, [{$r5}];", + read_only_args=[remote.ir_value(), a, b, c, d, remote_completion.ir_value()], + loc=loc, + ip=ip, + ) + + +@cute.jit +def convert_k_and_exchange_v( + kraw, + kdst, + vdst, + lane, + warp, + cta, + barrier, + done, + *, + fp8: cutlass.Constexpr[bool], + register_exchange: cutlass.Constexpr[bool], + rope_dim: cutlass.Constexpr[int], +): + """Convert eight owned token rows and distribute their V features across the CTA pair. + + Each lane handles sixteen columns. Its CTA keeps one 256-feature V + half and sends the other half to its peer. FP8 primary K is already in + place; NVFP4 is expanded once and reused for both QK and PV. Register + exchange permits unequal K/V stage counts; bulk exchange retains K + as the copy source until both V consumers finish.""" + stride = nvfp4_row_bytes(rope_dim) + for i in cutlass.range(8, unroll=2 if fp8 else 8): + row = warp * Int32(8) + Int32(i) + offset = (lane >> Int32(3)) * Int32(8192) + swizzle_byte_offset( + row * Int32(128) + (lane & Int32(7)) * Int32(16), 128 + ) + a, b, c, d = Int32(0), Int32(0), Int32(0), Int32(0) + if cutlass.const_expr(fp8): + k8 = load_shared_words(kdst.toint(Int32) + offset, 4) + a, b, c, d = k8[0], k8[1], k8[2], k8[3] + else: + k4 = load_shared_words( + kraw.toint(Int32) + Int32(i * stride) + lane * Int32(8), 2 + ) + sf = Int32( + shared_byte_view(kraw, 8 * stride)[ + Int32(i * stride + 256 + rope_dim) + + nvfp4_scale_byte_offset(lane, rope_dim) + ] + ) + a, b = convert_nvfp4x8(k4[0], sf) + c, d = convert_nvfp4x8(k4[1], sf) + # FP8 panel TMA already populated K. Only NVFP4 needs a K write; + # both formats reuse these registers for local/remote V publication. + if cutlass.const_expr(not fp8): + store_shared_u32x4(kdst, offset, a, b, c, d) + owner = (lane >> Int32(3)) & Int32(1) + voffset = (cta * Int32(2) + (lane >> Int32(4))) * Int32(8192) + voffset += swizzle_byte_offset( + row * Int32(128) + (lane & Int32(7)) * Int32(16), 128 + ) + if owner == cta: + store_shared_u32x4(vdst, voffset, a, b, c, d) + elif cutlass.const_expr(register_exchange): + store_v_remote(vdst, voffset, done, owner, a, b, c, d) + if cutlass.const_expr(rope_dim): + r, c = lane >> Int32(2), lane & Int32(3) + wr = load_shared_words( + kraw.toint(Int32) + Int32(256) + r * Int32(stride) + c * Int32(16), 4 + ) + row = warp * Int32(8) + r + store_shared_u32x4( + kdst, + Int32(32768) + swizzle_byte_offset(row * Int32(64) + c * Int32(16), 64), + wr[0], + wr[1], + wr[2], + wr[3], + ) + publish_nvfp4_kv(barrier, 8 * (512 + rope_dim), cta_group=2) + if cutlass.const_expr(not register_exchange): + # This warp owns all eight source rows; K publication supplied its + # proxy fence and warp join. Coupled K/V slots retain the copy source. + if prims.elect_sync(): + peer = cta ^ Int32(1) + for part in cutlass.range_constexpr(2): + source = ( + kdst + (peer + Int32(part * 2)) * Int32(8192) + warp * Int32(1024) + ) + destination = ( + vdst + + (cta * Int32(2) + Int32(part)) * Int32(8192) + + warp * Int32(1024) + ) + prims.cp_async_bulk_shared_cluster_shared_cta( + prims.mapa(destination, peer), + source, + prims.mapa(done, peer), + Int32(1024), + ) + + +@cute.jit +def load_nvfp4_source_stage( + k, v, stage_info, v_stage, page_stage, *, fp8_source: cutlass.Constexpr[bool] +): + """Specialize the complete producer so source-format branches stay together.""" + cfg = k.cfg + assert cfg.load_num_warps == 8 + assert cfg.sparse_offset_stages >= 2 + register_exchange = cfg.load_v_stage > cfg.load_k_stage + lane = cute.arch.thread_idx()[0] & Int32(31) + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - Int32(cfg.load_k_warp_id) + cta = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + kdst = k.smem_k.data_ptr(stage_info.stage_idx * cfg.smem_k_stage_elems) + vdst = v.smem_v.data_ptr(v_stage * cfg.smem_v_stage_elems) + raw_bytes = 8 * nvfp4_row_bytes(cfg.rope_dim) + raw = k.smem_raw.data_ptr(warp * Int32(raw_bytes)) + raw_bar = k.raw_mbar.data_ptr(warp) + done = k.raw_mbar.data_ptr(Int32(cfg.load_num_warps) + v_stage) + phase = v.current_slot[1] + # The group-2 V empty barrier completes only after both CTAs' V readers + # retire. Local acquire therefore permits writes to either V destination. + # Completion may precede expect_tx: the pending arrival keeps this phase + # open, and the signed transaction counter permits that ordering (PTX9.7.15.16). + if warp == Int32(0) and prims.elect_sync(): + prims.mbarrier_arrive_expect_tx(done, Int32(16384)) + base = page_stage * Int32(128) + cta * Int32(64) + loop = Int32(stage_info.loop_offset) + if fp8_source or loop == Int32(0): + if prims.elect_sync(): + if fp8_source: + arm_raw_load(raw_bar, 8 * 512, loop == Int32(0)) + else: + arm_raw_load(raw_bar, raw_bytes, True) + for q in cutlass.range_constexpr(2): + rows = load_shared_row_quad( + k.smem_sparse_offsets, base + warp * Int32(8) + Int32(q * 4) + ) + issue_k_gathers(k, raw, kdst, raw_bar, rows, q, warp, fp8_source) + has_next = loop + Int32(1) < Int32(stage_info.work_tile.k_tile_count) + next_stage = Int32(0) + if has_next: + # Read ahead under the same future consumer token. This wait does not + # advance/release it; the normal next iteration still owns its release. + # Construct new scalar values: DSL 4.7 can alias the staged fields + # returned by clone(), so advance() also advances the live consumer + # and releases the next ring slot instead of the one being read. + current = k.page_window.state_src.consumer_state + next_stage = (current.index + Int32(1)) % Int32(cfg.sparse_offset_stages) + upcoming = PipelineState( + cfg.sparse_offset_stages, + current.count + Int32(1), + next_stage, + current.phase ^ Int32(next_stage == Int32(0)), + ) + k.page_window.pipeline.consumer_wait(upcoming) + wait_raw_load(raw_bar, loop & Int32(1)) + if fp8_source: + convert_k_and_exchange_v( + raw, + kdst, + vdst, + lane, + warp, + cta, + stage_info.barrier, + done, + fp8=True, + register_exchange=register_exchange, + rope_dim=cfg.rope_dim, + ) + else: + convert_k_and_exchange_v( + raw, + kdst, + vdst, + lane, + warp, + cta, + stage_info.barrier, + done, + fp8=False, + register_exchange=register_exchange, + rope_dim=cfg.rope_dim, + ) + + # Prefetch only raw NVFP4 K; per-tensor FP8 writes the next acquired K stage directly. + if has_next: + # K publication already joined all raw readers. + next_base = next_stage * Int32(128) + cta * Int32(64) + warp * Int32(8) + next_fp8 = tile_is_fp8(k, stage_info, loop + Int32(1)) + if not next_fp8 and prims.elect_sync(): + arm_raw_load(raw_bar, raw_bytes) + rows0 = load_shared_row_quad(k.smem_sparse_offsets, next_base) + rows1 = load_shared_row_quad(k.smem_sparse_offsets, next_base + Int32(4)) + issue_k_gathers(k, raw, kdst, raw_bar, rows0, 0, warp, False) + issue_k_gathers(k, raw, kdst, raw_bar, rows1, 1, warp, False) + # For bulk copies, the next coupled K/V acquisitions protect K until both + # V destinations have consumed the previous copies. Register stores need + # no retained K source after their payload has been issued. + # Only one publisher needs to gate the V full barrier on incoming data. + # Other publishers can advance; the group-2 MMA cannot start until this + # publisher contributes its final bytes after observing completion. + if warp == Int32(0): + while not prims.mbarrier_try_wait_parity( + done, phase, scope=prims.MBarrierScope.CLUSTER, order=prims.MemOrder.ACQUIRE + ): + pass + + +@cute.jit +def publish_nvfp4_v(barrier, count): + # K publication has already proxy-fenced and warp-joined every ordinary + # local V store. The remaining writes are async DSMEM transactions. + # Warp0's incoming completion acquire is ordered before this release; + # its withheld credit keeps V-full closed until all received bytes arrive. + """Publish V only after local stores and incoming DSMEM transfers are complete. + + K publication already fenced the local writes. The designated publisher + acquired peer completion before releasing the cluster-wide V-full credit.""" + if prims.elect_sync(): + prims.fence_sync_restrict(prims.MemOrder.RELEASE) + prims.mbarrier_complete_tx( + prims.mapa(barrier, Int32(0)), Int32(count), scope=prims.MemScope.CLUSTER + ) + + +@cute.jit +def tile_is_fp8(k, stage_info, offset): + """Select the primary FP8 producer for the current or prefetched KV tile. + + Compare logical tile position with the live primary boundary cached + by the metadata task. The extra source always selects NVFP4.""" + cfg = k.cfg + if cutlass.const_expr(cfg.kv_cache_format != "fp8"): + return cutlass.Boolean(False) + extra = cutlass.Boolean(False) + if cutlass.const_expr(k.page_offsets.has_extra): + primary_tiles = Int32(0) + if cutlass.const_expr(k.page_offsets.is_bound): + primary_tiles = ( + Int32(k.page_offsets.primary_length) + Int32(127) + ) // Int32(128) + else: + primary_tiles = k.request_lengths[2] + extra = Int32(stage_info.work_tile.k_index_base) + offset >= primary_tiles + return not extra + + +@cute.jit +def load_nvfp4_kv_stage(k, v, stage_info, v_stage, page_stage): + """Select the format once before the specialized gather/conversion path.""" + if tile_is_fp8(k, stage_info, Int32(stage_info.loop_offset)): + load_nvfp4_source_stage(k, v, stage_info, v_stage, page_stage, fp8_source=True) + else: + load_nvfp4_source_stage(k, v, stage_info, v_stage, page_stage, fp8_source=False) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/parallel_reduction.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/parallel_reduction.py index 41b61707a4c4..1df159a0869a 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/parallel_reduction.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/parallel_reduction.py @@ -24,7 +24,8 @@ normalized_lse_weight, attention_sink_scale, ) -from ..helpers.constants import LN_2, SPLIT_REDUCTION_SCALE_BARRIER_ID +from ..helpers.constants import SPLIT_REDUCTION_SCALE_BARRIER_ID +from ..helpers.math import final_lse from ..helpers.mask import MaskType, mask_visible_k_length from ..helpers.math import ceil_div from ..helpers.ops import fmax_f32, warp_reduce_max_f32, warp_reduce_sum_f32 @@ -49,7 +50,7 @@ def _parallel_reduction_row_state( tile_size_q: cutlass.Constexpr[int], num_heads: cutlass.Constexpr[int], - seq_len_q: cutlass.Constexpr[int], + seq_len_q: Int32, is_var_split_kv: cutlass.Constexpr[bool], cache_seqs, cu_seqlens_q, @@ -59,6 +60,7 @@ def _parallel_reduction_row_state( query_tile_idx, batch_idx, cfg, + split_divisor=None, ): """Return logical row coordinates and its runtime active split count. @@ -93,8 +95,10 @@ def _parallel_reduction_row_state( batch_idx, ) tile_k = cache_seqs[batch_idx] + query_is_valid = query_is_valid & (tile_k > Int32(0)) + tile_k = cute.math.max(tile_k, Int32(0)) row_k = tile_k - if cutlass.const_expr(cfg.mask_type == MaskType.CAUSAL.value and seq_len_q > 1): + if cutlass.const_expr(cfg.mask_type == MaskType.CAUSAL.value): _, logical_seq_len_q = query_batch_bounds( cu_seqlens_q, batch_idx, @@ -127,6 +131,7 @@ def _parallel_reduction_row_state( row_k_tile_total, tile_k_tile_total, split_kv_cap, + split_divisor=split_divisor, ) return ( logical_head_idx, @@ -154,16 +159,17 @@ def _store_parallel_reduction_result( finalize_attention: cutlass.Constexpr[bool] = False, output_quant: cutlass.Constexpr = None, output_quant_args=None, + base_on_e: cutlass.Constexpr[bool] = False, ): """Publish one normalized FP32 output fragment and final LSE.""" - if cutlass.const_expr(finalize_attention): - global_lse *= Float32(LN_2) - if tidx == Int32(0) and query_is_valid: - if cutlass.const_expr(cu_seqlens_q is not None): - lse[logical_head_idx, storage_q_idx] = global_lse - else: - lse[logical_head_idx, logical_q_idx, batch_idx] = global_lse + if cutlass.const_expr(lse is not None): + if tidx == Int32(0) and query_is_valid: + value = final_lse(global_lse, base_on_e or finalize_attention) + if cutlass.const_expr(cu_seqlens_q is not None): + lse[logical_head_idx, storage_q_idx] = value + else: + lse[logical_head_idx, logical_q_idx, batch_idx] = value if cutlass.const_expr(output_quant is not None): if query_is_valid: @@ -195,10 +201,63 @@ def _store_parallel_reduction_result( ) +@cute.jit +def accumulate_parallel_reduction_slot( + acc_output, + smem_local_scale, + local_output, + row_in_tile, + query_tile_idx, + batch_idx, + element_idx, + slot_i, + rank_split_begin, + query_is_valid, + active_split_kv, + actual_splits: cutlass.Constexpr[int], + finalize_attention: cutlass.Constexpr[bool], +): + """Add one valid split's BF16 vec4 partial into the FP32 register fragment. + + ``slot_i`` indexes rank-local scales; ``rank_split_begin`` maps it into the + producer workspace. Inactive slots never load workspace, and sparse + finalization skips zero-weight partials. The updated fragment is returned + so both static and runtime merge loops share identical accumulation. + """ + acc_output_ptr = acc_output.iterator.raw_ptr() + split_idx = rank_split_begin + Int32(slot_i) + active_slot = ( + query_is_valid + & (split_idx < Int32(actual_splits)) + & (split_idx < active_split_kv) + ) + if active_slot: + split_scale = Float32(smem_local_scale[slot_i]) + if not finalize_attention or split_scale > Float32(0): + partial_offset = Int64( + row_in_tile * acc_output.stride[0] + + split_idx * acc_output.stride[1] + + element_idx * acc_output.stride[2] + + query_tile_idx * acc_output.stride[3] + + batch_idx * acc_output.stride[4] + ) + partial_output = ( + (acc_output_ptr + partial_offset) + .load( + count=PARALLEL_REDUCTION_ELEMENTS_PER_THREAD, + alignment=8, + ) + .to(Float32) + ) + for j in cutlass.range_constexpr(PARALLEL_REDUCTION_ELEMENTS_PER_THREAD): + local_output[j] += partial_output[j] * split_scale + return local_output + + @cute.jit def run_parallel_reduction_kernel( num_heads: cutlass.Constexpr[int], - seq_len_q: cutlass.Constexpr[int], + seq_len_q: Int32, is_var_split_kv: cutlass.Constexpr[bool], output, lse, @@ -212,15 +271,17 @@ def run_parallel_reduction_kernel( actual_splits: cutlass.Constexpr[int], cluster_size: cutlass.Constexpr[int], slots_per_rank: cutlass.Constexpr[int], - atten_sinks=None, + attention_sinks=None, finalize_attention: cutlass.Constexpr[bool] = False, output_quant: cutlass.Constexpr = None, output_quant_args=None, + split_divisor=None, ): """Reduce one D=512 row cooperatively across a padded CTA cluster. - Rank ``r`` owns ``slots_per_rank`` contiguous split slots. Slots beyond - ``actual_splits`` and row-specific inactive splits perform no GMEM access + Rank ``r`` owns at most ``slots_per_rank`` contiguous split slots, balanced + over the live runtime prefix. Slots beyond ``actual_splits`` and + row-specific inactive splits perform no GMEM access or arithmetic. G1 writes its result directly; for G2/G4/G8 every rank publishes a neutral-or-valid ``(FP32 LSE, BF16 O[512])`` state to DSMEM and rank zero performs the final merge. Arithmetic remains FP32. @@ -241,7 +302,9 @@ def run_parallel_reduction_kernel( # A sparse query is already one virtual request. The producer writes # neutral LSE for skipped splits, so the compiled slot count suffices; # live length loads and dynamic split division would repeat per head. - if cutlass.const_expr(seq_len_q != 1 or cu_seqlens_q is not None): + if cutlass.const_expr( + not isinstance(seq_len_q, int) or seq_len_q != 1 or cu_seqlens_q is not None + ): raise ValueError("final sparse reduction requires virtual query rows") logical_head_idx = query_tile_idx * Int32(cfg.mma_qk_tiler[0]) + row_in_tile logical_q_idx = Int32(0) @@ -268,6 +331,7 @@ def run_parallel_reduction_kernel( query_tile_idx, batch_idx, cfg, + split_divisor=split_divisor, ) element_idx = tidx * Int32(PARALLEL_REDUCTION_ELEMENTS_PER_THREAD) @@ -286,6 +350,19 @@ def run_parallel_reduction_kernel( alignment=16, ) + # Capacity sizes SMEM only. Partition the live prefix so a smaller runtime + # count cannot strand ranks while the remaining ranks do extra work. + live_splits = cute.math.min(active_split_kv, Int32(actual_splits)) + live_slots_per_rank = cute.ceil_div(live_splits, cluster_size) + rank_split_begin = cluster_rank * live_slots_per_rank + live_rank_splits = cute.math.max( + cute.math.min( + live_splits - rank_split_begin, + live_slots_per_rank, + ), + Int32(0), + ) + # One warp computes the rank-local softmax state for this row. All output # threads consume the same scales, so repeating LSE loads and SFU work in # every vector lane only adds instructions and register pressure. @@ -297,25 +374,26 @@ def run_parallel_reduction_kernel( ) local_lse_max = neg_inf for lane_slot_i in cutlass.range_constexpr(lse_slots_per_lane): - local_slot_idx = lane_idx + Int32(lane_slot_i * cfg.threads_per_warp) - split_idx = cluster_rank * Int32(slots_per_rank) + local_slot_idx - active_slot = ( - query_is_valid - & (local_slot_idx < Int32(slots_per_rank)) - & (split_idx < Int32(actual_splits)) - & (split_idx < active_split_kv) - ) - lane_lse[lane_slot_i] = neg_inf - # Keep the workspace access inside the dynamic predicate. Padded - # ranks must not form a load from the unpadded producer allocation. - if active_slot: - lane_lse[lane_slot_i] = Float32( - acc_lse[row_in_tile, split_idx, query_tile_idx, batch_idx] + if Int32(lane_slot_i * cfg.threads_per_warp) < live_rank_splits: + local_slot_idx = lane_idx + Int32(lane_slot_i * cfg.threads_per_warp) + split_idx = rank_split_begin + local_slot_idx + active_slot = ( + query_is_valid + & (local_slot_idx < live_rank_splits) + & (split_idx < Int32(actual_splits)) + & (split_idx < active_split_kv) + ) + lane_lse[lane_slot_i] = neg_inf + # Keep the workspace access inside the dynamic predicate. Padded + # ranks must not form a load from the unpadded producer allocation. + if active_slot: + lane_lse[lane_slot_i] = Float32( + acc_lse[row_in_tile, split_idx, query_tile_idx, batch_idx] + ) + local_lse_max = fmax_f32( + local_lse_max, + lane_lse[lane_slot_i], ) - local_lse_max = fmax_f32( - local_lse_max, - lane_lse[lane_slot_i], - ) local_lse_max = warp_reduce_max_f32(local_lse_max) local_exp_frame = local_lse_max if local_lse_max != neg_inf else Float32(0.0) @@ -326,33 +404,37 @@ def run_parallel_reduction_kernel( ) local_sum_lse = Float32(0.0) for lane_slot_i in cutlass.range_constexpr(lse_slots_per_lane): - lane_exp[lane_slot_i] = Float32( - cute.math.exp2( - lane_lse[lane_slot_i] - local_exp_frame, - fastmath=True, + if Int32(lane_slot_i * cfg.threads_per_warp) < live_rank_splits: + lane_exp[lane_slot_i] = Float32( + cute.math.exp2( + lane_lse[lane_slot_i] - local_exp_frame, + fastmath=True, + ) ) - ) - local_sum_lse += lane_exp[lane_slot_i] + local_sum_lse += lane_exp[lane_slot_i] local_sum_lse = warp_reduce_sum_f32(local_sum_lse) # Keep this warp-local value distinct from the post-barrier SMEM load. stats_lse_value = finalize_log2_sum_exp(local_exp_frame, local_sum_lse) sink_scale = Float32(1) - if cutlass.const_expr(finalize_attention and atten_sinks is not None): + if cutlass.const_expr(finalize_attention and attention_sinks is not None): sink_scale = Float32(0) if query_is_valid: sink_scale = attention_sink_scale( - stats_lse_value, Float32(atten_sinks[logical_head_idx]) + stats_lse_value, Float32(attention_sinks[logical_head_idx]) ) if lane_idx == Int32(0): smem_local_lse[0] = stats_lse_value for lane_slot_i in cutlass.range_constexpr(lse_slots_per_lane): - local_slot_idx = lane_idx + Int32(lane_slot_i * cfg.threads_per_warp) - if local_slot_idx < Int32(slots_per_rank): - weight = normalized_lse_weight(lane_lse[lane_slot_i], stats_lse_value) - if cutlass.const_expr(finalize_attention): - weight *= sink_scale - smem_local_scale[local_slot_idx] = weight + if Int32(lane_slot_i * cfg.threads_per_warp) < live_rank_splits: + local_slot_idx = lane_idx + Int32(lane_slot_i * cfg.threads_per_warp) + if local_slot_idx < live_rank_splits: + weight = normalized_lse_weight( + lane_lse[lane_slot_i], stats_lse_value + ) + if cutlass.const_expr(finalize_attention): + weight *= sink_scale + smem_local_scale[local_slot_idx] = weight prims.barrier_cta_sync(SPLIT_REDUCTION_SCALE_BARRIER_ID) @@ -367,36 +449,43 @@ def run_parallel_reduction_kernel( for j in cutlass.range_constexpr(PARALLEL_REDUCTION_ELEMENTS_PER_THREAD): local_output[j] = Float32(0.0) - acc_output_ptr = acc_output.iterator.raw_ptr() - for slot_i in cutlass.range_constexpr(slots_per_rank): - split_idx = cluster_rank * Int32(slots_per_rank) + Int32(slot_i) - active_slot = ( - query_is_valid - & (split_idx < Int32(actual_splits)) - & (split_idx < active_split_kv) - ) - if active_slot: - split_scale = Float32(smem_local_scale[slot_i]) - if not finalize_attention or split_scale > Float32(0): - partial_offset = Int64( - row_in_tile * acc_output.stride[0] - + split_idx * acc_output.stride[1] - + element_idx * acc_output.stride[2] - + query_tile_idx * acc_output.stride[3] - + batch_idx * acc_output.stride[4] - ) - partial_output = ( - (acc_output_ptr + partial_offset) - .load( - count=PARALLEL_REDUCTION_ELEMENTS_PER_THREAD, - alignment=8, - ) - .to(Float32) - ) - for j in cutlass.range_constexpr( - PARALLEL_REDUCTION_ELEMENTS_PER_THREAD - ): - local_output[j] += partial_output[j] * split_scale + if cutlass.const_expr(finalize_attention): + # Sparse plans have fixed split capacity; stage each merge slot directly. + # Resolve fixed slot addresses before generating the staged merge loop. + for slot_i in cutlass.range_constexpr(slots_per_rank): + local_output = accumulate_parallel_reduction_slot( + acc_output, + smem_local_scale, + local_output, + row_in_tile, + query_tile_idx, + batch_idx, + element_idx, + slot_i, + rank_split_begin, + query_is_valid, + active_split_kv, + actual_splits, + finalize_attention, + ) + else: + # Dense plans keep runtime split pruning within their prepared capacity. + for slot_i in cutlass.range(live_rank_splits, unroll=4): + local_output = accumulate_parallel_reduction_slot( + acc_output, + smem_local_scale, + local_output, + row_in_tile, + query_tile_idx, + batch_idx, + element_idx, + slot_i, + rank_split_begin, + query_is_valid, + active_split_kv, + actual_splits, + finalize_attention, + ) local_lse_value = smem_local_lse[0] @@ -419,6 +508,7 @@ def run_parallel_reduction_kernel( finalize_attention, output_quant, output_quant_args, + base_on_e=cfg.lse_base_on_e, ) return @@ -536,6 +626,7 @@ def run_parallel_reduction_kernel( tidx, batch_idx, cu_seqlens_q, + base_on_e=cfg.lse_base_on_e, ) prims.barrier_cluster_arrive_relaxed() diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/reduction.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/reduction.py index 5de512dea6d0..6310300e8401 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/reduction.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/reduction.py @@ -19,6 +19,7 @@ from cutlass import Float32, Int32, Int64 from cutlass.experimental import primitives as prims +from ..helpers.math import final_lse from .config import ( REDUCTION_THREADS_PER_ROW, REDUCTION_VALUES_PER_THREAD, @@ -54,6 +55,8 @@ def run_reduction_kernel( cfg, max_splits: cutlass.Constexpr[int], rows_per_cta: cutlass.Constexpr[int], + split_divisor=None, + runtime_seq_len_q=None, ): """Combine consecutive logical split rows in one reference-reducer CTA. @@ -65,6 +68,7 @@ def run_reduction_kernel( inactive row groups; those groups still synchronize but never access the workspace or publish public output. """ + seq_len_q = kernel.seq_len_q if runtime_seq_len_q is None else runtime_seq_len_q reduction_group_idx, _, batch_idx = cute.arch.block_idx() tidx, _, _ = cute.arch.thread_idx() row_in_cta = tidx // Int32(REDUCTION_THREADS_PER_ROW) @@ -75,7 +79,7 @@ def run_reduction_kernel( lane_idx = row_thread_idx % Int32(cfg.threads_per_warp) physical_tile_rows = Int32(cfg.mma_qk_tiler[0]) - logical_query_rows = Int32(kernel.num_heads * kernel.seq_len_q) + logical_query_rows = Int32(kernel.num_heads * seq_len_q) local_flat_query_row = reduction_group_idx * Int32(rows_per_cta) + row_in_cta row_is_valid = local_flat_query_row < logical_query_rows # Clamp the final CTA's inactive row groups before decomposing the physical @@ -84,7 +88,9 @@ def run_reduction_kernel( safe_local_flat_query_row = cute.math.min( local_flat_query_row, logical_query_rows - Int32(1) ) - if cutlass.const_expr(kernel.num_heads * kernel.seq_len_q <= cfg.mma_qk_tiler[0]): + if cutlass.const_expr( + kernel.fixed_seqlen and kernel.num_heads * seq_len_q <= cfg.mma_qk_tiler[0] + ): query_tile_idx = Int32(0) row_in_tile = safe_local_flat_query_row else: @@ -108,7 +114,7 @@ def run_reduction_kernel( query_tile_idx, cfg.mma_qk_tiler[0], kernel.num_heads, - kernel.seq_len_q, + seq_len_q, cu_seqlens_q, batch_idx, ) @@ -116,7 +122,7 @@ def run_reduction_kernel( public_flat_query_row = storage_flat_query_row if cutlass.const_expr(cu_seqlens_q is None): public_flat_query_row = public_flat_query_row + batch_idx * Int32( - kernel.seq_len_q * kernel.num_heads + seq_len_q * kernel.num_heads ) # The scalar split count remains the grid/workspace capacity. A variable- @@ -135,21 +141,24 @@ def run_reduction_kernel( # Producer splits are sized from the group's largest K domain, while each # logical row consumes only the prefix containing its visible K tiles. tile_k = cache_seqs[batch_idx] + query_is_valid = query_is_valid & (tile_k > Int32(0)) + tile_k = cute.math.max(tile_k, Int32(0)) row_k = tile_k if cutlass.const_expr( - cfg.mask_type == MaskType.CAUSAL.value and kernel.seq_len_q > 1 + cfg.mask_type == MaskType.CAUSAL.value + and (not kernel.fixed_seqlen or seq_len_q > 1) ): _, logical_seq_len_q = query_batch_bounds( cu_seqlens_q, batch_idx, - kernel.seq_len_q, + seq_len_q, ) _, _, tile_last_logical_q_idx, _, _ = flat_query_row_state( Int32(cfg.mma_qk_tiler[0] - 1), query_tile_idx, cfg.mma_qk_tiler[0], kernel.num_heads, - kernel.seq_len_q, + seq_len_q, cu_seqlens_q, batch_idx, ) @@ -173,6 +182,7 @@ def run_reduction_kernel( row_k_tile_total, tile_k_tile_total, split_kv_cap, + split_divisor=split_divisor, ) smem_lse_scale = cutlass.Array( @@ -191,23 +201,25 @@ def run_reduction_kernel( local_lse = cutlass.Array(kernel.lse_dtype, lse_per_thread) lse_max = kernel.lse_dtype(-kernel.lse_dtype.inf) for i in cutlass.range_constexpr(lse_per_thread): - split_kv_idx = lane_idx + i * cfg.threads_per_warp - active_slot = query_is_valid & cute.elem_less( - split_kv_idx, - local_split_kv, - ) - local_lse[i] = -kernel.lse_dtype.inf - # Keep the workspace access inside the dynamic predicate. Split - # capacities need not be warp-aligned, so padded lanes must not - # form an address beyond the producer allocation. - if active_slot: - local_lse[i] = acc_lse_tile[split_kv_idx] - lse_max = fmax_f32(lse_max, local_lse[i]) + if Int32(i * cfg.threads_per_warp) < local_split_kv: + split_kv_idx = lane_idx + i * cfg.threads_per_warp + active_slot = query_is_valid & cute.elem_less( + split_kv_idx, + local_split_kv, + ) + local_lse[i] = -kernel.lse_dtype.inf + # Keep the workspace access inside the dynamic predicate. Split + # capacities need not be warp-aligned, so padded lanes must not + # form an address beyond the producer allocation. + if active_slot: + local_lse[i] = acc_lse_tile[split_kv_idx] + lse_max = fmax_f32(lse_max, local_lse[i]) lse_max = warp_reduce_max_f32(lse_max) lse_max = lse_max if lse_max != -kernel.lse_dtype.inf else 0.0 sum_lse = kernel.lse_dtype(0.0) for i in cutlass.range_constexpr(lse_per_thread): - sum_lse += cute.math.exp2(local_lse[i] - lse_max, fastmath=True) + if Int32(i * cfg.threads_per_warp) < local_split_kv: + sum_lse += cute.math.exp2(local_lse[i] - lse_max, fastmath=True) sum_lse = warp_reduce_sum_f32(sum_lse) has_finite_mass = sum_lse == sum_lse and sum_lse != kernel.lse_dtype(0.0) global_lse = ( @@ -216,15 +228,19 @@ def run_reduction_kernel( else -kernel.lse_dtype.inf ) if lane_idx == 0 and query_is_valid: - (lse.iterator.raw_ptr() + Int64(public_flat_query_row)).store(global_lse) - for i in cutlass.range_constexpr(lse_per_thread): - split_kv_idx = lane_idx + i * cfg.threads_per_warp - if cute.elem_less(split_kv_idx, local_split_kv): - smem_lse_scale[row_scale_offset + split_kv_idx] = ( - cute.math.exp2(local_lse[i] - global_lse, fastmath=True) - if has_finite_mass - else kernel.lse_dtype(0.0) + if cutlass.const_expr(lse is not None): + (lse.iterator.raw_ptr() + Int64(public_flat_query_row)).store( + final_lse(global_lse, cfg.lse_base_on_e) ) + for i in cutlass.range_constexpr(lse_per_thread): + if Int32(i * cfg.threads_per_warp) < local_split_kv: + split_kv_idx = lane_idx + i * cfg.threads_per_warp + if cute.elem_less(split_kv_idx, local_split_kv): + smem_lse_scale[row_scale_offset + split_kv_idx] = ( + cute.math.exp2(local_lse[i] - global_lse, fastmath=True) + if has_finite_mass + else kernel.lse_dtype(0.0) + ) # Independent writer warps publish scales before any row consumes # them. Invalid/padded rows participate so this remains a full CTA barrier. @@ -247,16 +263,9 @@ def run_reduction_kernel( acc_output_ptr = acc_output.iterator.raw_ptr() if query_is_valid: partial_elem_offset_base = ( - Int64(batch_idx) - * Int64(cute.size(acc_output.shape[3])) - * Int64(physical_tile_rows) - * Int64(split_kv) - * Int64(cfg.latent_dim) - + Int64(query_tile_idx) - * Int64(physical_tile_rows) - * Int64(split_kv) - * Int64(cfg.latent_dim) - + Int64(row_in_tile) * Int64(split_kv) * Int64(cfg.latent_dim) + Int64(batch_idx) * Int64(acc_output.stride[4]) + + Int64(query_tile_idx) * Int64(acc_output.stride[3]) + + Int64(row_in_tile) * Int64(acc_output.stride[0]) + Int64(element_idx) ) partial_output_ptr = acc_output_ptr + partial_elem_offset_base @@ -277,7 +286,7 @@ def run_reduction_kernel( scale = Float32(smem_lse_scale[row_scale_offset + Int32(i)]) acc_vec = acc_vec + partial_vec * scale else: - for i in range(local_split_kv): + for i in cutlass.range(local_split_kv, unroll=4): partial_vec = ( (partial_output_ptr + Int64(i) * Int64(cfg.latent_dim)) .load( diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/resources.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/resources.py index 4e3d24d11fd6..b655e8a73a09 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/resources.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/resources.py @@ -37,7 +37,7 @@ - GmemOResource : Direct stores, or Async(2 stages) Correction -> TMA -> GMEM """ -from dataclasses import dataclass, field +from dataclasses import dataclass, field, fields as dataclass_fields from typing import Any, ClassVar from ..helpers.constants import LN_2, LOG2_E @@ -46,6 +46,7 @@ import cutlass.cute as cute from cutlass.experimental import primitives as prims from cutlass import Boolean, Float32, Int16, Int32, Int64 +from cutlass.cutlass_dsl import dsl_user_op, if_generate from cutlass.experimental.task_scheduling.resources import ( MemoryResource, @@ -55,6 +56,7 @@ ) from cutlass.experimental.task_scheduling.resources import consumer_work, producer_work from cutlass.experimental.task_scheduling.enums import WorkAttr +from ..helpers.math import final_lse from ...mask import kv_tile_needs_right_mask from ...tensor_map import transform_ragged_coords @@ -69,6 +71,9 @@ CachedSparseMask, ) +from .nvfp4_kv_load import load_nvfp4_kv_stage, publish_nvfp4_v +from ..helpers.sparse_indices import cache_sparse_source_lengths +from ..helpers.nvfp4_kv import prefetch_nvfp4_kv_row from .config import ( OUTPUT_TMA_STAGE_BYTES, V_SMEM_K_BLOCK_TOKENS, @@ -116,13 +121,14 @@ qkv_major_k_stride_bytes_for, ) from ..helpers.ops import ( - fp8_log2_quant_scale, - fp8_quant_scale_rcp, + fp8_log2_p_scale, + fp8_p_scale_rcp, fmax_f32, pack_float4_to_fp8_e4m3, ) from ..helpers.mask import MaskType, mask_visible_k_length from ..helpers.query import ( + copy_q_fp8_to_tmem, flat_query_row_state, query_batch_bounds, runtime_flat_query_tile_has_rows, @@ -133,6 +139,52 @@ ) +# Both CTAs publish converted payloads to the leader with a cluster release. +# The consumer must acquire at the same scope before reading either payload. +def with_cluster_acquire_waits(base: object) -> object: + """Return ``base`` (a TmaUmma pipeline) with consumer waits/try-waits on the full barriers using acquire at cluster scope.""" + cls = type(base) + + class ClusterAcquirePipeline(cls): # type: ignore[valid-type, misc] + @cute.jit + def try_cluster_acquire(self, state) -> Boolean: + """Try the full-stage barrier with cluster acquire ordering for both CTAs' payloads.""" + bar = self.sync_object_full.get_barrier(state.index) + return prims.mbarrier_wait_parity( + bar, + state.phase, + prims.MBarrierWait.TRY, + scope=prims.MBarrierScope.CLUSTER, + order=prims.MemOrder.ACQUIRE, + ) + + @cute.jit + def wait_cluster_acquire(self, state) -> None: + """Wait until the full stage is visible under the producer's cluster release.""" + while not self.try_cluster_acquire(state): + pass + + @dsl_user_op + def consumer_try_wait(self, state, *, loc=None, ip=None) -> Boolean: + """Adapt the pipeline try-wait hook to the cluster-scope payload acquire.""" + return self.try_cluster_acquire(state) + + @dsl_user_op + def consumer_wait( + self, state, try_wait_token=None, *, loc=None, ip=None + ) -> None: + """Complete the cluster acquire unless an earlier successful try-wait already did so.""" + if_generate( + try_wait_token is None or try_wait_token == 0, + lambda: self.wait_cluster_acquire(state), + loc=loc, + ip=ip, + ) + + values = {f.name: getattr(base, f.name) for f in dataclass_fields(base)} + return ClusterAcquirePipeline(**values) + + def _install_task_local_specs(resource: object, specs: tuple[tuple, ...]) -> None: """Install TaskLocalVariable fields declared by resource classes.""" for spec in specs: @@ -311,7 +363,7 @@ class MlaWorkQueue(WorkQueue): static_seq_len_k: cutlass.Constexpr = None cu_seqlens_q: Any = None logical_num_heads_q: cutlass.Constexpr[int] = 128 - logical_seq_len_q: cutlass.Constexpr[int] = 1 + logical_seq_len_q: Int32 = 1 static_problem_shape_b: cutlass.Constexpr[int] = None static_problem_shape_s: cutlass.Constexpr[int] = None use_clc_dynamic: cutlass.Constexpr[bool] = False @@ -443,6 +495,8 @@ def _work_tile_from_clc_tile(self, work_tile): work_tile.is_valid_tile, ) + # Bind a scheduler tile to its live causal KV range and split; inactive Q gets zero + # tiles. @cute.jit def _make_work_tile_info( self, @@ -476,9 +530,7 @@ def _make_work_tile_info( self.cu_seqlens_q, safe_b_idx, ) - if cutlass.const_expr( - cfg.mask_type == MaskType.CAUSAL.value and self.logical_seq_len_q > 1 - ): + if cutlass.const_expr(cfg.mask_type == MaskType.CAUSAL.value): # Split partitioning owns the mask-visible CTA domain. The final # physical row resolves to the tile's latest valid Q # token. Row-causal softmax narrows earlier rows without changing @@ -521,6 +573,9 @@ def _make_work_tile_info( k_tile_total, split_kv_cap, split_kv_idx, + split_divisor=params.split_kv_fdd + if self.static_split_kv is None and not self.is_var_split_kv + else None, ) return MlaTsWorkTileInfo( tile_idx, @@ -646,6 +701,8 @@ class PageOffsetWindowResource(HighThroughputMlaResource): """ request_lengths: Any = field(init=False, default=None) + primary_cache: Any = None + extra_cache: Any = None page_offsets: Any = None # GMEM page-offset tensor smem_sparse_offsets: Any = None cfg: cutlass.Constexpr = field(default_factory=MlaDecodeConfig) @@ -664,6 +721,9 @@ class PageOffsetWindowResource(HighThroughputMlaResource): cached_window_page: cutlass.Constexpr[TaskLocalVariable] = ( TaskLocalVariable.uninitialized() ) + cached_next_window_page: cutlass.Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) cached_mask_low: cutlass.Constexpr[TaskLocalVariable] = ( TaskLocalVariable.uninitialized() ) @@ -703,6 +763,14 @@ class PageOffsetWindowResource(HighThroughputMlaResource): Int32(0), "Per-lane page id retained across one 32-entry page-table window.", ), + ( + "cached_next_window_page", + Int32, + Int32(0), + "This lane's page id in the next 32-entry window (dense FP8 " + "loaders), loaded one window ahead so that a window boundary " + "does not wait on the page table.", + ), ) @producer_work(work_attrs=WorkAttr.AUXILIARY) @@ -726,6 +794,12 @@ def bind_sparse_request(self, stage_info: StageInfo): @producer_work @cute.jit def load_sparse_offsets(self, stage_info: StageInfo): + """Stage source-tagged physical row IDs for the next K/V tile. + + With assume_valid_prefix, entries before each live length contain no + -1 holes, so each CTA can retain only its own 64 rows. Otherwise keep + the full 128-row metadata tile for score masking. Valid-prefix does + not imply contiguous physical rows, sorted indices, or no tail padding.""" lane = cute.arch.thread_idx()[0] % 32 k_index = stage_info.work_tile.k_index_base + Int32(stage_info.loop_offset) token = k_index * Int32(128) + lane * Int32(4) @@ -739,11 +813,28 @@ def load_sparse_offsets(self, stage_info: StageInfo): ) else: routes = self.page_offsets - quad = load_sparse_quad(routes, row, token) - self.smem_sparse_offsets.data_ptr(stage_info.stage_idx * 128 + lane * 4).store( - quad.load(0, 4), alignment=16 - ) - cute.arch.fence_view_async_shared() + active = cutlass.Boolean(True) + if cutlass.const_expr( + self.cfg.has_nvfp4_kv + and getattr(self.page_offsets, "assume_valid_prefix", False) + ): + active = lane // Int32(16) == cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + if active: + quad = load_sparse_quad(routes, row, token) + if cutlass.const_expr(self.cfg.has_nvfp4_kv): + cta = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + if lane // Int32(16) == cta: + for j in cutlass.range_constexpr(4): + prefetch_nvfp4_kv_row( + self.primary_cache, self.extra_cache, quad[j], self.cfg + ) + + self.smem_sparse_offsets.data_ptr( + stage_info.stage_idx * 128 + lane * 4 + ).store(quad.load(0, 4), alignment=16) + cute.arch.fence_view_async_shared() @consumer_work(returns=(cached_mask_low, cached_mask_high)) @cute.jit @@ -1020,6 +1111,159 @@ def read_page_offset_window( ) return cached_k, cached_v, cached_next_v, cached_window_page + @cute.jit + def advance_dense_page_window( + self, stage_info: StageInfo, cached_window_page, cached_next_window_page + ): + """Advance the FP8 dense loader's prefetched page-ID window. + + A window is one page ID per lane (32 per warp). Each lane holds its ID + of the current window and of the next. The next window's load is + issued at the start of the current window (``32 / pages_per_k_tile`` K + tiles before it is used); ptxas copies the loop-carried register at the + next K tile, so the load has at least one K tile to land. Only the + first window of a work tile is loaded synchronously, so no page-table + read sits between a freed K/V slot and its TMA. Indices clamp to the + last table column; fragments past the sequence are masked by the + runtime length. + """ + cfg = self.cfg + pages_per_tile = cutlass.const_expr(cfg.pages_per_k_tile) + window_tiles = cutlass.const_expr(WARP_LANES // pages_per_tile) + work_tile = stage_info.work_tile + local_k_index = Int32(stage_info.loop_offset) + global_k_index = work_tile.k_index_base + local_k_index + page_offsets_batch = self.page_offsets[None, work_tile.tile_idx[2]] + last_page_col = Int32(page_offsets_batch.shape[0] - 1) + lane_idx = cute.arch.thread_idx()[0] & Int32(31) + window_pos = local_k_index & Int32(window_tiles - 1) + if local_k_index == Int32(0): + cached_window_page = Int32( + page_offsets_batch[ + cute.math.min( + global_k_index * Int32(pages_per_tile) + lane_idx, + last_page_col, + ) + ] + ) + if window_pos == Int32(0): + if local_k_index != Int32(0): + cached_window_page = cached_next_window_page + # Short split-KV work tiles fit in one window; skip a dead prefetch. + if local_k_index + Int32(window_tiles) < Int32(work_tile.k_tile_count): + cached_next_window_page = Int32( + page_offsets_batch[ + cute.math.min( + (global_k_index + Int32(window_tiles)) + * Int32(pages_per_tile) + + lane_idx, + last_page_col, + ) + ] + ) + return cached_window_page, cached_next_window_page, window_pos + + @consumer_work( + work_attrs=WorkAttr.AUXILIARY, + returns=(cached_k_pages, cached_window_page, cached_next_window_page), + ) + @cute.jit + def init_dense_k_window(self, stage_info: StageInfo): + """Create the FP8 dense K loader's page-ID window state.""" + del stage_info + return ( + cutlass.Array( + Int32, self.cfg.pages_per_k_cta, space=cutlass.AddressSpace.rmem + ), + Int32(0), + Int32(0), + ) + + @consumer_work( + returns=(cached_k_pages, cached_window_page, cached_next_window_page) + ) + @cute.jit + def read_dense_k_window( + self, + stage_info: StageInfo, + *, + cached_k_pages, + cached_window_page, + cached_next_window_page, + ): + """Return this CTA's K page IDs for the current tile from the window.""" + cfg = self.cfg + cached_window_page, cached_next_window_page, window_pos = ( + self.advance_dense_page_window( + stage_info, cached_window_page, cached_next_window_page + ) + ) + cta_v = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + for pk in cutlass.range_constexpr(cfg.pages_per_k_cta): + # Same logical page order as ``SmemKResource.tma_load_paged``. + source_lane = window_pos * Int32(cfg.pages_per_k_tile) + ( + cta_v * Int32(cfg.pages_per_k_cta) + Int32(pk) + if cfg.pages_per_k_tile > 1 + else Int32(0) + ) + cached_k_pages[pk] = Int32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=cached_window_page, + offset=source_lane, + mask_and_clamp=0x1F, + kind=prims.Shfl.IDX, + ) + ) + return cached_k_pages, cached_window_page, cached_next_window_page + + @consumer_work( + work_attrs=WorkAttr.AUXILIARY, + returns=(cached_v_pages, cached_window_page, cached_next_window_page), + ) + @cute.jit + def init_dense_v_window(self, stage_info: StageInfo): + """Create the FP8 dense V loader's page-ID window state.""" + del stage_info + return ( + cutlass.Array( + Int32, self.cfg.pages_per_v_tile, space=cutlass.AddressSpace.rmem + ), + Int32(0), + Int32(0), + ) + + @consumer_work( + returns=(cached_v_pages, cached_window_page, cached_next_window_page) + ) + @cute.jit + def read_dense_v_window( + self, + stage_info: StageInfo, + *, + cached_v_pages, + cached_window_page, + cached_next_window_page, + ): + """Return the V page IDs of the current tile from the window.""" + cfg = self.cfg + cached_window_page, cached_next_window_page, window_pos = ( + self.advance_dense_page_window( + stage_info, cached_window_page, cached_next_window_page + ) + ) + for pk in cutlass.range_constexpr(cfg.pages_per_v_tile): + cached_v_pages[pk] = Int32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=cached_window_page, + offset=window_pos * Int32(cfg.pages_per_v_tile) + Int32(pk), + mask_and_clamp=0x1F, + kind=prims.Shfl.IDX, + ) + ) + return cached_v_pages, cached_window_page, cached_next_window_page + @consumer_work( work_attrs=WorkAttr.AUXILIARY, returns=(cached_k_pages, cached_v_pages, cached_next_v_pages), @@ -1052,12 +1296,13 @@ class SmemQResource(HighThroughputMlaResource): """ smem_q_latent: Any = None # SMEM Q latent array + tmem_base_addr: Any = None smem_q_rope: Any = None # SMEM Q rope array tma_desc_q_latent: Any = None tma_desc_q_rope: Any = None cu_seqlens_q: Any = None logical_num_heads_q: cutlass.Constexpr[int] = 128 - logical_seq_len_q: cutlass.Constexpr[int] = 1 + logical_seq_len_q: Int32 = 1 cfg: cutlass.Constexpr = field(default_factory=MlaDecodeConfig) cta_rank: Any = field(init=False, default=None) is_leader: Any = field(init=False, default=None) @@ -1179,8 +1424,15 @@ def _tma_load(self, stage_info: StageInfo) -> None: @consumer_work @cute.jit def q_desc(self, stage_info: StageInfo) -> None: - """Schedule marker after the Q SMEM stage is waited.""" + """Stage64 local Q rows into the replicated group-2 TMEM A layout.""" del stage_info + if cutlass.const_expr(self.cfg.has_nvfp4_kv): + if prims.elect_sync(): + copy_q_fp8_to_tmem( + self.smem_q_latent.data_ptr(), + self.tmem_base_addr + Int32(384), + cta_group=2, + ) # ===================================================================== @@ -1199,6 +1451,8 @@ class SmemKVResource(HighThroughputMlaResource): smem_kv: Any = None tma_desc_c_latent: Any = None tma_desc_c_rope: Any = None + tma_desc_sparse_rope: Any = None + tma_desc_sparse_extra_rope: Any = None tma_desc_c_transpose: Any = None tma_desc_extra_v: Any = None cfg: cutlass.Constexpr = field(default_factory=MlaDecodeConfig) @@ -1246,6 +1500,9 @@ def tma_load( warp = ( cute.arch.make_warp_uniform(cute.arch.warp_idx()) - cfg.load_tma_warp_id ) + rope = not is_v and subtile_idx >= cfg.iterations_qk_latent_stages + primary_k = self.tma_desc_sparse_rope if rope else self.tma_desc_c_latent + extra_k = self.tma_desc_sparse_extra_rope if rope else self.tma_desc_c_rope # Each issuer owns disjoint four-row quads. K covers 64 rows per # CTA (16 quads); V copies 32-token subtiles (8 quads). BF16 V is # blocked as 32x64 panels: two 64-feature/128-byte slices cover @@ -1282,10 +1539,12 @@ def tma_load( self.smem_kv.data_ptr( base + group * 4 * cfg.mma_qk_tiler_k ), - self.tma_desc_c_latent, - self.tma_desc_c_rope, + primary_k, + extra_k, Int32( - subtile_idx + 0 + if rope + else subtile_idx * cfg.kv_subtiles_per_stage * cfg.mma_qk_tiler_k ), @@ -1293,7 +1552,7 @@ def tma_load( stage_info.barrier, num_slices=cfg.kv_subtiles_per_stage, dst_stride=64 * cfg.mma_qk_tiler_k * cfg.qkv_dtype_bytes, - col_stride=cfg.mma_qk_tiler_k, + col_stride=0 if rope else cfg.mma_qk_tiler_k, ) else: # PV consumes the delayed current-V window; the tail can @@ -1346,7 +1605,9 @@ def tma_load( if cutlass.const_expr(not is_v): for part in cutlass.range_constexpr(cfg.kv_subtiles_per_stage): column = Int32( - (subtile_idx * cfg.kv_subtiles_per_stage + part) + 0 + if rope + else (subtile_idx * cfg.kv_subtiles_per_stage + part) * cfg.mma_qk_tiler_k ) for local_group in cutlass.range( @@ -1359,8 +1620,8 @@ def tma_load( + part * 64 * cfg.mma_qk_tiler_k + group * 4 * cfg.mma_qk_tiler_k ), - self.tma_desc_c_latent, - self.tma_desc_c_rope, + primary_k, + extra_k, column, cached_k_pages, group * 4, @@ -1657,12 +1918,19 @@ class SmemKResource(HighThroughputMlaResource): all latent K sub-tiles and optional RoPE channels for one logical K tile. """ + # Source metadata for NVFP4 gathers; dense loaders consume the page window. + page_offsets: Any = None smem_k: Any = None smem_sparse_offsets: Any = None - page_offsets: Any = None tma_desc_c_latent: Any = None tma_desc_c_rope: Any = None - logical_seq_len_q: cutlass.Constexpr[int] = 1 + tma_desc_sparse_rope: Any = None + tma_desc_sparse_extra_rope: Any = None + paired_v: cutlass.Constexpr = None + smem_raw: Any = None + raw_mbar: Any = None + request_lengths: Any = field(init=False, default=None) + logical_seq_len_q: Int32 = 1 cfg: cutlass.Constexpr = field(default_factory=MlaDecodeConfig) desc_k_base: cutlass.Constexpr[TaskLocalVariable] = ( TaskLocalVariable.uninitialized() @@ -1676,12 +1944,46 @@ class SmemKResource(HighThroughputMlaResource): def init_load_state(self, stage_info: StageInfo) -> None: """Initialize per-warp state used by K TMA loads.""" del stage_info + if cutlass.const_expr(self.cfg.has_nvfp4_kv): + self.request_lengths = cutlass.Array( + Int32, + 3, + space=cutlass.AddressSpace.rmem, + ) + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def bind_sparse_request(self, stage_info: StageInfo): + """Cache live source lengths for this query before its KV/softmax loop.""" + if cutlass.const_expr(self.cfg.has_nvfp4_kv): + cache_sparse_source_lengths( + self.page_offsets, + stage_info.work_tile.tile_idx[2], + self.request_lengths, + include_primary_tiles=True, + ) + + def create_pipeline(self, pipeline_config): + """Use cluster-scope consumer acquires for K published by both conversion CTAs.""" + base = super().create_pipeline(pipeline_config) + if self.cfg.has_nvfp4_kv: + return with_cluster_acquire_waits(base) + return base @producer_work @cute.jit def tma_load_shared(self, stage_info: StageInfo, *, page_offset_stage): """Read one warp-uniform SMEM quad and reuse it across FP8 feature slices.""" cfg = self.cfg + if cutlass.const_expr(cfg.has_nvfp4_kv): + load_nvfp4_kv_stage( + self, + self.paired_v, + stage_info, + self.paired_v.current_slot[0], + page_offset_stage, + ) + return warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - cfg.load_k_warp_id base = ( cute.arch.make_warp_uniform(stage_info.stage_idx) * cfg.smem_k_stage_elems @@ -1710,20 +2012,33 @@ def tma_load_shared(self, stage_info: StageInfo, *, page_offset_stage): dst_stride=64 * cfg.mma_qk_tiler_k, col_stride=cfg.mma_qk_tiler_k, ) + if cutlass.const_expr(cfg.rope_dim): + gather4_uniform_quad_slices( + self.smem_k.data_ptr( + base + 64 * cfg.latent_dim + quad * 4 * cfg.rope_dim + ), + self.tma_desc_sparse_rope, + self.tma_desc_sparse_extra_rope, + Int32(0), + raw, + stage_info.barrier, + num_slices=1, + dst_stride=64 * cfg.rope_dim, + col_stride=0, + ) @producer_work @cute.jit - def tma_load_paged(self, stage_info: StageInfo) -> None: - """Load a paged K tile; page-size-one inputs use the offset ring.""" + def tma_load_paged(self, stage_info: StageInfo, *, cached_page_ids) -> None: + """Load a paged K tile with this CTA's page IDs from the page window. + + ``PageOffsetWindowResource.read_dense_k_window`` supplies the IDs in + the logical order ``k_index * pages_per_k_tile + cta * pages_per_k_cta + + pk`` (or ``k_index`` when one page spans the tile). + """ cfg = self.cfg stage_idx = stage_info.stage_idx stage_base = stage_idx * cfg.smem_k_stage_elems - work_tile = stage_info.work_tile - blk_coord = work_tile.tile_idx - k_index = work_tile.k_index_base + Int32(stage_info.loop_offset) - - page_row_idx = blk_coord[2] - page_offsets_batch = self.page_offsets[None, page_row_idx] kv_mbar_arr = cutlass.Array(stage_info.barrier.data_ptr(), dtype=Int64) pages_per_k_cta = cfg.pages_per_k_cta kc_page_smem_elems = cfg.kc_page_tile_size * cfg.mma_qk_tiler_k @@ -1735,24 +2050,14 @@ def tma_load_paged(self, stage_info: StageInfo) -> None: k_latent_subtile_elems = cutlass.const_expr( cfg.mma_qk_tiler[1] // cfg.num_mma_ctas * cfg.mma_qk_tiler_k ) - # Use page 0 for masked fragments beyond a compact page table. A direct - # bounds predicate avoids carrying runtime K-length division or a full - # page-ID register array through this producer warp. + # Fragments past the page table use the window's clamped last column; + # the runtime sequence length masks them. for k_call in cutlass.range_constexpr(cfg.iterations_qk_latent): coord_kcl = cutlass.Int32(k_call * cfg.mma_qk_tiler_k) subtile_base = stage_base + k_call * k_latent_subtile_elems for pk in cutlass.range_constexpr(pages_per_k_cta): - logical_page_idx = ( - k_index - if cfg.pages_per_k_tile == 1 - else (k_index * Int32(cfg.num_mma_ctas) + cta_v) - * Int32(pages_per_k_cta) - + Int32(pk) - ) - page_idx = Int32(0) - if cute.elem_less(logical_page_idx, page_offsets_batch.shape[0]): - page_idx = page_offsets_batch[logical_page_idx] + page_idx = Int32(cached_page_ids[pk]) if prims.elect_sync(): kcl_smem = cutlass.Array( self.smem_k.data_ptr(subtile_base + pk * kc_page_smem_elems), @@ -1777,16 +2082,7 @@ def tma_load_paged(self, stage_info: StageInfo) -> None: coord_kcr = cutlass.Int32(rope_idx * cfg.mma_qk_rope_tiler[2]) subtile_base = rope_stage_base + rope_idx * k_rope_subtile_elems for pk in cutlass.range_constexpr(pages_per_k_cta): - logical_page_idx = ( - k_index - if cfg.pages_per_k_tile == 1 - else (k_index * Int32(cfg.num_mma_ctas) + cta_v) - * Int32(pages_per_k_cta) - + Int32(pk) - ) - page_idx = Int32(0) - if cute.elem_less(logical_page_idx, page_offsets_batch.shape[0]): - page_idx = page_offsets_batch[logical_page_idx] + page_idx = Int32(cached_page_ids[pk]) if prims.elect_sync(): kcr_smem = cutlass.Array( self.smem_k.data_ptr(subtile_base + pk * rope_page_smem_elems), @@ -1868,14 +2164,18 @@ class SmemVResource(HighThroughputMlaResource): smem_v: Any = None smem_sparse_offsets: Any = None - page_offsets: Any = None tma_desc_c_transpose: Any = None tma_desc_extra_v: Any = None - logical_seq_len_q: cutlass.Constexpr[int] = 1 + tma_desc_sparse_rope: Any = None + tma_desc_sparse_extra_rope: Any = None + smem_raw: Any = None + raw_mbar: Any = None + logical_seq_len_q: Int32 = 1 cfg: cutlass.Constexpr = field(default_factory=MlaDecodeConfig) desc_v_base: cutlass.Constexpr[TaskLocalVariable] = ( TaskLocalVariable.uninitialized() ) + current_slot: Any = None _task_local_specs: ClassVar[tuple[tuple, ...]] = ( ("desc_v_base", Int64, Int64(0), "SMEM descriptor for staged V."), ) @@ -1885,12 +2185,35 @@ class SmemVResource(HighThroughputMlaResource): def init_load_state(self, stage_info: StageInfo) -> None: """Initialize per-warp state used by V TMA loads.""" del stage_info + if cutlass.const_expr(self.cfg.has_nvfp4_kv): + self.current_slot = cutlass.Array(Int32, 2, space=cutlass.AddressSpace.rmem) + + def create_pipeline(self, pipeline_config): + """Use cluster-scope consumer acquires for V assembled from local and peer data.""" + base = super().create_pipeline(pipeline_config) + if self.cfg.has_nvfp4_kv: + return with_cluster_acquire_waits(base) + return base + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def combined_stage(self, stage_info: StageInfo): + """Expose the acquired V stage index and phase to the paired K/V producer.""" + self.current_slot[0] = stage_info.stage_idx + self.current_slot[1] = self.producer_state.phase ^ Int32(1) @producer_work @cute.jit def tma_load_shared(self, stage_info: StageInfo, *, page_offset_stage): """Read one warp-uniform SMEM quad and reuse it across FP8 V panels.""" cfg = self.cfg + if cutlass.const_expr(cfg.has_nvfp4_kv): + # Match this warp's share of the cluster-wide V pipeline credit. + publish_nvfp4_v( + stage_info.barrier, + cfg.tma_copy_v_tile_bytes // cfg.num_mma_ctas // cfg.load_num_warps, + ) + return warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - cfg.load_v_warp_id cta = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) base = ( @@ -1923,17 +2246,15 @@ def tma_load_shared(self, stage_info: StageInfo, *, page_offset_stage): @producer_work @cute.jit - def tma_load_paged(self, stage_info: StageInfo) -> None: - """Load a paged V tile; page-size-one inputs use the offset ring.""" + def tma_load_paged(self, stage_info: StageInfo, *, cached_page_ids) -> None: + """Load a paged V tile with the tile's page IDs from the page window. + + ``PageOffsetWindowResource.read_dense_v_window`` supplies the IDs in + the logical order ``k_index * pages_per_v_tile + page_offset``. + """ cfg = self.cfg stage_idx = stage_info.stage_idx stage_base = stage_idx * cfg.smem_v_stage_elems - work_tile = stage_info.work_tile - blk_coord = work_tile.tile_idx - k_index = work_tile.k_index_base + Int32(stage_info.loop_offset) - - page_row_idx = blk_coord[2] - page_offsets_batch = self.page_offsets[None, page_row_idx] v_mbar_arr = cutlass.Array(stage_info.barrier.data_ptr(), dtype=Int64) cta_v = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) @@ -1946,8 +2267,9 @@ def tma_load_paged(self, stage_info: StageInfo) -> None: svc_copy_elems = cutlass.const_expr( cfg.mma_pv_tiler[1] // cfg.num_mma_ctas * cfg.mma_pv_tiler[2] ) - # As for K, out-of-table fragments use page 0 and are masked by the - # runtime sequence length. p128 still shares one page ID across CTAs. + # As for K, out-of-table fragments use the clamped last column and are + # masked by the runtime sequence length. p128 shares one page ID across + # CTAs. for v_call in cutlass.range_constexpr( cfg.iterations_pv_k * cfg.iterations_pv_n @@ -1965,14 +2287,9 @@ def tma_load_paged(self, stage_info: StageInfo) -> None: tile_token_offset = pv_i * cfg.mma_pv_tiler[2] + local_token_offset page_offset = tile_token_offset // cfg.page_size coord_k_v = Int32(tile_token_offset % cfg.page_size) - logical_page_idx = ( - k_index - if pages_per_v_tile == 1 - else k_index * Int32(pages_per_v_tile) + Int32(page_offset) + page_idx = Int32( + cached_page_ids[page_offset if pages_per_v_tile > 1 else 0] ) - page_idx = Int32(0) - if cute.elem_less(logical_page_idx, page_offsets_batch.shape[0]): - page_idx = page_offsets_batch[logical_page_idx] if prims.elect_sync(): smem_k_block_idx = local_token_offset // V_SMEM_K_BLOCK_TOKENS token_offset_in_k_block = local_token_offset % V_SMEM_K_BLOCK_TOKENS @@ -2113,13 +2430,13 @@ class TmemSResource(HighThroughputMlaResource): smem_q_rope: Any = None smem_p: Any = None # SMEM P array smem_exchange: Any = None # SMEM array for cross-warp max exchange - softmax_scale_log2: Any = None # softmax_scale * log2(e) + scale_softmax_log2: Any = None # softmax_scale * log2(e) page_offsets: Any = None cache_seqs: Any = None # per-batch valid K length cu_seqlens_q: Any = None # cumulative compact-Q offsets, or None for fixed Q split_kv: Any = None # per-work-tile split count logical_num_heads_q: cutlass.Constexpr[int] = 128 - logical_seq_len_q: cutlass.Constexpr[int] = 1 + logical_seq_len_q: Int32 = 1 cfg: cutlass.Constexpr = field(default_factory=MlaDecodeConfig) tiled_mma_qk: Any = None cta_rank: Any = field(init=False, default=None) @@ -2285,14 +2602,14 @@ def init_softmax_state_odd(self, stage_info: StageInfo): @consumer_work(work_attrs=WorkAttr.AUXILIARY) @cute.jit def bind_sparse_request(self, stage_info: StageInfo): - """Cache live source lengths once per query tile, including graph replay.""" - if cutlass.const_expr( - isinstance(self.page_offsets, SparseKVIndices) - and not self.page_offsets.is_bound - ): - row = stage_info.work_tile.tile_idx[2] - self.request_lengths[0] = Int32(self.page_offsets.primary_lengths[row]) - self.request_lengths[1] = Int32(self.page_offsets.extra_lengths[row]) + """Cache live source lengths for this query before its KV/softmax loop.""" + if cutlass.const_expr(isinstance(self.page_offsets, SparseKVIndices)): + cache_sparse_source_lengths( + self.page_offsets, + stage_info.work_tile.tile_idx[2], + self.request_lengths, + include_primary_tiles=False, + ) @producer_work @cute.jit @@ -2366,15 +2683,36 @@ def qk_mma( logical_call_idx > 0 or k_block > 0 ) if prims.elect_sync(): - prims.tcgen05_mma( - mma_kind, - cta_group, - tmem_s_ptr, - desc_q + k_block * 2, - desc_k_subtile + k_block * 2, - idesc_qk, - Boolean(scale_d), - ) + if cutlass.const_expr(cfg.has_nvfp4_kv): + # Q NoPE occupies TMEM columns 384:512: + # 512 E4M3 bytes / 4 bytes per column = 128. + # Each D128 slice advances 32 columns; each + # K32 MMA step advances another 8 columns. + prims.tcgen05_mma( + mma_kind, + cta_group, + tmem_s_ptr, + prims.make_tmem_ptr( + self.tmem_base_addr + + Int32( + 384 + logical_call_idx * 32 + k_block * 8 + ), + qkv_dtype(cfg), + ), + desc_k_subtile + k_block * 2, + idesc_qk, + Boolean(scale_d), + ) + else: + prims.tcgen05_mma( + mma_kind, + cta_group, + tmem_s_ptr, + desc_q + k_block * 2, + desc_k_subtile + k_block * 2, + idesc_qk, + Boolean(scale_d), + ) elif cutlass.const_expr(call_idx < cfg.iterations_qk_stages): # K rope: build Q rope descriptor unconditionally @@ -2488,7 +2826,7 @@ def _load_s_impl( k_index = work_tile.k_index_base + Int32(stage_info.loop_offset) tile_offset_k = k_index * Int32(cfg.mma_qk_tiler[1]) needs_row_causal_mask = cutlass.const_expr( - cfg.mask_type == MaskType.CAUSAL.value and self.logical_seq_len_q > 1 + cfg.mask_type == MaskType.CAUSAL.value ) if cutlass.const_expr(needs_row_causal_mask): batch_idx = Int32(work_tile.tile_idx[2]) @@ -2544,6 +2882,15 @@ def _load_s_impl( ) else: mask_page_offsets = self.page_offsets + if cutlass.const_expr(cfg.has_fp8_primary): + scale = mask_page_offsets.tensor_scale(tile_offset_k) + if scale != cutlass.Float32(1): + for i in cutlass.range_constexpr(0, 64, 2): + pair = mul_packed_f32x2( + (qk_acc_regs[i], qk_acc_regs[i + 1]), (scale, scale) + ) + qk_acc_regs[i], qk_acc_regs[i + 1] = pair[0], pair[1] + if cutlass.const_expr(cached_mask_low is not None): mask_routes = CachedSparseMask((cached_mask_low, cached_mask_high), False) # The cached invalid bits already include the live K boundary. @@ -2568,13 +2915,20 @@ def _load_s_impl( column_half=local_tidx >> EPILOGUE_COLUMN_GROUP_SHIFT, ) if cutlass.const_expr( - cfg.page_size == 1 and self.logical_seq_len_q == 1 + cfg.page_size == 1 + and ( + isinstance(self.logical_seq_len_q, int) + and self.logical_seq_len_q == 1 + ) ) else raw_mask_routes ) if cutlass.const_expr( cfg.page_size == 1 - and self.logical_seq_len_q == 1 + and ( + isinstance(self.logical_seq_len_q, int) + and self.logical_seq_len_q == 1 + ) and not needs_row_causal_mask ): group_needs_mask = (mask_routes.low | mask_routes.high) != Int32(0) @@ -2838,6 +3192,16 @@ def load_peer_state(): row_max_prev, row_sum_prev = load_peer_state() row_max_new = fmax_f32(row_max_prev, row_max_tile) + if cutlass.const_expr(cfg.softmax_rescale_threshold_log2 > 0.0): + # Freeze the running row max while this K tile raises it by at most + # the threshold. P then grows by at most 2^T, and the unchanged max + # keeps this row's O rescale on the existing exact-equality skip + # path below. The first tile (row_max_prev == -inf) yields an + # infinite or NaN increase and never freezes. Compare in log2 units with + # the threshold as an immediate: the scale is already live here. + max_increase_log2 = (row_max_new - row_max_prev) * self.scale_softmax_log2 + if max_increase_log2 <= Float32(cfg.softmax_rescale_threshold_log2): + row_max_new = row_max_prev row_has_values = row_max_new != neg_inf safe_row_max_prev = row_max_prev if row_has_values else Float32(0) safe_row_max_new = row_max_new if row_has_values else Float32(0) @@ -2847,17 +3211,17 @@ def load_peer_state(): correction_factor = Float32(1) if max_changed: correction_factor = cute.math.exp2( - (safe_row_max_prev - safe_row_max_new) * self.softmax_scale_log2, + (safe_row_max_prev - safe_row_max_new) * self.scale_softmax_log2, fastmath=True, ) no_correction = Int32(not max_changed) - fma_b = self.softmax_scale_log2 - fma_c = Float32(0) - safe_row_max_new * self.softmax_scale_log2 + fma_b = self.scale_softmax_log2 + fma_c = Float32(0) - safe_row_max_new * self.scale_softmax_log2 if cutlass.const_expr(cfg.is_fp8_qkv()): - # Use the same selected FP8 probability scale as output and LSE. - # Native FP8 uses exact maxima and scale 448. - fma_c = fma_c + fp8_log2_quant_scale() + # Use the same selected FP8 probability scale as output and LSE: + # 448 with exact maxima, 448 * 2**-T with a rescale threshold T. + fma_c = fma_c + fp8_log2_p_scale(cfg) for i in cutlass.range_constexpr(0, 64, 2): fma_result = fma_packed_f32x2( (qk_acc_regs[i], qk_acc_regs[i + 1]), @@ -3144,6 +3508,8 @@ class TmemCorrResource(HighThroughputMlaResource): cfg: cutlass.Constexpr = field(default_factory=MlaDecodeConfig) cta_rank: Any = field(init=False, default=None) final_row_stats: Any = field(init=False, default=None) + score_resource: Any = None + page_offsets: Any = None row_sum: cutlass.Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() row_max: cutlass.Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() epilogue_row_sum: cutlass.Constexpr[TaskLocalVariable] = ( @@ -3219,6 +3585,27 @@ def _store_corr_impl( correction_regs = cutlass.Array(Float32, 4, space=cutlass.AddressSpace.rmem) correction_regs[0] = row_sum_out correction_regs[1] = row_max_new + if cutlass.const_expr(cfg.has_fp8_primary): + work = stage_info.work_tile + current = (work.k_index_base + Int32(stage_info.loop_offset)) * Int32( + cfg.mma_qk_tiler[1] + ) + if cutlass.const_expr(self.page_offsets.is_bound): + routes = self.page_offsets + else: + routes = self.page_offsets.bind_query_row( + work.tile_idx[2], + self.score_resource.request_lengths[0], + self.score_resource.request_lengths[1], + ) + ratio = routes.tensor_scale_ratio( + current, current - Int32(cfg.mma_qk_tiler[1]) + ) + if stage_info.loop_offset != Int32(0): + correction_factor_out *= ratio + no_correction_out = Int32( + no_correction_out != 0 and ratio == cutlass.Float32(1) + ) correction_regs[2] = correction_factor_out correction_regs[3] = prims.mov_b32(no_correction_out, target_type=Float32) @@ -3636,16 +4023,17 @@ class GmemOResource(HighThroughputMlaResource): tmem_o_ref: Any = None # Reference to TmemOResource for tmem_base_addr tmem_corr_ref: Any = None # Reference to TmemCorrResource for correction data output_scale: Any = None - atten_sinks: Any = None + page_offsets: Any = None + attention_sinks: Any = None output_quant: cutlass.Constexpr = None output_quant_args: Any = None lse_in_natural_log: cutlass.Constexpr[bool] = False - softmax_scale_log2: Any = None + scale_softmax_log2: Any = None smem_exchange: Any = None # SMEM for row_sum exchange (as Int32 base addr) split_kv: Any = None cu_seqlens_q: Any = None logical_num_heads_q: cutlass.Constexpr[int] = 128 - logical_seq_len_q: cutlass.Constexpr[int] = 1 + logical_seq_len_q: Int32 = 1 cfg: cutlass.Constexpr = field(default_factory=MlaDecodeConfig) @cute.jit @@ -3666,18 +4054,29 @@ def _prefetch_output_rope(self, flat_row, batch_idx, local_tidx, head, valid): prims.prefetch_l1(cache_row) prims.prefetch_l1(cache_row + 32) + # Normalize one output row, including its head's sink mass and any FP8 probability + # scale. @cute.jit - def _normalization_scale(self, row_sum, row_max, local_tidx, blk_coord): + def _normalization_scale(self, row_sum, row_max, local_tidx, blk_coord, stage_info): + """Compute output normalization including sink mass and the final source's V units. + + Row sums retain the FP8 probability scale. Apply tensor-FP8 V descale + only to O, after denominator/sink normalization; it must not change LSE.""" safe_sum = row_sum if row_sum > Float32(0) else Float32(1) scale = self.output_scale * cute.math.rcp(safe_sum, approx=True) - if cutlass.const_expr(self.atten_sinks is not None): + if cutlass.const_expr(self.attention_sinks is not None): head = ( Int32(blk_coord[1]) * self.cfg.mma_qk_tiler[0] + Int32(blk_coord[0]) * (self.cfg.mma_qk_tiler[0] // self.cfg.num_mma_ctas) + (local_tidx & Int32(63)) ) - if cutlass.const_expr(self.logical_seq_len_q == 1): + if cutlass.const_expr( + ( + isinstance(self.logical_seq_len_q, int) + and self.logical_seq_len_q == 1 + ) + ): # Preserve the cheap single-query mapping, including padded heads. head = cute.math.min(head, Int32(self.logical_num_heads_q - 1)) else: @@ -3687,15 +4086,15 @@ def _normalization_scale(self, row_sum, row_max, local_tidx, blk_coord): _, head, _, _, _ = self._query_row_state( row, Int32(blk_coord[1]), Int32(blk_coord[2]) ) - sink = Float32(self.atten_sinks[head]) + sink = Float32(self.attention_sinks[head]) scale = Float32(0) if row_sum > Float32(0) and sink != Float32(Float32.inf): - max_log2 = row_max * self.softmax_scale_log2 + max_log2 = row_max * self.scale_softmax_log2 sink_log2 = sink * Float32(LOG2_E) norm_max = cute.math.max(max_log2, sink_log2) rescale = cute.math.exp2(max_log2 - norm_max, fastmath=True) inv_quant = Float32( - fp8_quant_scale_rcp() if self.cfg.is_fp8_qkv() else 1.0 + fp8_p_scale_rcp(self.cfg) if self.cfg.is_fp8_qkv() else 1.0 ) denominator = row_sum * inv_quant * rescale + cute.math.exp2( sink_log2 - norm_max, fastmath=True @@ -3706,6 +4105,12 @@ def _normalization_scale(self, row_sum, row_max, local_tidx, blk_coord): * rescale * cute.math.rcp(denominator, approx=True) ) + if cutlass.const_expr(self.cfg.has_fp8_primary): + work = stage_info.work_tile + last = (work.k_index_base + work.k_tile_count - Int32(1)) * Int32( + self.cfg.mma_qk_tiler[1] + ) + scale *= self.page_offsets.tensor_scale(last, work.tile_idx[2]) return scale @cute.jit @@ -3906,7 +4311,9 @@ def epilogue_store(self, stage_info: StageInfo) -> None: # causal query rows have no visible K values in this split. Store zero # O and -inf LSE so split-KV reduction gives those rows zero weight. row_has_values = row_sum > Float32(0) - norm_scale = self._normalization_scale(row_sum, row_max, local_tidx, blk_coord) + norm_scale = self._normalization_scale( + row_sum, row_max, local_tidx, blk_coord, stage_info + ) for iter_n in cutlass.range_constexpr(cfg.iterations_pv_n): # Load O from TMEM @@ -4071,10 +4478,10 @@ def epilogue_store(self, stage_info: StageInfo) -> None: # Compute and store LSE in the same row-sum domain used by P. lse_row_sum = row_sum if cutlass.const_expr(cfg.is_fp8_qkv()): - lse_row_sum = lse_row_sum * fp8_quant_scale_rcp() + lse_row_sum = lse_row_sum * fp8_p_scale_rcp(cfg) lse = ( cute.math.log2(lse_row_sum, fastmath=True) - + self.softmax_scale_log2 * row_max + + self.scale_softmax_log2 * row_max if row_has_values else Float32(-Float32.inf) ) @@ -4130,7 +4537,7 @@ def epilogue_store(self, stage_info: StageInfo) -> None: + Int64(logical_lse_q_idx) * Int64(logical_num_heads_q) + Int64(batch_idx) * Int64(logical_num_heads_q) * Int64(S_q) ) - lse_base_ptr.store(lse) + lse_base_ptr.store(final_lse(lse, cfg.lse_base_on_e)) prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) cute.arch.fence_view_async_tmem_load() @@ -4165,7 +4572,9 @@ def epilogue_store_slice( num_tmem_loads = 4 qk_acc_regs = cutlass.Array(Float32, 128, space=cutlass.AddressSpace.rmem) row_has_values = row_sum > Float32(0) - norm_scale = self._normalization_scale(row_sum, row_max, local_tidx, blk_coord) + norm_scale = self._normalization_scale( + row_sum, row_max, local_tidx, blk_coord, stage_info + ) tile_h = cfg.mma_pv_tiler[0] // cfg.num_mma_ctas tile_d = cfg.mma_pv_tiler[1] @@ -4350,10 +4759,10 @@ def epilogue_store_slice( if cutlass.const_expr(iter_n == 0): lse_row_sum = row_sum if cutlass.const_expr(cfg.is_fp8_qkv()): - lse_row_sum = lse_row_sum * fp8_quant_scale_rcp() + lse_row_sum = lse_row_sum * fp8_p_scale_rcp(cfg) lse = ( cute.math.log2(lse_row_sum, fastmath=True) - + self.softmax_scale_log2 * row_max + + self.scale_softmax_log2 * row_max if row_has_values else Float32(-Float32.inf) ) @@ -4408,7 +4817,7 @@ def epilogue_store_slice( * Int64(logical_num_heads_q) * Int64(S_q) ) - lse_base_ptr.store(lse) + lse_base_ptr.store(final_lse(lse, cfg.lse_base_on_e)) prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) cute.arch.fence_view_async_tmem_load() diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/tasks.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/tasks.py index 0a98fca8ea6a..d3b8ea9cbd80 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/tasks.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/tasks.py @@ -508,10 +508,58 @@ def offsets_schedule(page_offsets, work_queue=None): ) +def create_load_combined_task( + smem_q, + smem_k, + smem_v, + page_offsets, + work_queue=None, + task_class=MlaTask, + **task_kwargs, +): + """Acquire K/V together and reuse each CTA's local K values for V.""" + loop_start, loop_end, loop_step = captured_loop_bounds(task_kwargs, 0) + + @schedule + def combined_schedule(smem_q, smem_k, smem_v, page_offsets, work_queue=None): + """Acquire paired K/V stages, convert once, and publish V after peer DSMEM completion.""" + smem_q.init_load_state() + smem_k.init_load_state() + smem_v.init_load_state() + smem_k.bind_sparse_request() + smem_q.acquire() + smem_q.tma_load() + smem_q.commit() + with domain_loop(loop_start, loop_end, loop_step): + page_offsets.wait() + page_stage = page_offsets.sparse_stage() + smem_k.acquire() + smem_v.acquire() + smem_v.combined_stage() + smem_k.tma_load_shared(page_offset_stage=page_stage) + smem_k.commit() + smem_v.tma_load_shared(page_offset_stage=page_stage) + smem_v.commit() + page_offsets.release() + work_queue_tail(work_queue, advance_label="advance_tile") + + captured = combined_schedule(smem_q, smem_k, smem_v, page_offsets, work_queue) + return task_class( + src_resources=[page_offsets] + ([work_queue] if work_queue is not None else []), + dst_resources=[smem_q, smem_k, smem_v], + warp_idx=smem_k.cfg.load_k_warp_id, + num_warps=8, + schedule=captured, + name="LoadCombinedTask", + **task_kwargs, + ) + + def create_load_k_task( smem_q: SmemQResource, smem_k: SmemKResource, page_offsets=None, + page_window=None, work_queue: MlaWorkQueue = None, task_class: type = MlaTask, **task_kwargs, @@ -520,14 +568,21 @@ def create_load_k_task( Q is loaded once before the loop. K uses one whole-tile pipeline stage per logical K tile. Sparse indices are prefetched through the shared offset - ring; ordinary paged inputs read page offsets directly from GMEM. + ring (``page_offsets``); ordinary paged inputs take page IDs from the + prefetched register window of ``page_window``. Pass exactly one of them. """ + if (page_offsets is None) == (page_window is None): + raise ValueError("pass exactly one of page_offsets or page_window") loop_start, loop_end, loop_step = captured_loop_bounds(task_kwargs, 0) - def load_k_body(smem_q, smem_k, page_offsets=None, work_queue=None): + def load_k_body( + smem_q, smem_k, page_offsets=None, page_window=None, work_queue=None + ): """Load Q once and then publish one K stage per loop tile.""" smem_q.init_load_state() smem_k.init_load_state() + if page_offsets is None: + k_pages, window_page, next_window_page = page_window.init_dense_k_window() smem_q.acquire() smem_q.tma_load() smem_q.commit() @@ -543,24 +598,33 @@ def load_k_body(smem_q, smem_k, page_offsets=None, work_queue=None): # may the metadata producer overwrite this shared ring slot. page_offsets.release() else: + # Read the page IDs before waiting for the free slot so the + # window refresh overlaps the wait instead of following it. + k_pages, window_page, next_window_page = ( + page_window.read_dense_k_window( + cached_k_pages=k_pages, + cached_window_page=window_page, + cached_next_window_page=next_window_page, + ) + ) smem_k.acquire() - smem_k.tma_load_paged() + smem_k.tma_load_paged(cached_page_ids=k_pages) smem_k.commit() work_queue_tail(work_queue, advance_label="advance_tile") @schedule def load_k_schedule(smem_q, smem_k, page_offsets, work_queue=None): - load_k_body(smem_q, smem_k, page_offsets, work_queue) + load_k_body(smem_q, smem_k, page_offsets, None, work_queue) @schedule - def load_k_dense_schedule(smem_q, smem_k, work_queue=None): - load_k_body(smem_q, smem_k, None, work_queue) + def load_k_dense_schedule(smem_q, smem_k, page_window, work_queue=None): + load_k_body(smem_q, smem_k, None, page_window, work_queue) if page_offsets is None: schedule_result = ( - load_k_dense_schedule(smem_q, smem_k) + load_k_dense_schedule(smem_q, smem_k, page_window) if work_queue is None - else load_k_dense_schedule(smem_q, smem_k, work_queue) + else load_k_dense_schedule(smem_q, smem_k, page_window, work_queue) ) else: schedule_result = ( @@ -569,7 +633,7 @@ def load_k_dense_schedule(smem_q, smem_k, work_queue=None): else load_k_schedule(smem_q, smem_k, page_offsets, work_queue) ) - src = ([page_offsets] if page_offsets is not None else []) + ( + src = [page_offsets if page_offsets is not None else page_window] + ( [work_queue] if work_queue is not None else [] ) return task_class( @@ -586,16 +650,25 @@ def load_k_dense_schedule(smem_q, smem_k, work_queue=None): def create_load_v_task( smem_v: SmemVResource, page_offsets=None, + page_window=None, work_queue: MlaWorkQueue = None, task_class: type = MlaTask, **task_kwargs, ) -> Task: - """Create the FP8 V TMA task on its configured loader warps.""" + """Create the FP8 V TMA task on its configured loader warps. + + Sparse indices come from the shared offset ring (``page_offsets``); ordinary + paged inputs take page IDs from ``page_window``. Pass exactly one of them. + """ + if (page_offsets is None) == (page_window is None): + raise ValueError("pass exactly one of page_offsets or page_window") loop_start, loop_end, loop_step = captured_loop_bounds(task_kwargs, 0) - def load_v_body(smem_v, page_offsets=None, work_queue=None): + def load_v_body(smem_v, page_offsets=None, page_window=None, work_queue=None): """Publish one V stage per logical K tile.""" smem_v.init_load_state() + if page_offsets is None: + v_pages, window_page, next_window_page = page_window.init_dense_v_window() with domain_loop(loop_start, loop_end, loop_step): if page_offsets is not None: page_offsets.wait() @@ -606,24 +679,31 @@ def load_v_body(smem_v, page_offsets=None, work_queue=None): # Keep the ring slot protected through every shared read. page_offsets.release() else: + v_pages, window_page, next_window_page = ( + page_window.read_dense_v_window( + cached_v_pages=v_pages, + cached_window_page=window_page, + cached_next_window_page=next_window_page, + ) + ) smem_v.acquire() - smem_v.tma_load_paged() + smem_v.tma_load_paged(cached_page_ids=v_pages) smem_v.commit() work_queue_tail(work_queue, advance_label="advance_tile") @schedule def load_v_schedule(smem_v, page_offsets, work_queue=None): - load_v_body(smem_v, page_offsets, work_queue) + load_v_body(smem_v, page_offsets, None, work_queue) @schedule - def load_v_dense_schedule(smem_v, work_queue=None): - load_v_body(smem_v, None, work_queue) + def load_v_dense_schedule(smem_v, page_window, work_queue=None): + load_v_body(smem_v, None, page_window, work_queue) if page_offsets is None: schedule_result = ( - load_v_dense_schedule(smem_v) + load_v_dense_schedule(smem_v, page_window) if work_queue is None - else load_v_dense_schedule(smem_v, work_queue) + else load_v_dense_schedule(smem_v, page_window, work_queue) ) else: schedule_result = ( @@ -632,7 +712,7 @@ def load_v_dense_schedule(smem_v, work_queue=None): else load_v_schedule(smem_v, page_offsets, work_queue) ) - src = ([page_offsets] if page_offsets is not None else []) + ( + src = [page_offsets if page_offsets is not None else page_window] + ( [work_queue] if work_queue is not None else [] ) return task_class( diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/work_partition.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/work_partition.py index a9448f32c431..6aa5da0a3e23 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/work_partition.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_2cta/work_partition.py @@ -108,14 +108,15 @@ def runtime_active_split_count(k_tile_total, split_kv_cap): @cute.jit -def runtime_split_tile_range(k_tile_total, split_kv_cap, split_idx): +def runtime_split_tile_range(k_tile_total, split_kv_cap, split_idx, split_divisor=None): """Device form of :func:`split_tile_range`.""" k_tile_total = cute.math.max(Int32(k_tile_total), Int32(0)) split_idx = Int32(split_idx) split_kv_cap = cute.math.max(Int32(split_kv_cap), Int32(1)) + divisor = split_kv_cap if split_divisor is None else split_divisor tiles_per_split = cute.math.max( - (k_tile_total + split_kv_cap - Int32(1)) // split_kv_cap, + (k_tile_total + split_kv_cap - Int32(1)) // divisor, Int32(1), ) start = split_idx * tiles_per_split @@ -132,6 +133,7 @@ def runtime_row_prefix_active_split_count( row_k_tile_total, group_k_tile_total, split_kv_cap, + split_divisor=None, ): """Device form of :func:`row_prefix_active_split_count`.""" @@ -141,8 +143,9 @@ def runtime_row_prefix_active_split_count( Int32(0), ) split_kv_cap = cute.math.max(Int32(split_kv_cap), Int32(1)) + divisor = split_kv_cap if split_divisor is None else split_divisor tiles_per_split = cute.math.max( - (group_k_tile_total + split_kv_cap - Int32(1)) // split_kv_cap, + (group_k_tile_total + split_kv_cap - Int32(1)) // divisor, Int32(1), ) row_active_splits = ( diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/config.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/config.py index 7657de52608d..ed5333e34877 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/config.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/config.py @@ -31,7 +31,11 @@ from typing import Any from ....split_kv_mode_policy import select_split_kv_modes -from ..helpers.constants import MAX_MLA_SPLITS_KV, SUPPORTED_MLA_PAGE_SIZES +from ..helpers.constants import SUPPORTED_MLA_PAGE_SIZES +from ..helpers.nvfp4_kv import ( + Nvfp4KVConfig, + nvfp4_row_bytes, +) from ..helpers.mask import MaskType, normalize_mask_type from ..helpers.query import FlatQueryTileLayout @@ -78,7 +82,7 @@ def estimated_throughput_latency_smem_data_bytes(cfg: "MlaConfig") -> int: @dataclass(frozen=True) -class MlaConfig: +class MlaConfig(Nvfp4KVConfig): """Kernel traits for one throughput-latency 1CTA MLA schedule. The dataclass groups shape dimensions, tile sizes, CTA decomposition, @@ -92,12 +96,20 @@ class MlaConfig: batch_size: int = 1 num_heads_q: int = 16 seq_len_q: int = 4 - seq_len_kv: int = 4096 + seq_len_kv: int | None = 4096 logical_num_heads_q: int = 16 logical_seq_len_q: int = 4 + # True: the Q lengths are compile-time. Sparse flattened Q is always fixed. + fixed_seqlen: bool = False # Causal is bottom-right aligned for speculative decode. Dense keeps only # the ordinary per-batch KV-tail predicate. mask_type: str = MaskType.CAUSAL.value + # Threshold skip correction in log2 units (binades; 0 disables it): freeze + # the running row max while a K tile raises it by at most this + # compile-time amount and skip that tile's O rescale. E4M3 P then uses + # the scale ``448 * 2**-T`` instead of 448. + softmax_rescale_threshold_log2: float = 0.0 + lse_base_on_e: bool = False head_dim_qk: int = 576 head_dim_v: int = 512 latent_dim: int = 512 @@ -165,6 +177,18 @@ class MlaConfig: correction_num_warps: int = 4 mma_num_warps: int = 1 load_num_warps: int = 1 + whole_kv_tile: bool = False + single_raw_prefetch: bool = False + compact_nv_rope: bool = False + q_tma_stage: bool = False + q_nope_tmem: bool = False + # Softmax/correction already use the native WS Layout-E row ownership. + ws_native_e: bool = False + ws_qk: bool = False + # Use M64 tcgen05.mma.ws for PV; P is staged in SMEM, O follows Layout E. + ws_pv: bool = False + kv_cache_format: str = "native" + extra_kv_cache_format: str = "native" sparse_reuse_kv: bool = False paired_sparse_correction: bool = False defer_sparse_max_update: bool = False @@ -173,9 +197,11 @@ class MlaConfig: fuse_sparse_cluster_epilogue: bool = False fuse_sparse_reduction: bool = False sparse_direct: bool = False - sparse_static_scales: bool = False sparse_valid_prefix: bool = False sparse_direct_capacities: tuple[int, int] = (0, 0) + # Planned KV tiles per instruction group (1 Keeps, 2 Swaps); zero selects + # a runtime loop domain, while None derives the group from static lengths. + sparse_instruction_group: int | None = None page_offsets_num_warps: int = 1 scheduler_num_warps: int = 1 clc_padding_warp_idx: int = 14 @@ -204,11 +230,64 @@ class MlaConfig: use_sliding_window_causal: int = 0 attention_window_size: int = 0 + @property + def sparse_d512(self) -> bool: + """Two-source D512 sparse MLA: one token row per index (page size + one), RoPE folded into the 512 columns.""" + return self.num_tokens_per_page == 1 and self.rope_dim == 0 + + @property + def sparse_d576_reload(self) -> bool: + """Use the native FP8 q64 loader that releases K before reloading V. + + Retained native tiles and packed caches share D576 geometry but own + different descriptors, coordinate caches and output finalization. + The existing format and pipeline traits identify this reload schedule. + """ + return ( + self.num_tokens_per_page == 1 + and self.rope_dim == 64 + and self.qkv_dtype == "e4m3" + and self.cache_uniform_sparse_quads + and not self.sparse_reuse_kv + and not self.has_nvfp4_kv + ) + + @property + def kv_load_regs(self) -> int: + """Give raw-record loaders 96 registers without raising the MMA warp budget. + + Loader warps retain packed values/scales across raw-buffer release. + MMA/page warps keep mma_load_regs; raising both exceeds the CTA budget. + This is derived from the existing pipeline choice, not a tuning field. + """ + return 96 if self.async_nvfp4_kv_gather else self.mma_load_regs + + @property + def nvfp4_row_bytes(self) -> int: + """Return packed bytes per token, including scales and the optional FP8 RoPE tail.""" + return nvfp4_row_bytes(self.rope_dim) + + @property + def whole_kv_tile_bytes(self) -> int: + """Return SMEM bytes retained for one complete converted K/V tile.""" + return self.tile_size_kv * self.head_dim_qk * self.qkv_dtype_bytes + + @property + def async_nvfp4_kv_gather(self) -> bool: + """Use independent gather/transform warps with two retained tiles.""" + return self.has_nvfp4_kv and self.whole_kv_tile + @property def one_insts_kv_swap(self) -> bool: """One KV stream with swapped MMA operands (P remains in SMEM).""" return self.num_insts_kv == 1 and self.kernel_variant == "swaps_mma_ab" + @property + def compact_softmax_stats(self) -> bool: + """Store one warp-reduced statistic per Swap-AB column group.""" + return self.whole_kv_tile and self.one_insts_kv_swap + @property def single_kv_pipe(self) -> bool: """Use one softmax/P stream while retaining two score stages.""" @@ -279,16 +358,22 @@ def pages_per_kv_tile(self) -> int: return ceil(self.tile_size_kv / self.num_tokens_per_page) + # Return the static KV tile count, or zero when live metadata owns the domain. @property def total_kv_tiles(self) -> int: + if self.seq_len_kv is None: + return 0 return ceil(self.seq_len_kv / self.tile_size_kv) @property def kv_tiles_per_multi_cta_group(self) -> int: return self.num_ctas_per_seq_kv * self.num_insts_kv + # Return static loop steps per split CTA, or zero for a runtime KV domain. @property def num_steps_per_cta_kv(self) -> int: + if self.seq_len_kv is None: + return 0 tokens_per_multi_cta_step = ( self.num_ctas_per_seq_kv * self.num_insts_kv * self.tile_size_kv ) @@ -429,6 +514,8 @@ def fixed_sparse_kv_tiles(self) -> int: If the maximum planned split fits that group, every shorter valid prefix already has the same partition. Live lengths still mask data. """ + if self.sparse_instruction_group is not None: + return self.sparse_instruction_group tiles = self.local_kv_tiles(self.total_kv_tiles) return ( self.num_insts_kv @@ -627,10 +714,15 @@ def q_tile_work_count( def automatic_split_kv_step_tokens(tile_size_q: int) -> int: - """Return one steady-state K step for the selected task graph.""" + """Return a profitable K span per split, at least two physical KV tiles. + + One-tile splits do not amortize CTA setup and final reduction. The + one-instance Q64+ schedules retain two tiles as a policy work bound, + while the kernel and its hardware ceiling still accept larger counts. + """ num_kv_insts = 1 if tile_size_q >= 64 else MlaConfig.num_insts_kv - return MlaConfig.tile_size_kv * num_kv_insts + return MlaConfig.tile_size_kv * max(num_kv_insts, 2) def select_auto_split_kv( @@ -647,7 +739,6 @@ def select_auto_split_kv( split_kv = min( max_split_kv, max(1, target_work // base_work), - MAX_MLA_SPLITS_KV, ) local_k_steps = ceil(max_split_kv / split_kv) return ceil(max_split_kv / local_k_steps) @@ -734,7 +825,12 @@ def correction_register_budget(tile_size_q: int) -> int: return MlaConfig.correction_regs -def keeps_mma_ab_config_kwargs(profile: MlaProfile, qkv_dtype: str) -> dict[str, int]: +def keeps_mma_ab_config_kwargs( + profile: MlaProfile, + qkv_dtype: str, + *, + sparse_d576: bool = False, +) -> dict[str, int]: """Return keeps-MMA-AB scheduler, pipeline, and register traits.""" kv_stages = MlaConfig.kv_stages @@ -749,7 +845,7 @@ def keeps_mma_ab_config_kwargs(profile: MlaProfile, qkv_dtype: str) -> dict[str, use_persistent_scheduler = 0 use_clc_dynamic_persistent_scheduler = 0 - return { + kwargs = { # Keeps-MMA-AB has one QK/PV pipe. Raising this value would change the # KV work partition without adding the second stream used by the swaps # schedule, and would therefore skip every other KV tile. @@ -771,6 +867,32 @@ def keeps_mma_ab_config_kwargs(profile: MlaProfile, qkv_dtype: str) -> dict[str, "use_persistent_scheduler": use_persistent_scheduler, "use_clc_dynamic_persistent_scheduler": use_clc_dynamic_persistent_scheduler, } + if sparse_d576: + # D576 sparse: one warp cannot issue 32 Gather4 per 16 KB stage fast + # enough (~0.6 us per stage measured), so the K/V loader spans warps + # 12-15 (8 gathers each, issued from warp-uniform cached rows) and + # warp 9 pads. Register budget for 16 warps: 4x168 + 4x160 + 8x80 = + # 62464 <= 65536. QK(t + 1) is issued before PV(t) and V is gathered + # again from L2 (trtllm-gen's order): softmax(t) overlaps QK(t + 1). + # Retained K/V schedules can also issue QK(t + 1) before PV(t), but + # retain the current tile's stages through its V use. Releasing K + # early here changes stage availability for the next prefetch. + # One CTA runs one work tile, so a single Q stage suffices and its second 36 KB pays + # for two more KV stages. + kwargs.update( + { + "q_stages": 1, + "kv_stages": 10, + "cache_uniform_sparse_quads": True, + "load_warp_idx": 12, + "load_num_warps": 4, + "threads_per_cta": 512, + "softmax_regs": 168, + "correction_regs": 160, + "mma_load_regs": 80, + } + ) + return kwargs def resolve_split_kv_reduction_policy( @@ -862,6 +984,8 @@ def resolve_runtime_cluster_reduction_mode( cfg: MlaConfig | None, *, reduction_mode: str | None, + hardware_info=None, + stream=None, log=None, ) -> str | None: """Use cached device occupancy to resolve the runtime reduction mode.""" @@ -877,7 +1001,11 @@ def resolve_runtime_cluster_reduction_mode( cluster_size = cluster_shape[0] * cluster_shape[1] * cluster_shape[2] from ...._hardware import get_max_active_clusters - active_clusters = get_max_active_clusters(cluster_size) + active_clusters = ( + get_max_active_clusters(cluster_size) + if hardware_info is None + else hardware_info.get_max_active_clusters(cluster_size, stream) + ) cluster_count = cluster_reduction_cluster_count(cfg) if log is not None: log( @@ -989,7 +1117,10 @@ def persistent_override_profile( return profile if profile.use_multi_ctas_kv == 1: if explicit_persistent: - raise ValueError("persistent scheduling is not supported with split-KV") + raise ValueError( + "1CTA persistent scheduling is not supported with split-KV; " + "use split_kv=1 or nonpersistent scheduling" + ) return profile if explicit_persistent: return replace( @@ -1305,7 +1436,7 @@ def is_throughput_latency_mla_supported_shape( and is_power_of_two(num_heads_q) and ( num_tokens_per_page in SUPPORTED_MLA_PAGE_SIZES - or (num_tokens_per_page == 1 and rope_dim == 0) + or (num_tokens_per_page == 1 and rope_dim in (0, 64)) ) and seq_len_q >= 1 and seq_len_kv >= 128 @@ -1349,7 +1480,10 @@ def enumerate_throughput_latency_mla_profiles( ) if explicit_split_kv is not None and explicit_split_kv > 0: if explicit_split_kv > 1 and explicit_persistent is True: - raise ValueError("persistent scheduling is not supported with split-KV") + raise ValueError( + "1CTA persistent scheduling is not supported with split-KV; " + "use split_kv=1 or nonpersistent scheduling" + ) if explicit_split_kv == 1 and explicit_persistent is not None: if selected_tile_size_q >= 64: profile = keeps_mma_ab_forced_persistent_profile( @@ -1518,8 +1652,15 @@ def make_throughput_latency_mla_config( explicit_split_kv: int | None = None, explicit_persistent: bool | None = None, mask_type: MaskType | str = MaskType.CAUSAL, + softmax_rescale_threshold_log2: float = 0.0, + sparse_profile=None, ) -> MlaConfig: - """Return throughput-latency 1CTA MLA traits for a concrete profile.""" + """Return throughput-latency 1CTA MLA traits for a concrete profile. + + ``sparse_profile`` is the already-selected general sparse pipeline. Its + presence excludes native D576 reload defaults; it adds no independent + tuning or compile-time choice. + """ mask_type = normalize_mask_type(mask_type) @@ -1549,8 +1690,14 @@ def make_throughput_latency_mla_config( raise ValueError("latent_dim must be positive") if rope_dim < 0: raise ValueError("rope_dim must be non-negative") + sparse_d576 = ( + num_tokens_per_page == 1 + and rope_dim == 64 + and qkv_dtype == "e4m3" + and sparse_profile is None + ) if num_tokens_per_page not in SUPPORTED_MLA_PAGE_SIZES and not ( - num_tokens_per_page == 1 and rope_dim == 0 + num_tokens_per_page == 1 and rope_dim in (0, 64) ): raise ValueError( "num_tokens_per_page must be one of " @@ -1566,6 +1713,9 @@ def make_throughput_latency_mla_config( metadata_capacity = ( 128 if num_tokens_per_page == 1 else MlaConfig.page_offsets_entries_per_stage ) + page_offsets_stages = MlaConfig.page_offsets_stages + if sparse_d576: + page_offsets_stages = 4 if pages_per_kv_tile > metadata_capacity: raise ValueError( "page-offset staging capacity is smaller than one KV tile: " @@ -1592,7 +1742,9 @@ def make_throughput_latency_mla_config( seq_len_kv=seq_len_kv, latent_dim=latent_dim, rope_dim=rope_dim, - num_tokens_per_page=num_tokens_per_page, + # Profile enumeration only checks that the page size is a dense one; + # the D576 sparse page size of one is validated above. + num_tokens_per_page=(16 if sparse_d576 else num_tokens_per_page), max_active_clusters=max_active_clusters, qkv_dtype=qkv_dtype, tile_size_q=tile_size_q, @@ -1607,7 +1759,10 @@ def make_throughput_latency_mla_config( # Persistent work queues enumerate Q/head-dimension tiles, not KV # splits. Direct MlaProfile callers must obey the same restriction # as persistent_override_profile, or only split zero gets produced. - raise ValueError("persistent scheduling is not supported with split-KV") + raise ValueError( + "1CTA persistent scheduling is not supported with split-KV; " + "use split_kv=1 or nonpersistent scheduling" + ) tile_size_q = tile_size_q_for_profile( selected_profile, num_heads_q, seq_len_q, tile_size_q ) @@ -1674,6 +1829,8 @@ def make_throughput_latency_mla_config( logical_num_heads_q=logical_num_heads_q, logical_seq_len_q=logical_seq_len_q, mask_type=mask_type, + # Range-checked by ThroughputLatencyMlaDecodeTs and plan(). + softmax_rescale_threshold_log2=softmax_rescale_threshold_log2, head_dim_qk=head_dim_qk, head_dim_v=latent_dim, latent_dim=latent_dim, @@ -1693,6 +1850,7 @@ def make_throughput_latency_mla_config( correction_regs=correction_register_budget(tile_size_q), num_tokens_per_page=num_tokens_per_page, page_offsets_entries_per_stage=metadata_capacity, + page_offsets_stages=page_offsets_stages, max_num_pages_per_seq_kv=max(1, ceil(seq_len_kv / num_tokens_per_page)), tmem_s_cols=tile_size_q, tmem_stats_cols=MlaConfig.tmem_stats_cols, @@ -1708,7 +1866,13 @@ def make_throughput_latency_mla_config( ) config_kwargs.update(dtype_config_kwargs(qkv_dtype, o_dtype)) if selected_profile.kernel_variant == "keeps_mma_ab": - config_kwargs.update(keeps_mma_ab_config_kwargs(selected_profile, qkv_dtype)) + config_kwargs.update( + keeps_mma_ab_config_kwargs( + selected_profile, + qkv_dtype, + sparse_d576=sparse_d576, + ) + ) cfg = MlaConfig(**config_kwargs) validate_explicit_cluster_reduction_config(cfg, reduction_mode=reduction_mode) cfg = resolve_auto_cluster_reduction_config( @@ -1725,7 +1889,6 @@ def configure_sparse_mla( finalize_output=False, direct_inputs=False, assume_valid_prefix=False, - static_scales=False, source_capacities=(0, 0), ) -> MlaConfig: """Resolve one selected sparse profile without order-dependent overrides. @@ -1738,7 +1901,20 @@ def configure_sparse_mla( reuse = profile.reuse_kv uniform = profile.uniform_offset_cache single = profile.single_kv_stream - if cfg.num_tokens_per_page != 1 or cfg.rope_dim != 0: + whole_nvfp4_kv = profile.whole_kv_tile + if whole_nvfp4_kv and not ( + cfg.has_nvfp4_kv + and cfg.is_fp8_qkv() + and cfg.rope_dim in (0, 64) + and reuse + and (keep or single) + and issuers == 8 + and profile.kv_tile_size == 128 + ): + raise ValueError( + "whole scaled KV tiles require D512/D576, one KV stream, BK128 and eight converters" + ) + if cfg.num_tokens_per_page != 1 or cfg.rope_dim not in (0, 64): raise ValueError("sparse profiles require native page-size-one MLA") if issuers not in (1, 2, 4, 8): raise ValueError("sparse profiles require one, two, four or eight issuers") @@ -1748,14 +1924,14 @@ def configure_sparse_mla( raise ValueError("unsupported sparse query tile") if single and ( keep - or cfg.tile_size_q not in (8, 16) + or cfg.tile_size_q not in ((8, 16, 32) if whole_nvfp4_kv else (8, 16)) or cfg.num_ctas_per_head_dim != 1 - or issuers != 4 + or issuers != (8 if whole_nvfp4_kv else 4) or cfg.use_cluster_reduction or (cfg.num_ctas_per_seq_kv > 1 and cfg.use_persistent_scheduler) ): raise ValueError( - "one-instance Swap-AB requires M8/M16, V1, four issuers and no cluster reduction" + "one-instance Swap-AB requires a supported tile/loader layout, one V owner and no cluster reduction" ) if reuse and cfg.num_ctas_per_head_dim != 1: raise ValueError("retained KV requires one V owner") @@ -1792,6 +1968,13 @@ def configure_sparse_mla( raise ValueError("fused output requires one logical query per virtual request") if direct_inputs and (not finalize_output or cfg.qkv_dtype not in ("bf16", "e4m3")): raise ValueError("separate sparse indices require native BF16/FP8 final output") + if profile.defer_max_update and cfg.softmax_rescale_threshold_log2 > 0.0: + # The sparse fixed-bound deferral takes precedence in update_softmax; + # a threshold would then only change the FP8 P scale. + raise ValueError( + "softmax_rescale_threshold_log2 is not supported with deferred " + "sparse max updates" + ) cfg = dataclass_replace( cfg, num_insts_kv=1 if keep or single else 2, @@ -1802,7 +1985,6 @@ def configure_sparse_mla( paired_sparse_correction=profile.paired_correction, defer_sparse_max_update=profile.defer_max_update, sparse_direct=direct_inputs, - sparse_static_scales=static_scales, sparse_valid_prefix=assume_valid_prefix, sparse_direct_capacities=source_capacities if direct_inputs else (0, 0), fuse_sparse_epilogue=direct_finish, @@ -1815,23 +1997,32 @@ def configure_sparse_mla( issuers == 8 and not keep and ( - cfg.tile_size_q not in (8, 16) + (cfg.tile_size_q not in (8, 16) and not whole_nvfp4_kv) or cfg.local_kv_tiles(cfg.total_kv_tiles) <= cfg.num_insts_kv ) ): raise ValueError("eight Swap-AB issuers require a steady-state loop") + # Retain complete converted FP8 K tiles through every PV feature pass, + # avoiding a second NVFP4 gather/conversion for V. Eight loaders and an + # independent raw-prefetch buffer provide room for two complete payloads. + retained_nvfp4_kv = ( + cfg.has_nvfp4_kv and keep and issuers == 8 and reuse + ) or whole_nvfp4_kv + ws_mma = retained_nvfp4_kv and keep # Four D128 slices make one D512 tile. FP8/BK64 reuse keeps two tiles; # BF16 BK128 retains one, with additional stages used for prefetch. kv_stages = ( profile.reuse_kv_stages if reuse and profile.reuse_kv_stages else cfg.kv_stages ) - if reuse and profile.reuse_kv_stages: - retained_tiles = 2 if cfg.is_fp8_qkv() or cfg.tile_size_kv == 64 else 1 + if reuse and profile.reuse_kv_stages and not retained_nvfp4_kv: + retained_tiles = ( + 2 + if cfg.is_fp8_qkv() or (cfg.tile_size_kv == 64 and not cfg.rope_dim) + else 1 + ) if kv_stages < retained_tiles * cfg.qk_head_dim_stages: raise ValueError("insufficient stages to retain KV through PV") - if cfg.tile_size_kv == 64 and profile.reuse_kv_stages < 2 * cfg.qk_head_dim_stages: - raise ValueError("BK64 requires two retained K tiles") page_stages = cfg.page_offsets_stages if ( reuse @@ -1842,6 +2033,8 @@ def configure_sparse_mla( page_stages = 1 if profile.page_pipeline_stages: page_stages = profile.page_pipeline_stages + if retained_nvfp4_kv: + page_stages = 2 if single and reuse and kv_stages >= 6 and not cfg.is_fp8_qkv(): page_stages = 1 @@ -1904,10 +2097,31 @@ def configure_sparse_mla( if balanced else cfg.scheduler_regs, ) + if retained_nvfp4_kv: + roles.update( + threads_per_cta=640, + page_offsets_warp_idx=9, + page_offsets_num_warps=2, + softmax_regs=80 if ws_mma else 96, + correction_regs=144 if ws_mma else 128, + ws_pv=ws_mma, + ws_qk=ws_mma, + ws_native_e=ws_mma, + q_nope_tmem=ws_mma, + q_tma_stage=ws_mma, + tmem_s_cols=64 if ws_mma else cfg.tile_size_q, + compact_nv_rope=bool(cfg.rope_dim), + single_raw_prefetch=True, + page_offsets_entries_per_stage=132, + whole_kv_tile=True, + ) + kv_stages = 2 return dataclass_replace( cfg, **roles, kv_stages=kv_stages, page_offsets_stages=page_stages, - q_stages=1 if single or (reuse and profile.reuse_kv_stages) else cfg.q_stages, + q_stages=1 + if retained_nvfp4_kv or single or (reuse and profile.reuse_kv_stages) + else cfg.q_stages, ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/kernel.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/kernel.py index 130297f3b4a9..1ebd025c3e98 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/kernel.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/kernel.py @@ -50,7 +50,7 @@ WorkQueue, WorkTileInfo, ) -from cutlass.experimental.task_scheduling.task_manager import TaskManager +from ..helpers.diagnostics import MlaTaskManager as TaskManager from .config import ( MlaConfig, @@ -95,9 +95,18 @@ supports_parallel_gmem_reduction, ) from .reduction import gmem_reduction_launch_shape, run_gmem_reduction_kernel +from ..helpers.gather import make_sparse_rope_map +from ..helpers.nvfp4_kv import ( + validate_nvfp4_kv_formats, + Nvfp4KVConfig, + make_nvfp4_kv_maps, +) from ..helpers.constants import LOG2_E, TMEM_LIFECYCLE_BARRIER_ID, TMEM_READY_BARRIER_ID +from ..helpers.math import canonicalize_softmax_rescale_threshold_log2 from ..helpers.mask import MaskType, normalize_mask_type from ..helpers.tile import ( + runtime_request_publishes, + runtime_split_count, runtime_query_tile_is_active, runtime_split_pruning_is_profitable, runtime_split_tile_is_active, @@ -244,7 +253,7 @@ def __init__( self.batch_size = batch_size self.cache_seqs = cache_seqs self.cu_seqlens_q = cu_seqlens_q - self.enable_runtime_skip = cu_seqlens_q is not None + self.enable_runtime_skip = cache_seqs is not None or cu_seqlens_q is not None @cute.jit def skip_work_tile_if(self, work_tile: WorkTileInfo): @@ -327,7 +336,7 @@ def __init__( self.cfg = cfg self.cache_seqs = cache_seqs self.cu_seqlens_q = cu_seqlens_q - self.enable_runtime_skip = cu_seqlens_q is not None + self.enable_runtime_skip = cache_seqs is not None or cu_seqlens_q is not None self._clc_response_alloc = SmemAllocation( self.name + "_clc_response", dtype=cutlass.Int128, count=2, alignment=16 ) @@ -462,7 +471,12 @@ def build_throughput_latency_mla_task_manager( tma_desc_c_latent=None, tma_desc_c_rope=None, tma_desc_v=None, + tma_desc_bulk_extra=None, + tma_desc_sparse_rope=None, + tma_desc_sparse_extra_rope=None, c_rope_tensor=None, + c_latent_tensor=None, + q_tensor=None, page_offsets=None, cache_seqs=None, cu_seqlens_q=None, @@ -486,6 +500,7 @@ def build_throughput_latency_mla_task_manager( use_static_persistent_scheduler: bool = False, verbose: bool = False, exhaustive_deadlock_race_check: bool = False, + requested_splits=None, ) -> tuple[TaskManager, object]: """Build the throughput-latency 1CTA MLA TaskManager for validation or execution. @@ -514,7 +529,12 @@ def build_throughput_latency_mla_task_manager( tma_desc_c_latent=tma_desc_c_latent, tma_desc_c_rope=tma_desc_c_rope, tma_desc_v=tma_desc_v, + tma_desc_bulk_extra=tma_desc_bulk_extra, + tma_desc_sparse_rope=tma_desc_sparse_rope, + tma_desc_sparse_extra_rope=tma_desc_sparse_extra_rope, c_rope_tensor=c_rope_tensor, + c_latent_tensor=c_latent_tensor, + q_tensor=q_tensor, page_offsets=page_offsets, cache_seqs=cache_seqs, cu_seqlens_q=cu_seqlens_q, @@ -538,6 +558,7 @@ def build_throughput_latency_mla_task_manager( use_static_persistent_scheduler=use_static_persistent_scheduler, verbose=verbose, exhaustive_deadlock_race_check=exhaustive_deadlock_race_check, + requested_splits=requested_splits, ) agent = pipeline.Agent @@ -545,7 +566,9 @@ def build_throughput_latency_mla_task_manager( tma_group = pipeline.CooperativeGroup(agent.Thread) umma_group = pipeline.CooperativeGroup(agent.Thread) load_group = pipeline.CooperativeGroup(agent.Thread, cfg.load_num_warps * 32) - page_group = pipeline.CooperativeGroup(agent.Thread, 32) + page_group = pipeline.CooperativeGroup( + agent.Thread, 32 * cfg.page_offsets_num_warps + ) scheduler_group = pipeline.CooperativeGroup( agent.Thread, cfg.scheduler_num_warps * 32, @@ -599,9 +622,12 @@ def build_throughput_latency_mla_task_manager( ) smem_kv_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( num_stages=cfg.kv_stages, - num_bytes=cfg.kv_smem_tile_bytes, + num_bytes=cfg.whole_kv_tile_bytes + if cfg.whole_kv_tile + else cfg.kv_smem_tile_bytes, num_bytes_per_warp_per_cta=( - cfg.kv_smem_tile_bytes // cfg.load_num_warps + (cfg.whole_kv_tile_bytes if cfg.whole_kv_tile else cfg.kv_smem_tile_bytes) + // cfg.load_num_warps if cfg.load_num_warps > 1 else None ), @@ -722,8 +748,11 @@ def build_throughput_latency_mla_task_manager( tma_desc_c_latent=tma_desc_c_latent, tma_desc_c_rope=tma_desc_c_rope, tma_desc_v=tma_desc_v, - tma_desc_bulk_extra=tma_desc_q_rope, + tma_desc_bulk_extra=tma_desc_bulk_extra, + tma_desc_sparse_rope=tma_desc_sparse_rope, + tma_desc_sparse_extra_rope=tma_desc_sparse_extra_rope, c_rope_tensor=c_rope_tensor, + c_latent_tensor=c_latent_tensor, page_offsets_kv=smem_page_offsets if use_page_offsets else None, page_offsets=page_offsets, cache_seqs=cache_seqs, @@ -818,10 +847,11 @@ def build_throughput_latency_mla_task_manager( global1 = TmemSoftmaxGlobalResource(cfg=cfg, inst_id=1, name="ll_mla_global1") corr0 = TmemCorrResource( cfg=cfg, + page_offsets=page_offsets, inst_id=0, scale_softmax_log2=scale_softmax_log2, output_scale=output_scale, - atten_sinks=sparse_epilogue_params, + attention_sinks=sparse_epilogue_params, output_quant=output_quant, output_quant_args=output_quant_args, o_tensor=o_tensor, @@ -839,10 +869,11 @@ def build_throughput_latency_mla_task_manager( ) corr1 = TmemCorrResource( cfg=cfg, + page_offsets=page_offsets, inst_id=1, scale_softmax_log2=scale_softmax_log2, output_scale=output_scale, - atten_sinks=sparse_epilogue_params, + attention_sinks=sparse_epilogue_params, output_quant=output_quant, output_quant_args=output_quant_args, o_tensor=o_tensor, @@ -1092,6 +1123,7 @@ def build_throughput_latency_mla_task_manager( tmem_allocator.add_resource(local0) tmem_allocator.add_resource(local1) tmem_allocator.add_resource(tmem_o) + tmem_allocator.add_resource(smem_q) tmem_allocator.compute_layout() if cfg.short_merged_softmax: @@ -1128,6 +1160,11 @@ def build_throughput_latency_mla_task_manager( ) idle.name = "MergedLoadGroupPadding" tasks.append(idle) + for task in tasks: + task.runtime_split_kv = requested_splits + for resource in deps: + if hasattr(resource, "runtime_split_kv"): + resource.runtime_split_kv = requested_splits task_manager = TaskManager( tasks=tasks, resource_dependency_graph=deps, @@ -1153,7 +1190,12 @@ def _make_single_kv_pipe_task_graph( tma_desc_c_latent=None, tma_desc_c_rope=None, tma_desc_v=None, + tma_desc_bulk_extra=None, + tma_desc_sparse_rope=None, + tma_desc_sparse_extra_rope=None, c_rope_tensor=None, + c_latent_tensor=None, + q_tensor=None, page_offsets=None, cache_seqs=None, cu_seqlens_q=None, @@ -1177,6 +1219,7 @@ def _make_single_kv_pipe_task_graph( use_static_persistent_scheduler: bool = False, verbose: bool = False, exhaustive_deadlock_race_check: bool = False, + requested_splits=None, ) -> tuple[TaskManager, object]: """Build the single-stream 1CTA task graph. @@ -1197,7 +1240,9 @@ def _make_single_kv_pipe_task_graph( tma_group = pipeline.CooperativeGroup(agent.Thread) umma_group = pipeline.CooperativeGroup(agent.Thread) load_group = pipeline.CooperativeGroup(agent.Thread, cfg.load_num_warps * 32) - page_group = pipeline.CooperativeGroup(agent.Thread, 32) + page_group = pipeline.CooperativeGroup( + agent.Thread, 32 * cfg.page_offsets_num_warps + ) softmax_group = pipeline.CooperativeGroup( agent.Thread, cfg.softmax_num_warps * 32, @@ -1241,9 +1286,12 @@ def _make_single_kv_pipe_task_graph( ) smem_kv_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( num_stages=cfg.kv_stages, - num_bytes=cfg.kv_smem_tile_bytes, + num_bytes=cfg.whole_kv_tile_bytes + if cfg.whole_kv_tile + else cfg.kv_smem_tile_bytes, num_bytes_per_warp_per_cta=( - cfg.kv_smem_tile_bytes // cfg.load_num_warps + (cfg.whole_kv_tile_bytes if cfg.whole_kv_tile else cfg.kv_smem_tile_bytes) + // cfg.load_num_warps if cfg.load_num_warps > 1 else None ), @@ -1329,6 +1377,7 @@ def _make_single_kv_pipe_task_graph( pipeline_config=smem_q_cfg, tma_desc_q_latent=tma_desc_q_latent, tma_desc_q_rope=tma_desc_q_rope, + q_tensor=q_tensor, cu_seqlens_q=cu_seqlens_q, head_idx=head_idx, batch_idx=batch_idx, @@ -1341,8 +1390,11 @@ def _make_single_kv_pipe_task_graph( tma_desc_c_latent=tma_desc_c_latent, tma_desc_c_rope=tma_desc_c_rope, tma_desc_v=tma_desc_v, - tma_desc_bulk_extra=tma_desc_q_rope, + tma_desc_bulk_extra=tma_desc_bulk_extra, + tma_desc_sparse_rope=tma_desc_sparse_rope, + tma_desc_sparse_extra_rope=tma_desc_sparse_extra_rope, c_rope_tensor=c_rope_tensor, + c_latent_tensor=c_latent_tensor, page_offsets_kv=smem_page_offsets if use_page_offsets else None, page_offsets=page_offsets, cache_seqs=cache_seqs, @@ -1354,6 +1406,10 @@ def _make_single_kv_pipe_task_graph( cta_idx_head_dim_v=cta_idx_head_dim_v, name="ll_mla_q64_smem_kv", ) + if cfg.q_tma_stage: + smem_q.kv_alias_ref = smem_kv + if cfg.async_nvfp4_kv_gather: + smem_page_offsets.nv_target = smem_kv tmem_s = (TmemSResource if cfg.one_insts_kv_swap else TmemSKeepsResource)( cfg=cfg, page_offsets=page_offsets, @@ -1406,7 +1462,7 @@ def _make_single_kv_pipe_task_graph( local_resource = ( SmemSoftmaxLocalResource - if cfg.num_tokens_per_page == 1 + if (cfg.num_tokens_per_page == 1 and not cfg.sparse_d576_reload) else TmemSoftmaxLocalResource ) local = local_resource( @@ -1421,10 +1477,11 @@ def _make_single_kv_pipe_task_graph( ) corr = TmemCorrResource( cfg=cfg, + page_offsets=page_offsets, inst_id=1, scale_softmax_log2=scale_softmax_log2, output_scale=output_scale, - atten_sinks=sparse_epilogue_params, + attention_sinks=sparse_epilogue_params, output_quant=output_quant, output_quant_args=output_quant_args, o_tensor=o_tensor, @@ -1514,7 +1571,25 @@ def _make_single_kv_pipe_task_graph( ), ] ) - if not use_clc_dynamic: + if cfg.sparse_d576_reload: + # The D576 sparse layout moves the loader to W12-15 with page + # offsets on W10; pad W11 unless it loads or stages page offsets, and + # whichever of warps 9-10 stays unused. + used = set(range(cfg.load_warp_idx, cfg.load_warp_idx + cfg.load_num_warps)) + used.add(cfg.page_offsets_warp_idx) + pad_warps = [11] if not use_clc_dynamic and 11 not in used else [] + pad_warps += [w for w in (9, 10) if w not in used] + for pad_warp in pad_warps: + tasks.append( + create_padding_task( + cfg, + work_queue, + warp_idx=pad_warp, + num_warps=1, + task_class=MlaDecodeTask, + ) + ) + elif not use_clc_dynamic: tasks.append( create_padding_task( cfg, @@ -1544,7 +1619,11 @@ def _make_single_kv_pipe_task_graph( ) ) - if cfg.load_warp_idx == 12: + if ( + cfg.load_warp_idx == 12 + and cfg.page_offsets_warp_idx != 9 + and not cfg.sparse_d576_reload + ): idle = create_padding_task( cfg, work_queue, warp_idx=9, num_warps=1, task_class=MlaDecodeTask ) @@ -1633,8 +1712,14 @@ def _make_single_kv_pipe_task_graph( tmem_allocator.add_resource(tmem_s) tmem_allocator.add_resource(local) tmem_allocator.add_resource(tmem_o) + tmem_allocator.add_resource(smem_q) tmem_allocator.compute_layout() + for task in tasks: + task.runtime_split_kv = requested_splits + for resource in deps: + if hasattr(resource, "runtime_split_kv"): + resource.runtime_split_kv = requested_splits task_manager = TaskManager( tasks=tasks, resource_dependency_graph=deps, @@ -1650,9 +1735,13 @@ def _make_single_kv_pipe_task_graph( return task_manager, corr -class ThroughputLatencyMlaDecodeTs: +class ThroughputLatencyMlaDecodeTs(Nvfp4KVConfig): """Throughput-latency 1CTA MLA over dense pages or prepared sparse rows.""" + lse_base_on_e = False + runtime_split_kv = False + runtime_kv_lengths = False + def __init__( self, *, @@ -1675,12 +1764,22 @@ def __init__( logical_seq_len_q: int, tile_size_q: int, explicit_split_kv: int | None = None, + split_kv_capacity: int | None = None, + reduction_policy_splits: int | None = None, explicit_persistent: bool | None = None, mask_type: MaskType | str = MaskType.CAUSAL, sparse_profile=None, finalize_output: bool = False, + softmax_rescale_threshold_log2: float = 0.0, + kv_cache_format: str = "native", + extra_kv_cache_format: str = "native", + fixed_seqlen: bool = False, ): - """Initialize one selected physical tile profile over logical flat Q rows.""" + """Initialize one selected physical tile profile over logical flat Q rows. + + ``fixed_seqlen=False`` (default) makes the KV and Q lengths runtime + launch operands; ``True`` compiles for the exact static lengths. + """ import cutlass as _cutlass if acc_dtype is None: @@ -1705,16 +1804,24 @@ def __init__( self.logical_seq_len_q = logical_seq_len_q self.tile_size_q = tile_size_q self.explicit_split_kv = explicit_split_kv + self.split_kv_capacity = split_kv_capacity self.explicit_persistent = explicit_persistent self.mask_type = normalize_mask_type(mask_type) self.sparse_profile = sparse_profile + self.kv_cache_format = kv_cache_format + self.extra_kv_cache_format = extra_kv_cache_format self.finalize_output = finalize_output + self.softmax_rescale_threshold_log2 = ( + canonicalize_softmax_rescale_threshold_log2( + softmax_rescale_threshold_log2, qkv_dtype=qkv_dtype + ) + ) + self.fixed_seqlen = fixed_seqlen self.output_quant = None # Sparse attention binds separate source lists; dense attention uses # the ordinary page table. Both reuse the staged KV loaders. self.direct_sparse = False self.assume_valid_prefix = False - self.direct_static_scales = False self.direct_sparse_capacities = (0, 0) cfg = self._make_config() @@ -1750,7 +1857,9 @@ def __init__( ) cluster_size = ( choose_q64_parallel_reducer_cluster_size( - cfg.num_ctas_per_seq_kv, + cfg.num_ctas_per_seq_kv + if reduction_policy_splits is None + else reduction_policy_splits, base_clusters=base_clusters, sm_count=self.max_active_clusters, ) @@ -1766,21 +1875,45 @@ def __init__( self.acc_dtype = acc_dtype self.lse_dtype = lse_dtype + @property + def runtime_query_lengths(self) -> bool: + """Q lengths are runtime launch operands unless ``fixed_seqlen``. + + Sparse MLA flattens Q rows into virtual batch items, so its physical Q + length is one and always stays compile-time; only KV is runtime. + """ + return not self.fixed_seqlen and self.sparse_profile is None + def _parallel_reduction_setup(self, cfg): + """Select reducer geometry for complete D512 heads and the planned split capacity. + + Unquantized retained-KV M64 uses two heads per block; quantized output + keeps one head so its group/block scale ownership stays unchanged.""" if cfg.fuse_sparse_reduction: if cfg.logical_seq_len_q != 1: raise ValueError( "final sparse reduction requires one virtual query per request" ) - # Partials already contain every V partition. Finalize one whole - # D512 head per CTA, sharing its statistics across four vec4 warps. + # Partials already contain every V partition. Each slice covers + # complete D512 heads and shares their statistics across its warps. reducer_cfg = dataclass_replace( cfg, num_ctas_per_head_dim=1, head_dim_per_cta_v=cfg.head_dim_v ) topology = make_balanced_parallel_reduction_topology( cfg.num_ctas_per_seq_kv, cluster_size=1 ) - return reducer_cfg, topology, cfg.head_dim_v + # Two complete heads per block amortize sparse reduction setup + # and use vec8 loads/stores. Keep quantized output on its existing + # one-head path until its wider-fragment contract is qualified. + heads_per_cta = ( + 2 + if cfg.has_nvfp4_kv + and cfg.whole_kv_tile + and cfg.tile_size_q == 64 + and self.output_quant is None + else 1 + ) + return reducer_cfg, topology, cfg.head_dim_v * heads_per_cta return ( cfg, self.parallel_reduction_topology, @@ -1805,7 +1938,7 @@ def compile_topology_signature(self) -> tuple[object, ...]: ) return ( "reference", - grid[0], + None if self.runtime_query_lengths else grid[0], grid[1], smem, threads, @@ -1816,10 +1949,15 @@ def compile_signature(self) -> tuple[object, ...]: """Return the complete batch-independent JIT identity.""" cfg = dataclass_replace(self._make_config(), batch_size=1) + if self.runtime_query_lengths: + cfg = dataclass_replace( + cfg, seq_len_q=0, logical_seq_len_q=0, num_ctas_per_seq_q=0 + ) return ( cfg, self.acc_dtype, self.lse_dtype, + self.runtime_split_kv, self.use_parallel_reduction, self.parallel_reduction_topology, self.parallel_reduction_elements_per_slice, @@ -1829,6 +1967,13 @@ def compile_signature(self) -> tuple[object, ...]: def _make_config(self): """Create the static throughput-latency MLA config for this wrapper.""" + validate_nvfp4_kv_formats( + self.kv_cache_format, + self.extra_kv_cache_format, + self.rope_dim, + self.qkv_dtype, + self.page_size, + ) cfg = make_throughput_latency_mla_config( batch_size=self.batch_size, num_heads_q=self.num_heads, @@ -1849,6 +1994,13 @@ def _make_config(self): explicit_split_kv=self.explicit_split_kv, explicit_persistent=self.explicit_persistent, mask_type=self.mask_type, + softmax_rescale_threshold_log2=self.softmax_rescale_threshold_log2, + sparse_profile=self.sparse_profile, + ) + cfg = dataclass_replace( + cfg, + kv_cache_format=self.kv_cache_format, + extra_kv_cache_format=self.extra_kv_cache_format, ) if self.sparse_profile is not None: cfg = configure_sparse_mla( @@ -1857,9 +2009,51 @@ def _make_config(self): finalize_output=self.finalize_output, direct_inputs=self.direct_sparse, assume_valid_prefix=self.assume_valid_prefix, - static_scales=self.direct_static_scales, source_capacities=self.direct_sparse_capacities, ) + if self.split_kv_capacity is not None and cfg.use_multi_ctas_kv == 1: + cfg = dataclass_replace(cfg, num_ctas_per_seq_kv=self.split_kv_capacity) + if not self.fixed_seqlen: + group = cfg.fixed_sparse_kv_tiles + cfg = dataclass_replace( + cfg, + seq_len_kv=None, + max_num_pages_per_seq_kv=0, + sparse_instruction_group=group, + sparse_direct_capacities=tuple( + int(x > 0) for x in cfg.sparse_direct_capacities + ), + ) + return dataclass_replace( + cfg, + fixed_seqlen=not self.runtime_query_lengths, + lse_base_on_e=self.lse_base_on_e, + ) + + def runtime_config(self, max_seq_len_q): + """Bind live Q dimensions without changing the selected resource topology.""" + cfg = self._make_config() + if self.runtime_query_lengths: + if cfg.tile_size_q % cfg.logical_num_heads_q == 0: + queries_per_tile = cfg.tile_size_q // cfg.logical_num_heads_q + # Q bounds are nonnegative. Factor the static head ratio and + # use unsigned ceil division without a signed adjustment. + q_tiles = Int32( + (cutlass.Uint32(max_seq_len_q) + queries_per_tile - 1) + // queries_per_tile + ) + elif cfg.logical_num_heads_q % cfg.tile_size_q == 0: + q_tiles = max_seq_len_q * (cfg.logical_num_heads_q // cfg.tile_size_q) + else: + q_tiles = ( + max_seq_len_q * cfg.logical_num_heads_q + cfg.tile_size_q - 1 + ) // cfg.tile_size_q + cfg = dataclass_replace( + cfg, + logical_seq_len_q=max_seq_len_q, + seq_len_q=q_tiles, + num_ctas_per_seq_q=q_tiles, + ) return cfg def validate_split_kv_launch(self, split_kv: int, workspace) -> None: @@ -1874,10 +2068,8 @@ def validate_split_kv_launch(self, split_kv: int, workspace) -> None: raise ValueError( "multi-CTA-KV throughput-latency 1CTA MLA requires workspace" ) - if cfg.use_multi_ctas_kv == 1 and split_kv < cfg.num_ctas_per_seq_kv: - raise ValueError( - "split_kv is smaller than the configured multi-CTA-KV split count" - ) + if cfg.use_multi_ctas_kv == 1 and not 1 <= split_kv <= cfg.num_ctas_per_seq_kv: + raise ValueError("split_kv must be within the prepared split capacity") def initialize_workspace( self, @@ -1935,9 +2127,10 @@ def __call__( output_scale: cutlass.Float32, stream: object, output_quant_args=None, + max_q_len: cutlass.Int32 = 0, ): """Execute throughput-latency 1CTA MLA with fixed or compact ragged Q.""" - cfg = self._make_config() + cfg = self.runtime_config(max_q_len) # Public fixed tensors use [H,D,SQ,B]/[H,SQ,B]; compact ragged tensors # use [H,D,totalQ]/[H,totalQ]. Both flatten H x Q for TMA because the @@ -1972,18 +2165,19 @@ def __call__( o.stride[3] == o.shape[2] * o.stride[2], "o must be compact from the query axis into the batch axis", ) - if cutlass.const_expr(lse.stride[0] != 1): - raise ValueError("lse must have leading dimension 0") - runtime_assert( - lse.stride[1] == lse.shape[0] * lse.stride[0], - "lse must be compact from the head axis into the query axis", - ) - if cutlass.const_expr(cu_seqlens_q is None): + if cutlass.const_expr(lse is not None): + if cutlass.const_expr(lse.stride[0] != 1): + raise ValueError("lse must have leading dimension 0") runtime_assert( - lse.stride[2] == lse.shape[1] * lse.stride[1], - "lse must be compact from the query axis into the batch axis", + lse.stride[1] == lse.shape[0] * lse.stride[0], + "lse must be compact from the head axis into the query axis", ) - else: + if cutlass.const_expr(cu_seqlens_q is None): + runtime_assert( + lse.stride[2] == lse.shape[1] * lse.stride[1], + "lse must be compact from the query axis into the batch axis", + ) + if cutlass.const_expr(cu_seqlens_q is not None): runtime_assert( cute.size(cu_seqlens_q) == cute.size(cache_seqs) + Int32(1), "cu_seqlens_q must contain batch_size + 1 offsets", @@ -2000,8 +2194,8 @@ def __call__( if cutlass.const_expr(cfg.use_multi_ctas_kv == 1): workspace_split_kv = cutlass.Int32(cfg.num_ctas_per_seq_kv) runtime_assert( - split_kv >= workspace_split_kv, - "split_kv is smaller than the configured multi-CTA-KV split count", + split_kv >= Int32(1) and split_kv <= workspace_split_kv, + "split_kv must be within the prepared split capacity", ) if cutlass.const_expr(cu_seqlens_q is not None): @@ -2016,6 +2210,7 @@ def __call__( q_latent_tma, box_dims=(tma_box0, cfg.tile_size_q), ragged_dim=1, + use_blackwell_tma_oob_workaround=True, stride_order=(0, 1), swizzle=cuda.TensorMapSwizzle.s128b, l2_promotion=cuda.TensorMapL2Promotion.l2_128b, @@ -2060,6 +2255,7 @@ def __call__( q_rope_tma, box_dims=(min(tma_box0, cfg.rope_dim), cfg.tile_size_q), ragged_dim=1, + use_blackwell_tma_oob_workaround=True, stride_order=(0, 1), swizzle=q_rope_swizzle, l2_promotion=cuda.TensorMapL2Promotion.l2_128b, @@ -2095,8 +2291,9 @@ def __call__( cute.select(c_latent.layout, mode=[1, 0, 2]), ) if cutlass.const_expr(cfg.num_tokens_per_page == 1): - # Sparse routes address individual rows; the unused D512 RoPE - # descriptor slot carries the independent compressed-pool map. + # Sparse routes address individual rows (Gather4 row maps). D512: + # the unused RoPE descriptor slot carries the independent + # compressed-pool map. D576: it carries the RoPE columns. c_rows = cute.make_tensor( c_latent.iterator, cute.make_layout( @@ -2120,9 +2317,13 @@ def __call__( ) tma_desc_c_rope = create_tensor_map_tiled_from_view( c_extra_rows, - box_dims=(tma_box0, 1), + box_dims=(cfg.rope_dim if cfg.sparse_d576_reload else tma_box0, 1), stride_order=(0, 1), - swizzle=cuda.TensorMapSwizzle.s128b, + swizzle=( + cuda.TensorMapSwizzle.s64b + if cfg.sparse_d576_reload + else cuda.TensorMapSwizzle.s128b + ), l2_promotion=cuda.TensorMapL2Promotion.l2_128b, ) else: @@ -2134,7 +2335,7 @@ def __call__( l2_promotion=cuda.TensorMapL2Promotion.l2_128b, ) - if cutlass.const_expr(cfg.rope_dim > 0): + if cutlass.const_expr(cfg.rope_dim > 0 and cfg.num_tokens_per_page != 1): c_rope_tma = cute.make_tensor( c_rope.iterator, cute.select(c_rope.layout, mode=[1, 0, 2]), @@ -2153,9 +2354,14 @@ def __call__( tma_desc_c_rope = tma_desc_c_latent tma_desc_v = tma_desc_c_latent - if cutlass.const_expr(cfg.num_tokens_per_page == 1): - # D512 does not use Q-rope; reuse its descriptor slot for the - # second pool's contiguous-run map. Gather maps stay separate. + tma_desc_bulk_extra = tma_desc_c_rope + tma_desc_sparse_rope = tma_desc_c_latent + tma_desc_sparse_extra_rope = tma_desc_c_rope + if cutlass.const_expr( + cfg.num_tokens_per_page == 1 and not cfg.sparse_d576_reload + ): + # Generic sparse bulk maps are separate from query RoPE. The + # native reload schedule uses its compact 64-column RoPE map. rows_per_issuer = 32 tma_desc_v = create_tensor_map_tiled_from_view( c_rows, @@ -2164,7 +2370,7 @@ def __call__( swizzle=cuda.TensorMapSwizzle.s128b, l2_promotion=cuda.TensorMapL2Promotion.l2_128b, ) - tma_desc_q_rope = create_tensor_map_tiled_from_view( + tma_desc_bulk_extra = create_tensor_map_tiled_from_view( c_extra_rows, box_dims=(tma_box0, rows_per_issuer), stride_order=(0, 1), @@ -2172,7 +2378,35 @@ def __call__( l2_promotion=cuda.TensorMapL2Promotion.l2_128b, ) - softmax_scale_log2 = softmax_scale * LOG2_E + if cutlass.const_expr( + cfg.num_tokens_per_page == 1 + and cfg.rope_dim + and not cfg.has_nvfp4_kv + and not cfg.sparse_d576_reload + ): + tma_desc_sparse_rope = make_sparse_rope_map(c_latent, cfg.qkv_dtype_bytes) + tma_desc_sparse_extra_rope = make_sparse_rope_map( + c_rope, cfg.qkv_dtype_bytes + ) + + if cutlass.const_expr(cfg.has_nvfp4_kv): + ( + tma_desc_c_latent, + tma_desc_c_rope, + tma_desc_v, + tma_desc_bulk_extra, + tma_desc_sparse_rope, + tma_desc_sparse_extra_rope, + ) = make_nvfp4_kv_maps( + c_latent, + c_rope, + cfg.rope_dim, + whole_record=cfg.async_nvfp4_kv_gather, + primary_format=cfg.kv_cache_format, + extra_format=cfg.extra_kv_cache_format, + ) + + scale_softmax_log2 = softmax_scale * LOG2_E use_gmem_reduction = cutlass.const_expr( cfg.use_multi_ctas_kv == 1 and cfg.use_cluster_reduction != 1 ) @@ -2223,16 +2457,23 @@ def __call__( self.max_active_clusters, ) else: + launch_splits = Int32(cfg.num_ctas_per_seq_kv) + if cutlass.const_expr(self.runtime_split_kv and use_gmem_reduction): + launch_splits = split_kv grid = ( - cfg.num_ctas_for_all_heads - * cfg.num_ctas_per_seq_q - * cfg.num_ctas_per_seq_kv, + cfg.num_ctas_for_all_heads * cfg.num_ctas_per_seq_q * launch_splits, cfg.num_ctas_per_head_dim, batch_size, ) cluster_shape = None if cutlass.const_expr(cfg.use_cluster_reduction == 1): cluster_shape = (cfg.num_ctas_per_seq_kv, 1, 1) + requested_splits = None + if cutlass.const_expr(self.runtime_split_kv): + requested_splits = ( + split_kv, + cute.fast_divmod_create_divisor(split_kv * Int32(cfg.num_insts_kv)), + ) self.dense_kernel( tma_desc_q_latent, tma_desc_q_rope, @@ -2247,11 +2488,18 @@ def __call__( acc_lse, cache_seqs, cu_seqlens_q, - softmax_scale_log2, + scale_softmax_log2, output_scale, tile_sched_params, block_split_kvs, output_quant_args, + tma_desc_bulk_extra, + tma_desc_sparse_rope, + tma_desc_sparse_extra_rope, + c_latent, + q_latent, + requested_splits=requested_splits, + max_q_len=max_q_len, ).launch( grid=grid, block=[cfg.threads_per_cta, 1, 1], @@ -2265,9 +2513,9 @@ def __call__( self.use_parallel_reduction or cfg.fuse_sparse_reduction ): reducer_cfg, topology, elements = self._parallel_reduction_setup(cfg) - atten_sinks = None + attention_sinks = None if cutlass.const_expr(cfg.fuse_sparse_reduction): - atten_sinks = block_split_kvs[-1] + attention_sinks = block_split_kvs[4] reduction_grid, reduction_cluster = ( parallel_gmem_reduction_launch_shape( reducer_cfg, @@ -2281,7 +2529,12 @@ def __call__( batch_size, ) if cutlass.const_expr(cfg.fuse_sparse_reduction): - reduction_grid = (cfg.logical_num_heads_q, 1, batch_size) + heads_per_cta = elements // cfg.head_dim_v + reduction_grid = ( + (cfg.logical_num_heads_q + heads_per_cta - 1) // heads_per_cta, + 1, + batch_size, + ) reduction_threads = parallel_gmem_reduction_threads(elements) parallel_reducer = self.parallel_gmem_reduction_kernel( o, @@ -2290,8 +2543,10 @@ def __call__( acc_lse, cache_seqs, cu_seqlens_q, - atten_sinks, + attention_sinks, output_quant_args, + requested_splits=requested_splits, + max_q_len=max_q_len, ) if cutlass.const_expr(topology.cluster_size == 1): parallel_reducer.launch( @@ -2336,6 +2591,8 @@ def __call__( cache_seqs, cu_seqlens_q, cutlass.Int32(reduction_ctas), + requested_splits=requested_splits, + max_q_len=max_q_len, ).launch( grid=reduction_grid, block=[reduction_threads, 1, 1], @@ -2361,36 +2618,47 @@ def dense_kernel( acc_lse: cute.Tensor, cache_seqs: cute.Tensor, cu_seqlens_q: cute.Tensor, - softmax_scale_log2: cutlass.Float32, + scale_softmax_log2: cutlass.Float32, output_scale: cutlass.Float32, tile_sched_params: object, sparse_inputs: object, output_quant_args=None, + tma_desc_bulk_extra: cutlass.GridConstant[cuda.TensorMap] = None, + tma_desc_sparse_rope: cutlass.GridConstant[cuda.TensorMap] = None, + tma_desc_sparse_extra_rope: cutlass.GridConstant[cuda.TensorMap] = None, + c_latent: cute.Tensor = None, + q_nope_tensor: cute.Tensor = None, + requested_splits=None, + max_q_len: cutlass.Int32 = 0, ): """Execute one flat-Q, batch, KV-split, and V head-dimension tile.""" - cfg = self._make_config() + cfg = self.runtime_config(max_q_len) + if cutlass.const_expr(acc_o is not None): + # Match the PDL launch before reading live metadata or sparse scales. + prims.griddepcontrol(kind=prims.GridDepAction.WAIT) if cutlass.const_expr(cfg.sparse_direct): from ..helpers.sparse_indices import bind_sparse_kv_indices - si, ci, sl, cl, sm, os, sinks = sparse_inputs + si, ci, sl, cl, sinks, ks = sparse_inputs request = None if cutlass.const_expr(cfg.use_persistent_scheduler != 1): request = cutlass.Int64(cute.arch.block_idx()[2]) + capacity = ((cutlass.Int32(si.shape[-1]) + 127) // 128) * 128 + if cutlass.const_expr(cfg.sparse_direct_capacities[1] > 0): + capacity = capacity + ((cutlass.Int32(ci.shape[-1]) + 127) // 128) * 128 + capacity = cute.math.max(cutlass.Int32(256), capacity) page_offsets, cache_seqs = bind_sparse_kv_indices( si, ci, sl, cl, - cfg.total_kv_tiles * cfg.tile_size_kv, + capacity, query_row=request, assume_valid_prefix=cfg.sparse_valid_prefix, has_extra=cfg.sparse_direct_capacities[1] > 0, + primary_tensor_scale=ks if cfg.kv_cache_format == "fp8" else None, ) - if cutlass.const_expr(not cfg.sparse_static_scales): - # Caller scales already include Q/K/V descales. - softmax_scale_log2 = cutlass.Float32(sm[0]) * LOG2_E - output_scale = cutlass.Float32(os[0]) sparse_epilogue_params = sinks else: sparse_epilogue_params = None @@ -2399,18 +2667,41 @@ def dense_kernel( # once here; Q/O resources retain responsibility for logical-row # mapping, compact-ragged offsets, and padded-tail publication. cta_idx_x, cta_idx_head_dim_v, batch_idx = cute.arch.block_idx() + # Flat Q rows already include every logical head. The config requires + # one physical head tile, so decoding X never needs a runtime divide + # by the Q-tile count. + head_idx = Int32(0) + cta_idx_head_q = Int32(0) if cutlass.const_expr(cfg.use_cluster_reduction == 1): - cta_idx_x = cta_idx_x // Int32(cfg.num_ctas_per_seq_kv) - ctas_per_head_tile = Int32(cfg.num_ctas_per_seq_q) - cta_idx_head_q = cta_idx_x // ctas_per_head_tile - cta_idx_q = cta_idx_x - cta_idx_head_q * ctas_per_head_tile + cta_idx_q = cta_idx_x // Int32(cfg.num_ctas_per_seq_kv) + if cutlass.const_expr(not self.runtime_query_lengths): + cta_idx_head_q = cta_idx_q // Int32(cfg.num_ctas_per_seq_q) + cta_idx_q = cta_idx_q - cta_idx_head_q * Int32(cfg.num_ctas_per_seq_q) cta_idx_kv = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + elif cutlass.const_expr( + not self.runtime_query_lengths + and cfg.num_ctas_for_all_heads == 1 + and cfg.num_ctas_per_seq_q == 1 + ): + cta_idx_head_q = Int32(0) + cta_idx_q = Int32(0) + cta_idx_kv = cta_idx_x else: - ctas_per_head_tile = Int32(cfg.num_ctas_per_seq_q * cfg.num_ctas_per_seq_kv) - cta_idx_head_q = cta_idx_x // ctas_per_head_tile - cta_idx_x_in_head = cta_idx_x - cta_idx_head_q * ctas_per_head_tile - cta_idx_q = cta_idx_x_in_head // Int32(cfg.num_ctas_per_seq_kv) - cta_idx_kv = cta_idx_x_in_head - cta_idx_q * Int32(cfg.num_ctas_per_seq_kv) + grid_splits = ( + Int32(cfg.num_ctas_per_seq_kv) + if requested_splits is None + else runtime_split_count(requested_splits) + ) + if cutlass.const_expr(self.runtime_query_lengths): + cta_idx_head_q = Int32(0) + cta_idx_q = cta_idx_x // grid_splits + cta_idx_kv = cta_idx_x - cta_idx_q * grid_splits + else: + ctas_per_head_tile = Int32(cfg.num_ctas_per_seq_q) * grid_splits + cta_idx_head_q = cta_idx_x // ctas_per_head_tile + cta_idx_x_in_head = cta_idx_x - cta_idx_head_q * ctas_per_head_tile + cta_idx_q = cta_idx_x_in_head // grid_splits + cta_idx_kv = cta_idx_x_in_head - cta_idx_q * grid_splits head_idx = cta_idx_head_q * Int32(cfg.tile_size_q) run_task_graph = cutlass.Boolean(True) query_tile_is_active = cutlass.Boolean(True) @@ -2419,9 +2710,7 @@ def dense_kernel( cfg.use_multi_ctas_kv == 1 and runtime_split_pruning_is_profitable(cfg.num_ctas_per_seq_kv) ) - has_runtime_activity_guard = tile_sched_params is None and ( - cu_seqlens_q is not None or enable_runtime_split_pruning - ) + has_runtime_activity_guard = tile_sched_params is None if cutlass.const_expr(has_runtime_activity_guard): # Decode launch coordinates against the configured maximum grid, # then independently drop compact-Q padding and profitable split @@ -2444,8 +2733,14 @@ def dense_kernel( batch_idx, cta_idx_q, cta_idx_kv, + requested_splits=requested_splits, ) split_tile_is_active = cute.arch.make_warp_uniform(split_tile_is_active) + # An empty request publishes neither O/LSE nor split scratch. + # Cluster peers agree on this batch-owned predicate. + query_tile_is_active = query_tile_is_active & ( + runtime_request_publishes(cfg, cache_seqs, batch_idx) + ) run_task_graph = query_tile_is_active and split_tile_is_active if cutlass.const_expr(tile_sched_params is not None): cta_idx_q = None @@ -2489,7 +2784,12 @@ def dense_kernel( tma_desc_c_latent=tma_desc_c_latent.get_ptr(), tma_desc_c_rope=tma_desc_c_rope.get_ptr(), tma_desc_v=tma_desc_v.get_ptr(), + tma_desc_bulk_extra=tma_desc_bulk_extra.get_ptr(), + tma_desc_sparse_rope=tma_desc_sparse_rope.get_ptr(), + tma_desc_sparse_extra_rope=tma_desc_sparse_extra_rope.get_ptr(), c_rope_tensor=c_rope, + c_latent_tensor=c_latent, + q_tensor=q_nope_tensor, page_offsets=page_offsets, cache_seqs=cache_seqs, cu_seqlens_q=cu_seqlens_q, @@ -2498,7 +2798,7 @@ def dense_kernel( cta_idx_q=cta_idx_q, cta_idx_kv=cta_idx_kv, cta_idx_head_dim_v=cta_idx_head_dim_v, - scale_softmax_log2=softmax_scale_log2, + scale_softmax_log2=scale_softmax_log2, output_scale=output_scale, sparse_epilogue_params=sparse_epilogue_params, output_quant=self.output_quant @@ -2513,6 +2813,7 @@ def dense_kernel( clc_response_ptr=clc_response_ptr, use_clc_dynamic_scheduler=use_clc_dynamic, use_static_persistent_scheduler=use_static_persistent, + requested_splits=requested_splits, ) task_manager.setup_resources_and_tasks() smem_allocator = task_manager.smem_allocator @@ -2631,9 +2932,11 @@ def gmem_reduction_kernel( cache_seqs: cute.Tensor, cu_seqlens_q: cute.Tensor, num_reduction_ctas: cutlass.Int32, + requested_splits=None, + max_q_len: cutlass.Int32 = 0, ): """Dispatch the throughput-latency split-KV reduction body.""" - cfg = self._make_config() + cfg = self.runtime_config(max_q_len) run_gmem_reduction_kernel( self, output, @@ -2644,6 +2947,7 @@ def gmem_reduction_kernel( cu_seqlens_q, cfg, num_reduction_ctas, + requested_splits=requested_splits, ) @cute.kernel @@ -2655,12 +2959,14 @@ def parallel_gmem_reduction_kernel( acc_lse: cute.Tensor, cache_seqs: cute.Tensor, cu_seqlens_q: cute.Tensor, - atten_sinks: cute.Tensor | None = None, + attention_sinks: cute.Tensor | None = None, output_quant_args=None, + requested_splits=None, + max_q_len: cutlass.Int32 = 0, ): """Dispatch the automatically selected parallel standalone reducer.""" - cfg = self._make_config() + cfg = self.runtime_config(max_q_len) cfg, topology, elements = self._parallel_reduction_setup(cfg) run_parallel_gmem_reduction_kernel( output, @@ -2674,8 +2980,9 @@ def parallel_gmem_reduction_kernel( topology.slots_per_rank, topology.actual_splits, elements, - atten_sinks, + attention_sinks, cfg.fuse_sparse_reduction, self.output_quant, output_quant_args, + requested_splits=requested_splits, ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/parallel_reduction.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/parallel_reduction.py index 11a40b52bfd1..373bb2bea546 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/parallel_reduction.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/parallel_reduction.py @@ -24,14 +24,16 @@ normalized_lse_weight, attention_sink_scale, ) -from ..helpers.constants import LN_2, SPLIT_REDUCTION_SCALE_BARRIER_ID -from ..helpers.math import ceil_div +from ..helpers.constants import SPLIT_REDUCTION_SCALE_BARRIER_ID +from ..helpers.tile import runtime_split_count +from ..helpers.math import ceil_div, final_lse, ffma2 from ..helpers.ops import ( fmax_f32, warp_reduce_max_f32, warp_reduce_sum_f32, ) from ..helpers.query import flat_query_row_state, public_query_flat_row +from ..helpers.tile import runtime_request_publishes from ..parallel_reduction_topology import ParallelReductionTopology from .config import MlaConfig @@ -102,7 +104,7 @@ def supports_parallel_gmem_reduction(cfg: MlaConfig) -> bool: and cfg.use_cluster_reduction != 1 and cfg.head_dim_per_cta_v in (128, 256, 512) and cfg.tile_size_q in (8, 16, 32, 64) - and 2 <= cfg.num_ctas_per_seq_kv <= 128 + and cfg.num_ctas_per_seq_kv >= 2 and ( (cfg.o_dtype == "bf16" and cfg.use_bf16_output == 1) or (cfg.o_dtype == "e4m3" and cfg.use_fp8_output == 1) @@ -201,22 +203,32 @@ def _run_parallel_gmem_reduction_g1_shared_stats( cu_seqlens_q, cfg, elements_per_slice: cutlass.Constexpr[int], - atten_sinks=None, + attention_sinks=None, finalize_attention: cutlass.Constexpr[bool] = False, output_quant: cutlass.Constexpr = None, output_quant_args=None, + requested_splits=None, + request_is_active=True, ): """Reduce G1 partials with the compact row-shared schedule.""" + split_limit = ( + Int32(cfg.num_ctas_per_seq_kv) + if requested_splits is None + else runtime_split_count(requested_splits) + ) + slice_idx, block_idx_y, batch_idx = cute.arch.block_idx() thread_idx, _, _ = cute.arch.thread_idx() warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) lane_idx = thread_idx % Int32(GMEM_REDUCTION_WARP_LANES) if cutlass.const_expr(finalize_attention): - # The final sparse grid is one CTA per logical head. Decode its - # physical producer tile without launching padded head/V slices. - q_idx = slice_idx // Int32(cfg.tile_size_q) - slice_idx = slice_idx % Int32(cfg.tile_size_q) + # Sparse slices contain complete heads. Map the head group back + # to its physical producer tile; row validity masks a final odd head. + heads_per_slice = elements_per_slice // cfg.head_dim_v + slices_per_q_tile = cfg.tile_size_q // heads_per_slice + q_idx = slice_idx // Int32(slices_per_q_tile) + slice_idx = slice_idx % Int32(slices_per_q_tile) head_dim_cta_idx = Int32(0) else: head_dim_cta_idx = block_idx_y % Int32(cfg.num_ctas_per_head_dim) @@ -253,7 +265,10 @@ def _run_parallel_gmem_reduction_g1_shared_stats( cu_seqlens_q=cu_seqlens_q, batch_idx=batch_idx, ) - if stats_head_idx < Int32(cfg.num_heads_q) and stats_query_row[4]: + valid_stats = stats_head_idx < Int32(cfg.num_heads_q) and stats_query_row[4] + if cutlass.const_expr(not finalize_attention): + valid_stats = valid_stats & request_is_active + if valid_stats: row_lse = acc_lse[stats_head_idx, None, q_idx, batch_idx] lse_per_lane = ceil_div( cfg.num_ctas_per_seq_kv, @@ -266,30 +281,38 @@ def _run_parallel_gmem_reduction_g1_shared_stats( ) lse_max = Float32(-Float32.inf) for lane_slot_i in cutlass.range_constexpr(lse_per_lane): - stats_split_idx = lane_idx + Int32( - lane_slot_i * GMEM_REDUCTION_WARP_LANES - ) - lane_lse[lane_slot_i] = ( - Float32(row_lse[stats_split_idx]) - if stats_split_idx < Int32(cfg.num_ctas_per_seq_kv) - else Float32(-Float32.inf) - ) - lse_max = fmax_f32(lse_max, lane_lse[lane_slot_i]) + if ( + cutlass.const_expr(lane_slot_i == 0) + or Int32(lane_slot_i * GMEM_REDUCTION_WARP_LANES) < split_limit + ): + stats_split_idx = lane_idx + Int32( + lane_slot_i * GMEM_REDUCTION_WARP_LANES + ) + lane_lse[lane_slot_i] = ( + Float32(row_lse[stats_split_idx]) + if stats_split_idx < split_limit + else Float32(-Float32.inf) + ) + lse_max = fmax_f32(lse_max, lane_lse[lane_slot_i]) lse_max = warp_reduce_max_f32(lse_max) lse_max = lse_max if lse_max != Float32(-Float32.inf) else Float32(0.0) lse_sum = Float32(0.0) for lane_slot_i in cutlass.range_constexpr(lse_per_lane): - lse_sum += cute.math.exp2( - lane_lse[lane_slot_i] - lse_max, - fastmath=True, - ) + if ( + cutlass.const_expr(lane_slot_i == 0) + or Int32(lane_slot_i * GMEM_REDUCTION_WARP_LANES) < split_limit + ): + lse_sum += cute.math.exp2( + lane_lse[lane_slot_i] - lse_max, + fastmath=True, + ) lse_sum = warp_reduce_sum_f32(lse_sum) global_lse = finalize_log2_sum_exp(lse_max, lse_sum) stats_sink_scale = Float32(1) - if cutlass.const_expr(finalize_attention and atten_sinks is not None): + if cutlass.const_expr(finalize_attention and attention_sinks is not None): stats_sink_scale = attention_sink_scale( - global_lse, Float32(atten_sinks[stats_query_row[1]]) + global_lse, Float32(attention_sinks[stats_query_row[1]]) ) if ( lane_idx == Int32(0) @@ -302,23 +325,28 @@ def _run_parallel_gmem_reduction_g1_shared_stats( batch_idx, cu_seqlens_q, ) - (lse.iterator.raw_ptr() + stats_output_query_row).store( - global_lse * Float32(LN_2) if finalize_attention else global_lse - ) + if cutlass.const_expr(lse is not None): + (lse.iterator.raw_ptr() + stats_output_query_row).store( + final_lse(global_lse, cfg.lse_base_on_e or finalize_attention) + ) for lane_slot_i in cutlass.range_constexpr(lse_per_lane): - stats_split_idx = lane_idx + Int32( - lane_slot_i * GMEM_REDUCTION_WARP_LANES - ) - if stats_split_idx < Int32(cfg.num_ctas_per_seq_kv): - stats_weight = normalized_lse_weight( - lane_lse[lane_slot_i], global_lse + if ( + cutlass.const_expr(lane_slot_i == 0) + or Int32(lane_slot_i * GMEM_REDUCTION_WARP_LANES) < split_limit + ): + stats_split_idx = lane_idx + Int32( + lane_slot_i * GMEM_REDUCTION_WARP_LANES ) - if cutlass.const_expr(finalize_attention): - stats_weight *= stats_sink_scale - smem_scale[ - stats_row * Int32(cfg.num_ctas_per_seq_kv) + stats_split_idx - ] = stats_weight + if stats_split_idx < split_limit: + stats_weight = normalized_lse_weight( + lane_lse[lane_slot_i], global_lse + ) + if cutlass.const_expr(finalize_attention): + stats_weight *= stats_sink_scale + smem_scale[ + stats_row * Int32(cfg.num_ctas_per_seq_kv) + stats_split_idx + ] = stats_weight prims.barrier_cta_sync( barrier_id=SPLIT_REDUCTION_SCALE_BARRIER_ID, @@ -346,6 +374,10 @@ def _run_parallel_gmem_reduction_g1_shared_stats( and dim_idx < Int32(cfg.head_dim_v) and output_query_row_state[4] ) + # Sparse finalization must publish even if the primary source is empty: + # the extra source may be live, and two empty sources produce zero output. + if cutlass.const_expr(not finalize_attention): + valid_output = valid_output & request_is_active if valid_output: output_acc = cutlass.Array( Float32, @@ -354,9 +386,10 @@ def _run_parallel_gmem_reduction_g1_shared_stats( ) for elem_i in cutlass.range_constexpr(output_elements_per_thread): output_acc[elem_i] = Float32(0.0) - # Form the 64-bit row base once. Split strides are compile-time - # constants in this fixed-S kernel, avoiding repeated wide integer - # linearization in the hot accumulation loop. + # Scratch strides remain static for the prepared capacity. Form the + # 64-bit row base once to avoid repeated address linearization. + # A full prefix keeps the unrolled path without one branch per slot; + # a shorter prefix uses a runtime loop and never reads padding. acc_row_ptr = acc_output.iterator.raw_ptr() + _split_o_row_base_offset( cfg, batch_idx, @@ -364,25 +397,65 @@ def _run_parallel_gmem_reduction_g1_shared_stats( head_idx, dim_idx, ) - for split_i in cutlass.range_constexpr(cfg.num_ctas_per_seq_kv): - split_idx = Int32(split_i) - scale = Float32( - smem_scale[row_in_slice * Int32(cfg.num_ctas_per_seq_kv) + split_idx] - ) - # Sparse masked splits can leave their vector unused. Do not - # multiply an uninitialized/NaN partial by a zero statistical weight. - if not finalize_attention or scale > Float32(0): - split_ptr = acc_row_ptr + Int64(split_i * cfg.head_dim_v) - if cutlass.const_expr(output_elements_per_thread == 1): - output_acc[0] += Float32(split_ptr.load()) * scale - else: - partial_o = split_ptr.load( - count=output_elements_per_thread, - alignment=output_elements_per_thread - * cfg.partial_o_dtype_bytes, - ).to(Float32) - for elem_i in cutlass.range_constexpr(output_elements_per_thread): - output_acc[elem_i] += partial_o[elem_i] * scale + # Pair FP32 channels while preserving each channel's split order. + unrolled_splits = ( + cfg.num_ctas_per_seq_kv + if requested_splits is None + else GMEM_REDUCTION_WARP_LANES + ) + if split_limit == Int32(unrolled_splits): + for split_i in cutlass.range_constexpr(unrolled_splits): + split_idx = Int32(split_i) + scale = Float32( + smem_scale[ + row_in_slice * Int32(cfg.num_ctas_per_seq_kv) + split_idx + ] + ) + # Sparse inactive splits may leave their O partial untouched. + if not finalize_attention or scale > Float32(0): + split_ptr = acc_row_ptr + Int64(split_i * cfg.head_dim_v) + if cutlass.const_expr(output_elements_per_thread == 1): + output_acc[0] += Float32(split_ptr.load()) * scale + else: + partial_o = split_ptr.load( + count=output_elements_per_thread, + alignment=output_elements_per_thread + * cfg.partial_o_dtype_bytes, + ).to(Float32) + for elem_i in cutlass.range_constexpr( + 0, output_elements_per_thread, 2 + ): + output_acc[elem_i], output_acc[elem_i + 1] = ffma2( + (partial_o[elem_i], partial_o[elem_i + 1]), + (scale, scale), + (output_acc[elem_i], output_acc[elem_i + 1]), + ) + else: + for split_i in cutlass.range(split_limit, unroll=4): + split_idx = Int32(split_i) + scale = Float32( + smem_scale[ + row_in_slice * Int32(cfg.num_ctas_per_seq_kv) + split_idx + ] + ) + if not finalize_attention or scale > Float32(0): + split_ptr = acc_row_ptr + Int64(split_i) * Int64(cfg.head_dim_v) + if cutlass.const_expr(output_elements_per_thread == 1): + output_acc[0] += Float32(split_ptr.load()) * scale + else: + partial_o = split_ptr.load( + count=output_elements_per_thread, + alignment=output_elements_per_thread + * cfg.partial_o_dtype_bytes, + ).to(Float32) + for elem_i in cutlass.range_constexpr( + 0, output_elements_per_thread, 2 + ): + output_acc[elem_i], output_acc[elem_i + 1] = ffma2( + (partial_o[elem_i], partial_o[elem_i + 1]), + (scale, scale), + (output_acc[elem_i], output_acc[elem_i + 1]), + ) output_query_row = public_query_flat_row( cfg, @@ -427,6 +500,8 @@ def _run_parallel_gmem_reduction_shared_stats( slots_per_rank: cutlass.Constexpr[int], actual_splits: cutlass.Constexpr[int], elements_per_slice: cutlass.Constexpr[int], + requested_splits=None, + request_is_active=True, ): """Reduce split-KV with one cooperative statistics warp per output row. @@ -436,6 +511,12 @@ def _run_parallel_gmem_reduction_shared_stats( zero repeats the same row-shared scheme for the final DSMEM merge. """ + split_limit = ( + Int32(cfg.num_ctas_per_seq_kv) + if requested_splits is None + else runtime_split_count(requested_splits) + ) + block_idx_x, block_idx_y, batch_idx = cute.arch.block_idx() thread_idx, _, _ = cute.arch.thread_idx() cluster_rank = cute.arch.block_idx_in_cluster() @@ -490,6 +571,7 @@ def _run_parallel_gmem_reduction_shared_stats( batch_idx=batch_idx, ) valid_stats = stats_head_idx < Int32(cfg.num_heads_q) and valid_stats_row + valid_stats = valid_stats & request_is_active lse_per_lane = ceil_div(local_slots, GMEM_REDUCTION_WARP_LANES) lane_lse = cutlass.Array( Float32, @@ -515,8 +597,10 @@ def _run_parallel_gmem_reduction_shared_stats( lane_slot_i * GMEM_REDUCTION_WARP_LANES ) split_idx = local_slot_idx * Int32(cluster_size) + cluster_rank - if local_slot_idx < Int32(local_slots) and split_idx < Int32( - actual_splits + if ( + local_slot_idx < Int32(local_slots) + and split_idx < Int32(actual_splits) + and split_idx < split_limit ): lane_lse[lane_slot_i] = Float32(row_lse[split_idx]) local_max = fmax_f32(local_max, lane_lse[lane_slot_i]) @@ -577,6 +661,7 @@ def _run_parallel_gmem_reduction_shared_stats( and dim_idx < Int32(cfg.head_dim_v) and valid_output_row ) + valid_output = valid_output & request_is_active local_output = cutlass.Array( Float32, output_elements_per_thread, @@ -599,27 +684,28 @@ def _run_parallel_gmem_reduction_shared_stats( ) + Int64(cluster_rank) * Int64(cfg.head_dim_v) ) - for local_slot_i in cutlass.range_constexpr(local_slots): + # Each rank traverses its live cyclic split prefix, independent of capacity. + for local_slot_i in cutlass.range( + cute.ceil_div(split_limit - cluster_rank, cluster_size), unroll=32 + ): split_idx = Int32(local_slot_i * cluster_size) + cluster_rank - if split_idx < Int32(actual_splits): - scale = Float32( - smem_local_scale[ - row_in_slice * Int32(local_slots) + Int32(local_slot_i) - ] - ) - split_ptr = acc_row_ptr + Int64( - local_slot_i * cluster_size * cfg.head_dim_v - ) - if cutlass.const_expr(output_elements_per_thread == 1): - local_output[0] += Float32(split_ptr.load()) * scale - else: - partial_o = split_ptr.load( - count=output_elements_per_thread, - alignment=output_elements_per_thread - * cfg.partial_o_dtype_bytes, - ).to(Float32) - for elem_i in cutlass.range_constexpr(output_elements_per_thread): - local_output[elem_i] += partial_o[elem_i] * scale + scale = Float32( + smem_local_scale[ + row_in_slice * Int32(local_slots) + Int32(local_slot_i) + ] + ) + split_ptr = acc_row_ptr + Int64( + local_slot_i * cluster_size * cfg.head_dim_v + ) + if cutlass.const_expr(output_elements_per_thread == 1): + local_output[0] += Float32(split_ptr.load()) * scale + else: + partial_o = split_ptr.load( + count=output_elements_per_thread, + alignment=output_elements_per_thread * cfg.partial_o_dtype_bytes, + ).to(Float32) + for elem_i in cutlass.range_constexpr(output_elements_per_thread): + local_output[elem_i] += partial_o[elem_i] * scale # Publish normalized O in BF16 while retaining FP32 accumulation. smem_o = cutlass.Array( @@ -686,6 +772,7 @@ def _run_parallel_gmem_reduction_shared_stats( batch_idx=batch_idx, ) valid_stats = stats_head_idx < Int32(cfg.num_heads_q) and valid_stats_row + valid_stats = valid_stats & request_is_active peer_lse = neg_inf for peer_rank_i in cutlass.range_constexpr(cluster_size): if valid_stats and lane_idx == Int32(peer_rank_i): @@ -760,9 +847,10 @@ def _run_parallel_gmem_reduction_shared_stats( cu_seqlens_q, ) if head_dim_cta_idx == Int32(0) and dim_in_cta == Int32(0): - (lse.iterator.raw_ptr() + output_query_row).store( - smem_global_lse[row_in_slice] - ) + if cutlass.const_expr(lse is not None): + (lse.iterator.raw_ptr() + output_query_row).store( + final_lse(smem_global_lse[row_in_slice], cfg.lse_base_on_e) + ) out_elem_offset = _output_element_offset(cfg, output_query_row, dim_idx) if cutlass.const_expr(output_elements_per_thread == 1): (output.iterator.raw_ptr() + out_elem_offset).store( @@ -802,17 +890,28 @@ def run_parallel_gmem_reduction_kernel( slots_per_rank: cutlass.Constexpr[int], actual_splits: cutlass.Constexpr[int], elements_per_slice: cutlass.Constexpr[int], - atten_sinks=None, + attention_sinks=None, finalize_attention: cutlass.Constexpr[bool] = False, output_quant: cutlass.Constexpr = None, output_quant_args=None, + requested_splits=None, ): """Reduce split-KV partials with row-shared local and peer statistics.""" if cutlass.const_expr(finalize_attention and cluster_size != 1): - raise ValueError("final attention output uses a single-CTA reducer per row") + raise ValueError( + "final attention output requires single-CTA reduction clusters" + ) prims.griddepcontrol(kind=prims.GridDepAction.WAIT) + # Reuse the CTA's predicate across the statistics barrier instead of + # reloading metadata. Sparse finalization can have a live extra source. + request_is_active = cutlass.Boolean(True) + if cutlass.const_expr(not finalize_attention): + request_is_active = runtime_request_publishes( + cfg, cache_seqs, cute.arch.block_idx()[2] + ) + if cutlass.const_expr(cluster_size == 1): _run_parallel_gmem_reduction_g1_shared_stats( output, @@ -823,10 +922,12 @@ def run_parallel_gmem_reduction_kernel( cu_seqlens_q, cfg, elements_per_slice, - atten_sinks, + attention_sinks, finalize_attention, output_quant, output_quant_args, + requested_splits=requested_splits, + request_is_active=request_is_active, ) return @@ -842,4 +943,6 @@ def run_parallel_gmem_reduction_kernel( slots_per_rank, actual_splits, elements_per_slice, + requested_splits=requested_splits, + request_is_active=request_is_active, ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/reduction.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/reduction.py index c5a694dbdba1..3b1650c3a75e 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/reduction.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/reduction.py @@ -19,6 +19,7 @@ from cutlass.experimental import primitives as prims from cutlass import Float32, Int32, Int64 +from ..helpers.math import final_lse from .config import MlaConfig from ..helpers.math import ceil_div from ..helpers.ops import ( @@ -32,7 +33,7 @@ public_query_flat_row, split_o_element_offset, ) -from ..helpers.tile import runtime_seq_len_kv_for_q +from ..helpers.tile import runtime_seq_len_kv_for_q, runtime_split_count # Baseline split-KV reduction uses one 128-thread CTA. Each thread owns eight @@ -83,7 +84,8 @@ def uses_two_head_gmem_reduction(cfg: MlaConfig) -> bool: def uses_q_row_gmem_reduction(cfg: MlaConfig) -> bool: """Return whether one reducer CTA covers multiple Q rows for one head.""" return ( - cfg.seq_len_q >= 2 * GMEM_REDUCTION_Q_ROWS_PER_CTA + cfg.fixed_seqlen + and cfg.seq_len_q >= 2 * GMEM_REDUCTION_Q_ROWS_PER_CTA and cfg.seq_len_q % GMEM_REDUCTION_Q_ROWS_PER_CTA == 0 and cfg.head_dim_v == GMEM_REDUCTION_VECTOR_HEAD_DIM and cfg.o_dtype_bytes == 2 @@ -131,6 +133,8 @@ def slice_split_num_reduction_ctas( ) -> int: """Return reducer CTAs per Q/head tile, capped to roughly two SM waves.""" + if not cfg.fixed_seqlen: + return 1 base_ctas = ( int(seq_len_q) * cfg.num_ctas_for_all_heads @@ -161,9 +165,9 @@ def gmem_reduction_launch_shape( ) return ( ( - int(seq_len_q) * num_reduction_ctas, + seq_len_q * num_reduction_ctas, cfg.num_ctas_for_all_heads * cfg.num_ctas_per_head_dim, - int(batch_size), + batch_size, ), 0, GMEM_REDUCTION_SLICE_THREADS, @@ -235,6 +239,38 @@ def run_gmem_reduction_kernel( cu_seqlens_q, cfg, num_reduction_ctas, + requested_splits=None, +): + """Preserve inactive output and never consume stale split scratch.""" + prims.griddepcontrol(kind=prims.GridDepAction.WAIT) + _, _, batch_idx = cute.arch.block_idx() + if cache_seqs[batch_idx] > Int32(0): + _run_active_gmem_reduction_kernel( + kernel, + output, + lse, + acc_output, + acc_lse, + cache_seqs, + cu_seqlens_q, + cfg, + num_reduction_ctas, + requested_splits, + ) + + +@cute.jit +def _run_active_gmem_reduction_kernel( + kernel, + output, + lse, + acc_output, + acc_lse, + cache_seqs, + cu_seqlens_q, + cfg, + num_reduction_ctas, + requested_splits=None, ): """Reduce split-KV partial O/LSE rows written by the 1CTA main kernel. @@ -251,8 +287,12 @@ def run_gmem_reduction_kernel( public storage; inactive tail rows may synchronize and consume workspace slots but never publish O or LSE. """ - # Pair with the dense kernel PDL signal before reading partials. - prims.griddepcontrol(kind=prims.GridDepAction.WAIT) + + split_limit = ( + Int32(cfg.num_ctas_per_seq_kv) + if requested_splits is None + else runtime_split_count(requested_splits) + ) q_idx, head_dim_tile_idx, batch_idx = cute.arch.block_idx() tidx, _, _ = cute.arch.thread_idx() warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) @@ -311,9 +351,21 @@ def run_gmem_reduction_kernel( space=cutlass.AddressSpace.rmem, ) lse_max = kernel.lse_dtype(-kernel.lse_dtype.inf) - for split_idx in cutlass.range_constexpr(cfg.num_ctas_per_seq_kv): - local_lse[split_idx] = row_lse[split_idx] - lse_max = fmax_f32(lse_max, local_lse[split_idx]) + for split_idx in cutlass.range( + split_limit, unroll_full=requested_splits is None + ): + if Int32(split_idx) < split_limit: + local_lse[split_idx] = ( + (row_lse[split_idx]) + if Int32(split_idx) + < ( + Int32(cfg.num_ctas_per_seq_kv) + if requested_splits is None + else runtime_split_count(requested_splits) + ) + else Float32(-Float32.inf) + ) + lse_max = fmax_f32(lse_max, local_lse[split_idx]) lse_max = ( lse_max @@ -321,11 +373,14 @@ def run_gmem_reduction_kernel( else kernel.lse_dtype(0.0) ) sum_lse = kernel.lse_dtype(0.0) - for split_idx in cutlass.range_constexpr(cfg.num_ctas_per_seq_kv): - sum_lse += cute.math.exp2( - local_lse[split_idx] - lse_max, - fastmath=True, - ) + for split_idx in cutlass.range( + split_limit, unroll_full=requested_splits is None + ): + if Int32(split_idx) < split_limit: + sum_lse += cute.math.exp2( + local_lse[split_idx] - lse_max, + fastmath=True, + ) has_finite_mass = sum_lse == sum_lse and sum_lse != kernel.lse_dtype( 0.0 ) @@ -341,7 +396,10 @@ def run_gmem_reduction_kernel( batch_idx, cu_seqlens_q, ) - (lse.iterator.raw_ptr() + output_query_row).store(global_lse) + if cutlass.const_expr(lse is not None): + (lse.iterator.raw_ptr() + output_query_row).store( + final_lse(global_lse, cfg.lse_base_on_e) + ) acc_vec = vector_from_scalars( ( @@ -356,32 +414,35 @@ def run_gmem_reduction_kernel( ), dtype=Float32, ) - for split_idx in cutlass.range_constexpr(cfg.num_ctas_per_seq_kv): - scale = Float32( - cute.math.exp2( - local_lse[split_idx] - global_lse, - fastmath=True, + for split_idx in cutlass.range( + split_limit, unroll_full=requested_splits is None + ): + if Int32(split_idx) < split_limit: + scale = Float32( + cute.math.exp2( + local_lse[split_idx] - global_lse, + fastmath=True, + ) + if has_finite_mass + else kernel.acc_dtype(0.0) ) - if has_finite_mass - else kernel.acc_dtype(0.0) - ) - acc_elem_offset = split_o_element_offset( - cfg, - batch_idx, - cta_idx_q, - head_idx, - Int32(split_idx), - dim_idx, - ) - partial_vec = ( - (acc_ptr + acc_elem_offset) - .load( - count=GMEM_REDUCTION_ELEMENTS_PER_THREAD, - alignment=GMEM_REDUCTION_VECTOR_BYTES, + acc_elem_offset = split_o_element_offset( + cfg, + batch_idx, + cta_idx_q, + head_idx, + Int32(split_idx), + dim_idx, ) - .to(Float32) - ) - acc_vec = acc_vec + partial_vec * scale + partial_vec = ( + (acc_ptr + acc_elem_offset) + .load( + count=GMEM_REDUCTION_ELEMENTS_PER_THREAD, + alignment=GMEM_REDUCTION_VECTOR_BYTES, + ) + .to(Float32) + ) + acc_vec = acc_vec + partial_vec * scale output_query_row = public_query_flat_row( cfg, @@ -453,12 +514,24 @@ def run_gmem_reduction_kernel( space=cutlass.AddressSpace.rmem, ) lse_max = kernel.lse_dtype(-kernel.lse_dtype.inf) - for i in cutlass.range_constexpr(lse_per_thread): + for i in cutlass.range( + ceil_div(split_limit, GMEM_REDUCTION_WARP_LANES), + unroll_full=requested_splits is None, + ): split_idx = lane_idx + Int32(i * GMEM_REDUCTION_WARP_LANES) local_lse[i] = ( - row_lse[split_idx] - if cute.elem_less(split_idx, Int32(cfg.num_ctas_per_seq_kv)) - else -kernel.lse_dtype.inf + ( + row_lse[split_idx] + if cute.elem_less(split_idx, Int32(cfg.num_ctas_per_seq_kv)) + else -kernel.lse_dtype.inf + ) + if Int32(split_idx) + < ( + Int32(cfg.num_ctas_per_seq_kv) + if requested_splits is None + else runtime_split_count(requested_splits) + ) + else Float32(-Float32.inf) ) lse_max = fmax_f32(lse_max, local_lse[i]) lse_max = warp_reduce_max_f32(lse_max) @@ -468,7 +541,10 @@ def run_gmem_reduction_kernel( else kernel.lse_dtype(0.0) ) sum_lse = kernel.lse_dtype(0.0) - for i in cutlass.range_constexpr(lse_per_thread): + for i in cutlass.range( + ceil_div(split_limit, GMEM_REDUCTION_WARP_LANES), + unroll_full=requested_splits is None, + ): sum_lse += cute.math.exp2(local_lse[i] - lse_max, fastmath=True) sum_lse = warp_reduce_sum_f32(sum_lse) has_finite_mass = sum_lse == sum_lse and sum_lse != kernel.lse_dtype( @@ -486,8 +562,14 @@ def run_gmem_reduction_kernel( batch_idx, cu_seqlens_q, ) - (lse.iterator.raw_ptr() + output_query_row).store(global_lse) - for i in cutlass.range_constexpr(lse_per_thread): + if cutlass.const_expr(lse is not None): + (lse.iterator.raw_ptr() + output_query_row).store( + final_lse(global_lse, cfg.lse_base_on_e) + ) + for i in cutlass.range( + ceil_div(split_limit, GMEM_REDUCTION_WARP_LANES), + unroll_full=requested_splits is None, + ): split_idx = lane_idx + Int32(i * GMEM_REDUCTION_WARP_LANES) if cute.elem_less(split_idx, Int32(cfg.num_ctas_per_seq_kv)): smem_lse_scale[ @@ -541,30 +623,33 @@ def run_gmem_reduction_kernel( ), dtype=Float32, ) - for split_idx in cutlass.range_constexpr(cfg.num_ctas_per_seq_kv): - scale = Float32( - smem_lse_scale[ - local_head_idx * Int32(cfg.num_ctas_per_seq_kv) - + Int32(split_idx) - ] - ) - acc_elem_offset = split_o_element_offset( - cfg, - batch_idx, - q_idx, - global_head_idx, - Int32(split_idx), - dim_idx, - ) - partial_vec = ( - (acc_ptr + acc_elem_offset) - .load( - count=GMEM_REDUCTION_ELEMENTS_PER_THREAD, - alignment=16, + for split_idx in cutlass.range( + split_limit, unroll_full=requested_splits is None + ): + if Int32(split_idx) < split_limit: + scale = Float32( + smem_lse_scale[ + local_head_idx * Int32(cfg.num_ctas_per_seq_kv) + + Int32(split_idx) + ] ) - .to(Float32) - ) - acc_vec = acc_vec + partial_vec * scale + acc_elem_offset = split_o_element_offset( + cfg, + batch_idx, + q_idx, + global_head_idx, + Int32(split_idx), + dim_idx, + ) + partial_vec = ( + (acc_ptr + acc_elem_offset) + .load( + count=GMEM_REDUCTION_ELEMENTS_PER_THREAD, + alignment=16, + ) + .to(Float32) + ) + acc_vec = acc_vec + partial_vec * scale output_query_row = public_query_flat_row( cfg, storage_flat_query_row, @@ -620,12 +705,24 @@ def run_gmem_reduction_kernel( space=cutlass.AddressSpace.rmem, ) lse_max = kernel.lse_dtype(-kernel.lse_dtype.inf) - for i in cutlass.range_constexpr(lse_per_thread): + for i in cutlass.range( + ceil_div(split_limit, GMEM_REDUCTION_WARP_LANES), + unroll_full=requested_splits is None, + ): split_idx = lane_idx + Int32(i * GMEM_REDUCTION_WARP_LANES) local_lse[i] = ( - row_lse[split_idx] - if cute.elem_less(split_idx, Int32(cfg.num_ctas_per_seq_kv)) - else -kernel.lse_dtype.inf + ( + row_lse[split_idx] + if cute.elem_less(split_idx, Int32(cfg.num_ctas_per_seq_kv)) + else -kernel.lse_dtype.inf + ) + if Int32(split_idx) + < ( + Int32(cfg.num_ctas_per_seq_kv) + if requested_splits is None + else runtime_split_count(requested_splits) + ) + else Float32(-Float32.inf) ) lse_max = fmax_f32(lse_max, local_lse[i]) lse_max = warp_reduce_max_f32(lse_max) @@ -635,7 +732,10 @@ def run_gmem_reduction_kernel( else kernel.lse_dtype(0.0) ) sum_lse = kernel.lse_dtype(0.0) - for i in cutlass.range_constexpr(lse_per_thread): + for i in cutlass.range( + ceil_div(split_limit, GMEM_REDUCTION_WARP_LANES), + unroll_full=requested_splits is None, + ): sum_lse += cute.math.exp2(local_lse[i] - lse_max, fastmath=True) sum_lse = warp_reduce_sum_f32(sum_lse) has_finite_mass = sum_lse == sum_lse and sum_lse != kernel.lse_dtype( @@ -653,8 +753,14 @@ def run_gmem_reduction_kernel( batch_idx, cu_seqlens_q, ) - (lse.iterator.raw_ptr() + output_query_row).store(global_lse) - for i in cutlass.range_constexpr(lse_per_thread): + if cutlass.const_expr(lse is not None): + (lse.iterator.raw_ptr() + output_query_row).store( + final_lse(global_lse, cfg.lse_base_on_e) + ) + for i in cutlass.range( + ceil_div(split_limit, GMEM_REDUCTION_WARP_LANES), + unroll_full=requested_splits is None, + ): split_idx = lane_idx + Int32(i * GMEM_REDUCTION_WARP_LANES) if cute.elem_less(split_idx, Int32(cfg.num_ctas_per_seq_kv)): smem_lse_scale[ @@ -704,29 +810,32 @@ def run_gmem_reduction_kernel( dtype=Float32, ) acc_ptr = acc_output.iterator.raw_ptr() - for split_idx in cutlass.range_constexpr(cfg.num_ctas_per_seq_kv): - scale = Float32( - smem_lse_scale[ - row_in_slice * Int32(cfg.num_ctas_per_seq_kv) + split_idx - ] - ) - acc_elem_offset = split_o_element_offset( - cfg, - batch_idx, - q_idx, - head_idx, - Int32(split_idx), - dim_idx, - ) - partial_vec = ( - (acc_ptr + acc_elem_offset) - .load( - count=GMEM_REDUCTION_ELEMENTS_PER_THREAD, - alignment=16, + for split_idx in cutlass.range( + split_limit, unroll_full=requested_splits is None + ): + if Int32(split_idx) < split_limit: + scale = Float32( + smem_lse_scale[ + row_in_slice * Int32(cfg.num_ctas_per_seq_kv) + split_idx + ] ) - .to(Float32) - ) - acc_vec = acc_vec + partial_vec * scale + acc_elem_offset = split_o_element_offset( + cfg, + batch_idx, + q_idx, + head_idx, + Int32(split_idx), + dim_idx, + ) + partial_vec = ( + (acc_ptr + acc_elem_offset) + .load( + count=GMEM_REDUCTION_ELEMENTS_PER_THREAD, + alignment=16, + ) + .to(Float32) + ) + acc_vec = acc_vec + partial_vec * scale output_query_row = public_query_flat_row( cfg, storage_flat_query_row, @@ -781,19 +890,34 @@ def run_gmem_reduction_kernel( space=cutlass.AddressSpace.rmem, ) lse_max = kernel.lse_dtype(-kernel.lse_dtype.inf) - for i in cutlass.range_constexpr(lse_per_thread): + for i in cutlass.range( + ceil_div(split_limit, GMEM_REDUCTION_WARP_LANES), + unroll_full=requested_splits is None, + ): split_idx = lane_idx + Int32(i * GMEM_REDUCTION_WARP_LANES) local_lse[i] = ( - g_lse[split_idx] - if valid_output_row - and cute.elem_less(split_idx, Int32(cfg.num_ctas_per_seq_kv)) - else -kernel.lse_dtype.inf + ( + g_lse[split_idx] + if valid_output_row + and cute.elem_less(split_idx, Int32(cfg.num_ctas_per_seq_kv)) + else -kernel.lse_dtype.inf + ) + if Int32(split_idx) + < ( + Int32(cfg.num_ctas_per_seq_kv) + if requested_splits is None + else runtime_split_count(requested_splits) + ) + else Float32(-Float32.inf) ) lse_max = fmax_f32(lse_max, local_lse[i]) lse_max = warp_reduce_max_f32(lse_max) lse_max = lse_max if lse_max != -kernel.lse_dtype.inf else kernel.lse_dtype(0.0) sum_lse = kernel.lse_dtype(0.0) - for i in cutlass.range_constexpr(lse_per_thread): + for i in cutlass.range( + ceil_div(split_limit, GMEM_REDUCTION_WARP_LANES), + unroll_full=requested_splits is None, + ): sum_lse += cute.math.exp2(local_lse[i] - lse_max, fastmath=True) sum_lse = warp_reduce_sum_f32(sum_lse) has_finite_mass = sum_lse == sum_lse and sum_lse != kernel.lse_dtype(0.0) @@ -810,8 +934,14 @@ def run_gmem_reduction_kernel( batch_idx, cu_seqlens_q, ) - (lse.iterator.raw_ptr() + output_query_row).store(global_lse) - for i in cutlass.range_constexpr(lse_per_thread): + if cutlass.const_expr(lse is not None): + (lse.iterator.raw_ptr() + output_query_row).store( + final_lse(global_lse, cfg.lse_base_on_e) + ) + for i in cutlass.range( + ceil_div(split_limit, GMEM_REDUCTION_WARP_LANES), + unroll_full=requested_splits is None, + ): split_idx = lane_idx + Int32(i * GMEM_REDUCTION_WARP_LANES) if cute.elem_less(split_idx, Int32(cfg.num_ctas_per_seq_kv)): smem_lse_scale[split_idx] = ( @@ -832,10 +962,15 @@ def run_gmem_reduction_kernel( r_o = cutlass.Array(out_element_dtype, 1, space=cutlass.AddressSpace.rmem) r_acc_o[0] = kernel.acc_dtype(0.0) if valid_output_row: - for split_idx in cutlass.range_constexpr(cfg.num_ctas_per_seq_kv): - scale = Float32(smem_lse_scale[split_idx]) - if dim_idx < Int32(cfg.head_dim_v): - r_acc_o[0] = r_acc_o[0] + Float32(g_acc_o[split_idx, dim_idx]) * scale + for split_idx in cutlass.range( + split_limit, unroll_full=requested_splits is None + ): + if Int32(split_idx) < split_limit: + scale = Float32(smem_lse_scale[split_idx]) + if dim_idx < Int32(cfg.head_dim_v): + r_acc_o[0] = ( + r_acc_o[0] + Float32(g_acc_o[split_idx, dim_idx]) * scale + ) r_o.store(r_acc_o.load(0, 1).to(out_element_dtype), 0) if dim_idx < Int32(cfg.head_dim_v) and valid_output_row: output_query_row = public_query_flat_row( diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/common.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/common.py index b8a7aa203e3a..2d95f3877b83 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/common.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/common.py @@ -67,6 +67,7 @@ class MlaResource(MemoryResource): cfg: cutlass.Constexpr[MlaConfig] = None cu_seqlens_q: object = None + runtime_split_kv: object = None _alloc: cutlass.Constexpr[Optional[SmemAllocation]] = None _tmem_alloc: cutlass.Constexpr[Optional[TmemAllocation]] = None _tmem_base_addr: object = None diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_p.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_p.py index 14fd5758621b..5c78358858ab 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_p.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_p.py @@ -56,7 +56,7 @@ qkv_dtype, ) from ...helpers.ops import ( - fp8_log2_quant_scale, + fp8_log2_p_scale, pack_float4_to_fp8_e4m3, store_transposed_smem8b_x2, store_transposed_smem8b_x4, @@ -248,7 +248,7 @@ def materialize_p( safe_new_max = Float32(0.0) neg_scaled_max[idx] = -self.scale_softmax_log2 * safe_new_max if cutlass.const_expr(cfg.is_fp8_qkv()): - neg_scaled_max[idx] += fp8_log2_quant_scale() + neg_scaled_max[idx] += fp8_log2_p_scale(cfg) local_sums[idx] = Float32(0.0) # Convert S registers to P and accumulate local softmax sums. FP8 packs diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_resources.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_resources.py index 590173465f44..2ee2d9474c21 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_resources.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/smem_resources.py @@ -27,7 +27,7 @@ from cutlass.pipeline import PipelineAsync, PipelineState from cutlass.experimental import primitives as cprims from cutlass.experimental.task_scheduling.enums import WorkAttr -from cutlass.experimental.task_scheduling.memory import SmemAllocation +from cutlass.experimental.task_scheduling.memory import SmemAllocation, TmemAllocation from cutlass.experimental.task_scheduling.resources import ( StageInfo, TaskLocalVariable, @@ -35,7 +35,22 @@ producer_work, ) -from ...helpers.gather import gather4, gather4_cached_warp, select_gather_map +from ...helpers.sparse_indices import cache_sparse_source_lengths +from ...helpers.nvfp4_kv import ( + publish_nvfp4_kv, + source_is_fp8, + load_nvfp4_panel, + prefetch_nvfp4_kv_row, + copy_fp8_record, + load_nvfp4_panels, + issue_whole_record, + convert_nvfp4_record, +) +from ...helpers.gather import ( + gather4, + gather4_cached_warp, + select_gather_map, +) from ...helpers.constants import ( CP_ASYNC_CACHE_CA, @@ -45,6 +60,7 @@ from ...helpers.layout import ( decode_gen_task_cache, + _TASK_CACHE_TMEM_BASE_OFFSET, head_dim_cta_offset_v, q_stage_smem_element_offset, smem_array, @@ -59,7 +75,7 @@ from ...helpers.ops import ( lane_idx_from_thread, ) -from ...helpers.query import query_batch_bounds +from ...helpers.query import query_batch_bounds, copy_q_fp8_to_tmem from ...helpers.stage import MlaStage from ...helpers.tile import ( batch_idx_for_stage_cfg, @@ -105,6 +121,10 @@ class SmemQResource(MlaResource): "q_desc_rope", ), ) + q_temp_alloc: object = None + q_temp_state: object = None + kv_alias_ref: object = None + q_tensor: object = None tma_desc_q_latent: object = None tma_desc_q_rope: object = None head_idx: object = None @@ -117,22 +137,45 @@ class SmemQResource(MlaResource): TaskLocalVariable.uninitialized() ) + def get_tmem_requirements(self): + """Reserve NoPE query columns when WS MMA reads Q from TMEM.""" + if not self.cfg.q_nope_tmem: + return [] + if self._tmem_alloc is None: + self._tmem_alloc = TmemAllocation( + name=f"{self.name}_nope_tmem", num_columns=128 + ) + return [self._tmem_alloc] + def get_smem_requirements(self): """Return the SMEM allocation used for staged Q tiles.""" if self._alloc is None: self._alloc = SmemAllocation( name=f"{self.name}", - size_bytes=self.cfg.qk_smem_tile_bytes * self.cfg.q_stages, + size_bytes=( + 4096 if self.cfg.q_nope_tmem else self.cfg.qk_smem_tile_bytes + ) + * self.cfg.q_stages, alignment=self.cfg.stensor_align, ) - return [self._alloc] + allocations = [self._alloc] + if self.cfg.q_tma_stage: + if self.q_temp_alloc is None: + self.q_temp_alloc = SmemAllocation( + name=f"{self.name}_q_temp_state", size_bytes=32, alignment=8 + ) + allocations.append(self.q_temp_alloc) + return allocations @cute.jit def _init_smem_state(self, stage_info: StageInfo) -> None: """Create the Q SMEM view and base descriptor state.""" context = stage_info.context + if cutlass.const_expr(self.cfg.q_tma_stage): + self.q_temp_state = smem_array(context, self.q_temp_alloc, Int64, 4) elem_count = ( - self.cfg.qk_smem_tile_bytes * self.cfg.q_stages + (4096 if self.cfg.q_nope_tmem else self.cfg.qk_smem_tile_bytes) + * self.cfg.q_stages ) // self.cfg.qkv_dtype_bytes self._smem_q = smem_array(context, self._alloc, qkv_dtype(self.cfg), elem_count) if self._smem_q is not None: @@ -171,6 +214,14 @@ def init_load_state(self, stage_info: StageInfo) -> None: # Producer aux work runs before the load task starts issuing Q TMA # copies. It creates the SMEM view used by producer_work(load_q). self._init_smem_state(stage_info) + if cutlass.const_expr(self.cfg.q_tma_stage): + relative = cute.arch.thread_idx()[0] - Int32(self.cfg.load_warp_idx * 32) + if relative == Int32(0): + prims.mbarrier_init(self.q_temp_state.data_ptr(), 1) + prims.mbarrier_init(self.q_temp_state.data_ptr(2), 1) + self.q_temp_state[3] = Int64(0) + cute.arch.mbarrier_init_fence() + prims.barrier_cta_sync(barrier_id=12, thread_count=256) @consumer_work(work_attrs=WorkAttr.AUXILIARY) @cute.jit @@ -238,6 +289,8 @@ def _load_q_stage( == qk_stage_idx % self.cfg.load_num_warps ): smem_offset = Int32(q_stage_smem_element_offset(self.cfg, qk_stage_idx)) + if cutlass.const_expr(self.cfg.q_nope_tmem): + smem_offset = Int32(0) prims.cp_async_bulk_tensor_shared_cta_global( stage_base.data_ptr(smem_offset), tma_desc, @@ -259,6 +312,147 @@ def _load_q_stage( stage_info.barrier, ) + @cute.jit + def copy_nope_to_tmem(self, stage_info): + """Copy Q NoPE through registers from GMEM to TMEM without a TMA scratch stage. + + Four warps write repeated M64 rows for the two WS column halves. This + is the direct-copy alternative to stage_q_via_tma, which uses eight + load warps and temporarily borrows an empty KV slot for TMA staging.""" + cfg = self.cfg + assert ( + cfg.tile_size_q == 64 + and cfg.latent_dim == 512 + and cfg.is_fp8_qkv() + and cfg.q_stages == 1 + ) + relative = cute.arch.thread_idx()[0] - Int32(cfg.load_warp_idx * 32) + # One warpgroup covers128 DP rows; repeat logical64 Q rows for the + # two WS column halves. The other load warps issue the RoPE TMA. + if relative < Int32(128): + row = relative % Int32(64) + batch = batch_idx_for_stage_cfg(self.batch_idx, cfg, stage_info) + local_flat = ( + cta_idx_q_for_stage(self.cta_idx_q, stage_info) * Int32(cfg.num_heads_q) + + head_idx_for_stage(self.head_idx, cfg, stage_info) + + row + ) + storage_flat = local_flat + source_batch = batch + extent = Int32(self.q_tensor.shape[0] * self.q_tensor.shape[2]) + if cutlass.const_expr(self.cu_seqlens_q is not None): + start, length = query_batch_bounds( + self.cu_seqlens_q, batch, cfg.logical_seq_len_q + ) + storage_flat = start * Int32(cfg.logical_num_heads_q) + local_flat + extent = length * Int32(cfg.logical_num_heads_q) + source_batch = Int32(0) + address = ( + self.q_tensor.iterator.toint() + + Int64(source_batch) * Int64(self.q_tensor.stride[3]) + + Int64(storage_flat) * Int64(self.q_tensor.stride[0]) + ) + cache = decode_gen_task_cache(stage_info) + tbase = cache[_TASK_CACHE_TMEM_BASE_OFFSET] + Int32(self._tmem_alloc.offset) + for part in cutlass.range_constexpr(4): + words = cutlass.Array(Int32, 32, space=cutlass.AddressSpace.rmem) + for j in cutlass.range_constexpr(32): + words[j] = Int32(0) + if local_flat < extent: + ptr = cutlass.inttoptr( + address + Int64(part * 128), mem_space=1, dtype=Int32 + ) + loaded = ptr.load(count=32, alignment=16) + for j in cutlass.range_constexpr(32): + words[j] = loaded[j] + prims.tcgen05_st( + prims.Tcgen05LdStShape.SHAPE_32X32B, + prims.make_tmem_ptr(tbase + Int32(part * 32), Int32), + words.load(0, 32), + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + # Keep the existing36KiB Q transaction accounting: four warp + # completions cover32KiB logical NoPE, RoPE TMA covers4KiB. + publish_nvfp4_kv(stage_info.barrier, 8192) + + @cute.jit + def stage_q_via_tma(self, stage_info, rope_stage_base): + """Stage Q NoPE with TMA in an empty KV slot, then copy it from SMEM to TMEM. + + Unlike copy_nope_to_tmem, this path avoids wide per-thread GMEM loads. + It requires eight loaders and single-raw-prefetch geometry. Acquiring + scratch does not advance the KV pipeline cursor; Q finishes before + normal KV loading reuses the slot. D576 RoPE keeps its own Q stage.""" + cfg = self.cfg + assert cfg.q_nope_tmem and cfg.load_num_warps == 8 and cfg.single_raw_prefetch + relative = cute.arch.thread_idx()[0] - Int32(cfg.load_warp_idx * 32) + kv = self.kv_alias_ref + # Form a local scratch view. Reinitializing the KV resource here would + # retain descriptors defined inside the runtime-active Q branch and + # leak those staged values into the subsequent KV task loop. + kv_scratch = smem_array( + stage_info.context, + kv._alloc, + qkv_dtype(cfg), + cfg.whole_kv_tile_bytes * cfg.kv_stages // cfg.qkv_dtype_bytes, + ) + first = kv.producer_state.index + phase = kv.producer_state.phase + # This is only scratch reuse: no KV acquire/commit cursor is changed. + # The subsequent normal KV producer still acquires its own stages. + for i in cutlass.range_constexpr(1 if cfg.whole_kv_tile else 2): + kv.pipeline.sync_object_empty.wait(first + Int32(i), phase) + temp = kv_scratch.data_ptr(kv.kv_stage_offset(first)) + if relative == Int32(0): + token = prims.mbarrier_arrive_expect_tx(self.q_temp_state.data_ptr(), 32768) + self.q_temp_state[1] = token + prims.barrier_cta_sync(barrier_id=12, thread_count=256) + qpart = relative // Int32(32) + if qpart < Int32(4) and prims.elect_sync(): + head = head_idx_for_stage(self.head_idx, cfg, stage_info) + query = cta_idx_q_for_stage(self.cta_idx_q, stage_info) + batch = batch_idx_for_stage_cfg(self.batch_idx, cfg, stage_info) + coords = self._query_tma_coords( + qpart * Int32(128), query * Int32(cfg.num_heads_q) + head, batch + ) + prims.cp_async_bulk_tensor_shared_cta_global( + temp + qpart * Int32(8192), + self.tma_desc_q_latent, + coords, + self.q_temp_state.data_ptr(), + ) + if cutlass.const_expr(cfg.rope_dim): + self._load_q_stage( + stage_info, rope_stage_base, 4, Int32(0), self.tma_desc_q_rope + ) + token = self.q_temp_state[1] + while not prims.mbarrier_try_wait(self.q_temp_state.data_ptr(), token): + pass + if relative < Int32(32): + cache = decode_gen_task_cache(stage_info) + tbase = cache[_TASK_CACHE_TMEM_BASE_OFFSET] + Int32(self._tmem_alloc.offset) + copy_phase = Int32(self.q_temp_state[3]) + if prims.elect_sync(): + copy_q_fp8_to_tmem(temp, tbase, cta_group=1) + prims.tcgen05_commit( + self.q_temp_state.data_ptr(2), group=prims.CTAGroup.CTA_1 + ) + while not prims.mbarrier_try_wait_parity( + self.q_temp_state.data_ptr(2), copy_phase + ): + pass + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + cute.arch.sync_warp() + if prims.elect_sync(): + self.q_temp_state[3] = Int64(copy_phase ^ Int32(1)) + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + # One warp now publishes the whole32KiB logical NoPE payload, + # after its TCGEN copy commit has completed. + publish_nvfp4_kv(stage_info.barrier, 32768) + # Every Q staging read ends before any loader may expand KV over it. + prims.barrier_cta_sync(barrier_id=12, thread_count=256) + @producer_work @cute.jit def load_q(self, stage_info: StageInfo): @@ -273,15 +467,21 @@ def load_q(self, stage_info: StageInfo): stage_elems = cfg.q_smem_tile_elements stage_base = self._smem_q.subview(stage_info.stage_idx * Int32(stage_elems)) latent_stages = cfg.latent_dim // cfg.head_dim_per_stage_kv - for qk_stage_idx in cutlass.range_constexpr(latent_stages): - dim_offset = Int32(qk_stage_idx * cfg.head_dim_per_stage_kv) - self._load_q_stage( - stage_info, - stage_base, - qk_stage_idx, - dim_offset, - self.tma_desc_q_latent, - ) + if cutlass.const_expr(cfg.q_tma_stage): + self.stage_q_via_tma(stage_info, stage_base) + return + if cutlass.const_expr(cfg.q_nope_tmem): + self.copy_nope_to_tmem(stage_info) + else: + for qk_stage_idx in cutlass.range_constexpr(latent_stages): + dim_offset = Int32(qk_stage_idx * cfg.head_dim_per_stage_kv) + self._load_q_stage( + stage_info, + stage_base, + qk_stage_idx, + dim_offset, + self.tma_desc_q_latent, + ) if cutlass.const_expr(cfg.rope_dim > 0): self._load_q_stage( stage_info, @@ -330,6 +530,8 @@ def q_desc(self, stage_info: StageInfo): rope_stage_offset = Int32( q_stage_smem_element_offset(self.cfg, rope_stage_idx) ) + if cutlass.const_expr(self.cfg.q_nope_tmem): + rope_stage_offset = Int32(0) desc_q_rope = cprims.Tcgen05SmemDesc.build( stage_base.subview(rope_stage_offset), leading_byte_offset=Int32( @@ -340,6 +542,12 @@ def q_desc(self, stage_info: StageInfo): ), layout=qkv_smem_swizzle_for_head_dim(self.cfg, self.cfg.rope_dim), ) + if cutlass.const_expr(self.cfg.q_nope_tmem): + prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) + cache = decode_gen_task_cache(stage_info) + desc_q = Int64( + cache[_TASK_CACHE_TMEM_BASE_OFFSET] + Int32(self._tmem_alloc.offset) + ) return desc_q, desc_q_rope @@ -420,6 +628,8 @@ class SmemPageOffsetsResource(MlaResource): cta_idx_q: object = None cta_idx_kv: object = None _smem_page_offsets: object = None + nv_target: cutlass.Constexpr = None + request_lengths: object = None cached_page_ids: cutlass.Constexpr[TaskLocalVariable] = ( TaskLocalVariable.uninitialized() ) @@ -442,6 +652,8 @@ def create_pipeline(self, pipeline_config): def get_smem_requirements(self): """Return the SMEM allocation for cached page offsets.""" + if self.cfg.single_raw_prefetch: + assert self.cfg.page_offsets_entries_per_stage % 4 == 0 if self._alloc is None: self._alloc = SmemAllocation( name=f"{self.name}", @@ -474,6 +686,29 @@ def init_load_state(self, stage_info: StageInfo) -> None: # Producer aux work creates the page-offset SMEM view for the page load # warp before it issues cp.async copies. self._init_smem_state(stage_info) + if cutlass.const_expr(self.cfg.sparse_direct): + self.request_lengths = cutlass.Array( + Int32, 2, space=cutlass.AddressSpace.rmem + ) + if cutlass.const_expr(self.cfg.async_nvfp4_kv_gather): + self.nv_target._init_smem_state(stage_info) + prims.barrier_cta_sync( + barrier_id=13, + thread_count=32 + * (self.cfg.load_num_warps + self.cfg.page_offsets_num_warps), + ) + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def bind_sparse_request(self, stage_info: StageInfo): + """Cache live source lengths for this query before its KV/softmax loop.""" + if cutlass.const_expr(self.cfg.sparse_direct): + cache_sparse_source_lengths( + self.page_offsets, + batch_idx_for_stage_cfg(self.batch_idx, self.cfg, stage_info), + self.request_lengths, + include_primary_tiles=False, + ) @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=cached_page_ids) @cute.jit @@ -533,31 +768,106 @@ def _producer_load_page_offsets( batch_idx, ) local_tile_idx = local_kv_tile_idx( - cfg, stage_info, inst_id, is_v, section=section + cfg, + stage_info, + inst_id, + is_v, + section=section, + requested_splits=self.runtime_split_kv, + ) + tile_idx = global_kv_tile_idx( + cfg, + local_tile_idx, + seq_len_kv, + cta_idx_kv, + requested_splits=self.runtime_split_kv, + task_cache=decode_gen_task_cache(stage_info), ) - tile_idx = global_kv_tile_idx(cfg, local_tile_idx, seq_len_kv, cta_idx_kv) tile_start = tile_idx * pages_per_tile - routes = self.page_offsets.select_source_tile(tile_start, batch_idx) + rows = self.page_offsets.bind_query_row( + batch_idx, self.request_lengths[0], self.request_lengths[1] + ) + routes = rows.select_source_tile(tile_start) position = lane_idx * Int32(4) + if cutlass.const_expr(cfg.async_nvfp4_kv_gather): + issuer = cute.arch.thread_idx()[0] // 32 - cfg.page_offsets_warp_idx + position = issuer * Int32(64) + (lane_idx % Int32(16)) * Int32(4) smem_offset = ( stage_info.stage_idx * Int32(cfg.page_offsets_entries_per_stage) + position ) - if position < pages_per_tile: - # Source selection is tile-wide. Hoist the bounds decision out - # of individual quads so full tiles take one uniform branch; - # index values still mask holes even in an in-bounds tile. + raw = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) + for j in cutlass.range_constexpr(4): + raw[j] = Int32(0x7FFFFFFF) + if position < pages_per_tile and ( + not cfg.async_nvfp4_kv_gather or lane_idx < 16 + ): + # Fetch four physical storage-row IDs for this lane's source. + # The same quad fills the metadata ring and issues the raw TMA + # gather, avoiding another global index read in KV loaders. + # Full tiles can use unmasked reads; live-length tails and + # holes preserve the source tag on synthesized invalid rows. if tile_start + pages_per_tile <= routes.origin + routes.length: quad = routes.load_full_quad(tile_start + position) - self._smem_page_offsets.data_ptr(smem_offset).store( - quad.load(0, 4), alignment=16 - ) + for j in cutlass.range_constexpr(4): + raw[j] = quad[j] else: quad = routes.load_quad(tile_start + position) - self._smem_page_offsets.data_ptr(smem_offset).store( - quad.load(0, 4), alignment=16 - ) + for j in cutlass.range_constexpr(4): + raw[j] = quad[j] + self._smem_page_offsets.data_ptr(smem_offset).store( + raw.load(0, 4), alignment=16 + ) cute.arch.fence_view_async_shared() + if cutlass.const_expr(cfg.async_nvfp4_kv_gather): + target = self.nv_target + if lane_idx < Int32(16): + for j in cutlass.range_constexpr(4): + prefetch_nvfp4_kv_row( + target.c_latent_tensor, + target.c_rope_tensor, + Int32(raw[j]), + cfg, + ) + if cutlass.const_expr(cfg.single_raw_prefetch): + # PG has two stages; its index alternates once per tile, + # exactly matching the one raw-empty barrier's epochs. + while not prims.mbarrier_try_wait_parity( + target._nv_barriers.data_ptr(2), stage_info.stage_idx ^ Int32(1) + ): + pass + cute.arch.fence_view_async_shared() + destinations = (target._nv_raw.data_ptr(),) + if issuer == Int32(0) and prims.elect_sync(): + header = stage_info.stage_idx * Int32( + cfg.page_offsets_entries_per_stage + ) + Int32(128) + self._smem_page_offsets[header] = stage_info.stage_idx + self._smem_page_offsets[header + Int32(1)] = ( + self.producer_state.phase ^ Int32(1) + ) + else: + destinations = () + for part in cutlass.range_constexpr(5): + slot = stage_info.stage_idx * Int32(5) + Int32(part) + target.pipeline.sync_object_empty.wait( + slot, self.producer_state.phase + ) + destinations += ( + target._smem_kv.data_ptr(target.kv_stage_offset(slot)), + ) + record_bytes = Int32(cfg.nvfp4_row_bytes) + if source_is_fp8(cfg, raw[0] < Int32(0)): + record_bytes = Int32(512) + issue_whole_record( + destinations[0], + target._nv_barriers.data_ptr(stage_info.stage_idx), + raw, + target.tma_desc_c_latent, + target.tma_desc_c_rope, + issuer=issuer, + row_bytes=record_bytes, + ) return for fragment in cutlass.range_constexpr((cfg.pages_per_kv_tile + 31) // 32): fragment_lane = lane_idx + Int32(fragment * 32) @@ -573,10 +883,20 @@ def _producer_load_page_offsets( batch_idx, ) local_tile_idx = local_kv_tile_idx( - cfg, stage_info, inst_id, is_v, section=section + cfg, + stage_info, + inst_id, + is_v, + section=section, + requested_splits=self.runtime_split_kv, ) tile_idx = global_kv_tile_idx( - cfg, local_tile_idx, seq_len_kv, cta_idx_kv + cfg, + local_tile_idx, + seq_len_kv, + cta_idx_kv, + requested_splits=self.runtime_split_kv, + task_cache=decode_gen_task_cache(stage_info), ) last_valid_page = cute.math.max( (seq_len_kv + Int32(cfg.num_tokens_per_page - 1)) @@ -648,7 +968,28 @@ def read_offsets( """Cache staged page IDs in the selected K stream's register slot.""" # The load task waits on the page-offset stage, snapshots the page ids # into registers, and keeps swaps K0/K1 live until their delayed V use. - del stage_info + if cutlass.const_expr(self.cfg.sparse_d576_reload): + assert self.cfg.cache_uniform_sparse_quads + # Snapshot this warp's eight quads once per K / V tile as + # warp-uniform coordinates; every head-dimension stage of the + # tile then issues them from one elected lane without the + # per-lane R2UR waterfall. + quads = self.cfg.tile_size_kv // (4 * self.cfg.load_num_warps) + warp = cute.arch.make_warp_uniform(cute.arch.warp_idx()) - Int32( + self.cfg.load_warp_idx + ) + row_base = self.consumer_work_stage * Int32( + self.cfg.page_offsets_entries_per_stage + ) + warp * Int32(4 * quads) + for quad in cutlass.range_constexpr(quads): + rows = self._smem_page_offsets.data_ptr( + row_base + Int32(4 * quad) + ).load(count=4, alignment=16) + for j in cutlass.range_constexpr(4): + cached_page_ids[6 + quad * 4 + j] = cute.arch.make_warp_uniform( + Int32(rows[j]) + ) + return cached_page_ids if cutlass.const_expr(self.cfg.num_tokens_per_page == 1): lane = lane_idx_from_thread(cute.arch.thread_idx()[0]) warp = cute.arch.thread_idx()[0] // 32 - self.cfg.load_warp_idx @@ -733,6 +1074,16 @@ def read_offsets( # Constants avoid carrying the previous tile's uniform cache. for j in cutlass.range_constexpr(self.cfg.sparse_cache_words - 6): cached_page_ids[base + 6 + j] = Int32(0) + if cutlass.const_expr(self.cfg.single_raw_prefetch): + header = stage_info.stage_idx * Int32( + self.cfg.page_offsets_entries_per_stage + ) + Int32(128) + cached_page_ids[cache_slot * self.cfg.sparse_cache_words + 4] = ( + self._smem_page_offsets[header] + ) + cached_page_ids[cache_slot * self.cfg.sparse_cache_words + 5] = ( + self._smem_page_offsets[header + Int32(1)] + ) return cached_page_ids @@ -772,7 +1123,10 @@ class SmemKvResource(MlaResource): tma_desc_c_rope: object = None tma_desc_v: object = None tma_desc_bulk_extra: object = None + tma_desc_sparse_rope: object = None + tma_desc_sparse_extra_rope: object = None c_rope_tensor: object = None + c_latent_tensor: object = None page_offsets_kv: object = None page_offsets: object = None cache_seqs: object = None @@ -794,6 +1148,13 @@ class SmemKvResource(MlaResource): TaskLocalVariable.uninitialized() ) + _nv_raw_alloc: object = None + _nv_raw: object = None + _nv_scratch_alloc: object = None + _nv_barrier_alloc: object = None + _nv_scratch: object = None + _nv_barriers: object = None + def get_smem_requirements(self): """Return the SMEM allocation used for staged K/V tiles.""" num_stages = ( @@ -804,10 +1165,42 @@ def get_smem_requirements(self): if self._alloc is None: self._alloc = SmemAllocation( name=f"{self.name}", - size_bytes=self.cfg.kv_smem_tile_bytes * num_stages, + size_bytes=( + self.cfg.whole_kv_tile_bytes * num_stages + if self.cfg.whole_kv_tile + else self.cfg.kv_smem_tile_bytes * num_stages + - (16384 if self.cfg.compact_nv_rope else 0) + ), alignment=self.cfg.stensor_align, ) - return [self._alloc] + allocations = [self._alloc] + if self.cfg.has_nvfp4_kv: + # Whole-record gather carries SF and RoPE inside the raw record; + # it never reads the older separate scale scratch allocation. + if not self.cfg.async_nvfp4_kv_gather: + if self._nv_scratch_alloc is None: + self._nv_scratch_alloc = SmemAllocation( + name=f"{self.name}_nv_scratch", + size_bytes=self.cfg.tile_size_kv * 64, + alignment=1024, + ) + allocations.append(self._nv_scratch_alloc) + if self._nv_barrier_alloc is None: + self._nv_barrier_alloc = SmemAllocation( + name=f"{self.name}_nv_barriers", + size_bytes=self.cfg.load_num_warps * 8, + alignment=8, + ) + allocations.append(self._nv_barrier_alloc) + if self.cfg.single_raw_prefetch: + if self._nv_raw_alloc is None: + self._nv_raw_alloc = SmemAllocation( + name=f"{self.name}_raw_record", + size_bytes=128 * self.cfg.raw_record_stride, + alignment=1024, + ) + allocations.append(self._nv_raw_alloc) + return allocations @cute.jit def _init_smem_state(self, stage_info: StageInfo) -> None: @@ -819,11 +1212,31 @@ def _init_smem_state(self, stage_info: StageInfo) -> None: else self.cfg.kv_stages ) elem_count = ( - self.cfg.kv_smem_tile_bytes * num_stages + self.cfg.whole_kv_tile_bytes * num_stages + if self.cfg.whole_kv_tile + else self.cfg.kv_smem_tile_bytes * num_stages + - (16384 if self.cfg.compact_nv_rope else 0) ) // self.cfg.qkv_dtype_bytes self._smem_kv = smem_array( context, self._alloc, qkv_dtype(self.cfg), elem_count ) + if cutlass.const_expr(self.cfg.single_raw_prefetch): + self._nv_raw = smem_array( + context, + self._nv_raw_alloc, + cutlass.Uint8, + 128 * self.cfg.raw_record_stride, + ) + if cutlass.const_expr(self.cfg.has_nvfp4_kv): + self._nv_scratch = smem_array( + context, + self._nv_scratch_alloc, + cutlass.Uint8, + self.cfg.tile_size_kv * 64, + ) + self._nv_barriers = smem_array( + context, self._nv_barrier_alloc, Int64, self.cfg.load_num_warps + ) if self._smem_kv is not None: # FP8 descriptors use dtype-specific major-K strides rather than # the BF16 128B-row swizzle group. @@ -862,6 +1275,28 @@ def init_load_state(self, stage_info: StageInfo) -> None: # Producer aux work initializes K/V SMEM views for the load task. self._init_smem_state(stage_info) + if cutlass.const_expr(self.cfg.has_nvfp4_kv): + warp = cute.arch.thread_idx()[0] // 32 - self.cfg.load_warp_idx + if prims.elect_sync(): + count = Int32( + self.cfg.page_offsets_num_warps + if self.cfg.async_nvfp4_kv_gather + else 1 + ) + if cutlass.const_expr(self.cfg.single_raw_prefetch): + if warp == Int32(2): + count = Int32(self.cfg.load_num_warps) + prims.mbarrier_init(self._nv_barriers.data_ptr(warp), count) + cute.arch.mbarrier_init_fence() + cute.arch.sync_warp() + if cutlass.const_expr(self.cfg.async_nvfp4_kv_gather): + # Raw barriers are shared by the gather and transform warps. + # IDs 14/15 belong to TMEM; 13 is reserved for this load group. + prims.barrier_cta_sync( + barrier_id=13, + thread_count=32 + * (self.cfg.load_num_warps + self.cfg.page_offsets_num_warps), + ) @consumer_work(work_attrs=WorkAttr.AUXILIARY) @cute.jit @@ -871,12 +1306,28 @@ def init_descriptor_state(self, stage_info: StageInfo) -> None: # Consumer aux work initializes K/V SMEM views for the MMA task. self._init_smem_state(stage_info) + @cute.jit + def kv_stage_offset(self, stage): + """Return the SMEM element offset for a payload slot, accounting for compact RoPE storage.""" + if cutlass.const_expr(self.cfg.whole_kv_tile): + return stage * Int32(self.cfg.whole_kv_tile_bytes) + offset = stage * Int32(self.cfg.kv_smem_stage_elements) + if cutlass.const_expr(self.cfg.compact_nv_rope): + assert ( + self.cfg.kv_stages == 10 + and self.cfg.tile_size_kv == 128 + and self.cfg.rope_dim == 64 + ) + # Each five-stage group has four16KiB NoPE panels and one8KiB + # RoPE panel. The second group's base is therefore72KiB. + offset -= Int32(stage >= Int32(5)) * Int32(8192) + return offset + @cute.jit def _stage_base(self, stage_info: StageInfo): """Return the base of the current K/V SMEM pipeline stage.""" - stage_elems = self.cfg.kv_smem_stage_elements - return self._smem_kv.subview(stage_info.stage_idx * Int32(stage_elems)) + return self._smem_kv.subview(self.kv_stage_offset(stage_info.stage_idx)) @cute.jit def _producer_copy_rope_stage( @@ -909,6 +1360,159 @@ def _producer_copy_rope_stage( stage_info.barrier, Int32(cfg.kv_smem_tile_bytes) ) + @cute.jit + def load_nvfp4_stage_payload( + self, + stage_info, + stage_base, + cached_page_ids, + dim_offset, + warp, + *, + page_id_slot: cutlass.Constexpr[int], + qk_stage_idx: cutlass.Constexpr[int], + is_v: cutlass.Constexpr[int], + active_width: cutlass.Constexpr[int], + ): + """Fill a KV stage using whole-record reuse or the shared panel loader. + + The caller owns the stage's full-barrier publication. Whole-record + mode waits for metadata-issued TMA and expands/copies only the first + feature stage; later stages reuse that payload through PV. Panel mode + uses the same conversion helpers for separately staged K or V slices. + """ + cfg = self.cfg + rows_per_warp = cfg.tile_size_kv // cfg.load_num_warps + if cutlass.const_expr(cfg.async_nvfp4_kv_gather): + if cutlass.const_expr(qk_stage_idx == 0): + raw_slot = stage_info.stage_idx // Int32(5) + raw_phase = self.producer_state.phase ^ Int32(1) + if cutlass.const_expr(cfg.single_raw_prefetch): + raw_slot = cached_page_ids[ + page_id_slot * cfg.sparse_cache_words + 4 + ] + raw_phase = cached_page_ids[ + page_id_slot * cfg.sparse_cache_words + 5 + ] + # Normal acquire already claimed the first stage. + # Retain all other expanded panels until old PV ends. + if cutlass.const_expr(not cfg.whole_kv_tile): + for part in cutlass.range_constexpr(1, 5): + self.pipeline.sync_object_empty.wait( + stage_info.stage_idx + Int32(part), + self.producer_state.phase, + ) + raw_barrier = self._nv_barriers.data_ptr(raw_slot) + destinations = () + for part in cutlass.range_constexpr(5): + slot = stage_info.stage_idx + Int32(part) + destinations += ( + self._smem_kv.data_ptr( + self.kv_stage_offset(stage_info.stage_idx) + + Int32(part * 16384) + if cfg.whole_kv_tile + else self.kv_stage_offset(slot) + ), + ) + # The metadata warp issues gathers independently, so + # all eight original load warps can convert in parallel. + while not prims.mbarrier_try_wait_parity(raw_barrier, raw_phase): + pass + if source_is_fp8( + cfg, + cached_page_ids[page_id_slot * cfg.sparse_cache_words] < Int32(0), + ): + copy_fp8_record( + destinations[0], + self._nv_raw.data_ptr(), + cute.arch.thread_idx()[0] - Int32(cfg.load_warp_idx * 32), + self._nv_barriers.data_ptr(2), + ) + elif cutlass.const_expr(cfg.single_raw_prefetch): + convert_nvfp4_record( + destinations[0], + cute.arch.thread_idx()[0] - Int32(cfg.load_warp_idx * 32), + stage_bytes=cfg.kv_smem_tile_bytes, + rope_dim=cfg.rope_dim, + raw=self._nv_raw.data_ptr(), + raw_empty=self._nv_barriers.data_ptr(2), + ) + else: + convert_nvfp4_record( + destinations[0], + cute.arch.thread_idx()[0] - Int32(cfg.load_warp_idx * 32), + stage_bytes=cfg.kv_smem_tile_bytes, + rope_dim=cfg.rope_dim, + ) + return + if cutlass.const_expr(cfg.sparse_reuse_kv and not is_v): + if cutlass.const_expr(qk_stage_idx == 0): + # Retained-K schedules have room for two complete K + # tiles. Claim each payload's existing empty barrier + # before batching its gathers. The captured schedule + # still arms/publishes each full barrier individually. + destinations = () + for part in cutlass.range_constexpr(cfg.qk_head_dim_stages): + future = stage_info.stage_idx + part + slot = future - cfg.kv_stages if future >= cfg.kv_stages else future + if cutlass.const_expr(part > 0): + phase = self.producer_state.phase ^ Int32( + future >= cfg.kv_stages + ) + self.pipeline.sync_object_empty.wait(slot, phase) + width = 64 if part == 4 else 128 + destinations += ( + self._smem_kv.data_ptr( + self.kv_stage_offset(slot) + warp * rows_per_warp * width + ), + ) + load_nvfp4_panels( + destinations, + (Int32(0), Int32(128), Int32(256), Int32(384)), + self._nv_scratch.data_ptr(warp * rows_per_warp * 64), + self._nv_barriers.data_ptr(warp), + cached_page_ids, + page_id_slot * cfg.sparse_cache_words, + self.tma_desc_c_latent, + self.tma_desc_c_rope, + self.tma_desc_v, + self.tma_desc_bulk_extra, + self.tma_desc_sparse_rope, + self.tma_desc_sparse_extra_rope, + rows=rows_per_warp, + rope_dim=cfg.rope_dim, + separate_rope=cfg.rope_dim > 0, + fp8=source_is_fp8( + cfg, + cached_page_ids[page_id_slot * cfg.sparse_cache_words] + < Int32(0), + ), + ) + return + load_nvfp4_panel( + stage_base.data_ptr(warp * rows_per_warp * active_width), + self._nv_scratch.data_ptr(warp * rows_per_warp * 64), + self._nv_barriers.data_ptr(warp), + cached_page_ids, + page_id_slot * cfg.sparse_cache_words, + dim_offset, + self.tma_desc_c_latent, + self.tma_desc_c_rope, + self.tma_desc_v, + self.tma_desc_bulk_extra, + self.tma_desc_sparse_rope, + self.tma_desc_sparse_extra_rope, + rows=rows_per_warp, + rope_dim=cfg.rope_dim, + rope_only=not is_v and active_width == 64, + load_scales=True, + fp8=source_is_fp8( + cfg, + cached_page_ids[page_id_slot * cfg.sparse_cache_words] < Int32(0), + ), + ) + return + @cute.jit def _producer_load_kv_stage( self, @@ -958,7 +1562,12 @@ def _producer_load_kv_stage( if cutlass.const_expr(tma_desc is None): return local_tile_idx = local_kv_tile_idx( - cfg, stage_info, inst_id, is_v, section=section + cfg, + stage_info, + inst_id, + is_v, + section=section, + requested_splits=self.runtime_split_kv, ) batch_idx = batch_idx_for_stage_cfg(self.batch_idx, cfg, stage_info) cta_idx_q = cta_idx_q_for_stage(self.cta_idx_q, stage_info) @@ -970,11 +1579,18 @@ def _producer_load_kv_stage( self.cu_seqlens_q, batch_idx, ) - tile_idx = global_kv_tile_idx(cfg, local_tile_idx, seq_len_kv, cta_idx_kv) + tile_idx = global_kv_tile_idx( + cfg, + local_tile_idx, + seq_len_kv, + cta_idx_kv, + requested_splits=self.runtime_split_kv, + task_cache=decode_gen_task_cache(stage_info), + ) stage_base = self._stage_base(stage_info) inner_width = min(tma_inner_dim_elems(cfg), active_width) uses_compact_fp8_rope_stage = ( - cfg.load_num_warps == 1 + (cfg.load_num_warps == 1 or cfg.sparse_d576_reload) and cfg.is_fp8_qkv() and not is_v and cfg.rope_dim == 64 @@ -982,11 +1598,121 @@ def _producer_load_kv_stage( and cfg.head_dim_per_stage_kv == 128 ) + if cutlass.const_expr(cfg.sparse_d576_reload): + # D576 sparse K/V stage: each loader warp issues its quads (one + # Gather4 of four token rows each) from the warp-uniform rows that + # read_offsets cached for this tile. The four rows land at + # consecutive SMEM rows of the stage, so the TMA swizzle (SW128 + # latent / SW64 RoPE) is the one the UMMA descriptors expect. Rows + # addressed outside the pool are zero-filled by the TMA unit and + # still count toward the stage's transaction bytes; only -1 slots + # are also masked in softmax. + assert cfg.cache_uniform_sparse_quads, ( + "D576 sparse loads issue from the cached uniform quads" + ) + is_rope_stage = (not is_v) and ( + qk_stage_idx >= cfg.latent_dim // cfg.head_dim_per_stage_kv + ) + num_load_warps = cfg.load_num_warps + warp_in_task = Int32(0) + if cutlass.const_expr(num_load_warps > 1): + warp_in_task = cute.arch.make_warp_uniform( + cute.arch.warp_idx() + ) - Int32(cfg.load_warp_idx) + # Row maps: latent D128 slices (SW128) or the RoPE columns (SW64). + descriptor = ( + self.tma_desc_c_rope if is_rope_stage else self.tma_desc_c_latent + ) + quads_per_warp = cfg.tile_size_kv // (4 * num_load_warps) + gather4_cached_warp( + stage_base.data_ptr( + warp_in_task * Int32(quads_per_warp * 4 * active_width) + ), + descriptor, + Int32(0) if is_rope_stage else dim_offset, + cached_page_ids, + stage_info.barrier, + cache_start=6, + num_quads=quads_per_warp, + num_slices=1, + quad_stride=4 * active_width, + slice_stride=0, + col_stride=0, + ) + if cutlass.const_expr(uses_compact_fp8_rope_stage): + if warp_in_task == Int32(0): + if prims.elect_sync(): + missing_stage_bytes = ( + cfg.tile_size_kv + * (cfg.head_dim_per_stage_kv - active_width) + * cfg.qkv_dtype_bytes + ) + prims.mbarrier_complete_tx( + stage_info.barrier, Int32(missing_stage_bytes) + ) + return # Each lane owns four adjacent index-list entries. Bulk TMA requires # a verified 32-row physical run within one source; irregular rows # use gather4. The separate empty-window path only zero-fills padding. if cutlass.const_expr(cfg.num_tokens_per_page == 1): warp = cute.arch.thread_idx()[0] // 32 - cfg.load_warp_idx + if cutlass.const_expr(cfg.has_nvfp4_kv): + self.load_nvfp4_stage_payload( + stage_info, + stage_base, + cached_page_ids, + dim_offset, + warp, + page_id_slot=page_id_slot, + qk_stage_idx=qk_stage_idx, + is_v=is_v, + active_width=active_width, + ) + publish_nvfp4_kv( + stage_info.barrier, + ( + cfg.whole_kv_tile_bytes + if cfg.whole_kv_tile + else cfg.kv_smem_tile_bytes + ) + // cfg.load_num_warps, + ) + return + if cutlass.const_expr( + not is_v and active_width == 64 and cfg.rope_dim == 64 + ): + # The QK-only tail uses its own 64-wide swizzle. Retire the + # unused half of the fixed-width pipeline transaction once. + groups_per_warp = cfg.tile_size_kv // (4 * cfg.load_num_warps) + lane = cute.arch.thread_idx()[0] % 32 + base = page_id_slot * cfg.sparse_cache_words + if lane < groups_per_warp: + raw = cached_page_ids[base] + desc = select_gather_map( + self.tma_desc_sparse_rope, + self.tma_desc_sparse_extra_rope, + raw < 0, + ) + coords = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) + for j in cutlass.range_constexpr(4): + row = cached_page_ids[base + j] & Int32(0x7FFFFFFF) + coords[j] = row if row != Int32(0x7FFFFFFF) else Int32(-1) + gather4( + stage_base.data_ptr((warp * groups_per_warp + lane) * 4 * 64), + desc, + Int32(0), + coords[0], + coords[1], + coords[2], + coords[3], + stage_info.barrier, + ) + if warp == 0 and prims.elect_sync(): + prims.mbarrier_complete_tx( + stage_info.barrier, + Int32(cfg.tile_size_kv * 64 * cfg.qkv_dtype_bytes), + ) + return groups_per_warp = cfg.tile_size_kv // (4 * cfg.load_num_warps) base = page_id_slot * cfg.sparse_cache_words for chunk in cutlass.range_constexpr( @@ -1314,9 +2040,10 @@ def _set_desc( # K waits advance independently of releases. Keep its payload # live until the corresponding delayed PV has consumed it. stage_base = self._smem_kv.subview( - self.consumer_release_state.index - * Int32(self.cfg.kv_smem_stage_elements) + self.kv_stage_offset(self.consumer_release_state.index) ) + if cutlass.const_expr(self.cfg.whole_kv_tile): + stage_base = stage_base.subview(Int32(k_subtile_idx * 16384)) v_leading_byte_offset = ( Int32(0) if self.cfg.head_dim_per_stage_v == TCGEN05_BF16_K_BLOCK_WIDTH @@ -1324,6 +2051,10 @@ def _set_desc( ) if cutlass.const_expr(self.cfg.is_fp8_qkv()): v_leading_byte_offset = Int32(0) + if cutlass.const_expr(self.cfg.ws_pv): + # B is N-major: next128-byte output-feature atom is the + # next retained128x128 panel, not the next KV row. + v_leading_byte_offset = Int32(self.cfg.kv_smem_tile_bytes) desc = cprims.Tcgen05SmemDesc.build( stage_base, leading_byte_offset=v_leading_byte_offset, @@ -1347,7 +2078,11 @@ def _set_desc( * self.cfg.qkv_dtype_bytes ) return cprims.Tcgen05SmemDesc.build( - self._stage_base(stage_info), + self._stage_base(stage_info).subview( + Int32(k_subtile_idx * 16384) + if self.cfg.whole_kv_tile + else Int32(0) + ), leading_byte_offset=k_leading_byte_offset, stride_byte_offset=Int32( qkv_major_k_stride_bytes_for(self.cfg, self.cfg.rope_dim) @@ -1357,7 +2092,9 @@ def _set_desc( ), ) desc = cprims.Tcgen05SmemDesc.build( - self._stage_base(stage_info), + self._stage_base(stage_info).subview( + Int32(k_subtile_idx * 16384) if self.cfg.whole_kv_tile else Int32(0) + ), leading_byte_offset=k_leading_byte_offset, stride_byte_offset=stride_byte_offset, layout=qkv_smem_swizzle(self.cfg), @@ -1392,24 +2129,29 @@ def v_desc_reused( self, stage_info: StageInfo, *, v_subtile_idx: cutlass.Constexpr[int] ): """Read delayed V from the unreleased K payload in the same SMEM slot.""" - del v_subtile_idx - return self._set_desc(stage_info, 1, 0, reuse_k=True) + return self._set_desc( + stage_info, 1, 0, k_subtile_idx=v_subtile_idx, reuse_k=True + ) @consumer_work(returns=("v_desc_0",)) @cute.jit def v_desc_reused_0( self, stage_info: StageInfo, *, v_subtile_idx: cutlass.Constexpr[int] ): - del v_subtile_idx - return self._set_desc(stage_info, 1, 0, reuse_k=True) + """Bind V consumer zero to the selected feature slice of its retained K payload.""" + return self._set_desc( + stage_info, 1, 0, k_subtile_idx=v_subtile_idx, reuse_k=True + ) @consumer_work(returns=("v_desc_1",)) @cute.jit def v_desc_reused_1( self, stage_info: StageInfo, *, v_subtile_idx: cutlass.Constexpr[int] ): - del v_subtile_idx - return self._set_desc(stage_info, 1, 1, reuse_k=True) + """Bind V consumer one to the selected feature slice of its retained K payload.""" + return self._set_desc( + stage_info, 1, 1, k_subtile_idx=v_subtile_idx, reuse_k=True + ) @consumer_work(returns=("v_desc_1",)) @cute.jit diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_corr.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_corr.py index cf547165e9e1..a09f9d5cc4ea 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_corr.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_corr.py @@ -32,6 +32,7 @@ producer_work, ) +from ...helpers.math import final_lse from ...helpers.constants import ( SMEM_WORD_BYTE_SHIFT, SMEM_WORD_BYTES, @@ -70,8 +71,8 @@ partial_output_dtype, ) from ...helpers.ops import ( - fp8_quant_scale, - fp8_quant_scale_rcp, + fp8_p_scale, + fp8_p_scale_rcp, pack_float4_to_fp8_e4m3, store_transposed_smem8b_x2, store_transposed_smem8b_x4, @@ -92,6 +93,8 @@ cta_idx_kv_for_stage, cta_idx_q_for_stage, head_idx_for_stage, + global_kv_tile_idx, + runtime_seq_len_kv_from_task_cache, ) from .common import ( @@ -116,7 +119,7 @@ class TmemCorrResource(MlaResource): inst_id: cutlass.Constexpr[int] = 0 scale_softmax_log2: Float32 = None output_scale: Float32 = None - atten_sinks: object = None + attention_sinks: object = None output_quant: cutlass.Constexpr = None output_quant_args: object = None o_tensor: object = None @@ -124,6 +127,7 @@ class TmemCorrResource(MlaResource): acc_o_tensor: object = None acc_lse_tensor: object = None cache_seqs: object = None + page_offsets: object = None head_idx: object = None batch_idx: object = None cta_idx_q: object = None @@ -131,6 +135,10 @@ class TmemCorrResource(MlaResource): cta_idx_head_dim_v: object = None store_barrier_id: cutlass.Constexpr[int] = 6 sum_barrier_id: cutlass.Constexpr[int] = 7 + # WS means tcgen05.mma.ws (M64 warp-specialized PV). This scratch either + # joins the two Layout-E N-column sums (128 FP32 values, 512 bytes) or + # relays Layout-F row factors to WS consumers (64 FP32 values, 256 bytes). + _ws_scale_alloc: cutlass.Constexpr[Optional[SmemAllocation]] = None _sum_alloc: cutlass.Constexpr[Optional[SmemAllocation]] = None _cluster_reduction_alloc: cutlass.Constexpr[Optional[SmemAllocation]] = None _cluster_reduction_barrier_alloc: cutlass.Constexpr[Optional[SmemAllocation]] = None @@ -143,12 +151,21 @@ class TmemCorrResource(MlaResource): def _keeps_direct_output(self): return ( self.cfg.kernel_variant == "keeps_mma_ab" - and self.cfg.num_tokens_per_page == 1 + and (self.cfg.num_tokens_per_page == 1 and not self.cfg.sparse_d576_reload) and self.cfg.use_cluster_reduction != 1 ) def get_smem_requirements(self): """Return correction/output SMEM scratch and optional cluster buffers.""" + if self.cfg.ws_pv and self.inst_id == 1: + if self._ws_scale_alloc is None: + self._ws_scale_alloc = SmemAllocation( + name=f"{self.name}_ws_scale", + size_bytes=512 if self.cfg.ws_native_e else 256, + alignment=16, + ) + return [self._ws_scale_alloc] + if self.inst_id != 1 or self._keeps_direct_output(): # Keeps-AB normalizes through warp shuffles and stores from TMEM # directly to GMEM. Its output and sum scratch are never used. @@ -494,7 +511,11 @@ def _store_lse_to_gmem( self, stage_info: StageInfo, task_cache, final_max, reduced_sum ): """Store final LSE or split-KV partial LSE for one corrected tile.""" - if cutlass.const_expr(self.lse_tensor is None and self.acc_lse_tensor is None): + if cutlass.const_expr( + self.lse_tensor is None + and self.acc_lse_tensor is None + and self.cfg.use_cluster_reduction != 1 + ): return batch_idx = batch_idx_for_stage_cfg(self.batch_idx, self.cfg, stage_info) head_base_idx = head_idx_for_stage(self.head_idx, self.cfg, stage_info) @@ -508,6 +529,9 @@ def _store_lse_to_gmem( if cutlass.const_expr(self.cfg.kernel_variant == "keeps_mma_ab"): should_store_lse = lane_idx < Int32(16) local_row_idx = warp_idx * Int32(16) + (lane_idx & Int32(0xF)) + if cutlass.const_expr(self.cfg.ws_native_e): + should_store_lse = warp_idx < Int32(2) + local_row_idx = (warp_idx % Int32(2)) * Int32(32) + lane_idx head_idx = head_base_idx + local_row_idx scale_idx = Int32(0) else: @@ -521,7 +545,8 @@ def _store_lse_to_gmem( ) head_idx = head_base_idx + local_row_idx if cutlass.const_expr( - self.cfg.num_tokens_per_page == 1 and self.cfg.use_cluster_reduction != 1 + (self.cfg.num_tokens_per_page == 1 and not self.cfg.sparse_d576_reload) + and self.cfg.use_cluster_reduction != 1 ): should_store_lse = should_store_lse and cta_idx_head_dim_v == Int32(0) if should_store_lse: @@ -555,7 +580,7 @@ def _store_lse_to_gmem( if cutlass.const_expr(self.cfg.is_fp8_qkv()): # P and the online sum use the same FP8 quantization scale. # Undo that scale for the externally visible log-sum-exp. - lse_sum = lse_sum * fp8_quant_scale_rcp() + lse_sum = lse_sum * fp8_p_scale_rcp(self.cfg) lse_val = ( cute.math.log2(lse_sum, fastmath=True) + self.scale_softmax_log2 * selected_max @@ -592,7 +617,7 @@ def _store_lse_to_gmem( (self.acc_lse_tensor.iterator.raw_ptr() + elem_offset).store( lse_val ) - else: + elif cutlass.const_expr(self.lse_tensor is not None): if valid_output_row: elem_offset = public_query_flat_row( self.cfg, @@ -601,7 +626,7 @@ def _store_lse_to_gmem( self.cu_seqlens_q, ) (self.lse_tensor.iterator.raw_ptr() + elem_offset).store( - lse_val + final_lse(lse_val, self.cfg.lse_base_on_e) ) @cute.jit @@ -846,7 +871,7 @@ def _reduce_cluster_partials_and_store_for_tile( head = storage_flat_query_row % Int32( cfg.logical_num_heads_q ) - sink = Float32(self.atten_sinks[head]) + sink = Float32(self.attention_sinks[head]) sink_scale = Float32(0) if has_finite_mass and sink != Float32(Float32.inf): delta = sink * Float32(LOG2_E) - global_lse @@ -873,7 +898,7 @@ def _reduce_cluster_partials_and_store_for_tile( ) ( self.lse_tensor.iterator.raw_ptr() + lse_offset - ).store(published_lse) + ).store(final_lse(published_lse, cfg.lse_base_on_e)) acc_vec = vector_from_scalars( ( @@ -1031,6 +1056,39 @@ def correct_tail_and_store( tail_o_stage_idx_1=tail_o_stage_idx_1, ) + @cute.jit + def ws_sum_reduce(self, stage_info, value): + """Join the two N-column partial softmax sums for each Layout-E M64 query row. + + Threads tid and tid^64 represent the same row in the two column halves + of the WS data-path layout. Their sums must be added once before O + normalization. The barriers publish scratch and protect it from reuse. + PTX: https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-data-path-layout-e""" + scratch = smem_array(stage_info.context, self._ws_scale_alloc, Float32, 128) + tid = cute.arch.thread_idx()[0] % Int32(128) + scratch[tid] = value + prims.barrier_cta_sync(barrier_id=self.store_barrier_id, thread_count=128) + result = value + scratch[tid ^ Int32(64)] + prims.barrier_cta_sync(barrier_id=self.store_barrier_id, thread_count=128) + return result + + @cute.jit + def ws_row_scale(self, stage_info, scale): + # Existing stats arrive in Layout-F thread order. One writer per row; + # every correction thread participates before warp-local skip logic. + """Remap legacy Layout-F row factors to the Layout-E WS PV consumer order. + + One writer per row publishes a factor; all correction threads read the + matching row, with barriers before reading and before scratch reuse.""" + scratch = smem_array(stage_info.context, self._ws_scale_alloc, Float32, 64) + tid = cute.arch.thread_idx()[0] % 128 + if tid % 32 < 16: + scratch[(tid // 32) * 16 + tid % 16] = scale + prims.barrier_cta_sync(barrier_id=self.store_barrier_id, thread_count=128) + result = scratch[tid % 64] + prims.barrier_cta_sync(barrier_id=self.store_barrier_id, thread_count=128) + return result + @cute.jit def _apply_loop_correction_in_tmem( self, @@ -1048,6 +1106,29 @@ def _apply_loop_correction_in_tmem( o_reg_pair_count = num_o_reg_pairs(cfg) num_o_tmem_loads = num_o_tmem_loads_per_stage(cfg) + if cutlass.const_expr(cfg.ws_pv): + for tile in cutlass.range_constexpr(4): + addr = tmem_row_base + self._o_base_col() + Int32(tile * 64) + loaded = prims.tcgen05_ld( + prims.Tcgen05LdStShape.SHAPE_32X32B, + prims.make_tmem_ptr(addr, Float32), + num=64, + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + cute.arch.fence_view_async_tmem_load() + scaled = vector_from_scalars( + tuple( + loaded[i].bitcast(Float32) * scale_vals[0] for i in range(64) + ), + Float32, + ) + prims.tcgen05_st( + prims.Tcgen05LdStShape.SHAPE_32X32B, + prims.make_tmem_ptr(addr, Float32), + scaled, + ) + return + if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): for v_stage_idx in cutlass.range_constexpr(cfg.v_head_dim_stages): base_addr = ( @@ -1208,6 +1289,23 @@ def _apply_loop_correction_in_tmem( ) return + @cute.jit + def tensor_scale_tile(self, stage_info, local_tile): + """Map one local pipeline-instance tile to its global split-KV offset and query row.""" + cfg = self.cfg + batch = batch_idx_for_stage_cfg(self.batch_idx, cfg, stage_info) + length = runtime_seq_len_kv_from_task_cache( + cfg, + decode_gen_task_cache(stage_info), + cta_idx_q_for_stage(self.cta_idx_q, stage_info), + self.cu_seqlens_q, + batch, + ) + tile = global_kv_tile_idx( + cfg, local_tile, length, cta_idx_kv_for_stage(self.cta_idx_kv, stage_info) + ) + return tile * Int32(cfg.tile_size_kv), batch + @cute.jit def _rescale_loop_o( self, @@ -1269,6 +1367,28 @@ def _rescale_loop_o( lane_scales_are_identity & scale_0_is_identity & scale_1_is_identity ) + if cutlass.const_expr(cfg.has_fp8_primary): + # O is still expressed in the previous source's V units. Apply + # its descale/current descale together with the normal softmax + # correction when this stream crosses primary -> extra. This is + # independent of the denominator update, which uses softmax only. + local = (stage_info.loop_offset + Int32(1)) * Int32( + cfg.num_insts_kv + ) + Int32(self.inst_id % cfg.num_insts_kv) + current, batch = self.tensor_scale_tile(stage_info, local) + ratio = self.page_offsets.tensor_scale_ratio( + current, current - Int32(cfg.num_insts_kv * cfg.tile_size_kv), batch + ) + for i in cutlass.range_constexpr(num_scale_groups): + scale_vals[i] *= ratio + lane_scales_are_identity = lane_scales_are_identity and ratio == Float32(1) + + if cutlass.const_expr(cfg.ws_pv and not cfg.ws_native_e): + # Legacy non-WS QK produces Layout-F row factors, while WS PV/O + # uses Layout E. Relay them before the warp-wide identity test; + # every correction lane must participate in the scratch barriers. + scale_vals[0] = self.ws_row_scale(stage_info, scale_vals[0]) + lane_scales_are_identity = scale_vals[0] == Float32(1.0) skip_rescale = prims.vote_sync( cute.arch.FULL_MASK, lane_scales_are_identity, @@ -1338,15 +1458,20 @@ def _compute_tail_softmax_scales( exp_scale0[scale_idx] = Float32(1.0) exp_scale1[scale_idx] = Float32(0.0) final_sum[scale_idx] = inst_sum - final_sum[scale_idx] += Float32( - cprims.shfl_sync( - thread_mask=0xFFFFFFFF, - val=final_sum[scale_idx], - offset=16, - mask_and_clamp=0x1F, - kind=cprims.Shfl.BFLY, + if cutlass.const_expr(cfg.ws_native_e): + final_sum[scale_idx] = self.ws_sum_reduce( + stage_info, final_sum[scale_idx] + ) + else: + final_sum[scale_idx] += Float32( + cprims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=final_sum[scale_idx], + offset=16, + mask_and_clamp=0x1F, + kind=cprims.Shfl.BFLY, + ) ) - ) reduced_sum[scale_idx] = final_sum[scale_idx] else: inst0_sum = inst0_sum_arr[scale_idx] @@ -1376,33 +1501,34 @@ def _compute_tail_softmax_scales( inst0_sum * exp_scale0[scale_idx] + inst1_sum * exp_scale1[scale_idx] ) - final_sum[scale_idx] += Float32( - cprims.shfl_sync( - thread_mask=0xFFFFFFFF, - val=final_sum[scale_idx], - offset=16, - mask_and_clamp=0x1F, - kind=cprims.Shfl.BFLY, + if cutlass.const_expr(not cfg.compact_softmax_stats): + final_sum[scale_idx] += Float32( + cprims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=final_sum[scale_idx], + offset=16, + mask_and_clamp=0x1F, + kind=cprims.Shfl.BFLY, + ) ) - ) - final_sum[scale_idx] += Float32( - cprims.shfl_sync( - thread_mask=0xFFFFFFFF, - val=final_sum[scale_idx], - offset=8, - mask_and_clamp=0x1F, - kind=cprims.Shfl.BFLY, + final_sum[scale_idx] += Float32( + cprims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=final_sum[scale_idx], + offset=8, + mask_and_clamp=0x1F, + kind=cprims.Shfl.BFLY, + ) ) - ) - final_sum[scale_idx] += Float32( - cprims.shfl_sync( - thread_mask=0xFFFFFFFF, - val=final_sum[scale_idx], - offset=4, - mask_and_clamp=0x1F, - kind=cprims.Shfl.BFLY, + final_sum[scale_idx] += Float32( + cprims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=final_sum[scale_idx], + offset=4, + mask_and_clamp=0x1F, + kind=cprims.Shfl.BFLY, + ) ) - ) if cutlass.const_expr(cfg.kernel_variant != "keeps_mma_ab"): # The two softmax pipes first reduce within their owning warp, then @@ -1453,13 +1579,20 @@ def _compute_tail_softmax_scales( else col_group_idx ) norm_lanes = 16 if cfg.kernel_variant == "keeps_mma_ab" else 4 + if cutlass.const_expr(cfg.ws_native_e): + norm_lane = lane_idx + norm_lanes = 32 if lane_idx < Int32(norm_lanes): logical_head = ( cta_idx_q_for_stage(self.cta_idx_q, stage_info) * cfg.tile_size_q + head_idx_for_stage(self.head_idx, cfg, stage_info) + ( - warp_idx * Int32(16) + norm_lane + ( + (warp_idx % Int32(2)) * Int32(32) + norm_lane + if cfg.ws_native_e + else warp_idx * Int32(16) + norm_lane + ) if cfg.kernel_variant == "keeps_mma_ab" else local_q_head_idx_for_scale( cfg, col_group_idx, scale_idx @@ -1469,7 +1602,7 @@ def _compute_tail_softmax_scales( count = Float32(final_max[scale_idx] != neg_max_f32()) norm_scale = Float32(0) if logical_head < cfg.logical_num_heads_q and count > Float32(0): - sink = Float32(self.atten_sinks[logical_head]) + sink = Float32(self.attention_sinks[logical_head]) if sink == Float32(-Float32.inf): norm_scale = self.output_scale / reduced_sum[scale_idx] elif sink != Float32(Float32.inf): @@ -1484,7 +1617,7 @@ def _compute_tail_softmax_scales( cute.math.min(delta, Float32(0)), approx=True ) if cutlass.const_expr(cfg.is_fp8_qkv()): - sink_mass *= fp8_quant_scale() + sink_mass *= fp8_p_scale(cfg) norm_scale = ( self.output_scale * attention_rescale @@ -1505,6 +1638,17 @@ def _compute_tail_softmax_scales( ) final_scale0[scale_idx] = norm_scale * exp_scale0[scale_idx] final_scale1[scale_idx] = norm_scale * exp_scale1[scale_idx] + if cutlass.const_expr(cfg.has_fp8_primary): + last0, batch = self.tensor_scale_tile( + stage_info, stage_info.loop_end * Int32(cfg.num_insts_kv) + ) + scale0 = self.page_offsets.tensor_scale(last0, batch) + scale1 = self.page_offsets.tensor_scale( + last0 + Int32(cfg.tile_size_kv), batch + ) + for i in cutlass.range_constexpr(num_scale_groups): + final_scale0[i] *= scale0 + final_scale1[i] *= scale1 return final_max, reduced_sum, final_scale0, final_scale1 @cute.jit @@ -1520,6 +1664,11 @@ def _store_keeps_tail_o_stage( ): """Normalize and directly store one keeps-MMA-AB O stage.""" cfg = self.cfg + if cutlass.const_expr(cfg.ws_pv and v_stage_idx % 2 != 0): + return + store_scale = final_scale0[0] + if cutlass.const_expr(cfg.ws_pv and not cfg.ws_native_e): + store_scale = self.ws_row_scale(stage_info, store_scale) batch_idx = batch_idx_for_stage_cfg(self.batch_idx, cfg, stage_info) head_base_idx = head_idx_for_stage(self.head_idx, cfg, stage_info) cta_idx_q = cta_idx_q_for_stage(self.cta_idx_q, stage_info) @@ -1529,6 +1678,8 @@ def _store_keeps_tail_o_stage( # Keeps-MMA-AB stores one head row per half-warp. The low 16 lanes own # the row index; the upper lane bit selects the head-dim half. local_row_idx = warp_idx * Int32(16) + (lane_idx & Int32(0xF)) + if cutlass.const_expr(cfg.ws_pv): + local_row_idx = (warp_idx % Int32(2)) * Int32(32) + lane_idx global_head_idx = head_base_idx + local_row_idx storage_flat_query_row, _, _, _, valid_output_row = flat_query_row_state( global_head_idx, @@ -1543,23 +1694,34 @@ def _store_keeps_tail_o_stage( cfg.head_dim_per_stage_v // 2 ) + if cutlass.const_expr(cfg.ws_pv): + half_warp_col_offset = (warp_idx // Int32(2)) * Int32(128) + if cutlass.const_expr(self.output_quant is not None): - # Each lane owns half a head's D stage. K128 can span the two - # lanes separated by 16; K32 blocks remain entirely lane-local. - width = min(self.output_quant.block_size, cfg.head_dim_per_stage_v // 2) - for stripe in cutlass.range_constexpr( - 0, cfg.head_dim_per_stage_v // 2, width - ): + # Non-WS Layout F gives each lane half a D128 stage; block128 + # spans lanes separated by16. WS Layout E gives a lane128 values + # in a D256 PV pass, so block32/block128 are both lane-local. + # Match the TMEM load shape and extent to the selected layout. + lane_columns = 128 if cfg.ws_pv else cfg.head_dim_per_stage_v // 2 + width = min(self.output_quant.block_size, lane_columns) + for stripe in cutlass.range_constexpr(0, lane_columns, width): values = cutlass.Array(Float32, width, space=cutlass.AddressSpace.rmem) - loaded = tcgen05_ld_16x32bx2_f32( - prims.make_tmem_ptr(base_addr0 + Int32(stripe), Float32), - num=width, - offset=Int32(cfg.head_dim_per_stage_v // 2), - ) + if cutlass.const_expr(cfg.ws_pv): + loaded = prims.tcgen05_ld( + prims.Tcgen05LdStShape.SHAPE_32X32B, + prims.make_tmem_ptr(base_addr0 + Int32(stripe), Float32), + num=width, + ) + else: + loaded = tcgen05_ld_16x32bx2_f32( + prims.make_tmem_ptr(base_addr0 + Int32(stripe), Float32), + num=width, + offset=Int32(cfg.head_dim_per_stage_v // 2), + ) prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) cute.arch.fence_view_async_tmem_load() for i in cutlass.range_constexpr(width): - values[i] = loaded[i] * final_scale0[0] + values[i] = loaded[i] * store_scale if global_head_idx < Int32(cfg.num_heads_q) and valid_output_row: row = public_query_flat_row( cfg, storage_flat_query_row, batch_idx, self.cu_seqlens_q @@ -1572,97 +1734,122 @@ def _store_keeps_tail_o_stage( + v_stage_idx * cfg.head_dim_per_stage_v + half_warp_col_offset + stripe, - lane_stride=16, - rotary=v_stage_idx + 1 == cfg.v_head_dim_stages, + lane_stride=1 if cfg.ws_pv else 16, + # WS consumes two D128 descriptor slices per PV pass; + # its last even stage owns the final64 rotary channels. + rotary=v_stage_idx + (2 if cfg.ws_pv else 1) + == cfg.v_head_dim_stages, ) return + store_col_base = ( + head_dim_offset + + Int32(v_stage_idx * cfg.head_dim_per_stage_v) + + half_warp_col_offset + ) + if cutlass.const_expr(self.acc_o_tensor is not None): + base_elem_offset = split_o_element_offset( + cfg, + batch_idx, + cta_idx_q, + global_head_idx, + cta_idx_kv, + store_col_base, + ) + base_ptr = self.acc_o_tensor.iterator.raw_ptr().toint(Int64) + element_bytes = cfg.partial_o_dtype_bytes + else: + output_query_row = public_query_flat_row( + cfg, + storage_flat_query_row, + batch_idx, + self.cu_seqlens_q, + ) + base_elem_offset = Int64(output_query_row) * Int64(cfg.head_dim_v) + Int64( + store_col_base + ) + base_ptr = self.o_tensor.iterator.raw_ptr().toint(Int64) + element_bytes = cfg.o_dtype_bytes + row_address = base_ptr + base_elem_offset * Int64(element_bytes) + valid_store = (global_head_idx < Int32(cfg.num_heads_q)) & valid_output_row + + # Compute the row address once, outside the validity guard. Keeping + # only each store predicated lets subsequent TMEM loads overlap it. # ``tcgen05.ld.sync.aligned`` must remain convergent across the warp. # Padded query rows still own valid physical TMEM rows, so load and # normalize every row and predicate only the externally visible store. - for chunk_idx in cutlass.range_constexpr(0, cfg.head_dim_per_stage_v // 2, 8): - o_loaded = tcgen05_ld_16x32bx2_f32( - prims.make_tmem_ptr(base_addr0 + Int32(chunk_idx), Float32), - num=8, - offset=Int32(cfg.head_dim_per_stage_v // 2), - ) - prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) - cute.arch.fence_view_async_tmem_load() - if cutlass.const_expr( - cfg.use_fp8_output == 1 - and self.acc_o_tensor is None - and cfg.use_cluster_reduction != 1 - ): - final_vals = cutlass.Array(Float32, 8, space=cutlass.AddressSpace.rmem) - packed_o = cutlass.Array(Int32, 2, space=cutlass.AddressSpace.rmem) + for chunk_idx in cutlass.range_constexpr( + 0, + 128 if cfg.ws_pv else cfg.head_dim_per_stage_v // 2, + 32 if cfg.ws_pv else 8, + ): + if cutlass.const_expr(cfg.ws_pv): + o_loaded = prims.tcgen05_ld( + prims.Tcgen05LdStShape.SHAPE_32X32B, + prims.make_tmem_ptr(base_addr0 + Int32(chunk_idx), Float32), + num=32, + ) else: - packed_o = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) - for pair_idx in cutlass.range_constexpr(4): - reg_base = pair_idx * 2 - final_pair = fmul2( - (final_scale0[0], final_scale0[0]), - (o_loaded[reg_base], o_loaded[reg_base + 1]), + o_loaded = tcgen05_ld_16x32bx2_f32( + prims.make_tmem_ptr(base_addr0 + Int32(chunk_idx), Float32), + num=8, + offset=Int32(cfg.head_dim_per_stage_v // 2), ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + cute.arch.fence_view_async_tmem_load() + # Load32 columns per lane, then retain the same aligned16B + # stores and exact BF16 rounding in four8-value subchunks. + for store_sub in cutlass.range_constexpr(0, 32 if cfg.ws_pv else 8, 8): if cutlass.const_expr( cfg.use_fp8_output == 1 and self.acc_o_tensor is None and cfg.use_cluster_reduction != 1 ): - final_vals[reg_base] = final_pair[0] - final_vals[reg_base + 1] = final_pair[1] + final_vals = cutlass.Array( + Float32, 8, space=cutlass.AddressSpace.rmem + ) + packed_o = cutlass.Array(Int32, 2, space=cutlass.AddressSpace.rmem) else: - packed_o[pair_idx] = pack_float2_to_bf16( - final_pair[0], final_pair[1] + packed_o = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) + for pair_idx in cutlass.range_constexpr(4): + reg_base = pair_idx * 2 + final_pair = fmul2( + (store_scale, store_scale), + ( + o_loaded[store_sub + reg_base], + o_loaded[store_sub + reg_base + 1], + ), ) - if cutlass.const_expr( - cfg.use_fp8_output == 1 - and self.acc_o_tensor is None - and cfg.use_cluster_reduction != 1 - ): - packed_o[0] = pack_float4_to_fp8_e4m3( - final_vals[0], - final_vals[1], - final_vals[2], - final_vals[3], - ) - packed_o[1] = pack_float4_to_fp8_e4m3( - final_vals[4], - final_vals[5], - final_vals[6], - final_vals[7], - ) - store_col_offset = ( - head_dim_offset - + Int32(v_stage_idx * cfg.head_dim_per_stage_v) - + half_warp_col_offset - + Int32(chunk_idx) - ) - if global_head_idx < Int32(cfg.num_heads_q) and valid_output_row: - if cutlass.const_expr(self.acc_o_tensor is not None): - base_elem_offset = split_o_element_offset( - cfg, - batch_idx, - cta_idx_q, - global_head_idx, - cta_idx_kv, - store_col_offset, + if cutlass.const_expr( + cfg.use_fp8_output == 1 + and self.acc_o_tensor is None + and cfg.use_cluster_reduction != 1 + ): + final_vals[reg_base] = final_pair[0] + final_vals[reg_base + 1] = final_pair[1] + else: + packed_o[pair_idx] = pack_float2_to_bf16( + final_pair[0], final_pair[1] + ) + if cutlass.const_expr( + cfg.use_fp8_output == 1 + and self.acc_o_tensor is None + and cfg.use_cluster_reduction != 1 + ): + packed_o[0] = pack_float4_to_fp8_e4m3( + final_vals[0], + final_vals[1], + final_vals[2], + final_vals[3], ) - base_ptr = self.acc_o_tensor.iterator.raw_ptr().toint(Int64) - element_bytes = cfg.partial_o_dtype_bytes - else: - output_query_row = public_query_flat_row( - cfg, - storage_flat_query_row, - batch_idx, - self.cu_seqlens_q, + packed_o[1] = pack_float4_to_fp8_e4m3( + final_vals[4], + final_vals[5], + final_vals[6], + final_vals[7], ) - base_elem_offset = Int64(output_query_row) * Int64( - cfg.head_dim_v - ) + Int64(store_col_offset) - base_ptr = self.o_tensor.iterator.raw_ptr().toint(Int64) - element_bytes = cfg.o_dtype_bytes dst_ptr = cutlass.inttoptr( - base_ptr + Int64(base_elem_offset) * Int64(element_bytes), + row_address + Int64((chunk_idx + store_sub) * element_bytes), mem_space=1, dtype=Int32, ) @@ -1671,9 +1858,11 @@ def _store_keeps_tail_o_stage( and self.acc_o_tensor is None and cfg.use_cluster_reduction != 1 ): - dst_ptr.store(packed_o.load(0, 2), alignment=8) + if valid_store: + dst_ptr.store(packed_o.load(0, 2), alignment=8) else: - dst_ptr.store(packed_o.load(0, 4), alignment=16) + if valid_store: + dst_ptr.store(packed_o.load(0, 4), alignment=16) @cute.jit def _stage_tile32_tail_o_to_smem( diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_o.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_o.py index 6e66442172e4..3a46e4622db9 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_o.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_o.py @@ -32,15 +32,18 @@ producer_work, ) +from ....tcgen05_compat import tcgen05_mma_ws from ...helpers.constants import ( WARP_LANES, ) from ...helpers.layout import ( + _TASK_CACHE_LOCAL_KV_TILES, _TASK_CACHE_TMEM_BASE_OFFSET, SMEM_ROW_BYTES, decode_gen_task_cache, o_stage_tmem_col_offset, + smem_array, q_p_desc_k_block_wrap_bytes, ) from ...helpers.math import ( @@ -143,6 +146,10 @@ def init_stage_work_tile_state(self, stage_info: StageInfo): @cute.jit def _tail_has_prior_o(self, stage_info: StageInfo): """Return whether tail PV should accumulate into existing O state.""" + if cutlass.const_expr(self.runtime_split_kv is not None): + return Int32( + decode_gen_task_cache(stage_info)[_TASK_CACHE_LOCAL_KV_TILES] + ) > Int32(self.cfg.num_insts_kv) batch_idx = batch_idx_for_stage_cfg(self.batch_idx, self.cfg, stage_info) cta_idx_q = cta_idx_q_for_stage(self.cta_idx_q, stage_info) seq_len_kv = runtime_seq_len_kv_from_task_cache( @@ -152,9 +159,9 @@ def _tail_has_prior_o(self, stage_info: StageInfo): self.cu_seqlens_q, batch_idx, ) - return runtime_local_kv_tiles(self.cfg, seq_len_kv) > Int32( - self.cfg.num_insts_kv - ) + return runtime_local_kv_tiles( + self.cfg, seq_len_kv, requested_splits=self.runtime_split_kv + ) > Int32(self.cfg.num_insts_kv) @cute.jit def _p_desc_for_inst(self, p_inst: int): @@ -370,6 +377,47 @@ def _issue_tmem_p_pv_mma( ): """Issue PV MMA for keeps-MMA-AB with P sourced from TMEM.""" cfg = self.cfg + if cutlass.const_expr(cfg.ws_pv): + assert ( + cfg.async_nvfp4_kv_gather + and cfg.tile_size_q == 64 + and cfg.tile_size_kv == 128 + ) + p_smem = smem_array( + stage_info.context, self.p_tmem_ref._ws_p_alloc, Int32, 4096 + ) + p_desc = cprims.Tcgen05SmemDesc.build( + p_smem.data_ptr(p_stage_idx * Int32(2048)), + leading_byte_offset=16, + stride_byte_offset=1024, + layout=qkv_smem_swizzle(cfg), + ) + idesc = cprims.Tcgen05InstrDesc.build( + c_dtype=Float32, + a_dtype=qkv_dtype(cfg), + b_dtype=qkv_dtype(cfg), + a_major=0, + b_major=1, + n_dim=256, + m_dim=64, + ) + if prims.elect_sync(): + base_addr = self._pv_o_base_addr( + stage_info, v_subtile_idx=v_subtile_idx + ) + for _kb in cutlass.range_constexpr(4): + tcgen05_mma_ws( + mma_kind_for_qkv(cfg), + prims.make_tmem_ptr(base_addr, Float32), + p_desc, + v_desc, + idesc, + scale_d, + ) + scale_d = Boolean(True) + p_desc = p_desc + Int32(2) + v_desc = v_desc + Int32(256) + return v_stage_idx = v_subtile_idx % cfg.v_head_dim_stages task_cache = decode_gen_task_cache(stage_info) p_addr = ( diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_p.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_p.py index e867ba4e7c90..48994d5ca62e 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_p.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_p.py @@ -23,6 +23,7 @@ import cutlass.cute as cute from cutlass import Float32, Int32 from cutlass.experimental.task_scheduling.enums import WorkAttr +from cutlass.experimental.task_scheduling.memory import SmemAllocation from cutlass.experimental.task_scheduling.resources import ( StageInfo, TaskLocalVariable, @@ -40,6 +41,7 @@ decode_gen_task_cache, num_packed_p_regs, num_softmax_scale_groups, + smem_array, ) from ...helpers.math import ( fadd2, @@ -50,7 +52,7 @@ ) from ...helpers.ops import ( float_to_u32_bits, - fp8_log2_quant_scale, + fp8_log2_p_scale, pack_float4_to_fp8_e4m3, softmax_sum_state_ptr, tcgen05_second_panel_addr, @@ -87,6 +89,21 @@ class TmemPResource(MlaResource): TaskLocalVariable.uninitialized() ) + # Two ping-pong M64 x N128 E4M3 probability tiles in SMEM. WS PV reads + # this packed A operand rather than the usual A-from-TMEM score alias; + # P's pipeline protects each slot while softmax fills the next one. + _ws_p_alloc: cutlass.Constexpr[Optional[SmemAllocation]] = None + + def get_smem_requirements(self): + """Reserve two FP8 probability panels when WS PV reads P from shared memory.""" + if not self.cfg.ws_pv: + return [] + if self._ws_p_alloc is None: + self._ws_p_alloc = SmemAllocation( + name=f"{self.name}_ws_p", size_bytes=2 * 64 * 128, alignment=1024 + ) + return [self._ws_p_alloc] + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=p_stage_idx) @cute.jit def init_stage_state(self, stage_info: StageInfo): @@ -160,8 +177,8 @@ def materialize_p( ): log2_scale_pair = (self.scale_softmax_log2, self.scale_softmax_log2) neg_scaled_pair = ( - neg_scaled_max[0] + fp8_log2_quant_scale(), - neg_scaled_max[0] + fp8_log2_quant_scale(), + neg_scaled_max[0] + fp8_log2_p_scale(cfg), + neg_scaled_max[0] + fp8_log2_p_scale(cfg), ) has_finite_max = new_max_arr[0] != neg_max_f32() for packed_idx in cutlass.range_constexpr(packed_p_reg_count): @@ -225,7 +242,7 @@ def materialize_p( p_val = cute.math.exp2( s_arr[s_idx] * self.scale_softmax_log2 + neg_scaled_max[scale_idx] - + fp8_log2_quant_scale(), + + fp8_log2_p_scale(cfg), fastmath=True, ) p_vals[elem_idx] = p_val @@ -323,6 +340,32 @@ def materialize_p( alignment=4, ) + if cutlass.const_expr(cfg.ws_pv): + # Materialize P in K-major SMEM for WS MMA A. Native WS QK uses + # Layout E; the non-WS M64 path uses Layout F and remaps its rows. + # These are PTX data-path layouts, not global-cache layouts: + # https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-data-path-layout-e + # https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-data-path-layout-f + dst = smem_array(stage_info.context, self._ws_p_alloc, Int32, 4096) + tid = cute.arch.thread_idx()[0] + row = (tid // 32) * 16 + tid % 16 + col = (tid % 32 // 16) * 16 + if cutlass.const_expr(cfg.ws_native_e): + row = tid % Int32(64) + col = (tid // Int32(64)) * Int32(16) + for part in cutlass.range_constexpr(4): + offset = ( + stage_info.stage_idx * 2048 + + row * 32 + + ((col + part * 4) ^ ((row % 8) * 4)) + ) + dst.data_ptr(offset).store(regs_p.load(part * 4, 4), alignment=16) + cute.arch.fence_view_async_shared() + prims.barrier_cta_sync(barrier_id=0, thread_count=128) + if cutlass.const_expr(cfg.ws_native_e): + # PV consumes shared A; no redundant F-layout TMEM P copies. + return + p_base = self._p_base_addr(stage_info, task_cache) if cutlass.const_expr(cfg.is_fp8_qkv()): tcgen05_store_p_fp8_16x32bx2_x16(p_base, regs_p) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_s.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_s.py index 9831ab71e5d3..e2ed69131b79 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_s.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_s.py @@ -32,6 +32,7 @@ producer_work, ) +from ....tcgen05_compat import tcgen05_mma_ws from ...helpers.constants import ( OCTET_LANES, QUAD_LANE_MASK, @@ -130,6 +131,7 @@ class TmemSResource(MlaResource): new_max_state: object = None sum_state: object = None _scratch_alloc: cutlass.Constexpr[Optional[SmemAllocation]] = None + _ws_bridge_alloc: cutlass.Constexpr[Optional[SmemAllocation]] = None _softmax_scratch: object = None _p_local_sum_arr: object = None old_max_arr: cutlass.Constexpr[TaskLocalVariable] = ( @@ -164,6 +166,13 @@ def _reset_softmax_state(self): def get_smem_requirements(self): """Return the softmax scratch SMEM allocation.""" + if self.cfg.ws_qk: + if self._ws_bridge_alloc is None: + self._ws_bridge_alloc = SmemAllocation( + name=f"{self.name}_ws_layout_bridge", size_bytes=512, alignment=128 + ) + return [self._ws_bridge_alloc] + if ( self.cfg.kernel_variant == "keeps_mma_ab" and self.cfg.num_tokens_per_page == 1 @@ -242,7 +251,10 @@ def init_softmax_state(self, stage_info: StageInfo): if cutlass.const_expr( not ( self.cfg.kernel_variant == "keeps_mma_ab" - and self.cfg.num_tokens_per_page == 1 + and ( + self.cfg.num_tokens_per_page == 1 + and not self.cfg.sparse_d576_reload + ) ) ): self._softmax_scratch = smem_array( @@ -251,8 +263,6 @@ def init_softmax_state(self, stage_info: StageInfo): Uint32, softmax_scratch_words(self.cfg), ) - self.q_desc_current = Int64(0) - self.q_desc_rope_current = Int64(0) if cutlass.const_expr( self.cfg.kernel_variant == "keeps_mma_ab" and self.cfg.num_tokens_per_page != 1 @@ -292,6 +302,16 @@ def reset_softmax_work_tile_state(self, stage_info: StageInfo) -> None: del stage_info self._make_initial_softmax_vars() + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_descriptor_state(self, stage_info: StageInfo): + """Initialize cached Q descriptors in the producer task's scope.""" + # Persistent HEAD and LOOP have separate runtime skip guards. Seed the + # descriptors before those guards so their values dominate every use. + del stage_info + self.q_desc_current = Int64(0) + self.q_desc_rope_current = Int64(0) + @producer_work(work_attrs=WorkAttr.AUXILIARY) @cute.jit def set_q_desc(self, stage_info: StageInfo, *, q_desc, q_desc_rope): @@ -333,7 +353,10 @@ def qk_mma( or cfg.rope_dim != 64 or qk_stage_idx != cfg.latent_dim // cfg.head_dim_per_stage_kv ): - q_desc = q_desc + (q_stage_offset_bytes >> 4) + if cutlass.const_expr(cfg.q_nope_tmem): + q_desc = q_desc + Int64(qk_stage_idx * 32) + else: + q_desc = q_desc + (q_stage_offset_bytes >> 4) task_cache = decode_gen_task_cache(stage_info) stage_col_offset = Int32(0) @@ -366,15 +389,35 @@ def qk_mma( scale_d = Boolean(qk_stage_idx != 0) for k_block in cutlass.range_constexpr(k_block_count): if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): - prims.tcgen05_mma( - mma_kind_for_qkv(cfg), - prims.CTAGroup.CTA_1, - tmem_ptr, - q_desc, - kv_desc, - idesc, - scale_d, - ) + if cutlass.const_expr(cfg.ws_qk): + if cutlass.const_expr(cfg.q_nope_tmem and qk_stage_idx < 4): + tcgen05_mma_ws( + mma_kind_for_qkv(cfg), + tmem_ptr, + prims.make_tmem_ptr(Int32(q_desc), qkv_dtype(cfg)), + kv_desc, + idesc, + scale_d, + ) + else: + tcgen05_mma_ws( + mma_kind_for_qkv(cfg), + tmem_ptr, + q_desc, + kv_desc, + idesc, + scale_d, + ) + else: + prims.tcgen05_mma( + mma_kind_for_qkv(cfg), + prims.CTAGroup.CTA_1, + tmem_ptr, + q_desc, + kv_desc, + idesc, + scale_d, + ) else: prims.tcgen05_mma( mma_kind_for_qkv(cfg), @@ -396,7 +439,11 @@ def qk_mma( q_desc = q_desc + Int32(q_p_desc_k_block_wrap_units(cfg)) else: kv_desc = kv_desc + Int32(TCGEN05_DESC_NEXT_K_BLOCK_UNITS) - q_desc = q_desc + Int32(TCGEN05_DESC_NEXT_K_BLOCK_UNITS) + q_desc = q_desc + Int32( + 8 + if cfg.q_nope_tmem and qk_stage_idx < 4 + else TCGEN05_DESC_NEXT_K_BLOCK_UNITS + ) @consumer_work(returns=(old_max_arr, sum_arr, new_max_arr, local_sum_arr, s_arr)) @cute.jit @@ -468,7 +515,14 @@ def update_softmax( batch_idx, ) local_tile_idx = softmax_kv_tile_idx(cfg, stage_info, self.inst_id) - tile_idx = global_kv_tile_idx(cfg, local_tile_idx, seq_len_kv, cta_idx_kv) + tile_idx = global_kv_tile_idx( + cfg, + local_tile_idx, + seq_len_kv, + cta_idx_kv, + requested_splits=self.runtime_split_kv, + task_cache=task_cache, + ) tile_offset_k = tile_idx * Int32(cfg.tile_size_kv) next_tile_offset_k = tile_offset_k + Int32(cfg.tile_size_kv) raw_mask_routes = ( @@ -486,9 +540,19 @@ def update_softmax( batch_idx, swaps=cfg.kernel_variant != "keeps_mma_ab", tile_size_kv=cfg.tile_size_kv, + column_half=(warp_grp_thread_idx // Int32(64)) + if cfg.ws_native_e + else None, ) + # D576 query rows of one CTA are heads of one token: one mask. if cutlass.const_expr( - cfg.num_tokens_per_page == 1 and cfg.logical_seq_len_q == 1 + ( + cfg.num_tokens_per_page == 1 + and not cfg.sparse_d576_reload + and cfg.fixed_seqlen + and cfg.logical_seq_len_q == 1 + ) + or cfg.sparse_d576_reload ) else raw_mask_routes ) @@ -505,14 +569,25 @@ def update_softmax( + stage_col_offset ) if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): - loaded = tcgen05_ld_16x32bx2_f32( - prims.make_tmem_ptr(base_addr, Float32), - num=cfg.tile_size_kv // 2, - offset=Int32(cfg.tile_size_kv // 2), - ) - prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) - for reg_idx in cutlass.range_constexpr(num_s_regs_per_thread(cfg)): - s_vals[reg_idx] = loaded[reg_idx] + if cutlass.const_expr(cfg.ws_qk): + # Native E: each thread retains one row and half of N. + loaded = prims.tcgen05_ld( + prims.Tcgen05LdStShape.SHAPE_32X32B, + prims.make_tmem_ptr(base_addr, Float32), + num=64, + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + for reg_idx in cutlass.range_constexpr(64): + s_vals[reg_idx] = loaded[reg_idx].bitcast(Float32) + else: + loaded = tcgen05_ld_16x32bx2_f32( + prims.make_tmem_ptr(base_addr, Float32), + num=cfg.tile_size_kv // 2, + offset=Int32(cfg.tile_size_kv // 2), + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + for reg_idx in cutlass.range_constexpr(num_s_regs_per_thread(cfg)): + s_vals[reg_idx] = loaded[reg_idx] else: loaded0 = prims.tcgen05_ld( TCGEN05_16X256B_SHAPE, @@ -536,11 +611,18 @@ def update_softmax( s_vals[q_repeats * 4 + ld_base + 2] = loaded1[ld_base + 2] s_vals[q_repeats * 4 + ld_base + 3] = loaded1[ld_base + 3] + if cutlass.const_expr(cfg.has_fp8_primary): + scale = self.page_offsets.tensor_scale(tile_offset_k, batch_idx) + if scale != cutlass.Float32(1): + for i in cutlass.range_constexpr(num_s_regs_per_thread(cfg)): + s_vals[i] *= scale + should_apply_dense_mask = ( (seq_len_kv % Int32(cfg.tile_size_kv)) != Int32(0) ) or (next_tile_offset_k > seq_len_kv) needs_row_causal_mask = cutlass.const_expr( - cfg.mask_type == MaskType.CAUSAL.value and cfg.logical_seq_len_q > 1 + cfg.mask_type == MaskType.CAUSAL.value + and (not cfg.fixed_seqlen or cfg.logical_seq_len_q > 1) ) if cutlass.const_expr(needs_row_causal_mask): min_seq_len_kv = runtime_seq_len_kv_for_query_row( @@ -555,7 +637,10 @@ def update_softmax( next_tile_offset_k > min_seq_len_kv ) if cutlass.const_expr(cfg.num_tokens_per_page == 1): - if cutlass.const_expr(cfg.logical_seq_len_q == 1): + if cutlass.const_expr( + (cfg.fixed_seqlen and cfg.logical_seq_len_q == 1) + or cfg.sparse_d576_reload + ): # Cached route bits include the live prefix boundary. should_apply_dense_mask = (mask_routes.low | mask_routes.high) != Int32( 0 @@ -570,6 +655,9 @@ def update_softmax( if cutlass.const_expr(cfg.kernel_variant == "keeps_mma_ab"): local_col_base = (lane_idx >> Int32(4)) * Int32(cfg.tile_size_kv // 2) local_row_idx = warp_idx * Int32(16) + (lane_idx & Int32(0xF)) + if cutlass.const_expr(cfg.ws_native_e): + local_row_idx = warp_grp_thread_idx % Int32(64) + local_col_base = (warp_grp_thread_idx // Int32(64)) * Int32(64) row_seq_len_kv = seq_len_kv if cutlass.const_expr(needs_row_causal_mask): row_seq_len_kv = runtime_seq_len_kv_for_query_row( @@ -728,19 +816,42 @@ def update_softmax( else: for reg_idx in cutlass.range_constexpr(num_s_regs_per_thread(cfg)): local_max = cute.math.max(local_max, s_vals[reg_idx], ftz=True) - local_max = cute.math.max( - local_max, - Float32( - prims.shfl_sync( - thread_mask=0xFFFFFFFF, - val=local_max, - offset=16, - mask_and_clamp=0x1F, - kind=prims.Shfl.BFLY, - ) - ), - ftz=True, - ) + if cutlass.const_expr(cfg.ws_native_e): + scratch = smem_array( + stage_info.context, self._ws_bridge_alloc, Float32, 128 + ) + tid = warp_grp_thread_idx + scratch[tid] = local_max + prims.barrier_cta_sync( + barrier_id=self.sync_barrier_id, thread_count=128 + ) + local_max = cute.math.max(local_max, scratch[tid ^ Int32(64)], ftz=True) + prims.barrier_cta_sync( + barrier_id=self.sync_barrier_id, thread_count=128 + ) + else: + local_max = cute.math.max( + local_max, + Float32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=local_max, + offset=16, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + ), + ftz=True, + ) + if cutlass.const_expr(cfg.softmax_rescale_threshold_log2 > 0.0): + # Keeps/WS uses one row anchor, unlike the grouped Swaps + # layout below. Apply the same opt-in deferral after both + # column halves have contributed their maximum. + old_anchor = old_max_vals[0] + if old_anchor != neg_max_f32(): + increase_log2 = (local_max - old_anchor) * self.scale_softmax_log2 + if increase_log2 <= Float32(cfg.softmax_rescale_threshold_log2): + local_max = old_anchor new_max_vals[0] = local_max self.new_max_state[0] = local_max else: @@ -844,7 +955,7 @@ def update_softmax( candidate = new_max_vals[scale_idx] # Any common normalization anchor gives the same attention. # Match FlashMLA's six-log2 bound (weights up to 64). - # Native FP8 retains exact maxima and scale 448. + # This fixed deferral is for BF16; explicit tuning replaces it. # Initialize normally and update on larger score increases. if old_anchor != neg_max_f32() and self.scale_softmax_log2 > Float32( 0.0 @@ -854,6 +965,20 @@ def update_softmax( if increase_log2 <= limit: new_max_vals[scale_idx] = old_anchor self.new_max_state[scale_idx] = new_max_vals[scale_idx] + elif cutlass.const_expr(cfg.softmax_rescale_threshold_log2 > 0.0): + # Threshold skip correction: keep the previous row max while + # this K tile raises it by at most the compile-time threshold + # (log2 units, like the sparse bound above). P grows by at most + # 2**threshold and the unchanged max lets correction skip the + # O rescale. The first tile always adopts the new max. + old_anchor = old_max_vals[scale_idx] + if old_anchor != neg_max_f32(): + increase_log2 = ( + new_max_vals[scale_idx] - old_anchor + ) * self.scale_softmax_log2 + if increase_log2 <= Float32(cfg.softmax_rescale_threshold_log2): + new_max_vals[scale_idx] = old_anchor + self.new_max_state[scale_idx] = new_max_vals[scale_idx] old_max_arr[scale_idx] = old_max_vals[scale_idx] sum_arr[scale_idx] = sum_vals[scale_idx] new_max_arr[scale_idx] = new_max_vals[scale_idx] diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_softmax_stats.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_softmax_stats.py index f6bec6761305..7a08062f67d6 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_softmax_stats.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/resources/tmem_softmax_stats.py @@ -575,14 +575,16 @@ class SmemSoftmaxLocalResource(TmemSoftmaxLocalResource): """ def get_tmem_requirements(self): + """Keep correction statistics in owned SMEM slots rather than aliasing recyclable score TMEM.""" return [] def get_smem_requirements(self): + """Allocate per-stage old/new statistics for full lanes or compact column-group leaders.""" if self._alloc is None: self._alloc = SmemAllocation( name=self.name + "_stats", size_bytes=self.pipeline_config.num_stages - * 128 + * (16 if self.cfg.compact_softmax_stats else 128) * 2 * num_softmax_scale_groups(self.cfg) * 4, @@ -592,18 +594,24 @@ def get_smem_requirements(self): @cute.jit def _stats_pointer(self, stage_info): + """Address this thread's statistics in the current owned SMEM pipeline slot.""" storage = smem_array( stage_info.context, self._alloc, Float32, self.pipeline_config.num_stages - * 128 + * (16 if self.cfg.compact_softmax_stats else 128) * 2 * num_softmax_scale_groups(self.cfg), ) local_thread = cute.arch.thread_idx()[0] & Int32(127) + threads = 16 if self.cfg.compact_softmax_stats else 128 + if cutlass.const_expr(self.cfg.compact_softmax_stats): + local_thread = (local_thread // Int32(32)) * Int32( + 4 + ) + local_thread % Int32(4) return storage.data_ptr( - (stage_info.stage_idx * Int32(128) + local_thread) + (stage_info.stage_idx * Int32(threads) + local_thread) * Int32(2 * num_softmax_scale_groups(self.cfg)) ) @@ -618,15 +626,33 @@ def store_loop_stats( sum_arr, inst_idx: cutlass.Constexpr[int], ): + """Publish old maxima or reduced sums without reusing the slot before correction consumes it.""" n = num_softmax_scale_groups(self.cfg) - values = tuple( - old_max_arr[i] if inst_idx == 0 else sum_arr[i] for i in range(n) - ) + first = cutlass.Array(Float32, n, space=cutlass.AddressSpace.rmem) + for i in cutlass.range_constexpr(n): + value = old_max_arr[i] if inst_idx == 0 else sum_arr[i] + if cutlass.const_expr(self.cfg.compact_softmax_stats and inst_idx != 0): + # The correction task normally performs these exact reductions. + # Move them before publication so only four lanes per warp store + # partial sums. Maxima already agree within each column group. + for delta in cutlass.range_constexpr(3): + value += Float32( + prims.shfl_sync( + thread_mask=0xFFFFFFFF, + val=value, + offset=16 >> delta, + mask_and_clamp=0x1F, + kind=prims.Shfl.BFLY, + ) + ) + first[i] = value + values = tuple(first[i] for i in range(n)) values += tuple(new_max_arr[i] for i in range(n)) - self._stats_pointer(stage_info).store( - vector_from_scalars(values, Float32), - alignment=16 if n > 1 else 8, - ) + if not self.cfg.compact_softmax_stats or cute.arch.thread_idx()[0] % 32 < 4: + self._stats_pointer(stage_info).store( + vector_from_scalars(values, Float32), + alignment=16 if n > 1 else 8, + ) cute.arch.fence_view_async_shared() @cute.jit diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/tasks.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/tasks.py index 1d75ee237638..1619c3d47f35 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/tasks.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/mla_decode/throughput_latency_1cta/tasks.py @@ -46,6 +46,7 @@ ) from ..helpers.stage import MlaStage from ..helpers.tile import ( + runtime_split_count, runtime_base_seq_len_kv, runtime_local_kv_tiles, runtime_seq_len_kv_from_task_cache, @@ -58,6 +59,7 @@ class MlaDecodeTask(Task): def __init__(self, **kwargs): """Initialize MLA-specific task parameters and cached thread indices.""" self.cfg = kwargs.pop("cfg", None) + self.runtime_split_kv = kwargs.pop("runtime_split_kv", None) self.seqlens_kv = kwargs.pop("seqlens_kv", None) self.cu_seqlens_q = kwargs.pop("cu_seqlens_q", None) self.domain_bias = kwargs.pop("domain_bias", 0) @@ -67,6 +69,7 @@ def __init__(self, **kwargs): self._local_warp_idx = Int32(0) self._lane_idx = Int32(0) self._seq_len_kv = Int32(0) + self._local_kv_tiles = Int32(1) @cute.jit def init_variables(self, context=None): @@ -116,7 +119,7 @@ def make_task_cache(self): self._local_warp_idx, self._lane_idx, self._seq_len_kv, - Int32(0), + self._local_kv_tiles, Int32(0), Int32(0), ) @@ -126,6 +129,9 @@ def get_domain(self, tile_coord): """Return the runtime loop domain for the current 1CTA work tile.""" if cutlass.const_expr(self.cfg is None): return self.domain + if cutlass.const_expr(self.seqlens_kv is None and self.cfg.seq_len_kv is None): + # Padding and scheduler tasks have no K/V metadata or mainloop. + return self.domain # Persistent 1CTA coordinates combine batch and head tile in z; # non-persistent grids keep z as batch. cache_seqs remains batch-owned. @@ -145,13 +151,23 @@ def get_domain(self, tile_coord): # raw batch length for causal can move an all-masked tile into tail and # replace a real loop accumulator at a tile boundary. cta_idx_q = tile_coord[0] - if cutlass.const_expr(self.cfg.use_persistent_scheduler != 1): - # Nonpersistent grid X combines (head tile, Q tile, KV split), with - # KV split innermost. Persistent work queues already expose Q as - # tile coordinate 0 and must not be decoded a second time. - cta_idx_q = (tile_coord[0] // Int32(self.cfg.num_ctas_per_seq_kv)) % Int32( - self.cfg.num_ctas_per_seq_q - ) + if cutlass.const_expr( + self.cfg.fixed_seqlen + and self.cfg.use_persistent_scheduler != 1 + and self.cfg.num_ctas_per_seq_q == 1 + ): + cta_idx_q = Int32(0) + elif cutlass.const_expr(self.cfg.use_persistent_scheduler != 1): + grid_splits = Int32(self.cfg.num_ctas_per_seq_kv) + if cutlass.const_expr( + self.runtime_split_kv is not None + and self.cfg.use_cluster_reduction != 1 + ): + grid_splits = runtime_split_count(self.runtime_split_kv) + # Heads are already flattened into Q; grid-X packs Q and KV split. + cta_idx_q = tile_coord[0] // grid_splits + if cutlass.const_expr(self.cfg.fixed_seqlen): + cta_idx_q = cta_idx_q % Int32(self.cfg.num_ctas_per_seq_q) # Cache the raw batch length once per task/work tile. Resource work # methods derive their CTA-visible length from this task cache instead # of independently reloading cache_seqs throughout the hot loop. @@ -169,7 +185,10 @@ def get_domain(self, tile_coord): ) if cutlass.const_expr(self.cfg.fixed_sparse_kv_tiles > 0): return Int32(self.domain_bias) - total_kv_tiles = runtime_local_kv_tiles(self.cfg, seq_len_kv) + total_kv_tiles = runtime_local_kv_tiles( + self.cfg, seq_len_kv, requested_splits=self.runtime_split_kv + ) + self._local_kv_tiles = total_kv_tiles remaining_kv_tiles = cute.math.max( total_kv_tiles - Int32(self.cfg.num_insts_kv), Int32(0) ) @@ -1096,6 +1115,7 @@ def page_offsets_schedule(smem_page_offsets, work_queue=None): smem_page_offsets.init_load_state() with work_tile_schedule_loop(work_queue, skip_if=work_tile_skip_if): + smem_page_offsets.bind_sparse_request() if cfg.single_kv_pipe: page_offsets_produce( smem_page_offsets, "load_k0", section=MlaStage.Head @@ -1257,6 +1277,7 @@ def load_kv( else: cached_page_ids = staged_kv_load( smem_kv, + whole_kv=cfg.whole_kv_tile, head_dim_stages=head_dim_stages, producer_label=producer_label, section=section, @@ -1492,7 +1513,7 @@ def load_page_offsets_wq_throttle_schedule( warp_idx=cfg.load_warp_idx if warp_idx is None else warp_idx, num_warps=cfg.load_num_warps if num_warps is None else num_warps, schedule=captured_schedule, - num_registers=cfg.mma_load_regs, + num_registers=cfg.kv_load_regs, name="LoadTmaTask", **kw, ) @@ -1633,6 +1654,8 @@ def mma_schedule( smem_kv.init_descriptor_state() smem_p0.init_descriptor_state() smem_p1.init_descriptor_state() + tmem_s0.init_descriptor_state() + tmem_s1.init_descriptor_state() tmem_o.init_mma_state() with work_tile_schedule_loop(work_queue, skip_if=work_tile_skip_if): @@ -1659,6 +1682,7 @@ def mma_schedule( head_dim_stages=cfg.v_head_dim_stages, consumer_label="v_desc_reused_0", reuse_k=True, + retained_k_stages=cfg.qk_head_dim_stages, producer_label="pv_mma_loop_0", ) staged_qk_mma( @@ -1675,6 +1699,7 @@ def mma_schedule( head_dim_stages=cfg.v_head_dim_stages, consumer_label="v_desc_reused_1", reuse_k=True, + retained_k_stages=cfg.qk_head_dim_stages, producer_label="pv_mma_loop_1", ) staged_qk_mma( @@ -1691,6 +1716,7 @@ def mma_schedule( head_dim_stages=cfg.v_head_dim_stages, consumer_label="v_desc_reused_0", reuse_k=True, + retained_k_stages=cfg.qk_head_dim_stages, producer_label="pv_mma_tail_0", ) staged_qk_mma( @@ -1707,6 +1733,7 @@ def mma_schedule( head_dim_stages=cfg.v_head_dim_stages, consumer_label="v_desc_reused_1", reuse_k=True, + retained_k_stages=cfg.qk_head_dim_stages, producer_label="pv_mma_tail_1", ) else: @@ -1735,6 +1762,7 @@ def mma_schedule( "v_desc_reused" if cfg.sparse_reuse_kv else "v_desc_0" ), reuse_k=cfg.sparse_reuse_kv, + retained_k_stages=cfg.qk_head_dim_stages, producer_label="pv_mma_loop_0", ) staged_qk_mma( @@ -1755,6 +1783,7 @@ def mma_schedule( "v_desc_reused" if cfg.sparse_reuse_kv else "v_desc_1" ), reuse_k=cfg.sparse_reuse_kv, + retained_k_stages=cfg.qk_head_dim_stages, producer_label="pv_mma_loop_1", ) staged_qk_mma( @@ -1774,6 +1803,7 @@ def mma_schedule( "v_desc_reused" if cfg.sparse_reuse_kv else "v_desc_0" ), reuse_k=cfg.sparse_reuse_kv, + retained_k_stages=cfg.qk_head_dim_stages, producer_label="pv_mma_tail_0", ) staged_pv_mma( @@ -1785,6 +1815,7 @@ def mma_schedule( "v_desc_reused" if cfg.sparse_reuse_kv else "v_desc_1" ), reuse_k=cfg.sparse_reuse_kv, + retained_k_stages=cfg.qk_head_dim_stages, producer_label="pv_mma_tail_1", ) smem_q.release() @@ -1849,6 +1880,9 @@ def create_single_kv_pipe_mma_task( work_tile_skip_if = runtime_work_tile_skip_if(work_queue) swap = cfg.one_insts_kv_swap pv_mma = staged_pv_mma if swap else staged_pv_mma_tmem_p + pv_kwargs = {"whole_kv": cfg.whole_kv_tile} + if not swap: + pv_kwargs["ws_pair"] = cfg.ws_pv pv_loop = "pv_mma_loop_0" if swap else "pv_mma_loop_tmem_p" pv_tail = "pv_mma_tail_0" if swap else "pv_mma_tail_tmem_p" v_desc = "v_desc_reused_0" if swap else "v_desc_reused" @@ -1856,8 +1890,8 @@ def create_single_kv_pipe_mma_task( v_desc = "v_desc_0" pv_before_qk = ( cfg.sparse_reuse_kv - and not cfg.is_fp8_qkv() - and (swap or cfg.tile_size_kv == 128) + and not cfg.whole_kv_tile + and cfg.kv_stages < 2 * cfg.qk_head_dim_stages ) @schedule @@ -1872,6 +1906,7 @@ def mma_schedule( """Capture identical pipeline ownership for the two P layouts.""" smem_q.init_descriptor_state() smem_kv.init_descriptor_state() + tmem_s.init_descriptor_state() tmem_o.init_mma_state() if swap: p_resource.init_descriptor_state() @@ -1890,6 +1925,7 @@ def mma_schedule( smem_kv, tmem_s, head_dim_stages=cfg.qk_head_dim_stages, + whole_kv=cfg.whole_kv_tile, consumer_label="k_desc_0", release_k=not cfg.sparse_reuse_kv, ) @@ -1902,12 +1938,15 @@ def mma_schedule( head_dim_stages=cfg.v_head_dim_stages, consumer_label=v_desc, reuse_k=True, + retained_k_stages=cfg.qk_head_dim_stages, producer_label=pv_loop, + **pv_kwargs, ) staged_qk_mma( smem_kv, tmem_s, head_dim_stages=cfg.qk_head_dim_stages, + whole_kv=cfg.whole_kv_tile, consumer_label="k_desc_0", release_k=not cfg.sparse_reuse_kv, ) @@ -1919,7 +1958,9 @@ def mma_schedule( head_dim_stages=cfg.v_head_dim_stages, consumer_label=v_desc, reuse_k=cfg.sparse_reuse_kv, + retained_k_stages=cfg.qk_head_dim_stages, producer_label=pv_loop, + **pv_kwargs, ) pv_mma( smem_kv, @@ -1928,7 +1969,9 @@ def mma_schedule( head_dim_stages=cfg.v_head_dim_stages, consumer_label=v_desc, reuse_k=cfg.sparse_reuse_kv, + retained_k_stages=cfg.qk_head_dim_stages, producer_label=pv_tail, + **pv_kwargs, ) smem_q.release() diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/tensor_map.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/tensor_map.py index 1501d57317e3..b88ab4c5923b 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/tensor_map.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/tensor_map.py @@ -22,6 +22,7 @@ from typing import Sequence, Tuple, Type import cutlass.cute as cute +from cutlass import Pointer, Uint64 from cutlass.cute import depth, leading_dim from cutlass.cute.testing import assert_ as runtime_assert from cutlass.cutlass_dsl import ( @@ -280,6 +281,7 @@ def create_tensor_map_ragged_from_tensor( swizzle: TensorMapSwizzle | None = None, l2_promotion: TensorMapL2Promotion | None = None, oob_fill: TensorMapFloatOOBFill | None = None, + use_blackwell_tma_oob_workaround: bool = False, tma_format: TensorMapDataType | None = None, loc=None, ip=None, @@ -316,6 +318,8 @@ def create_tensor_map_ragged_from_tensor( :param l2_promotion: L2 promotion hint, defaults to ``none``. :param oob_fill: OOB fill mode; see :func:`create_tensor_map_ragged` for the format-aware default. + :param use_blackwell_tma_oob_workaround: Disable speculative OOB addressing + when the ragged view cannot establish a mapped 128 KiB prefix. :param tma_format: Override the element data type, defaults to auto-detect from ``tensor.element_type``. :raises ValueError: For unsupported rank (``R ∉ {2, 3}``), @@ -418,7 +422,7 @@ def create_tensor_map_ragged_from_tensor( else get_dsl_type_to_tensormap_type(tensor.element_type) ) - return create_tensor_map_ragged( + descriptor = create_tensor_map_ragged( global_address=tensor.iterator.toint(), tma_format=fmt, global_dims=tma_dims, @@ -432,6 +436,11 @@ def create_tensor_map_ragged_from_tensor( loc=loc, ip=ip, ) + if use_blackwell_tma_oob_workaround: + # Match context's Blackwell descriptor-bit-85 workaround for short views. + word = Pointer(descriptor.get_ptr(), dtype=Uint64) + 1 + word.store(word.load() & ~Uint64(1 << 21)) + return descriptor def transform_ragged_coords( diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py b/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py index 7a305b1ec6f2..e60454d40ec7 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py @@ -23,14 +23,24 @@ import functools from typing import Any, Literal, Optional, cast +import torch + from ._compile_cache import build_and_load_attention_kernel from ._hardware import get_max_active_clusters -import torch + +from .kernels.mla_decode.helpers.nvfp4_kv import ( + cache_record_size, + validate_nvfp4_kv_formats, +) from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api +from flashinfer.utils import get_device_sm_count + from .decode import ( + prepare_lse, + validate_lse_mode, _WorkspaceSection, _align_up, _append_workspace_section, @@ -40,14 +50,18 @@ _validate_mask, _validate_page_size, _validate_positive_int, - _validate_runtime_device, _validate_scale, + _validate_runtime_device, _validate_workspace_buffer, _workspace_section_view, + resolve_max_sequence_length, ) -from .kernels.mla_decode.sparse_policy import select_sparse_mla_profile +from .kernels.mla_decode.sparse_policy import ( + prefer_sparse_d576_reload, + select_sparse_mla_profile, +) _COMPILE_OPTIONS = "--enable-tvm-ffi --opt-level 2" _MLA_LATENT_DIM = 512 @@ -56,11 +70,11 @@ _SUPPORTED_INPUT_DTYPES = (torch.bfloat16, torch.float8_e4m3fn) _SUPPORTED_OUTPUT_DTYPES = (torch.bfloat16,) _INT32_MAX = 2**31 - 1 -# The largest public 1CTA schedule can pad one K/V split group across 128 -# splits, two 128-token K/V instructions apiece. Reserve that complete span so -# every padded tile boundary remains representable as signed Int32. -_MLA_MAX_KV_COORDINATE_SPAN = 128 * 2 * 128 -_MLA_MAX_KV_LEN = _INT32_MAX - (_MLA_MAX_KV_COORDINATE_SPAN - 1) +# Retain the existing broad input guard; after policy resolution, validate +# the complete padded endpoint from its actual capacity (including >128). +# This legacy coordinate reserve is not a split-count ceiling. +_MLA_LEGACY_KV_COORDINATE_RESERVE = 128 * 2 * 128 +_MLA_MAX_KV_LEN = _INT32_MAX - (_MLA_LEGACY_KV_COORDINATE_RESERVE - 1) @dataclass(frozen=True) @@ -79,17 +93,20 @@ class _MLADecodeCompileSpec: device_index: int kernel_signature: tuple[object, ...] - num_heads: int + num_qo_heads: int kv_lora_rank: int qk_rope_head_dim: int page_size: int q_dtype_key: str - output_dtype_key: str - max_seq_len_q: int + out_dtype_key: str + max_q_len: int | None packed_query: bool has_kernel_workspace: bool - split_kv: int + # Active split count is a runtime operand; kernel_signature carries capacity. + split_kv: int = field(compare=False, hash=False) kernel: Any = field(compare=False, hash=False, repr=False) + write_final_lse: bool = True + lse_mode: str = "none" @dataclass(frozen=True) @@ -104,7 +121,7 @@ class _MLAWorkspaceLayout: @dataclass(frozen=True) class _MLAWorkspaceViews: kernel_workspace: Optional[torch.Tensor] - lse: torch.Tensor + lse: Optional[torch.Tensor] @dataclass(frozen=True) @@ -113,15 +130,15 @@ class _MLADecodePlanState: device: torch.device batch_size: int - num_heads: int - max_seq_len_q: int + num_qo_heads: int + max_q_len: int packed_query: bool kv_lora_rank: int qk_rope_head_dim: int page_size: int q_dtype: torch.dtype kv_dtype: torch.dtype - output_dtype: torch.dtype + out_dtype: torch.dtype mask_type: Literal["dense", "causal"] max_kv_len: int required_page_columns: int @@ -131,6 +148,31 @@ class _MLADecodePlanState: compiled: Callable[..., object] policy: tuple[tuple[str, object], ...] split_kv: int + softmax_rescale_threshold_log2: float = 0.0 + lse_mode: str = "none" + + +@dataclass(frozen=True) +class MLAAttentionTuning: + """Compile-time tuning of one PrimTS MLA attention plan. + + It applies to every query length the MLA wrappers serve, context + (multi-token Q) as well as generation. ``softmax_rescale_threshold_log2`` + is in binades. ``None`` keeps the entry point's default: exact rescale for + dense, and each schedule's existing deferral for sparse. An explicit + value, including ``0.0``, is compiled into the kernel; see + :meth:`BatchMLADecodePagedTSWrapper.plan`. + """ + + softmax_rescale_threshold_log2: Optional[float] = None + + +def _tuning_rescale_threshold(tuning: object) -> Optional[float]: + if tuning is None: + return None + if not isinstance(tuning, MLAAttentionTuning): + raise TypeError("tuning must be an MLAAttentionTuning or None") + return tuning.softmax_rescale_threshold_log2 @dataclass(frozen=True) @@ -139,24 +181,29 @@ class _MLARuntime: normalized_cache: torch.Tensor out: torch.Tensor num_physical_pages: int - bmm1_scale: float - bmm2_scale: float + softmax_scale: float + output_scale: float extra_cache: Optional[torch.Tensor] = None sparse_inputs: Optional[tuple] = None output_quant_args: Optional[tuple] = None +# Pack kernel scratch bytes and [batch, max Q tokens, heads] FP32 LSE storage. def _make_mla_workspace_layout( kernel_workspace_bytes: int, batch_size: int, - num_heads: int, - max_seq_len_q: int = 1, + num_qo_heads: int, + max_q_len: int = 1, + *, + include_lse: bool = True, ) -> _MLAWorkspaceLayout: kernel_workspace, byte_end = _append_workspace_section( 0, (kernel_workspace_bytes,), torch.int8 ) lse, byte_end = _append_workspace_section( - byte_end, (batch_size, max_seq_len_q, num_heads), torch.float32 + byte_end, + (batch_size, max_q_len, num_qo_heads) if include_lse else (0,), + torch.float32, ) return _MLAWorkspaceLayout( kernel_workspace=kernel_workspace, @@ -175,7 +222,9 @@ def _bind_mla_workspace( ) return _MLAWorkspaceViews( kernel_workspace=kernel_workspace, - lse=_workspace_section_view(workspace_buffer, layout.lse), + lse=_workspace_section_view(workspace_buffer, layout.lse) + if layout.lse.byte_size + else None, ) @@ -223,47 +272,48 @@ def _validate_mla_nonnegative_int32_extent(value: int, name: str) -> int: def _validate_mla_query_head_extent( *, batch_size: int, - num_heads: int, - max_seq_len_q: int, + num_qo_heads: int, + max_q_len: int, total_q: Optional[int] = None, ) -> None: """Keep fixed-capacity and packed query-head coordinates in signed Int32.""" _validate_mla_int32_extent( - batch_size * max_seq_len_q * num_heads, - "batch_size * max_seq_len_q * num_heads", + batch_size * max_q_len * num_qo_heads, + "batch_size * max_q_len * num_qo_heads", ) if total_q is not None: _validate_mla_nonnegative_int32_extent( - total_q * num_heads, - "total_q * num_heads", + total_q * num_qo_heads, + "total_q * num_qo_heads", ) def _validate_mla_policy_coordinate_span( policy: tuple[tuple[str, object], ...], + *, + max_kv_len: int = 1, ) -> None: """Keep the host K/V bound coupled to the automatically selected policy.""" resolved = dict(policy) span = ( int(cast(int, resolved["tile_size_kv"])) * int(cast(int, resolved["num_insts_kv"])) - * max(int(cast(int, resolved["split_kv"])), 1) + * max(int(cast(int, resolved.get("split_capacity", resolved["split_kv"]))), 1) ) - if span > _MLA_MAX_KV_COORDINATE_SPAN: - raise RuntimeError( - "MLA Int32 extent safety assumes a padded K/V coordinate span no " - f"larger than {_MLA_MAX_KV_COORDINATE_SPAN}, got {span}" + if max_kv_len + span - 1 > _INT32_MAX: + raise NotImplementedError( + f"MLA padded K/V extent must fit signed Int32: max_kv_len={max_kv_len}, span={span}" ) def _validate_mla_dtype_pair( q_dtype: torch.dtype, kv_dtype: torch.dtype, - output_dtype: torch.dtype, + out_dtype: torch.dtype, ) -> None: _dtype_key(q_dtype) _dtype_key(kv_dtype) - _dtype_key(output_dtype) + _dtype_key(out_dtype) if q_dtype != kv_dtype: raise NotImplementedError( "attention-ts MLA decode requires query and KV cache to use the " @@ -273,10 +323,10 @@ def _validate_mla_dtype_pair( raise NotImplementedError( f"attention-ts MLA decode supports BF16 and FP8-E4M3 input; got {q_dtype}" ) - if output_dtype not in _SUPPORTED_OUTPUT_DTYPES: + if out_dtype not in _SUPPORTED_OUTPUT_DTYPES: raise NotImplementedError( "attention-ts MLA decode currently supports BF16 output only; " - f"got {output_dtype}" + f"got {out_dtype}" ) @@ -304,7 +354,7 @@ def _validate_int32_cuda_tensor( def _validate_mla_metadata( block_tables: torch.Tensor, - seq_lens: torch.Tensor, + seq_lens_kv: torch.Tensor, ) -> tuple[torch.device, int, int]: _validate_int32_cuda_tensor( block_tables, @@ -314,16 +364,16 @@ def _validate_mla_metadata( require_16byte_alignment=False, ) _validate_int32_cuda_tensor( - seq_lens, - "seq_lens", + seq_lens_kv, + "seq_lens_kv", ndim=1, require_16byte_alignment=False, ) - if block_tables.device != seq_lens.device: - raise ValueError("block_tables and seq_lens must be on the same device") - batch_size = int(seq_lens.numel()) + if block_tables.device != seq_lens_kv.device: + raise ValueError("block_tables and seq_lens_kv must be on the same device") + batch_size = int(seq_lens_kv.numel()) if batch_size <= 0: - raise ValueError("seq_lens must contain at least one request") + raise ValueError("seq_lens_kv must contain at least one request") if block_tables.shape[0] != batch_size: raise ValueError( "block_tables must have one row per request: expected " @@ -341,11 +391,11 @@ def _validate_mla_metadata( ) if block_tables.data_ptr() % 4 != 0: raise ValueError("block_tables must be 4-byte aligned") - if seq_lens.data_ptr() % 4 != 0: - raise ValueError("seq_lens must be 4-byte aligned") + if seq_lens_kv.data_ptr() % 4 != 0: + raise ValueError("seq_lens_kv must be 4-byte aligned") _validate_mla_int32_extent(batch_size, "batch_size") _validate_mla_int32_extent(int(block_tables.numel()), "block_tables elements") - return seq_lens.device, batch_size, max_num_pages + return seq_lens_kv.device, batch_size, max_num_pages def _validate_qo_indptr( @@ -372,6 +422,15 @@ def _validate_qo_indptr( ) +def has_runtime_query_lengths(kernel) -> bool: + """Whether the kernel takes its Q length at launch. + + Sparse MLA flattens Q rows into virtual batch items, so its physical Q length + stays compile-time even when the KV lengths are runtime operands. + """ + return not kernel.fixed_seqlen and getattr(kernel, "sparse_profile", None) is None + + def _derive_max_seq_len_q( qo_indptr: torch.Tensor, *, @@ -399,18 +458,18 @@ def _derive_max_seq_len_q( def _validate_mla_run_metadata( runtime: _MLARuntime, block_tables: torch.Tensor, - seq_lens: torch.Tensor, + seq_lens_kv: torch.Tensor, qo_indptr: Optional[torch.Tensor], *, device: torch.device, batch_size: int, required_page_columns: int, packed_query: bool, - max_seq_len_q: int, + max_q_len: int, max_kv_len: int, page_size: int, mask_type: str, -) -> None: +) -> int: """Validate per-run request metadata against one static MLA plan. Value checks synchronize CUDA metadata with the host. ``run(validate=False)`` @@ -418,7 +477,7 @@ def _validate_mla_run_metadata( """ metadata_device, metadata_batch_size, max_num_pages = _validate_mla_metadata( - block_tables, seq_lens + block_tables, seq_lens_kv ) if metadata_device != device: raise ValueError( @@ -452,20 +511,19 @@ def _validate_mla_run_metadata( "qo_indptr must end at the packed query row count " f"({runtime.query.shape[0]}), got {total_q}" ) - if runtime_max_seq_len_q > max_seq_len_q: + if runtime_max_seq_len_q > max_q_len: raise ValueError( "qo_indptr contains a per-request Q length larger than " - f"max_seq_len_q ({max_seq_len_q}): got " + f"max_q_len ({max_q_len}): got " f"{runtime_max_seq_len_q}" ) else: if qo_indptr is not None: raise ValueError("fixed-query MLA plan does not accept qo_indptr") - q_lengths = (max_seq_len_q,) * batch_size + runtime_max_seq_len_q = int(runtime.query.shape[1]) + q_lengths = (runtime_max_seq_len_q,) * batch_size - seq_lens_host = tuple(int(value) for value in seq_lens.tolist()) - if any(seq_len <= 0 for seq_len in seq_lens_host): - raise ValueError("every runtime request must contain at least one KV token") + seq_lens_host = tuple(int(value) for value in seq_lens_kv.tolist()) runtime_max_kv_len = max(seq_lens_host) if runtime_max_kv_len > max_kv_len: raise ValueError( @@ -476,7 +534,7 @@ def _validate_mla_run_metadata( for request_idx, (row, seq_len) in enumerate( zip(block_table_rows, seq_lens_host, strict=True) ): - required_pages = _ceil_div(seq_len, page_size) + required_pages = _ceil_div(max(seq_len, 0), page_size) if any( int(page_id) < 0 or int(page_id) >= runtime.num_physical_pages for page_id in row[:required_pages] @@ -490,18 +548,19 @@ def _validate_mla_run_metadata( for request_idx, (q_len, kv_len) in enumerate( zip(q_lengths, seq_lens_host, strict=True) ): - if q_len > kv_len: + if kv_len > 0 and q_len > kv_len: raise ValueError( "causal MLA decode requires every per-request Q length " "to be no greater than its K/V length; request " f"{request_idx} has Q={q_len} and K/V={kv_len}" ) + return runtime_max_seq_len_q def _resolve_max_seq_len_q_alias( *, seq_len_q: Optional[int], - max_seq_len_q: Optional[int], + max_q_len: Optional[int], default: Optional[int], ) -> Optional[int]: """Resolve the legacy fixed-Q name and the explicit static-bound name.""" @@ -512,8 +571,8 @@ def _resolve_max_seq_len_q_alias( else None ) explicit_bound = ( - _validate_positive_int(max_seq_len_q, "max_seq_len_q") - if max_seq_len_q is not None + _validate_positive_int(max_q_len, "max_q_len") + if max_q_len is not None else None ) if ( @@ -522,7 +581,7 @@ def _resolve_max_seq_len_q_alias( and legacy_bound != explicit_bound ): raise ValueError( - "seq_len_q and max_seq_len_q must agree when both are provided: " + "seq_len_q and max_q_len must agree when both are provided: " f"got {legacy_bound} and {explicit_bound}" ) if explicit_bound is not None: @@ -532,14 +591,15 @@ def _resolve_max_seq_len_q_alias( return default +# Validate compact fixed or packed CUDA Q storage against optional plan bounds. def _validate_query( query: torch.Tensor, *, packed_query: bool = False, device: Optional[torch.device] = None, batch_size: Optional[int] = None, - num_heads: Optional[int] = None, - max_seq_len_q: Optional[int] = None, + num_qo_heads: Optional[int] = None, + max_q_len: Optional[int] = None, q_dtype: Optional[torch.dtype] = None, query_dim: int = _MLA_QUERY_DIM, ) -> None: @@ -575,31 +635,31 @@ def _validate_query( f"query batch size must match the plan ({batch_size}), got {query.shape[0]}" ) head_axis = 1 if packed_query else 2 - if num_heads is not None and query.shape[head_axis] != num_heads: + if num_qo_heads is not None and query.shape[head_axis] != num_qo_heads: raise ValueError( "query head count must match the plan " - f"({num_heads}), got {query.shape[head_axis]}" + f"({num_qo_heads}), got {query.shape[head_axis]}" ) - if max_seq_len_q is not None: + if max_q_len is not None: if packed_query: if batch_size is None: raise ValueError("batch_size is required to validate packed query") total_q = int(query.shape[0]) - if total_q > batch_size * max_seq_len_q: + if total_q > batch_size * max_q_len: raise ValueError( "packed query total rows must be within " - f"[0, {batch_size * max_seq_len_q}], got {total_q}" + f"[0, {batch_size * max_q_len}], got {total_q}" ) - elif query.shape[1] != max_seq_len_q: + elif query.shape[1] != max_q_len: raise ValueError( - "fixed query length must equal the planned max_seq_len_q " - f"({max_seq_len_q}), got {query.shape[1]}" + "fixed query length must equal the planned max_q_len " + f"({max_q_len}), got {query.shape[1]}" ) - if batch_size is not None and num_heads is not None: + if batch_size is not None and num_qo_heads is not None: _validate_mla_query_head_extent( batch_size=batch_size, - num_heads=num_heads, - max_seq_len_q=max_seq_len_q, + num_qo_heads=num_qo_heads, + max_q_len=max_q_len, total_q=int(query.shape[0]) if packed_query else None, ) if q_dtype is not None and query.dtype != q_dtype: @@ -655,16 +715,17 @@ def _normalize_mla_kv_cache( return normalized, int(normalized.shape[0]), int(normalized.shape[1]) +# Validate caller-owned output shape, dtype, device and 16-byte alignment. def _validate_out( out: torch.Tensor, *, device: torch.device, batch_size: int, - num_heads: int, - max_seq_len_q: int, + num_qo_heads: int, + max_q_len: int, packed_query: bool, total_q: Optional[int] = None, - output_dtype: torch.dtype, + out_dtype: torch.dtype, ) -> None: if not isinstance(out, torch.Tensor): raise TypeError("out must be a torch.Tensor") @@ -672,15 +733,15 @@ def _validate_out( if total_q is None: raise ValueError("total_q is required to validate packed output") expected_shape: tuple[int, ...] - expected_shape = (total_q, num_heads, _MLA_LATENT_DIM) + expected_shape = (total_q, num_qo_heads, _MLA_LATENT_DIM) else: - expected_shape = (batch_size, max_seq_len_q, num_heads, _MLA_LATENT_DIM) + expected_shape = (batch_size, max_q_len, num_qo_heads, _MLA_LATENT_DIM) if out.shape != expected_shape: raise ValueError( f"out must have shape {expected_shape}, got {tuple(out.shape)}" ) - if out.dtype != output_dtype: - raise ValueError(f"out must have dtype {output_dtype}, got {out.dtype}") + if out.dtype != out_dtype: + raise ValueError(f"out must have dtype {out_dtype}, got {out.dtype}") if out.device != device: raise ValueError(f"out must be on {device}, got {out.device}") if not out.is_contiguous(): @@ -689,6 +750,28 @@ def _validate_out( _validate_16byte_alignment(out, "out") +def _check_softmax_rescale_threshold_log2_range( + threshold: object, *, q_dtype: torch.dtype +) -> float: + """Return the threshold as float32-rounded ``float``; raise if out of range. + + Values outside ``[0, bound]`` (E4M3: 8, BF16: 32) raise instead of being + clamped, as trtllm-gen's ``checkSkipCorrThreshold`` does. + """ + + from .kernels.mla_decode.helpers.math import ( + canonicalize_softmax_rescale_threshold_log2, + ) + + if threshold == 0.0 and type(threshold) in (int, float): + # Zero selects the default kernel for every dtype, including ones + # without a threshold bound. + return 0.0 + return canonicalize_softmax_rescale_threshold_log2( + threshold, qkv_dtype=_kernel_dtype_name(_dtype_key(q_dtype)) + ) + + def _kernel_dtype_name(dtype_key: str) -> str: names = { "bfloat16": "bf16", @@ -724,23 +807,27 @@ def _separate_reducer_provenance( return "reference", 1 +# CPU policies are small and tensor-free. Cache them across layers and decode +# steps; native compilation remains cached independently by its static signature. @functools.cache def _resolve_mla_decode_launch_spec( device_index: int, batch_size: int, - num_heads: int, + num_qo_heads: int, kv_lora_rank: int, qk_rope_head_dim: int, page_size: int, max_kv_len: int, q_dtype_key: str, kv_dtype_key: str, - output_dtype_key: str, + out_dtype_key: str, mask_type: str, seq_len_q: int = 1, + softmax_rescale_threshold_log2: float = 0.0, ): """Resolve and cache MLA policy/workspace without compiling.""" + policy: tuple[tuple[str, Any], ...] max_kv_len = _validate_mla_max_kv_len(max_kv_len, "max_kv_len") import cutlass @@ -760,6 +847,7 @@ def _resolve_mla_decode_launch_spec( q_tile_work_count, resolve_auto_flat_query_launch_shape, resolve_runtime_cluster_reduction_mode, + resolve_auto_cluster_reduction_config, select_auto_split_kv, ) from .kernels.mla_decode.throughput_latency_1cta.kernel import ( @@ -771,18 +859,18 @@ def _resolve_mla_decode_launch_spec( seq_len_q = _validate_positive_int(seq_len_q, "seq_len_q") _validate_mla_query_head_extent( batch_size=batch_size, - num_heads=num_heads, - max_seq_len_q=seq_len_q, + num_qo_heads=num_qo_heads, + max_q_len=seq_len_q, ) _validate_mla_dims(kv_lora_rank, qk_rope_head_dim) qkv_dtype_name = _kernel_dtype_name(q_dtype_key) - output_dtype_name = _kernel_dtype_name(output_dtype_key) + output_dtype_name = _kernel_dtype_name(out_dtype_key) with torch.cuda.device(device_index): max_active_one_cta_clusters = get_max_active_clusters(1) max_active_two_cta_clusters = get_max_active_clusters(2) one_cta_launch_shape = resolve_auto_flat_query_launch_shape( - num_heads_q=num_heads, + num_heads_q=num_qo_heads, seq_len_q=seq_len_q, ) one_cta_base_work = q_tile_work_count( @@ -797,7 +885,9 @@ def _resolve_mla_decode_launch_spec( base_work=one_cta_base_work, target_work=max_active_one_cta_clusters, ) - two_cta_launch_shape = FlatQueryTileLayout.for_tile(num_heads, seq_len_q, 128) + two_cta_launch_shape = FlatQueryTileLayout.for_tile( + num_qo_heads, seq_len_q, 128 + ) two_cta_split_kv = compute_split_kv( batch_size=batch_size, num_q_tiles=two_cta_launch_shape.num_tiles, @@ -805,9 +895,12 @@ def _resolve_mla_decode_launch_spec( mma_qk_tiler_mn=(128, 128), max_active_blocks=max_active_two_cta_clusters * 2, ) + sm_count = get_device_sm_count(torch.device("cuda", device_index)) + one_cta_split_kv = min(one_cta_split_kv, sm_count) + two_cta_split_kv = min(two_cta_split_kv, sm_count) requested_policy, policy_source = resolve_mla_kernel_policy( None, - num_heads, + num_qo_heads, seq_len_q, one_cta_split_kv=one_cta_split_kv, two_cta_split_kv=two_cta_split_kv, @@ -843,10 +936,24 @@ def _resolve_mla_decode_launch_spec( if use_throughput_latency: assert decision.config is not None - reduction_mode = resolve_runtime_cluster_reduction_mode( + # Cluster fanout is physical topology, not a rounded scratch bound. + prepared_cfg = resolve_auto_cluster_reduction_config( decision.config, reduction_mode=None, ) + reduction_mode = ( + "gmem_separate" + if decision.config.use_cluster_reduction + and not prepared_cfg.use_cluster_reduction + else resolve_runtime_cluster_reduction_mode( + prepared_cfg, reduction_mode=None + ) + ) + capacity = ( + one_cta_split_kv + if prepared_cfg.use_cluster_reduction or one_cta_split_kv == 1 + else sm_count + ) kernel = ThroughputLatencyMlaDecodeTs( batch_size=batch_size, num_heads=launch_shape.num_heads_q, @@ -862,15 +969,19 @@ def _resolve_mla_decode_launch_spec( out_dtype=output_dtype_name, profile=decision.profile_name, reduction_mode=reduction_mode, - logical_num_heads=num_heads, + logical_num_heads=num_qo_heads, logical_seq_len_q=seq_len_q, tile_size_q=launch_shape.tile_size_q, explicit_split_kv=None, + split_kv_capacity=capacity, + reduction_policy_splits=one_cta_split_kv, explicit_persistent=None, mask_type=mask_type, + softmax_rescale_threshold_log2=softmax_rescale_threshold_log2, ) + kernel.runtime_split_kv = True final_cfg = kernel._make_config() - split_kv = int(final_cfg.num_ctas_per_seq_kv) + split_kv = one_cta_split_kv workspace_size = compute_1cta_workspace_size( cfg=final_cfg, partial_o_dtype=cutlass.BFloat16, @@ -909,7 +1020,7 @@ def _resolve_mla_decode_launch_spec( decision = select_mla_ts_kernel( requested_policy=requested_policy, batch_size=batch_size, - num_heads=num_heads, + num_heads=num_qo_heads, seq_len_q=seq_len_q, seq_len_k=max_kv_len, latent_dim=kv_lora_rank, @@ -926,6 +1037,7 @@ def _resolve_mla_decode_launch_spec( if not decision.implementation_ready: raise NotImplementedError(decision.reason) split_kv = two_cta_split_kv + capacity = 1 if split_kv == 1 else sm_count work_clusters = batch_size * launch_shape.num_tiles * max(split_kv, 1) # Dynamic cluster stealing only helps once logical work exceeds a # resident wave. Within one wave every cluster already launches, @@ -941,22 +1053,25 @@ def _resolve_mla_decode_launch_spec( is_persistent=is_persistent, is_var_seq=False, is_var_split_kv=False, - static_split_kv=split_kv, + static_split_kv=1 if capacity == 1 else None, + split_kv_capacity=capacity, static_seq_len_k=None, qkv_dtype=qkv_dtype_name, out_dtype=output_dtype_name, rope_dim=qk_rope_head_dim, - num_heads=num_heads, + num_heads=num_qo_heads, seq_len_q=seq_len_q, batch_size=batch_size, mask_type=mask_type, + softmax_rescale_threshold_log2=softmax_rescale_threshold_log2, ) + kernel._configure_parallel_reduction_topology(split_kv) workspace_size = compute_2cta_workspace_size( tile_size_q=int(launch_shape.tile_size_q), num_q_tiles=int(launch_shape.num_tiles), latent_dim=kv_lora_rank, batch_size=batch_size, - split_kv=split_kv, + split_kv=capacity, partial_o_dtype=cutlass.BFloat16, lse_dtype=cutlass.Float32, ) @@ -984,7 +1099,15 @@ def _resolve_mla_decode_launch_spec( ("separate_reducer_impl", separate_reducer_impl), ("reducer_cluster_size", reducer_cluster_size), ) - _validate_mla_policy_coordinate_span(policy) + policy = ( + *policy, + ( + "split_capacity", + capacity, + ), + ("max_split_kv", sm_count), + ) + _validate_mla_policy_coordinate_span(policy, max_kv_len=max_kv_len) return _MLADecodeLaunchSpec( kernel=kernel, @@ -1016,34 +1139,49 @@ def _make_mla_decode_compile_spec( launch_spec: _MLADecodeLaunchSpec, *, device_index: int, - num_heads: int, + num_qo_heads: int, kv_lora_rank: int, qk_rope_head_dim: int, page_size: int, q_dtype_key: str, - output_dtype_key: str, - max_seq_len_q: int, + out_dtype_key: str, + max_q_len: int, packed_query: bool, + lse_mode: str = "none", + runtime_kv_lengths: bool = False, + write_final_lse: bool = True, ) -> _MLADecodeCompileSpec: """Keep policy resolution plan-specific and JIT identity batch-free.""" + kernel = copy.copy(launch_spec.kernel) + kernel.lse_base_on_e = False + kernel.runtime_kv_lengths = runtime_kv_lengths return _MLADecodeCompileSpec( device_index=device_index, - kernel_signature=_mla_kernel_compile_signature(launch_spec.kernel), - num_heads=num_heads, + kernel_signature=_mla_kernel_compile_signature(kernel), + num_qo_heads=num_qo_heads, kv_lora_rank=kv_lora_rank, qk_rope_head_dim=qk_rope_head_dim, page_size=page_size, q_dtype_key=q_dtype_key, - output_dtype_key=output_dtype_key, - max_seq_len_q=max_seq_len_q, + out_dtype_key=out_dtype_key, + max_q_len=( + None if has_runtime_query_lengths(launch_spec.kernel) else max_q_len + ), packed_query=packed_query, has_kernel_workspace=launch_spec.kernel_workspace_bytes > 0, split_kv=launch_spec.split_kv, - kernel=launch_spec.kernel, + kernel=kernel, + lse_mode=lse_mode, + write_final_lse=write_final_lse, ) +# Dense launch specs are cached; sparse plans construct transient kernel objects. +# Reuse dense signature construction without retaining each transient sparse plan. +_cached_mla_decode_compile_spec = functools.cache(_make_mla_decode_compile_spec) + + _COMPILED_MLA_DECODE: dict = {} @@ -1071,14 +1209,16 @@ def _compile_cache_is_placeholder(compiled) -> bool: return is_placeholder(compiled) -def _compile_mla_decode(compile_spec: _MLADecodeCompileSpec): - """Build (or load from the on-disk cache) one batch-dynamic TS MLA topology.""" +def _compile_mla_decode( + compile_spec: _MLADecodeCompileSpec, +): + """Compile and cache one batch-dynamic TS MLA topology.""" import cutlass import cutlass.cute as cute device_index = compile_spec.device_index - num_heads = compile_spec.num_heads + num_qo_heads = compile_spec.num_qo_heads kv_lora_rank = compile_spec.kv_lora_rank qk_rope_head_dim = compile_spec.qk_rope_head_dim query_dim = kv_lora_rank + qk_rope_head_dim @@ -1086,15 +1226,20 @@ def _compile_mla_decode(compile_spec: _MLADecodeCompileSpec): # alias in the callable ABI rather than a zero-extent tensor map. rope_view_dim = qk_rope_head_dim or kv_lora_rank page_size = compile_spec.page_size - max_seq_len_q = compile_spec.max_seq_len_q + max_q_len = ( + cute.sym_int() if compile_spec.max_q_len is None else compile_spec.max_q_len + ) packed_query = compile_spec.packed_query - kernel = compile_spec.kernel + # Keep the cached spec pristine if compiled-callable caches are cleared. + kernel = copy.copy(compile_spec.kernel) + sparse = getattr(kernel, "direct_sparse", False) + nvfp4_kv = getattr(kernel, "has_nvfp4_kv", False) dtype_map = { "bfloat16": cutlass.BFloat16, "float8_e4m3fn": cutlass.Float8E4M3FN, } qkv_dtype = dtype_map[compile_spec.q_dtype_key] - output_dtype = dtype_map[compile_spec.output_dtype_key] + out_dtype = dtype_map[compile_spec.out_dtype_key] physical_pages = cute.sym_int() batch_size = cute.sym_int() runtime_total_q = cute.sym_int() @@ -1102,21 +1247,21 @@ def _compile_mla_decode(compile_spec: _MLADecodeCompileSpec): # These fake tensors pin the public ABI while allowing runtime page counts, # table widths/row strides, and batch metadata pointers to vary. q_stride_h = query_dim - q_stride_q = num_heads * query_dim + q_stride_q = num_qo_heads * query_dim q_latent_shape: tuple[int, ...] q_rope_shape: tuple[int, ...] q_stride: tuple[int, ...] if packed_query: - q_latent_shape = (num_heads, kv_lora_rank, runtime_total_q) - q_rope_shape = (num_heads, rope_view_dim, runtime_total_q) + q_latent_shape = (num_qo_heads, kv_lora_rank, runtime_total_q) + q_rope_shape = (num_qo_heads, rope_view_dim, runtime_total_q) q_stride = (q_stride_h, 1, q_stride_q) else: - q_stride_batch = max_seq_len_q * q_stride_q - q_latent_shape = (num_heads, kv_lora_rank, max_seq_len_q, batch_size) + q_stride_batch = max_q_len * q_stride_q + q_latent_shape = (num_qo_heads, kv_lora_rank, max_q_len, batch_size) q_rope_shape = ( - num_heads, + num_qo_heads, rope_view_dim, - max_seq_len_q, + max_q_len, batch_size, ) q_stride = (q_stride_h, 1, q_stride_q, q_stride_batch) @@ -1126,22 +1271,35 @@ def _compile_mla_decode(compile_spec: _MLADecodeCompileSpec): q_rope_fake = cute.runtime.make_fake_tensor( qkv_dtype, q_rope_shape, stride=q_stride, assumed_align=16 ) - cache_token_stride = query_dim - cache_page_stride = page_size * query_dim + cache_dim = ( + cache_record_size(kernel.kv_cache_format, qk_rope_head_dim) + if nvfp4_kv + else query_dim + ) + extra_cache_dim = ( + cache_record_size(kernel.extra_kv_cache_format, qk_rope_head_dim) + if nvfp4_kv + else query_dim + ) + cache_dtype = cutlass.Uint8 if nvfp4_kv else qkv_dtype + cache_token_stride = cache_dim + cache_page_stride = page_size * cache_dim c_latent_fake = cute.runtime.make_fake_tensor( - qkv_dtype, - (page_size, kv_lora_rank, physical_pages), + cache_dtype, + (page_size, cache_dim if nvfp4_kv else kv_lora_rank, physical_pages), stride=(cache_token_stride, 1, cache_page_stride), assumed_align=16, ) c_rope_fake = cute.runtime.make_fake_tensor( - qkv_dtype, + cache_dtype, ( page_size, - rope_view_dim, + extra_cache_dim if sparse else rope_view_dim, cute.sym_int() if page_size == 1 else physical_pages, ), - stride=(cache_token_stride, 1, cache_page_stride), + stride=(extra_cache_dim, 1, page_size * extra_cache_dim) + if sparse + else (cache_token_stride, 1, cache_page_stride), assumed_align=16, ) runtime_page_columns = cute.sym_int() @@ -1152,27 +1310,31 @@ def _compile_mla_decode(compile_spec: _MLADecodeCompileSpec): stride=(1, runtime_page_row_stride), assumed_align=4, ) - out_stride_row = num_heads * kv_lora_rank + out_stride_row = num_qo_heads * kv_lora_rank out_shape: tuple[int, ...] out_stride: tuple[int, ...] lse_shape: tuple[int, ...] lse_stride: tuple[int, ...] if packed_query: - out_shape = (num_heads, kv_lora_rank, runtime_total_q) + out_shape = (num_qo_heads, kv_lora_rank, runtime_total_q) out_stride = (kv_lora_rank, 1, out_stride_row) - lse_shape = (num_heads, runtime_total_q) - lse_stride = (1, num_heads) + lse_shape = (num_qo_heads, runtime_total_q) + lse_stride = (1, num_qo_heads) else: - out_stride_batch = max_seq_len_q * out_stride_row - out_shape = (num_heads, kv_lora_rank, max_seq_len_q, batch_size) + out_stride_batch = max_q_len * out_stride_row + out_shape = (num_qo_heads, kv_lora_rank, max_q_len, batch_size) out_stride = (kv_lora_rank, 1, out_stride_row, out_stride_batch) - lse_shape = (num_heads, max_seq_len_q, batch_size) - lse_stride = (1, num_heads, max_seq_len_q * num_heads) + lse_shape = (num_qo_heads, max_q_len, batch_size) + lse_stride = (1, num_qo_heads, max_q_len * num_qo_heads) out_fake = cute.runtime.make_fake_tensor( - output_dtype, out_shape, stride=out_stride, assumed_align=16 + out_dtype, out_shape, stride=out_stride, assumed_align=16 ) - lse_fake = cute.runtime.make_fake_tensor( - cutlass.Float32, lse_shape, stride=lse_stride, assumed_align=16 + lse_fake = ( + cute.runtime.make_fake_tensor( + cutlass.Float32, lse_shape, stride=lse_stride, assumed_align=4 + ) + if compile_spec.write_final_lse + else None ) workspace_fake = None if compile_spec.has_kernel_workspace: @@ -1200,29 +1362,26 @@ def _compile_mla_decode(compile_spec: _MLADecodeCompileSpec): ) sparse_inputs_fake = None if getattr(kernel, "direct_sparse", False): - + # Describe runtime index-list extents; the plan capacity does not specialize + # this view. def sparse_indices(capacity): return cute.runtime.make_fake_tensor( cutlass.Int32, - (batch_size, max(1, capacity)), + (batch_size, cute.sym_int()), stride=(cute.sym_int64(), 1), assumed_align=4, ) - scalar = cute.runtime.make_fake_compact_tensor( - cutlass.Float32, (1,), stride_order=(0,), assumed_align=4 - ) sink = cute.runtime.make_fake_compact_tensor( - cutlass.Float32, (num_heads,), stride_order=(0,), assumed_align=4 + cutlass.Float32, (num_qo_heads,), stride_order=(0,), assumed_align=4 ) sparse_inputs_fake = ( sparse_indices(kernel.direct_sparse_capacities[0]), sparse_indices(kernel.direct_sparse_capacities[1]), cache_seqs_fake, cache_seqs_fake, - scalar, - scalar, sink, + cutlass.Float32(1.0), ) stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) @@ -1231,26 +1390,31 @@ def sparse_indices(capacity): # Task objects carry loop-local state through generated control flow, so # select the public staged frontend for this compilation. def compile_kernel(): - return cute.compile[cute.FrontendNext]( - kernel, - q_latent_fake, - q_rope_fake, - c_latent_fake, - c_rope_fake, - page_offsets_fake, - out_fake, - lse_fake, - workspace_fake, - cutlass.Int32(compile_spec.split_kv), - cache_seqs_fake, - qo_indptr_fake, - sparse_inputs_fake, - cutlass.Float32(1.0), - cutlass.Float32(1.0), - stream_fake, - output_quant_args=output_quant_args, - options=_COMPILE_OPTIONS, - ) + """Compile MLA with opt-in task-builder advisory diagnostics.""" + from .kernels.mla_decode.helpers.diagnostics import mla_build_diagnostics + + with mla_build_diagnostics(): + return cute.compile[cute.FrontendNext]( + kernel, + q_latent_fake, + q_rope_fake, + c_latent_fake, + c_rope_fake, + page_offsets_fake, + out_fake, + lse_fake, + workspace_fake, + cutlass.Int32(compile_spec.split_kv), + cache_seqs_fake, + qo_indptr_fake, + sparse_inputs_fake, + cutlass.Float32(1.0), + cutlass.Float32(1.0), + stream_fake, + output_quant_args=output_quant_args, + max_q_len=cutlass.Int32(0), + options=_COMPILE_OPTIONS, + ) with torch.cuda.device(device_index): compiled = build_and_load_attention_kernel( @@ -1266,14 +1430,14 @@ def compile_kernel(): def get_prims_ts_batch_mla_decode_workspace_size( batch_size: int, - num_heads: int, + num_qo_heads: int, kv_lora_rank: int, qk_rope_head_dim: int, page_size: int, - max_seq_len: int, + max_kv_len: int, *, seq_len_q: Optional[int] = None, - max_seq_len_q: Optional[int] = None, + max_q_len: Optional[int] = None, q_dtype: torch.dtype = torch.bfloat16, kv_dtype: Optional[torch.dtype] = None, out_dtype: torch.dtype = torch.bfloat16, @@ -1285,30 +1449,30 @@ def get_prims_ts_batch_mla_decode_workspace_size( The arguments define the static geometry used to resolve the same automatic policy and private scratch layout as :func:`batch_mla_decode_with_paged_kv_cache`, without compiling a kernel. - ``max_seq_len_q`` is the static per-request Q bound for both fixed and + ``max_q_len`` is the static per-request Q bound for both fixed and packed-query launches; ``seq_len_q`` remains a backward-compatible fixed-Q alias. If neither is - supplied, the bound is one. The returned byte count includes both split-KV - scratch and the internal FP32 LSE tensor. Allocate a contiguous - ``torch.int8`` or ``torch.uint8`` CUDA buffer; MLA does not require its + supplied, the bound is one. The returned byte count covers split-KV + scratch. Optional final LSE uses a separate caller-owned output. Allocate + a contiguous ``torch.int8`` or ``torch.uint8`` CUDA buffer; MLA does not require its contents to be initialized before first use. """ batch_size = _validate_positive_int(batch_size, "batch_size") - num_heads = _validate_positive_int(num_heads, "num_heads") + num_qo_heads = _validate_positive_int(num_qo_heads, "num_qo_heads") _validate_mla_dims(kv_lora_rank, qk_rope_head_dim) page_size = _validate_page_size(page_size) - max_seq_len = _validate_mla_max_kv_len(max_seq_len, "max_seq_len") - max_seq_len_q = _resolve_max_seq_len_q_alias( + max_kv_len = _validate_mla_max_kv_len(max_kv_len, "max_kv_len") + max_q_len = _resolve_max_seq_len_q_alias( seq_len_q=seq_len_q, - max_seq_len_q=max_seq_len_q, + max_q_len=max_q_len, default=1, ) - assert max_seq_len_q is not None + assert max_q_len is not None _validate_mla_query_head_extent( batch_size=batch_size, - num_heads=num_heads, - max_seq_len_q=max_seq_len_q, + num_qo_heads=num_qo_heads, + max_q_len=max_q_len, ) _validate_mask(mask_type) if kv_dtype is None: @@ -1320,19 +1484,26 @@ def get_prims_ts_batch_mla_decode_workspace_size( spec = _resolve_mla_decode_launch_spec( device_index, batch_size, - num_heads, + num_qo_heads, kv_lora_rank, qk_rope_head_dim, page_size, - max_seq_len, + max_kv_len, _dtype_key(q_dtype), _dtype_key(kv_dtype), _dtype_key(out_dtype), mask_type, - max_seq_len_q, + max_q_len, + # The skip-correction threshold does not change the workspace; pass the default + # positionally so a default plan() hits the same cached spec. + 0.0, ) return _make_mla_workspace_layout( - spec.kernel_workspace_bytes, batch_size, num_heads, max_seq_len_q + spec.kernel_workspace_bytes, + batch_size, + num_qo_heads, + max_q_len, + include_lse=False, ).total_bytes @@ -1342,16 +1513,16 @@ def _prepare_mla_runtime( *, device: torch.device, batch_size: int, - num_heads: int, - max_seq_len_q: int, + num_qo_heads: int, + max_q_len: int, packed_query: bool, qo_indptr: Optional[torch.Tensor], page_size: int, q_dtype: torch.dtype, kv_dtype: torch.dtype, - output_dtype: torch.dtype, - bmm1_scale: float, - bmm2_scale: float, + out_dtype: torch.dtype, + softmax_scale: float, + output_scale: float, out: Optional[torch.Tensor], validate: bool, ) -> _MLARuntime: @@ -1373,8 +1544,8 @@ def _prepare_mla_runtime( packed_query=packed_query, device=device, batch_size=batch_size, - num_heads=num_heads, - max_seq_len_q=max_seq_len_q, + num_qo_heads=num_qo_heads, + max_q_len=max_q_len, q_dtype=q_dtype, ) ( @@ -1392,39 +1563,39 @@ def _prepare_mla_runtime( f"kv_cache dtype must match the launch ({kv_dtype}), " f"got {normalized_cache.dtype}" ) - effective_bmm1_scale = _validate_scale(bmm1_scale, "bmm1_scale") - effective_bmm2_scale = _validate_scale(bmm2_scale, "bmm2_scale") + effective_softmax_scale = _validate_scale(softmax_scale, "softmax_scale") + effective_output_scale = _validate_scale(output_scale, "output_scale") else: normalized_cache = kv_cache[:, 0] if kv_cache.ndim == 4 else kv_cache num_physical_pages = int(normalized_cache.shape[0]) - effective_bmm1_scale = bmm1_scale - effective_bmm2_scale = bmm2_scale + effective_softmax_scale = softmax_scale + effective_output_scale = output_scale total_q = int(query.shape[0]) if packed_query else None if out is None: out_shape = ( - (total_q, num_heads, _MLA_LATENT_DIM) + (total_q, num_qo_heads, _MLA_LATENT_DIM) if packed_query - else (batch_size, max_seq_len_q, num_heads, _MLA_LATENT_DIM) + else (batch_size, max_q_len, num_qo_heads, _MLA_LATENT_DIM) ) - out = torch.empty(out_shape, device=device, dtype=output_dtype) + out = torch.empty(out_shape, device=device, dtype=out_dtype) elif validate: _validate_out( out, device=device, batch_size=batch_size, - num_heads=num_heads, - max_seq_len_q=max_seq_len_q, + num_qo_heads=num_qo_heads, + max_q_len=max_q_len, packed_query=packed_query, total_q=total_q, - output_dtype=output_dtype, + out_dtype=out_dtype, ) return _MLARuntime( query=query, normalized_cache=normalized_cache, out=out, num_physical_pages=num_physical_pages, - bmm1_scale=effective_bmm1_scale, - bmm2_scale=effective_bmm2_scale, + softmax_scale=effective_softmax_scale, + output_scale=effective_output_scale, ) @@ -1432,16 +1603,19 @@ def _launch_mla_decode( runtime: _MLARuntime, *, block_tables: torch.Tensor, - seq_lens: torch.Tensor, + seq_lens_kv: torch.Tensor, qo_indptr: Optional[torch.Tensor], packed_query: bool, kv_lora_rank: int, split_kv: int, workspace: _MLAWorkspaceViews, compiled: Callable[..., object], + lse: Optional[torch.Tensor] = None, + max_q_len: int, ) -> torch.Tensor: """Form the dimension-first views and launch one compiled MLA kernel.""" + lse_storage = workspace.lse if lse is None else lse if packed_query and int(runtime.query.shape[0]) == 0: return runtime.out if packed_query: @@ -1449,22 +1623,29 @@ def _launch_mla_decode( q_rope = runtime.query[..., kv_lora_rank:].permute(1, 2, 0) out_kernel = runtime.out.permute(1, 2, 0) total_q = int(runtime.query.shape[0]) - lse_kernel = workspace.lse.view(-1, workspace.lse.shape[-1])[ - :total_q - ].transpose(0, 1) + lse_kernel = ( + lse_storage.view(-1, lse_storage.shape[-1])[:total_q].transpose(0, 1) + if lse_storage is not None + else None + ) else: q_latent = runtime.query[..., :kv_lora_rank].permute(2, 3, 1, 0) q_rope = runtime.query[..., kv_lora_rank:].permute(2, 3, 1, 0) out_kernel = runtime.out.permute(2, 3, 1, 0) - lse_kernel = workspace.lse.permute(2, 1, 0) - c_latent = runtime.normalized_cache[..., :kv_lora_rank].permute(1, 2, 0) + lse_kernel = lse_storage.permute(2, 1, 0) if lse_storage is not None else None + c_latent = ( + runtime.normalized_cache + if runtime.normalized_cache.dtype == torch.uint8 + else runtime.normalized_cache[..., :kv_lora_rank] + ).permute(1, 2, 0) c_rope = runtime.normalized_cache[..., kv_lora_rank:].permute(1, 2, 0) if runtime.query.shape[-1] == kv_lora_rank: q_rope = q_latent c_rope = c_latent - if runtime.extra_cache is not None: - c_rope = runtime.extra_cache.permute(1, 2, 0) + if runtime.extra_cache is not None: + c_rope = runtime.extra_cache.permute(1, 2, 0) page_offsets = block_tables.transpose(0, 1) + # Persisted FFI callables require positional arguments, including None. compiled( q_latent, q_rope, @@ -1475,23 +1656,23 @@ def _launch_mla_decode( lse_kernel, workspace.kernel_workspace, split_kv, - seq_lens, + seq_lens_kv, qo_indptr, runtime.sparse_inputs, - runtime.bmm1_scale, - runtime.bmm2_scale, - # Passed positionally and also when None. The kernel was compiled with - # ``output_quant_args`` in its signature (``None`` for unquantized - # output), and a kernel reloaded from the persistent compile cache is - # a positional TVM-FFI entry with no Python defaults: it rejects both - # a missing parameter and a keyword argument. + runtime.softmax_scale, + runtime.output_scale, runtime.output_quant_args, + max_q_len, ) return runtime.out class BatchMLADecodePagedTSWrapper: - """Compile and reuse task-scheduled paged MLA decode launches.""" + """Reusable paged MLA with eager planning and automatic runtime splits. + + Plans select cached kernels from explicit problem bounds and retain scratch. + Each run supplies live tensors and request metadata. Replanning invalidates + graphs; concurrent execution requires independent wrappers/workspaces.""" @flashinfer_experimental_api def __init__(self) -> None: @@ -1503,20 +1684,22 @@ def plan( self, device: int | str | torch.device, batch_size: int, - num_heads: int, + num_qo_heads: int, kv_lora_rank: int, qk_rope_head_dim: int, page_size: int, max_kv_len: int, *, - max_seq_len_q: int, + max_q_len: int, packed_query: bool, - q_data_type: torch.dtype, - kv_data_type: torch.dtype, - o_data_type: torch.dtype, + q_dtype: torch.dtype, + kv_dtype: torch.dtype, + out_dtype: torch.dtype, mask_type: Literal["dense", "causal"] = "causal", + lse_mode: Literal["none", "base2"] = "none", workspace_buffer: Optional[torch.Tensor] = None, validate: bool = True, + tuning: Optional[MLAAttentionTuning] = None, ) -> None: """Compile one static MLA shape and bind its reusable workspace. @@ -1527,30 +1710,40 @@ def plan( leaves that state usable. ``packed_query=False`` selects fixed ``[B, SQ, H, 576]`` query storage, - where ``SQ`` is exactly ``max_seq_len_q``. ``packed_query=True`` selects + where ``SQ`` is exactly ``max_q_len``. ``packed_query=True`` selects ``[total_q, H, 576]`` storage with per-run cumulative offsets supplied to - every run. ``max_seq_len_q`` is then the per-request capacity. + every run. ``max_q_len`` is then the per-request capacity. Individual packed requests may be empty; an all-empty run returns its empty output without dispatching a GPU kernel. + The request count B remains fixed in both modes. Fixed mode also fixes + the total query rows B*SQ; packed mode accepts changing per-request + lengths within the planned capacity. Sparse independent rows can use + the capacity-plan recipe on :class:`BatchSparseMLADecodePagedTSWrapper`. + Q quantization and RoPE are caller-owned preparation; the kernel uses + Q in ``q_data_type`` directly. + If ``workspace_buffer`` is omitted, the plan allocates private scratch. A workspace is mutable and exclusive to one in-flight launch or graph replay. Warm the plan before graph capture, and call ``run`` with ``validate=False`` inside compiled or captured regions. + Equivalent plans/layers reuse CPU policy and compiled-kernel caches. + The active split count remains a runtime attention/reduction operand. + Parameters ---------- device : int, str, or torch.device CUDA device on which the plan will execute. batch_size : int Exact runtime request count. - num_heads, kv_lora_rank, qk_rope_head_dim, page_size : int + num_qo_heads, kv_lora_rank, qk_rope_head_dim, page_size : int Static MLA head geometry and K/V page size. - max_kv_len, max_seq_len_q : int + max_kv_len, max_q_len : int Static per-request K/V and Q capacities. packed_query : bool Select packed rather than fixed query storage. - q_data_type, kv_data_type, o_data_type : torch.dtype + q_dtype, kv_dtype, out_dtype : torch.dtype Query, K/V, and output dtypes used to compile the plan. mask_type : {"dense", "causal"} Attention mask mode. @@ -1562,8 +1755,33 @@ def plan( validate : bool Validate static geometry and caller scratch. Defaults to ``True``. Disable only for previously validated inputs and warmed topology. + tuning : MLAAttentionTuning, optional + Compile-time tuning; ``None`` keeps every default. + ``tuning.softmax_rescale_threshold_log2`` is the opt-in skip of + the online-softmax O rescale (correction), in binades: 0 to 8 for + FP8 QKV, 0 to 32 for BF16 QKV. The name and the FP8 P scale + follow the Prims-TS FMHA decode option of the same name; the + bounds follow the trtllm-gen FMHA ``skipCorrThreshold``. + ``None`` or ``0.0`` rescales on every new row maximum. A positive + value T keeps a row's softmax anchor while a K tile raises the row + maximum by at most T binades, which skips that tile's rescale of + the accumulated output. Probabilities then reach ``2**T``, so FP8 + P is quantized at ``448 * 2**-T`` instead of 448 (1.75 at T = 8, + the trtllm-gen scale). E4M3 P then loses precision from about + ``10.3 - 0.69 * T`` nats below the anchor and flushes to zero + below about ``13 - 0.69 * T`` nats (4.7 and 7.5 at T = 8); the + FP32 row sum keeps the exact values, so the tail's weight is + biased low. The loss covers the KV range + that shares the anchor's split, which is the whole row when + split-KV is not used (typically large batches). Rows whose largest + scores come first, such as attention sinks, with a long + low-probability tail are the most exposed, so check model accuracy + when enabling it for FP8. BF16 P is not scaled. The value is + compiled into the kernel, so each distinct value builds its own + kernel. It is range-checked even when ``validate`` is False. """ + lse_mode = validate_lse_mode(lse_mode) if not isinstance(validate, bool): raise TypeError("validate must be a bool") if validate: @@ -1572,17 +1790,24 @@ def plan( _validate_mask(mask_type) batch_size = _validate_positive_int(batch_size, "batch_size") _validate_mla_int32_extent(batch_size, "batch_size") - num_heads = _validate_positive_int(num_heads, "num_heads") + num_qo_heads = _validate_positive_int(num_qo_heads, "num_qo_heads") _validate_mla_dims(kv_lora_rank, qk_rope_head_dim) page_size = _validate_page_size(page_size) max_kv_len = _validate_mla_max_kv_len(max_kv_len, "max_kv_len") - max_seq_len_q = _validate_positive_int(max_seq_len_q, "max_seq_len_q") + max_q_len = _validate_positive_int(max_q_len, "max_q_len") _validate_mla_query_head_extent( batch_size=batch_size, - num_heads=num_heads, - max_seq_len_q=max_seq_len_q, + num_qo_heads=num_qo_heads, + max_q_len=max_q_len, ) - _validate_mla_dtype_pair(q_data_type, kv_data_type, o_data_type) + _validate_mla_dtype_pair(q_dtype, kv_dtype, out_dtype) + # Range-check the threshold even when validate is False: it selects a + # compiled kernel, and past the per-dtype bound FP8 P (scaled by + # 448 * 2**-T) would lose most of the row to E4M3 subnormals. + threshold = _tuning_rescale_threshold(tuning) + softmax_rescale_threshold_log2 = _check_softmax_rescale_threshold_log2_range( + 0.0 if threshold is None else threshold, q_dtype=q_dtype + ) device, device_index = _resolve_cuda_device(device) if validate: @@ -1593,33 +1818,42 @@ def plan( spec_key = ( device_index, batch_size, - num_heads, + num_qo_heads, kv_lora_rank, qk_rope_head_dim, page_size, max_kv_len, - _dtype_key(q_data_type), - _dtype_key(kv_data_type), - _dtype_key(o_data_type), + _dtype_key(q_dtype), + _dtype_key(kv_dtype), + _dtype_key(out_dtype), mask_type, - max_seq_len_q, + max_q_len, + ) + spec = _resolve_mla_decode_launch_spec( + *spec_key, softmax_rescale_threshold_log2 ) - spec = _resolve_mla_decode_launch_spec(*spec_key) - compile_spec = _make_mla_decode_compile_spec( + compile_spec = _cached_mla_decode_compile_spec( spec, device_index=device_index, - num_heads=num_heads, + num_qo_heads=num_qo_heads, kv_lora_rank=kv_lora_rank, qk_rope_head_dim=qk_rope_head_dim, page_size=page_size, - q_dtype_key=_dtype_key(q_data_type), - output_dtype_key=_dtype_key(o_data_type), - max_seq_len_q=max_seq_len_q, + q_dtype_key=_dtype_key(q_dtype), + out_dtype_key=_dtype_key(out_dtype), + max_q_len=max_q_len, packed_query=packed_query, + lse_mode=lse_mode, + runtime_kv_lengths=True, + write_final_lse=lse_mode != "none", ) policy = spec.policy workspace_layout = _make_mla_workspace_layout( - spec.kernel_workspace_bytes, batch_size, num_heads, max_seq_len_q + spec.kernel_workspace_bytes, + batch_size, + num_qo_heads, + max_q_len, + include_lse=False, ) if workspace_buffer is None: workspace_buffer = torch.empty( @@ -1639,15 +1873,15 @@ def plan( self._plan_state = _MLADecodePlanState( device=device, batch_size=batch_size, - num_heads=num_heads, - max_seq_len_q=max_seq_len_q, + num_qo_heads=num_qo_heads, + max_q_len=max_q_len, packed_query=packed_query, kv_lora_rank=kv_lora_rank, qk_rope_head_dim=qk_rope_head_dim, page_size=page_size, - q_dtype=q_data_type, - kv_dtype=kv_data_type, - output_dtype=o_data_type, + q_dtype=q_dtype, + kv_dtype=kv_dtype, + out_dtype=out_dtype, mask_type=mask_type, max_kv_len=max_kv_len, required_page_columns=required_page_columns, @@ -1656,7 +1890,9 @@ def plan( workspace_views=workspace, compiled=compiled, policy=policy, - split_kv=int(dict(policy)["split_kv"]), + split_kv=spec.split_kv, + lse_mode=lse_mode, + softmax_rescale_threshold_log2=softmax_rescale_threshold_log2, ) @flashinfer_experimental_api @@ -1665,103 +1901,105 @@ def run( query: torch.Tensor, kv_cache: torch.Tensor, block_tables: torch.Tensor, - seq_lens: torch.Tensor, + seq_lens_kv: torch.Tensor, *, qo_indptr: Optional[torch.Tensor] = None, - bmm1_scale: float = 1.0, - bmm2_scale: float = 1.0, + max_q_len: int | None = None, + softmax_scale: float = 1.0, + output_scale: float = 1.0, out: Optional[torch.Tensor] = None, + lse: Optional[torch.Tensor] = None, + return_lse: bool = False, validate: bool = True, - ) -> torch.Tensor: - """Launch the most recently planned MLA decode on the current stream. - - ``block_tables`` and ``seq_lens`` are required per-run bindings. - ``qo_indptr`` is required by a packed-query plan and rejected by a - fixed-query plan. With validation enabled, tensor structure, metadata - values, scales, and every static capacity are checked before - launch. These checks synchronize metadata to the host. Set - ``validate=False`` only after validating representative inputs, and use - it for ``torch.compile`` or CUDA graph capture. - - In either mode, output and workspace must be disjoint from each other - and from all inputs. Storage overlap is not checked. - - Parameters - ---------- - query : torch.Tensor - Runtime fixed or packed query tensor matching the plan. - kv_cache : torch.Tensor - Runtime compact paged latent K/V cache. - block_tables : torch.Tensor - Runtime physical-page table with one inner-contiguous, - non-overlapping row per request. Inter-row padding is accepted. - seq_lens : torch.Tensor - Runtime K/V lengths with one element per request. - qo_indptr : torch.Tensor, optional - Runtime cumulative query offsets for a packed-query plan. - bmm1_scale, bmm2_scale : float - QK and value/output scaling factors. - out : torch.Tensor, optional - Caller-owned output tensor. A new tensor is allocated when omitted. - validate : bool - Enable explicit runtime validation. Defaults to ``True``. + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Launch MLA using live query/cache, page tables and KV lengths. + + Uses the policy and compiled kernels prepared by plan(). Batch and query + shape must match that plan; live K/V lengths stay within its capacity. + Supply output/LSE buffers and validate=False during graph capture, and + retain the wrapper and captured buffers for every replay. + + An enabled lse_mode="base2" plan writes compact FP32 log2(sum(exp(scores))) + with the query shape minus its final dimension. return_lse controls tuple + return only; an omitted enabled buffer is allocated. Output/LSE/workspace must + be disjoint from inputs and live buffers. One wrapper serves one ordered lane. + + Fixed queries use [B, Sq, H, 576]; packed queries use [total_q, H, 576] and + qo_indptr with nonnegative per-request deltas. All-empty packed Q is a no-op. + softmax_scale and output_scale are finite positive host QK and V/output scales. + The dense MLA interface does not expose attention sinks. + max_q_len covers the live Q lengths within plan capacity. Validation + infers an omitted bound from existing metadata; trusted packed calls + default to plan capacity without a readback. This bound is runtime.""" - Returns - ------- - torch.Tensor - The fixed or packed MLA attention output. - """ + if not isinstance(validate, bool): + raise TypeError("validate must be a bool") state = self._plan_state if state is None: raise RuntimeError("plan() must be called before run()") - if not isinstance(validate, bool): - raise TypeError("validate must be a bool") runtime_qo_indptr = qo_indptr if state.packed_query else None runtime = _prepare_mla_runtime( query, kv_cache, device=state.device, batch_size=state.batch_size, - num_heads=state.num_heads, - max_seq_len_q=state.max_seq_len_q, + num_qo_heads=state.num_qo_heads, + max_q_len=state.max_q_len, packed_query=state.packed_query, qo_indptr=runtime_qo_indptr, page_size=state.page_size, q_dtype=state.q_dtype, kv_dtype=state.kv_dtype, - output_dtype=state.output_dtype, - bmm1_scale=bmm1_scale, - bmm2_scale=bmm2_scale, + out_dtype=state.out_dtype, + softmax_scale=softmax_scale, + output_scale=output_scale, out=out, validate=validate, ) if validate: - _validate_mla_run_metadata( + actual_q = _validate_mla_run_metadata( runtime, block_tables, - seq_lens, + seq_lens_kv, qo_indptr, device=state.device, batch_size=state.batch_size, required_page_columns=state.required_page_columns, packed_query=state.packed_query, - max_seq_len_q=state.max_seq_len_q, + max_q_len=state.max_q_len, max_kv_len=state.max_kv_len, page_size=state.page_size, mask_type=state.mask_type, ) - return _launch_mla_decode( + max_q_len = resolve_max_sequence_length( + max_q_len, actual_q, "max_q_len", state.max_q_len + ) + elif not state.packed_query: + max_q_len = int(query.shape[1]) + elif max_q_len is None: + max_q_len = state.max_q_len + lse = prepare_lse( + lse, + q=query, + mode=state.lse_mode, + return_lse=return_lse, + validate=validate, + ) + result = _launch_mla_decode( runtime, block_tables=block_tables, - seq_lens=seq_lens, + seq_lens_kv=seq_lens_kv, qo_indptr=runtime_qo_indptr, packed_query=state.packed_query, kv_lora_rank=state.kv_lora_rank, split_kv=state.split_kv, workspace=state.workspace_views, compiled=state.compiled, + lse=lse, + max_q_len=max_q_len, ) + return (result, lse) if return_lse else result @flashinfer_experimental_api @@ -1769,26 +2007,31 @@ def batch_mla_decode_with_paged_kv_cache( query: torch.Tensor, kv_cache: torch.Tensor, block_tables: torch.Tensor, - seq_lens: torch.Tensor, + seq_lens_kv: torch.Tensor, *, qo_indptr: Optional[torch.Tensor] = None, - max_seq_len_q: Optional[int] = None, + max_q_len: Optional[int] = None, kv_lora_rank: int = _MLA_LATENT_DIM, qk_rope_head_dim: int = _MLA_ROPE_DIM, mask_type: Literal["dense", "causal"] = "causal", max_kv_len: Optional[int] = None, - bmm1_scale: float = 1.0, - bmm2_scale: float = 1.0, + softmax_scale: float = 1.0, + output_scale: float = 1.0, out: Optional[torch.Tensor] = None, + lse: Optional[torch.Tensor] = None, + return_lse: bool = False, + lse_mode: Literal["base2"] = "base2", out_dtype: torch.dtype = torch.bfloat16, workspace_buffer: Optional[torch.Tensor] = None, validate: bool = True, -) -> torch.Tensor: + tuning: Optional[MLAAttentionTuning] = None, +) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """One-shot convenience wrapper for fixed or packed-query MLA decode. - Without caller scratch, this helper reads ``seq_lens`` and packed-Q - ``qo_indptr`` on the host to derive bounds and constructs a temporary - wrapper outside capture. See Notes for the explicit capture-safe mode. + Without caller scratch, this helper may read ``seq_lens_kv`` on the host + to derive the K/V bound and constructs a temporary wrapper outside capture. + Packed Q requires a caller-provided ``max_q_len`` in either mode. + See Notes for the explicit capture-safe mode. Parameters ---------- @@ -1799,13 +2042,14 @@ def batch_mla_decode_with_paged_kv_cache( block_tables : torch.Tensor Dense physical-page table for each request. Rows must be inner contiguous and non-overlapping, but may have padding between them. - seq_lens : torch.Tensor + seq_lens_kv : torch.Tensor Live K/V sequence lengths. qo_indptr : torch.Tensor, optional Cumulative query offsets selecting packed-query mode. - max_seq_len_q : int, optional - Per-request packed-query length capacity. For packed Q it defaults to - the maximum delta in ``qo_indptr``; an explicit value may be larger. + max_q_len : int, optional + Per-request packed-query length capacity. Required for packed Q; + it is never inferred from ``qo_indptr`` and may exceed the maximum + request length. For fixed Q it defaults to the query's sequence extent and, when provided, must equal that extent. kv_lora_rank, qk_rope_head_dim : int @@ -1814,19 +2058,24 @@ def batch_mla_decode_with_paged_kv_cache( Attention mask mode. max_kv_len : int, optional Static K/V length bound; defaults to the metadata maximum. - bmm1_scale, bmm2_scale : float + softmax_scale, output_scale : float QK and value/output scaling factors. out : torch.Tensor, optional Caller-owned output tensor. out_dtype : torch.dtype Output dtype. - workspace_buffer : torch.Tensor, optional Caller-owned byte scratch, exclusive to one in-flight launch/graph. validate : bool Validate tensors and metadata values (may synchronize), default True. False trusts the caller and requires workspace and explicit bounds; skips tensor and metadata validation. Invalid inputs have undefined behavior. + tuning : MLAAttentionTuning, optional + Compile-time tuning; see :meth:`BatchMLADecodePagedTSWrapper.plan`. + ``None`` (default) keeps the exact rescale. A positive + ``softmax_rescale_threshold_log2`` T quantizes FP8 P at + ``448 * 2**-T`` instead of 448, which costs E4M3 precision for + probabilities far below the anchor. Returns ------- @@ -1837,10 +2086,11 @@ def batch_mla_decode_with_paged_kv_cache( Both owned and caller-provided scratch use the same wrapper plan/run path. For CUDA Graph capture, supply ``workspace_buffer``, ``max_kv_len``, ``out``, and ``validate=False``; packed Q also needs an explicit - ``max_seq_len_q``. + ``max_q_len``. Warm this exact topology outside capture first. Retain stable tensor storage and mutate metadata only between completed launches/replays. All live K/V - lengths must be positive and within the static bound, active page IDs must + positive lengths must be within the static bound; nonpositive lengths mark + inactive requests whose output is untouched. Active page IDs must index the cache, and causal per-request Q lengths must not exceed K/V lengths. Packed offsets must start at zero, end at the query token count, and have nonnegative deltas within the Q bound. Scratch/output must not alias @@ -1848,6 +2098,18 @@ def batch_mla_decode_with_paged_kv_cache( explicit path. Kernel policy, necessary control resets, and output layout are unchanged. Missing output may be allocated only outside capture. + + ``return_lse=True`` returns ``(output, lse)``. ``lse_mode`` selects + ``"base2"`` (the only supported base); it is used only when returning LSE + or writing a caller-provided ``lse`` buffer. + + ``lse`` is an optional independent, compact float32 output with shape + ``query.shape[:-1]`` (``q.shape[:-1]`` for FMHA). An enabled LSE plan + allocates it when omitted, including when ``return_lse=False``. A supplied + buffer is written regardless of ``return_lse``; the flag only controls + the Python return type. LSE includes any sink mass and excludes output + scaling. Keep LSE disjoint from output, inputs, and workspace. Supply + both output buffers and disable validation for allocation-free capture. """ if not isinstance(validate, bool): @@ -1863,13 +2125,17 @@ def batch_mla_decode_with_paged_kv_cache( raise ValueError("validate=False requires workspace_buffer") if max_kv_len is None: raise ValueError("validate=False requires max_kv_len") - if qo_indptr is not None and max_seq_len_q is None: - raise ValueError("validate=False requires max_seq_len_q for packed Q") + if qo_indptr is not None and max_q_len is None: + raise ValueError("validate=False requires max_q_len for packed Q") packed_query = qo_indptr is not None + if packed_query and max_q_len is None: + raise ValueError("packed Q requires max_q_len") if validate: _validate_query(query, packed_query=packed_query) - metadata_device, batch_size, _ = _validate_mla_metadata(block_tables, seq_lens) + metadata_device, batch_size, _ = _validate_mla_metadata( + block_tables, seq_lens_kv + ) if metadata_device != query.device: raise ValueError( f"MLA metadata must be on {query.device}, got {metadata_device}" @@ -1879,86 +2145,65 @@ def batch_mla_decode_with_paged_kv_cache( ) if packed_query: _validate_qo_indptr(qo_indptr, device=query.device, batch_size=batch_size) - if max_seq_len_q is None: - max_seq_len_q, _, _ = _derive_max_seq_len_q( - qo_indptr, batch_size=batch_size - ) - if max_seq_len_q == 0: - raise ValueError( - "max_seq_len_q is required for an all-empty packed query" - ) - elif max_seq_len_q is not None and max_seq_len_q != int(query.shape[1]): + elif max_q_len is not None and max_q_len != int(query.shape[1]): raise ValueError( - "fixed query length must equal max_seq_len_q: " - f"got SQ={query.shape[1]} and max_seq_len_q={max_seq_len_q}" + "fixed query length must equal max_q_len: " + f"got SQ={query.shape[1]} and max_q_len={max_q_len}" ) else: - batch_size = int(seq_lens.shape[0]) + batch_size = int(seq_lens_kv.shape[0]) normalized_cache = kv_cache[:, 0] if kv_cache.ndim == 4 else kv_cache page_size = int(normalized_cache.shape[1]) - num_heads = int(query.shape[-2]) - if not packed_query and max_seq_len_q is None: - max_seq_len_q = int(query.shape[1]) + num_qo_heads = int(query.shape[-2]) + if not packed_query and max_q_len is None: + max_q_len = int(query.shape[1]) if max_kv_len is None: - max_kv_len = int(seq_lens.max().item()) - assert max_seq_len_q is not None + max_kv_len = int(seq_lens_kv.max().item()) + assert max_q_len is not None if validate and out is not None: _validate_out( out, device=query.device, batch_size=batch_size, - num_heads=num_heads, - max_seq_len_q=max_seq_len_q, + num_qo_heads=num_qo_heads, + max_q_len=max_q_len, packed_query=packed_query, total_q=int(query.shape[0]) if packed_query else None, - output_dtype=out_dtype, - ) - if workspace_buffer is None: - workspace_bytes = get_prims_ts_batch_mla_decode_workspace_size( - batch_size, - num_heads, - kv_lora_rank, - qk_rope_head_dim, - page_size, - max_kv_len, - max_seq_len_q=max_seq_len_q, - q_dtype=query.dtype, - kv_dtype=normalized_cache.dtype, out_dtype=out_dtype, - mask_type=mask_type, - device=query.device, ) - workspace_buffer = torch.empty( - workspace_bytes, dtype=torch.int8, device=query.device - ) - wrapper = BatchMLADecodePagedTSWrapper() wrapper.plan( query.device, - batch_size, - num_heads, - kv_lora_rank, - qk_rope_head_dim, - page_size, - max_kv_len, - max_seq_len_q=max_seq_len_q, + batch_size=batch_size, + max_q_len=max_q_len, + max_kv_len=max_kv_len, + num_qo_heads=num_qo_heads, + kv_lora_rank=kv_lora_rank, + qk_rope_head_dim=qk_rope_head_dim, + page_size=page_size, packed_query=packed_query, - q_data_type=query.dtype, - kv_data_type=normalized_cache.dtype, - o_data_type=out_dtype, + q_dtype=query.dtype, + kv_dtype=normalized_cache.dtype, + out_dtype=out_dtype, mask_type=mask_type, workspace_buffer=workspace_buffer, validate=validate, + tuning=tuning, + lse_mode=validate_lse_mode(lse_mode) + if return_lse or lse is not None + else "none", ) return wrapper.run( query, normalized_cache, block_tables, - seq_lens, + seq_lens_kv, qo_indptr=qo_indptr, - bmm1_scale=bmm1_scale, - bmm2_scale=bmm2_scale, + softmax_scale=softmax_scale, + output_scale=output_scale, out=out, + lse=lse, + return_lse=return_lse, validate=validate, ) @@ -1984,7 +2229,7 @@ def _make_sparse_output_quant(heads, block_size, scale_format, groups): raise ValueError("out_scale_format must be 'fp32' or 'ue8m0'") groups = _validate_positive_int(groups, "num_output_groups") if heads % groups: - raise ValueError("num_output_groups must divide num_heads") + raise ValueError("num_output_groups must divide num_qo_heads") from .kernels.mla_decode.helpers.output_quant import MlaOutputQuant return MlaOutputQuant(block_size, scale_format == "ue8m0", heads, groups) @@ -2125,20 +2370,25 @@ def _prepare_sparse_output_quant( def _resolve_sparse_mla_plan( device, batch_size, - num_heads, + num_qo_heads, *, max_topk, max_extra_topk=0, - max_seq_len_q=1, + max_q_len=1, packed_query=False, - q_data_type=torch.bfloat16, + q_dtype=torch.bfloat16, + qk_rope_head_dim=0, + kv_cache_format="native", + extra_kv_cache_format=None, kv_layout="NHD", - has_sinks=False, + use_attention_sinks=False, return_lse=False, assume_valid_prefix=False, output_quant_block_size=None, out_scale_format="ue8m0", num_output_groups=None, + tuning=None, + kv_lora_rank=512, ): """Resolve sparse geometry and scratch sizing without allocation/compilation.""" import cutlass @@ -2156,6 +2406,8 @@ def _resolve_sparse_mla_plan( MlaDecodeTs, ) + if isinstance(kv_lora_rank, bool) or kv_lora_rank != 512: + raise ValueError("sparse MLA requires kv_lora_rank=512") if not isinstance(assume_valid_prefix, bool): raise TypeError("assume_valid_prefix must be bool") device, _ = _resolve_cuda_device(device) @@ -2163,60 +2415,128 @@ def _resolve_sparse_mla_plan( raise NotImplementedError("sparse TS MLA requires SM100 or SM103") for name, value in ( ("batch_size", batch_size), - ("num_heads", num_heads), - ("max_seq_len_q", max_seq_len_q), + ("num_qo_heads", num_qo_heads), + ("max_q_len", max_q_len), ): if not isinstance(value, int) or isinstance(value, bool) or value <= 0: raise ValueError(f"{name} must be a positive integer") - if num_heads > 128: - raise ValueError("num_heads must be at most 128") + if num_qo_heads > 128: + raise ValueError("num_qo_heads must be at most 128") + if isinstance(qk_rope_head_dim, bool) or qk_rope_head_dim not in (0, 64): + raise ValueError("qk_rope_head_dim must be 0 or 64") + if qk_rope_head_dim and output_quant_block_size is not None: + raise ValueError("fused output RoPE requires D512 attention") output_quant = _make_sparse_output_quant( - num_heads, output_quant_block_size, out_scale_format, num_output_groups + num_qo_heads, output_quant_block_size, out_scale_format, num_output_groups ) for value in (max_topk, max_extra_topk): if not isinstance(value, int) or isinstance(value, bool) or value < 0: raise ValueError("source capacities must be nonnegative integers") if max_topk == 0: raise ValueError("the primary selected-slot capacity must be positive") - if q_data_type not in (torch.bfloat16, torch.float8_e4m3fn): + if q_dtype not in (torch.bfloat16, torch.float8_e4m3fn): raise ValueError("native BF16 or E4M3 is required") + if extra_kv_cache_format is None: + extra_kv_cache_format = "nvfp4" if kv_cache_format == "fp8" else kv_cache_format + validate_nvfp4_kv_formats( + kv_cache_format, + extra_kv_cache_format, + qk_rope_head_dim, + "e4m3" if q_dtype == torch.float8_e4m3fn else "bf16", + 1, + ) if kv_layout not in ("NHD", "HND"): raise ValueError("kv_layout must be NHD or HND") - max_rows = batch_size * max_seq_len_q + softmax_rescale_threshold_log2 = _tuning_rescale_threshold(tuning) + if softmax_rescale_threshold_log2 is not None: + softmax_rescale_threshold_log2 = _check_softmax_rescale_threshold_log2_range( + softmax_rescale_threshold_log2, q_dtype=q_dtype + ) + max_rows = batch_size * max_q_len capacity = max( 256, ((max_topk + 127) // 128 + (max_extra_topk + 127) // 128) * 128, ) - if max_rows * num_heads >= 2**31 or capacity >= 2**31 - 32768: + if max_rows * num_qo_heads >= 2**31 or capacity >= 2**31 - 32768: raise ValueError("query/route extent exceeds the int32 kernel domain") sm_count = torch.cuda.get_device_properties(device).multi_processor_count profile = select_sparse_mla_profile( rows=max_rows, - heads=num_heads, - query_length=max_seq_len_q, + heads=num_qo_heads, + query_length=max_q_len, capacity=capacity, - dtype="bf16" if q_data_type == torch.bfloat16 else "fp8", + dtype="bf16" if q_dtype == torch.bfloat16 else "fp8", sm_count=sm_count, + rope_dim=qk_rope_head_dim, + kv_cache_format=kv_cache_format, ) + # Reloading V from L2 shortens native FP8 decode's producer lifetime. + # Sustained grids amortize the retained-KV/CLC pipeline instead; packed + # caches and two-source/sink plans retain their format-specific schedules. + # The general path also exploits a hole-free prefix promise; the reload + # loader keeps arbitrary-hole masks and does not use that optimization. + if ( + qk_rope_head_dim == 64 + and kv_cache_format == "native" + and q_dtype == torch.float8_e4m3fn + and num_qo_heads in _D576_HEADS + and not max_extra_topk + and not use_attention_sinks + and not assume_valid_prefix + and prefer_sparse_d576_reload( + rows=max_rows, + heads=num_qo_heads, + query_length=max_q_len, + capacity=max_topk, + sm_count=sm_count, + general_profile=profile, + ) + ): + return resolve_sparse_d576_reload_plan( + device, + max_rows, + num_qo_heads, + max_topk=max_topk, + max_extra_topk=max_extra_topk, + q_dtype=q_dtype, + use_attention_sinks=use_attention_sinks, + assume_valid_prefix=assume_valid_prefix, + threshold=softmax_rescale_threshold_log2, + ) + kernel_threshold = 0.0 + if softmax_rescale_threshold_log2 is not None: + # An explicit threshold replaces the schedule's fixed six-binade + # deferral; zero restores exact maxima for every schedule. + profile = replace(profile, defer_max_update=False) + kernel_threshold = softmax_rescale_threshold_log2 family, tile = profile.family, profile.tile_size_q splits, head_dim_ctas = profile.split_kv, profile.head_dim_ctas - dtype_name = "bf16" if q_data_type == torch.bfloat16 else "e4m3" + dtype_name = "bf16" if q_dtype == torch.bfloat16 else "e4m3" if family == "2cta": + active_clusters = sm_count // 2 + # Keep the physical device limit batch-independent. Quantized + # persistent launches balance their grid from live tensor extents. kernel = MlaDecodeTs( page_size=1, - max_active_clusters=sm_count // 2, + max_active_clusters=active_clusters, is_persistent=profile.scheduler != "nonpersistent", is_var_seq=False, is_var_split_kv=False, static_split_kv=splits, qkv_dtype=dtype_name, out_dtype="bf16", - rope_dim=0, - num_heads=num_heads, + rope_dim=qk_rope_head_dim, + num_heads=num_qo_heads, seq_len_q=1, batch_size=max_rows, mask_type="dense", + softmax_rescale_threshold_log2=kernel_threshold, + # Sparse MLA flattens Q rows into virtual batch items, so the + # physical Q length stays compile-time. + fixed_seqlen=True, ) + kernel.kv_cache_format = kv_cache_format + kernel.extra_kv_cache_format = extra_kv_cache_format kernel.sparse_gather_warps = profile.gather_issue_warps kernel.sparse_kv_stages = profile.kv_pipeline_stages kernel.sparse_uniform_pages = profile.uniform_offset_cache @@ -2245,20 +2565,23 @@ def _resolve_sparse_mla_plan( kernel = ThroughputLatencyMlaDecodeTs( batch_size=max_rows, num_heads=tile, - seq_len_q=(num_heads + tile - 1) // tile, + seq_len_q=(num_qo_heads + tile - 1) // tile, seq_len_k=capacity, - rope_dim=0, + rope_dim=qk_rope_head_dim, page_size=1, max_active_clusters=sm_count, qkv_dtype=dtype_name, out_dtype="bf16", profile=kernel_profile, reduction_mode=profile.reduction, - logical_num_heads=num_heads, + logical_num_heads=num_qo_heads, logical_seq_len_q=1, tile_size_q=tile, sparse_profile=profile, mask_type="dense", + softmax_rescale_threshold_log2=kernel_threshold, + kv_cache_format=kv_cache_format, + extra_kv_cache_format=extra_kv_cache_format, ) core_bytes = compute_workspace_size( cfg=kernel._make_config(), @@ -2268,32 +2591,202 @@ def _resolve_sparse_mla_plan( spec = _make_mla_decode_compile_spec( _MLADecodeLaunchSpec(kernel, (), core_bytes, splits), device_index=device.index, - num_heads=num_heads, + num_qo_heads=num_qo_heads, kv_lora_rank=512, - qk_rope_head_dim=0, + qk_rope_head_dim=qk_rope_head_dim, page_size=1, - q_dtype_key=str(q_data_type).removeprefix("torch."), - output_dtype_key="float8_e4m3fn" if output_quant is not None else "bfloat16", - max_seq_len_q=1, + q_dtype_key=str(q_dtype).removeprefix("torch."), + out_dtype_key="float8_e4m3fn" if output_quant is not None else "bfloat16", + max_q_len=1, packed_query=False, ) sections = {} byte_end = 0 for name, shape, dtype in ( ("core", (core_bytes,), torch.int8), - ("public_lse", (max_rows, num_heads), torch.float32), + ("public_lse", (max_rows, num_qo_heads), torch.float32), ): sections[name], byte_end = _append_workspace_section(byte_end, shape, dtype) return device, profile, spec, sections, byte_end, capacity, output_quant -# Prepared launches kept per plan (one per layer sharing the plan; DeepSeek-V4 has 61). +_D576_HEAD_DIM = 576 + + +_D576_V_DIM = 512 + + +_D576_HEADS = (64, 128) + + +def _sparse_d576_workspace_bytes(max_rows: int, num_heads: int, split_kv: int) -> int: + """Split-KV partials of a D576 plan: BF16 O and FP32 LSE per split row.""" + if split_kv <= 1: + return 0 + return max_rows * num_heads * split_kv * (_D576_V_DIM * 2 + 4) + + +def resolve_sparse_d576_reload_plan( + device, + max_rows, + num_heads, + *, + max_topk, + max_extra_topk, + q_dtype, + use_attention_sinks, + assume_valid_prefix, + threshold, +): + """Plan the native D576 reload schedule and its split-KV scratch. + + One CTA per (token, 64 heads) on the 1CTA keeps-MMA-AB q64 schedule; 128 + heads run as two physical query positions of one token. The kernel is + compiled for the top-k capacity rounded up to whole 128-key tiles (it uses + the capacity only through its tile count). Below one wave the plan splits + each top-k list over ``split_kv`` CTAs (and the 512 V columns over + ``head_dim_ctas``); the parallel GMEM reducer, launched with PDL, merges + the partials and publishes natural-log LSE. The split is fixed per plan, + so a CUDA graph captured from it uses one compiled kernel. + """ + import cutlass + + from ._hardware import get_max_active_clusters + from .kernels.mla_decode.sparse_policy import ( + SparseMlaProfile, + select_sparse_mla_d576_split, + ) + from .kernels.mla_decode.throughput_latency_1cta.config import MlaProfile + from .kernels.mla_decode.throughput_latency_1cta.kernel import ( + ThroughputLatencyMlaDecodeTs, + ) + + if q_dtype != torch.float8_e4m3fn: + raise ValueError("D576 sparse MLA (qk_rope_head_dim=64) requires E4M3 Q / KV") + if num_heads not in _D576_HEADS: + raise ValueError( + f"D576 sparse MLA supports 64 or 128 query heads (64-row tiles), got {num_heads}" + ) + if max_extra_topk: + raise ValueError("D576 sparse MLA attends to one KV pool (max_extra_topk=0)") + if use_attention_sinks: + raise ValueError("D576 sparse MLA does not support attention sinks") + if max_rows > 65535: + # Each query token is one grid-z slice of the attention and reducer + # launches; CUDA caps grid z at 65535. + raise ValueError( + "D576 sparse MLA supports at most 65535 query tokens per plan " + f"(batch_size * max_q_len), got {max_rows}" + ) + split_kv, head_dim_ctas = select_sparse_mla_d576_split( + max_rows * (num_heads // 64), + max_topk, + torch.cuda.get_device_properties(device).multi_processor_count, + ) + profile = SparseMlaProfile( + family="keep", + tile_size_q=64, + split_kv=split_kv, + head_dim_ctas=head_dim_ctas, + uniform_offset_cache=True, + kv_pipeline_stages=10, + ) + capacity = -(-max_topk // 128) * 128 + phys_sq = num_heads // 64 + kernel_profile = "h64_keeps_mma_ab" + if split_kv > 1: + hdim = "" if head_dim_ctas == 1 else f"_hdim{_D576_V_DIM // head_dim_ctas}" + kernel_profile = MlaProfile( + name=f"h64_keeps_mma_ab_splitkv{split_kv}{hdim}_gmem", + num_ctas_per_seq_kv=split_kv, + num_ctas_per_head_dim=head_dim_ctas, + use_persistent_scheduler=0, + use_clc_dynamic_persistent_scheduler=0, + use_multi_ctas_kv=1, + use_cluster_reduction=0, + kernel_variant="keeps_mma_ab", + tile_size_q=64, + ) + with torch.cuda.device(device): + max_active_clusters = get_max_active_clusters(1) + kernel = ThroughputLatencyMlaDecodeTs( + batch_size=1, + num_heads=64, + seq_len_q=phys_sq, + seq_len_k=capacity, + latent_dim=512, + rope_dim=64, + page_size=1, + max_active_clusters=max_active_clusters, + acc_dtype=cutlass.Float32, + lse_dtype=cutlass.Float32, + qkv_dtype="e4m3", + out_dtype="bf16", + profile=kernel_profile, + reduction_mode=None if split_kv == 1 else "gmem_separate", + logical_num_heads=num_heads, + logical_seq_len_q=1, + fixed_seqlen=True, + tile_size_q=64, + explicit_split_kv=None, + explicit_persistent=None, + mask_type="dense", + softmax_rescale_threshold_log2=0.0 if threshold is None else threshold, + ) + # Preserve the sparse wrapper's natural-log LSE contract. + kernel.lse_base_on_e = True + cfg = kernel._make_config() + if ( + int(cfg.num_ctas_per_seq_kv) != split_kv + or int(cfg.num_ctas_per_head_dim) != head_dim_ctas + or int(cfg.use_cluster_reduction) != 0 + ): + raise NotImplementedError( + f"D576 sparse MLA expects {split_kv} KV CTAs, {head_dim_ctas} V " + f"CTAs and GMEM reduction; got {cfg.num_ctas_per_seq_kv}, " + f"{cfg.num_ctas_per_head_dim}, cluster={cfg.use_cluster_reduction}" + ) + if split_kv > 1 and not kernel.use_parallel_reduction: + # Only the parallel GMEM reducer publishes the D576 natural-log LSE. + raise NotImplementedError("D576 split-KV requires the parallel GMEM reducer") + spec = _MLADecodeCompileSpec( + device_index=device.index, + kernel_signature=_mla_kernel_compile_signature(kernel), + num_qo_heads=64, + kv_lora_rank=512, + qk_rope_head_dim=64, + page_size=1, + q_dtype_key="float8_e4m3fn", + out_dtype_key="bfloat16", + max_q_len=phys_sq, + packed_query=False, + has_kernel_workspace=split_kv > 1, + split_kv=split_kv, + kernel=kernel, + ) + sections = {} + byte_end = 0 + for name, shape, dtype in ( + ( + "core", + (_sparse_d576_workspace_bytes(max_rows, num_heads, split_kv),), + torch.int8, + ), + ("public_lse", (max_rows, num_heads), torch.float32), + ): + sections[name], byte_end = _append_workspace_section(byte_end, shape, dtype) + return device, profile, spec, sections, byte_end, capacity, None + + _PREPARED_LAUNCH_LIMIT = 256 + class BatchSparseMLADecodePagedTSWrapper: """Plan native BF16/E4M3 attention, then bind one or two source index lists. - Q is [B,Sq,H,512] or packed [total_q,H,512]. Each source supplies int32 + Q is [B,Sq,H,D] or packed [total_q,H,D], where D is 512 or 576. + D576 adds 64 QK-only RoPE channels; V and output remain D512. + Each source supplies int32 [rows, capacity] storage-row indices and int32 [rows] live lengths directly to run(). Indices include physical page strides; -1 is masked. Both source lists participate in one attention distribution. The caller folds the @@ -2316,17 +2809,27 @@ class BatchSparseMLADecodePagedTSWrapper: warmup, stable addresses, preallocated outputs, and validate=False. Input/output/workspace storage must not overlap; this is an unchecked caller precondition. Separate wrappers/workspaces are needed per stream. + + Planning compiles one kernel with host-runtime FP32 scale arguments. + Scale values may change between launches without recompiling. CUDA graphs + capture their values, so new values require graph parameter updates or + recapture. For a + fixed dispatch profile over a range of independent sparse query rows, + plan with ``batch_size=1``, ``packed_query=True`` and + ``max_q_len=row_capacity``. Run with ``[rows,H,D]`` queries and + ``qo_indptr=[0,rows]`` for any ``rows <= row_capacity``. + + Q quantization and query RoPE belong to the caller. For a serving hot path, + establish valid metadata upstream and pass ``run(validate=False)``. + The default ``validate=True`` reads CUDA metadata values on the host and + therefore synchronizes. Kernel-builder advisories are controlled separately + at plan/JIT time by ``FLASHINFER_PRIMS_TS_DEBUG_CHECKS=1``. """ @flashinfer_experimental_api def __init__(self, workspace_buffer=None): self._workspace_buffer = workspace_buffer self._state = None - self._scalar_cache = {} - # Prepared launches for the unfused validate=False path: everything that - # is static for one (plan, pools, scales, sinks) tuple, keyed by the pools' - # storage and the sinks address. One plan serves every layer of a model, - # and each layer brings its own sinks, so one entry per layer. self._prepared: dict[tuple, dict] = {} self.workspace_size_bytes = 0 @@ -2334,56 +2837,134 @@ def plan( self, device, batch_size, - num_heads, + num_qo_heads, *, max_topk, max_extra_topk=0, - max_seq_len_q=1, + max_q_len=1, packed_query=False, - q_data_type=torch.bfloat16, + q_dtype=torch.bfloat16, + qk_rope_head_dim=0, + kv_cache_format="native", + extra_kv_cache_format=None, kv_layout="NHD", - has_sinks=False, + use_attention_sinks=False, return_lse=False, assume_valid_prefix=False, output_quant_block_size=None, out_scale_format="ue8m0", num_output_groups=None, + tuning=None, + kv_lora_rank=512, ): - """Plan D512 attention over one primary and an optional extra pool. + """Plan D512/D576 attention over one primary and an optional extra pool. + + ``kv_cache_format="nvfp4"`` uses the DS v4.1 D512 format: 256 packed + E2M1 bytes followed by 32 E4M3 scales in linear block16 order. Every + QK/V channel, including rotated channels, uses the same storage. + Keep ``qk_rope_head_dim=0``: there is no additional QK-only tail. + + Set ``qk_rope_head_dim=64`` and ``kv_cache_format="nvfp4.fp8rope"`` + for D576: 256 packed E2M1 NoPE bytes, 64 unscaled E4M3 RoPE bytes, + then 32 E4M3 scales, with block b stored at 320 + 8*(b%4) + b//4. + ``kv_cache_format="fp8"`` selects D512 per-tensor E4M3 storage: + 512 bytes per token with no scale trailer. Pass its positive tensor + descale as ``kv_cache_scale`` to run(). Pair it with + ``extra_kv_cache_format="nvfp4"`` for FP8 SWA and compressed NVFP4 KV. + The extra format defaults to NVFP4 for an FP8 primary, otherwise to + the primary format. D512 supports NVFP4 in both pools or FP8 primary + with NVFP4 extra; native and mixed-RoPE formats require matching pools. + NVFP4-containing recipes require E4M3 Q. NVFP4 block scales are + applied in the loader before ordinary FP8 MMA, without a global + storage scale. Per-tensor FP8 bytes feed MMA directly; its scale is + applied to QK scores and FP32 output accumulation, avoiding another + E4M3 rounding. Caches are uint8 records; FP8 accepts E4M3 tensors too. + In NVFP4 bytes, even-indexed values occupy the low nibble. NVFP4 + expands to E4M3 in the loader; BF16 Q must be converted upstream for + these recipes. ``max_topk`` describes arbitrary selected primary entries, not an SWA window. Prepared indices already include the physical page stride, so the core uses page size one regardless of the pools' external pages. Planning compiles no metadata-preparation kernel. + Fixed queries use exactly ``batch_size * max_q_len`` rows. To reuse + one dispatch profile over a range of independent query-row counts, + use ``batch_size=1``, ``packed_query=True`` and + ``max_q_len=row_capacity``. Supply ``[rows,H,D]`` Q and CUDA INT32 + ``qo_indptr=[0,rows]`` to each run, with ``rows <= row_capacity``. + Index/length tensors and outputs use the same live row count. The + profile is chosen for the planned capacity; each captured graph keeps + its captured views and addresses. For example:: + + wrapper.plan("cuda", 1, heads, max_topk=topk, + max_q_len=row_capacity, packed_query=True, + q_data_type=torch.float8_e4m3fn, + kv_cache_format="nvfp4", return_lse=True) + # Q is already quantized/rotated; update metadata upstream. + wrapper.run(q, cache, indices, lengths, qo_indptr=qo_indptr, + out=out, lse=lse, validate=False) + + Whole-tile K/V reuse supports nonpersistent split-KV schedules. + The 1CTA persistent/CLC scheduler requires ``split_kv=1`` because its + work queue enumerates query/head tiles without a KV-split coordinate. + Low-level profile overrides must use a nonpersistent scheduler with + multiple KV splits; this restriction does not apply to 2CTA profiles. + + Routine task-builder overlap and unused-routing advisories are quiet + by default. Set ``FLASHINFER_PRIMS_TS_DEBUG_CHECKS=1`` before a cold + plan/JIT build to show them. This controls build diagnostics separately + from :meth:`run`'s runtime input validation; other warnings/errors are + retained. + Output quantization is disabled by default. When enabled, G must divide the local query-head count H. Blocks of 32 with UE8M0 scales implement MXFP8; FP32 scales use unrounded amax/448. Both use an amax floor of 1e-4. A block never crosses a head boundary. + + ``tuning`` (an :class:`MLAAttentionTuning`) has the meaning of + :meth:`BatchMLADecodePagedTSWrapper.plan`; its + ``softmax_rescale_threshold_log2`` takes 0 to 8 binades for E4M3 and + 0 to 32 for BF16. ``None`` (default) keeps each schedule's own + behavior: some BF16 schedules keep the row anchor while a K tile + raises it by at most 6 binades, and all others rescale on every new + maximum. An explicit value applies to every schedule; ``0.0`` + disables the skip (slower than ``None`` on the deferred BF16 + schedules), and a positive value T also quantizes E4M3 P at + ``448 * 2**-T`` instead of 448, with the precision cost described + there. """ (device, profile, spec, sections, byte_end, _, output_quant) = ( _resolve_sparse_mla_plan( device, batch_size, - num_heads, + num_qo_heads, max_topk=max_topk, max_extra_topk=max_extra_topk, - max_seq_len_q=max_seq_len_q, + max_q_len=max_q_len, packed_query=packed_query, - q_data_type=q_data_type, + q_dtype=q_dtype, + qk_rope_head_dim=qk_rope_head_dim, + kv_cache_format=kv_cache_format, + extra_kv_cache_format=extra_kv_cache_format, kv_layout=kv_layout, - has_sinks=has_sinks, + use_attention_sinks=use_attention_sinks, return_lse=return_lse, assume_valid_prefix=assume_valid_prefix, output_quant_block_size=output_quant_block_size, out_scale_format=out_scale_format, num_output_groups=num_output_groups, + tuning=tuning, + kv_lora_rank=kv_lora_rank, ) ) kernel = spec.kernel + reload_d576 = ( + profile.family == "keep" and kernel._make_config().sparse_d576_reload + ) family = profile.family splits = profile.split_kv - max_rows = batch_size * max_seq_len_q + max_rows = batch_size * max_q_len workspace = self._workspace_buffer if workspace is None: workspace = torch.empty(byte_end, device=device, dtype=torch.uint8) @@ -2394,101 +2975,100 @@ def plan( } # Finalize the joint distribution in the main kernel or its existing # split-KV reducer. No per-source outputs or finishing launch are needed. - kernel = copy.copy(kernel) - kernel.output_quant = output_quant - kernel.assume_valid_prefix = assume_valid_prefix - if family == "2cta": - kernel.fuse_sparse_epilogue = splits == 1 - kernel.fuse_sparse_reduction = splits > 1 - else: - kernel.finalize_output = True - kernel.direct_sparse = True - kernel.direct_sparse_capacities = (max_topk, max_extra_topk) - spec = replace( - spec, kernel=kernel, kernel_signature=_mla_kernel_compile_signature(kernel) - ) - compiled = _get_compiled_mla_decode(spec) - static_kernel = copy.copy(kernel) - static_kernel.direct_static_scales = True - compiled_static = _get_compiled_mla_decode( - replace( + if not reload_d576: + kernel = copy.copy(kernel) + kernel.output_quant = output_quant + kernel.assume_valid_prefix = assume_valid_prefix + if family == "2cta": + kernel.fuse_sparse_epilogue = splits == 1 + kernel.fuse_sparse_reduction = splits > 1 + else: + kernel.finalize_output = True + kernel.direct_sparse = True + kernel.direct_sparse_capacities = (max_topk, max_extra_topk) + spec = replace( spec, - kernel=static_kernel, - kernel_signature=_mla_kernel_compile_signature(static_kernel), + kernel=kernel, + kernel_signature=_mla_kernel_compile_signature(kernel), ) - ) + compiled = _get_compiled_mla_decode(spec) self._prepared.clear() self._state = dict( + reload_d576=reload_d576, device=device, - batch=batch_size, - heads=num_heads, - max_q=max_seq_len_q, + batch_size=batch_size, + num_qo_heads=num_qo_heads, + max_q_len=max_q_len, max_rows=max_rows, - packed=packed_query, - dtype=q_data_type, - ks=max_topk, - kc=max_extra_topk, + packed_query=packed_query, + q_dtype=q_dtype, + query_dim=512 + qk_rope_head_dim, + cache_dim=cache_record_size(kernel.kv_cache_format, qk_rope_head_dim), + extra_cache_dim=cache_record_size( + kernel.extra_kv_cache_format, qk_rope_head_dim + ), + cache_dtype=torch.uint8 if kernel.has_nvfp4_kv else q_dtype, + kv_cache_format=kernel.kv_cache_format, + extra_kv_cache_format=kernel.extra_kv_cache_format, + max_topk=max_topk, + max_extra_topk=max_extra_topk, kv_layout=kv_layout, - has_sinks=has_sinks, + use_attention_sinks=use_attention_sinks, return_lse=return_lse, output_quant=output_quant, buffers=buffers, workspace=workspace, compiled=compiled, - compiled_static=compiled_static, assume_valid_prefix=assume_valid_prefix, - splits=splits, + requested_splits=splits, default_cl=torch.full( (max_rows,), max_extra_topk, device=device, dtype=torch.int32 ), dummy_indices=torch.zeros((max_rows, 1), device=device, dtype=torch.int32), - dummy_cache=torch.zeros((1, 1, 512), device=device, dtype=q_data_type), - default_sinks=torch.full( - (num_heads,), -torch.inf, device=device, dtype=torch.float32 + dummy_cache=torch.zeros( + ( + 1, + 1, + cache_record_size(kernel.extra_kv_cache_format, qk_rope_head_dim), + ), + device=device, + dtype=torch.uint8 if kernel.has_nvfp4_kv else q_dtype, + ), + default_attention_sinks=torch.full( + (num_qo_heads,), -torch.inf, device=device, dtype=torch.float32 ), ) self.workspace_size_bytes = byte_end - self._scalar_cache.clear() - def _scalar(self, value, name, validate): - state = self._state - if isinstance(value, torch.Tensor): - if ( - value.dtype != torch.float32 - or value.device != state["device"] - or value.numel() != 1 - or not value.is_contiguous() - ): - raise ValueError(f"{name} must be a scalar CUDA FP32 tensor") - if validate and ( - not torch.isfinite(value).all().item() or (value <= 0).any().item() - ): - raise ValueError(f"{name} must be positive and finite") - return value.reshape(1) - if ( - isinstance(value, bool) - or not isinstance(value, (int, float)) - or not math.isfinite(value) - or value <= 0 - ): + @staticmethod + def host_scale(value, name): + """Validate and round a host scalar to the runtime kernel's FP32 ABI.""" + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise TypeError(f"{name} must be a Python int or float") + if not math.isfinite(value) or value <= 0: raise ValueError(f"{name} must be positive and finite") try: - key = struct.unpack("f", struct.pack("f", float(value)))[0] + rounded = struct.unpack("f", struct.pack("f", float(value)))[0] except OverflowError as error: raise ValueError(f"{name} must be representable in float32") from error - if not math.isfinite(key) or key <= 0: + if not math.isfinite(rounded) or rounded <= 0: raise ValueError(f"{name} must be positive finite float32") - if key not in self._scalar_cache: - self._scalar_cache[key] = torch.full( - (1,), key, device=state["device"], dtype=torch.float32 - ) - return self._scalar_cache[key] + return rounded def _pool(self, cache, name): + """Validate one cache pool and expose its physical rows without repacking storage.""" state = self._state + dim = state["extra_cache_dim" if name == "extra_kv_cache" else "cache_dim"] if not isinstance(cache, torch.Tensor): raise TypeError(f"{name} must be a tensor") + fmt = state[ + "extra_kv_cache_format" if name == "extra_kv_cache" else "kv_cache_format" + ] + if fmt == "fp8" and cache.dtype == torch.float8_e4m3fn: + cache = cache.view(torch.uint8) + if cache.ndim == 2: + cache = cache.unsqueeze(1) if cache.ndim == 4: axis = 1 if state["kv_layout"] == "HND" else 2 if cache.shape[axis] != 1: @@ -2497,30 +3077,30 @@ def _pool(self, cache, name): if ( cache.ndim != 3 or cache.shape[1] <= 0 - or cache.shape[2] != 512 + or cache.shape[2] != dim or cache.shape[0] <= 0 ): raise ValueError( - f"{name} must have shape [pages,page_size,512] with nonempty pages" + f"{name} must have shape [pages,page_size,{dim}] with nonempty pages" ) page_size = cache.shape[1] - if cache.device != state["device"] or cache.dtype != state["dtype"]: + if cache.device != state["device"] or cache.dtype != state["cache_dtype"]: raise ValueError(f"{name} must match the planned device and dtype") if ( cache.stride(-1) != 1 - or cache.stride(1) != 512 - or cache.stride(0) < page_size * 512 - or cache.stride(0) % 512 + or cache.stride(1) != dim + or cache.stride(0) < page_size * dim + or cache.stride(0) % dim or cache.data_ptr() % 16 ): raise ValueError( - f"{name} requires aligned compact rows and a page stride divisible by 512 elements" + f"{name} requires aligned compact rows and a page stride divisible by {dim} elements" ) - row_stride = cache.stride(0) // 512 + row_stride = cache.stride(0) // dim rows = (cache.shape[0] - 1) * row_stride + page_size if rows >= 0x7FFFFFFF: raise ValueError(f"{name} exceeds signed gather coordinate range") - return cache.as_strided((rows, 1, 512), (512, 512, 1)) + return cache.as_strided((rows, 1, dim), (dim, dim, 1)) def _source_metadata( self, indices, lengths, capacity, tokens, name, *, rows, validate @@ -2571,9 +3151,10 @@ def run( extra_indices=None, extra_lengths=None, qo_indptr=None, - softmax_scale=512**-0.5, + softmax_scale=None, output_scale=1.0, - sinks=None, + kv_cache_scale=1.0, + attention_sinks=None, out=None, lse=None, token_positions=None, @@ -2595,19 +3176,35 @@ def run( and live lengths need no 128-entry padding: the kernel aligns source execution spans internally and bounds all source reads. + ``validate=True`` (default) checks CUDA metadata contents using host + reads, which synchronize the device. Serving should establish index, + length, offset and position bounds upstream and use ``validate=False`` + on the hot path and during CUDA graph capture. Shape, dtype, device, + stride and alignment checks still run without reading CUDA values. + Q must already use the planned dtype and include the required query + RoPE; run() performs neither Q quantization nor query rotation. + ``softmax_scale`` and ``output_scale`` are positive finite Python - scalars or scalar CUDA FP32 tensors. They are already combined: + host scalars, passed as runtime FP32 kernel arguments. They are already combined: softmax_scale = attention_scale * q_descale * kv_descale; output_scale = output_multiplier * kv_descale. Both pools share that KV descale. These multiply QK scores and the weighted stored V values, respectively; sink logits are not multiplied by softmax_scale. The output multiplier applies before fused RoPE/quantization. ``out_scale`` instead holds the generated per-block quantization scales. - Combine dynamic scales in upstream work before each launch/replay and - keep their tensor addresses stable for CUDA Graphs. + Combine scales on the host before each launch. Their values are not + compilation keys. CUDA graphs capture the supplied values; changing + them for replay requires graph parameter updates or recapture. + + For ``kv_cache_format="fp8"``, ``kv_cache_scale`` is the independent + positive host-runtime per-tensor descale of the primary pool. It accepts + a Python scalar; leave it at the default 1.0 + for other formats. The extra pool is NVFP4 and has no tensor descale. + The primary scale is applied internally to both K and V; do not also + fold it into softmax_scale/output_scale. Caller-owned inputs and plan workspace must outlive graph replay, - including tensor scales and sinks, and must not alias writable output + including attention_sinks, and must not alias writable output or workspace storage. Q and native KV are used directly; flattening padded pages only creates a tensor view. Required attention reductions remain included. Caller-owned LSE must be contiguous FP32 and 16-byte @@ -2640,6 +3237,8 @@ def run( state = self._state if state is None: raise RuntimeError("plan() must be called before run()") + if softmax_scale is None: + softmax_scale = state["query_dim"] ** -0.5 lse_arg = lse prepared_key = ( kv_cache.data_ptr(), @@ -2654,7 +3253,7 @@ def run( extra_kv_cache.stride(), extra_kv_cache.dtype, ), - None if sinks is None else sinks.data_ptr(), + None if attention_sinks is None else attention_sinks.data_ptr(), softmax_scale, output_scale, ) @@ -2666,6 +3265,8 @@ def run( and token_positions is None and cos_sin_cache is None and out_scale is None + and type(kv_cache_scale) in (int, float) + and kv_cache_scale == 1.0 ): prepared = self._prepared.get(prepared_key) if prepared is not None: @@ -2674,32 +3275,36 @@ def run( ) _validate_query( query, - packed_query=state["packed"], + packed_query=state["packed_query"], device=state["device"], - batch_size=state["batch"], - num_heads=state["heads"], - max_seq_len_q=state["max_q"], - q_dtype=state["dtype"], - query_dim=512, + batch_size=state["batch_size"], + num_qo_heads=state["num_qo_heads"], + max_q_len=state["max_q_len"], + q_dtype=state["q_dtype"], + query_dim=state["query_dim"], ) prefix = tuple(query.shape[:-2]) rows = math.prod(prefix) - if state["packed"]: + if state["packed_query"]: _validate_qo_indptr( - qo_indptr, device=state["device"], batch_size=state["batch"] + qo_indptr, device=state["device"], batch_size=state["batch_size"] ) if validate: maximum, total, _ = _derive_max_seq_len_q( - qo_indptr, batch_size=state["batch"] + qo_indptr, batch_size=state["batch_size"] ) - if total != rows or maximum > state["max_q"]: + if total != rows or maximum > state["max_q_len"]: raise ValueError("invalid packed query offsets") elif qo_indptr is not None: raise ValueError("fixed queries do not accept qo_indptr") primary = self._pool(kv_cache, "kv_cache") if extra_kv_cache is None: - if state["kc"] or extra_indices is not None or extra_lengths is not None: + if ( + state["max_extra_topk"] + or extra_indices is not None + or extra_lengths is not None + ): raise ValueError("extra pool and indices are required by the plan") extra = state["dummy_cache"] else: @@ -2708,7 +3313,7 @@ def run( si, sl = self._source_metadata( indices, lengths, - state["ks"], + state["max_topk"], primary.shape[0], "primary", rows=rows, @@ -2717,32 +3322,30 @@ def run( ci, cl = self._source_metadata( extra_indices, extra_lengths, - state["kc"], + state["max_extra_topk"], extra.shape[0], "extra", rows=rows, validate=validate, ) - if state["has_sinks"] != (sinks is not None): - raise ValueError("sinks presence must match the plan") - if sinks is None: - sinks = state["default_sinks"] + if state["use_attention_sinks"] != (attention_sinks is not None): + raise ValueError("attention_sinks presence must match the plan") + if attention_sinks is None: + attention_sinks = state["default_attention_sinks"] if ( - sinks.shape != (state["heads"],) - or sinks.dtype != torch.float32 - or sinks.device != state["device"] - or not sinks.is_contiguous() + attention_sinks.shape != (state["num_qo_heads"],) + or attention_sinks.dtype != torch.float32 + or attention_sinks.device != state["device"] + or not attention_sinks.is_contiguous() ): - raise ValueError("sinks must be contiguous CUDA FP32[H]") - if validate and torch.isnan(sinks).any().item(): - raise ValueError("sinks must not contain NaN") - scale_tensors = [ - self._scalar(v, n, validate) - for v, n in ( - (softmax_scale, "softmax_scale"), - (output_scale, "output_scale"), - ) - ] + raise ValueError("attention_sinks must be contiguous CUDA FP32[H]") + if validate and torch.isnan(attention_sinks).any().item(): + raise ValueError("attention_sinks must not contain NaN") + softmax_scale = self.host_scale(softmax_scale, "softmax_scale") + output_scale = self.host_scale(output_scale, "output_scale") + kv_cache_scale = self.host_scale(kv_cache_scale, "kv_cache_scale") + if state["kv_cache_format"] != "fp8" and kv_cache_scale != 1.0: + raise ValueError("kv_cache_scale is only supported for an FP8 primary pool") quant = state["output_quant"] output_quant_args = None if quant is not None: @@ -2755,7 +3358,7 @@ def run( ) elif out is None: out = torch.empty( - (*prefix, state["heads"], 512), + (*prefix, state["num_qo_heads"], 512), device=state["device"], dtype=torch.bfloat16, ) @@ -2763,27 +3366,29 @@ def run( _validate_out( out, device=state["device"], - batch_size=state["batch"], - num_heads=state["heads"], - max_seq_len_q=state["max_q"], - packed_query=state["packed"], - total_q=rows if state["packed"] else None, - output_dtype=torch.bfloat16, + batch_size=state["batch_size"], + num_qo_heads=state["num_qo_heads"], + max_q_len=state["max_q_len"], + packed_query=state["packed_query"], + total_q=rows if state["packed_query"] else None, + out_dtype=torch.bfloat16, ) if lse is None: # A returned result must survive the next call on this wrapper. # Use plan-owned scratch only when the LSE is not exposed. lse = ( torch.empty( - (*prefix, state["heads"]), + (*prefix, state["num_qo_heads"]), device=state["device"], dtype=torch.float32, ) if state["return_lse"] - else state["buffers"]["public_lse"][:rows].view(*prefix, state["heads"]) + else state["buffers"]["public_lse"][:rows].view( + *prefix, state["num_qo_heads"] + ) ) if ( - lse.shape != (*prefix, state["heads"]) + lse.shape != (*prefix, state["num_qo_heads"]) or lse.dtype != torch.float32 or lse.device != state["device"] or not lse.is_contiguous() @@ -2792,66 +3397,86 @@ def run( raise ValueError( "lse must be contiguous, 16-byte-aligned FP32 with one value per query/head" ) - static_scales = all( - not isinstance(v, torch.Tensor) for v in (softmax_scale, output_scale) - ) - # Both paths consume the same effective scales; no Q/K/V scale - # composition or auxiliary preparation is performed here. - bmm1_scale = float(softmax_scale) if static_scales else 1.0 - bmm2_scale = float(output_scale) if static_scales else 1.0 sparse_inputs = ( si, ci, sl, cl, - scale_tensors[0], - scale_tensors[1], - sinks, - ) - if rows: + attention_sinks, + kv_cache_scale, + ) + if rows and state["reload_d576"]: + core = state["buffers"]["core"] + # One token per batch row through the 1CTA page-table ABI; 128 heads + # are two physical query positions of 64 heads. + phys = (rows, state["num_qo_heads"] // 64, 64) + with torch.cuda.device(state["device"]): + _launch_mla_decode( + _MLARuntime( + query.view(*phys, _D576_HEAD_DIM), + primary, + out.view(*phys, _D576_V_DIM), + primary.shape[0], + softmax_scale, + output_scale, + ), + block_tables=si, + seq_lens_kv=sl, + qo_indptr=None, + packed_query=False, + kv_lora_rank=512, + split_kv=state["requested_splits"], + workspace=_MLAWorkspaceViews( + core if core.numel() else None, lse.view(*phys) + ), + compiled=state["compiled"], + max_q_len=phys[1], + ) + elif rows: buffers = state["buffers"] # Quantized stores use output_quant_args, including group/bucket # strides. The core O view supplies compact launch geometry only. core_out = ( - out.view(rows, 1, state["heads"], 512) + out.view(rows, 1, state["num_qo_heads"], 512) if quant is None - else output_quant_args[0][: rows * state["heads"] * 512].view( - rows, 1, state["heads"], 512 + else output_quant_args[0][: rows * state["num_qo_heads"] * 512].view( + rows, 1, state["num_qo_heads"], 512 ) ) _launch_mla_decode( _MLARuntime( - query.view(rows, 1, state["heads"], 512), + query.view(rows, 1, state["num_qo_heads"], state["query_dim"]), primary, core_out, primary.shape[0], - bmm1_scale, - bmm2_scale, + softmax_scale, + output_scale, extra_cache=extra, sparse_inputs=sparse_inputs, output_quant_args=output_quant_args, ), block_tables=si, - seq_lens=sl, + seq_lens_kv=sl, qo_indptr=None, packed_query=False, kv_lora_rank=512, - split_kv=state["splits"], + split_kv=state["requested_splits"], workspace=_MLAWorkspaceViews( buffers["core"] if buffers["core"].numel() else None, - lse.view(rows, 1, state["heads"]), + lse.view(rows, 1, state["num_qo_heads"]), ), - compiled=state["compiled_static"] - if static_scales - else state["compiled"], + compiled=state["compiled"], + max_q_len=1, ) if ( rows and quant is None - and static_scales - and not state["packed"] + and not state["packed_query"] and not state["return_lse"] - and state["max_q"] == 1 + and state["max_q_len"] == 1 + and state["query_dim"] == 512 + and state["kv_cache_format"] == "native" + and state["extra_kv_cache_format"] == "native" and lse_arg is None ): # Static for this (plan, pools, scales): pool views, scalars, sinks, split, @@ -2865,10 +3490,10 @@ def run( primary_rows=primary.shape[0], c_latent=primary.permute(1, 2, 0), c_rope=extra.permute(1, 2, 0), - scalars=(scale_tensors[0], scale_tensors[1], sinks), - bmm=(bmm1_scale, bmm2_scale), - ks=state["ks"], - kc=state["kc"], + sinks=attention_sinks, + bmm=(softmax_scale, output_scale), + ks=state["max_topk"], + kc=state["max_extra_topk"], dummy_indices=state["dummy_indices"], default_cl=state["default_cl"], ) @@ -2886,10 +3511,10 @@ def _run_prepared(self, prepared, query, indices, lengths, extra_indices, extra_ full path for the same arguments. """ state = self._state - rows, heads, device = state["batch"], state["heads"], state["device"] + rows, heads, device = state["batch_size"], state["num_qo_heads"], state["device"] if ( query.shape != (rows, 1, heads, 512) - or query.dtype != state["dtype"] + or query.dtype != state["q_dtype"] or query.device != device or not query.is_contiguous() or query.data_ptr() % 16 @@ -2940,11 +3565,11 @@ def _run_prepared(self, prepared, query, indices, lengths, extra_indices, extra_ raise ValueError("prepared launch: extra indices without a planned extra pool") ci, cl = prepared["dummy_indices"][:rows], prepared["default_cl"][:rows] buffers = state["buffers"] - s0, s1, sinks = prepared["scalars"] + sinks = prepared["sinks"] q_kernel = query.permute(2, 3, 1, 0) lse_kernel = buffers["public_lse"][:rows].view(rows, 1, heads).permute(2, 1, 0) bmm1, bmm2 = prepared["bmm"] - state["compiled_static"]( + state["compiled"]( q_kernel, q_kernel, prepared["c_latent"], @@ -2953,13 +3578,14 @@ def _run_prepared(self, prepared, query, indices, lengths, extra_indices, extra_ out.view(rows, 1, heads, 512).permute(2, 3, 1, 0), lse_kernel, buffers["core"] if buffers["core"].numel() else None, - state["splits"], + state["requested_splits"], lengths, None, - (indices, ci, lengths, cl, s0, s1, sinks), + (indices, ci, lengths, cl, sinks, 1.0), bmm1, bmm2, None, + 1, ) return out @@ -2986,11 +3612,14 @@ def batch_sparse_mla_decode_with_paged_kv_cache( extra_indices=None, extra_lengths=None, qo_indptr=None, - max_seq_len_q=None, + max_q_len=None, kv_layout="NHD", - softmax_scale=512**-0.5, + kv_cache_format="native", + extra_kv_cache_format=None, + softmax_scale=None, output_scale=1.0, - sinks=None, + kv_cache_scale=1.0, + attention_sinks=None, out=None, lse=None, return_lse=False, @@ -3002,13 +3631,15 @@ def batch_sparse_mla_decode_with_paged_kv_cache( token_positions=None, cos_sin_cache=None, out_scale=None, + tuning=None, ): """Eager plan-and-run helper using separate source indices and lengths. Use a planned wrapper for CUDA Graph replay. Index conversion and scale composition remain external. softmax_scale includes attention/Q/K descales; output_scale includes the V descale and optional output multiplier. - Source contracts, output quantization arguments and return layouts follow + Source contracts, output quantization arguments, ``tuning`` and return + layouts follow :class:`BatchSparseMLADecodePagedTSWrapper`. """ if torch.cuda.is_current_stream_capturing(): @@ -3020,11 +3651,11 @@ def batch_sparse_mla_decode_with_paged_kv_cache( if qo_indptr is None: raise ValueError("packed queries require qo_indptr") batch = qo_indptr.numel() - 1 - if max_seq_len_q is None: - max_seq_len_q = max(1, int((qo_indptr[1:] - qo_indptr[:-1]).max().item())) + if max_q_len is None: + max_q_len = max(1, int((qo_indptr[1:] - qo_indptr[:-1]).max().item())) else: batch = query.shape[0] - max_seq_len_q = query.shape[1] if max_seq_len_q is None else max_seq_len_q + max_q_len = query.shape[1] if max_q_len is None else max_q_len wrapper = BatchSparseMLADecodePagedTSWrapper(workspace_buffer) wrapper.plan( query.device, @@ -3032,16 +3663,20 @@ def batch_sparse_mla_decode_with_paged_kv_cache( query.shape[-2], max_topk=indices.shape[-1], max_extra_topk=0 if extra_indices is None else extra_indices.shape[-1], - max_seq_len_q=max_seq_len_q, + max_q_len=max_q_len, packed_query=packed, - q_data_type=query.dtype, + q_dtype=query.dtype, + qk_rope_head_dim=query.shape[-1] - 512, + kv_cache_format=kv_cache_format, + extra_kv_cache_format=extra_kv_cache_format, kv_layout=kv_layout, - has_sinks=sinks is not None, + use_attention_sinks=attention_sinks is not None, return_lse=return_lse, assume_valid_prefix=assume_valid_prefix, output_quant_block_size=output_quant_block_size, out_scale_format=out_scale_format, num_output_groups=num_output_groups, + tuning=tuning, ) return wrapper.run( query, @@ -3054,7 +3689,8 @@ def batch_sparse_mla_decode_with_paged_kv_cache( qo_indptr=qo_indptr, softmax_scale=softmax_scale, output_scale=output_scale, - sinks=sinks, + kv_cache_scale=kv_cache_scale, + attention_sinks=attention_sinks, out=out, lse=lse, token_positions=token_positions, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/q_token_kv_block_sparse_metadata.py b/tensorrt_llm/_torch/attention/backends/prims_ts/q_token_kv_block_sparse_metadata.py index dff16046cf7c..142cdb1ec83a 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/q_token_kv_block_sparse_metadata.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/q_token_kv_block_sparse_metadata.py @@ -38,7 +38,7 @@ import functools import math import threading -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import Any, Literal, Optional import torch @@ -243,23 +243,25 @@ class _PrimsTSQTokenKvBlockSparseMetadataPlan: qo_indptr: Optional[torch.Tensor] metadata_run: Any group_size: int - use_packed_q: bool + packed_query: bool sparse_block_size: int storage_page_size: int max_seq_len_kv: int release_attention_pdl: bool locator_head_multiplier: int = 1 + count_current_block_in_topk: bool = True + fixed_query_len: int = 0 def run( self, block_indices: torch.Tensor, block_table: torch.Tensor, - token_to_request: torch.Tensor, + token_to_request: Optional[torch.Tensor], query_positions: torch.Tensor, ) -> None: """Launch the prepared CUDA metadata kernel without shape resolution.""" - if self.use_packed_q: + if self.packed_query: self.metadata_run( block_indices, block_table, @@ -275,6 +277,8 @@ def run( self.max_seq_len_kv, self.release_attention_pdl, self.locator_head_multiplier, + self.count_current_block_in_topk, + self.fixed_query_len, ) else: self.metadata_run( @@ -291,6 +295,8 @@ def run( self.max_seq_len_kv, self.release_attention_pdl, self.locator_head_multiplier, + self.count_current_block_in_topk, + self.fixed_query_len, ) @@ -319,7 +325,7 @@ class _PrimsTSQTokenKvBlockSparsePlan: _query: _QTokenKvBlockSparseTensorDescriptor _block_indices: _QTokenKvBlockSparseTensorDescriptor _block_table: _QTokenKvBlockSparseTensorDescriptor - _token_to_request: _QTokenKvBlockSparseTensorDescriptor + _token_to_request: Optional[_QTokenKvBlockSparseTensorDescriptor] _query_positions: _QTokenKvBlockSparseTensorDescriptor _out: _QTokenKvBlockSparseTensorDescriptor _fixed_query_group_size: Optional[int] = None @@ -329,7 +335,7 @@ def run( query: torch.Tensor, block_indices: torch.Tensor, block_table: torch.Tensor, - token_to_request: torch.Tensor, + token_to_request: Optional[torch.Tensor], query_positions: torch.Tensor, *, out: torch.Tensor, @@ -345,11 +351,17 @@ def run( _validate_q_token_kv_block_sparse_plan_tensor( block_table, "block_table", self._block_table ) - _validate_q_token_kv_block_sparse_plan_tensor( - token_to_request, - "token_to_request", - self._token_to_request, - ) + if self._token_to_request is None: + if token_to_request is not None: + raise ValueError("the fixed-Q plan derives request IDs") + else: + if token_to_request is None: + raise ValueError("the prepared plan requires explicit request IDs") + _validate_q_token_kv_block_sparse_plan_tensor( + token_to_request, + "token_to_request", + self._token_to_request, + ) _validate_q_token_kv_block_sparse_plan_tensor( query_positions, "query_positions", @@ -418,12 +430,13 @@ class _QTokenKvBlockSparsePlanConfig: max_seq_len_kv: int device: torch.device workspace_buffer: torch.Tensor - use_packed_q: bool + packed_query: bool split_kv: bool share_pattern_across_kv_heads: bool - q_data_type: torch.dtype - kv_data_type: torch.dtype - o_data_type: torch.dtype + count_current_block_in_topk: bool + q_dtype: torch.dtype + kv_dtype: torch.dtype + out_dtype: torch.dtype def _q_token_tensor_abi_key( @@ -484,13 +497,14 @@ def plan( device: torch.device | str | int, workspace_buffer: torch.Tensor, group_size: int = 1, - use_packed_q: bool = False, + packed_query: bool = False, split_kv: bool = True, share_pattern_across_kv_heads: bool = True, + count_current_block_in_topk: bool = True, mask_type: Literal["causal"] = "causal", - q_data_type: torch.dtype = torch.float16, - kv_data_type: Optional[torch.dtype] = None, - o_data_type: Optional[torch.dtype] = None, + q_dtype: torch.dtype = torch.float16, + kv_dtype: Optional[torch.dtype] = None, + out_dtype: Optional[torch.dtype] = None, ) -> None: """Publish a capacity-only QToken-KvBlock-Sparse-Attention plan. @@ -507,9 +521,20 @@ def plan( top-k ``[T,K]`` input. False requires ``[T,Hkv,K]`` and prepares one union per group and KV head. Q, K, V and output layouts are unchanged. + ``count_current_block_in_topk=True`` selects completed blocks and + appends only an incomplete current block. False excludes the current + block from the top-k budget and always appends its visible tokens, + including the entire block when the query is its last token. + + For example, with block size 8, top-k budget 2, and query position 23, + True selects two IDs from blocks 0/1/2 and appends nothing; IDs [0, 1] + attend to 16 tokens. False selects blocks [0, 1] and appends complete + current block 2 outside the budget, attending to 24 tokens. At position + 22 both modes append the seven visible tokens of partial block 2. + ``split_kv=True`` permits automatic useful split fanout; False forces one KV partition and omits split scratch. Pass False for prefill and - True for decode, independently of ``use_packed_q``. Workspace sizing + True for decode, independently of ``packed_query``. Workspace sizing must use the same value. Change it only by replanning outside capture. """ @@ -535,19 +560,19 @@ def plan( raise ValueError( "QToken-KvBlock-Sparse-Attention requires mask_type='causal'" ) - if not isinstance(use_packed_q, bool): - raise TypeError("use_packed_q must be a bool") + if not isinstance(packed_query, bool): + raise TypeError("packed_query must be a bool") for dtype, name in ( - (q_data_type, "q_data_type"), - (kv_data_type, "kv_data_type"), - (o_data_type, "o_data_type"), + (q_dtype, "q_dtype"), + (kv_dtype, "kv_dtype"), + (out_dtype, "out_dtype"), ): if dtype is not None and not isinstance(dtype, torch.dtype): raise TypeError(f"{name} must be a torch.dtype") - if kv_data_type is None: - kv_data_type = q_data_type - if o_data_type is None: - o_data_type = q_data_type + if kv_dtype is None: + kv_dtype = q_dtype + if out_dtype is None: + out_dtype = q_dtype planned_device, device_index = _resolve_cuda_device(device) _validate_runtime_device(planned_device) if not isinstance(workspace_buffer, torch.Tensor): @@ -566,15 +591,16 @@ def plan( num_qo_heads=num_qo_heads, num_kv_heads=num_kv_heads, head_dim=head_dim, - q_dtype=q_data_type, - kv_dtype=kv_data_type, - out_dtype=o_data_type, + q_dtype=q_dtype, + kv_dtype=kv_dtype, + out_dtype=out_dtype, device=planned_device, num_query_groups=batch_size, - use_packed_q=use_packed_q, + packed_query=packed_query, sparse_block_size=kv_block_size, split_kv=split_kv, share_pattern_across_kv_heads=share_pattern_across_kv_heads, + count_current_block_in_topk=count_current_block_in_topk, ) _validate_q_token_kv_block_sparse_attention_workspace( workspace_buffer, layout.total_bytes @@ -597,12 +623,13 @@ def plan( max_seq_len_kv=max_seq_len_kv, device=planned_device, workspace_buffer=workspace_buffer, - use_packed_q=use_packed_q, + packed_query=packed_query, split_kv=split_kv, share_pattern_across_kv_heads=share_pattern_across_kv_heads, - q_data_type=q_data_type, - kv_data_type=kv_data_type, - o_data_type=o_data_type, + count_current_block_in_topk=count_current_block_in_topk, + q_dtype=q_dtype, + kv_dtype=kv_dtype, + out_dtype=out_dtype, ) with self._plan_lock: self._config = candidate @@ -623,9 +650,10 @@ def run( paged_kv_cache: tuple[torch.Tensor, torch.Tensor], block_table: torch.Tensor, indexer_block_ids: torch.Tensor, - token_to_request: torch.Tensor, - query_positions: torch.Tensor, + token_to_request: Optional[torch.Tensor] = None, + query_positions: Optional[torch.Tensor] = None, *, + seq_lens: Optional[torch.Tensor] = None, qo_indptr: Optional[torch.Tensor] = None, sm_scale: Optional[float] = None, v_scale: Optional[float] = None, @@ -640,13 +668,21 @@ def run( contiguous last dimension. It never uses the BSR name ``block_indices``. K/V are separate HND caches shaped ``[num_pages, Hkv, page_size, D]``. - For every valid query at zero-based position ``p``, the first - ``min(block_topk, (p + 1) // kv_block_size)`` indexer entries must be - distinct completed-block IDs in ``[0, (p + 1) // kv_block_size)``. + Request-aligned fixed Q may omit both per-token metadata tensors and + supply ``seq_lens`` instead. For each request, all ``Nq * G`` queries + are consecutive and end at its KV length. Request IDs and positions + are derived inside the kernels; nonpositive lengths mark inactive + requests. Packed Q and arbitrary mappings retain explicit metadata. + + For every valid query at zero-based position ``p``, let ``bound`` be + ``(p + int(count_current_block_in_topk)) // kv_block_size``. The first + ``min(block_topk, bound)`` indexer entries must be distinct block IDs + in ``[0, bound)``. Their order is unrestricted; later entries are ignored. In particular, Q1 does not compact missing entries or arbitrary ``-1`` holes inside - this required prefix. The incomplete causal tail is derived from - ``query_positions`` and must not be inserted into the selected prefix. + this required prefix. The appended current block is derived from + ``query_positions`` and must not also appear in the selected prefix + when ``count_current_block_in_topk=False``. The KV-cache writer must zero-fill unused token slots in allocated pages before attention. TMA loads complete fragments, including the @@ -664,6 +700,8 @@ def run( q : torch.Tensor Packed ``[total_q, Hq, D]`` or fixed ``[B, Nq, G, Hq, D]`` query, with the planned device, dtype and head geometry. + Packed queries use a TensorMap bounded by the physical tensor; + partial final routes do not require extra tail storage. paged_kv_cache : tuple[torch.Tensor, torch.Tensor] Separate HND K/V caches shaped ``[num_pages, Hkv, page_size, D]``. @@ -675,10 +713,16 @@ def run( shares patterns, otherwise ``[T, Hkv, block_topk]``. The last dimension must be contiguous; each row supplies the distinct completed-block prefix above. The causal tail is added internally. - token_to_request : torch.Tensor - CUDA Int32 request ID for each flattened query token. - query_positions : torch.Tensor + token_to_request : torch.Tensor, optional + CUDA Int32 request ID for each flattened query token. Omit with + ``query_positions`` when supplying fixed-Q ``seq_lens``. + query_positions : torch.Tensor, optional CUDA Int32 or Int64 zero-based causal position for each query token. + seq_lens : torch.Tensor, optional + Contiguous CUDA Int32 or Int64 KV lengths, including all queries, + shaped ``[B]`` for request-aligned fixed Q. Mutually exclusive + with explicit per-token metadata. Active lengths must be at least + ``Nq * G`` and at most the planned model length. qo_indptr : torch.Tensor, optional CUDA Int32 packed-route offsets. Required for packed Q and omitted for fixed Q; each route must stay within one request and contain @@ -703,9 +747,9 @@ def run( config = self._require_config() if not isinstance(q, torch.Tensor): raise TypeError("q must be a torch.Tensor") - if q.device != config.device or q.dtype != config.q_data_type: - raise ValueError("q must use the planned device and q_data_type") - if config.use_packed_q: + if q.device != config.device or q.dtype != config.q_dtype: + raise ValueError("q must use the planned device and q_dtype") + if config.packed_query: if q.ndim != 3 or qo_indptr is None: raise ValueError("packed Q requires [total_q,Hq,D] and qo_indptr") _validate_q_token_kv_block_sparse_qo_indptr_tensor( @@ -730,7 +774,7 @@ def run( raise ValueError("q must use the planned query-head and head dimensions") num_query_tokens = ( int(q.shape[0]) - if config.use_packed_q + if config.packed_query else (config.batch_size * config.group_size) ) expected_ids_shape = ( @@ -743,24 +787,44 @@ def run( or indexer_block_ids.shape != expected_ids_shape ): raise ValueError(f"indexer_block_ids must have shape {expected_ids_shape}") + fixed_query_len = 0 + if seq_lens is not None: + if ( + config.packed_query + or token_to_request is not None + or query_positions is not None + ): + raise ValueError( + "seq_lens requires fixed Q without explicit token metadata" + ) + fixed_query_len = int(q.shape[1] * q.shape[2]) + query_positions = seq_lens for tensor, name in ( (token_to_request, "token_to_request"), (query_positions, "query_positions"), (block_table, "block_table"), ): + if name == "token_to_request" and fixed_query_len: + continue if not isinstance(tensor, torch.Tensor): raise TypeError(f"{name} must be a torch.Tensor") - if token_to_request.shape != (num_query_tokens,): + if token_to_request is not None and token_to_request.shape != ( + num_query_tokens, + ): raise ValueError("token_to_request must have one entry per query token") - if query_positions.shape != (num_query_tokens,): - raise ValueError("query_positions must have one entry per query token") + assert query_positions is not None + metadata_rows = int(q.shape[0]) if fixed_query_len else num_query_tokens + if query_positions.shape != (metadata_rows,): + raise ValueError( + "expected one KV length per request or one position per query" + ) k_cache, v_cache = _validate_q_token_kv_block_sparse_paged_kv_cache( paged_kv_cache ) if ( k_cache.device != config.device - or k_cache.dtype != config.kv_data_type + or k_cache.dtype != config.kv_dtype or tuple(k_cache.shape[1:]) != (config.num_kv_heads, config.page_size, config.head_dim) ): @@ -783,18 +847,19 @@ def run( raise RuntimeError( "run once eagerly before capture when out is not supplied" ) - self._default_out = torch.empty_like(q, dtype=config.o_data_type) + self._default_out = torch.empty_like(q, dtype=config.out_dtype) out = self._default_out elif not isinstance(out, torch.Tensor): raise TypeError("out must be a torch.Tensor") elif ( out.shape != q.shape or out.device != config.device - or out.dtype != config.o_data_type + or out.dtype != config.out_dtype ): raise ValueError("out must match q shape and the planned output dtype") prepared_key = ( + fixed_query_len, _q_token_tensor_abi_key(q), _q_token_tensor_abi_key(indexer_block_ids), _q_token_tensor_abi_key(block_table), @@ -828,6 +893,8 @@ def run( sparse_block_size=config.kv_block_size, split_kv=config.split_kv, share_pattern_across_kv_heads=config.share_pattern_across_kv_heads, + count_current_block_in_topk=config.count_current_block_in_topk, + fixed_query_len=fixed_query_len, ) self._prepared_key = prepared_key assert self._prepared_plan is not None @@ -940,10 +1007,11 @@ def _get_prims_ts_q_token_kv_block_sparse_workspace_layout( out_dtype: Optional[torch.dtype] = None, device: Optional[torch.device | str | int] = None, num_query_groups: Optional[int] = None, - use_packed_q: bool = False, + packed_query: bool = False, sparse_block_size: int = 4, split_kv: bool = True, share_pattern_across_kv_heads: bool = True, + count_current_block_in_topk: bool = True, ) -> _PrimsTSQTokenKvBlockSparseWorkspaceLayout: """Return the unified allocation layout for QToken-KvBlock-Sparse-Attention metadata and attention. @@ -996,6 +1064,8 @@ def _get_prims_ts_q_token_kv_block_sparse_workspace_layout( group_size, num_query_groups, ) + if not isinstance(count_current_block_in_topk, bool): + raise TypeError("count_current_block_in_topk must be a bool") if not isinstance(share_pattern_across_kv_heads, bool): raise TypeError("share_pattern_across_kv_heads must be a bool") pattern_heads = _num_sparse_pattern_heads( @@ -1010,7 +1080,9 @@ def _get_prims_ts_q_token_kv_block_sparse_workspace_layout( page_capacity = block_capacity * (sparse_block_size // fragment_size) # Padding is an allocation property, not additional selected KV work. max_seq_len = ( - block_topk * sparse_block_size + (sparse_block_size - 1) + block_topk * sparse_block_size + + sparse_block_size + - int(count_current_block_in_topk) if group_size == 1 else page_capacity * fragment_size ) @@ -1035,7 +1107,7 @@ def _get_prims_ts_q_token_kv_block_sparse_workspace_layout( _dtype_key(out_dtype), "HND", "causal", - use_packed_q, + packed_query, -1, storage_page_size, True, @@ -1048,7 +1120,7 @@ def _get_prims_ts_q_token_kv_block_sparse_workspace_layout( query_major_memberships = ( sparse_block_size == 4 and spec.config.use_fp8_qkv - and spec.config.supports_query_major_memberships + and replace(spec.config, fixed_seqlen=True).supports_query_major_memberships ) if query_major_memberships: page_capacity = (page_capacity + 31) // 32 * 32 @@ -1110,12 +1182,13 @@ def get_q_token_kv_block_sparse_workspace_size( *, block_topk: int, max_seq_len_kv: int, - o_data_type: Optional[torch.dtype] = None, + out_dtype: Optional[torch.dtype] = None, qo_indptr: Optional[torch.Tensor] = None, group_size: int = 1, kv_block_size: int = 4, split_kv: bool = True, share_pattern_across_kv_heads: bool = True, + count_current_block_in_topk: bool = True, ) -> int: """Return bytes for QToken-KvBlock-Sparse-Attention workspace storage. @@ -1155,7 +1228,7 @@ def get_q_token_kv_block_sparse_workspace_size( Static maximum visible logical K/V length in tokens, including current query or MTP tokens. It may be smaller than the reserved block-table capacity and must not change across replay of a prepared CUDA graph. - o_data_type : torch.dtype, optional + out_dtype : torch.dtype, optional Output dtype used to size attention partials. It defaults to the query dtype. qo_indptr : torch.Tensor, optional @@ -1175,6 +1248,12 @@ def get_q_token_kv_block_sparse_workspace_size( False uses independent per-KV-head patterns. Match this setting in workspace sizing and planning; it does not change the attention grid. + count_current_block_in_topk : bool + True (default) counts a completed current block in the selection budget. + False excludes it from selected IDs and always appends its visible tokens. + Match this setting in workspace sizing and planning; False can require + a larger workspace. + Returns ------- int @@ -1188,12 +1267,13 @@ def get_q_token_kv_block_sparse_workspace_size( block_table, block_topk, max_seq_len_kv=max_seq_len_kv, - out_dtype=o_data_type, + out_dtype=out_dtype, qo_indptr=qo_indptr, group_size=group_size, sparse_block_size=kv_block_size, split_kv=split_kv, share_pattern_across_kv_heads=share_pattern_across_kv_heads, + count_current_block_in_topk=count_current_block_in_topk, ).total_bytes @@ -1210,7 +1290,9 @@ def _get_prims_ts_q_token_kv_block_sparse_workspace_layout_from_tensors( sparse_block_size: int = 4, split_kv: bool = True, share_pattern_across_kv_heads: bool = True, + count_current_block_in_topk: bool = True, ) -> _PrimsTSQTokenKvBlockSparseWorkspaceLayout: + """Infer query geometry and size scratch for the selected sparse-block policy.""" sparse_block_size = _validate_sparse_block_size(sparse_block_size) if not isinstance(query, torch.Tensor): raise TypeError("query must be a torch.Tensor") @@ -1220,8 +1302,8 @@ def _get_prims_ts_q_token_kv_block_sparse_workspace_layout_from_tensors( raise TypeError("k_cache must be a torch.Tensor") if not isinstance(block_table, torch.Tensor): raise TypeError("block_table must be a torch.Tensor") - use_packed_q = qo_indptr is not None - if use_packed_q: + packed_query = qo_indptr is not None + if packed_query: if query.ndim != 3: raise ValueError( "packed QToken-KvBlock-Sparse-Attention query must have shape [total_q,Hq,D]" @@ -1260,7 +1342,7 @@ def _get_prims_ts_q_token_kv_block_sparse_workspace_layout_from_tensors( if out_dtype is None: out_dtype = query.dtype return _get_prims_ts_q_token_kv_block_sparse_workspace_layout( - int(query.shape[0]) if use_packed_q else groups * group_size, + int(query.shape[0]) if packed_query else groups * group_size, block_topk, block_table.shape[1], k_cache.shape[2], @@ -1274,10 +1356,11 @@ def _get_prims_ts_q_token_kv_block_sparse_workspace_layout_from_tensors( out_dtype=out_dtype, device=query.device, num_query_groups=groups, - use_packed_q=use_packed_q, + packed_query=packed_query, sparse_block_size=sparse_block_size, split_kv=split_kv, share_pattern_across_kv_heads=share_pattern_across_kv_heads, + count_current_block_in_topk=count_current_block_in_topk, ) @@ -1295,6 +1378,7 @@ def _build_prims_ts_q_token_kv_block_sparse_metadata( out: Optional[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = None, release_attention_pdl: bool = False, share_pattern_across_kv_heads: bool = True, + count_current_block_in_topk: bool = True, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Build dense QToken-KvBlock-Sparse-Attention indices and memberships from compact sparse-block IDs. @@ -1338,16 +1422,17 @@ def _build_prims_ts_q_token_kv_block_sparse_metadata( max_seq_len_kv, sparse_block_size, share_pattern_across_kv_heads, + count_current_block_in_topk, ) rows, block_topk = block_indices.shape[0], block_indices.shape[-1] pattern_heads = block_indices.shape[1] if block_indices.ndim == 3 else 1 - use_packed_q = qo_indptr is not None - if use_packed_q: + packed_query = qo_indptr is not None + if packed_query: _validate_q_token_kv_block_sparse_qo_indptr_tensor( qo_indptr, expected_device=block_indices.device, ) - groups = int(qo_indptr.numel()) - 1 if use_packed_q else None + groups = int(qo_indptr.numel()) - 1 if packed_query else None expected_shapes = _get_q_token_kv_block_sparse_metadata_output_shapes( rows, block_topk, @@ -1391,7 +1476,7 @@ def _build_prims_ts_q_token_kv_block_sparse_metadata( enable_pdl = device_support_pdl(block_indices.device) metadata_module = _get_prims_ts_q_token_kv_block_sparse_metadata_module() release_attention_pdl = enable_pdl and release_attention_pdl - if use_packed_q: + if packed_query: metadata_module.run_packed( block_indices, block_table, @@ -1407,6 +1492,8 @@ def _build_prims_ts_q_token_kv_block_sparse_metadata( max_seq_len_kv, release_attention_pdl, 1, + count_current_block_in_topk, + 0, ) else: metadata_module.run_fixed( @@ -1423,6 +1510,8 @@ def _build_prims_ts_q_token_kv_block_sparse_metadata( max_seq_len_kv, release_attention_pdl, 1, + count_current_block_in_topk, + 0, ) return outputs @@ -1440,6 +1529,7 @@ def _build_q_token_kv_block_sparse_metadata( qo_indptr: Optional[torch.Tensor] = None, out: Optional[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = None, share_pattern_across_kv_heads: bool = True, + count_current_block_in_topk: bool = True, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Build dense QToken-KvBlock-Sparse-Attention page indices and memberships from sparse-block IDs. @@ -1521,13 +1611,14 @@ def _build_q_token_kv_block_sparse_metadata( qo_indptr=qo_indptr, out=out, share_pattern_across_kv_heads=share_pattern_across_kv_heads, + count_current_block_in_topk=count_current_block_in_topk, ) def _prepare_prims_ts_q_token_kv_block_sparse_metadata_plan( block_indices: torch.Tensor, block_table: torch.Tensor, - token_to_request: torch.Tensor, + token_to_request: Optional[torch.Tensor], query_positions: torch.Tensor, q_token_kv_block_sparse_page_indices: torch.Tensor, q_token_kv_block_sparse_page_memberships: torch.Tensor, @@ -1540,6 +1631,8 @@ def _prepare_prims_ts_q_token_kv_block_sparse_metadata_plan( qo_indptr: Optional[torch.Tensor] = None, query_major_memberships: bool = False, locator_head_multiplier: int = 1, + count_current_block_in_topk: bool = True, + fixed_query_len: int = 0, ) -> _PrimsTSQTokenKvBlockSparseMetadataPlan: """Freeze validated metadata tensors and launch constants. @@ -1553,11 +1646,11 @@ def _prepare_prims_ts_q_token_kv_block_sparse_metadata_plan( sparse_block_size = _validate_sparse_block_size(sparse_block_size) rows = block_indices.shape[0] - use_packed_q = qo_indptr is not None + packed_query = qo_indptr is not None _resolve_num_query_groups( rows, group_size, - int(qo_indptr.numel()) - 1 if use_packed_q else None, + int(qo_indptr.numel()) - 1 if packed_query else None, ) max_seq_len_kv = _validate_q_token_kv_block_sparse_max_seq_len_kv( max_seq_len_kv, @@ -1568,12 +1661,12 @@ def _prepare_prims_ts_q_token_kv_block_sparse_metadata_plan( if query_major_memberships: metadata_run = ( metadata_module.run_packed_query_major - if use_packed_q + if packed_query else metadata_module.run_fixed_query_major ) else: metadata_run = ( - metadata_module.run_packed if use_packed_q else metadata_module.run_fixed + metadata_module.run_packed if packed_query else metadata_module.run_fixed ) return _PrimsTSQTokenKvBlockSparseMetadataPlan( q_token_kv_block_sparse_page_indices=q_token_kv_block_sparse_page_indices, @@ -1582,12 +1675,14 @@ def _prepare_prims_ts_q_token_kv_block_sparse_metadata_plan( qo_indptr=qo_indptr, metadata_run=metadata_run, group_size=group_size, - use_packed_q=use_packed_q, + packed_query=packed_query, sparse_block_size=sparse_block_size, storage_page_size=storage_page_size, max_seq_len_kv=max_seq_len_kv, release_attention_pdl=release_attention_pdl, locator_head_multiplier=locator_head_multiplier, + count_current_block_in_topk=count_current_block_in_topk, + fixed_query_len=fixed_query_len, ) @@ -1596,7 +1691,7 @@ def _prepare_q_token_kv_block_sparse_attention( paged_kv_cache: tuple[torch.Tensor, torch.Tensor], block_indices: torch.Tensor, block_table: torch.Tensor, - token_to_request: torch.Tensor, + token_to_request: Optional[torch.Tensor], query_positions: torch.Tensor, workspace_buffer: torch.Tensor, *, @@ -1609,6 +1704,8 @@ def _prepare_q_token_kv_block_sparse_attention( sparse_block_size: int = 4, split_kv: bool = True, share_pattern_across_kv_heads: bool = True, + count_current_block_in_topk: bool = True, + fixed_query_len: int = 0, ) -> _PrimsTSQTokenKvBlockSparsePlan: """Prepare one graph-stable metadata-plus-attention QToken-KvBlock-Sparse-Attention launch. @@ -1700,9 +1797,9 @@ def _prepare_q_token_kv_block_sparse_attention( if not isinstance(out, torch.Tensor): raise TypeError("out must be a caller-owned torch.Tensor") - use_packed_q = qo_indptr is not None + packed_query = qo_indptr is not None fixed_query_group_size = None - if use_packed_q: + if packed_query: if query.ndim != 3: raise ValueError( "packed QToken-KvBlock-Sparse-Attention query must have shape [total_q,Hq,D]" @@ -1747,6 +1844,8 @@ def _prepare_q_token_kv_block_sparse_attention( max_seq_len_kv, sparse_block_size, share_pattern_across_kv_heads, + count_current_block_in_topk, + fixed_query_len, ) if not share_pattern_across_kv_heads and block_indices.shape[1] != k_cache.shape[1]: raise ValueError("block_indices pattern head count must match the KV cache") @@ -1765,6 +1864,7 @@ def _prepare_q_token_kv_block_sparse_attention( sparse_block_size=sparse_block_size, split_kv=split_kv, share_pattern_across_kv_heads=share_pattern_across_kv_heads, + count_current_block_in_topk=count_current_block_in_topk, ) views = layout.bind(workspace_buffer) from .decode import _prepare_prims_ts_batch_decode_plan, _validate_scale @@ -1786,7 +1886,7 @@ def _prepare_q_token_kv_block_sparse_attention( out=attention_out, seq_len_q=group_size, qo_indptr=qo_indptr, - max_seq_len_q=group_size if use_packed_q else None, + max_q_len=group_size if packed_query else None, out_dtype=out.dtype, mask_type="causal", window_left=-1, @@ -1798,6 +1898,7 @@ def _prepare_q_token_kv_block_sparse_attention( use_pdl=True, split_kv=split_kv, share_pattern_across_kv_heads=share_pattern_across_kv_heads, + count_current_block_in_topk=count_current_block_in_topk, q_token_kv_block_sparse_page_memberships=( views.q_token_kv_block_sparse_page_memberships if group_size > 1 else None ), @@ -1807,6 +1908,7 @@ def _prepare_q_token_kv_block_sparse_attention( else () ), direct_q1_max_model_len=max_seq_len_kv if group_size == 1 else None, + direct_q1_fixed_query_len=fixed_query_len, allow_head_aware_locators=group_size > 1, ) # Preparation only binds shapes and storage; attention does not read these @@ -1826,6 +1928,8 @@ def _prepare_q_token_kv_block_sparse_attention( qo_indptr=qo_indptr, query_major_memberships=layout.query_major_memberships, locator_head_multiplier=attention_plan._flat_locator_heads, + count_current_block_in_topk=count_current_block_in_topk, + fixed_query_len=fixed_query_len, ) return _PrimsTSQTokenKvBlockSparsePlan( _metadata_plan=metadata_plan, @@ -1835,7 +1939,11 @@ def _prepare_q_token_kv_block_sparse_attention( _query=_describe_q_token_kv_block_sparse_tensor(query), _block_indices=_describe_q_token_kv_block_sparse_tensor(block_indices), _block_table=_describe_q_token_kv_block_sparse_tensor(block_table), - _token_to_request=_describe_q_token_kv_block_sparse_tensor(token_to_request), + _token_to_request=( + _describe_q_token_kv_block_sparse_tensor(token_to_request) + if token_to_request is not None + else None + ), _query_positions=_describe_q_token_kv_block_sparse_tensor(query_positions), _out=_describe_q_token_kv_block_sparse_tensor(out), _fixed_query_group_size=fixed_query_group_size, @@ -1848,8 +1956,8 @@ def q_token_kv_block_sparse_attention_with_paged_kv_cache( paged_kv_cache: tuple[torch.Tensor, torch.Tensor], block_table: torch.Tensor, indexer_block_ids: torch.Tensor, - token_to_request: torch.Tensor, - query_positions: torch.Tensor, + token_to_request: Optional[torch.Tensor], + query_positions: Optional[torch.Tensor], workspace_buffer: torch.Tensor, *, max_seq_len_kv: int, @@ -1859,10 +1967,12 @@ def q_token_kv_block_sparse_attention_with_paged_kv_cache( sm_scale: Optional[float] = None, v_scale: Optional[float] = None, out: Optional[torch.Tensor] = None, - o_data_type: Optional[torch.dtype] = None, + out_dtype: Optional[torch.dtype] = None, qo_indptr: Optional[torch.Tensor] = None, split_kv: bool = True, share_pattern_across_kv_heads: bool = True, + count_current_block_in_topk: bool = True, + seq_lens: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Plan and run one QToken-KvBlock-Sparse-Attention launch. @@ -1899,10 +2009,14 @@ def q_token_kv_block_sparse_attention_with_paged_kv_cache( dimension. Each valid row provides the completed-block prefix described by :meth:`QTokenKvBlockSparsePagedTSWrapper.run`; Q1 does not compact missing entries inside that prefix. Tail tokens are derived separately. - token_to_request : torch.Tensor - CUDA Int32 request ID for each flattened query token. - query_positions : torch.Tensor + token_to_request : torch.Tensor or None + CUDA Int32 request ID for each flattened query token. Pass None + with ``query_positions`` when using request-aligned fixed-Q lengths. + query_positions : torch.Tensor or None CUDA Int32 or Int64 causal position for each query token. + seq_lens : torch.Tensor, optional + Per-request fixed-Q KV lengths, as described by + :meth:`QTokenKvBlockSparsePagedTSWrapper.run`. workspace_buffer : torch.Tensor Caller-owned byte workspace sized by :func:`get_q_token_kv_block_sparse_workspace_size`. @@ -1925,7 +2039,7 @@ def q_token_kv_block_sparse_attention_with_paged_kv_cache( Caller-owned output with the same logical shape as ``q``. Must be disjoint from all inputs, including indexer/request metadata, and from ``workspace_buffer``. Storage overlap is not checked. - o_data_type : torch.dtype, optional + out_dtype : torch.dtype, optional Output dtype, defaulting to ``out.dtype`` or ``q.dtype``. qo_indptr : torch.Tensor, optional CUDA Int32 packed-route offsets. @@ -1938,6 +2052,10 @@ def q_token_kv_block_sparse_attention_with_paged_kv_cache( True (default) uses one pattern per query, shared by KV heads. False uses independent per-KV-head patterns. Match this setting in workspace sizing and planning; it does not change the attention grid. + count_current_block_in_topk : bool + True (default) counts a completed current block in the selection budget. + False excludes it from selected IDs and always appends its visible tokens. + Match this setting in workspace sizing and planning. Returns ------- @@ -1958,8 +2076,8 @@ def q_token_kv_block_sparse_attention_with_paged_kv_cache( ) if out is not None and not isinstance(out, torch.Tensor): raise TypeError("out must be a torch.Tensor") - if o_data_type is None: - o_data_type = out.dtype if out is not None else q.dtype + if out_dtype is None: + out_dtype = out.dtype if out is not None else q.dtype if qo_indptr is None: if q.ndim != 5: raise ValueError( @@ -1968,12 +2086,12 @@ def q_token_kv_block_sparse_attention_with_paged_kv_cache( if group_size != q.shape[2]: raise ValueError("fixed group_size must match q.shape[2]") batch_size = int(q.shape[0] * q.shape[1]) - use_packed_q = False + packed_query = False else: if q.ndim != 3: raise ValueError("packed q must have shape [total_q,Hq,D]") batch_size = int(qo_indptr.numel()) - 1 - use_packed_q = True + packed_query = True wrapper = QTokenKvBlockSparsePagedTSWrapper() wrapper.plan( @@ -1988,13 +2106,14 @@ def q_token_kv_block_sparse_attention_with_paged_kv_cache( device=q.device, workspace_buffer=workspace_buffer, group_size=group_size, - use_packed_q=use_packed_q, + packed_query=packed_query, mask_type=mask_type, split_kv=split_kv, share_pattern_across_kv_heads=share_pattern_across_kv_heads, - q_data_type=q.dtype, - kv_data_type=k_cache.dtype, - o_data_type=o_data_type, + count_current_block_in_topk=count_current_block_in_topk, + q_dtype=q.dtype, + kv_dtype=k_cache.dtype, + out_dtype=out_dtype, ) return wrapper.run( q, @@ -2003,6 +2122,7 @@ def q_token_kv_block_sparse_attention_with_paged_kv_cache( indexer_block_ids, token_to_request, query_positions, + seq_lens=seq_lens, qo_indptr=qo_indptr, sm_scale=sm_scale, v_scale=v_scale, @@ -2145,14 +2265,17 @@ def _align_up_q_token_kv_block_sparse_workspace(value: int) -> int: def _validate_inputs( block_indices: torch.Tensor, block_table: torch.Tensor, - token_to_request: torch.Tensor, + token_to_request: Optional[torch.Tensor], query_positions: torch.Tensor, group_size: int, storage_page_size: int, max_seq_len_kv: int, sparse_block_size: int, share_pattern_across_kv_heads: bool = True, + count_current_block_in_topk: bool = True, + fixed_query_len: int = 0, ) -> None: + """Validate tensor structure without copying live metadata values to the host.""" sparse_block_size = _validate_sparse_block_size(sparse_block_size) for name, tensor in ( ("block_indices", block_indices), @@ -2160,12 +2283,16 @@ def _validate_inputs( ("token_to_request", token_to_request), ("query_positions", query_positions), ): + if name == "token_to_request" and fixed_query_len and tensor is None: + continue if not isinstance(tensor, torch.Tensor): raise TypeError(f"{name} must be a torch.Tensor") if not block_indices.is_cuda: raise ValueError( "QToken-KvBlock-Sparse-Attention metadata inputs must be CUDA tensors" ) + if not isinstance(count_current_block_in_topk, bool): + raise TypeError("count_current_block_in_topk must be a bool") if not isinstance(share_pattern_across_kv_heads, bool): raise TypeError("share_pattern_across_kv_heads must be a bool") rank = 2 if share_pattern_across_kv_heads else 3 @@ -2189,14 +2316,31 @@ def _validate_inputs( max_seq_len_kv, block_table_token_capacity=block_table.shape[1] * storage_page_size, ) - if token_to_request.shape != (rows,) or token_to_request.dtype != torch.int32: - raise ValueError("token_to_request must be int32 with one value per row") - if query_positions.shape != (rows,) or query_positions.dtype not in ( + if fixed_query_len: + if ( + token_to_request is not None + or fixed_query_len < 0 + or fixed_query_len > 2**31 - 1 + or fixed_query_len % group_size + or rows != block_table.shape[0] * fixed_query_len + ): + raise ValueError("fixed Q must align complete query groups with requests") + metadata_rows = block_table.shape[0] + else: + assert token_to_request is not None + if token_to_request.shape != (rows,) or token_to_request.dtype != torch.int32: + raise ValueError("token_to_request must be int32 with one value per row") + metadata_rows = rows + if query_positions.shape != (metadata_rows,) or query_positions.dtype not in ( torch.int32, torch.int64, ): - raise ValueError("query_positions must be int32/int64 with one value per row") - tensors = (block_table, token_to_request, query_positions) + raise ValueError( + "positions/KV lengths must be int32/int64 with the expected extent" + ) + tensors: tuple[torch.Tensor, ...] = (block_table, query_positions) + if token_to_request is not None: + tensors += (token_to_request,) if any(tensor.device != block_indices.device for tensor in tensors): raise ValueError( "QToken-KvBlock-Sparse-Attention metadata inputs must share one CUDA device" @@ -2207,7 +2351,9 @@ def _validate_inputs( raise ValueError( "block_table must be a non-overlapping dense row-strided page table" ) - if token_to_request.stride(0) != 1 or query_positions.stride(0) != 1: + if ( + token_to_request is not None and token_to_request.stride(0) != 1 + ) or query_positions.stride(0) != 1: raise ValueError( "per-row QToken-KvBlock-Sparse-Attention metadata must be contiguous" ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/sage.py b/tensorrt_llm/_torch/attention/backends/prims_ts/sage.py deleted file mode 100644 index 8d224328a3d9..000000000000 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/sage.py +++ /dev/null @@ -1,480 +0,0 @@ -# Copyright (c) 2026 by FlashInfer team. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""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 -the trtllm-gen flat layout of TensorRT-LLM's ``sageQuant``: per head, sequence -``b`` starts at slot ``b * S // blk + b`` and token ``t`` uses slot -``t // blk`` inside it. -""" - -from __future__ import annotations - -from collections.abc import Mapping -from dataclasses import dataclass -from typing import TypeVar - -import torch - -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. -SAGE_ADAPTER_SLOTS = ("q_scale", "k_scale", "k_summary_scale", "v_scale", "v_mean") - -_T = TypeVar("_T") - - -@dataclass(frozen=True) -class SageAttentionConfig: - """Compile-time Sage attention recipe of one plan. - - ``q_block_size`` tokens share one Q scale and must be a power of two no - larger than the planned Q tile. ``k_block_size`` tokens share one K scale; - ``k_summary_block_size`` is the K block size of a proxy plan's summary - scales; construction sets it to ``k_block_size`` when it is omitted. - Both must be one of ``SAGE_K_BLOCK_SIZES``. ``v_mean`` says whether every - run supplies a per-channel V mean that is added back to the output. The - defaults are TensorRT-LLM's production recipe. - """ - - q_block_size: int = 1 - k_block_size: int = 16 - v_mean: bool = False - k_summary_block_size: int | None = None - - def __post_init__(self) -> None: - if self.k_summary_block_size is None: - object.__setattr__(self, "k_summary_block_size", self.k_block_size) - - -@dataclass(frozen=True) -class SageAttentionParams: - """Per-block Q/K scales and per-channel V scales of one run. - - All tensors are FP32 with the block sizes of the plan's - :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 ``[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 - them. Scales must be positive and finite; the kernel does not check them. - """ - - q_scale: torch.Tensor - k_scale: torch.Tensor - v_scale: torch.Tensor - k_summary_scale: torch.Tensor | None = None - v_mean: torch.Tensor | None = None - - -def flat_scale_numel(batch_size: int, seq_len: int, block_size: int) -> int: - """Return the per-head slot count of the flat scale layout.""" - return ceil_div(batch_size * seq_len, block_size) + batch_size - 1 - - -def is_power_of_two(value: int) -> bool: - """Return whether ``value`` is a positive power of two.""" - return value > 0 and value & (value - 1) == 0 - - -def log2_block_size(block_size: int) -> int: - """Return ``log2`` of a power-of-two scale block size.""" - if not is_power_of_two(block_size): - raise ValueError(f"block size must be a power of two, got {block_size}") - return block_size.bit_length() - 1 - - -def sage_adapter_slots(values: Mapping[str, _T | None]) -> tuple[_T | None, ...]: - """Arrange Sage values named by scale into the adapter slots. - - An absent name leaves its slot ``None``, which the adapter binds as a null - pointer. - """ - return tuple(values.get(name) for name in SAGE_ADAPTER_SLOTS) - - -def flat_scale_slot(batch_idx, token_idx, seq_len, log2_block: int): - """Return the per-head slot of one token in the flat scale layout. - - Sequence ``b`` starts at ``b * S // blk + b``; the extra ``b`` keeps its - last block apart from the next sequence's first block when ``blk`` does - not divide ``S``. The arithmetic serves host and device ``Int32`` values. - """ - return ((batch_idx * seq_len) >> log2_block) + batch_idx + (token_idx >> log2_block) - - -def _validate_scale_tensor( - tensor: torch.Tensor, - name: str, - *, - expected_shape: tuple[int, ...], - device: torch.device, -) -> None: - if not isinstance(tensor, torch.Tensor): - raise TypeError(f"{name} must be a torch.Tensor") - if tensor.dtype != torch.float32: - raise ValueError(f"{name} must have dtype float32, got {tensor.dtype}") - if tuple(tensor.shape) != expected_shape: - raise ValueError( - f"{name} must have shape {expected_shape}, got {tuple(tensor.shape)}" - ) - if not tensor.is_contiguous(): - raise ValueError(f"{name} must be contiguous") - if tensor.data_ptr() % 16 != 0: - raise ValueError(f"{name} data pointer must be 16-byte aligned") - if tensor.device != device: - raise ValueError(f"{name} must be on device {device}, got {tensor.device}") - - -def sage_scale_shapes( - config: SageAttentionConfig, - *, - batch_size: int, - seq_len_q: int, - seq_len_kv: int, - num_qo_heads: int, - num_kv_heads: int, - head_dim: int, - summary_seq_len: int | None = None, -) -> 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`` - otherwise. - """ - shapes = { - "q_scale": ( - num_qo_heads, - flat_scale_numel(batch_size, seq_len_q, config.q_block_size), - ), - "k_scale": ( - num_kv_heads, - flat_scale_numel(batch_size, seq_len_kv, config.k_block_size), - ), - "v_scale": (batch_size, num_kv_heads, head_dim), - } - if summary_seq_len is not None: - shapes["k_summary_scale"] = ( - num_kv_heads, - flat_scale_numel(batch_size, summary_seq_len, config.k_summary_block_size), - ) - if config.v_mean: - shapes["v_mean"] = (batch_size, num_kv_heads, head_dim) - return shapes - - -def validate_sage_params( - params: SageAttentionParams, - expected_shapes: Mapping[str, tuple[int, ...]], - *, - device: torch.device, -) -> None: - """Validate the scale tensors of one run against the plan's expected shapes. - - ``expected_shapes`` is the plan's ``sage_scale_shapes`` result; an - optional tensor is required exactly when its name is present. - """ - if not isinstance(params, SageAttentionParams): - raise TypeError("sage must be a SageAttentionParams instance") - for name, consumer in ( - ("v_mean", "a plan configured with v_mean=True"), - ("k_summary_scale", "block-sparse proxy routes"), - ): - if (getattr(params, name) is None) == (name in expected_shapes): - raise ValueError(f"{name} is required by {consumer} and rejected otherwise") - for name, expected_shape in expected_shapes.items(): - _validate_scale_tensor( - getattr(params, name), name, expected_shape=expected_shape, device=device - ) - - -_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", - "SageAttentionConfig", - "SageAttentionParams", - "flat_scale_numel", - "flat_scale_slot", - "is_power_of_two", - "log2_block_size", - "sage_adapter_slots", - "sage_scale_shapes", - "validate_sage_params", -] diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/sage/__init__.py b/tensorrt_llm/_torch/attention/backends/prims_ts/sage/__init__.py new file mode 100644 index 000000000000..7fe945dfbf93 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/sage/__init__.py @@ -0,0 +1,59 @@ +# 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. + +"""Sage attention shared by the PrimTS decode and context kernels. + +* ``config``: the public recipe (``SageAttentionConfig``) and per-run scale + tensors (``SageAttentionParams``), their flat layout, shapes and validation, + and the INT8 Q/K device gate. +* ``layout``: the word layout of the scale buffer a plan prepares before + attention (``_SageScaleLayout``). +* ``prepare``: the kernel that writes that buffer from the flat scales. +* ``index``: device-side index math and loads over the layout. +* ``modes``: the names of how a softmax obtains and applies the K scales. +* ``int32_scores``: INT8 Q/K accumulated as INT32 scores on a float bias. +* ``plan``: the host-side state of a plan's prepared buffer and its compiled + prepare launch, for plans that run the prepare as its own launch. + +Each attention kernel keeps its own resources, pipelines and softmax for +the scales; only these components are shared. This package imports the +host-side modules only; kernels import ``prepare``, ``index`` and +``int32_scores`` directly. +""" + +from .config import ( + SAGE_K_BLOCK_SIZES, + SageAttentionConfig, + SageAttentionParams, + is_power_of_two, + sage_flat_scale_numel, + sage_flat_scale_slot, + sage_log2_block_size, + sage_scale_shapes, + validate_sage_params, + validate_sage_recipe, +) + +__all__ = [ + "SAGE_K_BLOCK_SIZES", + "SageAttentionConfig", + "SageAttentionParams", + "is_power_of_two", + "sage_flat_scale_numel", + "sage_flat_scale_slot", + "sage_log2_block_size", + "sage_scale_shapes", + "validate_sage_params", + "validate_sage_recipe", +] diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/sage/config.py b/tensorrt_llm/_torch/attention/backends/prims_ts/sage/config.py new file mode 100644 index 000000000000..7b6b8afeef96 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/sage/config.py @@ -0,0 +1,267 @@ +# 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. + +"""Sage attention recipe and per-run scale tensors of the PrimTS attention plans. + +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 +the trtllm-gen flat layout of TensorRT-LLM's ``sageQuant``: per head, sequence +``b`` starts at slot ``b * S // blk + b`` and token ``t`` uses slot +``t // blk`` inside it. This module holds the recipe, the scale tensors, their +flat layout and validation, and the INT8 Q/K device gate; the plans of both +the decode and the context kernels use it. +""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass + +import torch + +from flashinfer.utils import ceil_div + +# Only SM100 has the INT8 tcgen05 MMA kind; later GPUs use E4M3 Q/K. +_INT8_QK_COMPUTE_CAPABILITIES = ((10, 0),) + + +def _validate_int8_qk_device(q_dtype: torch.dtype, device_index: int) -> None: + """Reject Int8 Q/K on GPUs without the INT8 tcgen05 MMA.""" + + if q_dtype != torch.int8: + return + capability = torch.cuda.get_device_capability(device_index) + if capability not in _INT8_QK_COMPUTE_CAPABILITIES: + raise NotImplementedError( + "Int8 Q/K Sage attention requires an SM100a/B200 GPU, whose tcgen05 " + "MMA supports INT8 operands; use the torch.float8_e4m3fn Q/K recipe " + f"on device cuda:{device_index} with compute capability {capability}" + ) + + +SAGE_K_BLOCK_SIZES = (1, 4, 16, 32, 64, 128, 256) + + +@dataclass(frozen=True) +class SageAttentionConfig: + """Compile-time Sage attention recipe of one plan. + + ``q_block_size`` tokens share one Q scale and must be a power of two; a + block may span several Q tiles, since every Q row reads its own prepared + scale. ``k_block_size`` tokens share one K scale; + ``k_summary_block_size`` is the K block size of a proxy plan's summary + scales; construction sets it to ``k_block_size`` when it is omitted. + Both must be one of ``SAGE_K_BLOCK_SIZES``. ``v_mean`` says whether every + run supplies a per-channel V mean that is added back to the output. The + defaults are TensorRT-LLM's production recipe. + """ + + q_block_size: int = 1 + k_block_size: int = 16 + v_mean: bool = False + k_summary_block_size: int | None = None + + def __post_init__(self) -> None: + if self.k_summary_block_size is None: + object.__setattr__(self, "k_summary_block_size", self.k_block_size) + + +@dataclass(frozen=True) +class SageAttentionParams: + """Per-block Q/K scales and per-channel V scales of one run. + + All tensors are FP32 with the block sizes of the plan's + :class:`SageAttentionConfig`: ``q_scale`` is + ``[Hq, sage_flat_scale_numel(B, Sq, q_block_size)]``, ``k_scale`` is + ``[Hkv, sage_flat_scale_numel(B, Skv, k_block_size)]``, and ``v_scale`` and + ``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, sage_flat_scale_numel(B, num_kv_blocks, k_summary_block_size)]``. + ``v_mean`` and ``k_summary_scale`` are present exactly when the plan uses + them. Scales must be positive and finite; the kernel does not check them. + """ + + q_scale: torch.Tensor + k_scale: torch.Tensor + v_scale: torch.Tensor + k_summary_scale: torch.Tensor | None = None + v_mean: torch.Tensor | None = None + + +def sage_flat_scale_numel(batch_size: int, seq_len: int, block_size: int) -> int: + """Return the per-head slot count of the flat scale layout.""" + return ceil_div(batch_size * seq_len, block_size) + batch_size - 1 + + +def is_power_of_two(value: int) -> bool: + """Return whether ``value`` is a positive power of two.""" + return value > 0 and value & (value - 1) == 0 + + +def sage_log2_block_size(block_size: int) -> int: + """Return ``log2`` of a power-of-two scale block size.""" + if not is_power_of_two(block_size): + raise ValueError(f"block size must be a power of two, got {block_size}") + return block_size.bit_length() - 1 + + +def sage_flat_scale_slot(batch_idx, token_idx, seq_len, log2_block: int): + """Return the per-head slot of one token in the flat scale layout. + + Sequence ``b`` starts at ``b * S // blk + b``; the extra ``b`` keeps its + last block apart from the next sequence's first block when ``blk`` does + not divide ``S``. The arithmetic serves host and device ``Int32`` values. + """ + return ((batch_idx * seq_len) >> log2_block) + batch_idx + (token_idx >> log2_block) + + +def _validate_scale_tensor( + tensor: torch.Tensor, + name: str, + *, + expected_shape: tuple[int, ...], + device: torch.device, +) -> None: + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tensor.dtype != torch.float32: + raise ValueError(f"{name} must have dtype float32, got {tensor.dtype}") + if tuple(tensor.shape) != expected_shape: + raise ValueError( + f"{name} must have shape {expected_shape}, got {tuple(tensor.shape)}" + ) + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + if tensor.data_ptr() % 16 != 0: + raise ValueError(f"{name} data pointer must be 16-byte aligned") + if tensor.device != device: + raise ValueError(f"{name} must be on device {device}, got {tensor.device}") + + +# Q/K and output dtypes of the Sage recipe. +_SAGE_QK_DTYPES = (torch.float8_e4m3fn, torch.int8) +_SAGE_OUTPUT_DTYPES = (torch.float16, torch.bfloat16) + + +def validate_sage_recipe( + config: object, + *, + q_dtype: torch.dtype, + k_dtype: torch.dtype, + v_dtype: torch.dtype, + output_dtype: torch.dtype, +) -> None: + """Reject a Sage recipe at plan time, the same way for every wrapper. + + The recipe quantizes Q and K to one 8-bit dtype, float8_e4m3fn or int8, + and V to float8_e4m3fn, and returns a float16 or bfloat16 output: an + 8-bit output would discard the precision the scales keep. The block + sizes follow ``SageAttentionConfig``. Int8 Q/K also need an SM100 GPU + (``_validate_int8_qk_device``), and each wrapper rejects the plan + geometries its kernel does not run. + """ + if not isinstance(config, SageAttentionConfig): + raise TypeError("sage must be a SageAttentionConfig") + if q_dtype != k_dtype: + raise ValueError( + "Sage attention requires Q and K in one dtype; " + f"got q {q_dtype} and k {k_dtype}" + ) + if q_dtype not in _SAGE_QK_DTYPES: + raise NotImplementedError( + f"Sage attention requires float8_e4m3fn or int8 Q/K; got {q_dtype}" + ) + if v_dtype != torch.float8_e4m3fn: + raise NotImplementedError( + f"Sage attention requires float8_e4m3fn V; got {v_dtype}" + ) + if output_dtype not in _SAGE_OUTPUT_DTYPES: + raise NotImplementedError( + f"Sage attention requires a float16 or bfloat16 output; got {output_dtype}" + ) + for name in ("k_block_size", "k_summary_block_size"): + block_size = getattr(config, name) + if block_size not in SAGE_K_BLOCK_SIZES: + raise ValueError( + f"sage.{name} must be one of {SAGE_K_BLOCK_SIZES}, got {block_size}" + ) + if not is_power_of_two(config.q_block_size): + raise ValueError( + f"sage.q_block_size must be a power of two, got {config.q_block_size}" + ) + + +def sage_scale_shapes( + config: SageAttentionConfig, + *, + batch_size: int, + seq_len_q: int, + seq_len_kv: int, + num_qo_heads: int, + num_kv_heads: int, + head_dim: int, + summary_seq_len: int | None = None, +) -> 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`` + otherwise. + """ + shapes = { + "q_scale": ( + num_qo_heads, + sage_flat_scale_numel(batch_size, seq_len_q, config.q_block_size), + ), + "k_scale": ( + num_kv_heads, + sage_flat_scale_numel(batch_size, seq_len_kv, config.k_block_size), + ), + "v_scale": (batch_size, num_kv_heads, head_dim), + } + if summary_seq_len is not None: + shapes["k_summary_scale"] = ( + num_kv_heads, + sage_flat_scale_numel( + batch_size, summary_seq_len, config.k_summary_block_size + ), + ) + if config.v_mean: + shapes["v_mean"] = (batch_size, num_kv_heads, head_dim) + return shapes + + +def validate_sage_params( + params: SageAttentionParams, + expected_shapes: Mapping[str, tuple[int, ...]], + *, + device: torch.device, +) -> None: + """Validate the scale tensors of one run against the plan's expected shapes. + + ``expected_shapes`` is the plan's ``sage_scale_shapes`` result; an + optional tensor is required exactly when its name is present. + """ + if not isinstance(params, SageAttentionParams): + raise TypeError("sage must be a SageAttentionParams instance") + for name, consumer in ( + ("v_mean", "a plan configured with v_mean=True"), + ("k_summary_scale", "block-sparse proxy routes"), + ): + if (getattr(params, name) is None) == (name in expected_shapes): + raise ValueError(f"{name} is required by {consumer} and rejected otherwise") + for name, expected_shape in expected_shapes.items(): + _validate_scale_tensor( + getattr(params, name), name, expected_shape=expected_shape, device=device + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/sage/index.py b/tensorrt_llm/_torch/attention/backends/prims_ts/sage/index.py new file mode 100644 index 000000000000..7d3dfffa2c2f --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/sage/index.py @@ -0,0 +1,134 @@ +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Device-side index math and loads over the Sage scale buffers. + +The helpers turn coordinates into words of the flat scale tensors +(``config.sage_flat_scale_slot``) and of the prepared buffer +(``layout._SageScaleLayout``): a sequence's first word, a section's first +word and a Q tile row's word. They also load single FP32 scales. They hold +no resource or pipeline state; each attention kernel moves the words it +reads with its own resources. +""" + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Int64 + +from .config import sage_flat_scale_slot +from .layout import _SageKSection, _SageQSection, _SageScaleLayout, _SageVSection + + +@cute.jit +def sage_sequence_words( + layout: cutlass.Constexpr[_SageScaleLayout], + *, + batch_idx: Int32, + kv_head_idx: Int32, + num_kv_heads: Int32, +) -> Int32: + """Return the first word of sequence ``b * Hkv + h`` in the prepared scales.""" + sequence_idx = batch_idx * num_kv_heads + kv_head_idx + return sequence_idx * Int32(layout.sequence_words) + + +@cute.jit +def sage_section_word( + sequence_words: Int32, + section: cutlass.Constexpr[_SageKSection | _SageQSection | _SageVSection], +) -> Int32: + """Return the first word of one section of a sequence's prepared scales.""" + return sequence_words + Int32(section.word_offset) + + +@cute.jit +def sage_q_word( + layout: cutlass.Constexpr[_SageScaleLayout], + *, + sequence_words: Int32, + q_tile_idx: Int32, + row_idx: Int32, +) -> Int32: + """Return the ``sfQ`` word of row ``row_idx`` of Q tile ``q_tile_idx``. + + ``q_tile_idx`` numbers the tiles as ``_SageQSection`` lays them out: the + decode kernel's grouped Q CTA, or the context kernel's + ``tile_index(local_head, token_tile)``. + """ + q_rows = layout.q_rows + return ( + sage_section_word(sequence_words, q_rows) + + q_tile_idx * Int32(q_rows.rows_per_tile) + + row_idx + ) + + +@cute.jit +def load_sage_q_scale( + layout: cutlass.Constexpr[_SageScaleLayout], + scales: cute.Pointer, + *, + sequence_words: Int32, + q_tile_idx: Int32, + row_idx: Int32, +) -> Float32: + """Return ``sfQ`` of row ``row_idx`` of Q tile ``q_tile_idx``. + + ``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 token's for a row past the sequence's end + (``_SageQSection``). A softmax lane reads it once per work tile. + """ + word = sage_q_word( + layout, sequence_words=sequence_words, q_tile_idx=q_tile_idx, row_idx=row_idx + ) + return Float32(scales[word]) + + +@cute.jit +def _load_sage_scale(scale_addr: Int64, index: Int32) -> Float32: + """Return one FP32 scale loaded from a global memory base address. + + One 32-bit load per call; nothing widens it. The callers read single + words: the prepare gathers each prepared word from its own flat source + word, and the decode kernel stages a sequence's ``sfV`` row one channel + per lane, so a warp's loads coalesce. ``sfK`` atoms move in 16-byte units + instead. + """ + value_ptr = cutlass.inttoptr( + scale_addr + Int64(index) * Int64(4), + mem_space=1, + dtype=Float32, + ) + return Float32(value_ptr.load(count=1, alignment=4)[0]) + + +@cute.jit +def load_sage_flat_scale( + scale_addr: Int64, + head_stride: Int32, + *, + head_idx: Int32, + batch_idx: Int32, + seq_len: Int32, + token_idx: Int32, + log2_block: Int32, +) -> Float32: + """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 = sage_flat_scale_slot(batch_idx, token_idx, seq_len, log2_block) + return _load_sage_scale(scale_addr, head_idx * head_stride + slot) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/sage/int32_scores.py b/tensorrt_llm/_torch/attention/backends/prims_ts/sage/int32_scores.py new file mode 100644 index 000000000000..37f234fa3a7c --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/sage/int32_scores.py @@ -0,0 +1,114 @@ +"""INT32 scores on a float bias: 8-bit integer QK MMA that the softmax reads as FP32. + +INT8 Q/K accumulate onto the bit pattern of ``1.5 * 2**23``. Integers in +``[2**23, 2**24)`` are FP32 numbers with unit spacing and every D = 128 INT8 +dot product has ``|score| <= 2**21``, so the INT32 accumulator is the exact +FP32 number ``12582912.0 + score`` and the softmax reads it without a +conversion, removing the bias where it applies its scales. + +One ``kind::f16`` MMA step writes the bias before the INT8 K steps: both +operands read one BF16 tile whose K = 16 rows repeat ``[1024, 1024, 1024, +0]``, so the twelve products ``2**20`` sum exactly to ``1.5 * 2**23``. The +tile is unswizzled K-major 32-byte rows with 128-byte LBO and 256-byte SBO +core-matrix strides. The decode and context kernels share the tile and the +seeding step through the helpers below. +""" + +from __future__ import annotations + +import cutlass +import cutlass.cute as cute +from cutlass import BFloat16, Float32, Int32, Uint32 +from cutlass.experimental import primitives as prims + +from ..kernels.tcgen05_compat import tcgen05_mma_ws + +INT32_SCORE_BIAS = 12582912.0 +INT32_SCORE_SEED_MMA_K = 16 +INT32_SCORE_SEED_TILE_WORDS = (0x44804480, 0x00004480) +INT32_SCORE_SEED_TILE_LBO = 128 +INT32_SCORE_SEED_TILE_SBO = 256 + + +def sage_score_seed_tile_bytes(mma_m: int, mma_n: int) -> int: + """Return the byte size of the seed MMA's operand tile. + + A reads the first M rows and B the first N rows of the same tile. + """ + return max(mma_m, mma_n) * INT32_SCORE_SEED_MMA_K * 2 + + +@cute.jit +def sage_score_seed_words(smem_base, offset: int, num_bytes: int) -> cutlass.Array: + """Return the seed MMA's operand tile as a 32-bit SMEM array.""" + return cutlass.Array( + smem_base.data_ptr() + offset, + dtype=Uint32, + shape=(num_bytes // 4,), + addrspace=3, + ) + + +@cute.jit +def fill_sage_score_seed_tile( + words: cutlass.Array, num_words: int, thread_idx: Int32, num_threads: int +) -> None: + """Fill the seed tile with ``num_threads`` threads, one word pair each per round. + + The fence makes the generic-proxy stores visible to the tensor core. + """ + even_word, odd_word = INT32_SCORE_SEED_TILE_WORDS + num_pairs = num_words // 2 + for round_idx in cutlass.range_constexpr( + (num_pairs + num_threads - 1) // num_threads + ): + pair_idx = thread_idx + Int32(round_idx * num_threads) + if pair_idx < Int32(num_pairs): + words[pair_idx * 2] = Uint32(even_word) + words[pair_idx * 2 + 1] = Uint32(odd_word) + cute.arch.fence_view_async_shared() + + +@cute.jit +def issue_sage_score_seed_mma( + tmem_s, + words: cutlass.Array, + mma_m: int, + mma_n: int, + *, + ws: cutlass.Constexpr[bool] = False, +) -> None: + """Write ``INT32_SCORE_BIAS`` into an S stage with one BF16 MMA step. + + The elected thread issues the step; it reads only the constant tile and no + accumulator, so the INT8 K steps issued afterwards accumulate onto the + bias bit pattern in order. + """ + if prims.elect_sync(): + tile_desc = prims.Tcgen05SmemDesc.build( + words, + leading_byte_offset=INT32_SCORE_SEED_TILE_LBO, + stride_byte_offset=INT32_SCORE_SEED_TILE_SBO, + layout=prims.Tcgen05SmemSwizzle.NONE, + ) + idesc = prims.Tcgen05InstrDesc.build( + c_dtype=Float32, + a_dtype=BFloat16, + b_dtype=BFloat16, + n_dim=mma_n, + m_dim=mma_m, + ) + if cutlass.const_expr(ws): + tcgen05_mma_ws( + prims.Tcgen05MMAKind.F16, tmem_s, tile_desc, tile_desc, idesc, False + ) + else: + prims.tcgen05_mma( + prims.Tcgen05MMAKind.F16, + prims.CTAGroup.CTA_1, + tmem_s, + tile_desc, + tile_desc, + idesc, + False, + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/sage/layout.py b/tensorrt_llm/_torch/attention/backends/prims_ts/sage/layout.py new file mode 100644 index 000000000000..672d100e49b6 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/sage/layout.py @@ -0,0 +1,288 @@ +# 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. + +"""Word layout of the FP32 scale buffer a Sage plan prepares before attention. + +The Sage prepare (``prepare.py``) gathers a plan's flat scale tensors into one +buffer in the order the attention kernels read them; ``_SageScaleLayout`` +describes that buffer and its per-sequence sections. Both the decode and the +context kernels read it; each kernel fixes the parameters of its own reads +(the Q tile rows, the K tile atoms). +""" + +from __future__ import annotations + +from dataclasses import dataclass, replace + +from flashinfer.utils import ceil_div, round_up + +_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. The atoms are padded to whole + groups of ``tile_atoms``, so a reader that addresses whole K tiles never + runs past the section; a pad atom repeats the last scale. + """ + + length: int + block_size: int + source_block_size: int + tile_atoms: int = 1 + # First word in the sequence's words; ``_SageScaleLayout`` places it. + word_offset: int = 0 + + @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, padded to whole ``tile_atoms`` groups.""" + + return round_up(ceil_div(self.length, _SAGE_ATOM_TOKENS), self.tile_atoms) + + @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 row of each tile. + + A Q tile is what one softmax instance reads: ``rows_per_tile`` rows, its + MMA M extent, row ``r`` of tile ``t`` at word ``word_offset + t * + rows_per_tile + r``, in the order the softmax lanes read them. That is + the decode kernel's ``tile_size_q``, whose two K/V instances read the + same rows, and half of the query-paired context kernel's 256-row CTA + tile (``tile_size_q``), whose two softmax instances each read one. With + ``heads_interleaved`` (the decode kernel's grouped tile) a tile holds + ``tokens_per_tile`` tokens of each of the ``heads_q_per_kv`` Q heads of + its KV head, row ``r`` being token ``r // heads_q_per_kv`` of local head + ``r % heads_q_per_kv``; when those rows do not fill the tile, its + remaining padding rows decode to the tokens after it. Otherwise (the + context kernel's tile) a tile holds ``rows_per_tile`` tokens of one Q + head: tile ``h * token_tiles + i`` is token range ``i`` of local head + ``h`` and row ``r`` its token ``r``. The word is that head's scale of the + token's ``q_block_size`` block; a row whose token lies past the + sequence's end takes the last token's, so a masked or padding row keeps + a finite positive scale. + """ + + seq_len_q: int + q_block_size: int + heads_q_per_kv: int + rows_per_tile: int + heads_interleaved: bool + # First word in the sequence's words; ``_SageScaleLayout`` places it. + word_offset: int = 0 + + def __post_init__(self) -> None: + if self.tokens_per_tile < 1 or self.rows_per_tile % _FLOAT4_WORDS: + raise ValueError("sfQ tiles hold whole float4s of at least one token") + + @property + def tokens_per_tile(self) -> int: + """Return the tokens of each Q head that one tile holds.""" + + if self.heads_interleaved: + return self.rows_per_tile // self.heads_q_per_kv + return self.rows_per_tile + + @property + def token_tiles(self) -> int: + """Return the ``tokens_per_tile`` token ranges of the sequence, the last possibly partial.""" + + return ceil_div(self.seq_len_q, self.tokens_per_tile) + + @property + def tiles(self) -> int: + """Return the Q tiles of one sequence.""" + + heads = 1 if self.heads_interleaved else self.heads_q_per_kv + return heads * self.token_tiles + + @property + def words(self) -> int: + """Return the words of the section: ``rows_per_tile`` per tile.""" + + return self.tiles * self.rows_per_tile + + def tile_index(self, local_head, token_tile): + """Return the tile holding token range ``token_tile`` of local Q head ``local_head``. + + The arithmetic serves host integers and device ``Int32`` values. + """ + + if self.heads_interleaved: + return token_tile + return local_head * self.token_tiles + token_tile + + def row_token_and_head(self, tile, row): + """Return the token and the local Q head of row ``row`` of tile ``tile``. + + The token is not clamped to the sequence. The arithmetic serves host + integers, tensors and device ``Int32`` values. + """ + + if self.heads_interleaved: + token_in_tile = row // self.heads_q_per_kv + local_head = row - token_in_tile * self.heads_q_per_kv + return tile * self.tokens_per_tile + token_in_tile, local_head + local_head = tile // self.token_tiles + token_tile = tile - local_head * self.token_tiles + return token_tile * self.tokens_per_tile + row, local_head + + +@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 + # First word in the sequence's words; ``_SageScaleLayout`` places it. + word_offset: int = 0 + + 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``. + """ + + k_tokens: _SageKSection + q_rows: _SageQSection + v_rows: _SageVSection + # A proxy plan's summary atoms, else ``None``. + summaries: _SageKSection | None = None + + def __post_init__(self) -> None: + # Place the sections in buffer order; each starts where the last ends. + word_offset = 0 + for name in ("k_tokens", "summaries", "q_rows", "v_rows"): + section = getattr(self, name) + if section is None: + continue + object.__setattr__(self, name, replace(section, word_offset=word_offset)) + word_offset += section.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 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" + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/sage/modes.py b/tensorrt_llm/_torch/attention/backends/prims_ts/sage/modes.py new file mode 100644 index 000000000000..76df02de8fa8 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/sage/modes.py @@ -0,0 +1,52 @@ +# 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. + +"""Names for how a Sage softmax obtains and applies the ``sfK`` scales. + +Each attention kernel picks one transport and one scaling per plan and keeps +its own selection, resources and code (the context kernel's +``sage_k_transport`` and ``sage_score_scaling`` in +``kernels/fmha_context/sage_scales.py``, the decode kernel's +``FmhaDecodeConfig.sage_score_scaling``); only the names are shared. +""" + +import enum + + +class SageKTransport(enum.Enum): + """Where the softmax reads a K tile's ``sfK`` words from. + + The choice follows the number of words one K tile carries. The context + kernel reads them from global memory when a 128-token tile has at most 32 + words (K blocks of four tokens or more) and stages the 128 words of + one-token blocks in shared memory (``sage_k_transport``). The decode + kernel stages every tile's words in shared memory, where its load warp + also gathers the atoms of block-sparse routes (``SageKScalesResource``). + """ + + # The softmax lanes load the words from the prepared buffer in global + # memory before the score wait. + GMEM = "gmem" + # The load warp stages the words in a shared-memory ring. + SMEM = "smem" + + +class SageScoreScaling(enum.Enum): + """How the softmax applies ``sfK`` to a K tile's scores.""" + + # The scores stay raw; each block's maximum and exp2 multiplier carry + # ``sfQ * sfK``. + BLOCK = "block" + # The max pass multiplies the scores by ``sfK`` in place. + DEQUANTIZED = "dequantized" diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/sage/plan.py b/tensorrt_llm/_torch/attention/backends/prims_ts/sage/plan.py new file mode 100644 index 000000000000..a910226b2566 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/sage/plan.py @@ -0,0 +1,140 @@ +# 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. + +"""Host-side state of a plan's prepared Sage scale buffer. + +A plan that runs the Sage prepare as its own launch before attention keeps +the buffer it writes (``_SageScaleLayout``), the scale shapes a run must +supply and the compiled prepare launch; the contiguous context plan does. +The decode kernel compiles its prepare into the attention launch instead. +""" + +import functools +from dataclasses import dataclass +from typing import Callable + +import torch + +from .config import SageAttentionConfig, sage_scale_shapes +from .layout import _SageScaleLayout + +_COMPILE_OPTIONS = "--enable-tvm-ffi --opt-level 3" + + +@dataclass(frozen=True) +class SagePreparedScales: + """A plan's prepared Sage scale buffer and the launch that writes it.""" + + layout: _SageScaleLayout + # Expected shape of every scale tensor a run supplies, by field name. + scale_shapes: dict[str, tuple[int, ...]] + # The buffer the Sage prepare writes before attention (``_SageScaleLayout``). + scales: torch.Tensor + prepare: Callable[..., None] + + +@functools.cache +def get_compiled_sage_prepare( + device_index: int, + layout: _SageScaleLayout, + num_qo_heads: int, + num_kv_heads: int, +) -> Callable[..., None]: + """Compile the launch that writes a plan's prepared Sage scales. + + The compiled callable takes ``(q_scale, k_scale, v_scale, v_mean, scales)`` + torch tensors (``v_mean`` is ``None`` for a recipe without one) and runs on + the current stream. The batch size is a runtime extent, as for the + attention kernel, so plans of one layout share the compilation. + """ + + import cutlass + import cutlass.cute as cute + from cuda.bindings import driver as cuda_drv + + from .prepare import _PrepareSageScales + + prepare = _PrepareSageScales(layout=layout, num_kv_heads=num_kv_heads) + + @cute.jit + def prepare_adapter( + q_scale: cute.Tensor, + k_scale: cute.Tensor, + v_scale: cute.Tensor, + v_mean: cute.Tensor | None, + scales: cute.Tensor, + stream: cuda_drv.CUstream, + ) -> None: + prepare(q_scale, k_scale, None, v_scale, v_mean, scales, stream) + + def fake_compact(shape): + return cute.runtime.make_fake_compact_tensor( + cutlass.Float32, + shape, + stride_order=tuple(reversed(range(len(shape)))), + assumed_align=16, + ) + + v_rows_shape = (cute.sym_int(), num_kv_heads, layout.v_rows.headdim) + with torch.cuda.device(device_index): + return cute.compile( + prepare_adapter, + fake_compact((num_qo_heads, cute.sym_int())), + fake_compact((num_kv_heads, cute.sym_int())), + fake_compact(v_rows_shape), + fake_compact(v_rows_shape) if layout.v_rows.has_mean else None, + fake_compact((cute.sym_int(),)), + cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True), + options=_COMPILE_OPTIONS, + ) + + +def build_sage_prepared_scales( + layout: _SageScaleLayout, + sage: SageAttentionConfig, + *, + batch_size: int, + seq_len_q: int, + seq_len_kv: int, + num_qo_heads: int, + num_kv_heads: int, + head_dim: int, + device: torch.device, + device_index: int, +) -> SagePreparedScales: + """Allocate a plan's prepared scale buffer and compile its prepare launch.""" + + scale_shapes = sage_scale_shapes( + sage, + batch_size=batch_size, + seq_len_q=seq_len_q, + seq_len_kv=seq_len_kv, + num_qo_heads=num_qo_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + ) + scales = torch.empty( + layout.total_words(batch_size * num_kv_heads), + dtype=torch.float32, + device=device, + ) + prepare = get_compiled_sage_prepare( + device_index, layout, num_qo_heads, num_kv_heads + ) + return SagePreparedScales( + layout=layout, + scale_shapes=scale_shapes, + scales=scales, + prepare=prepare, + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/sage_prepare.py b/tensorrt_llm/_torch/attention/backends/prims_ts/sage/prepare.py similarity index 86% rename from tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/sage_prepare.py rename to tensorrt_llm/_torch/attention/backends/prims_ts/sage/prepare.py index 88afcfccd3cc..13654c18198a 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/sage_prepare.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/sage/prepare.py @@ -32,19 +32,14 @@ from flashinfer.utils import ceil_div -from ...sage import ( +from .config import sage_log2_block_size +from .index import load_sage_flat_scale +from .layout 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 @@ -78,7 +73,7 @@ def create(section: _SageKSection, scales: cute.Tensor) -> "_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)), + log2_block=cutlass.Int32(sage_log2_block_size(section.source_block_size)), scale_addr=scales.iterator.toint(), head_stride=cutlass.Int32(scales.shape[1]), length=cutlass.Int32(section.length), @@ -188,7 +183,7 @@ def _write_k_float4( 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( + words[elem] = load_sage_flat_scale( section.scale_addr, section.head_stride, head_idx=head_idx, @@ -206,7 +201,6 @@ 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, @@ -216,42 +210,33 @@ def _write_q_float4( """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. + ``r`` of the tile takes the Q scale of the token and Q head the section's + row order maps it to (``_SageQSection.row_token_and_head``); + ``load_sage_flat_scale`` gives a token past the sequence's end the last token's. """ q_rows = layout.q_rows - tile_float4s = cutlass.Int32(q_rows.tile_size_q // _FLOAT4_WORDS) + tile_float4s = cutlass.Int32(q_rows.rows_per_tile // _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 + q_token_idx, local_head_idx = q_rows.row_token_and_head( + tile_idx, first_row + cutlass.Int32(elem) ) - words[elem] = load_flat_scale( + words[elem] = load_sage_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, + seq_len=cutlass.Int32(q_rows.seq_len_q), token_idx=q_token_idx, - log2_block=cutlass.Int32(log2_block_size(q_rows.q_block_size)), + log2_block=cutlass.Int32(sage_log2_block_size(q_rows.q_block_size)), ) _store_float4(scales, float4_idx, words) @@ -259,6 +244,9 @@ def _write_q_float4( class _PrepareSageScales: """Write a Sage plan's prepared scales from its flat scale tensors. + Serves the decode and the context kernels: the layout fixes the Q row + order, the K sections and the V rows. + 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 @@ -267,20 +255,10 @@ class _PrepareSageScales: route prepare in stream order. """ - def __init__( - self, - *, - cfg: FmhaDecodeConfig, - layout: _SageScaleLayout, - batch_size: int, - num_kv_heads: int, - ) -> None: - self.cfg = cfg + def __init__(self, *, layout: _SageScaleLayout, num_kv_heads: int) -> None: 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__( @@ -293,8 +271,10 @@ def __call__( scales: cute.Tensor, stream: cuda_drv.CUstream, ) -> None: + # The buffer's float4 count follows the run's batch size. + num_float4s = scales.shape[0] // _FLOAT4_WORDS 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], + grid=[ceil_div(num_float4s, _THREADS_PER_CTA), 1, 1], block=[_THREADS_PER_CTA, 1, 1], stream=stream, ) @@ -314,7 +294,7 @@ def kernel( 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): + if float4_idx < cutlass.Int32(scales.shape[0] // _FLOAT4_WORDS): sequence_idx = float4_idx // cutlass.Int32(self.sequence_float4s) float4_in_sequence = float4_idx - sequence_idx * cutlass.Int32( self.sequence_float4s @@ -353,7 +333,6 @@ def kernel( q_scale, scales, float4_idx, - self.cfg, self.layout, batch_idx=batch_idx, kv_head_idx=head_idx, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/split_kv_mode_policy.py b/tensorrt_llm/_torch/attention/backends/prims_ts/split_kv_mode_policy.py index 083cc9b68ef7..df88a2092010 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/split_kv_mode_policy.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/split_kv_mode_policy.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Shared split-KV reduction-mode policy for FMHA and MLA decode.""" +"""Shared split-KV reduction-mode policy for FMHA and MLA.""" from collections.abc import Iterable diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/vc_attention/__init__.py b/tensorrt_llm/_torch/attention/backends/prims_ts/vc_attention/__init__.py new file mode 100644 index 000000000000..20257ea3618f --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/vc_attention/__init__.py @@ -0,0 +1,311 @@ +# 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. + +"""VC-Attention operand contract of the prims_ts context kernel. + +VC-Attention (arXiv 2609.15810) adds two changes on the value side of +attention: a token permutation that groups similar value rows into the same +128-token K/V tile, and per-tile value smoothing. Each V tile is stored as E4M3 +residuals around its tile mean; the kernel restores the mean inside the online +softmax as ``O += rowsum(P_tile) * mean_tile`` (bf16 K=16 UMMA steps), so the +residual quantization error no longer carries the value DC component. P is +E4M3 through the ExpCast path in both recipes of the Q/K side: + +* VC-Attention-QK16 keeps Q and K in bf16 (plan ``q_dtype`` bf16); +* VC-Attention-QK8 (plan ``q_dtype`` E4M3) centres K by its per-head channel + mean, rotates Q and K by the normalised 128-point Hadamard matrix and + quantizes them to E4M3 with one scale per 128-token Q block and per + 128-token K tile, as the paper does. + +This module holds what the kernel and the plan need: the tile constants, the +mean operand shapes, ``VCAttentionConfig`` and ``VCAttentionParams`` with +their validation, the repair tile arithmetic and ``vc_reference``, the fp32 +attention over the dequantized operands. The quantizer that produces the +operands lives in ``quantization`` beside this module and its CuTe DSL +kernels in ``prepare``. + +The kernel takes the means and the per-channel V scale through +``BatchPrefillTSWrapper.run`` (``vc=VCAttentionParams(...)``), which folds the +scale into ``output_scale``. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from .quantization import VCAttentionOperands + + +# The head dimension the preparation kernels and the attention plans support. +VC_HEAD_DIM = 128 +# The K/V tile whose value mean the kernel restores, its only supported size. +VC_K_BLOCK_SIZE = 128 +# Q tokens per E4M3 Q scale of VC-Attention-QK8, powers of two up to the Q tile. +# K of the bf16 UMMA mean step, one kind::f16 K step. +VC_MEAN_MMA_K = 16 +# Mean UMMA steps issued back to back per group of tiles; each tile takes two of a +# step's 16 K slots (bf16 hi/lo row sum). +VC_MEAN_OPERANDS = 2 +VC_MEAN_GROUP_TILES = VC_MEAN_OPERANDS * VC_MEAN_MMA_K // 2 +E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max + + +def vc_mean_group_shape(head_dim: int) -> tuple[int, int]: + """Shape of the packed mean operands of one tile group: the ``VC_MEAN_OPERANDS`` + operands of ``vc_mean_operand_shape`` stacked along the row axis.""" + rows, cols = vc_mean_operand_shape(head_dim) + return (VC_MEAN_OPERANDS * rows, cols) + + +def vc_mean_operand_shape(head_dim: int) -> tuple[int, int]: + """Shape of one packed bf16 ``[head_dim x VC_MEAN_MMA_K]`` mean operand, in + rows of 256 elements (the kernel's 512-byte TMA box row).""" + return (head_dim * VC_MEAN_MMA_K // 256, 256) + + +@dataclass(frozen=True) +class VCAttentionConfig: + """Compile-time VC-Attention recipe of one plan. + + With demean on, the default, the kernel restores the tile means. Each + 128-token K/V tile stores E4M3 residuals around its bf16 mean and the + kernel adds the mean back. With demean off the means are zero and the + kernel skips that step. This is the plan for the steps after the V-Smooth + window. The demean flag changes per denoising step while a plan is + compiled once, so a caller holds one plan of each kind. + + A positive repair tile count selects V repair and needs demean off. The + preprocessor appends that many 128-row tiles of the largest quantization + residuals after the sequence and the kernel adds them to the output. Each + repair tile costs one extra K/V tile of attention work. Budgets of two to + five percent of the sequence were used so far and the count cannot exceed + the number of sequence tiles. + + Under VC-Attention-QK8, the plan with E4M3 Q and K, every 128-token Q + block shares one E4M3 Q scale and every 128-token K/V tile one K scale. + """ + + demean: bool = True + repair_tiles: int = 0 + + def __post_init__(self) -> None: + if self.repair_tiles < 0: + raise ValueError(f"repair_tiles must be >= 0, got {self.repair_tiles}") + if self.repair_tiles and self.demean: + raise ValueError("repair_tiles requires demean=False") + + +def vc_repair_tiles(seq_len_kv: int, budget: float) -> int: + """Repair tiles of a V repair run, whole 128-row tiles holding the selected + tokens, ``budget`` times the sequence length rounded and at least one.""" + if not (0.0 < budget < 1.0): + raise ValueError(f"budget must be in (0, 1), got {budget}") + return _blocks(max(1, round(budget * seq_len_kv)), VC_K_BLOCK_SIZE) + + +def vc_repair_kv_len(seq_len_kv: int, repair_tiles: int) -> int: + """K/V rows of a V repair run, the sequence zero-padded to whole tiles plus + the repair tiles.""" + return (_blocks(seq_len_kv, VC_K_BLOCK_SIZE) + repair_tiles) * VC_K_BLOCK_SIZE + + +@dataclass(frozen=True, kw_only=True) +class VCAttentionParams: + """Per-run VC-Attention operands beyond ``q``, the permuted ``k`` and the + E4M3 ``v`` residuals. + + VC-Attention-QK8 plans take ``q_scale`` ``[Hq, B, ceil(Sq / 128)]`` and + ``k_scale`` ``[Hkv, B, ceil(Skv / 128)]``, fp32, one scale per 128-token + block. QK16 plans take neither. ``v_scale`` is the ``[B, Hkv, D]`` fp32 + per-channel E4M3 residual scale. ``tile_means`` is the packed bf16 mean + operand ``[B, Hkv, ceil(num_kv_tiles / 16), 16, 256]`` the preparation + writes, the means already divided by ``v_scale``. The scale must be positive + and finite. The kernel does not check it. A plan with demean off takes no + ``tile_means``. + """ + + q_scale: torch.Tensor | None = None + k_scale: torch.Tensor | None = None + v_scale: torch.Tensor + tile_means: torch.Tensor | None + + +def vc_scale_shapes( + config: VCAttentionConfig, + *, + batch_size: int, + seq_len_kv: int, + num_kv_heads: int, + head_dim: int, + seq_len_q: int | None = None, + num_qo_heads: int | None = None, + qk_fp8: bool = False, +) -> dict[str, tuple[int, ...]]: + """Return the shape of every ``VCAttentionParams`` tensor a plan consumes, by + field name. ``qk_fp8`` adds the Q/K scales of VC-Attention-QK8, which need + ``seq_len_q`` and ``num_qo_heads``.""" + num_kv_tiles = _blocks(seq_len_kv, VC_K_BLOCK_SIZE) + if config.repair_tiles: + # The repair tiles follow the whole tiles. + seq_len_kv = vc_repair_kv_len(seq_len_kv, config.repair_tiles) + shapes: dict[str, tuple[int, ...]] = {} + if qk_fp8: + if seq_len_q is None or num_qo_heads is None: + raise ValueError("qk_fp8 scale shapes need seq_len_q and num_qo_heads") + shapes["q_scale"] = ( + num_qo_heads, + batch_size, + _blocks(seq_len_q, VC_K_BLOCK_SIZE), + ) + shapes["k_scale"] = ( + num_kv_heads, + batch_size, + _blocks(seq_len_kv, VC_K_BLOCK_SIZE), + ) + shapes["v_scale"] = (batch_size, num_kv_heads, head_dim) + if config.demean: + shapes["tile_means"] = ( + batch_size, + num_kv_heads, + _blocks(num_kv_tiles, VC_MEAN_GROUP_TILES), + *vc_mean_group_shape(head_dim), + ) + return shapes + + +_VC_PARAM_DTYPES = { + "q_scale": torch.float32, + "k_scale": torch.float32, + "v_scale": torch.float32, + "tile_means": torch.bfloat16, +} + + +def validate_vc_params( + params: VCAttentionParams, + expected_shapes: dict[str, tuple[int, ...]], + *, + device: torch.device, +) -> None: + """Validate the operands of one run against the plan's expected shapes. + + ``expected_shapes`` is the plan's ``vc_scale_shapes`` result. + """ + if not isinstance(params, VCAttentionParams): + raise TypeError("vc must be a VCAttentionParams instance") + if "tile_means" not in expected_shapes and params.tile_means is not None: + raise ValueError("tile_means must be None on a plan with demean off") + for name, expected_shape in expected_shapes.items(): + tensor = getattr(params, name) + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if tensor.dtype != _VC_PARAM_DTYPES[name]: + raise ValueError( + f"{name} must have dtype {_VC_PARAM_DTYPES[name]}, got {tensor.dtype}" + ) + if tuple(tensor.shape) != tuple(expected_shape): + raise ValueError( + f"{name} must have shape {tuple(expected_shape)}, got {tuple(tensor.shape)}" + ) + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous") + if tensor.data_ptr() % 16 != 0: + raise ValueError(f"{name} data pointer must be 16-byte aligned") + if tensor.device != device: + raise ValueError(f"{name} must be on device {device}, got {tensor.device}") + + +def _blocks(length: int, block: int) -> int: + return (length + block - 1) // block + + +def _per_token_scales(scale: torch.Tensor, s: int, block: int) -> torch.Tensor: + """Expand ``[B, H, nb]`` block scales to ``[B, S, H]``.""" + return scale.permute(0, 2, 1).repeat_interleave(block, dim=1)[:, :s] + + +@torch.no_grad() +def vc_reference( + q: torch.Tensor, ops: VCAttentionOperands, *, sm_scale: float | None = None +) -> torch.Tensor: # noqa: D401 + """fp32 attention of ``q`` over the dequantized VC operands (what the kernel + computes, up to E4M3 P rounding), as ``[B, S_q, H, D]`` fp32. Under + VC-Attention-QK8 the rotated E4M3 ``ops.q`` and ``ops.k`` are dequantized + with their scales and ``q`` is unused.""" + b, s_k, h, d = ops.k.shape + if sm_scale is None: + sm_scale = 1.0 / math.sqrt(d) + if ops.q_scale is not None: + s_q = ops.q.shape[1] + q_scale = ops.q_scale.permute(1, 0, 2) + k_scale = ops.k_scale.permute(1, 0, 2) + q = ops.q.float() * _per_token_scales(q_scale, s_q, VC_K_BLOCK_SIZE).unsqueeze( + -1 + ) + k = ops.k.float() * _per_token_scales(k_scale, s_k, VC_K_BLOCK_SIZE).unsqueeze( + -1 + ) + else: + q = q.float() + k = ops.k.float() + v = ops.v.float() * ops.v_scale.unsqueeze(1) + repairs_v = ops.repair_tiles > 0 + if repairs_v: + # Repair rows add to the numerator only, the zero padding of the partial + # last tile is masked, and the denominator runs over the original tokens. + repair = ops.repair_tiles * VC_K_BLOCK_SIZE + whole = ops.seq_len_kv // VC_K_BLOCK_SIZE * VC_K_BLOCK_SIZE + end = ops.seq_len_kv + repair + row = torch.arange(s_k, device=k.device) + original = (row < whole) | ((row >= whole + repair) & (row < end)) + scores = torch.einsum("bqhd,bkhd->bhqk", q, k) * sm_scale + scores = scores.masked_fill(row >= end, -torch.inf) + weights = torch.exp(scores - scores.amax(dim=-1, keepdim=True)) + out = torch.einsum("bhqk,bkhd->bqhd", weights, v) + denom = weights[..., original].sum(dim=-1).permute(0, 2, 1) # [B, Sq, H] + return out / denom.unsqueeze(-1) + # The kernel restores bf16(mean / v_scale) * v_scale, per channel. + vs = ops.v_scale.reshape(b, ops.k.shape[2], 1, d) # [B, H, 1, D] + mean = (ops.mean / vs).to(torch.bfloat16).float() * vs + mean_tokens = mean.permute(0, 2, 1, 3).repeat_interleave(VC_K_BLOCK_SIZE, dim=1)[ + :, :s_k + ] + v = v + mean_tokens + scores = torch.einsum("bqhd,bkhd->bhqk", q, k) * sm_scale + probs = torch.softmax(scores, dim=-1) + return torch.einsum("bhqk,bkhd->bqhd", probs, v) + + +__all__ = [ + "E4M3_MAX", + "VC_HEAD_DIM", + "VC_K_BLOCK_SIZE", + "VC_MEAN_MMA_K", + "VCAttentionConfig", + "VCAttentionParams", + "validate_vc_params", + "vc_repair_kv_len", + "vc_repair_tiles", + "vc_mean_group_shape", + "vc_mean_operand_shape", + "VC_MEAN_GROUP_TILES", + "VC_MEAN_OPERANDS", + "vc_reference", + "vc_scale_shapes", +] diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/vc_attention/prepare.py b/tensorrt_llm/_torch/attention/backends/prims_ts/vc_attention/prepare.py new file mode 100644 index 000000000000..c67acd80a71d --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/vc_attention/prepare.py @@ -0,0 +1,1498 @@ +"""VC-Attention operand preparation kernels (CuTe DSL). + +``_prepare`` turns bf16 or fp16 ``[B, S, H, D]`` K and V, and Q under +VC-Attention-QK8, into the prims_ts VC-Attention operands with one CTA per +batch, head and 128-token tile. Each input tensor is read once, V twice, and +each output written once. Shapes are compile-time constants of the kernels and +the launches take tensors only. An operand a recipe does not use is ``None`` +and the kernels take their mode from that. + +* ``VcKvPass1`` writes the V tile means and the per-channel residual amax, plus + the permuted K rows under QK16 or the K tile channel sums under QK8. +* ``VcKvPass2`` writes the V scale, the E4M3 V residuals and the packed bf16 + tile-mean UMMA operand, the K centre under QK8 and the residual energies + under repair. +* ``VcQKPass`` writes the centred, Hadamard-rotated E4M3 Q and K with one + scale per 128-token block in the flat per-sequence scale tables, QK8 only. +* ``VcSelectPass`` and ``VcRepairRowsPass`` pick the tokens of largest + residual energy and write their repair rows, repair only. + +Mean recipes take tokens in ``perm`` order. Repair recipes keep the input +order, move the partial last tile past the repair tiles and write each token's +residual energy instead of the mean operand. Following the paper, V residuals +get one E4M3 scale per batch, head and channel, and the tile means are stored +divided by that scale. Under QK16 Q and K stay in their input dtype. + +Every kernel runs 256 threads. Thread ``t`` owns head_dim columns +``(t % 16) * 8`` to ``+ 8``, one 16-byte vector, of rows ``t // 16 + 16 * i`` +for ``i < 8``. Rows stream through registers and the per-column state stays 8 +wide, with every load of a thread issued before the guarded math. The pre-pass +is memory-latency bound, so it needs that small register footprint for the +occupancy that keeps enough loads in flight. +""" + +from __future__ import annotations + +import functools +import math + +import cutlass +import cutlass.cute as cute +import torch +from cutlass import Float32, Int32, Int64 + +from . import ( + E4M3_MAX, + VC_HEAD_DIM, + VC_K_BLOCK_SIZE, + VC_MEAN_GROUP_TILES, + VC_MEAN_MMA_K, + VC_MEAN_OPERANDS, + vc_mean_group_shape, + vc_mean_operand_shape, + vc_repair_kv_len, +) +from .._compile_cache import build_and_load_attention_kernel +from ..kernels.fmha_decode.fmha_decode_resources.helpers_common import ( + _pack_float2_to_bf16, +) +from ..kernels.fmha_decode.fmha_decode_resources.helpers_softmax import ( + _pack_float4_to_fp8_e4m3, +) + +_COMPILE_OPTIONS = "--enable-tvm-ffi --opt-level 3" +_THREADS = 256 +_COLS_PER_THREAD = 8 +_COL_GROUPS = VC_HEAD_DIM // _COLS_PER_THREAD +_ROW_GROUPS = _THREADS // _COL_GROUPS +_ROWS_PER_THREAD = VC_K_BLOCK_SIZE // _ROW_GROUPS +_WORDS = _COLS_PER_THREAD // 2 +_MU_ELEMS = math.prod(vc_mean_operand_shape(VC_HEAD_DIM)) +_HADAMARD_NORM = 1.0 / math.sqrt(VC_HEAD_DIM) +_SELECT_BINS = 256 +assert _THREADS == _SELECT_BINS, "VcSelectPass owns one histogram bin per thread" + + +@cute.jit +def _hadamard128_row(vals, base: cutlass.Constexpr[int], cg: Int32): + """Unnormalised 128-point Walsh-Hadamard transform of one row in place. + + This thread holds 8 columns and the 15 lanes of its column group hold the + rest. The caller folds the normalisation into its scale. The butterflies run + as packed f32x2 pairs because the pass is issue-bound. + """ + for h in (1, 2, 4): + lo = [i for i in range(_COLS_PER_THREAD) if (i & h) == 0] + for q in cutlass.range_constexpr(0, len(lo), 2): + i0, i1 = lo[q], lo[q + 1] + a = (vals[base + i0], vals[base + i1]) + b = (vals[base + (i0 ^ h)], vals[base + (i1 ^ h)]) + vals[base + i0], vals[base + i1] = cute.arch.add_packed_f32x2(a, b) + vals[base + (i0 ^ h)], vals[base + (i1 ^ h)] = cute.arch.sub_packed_f32x2( + a, b + ) + for sbit in cutlass.range_constexpr(4): + stride = 1 << sbit + sign = Float32(1.0) - Float32(2.0) * Float32((cg >> sbit) & 1) + for j in cutlass.range_constexpr(0, _COLS_PER_THREAD, 2): + m0 = vals[base + j] + m1 = vals[base + j + 1] + o0 = cute.arch.shuffle_sync_bfly(m0, stride) + o1 = cute.arch.shuffle_sync_bfly(m1, stride) + vals[base + j], vals[base + j + 1] = cute.arch.fma_packed_f32x2( + (m0, m1), (sign, sign), (o0, o1) + ) + + +@cute.jit +def _block_max(v: Float32, red_smem, tidx: Int32) -> Float32: + """Max over the threads of the CTA. + + The pre-pass is a plain streaming kernel with no tensor memory allocation, so + the reduction is the usual warp shuffle tree followed by one word per warp in + shared memory. The attention kernel's TMEM reductions are not available here. + """ + for shift in cutlass.range_constexpr(5): + v = cute.arch.fmax(v, cute.arch.shuffle_sync_bfly(v, 1 << shift)) + if tidx % 32 == 0: + red_smem[tidx // 32] = v + cute.arch.sync_threads() + r = red_smem[0] + for w in cutlass.range_constexpr(1, _THREADS // 32): + r = cute.arch.fmax(r, red_smem[w]) + cute.arch.sync_threads() + return r + + +@cute.jit +def _load_raw(raw, base: cutlass.Constexpr[int], addr: Int64): + """Load this thread's 16 bytes at ``addr`` into four words of ``raw``.""" + regs = cutlass.inttoptr(addr, mem_space=1, dtype=Int32).load( + count=_WORDS, alignment=16 + ) + for q in cutlass.range_constexpr(_WORDS): + raw[base + q] = regs[q] + + +@cute.jit +def _row(raw, base: cutlass.Constexpr[int], is_bf16: cutlass.Constexpr[bool]): + """The 8 bf16 or fp16 values held in four words of ``raw`` as Float32.""" + regs = (raw.data_ptr() + base).load(count=_WORDS, alignment=16) + if cutlass.const_expr(is_bf16): + return regs.bitcast(cutlass.BFloat16).to(Float32) + return regs.bitcast(cutlass.Float16).to(Float32) + + +@cute.jit +def _load_row_f32(addr: Int64): + return cutlass.inttoptr(addr, mem_space=1, dtype=Float32).load( + count=_COLS_PER_THREAD, alignment=32 + ) + + +@cute.jit +def _pack_row_fp8(vals, base: cutlass.Constexpr[int], scale_inv: Float32): + """This thread's 8 values from ``base`` times ``scale_inv`` as two E4M3 words.""" + packed = cutlass.Array( + Int32, _COLS_PER_THREAD // 4, space=cutlass.AddressSpace.rmem + ) + for j in cutlass.range_constexpr(_COLS_PER_THREAD // 4): + packed[j] = _pack_float4_to_fp8_e4m3( + vals[base + 4 * j] * scale_inv, + vals[base + 4 * j + 1] * scale_inv, + vals[base + 4 * j + 2] * scale_inv, + vals[base + 4 * j + 3] * scale_inv, + ) + return packed.data_ptr().load(count=_COLS_PER_THREAD // 4, alignment=8) + + +@cute.jit +def _coords(num_heads: cutlass.Constexpr[int]): + """This CTA's tile, batch and head, this thread's row group and first column.""" + tidx, _, _ = cute.arch.thread_idx() + t, bh, _ = cute.arch.block_idx() + return ( + t, + bh, + bh // num_heads, + bh % num_heads, + tidx // _COL_GROUPS, + (tidx % _COL_GROUPS) * _COLS_PER_THREAD, + ) + + +@cute.jit +def _row_base( + base: Int64, + b: Int32, + rows: cutlass.Constexpr[int], + h: Int32, + col0: Int32, + elem_bytes: cutlass.Constexpr[int], + row_bytes: cutlass.Constexpr[int], +) -> Int64: + """Address of this thread's columns of row 0 of batch ``b`` in a ``[B, rows, H, D]`` tensor.""" + return ( + base + + Int64(b) * rows * row_bytes + + (Int64(h) * VC_HEAD_DIM + col0) * elem_bytes + ) + + +@cute.jit +def _store_tile_cols( + base: Int64, + bh: Int32, + num_tiles: cutlass.Constexpr[int], + t: Int32, + col0: Int32, + vals, +): + """Store this thread's 8 columns of tile ``t`` of a ``[B, H, T, D]`` fp32 tensor.""" + cutlass.inttoptr( + base + ((Int64(bh) * num_tiles + t) * VC_HEAD_DIM + col0) * 4, + mem_space=1, + dtype=Float32, + ).store(vals.data_ptr().load(count=_COLS_PER_THREAD, alignment=32), alignment=32) + + +@cute.jit +def _cta_col_reduce( + col_sums, vals, rg: Int32, col0: Int32, is_max: cutlass.Constexpr[bool] +): + """Reduce this thread's 8 column values over the row groups through shared memory.""" + cute.arch.sync_threads() + for j in cutlass.range_constexpr(_COLS_PER_THREAD): + col_sums[rg * VC_K_BLOCK_SIZE + col0 + j] = vals[j] + cute.arch.sync_threads() + for j in cutlass.range_constexpr(_COLS_PER_THREAD): + acc = col_sums[col0 + j] + for g in cutlass.range_constexpr(1, _ROW_GROUPS): + if cutlass.const_expr(is_max): + acc = cute.arch.fmax(acc, col_sums[g * VC_K_BLOCK_SIZE + col0 + j]) + else: + acc = acc + col_sums[g * VC_K_BLOCK_SIZE + col0 + j] + vals[j] = acc + + +@cute.jit +def _v_scale(vmax: Float32) -> Float32: + return cute.arch.fmax(vmax * Float32(1.0 / E4M3_MAX), Float32(1e-12)) + + +@cute.jit +def _stats_addr( + mStats: cute.Tensor, bh: Int32, slot: cutlass.Constexpr[int], col0: Int32 +) -> Int64: + """Slot 0 holds the K channel sums and slot 1 the V residual amax.""" + return mStats.iterator.toint() + ((Int64(bh) * 2 + slot) * VC_HEAD_DIM + col0) * 4 + + +@cute.jit +def _scale_slot( + bh: Int32, + block: Int32, + num_heads: cutlass.Constexpr[int], + batch: cutlass.Constexpr[int], + rows: cutlass.Constexpr[int], +) -> Int64: + """Element of block ``block`` of sequence ``bh`` in the ``[H, B, blocks]`` scale table.""" + blocks = (rows + VC_K_BLOCK_SIZE - 1) // VC_K_BLOCK_SIZE + return Int64(bh % num_heads) * (batch * blocks) + Int64( + (bh // num_heads) * blocks + block + ) + + +class _Shape: + """Compile-time geometry shared by the kernels of one preparation.""" + + def __init__( + self, + is_bf16: bool, + batch: int, + seq_len_kv: int, + num_heads: int, + repair_tiles: int, + ): + self.is_bf16 = is_bf16 + self.batch = batch + self.seq_len_kv = seq_len_kv + self.num_heads = num_heads + self.num_tiles = (seq_len_kv + VC_K_BLOCK_SIZE - 1) // VC_K_BLOCK_SIZE + self.repair_tiles = repair_tiles + self.repair = repair_tiles > 0 + self.rows = ( + vc_repair_kv_len(seq_len_kv, repair_tiles) if self.repair else seq_len_kv + ) + self.tail = ( + seq_len_kv // VC_K_BLOCK_SIZE * VC_K_BLOCK_SIZE + if self.repair + else seq_len_kv + ) + self.shift = repair_tiles * VC_K_BLOCK_SIZE + self.row_bytes = num_heads * VC_HEAD_DIM * 2 + self.row_bytes8 = num_heads * VC_HEAD_DIM + + +class VcKvPass1(_Shape): + def __init__( + self, is_bf16, batch, seq_len_kv, num_heads, repair_tiles, demean: bool + ): + super().__init__(is_bf16, batch, seq_len_kv, num_heads, repair_tiles) + self.demean = demean + + @cute.kernel + def kernel( + self, + mK: cute.Tensor, + mV: cute.Tensor, + mPerm: cute.Tensor | None, + mKPerm: cute.Tensor | None, + mMean: cute.Tensor, + mKSum: cute.Tensor | None, + mStats: cute.Tensor, + ): + """Pass 1 over one 128-token K/V tile of one batch head. + + Reads K and V once. With a permutation it writes the permuted K rows, + otherwise it accumulates the K channel sums the QK8 centre needs. From + V it writes the tile mean and folds the tile's residual amax into the + per-channel statistics that pass 2 turns into the V scale. + """ + # Thread t owns 8 columns of rows t // 16 + 16 i for i < 8. col_sums is + # the scratch of the CTA-wide per-column reductions. + t, bh, b, h, rg, col0 = _coords(self.num_heads) + smem = cutlass.utils.SmemAllocator() + col_sums = smem.allocate_array(Float32, _ROW_GROUPS * VC_K_BLOCK_SIZE) + k_base = _row_base( + mK.iterator.toint(), b, self.seq_len_kv, h, col0, 2, self.row_bytes + ) + v_base = _row_base( + mV.iterator.toint(), b, self.seq_len_kv, h, col0, 2, self.row_bytes + ) + if cutlass.const_expr(mKPerm is not None): + kp_base = _row_base( + mKPerm.iterator.toint(), b, self.rows, h, col0, 2, self.row_bytes + ) + + # Rows past the sequence read token 0 and skip the math, so every load + # of a thread is in flight together. + toks = cutlass.Array(Int32, _ROWS_PER_THREAD, space=cutlass.AddressSpace.rmem) + for i in cutlass.range_constexpr(_ROWS_PER_THREAD): + pos = t * VC_K_BLOCK_SIZE + rg + _ROW_GROUPS * i + tok = Int32(0) + if pos < self.seq_len_kv: + tok = pos + if cutlass.const_expr(mPerm is not None): + tok = cutlass.inttoptr( + mPerm.iterator.toint() + + (Int64(bh) * self.seq_len_kv + pos) * 4, + mem_space=1, + dtype=Int32, + ).load() + toks[i] = tok + raw = cutlass.Array( + Int32, _ROWS_PER_THREAD * _WORDS, space=cutlass.AddressSpace.rmem + ) + + # K. Either copy the permuted rows to their output slot, shifted past the + # repair tiles for the partial tail, or sum the channels for the centre. + if cutlass.const_expr(mKPerm is not None): + for i in cutlass.range_constexpr(_ROWS_PER_THREAD): + _load_raw(raw, i * _WORDS, k_base + Int64(toks[i]) * self.row_bytes) + for i in cutlass.range_constexpr(_ROWS_PER_THREAD): + pos = t * VC_K_BLOCK_SIZE + rg + _ROW_GROUPS * i + if pos < self.seq_len_kv: + out_row = pos + self.shift * Int32(pos >= self.tail) + cutlass.inttoptr( + kp_base + Int64(out_row) * self.row_bytes, + mem_space=1, + dtype=Int32, + ).store( + (raw.data_ptr() + i * _WORDS).load(count=_WORDS, alignment=16), + alignment=16, + ) + else: + # K channel sums of the tile in input order. Pass 2 sums the tiles + # in a fixed order so the centre is deterministic. + ksum = cutlass.Array( + Float32, _COLS_PER_THREAD, space=cutlass.AddressSpace.rmem + ) + for j in cutlass.range_constexpr(_COLS_PER_THREAD): + ksum[j] = Float32(0.0) + for i in cutlass.range_constexpr(_ROWS_PER_THREAD): + pos = t * VC_K_BLOCK_SIZE + rg + _ROW_GROUPS * i + kpos = Int32(0) + if pos < self.seq_len_kv: + kpos = pos + _load_raw(raw, i * _WORDS, k_base + Int64(kpos) * self.row_bytes) + for i in cutlass.range_constexpr(_ROWS_PER_THREAD): + pos = t * VC_K_BLOCK_SIZE + rg + _ROW_GROUPS * i + if pos < self.seq_len_kv: + row = _row(raw, i * _WORDS, self.is_bf16) + for j in cutlass.range_constexpr(_COLS_PER_THREAD): + ksum[j] = ksum[j] + row[j] + _cta_col_reduce(col_sums, ksum, rg, col0, False) + if rg == 0: + _store_tile_cols( + mKSum.iterator.toint(), bh, self.num_tiles, t, col0, ksum + ) + + # V. Per-column sum, max and min of the tile's rows give the mean and the + # residual amax exactly, since max |v - mean| is max(vmax - mean, mean - vmin). + psum = cutlass.Array(Float32, _COLS_PER_THREAD, space=cutlass.AddressSpace.rmem) + pmax = cutlass.Array(Float32, _COLS_PER_THREAD, space=cutlass.AddressSpace.rmem) + pmin = cutlass.Array(Float32, _COLS_PER_THREAD, space=cutlass.AddressSpace.rmem) + for j in cutlass.range_constexpr(_COLS_PER_THREAD): + psum[j] = Float32(0.0) + pmax[j] = -Float32.inf + pmin[j] = Float32.inf + for i in cutlass.range_constexpr(_ROWS_PER_THREAD): + _load_raw(raw, i * _WORDS, v_base + Int64(toks[i]) * self.row_bytes) + for i in cutlass.range_constexpr(_ROWS_PER_THREAD): + pos = t * VC_K_BLOCK_SIZE + rg + _ROW_GROUPS * i + if pos < self.seq_len_kv: + row = _row(raw, i * _WORDS, self.is_bf16) + for j in cutlass.range_constexpr(_COLS_PER_THREAD): + psum[j] = psum[j] + row[j] + pmax[j] = cute.arch.fmax(pmax[j], row[j]) + pmin[j] = cute.arch.fmin(pmin[j], row[j]) + valid_rows = self.seq_len_kv - t * VC_K_BLOCK_SIZE + if valid_rows > VC_K_BLOCK_SIZE: + valid_rows = Int32(VC_K_BLOCK_SIZE) + inv_count = Float32(1.0) / Float32(valid_rows) + _cta_col_reduce(col_sums, psum, rg, col0, False) + mean = cutlass.Array(Float32, _COLS_PER_THREAD, space=cutlass.AddressSpace.rmem) + for j in cutlass.range_constexpr(_COLS_PER_THREAD): + mean[j] = psum[j] * inv_count + if cutlass.const_expr(not self.demean): + mean[j] = Float32(0.0) + if rg == 0: + _store_tile_cols(mMean.iterator.toint(), bh, self.num_tiles, t, col0, mean) + # The residual amax of this tile joins the per-channel maximum over all + # tiles with one atomic per column. Pass 2 derives the V scale from it. + cmax = cutlass.Array(Float32, _COLS_PER_THREAD, space=cutlass.AddressSpace.rmem) + for j in cutlass.range_constexpr(_COLS_PER_THREAD): + cmax[j] = Float32(0.0) + if pmax[j] >= pmin[j]: + cmax[j] = cute.arch.fmax(pmax[j] - mean[j], mean[j] - pmin[j]) + _cta_col_reduce(col_sums, cmax, rg, col0, True) + if rg == 0: + for j in cutlass.range_constexpr(_COLS_PER_THREAD): + cute.arch.atomic_fmax( + cutlass.inttoptr( + _stats_addr(mStats, bh, 1, col0) + j * 4, + mem_space=1, + dtype=Float32, + ), + cmax[j], + sign_bit=False, + ) + + @cute.jit + def __call__(self, mK, mV, mPerm, mKPerm, mMean, mKSum, mStats, stream): + self.kernel(mK, mV, mPerm, mKPerm, mMean, mKSum, mStats).launch( + grid=[self.num_tiles, self.batch * self.num_heads, 1], + block=[_THREADS, 1, 1], + smem=_ROW_GROUPS * VC_K_BLOCK_SIZE * 4, + stream=stream, + ) + + +class VcKvPass2(_Shape): + @cute.kernel + def kernel( + self, + mV: cute.Tensor, + mPerm: cute.Tensor | None, + mMean: cute.Tensor, + mKSum: cute.Tensor | None, + mStats: cute.Tensor, + mVScale: cute.Tensor, + mV8: cute.Tensor, + mMu: cute.Tensor | None, + mEnergy: cute.Tensor | None, + ): + """Pass 2 over one 128-token K/V tile of one batch head. + + With the per-channel residual amax of pass 1 known, it derives the E4M3 + V scale, stores the E4M3 residuals of the permuted V rows, the per-token + residual energy for the repair selection when asked, and the packed bf16 + mean operand of the tile's group for the kernel's mean UMMA step. Tile 0 + also writes the V scale and finishes the QK8 K centre sum. + """ + tidx, _, _ = cute.arch.thread_idx() + t, bh, _ = cute.arch.block_idx() + b = bh // self.num_heads + h = bh % self.num_heads + cg = tidx % _COL_GROUPS + rg = tidx // _COL_GROUPS + col0 = cg * _COLS_PER_THREAD + v_base = ( + mV.iterator.toint() + + Int64(b) * self.seq_len_kv * self.row_bytes + + (Int64(h) * VC_HEAD_DIM + col0) * 2 + ) + mean_base = mMean.iterator.toint() + (Int64(bh) * self.num_tiles + t) * ( + VC_HEAD_DIM * 4 + ) + smem = cutlass.utils.SmemAllocator() + energy_part = smem.allocate_array(Float32, _THREADS * _ROWS_PER_THREAD) + + # Step 1. The V scale of this thread's channels from the residual amax. + vmax_base = _stats_addr(mStats, bh, 1, 0) + vmax = _load_row_f32(vmax_base + col0 * 4) + vs = cutlass.Array(Float32, _COLS_PER_THREAD, space=cutlass.AddressSpace.rmem) + vs_inv = cutlass.Array( + Float32, _COLS_PER_THREAD, space=cutlass.AddressSpace.rmem + ) + for j in cutlass.range_constexpr(_COLS_PER_THREAD): + vs[j] = _v_scale(vmax[j]) + vs_inv[j] = Float32(1.0) / vs[j] + if (t == 0) & (rg == 0): + cutlass.inttoptr( + mVScale.iterator.toint() + (Int64(bh) * VC_HEAD_DIM + col0) * 4, + mem_space=1, + dtype=Float32, + ).store( + vs.data_ptr().load(count=_COLS_PER_THREAD, alignment=32), alignment=32 + ) + # Tile 0 sums the per-tile K channel sums of pass 1 in a fixed order, so + # the QK8 centre is deterministic. + if cutlass.const_expr(mKSum is not None): + if (t == 0) & (tidx < VC_HEAD_DIM): + ksum_base = ( + mKSum.iterator.toint() + + (Int64(bh) * self.num_tiles * VC_HEAD_DIM + tidx) * 4 + ) + acc = Float32(0.0) + for i in cutlass.range(self.num_tiles): + acc = ( + acc + + cutlass.inttoptr( + ksum_base + Int64(i) * (VC_HEAD_DIM * 4), + mem_space=1, + dtype=Float32, + ).load() + ) + cutlass.inttoptr( + _stats_addr(mStats, bh, 0, 0) + tidx * 4, mem_space=1, dtype=Float32 + ).store(acc) + + # Step 2. Reload the permuted V rows, subtract the tile mean, quantize the + # residuals to E4M3 and store them at the shifted output row. + mean = _load_row_f32(mean_base + col0 * 4) + raw = cutlass.Array( + Int32, _ROWS_PER_THREAD * _WORDS, space=cutlass.AddressSpace.rmem + ) + for i in cutlass.range_constexpr(_ROWS_PER_THREAD): + pos = t * VC_K_BLOCK_SIZE + rg + _ROW_GROUPS * i + tok = Int32(0) + if pos < self.seq_len_kv: + tok = pos + if cutlass.const_expr(mPerm is not None): + tok = cutlass.inttoptr( + mPerm.iterator.toint() + + (Int64(bh) * self.seq_len_kv + pos) * 4, + mem_space=1, + dtype=Int32, + ).load() + _load_raw(raw, i * _WORDS, v_base + Int64(tok) * self.row_bytes) + for i in cutlass.range_constexpr(_ROWS_PER_THREAD): + pos = t * VC_K_BLOCK_SIZE + rg + _ROW_GROUPS * i + energy = Float32(0.0) + if pos < self.seq_len_kv: + row = _row(raw, i * _WORDS, self.is_bf16) + res = cutlass.Array( + Float32, _COLS_PER_THREAD, space=cutlass.AddressSpace.rmem + ) + for j in cutlass.range_constexpr(_COLS_PER_THREAD): + res[j] = (row[j] - mean[j]) * vs_inv[j] + out_row = pos + self.shift * Int32(pos >= self.tail) + out_addr = ( + mV8.iterator.toint() + + (Int64(b) * self.rows + out_row) * self.row_bytes8 + + Int64(h) * VC_HEAD_DIM + + col0 + ) + if cutlass.const_expr(mEnergy is not None): + # The squared E4M3 rounding error in value units ranks the + # tokens for repair. + words = _pack_row_fp8(res, 0, Float32(1.0)) + codes = words.bitcast(cutlass.Float8E4M3FN).to(Float32) + for c in cutlass.range_constexpr(_COLS_PER_THREAD): + err = (res[c] - codes[c]) * vs[c] + energy = energy + err * err + cutlass.inttoptr(out_addr, mem_space=1, dtype=Int32).store( + words, alignment=8 + ) + else: + cutlass.inttoptr(out_addr, mem_space=1, dtype=Int32).store( + _pack_row_fp8(res, 0, Float32(1.0)), alignment=8 + ) + if cutlass.const_expr(mEnergy is not None): + energy_part[(rg + _ROW_GROUPS * i) * _COL_GROUPS + cg] = energy + # Step 3. V repair only. Sum each row's partial energies over the 16 + # column groups and store one energy per token. + if cutlass.const_expr(mEnergy is not None): + cute.arch.sync_threads() + for i in cutlass.range_constexpr(_ROWS_PER_THREAD): + pos = t * VC_K_BLOCK_SIZE + rg + _ROW_GROUPS * i + if (cg == 0) & (pos < self.seq_len_kv): + acc = Float32(0.0) + for c in cutlass.range_constexpr(_COL_GROUPS): + acc = ( + acc + energy_part[(rg + _ROW_GROUPS * i) * _COL_GROUPS + c] + ) + cutlass.inttoptr( + mEnergy.iterator.toint() + + (Int64(bh) * self.seq_len_kv + pos) * 4, + mem_space=1, + dtype=Float32, + ).store(acc) + + if cutlass.const_expr(mMu is not None): + # Step 4. Means only. + # Mean operand of this tile's group in core-matrix order. Tile 8o+i + # of a group owns K slots 2i and 2i+1 of operand o. Each of the first + # 128 threads stores its two rows' slots, and the last tile stores + # zero slots for the tiles its group lacks. Means are stored divided + # by the per-channel scale. + k_per_core = VC_MEAN_MMA_K // 2 + idx0 = tidx * (2 * k_per_core) + g = idx0 // (8 * VC_MEAN_MMA_K) + c = (idx0 // (8 * k_per_core)) % 2 + d0 = g * 8 + (idx0 // k_per_core) % 8 + num_groups = ( + self.num_tiles + VC_MEAN_GROUP_TILES - 1 + ) // VC_MEAN_GROUP_TILES + group0 = t - t % VC_MEAN_GROUP_TILES + if tidx < 128: + m0 = cutlass.inttoptr( + mean_base + Int64(d0) * 4, mem_space=1, dtype=Float32 + ).load() + m1 = cutlass.inttoptr( + mean_base + Int64(d0 + 1) * 4, mem_space=1, dtype=Float32 + ).load() + s0 = _v_scale( + cutlass.inttoptr( + vmax_base + Int64(d0) * 4, mem_space=1, dtype=Float32 + ).load() + ) + s1 = _v_scale( + cutlass.inttoptr( + vmax_base + Int64(d0 + 1) * 4, mem_space=1, dtype=Float32 + ).load() + ) + own0 = _pack_float2_to_bf16(m0 / s0, m0 / s0) + own1 = _pack_float2_to_bf16(m1 / s1, m1 / s1) + for i in cutlass.range_constexpr(VC_MEAN_GROUP_TILES): + operand = i // k_per_core + slot = 2 * (i % k_per_core) + if cutlass.const_expr(slot // k_per_core == 0): + mine = c == 0 + else: + mine = c == 1 + tile = group0 + i + if mine & ((tile == t) | ((t == self.num_tiles - 1) & (tile > t))): + val0 = Int32(0) + val1 = Int32(0) + if tile == t: + val0 = own0 + val1 = own1 + dst = ( + mMu.iterator.toint() + + ( + ( + (Int64(bh) * num_groups + t // VC_MEAN_GROUP_TILES) + * VC_MEAN_OPERANDS + + operand + ) + * _MU_ELEMS + + idx0 + + slot % k_per_core + ) + * 2 + ) + cutlass.inttoptr(dst, mem_space=1, dtype=Int32).store(val0) + cutlass.inttoptr( + dst + k_per_core * 2, mem_space=1, dtype=Int32 + ).store(val1) + + @cute.jit + def __call__( + self, mV, mPerm, mMean, mKSum, mStats, mVScale, mV8, mMu, mEnergy, stream + ): + self.kernel(mV, mPerm, mMean, mKSum, mStats, mVScale, mV8, mMu, mEnergy).launch( + grid=[self.num_tiles, self.batch * self.num_heads, 1], + block=[_THREADS, 1, 1], + smem=_THREADS * _ROWS_PER_THREAD * 4, + stream=stream, + ) + + +class VcQKPass(_Shape): + """QK8 quantization of Q and K in one launch, one CTA per 128-token block. + + The first ``blocks_q`` blocks are Q, the rest are K through ``perm`` or the + repair row map, with the repair tiles through the selection index. K is + centred with the channel sums of the statistics. The rotated block stays + in registers between the amax and the store. + """ + + def __init__( + self, + is_bf16, + batch, + seq_len_kv, + num_heads, + repair_tiles, + seq_len_q: int, + repair_tokens: int = 0, + ): + super().__init__(is_bf16, batch, seq_len_kv, num_heads, repair_tiles) + self.seq_len_q = seq_len_q + self.repair_tokens = repair_tokens + self.blocks_q = (seq_len_q + VC_K_BLOCK_SIZE - 1) // VC_K_BLOCK_SIZE + + @cute.jit + def _block( + self, + nb: Int32, + co, + mX: cute.Tensor, + mPerm: cute.Tensor | None, + mIndex: cute.Tensor | None, + mStats: cute.Tensor | None, + mX8: cute.Tensor, + mScale: cute.Tensor, + red, + rows_out: cutlass.Constexpr[int], + src: cutlass.Constexpr[int], + ): + """One block of Q, of K when ``mStats`` is given, or of the repair tiles + when ``mIndex`` is given. ``co`` carries the CTA and thread coordinates.""" + # One CTA per 128-token block of one batch head. Thread t owns 8 columns of + # rows t // 16 + 16 i for i < 8, so the 256 threads cover the 128 x 128 block. + _, bh, b, h, rg, col0 = co + cg = col0 // _COLS_PER_THREAD + tidx = rg * _COL_GROUPS + cg + x_base = _row_base(mX.iterator.toint(), b, src, h, col0, 2, self.row_bytes) + centre = cutlass.Array( + Float32, _COLS_PER_THREAD, space=cutlass.AddressSpace.rmem + ) + for j in cutlass.range_constexpr(_COLS_PER_THREAD): + centre[j] = Float32(0.0) + # K only. The per head channel mean from the statistics pass, subtracted + # before the rotation. + if cutlass.const_expr(mStats is not None): + c = _load_row_f32(_stats_addr(mStats, bh, 0, col0)) + for j in cutlass.range_constexpr(_COLS_PER_THREAD): + centre[j] = c[j] * Float32(1.0 / src) + + # Step 1. Load this thread's 8 rows of the block, through the row map when + # gathering, centre them and keep them in registers for the whole kernel. + # The source token per row is -1 for rows outside the sequence and for + # zero rows. Loads read token 0 for those and the math is skipped. + toks = cutlass.Array(Int32, _ROWS_PER_THREAD, space=cutlass.AddressSpace.rmem) + for i in cutlass.range_constexpr(_ROWS_PER_THREAD): + pos = nb * VC_K_BLOCK_SIZE + rg + _ROW_GROUPS * i + tok = Int32(-1) + if pos < rows_out: + tok = pos + if cutlass.const_expr(mIndex is not None): + tok = Int32(-1) + sel = pos - self.tail + if (sel >= 0) & (sel < self.repair_tokens): + tok = cutlass.inttoptr( + mIndex.iterator.toint() + + (Int64(bh) * self.repair_tokens + sel) * 4, + mem_space=1, + dtype=Int32, + ).load() + elif cutlass.const_expr(mPerm is not None): + tok = cutlass.inttoptr( + mPerm.iterator.toint() + + (Int64(bh) * self.seq_len_kv + pos) * 4, + mem_space=1, + dtype=Int32, + ).load() + elif cutlass.const_expr(mStats is not None and self.repair): + if pos >= self.tail: + tok = Int32(-1) + if (pos >= self.tail + self.shift) & ( + pos < self.seq_len_kv + self.shift + ): + tok = pos - self.shift + toks[i] = tok + raw = cutlass.Array( + Int32, _ROWS_PER_THREAD * _WORDS, space=cutlass.AddressSpace.rmem + ) + for i in cutlass.range_constexpr(_ROWS_PER_THREAD): + tok = toks[i] + if tok < 0: + tok = Int32(0) + _load_raw(raw, i * _WORDS, x_base + Int64(tok) * self.row_bytes) + vals = cutlass.Array( + Float32, + _ROWS_PER_THREAD * _COLS_PER_THREAD, + space=cutlass.AddressSpace.rmem, + ) + for i in cutlass.range_constexpr(_ROWS_PER_THREAD): + if toks[i] >= 0: + row = _row(raw, i * _WORDS, self.is_bf16) + for j in cutlass.range_constexpr(_COLS_PER_THREAD): + vals[i * _COLS_PER_THREAD + j] = row[j] - centre[j] + else: + # The rotation shuffles need all 16 lanes of a row. + for j in cutlass.range_constexpr(_COLS_PER_THREAD): + vals[i * _COLS_PER_THREAD + j] = Float32(0.0) + # Step 2. Rotate every row by the 128 point Hadamard transform. The 16 + # threads of a row exchange values through warp shuffles, so this needs no + # shared memory and the rows stay in registers. + for i in cutlass.range_constexpr(_ROWS_PER_THREAD): + _hadamard128_row(vals, i * _COLS_PER_THREAD, cg) + # Step 3. Block amax of the rotated values over the valid rows, reduced + # across the CTA. One E4M3 scale per block follows from it. + amax = Float32(0.0) + for i in cutlass.range_constexpr(_ROWS_PER_THREAD): + pos = nb * VC_K_BLOCK_SIZE + rg + _ROW_GROUPS * i + if pos < rows_out: + for j in cutlass.range_constexpr(_COLS_PER_THREAD): + v = vals[i * _COLS_PER_THREAD + j] + amax = cute.arch.fmax(amax, cute.arch.fmax(v, -v)) + amax = _block_max(amax, red, tidx) + scale = cute.arch.fmax( + amax * Float32(_HADAMARD_NORM / E4M3_MAX), Float32(1e-12) + ) + # Step 4. Thread 0 writes the block scale to its flat slot in the table. + if tidx == 0: + cutlass.inttoptr( + mScale.iterator.toint() + + _scale_slot(bh, nb, self.num_heads, self.batch, rows_out) * 4, + mem_space=1, + dtype=Float32, + ).store(scale) + # Step 5. Quantize the rotated rows with that scale and store the E4M3 + # block in output row order. + scale_inv = Float32(_HADAMARD_NORM) / scale + x8_base = _row_base( + mX8.iterator.toint(), b, rows_out, h, col0, 1, self.row_bytes8 + ) + for i in cutlass.range_constexpr(_ROWS_PER_THREAD): + pos = nb * VC_K_BLOCK_SIZE + rg + _ROW_GROUPS * i + if pos < rows_out: + cutlass.inttoptr( + x8_base + Int64(pos) * self.row_bytes8, mem_space=1, dtype=Int32 + ).store( + _pack_row_fp8(vals, i * _COLS_PER_THREAD, scale_inv), alignment=8 + ) + + @cute.kernel + def kernel( + self, + mQ: cute.Tensor, + mK: cute.Tensor, + mPerm: cute.Tensor | None, + mIndex: cute.Tensor | None, + mStats: cute.Tensor, + mQ8: cute.Tensor, + mK8: cute.Tensor, + mQScale: cute.Tensor, + mKScale: cute.Tensor, + ): + """One launch over the Q blocks, then the K tiles, then the repair tiles. + + Block ``nb`` below ``blocks_q`` quantizes Q block ``nb``. Above it, K + tile ``nb - blocks_q`` is read through the permutation, except the repair + tiles after the whole tiles, which gather the selected tokens through the + index. Every block runs ``_block`` with its own source and row map. + """ + co = _coords(self.num_heads) + nb = co[0] + smem = cutlass.utils.SmemAllocator() + red = smem.allocate_array(Float32, _THREADS // 32) + kb = nb - self.blocks_q + if nb < self.blocks_q: + self._block( + nb, + co, + mQ, + None, + None, + None, + mQ8, + mQScale, + red, + self.seq_len_q, + self.seq_len_q, + ) + elif cutlass.const_expr(mIndex is not None) and ( + (kb >= self.tail // VC_K_BLOCK_SIZE) + & (kb < self.tail // VC_K_BLOCK_SIZE + self.repair_tiles) + ): + self._block( + kb, + co, + mK, + None, + mIndex, + mStats, + mK8, + mKScale, + red, + self.rows, + self.seq_len_kv, + ) + else: + self._block( + kb, + co, + mK, + mPerm, + None, + mStats, + mK8, + mKScale, + red, + self.rows, + self.seq_len_kv, + ) + + @cute.jit + def __call__( + self, mQ, mK, mPerm, mIndex, mStats, mQ8, mK8, mQScale, mKScale, stream + ): + self.kernel(mQ, mK, mPerm, mIndex, mStats, mQ8, mK8, mQScale, mKScale).launch( + grid=[ + self.blocks_q + (self.rows + VC_K_BLOCK_SIZE - 1) // VC_K_BLOCK_SIZE, + self.batch * self.num_heads, + 1, + ], + block=[_THREADS, 1, 1], + smem=(_THREADS // 32) * 4, + stream=stream, + ) + + +class VcSelectPass(_Shape): + """The ``repair_tokens`` tokens of largest residual energy per sequence. + + A radix select over the four bytes of the non-negative fp32 bits with one + histogram per warp, then the indices above the threshold and the first + ones at it, in arbitrary order. + """ + + def __init__( + self, is_bf16, batch, seq_len_kv, num_heads, repair_tiles, repair_tokens: int + ): + super().__init__(is_bf16, batch, seq_len_kv, num_heads, repair_tiles) + self.repair_tokens = repair_tokens + + @cute.kernel + def kernel(self, mEnergy: cute.Tensor, mIndex: cute.Tensor): + """Select the ``repair_tokens`` largest energies of one batch head. + + Four radix passes over the bytes of the fp32 bit pattern, most significant + first. Each pass histograms the energies that match the prefix found so + far, scans the bins from the top and keeps the bin where the wanted count + falls, which fixes one more byte of the threshold. A final sweep writes + the indices above the threshold and as many at it as are still wanted. + """ + tidx, _, _ = cute.arch.thread_idx() + bh, _, _ = cute.arch.block_idx() + warp = tidx // 32 + lane = tidx % 32 + smem = cutlass.utils.SmemAllocator() + hist = smem.allocate_array(Int32, (_THREADS // 32) * _SELECT_BINS) + warp_tot = smem.allocate_array(Int32, _THREADS // 32) + # The threshold prefix, the count still wanted at it, the count stored + # above it and the count seen at it. + state = smem.allocate_array(Int32, 4) + e_base = mEnergy.iterator.toint() + Int64(bh) * self.seq_len_kv * 4 + out_base = mIndex.iterator.toint() + Int64(bh) * self.repair_tokens * 4 + # One CTA streams the whole row, so four loads per thread and + # iteration set the memory parallelism. + chunks = (self.seq_len_kv + 4 * _THREADS - 1) // (4 * _THREADS) + if tidx == 0: + state[0] = Int32(0) + state[1] = Int32(self.repair_tokens) + state[2] = Int32(0) + state[3] = Int32(0) + cute.arch.sync_threads() + # Step 1. One radix pass per byte. Per-warp histograms avoid contention + # on the atomics, the scan adds them up. + for p in cutlass.range_constexpr(4): + shift = 24 - 8 * p + for w in cutlass.range_constexpr(_THREADS // 32): + hist[w * _SELECT_BINS + tidx] = Int32(0) + prefix = state[0] + rem = state[1] + cute.arch.sync_threads() + for it in cutlass.range(chunks): + vals = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) + for u in cutlass.range_constexpr(4): + i = tidx + (it * 4 + u) * _THREADS + vals[u] = Int32(-1) + if i < self.seq_len_kv: + vals[u] = cutlass.inttoptr( + e_base + Int64(i) * 4, mem_space=1, dtype=Int32 + ).load() + for u in cutlass.range_constexpr(4): + bits = vals[u] + hit = bits >= 0 + if cutlass.const_expr(p > 0): + hit = hit & ((bits >> (shift + 8)) == prefix) + if hit: + cute.arch.atomic_add( + hist + + ( + warp * _SELECT_BINS + + ((bits >> shift) & (_SELECT_BINS - 1)) + ), + Int32(1), + ) + cute.arch.sync_threads() + # Thread tidx owns bin 255 - tidx and the scan runs in descending bin order. + mybin = _SELECT_BINS - 1 - tidx + cnt = Int32(0) + for w in cutlass.range_constexpr(_THREADS // 32): + cnt = cnt + hist[w * _SELECT_BINS + mybin] + incl = cnt + for sbit in cutlass.range_constexpr(5): + other = cute.arch.shuffle_sync_up(incl, 1 << sbit) + if lane >= (1 << sbit): + incl = incl + other + if lane == 31: + warp_tot[warp] = incl + cute.arch.sync_threads() + excl = incl - cnt + for w in cutlass.range_constexpr(_THREADS // 32): + if w < warp: + excl = excl + warp_tot[w] + if (excl < rem) & (rem <= excl + cnt): + state[0] = (prefix << 8) | mybin + state[1] = rem - excl + cute.arch.sync_threads() + # Step 2. Emit the indices. Energies above the threshold all qualify, + # energies equal to it fill the remaining slots in arrival order. + threshold = state[0] + rem = state[1] + n_above = self.repair_tokens - rem + for it in cutlass.range(chunks): + vals = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) + for u in cutlass.range_constexpr(4): + i = tidx + (it * 4 + u) * _THREADS + vals[u] = Int32(-1) + if i < self.seq_len_kv: + vals[u] = cutlass.inttoptr( + e_base + Int64(i) * 4, mem_space=1, dtype=Int32 + ).load() + for u in cutlass.range_constexpr(4): + i = tidx + (it * 4 + u) * _THREADS + bits = vals[u] + if bits > threshold: + slot = cute.arch.atomic_add(state + 2, Int32(1)) + cutlass.inttoptr( + out_base + Int64(slot) * 4, mem_space=1, dtype=Int32 + ).store(i) + elif bits == threshold: + slot = cute.arch.atomic_add(state + 3, Int32(1)) + if slot < rem: + cutlass.inttoptr( + out_base + Int64(n_above + slot) * 4, + mem_space=1, + dtype=Int32, + ).store(i) + + @cute.jit + def __call__(self, mEnergy, mIndex, stream): + self.kernel(mEnergy, mIndex).launch( + grid=[self.batch * self.num_heads, 1, 1], + block=[_THREADS, 1, 1], + smem=((_THREADS // 32) * _SELECT_BINS + _THREADS // 32 + 4) * 4, + stream=stream, + ) + + +class VcRepairRowsPass(_Shape): + """The repair tiles and the padding after the shifted tail. + + One CTA per repair tile writes the selected tokens' E4M3 value residuals + and, when ``mKOut`` is given, their keys, with zero rows for unused slots. + One more CTA zeroes the padding rows. + """ + + def __init__( + self, is_bf16, batch, seq_len_kv, num_heads, repair_tiles, repair_tokens: int + ): + super().__init__(is_bf16, batch, seq_len_kv, num_heads, repair_tiles) + self.repair_tokens = repair_tokens + + @cute.kernel + def kernel( + self, + mK: cute.Tensor | None, + mV: cute.Tensor, + mIndex: cute.Tensor, + mStats: cute.Tensor, + mKOut: cute.Tensor | None, + mV8: cute.Tensor, + ): + """Fill repair tile ``j`` of one batch head, or the padding after the tail. + + Row ``rr`` of repair tile ``j`` holds selected token ``j * 128 + rr``. Its + value row is the E4M3 residual left by pass 2's code, so the kernel adds + back what the first code missed. Its key row is a copy of the key. Slots + past the selected count and the padding rows are zero. + """ + j, bh, b, h, rg, col0 = _coords(self.num_heads) + v_base = _row_base( + mV.iterator.toint(), b, self.seq_len_kv, h, col0, 2, self.row_bytes + ) + v8_base = _row_base( + mV8.iterator.toint(), b, self.rows, h, col0, 1, self.row_bytes8 + ) + if cutlass.const_expr(mKOut is not None): + k_base = _row_base( + mK.iterator.toint(), b, self.seq_len_kv, h, col0, 2, self.row_bytes + ) + kout_base = _row_base( + mKOut.iterator.toint(), b, self.rows, h, col0, 2, self.row_bytes + ) + vmax = _load_row_f32(_stats_addr(mStats, bh, 1, col0)) + vs_inv = cutlass.Array( + Float32, _COLS_PER_THREAD, space=cutlass.AddressSpace.rmem + ) + for c in cutlass.range_constexpr(_COLS_PER_THREAD): + vs_inv[c] = Float32(1.0) / _v_scale(vmax[c]) + zeros = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) + for c in cutlass.range_constexpr(4): + zeros[c] = Int32(0) + raw = cutlass.Array(Int32, _WORDS, space=cutlass.AddressSpace.rmem) + for i in cutlass.range_constexpr(_ROWS_PER_THREAD): + rr = rg + _ROW_GROUPS * i + tok = Int32(-1) + out_row = self.seq_len_kv + self.shift + rr + if j < self.repair_tiles: + out_row = self.tail + j * VC_K_BLOCK_SIZE + rr + sel = j * VC_K_BLOCK_SIZE + rr + if sel < self.repair_tokens: + tok = cutlass.inttoptr( + mIndex.iterator.toint() + + (Int64(bh) * self.repair_tokens + sel) * 4, + mem_space=1, + dtype=Int32, + ).load() + if out_row < self.rows: + v_dst = v8_base + Int64(out_row) * self.row_bytes8 + if tok >= 0: + src_off = Int64(tok) * self.row_bytes + _load_raw(raw, 0, v_base + src_off) + row = _row(raw, 0, self.is_bf16) + # The residual of the E4M3 code pass 2 stored for this token. + res = cutlass.Array( + Float32, _COLS_PER_THREAD, space=cutlass.AddressSpace.rmem + ) + for c in cutlass.range_constexpr(_COLS_PER_THREAD): + res[c] = row[c] * vs_inv[c] + codes = ( + _pack_row_fp8(res, 0, Float32(1.0)) + .bitcast(cutlass.Float8E4M3FN) + .to(Float32) + ) + for c in cutlass.range_constexpr(_COLS_PER_THREAD): + res[c] = res[c] - codes[c] + cutlass.inttoptr(v_dst, mem_space=1, dtype=Int32).store( + _pack_row_fp8(res, 0, Float32(1.0)), alignment=8 + ) + if cutlass.const_expr(mKOut is not None): + _load_raw(raw, 0, k_base + src_off) + cutlass.inttoptr( + kout_base + Int64(out_row) * self.row_bytes, + mem_space=1, + dtype=Int32, + ).store( + raw.data_ptr().load(count=_WORDS, alignment=16), + alignment=16, + ) + else: + cutlass.inttoptr(v_dst, mem_space=1, dtype=Int32).store( + zeros.data_ptr().load(count=2, alignment=8), alignment=8 + ) + if cutlass.const_expr(mKOut is not None): + cutlass.inttoptr( + kout_base + Int64(out_row) * self.row_bytes, + mem_space=1, + dtype=Int32, + ).store( + zeros.data_ptr().load(count=4, alignment=16), alignment=16 + ) + + @cute.jit + def __call__(self, mK, mV, mIndex, mStats, mKOut, mV8, stream): + self.kernel(mK, mV, mIndex, mStats, mKOut, mV8).launch( + grid=[self.repair_tiles + 1, self.batch * self.num_heads, 1], + block=[_THREADS, 1, 1], + smem=0, + stream=stream, + ) + + +def _fake1d(dtype, align=32): + return cute.runtime.make_fake_compact_tensor( + dtype, (cute.sym_int(),), assumed_align=align + ) + + +def _in_dtype(dtype: torch.dtype): + if dtype == torch.bfloat16: + return cutlass.BFloat16, True + if dtype == torch.float16: + return cutlass.Float16, False + raise TypeError(f"VC fused preparation needs bf16/fp16 inputs, got {dtype}") + + +@functools.lru_cache(maxsize=None) +def _compiled( + dtype: torch.dtype, + demean: bool, + batch: int, + seq_len_kv: int, + num_heads: int, + repair_tiles: int, + seq_len_q: int, + repair_tokens: int, +): + """The kernels of one geometry. ``seq_len_q`` is 0 under QK16 and ``repair_tokens`` is 0 + without repair. Absent operands are ``None`` at compile time as at run time. + + Each kernel goes through the PrimTS on-disk CuTe DSL cache like the context + kernels, keyed by this geometry and the attention source tree, so a process + loads a kernel it or an earlier process compiled on this machine. + """ + cdt, is_bf16 = _in_dtype(dtype) + spec = ( + str(dtype), + demean, + batch, + seq_len_kv, + num_heads, + repair_tiles, + seq_len_q, + repair_tokens, + ) + + def build(kind, compile_fn): + return build_and_load_attention_kernel( + "vc_prepare", kind, spec, _COMPILE_OPTIONS, compile_fn + ) + + stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + f32 = lambda: _fake1d(Float32, 64) # noqa: E731 + i32 = lambda: _fake1d(Int32, 16) # noqa: E731 + e4m3 = lambda: _fake1d(cutlass.Float8E4M3FN, 16) # noqa: E731 + x16 = lambda: _fake1d(cdt) # noqa: E731 + qk8, repair = seq_len_q > 0, repair_tiles > 0 + shape = (is_bf16, batch, seq_len_kv, num_heads, repair_tiles) + perm = None if repair else i32() + kperm = None if qk8 else x16() + ksum = f32() if qk8 else None + k = { + "p1": build( + "p1", + lambda: cute.compile( + VcKvPass1(*shape, demean), + x16(), + x16(), + perm, + kperm, + f32(), + ksum, + f32(), + stream, + options=_COMPILE_OPTIONS, + ), + ), + "p2": build( + "p2", + lambda: cute.compile( + VcKvPass2(*shape), + x16(), + perm, + f32(), + ksum, + f32(), + f32(), + e4m3(), + None if repair else _fake1d(cutlass.BFloat16), + f32() if repair else None, + stream, + options=_COMPILE_OPTIONS, + ), + ), + } + if repair: + k["select"] = build( + "select", + lambda: cute.compile( + VcSelectPass(*shape, repair_tokens), + f32(), + i32(), + stream, + options=_COMPILE_OPTIONS, + ), + ) + k["rows"] = build( + "rows", + lambda: cute.compile( + VcRepairRowsPass(*shape, repair_tokens), + None if qk8 else x16(), + x16(), + i32(), + f32(), + None if qk8 else x16(), + e4m3(), + stream, + options=_COMPILE_OPTIONS, + ), + ) + if qk8: + k["qk"] = build( + "qk", + lambda: cute.compile( + VcQKPass(*shape, seq_len_q, repair_tokens), + x16(), + x16(), + perm, + i32() if repair else None, + f32(), + e4m3(), + e4m3(), + f32(), + f32(), + stream, + options=_COMPILE_OPTIONS, + ), + ) + return k + + +@torch.no_grad() +def vc_prepare( + k: torch.Tensor, v: torch.Tensor, perm: torch.Tensor, *, demean: bool = True +): + """VC-Attention-QK16 operands ``(k_perm, v8, v_scale, mean, mu)`` of + contiguous ``[B, S, H, D]`` K and V in ``perm`` order.""" + return _prepare(None, k, v, perm, demean=demean, repair_tiles=0, repair_tokens=0)[ + :5 + ] + + +@torch.no_grad() +def vc_prepare_fp8( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + perm: torch.Tensor, + *, + demean: bool = True, +): + """VC-Attention-QK8 operands ``(q8, k8, v8, q_scale, k_scale, v_scale, mean, + mu)`` of contiguous ``[B, S, H, D]`` Q, K and V, K and V in ``perm`` order.""" + if q.shape[-1] != VC_HEAD_DIM or not q.is_contiguous(): + raise ValueError( + f"VC fused preparation needs contiguous head_dim {VC_HEAD_DIM} Q" + ) + k8, v8, v_scale, mean, mu, q8, q_scale, k_scale = _prepare( + q, k, v, perm, demean=demean, repair_tiles=0, repair_tokens=0 + ) + return q8, k8, v8, q_scale, k_scale, v_scale, mean, mu + + +@torch.no_grad() +def vc_prepare_repair( + q: torch.Tensor | None, + k: torch.Tensor, + v: torch.Tensor, + *, + repair_tiles: int, + repair_tokens: int, +): + """V repair operands ``(q8, k_all, v_all, q_scale, k_scale, v_scale)`` in + input order plus ``repair_tiles`` tiles of the ``repair_tokens`` largest-residual + tokens. The QK8 entries are ``None`` under QK16.""" + if q is not None and (q.shape[-1] != VC_HEAD_DIM or not q.is_contiguous()): + raise ValueError( + f"VC fused preparation needs contiguous head_dim {VC_HEAD_DIM} Q" + ) + if not 0 < repair_tokens <= min(repair_tiles * VC_K_BLOCK_SIZE, k.shape[1]): + raise ValueError("repair_tokens must fit the repair tiles and the sequence") + k_all, v_all, v_scale, _, _, q8, q_scale, k_scale = _prepare( + q, + k, + v, + None, + demean=False, + repair_tiles=repair_tiles, + repair_tokens=repair_tokens, + ) + return q8, k_all, v_all, q_scale, k_scale, v_scale + + +def _prepare(q, k, v, perm, *, demean: bool, repair_tiles: int, repair_tokens: int): + """Run the passes and return ``(k_out, v8, v_scale, mean, mu, q8, q_scale, + k_scale)``. ``q`` selects QK8 and the QK8 outputs are ``None`` without it.""" + b, s_k, h, d = k.shape + if d != VC_HEAD_DIM: + raise ValueError(f"VC fused preparation supports head_dim {VC_HEAD_DIM} only") + if not (k.is_contiguous() and v.is_contiguous()): + raise ValueError("k and v must be contiguous") + dev = k.device + t = (s_k + VC_K_BLOCK_SIZE - 1) // VC_K_BLOCK_SIZE + repair = repair_tiles > 0 + rows = vc_repair_kv_len(s_k, repair_tiles) if repair else s_k + s_q = 0 if q is None else q.shape[1] + with torch.cuda.device(dev): + ker = _compiled( + k.dtype, bool(demean), b, s_k, h, repair_tiles, s_q, repair_tokens + ) + if not repair and ( + perm is None or tuple(perm.shape) != (b, h, s_k) or perm.device != dev + ): + raise ValueError(f"perm must be a [B, H, S_k] index tensor on {dev}") + perm32 = None if repair else perm.to(torch.int32).contiguous() + stats = torch.zeros((b, h, 2, d), dtype=torch.float32, device=dev) + mean = torch.empty((b, h, t, d), dtype=torch.float32, device=dev) + v8 = torch.empty((b, rows, h, d), dtype=torch.float8_e4m3fn, device=dev) + v_scale = torch.empty((b, h, d), dtype=torch.float32, device=dev) + k_perm = k_sum = mu = energy = index = None + if q is None: + k_perm = torch.empty((b, rows, h, d), dtype=k.dtype, device=dev) + else: + k_sum = torch.empty((b, h, t, d), dtype=torch.float32, device=dev) + if repair: + energy = torch.empty((b, h, s_k), dtype=torch.float32, device=dev) + index = torch.empty((b, h, repair_tokens), dtype=torch.int32, device=dev) + else: + groups = (t + VC_MEAN_GROUP_TILES - 1) // VC_MEAN_GROUP_TILES + mu = torch.empty( + (b, h, groups, *vc_mean_group_shape(d)), dtype=torch.bfloat16, device=dev + ) + flat = lambda x: None if x is None else x.view(-1) # noqa: E731 + ker["p1"]( + flat(k), + flat(v), + flat(perm32), + flat(k_perm), + flat(mean), + flat(k_sum), + flat(stats), + ) + ker["p2"]( + flat(v), + flat(perm32), + flat(mean), + flat(k_sum), + flat(stats), + flat(v_scale), + flat(v8), + flat(mu), + flat(energy), + ) + if repair: + ker["select"](flat(energy), flat(index)) + ker["rows"]( + flat(k) if q is None else None, + flat(v), + flat(index), + flat(stats), + flat(k_perm), + flat(v8), + ) + if q is None: + return k_perm, v8, v_scale, mean, mu, None, None, None + k8 = torch.empty((b, rows, h, d), dtype=torch.float8_e4m3fn, device=dev) + q8 = torch.empty((b, s_q, h, d), dtype=torch.float8_e4m3fn, device=dev) + k_scale = torch.empty( + (h, b, (rows + VC_K_BLOCK_SIZE - 1) // VC_K_BLOCK_SIZE), + dtype=torch.float32, + device=dev, + ) + q_scale = torch.empty( + (h, b, (s_q + VC_K_BLOCK_SIZE - 1) // VC_K_BLOCK_SIZE), + dtype=torch.float32, + device=dev, + ) + ker["qk"]( + flat(q), + flat(k), + flat(perm32), + flat(index), + flat(stats), + flat(q8), + flat(k8), + flat(q_scale), + flat(k_scale), + ) + return k8, v8, v_scale, mean, mu, q8, q_scale, k_scale diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/vc_attention/quantization.py b/tensorrt_llm/_torch/attention/backends/prims_ts/vc_attention/quantization.py new file mode 100644 index 000000000000..3d42d1b718c3 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/vc_attention/quantization.py @@ -0,0 +1,503 @@ +# 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. + +"""VC-Attention quantization of Q, K and V for the prims_ts context kernel. + +VC-Attention (arXiv 2609.15810) adds two changes on the value side of +attention: a token permutation that groups similar value rows into the same +128-token K/V tile, and per-tile value smoothing. Each V tile is stored as E4M3 +residuals around its tile mean and the kernel restores the mean inside the +online softmax, so the residual quantization error no longer carries the value +DC component. P is E4M3 through the ExpCast path in both recipes of the Q/K +side. VC-Attention-QK16 keeps Q and K in bf16. VC-Attention-QK8 centres K by +its per-head channel mean, rotates Q and K by the normalised 128-point Hadamard +matrix and quantizes them to E4M3 with one scale per 128-token Q block and per +128-token K tile, as the paper does. + +* ``vc_token_permutation`` runs an online k-means over value rows, one + permutation per batch and head, applied to K and V. Attention is invariant + to a common key and value permutation. +* ``vc_quantize`` permutes K, subtracts the V tile means, quantizes the + residuals per channel to E4M3 and packs the kernel operands with the CuTe + DSL kernels of ``prepare``. +* ``vc_quantize_fp8`` does the same plus the E4M3 Q/K operands and scales. +* ``vc_quantize_repair`` writes V repair rows instead of tile means, for + either recipe. +* ``VCAttentionPreprocessor`` applies the paper's V-Smooth schedule over the + denoising steps on top of the two. + +The operand contract the kernel consumes is defined in +``flashinfer.attention.prims_ts.vc_attention``. +""" + +from __future__ import annotations + +import functools +import math +from dataclasses import dataclass +from typing import Any + +import torch + +from . import ( + VC_HEAD_DIM, + VCAttentionParams, + vc_repair_tiles, +) + +# V-Smooth schedule. The paper groups and demeans on the first quarter of the +# denoising steps and keeps the permutation afterwards. The refresh cadence +# inside that window and the k-means size are not fixed by the paper. 64 +# clusters sort the values into 128-token tiles of similar rows while keeping +# the online k-means a small fraction of the attention call, and 3 iterations +# are enough because every refresh warm-starts from the previous centroids. +VC_SMOOTH_STEP_FRACTION = 0.25 +VC_PERM_REFRESH_EVERY = 4 +VC_KMEANS_CLUSTERS = 64 +VC_KMEANS_ITERS = 3 + + +@dataclass(frozen=True) +class VCAttentionOperands: + """Everything one VC-Attention forward needs, in kernel layout. + + ``k`` is the input in permuted token order; ``v`` is E4M3 ``[B, S, H, D]`` + in the same order holding the tile residuals ``(V[perm] - mean) / v_scale``. + Under VC-Attention-QK8 ``q`` and ``k`` are the rotated E4M3 operands and + ``q_scale`` and ``k_scale`` their per-block scales; under QK16 those + three are ``None`` and the caller's bf16 ``q`` is used as is. With + ``repair_tiles`` the token order is the input order, ``mean``, ``mu`` and + ``perm`` are ``None`` and the K/V rows are the whole 128-token tiles, the + repair tiles of key copies and E4M3 value residuals, then the zero-padded + partial last tile. + """ + + k: torch.Tensor + v: torch.Tensor + v_scale: ( + torch.Tensor + ) # [B, H, D] fp32 per-channel E4M3 residual scale (feeds output_scale) + mean: ( + torch.Tensor | None + ) # [B, H, num_kv_tiles, D] fp32 tile means (permuted order) + mu: ( + torch.Tensor | None + ) # [B, H, ceil(num_kv_tiles / 16), 16, 256] bf16 packed kernel operands + perm: torch.Tensor | None # [B, H, S_k] int32/int64 token permutation + demean: bool = True # whether the means were subtracted (else zero) + repair_tiles: int = 0 + seq_len_kv: int | None = None # K/V tokens before the repair tiles and padding + q: torch.Tensor | None = None # [B, S_q, H, D] E4M3 rotated Q (QK8) + q_scale: torch.Tensor | None = None # [H, B, ceil(S_q / 128)] fp32 + k_scale: torch.Tensor | None = None # [H, B, ceil(S_k / 128)] fp32 + + @property + def params(self) -> "VCAttentionParams": + """The per-run operands ``BatchPrefillTSWrapper.run`` takes as ``vc``.""" + return VCAttentionParams( + q_scale=self.q_scale, + k_scale=self.k_scale, + v_scale=self.v_scale, + tile_means=self.mu, + ) + + +@functools.lru_cache(maxsize=None) +def _hadamard_matrix_cached(d: int, device: str) -> torch.Tensor: + h = torch.ones((1, 1), dtype=torch.float32) + while h.shape[0] < d: + h = torch.cat([torch.cat([h, h], 1), torch.cat([h, -h], 1)], 0) + return (h / math.sqrt(d)).to(device) + + +def hadamard_matrix(d: int, device: torch.device) -> torch.Tensor: + """Normalised Sylvester Hadamard matrix ``[d, d]`` (``d`` a power of two).""" + if d & (d - 1): + raise ValueError("head_dim must be a power of two for the Hadamard rotation") + return _hadamard_matrix_cached(d, str(device)) + + +@torch.no_grad() +def vc_token_permutation( + v: torch.Tensor, + *, + num_clusters: int | None = None, + iters: int = VC_KMEANS_ITERS, + generator: torch.Generator | None = None, + init_centroids: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Group similar value rows by k-means per batch and head, then argsort the labels. + + Returns the int64 permutation ``[B, H, S]`` and the ``[B*H, k, D]`` centroids. + ``V[b, perm[b, h], h]`` lists the tokens cluster by cluster, so 128-token K/V + tiles hold similar values and their tile means remove most of the value + energy. ``init_centroids`` from a previous call warm-starts the clustering, + as the paper does when the grouping is refreshed a few denoising steps later. + """ + b, s, h, d = v.shape + if num_clusters is None: + num_clusters = VC_KMEANS_CLUSTERS + # Seeding draws distinct value rows, so a short sequence caps the clusters. + num_clusters = min(num_clusters, s) + n_groups = b * h + # Work on a few heads at a time so the fp32 copy of V and the distance + # matrix each stay under 0.75 GB. + per_group = 4 * s * max(d, num_clusters) + chunk_h = max(1, min(h, int(0.75e9 // per_group))) + perm_out = torch.empty((b, h, s), dtype=torch.int64, device=v.device) + cent_out = torch.empty( + (n_groups, num_clusters, d), dtype=torch.float32, device=v.device + ) + warm = init_centroids is not None and tuple(init_centroids.shape) == ( + n_groups, + num_clusters, + d, + ) + for bi in range(b): + for h0 in range(0, h, chunk_h): + h1 = min(h, h0 + chunk_h) + gsz = h1 - h0 + g0 = bi * h + h0 + x = v[bi, :, h0:h1, :].permute(1, 0, 2).float() # [G, S, D] + if warm: + centroids = init_centroids[g0 : g0 + gsz].to(x.dtype) + else: + idx = torch.stack( + [ + torch.randperm(s, device=v.device, generator=generator)[ + :num_clusters + ] + for _ in range(gsz) + ] + ) # [G, k] + centroids = torch.gather( + x, 1, idx.unsqueeze(-1).expand(-1, -1, d) + ) # [G, k, D] + x_sq = (x * x).sum(-1, keepdim=True) # [G, S, 1] + labels = None + for _ in range(iters): + c_sq = (centroids * centroids).sum(-1).unsqueeze(1) # [G, 1, k] + dist = x_sq + c_sq - 2.0 * torch.bmm(x, centroids.transpose(1, 2)) + labels = dist.argmin(dim=-1) # [G, S] + del dist + counts = torch.zeros( + gsz, num_clusters, device=v.device, dtype=torch.float32 + ) + counts.scatter_add_( + 1, labels, torch.ones_like(labels, dtype=torch.float32) + ) + sums = torch.zeros_like(centroids) + sums.scatter_add_(1, labels.unsqueeze(-1).expand(-1, -1, d), x) + nonempty = counts > 0 + centroids = torch.where( + nonempty.unsqueeze(-1), + sums / counts.clamp_min(1.0).unsqueeze(-1), + centroids, + ) + assert labels is not None + perm_out[bi, h0:h1] = torch.argsort(labels, dim=-1, stable=True) + cent_out[g0 : g0 + gsz] = centroids + del x, x_sq + return perm_out, cent_out + + +# --------------------------------------------------------------------------- +# Preprocessing (CuTe DSL kernels in prepare.py) +# --------------------------------------------------------------------------- +def _check_prepare_inputs(x: torch.Tensor) -> None: + if ( + x.shape[-1] != VC_HEAD_DIM + or x.dtype not in (torch.bfloat16, torch.float16) + or not x.is_cuda + ): + raise ValueError( + f"VC-Attention preparation needs CUDA bf16/fp16 inputs with head_dim {VC_HEAD_DIM}, " + f"got {x.dtype} {tuple(x.shape)} on {x.device}" + ) + + +@torch.no_grad() +def vc_quantize( + k: torch.Tensor, + v: torch.Tensor, + *, + perm: torch.Tensor | None = None, + kmeans_iters: int = VC_KMEANS_ITERS, + generator: torch.Generator | None = None, + demean: bool = True, +) -> VCAttentionOperands: + """Turn ``[B, S, H, D]`` bf16/fp16 K/V into VC-Attention-QK16 kernel + operands with two CuTe DSL kernels (``prepare``). + + Keys are permuted with the values and kept in their dtype. Values are + permuted, split into 128-token tile means and residuals (``demean=False`` + keeps the means at zero), and the residuals are quantized to E4M3 with one + scale per (batch, head, channel), returned as ``v_scale`` for the run's + ``vc`` operands. Pass 1 gathers each permuted tile once, writes the + permuted K rows, the V tile mean and residual amax; pass 2 writes the E4M3 + residuals and the packed bf16 mean operand. Requires CUDA bf16/fp16 inputs + with head_dim 128. + """ + if k.dim() != 4 or k.shape != v.shape or k.dtype != v.dtype: + raise ValueError("k and v must be [B, S, H, D] with matching shapes and dtype") + _check_prepare_inputs(k) + from .prepare import vc_prepare + + if perm is None: + perm = vc_token_permutation(v, iters=kmeans_iters, generator=generator)[0] + k_p, v8, v_scale, mean, mu = vc_prepare( + k.contiguous(), + v.contiguous(), + perm, + demean=demean, + ) + return VCAttentionOperands( + k=k_p, + v=v8, + v_scale=v_scale, + mean=mean, + mu=mu if demean else None, + perm=perm, + demean=demean, + ) + + +@torch.no_grad() +def vc_quantize_fp8( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + perm: torch.Tensor | None = None, + kmeans_iters: int = VC_KMEANS_ITERS, + generator: torch.Generator | None = None, + demean: bool = True, +) -> VCAttentionOperands: + """Turn ``[B, S, H, D]`` bf16/fp16 Q/K/V into VC-Attention-QK8 kernel + operands with four CuTe DSL kernels (``prepare``). + + Keys are centred by their per-(batch, head) channel mean over tokens + (softmax-invariant), permuted with the values, rotated with Q by the + normalised Hadamard matrix and quantized to E4M3 with one scale per + 128-token tile; queries take one E4M3 scale per 128-token block. Values + are prepared as in ``vc_quantize``. Requires CUDA bf16/fp16 inputs with + head_dim 128. + """ + if ( + q.dim() != 4 + or k.shape != v.shape + or q.shape[0] != k.shape[0] + or q.shape[2:] != k.shape[2:] + or q.dtype != k.dtype + or k.dtype != v.dtype + ): + raise ValueError( + "q, k, v must be [B, S, H, D] with matching batch, heads, head_dim and dtype" + ) + _check_prepare_inputs(q) + from .prepare import vc_prepare_fp8 + + if perm is None: + perm = vc_token_permutation(v, iters=kmeans_iters, generator=generator)[0] + q8, k8, v8, q_scale, k_scale, v_scale, mean, mu = vc_prepare_fp8( + q.contiguous(), k.contiguous(), v.contiguous(), perm, demean=demean + ) + return VCAttentionOperands( + k=k8, + v=v8, + v_scale=v_scale, + mean=mean, + mu=mu if demean else None, + perm=perm, + demean=demean, + q=q8, + q_scale=q_scale, + k_scale=k_scale, + ) + + +@torch.no_grad() +def vc_quantize_repair( + k: torch.Tensor, + v: torch.Tensor, + *, + budget: float, + q: torch.Tensor | None = None, +) -> VCAttentionOperands: + """Turn ``[B, S, H, D]`` K/V into V repair operands, of VC-Attention-QK8 + with ``q``. + + Values are quantized to E4M3 in input order with one scale per batch, + head and channel. The ``budget`` fraction of tokens with the largest + residual energy per batch and head gets a repair row, a copy of the key + (under QK8 re-centred, rotated and quantized with the repair tile's own + scale) with the E4M3 residual of the value. + """ + if k.dim() != 4 or k.shape != v.shape or k.dtype != v.dtype: + raise ValueError("k and v must be [B, S, H, D] with matching shapes and dtype") + if q is not None and ( + q.dim() != 4 + or q.shape[0] != k.shape[0] + or q.shape[2:] != k.shape[2:] + or q.dtype != k.dtype + ): + raise ValueError( + "q must be [B, S, H, D] with the batch, heads, head_dim and dtype of k" + ) + _check_prepare_inputs(k) + from .prepare import vc_prepare_repair + + s_k = k.shape[1] + repair_tiles = vc_repair_tiles(s_k, budget) + repair_tokens = max(1, round(budget * s_k)) + q8, k_all, v_all, q_scale, k_scale, v_scale = vc_prepare_repair( + None if q is None else q.contiguous(), + k.contiguous(), + v.contiguous(), + repair_tiles=repair_tiles, + repair_tokens=repair_tokens, + ) + return VCAttentionOperands( + k=k_all, + v=v_all, + v_scale=v_scale, + mean=None, + mu=None, + perm=None, + demean=False, + repair_tiles=repair_tiles, + seq_len_kv=s_k, + q=q8, + q_scale=q_scale, + k_scale=k_scale, + ) + + +class VCAttentionPreprocessor: + """Caller-owned preparation of bf16 K/V with the paper's V-Smooth schedule. + + Value grouping and demeaning run on the first ``smooth_step_fraction`` of + the denoising steps. Inside that window the token permutation is + recomputed every ``perm_refresh_every`` steps (k-means warm-started from + the previous centroids) and kept afterwards, when the plain per-channel + E4M3 V runs with zero tile means. ``kmeans_clusters`` (``None`` = 64) and + ``kmeans_iters`` parametrize the online k-means. Without a denoise step + the grouping is computed once per K/V geometry and kept. With ``q`` the + preparation is the VC-Attention-QK8 one (one E4M3 Q scale per 128 tokens). + """ + + def __init__( + self, + *, + smooth_step_fraction: float = VC_SMOOTH_STEP_FRACTION, + perm_refresh_every: int = VC_PERM_REFRESH_EVERY, + kmeans_clusters: int | None = None, + kmeans_iters: int = VC_KMEANS_ITERS, + ) -> None: + if not (0.0 <= smooth_step_fraction <= 1.0): + raise ValueError("smooth_step_fraction must be in [0, 1]") + if perm_refresh_every < 1: + raise ValueError("perm_refresh_every must be at least 1") + if kmeans_clusters is not None and kmeans_clusters < 1: + raise ValueError("kmeans_clusters must be at least 1") + if kmeans_iters < 1: + raise ValueError("kmeans_iters must be at least 1") + self.smooth_step_fraction = smooth_step_fraction + self.perm_refresh_every = perm_refresh_every + self.kmeans_clusters = kmeans_clusters + self.kmeans_iters = kmeans_iters + # Per (batch, seq_len_k, num_kv_heads): perm, centroids, step of the last refresh. + self._groupings: dict[tuple[int, int, int], dict[str, Any]] = {} + + @torch.no_grad() + @torch.compiler.disable + def prepare( + self, + k: torch.Tensor, + v: torch.Tensor, + *, + q: torch.Tensor | None = None, + denoise_step: tuple[int, int] | None = None, + ) -> VCAttentionOperands: + """Return the operands of one run for ``[B, S, H, D]`` bf16 K/V, and Q + under VC-Attention-QK8. + + ``denoise_step`` is the sampler's ``(step_index, num_steps)``. Kept + out of torch.compile so the step index and the k-means refresh do not + recompile the caller per step. + """ + key = (k.shape[0], k.shape[1], k.shape[2]) + entry = self._groupings.get(key) + smooth = True + refresh = entry is None + if denoise_step is not None: + step_idx, num_steps = denoise_step + smooth = step_idx < math.ceil(self.smooth_step_fraction * num_steps) + # Step 0 starts a new generation, so it always regroups. + refresh = smooth and ( + entry is None + or step_idx == 0 + or ( + step_idx % self.perm_refresh_every == 0 + and entry["step"] != step_idx + ) + ) + if smooth: + if refresh: + perm, centroids = vc_token_permutation( + v, + num_clusters=self.kmeans_clusters, + iters=self.kmeans_iters, + init_centroids=None if entry is None else entry["centroids"], + ) + step = None if denoise_step is None else denoise_step[0] + if entry is None: + entry = { + "perm": perm.to(torch.int32), + "centroids": centroids, + "step": step, + } + self._groupings[key] = entry + else: + # The buffers keep their addresses, so a CUDA graph captured around one call + # reads the grouping that later calls or replays write into them. + entry["perm"].copy_(perm.to(torch.int32)) + entry["centroids"].copy_(centroids) + entry["step"] = step + perm = entry["perm"] + elif entry is not None: + # V-Smooth off (paper Section 3.4). The plain kernel runs on the last + # permutation with zero tile means. + perm = entry["perm"] + else: + perm = torch.arange(k.shape[1], device=k.device, dtype=torch.int32) + perm = perm.view(1, 1, -1).expand(k.shape[0], k.shape[2], -1).contiguous() + if q is not None: + return vc_quantize_fp8(q, k, v, perm=perm, demean=smooth) + return vc_quantize(k, v, perm=perm, demean=smooth) + + +__all__ = [ + "VC_KMEANS_CLUSTERS", + "VC_KMEANS_ITERS", + "VC_PERM_REFRESH_EVERY", + "VC_SMOOTH_STEP_FRACTION", + "VCAttentionOperands", + "VCAttentionPreprocessor", + "hadamard_matrix", + "vc_quantize", + "vc_quantize_fp8", + "vc_quantize_repair", + "vc_token_permutation", +] diff --git a/tests/unittest/_torch/attention/kernels/test_sage_quant.py b/tests/unittest/_torch/attention/kernels/test_sage_quant.py index 7db0a5da2473..f875d7f7c016 100644 --- a/tests/unittest/_torch/attention/kernels/test_sage_quant.py +++ b/tests/unittest/_torch/attention/kernels/test_sage_quant.py @@ -21,7 +21,10 @@ 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.backends.prims_ts.sage import ( + sage_flat_scale_numel, + sage_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") @@ -97,7 +100,7 @@ def _assert_summaries_match_reference( 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( + slots = sage_flat_scale_slot( torch.arange(batch, device=k.device)[:, None], torch.arange(blocks, device=k.device)[None, :], blocks, @@ -106,7 +109,7 @@ def _assert_summaries_match_reference( 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)) + assert actual_scales.shape == (heads, sage_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 ) @@ -223,7 +226,7 @@ def _check_block_quantized( """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 + expected_numel = sage_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): @@ -235,7 +238,9 @@ def _check_block_quantized( 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) + slots = begin // block_size + sage_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) diff --git a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py index 5beda56866d2..c7c4cf8dc833 100644 --- a/tests/unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py +++ b/tests/unittest/_torch/attention/sparse/deepseek_v4/test_prims_ts_sparse_mla_fmha.py @@ -405,7 +405,9 @@ def test_prepared_launch_matches_the_full_path_per_layer(): sinks_a = torch.randn(heads, dtype=torch.float32, device=device) sinks_b = torch.randn(heads, dtype=torch.float32, device=device) wrapper = mla_decode.BatchSparseMLADecodePagedTSWrapper() - wrapper.plan(device, rows, heads, max_topk=topk, max_extra_topk=extra_topk, has_sinks=True) + wrapper.plan( + device, rows, heads, max_topk=topk, max_extra_topk=extra_topk, use_attention_sinks=True + ) prepared_launches = [] original = wrapper._run_prepared @@ -427,7 +429,7 @@ def launch(sinks, kv=pool, validate=False, idx=0): extra_lengths=sets[idx][3], softmax_scale=LATENT**-0.5, output_scale=1.0, - sinks=sinks, + attention_sinks=sinks, out=out, validate=validate, ) @@ -453,7 +455,9 @@ def launch(sinks, kv=pool, validate=False, idx=0): launch(sinks_a.clone()) assert len(prepared_launches) == 3 and len(wrapper._prepared) == 4 # Re-planning drops every prepared launch. - wrapper.plan(device, rows, heads, max_topk=topk, max_extra_topk=extra_topk, has_sinks=True) + wrapper.plan( + device, rows, heads, max_topk=topk, max_extra_topk=extra_topk, use_attention_sinks=True + ) assert not wrapper._prepared @@ -649,7 +653,7 @@ def make_pool(n): # on the softmax, the KV descale on the output), as the adapter passes them. run_scales = dict(softmax_scale=softmax_scale * descale * descale, output_scale=descale) sources = dict(extra_indices=extra_indices, extra_lengths=extra_lengths) if extra_topk else {} - plan_args = dict(max_topk=topk, max_extra_topk=extra_topk, max_seq_len_q=1, q_data_type=dtype) + plan_args = dict(max_topk=topk, max_extra_topk=extra_topk, max_q_len=1, q_dtype=dtype) reference = BatchSparseMLADecodePagedTSWrapper() reference.plan(device, rows, heads, **plan_args) @@ -801,7 +805,7 @@ def test_speculative_decoding_rows_run_on_prims_ts(monkeypatch, generation_seq_l library default and the spy asserts that PrimTS served those steps with one row per query token (``rows = sum(generation_lens)``, not one per request), which is the per-token-row contract the kernel is planned for - (``max_seq_len_q = 1``). + (``max_q_len = 1``). """ _skip_unless_toolchain() monkeypatch.delenv("TLLM_FMHA_LIBS", raising=False) @@ -1192,7 +1196,7 @@ def probing_plan(self, device, rows, heads, **kwargs): payloads, workers = fanned[0] assert workers == 2 # capped by the env value, below the three variants assert [p["rows"] for p in payloads] == [20, 15, 7] - assert payloads[0]["plan_kwargs"]["q_data_type"] == "bfloat16" # JSON-safe dtype + assert payloads[0]["plan_kwargs"]["q_dtype"] == "bfloat16" # JSON-safe dtype assert payloads[0]["plan_kwargs"]["max_extra_topk"] == 32 assert payloads[0]["device_index"] == 0 import json @@ -1301,7 +1305,7 @@ def recording_plan(self, device, rows, heads, **kwargs): monkeypatch.setattr(cc, "expected_cache_meta", lambda: {"probe": "meta"}) payloads = [ pp.encode_payload( - 0, rows, 16, dict(max_topk=128, max_extra_topk=512, q_data_type=torch.bfloat16) + 0, rows, 16, dict(max_topk=128, max_extra_topk=512, q_dtype=torch.bfloat16) ) for rows in (20, 15, 7, 3) ] diff --git a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py index d0e61dfb5e54..d3c7b5ed8696 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py +++ b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py @@ -658,7 +658,7 @@ def compile_with_record(*args): wrapper.plan( device=device, batch_size=batch_size, - max_seq_len_q=3, + max_q_len=3, max_kv_len=max_seq_len, num_qo_heads=num_qo_heads, num_kv_heads=num_kv_heads, @@ -676,12 +676,12 @@ def compile_with_record(*args): assert len(compile_records) == 3 assert [ - (args[0].max_seq_len_q, args[0].max_kv_len, policy["scheduler"]) + (args[0].max_q_len, args[0].max_kv_len, args[0].runtime_lengths, policy["scheduler"]) for args, policy, _compiled in compile_records ] == [ - (96, 96, "clc_dynamic_persistent"), - (128, 128, "nonpersistent"), - (3, 64, "clc_dynamic_persistent"), + (None, None, True, "clc_dynamic_persistent"), + (None, None, True, "nonpersistent"), + (None, None, True, "clc_dynamic_persistent"), ] assert compile_records[0][0] != compile_records[2][0] assert len({id(recorded) for _args, _policy, recorded in compile_records}) == 3 @@ -801,11 +801,11 @@ def test_prims_ts_decode_live_wrapper_cuda_graph_replay() -> None: head_dim, page_size, max_seq_len, - max_seq_len_q=1, + max_q_len=1, packed_query=False, - q_data_type=dtype, - kv_data_type=dtype, - o_data_type=dtype, + q_dtype=dtype, + kv_dtype=dtype, + out_dtype=dtype, mask_type="causal", workspace_buffer=external_workspace, ) @@ -916,7 +916,7 @@ def test_prims_ts_mla_live_wrapper_cuda_graph_replay(max_seq_len: int) -> None: qk_rope_head_dim, page_size, max_seq_len, - max_seq_len_q=1, + max_q_len=1, q_dtype=dtype, kv_dtype=dtype, out_dtype=dtype, @@ -937,11 +937,11 @@ def test_prims_ts_mla_live_wrapper_cuda_graph_replay(max_seq_len: int) -> None: qk_rope_head_dim, page_size, max_seq_len, - max_seq_len_q=1, + max_q_len=1, packed_query=False, - q_data_type=dtype, - kv_data_type=dtype, - o_data_type=dtype, + q_dtype=dtype, + kv_dtype=dtype, + out_dtype=dtype, mask_type="causal", workspace_buffer=external_workspace, ) @@ -957,7 +957,7 @@ def test_prims_ts_mla_live_wrapper_cuda_graph_replay(max_seq_len: int) -> None: kv_cache, block_tables, seq_lens, - bmm1_scale=bmm1_scale, + softmax_scale=bmm1_scale, out=output, ) graph = torch.cuda.CUDAGraph() @@ -967,7 +967,7 @@ def test_prims_ts_mla_live_wrapper_cuda_graph_replay(max_seq_len: int) -> None: kv_cache, block_tables, seq_lens, - bmm1_scale=bmm1_scale, + softmax_scale=bmm1_scale, out=output, validate=False, ) @@ -985,8 +985,8 @@ def test_prims_ts_mla_live_wrapper_cuda_graph_replay(max_seq_len: int) -> None: kv_lora_rank=kv_lora_rank, qk_rope_head_dim=qk_rope_head_dim, max_kv_len=max_seq_len, - max_seq_len_q=1, - bmm1_scale=bmm1_scale, + max_q_len=1, + softmax_scale=bmm1_scale, out_dtype=dtype, out=torch.empty_like(output), mask_type="causal", @@ -1068,11 +1068,11 @@ def plan(batch_size: int, max_seq_len: int) -> None: head_dim, page_size, max_seq_len, - max_seq_len_q=1, + max_q_len=1, packed_query=False, - q_data_type=dtype, - kv_data_type=dtype, - o_data_type=dtype, + q_dtype=dtype, + kv_dtype=dtype, + out_dtype=dtype, mask_type="causal", workspace_buffer=shared_workspace, ) @@ -1227,11 +1227,11 @@ def test_prims_ts_decode_wrappers_share_workspace_across_serialized_layers() -> head_dim, page_size, max_seq_len, - max_seq_len_q=1, + max_q_len=1, packed_query=False, - q_data_type=dtype, - kv_data_type=dtype, - o_data_type=dtype, + q_dtype=dtype, + kv_dtype=dtype, + out_dtype=dtype, mask_type="causal", workspace_buffer=shared_workspace, ) diff --git a/tests/unittest/_torch/attention/test_prims_ts_context_cuda_scales.py b/tests/unittest/_torch/attention/test_prims_ts_context_cuda_scales.py new file mode 100644 index 000000000000..d0aa002b8991 --- /dev/null +++ b/tests/unittest/_torch/attention/test_prims_ts_context_cuda_scales.py @@ -0,0 +1,71 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Device quantization scales remain live across PrimTS context graph replay.""" + +import math + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.prims_ts.context import BatchPrefillTSWrapper + + +@pytest.mark.parametrize( + "qk_fp8,override_output_scale", [(False, True), (True, False), (True, True)] +) +def test_context_cuda_scales_graph_replay(qk_fp8: bool, override_output_scale: bool) -> None: + if not torch.cuda.is_available() or torch.cuda.get_device_capability() not in ( + (10, 0), + (10, 3), + ): + pytest.skip("PrimTS context validation requires SM100 or SM103") + torch.manual_seed(20261030) + q_dtype = torch.float8_e4m3fn if qk_fp8 else torch.bfloat16 + q = torch.randn(1, 32, 4, 128, device="cuda").to(q_dtype) + k = torch.randn(1, 64, 4, 128, device="cuda").to(q_dtype) + v = torch.randn(1, 64, 4, 128, device="cuda").to(torch.float8_e4m3fn) + out = torch.empty_like(q, dtype=torch.bfloat16) + softmax_scale = torch.tensor([math.log2(math.e) / math.sqrt(128)], device="cuda") + output_scale = torch.ones(1, device="cuda") + wrapper = BatchPrefillTSWrapper() + wrapper.plan( + device=q.device, + batch_size=1, + max_q_len=32, + max_kv_len=64, + num_qo_heads=4, + num_kv_heads=4, + head_dim=128, + q_dtype=q.dtype, + k_dtype=k.dtype, + v_dtype=v.dtype, + out_dtype=out.dtype, + cuda_scales=True, + ) + + def run() -> None: + wrapper.run( + q, + k, + v, + scale_softmax_log2=softmax_scale if qk_fp8 else None, + output_scale=output_scale if override_output_scale else None, + out=out, + validate=False, + ) + + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + for qk_scale, v_scale in ((0.5, 2.0), (1.5, 0.25)): + softmax_scale.fill_(qk_scale * math.log2(math.e) / math.sqrt(128)) + output_scale.fill_(v_scale) + graph.replay() + expected = torch.nn.functional.scaled_dot_product_attention( + q.float().transpose(1, 2), + k.float().transpose(1, 2), + v.float().transpose(1, 2), + scale=(qk_scale if qk_fp8 else 1.0) / math.sqrt(128), + ).transpose(1, 2) * (v_scale if override_output_scale else 1.0) + torch.testing.assert_close(out.float(), expected, rtol=0.04, atol=0.04) diff --git a/tests/unittest/_torch/attention/test_prims_ts_fmha.py b/tests/unittest/_torch/attention/test_prims_ts_fmha.py index c1d1fe078709..68801f0bbb40 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_fmha.py +++ b/tests/unittest/_torch/attention/test_prims_ts_fmha.py @@ -1032,7 +1032,7 @@ def test_context_wrapper_plans_once_and_reads_live_fixed_metadata( assert plan_kwargs == { "device": q_processed.device, "batch_size": 2, - "max_seq_len_q": 8, + "max_q_len": 8, "max_kv_len": 128, "num_qo_heads": attn.num_heads, "num_kv_heads": attn.num_kv_heads, @@ -1044,7 +1044,7 @@ def test_context_wrapper_plans_once_and_reads_live_fixed_metadata( "page_size": 32, "mask_type": "causal", "window_left": -1, - "sm_scale": pytest.approx(1.0 / math.sqrt(attn.head_dim)), + "softmax_scale": pytest.approx(1.0 / math.sqrt(attn.head_dim)), "output_scale": 1.0, "paged_v_tail_is_zero": True, } @@ -1250,11 +1250,11 @@ def test_generation_wrapper_plans_once_and_reads_live_fixed_metadata( assert decode_workspace.data_ptr() == params.workspace.data_ptr() assert decode_workspace.numel() == 64 assert {key: value for key, value in plan_kwargs.items() if key != "workspace_buffer"} == { - "max_seq_len_q": 1, + "max_q_len": 1, "packed_query": False, - "q_data_type": torch.bfloat16, - "kv_data_type": torch.bfloat16, - "o_data_type": torch.bfloat16, + "q_dtype": torch.bfloat16, + "kv_dtype": torch.bfloat16, + "out_dtype": torch.bfloat16, "mask_type": "causal", "window_left": 15, } @@ -1276,8 +1276,8 @@ def test_generation_wrapper_plans_once_and_reads_live_fixed_metadata( fixed_block_tables, torch.tensor([[0, 1, 2, 3], [2, 3, 4, 5]], dtype=torch.int32), ) - assert run_kwargs["bmm1_scale"] == pytest.approx(1.0 / math.sqrt(attn.head_dim)) - assert run_kwargs["bmm2_scale"] == 1.0 + assert run_kwargs["softmax_scale"] == pytest.approx(1.0 / math.sqrt(attn.head_dim)) + assert run_kwargs["output_scale"] == 1.0 assert run_kwargs["out"].shape == (2, attn.num_heads, attn.head_dim) assert run_kwargs["out"].data_ptr() == output.data_ptr() assert run_kwargs["validate"] is False @@ -1389,12 +1389,12 @@ def get_wrapper(batch_size: int) -> object: k_cache, v_cache, batch_size=batch_size, - max_seq_len_q=128, + max_q_len=128, max_kv_len=256, page_size=32, mask_type="causal", window_left=-1, - sm_scale=1.0 / math.sqrt(128), + softmax_scale=1.0 / math.sqrt(128), output_dtype=torch.bfloat16, ) @@ -1486,7 +1486,7 @@ def _get_test_mla_wrapper( kv_lora_rank=512, qk_rope_head_dim=64, page_size=32, - max_seq_len_q=1, + max_q_len=1, max_kv_len=96, q_dtype=torch.bfloat16, kv_dtype=torch.bfloat16, @@ -1592,11 +1592,11 @@ def test_mla_eager_wrapper_plans_once_and_reads_live_fixed_metadata( assert workspace_buffer.data_ptr() == params.workspace.data_ptr() assert workspace_buffer.numel() == 64 assert {key: value for key, value in plan_kwargs.items() if key != "workspace_buffer"} == { - "max_seq_len_q": 1, + "max_q_len": 1, "packed_query": False, - "q_data_type": kernel_dtype, - "kv_data_type": kernel_dtype, - "o_data_type": torch.bfloat16, + "q_dtype": kernel_dtype, + "kv_dtype": kernel_dtype, + "out_dtype": torch.bfloat16, "mask_type": "causal", } wrapper.run.assert_called_once() @@ -1614,12 +1614,12 @@ def test_mla_eager_wrapper_plans_once_and_reads_live_fixed_metadata( run_kwargs["block_tables"], torch.tensor([[0, 1, 2], [3, 4, 5]], dtype=torch.int32), ) - assert run_kwargs["seq_lens"].data_ptr() == sequence_lengths.data_ptr() - torch.testing.assert_close(run_kwargs["seq_lens"], sequence_lengths) + assert run_kwargs["seq_lens_kv"].data_ptr() == sequence_lengths.data_ptr() + torch.testing.assert_close(run_kwargs["seq_lens_kv"], sequence_lengths) assert run_kwargs["out"].shape == (2, 1, attn.num_heads, 512) assert run_kwargs["out"].data_ptr() == output.data_ptr() - assert run_kwargs["bmm1_scale"] == pytest.approx(0.125 if fp8 else 1.0 / math.sqrt(128 + 64)) - assert run_kwargs["bmm2_scale"] == (2.0 if fp8 else 1.0) + assert run_kwargs["softmax_scale"] == pytest.approx(0.125 if fp8 else 1.0 / math.sqrt(128 + 64)) + assert run_kwargs["output_scale"] == (2.0 if fp8 else 1.0) assert run_kwargs["validate"] is False block_tables[:, 0].add_(20) @@ -1636,7 +1636,7 @@ def test_mla_eager_wrapper_plans_once_and_reads_live_fixed_metadata( run_kwargs["block_tables"], torch.tensor([[20, 21, 22], [23, 24, 25]], dtype=torch.int32), ) - torch.testing.assert_close(run_kwargs["seq_lens"], torch.tensor([34, 65], dtype=torch.int32)) + torch.testing.assert_close(run_kwargs["seq_lens_kv"], torch.tensor([34, 65], dtype=torch.int32)) assert fmha._mla_decode_wrappers[2] is wrapper @@ -1910,12 +1910,12 @@ def test_mla_wrapper_receives_v2_bound_and_shared_workspace( assert run_kwargs["block_tables"].shape == (2, 3) assert run_kwargs["block_tables"].stride() == (6, 1) assert run_kwargs["block_tables"].data_ptr() == block_tables.data_ptr() - torch.testing.assert_close(run_kwargs["seq_lens"], params.sequence_lengths) - assert run_kwargs["seq_lens"].data_ptr() == params.sequence_lengths.data_ptr() + torch.testing.assert_close(run_kwargs["seq_lens_kv"], params.sequence_lengths) + assert run_kwargs["seq_lens_kv"].data_ptr() == params.sequence_lengths.data_ptr() assert run_kwargs["out"].shape == (2, 1, attn.num_heads, 512) assert run_kwargs["out"].data_ptr() == output.data_ptr() - assert run_kwargs["bmm1_scale"] == pytest.approx(1.0 / math.sqrt(128 + 64)) - assert run_kwargs["bmm2_scale"] == 1.0 + assert run_kwargs["softmax_scale"] == pytest.approx(1.0 / math.sqrt(128 + 64)) + assert run_kwargs["output_scale"] == 1.0 assert run_kwargs["validate"] is False builder_args = build_metadata.call_args.args assert builder_args[8] == total_num_blocks @@ -2372,7 +2372,7 @@ def test_cache_free_context_plans_contiguous_wrapper_once(monkeypatch: pytest.Mo wrapper.plan.assert_called_once() plan_kwargs = wrapper.plan.call_args.kwargs assert plan_kwargs["batch_size"] == batch_size - assert plan_kwargs["max_seq_len_q"] == plan_kwargs["max_kv_len"] == seq_len + assert plan_kwargs["max_q_len"] == plan_kwargs["max_kv_len"] == seq_len assert ( plan_kwargs["q_dtype"] == plan_kwargs["k_dtype"] == plan_kwargs["v_dtype"] == torch.bfloat16 )