From 24e54bdffa45f46853a249aff1a2b3917675699c Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Mon, 12 Oct 2026 10:24:43 +0000 Subject: [PATCH 1/3] [None][chore] Refresh PrimTS vendor source onto FlashInfer main Move to the dated canonical branch trtllm-prims-ts-dev-20261012 at ae312f7cc1723b06fda5168105ade0f328ee47d8. Merged FlashInfer PRs #4829, #4872, #5002, and #5144 cover the downstream series; #5478 supersedes the standalone CUTLASS DSL 4.8 fix. No downstream source commits remain to replay. Regenerate the TRT-LLM compatibility patch and bind the paged context adapter to separate K/V dtype arguments. Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../vendor_patches/flashinfer-prims-ts.patch | 226 +- 3rdparty/vendor_sources.lock.yaml | 8 +- .../attention/backends/fmha/prims_ts.py | 3 +- .../attention/backends/prims_ts/__init__.py | 25 +- .../backends/prims_ts/_block_sparse/common.py | 25 + .../prims_ts/_block_sparse/compiler.py | 631 ++-- .../backends/prims_ts/_block_sparse/config.py | 273 +- .../prims_ts/_block_sparse/inspection.py | 9 +- .../backends/prims_ts/_block_sparse/plan.py | 164 +- .../prims_ts/_block_sparse/runtime.py | 382 +-- .../_q_token_kv_block_sparse_policy.py | 235 ++ .../backends/prims_ts/_tensor_aliasing.py | 109 - .../backends/prims_ts/block_sparse.py | 358 ++- .../attention/backends/prims_ts/context.py | 519 +++- .../attention/backends/prims_ts/decode.py | 2767 ++++++++++++----- .../kernels/fmha_context/fmha_kernel.py | 648 +++- .../kernels/fmha_context/fmha_resources.py | 1209 +++++-- .../kernels/fmha_context/fmha_tasks.py | 1281 ++++++-- .../prims_ts/kernels/fmha_context/helpers.py | 23 +- .../fmha_decode/direct_sparse_metadata.py | 201 ++ .../kernels/fmha_decode/fmha_decode_config.py | 1768 +++++++++-- .../fmha_decode/fmha_decode_constants.py | 91 +- .../kernels/fmha_decode/fmha_decode_kernel.py | 1862 ++++++++--- .../fmha_decode_resources/__init__.py | 16 +- .../fmha_decode_resources/helpers_common.py | 169 +- .../helpers_kv_tile_idx.py | 123 + .../fmha_decode_resources/helpers_output.py | 2 +- .../fmha_decode_resources/helpers_softmax.py | 129 +- .../fmha_decode_resources/sage_scales.py | 825 +++++ .../smem_block_sparse_metadata.py | 144 +- .../fmha_decode_resources/smem_p.py | 953 +++--- .../fmha_decode_resources/smem_resources.py | 1368 ++++++-- .../smem_transformed_kv.py | 574 ++++ .../fmha_decode_resources/tmem_corr.py | 955 ++++-- .../fmha_decode_resources/tmem_o.py | 129 +- .../fmha_decode_resources/tmem_s.py | 1302 ++++++-- .../tmem_softmax_stats.py | 25 +- .../tmem_transformed_kv.py | 437 +++ .../kernels/fmha_decode/fmha_decode_tasks.py | 1492 +++++---- .../prims_ts/kernels/fmha_decode/reduction.py | 118 +- .../mla_decode/throughput_2cta/tasks.py | 37 +- .../parallel_reduction.py | 52 +- .../throughput_latency_1cta/tasks.py | 13 +- .../backends/prims_ts/kernels/task_compat.py | 39 + .../attention/backends/prims_ts/mla_decode.py | 574 +--- .../q_token_kv_block_sparse_metadata.py | 2199 +++++++++++++ .../attention/backends/prims_ts/sage.py | 220 ++ .../_torch/attention/test_prims_ts_fmha.py | 7 +- 48 files changed, 19239 insertions(+), 5480 deletions(-) create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/_q_token_kv_block_sparse_policy.py delete mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/_tensor_aliasing.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/direct_sparse_metadata.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/sage_scales.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_transformed_kv.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/kernels/task_compat.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/q_token_kv_block_sparse_metadata.py create mode 100644 tensorrt_llm/_torch/attention/backends/prims_ts/sage.py diff --git a/3rdparty/vendor_patches/flashinfer-prims-ts.patch b/3rdparty/vendor_patches/flashinfer-prims-ts.patch index c59a19e167d6..a4319901a3d5 100644 --- a/3rdparty/vendor_patches/flashinfer-prims-ts.patch +++ b/3rdparty/vendor_patches/flashinfer-prims-ts.patch @@ -1,5 +1,14 @@ +diff --git a/_q_token_kv_block_sparse_policy.py b/_q_token_kv_block_sparse_policy.py +index a0ea3425aab7a23e26b7292eaaf7f5fb772621e8..45f495c955d878625fd88e13cd86212fd85f3e1b 100644 +--- a/_q_token_kv_block_sparse_policy.py ++++ b/_q_token_kv_block_sparse_policy.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"); + # you may not use this file except in compliance with the License. diff --git a/block_sparse.py b/block_sparse.py -index c7e881ca73dfb95314da532912db39fe7b05eb94..e19dc7b1de7fca03c44b8e98217f8d10fb17e26b 100644 +index d4652b8b68958ee3dab29d062a7cd082ea604c47..bb56724e986d302854722bf33962bea063f9e6e4 100644 --- a/block_sparse.py +++ b/block_sparse.py @@ -1,3 +1,4 @@ @@ -7,70 +16,79 @@ index c7e881ca73dfb95314da532912db39fe7b05eb94..e19dc7b1de7fca03c44b8e98217f8d10 # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); -@@ -35,12 +36,6 @@ from typing import Literal +@@ -36,13 +37,7 @@ from typing import Literal + import torch - from flashinfer.api_logging import flashinfer_api +-from flashinfer.api_logging import flashinfer_experimental_api -from flashinfer.trace.templates.attention import ( - prims_ts_block_sparse_trace_dispatch, - prims_ts_block_sparse_wrapper_trace_dispatch, - prims_ts_paged_block_sparse_trace_dispatch, - prims_ts_paged_block_sparse_wrapper_trace_dispatch, -) ++from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api - from ._block_sparse.common import _validate_contiguous_route_mode - from ._block_sparse.config import _validate_block_sparse_static_profile -@@ -239,7 +234,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase): + from ._block_sparse.common import ( + _num_sparse_pattern_heads, +@@ -301,7 +296,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase): # previously published revision intact and runnable. self._plan_state = candidate -- @flashinfer_api(trace=prims_ts_block_sparse_wrapper_trace_dispatch) -+ @flashinfer_api +- @flashinfer_experimental_api(trace=prims_ts_block_sparse_wrapper_trace_dispatch) ++ @flashinfer_experimental_api def run( self, q: torch.Tensor, -@@ -361,7 +356,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase): +@@ -433,7 +428,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase): return self._launch_validated_run(state, run_args, run_stream) --@flashinfer_api(trace=prims_ts_block_sparse_trace_dispatch) -+@flashinfer_api +-@flashinfer_experimental_api(trace=prims_ts_block_sparse_trace_dispatch) ++@flashinfer_experimental_api def block_sparse_attention( q: torch.Tensor, k: torch.Tensor, -@@ -607,7 +602,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase): +@@ -726,9 +721,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase): ) self._plan_state = candidate -- @flashinfer_api(trace=prims_ts_paged_block_sparse_wrapper_trace_dispatch) -+ @flashinfer_api +- @flashinfer_experimental_api( +- trace=prims_ts_paged_block_sparse_wrapper_trace_dispatch +- ) ++ @flashinfer_experimental_api def run( self, q: torch.Tensor, -@@ -727,7 +722,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase): +@@ -850,7 +843,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase): return self._launch_validated_run(state, run_args, run_stream) --@flashinfer_api(trace=prims_ts_paged_block_sparse_trace_dispatch) -+@flashinfer_api +-@flashinfer_experimental_api(trace=prims_ts_paged_block_sparse_trace_dispatch) ++@flashinfer_experimental_api def block_sparse_attention_with_paged_kv_cache( q: torch.Tensor, paged_kv_cache: PagedKVCache, diff --git a/context.py b/context.py -index 4ada42762c38c1c7215d2dfbe4582926fd8f9701..678432780332b20a78fb553405ad1e29d692b430 100644 +index dd9a29c3e36111a0de9a57c487068b6b5caa4ea2..a8a337373ad1607ca12ac12b2bb3117e485704d8 100644 --- a/context.py +++ b/context.py -@@ -31,8 +31,7 @@ position is ``q + (S_kv - S_q)`` and ``window_left`` is measured from that - position. +@@ -1,4 +1,5 @@ + # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. ++# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + # Copyright (c) 2026 by FlashInfer team. + # + # Licensed under the Apache License, Version 2.0 (the "License"); +@@ -46,7 +47,7 @@ from typing import TYPE_CHECKING, Callable, Literal, Optional - PrimTS context entry points are intentionally excluded from ``fi_trace`` for --now; unlike the decode APIs, their ``@flashinfer_api`` decorators do not --register trace templates. -+now; their ``@flashinfer_api`` decorators do not register trace templates. - """ + import torch + +-from flashinfer.api_logging import flashinfer_experimental_api ++from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api - from dataclasses import dataclass -@@ -438,7 +437,7 @@ def _validate_device(device: torch.device) -> int: + + if TYPE_CHECKING: +@@ -488,7 +489,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): @@ -80,7 +98,7 @@ index 4ada42762c38c1c7215d2dfbe4582926fd8f9701..678432780332b20a78fb553405ad1e29 require_cute_dsl_arch(device_index) return device_index diff --git a/decode.py b/decode.py -index b47093bd298ea9b5514e6b25ad2a34a47ac3bff0..35236fcc7ce98bdee1f51f72dbd1f841cdbce7b2 100644 +index 41479a386d72c27eb0df76107318dc09256aff06..3fcec5efdc2a3825be0dc1d9a238aaaa6d76a609 100644 --- a/decode.py +++ b/decode.py @@ -1,3 +1,4 @@ @@ -88,19 +106,20 @@ index b47093bd298ea9b5514e6b25ad2a34a47ac3bff0..35236fcc7ce98bdee1f51f72dbd1f841 # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); -@@ -25,11 +26,6 @@ from typing import TYPE_CHECKING, Literal, Optional, Union +@@ -24,11 +25,7 @@ from typing import TYPE_CHECKING, Literal, NamedTuple, Optional, Union, cast + import torch - from flashinfer.api_logging import flashinfer_api +-from flashinfer.api_logging import flashinfer_experimental_api -from flashinfer.trace.templates.attention import ( - attention_ts_decode_trace_dispatch, -- prims_ts_decode_trace_dispatch, - prims_ts_decode_wrapper_trace_dispatch, -) ++from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api - from ._tensor_aliasing import ( - _validate_out_does_not_overlap_inputs, -@@ -702,7 +698,7 @@ def _validate_runtime_device(device: torch.device) -> int: + from . import _q_token_kv_block_sparse_policy as _sparse_policy + from ._block_sparse.common import _num_sparse_pattern_heads +@@ -789,7 +786,7 @@ def _validate_runtime_device(device: torch.device) -> None: # 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): @@ -108,36 +127,72 @@ index b47093bd298ea9b5514e6b25ad2a34a47ac3bff0..35236fcc7ce98bdee1f51f72dbd1f841 + from flashinfer.cute_dsl.utils import require_cute_dsl_arch require_cute_dsl_arch(device_index) - return device_index -@@ -2145,7 +2141,7 @@ def _validate_decode_run_metadata_values( - ) - --@flashinfer_api(trace=prims_ts_decode_trace_dispatch) -+@flashinfer_api - def prims_ts_batch_decode_with_kv_cache( - query: torch.Tensor, - kv_cache: PagedKVCache, -@@ -2708,7 +2704,7 @@ class BatchDecodePagedTSWrapper: +@@ -4046,7 +4043,7 @@ class BatchDecodePagedTSWrapper: # previous complete plan revision usable. self._plan_state = candidate -- @flashinfer_api(trace=prims_ts_decode_wrapper_trace_dispatch) -+ @flashinfer_api +- @flashinfer_experimental_api(trace=prims_ts_decode_wrapper_trace_dispatch) ++ @flashinfer_experimental_api def run( self, q: torch.Tensor, -@@ -2893,7 +2889,7 @@ class BatchDecodePagedTSWrapper: +@@ -4228,7 +4225,7 @@ class BatchDecodePagedTSWrapper: ) --@flashinfer_api(trace=attention_ts_decode_trace_dispatch) -+@flashinfer_api +-@flashinfer_experimental_api(trace=attention_ts_decode_trace_dispatch) ++@flashinfer_experimental_api def batch_decode_with_paged_kv_cache( q: torch.Tensor, paged_kv_cache: PagedKVCache, +diff --git a/kernels/fmha_decode/direct_sparse_metadata.py b/kernels/fmha_decode/direct_sparse_metadata.py +index 8bffb076c4a3ef56bc9e463733066abcaf0a7725..15d6f271029540de1e76633e1a81cd78bc553729 100644 +--- a/kernels/fmha_decode/direct_sparse_metadata.py ++++ b/kernels/fmha_decode/direct_sparse_metadata.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"); + # you may not use this file except in compliance with the License. +diff --git a/kernels/fmha_decode/fmha_decode_resources/sage_scales.py b/kernels/fmha_decode/fmha_decode_resources/sage_scales.py +index f2df9bbeafcbdf22e67f9a754d06144b8e3a1dab..4d862f9fa1f1be3e8830a7bbc502e0369237f0d4 100644 +--- a/kernels/fmha_decode/fmha_decode_resources/sage_scales.py ++++ b/kernels/fmha_decode/fmha_decode_resources/sage_scales.py +@@ -1,3 +1,4 @@ ++# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + # Copyright (c) 2026 by FlashInfer team. + # + # Licensed under the Apache License, Version 2.0 (the "License"); +diff --git a/kernels/fmha_decode/fmha_decode_resources/smem_transformed_kv.py b/kernels/fmha_decode/fmha_decode_resources/smem_transformed_kv.py +index 184326c482f21ecb1a81f8ba337a996f9ab2bea8..148ebcae6b7431985bb91c02e7fa6d5ddfe69253 100644 +--- a/kernels/fmha_decode/fmha_decode_resources/smem_transformed_kv.py ++++ b/kernels/fmha_decode/fmha_decode_resources/smem_transformed_kv.py +@@ -1,3 +1,4 @@ ++# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + # Copyright (c) 2026 by FlashInfer team. + # + # Licensed under the Apache License, Version 2.0 (the "License"); +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 18f2c272cf5c1e089193a8a5437560f45cbbbb8e..1d607a92e63df03c04465a9c7836422088c13d49 100644 +--- a/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py ++++ b/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py +@@ -1,3 +1,4 @@ ++# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + # Copyright (c) 2026 by FlashInfer team. + # + # Licensed under the Apache License, Version 2.0 (the "License"); +diff --git a/kernels/task_compat.py b/kernels/task_compat.py +index 80a80102c3ecf6a856d441c0dac2ca5c8215ad82..69f60317b172457e7fda347a3b935ae67dbdd892 100644 +--- a/kernels/task_compat.py ++++ b/kernels/task_compat.py +@@ -1,3 +1,4 @@ ++# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + # Copyright (c) 2026 by FlashInfer team. + # + # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/mla_decode.py b/mla_decode.py -index 8b7f19d1cc2baab6309018d46213f509c44d919a..cff2ee7d2ad051c8061ad396b87ee973eebed788 100644 +index 41ba3b9055d229d227ac10cfc6f7d0f403eae69f..6d69ae7b554afe1502f4c369e5ddc805847b876c 100644 --- a/mla_decode.py +++ b/mla_decode.py @@ -1,3 +1,4 @@ @@ -145,20 +200,73 @@ index 8b7f19d1cc2baab6309018d46213f509c44d919a..cff2ee7d2ad051c8061ad396b87ee973 # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); -@@ -25,7 +26,6 @@ from flashinfer.api_logging import flashinfer_api - from flashinfer.trace.templates.attention import ( - prims_ts_decode_mla_one_shot_trace_dispatch, - prims_ts_decode_mla_trace_dispatch, +@@ -21,11 +22,7 @@ from typing import Any, Literal, Optional, cast + + import torch + +-from flashinfer.api_logging import flashinfer_experimental_api +-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 ._tensor_aliasing import ( -@@ -1799,7 +1799,7 @@ class BatchMLADecodePagedTSWrapper: + from .decode import ( + _WorkspaceSection, +@@ -1576,7 +1573,7 @@ class BatchMLADecodePagedTSWrapper: split_kv=int(dict(policy)["split_kv"]), ) -- @flashinfer_api(trace=prims_ts_decode_mla_wrapper_trace_dispatch) -+ @flashinfer_api +- @flashinfer_experimental_api(trace=prims_ts_decode_mla_wrapper_trace_dispatch) ++ @flashinfer_experimental_api def run( self, query: torch.Tensor, +@@ -1681,7 +1678,7 @@ class BatchMLADecodePagedTSWrapper: + ) + + +-@flashinfer_experimental_api(trace=prims_ts_decode_mla_one_shot_trace_dispatch) ++@flashinfer_experimental_api + def batch_mla_decode_with_paged_kv_cache( + query: torch.Tensor, + kv_cache: torch.Tensor, +diff --git a/q_token_kv_block_sparse_metadata.py b/q_token_kv_block_sparse_metadata.py +index 35efc60b7ff3dac6af1c5c5e46631712a171fb72..0c5cf40dbd4e4708b8f74bfa698f3c0788f9f3a4 100644 +--- a/q_token_kv_block_sparse_metadata.py ++++ b/q_token_kv_block_sparse_metadata.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"); +@@ -40,9 +41,9 @@ from dataclasses import dataclass + from typing import Any, Literal, Optional + + import torch +-from flashinfer.api_logging import flashinfer_experimental_api ++from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api + +-from ...utils import device_support_pdl ++from flashinfer.utils import device_support_pdl + from ._block_sparse.common import _num_sparse_pattern_heads + + _Q_TOKEN_KV_BLOCK_SPARSE_MIN_BLOCK_SIZE = 4 +@@ -62,7 +63,7 @@ _Q_TOKEN_KV_BLOCK_SPARSE_WORKSPACE_ALIGNMENT = 256 + def _get_prims_ts_q_token_kv_block_sparse_metadata_module() -> Any: + """Build and cache the CUDA C++ QToken-KvBlock-Sparse-Attention metadata implementation.""" + +- from ...jit import gen_prims_ts_q_token_kv_block_sparse_metadata_module ++ from flashinfer.jit import gen_prims_ts_q_token_kv_block_sparse_metadata_module + + return gen_prims_ts_q_token_kv_block_sparse_metadata_module().build_and_load() + +diff --git a/sage.py b/sage.py +index 97708c79525ba840e4ebfd11876084eae952254d..dcf0d4ae05a0d1358cfa796a3e23dbe098fbab10 100644 +--- a/sage.py ++++ b/sage.py +@@ -1,3 +1,4 @@ ++# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + # Copyright (c) 2026 by FlashInfer team. + # + # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index cb0cfb444e45..9f5f510adefd 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -2,15 +2,15 @@ schema_version: 1 vendors: flashinfer-prims-ts: url: https://github.com/yuxianq/flashinfer.git - branch: trtllm-prims-ts-dev - commit: 6e0d64e79f2a87eeb539e2f40ea12d009679a8f5 + branch: trtllm-prims-ts-dev-20261012 + commit: ae312f7cc1723b06fda5168105ade0f328ee47d8 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:ba917f330d8281781ce30ddb89f1772ee8f9017264d948b588b881a92f511e45 - digest: sha256-tree-v1:fa65b674a9ec6121b4158069d5585cc56df28243d1ac1199eb97003d81135d76 + patch_digest: sha256:a9caa1909814f67df34fff35f21043d95c659ccbd86975924c9d8784adcdaf39 + digest: sha256-tree-v1:1f2f960dfe7b4f875b2e8d6a8b119fec9aae18a527e175b2cd4bde66393e16b0 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 4b68b9cc53aa..2e67616b5e9d 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py @@ -621,7 +621,8 @@ def _get_or_plan_context_wrapper( num_kv_heads=int(k_cache.shape[1]), head_dim=int(q.shape[-1]), q_dtype=q.dtype, - kv_dtype=k_cache.dtype, + k_dtype=k_cache.dtype, + v_dtype=v_cache.dtype, out_dtype=output_dtype, page_size=page_size, mask_type=mask_type, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/__init__.py b/tensorrt_llm/_torch/attention/backends/prims_ts/__init__.py index 3cb31497cefa..e40d8ce0d41c 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/__init__.py @@ -22,9 +22,13 @@ ) from .decode import ( BatchDecodePagedTSWrapper, + PrimsTSBatchDecodePlan, batch_decode_with_paged_kv_cache, get_prims_ts_batch_decode_workspace_size, - prims_ts_batch_decode_with_kv_cache, + make_q_token_kv_block_sparse_qo_indptr, + prepare_prims_ts_batch_decode_with_kv_cache, + suggest_q_token_kv_block_sparse_group_size, + validate_q_token_kv_block_sparse_group_size, ) from .context import ( BatchPrefillPagedTSWrapper, @@ -36,11 +40,18 @@ BatchMLADecodePagedTSWrapper, batch_mla_decode_with_paged_kv_cache, get_prims_ts_batch_mla_decode_workspace_size, - prims_ts_batch_mla_decode_with_kv_cache, ) +from .q_token_kv_block_sparse_metadata import ( + QTokenKvBlockSparsePagedTSWrapper, + get_q_token_kv_block_sparse_workspace_size, + q_token_kv_block_sparse_attention_with_paged_kv_cache, +) +from .sage import SageAttentionConfig, SageAttentionParams __all__ = [ "BlockSparseTSWrapper", + "SageAttentionConfig", + "SageAttentionParams", "BlockSparsePagedTSWrapper", "block_sparse_attention", "block_sparse_attention_with_paged_kv_cache", @@ -49,11 +60,17 @@ "batch_prefill", "batch_prefill_with_paged_kv_cache", "BatchDecodePagedTSWrapper", + "PrimsTSBatchDecodePlan", "batch_decode_with_paged_kv_cache", "get_prims_ts_batch_decode_workspace_size", - "prims_ts_batch_decode_with_kv_cache", + "make_q_token_kv_block_sparse_qo_indptr", + "prepare_prims_ts_batch_decode_with_kv_cache", + "suggest_q_token_kv_block_sparse_group_size", + "validate_q_token_kv_block_sparse_group_size", + "QTokenKvBlockSparsePagedTSWrapper", + "get_q_token_kv_block_sparse_workspace_size", + "q_token_kv_block_sparse_attention_with_paged_kv_cache", "BatchMLADecodePagedTSWrapper", "batch_mla_decode_with_paged_kv_cache", "get_prims_ts_batch_mla_decode_workspace_size", - "prims_ts_batch_mla_decode_with_kv_cache", ] 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 13cf6038efd9..f67fdba1602b 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,6 +14,8 @@ """Dependency-neutral semantic contract for PrimTS block-sparse attention.""" +from collections.abc import Iterable + _FINE_KV_BLOCK_SIZES = (8, 16, 32) _PREPARED_KV_ROUTE_SIZE = 128 _MAX_KV_ATOM_SIZE = 64 @@ -22,6 +24,16 @@ _BLOCK_SPARSE_MAX_HEADS_Q_PER_KV = 32 +def _num_sparse_pattern_heads( + num_kv_heads: int, share_pattern_across_kv_heads: bool +) -> int: + """Return metadata head count without changing physical KV-head geometry. + + The caller resolves its API's pattern-sharing default before calling. + """ + return 1 if share_pattern_across_kv_heads else num_kv_heads + + def _validate_contiguous_route_mode( sparse_format: object, use_proxy_routes: object, @@ -36,6 +48,19 @@ def _validate_contiguous_route_mode( raise TypeError("use_proxy_routes must be a bool") +def _validate_dense_contiguous_plan_inputs( + inputs: Iterable[tuple[str, object, object]], +) -> None: + """Reject block-sparse-only inputs given to a dense contiguous plan. + + Each entry is ``(name, value, unused value)``. + """ + + for name, value, unused in inputs: + if value != unused: + raise ValueError(f"{name} is unsupported by a dense contiguous plan") + + def _validate_sparse_q_block_size(value: object) -> int: """Return a positive semantic Q block size representable by the ABI.""" 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 fcaa0735436f..1a9d42857f94 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 @@ -21,6 +21,8 @@ from flashinfer.utils import ceil_div +from .common import _num_sparse_pattern_heads +from ..sage import sage_adapter_slots, sage_scale_shapes from .config import _BlockSparseCompileKey, _make_block_sparse_config @@ -28,7 +30,18 @@ def _compile_block_sparse(key: _BlockSparseCompileKey) -> Callable[..., object]: - """Compile one storage-specialized prepare-plus-attention adapter.""" + """Compile one contiguous or paged attention adapter for ``key``. + + Every contiguous plan, dense or block-sparse, exact or proxy, BSR or + bitmask, compiles the same adapter signature. The plan decides at compile + time which tensor slots it uses; the unused slots are ``None`` both for the + compile-time fakes and for every later call, and the adapter body reads a + slot only on the branch its configuration enables. A dense plan runs no + prepare kernel and reaches the decode kernel directly; a block-sparse plan + prepares its routes and launches the prepared-route attention. Paged + block-sparse plans keep their own adapter because their K/V and request + metadata differ in kind, not only in presence. + """ import cutlass import cutlass.cute as cute @@ -41,222 +54,72 @@ def _compile_block_sparse(key: _BlockSparseCompileKey) -> Callable[..., object]: ) from ..kernels.fmha_decode.fmha_decode_kernel import ( fmha_block_sparse_launch, + fmha_decode_launch, ) config = _make_block_sparse_config(key) - prepare_kwargs = { - "batch_size": key.batch_size, - "num_kv_heads": key.num_kv_heads, - "seq_len_q": key.seq_len_q, - "seq_len_kv": key.seq_len_kv, - "q_block_size": key.q_block_size, - "kv_block_size": key.kv_block_size, - "kv_route_size": key.kv_route_size, - "use_proxy_routes": key.use_proxy_routes, - "use_causal_mask": key.mask_type == "causal", - "apply_token_mask": key.use_kv_valid_bits, - "store_score_words": config.uses_prepared_score_keep_words, - } - if key.page_size is not None: - if key.sparse_format != "bsr" or key.use_proxy_routes: - raise AssertionError("paged block-sparse supports exact BSR routes only") - prepare_kwargs["page_size"] = key.page_size - - if key.sparse_format == "bsr": - prepare_routes = _PrepareBsrRoutes(**prepare_kwargs) - elif key.sparse_format == "bitmask": - prepare_routes = _PrepareBitmaskRoutes(**prepare_kwargs) - else: - raise AssertionError("sparse_format must be 'bsr' or 'bitmask'") + sparse_format = key.sparse_format + use_proxy_routes = key.use_proxy_routes + prepare_routes: _PrepareBsrRoutes | _PrepareBitmaskRoutes | None = None + route_metadata_base = 0 + pattern_heads = _num_sparse_pattern_heads( + key.num_kv_heads, key.share_pattern_across_kv_heads + ) + if key.use_block_sparse: + prepare_kwargs = { + "batch_size": key.batch_size, + "num_kv_heads": pattern_heads, + "seq_len_q": key.seq_len_q, + "seq_len_kv": key.seq_len_kv, + "q_block_size": key.q_block_size, + "kv_block_size": key.kv_block_size, + "kv_route_size": key.kv_route_size, + "use_proxy_routes": use_proxy_routes, + "use_causal_mask": key.mask_type == "causal", + "apply_token_mask": key.use_kv_valid_bits, + "store_score_words": config.uses_prepared_score_keep_words, + } + if key.page_size is not None: + if sparse_format != "bsr" or use_proxy_routes: + raise AssertionError( + "paged block-sparse supports exact BSR routes only" + ) + prepare_kwargs["page_size"] = key.page_size + if sparse_format == "bsr": + prepare_routes = _PrepareBsrRoutes(**prepare_kwargs) + elif sparse_format == "bitmask": + prepare_routes = _PrepareBitmaskRoutes(**prepare_kwargs) + else: + raise AssertionError("sparse_format must be 'bsr' or 'bitmask'") + route_metadata_base = prepare_routes.route_metadata_base_word_offset + elif key.page_size is not None: + raise AssertionError("dense attention requires contiguous K/V") - route_metadata_base = prepare_routes.route_metadata_base_word_offset Int32 = cutlass.Int32 Int64 = cutlass.Int64 Float32 = cutlass.Float32 @cute.jit - def exact_bsr_adapter( - q: cute.Tensor, - k: cute.Tensor, - v: cute.Tensor, - out: cute.Tensor, - block_indptr: cute.Tensor, - block_indices: cute.Tensor, - kv_valid_bits: cute.Tensor, - row_route_offsets: cute.Tensor, - route_workspace: cute.Tensor, - max_blocks_per_row: cutlass.Int32, - sm_scale: cutlass.Float32, - stream: cuda_drv.CUstream, - static_config: cutlass.Constexpr[FmhaDecodeConfig], - static_batch_size: cutlass.Constexpr[int], - static_seq_len_kv: cutlass.Constexpr[int], - static_num_qo_heads: cutlass.Constexpr[int], - static_num_kv_heads: cutlass.Constexpr[int], - static_head_dim: cutlass.Constexpr[int], - ) -> None: - prepare_routes( - block_indptr, - block_indices, - kv_valid_bits, - None, - None, - Int64(0), - Int64(0), - row_route_offsets, - route_workspace, - max_blocks_per_row, - stream, - ) - # Live per-row route counts occupy the first words of run scratch. - row_route_counts = route_workspace.iterator - route_metadata = route_workspace.iterator + Int32(route_metadata_base) - fmha_block_sparse_launch( - ( - Int32(static_batch_size), - Int32(static_num_qo_heads), - Int32(static_num_kv_heads), - Int32(static_seq_len_kv), - Int32(static_head_dim), - ), - q.iterator, - k.iterator, - v.iterator, - k.iterator, - v.iterator, - out.iterator, - row_route_offsets.iterator, - row_route_counts, - route_metadata, - sm_scale, - stream, - static_config, - static_seq_len_kv, - ) - - @cute.jit - def exact_bitmask_adapter( - q: cute.Tensor, - k: cute.Tensor, - v: cute.Tensor, - out: cute.Tensor, - exact_block_bits: cute.Tensor, - kv_valid_bits: cute.Tensor, - row_route_offsets: cute.Tensor, - route_workspace: cute.Tensor, - max_blocks_per_row: cutlass.Int32, - sm_scale: cutlass.Float32, - stream: cuda_drv.CUstream, - static_config: cutlass.Constexpr[FmhaDecodeConfig], - static_batch_size: cutlass.Constexpr[int], - static_seq_len_kv: cutlass.Constexpr[int], - static_num_qo_heads: cutlass.Constexpr[int], - static_num_kv_heads: cutlass.Constexpr[int], - static_head_dim: cutlass.Constexpr[int], - ) -> None: - prepare_routes( - exact_block_bits, - kv_valid_bits, - row_route_offsets, - route_workspace, - max_blocks_per_row, - stream, - ) - fmha_block_sparse_launch( - ( - Int32(static_batch_size), - Int32(static_num_qo_heads), - Int32(static_num_kv_heads), - Int32(static_seq_len_kv), - Int32(static_head_dim), - ), - q.iterator, - k.iterator, - v.iterator, - k.iterator, - v.iterator, - out.iterator, - row_route_offsets.iterator, - route_workspace.iterator, - route_workspace.iterator + Int32(route_metadata_base), - sm_scale, - stream, - static_config, - static_seq_len_kv, - ) - - @cute.jit - def proxy_bsr_adapter( - q: cute.Tensor, - k: cute.Tensor, - v: cute.Tensor, - k_summary: cute.Tensor, - v_summary: cute.Tensor, - out: cute.Tensor, - block_indptr: cute.Tensor, - block_indices: cute.Tensor, - kv_valid_bits: cute.Tensor, - row_route_offsets: cute.Tensor, - route_workspace: cute.Tensor, - max_blocks_per_row: cutlass.Int32, - sm_scale: cutlass.Float32, - stream: cuda_drv.CUstream, - static_config: cutlass.Constexpr[FmhaDecodeConfig], - static_batch_size: cutlass.Constexpr[int], - static_seq_len_kv: cutlass.Constexpr[int], - static_num_qo_heads: cutlass.Constexpr[int], - static_num_kv_heads: cutlass.Constexpr[int], - static_head_dim: cutlass.Constexpr[int], - ) -> None: - prepare_routes( - block_indptr, - block_indices, - kv_valid_bits, - None, - None, - Int64(0), - Int64(0), - row_route_offsets, - route_workspace, - max_blocks_per_row, - stream, - ) - fmha_block_sparse_launch( - ( - Int32(static_batch_size), - Int32(static_num_qo_heads), - Int32(static_num_kv_heads), - Int32(static_seq_len_kv), - Int32(static_head_dim), - ), - q.iterator, - k.iterator, - v.iterator, - k_summary.iterator, - v_summary.iterator, - out.iterator, - row_route_offsets.iterator, - route_workspace.iterator, - route_workspace.iterator + Int32(route_metadata_base), - sm_scale, - stream, - static_config, - static_seq_len_kv, - ) - - @cute.jit - def proxy_bitmask_adapter( + def contiguous_adapter( q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, - k_summary: cute.Tensor, - v_summary: cute.Tensor, + k_summary: cute.Tensor | None, + v_summary: cute.Tensor | None, out: cute.Tensor, - exact_block_bits: cute.Tensor, - kv_valid_bits: cute.Tensor, - row_route_offsets: cute.Tensor, - route_workspace: cute.Tensor, + block_indptr: cute.Tensor | None, + block_indices: cute.Tensor | None, + exact_block_bits: cute.Tensor | None, + kv_valid_bits: cute.Tensor | None, + row_route_offsets: cute.Tensor | None, + route_workspace: cute.Tensor | None, max_blocks_per_row: cutlass.Int32, + # The Sage scale slots, in ``SAGE_ADAPTER_SLOTS`` order. + q_scale: cute.Tensor | None, + k_scale: cute.Tensor | None, + k_summary_scale: cute.Tensor | None, + v_scale: cute.Tensor | None, + v_mean: cute.Tensor | None, sm_scale: cutlass.Float32, stream: cuda_drv.CUstream, static_config: cutlass.Constexpr[FmhaDecodeConfig], @@ -266,36 +129,131 @@ def proxy_bitmask_adapter( static_num_kv_heads: cutlass.Constexpr[int], static_head_dim: cutlass.Constexpr[int], ) -> None: - prepare_routes( - exact_block_bits, - kv_valid_bits, - row_route_offsets, - route_workspace, - max_blocks_per_row, - stream, - ) - fmha_block_sparse_launch( - ( - Int32(static_batch_size), - Int32(static_num_qo_heads), - Int32(static_num_kv_heads), - Int32(static_seq_len_kv), - Int32(static_head_dim), - ), - q.iterator, - k.iterator, - v.iterator, - k_summary.iterator, - v_summary.iterator, - out.iterator, - row_route_offsets.iterator, - route_workspace.iterator, - route_workspace.iterator + Int32(route_metadata_base), - sm_scale, - stream, - static_config, - static_seq_len_kv, + problem_shape = ( + Int32(static_batch_size), + Int32(static_num_qo_heads), + Int32(static_num_kv_heads), + Int32(static_seq_len_kv), + Int32(static_head_dim), ) + # Sage scale tensors are [heads, flat slots]; the head stride is the + # flat slot count of one head. A slot the recipe does not use (the V + # mean without a mean, the summary K scales without proxy routes) is + # ``None`` and stays unbound. + sage_kwargs = {} + if cutlass.const_expr(static_config.use_sage_attention): + sage_kwargs = { + "sage_q_scale_iter": q_scale.iterator, + "sage_k_scale_iter": k_scale.iterator, + "sage_v_scale_iter": v_scale.iterator, + "sage_q_scale_head_stride": Int32(q_scale.shape[1]), + "sage_k_scale_head_stride": Int32(k_scale.shape[1]), + } + if cutlass.const_expr(static_config.sage_v_mean): + sage_kwargs["sage_v_mean_iter"] = v_mean.iterator + if cutlass.const_expr(static_config.use_block_sparse_proxy_routes): + sage_kwargs["sage_k_summary_scale_iter"] = k_summary_scale.iterator + sage_kwargs["sage_k_summary_scale_head_stride"] = Int32( + k_summary_scale.shape[1] + ) + if cutlass.const_expr(static_config.use_block_sparse): + if cutlass.const_expr(sparse_format == "bsr"): + prepare_routes( + block_indptr, + block_indices, + kv_valid_bits, + None, + None, + Int64(0), + Int64(0), + row_route_offsets, + route_workspace, + max_blocks_per_row, + stream, + ) + else: + prepare_routes( + exact_block_bits, + kv_valid_bits, + row_route_offsets, + route_workspace, + max_blocks_per_row, + stream, + ) + # Exact routes address K/V through the summary TensorMaps as well; + # proxy routes address the caller's block summaries. + if cutlass.const_expr(static_config.use_block_sparse_proxy_routes): + k_summary_iter = k_summary.iterator + v_summary_iter = v_summary.iterator + else: + k_summary_iter = k.iterator + v_summary_iter = v.iterator + # Live per-row route counts occupy the first words of run scratch. + fmha_block_sparse_launch( + problem_shape, + q.iterator, + k.iterator, + v.iterator, + k_summary_iter, + v_summary_iter, + out.iterator, + row_route_offsets.iterator, + route_workspace.iterator, + route_workspace.iterator + Int32(route_metadata_base), + sm_scale, + stream, + static_config, + static_seq_len_kv, + **sage_kwargs, + ) + else: + null_i32 = cute.make_ptr(Int32, 0, mem_space=cutlass.AddressSpace.gmem) + null_f32 = cute.make_ptr(Float32, 0, mem_space=cutlass.AddressSpace.gmem) + # Only packed NVFP4 K/V read scale factors. + null_u8 = cute.make_ptr( + cutlass.Uint8, 0, mem_space=cutlass.AddressSpace.gmem + ) + fmha_decode_launch( + problem_shape, + q.iterator, + k.iterator, + v.iterator, + null_u8, # k_sf + null_u8, # v_sf + out.iterator, + null_i32, # seqlens_kv + null_i32, # cu_seqlens_q + Int32(0), # total_q_tokens + null_i32, # page_idx_kv + null_i32, # q_token_kv_block_sparse_page_memberships + out.iterator, # partial_o + null_f32, # partial_stats + null_i32, # split_kv_counter + null_f32, # attention_sinks + sm_scale, + Float32(1.0), # output_scale + Int32( + static_seq_len_kv * static_num_kv_heads * static_head_dim + ), # kv_b_stride + Int32(0), # max_active_clusters + stream, + static_config, + static_seq_len_kv, + # Contiguous K/V has no page table; the kernel parameters + # behind these launcher arguments are typed, so pass typed + # zeros rather than relying on the launcher's literal defaults. + page_table_stride=Int64(0), + page_table_capacity=Int32(0), + q_token_kv_block_sparse_page_membership_stride=Int32(0), + num_physical_kv_pages=Int64(0), + k_page_stride=Int64(0), + k_head_stride=Int64(0), + k_token_stride=Int64(0), + v_page_stride=Int64(0), + v_head_stride=Int64(0), + v_token_stride=Int64(0), + **sage_kwargs, + ) @cute.jit def paged_tensor_adapter( @@ -368,7 +326,7 @@ def paged_tensor_adapter( ) def fake_compact( - dtype: object, shape: tuple[object, ...], alignment: int + dtype: object, shape: tuple[object, ...], alignment: int = 16 ) -> object: return cute.runtime.make_fake_compact_tensor( dtype, @@ -377,34 +335,11 @@ def fake_compact( assumed_align=alignment, ) - logical_nnz = cute.sym_int() - logical_workspace_words = cute.sym_int() q_shape = (key.batch_size, key.seq_len_q, key.num_qo_heads, key.head_dim) num_q_blocks = ceil_div(key.seq_len_q, key.q_block_size) num_kv_blocks = ceil_div(key.seq_len_kv, key.kv_block_size) - indptr_fake = fake_compact( - Int32, - (key.batch_size, key.num_kv_heads, num_q_blocks + 1), - 4, - ) - indices_fake = fake_compact(Int32, (logical_nnz,), 4) - valid_bits_fake = fake_compact( - cutlass.Uint32, - (key.batch_size, ceil_div(key.seq_len_kv, 32)), - 4, - ) - row_route_offsets_fake = fake_compact( - Int32, - (key.batch_size * key.num_kv_heads * num_q_blocks + 1,), - 4, - ) - route_workspace_fake = fake_compact( - Int32, - (logical_workspace_words,), - 4, - ) - q_fake = fake_compact(config.q_dtype, q_shape, 16) - out_fake = fake_compact(config.out_dtype, q_shape, 16) + q_fake = fake_compact(config.q_dtype, q_shape) + out_fake = fake_compact(config.out_dtype, q_shape) stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) tensor_adapter: Callable[..., object] dynamic_args: tuple[object, ...] @@ -415,95 +350,95 @@ def fake_compact( key.num_kv_heads, key.head_dim, ) - k_fake = fake_compact(config.kv_dtype, kv_shape, 16) - v_fake = fake_compact(config.kv_dtype, kv_shape, 16) - exact_bits_fake = fake_compact( - cutlass.Uint32, - ( - key.batch_size, - key.num_kv_heads, - num_q_blocks, - ceil_div(num_kv_blocks, 32), - ), - 4, - ) - common_tail = ( - valid_bits_fake, - row_route_offsets_fake, - route_workspace_fake, - Int32(0), - Float32(1.0), - ) - if key.sparse_format == "bsr" and not key.use_proxy_routes: - tensor_adapter = exact_bsr_adapter - dynamic_args = ( - q_fake, - k_fake, - v_fake, - out_fake, - indptr_fake, - indices_fake, - *common_tail, - ) - elif key.sparse_format == "bitmask" and not key.use_proxy_routes: - tensor_adapter = exact_bitmask_adapter - dynamic_args = ( - q_fake, - k_fake, - v_fake, - out_fake, - exact_bits_fake, - *common_tail, - ) - elif key.sparse_format == "bsr" and key.use_proxy_routes: + k_fake = fake_compact(config.k_dtype, kv_shape) + v_fake = fake_compact(config.v_dtype, kv_shape) + k_summary_fake = None + v_summary_fake = None + indptr_fake = None + indices_fake = None + exact_bits_fake = None + valid_bits_fake = None + row_route_offsets_fake = None + route_workspace_fake = None + if use_proxy_routes: summary_shape = ( key.batch_size, num_kv_blocks, key.num_kv_heads, key.head_dim, ) - k_summary_fake = fake_compact(config.kv_dtype, summary_shape, 16) - v_summary_fake = fake_compact(config.kv_dtype, summary_shape, 16) - proxy_prefix = ( - q_fake, - k_fake, - v_fake, - k_summary_fake, - v_summary_fake, - out_fake, + k_summary_fake = fake_compact(config.k_dtype, summary_shape) + v_summary_fake = fake_compact(config.v_dtype, summary_shape) + if key.use_block_sparse: + if sparse_format == "bsr": + indptr_fake = fake_compact( + Int32, + (key.batch_size, pattern_heads, num_q_blocks + 1), + 4, + ) + indices_fake = fake_compact(Int32, (cute.sym_int(),), 4) + else: + exact_bits_fake = fake_compact( + cutlass.Uint32, + ( + key.batch_size, + pattern_heads, + num_q_blocks, + ceil_div(num_kv_blocks, 32), + ), + 4, + ) + valid_bits_fake = fake_compact( + cutlass.Uint32, + (key.batch_size, ceil_div(key.seq_len_kv, 32)), + 4, ) - tensor_adapter = proxy_bsr_adapter - dynamic_args = ( - *proxy_prefix, - indptr_fake, - indices_fake, - *common_tail, + row_route_offsets_fake = fake_compact( + Int32, + (key.batch_size * pattern_heads * num_q_blocks + 1,), + 4, ) - elif key.sparse_format == "bitmask" and key.use_proxy_routes: - summary_shape = ( - key.batch_size, - num_kv_blocks, - key.num_kv_heads, - key.head_dim, + route_workspace_fake = fake_compact(Int32, (cute.sym_int(),), 4) + sage_fakes: tuple[object | None, ...] = sage_adapter_slots({}) + if key.sage is not None: + # The plan's scale shapes fill the adapter slots through + # ``sage_adapter_slots``, the helper that binds each run's + # tensors: the summary K scales only for a proxy plan, and no V + # mean for a recipe without one. + shapes = sage_scale_shapes( + key.sage, + batch_size=key.batch_size, + seq_len_q=key.seq_len_q, + seq_len_kv=key.seq_len_kv, + num_qo_heads=key.num_qo_heads, + num_kv_heads=key.num_kv_heads, + head_dim=key.head_dim, + summary_seq_len=num_kv_blocks if use_proxy_routes else None, ) - k_summary_fake = fake_compact(config.kv_dtype, summary_shape, 16) - v_summary_fake = fake_compact(config.kv_dtype, summary_shape, 16) - tensor_adapter = proxy_bitmask_adapter - dynamic_args = ( - q_fake, - k_fake, - v_fake, - k_summary_fake, - v_summary_fake, - out_fake, - exact_bits_fake, - *common_tail, + sage_fakes = tuple( + None if shape is None else fake_compact(Float32, shape) + for shape in sage_adapter_slots(shapes) ) - else: - raise AssertionError("continuous sparse_format must be 'bsr' or 'bitmask'") + tensor_adapter = contiguous_adapter + dynamic_args = ( + q_fake, + k_fake, + v_fake, + k_summary_fake, + v_summary_fake, + out_fake, + indptr_fake, + indices_fake, + exact_bits_fake, + valid_bits_fake, + row_route_offsets_fake, + route_workspace_fake, + Int32(0), + *sage_fakes, + Float32(1.0), + ) else: page_size = key.page_size - assert page_size is not None physical_pages = cute.sym_int() runtime_page_columns = cute.sym_int() runtime_page_row_stride = cute.sym_int64(divisibility=1) @@ -516,7 +451,7 @@ def fake_compact( key.head_dim, ) k_fake = cute.runtime.make_fake_tensor( - config.kv_dtype, + config.k_dtype, kv_shape, stride=( k_outer_stride, @@ -527,7 +462,7 @@ def fake_compact( assumed_align=16, ) v_fake = cute.runtime.make_fake_tensor( - config.kv_dtype, + config.v_dtype, kv_shape, stride=( v_outer_stride, @@ -544,6 +479,23 @@ def fake_compact( assumed_align=4, ) seq_lens_kv_fake = fake_compact(Int32, (key.batch_size,), 4) + indptr_fake = fake_compact( + Int32, + (key.batch_size, pattern_heads, num_q_blocks + 1), + 4, + ) + indices_fake = fake_compact(Int32, (cute.sym_int(),), 4) + valid_bits_fake = fake_compact( + cutlass.Uint32, + (key.batch_size, ceil_div(key.seq_len_kv, 32)), + 4, + ) + row_route_offsets_fake = fake_compact( + Int32, + (key.batch_size * pattern_heads * num_q_blocks + 1,), + 4, + ) + route_workspace_fake = fake_compact(Int32, (cute.sym_int(),), 4) tensor_adapter = paged_tensor_adapter dynamic_args = ( q_fake, @@ -584,7 +536,10 @@ def fake_compact( def _get_compiled_block_sparse( key: _BlockSparseCompileKey, ) -> Callable[..., object]: - """Compile and cache one contiguous or paged prepare-plus-attention adapter.""" + """Compile and cache the adapter of one contiguous or paged plan. + + A contiguous plan is dense or block-sparse; a paged plan is block-sparse. + """ return _compile_block_sparse(key) 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 cc4458a4298a..7f6fc2fb15de 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 @@ -22,7 +22,8 @@ from flashinfer.utils import ceil_div -from ..decode import _dtype_key, _validate_mask, _validate_positive_int +from ..decode import _cutlass_dtype, _dtype_key, _validate_mask, _validate_positive_int +from ..sage import SageAttentionConfig from .common import ( _PREPARED_KV_ROUTE_SIZE, _SIGNED_INT32_MAX, @@ -64,7 +65,7 @@ @dataclass(frozen=True) class _BlockSparseCompileKey: - """Named, hashable inputs that determine one compiled sparse adapter.""" + """Named, hashable inputs that determine one compiled attention adapter.""" device_index: int batch_size: int @@ -77,6 +78,9 @@ class _BlockSparseCompileKey: kv_block_size: int kv_route_size: int dtype_key: str + # ``dtype_key`` names Q and K; a Sage plan may resolve other O and V dtypes. + out_dtype_key: str + v_dtype_key: str mask_type: Literal["dense", "causal"] use_kv_valid_bits: bool use_persistent_scheduler: bool @@ -84,11 +88,15 @@ class _BlockSparseCompileKey: sparse_format: Literal["bsr", "bitmask"] = "bsr" use_proxy_routes: bool = False page_size: int | None = None + share_pattern_across_kv_heads: bool = False + # A dense key ignores the routing fields above. + use_block_sparse: bool = True + sage: SageAttentionConfig | None = None @dataclass(frozen=True) class _BlockSparseLaunchSpec: - """Resolved launch policy and compiler key for one sparse plan.""" + """Resolved launch policy and compile key for one dense or block-sparse plan.""" policy: tuple[tuple[str, object], ...] compile_key: _BlockSparseCompileKey @@ -115,13 +123,19 @@ class _BlockSparseStaticProfile: q_tile_size: int kv_route_size: int dtype_key: str + out_dtype_key: str + v_dtype_key: str q_dtype: torch.dtype kv_dtype: torch.dtype + v_dtype: torch.dtype output_dtype: torch.dtype mask_type: Literal["dense", "causal"] use_kv_valid_bits: bool max_blocks_per_row: int | None page_size: int | None = None + share_pattern_across_kv_heads: bool = False + sage: SageAttentionConfig | None = None + use_block_sparse: bool = True def _select_block_sparse_kv_route_size( @@ -190,6 +204,7 @@ def _select_block_sparse_scheduler( device_index: int, batch_size: int, seq_len_q: int, + seq_len_kv: int, num_qo_heads: int, num_kv_heads: int, q_block_size: int, @@ -198,9 +213,11 @@ def _select_block_sparse_scheduler( mask_type: Literal["dense", "causal"], use_kv_valid_bits: bool, max_row_route_capacity: int, + use_block_sparse: bool, ) -> tuple[int, bool]: """Select the Q tile and scheduler without depending on KV storage. + A dense plan sizes the launch heuristic by the whole K/V sequence. Proxy routes add one summary route per row on top of the exact routes and see the same per-tile fixed cost the persistent scheduler amortizes, so the selection does not depend on the route kind. @@ -212,6 +229,19 @@ def _select_block_sparse_scheduler( heads_q_per_kv=heads_q_per_kv, kv_block_size=kv_block_size, ) + launch_args = { + "device_index": device_index, + "batch_size": batch_size, + "seq_len_q": seq_len_q, + "num_qo_heads": num_qo_heads, + "num_kv_heads": num_kv_heads, + "q_tile_size": q_tile_size, + "kv_route_size": kv_route_size, + } + if not use_block_sparse: + return q_tile_size, _select_persistent_launch( + seq_len_kv=seq_len_kv, **launch_args + ) if not _should_consider_clc( q_tile_size=q_tile_size, kv_block_size=kv_block_size, @@ -221,6 +251,41 @@ def _select_block_sparse_scheduler( ): return q_tile_size, False + return q_tile_size, _select_persistent_launch( + # Each task is sized by its prepared route capacity, not by the K/V + # extent. + seq_len_kv=max_row_route_capacity * kv_route_size, + persistent_min_waves=( + _CAUSAL_CLC_WAVE_THRESHOLD + if mask_type == "causal" + else _BLOCK_SPARSE_CLC_MIN_WAVES + ), + persistent_min_tiles_per_cta=( + _CAUSAL_CLC_MIN_MAX_ROW_ROUTES if mask_type == "causal" else 1 + ), + **launch_args, + ) + + +def _select_persistent_launch( + *, + device_index: int, + batch_size: int, + seq_len_q: int, + seq_len_kv: int, + num_qo_heads: int, + num_kv_heads: int, + q_tile_size: int, + kv_route_size: int, + persistent_min_waves: int = 1, + persistent_min_tiles_per_cta: int = 1, +) -> bool: + """Return whether the decode launch heuristic selects the CLC scheduler. + + ``seq_len_kv`` sizes each task. These plans never split K/V, so any other + outcome selects the static grid. + """ + from ..kernels.fmha_decode.fmha_decode_config import ( _select_auto_launch_mode, make_q_tile_geometry, @@ -228,35 +293,45 @@ def _select_block_sparse_scheduler( q_geometry = make_q_tile_geometry( rows_per_cta=q_tile_size, - heads_q_per_kv=heads_q_per_kv, + heads_q_per_kv=num_qo_heads // num_kv_heads, groups_tokens_heads_q=True, ) - scheduler_kv_capacity_tokens = max_row_route_capacity * kv_route_size with torch.cuda.device(device_index): mode = _select_auto_launch_mode( batch_size=batch_size, num_heads_kv=num_kv_heads, - seq_len_kv=scheduler_kv_capacity_tokens, + seq_len_kv=seq_len_kv, num_q_tiles=q_geometry.num_q_ctas(seq_len_q), tile_size_kv=kv_route_size, - persistent_min_waves=( - _CAUSAL_CLC_WAVE_THRESHOLD - if mask_type == "causal" - else _BLOCK_SPARSE_CLC_MIN_WAVES - ), - persistent_min_tiles_per_cta=( - _CAUSAL_CLC_MIN_MAX_ROW_ROUTES if mask_type == "causal" else 1 - ), + persistent_min_waves=persistent_min_waves, + persistent_min_tiles_per_cta=persistent_min_tiles_per_cta, + ) + return mode == "persistent" + + +def _validate_dense_contiguous_kv_extent( + *, + seq_len_kv: int, + num_kv_heads: int, + head_dim: int, +) -> None: + """Reject a contiguous K/V batch stride the launch ABI cannot express.""" + + if seq_len_kv * num_kv_heads * head_dim > _SIGNED_INT32_MAX: + raise OverflowError( + "dense contiguous K/V batch stride must fit in signed int32" ) - return q_tile_size, mode == "persistent" def _validate_matching_dtypes( q_dtype: torch.dtype, kv_dtype: torch.dtype, + v_dtype: torch.dtype, output_dtype: torch.dtype, ) -> str: - if not (q_dtype == kv_dtype == output_dtype): + """Validate the 16-bit dtype rule of a plan without Sage attention.""" + + if not (q_dtype == kv_dtype == v_dtype == output_dtype): raise ValueError("block-sparse requires matching Q, K/V, and output dtypes") if q_dtype not in _SUPPORTED_DTYPES: raise NotImplementedError( @@ -286,6 +361,47 @@ def _validate_max_blocks_per_row( return max_blocks_per_row +def _plan_dtype_key(dtype: torch.dtype) -> str: + """Return the compile-key name of a plan dtype; Int8 is a Sage Q/K dtype.""" + + return "int8" if dtype == torch.int8 else _dtype_key(dtype) + + +def _resolve_plan_dtypes( + q_dtype: torch.dtype, + kv_dtype: torch.dtype | None, + v_dtype: torch.dtype | None, + output_dtype: torch.dtype | None, + sage: SageAttentionConfig | None, +) -> tuple[torch.dtype, torch.dtype, torch.dtype, str]: + """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. + """ + + if kv_dtype is None: + kv_dtype = q_dtype + if v_dtype is None: + v_dtype = kv_dtype + if sage is None: + if output_dtype is None: + output_dtype = q_dtype + return ( + kv_dtype, + v_dtype, + 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 + return kv_dtype, v_dtype, output_dtype, _plan_dtype_key(q_dtype) + + def _validate_block_sparse_static_profile( *, batch_size: int, @@ -300,12 +416,18 @@ def _validate_block_sparse_static_profile( mask_type: Literal["dense", "causal"], q_dtype: torch.dtype, kv_dtype: torch.dtype | None, + v_dtype: torch.dtype | None = None, output_dtype: torch.dtype | None, max_blocks_per_row: object = _CAPACITY_UNSET, page_size: int | None = None, + share_pattern_across_kv_heads: bool = False, + sage: SageAttentionConfig | None = None, + use_block_sparse: bool = True, ) -> _BlockSparseStaticProfile: """Validate static policy before any device work or BSR inspection.""" + if not isinstance(share_pattern_across_kv_heads, bool): + raise TypeError("share_pattern_across_kv_heads must be a bool") batch_size = _validate_positive_int(batch_size, "batch_size") seq_len_q = _validate_positive_int(seq_len_q, "seq_len_q") seq_len_kv = _validate_positive_int(seq_len_kv, "seq_len_kv") @@ -352,11 +474,9 @@ def _validate_block_sparse_static_profile( raise ValueError("block-sparse requires head_dim=128") if mask_type == "causal" and seq_len_q > seq_len_kv: raise ValueError("causal block-sparse requires seq_len_q <= seq_len_kv") - if kv_dtype is None: - kv_dtype = q_dtype - if output_dtype is None: - output_dtype = q_dtype - dtype_key = _validate_matching_dtypes(q_dtype, kv_dtype, output_dtype) + kv_dtype, v_dtype, output_dtype, dtype_key = _resolve_plan_dtypes( + q_dtype, kv_dtype, v_dtype, output_dtype, sage + ) kv_route_size = _select_block_sparse_kv_route_size( q_tile_size=q_tile_size, kv_block_size=kv_block_size, @@ -384,28 +504,30 @@ def _validate_block_sparse_static_profile( q_tile_size=q_tile_size, kv_route_size=kv_route_size, dtype_key=dtype_key, + out_dtype_key=_plan_dtype_key(output_dtype), + v_dtype_key=_plan_dtype_key(v_dtype), q_dtype=q_dtype, kv_dtype=kv_dtype, + v_dtype=v_dtype, output_dtype=output_dtype, mask_type=mask_type, use_kv_valid_bits=use_kv_valid_bits, max_blocks_per_row=validated_max_blocks_per_row, page_size=page_size, + share_pattern_across_kv_heads=share_pattern_across_kv_heads, + sage=sage, + use_block_sparse=use_block_sparse, ) def _make_block_sparse_config(key: _BlockSparseCompileKey) -> "FmhaDecodeConfig": """Build one decode configuration from its exact compile cache key.""" - import cutlass - from ..kernels.fmha_decode.fmha_decode_config import make_decode_config - dtype_map = { - "float16": cutlass.Float16, - "bfloat16": cutlass.BFloat16, - } - dtype = dtype_map[key.dtype_key] + dtype = _cutlass_dtype(key.dtype_key) + out_dtype = _cutlass_dtype(key.out_dtype_key) + v_dtype = _cutlass_dtype(key.v_dtype_key) q_tile_size = _select_block_sparse_q_tile_size( q_block_size=key.q_block_size, heads_q_per_kv=key.num_qo_heads // key.num_kv_heads, @@ -417,16 +539,35 @@ def _make_block_sparse_config(key: _BlockSparseCompileKey) -> "FmhaDecodeConfig" "tile_size_q": q_tile_size, "tile_size_kv": key.kv_route_size, "groups_tokens_heads_q": True, - "use_block_sparse": True, - "q_block_size": key.q_block_size, - "kv_block_size": key.kv_block_size, - "use_kv_valid_bits": key.use_kv_valid_bits, - "use_parallel_sparse_kv_loads": key.use_parallel_sparse_kv_loads, + "use_block_sparse": key.use_block_sparse, } + if key.use_block_sparse: + config_args.update( + { + "q_block_size": key.q_block_size, + "kv_block_size": key.kv_block_size, + "use_kv_valid_bits": key.use_kv_valid_bits, + "use_parallel_sparse_kv_loads": key.use_parallel_sparse_kv_loads, + "share_pattern_across_kv_heads": key.share_pattern_across_kv_heads, + } + ) + if key.use_proxy_routes: + config_args["use_block_sparse_proxy_routes"] = True if key.use_persistent_scheduler: config_args["use_persistent_scheduler"] = True - if key.use_proxy_routes: - config_args["use_block_sparse_proxy_routes"] = True + if key.sage is not None: + config_args.update( + { + "sage_q_block_size": key.sage.q_block_size, + "sage_k_block_size": key.sage.k_block_size, + "sage_k_summary_block_size": ( + key.sage.k_summary_block_size + if key.use_proxy_routes + else key.sage.k_block_size + ), + "sage_v_mean": key.sage.v_mean, + } + ) layout_args: dict[str, object] if key.page_size is None: layout_args = {"qkv_layout": "contiguousKv"} @@ -443,10 +584,13 @@ def _make_block_sparse_config(key: _BlockSparseCompileKey) -> "FmhaDecodeConfig" batch_size=key.batch_size, num_heads_q=key.num_qo_heads, num_heads_kv=key.num_kv_heads, - qkv_dtype=dtype, - o_dtype=dtype, + q_dtype=dtype, + k_dtype=dtype, + v_dtype=v_dtype, + o_dtype=out_dtype, split_kv_mode="disabled", splits_kv=1, + split_kv=False, mask_type=key.mask_type, auto_tuner=False, **layout_args, @@ -472,6 +616,12 @@ def _resolve_block_sparse_launch_spec( sparse_format: Literal["bsr", "bitmask"] = "bsr", use_proxy_routes: bool = False, page_size: int | None = None, + share_pattern_across_kv_heads: bool = False, + *, + use_block_sparse: bool = True, + out_dtype_key: str, + v_dtype_key: str, + sage: SageAttentionConfig | None = None, ) -> _BlockSparseLaunchSpec: """Resolve and cache one validated static or CLC launch. @@ -486,6 +636,7 @@ def _resolve_block_sparse_launch_spec( device_index=device_index, 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, q_block_size=q_block_size, @@ -494,7 +645,23 @@ def _resolve_block_sparse_launch_spec( mask_type=mask_type, use_kv_valid_bits=use_kv_valid_bits, max_row_route_capacity=max_row_route_capacity, + use_block_sparse=use_block_sparse, ) + if not use_block_sparse: + _validate_dense_contiguous_kv_extent( + seq_len_kv=seq_len_kv, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + ) + + def parallel_loads(persistent: bool) -> bool: + return use_block_sparse and _select_parallel_sparse_kv_loads( + kv_block_size=kv_block_size, + use_kv_valid_bits=use_kv_valid_bits, + max_row_route_capacity=max_row_route_capacity, + use_persistent_scheduler=persistent, + ) + compile_key = _BlockSparseCompileKey( device_index=device_index, batch_size=batch_size, @@ -510,15 +677,15 @@ def _resolve_block_sparse_launch_spec( mask_type=mask_type, use_kv_valid_bits=use_kv_valid_bits, use_persistent_scheduler=use_persistent_scheduler, - use_parallel_sparse_kv_loads=_select_parallel_sparse_kv_loads( - kv_block_size=kv_block_size, - use_kv_valid_bits=use_kv_valid_bits, - max_row_route_capacity=max_row_route_capacity, - use_persistent_scheduler=use_persistent_scheduler, - ), + use_parallel_sparse_kv_loads=parallel_loads(use_persistent_scheduler), sparse_format=sparse_format, use_proxy_routes=use_proxy_routes, page_size=page_size, + use_block_sparse=use_block_sparse, + share_pattern_across_kv_heads=share_pattern_across_kv_heads, + out_dtype_key=out_dtype_key, + v_dtype_key=v_dtype_key, + sage=sage, ) try: config = _make_block_sparse_config(compile_key) @@ -528,15 +695,21 @@ def _resolve_block_sparse_launch_spec( compile_key = replace( compile_key, use_persistent_scheduler=False, - use_parallel_sparse_kv_loads=_select_parallel_sparse_kv_loads( - kv_block_size=kv_block_size, - use_kv_valid_bits=use_kv_valid_bits, - max_row_route_capacity=max_row_route_capacity, - use_persistent_scheduler=False, - ), + use_parallel_sparse_kv_loads=parallel_loads(False), ) config = _make_block_sparse_config(compile_key) + if compile_key.sage is not None: + # Key by the resolved summary block so equivalent recipes share one + # adapter. + compile_key = replace( + compile_key, + sage=replace( + compile_key.sage, + k_summary_block_size=config.sage_k_summary_block_size, + ), + ) + policy_entries: list[tuple[str, object]] = [ ("tile_size_q", q_tile_size), ("tile_size_kv", kv_route_size), @@ -573,8 +746,10 @@ def _resolve_block_sparse_launch_spec( "_resolve_block_sparse_launch_spec", "_select_block_sparse_kv_route_size", "_select_parallel_sparse_kv_loads", + "_select_persistent_launch", "_should_consider_clc", "_validate_block_sparse_static_profile", + "_validate_dense_contiguous_kv_extent", "_validate_matching_dtypes", "_validate_max_blocks_per_row", ] diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/inspection.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/inspection.py index 7abf22e52792..b8f0abe9820c 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/inspection.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/inspection.py @@ -25,6 +25,7 @@ import torch +from .common import _num_sparse_pattern_heads from .config import _BlockSparseStaticProfile @@ -126,7 +127,9 @@ def launch(summary: torch.Tensor, device_index: int) -> None: inspect_bsr = compile_block_sparse_inspection( device_index=device_index, batch_size=static.batch_size, - num_kv_heads=static.num_kv_heads, + num_kv_heads=_num_sparse_pattern_heads( + static.num_kv_heads, static.share_pattern_across_kv_heads + ), seq_len_q=static.seq_len_q, seq_len_kv=static.seq_len_kv, q_block_size=static.q_block_size, @@ -169,7 +172,9 @@ def launch(summary: torch.Tensor, device_index: int) -> None: inspect_metadata = compile_paged_block_sparse_metadata_inspection( device_index=device_index, batch_size=static.batch_size, - num_kv_heads=static.num_kv_heads, + num_kv_heads=_num_sparse_pattern_heads( + static.num_kv_heads, static.share_pattern_across_kv_heads + ), seq_len_q=static.seq_len_q, minimum_seq_len_kv=minimum_seq_len_kv, max_seq_len_kv=static.seq_len_kv, 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 453e8e35134a..4c721004d9d0 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,12 @@ from flashinfer.utils import ceil_div -from .common import _SIGNED_INT32_MAX, _block_sparse_proxy_summary_geometry +from ..sage import SageAttentionConfig, sage_scale_shapes +from .common import ( + _SIGNED_INT32_MAX, + _block_sparse_proxy_summary_geometry, + _num_sparse_pattern_heads, +) from .compiler import _get_compiled_block_sparse from .config import ( _BlockSparseStaticProfile, @@ -35,6 +40,23 @@ _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.""" @@ -62,9 +84,10 @@ def serialized(self: _S, /, *args: _P.args, **kwargs: _P.kwargs) -> _R: class _BlockSparsePlanState: """One complete launch state published by a block-sparse wrapper. - Every state executes one ``prepare -> prepared-route attention`` adapter, - including a pattern whose rows select every KV block. Caller BSR and token - mask tensors belong to individual runs and are never retained here. + Every block-sparse state executes one ``prepare -> prepared-route + attention`` adapter, including a pattern whose rows select every KV block; + a dense state skips route preparation. Caller BSR and token mask tensors + belong to individual runs and are never retained here. Runtime geometry, dtypes, the compiled launch, and the readiness event are published together. ``run()`` therefore sees either the complete old state @@ -85,6 +108,7 @@ class _BlockSparsePlanState: cannot prevent in-place modification or replace graph ownership. """ + use_block_sparse: bool sparse_format: Literal["bsr", "bitmask"] use_proxy_routes: bool device: torch.device @@ -98,19 +122,24 @@ class _BlockSparsePlanState: kv_block_size: int q_dtype: torch.dtype kv_dtype: torch.dtype + v_dtype: torch.dtype output_dtype: torch.dtype use_kv_valid_bits: bool page_size: int | None + sage: SageAttentionConfig | None + # Expected shape of every Sage scale tensor, by field name. + sage_scale_shapes: dict[str, tuple[int, int]] | None - # Only an unmasked specialization needs a shape-correct ABI placeholder. + # Only an unmasked specialization needs a shape-correct ABI placeholder; + # dense plans own neither it nor route storage. dummy_kv_valid_bits: torch.Tensor | None # Immutable row capacities and mutable per-run route payload. - row_route_offsets: torch.Tensor - route_workspace: torch.Tensor + row_route_offsets: torch.Tensor | None + route_workspace: torch.Tensor | None # Semantic row bound; unlike route capacity, this distinguishes # multiple semantic blocks packed into one prepared route. - max_blocks_per_row: int + max_blocks_per_row: int | None policy: tuple[tuple[str, object], ...] compiled: Callable[..., object] @@ -118,6 +147,7 @@ class _BlockSparsePlanState: # All plan-stream work happens-before run after waiting on this event. ready_event: torch.cuda.Event ready_stream_handle: int + share_pattern_across_kv_heads: bool = False def _allocate_dummy_kv_valid_bits( @@ -176,23 +206,22 @@ def _record_block_sparse_plan_ready_event( return event -def _build_block_sparse_plan_state( - static: _BlockSparseStaticProfile, - *, - device: torch.device, - device_index: int, - plan_stream: torch.cuda.Stream, - sparse_format: Literal["bsr", "bitmask"] = "bsr", - use_proxy_routes: bool = False, -) -> _BlockSparsePlanState: - """Build one format- and route-specialized plan atomically.""" +def _block_sparse_route_capacity( + static: _BlockSparseStaticProfile, *, use_proxy_routes: bool +) -> tuple[int, int]: + """Return ``(prepared rows, route capacity per row)`` of one plan. + A dense plan has none; a proxy plan adds its summary routes to every row. + """ + + if not static.use_block_sparse: + return 0, 0 assert static.max_blocks_per_row is not None - if static.page_size is not None: - assert sparse_format == "bsr" and not use_proxy_routes num_rows = ( static.batch_size - * static.num_kv_heads + * _num_sparse_pattern_heads( + static.num_kv_heads, static.share_pattern_across_kv_heads + ) * ceil_div(static.seq_len_q, static.q_block_size) ) if num_rows > _SIGNED_INT32_MAX: @@ -207,6 +236,42 @@ def _build_block_sparse_plan_state( static.kv_block_size, ) max_row_route_capacity += ceil_div(num_summaries, static.kv_route_size) + return num_rows, max_row_route_capacity + + +def _build_block_sparse_plan_state( + static: _BlockSparseStaticProfile, + *, + device: torch.device, + device_index: int, + plan_stream: torch.cuda.Stream, + sparse_format: Literal["bsr", "bitmask"] = "bsr", + use_proxy_routes: bool = False, +) -> _BlockSparsePlanState: + """Build one format- and route-specialized plan atomically.""" + + if static.page_size is not None: + assert sparse_format == "bsr" and not use_proxy_routes + _validate_int8_qk_device(static.q_dtype, device_index) + num_rows, max_row_route_capacity = _block_sparse_route_capacity( + static, use_proxy_routes=use_proxy_routes + ) + scale_shapes = None + if static.sage is not None: + scale_shapes = sage_scale_shapes( + static.sage, + batch_size=static.batch_size, + seq_len_q=static.seq_len_q, + seq_len_kv=static.seq_len_kv, + num_qo_heads=static.num_qo_heads, + num_kv_heads=static.num_kv_heads, + head_dim=static.head_dim, + summary_seq_len=( + ceil_div(static.seq_len_kv, static.kv_block_size) + if use_proxy_routes + else None + ), + ) with torch.cuda.device(device_index), torch.cuda.stream(plan_stream): spec = _resolve_block_sparse_launch_spec( device_index=device_index, @@ -215,6 +280,7 @@ def _build_block_sparse_plan_state( seq_len_kv=static.seq_len_kv, num_qo_heads=static.num_qo_heads, num_kv_heads=static.num_kv_heads, + share_pattern_across_kv_heads=static.share_pattern_across_kv_heads, head_dim=static.head_dim, q_block_size=static.q_block_size, kv_block_size=static.kv_block_size, @@ -226,37 +292,43 @@ def _build_block_sparse_plan_state( max_row_route_capacity=max_row_route_capacity, sparse_format=sparse_format, use_proxy_routes=use_proxy_routes, + use_block_sparse=static.use_block_sparse, + out_dtype_key=static.out_dtype_key, + v_dtype_key=static.v_dtype_key, + sage=static.sage, ) policy = ( *spec.policy, ("max_blocks_per_row", static.max_blocks_per_row), ) - route_layout = _BlockSparseRouteLayout.create( - kv_route_size=static.kv_route_size, - kv_block_size=static.kv_block_size, - page_size=static.page_size, - has_token_bits=spec.prepares_score_words, - route_metadata_capacity=num_rows * max_row_route_capacity, - num_rows=num_rows, - ) compiled = _get_compiled_block_sparse(spec.compile_key) - dummy_kv_valid_bits = ( - None - if static.use_kv_valid_bits - else _allocate_dummy_kv_valid_bits( - batch_size=static.batch_size, - seq_len_kv=static.seq_len_kv, + dummy_kv_valid_bits = None + row_route_offsets = None + route_workspace = None + if static.use_block_sparse: + route_layout = _BlockSparseRouteLayout.create( + kv_route_size=static.kv_route_size, + kv_block_size=static.kv_block_size, + page_size=static.page_size, + has_token_bits=spec.prepares_score_words, + route_metadata_capacity=num_rows * max_row_route_capacity, + num_rows=num_rows, + ) + if not static.use_kv_valid_bits: + dummy_kv_valid_bits = _allocate_dummy_kv_valid_bits( + batch_size=static.batch_size, + seq_len_kv=static.seq_len_kv, + device=device, + ) + row_route_offsets, route_workspace = _allocate_route_storage( device=device, + route_layout=route_layout, + uniform_row_route_capacity=max_row_route_capacity, ) - ) - row_route_offsets, route_workspace = _allocate_route_storage( - device=device, - route_layout=route_layout, - uniform_row_route_capacity=max_row_route_capacity, - ) ready_event = _record_block_sparse_plan_ready_event(plan_stream) return _BlockSparsePlanState( + use_block_sparse=static.use_block_sparse, sparse_format=sparse_format, use_proxy_routes=use_proxy_routes, device=device, @@ -270,9 +342,12 @@ def _build_block_sparse_plan_state( kv_block_size=static.kv_block_size, q_dtype=static.q_dtype, kv_dtype=static.kv_dtype, + v_dtype=static.v_dtype, output_dtype=static.output_dtype, use_kv_valid_bits=static.use_kv_valid_bits, page_size=static.page_size, + sage=static.sage, + sage_scale_shapes=scale_shapes, dummy_kv_valid_bits=dummy_kv_valid_bits, row_route_offsets=row_route_offsets, route_workspace=route_workspace, @@ -281,6 +356,7 @@ def _build_block_sparse_plan_state( compiled=compiled, ready_event=ready_event, ready_stream_handle=plan_stream.cuda_stream, + share_pattern_across_kv_heads=static.share_pattern_across_kv_heads, ) @@ -297,8 +373,10 @@ def _wait_and_record_block_sparse_plan( if state.dummy_kv_valid_bits is not None: state.dummy_kv_valid_bits.record_stream(stream) - state.row_route_offsets.record_stream(stream) - state.route_workspace.record_stream(stream) + if state.row_route_offsets is not None: + state.row_route_offsets.record_stream(stream) + if state.route_workspace is not None: + state.route_workspace.record_stream(stream) __all__ = [ 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 4f24524e3d4a..22027a9c5ef8 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 @@ -20,7 +20,6 @@ import torch -from .._tensor_aliasing import _validate_out_does_not_overlap_inputs from ..decode import ( PagedKVCache, _normalize_paged_kv_cache, @@ -29,7 +28,12 @@ _validate_exact_compact_strides, _validate_scale, ) -from .common import _SIGNED_INT32_MAX +from ..sage import SageAttentionParams, sage_adapter_slots, validate_sage_params +from .common import ( + _SIGNED_INT32_MAX, + _num_sparse_pattern_heads, + _validate_dense_contiguous_plan_inputs, +) if TYPE_CHECKING: from .plan import _BlockSparsePlanState @@ -64,6 +68,19 @@ class _PagedKVLaunchPayload: v_page_stride: int +@dataclass(frozen=True) +class _ValidatedPagedKV: + """Zero-copy paged K/V views, cache geometry, and launch metadata of one run.""" + + k: torch.Tensor + v: torch.Tensor + num_kv_heads: int + page_size: int + head_dim: int + table_capacity: int + payload: _PagedKVLaunchPayload + + @dataclass(frozen=True) class _BlockSparseRunArgs: """Validated launch arguments with optional values materialized once.""" @@ -74,13 +91,16 @@ class _BlockSparseRunArgs: out: torch.Tensor block_indptr: torch.Tensor | None block_indices: torch.Tensor | None - kv_valid_bits: torch.Tensor + # ``None`` for a dense run. + kv_valid_bits: torch.Tensor | None kv_valid_bits_is_live: bool sm_scale: float paged_kv: _PagedKVLaunchPayload | None exact_block_bits: torch.Tensor | None = None k_summary: torch.Tensor | None = None v_summary: torch.Tensor | None = None + # Sage scale tensors in adapter ABI order; ``None`` for unused slots. + sage_slots: tuple[torch.Tensor | None, ...] = sage_adapter_slots({}) def _validate_metadata_tensor( @@ -222,6 +242,75 @@ def _validate_bshd_tensor( _validate_16byte_alignment(tensor, name) +def validate_paged_kv_storage( + kv_storage: _PagedKVStorage, + *, + expected_device: torch.device, + expected_batch_size: int, + seq_len_kv: int, +) -> _ValidatedPagedKV: + """Validate paged K/V structure and page-table metadata for one run. + + The cache is normalized into zero-copy HND views, and the page table and + K/V lengths are checked structurally; both must sit on ``expected_device`` + with one row per request, and every row must cover ``seq_len_kv`` tokens. + Agreement of the cache geometry and dtype with a plan stays with the + caller, because the one-shot entry point derives its plan from the values + returned here. + """ + + ( + k, + v, + num_physical_kv_pages, + num_kv_heads, + page_size, + head_dim, + k_page_stride, + v_page_stride, + ) = _normalize_paged_kv_cache( + kv_storage.paged_kv_cache, + expected_device=expected_device, + ) + metadata_device, metadata_batch_size, table_capacity = ( + _validate_block_table_metadata( + kv_storage.block_tables, + kv_storage.seq_lens_kv, + ) + ) + if metadata_device != expected_device: + raise ValueError( + f"paged-KV metadata must be on {expected_device}, got {metadata_device}" + ) + if metadata_batch_size != expected_batch_size: + raise ValueError( + "seq_lens_kv must have one entry per request: " + f"expected {expected_batch_size}, got {metadata_batch_size}" + ) + if table_capacity * page_size < seq_len_kv: + raise ValueError( + "block_tables must cover the planned K/V capacity: expected at " + f"least {(seq_len_kv + page_size - 1) // page_size} columns, " + f"got {table_capacity}" + ) + return _ValidatedPagedKV( + k=k, + v=v, + num_kv_heads=num_kv_heads, + page_size=page_size, + head_dim=head_dim, + table_capacity=table_capacity, + payload=_PagedKVLaunchPayload( + block_tables=kv_storage.block_tables, + block_table_row_stride=kv_storage.block_tables.stride(0), + seq_lens_kv=kv_storage.seq_lens_kv, + num_physical_kv_pages=num_physical_kv_pages, + k_page_stride=k_page_stride, + v_page_stride=v_page_stride, + ), + ) + + def validate_block_sparse_run( q: torch.Tensor, kv_storage: _ContiguousKVStorage | _PagedKVStorage, @@ -235,6 +324,7 @@ def validate_block_sparse_run( kv_valid_bits: torch.Tensor | None, sm_scale: float | None, out: torch.Tensor | None, + sage: SageAttentionParams | None = None, ) -> _BlockSparseRunArgs: """Validate one run and allocate only an omitted output tensor. @@ -242,51 +332,73 @@ def validate_block_sparse_run( Contiguous K/V use compact ``[B, Skv, Hkv, D]`` directly. Paged K/V are normalized once into zero-copy HND cache views plus launch metadata. An explicit output is returned by identity and may not overlap any live launch - input. ``sm_scale=None`` is materialized as ``1 / sqrt(D)``. + input or plan-owned buffer. Storage overlap is a caller precondition and + is not checked. ``sm_scale=None`` is materialized as ``1 / sqrt(D)``. """ use_proxy_routes = state.use_proxy_routes num_kv_blocks = (state.seq_len_kv + state.kv_block_size - 1) // state.kv_block_size - validate_block_sparse_metadata( - sparse_format=state.sparse_format, - block_indptr=block_indptr, - block_indices=block_indices, - exact_block_bits=exact_block_bits, - kv_valid_bits=kv_valid_bits, - device=state.device, - batch_size=state.batch_size, - seq_len_q=state.seq_len_q, - seq_len_kv=state.seq_len_kv, - num_kv_heads=state.num_kv_heads, - q_block_size=state.q_block_size, - kv_block_size=state.kv_block_size, - use_kv_valid_bits=state.use_kv_valid_bits, - ) - - if use_proxy_routes: - if k_summary is None or v_summary is None: - raise ValueError("K/V summaries are required when proxy routes are enabled") - summary_shape = ( - state.batch_size, - num_kv_blocks, - state.num_kv_heads, - state.head_dim, + if not state.use_block_sparse: + # A dense plan takes no routing inputs. + _validate_dense_contiguous_plan_inputs( + (name, value, None) + for name, value in ( + ("block_indptr", block_indptr), + ("block_indices", block_indices), + ("exact_block_bits", exact_block_bits), + ("k_summary", k_summary), + ("v_summary", v_summary), + ("kv_valid_bits", kv_valid_bits), + ) ) - for tensor, name in ((k_summary, "k_summary"), (v_summary, "v_summary")): - _validate_bshd_tensor( - tensor, - name, - expected_shape=summary_shape, - expected_dtype=state.kv_dtype, - expected_device=state.device, + else: + validate_block_sparse_metadata( + sparse_format=state.sparse_format, + block_indptr=block_indptr, + block_indices=block_indices, + exact_block_bits=exact_block_bits, + kv_valid_bits=kv_valid_bits, + device=state.device, + batch_size=state.batch_size, + seq_len_q=state.seq_len_q, + seq_len_kv=state.seq_len_kv, + num_kv_heads=_num_sparse_pattern_heads( + state.num_kv_heads, state.share_pattern_across_kv_heads + ), + q_block_size=state.q_block_size, + kv_block_size=state.kv_block_size, + use_kv_valid_bits=state.use_kv_valid_bits, + ) + if use_proxy_routes: + if k_summary is None or v_summary is None: + raise ValueError( + "K/V summaries are required when proxy routes are enabled" + ) + summary_shape = ( + state.batch_size, + num_kv_blocks, + state.num_kv_heads, + state.head_dim, ) - elif k_summary is not None or v_summary is not None: - raise ValueError("summaries are valid only when proxy routes are enabled") - + for tensor, name, dtype in ( + (k_summary, "k_summary", state.kv_dtype), + (v_summary, "v_summary", state.v_dtype), + ): + _validate_bshd_tensor( + tensor, + name, + expected_shape=summary_shape, + expected_dtype=dtype, + expected_device=state.device, + ) + elif k_summary is not None or v_summary is not None: + raise ValueError("summaries are valid only when proxy routes are enabled") + + effective_kv_valid_bits: torch.Tensor | None = None if state.use_kv_valid_bits: assert kv_valid_bits is not None effective_kv_valid_bits = kv_valid_bits - else: + elif state.use_block_sparse: effective_kv_valid_bits = state.dummy_kv_valid_bits if effective_kv_valid_bits is None: raise RuntimeError("unmasked block-sparse plan is missing its dummy mask") @@ -300,7 +412,6 @@ def validate_block_sparse_run( expected_device=state.device, ) paged_kv: _PagedKVLaunchPayload | None = None - overlap_inputs: list[tuple[str, torch.Tensor]] if isinstance(kv_storage, _ContiguousKVStorage): if state.page_size is not None: raise TypeError("contiguous K/V storage requires a contiguous plan state") @@ -310,115 +421,59 @@ def validate_block_sparse_run( state.num_kv_heads, state.head_dim, ) - for tensor, name in ((kv_storage.k, "k"), (kv_storage.v, "v")): + for tensor, name, dtype in ( + (kv_storage.k, "k", state.kv_dtype), + (kv_storage.v, "v", state.v_dtype), + ): _validate_bshd_tensor( tensor, name, expected_shape=kv_shape, - expected_dtype=state.kv_dtype, + expected_dtype=dtype, expected_device=state.device, ) k = kv_storage.k v = kv_storage.v - overlap_inputs = [ - ("q", q), - ("k", k), - ("v", v), - ] elif isinstance(kv_storage, _PagedKVStorage): page_size = state.page_size if page_size is None: raise TypeError("paged K/V storage requires a paged plan state") - ( - k, - v, - num_physical_kv_pages, - runtime_num_kv_heads, - runtime_page_size, - runtime_head_dim, - k_page_stride, - v_page_stride, - ) = _normalize_paged_kv_cache( - kv_storage.paged_kv_cache, + paged = validate_paged_kv_storage( + kv_storage, expected_device=state.device, + expected_batch_size=state.batch_size, + seq_len_kv=state.seq_len_kv, ) - runtime_geometry = ( - runtime_num_kv_heads, - runtime_page_size, - runtime_head_dim, - ) - expected_geometry = ( - state.num_kv_heads, - page_size, - state.head_dim, - ) + runtime_geometry = (paged.num_kv_heads, paged.page_size, paged.head_dim) + expected_geometry = (state.num_kv_heads, page_size, state.head_dim) if runtime_geometry != expected_geometry: raise ValueError( "paged_kv_cache geometry does not match the plan: expected " f"Hkv/page/D={expected_geometry}, got {runtime_geometry}" ) - if k.dtype != state.kv_dtype: - raise ValueError( - f"K/V dtype must match the plan ({state.kv_dtype}), got {k.dtype}" - ) - metadata_device, metadata_batch_size, table_capacity = ( - _validate_block_table_metadata( - kv_storage.block_tables, - kv_storage.seq_lens_kv, - ) - ) - if metadata_device != state.device: - raise ValueError( - f"per-run metadata must be on {state.device}, got {metadata_device}" - ) - if metadata_batch_size != state.batch_size: - raise ValueError( - "per-run metadata batch size must match the plan " - f"({state.batch_size}), got {metadata_batch_size}" - ) - if table_capacity * page_size < state.seq_len_kv: + if paged.k.dtype != state.kv_dtype: raise ValueError( - "block_tables must cover the planned K/V capacity: expected at " - f"least {(state.seq_len_kv + page_size - 1) // page_size} columns, " - f"got {table_capacity}" + f"K/V dtype must match the plan ({state.kv_dtype}), got {paged.k.dtype}" ) - paged_kv = _PagedKVLaunchPayload( - block_tables=kv_storage.block_tables, - block_table_row_stride=kv_storage.block_tables.stride(0), - seq_lens_kv=kv_storage.seq_lens_kv, - num_physical_kv_pages=num_physical_kv_pages, - k_page_stride=k_page_stride, - v_page_stride=v_page_stride, - ) - overlap_inputs = [ - ("q", q), - ("k_cache", k), - ("v_cache", v), - ("block_tables", kv_storage.block_tables), - ("seq_lens_kv", kv_storage.seq_lens_kv), - ] + k = paged.k + v = paged.v + paged_kv = paged.payload else: raise TypeError("kv_storage must be _ContiguousKVStorage or _PagedKVStorage") - if block_indptr is not None and block_indices is not None: - overlap_inputs.extend( - (("block_indptr", block_indptr), ("block_indices", block_indices)) - ) - if exact_block_bits is not None: - overlap_inputs.append(("exact_block_bits", exact_block_bits)) - if k_summary is not None and v_summary is not None: - overlap_inputs.extend((("k_summary", k_summary), ("v_summary", v_summary))) - overlap_inputs.extend( - ( - ("kv_valid_bits", effective_kv_valid_bits), - ("row_route_offsets", state.row_route_offsets), - ("route_workspace", state.route_workspace), + 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: + assert state.sage_scale_shapes is not None + validate_sage_params(sage, state.sage_scale_shapes, device=state.device) - effective_scale = _validate_scale( - 1.0 / math.sqrt(state.head_dim) if sm_scale is None else sm_scale, - "sm_scale", + effective_scale = ( + 1.0 / math.sqrt(state.head_dim) + if sm_scale is None + else _validate_scale(sm_scale, "sm_scale") ) if out is None: out = torch.empty(q_shape, device=state.device, dtype=state.output_dtype) @@ -430,7 +485,6 @@ def validate_block_sparse_run( expected_dtype=state.output_dtype, expected_device=state.device, ) - _validate_out_does_not_overlap_inputs(out, *overlap_inputs) return _BlockSparseRunArgs( q=q, k=k, @@ -445,6 +499,7 @@ def validate_block_sparse_run( kv_valid_bits_is_live=state.use_kv_valid_bits, sm_scale=effective_scale, paged_kv=paged_kv, + sage_slots=sage_adapter_slots({} if sage is None else vars(sage)), ) @@ -461,6 +516,7 @@ def prepare_block_sparse_run_unchecked( kv_valid_bits: torch.Tensor | None, sm_scale: float | None, out: torch.Tensor | None, + sage: SageAttentionParams | None = None, ) -> _BlockSparseRunArgs: """Canonicalize one trusted run without invoking explicit validators. @@ -492,7 +548,8 @@ def prepare_block_sparse_run_unchecked( effective_kv_valid_bits = kv_valid_bits else: effective_kv_valid_bits = state.dummy_kv_valid_bits - assert effective_kv_valid_bits is not None + if state.use_block_sparse: + assert effective_kv_valid_bits is not None if out is None: out = torch.empty( (state.batch_size, state.seq_len_q, state.num_qo_heads, state.head_dim), @@ -515,6 +572,7 @@ def prepare_block_sparse_run_unchecked( if sm_scale is None else float(sm_scale), paged_kv=paged_kv, + sage_slots=sage_adapter_slots({} if sage is None else vars(sage)), ) @@ -535,11 +593,14 @@ def record_block_sparse_run_args( run_args.block_indptr.record_stream(stream) assert run_args.block_indices is not None run_args.block_indices.record_stream(stream) - else: - assert run_args.exact_block_bits is not None + elif run_args.exact_block_bits is not None: run_args.exact_block_bits.record_stream(stream) if run_args.kv_valid_bits_is_live: + assert run_args.kv_valid_bits is not None run_args.kv_valid_bits.record_stream(stream) + for tensor in run_args.sage_slots: + if tensor is not None: + tensor.record_stream(stream) if run_args.paged_kv is not None: run_args.paged_kv.block_tables.record_stream(stream) run_args.paged_kv.seq_lens_kv.record_stream(stream) @@ -550,10 +611,11 @@ def launch_block_sparse( *, state: "_BlockSparsePlanState", ) -> torch.Tensor: - """Invoke the layout- and route-specific ABI chosen by the frozen plan.""" + """Invoke the adapter ABI chosen by the frozen plan. + + Contiguous plans share one adapter; unused optional slots are ``None``. + """ - sparse_format = state.sparse_format - use_proxy_routes = state.use_proxy_routes if run_args.paged_kv is not None: assert run_args.block_indptr is not None assert run_args.block_indices is not None @@ -576,41 +638,7 @@ def launch_block_sparse( run_args.paged_kv.v_page_stride, run_args.sm_scale, ) - elif sparse_format == "bsr" and not use_proxy_routes: - assert run_args.block_indptr is not None - assert run_args.block_indices is not None - state.compiled( - run_args.q, - run_args.k, - run_args.v, - run_args.out, - run_args.block_indptr, - run_args.block_indices, - run_args.kv_valid_bits, - state.row_route_offsets, - state.route_workspace, - state.max_blocks_per_row, - run_args.sm_scale, - ) - elif sparse_format == "bitmask" and not use_proxy_routes: - assert run_args.exact_block_bits is not None - state.compiled( - run_args.q, - run_args.k, - run_args.v, - run_args.out, - run_args.exact_block_bits, - run_args.kv_valid_bits, - state.row_route_offsets, - state.route_workspace, - state.max_blocks_per_row, - run_args.sm_scale, - ) - elif sparse_format == "bsr" and use_proxy_routes: - assert run_args.block_indptr is not None - assert run_args.block_indices is not None - assert run_args.k_summary is not None - assert run_args.v_summary is not None + else: state.compiled( run_args.q, run_args.k, @@ -620,32 +648,14 @@ def launch_block_sparse( run_args.out, run_args.block_indptr, run_args.block_indices, - run_args.kv_valid_bits, - state.row_route_offsets, - state.route_workspace, - state.max_blocks_per_row, - run_args.sm_scale, - ) - elif sparse_format == "bitmask" and use_proxy_routes: - assert run_args.exact_block_bits is not None - assert run_args.k_summary is not None - assert run_args.v_summary is not None - state.compiled( - run_args.q, - run_args.k, - run_args.v, - run_args.k_summary, - run_args.v_summary, - run_args.out, run_args.exact_block_bits, run_args.kv_valid_bits, state.row_route_offsets, state.route_workspace, - state.max_blocks_per_row, + 0 if state.max_blocks_per_row is None else state.max_blocks_per_row, + *run_args.sage_slots, run_args.sm_scale, ) - else: - raise AssertionError("frozen block-sparse plan has an unsupported route mode") return run_args.out @@ -654,9 +664,11 @@ def launch_block_sparse( "_ContiguousKVStorage", "_PagedKVLaunchPayload", "_PagedKVStorage", + "_ValidatedPagedKV", "launch_block_sparse", "prepare_block_sparse_run_unchecked", "record_block_sparse_run_args", "validate_block_sparse_metadata", "validate_block_sparse_run", + "validate_paged_kv_storage", ] 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 new file mode 100644 index 000000000000..45f495c955d8 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_q_token_kv_block_sparse_policy.py @@ -0,0 +1,235 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 by FlashInfer team. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Host-only sparse policy shared by group recommendation and launch planning. + +Reuse common FMHA geometry, split and scheduler selection. Explicit G remains +caller-owned; only the recommendation evaluates alternative groups. +""" + +from dataclasses import dataclass + +MAX_GROUP_SIZE = 8 +SUPPORTED_GROUP_SIZES = tuple(range(1, MAX_GROUP_SIZE + 1)) +MAX_TILE_SIZE_Q = 128 +TILE_SIZE_KV = 128 +# Standalone reducer capacity; work and the first SM wave bound actual fanout. +MAX_SPLITS_KV = 128 + + +@dataclass(frozen=True) +class SparseLaunch: + tile_size_q: int + num_insts_kv: int + splits_kv: int + head_dim_per_stage_kv: int + o_stages: int + use_persistent_scheduler: bool + + @property + def use_keeps_mma_ab(self) -> bool: + return self.tile_size_q >= 64 + + +def select_sparse_launch( + *, + group_size: int, + heads_q_per_kv: int, + head_dim: int, + q_dtype_key: str, + num_routes: int, + num_kv_heads: int, + route_kv_tokens: int, + multi_processor_count: int, + split_kv: bool = True, +) -> SparseLaunch: + """Choose the smallest fitting dense Q recipe, then useful one-wave splits. + + G is already fixed by the caller. This function never changes it, queries + the device, reads tensors, or applies a batch/context tuning table. + """ + from .kernels.fmha_decode.fmha_decode_config import ( + MIN_LOOP_ITERS_PER_SPLIT, + _select_auto_launch_mode, + enumerate_grouped_q_mma_candidates, + select_splits_kv, + ) + + # Splits must also cover non-shared KV in the union. Retain the route's + # 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 + # qualified unstaged Swaps recipes. + minimum_q = ( + 64 + if head_dim == 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 + ) + + def recipe(minimum: int): + return next( + c + for c in candidates + if c.tile_size_q >= minimum and c.q_tokens_per_cta >= group_size + ) + + def fanout(instances: int) -> int: + maximum = MAX_SPLITS_KV + if group_size == 1: + # A partial final iteration must not buy an otherwise empty split. + maximum = min(maximum, max(1, work_tokens // (TILE_SIZE_KV * instances))) + return select_splits_kv( + seq_len_kv=work_tokens, + batch_size=num_routes, + num_heads_kv=num_kv_heads, + tile_size_kv=TILE_SIZE_KV, + 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, + split_kv=split_kv, + ) + + candidate = recipe(minimum_q) + instances = 1 if head_dim == 256 and candidate.variant == "keeps_mma_ab" else 2 + splits = fanout(instances) + # Other D256 FP8 groups may use Swaps only when split-KV is active. + if ( + head_dim == 256 + and q_dtype_key == "float8_e4m3fn" + and splits == 1 + and candidate.tile_size_q < 64 + ): + candidate = recipe(64) + instances = 1 + splits = fanout(instances) + return SparseLaunch( + tile_size_q=candidate.tile_size_q, + num_insts_kv=instances, + splits_kv=splits, + head_dim_per_stage_kv=128 if head_dim == 256 else 0, + o_stages=1 if head_dim == 256 and candidate.variant == "keeps_mma_ab" else 2, + 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, + ) + == "persistent" + ), + ) + + +def candidate_union_tokens( + group_size: int, selected_kv_tokens: int, block_size: int = 4 +) -> int: + """Bound a union in semantic blocks; G1 retains its exact tail length.""" + if group_size == 1: + return selected_kv_tokens + return ( + group_size * ((selected_kv_tokens + block_size - 1) // block_size) * block_size + ) + + +def suggest_sparse_group( + *, + batch_size: int, + seq_len_q: int, + selected_kv_tokens: int, + heads_q_per_kv: int, + num_kv_heads: int, + multi_processor_count: int, + head_dim: int, + q_dtype_key: str, + split_kv: bool, + block_size: int = 4, +) -> int: + """Prefer grouping, then splitting; shrink G only for better SM coverage. + + Bounds and dtype validation belong to the public caller. A final short Q + group contributes only its live rows. No SQ % G divisibility is required. + A fully-shared per-query KV proxy estimates useful active splits, rather + than counting every capacity-based split as active. Padding breaks ties + in wave coverage; it must not keep an otherwise under-filled GPU grouped. + TODO: model actual union overlap and scattered loading/metadata costs. + """ + from .kernels.fmha_decode.fmha_decode_config import ( + compute_runtime_active_splits_kv, + enumerate_grouped_q_mma_candidates, + ) + + largest = min(MAX_GROUP_SIZE, seq_len_q, MAX_TILE_SIZE_Q // heads_q_per_kv) + dense_candidates = enumerate_grouped_q_mma_candidates( + heads_q_per_kv=heads_q_per_kv, seq_len_q=largest + ) + # Retain the largest group fitting each dense Q tile. Smaller groups + # sharing that tile only add partially occupied CTAs. At the smallest + # tile, still allow further ungrouping to expose otherwise idle SMs. + groups = { + min(largest, candidate.q_tokens_per_cta) for candidate in dense_candidates + } + smallest_capacity = min( + candidate.q_tokens_per_cta for candidate in dense_candidates + ) + groups.update(range(1, min(largest, smallest_capacity) + 1)) + best_group, best_coverage, best_density = largest, -1.0, -1.0 + for group in sorted(groups, reverse=True): + full_groups, tail = divmod(seq_len_q, group) + routes = batch_size * (full_groups + int(tail > 0)) + length = candidate_union_tokens(group, selected_kv_tokens, block_size) + launch = select_sparse_launch( + group_size=group, + heads_q_per_kv=heads_q_per_kv, + head_dim=head_dim, + q_dtype_key=q_dtype_key, + num_routes=routes, + num_kv_heads=num_kv_heads, + route_kv_tokens=length, + multi_processor_count=multi_processor_count, + split_kv=split_kv, + ) + + # Score active work, not the configured union-capacity split grid. + active_splits = compute_runtime_active_splits_kv( + valid_k=length // group, + tile_size_kv=TILE_SIZE_KV, + num_insts_kv=launch.num_insts_kv, + configured_splits_kv=launch.splits_kv, + ) + per_request = (full_groups + int(tail > 0)) * active_splits + grid = batch_size * num_kv_heads * per_request + waves = (grid + multi_processor_count - 1) // multi_processor_count + coverage = grid / waves + # A short final group still pays for the full Q tile. Prefer less + # padding only when active SM coverage is equal. + live_rows = batch_size * num_kv_heads * seq_len_q * heads_q_per_kv + density = live_rows / (routes * num_kv_heads * launch.tile_size_q) + if coverage > best_coverage or ( + coverage == best_coverage and density > best_density + ): + best_group, best_coverage, best_density = group, coverage, density + if group == largest and grid >= multi_processor_count: + # The grouped Q grid already supplies a full first wave. + break + if coverage == multi_processor_count: + break + return best_group diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_tensor_aliasing.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_tensor_aliasing.py deleted file mode 100644 index ff4c55ede484..000000000000 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/_tensor_aliasing.py +++ /dev/null @@ -1,109 +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. - -"""Storage-alias checks shared by task-scheduled attention launchers.""" - -from typing import Optional - -import torch - - -def _tensor_byte_span(tensor: torch.Tensor) -> tuple[int, int]: - """Return a conservative half-open byte span covering a strided tensor. - - The interval includes stride holes. That deliberate over-approximation - keeps alias rejection safe for paged-cache views whose outer page stride is - larger than one compact page. ``data_ptr`` already includes the tensor's - storage offset, so no storage implementation API is needed. - """ - - if not isinstance(tensor, torch.Tensor): - raise TypeError("tensor must be a torch.Tensor") - if tensor.layout != torch.strided: - raise TypeError("tensor must have strided layout") - - byte_start = tensor.data_ptr() - numel = tensor.numel() - if numel == 0: - return byte_start, byte_start - - element_size = tensor.element_size() - if tensor.is_contiguous(): - return byte_start, byte_start + numel * element_size - - min_element_offset = 0 - max_element_offset = 0 - for extent, stride in zip(tensor.shape, tensor.stride(), strict=True): - last_offset = (int(extent) - 1) * int(stride) - min_element_offset += min(last_offset, 0) - max_element_offset += max(last_offset, 0) - - return ( - byte_start + min_element_offset * element_size, - byte_start + (max_element_offset + 1) * element_size, - ) - - -def _byte_spans_overlap(lhs_span: tuple[int, int], rhs_span: tuple[int, int]) -> bool: - """Return whether two half-open byte spans overlap.""" - - lhs_start, lhs_end = lhs_span - rhs_start, rhs_end = rhs_span - if lhs_start == lhs_end or rhs_start == rhs_end: - return False - return lhs_start < rhs_end and rhs_start < lhs_end - - -def _tensors_overlap(lhs: torch.Tensor, rhs: torch.Tensor) -> bool: - """Return whether conservative storage spans for two tensors overlap.""" - - if not isinstance(lhs, torch.Tensor) or not isinstance(rhs, torch.Tensor): - raise TypeError("overlap operands must be torch.Tensor") - if lhs.device != rhs.device: - return False - return _byte_spans_overlap(_tensor_byte_span(lhs), _tensor_byte_span(rhs)) - - -def _validate_out_does_not_overlap_inputs( - out: torch.Tensor, - *named_inputs: tuple[str, Optional[torch.Tensor]], -) -> None: - """Reject output aliasing with any named live launch input.""" - - _validate_tensor_does_not_overlap_inputs(out, "out", *named_inputs) - - -def _validate_tensor_does_not_overlap_inputs( - checked_tensor: torch.Tensor, - tensor_name: str, - *named_inputs: tuple[str, Optional[torch.Tensor]], -) -> None: - """Reject one named tensor aliasing with any named live launch input.""" - - tensor_span = _tensor_byte_span(checked_tensor) - for name, input_tensor in named_inputs: - if ( - input_tensor is not None - and input_tensor.device == checked_tensor.device - and _byte_spans_overlap(tensor_span, _tensor_byte_span(input_tensor)) - ): - raise ValueError(f"{tensor_name} must not overlap {name} storage") - - -__all__ = [ - "_tensor_byte_span", - "_tensors_overlap", - "_validate_out_does_not_overlap_inputs", - "_validate_tensor_does_not_overlap_inputs", -] 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 e19dc7b1de7f..bb56724e986d 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/block_sparse.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/block_sparse.py @@ -21,6 +21,8 @@ routing tensors, then enqueues route preparation followed by attention in one compiled adapter on the caller's CUDA stream. The one-shot entry point retains synchronous canonical-BSR inspection to derive its temporary plan capacity. +A dense plan (``use_block_sparse=False``) keeps the tile selection but attends +over the whole K/V sequence without route preparation. ``BlockSparsePagedTSWrapper`` shares the same scheduling policy and decode kernel, but plans only fixed Q geometry and maximum K/V capacity. Page spans, @@ -35,10 +37,17 @@ import torch -from flashinfer.api_logging import flashinfer_api +from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api -from ._block_sparse.common import _validate_contiguous_route_mode -from ._block_sparse.config import _validate_block_sparse_static_profile +from ._block_sparse.common import ( + _num_sparse_pattern_heads, + _validate_contiguous_route_mode, + _validate_dense_contiguous_plan_inputs, +) +from ._block_sparse.config import ( + _CAPACITY_UNSET, + _validate_block_sparse_static_profile, +) from ._block_sparse.inspection import ( _inspect_block_sparse_bsr, _inspect_paged_block_sparse_metadata, @@ -58,13 +67,15 @@ record_block_sparse_run_args as _record_block_sparse_run_args, validate_block_sparse_metadata as _validate_block_sparse_metadata, validate_block_sparse_run as _validate_block_sparse_run, + validate_paged_kv_storage as _validate_paged_kv_storage, ) from .decode import ( PagedKVCache, - _normalize_paged_kv_cache, _resolve_cuda_device, - _validate_block_table_metadata, + _resolve_kv_dtypes, + _validate_runtime_device, ) +from .sage import SageAttentionConfig, SageAttentionParams class _BlockSparseWrapperBase: @@ -138,14 +149,19 @@ def plan( kv_block_size: int, *, device: torch.device | str | int, - max_blocks_per_row: int, - use_kv_valid_bits: bool, + use_block_sparse: bool = True, + max_blocks_per_row: int | None = None, + use_kv_valid_bits: bool = False, sparse_format: Literal["bsr", "bitmask"] = "bsr", use_proxy_routes: bool = False, mask_type: Literal["dense", "causal"] = "dense", + share_pattern_across_kv_heads: bool = False, q_data_type: torch.dtype = torch.float16, + k_data_type: torch.dtype | None = None, + v_data_type: torch.dtype | None = None, kv_data_type: torch.dtype | None = None, o_data_type: torch.dtype | None = None, + sage_config: SageAttentionConfig | None = None, ) -> None: """Choose a legal profile and allocate reusable routing capacity. @@ -154,17 +170,25 @@ def plan( semantic ``kv_block_size`` blocks. ``sparse_format="bsr"`` consumes canonical CSR-style rows, while ``"bitmask"`` consumes packed exact- block bits. Enabling proxy routes represents unselected blocks through - caller-provided K/V summaries and currently requires + caller-provided K/V summaries (per-block means) and requires ``mask_type="dense"``. Exact-only plans continue to support causal masking. ``use_kv_valid_bits`` selects whether every :meth:`run` must supply the shared batch token mask. Callers may pass different routing tensor identities and index extents to each run as long as they fit this declared capacity. + ``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. + + ``share_pattern_across_kv_heads=False`` retains independent per-KV-head + patterns. True uses a singleton head axis in BSR/bitmask inputs and + prepares each pattern once. K/V values still have every physical head. + MHA, GQA, and MQA are supported with ``Hq / Hkv`` a power of two no - greater than 32 and ``D=128``. Q, K, V, and O use one matching - ``torch.float16`` or ``torch.bfloat16`` dtype. Runtime tensor shapes - are documented by :meth:`run`. + greater than 32 and ``D=128``. Without Sage attention, Q, K, V, and O + use one matching ``torch.float16`` or ``torch.bfloat16`` dtype. + Runtime tensor shapes are documented by :meth:`run`. ``q_block_size`` may be any positive signed-Int32 value for which a physical Q tile stays within one BSR row. Equivalently, @@ -174,14 +198,12 @@ def plan( multiple of 64. The Q tile groups complete Q-head groups and as many Q tokens as fit without crossing a semantic Q-block row, up to Q128; fine KV blocks cap this at a SWAPAB Q32 tile. Proxy routes reuse the - same Q-tile, KV-route, and MMA geometry as exact routes, but currently - use the direct scheduler because reusable planning cannot observe live - exact-route work. Every run prepares its selected BSR or bitmask into + same Q-tile, KV-route, MMA geometry, and scheduler selection as exact + routes. Every run prepares its selected BSR or bitmask into compact, profile-selected fixed-width route metadata, and the attention core consumes only that metadata. This remains true when every KV block is selected; - callers that know a pattern is dense should choose the dense FMHA API - explicitly. + callers that know a pattern is dense should plan the dense mode. Planning does not inspect routing values and does not synchronize the host. Reusable runs trust those values; assertion-enabled CuTe DSL @@ -194,9 +216,42 @@ def plan( One revision has one mutable route workspace, so its runs must be ordered on one stream or externally synchronized. Unordered concurrent runs require distinct wrappers. + + ``k_data_type`` defaults to ``q_data_type`` and ``v_data_type`` to + ``k_data_type``. ``kv_data_type`` is a compatibility alias that sets + both K and V; explicit K/V dtypes must agree with it. + + ``sage_config`` enables Sage attention on the Q64/KV256 and Q128/KV128 + Keeps profiles: Q and K share ``torch.float8_e4m3fn`` or + ``torch.int8``, V is ``torch.float8_e4m3fn`` (so INT8 plans set + ``v_data_type``), and the output is ``torch.bfloat16`` (default) or + ``torch.float16``. Every :meth:`run` then supplies the scale tensors as + :class:`SageAttentionParams`. """ - _validate_contiguous_route_mode(sparse_format, use_proxy_routes) + if use_block_sparse: + if max_blocks_per_row is None: + raise ValueError( + "max_blocks_per_row is required by a block-sparse plan" + ) + _validate_contiguous_route_mode(sparse_format, use_proxy_routes) + else: + _validate_dense_contiguous_plan_inputs( + ( + ("max_blocks_per_row", max_blocks_per_row, None), + ("use_kv_valid_bits", use_kv_valid_bits, False), + ("sparse_format", sparse_format, "bsr"), + ("use_proxy_routes", use_proxy_routes, False), + ( + "share_pattern_across_kv_heads", + share_pattern_across_kv_heads, + False, + ), + ) + ) + k_data_type, v_data_type = _resolve_kv_dtypes( + q_data_type, k_data_type, v_data_type, kv_data_type + ) static = _validate_block_sparse_static_profile( batch_size=batch_size, seq_len_q=seq_len_q, @@ -207,15 +262,22 @@ def plan( q_block_size=q_block_size, kv_block_size=kv_block_size, use_kv_valid_bits=use_kv_valid_bits, + share_pattern_across_kv_heads=share_pattern_across_kv_heads, mask_type=mask_type, q_dtype=q_data_type, - kv_dtype=kv_data_type, + kv_dtype=k_data_type, + v_dtype=v_data_type, output_dtype=o_data_type, - max_blocks_per_row=max_blocks_per_row, + max_blocks_per_row=( + max_blocks_per_row if use_block_sparse else _CAPACITY_UNSET + ), + sage=sage_config, + use_block_sparse=use_block_sparse, ) if use_proxy_routes and static.mask_type != "dense": raise ValueError("block-sparse proxy routes require mask_type='dense'") device, device_index = _resolve_cuda_device(device) + _validate_runtime_device(device) plan_stream = torch.cuda.current_stream(device) with torch.cuda.device(device_index), torch.cuda.stream(plan_stream): if torch.cuda.is_current_stream_capturing(): @@ -234,7 +296,7 @@ def plan( # previously published revision intact and runnable. self._plan_state = candidate - @flashinfer_api + @flashinfer_experimental_api def run( self, q: torch.Tensor, @@ -250,6 +312,7 @@ def run( sm_scale: float | None = None, out: torch.Tensor | None = None, validate: bool = True, + sage: SageAttentionParams | None = None, ) -> torch.Tensor: """Launch the current plan on the caller's current CUDA stream. @@ -261,12 +324,14 @@ def run( Only O is returned; this PrimTS API does not return LSE. The launch is enqueued asynchronously on the caller's current CUDA stream. - ``validate=True`` performs structural, plan-geometry, and alias - validation without reading tensor values; it is the safe public - default. ``validate=False`` treats every run argument as a trusted - binding and performs no explicit wrapper validation. K/V view - selection, scale forwarding, and optional output allocation are - unavoidable in both modes. + A dense plan rejects every routing argument below. + + ``validate=True`` performs structural and plan-geometry validation + without reading tensor values; it is the safe public default. + ``validate=False`` treats every run argument as a trusted binding and + performs no explicit wrapper validation. K/V view selection, scale + forwarding, and optional output allocation are unavoidable in both + modes. A BSR plan consumes compact Int32 ``block_indptr`` with shape ``[B, Hkv, ceil(Sq / q_block_size) + 1]`` and compact Int32 @@ -276,9 +341,10 @@ def run( Bit ``r`` of word ``w`` selects block ``32 * w + r``; final-word padding bits are ignored. A proxy plan additionally consumes compact ``k_summary`` and ``v_summary`` with shape - ``[B, num_kv_blocks, Hkv, D]``. K summaries are block means and V - summaries are block sums; the final partial block covers only its - structural tokens. + ``[B, num_kv_blocks, Hkv, D]``. K and V summaries are block means; the + final partial block covers only its structural tokens. With Sage + attention, K summaries use the K dtype and ``k_summary_scale``, and V + summaries are E4M3 and share ``v_scale``. Every row must fit the planned semantic-block capacity. Reusable runs trust routing values. CuTe DSL assertions can diagnose violations when @@ -299,11 +365,12 @@ def run( k : torch.Tensor Compact key tensor ``[B, Skv, Hkv, D]`` matching the plan. v : torch.Tensor - Compact value tensor with the same shape, dtype, and strides as - ``k``. + Compact value tensor with the same shape and strides as ``k`` and + the planned V dtype. block_indptr : torch.Tensor, optional Contiguous Int32 BSR row offsets with shape - ``[B, Hkv, ceil(Sq / q_block_size) + 1]``. Required by BSR plans. + ``[B, Hpattern, ceil(Sq / q_block_size) + 1]``, where Hpattern is one + when the plan shares patterns and Hkv otherwise. Required by BSR plans. block_indices : torch.Tensor, optional Contiguous Int32 semantic KV-block IDs referenced by ``block_indptr``. Required by BSR plans. @@ -312,7 +379,7 @@ def run( k_summary : torch.Tensor, optional Per-block mean K tensor required by proxy plans. v_summary : torch.Tensor, optional - Per-block summed V tensor required by proxy plans. + Per-block mean V tensor required by proxy plans. kv_valid_bits : torch.Tensor, optional Contiguous UInt32 token-validity bitmap ``[B, ceil(Skv / 32)]``. Supply it exactly when the plan enabled token validity bits. @@ -320,10 +387,14 @@ def run( Softmax scale. Defaults to ``1 / sqrt(D)``. out : torch.Tensor, optional Caller-owned compact output buffer ``[B, Sq, Hq, D]`` with the - planned output dtype. + planned output dtype. Must not overlap any live input or plan-owned + buffer; storage overlap is not checked. validate : bool - Whether to validate tensor structure, plan geometry, and aliasing - before launching. Defaults to ``True``. + Whether to validate tensor structure and plan geometry before + launching. Defaults to ``True``. + sage : SageAttentionParams, optional + Scale tensors of this run. Required by a Sage plan and rejected + otherwise. Returns ------- @@ -351,12 +422,13 @@ def run( kv_valid_bits=kv_valid_bits, sm_scale=sm_scale, out=out, + sage=sage, ) run_stream = torch.cuda.current_stream(state.device) return self._launch_validated_run(state, run_args, run_stream) -@flashinfer_api +@flashinfer_experimental_api def block_sparse_attention( q: torch.Tensor, k: torch.Tensor, @@ -370,13 +442,17 @@ def block_sparse_attention( k_summary: torch.Tensor | None = None, v_summary: torch.Tensor | None = None, kv_valid_bits: torch.Tensor | None = None, + use_block_sparse: bool = True, sparse_format: Literal["bsr", "bitmask"] = "bsr", use_proxy_routes: bool = False, mask_type: Literal["dense", "causal"] = "dense", + share_pattern_across_kv_heads: bool = False, sm_scale: float | None = None, out: torch.Tensor | None = None, + sage: SageAttentionParams | None = None, + sage_config: SageAttentionConfig | None = None, ) -> torch.Tensor: - """Plan and run one compact-BSHD block-sparse attention launch. + """Plan and run one compact-BSHD block-sparse or dense attention launch. This one-shot form creates a capacity-only plan and passes the original routing tensors to :meth:`BlockSparseTSWrapper.run`. BSR inputs are @@ -392,10 +468,12 @@ def block_sparse_attention( k : torch.Tensor Compact key tensor ``[B, Skv, Hkv, D]``. v : torch.Tensor - Compact value tensor with the same shape, dtype, and strides as ``k``. + Compact value tensor with the same shape and strides as ``k``; its + dtype matches ``k`` except for the INT8 Sage recipe's E4M3 V. block_indptr : torch.Tensor, optional Contiguous Int32 BSR row offsets with shape - ``[B, Hkv, ceil(Sq / q_block_size) + 1]``. Required in BSR mode. + ``[B, Hpattern, ceil(Sq / q_block_size) + 1]``. Hpattern is one for + shared patterns and Hkv otherwise. Required in BSR mode. block_indices : torch.Tensor, optional Contiguous Int32 semantic KV-block IDs referenced by ``block_indptr``. Required in BSR mode. @@ -411,9 +489,12 @@ def block_sparse_attention( k_summary : torch.Tensor, optional Per-block mean K tensor required when proxy routes are enabled. v_summary : torch.Tensor, optional - Per-block summed V tensor required when proxy routes are enabled. + Per-block mean V tensor required when proxy routes are enabled. kv_valid_bits : torch.Tensor, optional Contiguous UInt32 token-validity bitmap ``[B, ceil(Skv / 32)]``. + use_block_sparse : bool, optional + ``False`` runs dense attention over the whole K/V sequence; routing + arguments must then be ``None`` and route options keep their defaults. sparse_format : {"bsr", "bitmask"}, optional Runtime sparse representation. Defaults to ``"bsr"``. use_proxy_routes : bool, optional @@ -425,6 +506,18 @@ def block_sparse_attention( Softmax scale. Defaults to ``1 / sqrt(D)``. out : torch.Tensor, optional Caller-owned compact output buffer ``[B, Sq, Hq, D]``. + Must not overlap any live input; storage overlap is not checked. + Without ``out`` the output takes the Q dtype, or bfloat16 with + ``sage``. + sage : SageAttentionParams, optional + Scale tensors of a Sage launch; see :meth:`BlockSparseTSWrapper.run`. + sage_config : SageAttentionConfig, optional + The Sage recipe to plan. Defaults to :class:`SageAttentionConfig` with + ``v_mean`` set when ``sage.v_mean`` is present. + + share_pattern_across_kv_heads : bool + False (default) uses one pattern per KV head. True requires a singleton + pattern-head axis and reuses those rows across all physical KV heads. Returns ------- @@ -432,6 +525,10 @@ def block_sparse_attention( The compact output tensor; identical to ``out`` when provided. """ + if sage_config is not None and sage is None: + raise ValueError("sage_config requires the sage scale tensors of this call") + if sage is not None and sage_config is None: + sage_config = SageAttentionConfig(v_mean=sage.v_mean is not None) for tensor, name in ((q, "q"), (k, "k"), (v, "v")): if not isinstance(tensor, torch.Tensor): raise TypeError(f"{name} must be a torch.Tensor") @@ -458,42 +555,52 @@ def block_sparse_attention( q_block_size=q_block_size, kv_block_size=kv_block_size, use_kv_valid_bits=use_kv_valid_bits, + share_pattern_across_kv_heads=share_pattern_across_kv_heads, mask_type=mask_type, q_dtype=q.dtype, kv_dtype=k.dtype, - output_dtype=q.dtype if out is None else out.dtype, + v_dtype=v.dtype, + output_dtype=None if out is None else out.dtype, + sage=sage_config, + use_block_sparse=use_block_sparse, ) if use_proxy_routes and static.mask_type != "dense": raise ValueError("block-sparse proxy routes require mask_type='dense'") device, _ = _resolve_cuda_device(q.device) - _validate_block_sparse_metadata( - sparse_format=sparse_format, - block_indptr=block_indptr, - block_indices=block_indices, - exact_block_bits=exact_block_bits, - kv_valid_bits=kv_valid_bits, - device=device, - batch_size=static.batch_size, - seq_len_q=static.seq_len_q, - seq_len_kv=static.seq_len_kv, - num_kv_heads=static.num_kv_heads, - q_block_size=static.q_block_size, - kv_block_size=static.kv_block_size, - use_kv_valid_bits=static.use_kv_valid_bits, - ) - if sparse_format == "bsr": - max_blocks_per_row = _inspect_block_sparse_bsr( - block_indptr, - block_indices, - static=static, - stream=torch.cuda.current_stream(device), + _validate_runtime_device(device) + # A dense plan has no route capacity; its run rejects routing tensors. + max_blocks_per_row = None + if use_block_sparse: + _validate_block_sparse_metadata( + sparse_format=sparse_format, + block_indptr=block_indptr, + block_indices=block_indices, + exact_block_bits=exact_block_bits, + kv_valid_bits=kv_valid_bits, + device=device, + batch_size=static.batch_size, + seq_len_q=static.seq_len_q, + seq_len_kv=static.seq_len_kv, + num_kv_heads=_num_sparse_pattern_heads( + static.num_kv_heads, static.share_pattern_across_kv_heads + ), + q_block_size=static.q_block_size, + kv_block_size=static.kv_block_size, + use_kv_valid_bits=static.use_kv_valid_bits, ) - elif sparse_format == "bitmask": - max_blocks_per_row = ( - static.seq_len_kv + static.kv_block_size - 1 - ) // static.kv_block_size - else: - raise AssertionError(f"unsupported sparse format {sparse_format!r}") + if sparse_format == "bsr": + max_blocks_per_row = _inspect_block_sparse_bsr( + block_indptr, + block_indices, + static=static, + stream=torch.cuda.current_stream(device), + ) + elif sparse_format == "bitmask": + max_blocks_per_row = ( + static.seq_len_kv + static.kv_block_size - 1 + ) // static.kv_block_size + else: + raise AssertionError(f"unsupported sparse format {sparse_format!r}") wrapper = BlockSparseTSWrapper() wrapper.plan( @@ -506,14 +613,18 @@ def block_sparse_attention( static.q_block_size, static.kv_block_size, device=device, + use_block_sparse=use_block_sparse, max_blocks_per_row=max_blocks_per_row, use_kv_valid_bits=static.use_kv_valid_bits, sparse_format=sparse_format, use_proxy_routes=use_proxy_routes, mask_type=static.mask_type, q_data_type=static.q_dtype, - kv_data_type=static.kv_dtype, + k_data_type=static.kv_dtype, + v_data_type=static.v_dtype, o_data_type=static.output_dtype, + share_pattern_across_kv_heads=static.share_pattern_across_kv_heads, + sage_config=sage_config, ) return wrapper.run( q, @@ -527,6 +638,7 @@ def block_sparse_attention( kv_valid_bits=kv_valid_bits, sm_scale=sm_scale, out=out, + sage=sage, ) @@ -550,17 +662,22 @@ def plan( max_blocks_per_row: int, use_kv_valid_bits: bool, mask_type: Literal["dense", "causal"] = "dense", + share_pattern_across_kv_heads: bool = False, q_data_type: torch.dtype = torch.float16, kv_data_type: torch.dtype | None = None, o_data_type: torch.dtype | None = None, ) -> None: """Plan fixed-Q geometry and a maximum variable-K capacity. - The plan stores no request metadata. Every run supplies live page-table - offsets, page IDs, sequence lengths, sparse routes, and optional token - bits. ``max_seq_len_kv`` fixes compilation and mask capacity. Attention - consumes caller-owned live lengths directly, without a plan-owned copy - or device-to-host validation. + ``share_pattern_across_kv_heads=True`` accepts a singleton pattern-head + axis and prepares each BSR row once for all KV heads. False (default) + retains separate patterns. Cache head count and Q layout do not change. + + The plan stores no request metadata. Every run supplies a live fixed 2D + page table of physical page IDs, per-request K/V lengths, sparse routes, + and optional token bits. ``max_seq_len_kv`` fixes compilation and mask + capacity. Attention consumes caller-owned live lengths directly, without + a plan-owned copy or device-to-host validation. ``q_block_size`` may be any positive signed-Int32 value satisfying ``q_block_size * (Hq / Hkv) % 8 == 0``. This row-purity condition keeps @@ -580,6 +697,7 @@ def plan( kv_block_size=kv_block_size, page_size=page_size, use_kv_valid_bits=use_kv_valid_bits, + share_pattern_across_kv_heads=share_pattern_across_kv_heads, mask_type=mask_type, q_dtype=q_data_type, kv_dtype=kv_data_type, @@ -588,6 +706,7 @@ def plan( ) assert static.page_size is not None device, device_index = _resolve_cuda_device(device) + _validate_runtime_device(device) plan_stream = torch.cuda.current_stream(device) with torch.cuda.device(device_index), torch.cuda.stream(plan_stream): if torch.cuda.is_current_stream_capturing(): @@ -602,7 +721,7 @@ def plan( ) self._plan_state = candidate - @flashinfer_api + @flashinfer_experimental_api def run( self, q: torch.Tensor, @@ -624,12 +743,12 @@ def run( tuple whose members are ``[P, Hkv, page, D]`` with compact inner HND strides and arbitrary non-overlapping outer page strides. - ``validate=True`` performs structural, plan-geometry, and alias - validation without reading tensor values; it is the safe public - default. ``validate=False`` treats every run argument as a trusted - binding and performs no explicit wrapper validation. K/V view - selection, scale forwarding, and optional output allocation are - unavoidable in both modes. + ``validate=True`` performs structural and plan-geometry validation + without reading tensor values; it is the safe public default. + ``validate=False`` treats every run argument as a trusted binding and + performs no explicit wrapper validation. K/V view selection, scale + forwarding, and optional output allocation are unavoidable in both + modes. ``block_tables`` is Int32 ``[B, C]``, contiguous within each row but permitted to use a padded outer row stride; ``seq_lens_kv`` is compact @@ -673,7 +792,8 @@ def run( Values must satisfy the dense or causal bounds above. block_indptr : torch.Tensor Contiguous Int32 BSR row offsets with shape - ``[B, Hkv, ceil(Sq / q_block_size) + 1]``. + ``[B, Hpattern, ceil(Sq / q_block_size) + 1]``. Hpattern is one for + shared patterns and Hkv otherwise. block_indices : torch.Tensor Contiguous Int32 logical KV-block IDs referenced by ``block_indptr``. @@ -685,10 +805,11 @@ def run( Softmax scale. Defaults to ``1 / sqrt(D)``. out : torch.Tensor, optional Caller-owned compact output buffer ``[B, Sq, Hq, D]`` with the - planned output dtype. + planned output dtype. Must not overlap any live input or plan-owned + buffer; storage overlap is not checked. validate : bool - Whether to validate tensor structure, plan geometry, and aliasing - before launching. Defaults to ``True``. + Whether to validate tensor structure and plan geometry before + launching. Defaults to ``True``. Returns ------- @@ -722,7 +843,7 @@ def run( return self._launch_validated_run(state, run_args, run_stream) -@flashinfer_api +@flashinfer_experimental_api def block_sparse_attention_with_paged_kv_cache( q: torch.Tensor, paged_kv_cache: PagedKVCache, @@ -736,6 +857,7 @@ def block_sparse_attention_with_paged_kv_cache( max_seq_len_kv: int, kv_valid_bits: torch.Tensor | None = None, mask_type: Literal["dense", "causal"] = "dense", + share_pattern_across_kv_heads: bool = False, sm_scale: float | None = None, out: torch.Tensor | None = None, ) -> torch.Tensor: @@ -763,7 +885,8 @@ def block_sparse_attention_with_paged_kv_cache( Contiguous Int32 per-request logical KV lengths with shape ``[B]``. block_indptr : torch.Tensor Contiguous Int32 BSR row offsets with shape - ``[B, Hkv, ceil(Sq / q_block_size) + 1]``. + ``[B, Hpattern, ceil(Sq / q_block_size) + 1]``. Hpattern is one for + shared patterns and Hkv otherwise. block_indices : torch.Tensor Contiguous Int32 logical KV-block IDs referenced by ``block_indptr``. q_block_size : int @@ -784,6 +907,11 @@ def block_sparse_attention_with_paged_kv_cache( Softmax scale. Defaults to ``1 / sqrt(D)``. out : torch.Tensor, optional Caller-owned compact output buffer ``[B, Sq, Hq, D]``. + Must not overlap any live input; storage overlap is not checked. + + share_pattern_across_kv_heads : bool + False (default) uses one pattern per KV head. True requires a singleton + pattern-head axis and reuses those rows across all physical KV heads. Returns ------- @@ -800,31 +928,21 @@ def block_sparse_attention_with_paged_kv_cache( batch_size, seq_len_q, num_qo_heads, head_dim = map(int, q.shape) metadata_device, _ = _resolve_cuda_device(q.device) - table_device, table_batch_size, table_capacity = _validate_block_table_metadata( - block_tables, seq_lens_kv + _validate_runtime_device(metadata_device) + paged = _validate_paged_kv_storage( + _PagedKVStorage( + paged_kv_cache=paged_kv_cache, + block_tables=block_tables, + seq_lens_kv=seq_lens_kv, + ), + expected_device=q.device, + expected_batch_size=batch_size, + seq_len_kv=max_seq_len_kv, ) - if table_device != q.device: - raise ValueError(f"paged-KV metadata must be on {q.device}, got {table_device}") - if table_batch_size != batch_size: - raise ValueError( - "seq_lens_kv must have one entry per request: " - f"expected {batch_size}, got {table_batch_size}" - ) - - ( - k_cache, - _, - num_physical_kv_pages, - num_kv_heads, - page_size, - runtime_head_dim, - _, - _, - ) = _normalize_paged_kv_cache(paged_kv_cache, expected_device=q.device) - if runtime_head_dim != head_dim: + if paged.head_dim != head_dim: raise ValueError( "Q and paged K/V head dimensions must agree: " - f"expected {head_dim}, got {runtime_head_dim}" + f"expected {head_dim}, got {paged.head_dim}" ) use_kv_valid_bits = kv_valid_bits is not None @@ -833,23 +951,18 @@ def block_sparse_attention_with_paged_kv_cache( seq_len_q=seq_len_q, seq_len_kv=max_seq_len_kv, num_qo_heads=num_qo_heads, - num_kv_heads=num_kv_heads, + num_kv_heads=paged.num_kv_heads, head_dim=head_dim, q_block_size=q_block_size, kv_block_size=kv_block_size, - page_size=page_size, + page_size=paged.page_size, use_kv_valid_bits=use_kv_valid_bits, + share_pattern_across_kv_heads=share_pattern_across_kv_heads, mask_type=mask_type, q_dtype=q.dtype, - kv_dtype=k_cache.dtype, + kv_dtype=paged.k.dtype, output_dtype=q.dtype if out is None else out.dtype, ) - if table_capacity * page_size < static.seq_len_kv: - raise ValueError( - "block_tables must cover the planned K/V capacity: expected at " - f"least {(static.seq_len_kv + page_size - 1) // page_size} columns, " - f"got {table_capacity}" - ) _validate_block_sparse_metadata( sparse_format="bsr", block_indptr=block_indptr, @@ -860,7 +973,9 @@ def block_sparse_attention_with_paged_kv_cache( batch_size=static.batch_size, seq_len_q=static.seq_len_q, seq_len_kv=static.seq_len_kv, - num_kv_heads=static.num_kv_heads, + num_kv_heads=_num_sparse_pattern_heads( + static.num_kv_heads, static.share_pattern_across_kv_heads + ), q_block_size=static.q_block_size, kv_block_size=static.kv_block_size, use_kv_valid_bits=static.use_kv_valid_bits, @@ -871,7 +986,7 @@ def block_sparse_attention_with_paged_kv_cache( block_tables, seq_lens_kv, static=static, - num_physical_kv_pages=num_physical_kv_pages, + num_physical_kv_pages=paged.payload.num_physical_kv_pages, stream=torch.cuda.current_stream(metadata_device), ) @@ -894,6 +1009,7 @@ def block_sparse_attention_with_paged_kv_cache( q_data_type=static.q_dtype, kv_data_type=static.kv_dtype, o_data_type=static.output_dtype, + share_pattern_across_kv_heads=static.share_pattern_across_kv_heads, ) return wrapper.run( q, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/context.py b/tensorrt_llm/_torch/attention/backends/prims_ts/context.py index 678432780332..a8a337373ad1 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/context.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/context.py @@ -1,4 +1,5 @@ # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -31,11 +32,13 @@ position. PrimTS context entry points are intentionally excluded from ``fi_trace`` for -now; their ``@flashinfer_api`` decorators do not register trace templates. +now; unlike the decode APIs, their ``@flashinfer_experimental_api`` decorators do not +register trace templates. """ from dataclasses import dataclass import functools +import importlib.metadata import itertools import math import numbers @@ -44,9 +47,7 @@ import torch -from flashinfer.api_logging import flashinfer_api - -from ._tensor_aliasing import _validate_out_does_not_overlap_inputs +from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api if TYPE_CHECKING: @@ -101,7 +102,8 @@ class _ContextGeometry: num_qo_heads: int num_kv_heads: int head_dim: int - q_dtype: torch.dtype + qk_dtype: torch.dtype + pv_dtype: torch.dtype output_dtype: torch.dtype mask_type: str window_left: int @@ -112,6 +114,7 @@ class _ContextGeometry: packed_dense_k_mask: bool q_shape: tuple[int, ...] kv_shape: tuple[int, ...] + head_dim_vo: int | None = None @dataclass(frozen=True) @@ -127,8 +130,8 @@ class _ContextPlanGeometry: num_qo_heads: int num_kv_heads: int head_dim: int - q_dtype: torch.dtype - kv_dtype: torch.dtype + qk_dtype: torch.dtype + pv_dtype: torch.dtype output_dtype: torch.dtype mask_type: str window_left: int @@ -137,6 +140,7 @@ class _ContextPlanGeometry: has_q_offset: bool causal_single_kv_tile: bool packed_dense_k_mask: bool + head_dim_vo: int | None = None @dataclass(frozen=True) @@ -166,8 +170,8 @@ class _PagedContextPlanGeometry: num_qo_heads: int num_kv_heads: int head_dim: int - q_dtype: torch.dtype - kv_dtype: torch.dtype + qk_dtype: torch.dtype + pv_dtype: torch.dtype output_dtype: torch.dtype mask_type: str window_left: int @@ -215,7 +219,8 @@ class _ContextCompileSpec: num_qo_heads: int num_kv_heads: int head_dim: int - q_dtype_key: str + qk_dtype_key: str + pv_dtype_key: str output_dtype_key: str mask_type: str window_left: int @@ -226,6 +231,8 @@ class _ContextCompileSpec: causal_single_kv_tile: bool packed_dense_k_mask: bool scheduler: _ContextScheduler + head_dim_vo: int | None = None + two_cta_umma: bool = False @dataclass(frozen=True) @@ -239,7 +246,8 @@ class _PagedContextCompileSpec: num_qo_heads: int num_kv_heads: int head_dim: int - q_dtype_key: str + qk_dtype_key: str + pv_dtype_key: str output_dtype_key: str mask_type: str window_left: int @@ -253,9 +261,12 @@ class _PagedContextCompileSpec: def _make_context_kernel( *, - input_dtype, + device_index: int, + input_qk_dtype, + input_pv_dtype, output_dtype, head_dim: int, + head_dim_vo: int | None = None, mask_type: str, window_left: int, head_paired: bool, @@ -264,6 +275,9 @@ def _make_context_kernel( has_q_offset: bool, causal_single_kv_tile: bool, scheduler: _ContextScheduler, + uses_ldtm_stat: bool, + fp8_psmem_early_token: bool = False, + two_cta_umma: bool = False, page_size: int | None = None, max_kv_len: int | None = None, ): @@ -287,13 +301,17 @@ def _make_context_kernel( if use_paged_kv else {} ) - return FmhaTs( + fmha = FmhaTs( qk_acc_dtype=cutlass.Float32, pv_acc_dtype=cutlass.Float32, - in_dtype=input_dtype, + in_qk_dtype=input_qk_dtype, + in_pv_dtype=input_pv_dtype, out_dtype=output_dtype, d=head_dim, + d_v=head_dim_vo, is_persistent=is_persistent, + fp8_psmem_early_token=fp8_psmem_early_token, + two_cta_umma=two_cta_umma, is_causal=mask_type == "causal", has_variable_window=mask_type == "variable_window", balance_causal_workload=( @@ -309,9 +327,11 @@ def _make_context_kernel( window_size_left=window_left if window_left > 0 else 0, h_r=num_qo_heads // num_kv_heads, enable_skip_correction=True, + uses_ldtm_stat=uses_ldtm_stat, causal_single_kv_tile=(causal_single_kv_tile and not use_paged_kv), **paged_kwargs, ) + return fmha def _validate_tensor(tensor: torch.Tensor, name: str) -> None: @@ -360,10 +380,22 @@ def _dtype_key(dtype: torch.dtype) -> str: 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. + """ _dtype_key(q.dtype) - if k.dtype != q.dtype or v.dtype != q.dtype: + if q.dtype != k.dtype: raise NotImplementedError( - "attention-ts context requires Q, K, and V to use the same dtype; " + "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: + raise NotImplementedError( + "attention-ts context requires Q, K, and V to use the same dtype " + "except for QK-BF16/PV-FP8; " f"got Q {q.dtype}, K {k.dtype}, and V {v.dtype}" ) @@ -376,45 +408,65 @@ def _validate_output_dtype(output_dtype: torch.dtype) -> None: def _validate_context_dtype_pair( q_dtype: torch.dtype, - kv_dtype: torch.dtype, + k_dtype: torch.dtype, + v_dtype: torch.dtype, output_dtype: torch.dtype, ) -> None: """Validate the explicit dtypes of a contiguous-context specialization.""" for dtype, name in ( (q_dtype, "q_dtype"), - (kv_dtype, "kv_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) - if kv_dtype != q_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 + ): raise NotImplementedError( - "attention-ts context requires Q and K/V to use the same dtype; " - f"got q_dtype={q_dtype} and kv_dtype={kv_dtype}" + "attention-ts context requires Q, K, and V to use the same dtype " + "except for QK-BF16/PV-FP8; " + f"got q_dtype={q_dtype} and v_dtype={v_dtype}" ) def _validate_paged_dtype_pair( q_dtype: torch.dtype, - kv_dtype: torch.dtype, + k_dtype: torch.dtype, + v_dtype: torch.dtype, output_dtype: torch.dtype, ) -> None: """Validate the explicit dtypes of a paged-context specialization.""" for dtype, name in ( (q_dtype, "q_dtype"), - (kv_dtype, "kv_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) - if kv_dtype != q_dtype: + if k_dtype != q_dtype: raise NotImplementedError( - "attention-ts paged context requires Q and K/V to use the same " - f"dtype; got q_dtype={q_dtype} and kv_dtype={kv_dtype}" + "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 + ): + raise NotImplementedError( + "attention-ts paged context requires Q, K, and V to use the same " + "dtype except for QK-BF16/PV-FP8; " + f"got q_dtype={q_dtype} and v_dtype={v_dtype}" ) @@ -443,6 +495,66 @@ def _validate_device(device: torch.device) -> int: return device_index +@functools.cache +def _dsl_supports_ldtm_stat() -> bool: + """LDTM.STAT uses a native wrapper in DSL 4.8 or inline PTX in 4.7.""" + try: + dsl_version = importlib.metadata.version("nvidia-cutlass-dsl") + except importlib.metadata.PackageNotFoundError: + return False + from packaging import version as pkg_version + + try: + # Use .release so 4.7.0.dev* includes the inline-PTX path. + return pkg_version.Version(dsl_version).release >= (4, 7, 0) + except pkg_version.InvalidVersion: + return False + + +def _default_exp2_fma_pairs(device_index: int, cfg) -> int: + """FMA-pipe exp2 pairs per 16-pair softmax chunk on SM100 with 16-bit V. + + Measured on B200. The paired D128 dense path is MUFU bound and takes 4. + 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): + return 0 + if cfg.v_dtype.width != 16 or cfg.is_causal: + return 0 + if cfg.pv_half_overlap: + return 4 + return 2 if cfg.single_qkv_instance else 0 + + +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_uses_ldtm_stat(device_index: int) -> bool: + """Enable LDTM.STAT on SM103/SM107 with native or compatibility lowering.""" + if not _dsl_supports_ldtm_stat(): + return False + # tcgen05.ld.red.max (LDTM.STAT) is available on B300 (SM103) and Rubin + # (SM107), not B200 (SM100). + capability = torch.cuda.get_device_capability(device_index) + if capability == (10, 3): + return True + if capability == (10, 7): + # DSL 4.7's Rubin escape hatch targets sm_100f, which cannot emit + # LDTM.STAT. Keep Rubin on the native wrapper/target support path. + from cutlass.experimental import primitives as prims + + return hasattr(prims, "tcgen05_ld_red") + return False + + def _resolve_cuda_device( device: int | str | torch.device, ) -> tuple[torch.device, int]: @@ -804,6 +916,8 @@ def _validate_base_tensors( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + *, + allow_mla: bool = False, ) -> None: for tensor, name in ((q, "q"), (k, "k"), (v, "v")): _validate_tensor(tensor, name) @@ -814,7 +928,13 @@ def _validate_base_tensors( f"got {q.device}, {k.device}, and {v.device}" ) _validate_qkv_dtype(q, k, v) - if tuple(v.shape) != tuple(k.shape): + mla_shape = ( + allow_mla + and k.ndim >= 1 + and tuple(v.shape) == (*k.shape[:-1], 128) + and k.shape[-1] == 192 + ) + if tuple(v.shape) != tuple(k.shape) and not mla_shape: raise ValueError( f"v must have the same shape as k; got {tuple(v.shape)} and {tuple(k.shape)}" ) @@ -831,16 +951,25 @@ def _validate_head_geometry(num_qo_heads: int, num_kv_heads: int) -> int: return num_qo_heads // num_kv_heads -def _validate_head_dim(q_head_dim: int, kv_head_dim: int) -> int: +def _validate_head_dim( + q_head_dim: int, kv_head_dim: int, v_head_dim: int | None = None +) -> int: if q_head_dim != kv_head_dim: raise ValueError( - "Q and K/V head dimensions must match; " - f"got Q {q_head_dim} and K/V {kv_head_dim}" + "Q and K head dimensions must match; " + f"got Q {q_head_dim} and K {kv_head_dim}" + ) + if (q_head_dim, v_head_dim) == (192, 128): + return q_head_dim + if v_head_dim is not None and v_head_dim != q_head_dim: + raise NotImplementedError( + "separate QK/V dimensions require contiguous QK/V=(192, 128); " + f"got ({q_head_dim}, {v_head_dim})" ) if q_head_dim not in _SUPPORTED_HEAD_DIMS: raise NotImplementedError( "attention-ts context supports head_dim in " - f"{_SUPPORTED_HEAD_DIMS}; got {q_head_dim}" + f"{_SUPPORTED_HEAD_DIMS}, or contiguous QK/V=(192, 128); got {q_head_dim}" ) return q_head_dim @@ -1015,7 +1144,7 @@ def _resolve_geometry( ) -> _ContextGeometry: """Validate one-shot inputs and derive their static wrapper bounds.""" - _validate_base_tensors(q, k, v) + _validate_base_tensors(q, k, v, allow_mla=True) _validate_output_dtype(output_dtype) _validate_mask(mask_type) window_left = _validate_window_left(window_left, mask_type) @@ -1081,7 +1210,7 @@ def _resolve_geometry( q_shape = tuple(q.shape) kv_shape = tuple(k.shape) - _validate_head_dim(q_head_dim, kv_head_dim) + _validate_head_dim(q_head_dim, kv_head_dim, int(v.shape[-1])) head_ratio = _validate_head_geometry(num_qo_heads, num_kv_heads) if mask_type == "causal": for batch_idx, (q_length, k_length) in enumerate( @@ -1093,6 +1222,8 @@ def _resolve_geometry( f"request; got batch {batch_idx}: Sq={q_length}, Sk={k_length}" ) + if q_head_dim == 192 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): raise NotImplementedError( @@ -1129,7 +1260,9 @@ def _resolve_geometry( num_qo_heads=num_qo_heads, num_kv_heads=num_kv_heads, head_dim=q_head_dim, - q_dtype=q.dtype, + head_dim_vo=int(v.shape[-1]), + qk_dtype=q.dtype, + pv_dtype=v.dtype, output_dtype=output_dtype, mask_type=mask_type, window_left=window_left, @@ -1152,8 +1285,9 @@ def _resolve_context_plan_geometry( num_qo_heads: int, num_kv_heads: int, head_dim: int, - q_dtype: torch.dtype, - kv_dtype: torch.dtype, + head_dim_vo: int | None = None, + qk_dtype: torch.dtype, + pv_dtype: torch.dtype, packed: bool, mask_type: str, window_left: int, @@ -1161,7 +1295,7 @@ def _resolve_context_plan_geometry( ) -> _ContextPlanGeometry: """Validate explicit static bounds for a reusable contiguous plan.""" - _validate_context_dtype_pair(q_dtype, kv_dtype, output_dtype) + _validate_context_dtype_pair(qk_dtype, qk_dtype, pv_dtype, output_dtype) _validate_mask(mask_type) window_left = _validate_window_left(window_left, mask_type) if not isinstance(packed, bool): @@ -1177,6 +1311,11 @@ def _resolve_context_plan_geometry( 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") + head_dim_vo = ( + head_dim + if head_dim_vo is None + else _validate_static_extent(head_dim_vo, "head_dim_vo") + ) device, device_index = _resolve_cuda_device(device) _validate_padded_data_extent( @@ -1184,7 +1323,11 @@ def _resolve_context_plan_geometry( ) _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) + _validate_head_dim(head_dim, head_dim, head_dim_vo) + if head_dim == 192 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: raise ValueError( "bottom-right causal context requires max_seq_len_q <= max_kv_len; " @@ -1210,8 +1353,9 @@ def _resolve_context_plan_geometry( num_qo_heads=num_qo_heads, num_kv_heads=num_kv_heads, head_dim=head_dim, - q_dtype=q_dtype, - kv_dtype=kv_dtype, + head_dim_vo=head_dim_vo, + qk_dtype=qk_dtype, + pv_dtype=pv_dtype, output_dtype=output_dtype, mask_type=mask_type, window_left=window_left, @@ -1312,8 +1456,8 @@ def _resolve_paged_geometry( num_qo_heads=num_qo_heads, num_kv_heads=num_kv_heads, head_dim=q_head_dim, - q_dtype=q.dtype, - kv_dtype=k_cache.dtype, + qk_dtype=q.dtype, + pv_dtype=v_cache.dtype, page_size=page_size, mask_type=mask_type, window_left=window_left, @@ -1351,9 +1495,12 @@ def _make_context_scheduler_probe( torch.float8_e4m3fn: cutlass.Float8E4M3FN, } probe = _make_context_kernel( - input_dtype=dtype_map[geometry.q_dtype], + 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], head_dim=geometry.head_dim, + head_dim_vo=getattr(geometry, "head_dim_vo", None), mask_type=geometry.mask_type, window_left=geometry.window_left, head_paired=geometry.head_paired, @@ -1362,6 +1509,7 @@ 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), page_size=page_size, max_kv_len=max_kv_len, ) @@ -1388,8 +1536,8 @@ def _resolve_paged_plan_geometry( num_qo_heads: int, num_kv_heads: int, head_dim: int, - q_dtype: torch.dtype, - kv_dtype: torch.dtype, + qk_dtype: torch.dtype, + pv_dtype: torch.dtype, page_size: int, mask_type: str, window_left: int, @@ -1400,7 +1548,7 @@ def _resolve_paged_plan_geometry( ) -> _PagedContextPlanGeometry: """Validate explicit static bounds for a reusable paged specialization.""" - _validate_paged_dtype_pair(q_dtype, kv_dtype, output_dtype) + _validate_paged_dtype_pair(qk_dtype, qk_dtype, pv_dtype, output_dtype) _validate_mask(mask_type) if not isinstance(uniform_packed_lengths, bool): raise TypeError("uniform_packed_lengths must be a bool") @@ -1449,8 +1597,8 @@ def _resolve_paged_plan_geometry( num_qo_heads=num_qo_heads, num_kv_heads=num_kv_heads, head_dim=head_dim, - q_dtype=q_dtype, - kv_dtype=kv_dtype, + qk_dtype=qk_dtype, + pv_dtype=pv_dtype, output_dtype=output_dtype, mask_type=mask_type, window_left=window_left, @@ -1465,6 +1613,27 @@ def _resolve_paged_plan_geometry( ) +def _two_cta_umma_geometry_eligible(geometry: _ContextPlanGeometry) -> bool: + """Dense contiguous MHA or GQA at D=128 with bf16 QK runs the two-CTA UMMA + form, which pairs adjacent Q tiles of one head through the grid. + 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.""" + return ( + _default_two_cta_umma(geometry.device_index) + 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.head_paired + and torch.finfo(geometry.qk_dtype).bits in (8, 16) + and not ( + geometry.batch_size > _CUDA_GRID_YZ_MAX + or geometry.num_qo_heads > _CUDA_GRID_YZ_MAX + ) + ) + + def _resolve_context_scheduler(geometry: _ContextPlanGeometry) -> _ContextScheduler: """Select a scheduler from plan-time work while keeping batch out of JIT.""" @@ -1472,6 +1641,9 @@ def _resolve_context_scheduler(geometry: _ContextPlanGeometry) -> _ContextSchedu geometry, causal_single_kv_tile=geometry.causal_single_kv_tile, ) + if _two_cta_umma_geometry_eligible(geometry): + # Two-CTA pairs Q tiles through the grid, so no persistent scheduler. + return "nonpersistent" is_persistent = _contiguous_context_uses_persistent_scheduler( single_qkv_instance=probe.single_qkv_instance, head_paired=geometry.head_paired, @@ -1530,6 +1702,7 @@ def _resolve_paged_context_scheduler( def _context_compile_spec(geometry: _ContextPlanGeometry) -> _ContextCompileSpec: + """Build the contiguous compile key, including separate QK and PV dtype keys.""" return _ContextCompileSpec( device_index=geometry.device_index, max_seq_len_q=geometry.max_seq_len_q, @@ -1537,7 +1710,9 @@ def _context_compile_spec(geometry: _ContextPlanGeometry) -> _ContextCompileSpec num_qo_heads=geometry.num_qo_heads, num_kv_heads=geometry.num_kv_heads, head_dim=geometry.head_dim, - q_dtype_key=_dtype_key(geometry.q_dtype), + 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), mask_type=geometry.mask_type, window_left=geometry.window_left, @@ -1548,12 +1723,14 @@ def _context_compile_spec(geometry: _ContextPlanGeometry) -> _ContextCompileSpec causal_single_kv_tile=geometry.causal_single_kv_tile, packed_dense_k_mask=geometry.packed_dense_k_mask, scheduler=_resolve_context_scheduler(geometry), + two_cta_umma=_two_cta_umma_geometry_eligible(geometry), ) def _paged_context_compile_spec( geometry: _PagedContextPlanGeometry, ) -> _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, @@ -1562,7 +1739,8 @@ def _paged_context_compile_spec( num_qo_heads=geometry.num_qo_heads, num_kv_heads=geometry.num_kv_heads, head_dim=geometry.head_dim, - q_dtype_key=_dtype_key(geometry.q_dtype), + qk_dtype_key=_dtype_key(geometry.qk_dtype), + pv_dtype_key=_dtype_key(geometry.pv_dtype), output_dtype_key=_dtype_key(geometry.output_dtype), mask_type=geometry.mask_type, window_left=geometry.window_left, @@ -1587,7 +1765,9 @@ def _get_compiled_context( num_qo_heads = compile_spec.num_qo_heads num_kv_heads = compile_spec.num_kv_heads head_dim = compile_spec.head_dim - q_dtype_key = compile_spec.q_dtype_key + 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 mask_type = compile_spec.mask_type window_left = compile_spec.window_left @@ -1598,6 +1778,7 @@ def _get_compiled_context( causal_single_kv_tile = compile_spec.causal_single_kv_tile packed_dense_k_mask = compile_spec.packed_dense_k_mask scheduler = compile_spec.scheduler + two_cta_umma = compile_spec.two_cta_umma import cutlass import cutlass.cute as cute @@ -1609,15 +1790,19 @@ def _get_compiled_context( "bfloat16": cutlass.BFloat16, "float8_e4m3fn": cutlass.Float8E4M3FN, } - input_dtype = dtype_map[q_dtype_key] + input_qk_dtype = dtype_map[qk_dtype_key] + input_pv_dtype = dtype_map[pv_dtype_key] output_dtype = dtype_map[output_dtype_key] has_variable_window = mask_type == "variable_window" with torch.cuda.device(device_index): max_active_clusters = int(utils.HardwareInfo().get_max_active_clusters(1)) fmha = _make_context_kernel( - input_dtype=input_dtype, + device_index=device_index, + input_qk_dtype=input_qk_dtype, + input_pv_dtype=input_pv_dtype, output_dtype=output_dtype, head_dim=head_dim, + head_dim_vo=head_dim_vo, mask_type=mask_type, window_left=window_left, head_paired=head_paired, @@ -1626,6 +1811,9 @@ def _get_compiled_context( has_q_offset=has_q_offset, 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), + fp8_psmem_early_token=_default_fp8_psmem_early_token(device_index), ) fmha.cfg.has_varlen = packed fmha.cfg.has_uniform_varlen = uniform_packed_lengths @@ -1633,6 +1821,7 @@ def _get_compiled_context( fmha.cfg.uniform_seq_len_q = max_seq_len_q 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) if fmha.cfg.kv_tile_n != _CONTEXT_KV_TILE_N: raise RuntimeError( "context packed-K specialization assumes kv_tile_n=" @@ -1715,19 +1904,19 @@ def fake_compact(dtype, shape, assumed_align): runtime_num_k_offsets = cute.sym_int() q_shape = (runtime_total_q, num_qo_heads, head_dim) kv_shape = (runtime_total_k, num_kv_heads, head_dim) - out_shape = (runtime_total_q, num_qo_heads, head_dim) + out_shape = (runtime_total_q, num_qo_heads, head_dim_vo) qo_indptr_shape = (runtime_num_q_offsets,) 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) - out_shape = q_shape + out_shape = (*q_shape[:-1], head_dim_vo) qo_indptr_shape = (1,) kv_indptr_shape = (1,) - q_fake = fake_compact(input_dtype, q_shape, 16) - k_fake = fake_compact(input_dtype, kv_shape, 16) - v_fake = fake_compact(input_dtype, kv_shape, 16) + 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) @@ -1797,7 +1986,8 @@ def _get_compiled_paged_context( num_qo_heads = compile_spec.num_qo_heads num_kv_heads = compile_spec.num_kv_heads head_dim = compile_spec.head_dim - q_dtype_key = compile_spec.q_dtype_key + qk_dtype_key = compile_spec.qk_dtype_key + pv_dtype_key = compile_spec.pv_dtype_key output_dtype_key = compile_spec.output_dtype_key mask_type = compile_spec.mask_type window_left = compile_spec.window_left @@ -1818,12 +2008,15 @@ def _get_compiled_paged_context( "bfloat16": cutlass.BFloat16, "float8_e4m3fn": cutlass.Float8E4M3FN, } - input_dtype = dtype_map[q_dtype_key] + input_qk_dtype = dtype_map[qk_dtype_key] + input_pv_dtype = dtype_map[pv_dtype_key] output_dtype = dtype_map[output_dtype_key] with torch.cuda.device(device_index): max_active_clusters = int(utils.HardwareInfo().get_max_active_clusters(1)) fmha = _make_context_kernel( - input_dtype=input_dtype, + device_index=device_index, + input_qk_dtype=input_qk_dtype, + input_pv_dtype=input_pv_dtype, output_dtype=output_dtype, head_dim=head_dim, mask_type=mask_type, @@ -1834,6 +2027,8 @@ 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), + fp8_psmem_early_token=_default_fp8_psmem_early_token(device_index), page_size=page_size, max_kv_len=max_kv_len, ) @@ -1843,6 +2038,7 @@ def _get_compiled_paged_context( fmha.cfg.uniform_seq_len_q = max_seq_len_q 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.paged_v_tail_is_zero = paged_v_tail_is_zero if fmha.cfg.kv_tile_n != _CONTEXT_KV_TILE_N: raise RuntimeError( @@ -1896,15 +2092,15 @@ def fake_compact(dtype, shape, assumed_align): runtime_num_pages = cute.sym_int() batch_size = cute.sym_int() runtime_num_q_offsets = cute.sym_int() - q_fake = fake_compact(input_dtype, (runtime_total_q, num_qo_heads, head_dim), 16) + q_fake = fake_compact(input_qk_dtype, (runtime_total_q, num_qo_heads, head_dim), 16) kv_shape = ( runtime_num_pages, num_kv_heads, page_size, head_dim, ) - k_fake = fake_compact(input_dtype, kv_shape, 16) - v_fake = fake_compact(input_dtype, kv_shape, 16) + 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) @@ -1963,13 +2159,18 @@ def _validate_runtime_inputs( ) -> None: """Validate contiguous data and request metadata against one plan.""" - _validate_base_tensors(q, k, v) + _validate_base_tensors(q, k, v, allow_mla=True) if q.device != geometry.device: raise ValueError(f"q must be on {geometry.device}, got {q.device}") - if q.dtype != geometry.q_dtype: - raise ValueError(f"q must have dtype {geometry.q_dtype}, got {q.dtype}") - if k.dtype != geometry.kv_dtype: - raise ValueError(f"k and v must have dtype {geometry.kv_dtype}, got {k.dtype}") + if q.dtype != geometry.qk_dtype: + raise ValueError(f"q must have dtype {geometry.qk_dtype}, got {q.dtype}") + if k.dtype != geometry.qk_dtype: + raise ValueError(f"k must have dtype {geometry.qk_dtype}, got {k.dtype}") + if v.dtype != geometry.pv_dtype: + raise ValueError(f"v must have dtype {geometry.pv_dtype}, got {v.dtype}") + v_shape = (*k.shape[:-1], geometry.head_dim_vo or geometry.head_dim) + if tuple(v.shape) != v_shape: + raise ValueError(f"v must have shape {v_shape}, got {tuple(v.shape)}") if geometry.packed: if q.ndim != 3 or tuple(q.shape[1:]) != ( @@ -2094,11 +2295,15 @@ def _validate_paged_runtime_inputs( _validate_base_tensors(q, k_cache, v_cache) if q.device != geometry.device: raise ValueError(f"q must be on {geometry.device}, got {q.device}") - if q.dtype != geometry.q_dtype: - raise ValueError(f"q must have dtype {geometry.q_dtype}, got {q.dtype}") - if k_cache.dtype != geometry.kv_dtype: + if q.dtype != geometry.qk_dtype: + raise ValueError(f"q must have dtype {geometry.qk_dtype}, got {q.dtype}") + if k_cache.dtype != geometry.qk_dtype: + raise ValueError( + f"k_cache must have dtype {geometry.qk_dtype}, got {k_cache.dtype}" + ) + if v_cache.dtype != geometry.pv_dtype: raise ValueError( - f"k_cache must have dtype {geometry.kv_dtype}, got {k_cache.dtype}" + f"v_cache must have dtype {geometry.pv_dtype}, got {v_cache.dtype}" ) if q.ndim != 3 or tuple(q.shape[1:]) != ( geometry.num_qo_heads, @@ -2252,14 +2457,14 @@ def _prepare_out( *, q: torch.Tensor, output_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(tuple(q.shape), dtype=output_dtype, device=q.device) + return torch.empty(out_shape, dtype=output_dtype, device=q.device) _validate_tensor(out, "out") - if tuple(out.shape) != tuple(q.shape): - raise ValueError( - f"out must have shape {tuple(q.shape)}, got {tuple(out.shape)}" - ) + 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.device != q.device: @@ -2273,6 +2478,10 @@ def _prepare_out( 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 @@ -2294,12 +2503,12 @@ class BatchPrefillTSWrapper: dependency. """ - @flashinfer_api + @flashinfer_experimental_api def __init__(self) -> None: """Initialize an unplanned task-scheduled context-attention wrapper.""" self._plan_state: Optional[_ContextPlanState] = None - @flashinfer_api + @flashinfer_experimental_api def plan( self, *, @@ -2310,8 +2519,10 @@ def plan( num_qo_heads: int, num_kv_heads: int, head_dim: int, + head_dim_vo: int | None = None, q_dtype: torch.dtype, - kv_dtype: torch.dtype, + k_dtype: torch.dtype, + v_dtype: Optional[torch.dtype] = None, out_dtype: Optional[torch.dtype] = None, packed: bool = False, mask_type: Literal["dense", "causal", "variable_window"] = "dense", @@ -2344,11 +2555,18 @@ def plan( num_kv_heads : int Number of key/value heads. head_dim : int - Query, key, value, and output head dimension. + 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. - kv_dtype : torch.dtype - Key/value dtype. It must currently equal ``q_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 @@ -2365,6 +2583,12 @@ def plan( """ 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: + 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}" + ) geometry = _resolve_context_plan_geometry( device=device, batch_size=batch_size, @@ -2373,8 +2597,9 @@ def plan( num_qo_heads=num_qo_heads, num_kv_heads=num_kv_heads, head_dim=head_dim, - q_dtype=q_dtype, - kv_dtype=kv_dtype, + head_dim_vo=head_dim_vo, + qk_dtype=k_dtype, + pv_dtype=resolved_v_dtype, packed=packed, mask_type=mask_type, window_left=window_left, @@ -2431,7 +2656,7 @@ def plan( policy=policy, ) - @flashinfer_api + @flashinfer_experimental_api def run( self, q: torch.Tensor, @@ -2474,9 +2699,10 @@ def run( output_scale : torch.Tensor, optional Per-run one-element float32 output scale. Defaults to the plan. validate : bool - Validate the complete runtime contract. Disable only when the - caller guarantees dtype, device, shape, stride, alignment, values, - aliasing, and lifetime. Defaults to ``True``. + 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``. """ state = self._plan_state @@ -2572,46 +2798,18 @@ def run( device=geometry.device, ) - caller_provided_out = out is not None if out is None: out = torch.empty( - tuple(q.shape), dtype=geometry.output_dtype, device=q.device + (*q.shape[:-1], geometry.head_dim_vo or geometry.head_dim), + dtype=geometry.output_dtype, + device=q.device, ) elif validate: - out = _prepare_out(out, q=q, output_dtype=geometry.output_dtype) - if validate and caller_provided_out: - alias_inputs = [ - ("q", q), - ("k", k), - ("v", v), - ("scale_softmax_log2", scale_softmax_log2), - ("output_scale", output_scale), - ] - if geometry.packed: - alias_inputs.extend( - (("qo_indptr", runtime_qo_indptr), ("kv_indptr", runtime_kv_indptr)) - ) - if geometry.mask_type == "variable_window": - alias_inputs.extend( - ( - ("variable_window_token_starts", runtime_window_starts), - ("variable_window_token_ends", runtime_window_ends), - ("variable_window_cta_starts", runtime_window_cta_starts), - ) - ) - if ( - variable_window_cta_starts is None - and state.variable_window_padded_starts is not None - ): - alias_inputs.append( - ( - "variable_window_padded_starts", - state.variable_window_padded_starts, - ) - ) - _validate_out_does_not_overlap_inputs( + out = _prepare_out( out, - *alias_inputs, + q=q, + output_dtype=geometry.output_dtype, + head_dim_vo=geometry.head_dim_vo, ) state.compiled( q, @@ -2657,7 +2855,7 @@ class BatchPrefillPagedTSWrapper: planning stream, the caller must establish that dependency. """ - @flashinfer_api + @flashinfer_experimental_api def __init__(self, kv_layout: Literal["HND"] = "HND") -> None: """Create an unplanned paged-context wrapper. @@ -2671,7 +2869,7 @@ def __init__(self, kv_layout: Literal["HND"] = "HND") -> None: self._kv_layout = kv_layout self._plan_state: Optional[_PagedContextPlanState] = None - @flashinfer_api + @flashinfer_experimental_api def plan( self, *, @@ -2683,7 +2881,8 @@ def plan( num_kv_heads: int, head_dim: int, q_dtype: torch.dtype, - kv_dtype: torch.dtype, + k_dtype: torch.dtype, + v_dtype: Optional[torch.dtype] = None, out_dtype: Optional[torch.dtype] = None, page_size: int = _DEFAULT_PAGED_KV_PAGE_SIZE, mask_type: Literal["dense", "causal"] = "dense", @@ -2737,8 +2936,12 @@ def plan( Query, key, value, and output head dimension. q_dtype : torch.dtype Query dtype. - kv_dtype : torch.dtype - Key/value cache dtype. It must currently equal ``q_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 @@ -2766,6 +2969,12 @@ def plan( """ 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: + 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}" + ) geometry = _resolve_paged_plan_geometry( device=device, batch_size=batch_size, @@ -2774,8 +2983,8 @@ def plan( num_qo_heads=num_qo_heads, num_kv_heads=num_kv_heads, head_dim=head_dim, - q_dtype=q_dtype, - kv_dtype=kv_dtype, + qk_dtype=k_dtype, + pv_dtype=resolved_v_dtype, page_size=page_size, mask_type=mask_type, window_left=window_left, @@ -2812,7 +3021,7 @@ def plan( policy=policy, ) - @flashinfer_api + @flashinfer_experimental_api def run( self, q: torch.Tensor, @@ -2882,8 +3091,10 @@ def run( retained by the plan. validate : bool Run explicit storage, shape, dtype, device, scale, output, and - aliasing validators. Disable only when the caller guarantees the - complete runtime contract. Defaults to ``True``. + 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 ------- @@ -2922,25 +3133,12 @@ def run( output_scale, "output_scale", device=geometry.device ) - caller_provided_out = out is not None if out is None: out = torch.empty( tuple(q.shape), dtype=geometry.output_dtype, device=q.device ) elif validate: out = _prepare_out(out, q=q, output_dtype=geometry.output_dtype) - if validate and caller_provided_out: - _validate_out_does_not_overlap_inputs( - out, - ("q", q), - ("k_cache", k_cache), - ("v_cache", v_cache), - ("qo_indptr", qo_indptr), - ("block_tables", block_tables), - ("seq_lens_kv", seq_lens_kv), - ("scale_softmax_log2", scale_softmax_log2), - ("output_scale", output_scale), - ) state.compiled( q, k_cache, @@ -2955,7 +3153,7 @@ def run( return out -@flashinfer_api +@flashinfer_experimental_api def batch_prefill( q: torch.Tensor, k: torch.Tensor, @@ -2976,7 +3174,8 @@ def batch_prefill( Fixed tensors use ``[B, S, H, D]`` storage. Providing both cumulative int32 offset tensors selects packed ``[total_tokens, H, D]`` storage. - ``D`` may be 128 or 256. + Equal Q/K/V head dimensions may be 128 or 256. 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 and retains at most ``window_left + 1`` keys at each causal row, including @@ -2985,7 +3184,9 @@ def batch_prefill( Parameters ---------- q, k, v : torch.Tensor - Fixed or packed query, key, and value tensors. + Fixed or packed query, key, and value tensors. Q and K share one + dtype; V matches it or is ``torch.float8_e4m3fn`` with + ``torch.bfloat16`` Q/K. qo_indptr, kv_indptr : torch.Tensor, optional Cumulative query and K/V offsets for packed-ragged input. mask_type : {"dense", "causal", "variable_window"} @@ -3044,8 +3245,10 @@ def batch_prefill( num_qo_heads=geometry.num_qo_heads, num_kv_heads=geometry.num_kv_heads, head_dim=geometry.head_dim, - q_dtype=geometry.q_dtype, - kv_dtype=geometry.q_dtype, + head_dim_vo=geometry.head_dim_vo, + q_dtype=geometry.qk_dtype, + k_dtype=geometry.qk_dtype, + v_dtype=geometry.pv_dtype, out_dtype=geometry.output_dtype, packed=geometry.packed, mask_type=mask_type, @@ -3065,7 +3268,7 @@ def batch_prefill( ) -@flashinfer_api +@flashinfer_experimental_api def batch_prefill_with_paged_kv_cache( q: torch.Tensor, k_cache: torch.Tensor, @@ -3090,7 +3293,8 @@ def batch_prefill_with_paged_kv_cache( ``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 - 256; Q, K, and V must share one supported dtype. + 256. Q and K must share one supported dtype; V matches it or is + ``torch.float8_e4m3fn`` alongside ``torch.bfloat16`` Q/K. This convenience API reads request metadata on the host to derive exact plan bounds and is not CUDA-graph-capturable. Capture-sensitive callers @@ -3169,8 +3373,9 @@ def batch_prefill_with_paged_kv_cache( num_qo_heads=geometry.num_qo_heads, num_kv_heads=geometry.num_kv_heads, head_dim=geometry.head_dim, - q_dtype=geometry.q_dtype, - kv_dtype=geometry.kv_dtype, + q_dtype=geometry.qk_dtype, + k_dtype=geometry.qk_dtype, + v_dtype=geometry.pv_dtype, out_dtype=geometry.output_dtype, page_size=page_size, mask_type=mask_type, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py b/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py index 35236fcc7ce9..3fcec5efdc2a 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/decode.py @@ -21,25 +21,23 @@ import math import numbers import struct -from typing import TYPE_CHECKING, Literal, Optional, Union +from typing import TYPE_CHECKING, Literal, NamedTuple, Optional, Union, cast import torch -from flashinfer.api_logging import flashinfer_api - -from ._tensor_aliasing import ( - _validate_out_does_not_overlap_inputs, - _validate_tensor_does_not_overlap_inputs, -) +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 PagedKVCache = Union[torch.Tensor, tuple[torch.Tensor, torch.Tensor]] +PagedKVScaleFactors = tuple[torch.Tensor, torch.Tensor] if TYPE_CHECKING: from .kernels.fmha_decode.fmha_decode_config import FmhaDecodeConfig _SUPPORTED_HEAD_DIMS = (64, 128, 256) -_SUPPORTED_PAGE_SIZES = (16, 32, 64, 128) +_SUPPORTED_PAGE_SIZES = (4, 8, 16, 32, 64, 128) _MAX_INT32 = 2**31 - 1 # Decode K/V masks form an exclusive tile endpoint as # ``tile_offset_k + tile_size_kv`` in signed Int32. Public policies use at @@ -55,6 +53,7 @@ _COMPILE_OPTIONS = "--enable-tvm-ffi --opt-level 3" _WORKSPACE_ALIGNMENT = 256 _WORKSPACE_DTYPES = (torch.int8, torch.uint8) +_MAX_HEAD_RATIO = 128 @dataclass(frozen=True) @@ -76,6 +75,7 @@ class _DecodeWorkspaceLayout: split_kv_counter: _WorkspaceSection cu_seqlens_q: _WorkspaceSection attention_sinks: _WorkspaceSection + uses_split_kv: bool total_bytes: int @@ -100,6 +100,16 @@ class _DecodeLaunchSpec: scratch_shapes: tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]] +class _DirectQ1CompileSpec(NamedTuple): + """Static raw-index geometry; tensor values and strides remain runtime inputs.""" + + position_dtype_key: str + max_seq_len_kv: int + block_topk: int + kv_block_size: int + page_capacity: int + + @dataclass(frozen=True) class _DecodeCompileSpec: """Batch-independent static identity for one compiled decode callable.""" @@ -113,11 +123,14 @@ class _DecodeCompileSpec: 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 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 @dataclass(frozen=True) @@ -127,10 +140,16 @@ class _DecodeRuntime: q: torch.Tensor k_cache: torch.Tensor v_cache: torch.Tensor + k_sf_cache: torch.Tensor + v_sf_cache: torch.Tensor out: torch.Tensor num_physical_pages: int k_page_stride: int + k_head_stride: int + k_token_stride: int v_page_stride: int + v_head_stride: int + v_token_stride: int bmm1_scale: float bmm2_scale: float @@ -154,7 +173,8 @@ class _DecodePlanState: page_size: int max_kv_len: int q_dtype: torch.dtype - kv_dtype: torch.dtype + k_dtype: torch.dtype + v_dtype: torch.dtype output_dtype: torch.dtype mask_type: str window_left: int @@ -169,6 +189,7 @@ class _DecodePlanState: 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 def _decode_policy_from_config( @@ -201,6 +222,7 @@ def _decode_policy_from_config( ), ("tile_size_q", int(config.tile_size_q)), ("tile_size_kv", int(config.tile_size_kv)), + ("num_insts_kv", int(config.num_insts_kv)), ("use_split_kv", bool(config.use_split_kv)), ("splits_kv", int(config.splits_kv)), ("max_splits_kv", int(config.max_splits_kv)), @@ -340,6 +362,7 @@ def _make_decode_workspace_layout( output_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 @@ -369,6 +392,7 @@ def _make_decode_workspace_layout( split_kv_counter=split_kv_counter, cu_seqlens_q=cu_seqlens_q, attention_sinks=attention_sinks, + uses_split_kv=use_split_kv, total_bytes=_align_up(byte_end), ) @@ -605,9 +629,10 @@ def _validate_head_geometry(num_qo_heads: int, num_kv_heads: int) -> None: f"{num_qo_heads} and {num_kv_heads}" ) head_ratio = num_qo_heads // num_kv_heads - if head_ratio > 32: + if head_ratio > _MAX_HEAD_RATIO: raise ValueError( - f"attention-ts decode requires 1 <= Hq/Hkv <= 32, got {head_ratio}" + "attention-ts decode requires " + f"1 <= Hq/Hkv <= {_MAX_HEAD_RATIO}, got {head_ratio}" ) @@ -638,55 +663,118 @@ def _dtype_key(dtype: torch.dtype) -> str: torch.float16: "float16", torch.bfloat16: "bfloat16", torch.float8_e4m3fn: "float8_e4m3fn", + torch.uint8: "float4_e2m1fn", } try: return keys[dtype] except KeyError as error: raise NotImplementedError( "attention-ts decode supports torch.float16, torch.bfloat16, " - f"and torch.float8_e4m3fn; got {dtype}" + "torch.float8_e4m3fn, and packed NVFP4 K/V in torch.uint8; " + f"got {dtype}" ) from error +def _cutlass_dtype(dtype_key: str): + """Return the cutlass numeric type named by a dtype key.""" + + import cutlass + + return { + "float16": cutlass.Float16, + "bfloat16": cutlass.BFloat16, + "float8_e4m3fn": cutlass.Float8E4M3FN, + "float4_e2m1fn": cutlass.Float4E2M1FN, + "int8": cutlass.Int8, + }[dtype_key] + + def _validate_dtype_pair( q_dtype: torch.dtype, - kv_dtype: torch.dtype, + k_dtype: torch.dtype, + v_dtype: torch.dtype, output_dtype: torch.dtype, + *, + allow_fp8_bf16_output: bool = False, ) -> None: _dtype_key(q_dtype) - _dtype_key(kv_dtype) + _dtype_key(k_dtype) + _dtype_key(v_dtype) _dtype_key(output_dtype) - if q_dtype != kv_dtype: - raise NotImplementedError( - "attention-ts decode requires Q, K, and V to use the same dtype; " - f"got Q {q_dtype} and K/V {kv_dtype}" - ) + matching_kv = k_dtype == v_dtype supported = ( - (q_dtype == torch.float16 and output_dtype == torch.float16) - or (q_dtype == torch.bfloat16 and output_dtype == torch.bfloat16) + ( + q_dtype == torch.float16 + and k_dtype == v_dtype == torch.float16 + and output_dtype == torch.float16 + ) + or ( + q_dtype == torch.bfloat16 + and output_dtype == torch.bfloat16 + and ( + ( + matching_kv + and k_dtype in (torch.bfloat16, torch.float8_e4m3fn, torch.uint8) + ) + or (k_dtype == torch.bfloat16 and v_dtype == torch.float8_e4m3fn) + ) + ) or ( q_dtype == torch.float8_e4m3fn - and output_dtype in (torch.float16, torch.float8_e4m3fn) + and matching_kv + and k_dtype in (torch.float8_e4m3fn, torch.uint8) + and output_dtype + in ( + (torch.float16, torch.bfloat16, torch.float8_e4m3fn) + if allow_fp8_bf16_output + else (torch.float16, torch.float8_e4m3fn) + ) ) ) if not supported: raise NotImplementedError( "attention-ts decode supports FP16->FP16, BF16->BF16, " - "FP8-E4M3->FP16, and FP8-E4M3->FP8-E4M3; got " - f"{q_dtype}->{output_dtype}" + + ( + "FP8-E4M3->FP16/BF16, and FP8-E4M3->FP8-E4M3; " + if allow_fp8_bf16_output + else "FP8-E4M3->FP16, and FP8-E4M3->FP8-E4M3; " + ) + + "BF16 Q + FP8 K/V, BF16/FP8-E4M3 Q + NVFP4 K/V, " + + "and BF16 Q/K + FP8 V; got " + + f"Q_{q_dtype}_K_{k_dtype}_V_{v_dtype}->{output_dtype}" ) +def _resolve_kv_dtypes( + q_dtype: torch.dtype, + k_dtype: Optional[torch.dtype], + v_dtype: Optional[torch.dtype], + kv_dtype: Optional[torch.dtype], +) -> tuple[torch.dtype, torch.dtype]: + """Resolve the common-KV compatibility alias without hiding disagreement.""" + if kv_dtype is not None: + if (k_dtype is not None and k_dtype != kv_dtype) or ( + v_dtype is not None and v_dtype != kv_dtype + ): + raise ValueError("explicit K/V dtypes must agree with the common KV dtype") + k_dtype = v_dtype = kv_dtype + if k_dtype is None: + k_dtype = q_dtype + if v_dtype is None: + v_dtype = k_dtype + return k_dtype, v_dtype + + def _device_index(device: torch.device) -> int: if device.index is not None: return int(device.index) return int(torch.cuda.current_device()) -def _validate_runtime_device(device: torch.device) -> int: +def _validate_runtime_device(device: torch.device) -> None: if device.type != "cuda": raise ValueError("attention-ts decode tensors must be CUDA tensors") - device_index = _device_index(device) + device_index = device.index with torch.cuda.device(device_index): capability = torch.cuda.get_device_capability(device_index) if capability not in _SUPPORTED_COMPUTE_CAPABILITIES: @@ -701,12 +789,13 @@ def _validate_runtime_device(device: torch.device) -> int: from flashinfer.cute_dsl.utils import require_cute_dsl_arch require_cute_dsl_arch(device_index) - return device_index def _resolve_cuda_device( device: Optional[Union[int, str, torch.device]], ) -> tuple[torch.device, int]: + """Normalize a device specifier without checking CUDA support.""" + if device is None: resolved = torch.device("cuda", torch.cuda.current_device()) elif isinstance(device, int) and not isinstance(device, bool): @@ -718,8 +807,11 @@ def _resolve_cuda_device( raise TypeError("device must identify one CUDA device") from error if resolved.type == "cuda" and resolved.index is None: resolved = torch.device("cuda", torch.cuda.current_device()) - device_index = _validate_runtime_device(resolved) - return torch.device("cuda", device_index), device_index + # Preserve non-CUDA device types for the caller's explicit validation. + device_index = ( + _device_index(resolved) if resolved.type == "cuda" else (resolved.index or 0) + ) + return torch.device(resolved.type, device_index), device_index def _validate_q( @@ -887,66 +979,74 @@ def _normalize_paged_kv_cache( paged_kv_cache: PagedKVCache, *, expected_device: torch.device, + logical_head_dim: Optional[int] = None, ) -> tuple[torch.Tensor, torch.Tensor, int, int, int, int, int, int]: - """Return zero-copy HND K/V views and their runtime geometry.""" + """Return compact HND views for block-sparse and legacy paged paths.""" - if isinstance(paged_kv_cache, torch.Tensor): - if paged_kv_cache.ndim != 5 or paged_kv_cache.shape[1] != 2: + views = _normalize_paged_kv_cache_views( + paged_kv_cache, + expected_device=expected_device, + ) + k_cache, v_cache = views[:2] + inferred_head_dim = views[5] + if logical_head_dim is not None and inferred_head_dim != logical_head_dim: + raise ValueError( + "paged_kv_cache storage head dimension does not match the logical " + f"head dimension {logical_head_dim}: got {int(k_cache.shape[-1])}" + ) + k_page_stride = _validate_hnd_inner_strides(k_cache, "K cache") + v_page_stride = _validate_hnd_inner_strides(v_cache, "V cache") + return (*views, k_page_stride, v_page_stride) + + +def _normalize_paged_kv_scale_factors( + kv_scale_factors: Optional[PagedKVScaleFactors], + *, + k_cache: torch.Tensor, + logical_head_dim: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Validate NVFP4 scale tensors or create homogeneous-mode placeholders.""" + + if k_cache.dtype != torch.uint8: + if kv_scale_factors is not None: raise ValueError( - "combined paged_kv_cache must have shape " - "[num_pages, 2, Hkv, page_size, head_dim]" + "kv_scale_factors are accepted only with packed NVFP4 torch.uint8 K/V" ) - if paged_kv_cache.device != expected_device: + placeholder = k_cache[0, 0, 0, :1].view(torch.uint8)[:1] + return placeholder, placeholder + if kv_scale_factors is None: + raise ValueError( + "packed NVFP4 torch.uint8 K/V requires kv_scale_factors=(K_SF, V_SF)" + ) + if not isinstance(kv_scale_factors, tuple) or len(kv_scale_factors) != 2: + raise TypeError("kv_scale_factors must be a (K_SF, V_SF) tensor tuple") + k_sf_cache, v_sf_cache = kv_scale_factors + if not isinstance(k_sf_cache, torch.Tensor) or not isinstance( + v_sf_cache, torch.Tensor + ): + raise TypeError("kv_scale_factors tuple members must be torch.Tensor") + expected_shape = (*k_cache.shape[:-1], logical_head_dim // 16) + for scale, name in ( + (k_sf_cache, "K scale factors"), + (v_sf_cache, "V scale factors"), + ): + if tuple(scale.shape) != expected_shape: raise ValueError( - f"paged_kv_cache must be on {expected_device}, " - f"got {paged_kv_cache.device}" + f"{name} must have shape {expected_shape}, got {tuple(scale.shape)}" ) - k_cache = paged_kv_cache[:, 0] - v_cache = paged_kv_cache[:, 1] - elif isinstance(paged_kv_cache, tuple): - if len(paged_kv_cache) != 2: - raise ValueError("paged_kv_cache tuple must contain exactly (K, V)") - k_cache, v_cache = paged_kv_cache - if not isinstance(k_cache, torch.Tensor) or not isinstance( - v_cache, torch.Tensor - ): - raise TypeError("paged_kv_cache tuple members must be torch.Tensor") - if k_cache.ndim != 4 or v_cache.ndim != 4: + if scale.dtype != torch.float8_e4m3fn: raise ValueError( - "tuple K/V caches must each have shape " - "[num_pages, Hkv, page_size, head_dim]" + f"{name} must have dtype torch.float8_e4m3fn, got {scale.dtype}" ) - if k_cache.device != expected_device or v_cache.device != expected_device: - raise ValueError(f"tuple K/V caches must be on {expected_device}") - else: - raise TypeError( - "paged_kv_cache must be a combined torch.Tensor or a (K, V) tuple" + if scale.device != k_cache.device: + raise ValueError(f"{name} must be on {k_cache.device}, got {scale.device}") + _validate_exact_compact_strides( + scale, + name, + "[pages, Hkv, page_size, D/16]", ) - - if k_cache.ndim != 4 or v_cache.ndim != 4: - raise ValueError("K/V cache views must be rank-4 HND tensors") - if k_cache.shape != v_cache.shape: - raise ValueError("K and V cache views must have identical logical shapes") - if k_cache.dtype != v_cache.dtype: - raise ValueError("K and V cache views must have identical dtypes") - if k_cache.device != v_cache.device: - raise ValueError("K and V cache views must be on the same device") - - num_pages, num_kv_heads, page_size, head_dim = map(int, k_cache.shape) - if min(num_pages, num_kv_heads, page_size, head_dim) <= 0: - raise ValueError("paged_kv_cache dimensions must be positive") - k_page_stride = _validate_hnd_inner_strides(k_cache, "K cache") - v_page_stride = _validate_hnd_inner_strides(v_cache, "V cache") - return ( - k_cache, - v_cache, - num_pages, - num_kv_heads, - page_size, - head_dim, - k_page_stride, - v_page_stride, - ) + _validate_16byte_alignment(scale, name) + return k_sf_cache, v_sf_cache def _validate_block_tables( @@ -1019,69 +1119,6 @@ def _validate_block_table_metadata( return seq_lens.device, batch_size, table_capacity -def _csr_to_block_tables( - paged_kv_indices: torch.Tensor, - indptr: tuple[int, ...], - seq_lens: tuple[int, ...], - *, - page_size: int, -) -> torch.Tensor: - """Materialize canonical CSR page IDs as a native fixed page table. - - The caller has already synchronized to materialize and validate CSR row - offsets and logical lengths. Equal-width rows are exposed as a zero-copy - view only when their actual CSR extents equal the native table capacity. - Otherwise a temporary dense table copies each active prefix from its true - CSR row start. Its inactive tail is deliberately invalid; the native kernel - must bound every access by ``seq_lens``. - """ - - if paged_kv_indices.ndim != 1: - raise ValueError("CSR page IDs must be one-dimensional") - if paged_kv_indices.dtype != torch.int32: - raise TypeError("CSR page IDs must have dtype torch.int32") - if not paged_kv_indices.is_contiguous(): - raise ValueError("CSR page IDs must be contiguous") - page_counts = tuple((seq_len + page_size - 1) // page_size for seq_len in seq_lens) - batch_size = len(page_counts) - if len(indptr) != batch_size + 1: - raise ValueError("CSR indptr must contain one more offset than rows") - if indptr[0] != 0: - raise ValueError("CSR indptr must start at zero") - if any(end < begin for begin, end in zip(indptr[:-1], indptr[1:], strict=True)): - raise ValueError("CSR indptr offsets must be nondecreasing") - if indptr[-1] != int(paged_kv_indices.numel()): - raise ValueError("the final CSR indptr offset must equal the page-ID count") - table_capacity = max(page_counts) - csr_page_counts = tuple( - end - begin for begin, end in zip(indptr[:-1], indptr[1:], strict=True) - ) - if any( - active_page_count > csr_page_count - for active_page_count, csr_page_count in zip( - page_counts, csr_page_counts, strict=True - ) - ): - raise ValueError("CSR rows do not contain every active K/V page") - if all(csr_page_count == table_capacity for csr_page_count in csr_page_counts): - return paged_kv_indices.view(batch_size, table_capacity) - - block_tables = torch.full( - (batch_size, table_capacity), - -1, - dtype=torch.int32, - device=paged_kv_indices.device, - ) - for request_idx, (page_begin, page_count) in enumerate( - zip(indptr[:-1], page_counts, strict=True) - ): - page_end = page_begin + page_count - block_tables[request_idx, :page_count].copy_( - paged_kv_indices[page_begin:page_end] - ) - return block_tables - - def _decode_output_shape( *, batch_size: int, @@ -1208,12 +1245,18 @@ def _resolve_decode_launch_spec( max_kv_len: int, seq_len_q: int, q_dtype_key: str, - kv_dtype_key: str, + k_dtype_key: str, + v_dtype_key: str, output_dtype_key: str, kv_layout: str, mask_type: str, use_packed_q: 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, ) -> _DecodeLaunchSpec: """Resolve automatic policy and workspace geometry without compiling.""" @@ -1226,10 +1269,13 @@ def _resolve_decode_launch_spec( ) max_kv_len = _validate_max_kv_len(max_kv_len, "max_kv_len") window_left = _validate_window_left(window_left, mask_type) - - import cutlass + storage_page_size = _validate_storage_page_size( + page_size, + page_size if storage_page_size is None else storage_page_size, + ) from .kernels.fmha_decode.fmha_decode_config import ( + MIN_LOOP_ITERS_PER_SPLIT, get_max_active_clusters_for_cluster_size, make_decode_config, make_q_tile_geometry, @@ -1237,17 +1283,19 @@ def _resolve_decode_launch_spec( if kv_layout != "HND": raise ValueError("the cached TS decode compiler accepts HND only") - if q_dtype_key != kv_dtype_key: - raise ValueError("the cached TS decode compiler requires one QKV dtype") - dtype_map = { - "float16": cutlass.Float16, - "bfloat16": cutlass.BFloat16, - "float8_e4m3fn": cutlass.Float8E4M3FN, - } - qkv_dtype = dtype_map[q_dtype_key] - output_dtype = dtype_map[output_dtype_key] + 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) - def make_config(args: object | None = None) -> "FmhaDecodeConfig": + def make_config( + args: object | None = None, + *, + split_kv_mode: str = "disabled", + splits_kv: int = -1, + max_splits_kv: int | None = None, + min_loop_iters_per_split: int = MIN_LOOP_ITERS_PER_SPLIT, + ) -> "FmhaDecodeConfig": return make_decode_config( headdim=head_dim, args=args, @@ -1256,15 +1304,22 @@ def make_config(args: object | None = None) -> "FmhaDecodeConfig": batch_size=batch_size, num_heads_q=num_qo_heads, num_heads_kv=num_kv_heads, - qkv_dtype=qkv_dtype, + q_dtype=q_dtype, + k_dtype=k_dtype, + v_dtype=v_dtype, o_dtype=output_dtype, qkv_layout="pagedKv", num_tokens_per_page=page_size, - split_kv_mode="disabled", + storage_tokens_per_page=storage_page_size, + split_kv_mode=split_kv_mode, + splits_kv=splits_kv, + max_splits_kv=max_splits_kv, + min_loop_iters_per_split=min_loop_iters_per_split, mask_type=mask_type, sliding_window_causal=window_left >= 0, attention_window_size=window_left + 1 if window_left >= 0 else 0, auto_tuner=True, + split_kv=split_kv, ) def q_ctas(config: "FmhaDecodeConfig") -> int: @@ -1290,55 +1345,115 @@ def fits_one_service_wave(config: "FmhaDecodeConfig", num_q_ctas: int) -> bool: # Device capacity participates in automatic selection. Resolve it in the # target device context without introducing caller-visible policy knobs. with torch.cuda.device(device_index): - cfg = make_config({"use_variable_seqlens_q": True} if use_packed_q else None) max_active_clusters = get_max_active_clusters_for_cluster_size(1) - - # A grouped fixed-Q launch can leave most of the first service wave - # idle. In that regime, evaluate the narrowest supported Swaps head - # band. Keep it only when the extra head-band CTAs fit in the same - # resident wave without reducing KV fanout or changing the launch and - # reduction topology. The grouped cost selector excludes SQ1, so both - # configs below remain KV128; this legacy Q8-only adjustment cannot - # participate in or override automatic KV256 promotion. - if ( - seq_len_q == 1 - and not use_packed_q - and cfg.groups_tokens_heads_q - and not cfg.use_keeps_mma_ab - ): - grouped_q_ctas = q_ctas(cfg) - head_band_geometry = make_q_tile_geometry( - rows_per_cta=8, - heads_q_per_kv=num_qo_heads // num_kv_heads, - groups_tokens_heads_q=False, + if use_q_token_kv_block_sparse_route: + cfg = _resolve_q_token_kv_block_sparse_decode_config( + make_config, + batch_size=batch_size, + 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, + 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, + mask_type=mask_type, + use_packed_q=use_packed_q, + window_left=window_left, + max_active_clusters=max_active_clusters, + use_pdl=use_pdl, + split_kv=split_kv, + share_pattern_across_kv_heads=share_pattern_across_kv_heads, ) - head_band_q_ctas = max(int(head_band_geometry.num_q_ctas(seq_len_q)), 1) - head_band_cfg = None - if head_band_q_ctas > grouped_q_ctas and fits_one_service_wave( - cfg, head_band_q_ctas + else: + config_overrides = {"use_pdl": use_pdl} + if use_packed_q: + config_overrides["use_variable_seqlens_q"] = True + cfg = make_config(config_overrides) + + # A grouped fixed-Q launch can leave most of the first service wave + # idle. In that regime, evaluate the narrowest supported Swaps head + # band. Keep it only when the extra head-band CTAs fit in the same + # resident wave without reducing KV fanout or changing topology. + if ( + seq_len_q == 1 + and not use_packed_q + and cfg.groups_tokens_heads_q + and not cfg.use_keeps_mma_ab ): - try: - head_band_cfg = make_config( - {"groups_tokens_heads_q": False, "tile_size_q": 8} + grouped_q_ctas = q_ctas(cfg) + head_band_geometry = make_q_tile_geometry( + rows_per_cta=8, + heads_q_per_kv=num_qo_heads // num_kv_heads, + groups_tokens_heads_q=False, + ) + head_band_q_ctas = max(int(head_band_geometry.num_q_ctas(seq_len_q)), 1) + head_band_cfg = None + if head_band_q_ctas > grouped_q_ctas and fits_one_service_wave( + cfg, head_band_q_ctas + ): + try: + head_band_cfg = make_config( + { + "groups_tokens_heads_q": False, + "tile_size_q": 8, + "use_pdl": use_pdl, + } + ) + except ValueError: + head_band_cfg = None + if head_band_cfg is not None: + same_launch_topology = all( + getattr(cfg, field) == getattr(head_band_cfg, field) + for field in ( + "use_split_kv", + "splits_kv", + "max_splits_kv", + "use_cluster_smem_reduction", + "use_separate_reduction_kernel", + "use_persistent_scheduler", + ) ) - except ValueError: - head_band_cfg = None - if head_band_cfg is not None: - same_launch_topology = all( - getattr(cfg, field) == getattr(head_band_cfg, field) - for field in ( - "use_split_kv", - "splits_kv", - "max_splits_kv", - "use_cluster_smem_reduction", - "use_separate_reduction_kernel", - "use_persistent_scheduler", + if same_launch_topology and fits_one_service_wave( + head_band_cfg, q_ctas(head_band_cfg) + ): + cfg = head_band_cfg + + q_token_kv_block_sparse_scattered_fp8 = ( + page_size == 4 + and storage_page_size > page_size + and q_dtype_key == "float8_e4m3fn" + ) + unsafe_direct_swaps = not cfg.use_split_kv and not cfg.use_keeps_mma_ab + unsafe_split_publisher = ( + cfg.use_split_kv and not cfg.use_separate_reduction_kernel + ) + if q_token_kv_block_sparse_scattered_fp8 and ( + unsafe_direct_swaps or unsafe_split_publisher + ): + if not split_kv: + raise NotImplementedError( + "this FP8 encoded-subpage profile requires split-KV; " + "no qualified direct recipe is available with split_kv=False" ) + # Direct, fused-GMEM, and cluster-SMEM publication are not + # qualified with encoded subpage locators on the FP8 pipeline. + safe_splits = int(cfg.splits_kv) if cfg.use_split_kv else 2 + safe_tile_size_q = max(int(cfg.tile_size_q), 16) + cfg = make_config( + { + "groups_tokens_heads_q": cfg.groups_tokens_heads_q, + "tile_size_q": safe_tile_size_q, + "use_variable_seqlens_q": use_packed_q, + "use_pdl": use_pdl, + }, + split_kv_mode="gmem_reduction_with_separate_kernel", + splits_kv=safe_splits, + max_splits_kv=safe_splits, ) - if same_launch_topology and fits_one_service_wave( - head_band_cfg, q_ctas(head_band_cfg) - ): - cfg = head_band_cfg _validate_decode_policy_kv_tile_size(cfg) return _decode_launch_spec_from_config( @@ -1363,10 +1478,13 @@ def _make_decode_compile_spec( 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"], + direct_q1_spec: Optional[_DirectQ1CompileSpec] = None, ) -> _DecodeCompileSpec: """Freeze the resolved topology while leaving batch in runtime tensors.""" @@ -1380,11 +1498,14 @@ def _make_decode_compile_spec( 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, 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, ) @@ -1398,10 +1519,11 @@ 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 - page_size = compile_spec.page_size 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 max_active_clusters = compile_spec.max_active_clusters @@ -1421,15 +1543,18 @@ 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 - dtype_map = { - "float16": cutlass.Float16, - "bfloat16": cutlass.BFloat16, - "float8_e4m3fn": cutlass.Float8E4M3FN, - } - qkv_dtype = dtype_map[q_dtype_key] - output_dtype = dtype_map[output_dtype_key] + 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) 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 ( cutlass.BFloat16, @@ -1448,19 +1573,27 @@ def main_tensor_adapter( q: cute.Tensor, k_cache: cute.Tensor, v_cache: cute.Tensor, + k_sf_cache: cute.Tensor, + v_sf_cache: cute.Tensor, out: cute.Tensor, seq_lens: cute.Tensor, cu_seqlens_q: cute.Tensor, - block_tables: cute.Tensor, + block_table: cute.Tensor, + q_token_kv_block_sparse_page_memberships: cute.Tensor, partial_o: cute.Tensor, partial_stats: cute.Tensor, split_kv_counter: cute.Tensor, attention_sinks: cute.Tensor, num_physical_kv_pages: cutlass.Int64, k_page_stride: cutlass.Int64, + k_head_stride: cutlass.Int64, + k_token_stride: cutlass.Int64, v_page_stride: cutlass.Int64, + v_head_stride: cutlass.Int64, + v_token_stride: cutlass.Int64, bmm1_scale: cutlass.Float32, bmm2_scale: cutlass.Float32, + direct_q1_inputs: tuple, stream: cuda_drv.CUstream, static_cfg: cutlass.Constexpr[FmhaDecodeConfig], static_seq_len_q: cutlass.Constexpr[int], @@ -1472,9 +1605,12 @@ def main_tensor_adapter( static_full_split_prefix: cutlass.Constexpr[bool], static_native_uniform_kv: cutlass.Constexpr[bool], ) -> None: - """Adapt TVM-FFI tensors to the row-strided native launcher.""" + """Adapt TVM-FFI tensors to the dense block-table pointer launcher.""" - batch_size = Int32(cute.size(seq_lens)) + 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) q_offsets_iter = cu_seqlens_q.iterator total_q_tokens = batch_size * Int32(static_seq_len_q) if cutlass.const_expr(not static_cfg.use_variable_seqlens_q): @@ -1485,6 +1621,32 @@ def main_tensor_adapter( else: total_q_tokens = Int32(q.shape[0]) + lengths_iter = seq_lens.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, + page_size=storage_page_size, + page_capacity=direct_q1_spec.page_capacity, + sparse_block_size=direct_q1_spec.kv_block_size, + fragment_size=static_cfg.num_tokens_per_page, + lengths=True, + ) + pages_iter = DirectSparseMetadataView( + direct_q1_inputs, + cu_seqlens_q, + packed=static_cfg.use_variable_seqlens_q, + model_len=direct_q1_spec.max_seq_len_kv, + page_size=storage_page_size, + page_capacity=direct_q1_spec.page_capacity, + sparse_block_size=direct_q1_spec.kv_block_size, + fragment_size=static_cfg.num_tokens_per_page, + lengths=False, + ) + fmha_decode_launch( ( batch_size, @@ -1496,11 +1658,14 @@ def main_tensor_adapter( q.iterator, k_cache.iterator, v_cache.iterator, + k_sf_cache.iterator, + v_sf_cache.iterator, out.iterator, - seq_lens.iterator, + lengths_iter, q_offsets_iter, total_q_tokens, - block_tables.iterator, + pages_iter, + q_token_kv_block_sparse_page_memberships.iterator, partial_o.iterator, partial_stats.iterator, split_kv_counter.iterator, @@ -1514,12 +1679,16 @@ def main_tensor_adapter( static_max_kv_len, False, True, - block_tables.iterator, - Int32(block_tables.shape[1]), - Int64(block_tables.stride[0]), + Int64(block_table.stride[0]), + Int32(block_table.shape[1]), + Int32(q_token_kv_block_sparse_page_memberships.stride[0]), num_physical_kv_pages, k_page_stride, + k_head_stride, + k_token_stride, v_page_stride, + v_head_stride, + v_token_stride, static_full_split_prefix, static_native_uniform_kv, ) @@ -1540,6 +1709,7 @@ def reduction_tensor_adapter( attention_sinks: cute.Tensor, bmm1_scale: cutlass.Float32, bmm2_scale: cutlass.Float32, + direct_q1_inputs: tuple, stream: cuda_drv.CUstream, static_cfg: cutlass.Constexpr[FmhaDecodeConfig], static_num_qo_heads: cutlass.Constexpr[int], @@ -1550,11 +1720,27 @@ def reduction_tensor_adapter( ) -> None: """Adapt TVM-FFI tensors to the raw standalone split reducer.""" - batch_size = Int32(cute.size(seq_lens)) + 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) q_offsets_iter = cu_seqlens_q.iterator if cutlass.const_expr(not static_cfg.use_variable_seqlens_q): q_offsets_iter = cute.make_ptr(Int32, 0) + lengths_iter = seq_lens.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, + page_size=storage_page_size, + page_capacity=direct_q1_spec.page_capacity, + sparse_block_size=direct_q1_spec.kv_block_size, + fragment_size=static_cfg.num_tokens_per_page, + lengths=True, + ) fmha_decode_separate_reduction_launch( ( batch_size, @@ -1564,7 +1750,7 @@ def reduction_tensor_adapter( Int32(static_head_dim), ), out.iterator, - seq_lens.iterator, + lengths_iter, q_offsets_iter, partial_o.iterator, partial_stats.iterator, @@ -1577,10 +1763,15 @@ def reduction_tensor_adapter( ) physical_pages = cute.sym_int() - block_table_capacity = cute.sym_int() + logical_pages = cute.sym_int() block_table_row_stride = cute.sym_int64(divisibility=1) + membership_words = cute.sym_int() k_outer_stride = cute.sym_int64(divisibility=1) + k_head_stride = cute.sym_int64(divisibility=1) + k_token_stride = cute.sym_int64(divisibility=1) v_outer_stride = cute.sym_int64(divisibility=1) + v_head_stride = cute.sym_int64(divisibility=1) + v_token_stride = cute.sym_int64(divisibility=1) batch_size = cute.sym_int() total_q_tokens = cute.sym_int() runtime_num_q_offsets = cute.sym_int() @@ -1594,21 +1785,33 @@ def reduction_tensor_adapter( ) ) q_fake = cute.runtime.make_fake_compact_tensor( - qkv_dtype, + q_dtype, q_shape, stride_order=tuple(reversed(range(len(q_shape)))), assumed_align=16, ) + k_storage_head_dim = head_dim // 2 if k_dtype == cutlass.Float4E2M1FN else head_dim + v_storage_head_dim = head_dim // 2 if v_dtype == cutlass.Float4E2M1FN else head_dim k_fake = cute.runtime.make_fake_tensor( - qkv_dtype, - (physical_pages, num_kv_heads, page_size, head_dim), - stride=(k_outer_stride, page_size * head_dim, head_dim, 1), + k_storage_dtype, + (physical_pages, num_kv_heads, storage_page_size, k_storage_head_dim), + stride=( + k_outer_stride, + k_head_stride, + k_token_stride, + 1, + ), assumed_align=16, ) v_fake = cute.runtime.make_fake_tensor( - qkv_dtype, - (physical_pages, num_kv_heads, page_size, head_dim), - stride=(v_outer_stride, page_size * head_dim, head_dim, 1), + v_storage_dtype, + (physical_pages, num_kv_heads, storage_page_size, v_storage_head_dim), + stride=( + v_outer_stride, + v_head_stride, + v_token_stride, + 1, + ), assumed_align=16, ) out_shape = _decode_output_shape( @@ -1633,16 +1836,21 @@ def fake_compact(dtype, shape, assumed_align): assumed_align=assumed_align, ) - seq_lens_fake = fake_compact(Int32, (batch_size,), 4) + 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) cu_seqlens_q_fake = fake_compact( Int32, (runtime_num_q_offsets,) if use_packed_q else (1,), 4 ) - block_tables_fake = cute.runtime.make_fake_tensor( + block_table_fake = cute.runtime.make_fake_tensor( Int32, - (batch_size, block_table_capacity), + (metadata_rows, logical_pages), stride=(block_table_row_stride, 1), assumed_align=4, ) + q_token_kv_block_sparse_page_memberships_fake = fake_compact( + Int32, (metadata_rows, membership_words), 4 + ) partial_o_shape, partial_stats_shape, counter_shape = _decode_scratch_shapes( cfg, batch_size=batch_size, @@ -1655,7 +1863,45 @@ def fake_compact(dtype, shape, assumed_align): 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) + 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) + v_sf_fake = fake_compact(cutlass.Float8E4M3FN, sf_shape, 16) + else: + k_sf_fake = fake_compact(cutlass.Uint8, (1,), 1) + v_sf_fake = fake_compact(cutlass.Uint8, (1,), 1) stream_fake = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) + direct_q1_fake: tuple = () + if direct_q1_spec is not None: + rows = cute.sym_int() + direct_q1_fake = ( + cute.runtime.make_fake_tensor( + Int32, + ( + (rows, direct_q1_spec.block_topk) + if cfg.shares_sparse_pattern + else (rows, num_kv_heads, direct_q1_spec.block_topk) + ), + stride=( + (cute.sym_int64(), 1) + if cfg.shares_sparse_pattern + else (cute.sym_int64(), cute.sym_int64(), 1) + ), + assumed_align=4, + ), + cute.runtime.make_fake_tensor( + Int32, + (cute.sym_int(), cute.sym_int()), + stride=(cute.sym_int64(), 1), + assumed_align=4, + ), + fake_compact(Int32, (rows,), 4), + fake_compact( + Int64 if direct_q1_spec.position_dtype_key == "int64" else Int32, + (rows,), + 4, + ), + ) with torch.cuda.device(device_index): compiled_main = cute.compile( @@ -1663,10 +1909,13 @@ def fake_compact(dtype, shape, assumed_align): q_fake, k_fake, v_fake, + k_sf_fake, + v_sf_fake, out_fake, seq_lens_fake, cu_seqlens_q_fake, - block_tables_fake, + block_table_fake, + q_token_kv_block_sparse_page_memberships_fake, partial_o_fake, partial_stats_fake, counter_fake, @@ -1674,8 +1923,13 @@ def fake_compact(dtype, shape, assumed_align): 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, @@ -1701,6 +1955,7 @@ def fake_compact(dtype, shape, assumed_align): attention_sinks_fake, Float32(1.0), Float32(1.0), + direct_q1_fake, stream_fake, cfg, num_qo_heads, @@ -1726,17 +1981,21 @@ def get_prims_ts_batch_decode_workspace_size( qo_indptr: Optional[torch.Tensor] = None, max_seq_len_q: Optional[int] = 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", window_left: int = -1, kv_layout: Literal["HND"] = "HND", + storage_page_size: Optional[int] = None, device: Optional[Union[int, str, torch.device]] = None, + split_kv: bool = True, ) -> int: """Return caller-workspace bytes for one automatic FMHA policy. The arguments resolve the same policy and scratch layout as - :func:`prims_ts_batch_decode_with_kv_cache`, without compiling a kernel. + :func:`batch_decode_with_paged_kv_cache`, without compiling a kernel. Allocate at least the returned number of bytes as a contiguous ``torch.int8`` or ``torch.uint8`` CUDA tensor and zero it before its first FMHA launch. Re-zero a reused buffer whenever any workspace-layout input, @@ -1746,9 +2005,19 @@ def get_prims_ts_batch_decode_workspace_size( 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 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 + encoded ``(physical page, subpage)`` locators. This sizing helper validates that every cumulative-offset delta is positive and no larger than the bound. If ``device`` is omitted, it is inferred from - ``qo_indptr`` for a packed launch. + ``qo_indptr`` for a packed launch. ``k_dtype`` defaults to ``q_dtype`` and + ``v_dtype`` to ``k_dtype``; pass ``v_dtype=torch.float8_e4m3fn`` with + BF16 Q/K to size the QK-BF16/PV-FP8 launch. ``kv_dtype`` is a + compatibility alias for both K and V and must agree with explicit dtypes. + + ``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. """ batch_size = _validate_positive_int(batch_size, "batch_size") @@ -1768,21 +2037,24 @@ def get_prims_ts_batch_decode_workspace_size( ) head_dim = _validate_head_dim(head_dim) page_size = _validate_page_size(page_size) + storage_page_size = _validate_storage_page_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") _validate_layout(kv_layout) _validate_mask(mask_type) window_left = _validate_window_left(window_left, mask_type) - if kv_dtype is None: - kv_dtype = q_dtype + k_dtype, v_dtype = _resolve_kv_dtypes(q_dtype, k_dtype, v_dtype, kv_dtype) if out_dtype is None: out_dtype = q_dtype - _validate_dtype_pair(q_dtype, kv_dtype, out_dtype) + _validate_dtype_pair(q_dtype, k_dtype, v_dtype, out_dtype) inferred_device = ( qo_indptr.device if device is None and isinstance(qo_indptr, torch.Tensor) else device ) - resolved_device, device_index = _resolve_cuda_device(inferred_device) + resolved_device, _ = _resolve_cuda_device(inferred_device) if qo_indptr is not None: _validate_qo_indptr( qo_indptr, @@ -1794,8 +2066,7 @@ def get_prims_ts_batch_decode_workspace_size( max_seq_len_q=seq_len_q, ) - spec = _resolve_decode_launch_spec( - device_index, + return _resolve_decode_workspace_layout( batch_size, num_qo_heads, num_kv_heads, @@ -1803,24 +2074,24 @@ def get_prims_ts_batch_decode_workspace_size( page_size, max_seq_len, seq_len_q, - _dtype_key(q_dtype), - _dtype_key(kv_dtype), - _dtype_key(out_dtype), + q_dtype, + k_dtype, + v_dtype, + out_dtype, kv_layout, mask_type, use_packed_q, window_left, - ) - return _make_decode_workspace_layout( - spec.scratch_shapes, - out_dtype, - use_separate_reduction_kernel=spec.config.use_separate_reduction_kernel, + storage_page_size, + resolved_device, + split_kv=split_kv, ).total_bytes def _prepare_decode_runtime( q: torch.Tensor, paged_kv_cache: PagedKVCache, + kv_scale_factors: Optional[PagedKVScaleFactors], *, device: torch.device, batch_size: int, @@ -1831,7 +2102,8 @@ def _prepare_decode_runtime( head_dim: int, page_size: int, q_dtype: torch.dtype, - kv_dtype: torch.dtype, + k_dtype: torch.dtype, + v_dtype: torch.dtype, output_dtype: torch.dtype, bmm1_scale: Optional[float], bmm2_scale: float, @@ -1839,6 +2111,8 @@ def _prepare_decode_runtime( ) -> _DecodeRuntime: """Validate runtime tensors and normalize zero-copy K/V views.""" + storage_page_size = page_size + _validate_q( q, seq_len_q=seq_len_q, @@ -1849,30 +2123,34 @@ def _prepare_decode_runtime( head_dim=head_dim, q_dtype=q_dtype, ) - ( - k_cache, - v_cache, - num_physical_pages, - runtime_num_kv_heads, - runtime_page_size, - runtime_head_dim, - k_page_stride, - v_page_stride, - ) = _normalize_paged_kv_cache(paged_kv_cache, expected_device=device) + normalized_cache = _normalize_native_paged_kv_cache( + paged_kv_cache, + expected_device=device, + ) + k_cache = normalized_cache.k_cache + v_cache = normalized_cache.v_cache if ( - runtime_num_kv_heads != num_kv_heads - or runtime_page_size != page_size - or runtime_head_dim != head_dim + normalized_cache.num_kv_heads != num_kv_heads + or normalized_cache.storage_page_size != storage_page_size + or normalized_cache.head_dim != head_dim ): raise ValueError( "paged_kv_cache geometry does not match the launch: expected " - f"Hkv/page/D=({num_kv_heads}, {page_size}, {head_dim}), got " - f"({runtime_num_kv_heads}, {runtime_page_size}, {runtime_head_dim})" + f"Hkv/storage_page/D=({num_kv_heads}, {storage_page_size}, " + f"{head_dim}), got " + f"({normalized_cache.num_kv_heads}, " + f"{normalized_cache.storage_page_size}, {normalized_cache.head_dim})" ) - if k_cache.dtype != kv_dtype: + if k_cache.dtype != k_dtype or v_cache.dtype != v_dtype: raise ValueError( - f"K/V dtype must match the launch ({kv_dtype}), got {k_cache.dtype}" + f"K/V dtype must match the launch (K {k_dtype}, V {v_dtype}), got K " + f"{k_cache.dtype} and V {v_cache.dtype}" ) + k_sf_cache, v_sf_cache = _normalize_paged_kv_scale_factors( + kv_scale_factors, + 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", @@ -1900,37 +2178,21 @@ def _prepare_decode_runtime( q=q, k_cache=k_cache, v_cache=v_cache, + k_sf_cache=k_sf_cache, + v_sf_cache=v_sf_cache, out=out, - num_physical_pages=num_physical_pages, - k_page_stride=k_page_stride, - v_page_stride=v_page_stride, + num_physical_pages=normalized_cache.num_physical_pages, + k_page_stride=normalized_cache.k_page_stride, + k_head_stride=normalized_cache.k_head_stride, + k_token_stride=normalized_cache.k_token_stride, + 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, ) -def _validate_decode_output_aliasing( - runtime: _DecodeRuntime, - *, - seq_lens: torch.Tensor, - qo_indptr: Optional[torch.Tensor], - block_tables: torch.Tensor, - workspace_buffer: torch.Tensor, -) -> None: - """Keep output disjoint from every live FMHA decode allocation.""" - - _validate_out_does_not_overlap_inputs( - runtime.out, - ("query", runtime.q), - ("k_cache", runtime.k_cache), - ("v_cache", runtime.v_cache), - ("seq_lens", seq_lens), - ("qo_indptr", qo_indptr), - ("block_tables", block_tables), - ("workspace_buffer", workspace_buffer), - ) - - def _launch_decode( runtime: _DecodeRuntime, *, @@ -1948,19 +2210,27 @@ def _launch_decode( runtime.q, runtime.k_cache, runtime.v_cache, + runtime.k_sf_cache, + runtime.v_sf_cache, runtime.out, seq_lens, 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, runtime.num_physical_pages, runtime.k_page_stride, + runtime.k_head_stride, + runtime.k_token_stride, runtime.v_page_stride, + runtime.v_head_stride, + runtime.v_token_stride, runtime.bmm1_scale, runtime.bmm2_scale, + (), ) if compiled_reducer is not None: compiled_reducer( @@ -1972,6 +2242,7 @@ def _launch_decode( workspace.attention_sinks, runtime.bmm1_scale, runtime.bmm2_scale, + (), ) return runtime.out @@ -2028,8 +2299,9 @@ def _normalize_plan_seq_lens( def _prepare_decode_runtime_unchecked( q: torch.Tensor, paged_kv_cache: PagedKVCache, + kv_scale_factors: Optional[PagedKVScaleFactors], *, - state: _DecodePlanState, + output_dtype: torch.dtype, bmm1_scale: Optional[float], bmm2_scale: float, out: Optional[torch.Tensor], @@ -2041,55 +2313,67 @@ def _prepare_decode_runtime_unchecked( v_cache = paged_kv_cache[:, 1] else: k_cache, v_cache = paged_kv_cache - output_shape = _decode_output_shape( - batch_size=state.batch_size, - num_qo_heads=state.num_qo_heads, - seq_len_q=state.seq_len_q, - head_dim=state.head_dim, - total_q_tokens=int(q.shape[0]) if state.use_packed_q else None, - ) + if k_cache.dtype == torch.uint8: + k_sf_cache, v_sf_cache = cast(PagedKVScaleFactors, kv_scale_factors) + else: + placeholder = k_cache[0, 0, 0, :1].view(torch.uint8)[:1] + k_sf_cache, v_sf_cache = placeholder, placeholder if out is None: out = torch.empty( - output_shape, - device=state.device, - dtype=state.output_dtype, + q.shape, + device=q.device, + dtype=output_dtype, ) return _DecodeRuntime( q=q, k_cache=k_cache, v_cache=v_cache, + k_sf_cache=k_sf_cache, + v_sf_cache=v_sf_cache, out=out, num_physical_pages=int(k_cache.shape[0]), k_page_stride=int(k_cache.stride(0)), + k_head_stride=int(k_cache.stride(1)), + k_token_stride=int(k_cache.stride(2)), 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=( - 1.0 / math.sqrt(state.head_dim) if bmm1_scale is None else float(bmm1_scale) + 1.0 / math.sqrt(int(q.shape[-1])) + if bmm1_scale is None + else float(bmm1_scale) ), bmm2_scale=float(bmm2_scale), ) def _validate_decode_run_metadata_values( - state: _DecodePlanState, runtime: _DecodeRuntime, *, seq_lens: 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, + seq_len_q: int, + batch_size: int, + mask_type: str, ) -> None: """Synchronously validate per-run metadata values.""" - runtime_seq_lens = state.planned_seq_lens_host + 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()) table_capacity = int(block_tables.shape[1]) for request_idx, seq_len in enumerate(runtime_seq_lens): - if seq_len <= 0 or seq_len > state.max_kv_len: + if seq_len <= 0 or seq_len > max_kv_len: raise ValueError( - f"seq_lens values must be within [1, {state.max_kv_len}]; " + f"seq_lens values must be within [1, {max_kv_len}]; " f"request {request_idx} has {seq_len}" ) - required_pages = (seq_len + state.page_size - 1) // state.page_size + required_pages = (seq_len + page_size - 1) // page_size if table_capacity < required_pages: raise ValueError( "block_tables does not have enough columns for " @@ -2098,27 +2382,30 @@ def _validate_decode_run_metadata_values( ) block_table_rows = block_tables.tolist() + locator_capacity = runtime.num_physical_pages * ( + 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) ): - required_pages = (seq_len + state.page_size - 1) // state.page_size + required_pages = (seq_len + page_size - 1) // page_size if any( - int(page_id) < 0 or int(page_id) >= runtime.num_physical_pages + int(page_id) < 0 or int(page_id) >= locator_capacity for page_id in row[:required_pages] ): raise ValueError( "block_tables values for active pages must index the physical " - f"K/V cache in [0, {runtime.num_physical_pages}); request " + f"K/V cache through locators in [0, {locator_capacity}); request " f"{request_idx} contains an invalid page ID" ) - if state.use_packed_q: + if use_packed_q: assert qo_indptr is not None _, _, q_lengths = _read_packed_q_plan_metadata(qo_indptr) - if max(q_lengths) > state.seq_len_q: + if max(q_lengths) > seq_len_q: raise ValueError( "qo_indptr contains a per-request Q length larger than " - f"max_seq_len_q ({state.seq_len_q}): got {max(q_lengths)}" + f"max_seq_len_q ({seq_len_q}): got {max(q_lengths)}" ) total_q = sum(q_lengths) if total_q != int(runtime.q.shape[0]): @@ -2127,9 +2414,9 @@ def _validate_decode_run_metadata_values( f"expected {runtime.q.shape[0]}, got {total_q}" ) else: - q_lengths = (state.seq_len_q,) * state.batch_size + q_lengths = (seq_len_q,) * batch_size - if state.mask_type == "causal": + if mask_type == "causal": for request_idx, (q_len, kv_len) in enumerate( zip(q_lengths, runtime_seq_lens, strict=True) ): @@ -2141,119 +2428,939 @@ def _validate_decode_run_metadata_values( ) -@flashinfer_api -def prims_ts_batch_decode_with_kv_cache( - query: torch.Tensor, - kv_cache: PagedKVCache, - workspace_buffer: torch.Tensor, - block_tables: torch.Tensor, - seq_lens: torch.Tensor, - max_seq_len: int, - *, - seq_len_q: int = 1, - qo_indptr: Optional[torch.Tensor] = None, - max_seq_len_q: Optional[int] = None, - bmm1_scale: Optional[float] = None, - bmm2_scale: float = 1.0, - out: Optional[torch.Tensor] = None, - out_dtype: Optional[torch.dtype] = None, - mask_type: Literal["dense", "causal"] = "dense", - window_left: int = -1, - kv_layout: Literal["HND"] = "HND", -) -> torch.Tensor: - """Launch fixed or packed-Q native page-table decode with caller scratch. - - For ``seq_len_q=1``, ``query`` and the returned output both have shape - ``[B, Hq, D]``. For ``seq_len_q>1``, both use compact token-major - ``[B, SQ, Hq, D]`` storage. The kernel writes that layout directly; no - layout transpose is performed. When ``qo_indptr`` is supplied, Q and O use - packed ``[total_q, Hq, D]`` storage. Request ``b`` owns rows - ``qo_indptr[b]:qo_indptr[b+1]``; ``max_seq_len_q`` is only the static - workspace/JIT bound and is required for this standalone packed interface. - To keep this launch path free of device-to-host synchronization, callers - must ensure that packed offsets start at zero, are strictly increasing, end - at ``query.shape[0]``, and have every delta at most ``max_seq_len_q``. For - causal masking, every fixed or packed per-request Q length must also be no - greater than the corresponding live ``seq_lens`` value. - - ``kv_cache`` is either a combined - ``[pages, 2, Hkv, page_size, D]`` tensor or a ``(K, V)`` tuple of - ``[pages, Hkv, page_size, D]`` tensors. ``block_tables`` is an int32 CUDA - tensor with shape ``[B, C]``. Entries are contiguous within each row while - the outer row stride may be greater than ``C``. ``seq_lens`` is explicit - and ``max_seq_len`` is the exact static maximum used for automatic policy - selection and JIT caching. - It must be no larger than ``2,147,483,392`` so the padded 256-token K/V - tile endpoint remains representable as signed Int32. - The table must contain enough columns for every per-run length:: - - (seq_lens[b] + page_size - 1) // page_size <= C - - Every active page ID must index ``kv_cache``. Inactive tail entries at and - beyond ``ceil(seq_lens[b] / page_size)`` are ignored. - - ``workspace_buffer`` must be zero-initialized before its first use and - re-zeroed whenever any workspace-layout input, including ``batch_size``, - changes because the internal section offsets can move even when the - compiled callable is reused. It is exclusive to one in-flight launch or - captured graph and must not overlap query, K/V - cache, metadata, or output storage. Runtime sequence lengths must remain - positive and no larger than ``max_seq_len``; this hot path - deliberately does not read device metadata back to the host. Per-run table, - length, page-ID, and packed-Q values may change between completed launches - or graph replays only while all of their contracts remain valid. They must - not be mutated concurrently with a launch or replay that reads them. Warm - the planned topology before CUDA graph capture and provide ``out`` to avoid an - output allocation. Captured graphs must retain stable metadata storage; - ``qo_indptr`` values may change only while the packed-offset contract - remains valid, every delta stays within the compiled bound, and the final - offset continues to match the captured query/output extent. - ``window_left=-1`` disables the left window; a - non-negative value requires causal masking and includes the current token. - No backend fallback or scheduling knob is exposed. - - Parameters - ---------- - query : torch.Tensor - Fixed or packed query tensor. - kv_cache : torch.Tensor or tuple[torch.Tensor, torch.Tensor] - Combined or separate paged K/V storage. - workspace_buffer : torch.Tensor - Zero-initialized caller-owned byte workspace for this planned layout. - block_tables : torch.Tensor - CUDA int32 physical page IDs with shape ``[B, C]``, unit inner stride, - and non-overlapping row stride at least ``C``. - seq_lens : torch.Tensor - Live K/V sequence lengths for each request. - max_seq_len : int - Static maximum K/V length used for policy selection and JIT caching. - 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. - 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, - a non-default ``seq_len_q`` must supply the bound. In fixed-query mode, - it must equal ``seq_len_q``. - bmm1_scale, bmm2_scale : float, optional - QK and value/output scaling factors. - out : torch.Tensor, optional - Caller-owned output tensor. - out_dtype : torch.dtype, optional - Output dtype; defaults to ``out.dtype`` or the query dtype. - mask_type : {"dense", "causal"} - Attention mask mode. - window_left : int - Left sliding-window extent, or ``-1`` to disable the window. - kv_layout : {"HND"} - Layout of the paged K/V cache. - """ +@dataclass(frozen=True) +class _NativePagedKVCache: + """Zero-copy logical HND views and TensorMap strides for native decode.""" - _validate_layout(kv_layout) - _validate_mask(mask_type) - window_left = _validate_window_left(window_left, mask_type) + k_cache: torch.Tensor + v_cache: torch.Tensor + num_physical_pages: int + num_kv_heads: int + storage_page_size: int + head_dim: int + k_page_stride: int + k_head_stride: int + k_token_stride: int + v_page_stride: int + v_head_stride: int + v_token_stride: int + + +@dataclass(frozen=True) +class PrimsTSBatchDecodePlan: + """Validated dense-page-table PrimTS state for framework hot paths. + + The plan retains the K/V cache and scale tensors, dense block-table storage, compiled + callables, and typed workspace views validated at construction. Block-table + and sequence-length values may change between completed launches, which is + the contract needed by QToken-KvBlock-Sparse-Attention metadata builders. Query and output storage may + also change, but must preserve the exact shape, dtype, device, and strides + proven by the representative tensors passed to + :func:`prepare_prims_ts_batch_decode_with_kv_cache`. + + ``run`` deliberately omits repeated cache-stride validation, workspace + rebinding, and semantic-policy resolution. Storage overlap is not checked + during preparation or execution. The framework must keep retained tensors + alive, keep output/workspace disjoint from the inputs and each other, and + not mutate metadata concurrently with a launch or replay that reads it. + """ + + _query_shape: tuple[int, ...] + _query_stride: tuple[int, ...] + _output_shape: tuple[int, ...] + _output_stride: tuple[int, ...] + _device: torch.device + _q_dtype: torch.dtype + _output_dtype: torch.dtype + _head_dim: int + _cache: _NativePagedKVCache + _k_sf_cache: torch.Tensor + _v_sf_cache: torch.Tensor + _seq_lens: 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, ...] = () + + 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.""" + + if ( + query.shape != self._query_shape + or query.stride() != self._query_stride + or query.device != self._device + or query.dtype != self._q_dtype + ): + raise ValueError( + "query must preserve the shape, strides, device, and dtype " + "validated by the PrimTS plan" + ) + if ( + out.shape != self._output_shape + or out.stride() != self._output_stride + or out.device != self._device + or out.dtype != self._output_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", + ) + scale_v = _validate_scale(bmm2_scale, "bmm2_scale") + return self._run_unchecked(query, out, scale_qk, scale_v) + + def _run_unchecked( + self, + query: torch.Tensor, + out: torch.Tensor, + scale_qk: float, + scale_v: float, + direct_q1_inputs: Optional[tuple[torch.Tensor, ...]] = None, + ) -> torch.Tensor: + """Launch state already proven by a framework-owned outer plan.""" + + q_offsets = ( + self._workspace.cu_seqlens_q if self._qo_indptr is None else self._qo_indptr + ) + direct_inputs = ( + self._direct_q1_inputs if direct_q1_inputs is None else direct_q1_inputs + ) + self._compiled_main( + query, + self._cache.k_cache, + self._cache.v_cache, + self._k_sf_cache, + self._v_sf_cache, + out, + self._seq_lens, + 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._cache.num_physical_pages, + self._cache.k_page_stride, + self._cache.k_head_stride, + self._cache.k_token_stride, + self._cache.v_page_stride, + self._cache.v_head_stride, + self._cache.v_token_stride, + scale_qk, + scale_v, + direct_inputs, + ) + if self._compiled_reducer is not None: + self._compiled_reducer( + out, + self._seq_lens, + q_offsets, + self._workspace.partial_o, + self._workspace.partial_stats, + self._workspace.attention_sinks, + scale_qk, + scale_v, + direct_inputs, + ) + return out + + +def _resolve_decode_workspace_layout( + batch_size: int, + num_qo_heads: int, + num_kv_heads: int, + head_dim: int, + page_size: int, + max_seq_len: int, + seq_len_q: int, + q_dtype: torch.dtype, + k_dtype: torch.dtype, + v_dtype: torch.dtype, + out_dtype: torch.dtype, + kv_layout: str, + mask_type: str, + use_packed_q: bool, + window_left: int, + storage_page_size: int, + device: Optional[Union[int, str, torch.device]], + use_q_token_kv_block_sparse_route: bool = False, + use_pdl: bool = False, + split_kv: bool = True, + share_pattern_across_kv_heads: bool = True, +) -> _DecodeWorkspaceLayout: + """Resolve the byte layout for one already-validated semantic key.""" + + resolved_device, device_index = _resolve_cuda_device(device) + _validate_runtime_device(resolved_device) + spec = _resolve_decode_launch_spec( + device_index, + batch_size, + num_qo_heads, + num_kv_heads, + head_dim, + page_size, + max_seq_len, + seq_len_q, + _dtype_key(q_dtype), + _dtype_key(k_dtype), + _dtype_key(v_dtype), + _dtype_key(out_dtype), + kv_layout, + mask_type, + use_packed_q, + window_left, + storage_page_size, + use_q_token_kv_block_sparse_route, + use_pdl, + split_kv, + share_pattern_across_kv_heads, + ) + return _make_decode_workspace_layout( + spec.scratch_shapes, + out_dtype, + use_separate_reduction_kernel=spec.config.use_separate_reduction_kernel, + use_split_kv=spec.config.use_split_kv, + ) + + +def _validate_storage_page_size(page_size: int, storage_page_size: int) -> int: + """Require physical cache pages to contain whole logical fragments.""" + storage_page_size = _validate_positive_int(storage_page_size, "storage_page_size") + if storage_page_size % page_size != 0: + raise ValueError("storage_page_size must be divisible by page_size") + return storage_page_size + + +def _normalize_paged_kv_cache_views( + paged_kv_cache: PagedKVCache, + *, + expected_device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor, int, int, int, int]: + """Return validated zero-copy HND views and the logical head dimension.""" + + if isinstance(paged_kv_cache, torch.Tensor): + if paged_kv_cache.ndim != 5 or paged_kv_cache.shape[1] != 2: + raise ValueError( + "combined paged_kv_cache must have shape " + "[num_pages, 2, Hkv, page_size, head_dim]" + ) + if paged_kv_cache.device != expected_device: + raise ValueError( + f"paged_kv_cache must be on {expected_device}, " + f"got {paged_kv_cache.device}" + ) + k_cache = paged_kv_cache[:, 0] + v_cache = paged_kv_cache[:, 1] + elif isinstance(paged_kv_cache, tuple): + if len(paged_kv_cache) != 2: + raise ValueError("paged_kv_cache tuple must contain exactly (K, V)") + k_cache, v_cache = paged_kv_cache + if not isinstance(k_cache, torch.Tensor) or not isinstance( + v_cache, torch.Tensor + ): + raise TypeError("paged_kv_cache tuple members must be torch.Tensor") + if k_cache.ndim != 4 or v_cache.ndim != 4: + raise ValueError( + "tuple K/V caches must each have shape " + "[num_pages, Hkv, page_size, head_dim]" + ) + if k_cache.device != expected_device or v_cache.device != expected_device: + raise ValueError(f"tuple K/V caches must be on {expected_device}") + else: + raise TypeError( + "paged_kv_cache must be a combined torch.Tensor or a (K, V) tuple" + ) + + if k_cache.ndim != 4 or v_cache.ndim != 4: + raise ValueError("K/V cache views must be rank-4 HND tensors") + 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: + raise ValueError("K and V cache views must be on the same device") + + num_pages, num_kv_heads, page_size, storage_head_dim = map(int, k_cache.shape) + if min(num_pages, num_kv_heads, page_size, storage_head_dim) <= 0: + raise ValueError("paged_kv_cache dimensions must be positive") + # This API reserves uint8 K/V storage for two packed NVFP4 values per byte. + head_dim = ( + storage_head_dim * 2 if k_cache.dtype == torch.uint8 else storage_head_dim + ) + return ( + k_cache, + v_cache, + num_pages, + num_kv_heads, + page_size, + head_dim, + ) + + +def _validate_native_hnd_tma_strides( + tensor: torch.Tensor, + name: str, +) -> tuple[int, int, int]: + """Validate a non-overlapping logical HND view for TensorMap addressing.""" + + num_pages, num_kv_heads, storage_page_size, head_dim = map(int, tensor.shape) + page_stride, head_stride, token_stride, dim_stride = map(int, tensor.stride()) + if dim_stride != 1: + raise ValueError(f"{name} head-dimension stride must be one") + named_strides = ( + ("page", page_stride), + ("head", head_stride), + ("token", token_stride), + ) + if any(stride <= 0 for _, stride in named_strides): + raise ValueError(f"{name} TensorMap strides must be positive") + max_int64 = 2**63 - 1 + if any(stride > max_int64 for _, stride in named_strides): + raise ValueError(f"{name} TensorMap stride exceeds signed int64") + + element_size = tensor.element_size() + for stride_name, stride in named_strides: + if stride * element_size % 16 != 0: + raise ValueError(f"{name} {stride_name} stride must be 16-byte aligned") + _validate_16byte_alignment(tensor, name) + + # A gapped K or V slice of a packed 2D cache is not dense, so PyTorch's + # dense-layout predicates cannot establish safety. Prove non-overlap from + # the actual logical dimensions and strides instead. + span = 1 + dimensions = sorted( + ( + (1, head_dim, "head dimension"), + (token_stride, storage_page_size, "token"), + (head_stride, num_kv_heads, "head"), + (page_stride, num_pages, "page"), + ), + key=lambda item: item[0], + ) + for stride, extent, dimension_name in dimensions: + if extent <= 1: + continue + if stride < span: + raise ValueError( + f"{name} {dimension_name} dimension overlaps another dimension" + ) + span += (extent - 1) * stride + if span - 1 > max_int64: + raise ValueError(f"{name} address span exceeds signed int64") + return page_stride, head_stride, token_stride + + +def _normalize_native_paged_kv_cache( + paged_kv_cache: PagedKVCache, + *, + expected_device: torch.device, +) -> _NativePagedKVCache: + """Return native decode views supporting HND- or NHD-physical storage.""" + + ( + k_cache, + v_cache, + num_pages, + num_kv_heads, + storage_page_size, + head_dim, + ) = _normalize_paged_kv_cache_views( + paged_kv_cache, + expected_device=expected_device, + ) + k_page_stride, k_head_stride, k_token_stride = _validate_native_hnd_tma_strides( + k_cache, "K cache" + ) + v_page_stride, v_head_stride, v_token_stride = _validate_native_hnd_tma_strides( + v_cache, "V cache" + ) + return _NativePagedKVCache( + k_cache=k_cache, + v_cache=v_cache, + num_physical_pages=num_pages, + num_kv_heads=num_kv_heads, + storage_page_size=storage_page_size, + head_dim=head_dim, + k_page_stride=k_page_stride, + k_head_stride=k_head_stride, + k_token_stride=k_token_stride, + v_page_stride=v_page_stride, + v_head_stride=v_head_stride, + v_token_stride=v_token_stride, + ) + + +def _validate_q_token_kv_block_sparse_page_memberships( + q_token_kv_block_sparse_page_memberships: torch.Tensor, + *, + expected_device: torch.device, + batch_size: int, + max_num_pages: int, +) -> None: + """Validate the private grouped-QToken-KvBlock-Sparse-Attention packed-membership table.""" + + if not isinstance(q_token_kv_block_sparse_page_memberships, torch.Tensor): + raise TypeError( + "q_token_kv_block_sparse_page_memberships must be a torch.Tensor" + ) + if q_token_kv_block_sparse_page_memberships.ndim != 2: + raise ValueError( + f"q_token_kv_block_sparse_page_memberships must be rank 2, got rank {q_token_kv_block_sparse_page_memberships.ndim}" + ) + if q_token_kv_block_sparse_page_memberships.dtype != torch.int32: + raise TypeError( + "q_token_kv_block_sparse_page_memberships must have dtype torch.int32" + ) + if q_token_kv_block_sparse_page_memberships.device != expected_device: + raise ValueError( + f"q_token_kv_block_sparse_page_memberships must be on {expected_device}, got " + f"{q_token_kv_block_sparse_page_memberships.device}" + ) + if not q_token_kv_block_sparse_page_memberships.is_contiguous(): + raise ValueError("q_token_kv_block_sparse_page_memberships must be contiguous") + expected_shape = (batch_size, (max_num_pages + 3) // 4) + if tuple(q_token_kv_block_sparse_page_memberships.shape) != expected_shape: + raise ValueError( + "q_token_kv_block_sparse_page_memberships must have shape " + f"{expected_shape}, got {tuple(q_token_kv_block_sparse_page_memberships.shape)}" + ) + if q_token_kv_block_sparse_page_memberships.data_ptr() % 4 != 0: + raise ValueError( + "q_token_kv_block_sparse_page_memberships data pointer must be 4-byte aligned" + ) + + +def _resolve_q_token_kv_block_sparse_decode_config( + make_config: Callable[..., "FmhaDecodeConfig"], + *, + batch_size: int, + num_qo_heads: int, + 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, + mask_type: str, + use_packed_q: bool, + window_left: int, + max_active_clusters: int, + use_pdl: bool, + split_kv: bool = True, + share_pattern_across_kv_heads: bool = True, +) -> "FmhaDecodeConfig": + """Resolve the private encoded-page QToken-KvBlock-Sparse-Attention profile without compiling it.""" + + from .kernels.fmha_decode.fmha_decode_config import ( + MIN_LOOP_ITERS_PER_SPLIT, + ) + + heads_q_per_kv = num_qo_heads // num_kv_heads + _validate_prims_ts_q_token_kv_block_sparse_group_capacity( + seq_len_q, + num_qo_heads, + num_kv_heads, + ) + # This route has no QK-BF16/PV-FP8 kernel, so Q, K, and V must match. + q_token_kv_block_sparse_dtype_supported = ( + q_dtype_key == k_dtype_key == v_dtype_key + and ( + q_dtype_key == output_dtype_key + and q_dtype_key in ("float16", "bfloat16") + or ( + q_dtype_key == "float8_e4m3fn" + and output_dtype_key in ("float16", "bfloat16") + ) + ) + ) + if not ( + page_size in _SUPPORTED_PAGE_SIZES + and head_dim in _SUPPORTED_HEAD_DIMS + and q_token_kv_block_sparse_dtype_supported + and mask_type == "causal" + and window_left < 0 + ): + raise ValueError( + "PrimTS QToken-KvBlock-Sparse-Attention requires a supported sparse-fragment size, head_dim in {64,128,256}, " + "FP16/BF16 Q/K/V with matching output or FP8 Q/K/V with FP16/BF16 output, " + "and a causal non-windowed mask" + ) + + launch = _sparse_policy.select_sparse_launch( + group_size=seq_len_q, + heads_q_per_kv=heads_q_per_kv, + head_dim=head_dim, + q_dtype_key=q_dtype_key, + num_routes=batch_size, + num_kv_heads=num_kv_heads, + route_kv_tokens=max_kv_len, + multi_processor_count=max_active_clusters, + split_kv=split_kv, + ) + + q_token_kv_block_sparse_profile = { + "use_variable_seqlens_q": use_packed_q, + "use_q_token_kv_block_sparse_route": True, + "share_pattern_across_kv_heads": share_pattern_across_kv_heads, + "use_pdl": use_pdl, + "use_keeps_mma_ab": launch.use_keeps_mma_ab, + "groups_tokens_heads_q": True, + "tile_size_q": launch.tile_size_q, + "tile_size_kv": _sparse_policy.TILE_SIZE_KV, + "head_dim_per_stage_kv": launch.head_dim_per_stage_kv, + "num_insts_kv": launch.num_insts_kv, + "use_persistent_scheduler": launch.use_persistent_scheduler, + "correction_num_warps": 4, + "mma_warp_idx": 12, + # CLC uses the control slot beside MMA; page producers occupy the + # remaining slots. TMA issuers keep their own complete warp groups. + "scheduler_warp_idx": 13, + "clc_load_warp_idx": 16, + "page_offsets_warp_idx": 14, + "page_offsets_num_warps": 2 if seq_len_q == 1 else 1, + "load_warp_idx": 16, + # FP8's predicated TMA helper constructs coordinates in every load + # lane, so Q1 needs only one producer warpgroup. BF16 distributes page + # fragments across two producer warpgroups. Grouped routes keep + # their qualified two-warpgroup schedule for both dtypes. + "load_num_warps": ( + 4 if seq_len_q == 1 and q_dtype_key == "float8_e4m3fn" else 8 + ), + } + if launch.use_keeps_mma_ab: + q_token_kv_block_sparse_profile["o_stages"] = launch.o_stages + return make_config( + q_token_kv_block_sparse_profile, + split_kv_mode=( + "gmem_reduction_with_separate_kernel" + if launch.splits_kv > 1 + else "disabled" + ), + splits_kv=launch.splits_kv, + max_splits_kv=launch.splits_kv, + min_loop_iters_per_split=(1 if seq_len_q == 1 else MIN_LOOP_ITERS_PER_SPLIT), + ) + + +def _validate_prims_ts_q_token_kv_block_sparse_group_value(group_size: int) -> int: + """Validate one caller-selected QToken-KvBlock-Sparse-Attention grouping value.""" + + group_size = _validate_positive_int(group_size, "group_size") + if group_size not in _sparse_policy.SUPPORTED_GROUP_SIZES: + raise ValueError( + "QToken-KvBlock-Sparse-Attention group_size must be one of " + f"{_sparse_policy.SUPPORTED_GROUP_SIZES}, got {group_size}" + ) + return group_size + + +def _validate_prims_ts_q_token_kv_block_sparse_group_capacity( + group_size: int, + num_qo_heads: int, + num_kv_heads: int, +) -> int: + """Validate that one caller-selected sparse Q group fits TileQ128 and a membership byte.""" + + from .kernels.fmha_decode.fmha_decode_constants import ( + Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIP_BITS, + ) + + group_size = _validate_prims_ts_q_token_kv_block_sparse_group_value(group_size) + _validate_head_geometry(num_qo_heads, num_kv_heads) + heads_q_per_kv = num_qo_heads // num_kv_heads + max_group_size = min( + Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIP_BITS, + _sparse_policy.MAX_TILE_SIZE_Q // heads_q_per_kv, + ) + if group_size > max_group_size: + raise ValueError( + f"QToken-KvBlock-Sparse-Attention group_size={group_size} exceeds the TileQ128/head capacity " + f"of {max_group_size}" + ) + + return group_size + + +@flashinfer_experimental_api +def suggest_q_token_kv_block_sparse_group_size( + batch_size: int, + seq_len_q: int, + selected_seq_len_kv: int, + num_qo_heads: int, + num_kv_heads: int, + multi_processor_count: int, + *, + head_dim: int = 256, + q_data_type: torch.dtype = torch.bfloat16, + split_kv: bool = True, + kv_block_size: int = 4, +) -> int: + """Suggest a Q1--Q8 QToken-KvBlock-Sparse-Attention group. + + This pure host-side policy uses the caller-provided SM count; it never + queries device properties or reads tensors. It starts with the largest + legal group and the smallest fitting dense Q-tile recipe, then adds + useful one-wave split-KV parallelism. It considers the largest group for + each dense Q tile and further ungrouping at the smallest tile. If the + largest group cannot fill a wave, prefer better active SM coverage after + instruction-aligned split rounding. Ties prefer less Q padding, then the + larger group. Explicit G passed to plan/run is never overridden. + + Query lengths need not be divisible by the result. Packed prefill may end + each request with a shorter route. Fixed-shape decode may pad its final + ``[B, num_query_groups, G, Hq, D]`` route with consecutive semantic dummy + rows and discard their outputs. + + ``selected_seq_len_kv`` is the maximum candidate K/V-token count for one + query, including its causal tail. Configured split fanout uses union + capacity to accommodate non-shared KV. The recommendation scores useful + active splits using a fully-shared per-query proxy, rather than counting + every capacity-based split as active. Actual union-overlap and scattered-load + costs remain a TODO. + + Parameters + ---------- + batch_size : int + Number of requests in the launch. + seq_len_q : int + Maximum number of live query tokens per request. + selected_seq_len_kv : int + Maximum selected candidate K/V tokens per query, including the causal + tail. This is not the original context length or global cache capacity. + num_qo_heads : int + Number of query/output heads. + num_kv_heads : int + Number of key/value heads. + multi_processor_count : int + Number of SMs available to the launch. Frameworks should query it once + 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/K/V dtype used to select the matching dense MMA recipe. + split_kv : bool + Allow automatic split-KV fanout when True (default); False forces + one KV partition. Pass False for prefill and True for decode, + independently of packed or fixed Q layout. Match this value in + workspace sizing and attention planning. Ragged workloads use maximum + Q length as a capacity estimate; no GPU sequence lengths are inspected. + + kv_block_size : int + Semantic sparse block size: 4/8/16/32/64/128, default four. Used to + bound grouped union work; storage-page size does not affect this bound. + + Returns + ------- + int + One of Q1--Q8, subject to the TileQ128 head-capacity bound. + """ + + from .q_token_kv_block_sparse_metadata import _validate_sparse_block_size + + kv_block_size = _validate_sparse_block_size(kv_block_size) + batch_size = _validate_positive_int(batch_size, "batch_size") + seq_len_q = _validate_positive_int(seq_len_q, "seq_len_q") + selected_seq_len_kv = _validate_positive_int( + selected_seq_len_kv, "selected_seq_len_kv" + ) + multi_processor_count = _validate_positive_int( + multi_processor_count, "multi_processor_count" + ) + _validate_head_geometry(num_qo_heads, num_kv_heads) + 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 not isinstance(split_kv, bool): + raise TypeError("split_kv must be a bool") + return _sparse_policy.suggest_sparse_group( + batch_size=batch_size, + seq_len_q=seq_len_q, + selected_kv_tokens=selected_seq_len_kv, + heads_q_per_kv=heads_q_per_kv, + num_kv_heads=num_kv_heads, + multi_processor_count=multi_processor_count, + head_dim=head_dim, + q_dtype_key=_dtype_key(q_data_type), + split_kv=split_kv, + block_size=kv_block_size, + ) + + +def _validate_prims_ts_q_token_kv_block_sparse_group_layout( + group_size: int, + query_start_loc_cpu: Optional[torch.Tensor], + num_query_tokens: int, + num_qo_heads: int, + num_kv_heads: int, +) -> int: + """Validate one caller-selected QToken-KvBlock-Sparse-Attention group against the live query rows.""" + + group_size = _validate_prims_ts_q_token_kv_block_sparse_group_capacity( + group_size, + num_qo_heads, + num_kv_heads, + ) + num_query_tokens = _validate_positive_int(num_query_tokens, "num_query_tokens") + if group_size == 1: + return 1 + if ( + query_start_loc_cpu is None + or not isinstance(query_start_loc_cpu, torch.Tensor) + or query_start_loc_cpu.device.type != "cpu" + or query_start_loc_cpu.ndim != 1 + or query_start_loc_cpu.numel() < 2 + or query_start_loc_cpu.dtype not in (torch.int32, torch.int64) + ): + raise ValueError( + "grouped QToken-KvBlock-Sparse-Attention requires CPU int32/int64 query_start_loc with at " + "least two entries" + ) + + query_starts = [int(value) for value in query_start_loc_cpu.tolist()] + num_mapped_tokens = query_starts[-1] + if ( + query_starts[0] != 0 + or num_mapped_tokens < 0 + or num_mapped_tokens > num_query_tokens + ): + raise ValueError( + "grouped QToken-KvBlock-Sparse-Attention query boundaries do not cover valid rows" + ) + query_lengths = [ + end - begin for begin, end in zip(query_starts, query_starts[1:], strict=False) + ] + if any(length < 0 for length in query_lengths): + raise ValueError( + "grouped QToken-KvBlock-Sparse-Attention query boundaries must be nondecreasing" + ) + query_lengths = [length for length in query_lengths if length > 0] + if not query_lengths: + raise ValueError( + "grouped QToken-KvBlock-Sparse-Attention requires at least one nonempty request" + ) + + return group_size + + +@flashinfer_experimental_api +def validate_q_token_kv_block_sparse_group_size( + query_start_loc_cpu: Optional[torch.Tensor], + num_query_tokens: int, + num_qo_heads: int, + num_kv_heads: int, + *, + group_size: int, +) -> int: + """Validate a QToken-KvBlock-Sparse-Attention Q1--Q8 group. + + ``query_start_loc_cpu`` contains cumulative flattened-query offsets for + real requests. ``num_query_tokens`` includes any inert CUDA-graph padding + rows. Validation caps the maximum grouped query by both the kernel + membership representation and TileQ128 capacity. Request lengths do not + have to be divisible by ``group_size``: frameworks can partition each + request into packed routes of at most that size without crossing a request + or real/padding boundary. + + This function never chooses a group from workload size or occupancy. The + framework owns ``group_size``; attention split-KV fanout is selected + separately after the grouped launch grid is known. + + Parameters + ---------- + query_start_loc_cpu : torch.Tensor or None + CPU Int32 or Int64 cumulative offsets for the real requests. It is + required when ``group_size`` is greater than one; ``None`` is accepted + for Q1. + num_query_tokens : int + Total flattened query-row count, including any inert CUDA-graph + padding rows. + num_qo_heads : int + Number of query/output heads. + num_kv_heads : int + Number of key/value heads. + group_size : int + Caller-selected number of query rows per QToken-KvBlock-Sparse-Attention route. Supported values + are Q1--Q8, subject to the TileQ128 head-capacity bound. + + Returns + ------- + int + The validated ``group_size``. + """ + + return _validate_prims_ts_q_token_kv_block_sparse_group_layout( + group_size, + query_start_loc_cpu, + num_query_tokens, + num_qo_heads, + num_kv_heads, + ) + + +@flashinfer_experimental_api +def make_q_token_kv_block_sparse_qo_indptr( + query_start_loc_cpu: torch.Tensor, + num_query_tokens: int, + *, + group_size: int, + device: Optional[Union[int, str, torch.device]] = None, +) -> torch.Tensor: + """Build request-safe QToken-KvBlock-Sparse-Attention route offsets. + + ``query_start_loc_cpu`` contains cumulative offsets for real requests. + Every request is chunked independently, so its final route may contain + fewer than ``group_size`` rows. Any inert CUDA-graph padding suffix is + partitioned separately. The returned Int32 ``qo_indptr`` can be passed to + the packed-Q wrapper together with planned ``seq_len_q=group_size``. + + Parameters + ---------- + query_start_loc_cpu : torch.Tensor + One-dimensional CPU Int32 or Int64 cumulative offsets for the real + requests. The first entry must be zero and the offsets must be + nondecreasing. + num_query_tokens : int + Total flattened query-row count, including any inert CUDA-graph + padding suffix. + group_size : int + Maximum number of rows in each request-safe QToken-KvBlock-Sparse-Attention route. + device : int, str, torch.device, or None + Destination device for the returned offsets. By default, PyTorch's + default tensor device is used. + + Returns + ------- + torch.Tensor + Contiguous Int32 cumulative route offsets. Every nonempty route has at + most ``group_size`` rows and no route crosses a request boundary. + """ + + group_size = _validate_prims_ts_q_token_kv_block_sparse_group_value(group_size) + num_query_tokens = _validate_positive_int(num_query_tokens, "num_query_tokens") + if ( + not isinstance(query_start_loc_cpu, torch.Tensor) + or query_start_loc_cpu.device.type != "cpu" + or query_start_loc_cpu.ndim != 1 + or query_start_loc_cpu.numel() < 2 + or query_start_loc_cpu.dtype not in (torch.int32, torch.int64) + ): + raise ValueError( + "grouped QToken-KvBlock-Sparse-Attention requires CPU int32/int64 query_start_loc with at " + "least two entries" + ) + + request_offsets = [int(value) for value in query_start_loc_cpu.tolist()] + if ( + request_offsets[0] != 0 + or request_offsets[-1] < 0 + or request_offsets[-1] > num_query_tokens + or any( + end < begin + for begin, end in zip(request_offsets, request_offsets[1:], strict=False) + ) + ): + raise ValueError( + "grouped QToken-KvBlock-Sparse-Attention query boundaries do not cover valid rows" + ) + + route_offsets = [0] + for begin, end in zip(request_offsets, request_offsets[1:], strict=False): + next_offset = begin + while next_offset < end: + next_offset = min(next_offset + group_size, end) + route_offsets.append(next_offset) + + # Keep graph-padding rows out of the final real request's route even when + # that request ends with a partial group. + next_offset = request_offsets[-1] + while next_offset < num_query_tokens: + next_offset = min(next_offset + group_size, num_query_tokens) + route_offsets.append(next_offset) + + return torch.tensor( + _validate_q_token_kv_block_sparse_route_offsets_cpu( + route_offsets, + num_query_tokens=num_query_tokens, + group_size=group_size, + ), + dtype=torch.int32, + device=device, + ) + + +def _validate_q_token_kv_block_sparse_route_offsets_cpu( + route_offsets: list[int], + *, + num_query_tokens: int, + group_size: int, +) -> tuple[int, ...]: + """Validate CPU-generated QToken-KvBlock-Sparse-Attention routes before copying them to device storage.""" + + offsets = tuple(route_offsets) + if len(offsets) < 2 or offsets[0] != 0 or offsets[-1] != num_query_tokens: + raise ValueError( + "QToken-KvBlock-Sparse-Attention route offsets must cover exactly all query tokens" + ) + lengths = tuple( + end - begin for begin, end in zip(offsets, offsets[1:], strict=False) + ) + if any(length <= 0 or length > group_size for length in lengths): + raise ValueError( + "QToken-KvBlock-Sparse-Attention route offsets must describe nonempty routes no longer than group_size" + ) + return offsets + + +def _prepare_prims_ts_batch_decode_plan( + query: torch.Tensor, + kv_cache: PagedKVCache, + workspace_buffer: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + max_seq_len: int, + *, + seq_len_q: int, + qo_indptr: Optional[torch.Tensor], + max_seq_len_q: Optional[int], + out: Optional[torch.Tensor], + out_dtype: Optional[torch.dtype], + mask_type: Literal["dense", "causal"], + window_left: int, + kv_layout: Literal["HND"], + page_size: Optional[int], + kv_scale_factors: Optional[PagedKVScaleFactors] = None, + q_token_kv_block_sparse_page_memberships: Optional[torch.Tensor] = None, + 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_max_model_len: Optional[int] = None, + direct_q1_sparse_block_size: int = 4, + share_pattern_across_kv_heads: bool = True, +) -> tuple[PrimsTSBatchDecodePlan, torch.Tensor]: + """Validate and freeze one dense-block-table PrimTS launch contract.""" + + _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( seq_len_q=seq_len_q, qo_indptr=qo_indptr, @@ -2262,9 +3369,19 @@ def prims_ts_batch_decode_with_kv_cache( ) assert resolved_seq_len_q is not None seq_len_q = resolved_seq_len_q - metadata_device, batch_size, _ = _validate_block_table_metadata( - block_tables, seq_lens + metadata_device, metadata_rows, max_num_pages = _validate_block_table_metadata( + block_table, + seq_lens, + ) + normalized_cache = _normalize_native_paged_kv_cache( + kv_cache, expected_device=query.device ) + pattern_heads = _num_sparse_pattern_heads( + normalized_cache.num_kv_heads, share_pattern_across_kv_heads + ) + if metadata_rows % pattern_heads: + raise ValueError("metadata rows must be divisible by the pattern head count") + batch_size = metadata_rows // pattern_heads _validate_q( query, seq_len_q=seq_len_q, @@ -2272,30 +3389,35 @@ def prims_ts_batch_decode_with_kv_cache( device=metadata_device, batch_size=batch_size, ) - if metadata_device != query.device: - raise ValueError( - f"paged-KV metadata must be on {query.device}, got {metadata_device}" - ) if qo_indptr is not None: _validate_qo_indptr( qo_indptr, expected_device=query.device, batch_size=batch_size, ) - ( - k_cache, - _, - _, - num_kv_heads, - page_size, - head_dim, - _, - _, - ) = _normalize_paged_kv_cache(kv_cache, expected_device=query.device) + k_cache = normalized_cache.k_cache + v_cache = normalized_cache.v_cache + num_kv_heads = normalized_cache.num_kv_heads + storage_page_size = normalized_cache.storage_page_size + head_dim = normalized_cache.head_dim + if int(query.shape[-1]) != head_dim: + raise ValueError( + "paged_kv_cache logical head dimension must match the query: " + f"expected {int(query.shape[-1])}, got {head_dim}" + ) num_qo_heads = int(query.shape[-2]) _validate_head_geometry(num_qo_heads, num_kv_heads) - _validate_page_size(page_size) + page_size = _validate_page_size( + 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 + if max_num_pages < required_page_columns: + raise ValueError( + "block_table must have at least ceil(max_seq_len / page_size) " + f"columns ({required_page_columns}), got {max_num_pages}" + ) output_dtype = out_dtype if output_dtype is None: if out is not None and not isinstance(out, torch.Tensor): @@ -2303,8 +3425,20 @@ def prims_ts_batch_decode_with_kv_cache( output_dtype = out.dtype if out is not None else query.dtype elif not isinstance(output_dtype, torch.dtype): raise TypeError("out_dtype must be a torch.dtype") - _validate_dtype_pair(query.dtype, k_cache.dtype, output_dtype) - device_index = _validate_runtime_device(query.device) + _validate_dtype_pair( + query.dtype, + k_cache.dtype, + v_cache.dtype, + output_dtype, + allow_fp8_bf16_output=use_q_token_kv_block_sparse_route, + ) + k_sf_cache, v_sf_cache = _normalize_paged_kv_scale_factors( + kv_scale_factors, + k_cache=k_cache, + logical_head_dim=head_dim, + ) + resolved_device, device_index = _resolve_cuda_device(query.device) + _validate_runtime_device(resolved_device) policy_args = ( device_index, @@ -2317,60 +3451,77 @@ def prims_ts_batch_decode_with_kv_cache( seq_len_q, _dtype_key(query.dtype), _dtype_key(k_cache.dtype), + _dtype_key(v_cache.dtype), _dtype_key(output_dtype), kv_layout, mask_type, use_packed_q, window_left, + storage_page_size, + use_q_token_kv_block_sparse_route, + use_pdl, + split_kv, + share_pattern_across_kv_heads, ) spec = _resolve_decode_launch_spec(*policy_args) + if spec.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" + ) + _validate_q_token_kv_block_sparse_page_memberships( + q_token_kv_block_sparse_page_memberships, + expected_device=query.device, + batch_size=metadata_rows, + max_num_pages=max_num_pages, + ) + else: + if q_token_kv_block_sparse_page_memberships is not None: + raise ValueError( + "q_token_kv_block_sparse_page_memberships is valid only for grouped QToken-KvBlock-Sparse-Attention launches" + ) + # Keep one compiled tensor ABI without requiring Q1 or dense callers to + # allocate an unused table. The kernel constexpr-elides all reads. + q_token_kv_block_sparse_page_memberships = block_table layout = _make_decode_workspace_layout( spec.scratch_shapes, output_dtype, use_separate_reduction_kernel=spec.config.use_separate_reduction_kernel, + use_split_kv=spec.config.use_split_kv, ) _validate_workspace_buffer( workspace_buffer, device=query.device, required_bytes=layout.total_bytes, ) - caller_provided_out = out is not None - runtime = _prepare_decode_runtime( - query, - kv_cache, - device=query.device, + output_shape = _decode_output_shape( batch_size=batch_size, - seq_len_q=seq_len_q, - use_packed_q=use_packed_q, num_qo_heads=num_qo_heads, - num_kv_heads=num_kv_heads, + seq_len_q=seq_len_q, head_dim=head_dim, - page_size=page_size, - q_dtype=query.dtype, - kv_dtype=k_cache.dtype, - output_dtype=output_dtype, - bmm1_scale=bmm1_scale, - bmm2_scale=bmm2_scale, - out=out, + total_q_tokens=int(query.shape[0]) if use_packed_q else None, ) - _validate_tensor_does_not_overlap_inputs( - workspace_buffer, - "workspace_buffer", - ("query", runtime.q), - ("k_cache", runtime.k_cache), - ("v_cache", runtime.v_cache), - ("seq_lens", seq_lens), - ("qo_indptr", qo_indptr), - ("block_tables", block_tables), - ("out", runtime.out), - ) - if caller_provided_out: - _validate_decode_output_aliasing( - runtime, - seq_lens=seq_lens, - qo_indptr=qo_indptr, - block_tables=block_tables, - workspace_buffer=workspace_buffer, + if out is None: + out = torch.empty(output_shape, device=query.device, dtype=output_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, + ) + 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 + 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, + page_capacity=max_num_pages, ) compile_spec = _make_decode_compile_spec( spec, @@ -2382,22 +3533,150 @@ def prims_ts_batch_decode_with_kv_cache( 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", + direct_q1_spec=direct_q1_spec, ) compiled_main, compiled_reducer = _get_compiled_decode(compile_spec) workspace = _bind_decode_workspace(workspace_buffer, layout) - return _launch_decode( - runtime, - seq_lens=seq_lens, + # A fused split-KV launch uses a wrapping global completion counter. The + # final arriving CTA restores it to zero, so initialization belongs here, + # once, rather than on every prepared-plan run. The standalone reducer + # does not read this counter, but the same initialization is harmless and + # keeps one replay contract for every split policy. + if layout.uses_split_kv: + workspace.split_kv_counter.zero_() + plan = PrimsTSBatchDecodePlan( + _query_shape=tuple(query.shape), + _query_stride=tuple(query.stride()), + _output_shape=tuple(out.shape), + _output_stride=tuple(out.stride()), + _device=query.device, + _q_dtype=query.dtype, + _output_dtype=output_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, + _block_table=block_table, + _q_token_kv_block_sparse_page_memberships=q_token_kv_block_sparse_page_memberships, + _workspace=workspace, + _compiled_main=compiled_main, + _compiled_reducer=compiled_reducer, + _direct_q1_inputs=direct_q1_inputs, + ) + return plan, out + + +@flashinfer_experimental_api +def prepare_prims_ts_batch_decode_with_kv_cache( + query: torch.Tensor, + kv_cache: PagedKVCache, + workspace_buffer: torch.Tensor, + block_table: torch.Tensor, + seq_lens: torch.Tensor, + max_seq_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, + out_dtype: Optional[torch.dtype] = None, + mask_type: Literal["dense", "causal"] = "dense", + window_left: int = -1, + kv_layout: Literal["HND"] = "HND", + page_size: Optional[int] = None, + split_kv: bool = True, +) -> PrimsTSBatchDecodePlan: + """Validate and prepare a reusable dense-block-table PrimTS launch. + + This setup API is intended for frameworks that keep K/V, block-table + metadata, and workspace storage stable while changing their values between launches. + Call :meth:`PrimsTSBatchDecodePlan.run` on the hot path. The returned plan + is CUDA-graph compatible as long as captured tensor storage remains alive. + + Output and workspace must be disjoint from each other and from all inputs. + Storage overlap is a caller precondition and is not checked. + + Parameters + ---------- + query : torch.Tensor + Fixed ``[B, Hq, D]`` or ``[B, SQ, Hq, D]`` query, or packed + ``[total_q, Hq, D]`` query when ``qo_indptr`` is supplied. + kv_cache : torch.Tensor or tuple[torch.Tensor, torch.Tensor] + Combined or separate HND paged K/V storage. + workspace_buffer : torch.Tensor + Caller-owned byte workspace, zero-initialized before the first launch + for this semantic configuration. + block_table : torch.Tensor + Contiguous CUDA Int32 dense page table shaped ``[B, max_pages]``. + seq_lens : torch.Tensor + Positive live K/V sequence lengths, one per request. + max_seq_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. + kv_scale_factors : tuple[torch.Tensor, torch.Tensor], optional + Required for packed NVFP4 K/V stored as uint8. K and V scales are + FP8 tensors with width ``D / 16``; V scales use the TRTLLM-GEN + 4-token interleaved layout. The plan retains these tensors. + seq_len_q : int + 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 + Static maximum packed-query route length. + out_dtype : torch.dtype, optional + Output dtype; defaults to ``out.dtype``. + mask_type : {"dense", "causal"} + Attention mask mode. + window_left : int + Left sliding-window extent, or ``-1`` to disable the window. + kv_layout : {"HND"} + Layout of the paged K/V cache. + page_size : int, optional + Semantic page size; defaults to the physical cache-page extent. + split_kv : bool + Permit automatic split fanout (True, default), or force S1 (False). + Independent of packed/fixed Q. Match workspace sizing's value. + + Returns + ------- + PrimsTSBatchDecodePlan + Reusable prepared launch plan. Call :meth:`PrimsTSBatchDecodePlan.run` + after updating values in the prepared tensor storage. + """ + + if not isinstance(out, torch.Tensor): + raise TypeError("out must be a caller-owned torch.Tensor") + plan, _ = _prepare_prims_ts_batch_decode_plan( + query, + kv_cache, + workspace_buffer, + block_table, + seq_lens, + max_seq_len, + seq_len_q=seq_len_q, qo_indptr=qo_indptr, - block_tables=block_tables, - workspace=workspace, - compiled_main=compiled_main, - compiled_reducer=compiled_reducer, + max_seq_len_q=max_seq_len_q, + out=out, + out_dtype=out_dtype, + mask_type=mask_type, + window_left=window_left, + kv_layout=kv_layout, + page_size=page_size, + kv_scale_factors=kv_scale_factors, + split_kv=split_kv, + use_q_token_kv_block_sparse_route=False, ) + return plan class BatchDecodePagedTSWrapper: @@ -2413,7 +3692,7 @@ class BatchDecodePagedTSWrapper: graph replays require separate wrappers and workspace buffers. """ - @flashinfer_api + @flashinfer_experimental_api def __init__(self, kv_layout: Literal["HND"] = "HND") -> None: """Initialize an unplanned wrapper with one static K/V layout. @@ -2439,7 +3718,7 @@ def _policy(self) -> tuple[tuple[str, object], ...]: return self._require_plan_state().policy - @flashinfer_api + @flashinfer_experimental_api def plan( self, device: Union[int, str, torch.device], @@ -2453,12 +3732,18 @@ def plan( max_seq_len_q: int = 1, 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, mask_type: Literal["dense", "causal"] = "dense", window_left: int = -1, seq_lens: 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. @@ -2475,8 +3760,9 @@ def plan( every run must supply a CUDA length tensor. ``workspace_buffer`` is caller-owned scratch for this plan. It is - allocated when omitted, initialized during planning, retained by the - frozen plan state, and never reset by ``run``. + 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 ---------- @@ -2491,7 +3777,10 @@ def plan( head_dim : int Query, key, value, and output head dimension. page_size : int - Number of K/V tokens stored in each page. + 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. @@ -2505,8 +3794,16 @@ def plan( 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 - K/V dtype used to compile the plan. Defaults to ``q_data_type``. + 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``. @@ -2525,49 +3822,81 @@ def plan( must be 32-byte aligned and large enough for the selected plan. When omitted, planning allocates the buffer. The retained buffer is exclusive to one in-flight launch or graph replay. + 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. """ - if not isinstance(packed_query, bool): - raise TypeError("packed_query must be a bool") - batch_size = _validate_positive_int(batch_size, "batch_size") - head_dim = _validate_head_dim(head_dim) - page_size = _validate_page_size(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) - _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, + k_data_type, v_data_type = _resolve_kv_dtypes( + q_data_type, k_data_type, v_data_type, kv_data_type ) - _validate_mask(mask_type) - window_left = _validate_window_left(window_left, mask_type) - - if kv_data_type is None: - kv_data_type = q_data_type if o_data_type is None: o_data_type = q_data_type - _validate_dtype_pair(q_data_type, kv_data_type, o_data_type) - - specialization_seq_lens = _normalize_plan_seq_lens( - seq_lens, - batch_size=batch_size, - max_kv_len=max_kv_len, + seq_len_q = max_seq_len_q + storage_page_size = ( + page_size if storage_page_size is None else storage_page_size ) - if ( - specialization_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: - 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={seq_len_q} and K/V={kv_len}" - ) + if not isinstance(validate, bool): + raise TypeError("validate must be a bool") + if not isinstance(initialize_workspace, bool): + raise TypeError("initialize_workspace must be a bool") + if validate: + if not isinstance(packed_query, bool): + raise TypeError("packed_query must be a bool") + batch_size = _validate_positive_int(batch_size, "batch_size") + head_dim = _validate_head_dim(head_dim) + page_size = _validate_page_size(page_size) + storage_page_size = _validate_storage_page_size( + 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) + _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, + ) + _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) + + specialization_seq_lens = _normalize_plan_seq_lens( + seq_lens, + batch_size=batch_size, + max_kv_len=max_kv_len, + ) + if ( + specialization_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: + 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={seq_len_q} and K/V={kv_len}" + ) + + else: + specialization_seq_lens = _normalize_plan_seq_lens( + seq_lens, batch_size=batch_size, max_kv_len=max_kv_len + ) device, device_index = _resolve_cuda_device(device) + if validate: + _validate_runtime_device(device) policy_args = ( device_index, @@ -2579,14 +3908,16 @@ def plan( max_kv_len, seq_len_q, _dtype_key(q_data_type), - _dtype_key(kv_data_type), + _dtype_key(k_data_type), + _dtype_key(v_data_type), _dtype_key(o_data_type), self._kv_layout, mask_type, packed_query, window_left, + storage_page_size, ) - spec = _resolve_decode_launch_spec(*policy_args) + 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" @@ -2629,6 +3960,8 @@ def plan( 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, kv_prefix_mode=kv_prefix_mode, @@ -2638,28 +3971,32 @@ def plan( policy = spec.policy + ( ("kv_prefix_mode", kv_prefix_mode), ("kv_lengths_mode", kv_lengths_mode), + ("split_kv_allowed", split_kv), ) workspace_layout = _make_decode_workspace_layout( spec.scratch_shapes, o_data_type, use_separate_reduction_kernel=(spec.config.use_separate_reduction_kernel), + use_split_kv=spec.config.use_split_kv, ) if workspace_buffer is None: + initialize_workspace = True workspace_buffer = torch.empty( workspace_layout.total_bytes, device=device, dtype=torch.int8, ) - else: + elif validate: _validate_workspace_buffer( workspace_buffer, device=device, required_bytes=workspace_layout.total_bytes, ) workspace = _bind_decode_workspace(workspace_buffer, workspace_layout) - workspace.split_kv_counter.zero_() - workspace.cu_seqlens_q.zero_() - workspace.attention_sinks.zero_() + if initialize_workspace: + workspace.split_kv_counter.zero_() + workspace.cu_seqlens_q.zero_() + workspace.attention_sinks.zero_() # Materialize from the normalized tuple so the plan never aliases # caller-owned host or device storage. planned_seq_lens_device = ( @@ -2684,7 +4021,8 @@ def plan( page_size=page_size, max_kv_len=max_kv_len, q_dtype=q_data_type, - kv_dtype=kv_data_type, + k_dtype=k_data_type, + v_dtype=v_data_type, output_dtype=o_data_type, mask_type=mask_type, window_left=window_left, @@ -2699,12 +4037,13 @@ def plan( planned_seq_lens_host=specialization_seq_lens, planned_seq_lens_device=planned_seq_lens_device, policy=policy, + storage_page_size=storage_page_size, ) # This is the only wrapper mutation. Any failure above leaves the # previous complete plan revision usable. self._plan_state = candidate - @flashinfer_api + @flashinfer_experimental_api def run( self, q: torch.Tensor, @@ -2712,6 +4051,7 @@ def run( seq_lens: 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, @@ -2724,9 +4064,8 @@ def run( ``None`` and the plan's owned CUDA copy is used. Otherwise, ``seq_lens`` must be supplied on every run. This ownership check is unconditional. - ``validate=True`` performs structural, value, and alias validation. It - reads per-run metadata values back to the host. This is the safe public - default. + ``validate=True`` performs structural and value validation. It + reads per-run metadata values back to the host. ``validate=False`` treats every run argument as a trusted binding, performs no explicit wrapper validation, and remains free of metadata device-to-host synchronization. K/V view selection, scale forwarding, @@ -2736,6 +4075,9 @@ def run( plans require ``qo_indptr`` to be omitted. Per-run metadata tensors may change identity between ordered runs. + 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 @@ -2750,6 +4092,11 @@ def run( 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. @@ -2761,7 +4108,7 @@ def run( out : torch.Tensor, optional Caller-owned output tensor. A new tensor is allocated when omitted. validate : bool - Run explicit structural, value, and alias validation. Disable only + 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``. @@ -2790,7 +4137,6 @@ def run( effective_seq_lens = seq_lens runtime_qo_indptr = qo_indptr if state.use_packed_q else None - caller_provided_out = out is not None if validate: ( metadata_device, @@ -2827,6 +4173,7 @@ def run( runtime = _prepare_decode_runtime( q, paged_kv_cache, + kv_scale_factors, device=state.device, batch_size=state.batch_size, seq_len_q=state.seq_len_q, @@ -2834,45 +4181,34 @@ def run( num_qo_heads=state.num_qo_heads, num_kv_heads=state.num_kv_heads, head_dim=state.head_dim, - page_size=state.page_size, + page_size=state.storage_page_size or state.page_size, q_dtype=state.q_dtype, - kv_dtype=state.kv_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=out, ) _validate_decode_run_metadata_values( - state, runtime, + planned_seq_lens_host=state.planned_seq_lens_host, + max_kv_len=state.max_kv_len, + page_size=state.page_size, + use_packed_q=state.use_packed_q, + seq_len_q=state.seq_len_q, + batch_size=state.batch_size, + mask_type=state.mask_type, seq_lens=effective_seq_lens, block_tables=block_tables, qo_indptr=runtime_qo_indptr, ) - _validate_tensor_does_not_overlap_inputs( - state.workspace_buffer, - "workspace_buffer", - ("query", runtime.q), - ("k_cache", runtime.k_cache), - ("v_cache", runtime.v_cache), - ("seq_lens", effective_seq_lens), - ("qo_indptr", runtime_qo_indptr), - ("block_tables", block_tables), - ("out", runtime.out), - ) - if caller_provided_out: - _validate_decode_output_aliasing( - runtime, - seq_lens=effective_seq_lens, - qo_indptr=runtime_qo_indptr, - block_tables=block_tables, - workspace_buffer=state.workspace_buffer, - ) else: runtime = _prepare_decode_runtime_unchecked( q, paged_kv_cache, - state=state, + kv_scale_factors, + output_dtype=state.output_dtype, bmm1_scale=bmm1_scale, bmm2_scale=bmm2_scale, out=out, @@ -2889,13 +4225,14 @@ def run( ) -@flashinfer_api +@flashinfer_experimental_api def batch_decode_with_paged_kv_cache( q: torch.Tensor, paged_kv_cache: PagedKVCache, block_tables: torch.Tensor, seq_lens_kv: 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, @@ -2906,6 +4243,11 @@ def batch_decode_with_paged_kv_cache( bmm2_scale: float = 1.0, out: Optional[torch.Tensor] = None, 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: """One-shot fixed or packed-Q paged decode from fixed page tables. @@ -2915,23 +4257,27 @@ def batch_decode_with_paged_kv_cache( ``[total_q, Hq, D]`` query/output; the wrapper derives ``max_seq_len_q`` once when it is omitted. No transpose is hidden here. - The one-shot planner reads ``seq_lens_kv`` on the host to derive exact plan - bounds, so this convenience API is not CUDA-graph-capturable. Capture - callers should plan :class:`BatchDecodePagedTSWrapper` before capture and - either bind a stable runtime ``seq_lens`` tensor or let the plan retain - fixed sequence lengths before replay. + Without caller scratch, the convenience planner reads ``seq_lens_kv`` on + the host and retains the existing length-specialized kernel path. This mode + is not CUDA-graph-capturable. See Notes for the explicit capture-safe mode. Parameters ---------- q : torch.Tensor Fixed or packed query tensor. paged_kv_cache : torch.Tensor or tuple[torch.Tensor, torch.Tensor] - Combined or separate paged K/V storage. + Combined or separate paged K/V storage. K and V dtypes are taken from + the tensors; a ``torch.float8_e4m3fn`` V with ``torch.bfloat16`` Q/K + requires the separate ``(K, V)`` form. block_tables : torch.Tensor Fixed row-strided ``[B, C]`` page table. Rows may have padding between them, but each row must be contiguous. seq_lens_kv : torch.Tensor Per-request K/V sequence lengths with shape ``[B]``. + 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. 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 @@ -2954,151 +4300,176 @@ def batch_decode_with_paged_kv_cache( Caller-owned output tensor. out_dtype : torch.dtype, optional Output dtype; defaults to ``out.dtype`` or the query dtype. + page_size : int, optional + Semantic page-table size, defaulting to the physical cache extent. + A smaller supported value must divide the physical extent and uses + encoded subpage locators. + split_kv : bool + Allow automatic split fanout (True, default), or force S1 (False), + independently of packed/fixed query storage. + workspace_buffer : torch.Tensor, optional + Zero-initialized caller-owned byte scratch; re-zero when layout inputs + (including batch size) change. Exclusive to one in-flight launch/graph. + max_kv_len : int, optional + Static K/V capacity; defaults to the metadata maximum with validation. + 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. Returns ------- torch.Tensor The fixed or packed attention output. + Notes + ----- + 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). + Warm this exact topology outside capture first. Retain stable tensor storage + and mutate metadata only between completed launches/replays. All live K/V + lengths must be positive and within the static bound, active page IDs must + index the cache, and causal per-request Q lengths must not exceed K/V lengths. + 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 + any inputs or each other. No metadata is copied to the host on the trusted + explicit path. Kernel policy, necessary control resets, and output layout + are unchanged. Missing output may be allocated only outside capture. + """ - _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( - seq_len_q=seq_len_q, - qo_indptr=qo_indptr, - max_seq_len_q=max_seq_len_q, - require_packed_max=False, - ) - metadata_device, batch_size, _ = _validate_block_table_metadata( - block_tables, - seq_lens_kv, - ) - if metadata_device != q.device: - raise ValueError( - f"paged-KV metadata must be on {q.device}, got {metadata_device}" + if not isinstance(validate, bool): + raise TypeError("validate must be a bool") + if validate or workspace_buffer is None or out is None: + if torch.cuda.is_initialized() and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "CUDA graph capture requires workspace_buffer, max_kv_len, " + "out, validate=False, and explicit packed-Q bounds; warm up first" + ) + if not validate: + if workspace_buffer is None: + 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") + + allocate_workspace = workspace_buffer is None + if validate: + use_packed_q, 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, ) - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "batch_decode_with_paged_kv_cache cannot derive host plan bounds from " - "seq_lens_kv during CUDA graph capture; plan BatchDecodePagedTSWrapper " - "before capture" + metadata_device, batch_size, _ = _validate_block_table_metadata( + block_tables, seq_lens_kv ) - if qo_indptr is not None: - _validate_qo_indptr( - qo_indptr, - expected_device=q.device, + if metadata_device != q.device: + raise ValueError( + f"paged-KV metadata must be on {q.device}, got {metadata_device}" + ) + if qo_indptr is not None: + _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, + device=q.device, batch_size=batch_size, ) - derived_max_q, total_q, validation_q_lengths = _read_packed_q_plan_metadata( - qo_indptr + k_cache, v_cache, _, num_kv_heads, storage_page_size, head_dim = ( + _normalize_paged_kv_cache_views(paged_kv_cache, expected_device=q.device) ) - if resolved_seq_len_q is None: - validation_seq_len_q = _validate_seq_len_q(derived_max_q) - else: - validation_seq_len_q = resolved_seq_len_q - if derived_max_q > validation_seq_len_q: - raise ValueError( - "qo_indptr contains a per-request Q length larger than " - f"max_seq_len_q ({validation_seq_len_q}): got {derived_max_q}" - ) - if total_q != int(q.shape[0]): - raise ValueError( - "the final qo_indptr offset must equal the packed q token count: " - f"expected {q.shape[0]}, got {total_q}" - ) else: - assert resolved_seq_len_q is not None - validation_seq_len_q = resolved_seq_len_q - validation_q_lengths = (validation_seq_len_q,) * batch_size - _validate_q( - q, - seq_len_q=validation_seq_len_q, - use_packed_q=use_packed_q, - device=metadata_device, - batch_size=batch_size, - ) - ( - k_cache, - _, - _, - num_kv_heads, - page_size, - head_dim, - _, - _, - ) = _normalize_paged_kv_cache(paged_kv_cache, expected_device=q.device) - _validate_page_size(page_size) + use_packed_q = 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 + ) + batch_size = int(seq_lens_kv.shape[0]) + k_cache = ( + paged_kv_cache[:, 0] + if isinstance(paged_kv_cache, torch.Tensor) + else paged_kv_cache[0] + ) + v_cache = ( + paged_kv_cache[:, 1] + if isinstance(paged_kv_cache, torch.Tensor) + else paged_kv_cache[1] + ) + 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 + ) + + page_size = storage_page_size if page_size is None else page_size num_qo_heads = int(q.shape[-2]) - _validate_head_geometry(num_qo_heads, num_kv_heads) - output_dtype = out_dtype - if output_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 q.dtype - elif not isinstance(output_dtype, torch.dtype): - raise TypeError("out_dtype must be a torch.dtype") - if out is not None: + 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 + 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()) + ) + if validate and out is not None: _validate_out( out, q=q, expected_shape=_decode_output_shape( batch_size=batch_size, num_qo_heads=num_qo_heads, - seq_len_q=validation_seq_len_q, + seq_len_q=resolved_seq_len_q, head_dim=head_dim, total_q_tokens=int(q.shape[0]) if use_packed_q else None, ), - seq_len_q=validation_seq_len_q, + seq_len_q=resolved_seq_len_q, use_packed_q=use_packed_q, output_dtype=output_dtype, ) - _validate_dtype_pair( - q.dtype, - k_cache.dtype, - output_dtype, - ) - - seq_lens_host = tuple(int(value) for value in seq_lens_kv.tolist()) - for request_idx, seq_len in enumerate(seq_lens_host): - if seq_len <= 0: - raise ValueError( - "seq_lens_kv values must be positive; request " - f"{request_idx} has {seq_len}" - ) - max_kv_len = _validate_max_kv_len(max(seq_lens_host), "max(seq_lens_kv)") - table_capacity = int(block_tables.shape[1]) - block_table_rows = block_tables.tolist() - for request_idx, (row, q_len, kv_len) in enumerate( - zip( - block_table_rows, - validation_q_lengths, - seq_lens_host, - strict=True, + 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 ) - ): - required_pages = (kv_len + page_size - 1) // page_size - if table_capacity < required_pages: - raise ValueError( - "block_tables does not have enough columns for " - f"seq_lens_kv[{request_idx}]={kv_len}: requires " - f"{required_pages}, got {table_capacity}" - ) - if any( - int(page_id) < 0 or int(page_id) >= int(k_cache.shape[0]) - for page_id in row[:required_pages] - ): - raise ValueError( - "block_tables values for active pages must index the physical " - f"K/V cache in [0, {k_cache.shape[0]}); request {request_idx} " - "contains an invalid page ID" - ) - if mask_type == "causal" and 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}" - ) wrapper = BatchDecodePagedTSWrapper(kv_layout=kv_layout) wrapper.plan( @@ -3109,30 +4480,42 @@ def batch_decode_with_paged_kv_cache( head_dim, page_size, max_kv_len, - max_seq_len_q=validation_seq_len_q, + max_seq_len_q=resolved_seq_len_q, + storage_page_size=storage_page_size, + split_kv=split_kv, packed_query=use_packed_q, q_data_type=q.dtype, - kv_data_type=k_cache.dtype, + k_data_type=k_cache.dtype, + v_data_type=v_cache.dtype, o_data_type=output_dtype, mask_type=mask_type, window_left=window_left, seq_lens=seq_lens_host, + workspace_buffer=workspace_buffer, + validate=validate, + initialize_workspace=allocate_workspace, ) return wrapper.run( q, paged_kv_cache, - None, + None if seq_lens_host is not None else seq_lens_kv, block_tables, + kv_scale_factors=kv_scale_factors, qo_indptr=qo_indptr, bmm1_scale=bmm1_scale, bmm2_scale=bmm2_scale, out=out, + validate=validate, ) __all__ = [ "BatchDecodePagedTSWrapper", + "PrimsTSBatchDecodePlan", "batch_decode_with_paged_kv_cache", "get_prims_ts_batch_decode_workspace_size", - "prims_ts_batch_decode_with_kv_cache", + "make_q_token_kv_block_sparse_qo_indptr", + "suggest_q_token_kv_block_sparse_group_size", + "validate_q_token_kv_block_sparse_group_size", + "prepare_prims_ts_batch_decode_with_kv_cache", ] 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 bc2b51d5f114..c74fbcc1287c 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 @@ -43,7 +43,7 @@ | Variable seqlen | Supported through flattened tensors + `cum_seqlen_*` | | dtype | fp16/bf16/e4m3 Q/K/V/O, fp32 QK and PV accumulation | | Masking | non-causal and causal; causal left sliding window uses head-paired mode | - | head dimension | D=128 and D=256 | + | head dimension | QK/V=128/128, 192/128, and 256/256 | | `S_q` / `S_kv` | Query-paired causal requires `S_q <= S_kv`; arbitrary positive tails are supported | | GQA | Must satisfy `h_q % h_kv == 0`; causal GQA can use head-paired scheduling | | Sliding window | `mask_type="causal", window_left=N`; left window only | @@ -154,6 +154,7 @@ from cutlass.experimental.task_scheduling.resources import ( MemoryResource, PipelineConfig, + SignalingThreads, TileSchedulerConfig, WorkQueue, ) @@ -172,9 +173,11 @@ SmemQResource, TmemOResource, TmemPResource, + SmemPResource, TmemSPResource, TmemStatsResource, TmemStatsDoneResource, + TmemPPrefixReadyResource, ) from .fmha_tasks import ( PackedContextWorkQueue, @@ -690,6 +693,7 @@ def build_context_task_manager( list[MemoryResource], SmemAllocation, SmemAllocation, + SmemAllocation, WorkQueue | None, SmemAllocation | None, ]: @@ -747,7 +751,7 @@ def build_context_task_manager( ------- tuple ``(TaskManager, tmem_resources, tmem_ptr_alloc, dealloc_mbar_alloc, - work_queue, clc_response_alloc)``. + peer_dealloc_mbar_alloc, work_queue, clc_response_alloc)``. """ if cfg.use_paged_kv: if cfg.num_tokens_per_page not in _SUPPORTED_CONTEXT_PAGE_SIZES: @@ -762,9 +766,15 @@ def build_context_task_manager( # Cluster / CTA layout # --------------------------------------------------------------------------- # CTA layout in CuTe VMNK order. V is the CTA-group dimension used for - # 2-CTA cooperative MMA; MNK are the logical cluster axes. FMHA uses a - # single-CTA cluster here, so (V, M, N, K) = (1, 1, 1, 1). - cluster_shape_vmnk = (1, 1, 1, 1) + # 2-CTA cooperative MMA; MNK are the logical cluster axes. V is 2 in the + # two-CTA UMMA form and 1 otherwise. + cluster_shape_vmnk = (cfg.cta_group_size, 1, 1, 1) + two_cta = cfg.two_cta_umma + # With two CTAs, only the leader's UMMA consumes K/V, so both CTAs' TMA + # loads signal the leader's barrier. + tma_umma_leader_kwargs = ( + {"consumer_signaling_threads": SignalingThreads.CtaLeader} if two_cta else {} + ) # --------------------------------------------------------------------------- # Warp counts @@ -794,25 +804,42 @@ def build_context_task_manager( # advance_on_wait=True advances consumer_state on ConsumerWait. smem_q_pipeline_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( num_stages=cfg.q_stage, - num_bytes=cfg.tma_copy_q_bytes, + num_bytes=cfg.tma_copy_q_bytes * cfg.cta_group_size, 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.tma_copy_q_bytes if two_cta else None, + **tma_umma_leader_kwargs, ) # SmemKV: capacity-derived stages, Load -> MMA. # advance_on_wait=True advances consumer_state at ConsumerWait rather than # ConsumerRelease. This keeps the previous V tile live while MMA starts the # next QK tile, giving QK0 -> PV1(previous V) -> QK1 ordering without # releasing the previous V tile first. + # K and V share this pipeline unless their dtypes differ in width. smem_kv_pipeline_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( - num_stages=cfg.kv_stage, - num_bytes=cfg.tma_copy_kv_bytes, + num_stages=cfg.kv_stage_k, + num_bytes=cfg.tma_copy_kv_bytes * cfg.cta_group_size, 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.tma_copy_kv_bytes if two_cta else None, + **tma_umma_leader_kwargs, ) + smem_v_pipeline_cfg = None + if cfg.split_kv_pipelines: + smem_v_pipeline_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( + num_stages=cfg.kv_stage_v, + num_bytes=cfg.tma_copy_v_bytes * cfg.cta_group_size, + 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.tma_copy_v_bytes if two_cta else None, + **tma_umma_leader_kwargs, + ) # Page-offsets prefetch (paged-KV only): the auxiliary warp produces and # the load warp consumes. # Async-async because both are CUDA-thread groups (no TMA arrival barrier). @@ -873,18 +900,32 @@ def make_page_offsets_pipeline_cfg(num_stages: int) -> PipelineConfig: ) umma_hw_group = pipeline.CooperativeGroup(Agent.Thread) mma_group = pipeline.CooperativeGroup(Agent.Thread, num_mma_warps * warp_size) + # S/P and O barriers live on the leader CTA. Their arrival counts must cover + # the consumer threads of both CTAs, and only the leader's UMMA commits them. + softmax_group_umma = pipeline.CooperativeGroup( + Agent.Thread, + num_softmax_warps * warp_size * cfg.cta_group_size, + ) + correction_group_umma = pipeline.CooperativeGroup( + Agent.Thread, num_correction_warps * warp_size * cfg.cta_group_size + ) + umma_leader_kwargs = ( + {"producer_signaling_threads": SignalingThreads.CtaLeader} if two_cta else {} + ) tmem_sp0_pipeline_cfg = PipelineConfig.create_umma_async_pipeline_cfg( num_stages=cfg.mma_softmax_stage, producer_group=umma_hw_group, - consumer_group=softmax_group, + consumer_group=softmax_group_umma, cta_layout_vmnk=cluster_shape_vmnk, + **umma_leader_kwargs, ) tmem_sp1_pipeline_cfg = PipelineConfig.create_umma_async_pipeline_cfg( num_stages=cfg.mma_softmax_stage, producer_group=umma_hw_group, - consumer_group=softmax_group, + consumer_group=softmax_group_umma, cta_layout_vmnk=cluster_shape_vmnk, + **umma_leader_kwargs, ) tmem_p0_pipeline_cfg = PipelineConfig.create_async_umma_pipeline_cfg( num_stages=cfg.mma_softmax_stage, @@ -895,8 +936,9 @@ def make_page_offsets_pipeline_cfg(num_stages: int) -> PipelineConfig: tmem_o_pipeline_cfg = PipelineConfig.create_umma_async_pipeline_cfg( num_stages=cfg.mma_corr_stage, producer_group=umma_hw_group, - consumer_group=correction_group, + consumer_group=correction_group_umma, cta_layout_vmnk=cluster_shape_vmnk, + **umma_leader_kwargs, ) tmem_vec0_pipeline_cfg = PipelineConfig.create_async_async_pipeline_cfg( @@ -999,16 +1041,33 @@ def make_page_offsets_pipeline_cfg(num_stages: int) -> PipelineConfig: page_table_is_v=True, name="smem_page_offsets_v", ) - smem_kv = SmemKVResource( - tma_k_desc=tma_k_desc, - tma_v_desc=tma_v_desc, - pipeline_config=smem_kv_pipeline_cfg, + # smem_kv only stages K when we have mixed QK/PV dtypes, + # and K and V in alternating stages when QK/PV are the same dtype. + split_kv = cfg.split_kv_pipelines + kv_resource_kwargs = dict( cfg=cfg, page_offsets_kv=smem_page_offsets_kv, page_offsets_v=smem_page_offsets_v, block_tables=g_block_tables, - name="smem_kv", ) + smem_kv = SmemKVResource( + tma_k_desc=tma_k_desc, + tma_v_desc=None if split_kv else tma_v_desc, + pipeline_config=smem_kv_pipeline_cfg, + role="k" if split_kv else "kv", + name="smem_k" if split_kv else "smem_kv", + **kv_resource_kwargs, + ) + smem_v = None + if split_kv: + smem_v = SmemKVResource( + tma_k_desc=None, + tma_v_desc=tma_v_desc, + pipeline_config=smem_v_pipeline_cfg, + role="v", + name="smem_v", + **kv_resource_kwargs, + ) # WorkQueue: persistent tile scheduler state (static or CLC dynamic), not # Q/K/V/O dataflow. Non-persistent launches omit it so each CTA executes @@ -1083,6 +1142,28 @@ def make_page_offsets_pipeline_cfg(num_stages: int) -> PipelineConfig: else: work_queue = WorkQueue(**work_queue_kwargs) + # p_in_smem: one SMEM P tile per query group; softmax produces, UMMA consumes. + smem_p_resources: list[SmemPResource | None] = [None, None] + if cfg.p_in_smem: + smem_p_resources = [ + SmemPResource( + # Under two-CTA UMMA the leader's M=256 PV reads both CTAs' P + # tiles, so both softmax groups signal the leader's barrier and + # the UMMA commit releases both. + pipeline_config=PipelineConfig.create_async_umma_pipeline_cfg( + num_stages=1, + producer_group=softmax_group_umma if two_cta else softmax_group, + consumer_group=umma_hw_group, + cta_layout_vmnk=cluster_shape_vmnk, + **tma_umma_leader_kwargs, + ), + cfg=cfg, + group_idx=index, + name=f"smem_p{index}", + ) + for index in range(cfg.num_qkv_instances) + ] + smem_p0, smem_p1 = smem_p_resources tmem_sp0 = TmemSPResource( pipeline_config=tmem_sp0_pipeline_cfg, cfg=cfg, @@ -1098,6 +1179,7 @@ def make_page_offsets_pipeline_cfg(num_stages: int) -> PipelineConfig: 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, name="tmem_sp0", ) tmem_p0: TmemPResource | None = None @@ -1136,6 +1218,35 @@ def make_page_offsets_pipeline_cfg(num_stages: int) -> PipelineConfig: pipeline_config=tmem_stats_done_0_pipeline_cfg, name="tmem_stats_done_0", ) + # P-prefix barriers: Softmax (producer) arrives after storing the first 64 + # keys of P, MMA (consumer) starts PV on them. One per query group. + tmem_p_prefix_ready: list[TmemPPrefixReadyResource | None] = [None, None] + if cfg.pv_half_overlap: + tmem_p_prefix_ready = [ + TmemPPrefixReadyResource( + # With two CTAs the P-half barrier is consumed by the leader's UMMA, + # so it needs the AsyncUmma form. + pipeline_config=( + PipelineConfig.create_async_umma_pipeline_cfg( + num_stages=1, + producer_group=softmax_group_umma, + consumer_group=umma_hw_group, + cta_layout_vmnk=cluster_shape_vmnk, + **tma_umma_leader_kwargs, + ) + if two_cta + else PipelineConfig.create_async_async_pipeline_cfg( + num_stages=1, + producer_group=softmax_group, + consumer_group=mma_group, + cta_layout_vmnk=cluster_shape_vmnk, + ) + ), + name=f"tmem_p_prefix_ready_{index}", + ) + for index in range(2) + ] + tmem_p_prefix_ready_0, tmem_p_prefix_ready_1 = tmem_p_prefix_ready single_qkv_instance = cfg.single_qkv_instance tmem_sp1: TmemSPResource | None = None @@ -1161,6 +1272,7 @@ def make_page_offsets_pipeline_cfg(num_stages: int) -> PipelineConfig: 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, name="tmem_sp1", ) tmem_vec1 = TmemStatsResource( @@ -1204,6 +1316,8 @@ def make_page_offsets_pipeline_cfg(num_stages: int) -> PipelineConfig: tmem_o1_offset=cfg.tmem_o1_offset, tmem_vec0_resource=tmem_vec0, tmem_vec1_resource=tmem_vec1, + smem_p0_resource=smem_p0, + smem_p1_resource=smem_p1, name="tmem_o", **tmem_o_kwargs, ) @@ -1237,6 +1351,7 @@ def scheduler_deps(*resources: MemoryResource) -> list[MemoryResource]: gmem_qkv, smem_q, smem_kv, + smem_v, work_queue, smem_page_offsets_kv=smem_page_offsets_kv, smem_page_offsets_v=smem_page_offsets_v, @@ -1246,6 +1361,7 @@ def scheduler_deps(*resources: MemoryResource) -> list[MemoryResource]: gmem_qkv, smem_q, smem_kv, + smem_v, tmem_sp0, tmem_sp1, tmem_p0, @@ -1253,6 +1369,10 @@ def scheduler_deps(*resources: MemoryResource) -> list[MemoryResource]: tmem_stats_done_0, tmem_stats_done_1, work_queue, + tmem_p_prefix_ready_0=tmem_p_prefix_ready_0, + tmem_p_prefix_ready_1=tmem_p_prefix_ready_1, + smem_p0=smem_p0, + smem_p1=smem_p1, **mma_domain_kwargs, ) @@ -1265,6 +1385,8 @@ def scheduler_deps(*resources: MemoryResource) -> list[MemoryResource]: tmem_p0, s0s1_seq, work_queue, + tmem_p_prefix_ready=tmem_p_prefix_ready_0, + smem_p=smem_p0, **softmax0_domain_kwargs, ) softmax1_task: Task | None = None @@ -1278,6 +1400,8 @@ def scheduler_deps(*resources: MemoryResource) -> list[MemoryResource]: None, s0s1_seq, work_queue, + tmem_p_prefix_ready=tmem_p_prefix_ready_1, + smem_p=smem_p1, **softmax1_domain_kwargs, ) correction_task = create_correction_task( @@ -1378,34 +1502,56 @@ def scheduler_deps(*resources: MemoryResource) -> list[MemoryResource]: or tmem_stats_done_1 is None ): raise ValueError("paired resource graph requires peer-1 resources") - tmem_o_dependencies = scheduler_deps(tmem_sp0, tmem_sp1) + tmem_o_dependencies = scheduler_deps( + tmem_sp0, + tmem_sp1, + *([smem_p0, smem_p1] if cfg.p_in_smem else []), + ) smem_kv_deps: list[MemoryResource] = [gmem_qkv] if smem_page_offsets_kv is not None: smem_kv_deps.append(smem_page_offsets_kv) - if smem_page_offsets_v is not None: + if smem_page_offsets_v is not None and not cfg.split_kv_pipelines: smem_kv_deps.append(smem_page_offsets_v) stats_done_0_deps = [] if cfg.stats_via_smem else [tmem_stats_done_0] + # One entry when K and V share the buffer, and + # two otherwise (e.g. when K/V are different dtypes) + smem_kv_resources = (smem_kv, smem_v) if cfg.split_kv_pipelines else (smem_kv,) resource_dependency_graph: dict[MemoryResource, list[MemoryResource]] = { smem_q: scheduler_deps(gmem_qkv), smem_kv: scheduler_deps(*smem_kv_deps), - tmem_sp0: scheduler_deps(tmem_sp0, smem_q, smem_kv, *stats_done_0_deps), - tmem_vec0: scheduler_deps(tmem_sp0), - tmem_o: tmem_o_dependencies, - smem_o_0: scheduler_deps(tmem_vec0, tmem_o), - gmem_o_0: scheduler_deps(smem_o_0), } + if cfg.split_kv_pipelines: + smem_v_deps: list[MemoryResource] = [gmem_qkv] + 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) + resource_dependency_graph.update( + { + tmem_sp0: scheduler_deps( + tmem_sp0, smem_q, *smem_kv_resources, *stats_done_0_deps + ), + tmem_vec0: scheduler_deps(tmem_sp0), + tmem_o: tmem_o_dependencies, + smem_o_0: scheduler_deps(tmem_vec0, tmem_o), + gmem_o_0: scheduler_deps(smem_o_0), + } + ) if not cfg.stats_via_smem: resource_dependency_graph[tmem_stats_done_0] = [tmem_vec0] if tmem_p0 is not None: resource_dependency_graph[tmem_p0] = scheduler_deps(tmem_sp0) + if smem_p0 is not None: + resource_dependency_graph[smem_p0] = scheduler_deps(tmem_sp0) + if smem_p1 is not None and tmem_sp1 is not None: + resource_dependency_graph[smem_p1] = scheduler_deps(tmem_sp1) if not single_qkv_instance: resource_dependency_graph.update( { tmem_sp1: scheduler_deps( tmem_sp1, smem_q, - smem_kv, + *smem_kv_resources, s0s1_seq, *([] if cfg.stats_via_smem else [tmem_stats_done_1]), ), @@ -1417,6 +1563,9 @@ def scheduler_deps(*resources: MemoryResource) -> list[MemoryResource]: ) if not cfg.stats_via_smem: resource_dependency_graph[tmem_stats_done_1] = [tmem_vec1] + if cfg.pv_half_overlap: + resource_dependency_graph[tmem_p_prefix_ready_0] = scheduler_deps(tmem_sp0) + resource_dependency_graph[tmem_p_prefix_ready_1] = scheduler_deps(tmem_sp1) if work_queue is not None: resource_dependency_graph[work_queue] = [work_queue] if is_clc_dynamic else [] if smem_page_offsets_kv is not None: @@ -1436,6 +1585,8 @@ def add_smem_resource(resource: MemoryResource | None) -> None: add_smem_resource(smem_q) add_smem_resource(smem_kv) + if cfg.split_kv_pipelines: + add_smem_resource(smem_v) if smem_page_offsets_kv is not None: add_smem_resource(smem_page_offsets_kv) if smem_page_offsets_v is not None: @@ -1447,6 +1598,10 @@ def add_smem_resource(resource: MemoryResource | None) -> None: add_smem_resource(tmem_sp0) if tmem_p0 is not None: add_smem_resource(tmem_p0) + if smem_p0 is not None: + add_smem_resource(smem_p0) + if smem_p1 is not None: + add_smem_resource(smem_p1) add_smem_resource(tmem_vec0) add_smem_resource(tmem_o) if not cfg.stats_via_smem: @@ -1457,6 +1612,8 @@ def add_smem_resource(resource: MemoryResource | None) -> None: add_smem_resource(tmem_vec1) if s0s1_seq is not None: add_smem_resource(s0s1_seq) + add_smem_resource(tmem_p_prefix_ready_0) + add_smem_resource(tmem_p_prefix_ready_1) if tmem_stats_done_1 is not None and not cfg.stats_via_smem: add_smem_resource(tmem_stats_done_1) if work_queue is not None and work_queue.pipeline_config is not None: @@ -1493,6 +1650,9 @@ def add_smem_resource(resource: MemoryResource | None) -> None: dealloc_mbar_alloc = smem_allocator.add( SmemAllocation("tmem_dealloc_mbar", dtype=cutlass.Int64, alignment=8) ) + peer_dealloc_mbar_alloc = smem_allocator.add( + SmemAllocation("tmem_peer_dealloc_mbar", dtype=cutlass.Int64, alignment=8) + ) clc_response_alloc: SmemAllocation | None = None if is_clc_dynamic and clc_response_ptr is None: # Keep the CLC response inside the unified TS allocation. The kernel @@ -1509,11 +1669,12 @@ def add_smem_resource(resource: MemoryResource | None) -> None: ) ) smem_allocator.compute_layout() + kv_stages = cfg.kv_stage_k + (cfg.kv_stage_v if cfg.split_kv_pipelines else 0) expected_barrier_bytes = ( sum( _context_pipeline_stage_counts( cfg, - kv_stages=cfg.kv_stage, + kv_stages=kv_stages, is_clc_dynamic=is_clc_dynamic, ).values() ) @@ -1586,6 +1747,7 @@ def add_smem_resource(resource: MemoryResource | None) -> None: tmem_resources, tmem_ptr_alloc, dealloc_mbar_alloc, + peer_dealloc_mbar_alloc, work_queue, clc_response_alloc, ) @@ -1612,7 +1774,10 @@ def _context_pipeline_stage_counts( kv_stages: int, is_clc_dynamic: bool, ) -> dict[str, int]: - """Return every physical pipeline's mbarrier stage count.""" + """Return every physical pipeline's mbarrier stage count. + + ``kv_stages`` totals the K/V pipelines: one ring, or both when split. + """ counts = { "smem_q": cfg.q_stage, "smem_kv": kv_stages, @@ -1627,27 +1792,26 @@ def _context_pipeline_stage_counts( "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, "tmem_stats_done": 0 if cfg.stats_via_smem else cfg.num_qkv_instances, "work_queue": 1 if is_clc_dynamic else 0, } return {name: stages for name, stages in counts.items() if stages} -def _infer_single_instance_kv_stages( +def _kv_ring_smem_budget_bytes( cfg: FmhaConfig, *, is_clc_dynamic: bool, page_table_window_entries: int | None = None, - require_cadence: bool = True, ) -> int: - """Return the deepest K/V ring that fits the exact TS SMEM footprint. + """Return the SMEM left for the K/V ring(s) after every fixed allocation. - All terms come from resource topology or public CUTLASS hardware metadata: - Q, O, correction statistics, page-ID rings, fixed control records, and one - 16-byte pipeline barrier per physical stage. The task manager remains the - authoritative check and uses the same stage-count policy below. + Fixed terms are Q, O, correction statistics, page-ID rings, control + records, and one 16-byte barrier per stage. The task manager still checks. """ - q_row_bytes = (cfg.q_dtype.width * cfg.qk_mma_tiler[2] + 7) // 8 + q_row_bytes = (cfg.q_dtype.width * cfg.smem_q_head_dim + 7) // 8 q_row_bytes = ( (q_row_bytes + _Q_ROW_SMEM_ALIGNMENT_BYTES - 1) // _Q_ROW_SMEM_ALIGNMENT_BYTES @@ -1656,13 +1820,15 @@ def _infer_single_instance_kv_stages( q_tile_bytes = q_row_bytes * cfg.qk_mma_tiler[0] o_head_dim = ( - cfg.head_dim_per_stage_kv if cfg.stage_o_by_head_dim else cfg.epi_tile[1] + cfg.head_dim_per_stage_o if cfg.stage_o_by_head_dim else cfg.epi_tile[1] ) o_stage_bytes = (cfg.epi_tile[0] * o_head_dim * cfg.o_dtype.width + 7) // 8 stats_bytes = 0 if cfg.stats_via_smem: - stats_rows = len(cfg.softmax0_warp_ids) * cute.arch.WARP_SIZE + stats_rows = ( + cfg.num_qkv_instances * len(cfg.softmax0_warp_ids) * cute.arch.WARP_SIZE + ) stats_values_per_row = 2 stats_bytes = ( cfg.softmax_corr_stage @@ -1682,9 +1848,18 @@ def _infer_single_instance_kv_stages( page_offset_stages * page_table_window_entries * cutlass.Int32.width // 8 ) + # The TS allocator bump-allocates each data block at its own alignment + # (largest first), so a buffer-aligned block costs its size rounded up to + # cfg.buffer_align_bytes. The 4-32 byte records follow, then barriers. + align = cfg.buffer_align_bytes + smem_p_bytes = 0 + if cfg.p_in_smem: + smem_p_bytes = ( + (cfg.smem_p_bytes + align - 1) // align * align * cfg.num_qkv_instances + ) control_bytes = (2 * cutlass.Int32.width + cutlass.Int64.width) // 8 if is_clc_dynamic: - control_bytes += cutlass.Int128.width // 8 + control_bytes = (control_bytes + 15) // 16 * 16 + cutlass.Int128.width // 8 fixed_barrier_stages = sum( _context_pipeline_stage_counts( cfg, @@ -1693,45 +1868,123 @@ def _infer_single_instance_kv_stages( ).values() ) fixed_smem_bytes = ( - q_tile_bytes * cfg.q_stage - + o_stage_bytes - + stats_bytes - + page_offsets_bytes - + control_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 + + (control_bytes + 7) // 8 * 8 + fixed_barrier_stages * _PIPELINE_BARRIER_BYTES_PER_STAGE ) + return utils.get_smem_capacity_in_bytes("sm_100") - fixed_smem_bytes + + +def _infer_single_instance_kv_stages( + cfg: FmhaConfig, + *, + is_clc_dynamic: bool, + page_table_window_entries: int | None = None, + require_cadence: bool = True, +) -> int: + """Return the deepest K/V ring that fits the exact TS SMEM footprint. + + All terms come from resource topology or public CUTLASS hardware metadata: + Q, O, correction statistics, page-ID rings, fixed control records, and one + 16-byte pipeline barrier per physical stage. The task manager remains the + authoritative check and uses the same stage-count policy below. + """ + kv_budget_bytes = _kv_ring_smem_budget_bytes( + cfg, + is_clc_dynamic=is_clc_dynamic, + page_table_window_entries=page_table_window_entries, + ) kv_dtype_width = max(cfg.k_dtype.width, cfg.v_dtype.width) + kv_head_dim = ( + cfg.head_dim_per_stage_kv if cfg.stage_kv_by_head_dim else cfg.qk_mma_tiler[2] + ) kv_stage_bytes = ( - cfg.qk_mma_tiler[1] * cfg.head_dim_per_stage_kv * kv_dtype_width // 8 + 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 - kv_budget_bytes = utils.get_smem_capacity_in_bytes("sm_100") - fixed_smem_bytes 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: raise ValueError( - "single-instance context staging requires at least " + "context staging requires at least " f"{cadence_stages} K/V stages, but the shared-memory budget fits " f"only {memory_fit_stages}" ) return memory_fit_stages -def _configure_pipeline_stages(cfg: FmhaConfig, *, is_clc_dynamic: bool) -> None: +def _configure_kv_ring_depths(cfg: FmhaConfig, *, is_clc_dynamic: bool) -> None: + """Select one shared K/V ring or, for mixed dtypes, size K and V independently.""" + cfg.split_kv_pipelines = cfg.k_dtype.width != cfg.v_dtype.width + cfg.kv_stage_k = cfg.kv_stage + cfg.kv_stage_v = cfg.kv_stage + if not cfg.split_kv_pipelines: + # Same-width K and V share one ring: deepen it to what the SMEM budget + # allows after every other buffer, as FA4 does. + n_stages = max( + cfg.kv_stage, + _infer_single_instance_kv_stages( + cfg, is_clc_dynamic=is_clc_dynamic, require_cadence=False + ), + ) + cfg.kv_stage = n_stages + cfg.kv_stage_k = n_stages + cfg.kv_stage_v = n_stages + return + # Mixed dtypes require separate K and V rings. Size each ring from one + # cadence up to cfg.kv_stage against the shared SMEM budget. + budget_bytes = _kv_ring_smem_budget_bytes(cfg, is_clc_dynamic=is_clc_dynamic) + kv_head_dim = ( + cfg.head_dim_per_stage_kv if cfg.stage_kv_by_head_dim else cfg.qk_mma_tiler[2] + ) + kv_rows_per_cta = cfg.qk_mma_tiler[1] // cfg.cta_group_size + k_stage_footprint = ( + kv_rows_per_cta * kv_head_dim * cfg.k_dtype.width // 8 + + _PIPELINE_BARRIER_BYTES_PER_STAGE + ) + v_stage_footprint = ( + kv_rows_per_cta * kv_head_dim * cfg.v_dtype.width // 8 + + _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 + n_stages = budget_bytes // (k_stage_footprint + v_stage_footprint) + if n_stages < cadence: + raise ValueError( + f"split K/V staging requires at least {cadence} stages per ring " + f"({cadence * (k_stage_footprint + v_stage_footprint)} bytes), but the " + f"shared-memory budget fits only {budget_bytes} bytes" + ) + cfg.kv_stage_k = n_stages + cfg.kv_stage_v = n_stages + + +def _configure_pipeline_stages( + cfg: FmhaConfig, *, is_clc_dynamic: bool, is_persistent: bool +) -> None: """Set topology- and capacity-derived context pipeline stage counts.""" cfg.q_stage = cfg.num_qkv_instances - cfg.kv_stage = 3 + # Stages using Two-CTA method are half-tiles, where the peer CTA holds the + # other half. Double them to keep the same SMEM and tiles in flight. + cfg.kv_stage = 6 if cfg.two_cta_umma else 3 cfg.has_tmem_p_pipeline = _should_use_tmem_p_pipeline(cfg) cfg.stage_scoped_tmem_stats = cfg.has_tmem_p_pipeline cfg.mma_softmax_stage = 2 if cfg.has_tmem_p_pipeline else 1 cfg.softmax_corr_stage = 2 if cfg.stage_scoped_tmem_stats else 1 - # SMEM-backed D256 removes the independent StatsDone credit and always - # writes the same physical O0 accumulator. Its MMA->Correction handoff must - # therefore be single-stage so PV(i+1) cannot overwrite O0 before - # Correction consumes PV(i). TMEM-stats schedules retain their established - # two-stage O + StatsDone ordering. - cfg.mma_corr_stage = 1 if cfg.single_qkv_instance and cfg.stats_via_smem else 2 - if cfg.single_qkv_instance: + # The early-sum policy decides the P-prefix barriers, so settle it before + # the K/V ring is sized against the remaining SMEM. + _configure_early_tile_sum_policy(cfg, is_persistent=is_persistent) + # A single Q instance always reuses one physical O accumulator. Allow only + # one outstanding PV result so correction finishes before the next PV. + # StatsDone protects statistics, and is released before O correction; + # it cannot protect O from a second in-flight MMA. + cfg.mma_corr_stage = 1 if cfg.single_qkv_instance else 2 + if cfg.single_qkv_instance or cfg.logical_head_dim_qk == 192: natural_page_window_entries = cute.arch.WARP_SIZE cfg.page_table_window_entries = natural_page_window_entries candidate_page_window_entries = cfg.page_table_window_candidate_entries @@ -1749,6 +2002,7 @@ def _configure_pipeline_stages(cfg: FmhaConfig, *, is_clc_dynamic: bool) -> None cfg, is_clc_dynamic=is_clc_dynamic, ) + _configure_kv_ring_depths(cfg, is_clc_dynamic=is_clc_dynamic) # Dense work traverses the full K domain for every Q tile, so its persistent @@ -1760,6 +2014,9 @@ def _configure_pipeline_stages(cfg: FmhaConfig, *, is_clc_dynamic: bool) -> None _EARLY_TILE_SUM_DENSE_REGISTER_BUDGET = (176, 80, 80) _EARLY_TILE_SUM_CAUSAL_REGISTER_BUDGET = (184, 88, 56) +# Lazy-correction margin for BF16 V. cuDNN and FA4 use 8. +_CORR_SKIP_THRESHOLD_LOG2 = 8.0 + def _configure_early_tile_sum_policy( cfg: FmhaConfig, @@ -1796,6 +2053,13 @@ def _configure_early_tile_sum_policy( else _EARLY_TILE_SUM_DENSE_REGISTER_BUDGET ) + # MLA's additional live K descriptor benefits from the reference's + # softmax/correction split while staying within the same register budget. + if cfg.logical_head_dim_qk == 192 and cfg.is_causal and cfg.q_dtype.width == 8: + cfg.num_regs_softmax = 184 + cfg.num_regs_correction = 96 + cfg.num_regs_other = 48 + def _configure_smem_shapes(cfg: FmhaConfig) -> None: """Derive per-stage SMEM element counts from the configured tile shapes.""" @@ -1803,15 +2067,16 @@ def _configure_smem_shapes(cfg: FmhaConfig) -> None: cfg.head_dim_per_stage_kv if cfg.stage_kv_by_head_dim else cfg.qk_mma_tiler[2] ) o_head_dim = ( - cfg.head_dim_per_stage_kv if cfg.stage_o_by_head_dim else cfg.epi_tile[1] + cfg.head_dim_per_stage_o if cfg.stage_o_by_head_dim else cfg.epi_tile[1] ) cfg.sQ_shape = ( cfg.q_stage, - cfg.qk_mma_tiler[0] * cfg.qk_mma_tiler[2], + cfg.qk_mma_tiler[0] * cfg.smem_q_head_dim, ) + # In two-CTA, each CTA stages half the K rows and half the V columns of a tile. cfg.sK_shape = ( cfg.kv_stage, - cfg.qk_mma_tiler[1] * kv_head_dim, + cfg.qk_mma_tiler[1] * kv_head_dim // cfg.cta_group_size, ) cfg.sO_stage_elements = cfg.epi_tile[0] * o_head_dim @@ -1862,24 +2127,60 @@ def _configure_single_instance_warp_layout(cfg: FmhaConfig) -> None: cfg.num_regs_other = 112 +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 + 64-wide halves. The 192/128 path keeps 128-wide K/V staging instead, and an + 8-bit output has no 64-wide TMA store granule. Under two-CTA UMMA the + P-ready barrier is cluster-level like the P-prefix one.""" + return ( + not cfg.single_qkv_instance + and cfg.logical_head_dim_qk == 128 + and cfg.v_dtype.width == 8 + and cfg.o_dtype.width == 16 + and cfg.qk_mma_tiler[1] * cfg.v_dtype.width // 8 == 128 + and cfg.epi_tile[1] % 64 == 0 + and not has_variable_window + and not cfg.causal_single_kv_tile + ) + + def _configure_head_dim_staging(cfg: FmhaConfig) -> None: - """Split d>128 single-instance K/V/O staging into 128-wide slices.""" + """Stage K/V in 128-wide slices and size O staging to the paired footprint.""" cfg.head_dim_per_stage_kv = 0 cfg.num_head_dim_stages_k = 1 cfg.num_head_dim_stages_v = 1 cfg.num_o_head_dim_stages = 1 cfg.stage_kv_by_head_dim = False cfg.stage_o_by_head_dim = False - if cfg.num_qkv_instances != 1: + if cfg.p_in_smem: + # Stage O in 64-wide halves to free 32 KB of SMEM for the two P tiles. + cfg.head_dim_per_stage_o = 64 + cfg.num_o_head_dim_stages = cfg.epi_tile[1] // cfg.head_dim_per_stage_o + cfg.stage_o_by_head_dim = True + return + if cfg.num_qkv_instances != 1 and cfg.logical_head_dim_qk != 192: return cfg.head_dim_per_stage_kv = 128 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 - cfg.num_o_head_dim_stages = cfg.epi_tile[1] // cfg.head_dim_per_stage_kv + # K/V transfers remain 128-wide. 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 + 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.stage_kv_by_head_dim = True - cfg.stage_o_by_head_dim = True - # Stage K, V, and O as 128-wide head-dimension slices so the d>128 K/V - # pipeline can run deeper without exceeding Blackwell's SMEM budget. + 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 + # retain its K/V ring without exceeding Blackwell's SMEM budget. def _configure_head_paired_tilers( @@ -1903,6 +2204,7 @@ def _configure_head_paired_tma_copy_metadata( *, q_dtype: type, k_dtype: type, + v_dtype: type, o_dtype: type, ) -> None: """Derive TMA copy granularities for head-paired Q/K/V/O tensors.""" @@ -1933,6 +2235,17 @@ 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 + 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_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 cfg.tma_copy_o_iters = 1 if output_inner_dim_size % 128 == 0: @@ -1943,7 +2256,7 @@ def _configure_head_paired_tma_copy_metadata( ) cfg.tma_copy_o_granu_inner = cfg.epi_tile[1] // cfg.tma_copy_o_iters o_head_dim = ( - cfg.head_dim_per_stage_kv if cfg.stage_o_by_head_dim else cfg.epi_tile[1] + cfg.head_dim_per_stage_o if cfg.stage_o_by_head_dim else cfg.epi_tile[1] ) cfg.tma_copy_o_stage_iters = o_head_dim // cfg.tma_copy_o_granu_inner cfg.tma_copy_o_elements = cfg.epi_tile[0] * o_head_dim @@ -2210,11 +2523,23 @@ def _select_fmha_domain_policy( # peers traverse the full static K/V domain. domain_n_kwargs = {"domain": num_kv_tiles} domain_n_minus_1_kwargs = {"domain": num_kv_tiles - 1} + # Query-paired softmax handles the partial last K/V tile after its loop: + # N-1 unmasked iterations, then a masked tail. The single-QKV schedules + # mask the last tile inside the loop instead, so they keep the full N domain. + softmax_uses_dense_k_tail = ( + not cfg.single_qkv_instance + and not cfg.is_causal + and not cfg.has_varlen + and cfg.fixed_dense_k_tail > 0 + ) + softmax_domain_kwargs = ( + domain_n_minus_1_kwargs if softmax_uses_dense_k_tail else domain_n_kwargs + ) return FmhaDomainPolicy( domain_n_kwargs=domain_n_kwargs, domain_n_minus_1_kwargs=domain_n_minus_1_kwargs, - softmax0_domain_kwargs=domain_n_kwargs, - softmax1_domain_kwargs=domain_n_kwargs, + softmax0_domain_kwargs=softmax_domain_kwargs, + softmax1_domain_kwargs=softmax_domain_kwargs, ) @@ -2252,6 +2577,7 @@ def build_fmha_task_manager( list[MemoryResource], SmemAllocation, SmemAllocation, + SmemAllocation, WorkQueue | None, SmemAllocation | None, ]: @@ -2367,14 +2693,19 @@ class FmhaTs: Accumulator dtype for QK GEMM (default: Float32). pv_acc_dtype : type, optional Accumulator dtype for PV GEMM (default: Float32). - in_dtype : type, optional - Input tensor dtype for Q, K, V (default: Float16). + in_qk_dtype : type, optional + Input tensor dtype for Q and K (default: Float16). + in_pv_dtype : type, optional + Input tensor dtype for V (default: Float16). out_dtype : type, optional Output tensor dtype for O (default: Float16). mma_tiler_mn : Tuple[int, int], optional MMA tile shape (M, N) (default: (128, 128)). d : int, optional - Head dimension (default: 128). + Q/K head dimension (default: 128). + d_v : int, optional + Value/output head dimension, defaulting to d. Asymmetric support is + limited to separate contiguous QK=192, V=128. is_persistent : bool, optional Use persistent scheduling (default: True). is_causal : bool, optional @@ -2392,6 +2723,15 @@ class FmhaTs: grouped-query attention with an even repeat count. enable_skip_correction : bool, optional Enable skip-correction for softmax rescaling (default: True). + uses_ldtm_stat : bool, optional + Use ``tcgen05.ld.red.max`` (LDTM.STAT) to fuse the per-chunk row_max + 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). + 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 + bf16 or E4M3 QK, launched non-persistently. use_paged_kv : bool, optional Read K/V from a physical page pool through a fixed block table. num_tokens_per_page : int, optional @@ -2409,7 +2749,8 @@ def __init__( self, qk_acc_dtype: type | None = None, pv_acc_dtype: type | None = None, - in_dtype: type | None = None, + in_qk_dtype: type | None = None, + in_pv_dtype: type | None = None, out_dtype: type | None = None, mma_tiler_mn: Tuple[int, int] = (128, 128), d: int = 128, @@ -2422,11 +2763,15 @@ def __init__( has_variable_window: bool = False, h_r: int = 1, enable_skip_correction: bool = True, + uses_ldtm_stat: bool = False, + fp8_psmem_early_token: bool = False, + two_cta_umma: bool = False, use_paged_kv: bool = False, num_tokens_per_page: int = 32, max_kv_len: int = 1, causal_single_kv_tile: bool = False, exhaustive_deadlock_race_check: bool = True, + d_v: int | None = None, ) -> None: """Initialize mode-specific tiling, dtype, and schedule configuration.""" head_paired = resolve_head_paired_mode( @@ -2471,14 +2816,47 @@ def __init__( self.is_clc_dynamic = is_clc_dynamic self.exhaustive_deadlock_race_check = exhaustive_deadlock_race_check - q_dtype = in_dtype or cutlass.Float16 - k_dtype = in_dtype or cutlass.Float16 - v_dtype = in_dtype or cutlass.Float16 + q_dtype = in_qk_dtype or cutlass.Float16 + k_dtype = in_qk_dtype or cutlass.Float16 + v_dtype = in_pv_dtype or cutlass.Float16 + o_dtype = out_dtype or cutlass.Float16 + # 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): + raise ValueError("asymmetric context requires QK=192 and V=128") + if d == 192 and (head_paired or use_paged_kv): + raise NotImplementedError( + "192/128 requires contiguous query-paired context" + ) + # Round the MMA extent to the next power of two, with a 64-wide minimum. + # Global Q/K rows stay compact; TMA zero-fills partial shared tiles. + padded_d = 1 << (max(d, 64) - 1).bit_length() cfg = FmhaConfig() + cfg.logical_head_dim_qk = d + if two_cta_umma and ( + is_persistent + or is_causal + or has_variable_window + or head_paired + or use_paged_kv + or d != 128 + or d_v != 128 + or q_dtype.width not in (8, 16) + ): + raise ValueError( + "two-CTA UMMA requires the non-persistent dense contiguous " + "query-paired D128 context kernel with 16-bit or E4M3 QK" + ) + cfg.two_cta_umma = two_cta_umma + if two_cta_umma: + cfg.cluster_shape_mn = (2, 1) self.cfg = cfg - if d > 128: + # 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 + if d > 128 and not paired_mla: cfg.num_qkv_instances = 1 cfg.use_paged_kv = use_paged_kv single_instance_persistent = ( @@ -2504,6 +2882,10 @@ def __init__( cfg.num_regs_correction = 88 cfg.num_regs_other = 56 cfg.enable_skip_correction = enable_skip_correction + if enable_skip_correction: + 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.qk_acc_dtype = qk_acc_dtype or cutlass.Float32 cfg.pv_acc_dtype = pv_acc_dtype or cutlass.Float32 @@ -2521,7 +2903,9 @@ def __init__( if head_paired: _configure_head_paired_tilers(cfg, mma_tiler_mn=mma_tiler_mn, d=d) _configure_head_dim_staging(cfg) - _configure_pipeline_stages(cfg, is_clc_dynamic=is_clc_dynamic) + _configure_pipeline_stages( + cfg, is_clc_dynamic=is_clc_dynamic, is_persistent=is_persistent + ) if cfg.single_qkv_instance: _configure_single_instance_tmem_layout(cfg) _configure_single_instance_warp_layout(cfg) @@ -2531,6 +2915,7 @@ def __init__( cfg, q_dtype=q_dtype, k_dtype=k_dtype, + v_dtype=v_dtype, o_dtype=o_dtype, ) _configure_common_launch_flags( @@ -2542,16 +2927,18 @@ def __init__( window_size_left=window_size_left, has_variable_window=has_variable_window, ) - _configure_early_tile_sum_policy(cfg, is_persistent=is_persistent) return # MMA tiler: (M, N, K) = (128, 128, 128) - mma_tiler = (*mma_tiler_mn, d) + mma_tiler = (*mma_tiler_mn, padded_d) cfg.qk_mma_tiler = mma_tiler - cfg.pv_mma_tiler = (mma_tiler[0], mma_tiler[2], mma_tiler[1]) + 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) _configure_head_dim_staging(cfg) - _configure_pipeline_stages(cfg, is_clc_dynamic=is_clc_dynamic) + _configure_pipeline_stages( + cfg, is_clc_dynamic=is_clc_dynamic, is_persistent=is_persistent + ) if cfg.single_qkv_instance: _configure_single_instance_tmem_layout(cfg) _configure_single_instance_warp_layout(cfg) @@ -2561,7 +2948,7 @@ def __init__( _validate_tmem_columns(cfg) # TMA copy granularity for Q - qkv_tma_bits = cfg.qk_mma_tiler[2] * q_dtype.width + qkv_tma_bits = cfg.smem_q_head_dim * q_dtype.width if qkv_tma_bits % (128 * 8) != 0: raise ValueError( "FMHA TS requires a 128-byte aligned Q/K/V inner dimension, " @@ -2569,7 +2956,7 @@ def __init__( ) cfg.tma_copy_qkv_iters = qkv_tma_bits // (128 * 8) cfg.q_tile_m = cfg.qk_mma_tiler[0] - cfg.tma_copy_q_granu_inner = cfg.qk_mma_tiler[2] // cfg.tma_copy_qkv_iters + cfg.tma_copy_q_granu_inner = cfg.smem_q_head_dim // cfg.tma_copy_qkv_iters cfg.tma_copy_q_elements = cfg.sQ_shape[1] cfg.tma_copy_q_granu_elems = cfg.tma_copy_q_elements // cfg.tma_copy_qkv_iters cfg.tma_copy_q_bytes = cfg.tma_copy_q_elements * q_dtype.width // 8 @@ -2581,7 +2968,7 @@ def __init__( if cfg.stage_kv_by_head_dim else cfg.qk_mma_tiler[2] ) - cfg.tma_copy_kv_granu_inner = cfg.qk_mma_tiler[2] // cfg.tma_copy_qkv_iters + cfg.tma_copy_kv_granu_inner = cfg.tma_copy_q_granu_inner cfg.tma_copy_kv_elements = cfg.sK_shape[1] cfg.tma_copy_kv_stage_iters = kv_head_dim // cfg.tma_copy_kv_granu_inner cfg.tma_copy_kv_granu_elems = ( @@ -2589,11 +2976,29 @@ 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 + 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_stage_iters = kv_head_dim // cfg.tma_copy_v_granu_inner + if cfg.two_cta_umma: + cfg.tma_copy_v_iters = 1 + cfg.tma_copy_v_granu_inner = cfg.pv_n_per_cta + cfg.tma_copy_v_stage_iters = 1 + 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 + # 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 o_head_dim = ( - cfg.head_dim_per_stage_kv if cfg.stage_o_by_head_dim else cfg.epi_tile[1] + cfg.head_dim_per_stage_o if cfg.stage_o_by_head_dim else cfg.epi_tile[1] ) cfg.tma_copy_o_stage_iters = o_head_dim // cfg.tma_copy_o_granu_inner cfg.tma_copy_o_elements = cfg.epi_tile[0] * o_head_dim @@ -2610,7 +3015,6 @@ def __init__( window_size_left=window_size_left, has_variable_window=has_variable_window, ) - _configure_early_tile_sum_policy(cfg, is_persistent=is_persistent) # --------------------------------------------------------------------------- # Host entry point @@ -2703,13 +3107,22 @@ def __call__( elif cutlass.const_expr(output_inner_dim_size == 32): tma_o_swizzle = cuda.TensorMapSwizzle.s32b + v_inner_dim_size = cfg.tma_copy_v_granu_inner * cfg.v_dtype.width // 8 + tma_v_swizzle = cuda.TensorMapSwizzle.none + if cutlass.const_expr(v_inner_dim_size % 128 == 0): + tma_v_swizzle = cuda.TensorMapSwizzle.s128b + elif cutlass.const_expr(v_inner_dim_size == 64): + tma_v_swizzle = cuda.TensorMapSwizzle.s64b + elif cutlass.const_expr(v_inner_dim_size == 32): + tma_v_swizzle = cuda.TensorMapSwizzle.s32b + q_box_dims = (1, cfg.qk_mma_tiler[0], 1, cfg.tma_copy_q_granu_inner) - kv_box_dims = (1, cfg.qk_mma_tiler[1], 1, cfg.tma_copy_kv_granu_inner) + kv_box_dims = (1, cfg.kv_tile_rows_per_cta, 1, cfg.tma_copy_kv_granu_inner) v_box_dims = ( 1, cfg.pv_mma_tiler[2], 1, - cfg.pv_mma_tiler[1] // cfg.tma_copy_qkv_iters, + cfg.tma_copy_v_granu_inner, ) o_box_dims = (1, cfg.epi_tile[0], 1, cfg.tma_copy_o_granu_inner) stride_order = (3, 2, 1, 0) @@ -2729,7 +3142,7 @@ def __call__( v_box_dims = ( cfg.pv_mma_tiler[2], 1, - cfg.pv_mma_tiler[1] // cfg.tma_copy_qkv_iters, + cfg.pv_mma_tiler[1] // cfg.tma_copy_v_iters, ) kv_stride_order = stride_order if cutlass.const_expr(cum_seqlen_q is not None): @@ -2746,6 +3159,8 @@ def __call__( stride_order=stride_order, swizzle=tma_qkv_swizzle, l2_promotion=cuda.TensorMapL2Promotion.none, + # Zero-fill partial Q rows and head-dimension fragments. + oob_fill=cuda.TensorMapFloatOOBFill.none, ) else: tma_q_desc = cuda.create_tensor_map_tiled_from_view( @@ -2760,29 +3175,39 @@ def __call__( # Paged-KV path: K/V are compact pool tensors with shape # (total_pages, h_kv, num_tokens_per_page, d). The TMA box covers # one page × one d-fragment; the loader stitches pages and d-halves - # together via per-fragment coordinates. Keep both descriptors on - # the native rank-4 view. Synthetic ragged V maps with numeric-zero + # together via per-fragment coords. Keep both descriptors on the + # native rank-4 view. Synthetic ragged V maps with numeric-zero # OOB fill have triggered mixed-specialization SM100 aborts; their # special-NaN alternative is also invalid because tcgen05.mma # propagates the NaN instead of treating it as a zero operand. - paged_kv_box_dims = ( + # When K and V share a dtype, one box shape and swizzle covers + # both. Mixed dtypes (e.g. QK-BF16/PV-FP8) require separate + # descriptors with V's inner granularity and swizzle derived from + # v_dtype. + paged_k_box_dims = ( 1, 1, cfg.num_tokens_per_page, cfg.tma_copy_kv_granu_inner, ) + paged_v_box_dims = ( + 1, + 1, + cfg.num_tokens_per_page, + cfg.tma_copy_v_granu_inner, + ) tma_k_desc = cuda.create_tensor_map_tiled_from_view( k_cute, - box_dims=paged_kv_box_dims, + box_dims=paged_k_box_dims, stride_order=kv_stride_order, swizzle=tma_qkv_swizzle, l2_promotion=tma_kv_l2_promotion, ) tma_v_desc = cuda.create_tensor_map_tiled_from_view( v_cute, - box_dims=paged_kv_box_dims, + box_dims=paged_v_box_dims, stride_order=kv_stride_order, - swizzle=tma_qkv_swizzle, + swizzle=tma_v_swizzle, l2_promotion=tma_kv_l2_promotion, ) else: @@ -2797,7 +3222,7 @@ def __call__( v_cute, box_dims=v_box_dims, stride_order=kv_stride_order, - swizzle=tma_qkv_swizzle, + swizzle=tma_v_swizzle, l2_promotion=tma_kv_l2_promotion, ) @@ -2857,7 +3282,11 @@ def __call__( if cutlass.const_expr(cfg.uses_head_batch_seq_tile_order): problem_shape = (num_head_tiles, b, num_seq_tiles) else: - problem_shape = (num_seq_tiles, num_head_tiles, b) + grid_seq_tiles = num_seq_tiles + if cutlass.const_expr(cfg.two_cta_umma): + # Clusters pair Q tiles; an odd count gets one padding tile. + grid_seq_tiles = cute.ceil_div(num_seq_tiles, 2) * 2 + problem_shape = (grid_seq_tiles, num_head_tiles, b) if cutlass.const_expr(self.is_clc_dynamic): tile_sched_params = utils.ClcDynamicPersistentTileSchedulerParams( @@ -3004,6 +3433,7 @@ def kernel( tmem_resources, tmem_ptr_alloc, dealloc_mbar_alloc, + peer_dealloc_mbar_alloc, work_queue, clc_response_alloc, ) = build_fmha_task_manager( @@ -3064,8 +3494,10 @@ def kernel( smem_allocator = task_manager.smem_allocator tmem_ptr_i32 = smem_allocator.get(tmem_ptr_alloc) tmem_dealloc_mbar = smem_allocator.get(dealloc_mbar_alloc) + tmem_peer_dealloc_mbar = smem_allocator.get(peer_dealloc_mbar_alloc) - # 5. Initialize tmem dealloc barrier + # 5. Initialize tmem dealloc barriers. The local barrier counts this + # CTA's consumers. The peer MMA warp arrives on the second barrier. num_tmem_consumer_threads = cute.arch.WARP_SIZE * ( len(cfg.softmax0_warp_ids) + len(cfg.softmax1_warp_ids) @@ -3074,6 +3506,8 @@ def kernel( if warp_idx == cfg.empty_warp_id: if prims.elect_sync(): prims.mbarrier_init(tmem_dealloc_mbar, num_tmem_consumer_threads) + if cutlass.const_expr(cfg.two_cta_umma): + prims.mbarrier_init(tmem_peer_dealloc_mbar, 1) # Fence barrier inits from setup_resources_and_tasks() — pipeline # barriers need fencing before first use. @@ -3084,10 +3518,15 @@ def kernel( prims.barrier_cluster_wait() # 6. TMEM allocation (MMA warp) + tmem_cta_group = ( + prims.CTAGroup.CTA_2 + if cutlass.const_expr(cfg.two_cta_umma) + else prims.CTAGroup.CTA_1 + ) if warp_idx == cfg.mma_warp_id: tmem_alloc_cols = Int32(cfg.tmem_alloc_cols) - prims.tcgen05_alloc(tmem_ptr_i32, tmem_alloc_cols) - prims.tcgen05_relinquish_alloc_permit() + prims.tcgen05_alloc(tmem_ptr_i32, tmem_alloc_cols, group=tmem_cta_group) + prims.tcgen05_relinquish_alloc_permit(group=tmem_cta_group) # All-warp barrier to ensure TMEM allocation is visible. # Previously only correction+MMA warps synced; expanded to all @@ -3129,6 +3568,17 @@ def kernel( if warp_idx == cfg.mma_warp_id: while not prims.mbarrier_try_wait_parity(tmem_dealloc_mbar, 0): pass + if cutlass.const_expr(cfg.two_cta_umma): + # Wait for the peer after the local consumers are done so the + # paired TMEM and the leader-side barriers stay live until both + # CTAs have finished reading O and arriving on them. + cta_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + if prims.elect_sync(): + prims.mbarrier_arrive( + prims.mapa(tmem_peer_dealloc_mbar, cta_rank ^ 1) + ) + while not prims.mbarrier_try_wait_parity(tmem_peer_dealloc_mbar, 0): + pass tmem_alloc_cols = Int32(cfg.tmem_alloc_cols) tmem_ptr = prims.make_tmem_ptr(tmem_ptr_i32.load(), cutlass.Int8) - prims.tcgen05_dealloc(tmem_ptr, tmem_alloc_cols) + prims.tcgen05_dealloc(tmem_ptr, tmem_alloc_cols, group=tmem_cta_group) 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 6df566660606..c4d1b5512f92 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 @@ -27,8 +27,11 @@ SMEM resources (TMA pipelines) ------------------------------ - SmemQResource : SMEM Q buffer, TmaUmma pipeline. Load -> MMA. -- SmemKVResource : SMEM KV buffer, TmaUmma pipeline. Load -> MMA. The D256 - depth is derived from the public SM100 SMEM capacity. +- SmemKVResource : SMEM K/V buffer, TmaUmma pipeline. Load -> MMA. The D256 + depth is derived from the public SM100 SMEM capacity. One + instance holds both K and V when they share a dtype; a + QK/PV dtype mismatch instantiates one K-only and one V-only + instance, each sized from its own dtype. - SmemOResource : One SMEM O buffer per Q/O instance, AsyncAsync pipeline. Correction -> Epilogue. @@ -64,7 +67,7 @@ import cutlass import cutlass.cute as cute -from cutlass import Float32, Int32 +from cutlass import Float32, Int16, Int32 from ..tensor_map import transform_ragged_coords from cutlass.experimental.task_scheduling.enums import WorkAttr @@ -88,9 +91,11 @@ bottom_right_window_left_bound, bottom_right_window_tile_start, freeze_smem_descriptor, + load_tmem_32x32b_max, variable_window_cta_min_start, ) from cutlass.experimental import primitives as prims +from cutlass._mlir.dialects import arith as arith_dialect SmemDescOffsets: TypeAlias = tuple[int, int] TmemAddr: TypeAlias = int | Int32 @@ -105,6 +110,8 @@ # Stored at module level (not on self) to avoid adding a non-dynamic-expression # field to the dataclass, which breaks the framework's scf.if handling. _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] = {} @cute.jit @@ -161,24 +168,97 @@ def _pack_float4_to_fp8_e4m3( v2: Float32, v3: Float32, ) -> Int32: - """Pack four FP32 values with the public packed-conversion primitive.""" - lo = prims.cvt_packfloat_f32( - v1, - v0, - Int32(0), - prims.CVTPackFloat.E4M3X2, - rnd=prims.FPRoundingMode.RN, - sat=prims.SaturationModeKind.SATFINITE, - extract_hi=False, + """Pack four FP32 values into one e4m3x4 word. Two b16 conversions + are joined by mov.b32, which ptxas emits as F2FP pairs with no PRMT.""" + return cute.arch.inline_ptx( + """{ + .reg .b16 lo, hi; + cvt.rn.satfinite.e4m3x2.f32 lo, {$r1}, {$r0}; + cvt.rn.satfinite.e4m3x2.f32 hi, {$r3}, {$r2}; + mov.b32 {$w0}, {lo, hi}; + }""", + write_only_types=[Int32], + read_only_args=[v0, v1, v2, v3], + ) + + +@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 ) - return prims.cvt_packfloat_f32( - v3, - v2, - lo, - prims.CVTPackFloat.E4M3X2, - rnd=prims.FPRoundingMode.RN, - sat=prims.SaturationModeKind.SATFINITE, - extract_hi=True, + # 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)), + ) + + +@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)), ) @@ -221,6 +301,8 @@ class FmhaConfig: pv_acc_dtype: type | None = None # Tile shapes + logical_head_dim_qk: int = 128 + head_dim_per_stage_o: int = 128 qk_mma_tiler: tuple[int, int, int] = (128, 128, 64) pv_mma_tiler: tuple[int, int, int] = (128, 64, 128) epi_tile: tuple[int, int] = (128, 64) @@ -231,7 +313,19 @@ class FmhaConfig: # Pipeline stages q_stage: int = 2 kv_stage: int = 3 + # Hand the S0/S1 pacing token back before the exp2/P work on the fp8 P-in-SMEM path + # (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 + # 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. + split_kv_pipelines: bool = False + kv_stage_k: int = 3 + kv_stage_v: int = 3 mma_softmax_stage: int = 1 + # Stage fp8 P in SMEM so softmax releases the S stage right after loading it. + p_in_smem: bool = False # Use the two-stage loop-carried S/P schedule and an independent P-ready # handoff, allowing QK(i+1) and PV(i) to operate on opposite S/P stages. has_tmem_p_pipeline: bool = False @@ -250,6 +344,13 @@ class FmhaConfig: tma_copy_kv_elements: int = 0 tma_copy_kv_granu_elems: int = 0 tma_copy_kv_bytes: int = 0 + # V-specific TMA copy granularity. V may use a different dtype than K, + # e.g. QK-BF16/PV-FP8, but these fields are otherwise ignored. + tma_copy_v_iters: int = 1 + tma_copy_v_granu_inner: int = 128 + tma_copy_v_stage_iters: int = 0 + tma_copy_v_granu_elems: int = 0 + tma_copy_v_bytes: int = 0 tma_copy_o_iters: int = 1 tma_copy_o_granu_inner: int = 0 tma_copy_o_elements: int = 0 @@ -307,6 +408,11 @@ class FmhaConfig: num_seq_tiles: int | Int32 = 0 # Skip correction optimization: when True, skip rescale if old_max == new_max enable_skip_correction: bool = True + # 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 + # exp2 pairs per 16-pair softmax chunk computed on the FMA pipe. + exp2_fma_pairs: int = 0 # Variable sequence length mode stores Q/K/V/O as flattened # [sum_seqlen, head, dim] tensors and uses cum_seqlen_* for per-batch @@ -331,6 +437,11 @@ class FmhaConfig: seq_tile_n: int = 128 tmem_x_load_s: int = 32 + # Opt-in to `tcgen05.ld.red.max` (LDTM.STAT) for the non-masked + # `compute_row_max` path: fuses the per-chunk max into the TMEM load. + # Requires SM103+/SM110+ with tcgen05.ld.red support; the primitive is + # emitted with the 32dp x 32bit x 32rep shape only. + uses_ldtm_stat: bool = False # Causal S_q < S_kv shifts Q rows right by q_offset = S_kv - S_q. # This flag selects the shifted causal mask; there is no second causal mode. has_q_offset: bool = False @@ -379,11 +490,38 @@ class FmhaConfig: # 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 smem_q_head_dim(self) -> int: + """Q storage rounded to one 128-byte TMA fragment.""" + if self.logical_head_dim_qk == 192: + fragment_elements = 1024 // self.q_dtype.width + return ( + (self.logical_head_dim_qk + fragment_elements - 1) + // fragment_elements + * fragment_elements + ) + return self.qk_mma_tiler[2] + @property def single_qkv_instance(self) -> bool: """Return whether one work tile carries a single Q/KV/O instance.""" return self.num_qkv_instances == 1 + @property + def cta_group_size(self) -> int: + """CTAs cooperating on one UMMA: 2 in the two-CTA form, else 1.""" + return 2 if self.two_cta_umma else 1 + + @property + def kv_tile_rows_per_cta(self) -> int: + """K rows one CTA stages per K/V tile.""" + return self.kv_tile_n // self.cta_group_size + + @property + def pv_n_per_cta(self) -> int: + """V head-dim columns one CTA stages per tile.""" + return self.pv_mma_tiler[1] // self.cta_group_size + @property def uses_early_tile_sum(self) -> bool: """Return whether paired M128 geometry supports early V-tile reduction.""" @@ -411,11 +549,42 @@ def uses_d256_fp8_softmax_cadence(self) -> bool: and self.v_dtype == cutlass.Float8E4M3FN ) + @property + def smem_p_bytes(self) -> int: + """Bytes of one query group's SMEM P tile: q rows by k keys of the V dtype.""" + return self.qk_mma_tiler[0] * self.qk_mma_tiler[1] * self.v_dtype.width // 8 + + @property + def pv_half_overlap(self) -> bool: + """Publish P in two 64-key halves so PV can start on the first half. + Not used when P is staged in SMEM.""" + return ( + not self.single_qkv_instance + and self.enable_early_tile_sum + and not self.is_causal + and not self.has_varlen + and not self.has_tmem_p_pipeline + and not self.p_in_smem + and self.v_dtype is not None + and self.v_dtype.width in (8, 16) + and not self.uses_d128_fp8_softmax_cadence + and not self.uses_d256_fp8_softmax_cadence + # The head-dim-staged PV path issues every K slice per call. + and not self.stage_kv_by_head_dim + # Softmax halves the QK N tile and PV halves its K slices, so the + # two extents must match and split evenly. + and self.qk_mma_tiler[1] == self.pv_mma_tiler[2] + and (self.pv_mma_tiler[2] // (16 if self.v_dtype.width == 16 else 32)) % 2 + == 0 + ) + @property def uses_d128_fp8_softmax_cadence(self) -> bool: - """Return whether paired D128 FP8 uses interleaved softmax retirement.""" + """Paired D128 fp8 with P in TMEM retires softmax interleaved. Only the + shapes that do not stage P in SMEM take this path.""" return ( - not self.stage_kv_by_head_dim + not self.single_qkv_instance + and not self.p_in_smem and self.enable_early_tile_sum and self.q_dtype == cutlass.Float8E4M3FN and self.k_dtype == cutlass.Float8E4M3FN @@ -551,10 +720,13 @@ def uses_paired_fp8_head_batch_seq_tile_order(self) -> bool: Causal work uses this order for load balancing. Dense GQA uses it to keep Q-head groups that share the same K/V head adjacent. Dense MHA - has no cross-head K/V reuse and retains its sequence-local order. + has no cross-head K/V reuse and retains its sequence-local order. So + does the two-CTA form under GQA, whose clusters must pair adjacent Q + tiles of one head to share one K/V head. """ return ( (self.is_causal or self.h_r > 1) + and not self.two_cta_umma and not self.single_qkv_instance and self.q_dtype is not None and self.k_dtype is not None @@ -566,9 +738,10 @@ def uses_paired_fp8_head_batch_seq_tile_order(self) -> bool: @property def uses_head_batch_seq_tile_order(self) -> bool: - """Return whether work tiles use head_batch_seq coordinates.""" + """Return whether work tiles use head_batch_seq coordinates. Two-CTA + stays seq-first so the grid padding and the cluster axis agree.""" return self.uses_paired_fp8_head_batch_seq_tile_order or ( - self.is_causal and self.balance_causal_workload + self.is_causal and self.balance_causal_workload and not self.two_cta_umma ) @property @@ -582,9 +755,16 @@ def work_tile_coord_indices(self) -> tuple[int, int, int]: def pv_p_scale(self) -> float: """Return the P scale applied before PV MMA.""" if self.v_dtype is not None and self.v_dtype.width == 8: - # FP8 E4M3 has max finite magnitude 448; scaling P to that range - # before PV MMA preserves dynamic range. - return 448.0 + # E4M3 saturates at 448. Exact correction keeps P <= 1, so P scales + # by 448. Lazy correction lets the softmax row max lag the true one + # by up to corr_skip_threshold_log2, so P <= 2^threshold. Keep a + # power-of-two scale there. The largest with P * scale <= 448 is + # 2^floor(log2(448) - threshold), which is 1 at threshold 8. + if self.corr_skip_threshold_log2 == 0.0: + return 448.0 + return 2.0 ** max( + 0.0, math.floor(math.log2(448.0) - self.corr_skip_threshold_log2) + ) # Non-FP8 V uses P directly, so the PV-side P scale is identity. return 1.0 @@ -746,9 +926,10 @@ def consumer_wait( class S0S1SequenceResource(MemoryResource): """Sequence barrier between Softmax0 (producer) and Softmax1 (consumer). - Prevents both softmax groups from writing P to TMEM simultaneously. - Matches the handwritten FMHA kernel's ``s0_s1_sequence_mbar`` - PipelineAsync. + Paces peer iterations. The paired FP8 main loop hands the token back + before computing P, allowing overlap without letting one peer run ahead. + Each peer's SP pipeline separately protects its P stores and MMA reads. + Other paths retain serialized P computation. Single resource instance shared across tasks: - Softmax0's dst_resource (ProducerAcquire/Commit) @@ -789,6 +970,27 @@ class TmemStatsDoneResource(MemoryResource): is_barrier: cutlass.Constexpr[bool] = True +# --------------------------------------------------------------------------- +# TmemPPrefixReadyResource -- leading keys of P stored +# --------------------------------------------------------------------------- + + +@dataclass(kw_only=True) +class TmemPPrefixReadyResource(MemoryResource): + """Barrier: Softmax signals after storing the leading half of a P tile. + + Lets the MMA warp issue the PV MMAs for the first 64 keys while softmax + is still computing the rest. There is no matching barrier for the rest of + P; the existing SP stage guards the full tile and the S buffer reuse. + + Single resource instance shared across tasks: + - Softmax's dst_resource (ProducerAcquire/Commit) + - MMA's src_resource (ConsumerWait/Release) + """ + + is_barrier: cutlass.Constexpr[bool] = True + + # --------------------------------------------------------------------------- # GmemQKVResource -- global memory Q/K/V source (no pipeline) # --------------------------------------------------------------------------- @@ -1148,14 +1350,23 @@ def compute_page_coords(self, stage_info: StageInfo) -> tuple[Int32, Int32, Int3 ) -def _qkv_inner_dim_size_bytes(cfg: FmhaConfig) -> int: - """Return the byte width of one Q/K/V tile inner dimension.""" +def _qk_inner_dim_size_bytes(cfg: FmhaConfig) -> int: + """Return the byte width of one Q/K tile inner dimension.""" return cfg.qk_mma_tiler[2] * cfg.q_dtype.width // 8 -def _qkv_smem_layout(cfg: FmhaConfig) -> int: - """Return the tcgen05 descriptor layout selector for Q/K/V SMEM tiles.""" - inner_dim_size = _qkv_inner_dim_size_bytes(cfg) +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 + + +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 + + +def _smem_layout_for_inner_bytes(inner_dim_size: int) -> int: + """Return the tcgen05 descriptor swizzle selector for an SMEM row of this many bytes.""" if inner_dim_size % 128 == 0: return 2 if inner_dim_size == 64: @@ -1165,34 +1376,43 @@ def _qkv_smem_layout(cfg: FmhaConfig) -> int: raise RuntimeError(f"Unsupported inner dimension size: {inner_dim_size}") +def _qk_smem_layout(cfg: FmhaConfig) -> int: + """Return the tcgen05 descriptor layout selector for Q/K SMEM tiles.""" + return _smem_layout_for_inner_bytes(_qk_inner_dim_size_bytes(cfg)) + + +def _pv_smem_layout(cfg: FmhaConfig) -> int: + """Return the tcgen05 descriptor layout selector for V SMEM tiles.""" + return _smem_layout_for_inner_bytes(_pv_inner_dim_size_bytes(cfg)) + + def _qk_smem_desc_offsets(cfg: FmhaConfig) -> SmemDescOffsets: """Return Q/K descriptor leading and stride byte offsets.""" leading_byte_offset = 0 if cfg.head_paired else 16 - stride_byte_offset = ( - cfg.qk_mma_tiler[2] * cfg.q_dtype.width // cfg.tma_copy_qkv_iters - ) + stride_byte_offset = cfg.tma_copy_q_granu_inner * cfg.q_dtype.width return leading_byte_offset, stride_byte_offset def _pv_smem_desc_offsets(cfg: FmhaConfig) -> SmemDescOffsets: """Return V descriptor leading and stride byte offsets for PV MMA.""" + if cfg.two_cta_umma: + # One fragment of pv_n_per_cta columns per CTA. + return 0, cfg.tma_copy_v_granu_inner * cfg.v_dtype.width leading_byte_offset = 0 - if cfg.tma_copy_qkv_iters != 1: - tma_copy_kv_iters = ( - cfg.tma_copy_kv_stage_iters + if cfg.tma_copy_v_iters != 1: + tma_copy_v_iters = ( + cfg.tma_copy_v_stage_iters if cfg.stage_kv_by_head_dim - else cfg.tma_copy_qkv_iters + else cfg.tma_copy_v_iters ) - leading_byte_offset = cfg.tma_copy_kv_bytes // tma_copy_kv_iters - stride_byte_offset = ( - cfg.pv_mma_tiler[1] * cfg.v_dtype.width // cfg.tma_copy_qkv_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 return leading_byte_offset, stride_byte_offset -def _qkv_smem_swizzle(cfg: FmhaConfig) -> cutlass.Swizzle: - """Return the physical TMA swizzle used by Q/K/V SMEM fragments.""" - inner_dim_size = _qkv_inner_dim_size_bytes(cfg) +def _pv_smem_swizzle(cfg: FmhaConfig) -> cutlass.Swizzle: + """Return the physical TMA swizzle used by V SMEM fragments.""" + inner_dim_size = _pv_inner_dim_size_bytes(cfg) if inner_dim_size % 128 == 0: return cutlass.Swizzle(3, 4, 3) if inner_dim_size == 64: @@ -1204,7 +1424,14 @@ def _qkv_smem_swizzle(cfg: FmhaConfig) -> cutlass.Swizzle: def _smem_o_swizzle(cfg: FmhaConfig) -> cutlass.Swizzle: """Return the shared-memory swizzle used when staging O for TMA store.""" - return _qkv_smem_swizzle(cfg) + inner_dim_size = _o_inner_dim_size_bytes(cfg) + if inner_dim_size % 128 == 0: + return cutlass.Swizzle(3, 4, 3) + if inner_dim_size == 64: + return cutlass.Swizzle(2, 4, 3) + if inner_dim_size == 32: + return cutlass.Swizzle(1, 4, 3) + raise RuntimeError(f"Unsupported inner dimension size: {inner_dim_size}") # --------------------------------------------------------------------------- @@ -1325,12 +1552,32 @@ def tma_load( ragged_box_size=self.cfg.qk_mma_tiler[0], ragged_extent=q_seq_extent, ) - prims.cp_async_bulk_tensor_shared_cta_global( - sQ_curr.subview(i * self.cfg.tma_copy_q_granu_elems), - self.tma_q_desc, - q_coords, - stage_info.barrier, - ) + if cutlass.const_expr(self.cfg.two_cta_umma): + # Each CTA loads only its own 128 Q rows (multicast mask = own + # rank), but as a cta_group::2 copy, so the byte completion + # signals the mbarrier in the leader CTA's SMEM. The leader's + # MMA waits there for both CTAs' Q before the M=256 UMMA. + cta_rank = cute.arch.make_warp_uniform( + cute.arch.block_idx_in_cluster() + ) + prims.cp_async_bulk_tensor_shared_cluster_global( + sQ_curr.subview(i * self.cfg.tma_copy_q_granu_elems), + self.tma_q_desc, + q_coords, + 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( + sQ_curr.subview(i * self.cfg.tma_copy_q_granu_elems), + self.tma_q_desc, + q_coords, + stage_info.barrier, + ) def _build_q_descriptor(self, inst_idx: int) -> prims.Tcgen05SmemDesc: """Build SMEM descriptor for the current Q tile. @@ -1345,7 +1592,7 @@ def _build_q_descriptor(self, inst_idx: int) -> prims.Tcgen05SmemDesc: sQ_curr, leading_byte_offset=leading_byte_offset, stride_byte_offset=stride_byte_offset, - layout=_qkv_smem_layout(self.cfg), + layout=_qk_smem_layout(self.cfg), ) @consumer_work(returns=desc_q0_base) @@ -1605,17 +1852,16 @@ def dma_consumer_release_labels_for( ) -> set[str] | None: """Bind page-offset releases to the K/V TMA loads that consumed them.""" if isinstance(downstream, SmemKVResource): - if self.cfg.reuses_page_table_windows: - if self.page_table_is_v: - return {"v_load", "v_load_stage", "v_load_stage_cached"} - return {"k_load", "k_load_stage"} - return { - "k_load", - "v_load", - "k_load_stage", - "v_load_stage", - "v_load_stage_cached", - } + labels: set[str] = set() + binds_k = not (self.cfg.reuses_page_table_windows and self.page_table_is_v) + binds_v = not ( + self.cfg.reuses_page_table_windows and not self.page_table_is_v + ) + if downstream.holds_k and binds_k: + labels |= {"k_load", "k_load_stage"} + if downstream.holds_v and binds_v: + labels |= {"v_load", "v_load_stage", "v_load_stage_cached"} + return labels or None return None @@ -1631,6 +1877,10 @@ class SmemKVResource(MemoryResource): K and V tiles alternate in the pipeline stages: K0, V0, K1, V1, ... Producer: LoadTask (TMA loads K/V tiles). Consumer: MmaTask (builds SMEM descriptors for QK and PV MMAs). + + ``role="kv"`` shares one buffer and pipeline, which requires equal K and V + element widths. Mixed QK/PV dtypes instead build one ``role="k"`` and one + ``role="v"`` instance, each sized from its own dtype. """ sK_array: cutlass.Array = field(init=False, default=None) @@ -1643,9 +1893,11 @@ class SmemKVResource(MemoryResource): init=False, default=None ) block_tables: cute.Pointer | None = field(init=False, default=None) + role: Constexpr[str] = field(init=False, default="kv") cfg: Constexpr[FmhaConfig] = field(init=False, default=None) _alloc: Constexpr[Optional[SmemAllocation]] = field(init=False, default=None) desc_k_base: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + desc_k_stage_base: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() desc_v_base: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() def __init__( @@ -1657,10 +1909,19 @@ def __init__( page_offsets_kv: Optional["SmemPageOffsetsKvResource"] = None, page_offsets_v: Optional["SmemPageOffsetsKvResource"] = None, block_tables: cute.Pointer | None = None, + role: str = "kv", **kwargs: Any, ) -> None: - """Bind K/V TMA descriptors and reserve shared SMEM staging.""" + """Bind K/V TMA descriptors and reserve this role's SMEM staging.""" + if role not in ("kv", "k", "v"): + raise ValueError(f"unsupported SmemKVResource role: {role}") + if role == "kv" and cfg.k_dtype.width != cfg.v_dtype.width: + raise ValueError( + "a shared K/V buffer requires K and V to have the same element " + f"width; got {cfg.k_dtype.width} and {cfg.v_dtype.width}" + ) super().__init__(pipeline_config=pipeline_config, **kwargs) + self.role = role self.tma_k_desc = tma_k_desc self.tma_v_desc = tma_v_desc self.page_offsets_kv = page_offsets_kv @@ -1669,23 +1930,45 @@ def __init__( ) self.block_tables = block_tables self.cfg = cfg - total_elements = cfg.sK_shape[0] * cfg.sK_shape[1] - size_bytes = total_elements * cfg.k_dtype.width // 8 + # Every role stages pipeline_config.num_stages tiles; for "kv" that is + # cfg.kv_stage, which is cfg.sK_shape[0]. + total_elements = pipeline_config.num_stages * cfg.sK_shape[1] + size_bytes = total_elements * self._buffer_dtype(cfg, role).width // 8 self._alloc = SmemAllocation( - "smem_kv", size_bytes, alignment=cfg.buffer_align_bytes + f"smem_{role}", size_bytes, alignment=cfg.buffer_align_bytes ) - self.sK_array = _placeholder_smem_array(cfg.k_dtype) + self.sK_array = _placeholder_smem_array(self._buffer_dtype(cfg, role)) self.desc_k_base = TaskLocalVariable( dtype=cutlass.Int64, default=cutlass.Int64(0), docs="SMEM descriptor base for the current K tile.", ) + self.desc_k_stage_base = TaskLocalVariable( + dtype=cutlass.Int64, + default=cutlass.Int64(0), + docs="SMEM descriptor for the later K slice retained by both Q tiles.", + ) self.desc_v_base = TaskLocalVariable( dtype=cutlass.Int64, default=cutlass.Int64(0), docs="SMEM descriptor base for the current V tile.", ) + @staticmethod + def _buffer_dtype(cfg: FmhaConfig, role: str) -> type: + """Return the element type staged by a buffer with this role.""" + return cfg.v_dtype if role == "v" else cfg.k_dtype + + @property + def holds_k(self) -> bool: + """Return whether this buffer stages K tiles.""" + return self.role in ("kv", "k") + + @property + def holds_v(self) -> bool: + """Return whether this buffer stages V tiles.""" + return self.role in ("kv", "v") + def get_smem_requirements(self) -> list[SmemAllocation]: """Return the SMEM allocation required for staged K/V tiles.""" return [self._alloc] @@ -1694,10 +1977,10 @@ def get_smem_requirements(self) -> list[SmemAllocation]: def _init_smem_state(self, stage_info: StageInfo) -> None: """Materialize K/V SMEM storage and descriptor dataflow slots.""" smem_base = stage_info.context.smem_base - total_elements = self.cfg.sK_shape[0] * self.cfg.sK_shape[1] + total_elements = self.pipeline_config.num_stages * self.cfg.sK_shape[1] self.sK_array = cutlass.Array( smem_base.data_ptr() + self._alloc.offset, - dtype=self.cfg.k_dtype, + dtype=self._buffer_dtype(self.cfg, self.role), shape=(total_elements,), addrspace=3, ) @@ -1742,16 +2025,27 @@ def _tma_load( ) * self.cfg.kv_tile_n smem_stage_elements = self.cfg.tma_copy_kv_elements sK_curr = self.sK_array.subview(stage_info.stage_idx * smem_stage_elements) + # V carries its own copy granularity because v_dtype may be narrower. + # The two sets coincide when the widths match. + if cutlass.const_expr(is_v): + d_granu_inner = self.cfg.tma_copy_v_granu_inner + d_iter_elems = self.cfg.tma_copy_v_granu_elems + stage_iters = self.cfg.tma_copy_v_stage_iters + else: + d_granu_inner = self.cfg.tma_copy_kv_granu_inner + d_iter_elems = self.cfg.tma_copy_kv_granu_elems + stage_iters = self.cfg.tma_copy_kv_stage_iters if cutlass.const_expr(self.cfg.use_paged_kv): # Paged-KV path: read pre-staged page IDs and issue one TMA per - # (page fragment, d fragment). K and V share the native rank-4 - # descriptor coordinates (d_off, 0, kv_head_coord, page_id). + # (page fragment, d fragment). The descriptor is shaped + # (d_inner, num_tokens_per_page, h_kv, total_pages); coords are + # (d_off, 0, kv_head_coord, page_id). SMEM layout per stage matches + # the contiguous path: two d-halves (d_iter_elems each) with page + # fragments concatenated along the seq axis inside each d-half. tile_idx = kv_tile_start + stage_info.loop_offset + tile_offset pages_per_tile = self.cfg.kv_tile_n // self.cfg.num_tokens_per_page - d_granu_inner = self.cfg.tma_copy_kv_granu_inner page_d_elems = self.cfg.num_tokens_per_page * d_granu_inner - d_iter_elems = self.cfg.tma_copy_kv_granu_elems if prims.elect_sync(): # Only the elected TMA-issuing lane consumes page IDs. Loading # the vector outside this guard made every lane perform the @@ -1788,7 +2082,7 @@ def _tma_load( ) for frag in cutlass.range_constexpr(pages_per_tile): page_id = Int32(page_ids[frag]) - for i in cutlass.range_constexpr(self.cfg.tma_copy_kv_stage_iters): + for i in cutlass.range_constexpr(stage_iters): d_offset = Int32( head_dim_stage_idx * self.cfg.head_dim_per_stage_kv + i * d_granu_inner @@ -1803,22 +2097,49 @@ def _tma_load( return if prims.elect_sync(): - d_granu_inner = self.cfg.tma_copy_kv_granu_inner seq_coord_kv = cuseqlen_k + seq_offset - for i in cutlass.range_constexpr(self.cfg.tma_copy_kv_stage_iters): + d_base = Int32(0) + if cutlass.const_expr(self.cfg.two_cta_umma): + # This CTA stages its half of the K rows or V columns. + cta_rank = cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) + if cutlass.const_expr(is_v): + d_base = cta_rank * Int32(self.cfg.pv_n_per_cta) + else: + seq_coord_kv = seq_coord_kv + cta_rank * Int32( + self.cfg.kv_tile_rows_per_cta + ) + for i in cutlass.range_constexpr(stage_iters): d_offset = ( head_dim_stage_idx * self.cfg.head_dim_per_stage_kv + i * d_granu_inner ) - kv_coords = (d_offset, kv_head_coord, seq_coord_kv, batch_coord) - if cutlass.const_expr(self.cfg.has_varlen): - kv_coords = (d_offset, kv_head_coord, seq_coord_kv) - prims.cp_async_bulk_tensor_shared_cta_global( - sK_curr.subview(i * self.cfg.tma_copy_kv_granu_elems), - tma_desc, - kv_coords, - stage_info.barrier, + kv_coords = ( + d_offset + d_base, + kv_head_coord, + seq_coord_kv, + batch_coord, ) + if cutlass.const_expr(self.cfg.has_varlen): + kv_coords = (d_offset + d_base, kv_head_coord, seq_coord_kv) + if cutlass.const_expr(self.cfg.two_cta_umma): + prims.cp_async_bulk_tensor_shared_cluster_global( + sK_curr.subview(i * d_iter_elems), + tma_desc, + kv_coords, + 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( + sK_curr.subview(i * d_iter_elems), + tma_desc, + kv_coords, + stage_info.barrier, + ) @producer_work @cute.jit @@ -1993,10 +2314,24 @@ def k_desc(self, stage_info: StageInfo) -> prims.Tcgen05SmemDesc: sK_curr, leading_byte_offset=leading_byte_offset, stride_byte_offset=stride_byte_offset, - layout=_qkv_smem_layout(self.cfg), + layout=_qk_smem_layout(self.cfg), ) return desc_k_base + @consumer_work(returns=desc_k_stage_base) + @cute.jit + def k_stage_desc(self, stage_info: StageInfo) -> prims.Tcgen05SmemDesc: + """Keep the later K descriptor live alongside the first K and V.""" + smem_stage_elements = self.cfg.tma_copy_kv_elements + sK_curr = self.sK_array.subview(stage_info.stage_idx * smem_stage_elements) + leading_byte_offset, stride_byte_offset = _qk_smem_desc_offsets(self.cfg) + return prims.Tcgen05SmemDesc.build( + sK_curr, + leading_byte_offset=leading_byte_offset, + stride_byte_offset=stride_byte_offset, + layout=_qk_smem_layout(self.cfg), + ) + @cute.jit def _zero_paged_v_tail( self, @@ -2029,11 +2364,11 @@ def _zero_paged_v_tail( # iteration, and D iterations are concatenated within the stage. # Mirror that exact physical layout: a flat row-major clear would # target the wrong bytes under the s128b swizzle. - d_granu_inner = self.cfg.tma_copy_kv_granu_inner + d_granu_inner = self.cfg.tma_copy_v_granu_inner chunks_per_d_iter = d_granu_inner // 16 - chunks_per_v_row = self.cfg.tma_copy_kv_stage_iters * chunks_per_d_iter + chunks_per_v_row = self.cfg.tma_copy_v_stage_iters * chunks_per_d_iter page_d_elems = self.cfg.num_tokens_per_page * d_granu_inner - d_iter_elems = self.cfg.tma_copy_kv_granu_elems + d_iter_elems = self.cfg.tma_copy_v_granu_elems invalid_chunks = (Int32(self.cfg.kv_tile_n) - valid_rows) * Int32( chunks_per_v_row ) @@ -2068,7 +2403,7 @@ def _zero_paged_v_tail( sV_curr.subview(smem_offset).data_ptr().store_swizzled( zero_vec, alignment=16, - swizzle=_qkv_smem_swizzle(self.cfg), + swizzle=_pv_smem_swizzle(self.cfg), ) # v_desc is called only after this stage's skv.wait(), which makes @@ -2090,7 +2425,7 @@ def _build_v_descriptor(self, stage_info: StageInfo) -> prims.Tcgen05SmemDesc: sK_curr, leading_byte_offset=leading_byte_offset, stride_byte_offset=stride_byte_offset, - layout=_qkv_smem_layout(self.cfg), + layout=_pv_smem_layout(self.cfg), ) @consumer_work(returns=desc_v_base) @@ -2121,6 +2456,46 @@ def v_desc_paged( return self._build_v_descriptor(stage_info) +# --------------------------------------------------------------------------- +# SmemPResource -- fp8 P tile staged in SMEM per query group +# --------------------------------------------------------------------------- +@dataclass(kw_only=True) +class SmemPResource(MemoryResource): + """One K-major SW128 fp8 P tile per query group in SMEM: softmax produces, + the PV UMMA consumes, so softmax can release the S stage early.""" + + cfg: Constexpr[FmhaConfig] = field(init=False, default=None) + _alloc: Constexpr[Optional[SmemAllocation]] = field(init=False, default=None) + + def __init__( + self, + pipeline_config: PipelineConfig, + cfg: FmhaConfig, + group_idx: int, + **kwargs: Any, + ) -> None: + super().__init__(pipeline_config=pipeline_config, **kwargs) + self.cfg = cfg + self._alloc = SmemAllocation( + f"smem_p{group_idx}", cfg.smem_p_bytes, alignment=cfg.buffer_align_bytes + ) + + def get_smem_requirements(self) -> list[SmemAllocation]: + return [self._alloc] + + @property + def row_bytes(self) -> int: + return self.cfg.qk_mma_tiler[1] * self.cfg.v_dtype.width // 8 + + def descriptor_offsets(self) -> SmemDescOffsets: + """LBO and SBO of the K-major swizzled tile: eight rows per swizzle atom.""" + leading_byte_offset = 16 + return leading_byte_offset, 8 * self.row_bytes + + def descriptor_layout(self) -> int: + return _smem_layout_for_inner_bytes(self.row_bytes) + + # --------------------------------------------------------------------------- # TmemSPResource -- TMEM S/P ping-pong buffer with UmmaAsync pipeline # --------------------------------------------------------------------------- @@ -2167,8 +2542,11 @@ class TmemSPResource(MemoryResource): row_max: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() row_sum: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() p_chunk: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + p_lo: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() q_offset: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() seqlen_k: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + # p_in_smem: this group's SMEM P tile. + smem_p: Optional[SmemPResource] = field(init=False, default=None) variable_window_start: Constexpr[TaskLocalVariable] = ( TaskLocalVariable.uninitialized() ) @@ -2192,6 +2570,7 @@ def __init__( variable_window_cta_starts: cute.Tensor | None = None, variable_window_q_stride: int | Int32 = 0, scale_softmax_log2: cute.Tensor | None = None, + smem_p: Optional[SmemPResource] = None, **kwargs: Any, ) -> None: """Bind S/P TMEM offsets, Q peer index, and optional varlen metadata.""" @@ -2218,6 +2597,7 @@ def __init__( self.tmem_ptr_s_cached = _placeholder_tmem_ptr() self.tmem_s_addr_cached = Int32(0) self.tmem_p_addr_cached = Int32(0) + self.smem_p = smem_p self.old_row_max = TaskLocalVariable( dtype=Float32, default=Float32(-Float32.inf), @@ -2245,6 +2625,11 @@ def __init__( default_factory=lambda: _placeholder_softmax_chunks(cfg), docs="P fragments retained for post-release row-sum reduction.", ) + self.p_lo = TaskLocalVariable( + dtype=Float32, + default=Float32(0.0), + docs="FP32 sum of the first half of the current probability tile.", + ) self.q_offset = TaskLocalVariable( dtype=Int32, default=Int32(self.q_offset_default), @@ -2308,6 +2693,7 @@ def uses_fixed_dense_k_tail_mask(self) -> bool: return ( not self.cfg.is_causal and not self.cfg.has_varlen + and not self.cfg.has_variable_window and self.cfg.fixed_dense_k_tail > 0 ) @@ -2376,6 +2762,26 @@ def _stage_col_offset(self, stage_info: StageInfo) -> Int32 | int: stage_col_offset = stage_info.stage_idx * self.cfg.qk_mma_tiler[1] return stage_col_offset + @producer_work + @cute.jit + def qk_mma_stage( + self, + stage_info: StageInfo, + *, + desc_q_base: prims.Tcgen05SmemDesc, + desc_k_stage: prims.Tcgen05SmemDesc, + section: cutlass.Constexpr[FmhaStage], + head_dim_stage_idx: cutlass.Constexpr[int], + ) -> None: + """Accumulate a later K slice without replacing the first slice's binding.""" + self._qk_mma_impl( + stage_info, + desc_q_base=desc_q_base, + desc_k_base=desc_k_stage, + section=section, + head_dim_stage_idx=head_dim_stage_idx, + ) + @producer_work @cute.jit def qk_mma( @@ -2387,6 +2793,27 @@ def qk_mma( section: cutlass.Constexpr[FmhaStage], head_dim_stage_idx: cutlass.Constexpr[int] = 0, is_tail: cutlass.Constexpr[bool] = False, + ) -> None: + """Issue one QK slice, clipping MMA phases to the logical head dimension.""" + self._qk_mma_impl( + stage_info, + desc_q_base=desc_q_base, + desc_k_base=desc_k_base, + section=section, + head_dim_stage_idx=head_dim_stage_idx, + is_tail=is_tail, + ) + + @cute.jit + def _qk_mma_impl( + self, + stage_info: StageInfo, + *, + desc_q_base: prims.Tcgen05SmemDesc, + desc_k_base: prims.Tcgen05SmemDesc, + section: cutlass.Constexpr[FmhaStage], + head_dim_stage_idx: cutlass.Constexpr[int] = 0, + is_tail: cutlass.Constexpr[bool] = False, ) -> None: """QK MMA: compute Q*K -> S in TMEM. @@ -2428,21 +2855,36 @@ def qk_mma( a_dtype=ab_format, b_dtype=ab_format, n_dim=self.cfg.qk_mma_tiler[1], - m_dim=self.cfg.qk_mma_tiler[0], + m_dim=self.cfg.qk_mma_tiler[0] * self.cfg.cta_group_size, ) k_dim_per_mma = 16 if cutlass.const_expr(self.cfg.q_dtype.width != 16): k_dim_per_mma = 32 - num_kphases_qk = self.cfg.qk_mma_tiler[2] // k_dim_per_mma inc_bytes_qk = k_dim_per_mma * self.cfg.q_dtype.width // 8 - num_kphases_per_tma = num_kphases_qk // self.cfg.tma_copy_qkv_iters - chunk_bytes_qk = inc_bytes_qk * num_kphases_per_tma + num_kphases_per_tma = self.cfg.tma_copy_q_granu_inner // k_dim_per_mma + # Byte stride between TMA fragments of the Q and K tiles. They differ only + # in the two-CTA form, where a K stage holds half the rows. + chunk_bytes_q = inc_bytes_qk * num_kphases_per_tma + chunk_bytes_k = chunk_bytes_q if cutlass.const_expr(self.cfg.tma_copy_qkv_iters != 1): - chunk_bytes_qk = ( + chunk_bytes_q = ( + self.cfg.tma_copy_q_granu_elems * self.cfg.q_dtype.width // 8 + ) + chunk_bytes_k = ( self.cfg.tma_copy_kv_bytes // self.cfg.tma_copy_kv_stage_iters ) + cta_group = ( + prims.CTAGroup.CTA_2 + if cutlass.const_expr(self.cfg.two_cta_umma) + else prims.CTAGroup.CTA_1 + ) + issue_mma = cutlass.Boolean(True) + if cutlass.const_expr(self.cfg.two_cta_umma): + issue_mma = ( + cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) == 0 + ) num_tma_iters_qk = self.cfg.tma_copy_qkv_iters if cutlass.const_expr(self.cfg.stage_kv_by_head_dim): num_tma_iters_qk = self.cfg.tma_copy_kv_stage_iters @@ -2460,22 +2902,36 @@ def qk_mma( scale_d = 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_qk * q_tma_iter - k_tma_iter_offset = chunk_bytes_qk * tma_iter - for k_idx in cutlass.range_constexpr(num_kphases_per_tma): + q_tma_iter_offset = chunk_bytes_q * q_tma_iter + k_tma_iter_offset = chunk_bytes_k * tma_iter + # The final 128-wide K stage may contain only 64 logical + # elements (non-absorbed MLA). Do not issue MMAs on padding. + valid_kphases = min( + num_kphases_per_tma, + max( + 0, + ( + self.cfg.logical_head_dim_qk + - q_tma_iter * self.cfg.tma_copy_q_granu_inner + ) + // k_dim_per_mma, + ), + ) + for k_idx in cutlass.range_constexpr(valid_kphases): local_increment = inc_bytes_qk * k_idx dq = desc_q_base_ + ((local_increment + q_tma_iter_offset) >> 4) dk = desc_k_base_ + ((local_increment + k_tma_iter_offset) >> 4) - if prims.elect_sync(): - prims.tcgen05_mma( - mma_kind, - prims.CTAGroup.CTA_1, - tmem_ptr_s, - dq, - dk, - idesc_qk, - scale_d, - ) + if issue_mma: + if prims.elect_sync(): + prims.tcgen05_mma( + mma_kind, + cta_group, + tmem_ptr_s, + dq, + dk, + idesc_qk, + scale_d, + ) scale_d = True @producer_work @@ -2755,27 +3211,144 @@ def _reduce_row_max( row_max_safe = Float32(0.0) return old_row_max, row_max_safe + @cute.jit + def _load_s_chunks_and_reduce_row_max( + self, + stage_info: StageInfo, + row_max: SoftmaxScalar, + *, + causal_loop_mask: cutlass.Constexpr[bool] = False, + q_offset: Int32 = 0, + ) -> tuple[SoftmaxScalar, SoftmaxScalar]: + """LDTM.STAT: fuse S load and per-chunk row_max via tcgen05.ld.red.max. + + Instead of loading each S chunk with ``tcgen05.ld`` and then folding + a per-lane max in registers, use ``tcgen05.ld.red.max`` so the + reduction happens as part of the TMEM load. For causal loop masking, + only tiles crossing the runtime right bound need a software reduction + of masked scores; wholly visible tiles retain the hardware maximum. + """ + 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 + old_row_max = row_max + s_data: SoftmaxChunks = [None] * num_chunks + chunk_maxima: tuple[Any, ...] = () + for chunk_idx in cutlass.range_constexpr(num_chunks): + if cutlass.const_expr(hasattr(prims, "tcgen05_ld_red")): + loaded_words, red_word = prims.tcgen05_ld_red( + prims.Tcgen05LdStShape.SHAPE_32X32B, + prims.make_tmem_ptr( + tmem_s_addr + chunk_idx * tmem_x, self.cfg.qk_acc_dtype + ), + prims.ReductionKind.MAX, + num=tmem_x, + ) + s_data[chunk_idx] = cutlass.Vector.from_elements( + tuple( + loaded_words[i].bitcast(self.cfg.qk_acc_dtype) + for i in range(tmem_x) + ), + dtype=self.cfg.qk_acc_dtype, + ) + chunk_max = self.cfg.qk_acc_dtype( + arith_dialect.bitcast( + self.cfg.qk_acc_dtype.mlir_type, red_word.ir_value() + ) + ) + else: + # PTX ISA 8.8 supports LDTM.STAT on SM103 before the DSL 4.8 + # convenience wrapper is available. The geometry is fixed at + # 32 rows x 32 FP32 values for each context load fragment. + assert tmem_x == 32 + loaded = load_tmem_32x32b_max(tmem_s_addr + chunk_idx * tmem_x) + s_data[chunk_idx] = cutlass.Vector.from_elements( + loaded[:32], dtype=self.cfg.qk_acc_dtype + ) + chunk_max = loaded[32] + chunk_maxima += (chunk_max,) + cute.arch.fence_view_async_tmem_load() + # Reduction results are asynchronous, just like the loaded scores. + tile_max = cutlass.Vector.from_elements( + chunk_maxima, self.cfg.qk_acc_dtype + ).reduce("max") + row_max = cute.math.max(row_max, tile_max) + if cutlass.const_expr(causal_loop_mask): + seq_coord, _, _ = _resolve_work_tile_coords( + self.cfg, stage_info.work_tile.tile_idx + ) + q_min = ( + q_offset + + seq_coord * self.cfg.cta_tiler[0] + + self.q_half * self.cfg.q_tile_m + ) + kv_base = stage_info.loop_offset * self.cfg.kv_tile_n + k_max = kv_base + self.cfg.qk_mma_tiler[1] - Int32(1) + if q_min < k_max: + q_idx = ( + q_min + + (cute.arch.warp_idx() % 4) * cute.arch.WARP_SIZE + + cute.arch.lane_idx() + ) + # Discard the unmasked maxima when any key can be invisible. + masked_row_max = old_row_max + for chunk_idx in cutlass.range_constexpr(num_chunks): + num_valid = cute.math.min( + cute.math.max( + q_idx - kv_base - chunk_idx * tmem_x + Int32(1), + Int32(0), + ), + 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) + ) + s_data[chunk_idx] = cutlass.vector.where( + mask, s_data[chunk_idx], neg_inf + ) + masked_row_max = cute.math.max( + masked_row_max, s_data[chunk_idx].reduce("max") + ) + row_max = masked_row_max + _tmem_sp_sdata[id(self)] = s_data + row_max_safe = row_max + if row_max == -Float32.inf: + row_max_safe = Float32(0.0) + return old_row_max, row_max_safe + @cute.jit def _exp2_p_store( self, stage_col_offset: TmemAddr, row_max: SoftmaxScalar, scale_softmax_log2: SoftmaxScalar, + chunk_lo: cutlass.Constexpr[int] = 0, + chunk_hi: cutlass.Constexpr[int | None] = None, + sum_in: Float32 | None = None, ) -> SoftmaxRowSumContribution: - """Apply exp2 softmax P, fold the PV P scale, and store P to TMEM.""" + """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. + """ tmem_p_addr = self.tmem_p_addr_cached + stage_col_offset tmem_shape = "32x32b" tmem_x = self.cfg.tmem_x_load_s num_chunks = self.cfg.qk_mma_tiler[1] // tmem_x + if cutlass.const_expr(chunk_hi is None): + chunk_hi = num_chunks + 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.uses_d256_fp8_softmax_cadence): + 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, row_max, scale, ) - if cutlass.const_expr(self.cfg.uses_d128_fp8_softmax_cadence): + if cutlass.const_expr(not partial and self.cfg.uses_d128_fp8_softmax_cadence): return self._exp2_p_store_d128_fp8_cadence( tmem_p_addr, row_max, @@ -2789,14 +3362,19 @@ def _exp2_p_store( ) p_scale_log2 = Float32(self.cfg.pv_p_scale_log2) minus_row_max_scale = (Float32(0.0) - row_max) * scale + p_scale_log2 - s_data = _tmem_sp_sdata.pop(id(self)) - if cutlass.const_expr(self.enable_early_tile_sum): - # Keep four independent scalar dependency chains while expressing - # them as two packed float2 values. The explicit packed primitive - # lowers to FADD2 for D128 instead of two scalar FADDs per pair. + 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.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. local_sum_pair_0 = (Float32(0.0), Float32(0.0)) local_sum_pair_1 = (Float32(0.0), Float32(0.0)) - for chunk_idx in cutlass.range_constexpr(num_chunks): + local_sum_pair_2 = (Float32(0.0), Float32(0.0)) + local_sum_pair_3 = (Float32(0.0), Float32(0.0)) + for chunk_idx in cutlass.range_constexpr(chunk_lo, chunk_hi): p_vals = () for elem_idx in cutlass.range_constexpr(0, tmem_x, 2): fma_pair = cute.arch.fma_packed_f32x2( @@ -2809,39 +3387,63 @@ def _exp2_p_store( rnd="rn", ftz=False, ) - p0 = cute.math.exp2(fma_pair[0], fastmath=True) - p1 = cute.math.exp2(fma_pair[1], fastmath=True) - if cutlass.const_expr(self.enable_early_tile_sum): + if cutlass.const_expr( + elem_idx // 2 >= tmem_x // 2 - self.cfg.exp2_fma_pairs + ): + p0, p1 = _exp2_fma_packed(fma_pair[0], fma_pair[1]) + else: + p0 = cute.math.exp2(fma_pair[0], fastmath=True) + p1 = cute.math.exp2(fma_pair[1], fastmath=True) + if cutlass.const_expr(self.enable_early_tile_sum or partial): pair_idx = chunk_idx * (tmem_x // 2) + elem_idx // 2 - if cutlass.const_expr(pair_idx % 2 == 0): + if cutlass.const_expr(pair_idx % 4 == 0): local_sum_pair_0 = cute.arch.add_packed_f32x2( local_sum_pair_0, (p0, p1), rnd="rn", ftz=False, ) - else: + elif cutlass.const_expr(pair_idx % 4 == 1): local_sum_pair_1 = cute.arch.add_packed_f32x2( local_sum_pair_1, (p0, p1), rnd="rn", ftz=False, ) + elif cutlass.const_expr(pair_idx % 4 == 2): + local_sum_pair_2 = cute.arch.add_packed_f32x2( + local_sum_pair_2, (p0, p1), rnd="rn", ftz=False + ) + else: + local_sum_pair_3 = cute.arch.add_packed_f32x2( + local_sum_pair_3, (p0, p1), rnd="rn", ftz=False + ) p_vals += (p0, p1) s_data[chunk_idx] = cutlass.Vector.from_elements( p_vals, self.cfg.qk_acc_dtype ) use_fused_d128_fp8x4_pack = ( - not self.cfg.stage_kv_by_head_dim + not self.cfg.single_qkv_instance and self.cfg.v_dtype == cutlass.Float8E4M3FN ) - for pair_idx in cutlass.range_constexpr(num_chunks // p_packing_ratio): + # One word holds four fp8 P values, so a pair of chunks is tmem_x words + # and a chunk range inside a pair is a word sub-range of its store. + 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 + ) if cutlass.const_expr(use_fused_d128_fp8x4_pack): # Match the handwritten D128 pack: merge both FP8x2 # conversions in one side-effecting block so ptxas can retain # the 32-bit word without a PRMT between temporary vectors. packed_words: tuple[Any, ...] = () - for word_idx in cutlass.range_constexpr(tmem_x): + for word_idx in cutlass.range_constexpr(word_lo, word_hi): flat_idx = word_idx * 4 chunk_idx = pair_idx * p_packing_ratio + flat_idx // tmem_x elem_idx = flat_idx % tmem_x @@ -2853,6 +3455,8 @@ def _exp2_p_store( ) packed_words += (packed_word,) store_fragment = cutlass.Vector.from_elements(packed_words, Int32) + if cutlass.const_expr(self.cfg.p_in_smem): + _tmem_sp_pwords[id(self)] = packed_words else: for slice_idx in cutlass.range_constexpr(p_packing_ratio): chunk_idx = pair_idx * p_packing_ratio + slice_idx @@ -2863,19 +3467,29 @@ def _exp2_p_store( else: p_data_packed[slice_idx * tmem_x : tmem_x] = p_chunk_dtype store_fragment = p_data_f32[0:tmem_x] - prims.tcgen05_st( - tmem_shape, - prims.make_tmem_ptr(tmem_p_addr + pair_idx * tmem_x, cutlass.Int8), - store_fragment, + if cutlass.const_expr(not self.cfg.p_in_smem): + prims.tcgen05_st( + tmem_shape, + prims.make_tmem_ptr( + tmem_p_addr + pair_idx * tmem_x + word_lo, cutlass.Int8 + ), + store_fragment, + ) + if cutlass.const_expr(self.enable_early_tile_sum or partial): + local_sum_pair_0 = cute.arch.add_packed_f32x2( + local_sum_pair_0, local_sum_pair_2, rnd="rn", ftz=False + ) + local_sum_pair_1 = cute.arch.add_packed_f32x2( + local_sum_pair_1, local_sum_pair_3, rnd="rn", ftz=False ) - if cutlass.const_expr(self.enable_early_tile_sum): local_sum_pair = cute.arch.add_packed_f32x2( - local_sum_pair_0, - local_sum_pair_1, - rnd="rn", - ftz=False, + local_sum_pair_0, local_sum_pair_1, 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 ): @@ -2894,6 +3508,35 @@ def _exp2_p_store( result.append(s_data[chunk_idx]) return result + @cute.jit + def _store_p_row_smem( + self, stage_info: StageInfo, packed_words: tuple[Any, ...] + ) -> None: + """Write this thread's 128-byte P row in the SW128 K-major layout the PV + descriptor expects: 16-byte chunk c of row r lands at chunk c xor (r mod 8).""" + assert self.smem_p is not None + context = stage_info.context + assert context is not None and context.smem_base is not None + view = cutlass.Array( + context.smem_base.data_ptr() + self.smem_p._alloc.offset, + dtype=Int32, + shape=(self.cfg.smem_p_bytes // 4,), + addrspace=3, + ) + # One P row per softmax thread, indexed like its TMEM lane. + 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()) + row_words = row * Int32(32) + swz = row & Int32(7) + for chunk_idx in cutlass.range_constexpr(len(packed_words) // 4): + chunk_words = cutlass.Vector.from_elements( + packed_words[chunk_idx * 4 : chunk_idx * 4 + 4], Int32 + ) + phys_chunk = Int32(chunk_idx) ^ swz + view.subview(row_words + phys_chunk * Int32(4)).data_ptr().store( + chunk_words, alignment=16 + ) + @cute.jit def _exp2_p_store_d128_fp8_cadence( self, @@ -3098,7 +3741,7 @@ def _exp2_p_store_d256_fp8_cadence( packed_words: tuple[Any, ...] = () local_sum_pair_0 = (Float32(0.0), Float32(0.0)) local_sum_pair_1 = (Float32(0.0), Float32(0.0)) - use_two_sum_pairs = not self.cfg.stage_kv_by_head_dim + use_two_sum_pairs = not self.cfg.single_qkv_instance for flat_idx in cutlass.range_constexpr(0, num_values, 2): if cutlass.const_expr(flat_idx >= 8): delayed_idx = flat_idx - 8 @@ -3196,6 +3839,8 @@ def compute_row_max( 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) @@ -3206,13 +3851,22 @@ def fixed_dense_k_tail_masked_row_max( 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.""" + """Exclude TMA zero-fill lanes in a partial fixed dense K/V tile. + + ``section=Tail``: called once for the last tile, mask always applied. + ``section=Loop``: used in loop, masks 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 stage_info.loop_offset == stage_info.loop_end - Int32(1): + 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( [tmem_x], self.cfg.qk_acc_dtype(-Float32.inf), @@ -3407,12 +4061,28 @@ def loop_masked_row_max( q_offset: Int32, ) -> tuple[SoftmaxScalar, SoftmaxScalar]: """Loop stage: apply causal masking for mixed Q right-offset batches.""" + if cutlass.const_expr(self.cfg.uses_ldtm_stat): + return self._load_s_chunks_and_reduce_row_max( + stage_info, row_max, causal_loop_mask=True, q_offset=q_offset + ) s_data = self._load_s_chunks(stage_info) s_data = self._apply_causal_mask_for_kv_tile( stage_info, s_data, kv_tile_idx=stage_info.loop_offset, q_offset=q_offset ) return self._reduce_row_max(s_data, row_max) + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def store_p(self, stage_info: StageInfo) -> None: + """Store the P row from exp2_p into SMEM after the P-tile acquire; the + proxy fence orders the stores before the UMMA reads them. Auxiliary work + because the S stage is already released and P goes to the SMEM tile.""" + 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, + ) + @consumer_work(returns=p_chunk) @cute.jit def exp2_p( @@ -3427,6 +4097,58 @@ def exp2_p( self._stage_col_offset(stage_info), row_max, scale_softmax_log2 ) + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=p_chunk) + @cute.jit + def exp2_p_smem( + self, + stage_info: StageInfo, + *, + row_max: SoftmaxScalar, + scale_softmax_log2: SoftmaxScalar, + ) -> SoftmaxRowSumContribution: + """exp2_p with P in SMEM, run on the loaded S after the stage release.""" + return self._exp2_p_store( + self._stage_col_offset(stage_info), row_max, scale_softmax_log2 + ) + + @consumer_work(returns=p_lo) + @cute.jit + def exp2_p_lo( + self, + stage_info: StageInfo, + *, + row_max: SoftmaxScalar, + scale_softmax_log2: SoftmaxScalar, + ) -> Float32: + """First half of exp2_p: keys [0, N/2). Stores those P columns and + fences so the MMA can start PV on them. Returns their row sum.""" + return self._exp2_p_store( + self._stage_col_offset(stage_info), + row_max, + scale_softmax_log2, + chunk_hi=self.cfg.qk_mma_tiler[1] // self.cfg.tmem_x_load_s // 2, + sum_in=Float32(0.0), + ) + + @consumer_work(returns=p_chunk) + @cute.jit + def exp2_p_hi( + self, + stage_info: StageInfo, + *, + row_max: SoftmaxScalar, + scale_softmax_log2: SoftmaxScalar, + p_lo: Float32, + ) -> Float32: + """Second half of exp2_p: keys [N/2, N). Returns the full tile sum.""" + return self._exp2_p_store( + self._stage_col_offset(stage_info), + row_max, + scale_softmax_log2, + chunk_lo=self.cfg.qk_mma_tiler[1] // self.cfg.tmem_x_load_s // 2, + sum_in=p_lo, + ) + @consumer_work(returns=(old_row_max, row_max)) @cute.jit def right_masked_row_max( @@ -3599,6 +4321,20 @@ def masked_exp2_p( self._stage_col_offset(stage_info), row_max, scale_softmax_log2 ) + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=p_chunk) + @cute.jit + def masked_exp2_p_smem( + self, + stage_info: StageInfo, + *, + row_max: SoftmaxScalar, + scale_softmax_log2: SoftmaxScalar, + ) -> SoftmaxRowSumContribution: + """masked_exp2_p with P in SMEM, after the stage release.""" + return self._exp2_p_store( + self._stage_col_offset(stage_info), row_max, scale_softmax_log2 + ) + @consumer_work(returns=(old_row_max, row_max)) @cute.jit def invalid_row_max( @@ -3665,6 +4401,24 @@ def softmax_aux_identity( _ = stage_info return row_max + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=row_max) + @cute.jit + def freeze_row_max( + self, + stage_info: StageInfo, + *, + old_row_max: SoftmaxScalar, + row_max: SoftmaxScalar, + scale_softmax_log2: Float32, + ) -> SoftmaxScalar: + """Keep the old row max when the tile raises it by at most the threshold.""" + _ = stage_info + threshold = Float32(self.cfg.corr_skip_threshold_log2) + frozen = row_max + if (row_max - old_row_max) * scale_softmax_log2 <= threshold: + frozen = old_row_max + return frozen + @cute.jit def _row_sum_reduction( self, @@ -4141,12 +4895,12 @@ def read_vec( class TmemOResource(MemoryResource): """TMEM O accumulation with a topology-derived UmmaAsync pipeline. - Producer: MMA writes P*V -> O (double-buffered O0/O1). + Producer: MMA writes P*V into one O accumulator per Q instance. Consumer: Correction rescales O in-place. Paired schedules use two stages so MMA can commit O0, work on O1, - commit O1, then acquire O0. SMEM-stats D256 uses one stage because it - writes one physical O accumulator. + commit O1, then acquire O0. Single-instance schedules use one stage + because they write one physical O accumulator. """ cfg: Constexpr[FmhaConfig] = field(init=False, default=None) @@ -4160,6 +4914,9 @@ class TmemOResource(MemoryResource): tmem_o_addr_base_cached: TmemAddr | None = field(init=False, default=None) # P-stage base supplied by TmemPResource for split S/P scheduling. tmem_p_base_cached: TmemAddr | None = field(init=False, default=None) + # p_in_smem: per-group SMEM P tiles read by the SS PV MMA. + smem_p0_resource: Optional[SmemPResource] = field(init=False, default=None) + smem_p1_resource: Optional[SmemPResource] = field(init=False, default=None) # References to TmemStats resources for reading cached correction stats. # consumer_work reads stats from these instead of from TMEM, because # the stats TMEM region overlaps with S0/S1 and can be overwritten by @@ -4178,11 +4935,15 @@ def __init__( tmem_o1_offset: int, tmem_vec0_resource: TmemStatsResource | None = None, tmem_vec1_resource: TmemStatsResource | None = None, + smem_p0_resource: Optional[SmemPResource] = None, + smem_p1_resource: Optional[SmemPResource] = None, **kwargs: Any, ) -> None: - """Bind O TMEM offsets and correction-stat resources.""" + """Bind O TMEM offsets, correction-stat resources, and SMEM P tiles.""" super().__init__(pipeline_config=pipeline_config, **kwargs) self.cfg = cfg + self.smem_p0_resource = smem_p0_resource + self.smem_p1_resource = smem_p1_resource self.tmem_o0_offset = tmem_o0_offset self.tmem_o1_offset = tmem_o1_offset self.tmem_vec0_resource = tmem_vec0_resource @@ -4274,9 +5035,14 @@ def pv_mma( 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). + ``k_half`` issues only the first (0) or second (1) half of the K slices + so the schedule can run PV on the leading half of P before the rest is + stored. The second half accumulates onto the first. + Uses captured schedule section and call index to select O0/O1 and scale_d statically. Reads P from TMEM (via tmem_p offset on the corresponding SP resource) @@ -4287,7 +5053,12 @@ def pv_mma( The task domain pads partial final CTAs so this slot is always outside peer0's causal reach. """ - if cutlass.const_expr(section == FmhaStage.Head): + if cutlass.const_expr(self.cfg.p_in_smem): + # PV0(i), PV1(i) every iteration and in the tail, so inst_idx is the group. + writes_o0 = inst_idx == 0 + first_o0_write = False + first_o1_write_maybe = False + elif cutlass.const_expr(section == FmhaStage.Head): writes_o0 = True first_o0_write = True first_o1_write_maybe = False @@ -4300,15 +5071,17 @@ def pv_mma( first_o0_write = False first_o1_write_maybe = True - # In causal mode, check if O0 MMA should skip the last LOOP iteration. + # In causal mode, check if O0 MMA should skip peer 0's invalid last tile: + # the last loop iteration, or the tail when P in SMEM moves PV0 there. skip_o0_invalid = False - if cutlass.const_expr( - self.cfg.skip_causal_invalid_peer0 - and writes_o0 - and section == FmhaStage.Loop - ): - if not is_tail: - skip_o0_invalid = stage_info.loop_offset == (stage_info.loop_end - 1) + if cutlass.const_expr(self.cfg.skip_causal_invalid_peer0 and writes_o0): + if cutlass.const_expr(self.cfg.p_in_smem): + skip_o0_invalid = is_tail + elif cutlass.const_expr(section == FmhaStage.Loop): + if not is_tail: + skip_o0_invalid = stage_info.loop_offset == ( + stage_info.loop_end - 1 + ) if not skip_o0_invalid: tmem_ptr_raw = self.tmem_ptr_raw_cached @@ -4333,10 +5106,20 @@ def pv_mma( if self.cfg.single_qkv_instance and self.cfg.pv_mma_tiler[1] == 256 else self.cfg.pv_mma_tiler[1] ), - m_dim=self.cfg.pv_mma_tiler[0], + m_dim=self.cfg.pv_mma_tiler[0] * self.cfg.cta_group_size, # V is row-major / MN-major. b_major=1, ) + cta_group = ( + prims.CTAGroup.CTA_2 + if cutlass.const_expr(self.cfg.two_cta_umma) + else prims.CTAGroup.CTA_1 + ) + issue_mma = cutlass.Boolean(True) + if cutlass.const_expr(self.cfg.two_cta_umma): + issue_mma = ( + cute.arch.make_warp_uniform(cute.arch.block_idx_in_cluster()) == 0 + ) pv_n_dim = self.cfg.pv_mma_tiler[1] num_head_dim_stages = 1 @@ -4359,20 +5142,26 @@ def pv_mma( k_dim_per_mma * self.cfg.v_dtype.width // self.cfg.qk_acc_dtype.width ) tma_copy_iters_per_head_dim_stage = ( - self.cfg.tma_copy_qkv_iters // num_head_dim_stages + self.cfg.tma_copy_v_iters // num_head_dim_stages ) if cutlass.const_expr(self.cfg.stage_kv_by_head_dim): - tma_copy_iters_per_head_dim_stage = self.cfg.tma_copy_kv_stage_iters + tma_copy_iters_per_head_dim_stage = self.cfg.tma_copy_v_stage_iters inc_bytes_v = ( k_dim_per_mma * (pv_n_dim // tma_copy_iters_per_head_dim_stage) * self.cfg.v_dtype.width // 8 ) - kv_chunk_bytes = ( - self.cfg.tma_copy_kv_bytes // self.cfg.tma_copy_kv_stage_iters - ) - head_dim_stage_bytes_v = kv_chunk_bytes * tma_copy_iters_per_head_dim_stage + if cutlass.const_expr(self.cfg.two_cta_umma): + # One V fragment of tma_copy_v_granu_inner columns per CTA. + inc_bytes_v = ( + k_dim_per_mma + * self.cfg.tma_copy_v_granu_inner + * self.cfg.v_dtype.width + // 8 + ) + v_chunk_bytes = self.cfg.tma_copy_v_bytes // self.cfg.tma_copy_v_stage_iters + head_dim_stage_bytes_v = v_chunk_bytes * tma_copy_iters_per_head_dim_stage # Select O buffer and P offset at trace time (compile-time constant) if cutlass.const_expr(self.cfg.single_qkv_instance or writes_o0): @@ -4406,6 +5195,11 @@ def pv_mma( scale_d = stage_info.loop_offset > 0 else: scale_d = False + elif cutlass.const_expr(self.cfg.p_in_smem): + if cutlass.const_expr(is_tail): + scale_d = stage_info.loop_end > 0 + else: + scale_d = stage_info.loop_offset > 0 elif cutlass.const_expr(first_o0_write): # Head O0 is the first O0 write. scale_d = False @@ -4430,18 +5224,39 @@ def pv_mma( dp = tmem_ptr_raw.subview(tmem_p_base + k_idx * inc_tmem_p) increment = (inc_bytes_v * k_idx) >> 4 dv = desc_v_base_ + increment - if prims.elect_sync(): - prims.tcgen05_mma( - mma_kind, - prims.CTAGroup.CTA_1, - tmem_ptr_o_stage, - dp, - dv, - idesc_pv, - scale_d_stage, - ) + if issue_mma: + if prims.elect_sync(): + prims.tcgen05_mma( + mma_kind, + cta_group, + tmem_ptr_o_stage, + dp, + dv, + idesc_pv, + scale_d_stage, + ) scale_d_stage = True else: + if cutlass.const_expr(self.cfg.p_in_smem): + smem_p = ( + self.smem_p0_resource if writes_o0 else self.smem_p1_resource + ) + sP = cutlass.Array( + stage_info.context.smem_base.data_ptr() + smem_p._alloc.offset, + dtype=cutlass.Int8, + shape=(self.cfg.smem_p_bytes,), + addrspace=3, + ) + p_lbo, p_sbo = smem_p.descriptor_offsets() + desc_p_base_ = freeze_smem_descriptor( + prims.Tcgen05SmemDesc.build( + sP, + leading_byte_offset=p_lbo, + stride_byte_offset=p_sbo, + layout=smem_p.descriptor_layout(), + ) + ) + inc_bytes_p = k_dim_per_mma * self.cfg.v_dtype.width // 8 for head_dim_stage_idx in cutlass.range_constexpr( num_head_dim_stages_to_issue ): @@ -4449,21 +5264,30 @@ def pv_mma( v_stage_increment = ( head_dim_stage_bytes_v * head_dim_stage_idx ) >> 4 - scale_d_stage = scale_d - for k_idx in cutlass.range_constexpr(num_kphases_pv): - dp = tmem_ptr_raw.subview(tmem_p_base + k_idx * inc_tmem_p) + if cutlass.const_expr(k_half is None): + k_lo, k_hi = 0, num_kphases_pv + else: + k_lo = k_half * (num_kphases_pv // 2) + k_hi = k_lo + num_kphases_pv // 2 + scale_d_stage = scale_d if k_lo == 0 else True + for k_idx in cutlass.range_constexpr(k_lo, k_hi): + if cutlass.const_expr(self.cfg.p_in_smem): + dp = desc_p_base_ + ((inc_bytes_p * k_idx) >> 4) + else: + dp = tmem_ptr_raw.subview(tmem_p_base + k_idx * inc_tmem_p) increment = v_stage_increment + ((inc_bytes_v * k_idx) >> 4) dv = desc_v_base_ + increment - if prims.elect_sync(): - prims.tcgen05_mma( - mma_kind, - prims.CTAGroup.CTA_1, - tmem_ptr_o_stage, - dp, - dv, - idesc_pv, - scale_d_stage, - ) + if issue_mma: + if prims.elect_sync(): + prims.tcgen05_mma( + mma_kind, + cta_group, + tmem_ptr_o_stage, + dp, + dv, + idesc_pv, + scale_d_stage, + ) scale_d_stage = True @consumer_work @@ -4538,7 +5362,7 @@ def _correct_impl( skip_o0_invalid = stage_info.loop_offset == (stage_info.loop_end - 1) # Check if we should skip correction (when old_max == new_max) - should_rescale = True + 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) should_rescale = vote_ballot_cnt != Int32(0) @@ -4565,9 +5389,10 @@ def _correct_impl( tmem_o_addr = self.tmem_o_addr_base_cached + tmem_o_offset tmem_shape = "32x32b" - tmem_x = 16 + # Amortize TMEM handshakes with 32-value correction transfers. + tmem_x = 32 - num_iters = self.cfg.cta_tiler[2] // tmem_x + num_iters = self.cfg.pv_mma_tiler[1] // tmem_x for i in cutlass.range_constexpr(num_iters): tmem_tile_addr = tmem_o_addr + i * tmem_x tmem_ptr = cutlass.inttoptr( @@ -4762,7 +5587,7 @@ def _store_o( 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): - o_head_dim = self.cfg.head_dim_per_stage_kv + o_head_dim = self.cfg.head_dim_per_stage_o tma_copy_o_iters = self.cfg.tma_copy_o_stage_iters tmem_offset_o = ( @@ -4950,7 +5775,7 @@ def tma_store( if prims.elect_sync(): for i in cutlass.range_constexpr(tma_copy_o_iters): d_offset = ( - head_dim_stage_idx * self.cfg.head_dim_per_stage_kv + head_dim_stage_idx * self.cfg.head_dim_per_stage_o + i * self.cfg.tma_copy_o_granu_inner ) o_coords = (d_offset, head_coord, seq_offset_o, batch_coord) 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 72a71f89567d..2e2d8581378e 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 @@ -53,9 +53,11 @@ SmemQResource, TmemOResource, TmemPResource, + SmemPResource, TmemSPResource, TmemStatsResource, TmemStatsDoneResource, + TmemPPrefixReadyResource, ) @@ -180,7 +182,8 @@ def _captured_loop_bounds( def create_load_task( gmem_qkv: GmemQKVResource, smem_q: SmemQResource, - smem_kv: SmemKVResource, + smem_k_or_kv: SmemKVResource, + smem_v: SmemKVResource | None, work_queue: WorkQueue | None, task_class: type[Task] = Task, smem_page_offsets_kv: SmemPageOffsetsKvResource | None = None, @@ -196,20 +199,29 @@ def create_load_task( loop_start, loop_end, loop_step = _captured_loop_bounds(task_class, task_kwargs) skip_work_tile_if = _packed_context_skip_predicate(work_queue) src = _src_resources(gmem_qkv, work_queue=work_queue) - dst = [smem_q, smem_kv] + # One K/V resource when they share a buffer, two when their dtypes split it. + split_kv = smem_k_or_kv.cfg.split_kv_pipelines + if split_kv and smem_v is None: + 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_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_q.cfg.single_qkv_instance and smem_q.cfg.has_tmem_p_pipeline: - num_head_dim_stages_k = smem_kv.cfg.num_head_dim_stages_k - num_head_dim_stages_v = smem_kv.cfg.num_head_dim_stages_v + 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 if smem_page_offsets_v is not None: if smem_page_offsets_kv is None: raise ValueError("a V page window requires a matching K page window") - pages_per_tile = smem_kv.cfg.kv_tile_n // smem_kv.cfg.num_tokens_per_page - page_window_period = smem_kv.cfg.page_table_window_entries // pages_per_tile + pages_per_tile = ( + smem_k_or_kv.cfg.kv_tile_n // smem_k_or_kv.cfg.num_tokens_per_page + ) + page_window_period = ( + smem_k_or_kv.cfg.page_table_window_entries // pages_per_tile + ) if ( not isinstance(loop_start, int) or not isinstance(loop_end, int) @@ -227,14 +239,16 @@ def create_load_task( def load_reused_page_windows_schedule_body( gqkv: GmemQKVResource, sq: SmemQResource, - skv: SmemKVResource, + sk: SmemKVResource, + sv: SmemKVResource, spok: SmemPageOffsetsKvResource, spov: SmemPageOffsetsKvResource, wq: WorkQueue | None, ) -> None: """Load staged K/V while retaining each page-ID window.""" sq.init_load_state() - skv.init_load_state() + sk.init_load_state() + sv.init_load_state() spok.init_read_state() cached_v_page_ids = spov.init_cached_read_state() @@ -266,10 +280,11 @@ def load_reused_page_windows_schedule_body( sq.commit() def load_k_tile(*, tile_offset: int) -> None: + """Issue the K TMA loads for one tile across the K head-dim stages.""" for head_dim_stage_idx in range(num_head_dim_stages_k): - skv.try_acquire() - skv.acquire() - skv.k_load_stage( + sk.try_acquire() + sk.acquire() + sk.k_load_stage( stage_id=head_dim_stage_idx, tile_offset=tile_offset, kv_head_coord=kv_head_coord, @@ -280,7 +295,7 @@ def load_k_tile(*, tile_offset: int) -> None: kv_request_begin=kv_request_begin, kv_page_idx_ub=kv_page_idx_ub, ) - skv.commit() + sk.commit() def cache_v_tile(*, tile_offset: int) -> None: nonlocal cached_v_page_ids @@ -293,11 +308,16 @@ def cache_v_tile(*, tile_offset: int) -> None: def load_v_tile( *, tile_offset: int, reuse_cached_page_ids: bool = False ) -> None: + """Issue the V TMA loads for one tile across the V head-dim stages. + + With ``reuse_cached_page_ids`` the loads reuse page IDs staged by + ``cache_v_tile`` instead of re-reading the page table. + """ for head_dim_stage_idx in range(num_head_dim_stages_v): - skv.try_acquire() - skv.acquire() + sv.try_acquire() + sv.acquire() if reuse_cached_page_ids: - skv.v_load_stage_cached( + sv.v_load_stage_cached( cached_v_page_ids=cached_v_page_ids, stage_id=head_dim_stage_idx, tile_offset=tile_offset, @@ -310,7 +330,7 @@ def load_v_tile( kv_page_idx_ub=kv_page_idx_ub, ) else: - skv.v_load_stage( + sv.v_load_stage( stage_id=head_dim_stage_idx, tile_offset=tile_offset, kv_head_coord=kv_head_coord, @@ -321,7 +341,7 @@ def load_v_tile( kv_request_begin=kv_request_begin, kv_page_idx_ub=kv_page_idx_ub, ) - skv.commit() + sv.commit() # Window zero: K stays one tile ahead of V. Cache the final # V IDs before releasing the window because its last V @@ -378,13 +398,31 @@ def load_reused_page_windows_schedule( spov: SmemPageOffsetsKvResource, wq: WorkQueue | None = None, ) -> None: - load_reused_page_windows_schedule_body(gqkv, sq, skv, spok, spov, wq) + """Shared-buffer captured schedule.""" + load_reused_page_windows_schedule_body( + gqkv, sq, skv, skv, spok, spov, wq + ) + + @schedule + def load_reused_page_windows_split_schedule( + gqkv: GmemQKVResource, + sq: SmemQResource, + sk: SmemKVResource, + sv: SmemKVResource, + spok: SmemPageOffsetsKvResource, + spov: SmemPageOffsetsKvResource, + wq: WorkQueue | None = None, + ) -> None: + """Split K/V captured schedule.""" + load_reused_page_windows_schedule_body(gqkv, sq, sk, sv, spok, spov, wq) captured_schedule = _schedule_with_work_queue( - load_reused_page_windows_schedule, + load_reused_page_windows_split_schedule + if split_kv + else load_reused_page_windows_schedule, gmem_qkv, smem_q, - smem_kv, + *kv_resources, smem_page_offsets_kv, smem_page_offsets_v, work_queue=work_queue, @@ -392,10 +430,10 @@ def load_reused_page_windows_schedule( return task_class( src_resources=src, dst_resources=dst, - warp_idx=smem_kv.cfg.load_warp_id, + warp_idx=smem_k_or_kv.cfg.load_warp_id, num_warps=1, schedule=captured_schedule, - num_registers=smem_kv.cfg.num_regs_other, + num_registers=smem_k_or_kv.cfg.num_regs_other, name="LoadTask", **task_kwargs, ) @@ -403,12 +441,19 @@ def load_reused_page_windows_schedule( def load_schedule_body( gqkv: GmemQKVResource, sq: SmemQResource, - skv: SmemKVResource, + sk: SmemKVResource, + sv: SmemKVResource, spo: SmemPageOffsetsKvResource | None, wq: WorkQueue | None, ) -> None: + """Load-warp schedule: stage Q, then stream K and V tiles through their rings. + + K runs one tile ahead of V so QK(i+1) can start while PV(i) waits for V. + Page offsets are read once per tile when paged. + """ sq.init_load_state() - skv.init_load_state() + sk.init_load_state() + sv.init_load_state() if spo is not None: spo.init_read_state() with _work_tile_schedule_loop(wq, skip_if=skip_work_tile_if): @@ -440,12 +485,12 @@ def load_schedule_body( sq.commit() for head_dim_stage_idx in range(num_head_dim_stages_k): - skv.try_acquire() + sk.try_acquire() if spo is not None and head_dim_stage_idx == 0: spo.wait() spo.read_offsets() - skv.acquire() - skv.k_load_stage( + sk.acquire() + sk.k_load_stage( stage_id=head_dim_stage_idx, kv_head_coord=kv_head_coord, batch_coord=batch_coord, @@ -455,18 +500,18 @@ def load_schedule_body( kv_request_begin=kv_request_begin, kv_page_idx_ub=kv_page_idx_ub, ) - skv.commit() + sk.commit() if spo is not None: spo.release() with domain_loop(loop_start + 1, loop_end, loop_step): for head_dim_stage_idx in range(num_head_dim_stages_k): - skv.try_acquire() + sk.try_acquire() if spo is not None and head_dim_stage_idx == 0: spo.wait() spo.read_offsets() - skv.acquire() - skv.k_load_stage( + sk.acquire() + sk.k_load_stage( stage_id=head_dim_stage_idx, kv_head_coord=kv_head_coord, batch_coord=batch_coord, @@ -476,17 +521,17 @@ def load_schedule_body( kv_request_begin=kv_request_begin, kv_page_idx_ub=kv_page_idx_ub, ) - skv.commit() + sk.commit() if spo is not None: spo.release() for head_dim_stage_idx in range(num_head_dim_stages_v): - skv.try_acquire() + sv.try_acquire() if spo is not None and head_dim_stage_idx == 0: spo.wait() spo.read_offsets() - skv.acquire() - skv.v_load_stage( + sv.acquire() + sv.v_load_stage( stage_id=head_dim_stage_idx, previous=True, kv_head_coord=kv_head_coord, @@ -497,17 +542,17 @@ def load_schedule_body( kv_request_begin=kv_request_begin, kv_page_idx_ub=kv_page_idx_ub, ) - skv.commit() + sv.commit() if spo is not None: spo.release() for head_dim_stage_idx in range(num_head_dim_stages_v): - skv.try_acquire() + sv.try_acquire() if spo is not None and head_dim_stage_idx == 0: spo.wait() spo.read_offsets() - skv.acquire() - skv.v_load_stage( + sv.acquire() + sv.v_load_stage( stage_id=head_dim_stage_idx, previous=False, kv_head_coord=kv_head_coord, @@ -518,7 +563,7 @@ def load_schedule_body( kv_request_begin=kv_request_begin, kv_page_idx_ub=kv_page_idx_ub, ) - skv.commit() + sv.commit() if spo is not None: spo.release() @@ -529,7 +574,19 @@ def load_schedule( skv: SmemKVResource, wq: WorkQueue | None = None, ) -> None: - load_schedule_body(gqkv, sq, skv, None, wq) + """Shared-buffer captured schedule.""" + load_schedule_body(gqkv, sq, skv, skv, None, wq) + + @schedule + def load_split_schedule( + gqkv: GmemQKVResource, + sq: SmemQResource, + sk: SmemKVResource, + sv: SmemKVResource, + wq: WorkQueue | None = None, + ) -> None: + """Split K/V captured schedule.""" + load_schedule_body(gqkv, sq, sk, sv, None, wq) @schedule def load_page_offsets_schedule( @@ -539,28 +596,47 @@ def load_page_offsets_schedule( spo: SmemPageOffsetsKvResource, wq: WorkQueue | None = None, ) -> None: - load_schedule_body(gqkv, sq, skv, spo, wq) + """Shared-buffer captured schedule with a page-ID ring.""" + load_schedule_body(gqkv, sq, skv, skv, spo, wq) + + @schedule + def load_page_offsets_split_schedule( + gqkv: GmemQKVResource, + sq: SmemQResource, + sk: SmemKVResource, + sv: SmemKVResource, + spo: SmemPageOffsetsKvResource, + wq: WorkQueue | None = None, + ) -> None: + """Split K/V captured schedule with a page-ID ring.""" + load_schedule_body(gqkv, sq, sk, sv, spo, wq) if smem_page_offsets_kv is None: captured_schedule = _schedule_with_work_queue( - load_schedule, gmem_qkv, smem_q, smem_kv, work_queue=work_queue + load_split_schedule if split_kv else load_schedule, + gmem_qkv, + smem_q, + *kv_resources, + work_queue=work_queue, ) else: captured_schedule = _schedule_with_work_queue( - load_page_offsets_schedule, + load_page_offsets_split_schedule + if split_kv + else load_page_offsets_schedule, gmem_qkv, smem_q, - smem_kv, + *kv_resources, smem_page_offsets_kv, work_queue=work_queue, ) return task_class( src_resources=src, dst_resources=dst, - warp_idx=smem_kv.cfg.load_warp_id, + warp_idx=smem_k_or_kv.cfg.load_warp_id, num_warps=1, schedule=captured_schedule, - num_registers=smem_kv.cfg.num_regs_other, + num_registers=smem_k_or_kv.cfg.num_regs_other, name="LoadTask", **task_kwargs, ) @@ -573,15 +649,17 @@ def load_page_offsets_schedule( def load_schedule_body( gqkv: GmemQKVResource, sq: SmemQResource, - skv: SmemKVResource, + sk: SmemKVResource, + sv: SmemKVResource, wq: WorkQueue | None, ) -> None: """Load paired Q instances and their directly addressed K/V tiles.""" sq.init_load_state() - skv.init_load_state() + sk.init_load_state() + sv.init_load_state() 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 the SmemKV pipeline. + # only stream the next K/V tiles through their respective pipelines. with domain_loop(loop_start, loop_end, loop_step) as d: with d.first_iter(): ( @@ -610,37 +688,67 @@ def load_schedule_body( inst_idx=0, ) sq.commit() - # Throttle TMA before reserving a KV stage. - skv.try_acquire() - # Load Ki, with K0 handled by the first iteration. - skv.acquire() - skv.k_load( - kv_head_coord=kv_head_coord, - batch_coord=batch_coord, - cuseqlen_k=cuseqlen_k, - seqlen_k=seqlen_k, - kv_tile_start=kv_tile_start, - kv_request_begin=kv_request_begin, - kv_page_idx_ub=kv_page_idx_ub, - ) - skv.commit() - with d.first_iter(): - # Load Q1 for the second Q tile in this work tile. - sq.acquire() - sq.tma_load( - seq_coord_q=seq_coord_q, - head_coord=head_coord, + if smem_k_or_kv.cfg.stage_kv_by_head_dim: + with d.first_iter(): + # Load Q1 for the second Q tile in this work tile. + sq.acquire() + sq.tma_load( + seq_coord_q=seq_coord_q, + head_coord=head_coord, + batch_coord=batch_coord, + cuseqlen_q=cuseqlen_q, + seqlen_q=seqlen_q, + inst_idx=1, + ) + sq.commit() + for head_dim_stage_idx in range( + smem_k_or_kv.cfg.num_head_dim_stages_k + ): + sk.try_acquire() + sk.acquire() + sk.k_load_stage( + stage_id=head_dim_stage_idx, + kv_head_coord=kv_head_coord, + batch_coord=batch_coord, + cuseqlen_k=cuseqlen_k, + kv_tile_start=kv_tile_start, + seqlen_k=seqlen_k, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, + ) + sk.commit() + else: + # Throttle TMA before reserving a KV stage. + sk.try_acquire() + # Load Ki, with K0 handled by the first iteration. + sk.acquire() + sk.k_load( + kv_head_coord=kv_head_coord, batch_coord=batch_coord, - cuseqlen_q=cuseqlen_q, - seqlen_q=seqlen_q, - inst_idx=1, + cuseqlen_k=cuseqlen_k, + kv_tile_start=kv_tile_start, + seqlen_k=seqlen_k, + kv_request_begin=kv_request_begin, + kv_page_idx_ub=kv_page_idx_ub, ) - sq.commit() + sk.commit() + with d.first_iter(): + # Load Q1 for the second Q tile in this work tile. + sq.acquire() + sq.tma_load( + seq_coord_q=seq_coord_q, + head_coord=head_coord, + batch_coord=batch_coord, + cuseqlen_q=cuseqlen_q, + seqlen_q=seqlen_q, + inst_idx=1, + ) + sq.commit() # Throttle TMA before reserving a KV stage. - skv.try_acquire() + sv.try_acquire() # Load Vi, with V0 handled by the first iteration. - skv.acquire() - skv.v_load( + sv.acquire() + sv.v_load( kv_head_coord=kv_head_coord, batch_coord=batch_coord, cuseqlen_k=cuseqlen_k, @@ -649,7 +757,7 @@ def load_schedule_body( kv_request_begin=kv_request_begin, kv_page_idx_ub=kv_page_idx_ub, ) - skv.commit() + sv.commit() @schedule def load_schedule( @@ -658,20 +766,35 @@ def load_schedule( skv: SmemKVResource, wq: WorkQueue | None = None, ) -> None: - """Contiguous-KV captured schedule.""" + """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, wq) # type: ignore[call-arg] + load_schedule_body(gqkv, sq, skv, skv, wq) # type: ignore[call-arg] + + @schedule + def load_split_schedule( + gqkv: GmemQKVResource, + sq: SmemQResource, + sk: SmemKVResource, + sv: SmemKVResource, + 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] captured_schedule = _schedule_with_work_queue( - load_schedule, gmem_qkv, smem_q, smem_kv, work_queue=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, - warp_idx=smem_kv.cfg.load_warp_id, + warp_idx=smem_k_or_kv.cfg.load_warp_id, num_warps=1, schedule=captured_schedule, - num_registers=smem_kv.cfg.num_regs_other, + num_registers=smem_k_or_kv.cfg.num_regs_other, name="LoadTask", **task_kwargs, ) @@ -680,7 +803,8 @@ def load_schedule( def create_mma_task( gmem_qkv: GmemQKVResource, smem_q: SmemQResource, - smem_kv: SmemKVResource, + smem_k_or_kv: SmemKVResource, + smem_v: SmemKVResource | None, tmem_sp0: TmemSPResource, tmem_sp1: TmemSPResource | None, tmem_p0: TmemPResource | None, @@ -689,6 +813,10 @@ def create_mma_task( tmem_vec_done_1: TmemStatsDoneResource | None, work_queue: WorkQueue | None, task_class: type[Task] = Task, + tmem_p_prefix_ready_0: TmemPPrefixReadyResource | None = None, + tmem_p_prefix_ready_1: TmemPPrefixReadyResource | None = None, + smem_p0: SmemPResource | None = None, + smem_p1: SmemPResource | None = None, **task_kwargs: Any, ) -> Task: """Create the one-warp MMA compute task.""" @@ -696,7 +824,32 @@ def create_mma_task( skip_work_tile_if = _packed_context_skip_predicate(work_queue) # Only paged MMA schedules read request coordinates from global metadata. qkv_resources = [gmem_qkv] if smem_q.cfg.use_paged_kv else [] - src = _src_resources(*qkv_resources, smem_q, smem_kv, work_queue=work_queue) + split_kv = smem_k_or_kv.cfg.split_kv_pipelines + if split_kv and smem_v is None: + 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,) + pv_half_overlap = smem_q.cfg.pv_half_overlap + if pv_half_overlap and ( + tmem_p_prefix_ready_0 is None or tmem_p_prefix_ready_1 is None + ): + raise ValueError("pv_half_overlap requires both P-prefix barriers") + p_prefix_resources = ( + (tmem_p_prefix_ready_0, tmem_p_prefix_ready_1) if pv_half_overlap else () + ) + p_in_smem = smem_q.cfg.p_in_smem + 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 () + src = _src_resources( + *qkv_resources, + smem_q, + *kv_resources, + *p_prefix_resources, + *smem_p_resources, + work_queue=work_queue, + ) + 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 if ( smem_q.cfg.single_qkv_instance @@ -704,10 +857,10 @@ def create_mma_task( and tmem_p0 is not None ): split_src = _src_resources( - *qkv_resources, smem_q, smem_kv, tmem_p0, work_queue=work_queue + *qkv_resources, smem_q, *kv_resources, tmem_p0, work_queue=work_queue ) - num_head_dim_stages_k = smem_kv.cfg.num_head_dim_stages_k - num_head_dim_stages_v = smem_kv.cfg.num_head_dim_stages_v + 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 @@ -715,19 +868,25 @@ def create_mma_task( ): raise ValueError("loop-carried split S/P scheduling expects two K/V stages") - @schedule - def mma_schedule( + def mma_schedule_body( gqkv: GmemQKVResource, sq: SmemQResource, - skv: SmemKVResource, + sk: SmemKVResource, + sv: SmemKVResource, sp0: TmemSPResource, tp0: TmemPResource, to: TmemOResource, vd0: TmemStatsDoneResource, wq: WorkQueue | None = None, ) -> None: + """MMA-warp schedule with separate S and P TMEM pipelines. + + Prologue issues QK for the first tile; the steady-state loop then alternates + QK(i+1) into S with PV(i) from P, consuming K and V from independent rings. + """ sq.init_descriptor_state() - skv.init_descriptor_state() + sk.init_descriptor_state() + sv.init_descriptor_state() sp0.init_mma_state() to.init_mma_state() with _work_tile_schedule_loop(wq, skip_if=skip_work_tile_if): @@ -758,15 +917,15 @@ def mma_schedule( vd0.acquire() sp0.acquire() for head_dim_stage_idx in range(num_head_dim_stages_k): - skv.wait() - desc_k_base = skv.k_desc() + 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.Head, head_dim_stage_idx=head_dim_stage_idx, ) - skv.release() + sk.release() sp0.commit() if not smem_q.cfg.stats_via_smem: vd0.commit() @@ -778,25 +937,25 @@ def mma_schedule( with domain_loop(loop_start, loop_end, loop_step): if not smem_q.cfg.stats_via_smem: vd0.acquire() - skv.wait() - desc_k_base = skv.k_desc() + 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=0, ) - skv.release() + sk.release() - skv.wait() - desc_k_base = skv.k_desc() + 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, ) - skv.release() + sk.release() sp0.commit() if not smem_q.cfg.stats_via_smem: vd0.commit() @@ -807,37 +966,41 @@ def mma_schedule( tmem_p_base = tp0.p_base() to.set_p_base(tmem_p_base=tmem_p_base) - skv.wait() + # TODO: sv.wait() only depends on V's own TMA-load barrier, + # not on tp0 (softmax P). Now that K/V have independent + # pipelines, it could be issued earlier to overlap with the + # K/QK/softmax work above instead of waiting until here. + sv.wait() if cutlass.const_expr(smem_q.cfg.needs_paged_v_tail_clear): - desc_v_base = skv.v_desc_paged( + 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 = skv.v_desc() + desc_v_base = sv.v_desc() to.pv_mma( desc_v_base=desc_v_base, section=FmhaStage.Loop, head_dim_stage_idx=0, ) - skv.release() + sv.release() - skv.wait() + sv.wait() if cutlass.const_expr(smem_q.cfg.needs_paged_v_tail_clear): - desc_v_base = skv.v_desc_paged( + 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 = skv.v_desc() + desc_v_base = sv.v_desc() to.pv_mma( desc_v_base=desc_v_base, section=FmhaStage.Loop, head_dim_stage_idx=1, ) - skv.release() + sv.release() to.commit() tp0.release() @@ -847,22 +1010,22 @@ def mma_schedule( tmem_p_base = tp0.p_base() to.set_p_base(tmem_p_base=tmem_p_base) for head_dim_stage_idx in range(num_head_dim_stages_v): - skv.wait() + sv.wait() if cutlass.const_expr(smem_q.cfg.needs_paged_v_tail_clear): - desc_v_base = skv.v_desc_paged( + desc_v_base = sv.v_desc_paged( section=FmhaStage.Tail, seqlen_k=v_seqlen_k, kv_tile_start=v_kv_tile_start, ) else: - desc_v_base = skv.v_desc() + desc_v_base = sv.v_desc() to.pv_mma( desc_v_base=desc_v_base, section=FmhaStage.Tail, head_dim_stage_idx=head_dim_stage_idx, is_tail=True, ) - skv.release() + sv.release() to.commit() tp0.release() if not smem_q.cfg.stats_via_smem: @@ -873,11 +1036,40 @@ def mma_schedule( tp0.wait() tp0.release() + @schedule + def mma_schedule( + gqkv: GmemQKVResource, + sq: SmemQResource, + skv: SmemKVResource, + sp0: TmemSPResource, + tp0: TmemPResource, + to: TmemOResource, + vd0: TmemStatsDoneResource, + wq: WorkQueue | None = None, + ) -> None: + """Shared-buffer captured schedule.""" + mma_schedule_body(gqkv, sq, skv, skv, sp0, tp0, to, vd0, wq) + + @schedule + def mma_split_schedule( + gqkv: GmemQKVResource, + sq: SmemQResource, + sk: SmemKVResource, + sv: SmemKVResource, + sp0: TmemSPResource, + tp0: TmemPResource, + to: TmemOResource, + vd0: TmemStatsDoneResource, + wq: WorkQueue | None = None, + ) -> None: + """Split K/V captured schedule.""" + mma_schedule_body(gqkv, sq, sk, sv, sp0, tp0, to, vd0, wq) + captured_schedule = _schedule_with_work_queue( - mma_schedule, + mma_split_schedule if split_kv else mma_schedule, gmem_qkv, smem_q, - smem_kv, + *kv_resources, tmem_sp0, tmem_p0, tmem_o, @@ -897,21 +1089,24 @@ def mma_schedule( ) if smem_q.cfg.single_qkv_instance: - num_head_dim_stages_k = smem_kv.cfg.num_head_dim_stages_k - num_head_dim_stages_v = smem_kv.cfg.num_head_dim_stages_v + 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 - @schedule - def mma_schedule( + def mma_schedule_body( gqkv: GmemQKVResource, sq: SmemQResource, - skv: SmemKVResource, + sk: SmemKVResource, + sv: SmemKVResource, sp0: TmemSPResource, to: TmemOResource, vd0: TmemStatsDoneResource, wq: WorkQueue | None = None, ) -> None: desc_q0_base, _desc_q1_base = sq.create_function_variables() - desc_k_base, desc_v_base = skv.create_function_variables() + # Every instance declares both descriptor slots; a split buffer + # only ever fills the half matching its role. + desc_k_base, _desc_v_base = sk.create_function_variables() + _desc_k_base, desc_v_base = sv.create_function_variables() sp0.create_function_variables() to.create_function_variables() vd0.create_function_variables() @@ -945,15 +1140,15 @@ def mma_schedule( vd0.acquire() sp0.acquire() for head_dim_stage_idx in range(num_head_dim_stages_k): - skv.wait() - desc_k_base = skv.k_desc() + 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=head_dim_stage_idx, ) - skv.release() + sk.release() sp0.commit() with d.first_iter(): vd0.commit() @@ -961,31 +1156,58 @@ def mma_schedule( sp0.acquire() sp0.p_read() for head_dim_stage_idx in range(num_head_dim_stages_v): - skv.wait() + sv.wait() if cutlass.const_expr(smem_q.cfg.needs_paged_v_tail_clear): - desc_v_base = skv.v_desc_paged( + 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 = skv.v_desc() + desc_v_base = sv.v_desc() to.pv_mma( desc_v_base=desc_v_base, section=FmhaStage.Loop, head_dim_stage_idx=head_dim_stage_idx, ) - skv.release() + sv.release() to.commit() sq.release() sp0.commit() + @schedule + def mma_schedule( + gqkv: GmemQKVResource, + sq: SmemQResource, + skv: SmemKVResource, + sp0: TmemSPResource, + to: TmemOResource, + vd0: TmemStatsDoneResource, + wq: WorkQueue | None = None, + ) -> None: + """Shared-buffer captured schedule.""" + mma_schedule_body(gqkv, sq, skv, skv, sp0, to, vd0, wq) + + @schedule + def mma_split_schedule( + gqkv: GmemQKVResource, + sq: SmemQResource, + sk: SmemKVResource, + sv: SmemKVResource, + sp0: TmemSPResource, + to: TmemOResource, + vd0: TmemStatsDoneResource, + wq: WorkQueue | None = None, + ) -> None: + """Split K/V captured schedule.""" + mma_schedule_body(gqkv, sq, sk, sv, sp0, to, vd0, wq) + captured_schedule = _schedule_with_work_queue( - mma_schedule, + mma_split_schedule if split_kv else mma_schedule, gmem_qkv, smem_q, - smem_kv, + *kv_resources, tmem_sp0, tmem_o, tmem_vec_done_0, @@ -1005,21 +1227,93 @@ def mma_schedule( if tmem_sp1 is None or tmem_vec_done_1 is None: raise ValueError("paired MMA scheduling requires peer-1 resources") - @schedule - def mma_schedule( + # Padded QK is at most 256 for the admitted paired geometries. Keep each + # 128-wide slice in its own TS binding so later slices cannot replace it. + # The stage loops below handle both the full and partial final slice. + if num_head_dim_stages_k > 2: + raise ValueError("paired QK staging supports at most two 128-wide slices") + + def qk_mma_slice(sp, desc_q_base, desc_k, section, head_dim_stage_idx): + """Route a slice through its own descriptor binding to the shared QK body.""" + if head_dim_stage_idx == 0: + sp.qk_mma( + desc_q_base=desc_q_base, + desc_k_base=desc_k, + section=section, + head_dim_stage_idx=head_dim_stage_idx, + ) + else: + sp.qk_mma_stage( + desc_q_base=desc_q_base, + desc_k_stage=desc_k, + section=section, + head_dim_stage_idx=head_dim_stage_idx, + ) + + def load_k_stage(sk, head_dim_stage_idx): + """Retain the descriptor of each waited K slice for both query tiles.""" + sk.wait() + if head_dim_stage_idx == 0: + return sk.k_desc() + return sk.k_stage_desc() + + 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 = [] + for head_dim_stage_idx in range(num_head_dim_stages_k): + desc_k = ( + load_k_stage(sk, head_dim_stage_idx) + if descriptors is None + else descriptors[head_dim_stage_idx] + ) + qk_mma_slice(sp, desc_q_base, desc_k, section, head_dim_stage_idx) + stages.append(desc_k) + return stages + + def mma_schedule_body( gqkv: GmemQKVResource, sq: SmemQResource, - skv: SmemKVResource, + sk: SmemKVResource, + sv: SmemKVResource, sp0: TmemSPResource, sp1: TmemSPResource, to: TmemOResource, vd0: TmemStatsDoneResource, vd1: TmemStatsDoneResource, wq: WorkQueue | None = None, + pr0: TmemPPrefixReadyResource | None = None, + pr1: TmemPPrefixReadyResource | None = None, + pb0: SmemPResource | None = None, + pb1: SmemPResource | None = None, ) -> None: - """Captured schedule for interleaved QK and PV MMA work.""" + """Interleave paired QK/PV while retaining each K slice for both Qs.""" + + def pv( + sp: TmemSPResource, + pr: TmemPPrefixReadyResource | None, + pb: SmemPResource | None, + *, + group: int, + **kw: Any, + ) -> None: + """PV for query group ``group``. P-ready comes from the SMEM P tile with + P in SMEM, from the P-prefix barrier with the half overlap, else from the SP acquire.""" + if p_in_smem: + pb.wait() + to.pv_mma(inst_idx=group, **kw) + pb.release() + return + if pv_half_overlap: + pr.wait() + to.pv_mma(k_half=0, **kw) + pr.release() + sp.acquire() + sp.p_read() + to.pv_mma(k_half=1 if pv_half_overlap else None, **kw) + sq.init_descriptor_state() - skv.init_descriptor_state() + sk.init_descriptor_state() + sv.init_descriptor_state() sp0.init_mma_state() sp1.init_mma_state() to.init_mma_state() @@ -1049,131 +1343,316 @@ def mma_schedule( # Consume Q0, K0, then QK(Q0,K0)→S0. sq.wait() desc_q0_base = sq.q0_desc(inst_idx=0) - skv.wait() - desc_k_base = skv.k_desc() - if not smem_q.cfg.stats_via_smem: - vd0.acquire() - sp0.acquire() - sp0.qk_mma( - desc_q_base=desc_q0_base, - desc_k_base=desc_k_base, - section=FmhaStage.Head, - ) - sp0.commit() - if not smem_q.cfg.stats_via_smem: - vd0.commit() - # Consume Q1, then QK(Q1,K0)→S1. - sq.wait() - desc_q1_base = sq.q1_desc(inst_idx=1) - if not smem_q.cfg.stats_via_smem: - vd1.acquire() - sp1.acquire() - sp1.qk_mma( - desc_q_base=desc_q1_base, - desc_k_base=desc_k_base, - section=FmhaStage.Head, - ) + for head_dim_stage_idx in range(num_head_dim_stages_k): + desc_k = load_k_stage(sk, head_dim_stage_idx) + if head_dim_stage_idx == 0: + if not smem_q.cfg.stats_via_smem: + vd0.acquire() + sp0.acquire() + qk_mma_slice( + sp0, desc_q0_base, desc_k, FmhaStage.Head, head_dim_stage_idx + ) + if head_dim_stage_idx == num_head_dim_stages_k - 1: + sp0.commit() + if not smem_q.cfg.stats_via_smem: + vd0.commit() + if head_dim_stage_idx == 0: + sq.wait() + desc_q1_base = sq.q1_desc(inst_idx=1) + if not smem_q.cfg.stats_via_smem: + vd1.acquire() + sp1.acquire() + qk_mma_slice( + sp1, desc_q1_base, desc_k, FmhaStage.Head, head_dim_stage_idx + ) + if head_dim_stage_idx > 0: + sk.release() sp1.commit() if not smem_q.cfg.stats_via_smem: vd1.commit() # Q0/Q1 stay live because UMMA reads Q throughout the K-loop. # Release K0 (done with QK→S0 and QK→S1), then consume V0. - skv.release() - skv.wait() + sk.release() + # TODO: sv.wait() (V0) only depends on V's own TMA-load barrier, + # not on sk/sp0/sp1, so it could move earlier to overlap with the + # QK0/QK1 work above. + sv.wait() if cutlass.const_expr(smem_q.cfg.needs_paged_v_tail_clear): - desc_v_base = skv.v_desc_paged( + desc_v_base = sv.v_desc_paged( section=FmhaStage.Head, seqlen_k=v_seqlen_k, kv_tile_start=v_kv_tile_start, ) else: - desc_v_base = skv.v_desc() + desc_v_base = sv.v_desc() # Acquire O first (off critical path), then acquire SP0 and run PV→O0. - to.acquire() - sp0.acquire() - sp0.p_read() - to.pv_mma(desc_v_base=desc_v_base, section=FmhaStage.Head) - to.commit() - - # LOOP: interleave QK and PV work while preserving the previous V - # tile until its PV MMA has consumed it: - # QK0(deferred commit) -> PV1(V_prev, release V_prev) -> - # QK1(commit) -> release Ki+1 -> wait Vi+1 -> PV0(no commit) - with domain_loop(loop_start, loop_end, loop_step): - skv.wait() - desc_k_base = skv.k_desc() - # QK0: QK(Q0,Ki+1) → S0 (no acquire; handle held from PV0). - sp0.qk_mma( - desc_q_base=desc_q0_base, - desc_k_base=desc_k_base, - section=FmhaStage.Loop, - ) - sp0.commit() - # PV1(V_prev): P1 * V_prev → O1. + if not p_in_smem: to.acquire() - sp1.acquire() - sp1.p_read() - to.pv_mma(desc_v_base=desc_v_base, section=FmhaStage.Loop) - to.commit() - # Release V_prev after PV1 UMMA consumed SMEM data. - skv.release() - # QK1: QK(Q1,Ki+1) → S1 (no acquire; handle held from PV1). - sp1.qk_mma( - desc_q_base=desc_q1_base, - desc_k_base=desc_k_base, - section=FmhaStage.Loop, + pv( + sp0, + pr0, + pb0, + group=0, + desc_v_base=desc_v_base, + section=FmhaStage.Head, ) - sp1.commit() - # Release Ki+1, then wait Vi+1. - skv.release() - skv.wait() + to.commit() + + def wait_next_v() -> Any: + """Release Ki+1, then wait Vi+1 and build its descriptor.""" + for _ in range(num_head_dim_stages_k): + sk.release() + # TODO: sv.wait() (Vi+1) could move earlier in this iteration + # to overlap with QK0/PV1/QK1. + sv.wait() if cutlass.const_expr(smem_q.cfg.needs_paged_v_tail_clear): - desc_v_base = skv.v_desc_paged( + return sv.v_desc_paged( section=FmhaStage.Loop, tile_offset=1, seqlen_k=v_seqlen_k, kv_tile_start=v_kv_tile_start, ) - else: - desc_v_base = skv.v_desc() - # PV0: P0 * Vi+1 → O0. + return sv.v_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), + # QK1(i+1), PV1(i). Each QK acquires its SP slot here and each PV + # 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() + k_descriptors = qk_mma_stages(sk, sp0, desc_q0_base, FmhaStage.Loop) + sp0.commit() + to.acquire() + pv( + sp0, + pr0, + pb0, + group=0, + desc_v_base=desc_v_base, + section=FmhaStage.Loop, + ) + to.commit() + sp1.acquire() + qk_mma_stages( + sk, sp1, desc_q1_base, FmhaStage.Loop, descriptors=k_descriptors + ) + sp1.commit() + to.acquire() + pv( + sp1, + pr1, + pb1, + group=1, + desc_v_base=desc_v_base, + section=FmhaStage.Loop, + ) + to.commit() + sv.release() + desc_v_base = wait_next_v() + sq.release() + sq.release() to.acquire() - sp0.acquire() - sp0.p_read() - to.pv_mma( + pv( + sp0, + pr0, + pb0, + group=0, + desc_v_base=desc_v_base, + section=FmhaStage.Tail, + is_tail=True, + ) + to.commit() + to.acquire() + pv( + sp1, + pr1, + pb1, + group=1, + desc_v_base=desc_v_base, + section=FmhaStage.Tail, + is_tail=True, + ) + to.commit() + sv.release() + else: + # QK0 -> PV1 -> QK1 -> PV0: retain each K slice and the previous V + # until both peer query tiles have consumed the corresponding data. + with domain_loop(loop_start, loop_end, loop_step): + k_descriptors = qk_mma_stages(sk, sp0, desc_q0_base, FmhaStage.Loop) + sp0.commit() + to.acquire() + pv( + sp1, + pr1, + pb1, + group=1, + desc_v_base=desc_v_base, + section=FmhaStage.Loop, + ) + to.commit() + # Release V_prev after PV1 UMMA consumed SMEM data. + sv.release() + qk_mma_stages( + sk, sp1, desc_q1_base, FmhaStage.Loop, descriptors=k_descriptors + ) + sp1.commit() + desc_v_base = wait_next_v() + # PV0: P0 * Vi+1 → O0. + to.acquire() + pv( + sp0, + pr0, + pb0, + group=0, + desc_v_base=desc_v_base, + section=FmhaStage.Loop, + inst_idx=1, + ) + to.commit() + + sq.release() + sq.release() + sp0.commit() + to.acquire() + pv( + sp1, + pr1, + pb1, + group=1, desc_v_base=desc_v_base, - section=FmhaStage.Loop, - inst_idx=1, + section=FmhaStage.Tail, + is_tail=True, ) to.commit() + sv.release() + sp1.commit() - # TAIL: release Qs, close the deferred SP state, and run the final - # PV→O1 MMA. - sq.release() - sq.release() - sp0.commit() - to.acquire() - sp1.acquire() - sp1.p_read() - to.pv_mma( - desc_v_base=desc_v_base, - section=FmhaStage.Tail, - is_tail=True, - ) - to.commit() - skv.release() - sp1.commit() + @schedule + def mma_schedule( + gqkv: GmemQKVResource, + sq: SmemQResource, + skv: SmemKVResource, + sp0: TmemSPResource, + sp1: TmemSPResource, + to: TmemOResource, + vd0: TmemStatsDoneResource, + vd1: TmemStatsDoneResource, + wq: WorkQueue | None = None, + ) -> None: + """Shared-buffer captured schedule.""" + mma_schedule_body(gqkv, sq, skv, skv, sp0, sp1, to, vd0, vd1, wq) + @schedule + def mma_split_schedule( + gqkv: GmemQKVResource, + sq: SmemQResource, + sk: SmemKVResource, + sv: SmemKVResource, + sp0: TmemSPResource, + sp1: TmemSPResource, + to: TmemOResource, + vd0: TmemStatsDoneResource, + vd1: TmemStatsDoneResource, + wq: WorkQueue | None = None, + ) -> None: + """Split K/V captured schedule.""" + mma_schedule_body(gqkv, sq, sk, sv, sp0, sp1, to, vd0, vd1, wq) + + # Overlap variants carry the two P-prefix barriers. + @schedule + def mma_schedule_pr( + gqkv: GmemQKVResource, + sq: SmemQResource, + skv: SmemKVResource, + sp0: TmemSPResource, + sp1: TmemSPResource, + to: TmemOResource, + vd0: TmemStatsDoneResource, + vd1: TmemStatsDoneResource, + pr0: TmemPPrefixReadyResource, + pr1: TmemPPrefixReadyResource, + wq: WorkQueue | None = None, + ) -> None: + """Shared-buffer captured schedule with P-prefix overlap.""" + mma_schedule_body(gqkv, sq, skv, skv, sp0, sp1, to, vd0, vd1, wq, pr0, pr1) + + @schedule + def mma_split_schedule_pr( + gqkv: GmemQKVResource, + sq: SmemQResource, + sk: SmemKVResource, + sv: SmemKVResource, + sp0: TmemSPResource, + sp1: TmemSPResource, + to: TmemOResource, + vd0: TmemStatsDoneResource, + vd1: TmemStatsDoneResource, + pr0: TmemPPrefixReadyResource, + pr1: TmemPPrefixReadyResource, + wq: WorkQueue | None = None, + ) -> None: + """Split K/V captured schedule with P-prefix overlap.""" + mma_schedule_body(gqkv, sq, sk, sv, sp0, sp1, to, vd0, vd1, wq, pr0, pr1) + + # P-in-SMEM variants carry the two SMEM P tiles. + @schedule + def mma_split_schedule_pb( + gqkv: GmemQKVResource, + sq: SmemQResource, + sk: SmemKVResource, + sv: SmemKVResource, + sp0: TmemSPResource, + sp1: TmemSPResource, + to: TmemOResource, + vd0: TmemStatsDoneResource, + vd1: TmemStatsDoneResource, + pb0: SmemPResource, + pb1: SmemPResource, + wq: WorkQueue | None = None, + ) -> None: + """Split K/V captured schedule with P staged in SMEM.""" + mma_schedule_body( + gqkv, sq, sk, sv, sp0, sp1, to, vd0, vd1, wq, None, None, pb0, pb1 + ) + + @schedule + def mma_schedule_pb( + gqkv: GmemQKVResource, + sq: SmemQResource, + skv: SmemKVResource, + sp0: TmemSPResource, + sp1: TmemSPResource, + to: TmemOResource, + vd0: TmemStatsDoneResource, + vd1: TmemStatsDoneResource, + pb0: SmemPResource, + pb1: SmemPResource, + wq: WorkQueue | None = None, + ) -> None: + """Shared-buffer captured schedule with P staged in SMEM.""" + mma_schedule_body( + gqkv, sq, skv, skv, sp0, sp1, to, vd0, vd1, wq, None, None, pb0, pb1 + ) + + if 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 + else: + selected_mma_schedule = mma_split_schedule if split_kv else mma_schedule captured_schedule = _schedule_with_work_queue( - mma_schedule, + selected_mma_schedule, gmem_qkv, smem_q, - smem_kv, + *kv_resources, tmem_sp0, tmem_sp1, tmem_o, tmem_vec_done_0, tmem_vec_done_1, + *p_prefix_resources, + *smem_p_resources, work_queue=work_queue, ) return task_class( @@ -1197,6 +1676,8 @@ def create_softmax_task( s0s1_seq: S0S1SequenceResource | None, work_queue: WorkQueue | None, task_class: type[Task] = Task, + tmem_p_prefix_ready: TmemPPrefixReadyResource | None = None, + smem_p: SmemPResource | None = None, **task_kwargs: Any, ) -> Task: """Create a four-warp Softmax task. @@ -1280,6 +1761,12 @@ def softmax_schedule( ) else: old_row_max, row_max = sp.compute_row_max(row_max=row_max) + if tmem_sp.cfg.corr_skip_threshold_log2 > 0: + row_max = sp.freeze_row_max( + old_row_max=old_row_max, + row_max=row_max, + scale_softmax_log2=scale_softmax_log2, + ) # Stats(i): S(i) -> row max/sum for correction. vec.store_vec( old_row_max=old_row_max, @@ -1684,8 +2171,8 @@ def softmax_schedule( **task_kwargs, ) - # Paired QKV instances use separate SP resources. S0S1SequenceResource - # orders their P stores, so this path does not need the TMEM P handoff. + # Paired QKV instances use separate SP resources to protect P readiness. + # The sequence token paces peer progress; FP8 overlaps P computation. if s0s1_seq is not None and index == 1: src = _src_resources(tmem_sp, s0s1_seq, work_queue=work_queue) else: @@ -1693,13 +2180,24 @@ def softmax_schedule( dst = [tmem_vec] if s0s1_seq is not None and index == 0: dst.append(s0s1_seq) - - @schedule - def softmax_schedule( + pv_half_overlap = tmem_sp.cfg.pv_half_overlap + if pv_half_overlap: + if tmem_p_prefix_ready is None: + raise ValueError("pv_half_overlap requires the P-prefix barrier") + dst.append(tmem_p_prefix_ready) + p_in_smem = tmem_sp.cfg.p_in_smem + if p_in_smem: + if smem_p is None: + raise ValueError("p_in_smem requires the group's SMEM P resource") + dst.append(smem_p) + + def softmax_schedule_body( sp: TmemSPResource, vec: TmemStatsResource, seq: S0S1SequenceResource, wq: WorkQueue | None = None, + pr: TmemPPrefixReadyResource | None = None, + pb: SmemPResource | None = None, ) -> None: """Captured schedule for one softmax warp group.""" if tmem_sp.enable_early_tile_sum: @@ -1709,6 +2207,30 @@ def softmax_schedule( else: 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: + """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.""" + 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. + p_chunk = sp.exp2_p_smem( + row_max=row_max, scale_softmax_log2=scale_softmax_log2 + ) + pb.acquire() + sp.store_p() + pb.commit() + return p_chunk + if not pv_half_overlap: + return sp.exp2_p(row_max=row_max, scale_softmax_log2=scale_softmax_log2) + pr.acquire() + p_lo = sp.exp2_p_lo(row_max=row_max, scale_softmax_log2=scale_softmax_log2) + pr.commit() + return sp.exp2_p_hi( + row_max=row_max, scale_softmax_log2=scale_softmax_log2, p_lo=p_lo + ) + vec.init_store_state() with _work_tile_schedule_loop(wq, skip_if=skip_work_tile_if): # Recompute per-tile SP/Vec TMEM state. @@ -1749,10 +2271,6 @@ def softmax_schedule( row_max=row_max, q_offset=q_offset, ) - elif tmem_sp.uses_fixed_dense_k_tail_mask: - old_row_max, row_max = sp.fixed_dense_k_tail_masked_row_max( - row_max=row_max, - ) elif tmem_sp.uses_packed_dense_k_mask: old_row_max, row_max = sp.packed_dense_k_masked_row_max( row_max=row_max, @@ -1761,12 +2279,21 @@ def softmax_schedule( ) else: old_row_max, row_max = sp.compute_row_max(row_max=row_max) + if tmem_sp.cfg.corr_skip_threshold_log2 > 0: + row_max = sp.freeze_row_max( + old_row_max=old_row_max, + row_max=row_max, + scale_softmax_log2=scale_softmax_log2, + ) vec.store_vec( old_row_max=old_row_max, row_max=row_max, row_sum=row_sum, ) vec.commit() + if p_in_smem: + # S is in registers. Free the stage so QK(i+1) overlaps exp2. + sp.release() if s0s1_seq is None: pass elif index == 0: @@ -1775,18 +2302,30 @@ def softmax_schedule( else: # Softmax1 is the S0-S1 consumer: wait/release sequence. seq.wait() + # FP8 returns the pacing token before P work. Its SP release + # still follows every P store's completion. + early_token = tmem_sp.cfg.uses_d128_fp8_softmax_cadence or ( + p_in_smem and tmem_sp.cfg.fp8_psmem_early_token + ) + if early_token: + if index == 0: + seq.commit() + else: + seq.release() # Apply softmax and write P. - p_chunk = sp.exp2_p( + p_chunk = exp2_p( + sp, row_max=row_max, scale_softmax_log2=scale_softmax_log2, ) - if s0s1_seq is None: + if s0s1_seq is None or early_token: pass elif index == 0: seq.commit() else: seq.release() - sp.release() + if not p_in_smem: + sp.release() # Reduction. row_sum = sp.softmax_aux_reduce( old_row_max=old_row_max, @@ -1864,15 +2403,27 @@ def softmax_schedule( row_sum=row_sum, ) vec.commit() + if p_in_smem: + sp.release() if s0s1_seq is not None: seq.acquire() - p_chunk = sp.masked_exp2_p( - row_max=row_max, - scale_softmax_log2=scale_softmax_log2, - ) + if p_in_smem: + p_chunk = sp.masked_exp2_p_smem( + row_max=row_max, + scale_softmax_log2=scale_softmax_log2, + ) + pb.acquire() + sp.store_p() + pb.commit() + else: + p_chunk = sp.masked_exp2_p( + row_max=row_max, + scale_softmax_log2=scale_softmax_log2, + ) if s0s1_seq is not None: seq.commit() - sp.release() + if not p_in_smem: + sp.release() row_sum = sp.softmax_aux_reduce( old_row_max=old_row_max, row_max=row_max, @@ -1893,11 +2444,16 @@ def softmax_schedule( if s0s1_seq is not None: seq.acquire() sp.invalid_exp2_p(row_max=row_max) + if p_in_smem: + # MMA skips this PV but still waits on the P tile. + pb.acquire() + pb.commit() if s0s1_seq is not None: seq.commit() sp.release() - sp.wait() - sp.release() + if not p_in_smem: + sp.wait() + sp.release() old_row_max = sp.softmax_aux_identity(row_max=row_max) vec.acquire() vec.store_vec( @@ -1908,7 +2464,7 @@ def softmax_schedule( ) vec.commit() elif tmem_sp.cfg.is_causal: - # Causal softmax1 TAIL handles masked rows and cleanup. + # Causal tail: mask the last tile, then publish the final stats. sp.wait() old_row_max, row_max = sp.masked_row_max( row_max=row_max, @@ -1920,15 +2476,80 @@ def softmax_schedule( row_sum=row_sum, ) vec.commit() + if p_in_smem: + sp.release() if s0s1_seq is not None: seq.wait() - p_chunk = sp.masked_exp2_p( + if p_in_smem: + p_chunk = sp.masked_exp2_p_smem( + row_max=row_max, + scale_softmax_log2=scale_softmax_log2, + ) + pb.acquire() + sp.store_p() + pb.commit() + else: + p_chunk = sp.masked_exp2_p( + row_max=row_max, + scale_softmax_log2=scale_softmax_log2, + ) + if s0s1_seq is not None: + 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 s0s1_seq is not None: + if not p_in_smem: + sp.wait() + sp.release() + old_row_max = sp.softmax_aux_identity(row_max=row_max) + vec.acquire() + vec.store_vec( + old_row_max=old_row_max, + row_max=row_max, + row_sum=row_sum, + final_stats=True, + ) + vec.commit() + elif tmem_sp.uses_fixed_dense_k_tail_mask: + # Dense tail: one more loop step with the zero-filled lanes masked out. + sp.wait() + 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, + row_sum=row_sum, + ) + vec.commit() + if p_in_smem: + sp.release() + if s0s1_seq is None: + pass + elif index == 0: + seq.acquire() + else: + seq.wait() + p_chunk = exp2_p( + sp, + row_max=row_max, + scale_softmax_log2=scale_softmax_log2, + ) + if s0s1_seq is None: + pass + elif index == 0: + seq.commit() + else: seq.release() - sp.release() + if not p_in_smem: + sp.release() row_sum = sp.softmax_aux_reduce( old_row_max=old_row_max, row_max=row_max, @@ -1936,10 +2557,12 @@ def softmax_schedule( p_chunk=p_chunk, scale_softmax_log2=scale_softmax_log2, ) - sp.wait() - sp.release() - old_row_max = sp.softmax_aux_identity(row_max=row_max) 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, @@ -1948,10 +2571,10 @@ def softmax_schedule( ) vec.commit() else: - # Non-causal TAIL commits the reserved stats slot and lets MMA - # complete its cleanup path. - sp.wait() - sp.release() + # 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, @@ -1961,9 +2584,57 @@ def softmax_schedule( ) vec.commit() - captured_schedule = _schedule_with_work_queue( - softmax_schedule, tmem_sp, tmem_vec, s0s1_seq, work_queue=work_queue - ) + @schedule + def softmax_schedule( + sp: TmemSPResource, + vec: TmemStatsResource, + seq: S0S1SequenceResource, + wq: WorkQueue | None = None, + ) -> None: + softmax_schedule_body(sp, vec, seq, wq) + + @schedule + def softmax_schedule_pr( + sp: TmemSPResource, + vec: TmemStatsResource, + seq: S0S1SequenceResource, + pr: TmemPPrefixReadyResource, + wq: WorkQueue | None = None, + ) -> None: + softmax_schedule_body(sp, vec, seq, wq, pr) + + @schedule + def softmax_schedule_pb( + sp: TmemSPResource, + vec: TmemStatsResource, + seq: S0S1SequenceResource, + pb: SmemPResource, + wq: WorkQueue | None = None, + ) -> None: + softmax_schedule_body(sp, vec, seq, wq, None, pb) + + if p_in_smem: + captured_schedule = _schedule_with_work_queue( + softmax_schedule_pb, + tmem_sp, + tmem_vec, + s0s1_seq, + smem_p, + work_queue=work_queue, + ) + elif pv_half_overlap: + captured_schedule = _schedule_with_work_queue( + softmax_schedule_pr, + tmem_sp, + tmem_vec, + s0s1_seq, + tmem_p_prefix_ready, + work_queue=work_queue, + ) + else: + captured_schedule = _schedule_with_work_queue( + softmax_schedule, tmem_sp, tmem_vec, s0s1_seq, work_queue=work_queue + ) return task_class( src_resources=src, dst_resources=dst, @@ -2117,6 +2788,9 @@ def correction_schedule( def _create_paired_task() -> Task: if tmem_vec1 is None or smem_o_1 is None or tmem_vec_done_1 is None: raise ValueError("paired correction scheduling requires peer-1 resources") + release_stats_after_read = tmem_vec0.cfg.p_in_smem + if release_stats_after_read and not tmem_vec0.cfg.stats_via_smem: + raise ValueError("early stats release requires SMEM-staged stats") src = _src_resources( tmem_vec0, @@ -2163,14 +2837,19 @@ def correction_schedule( v0.wait() v0.release() v1.wait() + if release_stats_after_read: + v1.release() # The correction loop consumes vec/O pairs in alternating order so - # each half can unblock the other half's next producer. + # each half can unblock the other half's next producer. With P in + # SMEM each stats slot is released right after read_vec instead. with domain_loop(loop_start, loop_end, loop_step): # Part 1: consume TmemStats0 + O0, release TmemStats1. v0.wait() vec_old_max, vec_new_max, _, vec_scale = v0.read_vec( scale_softmax_log2=scale_softmax_log2_v0, ) + if release_stats_after_read: + v0.release() to.wait() to.correct( vec_old_max=vec_old_max, @@ -2178,13 +2857,16 @@ def correction_schedule( vec_scale=vec_scale, inst_idx=0, ) - v1.release() + if not release_stats_after_read: + v1.release() to.release() # Part 2: consume TmemStats1 + O1, release TmemStats0. v1.wait() vec_old_max, vec_new_max, _, vec_scale = v1.read_vec( scale_softmax_log2=scale_softmax_log2_v1, ) + if release_stats_after_read: + v1.release() to.wait() to.correct( vec_old_max=vec_old_max, @@ -2192,11 +2874,12 @@ def correction_schedule( vec_scale=vec_scale, inst_idx=1, ) - v0.release() + if not release_stats_after_read: + v0.release() to.release() - # TAIL: consume remaining stats, release tmem-stats-done gates, and - # stage corrected O0/O1 into SMEM for the epilogue task. - v1.release() + # Tail: read the final stats, then write corrected O0/O1 to smem. + if not release_stats_after_read: + v1.release() v0.wait() _, _, vec_row_sum, vec_scale = v0.read_vec( scale_softmax_log2=scale_softmax_log2_v0, @@ -2207,13 +2890,15 @@ def correction_schedule( vd0.release() v0.release() to.wait() - so0.acquire() - so0.store_o( - vec_row_sum=vec_row_sum, - vec_scale=vec_scale, - output_scale=output_scale0, - ) - so0.commit() + for head_dim_stage_idx in range(smem_o_0.cfg.num_o_head_dim_stages): + so0.acquire() + so0.store_o( + vec_row_sum=vec_row_sum, + vec_scale=vec_scale, + output_scale=output_scale0, + head_dim_stage_idx=head_dim_stage_idx, + ) + so0.commit() to.release() v1.wait() _, _, vec_row_sum, vec_scale = v1.read_vec( @@ -2225,13 +2910,15 @@ def correction_schedule( vd1.release() v1.release() to.wait() - so1.acquire() - so1.store_o( - vec_row_sum=vec_row_sum, - vec_scale=vec_scale, - output_scale=output_scale1, - ) - so1.commit() + for head_dim_stage_idx in range(smem_o_0.cfg.num_o_head_dim_stages): + so1.acquire() + so1.store_o( + vec_row_sum=vec_row_sum, + vec_scale=vec_scale, + output_scale=output_scale1, + head_dim_stage_idx=head_dim_stage_idx, + ) + so1.commit() to.release() captured_schedule = _schedule_with_work_queue( @@ -2344,23 +3031,27 @@ def epilogue_schedule( with domain_loop(loop_start, loop_end, loop_step): pass # Store the first corrected O tile through gmem_o_0. - so0.wait() - head_coord, batch_coord, seq_coord_q = so0.compute_output_coords() - go0.tma_store( - head_coord=head_coord, - batch_coord=batch_coord, - seq_coord_q=seq_coord_q, - ) - so0.release() + for head_dim_stage_idx in range(smem_o_0.cfg.num_o_head_dim_stages): + so0.wait() + head_coord, batch_coord, seq_coord_q = so0.compute_output_coords() + go0.tma_store( + head_coord=head_coord, + batch_coord=batch_coord, + seq_coord_q=seq_coord_q, + head_dim_stage_idx=head_dim_stage_idx, + ) + so0.release() # Store the second corrected O tile through gmem_o_1. - so1.wait() - head_coord, batch_coord, seq_coord_q = so1.compute_output_coords() - go1.tma_store( - head_coord=head_coord, - batch_coord=batch_coord, - seq_coord_q=seq_coord_q, - ) - so1.release() + for head_dim_stage_idx in range(smem_o_0.cfg.num_o_head_dim_stages): + so1.wait() + head_coord, batch_coord, seq_coord_q = so1.compute_output_coords() + go1.tma_store( + head_coord=head_coord, + batch_coord=batch_coord, + seq_coord_q=seq_coord_q, + head_dim_stage_idx=head_dim_stage_idx, + ) + so1.release() captured_schedule = _schedule_with_work_queue( epilogue_schedule, 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..a66a6385baa7 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/helpers.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_context/helpers.py @@ -17,7 +17,7 @@ import cutlass import cutlass.cute as cute -from cutlass import Int32, Int64 +from cutlass import Float32, Int32, Int64 @cute.jit @@ -94,3 +94,24 @@ def freeze_smem_descriptor(desc): write_only_types=[Int64], read_only_args=[desc], ) + + +_LDTM32_MAX_PTX = ( + "tcgen05.ld.red.sync.aligned.32x32b.x32.f32.max {" + + ", ".join(f"{{$w{i}}}" for i in range(32)) + + "}, {$w32}, [{$r0}];" +) + + +@cute.jit +def load_tmem_32x32b_max(tmem_addr: Int32) -> tuple: + """Load 32 FP32 scores and their maximum on SM103/SM107 (PTX 8.8). + + Compatibility wrapper for CUTLASS DSL 4.7. The caller gates the GPU + architecture and completes the TMEM loads before publishing/reusing data. + """ + return cute.arch.inline_ptx( + _LDTM32_MAX_PTX, + write_only_types=[Float32] * 33, + read_only_args=[tmem_addr], + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/direct_sparse_metadata.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/direct_sparse_metadata.py new file mode 100644 index 000000000000..15d6f2710295 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/direct_sparse_metadata.py @@ -0,0 +1,201 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 by FlashInfer team. +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Lazy sparse metadata views shared by attention and split reduction.""" + +import cutlass +import cutlass.cute as cute +from cutlass import Int32, Int64 + + +class HeadIndexedMetadataView: + """Read one KV head's values from interleaved [route, pattern_head] rows.""" + + def __init__(self, values, heads, head): + self.values = (values, heads, head) + + def __extract_mlir_values__(self): + return cutlass.extract_mlir_values(self.values) + + def __new_from_mlir_values__(self, values): + return HeadIndexedMetadataView( + *cutlass.new_from_mlir_values(self.values, values) + ) + + @cute.jit + def __getitem__(self, route): + values, heads, head = self.values + return values[Int64(route) * Int64(heads) + Int64(head)] + + +class DirectSparseMetadataView: + """Lazy sparse metadata resolved from raw inputs for one query per route. + + The combined G1 path skips metadata construction: attention resolves block + IDs, causal tails, physical locators and lengths directly from the inputs. + G1 means one token per route, not one token per request; packed and fixed + requests may each contain several routes. Split reduction can still run. + """ + + def __init__( + self, + inputs, + qo_indptr, + *, + packed, + model_len, + page_size, + page_capacity, + lengths, + sparse_block_size=4, + fragment_size=4, + head=0, + ): + self.values = (*inputs, qo_indptr, head) + 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.fragment_size = fragment_size + self.fragments_per_block = sparse_block_size // fragment_size + + def __extract_mlir_values__(self): + return cutlass.extract_mlir_values(self.values) + + 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, + page_size=self.page_size, + page_capacity=self.page_capacity, + lengths=self.lengths, + sparse_block_size=self.sparse_block_size, + fragment_size=self.fragment_size, + head=new[5], + ) + + def with_head(self, head): + return DirectSparseMetadataView( + self.values[:4], + self.values[4], + packed=self.packed, + model_len=self.model_len, + page_size=self.page_size, + page_capacity=self.page_capacity, + lengths=self.lengths, + sparse_block_size=self.sparse_block_size, + fragment_size=self.fragment_size, + head=head, + ) + + @cute.jit + def resolve_route(self, route): + blocks, table, requests, positions, offsets, _ = self.values + row = route + valid = (row >= Int32(0)) & (row < Int32(blocks.shape[0])) + if cutlass.const_expr(self.packed): + row = Int32(offsets[route]) + end = Int32(offsets[route + Int32(1)]) + valid = ( + (row >= Int32(0)) + & (row < Int32(blocks.shape[0])) + & (end == row + Int32(1)) + ) + request = Int32(-1) + visible = Int32(0) + if valid: + request = Int32(requests[row]) + position = Int64(positions[row]) + valid = ( + (request >= Int32(0)) + & (request < Int32(table.shape[0])) + & (position >= Int64(0)) + & (position < Int64(self.model_len)) + ) + if valid: + visible = Int32(position + Int64(1)) + return row, request, visible + + @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]) + ) + logical = Int32(-1) + candidate = slot // Int32(self.fragments_per_block) + if candidate < complete: + if cutlass.const_expr(len(blocks.shape) == 3): + logical = Int32(blocks[row, self.values[5], candidate]) + else: + logical = Int32(blocks[row, candidate]) + return self.map_page(route_state, slot, logical) + + @cute.jit + def map_page(self, route_state, slot, logical): + blocks, table, _, _, _, _ = self.values + _, request, visible = route_state + complete = cute.math.min( + visible // Int32(self.sparse_block_size), Int32(blocks.shape[-1]) + ) + 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) + ): + logical = Int32(-1) + elif (candidate == complete) & ( + slot % Int32(self.fragments_per_block) * Int32(self.fragment_size) < tail + ): + logical = visible // Int32(self.sparse_block_size) + else: + logical = Int32(-1) + if logical >= Int32(0): + fragment = logical * Int32(self.fragments_per_block) + slot % Int32( + self.fragments_per_block + ) + fragments_per_page = Int32(self.page_size // self.fragment_size) + storage_page = fragment // fragments_per_page + subpage = fragment % fragments_per_page + if storage_page < Int32(table.shape[1]): + physical = Int32(table[request, storage_page]) + if physical >= Int32(0): + result = physical * fragments_per_page + subpage + return result + + @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), + Int32(self.values[0].shape[-1]), + ) + tail = visible % Int32(self.sparse_block_size) + result = cute.math.max( + complete * Int32(self.sparse_block_size) + tail, Int32(1) + ) + else: + route = Int32(index // Int64(self.page_capacity)) + slot = Int32(index % Int64(self.page_capacity)) + result = self.load_page(self.resolve_route(route), slot) + return result 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 0626410d1f10..b0efb7adc248 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 @@ -22,12 +22,16 @@ from collections.abc import Mapping, Sequence from copy import deepcopy from dataclasses import dataclass, fields, replace +from functools import wraps +import math import cutlass.utils as utils -from cutlass import BFloat16, Float16, Float32, Float8E4M3FN +from cutlass import BFloat16, Float16, Float32, Float4E2M1FN, Float8E4M3FN, Int8 +from ...sage 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, ) @@ -40,10 +44,9 @@ FP32_BYTES, FP8_OUTPUT_ELEMENTS_PER_REG_GROUP, FP8_P_PACKED_REGS_PER_Q_REPEAT, - FP8_VALUES_PER_REG, FP16_OUTPUT_ELEMENTS_PER_REG_GROUP, FP16_P_PACKED_REGS_PER_Q_REPEAT, - FP16_VALUES_PER_REG, + KV_TILE_256_BYTE_WIDE_SHARED_FIFO_STAGES, KV_TILE_256_REGISTER_REALLOCATION_MIN_TILES, KV_TILE_256_SHARED_FIFO_STAGES, MAX_CLUSTER_DIM_X, @@ -51,23 +54,29 @@ MAX_KV_STAGE_SMEM_KIB, MAX_WARP_GROUPS, MIN_LOOP_ITERS_PER_SPLIT, + MMA_K_STEP_BYTES, + PACKED_REGISTER_BYTES, PARALLEL_REDUCTION_BYTES_PER_SLICE, PARALLEL_REDUCTION_THREADS_PER_CTA, PARTIAL_O_ELEMENT_BYTES, PARTIAL_STATS_VALUES_PER_ROW, + Q_TOKEN_KV_BLOCK_SPARSE_HELD_LOCATOR_MAX_TILES, + Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIP_BITS, Q_REPETITION_GROUP_HEADS, Q_ROW_ALIGNMENT_BYTES, REDUCTION_BYTES_PER_SLICE, REDUCTION_THREADS_PER_CTA, SPLIT_KV_MIN_TOKENS_PER_CTA, + SWIZZLE_128B_ROW_BYTES, TMEM_COLUMNS_PER_ROW, + TMEM_MAX_ALLOCATION_COLUMNS, TMEM_ROW_STRIDE, TOTAL_SMEM_BUDGET_KIB, WARP_THREADS, ) ConfigValue = int | float | bool | str | type | None -_GroupedKeepsProfileKey = tuple[type, type, type, int, int, int, int] +_GroupedKeepsProfileKey = tuple[type, type, type, type, int, int, int, int] # Public APIs use the strings ``dense`` and ``causal``. Keep the value carried # through FmhaDecodeConfig as a small integer so mask selection remains a @@ -84,23 +93,49 @@ Float16, Float16, Float16, + Float16, 128, 0, 2, 2, ) -# Block-sparse feature compatibility is validated by -# ``validate_block_sparse_profile``. This set only selects the Keeps MMA -# resource recipes qualified for that already-validated launch domain. -_BLOCK_SPARSE_GROUPED_KEEPS_PROFILES = { +# Keeps MMA resource recipes qualified for contiguous K/V, dense or +# block-sparse. Block-sparse feature compatibility is validated by +# ``validate_block_sparse_profile``. +_CONTIGUOUS_GROUPED_KEEPS_PROFILES = { _GROUPED_KEEPS_MAIN_PROFILE, - (BFloat16, BFloat16, BFloat16, 128, 0, 2, 2), + (BFloat16, BFloat16, BFloat16, BFloat16, 128, 0, 2, 2), +} +# 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. +_SAGE_GROUPED_KEEPS_PROFILES = { + (qk_dtype, qk_dtype, Float8E4M3FN, out_dtype, 128, 0, 2, 2) + for qk_dtype in (Float8E4M3FN, Int8) + for out_dtype in (Float16, BFloat16) +} +# (Q, K, output) dtype triples of the Sage recipes; V is E4M3 in every one. +_SAGE_DTYPE_RECIPES = { + (profile[0], profile[1], profile[3]) for profile in _SAGE_GROUPED_KEEPS_PROFILES } _GROUPED_KEEPS_STATIC_ONLY_PROFILES = { - (Float8E4M3FN, Float8E4M3FN, Float16, 128, 0, 2, 2), - (BFloat16, BFloat16, BFloat16, 64, 0, 2, 2), - (Float16, Float16, Float16, 256, 128, 1, 1), + (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, Float16, 128, 0, 2, 2), + (BFloat16, BFloat16, BFloat16, BFloat16, 64, 0, 2, 2), + (Float16, Float16, Float16, Float16, 256, 128, 1, 1), } +# This profile qualifies the Keeps MMA resource recipe for +# ``supports_grouped_keeps``. Mixed k_dtype != v_dtype itself is validated by +# ``_validate_mixed_kv_dtype_profile``. +_GROUPED_KEEPS_MIXED_KV_DTYPE_PROFILE: _GroupedKeepsProfileKey = ( + BFloat16, + BFloat16, + Float8E4M3FN, + BFloat16, + 128, + 0, + 2, + 2, +) # Per-thread register budgets for the Q64/KV256 warp groups once the launch # bound enables ``setmaxnreg``. The MMA, load, and scheduler warps keep 56, so @@ -112,6 +147,28 @@ KV_TILE_256_SOFTMAX_TASK_REGISTERS = 152 KV_TILE_256_CORRECTION_TASK_REGISTERS = 152 +# Sparse route unions reuse the dense Keeps staging: D64/D128 use two +# unstaged K/V instances; D256 uses one instance with D128 head bands. +_Q_TOKEN_KV_BLOCK_SPARSE_GROUPED_KEEPS_PROFILES = { + ( + dtype, + dtype, + dtype, + output, + dim, + 128 if dim == 256 else 0, + 1 if dim == 256 else 2, + 1 if dim == 256 else 2, + ) + for dim in (64, 128, 256) + for dtype, output in ( + (Float16, Float16), + (BFloat16, BFloat16), + (Float8E4M3FN, Float16), + (Float8E4M3FN, BFloat16), + ) +} + _KV_TILE_256_PHYSICAL_DEFAULTS: Mapping[str, ConfigValue] = { "tmem_s_cols": 128, "tmem_stats_cols": 32, @@ -122,7 +179,6 @@ "mma_tile_m_bmm2": 64, "mma_tile_n_bmm2": 256, "q_stages": 1, - "kv_stages": KV_TILE_256_SHARED_FIFO_STAGES, "head_dim_per_stage_kv": 0, "num_insts_kv": 2, "o_stages": 2, @@ -148,18 +204,32 @@ "scheduler_num_warps": 1, } -_KV_TILE_256_TUNABLE_FIELDS = frozenset(("kv_stages",)) -# Public cost-model collection uses the FP8 proxy for every source dtype. The -# original fixed-Q1 ratio-32 requests exercise a partial grouped-Q tile; the -# shape-aware path also admits complete fixed multi-Q tiles at any legal head -# ratio. These are profile families rather than shape exceptions: batch size, -# KV length, tile choice, and legal GMEM split fanout remain unrestricted by -# this declaration. +def _dtype_bytes(dtype: type) -> int: + """Return the storage bytes of one CUTLASS element; packed NVFP4 counts one.""" + return max(dtype.width // BITS_PER_BYTE, 1) + + +# ``(Q, K/V, O)`` dtype recipes that qualify the Q64/KV256 profile without Sage +# attention. KV128 Keeps profiles qualify on the full profile keys above. +_KV256_DTYPE_RECIPES = { + (Float16, Float16, Float16), + (BFloat16, BFloat16, BFloat16), + (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN), + (Float8E4M3FN, Float8E4M3FN, Float16), +} + + +# Public cost-model collection uses the FP8 proxy for every source dtype. The +# fixed-Q1 path admits every head ratio covered by its Q64/Q128 tile, while the +# shape-aware path also admits complete fixed multi-Q tiles. These are profile +# families rather than shape exceptions: batch size, KV length, tile choice, +# and legal GMEM split fanout remain unrestricted by this declaration. _GROUPED_KEEPS_PAGED_FP8_PROFILES = { - (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, 64, 0, 2, 2), - (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, 128, 0, 2, 2), - (Float8E4M3FN, Float8E4M3FN, Float16, 256, 128, 1, 1), + (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, 64, 0, 2, 2), + (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, 128, 0, 2, 2), + (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, 256, 128, 1, 1), + (Float8E4M3FN, Float8E4M3FN, Float8E4M3FN, Float16, 256, 128, 1, 1), } @@ -330,7 +400,7 @@ class GroupedQLaunchCandidate: def enumerate_grouped_q_mma_candidates( - *, heads_q_per_kv: int, seq_len_q: int + *, heads_q_per_kv: int, seq_len_q: int, swaps_only: bool = False ) -> tuple[GroupedQMmaCandidate, ...]: """Return grouped-Q candidates without applying launch policy. @@ -338,7 +408,8 @@ def enumerate_grouped_q_mma_candidates( leave structural padding rows and the final CTA may own fewer tokens than its capacity; the common Q geometry and row masks already represent both cases. This helper deliberately does not inspect dtypes, reduction modes, - scheduler state, or mutate ``FmhaDecodeConfig``. + scheduler state, or mutate ``FmhaDecodeConfig``. ``swaps_only`` excludes + KeepsMmaAb tiles for profiles such as transformed mixed-precision K/V. """ if heads_q_per_kv <= 0: raise ValueError("heads_q_per_kv must be positive") @@ -347,6 +418,8 @@ def enumerate_grouped_q_mma_candidates( candidates = [] for variant, tile_size_q in _GROUPED_Q_MMA_TILES: + if swaps_only and variant != "swaps_mma_ab": + continue if tile_size_q < heads_q_per_kv: continue q_tokens_per_cta = tile_size_q // heads_q_per_kv @@ -482,6 +555,26 @@ def make_q_tile_geometry( ) +def _accept_legacy_dtype_aliases(cls): + """Keep legacy constructor keywords outside the canonical dataclass fields.""" + init = cls.__init__ + + @wraps(init) + def init_with_dtype_aliases(self, *args, qkv_dtype=None, kv_dtype=None, **kwargs): + if qkv_dtype is not None: + kwargs.setdefault("q_dtype", qkv_dtype) + kwargs.setdefault("k_dtype", kv_dtype or qkv_dtype) + kwargs.setdefault("v_dtype", kv_dtype or qkv_dtype) + elif kv_dtype is not None: + kwargs.setdefault("k_dtype", kv_dtype) + kwargs.setdefault("v_dtype", kv_dtype) + init(self, *args, **kwargs) + + cls.__init__ = init_with_dtype_aliases + return cls + + +@_accept_legacy_dtype_aliases @dataclass class FmhaDecodeConfig: """ @@ -505,6 +598,17 @@ 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 + # 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; + # Q1 uses the same scattered page route without membership masking. + use_q_token_kv_block_sparse_route: bool = False + # None preserves each sparse family's default in raw/static configs too. + share_pattern_across_kv_heads: bool | None = None + # Allow the attention grid to acquire a programmatic launch dependency. + # Callers must opt in only when a producer was launched immediately before + # attention on the same stream and will release that dependency. + use_pdl: bool = False heads_q_per_kv: int = 0 groups_tokens_heads_q: bool = False # K/V tokens per tile along the K-sequence dimension; also the MMA "M" @@ -538,18 +642,19 @@ class FmhaDecodeConfig: # ------------------------------------------------------------------ # Data types # ------------------------------------------------------------------ - # Q element type. One of Float16 / BFloat16 / Float8E4M3FN. Must equal - # kv_dtype — mixed Q/KV element types are not supported yet; enforced by - # the guard in make_decode_config. + # Q element type. One of Float16 / BFloat16 / Float8E4M3FN, or Int8 with + # Sage attention. q_dtype: type = Float16 - # K and V element type. One of Float16 / BFloat16 / Float8E4M3FN. - kv_dtype: type = Float16 + # K/V element types. K may be Int8 with Sage attention. Raw K/V storage may + # also use packed Float4E2M1FN (NVFP4), which the transform task converts + # to q_dtype before MMA consumption. + k_dtype: type = Float16 + v_dtype: type = Float16 # Output O element type. One of Float16 / BFloat16 / Float8E4M3FN. out_dtype: type = Float16 # Accumulator type (BMM accumulators and softmax stats), always Float32 # in the currently supported recipes. acc_dtype: type = Float32 - # ------------------------------------------------------------------ # Software pipeline depths # ------------------------------------------------------------------ @@ -559,6 +664,13 @@ class FmhaDecodeConfig: # K/V TMA pipeline depth. Deeper KV staging hides the long # GMEM→SMEM latency in the BMM1↔BMM2 chain. kv_stages: int = 4 + # Transformed K/V pipeline depth for mixed Q/KV modes. Raw K/V is loaded + # into SmemKv, converted to the Q-side MMA input type, then consumed by MMA. + # Zero selects the profile-derived stage count. + transformed_kv_stages: int = 0 + # Store the transformed NVFP4 K/V operand ring in TMEM. The supported + # FP8-Q Swaps profile defaults to this path; other mixed modes use SMEM. + store_transformed_kv_in_tmem: bool = False # ------------------------------------------------------------------ # TMEM column counts @@ -591,7 +703,8 @@ class FmhaDecodeConfig: mma_tile_n_bmm2: int = 8 # ------------------------------------------------------------------ - # Warp specialization layout (4 warp groups × 4 warps = 16 warps total) + # Warp specialization layout (normally 4 warp groups / 16 warps; a static + # sparse-attention throughput profile may add two producer-only groups) # NOTE: please update `_active_warp_roles` after new roles are added. # ------------------------------------------------------------------ # Softmax0Task: WG0 (warps 0–3) handles even K/V instances (K0/V0). @@ -607,11 +720,11 @@ class FmhaDecodeConfig: # MmaTask: a single warp in WG3 issues tcgen05 MMA instructions for BMM1/BMM2. mma_warp_idx: int = 12 # WG3: warp 12 mma_num_warps: int = 1 - # LoadTask: a single warp in WG3 issues TMA loads for Q/K/V. + # LoadTask: one or more independently configured warps issue Q/K/V TMA. load_warp_idx: int = 13 # WG3: warp 13 load_num_warps: int = 1 - # PageTableTask (paged-KV only): warp 14 prefetches logical→physical - # page IDs that LoadTask consumes when issuing the TMA copies. + # PageTableTask (paged-KV only): prefetch logical→physical page IDs for + # LoadTask. Its warp count is independent of num_insts_kv and TMA issuers. page_offsets_warp_idx: int = 14 # WG3: warp 14 for paged-KV page table prefetch page_offsets_num_warps: int = 1 # SMEM pipeline depth for the prefetched page-offset table. @@ -623,10 +736,14 @@ class FmhaDecodeConfig: # ClcLoadTask: under persistent scheduling, warp 15 issues the CLC # response loads. clc_load_warp_idx: int = 15 + # TransformKvTask (mixed Q/KV): converts raw K/V into the Q-side MMA input + # type. Finalization preserves its configured role order while compacting. + transform_kv_warp_idx: int = 14 + transform_kv_num_warps: int = 2 # PaddingTask placement is derived after selecting the active task roles. # Each active warp group is compacted first, then its unused tail warps are # assigned to the corresponding padding task. Persistent layouts retain - # all four warp groups even when the last group contains only padding. + # all four original warp groups even when the last contains only padding. wg0_padding_warp_idx: int = 4 wg0_padding_num_warps: int = 0 wg1_padding_warp_idx: int = 8 @@ -635,6 +752,10 @@ class FmhaDecodeConfig: wg2_padding_num_warps: int = 0 wg3_padding_warp_idx: int = 16 wg3_padding_num_warps: int = 0 + wg4_padding_warp_idx: int = 20 + wg4_padding_num_warps: int = 0 + wg5_padding_warp_idx: int = 24 + wg5_padding_num_warps: int = 0 # ------------------------------------------------------------------ # Task-local register allocation @@ -642,33 +763,159 @@ class FmhaDecodeConfig: # The full TileQ128 Keeps graph has the largest live softmax/correction # fragments and needs registers moved from its descriptor-only task group. # A long KV256 graph instead moves a smaller share from Softmax to its - # heavier correction tail. Short KV256 loops avoid the fixed hand-off cost. + # heavier correction tail. Mixed Q/KV graphs also reserve registers for + # TransformKvTask. Short KV256 loops avoid the fixed hand-off cost. + @property + def effective_storage_tokens_per_page(self) -> int: + """Return the physical cache-page extent used by the TensorMap.""" + return self.storage_tokens_per_page or self.num_tokens_per_page + + @property + def has_storage_subpages(self) -> bool: + """Whether one physical cache page contains multiple semantic pages.""" + return self.effective_storage_tokens_per_page != self.num_tokens_per_page + + @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 + + @property + def shares_sparse_pattern(self) -> bool: + """Q-block routes default to per-head patterns; token routes share.""" + if self.share_pattern_across_kv_heads is None: + return not self.use_block_sparse + return self.share_pattern_across_kv_heads + + @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 + + @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 + + @property + def grouped_q_fits_tile(self) -> bool: + """Whether one complete query group fits in the physical Q tile.""" + return ( + self.groups_tokens_heads_q + and self.grouped_q_rows > 0 + and self.grouped_q_rows <= self.tile_size_q + ) + + @property + def uses_q_token_kv_block_sparse_page_route(self) -> bool: + """Whether this grouped-Q route consumes sparse-block locators.""" + return ( + self.use_paged_kv + and not self.use_block_sparse + and self.uses_scattered_page_route + 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 self.grouped_q_fits_tile + ) + + @property + def uses_held_encoded_locator_window(self) -> bool: + """Whether encoded paging keeps the complete CTA route in SMEM.""" + + return ( + self.use_paged_kv + and self.uses_scattered_page_route + and self.tile_size_kv == 128 + and not self.use_sliding_window_causal + and self.static_local_kv_tiles + <= Q_TOKEN_KV_BLOCK_SPARSE_HELD_LOCATOR_MAX_TILES + ) + @property def uses_task_register_reallocation(self) -> bool: - return self.use_keeps_mma_ab and ( - self.tile_size_q == 128 + if ( + self.use_transform_kv + or (self.use_keeps_mma_ab and self.tile_size_q == 128) + # Persistent grouped masks keep additional routing/softmax state + # live. Reuse the same CTA-pool budgets to give it register headroom. or ( - self.tile_size_q == 64 + self.use_keeps_mma_ab + and self.use_persistent_scheduler + and self.uses_q_token_kv_block_sparse_page_membership + ) + or ( + 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 ) - ) + ): + return True + else: + return False @property def softmax_task_num_registers(self) -> int | None: if not self.uses_task_register_reallocation: return None - return KV_TILE_256_SOFTMAX_TASK_REGISTERS if self.tile_size_kv == 256 else 184 + if self.use_transform_kv: + preferred = 168 if self.headdim == 64 and self.num_insts_kv != 1 else 184 + else: + preferred = ( + KV_TILE_256_SOFTMAX_TASK_REGISTERS if self.tile_size_kv == 256 else 184 + ) + # setmaxnreg redistributes the CTA's initial register pool, not the + # whole SM register file. Complete 8-register/thread quanta leave + # unused SM registers when a CTA has more than 16 warps. + total_warps = self.threads_per_cta // WARP_THREADS + initial_regs = 65536 // (self.threads_per_cta * 8) * 8 + softmax_warps = self.softmax0_num_warps + ( + self.softmax1_num_warps if self.num_insts_kv != 1 else 0 + ) + other_warps = total_warps - softmax_warps - self.correction_num_warps + correction_regs = self.correction_task_num_registers + producer_regs = self.mma_load_task_num_registers + assert correction_regs is not None and producer_regs is not None + available = ( + total_warps * initial_regs + - self.correction_num_warps * correction_regs + - other_warps * producer_regs + ) + if self.use_transform_kv: + # Transform warps can retain more registers than descriptor-only + # producers; reserve their actual budget before growing Softmax. + transform_regs = self.transform_kv_task_num_registers + assert transform_regs is not None + available -= self.transform_kv_num_warps * (transform_regs - producer_regs) + return min(preferred, available // (softmax_warps * 8) * 8) @property def correction_task_num_registers(self) -> int | None: if not self.uses_task_register_reallocation: return None + if self.use_transform_kv: + return 88 return KV_TILE_256_CORRECTION_TASK_REGISTERS if self.tile_size_kv == 256 else 88 @property def mma_load_task_num_registers(self) -> int | None: - return 56 if self.uses_task_register_reallocation else None + if not self.uses_task_register_reallocation: + return None + if self.use_transform_kv and self.headdim == 64 and self.num_insts_kv != 1: + return 88 + return 56 + + @property + def transform_kv_task_num_registers(self) -> int | None: + if not self.use_transform_kv: + return None + if self.num_insts_kv == 1 or self.store_transformed_kv_in_tmem: + return 184 + if self.headdim == 64: + return 88 + return 56 # ------------------------------------------------------------------ # SMEM allocation alignment @@ -733,8 +980,73 @@ def q_manual_padding_rows(self) -> int: @property def smem_kv_tile_bytes(self) -> int: - """SMEM bytes for one staged K or V tile.""" - return self.tile_size_kv * self.head_dim_kv_stage * self.kv_dtype_bytes + """TMA bytes for a shared K/V stage; unequal dtypes use separate rings.""" + return max(self.smem_k_tile_bytes, self.smem_v_tile_bytes) + + @property + def smem_k_tile_bytes(self) -> int: + """TMA transaction bytes for one raw K tile, before NVFP4 expansion.""" + bits = 4 if self.k_dtype == Float4E2M1FN else self.k_dtype_bytes * 8 + return self.tile_size_kv * self.head_dim_kv_stage * bits // 8 + + @property + def smem_v_tile_bytes(self) -> int: + """TMA transaction bytes for one raw V tile, before NVFP4 expansion.""" + bits = 4 if self.v_dtype == Float4E2M1FN else self.v_dtype_bytes * 8 + return self.tile_size_kv * self.head_dim_kv_stage * bits // 8 + + @property + def smem_k_storage_tile_bytes(self) -> int: + """SMEM allocation bytes for one raw K tile.""" + return self.smem_k_tile_elements * self.k_dtype_bytes + + @property + def smem_v_storage_tile_bytes(self) -> int: + """SMEM allocation bytes for one raw V tile.""" + return self.smem_v_tile_elements * self.v_dtype_bytes + + @property + def smem_kv_storage_tile_bytes(self) -> int: + """SMEM bytes for storing one K or V stage.""" + return self.smem_kv_tile_elements * self.kv_dtype_bytes + + @property + def smem_kv_sf_bytes_per_token(self) -> int: + """SMEM bytes for one staged K or V of a token.""" + if self.use_nvfp4_kv: + sf_bytes_per_token = self.head_dim_kv_stage // 16 + if self.num_head_dim_stages_kv > 1: + # A split head-dimension stage cannot fold scale bytes from + # adjacent tokens because its slice is not contiguous in the + # source tensor. Keep its inner box at TMA's 16-byte minimum. + sf_bytes_per_token = max(sf_bytes_per_token, 16) + return sf_bytes_per_token + return 0 + + @property + def smem_kv_sf_tile_bytes(self) -> int: + """SMEM bytes for one staged K or V NVFP4 scale-factor tile.""" + return self.tile_size_kv * self.smem_kv_sf_bytes_per_token + + @property + def sf_tma_reshape_factor(self) -> int: + """Token-fold factor that widens the NVFP4 SF TMA inner box to 128 B.""" + sf_per_token = self.headdim // 16 + return min(max(1, 128 // sf_per_token), self.num_tokens_per_page) + + @property + def smem_transformed_kv_tile_bytes(self) -> int: + """SMEM bytes for one transformed K or V tile consumed by MMA.""" + return self.tile_size_kv * self.head_dim_kv_stage * self.q_dtype_bytes + + @property + def tmem_transformed_kv_stage_cols(self) -> int: + """TMEM columns occupied by one transformed K or V stage.""" + if not self.store_transformed_kv_in_tmem: + return 0 + return ( + self.tile_size_kv * self.head_dim_kv_stage * self.q_dtype_bytes // (128 * 4) + ) @property def head_dim_kv_stage(self) -> int: @@ -802,34 +1114,40 @@ def pv_head_dim_stage_tmem_offset(self, head_dim_stage_idx: int) -> int: @property def smem_p_tile_bytes(self) -> int: """P stored in SMEM for the SwapsMmaAb BMM2 B operand.""" - return self.tile_size_kv * self.tile_size_q * self.q_dtype_bytes + return self.tile_size_kv * self.tile_size_q * self.pv_dtype_bytes @property def tmem_total_cols(self) -> int: - """Sum of TMEM columns used by the kernel: 2× S, 2× softmax stats, - and o_stages× O accumulators (the factor 2 is the two softmax - instances K0/V0 and K1/V1).""" + """Sum of TMEM columns used by the kernel: (S + softmax) * num_insts_kv + + o_stages * O accumulators + transformed K/V stages (if any).""" + total_cols = 0 if self.use_keeps_mma_ab: if self.num_insts_kv == 1: - return ( + total_cols = ( self.tmem_s_cols + self.tmem_stats_cols + self.tmem_p_cols + self.tmem_o_stage_cols * self.o_stages ) - return ( - 2 * self.tmem_s_cols - + ( - 2 * self.tmem_stats_cols - if self.keeps_separates_tmem_s_and_stats - else 0 + else: + total_cols = ( + 2 * self.tmem_s_cols + + ( + 2 * self.tmem_stats_cols + if self.keeps_separates_tmem_s_and_stats + else 0 + ) + + self.tmem_o_stage_cols * self.o_stages ) + else: + total_cols = ( + self.num_insts_kv * self.tmem_s_cols + + self.num_insts_kv * self.tmem_stats_cols + self.tmem_o_stage_cols * self.o_stages ) return ( - 2 * self.tmem_s_cols - + 2 * self.tmem_stats_cols - + self.tmem_o_stage_cols * self.o_stages + total_cols + + self.tmem_transformed_kv_stage_cols * self.transformed_kv_stages ) @property @@ -850,25 +1168,70 @@ def threads_per_cta(self) -> int: ) # ------------------------------------------------------------------ - # Inferred dtype attributes (derived from q_dtype / kv_dtype / out_dtype) + # Inferred dtype attributes (derived from q_dtype / k_dtype / v_dtype / out_dtype) # ------------------------------------------------------------------ @property def q_dtype_bytes(self) -> int: - """Byte width of one Q element (fp16/bf16=2, e4m3=1). - - Also the byte width used by anything that feeds the MMA on the Q/S/P - side (softmax stats, P tile, MMA operand descriptors).""" - return 1 if self.q_dtype == Float8E4M3FN else 2 + """Byte width of one Q element (fp16/bf16=2, e4m3/int8=1).""" + return _dtype_bytes(self.q_dtype) @property def kv_dtype_bytes(self) -> int: - """Byte width of one K/V element (fp16/bf16=2, e4m3=1).""" - return 1 if self.kv_dtype == Float8E4M3FN else 2 + """Integral byte width used for raw K/V storage bookkeeping.""" + return max(self.k_dtype_bytes, self.v_dtype_bytes) + + @property + def k_dtype_bytes(self) -> int: + """Raw K storage unit in bytes (fp16/bf16=2, e4m3/int8=1, packed NVFP4=1).""" + return _dtype_bytes(self.k_dtype) + + @property + def v_dtype_bytes(self) -> int: + """Raw V storage unit in bytes (fp16/bf16=2, e4m3=1, packed NVFP4=1).""" + return _dtype_bytes(self.v_dtype) + + @property + def qk_dtype(self) -> type: + """Effective input precision for the QK MMA.""" + return self.q_dtype + + @property + def pv_dtype(self) -> type: + """Effective input precision for P and the transformed/raw V operand.""" + return self.q_dtype if self.use_transform_kv else self.v_dtype + + @property + def qk_dtype_bytes(self) -> int: + return self.q_dtype_bytes + + @property + def pv_dtype_bytes(self) -> int: + return 1 if self.use_fp8_pv else 2 + + @property + def use_fp8_qk(self) -> bool: + return self.qk_dtype == Float8E4M3FN + + @property + def use_fp8_pv(self) -> bool: + return self.pv_dtype == Float8E4M3FN + + @property + def use_bf16_pv(self) -> bool: + return self.pv_dtype == BFloat16 + + @property + def use_bf16_qkv(self) -> bool: + return self.q_dtype == self.k_dtype == self.v_dtype == BFloat16 + + @property + def use_fp8_qkv(self) -> bool: + return self.q_dtype == self.k_dtype == self.v_dtype == Float8E4M3FN @property def o_dtype_bytes(self) -> int: """Byte width of one O element (fp16/bf16=2, e4m3=1).""" - return 1 if self.out_dtype == Float8E4M3FN else 2 + return _dtype_bytes(self.out_dtype) @property def acc_dtype_bytes(self) -> int: @@ -876,9 +1239,9 @@ def acc_dtype_bytes(self) -> int: return 4 if self.acc_dtype == Float32 else 2 @property - def use_bf16_qkv(self) -> bool: - """Whether Q/K/V use BF16 storage and MMA inputs.""" - return self.kv_dtype == BFloat16 + def use_bf16_q(self) -> bool: + """Whether Q uses BF16 storage and MMA inputs.""" + return self.q_dtype == BFloat16 @property def use_bf16_output(self) -> bool: @@ -886,9 +1249,14 @@ def use_bf16_output(self) -> bool: return self.out_dtype == BFloat16 @property - def use_fp8_qkv(self) -> bool: - """fp8 (E4M3) Q/K/V path: switches MMA kind and P-quantization.""" - return self.kv_dtype == Float8E4M3FN + def use_fp8_q(self) -> bool: + """Whether Q uses FP8 storage and QK MMA inputs.""" + return self.q_dtype == Float8E4M3FN + + @property + def kv_swizzle_row_elems(self) -> int: + """Return the K/V elements of one 128-byte swizzle row: 64 or 128.""" + return SWIZZLE_128B_ROW_BYTES // self.kv_dtype_bytes @property def use_fp8_output(self) -> bool: @@ -911,15 +1279,45 @@ def supports_reduction_dtypes(self) -> bool: self.q_dtype in (Float16, BFloat16) and self.out_dtype in (Float16, BFloat16) ) or ( - self.q_dtype == Float8E4M3FN and self.out_dtype in (Float16, Float8E4M3FN) + self.q_dtype == Float8E4M3FN + and ( + self.out_dtype in (Float16, Float8E4M3FN) + or ( + self.uses_q_token_kv_block_sparse_page_route + and self.out_dtype == BFloat16 + ) + ) ) + @property + def kv_dtype(self) -> type | None: + """Legacy shared-K/V alias; independent storage has no common dtype.""" + return self.k_dtype if self.k_dtype == self.v_dtype else None + + @kv_dtype.setter + def kv_dtype(self, dtype: type) -> None: + self.k_dtype = self.v_dtype = dtype + + @property + def use_transform_kv(self) -> bool: + """Whether raw K/V must be transformed before MMA consumes it.""" + return self.k_dtype == self.v_dtype and self.q_dtype != self.k_dtype + + @property + def use_nvfp4_kv(self) -> bool: + """Whether the raw KV cache uses packed NVFP4 storage.""" + return self.k_dtype == self.v_dtype == Float4E2M1FN + + @property + def use_fp8_kv(self) -> bool: + """Whether the raw KV cache uses FP8 E4M3 storage.""" + return self.k_dtype == self.v_dtype == Float8E4M3FN + # ------------------------------------------------------------------ # Feature flags # ------------------------------------------------------------------ - # Enable persistent scheduling: work tiles are fetched from a CLC response - # queue at runtime instead of mapping one tile per CTA in the launch grid. - # Mutually exclusive with split-KV mode. + # Enable persistent scheduling: one physical worker may consume multiple + # logical work tiles. Mutually exclusive with split-KV mode. use_persistent_scheduler: bool = False # Split-KV: split the K-sequence across several CTAs that produce partial # O/stats, with a GMEM reduction epilogue. @@ -934,9 +1332,14 @@ def supports_reduction_dtypes(self) -> bool: # Paged-KV cache layout: K/V live in fixed-size pages and the kernel # follows a logical→physical page index table per request. use_paged_kv: bool = False - # Page size (tokens per page) when use_paged_kv is enabled. Must be one of - # 16 / 32 / 64 / 128 and must divide the 128-token KV tile. + # Semantic page size (tokens represented by one page-table entry) when + # use_paged_kv is enabled. It must divide the 128-token KV tile. num_tokens_per_page: int = 32 + # Physical tokens stored in one cache page. Zero selects the semantic page + # size. A larger value enables encoded subpage locators for a native + # page-four table: + # locator = physical_page * subpages_per_storage_page + subpage. + storage_tokens_per_page: int = 0 # Maximum number of pages per (batch, head_kv) — sizes the page index # table stride. max_num_pages_per_seq_kv: int = 1 @@ -968,6 +1371,14 @@ def supports_reduction_dtypes(self) -> bool: use_attention_sinks: bool = False # Optional profile and reduction knobs selected by launcher policy or tests. use_keeps_mma_ab: bool = False + # Sage attention: a nonzero ``sage_k_block_size`` enables per-block Q/K + # scales and per-channel V scales. ``sage_k_summary_block_size`` is the K + # block size of proxy summary scales and equals ``sage_k_block_size`` + # without proxy routes. ``sage_v_mean`` adds a per-channel V mean back. + sage_q_block_size: int = 0 + sage_k_block_size: int = 0 + sage_k_summary_block_size: int = 0 + sage_v_mean: bool = False # Nonzero means each K/V stage covers only this many head-dim columns. # H256 SwapsMmaAb uses 128-column stages to keep TMA and TMEM layouts valid. head_dim_per_stage_kv: int = 0 @@ -1057,9 +1468,29 @@ def parallel_reduction_bytes_per_slice(self) -> int: @property def smem_kv_tile_elements(self) -> int: - """Return K or V elements in one staged SMEM tile.""" + """Return physical K or V elements in one staged SMEM tile. + + When using nvfp4 KV with TmemTransformedKvResource, TMA uses ``B4X16_P64`` format. + It adds 4 padding bits to each 4-bit e2m1 value. + """ + if self.store_transformed_kv_in_tmem and self.use_nvfp4_kv: + return self.tile_size_kv * self.head_dim_kv_stage return self.smem_kv_tile_bytes // self.kv_dtype_bytes + @property + def smem_k_tile_elements(self) -> int: + """Return K elements in one staged SMEM tile.""" + if self.store_transformed_kv_in_tmem and self.k_dtype == Float4E2M1FN: + return self.tile_size_kv * self.head_dim_kv_stage + return self.smem_k_tile_bytes // self.k_dtype_bytes + + @property + def smem_v_tile_elements(self) -> int: + """Return V elements in one staged SMEM tile.""" + if self.store_transformed_kv_in_tmem and self.v_dtype == Float4E2M1FN: + return self.tile_size_kv * self.head_dim_kv_stage + return self.smem_v_tile_bytes // self.v_dtype_bytes + @property def num_softmax_scale_groups(self) -> int: """Return independent max/sum groups tracked by each softmax lane.""" @@ -1095,6 +1526,31 @@ def num_softmax_score_fragments(self) -> int: """Return score fragments used to cover one logical KV tile.""" return self.num_s_regs_per_thread // self.softmax_score_fragment_regs + @property + def p_values_per_reg(self) -> int: + """Return the probabilities packed into one 32-bit P register.""" + return PACKED_REGISTER_BYTES // self.pv_dtype_bytes + + @property + def fragment_p_packed_cols(self) -> int: + """Return the packed TMEM P columns published by one score fragment.""" + return self.softmax_score_fragment_regs // self.p_values_per_reg + + @property + def pv_mma_k_step(self) -> int: + """Return the K elements one PV MMA instruction consumes.""" + return MMA_K_STEP_BYTES // self.pv_dtype_bytes + + @property + def p_cols_per_mma_k_step(self) -> int: + """Return the packed TMEM P columns one PV MMA K step advances.""" + return MMA_K_STEP_BYTES // PACKED_REGISTER_BYTES + + @property + def pv_mma_steps_per_fragment(self) -> int: + """Return the PV MMA instructions consumed by one score fragment.""" + return self.softmax_score_fragment_regs // self.pv_mma_k_step + @property def block_sparse_kv_atom_size(self) -> int: """Return the K token span of one block-sparse route origin.""" @@ -1119,18 +1575,115 @@ def uses_ws_2x2_datapath(self) -> bool: """ return self.tile_size_kv == 256 + @property + def keeps_spatial_halves(self) -> int: + """Return the Keeps lane halves that own distinct KV tokens of one row. + + KV256's 2x2 datapath splits a row between threads ``[0, 64)`` (K64 + atoms 0 and 2) and ``[64, 128)`` (atoms 1 and 3); every other Keeps + profile has one half. + """ + return 2 if self.uses_ws_2x2_datapath else 1 + + def keeps_route_atom_owner(self, atom: int) -> tuple[int, int]: + """Return ``(half, position)`` of a route atom. + + Half ``h`` owns atoms ``h``, ``h + halves``, ... in order; + ``_keeps_route_atom`` in ``helpers_common`` is the inverse. + """ + halves = self.keeps_spatial_halves + return atom % halves, atom // halves + + @property + def keeps_fragments_per_atom(self) -> int: + """Return the K32 score fragments 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 + tokens. + """ + return 2 + + @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 + + @property + def proxy_summary_geometry(self) -> tuple[int, int]: + """Return ``(summary count, final summary's token mass)`` of proxy routes. + + One summary stands for each KV block; the final block may be ragged. + """ + return _block_sparse_proxy_summary_geometry( + self.static_seq_len_kv, self.kv_block_size + ) + + @property + def num_proxy_summaries(self) -> int: + """Return the length of the summary sequence proxy routes address.""" + num_summaries, _ = self.proxy_summary_geometry + return num_summaries + + @property + def proxy_log2_block_mass(self) -> float: + """Return ``log2`` of the tokens one full proxy summary stands for. + + A summary is its block's mean K/V, so it acts as that many identical + tokens: the mass enters its log2-domain logit additively. + """ + return math.log2(self.kv_block_size) + + @property + def proxy_tail_summary(self) -> tuple[int, float]: + """Return the final summary index and its ``log2(mass)`` shortfall. + + The shortfall is non-positive and zero for a sequence that ends on a + block boundary. + """ + num_summaries, tail_mass = self.proxy_summary_geometry + return num_summaries - 1, math.log2(tail_mass) - math.log2(self.kv_block_size) + + @property + def num_proxy_groups(self) -> int: + """Return the fixed summary groups that make up the proxy routes. + + The prepare kernel emits one proxy record per ``tile_size_kv`` + consecutive summaries, whatever blocks the row selects. + """ + return (self.num_proxy_summaries + self.tile_size_kv - 1) // self.tile_size_kv + + @property + def proxy_static_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 + position ``tail_summary_idx % tile_size_kv`` of the last group. Only + meaningful when ``proxy_tail_summary`` reports a non-zero shortfall. + """ + atom_size = self.block_sparse_kv_atom_size + fragment_regs = self.softmax_score_fragment_regs + assert atom_size == 2 * fragment_regs, ( + "the proxy tail location requires the K64 Keeps atom" + ) + tail_summary_idx, _ = self.proxy_tail_summary + in_group_idx = tail_summary_idx % self.tile_size_kv + half, position = self.keeps_route_atom_owner(in_group_idx // atom_size) + fragment_idx = ( + position * self.softmax_fragments_per_route_atom + + (in_group_idx % atom_size) // fragment_regs + ) + return half, fragment_idx + @property def num_packed_p_regs(self) -> int: """Return packed P registers stored by each softmax producer lane.""" if self.use_keeps_mma_ab: - values_per_reg = ( - FP8_VALUES_PER_REG if self.use_fp8_qkv else FP16_VALUES_PER_REG - ) - return max(self.num_s_regs_per_thread // values_per_reg, 1) + return max(self.num_s_regs_per_thread // self.p_values_per_reg, 1) q_repeats = max(self.tile_size_q // Q_REPETITION_GROUP_HEADS, 1) regs_per_repeat = ( FP8_P_PACKED_REGS_PER_Q_REPEAT - if self.use_fp8_qkv + if self.use_fp8_pv else FP16_P_PACKED_REGS_PER_Q_REPEAT ) return regs_per_repeat * q_repeats @@ -1139,12 +1692,12 @@ def num_packed_p_regs(self) -> int: def tmem_p_cols_per_inst(self) -> int: """Return the TMEM-P columns owned by one K/V instance. - Two-instance Keeps publishes one packed row per producer lane, with one + Two-instance Keeps stores one packed row per producer lane, with one packed 32-bit register slot per TMEM column. Therefore its footprint is - ``num_s_regs_per_thread / values_per_reg == num_packed_p_regs``, rather - than a function of the complete logical KV tile width. Q128/KV128 and - Q64/KV256 both own 128 packed 16-bit P values per lane and need 64 - columns per instance. + ``num_s_regs_per_thread / p_values_per_reg == num_packed_p_regs``, + rather than a function of the complete logical KV tile width. + Q128/KV128 and Q64/KV256 both own 128 P values per lane: 64 columns + per instance for 16-bit P and 32 for FP8 P. """ if self.uses_two_inst_tmem_p: return self.num_packed_p_regs @@ -1181,7 +1734,7 @@ def correction_barrier_threads(self) -> int: @property def keeps_p_smem_vector_elements(self) -> int: """Return P elements in one aligned 16-byte SMEM store.""" - return 16 // self.q_dtype_bytes + return 16 // self.pv_dtype_bytes @property def static_local_kv_tiles(self) -> int: @@ -1197,38 +1750,106 @@ def static_local_kv_tiles(self) -> int: @property def inferred_kv_stages(self) -> int: """Return the deepest K/V ring that fits the shared-memory budget.""" - q_dtype_bits = 8 if self.q_dtype == Float8E4M3FN else 16 + 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_ROW_ALIGNMENT_BYTES ) * Q_ROW_ALIGNMENT_BYTES - q_tile_kib = q_row_bytes * self.tile_size_q // BYTES_PER_KIB - kv_budget_kib = min( - MAX_KV_STAGE_SMEM_KIB, - TOTAL_SMEM_BUDGET_KIB - q_tile_kib * self.q_stages, + fixed_bytes = ( + q_row_bytes * self.tile_size_q * self.q_stages + + self.correction_sum_scratch_entries * FP32_BYTES ) - kv_stage_head_dim = self.head_dim_per_stage_kv or self.headdim - kv_tile_bits = q_dtype_bits * self.tile_size_kv * kv_stage_head_dim - return max( - 1, - kv_budget_kib * BYTES_PER_KIB * BITS_PER_BYTE // kv_tile_bits, + 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) + if self.use_paged_kv: + page_ids_per_stage = WARP_THREADS + if self.uses_held_encoded_locator_window: + page_ids_per_stage = max( + page_ids_per_stage, + self.static_local_kv_tiles + * self.tile_size_kv + // self.num_tokens_per_page, + ) + fixed_bytes += self.page_offsets_stages * page_ids_per_stage * 4 + if self.use_cluster_smem_reduction: + fixed_bytes += self.cluster_transaction_bytes + + # Separate K/V rings each contain half of kv_stages slots. Their + # average stage width accounts for the BF16-K/FP8-V profile exactly. + kv_tile_bits = (self.smem_k_tile_bytes + self.smem_v_tile_bytes) * 4 + if self.use_transform_kv: + # Raw storage and converted operands coexist. When the converted + # depth is still implicit, budget both rings per stage; otherwise + # reserve the configured converted ring before sizing raw storage. + kv_tile_bits = ( + self.smem_kv_storage_tile_bytes + self.smem_kv_sf_tile_bytes + ) * BITS_PER_BYTE + if not self.store_transformed_kv_in_tmem: + if self.transformed_kv_stages > 0: + fixed_bytes += ( + self.transformed_kv_stages * self.smem_transformed_kv_tile_bytes + ) + else: + kv_tile_bits += self.smem_transformed_kv_tile_bytes * BITS_PER_BYTE + + # Reserve common buffers for every profile, leaving hardware SMEM + # above this budget for alignment, barriers, and scheduler state. + kv_budget_bytes = min( + MAX_KV_STAGE_SMEM_KIB * BYTES_PER_KIB, + TOTAL_SMEM_BUDGET_KIB * BYTES_PER_KIB - fixed_bytes, ) + inferred_stages = max(1, kv_budget_bytes * BITS_PER_BYTE // kv_tile_bits) + return inferred_stages def validate_dtypes(self) -> None: - """Validate decode input, output, and accumulator dtypes.""" + """Validate homogeneous, transformed K/V, QK-BF16/PV-FP8 and Sage IO. + + Sage attention takes 8-bit Q/K of one dtype, E4M3 V and a 16-bit + output; Int8 Q/K exists only in that recipe. + """ for name, dtype, supported in ( - ("q_dtype", self.q_dtype, SUPPORTED_IO_DTYPES), - ("kv_dtype", self.kv_dtype, SUPPORTED_IO_DTYPES), + ("q_dtype", self.q_dtype, SUPPORTED_QK_DTYPES), + ("k_dtype", self.k_dtype, SUPPORTED_KV_DTYPES | SUPPORTED_QK_DTYPES), + ("v_dtype", self.v_dtype, SUPPORTED_KV_DTYPES), ("out_dtype", self.out_dtype, SUPPORTED_IO_DTYPES), ("acc_dtype", self.acc_dtype, SUPPORTED_ACC_DTYPES), ): if dtype not in supported: raise ValueError(f"Unsupported {name}: {dtype}") - if self.q_dtype != self.kv_dtype: - raise ValueError( - f"q_dtype ({self.q_dtype}) != kv_dtype ({self.kv_dtype}): " - "mixed Q/KV element types are not supported" - ) + if self.use_sage_attention: + if (self.q_dtype, self.k_dtype, self.out_dtype) not in _SAGE_DTYPE_RECIPES: + raise ValueError( + "Sage attention requires Float8E4M3FN or Int8 Q and K with " + "Float16 or BFloat16 output" + ) + if self.v_dtype != Float8E4M3FN: + raise ValueError( + "Sage attention requires Float8E4M3FN V, " + f"got v_dtype {self.v_dtype}" + ) + return + if Int8 in (self.q_dtype, self.k_dtype): + raise ValueError("Int8 Q/K requires Sage attention") + if self.q_dtype == self.k_dtype == self.v_dtype: + return + if self.k_dtype != self.v_dtype: + if ( + self.q_dtype == self.k_dtype == BFloat16 + and self.v_dtype == Float8E4M3FN + ): + return + raise ValueError("k_dtype != v_dtype is supported only for QK-BF16/PV-FP8") + if (self.q_dtype == BFloat16 and self.k_dtype == Float8E4M3FN) or ( + self.q_dtype in (BFloat16, Float8E4M3FN) and self.k_dtype == Float4E2M1FN + ): + return + raise ValueError( + f"q_dtype ({self.q_dtype}) != kv_dtype ({self.kv_dtype}): " + "only BF16 Q with E4M3 K/V or BF16/FP8 Q with NVFP4 K/V " + "mixed modes are supported" + ) def validate_boolean_fields(self) -> None: """Require every boolean config field to carry a real Python bool.""" @@ -1243,8 +1864,16 @@ def validate_paged_kv_staging_config(self) -> None: if not self.use_paged_kv: raise ValueError("paged-KV staging requires use_paged_kv=True") validate_page_size(self.num_tokens_per_page) + validate_storage_page_size( + self.num_tokens_per_page, + self.effective_storage_tokens_per_page, + ) if self.use_block_sparse: + if self.uses_scattered_page_route: + raise ValueError( + "scattered page routes are supported only by dense paged-KV" + ) if self.tile_size_kv not in (128, 256): raise ValueError( "paged block-sparse supports only KV128 or KV256 routes" @@ -1273,14 +1902,72 @@ def validate_paged_kv_staging_config(self) -> None: raise ValueError( "paged-KV num_tokens_per_page must divide tile_size_kv exactly" ) + if self.uses_scattered_page_route and self.tile_size_kv != 128: + raise ValueError("scattered page routes currently require tile_size_kv=128") pages_per_tile = self.tile_size_kv // self.num_tokens_per_page - if pages_per_tile not in (1, 2, 4, 8, 16): + if pages_per_tile not in (1, 2, 4, 8, 16, 32): raise ValueError( - "paged-KV staging supports 1, 2, 4, 8, or 16 pages per KV tile" + "paged-KV staging supports 1, 2, 4, 8, 16, or 32 pages per KV tile" ) - if self.page_offsets_num_warps != 1: + max_page_producers = 2 if self.uses_q_token_kv_block_sparse_page_route else 1 + if not 1 <= self.page_offsets_num_warps <= max_page_producers: raise ValueError( - "paged-KV page-offset staging requires exactly one producer warp" + "paged-KV page-offset staging requires one producer warp, " + "or up to two for QToken-KvBlock-Sparse-Attention" + ) + + def validate_sage_profile(self) -> None: + """Validate the Sage scale block sizes and tile profile. + + ``validate_dtypes`` owns the Sage dtype recipe. + """ + if not self.use_sage_attention: + if ( + self.sage_q_block_size != 0 + or self.sage_k_summary_block_size != 0 + or self.sage_v_mean + ): + raise ValueError( + "sage_q_block_size, sage_k_summary_block_size and sage_v_mean " + "require sage_k_block_size" + ) + return + if self.sage_k_block_size not in SAGE_K_BLOCK_SIZES: + raise ValueError( + f"sage_k_block_size must be one of {SAGE_K_BLOCK_SIZES}, " + f"got {self.sage_k_block_size}" + ) + if self.sage_k_summary_block_size not in SAGE_K_BLOCK_SIZES: + raise ValueError( + f"sage_k_summary_block_size must be one of {SAGE_K_BLOCK_SIZES}, " + 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: + 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}" + ) + if not self.streams_tmem_p_fragments: + raise ValueError( + "Sage attention requires a two-instance Keeps profile: " + "Q64/KV256 or Q128/KV128" + ) + if self.use_paged_kv or self.headdim != 128: + raise ValueError("Sage attention requires contiguous K/V with headdim=128") + # Only the direct output store applies the V channel scales and means; + # split-KV partials and reductions do not. + if ( + self.use_variable_seqlens_q + or self.use_sliding_window_causal + or self.use_attention_sinks + or self.use_split_kv + or self.use_cluster_smem_reduction + or self.use_separate_reduction_kernel + ): + raise ValueError( + "Sage attention supports the direct-output grid only, without " + "split-KV, variable-Q, sliding-window, or attention-sink features" ) def validate_block_sparse_profile(self, *, heads_q_per_kv: int) -> None: @@ -1296,12 +1983,16 @@ def validate_block_sparse_profile(self, *, heads_q_per_kv: int) -> None: ) return - if not (self.q_dtype == self.kv_dtype == self.out_dtype): - raise ValueError("block-sparse requires q_dtype == kv_dtype == out_dtype") - if self.q_dtype not in (Float16, BFloat16): - raise ValueError( - "block-sparse supports only matching Float16 or BFloat16 IO" - ) + # ``validate_dtypes`` owns the Sage dtype recipe. + if not self.use_sage_attention: + if not (self.q_dtype == self.k_dtype == self.v_dtype == self.out_dtype): + raise ValueError( + "block-sparse requires q_dtype == k_dtype == v_dtype == out_dtype" + ) + if self.q_dtype not in (Float16, BFloat16): + raise ValueError( + "block-sparse supports only matching Float16 or BFloat16 IO" + ) kv_block_size = _validate_sparse_kv_block_size(self.kv_block_size) selected_q_tile = _select_block_sparse_q_tile_size( @@ -1323,15 +2014,15 @@ def validate_block_sparse_profile(self, *, heads_q_per_kv: int) -> None: and not self.use_parallel_sparse_kv_loads ): raise ValueError( - "block-sparse tile_size_kv=256 requires the Q64 16-bit Keeps " - "profile with coarse KV blocks and one load task" + "block-sparse tile_size_kv=256 requires the Q64 Keeps profile " + "with coarse KV blocks and one load task" ) if self.use_keeps_mma_ab and not self.streams_tmem_p_fragments: # The block-sparse Keeps softmax and P passes exist only in their # streamed K32-fragment form. raise ValueError( "block-sparse KeepsMmaAb requires a streamed TMEM-P profile " - "(Q64/KV256 or 16-bit Q128/KV128)" + "(Q64/KV256 or Q128/KV128)" ) if self.tile_size_q != selected_q_tile: raise ValueError( @@ -1452,7 +2143,8 @@ def _grouped_keeps_profile_key(self) -> _GroupedKeepsProfileKey: """Return the dtype, shape, and staging key used by Keeps recipe tables.""" return ( self.q_dtype, - self.kv_dtype, + self.k_dtype, + self.v_dtype, self.out_dtype, self.headdim, self.head_dim_per_stage_kv, @@ -1460,6 +2152,23 @@ def _grouped_keeps_profile_key(self) -> _GroupedKeepsProfileKey: self.o_stages, ) + @property + def _uses_contiguous_grouped_keeps_recipe(self) -> bool: + """Whether the dtype and staging key is a contiguous grouped-Keeps recipe.""" + profiles = ( + _SAGE_GROUPED_KEEPS_PROFILES + if self.use_sage_attention + else _CONTIGUOUS_GROUPED_KEEPS_PROFILES + ) + return self._grouped_keeps_profile_key in profiles + + @property + def _kv256_dtypes_qualified(self) -> bool: + """Whether the Q64/KV256 profile is qualified for this dtype combination.""" + if self.use_sage_attention: + return self._uses_contiguous_grouped_keeps_recipe + return (self.q_dtype, self.k_dtype, self.out_dtype) in _KV256_DTYPE_RECIPES + @property def uses_guarded_fixed_q1_grouped_keeps(self) -> bool: """Whether inactive grouped rows are excluded solely at output stores. @@ -1640,11 +2349,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. - FP8 Q128/KV128 and dense 16-bit Q128/KV128 publish a complete - packed-P row per pipeline token. Q64/KV256 and block-sparse 16-bit - Q128/KV128 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. + 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. """ # 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 @@ -1673,18 +2382,26 @@ def streams_tmem_p_fragments(self) -> bool: the MMA warp start its PV k-slice before the row is complete, at the cost of one barrier round per fragment. - Streaming is limited to the 16-bit two-instance profiles whose route - loop waits on the K/V loads, where the earlier PV start hides load - latency: Q64/KV256 and block-sparse Q128/KV128. Dense Q128/KV128 - keeps the complete row because its route loop is not load-bound, so - the per-fragment barriers are not compensated. FP8 Q128 keeps the - complete row because its P publication packs four values per column - into one store. + 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. """ + if not self.uses_two_inst_tmem_p: + return False return ( - self.uses_two_inst_tmem_p - and not self.use_fp8_qkv - and (self.tile_size_kv == 256 or self.use_block_sparse) + self.tile_size_kv == 256 + or self.use_block_sparse + or ( + self.k_dtype.width == 8 + and self.v_dtype.width == 8 + and self.keeps_stats_via_smem + and not self.uses_q_token_kv_block_sparse_page_membership + ) ) @property @@ -1696,15 +2413,89 @@ def defers_softmax_anchor_updates(self) -> bool: ``SOFTMAX_RESCALE_THRESHOLD_LOG2`` trades a bounded 16-bit P range (2**8) for fewer TMEM rescales. The profiles listed here are the ones where that trade was measured to pay: KV256 tiles and block-sparse - routes, whose row maximum moves often but rarely by much. + routes, whose row maximum moves often but rarely by much. FP8 P uses + the static 448 scale, which needs ``p <= 1``, so it always anchors on + the exact row maximum. """ - return self.use_keeps_mma_ab and ( - self.tile_size_kv == 256 or self.use_block_sparse + return ( + self.use_keeps_mma_ab + and not self.use_fp8_pv + and (self.tile_size_kv == 256 or self.use_block_sparse) ) + @property + def use_sage_attention(self) -> bool: + """Whether Q/K scores carry per-block scales and V per-channel scales.""" + return self.sage_k_block_size > 0 + + @property + def uses_int32_scores(self) -> bool: + """Whether BMM1 accumulates INT8 Q/K into INT32 scores.""" + return self.qk_dtype == Int8 + + def softmax_num_warps(self, inst_id: int) -> int: + """Return the warps of one softmax instance.""" + return self.softmax0_num_warps if inst_id == 0 else self.softmax1_num_warps + + def sage_k_groups_for_block(self, k_block_size: int) -> int: + """Return the K scale groups inside one streamed K32 score fragment.""" + return max(1, self.softmax_score_fragment_regs // k_block_size) + + @property + def sage_k_groups_per_fragment(self) -> int: + """Return the scale groups per fragment of an exact route or dense tile.""" + if not self.use_sage_attention: + return 1 + return self.sage_k_groups_for_block(self.sage_k_block_size) + + @property + def sage_summary_k_groups_per_fragment(self) -> int: + """Return the scale groups per fragment of a proxy route's summaries.""" + if not self.use_sage_attention: + return 1 + return self.sage_k_groups_for_block(self.sage_k_summary_block_size) + + def sage_k_groups_per_fragment_for(self, proxy: bool) -> int: + """Return the scale groups per fragment of one route kind.""" + if proxy: + return self.sage_summary_k_groups_per_fragment + return self.sage_k_groups_per_fragment + + @property + def sage_mixed_k_geometry(self) -> bool: + """Whether exact and proxy routes use different scale-group geometries.""" + return ( + self.use_sage_attention + and self.sage_summary_k_groups_per_fragment + != self.sage_k_groups_per_fragment + ) + + def sage_scores_dequantized_for(self, proxy: bool) -> bool: + """Whether the max pass writes one route kind's scores back dequantized. + + With one scale per score (the one-token K block), the P pass then + reads neither ``sfK`` nor the INT32 bias. Coarser geometries keep the + scores quantized to avoid a TMEM round trip per tile. + """ + return ( + self.use_sage_attention + and self.sage_k_groups_per_fragment_for(proxy) + == self.softmax_score_fragment_regs + ) + + def sage_k_scales_in_smem_for(self, groups: int) -> bool: + """Whether a tile with ``groups`` scale groups per fragment keeps ``sfK`` in SMEM. + + Four or more groups per K32 fragment (K blocks of 4 or 1 token) would + need 16 or more registers per lane, so those tiles read ``sfK`` from + SMEM; larger blocks keep a rotating register array. + """ + return self.use_sage_attention and groups >= 4 + @property def matches_kv256_task_topology(self) -> bool: """Whether task roles match KV256's validated 16-warp layout.""" + return all( getattr(self, field) == expected for field, expected in _KV_TILE_256_TASK_TOPOLOGY_DEFAULTS.items() @@ -1715,7 +2506,8 @@ def keeps_separates_tmem_s_and_stats(self) -> bool: """Whether two-instance Keeps has room for standalone stats tiles.""" if not ( self.use_keeps_mma_ab - and self.use_fp8_qkv + and self.k_dtype.width == 8 + and self.v_dtype.width == 8 and self.tile_size_kv == 128 and self.head_dim_per_stage_kv == 0 and self.num_insts_kv == 2 @@ -1742,6 +2534,20 @@ def keeps_stats_via_smem(self) -> bool: """ return self.use_keeps_mma_ab and not self.keeps_separates_tmem_s_and_stats + @property + def splits_kv_tile_256_tail_columns(self) -> bool: + """Whether the KV256 tail splits the output columns between spatial halves. + + Each correction lane then merges and stores its own 64 of the 128 + output columns, so all four correction warps store. The per-lane + exchange needs SMEM that only the byte-wide K/V ring leaves free. + """ + return ( + self.tile_size_kv == 256 + and self.k_dtype.width == 8 + and self.v_dtype.width == 8 + ) + @property def keeps_loop_correction_chunk_regs(self) -> int: """Return FP32 registers corrected by one Keeps TMEM pair.""" @@ -1874,6 +2680,45 @@ def max_runtime_row_split_segments(self) -> int: rows_per_slice = self.split_reduction_rows_per_slice return max((self.tile_size_q + rows_per_slice - 1) // rows_per_slice, 1) + @property + def uses_direct_kv_launch(self) -> bool: + """Whether the attention CTA publishes final output directly.""" + return not (self.use_split_kv or self.use_separate_reduction_kernel) + + @property + def has_valid_split_kv_fanout(self) -> bool: + """Whether the configured split fanout is internally consistent.""" + return ( + self.use_split_kv + and self.splits_kv > 1 + and self.max_splits_kv >= self.splits_kv + ) + + def validate_q_token_kv_block_sparse_grouped_keeps_profile(self) -> None: + """Validate the grouped-Keeps profiles for token-query sparse routes.""" + common_profile = ( + self.use_keeps_mma_ab + and self.groups_tokens_heads_q + and self.tile_size_q in (64, 128) + and self.uses_q_token_kv_block_sparse_page_route + and self.mask_type == CAUSAL + and not self.use_cluster_smem_reduction + and not self.use_attention_sinks + ) + kv128_profile = ( + self.tile_size_kv == 128 + and self._grouped_keeps_profile_key + in _Q_TOKEN_KV_BLOCK_SPARSE_GROUPED_KEEPS_PROFILES + and not self.use_sliding_window_causal + and (self.uses_direct_kv_launch or self.has_valid_split_kv_fanout) + ) + if not (common_profile and kv128_profile): + raise ValueError( + "QToken-KvBlock-Sparse-Attention grouped KeepsMmaAb requires an encoded sparse Q group " + "that fits Q64/Q128, a supported D64/D128/D256 KV128 recipe, " + "and a supported causal direct, persistent or split launch" + ) + @property def supports_grouped_keeps(self) -> bool: """Whether a validated config uses a qualified grouped-Keeps recipe.""" @@ -1887,8 +2732,14 @@ def supports_grouped_keeps(self) -> bool: or not self.groups_tokens_heads_q or self.tile_size_q != 64 or self.headdim != 128 - or self.q_dtype not in (Float16, BFloat16) - or not (self.q_dtype == self.kv_dtype == self.out_dtype) + or not self._kv256_dtypes_qualified + # FP8 Keeps recipes exclude attention sinks, as the paged FP8 + # Q64/Q128 profiles do. + or ( + self.k_dtype.width == 8 + and self.v_dtype.width == 8 + and self.use_attention_sinks + ) or self.use_cluster_smem_reduction or not self.matches_kv256_task_topology ): @@ -1915,25 +2766,40 @@ def supports_grouped_keeps(self) -> bool: return False profile = self._grouped_keeps_profile_key + direct = not (self.use_split_kv or self.use_separate_reduction_kernel) - # Keep block-sparse qualification separate from the dense/paged - # profile matrix below. Its structural, masking, and reduction - # constraints are validated separately; both scheduler modes use the - # same qualified recipe keys. + # Contiguous K/V launches share one qualified recipe set, apart from the + # dense/paged profile matrix below. Block-sparse constraints are + # validated separately; a dense contiguous launch must be direct. if self.use_block_sparse: - return profile in _BLOCK_SPARSE_GROUPED_KEEPS_PROFILES - - direct = not (self.use_split_kv or self.use_separate_reduction_kernel) + return self._uses_contiguous_grouped_keeps_recipe + if ( + self._uses_contiguous_grouped_keeps_recipe + and direct + and not self.use_paged_kv + and not any( + ( + self.use_variable_seqlens_q, + self.use_sliding_window_causal, + self.use_attention_sinks, + ) + ) + ): + return True + if self.use_sage_attention: + return False if profile in _GROUPED_KEEPS_PAGED_FP8_PROFILES: - fixed_q1_ratio32 = self.max_seq_len_q == 1 and self.heads_q_per_kv == 32 + 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 return ( - (fixed_q1_ratio32 or fixed_grouped_q) + (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 == CAUSAL + and self.mask_type in (DENSE, CAUSAL) and not any( ( self.use_cluster_smem_reduction, @@ -1966,7 +2832,10 @@ def supports_grouped_keeps(self) -> bool: ) ) ) - if profile != _GROUPED_KEEPS_MAIN_PROFILE: + if profile not in ( + _GROUPED_KEEPS_MAIN_PROFILE, + _GROUPED_KEEPS_MIXED_KV_DTYPE_PROFILE, + ): return False if self.tile_size_q == 128: @@ -2013,6 +2882,9 @@ def supports_grouped_keeps(self) -> bool: SUPPORTED_IO_DTYPES = {Float16, BFloat16, Float8E4M3FN} +SUPPORTED_KV_DTYPES = SUPPORTED_IO_DTYPES | {Float4E2M1FN} +# Q and K also admit Int8, which only the Sage recipe uses. +SUPPORTED_QK_DTYPES = SUPPORTED_IO_DTYPES | {Int8} SUPPORTED_ACC_DTYPES = {Float32} @@ -2045,7 +2917,19 @@ def _apply_config_source(cfg: FmhaDecodeConfig, source: object) -> set[str]: field_names = FmhaDecodeConfig.__dataclass_fields__ explicit_fields: set[str] = set() for source_item in _iter_config_sources(source): - for key, value in _decode_config_items(source_item): + source_fields = dict(_decode_config_items(source_item)) + # Normalize legacy shared dtypes before individual fields so explicit + # q_dtype/k_dtype/v_dtype values retain their normal precedence. + for alias, names in ( + ("qkv_dtype", ("q_dtype", "k_dtype", "v_dtype")), + ("kv_dtype", ("k_dtype", "v_dtype")), + ): + value = source_fields.get(alias) + if value is not None: + for name in names: + setattr(cfg, name, value) + explicit_fields.add(name) + for key, value in source_fields.items(): if key == "use_causal_spec_decoding": raise ValueError( "use_causal_spec_decoding was removed; use " @@ -2185,6 +3069,9 @@ def _finalize_static_decode_config( if not use_keeps_mma_ab and cfg.headdim > 128: _set_if_implicit(cfg, "head_dim_per_stage_kv", 128, explicit_fields) _set_if_implicit(cfg, "num_insts_kv", 2, explicit_fields) + if cfg.use_transform_kv and cfg.use_paged_kv: + # Force using 1 instances kv to allow more warps for transform kv tasks. + _set_if_implicit(cfg, "num_insts_kv", 1, explicit_fields) if use_keeps_mma_ab: tile_size_q = cfg.tile_size_q if "tile_size_q" in explicit_fields else 64 @@ -2200,6 +3087,15 @@ def _finalize_static_decode_config( # effective configuration is supported by the kernel. for field_name, value in _KV_TILE_256_PHYSICAL_DEFAULTS.items(): _set_if_implicit(cfg, field_name, value, explicit_fields) + # The K/V ring depth follows the element width and stays tunable. + _set_if_implicit( + cfg, + "kv_stages", + KV_TILE_256_BYTE_WIDE_SHARED_FIFO_STAGES + if cfg.k_dtype.width == 8 and cfg.v_dtype.width == 8 + else KV_TILE_256_SHARED_FIFO_STAGES, + explicit_fields, + ) for field_name, value in _KV_TILE_256_TASK_TOPOLOGY_DEFAULTS.items(): _set_if_implicit(cfg, field_name, value, explicit_fields) else: @@ -2235,8 +3131,54 @@ def _finalize_static_decode_config( _set_if_implicit(cfg, "q_stages", 1, explicit_fields) _set_if_implicit(cfg, "ordered_softmax_barrier_mode", 1, explicit_fields) - if cfg.tile_size_kv != 256 and "kv_stages" not in explicit_fields: - cfg.kv_stages = cfg.inferred_kv_stages + # Derive transform placement after the profile has selected its KV + # instance count: one-instance mixed decode reuses Softmax1's warp group. + uses_full_wg_transform = cfg.use_transform_kv and cfg.num_insts_kv == 1 + if uses_full_wg_transform: + _set_if_implicit(cfg, "transform_kv_warp_idx", 4, explicit_fields) + _set_if_implicit(cfg, "transform_kv_num_warps", 4, explicit_fields) + + # MmaTask publishes one TMEM O stage per K/V instance in each loop group, + # and Correction uses those stage slots as the corresponding accumulators. + _set_if_implicit(cfg, "o_stages", cfg.num_insts_kv, explicit_fields) + + _set_if_implicit( + cfg, + "store_transformed_kv_in_tmem", + ( + cfg.use_nvfp4_kv + and cfg.use_fp8_q + and not cfg.use_keeps_mma_ab + and cfg.head_dim_kv_stage == 128 + and cfg.tile_size_kv == 128 + ), + explicit_fields, + ) + if cfg.store_transformed_kv_in_tmem: + # Allocate the remaining TMEM columns to transformed K/V eagerly. + if ( + "transformed_kv_stages" not in explicit_fields + or cfg.transformed_kv_stages == 0 + ): + cfg.transformed_kv_stages = ( + TMEM_MAX_ALLOCATION_COLUMNS + - cfg.tmem_total_cols + + cfg.tmem_transformed_kv_stage_cols * cfg.transformed_kv_stages + # tmem_total_cols already includes the tmem of transform kv task. drop it + ) // cfg.tmem_transformed_kv_stage_cols + + if cfg.tile_size_kv != 256: + _set_if_implicit(cfg, "kv_stages", cfg.inferred_kv_stages, explicit_fields) + + if ( + cfg.use_transform_kv + and not cfg.store_transformed_kv_in_tmem + and cfg.transformed_kv_stages == 0 + ): + # SMEM-backed mixed modes still need a real number of stages. + _set_if_implicit( + cfg, "transformed_kv_stages", min(cfg.kv_stages, 4), explicit_fields + ) # Split-KV mode forbids persistent scheduling. Canonicalize here so # downstream consumers can gate on use_persistent_scheduler alone without @@ -2246,7 +3188,11 @@ def _finalize_static_decode_config( def _validate_kv256_static_config(cfg: FmhaDecodeConfig) -> None: - """Validate the effective KV256 profile after implicit defaults are filled.""" + """Validate the effective KV256 profile after implicit defaults are filled. + + The SMEM check covers the Q and K/V pipelines; the schedule builder checks + the complete layout against the SM capacity. + """ if cfg.tile_size_kv != 256: return @@ -2261,10 +3207,9 @@ def _require_python_int(field_name: str) -> int: for field_name, expected in _KV_TILE_256_PHYSICAL_DEFAULTS.items(): actual = _require_python_int(field_name) - if field_name in _KV_TILE_256_TUNABLE_FIELDS: - continue if actual != expected: raise ValueError(f"KV256 requires {field_name}={expected}, got {actual}") + _require_python_int("kv_stages") for field_name, expected in _KV_TILE_256_TASK_TOPOLOGY_DEFAULTS.items(): actual = _require_python_int(field_name) if actual != expected: @@ -2289,8 +3234,7 @@ def _require_python_int(field_name: str) -> int: ) if not cfg.supports_grouped_keeps: raise ValueError( - "KV256 currently supports only the qualified Q64 FP16/BF16/D128 " - "grouped Keeps profile" + "KV256 supports only the qualified Q64 D128 grouped Keeps profile" ) @@ -2315,24 +3259,45 @@ def _append_padding_warp_roles(cfg: FmhaDecodeConfig, roles: list[_WarpRole]) -> preferred_wg = role.preferred_warp_idx // 4 preferred_wg_end = (role.preferred_warp_idx + role.num_warps - 1) // 4 if preferred_wg != preferred_wg_end: - raise ValueError( - f"{role.name} spans warp groups {preferred_wg} and {preferred_wg_end}; " - "each task role must fit within one warp group" + # A producer-only LoadTask may occupy complete consecutive warp + # groups. Keeping each group full preserves warpgroup register + # reallocation participation and needs no cross-group padding. + spans_complete_wgs = ( + role.name == "load" + and role.preferred_warp_idx % 4 == 0 + and role.num_warps % 4 == 0 ) - if preferred_wg < 0 or preferred_wg >= MAX_WARP_GROUPS: + if not spans_complete_wgs: + raise ValueError( + f"{role.name} spans warp groups {preferred_wg} and " + f"{preferred_wg_end}; only a complete-group load task may " + "span warp groups" + ) + if preferred_wg < 0 or preferred_wg_end >= MAX_WARP_GROUPS: raise ValueError( - f"{role.name} is assigned to unsupported warp group {preferred_wg}" + f"{role.name} is assigned through unsupported warp group " + f"{preferred_wg_end}" ) - num_warps_per_wg[preferred_wg] += role.num_warps - total_num_wgs = max(total_num_wgs, preferred_wg + 1) + role_warps_remaining = role.num_warps + role_wg = preferred_wg + while role_warps_remaining > 0: + warps_in_group = min(role_warps_remaining, 4) + num_warps_per_wg[role_wg] += warps_in_group + role_warps_remaining -= warps_in_group + role_wg += 1 + total_num_wgs = max(total_num_wgs, preferred_wg_end + 1) for wg_idx, num_warps in enumerate(num_warps_per_wg): if num_warps > 4: raise ValueError( f"warp group {wg_idx} has {num_warps} active warps; at most 4 are allowed" ) + # When using persistent scheduling, force 4 warp groups if cfg.use_persistent_scheduler: - total_num_wgs = MAX_WARP_GROUPS + # Ordinary persistent decode retains its established four-group + # contract. A profile that explicitly places a role in WG4 keeps that + # fifth group instead of truncating the task from the CTA. + total_num_wgs = max(total_num_wgs, 4) for wg_idx, active_warps in enumerate(num_warps_per_wg): padding_index_field = f"wg{wg_idx}_padding_warp_idx" @@ -2379,6 +3344,15 @@ def _active_warp_roles(cfg: FmhaDecodeConfig) -> list[_WarpRole]: cfg.softmax1_num_warps, ) ) + if cfg.use_transform_kv: + roles.append( + _WarpRole( + "transform_kv", + "transform_kv_warp_idx", + cfg.transform_kv_warp_idx, + cfg.transform_kv_num_warps, + ) + ) if cfg.use_persistent_scheduler: roles.extend( @@ -2461,6 +3435,7 @@ def _make_static_decode_config( explicit_fields=explicit_fields, ) _finalize_static_decode_config(cfg, explicit_fields) + _validate_mixed_kv_dtype_profile(cfg) _validate_kv256_static_config(cfg) _finalize_warp_roles(cfg) return cfg @@ -2591,9 +3566,12 @@ def _max_splits_kv_by_work( tile_size_kv: int, num_insts_kv: int, max_splits_kv: int | None = None, + min_loop_iters_per_split: int = MIN_LOOP_ITERS_PER_SPLIT, ) -> int: """Return the fanout cap that retains useful KV work per CTA.""" - tile_size_per_cta_kv = tile_size_kv * num_insts_kv * MIN_LOOP_ITERS_PER_SPLIT + if min_loop_iters_per_split <= 0: + raise ValueError("min_loop_iters_per_split must be positive") + tile_size_per_cta_kv = tile_size_kv * num_insts_kv * min_loop_iters_per_split max_by_seq = max( 1, (seq_len_kv + tile_size_per_cta_kv - 1) // tile_size_per_cta_kv, @@ -2612,6 +3590,7 @@ def enumerate_auto_splits_kv( num_insts_kv: int, num_q_tiles: int, service_capacity: int, + split_kv: bool = True, ) -> tuple[int, ...]: """Enumerate direct and useful split fanouts for an under-filled Q grid. @@ -2619,6 +3598,8 @@ def enumerate_auto_splits_kv( the empirical score. The latter is important when a partially filled wave cannot be completed by any uniform integer fanout. """ + if not split_kv: + return (1,) if num_q_tiles <= 0: raise ValueError("num_q_tiles must be positive") if service_capacity <= 0: @@ -2650,15 +3631,21 @@ def select_splits_kv( service_capacity: int | None = None, requested_splits_kv: int = -1, max_splits_kv: int | None = None, + min_loop_iters_per_split: int = MIN_LOOP_ITERS_PER_SPLIT, + split_kv: bool = True, ) -> int: """ Select a Q-grid-aware split-KV fanout. Every legal TileQ is considered with its actual number of Q CTAs. The automatic fanout fills otherwise idle cluster-size-one service slots while - retaining at least ``MIN_LOOP_ITERS_PER_SPLIT`` KV iterations per CTA. - Positive caller fanouts remain pinned subject only to the KV-work cap. + retaining at least ``min_loop_iters_per_split`` KV iterations per CTA. + Positive caller fanouts remain pinned subject only to the KV-work cap. The + default preserves the dense-decode policy; specialized bounded-work routes + may choose a lower positive iteration floor explicitly. """ + if not split_kv: + return 1 if num_q_tiles <= 0: raise ValueError("num_q_tiles must be positive") max_by_seq = _max_splits_kv_by_work( @@ -2666,6 +3653,7 @@ def select_splits_kv( tile_size_kv=tile_size_kv, num_insts_kv=num_insts_kv, max_splits_kv=max_splits_kv, + min_loop_iters_per_split=min_loop_iters_per_split, ) if requested_splits_kv > 0: return max(1, min(max_by_seq, requested_splits_kv)) @@ -2700,6 +3688,8 @@ def _select_auto_launch_mode( tile_size_kv: int = AUTO_LAUNCH_TILE_SIZE_KV, persistent_min_waves: int = 1, persistent_min_tiles_per_cta: int = 1, + split_kv: bool = True, + service_capacity: int | None = None, ) -> str: """Pick the launch mode that best matches the kernel's parallelism budget. @@ -2733,14 +3723,15 @@ def _select_auto_launch_mode( """ if seq_len_kv <= 0 or batch_size <= 0 or num_heads_kv <= 0 or num_q_tiles <= 0: return "static" - hardware_info = utils.HardwareInfo() - sm_count = hardware_info.get_device_multiprocessor_count() + sm_count = service_capacity + if sm_count is None: + sm_count = utils.HardwareInfo().get_device_multiprocessor_count() sm_count = FALLBACK_SM_COUNT_B200 if sm_count <= 0 else sm_count ctas = batch_size * num_heads_kv * num_q_tiles waves = ctas / sm_count tiles_per_cta = (seq_len_kv + tile_size_kv - 1) // tile_size_kv kv_tokens_per_cta = tiles_per_cta * tile_size_kv - if waves < 1 and kv_tokens_per_cta >= SPLIT_KV_MIN_TOKENS_PER_CTA: + if split_kv and waves < 1 and kv_tokens_per_cta >= SPLIT_KV_MIN_TOKENS_PER_CTA: return "gmem_reduction" if ( ctas > persistent_min_waves * sm_count @@ -2916,6 +3907,70 @@ def _apply_swaps_tile_config( } +def _try_apply_default_wide_keeps_config( + cfg: FmhaDecodeConfig, + *, + explicit_fields: set[str], + seq_len_q: int, + num_heads_q: int, + num_heads_kv: int, +) -> bool: + """Select Q64/Q128 Keeps for an unpinned head ratio above 32. + + Prefer a grouped tile so a partial GQA group can occupy the next supported + MMA width. If that profile family is not qualified, an exact 64- or + 128-head group may still use the established ungrouped Keeps path. The + caller falls back to ungrouped Q16 Swaps head bands for other profiles. + """ + + heads_q_per_kv = num_heads_q // num_heads_kv + if ( + heads_q_per_kv <= 32 + or heads_q_per_kv > 128 + or bool(_MMA_SELECTION_FIELDS & explicit_fields) + ): + return False + + tile_size_q = 64 if heads_q_per_kv <= 64 else 128 + grouping_was_explicit = "groups_tokens_heads_q" in explicit_fields + grouping_candidates = (cfg.groups_tokens_heads_q,) + if ( + not grouping_was_explicit + and cfg.groups_tokens_heads_q + and heads_q_per_kv == tile_size_q + ): + grouping_candidates += (False,) + + for groups_tokens_heads_q in grouping_candidates: + probe = deepcopy(cfg) + probe.groups_tokens_heads_q = groups_tokens_heads_q + probe.use_keeps_mma_ab = True + probe.tile_size_q = tile_size_q + try: + _finalize_static_decode_config( + probe, + explicit_fields | {"tile_size_q"}, + ) + _validate_profile_support( + cfg=probe, + seq_len_q=seq_len_q, + num_heads_q=num_heads_q, + num_heads_kv=num_heads_kv, + split_kv_mode="disabled", + ) + except ValueError: + continue + + cfg.groups_tokens_heads_q = groups_tokens_heads_q + cfg.use_keeps_mma_ab = True + cfg.tile_size_q = tile_size_q + # Keeps finalization otherwise canonicalizes an implicit tile to Q64. + explicit_fields.add("tile_size_q") + return True + + return False + + def _try_apply_auto_kv256_profile( cfg: FmhaDecodeConfig, *, @@ -2953,7 +4008,7 @@ def _try_apply_auto_kv256_profile( and cfg.groups_tokens_heads_q and cfg.headdim == 128 and cfg.q_dtype in (Float16, BFloat16) - and cfg.q_dtype == cfg.kv_dtype == cfg.out_dtype + and cfg.q_dtype == cfg.k_dtype == cfg.v_dtype == cfg.out_dtype and num_heads_q // num_heads_kv <= 64 ) if not profile_is_eligible: @@ -2989,6 +4044,7 @@ def _resolve_grouped_q_launch_candidates( num_heads_q: int, num_heads_kv: int, service_capacity: int, + split_kv: bool = True, ) -> tuple[GroupedQLaunchCandidate, ...]: """Resolve one TileQ with Q-grid-aware direct and split recipes. @@ -3008,7 +4064,7 @@ def _resolve_grouped_q_launch_candidates( and candidate.tile_size_q == 64 and probe.headdim == 128 and probe.q_dtype == BFloat16 - and probe.q_dtype == probe.kv_dtype == probe.out_dtype + and probe.q_dtype == probe.k_dtype == probe.v_dtype == probe.out_dtype ): # TileQ selection uses a common KV128 cost basis and is independent of # the later KV-tile decision. BF16 Q64 has no final KV128 profile, but @@ -3016,7 +4072,8 @@ def _resolve_grouped_q_launch_candidates( # logical candidate. The KV selector materializes and validates the # actual BF16 profile only after the Q winner is known. probe.q_dtype = Float16 - probe.kv_dtype = Float16 + probe.k_dtype = Float16 + probe.v_dtype = Float16 probe.out_dtype = Float16 try: _finalize_static_decode_config( @@ -3042,6 +4099,7 @@ def _resolve_grouped_q_launch_candidates( num_insts_kv=cost_num_insts_kv, num_q_tiles=candidate.q_tiles, service_capacity=service_capacity, + split_kv=split_kv, ) recipes = [] for candidate_splits_kv in split_candidates: @@ -3088,12 +4146,14 @@ def _apply_auto_grouped_q_mma_config( num_heads_kv: int, splits_kv: int, max_splits_kv: int | None, + split_kv: bool = True, ) -> GroupedQLaunchCandidate | None: """Select a legal fixed multi-Q MMA and KV-split recipe when unpinned. 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. + Mixed-precision candidates use SwapsMmaAb since Keeps is not supported yet. """ if ( not auto_tuner @@ -3116,6 +4176,7 @@ def _apply_auto_grouped_q_mma_config( candidates = enumerate_grouped_q_mma_candidates( heads_q_per_kv=num_heads_q // num_heads_kv, seq_len_q=seq_len_q, + swaps_only=cfg.use_transform_kv, ) if not candidates: return None @@ -3133,6 +4194,7 @@ def _apply_auto_grouped_q_mma_config( num_heads_q=num_heads_q, num_heads_kv=num_heads_kv, service_capacity=service_capacity, + split_kv=split_kv, ) ) if not supported: @@ -3186,14 +4248,56 @@ def _apply_auto_grouped_q_mma_config( return selected +def _apply_dtype_config( + cfg: FmhaDecodeConfig, + *, + explicit_fields: set[str], + q_dtype: type, + k_dtype: type, + v_dtype: type, + o_dtype: type, +) -> None: + """Apply independent Q/K/V storage dtypes before selecting MMA geometry.""" + for name, dtype in ( + ("q_dtype", q_dtype), + ("k_dtype", k_dtype), + ("v_dtype", v_dtype), + ("out_dtype", o_dtype), + ): + _set_if_implicit(cfg, name, dtype, explicit_fields) + + +def _resolve_decode_dtype_aliases( + *, + q_dtype: type | None, + k_dtype: type | None, + v_dtype: type | None, + qkv_dtype: type | None, + kv_dtype: type | None, +) -> tuple[type, type, type]: + """Resolve legacy defaults while retaining independent canonical dtypes.""" + if kv_dtype is not None: + for name, dtype in (("k_dtype", k_dtype), ("v_dtype", v_dtype)): + if dtype is not None and dtype != kv_dtype: + raise ValueError(f"{name} conflicts with kv_dtype") + default_dtype = qkv_dtype if qkv_dtype is not None else Float16 + default_kv_dtype = kv_dtype if kv_dtype is not None else default_dtype + return ( + q_dtype if q_dtype is not None else default_dtype, + k_dtype if k_dtype is not None else default_kv_dtype, + v_dtype if v_dtype is not None else default_kv_dtype, + ) + + def _apply_default_q_grouping( cfg: FmhaDecodeConfig, *, explicit_fields: set[str], + heads_q_per_kv: int, ) -> None: - """Enable token/head grouping unless the caller explicitly opts out.""" + """Group complete head sets when they fit one supported Q tile.""" if "groups_tokens_heads_q" not in explicit_fields: - cfg.groups_tokens_heads_q = True + cfg.groups_tokens_heads_q = heads_q_per_kv <= 128 def _apply_layout_config( @@ -3201,14 +4305,19 @@ def _apply_layout_config( *, qkv_layout: str, num_tokens_per_page: int, + storage_tokens_per_page: int | None, seq_len_kv: int, ) -> str: """Apply contiguous/paged-KV layout fields and return the canonical layout.""" qkv_layout = normalize_qkv_layout(qkv_layout) if qkv_layout == "pagedKv": validate_page_size(num_tokens_per_page) + if storage_tokens_per_page is None: + storage_tokens_per_page = cfg.storage_tokens_per_page or num_tokens_per_page + validate_storage_page_size(num_tokens_per_page, storage_tokens_per_page) cfg.use_paged_kv = True cfg.num_tokens_per_page = num_tokens_per_page + cfg.storage_tokens_per_page = storage_tokens_per_page cfg.max_num_pages_per_seq_kv = ( seq_len_kv + num_tokens_per_page - 1 ) // num_tokens_per_page @@ -3227,6 +4336,7 @@ def _apply_feature_config( use_attention_sinks: bool, ) -> None: """Apply feature flags that affect config construction.""" + if sliding_window_causal: cfg.use_sliding_window_causal = True cfg.attention_window_size = attention_window_size @@ -3283,6 +4393,7 @@ def _apply_auto_launch_mode( num_heads_kv: int, seq_len_kv: int, seq_len_q: int, + split_kv: bool = True, ) -> str: """Apply the static/persistent/split-KV launch heuristic when it is allowed.""" if not _should_auto_select_launch_mode( @@ -3299,14 +4410,11 @@ def _apply_auto_launch_mode( seq_len_kv=seq_len_kv, num_q_tiles=_num_q_tiles_for_launch(cfg), tile_size_kv=cfg.tile_size_kv, + split_kv=split_kv, ) - if (cfg.use_variable_seqlens_q or cfg.use_sliding_window_causal) and mode == ( - "gmem_reduction" - ): - # Runtime Q offsets and sliding-window bounds are compatible with CLC - # work discovery, but they deliberately remain nonsplit. Underfilled - # grids therefore stay direct while grids above one resident wave use - # the same structural persistence rule as fixed-Q decode. + if cfg.use_sliding_window_causal and mode == ("gmem_reduction"): + # Automatic sliding-window split recipes are not qualified. Query + # storage layout alone does not restrict split selection. return split_kv_mode if mode not in ("gmem_reduction", "persistent"): return split_kv_mode @@ -3362,6 +4470,7 @@ def _apply_split_kv_config( max_splits_kv: int | None, sliding_window_causal: bool, attention_window_size: int, + min_loop_iters_per_split: int = MIN_LOOP_ITERS_PER_SPLIT, ) -> None: """Resolve split-KV fanout and config flags for the selected reduction mode.""" if split_kv_mode == "disabled": @@ -3382,6 +4491,7 @@ def _apply_split_kv_config( num_q_tiles=_num_q_tiles_for_launch(cfg), requested_splits_kv=splits_kv, max_splits_kv=max_splits_kv, + min_loop_iters_per_split=min_loop_iters_per_split, ) if selected_splits_kv <= 1: return @@ -3487,6 +4597,24 @@ def _select_auto_split_kv_reduction_mode( return _config_with_split_kv_mode(cfg, "gmem_reduction"), "gmem_reduction" +def _validate_mixed_kv_dtype_profile(cfg: FmhaDecodeConfig) -> None: + """Reject k_dtype != v_dtype profiles outside the supported combinations. + + Sage pairs Int8 K with E4M3 V; 16-bit K with E4M3 V is dense KV128 only. + """ + if cfg.k_dtype == cfg.v_dtype: + return + if cfg.use_sage_attention and cfg.k_dtype == Int8 and cfg.v_dtype == Float8E4M3FN: + return + if cfg.use_block_sparse or cfg.tile_size_kv == 256: + raise ValueError( + "k_dtype != v_dtype requires use_block_sparse=False and " + "tile_size_kv=128; got " + f"use_block_sparse={cfg.use_block_sparse}, " + f"tile_size_kv={cfg.tile_size_kv}" + ) + + def _validate_profile_support( *, cfg: FmhaDecodeConfig, @@ -3502,6 +4630,10 @@ def _validate_profile_support( use_groups_tokens_heads_q = cfg.groups_tokens_heads_q tile_size_q = cfg.tile_size_q cfg.validate_boolean_fields() + cfg.validate_dtypes() + _validate_mixed_kv_dtype_profile(cfg) + if cfg.o_stages != cfg.num_insts_kv: + raise ValueError("fmha_decode requires o_stages == num_insts_kv") _validate_kv256_static_config(cfg) if cfg.mask_type not in (DENSE, CAUSAL): raise ValueError("mask_type must be DENSE or CAUSAL") @@ -3514,6 +4646,7 @@ def _validate_profile_support( raise ValueError( "heads_q_per_kv metadata must equal num_heads_q / num_heads_kv" ) + cfg.validate_sage_profile() cfg.validate_block_sparse_profile(heads_q_per_kv=heads_q_per_kv) if use_groups_tokens_heads_q: make_q_tile_geometry( @@ -3530,7 +4663,8 @@ def _validate_profile_support( and cfg.use_separate_reduction_kernel and not cfg.use_cluster_smem_reduction and cfg.q_dtype == Float8E4M3FN - and cfg.kv_dtype == Float8E4M3FN + and cfg.k_dtype == Float8E4M3FN + and cfg.v_dtype == Float8E4M3FN and ( (headdim == 128 and cfg.out_dtype == Float8E4M3FN) or ( @@ -3592,26 +4726,51 @@ def _validate_profile_support( "parallel separate reduction cluster size " f"{cluster_size} is not supported on this device" ) + is_q_token_kv_block_sparse_grouped_keeps = ( + use_keeps_mma_ab + and use_groups_tokens_heads_q + and cfg.uses_q_token_kv_block_sparse_page_route + ) + if ( + cfg.use_pdl + and cfg.use_split_kv + and ( + not cfg.use_separate_reduction_kernel + or cfg.use_cluster_smem_reduction + or cfg.use_persistent_scheduler + ) + ): + raise ValueError( + "split PDL requires nonpersistent attention with a standalone GMEM reducer" + ) + if is_q_token_kv_block_sparse_grouped_keeps: + cfg.validate_q_token_kv_block_sparse_grouped_keeps_profile() supports_grouped_keeps = cfg.supports_grouped_keeps if cfg.tile_size_kv != 128 and not ( cfg.tile_size_kv == 256 and supports_grouped_keeps ): raise ValueError( - "wide KeepsMmaAb is enabled only for the qualified FP16/BF16/D128 " - "KV256 native warp-specialized profile" + "wide KeepsMmaAb is enabled only for the qualified D128 KV256 " + "native warp-specialized profile" ) if use_keeps_mma_ab and use_groups_tokens_heads_q: - if not supports_grouped_keeps: + if not (is_q_token_kv_block_sparse_grouped_keeps or supports_grouped_keeps): raise ValueError( "grouped KeepsMmaAb currently supports only validated narrow " "profiles; pass groups_tokens_heads_q=False to use a supported " "ungrouped Keeps profile" ) if cfg.use_variable_seqlens_q: + if cfg.use_transform_kv: + raise ValueError( + "packed variable-Q does not support mixed Q/KV dtypes or " + "transformed K/V" + ) # Packed Q supports the broad grouped Swaps matrix plus the narrow # grouped Keeps direct profile validated above. if use_keeps_mma_ab and not ( - use_groups_tokens_heads_q and supports_grouped_keeps + use_groups_tokens_heads_q + and (is_q_token_kv_block_sparse_grouped_keeps or supports_grouped_keeps) ): raise ValueError( "packed variable-Q KeepsMmaAb requires its supported grouped " @@ -3622,7 +4781,41 @@ 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" + ) + 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 + elems_per_iteration = cfg.transform_kv_num_warps * 32 * 8 + if elems_per_iteration <= 0 or total_elems % elems_per_iteration != 0: + raise ValueError( + "SMEM transform-KV requires tile_size_kv * head_dim_kv_stage " + f"({total_elems}) to be divisible by a positive " + "transform_kv_num_warps * 32 * 8 " + f"({elems_per_iteration}); got " + f"transform_kv_num_warps={cfg.transform_kv_num_warps}" + ) + if cfg.store_transformed_kv_in_tmem and not ( + cfg.use_nvfp4_kv + and cfg.use_fp8_q + and not cfg.use_keeps_mma_ab + and cfg.head_dim_kv_stage == 128 + and cfg.tile_size_kv == 128 + and cfg.transform_kv_num_warps == 4 + and cfg.transformed_kv_stages > 0 + and cfg.tmem_total_cols <= TMEM_MAX_ALLOCATION_COLUMNS + ): + raise ValueError( + "store_transformed_kv_in_tmem requires FP8 Q, NVFP4 K/V, " + "SwapsMmaAb, 128-wide K/V head-dimension stages and KV tiles, " + "and a four-warp TransformKvTask" + ) 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 tile_size_q not in (64, 128): @@ -3656,9 +4849,16 @@ def _validate_profile_support( raise ValueError( "fmha_decode keepsMmaAb requires numHeadsQPerKv == tile_size_q" ) - if cfg.q_dtype == Float8E4M3FN and cfg.out_dtype not in ( - Float16, - Float8E4M3FN, + fp8_q_token_kv_block_sparse_bf16_output = ( + cfg.out_dtype == BFloat16 + and use_groups_tokens_heads_q + and is_q_token_kv_block_sparse_grouped_keeps + ) + if ( + cfg.q_dtype == Float8E4M3FN + and cfg.out_dtype not in (Float16, Float8E4M3FN) + and not fp8_q_token_kv_block_sparse_bf16_output + and not cfg.use_sage_attention ): raise ValueError( "fmha_decode keepsMmaAb fp8 qkv path supports fp16 or fp8 output" @@ -3678,7 +4878,10 @@ def _validate_profile_support( not use_groups_tokens_heads_q and heads_q_per_kv == tile_size_q ) or use_groups_tokens_heads_q separate_reduction_unstaged_supported = ( - headdim == 128 + ( + headdim == 128 + or (headdim == 64 and cfg.uses_q_token_kv_block_sparse_page_route) + ) and tile_size_q in (64, 128) and effective_head_dim_stage == 0 and effective_num_insts_kv == 2 @@ -3686,7 +4889,14 @@ def _validate_profile_support( ) separate_reduction_h256_supported = ( headdim == 256 - and tile_size_q == 128 + 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 effective_head_dim_stage == 128 and effective_num_insts_kv == 1 and effective_o_stages == 1 @@ -3700,7 +4910,8 @@ def _validate_profile_support( and cfg.use_paged_kv and cfg.num_tokens_per_page == 32 and cfg.q_dtype == Float8E4M3FN - and cfg.kv_dtype == Float8E4M3FN + and cfg.k_dtype == Float8E4M3FN + and cfg.v_dtype == Float8E4M3FN and cfg.out_dtype == (Float16 if headdim == 256 else Float8E4M3FN) and cfg.splits_kv == cfg.max_splits_kv and 2 <= cfg.max_splits_kv <= 128 @@ -3718,7 +4929,8 @@ def _validate_profile_support( ): raise ValueError( "separate reduction keepsMmaAb profiles require " - "the established D128/Q64-Q128 or D256/Q128 profiles, or a " + "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 " "reduction dtypes, Q layout, and static split-KV are required" ) @@ -3728,12 +4940,10 @@ 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 - effective_num_insts_kv = cfg.num_insts_kv if headdim == 256: - if effective_head_dim_stage != 128 or effective_num_insts_kv != 2: + if effective_head_dim_stage != 128: raise ValueError( - "fmha_decode SwapsMmaAb headDim=256 requires " - "head_dim_per_stage_kv=128 and num_insts_kv=2" + "fmha_decode SwapsMmaAb headDim=256 requires head_dim_per_stage_kv=128" ) elif effective_head_dim_stage != 0: raise ValueError( @@ -3765,14 +4975,20 @@ def _validate_profile_support( use_split_kv and cfg.supports_reduction_dtypes and (seq_len_q == 1 or cfg.use_variable_seqlens_q or use_groups_tokens_heads_q) - and (tile_size_q in (16, 32) or qualified_fp8_q8_separate_reduction_supported) + and ( + tile_size_q in (16, 32) + or qualified_fp8_q8_separate_reduction_supported + # The sparse publisher and standalone reducer both index + # actual logical rows, so a partial TileQ8 is also valid. + or (tile_size_q == 8 and cfg.uses_q_token_kv_block_sparse_page_route) + ) and headdim >= 64 ) if use_separate_reduction_kernel and not separate_reduction_supported: raise ValueError( "separate reduction SwapsMmaAb profiles require fixed SQ=1, " - "fixed grouped Q, or packed variable Q; tile_size_q in {16,32} " - "or a fixed FP8 Q8/HqPerKv8 profile (legacy D128 with FP8 output, " + "fixed grouped Q, or packed variable Q; tile_size_q in {16,32}, " + "QToken-KvBlock-Sparse-Attention TileQ8, or a fixed FP8 Q8/HqPerKv8 profile (legacy D128 with FP8 output, " "plus grouped paged-KV/page-32 D64 with FP8 output or D256 with " "FP16 output); valid reduction dtypes and static split-KV are required" ) @@ -3782,8 +4998,14 @@ def _validate_profile_support( raise ValueError( "fmha_decode SwapsMmaAb supports headDim in " "{64,128,256}. headDim=256 uses the staged profile with " - "head_dim_per_stage_kv=128 and num_insts_kv=2." + "head_dim_per_stage_kv=128." ) + # Validate constraints imposed by mixed precision KV. + if cfg.use_nvfp4_kv: + if not cfg.use_paged_kv: + raise ValueError("NVFP4 KV cache is supported only with pagedKv layout") + if cfg.headdim % 16 != 0: + raise ValueError("NVFP4 KV cache requires headdim to be a multiple of 16") def normalize_qkv_layout(qkv_layout: str) -> str: @@ -3798,12 +5020,29 @@ def normalize_qkv_layout(qkv_layout: str) -> str: def validate_page_size(num_tokens_per_page: int) -> None: """Validate a paged-KV page size against supported tile shapes.""" - if num_tokens_per_page not in (16, 32, 64, 128): - raise ValueError("num_tokens_per_page must be one of 16, 32, 64, or 128") + if num_tokens_per_page not in (4, 8, 16, 32, 64, 128): + raise ValueError("num_tokens_per_page must be one of 4, 8, 16, 32, 64, or 128") if 128 % num_tokens_per_page != 0: raise ValueError("num_tokens_per_page must divide the 128-token KV tile") +def validate_storage_page_size( + num_tokens_per_page: int, + storage_tokens_per_page: int, +) -> None: + """Validate native page-4 subpage-locator storage geometry.""" + if isinstance(storage_tokens_per_page, bool) or not isinstance( + storage_tokens_per_page, int + ): + raise TypeError("storage_tokens_per_page must be an integer") + if storage_tokens_per_page <= 0: + raise ValueError("storage_tokens_per_page must be positive") + if storage_tokens_per_page % num_tokens_per_page != 0: + raise ValueError( + "storage_tokens_per_page must be divisible by num_tokens_per_page" + ) + + def make_decode_config( headdim: int = 128, args: object | None = None, @@ -3813,18 +5052,25 @@ def make_decode_config( batch_size: int | None = None, num_heads_q: int | None = None, num_heads_kv: int | None = None, - qkv_dtype: type = Float16, + q_dtype: type | None = None, + k_dtype: type | None = None, + v_dtype: type | None = None, + qkv_dtype: type | None = None, + kv_dtype: type | None = None, o_dtype: type = Float16, qkv_layout: str = "contiguousKv", num_tokens_per_page: int = 32, + storage_tokens_per_page: int | None = None, split_kv_mode: str = "disabled", splits_kv: int = -1, max_splits_kv: int | None = None, + min_loop_iters_per_split: int = MIN_LOOP_ITERS_PER_SPLIT, sliding_window_causal: bool = False, attention_window_size: int = 0, mask_type: str | None = None, use_attention_sinks: bool = False, auto_tuner: bool = True, + split_kv: bool = True, ) -> FmhaDecodeConfig: """Build the static decode kernel config and apply auto-selection policy. @@ -3833,6 +5079,11 @@ def make_decode_config( applies static profile defaults. When launch-shape inputs are supplied, it additionally runs the decode kernel-selection workflow below. + ``q_dtype``, ``k_dtype``, and ``v_dtype`` select independent storage types. + The legacy ``qkv_dtype`` sets their common default, while ``kv_dtype`` + sets a common K/V default. Mixed-profile scheduling derives from the + effective dtypes before Q-tile or launch-mode selection. + Workflow: 1. Create a default ``FmhaDecodeConfig`` and apply caller-supplied config fields directly onto it, remembering which profile fields were explicit. @@ -3850,12 +5101,14 @@ def make_decode_config( KV128. The final KV width re-derives launch policy instead of reusing a fanout scored for KV128. Explicit policies remain caller-controlled. TileQ128 remains automatic over TileQ64 only for staged D256. SQ1 is - outside this grouped-Q cost model and therefore remains KV128. + outside this cost model and remains KV128, but ratios above 32 select + the smallest qualified Q64/Q128 Keeps tile. Other profile families use + exact-width ungrouped Keeps or ungrouped Swaps head bands as a fallback. 3. Shapes outside that qualified Q/launch selector retain the general launch - policy: under-filled fixed-Q long-sequence grids use split-KV GMEM + policy: under-filled long-sequence grids may use split-KV GMEM reduction, direct grids above one resident wave use persistent - scheduling, and the rest stay static. Packed-Q and sliding-window grids - remain nonsplit but use the same structural persistence boundary. An + scheduling, and the rest stay static. Query packing does not disable + splits; automatic sliding-window splits remain unqualified. An unsupported automatic mode falls back to direct. 4. If split-KV is selected or requested, compute the split fanout from the effective KV length, requested split count, max split cap, SM count, and @@ -3869,15 +5122,49 @@ def make_decode_config( Attention-sink paths skip automatic launch-mode selection. 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 + to nonsplit recipes without disabling Q-tile or persistence selection. + It is independent of fixed/packed Q. Pass False for prefill and True for + decode. Explicit split modes/fanouts conflict with False. The resulting + ``cfg.use_split_kv`` describes the selected launch, not this permission. """ + if not isinstance(split_kv, bool): + raise TypeError("split_kv must be a bool") + if not split_kv and (split_kv_mode != "disabled" or splits_kv > 1): + raise ValueError("split_kv=False conflicts with an explicit split mode/fanout") + dtype_keywords_supplied = any( + dtype is not None for dtype in (q_dtype, k_dtype, v_dtype, qkv_dtype, kv_dtype) + ) + q_dtype, k_dtype, v_dtype = _resolve_decode_dtype_aliases( + q_dtype=q_dtype, + k_dtype=k_dtype, + v_dtype=v_dtype, + qkv_dtype=qkv_dtype, + kv_dtype=kv_dtype, + ) shape_values = (seq_len_kv, batch_size, num_heads_q, num_heads_kv) if all(value is None for value in shape_values): - return _make_static_decode_config( + if dtype_keywords_supplied: + args = ( + { + "q_dtype": q_dtype, + "k_dtype": k_dtype, + "v_dtype": v_dtype, + "out_dtype": o_dtype, + }, + args, + ) + cfg = _make_static_decode_config( headdim, args, mask_type=mask_type, sliding_window_causal=sliding_window_causal, ) + if not split_kv and cfg.use_split_kv: + raise ValueError("split_kv=False conflicts with use_split_kv=True") + return cfg if any(value is None for value in shape_values): raise ValueError( "seq_len_kv, batch_size, num_heads_q, and num_heads_kv are all " @@ -3887,6 +5174,8 @@ 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) + 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( cfg, explicit_fields=explicit_fields, @@ -3913,17 +5202,23 @@ def make_decode_config( _apply_default_q_grouping( cfg, explicit_fields=explicit_fields, + heads_q_per_kv=num_heads_q // num_heads_kv, ) - cfg.q_dtype = qkv_dtype - cfg.kv_dtype = qkv_dtype - cfg.out_dtype = o_dtype - qkv_layout = _apply_layout_config( cfg, qkv_layout=qkv_layout, num_tokens_per_page=num_tokens_per_page, + storage_tokens_per_page=storage_tokens_per_page, seq_len_kv=seq_len_kv, ) + _apply_dtype_config( + cfg, + explicit_fields=explicit_fields, + q_dtype=q_dtype, + k_dtype=k_dtype, + v_dtype=v_dtype, + o_dtype=o_dtype, + ) _apply_feature_config( cfg, explicit_fields=explicit_fields, @@ -3954,14 +5249,32 @@ def make_decode_config( num_heads_kv=num_heads_kv, splits_kv=splits_kv, max_splits_kv=max_splits_kv, + split_kv=split_kv, ) if selected_grouped_q_recipe is None: - _apply_swaps_tile_config( + selected_wide_keeps = _try_apply_default_wide_keeps_config( cfg, explicit_fields=explicit_fields, + seq_len_q=seq_len_q, num_heads_q=num_heads_q, num_heads_kv=num_heads_kv, ) + if not selected_wide_keeps: + heads_q_per_kv = num_heads_q // num_heads_kv + if ( + heads_q_per_kv > 32 + and "groups_tokens_heads_q" not in explicit_fields + and not (_MMA_SELECTION_FIELDS & auto_selection_explicit_fields) + ): + # Profiles outside the grouped-Keeps matrix remain valid via + # the established ungrouped Swaps head bands. + cfg.groups_tokens_heads_q = False + _apply_swaps_tile_config( + cfg, + explicit_fields=explicit_fields, + num_heads_q=num_heads_q, + num_heads_kv=num_heads_kv, + ) kv_tile_was_promoted = _try_apply_auto_kv256_profile( cfg, q_candidate=( @@ -3983,6 +5296,7 @@ def make_decode_config( # below from the final KV256 work granularity. cfg.use_persistent_scheduler = False selected_grouped_q_recipe = None + _finalize_static_decode_config(cfg, explicit_fields) if not cfg.use_variable_seqlens_q: validate_causal_decode_lengths( @@ -3992,8 +5306,8 @@ def make_decode_config( ) # Auto launch-mode selection. Only kicks in if the caller has not already - # opted into a specific mode. Packed-Q and sliding-window shapes may select - # CLC persistence, but remain nonsplit. + # 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.splits_kv > 1: split_kv_mode = selected_grouped_q_recipe.split_kv_mode @@ -4012,6 +5326,7 @@ def make_decode_config( num_heads_kv=num_heads_kv, seq_len_kv=seq_len_kv, seq_len_q=seq_len_q, + split_kv=split_kv, ) reduction_mode_preconfigured = ( @@ -4022,7 +5337,6 @@ def make_decode_config( and launch_mode_was_auto and split_kv_mode == "gmem_reduction" and not reduction_mode_preconfigured - and not cfg.use_variable_seqlens_q and not cfg.use_sliding_window_causal and not cfg.use_attention_sinks ) @@ -4038,6 +5352,7 @@ def make_decode_config( max_splits_kv=max_splits_kv, sliding_window_causal=sliding_window_causal, attention_window_size=attention_window_size, + min_loop_iters_per_split=min_loop_iters_per_split, ) if auto_split_kv_selected and cfg.use_split_kv: @@ -4052,6 +5367,7 @@ def make_decode_config( ), ) + _validate_mixed_kv_dtype_profile(cfg) _validate_kv256_static_config(cfg) _finalize_warp_roles(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 8bf94760cf31..843b7ffb44b0 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 @@ -12,12 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Integer constants shared by the FMHA decode TS implementation. +"""Constants shared by the FMHA decode TS implementation. -Keep non-obvious integer constants here with their rationale so config, -resource, and reduction code can use named values without duplicating comments. +Keep non-obvious constants here with their rationale so config, resource, and +reduction code can use named values without duplicating comments. """ +import math + # B200 has 148 SMs. Use this only when the runtime SM query is unavailable, # so auto split-KV selection remains deterministic in offline/test flows. FALLBACK_SM_COUNT_B200 = 148 @@ -28,18 +30,27 @@ TOTAL_SMEM_BUDGET_KIB = 218 MAX_KV_STAGE_SMEM_KIB = 144 BYTES_PER_KIB = 1024 - -# The supported BF16 M64N256 profile uses a 64-KiB shared K/V stage. Three -# stages occupy 192 KiB; its 16-KiB Q stage and small metadata/barrier -# allocations fit in the remaining SM100 budget. Persistent direct-output tail -# correction rotates a compact 35,840-byte exchange payload over one drained -# 64-KiB stage in this ring; split-KV keeps its fixed full exchange allocation. -# Keep this exact-profile override separate from the conservative, -# topology-independent MAX_KV_STAGE_SMEM_KIB inference above. +# Dynamic SMEM one CTA can use on the SM100 family: 228 KiB per SM minus the +# 1 KiB the system reserves. The complete decode layout (pipelines, barriers, +# metadata and scratch) must fit it; TOTAL_SMEM_BUDGET_KIB sizes the Q and +# K/V pipelines below it. +SMEM_CAPACITY_KIB = 227 + +# The M64N256 profile stages one 256-row K or V tile per shared-ring slot. A +# 16-bit tile is 64 KiB, so three stages, the 16-KiB Q stage, the tail +# exchange and the metadata fill the SM100 carveout. A byte-wide tile is +# 32 KiB; with four stages every K load takes the slot the previous QK freed +# instead of waiting for a V slot that its tile's PV still holds. Keep both +# exact-profile depths separate from the conservative, topology-independent +# MAX_KV_STAGE_SMEM_KIB inference above. KV_TILE_256_SHARED_FIFO_STAGES = 3 +KV_TILE_256_BYTE_WIDE_SHARED_FIFO_STAGES = 4 -# The four semantic K64 atoms are stored in the physical K slots consumed by -# the two interleaved QK instructions in this order. +# The four semantic K64 atoms of a KV256 tile are staged in these physical K +# slots. On the WS 2x2 datapath each TMEM lane half computes 128 contiguous QK +# columns, while PV pairs KV64 blocks 2j and 2j+1 per K step; staging K as +# (0, 2, 1, 3) lets P alias S unchanged with V in natural order. The score +# column mapping undoes the permutation for masking and Sage scales. KV_TILE_256_K_SLOT_FOR_SEMANTIC_ATOM = (0, 2, 1, 3) # Keep the old maximum as the exponent reference while a new maximum is at @@ -84,8 +95,27 @@ # overhead does not dominate tiny per-split K ranges. MIN_LOOP_ITERS_PER_SPLIT = 2 -# The maximum number of warp groups per CTA. -MAX_WARP_GROUPS = 4 +# Grouped sparse attention may add two producer-only groups to the common +# four-warpgroup decode topology. +MAX_WARP_GROUPS = 6 + +# Grouped sparse routes keep one membership byte per selected page in a table +# separate from the plain Int32 locators. Four bytes share one Int32 word. Q1 does +# not consume membership; Q2--Q8 use the packed table for their per-query +# softmax mask. A byte supports all eight query positions without changing the ABI. +Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIP_BITS = 8 +Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIP_MASK = ( + 1 << Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIP_BITS +) - 1 +Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIPS_PER_WORD = ( + 32 // Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIP_BITS +) +# Holding a complete locator window removes duplicate K/V metadata traffic, +# but a graph-safe Q2/Q4 worst-case bound can otherwise reserve 8 KiB or more +# of SMEM for rows that commonly use only three KV tiles. Beyond this crossover, +# stage locators one tile at a time while retaining the compact packed +# membership row separately. +Q_TOKEN_KV_BLOCK_SPARSE_HELD_LOCATOR_MAX_TILES = 8 # Default used only when the launch helper has no resolved decode config. # Config-aware FMHA and block-sparse callers pass their selected KV tile. @@ -113,6 +143,7 @@ # TMEM column layout for staged SwapsMmaAb O. A TMEM row holds 256 columns, and # the tcgen05 descriptor encodes a 16-row jump in the high 16 bits. TMEM_COLUMNS_PER_ROW = 256 +TMEM_MAX_ALLOCATION_COLUMNS = 512 TMEM_ROW_STRIDE = 16 << 16 # Number of scalar softmax/output values packed in one register for the @@ -121,6 +152,20 @@ FP8_VALUES_PER_REG = 4 FP16_VALUES_PER_REG = 2 +# Bytes of the K-major operand row one tcgen05 MMA instruction consumes: 32 +# one-byte or 16 two-byte K elements. +MMA_K_STEP_BYTES = 32 + +# FP8 probabilities are quantized as 448 * p (the E4M3 maximum); row sums and +# sink terms share the scale, and the output normalization divides it out. +FP8_P_QUANT_SCALE = 448.0 +FP8_P_QUANT_LOG2_SCALE = math.log2(FP8_P_QUANT_SCALE) + +# tcgen05 SMEM descriptor geometry: address offsets count 16-byte units and a +# 128-byte swizzle atom spans one 128-byte row per K or MN index. +SMEM_DESC_UNIT_BYTES = 16 +SWIZZLE_128B_ROW_BYTES = 128 + # Per-lane register ownership denominators for packed output fragments. FP8_OUTPUT_ELEMENTS_PER_REG_GROUP = 512 FP16_OUTPUT_ELEMENTS_PER_REG_GROUP = 256 @@ -154,3 +199,19 @@ 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 d2540ec61920..6e03c2a0c12b 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 @@ -34,12 +34,13 @@ import cutlass import cutlass.experimental.cuda as cuda import cutlass.cute as cute +from .direct_sparse_metadata import DirectSparseMetadataView, HeadIndexedMetadataView import cutlass.pipeline as pipeline import cutlass.utils as utils from cuda.bindings import driver as cuda_drv from cutlass import Float32, Int32, Int64 from cutlass.experimental import primitives as prims -from cutlass.experimental.task_scheduling.enums import PipelineType +from cutlass.experimental.task_scheduling.enums import PipelineType, SignalingThreads from cutlass.experimental.task_scheduling.memory import ( ResourceContext, SmemAllocation, @@ -54,6 +55,8 @@ ) from cutlass.experimental.task_scheduling.task import Task from cutlass.experimental.task_scheduling.task_manager import TaskManager +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import T as mlir_T from ..._block_sparse.common import ( _block_sparse_contiguous_kv_copy_geometry, @@ -67,9 +70,12 @@ ) from .fmha_decode_config import FmhaDecodeConfig from .fmha_decode_constants import ( + BYTES_PER_KIB, KV_KIND_K, KV_KIND_V, KV_TILE_256_REGISTER_REALLOCATION_MIN_TILES, + MAX_KV_STAGE_SMEM_KIB, + SMEM_CAPACITY_KIB, ) from .fmha_decode_resources import ( SmemBlockSparseKvMetadataResource, @@ -78,6 +84,7 @@ SmemKvResource, SmemPageOffsetsKvResource, SmemQResource, + SmemTransformedKvResource, TmemCorrResource, TmemOResource, SmemPResource, @@ -86,15 +93,24 @@ TmemSoftmaxGlobalResource, TmemSoftmaxLocalResource, TmemSoftmaxOrderResource, + TmemTransformedKvResource, ) from .fmha_decode_resources.helpers_common import ( _q_group_token_base, _q_seq_bounds, ) from .fmha_decode_resources.helpers_kv_tile_idx import _runtime_active_splits_kv +from .fmha_decode_resources.sage_scales import ( + SAGE_K_SCALES_RING_STAGES, + SageKScalesResource, + SageScaleTensors, + SageVScalesResource, +) from .fmha_decode_tasks import ( PackedDecodeWorkQueue, ScheduleTokenThrottleResource, + SparseMembershipLifetimeResource, + _can_hold_native_page_window, _prefetch_prepared_sparse_row, create_block_sparse_load_tasks_per_inst, create_correction_task, @@ -112,6 +128,7 @@ create_scheduler_task, create_softmax0_task, create_softmax1_task, + create_transform_kv_task, ) from .reduction import ( # noqa: F401 decode_gen_separate_reduction_kernel, @@ -119,6 +136,10 @@ ) _PERSISTENT_SCHEDULE_TOKEN_STAGES = 2 +# Response slots of the CLC work-queue fetch pipeline. Every role consumes +# the queue one tile behind the scheduler warp, so a second slot adds no +# lookahead, only a loop-carried stage index that spills to local memory. +_WORK_QUEUE_STAGES = 1 def _block_sparse_bshd_tma_strides( @@ -128,6 +149,7 @@ def _block_sparse_bshd_tma_strides( h_k: cutlass.Integer | int, s_k: cutlass.Integer | int, d: cutlass.Integer | int, + element_bytes: int, ) -> tuple[ tuple[cutlass.Integer | int, ...], tuple[cutlass.Integer | int, ...], @@ -135,13 +157,13 @@ def _block_sparse_bshd_tma_strides( """Build BSHD TensorMap strides in 16-byte units using Int64 math. The raw TensorMap API omits the implicit contiguous stride and takes the - remaining strides in 16-byte units. Block-sparse attention supports only - 16-bit Q/K/V with headDim=128, so one stride unit contains eight elements. - Keep every returned value in Int64: the outer batch stride can exceed the + remaining strides in 16-byte units, so one unit holds ``16 / + element_bytes`` elements of the 16-bit or 8-bit Q/K/V (headDim=128). Keep + every returned value in Int64: the outer batch stride can exceed the signed Int32 range even though each public tensor dimension is Int32. """ - elements_per_stride_unit = 8 + elements_per_stride_unit = 16 // element_bytes d_units = Int64(d // elements_per_stride_unit) h_r = h_q // h_k return ( @@ -244,7 +266,8 @@ def _decode_min_blocks_per_mp(cfg: FmhaDecodeConfig, seq_len_kv: int) -> int: >= KV_TILE_256_REGISTER_REALLOCATION_MIN_TILES ) return int( - kv256_reallocation + cfg.use_transform_kv + or kv256_reallocation or cfg.tile_size_q == 8 or (cfg.tile_size_q == 16 and cfg.q_dtype_bytes == 1) or (cfg.use_keeps_mma_ab and cfg.tile_size_q == 128) @@ -319,6 +342,8 @@ def _build_decode_gen_schedule( tma_desc_q: cutlass.Pointer | None = None, tma_desc_k: cutlass.Pointer | None = None, tma_desc_v: cutlass.Pointer | None = None, + tma_desc_k_sf: cutlass.Pointer | None = None, + tma_desc_v_sf: cutlass.Pointer | None = None, tma_desc_k_atom: cutlass.Pointer | None = None, tma_desc_v_atom: cutlass.Pointer | None = None, tma_desc_k_summary: cutlass.Pointer | None = None, @@ -326,6 +351,10 @@ def _build_decode_gen_schedule( tma_desc_k_summary_atom: cutlass.Pointer | None = None, tma_desc_v_summary_atom: cutlass.Pointer | None = None, page_idx_kv: cute.Pointer | None = None, + page_table_stride: Int64 | None = None, + page_table_capacity: Int32 | None = None, + q_token_kv_block_sparse_page_memberships: cute.Pointer | None = None, + q_token_kv_block_sparse_page_membership_stride: Int32 | None = None, h_k_idx: Int32 | None = None, b_idx: Int32 | None = None, q_group_idx: Int32 | None = None, @@ -334,18 +363,22 @@ def _build_decode_gen_schedule( active_splits_kv: Int32 | None = None, static_full_split_prefix: bool = False, tile_sched_params: utils.ClcDynamicPersistentTileSchedulerParams | None = None, - clc_response_ptr: cute.Pointer | None = None, use_variable_seqlens_kv: bool = False, use_native_paged_kv: bool = False, use_static_native_seqlens_kv: bool = False, - block_tables: cute.Pointer | None = None, - block_table_capacity: Int32 | None = None, - block_table_row_stride: Int64 | None = None, sparse_row_route_offsets: cute.Pointer | None = None, sparse_row_route_counts: cute.Pointer | None = None, sparse_route_metadata: cute.Pointer | None = None, sparse_row_route_begin: Int32 | None = None, sparse_route_count: Int32 | None = None, + sage_q_scale_ptr: cute.Pointer | None = None, + sage_k_scale_ptr: cute.Pointer | None = None, + sage_k_summary_scale_ptr: cute.Pointer | None = None, + sage_v_scale_ptr: cute.Pointer | None = None, + sage_v_mean_ptr: cute.Pointer | None = None, + sage_q_scale_head_stride: Int32 | None = None, + sage_k_scale_head_stride: Int32 | None = None, + sage_k_summary_scale_head_stride: Int32 | None = None, ) -> tuple[ list[Task], dict[MemoryResource, list[MemoryResource]], @@ -375,6 +408,8 @@ def _build_decode_gen_schedule( (task_list, dep_graph, dma labels, smem_allocator, tmem_allocator, eager_init_resources) """ + if cfg.use_block_sparse and cfg.num_insts_kv != 2: + raise ValueError("block-sparse attention requires num_insts_kv == 2") if cfg.use_keeps_mma_ab and cfg.num_insts_kv == 1 and not cfg.uses_tmem_p: raise ValueError( "one-instance KeepsMmaAb is enabled only for the staged headDim=256 " @@ -383,13 +418,21 @@ def _build_decode_gen_schedule( if use_native_paged_kv and not cfg.use_paged_kv: raise ValueError("native paged-KV ABI requires cfg.use_paged_kv=True") if use_native_paged_kv and ( - block_tables is None - or block_table_capacity is None - or block_table_row_stride is None + page_idx_kv is None or page_table_capacity is None or page_table_stride is None ): raise ValueError( "native paged-KV ABI requires block tables, capacity, and row stride" ) + if cfg.uses_scattered_page_route and not use_native_paged_kv: + raise ValueError("scattered page routes require the native paged-KV ABI") + if ( + cfg.uses_q_token_kv_block_sparse_page_membership + and tma_desc_q is not None + and q_token_kv_block_sparse_page_memberships is None + ): + raise ValueError( + "grouped QToken-KvBlock-Sparse-Attention kernel construction requires q_token_kv_block_sparse_page_memberships" + ) if cfg.use_paged_kv: cfg.validate_paged_kv_staging_config() if cfg.use_block_sparse: @@ -456,6 +499,10 @@ def _build_decode_gen_schedule( Agent.Thread, cfg.page_offsets_num_warps * WARP_SIZE ) load_grp = pipeline.CooperativeGroup(Agent.Thread, cfg.load_num_warps * WARP_SIZE) + if cfg.use_transform_kv: + transform_kv_grp = pipeline.CooperativeGroup( + Agent.Thread, cfg.transform_kv_num_warps * WARP_SIZE + ) umma_hw = pipeline.CooperativeGroup(Agent.Thread) # The staged one-instance S/P overlay uses this group for overwrite credit. mma_grp = pipeline.CooperativeGroup(Agent.Thread, cfg.mma_num_warps * WARP_SIZE) @@ -471,6 +518,14 @@ def _build_decode_gen_schedule( scheduler_grp = pipeline.CooperativeGroup( Agent.Thread, cfg.scheduler_num_warps * WARP_SIZE ) + tma_producer_signaling = ( + SignalingThreads.TaskWarpLeader + if cfg.load_num_warps > 1 + else SignalingThreads.All + ) + q_bytes_per_load_warp = ( + cfg.smem_q_tile_bytes // cfg.load_num_warps if cfg.load_num_warps > 1 else None + ) # ------------------------------------------------------------------ # Pipeline configs @@ -479,8 +534,15 @@ def _build_decode_gen_schedule( # into the unified block. Separate barrier arrays create an alignment gap # before the 1024-byte-aligned data block and overflow near-capacity Q128. use_paged_kv = cfg.use_paged_kv + use_one_inst_kv = cfg.num_insts_kv == 1 use_dense_page_offsets = use_paged_kv and not cfg.use_block_sparse use_one_inst_qkv = cfg.use_keeps_mma_ab and cfg.num_insts_kv == 1 + # K and V of different byte widths take the split-resource paths with one + # K ring and one V ring, each shared by both K/V instances; equal widths + # (including Int8 K with E4M3 V) share one ring. + use_shared_inst_kv_rings = ( + cfg.k_dtype_bytes != cfg.v_dtype_bytes and cfg.tile_size_kv != 256 + ) one_inst_tmem_stages = 2 if use_one_inst_qkv else 1 one_inst_kv_stages = cfg.num_head_dim_stages_kv if use_one_inst_qkv else 1 use_distributed_split_kv_stages = not use_one_inst_qkv @@ -526,25 +588,39 @@ def _build_decode_gen_schedule( if use_one_inst_qkv else max((split_total_v_stages + cfg.num_insts_kv - 2) // cfg.num_insts_kv, 1) ) + if use_shared_inst_kv_rings and not use_one_inst_qkv: + # One K and one V stage per unit of depth. + shared_inst_kv_stages = max( + (MAX_KV_STAGE_SMEM_KIB * BYTES_PER_KIB) + // (cfg.smem_k_tile_bytes + cfg.smem_v_tile_bytes), + 1, + ) + split_k0_stages = shared_inst_kv_stages + split_v0_stages = shared_inst_kv_stages use_ordered_softmax_barrier = ( - not use_one_inst_qkv and cfg.uses_ordered_softmax_barrier + not use_one_inst_kv and cfg.uses_ordered_softmax_barrier ) # A two-inst Keeps profile can use the deeper shared K/V FIFO when stats # are standalone and P remains in SMEM. Keep instruction-local FIFOs when # stats or TMEM-P alias S: their overwrite-credit cadence is tied to each # instruction. Dense Swaps uses the shared FIFO, including staged H256. # Sparse KV128 keeps instruction-local rings in either MMA orientation; - # sparse KV256 reuses its only feasible three-stage shared data ring while - # retaining instruction-local route metadata. The load warp issues V(route - # R) before replacing that metadata with route R+1, so its lifetime remains - # independent of the K/V data-ring depth. + # sparse KV256 reuses the shared data ring at its element-width-derived + # depth while retaining instruction-local route metadata. The load warp + # issues V(route R) before replacing that metadata with route R+1, so its + # lifetime remains independent of the K/V data-ring depth. # With cfg.keeps_stats_via_smem the stats-alias justification no longer # applies, but the shared FIFO still causes a material Q128 regression, so # the instruction-local FIFO gate remains part of that kernel policy. - use_per_inst_kv_resources = (cfg.use_block_sparse and cfg.tile_size_kv != 256) or ( - cfg.use_keeps_mma_ab - and cfg.tile_size_kv != 256 - and (not cfg.keeps_separates_tmem_s_and_stats or cfg.uses_two_inst_tmem_p) + use_per_inst_kv_resources = ( + use_one_inst_qkv + or (cfg.use_block_sparse and cfg.tile_size_kv != 256) + or use_shared_inst_kv_rings + or ( + cfg.use_keeps_mma_ab + and cfg.tile_size_kv != 256 + and (not cfg.keeps_separates_tmem_s_and_stats or cfg.uses_two_inst_tmem_p) + ) ) # B8/B16 issue enough fine-grained TMA copies to benefit from reusing a # padding warp as a second issuer. The host policy applies one KV-side @@ -556,8 +632,18 @@ def _build_decode_gen_schedule( and cfg.kv_block_size in (8, 16) and cfg.num_insts_kv == 2 ) + # Independent K0/K1/V0/V1 data rings share identical native sparse locators. + # A held route is published once and retained through the final V issue; + # only schedules that replace locators per tile need separate K/V state. use_separate_kv_page_offset_resources = ( - use_dense_page_offsets and use_per_inst_kv_resources and not use_one_inst_qkv + use_dense_page_offsets + and use_per_inst_kv_resources + and not use_one_inst_qkv + and not ( + use_native_paged_kv + and cfg.uses_q_token_kv_block_sparse_page_route + and cfg.uses_held_encoded_locator_window + ) ) # Paired resources publish independent K0/K1 and V0/V1 stages. Shared # split-KV retains the aligned 32-ID representation for its optional @@ -572,49 +658,58 @@ def _build_decode_gen_schedule( producer_group=tma_producer, consumer_group=umma_hw, cta_layout_vmnk=cta_layout, - advance_on_wait=True, - ) - smem_kv_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( - num_stages=cfg.kv_stages, - num_bytes=cfg.smem_kv_tile_bytes, - producer_group=tma_producer, - consumer_group=umma_hw, - cta_layout_vmnk=cta_layout, - advance_on_wait=True, - ) - smem_k0_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( - num_stages=split_k0_stages, - num_bytes=cfg.smem_kv_tile_bytes, - producer_group=tma_producer, - consumer_group=umma_hw, - cta_layout_vmnk=cta_layout, - advance_on_wait=True, - ) - smem_k1_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( - num_stages=split_k1_stages, - num_bytes=cfg.smem_kv_tile_bytes, - producer_group=tma_producer, - consumer_group=umma_hw, - cta_layout_vmnk=cta_layout, - advance_on_wait=True, - ) - smem_v0_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( - num_stages=split_v0_stages, - num_bytes=cfg.smem_kv_tile_bytes, - producer_group=tma_producer, - consumer_group=umma_hw, - cta_layout_vmnk=cta_layout, - advance_on_wait=True, - ) - smem_v1_cfg = PipelineConfig.create_tma_umma_pipeline_cfg( - num_stages=split_v1_stages, - num_bytes=cfg.smem_kv_tile_bytes, - producer_group=tma_producer, - consumer_group=umma_hw, - cta_layout_vmnk=cta_layout, + producer_signaling_threads=tma_producer_signaling, + num_bytes_per_warp_per_cta=q_bytes_per_load_warp, advance_on_wait=True, ) + def _make_raw_kv_cfg(num_stages: int, tile_bytes: int) -> PipelineConfig: + """Create a per-operand raw pipeline with its actual TMA byte count.""" + num_bytes = tile_bytes + cfg.smem_kv_sf_tile_bytes + bytes_per_load_warp = ( + num_bytes // cfg.load_num_warps if cfg.load_num_warps > 1 else None + ) + if cfg.use_transform_kv: + # NVFP4 payload and scale factors complete the same raw barrier. + return PipelineConfig( + num_stages=num_stages, + num_bytes=num_bytes, + producer_group=tma_producer, + consumer_group=transform_kv_grp, + pipeline_type=PipelineType.TmaAsync, + cta_layout_vmnk=cta_layout, + producer_signaling_threads=tma_producer_signaling, + num_bytes_per_warp_per_cta=bytes_per_load_warp, + advance_on_wait=True, + ) + return PipelineConfig.create_tma_umma_pipeline_cfg( + num_stages=num_stages, + num_bytes=num_bytes, + producer_group=tma_producer, + consumer_group=umma_hw, + cta_layout_vmnk=cta_layout, + producer_signaling_threads=tma_producer_signaling, + num_bytes_per_warp_per_cta=bytes_per_load_warp, + advance_on_wait=True, + ) + + smem_kv_cfg = None + if not use_per_inst_kv_resources: + smem_kv_cfg = _make_raw_kv_cfg(cfg.kv_stages, cfg.smem_kv_tile_bytes) + smem_transformed_kv_cfg = None + if cfg.use_transform_kv: + smem_transformed_kv_cfg = PipelineConfig.create_async_umma_pipeline_cfg( + num_stages=cfg.transformed_kv_stages, + producer_group=transform_kv_grp, + consumer_group=umma_hw, + cta_layout_vmnk=cta_layout, + advance_on_wait=True, + ) + smem_k0_cfg = _make_raw_kv_cfg(split_k0_stages, cfg.smem_k_tile_bytes) + smem_k1_cfg = _make_raw_kv_cfg(split_k1_stages, cfg.smem_k_tile_bytes) + smem_v0_cfg = _make_raw_kv_cfg(split_v0_stages, cfg.smem_v_tile_bytes) + smem_v1_cfg = _make_raw_kv_cfg(split_v1_stages, cfg.smem_v_tile_bytes) + def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: """Create the async page-offsets pipeline for the selected stage count.""" if num_stages is None: @@ -635,6 +730,11 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: page_offsets_stages = ( 3 if use_separate_kv_page_offset_resources else cfg.page_offsets_stages ) + if use_native_paged_kv and cfg.uses_held_encoded_locator_window: + # QToken-KvBlock-Sparse-Attention publishes its complete CTA-local locator span once per work + # tile and holds it through the K/V tail. Only one stage is live + # for direct, persistent, and split-KV routes alike. + page_offsets_stages = 1 smem_page_offsets_cfg = _make_page_offsets_cfg(page_offsets_stages) if use_separate_kv_page_offset_resources: smem_page_offsets_v_cfg = _make_page_offsets_cfg(page_offsets_stages) @@ -775,9 +875,9 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: per_inst_block_sparse_load_topology is not None ) if use_clc_dynamic: - num_consumer_threads = 16 * WARP_SIZE + num_consumer_threads = cfg.threads_per_cta wq_pipeline_config = PipelineConfig.create_clc_fetch_async_pipeline_cfg( - num_stages=_PERSISTENT_SCHEDULE_TOKEN_STAGES, + num_stages=_WORK_QUEUE_STAGES, num_bytes=16, producer_group=pipeline.CooperativeGroup(Agent.Thread), consumer_group=pipeline.CooperativeGroup( @@ -785,13 +885,10 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: ), cta_layout_vmnk=cta_layout, ) + # The scheduler config needs the CLC response slots' address; it is + # bound once the unified SMEM layout below is computed. work_queue_kwargs = { - "tile_scheduler_config": ( - TileSchedulerConfig.create_clc_dynamic_persistent_tile_scheduler_params( - tile_scheduler_params=tile_sched_params, - response_ptr=clc_response_ptr, - ) - ), + "tile_scheduler_config": None, "pipeline_config": wq_pipeline_config, "name": "work_queue", } @@ -803,6 +900,7 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: ) else: work_queue = WorkQueue(**work_queue_kwargs) + schedule_token_throttle = ScheduleTokenThrottleResource( pipeline_config=PipelineConfig.create_async_async_pipeline_cfg( num_stages=_PERSISTENT_SCHEDULE_TOKEN_STAGES, @@ -812,6 +910,30 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: ), name="schedule_token_throttle", ) + membership_lifetime = None + if use_clc_dynamic and cfg.uses_q_token_kv_block_sparse_page_membership: + # Page IDs are consumed by Load, but membership bytes are consumed by + # Softmax. Do not overwrite the next work item's membership row until + # every score stream has finished reading the current one. + membership_lifetime = SparseMembershipLifetimeResource( + name="sparseMembershipLifetime", + pipeline_config=PipelineConfig( + num_stages=1, + num_bytes=0, + producer_group=page_offsets_grp, + consumer_group=pipeline.CooperativeGroup( + Agent.Thread, + WARP_SIZE + * ( + cfg.softmax0_num_warps + + (0 if use_one_inst_kv else cfg.softmax1_num_warps) + ), + ), + pipeline_type=PipelineType.AsyncAsync, + cta_layout_vmnk=cta_layout, + advance_on_wait=True, + ), + ) smem_q = SmemQResource( pipeline_config=smem_q_cfg, cfg=cfg, @@ -823,9 +945,9 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: seq_len_q=seq_len_q, name="smemQ", ) - # Native callers provide one canonical sequence-length tensor alongside - # their fixed page table, so native mode reuses the existing variable-length - # domain, split, sliding-window, and masking paths. + # Native dense page tables pair fixed-capacity rows with one canonical + # sequence-length tensor, so native mode reuses the existing + # variable-length domain, split, sliding-window, and masking paths. use_runtime_seqlens_kv = use_variable_seqlens_kv or ( use_native_paged_kv and not use_static_native_seqlens_kv ) @@ -838,10 +960,12 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: cfg=cfg, stage_page_ids_per_tile=stage_page_ids_per_tile, page_idx_kv=page_idx_kv, + page_table_stride=page_table_stride, + num_heads_kv=num_heads_kv, + q_token_kv_block_sparse_page_memberships=q_token_kv_block_sparse_page_memberships, + q_token_kv_block_sparse_page_membership_stride=q_token_kv_block_sparse_page_membership_stride, seqlens_kv=kv_seqlens, use_native_paged_kv=use_native_paged_kv, - block_tables=block_tables, - block_table_row_stride=block_table_row_stride, max_seq_len_kv=max_seq_len_kv, h_k_idx=h_k_idx, b_idx=b_idx, @@ -858,11 +982,15 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: pipeline_config=smem_page_offsets_v_cfg, cfg=cfg, stage_page_ids_per_tile=stage_page_ids_per_tile, + # Softmax consumes only the K-side membership view. + cache_memberships_in_smem=False, page_idx_kv=page_idx_kv, + page_table_stride=page_table_stride, + num_heads_kv=num_heads_kv, + q_token_kv_block_sparse_page_memberships=q_token_kv_block_sparse_page_memberships, + q_token_kv_block_sparse_page_membership_stride=q_token_kv_block_sparse_page_membership_stride, seqlens_kv=kv_seqlens, use_native_paged_kv=use_native_paged_kv, - block_tables=block_tables, - block_table_row_stride=block_table_row_stride, max_seq_len_kv=max_seq_len_kv, h_k_idx=h_k_idx, b_idx=b_idx, @@ -874,6 +1002,18 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: sparse_kv_metadata1 = None sparse_softmax_metadata0 = None sparse_softmax_metadata1 = None + sage_scale_tensors = None + if cfg.use_sage_attention: + sage_scale_tensors = SageScaleTensors( + q_scale_ptr=sage_q_scale_ptr, + q_scale_head_stride=sage_q_scale_head_stride, + k_scale_ptr=sage_k_scale_ptr, + k_scale_head_stride=sage_k_scale_head_stride, + k_summary_scale_ptr=sage_k_summary_scale_ptr, + k_summary_scale_head_stride=sage_k_summary_scale_head_stride, + v_scale_ptr=sage_v_scale_ptr, + v_mean_ptr=sage_v_mean_ptr, + ) if cfg.use_block_sparse: # This selects the prepared-record storage ABI. Causal consumers still # intersect these column-validity words with each Q row's causal mask. @@ -910,6 +1050,9 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: route_metadata=sparse_route_metadata, route_layout=prepared_route_layout, name="smemBlockSparseSoftmaxMetadata0", + h_k_idx=h_k_idx, + b_idx=b_idx, + scale_tensors=sage_scale_tensors, ) sparse_softmax_metadata1 = SmemBlockSparseSoftmaxMetadataResource( pipeline_config=sparse_softmax_metadata1_cfg, @@ -918,6 +1061,9 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: route_metadata=sparse_route_metadata, route_layout=prepared_route_layout, name="smemBlockSparseSoftmaxMetadata1", + h_k_idx=h_k_idx, + b_idx=b_idx, + scale_tensors=sage_scale_tensors, ) smem_kv = None smem_k0 = None @@ -930,6 +1076,8 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: cfg=cfg, tma_desc_k=tma_desc_k, tma_desc_v=tma_desc_v, + tma_desc_k_sf=tma_desc_k_sf, + tma_desc_v_sf=tma_desc_v_sf, tma_desc_k_atom=tma_desc_k_atom, tma_desc_v_atom=tma_desc_v_atom, tma_desc_k_summary=tma_desc_k_summary, @@ -948,34 +1096,39 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: kv_kind=KV_KIND_K, name="smemK0", ) - smem_k1 = SmemKvTileResource( - pipeline_config=smem_k1_cfg, - cfg=cfg, - tma_desc_k=tma_desc_k, - tma_desc_v=tma_desc_v, - tma_desc_k_atom=tma_desc_k_atom, - tma_desc_v_atom=tma_desc_v_atom, - tma_desc_k_summary=tma_desc_k_summary, - tma_desc_v_summary=tma_desc_v_summary, - tma_desc_k_summary_atom=tma_desc_k_summary_atom, - tma_desc_v_summary_atom=tma_desc_v_summary_atom, - sparse_kv_metadata=sparse_kv_metadata1, - page_offsets_kv=smem_page_offsets, - seqlens_kv=kv_seqlens, - max_seq_len_kv=max_seq_len_kv, - h_k_idx=h_k_idx, - b_idx=b_idx, - q_group_idx=q_group_idx, - seq_len_q=seq_len_q, - inst_id=1, - kv_kind=KV_KIND_K, - name="smemK1", - ) + if not use_shared_inst_kv_rings: + smem_k1 = SmemKvTileResource( + pipeline_config=smem_k1_cfg, + cfg=cfg, + tma_desc_k=tma_desc_k, + tma_desc_v=tma_desc_v, + tma_desc_k_sf=tma_desc_k_sf, + tma_desc_v_sf=tma_desc_v_sf, + tma_desc_k_atom=tma_desc_k_atom, + tma_desc_v_atom=tma_desc_v_atom, + tma_desc_k_summary=tma_desc_k_summary, + tma_desc_v_summary=tma_desc_v_summary, + tma_desc_k_summary_atom=tma_desc_k_summary_atom, + tma_desc_v_summary_atom=tma_desc_v_summary_atom, + sparse_kv_metadata=sparse_kv_metadata1, + page_offsets_kv=smem_page_offsets, + seqlens_kv=kv_seqlens, + max_seq_len_kv=max_seq_len_kv, + h_k_idx=h_k_idx, + b_idx=b_idx, + q_group_idx=q_group_idx, + seq_len_q=seq_len_q, + inst_id=1, + kv_kind=KV_KIND_K, + name="smemK1", + ) smem_v0 = SmemKvTileResource( pipeline_config=smem_v0_cfg, cfg=cfg, tma_desc_k=tma_desc_k, tma_desc_v=tma_desc_v, + tma_desc_k_sf=tma_desc_k_sf, + tma_desc_v_sf=tma_desc_v_sf, tma_desc_k_atom=tma_desc_k_atom, tma_desc_v_atom=tma_desc_v_atom, tma_desc_k_summary=tma_desc_k_summary, @@ -994,52 +1147,185 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: kv_kind=KV_KIND_V, name="smemV0", ) - smem_v1 = SmemKvTileResource( - pipeline_config=smem_v1_cfg, + if not use_shared_inst_kv_rings: + smem_v1 = SmemKvTileResource( + pipeline_config=smem_v1_cfg, + cfg=cfg, + tma_desc_k=tma_desc_k, + tma_desc_v=tma_desc_v, + tma_desc_k_sf=tma_desc_k_sf, + tma_desc_v_sf=tma_desc_v_sf, + tma_desc_k_atom=tma_desc_k_atom, + tma_desc_v_atom=tma_desc_v_atom, + tma_desc_k_summary=tma_desc_k_summary, + tma_desc_v_summary=tma_desc_v_summary, + tma_desc_k_summary_atom=tma_desc_k_summary_atom, + tma_desc_v_summary_atom=tma_desc_v_summary_atom, + sparse_kv_metadata=sparse_kv_metadata1, + page_offsets_kv=smem_page_offsets_v or smem_page_offsets, + seqlens_kv=kv_seqlens, + max_seq_len_kv=max_seq_len_kv, + h_k_idx=h_k_idx, + b_idx=b_idx, + q_group_idx=q_group_idx, + seq_len_q=seq_len_q, + inst_id=1, + kv_kind=KV_KIND_V, + name="smemV1", + ) + else: + smem_kv = SmemKvResource( + pipeline_config=smem_kv_cfg, cfg=cfg, tma_desc_k=tma_desc_k, tma_desc_v=tma_desc_v, + tma_desc_k_sf=tma_desc_k_sf, + tma_desc_v_sf=tma_desc_v_sf, tma_desc_k_atom=tma_desc_k_atom, tma_desc_v_atom=tma_desc_v_atom, tma_desc_k_summary=tma_desc_k_summary, tma_desc_v_summary=tma_desc_v_summary, tma_desc_k_summary_atom=tma_desc_k_summary_atom, tma_desc_v_summary_atom=tma_desc_v_summary_atom, - sparse_kv_metadata=sparse_kv_metadata1, - page_offsets_kv=smem_page_offsets_v or smem_page_offsets, + sparse_kv_metadata0=sparse_kv_metadata0, + sparse_kv_metadata1=sparse_kv_metadata1, + page_offsets_kv=smem_page_offsets, seqlens_kv=kv_seqlens, max_seq_len_kv=max_seq_len_kv, h_k_idx=h_k_idx, b_idx=b_idx, q_group_idx=q_group_idx, seq_len_q=seq_len_q, - inst_id=1, - kv_kind=KV_KIND_V, - name="smemV1", + name="smemKv", ) - else: - smem_kv = SmemKvResource( - pipeline_config=smem_kv_cfg, + transformed_kv = None + mma_smem_kv = smem_kv + if cfg.use_transform_kv: + if cfg.store_transformed_kv_in_tmem: + assert smem_kv is not None + transformed_kv = TmemTransformedKvResource( + pipeline_config=smem_transformed_kv_cfg, + cfg=cfg, + src_smem_kv=smem_kv, + name="tmemTransformedKv", + ) + else: + transformed_kv = SmemTransformedKvResource( + pipeline_config=smem_transformed_kv_cfg, + cfg=cfg, + src_smem_kv=smem_kv, + src_smem_k0=smem_k0, + src_smem_k1=smem_k1, + src_smem_v0=smem_v0, + src_smem_v1=smem_v1, + page_idx_kv=page_idx_kv, + num_heads_kv=num_heads_kv, + name="smemTransformedKv", + ) + mma_smem_kv = transformed_kv + + # The correction task stages and reads the work tile's V scales itself. + sage_v_scales = None + if cfg.use_sage_attention: + sage_v_scales = SageVScalesResource( + pipeline_config=PipelineConfig( + num_stages=1, + num_bytes=0, + producer_group=correction_grp, + consumer_group=correction_grp, + pipeline_type=PipelineType.AsyncAsync, + cta_layout_vmnk=cta_layout, + advance_on_wait=True, + ), + cfg=cfg, + scale_tensors=sage_scale_tensors, + h_k_idx=h_k_idx, + b_idx=b_idx, + name="sageVScales", + ) + # Each softmax instance fills and reads its own ``sfK`` words; only the + # SMEM form carries a pipeline. A mixed-geometry plan adds a second + # resource for the proxy summaries. + sage_k_scales0 = None + sage_k_scales1 = None + sage_summary_k_scales0 = None + sage_summary_k_scales1 = None + if cfg.use_sage_attention: + + def _sage_k_scales_cfg(groups: int, softmax_grp) -> PipelineConfig | None: + if not cfg.sage_k_scales_in_smem_for(groups): + return None + return PipelineConfig( + num_stages=SAGE_K_SCALES_RING_STAGES, + num_bytes=0, + producer_group=softmax_grp, + consumer_group=softmax_grp, + pipeline_type=PipelineType.AsyncAsync, + cta_layout_vmnk=cta_layout, + advance_on_wait=True, + ) + + token_groups = cfg.sage_k_groups_per_fragment + sage_k_scales0 = SageKScalesResource( + pipeline_config=_sage_k_scales_cfg(token_groups, softmax0_grp), + inst_id=0, + route_metadata=sparse_softmax_metadata0, + name="sageKScales0", cfg=cfg, - tma_desc_k=tma_desc_k, - tma_desc_v=tma_desc_v, - tma_desc_k_atom=tma_desc_k_atom, - tma_desc_v_atom=tma_desc_v_atom, - tma_desc_k_summary=tma_desc_k_summary, - tma_desc_v_summary=tma_desc_v_summary, - tma_desc_k_summary_atom=tma_desc_k_summary_atom, - tma_desc_v_summary_atom=tma_desc_v_summary_atom, - sparse_kv_metadata0=sparse_kv_metadata0, - sparse_kv_metadata1=sparse_kv_metadata1, - page_offsets_kv=smem_page_offsets, seqlens_kv=kv_seqlens, max_seq_len_kv=max_seq_len_kv, + q_group_idx=q_group_idx, + seq_len_q=seq_len_q, h_k_idx=h_k_idx, b_idx=b_idx, + scale_tensors=sage_scale_tensors, + ) + sage_k_scales1 = SageKScalesResource( + pipeline_config=_sage_k_scales_cfg(token_groups, softmax1_grp), + inst_id=1, + route_metadata=sparse_softmax_metadata1, + name="sageKScales1", + cfg=cfg, + seqlens_kv=kv_seqlens, + max_seq_len_kv=max_seq_len_kv, q_group_idx=q_group_idx, seq_len_q=seq_len_q, - name="smemKv", + h_k_idx=h_k_idx, + b_idx=b_idx, + scale_tensors=sage_scale_tensors, ) + if cfg.sage_mixed_k_geometry: + summary_groups = cfg.sage_summary_k_groups_per_fragment + sage_summary_k_scales0 = SageKScalesResource( + pipeline_config=_sage_k_scales_cfg(summary_groups, softmax0_grp), + inst_id=0, + summary=True, + route_metadata=sparse_softmax_metadata0, + name="sageSummaryKScales0", + cfg=cfg, + seqlens_kv=kv_seqlens, + max_seq_len_kv=max_seq_len_kv, + q_group_idx=q_group_idx, + seq_len_q=seq_len_q, + h_k_idx=h_k_idx, + b_idx=b_idx, + scale_tensors=sage_scale_tensors, + ) + sage_summary_k_scales1 = SageKScalesResource( + pipeline_config=_sage_k_scales_cfg(summary_groups, softmax1_grp), + inst_id=1, + summary=True, + route_metadata=sparse_softmax_metadata1, + name="sageSummaryKScales1", + cfg=cfg, + seqlens_kv=kv_seqlens, + max_seq_len_kv=max_seq_len_kv, + q_group_idx=q_group_idx, + seq_len_q=seq_len_q, + h_k_idx=h_k_idx, + b_idx=b_idx, + scale_tensors=sage_scale_tensors, + ) tmem_s0 = TmemSResource( inst_id=0, @@ -1051,28 +1337,47 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: h_r=h_r, q_group_idx=q_group_idx, seq_len_q=seq_len_q, + h_k_idx=h_k_idx, + b_idx=b_idx, sync_barrier_id=0, + sage_k_scales=sage_k_scales0, + sage_summary_k_scales=sage_summary_k_scales0 or sage_k_scales0, name="tmemS0", + scale_tensors=sage_scale_tensors, ) - tmem_s1 = TmemSResource( - inst_id=1, - pipeline_config=tmem_s1_cfg, - cfg=cfg, - scale_softmax_log2=scale_softmax_log2, - seqlens_kv=kv_seqlens, - max_seq_len_kv=max_seq_len_kv, - h_r=h_r, - q_group_idx=q_group_idx, - seq_len_q=seq_len_q, - sync_barrier_id=1, - name="tmemS1", - ) + tmem_s1 = None + if not use_one_inst_kv: + tmem_s1 = TmemSResource( + inst_id=1, + pipeline_config=tmem_s1_cfg, + cfg=cfg, + scale_softmax_log2=scale_softmax_log2, + seqlens_kv=kv_seqlens, + max_seq_len_kv=max_seq_len_kv, + h_r=h_r, + q_group_idx=q_group_idx, + seq_len_q=seq_len_q, + h_k_idx=h_k_idx, + b_idx=b_idx, + sync_barrier_id=1, + score_seed_owner=tmem_s0, + sage_k_scales=sage_k_scales1, + sage_summary_k_scales=sage_summary_k_scales1 or sage_k_scales1, + name="tmemS1", + scale_tensors=sage_scale_tensors, + ) # Packed persistent QK derives the descriptor from Q's just-waited # consumer stage, avoiding a routed HEAD-to-LOOP descriptor value across # the guarded work-tile region. Fixed/static schedules keep their existing # explicit descriptor route. tmem_s0.q_ref = smem_q - tmem_s1.q_ref = smem_q + if tmem_s1 is not None: + tmem_s1.q_ref = smem_q + if cfg.uses_q_token_kv_block_sparse_page_membership: + assert smem_page_offsets is not None + tmem_s0.page_offsets_ref = smem_page_offsets + if tmem_s1 is not None: + tmem_s1.page_offsets_ref = smem_page_offsets smem_p0 = SmemPResource( inst_id=0, @@ -1080,16 +1385,22 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: cfg=cfg, scale_softmax_log2=scale_softmax_log2, use_variable_seqlens_kv=use_runtime_seqlens_kv, + sage_k_scales=sage_k_scales0, + sage_summary_k_scales=sage_summary_k_scales0 or sage_k_scales0, name="smemP0", ) - smem_p1 = SmemPResource( - inst_id=1, - pipeline_config=smem_p1_cfg, - cfg=cfg, - scale_softmax_log2=scale_softmax_log2, - use_variable_seqlens_kv=use_runtime_seqlens_kv, - name="smemP1", - ) + smem_p1 = None + if not use_one_inst_kv: + smem_p1 = SmemPResource( + inst_id=1, + pipeline_config=smem_p1_cfg, + cfg=cfg, + scale_softmax_log2=scale_softmax_log2, + use_variable_seqlens_kv=use_runtime_seqlens_kv, + sage_k_scales=sage_k_scales1, + sage_summary_k_scales=sage_summary_k_scales1 or sage_k_scales1, + name="smemP1", + ) tmem_o = TmemOResource( pipeline_config=tmem_o_cfg, @@ -1104,12 +1415,14 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: cfg=cfg, name="tmemSoftmaxLocal0", ) - tmem_softmax_local1 = TmemSoftmaxLocalResource( - inst_id=1, - pipeline_config=softmax_local1_cfg, - cfg=cfg, - name="tmemSoftmaxLocal1", - ) + tmem_softmax_local1 = None + if not use_one_inst_kv: + tmem_softmax_local1 = TmemSoftmaxLocalResource( + inst_id=1, + pipeline_config=softmax_local1_cfg, + cfg=cfg, + name="tmemSoftmaxLocal1", + ) tmem_stats_done0 = ( TmemStatsDoneResource( pipeline_config=stats_done0_cfg, @@ -1133,26 +1446,29 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: sum_barrier_id=2, name="tmemSoftmaxGlobal0", ) - tmem_softmax_global1 = TmemSoftmaxGlobalResource( - inst_id=1, - cfg=cfg, - scale_softmax_log2=scale_softmax_log2, - sum_barrier_id=3, - name="tmemSoftmaxGlobal1", - ) + tmem_softmax_global1 = None + if not use_one_inst_kv: + tmem_softmax_global1 = TmemSoftmaxGlobalResource( + inst_id=1, + cfg=cfg, + scale_softmax_log2=scale_softmax_log2, + sum_barrier_id=3, + name="tmemSoftmaxGlobal1", + ) tmem_softmax_order = ( TmemSoftmaxOrderResource(cfg=cfg, name="tmemSoftmaxOrder") if use_ordered_softmax_barrier else None ) smem_p0.tmem_s_ref = tmem_s0 - smem_p1.tmem_s_ref = tmem_s1 smem_p0.tmem_o_ref = tmem_o - smem_p1.tmem_o_ref = tmem_o tmem_softmax_global0.p_ref = smem_p0 - tmem_softmax_global1.p_ref = smem_p1 tmem_softmax_global0.tmem_s_ref = tmem_s0 - tmem_softmax_global1.tmem_s_ref = tmem_s1 + if not use_one_inst_kv: + smem_p1.tmem_s_ref = tmem_s1 + smem_p1.tmem_o_ref = tmem_o + tmem_softmax_global1.p_ref = smem_p1 + tmem_softmax_global1.tmem_s_ref = tmem_s1 tmem_corr0 = TmemCorrResource( inst_id=0, @@ -1176,43 +1492,168 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: active_splits_kv=active_splits_kv, static_full_split_prefix=static_full_split_prefix, name="tmemCorr0", + sage_v_scales=sage_v_scales, ) - tmem_corr1 = TmemCorrResource( - inst_id=1, - cfg=cfg, - scale_softmax_log2=scale_softmax_log2, - output_scale=output_scale, - o_ptr=o_ptr, - partial_o_ptr=partial_o_ptr, - partial_stats_ptr=partial_stats_ptr, - split_kv_counter_ptr=split_kv_counter_ptr, - attention_sinks_ptr=attention_sinks_ptr, - seqlens_kv=kv_seqlens, - max_seq_len_kv=corr_max_seq_len_kv, - num_heads_kv=num_heads_kv, - h_r=h_r, - h_k_idx=h_k_idx, - b_idx=b_idx, - q_group_idx=q_group_idx, - q_token_offset=q_token_offset, - seq_len_q=seq_len_q, - active_splits_kv=active_splits_kv, - static_full_split_prefix=static_full_split_prefix, - name="tmemCorr1", - ) - tmem_corr1.smem_p0_ref = smem_p0 - tmem_corr1.smem_p1_ref = smem_p1 tmem_corr0.tmem_o_ref = tmem_o - tmem_corr1.tmem_o_ref = tmem_o tmem_corr0.softmax_local0_ref = tmem_softmax_local0 - tmem_corr1.softmax_local0_ref = tmem_softmax_local0 - if use_one_inst_qkv: - tmem_corr0.softmax_local1_ref = None - tmem_corr1.softmax_local1_ref = None - else: + tmem_corr0.softmax_local1_ref = None if use_one_inst_kv else tmem_softmax_local1 + tmem_corr1 = None + if not use_one_inst_kv: + tmem_corr1 = TmemCorrResource( + inst_id=1, + cfg=cfg, + scale_softmax_log2=scale_softmax_log2, + output_scale=output_scale, + o_ptr=o_ptr, + partial_o_ptr=partial_o_ptr, + partial_stats_ptr=partial_stats_ptr, + split_kv_counter_ptr=split_kv_counter_ptr, + attention_sinks_ptr=attention_sinks_ptr, + seqlens_kv=kv_seqlens, + max_seq_len_kv=corr_max_seq_len_kv, + num_heads_kv=num_heads_kv, + h_r=h_r, + h_k_idx=h_k_idx, + b_idx=b_idx, + q_group_idx=q_group_idx, + q_token_offset=q_token_offset, + seq_len_q=seq_len_q, + active_splits_kv=active_splits_kv, + static_full_split_prefix=static_full_split_prefix, + name="tmemCorr1", + sage_v_scales=sage_v_scales, + ) + tmem_corr1.smem_p0_ref = smem_p0 + tmem_corr1.smem_p1_ref = smem_p1 + tmem_corr1.tmem_o_ref = tmem_o + tmem_corr1.softmax_local0_ref = tmem_softmax_local0 tmem_corr0.softmax_local1_ref = tmem_softmax_local1 tmem_corr1.softmax_local1_ref = tmem_softmax_local1 + # ------------------------------------------------------------------ + # Finalize SMEM and membership mode before tracing tasks. + # ------------------------------------------------------------------ + smem_resources = [] + if work_queue is not None: + smem_resources.append(work_queue) + if schedule_token_throttle is not None: + smem_resources.append(schedule_token_throttle) + if membership_lifetime is not None: + smem_resources.append(membership_lifetime) + smem_resources.append(smem_q) + if smem_page_offsets is not None: + smem_resources.append(smem_page_offsets) + if smem_page_offsets_v is not None: + smem_resources.append(smem_page_offsets_v) + if sparse_kv_metadata0 is not None: + smem_resources.append(sparse_kv_metadata0) + smem_resources.append(sparse_kv_metadata1) + if sparse_softmax_metadata0 is not None: + smem_resources.append(sparse_softmax_metadata0) + smem_resources.append(sparse_softmax_metadata1) + if use_one_inst_qkv: + smem_resources.append(smem_k0) + smem_resources.append(smem_v0) + elif use_per_inst_kv_resources: + for resource in (smem_k0, smem_k1, smem_v0, smem_v1): + if resource is not None: + smem_resources.append(resource) + else: + smem_resources.append(smem_kv) + if transformed_kv is not None: + smem_resources.append(transformed_kv) + smem_resources.append(smem_p0) + if not use_one_inst_kv: + smem_resources.append(smem_p1) + # Each instance's scale rings follow its S resource in the layout. + smem_resources.append(tmem_s0) + smem_resources.extend( + r for r in (sage_k_scales0, sage_summary_k_scales0) if r is not None + ) + if not use_one_inst_kv: + smem_resources.append(tmem_s1) + smem_resources.extend( + r for r in (sage_k_scales1, sage_summary_k_scales1) if r is not None + ) + smem_resources.append(tmem_o) + smem_resources.append(tmem_softmax_local0) + if not use_one_inst_kv: + smem_resources.append(tmem_softmax_local1) + smem_resources.append(tmem_softmax_global0) + if not use_one_inst_kv: + smem_resources.append(tmem_softmax_global1) + smem_resources.append(tmem_corr0) + if not use_one_inst_kv: + smem_resources.append(tmem_corr1) + if sage_v_scales is not None: + smem_resources.append(sage_v_scales) + + def allocate_smem() -> tuple[SmemAllocator, SmemAllocation | None]: + allocator = SmemAllocator() + clc_response_alloc = None + for resource in smem_resources: + allocator.add_resource(resource) + if resource is work_queue and tile_sched_params is not None: + # A compile-time offset from the unified base spares every + # role a separate pointer that ptxas spills to local memory. + clc_response_alloc = allocator.add( + SmemAllocation( + "clc_response", + dtype=cutlass.Int128, + count=_WORK_QUEUE_STAGES, + alignment=16, + ) + ) + allocator.add_tmem_ptr( + SmemAllocation("fmha_tmem_ptr_i32", dtype=cutlass.Int32, alignment=4) + ) + allocator.compute_layout() + return allocator, clc_response_alloc + + def smem_bytes(allocator: SmemAllocator) -> int: + # Include barriers and conservatively round data to tensor alignment. + return ( + allocator.total_smem_bytes + cfg.stensor_align - 1 + ) // cfg.stensor_align * cfg.stensor_align + allocator.barrier_smem_bytes + + # The static profile checks bound only the Q and K/V pipelines; check the + # complete layout against the SM100-family SMEM capacity. + smem_capacity_bytes = SMEM_CAPACITY_KIB * BYTES_PER_KIB + smem_allocator, clc_response_alloc = allocate_smem() + if ( + cfg.uses_q_token_kv_block_sparse_page_membership + and smem_bytes(smem_allocator) > smem_capacity_bytes + ): + # A full union can be much larger than a KV tile. Keep small routes + # cached, but let Softmax read immutable packed GMEM words on demand + # when the complete resource layout cannot hold the membership row. + # This also accounts for the fixed K/V rings of one-inst D256. + assert smem_page_offsets is not None + smem_page_offsets.cache_memberships_in_smem = False + smem_page_offsets._init_placeholder_state() + smem_allocator, clc_response_alloc = allocate_smem() + launch_smem_bytes = smem_bytes(smem_allocator) + if launch_smem_bytes > smem_capacity_bytes: + raise ValueError( + "decode resources exceed the SM shared-memory capacity: " + f"q_stages={cfg.q_stages}, kv_stages={cfg.kv_stages} need " + f"{launch_smem_bytes} bytes, capacity is {smem_capacity_bytes} bytes" + ) + if clc_response_alloc is not None: + # Bind the scheduler config before the task manager creates the queue. + smem_allocator.allocate() + work_queue.tile_scheduler_config = ( + TileSchedulerConfig.create_clc_dynamic_persistent_tile_scheduler_params( + tile_scheduler_params=tile_sched_params, + response_ptr=cute.make_ptr( + cutlass.Int128, + smem_allocator.get(clc_response_alloc).data_ptr(), + mem_space=cutlass.AddressSpace.smem, + assumed_align=16, + ), + ) + ) + # ------------------------------------------------------------------ # Domain computation # ------------------------------------------------------------------ @@ -1334,9 +1775,10 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: domain_bias=0, warp_idx=page_offsets_warp_idx, num_warps=cfg.page_offsets_num_warps, - block_table_capacity=( - block_table_capacity if use_native_paged_kv else None - ), + membership_lifetime=membership_lifetime, + block_table_capacity=page_table_capacity + if use_native_paged_kv + else None, **task_runtime_kwargs, ) else: @@ -1353,9 +1795,10 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: domain_bias=0, warp_idx=page_offsets_warp_idx, num_warps=cfg.page_offsets_num_warps, - block_table_capacity=( - block_table_capacity if use_native_paged_kv else None - ), + membership_lifetime=membership_lifetime, + block_table_capacity=page_table_capacity + if use_native_paged_kv + else None, **task_runtime_kwargs, ) if use_one_inst_qkv: @@ -1373,7 +1816,7 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: domain_bias=0, **task_runtime_kwargs, ) - elif use_per_inst_kv_resources: + elif use_per_inst_kv_resources and not cfg.use_transform_kv: mma_task = create_mma_task_split_kv( smem_q, smem_k0, @@ -1396,7 +1839,7 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: else: mma_task = create_mma_task( smem_q, - smem_kv, + mma_smem_kv, tmem_s0, tmem_s1, smem_p0, @@ -1408,6 +1851,21 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: domain_bias=0, **task_runtime_kwargs, ) + transform_kv_task = None + if cfg.use_transform_kv: + transform_kv_task = create_transform_kv_task( + smem_kv, + smem_k0, + smem_k1, + smem_v0, + smem_v1, + transformed_kv, + work_queue, + cfg, + domain=load_domain, + domain_bias=0, + **task_runtime_kwargs, + ) softmax0_task = create_softmax0_task( tmem_s0, tmem_softmax_local0, @@ -1419,10 +1877,13 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: cfg, domain=softmax_domain, domain_bias=1, + membership_lifetime=membership_lifetime, + sage_k_scales=sage_k_scales0, + sage_summary_k_scales=sage_summary_k_scales0, **task_runtime_kwargs, ) softmax1_task = None - if not use_one_inst_qkv: + if not use_one_inst_kv: softmax1_task = create_softmax1_task( tmem_s1, tmem_softmax_local1, @@ -1434,9 +1895,12 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: cfg, domain=softmax_domain, domain_bias=1, + membership_lifetime=membership_lifetime, + sage_k_scales=sage_k_scales1, + sage_summary_k_scales=sage_summary_k_scales1, **task_runtime_kwargs, ) - if use_one_inst_qkv: + if use_one_inst_kv: correction_task = create_correction_task_one_inst_qkv( tmem_softmax_local0, tmem_o, @@ -1444,6 +1908,7 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: work_queue, cfg, tmem_stats_done=tmem_stats_done0, + sage_v_scales=sage_v_scales, domain=corr_domain, domain_bias=0, **task_runtime_kwargs, @@ -1460,6 +1925,7 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: domain=corr_domain, tmem_stats_done0=tmem_stats_done0, tmem_stats_done1=tmem_stats_done1, + sage_v_scales=sage_v_scales, domain_bias=0, **task_runtime_kwargs, ) @@ -1475,6 +1941,8 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: (cfg.wg1_padding_warp_idx, cfg.wg1_padding_num_warps), (cfg.wg2_padding_warp_idx, cfg.wg2_padding_num_warps), (cfg.wg3_padding_warp_idx, cfg.wg3_padding_num_warps), + (cfg.wg4_padding_warp_idx, cfg.wg4_padding_num_warps), + (cfg.wg5_padding_warp_idx, cfg.wg5_padding_num_warps), ) padding_tasks = [ create_padding_task( @@ -1494,6 +1962,8 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: if page_offsets_task is not None: task_list.append(page_offsets_task) task_list.extend(load_tasks) + if transform_kv_task is not None: + task_list.append(transform_kv_task) if use_one_inst_qkv and not use_clc_dynamic: task_list.extend([correction_task, mma_task]) task_list.append(softmax0_task) @@ -1527,6 +1997,47 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: tmem_o: [smem_p0, smem_v0], tmem_corr0: [tmem_softmax_local0, tmem_o], } + elif cfg.use_transform_kv and use_per_inst_kv_resources: + resource_dependency_graph = { + smem_q: [], + smem_k0: smem_k_deps, + smem_k1: smem_k_deps, + smem_v0: smem_v_deps, + smem_v1: smem_v_deps, + transformed_kv: [smem_k0, smem_k1, smem_v0, smem_v1], + tmem_s0: [transformed_kv, smem_q], + smem_p0: [tmem_s0], + tmem_softmax_local0: [tmem_s0], + tmem_softmax_global0: [tmem_s0], + tmem_o: [smem_p0, transformed_kv], + tmem_corr0: [tmem_softmax_local0, tmem_o], + } + if not use_one_inst_kv: + resource_dependency_graph.update( + { + tmem_s1: [transformed_kv, smem_q], + smem_p1: [tmem_s1], + tmem_softmax_local1: [tmem_s1], + tmem_softmax_global1: [tmem_s1], + tmem_o: [smem_p0, smem_p1, transformed_kv], + tmem_corr1: [ + tmem_softmax_local0, + tmem_softmax_local1, + tmem_o, + ], + } + ) + elif use_one_inst_kv: + resource_dependency_graph = { + smem_q: [], + smem_kv: smem_kv_deps, + tmem_s0: [mma_smem_kv, smem_q], + smem_p0: [tmem_s0], + tmem_softmax_local0: [tmem_s0], + tmem_softmax_global0: [tmem_s0], + tmem_o: [smem_p0, mma_smem_kv], + tmem_corr0: [tmem_softmax_local0, tmem_o], + } elif use_per_inst_kv_resources: resource_dependency_graph = { **( @@ -1548,21 +2059,33 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: smem_q: [], smem_k0: smem_k_deps + ([sparse_kv_metadata0] if sparse_kv_metadata0 is not None else []), - smem_k1: smem_k_deps - + ([sparse_kv_metadata1] if sparse_kv_metadata1 is not None else []), smem_v0: smem_v_deps + ([sparse_kv_metadata0] if sparse_kv_metadata0 is not None else []), - smem_v1: smem_v_deps - + ([sparse_kv_metadata1] if sparse_kv_metadata1 is not None else []), + **( + { + smem_k1: smem_k_deps + + ( + [sparse_kv_metadata1] if sparse_kv_metadata1 is not None else [] + ), + smem_v1: smem_v_deps + + ( + [sparse_kv_metadata1] if sparse_kv_metadata1 is not None else [] + ), + } + if smem_k1 is not None + else {} + ), tmem_s0: [smem_k0, smem_q], - tmem_s1: [smem_k1, smem_q], + tmem_s1: [smem_k0 if smem_k1 is None else smem_k1, smem_q], smem_p0: [tmem_s0], smem_p1: [tmem_s1], tmem_softmax_local0: [tmem_s0], tmem_softmax_local1: [tmem_s1], tmem_softmax_global0: [tmem_s0], tmem_softmax_global1: [tmem_s1], - tmem_o: [smem_p0, smem_p1, smem_v0, smem_v1], + tmem_o: [smem_p0, smem_p1, smem_v0] + if smem_v1 is None + else [smem_p0, smem_p1, smem_v0, smem_v1], tmem_corr0: [tmem_softmax_local0, tmem_o], tmem_corr1: [tmem_softmax_local0, tmem_softmax_local1, tmem_o], } @@ -1585,18 +2108,20 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: if sparse_kv_metadata0 is not None else [] ), - tmem_s0: [smem_kv, smem_q], - tmem_s1: [smem_kv, smem_q], + tmem_s0: [mma_smem_kv, smem_q], + tmem_s1: [mma_smem_kv, smem_q], smem_p0: [tmem_s0], smem_p1: [tmem_s1], tmem_softmax_local0: [tmem_s0], tmem_softmax_local1: [tmem_s1], tmem_softmax_global0: [tmem_s0], tmem_softmax_global1: [tmem_s1], - tmem_o: [smem_p0, smem_p1, smem_kv], + tmem_o: [smem_p0, smem_p1, mma_smem_kv], tmem_corr0: [tmem_softmax_local0, tmem_o], tmem_corr1: [tmem_softmax_local0, tmem_softmax_local1, tmem_o], } + if transformed_kv is not None and not use_per_inst_kv_resources: + resource_dependency_graph[transformed_kv] = [smem_kv] if sparse_softmax_metadata0 is not None: resource_dependency_graph[smem_p0].append(sparse_softmax_metadata0) resource_dependency_graph[tmem_softmax_local0].append(sparse_softmax_metadata0) @@ -1605,12 +2130,51 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: resource_dependency_graph[smem_p1].append(sparse_softmax_metadata1) resource_dependency_graph[tmem_softmax_local1].append(sparse_softmax_metadata1) resource_dependency_graph[tmem_softmax_global1].append(sparse_softmax_metadata1) + if membership_lifetime is not None: + resource_dependency_graph[membership_lifetime] = [] + for resource in (smem_p0, tmem_softmax_local0, tmem_softmax_global0): + resource_dependency_graph[resource].append(membership_lifetime) + if not use_one_inst_kv: + for resource in (smem_p1, tmem_softmax_local1, tmem_softmax_global1): + resource_dependency_graph[resource].append(membership_lifetime) if tmem_stats_done0 is not None: resource_dependency_graph[tmem_s0].append(tmem_stats_done0) resource_dependency_graph[tmem_stats_done0] = [tmem_softmax_local0] if tmem_stats_done1 is not None: + assert tmem_s1 is not None + assert tmem_softmax_local1 is not None resource_dependency_graph[tmem_s1].append(tmem_stats_done1) resource_dependency_graph[tmem_stats_done1] = [tmem_softmax_local1] + # A softmax instance's outputs read its K scale words; an SMEM-form + # resource fills them from the route's staged metadata, while a + # register-form one hands them straight to the outputs. The correction + # epilogue reads the V channel scales. + for scales_pair, route_metadata, outputs in ( + ( + (sage_k_scales0, sage_summary_k_scales0), + sparse_softmax_metadata0, + (smem_p0, tmem_softmax_local0, tmem_softmax_global0), + ), + ( + (sage_k_scales1, sage_summary_k_scales1), + sparse_softmax_metadata1, + (smem_p1, tmem_softmax_local1, tmem_softmax_global1), + ), + ): + for scales in scales_pair: + if scales is None: + continue + resource_dependency_graph[scales] = ( + [route_metadata] + if route_metadata is not None and scales.in_smem + else [] + ) + for resource in outputs: + resource_dependency_graph[resource].append(scales) + if sage_v_scales is not None: + resource_dependency_graph[sage_v_scales] = [] + for resource in (tmem_corr0, tmem_corr1): + resource_dependency_graph[resource].append(sage_v_scales) if cutlass.const_expr(use_ordered_softmax_barrier): resource_dependency_graph[tmem_softmax_order] = [tmem_s0] resource_dependency_graph[smem_p1] = [ @@ -1634,16 +2198,59 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: dma_consumer_release_labels: dict[ tuple[MemoryResource, MemoryResource], set[str] ] = {} + if transformed_kv is not None: + dma_consumer_release_labels[(transformed_kv, tmem_s0)] = {"k_desc_0"} + dma_consumer_release_labels[(transformed_kv, tmem_o)] = {"v_desc_0"} + if not use_one_inst_kv: + dma_consumer_release_labels[(transformed_kv, tmem_s1)] = {"k_desc_1"} + dma_consumer_release_labels[(transformed_kv, tmem_o)].add("v_desc_1") if smem_page_offsets is not None: if use_one_inst_qkv: - dma_consumer_release_labels.update( - { - (smem_page_offsets, smem_k0): {"read_offsets_k0"}, - (smem_page_offsets, smem_v0): {"read_offsets_v0"}, - } - ) + if _can_hold_native_page_window(cfg, smem_page_offsets): + # One consumer stage remains live across the complete K/V + # cadence, so both DMA edges share its final release label. + dma_consumer_release_labels.update( + { + (smem_page_offsets, smem_k0): {"read_offsets"}, + (smem_page_offsets, smem_v0): {"read_offsets"}, + } + ) + else: + dma_consumer_release_labels.update( + { + (smem_page_offsets, smem_k0): {"read_offsets_k0"}, + (smem_page_offsets, smem_v0): {"read_offsets_v0"}, + } + ) elif use_per_inst_kv_resources: - if smem_page_offsets_v is not None: + if ( + smem_page_offsets_v is None + and smem_page_offsets.holds_encoded_locator_window + ): + # Independent K/V data FIFOs can still share a single + # read-only locator window until the last V tile is issued. + dma_consumer_release_labels.update( + { + (smem_page_offsets, resource): {"read_offsets"} + for resource in (smem_k0, smem_k1, smem_v0, smem_v1) + } + ) + elif smem_k1 is None: + # One ring per operand serves both instances, so the single key + # carries the release labels of both. + dma_consumer_release_labels.update( + { + (smem_page_offsets, smem_k0): { + "read_offsets_k0", + "read_offsets_k1", + }, + (smem_page_offsets_v, smem_v0): { + "read_offsets_v0", + "read_offsets_v1", + }, + } + ) + elif smem_page_offsets_v is not None: dma_consumer_release_labels.update( { (smem_page_offsets, smem_k0): {"read_offsets_k0"}, @@ -1662,67 +2269,33 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: } ) else: - dma_consumer_release_labels[(smem_page_offsets, smem_kv)] = { - "cache_page_ids" if cfg.num_head_dim_stages_kv > 1 else "read_offsets" - } - if smem_kv is not None: - dma_consumer_release_labels.update( - { - (smem_kv, tmem_s0): {"k_desc_0"}, - (smem_kv, tmem_s1): {"k_desc_1"}, - (smem_kv, tmem_o): {"v_desc_0", "v_desc_1"}, - } - ) + hold_page_window = _can_hold_native_page_window(cfg, smem_page_offsets) + if cfg.num_head_dim_stages_kv > 1 and not cfg.uses_scattered_page_route: + page_offset_labels = {"cache_page_ids"} + elif hold_page_window: + page_offset_labels = {"read_offsets"} + elif use_one_inst_kv: + page_offset_labels = {"read_offsets_k0", "read_offsets_v0"} + else: + page_offset_labels = { + "read_offsets_k0", + "read_offsets_k1", + "read_offsets_v0", + "read_offsets_v1", + } + dma_consumer_release_labels[(smem_page_offsets, smem_kv)] = ( + page_offset_labels + ) + if smem_kv is not None and transformed_kv is None: + dma_consumer_release_labels[(smem_kv, tmem_s0)] = {"k_desc_0"} + dma_consumer_release_labels[(smem_kv, tmem_o)] = {"v_desc_0"} + if not use_one_inst_kv: + dma_consumer_release_labels[(smem_kv, tmem_s1)] = {"k_desc_1"} + dma_consumer_release_labels[(smem_kv, tmem_o)].add("v_desc_1") # ------------------------------------------------------------------ - # SMEM / TMEM allocators + # TMEM allocator # ------------------------------------------------------------------ - smem_allocator = SmemAllocator() - if work_queue is not None: - smem_allocator.add_resource(work_queue) - if schedule_token_throttle is not None: - smem_allocator.add_resource(schedule_token_throttle) - smem_allocator.add_resource(smem_q) - if smem_page_offsets is not None: - smem_allocator.add_resource(smem_page_offsets) - if smem_page_offsets_v is not None: - smem_allocator.add_resource(smem_page_offsets_v) - if sparse_kv_metadata0 is not None: - smem_allocator.add_resource(sparse_kv_metadata0) - smem_allocator.add_resource(sparse_kv_metadata1) - if sparse_softmax_metadata0 is not None: - smem_allocator.add_resource(sparse_softmax_metadata0) - smem_allocator.add_resource(sparse_softmax_metadata1) - if use_one_inst_qkv: - smem_allocator.add_resource(smem_k0) - smem_allocator.add_resource(smem_v0) - elif use_per_inst_kv_resources: - smem_allocator.add_resource(smem_k0) - smem_allocator.add_resource(smem_k1) - smem_allocator.add_resource(smem_v0) - smem_allocator.add_resource(smem_v1) - else: - smem_allocator.add_resource(smem_kv) - smem_allocator.add_resource(smem_p0) - if not use_one_inst_qkv: - smem_allocator.add_resource(smem_p1) - smem_allocator.add_resource(tmem_s0) - if not use_one_inst_qkv: - smem_allocator.add_resource(tmem_s1) - smem_allocator.add_resource(tmem_o) - smem_allocator.add_resource(tmem_softmax_local0) - if not use_one_inst_qkv: - smem_allocator.add_resource(tmem_softmax_local1) - smem_allocator.add_resource(tmem_softmax_global0) - if not use_one_inst_qkv: - smem_allocator.add_resource(tmem_softmax_global1) - smem_allocator.add_resource(tmem_corr0) - if not use_one_inst_qkv: - smem_allocator.add_resource(tmem_corr1) - smem_allocator.add_tmem_ptr( - SmemAllocation("fmha_tmem_ptr_i32", dtype=cutlass.Int32, alignment=4) - ) - smem_allocator.compute_layout() tmem_allocator = TmemAllocator() if cfg.use_keeps_mma_ab: if use_one_inst_qkv: @@ -1749,10 +2322,14 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: tmem_allocator.add_resource(tmem_softmax_local1) else: tmem_allocator.add_resource(tmem_s0) - tmem_allocator.add_resource(tmem_s1) tmem_allocator.add_resource(tmem_softmax_local0) - tmem_allocator.add_resource(tmem_softmax_local1) + if not use_one_inst_kv: + tmem_allocator.add_resource(tmem_s1) + tmem_allocator.add_resource(tmem_softmax_local1) tmem_allocator.add_resource(tmem_o) + if cfg.store_transformed_kv_in_tmem: + assert transformed_kv is not None + 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 @@ -1841,13 +2418,14 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: tmem_s_alloc.offset + cfg.tmem_stats_cols ) - eager_init_resources = ( - [tmem_corr0] if use_one_inst_qkv else [tmem_corr0, tmem_corr1] - ) + eager_init_resources = [tmem_corr0] if use_one_inst_kv else [tmem_corr0, tmem_corr1] if cfg.streams_tmem_p_fragments: # Streamed TMEM P operands use one-way per-fragment ready barriers. # Initialize them beside correction's manually managed SMEM state. eager_init_resources.extend([smem_p0, smem_p1]) + if cfg.uses_int32_scores: + # The score-seed operand tile is written once by its owning instance. + eager_init_resources.append(tmem_s0) return ( task_list, @@ -1948,11 +2526,93 @@ def build_decode_task_manager( # ===================================================================== +@cute.jit +def _deallocate_decode_tmem( + tmem_ptr_i32: cute.Pointer, + tmem_alloc_cols: Int32, + cfg: cutlass.Constexpr[FmhaDecodeConfig], +) -> None: + """Retire one CTA's TMEM after all of its task users have stopped.""" + warp_size = 32 + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + dealloc_barrier_id = 13 + if cutlass.const_expr(cfg.use_keeps_mma_ab): + prims.barrier_cta_sync(dealloc_barrier_id) + if ( + warp_idx >= cfg.correction_warp_idx + and warp_idx < cfg.correction_warp_idx + cfg.correction_num_warps + ): + tidx, _, _ = cute.arch.thread_idx() + correction_thread_idx = tidx - cfg.correction_warp_idx * warp_size + if correction_thread_idx < warp_size: + tmem_ptr = prims.make_tmem_ptr(tmem_ptr_i32.load(), cutlass.Int8) + prims.tcgen05_dealloc(tmem_ptr, tmem_alloc_cols) + else: + if ( + warp_idx >= cfg.correction_warp_idx + and warp_idx < cfg.correction_warp_idx + cfg.correction_num_warps + ): + prims.barrier_cta_sync( + dealloc_barrier_id, + thread_count=cfg.correction_num_warps * warp_size, + ) + tidx, _, _ = cute.arch.thread_idx() + correction_thread_idx = tidx - cfg.correction_warp_idx * warp_size + if correction_thread_idx < warp_size: + tmem_ptr = prims.make_tmem_ptr(tmem_ptr_i32.load(), cutlass.Int8) + prims.tcgen05_dealloc(tmem_ptr, tmem_alloc_cols) + + +@cute.jit +def _release_decode_dependents( + cfg: cutlass.Constexpr[FmhaDecodeConfig], +) -> None: + """Uniformly release a following PDL grid after one CTA's true tail.""" + if cutlass.const_expr(cfg.use_parallel_separate_reduction_pdl): + prims.barrier_cta_sync(14) + prims.griddepcontrol(kind=prims.GridDepAction.LAUNCH_DEPENDENTS) + + +@cute.jit +def _exit_cta_if_inactive(is_active: cutlass.Boolean) -> None: + """Retire all threads in a CTA when a uniform activity predicate is false.""" + active_i32 = Int32(is_active) + llvm.inline_asm( + mlir_T.i32(), + [active_i32.ir_value()], + "{ .reg .pred _pexit; setp.eq.s32 _pexit, $1, 0;" + " @_pexit exit; mov.u32 $0, 0; }", + "=r,r", + has_side_effects=True, + is_align_stack=False, + asm_dialect=llvm.AsmDialect.AD_ATT, + ) + + +@cute.jit +def _retire_inactive_pdl_split( + is_active: cutlass.Boolean, + tmem_ptr_i32: cute.Pointer, + tmem_alloc_cols: Int32, + cfg: cutlass.Constexpr[FmhaDecodeConfig], +) -> None: + """Cleanly retire a post-acquire split without staging TaskManager.""" + if not is_active: + _deallocate_decode_tmem(tmem_ptr_i32, tmem_alloc_cols, cfg) + _release_decode_dependents(cfg) + # A runtime branch around ``TaskManager.run`` would make the DSL carry the + # Python TaskManager meta object through ``scf.if``. Retiring only the + # inactive, CTA-uniform path here leaves the active task graph straight-line. + _exit_cta_if_inactive(is_active) + + @cute.jit def _run_decode_gen_active( tma_desc_q: cutlass.GridConstant[cuda.TensorMap], tma_desc_k: cutlass.GridConstant[cuda.TensorMap], tma_desc_v: cutlass.GridConstant[cuda.TensorMap], + tma_desc_k_sf: cutlass.GridConstant[cuda.TensorMap], + tma_desc_v_sf: cutlass.GridConstant[cuda.TensorMap], tma_desc_k_atom: cutlass.GridConstant[cuda.TensorMap], tma_desc_v_atom: cutlass.GridConstant[cuda.TensorMap], o_iter: cute.Pointer, @@ -1960,9 +2620,15 @@ def _run_decode_gen_active( g_h_k: Int32, g_scale_s_log2_e: Float32, g_output_scale: Float32, - g_seqlens_kv: cute.Pointer, + g_seqlens_kv: cute.Pointer | DirectSparseMetadataView, g_cu_seqlens_q: cute.Pointer, - g_page_idx_kv: cute.Pointer, + g_page_idx_kv: cute.Pointer | DirectSparseMetadataView, + g_k_sf: cute.Pointer, + g_v_sf: cute.Pointer, + g_page_table_stride: Int64, + g_page_table_capacity: Int32, + g_q_token_kv_block_sparse_page_memberships: cute.Pointer, + g_q_token_kv_block_sparse_page_membership_stride: Int32, g_partial_o: cute.Pointer, g_partial_stats: cute.Pointer, g_split_kv_counter: cute.Pointer, @@ -1975,15 +2641,13 @@ def _run_decode_gen_active( seq_len_q: Int32, active_splits_kv: Int32, static_full_split_prefix: cutlass.Constexpr[bool], + defer_runtime_split_pruning: cutlass.Constexpr[bool], tile_sched_params: utils.ClcDynamicPersistentTileSchedulerParams | None, cfg: cutlass.Constexpr[FmhaDecodeConfig], seq_len_kv: cutlass.Constexpr[int] = 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, - g_block_tables: cute.Pointer | None = None, - g_block_table_capacity: Int32 | None = None, - g_block_table_row_stride: Int64 | None = None, g_sparse_row_route_offsets: cute.Pointer | None = None, g_sparse_row_route_counts: cute.Pointer | None = None, g_sparse_route_metadata: cute.Pointer | None = None, @@ -1991,6 +2655,14 @@ def _run_decode_gen_active( tma_desc_v_summary: cutlass.GridConstant[cuda.TensorMap] | None = None, tma_desc_k_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, tma_desc_v_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, + g_sage_q_scale: cute.Pointer | None = None, + g_sage_k_scale: cute.Pointer | None = None, + g_sage_k_summary_scale: cute.Pointer | None = None, + g_sage_v_scale: cute.Pointer | None = None, + g_sage_v_mean: cute.Pointer | None = None, + g_sage_q_scale_head_stride: Int32 | None = None, + g_sage_k_scale_head_stride: Int32 | None = None, + g_sage_k_summary_scale_head_stride: Int32 | None = None, ) -> None: """Run the complete decode body for one runtime-valid Q tile. @@ -2002,8 +2674,6 @@ def _run_decode_gen_active( dynamic branch. """ - WARP_SIZE = 32 - warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) bias_static_sliding_kv_tma = ( @@ -2029,7 +2699,6 @@ def _run_decode_gen_active( if cutlass.const_expr(use_runtime_seqlens_kv) else Int32(cfg.static_seq_len_kv) ) - use_clc_dynamic_scheduler = cfg.use_persistent_scheduler tma_desc_k_summary_ptr = None tma_desc_v_summary_ptr = None tma_desc_k_summary_atom_ptr = None @@ -2056,6 +2725,9 @@ def _run_decode_gen_active( prims.prefetch_tensormap(tma_desc_q.get_ptr()) prims.prefetch_tensormap(tma_desc_k.get_ptr()) prims.prefetch_tensormap(tma_desc_v.get_ptr()) + if cutlass.const_expr(cfg.use_nvfp4_kv): + prims.prefetch_tensormap(tma_desc_k_sf.get_ptr()) + prims.prefetch_tensormap(tma_desc_v_sf.get_ptr()) if cutlass.const_expr(cfg.use_block_sparse and uses_atom_desc): # KV256 and non-aligned coarse KV128 may select the exact atom maps. prims.prefetch_tensormap(tma_desc_k_atom.get_ptr()) @@ -2069,12 +2741,6 @@ def _run_decode_gen_active( prims.prefetch_tensormap(tma_desc_v_summary_atom_ptr) init_warp += 1 - clc_response_ptr = None - if cutlass.const_expr(use_clc_dynamic_scheduler): - clc_response_ptr = cute.arch.alloc_smem( - cutlass.Int128, _PERSISTENT_SCHEDULE_TOKEN_STAGES - ) - q_output_rows = g_h_r if cutlass.const_expr(cfg.max_seq_len_q > 1): q_output_rows = g_h_r * Int32(cfg.max_seq_len_q) @@ -2126,28 +2792,38 @@ def _run_decode_gen_active( tma_desc_q=tma_desc_q.get_ptr(), tma_desc_k=tma_desc_k.get_ptr(), tma_desc_v=tma_desc_v.get_ptr(), + tma_desc_k_sf=tma_desc_k_sf.get_ptr(), + tma_desc_v_sf=tma_desc_v_sf.get_ptr(), tma_desc_k_atom=tma_desc_k_atom.get_ptr(), tma_desc_v_atom=tma_desc_v_atom.get_ptr(), tma_desc_k_summary=tma_desc_k_summary_ptr, tma_desc_v_summary=tma_desc_v_summary_ptr, tma_desc_k_summary_atom=tma_desc_k_summary_atom_ptr, tma_desc_v_summary_atom=tma_desc_v_summary_atom_ptr, + sage_q_scale_ptr=g_sage_q_scale, + sage_k_scale_ptr=g_sage_k_scale, + sage_k_summary_scale_ptr=g_sage_k_summary_scale, + sage_v_scale_ptr=g_sage_v_scale, + sage_v_mean_ptr=g_sage_v_mean, + sage_q_scale_head_stride=g_sage_q_scale_head_stride, + sage_k_scale_head_stride=g_sage_k_scale_head_stride, + sage_k_summary_scale_head_stride=g_sage_k_summary_scale_head_stride, page_idx_kv=g_page_idx_kv, + page_table_stride=g_page_table_stride, + page_table_capacity=g_page_table_capacity, + q_token_kv_block_sparse_page_memberships=g_q_token_kv_block_sparse_page_memberships, + q_token_kv_block_sparse_page_membership_stride=g_q_token_kv_block_sparse_page_membership_stride, h_k_idx=h_k_idx, b_idx=b_idx, q_group_idx=q_group_idx, q_token_offset=q_token_offset, seq_len_q=seq_len_q, - active_splits_kv=active_splits_kv, + active_splits_kv=(None if defer_runtime_split_pruning else active_splits_kv), static_full_split_prefix=static_full_split_prefix, tile_sched_params=tile_sched_params, - clc_response_ptr=clc_response_ptr, use_variable_seqlens_kv=use_variable_seqlens_kv, use_native_paged_kv=use_native_paged_kv, use_static_native_seqlens_kv=use_static_native_seqlens_kv, - block_tables=g_block_tables, - block_table_capacity=g_block_table_capacity, - block_table_row_stride=g_block_table_row_stride, sparse_row_route_offsets=g_sparse_row_route_offsets, sparse_row_route_counts=g_sparse_row_route_counts, sparse_route_metadata=g_sparse_route_metadata, @@ -2173,7 +2849,14 @@ def _run_decode_gen_active( dma_consumer_release_labels=dma_consumer_release_labels, skip_validation=True, verbose=False, - exhaustive_deadlock_race_check=not _has_unmodeled_tmem_p_alias_protocol(cfg), + # Kernel construction is a latency-sensitive JIT path and may run once + # per CUDA-graph bucket on every TP rank. Keep the cheap structural + # checks here, but leave the exhaustive interleaving proof to + # build_decode_task_manager() and its offline tests. `skip_validation` + # only changes failures into warnings; it does not skip the expensive + # state-space search. + exhaustive_deadlock_race_check=False, + assume_pdl_wait_completed=cfg.use_pdl, smem_allocator=smem_allocator, tmem_allocator=tmem_allocator, ) @@ -2197,15 +2880,42 @@ def _run_decode_gen_active( # After this, peers may safely signal an owner CTA's mbarrier via mapa. prims.barrier_cluster_arrive_relaxed() prims.barrier_cluster_wait() + if cutlass.const_expr(cfg.use_pdl): + # Delay the acquire until every CTA-local pipeline barrier and + # SMEM/TMEM resource is initialized, but keep it before TaskManager + # resolves any value published by the preceding producer grid. + prims.griddepcontrol(kind=prims.GridDepAction.WAIT) + if cutlass.const_expr(defer_runtime_split_pruning): + # A split launch cannot read its producer-owned sequence length until + # the dependency above is acquired. Every physical split CTA + # initializes the same static task resources first; only the useful + # runtime prefix enters TaskManager. Pruned CTAs take the equivalent + # TMEM teardown and dependent-release helpers, then retire before the + # straight-line task graph. + runtime_seq_len_kv = g_s_k + if cutlass.const_expr( + use_variable_seqlens_kv + 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 + ) + _retire_inactive_pdl_split( + run_task_graph, + tmem_ptr_i32, + Int32(tmem_alloc_cols), + cfg, + ) task_manager.run() - if cutlass.const_expr(cfg.use_parallel_separate_reduction_pdl): - # Publish the dependent reducer only after this producer CTA has made - # its normalized partial O and log2-LSE stores visible. - thread_idx, _, _ = cute.arch.thread_idx() - if thread_idx == Int32(0): - prims.griddepcontrol(kind=prims.GridDepAction.LAUNCH_DEPENDENTS) - # Every peer async-stores into an owner CTA's distributed SMEM and charges # the bytes to that owner's transaction mbarrier. Producer-only CTAs are # never remote-SMEM targets, so they may retire after their store issues; @@ -2213,37 +2923,11 @@ def _run_decode_gen_active( # reduction consumes all delivered partials. A final cluster rendezvous # would therefore only make completed producers wait for the owners. - # ── TMEM cleanup ── - dealloc_barrier_id = 13 - if cutlass.const_expr(cfg.use_keeps_mma_ab): - # Keeps aliases score/stat TMEM columns across task warp groups. Wait - # for every task's producer tail before the correction owner warp - # releases the allocation; a correction-only barrier can race those - # tails and leave tcgen05.dealloc waiting on live TMEM users. - prims.barrier_cta_sync(dealloc_barrier_id) - if ( - warp_idx >= cfg.correction_warp_idx - and warp_idx < cfg.correction_warp_idx + cfg.correction_num_warps - ): - tidx, _, _ = cute.arch.thread_idx() - correction_thread_idx = tidx - cfg.correction_warp_idx * WARP_SIZE - if correction_thread_idx < WARP_SIZE: - tmem_ptr = prims.make_tmem_ptr(tmem_ptr_i32.load(), cutlass.Int8) - prims.tcgen05_dealloc(tmem_ptr, Int32(tmem_alloc_cols)) - else: - if ( - warp_idx >= cfg.correction_warp_idx - and warp_idx < cfg.correction_warp_idx + cfg.correction_num_warps - ): - prims.barrier_cta_sync( - dealloc_barrier_id, - thread_count=cfg.correction_num_warps * WARP_SIZE, - ) - tidx, _, _ = cute.arch.thread_idx() - correction_thread_idx = tidx - cfg.correction_warp_idx * WARP_SIZE - if correction_thread_idx < WARP_SIZE: - tmem_ptr = prims.make_tmem_ptr(tmem_ptr_i32.load(), cutlass.Int8) - prims.tcgen05_dealloc(tmem_ptr, Int32(tmem_alloc_cols)) + # Keeps aliases score/stat TMEM columns across task warp groups, while the + # swaps path needs only its correction-warp rendezvous. In both cases the + # following PDL grid is released only after deallocation completes. + _deallocate_decode_tmem(tmem_ptr_i32, Int32(tmem_alloc_cols), cfg) + _release_decode_dependents(cfg) @cute.jit @@ -2269,11 +2953,13 @@ def _run_decode_gen_inactive_cluster_rank() -> None: @cute.jit def _signal_padded_pdl_producer(cfg: cutlass.Constexpr[FmhaDecodeConfig]) -> None: - """Release the dependent reducer launch from a zero-work producer CTA.""" + """Preserve PDL ordering through a zero-work producer CTA.""" + if cutlass.const_expr(cfg.use_pdl): + # This CTA has no task resources or producer-owned work, but it must + # acquire the incoming dependency before retiring. + prims.griddepcontrol(kind=prims.GridDepAction.WAIT) if cutlass.const_expr(cfg.use_parallel_separate_reduction_pdl): - thread_idx, _, _ = cute.arch.thread_idx() - if thread_idx == Int32(0): - prims.griddepcontrol(kind=prims.GridDepAction.LAUNCH_DEPENDENTS) + prims.griddepcontrol(kind=prims.GridDepAction.LAUNCH_DEPENDENTS) @cute.jit @@ -2281,6 +2967,8 @@ def _run_decode_gen_runtime_prefix( tma_desc_q: cutlass.GridConstant[cuda.TensorMap], tma_desc_k: cutlass.GridConstant[cuda.TensorMap], tma_desc_v: cutlass.GridConstant[cuda.TensorMap], + tma_desc_k_sf: cutlass.GridConstant[cuda.TensorMap], + tma_desc_v_sf: cutlass.GridConstant[cuda.TensorMap], tma_desc_k_atom: cutlass.GridConstant[cuda.TensorMap], tma_desc_v_atom: cutlass.GridConstant[cuda.TensorMap], o_iter: cute.Pointer, @@ -2288,9 +2976,15 @@ def _run_decode_gen_runtime_prefix( g_h_k: Int32, g_scale_s_log2_e: Float32, g_output_scale: Float32, - g_seqlens_kv: cute.Pointer, + g_seqlens_kv: cute.Pointer | DirectSparseMetadataView, g_cu_seqlens_q: cute.Pointer, - g_page_idx_kv: cute.Pointer, + g_page_idx_kv: cute.Pointer | DirectSparseMetadataView, + g_k_sf: cute.Pointer, + g_v_sf: cute.Pointer, + g_page_table_stride: Int64, + g_page_table_capacity: Int32, + g_q_token_kv_block_sparse_page_memberships: cute.Pointer, + g_q_token_kv_block_sparse_page_membership_stride: Int32, g_partial_o: cute.Pointer, g_partial_stats: cute.Pointer, g_split_kv_counter: cute.Pointer, @@ -2309,9 +3003,6 @@ def _run_decode_gen_runtime_prefix( use_variable_seqlens_kv: cutlass.Constexpr[bool], use_native_paged_kv: cutlass.Constexpr[bool], use_static_native_seqlens_kv: cutlass.Constexpr[bool], - g_block_tables: cute.Pointer | None, - g_block_table_capacity: Int32 | None, - g_block_table_row_stride: Int64 | None, g_sparse_row_route_offsets: cute.Pointer | None, g_sparse_row_route_counts: cute.Pointer | None, g_sparse_route_metadata: cute.Pointer | None, @@ -2319,6 +3010,14 @@ def _run_decode_gen_runtime_prefix( tma_desc_v_summary: cutlass.GridConstant[cuda.TensorMap] | None = None, tma_desc_k_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, tma_desc_v_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, + g_sage_q_scale: cute.Pointer | None = None, + g_sage_k_scale: cute.Pointer | None = None, + g_sage_k_summary_scale: cute.Pointer | None = None, + g_sage_v_scale: cute.Pointer | None = None, + g_sage_v_mean: cute.Pointer | None = None, + g_sage_q_scale_head_stride: Int32 | None = None, + g_sage_k_scale_head_stride: Int32 | None = None, + g_sage_k_summary_scale_head_stride: Int32 | None = None, ) -> None: """Run the general runtime split-prefix producer or retire its suffix.""" @@ -2340,6 +3039,16 @@ def _run_decode_gen_runtime_prefix( if cutlass.const_expr(not cfg.use_separate_reduction_kernel): # Preserve one neutral producer for fused empty-K semantics. active_splits_kv = cute.math.max(active_splits_kv, Int32(1)) + if cutlass.const_expr( + cfg.uses_q_token_kv_block_sparse_page_route + and not cfg.supports_cluster_smem_reduction + ): + # When only the final configured split is empty, retaining its + # neutral producer lets the fused GMEM reducer use the fully + # unrolled prefix instead of its dynamically guarded fallback. + # 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_is_active = split_idx < active_splits_kv @@ -2349,6 +3058,8 @@ def _run_decode_gen_runtime_prefix( tma_desc_q, tma_desc_k, tma_desc_v, + tma_desc_k_sf, + tma_desc_v_sf, tma_desc_k_atom, tma_desc_v_atom, o_iter, @@ -2359,6 +3070,12 @@ def _run_decode_gen_runtime_prefix( g_seqlens_kv, g_cu_seqlens_q, g_page_idx_kv, + g_k_sf, + g_v_sf, + g_page_table_stride, + g_page_table_capacity, + g_q_token_kv_block_sparse_page_memberships, + g_q_token_kv_block_sparse_page_membership_stride, g_partial_o, g_partial_stats, g_split_kv_counter, @@ -2371,15 +3088,13 @@ def _run_decode_gen_runtime_prefix( seq_len_q, active_splits_kv, False, + False, tile_sched_params, cfg, seq_len_kv, use_variable_seqlens_kv, use_native_paged_kv, use_static_native_seqlens_kv, - g_block_tables, - g_block_table_capacity, - g_block_table_row_stride, g_sparse_row_route_offsets, g_sparse_row_route_counts, g_sparse_route_metadata, @@ -2387,6 +3102,14 @@ def _run_decode_gen_runtime_prefix( tma_desc_v_summary=tma_desc_v_summary, tma_desc_k_summary_atom=tma_desc_k_summary_atom, tma_desc_v_summary_atom=tma_desc_v_summary_atom, + g_sage_q_scale=g_sage_q_scale, + g_sage_k_scale=g_sage_k_scale, + g_sage_k_summary_scale=g_sage_k_summary_scale, + g_sage_v_scale=g_sage_v_scale, + g_sage_v_mean=g_sage_v_mean, + g_sage_q_scale_head_stride=g_sage_q_scale_head_stride, + g_sage_k_scale_head_stride=g_sage_k_scale_head_stride, + g_sage_k_summary_scale_head_stride=g_sage_k_summary_scale_head_stride, ) else: _run_decode_gen_inactive_cluster_rank() @@ -2396,6 +3119,8 @@ def _run_decode_gen_runtime_prefix( tma_desc_q, tma_desc_k, tma_desc_v, + tma_desc_k_sf, + tma_desc_v_sf, tma_desc_k_atom, tma_desc_v_atom, o_iter, @@ -2406,6 +3131,12 @@ def _run_decode_gen_runtime_prefix( g_seqlens_kv, g_cu_seqlens_q, g_page_idx_kv, + g_k_sf, + g_v_sf, + g_page_table_stride, + g_page_table_capacity, + g_q_token_kv_block_sparse_page_memberships, + g_q_token_kv_block_sparse_page_membership_stride, g_partial_o, g_partial_stats, g_split_kv_counter, @@ -2418,15 +3149,13 @@ def _run_decode_gen_runtime_prefix( seq_len_q, active_splits_kv, False, + False, tile_sched_params, cfg, seq_len_kv, use_variable_seqlens_kv, use_native_paged_kv, use_static_native_seqlens_kv, - g_block_tables, - g_block_table_capacity, - g_block_table_row_stride, g_sparse_row_route_offsets, g_sparse_row_route_counts, g_sparse_route_metadata, @@ -2434,6 +3163,14 @@ def _run_decode_gen_runtime_prefix( tma_desc_v_summary=tma_desc_v_summary, tma_desc_k_summary_atom=tma_desc_k_summary_atom, tma_desc_v_summary_atom=tma_desc_v_summary_atom, + g_sage_q_scale=g_sage_q_scale, + g_sage_k_scale=g_sage_k_scale, + g_sage_k_summary_scale=g_sage_k_summary_scale, + g_sage_v_scale=g_sage_v_scale, + g_sage_v_mean=g_sage_v_mean, + g_sage_q_scale_head_stride=g_sage_q_scale_head_stride, + g_sage_k_scale_head_stride=g_sage_k_scale_head_stride, + g_sage_k_summary_scale_head_stride=g_sage_k_summary_scale_head_stride, ) else: _signal_padded_pdl_producer(cfg) @@ -2444,6 +3181,8 @@ def decode_gen_kernel( tma_desc_q: cutlass.GridConstant[cuda.TensorMap], tma_desc_k: cutlass.GridConstant[cuda.TensorMap], tma_desc_v: cutlass.GridConstant[cuda.TensorMap], + tma_desc_k_sf: cutlass.GridConstant[cuda.TensorMap], + tma_desc_v_sf: cutlass.GridConstant[cuda.TensorMap], tma_desc_k_atom: cutlass.GridConstant[cuda.TensorMap], tma_desc_v_atom: cutlass.GridConstant[cuda.TensorMap], o_iter: cute.Pointer, @@ -2451,9 +3190,15 @@ def decode_gen_kernel( g_h_k: Int32, g_scale_s_log2_e: Float32, g_output_scale: Float32, - g_seqlens_kv: cute.Pointer, + g_seqlens_kv: cute.Pointer | DirectSparseMetadataView, g_cu_seqlens_q: cute.Pointer, - g_page_idx_kv: cute.Pointer, + g_page_idx_kv: cute.Pointer | DirectSparseMetadataView, + g_k_sf: cute.Pointer, + g_v_sf: cute.Pointer, + g_page_table_stride: Int64, + g_page_table_capacity: Int32, + g_q_token_kv_block_sparse_page_memberships: cute.Pointer, + g_q_token_kv_block_sparse_page_membership_stride: Int32, g_partial_o: cute.Pointer, g_partial_stats: cute.Pointer, g_split_kv_counter: cute.Pointer, @@ -2465,9 +3210,6 @@ def decode_gen_kernel( 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, - g_block_tables: cute.Pointer | None = None, - g_block_table_capacity: Int32 | None = None, - g_block_table_row_stride: Int64 | None = None, g_sparse_row_route_offsets: cute.Pointer | None = None, g_sparse_row_route_counts: cute.Pointer | None = None, g_sparse_route_metadata: cute.Pointer | None = None, @@ -2476,9 +3218,36 @@ def decode_gen_kernel( tma_desc_v_summary: cutlass.GridConstant[cuda.TensorMap] | None = None, tma_desc_k_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, tma_desc_v_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, + g_sage_q_scale: cute.Pointer | None = None, + g_sage_k_scale: cute.Pointer | None = None, + g_sage_k_summary_scale: cute.Pointer | None = None, + g_sage_v_scale: cute.Pointer | None = None, + g_sage_v_mean: cute.Pointer | None = None, + g_sage_q_scale_head_stride: Int32 | None = None, + g_sage_k_scale_head_stride: Int32 | None = None, + g_sage_k_summary_scale_head_stride: Int32 | None = None, ) -> 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( + cfg.use_q_token_kv_block_sparse_route + and not cfg.shares_sparse_pattern + and not cfg.use_persistent_scheduler + ): + if cutlass.const_expr(isinstance(g_page_idx_kv, DirectSparseMetadataView)): + g_page_idx_kv = g_page_idx_kv.with_head(h_k_idx) + else: + g_seqlens_kv = HeadIndexedMetadataView(g_seqlens_kv, g_h_k, h_k_idx) + g_page_idx_kv = g_page_idx_kv + Int64(h_k_idx) * g_page_table_stride + g_page_table_stride = g_page_table_stride * Int64(g_h_k) + g_q_token_kv_block_sparse_page_memberships = ( + g_q_token_kv_block_sparse_page_memberships + + Int64(h_k_idx) + * Int64(g_q_token_kv_block_sparse_page_membership_stride) + ) + g_q_token_kv_block_sparse_page_membership_stride = ( + 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. @@ -2503,6 +3272,57 @@ def decode_gen_kernel( tma_desc_q, tma_desc_k, tma_desc_v, + tma_desc_k_sf, + tma_desc_v_sf, + tma_desc_k_atom, + tma_desc_v_atom, + o_iter, + g_s_k, + g_h_k, + g_scale_s_log2_e, + g_output_scale, + g_seqlens_kv, + g_cu_seqlens_q, + g_page_idx_kv, + g_k_sf, + g_v_sf, + g_page_table_stride, + g_page_table_capacity, + g_q_token_kv_block_sparse_page_memberships, + g_q_token_kv_block_sparse_page_membership_stride, + g_partial_o, + g_partial_stats, + g_split_kv_counter, + g_attention_sinks, + g_h_r, + q_group_idx, + h_k_idx, + b_idx, + q_token_offset, + seq_len_q, + Int32(cfg.splits_kv), + True, + False, + tile_sched_params, + cfg, + seq_len_kv, + use_variable_seqlens_kv, + use_native_paged_kv, + use_static_native_seqlens_kv, + g_sparse_row_route_offsets, + g_sparse_row_route_counts, + g_sparse_route_metadata, + ) + elif cutlass.const_expr(cfg.use_split_kv and cfg.use_pdl): + # The immediately preceding grid may produce seq_lens. Send every + # configured split through resource initialization and defer + # useful-prefix pruning until after the PDL acquire. + _run_decode_gen_active( + tma_desc_q, + tma_desc_k, + tma_desc_v, + tma_desc_k_sf, + tma_desc_v_sf, tma_desc_k_atom, tma_desc_v_atom, o_iter, @@ -2513,6 +3333,12 @@ def decode_gen_kernel( g_seqlens_kv, g_cu_seqlens_q, g_page_idx_kv, + g_k_sf, + g_v_sf, + g_page_table_stride, + g_page_table_capacity, + g_q_token_kv_block_sparse_page_memberships, + g_q_token_kv_block_sparse_page_membership_stride, g_partial_o, g_partial_stats, g_split_kv_counter, @@ -2524,6 +3350,7 @@ def decode_gen_kernel( q_token_offset, seq_len_q, Int32(cfg.splits_kv), + False, True, tile_sched_params, cfg, @@ -2531,9 +3358,6 @@ def decode_gen_kernel( use_variable_seqlens_kv, use_native_paged_kv, use_static_native_seqlens_kv, - g_block_tables, - g_block_table_capacity, - g_block_table_row_stride, g_sparse_row_route_offsets, g_sparse_row_route_counts, g_sparse_route_metadata, @@ -2541,12 +3365,22 @@ def decode_gen_kernel( tma_desc_v_summary=tma_desc_v_summary, tma_desc_k_summary_atom=tma_desc_k_summary_atom, tma_desc_v_summary_atom=tma_desc_v_summary_atom, + g_sage_q_scale=g_sage_q_scale, + g_sage_k_scale=g_sage_k_scale, + g_sage_k_summary_scale=g_sage_k_summary_scale, + g_sage_v_scale=g_sage_v_scale, + g_sage_v_mean=g_sage_v_mean, + g_sage_q_scale_head_stride=g_sage_q_scale_head_stride, + g_sage_k_scale_head_stride=g_sage_k_scale_head_stride, + g_sage_k_summary_scale_head_stride=g_sage_k_summary_scale_head_stride, ) else: _run_decode_gen_runtime_prefix( tma_desc_q, tma_desc_k, tma_desc_v, + tma_desc_k_sf, + tma_desc_v_sf, tma_desc_k_atom, tma_desc_v_atom, o_iter, @@ -2557,6 +3391,12 @@ def decode_gen_kernel( g_seqlens_kv, g_cu_seqlens_q, g_page_idx_kv, + g_k_sf, + g_v_sf, + g_page_table_stride, + g_page_table_capacity, + g_q_token_kv_block_sparse_page_memberships, + g_q_token_kv_block_sparse_page_membership_stride, g_partial_o, g_partial_stats, g_split_kv_counter, @@ -2575,9 +3415,6 @@ def decode_gen_kernel( use_variable_seqlens_kv, use_native_paged_kv, use_static_native_seqlens_kv, - g_block_tables, - g_block_table_capacity, - g_block_table_row_stride, g_sparse_row_route_offsets, g_sparse_row_route_counts, g_sparse_route_metadata, @@ -2585,6 +3422,14 @@ def decode_gen_kernel( tma_desc_v_summary=tma_desc_v_summary, tma_desc_k_summary_atom=tma_desc_k_summary_atom, tma_desc_v_summary_atom=tma_desc_v_summary_atom, + g_sage_q_scale=g_sage_q_scale, + g_sage_k_scale=g_sage_k_scale, + g_sage_k_summary_scale=g_sage_k_summary_scale, + g_sage_v_scale=g_sage_v_scale, + g_sage_v_mean=g_sage_v_mean, + g_sage_q_scale_head_stride=g_sage_q_scale_head_stride, + g_sage_k_scale_head_stride=g_sage_k_scale_head_stride, + g_sage_k_summary_scale_head_stride=g_sage_k_summary_scale_head_stride, ) else: # Packed-Q grids use a batch-wide maximum envelope. These Q CTAs own no @@ -2598,11 +3443,14 @@ def fmha_decode_launch( q_iter: cute.Pointer, k_iter: cute.Pointer, v_iter: cute.Pointer, + k_sf_iter: cute.Pointer, + v_sf_iter: cute.Pointer, o_iter: cute.Pointer, - seqlens_kv_iter: cute.Pointer, + seqlens_kv_iter: cute.Pointer | DirectSparseMetadataView, cu_seqlens_q_iter: cute.Pointer, total_q_tokens: Int32, - page_idx_kv_iter: cute.Pointer, + page_idx_kv_iter: cute.Pointer | DirectSparseMetadataView, + q_token_kv_block_sparse_page_memberships_iter: cute.Pointer, partial_o_iter: cute.Pointer, partial_stats_iter: cute.Pointer, split_kv_counter_iter: cute.Pointer, @@ -2616,16 +3464,29 @@ def fmha_decode_launch( seq_len_kv: cutlass.Constexpr[int] = 2048, use_variable_seqlens_kv: cutlass.Constexpr[bool] = False, use_native_paged_kv: cutlass.Constexpr[bool] = False, - block_tables_iter: cute.Pointer | None = None, - block_table_capacity: Int32 = 0, - block_table_row_stride: Int64 = 0, + page_table_stride: Int64 = 0, + page_table_capacity: Int32 = 0, + q_token_kv_block_sparse_page_membership_stride: Int32 = 0, num_physical_kv_pages: Int64 = 0, k_page_stride: Int64 = 0, + k_head_stride: Int64 = 0, + k_token_stride: Int64 = 0, v_page_stride: Int64 = 0, + v_head_stride: Int64 = 0, + v_token_stride: Int64 = 0, static_full_split_prefix: cutlass.Constexpr[bool] = False, use_static_native_seqlens_kv: cutlass.Constexpr[bool] = False, + sage_q_scale_iter: cute.Pointer | None = None, + sage_k_scale_iter: cute.Pointer | None = None, + sage_v_scale_iter: cute.Pointer | None = None, + sage_v_mean_iter: cute.Pointer | None = None, + sage_q_scale_head_stride: Int32 | None = None, + sage_k_scale_head_stride: Int32 | None = None, ) -> None: - """Standalone JIT launcher for FMHA decode TS.""" + """Standalone JIT launcher for FMHA decode TS. + + The ``sage_*`` scale arguments are used only by a Sage attention config. + """ log2_e = math.log2(math.e) b, h_q, h_k, s_k, d = problem_shape h_r = h_q // h_k @@ -2642,29 +3503,51 @@ def fmha_decode_launch( q_seq = Int32(cfg.max_seq_len_q) if cutlass.const_expr(cfg.use_paged_kv): + kv_tma_d = d + if cutlass.const_expr( + cfg.use_nvfp4_kv and not cfg.store_transformed_kv_in_tmem + ): + kv_tma_d = d // Int32(2) if cutlass.const_expr(use_native_paged_kv): + total_pages = num_physical_kv_pages + storage_tokens_per_page = Int32(cfg.effective_storage_tokens_per_page) kv_shape = ( - d, - Int32(cfg.num_tokens_per_page), + kv_tma_d, + storage_tokens_per_page, h_k, - num_physical_kv_pages, + total_pages, ) + k_tma_page_stride = k_page_stride + k_tma_head_stride = k_head_stride + k_tma_token_stride = k_token_stride + v_tma_page_stride = v_page_stride + v_tma_head_stride = v_head_stride + v_tma_token_stride = v_token_stride + if cutlass.const_expr(cfg.store_transformed_kv_in_tmem): + # B4X16_P64 tensor-map strides are expressed in logical FP4 + # lanes while the native cache reports packed-byte strides. + k_tma_page_stride = k_page_stride * Int64(2) + k_tma_head_stride = k_head_stride * Int64(2) + k_tma_token_stride = k_token_stride * Int64(2) + v_tma_page_stride = v_page_stride * Int64(2) + v_tma_head_stride = v_head_stride * Int64(2) + v_tma_token_stride = v_token_stride * Int64(2) k_layout = cute.make_layout( kv_shape, stride=( 1, - d, - d * Int32(cfg.num_tokens_per_page), - k_page_stride, + k_tma_token_stride, + k_tma_head_stride, + k_tma_page_stride, ), ) v_layout = cute.make_layout( kv_shape, stride=( 1, - d, - d * Int32(cfg.num_tokens_per_page), - v_page_stride, + v_tma_token_stride, + v_tma_head_stride, + v_tma_page_stride, ), ) k_tma = cute.make_tensor(k_iter, k_layout) @@ -2672,12 +3555,12 @@ def fmha_decode_launch( else: total_pages = b * Int32(cfg.max_num_pages_per_seq_kv) kv_layout = cute.make_layout( - (d, Int32(cfg.num_tokens_per_page), h_k, total_pages), + (kv_tma_d, Int32(cfg.num_tokens_per_page), h_k, total_pages), stride=( 1, - d, - d * Int32(cfg.num_tokens_per_page), - d * Int32(cfg.num_tokens_per_page) * h_k, + kv_tma_d, + kv_tma_d * Int32(cfg.num_tokens_per_page), + kv_tma_d * Int32(cfg.num_tokens_per_page) * h_k, ), ) k_tma_iter = k_iter @@ -2692,12 +3575,14 @@ def fmha_decode_launch( skipped_tokens = Int32( _compute_static_num_skipped_kv_tiles(cfg, seq_len_kv) * cfg.tile_size_kv ) - skipped_elems = skipped_tokens * d + skipped_elems = skipped_tokens * d * h_k k_tma_iter = k_iter + skipped_elems v_tma_iter = v_iter + skipped_elems kv_s_for_tma = Int32(effective_seq_len_kv) + # Contiguous K/V are compact BSHD, matching the public PrimTS tensor + # contract and the Q layout above. kv_layout = cute.make_layout( - (d, kv_s_for_tma, h_k, b), stride=(1, d, d * s_k, kv_b_stride) + (d, kv_s_for_tma, h_k, b), stride=(1, d * h_k, d, kv_b_stride) ) k_tma = cute.make_tensor(k_tma_iter, kv_layout) v_tma = cute.make_tensor(v_tma_iter, kv_layout) @@ -2705,10 +3590,30 @@ def fmha_decode_launch( # Keep the TMA inner box at 128B when possible, but never exceed headDim. # box_dim is expressed in elements of the source dtype, not bytes. tma_box0_q = min(128 // cfg.q_dtype_bytes, cfg.headdim) - tma_box0_kv = min(128 // cfg.kv_dtype_bytes, cfg.headdim) - tma_swizzle = cuda.TensorMapSwizzle.s128b - if cutlass.const_expr(cfg.use_fp8_qkv and cfg.headdim == 64): - tma_swizzle = cuda.TensorMapSwizzle.s64b + tma_box0_k = min(128 // cfg.k_dtype_bytes, cfg.headdim) + tma_box0_v = min(128 // cfg.v_dtype_bytes, cfg.headdim) + tma_swizzle_q = cuda.TensorMapSwizzle.s128b + if cutlass.const_expr(cfg.q_dtype_bytes == 1 and cfg.headdim == 64): + tma_swizzle_q = cuda.TensorMapSwizzle.s64b + tma_swizzle_k = cuda.TensorMapSwizzle.s128b + if cutlass.const_expr(cfg.k_dtype_bytes == 1 and cfg.headdim == 64): + tma_swizzle_k = cuda.TensorMapSwizzle.s64b + tma_swizzle_v = cuda.TensorMapSwizzle.s128b + if cutlass.const_expr(cfg.v_dtype_bytes == 1 and cfg.headdim == 64): + tma_swizzle_v = cuda.TensorMapSwizzle.s64b + if cutlass.const_expr(cfg.use_nvfp4_kv): + if cutlass.const_expr(cfg.store_transformed_kv_in_tmem): + # Unpacking TMA consumes 128 logical FP4 lanes and materializes + # one s128b-swizzled byte per lane in the raw SMEM stage. + tma_box0_k = cfg.head_dim_kv_stage + tma_box0_v = cfg.head_dim_kv_stage + tma_swizzle_k = cuda.TensorMapSwizzle.s128b + tma_swizzle_v = cuda.TensorMapSwizzle.s128b + else: + tma_box0_k = cfg.head_dim_kv_stage // 2 + tma_box0_v = cfg.head_dim_kv_stage // 2 + tma_swizzle_k = cuda.TensorMapSwizzle.none + tma_swizzle_v = cuda.TensorMapSwizzle.none if cutlass.const_expr(cfg.tile_size_kv == 256): # The 2x2 datapath consumes K in a (0, 2, 1, 3) KV64 permutation. # A KV64 TensorMap atom lets the shared load resource place each @@ -2756,7 +3661,7 @@ def fmha_decode_launch( box_dims=q_box_dims, ragged_dim=2, stride_order=(0, 1, 2), - swizzle=tma_swizzle, + swizzle=tma_swizzle_q, ) else: q_tma = cute.make_tensor( @@ -2786,20 +3691,79 @@ def fmha_decode_launch( q_tma, box_dims=q_box_dims, stride_order=(0, 1, 2, 3, 4), - swizzle=tma_swizzle, + swizzle=tma_swizzle_q, + ) + if cutlass.const_expr(cfg.store_transformed_kv_in_tmem): + tma_desc_k = create_tensor_map_tiled_from_view( + k_tma, + box_dims=(tma_box0_k, tma_kv_tokens, 1, 1), + stride_order=(0, 1, 2, 3), + swizzle=tma_swizzle_k, + dtype=cutlass.Float4E2M1FNx2, + tma_format=cuda.TensorMapDataFormat.B4X16_P64, + ) + tma_desc_v = create_tensor_map_tiled_from_view( + v_tma, + box_dims=(tma_box0_v, tma_kv_tokens, 1, 1), + stride_order=(0, 1, 2, 3), + swizzle=tma_swizzle_v, + dtype=cutlass.Float4E2M1FNx2, + tma_format=cuda.TensorMapDataFormat.B4X16_P64, + ) + else: + tma_desc_k = create_tensor_map_tiled_from_view( + k_tma, + box_dims=(tma_box0_k, tma_kv_tokens, 1, 1), + stride_order=(0, 1, 2, 3), + swizzle=tma_swizzle_k, + ) + tma_desc_v = create_tensor_map_tiled_from_view( + v_tma, + box_dims=(tma_box0_v, tma_kv_tokens, 1, 1), + stride_order=(0, 1, 2, 3), + swizzle=tma_swizzle_v, + ) + + # NVFP4 scale factors are staged into SMEM by the same SmemKv TMA pipeline + # as the packed K/V. The SF tensor is logically + # (headdim // 16, storage_tokens_per_page, h_k, total_pages) E4M3. K uses + # token-major layout and V uses TRT-LLM's 4-token interleaved layout. The + # inner SF box (headdim // 16 = 8 B) is below TMA's + # 16 B minimum, so fold `r` tokens into dim0 for up to a 128 B inner box, + # without spanning semantic page fragments. + # This reshape is a pure reinterpretation of the same contiguous bytes. + tma_desc_k_sf = tma_desc_k + tma_desc_v_sf = tma_desc_v + if cutlass.const_expr(cfg.use_nvfp4_kv): + sf_per_token = cfg.headdim // 16 + sf_r = cfg.sf_tma_reshape_factor + sf_inner = sf_per_token * sf_r + sf_storage_tokens_per_page = cfg.num_tokens_per_page + if cutlass.const_expr(use_native_paged_kv): + sf_storage_tokens_per_page = cfg.effective_storage_tokens_per_page + sf_tokens_outer = sf_storage_tokens_per_page // sf_r + # The descriptor spans a physical page, but each transaction copies + # only the semantic fragment selected by the decoded page locator. + sf_box_tokens_outer = cfg.num_tokens_per_page // sf_r + sf_page_elems = Int32(sf_per_token * sf_storage_tokens_per_page) + sf_layout = cute.make_layout( + (Int32(sf_inner), Int32(sf_tokens_outer), h_k, total_pages), + stride=(1, Int32(sf_inner), sf_page_elems, sf_page_elems * h_k), + ) + k_sf_tma = cute.make_tensor(k_sf_iter, sf_layout) + v_sf_tma = cute.make_tensor(v_sf_iter, sf_layout) + tma_desc_k_sf = cuda.create_tensor_map_tiled_from_view( + k_sf_tma, + box_dims=(sf_inner, sf_box_tokens_outer, 1, 1), + stride_order=(0, 1, 2, 3), + swizzle=cuda.TensorMapSwizzle.none, + ) + tma_desc_v_sf = cuda.create_tensor_map_tiled_from_view( + v_sf_tma, + box_dims=(sf_inner, sf_box_tokens_outer, 1, 1), + stride_order=(0, 1, 2, 3), + swizzle=cuda.TensorMapSwizzle.none, ) - tma_desc_k = create_tensor_map_tiled_from_view( - k_tma, - box_dims=(tma_box0_kv, tma_kv_tokens, 1, 1), - stride_order=(0, 1, 2, 3), - swizzle=tma_swizzle, - ) - tma_desc_v = create_tensor_map_tiled_from_view( - v_tma, - box_dims=(tma_box0_kv, tma_kv_tokens, 1, 1), - stride_order=(0, 1, 2, 3), - swizzle=tma_swizzle, - ) grid_x = q_groups if cutlass.const_expr(cfg.use_split_kv): @@ -2829,6 +3793,8 @@ def fmha_decode_launch( 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, @@ -2839,6 +3805,12 @@ def fmha_decode_launch( 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, @@ -2850,20 +3822,26 @@ def fmha_decode_launch( use_variable_seqlens_kv, use_native_paged_kv, use_static_native_seqlens_kv, - block_tables_iter, - block_table_capacity, - block_table_row_stride, null_sparse_route_ptr, null_sparse_route_ptr, null_sparse_route_ptr, static_full_split_prefix, + g_sage_q_scale=sage_q_scale_iter, + g_sage_k_scale=sage_k_scale_iter, + g_sage_v_scale=sage_v_scale_iter, + g_sage_v_mean=sage_v_mean_iter, + g_sage_q_scale_head_stride=sage_q_scale_head_stride, + g_sage_k_scale_head_stride=sage_k_scale_head_stride, ).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), - use_pdl=cfg.use_parallel_separate_reduction_pdl, + # 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, ) @@ -2888,6 +3866,14 @@ def fmha_block_sparse_launch( num_physical_kv_pages: Int64 = 0, k_page_stride: Int64 = 0, v_page_stride: Int64 = 0, + sage_q_scale_iter: cute.Pointer | None = None, + sage_k_scale_iter: cute.Pointer | None = None, + sage_k_summary_scale_iter: cute.Pointer | None = None, + sage_v_scale_iter: cute.Pointer | None = None, + sage_v_mean_iter: cute.Pointer | None = None, + sage_q_scale_head_stride: Int32 | None = None, + sage_k_scale_head_stride: Int32 | None = None, + sage_k_summary_scale_head_stride: Int32 | None = None, ) -> None: """Launch attention over exact and typed exact/proxy prepared KV routes. @@ -2896,6 +3882,8 @@ def fmha_block_sparse_launch( words. Exact routes address K/V; proxy routes address summary K/V. Both layouts execute the same ``decode_gen_kernel`` schedule and physical copy policy. Exact builds constexpr-elide summary TensorMaps. + The ``sage_*`` scale arguments are used only by a Sage attention config, + the summary K scales only by its proxy routes. """ if cutlass.const_expr(not cfg.use_block_sparse): raise ValueError("fmha_block_sparse_launch requires block-sparse config") @@ -2906,16 +3894,26 @@ def fmha_block_sparse_launch( b, h_q, h_k, s_k, d = problem_shape h_r = h_q // h_k q_seq = Int32(cfg.max_seq_len_q) - q_strides, kv_strides = _block_sparse_bshd_tma_strides( + q_strides, _ = _block_sparse_bshd_tma_strides( + q_seq=q_seq, + h_q=h_q, + h_k=h_k, + s_k=s_k, + d=d, + element_bytes=cfg.q_dtype_bytes, + ) + _, kv_strides = _block_sparse_bshd_tma_strides( q_seq=q_seq, h_q=h_q, h_k=h_k, s_k=s_k, d=d, + element_bytes=cfg.kv_dtype_bytes, ) - # FP16/BF16 H128 uses a 64-element (128-byte) inner box. The sparse - # profile validator rejects other element widths and head dimensions. + # H128 uses a 128-byte inner box: 64 two-byte or 128 one-byte elements. + # The sparse profile validator rejects other head dimensions. + tma_box0_q = min(128 // cfg.q_dtype_bytes, cfg.headdim) tma_box0 = min(128 // cfg.kv_dtype_bytes, cfg.headdim) tma_swizzle = cuda.TensorMapSwizzle.s128b # Public tensors are contiguous BSHD. Q is factored into (Hr, Hkv) so the @@ -2926,7 +3924,7 @@ def fmha_block_sparse_launch( global_dims=(d, h_r, h_k, q_seq, b), global_strides=q_strides, box_dims=( - tma_box0, + tma_box0_q, cfg.heads_q_per_kv, 1, cfg.q_tokens_per_cta, @@ -2997,7 +3995,7 @@ def fmha_block_sparse_launch( kv_dims = (d, s_k, h_k, b) k_desc_primary = create_tensor_map_tiled( global_address=k_iter.toint(), - dtype=cfg.kv_dtype, + dtype=cfg.k_dtype, global_dims=kv_dims, global_strides=kv_strides, box_dims=(tma_box0, primary_kv_box_size, 1, 1), @@ -3005,7 +4003,7 @@ def fmha_block_sparse_launch( ) v_desc_primary = create_tensor_map_tiled( global_address=v_iter.toint(), - dtype=cfg.kv_dtype, + dtype=cfg.v_dtype, global_dims=kv_dims, global_strides=kv_strides, box_dims=(tma_box0, primary_kv_box_size, 1, 1), @@ -3018,7 +4016,7 @@ def fmha_block_sparse_launch( # map only when a route may join unrelated BSR entries. k_desc_atom = create_tensor_map_tiled( global_address=k_iter.toint(), - dtype=cfg.kv_dtype, + dtype=cfg.k_dtype, global_dims=kv_dims, global_strides=kv_strides, box_dims=(tma_box0, kv_atom_size, 1, 1), @@ -3026,7 +4024,7 @@ def fmha_block_sparse_launch( ) v_desc_atom = create_tensor_map_tiled( global_address=v_iter.toint(), - dtype=cfg.kv_dtype, + dtype=cfg.v_dtype, global_dims=kv_dims, global_strides=kv_strides, box_dims=(tma_box0, kv_atom_size, 1, 1), @@ -3044,11 +4042,12 @@ def fmha_block_sparse_launch( h_k=h_k, s_k=num_kv_blocks, d=d, + element_bytes=cfg.kv_dtype_bytes, ) summary_dims = (d, num_kv_blocks, h_k, b) k_desc_summary_primary = create_tensor_map_tiled( global_address=k_summary_iter.toint(), - dtype=cfg.kv_dtype, + dtype=cfg.k_dtype, global_dims=summary_dims, global_strides=summary_kv_strides, box_dims=(tma_box0, primary_kv_box_size, 1, 1), @@ -3056,7 +4055,7 @@ def fmha_block_sparse_launch( ) v_desc_summary_primary = create_tensor_map_tiled( global_address=v_summary_iter.toint(), - dtype=cfg.kv_dtype, + dtype=cfg.v_dtype, global_dims=summary_dims, global_strides=summary_kv_strides, box_dims=(tma_box0, primary_kv_box_size, 1, 1), @@ -3067,7 +4066,7 @@ def fmha_block_sparse_launch( if cutlass.const_expr(uses_atom_desc): k_desc_summary_atom = create_tensor_map_tiled( global_address=k_summary_iter.toint(), - dtype=cfg.kv_dtype, + dtype=cfg.k_dtype, global_dims=summary_dims, global_strides=summary_kv_strides, box_dims=(tma_box0, kv_atom_size, 1, 1), @@ -3075,7 +4074,7 @@ def fmha_block_sparse_launch( ) v_desc_summary_atom = create_tensor_map_tiled( global_address=v_summary_iter.toint(), - dtype=cfg.kv_dtype, + dtype=cfg.v_dtype, global_dims=summary_dims, global_strides=summary_kv_strides, box_dims=(tma_box0, kv_atom_size, 1, 1), @@ -3112,6 +4111,8 @@ def fmha_block_sparse_launch( q_desc, k_desc_primary, v_desc_primary, + k_desc_primary, + v_desc_primary, k_desc_atom, v_desc_atom, o_iter, @@ -3122,6 +4123,12 @@ def fmha_block_sparse_launch( seqlens_kv_iter, null_i32_ptr, null_i32_ptr, + null_f32_ptr, + null_f32_ptr, + Int64(0), # g_page_table_stride + Int32(0), # g_page_table_capacity + null_i32_ptr, + Int32(0), # g_q_token_kv_block_sparse_page_membership_stride o_iter, null_f32_ptr, null_i32_ptr, @@ -3133,9 +4140,6 @@ def fmha_block_sparse_launch( use_variable_seqlens_kv, False, # use_native_paged_kv False, # use_static_native_seqlens_kv - null_i32_ptr, # g_block_tables - Int32(0), # g_block_table_capacity - Int64(0), # g_block_table_row_stride row_route_offsets_iter, row_route_counts_iter, route_metadata_iter, @@ -3144,6 +4148,14 @@ def fmha_block_sparse_launch( tma_desc_v_summary=v_desc_summary_primary, tma_desc_k_summary_atom=k_desc_summary_atom, tma_desc_v_summary_atom=v_desc_summary_atom, + g_sage_q_scale=sage_q_scale_iter, + g_sage_k_scale=sage_k_scale_iter, + g_sage_k_summary_scale=sage_k_summary_scale_iter, + g_sage_v_scale=sage_v_scale_iter, + g_sage_v_mean=sage_v_mean_iter, + g_sage_q_scale_head_stride=sage_q_scale_head_stride, + g_sage_k_scale_head_stride=sage_k_scale_head_stride, + g_sage_k_summary_scale_head_stride=sage_k_summary_scale_head_stride, ).launch( grid=grid, block=[cfg.threads_per_cta, 1, 1], @@ -3152,5 +4164,5 @@ def fmha_block_sparse_launch( # Reuse dense decode's entry-occupancy contract, including the long # KV256 profiles that execute dynamic register reallocation. min_blocks_per_mp=_decode_min_blocks_per_mp(cfg, seq_len_kv), - use_pdl=cfg.use_parallel_separate_reduction_pdl, + use_pdl=False, ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/__init__.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/__init__.py index 60360f90de17..664566125c43 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/__init__.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/__init__.py @@ -53,6 +53,11 @@ - SmemKvTileResource : Dedicated SMEM tile used by split-head-dimension profiles for one K or V producer instance. +- SmemTransformedKvResource : Mixed Q/KV shadow of SmemKvResource. Producer + (TransformKvTask): converts raw K/V into the + Q-side MMA input type. Consumer (MmaTask): uses the + same descriptor labels as SmemKvResource. + - SmemPResource : P operand for BMM2. The validated one-instance staged-D256 Keeps profile places P in two TMEM views aliased with the matching S stages; other profiles @@ -94,12 +99,17 @@ overwriting S; Correction returns it after loading the matching stats into registers. +- TmemTransformedKvResource : Mixed FP8-Q/NVFP4 K/V operand ring. Producer + (TransformKvTask) dequantizes unpacking-TMA SMEM + stages into TMEM; BMM1 and BMM2 consume K and + transposed V as operand A. + - TmemSoftmaxGlobalResource : FP8 sum-correction helper. Producer-only resource that, after P quantization, reapplies the running-max correction to the running denominator (using the TmemS local-sum array) and publishes the corrected sums back through TmemS. Inactive when - non-FP8 Q/K/V. + non-FP8 Q. - TmemCorrResource : Correction and output resource. LOOP stages rescale an in-flight O tile when the running max @@ -125,6 +135,7 @@ SmemBlockSparseKvMetadataResource, SmemBlockSparseSoftmaxMetadataResource, ) +from .smem_transformed_kv import SmemTransformedKvResource from .tmem_corr import TmemCorrResource from .tmem_o import TmemOResource from .smem_p import SmemPResource @@ -135,6 +146,7 @@ TmemSoftmaxLocalResource, TmemSoftmaxOrderResource, ) +from .tmem_transformed_kv import TmemTransformedKvResource __all__ = [ "DecodeGenResourceBase", @@ -145,6 +157,7 @@ "SmemQResource", "SmemBlockSparseKvMetadataResource", "SmemBlockSparseSoftmaxMetadataResource", + "SmemTransformedKvResource", "TmemCorrResource", "TmemOResource", "TmemSResource", @@ -152,4 +165,5 @@ "TmemSoftmaxGlobalResource", "TmemSoftmaxLocalResource", "TmemSoftmaxOrderResource", + "TmemTransformedKvResource", ] 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 77412cddca7b..23fc82d712d5 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 @@ -22,6 +22,7 @@ from functools import partial from typing import ClassVar + import cutlass import cutlass.cute as cute from cutlass import BFloat16, Float16, Float32, Int32, Int64, Uint32 @@ -55,7 +56,7 @@ Int32, Int32, ] -DescriptorValue = prims.Tcgen05SmemDesc | cutlass.Int64 +DescriptorValue = prims.Tcgen05SmemDesc | cutlass.Int64 | Int32 ResourceVarValue = ( Int32 | Float32 | Uint32 | cutlass.Int64 | cutlass.Array | DescriptorValue ) @@ -154,14 +155,44 @@ def _swaps_routed_coordinate( return atom_origin, atom_origin + Int32(token_offset) + lane_k_offset -def _mma_kind_for_qkv(cfg: FmhaDecodeConfig) -> prims.Tcgen05MMAKind: - """Select the tcgen05 MMA opcode family used for Q/K/V operands.""" - return prims.Tcgen05MMAKind.F8F6F4 if cfg.use_fp8_qkv else prims.Tcgen05MMAKind.F16 +def _mma_kind_for_qk(cfg: FmhaDecodeConfig) -> prims.Tcgen05MMAKind: + """Select the tcgen05 MMA kind of the QK GEMM from ``qk_dtype``. + + Int8 accumulates INT32 scores; E4M3 and the 16-bit types accumulate FP32. + The INT8 kind exists on SM100 alone, so planning admits Int8 Q/K only on + that architecture (``_validate_int8_qk_device``). + """ + if cfg.uses_int32_scores: + return prims.Tcgen05MMAKind.INT8 + if cfg.use_fp8_qk: + return prims.Tcgen05MMAKind.F8F6F4 + return prims.Tcgen05MMAKind.F16 + + +def _mma_k_step_qk(cfg: FmhaDecodeConfig) -> int: + """Return the K dimension advanced by one QK-GEMM MMA instruction.""" + return 32 if cfg.qk_dtype_bytes == 1 else 16 + +def _mma_kind_for_pv(cfg: FmhaDecodeConfig) -> prims.Tcgen05MMAKind: + """Select the tcgen05 MMA opcode family for the PV GEMM.""" + if cfg.use_fp8_pv: + return prims.Tcgen05MMAKind.F8F6F4 + return prims.Tcgen05MMAKind.F16 -def _mma_k_step(cfg: FmhaDecodeConfig) -> int: - """Return the K dimension advanced by one tcgen05 MMA instruction.""" - return 32 if cfg.use_fp8_qkv else 16 + +def _mma_k_step_pv(cfg: FmhaDecodeConfig) -> int: + """Return the K dimension advanced by one PV-GEMM MMA instruction.""" + return cfg.pv_mma_k_step + + +def _qk_accumulator_dtype(cfg: FmhaDecodeConfig) -> type: + """Return the type the QK GEMM accumulates one score as: Int32 for Int8 Q/K. + + INT32 scores are accumulated on top of ``INT32_SCORE_BIAS``, so the + softmax reads every score tile as FP32. + """ + return Int32 if cfg.uses_int32_scores else Float32 @cute.jit @@ -241,6 +272,32 @@ def _keeps_q64_col_base(lane_idx: Int32, half_cols: int) -> Int32: return (lane_idx >> Int32(4)) * Int32(half_cols) +@cute.jit +def _keeps_spatial_half( + cfg: Constexpr[FmhaDecodeConfig], warp_grp_thread_idx: Int32 +) -> Int32: + """Return the spatial half a Keeps warp-group thread belongs to. + + KV256 threads ``[0, 64)`` and ``[64, 128)`` own the two spatial KV128 + partials of one logical Q row (``FmhaDecodeConfig.keeps_spatial_halves``); + every other Keeps profile has a single half. + """ + if cutlass.const_expr(cfg.keeps_spatial_halves == 2): + return warp_grp_thread_idx >> Int32(6) + return Int32(0) + + +@cute.jit +def _keeps_route_atom(cfg: Constexpr[FmhaDecodeConfig], half: Int32, position) -> Int32: + """Return the route atom a Keeps spatial half owns at ``position``. + + The halves interleave over the route's atoms: half ``h`` owns atoms ``h``, + ``h + halves``, ... in order, so its ``position``-th atom is ``position * + halves + h`` (``FmhaDecodeConfig.keeps_route_atom_owner`` is the inverse). + """ + return position * Int32(cfg.keeps_spatial_halves) + half + + @cute.jit def _keeps_row_idx(cfg: Constexpr[FmhaDecodeConfig], warp_grp_thread_idx: Int32): """Map a correction lane to the logical output row it owns.""" @@ -274,16 +331,18 @@ def _keeps_score_col( ) -> Int32: """Return the semantic KV column represented by one Keeps score register.""" if cutlass.const_expr(cfg.tile_size_kv == 256): - # A physical thread owns four K32 fragments. The spatial half selects - # alternating KV64 blocks; the temporal fragment selects the low/high - # K32 sub-block inside that semantic KV64 block. - spatial = warp_grp_thread_idx >> Int32(6) - fragment = reg_idx // 32 - semantic_block = Int32(2 * (fragment // 2)) + spatial + # A physical thread owns four K32 fragments: ``keeps_fragments_per_atom`` + # per layout atom of its spatial half (``_keeps_route_atom``), each + # fragment one K32 sub-block of that atom. + fragment_regs = cfg.softmax_score_fragment_regs + fragments_per_atom = cfg.keeps_fragments_per_atom + spatial = _keeps_spatial_half(cfg, warp_grp_thread_idx) + fragment = reg_idx // fragment_regs + semantic_block = _keeps_route_atom(cfg, spatial, fragment // fragments_per_atom) return ( - semantic_block * Int32(64) - + Int32((fragment % 2) * 32) - + Int32(reg_idx % 32) + semantic_block * Int32(cfg.keeps_atom_tokens) + + Int32((fragment % fragments_per_atom) * fragment_regs) + + Int32(reg_idx % fragment_regs) ) return col_base + Int32(reg_idx) @@ -352,9 +411,30 @@ def _pack_float2_to_bf16(v0: Float32, v1: Float32) -> Int32: ) -def _qkv_smem_swizzle(cfg: FmhaDecodeConfig) -> prims.Tcgen05SmemSwizzle: - """Select the tcgen05 SMEM swizzle for staged Q/K/V tiles.""" - if cfg.use_fp8_qkv and cfg.headdim == 64: +def _q_dtype_elements(cfg: FmhaDecodeConfig) -> int: + """Return element count in one staged Q SMEM allocation.""" + return cfg.smem_q_tile_bytes // cfg.q_dtype_bytes + + +def _kv_dtype_elements(cfg: FmhaDecodeConfig) -> int: + """Return element count in one staged K/V SMEM allocation.""" + return cfg.smem_kv_tile_elements + + +def _qkv_smem_swizzle( + cfg: FmhaDecodeConfig, +) -> prims.Tcgen05SmemSwizzle: + """Select the tcgen05 SMEM swizzle for staged Q tiles.""" + if cfg.q_dtype_bytes == 1 and cfg.headdim == 64: + return prims.Tcgen05SmemSwizzle.SWIZZLE_64B + return prims.Tcgen05SmemSwizzle.SWIZZLE_128B + + +def _kv_smem_swizzle( + cfg: FmhaDecodeConfig, +) -> prims.Tcgen05SmemSwizzle: + """Select the tcgen05 SMEM swizzle for raw staged K/V tiles.""" + if cfg.k_dtype.width == 8 and cfg.v_dtype.width == 8 and cfg.headdim == 64: return prims.Tcgen05SmemSwizzle.SWIZZLE_64B return prims.Tcgen05SmemSwizzle.SWIZZLE_128B @@ -376,17 +456,38 @@ def _major_k_stride_bytes(dtype_bytes: int, headdim: int) -> int: return 128 * rows_per_swizzle_blk -@cute.jit -def _fp8_log2_quant_scale() -> Float32: - """Return log2 scaling used by FP8 probability quantization.""" - return Float32(8.8073549) - - def _neg_max_f32() -> Float32: """Return the negative sentinel used for running softmax maxima.""" return Float32(NEG_FLT_MAX) +@cute.jit +def _masked_weighted_sum(terms: tuple) -> Float32: + """Return the sum of ``value * weight`` over four ``(uses, value, weight)`` terms. + + A term whose ``uses`` is false contributes zero, whatever its value. The + fixed term order makes lanes that share a row agree bitwise. + """ + assert len(terms) == 4 + (uses0, value0, weight0), (uses1, value1, weight1) = terms[0], terms[1] + (uses2, value2, weight2), (uses3, value3, weight3) = terms[2], terms[3] + masked0 = Float32(0.0) + masked1 = Float32(0.0) + masked2 = Float32(0.0) + masked3 = Float32(0.0) + if uses0: + masked0 = value0 + if uses1: + masked1 = value1 + if uses2: + masked2 = value2 + if uses3: + masked3 = value3 + lanes = ffma2((masked0, masked2), (weight0, weight2), (Float32(0.0), Float32(0.0))) + lanes = ffma2((masked1, masked3), (weight1, weight3), lanes) + return lanes[0] + lanes[1] + + def _softmax_tile_idx( cfg: FmhaDecodeConfig, stage_info: StageInfo, inst_id: int ) -> Int32: @@ -781,3 +882,21 @@ def __post_init__(self) -> None: def _init_placeholder_state(self) -> None: """Hook for subclasses to install placeholder arrays before tracing.""" return + + +# Keeps softmax staging reserves bit 5 for the prepared route kind; the low +# four structural validity bits and bit 4 keep their meaning. +_SOFTMAX_ROUTE_IS_PROXY_FLAG = 1 << 5 + + +@cute.jit +def _route_is_proxy(route_flags: Int32) -> cutlass.Boolean: + """Return whether staged Softmax route flags mark a proxy route. + + Keeps stages the flags as an ``Int32`` word. SWAP forwards them as a + bit-preserving ``Uint32`` dataflow token, which the consumer bitcasts + before calling. + """ + return cutlass.Boolean( + (route_flags & Int32(_SOFTMAX_ROUTE_IS_PROXY_FLAG)) != Int32(0) + ) 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 538a360badbb..3a795a73a4a2 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 @@ -28,9 +28,18 @@ from ...mask import kv_tile_is_fully_visible from ..fmha_decode_config import CAUSAL, FmhaDecodeConfig from .helpers_common import ( + Constexpr, + _TASK_CACHE_KV_RAW_TILE_BASE, + _TASK_CACHE_KV_VALID_TILE_END, + _TASK_CACHE_KV_WINDOW_START, _TASK_CACHE_SEQ_LEN_KV, + _clamp_valid_tile_idx, _decode_gen_task_cache, + _is_last_loop_iteration, _logical_head_batch, + _logical_q_group_idx, + _q_group_token_base, + _softmax_tile_idx, ) @@ -337,3 +346,117 @@ def _runtime_split_kv_global_tile_idx( * _runtime_local_kv_tiles(cfg, seq_len_kv, seq_len_q, q_token_base) + local_tile_idx ) + + +@cute.jit +def resolve_keeps_tile_context( + cfg: Constexpr[FmhaDecodeConfig], + stage_info: StageInfo, + *, + inst_id: Constexpr[int], + seqlens_kv: cute.Pointer | None, + max_seq_len_kv: Int32, + seq_len_q: Int32, + q_group_idx: Int32 | None, +): + """Resolve one score tile's logical position and boundary-mask state. + + Every consumer of a softmax instance's tiles (the S resource and the + Sage K scale resource) derives the tile's token range from the same + sequence-length sources, so the resolution lives outside the S resource. + """ + task_cache = _decode_gen_task_cache(stage_info) + if cutlass.const_expr(seqlens_kv is None): + seq_len_kv = Int32(max_seq_len_kv) + else: + seq_len_kv = _load_runtime_seq_len_kv( + seqlens_kv, + max_seq_len_kv, + stage_info, + Int32(0), + Int32(0), + ) + + logical_q_group_idx = _logical_q_group_idx(cfg, stage_info, q_group_idx) + q_token_base = _q_group_token_base(cfg, logical_q_group_idx) + element_mask_end_idx = seq_len_kv + 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 + local_tile_idx = _softmax_tile_idx(cfg, stage_info, inst_id) + if cutlass.const_expr(not use_runtime_kv_domain): + effective_tile_idx = _static_split_kv_global_tile_idx( + cfg, stage_info, local_tile_idx + ) + effective_total_kv_tiles = Int32(cfg.total_kv_tiles) + tile_idx = _clamp_valid_tile_idx(cfg, effective_tile_idx) + tile_idx = tile_idx + Int32(cfg.static_num_skipped_kv_tiles) + window_start_idx = Int32(cfg.static_window_start_idx) + elif cutlass.const_expr(cfg.use_paged_kv and not cfg.use_split_kv): + effective_tile_idx = ( + Int32(task_cache[_TASK_CACHE_KV_RAW_TILE_BASE]) + local_tile_idx + ) + effective_total_kv_tiles = Int32(task_cache[_TASK_CACHE_KV_VALID_TILE_END]) + tile_idx = effective_tile_idx + window_start_idx = Int32(task_cache[_TASK_CACHE_KV_WINDOW_START]) + else: + effective_tile_idx = _runtime_split_kv_global_tile_idx( + cfg, + stage_info, + local_tile_idx, + seq_len_kv, + seq_len_q, + q_token_base, + ) + effective_total_kv_tiles = _runtime_total_kv_tiles( + cfg, seq_len_kv, seq_len_q, q_token_base + ) + tile_idx = _runtime_clamp_valid_tile_idx( + cfg, + effective_tile_idx, + seq_len_kv, + seq_len_q, + q_token_base, + ) + tile_idx = tile_idx + _num_skipped_kv_tiles( + cfg, seq_len_kv, seq_len_q, q_token_base + ) + window_start_idx = _sliding_window_start_idx( + cfg, seq_len_kv, seq_len_q, q_token_base + ) + + tile_offset_k = tile_idx * Int32(cfg.tile_size_kv) + is_valid_effective_tile = effective_tile_idx < effective_total_kv_tiles + is_masked_final_wave = False + if cutlass.const_expr(not use_runtime_kv_domain and not cfg.use_split_kv): + if cutlass.const_expr(cfg.has_odd_kv_tail and inst_id == 1): + is_masked_final_wave = _is_last_loop_iteration(stage_info) + tile_has_valid_scores = ( + is_valid_effective_tile + and (tile_offset_k < seq_len_kv) + and not is_masked_final_wave + ) + tile_is_unmasked = _kv_tile_is_fully_unmasked_for_q_group( + cfg, + tile_offset_k, + seq_len_kv, + seq_len_q, + q_token_base, + tile_has_valid_scores, + ) + # Flattened SQ=1 QToken-KvBlock-Sparse-Attention routes use their compacted causal length as + # seq_len_kv. At or above the 2048-token budget this marks all first + # sixteen KV128 tiles unmasked; only the optional 0--3-token tail in a + # seventeenth tile reaches the boundary-mask specialization. + return ( + seq_len_kv, + logical_q_group_idx, + element_mask_end_idx, + tile_offset_k, + window_start_idx, + is_valid_effective_tile, + is_masked_final_wave, + tile_is_unmasked, + tile_has_valid_scores, + ) 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 db6bd6b8de8c..24a7a4c8b29e 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 @@ -37,7 +37,7 @@ def _keeps_p_smem_block_offset_bytes( cfg: Constexpr[FmhaDecodeConfig], row_idx: Int32, col_idx: Int32 ) -> Int32: """JIT form of the Keeps P SWIZZLE_128B vector-block address.""" - dtype_bytes = cfg.q_dtype_bytes + dtype_bytes = cfg.v_dtype_bytes chunk_cols = 128 // dtype_bytes chunk_idx = col_idx // Int32(chunk_cols) col_in_chunk = col_idx - chunk_idx * Int32(chunk_cols) 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 2dbe3c7e7ff8..ca7b8a662cec 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_softmax.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_softmax.py @@ -26,12 +26,11 @@ from cutlass.experimental import primitives as prims from ..fmha_decode_config import FmhaDecodeConfig +from ..fmha_decode_constants import FP8_P_QUANT_SCALE from .helpers_common import ( Constexpr, fadd2, ffma2, - fmul2, - _fp8_log2_quant_scale, _neg_max_f32, _q_row_token_and_local_head, ) @@ -138,8 +137,8 @@ def _ex2_emulation_packed_f32x2(x: Float32, y: Float32) -> tuple[Float32, Float3 @cute.jit def _compute_fp8_p_regs_and_local_sums( scale_softmax_log2: Float32, - new_max_0: Float32, - new_max_1: Float32, + addend_0: Float32, + addend_1: Float32, s0: Float32, s1: Float32, s2: Float32, @@ -149,40 +148,29 @@ def _compute_fp8_p_regs_and_local_sums( s6: Float32, s7: Float32, ) -> tuple[Int32, Int32, Float32, Float32]: - """Compute masked FP8 P registers and local softmax sums.""" - # Safe path: masked tiles can produce NEG_FLT_MAX as new_max. Treat that - # as zero for the exponent offset so invalid rows generate zero P instead - # of NaNs while local sums remain zero. - safe_new_max_0 = new_max_0 - safe_new_max_1 = new_max_1 - if safe_new_max_0 == _neg_max_f32(): - safe_new_max_0 = Float32(0.0) - if safe_new_max_1 == _neg_max_f32(): - safe_new_max_1 = Float32(0.0) - neg_scaled_max_pair = ffma2( - (safe_new_max_0, safe_new_max_1), - (-scale_softmax_log2, -scale_softmax_log2), - (_fp8_log2_quant_scale(), _fp8_log2_quant_scale()), - ) + """Exponentiate two scale groups of four scores into packed FP8 P and sums. - # Scale S by log2(e) * softmax_scale and include the FP8 quantization - # shift. The packed f32x2 operations keep paired scale groups aligned. + Every probability is ``exp2(c * s + addend)`` with the group's exponent + addend (``SmemPResource._exponent_addend``), so no mask logic is needed: + a masked score at ``-FLT_MAX`` exponentiates to zero. + """ + addend_pair = (addend_0, addend_1) scaled_pair_01 = ffma2( (s0, s1), (scale_softmax_log2, scale_softmax_log2), - neg_scaled_max_pair, + addend_pair, ) scaled_pair_23 = ffma2( (s2, s3), (scale_softmax_log2, scale_softmax_log2), - neg_scaled_max_pair, + addend_pair, ) prob0 = cute.math.exp2(scaled_pair_01[0], fastmath=True) prob1 = cute.math.exp2(scaled_pair_01[1], fastmath=True) scaled_pair_45 = ffma2( (s4, s5), (scale_softmax_log2, scale_softmax_log2), - neg_scaled_max_pair, + addend_pair, ) prob2 = cute.math.exp2(scaled_pair_23[0], fastmath=True) prob3 = cute.math.exp2(scaled_pair_23[1], fastmath=True) @@ -192,7 +180,7 @@ def _compute_fp8_p_regs_and_local_sums( scaled_pair_67 = ffma2( (s6, s7), (scale_softmax_log2, scale_softmax_log2), - neg_scaled_max_pair, + addend_pair, ) prob4 = cute.math.exp2(scaled_pair_45[0], fastmath=True) prob5 = cute.math.exp2(scaled_pair_45[1], fastmath=True) @@ -206,69 +194,10 @@ def _compute_fp8_p_regs_and_local_sums( @cute.jit -def _compute_fp8_p_regs_and_local_sums_dense( +def _compute_p_values_and_local_sums( scale_softmax_log2: Float32, - new_max_0: Float32, - new_max_1: Float32, - s0: Float32, - s1: Float32, - s2: Float32, - s3: Float32, - s4: Float32, - s5: Float32, - s6: Float32, - s7: Float32, -) -> tuple[Int32, Int32, Float32, Float32]: - """Compute dense FP8 P registers and local softmax sums.""" - # Dense path: all S entries are valid, so no NEG_FLT_MAX guard is needed. - neg_scaled_max_pair = ffma2( - (new_max_0, new_max_1), - (-scale_softmax_log2, -scale_softmax_log2), - (_fp8_log2_quant_scale(), _fp8_log2_quant_scale()), - ) - - scaled_pair_01 = ffma2( - (s0, s1), - (scale_softmax_log2, scale_softmax_log2), - neg_scaled_max_pair, - ) - scaled_pair_23 = ffma2( - (s2, s3), - (scale_softmax_log2, scale_softmax_log2), - neg_scaled_max_pair, - ) - prob0 = cute.math.exp2(scaled_pair_01[0], fastmath=True) - prob1 = cute.math.exp2(scaled_pair_01[1], fastmath=True) - scaled_pair_45 = ffma2( - (s4, s5), - (scale_softmax_log2, scale_softmax_log2), - neg_scaled_max_pair, - ) - prob2 = cute.math.exp2(scaled_pair_23[0], fastmath=True) - prob3 = cute.math.exp2(scaled_pair_23[1], fastmath=True) - local_sum_pair_01 = fadd2((prob0, prob1), (prob2, prob3)) - packed_p_0 = _pack_float4_to_fp8_e4m3(prob0, prob1, prob2, prob3) - scaled_pair_67 = ffma2( - (s6, s7), - (scale_softmax_log2, scale_softmax_log2), - neg_scaled_max_pair, - ) - prob4 = cute.math.exp2(scaled_pair_45[0], fastmath=True) - prob5 = cute.math.exp2(scaled_pair_45[1], fastmath=True) - prob6 = cute.math.exp2(scaled_pair_67[0], fastmath=True) - prob7 = cute.math.exp2(scaled_pair_67[1], fastmath=True) - local_sum_pair_45 = fadd2((prob4, prob5), (prob6, prob7)) - packed_p_1 = _pack_float4_to_fp8_e4m3(prob4, prob5, prob6, prob7) - local_sum_pair = fadd2(local_sum_pair_01, local_sum_pair_45) - - return packed_p_0, packed_p_1, local_sum_pair[0], local_sum_pair[1] - - -@cute.jit -def _compute_p_values_and_local_sums_dense( - scale_softmax_log2: Float32, - new_max_0: Float32, - new_max_1: Float32, + addend_0: Float32, + addend_1: Float32, s0: Float32, s1: Float32, s2: Float32, @@ -289,36 +218,36 @@ def _compute_p_values_and_local_sums_dense( Float32, Float32, ]: - """Compute dense 16-bit P values and paired local softmax sums.""" - # Dense 16-bit path: compute eight P values and the two local sums without - # per-row validity checks. - neg_scaled_max_pair = fmul2( - (new_max_0, new_max_1), - (-scale_softmax_log2, -scale_softmax_log2), - ) + """Exponentiate two scale groups of four scores into FP32 P and paired sums. + + Straight-line twin of ``_compute_fp8_p_regs_and_local_sums`` for 16-bit + P: the same ``exp2(c * s + addend)`` per score, returned unpacked so the + caller can choose the P storage format. + """ + addend_pair = (addend_0, addend_1) scaled_pair_01 = ffma2( (s0, s1), (scale_softmax_log2, scale_softmax_log2), - neg_scaled_max_pair, + addend_pair, ) scaled_pair_23 = ffma2( (s2, s3), (scale_softmax_log2, scale_softmax_log2), - neg_scaled_max_pair, + addend_pair, ) prob0 = cute.math.exp2(scaled_pair_01[0], fastmath=True) prob1 = cute.math.exp2(scaled_pair_01[1], fastmath=True) scaled_pair_45 = ffma2( (s4, s5), (scale_softmax_log2, scale_softmax_log2), - neg_scaled_max_pair, + addend_pair, ) prob2 = cute.math.exp2(scaled_pair_23[0], fastmath=True) prob3 = cute.math.exp2(scaled_pair_23[1], fastmath=True) scaled_pair_67 = ffma2( (s6, s7), (scale_softmax_log2, scale_softmax_log2), - neg_scaled_max_pair, + addend_pair, ) prob4 = cute.math.exp2(scaled_pair_45[0], fastmath=True) prob5 = cute.math.exp2(scaled_pair_45[1], fastmath=True) @@ -482,8 +411,8 @@ def _attention_sink_for_local_head( sink_val * Float32(1.4426950408889634) - max_val * scale_softmax_log2, fastmath=True, ) - if cutlass.const_expr(cfg.use_fp8_qkv): - sink_exp = sink_exp * Float32(448.0) + if cutlass.const_expr(cfg.use_fp8_pv): + sink_exp = sink_exp * Float32(FP8_P_QUANT_SCALE) return sink_exp diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/sage_scales.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/sage_scales.py new file mode 100644 index 000000000000..4d862f9fa1f1 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/sage_scales.py @@ -0,0 +1,825 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Sage attention scale addressing for the decode softmax and epilogue. + +Every scale load of the kernel goes through this module. ``sfQ`` and ``sfK`` +use the trtllm-gen flat layout (``flat_scale_slot`` of +:mod:`flashinfer.attention.prims_ts.sage`). Tokens past the sequence end and +Q rows past the valid row count clamp to the last valid slot, so masked +scores keep a finite scale and exponentiate to zero. + +:class:`SageKScalesResource` holds one softmax instance's ``sfK`` words for +one KV tile. The instance produces and consumes it: it fills the words ahead +of the tile's score wait, and both softmax passes read them one K32 fragment +at a time. The scale groups per fragment select the storage form +(``FmhaDecodeConfig.sage_k_scales_in_smem_for``): + +* Register form (K blocks of 16 tokens or more, at most two groups per + fragment): each lane gathers its own words into a rotating register array; + there is no pipeline. +* SMEM form (K blocks of 4 and 1 token, eight or 32 groups per fragment): + a two-slot SMEM buffer in the ``sage_k_scale_words`` layout, pipelined + between the instance's own warps. The passes read a fragment's groups with + 16-byte broadcast loads; the routed register array is a placeholder. + +A dense tile gathers ``k_scale`` over its token range, a block-sparse route +takes the words the load warp staged with its metadata, and a proxy route of +a mixed-geometry plan (``sage_mixed_k_geometry``) gathers ``k_summary_scale`` +from the route's atom origins. A mixed plan holds one resource per geometry; +the CTA-uniform route kind selects the one that serves the tile. + +:class:`SageVScalesResource` stages the work tile's per-channel V scales +(and means) in SMEM once per tile, while no output is pending, so the +epilogue scales column ``c`` by ``sfV[c]`` and adds ``v_mean[c]`` without +waiting on global loads. +""" + +from collections.abc import Callable +from dataclasses import dataclass +from typing import TYPE_CHECKING, ClassVar + +import cutlass +import cutlass.cute as cute +from cutlass import Float32, Int32, Int64 +from cutlass.experimental.task_scheduling.enums import WorkAttr +from cutlass.experimental.task_scheduling.memory import ( + ResourceContext, + SmemAllocation, + TmemAllocation, +) +from cutlass.experimental.task_scheduling.resources import ( + StageInfo, + TaskLocalVariable, + consumer_work, + producer_work, +) + +from ....sage import flat_scale_slot, log2_block_size +from ...placeholder_helpers import _placeholder_local_array +from ..fmha_decode_config import FmhaDecodeConfig +from .helpers_common import ( + DecodeGenResourceBase, + _TASK_CACHE_WARP_GRP_THREAD_IDX, + _decode_gen_task_cache, + _keeps_route_atom, + _keeps_spatial_half, + _logical_head_batch, + _route_is_proxy, + fmul2, +) +from .helpers_kv_tile_idx import resolve_keeps_tile_context + +if TYPE_CHECKING: + from .smem_block_sparse_metadata import SmemBlockSparseSoftmaxMetadataResource + +Constexpr = cutlass.Constexpr + + +@dataclass(eq=False) +class SageScaleTensors: + """The plan's Sage scale tensors, shared by every resource that reads one. + + ``sfQ``, ``sfK`` and ``k_summary_scale`` are ``[heads, slots]`` FP32 + arrays in the flat layout; the V scales and means are ``[Hkv, D]`` FP32. + ``k_summary_scale`` is ``None`` without proxy routes and ``v_mean`` + without a channel mean. The dataclass is not frozen because the DSL + replaces frozen dataclasses with proxies inside traced dynamic branches, + which changes the traced structure of the holding resource. + """ + + q_scale_ptr: cute.Pointer + q_scale_head_stride: Int32 + k_scale_ptr: cute.Pointer + k_scale_head_stride: Int32 + k_summary_scale_ptr: cute.Pointer | None + k_summary_scale_head_stride: Int32 | None + v_scale_ptr: cute.Pointer + v_mean_ptr: cute.Pointer | None + + +@cute.jit +def _load_scale(scale_addr: Int64, index: Int32) -> Float32: + """Return one FP32 scale loaded from a global memory base address.""" + 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_q_scale( + cfg: Constexpr[FmhaDecodeConfig], + q_scale_addr: Int64, + q_scale_head_stride: Int32, + *, + kv_head_idx: Int32, + local_head_idx: Int32, + batch_idx: Int32, + q_token_idx: Int32, +) -> Float32: + """Return ``sfQ`` for one Q row (token and Q head of one KV head group).""" + heads_q_per_kv = Int32(cfg.heads_q_per_kv) + q_head_idx = kv_head_idx * heads_q_per_kv + cute.math.min( + local_head_idx, heads_q_per_kv - Int32(1) + ) + q_token_idx = cute.math.min(q_token_idx, Int32(cfg.max_seq_len_q - 1)) + slot = flat_scale_slot( + batch_idx, + q_token_idx, + Int32(cfg.max_seq_len_q), + log2_block_size(cfg.sage_q_block_size), + ) + return _load_scale(q_scale_addr, q_head_idx * q_scale_head_stride + slot) + + +@cute.jit +def load_k_scale( + cfg: Constexpr[FmhaDecodeConfig], + k_scale_addr: Int64, + k_scale_head_stride: Int32, + *, + kv_head_idx: Int32, + batch_idx: Int32, + seq_len_kv: Int32, + kv_token_idx: Int32, + log2_k_block_size: Int32 | None = None, +) -> Float32: + """Return ``sfK`` for one token of one KV head from a flat scale array. + + ``k_scale_addr`` is the base address of ``k_scale`` or ``k_summary_scale``, + ``seq_len_kv`` the length of the sequence it covers and + ``log2_k_block_size`` the log2 of its K block size, the recipe's token + block by default. + """ + log2_block = log2_k_block_size + if cutlass.const_expr(log2_block is None): + log2_block = Int32(log2_block_size(cfg.sage_k_block_size)) + kv_token_idx = cute.math.min(kv_token_idx, seq_len_kv - Int32(1)) + slot = flat_scale_slot(batch_idx, kv_token_idx, seq_len_kv, log2_block) + return _load_scale(k_scale_addr, kv_head_idx * k_scale_head_stride + slot) + + +def sage_scale_arr_size(cfg: FmhaDecodeConfig, groups: int) -> int: + """Return the number of ``sfK`` words one softmax lane holds per tile. + + ``groups`` is the scale groups per fragment of the tile's route kind + (``FmhaDecodeConfig.sage_k_groups_per_fragment_for``). + """ + return cfg.num_softmax_score_fragments * groups + + +def sage_k_scale_words(cfg: FmhaDecodeConfig, groups: int) -> int: + """Return the number of ``sfK`` words that cover one KV tile for every spatial half. + + Zero without Sage attention. Word ``half * arr_size + f * groups + g`` + holds group ``g`` of fragment ``f`` of that half's lanes. Route staging + and the SMEM form of :class:`SageKScalesResource` share this layout, so a + softmax thread reads its half with contiguous vector loads. + """ + if not cfg.use_sage_attention: + return 0 + return cfg.keeps_spatial_halves * sage_scale_arr_size(cfg, groups) + + +def sage_staged_k_scale_words(cfg: FmhaDecodeConfig) -> int: + """Return the ``sfK`` words a block-sparse route stages. + + Exact routes stage their words; proxy routes of a mixed-geometry plan + gather theirs in the softmax warps, so the staged area covers the exact + geometry alone. + """ + return sage_k_scale_words(cfg, cfg.sage_k_groups_per_fragment) + + +@cute.jit +def sage_word_position( + cfg: Constexpr[FmhaDecodeConfig], + half: Int32, + lane_entry, + groups: Constexpr[int], +) -> tuple[Int32, Int32]: + """Return ``(atom, token offset in the atom)`` of one ``sfK`` word. + + ``lane_entry = f * groups + g`` names group ``g`` of fragment ``f`` of + spatial half ``half``. Fragment ``f`` reads the half's Keeps layout atom + ``f // fragments_per_atom`` (``_keeps_route_atom``), which block-sparse + routes share, from token ``(f % fragments_per_atom) * fragment_regs``; + group ``g`` starts ``g * group_tokens`` later. + """ + fragments_per_atom = cfg.keeps_fragments_per_atom + fragment_regs = cfg.softmax_score_fragment_regs + group_tokens = fragment_regs // groups + fragment_idx = lane_entry // groups + group_idx = lane_entry % groups + atom_idx = _keeps_route_atom(cfg, half, fragment_idx // fragments_per_atom) + token_offset = ( + fragment_idx % fragments_per_atom + ) * fragment_regs + group_idx * group_tokens + return atom_idx, Int32(token_offset) + + +@cute.jit +def dense_k_scale_token( + cfg: Constexpr[FmhaDecodeConfig], half: Int32, lane_entry, tile_offset_k: Int32 +) -> Int32: + """Return the first KV token covered by one ``sfK`` word of a dense tile. + + A dense Keeps tile lays out its atoms in order (``_keeps_score_col``). + """ + atom_idx, token_offset = sage_word_position( + cfg, half, lane_entry, cfg.sage_k_groups_per_fragment + ) + return tile_offset_k + atom_idx * Int32(cfg.keeps_atom_tokens) + token_offset + + +@cute.jit +def scale_pairs_in_place( + values: cutlass.Array, factor: Float32, count: Constexpr[int] +) -> None: + """Multiply ``count`` values in place by ``factor``, two at a time.""" + assert count % 2 == 0 + for pair_base in cutlass.range_constexpr(0, count, 2): + values[pair_base], values[pair_base + 1] = fmul2( + (factor, factor), + (Float32(values[pair_base]), Float32(values[pair_base + 1])), + ) + + +@cute.jit +def block_sparse_k_scale_source( + cfg: Constexpr[FmhaDecodeConfig], + scales: SageScaleTensors, + *, + seq_len_kv: Int32, + route_is_proxy: cutlass.Boolean, +) -> tuple[Int64, Int32, Int32, Int32]: + """Return ``(address, head stride, sequence length, log2 block size)``. + + Exact routes read ``k_scale`` over the KV tokens; proxy routes read + ``k_summary_scale`` over the summary sequence, one summary per KV block. + The result feeds :func:`load_k_scale`. + """ + scale_addr = scales.k_scale_ptr.toint() + head_stride = scales.k_scale_head_stride + seq_len = seq_len_kv + log2_block = Int32(log2_block_size(cfg.sage_k_block_size)) + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): + if route_is_proxy: + scale_addr = scales.k_summary_scale_ptr.toint() + head_stride = scales.k_summary_scale_head_stride + seq_len = Int32(cfg.num_proxy_summaries) + log2_block = Int32(log2_block_size(cfg.sage_k_summary_block_size)) + return scale_addr, head_stride, seq_len, log2_block + + +def staged_v_channel_scale_entries(cfg: FmhaDecodeConfig) -> int: + """Return the number of SMEM floats holding one KV head's V scales and means. + + The scales occupy ``[0, headdim)``; with ``sage_v_mean`` the means follow + at ``[headdim, 2 * headdim)``. + """ + return cfg.headdim * (2 if cfg.sage_v_mean else 1) + + +@cute.jit +def stage_v_channel_scales( + cfg: Constexpr[FmhaDecodeConfig], + scales: SageScaleTensors, + staged: cutlass.Array, + *, + kv_head_idx: Int32, + thread_idx: Int32, + num_threads: Constexpr[int], +) -> None: + """Copy one KV head's per-channel V scales and means into SMEM. + + Each of the ``num_threads`` callers moves ``headdim / num_threads`` + channels. Callers order these writes before the epilogue reads, and the + next tile's writes after the last read. + """ + assert cfg.headdim % num_threads == 0 + v_scale_addr = scales.v_scale_ptr.toint() + head_base = kv_head_idx * Int32(cfg.headdim) + for base in cutlass.range_constexpr(0, cfg.headdim, num_threads): + channel = Int32(base) + thread_idx + staged[channel] = _load_scale(v_scale_addr, head_base + channel) + if cutlass.const_expr(cfg.sage_v_mean): + staged[channel + Int32(cfg.headdim)] = _load_scale( + scales.v_mean_ptr.toint(), head_base + channel + ) + + +@cute.jit +def load_staged_v_channel_scales( + cfg: Constexpr[FmhaDecodeConfig], + staged: cutlass.Array, + *, + first_col: Int32, + count: Constexpr[int], +) -> tuple[cutlass.Array, cutlass.Array]: + """Return ``count`` staged V scales and means from ``first_col``. + + Without ``sage_v_mean`` the returned mean array is zero so callers can + apply one fused multiply-add. + """ + assert count % 4 == 0 + scales = cutlass.Array(Float32, count, space=cutlass.AddressSpace.rmem) + means = cutlass.Array(Float32, count, space=cutlass.AddressSpace.rmem) + for chunk in cutlass.range_constexpr(0, count, 4): + chunk_col = first_col + Int32(chunk) + scale_vec = (staged.data_ptr() + chunk_col).load(count=4, alignment=16) + for elem in cutlass.range_constexpr(4): + scales[chunk + elem] = Float32(scale_vec[elem]) + if cutlass.const_expr(cfg.sage_v_mean): + mean_vec = (staged.data_ptr() + chunk_col + Int32(cfg.headdim)).load( + count=4, alignment=16 + ) + for elem in cutlass.range_constexpr(4): + means[chunk + elem] = Float32(mean_vec[elem]) + else: + for elem in cutlass.range_constexpr(4): + means[chunk + elem] = Float32(0.0) + return scales, means + + +WordSource = Callable[[Int32, Int32, object], Float32] + +SAGE_K_SCALES_RING_STAGES = 2 + + +@dataclass(kw_only=True) +class SageKScalesResource(DecodeGenResourceBase): + """One softmax instance's ``sfK`` words for one scale-group geometry. + + ``summary`` marks the resource of a mixed-geometry plan that serves proxy + routes from ``k_summary_scale``. The SMEM form takes the instance's + two-stage pipeline as ``pipeline_config``; the register form takes none. + ``route_metadata`` is the instance's block-sparse softmax metadata + resource, whose held stage carries the route's staged words and atom + origins; a dense plan resolves the tile's tokens from the sequence-length + fields instead. + """ + + _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ( + "sage_scale_arr", + cutlass.Array, + None, + "Raw sfK words of the lane's fragments for the current tile.", + ), + ) + cfg: Constexpr[FmhaDecodeConfig] = None + inst_id: Constexpr[int] = 0 + summary: Constexpr[bool] = False + route_metadata: "SmemBlockSparseSoftmaxMetadataResource | None" = None + seqlens_kv: cute.Pointer | None = None + max_seq_len_kv: Int32 = None + seq_len_q: Int32 = None + q_group_idx: Int32 | None = None + h_k_idx: Int32 | None = None + b_idx: Int32 | None = None + scale_tensors: SageScaleTensors | None = None + sage_scale_arr: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + _alloc: Constexpr[SmemAllocation | None] = None + + def __post_init__(self) -> None: + assert (self.pipeline_config is not None) == self.in_smem, ( + "the SMEM form of the Sage K scales carries the instance's pipeline " + "and the register form carries none" + ) + super().__post_init__() + + @property + def groups(self) -> int: + """Return the scale groups per K32 fragment of the served geometry.""" + if self.summary: + return self.cfg.sage_summary_k_groups_per_fragment + return self.cfg.sage_k_groups_per_fragment + + @property + def in_smem(self) -> bool: + """Whether the tile's words live in the SMEM ring rather than registers.""" + return self.cfg.sage_k_scales_in_smem_for(self.groups) + + @property + def arr_size(self) -> int: + """Return the ``sfK`` words one lane holds per tile in the register form.""" + return sage_scale_arr_size(self.cfg, self.groups) + + @property + def tile_words(self) -> int: + """Return the ``sfK`` words of one tile slot of the SMEM ring.""" + return sage_k_scale_words(self.cfg, self.groups) + + @property + def routed_words(self) -> int: + """Return the length of the routed array: the lane's words, or one placeholder.""" + return 1 if self.in_smem else self.arr_size + + def _init_placeholder_state(self) -> None: + """Create the placeholder routed array for task-graph tracing.""" + self.sage_scale_arr.default = _placeholder_local_array( + Float32, self.routed_words + ) + + def get_smem_requirements(self) -> list[SmemAllocation]: + """Allocate the ring of the SMEM form; the register form needs none.""" + if not self.in_smem: + return [] + if self._alloc is None: + self._alloc = SmemAllocation( + name=self.name, + size_bytes=SAGE_K_SCALES_RING_STAGES * self.tile_words * 4, + alignment=16, + ) + return [self._alloc] + + def get_tmem_requirements(self) -> list[TmemAllocation]: + """The words live in SMEM or registers only.""" + return [] + + # -- schedule steps ---------------------------------------------------- + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def publish_tile(self, stage_info: StageInfo) -> None: + """Fill the acquired ring slot with a dense tile's words (SMEM form).""" + self._fill_ring(stage_info, cutlass.Boolean(False)) + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def publish_route_tile(self, stage_info: StageInfo, *, route_flags: Int32) -> None: + """Fill the acquired ring slot with a block-sparse route's words (SMEM form). + + ``route_flags`` comes from the preceding ``load_route`` of the route + metadata resource, whose stage must stay held until this step and + ``take_route_tile`` have run. + """ + self._fill_ring(stage_info, self._route_is_proxy(route_flags)) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=sage_scale_arr) + @cute.jit + def take_tile(self, stage_info: StageInfo) -> cutlass.Array: + """Return the routed array of a dense tile.""" + return self._take(stage_info, cutlass.Boolean(False)) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=sage_scale_arr) + @cute.jit + def take_route_tile( + self, stage_info: StageInfo, *, route_flags: Int32 + ) -> cutlass.Array: + """Return the routed array of a block-sparse route.""" + return self._take(stage_info, self._route_is_proxy(route_flags)) + + @cute.jit + def _route_is_proxy(self, route_flags: Int32) -> cutlass.Boolean: + """Return the route kind, resolved only where it selects a geometry.""" + route_is_proxy = cutlass.Boolean(False) + if cutlass.const_expr(self._selects_by_route): + route_is_proxy = cute.arch.make_warp_uniform(_route_is_proxy(route_flags)) + return route_is_proxy + + @cute.jit + def _fill_ring( + self, stage_info: StageInfo, route_is_proxy: cutlass.Boolean + ) -> None: + """Store the tile's words into the acquired slot (SMEM form). + + In a mixed-geometry plan only the resource of the route's kind fills; + the branch is CTA-uniform. + """ + assert self.in_smem + if self._serves(route_is_proxy): + self._store_tile_words(stage_info) + + @cute.jit + def _take( + self, stage_info: StageInfo, route_is_proxy: cutlass.Boolean + ) -> cutlass.Array: + """Return the routed array: the lane's words in the register form, else ones.""" + words = cutlass.Array( + Float32, self.routed_words, space=cutlass.AddressSpace.rmem + ) + for entry in cutlass.range_constexpr(self.routed_words): + words[entry] = Float32(1.0) + if cutlass.const_expr(not self.in_smem): + if self._serves(route_is_proxy): + self._gather_lane_words(stage_info, words) + return words + + @property + def _selects_by_route(self) -> bool: + """Whether the route kind selects between two geometries of this plan.""" + return self.cfg.sage_mixed_k_geometry and self.route_metadata is not None + + @cute.jit + def _serves(self, route_is_proxy: cutlass.Boolean) -> cutlass.Boolean: + """Whether this resource holds the route's geometry; constant for one geometry.""" + if cutlass.const_expr(not self._selects_by_route): + return True + if cutlass.const_expr(self.summary): + return route_is_proxy + return not route_is_proxy + + # -- storage fills ----------------------------------------------------- + + @cute.jit + def _gather_lane_words(self, stage_info: StageInfo, words: cutlass.Array) -> None: + """Take the lane's spatial half of the tile's words into ``words``. + + Every entry index is a constant, so no pass indexes registers at run + time. + """ + word_value = self._word_source(stage_info) + arr_size = self.arr_size + warp_grp_thread_idx = Int32( + _decode_gen_task_cache(stage_info)[_TASK_CACHE_WARP_GRP_THREAD_IDX] + ) + half = _keeps_spatial_half(self.cfg, warp_grp_thread_idx) + half_base = half * Int32(arr_size) + for entry in cutlass.range_constexpr(arr_size): + words[entry] = word_value(half_base + Int32(entry), half, entry) + + @cute.jit + def _store_tile_words(self, stage_info: StageInfo) -> None: + """Store the tile's words into the slot of the acquired producer stage.""" + word_value = self._word_source(stage_info) + num_words = self.tile_words + arr_size = self.arr_size + threads = 32 * self.cfg.softmax_num_warps(self.inst_id) + warp_grp_thread_idx = Int32( + _decode_gen_task_cache(stage_info)[_TASK_CACHE_WARP_GRP_THREAD_IDX] + ) + buffer = self._ring(stage_info.context) + slot_base = Int32(stage_info.stage_idx) * Int32(num_words) + for round_idx in cutlass.range_constexpr((num_words + threads - 1) // threads): + word_idx = warp_grp_thread_idx + Int32(round_idx * threads) + if word_idx < Int32(num_words): + buffer[slot_base + word_idx] = word_value( + word_idx, word_idx // Int32(arr_size), word_idx % Int32(arr_size) + ) + + @cute.jit + def _ring(self, context: ResourceContext) -> cutlass.Array: + """Return the ring as an FP32 SMEM array.""" + return cutlass.Array( + context.smem_base.data_ptr() + self._alloc.offset, + dtype=Float32, + shape=(SAGE_K_SCALES_RING_STAGES * self.tile_words,), + addrspace=3, + ) + + # -- word sources ------------------------------------------------------ + + @cute.jit + def _word_source(self, stage_info: StageInfo) -> WordSource: + """Return ``word_value(word_idx, half, lane_entry)`` for the current tile. + + ``word_idx = half * arr_size + lane_entry`` in the + ``sage_k_scale_words`` layout. + """ + cfg = self.cfg + if cutlass.const_expr(self.route_metadata is None): + # A dense tile gathers ``k_scale`` over its token range; the load + # clamps tokens past the sequence end. + ( + seq_len_kv, + _q_group_idx, + _element_mask_end_idx, + tile_offset_k, + _window_start_idx, + _is_valid_effective_tile, + _is_masked_final_wave, + _tile_is_unmasked, + _rows_are_active, + ) = resolve_keeps_tile_context( + cfg, + stage_info, + inst_id=self.inst_id, + seqlens_kv=self.seqlens_kv, + max_seq_len_kv=self.max_seq_len_kv, + seq_len_q=self.seq_len_q, + q_group_idx=self.q_group_idx, + ) + kv_head_idx, batch_idx = _logical_head_batch( + stage_info, self.h_k_idx, self.b_idx + ) + k_scale_addr = self.scale_tensors.k_scale_ptr.toint() + + def dense_word(word_idx: Int32, half: Int32, lane_entry) -> Float32: + _ = word_idx + return load_k_scale( + cfg, + k_scale_addr, + self.scale_tensors.k_scale_head_stride, + kv_head_idx=kv_head_idx, + batch_idx=batch_idx, + seq_len_kv=seq_len_kv, + kv_token_idx=dense_k_scale_token( + cfg, half, lane_entry, tile_offset_k + ), + ) + + return dense_word + + metadata = self.route_metadata + stage_base = metadata._consumer_stage_base() + if cutlass.const_expr(not self.summary): + # The load warp staged the route's words in the + # ``sage_k_scale_words`` layout behind the route record. + assert metadata.staging_layout.sage_scale_words_word_offset is not None + source = stage_base + Int32( + metadata.staging_layout.sage_scale_words_word_offset + ) + staged = metadata._smem_scales + + def staged_word(word_idx: Int32, half: Int32, lane_entry) -> Float32: + _ = half, lane_entry + return Float32(staged[source + word_idx]) + + return staged_word + + # A proxy route of a mixed-geometry plan stages no scale words; gather + # each summary scale from the route's staged atom origins. + groups = self.groups + num_origins = metadata.staging_layout.num_origin_words + kv_head_idx, batch_idx = _logical_head_batch( + stage_info, self.h_k_idx, self.b_idx + ) + scale_addr = self.scale_tensors.k_summary_scale_ptr.toint() + scale_head_stride = self.scale_tensors.k_summary_scale_head_stride + summary_seq_len = Int32(cfg.num_proxy_summaries) + log2_k_block_size = Int32(log2_block_size(cfg.sage_k_summary_block_size)) + origins = metadata._smem_words + + def summary_word(word_idx: Int32, half: Int32, lane_entry) -> Float32: + _ = word_idx + atom_idx, token_offset = sage_word_position(cfg, half, lane_entry, groups) + atom_idx = cute.math.min(atom_idx, Int32(num_origins - 1)) + origin = Int32(origins[stage_base + atom_idx]) + return load_k_scale( + cfg, + scale_addr, + scale_head_stride, + kv_head_idx=kv_head_idx, + batch_idx=batch_idx, + seq_len_kv=summary_seq_len, + kv_token_idx=cute.math.max(origin, Int32(0)) + token_offset, + log2_k_block_size=log2_k_block_size, + ) + + return summary_word + + # -- pass-side reads --------------------------------------------------- + + @cute.jit + def open( + self, + stage_info: StageInfo, + scale_arr: cutlass.Array, + factor: Float32 | None, + ): + """Return a pass's view of the tile's words with ``factor`` applied. + + Register form: a rotating copy of the routed array, scaled once. + SMEM form: the lane's half pointer into the waited slot and the + factor, applied per fragment. + """ + if cutlass.const_expr(self.in_smem): + warp_grp_thread_idx = Int32( + _decode_gen_task_cache(stage_info)[_TASK_CACHE_WARP_GRP_THREAD_IDX] + ) + half_base = Int32(self.consumer_work_stage) * Int32( + self.tile_words + ) + _keeps_spatial_half(self.cfg, warp_grp_thread_idx) * Int32( + self.arr_size + ) + half_ptr = self._ring(stage_info.context).data_ptr() + half_base + return half_ptr, factor + arr_size = self.arr_size + words = cutlass.Array(Float32, arr_size, space=cutlass.AddressSpace.rmem) + for entry in cutlass.range_constexpr(arr_size): + words[entry] = Float32(scale_arr[entry]) + if cutlass.const_expr(factor is not None): + scale_pairs_in_place(words, factor, arr_size) + return words + + @cute.jit + def fragment(self, view, fragment_idx: Int32) -> cutlass.Array: + """Return ``factor * sfK`` of one fragment's ``groups`` words. + + The SMEM form uses 16-byte loads, which callers issue ahead of the + fragment's score wait; the register form reads the leading entries + of the rotating array. + """ + groups = self.groups + values = cutlass.Array(Float32, groups, space=cutlass.AddressSpace.rmem) + if cutlass.const_expr(self.in_smem): + half_ptr, factor = view + assert groups % 4 == 0 + words_ptr = half_ptr + fragment_idx * Int32(groups) + for chunk in cutlass.range_constexpr(0, groups, 4): + loaded = (words_ptr + Int32(chunk)).load(count=4, alignment=16) + for elem in cutlass.range_constexpr(4): + values[chunk + elem] = Float32(loaded[elem]) + if cutlass.const_expr(factor is not None): + scale_pairs_in_place(values, factor, groups) + else: + for group_idx in cutlass.range_constexpr(groups): + values[group_idx] = Float32(view[group_idx]) + return values + + @cute.jit + def advance(self, view) -> None: + """Rotate the register array down by one fragment; the SMEM form reads in place.""" + if cutlass.const_expr(not self.in_smem): + groups = self.groups + for entry in cutlass.range_constexpr(self.arr_size - groups): + view[entry] = Float32(view[entry + groups]) + + +@dataclass(kw_only=True) +class SageVScalesResource(DecodeGenResourceBase): + """The staged ``sfV`` and ``v_mean`` of the work tile's KV head. + + The pipeline is a one-stage pipeline of the correction warps. + """ + + cfg: Constexpr[FmhaDecodeConfig] = None + scale_tensors: SageScaleTensors | None = None + h_k_idx: Int32 | None = None + b_idx: Int32 | None = None + _alloc: Constexpr[SmemAllocation | None] = None + + @property + def entries(self) -> int: + """Return the FP32 entries of the block: the scales, then the means.""" + return staged_v_channel_scale_entries(self.cfg) + + def get_smem_requirements(self) -> list[SmemAllocation]: + """Allocate the one-stage block of scales and means.""" + if self._alloc is None: + self._alloc = SmemAllocation( + name=self.name, + size_bytes=self.entries * 4, + alignment=16, + ) + return [self._alloc] + + def get_tmem_requirements(self) -> list[TmemAllocation]: + """The scales live in SMEM only.""" + return [] + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def stage_tile(self, stage_info: StageInfo) -> None: + """Copy the work tile's KV head scales and means into the acquired block.""" + cfg = self.cfg + task_cache = _decode_gen_task_cache(stage_info) + kv_head_idx, _ = _logical_head_batch(stage_info, self.h_k_idx, self.b_idx) + stage_v_channel_scales( + cfg, + self.scale_tensors, + self.staged(stage_info.context), + kv_head_idx=kv_head_idx, + thread_idx=task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX], + num_threads=cfg.correction_barrier_threads, + ) + + @cute.jit + def staged(self, context: ResourceContext) -> cutlass.Array: + """Return the block as an FP32 SMEM array.""" + return cutlass.Array( + context.smem_base.data_ptr() + self._alloc.offset, + dtype=Float32, + shape=(self.entries,), + addrspace=3, + ) + + @cute.jit + def channel_scales( + self, staged: cutlass.Array, *, first_col: Int32, count: Constexpr[int] + ) -> tuple[cutlass.Array, cutlass.Array]: + """Return ``count`` scales and means from ``first_col`` of the waited block. + + Without ``sage_v_mean`` the means are zero so callers apply one fused + multiply-add. + """ + return load_staged_v_channel_scales( + self.cfg, staged, first_col=first_col, count=count + ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_block_sparse_metadata.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_block_sparse_metadata.py index f79f3a1e0566..22ca7169cf47 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_block_sparse_metadata.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_block_sparse_metadata.py @@ -27,7 +27,7 @@ import cutlass import cutlass.cute as cute -from cutlass import Int32, Uint32 +from cutlass import Float32, Int32, Uint32 from cutlass.experimental.task_scheduling.enums import WorkAttr from cutlass.experimental.task_scheduling.memory import ( ResourceContext, @@ -50,6 +50,7 @@ from ...stage import FmhaStage from ..fmha_decode_config import FmhaDecodeConfig from .helpers_common import ( + _SOFTMAX_ROUTE_IS_PROXY_FLAG, _TASK_CACHE_SEQ_LEN_KV, _TASK_CACHE_WARP_IDX, _TASK_CACHE_WARP_GRP_THREAD_IDX, @@ -57,10 +58,21 @@ DecodeGenResourceBase, ResourceVars, _decode_gen_task_cache, + _keeps_spatial_half, + _logical_head_batch, _sparse_task_cache_route_begin, _sparse_task_cache_route_count, _warp_broadcast_i32, ) +from .sage_scales import ( + SageScaleTensors, + block_sparse_k_scale_source, + load_k_scale, + sage_scale_arr_size, + sage_word_position, + sage_k_scale_words, + sage_staged_k_scale_words, +) # Keeps staging uses the low four bits for structural KV64 validity. Bit 4 @@ -69,9 +81,7 @@ # max pass derives its keep words from the token words directly, so the bit is # currently staged for the consumer but not read. _SOFTMAX_TOKEN_MASK_IS_FULL_FLAG = 1 << 4 -# Keeps reserves bit 5 for the prepared route kind. The low four structural -# validity bits and bit 4 keep their existing meaning. -_SOFTMAX_ROUTE_IS_PROXY_FLAG = 1 << 5 + # SWAP origins are at least eight-token aligned, so their low two bits are free # while the route is in Softmax's private staging payload. Reusing them avoids @@ -149,12 +159,13 @@ def _kv_retained_route_words( class _BlockSparseSoftmaxStagingLayout: """Layout of the staged cross-warp Softmax metadata payload. - Keeps retains all route origins, a flags word, alignment padding, and the - optional K32 token words. KV256 consumers then select the four words owned - by their spatial half. SWAP stores execution-ordered origins followed by - optional logical K32 token words, one for each consumer warp. Selected - prepared route flags travel in the otherwise-zero low bits of each warp's - first aligned origin. + Keeps retains all route origins, a flags word, alignment padding, the + optional K32 token words and, for Sage attention, the route's ``sfK`` + words on their own aligned boundary. KV256 consumers then select the four + words owned by their spatial half. SWAP stores execution-ordered origins + followed by optional logical K32 token words, one for each consumer warp. + Selected prepared route flags travel in the otherwise-zero low bits of + each warp's first aligned origin. """ # Logical-origin scalars staged for one complete KV route. @@ -163,6 +174,7 @@ class _BlockSparseSoftmaxStagingLayout: origins_per_warp: int route_flags_word_offset: int | None token_words_word_offset: int | None + sage_scale_words_word_offset: int | None stage_stride_words: int total_words: int @@ -178,12 +190,14 @@ def create( use_keeps_mma_ab: bool, route_layout: _BlockSparseRouteLayout, num_stages: int, + sage_scale_words: int = 0, ) -> "_BlockSparseSoftmaxStagingLayout": """Build a stage-count-dependent layout for Softmax metadata.""" kv_route_size = route_layout.kv_route_size atom_size = route_layout.atom_size has_token_bits = route_layout.has_token_bits + sage_scale_words_word_offset = None if use_keeps_mma_ab: num_origin_words = kv_route_size // atom_size assert num_origin_words <= 4, ( @@ -195,7 +209,11 @@ def create( token_words_word_offset = aligned_payload_words if has_token_bits else None token_words = kv_route_size // 32 if has_token_bits else 0 stage_stride_words = ((aligned_payload_words + token_words + 3) // 4) * 4 + if sage_scale_words > 0: + sage_scale_words_word_offset = stage_stride_words + stage_stride_words += ((sage_scale_words + 3) // 4) * 4 else: + assert sage_scale_words == 0, "Sage scales are staged for Keeps only" softmax_atom_size = min(atom_size, 32) num_origin_words = kv_route_size // softmax_atom_size origins_per_warp = 32 // softmax_atom_size @@ -209,6 +227,7 @@ def create( origins_per_warp=origins_per_warp, route_flags_word_offset=route_flags_word_offset, token_words_word_offset=token_words_word_offset, + sage_scale_words_word_offset=sage_scale_words_word_offset, stage_stride_words=stage_stride_words, total_words=num_stages * stage_stride_words, ) @@ -706,6 +725,10 @@ class SmemBlockSparseSoftmaxMetadataResource(DecodeGenResourceBase): data-dependent branch; each consumer receives at most four words through the stable task-local ABI. Runtime route flags carry the conservative FULL summary and, for proxy-capable builds, the route source kind. + + Sage attention adds the route's ``sfK`` words to the staged payload, so + the softmax threads read them from SMEM instead of issuing scattered + global loads before the score wait. """ _task_local_specs: ClassVar[tuple[tuple, ...]] = ( @@ -747,8 +770,12 @@ class SmemBlockSparseSoftmaxMetadataResource(DecodeGenResourceBase): route_metadata: cute.Pointer | None = None staging_layout: Constexpr[_BlockSparseSoftmaxStagingLayout | None] = None route_layout: Constexpr[_BlockSparseRouteLayout | None] = None + h_k_idx: Int32 | None = None + b_idx: Int32 | None = None + scale_tensors: SageScaleTensors | None = None _alloc: Constexpr[SmemAllocation | None] = None _smem_words: cutlass.Array = None + _smem_scales: cutlass.Array = None softmax_origin0_slot: Constexpr[TaskLocalVariable] = ( TaskLocalVariable.uninitialized() ) @@ -782,6 +809,7 @@ def __post_init__(self) -> None: use_keeps_mma_ab=self.cfg.use_keeps_mma_ab, route_layout=self.route_layout, num_stages=self.pipeline_config.num_stages, + sage_scale_words=sage_staged_k_scale_words(self.cfg), ) super().__post_init__() @@ -791,6 +819,9 @@ def _init_placeholder_state(self) -> None: self._smem_words = _placeholder_smem_array( Int32, self.staging_layout.total_words ) + self._smem_scales = _placeholder_smem_array( + Float32, self.staging_layout.total_words + ) def get_smem_requirements(self) -> list[SmemAllocation]: """Allocate one metadata payload per configured pipeline stage.""" @@ -821,6 +852,13 @@ def _create_initial_task_locals( shape=(self.staging_layout.total_words,), addrspace=3, ) + # Float32 view of the same allocation for the staged sfK words. + self._smem_scales = cutlass.Array( + context.smem_base.data_ptr() + self._alloc.offset, + dtype=Float32, + shape=(self.staging_layout.total_words,), + addrspace=3, + ) return {} @producer_work(work_attrs=WorkAttr.AUXILIARY) @@ -1100,6 +1138,90 @@ def _store_route_keeps( + Int32(self.staging_layout.token_words_word_offset) + lane_idx ] = Int32(token_word) + if cutlass.const_expr(self.cfg.use_sage_attention): + # Stage the route's ``sfK`` words in the exact geometry. Lane ``l`` + # owns words ``l * words_per_lane`` onward, which lie in one + # fragment of one spatial half, so it resolves one atom origin and + # loads its words at the group stride. Invalid atoms (origin + # ``-1``) are masked; the clamp to token zero only keeps the load + # in bounds. Proxy routes of a mixed-geometry plan stage nothing: + # the softmax warps gather their summary scales behind the score + # wait. + cfg = self.cfg + assert self.staging_layout.sage_scale_words_word_offset is not None + route_is_proxy = cutlass.Boolean(False) + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): + route_is_proxy = cutlass.Boolean( + ( + _warp_broadcast_i32( + resolved_record_word, + self.route_layout.route_flags_word_offset, + ) + & Int32(_PREPARED_ROUTE_IS_PROXY_FLAG) + ) + != Int32(0) + ) + task_cache = _decode_gen_task_cache(stage_info) + scale_addr, scale_head_stride, scale_seq_len, log2_k_block_size = ( + block_sparse_k_scale_source( + cfg, + self.scale_tensors, + seq_len_kv=Int32(task_cache[_TASK_CACHE_SEQ_LEN_KV]), + route_is_proxy=route_is_proxy, + ) + ) + kv_head_idx, batch_idx = _logical_head_batch( + stage_info, self.h_k_idx, self.b_idx + ) + groups: Constexpr[int] = cfg.sage_k_groups_per_fragment + num_words: Constexpr[int] = sage_k_scale_words(cfg, groups) + words_per_lane: Constexpr[int] = max(1, num_words // 32) + assert words_per_lane * 32 >= num_words and groups % words_per_lane == 0 + group_tokens: Constexpr[int] = cfg.softmax_score_fragment_regs // groups + arr_size: Constexpr[int] = sage_scale_arr_size(cfg, groups) + # Build the view outside the dynamic branch so the store reuses + # its base pointer. + staged = cutlass.Array( + self._smem_scales.data_ptr(), + dtype=Float32, + shape=(self.staging_layout.total_words,), + addrspace=3, + ) + if cutlass.const_expr(not cfg.sage_mixed_k_geometry) or not route_is_proxy: + first_word = lane_idx * Int32(words_per_lane) + atom_idx, token_offset = sage_word_position( + cfg, + first_word // Int32(arr_size), + first_word % Int32(arr_size), + groups, + ) + atom_idx = cute.math.min(atom_idx, Int32(num_origins - 1)) + if cutlass.const_expr(num_origins == 2 and not uses_one_warp_transport): + origin = Int32(resolved_record_word) + if atom_idx != Int32(0): + origin = Int32(resolved_origin1) + else: + origin = Int32( + cute.arch.shuffle_sync(Int32(resolved_record_word), atom_idx) + ) + if first_word < Int32(num_words): + first_token = cute.math.max(origin, Int32(0)) + token_offset + word_base = ( + stage_base + + Int32(self.staging_layout.sage_scale_words_word_offset) + + first_word + ) + for entry in cutlass.range_constexpr(words_per_lane): + staged[word_base + Int32(entry)] = load_k_scale( + cfg, + scale_addr, + scale_head_stride, + kv_head_idx=kv_head_idx, + batch_idx=batch_idx, + seq_len_kv=scale_seq_len, + kv_token_idx=first_token + Int32(entry * group_tokens), + log2_k_block_size=log2_k_block_size, + ) cute.arch.sync_warp() @producer_work @@ -1252,7 +1374,7 @@ def load_route( warp_grp_thread_idx = Int32( _decode_gen_task_cache(stage_info)[_TASK_CACHE_WARP_GRP_THREAD_IDX] ) - spatial = warp_grp_thread_idx >> Int32(6) + spatial = _keeps_spatial_half(self.cfg, warp_grp_thread_idx) origin0_idx = spatial origin1_idx = spatial + Int32(2) valid0 = (stored_route_flags >> origin0_idx) & Int32(1) 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 330f15b05691..80254f62ea52 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,20 +42,22 @@ ) from ..fmha_decode_config import FmhaDecodeConfig -from ...._block_sparse.common import _block_sparse_proxy_summary_geometry +from ..fmha_decode_constants import ( + FP8_P_QUANT_LOG2_SCALE, + INT32_SCORE_BIAS, +) from ...placeholder_helpers import _placeholder_smem_array from .helpers_common import ( + _route_is_proxy, Constexpr, DecodeGenResourceBase, ResourceVars, fadd2, ffma2, - fmul2, _TASK_CACHE_LANE_IDX, _TASK_CACHE_WARP_GRP_THREAD_IDX, _TASK_CACHE_WARP_IDX, _decode_gen_task_cache, - _fp8_log2_quant_scale, _is_last_loop_iteration, _keeps_col_base, _keeps_row_idx, @@ -65,7 +67,6 @@ _neg_max_f32, _pack_float2_to_bf16, _pack_float2_to_fp16, - _swaps_routed_coordinate, _wait_for_mbarrier_phase, ) from .helpers_output import ( @@ -77,13 +78,12 @@ ) from .helpers_softmax import ( _compute_fp8_p_regs_and_local_sums, - _compute_fp8_p_regs_and_local_sums_dense, - _compute_p_values_and_local_sums_dense, + _compute_p_values_and_local_sums, _ex2_emulation_packed_f32x2, _pack_float4_to_fp8_e4m3, _pack_float4_to_fp8_e4m3_inline, ) -from .smem_block_sparse_metadata import _SOFTMAX_ROUTE_IS_PROXY_FLAG +from .sage_scales import SageKScalesResource from .tmem_s import TmemSResource # Tunable: number of score pairs per streamed fragment whose exponentials run @@ -144,6 +144,10 @@ class SmemPResource(DecodeGenResourceBase): use_variable_seqlens_kv: Constexpr[bool] = False tmem_s_ref: Constexpr[TmemSResource] = None tmem_o_ref: Constexpr[object] = None + # The instance's ``sfK`` resources, shared with the S resource; proxy + # routes read ``sage_summary_k_scales``, everything else ``sage_k_scales``. + sage_k_scales: SageKScalesResource | None = None + sage_summary_k_scales: SageKScalesResource | None = None _alloc: Constexpr[SmemAllocation | None] = None _fragment_ready_alloc: Constexpr[SmemAllocation | None] = None _tmem_alloc: Constexpr[TmemAllocation | None] = None @@ -160,8 +164,8 @@ def _init_placeholder_state(self) -> None: """Create placeholder P storage state.""" self._tmem_base_addr = Int32(0) self._smem_base_p = _placeholder_smem_array( - self.cfg.q_dtype, - self.cfg.smem_p_tile_bytes // self.cfg.q_dtype_bytes, + self.cfg.pv_dtype, + self.cfg.smem_p_tile_bytes // self.cfg.pv_dtype_bytes, ) self._smem_base_p_i32 = _placeholder_smem_array( Int32, self.cfg.smem_p_tile_bytes // 4 @@ -266,8 +270,8 @@ def _create_initial_task_locals( smem_base_ptr = context.smem_base.data_ptr() + self._alloc.offset self._smem_base_p = cutlass.Array( smem_base_ptr, - dtype=self.cfg.q_dtype, - shape=(self.cfg.smem_p_tile_bytes // self.cfg.q_dtype_bytes,), + dtype=self.cfg.pv_dtype, + shape=(self.cfg.smem_p_tile_bytes // self.cfg.pv_dtype_bytes,), addrspace=3, ) self._smem_base_p_i32 = cutlass.Array( @@ -319,26 +323,36 @@ def init_descriptor_state(self, stage_info: StageInfo) -> None: self._create_initial_task_locals(stage_info.context) @cute.jit - def _apply_proxy_route_denominator_mass( + def _exponent_addend( self, - local_sum: Float32, - tail_p: Float32, - route_is_proxy: Int32, + new_max: Float32, + route_is_proxy: cutlass.Boolean, + *, + guards_masked_rows: Constexpr[bool] = True, ) -> Float32: - """Weight only a proxy route's softmax denominator by block mass.""" - - if cutlass.const_expr(not self.cfg.use_block_sparse_proxy_routes): - return local_sum - if route_is_proxy != Int32(0): - _, tail_len = _block_sparse_proxy_summary_geometry( - self.cfg.static_seq_len_kv, - self.cfg.kv_block_size, - ) - local_sum *= Float32(self.cfg.kv_block_size) - tail_delta = tail_len - self.cfg.kv_block_size - if cutlass.const_expr(tail_delta != 0): - local_sum += Float32(tail_delta) * tail_p - return local_sum + """Return ``addend`` of one row's P exponent ``exp2(c * s + addend)``. + + - ``-c * max`` anchors the row; a fully masked row (``-FLT_MAX`` + maximum) anchors on zero so its scores exponentiate to zero, not NaN. + Callers whose masked rows are skipped or discarded downstream pass + ``guards_masked_rows=False``. + - Byte-wide P adds ``log2(448)``, its static FP8 quantization scale. + - A proxy route adds ``log2`` of its block mass, as the max pass does, + so each summary probability carries the tokens it stands for. + """ + cfg = self.cfg + safe_new_max = new_max + if cutlass.const_expr(guards_masked_rows): + if safe_new_max == _neg_max_f32(): + safe_new_max = Float32(0.0) + # Not ``cute.math.fma``: the fused form makes ptxas spill in the callers. + addend = Float32(-self.scale_softmax_log2 * safe_new_max) + if cutlass.const_expr(self.cfg.use_fp8_pv): + addend += Float32(FP8_P_QUANT_LOG2_SCALE) + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): + if route_is_proxy: + addend += Float32(cfg.proxy_log2_block_mass) + return addend @producer_work @cute.jit @@ -352,9 +366,29 @@ def compute_p_fragments( self._compute_p_fragments_impl( stage_info, new_max_arr=new_max_arr, - route_is_proxy=Int32(0), - route_origin0=Int32(0), - route_origin1=Int32(0), + route_is_proxy=cutlass.Boolean(False), + ) + + @producer_work + @cute.jit + def compute_sage_p_fragments( + self, + stage_info: StageInfo, + *, + new_max_arr: cutlass.Array, + sage_q_scale: Float32, + sage_scale_arr: cutlass.Array, + sage_summary_scale_arr: cutlass.Array, + ) -> None: + """Stream every K32 fragment with per-group Sage dequantization scales.""" + assert self.cfg.use_sage_attention + self._compute_p_fragments_impl( + stage_info, + new_max_arr=new_max_arr, + route_is_proxy=cutlass.Boolean(False), + sage_q_scale=sage_q_scale, + sage_scale_arr=sage_scale_arr, + sage_summary_scale_arr=sage_summary_scale_arr, ) @producer_work @@ -365,20 +399,36 @@ def compute_proxy_route_p_fragments( *, new_max_arr: cutlass.Array, route_flags: Int32, - route_origin0: Int32, - route_origin1: Int32, ) -> None: - """Stream every proxy-capable KV256 K32 fragment from one rolled loop.""" + """Stream every exact or proxy K32 fragment from one rolled loop.""" assert self.cfg.use_block_sparse_proxy_routes - route_is_proxy = Int32( - (route_flags & Int32(_SOFTMAX_ROUTE_IS_PROXY_FLAG)) != Int32(0) + self._compute_p_fragments_impl( + stage_info, + new_max_arr=new_max_arr, + route_is_proxy=_route_is_proxy(route_flags), ) + + @producer_work + @cute.jit + def compute_sage_proxy_route_p_fragments( + self, + stage_info: StageInfo, + *, + new_max_arr: cutlass.Array, + route_flags: Int32, + sage_q_scale: Float32, + sage_scale_arr: cutlass.Array, + sage_summary_scale_arr: cutlass.Array, + ) -> None: + """Stream exact or proxy K32 fragments with per-group Sage scales.""" + assert self.cfg.use_block_sparse_proxy_routes and self.cfg.use_sage_attention self._compute_p_fragments_impl( stage_info, new_max_arr=new_max_arr, - route_is_proxy=route_is_proxy, - route_origin0=route_origin0, - route_origin1=route_origin1, + route_is_proxy=_route_is_proxy(route_flags), + sage_q_scale=sage_q_scale, + sage_scale_arr=sage_scale_arr, + sage_summary_scale_arr=sage_summary_scale_arr, ) @cute.jit @@ -387,75 +437,192 @@ def _compute_p_fragments_impl( stage_info: StageInfo, *, new_max_arr: cutlass.Array, - route_is_proxy: Int32, - route_origin0: Int32, - route_origin1: Int32, + route_is_proxy: cutlass.Boolean, + sage_q_scale: Float32 | None = None, + sage_scale_arr: cutlass.Array | None = None, + sage_summary_scale_arr: cutlass.Array | None = None, ) -> None: """Reload, exponentiate, and publish all K32 fragments in a rolled loop. The fragment index is a runtime loop variable, so the exponentiation body exists once in the instruction stream and only the TMEM column - offset, the fragment barrier, and the proxy tail bookkeeping depend on - it. Unrolling the fragments would replicate that body for every - fragment and both softmax instances and leave the softmax warps - instruction-fetch bound. The max pass has already written masked - scores back to TMEM, so the reload needs no mask logic of its own. + offset and the fragment barrier depend on it. Unrolling the fragments + would replicate that body for every fragment and both softmax + instances and leave the softmax warps instruction-fetch bound. The max + pass has already written masked (and mass-shifted) scores back to + TMEM, so the reload needs no mask or route logic beyond the proxy + addend. Sage attention changes only the per-group multipliers + ``c * sfQ * sfK_g``, biased INT32 scores only the per-group addends. """ - _ = stage_info cfg = self.cfg assert cfg.streams_tmem_p_fragments assert self._tmem_alloc.offset == self.tmem_s_ref._alloc.offset - # One FP32 score per column, two packed 16-bit probabilities per column. + # One FP32 score per column; two 16-bit or four FP8 probabilities per + # packed column. fragment_regs = cfg.softmax_score_fragment_regs - fragment_cols = fragment_regs // 2 + fragment_cols = cfg.fragment_p_packed_cols - new_max = new_max_arr[0] - safe_new_max = new_max - if safe_new_max == _neg_max_f32(): - safe_new_max = Float32(0.0) - minus_max_scale = Float32(-self.scale_softmax_log2 * safe_new_max) + exponent_addend = self._exponent_addend(new_max_arr[0], route_is_proxy) tmem_base = self._tmem_base_addr + Int32(self._tmem_alloc.offset) tidx, _, _ = cute.arch.thread_idx() publishes_fragment = (tidx & Int32(31)) == Int32(0) - total_sum = Float32(0.0) + # ``c * sfQ`` is one factor per lane and tile: the only multiplier of a + # route kind whose max pass returns dequantized scores (no bias), + # otherwise folded into the kind's raw ``sfK`` words. + exact_dequantized: Constexpr[bool] = cfg.sage_scores_dequantized_for( + proxy=False + ) + summary_dequantized: Constexpr[bool] = cfg.sage_scores_dequantized_for( + proxy=True + ) + exact_scales_in_smem: Constexpr[bool] = cfg.sage_k_scales_in_smem_for( + cfg.sage_k_groups_per_fragment + ) + # A mixed geometry selects the resource per fragment on the CTA-uniform + # route kind, unless proxy scores are dequantized and the exact words + # sit in registers: the exact resource then serves both kinds, routing + # ones on a proxy tile. + unified_scales: Constexpr[bool] = ( + cfg.sage_mixed_k_geometry + and summary_dequantized + and not exact_scales_in_smem + ) + mixed_scales: Constexpr[bool] = cfg.sage_mixed_k_geometry and not unified_scales + exp_scale = None + scales_view = None + summary_view = None + score_bias = None + row_multipliers = None + row_addends = None + if cutlass.const_expr(cfg.use_sage_attention): + exp_scale = self.scale_softmax_log2 * sage_q_scale + if cutlass.const_expr(exact_dequantized or summary_dequantized): + row_multipliers = cutlass.Array( + Float32, 1, space=cutlass.AddressSpace.rmem + ) + row_addends = cutlass.Array(Float32, 1, space=cutlass.AddressSpace.rmem) + row_multipliers[0] = exp_scale + row_addends[0] = exponent_addend + if cutlass.const_expr(not exact_dequantized): + scales_view = self.sage_k_scales.open( + stage_info, sage_scale_arr, exp_scale + ) + if cutlass.const_expr(mixed_scales and not summary_dequantized): + summary_view = self.sage_summary_k_scales.open( + stage_info, sage_summary_scale_arr, exp_scale + ) + if cutlass.const_expr(unified_scales and cfg.uses_int32_scores): + score_bias = Float32(INT32_SCORE_BIAS) + if route_is_proxy: + score_bias = Float32(0.0) + + # A fragment is scaled before its first exponential, and the next + # fragment's TMEM load issues once the scale FFMAs have consumed the + # scores, so it reuses their registers and hides behind the + # exponentials. The running sum stays a packed pair until the loop ends. + num_fragments = cfg.num_softmax_score_fragments + last_fragment = Int32(num_fragments - 1) + total_sum_pair = (Float32(0.0), Float32(0.0)) + s_arr = cutlass.Array(Float32, fragment_regs, space=cutlass.AddressSpace.rmem) + pending_scores = cutlass.Array( + Float32, fragment_regs, space=cutlass.AddressSpace.rmem + ) + self._load_score_fragment(tmem_base, Int32(0), pending_scores) for fragment_idx in cutlass.range(cfg.num_softmax_score_fragments, unroll=1): fragment = Int32(fragment_idx) - loaded = _keeps_tcgen05_ld( - cfg, - prims.make_tmem_ptr( - tmem_base + fragment * Int32(fragment_regs), Float32 - ), - num=fragment_regs, - offset=cfg.tile_size_kv // 2, - ) prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) - s_arr = cutlass.Array( - Float32, fragment_regs, space=cutlass.AddressSpace.rmem - ) for score_idx in cutlass.range_constexpr(fragment_regs): - s_arr[score_idx] = loaded[score_idx] - - local_sum = self._exponentiate_fragment_pairs(s_arr, minus_max_scale) - if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): - if route_is_proxy != Int32(0): - local_sum = self._proxy_fragment_sum( - local_sum, - s_arr, - fragment_origin=self._runtime_fragment_origin( - fragment, route_origin0, route_origin1 - ), + s_arr[score_idx] = pending_scores[score_idx] + + if cutlass.const_expr(mixed_scales): + # Each route kind scales the fragment in place with its own + # resource, so only the scores cross the branch. + if route_is_proxy: + if cutlass.const_expr(summary_dequantized): + self._scale_fragment_pairs( + s_arr, row_addends, row_multipliers, groups=1 + ) + else: + self._scale_fragment_with( + self.sage_summary_k_scales, + summary_view, + fragment, + s_arr, + exponent_addend, + ) + else: + if cutlass.const_expr(exact_dequantized): + self._scale_fragment_pairs( + s_arr, row_addends, row_multipliers, groups=1 + ) + else: + self._scale_fragment_with( + self.sage_k_scales, + scales_view, + fragment, + s_arr, + exponent_addend, + ) + elif cutlass.const_expr(exact_dequantized): + self._scale_fragment_pairs( + s_arr, row_addends, row_multipliers, groups=1 + ) + else: + fragment_multipliers = None + if cutlass.const_expr(cfg.use_sage_attention): + fragment_multipliers = self.sage_k_scales.fragment( + scales_view, fragment ) + self.sage_k_scales.advance(scales_view) + group_multipliers, group_addends = self._fragment_exponent_terms( + fragment_multipliers, + exponent_addend, + groups=cfg.sage_k_groups_per_fragment, + score_bias=score_bias, + ) + self._scale_fragment_pairs( + s_arr, + group_addends, + group_multipliers, + groups=cfg.sage_k_groups_per_fragment, + ) + # The last iteration reloads its own fragment so the loop body + # stays branch-free; the wait after the loop retires it. + next_fragment = fragment + Int32(1) + if next_fragment > last_fragment: + next_fragment = last_fragment + self._load_score_fragment(tmem_base, next_fragment, pending_scores) + local_sum_pair = self._exponentiate_fragment_pairs(s_arr) + total_sum_pair = cute.arch.add_packed_f32x2(total_sum_pair, local_sum_pair) - packed_p = ( - s_arr.data_ptr() - .load(count=fragment_regs, alignment=4) - .to(cfg.q_dtype) - .bitcast(Int32) - ) + if cutlass.const_expr(cfg.use_fp8_pv): + packed_regs = cutlass.Array( + Int32, fragment_cols, space=cutlass.AddressSpace.rmem + ) + for col_idx in cutlass.range_constexpr(fragment_cols): + value_base = col_idx * 4 + packed_regs[col_idx] = _pack_float4_to_fp8_e4m3_inline( + Float32(s_arr[value_base]), + Float32(s_arr[value_base + 1]), + Float32(s_arr[value_base + 2]), + Float32(s_arr[value_base + 3]), + ) + packed_p = packed_regs.data_ptr().load(count=fragment_cols, alignment=4) + else: + packed_p = ( + s_arr.data_ptr() + .load(count=fragment_regs, alignment=4) + .to(cfg.pv_dtype) + .bitcast(Int32) + ) + # Fence the packed TMEM store before one lane per warp arrives, so + # the MMA warp can start this fragment's PV k-slice. _keeps_tcgen05_st( cfg, - prims.make_tmem_ptr(tmem_base + fragment * Int32(fragment_cols), Int32), + prims.make_tmem_ptr( + tmem_base + fragment * Int32(cfg.fragment_p_packed_cols), Int32 + ), packed_p, offset=cfg.tmem_p_cols_per_inst, ) @@ -463,51 +630,154 @@ def _compute_p_fragments_impl( prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) if publishes_fragment: prims.mbarrier_arrive(self._fragment_ready.data_ptr() + fragment) - total_sum += local_sum + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + total_sum = Float32(total_sum_pair[0] + total_sum_pair[1]) self.tmem_s_ref.store_p_local_sum(0, total_sum) @cute.jit - def _runtime_fragment_origin( - self, fragment: Int32, route_origin0: Int32, route_origin1: Int32 - ) -> Int32: - """Return the token origin of a fragment selected at runtime. + def _scale_fragment_with( + self, + scales: Constexpr[SageKScalesResource], + view, + fragment: Int32, + s_arr: cutlass.Array, + exponent_addend: Float32, + ) -> None: + """Turn one fragment's scores into log2 exponents with one route kind's resource.""" + groups = scales.groups + fragment_multipliers = scales.fragment(view, fragment) + scales.advance(view) + group_multipliers, group_addends = self._fragment_exponent_terms( + fragment_multipliers, exponent_addend, groups=groups + ) + self._scale_fragment_pairs( + s_arr, group_addends, group_multipliers, groups=groups + ) + + @cute.jit + def _load_score_fragment( + self, tmem_base: Int32, fragment: Int32, scores: cutlass.Array + ) -> None: + """Issue the TMEM load of one K32 score fragment into ``scores``. - Each lane's fragments cover two K64 route atoms in order: the first - atom's fragments start at ``route_origin0``, the second atom's at - ``route_origin1``, and consecutive fragments within an atom advance by - one fragment width. + The caller waits for the load before reading ``scores``. The copy runs + before that wait, so it must stay a constant-indexed register rename + without optimization barriers, or it would read the destination early. """ cfg = self.cfg fragment_regs = cfg.softmax_score_fragment_regs - fragments_per_origin = cfg.softmax_fragments_per_route_atom - fragment_origin = Int32(route_origin0) - if fragment >= Int32(fragments_per_origin): - fragment_origin = Int32(route_origin1) - return fragment_origin + (fragment % Int32(fragments_per_origin)) * Int32( - fragment_regs + loaded = _keeps_tcgen05_ld( + cfg, + prims.make_tmem_ptr(tmem_base + fragment * Int32(fragment_regs), Float32), + num=fragment_regs, + offset=cfg.tile_size_kv // 2, + ) + for score_idx in cutlass.range_constexpr(fragment_regs): + scores[score_idx] = loaded[score_idx] + + @cute.jit + def _scale_fragment_pairs( + self, + s_arr: cutlass.Array, + group_addends: cutlass.Array, + group_multipliers: cutlass.Array, + groups: Constexpr[int], + ) -> None: + """Turn one fragment of scores into log2 exponents in place. + + Each score takes the multiplier and addend of its compile-time scale + group (``groups`` per fragment); the two scores of a pair differ in + group only for the one-token K block. + """ + fragment_regs = self.cfg.softmax_score_fragment_regs + group_regs = fragment_regs // groups + for value_idx in cutlass.range_constexpr(0, fragment_regs, 2): + group0: Constexpr[int] = value_idx // group_regs + group1: Constexpr[int] = (value_idx + 1) // group_regs + s_arr[value_idx], s_arr[value_idx + 1] = cute.arch.fma_packed_f32x2( + (Float32(s_arr[value_idx]), Float32(s_arr[value_idx + 1])), + ( + Float32(group_multipliers[group0]), + Float32(group_multipliers[group1]), + ), + (Float32(group_addends[group0]), Float32(group_addends[group1])), + ) + + @cute.jit + def _fragment_exponent_terms( + self, + fragment_multipliers: cutlass.Array | None, + exponent_addend: Float32, + groups: Constexpr[int], + score_bias: Float32 | None = None, + ) -> tuple[cutlass.Array, cutlass.Array]: + """Return one fragment's exponent multipliers and addends per scale group. + + The multipliers are the resource's ``c * sfQ * sfK_g`` under Sage, else + the softmax scale. Biased INT32 scores subtract ``bias * multiplier`` + from the row's addend; ``score_bias`` replaces the constant bias when + given (zero on a tile whose scores come back dequantized). + """ + cfg = self.cfg + neg_bias = Float32(-INT32_SCORE_BIAS) + if cutlass.const_expr(score_bias is not None): + neg_bias = -score_bias + group_multipliers = cutlass.Array( + Float32, groups, space=cutlass.AddressSpace.rmem ) + group_addends = cutlass.Array(Float32, groups, space=cutlass.AddressSpace.rmem) + if cutlass.const_expr(cfg.use_sage_attention): + for group_idx in cutlass.range_constexpr(groups): + group_multipliers[group_idx] = Float32(fragment_multipliers[group_idx]) + else: + group_multipliers[0] = self.scale_softmax_log2 + if cutlass.const_expr(cfg.uses_int32_scores): + # One rounding per group instead of one conversion per score; the + # error stays below one quantized score unit. + bias_pair = (neg_bias, neg_bias) + for group_base in cutlass.range_constexpr(0, groups - groups % 2, 2): + group_addends[group_base], group_addends[group_base + 1] = ( + cute.arch.fma_packed_f32x2( + bias_pair, + ( + Float32(group_multipliers[group_base]), + Float32(group_multipliers[group_base + 1]), + ), + (exponent_addend, exponent_addend), + ) + ) + if cutlass.const_expr(groups % 2 == 1): + group_addends[groups - 1] = Float32( + cute.math.fma( + neg_bias, + Float32(group_multipliers[groups - 1]), + exponent_addend, + ) + ) + else: + for group_idx in cutlass.range_constexpr(groups): + group_addends[group_idx] = exponent_addend + return group_multipliers, group_addends @cute.jit def _exponentiate_fragment_pairs( - self, s_arr: cutlass.Array, minus_max_scale: Float32 - ) -> Float32: - """Turn one fragment of scaled scores into probabilities in place. + self, s_arr: cutlass.Array + ) -> tuple[Float32, Float32]: + """Turn one fragment of log2 exponents into probabilities in place. - Returns the fragment's probability sum. Eight independent chains keep - the denominator update off one long dependency chain, and a configurable - subset of pairs runs its exponentials on the FMA pipe. + Returns the fragment's probability sum as a packed pair. Eight + independent chains keep the denominator update off one long dependency + chain; the first four pairs seed them, since an add to zero is a real + FADD. A configurable subset of pairs runs its exponentials on the FMA + pipe. """ pairs_per_fragment = self.cfg.softmax_score_fragment_regs // 2 + assert pairs_per_fragment >= 4 sum_chains = cutlass.Array(Float32, 8, space=cutlass.AddressSpace.rmem) - for chain_idx in cutlass.range_constexpr(8): - sum_chains[chain_idx] = Float32(0.0) for pair_idx in cutlass.range_constexpr(pairs_per_fragment): value_idx = pair_idx * 2 - p0, p1 = cute.arch.fma_packed_f32x2( - (Float32(s_arr[value_idx]), Float32(s_arr[value_idx + 1])), - (self.scale_softmax_log2, self.scale_softmax_log2), - (minus_max_scale, minus_max_scale), - ) + p0 = Float32(s_arr[value_idx]) + p1 = Float32(s_arr[value_idx + 1]) if cutlass.const_expr( _pair_uses_ex2_emulation(pair_idx, pairs_per_fragment) ): @@ -518,12 +788,16 @@ def _exponentiate_fragment_pairs( s_arr[value_idx] = p0 s_arr[value_idx + 1] = p1 chain_idx = (pair_idx & 3) * 2 - sum_chains[chain_idx], sum_chains[chain_idx + 1] = ( - cute.arch.add_packed_f32x2( - (sum_chains[chain_idx], sum_chains[chain_idx + 1]), - (p0, p1), + if cutlass.const_expr(pair_idx < 4): + sum_chains[chain_idx] = p0 + sum_chains[chain_idx + 1] = p1 + else: + sum_chains[chain_idx], sum_chains[chain_idx + 1] = ( + cute.arch.add_packed_f32x2( + (sum_chains[chain_idx], sum_chains[chain_idx + 1]), + (p0, p1), + ) ) - ) sum01 = cute.arch.add_packed_f32x2( (sum_chains[0], sum_chains[1]), (sum_chains[2], sum_chains[3]), @@ -532,44 +806,7 @@ def _exponentiate_fragment_pairs( (sum_chains[4], sum_chains[5]), (sum_chains[6], sum_chains[7]), ) - total_pair = cute.arch.add_packed_f32x2(sum01, sum23) - return Float32(total_pair[0] + total_pair[1]) - - @cute.jit - def _proxy_fragment_sum( - self, - local_sum: Float32, - s_arr: cutlass.Array, - *, - fragment_origin: Int32, - ) -> Float32: - """Weight a proxy fragment's sum by the token mass each summary stands for. - - KC stores one mean K vector per semantic KV block while VC stores its V - sum. P itself stays unweighted for PV; only the denominator accounts - for the represented token count, with the final summary covering the - shorter tail block. - """ - cfg = self.cfg - fragment_regs = cfg.softmax_score_fragment_regs - num_summaries, tail_len = _block_sparse_proxy_summary_geometry( - cfg.static_seq_len_kv, - cfg.kv_block_size, - ) - local_sum *= Float32(cfg.kv_block_size) - tail_delta = tail_len - cfg.kv_block_size - if cutlass.const_expr(tail_delta != 0): - final_summary_idx = num_summaries - 1 - final_summary_offset = Int32(final_summary_idx) - fragment_origin - if final_summary_offset >= Int32(0) and final_summary_offset < Int32( - fragment_regs - ): - # Proxy fragment origins are fragment-aligned in summary - # coordinates, so the tail's in-fragment lane is a compile-time - # constant even though route ownership is decided at runtime. - tail_lane = final_summary_idx % fragment_regs - local_sum += Float32(tail_delta) * Float32(s_arr[tail_lane]) - return local_sum + return cute.arch.add_packed_f32x2(sum01, sum23) @cute.jit def _compute_keeps_p( @@ -578,6 +815,7 @@ def _compute_keeps_p( *, new_max_arr: cutlass.Array, s_arr: cutlass.Array, + route_is_proxy: cutlass.Boolean, ) -> None: """Materialize one complete-row Keeps probability tile. @@ -607,13 +845,7 @@ def _compute_keeps_p( + stage_info.stage_idx * cfg.tmem_s_cols ) - new_max = new_max_arr[0] - safe_new_max = new_max - if safe_new_max == _neg_max_f32(): - safe_new_max = Float32(0.0) - neg_scaled_max = -self.scale_softmax_log2 * safe_new_max - if cutlass.const_expr(cfg.use_fp8_qkv): - neg_scaled_max += _fp8_log2_quant_scale() + exponent_addend = self._exponent_addend(new_max_arr[0], route_is_proxy) # Preserve four independent modulo-4 sum chains as two packed pairs, # without keeping a second 16-value P array live beside the S row. @@ -626,7 +858,7 @@ def _compute_keeps_p( for block_idx in cutlass.range_constexpr(num_vector_blocks): s_base = block_idx * vector_elements packed_base = block_idx * 4 if cfg.uses_two_inst_tmem_p else 0 - if cutlass.const_expr(cfg.use_fp8_qkv): + if cutlass.const_expr(cfg.use_fp8_pv): for packed_idx in cutlass.range_constexpr(4): val_base = packed_idx * 4 scaled_pair_01 = ffma2( @@ -635,7 +867,7 @@ def _compute_keeps_p( s_arr[s_base + val_base + 1], ), (self.scale_softmax_log2, self.scale_softmax_log2), - (neg_scaled_max, neg_scaled_max), + (exponent_addend, exponent_addend), ) p_pair_01 = ( cute.math.exp2(scaled_pair_01[0], fastmath=True), @@ -648,7 +880,7 @@ def _compute_keeps_p( s_arr[s_base + val_base + 3], ), (self.scale_softmax_log2, self.scale_softmax_log2), - (neg_scaled_max, neg_scaled_max), + (exponent_addend, exponent_addend), ) p_pair_23 = ( cute.math.exp2(scaled_pair_23[0], fastmath=True), @@ -672,7 +904,7 @@ def _compute_keeps_p( s_arr[s_base + val_base + 1], ), (self.scale_softmax_log2, self.scale_softmax_log2), - (neg_scaled_max, neg_scaled_max), + (exponent_addend, exponent_addend), ) p_pair = ( cute.math.exp2(scaled_pair[0], fastmath=True), @@ -682,7 +914,7 @@ def _compute_keeps_p( local_sum_pair_01 = fadd2(local_sum_pair_01, p_pair) else: local_sum_pair_23 = fadd2(local_sum_pair_23, p_pair) - if cutlass.const_expr(cfg.use_bf16_qkv): + if cutlass.const_expr(cfg.use_bf16_pv): packed_p[packed_base + packed_idx] = _pack_float2_to_bf16( p_pair[0], p_pair[1] ) @@ -717,11 +949,11 @@ def _compute_keeps_p( packed_p.data_ptr().load(count=4, alignment=4), alignment=16 ) if cutlass.const_expr(cfg.uses_two_inst_tmem_p): - # FP8 publishes the complete row with one x32 STTM. Dense 16-bit - # Q128/KV128 uses x16 slices to limit Softmax register pressure. - # Block-sparse two-instance profiles stream K32 fragments instead. - assert cfg.num_packed_p_regs in (32, 64) - regs_per_store = cfg.num_packed_p_regs if cfg.use_fp8_qkv else 16 + # FP8 publishes a complete row with one x16/x32 STTM. FP16/BF16 + # uses x16 slices to limit Softmax register pressure. This is the + # complete-row Q128/KV128 path; KV256 publishes K32 fragments. + assert cfg.num_packed_p_regs in (16, 32, 64) + regs_per_store = cfg.num_packed_p_regs if cfg.use_fp8_pv else 16 assert cfg.num_packed_p_regs % regs_per_store == 0 for store_idx in cutlass.range_constexpr( cfg.num_packed_p_regs // regs_per_store @@ -777,11 +1009,7 @@ def _compute_p_impl( *, new_max_arr: cutlass.Array, s_arr: cutlass.Array, - route_is_proxy: Int32, - route_origin0: Int32, - route_origin1: Int32, - route_origin2: Int32, - route_origin3: Int32, + route_is_proxy: cutlass.Boolean, ) -> None: """Compute P from S, stage its BMM2 operand, and publish local sums.""" cfg = self.cfg @@ -790,6 +1018,7 @@ def _compute_p_impl( stage_info, new_max_arr=new_max_arr, s_arr=s_arr, + route_is_proxy=route_is_proxy, ) return # ProdWork: transform the softmax S registers into the P operand layout @@ -800,7 +1029,16 @@ def _compute_p_impl( task_cache = _decode_gen_task_cache(stage_info) warp_grp_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] lane_idx = Int32(task_cache[_TASK_CACHE_LANE_IDX]) - if cutlass.const_expr(cfg.tile_size_q == 32 and cfg.use_fp8_qkv): + # Static launches whose KV tiles split evenly between the instances + # exponentiate unmasked tiles; a row that stays at -FLT_MAX there is + # discarded by reduce_sums and the correction gates. + static_even_tiles = ( + not self.use_variable_seqlens_kv + and cfg.total_kv_tiles > 0 + and cfg.total_kv_tiles % cfg.num_insts_kv == 0 + ) + guards_fp8_rows = not (static_even_tiles and cfg.q_tiles_are_full) + if cutlass.const_expr(cfg.tile_size_q == 32 and cfg.use_fp8_pv): # Tile-Q=32 FP8 fast path: compute E4M3 P registers in the # same order consumed by the STSM helper, while also capturing # one local denominator sum per softmax scale group. @@ -812,59 +1050,34 @@ def _compute_p_impl( cfg.num_softmax_scale_groups, space=cutlass.AddressSpace.rmem, ) - if cutlass.const_expr( - not self.use_variable_seqlens_kv - and cfg.total_kv_tiles > 0 - and (cfg.total_kv_tiles % cfg.num_insts_kv) == 0 - and cfg.q_tiles_are_full - ): - # Dense full tiles have no tail/window masking. Use the - # straight-line helper to keep this path compact. - for scale_pair_idx in cutlass.range_constexpr(4): - scale_base = scale_pair_idx * 2 - q32_s_base = scale_pair_idx * 4 - p_lo, p_hi, sum_lo, sum_hi = ( - _compute_fp8_p_regs_and_local_sums_dense( - self.scale_softmax_log2, - new_max_arr[scale_base], - new_max_arr[scale_base + 1], - s_arr[q32_s_base], - s_arr[q32_s_base + 1], - s_arr[q32_s_base + 2], - s_arr[q32_s_base + 3], - s_arr[q32_s_base + 16], - s_arr[q32_s_base + 17], - s_arr[q32_s_base + 18], - s_arr[q32_s_base + 19], - ) - ) - packed_p[scale_pair_idx] = p_lo - packed_p[scale_pair_idx + 4] = p_hi - local_sum[scale_base] = sum_lo - local_sum[scale_base + 1] = sum_hi - else: - # General path handles masked S values from variable - # seqlens, odd tail waves, sliding window, or split-KV. - for scale_pair_idx in cutlass.range_constexpr(4): - scale_base = scale_pair_idx * 2 - q32_s_base = scale_pair_idx * 4 - p_lo, p_hi, sum_lo, sum_hi = _compute_fp8_p_regs_and_local_sums( - self.scale_softmax_log2, + for scale_pair_idx in cutlass.range_constexpr(4): + scale_base = scale_pair_idx * 2 + q32_s_base = scale_pair_idx * 4 + p_lo, p_hi, sum_lo, sum_hi = _compute_fp8_p_regs_and_local_sums( + self.scale_softmax_log2, + self._exponent_addend( new_max_arr[scale_base], + route_is_proxy, + guards_masked_rows=guards_fp8_rows, + ), + self._exponent_addend( new_max_arr[scale_base + 1], - s_arr[q32_s_base], - s_arr[q32_s_base + 1], - s_arr[q32_s_base + 2], - s_arr[q32_s_base + 3], - s_arr[q32_s_base + 16], - s_arr[q32_s_base + 17], - s_arr[q32_s_base + 18], - s_arr[q32_s_base + 19], - ) - packed_p[scale_pair_idx] = p_lo - packed_p[scale_pair_idx + 4] = p_hi - local_sum[scale_base] = sum_lo - local_sum[scale_base + 1] = sum_hi + route_is_proxy, + guards_masked_rows=guards_fp8_rows, + ), + s_arr[q32_s_base], + s_arr[q32_s_base + 1], + s_arr[q32_s_base + 2], + s_arr[q32_s_base + 3], + s_arr[q32_s_base + 16], + s_arr[q32_s_base + 17], + s_arr[q32_s_base + 18], + s_arr[q32_s_base + 19], + ) + packed_p[scale_pair_idx] = p_lo + packed_p[scale_pair_idx + 4] = p_hi + local_sum[scale_base] = sum_lo + local_sum[scale_base + 1] = sum_hi # Publish the local denominator contribution before committing P so # TmemS can update the online-softmax sum after P materialization. for scale_idx in cutlass.range_constexpr(cfg.num_softmax_scale_groups): @@ -899,7 +1112,7 @@ def _compute_p_impl( prims.barrier_cta_sync(4 + self.inst_id, thread_count=128) return - if cutlass.const_expr(cfg.tile_size_q == 16 and cfg.use_fp8_qkv): + if cutlass.const_expr(cfg.tile_size_q == 16 and cfg.use_fp8_pv): # Tile-Q=16 FP8 fast path: each helper call handles two softmax # scale groups and returns the low/high K halves already packed for # STSM. This avoids keeping all 16 FP32 P values live and packing @@ -915,41 +1128,27 @@ def _compute_p_impl( for scale_pair_idx in cutlass.range_constexpr(2): scale_base = scale_pair_idx * 2 s_base = scale_pair_idx * 4 - if cutlass.const_expr( - not self.use_variable_seqlens_kv - and cfg.total_kv_tiles > 0 - and (cfg.total_kv_tiles % cfg.num_insts_kv) == 0 - and cfg.q_tiles_are_full - ): - p_lo, p_hi, sum_lo, sum_hi = ( - _compute_fp8_p_regs_and_local_sums_dense( - self.scale_softmax_log2, - new_max_arr[scale_base], - new_max_arr[scale_base + 1], - s_arr[s_base], - s_arr[s_base + 1], - s_arr[s_base + 2], - s_arr[s_base + 3], - s_arr[s_base + 8], - s_arr[s_base + 9], - s_arr[s_base + 10], - s_arr[s_base + 11], - ) - ) - else: - p_lo, p_hi, sum_lo, sum_hi = _compute_fp8_p_regs_and_local_sums( - self.scale_softmax_log2, + p_lo, p_hi, sum_lo, sum_hi = _compute_fp8_p_regs_and_local_sums( + self.scale_softmax_log2, + self._exponent_addend( new_max_arr[scale_base], + route_is_proxy, + guards_masked_rows=guards_fp8_rows, + ), + self._exponent_addend( new_max_arr[scale_base + 1], - s_arr[s_base], - s_arr[s_base + 1], - s_arr[s_base + 2], - s_arr[s_base + 3], - s_arr[s_base + 8], - s_arr[s_base + 9], - s_arr[s_base + 10], - s_arr[s_base + 11], - ) + route_is_proxy, + guards_masked_rows=guards_fp8_rows, + ), + s_arr[s_base], + s_arr[s_base + 1], + s_arr[s_base + 2], + s_arr[s_base + 3], + s_arr[s_base + 8], + s_arr[s_base + 9], + s_arr[s_base + 10], + s_arr[s_base + 11], + ) packed_p[scale_pair_idx] = p_lo packed_p[scale_pair_idx + 2] = p_hi local_sum[scale_base] = sum_lo @@ -980,25 +1179,16 @@ def _compute_p_impl( local_sums = cutlass.Array( Float32, num_scale_groups, space=cutlass.AddressSpace.rmem ) - proxy_tail_p = cutlass.Array( - Float32, num_scale_groups, space=cutlass.AddressSpace.rmem - ) for idx in cutlass.range_constexpr(num_s_regs): p_vals[idx] = Float32(0.0) for idx in cutlass.range_constexpr(num_scale_groups): local_sums[idx] = Float32(0.0) - proxy_tail_p[idx] = Float32(0.0) for scale_idx in cutlass.range_constexpr(num_scale_groups): # Convert each softmax scale group from S to P. Masked rows have # new_max == -inf and keep their initialized zero P/local_sum. new_max = new_max_arr[scale_idx] - safe_new_max = new_max - if safe_new_max == _neg_max_f32(): - safe_new_max = Float32(0.0) - neg_scaled_max = -self.scale_softmax_log2 * safe_new_max - if cutlass.const_expr(cfg.use_fp8_qkv): - neg_scaled_max += _fp8_log2_quant_scale() + exponent_addend = self._exponent_addend(new_max, route_is_proxy) if new_max != _neg_max_f32(): repeat_idx = scale_idx // 2 pair_idx = scale_idx % 2 @@ -1011,40 +1201,17 @@ def _compute_p_impl( generic_s_base + q_repeats * 4 + (k_pair_idx - 2) * 2 ) p_val = cute.math.exp2( - s_arr[s_idx] * self.scale_softmax_log2 + neg_scaled_max, + s_arr[s_idx] * self.scale_softmax_log2 + exponent_addend, fastmath=True, ) p_vals[s_idx] = p_val local_sums[scale_idx] += p_val - if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): - num_summaries, _ = _block_sparse_proxy_summary_geometry( - cfg.static_seq_len_kv, - cfg.kv_block_size, - ) - atom_origin, logical_summary = _swaps_routed_coordinate( - cfg, - lane_idx >> Int32(2), - route_origin0, - route_origin1, - route_origin2, - route_origin3, - token_group_idx=k_pair_idx, - ) - if atom_origin >= Int32(0) and logical_summary == Int32( - num_summaries - 1 - ): - proxy_tail_p[scale_idx] = p_val # Hand off denominator contributions through TmemS. P remains a pure # MMA operand in SMEM; sums are not reloaded from the P tile. for scale_idx in cutlass.range_constexpr(num_scale_groups): - local_sums[scale_idx] = self._apply_proxy_route_denominator_mass( - local_sums[scale_idx], - proxy_tail_p[scale_idx], - route_is_proxy, - ) self.tmem_s_ref.store_p_local_sum(scale_idx, local_sums[scale_idx]) - if cutlass.const_expr(cfg.use_fp8_qkv): + if cutlass.const_expr(cfg.use_fp8_pv): # FP8 P is packed four values per register and stored with # transposed 8-bit helpers so BMM2 sees the tcgen05 layout. packed_p = cutlass.Array( @@ -1081,7 +1248,7 @@ def _compute_p_impl( ) for pair_idx in cutlass.range_constexpr(cfg.num_packed_p_regs): val_base = pair_idx * 2 - if cutlass.const_expr(cfg.use_bf16_qkv): + if cutlass.const_expr(cfg.use_bf16_pv): regs_p[pair_idx] = _pack_float2_to_bf16( p_vals[val_base], p_vals[val_base + 1] ) @@ -1107,7 +1274,7 @@ def _compute_p_impl( cute.arch.fence_view_async_shared() return - if cutlass.const_expr(cfg.use_fp8_qkv): + if cutlass.const_expr(cfg.use_fp8_pv): # Tile-Q=8 FP8 path: compute packed P and local sums directly # from the eight S registers owned by this lane. packed_p = cutlass.Array( @@ -1118,6 +1285,29 @@ def _compute_p_impl( cfg.num_softmax_scale_groups, space=cutlass.AddressSpace.rmem, ) + packed_p[0], packed_p[1], local_sum[0], local_sum[1] = ( + _compute_fp8_p_regs_and_local_sums( + self.scale_softmax_log2, + self._exponent_addend( + new_max_arr[0], + route_is_proxy, + guards_masked_rows=not static_even_tiles, + ), + self._exponent_addend( + new_max_arr[1], + route_is_proxy, + guards_masked_rows=not static_even_tiles, + ), + s_arr[0], + s_arr[1], + s_arr[2], + s_arr[3], + s_arr[4], + s_arr[5], + s_arr[6], + s_arr[7], + ) + ) if cutlass.const_expr( not self.use_variable_seqlens_kv and not cfg.uses_runtime_q_kv_union @@ -1133,62 +1323,6 @@ def _compute_p_impl( packed_p[1] = Int32(0) local_sum[0] = Float32(0.0) local_sum[1] = Float32(0.0) - else: - packed_p[0], packed_p[1], local_sum[0], local_sum[1] = ( - _compute_fp8_p_regs_and_local_sums( - self.scale_softmax_log2, - new_max_arr[0], - new_max_arr[1], - s_arr[0], - s_arr[1], - s_arr[2], - s_arr[3], - s_arr[4], - s_arr[5], - s_arr[6], - s_arr[7], - ) - ) - elif cutlass.const_expr( - not self.use_variable_seqlens_kv - and cfg.total_kv_tiles > 0 - and (cfg.total_kv_tiles % cfg.num_insts_kv) == 0 - ): - # Dense full-tile FP8 can produce packed P and sums in one - # straight-line helper because no S entry is masked. - packed_p[0], packed_p[1], local_sum[0], local_sum[1] = ( - _compute_fp8_p_regs_and_local_sums_dense( - self.scale_softmax_log2, - new_max_arr[0], - new_max_arr[1], - s_arr[0], - s_arr[1], - s_arr[2], - s_arr[3], - s_arr[4], - s_arr[5], - s_arr[6], - s_arr[7], - ) - ) - else: - # General FP8 path preserves masks from softmax by letting the - # helper suppress entries whose S value is -inf. - packed_p[0], packed_p[1], local_sum[0], local_sum[1] = ( - _compute_fp8_p_regs_and_local_sums( - self.scale_softmax_log2, - new_max_arr[0], - new_max_arr[1], - s_arr[0], - s_arr[1], - s_arr[2], - s_arr[3], - s_arr[4], - s_arr[5], - s_arr[6], - s_arr[7], - ) - ) # Publish sums through TmemS; packed E4M3 P bytes use the # transposed SMEM tile consumed by BMM2. for scale_idx in cutlass.range_constexpr(cfg.num_softmax_scale_groups): @@ -1214,6 +1348,16 @@ def _compute_p_impl( local_sum[scale_idx] = Float32(0.0) for p_idx in cutlass.range_constexpr(8): p_vals[p_idx] = Float32(0.0) + exponent_addends = cutlass.Array( + Float32, + cfg.num_softmax_scale_groups, + space=cutlass.AddressSpace.rmem, + ) + # Both tile-Q=8 16-bit paths below discard or skip -FLT_MAX rows. + for scale_idx in cutlass.range_constexpr(cfg.num_softmax_scale_groups): + exponent_addends[scale_idx] = self._exponent_addend( + new_max_arr[scale_idx], route_is_proxy, guards_masked_rows=False + ) if cutlass.const_expr( not self.use_variable_seqlens_kv and not cfg.use_split_kv @@ -1225,10 +1369,10 @@ def _compute_p_impl( # intentionally safe: reduce_sums' guarded rescale and the # correction path's uses_instN gate both key on that sentinel # and discard the instance before its P/O contribution is visible. - p_result = _compute_p_values_and_local_sums_dense( + p_result = _compute_p_values_and_local_sums( self.scale_softmax_log2, - new_max_arr[0], - new_max_arr[1], + exponent_addends[0], + exponent_addends[1], s_arr[0], s_arr[1], s_arr[2], @@ -1246,19 +1390,17 @@ def _compute_p_impl( # General tile-Q=8 path preserves masked S entries as zero P # contribution by skipping groups whose new max stayed -inf. for scale_idx in cutlass.range_constexpr(cfg.num_softmax_scale_groups): - new_max = new_max_arr[scale_idx] - if new_max != _neg_max_f32(): + exponent_addend = exponent_addends[scale_idx] + if new_max_arr[scale_idx] != _neg_max_f32(): for pair_idx in cutlass.range_constexpr(2): p_base = scale_idx + pair_idx * 2 - scaled_pair = fmul2( + scaled_pair = ffma2( + (s_arr[p_base], s_arr[p_base + 4]), ( self.scale_softmax_log2, self.scale_softmax_log2, ), - fadd2( - (s_arr[p_base], s_arr[p_base + 4]), - (-new_max, -new_max), - ), + (exponent_addend, exponent_addend), ) p_pair = ( cute.math.exp2(scaled_pair[0], fastmath=True), @@ -1268,38 +1410,11 @@ def _compute_p_impl( p_vals[p_base + 4] = p_pair[1] local_sum[scale_idx] += p_pair[0] local_sum[scale_idx] += p_pair[1] - if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): - num_summaries, _ = _block_sparse_proxy_summary_geometry( - cfg.static_seq_len_kv, - cfg.kv_block_size, - ) - for scale_idx in cutlass.range_constexpr(cfg.num_softmax_scale_groups): - proxy_tail_p = Float32(0.0) - for token_group_idx in cutlass.range_constexpr(4): - atom_origin, logical_summary = _swaps_routed_coordinate( - cfg, - lane_idx >> Int32(2), - route_origin0, - route_origin1, - route_origin2, - route_origin3, - token_group_idx=token_group_idx, - ) - if atom_origin >= Int32(0) and logical_summary == Int32( - num_summaries - 1 - ): - p_idx = scale_idx + token_group_idx * 2 - proxy_tail_p = Float32(p_vals[p_idx]) - local_sum[scale_idx] = self._apply_proxy_route_denominator_mass( - local_sum[scale_idx], - proxy_tail_p, - route_is_proxy, - ) # Pack the P scalars to match the dtype consumed by BMM2. regs_p = cutlass.Array( Int32, cfg.num_packed_p_regs, space=cutlass.AddressSpace.rmem ) - if cutlass.const_expr(cfg.use_bf16_qkv): + if cutlass.const_expr(cfg.use_bf16_pv): for reg_idx in cutlass.range_constexpr(cfg.num_packed_p_regs): val_base = reg_idx * 2 regs_p[reg_idx] = _pack_float2_to_bf16( @@ -1339,7 +1454,7 @@ def _compute_p_impl( shape=prims.StoreShape.M8N8, ) cute.arch.fence_view_async_shared() - if cutlass.const_expr(cfg.use_fp8_qkv): + if cutlass.const_expr(cfg.use_fp8_pv): # FP8 P uses inline STSM stores. Synchronize the producer # warpgroup before the UMMA-consumer pipeline is committed so # BMM2 cannot observe a partially written P tile. @@ -1360,11 +1475,7 @@ def compute_p( stage_info, new_max_arr=new_max_arr, s_arr=s_arr, - route_is_proxy=Int32(0), - route_origin0=Int32(0), - route_origin1=Int32(0), - route_origin2=Int32(0), - route_origin3=Int32(0), + route_is_proxy=cutlass.Boolean(False), ) @producer_work @@ -1375,45 +1486,26 @@ def compute_proxy_route_p( *, new_max_arr: cutlass.Array, s_arr: cutlass.Array, - route_origin0: Int32, - route_origin1: Int32, - keeps_route_flags_or_swaps_origin2: Int32, - swaps_route_origin3_bits: Uint32, + keeps_route_flags: Int32, swaps_route_flags: Uint32, ) -> None: - """Normalize the active Keeps/SWAP metadata view and compute P. + """Read the route kind from the active Keeps/SWAP metadata view and compute P. - The shared Int32 input is Keeps route flags or SWAP origin2. SWAP's - origin3 and flags stay bit-preserving Uint32 values until this work - boundary because schedule-level dataflow tokens cannot be cast. + Keeps carries the kind in its Int32 route flags, SWAP in its + bit-preserving Uint32 flags word. The kind selects only the exponent + addend; the max pass has already applied the proxy mass to the scores. """ assert self.cfg.use_block_sparse_proxy_routes - route_origin2 = Int32(0) - route_origin3 = Int32(0) if cutlass.const_expr(self.cfg.use_keeps_mma_ab): - route_is_proxy = Int32( - ( - keeps_route_flags_or_swaps_origin2 - & Int32(_SOFTMAX_ROUTE_IS_PROXY_FLAG) - ) - != Int32(0) - ) + route_is_proxy = _route_is_proxy(keeps_route_flags) else: - route_is_proxy = Int32( - (swaps_route_flags & Uint32(_SOFTMAX_ROUTE_IS_PROXY_FLAG)) != Uint32(0) - ) - route_origin2 = keeps_route_flags_or_swaps_origin2 - route_origin3 = swaps_route_origin3_bits.bitcast(Int32) + route_is_proxy = _route_is_proxy(swaps_route_flags.bitcast(Int32)) self._compute_p_impl( stage_info, new_max_arr=new_max_arr, s_arr=s_arr, route_is_proxy=route_is_proxy, - route_origin0=route_origin0, - route_origin1=route_origin1, - route_origin2=route_origin2, - route_origin3=route_origin3, ) @consumer_work( @@ -1447,7 +1539,7 @@ def p_operands( if cutlass.const_expr(cfg.streams_tmem_p_fragments): # A streamed profile's four pipeline stages are K32 fragments of one P # operand, not four independent full S/P stages. - p_stage_cols = cfg.softmax_score_fragment_regs // 2 + p_stage_cols = cfg.fragment_p_packed_cols p_tmem_addr = self._tmem_base_addr + Int32( self._tmem_alloc.offset + stage_info.stage_idx * p_stage_cols ) @@ -1489,8 +1581,7 @@ def wait_p_fragment( ) prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) - fragment_cols = cfg.softmax_score_fragment_regs // 2 p_tmem_addr = self._tmem_base_addr + Int32( - self._tmem_alloc.offset + fragment_idx * fragment_cols + self._tmem_alloc.offset + fragment_idx * cfg.fragment_p_packed_cols ) return p_tmem_addr 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 173a13cb5447..133589603fd0 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 @@ -24,7 +24,7 @@ import cutlass import cutlass.cute as cute -from cutlass import Int32 +from cutlass import Int32, Uint32, Uint8 from cutlass.experimental import primitives as prims from cutlass.experimental.task_scheduling.enums import WorkAttr @@ -45,12 +45,17 @@ _prepared_kv_routes_are_block_aligned, ) from ..fmha_decode_config import FmhaDecodeConfig +from ..direct_sparse_metadata import DirectSparseMetadataView from ..fmha_decode_constants import ( KV_INST0, KV_INST1, KV_KIND_K, KV_KIND_V, KV_TILE_256_K_SLOT_FOR_SEMANTIC_ATOM, + Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIP_BITS, + Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIP_MASK, + Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIPS_PER_WORD, + SWIZZLE_128B_ROW_BYTES, ) from ...stage import FmhaStage from ...tensor_map import transform_ragged_coords @@ -61,11 +66,13 @@ from .helpers_common import ( _TASK_CACHE_KV_PAGE_IDX_UB, _TASK_CACHE_KV_RAW_TILE_BASE, + _TASK_CACHE_SEQ_LEN_KV, Constexpr, DecodeGenResourceBase, ResourceVars, _clamp_valid_tile_idx, _decode_gen_task_cache, + _kv_smem_swizzle, _logical_head_batch, _logical_q_group_idx, _major_k_stride_bytes, @@ -77,6 +84,7 @@ _num_skipped_kv_tiles, _runtime_clamp_valid_tile_idx, _runtime_last_valid_page_idx, + _runtime_local_kv_tiles, _runtime_split_kv_global_tile_idx, _static_split_kv_global_tile_idx, ) @@ -138,6 +146,70 @@ def _cp_async_bulk_tensor_4d_shared_cta_global_predicated( ) +def _kv_dtype_bytes_for_kind( + cfg: Constexpr[FmhaDecodeConfig], kv_kind: Constexpr[int] +) -> int: + """Return K's or V's byte width, selected by kv_kind.""" + return cfg.v_dtype_bytes if kv_kind == KV_KIND_V else cfg.k_dtype_bytes + + +@cute.jit +def _decode_native_page_locator( + cfg: Constexpr[FmhaDecodeConfig], + page_locator: Int32, +) -> tuple[Int32, Int32]: + """Return ``(token_offset, physical_page)`` for one native table entry.""" + token_offset = Int32(0) + physical_page = page_locator + if cutlass.const_expr(cfg.has_storage_subpages): + subpages_per_storage_page = ( + cfg.effective_storage_tokens_per_page // cfg.num_tokens_per_page + ) + if cutlass.const_expr( + subpages_per_storage_page > 0 + and subpages_per_storage_page & (subpages_per_storage_page - 1) == 0 + ): + # Supported physical storage pages are powers of two. Decode their + # locator with one arithmetic shift and mask. This + # is also branch-free for the inert/padded -1 locator: arithmetic + # right shift preserves physical_page=-1 while the masked token + # offset remains in bounds, so TensorMap OOB zero-fill still + # supplies the invalid fragment. + locator_shift = subpages_per_storage_page.bit_length() - 1 + physical_page = page_locator >> Int32(locator_shift) + subpage = page_locator & Int32(subpages_per_storage_page - 1) + token_offset = subpage * Int32(cfg.num_tokens_per_page) + else: + physical_page = Int32(-1) + if page_locator >= Int32(0): + subpages_per_storage_page_i32 = Int32(subpages_per_storage_page) + physical_page = page_locator // subpages_per_storage_page_i32 + subpage = page_locator - physical_page * subpages_per_storage_page_i32 + token_offset = subpage * Int32(cfg.num_tokens_per_page) + return token_offset, physical_page + + +@cute.jit +def _load_task_warp_rank(cfg: Constexpr[FmhaDecodeConfig]) -> Int32: + """Return this warp's zero-based rank inside the contiguous load task.""" + load_warp_base = ( + cfg.clc_load_warp_idx + if cutlass.const_expr(cfg.use_persistent_scheduler) + else cfg.load_warp_idx + ) + warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx()) + return Int32(warp_idx - Int32(load_warp_base)) + + +@cute.jit +def _is_load_task_warp_leader(cfg: Constexpr[FmhaDecodeConfig]): + """Elect one lane from the first warp of a possibly multi-warp task.""" + elected = prims.elect_sync() + if cutlass.const_expr(cfg.load_num_warps == 1): + return elected + return elected and _load_task_warp_rank(cfg) == Int32(0) + + @cute.jit def _local_kv_tile_idx_for_section( cfg: Constexpr[FmhaDecodeConfig], @@ -166,6 +238,55 @@ def _local_kv_tile_idx_for_section( return stage_info.loop_end * num_insts_kv + Int32(inst_id) +@cute.jit +def _issue_sparse_page_copies( + cfg: Constexpr[FmhaDecodeConfig], + page_offsets: Constexpr, + stage_base, + tma_desc, + barrier, + tile_idx, + local_tile_idx, + kv_head, + head_dim_stage: Constexpr[int], + head_dim_stage_offset: Constexpr[int], +): + """Partition (D64 chunk, physical fragment) copies across loader lanes. + + The cache writer zero-fills page padding: full-fragment tail loads rely + on this because score masking alone does not sanitize NaN/Inf in V. + """ + chunk_hd = min(head_dim_stage, 64) + chunks = head_dim_stage // chunk_hd + fragments = cfg.tile_size_kv // cfg.num_tokens_per_page + copies = fragments * chunks + copies_per_warp = (copies + cfg.load_num_warps - 1) // cfg.load_num_warps + lane = cute.arch.thread_idx()[0] & Int32(31) + warp = _load_task_warp_rank(cfg) + for iteration in cutlass.range_constexpr((copies_per_warp + 31) // 32): + local_copy = lane + Int32(iteration * 32) + copy = warp * Int32(copies_per_warp) + local_copy + if local_copy < Int32(copies_per_warp) and copy < Int32(copies): + chunk = copy // Int32(fragments) + fragment = copy % Int32(fragments) + locator = page_offsets.page_id(tile_idx, local_tile_idx, fragment) + token_offset, physical_page = _decode_native_page_locator(cfg, locator) + smem_offset = chunk * Int32(chunk_hd * cfg.tile_size_kv) + fragment * Int32( + chunk_hd * cfg.num_tokens_per_page + ) + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.subview(smem_offset), + tma_desc, + ( + Int32(head_dim_stage_offset) + chunk * Int32(chunk_hd), + token_offset, + kv_head, + physical_page, + ), + barrier, + ) + + @dataclass(kw_only=True) class SmemQResource(DecodeGenResourceBase): """Q tile in SMEM. Loaded once in HEAD.""" @@ -243,7 +364,7 @@ def _create_initial_task_locals( ) leading_byte_offset = q_leading_bytes stride_byte_offset = Int32(1024) - if cutlass.const_expr(self.cfg.use_fp8_qkv): + if cutlass.const_expr(self.cfg.q_dtype_bytes == 1): q_head_dim_stage = self.cfg.head_dim_kv_stage q_tile_bytes = Int32( self.cfg.tile_size_q * q_head_dim_stage * self.cfg.q_dtype_bytes @@ -390,8 +511,8 @@ def tma_load(self, stage_info: StageInfo) -> None: # barrier attached to stage_info protects this stage until QK consumes it. stage_elems = cfg.smem_q_tile_elements stage_base = self._smem_base_q.subview(stage_info.stage_idx * stage_elems) - if cutlass.const_expr(cfg.use_fp8_qkv): - if prims.elect_sync(): + if cutlass.const_expr(cfg.q_dtype_bytes == 1): + if _is_load_task_warp_leader(cfg): if cutlass.const_expr(cfg.num_head_dim_stages_kv == 1): # FP8 with one head-dim stage is one tensor copy into the # complete staged Q tile. @@ -434,23 +555,40 @@ def tma_load(self, stage_info: StageInfo) -> None: num_chunks = cfg.headdim // chunk_hd chunk_elems = chunk_hd * cfg.tile_size_q if prims.elect_sync(): - # One elected lane issues all TMA chunks for this Q tile; the - # async pipeline tracks completion through stage_info.barrier. - for chunk_idx in cutlass.range_constexpr(num_chunks): - head_dim_offset = chunk_idx * chunk_hd - smem_offset = chunk_idx * chunk_elems - prims.cp_async_bulk_tensor_shared_cta_global( - stage_base.subview(smem_offset), - self.tma_desc_q, - self._tma_coords( - logical_q_group_idx, - logical_h_k_idx, - logical_b_idx, - Int32(head_dim_offset), - ), - stage_info.barrier, - ) - self._complete_grouped_q_padding(stage_info) + # Distribute D64 chunks across as many load-task warps as Q can + # use. Any additional K/V loader warps remain idle because + # each Q chunk must be issued exactly once. + q_chunk_begin = 0 + 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) + 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) + q_chunk_count = num_chunks // q_issuer_warps + q_chunk_begin = load_warp_rank * Int32(q_chunk_count) + issues_q_chunk = load_warp_rank < Int32(q_issuer_warps) + for local_chunk_idx in cutlass.range_constexpr(q_chunk_count): + if issues_q_chunk: + chunk_idx = Int32(q_chunk_begin + local_chunk_idx) + head_dim_offset = chunk_idx * chunk_hd + smem_offset = chunk_idx * chunk_elems + prims.cp_async_bulk_tensor_shared_cta_global( + stage_base.subview(smem_offset), + self.tma_desc_q, + self._tma_coords( + logical_q_group_idx, + logical_h_k_idx, + logical_b_idx, + Int32(head_dim_offset), + ), + stage_info.barrier, + ) + if cutlass.const_expr(cfg.load_num_warps == 1) or _load_task_warp_rank( + cfg + ) == Int32(0): + self._complete_grouped_q_padding(stage_info) @consumer_work(returns=q_desc_slot) @cute.jit @@ -476,7 +614,12 @@ def current_consumer_q_desc(self) -> prims.Tcgen05SmemDesc: @dataclass(kw_only=True) class SmemKvTileResource(DecodeGenResourceBase): - """Single K or V tile in SMEM for the split K/V decode schedule.""" + """Single raw K or V tile in SMEM for the split K/V decode schedule. + + When Q's dtype is different from KV's dtype, `SmemTransformedKvResource` + will be added to transform raw KV and stored in SMEM. Otherwise, + `SmemKvTileResource` will be used directly for MMA. + """ _task_local_specs: ClassVar[tuple[tuple, ...]] = ( ( @@ -495,6 +638,8 @@ class SmemKvTileResource(DecodeGenResourceBase): cfg: Constexpr[FmhaDecodeConfig] = None tma_desc_k: cutlass.Pointer | None = None tma_desc_v: cutlass.Pointer | None = None + tma_desc_k_sf: cutlass.Pointer | None = None + tma_desc_v_sf: cutlass.Pointer | None = None tma_desc_k_atom: cutlass.Pointer | None = None tma_desc_v_atom: cutlass.Pointer | None = None tma_desc_k_summary: cutlass.Pointer | None = None @@ -512,7 +657,9 @@ class SmemKvTileResource(DecodeGenResourceBase): inst_id: Constexpr[int] = 0 kv_kind: Constexpr[int] = KV_KIND_K _alloc: Constexpr[SmemAllocation | None] = None + _sf_alloc: Constexpr[SmemAllocation | None] = None _smem_base_kv: cutlass.Array = None + _smem_base_kv_sf: cutlass.Array = None _desc_base: prims.Tcgen05SmemDesc = None kv_desc_slot: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() v_desc_slot: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() @@ -522,10 +669,25 @@ def _init_placeholder_state(self) -> None: num_stages = ( self.pipeline_config.num_stages if self.pipeline_config is not None else 1 ) + smem_dtype = ( + Uint8 + if self.cfg.use_nvfp4_kv + else (self.cfg.v_dtype if self.kv_kind == KV_KIND_V else self.cfg.k_dtype) + ) self._smem_base_kv = _placeholder_smem_array( - self.cfg.kv_dtype, - self.cfg.smem_kv_tile_elements * num_stages, + smem_dtype, + ( + self.cfg.smem_v_tile_elements + if self.kv_kind == KV_KIND_V + else self.cfg.smem_k_tile_elements + ) + * num_stages, ) + if self.cfg.use_nvfp4_kv: + self._smem_base_kv_sf = _placeholder_smem_array( + Uint8, + self.cfg.smem_kv_sf_tile_bytes * num_stages, + ) self._desc_base = prims.Tcgen05SmemDesc(0) def get_smem_requirements(self) -> list[SmemAllocation]: @@ -534,12 +696,26 @@ def get_smem_requirements(self) -> list[SmemAllocation]: self.pipeline_config.num_stages if self.pipeline_config is not None else 1 ) if self._alloc is None: + tile_bytes = ( + self.cfg.smem_v_storage_tile_bytes + if self.kv_kind == KV_KIND_V + else self.cfg.smem_k_storage_tile_bytes + ) self._alloc = SmemAllocation( name=f"{self.name}", - size_bytes=self.cfg.smem_kv_tile_bytes * num_stages, + size_bytes=tile_bytes * num_stages, alignment=self.cfg.stensor_align, ) - return [self._alloc] + allocs = [self._alloc] + if self.cfg.use_nvfp4_kv: + if self._sf_alloc is None: + self._sf_alloc = SmemAllocation( + name=f"{self.name}Sf", + size_bytes=self.cfg.smem_kv_sf_tile_bytes * num_stages, + alignment=self.cfg.stensor_align, + ) + allocs.append(self._sf_alloc) + return allocs def get_tmem_requirements(self) -> list[TmemAllocation]: """Split K/V staging uses SMEM only.""" @@ -556,40 +732,62 @@ def _create_initial_task_locals( if self.pipeline_config is not None else 1 ) + inst_dtype_bytes = _kv_dtype_bytes_for_kind(self.cfg, self.kv_kind) self._smem_base_kv = cutlass.Array( context.smem_base.data_ptr() + self._alloc.offset, - dtype=self.cfg.kv_dtype, - shape=(self.cfg.smem_kv_tile_elements * num_stages,), + dtype=Uint8 + if self.cfg.use_nvfp4_kv + else ( + self.cfg.v_dtype if self.kv_kind == KV_KIND_V else self.cfg.k_dtype + ), + shape=( + ( + self.cfg.smem_v_tile_elements + if self.kv_kind == KV_KIND_V + else self.cfg.smem_k_tile_elements + ) + * num_stages, + ), addrspace=3, ) + if cutlass.const_expr(self.cfg.use_nvfp4_kv): + self._smem_base_kv_sf = cutlass.Array( + context.smem_base.data_ptr() + self._sf_alloc.offset, + dtype=Uint8, + shape=(self.cfg.smem_kv_sf_tile_bytes * num_stages,), + addrspace=3, + ) + # NVFP4 KV is never consumed by MMA. The supported MMA operand types are + # FP16, BF16, and FP8. + self._desc_base = prims.Tcgen05SmemDesc(0) + return {"kv_desc": cutlass.Int64(0), "v_desc": cutlass.Int64(0)} kv_tile_bytes = Int32( - self.cfg.tile_size_kv - * self.cfg.head_dim_kv_stage - * self.cfg.kv_dtype_bytes + self.cfg.tile_size_kv * self.cfg.head_dim_kv_stage * inst_dtype_bytes ) leading_byte_offset = Int32( self.cfg.tile_size_kv * min(self.cfg.head_dim_kv_stage, 64) - * self.cfg.kv_dtype_bytes + * inst_dtype_bytes ) stride_byte_offset = Int32(1024) - if cutlass.const_expr(self.cfg.use_fp8_qkv): + if cutlass.const_expr(inst_dtype_bytes == 1): leading_byte_offset = kv_tile_bytes stride_byte_offset = Int32( - _major_k_stride_bytes( - self.cfg.kv_dtype_bytes, self.cfg.head_dim_kv_stage - ) + _major_k_stride_bytes(inst_dtype_bytes, self.cfg.head_dim_kv_stage) ) if cutlass.const_expr( self.kv_kind == KV_KIND_V - and (self.cfg.use_fp8_qkv or self.cfg.headdim == 64) + and (inst_dtype_bytes == 1 or self.cfg.headdim == 64) ): leading_byte_offset = Int32(0) + inst_swizzle = prims.Tcgen05SmemSwizzle.SWIZZLE_128B + if cutlass.const_expr(inst_dtype_bytes == 1 and self.cfg.headdim == 64): + inst_swizzle = prims.Tcgen05SmemSwizzle.SWIZZLE_64B self._desc_base = prims.Tcgen05SmemDesc.build( self._smem_base_kv, leading_byte_offset=leading_byte_offset, stride_byte_offset=stride_byte_offset, - layout=_qkv_smem_swizzle(self.cfg), + layout=inst_swizzle, ) return {"kv_desc": cutlass.Int64(0), "v_desc": cutlass.Int64(0)} @@ -620,16 +818,29 @@ def init_descriptor_state(self, stage_info: StageInfo) -> None: @cute.jit def _stage_base(self, stage_info: StageInfo) -> cutlass.Array: """Return the SMEM base for the current split K/V pipeline stage.""" - stage_elems = self.cfg.smem_kv_tile_bytes // self.cfg.kv_dtype_bytes + stage_elems = ( + self.cfg.smem_v_tile_elements + if self.kv_kind == KV_KIND_V + else self.cfg.smem_k_tile_elements + ) return self._smem_base_kv.subview(stage_info.stage_idx * stage_elems) + @cute.jit + def _sf_stage_base(self, stage_idx: Int32) -> cutlass.Array: + """Return the NVFP4 scale-factor SMEM slice for one pipeline stage.""" + sf_stage_elems = Int32(self.cfg.smem_kv_sf_tile_bytes) + return self._smem_base_kv_sf.subview(stage_idx * sf_stage_elems) + @cute.jit def _local_tile_idx( - self, stage_info: StageInfo, section: Constexpr[FmhaStage] + self, + stage_info: StageInfo, + inst_id: Constexpr[int], + section: Constexpr[FmhaStage], ) -> Int32: """Map the schedule phase to the local K or V tile index.""" return _local_kv_tile_idx_for_section( - self.cfg, stage_info, self.inst_id, self.kv_kind, section + self.cfg, stage_info, inst_id, self.kv_kind, section ) @cute.jit @@ -686,6 +897,7 @@ def _maybe_runtime_tile_idx(self, stage_info: StageInfo, tile_idx: Int32) -> Int def _producer_load( self, stage_info: StageInfo, + inst_id: Constexpr[int], section: Constexpr[FmhaStage], head_dim_stage_idx: Constexpr[int], ) -> None: @@ -693,7 +905,7 @@ def _producer_load( cfg = self.cfg # Resolve the schedule-local K/V tile, logical head/batch coordinates, # and descriptor kind before selecting paged or dense addressing. - local_tile_idx = self._local_tile_idx(stage_info, section) + local_tile_idx = self._local_tile_idx(stage_info, inst_id, section) logical_h_k_idx, logical_b_idx = _logical_head_batch( stage_info, self.h_k_idx, self.b_idx ) @@ -728,7 +940,8 @@ def _producer_load( kv_atom_size = _block_sparse_kv_atom_size(cfg.kv_block_size) head_dim_stage = cfg.head_dim_kv_stage head_dim_stage_offset = head_dim_stage_idx * head_dim_stage - chunk_hd = min(head_dim_stage, 64) + # One chunk is one swizzled 128-byte row, the TensorMap inner box. + chunk_hd = min(head_dim_stage, cfg.kv_swizzle_row_elems) num_chunks = head_dim_stage // chunk_hd tile_chunk_elems = chunk_hd * cfg.tile_size_kv if cutlass.const_expr(cfg.use_paged_kv): @@ -746,12 +959,20 @@ def _producer_load( tile_size_kv=cfg.tile_size_kv, kv_atom_size=kv_atom_size, head_dim_stage=head_dim_stage, - kv_dtype_bytes=cfg.kv_dtype_bytes, + kv_dtype_bytes=_kv_dtype_bytes_for_kind(cfg, self.kv_kind), + ) + assert total_transaction_bytes == ( + cfg.smem_v_tile_bytes + if self.kv_kind == KV_KIND_V + else cfg.smem_k_tile_bytes ) - assert total_transaction_bytes == cfg.smem_kv_tile_bytes num_chunks = transactions_per_load // atoms_per_route - chunk_hd = transaction_bytes // (kv_atom_size * cfg.kv_dtype_bytes) - atom_chunk_elems = transaction_bytes // cfg.kv_dtype_bytes + chunk_hd = transaction_bytes // ( + kv_atom_size * _kv_dtype_bytes_for_kind(cfg, self.kv_kind) + ) + atom_chunk_elems = transaction_bytes // _kv_dtype_bytes_for_kind( + cfg, self.kv_kind + ) tile_chunk_elems = chunk_hd * cfg.tile_size_kv if prims.elect_sync(): stage_base = self._stage_base(stage_info) @@ -956,46 +1177,115 @@ def _producer_load( head_dim_stage_offset = head_dim_stage_idx * head_dim_stage page_fragments = cfg.tile_size_kv // cfg.num_tokens_per_page tile_idx = self._maybe_runtime_tile_idx(stage_info, local_tile_idx) - if cutlass.const_expr(cfg.use_fp8_qkv): - if prims.elect_sync(): - # FP8 pages are copied as one contiguous head-dim stage per - # page fragment. + inst_dtype_bytes = _kv_dtype_bytes_for_kind(self.cfg, self.kv_kind) + if cutlass.const_expr(inst_dtype_bytes == 1): + if prims.elect_sync() and _load_task_warp_rank(cfg) < Int32( + min(page_fragments, cfg.load_num_warps) + ): + # NVFP4/FP8 pages are copied as one contiguous head-dim stage per + # page fragment. Partition the fragments across loader + # warps: elect_sync selects one issuer per warp, not one + # issuer for the complete cooperative group. Letting every + # elected lane replay the full loop over-completes the TMA + # transaction barrier when load_num_warps > 1. stage_base = self._stage_base(stage_info) - page_ids = self.page_offsets_kv.page_ids(tile_idx) - for page_frag in cutlass.range_constexpr(page_fragments): - page_id = Int32(page_ids[page_frag]) + page_ids = self.page_offsets_kv.page_ids(tile_idx, local_tile_idx) + sf_desc = ( + self.tma_desc_v_sf + if cutlass.const_expr(self.kv_kind == KV_KIND_V) + else self.tma_desc_k_sf + ) + if cutlass.const_expr(cfg.use_nvfp4_kv): + sf_stage_base = self._sf_stage_base(stage_info.stage_idx) + sf_page_elems = Int32( + cfg.num_tokens_per_page * cfg.smem_kv_sf_bytes_per_token + ) + active_load_warps = min(page_fragments, cfg.load_num_warps) + transactions_per_warp = page_fragments // active_load_warps + first_page_frag = _load_task_warp_rank(cfg) * Int32( + transactions_per_warp + ) + for warp_transaction in cutlass.range_constexpr( + transactions_per_warp + ): + page_frag = first_page_frag + Int32(warp_transaction) + token_offset, physical_page = _decode_native_page_locator( + cfg, Int32(page_ids[page_frag]) + ) smem_page_offset = Int32( - page_frag * cfg.num_tokens_per_page * head_dim_stage + page_frag + * cfg.num_tokens_per_page + * ( + head_dim_stage // 2 + if cfg.use_nvfp4_kv + and not cfg.store_transformed_kv_in_tmem + else head_dim_stage + ) ) prims.cp_async_bulk_tensor_shared_cta_global( stage_base.subview(smem_page_offset), tma_desc, ( - Int32(head_dim_stage_offset), - Int32(0), + Int32( + head_dim_stage_offset // 2 + if cfg.use_nvfp4_kv + and not cfg.store_transformed_kv_in_tmem + else head_dim_stage_offset + ), + token_offset, logical_h_k_idx, - page_id, + physical_page, ), stage_info.barrier, ) + if cutlass.const_expr(cfg.use_nvfp4_kv): + # SF TensorMaps fold multiple tokens into dim0. + prims.cp_async_bulk_tensor_shared_cta_global( + sf_stage_base.subview(page_frag * sf_page_elems), + sf_desc, + ( + Int32(0), + token_offset // Int32(cfg.sf_tma_reshape_factor), + logical_h_k_idx, + physical_page, + ), + stage_info.barrier, + ) else: - # 16-bit pages are copied in 64-column chunks so the SMEM - # layout matches the K/V tcgen05 descriptor swizzle. - chunk_hd = min(head_dim_stage, 64) - num_chunks = head_dim_stage // chunk_hd + # 128B TMA swizzling permits 64 columns per copy for 16-bit K/V. + # Store each chunk as a separate slab matching the tcgen05 layout. + chunk_hd = min(head_dim_stage, 128 // inst_dtype_bytes) + num_head_dim_tma_chunks = head_dim_stage // chunk_hd tile_chunk_elems = chunk_hd * cfg.tile_size_kv page_chunk_elems = chunk_hd * cfg.num_tokens_per_page + if cutlass.const_expr(cfg.uses_scattered_page_route): + _issue_sparse_page_copies( + cfg, + self.page_offsets_kv, + self._stage_base(stage_info), + tma_desc, + stage_info.barrier, + tile_idx, + local_tile_idx, + logical_h_k_idx, + head_dim_stage, + head_dim_stage_offset, + ) + return + # Bind the Array on every lane before entering staged control flow. + stage_base = self._stage_base(stage_info) if prims.elect_sync(): - stage_base = self._stage_base(stage_info) - page_ids = self.page_offsets_kv.page_ids(tile_idx) - for chunk_idx in cutlass.range_constexpr(num_chunks): + page_ids = self.page_offsets_kv.page_ids(tile_idx, local_tile_idx) + for chunk_idx in cutlass.range_constexpr(num_head_dim_tma_chunks): local_head_dim_offset = chunk_idx * chunk_hd global_head_dim_offset = ( head_dim_stage_offset + local_head_dim_offset ) local_tile_offset = chunk_idx * tile_chunk_elems for page_frag in cutlass.range_constexpr(page_fragments): - page_id = Int32(page_ids[page_frag]) + token_offset, physical_page = _decode_native_page_locator( + cfg, Int32(page_ids[page_frag]) + ) smem_page_offset = Int32( local_tile_offset + page_frag * page_chunk_elems ) @@ -1004,9 +1294,9 @@ def _producer_load( tma_desc, ( Int32(global_head_dim_offset), - Int32(0), + token_offset, logical_h_k_idx, - page_id, + physical_page, ), stage_info.barrier, ) @@ -1017,7 +1307,8 @@ def _producer_load( tile_offset = tile_idx * Int32(cfg.tile_size_kv) head_dim_stage = cfg.head_dim_kv_stage head_dim_stage_offset = head_dim_stage_idx * head_dim_stage - if cutlass.const_expr(cfg.use_fp8_qkv): + inst_dtype_bytes = _kv_dtype_bytes_for_kind(self.cfg, self.kv_kind) + if cutlass.const_expr(inst_dtype_bytes == 1): if prims.elect_sync(): # FP8 dense K/V needs one tensor copy for the active # head-dim stage. @@ -1038,8 +1329,9 @@ def _producer_load( chunk_hd = min(head_dim_stage, 64) num_chunks = head_dim_stage // chunk_hd tile_chunk_elems = chunk_hd * cfg.tile_size_kv + # Keep the Array type stable across the elected-lane control-flow join. + stage_base = self._stage_base(stage_info) if prims.elect_sync(): - stage_base = self._stage_base(stage_info) for chunk_idx in cutlass.range_constexpr(num_chunks): local_head_dim_offset = chunk_idx * chunk_hd global_head_dim_offset = ( @@ -1070,7 +1362,7 @@ def load_k0( """Produce the first split K tile for this schedule phase.""" # ProdWork: K0 uses inst slot 0; the section selects HEAD/LOOP/TAIL # tile numbering and head_dim_stage_idx selects the H256 slice. - self._producer_load(stage_info, section, head_dim_stage_idx) + self._producer_load(stage_info, KV_INST0, section, head_dim_stage_idx) @producer_work @cute.jit @@ -1084,7 +1376,7 @@ def load_k1( """Produce the second split K tile for this schedule phase.""" # ProdWork: K1 uses inst slot 1 but otherwise shares the same staged # K/V TMA path as K0. - self._producer_load(stage_info, section, head_dim_stage_idx) + self._producer_load(stage_info, KV_INST1, section, head_dim_stage_idx) @producer_work @cute.jit @@ -1098,7 +1390,7 @@ def load_v0( """Produce the first split V tile for this schedule phase.""" # ProdWork: V0 publishes the first V descriptor stream consumed by the # corresponding PV MMA call. - self._producer_load(stage_info, section, head_dim_stage_idx) + self._producer_load(stage_info, KV_INST0, section, head_dim_stage_idx) @producer_work @cute.jit @@ -1112,12 +1404,17 @@ def load_v1( """Produce the second split V tile for this schedule phase.""" # ProdWork: V1 publishes the second V descriptor stream consumed by the # corresponding PV MMA call. - self._producer_load(stage_info, section, head_dim_stage_idx) + self._producer_load(stage_info, KV_INST1, section, head_dim_stage_idx) @cute.jit def _build_desc(self, stage_info: StageInfo) -> prims.Tcgen05SmemDesc: """Advance the split K/V base descriptor to the committed stage.""" - stage_offset_bytes = stage_info.stage_idx * Int32(self.cfg.smem_kv_tile_bytes) + tile_bytes = ( + self.cfg.smem_v_storage_tile_bytes + if self.kv_kind == KV_KIND_V + else self.cfg.smem_k_storage_tile_bytes + ) + stage_offset_bytes = stage_info.stage_idx * Int32(tile_bytes) return self._desc_base.advance_start_address(stage_offset_bytes) @consumer_work(returns=kv_desc_slot) @@ -1139,31 +1436,41 @@ def v_desc(self, stage_info: StageInfo) -> prims.Tcgen05SmemDesc: @dataclass(kw_only=True) class SmemPageOffsetsKvResource(DecodeGenResourceBase): - """Paged-KV logical-to-physical page IDs staged in SMEM. - - A dedicated warp prefetches the page table entries for the next K/V tile. - The TMA load warp then reads these SMEM-cached offsets when issuing the - page-sized TMA copies, matching the split producer layout. - - Paired K0/K1/V0/V1 schedules publish one stage per logical tile and store - exactly that tile's page IDs. Shared-offset schedules retain a warp-aligned - 32-ID window so one coalesced load can serve adjacent logical tiles. + """SMEM page locators shared by page-table producers and K/V TMA issuers. + + Held native sparse routes publish the complete CTA-local locator span in + one offset stage. Producer warps stripe logical KV tiles; their count is + independent of the number of KV instances. Consumers address each tile's + slice and reuse it for K, V and all head-dimension stages through the V + tail. K/V data buffers advance independently of this held offset stage. + Native sparse routes use int32 cp.async copies, with in-place mapping for + direct raw IDs before the producer publishes the stage. + + Other paged-KV schedules stage one tile or an aligned 32-ID window. Native + page rows have a fixed dense stride and a sequence-length-bounded prefix. + Membership rows are cached only when the complete SMEM layout fits; larger + rows are read one packed word at a time from immutable GMEM by Softmax. """ cfg: Constexpr[FmhaDecodeConfig] = None stage_page_ids_per_tile: Constexpr[bool] = False - page_idx_kv: cute.Pointer | None = None + cache_memberships_in_smem: Constexpr[bool] = True + page_idx_kv: cute.Pointer | DirectSparseMetadataView | None = None + q_token_kv_block_sparse_page_memberships: cute.Pointer | None = None seqlens_kv: cute.Pointer | None = None use_native_paged_kv: Constexpr[bool] = False - block_tables: cute.Pointer | None = None - block_table_row_stride: cutlass.Int64 = None + page_table_stride: cutlass.Int64 = None + num_heads_kv: Int32 = None + q_token_kv_block_sparse_page_membership_stride: Int32 = None max_seq_len_kv: Int32 = None h_k_idx: Int32 = None b_idx: Int32 = None q_group_idx: Int32 = None seq_len_q: Int32 = None _alloc: Constexpr[SmemAllocation | None] = None + _membership_alloc: Constexpr[SmemAllocation | None] = None _smem_page_offsets: cutlass.Array = None + _smem_q_token_kv_block_sparse_memberships: cutlass.Array = None cached_page_ids: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() def __post_init__(self) -> None: @@ -1183,15 +1490,76 @@ def __post_init__(self) -> None: ) self._init_placeholder_state() + @property + def page_ids_per_stage(self) -> int: + """Return the locator capacity published by one pipeline stage. + + Held sparse routes keep all locators owned by the CTA work tile, not + one tile per producer warp. K and V reuse the same dense-table row + across the complete instruction cadence. + """ + pages_per_tile = self.cfg.tile_size_kv // self.cfg.num_tokens_per_page + if self.holds_encoded_locator_window: + return self.encoded_locator_window_tiles * pages_per_tile + if self.stages_encoded_locators_per_tile: + return pages_per_tile + return pages_per_tile if self.stage_page_ids_per_tile else 32 + + @property + def stages_encoded_locators_per_tile(self) -> bool: + """Whether a long graph-stable QToken-KvBlock-Sparse-Attention route stages one locator tile.""" + + return ( + self.use_native_paged_kv + and self.cfg.uses_scattered_page_route + and not self.cfg.use_sliding_window_causal + and not self.cfg.uses_held_encoded_locator_window + ) + + @property + def holds_encoded_locator_window(self) -> bool: + """Whether one offset stage retains this CTA's complete sparse route.""" + + return self.use_native_paged_kv and self.cfg.uses_held_encoded_locator_window + + @property + def q_token_kv_block_sparse_membership_entries(self) -> int: + """Return packed membership words retained by grouped-Q routes.""" + + if ( + not self.cfg.uses_q_token_kv_block_sparse_page_membership + or not self.cache_memberships_in_smem + ): + return 0 + pages_per_tile = self.cfg.tile_size_kv // self.cfg.num_tokens_per_page + page_entries = self.encoded_locator_window_tiles * pages_per_tile + return ( + page_entries + Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIPS_PER_WORD - 1 + ) // Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIPS_PER_WORD + + @property + def encoded_locator_window_tiles(self) -> int: + """Return the instruction-aligned KV-tile capacity of a held window.""" + # The two-instance schedule executes an even final group. Include its + # inert partner in odd-tail direct launches so lane-partitioned locator + # reads remain in bounds; the producer fills that partner with -1. + local_tiles = self.cfg.static_local_kv_tiles + return ( + (local_tiles + self.cfg.num_insts_kv - 1) + // self.cfg.num_insts_kv + * self.cfg.num_insts_kv + ) + def _init_placeholder_state(self) -> None: """Create placeholder storage for per-stage page-offset windows.""" num_stages = ( self.pipeline_config.num_stages if self.pipeline_config is not None else 1 ) - pages_per_tile = self.cfg.tile_size_kv // self.cfg.num_tokens_per_page - page_ids_per_stage = pages_per_tile if self.stage_page_ids_per_tile else 32 self._smem_page_offsets = _placeholder_smem_array( - Int32, num_stages * page_ids_per_stage + Int32, num_stages * self.page_ids_per_stage + ) + self._smem_q_token_kv_block_sparse_memberships = _placeholder_smem_array( + Uint32, self.q_token_kv_block_sparse_membership_entries ) def get_smem_requirements(self) -> list[SmemAllocation]: @@ -1199,15 +1567,22 @@ def get_smem_requirements(self) -> list[SmemAllocation]: num_stages = ( self.pipeline_config.num_stages if self.pipeline_config is not None else 1 ) - pages_per_tile = self.cfg.tile_size_kv // self.cfg.num_tokens_per_page - page_ids_per_stage = pages_per_tile if self.stage_page_ids_per_tile else 32 if self._alloc is None: self._alloc = SmemAllocation( name=f"{self.name}", - size_bytes=num_stages * page_ids_per_stage * 4, + size_bytes=num_stages * self.page_ids_per_stage * 4, alignment=16, ) - return [self._alloc] + allocations = [self._alloc] + if self.q_token_kv_block_sparse_membership_entries: + if self._membership_alloc is None: + self._membership_alloc = SmemAllocation( + name=f"{self.name}_q_token_kv_block_sparse_membership", + size_bytes=self.q_token_kv_block_sparse_membership_entries * 4, + alignment=16, + ) + allocations.append(self._membership_alloc) + return allocations def get_tmem_requirements(self) -> list[TmemAllocation]: """Page-offset staging uses SMEM only.""" @@ -1224,14 +1599,19 @@ def _create_initial_task_locals( if self.pipeline_config is not None else 1 ) - pages_per_tile = self.cfg.tile_size_kv // self.cfg.num_tokens_per_page - page_ids_per_stage = pages_per_tile if self.stage_page_ids_per_tile else 32 self._smem_page_offsets = cutlass.Array( context.smem_base.data_ptr() + self._alloc.offset, dtype=cutlass.Int32, - shape=(num_stages * page_ids_per_stage,), + shape=(num_stages * self.page_ids_per_stage,), addrspace=3, ) + if cutlass.const_expr(self.q_token_kv_block_sparse_membership_entries): + self._smem_q_token_kv_block_sparse_memberships = cutlass.Array( + context.smem_base.data_ptr() + self._membership_alloc.offset, + dtype=Uint32, + shape=(self.q_token_kv_block_sparse_membership_entries,), + addrspace=3, + ) return {} @producer_work(work_attrs=WorkAttr.AUXILIARY) @@ -1262,23 +1642,30 @@ def init_cached_read_state(self, stage_info: StageInfo) -> cutlass.Array: ) @cute.jit - def page_ids(self, tile_idx: Int32) -> cutlass.Array: + def page_ids(self, tile_idx: Int32, local_tile_idx: Int32) -> cutlass.Array: """Load the tile's page IDs from its staged cache entry. - Single-tile stages begin at offset zero. Multi-tile stages use the - runtime-resolved tile index to select from their aligned 32-ID window. + Held encoded windows use the CTA-local tile index. Single-tile stages + begin at offset zero, while shared multi-tile stages use the runtime + tile index to select from their aligned 32-ID window. """ cfg = self.cfg pages_per_tile = cfg.tile_size_kv // cfg.num_tokens_per_page - if cutlass.const_expr(self.stage_page_ids_per_tile): + if cutlass.const_expr(self.holds_encoded_locator_window): + offset = self.consumer_work_stage * Int32( + self.page_ids_per_stage + ) + local_tile_idx * Int32(pages_per_tile) + elif cutlass.const_expr( + self.stage_page_ids_per_tile or self.stages_encoded_locators_per_tile + ): offset = self.consumer_work_stage * Int32(pages_per_tile) else: group_page_idx = (tile_idx * Int32(pages_per_tile)) & Int32(31) offset = self.consumer_work_stage * Int32(32) + group_page_idx - if cutlass.const_expr(pages_per_tile in (8, 16)): + if cutlass.const_expr(pages_per_tile in (8, 16, 32)): # Native shared-memory vector loads top out at four Int32 values. - # Wide page-16 tiles therefore consume their IDs as independently - # aligned 16-byte loads from the same 32-ID cache window. + # Wide tiles therefore consume their IDs as independently aligned + # 16-byte loads from the same 32-ID cache window. page_ids = cutlass.Array( Int32, pages_per_tile, space=cutlass.AddressSpace.rmem ) @@ -1297,6 +1684,142 @@ def page_ids(self, tile_idx: Int32) -> cutlass.Array: return self._smem_page_offsets.load(offset, vector_size=2, alignment=8) return self._smem_page_offsets.load(offset, vector_size=1, alignment=4) + @cute.jit + def page_id( + self, tile_idx: Int32, local_tile_idx: Int32, page_frag: Int32 + ) -> Int32: + """Load a tile's locator; offset stages are independent of K/V data stages.""" + cfg = self.cfg + pages_per_tile = cfg.tile_size_kv // cfg.num_tokens_per_page + if cutlass.const_expr(self.holds_encoded_locator_window): + offset = self.consumer_work_stage * Int32( + self.page_ids_per_stage + ) + local_tile_idx * Int32(pages_per_tile) + elif cutlass.const_expr( + self.stage_page_ids_per_tile or self.stages_encoded_locators_per_tile + ): + offset = self.consumer_work_stage * Int32(pages_per_tile) + else: + group_page_idx = (tile_idx * Int32(pages_per_tile)) & Int32(31) + offset = self.consumer_work_stage * Int32(32) + group_page_idx + return Int32(self._smem_page_offsets[offset + page_frag]) + + @cute.jit + def _q_token_kv_block_sparse_membership_word( + self, stage_info: StageInfo, membership_idx: Int32 + ) -> tuple[Uint32, Int32]: + """Read a cached word, or fetch only the current tile's word from GMEM. + + Large routes retain metadata in its existing immutable GMEM buffer. + Both paths return the index in their word coordinate system, including + an unaligned split's byte offset. Never expose poisoned padding bytes. + """ + if cutlass.const_expr(self.cache_memberships_in_smem): + self._create_initial_task_locals(stage_info.context) + word = Uint32( + self._smem_q_token_kv_block_sparse_memberships[ + membership_idx + // Int32(Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIPS_PER_WORD) + ] + ) + else: + assert self.q_token_kv_block_sparse_page_memberships is not None + cfg = self.cfg + task_cache = _decode_gen_task_cache(stage_info) + membership_idx += Int32(task_cache[_TASK_CACHE_KV_RAW_TILE_BASE]) * Int32( + cfg.tile_size_kv // cfg.num_tokens_per_page + ) + word_idx = membership_idx // Int32( + Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIPS_PER_WORD + ) + first_page = word_idx * Int32( + Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIPS_PER_WORD + ) + last_page = Int32(task_cache[_TASK_CACHE_KV_PAGE_IDX_UB]) + word = Uint32(0) + if first_page <= last_page: + head, batch = _logical_head_batch(stage_info, self.h_k_idx, self.b_idx) + row = cutlass.Int64(batch) + if cutlass.const_expr( + cfg.use_persistent_scheduler and not cfg.shares_sparse_pattern + ): + row = row * cutlass.Int64(self.num_heads_kv) + cutlass.Int64(head) + word = Uint32( + self.q_token_kv_block_sparse_page_memberships[ + row + * cutlass.Int64( + self.q_token_kv_block_sparse_page_membership_stride + ) + + cutlass.Int64(word_idx) + ] + ) + live_bytes = last_page - first_page + Int32(1) + if live_bytes < Int32( + Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIPS_PER_WORD + ): + word &= ( + Uint32(1) + << ( + live_bytes + * Int32(Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIP_BITS) + ) + ) - Uint32(1) + return word, membership_idx + + @cute.jit + def q_token_kv_block_sparse_page_membership( + self, + stage_info: StageInfo, + local_tile_idx: Int32, + page_frag: Int32, + ) -> Uint32: + """Load one grouped-Q membership byte for a score fragment.""" + assert self.cfg.uses_q_token_kv_block_sparse_page_membership + pages_per_tile = Int32(self.cfg.tile_size_kv // self.cfg.num_tokens_per_page) + membership_idx = local_tile_idx * pages_per_tile + page_frag + membership_word, membership_idx = self._q_token_kv_block_sparse_membership_word( + stage_info, membership_idx + ) + membership_shift = ( + membership_idx % Int32(Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIPS_PER_WORD) + ) * Int32(Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIP_BITS) + return (membership_word >> membership_shift) & Uint32( + Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIP_MASK + ) + + @cute.jit + def q_token_kv_block_sparse_page_memberships4( + self, + stage_info: StageInfo, + local_tile_idx: Int32, + page_frag: Int32, + ) -> cutlass.Array: + """Load four aligned grouped-Q membership bytes from one packed word. + + The ``tmem_s`` Keeps Softmax path calls this helper four times to + assemble each lane-local 16-page membership word. + """ + + assert self.cfg.uses_q_token_kv_block_sparse_page_membership + pages_per_tile = Int32(self.cfg.tile_size_kv // self.cfg.num_tokens_per_page) + membership_idx = local_tile_idx * pages_per_tile + page_frag + memberships = cutlass.Array( + Uint32, + 4, + space=cutlass.AddressSpace.rmem, + ) + membership_word, _ = self._q_token_kv_block_sparse_membership_word( + stage_info, membership_idx + ) + for elem_idx in cutlass.range_constexpr(4): + memberships[elem_idx] = ( + membership_word + >> Int32(Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIP_BITS * elem_idx) + ) & Uint32( + Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIP_MASK, + ) + return memberships + @consumer_work(returns=cached_page_ids) @cute.jit def cache_page_ids( @@ -1319,19 +1842,101 @@ def cache_page_ids( ) pages_per_tile = cfg.tile_size_kv // cfg.num_tokens_per_page + inst_dtype_bytes = _kv_dtype_bytes_for_kind(cfg, kv_kind) # BF16 TMA is issued only by the elected lane, so only that lane needs # the register cache. FP8's predicated helper builds coordinates in # every lane and therefore keeps the existing all-lane semantics. - if cutlass.const_expr(cfg.use_fp8_qkv): - fp8_page_ids = self.page_ids(tile_idx) + if cutlass.const_expr(inst_dtype_bytes == 1): + fp8_page_ids = self.page_ids(tile_idx, local_tile_idx) for page_frag in cutlass.range_constexpr(pages_per_tile): cached_page_ids[page_frag] = Int32(fp8_page_ids[page_frag]) elif prims.elect_sync(): - bf16_page_ids = self.page_ids(tile_idx) + bf16_page_ids = self.page_ids(tile_idx, local_tile_idx) for page_frag in cutlass.range_constexpr(pages_per_tile): cached_page_ids[page_frag] = Int32(bf16_page_ids[page_frag]) return cached_page_ids + @cute.jit + def _prefetch_page_offsets_async( + self, + stage_info, + source, + table_offset, + first, + last, + lane, + warp_rank, + route_state, + ): + """Stage int32 IDs with cp.async and resolve direct inputs in place. + + Copy and mapping use the same lane ownership, so the per-thread + wait suffices before in-place mapping. The existing TS full barrier + then publishes all lanes' final locators to the TMA load task. + Materialized locators need no translation. Direct inputs contain only + selected block IDs; their causal tail locators are synthesized here. + """ + entries = self.page_ids_per_stage + stage_base = stage_info.stage_idx * Int32(entries) + source_last = last + if cutlass.const_expr(isinstance(source, DirectSparseMetadataView)): + blocks = source.values[0] + row, _, visible = route_state + safe_row = cute.math.min( + cute.math.max(row, Int32(0)), Int32(blocks.shape[0]) - Int32(1) + ) + src = blocks.iterator + cutlass.Int64(safe_row) * cutlass.Int64( + blocks.stride[0] + ) + if cutlass.const_expr(len(blocks.shape) == 3): + src = src + cutlass.Int64(source.values[5]) * cutlass.Int64( + blocks.stride[1] + ) + complete = cute.math.min( + visible // Int32(source.sparse_block_size), Int32(blocks.shape[-1]) + ) + # The tail is synthesized after copying; never read its raw slot. + source_last = cute.math.min( + last, complete * Int32(source.fragments_per_block) - Int32(1) + ) + else: + src = source + table_offset + stride = self.cfg.page_offsets_num_warps * 32 + owner = warp_rank * Int32(32) + lane + for group in cutlass.range_constexpr((entries + stride - 1) // stride): + offset = owner + Int32(group * stride) + if offset < Int32(entries): + slot = first + offset + source_slot = Int32(0) + copy_bytes = Int32(0) + if slot <= source_last: + source_slot = slot + if cutlass.const_expr(isinstance(source, DirectSparseMetadataView)): + source_slot = slot // Int32(source.fragments_per_block) + copy_bytes = Int32(4) + prims.cp_async_shared_global( + self._smem_page_offsets.data_ptr(stage_base + offset), + src + cutlass.Int64(source_slot), + 4, + prims.LoadCacheModifier.CA, + cp_size=copy_bytes, + ) + prims.cp_async_commit_group() + prims.cp_async_wait_group(0) + for group in cutlass.range_constexpr((entries + stride - 1) // stride): + offset = owner + Int32(group * stride) + if offset < Int32(entries): + slot = first + offset + if cutlass.const_expr(isinstance(source, DirectSparseMetadataView)): + locator = Int32(-1) + if slot <= last: + logical = Int32(self._smem_page_offsets[stage_base + offset]) + locator = source.map_page(route_state, slot, logical) + self._smem_page_offsets[stage_base + offset] = locator + elif slot > last: + # cp.async zero-fill is not the invalid-locator sentinel. + self._smem_page_offsets[stage_base + offset] = Int32(-1) + @cute.jit def _producer_load_page_offsets( self, @@ -1352,15 +1957,24 @@ def _producer_load_page_offsets( Int32(_decode_gen_task_cache(stage_info)[_TASK_CACHE_KV_RAW_TILE_BASE]) + local_tile_idx ) - _, logical_b_idx = _logical_head_batch(stage_info, self.h_k_idx, self.b_idx) + logical_h_idx, logical_b_idx = _logical_head_batch( + stage_info, self.h_k_idx, self.b_idx + ) + metadata_row = cutlass.Int64(logical_b_idx) + if cutlass.const_expr( + cfg.uses_q_token_kv_block_sparse_page_route + and cfg.use_persistent_scheduler + and not cfg.shares_sparse_pattern + ): + metadata_row = metadata_row * cutlass.Int64( + self.num_heads_kv + ) + cutlass.Int64(logical_h_idx) pages_per_tile = Int32(cfg.tile_size_kv // 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]) - page_table_offset = ( - cutlass.Int64(logical_b_idx) * self.block_table_row_stride - ) - page_idx_kv = self.block_tables + page_table_offset = metadata_row * self.page_table_stride + page_idx_kv = self.page_idx_kv else: if cutlass.const_expr(self.seqlens_kv is None): page_idx_ub = Int32(cfg.max_num_pages_per_seq_kv - 1) @@ -1382,9 +1996,229 @@ def _producer_load_page_offsets( page_table_offset += Int32(cfg.max_num_pages_per_seq_kv) page_idx_kv = self.page_idx_kv smem_page_offsets = self._smem_page_offsets + direct_route = None + if cutlass.const_expr(isinstance(page_idx_kv, DirectSparseMetadataView)): + if cutlass.const_expr( + cfg.use_persistent_scheduler and not cfg.shares_sparse_pattern + ): + page_idx_kv = page_idx_kv.with_head(logical_h_idx) + direct_route = page_idx_kv.resolve_route(logical_b_idx) + if cutlass.const_expr(self.holds_encoded_locator_window): + # The allocation covers the planned maximum, but a shorter + # runtime split may consume fewer tiles. Do not map a + # neighboring CTA's candidates just to fill unused SMEM. + runtime_q_group = _logical_q_group_idx( + cfg, stage_info, self.q_group_idx + ) + local_tiles = _runtime_local_kv_tiles( + cfg, + Int32(task_cache[_TASK_CACHE_SEQ_LEN_KV]), + self.seq_len_q, + _q_group_token_base(cfg, runtime_q_group), + ) + local_page_end = ( + Int32(task_cache[_TASK_CACHE_KV_RAW_TILE_BASE]) + local_tiles + ) * pages_per_tile - Int32(1) + page_idx_ub = cute.math.min(page_idx_ub, local_page_end) lane_idx = cute.arch.thread_idx()[0] & Int32(0x1F) - if cutlass.const_expr(self.stage_page_ids_per_tile): - if lane_idx < pages_per_tile: + page_warp_rank = Int32(0) + if cutlass.const_expr(cfg.page_offsets_num_warps > 1): + page_warp_rank = (cute.arch.thread_idx()[0] >> Int32(5)) - Int32( + cfg.page_offsets_warp_idx + ) + if cutlass.const_expr( + self.q_token_kv_block_sparse_membership_entries + and section == FmhaStage.Head + and inst_id == KV_INST0 + ): + # Page-offset producer warps cooperatively stage the separate + # packed membership row. Each Uint32 supplies four page bytes to + # Softmax, independent of the plain Int32 locator table below. + assert self.q_token_kv_block_sparse_page_memberships is not None + membership_words = self.q_token_kv_block_sparse_membership_entries + raw_tile_base = Int32( + _decode_gen_task_cache(stage_info)[_TASK_CACHE_KV_RAW_TILE_BASE] + ) + membership_word_base = (raw_tile_base * pages_per_tile) // Int32( + Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIPS_PER_WORD + ) + membership_table_offset = metadata_row * cutlass.Int64( + self.q_token_kv_block_sparse_page_membership_stride + ) + membership_threads = cfg.page_offsets_num_warps * 32 + for membership_vector_idx in cutlass.range_constexpr( + (membership_words + membership_threads - 1) // membership_threads + ): + membership_word_idx = ( + page_warp_rank * Int32(32) + + lane_idx + + Int32(membership_vector_idx * membership_threads) + ) + if membership_word_idx < Int32(membership_words): + membership_word = Uint32(0) + first_logical_page_idx = ( + raw_tile_base * pages_per_tile + + membership_word_idx + * Int32(Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIPS_PER_WORD) + ) + if first_logical_page_idx <= page_idx_ub: + membership_word = Uint32( + self.q_token_kv_block_sparse_page_memberships[ + membership_table_offset + + membership_word_base + + membership_word_idx + ] + ) + if cutlass.const_expr( + (cfg.tile_size_kv // cfg.num_tokens_per_page) + % Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIPS_PER_WORD + != 0 + ): + # A tile boundary need not align with a packed + # membership word. Rebase to this CTA's first page + # without reading beyond the live byte prefix. + byte_offset = (raw_tile_base * pages_per_tile) & Int32( + Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIPS_PER_WORD - 1 + ) + if byte_offset != Int32(0): + membership_word >>= byte_offset * Int32( + Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIP_BITS + ) + if ( + first_logical_page_idx + + Int32( + Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIPS_PER_WORD + ) + - byte_offset + <= page_idx_ub + ): + following = Uint32( + self.q_token_kv_block_sparse_page_memberships[ + membership_table_offset + + membership_word_base + + membership_word_idx + + Int32(1) + ] + ) + membership_word |= following << ( + ( + Int32( + Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIPS_PER_WORD + ) + - byte_offset + ) + * Int32( + Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIP_BITS + ) + ) + # The producer owns only the live byte prefix. Mask the + # final word here so padding bytes remain unobservable + # even when the caller reuses an uninitialized workspace. + live_memberships = ( + page_idx_ub - first_logical_page_idx + Int32(1) + ) + if live_memberships < Int32( + Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIPS_PER_WORD + ): + live_bits = live_memberships * Int32( + Q_TOKEN_KV_BLOCK_SPARSE_PAGE_MEMBERSHIP_BITS + ) + membership_word &= (Uint32(1) << live_bits) - Uint32(1) + self._smem_q_token_kv_block_sparse_memberships[ + membership_word_idx + ] = membership_word + if cutlass.const_expr( + self.use_native_paged_kv and cfg.uses_q_token_kv_block_sparse_page_route + ): + first_tile = tile_idx + if cutlass.const_expr(self.holds_encoded_locator_window): + first_tile = Int32( + _decode_gen_task_cache(stage_info)[_TASK_CACHE_KV_RAW_TILE_BASE] + ) + self._prefetch_page_offsets_async( + stage_info, + page_idx_kv, + page_table_offset, + first_tile * pages_per_tile, + page_idx_ub, + lane_idx, + page_warp_rank, + direct_route, + ) + elif cutlass.const_expr(self.holds_encoded_locator_window): + # Stripe whole KV tiles: with two producers, warps own even/odd + # local tiles. Each tile has one writer; runtime-short tails fill + # allocated padding with -1. All producer threads publish the one + # held stage, which K and V retain through the final V use. + raw_tile_base = Int32( + _decode_gen_task_cache(stage_info)[_TASK_CACHE_KV_RAW_TILE_BASE] + ) + held_pages_per_stage = Int32(self.page_ids_per_stage) + for held_tile_group in cutlass.range_constexpr( + (self.encoded_locator_window_tiles + cfg.page_offsets_num_warps - 1) + // cfg.page_offsets_num_warps + ): + held_tile_idx = ( + Int32(held_tile_group * cfg.page_offsets_num_warps) + page_warp_rank + ) + for locator_vector_idx in cutlass.range_constexpr( + (cfg.tile_size_kv // cfg.num_tokens_per_page + 31) // 32 + ): + page_frag = lane_idx + Int32(locator_vector_idx * 32) + grouped_logical_page_idx = ( + raw_tile_base + held_tile_idx + ) * pages_per_tile + page_frag + # Keep the decoded physical page at -1 for TensorMap OOB + # zero-fill. Grouped-Q membership is carried separately. + page_locator = Int32(-1) + if ( + (held_tile_idx < Int32(self.encoded_locator_window_tiles)) + & (page_frag < pages_per_tile) + & (grouped_logical_page_idx <= page_idx_ub) + ): + if cutlass.const_expr( + isinstance(page_idx_kv, DirectSparseMetadataView) + ): + page_locator = page_idx_kv.load_page( + direct_route, grouped_logical_page_idx + ) + else: + page_locator = Int32( + page_idx_kv[ + page_table_offset + grouped_logical_page_idx + ] + ) + held_smem_offset = ( + stage_info.stage_idx * held_pages_per_stage + + held_tile_idx * pages_per_tile + + page_frag + ) + if (held_tile_idx < Int32(self.encoded_locator_window_tiles)) & ( + page_frag < pages_per_tile + ): + smem_page_offsets[held_smem_offset] = page_locator + elif cutlass.const_expr(self.stages_encoded_locators_per_tile): + locator_vectors = (cfg.tile_size_kv // cfg.num_tokens_per_page + 31) // 32 + for locator_vector_idx in cutlass.range_constexpr(locator_vectors): + page_frag = lane_idx + Int32(locator_vector_idx * 32) + logical_page_idx = tile_idx * pages_per_tile + page_frag + page_locator = Int32(-1) + if (page_warp_rank == Int32(0)) & (logical_page_idx <= page_idx_ub): + if cutlass.const_expr( + isinstance(page_idx_kv, DirectSparseMetadataView) + ): + page_locator = page_idx_kv.load_page( + direct_route, logical_page_idx + ) + else: + page_locator = Int32( + page_idx_kv[page_table_offset + logical_page_idx] + ) + if (page_warp_rank == Int32(0)) & (page_frag < pages_per_tile): + smem_offset = stage_info.stage_idx * pages_per_tile + page_frag + smem_page_offsets[smem_offset] = page_locator + elif cutlass.const_expr(self.stage_page_ids_per_tile): + if (page_warp_rank == Int32(0)) & (lane_idx < pages_per_tile): logical_page_idx = cute.math.min( tile_idx * pages_per_tile + lane_idx, page_idx_ub ) @@ -1398,13 +2232,14 @@ def _producer_load_page_offsets( grouped_base_page_idx = ((tile_idx * pages_per_tile) >> Int32(5)) << Int32( 5 ) - 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] - ) + 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] + ) @producer_work @cute.jit @@ -1508,6 +2343,8 @@ class SmemKvResource(DecodeGenResourceBase): cfg: Constexpr[FmhaDecodeConfig] = None tma_desc_k: cutlass.Pointer | None = None tma_desc_v: cutlass.Pointer | None = None + tma_desc_k_sf: cutlass.Pointer | None = None + tma_desc_v_sf: cutlass.Pointer | None = None tma_desc_k_atom: cutlass.Pointer | None = None tma_desc_v_atom: cutlass.Pointer | None = None tma_desc_k_summary: cutlass.Pointer | None = None @@ -1524,7 +2361,9 @@ class SmemKvResource(DecodeGenResourceBase): q_group_idx: Int32 = None seq_len_q: Int32 = None _alloc: Constexpr[SmemAllocation | None] = None + _sf_alloc: Constexpr[SmemAllocation | None] = None _smem_base_kv: cutlass.Array = None + _smem_base_kv_sf: cutlass.Array = None _k_desc_base: prims.Tcgen05SmemDesc = None _v_desc_base: prims.Tcgen05SmemDesc = None kv_desc_slot: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() @@ -1538,27 +2377,43 @@ def _init_placeholder_state(self) -> None: if self.pipeline_config is not None else self.cfg.kv_stages ) + # The shared ring has one allocation for both operands, so K and V + # must have one byte width; their element types may differ (Int8 K + # with E4M3 V), each operand's descriptor carries its own type. + assert self.cfg.k_dtype_bytes == self.cfg.v_dtype_bytes self._smem_base_kv = _placeholder_smem_array( - self.cfg.kv_dtype, + Uint8 if self.cfg.use_nvfp4_kv else self.cfg.k_dtype, self.cfg.smem_kv_tile_elements * num_stages, ) + if self.cfg.use_nvfp4_kv: + self._smem_base_kv_sf = _placeholder_smem_array( + Uint8, + self.cfg.smem_kv_sf_tile_bytes * num_stages, + ) self._k_desc_base = prims.Tcgen05SmemDesc(0) self._v_desc_base = prims.Tcgen05SmemDesc(0) def get_smem_requirements(self) -> list[SmemAllocation]: """Allocate the shared K/V staged SMEM ring.""" - num_stages = ( - self.pipeline_config.num_stages - if self.pipeline_config is not None - else self.cfg.kv_stages - ) + num_stages = self.pipeline_config.num_stages if self._alloc is None: self._alloc = SmemAllocation( name=f"{self.name}", - size_bytes=self.cfg.smem_kv_tile_bytes * num_stages, + size_bytes=self.cfg.smem_kv_storage_tile_bytes * num_stages, alignment=self.cfg.stensor_align, ) - return [self._alloc] + allocs = [self._alloc] + # NVFP4 scale factors ride the same TMA pipeline; give them a separate + # SMEM ring of the same stage depth, read by TransformKvTask. + if self.cfg.use_nvfp4_kv: + if self._sf_alloc is None: + self._sf_alloc = SmemAllocation( + name=f"{self.name}Sf", + size_bytes=self.cfg.smem_kv_sf_tile_bytes * num_stages, + alignment=self.cfg.stensor_align, + ) + allocs.append(self._sf_alloc) + return allocs def get_tmem_requirements(self) -> list[TmemAllocation]: """Shared K/V staging uses SMEM only.""" @@ -1571,7 +2426,9 @@ def _create_initial_task_locals( """Bind the shared K/V ring and build K/V base descriptors.""" if cutlass.const_expr(context is not None and context.smem_base is not None): # Bind the shared K/V SMEM ring. K and V descriptors use the same - # allocation but may differ in leading-byte offset. + # allocation but may differ in leading-byte offset. The ring holds + # one byte width for both operands. + assert self.cfg.k_dtype_bytes == self.cfg.v_dtype_bytes num_stages = ( self.pipeline_config.num_stages if self.pipeline_config is not None @@ -1579,10 +2436,26 @@ def _create_initial_task_locals( ) self._smem_base_kv = cutlass.Array( context.smem_base.data_ptr() + self._alloc.offset, - dtype=self.cfg.kv_dtype, + dtype=Uint8 if self.cfg.use_nvfp4_kv else self.cfg.k_dtype, shape=(self.cfg.smem_kv_tile_elements * num_stages,), addrspace=3, ) + if cutlass.const_expr(self.cfg.use_nvfp4_kv): + # NVFP4 KV is never consumed by MMA. The supported MMA operand types are + # FP16, BF16, and FP8. + self._smem_base_kv_sf = cutlass.Array( + context.smem_base.data_ptr() + self._sf_alloc.offset, + dtype=Uint8, + shape=(self.cfg.smem_kv_sf_tile_bytes * num_stages,), + addrspace=3, + ) + self._k_desc_base = prims.Tcgen05SmemDesc(0) + self._v_desc_base = prims.Tcgen05SmemDesc(0) + return { + "kv_desc": cutlass.Int64(0), + "v_desc_0": cutlass.Int64(0), + "v_desc_1": cutlass.Int64(0), + } kv_tile_bytes = Int32( self.cfg.tile_size_kv * self.cfg.head_dim_kv_stage @@ -1596,7 +2469,9 @@ def _create_initial_task_locals( * self.cfg.kv_dtype_bytes ) stride_byte_offset = Int32(1024) - if cutlass.const_expr(self.cfg.use_fp8_qkv): + if cutlass.const_expr( + self.cfg.k_dtype.width == 8 and self.cfg.v_dtype.width == 8 + ): k_leading_byte_offset = kv_tile_bytes stride_byte_offset = Int32( _major_k_stride_bytes( @@ -1605,11 +2480,16 @@ def _create_initial_task_locals( ) v_leading_byte_offset = k_leading_byte_offset if cutlass.const_expr(self.cfg.tile_size_kv == 256): - # K spans the complete KV256 row between D64 halves. V is - # staged as four semantic KV64 blocks, each with D/64 adjacent - # D64 halves, so its MMA-K leading step is one KV64 block. - v_leading_byte_offset = Int32(64 * 64 * self.cfg.kv_dtype_bytes) - if cutlass.const_expr(self.cfg.use_fp8_qkv or self.cfg.headdim == 64): + # K spans the complete KV256 row between head-dim chunks. V + # is staged as four semantic KV64 blocks, each holding its + # 128-byte head-dim chunks side by side, so the MMA-N leading + # step is one KV64 block of 64 swizzled 128-byte rows for + # every element width. + v_leading_byte_offset = Int32(64 * SWIZZLE_128B_ROW_BYTES) + elif cutlass.const_expr( + (self.cfg.k_dtype.width == 8 and self.cfg.v_dtype.width == 8) + or self.cfg.headdim == 64 + ): v_leading_byte_offset = Int32(0) # Descriptor bases are advanced per stage at consumption time; the # swizzle parameters are invariant for the resource. @@ -1617,13 +2497,13 @@ def _create_initial_task_locals( self._smem_base_kv, leading_byte_offset=k_leading_byte_offset, stride_byte_offset=stride_byte_offset, - layout=_qkv_smem_swizzle(self.cfg), + layout=_kv_smem_swizzle(self.cfg), ) self._v_desc_base = prims.Tcgen05SmemDesc.build( self._smem_base_kv, leading_byte_offset=v_leading_byte_offset, stride_byte_offset=stride_byte_offset, - layout=_qkv_smem_swizzle(self.cfg), + layout=_kv_smem_swizzle(self.cfg), ) return { "kv_desc": cutlass.Int64(0), @@ -1663,9 +2543,18 @@ def init_descriptor_state(self, stage_info: StageInfo) -> None: def _stage_base(self, stage_info: StageInfo) -> cutlass.Array: """Return the SMEM base for the current shared K/V pipeline stage.""" # Return the base pointer for the producer stage selected by TS. - stage_elems = self.cfg.smem_kv_tile_bytes // self.cfg.kv_dtype_bytes + # return self._smem_base_kv.subview( + # stage_info.stage_idx * self.cfg.smem_kv_tile_elements + # ) + stage_elems = self.cfg.smem_kv_storage_tile_bytes // self.cfg.kv_dtype_bytes return self._smem_base_kv.subview(stage_info.stage_idx * stage_elems) + @cute.jit + def _sf_stage_base(self, stage_idx: Int32) -> cutlass.Array: + # Uint8 base pointer for the NVFP4 scale-factor slice of one stage. + sf_stage_elems = Int32(self.cfg.smem_kv_sf_tile_bytes) + return self._smem_base_kv_sf.subview(stage_idx * sf_stage_elems) + @cute.jit def _logical_coords( self, stage_info: StageInfo, tile_idx: Int32 @@ -1793,39 +2682,113 @@ def _producer_load_kv( # runtime-resolved tile_idx so the right per-tile slice is # selected from the shared window. page_fragments = cfg.tile_size_kv // cfg.num_tokens_per_page - if cutlass.const_expr(cfg.use_fp8_qkv): - # FP8 pages are contiguous across the staged head dimension. - fp8_stage_base = self._stage_base(stage_info) - if cutlass.const_expr(cached_page_ids is None): + grouped_tile_idx = self._maybe_runtime_tile_idx(stage_info, local_tile_idx) + if cutlass.const_expr(cfg.use_fp8_kv or cfg.use_nvfp4_kv): + # NVFP4/FP8 pages are contiguous across the staged head dimension. + stage_base = self._stage_base(stage_info) + page_ids = ( + self.page_offsets_kv.page_ids(grouped_tile_idx, local_tile_idx) + if cutlass.const_expr(cached_page_ids is None) + else cached_page_ids + ) + sf_desc = ( + self.tma_desc_v_sf + if cutlass.const_expr(kv_kind == KV_KIND_V) + else self.tma_desc_k_sf + ) + if cutlass.const_expr(cfg.use_nvfp4_kv): + sf_stage_base = self._sf_stage_base(stage_info.stage_idx) + sf_page_elems = Int32( + cfg.num_tokens_per_page * cfg.smem_kv_sf_bytes_per_token + ) + # Each active warp owns disjoint page fragments. Keep NVFP4 + # payload and scale copies together on the same barrier. + active_load_warps = min(page_fragments, cfg.load_num_warps) + assert page_fragments % active_load_warps == 0 + pages_per_warp = page_fragments // active_load_warps + load_warp_rank = Int32(0) + if cutlass.const_expr(cfg.load_num_warps > 1): + load_warp_rank = _load_task_warp_rank(cfg) + if cutlass.const_expr( + cfg.load_num_warps == 1 + ) or load_warp_rank < Int32(active_load_warps): + for local_page_idx in cutlass.range_constexpr(pages_per_warp): + page_frag = local_page_idx + if cutlass.const_expr(cfg.load_num_warps > 1): + page_frag = load_warp_rank * Int32(pages_per_warp) + Int32( + local_page_idx + ) + token_offset, physical_page = _decode_native_page_locator( + cfg, Int32(page_ids[page_frag]) + ) + # Packed E2M1 stores two head-dimension values per byte. + smem_page_offset = Int32( + page_frag + * cfg.num_tokens_per_page + * ( + head_dim_stage // 2 + if cfg.use_nvfp4_kv + and not cfg.store_transformed_kv_in_tmem + else head_dim_stage + ) + ) + _cp_async_bulk_tensor_4d_shared_cta_global_predicated( + stage_base.subview(smem_page_offset), + tma_desc, + ( + Int32( + head_dim_stage_offset // 2 + if cfg.use_nvfp4_kv + and not cfg.store_transformed_kv_in_tmem + else head_dim_stage_offset + ), + token_offset, + logical_h_k_idx, + physical_page, + ), + stage_info.barrier, + ) + # Scale factors for the same page use the same transaction + # barrier as the packed K/V payload. Their TensorMaps + # fold multiple tokens into dim0. + if cutlass.const_expr(cfg.use_nvfp4_kv): + _cp_async_bulk_tensor_4d_shared_cta_global_predicated( + sf_stage_base.subview(page_frag * sf_page_elems), + sf_desc, + ( + Int32(0), + token_offset // Int32(cfg.sf_tma_reshape_factor), + logical_h_k_idx, + physical_page, + ), + stage_info.barrier, + ) + else: + # 128B TMA swizzling permits 64 columns per copy for 16-bit K/V. + # Store each chunk as a separate slab matching the tcgen05 layout. + chunk_hd = min(head_dim_stage, 128 // cfg.kv_dtype_bytes) + num_head_dim_tma_chunks = head_dim_stage // chunk_hd + tile_chunk_elems = chunk_hd * cfg.tile_size_kv + page_chunk_elems = chunk_hd * cfg.num_tokens_per_page + stage_base = self._stage_base(stage_info) + if cutlass.const_expr(cfg.uses_scattered_page_route): + assert cached_page_ids is None grouped_tile_idx = self._maybe_runtime_tile_idx( stage_info, local_tile_idx ) - fp8_page_ids = self.page_offsets_kv.page_ids(grouped_tile_idx) - else: - fp8_page_ids = cached_page_ids - for fp8_page_frag in cutlass.range_constexpr(page_fragments): - fp8_page_id = Int32(fp8_page_ids[fp8_page_frag]) - fp8_smem_page_offset = Int32( - fp8_page_frag * cfg.num_tokens_per_page * head_dim_stage - ) - _cp_async_bulk_tensor_4d_shared_cta_global_predicated( - fp8_stage_base.subview(fp8_smem_page_offset), + _issue_sparse_page_copies( + cfg, + self.page_offsets_kv, + stage_base, tma_desc, - ( - Int32(head_dim_stage_offset), - Int32(0), - logical_h_k_idx, - fp8_page_id, - ), stage_info.barrier, + grouped_tile_idx, + local_tile_idx, + logical_h_k_idx, + head_dim_stage, + head_dim_stage_offset, ) - else: - # 16-bit pages are split into 64-column chunks inside the - # staged head-dim slice. - chunk_hd = min(head_dim_stage, 64) - num_chunks = head_dim_stage // chunk_hd - tile_chunk_elems = chunk_hd * cfg.tile_size_kv - page_chunk_elems = chunk_hd * cfg.num_tokens_per_page + return if cutlass.const_expr(cached_page_ids is None): # Resolve the tile on every lane before the elected-lane # branch. The release compiler rejects a local that is @@ -1834,9 +2797,10 @@ def _producer_load_kv( stage_info, local_tile_idx ) if prims.elect_sync(): - stage_base = self._stage_base(stage_info) if cutlass.const_expr(cached_page_ids is None): - page_ids = self.page_offsets_kv.page_ids(grouped_tile_idx) + page_ids = self.page_offsets_kv.page_ids( + grouped_tile_idx, local_tile_idx + ) else: page_ids = cached_page_ids # Consume each cached page ID across every head-dimension @@ -1844,8 +2808,12 @@ def _producer_load_kv( # this order bounds coordinate live ranges in the unrolled # TMA sequence for every supported page size. for page_frag in cutlass.range_constexpr(page_fragments): - page_id = Int32(page_ids[page_frag]) - for chunk_idx in cutlass.range_constexpr(num_chunks): + token_offset, physical_page = _decode_native_page_locator( + cfg, Int32(page_ids[page_frag]) + ) + for chunk_idx in cutlass.range_constexpr( + num_head_dim_tma_chunks + ): local_head_dim_offset = chunk_idx * chunk_hd global_head_dim_offset = ( head_dim_stage_offset + local_head_dim_offset @@ -1859,14 +2827,14 @@ def _producer_load_kv( tma_desc, ( Int32(global_head_dim_offset), - Int32(0), + token_offset, logical_h_k_idx, - page_id, + physical_page, ), stage_info.barrier, ) - elif cutlass.const_expr(cfg.use_fp8_qkv): - # Dense FP8 path: one tensor TMA loads the staged K or V tile. + elif cutlass.const_expr(cfg.k_dtype.width == 8 and cfg.v_dtype.width == 8): + # Dense byte-wide path: one tensor TMA loads the whole K or V tile. tile_idx = self._maybe_runtime_tile_idx(stage_info, local_tile_idx) tile_offset = tile_idx * Int32(cfg.tile_size_kv) if prims.elect_sync(): @@ -1890,8 +2858,9 @@ def _producer_load_kv( tile_chunk_elems = chunk_hd * cfg.tile_size_kv tile_idx = self._maybe_runtime_tile_idx(stage_info, local_tile_idx) tile_offset = tile_idx * Int32(cfg.tile_size_kv) + # Keep the Array type stable across the elected-lane control-flow join. + stage_base = self._stage_base(stage_info) if prims.elect_sync(): - stage_base = self._stage_base(stage_info) for chunk_idx in cutlass.range_constexpr(num_chunks): local_head_dim_offset = chunk_idx * chunk_hd global_head_dim_offset = ( @@ -1926,12 +2895,15 @@ def _producer_load_kv_tile_256( The public decode TensorMaps expose KV64 (or one smaller page) fragments. K places semantic KV64 blocks in physical order - ``(0, 2, 1, 3)`` while V keeps semantic block order with adjacent D64 - halves. Dense and paged profiles derive those fragments from one - contiguous tile; block-sparse profiles consume four prepared KV64 - origins retained by the instruction-local metadata resource. + ``(0, 2, 1, 3)`` while V keeps semantic block order with adjacent + head-dim chunks. A chunk is one swizzled 128-byte row: 64 two-byte or + 128 one-byte elements. Dense and paged profiles derive those fragments + from one contiguous tile; block-sparse profiles consume four prepared + KV64 origins retained by the instruction-local metadata resource. """ cfg = self.cfg + chunk_hd = min(cfg.headdim, cfg.kv_swizzle_row_elems) + num_chunks = cfg.headdim // chunk_hd grouped_tile_idx = Int32(0) if cutlass.const_expr(not cfg.use_block_sparse): grouped_tile_idx = self._maybe_runtime_tile_idx(stage_info, local_tile_idx) @@ -1948,7 +2920,9 @@ def _producer_load_kv_tile_256( and cached_page_ids is None ): assert self.page_offsets_kv is not None - dense_page_ids = self.page_offsets_kv.page_ids(grouped_tile_idx) + dense_page_ids = self.page_offsets_kv.page_ids( + grouped_tile_idx, local_tile_idx + ) # Select the logical source before the constexpr 4 x 2 loop so K/V # and summary routes retain one physical KV256 staging body. route_tma_desc = tma_desc @@ -1989,14 +2963,16 @@ def _producer_load_kv_tile_256( physical_block = KV_TILE_256_K_SLOT_FOR_SEMANTIC_ATOM[ semantic_block ] - for dim_half in cutlass.range_constexpr(2): + for dim_chunk in cutlass.range_constexpr(num_chunks): if cutlass.const_expr(kv_kind == KV_KIND_K): block_base = ( - dim_half * cfg.tile_size_kv * 64 + physical_block * 64 * 64 + dim_chunk * cfg.tile_size_kv * chunk_hd + + physical_block * 64 * chunk_hd ) else: block_base = ( - semantic_block * cfg.headdim * 64 + dim_half * 64 * 64 + semantic_block * cfg.headdim * 64 + + dim_chunk * 64 * chunk_hd ) if cutlass.const_expr(cfg.use_block_sparse): @@ -2004,7 +2980,7 @@ def _producer_load_kv_tile_256( stage_base.subview(block_base), route_tma_desc, ( - Int32(dim_half * 64), + Int32(dim_chunk * chunk_hd), token_coord, logical_h_k_idx, storage_coord, @@ -2020,16 +2996,22 @@ def _producer_load_kv_tile_256( ) logical_page = token_in_tile // cfg.num_tokens_per_page token_in_page = token_in_tile % cfg.num_tokens_per_page - page_id = Int32(dense_page_ids[logical_page]) - smem_offset = block_base + fragment * fragment_tokens * 64 + assert dense_page_ids is not None + page_locator = Int32(dense_page_ids[logical_page]) + token_offset, physical_page = _decode_native_page_locator( + cfg, page_locator + ) + smem_offset = ( + block_base + fragment * fragment_tokens * chunk_hd + ) prims.cp_async_bulk_tensor_shared_cta_global( stage_base.subview(smem_offset), tma_desc, ( - Int32(dim_half * 64), - Int32(token_in_page), + Int32(dim_chunk * chunk_hd), + token_offset + Int32(token_in_page), logical_h_k_idx, - page_id, + physical_page, ), stage_info.barrier, ) @@ -2039,7 +3021,7 @@ def _producer_load_kv_tile_256( stage_base.subview(block_base), tma_desc, ( - Int32(dim_half * 64), + Int32(dim_chunk * chunk_hd), tile_offset + Int32(semantic_block * 64), logical_h_k_idx, logical_b_idx, @@ -2194,7 +3176,9 @@ def _build_kv_desc( """Advance the shared K or V descriptor to the committed stage.""" # Consumers see the same descriptor layout for every stage; only the # base address advances by the committed SMEM stage index. - stage_offset_bytes = stage_info.stage_idx * Int32(self.cfg.smem_kv_tile_bytes) + stage_offset_bytes = stage_info.stage_idx * Int32( + self.cfg.smem_kv_storage_tile_bytes + ) return ( self._v_desc_base.advance_start_address(stage_offset_bytes) if cutlass.const_expr(kv_kind == KV_KIND_V) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_transformed_kv.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_transformed_kv.py new file mode 100644 index 000000000000..148ebcae6b74 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_transformed_kv.py @@ -0,0 +1,574 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Transformed K/V SMEM resource for mixed-dtype FMHA decode.""" + +from dataclasses import dataclass +from typing import ClassVar + +import cutlass +import cutlass.cute as cute +from cutlass import Int32 +from cutlass.experimental import primitives as prims +from cutlass.experimental.task_scheduling.enums import WorkAttr +from cutlass.experimental.task_scheduling.memory import ( + ResourceContext, + SmemAllocation, + TmemAllocation, +) +from cutlass.experimental.task_scheduling.resources import ( + StageInfo, + TaskLocalVariable, + consumer_work, + producer_work, +) + +from ...placeholder_helpers import _placeholder_smem_array +from ..fmha_decode_config import FmhaDecodeConfig +from ..fmha_decode_constants import KV_INST0, KV_KIND_K, KV_KIND_V +from .helpers_common import ( + _TASK_CACHE_WARP_GRP_THREAD_IDX, + Constexpr, + DecodeGenResourceBase, + ResourceVars, + _decode_gen_task_cache, + _major_k_stride_bytes, + _qkv_smem_swizzle, +) +from .smem_resources import SmemKvResource, SmemKvTileResource + + +@cute.jit +def _cvt_e4m3x2_word_to_bf16x2(packed_word: Int32, shift: Constexpr[int]) -> Int32: + """Extract an E4M3x2 half-word and convert it to BF16x2.""" + packed_fp8 = (packed_word >> Int32(shift)) & Int32(0xFFFF) + return cute.arch.inline_ptx( + "{ .reg .b8 f0, f1, u0, u1; .reg .b16 fp8, h0, h1, b0, b1; " + ".reg .b32 half2; " + "mov.b32 {f0, f1, u0, u1}, {$r0}; " + "mov.b16 fp8, {f0, f1}; cvt.rn.f16x2.e4m3x2 half2, fp8; " + "mov.b32 {h0, h1}, half2; " + "cvt.rn.bf16.f16 b0, h0; cvt.rn.bf16.f16 b1, h1; " + "mov.b32 {$w0}, {b0, b1}; }", + write_only_types=[Int32], + read_only_args=[packed_fp8], + ) + + +@cute.jit +def _cvt_e4m3x2_byte_to_bf16x2(packed_byte: Int32) -> Int32: + """Broadcast one E4M3 byte and convert the pair to BF16x2.""" + sf_pair = (packed_byte & Int32(0xFF)) * Int32(0x101) + return cute.arch.inline_ptx( + "{ .reg .b8 f0, f1, u0, u1; .reg .b16 fp8, h0, h1, b0, b1; " + ".reg .b32 half2; " + "mov.b32 {f0, f1, u0, u1}, {$r0}; " + "mov.b16 fp8, {f0, f1}; cvt.rn.f16x2.e4m3x2 half2, fp8; " + "mov.b32 {h0, h1}, half2; " + "cvt.rn.bf16.f16 b0, h0; cvt.rn.bf16.f16 b1, h1; " + "mov.b32 {$w0}, {b0, b1}; }", + write_only_types=[Int32], + read_only_args=[sf_pair], + ) + + +@cute.jit +def _cvt_e2m1_word_to_e4m3_words( + packed_word: Int32, sf_byte: Int32 +) -> tuple[Int32, Int32]: + """Scale eight dense-packed E2M1 values and return two E4M3 words.""" + sf_pair = (sf_byte & Int32(0xFF)) * Int32(0x101) + return cute.arch.inline_ptx( + """ + { + .reg .b16 sf; + .reg .b32 sfFp16x2; + .reg .b8 b0, b1, b2, b3; + .reg .b32 h0, h1, h2, h3; + .reg .b16 e0, e1, e2, e3; + mov.b32 {sf, _}, {$r1}; + cvt.rn.f16x2.e4m3x2 sfFp16x2, sf; + mov.b32 {b0, b1, b2, b3}, {$r0}; + cvt.rn.f16x2.e2m1x2 h0, b0; + cvt.rn.f16x2.e2m1x2 h1, b1; + cvt.rn.f16x2.e2m1x2 h2, b2; + cvt.rn.f16x2.e2m1x2 h3, b3; + mul.rn.f16x2 h0, h0, sfFp16x2; + mul.rn.f16x2 h1, h1, sfFp16x2; + mul.rn.f16x2 h2, h2, sfFp16x2; + mul.rn.f16x2 h3, h3, sfFp16x2; + cvt.rn.satfinite.e4m3x2.f16x2 e0, h0; + cvt.rn.satfinite.e4m3x2.f16x2 e1, h1; + cvt.rn.satfinite.e4m3x2.f16x2 e2, h2; + cvt.rn.satfinite.e4m3x2.f16x2 e3, h3; + mov.b32 {$w0}, {e0, e1}; + mov.b32 {$w1}, {e2, e3}; + } + """, + write_only_types=[Int32, Int32], + read_only_args=[packed_word, sf_pair], + ) + + +@cute.jit +def _cvt_e2m1x2_byte_to_bf16x2(packed_byte: Int32) -> Int32: + """Convert a packed E2M1x2 byte to BF16x2.""" + return cute.arch.inline_ptx( + "{ .reg .b8 byte, u0, u1, u2; .reg .b16 h0, h1, b0, b1; " + ".reg .b32 h; mov.b32 {byte, u0, u1, u2}, {$r0}; " + "cvt.rn.f16x2.e2m1x2 h, byte; mov.b32 {h0, h1}, h; " + "cvt.rn.bf16.f16 b0, h0; cvt.rn.bf16.f16 b1, h1; " + "mov.b32 {$w0}, {b0, b1}; }", + write_only_types=[Int32], + read_only_args=[packed_byte & Int32(0xFF)], + ) + + +@dataclass(kw_only=True) +class SmemTransformedKvResource(DecodeGenResourceBase): + """Shared transformed K/V staging resource for mixed Q/KV decode. + + The upstream ``SmemKvResource`` or ``SmemKvTileResource`` hold raw K/V + in their GMEM dtype. Mixed precision KV introduces this async producer + between raw K/V TMA and MMA so BMM1/BMM2 consume the Q-side MMA input dtype. + """ + + _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ( + "kv_desc_slot", + prims.Tcgen05SmemDesc, + prims.Tcgen05SmemDesc(0), + "SMEM descriptor for transformed K loads consumed by QK MMA.", + ), + ( + "v_desc_0_slot", + prims.Tcgen05SmemDesc, + prims.Tcgen05SmemDesc(0), + "First transformed V descriptor consumed by VP MMA.", + ), + ( + "v_desc_1_slot", + prims.Tcgen05SmemDesc, + prims.Tcgen05SmemDesc(0), + "Second transformed V descriptor consumed by VP MMA.", + ), + ) + cfg: Constexpr[FmhaDecodeConfig] = None + src_smem_kv: Constexpr[SmemKvResource | None] = None + src_smem_k0: Constexpr[SmemKvTileResource | None] = None + src_smem_k1: Constexpr[SmemKvTileResource | None] = None + src_smem_v0: Constexpr[SmemKvTileResource | None] = None + src_smem_v1: Constexpr[SmemKvTileResource | None] = None + page_idx_kv: cute.Pointer | None = None + num_heads_kv: Int32 | None = None + _alloc: Constexpr[SmemAllocation | None] = None + _smem_base_kv: cutlass.Array = None + _k_desc_base: prims.Tcgen05SmemDesc = None + _v_desc_base: prims.Tcgen05SmemDesc = None + kv_desc_slot: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + v_desc_0_slot: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + v_desc_1_slot: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + + def _init_placeholder_state(self) -> None: + """Create placeholder state for the shared K/V SMEM ring.""" + num_stages = self.pipeline_config.num_stages + self._smem_base_kv = _placeholder_smem_array( + self.cfg.qk_dtype, + (self.cfg.smem_transformed_kv_tile_bytes // self.cfg.qk_dtype_bytes) + * num_stages, + ) + self._k_desc_base = prims.Tcgen05SmemDesc(0) + self._v_desc_base = prims.Tcgen05SmemDesc(0) + + def get_smem_requirements(self) -> list[SmemAllocation]: + """Allocate the shared K/V staged SMEM ring.""" + num_stages = ( + self.pipeline_config.num_stages + if self.pipeline_config is not None + else self.cfg.transformed_kv_stages + ) + if self._alloc is None: + self._alloc = SmemAllocation( + name=f"{self.name}", + size_bytes=self.cfg.smem_transformed_kv_tile_bytes * num_stages, + alignment=self.cfg.stensor_align, + ) + return [self._alloc] + + def get_tmem_requirements(self) -> list[TmemAllocation]: + """Return no TMEM requirements because this resource uses only SMEM.""" + return [] + + @cute.jit + def _create_initial_task_locals( + self, context: ResourceContext | None = None + ) -> ResourceVars: + """Bind the shared K/V ring and build K/V base descriptors.""" + if cutlass.const_expr(context is not None and context.smem_base is not None): + num_stages = self.pipeline_config.num_stages + stage_elems = ( + self.cfg.smem_transformed_kv_tile_bytes // self.cfg.qk_dtype_bytes + ) + self._smem_base_kv = cutlass.Array( + context.smem_base.data_ptr() + self._alloc.offset, + dtype=self.cfg.qk_dtype, + shape=(stage_elems * num_stages,), + addrspace=3, + ) + kv_tile_bytes = Int32(self.cfg.smem_transformed_kv_tile_bytes) + k_leading_byte_offset = Int32(16384) + stride_byte_offset = Int32(1024) + if cutlass.const_expr(self.cfg.use_fp8_qk): + k_leading_byte_offset = kv_tile_bytes + stride_byte_offset = Int32( + _major_k_stride_bytes( + self.cfg.qk_dtype_bytes, self.cfg.head_dim_kv_stage + ) + ) + v_leading_byte_offset = k_leading_byte_offset + if cutlass.const_expr( + self.cfg.use_fp8_qk or self.cfg.head_dim_kv_stage == 64 + ): + v_leading_byte_offset = Int32(0) + self._k_desc_base = prims.Tcgen05SmemDesc.build( + self._smem_base_kv, + leading_byte_offset=k_leading_byte_offset, + stride_byte_offset=stride_byte_offset, + layout=_qkv_smem_swizzle(self.cfg), + ) + self._v_desc_base = prims.Tcgen05SmemDesc.build( + self._smem_base_kv, + leading_byte_offset=v_leading_byte_offset, + stride_byte_offset=stride_byte_offset, + layout=_qkv_smem_swizzle(self.cfg), + ) + return { + "kv_desc": cutlass.Int64(0), + "v_desc_0": cutlass.Int64(0), + "v_desc_1": cutlass.Int64(0), + } + + @cute.jit + def _create_work_tile_task_locals( + self, context: ResourceContext | None = None + ) -> ResourceVars: + """Provide shared K/V descriptor slots for one work tile.""" + _ = context + return { + "kv_desc": cutlass.Int64(0), + "v_desc_0": cutlass.Int64(0), + "v_desc_1": cutlass.Int64(0), + } + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_load_state(self, stage_info: StageInfo) -> None: + """Initialize producer-side shared K/V SMEM state.""" + # ProdAuxWork: bind the shared K/V ring and descriptor bases before + # the transform task alternates K and V stages through it. + self._create_initial_task_locals(stage_info.context) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_descriptor_state(self, stage_info: StageInfo) -> None: + """Initialize consumer-side shared K/V descriptor state.""" + # ConsAuxWork: initialize descriptor slots for both K and V consumers + # of the shared K/V ring. + self._create_initial_task_locals(stage_info.context) + + @cute.jit + def _dst_q_smem_byte_offset(self, elem_idx: Int32) -> Int32: + """Map a logical transformed-K/V element to its swizzled SMEM byte offset.""" + cfg = self.cfg + tile_row_idx = elem_idx // Int32(cfg.head_dim_kv_stage) + tile_col_idx = elem_idx % Int32(cfg.head_dim_kv_stage) + if cutlass.const_expr(cfg.use_fp8_qk and cfg.head_dim_kv_stage == 64): + swizzle_bytes = 64 + else: + swizzle_bytes = 128 + q_elems_per_smem_row = swizzle_bytes // cfg.qk_dtype_bytes + smem_slice_idx = tile_col_idx // Int32(q_elems_per_smem_row) + smem_col_idx = tile_col_idx % Int32(q_elems_per_smem_row) + smem_row_idx = smem_slice_idx * Int32(cfg.tile_size_kv) + tile_row_idx + dst_byte_offset = smem_row_idx * Int32(swizzle_bytes) + smem_col_idx * Int32( + cfg.qk_dtype_bytes + ) + if cutlass.const_expr(cfg.use_fp8_qk and cfg.head_dim_kv_stage == 64): + return dst_byte_offset ^ ( + ((dst_byte_offset >> Int32(7)) & Int32(3)) << Int32(4) + ) + return dst_byte_offset ^ ((smem_row_idx & Int32(7)) << Int32(4)) + + @cute.jit + def _transform_load_fp8_chunk( + self, + chunk_idx: Int32, + elem_idx: Int32, + src_word_base: cutlass.Array, + dst_pair_base: cutlass.Array, + ) -> None: + """Convert one eight-element FP8 chunk to BF16 in transformed SMEM.""" + cfg = self.cfg + src_word_idx = chunk_idx * Int32(2) + if cutlass.const_expr(cfg.head_dim_kv_stage == 64): + # H64 transform chunks always cover eight consecutive FP8 + # values, so derive the TMA-swizzled source word directly from + # the chunk index instead of expanding to elem row/column first. + src_row_idx = chunk_idx >> Int32(3) + src_sector_idx = (src_row_idx << Int32(2)) + ( + (chunk_idx & Int32(7)) >> Int32(1) + ) + src_swizzled_sector_idx = src_sector_idx ^ ( + (src_row_idx >> Int32(1)) & Int32(3) + ) + src_word_idx = src_swizzled_sector_idx * Int32(4) + ( + chunk_idx & Int32(1) + ) * Int32(2) + else: + # The 128-byte TMA swizzle XORs bits [6:4] with the logical row + # modulo eight. Convert that byte-address mapping to an Int32 word + # index before loading the raw FP8 chunk. + src_row_idx = elem_idx // Int32(cfg.head_dim_kv_stage) + src_word_idx = src_word_idx ^ ((src_row_idx & Int32(7)) << Int32(2)) + packed_fp8_word0 = src_word_base[src_word_idx] + packed_fp8_word1 = src_word_base[src_word_idx + Int32(1)] + dst_pair_idx = self._dst_q_smem_byte_offset(elem_idx) >> Int32(2) + bf16_pair0 = _cvt_e4m3x2_word_to_bf16x2(packed_fp8_word0, 0) + bf16_pair1 = _cvt_e4m3x2_word_to_bf16x2(packed_fp8_word0, 16) + bf16_pair2 = _cvt_e4m3x2_word_to_bf16x2(packed_fp8_word1, 0) + bf16_pair3 = _cvt_e4m3x2_word_to_bf16x2(packed_fp8_word1, 16) + dst_pair_base[dst_pair_idx] = bf16_pair0 + dst_pair_base[dst_pair_idx + Int32(1)] = bf16_pair1 + dst_pair_base[dst_pair_idx + Int32(2)] = bf16_pair2 + dst_pair_base[dst_pair_idx + Int32(3)] = bf16_pair3 + + @cute.jit + def _load_nvfp4_scale_byte( + self, + src_sf_base: cutlass.Array, + head_dim_stage_idx: Constexpr[int], + is_v: int, + elem_idx: Int32, + ) -> Int32: + """Read one E4M3 scale-factor byte from the SMEM-staged SF tile. + + The SF for the raw K/V tile is TMA'd into ``SmemKvResource``'s SF ring + (``src_sf_base`` is the matching stage's Uint8 slice). K uses + token-major, SF-minor order; V uses TRT-LLM's 4-token interleaved + layout. ``sf_idx`` includes the logical staged head-slice offset. + """ + cfg = self.cfg + sf_bytes_per_token = Int32(cfg.smem_kv_sf_bytes_per_token) + row_idx = elem_idx // Int32(cfg.head_dim_kv_stage) + sf_idx = Int32(head_dim_stage_idx * (cfg.head_dim_kv_stage // 16)) + ( + elem_idx % Int32(cfg.head_dim_kv_stage) + ) // Int32(16) + sf_offset = row_idx * sf_bytes_per_token + sf_idx + if cutlass.const_expr(is_v): + sf_offset = ( + (row_idx // Int32(4)) * Int32(4) * sf_bytes_per_token + + sf_idx * Int32(4) + + row_idx % Int32(4) + ) + return Int32(src_sf_base[sf_offset]) + + @cute.jit + def _transform_load_nvfp4_chunk( + self, + stage_info: StageInfo, + head_dim_stage_idx: Constexpr[int], + inst_id: int, + is_v: int, + chunk_idx: Int32, + elem_idx: Int32, + src_word_base: cutlass.Array, + src_sf_base: cutlass.Array, + dst_word_base: cutlass.Array, + dst_pair_base: cutlass.Array, + ) -> None: + """Dequantize one eight-element NVFP4 chunk to the effective MMA dtype.""" + cfg = self.cfg + packed_word = src_word_base[chunk_idx] + sf_byte = self._load_nvfp4_scale_byte( + src_sf_base, head_dim_stage_idx, is_v, elem_idx + ) + dst_word_idx = self._dst_q_smem_byte_offset(elem_idx) >> Int32(2) + if cutlass.const_expr(cfg.use_fp8_qk): + fp8_word0, fp8_word1 = _cvt_e2m1_word_to_e4m3_words(packed_word, sf_byte) + dst_word_base[dst_word_idx] = fp8_word0 + dst_word_base[dst_word_idx + Int32(1)] = fp8_word1 + else: + sf_bf16x2 = _cvt_e4m3x2_byte_to_bf16x2(sf_byte) + for pair_idx in cutlass.range_constexpr(4): + packed_byte = (packed_word >> Int32(pair_idx * 8)) & Int32(0xFF) + bf16_pair = _cvt_e2m1x2_byte_to_bf16x2(packed_byte) + dst_pair_base[dst_word_idx + Int32(pair_idx)] = prims.mul_bf16x2( + bf16_pair, sf_bf16x2 + ) + + @cute.jit + def _transform_load( + self, + stage_info: StageInfo, + inst_id: int, + is_v: int, + *, + head_dim_stage_idx: Constexpr[int], + ) -> None: + """Transform one raw K or V pipeline stage into MMA-ready SMEM.""" + cfg = self.cfg + task_cache = _decode_gen_task_cache(stage_info) + warp_grp_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] + transform_thread_idx = warp_grp_thread_idx - Int32( + (cfg.transform_kv_warp_idx % 4) * 32 + ) + + dst_stage_elems = cfg.smem_transformed_kv_tile_bytes // cfg.qk_dtype_bytes + dst_base = self._smem_base_kv.subview(stage_info.stage_idx * dst_stage_elems) + dst_word_base = cutlass.Array( + dst_base.data_ptr(), + dtype=Int32, + shape=(cfg.smem_transformed_kv_tile_bytes // 4,), + addrspace=3, + ) + dst_pair_base = cutlass.Array( + dst_base.data_ptr(), + dtype=Int32, + shape=(cfg.smem_transformed_kv_tile_bytes // 4,), + addrspace=3, + ) + + src = self.src_smem_kv + if cutlass.const_expr(self.src_smem_kv is None): + if cutlass.const_expr(is_v): + src = self.src_smem_v0 if inst_id == KV_INST0 else self.src_smem_v1 + else: + src = self.src_smem_k0 if inst_id == KV_INST0 else self.src_smem_k1 + assert src is not None + src_stage_idx = src.state_src.consumer_work_stage + src_stage_elems = cfg.smem_kv_tile_elements + src_base = src._smem_base_kv.subview(src_stage_idx * src_stage_elems) + src_word_base = cutlass.Array( + src_base.data_ptr(), + dtype=Int32, + shape=(cfg.smem_kv_storage_tile_bytes // 4,), + addrspace=3, + ) + # SF for this raw-KV stage, staged into SmemKv's SF ring by the loader. + src_sf_base = None + if cutlass.const_expr(cfg.use_nvfp4_kv): + src_sf_base = src._sf_stage_base(src_stage_idx) + total_elems = cfg.tile_size_kv * cfg.head_dim_kv_stage + num_threads = cfg.transform_kv_num_warps * 32 + chunks_per_thread = total_elems // (num_threads * 8) + for e in cutlass.range(chunks_per_thread, unroll=8): + chunk_idx = Int32(e * num_threads) + transform_thread_idx + elem_idx = chunk_idx * Int32(8) + if cutlass.const_expr(cfg.use_nvfp4_kv): + self._transform_load_nvfp4_chunk( + stage_info, + head_dim_stage_idx, + inst_id, + is_v, + chunk_idx, + elem_idx, + src_word_base, + src_sf_base, + dst_word_base, + dst_pair_base, + ) + else: + self._transform_load_fp8_chunk( + chunk_idx, + elem_idx, + src_word_base, + dst_pair_base, + ) + + cute.arch.fence_view_async_shared() + + @producer_work + @cute.jit + def transform_k0( + self, stage_info: StageInfo, *, head_dim_stage_idx: Constexpr[int] + ) -> None: + """Transform K instance 0 into the acquired destination stage.""" + self._transform_load(stage_info, 0, 0, head_dim_stage_idx=head_dim_stage_idx) + + @producer_work + @cute.jit + def transform_k1( + self, stage_info: StageInfo, *, head_dim_stage_idx: Constexpr[int] + ) -> None: + """Transform K instance 1 into the acquired destination stage.""" + self._transform_load(stage_info, 1, 0, head_dim_stage_idx=head_dim_stage_idx) + + @producer_work + @cute.jit + def transform_v0( + self, stage_info: StageInfo, *, head_dim_stage_idx: Constexpr[int] + ) -> None: + """Transform V instance 0 into the acquired destination stage.""" + self._transform_load(stage_info, 0, 1, head_dim_stage_idx=head_dim_stage_idx) + + @producer_work + @cute.jit + def transform_v1( + self, stage_info: StageInfo, *, head_dim_stage_idx: Constexpr[int] + ) -> None: + """Transform V instance 1 into the acquired destination stage.""" + self._transform_load(stage_info, 1, 1, head_dim_stage_idx=head_dim_stage_idx) + + @cute.jit + def _build_kv_desc(self, stage_info: StageInfo, is_v: int) -> prims.Tcgen05SmemDesc: + """Advance the shared K or V descriptor to the committed stage.""" + # Consumers see the same descriptor layout for every stage; only the + # base address advances by the committed SMEM stage index. + stage_offset_bytes = stage_info.stage_idx * Int32( + self.cfg.smem_transformed_kv_tile_bytes + ) + return ( + self._v_desc_base.advance_start_address(stage_offset_bytes) + if cutlass.const_expr(is_v) + else self._k_desc_base.advance_start_address(stage_offset_bytes) + ) + + @consumer_work(returns=kv_desc_slot) + @cute.jit + def k_desc_0(self, stage_info: StageInfo) -> prims.Tcgen05SmemDesc: + """Publish the first K descriptor consumed by QK MMA.""" + # ConsWork: expose the committed shared-ring K stage for QK instance 0. + return self._build_kv_desc(stage_info, KV_KIND_K) + + @consumer_work(returns=kv_desc_slot) + @cute.jit + def k_desc_1(self, stage_info: StageInfo) -> prims.Tcgen05SmemDesc: + """Publish the second K descriptor consumed by QK MMA.""" + # ConsWork: expose the committed shared-ring K stage for QK instance 1. + return self._build_kv_desc(stage_info, KV_KIND_K) + + @consumer_work(returns=v_desc_0_slot) + @cute.jit + def v_desc_0(self, stage_info: StageInfo) -> prims.Tcgen05SmemDesc: + """Publish the first V descriptor consumed by PV MMA.""" + # ConsWork: expose the committed shared-ring V stage for PV instance 0. + return self._build_kv_desc(stage_info, KV_KIND_V) + + @consumer_work(returns=v_desc_1_slot) + @cute.jit + def v_desc_1(self, stage_info: StageInfo) -> prims.Tcgen05SmemDesc: + """Publish the second V descriptor consumed by PV MMA.""" + # ConsWork: expose the committed shared-ring V stage for PV instance 1. + return self._build_kv_desc(stage_info, KV_KIND_V) 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 11b384a103bd..1a2a04fbd976 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 @@ -38,25 +38,23 @@ producer_work, ) -from ..fmha_decode_config import FmhaDecodeConfig from ...placeholder_helpers import _placeholder_smem_array +from ..fmha_decode_config import FmhaDecodeConfig +from ..fmha_decode_constants import FP8_P_QUANT_SCALE from .helpers_common import ( - Constexpr, - DecodeGenResourceBase, - ResourceVars, - fadd2, - ffma2, - fmul2, _TASK_CACHE_LANE_IDX, _TASK_CACHE_TMEM_BASE_OFFSET, _TASK_CACHE_WARP_GRP_THREAD_IDX, _TASK_CACHE_WARP_IDX, + Constexpr, + DecodeGenResourceBase, + ResourceVars, + _attention_sink_head_stride, _decode_gen_task_cache, _keeps_col_base, _keeps_row_idx, _keeps_tcgen05_ld, _keeps_tcgen05_st, - _attention_sink_head_stride, _local_head_from_q_output_row, _logical_head_batch, _logical_q_group_idx, @@ -66,9 +64,13 @@ _q_row_is_valid_for_seq, _q_tile_output_row_base, _q_tile_valid_rows_for_seq, + _masked_weighted_sum, _neg_max_f32, _pack_float2_to_bf16, _pack_float2_to_fp16, + fadd2, + ffma2, + fmul2, ) from .helpers_kv_tile_idx import ( _load_runtime_seq_len_kv, @@ -86,13 +88,15 @@ _attention_sink_for_scale_idx, _pack_float4_to_fp8_e4m3, ) +from .sage_scales import SageVScalesResource from .smem_p import SmemPResource from .tmem_softmax_stats import TmemSoftmaxLocalResource _KV_TILE_256_CORRECTION_THREADS = 128 _KV_TILE_256_LOGICAL_OUTPUT_ROWS = 64 -# One D32 fragment per logical output row, padded by four floats so adjacent -# rows fall on different bank groups. +# One D32 fragment slot, padded by four floats so adjacent slots fall on +# different bank groups: one per logical output row (row-owner exchange) or +# per lane (column-owner exchange). _KV_TILE_256_EXCHANGE_FRAGMENT_STRIDE = 36 _KV_TILE_256_STATS_PER_THREAD = 4 @@ -118,6 +122,9 @@ class TmemCorrResource(DecodeGenResourceBase): partial_stats_ptr: cute.Pointer = None split_kv_counter_ptr: cute.Pointer = None attention_sinks_ptr: cute.Pointer = None + # The correction task's staged ``sfV`` / ``v_mean`` resource; ``None`` + # without Sage attention. + sage_v_scales: SageVScalesResource | None = None seqlens_kv: cute.Pointer = None max_seq_len_kv: Constexpr[int] = 0 seq_len_q: Int32 = None @@ -148,6 +155,7 @@ class TmemCorrResource(DecodeGenResourceBase): _cluster_mbarrier: cutlass.Array = None _kv_tile_256_exchange_alloc: Constexpr[SmemAllocation | None] = None _kv_tile_256_exchange: cutlass.Array = None + _sage_v_scales: cutlass.Array = None def get_o_stage_dtype_bytes(self) -> int: """Return the element width used by the final O staging buffer.""" @@ -159,10 +167,20 @@ def get_o_stage_dtype_bytes(self) -> int: ) def _kv_tile_256_exchange_entries(self) -> int: - """Return 128 lane-local stats plus one D32 fragment per output row.""" + """Return the KV256 exchange floats. + + Every lane owns four stats slots, then one padded D32 fragment slot per + logical output row (the row-owner exchange) or per lane (the + column-owner exchange of the byte-wide tail). + """ + fragment_slots = ( + _KV_TILE_256_CORRECTION_THREADS + if self.cfg.splits_kv_tile_256_tail_columns + else _KV_TILE_256_LOGICAL_OUTPUT_ROWS + ) return ( _KV_TILE_256_CORRECTION_THREADS * _KV_TILE_256_STATS_PER_THREAD - + _KV_TILE_256_LOGICAL_OUTPUT_ROWS * _KV_TILE_256_EXCHANGE_FRAGMENT_STRIDE + + fragment_slots * _KV_TILE_256_EXCHANGE_FRAGMENT_STRIDE ) def _init_placeholder_state(self) -> None: @@ -197,6 +215,10 @@ def _init_placeholder_state(self) -> None: Float32, self._kv_tile_256_exchange_entries(), ) + self._sage_v_scales = _placeholder_smem_array( + Float32, + self.sage_v_scales.entries if self.sage_v_scales is not None else 1, + ) def _owns_final_epilogue(self) -> bool: """Whether this correction instance owns final normalization/output.""" @@ -307,12 +329,10 @@ def get_smem_requirements(self) -> list[SmemAllocation]: alignment=8, ) if self.cfg.tile_size_kv == 256 and self._kv_tile_256_exchange_alloc is None: - # Tail correction exchanges all lane-local stats, then pipelines - # D32 fragments through 64 logical output rows one fragment at a - # time. Upper lanes publish one spatial half while lower lanes - # retain the matching fragment in registers. The buffer is - # dedicated, so the shared KV ring keeps streaming the next tile's - # routes while the tail runs. + # Tail correction exchanges all lane-local stats, then D32 + # fragments one at a time (see ``_kv_tile_256_tail_epilogue``). + # The buffer is dedicated, so the shared KV ring keeps streaming + # the next tile's routes while the tail runs. self._kv_tile_256_exchange_alloc = SmemAllocation( name=f"{self.name}_kvTile256Exchange", size_bytes=self._kv_tile_256_exchange_entries() * 4, @@ -571,6 +591,24 @@ def _separate_partial_norm_scale(self, sum_val: Float32) -> Float32: # every normalized partial so the shared reducer only merges states. return self.output_scale * self._safe_norm_rcp(sum_val) + @cute.jit + def _split_partial_scale(self, denominator: Float32) -> Float32: + """Return the scale applied to a split-KV partial O before it is stored. + + The separate reduction kernel receives normalized partials. Fused GMEM + and cluster reductions receive unnormalized partials, with the static + 448 scale of E4M3 P divided out before the narrowing to 16 bits; the + reducer restores it. The 448 scale follows ``pv_dtype``, not the + Q/K dtype. + """ + cfg = self.cfg + partial_scale = Float32(1.0) + if cutlass.const_expr(cfg.use_separate_reduction_kernel): + partial_scale = self._separate_partial_norm_scale(denominator) + elif cutlass.const_expr(cfg.use_fp8_pv): + partial_scale = Float32(1.0 / FP8_P_QUANT_SCALE) + return partial_scale + @cute.jit def _separate_partial_lse(self, max_val: Float32, sum_val: Float32) -> Float32: """Convert one local max/sum pair to the shared log2-LSE contract.""" @@ -671,14 +709,15 @@ def _swaps_load_two_o_stage_chunks( num_o_chunks: Constexpr[int], output_f32_regs: Constexpr[int], ) -> tuple[cutlass.Array, cutlass.Array]: - """Load matching chunks from the two final Swaps O stages.""" + """Load matching chunks from the active final Swaps O stages.""" cfg = self.cfg o0_vals = cutlass.Array( Float32, output_f32_regs, space=cutlass.AddressSpace.rmem ) - o1_vals = cutlass.Array( - Float32, output_f32_regs, space=cutlass.AddressSpace.rmem - ) + if cutlass.const_expr(cfg.num_insts_kv != 1): + o1_vals = cutlass.Array( + Float32, output_f32_regs, space=cutlass.AddressSpace.rmem + ) for chunk_idx in cutlass.range_constexpr(num_o_chunks): o0_loaded = prims.tcgen05_ld( "16x256b", @@ -688,18 +727,23 @@ def _swaps_load_two_o_stage_chunks( ), num=q_repeats, ) - o1_loaded = prims.tcgen05_ld( - "16x256b", - prims.make_tmem_ptr( - base_addr1 + Int32(cfg.swaps_o_chunk_tmem_offset(chunk_idx)), - Float32, - ), - num=q_repeats, - ) + if cutlass.const_expr(cfg.num_insts_kv != 1): + o1_loaded = prims.tcgen05_ld( + "16x256b", + prims.make_tmem_ptr( + base_addr1 + Int32(cfg.swaps_o_chunk_tmem_offset(chunk_idx)), + Float32, + ), + num=q_repeats, + ) for reg_idx in cutlass.range_constexpr(4 * q_repeats): o0_vals[chunk_idx * 4 * q_repeats + reg_idx] = o0_loaded[reg_idx] - o1_vals[chunk_idx * 4 * q_repeats + reg_idx] = o1_loaded[reg_idx] + if cutlass.const_expr(cfg.num_insts_kv != 1): + o1_vals[chunk_idx * 4 * q_repeats + reg_idx] = o1_loaded[reg_idx] cute.arch.fence_view_async_tmem_load() + if cutlass.const_expr(cfg.num_insts_kv == 1): + # Preserve the return shape; one-instance epilogues only consume o0. + return o0_vals, o0_vals return o0_vals, o1_vals @cute.jit @@ -894,12 +938,56 @@ def _merge_kv_tile_256_peer_fragment( merged_vals[elem + 1] = merged[1] return merged_vals + @cute.jit + def _apply_sage_channel_scales( + self, + output_vals: cutlass.Array, + norm_scale: Float32 | None, + *, + first_col: Int32, + count: Constexpr[int], + row_has_mass: cutlass.Boolean, + ) -> Float32 | None: + """Dequantize ``count`` output columns per V channel in place. + + Column ``c`` becomes ``value * norm_scale * sfV[c] + v_mean[c]``, read + from the work tile's staged ``SageVScalesResource`` block, so the + normalization is consumed here and the caller continues with the unit + scale. A row without mass has no V contribution and keeps a zero + mean. ``None`` is the unit scale on both sides: the packed multiply + does not fold a constant one, so it is skipped explicitly. Without + Sage attention the values and ``norm_scale`` pass through. + """ + cfg = self.cfg + if cutlass.const_expr(not cfg.use_sage_attention): + return norm_scale + v_scales, v_means = self.sage_v_scales.channel_scales( + self._sage_v_scales, first_col=first_col, count=count + ) + if cutlass.const_expr(cfg.sage_v_mean): + if not row_has_mass: + for value_idx in cutlass.range_constexpr(count): + v_means[value_idx] = Float32(0.0) + for pair_idx in cutlass.range_constexpr(count // 2): + val_base = pair_idx * 2 + column_scales = (v_scales[val_base], v_scales[val_base + 1]) + if cutlass.const_expr(norm_scale is not None): + column_scales = fmul2((norm_scale, norm_scale), column_scales) + scaled = ffma2( + (output_vals[val_base], output_vals[val_base + 1]), + column_scales, + (v_means[val_base], v_means[val_base + 1]), + ) + output_vals[val_base] = scaled[0] + output_vals[val_base + 1] = scaled[1] + return None + @cute.jit def _store_final_o_columns( self, final_o_dst, output_vals: cutlass.Array, - norm_scale: Float32, + norm_scale: Float32 | None, *, count: Constexpr[int], sector_aligned: cutlass.Boolean, @@ -910,9 +998,10 @@ def _store_final_o_columns( 8-byte store. 16-bit output packs pairs; sixteen columns fill one 32-byte sector and go out as a single 256-bit store when the destination is sector aligned, otherwise every eight columns use one - 16-byte store. Callers choose ``count`` per path: the KV256 tail owns - whole rows per lane and pays for half-written sectors, while the - split-KV reducers write eight-column fragments. + 16-byte store. Callers choose ``count`` per path: the KV256 tail lanes + own whole rows (16-bit K/V) or 64-column half rows (byte-wide K/V) and + pay for half-written sectors, while the split-KV reducers write + eight-column fragments. A ``None`` scale packs the values as they are. """ cfg = self.cfg assert count % 8 == 0 @@ -958,20 +1047,19 @@ def _store_16bit_output_chunks( def _pack_fp8_output_quads( self, output_vals: cutlass.Array, - norm_scale: Float32, + norm_scale: Float32 | None, first_col: Constexpr[int], ) -> cutlass.Array: """Scale eight output columns from ``first_col`` into two FP8 registers.""" final_pairs = cutlass.Array(Float32, 8, space=cutlass.AddressSpace.rmem) for pair_idx in cutlass.range_constexpr(4): val_base = pair_idx * 2 - pair = fmul2( - (norm_scale, norm_scale), - ( - output_vals[first_col + val_base], - output_vals[first_col + val_base + 1], - ), + pair = ( + output_vals[first_col + val_base], + output_vals[first_col + val_base + 1], ) + if cutlass.const_expr(norm_scale is not None): + pair = fmul2((norm_scale, norm_scale), pair) final_pairs[val_base] = pair[0] final_pairs[val_base + 1] = pair[1] fp8_regs = cutlass.Array(Int32, 2, space=cutlass.AddressSpace.rmem) @@ -987,20 +1075,16 @@ def _pack_fp8_output_quads( def _pack_final_o_regs( self, output_vals: cutlass.Array, - norm_scale: Float32, + norm_scale: Float32 | None, count: Constexpr[int], ) -> cutlass.Array: """Scale ``count`` output columns and pack them as 16-bit pairs.""" cfg = self.cfg final_regs = cutlass.Array(Int32, count // 2, space=cutlass.AddressSpace.rmem) for reg_idx in cutlass.range_constexpr(count // 2): - pair = fmul2( - (norm_scale, norm_scale), - ( - output_vals[reg_idx * 2], - output_vals[reg_idx * 2 + 1], - ), - ) + pair = (output_vals[reg_idx * 2], output_vals[reg_idx * 2 + 1]) + if cutlass.const_expr(norm_scale is not None): + pair = fmul2((norm_scale, norm_scale), pair) if cutlass.const_expr(cfg.use_bf16_output): final_regs[reg_idx] = _pack_float2_to_bf16(pair[0], pair[1]) else: @@ -1049,7 +1133,11 @@ def _softmax_output_row_state( sum_val: Float32, max_val: Float32, ) -> tuple[Int64, Float32]: - """Resolve one logical row's output address and normalization once.""" + """Resolve one logical row's output address and normalization once. + + With E4M3 P the O accumulator and the row sum both carry the 448 P + scale, which cancels in a direct tail. + """ 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( @@ -1069,11 +1157,10 @@ def _softmax_output_row_state( # 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) - if cutlass.const_expr(cfg.use_fp8_qkv): - # Since P is scaled to [0, 448] for Fused GMEM/cluster FP8-Q, - # divide the partial O by 448 before narrowing it to 16 bits, - # and restore after the partials have been reduced in FP32. - norm_scale *= Float32(448.0) + if cutlass.const_expr(cfg.use_fp8_pv and cfg.use_split_kv): + # Fused split partials carry O divided by 448 (see + # ``_split_partial_scale``); restore it after the FP32 reduction. + norm_scale *= Float32(FP8_P_QUANT_SCALE) physical_dst_row_idx = _q_physical_output_row_from_logical( cfg, self.h_r, @@ -1243,6 +1330,12 @@ def create_function_variables( shape=(self._kv_tile_256_exchange_entries(),), addrspace=3, ) + if cutlass.const_expr( + context is not None + and context.smem_base is not None + and self.sage_v_scales is not None + ): + self._sage_v_scales = self.sage_v_scales.staged(context) return {} @producer_work(work_attrs=WorkAttr.AUXILIARY) @@ -2365,11 +2458,11 @@ def _kv_tile_256_temporal_fragment( *, base_addr0: Int32, base_addr1: Int32, - fragment_col: Constexpr[int], + fragment_col: Int32, weight00: Float32, weight10: Float32, ): - """Load and combine one D32 fragment from the two temporal stages.""" + """Load one D32 fragment and combine active temporal stages.""" cfg = self.cfg o0_vals = _keeps_tcgen05_ld( cfg, @@ -2402,9 +2495,10 @@ def _store_kv_tile_256_direct_fragment( own_vals: cutlass.Array, peer_vals, *, - fragment_col: Constexpr[int], + fragment_col: Int32, dst_row_base: Int64, norm_scale: Float32, + row_has_mass: cutlass.Boolean, valid_output_row: cutlass.Boolean, o_is_32b_aligned: cutlass.Boolean, ) -> None: @@ -2431,10 +2525,17 @@ def _store_kv_tile_256_direct_fragment( mem_space=1, dtype=Int32, ) + column_norm_scale = self._apply_sage_channel_scales( + merged_vals, + norm_scale, + first_col=Int32(output_col), + count=16, + row_has_mass=row_has_mass, + ) self._store_final_o_columns( final_o_dst, merged_vals, - norm_scale, + column_norm_scale, count=16, sector_aligned=o_is_32b_aligned, ) @@ -2445,7 +2546,7 @@ def _store_kv_tile_256_partial_fragment( own_vals: cutlass.Array, peer_vals, *, - fragment_col: Constexpr[int], + fragment_col: Int32, partial_row_base: Int64, partial_scale: Float32, valid_output_row: cutlass.Boolean, @@ -2489,6 +2590,56 @@ def _store_kv_tile_256_partial_fragment( ) partial_o_dst.store(packed, alignment=16) + @cute.jit + def _kv_tile_256_output_row_state( + self, + *, + resolves: cutlass.Boolean, + logical_h_k_idx: Int32, + logical_b_idx: Int32, + logical_kv_idx: Int32, + cta_idx_kv: Int32, + logical_q_group_idx: Int32, + exchange_row_idx: Int32, + logical_output_row_idx: Int32, + denominator: Float32, + global_max: Float32, + ) -> tuple[cutlass.Boolean, Int64, Float32]: + """Return ``(valid_output_row, row_base, row_scale)`` of a lane's output row. + + Lanes outside ``resolves`` return the defaults. A split-KV row gets its + partial row byte offset and the partial scale, a direct row its + destination byte offset and the normalization scale. + """ + cfg = self.cfg + valid_output_row = cutlass.Boolean(False) + row_base = Int64(0) + row_scale = Float32(1.0) + if resolves: + valid_output_row = _q_row_is_valid_for_seq( + cfg, + self.h_r, + logical_q_group_idx, + exchange_row_idx, + self.seq_len_q, + ) + if cutlass.const_expr(cfg.use_split_kv): + row_scale = self._split_partial_scale(denominator) + row_base = self._gmem_partial_row_offset( + logical_kv_idx, + cta_idx_kv, + logical_output_row_idx, + ) * Int64(cfg.headdim * 2) + else: + row_base, row_scale = self._softmax_output_row_state( + logical_h_k_idx, + logical_b_idx, + logical_output_row_idx, + denominator, + global_max, + ) + return valid_output_row, row_base, row_scale + @cute.jit def _kv_tile_256_merge_spatial_output( self, @@ -2526,47 +2677,21 @@ def _kv_tile_256_merge_spatial_output( _KV_TILE_256_EXCHANGE_FRAGMENT_STRIDE ) logical_output_row_idx = q_row_offset + exchange_row_idx - valid_output_row = cutlass.Boolean(False) - - if cutlass.const_expr(cfg.use_split_kv): - partial_scale = Float32(1.0) - partial_row_base = Int64(0) - if output_lane: - valid_output_row = _q_row_is_valid_for_seq( - cfg, - self.h_r, - logical_q_group_idx, - exchange_row_idx, - self.seq_len_q, - ) - if cutlass.const_expr(cfg.use_separate_reduction_kernel): - partial_scale = self._separate_partial_norm_scale(denominator) - partial_row_base = self._gmem_partial_row_offset( - logical_kv_idx, - cta_idx_kv, - logical_output_row_idx, - ) * Int64(cfg.headdim * 2) - else: - dst_row_base = Int64(0) - norm_scale = Float32(1.0) - # Row strides are multiples of 32 bytes for D128, so sector-wide - # stores are legal exactly when the output base pointer is. - o_is_32b_aligned = (self.o_ptr.toint() & Int64(31)) == Int64(0) - if output_lane: - valid_output_row = _q_row_is_valid_for_seq( - cfg, - self.h_r, - logical_q_group_idx, - exchange_row_idx, - self.seq_len_q, - ) - dst_row_base, norm_scale = self._softmax_output_row_state( - logical_h_k_idx, - logical_b_idx, - logical_output_row_idx, - denominator, - global_max, - ) + valid_output_row, row_base, row_scale = self._kv_tile_256_output_row_state( + resolves=output_lane, + logical_h_k_idx=logical_h_k_idx, + logical_b_idx=logical_b_idx, + logical_kv_idx=logical_kv_idx, + cta_idx_kv=cta_idx_kv, + logical_q_group_idx=logical_q_group_idx, + exchange_row_idx=exchange_row_idx, + logical_output_row_idx=logical_output_row_idx, + denominator=denominator, + global_max=global_max, + ) + # Row strides are multiples of 32 bytes for D128, so sector-wide + # stores are legal exactly when the output base pointer is. + o_is_32b_aligned = (self.o_ptr.toint() & Int64(31)) == Int64(0) for fragment in cutlass.range_constexpr(cfg.headdim // 32): fragment_col = fragment * 32 @@ -2606,8 +2731,9 @@ def _kv_tile_256_merge_spatial_output( own_vals, peer_vals, fragment_col=fragment_col, - dst_row_base=dst_row_base, - norm_scale=norm_scale, + dst_row_base=row_base, + norm_scale=row_scale, + row_has_mass=denominator > Float32(0.0), valid_output_row=valid_output_row, o_is_32b_aligned=o_is_32b_aligned, ) @@ -2616,46 +2742,59 @@ def _kv_tile_256_merge_spatial_output( own_vals, peer_vals, fragment_col=fragment_col, - partial_row_base=partial_row_base, - partial_scale=partial_scale, + partial_row_base=row_base, + partial_scale=row_scale, valid_output_row=valid_output_row, ) if cutlass.const_expr(cfg.use_split_kv): if valid_output_row: - if cutlass.const_expr(cfg.use_separate_reduction_kernel): - stats_row = self._gmem_partial_row_offset( - logical_kv_idx, - cta_idx_kv, - logical_output_row_idx, - ) - stats_ptr = cutlass.inttoptr( - self.partial_stats_ptr.toint() + stats_row * Int64(4), - mem_space=1, - dtype=Float32, - ) - stats_ptr.store( - self._separate_partial_lse(global_max, denominator), - alignment=4, - ) - else: - stats_row = self._gmem_partial_row_offset( - logical_kv_idx, - cta_idx_kv, - logical_output_row_idx, - ) - stats_ptr = cutlass.inttoptr( - self.partial_stats_ptr.toint() + stats_row * Int64(2 * 4), - mem_space=1, - dtype=Float32, - ) - stats_ptr.store( - cutlass.Vector.from_elements( - (global_max, denominator), - Float32, - ), - alignment=8, - ) + self._store_kv_tile_256_partial_stats( + logical_kv_idx, + cta_idx_kv, + logical_output_row_idx, + global_max, + denominator, + ) + + @cute.jit + def _store_kv_tile_256_partial_stats( + self, + logical_kv_idx: Int32, + cta_idx_kv: Int32, + logical_output_row_idx: Int32, + global_max: Float32, + denominator: Float32, + ) -> None: + """Write one logical row's split-KV stats record.""" + stats_row = self._gmem_partial_row_offset( + logical_kv_idx, + cta_idx_kv, + logical_output_row_idx, + ) + if cutlass.const_expr(self.cfg.use_separate_reduction_kernel): + stats_ptr = cutlass.inttoptr( + self.partial_stats_ptr.toint() + stats_row * Int64(4), + mem_space=1, + dtype=Float32, + ) + stats_ptr.store( + self._separate_partial_lse(global_max, denominator), + alignment=4, + ) + else: + stats_ptr = cutlass.inttoptr( + self.partial_stats_ptr.toint() + stats_row * Int64(2 * 4), + mem_space=1, + dtype=Float32, + ) + stats_ptr.store( + cutlass.Vector.from_elements( + (global_max, denominator), + Float32, + ), + alignment=8, + ) @cute.jit def _kv_tile_256_tail_epilogue( @@ -2676,9 +2815,11 @@ def _kv_tile_256_tail_epilogue( The standard decode schedule still owns the two temporal instances. KV256 adds one physical spatial split per instance. Correction exchanges - their stats, stages one spatial half through its dedicated SMEM - exchange one D32 fragment at a time, then publishes the ordinary - logical Q64xD128 output. + their stats, then merges the halves through its dedicated SMEM exchange + one D32 fragment at a time into the logical Q64xD128 output. With + 16-bit K/V the upper lanes publish and the lower lanes store whole + rows (row-owner mode); with byte-wide K/V every lane stores its own 64 + columns (column-owner mode, ``cfg.splits_kv_tile_256_tail_columns``). """ cfg = self.cfg assert cfg.headdim == 128 @@ -2722,29 +2863,40 @@ def _kv_tile_256_tail_epilogue( weight00 = Float32(0.0) weight10 = Float32(0.0) - denominator = Float32(0.0) + weight01 = Float32(0.0) + weight11 = Float32(0.0) if uses00: weight00 = cute.math.exp2( self.scale_softmax_log2 * (max00 - global_max), fastmath=True, ) - denominator += sum00 * weight00 if uses10: weight10 = cute.math.exp2( self.scale_softmax_log2 * (max10 - global_max), fastmath=True, ) - denominator += sum10 * weight10 if uses01: - denominator += sum01 * cute.math.exp2( + weight01 = cute.math.exp2( self.scale_softmax_log2 * (max01 - global_max), fastmath=True, ) if uses11: - denominator += sum11 * cute.math.exp2( + weight11 = cute.math.exp2( self.scale_softmax_log2 * (max11 - global_max), fastmath=True, ) + # Sum in one fixed order per row (lower spatial half first), so both + # column-owner lanes of a row get a bitwise-identical denominator. + own_terms = ((uses00, sum00, weight00), (uses10, sum10, weight10)) + peer_terms = ((uses01, sum01, weight01), (uses11, sum11, weight11)) + own_first = cutlass.Boolean(True) + if cutlass.const_expr(cfg.splits_kv_tile_256_tail_columns): + own_first = exchange_idx < Int32(_KV_TILE_256_LOGICAL_OUTPUT_ROWS) + denominator = Float32(0.0) + if own_first: + denominator = _masked_weighted_sum(own_terms + peer_terms) + else: + denominator = _masked_weighted_sum(peer_terms + own_terms) base_addr0 = tmem_row_base + Int32( o_base_col + tail_o_stage_idx_0 * cfg.tmem_o_stage_cols @@ -2764,7 +2916,12 @@ def _kv_tile_256_tail_epilogue( splits_kv = self._runtime_splits_kv(stage_info) cta_idx_kv = _logical_cta_kv_idx(cfg, stage_info) - self._kv_tile_256_merge_spatial_output( + merge_spatial_halves = ( + self._kv_tile_256_merge_owned_columns + if cfg.splits_kv_tile_256_tail_columns + else self._kv_tile_256_merge_spatial_output + ) + merge_spatial_halves( exchange=exchange, exchange_idx=exchange_idx, base_addr0=base_addr0, @@ -2785,7 +2942,7 @@ def _kv_tile_256_tail_epilogue( cfg.use_split_kv and not cfg.use_separate_reduction_kernel ): # Every correction lane joins the generic fused-GMEM completion - # protocol after the low 64 lanes publish one full D128 row each. + # protocol after every logical row's partial has been published. counter_q_groups = self._multi_cta_counter_q_groups(self.h_r) counter_group_idx = logical_kv_idx * counter_q_groups + logical_q_group_idx self._reduce_fused_gmem_partials( @@ -2799,11 +2956,142 @@ def _kv_tile_256_tail_epilogue( warp_grp_thread_idx, ) - # The next persistent work tile may reuse this exchange allocation. - prims.barrier_cta_sync( - self.store_barrier_id, - thread_count=cfg.correction_barrier_threads, + if cutlass.const_expr(not cfg.splits_kv_tile_256_tail_columns): + # The next persistent work tile may reuse this exchange allocation. + # The column-owner tail closed its last round with this barrier. + prims.barrier_cta_sync( + self.store_barrier_id, + thread_count=cfg.correction_barrier_threads, + ) + + @cute.jit + def _kv_tile_256_merge_owned_columns( + self, + *, + exchange: cutlass.Array, + exchange_idx: Int32, + base_addr0: Int32, + base_addr1: Int32, + logical_h_k_idx: Int32, + logical_b_idx: Int32, + logical_kv_idx: Int32, + cta_idx_kv: Int32, + logical_q_group_idx: Int32, + q_row_offset: Int32, + weight00: Float32, + weight10: Float32, + denominator: Float32, + global_max: Float32, + ) -> None: + """Merge and store the 64 output columns each lane owns, one D32 at a time. + + Lanes below 64 own columns [0, 64) of their logical row, lanes 64..127 + own [64, 128); the peer lane (index xor 64) holds the other spatial + partial of the row. Each round a lane publishes the temporal + combination of the D32 its peer owns, combines its own D32 from TMEM, + adds the peer's published fragment and stores the direct or split-KV + output. + """ + cfg = self.cfg + output_exchange_base = Int32( + _KV_TILE_256_CORRECTION_THREADS * _KV_TILE_256_STATS_PER_THREAD + ) + # Lane bit 6 selects the spatial half and, since the logical row count + # equals headdim / 2, is also the first column of the lane's half. + half_cols = cfg.headdim // 2 + assert half_cols == _KV_TILE_256_LOGICAL_OUTPUT_ROWS + exchange_row_idx = exchange_idx & Int32(_KV_TILE_256_LOGICAL_OUTPUT_ROWS - 1) + peer_idx = exchange_idx ^ Int32(_KV_TILE_256_LOGICAL_OUTPUT_ROWS) + owned_col_base = exchange_idx & Int32(half_cols) + peer_col_base = owned_col_base ^ Int32(half_cols) + own_slot = output_exchange_base + exchange_idx * Int32( + _KV_TILE_256_EXCHANGE_FRAGMENT_STRIDE + ) + peer_slot = output_exchange_base + peer_idx * Int32( + _KV_TILE_256_EXCHANGE_FRAGMENT_STRIDE ) + 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), + logical_h_k_idx=logical_h_k_idx, + logical_b_idx=logical_b_idx, + logical_kv_idx=logical_kv_idx, + cta_idx_kv=cta_idx_kv, + logical_q_group_idx=logical_q_group_idx, + exchange_row_idx=exchange_row_idx, + logical_output_row_idx=logical_output_row_idx, + denominator=denominator, + global_max=global_max, + ) + # Row strides are multiples of 32 bytes for D128, so sector-wide + # stores are legal exactly when the output base pointer is. + o_is_32b_aligned = (self.o_ptr.toint() & Int64(31)) == Int64(0) + + # A rolled loop keeps one copy of the round body in the instruction + # stream, so it does not compete with the softmax warps' rolled + # fragment loops for instruction fetch. + for round_idx in cutlass.range(cfg.headdim // 64, unroll=1): + fragment_offset = Int32(round_idx) * Int32(32) + published_vals = self._kv_tile_256_temporal_fragment( + base_addr0=base_addr0, + base_addr1=base_addr1, + fragment_col=peer_col_base + fragment_offset, + weight00=weight00, + weight10=weight10, + ) + (exchange.data_ptr() + own_slot).store(published_vals, alignment=16) + own_vals = self._kv_tile_256_temporal_fragment( + base_addr0=base_addr0, + base_addr1=base_addr1, + fragment_col=owned_col_base + fragment_offset, + weight00=weight00, + weight10=weight10, + ) + # Every lane's published fragment becomes visible to its peer. + prims.barrier_cta_sync( + self.store_barrier_id, + thread_count=cfg.correction_barrier_threads, + ) + peer_vals = (exchange.data_ptr() + peer_slot).load(count=32, alignment=16) + if cutlass.const_expr(not cfg.use_split_kv): + self._store_kv_tile_256_direct_fragment( + own_vals, + peer_vals, + fragment_col=owned_col_base + fragment_offset, + dst_row_base=row_base, + norm_scale=row_scale, + row_has_mass=denominator > Float32(0.0), + valid_output_row=valid_output_row, + o_is_32b_aligned=o_is_32b_aligned, + ) + else: + self._store_kv_tile_256_partial_fragment( + own_vals, + peer_vals, + fragment_col=owned_col_base + fragment_offset, + partial_row_base=row_base, + partial_scale=row_scale, + valid_output_row=valid_output_row, + ) + # Every lane finishes reading before the next round's publish or + # the next work tile's stats exchange reuses the exchange. + prims.barrier_cta_sync( + self.store_barrier_id, + thread_count=cfg.correction_barrier_threads, + ) + + if cutlass.const_expr(cfg.use_split_kv): + # One stats record per logical row: the lower spatial half writes it. + if valid_output_row and exchange_idx < Int32( + _KV_TILE_256_LOGICAL_OUTPUT_ROWS + ): + self._store_kv_tile_256_partial_stats( + logical_kv_idx, + cta_idx_kv, + logical_output_row_idx, + global_max, + denominator, + ) @cute.jit def _keeps_tail_epilogue( @@ -2860,7 +3148,7 @@ def _keeps_tail_epilogue( ) if cutlass.const_expr( - cfg.use_fp8_qkv + cfg.use_fp8_pv and cfg.use_fp8_output and cfg.tile_size_q == 128 and cfg.headdim == 128 @@ -2934,11 +3222,7 @@ def _keeps_tail_epilogue( if cutlass.const_expr(cfg.num_insts_kv == 1): base_addr1 = base_addr0 partial_dst_col_offset = col_base * Int32(2) - partial_norm_scale = Float32(1.0) - if cutlass.const_expr(cfg.use_separate_reduction_kernel): - partial_norm_scale = self._separate_partial_norm_scale(reduced_sum_0) - elif cutlass.const_expr(cfg.use_fp8_qkv): - partial_norm_scale = Float32(1.0 / 448.0) + partial_norm_scale = self._split_partial_scale(reduced_sum_0) regs_o_chunk = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) partial_o_row_base = self._gmem_partial_row_offset( logical_kv_idx, @@ -2987,7 +3271,7 @@ def _keeps_tail_epilogue( ), ) if cutlass.const_expr( - cfg.use_separate_reduction_kernel or cfg.use_fp8_qkv + cfg.use_separate_reduction_kernel or cfg.use_fp8_pv ): partial_pair = fmul2( (partial_norm_scale, partial_norm_scale), partial_pair @@ -3264,6 +3548,7 @@ def _keeps_tail_epilogue( final_pair1[1], ) else: + final_vals = cutlass.Array(Float32, 8, space=cutlass.AddressSpace.rmem) for chunk_idx in cutlass.range_constexpr(4): reg_base = chunk_idx * 2 if cutlass.const_expr(cfg.num_insts_kv == 1): @@ -3280,14 +3565,17 @@ def _keeps_tail_epilogue( (o0_vals[reg_base], o0_vals[reg_base + 1]), ), ) - if cutlass.const_expr(cfg.use_bf16_output): - regs_o_chunk[chunk_idx] = _pack_float2_to_bf16( - final_pair[0], final_pair[1] - ) - else: - regs_o_chunk[chunk_idx] = _pack_float2_to_fp16( - final_pair[0], final_pair[1] - ) + final_vals[reg_base] = final_pair[0] + final_vals[reg_base + 1] = final_pair[1] + # The per-instance scales above already normalized the row. + self._apply_sage_channel_scales( + final_vals, + None, + first_col=col_base + Int32(chunk_col), + count=8, + row_has_mass=reduced_sum_0 > Float32(0.0), + ) + regs_o_chunk = self._pack_final_o_regs(final_vals, None, 8) dst_ptr = cutlass.inttoptr( self.o_ptr.toint() + dst_row_base @@ -3373,7 +3661,9 @@ def _swaps_wide_tail_epilogue( inst0_max = inst0_new_max_arr[scale_idx] inst1_max = inst1_new_max_arr[scale_idx] uses_inst0 = inst0_max != _neg_max_f32() - uses_inst1 = inst1_max != _neg_max_f32() + uses_inst1 = False + if cutlass.const_expr(cfg.num_insts_kv != 1): + uses_inst1 = inst1_max != _neg_max_f32() final_max_val = _neg_max_f32() if uses_inst0: @@ -3500,18 +3790,29 @@ def _swaps_wide_tail_epilogue( base_addr1 = self._swaps_o_stage_base_addr( tmem_row_base, o_base_col, tail_o_stage_idx_1 ) - o0_vals, o1_vals = self._swaps_load_two_o_stage_chunks( - base_addr0, - base_addr1, - q_repeats=q_repeats, - num_o_chunks=num_o_chunks, - output_f32_regs=output_f32_regs, - ) + if cutlass.const_expr(cfg.num_insts_kv == 1): + o0_vals = self._swaps_load_o_stage_chunks( + base_addr0, + q_repeats=q_repeats, + num_o_chunks=num_o_chunks, + output_f32_regs=output_f32_regs, + ) + o1_vals = o0_vals + else: + o0_vals, o1_vals = self._swaps_load_two_o_stage_chunks( + base_addr0, + base_addr1, + q_repeats=q_repeats, + num_o_chunks=num_o_chunks, + output_f32_regs=output_f32_regs, + ) if cutlass.const_expr(cfg.use_split_kv): # Split-KV tail: separate reduction stores normalized 16-bit O # plus log2-LSE; fused GMEM/cluster retains unnormalized O plus # max/sum state. + # Fused GMEM/CGA divides FP8-Q partials by the E4M3 P quantization + # scale before narrowing; other fused partials remain unnormalized. logical_h_k_idx, logical_b_idx = _logical_head_batch( stage_info, self.h_k_idx, self.b_idx ) @@ -3523,13 +3824,9 @@ def _swaps_wide_tail_epilogue( splits_kv = self._runtime_splits_kv(stage_info) cta_idx_kv = _logical_cta_kv_idx(cfg, stage_info) - if cutlass.const_expr(cfg.use_separate_reduction_kernel or cfg.use_fp8_qkv): + if cutlass.const_expr(cfg.use_separate_reduction_kernel or cfg.use_fp8_pv): for scale_idx in cutlass.range_constexpr(num_scale_groups): - norm_scale = Float32(1.0 / 448.0) - if cutlass.const_expr(cfg.use_separate_reduction_kernel): - norm_scale = self._separate_partial_norm_scale( - reduced_sum[scale_idx] - ) + norm_scale = self._split_partial_scale(reduced_sum[scale_idx]) final_scale0[scale_idx] = norm_scale * exp_scale0[scale_idx] final_scale1[scale_idx] = norm_scale * exp_scale1[scale_idx] @@ -3539,30 +3836,38 @@ def _swaps_wide_tail_epilogue( for pair_idx in cutlass.range_constexpr(output_pair_regs): # Separate reduction includes this split's reciprocal sum; # fused reduction delays normalization until the final merge. + # FP8-P fused reduction removes the 448x P quantization scale + # so the unnormalized numerator fits the 16-bit scratch. scale_base = ((pair_idx % (2 * q_repeats)) // 2) * 2 reg_base = pair_idx * 2 partial_scale0 = ( (final_scale0[scale_base], final_scale0[scale_base + 1]) if cutlass.const_expr( - cfg.use_separate_reduction_kernel or cfg.use_fp8_qkv + cfg.use_separate_reduction_kernel or cfg.use_fp8_pv ) else (exp_scale0[scale_base], exp_scale0[scale_base + 1]) ) partial_scale1 = ( (final_scale1[scale_base], final_scale1[scale_base + 1]) if cutlass.const_expr( - cfg.use_separate_reduction_kernel or cfg.use_fp8_qkv + cfg.use_separate_reduction_kernel or cfg.use_fp8_pv ) else (exp_scale1[scale_base], exp_scale1[scale_base + 1]) ) - partial_pair = ffma2( - partial_scale0, - (o0_vals[reg_base], o0_vals[reg_base + 1]), - fmul2( - partial_scale1, - (o1_vals[reg_base], o1_vals[reg_base + 1]), - ), - ) + if cutlass.const_expr(cfg.num_insts_kv == 1): + partial_pair = fmul2( + partial_scale0, + (o0_vals[reg_base], o0_vals[reg_base + 1]), + ) + else: + partial_pair = ffma2( + partial_scale0, + (o0_vals[reg_base], o0_vals[reg_base + 1]), + fmul2( + partial_scale1, + (o1_vals[reg_base], o1_vals[reg_base + 1]), + ), + ) if cutlass.const_expr(cfg.use_separate_reduction_kernel): regs_partial_o[pair_idx] = self._pack_separate_partial_o_pair( partial_pair[0], partial_pair[1] @@ -3677,28 +3982,40 @@ def _swaps_wide_tail_epilogue( scale_base1 = ((pair_idx1 % (2 * q_repeats)) // 2) * 2 reg_base0 = pair_idx0 * 2 reg_base1 = pair_idx1 * 2 - final_pair0 = ffma2( - (final_scale0[scale_base0], final_scale0[scale_base0 + 1]), - (o0_vals[reg_base0], o0_vals[reg_base0 + 1]), - fmul2( - ( - final_scale1[scale_base0], - final_scale1[scale_base0 + 1], + if cutlass.const_expr(cfg.num_insts_kv == 1): + final_pair0 = fmul2( + (final_scale0[scale_base0], final_scale0[scale_base0 + 1]), + (o0_vals[reg_base0], o0_vals[reg_base0 + 1]), + ) + else: + final_pair0 = ffma2( + (final_scale0[scale_base0], final_scale0[scale_base0 + 1]), + (o0_vals[reg_base0], o0_vals[reg_base0 + 1]), + fmul2( + ( + final_scale1[scale_base0], + final_scale1[scale_base0 + 1], + ), + (o1_vals[reg_base0], o1_vals[reg_base0 + 1]), ), - (o1_vals[reg_base0], o1_vals[reg_base0 + 1]), - ), - ) - final_pair1 = ffma2( - (final_scale0[scale_base1], final_scale0[scale_base1 + 1]), - (o0_vals[reg_base1], o0_vals[reg_base1 + 1]), - fmul2( - ( - final_scale1[scale_base1], - final_scale1[scale_base1 + 1], + ) + if cutlass.const_expr(cfg.num_insts_kv == 1): + final_pair1 = fmul2( + (final_scale0[scale_base1], final_scale0[scale_base1 + 1]), + (o0_vals[reg_base1], o0_vals[reg_base1 + 1]), + ) + else: + final_pair1 = ffma2( + (final_scale0[scale_base1], final_scale0[scale_base1 + 1]), + (o0_vals[reg_base1], o0_vals[reg_base1 + 1]), + fmul2( + ( + final_scale1[scale_base1], + final_scale1[scale_base1 + 1], + ), + (o1_vals[reg_base1], o1_vals[reg_base1 + 1]), ), - (o1_vals[reg_base1], o1_vals[reg_base1 + 1]), - ), - ) + ) # bmm2_scale is already folded into ``final_scale*`` above. regs_o[packed_idx] = _pack_float4_to_fp8_e4m3( final_pair0[0], @@ -3712,17 +4029,23 @@ def _swaps_wide_tail_epilogue( # to the final output dtype. scale_base = ((pair_idx % (2 * q_repeats)) // 2) * 2 reg_base = pair_idx * 2 - final_pair = ffma2( - (final_scale0[scale_base], final_scale0[scale_base + 1]), - (o0_vals[reg_base], o0_vals[reg_base + 1]), - fmul2( - ( - final_scale1[scale_base], - final_scale1[scale_base + 1], + if cutlass.const_expr(cfg.num_insts_kv == 1): + final_pair = fmul2( + (final_scale0[scale_base], final_scale0[scale_base + 1]), + (o0_vals[reg_base], o0_vals[reg_base + 1]), + ) + else: + final_pair = ffma2( + (final_scale0[scale_base], final_scale0[scale_base + 1]), + (o0_vals[reg_base], o0_vals[reg_base + 1]), + fmul2( + ( + final_scale1[scale_base], + final_scale1[scale_base + 1], + ), + (o1_vals[reg_base], o1_vals[reg_base + 1]), ), - (o1_vals[reg_base], o1_vals[reg_base + 1]), - ), - ) + ) if cutlass.const_expr(cfg.use_bf16_output): regs_o[pair_idx] = _pack_float2_to_bf16( final_pair[0], final_pair[1] @@ -3852,7 +4175,9 @@ def _swaps_q8_tail_epilogue( inst0_max = inst0_new_max_arr[scale_idx] inst1_max = inst1_new_max_arr[scale_idx] uses_inst0 = inst0_max != _neg_max_f32() - uses_inst1 = inst1_max != _neg_max_f32() + uses_inst1 = False + if cutlass.const_expr(cfg.num_insts_kv != 1): + uses_inst1 = inst1_max != _neg_max_f32() final_max[scale_idx] = _neg_max_f32() if uses_inst0: final_max[scale_idx] = inst0_max @@ -3876,14 +4201,20 @@ def _swaps_q8_tail_epilogue( final_scale0[scale_idx] = Float32(0.0) final_scale1[scale_idx] = Float32(0.0) - final_sums = ffma2( - (exp_scale0[0], exp_scale0[1]), - (inst0_sum_arr[0], inst0_sum_arr[1]), - fmul2( - (exp_scale1[0], exp_scale1[1]), - (inst1_sum_arr[0], inst1_sum_arr[1]), - ), - ) + if cutlass.const_expr(cfg.num_insts_kv == 1): + final_sums = fmul2( + (exp_scale0[0], exp_scale0[1]), + (inst0_sum_arr[0], inst0_sum_arr[1]), + ) + else: + final_sums = ffma2( + (exp_scale0[0], exp_scale0[1]), + (inst0_sum_arr[0], inst0_sum_arr[1]), + fmul2( + (exp_scale1[0], exp_scale1[1]), + (inst1_sum_arr[0], inst1_sum_arr[1]), + ), + ) for scale_idx in cutlass.range_constexpr(num_scale_groups): final_sum[scale_idx] = final_sums[scale_idx] @@ -3959,19 +4290,30 @@ def _swaps_q8_tail_epilogue( base_addr0 = self._swaps_o_stage_base_addr( tmem_row_base, o_base_col, tail_o_stage_idx_0 ) - base_addr1 = self._swaps_o_stage_base_addr( - tmem_row_base, o_base_col, tail_o_stage_idx_1 - ) + base_addr1 = base_addr0 + if cutlass.const_expr(cfg.num_insts_kv != 1): + base_addr1 = self._swaps_o_stage_base_addr( + tmem_row_base, o_base_col, tail_o_stage_idx_1 + ) output_pair_regs = cfg.num_fp16_output_regs output_f32_regs = output_pair_regs * 2 num_o_chunks = cfg.headdim // 64 - o0_vals, o1_vals = self._swaps_load_two_o_stage_chunks( - base_addr0, - base_addr1, - q_repeats=1, - num_o_chunks=num_o_chunks, - output_f32_regs=output_f32_regs, - ) + if cutlass.const_expr(cfg.num_insts_kv == 1): + o0_vals = self._swaps_load_o_stage_chunks( + base_addr0, + q_repeats=1, + num_o_chunks=num_o_chunks, + output_f32_regs=output_f32_regs, + ) + o1_vals = o0_vals + else: + o0_vals, o1_vals = self._swaps_load_two_o_stage_chunks( + base_addr0, + base_addr1, + q_repeats=1, + num_o_chunks=num_o_chunks, + output_f32_regs=output_f32_regs, + ) if cutlass.const_expr(cfg.use_split_kv): # Publish this CTA's partial output and statistics. Standalone @@ -3988,18 +4330,16 @@ def _swaps_q8_tail_epilogue( splits_kv = self._runtime_splits_kv(stage_info) cta_idx_kv = _logical_cta_kv_idx(cfg, stage_info) - if cutlass.const_expr(cfg.use_separate_reduction_kernel or cfg.use_fp8_qkv): + if cutlass.const_expr(cfg.use_separate_reduction_kernel or cfg.use_fp8_pv): for scale_idx in cutlass.range_constexpr(num_scale_groups): - norm_scale = Float32(1.0 / 448.0) - if cutlass.const_expr(cfg.use_separate_reduction_kernel): - norm_scale = self._separate_partial_norm_scale( - reduced_sum[scale_idx] - ) + norm_scale = self._split_partial_scale(reduced_sum[scale_idx]) final_scale0[scale_idx] = norm_scale * exp_scale0[scale_idx] final_scale1[scale_idx] = norm_scale * exp_scale1[scale_idx] # Store normalized 16-bit O for the standalone reducer, or preserve # unnormalized 16-bit O for fused GMEM/cluster reduction. + # FP8-P fused reduction removes the 448x P quantization scale before + # storing the unnormalized O numerator in the 16-bit scratch. regs_partial_o = cutlass.Array( Int32, cfg.num_fp16_output_regs, @@ -4008,27 +4348,33 @@ def _swaps_q8_tail_epilogue( partial_scale0_pair = ( (final_scale0[0], final_scale0[1]) if cutlass.const_expr( - cfg.use_separate_reduction_kernel or cfg.use_fp8_qkv + cfg.use_separate_reduction_kernel or cfg.use_fp8_pv ) else (exp_scale0[0], exp_scale0[1]) ) partial_scale1_pair = ( (final_scale1[0], final_scale1[1]) if cutlass.const_expr( - cfg.use_separate_reduction_kernel or cfg.use_fp8_qkv + cfg.use_separate_reduction_kernel or cfg.use_fp8_pv ) else (exp_scale1[0], exp_scale1[1]) ) for reg_idx in cutlass.range_constexpr(cfg.num_fp16_output_regs): reg_base = reg_idx * 2 - partial_pair = ffma2( - partial_scale0_pair, - (o0_vals[reg_base], o0_vals[reg_base + 1]), - fmul2( - partial_scale1_pair, - (o1_vals[reg_base], o1_vals[reg_base + 1]), - ), - ) + if cutlass.const_expr(cfg.num_insts_kv == 1): + partial_pair = fmul2( + partial_scale0_pair, + (o0_vals[reg_base], o0_vals[reg_base + 1]), + ) + else: + partial_pair = ffma2( + partial_scale0_pair, + (o0_vals[reg_base], o0_vals[reg_base + 1]), + fmul2( + partial_scale1_pair, + (o1_vals[reg_base], o1_vals[reg_base + 1]), + ), + ) if cutlass.const_expr(cfg.use_separate_reduction_kernel): regs_partial_o[reg_idx] = self._pack_separate_partial_o_pair( partial_pair[0], partial_pair[1] @@ -4140,22 +4486,34 @@ def _swaps_q8_tail_epilogue( pair_idx1 = pair_idx0 + 1 src_idx0 = pair_idx0 * 2 src_idx1 = pair_idx1 * 2 - final_pair0 = ffma2( - (final_scale0[0], final_scale0[1]), - (o0_vals[src_idx0], o0_vals[src_idx0 + 1]), - fmul2( - (final_scale1[0], final_scale1[1]), - (o1_vals[src_idx0], o1_vals[src_idx0 + 1]), - ), - ) - final_pair1 = ffma2( - (final_scale0[0], final_scale0[1]), - (o0_vals[src_idx1], o0_vals[src_idx1 + 1]), - fmul2( - (final_scale1[0], final_scale1[1]), - (o1_vals[src_idx1], o1_vals[src_idx1 + 1]), - ), - ) + if cutlass.const_expr(cfg.num_insts_kv == 1): + final_pair0 = fmul2( + (final_scale0[0], final_scale0[1]), + (o0_vals[src_idx0], o0_vals[src_idx0 + 1]), + ) + else: + final_pair0 = ffma2( + (final_scale0[0], final_scale0[1]), + (o0_vals[src_idx0], o0_vals[src_idx0 + 1]), + fmul2( + (final_scale1[0], final_scale1[1]), + (o1_vals[src_idx0], o1_vals[src_idx0 + 1]), + ), + ) + if cutlass.const_expr(cfg.num_insts_kv == 1): + final_pair1 = fmul2( + (final_scale0[0], final_scale0[1]), + (o0_vals[src_idx1], o0_vals[src_idx1 + 1]), + ) + else: + final_pair1 = ffma2( + (final_scale0[0], final_scale0[1]), + (o0_vals[src_idx1], o0_vals[src_idx1 + 1]), + fmul2( + (final_scale1[0], final_scale1[1]), + (o1_vals[src_idx1], o1_vals[src_idx1 + 1]), + ), + ) regs_o[packed_idx] = _pack_float4_to_fp8_e4m3( final_pair0[0], final_pair0[1], @@ -4165,14 +4523,20 @@ def _swaps_q8_tail_epilogue( else: for reg_idx in cutlass.range_constexpr(cfg.num_fp16_output_regs): reg_base = reg_idx * 2 - final_pair = ffma2( - (final_scale0[0], final_scale0[1]), - (o0_vals[reg_base], o0_vals[reg_base + 1]), - fmul2( - (final_scale1[0], final_scale1[1]), - (o1_vals[reg_base], o1_vals[reg_base + 1]), - ), - ) + if cutlass.const_expr(cfg.num_insts_kv == 1): + final_pair = fmul2( + (final_scale0[0], final_scale0[1]), + (o0_vals[reg_base], o0_vals[reg_base + 1]), + ) + else: + final_pair = ffma2( + (final_scale0[0], final_scale0[1]), + (o0_vals[reg_base], o0_vals[reg_base + 1]), + fmul2( + (final_scale1[0], final_scale1[1]), + (o1_vals[reg_base], o1_vals[reg_base + 1]), + ), + ) if cutlass.const_expr(cfg.use_bf16_output): regs_o[reg_idx] = _pack_float2_to_bf16(final_pair[0], final_pair[1]) else: @@ -4278,17 +4642,17 @@ def correction_loop_epilogue( # arrive. Rescale it in place so later PV waves accumulate in the # updated online-softmax frame. cfg = self.cfg - # Resolve the live TMEM O stage and default SwapsMmaAb column base. The - # KeepsMmaAb path overrides the base because O is allocated by TmemO. + # Resolve the live TMEM O stage from the allocator. This is equivalent + # to the historical fixed two-inst Swaps offset, and also remains + # correct when a one-inst pipeline changes the preceding S/stats rings. task_cache = _decode_gen_task_cache(stage_info) tmem_row_base = task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] - o_base_col = 2 * cfg.tmem_s_cols + 2 * cfg.tmem_stats_cols + o_base_col = self.tmem_o_ref._alloc.offset if cutlass.const_expr(cfg.use_keeps_mma_ab): # KeepsMmaAb has one softmax scale group for this path. Compute the # rescale once, then apply it independently to every P-by-V # head-dimension slice in TMEM. - o_base_col = self.tmem_o_ref._alloc.offset corr_chunk_regs = cfg.keeps_loop_correction_chunk_regs old_max_0 = old_max_arr[0] @@ -4347,11 +4711,10 @@ def correction_loop_epilogue( ), offset=keeps_o_ldst_offset, ) - prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) - if skip_correction: - # Preserve the correction task's TMEM ordering when this warp - # issues no correction transaction. - prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) + # The chunks store to disjoint TMEM columns, so one wait retires + # them all before the O stage returns to PV; it also keeps the + # task's TMEM ordering when this warp issued no correction store. + prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) return if cutlass.const_expr(cfg.tile_size_q == 16): @@ -4515,9 +4878,12 @@ def _correction_tail_epilogue_impl( # private resource attributes across a persistent work-tile loop. task_cache = _decode_gen_task_cache(stage_info) tmem_row_base = task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] - o_base_col = 2 * cfg.tmem_s_cols + 2 * cfg.tmem_stats_cols + o_base_col = self.tmem_o_ref._alloc.offset warp_grp_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] + if cutlass.const_expr(not self._owns_final_epilogue()): + return + if cutlass.const_expr(cfg.use_keeps_mma_ab): # KeepsMmaAb finalization is owned by the last active K/V instance. # Earlier instances publish stats but do not write the final output. @@ -4527,7 +4893,6 @@ def _correction_tail_epilogue_impl( ): # Derive the per-lane output row/column ownership before # entering the common Keeps tail helper. - o_base_col = self.tmem_o_ref._alloc.offset output_f32_regs = cfg.keeps_output_f32_regs output_pair_regs = output_f32_regs // 2 keeps_o_ldst_offset = cfg.headdim // 2 @@ -4566,9 +4931,9 @@ def _correction_tail_epilogue_impl( ) return - if cutlass.const_expr(self.inst_id != 1): - # SwapsMmaAb tail uses instance 1 to combine inst0/inst1 final O - # stages. Instance 0 exits after publishing its stats. + if cutlass.const_expr(not self._owns_final_epilogue()): + # Two-inst Swaps uses instance 1 for the final merge; one-inst + # Swaps is finalized directly by its sole correction instance. return if cutlass.const_expr(cfg.tile_size_q in (16, 32)): 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 fa8928695a98..e75003484df3 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 @@ -26,7 +26,6 @@ import cutlass.cute as cute 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 SmemAllocation, TmemAllocation from cutlass.experimental.task_scheduling.resources import ( @@ -36,18 +35,22 @@ producer_work, ) -from ..fmha_decode_constants import KV_INST0 from ..fmha_decode_config import FmhaDecodeConfig +from ..fmha_decode_constants import ( + KV_INST0, + SMEM_DESC_UNIT_BYTES, + SWIZZLE_128B_ROW_BYTES, +) from ...tcgen05_compat import tcgen05_mma_ws from .helpers_common import ( + _TASK_CACHE_TMEM_BASE_OFFSET, Constexpr, DecodeGenResourceBase, DescriptorValue, - _TASK_CACHE_TMEM_BASE_OFFSET, _decode_gen_task_cache, _freeze_smem_descriptor, - _mma_k_step, - _mma_kind_for_qkv, + _mma_k_step_pv, + _mma_kind_for_pv, ) @@ -65,6 +68,9 @@ def _pv_mma_operand_contract_for_config( # halves after the two temporal decode streams are complete. return True, cfg.tile_size_q, cfg.tile_size_kv, 0, 1 return True, cfg.tile_size_q, active_head_dim, 0, 1 + if cfg.store_transformed_kv_in_tmem: + # Transposed TMEM V is operand A in K-major layout; P remains SMEM B. + return False, active_head_dim, cfg.tile_size_q, 0, 0 return False, active_head_dim, cfg.tile_size_q, 1, 0 @@ -234,13 +240,16 @@ def _vp_mma_fragment( fragment_idx: Constexpr[int], initial_scale_d, ) -> None: - """Issue the two MMA K-steps covered by one streamed P fragment. + """Issue the MMA K-steps covered by one streamed P fragment. ``p_tmem_addr`` is already the base of the fragment selected by - ``wait_p_fragment``. Only the two local K-step offsets are added here; - ``fragment_idx`` must not be applied to the TMEM address again. KV256 - issues the WS 2x2 instruction over its two spatial halves; KV128 issues - the plain M=128 instruction and advances V by one K16 slice per step. + ``wait_p_fragment``. Only the local K-step offsets are added here; + ``fragment_idx`` must not be applied to the TMEM address again. A K32 + fragment is two K16 steps for 16-bit operands and one K32 step for + FP8. KV256 issues the WS 2x2 instruction over its two spatial halves, + whose V is staged as KV64 blocks; KV128 issues the plain M=128 + instruction over one contiguous V tile. Every K step spans + ``mma_k_step`` swizzled 128-byte V rows. """ cfg = self.cfg assert cfg.streams_tmem_p_fragments @@ -254,31 +263,42 @@ def _vp_mma_fragment( Float32, ) _, mma_m, mma_n, a_major, b_major = _pv_mma_operand_contract_for_config(cfg) + # P and V use the effective PV precision after any KV transformation. idesc = prims.Tcgen05InstrDesc.build( c_dtype=Float32, - a_dtype=cfg.q_dtype, - b_dtype=cfg.q_dtype, + a_dtype=cfg.pv_dtype, + b_dtype=cfg.pv_dtype, a_major=a_major, b_major=b_major, n_dim=mma_n, m_dim=mma_m, ) - first_k_step = fragment_idx * 2 + mma_k_step = _mma_k_step_pv(cfg) + k_steps_per_fragment = cfg.pv_mma_steps_per_fragment + first_k_step = fragment_idx * k_steps_per_fragment + # One K step covers ``mma_k_step`` swizzled V rows. + k_step_desc_units = mma_k_step * SWIZZLE_128B_ROW_BYTES // SMEM_DESC_UNIT_BYTES + p_cols_per_k_step = cfg.p_cols_per_mma_k_step if prims.elect_sync(): - for local_k_step in cutlass.range_constexpr(2): + for local_k_step in cutlass.range_constexpr(k_steps_per_fragment): k_step = first_k_step + local_k_step p_operand = prims.make_tmem_ptr( - p_tmem_addr + Int32(local_k_step * 8), Int32 + p_tmem_addr + Int32(local_k_step * p_cols_per_k_step), Int32 ) scale_d = initial_scale_d or fragment_idx != 0 or local_k_step != 0 if cutlass.const_expr(cfg.uses_ws_2x2_datapath): - # V holds four K64 atoms; jump between atoms every four - # K16 steps. + # V holds four KV64 blocks in two spatial pairs; jump to + # the next pair once the K steps have covered 64 tokens. + k_steps_per_block = 64 // mma_k_step + block_pair_desc_units = ( + 2 * 64 * cfg.headdim * cfg.v_dtype_bytes // SMEM_DESC_UNIT_BYTES + ) iter_v_desc = v_desc + Int32( - (k_step // 4) * cfg.headdim * 16 + (k_step % 4) * 128 + (k_step // k_steps_per_block) * block_pair_desc_units + + (k_step % k_steps_per_block) * k_step_desc_units ) tcgen05_mma_ws( - _mma_kind_for_qkv(cfg), + _mma_kind_for_pv(cfg), tmem_col, p_operand, iter_v_desc, @@ -286,9 +306,9 @@ def _vp_mma_fragment( scale_d, ) else: - iter_v_desc = v_desc + Int32(k_step * 128) + iter_v_desc = v_desc + Int32(k_step * k_step_desc_units) prims.tcgen05_mma( - _mma_kind_for_qkv(cfg), + _mma_kind_for_pv(cfg), prims.CTAGroup.CTA_1, tmem_col, p_operand, @@ -328,7 +348,10 @@ def _vp_mma( v_desc = v_desc_1 p_desc = p_desc_1 p_tmem_addr = p_tmem_addr_1 - v_desc = _freeze_smem_descriptor(v_desc) + if cutlass.const_expr(cfg.store_transformed_kv_in_tmem): + v_operand = prims.make_tmem_ptr(v_desc, Int32) + else: + v_operand = _freeze_smem_descriptor(v_desc) if cutlass.const_expr(not cfg.uses_tmem_p): p_desc = _freeze_smem_descriptor(p_desc) task_cache = _decode_gen_task_cache(stage_info) @@ -349,8 +372,8 @@ def _vp_mma( ) idesc = prims.Tcgen05InstrDesc.build( c_dtype=Float32, - a_dtype=cfg.q_dtype, - b_dtype=cfg.q_dtype, + a_dtype=cfg.pv_dtype, + b_dtype=cfg.pv_dtype, a_major=a_major, b_major=b_major, n_dim=mma_n, @@ -361,7 +384,7 @@ def _vp_mma( # Accumulate into O after the first K/V tile for this output # stage; the first wave overwrites the TMEM O stage. scale_d = initial_scale_d - pv_k_steps = cfg.tile_size_kv // _mma_k_step(cfg) + pv_k_steps = cfg.tile_size_kv // _mma_k_step_pv(cfg) for ki in cutlass.range_constexpr(pv_k_steps): # Keeps computes P x V (A=P, B=V); Swaps computes the # transposed V^T x P^T tile (A=V, B=P). @@ -369,19 +392,18 @@ def _vp_mma( # TMEM P stores two 16-bit values per column (or four # FP8 values), so each 16-wide MMA-K step advances by # the corresponding packed-column count. - p_cols_per_k_step = _mma_k_step(cfg) * cfg.q_dtype_bytes // 4 p_operand = prims.make_tmem_ptr( - p_tmem_addr + Int32(ki * p_cols_per_k_step), + p_tmem_addr + Int32(ki * cfg.p_cols_per_mma_k_step), Int32, ) else: p_operand = p_desc if cutlass.const_expr(p_is_a): - a_desc, b_desc = p_operand, v_desc + a_desc, b_desc = p_operand, v_operand else: - a_desc, b_desc = v_desc, p_operand + a_desc, b_desc = v_operand, p_operand prims.tcgen05_mma( - _mma_kind_for_qkv(cfg), + _mma_kind_for_pv(cfg), prims.CTAGroup.CTA_1, tmem_col, a_desc, @@ -394,12 +416,18 @@ def _vp_mma( # Advance V and P descriptors to the next MMA-K # slice, including the 16-bit 128-token jump across # split SMEM rows. - v_desc = v_desc + Int32( - (cfg.headdim * 2) if cfg.use_fp8_qkv else 128 - ) + if cutlass.const_expr(cfg.store_transformed_kv_in_tmem): + v_operand = prims.make_tmem_ptr( + v_desc + Int32((ki + 1) * _mma_k_step_pv(cfg) // 4), + Int32, + ) + else: + v_operand = v_operand + Int32( + (cfg.headdim * 2) if cfg.use_fp8_pv else 128 + ) if cutlass.const_expr(not cfg.uses_tmem_p): if cutlass.const_expr( - not cfg.use_fp8_qkv + not cfg.use_fp8_pv and cfg.tile_size_kv == 128 and ki == 3 ): @@ -429,8 +457,8 @@ def _vp_mma( ) idesc = prims.Tcgen05InstrDesc.build( c_dtype=Float32, - a_dtype=cfg.q_dtype, - b_dtype=cfg.q_dtype, + a_dtype=cfg.pv_dtype, + b_dtype=cfg.pv_dtype, a_major=a_major, b_major=b_major, n_dim=mma_n, @@ -441,21 +469,22 @@ def _vp_mma( # Issue one MMA per K step. The first wave may overwrite the # slice, while later loop/tail waves accumulate into it. scale_d = initial_scale_d - for ki in cutlass.range_constexpr(cfg.tile_size_kv // _mma_k_step(cfg)): + for ki in cutlass.range_constexpr( + cfg.tile_size_kv // _mma_k_step_pv(cfg) + ): if cutlass.const_expr(cfg.uses_tmem_p): - p_cols_per_k_step = _mma_k_step(cfg) * cfg.q_dtype_bytes // 4 p_operand = prims.make_tmem_ptr( - p_tmem_addr + Int32(ki * p_cols_per_k_step), + p_tmem_addr + Int32(ki * cfg.p_cols_per_mma_k_step), Int32, ) else: p_operand = p_desc if cutlass.const_expr(p_is_a): - a_desc, b_desc = p_operand, v_desc + a_desc, b_desc = p_operand, v_operand else: - a_desc, b_desc = v_desc, p_operand + a_desc, b_desc = v_operand, p_operand prims.tcgen05_mma( - _mma_kind_for_qkv(cfg), + _mma_kind_for_pv(cfg), prims.CTAGroup.CTA_1, tmem_col, a_desc, @@ -465,16 +494,22 @@ def _vp_mma( ) scale_d = True if cutlass.const_expr( - ki + 1 < cfg.tile_size_kv // _mma_k_step(cfg) + ki + 1 < cfg.tile_size_kv // _mma_k_step_pv(cfg) ): # Advance V and P to the next MMA-K slice inside the # staged head-dim tile. - v_desc = v_desc + Int32( - (cfg.head_dim_kv_stage * 2) if cfg.use_fp8_qkv else 128 - ) + if cutlass.const_expr(cfg.store_transformed_kv_in_tmem): + v_operand = prims.make_tmem_ptr( + v_desc + Int32((ki + 1) * _mma_k_step_pv(cfg) // 4), + Int32, + ) + else: + v_operand = v_operand + Int32( + (cfg.head_dim_kv_stage * 2) if cfg.use_fp8_pv else 128 + ) if cutlass.const_expr(not cfg.uses_tmem_p): if cutlass.const_expr( - not cfg.use_fp8_qkv + not cfg.use_fp8_pv and cfg.tile_size_kv == 128 and ki == 3 ): 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 b010dfdb9d21..78317931721a 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 Boolean, Float32, Int32, Uint32 +from cutlass import BFloat16, Boolean, Float32, Int32, Uint32 from cutlass.experimental import primitives as prims from cutlass.experimental.task_scheduling.enums import WorkAttr @@ -43,17 +43,20 @@ from ...._block_sparse.prepared import _PREPARED_ROUTE_IS_FULL_FLAG from ..fmha_decode_config import CAUSAL, FmhaDecodeConfig -from ..fmha_decode_constants import SOFTMAX_RESCALE_THRESHOLD_LOG2 +from ..fmha_decode_constants import ( + INT32_SCORE_BIAS, + INT32_SCORE_SEED_MMA_K, + INT32_SCORE_SEED_TILE_LBO, + INT32_SCORE_SEED_TILE_SBO, + INT32_SCORE_SEED_TILE_WORDS, + SOFTMAX_RESCALE_THRESHOLD_LOG2, +) from ...tcgen05_compat import tcgen05_mma_ws from ...placeholder_helpers import ( _placeholder_local_array, _placeholder_smem_array, ) from .helpers_common import ( - Constexpr, - DecodeGenResourceBase, - ResourceVars, - ffma2, _TASK_CACHE_LANE_IDX, _TASK_CACHE_KV_RAW_TILE_BASE, _TASK_CACHE_KV_VALID_TILE_END, @@ -61,6 +64,10 @@ _TASK_CACHE_TMEM_BASE_OFFSET, _TASK_CACHE_WARP_GRP_THREAD_IDX, _TASK_CACHE_WARP_IDX, + Constexpr, + DecodeGenResourceBase, + DescriptorValue, + ResourceVars, _clamp_valid_tile_idx, _decode_gen_task_cache, _freeze_smem_descriptor, @@ -68,24 +75,31 @@ _keeps_col_base, _keeps_row_idx, _keeps_score_col, + _keeps_spatial_half, _keeps_tcgen05_ld, _keeps_tcgen05_st, + _logical_head_batch, _logical_q_group_idx, - _mma_k_step, - _mma_kind_for_qkv, + _mma_k_step_qk, + _mma_kind_for_qk, _neg_max_f32, + _qk_accumulator_dtype, _softmax_scale_pair_width, _swaps_routed_coordinate, _q_row_is_valid_for_seq, _q_row_token_and_local_head, _q_group_token_base, + _route_is_proxy, _softmax_tile_idx, + ffma2, + fmul2, ) from .smem_block_sparse_metadata import ( - _SOFTMAX_ROUTE_IS_PROXY_FLAG, _swaps_forwards_packed_route_full, ) +from .sage_scales import SageKScalesResource, SageScaleTensors, load_q_scale from .helpers_kv_tile_idx import ( + resolve_keeps_tile_context, _kv_tile_is_fully_unmasked_for_q_group, _load_runtime_seq_len_kv, _num_skipped_kv_tiles, @@ -236,6 +250,12 @@ class TmemSResource(DecodeGenResourceBase): "Current softmax anchor (normally the running row maximum).", ), ("s_arr", cutlass.Array, None, "Loaded S scores for the current tile."), + ( + "sage_q_scale", + Float32, + Float32(1.0), + "Sage Q scale of the lane's row, loaded once per work tile.", + ), ) inst_id: Constexpr[int] = 0 cfg: Constexpr[FmhaDecodeConfig] = None @@ -244,18 +264,32 @@ class TmemSResource(DecodeGenResourceBase): max_seq_len_kv: Int32 = None seq_len_q: Int32 = None h_r: Int32 | None = None + h_k_idx: Int32 | None = None + b_idx: Int32 | None = None q_group_idx: Int32 | None = None + scale_tensors: SageScaleTensors | None = None q_ref: Constexpr[MemoryResource | None] = None + page_offsets_ref: Constexpr[MemoryResource | None] = None _p_local_sum_arr: cutlass.Array | None = None _global_sum_arr: cutlass.Array | None = None _alloc: Constexpr[TmemAllocation | None] = None sync_barrier_id: Constexpr[int] = 0 _scratch_alloc: Constexpr[SmemAllocation | None] = None _softmax_scratch_u32: cutlass.Array = None + # Constant BF16 operand tile of the MMA step that seeds INT32 scores. One + # instance owns the tile; the other reads it through this reference. + score_seed_owner: Constexpr["TmemSResource | None"] = None + _seed_alloc: Constexpr[SmemAllocation | None] = None + # The instance's ``sfK`` resources. Proxy routes read + # ``sage_summary_k_scales``, which is ``sage_k_scales`` unless the summary + # scales have their own K block size. + sage_k_scales: SageKScalesResource | None = None + sage_summary_k_scales: SageKScalesResource | None = None old_max_arr: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() sum_arr: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() new_max_arr: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() s_arr: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + sage_q_scale: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() def _init_placeholder_state(self) -> None: """Create placeholder register and scratch state for softmax.""" @@ -305,7 +339,71 @@ def get_smem_requirements(self) -> list[SmemAllocation]: size_bytes=scratch_entries * 4, alignment=16, ) - return [self._scratch_alloc] + allocations = [self._scratch_alloc] + if self.cfg.uses_int32_scores and self.score_seed_owner is None: + if self._seed_alloc is None: + self._seed_alloc = SmemAllocation( + name=f"{self.name}_scoreSeed", + size_bytes=self._seed_mma_tile_bytes(), + alignment=128, + ) + allocations.append(self._seed_alloc) + 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. + """ + _, mma_m, mma_n = _qk_mma_operand_contract_for_config(self.cfg) + return max(mma_m, mma_n) * INT32_SCORE_SEED_MMA_K * 2 + + @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.""" + 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, + ) + + @cute.jit + def create_function_variables( + self, context: ResourceContext | None = None + ) -> ResourceVars: + """Write the score-seed operand tile; the prologue barrier publishes it.""" + if cutlass.const_expr( + self.cfg.uses_int32_scores + and context is not None + and context.smem_base is not None + ): + self.fill_score_seed_tiles(context) + return {} + + @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. + """ + 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() def get_tmem_requirements(self) -> list[TmemAllocation]: """Allocate TMEM S score columns for QK MMA output.""" @@ -359,7 +457,7 @@ def _advance_qk_descs_after_mma_k( MMA-K slice. """ cfg = self.cfg - if cutlass.const_expr(not cfg.use_fp8_qkv and crosses_64b_chunk): + if cutlass.const_expr(cfg.qk_dtype_bytes != 1 and crosses_64b_chunk): k_desc = k_desc + Int32(8 * cfg.tile_size_kv - 6) if cutlass.const_expr(cfg.tile_size_q >= 16): q_desc = q_desc + Int32(8 * cfg.tile_size_q - 6) @@ -518,6 +616,81 @@ def init_softmax_state( result["s_arr"], ) + @cute.jit + def _qk_tmem_col(self, stage_info: StageInfo, stage_slot_offset: Int32): + """Return the TMEM pointer of one producer S slot. + + cutlass's tcgen05_alloc returns the base in tmem_ptr_i32, so add the + per-resource column offset before issuing MMA. + """ + task_cache = _decode_gen_task_cache(stage_info) + return prims.make_tmem_ptr( + task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + + Int32(self._alloc.offset) + + stage_slot_offset, + Float32, + ) + + @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. + """ + 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, + ) + + @producer_work + @cute.jit + def seed_scores_head(self, stage_info: StageInfo) -> None: + """Seed the initial S slot before HEAD QK awaits its K tile.""" + self._seed_scores(stage_info, self._qk_head_stage_slot_offset(stage_info)) + + @producer_work + @cute.jit + def seed_scores_loop(self, stage_info: StageInfo) -> None: + """Seed the next producer S slot before LOOP QK awaits its K tile.""" + self._seed_scores(stage_info, self._qk_loop_stage_slot_offset(stage_info)) + @producer_work @cute.jit def qk_mma_head( @@ -525,7 +698,7 @@ def qk_mma_head( stage_info: StageInfo, *, q_desc: prims.Tcgen05SmemDesc, - kv_desc: prims.Tcgen05SmemDesc, + kv_desc: DescriptorValue, head_dim_stage_idx: Constexpr[int], ) -> None: """Issue HEAD QK MMA into the initial S slot.""" @@ -546,7 +719,7 @@ def qk_mma_loop( stage_info: StageInfo, *, q_desc: prims.Tcgen05SmemDesc, - kv_desc: prims.Tcgen05SmemDesc, + kv_desc: DescriptorValue, head_dim_stage_idx: Constexpr[int], ) -> None: """Issue LOOP QK MMA into the next producer S slot.""" @@ -566,7 +739,7 @@ def qk_mma_head_from_q_ref( self, stage_info: StageInfo, *, - kv_desc: prims.Tcgen05SmemDesc, + kv_desc: DescriptorValue, head_dim_stage_idx: Constexpr[int], ) -> None: """Issue guarded persistent HEAD QK without a routed Q descriptor.""" @@ -585,7 +758,7 @@ def qk_mma_loop_from_q_ref( self, stage_info: StageInfo, *, - kv_desc: prims.Tcgen05SmemDesc, + kv_desc: DescriptorValue, head_dim_stage_idx: Constexpr[int], ) -> None: """Issue guarded persistent LOOP QK without a routed Q descriptor.""" @@ -604,7 +777,7 @@ def _qk_mma( stage_info: StageInfo, *, q_desc: prims.Tcgen05SmemDesc, - kv_desc: prims.Tcgen05SmemDesc, + kv_desc: DescriptorValue, stage_slot_offset: Int32, head_dim_stage_idx: Constexpr[int], ) -> None: @@ -615,46 +788,43 @@ def _qk_mma( local descriptor helpers below. """ cfg = self.cfg - k_desc = _freeze_smem_descriptor(kv_desc) + if cutlass.const_expr(cfg.store_transformed_kv_in_tmem): + k_desc = prims.make_tmem_ptr(kv_desc, Int32) + else: + k_desc = _freeze_smem_descriptor(kv_desc) q_desc = _freeze_smem_descriptor(q_desc) q_desc, head_dim_stage_idx = self._q_desc_for_head_dim_stage( q_desc, head_dim_stage_idx ) - # TMEM destination: addrspace-6 pointer from base + alloc offset. - # cutlass's tcgen05_alloc returns the base in tmem_ptr_i32, so add the - # per-resource column offset before issuing MMA. - task_cache = _decode_gen_task_cache(stage_info) - tmem_col = prims.make_tmem_ptr( - task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] - + Int32(self._alloc.offset) - + stage_slot_offset, - Float32, - ) + tmem_col = self._qk_tmem_col(stage_info, stage_slot_offset) q_is_a, mma_m, mma_n = _qk_mma_operand_contract_for_config(cfg) + # Q and K use the effective QK precision after any KV transformation. idesc = prims.Tcgen05InstrDesc.build( - c_dtype=Float32, - a_dtype=cfg.q_dtype, - b_dtype=cfg.q_dtype, + c_dtype=_qk_accumulator_dtype(cfg), + a_dtype=cfg.qk_dtype, + b_dtype=cfg.qk_dtype, n_dim=mma_n, m_dim=mma_m, ) if cutlass.const_expr(cfg.head_dim_per_stage_kv == 0): if prims.elect_sync(): - scale_d = False - for ki in cutlass.range_constexpr(cfg.headdim // _mma_k_step(cfg)): + # INT32 scores accumulate onto the seeded bias; FP32 scores + # overwrite S with the first slice. + scale_d = cfg.uses_int32_scores + for ki in cutlass.range_constexpr(cfg.headdim // _mma_k_step_qk(cfg)): # Keeps computes Q x K^T (A=Q, B=K); Swaps computes the - # transposed K x Q^T tile (A=K, B=Q). The first - # instruction overwrites S and later slices accumulate. + # transposed K x Q^T tile (A=K, B=Q). Later slices + # accumulate. if cutlass.const_expr(q_is_a): a_desc, b_desc = q_desc, k_desc else: a_desc, b_desc = k_desc, q_desc if cutlass.const_expr(cfg.uses_ws_2x2_datapath): tcgen05_mma_ws( - _mma_kind_for_qkv(cfg), + _mma_kind_for_qk(cfg), tmem_col, a_desc, b_desc, @@ -663,7 +833,7 @@ def _qk_mma( ) else: prims.tcgen05_mma( - _mma_kind_for_qkv(cfg), + _mma_kind_for_qk(cfg), prims.CTAGroup.CTA_1, tmem_col, a_desc, @@ -672,14 +842,24 @@ def _qk_mma( scale_d, ) scale_d = True - if cutlass.const_expr(ki + 1 < cfg.headdim // _mma_k_step(cfg)): - k_desc, q_desc = self._advance_qk_descs_after_mma_k( - k_desc, - q_desc, - crosses_64b_chunk=cfg.headdim == 128 and ki == 3, - ) + if cutlass.const_expr(ki + 1 < cfg.headdim // _mma_k_step_qk(cfg)): + if cutlass.const_expr(cfg.store_transformed_kv_in_tmem): + k_desc = prims.make_tmem_ptr( + kv_desc + Int32((ki + 1) * _mma_k_step_qk(cfg) // 4), + Int32, + ) + q_desc = q_desc + Int32(2) + else: + k_desc, q_desc = self._advance_qk_descs_after_mma_k( + k_desc, + q_desc, + crosses_64b_chunk=cfg.headdim == 128 and ki == 3, + ) else: - mma_k_steps = cfg.head_dim_kv_stage // _mma_k_step(cfg) + 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) if prims.elect_sync(): # Peel the first MMA so overwrite-vs-accumulate remains a # compile-time value rather than loop-carried state. @@ -688,7 +868,7 @@ def _qk_mma( else: first_a_desc, first_b_desc = k_desc, q_desc prims.tcgen05_mma( - _mma_kind_for_qkv(cfg), + _mma_kind_for_qk(cfg), prims.CTAGroup.CTA_1, tmem_col, first_a_desc, @@ -703,11 +883,17 @@ def _qk_mma( # closed form therefore adds each boundary jump minus that +2. # Keeping descriptors out of iter_args avoids staged-D256 spills. for ki in cutlass.range(1, mma_k_steps, 1, unroll=1): - if cutlass.const_expr(cfg.use_fp8_qkv): + if cutlass.const_expr(cfg.store_transformed_kv_in_tmem): + iter_k_desc = prims.make_tmem_ptr( + kv_desc + ki * Int32(_mma_k_step_qk(cfg) // 4), Int32 + ) + q_desc_offset = ki * Int32(2) + elif cutlass.const_expr(cfg.qk_dtype_bytes == 1): k_desc_offset = ki * Int32(2) q_desc_offset = ki * Int32(2) + iter_k_desc = k_desc + k_desc_offset else: - chunk_idx = (ki * Int32(_mma_k_step(cfg))) // Int32(64) + chunk_idx = (ki * Int32(_mma_k_step_qk(cfg))) // Int32(64) k_desc_offset = ki * Int32(2) + chunk_idx * Int32(1016) q_chunk_extra = ( 8 * cfg.tile_size_q - 8 @@ -715,7 +901,7 @@ def _qk_mma( else 56 ) q_desc_offset = ki * Int32(2) + chunk_idx * Int32(q_chunk_extra) - iter_k_desc = k_desc + k_desc_offset + iter_k_desc = k_desc + k_desc_offset iter_q_desc = q_desc + q_desc_offset if prims.elect_sync(): if cutlass.const_expr(q_is_a): @@ -723,7 +909,7 @@ def _qk_mma( else: iter_a_desc, iter_b_desc = iter_k_desc, iter_q_desc prims.tcgen05_mma( - _mma_kind_for_qkv(cfg), + _mma_kind_for_qk(cfg), prims.CTAGroup.CTA_1, tmem_col, iter_a_desc, @@ -735,99 +921,14 @@ def _qk_mma( @cute.jit def _resolve_keeps_tile_context(self, stage_info: StageInfo): """Resolve one score tile's logical position and boundary-mask state.""" - cfg = self.cfg - task_cache = _decode_gen_task_cache(stage_info) - if cutlass.const_expr(self.seqlens_kv is None): - seq_len_kv = Int32(self.max_seq_len_kv) - else: - seq_len_kv = _load_runtime_seq_len_kv( - self.seqlens_kv, - self.max_seq_len_kv, - stage_info, - Int32(0), - Int32(0), - ) - - logical_q_group_idx = _logical_q_group_idx(cfg, stage_info, self.q_group_idx) - q_token_base = _q_group_token_base(cfg, logical_q_group_idx) - element_mask_end_idx = seq_len_kv - if cutlass.const_expr(cfg.uses_uniform_causal_mask): - element_mask_end_idx = seq_len_kv - self.seq_len_q + q_token_base + Int32(1) - - use_runtime_kv_domain = ( - self.seqlens_kv is not None 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): - effective_tile_idx = _static_split_kv_global_tile_idx( - cfg, stage_info, local_tile_idx - ) - effective_total_kv_tiles = Int32(cfg.total_kv_tiles) - tile_idx = _clamp_valid_tile_idx(cfg, effective_tile_idx) - tile_idx = tile_idx + Int32(cfg.static_num_skipped_kv_tiles) - window_start_idx = Int32(cfg.static_window_start_idx) - elif cutlass.const_expr(cfg.use_paged_kv and not cfg.use_split_kv): - effective_tile_idx = ( - Int32(task_cache[_TASK_CACHE_KV_RAW_TILE_BASE]) + local_tile_idx - ) - effective_total_kv_tiles = Int32(task_cache[_TASK_CACHE_KV_VALID_TILE_END]) - tile_idx = effective_tile_idx - window_start_idx = Int32(task_cache[_TASK_CACHE_KV_WINDOW_START]) - else: - effective_tile_idx = _runtime_split_kv_global_tile_idx( - cfg, - stage_info, - local_tile_idx, - seq_len_kv, - self.seq_len_q, - q_token_base, - ) - effective_total_kv_tiles = _runtime_total_kv_tiles( - cfg, seq_len_kv, self.seq_len_q, q_token_base - ) - tile_idx = _runtime_clamp_valid_tile_idx( - cfg, - effective_tile_idx, - seq_len_kv, - self.seq_len_q, - q_token_base, - ) - tile_idx = tile_idx + _num_skipped_kv_tiles( - cfg, seq_len_kv, self.seq_len_q, q_token_base - ) - window_start_idx = _sliding_window_start_idx( - cfg, seq_len_kv, self.seq_len_q, q_token_base - ) - - tile_offset_k = tile_idx * Int32(cfg.tile_size_kv) - is_valid_effective_tile = effective_tile_idx < effective_total_kv_tiles - is_masked_final_wave = False - if cutlass.const_expr(not use_runtime_kv_domain and not cfg.use_split_kv): - if cutlass.const_expr(cfg.has_odd_kv_tail and self.inst_id == 1): - is_masked_final_wave = _is_last_loop_iteration(stage_info) - tile_has_valid_scores = ( - is_valid_effective_tile - and (tile_offset_k < seq_len_kv) - and not is_masked_final_wave - ) - tile_is_unmasked = _kv_tile_is_fully_unmasked_for_q_group( - cfg, - tile_offset_k, - seq_len_kv, - self.seq_len_q, - q_token_base, - tile_has_valid_scores, - ) - return ( - seq_len_kv, - logical_q_group_idx, - element_mask_end_idx, - tile_offset_k, - window_start_idx, - is_valid_effective_tile, - is_masked_final_wave, - tile_is_unmasked, - tile_has_valid_scores, + return resolve_keeps_tile_context( + self.cfg, + stage_info, + inst_id=self.inst_id, + seqlens_kv=self.seqlens_kv, + max_seq_len_kv=self.max_seq_len_kv, + seq_len_q=self.seq_len_q, + q_group_idx=self.q_group_idx, ) @cute.jit @@ -888,6 +989,29 @@ def _publish_keeps_softmax_state( for reg_idx in cutlass.range_constexpr(self.cfg.num_s_regs_per_thread): s_arr[reg_idx] = s_vals[reg_idx] + @cute.jit + def _proxy_score_shifts( + self, route_is_proxy: cutlass.Boolean + ) -> tuple[Float32, Float32]: + """Return a route's ``(proxy_max_shift, tail_shift)`` in score units. + + A proxy route carries its token mass in the logit: the tile maximum + moves by the full block's ``log2`` mass and the ragged final summary's + score by its ``log2`` shortfall, both divided by ``c`` so they apply + to scores before the anchor. Exact routes get zero shifts. Callers + pass a warp-uniform route kind so the branch stays uniform. + """ + proxy_max_shift = Float32(0.0) + tail_shift = Float32(0.0) + if route_is_proxy: + inv_scale_softmax_log2 = cute.math.rcp(self.scale_softmax_log2) + _, tail_log2_delta = self.cfg.proxy_tail_summary + proxy_max_shift = ( + Float32(self.cfg.proxy_log2_block_mass) * inv_scale_softmax_log2 + ) + tail_shift = Float32(tail_log2_delta) * inv_scale_softmax_log2 + return proxy_max_shift, tail_shift + @cute.jit def _softmax_anchor(self, old_max: Float32, tile_max: Float32) -> Float32: """Return the exponent reference max for the tile's P pass. @@ -909,6 +1033,67 @@ def _softmax_anchor(self, old_max: Float32, tile_max: Float32) -> Float32: new_max = old_max return new_max + @cute.jit + def _q_token_kv_block_sparse_keeps_kv128_membership_word( + self, + stage_info: StageInfo, + logical_q_group_idx: Int32, + warp_grp_thread_idx: Int32, + tile_row_idx: Int32, + ) -> Uint32: + """Return a lane-local 16/32-page keep word for Q64/Q128 with KV128.""" + + cfg = self.cfg + assert cfg.tile_size_q in (64, 128) + assert cfg.tile_size_kv == 128 + assert cfg.uses_q_token_kv_block_sparse_page_membership + assert self.page_offsets_ref is not None + q_token_idx, _ = _q_row_token_and_local_head( + cfg, + self.h_r, + logical_q_group_idx, + tile_row_idx, + ) + membership_bit = Uint32(1) << q_token_idx + local_tile_idx = _softmax_tile_idx(cfg, stage_info, self.inst_id) + lane_idx = warp_grp_thread_idx & Int32(31) + col_base = _keeps_col_base( + cfg, + lane_idx, + cfg.softmax_score_fragment_regs, + ) + keep_word = Uint32(0) + page_span = min(cfg.num_tokens_per_page, cfg.num_s_regs_per_thread) + pages_per_lane = cfg.num_s_regs_per_thread // page_span + for page_vector_idx in cutlass.range_constexpr(pages_per_lane // 4): + memberships = ( + self.page_offsets_ref.q_token_kv_block_sparse_page_memberships4( + stage_info, + local_tile_idx, + col_base // Int32(cfg.num_tokens_per_page) + + Int32(page_vector_idx * 4), + ) + ) + for vector_elem_idx in cutlass.range_constexpr(4): + local_page_idx = page_vector_idx * 4 + vector_elem_idx + page_is_member = Uint32( + (memberships[vector_elem_idx] & membership_bit) != Uint32(0) + ) + keep_word = keep_word | (page_is_member << Uint32(local_page_idx)) + if cutlass.const_expr(pages_per_lane < 4): + for local_page_idx in cutlass.range_constexpr(pages_per_lane): + membership = ( + self.page_offsets_ref.q_token_kv_block_sparse_page_membership( + stage_info, + local_tile_idx, + (col_base + Int32(local_page_idx * page_span)) + // Int32(cfg.num_tokens_per_page), + ) + ) + member = Uint32((membership & membership_bit) != Uint32(0)) + keep_word = keep_word | (member << Uint32(local_page_idx)) + return keep_word + @cute.jit def _load_keeps_fragment_impl( self, @@ -921,6 +1106,7 @@ def _load_keeps_fragment_impl( logical_q_group_idx: Int32, is_valid_effective_tile: cutlass.Boolean, is_masked_final_wave: cutlass.Boolean, + membership_keep_word: Uint32, *, apply_boundary_mask: Constexpr[bool], ) -> None: @@ -1012,18 +1198,77 @@ def _load_keeps_fragment_impl( ) causal_start_rel = causal_start - tile_offset_k causal_end_rel = causal_end - tile_offset_k - for reg_idx in cutlass.range_constexpr(num_s_regs): - score_col = _keeps_score_col( - cfg, - warp_grp_thread_idx, - reg_idx, - col_base, - ) - if cutlass.const_expr(cfg.use_sliding_window_causal): - if score_col < causal_start_rel: + if cutlass.const_expr(cfg.uses_q_token_kv_block_sparse_page_membership): + assert cfg.tile_size_kv == 128 + # Membership removes complete future pages; only tokens + # beyond the causal endpoint within its page need masking. + causal_tail_tokens = causal_end & Int32(cfg.num_tokens_per_page - 1) + causal_tail_page_rel = causal_end_rel - causal_tail_tokens + page_span = min(num_s_regs, cfg.num_tokens_per_page) + for local_page_idx in cutlass.range_constexpr( + num_s_regs // page_span + ): + page_score_col = _keeps_score_col( + cfg, + warp_grp_thread_idx, + local_page_idx * page_span, + col_base, + ) + page_origin = ( + page_score_col // Int32(cfg.num_tokens_per_page) + ) * Int32(cfg.num_tokens_per_page) + page_is_causal_tail = ( + causal_tail_tokens != Int32(0) + and page_origin == causal_tail_page_rel + ) + # A KV128 block may span two lane-local 64-column halves. + # The second half's first register can already be masked. + for token_in_span in cutlass.range_constexpr( + 0 if cfg.num_tokens_per_page > num_s_regs else 1, page_span + ): + token_offset = ( + page_score_col - page_origin + Int32(token_in_span) + ) + valid = not ( + page_is_causal_tail + and token_offset >= causal_tail_tokens + ) + reg = local_page_idx * page_span + token_in_span + s_vals[reg] = cutlass.select_( + valid, s_vals[reg], _neg_max_f32() + ) + else: + for reg_idx in cutlass.range_constexpr(num_s_regs): + score_col = _keeps_score_col( + cfg, + warp_grp_thread_idx, + reg_idx, + col_base, + ) + if cutlass.const_expr(cfg.use_sliding_window_causal): + if score_col < causal_start_rel: + s_vals[reg_idx] = _neg_max_f32() + if score_col >= causal_end_rel: s_vals[reg_idx] = _neg_max_f32() - if score_col >= causal_end_rel: - s_vals[reg_idx] = _neg_max_f32() + + 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 + # complete 128-column row. The preloaded 16/32-page keep word + # avoids carrying SMEM values through this dynamic + # masked/unmasked loader specialization. + assert cfg.tile_size_kv == 128 + page_span = min(num_s_regs, cfg.num_tokens_per_page) + for local_page_idx in cutlass.range_constexpr(num_s_regs // page_span): + page_is_member = ( + (membership_keep_word >> Uint32(local_page_idx)) & Uint32(1) + ) != Uint32(0) + for token_in_page in cutlass.range_constexpr(page_span): + membership_reg_idx = local_page_idx * page_span + token_in_page + s_vals[membership_reg_idx] = cutlass.select_( + page_is_member, + s_vals[membership_reg_idx], + _neg_max_f32(), + ) if cutlass.const_expr(cfg.q_score_rows_need_mask): if not _q_row_is_valid_for_seq( @@ -1033,8 +1278,8 @@ def _load_keeps_fragment_impl( tile_row_idx, self.seq_len_q, ): - for reg_idx in cutlass.range_constexpr(num_s_regs): - s_vals[reg_idx] = _neg_max_f32() + for invalid_q_reg_idx in cutlass.range_constexpr(num_s_regs): + s_vals[invalid_q_reg_idx] = _neg_max_f32() @cute.jit def _load_keeps_fragment( @@ -1057,6 +1302,19 @@ def _load_keeps_fragment( lets the unmasked specialization erase boundary-mask instructions and avoids carrying the loaded score registers through a post-LDTM branch. """ + membership_keep_word = Uint32(0xFFFFFFFF) + if cutlass.const_expr(self.cfg.uses_q_token_kv_block_sparse_page_membership): + task_cache = _decode_gen_task_cache(stage_info) + warp_grp_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] + tile_row_idx = _keeps_row_idx(self.cfg, warp_grp_thread_idx) + membership_keep_word = ( + self._q_token_kv_block_sparse_keeps_kv128_membership_word( + stage_info, + logical_q_group_idx, + warp_grp_thread_idx, + tile_row_idx, + ) + ) if tile_is_unmasked: self._load_keeps_fragment_impl( stage_info, @@ -1068,6 +1326,7 @@ def _load_keeps_fragment( logical_q_group_idx, is_valid_effective_tile, is_masked_final_wave, + membership_keep_word, apply_boundary_mask=False, ) else: @@ -1081,6 +1340,7 @@ def _load_keeps_fragment( logical_q_group_idx, is_valid_effective_tile, is_masked_final_wave, + membership_keep_word, apply_boundary_mask=True, ) @@ -1112,13 +1372,6 @@ def _compute_softmax_loop_keeps( new_max_arr=new_max_arr, s_arr=s_arr, use_sparse=False, - sparse_origin0=Int32(0), - sparse_origin1=Int32(0), - sparse_route_flags=Int32(0), - sparse_token_word0=Uint32(0xFFFFFFFF), - sparse_token_word1=Uint32(0xFFFFFFFF), - sparse_token_word2=Uint32(0xFFFFFFFF), - sparse_token_word3=Uint32(0xFFFFFFFF), ) task_cache = _decode_gen_task_cache(stage_info) num_s_regs = cfg.num_s_regs_per_thread @@ -1267,6 +1520,64 @@ def _compute_softmax_loop_keeps( ) return old_max_arr, sum_arr, new_max_arr, s_arr + @cute.jit + def _apply_q_token_kv_block_sparse_swaps_page_membership_mask( + self, + stage_info: StageInfo, + s_vals: cutlass.Array, + task_cache: cutlass.Array, + logical_q_group_idx: Int32, + local_tile_idx: Int32, + ) -> None: + """Mask Swaps scores whose page is absent from a grouped Q row.""" + + cfg = self.cfg + assert not cfg.use_keeps_mma_ab + assert cfg.uses_q_token_kv_block_sparse_page_membership + assert self.page_offsets_ref is not None + num_scale_groups = cfg.num_softmax_scale_groups + q_repeats = max(cfg.tile_size_q // 8, 1) + warp_idx = task_cache[_TASK_CACHE_WARP_IDX] + lane_idx = task_cache[_TASK_CACHE_LANE_IDX] + col_group_idx = lane_idx & Int32(0x3) + local_idx_k0 = warp_idx * Int32(32) + (lane_idx >> Int32(2)) + + # Each held page-4 slot has a separate query-membership byte; locator + # values never carry membership bits. One lookup serves the pair of Q + # rows represented by a Swaps score register pair; causal masking still + # handles the zero-to-three-token tail within a member page. + for scale_idx in cutlass.range_constexpr(num_scale_groups): + repeat_idx = scale_idx // 2 + pair_idx = scale_idx % 2 + tile_row_idx = ( + Int32(repeat_idx * 8) + col_group_idx * Int32(2) + Int32(pair_idx) + ) + q_token_idx, _ = _q_row_token_and_local_head( + cfg, + self.h_r, + logical_q_group_idx, + tile_row_idx, + ) + membership_bit = Uint32(1) << q_token_idx + for token_group_idx in cutlass.range_constexpr(4): + local_token_idx = local_idx_k0 + Int32(token_group_idx * 8) + page_frag = local_token_idx // Int32(cfg.num_tokens_per_page) + membership = ( + self.page_offsets_ref.q_token_kv_block_sparse_page_membership( + stage_info, + local_tile_idx, + page_frag, + ) + ) + s_idx = ( + repeat_idx * 4 + + pair_idx + + (token_group_idx & 1) * 2 + + (token_group_idx >> 1) * q_repeats * 4 + ) + if (membership & membership_bit) == Uint32(0): + s_vals[s_idx] = _neg_max_f32() + @cute.jit def _compute_softmax_loop_swaps( self, @@ -1462,9 +1773,7 @@ def _compute_softmax_loop_swaps( route_is_proxy = cutlass.Boolean(False) if cutlass.const_expr(use_sparse and cfg.use_block_sparse_proxy_routes): - route_is_proxy = cutlass.Boolean( - (sparse_route_flags & Uint32(_SOFTMAX_ROUTE_IS_PROXY_FLAG)) != Uint32(0) - ) + route_is_proxy = _route_is_proxy(sparse_route_flags.bitcast(Int32)) if cutlass.const_expr(use_sparse): # Route, KV-tail, uniform-causal, and token validity depend only on # K, so one predicate masks the adjacent pair of Q-row registers. @@ -1617,6 +1926,15 @@ def _compute_softmax_loop_swaps( s_vals[s_base + 0] = _neg_max_f32() s_vals[s_base + 1] = _neg_max_f32() + if cutlass.const_expr(cfg.uses_q_token_kv_block_sparse_page_membership): + self._apply_q_token_kv_block_sparse_swaps_page_membership_mask( + stage_info, + s_vals, + task_cache, + logical_q_group_idx, + local_tile_idx, + ) + if cutlass.const_expr(cfg.uses_per_row_causal_mask): apply_per_row_causal_mask = cutlass.Boolean(True) if cutlass.const_expr(not use_sparse): @@ -1639,16 +1957,18 @@ def _compute_softmax_loop_swaps( # Grouped causal decode has a distinct causal/window bound for # every Q token. Sparse routes always use their logical K; # dense routes retain the boundary-tile fast path above. - warp_idx = task_cache[_TASK_CACHE_WARP_IDX] - lane_idx = task_cache[_TASK_CACHE_LANE_IDX] - col_group_idx = lane_idx & Int32(0x3) - local_idx_k0 = warp_idx * Int32(32) + (lane_idx >> Int32(2)) + causal_warp_idx = task_cache[_TASK_CACHE_WARP_IDX] + causal_lane_idx = task_cache[_TASK_CACHE_LANE_IDX] + causal_col_group_idx = causal_lane_idx & Int32(0x3) + causal_local_idx_k0 = causal_warp_idx * Int32(32) + ( + causal_lane_idx >> Int32(2) + ) for scale_idx in cutlass.range_constexpr(num_scale_groups): repeat_idx = scale_idx // 2 pair_idx = scale_idx % 2 tile_row_idx = ( Int32(repeat_idx * 8) - + col_group_idx * Int32(2) + + causal_col_group_idx * Int32(2) + Int32(pair_idx) ) q_token_idx, _ = _q_row_token_and_local_head( @@ -1665,12 +1985,14 @@ def _compute_softmax_loop_swaps( ) for token_group_idx in cutlass.range_constexpr(4): token_idx = ( - tile_offset_k + local_idx_k0 + Int32(token_group_idx * 8) + tile_offset_k + + causal_local_idx_k0 + + Int32(token_group_idx * 8) ) if cutlass.const_expr(use_sparse): _, token_idx = _swaps_routed_coordinate( cfg, - lane_idx >> Int32(2), + causal_lane_idx >> Int32(2), sparse_origin0, sparse_origin1, sparse_origin2, @@ -1714,6 +2036,45 @@ def _compute_softmax_loop_swaps( s_vals[s_base_hi + 2] = _neg_max_f32() lane_idx = task_cache[_TASK_CACHE_LANE_IDX] + # Proxy mass (``_proxy_score_shifts``): the ragged final summary's + # scores shift before the reduction so the P pass sees its true mass. + proxy_max_shift = Float32(0.0) + if cutlass.const_expr(use_sparse and cfg.use_block_sparse_proxy_routes): + tail_summary_idx, tail_log2_delta = cfg.proxy_tail_summary + route_is_proxy_uniform = cute.arch.make_warp_uniform(route_is_proxy) + proxy_max_shift, tail_shift = self._proxy_score_shifts( + route_is_proxy_uniform + ) + if cutlass.const_expr(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( + cfg, + Int32(lane_idx) >> Int32(2), + sparse_origin0, + sparse_origin1, + sparse_origin2, + sparse_origin3, + token_group_idx=token_group_idx, + ) + # Invalid atoms never hold the tail. + if atom_origin >= Int32(0) and logical_summary == Int32( + tail_summary_idx + ): + for tail_repeat_idx in cutlass.range_constexpr(q_repeats): + tail_s_base = tail_repeat_idx * 4 + token_group_idx * 2 + if cutlass.const_expr(token_group_idx >= 2): + tail_s_base = ( + q_repeats * 4 + + tail_repeat_idx * 4 + + (token_group_idx - 2) * 2 + ) + s_vals[tail_s_base + 0] = ( + s_vals[tail_s_base + 0] + tail_shift + ) + s_vals[tail_s_base + 1] = ( + s_vals[tail_s_base + 1] + tail_shift + ) for scale_idx in cutlass.range_constexpr(num_scale_groups): # Reduce this lane's S registers to one candidate per scale group. repeat_idx = scale_idx // 2 @@ -1725,6 +2086,9 @@ def _compute_softmax_loop_swaps( cute.math.max(s_vals[s_base_hi + 0], s_vals[s_base_hi + 2], ftz=True), ftz=True, ) + if cutlass.const_expr(use_sparse and cfg.use_block_sparse_proxy_routes): + # The mass enters before the anchor, so the P range is unchanged. + local_max = local_max + proxy_max_shift local_max = cute.math.max(local_max, old_max_vals[scale_idx], ftz=True) local_max_vals[scale_idx] = local_max @@ -1866,6 +2230,62 @@ def compute_softmax_loop( use_sparse=False, ) + @consumer_work(work_attrs=WorkAttr.AUXILIARY, returns=sage_q_scale) + @cute.jit + def load_sage_q_scale(self, stage_info: StageInfo) -> Float32: + """Load ``sfQ`` of the lane's Q row, fixed for the whole work tile.""" + cfg = self.cfg + assert cfg.use_sage_attention + task_cache = _decode_gen_task_cache(stage_info) + warp_grp_thread_idx = Int32(task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX]) + kv_head_idx, batch_idx = _logical_head_batch( + stage_info, self.h_k_idx, self.b_idx + ) + q_token_idx, local_head_idx = _q_row_token_and_local_head( + cfg, + self.h_r, + _logical_q_group_idx(cfg, stage_info, self.q_group_idx), + _keeps_row_idx(cfg, warp_grp_thread_idx), + ) + return load_q_scale( + cfg, + self.scale_tensors.q_scale_ptr.toint(), + self.scale_tensors.q_scale_head_stride, + kv_head_idx=kv_head_idx, + local_head_idx=local_head_idx, + batch_idx=batch_idx, + q_token_idx=q_token_idx, + ) + + @consumer_work(returns=("old_max_arr", "sum_arr", "new_max_arr", "s_arr")) + @cute.jit + def compute_sage_softmax_loop( + self, + stage_info: StageInfo, + *, + old_max_arr: cutlass.Array, + sum_arr: cutlass.Array, + new_max_arr: cutlass.Array, + s_arr: cutlass.Array, + sage_q_scale: Float32, + sage_scale_arr: cutlass.Array, + sage_summary_scale_arr: cutlass.Array, + ) -> tuple[object, object, object, object]: + """Consume S with per-group Sage scales and publish the dequantized max.""" + + assert self.cfg.use_sage_attention + return self._compute_softmax_loop_keeps_fragments( + 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_scale_arr=sage_scale_arr, + sage_summary_scale_arr=sage_summary_scale_arr, + ) + @consumer_work( returns=sum_arr, work_attrs=WorkAttr.AUXILIARY, @@ -1882,8 +2302,8 @@ def reduce_sums( """Fold local P sums into the running online-softmax denominators.""" cfg = self.cfg # ConsTailWork: denominator update runs after P has been materialized, - # so the resource-owned local sum matches the P payload consumed by BMM2. - if cutlass.const_expr(cfg.use_fp8_qkv): + # so local_sum_arr represents the exact P payload consumed by BMM2. + if cutlass.const_expr(cfg.use_fp8_pv): # FP8 uses TmemSoftmaxGlobal to update sums after P # quantization, so this stage only copies the corrected sums # back into the running state. This keeps the denominator @@ -1928,7 +2348,7 @@ def reduce_sums( cfg.has_static_dense_full_kv_tiles and cfg.tile_size_q in (16, 32) and not cfg.use_keeps_mma_ab - and not cfg.use_fp8_qkv + and not cfg.use_fp8_pv and cfg.q_tiles_are_full ): for pair_idx in cutlass.range_constexpr(pair_width): @@ -1968,13 +2388,16 @@ def _compute_softmax_loop_keeps_fragments( new_max_arr: cutlass.Array, s_arr: cutlass.Array, use_sparse: Constexpr[bool], - sparse_origin0: Int32, - sparse_origin1: Int32, - sparse_route_flags: Int32, - sparse_token_word0: Uint32, - sparse_token_word1: Uint32, - sparse_token_word2: Uint32, - sparse_token_word3: Uint32, + sparse_origin0: Int32 | None = None, + sparse_origin1: Int32 | None = None, + sparse_route_flags: Int32 | None = None, + sparse_token_word0: Uint32 | None = None, + sparse_token_word1: Uint32 | None = None, + sparse_token_word2: Uint32 | None = None, + sparse_token_word3: Uint32 | None = None, + sage_q_scale: Float32 | None = None, + sage_scale_arr: cutlass.Array | None = None, + sage_summary_scale_arr: cutlass.Array | None = None, ) -> tuple[object, object, object, object]: """Mask streamed K32 score fragments in place and reduce their max. @@ -1982,8 +2405,13 @@ def _compute_softmax_loop_keeps_fragments( their two K64 atom origins, validity flags and prepared token words; dense tiles derive it from the tile's visible token range (sequence end, uniform or per-row causal end, sliding-window start) and the Q - row's validity. Masked fragments are written back to TMEM so the P - pass can reload them without any mask logic. + row's validity, and leave the route arguments unset. Masked + fragments are written back to TMEM so the P pass can reload them + without any mask logic. With Sage attention the scores stay quantized + unless the geometry has one scale per score: each fragment folds its + group maxima with ``sfK`` and ``sfQ`` scales the tile maximum once. + Biased INT32 scores (``INT32_SCORE_BIAS``) share this path; the bias + leaves with the group maxima. """ cfg = self.cfg assert cfg.streams_tmem_p_fragments @@ -1995,8 +2423,8 @@ def _compute_softmax_loop_keeps_fragments( keep_words = cutlass.Array( Uint32, num_fragments, space=cutlass.AddressSpace.rmem ) - warp_group_thread_idx = Int32(task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX]) - tile_row_idx = _keeps_row_idx(cfg, warp_group_thread_idx) + warp_grp_thread_idx = Int32(task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX]) + tile_row_idx = _keeps_row_idx(cfg, warp_grp_thread_idx) if cutlass.const_expr(use_sparse): token_words = ( sparse_token_word0, @@ -2112,33 +2540,156 @@ def _compute_softmax_loop_keeps_fragments( # The load/store branch must be uniform for each participating warp. warp_scores_are_unmasked = cute.arch.vote_all_sync(warp_scores_are_unmasked) + # Proxy mass (``_proxy_score_shifts``): the ragged final summary's + # score shifts before the fold and the TMEM write-back, so the P pass + # reloads the shifted score. + proxy_max_shift = Float32(0.0) + tail_shift = Float32(0.0) + tail_fragment_mask = Int32(0) + tail_lane: Constexpr[int] = 0 + shifts_tail: Constexpr[bool] = False + route_is_proxy = cutlass.Boolean(False) + # A mixed plan quantizes proxy summaries with their own K block size, + # so exact and proxy tiles read different ``sfK`` geometries; the + # CTA-uniform route kind selects one per tile. + mixed_scales: Constexpr[bool] = use_sparse and cfg.sage_mixed_k_geometry + # The unrolled unmasked pass serves the exact geometry only. + exact_dequantized: Constexpr[bool] = cfg.sage_scores_dequantized_for( + proxy=False + ) + if cutlass.const_expr(use_sparse and cfg.use_block_sparse_proxy_routes): + # A warp-uniform route kind keeps the fold and the load/store + # branch below on the uniform datapath. + route_is_proxy = cute.arch.make_warp_uniform( + _route_is_proxy(sparse_route_flags) + ) + tail_summary_idx, tail_log2_delta = cfg.proxy_tail_summary + tail_lane = tail_summary_idx % fragment_regs + shifts_tail = 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 + # ``sfQ`` here and by the tail group's ``sfK`` in the fragment. + 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 + # 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 + 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): + # The half's first atom is atom ``tail_half`` of the group. + tail_group_origin0 = ( + cfg.num_proxy_groups - 1 + ) * cfg.tile_size_kv + tail_half * cfg.block_sparse_kv_atom_size + holds_tail = cutlass.Boolean( + holds_tail and origin0 == Int32(tail_group_origin0) + ) + if holds_tail: + tail_fragment_mask = Int32(1 << tail_fragment) + tail_fragment_mask = cute.arch.make_warp_uniform(tail_fragment_mask) + # The shifted tail score must reach TMEM, so the route takes + # the store branch. + warp_scores_are_unmasked = cute.arch.vote_all_sync( + cutlass.Boolean( + warp_scores_are_unmasked and tail_fragment_mask == Int32(0) + ) + ) + if cutlass.const_expr(mixed_scales): + # Proxy tiles take the rolled pass, so the summary geometry's fold + # is emitted once rather than per unrolled fragment. + warp_scores_are_unmasked = cutlass.Boolean( + warp_scores_are_unmasked and not route_is_proxy + ) + score_tmem_addr = ( task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + Int32(self._alloc.offset) + self._softmax_loop_stage_slot_offset(stage_info) ) max_chains = cutlass.Array(Float32, 4, space=cutlass.AddressSpace.rmem) + max_identity = _neg_max_f32() + if cutlass.const_expr(cfg.use_sage_attention): + max_identity = Float32(float("-inf")) for chain_idx in cutlass.range_constexpr(4): - max_chains[chain_idx] = _neg_max_f32() + max_chains[chain_idx] = max_identity + scales_view = None + summary_view = None + if cutlass.const_expr(cfg.use_sage_attention): + scales_view = self.sage_k_scales.open(stage_info, sage_scale_arr, None) + if cutlass.const_expr(mixed_scales): + summary_view = self.sage_summary_k_scales.open( + stage_info, sage_summary_scale_arr, None + ) if warp_scores_are_unmasked: for fragment_idx in cutlass.range_constexpr(num_fragments): + if cutlass.const_expr(cfg.use_sage_attention): + # Issued ahead of the score load so an SMEM read hides + # behind the TMEM wait. + fragment_scales = self.sage_k_scales.fragment( + scales_view, Int32(fragment_idx) + ) loaded = _keeps_tcgen05_ld( cfg, prims.make_tmem_ptr( - score_tmem_addr + Int32(fragment_idx * fragment_regs), Float32 + score_tmem_addr + Int32(fragment_idx * fragment_regs), + Float32, ), num=fragment_regs, offset=cfg.tile_size_kv // 2, ) prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) - for score_idx in cutlass.range_constexpr(fragment_regs): - chain_idx: Constexpr[int] = score_idx % 4 - max_chains[chain_idx] = cute.math.max( - max_chains[chain_idx], - Float32(loaded[score_idx]), - ftz=True, + if cutlass.const_expr(exact_dequantized): + scores = cutlass.Array( + Float32, fragment_regs, space=cutlass.AddressSpace.rmem + ) + for score_idx in cutlass.range_constexpr(fragment_regs): + scores[score_idx] = Float32(loaded[score_idx]) + self._dequantize_fragment_scores( + scores, fragment_scales, cfg.sage_k_groups_per_fragment ) + self.sage_k_scales.advance(scales_view) + for score_idx in cutlass.range_constexpr(fragment_regs): + chain_idx: Constexpr[int] = score_idx % 4 + max_chains[chain_idx] = cute.math.max( + max_chains[chain_idx], Float32(scores[score_idx]), ftz=True + ) + _keeps_tcgen05_st( + cfg, + prims.make_tmem_ptr( + score_tmem_addr + Int32(fragment_idx * fragment_regs), + Float32, + ), + scores.data_ptr().load(count=fragment_regs, alignment=4), + offset=cfg.tile_size_kv // 2, + ) + elif cutlass.const_expr(cfg.use_sage_attention): + self._fold_sage_fragment_max( + max_chains, + loaded, + fragment_scales=fragment_scales, + chain_base=fragment_idx * cfg.sage_k_groups_per_fragment, + groups=cfg.sage_k_groups_per_fragment, + ) + self.sage_k_scales.advance(scales_view) + else: + for score_idx in cutlass.range_constexpr(fragment_regs): + chain_idx: Constexpr[int] = score_idx % 4 + max_chains[chain_idx] = cute.math.max( + max_chains[chain_idx], + Float32(loaded[score_idx]), + ftz=True, + ) + if cutlass.const_expr(exact_dequantized): + prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) + cute.arch.fence_view_async_tmem_store() else: if cutlass.const_expr(not use_sparse): lane_idx = Int32(task_cache[_TASK_CACHE_LANE_IDX]) @@ -2146,7 +2697,7 @@ def _compute_softmax_loop_keeps_fragments( for fragment_idx in cutlass.range_constexpr(num_fragments): fragment_token_base = tile_offset_k + _keeps_score_col( cfg, - warp_group_thread_idx, + warp_grp_thread_idx, fragment_idx * fragment_regs, col_base, ) @@ -2156,35 +2707,45 @@ def _compute_softmax_loop_keeps_fragments( visible_end - fragment_token_base, fragment_regs=fragment_regs, ) - for fragment_idx in cutlass.range_constexpr(num_fragments): - fragment_addr = score_tmem_addr + Int32(fragment_idx * fragment_regs) - loaded = _keeps_tcgen05_ld( - cfg, - prims.make_tmem_ptr(fragment_addr, Float32), - num=fragment_regs, - offset=cfg.tile_size_kv // 2, - ) - prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) - masked_scores = cutlass.Array( - Float32, fragment_regs, space=cutlass.AddressSpace.rmem - ) - for score_idx in cutlass.range_constexpr(fragment_regs): - score = Float32(loaded[score_idx]) - score_is_kept = ( - (keep_words[fragment_idx] >> Int32(score_idx)) & Uint32(1) - ) != Uint32(0) - if not score_is_kept: - score = _neg_max_f32() - masked_scores[score_idx] = score - chain_idx: Constexpr[int] = score_idx % 4 - max_chains[chain_idx] = cute.math.max( - max_chains[chain_idx], score, ftz=True + # The masked pass is one rolled loop per route kind. A mixed + # geometry selects the loop per tile: a per-fragment selection + # would be if-converted and run both folds on every tile. + if cutlass.const_expr(mixed_scales): + if route_is_proxy: + self._mask_score_fragments( + score_tmem_addr, + keep_words, + max_chains, + proxy_kind=True, + scales_view=summary_view, + may_hold_tail=shifts_tail, + tail_fragment_mask=tail_fragment_mask, + tail_shift=tail_shift, + tail_lane=tail_lane, ) - _keeps_tcgen05_st( - cfg, - prims.make_tmem_ptr(fragment_addr, Float32), - masked_scores.data_ptr().load(count=fragment_regs, alignment=4), - offset=cfg.tile_size_kv // 2, + else: + self._mask_score_fragments( + score_tmem_addr, + keep_words, + max_chains, + proxy_kind=False, + scales_view=scales_view, + may_hold_tail=shifts_tail, + tail_fragment_mask=tail_fragment_mask, + tail_shift=tail_shift, + tail_lane=tail_lane, + ) + else: + self._mask_score_fragments( + score_tmem_addr, + keep_words, + max_chains, + proxy_kind=False, + scales_view=scales_view, + may_hold_tail=shifts_tail, + tail_fragment_mask=tail_fragment_mask, + tail_shift=tail_shift, + tail_lane=tail_lane, ) prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) cute.arch.fence_view_async_tmem_store() @@ -2194,12 +2755,309 @@ def _compute_softmax_loop_keeps_fragments( cute.math.max(max_chains[2], max_chains[3], ftz=True), ftz=True, ) + if cutlass.const_expr(cfg.use_sage_attention): + # The chains hold ``gmax * sfK``; ``sfQ`` applies once here. Only + # the ``-inf`` identity maps to the finite sentinel: a real + # pre-``sfQ`` maximum of ``-FLT_MAX`` is an ordinary score. + if tile_max == Float32(float("-inf")): + tile_max = _neg_max_f32() + else: + tile_max = tile_max * sage_q_scale + if cutlass.const_expr(use_sparse and cfg.use_block_sparse_proxy_routes): + # The mass enters before the anchor, so the P range is unchanged; a + # fully masked route's sentinel absorbs the shift in FP32. + tile_max = tile_max + proxy_max_shift old_max = new_max_arr[0] new_max = self._softmax_anchor(old_max, tile_max) old_max_arr[0] = old_max new_max_arr[0] = new_max return old_max_arr, sum_arr, new_max_arr, s_arr + @cute.jit + def _mask_score_fragments( + self, + score_tmem_addr: Int32, + keep_words: cutlass.Array, + max_chains: cutlass.Array, + *, + proxy_kind: Constexpr[bool], + scales_view, + may_hold_tail: Constexpr[bool], + tail_fragment_mask: Int32, + tail_shift: Float32, + tail_lane: Constexpr[int], + ) -> None: + """Mask, fold and write back every fragment with one route kind's scales. + + ``scales_view`` is the opened view of the resource ``proxy_kind`` + selects; the resource is read from ``self`` because a local would be + flattened as a loop-carried value. The keep words rotate down one + entry per fragment, so the body reads ``keep_words[0]``. + + Cleared keep bits become ``-inf`` with Sage, else ``-FLT_MAX``. The + fragment holding a proxy route's ragged final summary adds the mass + shortfall to that score; Sage divides the shift (already divided by + ``sfQ``) by the tail group's ``sfK`` with an approximate reciprocal, + ample for a shift far below the score resolution. A route kind whose + scores are dequantized writes the dequantized scores back. + """ + cfg = self.cfg + num_fragments = cfg.num_softmax_score_fragments + fragment_regs = cfg.softmax_score_fragment_regs + groups: Constexpr[int] = cfg.sage_k_groups_per_fragment_for(proxy=proxy_kind) + dequantize_scores: Constexpr[bool] = cfg.sage_scores_dequantized_for( + proxy=proxy_kind + ) + for fragment in cutlass.range(num_fragments, unroll=1): + fragment_scales = None + if cutlass.const_expr(cfg.use_sage_attention and proxy_kind): + fragment_scales = self.sage_summary_k_scales.fragment( + scales_view, Int32(fragment) + ) + elif cutlass.const_expr(cfg.use_sage_attention): + fragment_scales = self.sage_k_scales.fragment( + scales_view, Int32(fragment) + ) + keep_word = Uint32(keep_words[0]) + fragment_addr = score_tmem_addr + Int32(fragment) * Int32(fragment_regs) + loaded = _keeps_tcgen05_ld( + cfg, + prims.make_tmem_ptr(fragment_addr, Float32), + num=fragment_regs, + offset=cfg.tile_size_kv // 2, + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + masked_scores = cutlass.Array( + Float32, fragment_regs, space=cutlass.AddressSpace.rmem + ) + for score_idx in cutlass.range_constexpr(fragment_regs): + score = Float32(loaded[score_idx]) + score_is_kept = ((keep_word >> Int32(score_idx)) & Uint32(1)) != Uint32( + 0 + ) + if not score_is_kept: + if cutlass.const_expr(cfg.use_sage_attention): + 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): + lane_shift = tail_shift + if cutlass.const_expr(cfg.use_sage_attention): + tail_group: Constexpr[int] = tail_lane // ( + fragment_regs // groups + ) + lane_shift = tail_shift * cute.math.rcp( + Float32(fragment_scales[tail_group]), approx=True + ) + score = score + lane_shift + masked_scores[score_idx] = score + if cutlass.const_expr(not cfg.use_sage_attention): + chain_idx: Constexpr[int] = score_idx % 4 + max_chains[chain_idx] = cute.math.max( + max_chains[chain_idx], score, ftz=True + ) + if cutlass.const_expr(dequantize_scores): + self._dequantize_fragment_scores(masked_scores, fragment_scales, groups) + for score_idx in cutlass.range_constexpr(fragment_regs): + chain_idx: Constexpr[int] = score_idx % 4 + max_chains[chain_idx] = cute.math.max( + max_chains[chain_idx], + Float32(masked_scores[score_idx]), + ftz=True, + ) + elif cutlass.const_expr(cfg.use_sage_attention): + self._fold_sage_fragment_max( + max_chains, + masked_scores, + fragment_scales=fragment_scales, + chain_base=0, + groups=groups, + ) + _keeps_tcgen05_st( + cfg, + prims.make_tmem_ptr(fragment_addr, Float32), + masked_scores.data_ptr().load(count=fragment_regs, alignment=4), + offset=cfg.tile_size_kv // 2, + ) + if cutlass.const_expr(cfg.use_sage_attention and proxy_kind): + self.sage_summary_k_scales.advance(scales_view) + elif cutlass.const_expr(cfg.use_sage_attention): + self.sage_k_scales.advance(scales_view) + for entry in cutlass.range_constexpr(num_fragments - 1): + keep_words[entry] = Uint32(keep_words[entry + 1]) + + @cute.jit + def _dequantize_fragment_scores( + self, + scores: cutlass.Array, + fragment_scales: cutlass.Array, + groups: Constexpr[int], + ) -> None: + """Replace one fragment's quantized scores by ``(s - bias) * sfK`` in place. + + Biased INT32 scores fold ``-bias * sfK`` into the same packed FMA. A + masked ``-inf`` score stays ``-inf`` under any positive scale. + """ + cfg = self.cfg + fragment_regs = cfg.softmax_score_fragment_regs + group_regs = fragment_regs // groups + group_addends = None + if cutlass.const_expr(cfg.uses_int32_scores): + group_addends = cutlass.Array( + Float32, groups, space=cutlass.AddressSpace.rmem + ) + neg_bias = Float32(-INT32_SCORE_BIAS) + for group_base in cutlass.range_constexpr(0, groups - groups % 2, 2): + group_addends[group_base], group_addends[group_base + 1] = fmul2( + (neg_bias, neg_bias), + ( + Float32(fragment_scales[group_base]), + Float32(fragment_scales[group_base + 1]), + ), + ) + if cutlass.const_expr(groups % 2 == 1): + group_addends[groups - 1] = neg_bias * Float32( + fragment_scales[groups - 1] + ) + for pair_base in cutlass.range_constexpr(0, fragment_regs, 2): + group0: Constexpr[int] = pair_base // group_regs + group1: Constexpr[int] = (pair_base + 1) // group_regs + pair = (Float32(scores[pair_base]), Float32(scores[pair_base + 1])) + scales = ( + Float32(fragment_scales[group0]), + Float32(fragment_scales[group1]), + ) + if cutlass.const_expr(cfg.uses_int32_scores): + pair = cute.arch.fma_packed_f32x2( + pair, + scales, + (Float32(group_addends[group0]), Float32(group_addends[group1])), + ) + else: + pair = fmul2(pair, scales) + scores[pair_base] = pair[0] + scores[pair_base + 1] = pair[1] + + @cute.jit + def _fold_sage_fragment_max( + self, + max_chains: cutlass.Array, + scores, + *, + fragment_scales: cutlass.Array, + chain_base: Constexpr[int], + groups: Constexpr[int], + ) -> None: + """Fold one fragment's dequantized group maxima into the max chains. + + Each scale group is reduced on the quantized scores, then + ``(group_max - bias) * sfK_g`` is folded; the caller applies ``sfQ`` + to the tile maximum. Groups of four or more scores reduce over four + chains; masked ``-inf`` scores need no special case. Group maxima + leave in pairs through a packed bias add (exact on the biased scores' + unit spacing) and a packed scale multiply. Group ``g`` folds into + chain ``(chain_base + g) % 4``. + """ + cfg = self.cfg + group_regs = cfg.softmax_score_fragment_regs // groups + width: Constexpr[int] = 1 if groups == 1 else 2 + assert groups % width == 0 + for group_base in cutlass.range_constexpr(0, groups, width): + maxima: tuple = () + for elem in cutlass.range_constexpr(width): + first: Constexpr[int] = (group_base + elem) * group_regs + if cutlass.const_expr(group_regs < 4): + group_max = Float32(scores[first]) + for score_elem in cutlass.range_constexpr(1, group_regs): + group_max = cute.math.max( + group_max, Float32(scores[first + score_elem]), ftz=True + ) + else: + chains = cutlass.Array(Float32, 4, space=cutlass.AddressSpace.rmem) + for chain_idx in cutlass.range_constexpr(4): + chains[chain_idx] = Float32(scores[first + chain_idx]) + for score_elem in cutlass.range_constexpr(4, group_regs): + chain_idx: Constexpr[int] = score_elem % 4 + chains[chain_idx] = cute.math.max( + chains[chain_idx], + Float32(scores[first + score_elem]), + ftz=True, + ) + group_max = cute.math.max( + cute.math.max(chains[0], chains[1], ftz=True), + cute.math.max(chains[2], chains[3], ftz=True), + ftz=True, + ) + maxima += (group_max,) + if cutlass.const_expr(width == 1): + scaled_one = Float32(maxima[0]) + if cutlass.const_expr(cfg.uses_int32_scores): + scaled_one = scaled_one - Float32(INT32_SCORE_BIAS) + scaled = (scaled_one * Float32(fragment_scales[group_base]),) + else: + scaled = maxima + if cutlass.const_expr(cfg.uses_int32_scores): + scaled = cute.arch.add_packed_f32x2( + scaled, (Float32(-INT32_SCORE_BIAS), Float32(-INT32_SCORE_BIAS)) + ) + scaled = fmul2( + scaled, + ( + Float32(fragment_scales[group_base]), + Float32(fragment_scales[group_base + 1]), + ), + ) + for elem in cutlass.range_constexpr(width): + scaled_max = Float32(scaled[elem]) + fold_chain: Constexpr[int] = (chain_base + group_base + elem) % 4 + max_chains[fold_chain] = cute.math.max( + max_chains[fold_chain], scaled_max, ftz=True + ) + + @consumer_work(returns=("old_max_arr", "sum_arr", "new_max_arr", "s_arr")) + @cute.jit + def compute_sage_block_sparse_softmax_loop( + self, + stage_info: StageInfo, + *, + old_max_arr: cutlass.Array, + sum_arr: cutlass.Array, + new_max_arr: cutlass.Array, + s_arr: cutlass.Array, + sparse_origin0: Int32, + sparse_origin1: Int32, + sparse_route_flags: Int32, + sparse_token_word0: Uint32, + sparse_token_word1: Uint32, + sparse_token_word2: Uint32, + sparse_token_word3: Uint32, + sage_q_scale: Float32, + sage_scale_arr: cutlass.Array, + sage_summary_scale_arr: cutlass.Array, + ) -> tuple[object, object, object, object]: + """Consume one routed S payload with per-group Sage scales.""" + + assert self.cfg.use_block_sparse and self.cfg.use_sage_attention + return self._compute_softmax_loop_keeps_fragments( + stage_info, + old_max_arr=old_max_arr, + sum_arr=sum_arr, + new_max_arr=new_max_arr, + s_arr=s_arr, + use_sparse=True, + sparse_origin0=sparse_origin0, + sparse_origin1=sparse_origin1, + sparse_route_flags=sparse_route_flags, + sparse_token_word0=sparse_token_word0, + sparse_token_word1=sparse_token_word1, + sparse_token_word2=sparse_token_word2, + sparse_token_word3=sparse_token_word3, + sage_q_scale=sage_q_scale, + sage_scale_arr=sage_scale_arr, + sage_summary_scale_arr=sage_summary_scale_arr, + ) + @consumer_work(returns=("old_max_arr", "sum_arr", "new_max_arr", "s_arr")) @cute.jit def compute_block_sparse_softmax_loop( diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_softmax_stats.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_softmax_stats.py index 426cbdd7c260..d40abf33be04 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_softmax_stats.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_softmax_stats.py @@ -250,12 +250,33 @@ def get_tmem_requirements(self) -> list[TmemAllocation]: # below use only the SMEM ring for this profile. return [] if self._alloc is None: + num_stages = ( + self.pipeline_config.num_stages + if self.pipeline_config is not None + else 1 + ) self._alloc = TmemAllocation( name=f"{self.name}", - num_columns=self.cfg.tmem_stats_cols, + # Stage addresses retain the S-column stride so the producer + # and consumer can use the same pipeline stage index. Include + # that stride in the allocation for the staged one-inst path. + num_columns=(num_stages - 1) * self.cfg.tmem_s_cols + + self.cfg.tmem_stats_cols, ) return [self._alloc] + def get_producer_requirements(self) -> list[SmemAllocation | TmemAllocation]: + """Return the allocation written by the active stats handoff path.""" + if self.cfg.keeps_stats_via_smem: + return self.get_smem_requirements() + return self.get_tmem_requirements() + + def get_consumer_requirements(self) -> list[SmemAllocation | TmemAllocation]: + """Return the allocation read by the active stats handoff path.""" + if self.cfg.keeps_stats_via_smem: + return self.get_smem_requirements() + return self.get_tmem_requirements() + @cute.jit def _create_initial_task_locals( self, context: ResourceContext | None = None @@ -743,7 +764,7 @@ def global_correction( ) -> None: """Apply FP8 P-quantization denominator correction through TmemS.""" cfg = self.cfg - if cutlass.const_expr(not cfg.use_fp8_qkv): + if cutlass.const_expr(not cfg.use_fp8_pv): return num_scale_groups = cfg.num_softmax_scale_groups diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py new file mode 100644 index 000000000000..1d607a92e63d --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py @@ -0,0 +1,437 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Transformed NVFP4 K/V resource backed by tensor memory.""" + +from dataclasses import dataclass +from typing import ClassVar + +import cutlass +import cutlass.cute as cute +from cutlass import Int32, Int64, Uint8 +from cutlass.experimental import primitives as prims +from cutlass.experimental.task_scheduling.enums import WorkAttr +from cutlass.experimental.task_scheduling.memory import ( + ResourceContext, + SmemAllocation, + TmemAllocation, +) +from cutlass.experimental.task_scheduling.resources import ( + StageInfo, + TaskLocalVariable, + consumer_work, + producer_work, +) + +from ..fmha_decode_config import FmhaDecodeConfig +from .helpers_common import ( + _TASK_CACHE_TMEM_BASE_OFFSET, + _TASK_CACHE_WARP_GRP_THREAD_IDX, + Constexpr, + DecodeGenResourceBase, + ResourceVars, + _decode_gen_task_cache, +) +from .smem_resources import SmemKvResource + + +@cute.jit +def _mul_e2m1x4_e4m3x4(packed_fp4: Int32, packed_sf: Int32) -> Int32: + """Apply four E4M3 scale factors to four unpacked E2M1 values. + + PTX 9.4 (CUDA 13.4) supports the byte-padded E2M1 operand directly. + Each FP4 value occupies the low nibble of one byte in the 32-bit input. + """ + if cutlass.const_expr(cutlass.target_version(min_version="13.4")): + return cute.arch.inline_ptx( + "mul.e4m3x4.e2m1p4x4.e4m3x4.satfinite {$w0}, {$r0}, {$r1};", + write_only_types=[Int32], + read_only_args=[packed_fp4, packed_sf], + ) + else: + return cute.arch.inline_ptx( + """ + { + .reg .b8 fp4_01, fp4_23, unused0, unused1, unused2; + .reg .b16 sf_01, sf_23, e4m3_01, e4m3_23; + .reg .b32 tmp, h_01, h_23, sf_h_01, sf_h_23; + and.b32 tmp, {$r0}, 0x0000000f; + bfe.u32 h_01, {$r0}, 8, 4; + shl.b32 h_01, h_01, 4; + or.b32 h_01, h_01, tmp; + mov.b32 {fp4_01, unused0, unused1, unused2}, h_01; + bfe.u32 h_23, {$r0}, 16, 4; + bfe.u32 tmp, {$r0}, 24, 4; + shl.b32 tmp, tmp, 4; + or.b32 h_23, h_23, tmp; + mov.b32 {fp4_23, unused0, unused1, unused2}, h_23; + mov.b32 {sf_01, sf_23}, {$r1}; + cvt.rn.f16x2.e2m1x2 h_01, fp4_01; + cvt.rn.f16x2.e2m1x2 h_23, fp4_23; + cvt.rn.f16x2.e4m3x2 sf_h_01, sf_01; + cvt.rn.f16x2.e4m3x2 sf_h_23, sf_23; + mul.rn.f16x2 h_01, h_01, sf_h_01; + mul.rn.f16x2 h_23, h_23, sf_h_23; + cvt.rn.satfinite.e4m3x2.f16x2 e4m3_01, h_01; + cvt.rn.satfinite.e4m3x2.f16x2 e4m3_23, h_23; + mov.b32 {$w0}, {e4m3_01, e4m3_23}; + } + """, + write_only_types=[Int32], + read_only_args=[packed_fp4, packed_sf], + ) + + +@cute.jit +def _mul_e2m1x4_e4m3x4_bx( + packed_fp4: Int32, packed_sf: Int32, sf_byte_idx: Constexpr[int] +) -> Int32: + """Apply one selected E4M3 scale byte to four unpacked E2M1 values.""" + imm = "0x0000" + if cutlass.const_expr(sf_byte_idx == 1): + imm = "0x1111" + elif cutlass.const_expr(sf_byte_idx == 2): + imm = "0x2222" + elif cutlass.const_expr(sf_byte_idx == 3): + imm = "0x3333" + broadcast_sf = cute.arch.inline_ptx( + f"prmt.b32 {{$w0}}, {{$r0}}, {{$r0}}, {imm};", + write_only_types=[Int32], + read_only_args=[packed_sf], + ) + return _mul_e2m1x4_e4m3x4(packed_fp4, broadcast_sf) + + +@dataclass(kw_only=True) +class TmemTransformedKvResource(DecodeGenResourceBase): + """Pipeline transformed NVFP4 K/V through TMEM. + + The transform warp group dequantizes an unpacking-TMA NVFP4 stage from + :class:`SmemKvResource` into E4M3 registers. K is stored token-major and V + is transposed, so both BMMs consume the transformed tile as operand A. + """ + + _task_local_specs: ClassVar[tuple[tuple, ...]] = ( + ( + "kv_desc_slot", + Int32, + Int32(0), + "TMEM address for transformed K consumed by QK MMA.", + ), + ( + "v_desc_0_slot", + Int32, + Int32(0), + "TMEM address for transformed V consumed by VP MMA instance 0.", + ), + ( + "v_desc_1_slot", + Int32, + Int32(0), + "TMEM address for transformed V consumed by VP MMA instance 1.", + ), + ) + cfg: Constexpr[FmhaDecodeConfig] = None + src_smem_kv: Constexpr[SmemKvResource] = None + _alloc: Constexpr[TmemAllocation | None] = None + kv_desc_slot: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + v_desc_0_slot: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + v_desc_1_slot: Constexpr[TaskLocalVariable] = TaskLocalVariable.uninitialized() + + def get_smem_requirements(self) -> list[SmemAllocation]: + """Return no data allocation; the pipeline barriers still use SMEM.""" + return [] + + def get_tmem_requirements(self) -> list[TmemAllocation]: + """Allocate the staged transformed K/V TMEM ring.""" + if self._alloc is None: + self._alloc = TmemAllocation( + name=f"{self.name}", + num_columns=( + self.cfg.tmem_transformed_kv_stage_cols + * self.pipeline_config.num_stages + ), + ) + return [self._alloc] + + @cute.jit + def _create_initial_task_locals( + self, context: ResourceContext | None = None + ) -> ResourceVars: + """Initialize the stage-address task locals.""" + _ = context + return { + "kv_desc": Int32(0), + "v_desc_0": Int32(0), + "v_desc_1": Int32(0), + } + + @cute.jit + def _create_work_tile_task_locals( + self, context: ResourceContext | None = None + ) -> ResourceVars: + """Initialize one work tile's stage-address task locals.""" + _ = context + return { + "kv_desc": Int32(0), + "v_desc_0": Int32(0), + "v_desc_1": Int32(0), + } + + @producer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_load_state(self, stage_info: StageInfo) -> None: + """Initialize producer-side TMEM stage state.""" + self._create_initial_task_locals(stage_info.context) + + @consumer_work(work_attrs=WorkAttr.AUXILIARY) + @cute.jit + def init_descriptor_state(self, stage_info: StageInfo) -> None: + """Initialize consumer-side TMEM stage state.""" + self._create_initial_task_locals(stage_info.context) + + @cute.jit + def _source_stage(self) -> tuple[cutlass.Array, cutlass.Array]: + """Return unpacking-TMA NVFP4 and scale arrays for the raw stage.""" + cfg = self.cfg + src_stage_idx = self.src_smem_kv.state_src.consumer_work_stage + src_byte_base = self.src_smem_kv._smem_base_kv.subview( + src_stage_idx * cfg.smem_kv_tile_elements + ) + src_bytes = cutlass.Array( + src_byte_base.data_ptr(), + dtype=Uint8, + shape=(cfg.smem_kv_tile_elements,), + addrspace=3, + ) + return src_bytes, self.src_smem_kv._sf_stage_base(src_stage_idx) + + @cute.jit + def _destination_stage_addr(self, stage_info: StageInfo) -> Int32: + """Return the raw TMEM address of one transformed pipeline stage.""" + task_cache = _decode_gen_task_cache(stage_info) + return ( + task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + + Int32(self._alloc.offset) + + stage_info.stage_idx * Int32(self.cfg.tmem_transformed_kv_stage_cols) + ) + + @cute.jit + def _store_k( + self, stage_info: StageInfo, head_dim_stage_idx: Constexpr[int] + ) -> None: + """Dequantize K with unpacking LDSM and store token-major TMEM.""" + cfg = self.cfg + task_cache = _decode_gen_task_cache(stage_info) + transform_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] - Int32( + (cfg.transform_kv_warp_idx % 4) * 32 + ) + warp_idx = transform_thread_idx >> Int32(5) + lane_idx = transform_thread_idx & Int32(31) + src_bytes, src_sf = self._source_stage() + src_sf_qwords = cutlass.Array( + src_sf.data_ptr(), + dtype=Int64, + shape=(cfg.smem_kv_sf_tile_bytes // 8,), + addrspace=3, + ) + + # Each lane owns four K rows. H256 stages the complete 16-byte SF row, + # so select the qword matching the active 128-column head slice. + sf_rows = [] + sf_row_base = warp_idx * Int32(32) + lane_idx // Int32(4) + sf_qwords_per_row = Int32(cfg.smem_kv_sf_bytes_per_token // 8) + for row_group in cutlass.range_constexpr(4): + sf_row = sf_row_base + Int32(row_group * 8) + sf_qword = src_sf_qwords[ + sf_row * sf_qwords_per_row + Int32(head_dim_stage_idx) + ] + sf_rows.append((Int32(sf_qword), Int32(sf_qword >> Int64(32)))) + + thread_row = warp_idx * Int32(32) + lane_idx + swizzle_mask = (thread_row & Int32(7)) * Int32(16) + stage_addr = self._destination_stage_addr(stage_info) + for col_stage in cutlass.range_constexpr(8): + dst_regs = [] + for row_block in cutlass.range_constexpr(2): + local_offset = Int32( + row_block * 16 * cfg.head_dim_kv_stage + col_stage * 16 + ) + src_ptr = src_bytes.data_ptr() + ( + thread_row * Int32(cfg.head_dim_kv_stage) + + (local_offset ^ swizzle_mask) + ) + unpacked = prims.ldmatrix( + src_ptr, + 2, + prims.MMALayout.ROW, + shape=prims.LoadShape.M8N16, + src_format=prims.LoadSrcFormat.B4X16_P64, + ) + sf_word_idx = col_stage // 4 + sf_byte_idx = col_stage % 4 + dst_regs.append( + _mul_e2m1x4_e4m3x4_bx( + unpacked[0], + sf_rows[row_block * 2][sf_word_idx], + sf_byte_idx, + ) + ) + dst_regs.append( + _mul_e2m1x4_e4m3x4_bx( + unpacked[1], + sf_rows[row_block * 2 + 1][sf_word_idx], + sf_byte_idx, + ) + ) + + tmem_addr = stage_addr + Int32(col_stage * 4) + prims.tcgen05_st( + "16x128b", + prims.make_tmem_ptr(tmem_addr, Int32), + cutlass.Vector.from_elements(tuple(dst_regs[:2]), dtype=Int32), + ) + prims.tcgen05_st( + "16x128b", + prims.make_tmem_ptr(tmem_addr + (Int32(16) << Int32(16)), Int32), + cutlass.Vector.from_elements(tuple(dst_regs[2:]), dtype=Int32), + ) + cute.arch.fence_view_async_tmem_store() + + @cute.jit + def _store_v( + self, stage_info: StageInfo, head_dim_stage_idx: Constexpr[int] + ) -> None: + """Dequantize and transpose V with LDSM into TMEM operand A.""" + cfg = self.cfg + task_cache = _decode_gen_task_cache(stage_info) + transform_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] - Int32( + (cfg.transform_kv_warp_idx % 4) * 32 + ) + warp_idx = transform_thread_idx >> Int32(5) + lane_idx = transform_thread_idx & Int32(31) + src_bytes, src_sf = self._source_stage() + src_sf_words = cutlass.Array( + src_sf.data_ptr(), + dtype=Int32, + shape=(cfg.smem_kv_sf_tile_bytes // 4,), + addrspace=3, + ) + + swizzle_mask = (lane_idx & Int32(7)) * Int32(16) + stage_addr = self._destination_stage_addr(stage_info) + sf_bytes_per_token = cfg.smem_kv_sf_bytes_per_token + sf_head_stage_offset = head_dim_stage_idx * (cfg.head_dim_kv_stage // 16) + for token_stage in cutlass.range_constexpr(8): + # V scale factors use TRT-LLM's 4-token interleaved layout, so the + # two packed QMUL4 vectors are adjacent 32-bit words. + sf_byte_offset = ( + Int32(token_stage * 16 * sf_bytes_per_token) + + (lane_idx % Int32(4)) * Int32(4 * sf_bytes_per_token) + + Int32(sf_head_stage_offset * 4) + + warp_idx * Int32(8) + ) + sf_word_idx = sf_byte_offset // Int32(4) + sf_packed0 = src_sf_words[sf_word_idx] + sf_packed1 = src_sf_words[sf_word_idx + Int32(1)] + + dst_regs = [] + for col_block in cutlass.range_constexpr(2): + local_offset = Int32( + token_stage * 16 * cfg.head_dim_kv_stage + col_block * 16 + ) + warp_idx * Int32(32) + src_ptr = src_bytes.data_ptr() + ( + lane_idx * Int32(cfg.head_dim_kv_stage) + + (local_offset ^ swizzle_mask) + ) + unpacked = prims.ldmatrix( + src_ptr, + 2, + prims.MMALayout.COL, + shape=prims.LoadShape.M16N16, + src_format=prims.LoadSrcFormat.B4X16_P64, + ) + packed_sf = sf_packed0 if col_block == 0 else sf_packed1 + dst_regs.append(_mul_e2m1x4_e4m3x4(unpacked[0], packed_sf)) + dst_regs.append(_mul_e2m1x4_e4m3x4(unpacked[1], packed_sf)) + + tmem_addr = stage_addr + Int32(token_stage * 4) + prims.tcgen05_st( + "16x128b", + prims.make_tmem_ptr(tmem_addr, Int32), + cutlass.Vector.from_elements(tuple(dst_regs[:2]), dtype=Int32), + ) + prims.tcgen05_st( + "16x128b", + prims.make_tmem_ptr(tmem_addr + (Int32(16) << Int32(16)), Int32), + cutlass.Vector.from_elements(tuple(dst_regs[2:]), dtype=Int32), + ) + cute.arch.fence_view_async_tmem_store() + + @producer_work + @cute.jit + def transform_k0( + self, stage_info: StageInfo, *, head_dim_stage_idx: Constexpr[int] + ) -> None: + """Transform K instance 0 into the acquired TMEM stage.""" + self._store_k(stage_info, head_dim_stage_idx) + + @producer_work + @cute.jit + def transform_k1( + self, stage_info: StageInfo, *, head_dim_stage_idx: Constexpr[int] + ) -> None: + """Transform K instance 1 into the acquired TMEM stage.""" + self._store_k(stage_info, head_dim_stage_idx) + + @producer_work + @cute.jit + def transform_v0( + self, stage_info: StageInfo, *, head_dim_stage_idx: Constexpr[int] + ) -> None: + """Transform V instance 0 into the acquired TMEM stage.""" + self._store_v(stage_info, head_dim_stage_idx) + + @producer_work + @cute.jit + def transform_v1( + self, stage_info: StageInfo, *, head_dim_stage_idx: Constexpr[int] + ) -> None: + """Transform V instance 1 into the acquired TMEM stage.""" + self._store_v(stage_info, head_dim_stage_idx) + + @consumer_work(returns=kv_desc_slot) + @cute.jit + def k_desc_0(self, stage_info: StageInfo) -> Int32: + """Publish the transformed K TMEM address for QK instance 0.""" + return self._destination_stage_addr(stage_info) + + @consumer_work(returns=kv_desc_slot) + @cute.jit + def k_desc_1(self, stage_info: StageInfo) -> Int32: + """Publish the transformed K TMEM address for QK instance 1.""" + return self._destination_stage_addr(stage_info) + + @consumer_work(returns=v_desc_0_slot) + @cute.jit + def v_desc_0(self, stage_info: StageInfo) -> Int32: + """Publish the transformed V TMEM address for VP MMA instance 0.""" + return self._destination_stage_addr(stage_info) + + @consumer_work(returns=v_desc_1_slot) + @cute.jit + def v_desc_1(self, stage_info: StageInfo) -> Int32: + """Publish the transformed V TMEM address for VP MMA instance 1.""" + return self._destination_stage_addr(stage_info) 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 0483ca9b3519..49cd91cde855 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 @@ -52,6 +52,8 @@ from cutlass.experimental.task_scheduling.task import Task from ..stage import FmhaStage +from ..task_compat import task_context_kwargs +from .direct_sparse_metadata import DirectSparseMetadataView from .fmha_decode_config import FmhaDecodeConfig from .fmha_decode_constants import ( KV_INST0, @@ -71,6 +73,14 @@ _sliding_window_start_idx, ) +LocalStatsState = tuple[ + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, + cutlass.Array, +] QBoundBinding = tuple[MemoryResource, bool] TaskKwarg = ( int @@ -135,6 +145,11 @@ def _loop_domain_after_head( return (remaining + insts - cutlass.Int32(1)) // insts +@dataclass(kw_only=True) +class SparseMembershipLifetimeResource(MemoryResource): + """Keep the membership row live until every softmax stream releases it.""" + + @dataclass(kw_only=True) class ScheduleTokenThrottleResource(MemoryResource): """Order persistent load consumption before the scheduler reuses a slot.""" @@ -217,7 +232,7 @@ def _page_offsets_produce( smem_page_offsets.commit() -def _can_hold_native_split_page_window( +def _can_hold_native_page_window( cfg: FmhaDecodeConfig, smem_page_offsets: MemoryResource | None, ) -> bool: @@ -225,11 +240,19 @@ def _can_hold_native_split_page_window( if ( smem_page_offsets is None or not smem_page_offsets.use_native_paged_kv - or not cfg.use_split_kv or cfg.use_sliding_window_causal ): return False 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 + # instruction-aligned local span. Holding that stage lets + # K and V reuse the same dense-row locators, including -1 padding for + # a short or odd tail. This applies to direct and persistent kernels as + # well as split-KV: each work tile stages its own slice of that row. + return smem_page_offsets.holds_encoded_locator_window + if not cfg.use_split_kv: + return False # Runtime ragged lengths can reduce the split-local span in # ``num_insts_kv`` increments and therefore change every split rank's # aligned starting page. Hold one stage only when every possible span @@ -263,7 +286,11 @@ def _staged_kv_load( same token tile and therefore the same page IDs. Keep one page-offset consumer stage live across the complete logical K/V load. """ - reuse_page_ids = smem_page_offsets is not None and cfg.num_head_dim_stages_kv > 1 + reuse_page_ids = ( + smem_page_offsets is not None + and cfg.num_head_dim_stages_kv > 1 + and not cfg.uses_scattered_page_route + ) # Optional ConsWait/ConsWork: fetch page IDs for this logical K/V tile. # The cached D256 path performs its ConsumerWork below while materializing # the register array; single-stage kernels keep the original no-cache path. @@ -315,6 +342,21 @@ def _produce_staged_page_offsets( _page_offsets_produce(smem_page_offsets, label, section) +def _issue_score_seed( + tmem_s: MemoryResource, section: FmhaStage, cfg: FmhaDecodeConfig +) -> None: + """Seed the acquired S slot of an INT32-score QK wave before its K wait. + + The seed MMA reads only constant tiles, so its issue overlaps the K wait. + """ + if not cfg.uses_int32_scores: + return + if section == FmhaStage.Head: + tmem_s.seed_scores_head() + else: + tmem_s.seed_scores_loop() + + def _consume_staged_qk_mma( smem_kv: MemoryResource, tmem_s: MemoryResource, @@ -332,6 +374,7 @@ def _consume_staged_qk_mma( MMA's accumulator write, so no completion wait is needed before QK. """ tmem_s.acquire() + _issue_score_seed(tmem_s, section, cfg) for head_dim_stage_idx in range(cfg.num_head_dim_stages_kv): smem_kv.wait() kv_desc = getattr(smem_kv, k_desc_label)() @@ -560,6 +603,8 @@ def _prepared_sparse_row_address( q_token_base = _q_group_token_base(cfg, q_group_idx) q_block = q_token_base // cutlass.Int32(cfg.q_block_size) num_q_blocks = (cfg.max_seq_len_q + cfg.q_block_size - 1) // cfg.q_block_size + if cutlass.const_expr(cfg.shares_sparse_pattern): + return b_idx * cutlass.Int32(num_q_blocks) + q_block return (b_idx * num_heads_kv + h_idx) * cutlass.Int32(num_q_blocks) + q_block @@ -717,6 +762,7 @@ def _refresh_packed_q_bounds(self, work_tile: Any) -> None: def _run_packed_skip_iteration( self, work_tile: Any, + context: ResourceContext | None = None, ) -> None: """Advance one inactive tile through WorkQueue bookkeeping only.""" # Packed schedules place every data-path entry inside ``skippable()``; @@ -731,6 +777,7 @@ def _run_packed_skip_iteration( head_entries, work_tile, bookkeeping_domain, + **task_context_kwargs(context), ) for is_skippable_tail, tail_entries in self._tail_exec_groups: if cutlass.const_expr(not is_skippable_tail): @@ -738,6 +785,7 @@ def _run_packed_skip_iteration( tail_entries, work_tile, bookkeeping_domain, + **task_context_kwargs(context), ) @cute.jit @@ -745,12 +793,14 @@ def _run_task_body_impl( self, work_tile: cute.Coord, skip_work_tile: Any = None, + context: ResourceContext | None = None, ) -> None: """Run one ordinary task tile and synchronize attention-sink tails.""" Task._run_task_body_impl( self, work_tile, skip_work_tile, + **task_context_kwargs(context), ) if cutlass.const_expr( self.cfg is not None @@ -778,7 +828,10 @@ def _run_task_body_impl( prims.barrier_cta_sync(12, thread_count=16 * 32) @cute.jit - def _run_task_body_persistent(self) -> None: + def _run_task_body_persistent( + self, + context: ResourceContext | None = None, + ) -> None: """Drain inactive packed tiles before each unconditional active body.""" use_packed_early_stop = ( self.cfg is not None @@ -787,14 +840,14 @@ def _run_task_body_persistent(self) -> None: and self._has_skip_if ) if cutlass.const_expr(not use_packed_early_stop): - Task._run_task_body_persistent(self) + Task._run_task_body_persistent(self, **task_context_kwargs(context)) return assert self.work_queue is not None work_tile = self.work_queue.initial_work_tile_info() self.work_queue._set_consumer_var_from_ts("work_tile", work_tile) - self._run_pre_work_loop_entries(work_tile) + self._run_pre_work_loop_entries(work_tile, **task_context_kwargs(context)) work_tile = self.work_queue._get_consumer_var_from_ts("work_tile") for resource in self.dst_resources: if cutlass.const_expr( @@ -808,7 +861,7 @@ def _run_task_body_persistent(self) -> None: # inner loop executes only the non-skippable WorkQueue tail, so no TMA, # descriptor, pipeline, task data, or sink barrier is issued. while work_tile.is_valid_tile and self._should_skip_work_tile(work_tile): - self._run_packed_skip_iteration(work_tile) + self._run_packed_skip_iteration(work_tile, context) work_tile = self.work_queue._get_consumer_var_from_ts("work_tile") self.dummy = cutlass.Boolean(True) @@ -817,18 +870,20 @@ def _run_task_body_persistent(self) -> None: # 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. - Task._run_task_body_impl(self, work_tile, None) + Task._run_task_body_impl( + self, work_tile, None, **task_context_kwargs(context) + ) if cutlass.const_expr(self.cfg.use_attention_sinks): prims.barrier_cta_sync(12, thread_count=16 * 32) work_tile = self.work_queue._get_consumer_var_from_ts("work_tile") self.dummy = cutlass.Boolean(True) while work_tile.is_valid_tile and self._should_skip_work_tile(work_tile): - self._run_packed_skip_iteration(work_tile) + self._run_packed_skip_iteration(work_tile, context) work_tile = self.work_queue._get_consumer_var_from_ts("work_tile") self.dummy = cutlass.Boolean(True) - self._run_post_work_loop_entries(work_tile) + self._run_post_work_loop_entries(work_tile, **task_context_kwargs(context)) for resource in self.dst_resources: if cutlass.const_expr( resource.pipeline_config is not None @@ -910,20 +965,31 @@ def get_domain(self, tile_coord: cute.Coord) -> cutlass.Int32 | int: # can use the configured max length; variable-seqlen kernels read the # batch-specific length from GMEM. b_idx = cutlass.Int32(tile_coord[2]) - if cutlass.const_expr(self.block_table_capacity is not None): - self._kv_page_idx_ub = cutlass.Int32( - self.block_table_capacity - ) - cutlass.Int32(1) if self.seqlens_kv is None: seq_len_kv = cutlass.Int32(self.max_seq_len_kv) else: - seq_len_kv = cutlass.Int32(self.seqlens_kv[b_idx]) + metadata_row = cutlass.Int64(b_idx) + if cutlass.const_expr( + self.cfg.uses_q_token_kv_block_sparse_page_route + and self.cfg.use_persistent_scheduler + and not self.cfg.shares_sparse_pattern + and not isinstance(self.seqlens_kv, DirectSparseMetadataView) + ): + metadata_row = metadata_row * cutlass.Int64( + self.num_heads_kv + ) + cutlass.Int64(tile_coord[1]) + seq_len_kv = cutlass.Int32(self.seqlens_kv[metadata_row]) self._seq_len_kv = seq_len_kv - if cutlass.const_expr(self.block_table_capacity is not None): - self._kv_page_idx_ub = cute.math.min( - self._kv_page_idx_ub, - _runtime_last_valid_page_idx(self.cfg, seq_len_kv), - ) + if cutlass.const_expr(self.cfg.use_paged_kv): + # Native paged KV uses a fixed-stride dense page table. The row's + # live extent therefore comes from seq_lens; padded entries beyond + # this bound must never be staged. + self._kv_page_idx_ub = _runtime_last_valid_page_idx(self.cfg, seq_len_kv) + if cutlass.const_expr(self.block_table_capacity is not None): + self._kv_page_idx_ub = cute.math.min( + self._kv_page_idx_ub, + cutlass.Int32(self.block_table_capacity) - cutlass.Int32(1), + ) if ( self.seqlens_kv is None and not self.cfg.use_split_kv @@ -1012,7 +1078,7 @@ def get_domain(self, tile_coord: cute.Coord) -> cutlass.Int32 | int: # ====================================================================== -# LoadTask — warp 13 (or warp 15 under CLC persistent), 1 warp +# LoadTask — independently configured Q/K/V TMA issuer warps # K and V share a single SmemKv ring; loads alternate K and V tiles. # HEAD: Q + K0 + K1 # LOOP[i]: K(i+2) + V(i) @@ -1126,7 +1192,9 @@ def create_load_task( **kw: TaskKwarg, ) -> Task: """Create the shared-KV load task and optional page-offset dependency.""" - hold_page_window = _can_hold_native_split_page_window(cfg, smem_page_offsets) + hold_page_window = _can_hold_native_page_window(cfg, smem_page_offsets) + head_labels = ("load_k0",) if cfg.num_insts_kv == 1 else ("load_k0", "load_k1") + tail_labels = ("load_v0",) if cfg.num_insts_kv == 1 else ("load_v0", "load_v1") def load_schedule_body( smem_q: MemoryResource, @@ -1151,7 +1219,7 @@ def load_schedule_body( sparse_resource.init_load_state() cached_page_ids = None if smem_page_offsets is not None: - if cfg.num_head_dim_stages_kv > 1: + if cfg.num_head_dim_stages_kv > 1 and not cfg.uses_scattered_page_route: cached_page_ids = smem_page_offsets.init_cached_read_state() else: smem_page_offsets.init_read_state() @@ -1165,56 +1233,59 @@ def _kv_load(label: str, section: FmhaStage) -> None: section, cfg, smem_page_offsets=smem_page_offsets, + page_offsets_label=label.replace("load_", "read_offsets_"), manage_page_offsets=not hold_page_window, cached_page_ids=cached_page_ids, ) - # HEAD: load Q once, then prefetch the first two K tiles. + route_metadata_by_label = { + "load_k0": (sparse_kv_metadata0, sparse_softmax_metadata0), + "load_k1": (sparse_kv_metadata1, sparse_softmax_metadata1), + } + + # HEAD: load Q once, then prefetch K for each active instance. smem_q.acquire() smem_q.tma_load() smem_q.commit() if hold_page_window: # The native K/V caches share one page table. Keep its single - # 32-ID consumer stage live across every K/V and head-dimension - # load owned by this split CTA, matching the reference page-window - # lifetime and avoiding redundant pipeline handoffs. - if cfg.num_head_dim_stages_kv > 1: + # consumer stage live across every K/V and head-dimension load + # owned by this CTA work tile. Sparse routes size it for the full local + # span; other native split routes retain the aligned 32-ID window. + if cfg.num_head_dim_stages_kv > 1 and not cfg.uses_scattered_page_route: smem_page_offsets.wait() else: _page_offsets_consume(smem_page_offsets) # Dense profiles have no route metadata: the resolve and publish # helpers are no-ops for ``None`` resources, so one cadence serves # both dense and block-sparse loads. - route0, prefetch0 = _resolve_and_store_sparse_route( - sparse_kv_metadata0, FmhaStage.Head - ) - _kv_load("load_k0", FmhaStage.Head) - route1, prefetch1 = _resolve_and_store_sparse_route( - sparse_kv_metadata1, FmhaStage.Head - ) - _kv_load("load_k1", FmhaStage.Head) - # Issue both K tiles before either metadata FIFO can backpressure - # the load warp, matching the split-resource sparse cadence. - _publish_sparse_softmax_route(sparse_softmax_metadata0, route0) - _publish_sparse_softmax_route(sparse_softmax_metadata1, route1) - prefetch_by_label = {"load_k0": prefetch0, "load_k1": prefetch1} + prefetch_by_label = {} + head_routes = [] + for label in head_labels: + kv_metadata, softmax_metadata = route_metadata_by_label[label] + route, prefetch_by_label[label] = _resolve_and_store_sparse_route( + kv_metadata, FmhaStage.Head + ) + _kv_load(label, FmhaStage.Head) + if route is not None: + head_routes.append((softmax_metadata, route)) + # Issue all K tiles before a metadata FIFO can backpressure the load warp. + for softmax_metadata, route in head_routes: + _publish_sparse_softmax_route(softmax_metadata, route) # LOOP: each iter prefetches the full ``num_insts_kv`` K/V pair set. # When P aliases the consumed S columns, MMA must consume each V/P pair # before the following same-instance QK overwrites S. Keep the # shared-ring producer order identical to that consumer order. - loop_labels = ( - ("load_v0", "load_k0", "load_v1", "load_k1") - if cfg.uses_two_inst_tmem_p - else ("load_k0", "load_v0", "load_k1", "load_v1") - ) + if cfg.num_insts_kv == 1: + loop_labels = ("load_k0", "load_v0") + elif cfg.uses_two_inst_tmem_p: + loop_labels = ("load_v0", "load_k0", "load_v1", "load_k1") + else: + loop_labels = ("load_k0", "load_v0", "load_k1", "load_v1") with domain_loop(0, domain, 1, unroll=1): # Generic V-first profiles consume their retained route before the # matching K label replaces it. - route_metadata_by_label = { - "load_k0": (sparse_kv_metadata0, sparse_softmax_metadata0), - "load_k1": (sparse_kv_metadata1, sparse_softmax_metadata1), - } loop_routes = [] for label in loop_labels: kv_metadata, softmax_metadata = route_metadata_by_label.get( @@ -1229,9 +1300,8 @@ def _kv_load(label: str, section: FmhaStage) -> None: for sparse_softmax_metadata, route in loop_routes: _publish_sparse_softmax_route(sparse_softmax_metadata, route) - # TAIL: after no more future K tiles are needed, load the final two V - # tiles consumed by the final BMM2 calls. - for label in ("load_v0", "load_v1"): + # TAIL: load the final V tile for each active instance's last BMM2. + for label in tail_labels: _kv_load(label, FmhaStage.Tail) if hold_page_window: _page_offsets_release(smem_page_offsets) @@ -1326,27 +1396,36 @@ def create_page_offsets_task( cfg: FmhaDecodeConfig, *, domain: int | cutlass.Int32, + membership_lifetime: MemoryResource | None = None, warp_idx: int | None = None, num_warps: int | None = None, task_class: type[DecodeGenTask] = DecodeGenTask, **kw: TaskKwarg, ) -> Task: - """Prefetch page-table entries that LoadTask consumes for paged KV. + """Prefetch page-table entries for the shared-KV load task. - The schedule matches LoadTask's K/V cadence exactly so the page-offsets - ring and the SmemKv ring stay aligned. + A held native window is published once per CTA work item and retained + through the V tail. Other profiles mirror LoadTask's K/V cadence. """ - hold_page_window = _can_hold_native_split_page_window(cfg, smem_page_offsets) + hold_page_window = _can_hold_native_page_window(cfg, smem_page_offsets) + + head_labels = ("load_k0",) if cfg.num_insts_kv == 1 else ("load_k0", "load_k1") + tail_labels = ("load_v0",) if cfg.num_insts_kv == 1 else ("load_v0", "load_v1") def page_offsets_schedule_body( smem_page_offsets: MemoryResource, + membership_lifetime: MemoryResource | None, ) -> None: """Schedule page-offset prefetches for the shared-KV load cadence.""" + if membership_lifetime is not None: + membership_lifetime.acquire() smem_page_offsets.init_load_state() if hold_page_window: - # K0's aligned 32-ID window covers every contiguous tile assigned - # to this split CTA, and the native table uses those IDs for V too. + # K0's held window covers every tile assigned to this CTA work + # item, and native paged KV uses those same locators for V. _page_offsets_produce(smem_page_offsets, "load_k0", FmhaStage.Head) + if membership_lifetime is not None: + membership_lifetime.commit() # Preserve the runtime domain contract even though this fast path # needs no per-iteration page-window work. with domain_loop(0, domain, 1, unroll=1): @@ -1355,15 +1434,18 @@ def page_offsets_schedule_body( # Page offsets are staged by a separate producer so the load warp can # issue K/V TMA copies without reading page tables itself. # HEAD: produce page IDs for the two prefetched K tiles. - _produce_staged_page_offsets(smem_page_offsets, "load_k0", FmhaStage.Head, cfg) - _produce_staged_page_offsets(smem_page_offsets, "load_k1", FmhaStage.Head, cfg) + for label in head_labels: + _produce_staged_page_offsets(smem_page_offsets, label, FmhaStage.Head, cfg) + if label == "load_k0" and membership_lifetime is not None: + membership_lifetime.commit() # LOOP: mirror LoadTask's K/V production cadence exactly. - loop_labels = ( - ("load_v0", "load_k0", "load_v1", "load_k1") - if cfg.uses_two_inst_tmem_p - else ("load_k0", "load_v0", "load_k1", "load_v1") - ) + if cfg.num_insts_kv == 1: + loop_labels = ("load_k0", "load_v0") + elif cfg.uses_two_inst_tmem_p: + loop_labels = ("load_v0", "load_k0", "load_v1", "load_k1") + else: + loop_labels = ("load_k0", "load_v0", "load_k1", "load_v1") with domain_loop(0, domain, 1, unroll=1): for label in loop_labels: _produce_staged_page_offsets( @@ -1371,32 +1453,32 @@ def page_offsets_schedule_body( ) # TAIL: produce page IDs for the final two V loads. - for label in ("load_v0", "load_v1"): + for label in tail_labels: _produce_staged_page_offsets(smem_page_offsets, label, FmhaStage.Tail, cfg) - @schedule + @_schedule_with_optional_resources def page_offsets_schedule( smem_page_offsets: MemoryResource, + membership_lifetime: MemoryResource | None, work_queue: WorkQueue | None = None, ) -> None: """Wrap page-offset data work in packed persistent skip handling.""" _decode_work_tile_schedule( cfg, work_queue, - lambda: page_offsets_schedule_body(smem_page_offsets), + lambda: page_offsets_schedule_body(smem_page_offsets, membership_lifetime), ) - captured_schedule = ( - page_offsets_schedule(smem_page_offsets) - if work_queue is None - else page_offsets_schedule(smem_page_offsets, work_queue) + captured_schedule = page_offsets_schedule( + smem_page_offsets, membership_lifetime, work_queue ) src = [] if work_queue is not None: src.append(work_queue) return task_class( src_resources=src, - dst_resources=[smem_page_offsets], + dst_resources=[smem_page_offsets] + + ([membership_lifetime] if membership_lifetime is not None else []), q_bound_resources=((smem_page_offsets, False),), cfg=cfg, warp_idx=cfg.page_offsets_warp_idx if warp_idx is None else warp_idx, @@ -1415,6 +1497,7 @@ def create_page_offsets_task_split_kv( cfg: FmhaDecodeConfig, *, domain: int | cutlass.Int32, + membership_lifetime: MemoryResource | None = None, warp_idx: int | None = None, num_warps: int | None = None, task_class: type[DecodeGenTask] = DecodeGenTask, @@ -1425,13 +1508,18 @@ def create_page_offsets_task_split_kv( def page_offsets_schedule_body( smem_page_offsets_k: MemoryResource, smem_page_offsets_v: MemoryResource, + membership_lifetime: MemoryResource | None, ) -> None: """Publish one page-offset stage for every paired K/V tile.""" + if membership_lifetime is not None: + membership_lifetime.acquire() smem_page_offsets_k.init_load_state() smem_page_offsets_v.init_load_state() # HEAD: publish the initial K0/K1 pair as two independent stages. _page_offsets_produce(smem_page_offsets_k, "load_k0", FmhaStage.Head) + if membership_lifetime is not None: + membership_lifetime.commit() _page_offsets_produce(smem_page_offsets_k, "load_k1", FmhaStage.Head) # LOOP: mirror LoadTask's cross-resource consumption order exactly. @@ -1445,10 +1533,11 @@ def page_offsets_schedule_body( _page_offsets_produce(smem_page_offsets_v, "load_v0", FmhaStage.Tail) _page_offsets_produce(smem_page_offsets_v, "load_v1", FmhaStage.Tail) - @schedule + @_schedule_with_optional_resources def page_offsets_schedule( smem_page_offsets_k: MemoryResource, smem_page_offsets_v: MemoryResource, + membership_lifetime: MemoryResource | None, work_queue: WorkQueue | None = None, ) -> None: """Wrap split page-offset work in packed persistent skip handling.""" @@ -1458,20 +1547,20 @@ def page_offsets_schedule( lambda: page_offsets_schedule_body( smem_page_offsets_k, smem_page_offsets_v, + membership_lifetime, ), ) - schedule_result = ( - page_offsets_schedule(smem_page_offsets_k, smem_page_offsets_v) - if work_queue is None - else page_offsets_schedule(smem_page_offsets_k, smem_page_offsets_v, work_queue) + schedule_result = page_offsets_schedule( + smem_page_offsets_k, smem_page_offsets_v, membership_lifetime, work_queue ) src = [] if work_queue is not None: src.append(work_queue) return task_class( src_resources=src, - dst_resources=[smem_page_offsets_k, smem_page_offsets_v], + dst_resources=[smem_page_offsets_k, smem_page_offsets_v] + + ([membership_lifetime] if membership_lifetime is not None else []), q_bound_resources=( (smem_page_offsets_k, False), (smem_page_offsets_v, False), @@ -1492,48 +1581,67 @@ def create_page_offsets_task_one_inst_qkv( cfg: FmhaDecodeConfig, *, domain: int | cutlass.Int32, + membership_lifetime: MemoryResource | None = None, warp_idx: int | None = None, num_warps: int | None = None, task_class: type[DecodeGenTask] = DecodeGenTask, **kw: TaskKwarg, ) -> Task: - """Prefetch page-table entries for the one-inst keepsMmaAb QKV path.""" + """Prefetch offsets for one-instance Keeps, using the configured producer warps.""" + hold_page_window = _can_hold_native_page_window(cfg, smem_page_offsets) def page_offsets_schedule_body( smem_page_offsets: MemoryResource, + membership_lifetime: MemoryResource | None, ) -> None: """Schedule page-offset prefetches for the one-inst QKV load cadence.""" + if membership_lifetime is not None: + membership_lifetime.acquire() smem_page_offsets.init_load_state() + if hold_page_window: + # Publish the complete CTA-local physical-KV128 locator span once. + # The load task retains this stage from its K head through V tail. + _page_offsets_produce(smem_page_offsets, "load_k0", FmhaStage.Head) + if membership_lifetime is not None: + membership_lifetime.commit() + # Keep the task domain identical to the ordinary cadence even + # though the held stage needs no per-iteration producer work. + with domain_loop(0, domain, 1, unroll=1): + pass + return + _page_offsets_produce(smem_page_offsets, "load_k0", FmhaStage.Head) + if membership_lifetime is not None: + membership_lifetime.commit() with domain_loop(0, domain, 1, unroll=1): for label in ("load_k0", "load_v0"): _page_offsets_produce(smem_page_offsets, label, FmhaStage.Loop) _page_offsets_produce(smem_page_offsets, "load_v0", FmhaStage.Tail) - @schedule + @_schedule_with_optional_resources def page_offsets_schedule( smem_page_offsets: MemoryResource, + membership_lifetime: MemoryResource | None, work_queue: WorkQueue | None = None, ) -> None: """Wrap one-inst page-offset work in packed persistent skip handling.""" _decode_work_tile_schedule( cfg, work_queue, - lambda: page_offsets_schedule_body(smem_page_offsets), + lambda: page_offsets_schedule_body(smem_page_offsets, membership_lifetime), ) - schedule_result = ( - page_offsets_schedule(smem_page_offsets) - if work_queue is None - else page_offsets_schedule(smem_page_offsets, work_queue) + schedule_result = page_offsets_schedule( + smem_page_offsets, membership_lifetime, work_queue ) src = [] if work_queue is not None: src.append(work_queue) return task_class( src_resources=src, - dst_resources=[smem_page_offsets], + dst_resources=[smem_page_offsets] + + ([membership_lifetime] if membership_lifetime is not None else []), q_bound_resources=((smem_page_offsets, False),), cfg=cfg, warp_idx=cfg.page_offsets_warp_idx if warp_idx is None else warp_idx, @@ -1580,6 +1688,9 @@ def create_load_task_split_kv( raise ValueError("each split K resource requires its matching V resource") if smem_k0 is None and smem_k1 is None: raise ValueError("at least one split K/V instance is required") + hold_page_window = smem_page_offsets_v is None and _can_hold_native_page_window( + cfg, smem_page_offsets + ) def load_schedule_body( smem_q: MemoryResource | None, @@ -1608,8 +1719,8 @@ def load_schedule_body( "load_v0", ), ( - smem_k1, - smem_v1, + smem_k0 if smem_k1 is None and not cfg.use_block_sparse else smem_k1, + smem_v0 if smem_v1 is None and not cfg.use_block_sparse else smem_v1, sparse_kv_metadata1, sparse_softmax_metadata1, "load_k1", @@ -1618,12 +1729,9 @@ def load_schedule_body( ) # Preserve the original full-resource lowering order while allowing a # per-instance task to omit the other stream's resources. - for smem_k, _, _, _, _, _ in active_instances: - if smem_k is not None: - smem_k.init_load_state() - for _, smem_v, _, _, _, _ in active_instances: - if smem_v is not None: - smem_v.init_load_state() + for resource in (smem_k0, smem_k1, smem_v0, smem_v1): + if resource is not None: + resource.init_load_state() for _, _, sparse_kv_metadata, _, _, _ in active_instances: if sparse_kv_metadata is not None: sparse_kv_metadata.init_load_state() @@ -1649,20 +1757,27 @@ def load_tile( section: FmhaStage, ) -> None: """Acquire, load all head-dim stages, and release page offsets.""" - _page_offsets_consume(offsets, label.replace("load", "read_offsets")) + if not hold_page_window: + _page_offsets_consume(offsets, label.replace("load", "read_offsets")) for head_dim_stage_idx in range(cfg.num_head_dim_stages_kv): resource.acquire() getattr(resource, label)( section=section, head_dim_stage_idx=head_dim_stage_idx ) resource.commit() - _page_offsets_release(offsets) + if not hold_page_window: + _page_offsets_release(offsets) if smem_q is not None: smem_q.acquire() smem_q.tma_load() smem_q.commit() + if hold_page_window: + # Both instruction streams use the same immutable locator span. + # Keep its stage live across all K/V tiles, not one stage per load. + _page_offsets_consume(smem_page_offsets) + head_routes = [] for ( smem_k, @@ -1720,6 +1835,8 @@ def load_tile( smem_page_offsets_v_local, FmhaStage.Tail, ) + if hold_page_window: + _page_offsets_release(smem_page_offsets) @_schedule_with_optional_resources def load_schedule( @@ -1917,6 +2034,7 @@ def create_load_task_one_inst_qkv( **kw: TaskKwarg, ) -> Task: """Load schedule for the one-inst keepsMmaAb QKV path.""" + hold_page_window = _can_hold_native_page_window(cfg, smem_page_offsets) def load_schedule_body( smem_q: MemoryResource, @@ -1933,21 +2051,25 @@ def load_schedule_body( smem_page_offsets.init_read_state() def load_tile(resource: MemoryResource, label: str, section: FmhaStage) -> None: - """Acquire, load all head-dim stages, and release one page window.""" - _page_offsets_consume( - smem_page_offsets, label.replace("load", "read_offsets") - ) + """Load every head-dim stage from the active page window.""" + if not hold_page_window: + _page_offsets_consume( + smem_page_offsets, label.replace("load", "read_offsets") + ) for head_dim_stage_idx in range(cfg.num_head_dim_stages_kv): resource.acquire() getattr(resource, label)( section=section, head_dim_stage_idx=head_dim_stage_idx ) resource.commit() - _page_offsets_release(smem_page_offsets) + if not hold_page_window: + _page_offsets_release(smem_page_offsets) smem_q.acquire() smem_q.tma_load() smem_q.commit() + if hold_page_window: + _page_offsets_consume(smem_page_offsets) load_tile(smem_k, "load_k0", FmhaStage.Head) with domain_loop(0, domain, 1, unroll=1): @@ -1955,6 +2077,8 @@ def load_tile(resource: MemoryResource, label: str, section: FmhaStage) -> None: load_tile(smem_v, "load_v0", FmhaStage.Loop) load_tile(smem_v, "load_v0", FmhaStage.Tail) + if hold_page_window: + _page_offsets_release(smem_page_offsets) @schedule def load_schedule( @@ -2100,10 +2224,10 @@ def qk_mma( section: FmhaStage, ) -> None: """Issue one scheduled QK wave using the selected phase work.""" - _ = section if tmem_stats_done is not None: tmem_stats_done.acquire() tmem_s.acquire() + _issue_score_seed(tmem_s, section, cfg) for head_dim_stage_idx in range(cfg.num_head_dim_stages_kv): smem_kv.wait() kv_desc = smem_kv.kv_desc() @@ -2187,10 +2311,9 @@ def mma_schedule_prelude( ) -> None: """Initialize invariant split-resource descriptor slots.""" smem_q.init_descriptor_state() - smem_k0.init_descriptor_state() - smem_k1.init_descriptor_state() - smem_v0.init_descriptor_state() - smem_v1.init_descriptor_state() + for resource in (smem_k0, smem_k1, smem_v0, smem_v1): + if resource is not None: + resource.init_descriptor_state() smem_p0.init_descriptor_state() smem_p1.init_descriptor_state() @@ -2304,9 +2427,9 @@ def mma_keeps_schedule( mma_schedule( smem_q, smem_k0, - smem_k1, + smem_k0 if smem_k1 is None else smem_k1, smem_v0, - smem_v1, + smem_v0 if smem_v1 is None else smem_v1, tmem_s0, tmem_s1, smem_p0, @@ -2317,9 +2440,9 @@ def mma_keeps_schedule( else mma_schedule( smem_q, smem_k0, - smem_k1, + smem_k0 if smem_k1 is None else smem_k1, smem_v0, - smem_v1, + smem_v0 if smem_v1 is None else smem_v1, tmem_s0, tmem_s1, smem_p0, @@ -2334,9 +2457,9 @@ def mma_keeps_schedule( mma_keeps_schedule( smem_q, smem_k0, - smem_k1, + smem_k0 if smem_k1 is None else smem_k1, smem_v0, - smem_v1, + smem_v0 if smem_v1 is None else smem_v1, tmem_s0, tmem_s1, smem_p0, @@ -2349,9 +2472,9 @@ def mma_keeps_schedule( else mma_keeps_schedule( smem_q, smem_k0, - smem_k1, + smem_k0 if smem_k1 is None else smem_k1, smem_v0, - smem_v1, + smem_v0 if smem_v1 is None else smem_v1, tmem_s0, tmem_s1, smem_p0, @@ -2369,7 +2492,11 @@ def mma_keeps_schedule( tmem_stats_done0, tmem_stats_done1, ] - src = [smem_q, smem_k0, smem_k1, smem_v0, smem_v1, smem_p0, smem_p1] + src = [ + resource + for resource in (smem_q, smem_k0, smem_k1, smem_v0, smem_v1, smem_p0, smem_p1) + if resource is not None + ] if work_queue is not None: src.append(work_queue) return task_class( @@ -2418,9 +2545,9 @@ def mma_schedule_body( def qk_mma(q_desc, qk_mma_label: str, section: FmhaStage) -> None: """Issue one scheduled single-instance QK wave.""" - _ = section tmem_stats_done.acquire() tmem_s.acquire() + _issue_score_seed(tmem_s, section, cfg) for head_dim_stage_idx in range(cfg.num_head_dim_stages_kv): smem_k.wait() kv_desc = smem_k.kv_desc() @@ -2551,6 +2678,151 @@ def mma_schedule( ) +# ====================================================================== +# TransformKvTask — warps 14-15 in mixed Q/KV mode +# Mirrors LoadTask's shared-KV cadence and converts raw K/V into the Q-side +# MMA input type before MmaTask consumes it. +# ====================================================================== +def create_transform_kv_task( + smem_kv: MemoryResource | None, + smem_k0: MemoryResource | None, + smem_k1: MemoryResource | None, + smem_v0: MemoryResource | None, + smem_v1: MemoryResource | None, + smem_transformed_kv: MemoryResource, + work_queue: WorkQueue | None, + cfg: FmhaDecodeConfig, + *, + domain: int | cutlass.Int32, + task_class: type[DecodeGenTask] = DecodeGenTask, + **kw: TaskKwarg, +) -> Task: + """Create the K/V transform task. + + Consume one shared ``smem_kv``, or the four split `smem_k0/k1/v0/v1` + and publish the same ordered K0/V0/K1/V1 stream into ``smem_transformed_kv``. + """ + head_labels = ( + ("transform_k0",) if cfg.num_insts_kv == 1 else ("transform_k0", "transform_k1") + ) + loop_labels = ( + ("transform_k0", "transform_v0") + if cfg.num_insts_kv == 1 + else ("transform_k0", "transform_v0", "transform_k1", "transform_v1") + ) + tail_labels = ( + ("transform_v0",) if cfg.num_insts_kv == 1 else ("transform_v0", "transform_v1") + ) + + split_sources = (smem_k0, smem_k1, smem_v0, smem_v1) + use_split_sources = smem_kv is None + if use_split_sources and any(resource is None for resource in split_sources): + raise ValueError( + "create_transform_kv_task requires either smem_kv, or smem_k0, smem_k1, smem_v0, and smem_v1" + ) + + def transform_kv_schedule_body( + smem_transformed_kv: MemoryResource, + source_for_label: Callable[[str], MemoryResource], + ) -> None: + """Transform the HEAD/LOOP/TAIL stream from the selected raw rings.""" + smem_transformed_kv.init_load_state() + + def transform_one(label: str) -> None: + source = source_for_label(label) + for head_dim_stage_idx in range(cfg.num_head_dim_stages_kv): + smem_transformed_kv.acquire() + source.wait() + getattr(smem_transformed_kv, label)( + head_dim_stage_idx=head_dim_stage_idx + ) + smem_transformed_kv.commit() + source.release() + + for label in head_labels: + transform_one(label) + + # Keep the transformed-K/V steady state as one compact runtime loop. + # Static decode domains can be large, and the compiler default may + # otherwise fully unroll this transform body once per K/V tile. + with domain_loop(0, domain, 1, unroll=1): + for label in loop_labels: + transform_one(label) + + for label in tail_labels: + transform_one(label) + + @schedule + def transform_shared_kv_schedule( + smem_kv: MemoryResource, + smem_transformed_kv: MemoryResource, + work_queue: WorkQueue | None = None, + ) -> None: + smem_kv.init_descriptor_state() + _decode_work_tile_schedule( + cfg, + work_queue, + lambda: transform_kv_schedule_body(smem_transformed_kv, lambda _: smem_kv), + ) + + @schedule + def transform_split_kv_schedule( + smem_k0: MemoryResource, + smem_k1: MemoryResource, + smem_v0: MemoryResource, + smem_v1: MemoryResource, + smem_transformed_kv: MemoryResource, + work_queue: WorkQueue | None = None, + ) -> None: + smem_k0.init_descriptor_state() + smem_k1.init_descriptor_state() + smem_v0.init_descriptor_state() + smem_v1.init_descriptor_state() + + def source_for_label(label: str) -> MemoryResource: + if label == "transform_k0": + return smem_k0 + if label == "transform_k1": + return smem_k1 + if label == "transform_v0": + return smem_v0 + return smem_v1 + + _decode_work_tile_schedule( + cfg, + work_queue, + lambda: transform_kv_schedule_body(smem_transformed_kv, source_for_label), + ) + + if use_split_sources: + split_schedule_args = [*split_sources, smem_transformed_kv] + if work_queue is not None: + split_schedule_args.append(work_queue) + captured_schedule = transform_split_kv_schedule(*split_schedule_args) + src = list(split_sources) + else: + assert smem_kv is not None + captured_schedule = ( + transform_shared_kv_schedule(smem_kv, smem_transformed_kv) + if work_queue is None + else transform_shared_kv_schedule(smem_kv, smem_transformed_kv, work_queue) + ) + src = [smem_kv] + if work_queue is not None: + src.append(work_queue) + return task_class( + src_resources=src, + dst_resources=[smem_transformed_kv], + cfg=cfg, + warp_idx=cfg.transform_kv_warp_idx, + num_warps=cfg.transform_kv_num_warps, + schedule=captured_schedule, + num_registers=cfg.transform_kv_task_num_registers, + name="TransformKvTask", + **kw, + ) + + # ====================================================================== # MmaTask — warp 12, 1 warp # K and V share a single SmemKv ring; each MMA loop iter consumes 4 stages @@ -2563,9 +2835,9 @@ def create_mma_task( smem_q: MemoryResource, smem_kv: MemoryResource, tmem_s0: MemoryResource, - tmem_s1: MemoryResource, + tmem_s1: MemoryResource | None, smem_p0: MemoryResource, - smem_p1: MemoryResource, + smem_p1: MemoryResource | None, tmem_o: MemoryResource, work_queue: WorkQueue | None, cfg: FmhaDecodeConfig, @@ -2580,14 +2852,17 @@ def mma_schedule_body( smem_q: MemoryResource, smem_kv: MemoryResource, tmem_s0: MemoryResource, - tmem_s1: MemoryResource, + tmem_s1: MemoryResource | None, smem_p0: MemoryResource, - smem_p1: MemoryResource, + smem_p1: MemoryResource | None, tmem_o: MemoryResource, q_desc: Any, ) -> None: """Schedule shared-KV QK and PV waves across HEAD/LOOP/TAIL.""" - # HEAD: consume Q once and launch the two initial BMM1 waves. + if cfg.num_insts_kv != 1: + assert smem_p1 is not None + assert tmem_s1 is not None + # HEAD: consume Q once and launch BMM1 for each active instance. _consume_staged_qk_mma( smem_kv, tmem_s0, @@ -2597,15 +2872,16 @@ def mma_schedule_body( FmhaStage.Head, cfg, ) - _consume_staged_qk_mma( - smem_kv, - tmem_s1, - q_desc, - "k_desc_1", - "qk_mma_head", - FmhaStage.Head, - cfg, - ) + if cfg.num_insts_kv != 1: + _consume_staged_qk_mma( + smem_kv, + tmem_s1, + q_desc, + "k_desc_1", + "qk_mma_head", + FmhaStage.Head, + cfg, + ) # LOOP: consume aliased TMEM P before the next same-instance QK # overwrites its S columns. Full-SMEM P remains score-dependent during @@ -2643,37 +2919,38 @@ def mma_schedule_body( FmhaStage.Loop, cfg, ) - if cfg.uses_two_inst_tmem_p: - _consume_staged_pv_mma( - smem_kv, - smem_p1, - tmem_o, - "v_desc_1", - "vp_mma_loop", - KV_INST1, - FmhaStage.Loop, - cfg, - ) - _consume_staged_qk_mma( - smem_kv, - tmem_s1, - q_desc, - "k_desc_1", - "qk_mma_loop", - FmhaStage.Loop, - cfg, - ) - if not cfg.uses_two_inst_tmem_p: - _consume_staged_pv_mma( + if cfg.num_insts_kv != 1: + if cfg.uses_two_inst_tmem_p: + _consume_staged_pv_mma( + smem_kv, + smem_p1, + tmem_o, + "v_desc_1", + "vp_mma_loop", + KV_INST1, + FmhaStage.Loop, + cfg, + ) + _consume_staged_qk_mma( smem_kv, - smem_p1, - tmem_o, - "v_desc_1", - "vp_mma_loop", - KV_INST1, + tmem_s1, + q_desc, + "k_desc_1", + "qk_mma_loop", FmhaStage.Loop, cfg, ) + if not cfg.uses_two_inst_tmem_p: + _consume_staged_pv_mma( + smem_kv, + smem_p1, + tmem_o, + "v_desc_1", + "vp_mma_loop", + KV_INST1, + FmhaStage.Loop, + cfg, + ) # Q is live for every BMM1 call, and the last BMM1 has been issued once # the loop ends. Releasing here commits after those MMAs complete, so @@ -2681,7 +2958,7 @@ def mma_schedule_body( # instead of waiting for them. smem_q.release() - # TAIL: no future K tiles remain, so only the final two BMM2 waves run. + # TAIL: no future K tiles remain, so only the final BMM2 waves run. _consume_staged_pv_mma( smem_kv, smem_p0, @@ -2692,45 +2969,53 @@ def mma_schedule_body( FmhaStage.Tail, cfg, ) - _consume_staged_pv_mma( - smem_kv, - smem_p1, - tmem_o, - "v_desc_1", - "vp_mma_tail", - KV_INST1, - FmhaStage.Tail, - cfg, - ) + if cfg.num_insts_kv != 1: + _consume_staged_pv_mma( + smem_kv, + smem_p1, + tmem_o, + "v_desc_1", + "vp_mma_tail", + KV_INST1, + FmhaStage.Tail, + cfg, + ) def mma_schedule_prelude( smem_q: MemoryResource, smem_kv: MemoryResource, smem_p0: MemoryResource, - smem_p1: MemoryResource, + smem_p1: MemoryResource | None, ) -> None: """Initialize invariant shared-ring descriptor slots.""" smem_q.init_descriptor_state() smem_kv.init_descriptor_state() smem_p0.init_descriptor_state() - smem_p1.init_descriptor_state() + if cfg.num_insts_kv != 1: + assert smem_p1 is not None + smem_p1.init_descriptor_state() - @schedule + @_schedule_with_optional_resources def mma_schedule( smem_q: MemoryResource, smem_kv: MemoryResource, tmem_s0: MemoryResource, - tmem_s1: MemoryResource, + tmem_s1: MemoryResource | None, smem_p0: MemoryResource, - smem_p1: MemoryResource, + smem_p1: MemoryResource | None, tmem_o: MemoryResource, - work_queue: WorkQueue | None = None, + work_queue: WorkQueue | None, ) -> None: """Wrap shared-ring MMA work in packed persistent skip handling.""" _decode_work_tile_schedule_with_invariant_bridge( cfg, work_queue, - lambda: mma_schedule_prelude(smem_q, smem_kv, smem_p0, smem_p1), + lambda: mma_schedule_prelude( + smem_q, + smem_kv, + smem_p0, + smem_p1, + ), lambda: smem_q.q_desc(), lambda: smem_q.wait(), lambda q_desc: mma_schedule_body( @@ -2745,34 +3030,26 @@ def mma_schedule( ), ) - captured_schedule = ( - mma_schedule( - smem_q, - smem_kv, - tmem_s0, - tmem_s1, - smem_p0, - smem_p1, - tmem_o, - ) - if work_queue is None - else mma_schedule( - smem_q, - smem_kv, - tmem_s0, - tmem_s1, - smem_p0, - smem_p1, - tmem_o, - work_queue, - ) + captured_schedule = mma_schedule( + smem_q, + smem_kv, + tmem_s0, + tmem_s1, + smem_p0, + smem_p1, + tmem_o, + work_queue, ) - src = [smem_q, smem_kv, smem_p0, smem_p1] + src = [smem_q, smem_kv, smem_p0] + dst = [tmem_s0, tmem_o] + if cfg.num_insts_kv != 1: + src.append(smem_p1) + dst.insert(1, tmem_s1) if work_queue is not None: src.append(work_queue) return task_class( src_resources=src, - dst_resources=[tmem_s0, tmem_s1, tmem_o], + dst_resources=dst, cfg=cfg, warp_idx=cfg.mma_warp_idx, num_warps=cfg.mma_num_warps, @@ -2783,6 +3060,247 @@ def mma_schedule( ) +def _softmax_schedule_body( + cfg: FmhaDecodeConfig, + inst_id: int, + domain: int | cutlass.Int32, + tmem_s: MemoryResource, + tmem_softmax_local: MemoryResource, + smem_p: MemoryResource, + tmem_softmax_global: MemoryResource, + tmem_softmax_order: MemoryResource | None, + sparse_softmax_metadata: MemoryResource | None, + membership_lifetime: MemoryResource | None, + sage_k_scales: MemoryResource | None, + sage_summary_k_scales: MemoryResource | None, + sage_scales_in_smem: tuple[bool, ...] = (), +) -> None: + """Build one softmax instance's loop, P publication, and final stats handoff. + + Both instances run this body and differ only in their side of the P + baton. ``sage_scales_in_smem`` flags the scale resources in SMEM form, + which carry a pipeline; schedule proxies cannot answer it themselves. + """ + if membership_lifetime is not None: + membership_lifetime.wait() + ( + old_max_arr, + sum_arr, + new_max_arr, + s_arr, + ) = tmem_s.init_softmax_state() + smem_p.init_compute_state() + if sparse_softmax_metadata is not None: + sparse_softmax_metadata.init_read_state() + if cutlass.const_expr(cfg.use_sage_attention): + # The lane's Q row is fixed for the work tile. + sage_q_scale = tmem_s.load_sage_q_scale() + # The work framework rejects ``None`` arguments, so each Sage/sparse + # combination is its own work callable. A single-geometry plan passes + # its one routed ``sfK`` array under both names. + sage_scales = [r for r in (sage_k_scales, sage_summary_k_scales) if r is not None] + sage_smem_scales = [ + r + for r, in_smem in zip(sage_scales, sage_scales_in_smem, strict=True) + if in_smem + ] + use_sparse = sparse_softmax_metadata is not None + use_sage = cfg.use_sage_attention + compute_softmax_loop = { + (False, False): tmem_s.compute_softmax_loop, + (False, True): tmem_s.compute_sage_softmax_loop, + (True, False): tmem_s.compute_block_sparse_softmax_loop, + (True, True): tmem_s.compute_sage_block_sparse_softmax_loop, + }[(use_sparse, use_sage)] + compute_p_fragments = { + (False, False): smem_p.compute_p_fragments, + (False, True): smem_p.compute_sage_p_fragments, + (True, False): smem_p.compute_proxy_route_p_fragments, + (True, True): smem_p.compute_sage_proxy_route_p_fragments, + }[(cfg.use_block_sparse_proxy_routes, use_sage)] + + with domain_loop(0, domain, 1, unroll=1) as d: + if use_sparse: + # Consume the independent metadata stream first so its SMEM + # loads and release can overlap the subsequent score wait. + sparse_softmax_metadata.wait() + # Copy the complete payload to registers before release, so + # masking cannot race the producer's next SMEM-stage reuse. + ( + sparse_origin0, + sparse_origin1, + sparse_route_flags, + sparse_token_word0, + sparse_token_word1, + sparse_token_word2, + sparse_token_word3, + ) = sparse_softmax_metadata.load_route() + if cutlass.const_expr(cfg.use_sage_attention): + # Take the route's ``sfK`` words before the metadata stage + # is released. Each iteration fills and publishes this tile's + # words only; the second stage lets the fill overlap the + # previous tile's P pass on the instance's other warps. + for scales in sage_smem_scales: + scales.acquire() + scales.publish_route_tile(route_flags=sparse_route_flags) + scales.commit() + scales.wait() + sage_scale_arr = sage_k_scales.take_route_tile( + route_flags=sparse_route_flags + ) + sage_summary_scale_arr = sage_scale_arr + if sage_summary_k_scales is not None: + sage_summary_scale_arr = sage_summary_k_scales.take_route_tile( + route_flags=sparse_route_flags + ) + sparse_softmax_metadata.release() + if cutlass.const_expr(cfg.use_sage_attention and not use_sparse): + # Issue the dense tile's ``sfK`` loads before the score wait so + # they hide behind the QK MMA. + for scales in sage_smem_scales: + scales.acquire() + scales.publish_tile() + scales.commit() + scales.wait() + sage_scale_arr = sage_k_scales.take_tile() + sage_summary_scale_arr = sage_scale_arr + # ConsWait/ConsWork: load S from TMEM and compute the tile max. + tmem_s.wait() + softmax_loop_kwargs = dict( + old_max_arr=old_max_arr, + sum_arr=sum_arr, + new_max_arr=new_max_arr, + s_arr=s_arr, + ) + if cutlass.const_expr(cfg.use_sage_attention): + softmax_loop_kwargs["sage_q_scale"] = sage_q_scale + softmax_loop_kwargs["sage_scale_arr"] = sage_scale_arr + softmax_loop_kwargs["sage_summary_scale_arr"] = sage_summary_scale_arr + if use_sparse: + softmax_loop_kwargs.update( + sparse_origin0=sparse_origin0, + sparse_origin1=sparse_origin1, + sparse_route_flags=sparse_route_flags, + sparse_token_word0=sparse_token_word0, + sparse_token_word1=sparse_token_word1, + sparse_token_word2=sparse_token_word2, + sparse_token_word3=sparse_token_word3, + ) + old_max_arr, sum_arr, new_max_arr, s_arr = compute_softmax_loop( + **softmax_loop_kwargs + ) + if cutlass.const_expr(not cfg.use_keeps_mma_ab or not cfg.uses_tmem_p): + # ConsRelease: free S once the scores are in registers unless + # a Keeps TMEM-P operand still aliases the consumed columns. + tmem_s.release() + # Publish old/new max before the P path so correction can observe + # the same stats order as the decode pipeline. + # ProdWork: store old/new max for correction's in-loop O update. + tmem_softmax_local.acquire() + tmem_softmax_local.store_loop_old_new_stats( + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + ) + tmem_softmax_local.commit() + if cutlass.const_expr(cfg.streams_tmem_p_fragments): + # One rolled loop streams every K32 probability fragment; the + # fragment body exists once in the instruction stream. + p_fragments_kwargs = dict(new_max_arr=new_max_arr) + if cutlass.const_expr(cfg.use_sage_attention): + p_fragments_kwargs["sage_q_scale"] = sage_q_scale + p_fragments_kwargs["sage_scale_arr"] = sage_scale_arr + p_fragments_kwargs["sage_summary_scale_arr"] = sage_summary_scale_arr + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): + p_fragments_kwargs["route_flags"] = sparse_route_flags + compute_p_fragments(**p_fragments_kwargs) + else: + # Wait for a free P stage before entering the ordered window so + # BMM2 backpressure on this group's P pipeline cannot extend the + # baton hold and stall the partner softmax group. + smem_p.acquire() + if tmem_softmax_order is not None: + if inst_id == 0: + tmem_softmax_order.wait_softmax0() + else: + tmem_softmax_order.wait_softmax1() + # ProdWork: compute P=exp(S-new_max), store it in the profile's + # SMEM or staged-TMEM operand, and record local sums for the + # running softmax sum update. + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): + smem_p.compute_proxy_route_p( + new_max_arr=new_max_arr, + s_arr=s_arr, + keeps_route_flags=sparse_route_flags, + swaps_route_flags=sparse_token_word2, + ) + else: + smem_p.compute_p( + new_max_arr=new_max_arr, + s_arr=s_arr, + ) # publishes the local denominator through tmem_s + smem_p.commit() + if tmem_softmax_order is not None: + if inst_id == 0: + tmem_softmax_order.release_softmax1() + else: + tmem_softmax_order.release_softmax0() + if cutlass.const_expr(cfg.use_keeps_mma_ab and cfg.uses_tmem_p): + # TMEM-P has consumed the aliased S columns, so the next QK + # wave can now overwrite S. + tmem_s.release() + if cutlass.const_expr(cfg.use_sage_attention): + # ConsRelease: both passes have read this tile's ``sfK`` slot. + for scales in sage_smem_scales: + scales.release() + # ProdWork: FP8 path applies the cross-resource sum correction + # before TmemS.reduce_sums publishes the new running sums. + tmem_softmax_global.global_correction( + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + ) # publishes the corrected denominator through tmem_s + # ConsTailWork: update the running sum after P is available. + sum_arr = tmem_s.reduce_sums( + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + ) + if cutlass.const_expr(not cfg.use_keeps_mma_ab): + with d.last_iter(): + # LastIter ProdWork: store the final sums for correction's + # normalization/output tail. + tmem_softmax_local.acquire() + tmem_softmax_local.store_loop_sum_new_stats( + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + ) + tmem_softmax_local.commit() + if cutlass.const_expr(cfg.use_keeps_mma_ab): + # The final sum payload has no matching QK/StatsDone token. + tmem_softmax_local.acquire() + tmem_softmax_local.store_tail_sum_new_stats( + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + sum_arr=sum_arr, + ) + tmem_softmax_local.commit() + if cutlass.const_expr( + inst_id == 0 + and cfg.use_persistent_scheduler + and cfg.uses_staged_one_inst_tmem_p + ): + # Tail stats add one stage beyond the S cadence. Advance the + # second stats slot so both pipelines start the next work tile + # on the same physical TMEM stage. A one-instance profile has + # only the first instance. + tmem_softmax_local.acquire() + tmem_softmax_local.commit() + if membership_lifetime is not None: + membership_lifetime.release() + + # ====================================================================== # Softmax0Task — warps 0-3, 4 warps, profile-selected register budget # LOOP: consume S, produce stats + P + running sum @@ -2799,171 +3317,16 @@ def create_softmax0_task( cfg: FmhaDecodeConfig, *, domain: int | cutlass.Int32, + membership_lifetime: MemoryResource | None = None, + sage_k_scales: MemoryResource | None = None, + sage_summary_k_scales: MemoryResource | None = None, task_class: type[DecodeGenTask] = DecodeGenTask, **kw: TaskKwarg, ) -> Task: """Create the first softmax task, including optional ordered publication.""" - def softmax0_schedule_body( - tmem_s0: MemoryResource, - tmem_softmax_local0: MemoryResource, - smem_p0: MemoryResource, - tmem_softmax_global0: MemoryResource, - tmem_softmax_order: MemoryResource | None, - sparse_softmax_metadata: MemoryResource | None, - ) -> None: - """Build the softmax0 loop, P publication, and final stats handoff.""" - ( - old_max_arr, - sum_arr, - new_max_arr, - s_arr, - ) = tmem_s0.init_softmax_state() - smem_p0.init_compute_state() - if sparse_softmax_metadata is not None: - sparse_softmax_metadata.init_read_state() - - with domain_loop(0, domain, 1, unroll=1) as d: - if sparse_softmax_metadata is not None: - # Consume the independent metadata stream first so its SMEM - # loads and release can overlap the subsequent score wait. - sparse_softmax_metadata.wait() - # Copy the complete payload to registers before release, so - # masking cannot race the producer's next SMEM-stage reuse. - ( - sparse_origin0, - sparse_origin1, - sparse_route_flags, - sparse_token_word0, - sparse_token_word1, - sparse_token_word2, - sparse_token_word3, - ) = sparse_softmax_metadata.load_route() - sparse_softmax_metadata.release() - # ConsWait/ConsWork: load S from TMEM and compute the tile max. - tmem_s0.wait() - if sparse_softmax_metadata is not None: - old_max_arr, sum_arr, new_max_arr, s_arr = ( - tmem_s0.compute_block_sparse_softmax_loop( - old_max_arr=old_max_arr, - sum_arr=sum_arr, - new_max_arr=new_max_arr, - s_arr=s_arr, - sparse_origin0=sparse_origin0, - sparse_origin1=sparse_origin1, - sparse_route_flags=sparse_route_flags, - sparse_token_word0=sparse_token_word0, - sparse_token_word1=sparse_token_word1, - sparse_token_word2=sparse_token_word2, - sparse_token_word3=sparse_token_word3, - ) - ) - else: - old_max_arr, sum_arr, new_max_arr, s_arr = tmem_s0.compute_softmax_loop( - old_max_arr=old_max_arr, - sum_arr=sum_arr, - new_max_arr=new_max_arr, - s_arr=s_arr, - ) - if cutlass.const_expr(not cfg.use_keeps_mma_ab or not cfg.uses_tmem_p): - # ConsRelease: free S once the scores are in registers unless - # a Keeps TMEM-P operand still aliases the consumed columns. - tmem_s0.release() - # Publish old/new max before the P path so correction can observe - # the same stats order as the decode pipeline. - # ProdWork: store old/new max for correction's in-loop O update. - tmem_softmax_local0.acquire() - tmem_softmax_local0.store_loop_old_new_stats( - old_max_arr=old_max_arr, - new_max_arr=new_max_arr, - sum_arr=sum_arr, - ) - tmem_softmax_local0.commit() - if cutlass.const_expr(cfg.streams_tmem_p_fragments): - # One rolled loop streams every K32 probability fragment; the - # fragment body exists once in the instruction stream. - if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): - smem_p0.compute_proxy_route_p_fragments( - new_max_arr=new_max_arr, - route_flags=sparse_route_flags, - route_origin0=sparse_origin0, - route_origin1=sparse_origin1, - ) - else: - smem_p0.compute_p_fragments(new_max_arr=new_max_arr) - else: - # Wait for a free P stage before entering the ordered window so - # BMM2 backpressure on this group's P pipeline cannot extend the - # baton hold and stall the partner softmax group. - smem_p0.acquire() - if tmem_softmax_order is not None: - tmem_softmax_order.wait_softmax0() - # ProdWork: compute P=exp(S-new_max), store it in the profile's - # SMEM or staged-TMEM operand, and record local sums for the - # running softmax sum update. - if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): - smem_p0.compute_proxy_route_p( - new_max_arr=new_max_arr, - s_arr=s_arr, - route_origin0=sparse_origin0, - route_origin1=sparse_origin1, - keeps_route_flags_or_swaps_origin2=sparse_route_flags, - swaps_route_origin3_bits=sparse_token_word0, - swaps_route_flags=sparse_token_word2, - ) - else: - smem_p0.compute_p( - new_max_arr=new_max_arr, - s_arr=s_arr, - ) # publishes the local denominator through tmem_s0 - smem_p0.commit() - if tmem_softmax_order is not None: - tmem_softmax_order.release_softmax1() - if cutlass.const_expr(cfg.use_keeps_mma_ab and cfg.uses_tmem_p): - # TMEM-P has consumed the aliased S columns, so the next QK - # wave can now overwrite S. - tmem_s0.release() - # ProdWork: FP8 path applies the cross-resource sum correction - # before TmemS.reduce_sums publishes the new running sums. - tmem_softmax_global0.global_correction( - old_max_arr=old_max_arr, - new_max_arr=new_max_arr, - sum_arr=sum_arr, - ) # publishes the corrected denominator through tmem_s0 - # ConsTailWork: update the running sum after P is available. - sum_arr = tmem_s0.reduce_sums( - old_max_arr=old_max_arr, - new_max_arr=new_max_arr, - sum_arr=sum_arr, - ) - if cutlass.const_expr(not cfg.use_keeps_mma_ab): - with d.last_iter(): - # LastIter ProdWork: store the final sums for correction's - # normalization/output tail. - tmem_softmax_local0.acquire() - tmem_softmax_local0.store_loop_sum_new_stats( - old_max_arr=old_max_arr, - new_max_arr=new_max_arr, - sum_arr=sum_arr, - ) - tmem_softmax_local0.commit() - if cutlass.const_expr(cfg.use_keeps_mma_ab): - # The final sum payload has no matching QK/StatsDone token. - tmem_softmax_local0.acquire() - tmem_softmax_local0.store_tail_sum_new_stats( - old_max_arr=old_max_arr, - new_max_arr=new_max_arr, - sum_arr=sum_arr, - ) - tmem_softmax_local0.commit() - if cutlass.const_expr( - cfg.use_persistent_scheduler and cfg.uses_staged_one_inst_tmem_p - ): - # Tail stats add one stage beyond the S cadence. Advance the - # second stats slot so both pipelines start the next work tile - # on the same physical TMEM stage. - tmem_softmax_local0.acquire() - tmem_softmax_local0.commit() + sage_scales = [r for r in (sage_k_scales, sage_summary_k_scales) if r is not None] + sage_scales_in_smem = tuple(r.in_smem for r in sage_scales) @_schedule_with_optional_resources def softmax0_schedule( @@ -2973,19 +3336,29 @@ def softmax0_schedule( tmem_softmax_global0: MemoryResource, tmem_softmax_order: MemoryResource | None, sparse_softmax_metadata: MemoryResource | None, + membership_lifetime: MemoryResource | None, + sage_k_scales: MemoryResource | None, + sage_summary_k_scales: MemoryResource | None, work_queue: WorkQueue | None, ) -> None: """Schedule softmax0 with the supplied order and sparse resources.""" _decode_work_tile_schedule( cfg, work_queue, - lambda: softmax0_schedule_body( + lambda: _softmax_schedule_body( + cfg, + 0, + domain, tmem_s0, tmem_softmax_local0, smem_p0, tmem_softmax_global0, tmem_softmax_order, sparse_softmax_metadata, + membership_lifetime, + sage_k_scales, + sage_summary_k_scales, + sage_scales_in_smem, ), ) @@ -2996,14 +3369,22 @@ def softmax0_schedule( tmem_softmax_global0, tmem_softmax_order, sparse_softmax_metadata, + membership_lifetime, + sage_k_scales, + sage_summary_k_scales, work_queue, ) src = [tmem_s0] if sparse_softmax_metadata is not None: src.append(sparse_softmax_metadata) + if membership_lifetime is not None: + src.append(membership_lifetime) + # The scale resources are produced and consumed by this task. + src.extend(sage_scales) if work_queue is not None: src.append(work_queue) dst = [tmem_softmax_local0, smem_p0, tmem_softmax_global0] + dst.extend(r for r in sage_scales if r.in_smem) if tmem_softmax_order is not None: dst.append(tmem_softmax_order) return task_class( @@ -3034,154 +3415,16 @@ def create_softmax1_task( cfg: FmhaDecodeConfig, *, domain: int | cutlass.Int32, + membership_lifetime: MemoryResource | None = None, + sage_k_scales: MemoryResource | None = None, + sage_summary_k_scales: MemoryResource | None = None, task_class: type[DecodeGenTask] = DecodeGenTask, **kw: TaskKwarg, ) -> Task: """Create the second softmax task, including optional ordered publication.""" - def softmax1_schedule_body( - tmem_s1: MemoryResource, - tmem_softmax_local1: MemoryResource, - smem_p1: MemoryResource, - tmem_softmax_global1: MemoryResource, - tmem_softmax_order: MemoryResource | None, - sparse_softmax_metadata: MemoryResource | None, - ) -> None: - """Build the softmax1 loop, P publication, and final stats handoff.""" - ( - old_max_arr, - sum_arr, - new_max_arr, - s_arr, - ) = tmem_s1.init_softmax_state() - smem_p1.init_compute_state() - if sparse_softmax_metadata is not None: - sparse_softmax_metadata.init_read_state() - - with domain_loop(0, domain, 1, unroll=1) as d: - if sparse_softmax_metadata is not None: - # Consume the independent metadata stream first so its SMEM - # loads and release can overlap the subsequent score wait. - sparse_softmax_metadata.wait() - # Copy to registers before release so the producer can reuse - # the SMEM stage while this warp group applies the masks. - ( - sparse_origin0, - sparse_origin1, - sparse_route_flags, - sparse_token_word0, - sparse_token_word1, - sparse_token_word2, - sparse_token_word3, - ) = sparse_softmax_metadata.load_route() - sparse_softmax_metadata.release() - # ConsWait/ConsWork: load the second S instance and compute max. - tmem_s1.wait() - if sparse_softmax_metadata is not None: - old_max_arr, sum_arr, new_max_arr, s_arr = ( - tmem_s1.compute_block_sparse_softmax_loop( - old_max_arr=old_max_arr, - sum_arr=sum_arr, - new_max_arr=new_max_arr, - s_arr=s_arr, - sparse_origin0=sparse_origin0, - sparse_origin1=sparse_origin1, - sparse_route_flags=sparse_route_flags, - sparse_token_word0=sparse_token_word0, - sparse_token_word1=sparse_token_word1, - sparse_token_word2=sparse_token_word2, - sparse_token_word3=sparse_token_word3, - ) - ) - else: - old_max_arr, sum_arr, new_max_arr, s_arr = tmem_s1.compute_softmax_loop( - old_max_arr=old_max_arr, - sum_arr=sum_arr, - new_max_arr=new_max_arr, - s_arr=s_arr, - ) - if cutlass.const_expr(not cfg.use_keeps_mma_ab or not cfg.uses_tmem_p): - # ConsRelease: SMEM-P Keeps and Swaps no longer need S after - # the score fragment has been loaded into registers. - tmem_s1.release() - # ProdWork: store old/new max for the correction task. - tmem_softmax_local1.acquire() - tmem_softmax_local1.store_loop_old_new_stats( - old_max_arr=old_max_arr, - new_max_arr=new_max_arr, - sum_arr=sum_arr, - ) - tmem_softmax_local1.commit() - if cutlass.const_expr(cfg.streams_tmem_p_fragments): - if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): - smem_p1.compute_proxy_route_p_fragments( - new_max_arr=new_max_arr, - route_flags=sparse_route_flags, - route_origin0=sparse_origin0, - route_origin1=sparse_origin1, - ) - else: - smem_p1.compute_p_fragments(new_max_arr=new_max_arr) - else: - # Wait for a free P stage before entering the ordered window so - # BMM2 backpressure on this group's P pipeline cannot extend the - # baton hold and stall the partner softmax group. - smem_p1.acquire() - if tmem_softmax_order is not None: - tmem_softmax_order.wait_softmax1() - # ProdWork: compute and publish P1 for BMM2. - if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): - smem_p1.compute_proxy_route_p( - new_max_arr=new_max_arr, - s_arr=s_arr, - route_origin0=sparse_origin0, - route_origin1=sparse_origin1, - keeps_route_flags_or_swaps_origin2=sparse_route_flags, - swaps_route_origin3_bits=sparse_token_word0, - swaps_route_flags=sparse_token_word2, - ) - else: - smem_p1.compute_p(new_max_arr=new_max_arr, s_arr=s_arr) - smem_p1.commit() - if tmem_softmax_order is not None: - tmem_softmax_order.release_softmax0() - if cutlass.const_expr(cfg.use_keeps_mma_ab and cfg.uses_tmem_p): - # The TMEM-P store has consumed the aliased S columns, so the - # next QK wave can now overwrite them. - tmem_s1.release() - # ProdWork: update FP8 global sums if the configuration needs the - # split softmax-sum correction path. - tmem_softmax_global1.global_correction( - old_max_arr=old_max_arr, - new_max_arr=new_max_arr, - sum_arr=sum_arr, - ) - # ConsTailWork: fold local sums into the running sums. - sum_arr = tmem_s1.reduce_sums( - old_max_arr=old_max_arr, - new_max_arr=new_max_arr, - sum_arr=sum_arr, - ) - if cutlass.const_expr(not cfg.use_keeps_mma_ab): - with d.last_iter(): - # LastIter ProdWork: publish final sums for output - # normalization. - tmem_softmax_local1.acquire() - tmem_softmax_local1.store_loop_sum_new_stats( - old_max_arr=old_max_arr, - new_max_arr=new_max_arr, - sum_arr=sum_arr, - ) - tmem_softmax_local1.commit() - if cutlass.const_expr(cfg.use_keeps_mma_ab): - # The final sum payload has no matching QK/StatsDone token. - tmem_softmax_local1.acquire() - tmem_softmax_local1.store_tail_sum_new_stats( - old_max_arr=old_max_arr, - new_max_arr=new_max_arr, - sum_arr=sum_arr, - ) - tmem_softmax_local1.commit() + sage_scales = [r for r in (sage_k_scales, sage_summary_k_scales) if r is not None] + sage_scales_in_smem = tuple(r.in_smem for r in sage_scales) @_schedule_with_optional_resources def softmax1_schedule( @@ -3191,6 +3434,9 @@ def softmax1_schedule( tmem_softmax_global1: MemoryResource, tmem_softmax_order: MemoryResource | None, sparse_softmax_metadata: MemoryResource | None, + membership_lifetime: MemoryResource | None, + sage_k_scales: MemoryResource | None, + sage_summary_k_scales: MemoryResource | None, work_queue: WorkQueue | None, ) -> None: """Schedule softmax1 with the supplied order and sparse resources.""" @@ -3203,13 +3449,20 @@ def softmax1_schedule( _decode_work_tile_schedule( cfg, work_queue, - lambda: softmax1_schedule_body( + lambda: _softmax_schedule_body( + cfg, + 1, + domain, tmem_s1, tmem_softmax_local1, smem_p1, tmem_softmax_global1, tmem_softmax_order, sparse_softmax_metadata, + membership_lifetime, + sage_k_scales, + sage_summary_k_scales, + sage_scales_in_smem, ), ) @@ -3220,6 +3473,9 @@ def softmax1_schedule( tmem_softmax_global1, tmem_softmax_order, sparse_softmax_metadata, + membership_lifetime, + sage_k_scales, + sage_summary_k_scales, work_queue, ) src = [tmem_s1] @@ -3227,11 +3483,20 @@ def softmax1_schedule( src.append(sparse_softmax_metadata) if tmem_softmax_order is not None: src.append(tmem_softmax_order) + if membership_lifetime is not None: + src.append(membership_lifetime) + # The scale resources are produced and consumed by this task. + src.extend(sage_scales) if work_queue is not None: src.append(work_queue) return task_class( src_resources=src, - dst_resources=[tmem_softmax_local1, smem_p1, tmem_softmax_global1], + dst_resources=[ + tmem_softmax_local1, + smem_p1, + tmem_softmax_global1, + *(r for r in sage_scales if r.in_smem), + ], q_bound_resources=((tmem_s1, False),), cfg=cfg, warp_idx=cfg.softmax1_warp_idx, @@ -3261,6 +3526,7 @@ def create_correction_task( domain: int | cutlass.Int32, tmem_stats_done0: MemoryResource | None = None, tmem_stats_done1: MemoryResource | None = None, + sage_v_scales: MemoryResource | None = None, task_class: type[DecodeGenTask] = DecodeGenTask, **kw: TaskKwarg, ) -> Task: @@ -3274,6 +3540,7 @@ def correction_schedule_body( tmem_corr1: MemoryResource, tmem_stats_done0: MemoryResource | None, tmem_stats_done1: MemoryResource | None, + sage_v_scales: MemoryResource | None, ) -> None: """Schedule two-instance O correction and final output normalization.""" @@ -3330,25 +3597,11 @@ def correct_o( tmem_softmax_local: MemoryResource, tmem_stats_done: MemoryResource | None, tmem_corr: MemoryResource, - local_state: tuple[ - cutlass.Array, - cutlass.Array, - cutlass.Array, - cutlass.Array, - cutlass.Array, - cutlass.Array, - ], + local_state: LocalStatsState, tail_0: cutlass.Int32, tail_1: cutlass.Int32, ) -> tuple[ - tuple[ - cutlass.Array, - cutlass.Array, - cutlass.Array, - cutlass.Array, - cutlass.Array, - cutlass.Array, - ], + LocalStatsState, cutlass.Int32, cutlass.Int32, ]: @@ -3385,6 +3638,12 @@ def correct_o( _, tail_0, tail_1 = tmem_o.init_stage_state() tmem_corr0.init_epilogue_state() tmem_corr1.init_epilogue_state() + if sage_v_scales is not None: + # ProdWork: stage the work tile's V channel scales while no + # output is pending. + sage_v_scales.acquire() + sage_v_scales.stage_tile() + sage_v_scales.commit() # HEAD: drain the first softmax-local handoff. No O is corrected here; # this aligns the stats pipeline before loop work starts. @@ -3467,6 +3726,9 @@ def correct_o( tail_o_stage_idx_1=tail_1, inst_idx=KV_INST1, ) + if sage_v_scales is not None: + # ConsWait: every lane reads the scales the other lanes staged. + sage_v_scales.wait() tmem_corr1.correction_tail_epilogue( o_stage_idx=o_stage_idx, tail_o_stage_idx_0=tail_0, @@ -3478,12 +3740,16 @@ def correct_o( inst1_new_max_arr=inst1_new_max_arr, inst1_sum_arr=inst1_sum_arr, ) + if sage_v_scales is not None: + # ConsRelease: the next work tile may restage the block. + sage_v_scales.release() # Inst1 final reduction consumes both O0 and O1, so defer O0 release # until after inst1 has finished reading it. tmem_o.release() tmem_o.release() - def run_correction_schedule( + @_schedule_with_optional_resources + def correction_schedule( tmem_softmax_local0: MemoryResource, tmem_softmax_local1: MemoryResource, tmem_o: MemoryResource, @@ -3491,9 +3757,10 @@ def run_correction_schedule( tmem_corr1: MemoryResource, tmem_stats_done0: MemoryResource | None, tmem_stats_done1: MemoryResource | None, + sage_v_scales: MemoryResource | None, work_queue: WorkQueue | None, ) -> None: - """Wrap correction with optional stats lifetime gates.""" + """Capture correction with its optional stats gates and scale staging.""" _decode_work_tile_schedule( cfg, work_queue, @@ -3505,106 +3772,38 @@ def run_correction_schedule( tmem_corr1, tmem_stats_done0, tmem_stats_done1, + sage_v_scales, ), ) - @schedule - def correction_schedule( - tmem_softmax_local0: MemoryResource, - tmem_softmax_local1: MemoryResource, - tmem_o: MemoryResource, - tmem_corr0: MemoryResource, - tmem_corr1: MemoryResource, - work_queue: WorkQueue | None = None, - ) -> None: - """Capture the Swaps correction schedule.""" - run_correction_schedule( - tmem_softmax_local0, - tmem_softmax_local1, - tmem_o, - tmem_corr0, - tmem_corr1, - None, - None, - work_queue, - ) - - @schedule - def correction_keeps_schedule( - tmem_softmax_local0: MemoryResource, - tmem_softmax_local1: MemoryResource, - tmem_o: MemoryResource, - tmem_corr0: MemoryResource, - tmem_corr1: MemoryResource, - tmem_stats_done0: MemoryResource, - tmem_stats_done1: MemoryResource, - work_queue: WorkQueue | None = None, - ) -> None: - """Capture Keeps correction with explicit stats lifetime gates.""" - run_correction_schedule( - tmem_softmax_local0, - tmem_softmax_local1, - tmem_o, - tmem_corr0, - tmem_corr1, - tmem_stats_done0, - tmem_stats_done1, - work_queue, - ) - + # The Keeps stats lifetime gates come as a pair; the Swaps path has none. if tmem_stats_done0 is None or tmem_stats_done1 is None: - if work_queue is None: - captured_schedule = correction_schedule( - tmem_softmax_local0, - tmem_softmax_local1, - tmem_o, - tmem_corr0, - tmem_corr1, - ) - else: - captured_schedule = correction_schedule( - tmem_softmax_local0, - tmem_softmax_local1, - tmem_o, - tmem_corr0, - tmem_corr1, - work_queue, - ) - src = [tmem_softmax_local0, tmem_softmax_local1, tmem_o] - else: - if work_queue is None: - captured_schedule = correction_keeps_schedule( - tmem_softmax_local0, - tmem_softmax_local1, - tmem_o, - tmem_corr0, - tmem_corr1, - tmem_stats_done0, - tmem_stats_done1, - ) - else: - captured_schedule = correction_keeps_schedule( - tmem_softmax_local0, - tmem_softmax_local1, - tmem_o, - tmem_corr0, - tmem_corr1, - tmem_stats_done0, - tmem_stats_done1, - work_queue, - ) - src = [ - tmem_softmax_local0, - tmem_softmax_local1, - tmem_o, - tmem_stats_done0, - tmem_stats_done1, - ] + tmem_stats_done0 = None + tmem_stats_done1 = None + captured_schedule = correction_schedule( + tmem_softmax_local0, + tmem_softmax_local1, + tmem_o, + tmem_corr0, + tmem_corr1, + tmem_stats_done0, + tmem_stats_done1, + sage_v_scales, + work_queue, + ) + src = [tmem_softmax_local0, tmem_softmax_local1, tmem_o] + if tmem_stats_done0 is not None: + src.extend([tmem_stats_done0, tmem_stats_done1]) + dst = [tmem_corr0, tmem_corr1] + if sage_v_scales is not None: + # The scale block is produced and consumed by this task. + src.append(sage_v_scales) + dst.append(sage_v_scales) if work_queue is not None: src.append(work_queue) return task_class( src_resources=src, - dst_resources=[tmem_corr0, tmem_corr1], + dst_resources=dst, q_bound_resources=((tmem_corr0, True), (tmem_corr1, True)), cfg=cfg, warp_idx=cfg.correction_warp_idx, @@ -3623,8 +3822,9 @@ def create_correction_task_one_inst_qkv( work_queue: WorkQueue | None, cfg: FmhaDecodeConfig, *, - tmem_stats_done: MemoryResource, + tmem_stats_done: MemoryResource | None, domain: int | cutlass.Int32, + sage_v_scales: MemoryResource | None = None, task_class: type[DecodeGenTask] = DecodeGenTask, **kw: TaskKwarg, ) -> Task: @@ -3634,7 +3834,8 @@ def correction_schedule_body( tmem_softmax_local: MemoryResource, tmem_o: MemoryResource, tmem_corr: MemoryResource, - tmem_stats_done: MemoryResource, + tmem_stats_done: MemoryResource | None, + sage_v_scales: MemoryResource | None, ) -> None: """Schedule one-inst O correction and final output normalization.""" @@ -3677,7 +3878,7 @@ def consume_local_with_load( inst_new_max_arr=inst_new_max_arr, inst_sum_arr=inst_sum_arr, ) - if release_stats_done: + if release_stats_done and tmem_stats_done is not None: # Return one S-overwrite credit for this QK-derived payload. tmem_stats_done.wait() tmem_stats_done.release() @@ -3740,6 +3941,12 @@ def correct_o( local_state = tmem_softmax_local.init_stats_state() _, tail_0, tail_1 = tmem_o.init_stage_state() tmem_corr.init_epilogue_state() + if sage_v_scales is not None: + # ProdWork: stage the work tile's V channel scales while no + # output is pending. + sage_v_scales.acquire() + sage_v_scales.stage_tile() + sage_v_scales.commit() # HEAD: drain the first stats handoff. The first payload only seeds # the correction state; there is no prior O stage to rescale. @@ -3777,6 +3984,9 @@ def correct_o( tail_o_stage_idx_1=tail_1, inst_idx=KV_INST0, ) + if sage_v_scales is not None: + # ConsWait: every lane reads the scales the other lanes staged. + sage_v_scales.wait() # ProdWork: tail phase normalizes and stores final O. tmem_corr.correction_tail_epilogue( o_stage_idx=o_stage_idx, @@ -3789,16 +3999,20 @@ def correct_o( inst1_new_max_arr=inst_new_max_arr, inst1_sum_arr=inst_sum_arr, ) + if sage_v_scales is not None: + # ConsRelease: the next work tile may restage the block. + sage_v_scales.release() # ConsRelease: the final O stage is no longer needed. tmem_o.release() - @schedule + @_schedule_with_optional_resources def correction_schedule( tmem_softmax_local: MemoryResource, tmem_o: MemoryResource, tmem_corr: MemoryResource, - tmem_stats_done: MemoryResource, - work_queue: WorkQueue | None = None, + tmem_stats_done: MemoryResource | None, + sage_v_scales: MemoryResource | None, + work_queue: WorkQueue | None, ) -> None: """Wrap one-inst correction in packed persistent skip handling.""" _decode_work_tile_schedule( @@ -3809,31 +4023,31 @@ def correction_schedule( tmem_o, tmem_corr, tmem_stats_done, + sage_v_scales, ), ) - schedule_result = ( - correction_schedule( - tmem_softmax_local, - tmem_o, - tmem_corr, - tmem_stats_done, - ) - if work_queue is None - else correction_schedule( - tmem_softmax_local, - tmem_o, - tmem_corr, - tmem_stats_done, - work_queue, - ) + schedule_result = correction_schedule( + tmem_softmax_local, + tmem_o, + tmem_corr, + tmem_stats_done, + sage_v_scales, + work_queue, ) - src = [tmem_softmax_local, tmem_o, tmem_stats_done] + src = [tmem_softmax_local, tmem_o] + if tmem_stats_done is not None: + src.append(tmem_stats_done) + dst = [tmem_corr] + if sage_v_scales is not None: + # The scale block is produced and consumed by this task. + src.append(sage_v_scales) + dst.append(sage_v_scales) if work_queue is not None: src.append(work_queue) return task_class( src_resources=src, - dst_resources=[tmem_corr], + dst_resources=dst, q_bound_resources=((tmem_corr, True),), cfg=cfg, warp_idx=cfg.correction_warp_idx, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/reduction.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/reduction.py index 673d3205aa39..ebbdc53fdfbf 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 @@ -33,16 +33,16 @@ only split-KV is reduced here. """ -import math - import cutlass import cutlass.cute as cute +from .direct_sparse_metadata import DirectSparseMetadataView, HeadIndexedMetadataView from cuda.bindings import driver as cuda_drv from cutlass import Float32, Int32, Int64 from cutlass.experimental import primitives as prims from .fmha_decode_constants import ( FP32_BYTES, + FP8_P_QUANT_LOG2_SCALE, FP8_PACKED_OUTPUT_REGS_PER_THREAD, FP8_VALUES_PER_REG, FP16_VALUES_PER_REG, @@ -120,7 +120,7 @@ def _reduction_q_group_idx( @cute.jit def _reduction_active_splits_kv( cfg: cutlass.Constexpr[FmhaDecodeConfig], - g_seqlens_kv: cute.Pointer, + g_seqlens_kv: cute.Pointer | DirectSparseMetadataView, b_idx: Int32, h_r: Int32, logical_output_row_idx: Int32, @@ -139,7 +139,7 @@ def _reduction_active_splits_kv( @cute.jit def _reduce_exact_splits_body( o_iter: cute.Pointer, - g_seqlens_kv: cute.Pointer, + g_seqlens_kv: cute.Pointer | DirectSparseMetadataView, g_cu_seqlens_q: cute.Pointer, g_partial_o: cute.Pointer, g_partial_stats: cute.Pointer, @@ -147,18 +147,26 @@ def _reduce_exact_splits_body( g_q_output_rows: Int32, cfg: cutlass.Constexpr[FmhaDecodeConfig], static_full_split_prefix: cutlass.Constexpr[bool] = False, + wait_for_pdl_producer: cutlass.Constexpr[bool] = False, ) -> None: """Fold every split in one 512-thread CTA over one 8 KiB output slice. ``g_partial_stats`` stores one log2-LSE scalar for each split and output row. ``g_partial_o`` stores the corresponding normalized 16-bit O fragment. This body is shared by the serial reference kernel and the compact S2-S4 - production schedule; PDL ordering remains in the production outer kernel. + production schedule. The production caller requests its PDL acquire after + this body's local register state is initialized and before partial reads. """ thread_idx, _, _ = cute.arch.thread_idx() slice_idx, h_k_idx, b_idx = cute.arch.block_idx() _, grid_h_k, _ = cute.arch.grid_dim() + if cutlass.const_expr( + cfg.use_q_token_kv_block_sparse_route + and not cfg.shares_sparse_pattern + and not isinstance(g_seqlens_kv, DirectSparseMetadataView) + ): + g_seqlens_kv = HeadIndexedMetadataView(g_seqlens_kv, grid_h_k, h_k_idx) # Flatten batch and KV-head so the partial buffers use one contiguous # logical tile index independent of the reducer launch grid layout. logical_kv_idx = Int64(b_idx) * Int64(grid_h_k) + Int64(h_k_idx) @@ -177,15 +185,6 @@ def _reduce_exact_splits_body( ) q_token_offset, seq_len_q = _q_seq_bounds(cfg, g_cu_seqlens_q, b_idx) active_splits_kv = Int32(cfg.splits_kv) - if cutlass.const_expr(not static_full_split_prefix): - active_splits_kv = _reduction_active_splits_kv( - cfg, - g_seqlens_kv, - b_idx, - g_q_output_rows, - reduce_row_idx, - seq_len_q, - ) valid_reduce_row = _q_logical_output_row_is_valid_for_seq( cfg, g_q_output_rows, @@ -215,6 +214,22 @@ def _reduce_exact_splits_body( global_lse = Float32(-Float32.inf) + if cutlass.const_expr(wait_for_pdl_producer): + prims.griddepcontrol(kind=prims.GridDepAction.WAIT) + + if cutlass.const_expr(not static_full_split_prefix): + # Sparse-route seq_lens may originate two PDL stages upstream. Keep this read + # behind the attention-producer acquire together with partial O/stats; + # ordinary reducers preserve the same instruction ordering harmlessly. + active_splits_kv = _reduction_active_splits_kv( + cfg, + g_seqlens_kv, + b_idx, + g_q_output_rows, + reduce_row_idx, + seq_len_q, + ) + if valid_reduce_row: # The workspace retains configured-max strides, but only the runtime # active prefix was published by producer CTAs. @@ -287,7 +302,7 @@ def _reduce_exact_splits_body( @cute.kernel def decode_gen_separate_reduction_kernel( o_iter: cute.Pointer, - g_seqlens_kv: cute.Pointer, + g_seqlens_kv: cute.Pointer | DirectSparseMetadataView, g_cu_seqlens_q: cute.Pointer, g_partial_o: cute.Pointer, g_partial_stats: cute.Pointer, @@ -334,8 +349,8 @@ def _attention_sink_log2_lse( dtype=Float32, ) sink_lse = sink_ptr.load() * Float32(1.4426950408889634) - if cutlass.const_expr(cfg.use_fp8_qkv): - sink_lse += Float32(math.log2(448.0)) + if cutlass.const_expr(cfg.use_fp8_pv): + sink_lse += Float32(FP8_P_QUANT_LOG2_SCALE) return sink_lse @@ -433,7 +448,7 @@ def _store_parallel_reduction_output( @cute.kernel def decode_gen_parallel_separate_reduction_kernel( o_iter: cute.Pointer, - g_seqlens_kv: cute.Pointer, + g_seqlens_kv: cute.Pointer | DirectSparseMetadataView, g_cu_seqlens_q: cute.Pointer, g_partial_o: cute.Pointer, g_partial_stats: cute.Pointer, @@ -450,11 +465,6 @@ def decode_gen_parallel_separate_reduction_kernel( G16 uses a two-level 4x4 merge, and G2/G4/G8 finalize through rank zero. """ - # Pair with the producer's launch-dependents signal before reading any - # partial GMEM. The wait is CTA-convergent and stays outside both schedules. - if cutlass.const_expr(cfg.use_parallel_separate_reduction_pdl): - prims.griddepcontrol(kind=prims.GridDepAction.WAIT) - if cutlass.const_expr(cfg.use_compact_parallel_reduction): _reduce_exact_splits_body( o_iter, @@ -466,12 +476,19 @@ def decode_gen_parallel_separate_reduction_kernel( g_q_output_rows, cfg, static_full_split_prefix, + cfg.use_parallel_separate_reduction_pdl, ) return thread_idx, _, _ = cute.arch.thread_idx() block_idx_x, h_k_idx, b_idx = cute.arch.block_idx() _, grid_h_k, _ = cute.arch.grid_dim() + if cutlass.const_expr( + cfg.use_q_token_kv_block_sparse_route + and not cfg.shares_sparse_pattern + and not isinstance(g_seqlens_kv, DirectSparseMetadataView) + ): + g_seqlens_kv = HeadIndexedMetadataView(g_seqlens_kv, grid_h_k, h_k_idx) cluster_rank = cute.arch.block_idx_in_cluster() logical_kv_idx = Int64(b_idx) * Int64(grid_h_k) + Int64(h_k_idx) attention_sink_h_r = _attention_sink_head_stride(cfg, g_q_output_rows) @@ -487,15 +504,6 @@ def decode_gen_parallel_separate_reduction_kernel( ) q_token_offset, seq_len_q = _q_seq_bounds(cfg, g_cu_seqlens_q, b_idx) active_splits_kv = Int32(cfg.splits_kv) - if cutlass.const_expr(not static_full_split_prefix): - active_splits_kv = _reduction_active_splits_kv( - cfg, - g_seqlens_kv, - b_idx, - g_q_output_rows, - reduce_row_idx, - seq_len_q, - ) valid_reduce_row = _q_logical_output_row_is_valid_for_seq( cfg, g_q_output_rows, @@ -536,6 +544,37 @@ def decode_gen_parallel_separate_reduction_kernel( 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. + smem_lse = cutlass.Array( + Float32, + SEPARATE_REDUCTION_LSE_VALUES_PER_ROW * PARALLEL_REDUCTION_THREADS_PER_CTA, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + smem_partial_o = cutlass.Array( + Int32, + PACKED_OUTPUT_REGS_PER_THREAD * PARALLEL_REDUCTION_THREADS_PER_CTA, + space=cutlass.AddressSpace.smem, + alignment=16, + ) + + # Pair with the producer only after local state is ready, but before the + # first producer-dependent partial-O or partial-statistics read. + if cutlass.const_expr(cfg.use_parallel_separate_reduction_pdl): + prims.griddepcontrol(kind=prims.GridDepAction.WAIT) + + if cutlass.const_expr(not static_full_split_prefix): + active_splits_kv = _reduction_active_splits_kv( + cfg, + g_seqlens_kv, + b_idx, + g_q_output_rows, + reduce_row_idx, + seq_len_q, + ) + for split_base_i in cutlass.range_constexpr( 0, local_splits, PARALLEL_REDUCTION_LOAD_BATCH ): @@ -635,19 +674,6 @@ def decode_gen_parallel_separate_reduction_kernel( # profile's selected 16-bit partial type. Corresponding threads in every # rank map to the same output row, so row validity is cluster-uniform. A # rank with only padded split slots publishes the neutral ``(-inf, 0)``. - smem_lse = cutlass.Array( - Float32, - SEPARATE_REDUCTION_LSE_VALUES_PER_ROW * PARALLEL_REDUCTION_THREADS_PER_CTA, - space=cutlass.AddressSpace.smem, - alignment=16, - ) - smem_partial_o = cutlass.Array( - Int32, - PACKED_OUTPUT_REGS_PER_THREAD * PARALLEL_REDUCTION_THREADS_PER_CTA, - space=cutlass.AddressSpace.smem, - alignment=16, - ) - stats_smem_offset = thread_idx * Int32(SEPARATE_REDUCTION_LSE_VALUES_PER_ROW) partial_o_smem_offset = thread_idx * Int32(PACKED_OUTPUT_REGS_PER_THREAD) if valid_reduce_row: @@ -808,7 +834,7 @@ def decode_gen_parallel_separate_reduction_kernel( def fmha_decode_separate_reduction_launch( problem_shape: tuple[Int32, Int32, Int32, Int32, Int32], o_iter: cute.Pointer, - seqlens_kv_iter: cute.Pointer, + seqlens_kv_iter: cute.Pointer | DirectSparseMetadataView, cu_seqlens_q_iter: cute.Pointer, partial_o_iter: cute.Pointer, partial_stats_iter: cute.Pointer, 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 cadd4fb66b5f..0ae3fa0f8534 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 @@ -43,6 +43,7 @@ import cutlass import cutlass.cute as cute +from cutlass.experimental.task_scheduling.memory import ResourceContext from cutlass.experimental.task_scheduling.schedule_builder import ( domain_loop, schedule, @@ -51,6 +52,7 @@ from cutlass.experimental.task_scheduling.resources import StageInfo from cutlass.experimental.task_scheduling.task import Task +from ...task_compat import task_context_kwargs from .resources import ( MlaWorkQueue, WorkThrottleBarrierResource, @@ -149,14 +151,18 @@ class MlaTask(Task): """Task subclass that recomputes MLA k-domain per persistent work tile.""" @cute.jit - def _run_one_mla_work_tile(self, work_tile) -> None: + def _run_one_mla_work_tile( + self, + work_tile, + context: ResourceContext | None = None, + ) -> None: """Run a persistent work tile using the cached split-KV domain.""" # WorkQueue decomposes the MLA persistent tile and caches the K-domain. # Keep task bodies on that cached value so page-offset/TMA/MMA/softmax # paths do not each rebuild the same split-KV arithmetic. self.domain = work_tile.k_tile_count - self._run_task_body_impl(work_tile) + self._run_task_body_impl(work_tile, **task_context_kwargs(context)) @cute.jit def _drain_mla_work_tile_tails(self) -> None: @@ -188,7 +194,10 @@ def _drain_mla_work_tile_tails(self) -> None: self.work_queue.producer_tail() @cute.jit - def _run_task_body_persistent(self) -> None: + def _run_task_body_persistent( + self, + context: ResourceContext | None = None, + ) -> None: """Schedule one or more persistent work tiles for this task instance.""" params = self.work_queue.tile_sched_params @@ -196,13 +205,13 @@ def _run_task_body_persistent(self) -> None: if cutlass.const_expr(not params.is_persistent): work_tile = self.work_queue._work_tile_from_block_idx(cute.arch.block_idx()) self.work_queue._set_consumer_var_from_ts("work_tile", work_tile) - self._run_pre_work_loop_entries(work_tile) + self._run_pre_work_loop_entries(work_tile, **task_context_kwargs(context)) # Runtime K/Q metadata can make a statically launched split empty. # Keep the CTA on the ordinary initialized-pipeline path, but skip # its captured HEAD/LOOP/TAIL data work when the domain is zero. if work_tile.k_tile_count > cutlass.Int32(0): - self._run_one_mla_work_tile(work_tile) - self._run_post_work_loop_entries(work_tile) + self._run_one_mla_work_tile(work_tile, context) + self._run_post_work_loop_entries(work_tile, **task_context_kwargs(context)) self._drain_mla_work_tile_tails() return @@ -216,7 +225,7 @@ def _run_task_body_persistent(self) -> None: work_tile = self.work_queue._work_tile_from_linear_idx(current_work_linear_idx) self.work_queue._set_consumer_var_from_ts("work_tile", work_tile) - self._run_pre_work_loop_entries(work_tile) + self._run_pre_work_loop_entries(work_tile, **task_context_kwargs(context)) while current_work_linear_idx < num_blocks: work_tile.update_from( self.work_queue._work_tile_from_linear_idx(current_work_linear_idx) @@ -227,7 +236,7 @@ def _run_task_body_persistent(self) -> None: # splits empty. Skip them and continue grid-striding rather than # running a captured HEAD/TAIL sequence with domain zero. if work_tile.k_tile_count > cutlass.Int32(0): - self._run_one_mla_work_tile(work_tile) + self._run_one_mla_work_tile(work_tile, context) # Each warp branch advances from the same scalar tile id, keeping # the persistent loop state compact across task bodies. @@ -239,7 +248,7 @@ def _run_task_body_persistent(self) -> None: self.work_queue._work_tile_from_linear_idx(current_work_linear_idx) ) self.work_queue._set_consumer_var_from_ts("work_tile", work_tile) - self._run_post_work_loop_entries(work_tile) + self._run_post_work_loop_entries(work_tile, **task_context_kwargs(context)) self._drain_mla_work_tile_tails() @@ -256,7 +265,11 @@ def __init__(self, *args, **kwargs) -> None: self._fixed_loop_end = cutlass.Int32(self.domain_start) @cute.jit - def _run_one_mla_work_tile(self, work_tile) -> None: + def _run_one_mla_work_tile( + self, + work_tile, + context: ResourceContext | None = None, + ) -> None: """Run one tile by mapping its local offsets onto this fixed lane.""" fixed_lane = cutlass.Int32(self.domain_start) @@ -269,7 +282,7 @@ def _run_one_mla_work_tile(self, work_tile) -> None: ) // cutlass.Int32(2) self._fixed_loop_end = fixed_lane + lane_iterations * cutlass.Int32(2) self.domain = self._fixed_loop_end - self._run_task_body_impl(work_tile) + self._run_task_body_impl(work_tile, **task_context_kwargs(context)) self._cumulative_k_parity = ( self._cumulative_k_parity + self._actual_domain ) % cutlass.Int32(2) @@ -285,6 +298,7 @@ def _create_stage_info( label=None, schedule_stage=None, routing_slot=None, + context: ResourceContext | None = None, ) -> StageInfo: """Return base pipeline state with the work tile's actual K offset.""" @@ -298,6 +312,7 @@ def _create_stage_info( label, schedule_stage, routing_slot, + **task_context_kwargs(context), ) actual_loop_offset = self._mapped_domain_start + ( cutlass.Int32(base_info.loop_offset) - cutlass.Int32(self.domain_start) 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 bee9c7bfa7a7..11d39bdec7e6 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 @@ -223,18 +223,14 @@ def _run_parallel_gmem_reduction_g1_shared_stats( # One statistics warp per covered row stripes over split slots, then # publishes normalized scales for all output threads to reuse. + # Keep branch-local names separate from the output phase below so CuTe + # DSL 4.8 does not thread undefined values through inactive stats warps. if warp_idx < Int32(rows_per_slice): stats_row = warp_idx stats_element = slice_element_base + stats_row * Int32(cfg.head_dim_per_cta_v) - head_idx = stats_element // Int32(cfg.head_dim_per_cta_v) - ( - storage_flat_query_row, - _, - _, - _, - valid_output_row, - ) = flat_query_row_state( - head_idx, + stats_head_idx = stats_element // Int32(cfg.head_dim_per_cta_v) + stats_query_row = flat_query_row_state( + stats_head_idx, q_idx, cfg.tile_size_q, cfg.logical_num_heads_q, @@ -242,8 +238,8 @@ def _run_parallel_gmem_reduction_g1_shared_stats( cu_seqlens_q=cu_seqlens_q, batch_idx=batch_idx, ) - if head_idx < Int32(cfg.num_heads_q) and valid_output_row: - row_lse = acc_lse[head_idx, None, q_idx, batch_idx] + if stats_head_idx < Int32(cfg.num_heads_q) and stats_query_row[4]: + row_lse = acc_lse[stats_head_idx, None, q_idx, batch_idx] lse_per_lane = ceil_div( cfg.num_ctas_per_seq_kv, GMEM_REDUCTION_WARP_LANES, @@ -255,10 +251,12 @@ def _run_parallel_gmem_reduction_g1_shared_stats( ) lse_max = Float32(-Float32.inf) for lane_slot_i in cutlass.range_constexpr(lse_per_lane): - split_idx = lane_idx + Int32(lane_slot_i * GMEM_REDUCTION_WARP_LANES) + stats_split_idx = lane_idx + Int32( + lane_slot_i * GMEM_REDUCTION_WARP_LANES + ) lane_lse[lane_slot_i] = ( - Float32(row_lse[split_idx]) - if split_idx < Int32(cfg.num_ctas_per_seq_kv) + 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]) @@ -278,19 +276,21 @@ def _run_parallel_gmem_reduction_g1_shared_stats( and head_dim_cta_idx == Int32(0) and stats_element % Int32(cfg.head_dim_per_cta_v) == Int32(0) ): - output_query_row = public_query_flat_row( + stats_output_query_row = public_query_flat_row( cfg, - storage_flat_query_row, + stats_query_row[0], batch_idx, cu_seqlens_q, ) - (lse.iterator.raw_ptr() + output_query_row).store(global_lse) + (lse.iterator.raw_ptr() + stats_output_query_row).store(global_lse) for lane_slot_i in cutlass.range_constexpr(lse_per_lane): - split_idx = lane_idx + Int32(lane_slot_i * GMEM_REDUCTION_WARP_LANES) - if split_idx < Int32(cfg.num_ctas_per_seq_kv): + stats_split_idx = lane_idx + Int32( + lane_slot_i * GMEM_REDUCTION_WARP_LANES + ) + if stats_split_idx < Int32(cfg.num_ctas_per_seq_kv): smem_scale[ - stats_row * Int32(cfg.num_ctas_per_seq_kv) + split_idx + stats_row * Int32(cfg.num_ctas_per_seq_kv) + stats_split_idx ] = normalized_lse_weight(lane_lse[lane_slot_i], global_lse) prims.barrier_cta_sync( @@ -304,13 +304,7 @@ def _run_parallel_gmem_reduction_g1_shared_stats( dim_in_cta = flat_element - head_idx * Int32(cfg.head_dim_per_cta_v) row_in_slice = (flat_element - slice_element_base) // Int32(cfg.head_dim_per_cta_v) dim_idx = head_dim_cta_idx * Int32(cfg.head_dim_per_cta_v) + dim_in_cta - ( - storage_flat_query_row, - _, - _, - _, - valid_output_row, - ) = flat_query_row_state( + output_query_row_state = flat_query_row_state( head_idx, q_idx, cfg.tile_size_q, @@ -323,7 +317,7 @@ def _run_parallel_gmem_reduction_g1_shared_stats( thread_elem_offset < Int32(elements_per_slice) and head_idx < Int32(cfg.num_heads_q) and dim_idx < Int32(cfg.head_dim_v) - and valid_output_row + and output_query_row_state[4] ) if valid_output: output_acc = cutlass.Array( @@ -361,7 +355,7 @@ def _run_parallel_gmem_reduction_g1_shared_stats( output_query_row = public_query_flat_row( cfg, - storage_flat_query_row, + output_query_row_state[0], batch_idx, cu_seqlens_q, ) 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 e3d2e64eb937..bf7715ce4ddf 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 @@ -14,6 +14,7 @@ """Captured task schedules for the throughput-latency 1CTA MLA TS path.""" +import cutlass import cutlass.cute as cute from cutlass import Int32 from cutlass.experimental import primitives as prims @@ -67,6 +68,7 @@ def __init__(self, **kwargs): self._lane_idx = Int32(0) self._seq_len_kv = Int32(0) + @cute.jit def init_variables(self, context=None): """Initialize per-task thread cache and TMEM base state.""" super().init_variables(context) @@ -103,21 +105,22 @@ def make_task_cache(self): Int32(0), ) + @cute.jit def get_domain(self, tile_coord): """Return the runtime loop domain for the current 1CTA work tile.""" - if self.cfg is None: + if cutlass.const_expr(self.cfg is None): 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. - if isinstance(tile_coord[2], int): - if self.cfg.use_persistent_scheduler == 1: + if cutlass.const_expr(isinstance(tile_coord[2], int)): + if cutlass.const_expr(self.cfg.use_persistent_scheduler == 1): batch_idx = tile_coord[2] // self.cfg.num_ctas_for_all_heads else: batch_idx = tile_coord[2] else: batch_idx = Int32(tile_coord[2]) - if self.cfg.use_persistent_scheduler == 1: + if cutlass.const_expr(self.cfg.use_persistent_scheduler == 1): batch_idx = batch_idx // Int32(self.cfg.num_ctas_for_all_heads) # Task-domain ownership must match K/V, softmax, and reduction: dense @@ -126,7 +129,7 @@ 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 self.cfg.use_persistent_scheduler != 1: + 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. diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/task_compat.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/task_compat.py new file mode 100644 index 000000000000..69f60317b172 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/task_compat.py @@ -0,0 +1,39 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Compatibility for the CuTe DSL 4.7/4.8 task execution APIs.""" + +import inspect + +from cutlass.experimental.task_scheduling.memory import ResourceContext +from cutlass.experimental.task_scheduling.task import Task + + +_TASK_ACCEPTS_CONTEXT = ( + "context" in inspect.signature(Task._run_task_body_impl).parameters +) + + +def task_context_kwargs( + context: ResourceContext | None, +) -> dict[str, ResourceContext | None]: + """Supply the explicit execution context only on the 4.7 task API. + + DSL 4.8 removed this argument from task execution and stage-info methods: + ``init_variables`` now stores it on the task. Keep the 4.7 propagation, + including ``None``, while allowing both APIs to call PrimTS overrides. + Signature detection happens once at import; this helper runs at trace time. + """ + return {"context": context} if _TASK_ACCEPTS_CONTEXT else {} 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 cff2ee7d2ad0..6d69ae7b554a 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py @@ -22,16 +22,8 @@ import torch -from flashinfer.api_logging import flashinfer_api -from flashinfer.trace.templates.attention import ( - prims_ts_decode_mla_one_shot_trace_dispatch, - prims_ts_decode_mla_trace_dispatch, -) +from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api -from ._tensor_aliasing import ( - _validate_out_does_not_overlap_inputs, - _validate_tensor_does_not_overlap_inputs, -) from .decode import ( _WorkspaceSection, _align_up, @@ -394,11 +386,19 @@ def _derive_max_seq_len_q( def _validate_mla_run_metadata( - state: _MLADecodePlanState, runtime: _MLARuntime, block_tables: torch.Tensor, seq_lens: 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_kv_len: int, + page_size: int, + mask_type: str, ) -> None: """Validate per-run request metadata against one static MLA plan. @@ -406,64 +406,66 @@ def _validate_mla_run_metadata( is the synchronization-free path for compilation and graph capture. """ - device, batch_size, max_num_pages = _validate_mla_metadata(block_tables, seq_lens) - if device != state.device: + metadata_device, metadata_batch_size, max_num_pages = _validate_mla_metadata( + block_tables, seq_lens + ) + if metadata_device != device: raise ValueError( - f"MLA metadata must be on the planned device {state.device}, got {device}" + f"MLA metadata must be on the planned device {device}, got {metadata_device}" ) - if batch_size != state.batch_size: + if metadata_batch_size != batch_size: raise ValueError( "MLA metadata batch size must match the plan " - f"({state.batch_size}), got {batch_size}" + f"({batch_size}), got {metadata_batch_size}" ) - if max_num_pages < state.required_page_columns: + if max_num_pages < required_page_columns: raise ValueError( "block_tables must have at least ceil(max_kv_len / page_size) " - f"columns ({state.required_page_columns}), got {max_num_pages}" + f"columns ({required_page_columns}), got {max_num_pages}" ) - if state.packed_query: + if packed_query: if qo_indptr is None: raise ValueError("packed-query MLA run requires qo_indptr") _validate_qo_indptr( qo_indptr, - device=state.device, - batch_size=state.batch_size, + device=device, + batch_size=batch_size, ) runtime_max_seq_len_q, total_q, q_lengths = _derive_max_seq_len_q( qo_indptr, - batch_size=state.batch_size, + batch_size=batch_size, ) if total_q != int(runtime.query.shape[0]): raise ValueError( "qo_indptr must end at the packed query row count " f"({runtime.query.shape[0]}), got {total_q}" ) - if runtime_max_seq_len_q > state.max_seq_len_q: + if runtime_max_seq_len_q > max_seq_len_q: raise ValueError( "qo_indptr contains a per-request Q length larger than " - f"max_seq_len_q ({state.max_seq_len_q}): got " + f"max_seq_len_q ({max_seq_len_q}): 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 = (state.max_seq_len_q,) * state.batch_size + q_lengths = (max_seq_len_q,) * batch_size seq_lens_host = tuple(int(value) for value in seq_lens.tolist()) if any(seq_len <= 0 for seq_len in seq_lens_host): raise ValueError("every runtime request must contain at least one KV token") runtime_max_kv_len = max(seq_lens_host) - if runtime_max_kv_len > state.max_kv_len: + if runtime_max_kv_len > max_kv_len: raise ValueError( "runtime KV metadata contains a request longer than " - f"max_kv_len ({state.max_kv_len}): got {runtime_max_kv_len}" + f"max_kv_len ({max_kv_len}): got {runtime_max_kv_len}" ) block_table_rows = block_tables.tolist() for request_idx, (row, seq_len) in enumerate( zip(block_table_rows, seq_lens_host, strict=True) ): - required_pages = _ceil_div(seq_len, state.page_size) + required_pages = _ceil_div(seq_len, page_size) if any( int(page_id) < 0 or int(page_id) >= runtime.num_physical_pages for page_id in row[:required_pages] @@ -473,7 +475,7 @@ def _validate_mla_run_metadata( f"K/V cache in [0, {runtime.num_physical_pages}); request " f"{request_idx} contains an invalid page ID" ) - if state.mask_type == "causal": + if mask_type == "causal": for request_idx, (q_len, kv_len) in enumerate( zip(q_lengths, seq_lens_host, strict=True) ): @@ -1207,7 +1209,7 @@ 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:`prims_ts_batch_mla_decode_with_kv_cache`, without compiling a kernel. + :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 packed-query launches; ``seq_len_q`` remains a backward-compatible fixed-Q alias. If neither is @@ -1237,7 +1239,8 @@ def get_prims_ts_batch_mla_decode_workspace_size( if kv_dtype is None: kv_dtype = q_dtype _validate_mla_dtype_pair(q_dtype, kv_dtype, out_dtype) - _, device_index = _resolve_cuda_device(device) + resolved_device, device_index = _resolve_cuda_device(device) + _validate_runtime_device(resolved_device) spec = _resolve_mla_decode_launch_spec( device_index, @@ -1350,27 +1353,6 @@ def _prepare_mla_runtime( ) -def _validate_mla_output_aliasing( - runtime: _MLARuntime, - *, - block_tables: torch.Tensor, - seq_lens: torch.Tensor, - qo_indptr: Optional[torch.Tensor], - workspace_buffer: torch.Tensor, -) -> None: - """Keep output disjoint from every live MLA decode allocation.""" - - _validate_out_does_not_overlap_inputs( - runtime.out, - ("query", runtime.query), - ("kv_cache", runtime.normalized_cache), - ("block_tables", block_tables), - ("seq_lens", seq_lens), - ("qo_indptr", qo_indptr), - ("workspace_buffer", workspace_buffer), - ) - - def _launch_mla_decode( runtime: _MLARuntime, *, @@ -1422,234 +1404,15 @@ def _launch_mla_decode( return runtime.out -@flashinfer_api(trace=prims_ts_decode_mla_trace_dispatch) -def prims_ts_batch_mla_decode_with_kv_cache( - query: torch.Tensor, - kv_cache: torch.Tensor, - workspace_buffer: torch.Tensor, - kv_lora_rank: int, - qk_rope_head_dim: int, - block_tables: torch.Tensor, - seq_lens: torch.Tensor, - max_seq_len: int, - *, - qo_indptr: Optional[torch.Tensor] = None, - max_seq_len_q: Optional[int] = None, - out: Optional[torch.Tensor] = None, - bmm1_scale: float = 1.0, - bmm2_scale: float = 1.0, - mask_type: Literal["dense", "causal"] = "causal", - out_dtype: torch.dtype = torch.bfloat16, -) -> torch.Tensor: - """Launch fixed or packed-query paged MLA decode with caller-owned scratch. - - With ``qo_indptr=None``, ``query`` has fixed shape ``[B, SQ, H, 576]``. - Otherwise ``query`` has compact shape ``[total_q, H, 576]`` and - ``qo_indptr`` contains the ``B + 1`` cumulative Q offsets. Runtime Q - lengths are exclusively ``qo_indptr[b + 1] - qo_indptr[b]``; - ``max_seq_len_q`` is only the static policy, JIT, and workspace bound and - is required for compact launches. Individual packed requests may be empty, - and an all-empty launch returns its empty output without dispatching a GPU - kernel. The last query dimension concatenates the 512 latent and 64 RoPE - dimensions. ``kv_cache`` accepts compact rank-3 - ``[pages, page_size, 576]`` or rank-4 ``[pages, 1, page_size, 576]`` - storage. ``block_tables`` and ``seq_lens`` follow FlashInfer's native dense - paged-cache ABI. The table is contiguous within each row and may have - padding between rows; ``max_seq_len`` is the exact static policy/JIT - maximum. - Causal masking is bottom-right aligned: query row ``i`` can attend through - KV row ``seq_lens[b] - q_len[b] + i`` for request ``b``. - - The workspace is exclusive to one in-flight launch or captured graph and - must not overlap query, K/V cache, metadata, or output storage. - Runtime K/V lengths must remain positive and no larger than ``max_seq_len``; - this hot path deliberately performs no device-to-host metadata reads. For - packed launches, callers must ensure that offsets start at zero, are - nondecreasing, end at ``query.shape[0]``, and have every delta no - larger than ``max_seq_len_q``. For causal masking, every fixed or packed - per-request Q length must also be no greater than the corresponding live - ``seq_lens`` value. Warm the planned topology before CUDA graph - capture and provide ``out`` to avoid an output allocation. Captured graphs - must retain stable ``block_tables``, ``seq_lens``, and, for packed Q, - ``qo_indptr`` storage. Values may change only between completed replays - while the runtime metadata contracts and captured query/output extents - remain valid. No backend fallback or scheduling knob is exposed. - - Parameters - ---------- - query : torch.Tensor - Fixed or packed query tensor with concatenated latent and RoPE heads. - kv_cache : torch.Tensor - Compact paged latent K/V cache. - workspace_buffer : torch.Tensor - Caller-owned byte workspace for this planned layout. - kv_lora_rank, qk_rope_head_dim : int - Latent and RoPE dimensions. - 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 - Live K/V sequence lengths. - max_seq_len : int - Static maximum K/V length used for policy selection and JIT caching. - qo_indptr : torch.Tensor, optional - Cumulative query offsets selecting packed-query mode. - max_seq_len_q : int, optional - Static packed-query length bound. - out : torch.Tensor, optional - Caller-owned output tensor. - bmm1_scale, bmm2_scale : float - QK and value/output scaling factors. - mask_type : {"dense", "causal"} - Attention mask mode. - out_dtype : torch.dtype - Output dtype. - """ - - packed_query = qo_indptr is not None - _validate_query(query, packed_query=packed_query) - metadata_device, batch_size, max_num_pages = _validate_mla_metadata( - block_tables, seq_lens - ) - if metadata_device != query.device: - raise ValueError( - f"MLA metadata must be on {query.device}, got {metadata_device}" - ) - normalized_cache, _, page_size = _normalize_mla_kv_cache( - kv_cache, expected_device=query.device - ) - if packed_query: - _validate_qo_indptr( - qo_indptr, - device=query.device, - batch_size=batch_size, - ) - if max_seq_len_q is None: - raise ValueError( - "max_seq_len_q is required when qo_indptr selects packed query" - ) - max_seq_len_q = _validate_positive_int(max_seq_len_q, "max_seq_len_q") - num_heads = int(query.shape[1]) - else: - fixed_seq_len_q = int(query.shape[1]) - if max_seq_len_q is None: - max_seq_len_q = fixed_seq_len_q - else: - max_seq_len_q = _validate_positive_int(max_seq_len_q, "max_seq_len_q") - if max_seq_len_q != fixed_seq_len_q: - raise ValueError( - "fixed query length must equal max_seq_len_q: " - f"got SQ={fixed_seq_len_q} and max_seq_len_q={max_seq_len_q}" - ) - num_heads = int(query.shape[2]) - _validate_mla_dims(kv_lora_rank, qk_rope_head_dim) - _validate_page_size(page_size) - max_seq_len = _validate_mla_max_kv_len(max_seq_len, "max_seq_len") - required_page_columns = _ceil_div(max_seq_len, page_size) - if max_num_pages < required_page_columns: - raise ValueError( - "block_tables must have at least ceil(max_seq_len / page_size) " - f"columns ({required_page_columns}), got {max_num_pages}" - ) - _validate_mask(mask_type) - _validate_mla_dtype_pair(query.dtype, normalized_cache.dtype, out_dtype) - device_index = _validate_runtime_device(query.device) - spec_key = ( - device_index, - batch_size, - num_heads, - kv_lora_rank, - qk_rope_head_dim, - page_size, - max_seq_len, - _dtype_key(query.dtype), - _dtype_key(normalized_cache.dtype), - _dtype_key(out_dtype), - mask_type, - max_seq_len_q, - ) - spec = _resolve_mla_decode_launch_spec(*spec_key) - layout = _make_mla_workspace_layout( - spec.kernel_workspace_bytes, batch_size, num_heads, max_seq_len_q - ) - _validate_workspace_buffer( - workspace_buffer, - device=query.device, - required_bytes=layout.total_bytes, - ) - caller_provided_out = out is not None - runtime = _prepare_mla_runtime( - query, - normalized_cache, - device=query.device, - batch_size=batch_size, - num_heads=num_heads, - max_seq_len_q=max_seq_len_q, - packed_query=packed_query, - qo_indptr=qo_indptr, - page_size=page_size, - q_dtype=query.dtype, - kv_dtype=normalized_cache.dtype, - output_dtype=out_dtype, - bmm1_scale=bmm1_scale, - bmm2_scale=bmm2_scale, - out=out, - validate=True, - ) - _validate_tensor_does_not_overlap_inputs( - workspace_buffer, - "workspace_buffer", - ("query", runtime.query), - ("kv_cache", runtime.normalized_cache), - ("block_tables", block_tables), - ("seq_lens", seq_lens), - ("qo_indptr", qo_indptr), - ("out", runtime.out), - ) - if caller_provided_out: - _validate_mla_output_aliasing( - runtime, - block_tables=block_tables, - seq_lens=seq_lens, - qo_indptr=qo_indptr, - workspace_buffer=workspace_buffer, - ) - compile_spec = _make_mla_decode_compile_spec( - spec, - device_index=device_index, - num_heads=num_heads, - kv_lora_rank=kv_lora_rank, - qk_rope_head_dim=qk_rope_head_dim, - page_size=page_size, - q_dtype_key=_dtype_key(query.dtype), - output_dtype_key=_dtype_key(out_dtype), - max_seq_len_q=max_seq_len_q, - packed_query=packed_query, - ) - compiled = _get_compiled_mla_decode(compile_spec) - workspace = _bind_mla_workspace(workspace_buffer, layout) - return _launch_mla_decode( - runtime, - block_tables=block_tables, - seq_lens=seq_lens, - qo_indptr=qo_indptr, - packed_query=packed_query, - kv_lora_rank=kv_lora_rank, - split_kv=spec.split_kv, - workspace=workspace, - compiled=compiled, - ) - - class BatchMLADecodePagedTSWrapper: """Compile and reuse task-scheduled paged MLA decode launches.""" - @flashinfer_api + @flashinfer_experimental_api def __init__(self) -> None: """Initialize an unplanned task-scheduled paged-MLA wrapper.""" self._plan_state: Optional[_MLADecodePlanState] = None - @flashinfer_api + @flashinfer_experimental_api def plan( self, device: int | str | torch.device, @@ -1667,6 +1430,7 @@ def plan( o_data_type: torch.dtype, mask_type: Literal["dense", "causal"] = "causal", workspace_buffer: Optional[torch.Tensor] = None, + validate: bool = True, ) -> None: """Compile one static MLA shape and bind its reusable workspace. @@ -1709,25 +1473,35 @@ def plan( must be 32-byte aligned and large enough for the selected plan. When omitted, planning allocates the buffer. The retained buffer is exclusive to one in-flight launch or graph replay. + validate : bool + Validate static geometry and caller scratch. Defaults to ``True``. + Disable only for previously validated inputs and warmed topology. """ - if not isinstance(packed_query, bool): - raise TypeError("packed_query must be a bool") - _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") - _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") - _validate_mla_query_head_extent( - batch_size=batch_size, - num_heads=num_heads, - max_seq_len_q=max_seq_len_q, - ) - _validate_mla_dtype_pair(q_data_type, kv_data_type, o_data_type) + if not isinstance(validate, bool): + raise TypeError("validate must be a bool") + if validate: + if not isinstance(packed_query, bool): + raise TypeError("packed_query must be a bool") + _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") + _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") + _validate_mla_query_head_extent( + batch_size=batch_size, + num_heads=num_heads, + max_seq_len_q=max_seq_len_q, + ) + _validate_mla_dtype_pair(q_data_type, kv_data_type, o_data_type) + device, device_index = _resolve_cuda_device(device) + if validate: + _validate_runtime_device(device) + required_page_columns = _ceil_div(max_kv_len, page_size) spec_key = ( @@ -1765,7 +1539,7 @@ def plan( workspace_buffer = torch.empty( workspace_layout.total_bytes, device=device, dtype=torch.int8 ) - else: + elif validate: _validate_workspace_buffer( workspace_buffer, device=device, @@ -1799,7 +1573,7 @@ def plan( split_kv=int(dict(policy)["split_kv"]), ) - @flashinfer_api + @flashinfer_experimental_api def run( self, query: torch.Tensor, @@ -1818,11 +1592,14 @@ def run( ``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, aliases, and every static capacity are checked before + 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 @@ -1855,7 +1632,6 @@ def run( if not isinstance(validate, bool): raise TypeError("validate must be a bool") runtime_qo_indptr = qo_indptr if state.packed_query else None - caller_provided_out = out is not None runtime = _prepare_mla_runtime( query, kv_cache, @@ -1876,30 +1652,19 @@ def run( ) if validate: _validate_mla_run_metadata( - state, runtime, block_tables, seq_lens, 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_kv_len=state.max_kv_len, + page_size=state.page_size, + mask_type=state.mask_type, ) - _validate_tensor_does_not_overlap_inputs( - state.workspace_buffer, - "workspace_buffer", - ("query", runtime.query), - ("kv_cache", runtime.normalized_cache), - ("block_tables", block_tables), - ("seq_lens", seq_lens), - ("qo_indptr", runtime_qo_indptr), - ("out", runtime.out), - ) - if caller_provided_out: - _validate_mla_output_aliasing( - runtime, - block_tables=block_tables, - seq_lens=seq_lens, - qo_indptr=runtime_qo_indptr, - workspace_buffer=state.workspace_buffer, - ) return _launch_mla_decode( runtime, block_tables=block_tables, @@ -1913,7 +1678,7 @@ def run( ) -@flashinfer_api(trace=prims_ts_decode_mla_one_shot_trace_dispatch) +@flashinfer_experimental_api def batch_mla_decode_with_paged_kv_cache( query: torch.Tensor, kv_cache: torch.Tensor, @@ -1930,13 +1695,14 @@ def batch_mla_decode_with_paged_kv_cache( bmm2_scale: float = 1.0, out: Optional[torch.Tensor] = None, out_dtype: torch.dtype = torch.bfloat16, + workspace_buffer: Optional[torch.Tensor] = None, + validate: bool = True, ) -> torch.Tensor: """One-shot convenience wrapper for fixed or packed-query MLA decode. - This helper reads ``seq_lens`` and, for packed Q, ``qo_indptr`` on the host - to derive plan bounds, then constructs a temporary wrapper. Invoke it - outside CUDA Graph capture. Capture-sensitive callers should pre-plan - :class:`BatchMLADecodePagedTSWrapper` and use ``run(validate=False)``. + 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. Parameters ---------- @@ -1969,106 +1735,88 @@ def batch_mla_decode_with_paged_kv_cache( 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. + Returns ------- torch.Tensor The fixed or packed MLA attention output. + Notes + ----- + 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``. + Warm this exact topology outside capture first. Retain stable tensor storage + and mutate metadata only between completed launches/replays. All live K/V + lengths must be positive and within the static bound, active page IDs must + index the cache, and causal per-request Q lengths must not exceed K/V lengths. + Packed offsets must start at zero, end at the query token count, and have + nonnegative deltas within the Q bound. Scratch/output must not alias + any inputs or each other. No metadata is copied to the host on the trusted + explicit path. Kernel policy, necessary control resets, and output layout + are unchanged. Missing output may be allocated only outside capture. + """ + if not isinstance(validate, bool): + raise TypeError("validate must be a bool") + if validate or workspace_buffer is None or out is None: + if torch.cuda.is_initialized() and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "CUDA graph capture requires workspace_buffer, max_kv_len, " + "out, validate=False, and explicit packed-Q bounds; warm up first" + ) + if not validate: + if workspace_buffer is None: + 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") + packed_query = qo_indptr is not None - _validate_query(query, packed_query=packed_query) - metadata_device, batch_size, _ = _validate_mla_metadata(block_tables, seq_lens) - if metadata_device != query.device: - raise ValueError( - f"MLA metadata must be on {query.device}, got {metadata_device}" - ) - normalized_cache, _, page_size = _normalize_mla_kv_cache( - kv_cache, expected_device=query.device - ) - _validate_mla_dims(kv_lora_rank, qk_rope_head_dim) - _validate_page_size(page_size) - _validate_mla_dtype_pair(query.dtype, normalized_cache.dtype, out_dtype) - if packed_query: - _validate_qo_indptr( - qo_indptr, - device=query.device, - batch_size=batch_size, - ) - num_heads = int(query.shape[1]) - derived_max_seq_len_q, total_q, runtime_q_lengths = _derive_max_seq_len_q( - qo_indptr, - batch_size=batch_size, - ) - if total_q != int(query.shape[0]): + if validate: + _validate_query(query, packed_query=packed_query) + metadata_device, batch_size, _ = _validate_mla_metadata(block_tables, seq_lens) + if metadata_device != query.device: raise ValueError( - "qo_indptr must end at the packed query row count " - f"({query.shape[0]}), got {total_q}" + f"MLA metadata must be on {query.device}, got {metadata_device}" ) - if max_seq_len_q is None: - if derived_max_seq_len_q == 0: - raise ValueError( - "max_seq_len_q is required for an all-empty packed query" - ) - max_seq_len_q = derived_max_seq_len_q - else: - max_seq_len_q = _validate_positive_int(max_seq_len_q, "max_seq_len_q") - if derived_max_seq_len_q > max_seq_len_q: - raise ValueError( - "qo_indptr contains a per-request Q length larger than " - f"max_seq_len_q ({max_seq_len_q}): got " - f"{derived_max_seq_len_q}" - ) - _validate_mla_query_head_extent( - batch_size=batch_size, - num_heads=num_heads, - max_seq_len_q=max_seq_len_q, - total_q=int(query.shape[0]), + normalized_cache, _, page_size = _normalize_mla_kv_cache( + kv_cache, expected_device=query.device ) - else: - num_heads = int(query.shape[2]) - fixed_seq_len_q = int(query.shape[1]) - if max_seq_len_q is None: - max_seq_len_q = fixed_seq_len_q - else: - max_seq_len_q = _validate_positive_int(max_seq_len_q, "max_seq_len_q") - if max_seq_len_q != fixed_seq_len_q: - raise ValueError( - "fixed query length must equal max_seq_len_q: " - f"got SQ={fixed_seq_len_q} and " - f"max_seq_len_q={max_seq_len_q}" + 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 ) - _validate_mla_query_head_extent( - batch_size=batch_size, - num_heads=num_heads, - max_seq_len_q=max_seq_len_q, - ) - runtime_q_lengths = (max_seq_len_q,) * batch_size - seq_lens_host = tuple(int(value) for value in seq_lens.tolist()) - if any(seq_len <= 0 for seq_len in seq_lens_host): - raise ValueError("every runtime request must contain at least one KV token") - metadata_max_kv_len = max(seq_lens_host) - if max_kv_len is None: - max_kv_len = metadata_max_kv_len - else: - max_kv_len = _validate_mla_max_kv_len(max_kv_len, "max_kv_len") - if metadata_max_kv_len > max_kv_len: + 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]): raise ValueError( - "runtime KV metadata contains a request longer than " - f"max_kv_len ({max_kv_len}): got {metadata_max_kv_len}" + "fixed query length must equal max_seq_len_q: " + f"got SQ={query.shape[1]} and max_seq_len_q={max_seq_len_q}" ) - if mask_type == "causal": - for request_idx, (q_len, kv_len) in enumerate( - zip(runtime_q_lengths, seq_lens_host, strict=True) - ): - if 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}" - ) + else: + batch_size = int(seq_lens.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]) + if max_kv_len is None: + max_kv_len = int(seq_lens.max().item()) assert max_seq_len_q is not None - assert max_kv_len is not None - if out is not None: + if validate and out is not None: _validate_out( out, device=query.device, @@ -2079,6 +1827,24 @@ def batch_mla_decode_with_paged_kv_cache( 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( @@ -2095,6 +1861,8 @@ def batch_mla_decode_with_paged_kv_cache( kv_data_type=normalized_cache.dtype, o_data_type=out_dtype, mask_type=mask_type, + workspace_buffer=workspace_buffer, + validate=validate, ) return wrapper.run( query, @@ -2105,6 +1873,7 @@ def batch_mla_decode_with_paged_kv_cache( bmm1_scale=bmm1_scale, bmm2_scale=bmm2_scale, out=out, + validate=validate, ) @@ -2112,5 +1881,4 @@ def batch_mla_decode_with_paged_kv_cache( "BatchMLADecodePagedTSWrapper", "batch_mla_decode_with_paged_kv_cache", "get_prims_ts_batch_mla_decode_workspace_size", - "prims_ts_batch_mla_decode_with_kv_cache", ] 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 new file mode 100644 index 000000000000..0c5cf40dbd4e --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/q_token_kv_block_sparse_metadata.py @@ -0,0 +1,2199 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""CUDA-graph-safe compact QToken-KvBlock-Sparse-Attention metadata construction. + +The public attention kernel consumes dense sparse-page indices and, for +grouped routes, a parallel packed-membership table. Q1 rows map their compact +selected logical blocks directly to encoded physical +subpage locators. Q2--Q8 rows additionally union the selected blocks of +every adjacent query and write per-page query-membership masks into a separate +dense table, packed four 8-bit masks per Int32 word. The public API names the +sparse block size explicitly: 4, 8, 16, 32, 64, or 128 tokens. +Patterns can be shared across KV heads or supplied independently for each head. + +Construction uses one CUDA C++ CTA per route. Q1 maps selected and tail blocks +directly. Q2--Q8 radix-sort at most ``G * (topk + 1)`` selected/tail IDs, +segmented-OR equal-key memberships, and emit only unique pages; work and +temporary storage are independent of the model context length. On SM90 and +newer, the metadata grid releases the prepared attention grid through +programmatic dependent launch (PDL). +""" + +from __future__ import annotations + +import functools +import math +import threading +from dataclasses import dataclass +from typing import Any, Literal, Optional + +import torch +from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api + +from flashinfer.utils import device_support_pdl +from ._block_sparse.common import _num_sparse_pattern_heads + +_Q_TOKEN_KV_BLOCK_SPARSE_MIN_BLOCK_SIZE = 4 +_Q_TOKEN_KV_BLOCK_SPARSE_BLOCK_SIZES = (4, 8, 16, 32, 64, 128) +_Q_TOKEN_KV_BLOCK_SPARSE_MAX_BLOCK_TOPK = 512 +# Reserve three high key bits for grouped-query membership tags and retain an +# out-of-range radix sentinel in the CUDA/CUB touched-union representation. +_Q_TOKEN_KV_BLOCK_SPARSE_MAX_SEQ_LEN_KV = ( + (1 << 29) - 1 +) * _Q_TOKEN_KV_BLOCK_SPARSE_MIN_BLOCK_SIZE +_Q_TOKEN_KV_BLOCK_SPARSE_MEMBERSHIPS_PER_WORD = 4 +_Q_TOKEN_KV_BLOCK_SPARSE_INT32_LOCATOR_CAPACITY = 1 << 31 +_Q_TOKEN_KV_BLOCK_SPARSE_WORKSPACE_ALIGNMENT = 256 + + +@functools.cache +def _get_prims_ts_q_token_kv_block_sparse_metadata_module() -> Any: + """Build and cache the CUDA C++ QToken-KvBlock-Sparse-Attention metadata implementation.""" + + from flashinfer.jit import gen_prims_ts_q_token_kv_block_sparse_metadata_module + + return gen_prims_ts_q_token_kv_block_sparse_metadata_module().build_and_load() + + +@dataclass(frozen=True) +class _PrimsTSQTokenKvBlockSparseWorkspaceViews: + """Typed views bound to one caller-owned QToken-KvBlock-Sparse-Attention attention workspace.""" + + q_token_kv_block_sparse_page_indices: torch.Tensor + q_token_kv_block_sparse_page_memberships: torch.Tensor + seq_lens: torch.Tensor + attention_workspace_buffer: torch.Tensor + + +@dataclass(frozen=True) +class _QTokenKvBlockSparseTensorDescriptor: + """Structural replacement-storage contract for a prepared QToken-KvBlock-Sparse-Attention plan.""" + + shape: tuple[int, ...] + stride: tuple[int, ...] + device: torch.device + dtype: torch.dtype + + +def _describe_q_token_kv_block_sparse_tensor( + tensor: torch.Tensor, +) -> _QTokenKvBlockSparseTensorDescriptor: + return _QTokenKvBlockSparseTensorDescriptor( + shape=tuple(tensor.shape), + stride=tuple(tensor.stride()), + device=tensor.device, + dtype=tensor.dtype, + ) + + +def _validate_q_token_kv_block_sparse_plan_tensor( + tensor: torch.Tensor, + name: str, + descriptor: _QTokenKvBlockSparseTensorDescriptor, +) -> None: + """Check one replacement tensor without inspecting device values.""" + + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if ( + tuple(tensor.shape) != descriptor.shape + or tuple(tensor.stride()) != descriptor.stride + or tensor.device != descriptor.device + or tensor.dtype != descriptor.dtype + ): + raise ValueError( + f"{name} must preserve the shape, strides, device, and dtype " + "validated by the QToken-KvBlock-Sparse-Attention plan" + ) + + +def _validate_q_token_kv_block_sparse_paged_kv_cache( + paged_kv_cache: object, +) -> tuple[torch.Tensor, torch.Tensor]: + """Validate and return the two rank-four QToken-KvBlock-Sparse-Attention cache tensors.""" + + if ( + not isinstance(paged_kv_cache, tuple) + or len(paged_kv_cache) != 2 + or not all(isinstance(cache, torch.Tensor) for cache in paged_kv_cache) + ): + raise TypeError("paged_kv_cache must be a (k_cache, v_cache) tuple") + k_cache, v_cache = paged_kv_cache + if k_cache.ndim != 4: + raise ValueError( + "K and V cache tensors must have shape [pages,Hkv,storage_page_size,D]" + ) + if ( + v_cache.shape != k_cache.shape + or v_cache.device != k_cache.device + or v_cache.dtype != k_cache.dtype + ): + raise ValueError( + "K and V cache tensors must have matching shapes, devices, and dtypes" + ) + return k_cache, v_cache + + +def _validate_q_token_kv_block_sparse_locator_capacity( + k_cache: torch.Tensor, + *, + sparse_block_size: int, +) -> None: + """Keep every valid encoded cache locator nonnegative in Int32.""" + + storage_page_size = int(k_cache.shape[2]) + _validate_storage_page_size(storage_page_size) + fragment_size = math.gcd(storage_page_size, sparse_block_size) + num_encoded_locators = int(k_cache.shape[0]) * storage_page_size // fragment_size + if num_encoded_locators > _Q_TOKEN_KV_BLOCK_SPARSE_INT32_LOCATOR_CAPACITY: + raise NotImplementedError( + "the encoded QToken-KvBlock-Sparse-Attention cache-locator extent must fit in nonnegative " + "signed int32: got " + f"{num_encoded_locators}, limit is {_Q_TOKEN_KV_BLOCK_SPARSE_INT32_LOCATOR_CAPACITY}" + ) + + +@dataclass(frozen=True) +class _PrimsTSQTokenKvBlockSparseWorkspaceLayout: + """Layout of QToken-KvBlock-Sparse-Attention metadata outputs and disjoint kernel scratch. + + The dense page-index and packed-membership tables plus compact sequence + lengths remain live from metadata construction through the attention + launch. A decode split-KV plan must + preserve its partials and self-resetting completion counters independently + of metadata. A direct prefill plan has no split-KV storage; its small + attention region contains only uniform-call-ABI placeholder tensors. + Semantic inputs such as block tables and query mappings remain outside. + """ + + q_token_kv_block_sparse_page_indices_shape: tuple[int, int] + q_token_kv_block_sparse_page_indices_bytes: int + q_token_kv_block_sparse_page_memberships_shape: tuple[int, int] + q_token_kv_block_sparse_page_memberships_byte_offset: int + q_token_kv_block_sparse_page_memberships_bytes: int + seq_lens_byte_offset: int + seq_lens_bytes: int + attention_workspace_byte_offset: int + attention_scratch_bytes: int + uses_split_kv: bool + max_seq_len: int + total_bytes: int + + def bind( + self, workspace_buffer: torch.Tensor + ) -> _PrimsTSQTokenKvBlockSparseWorkspaceViews: + """Return zero-copy typed views over a validated byte workspace.""" + + _validate_q_token_kv_block_sparse_attention_workspace( + workspace_buffer, self.total_bytes + ) + workspace_bytes = workspace_buffer.reshape(-1).view(torch.uint8) + q_token_kv_block_sparse_page_indices = ( + workspace_bytes[: self.q_token_kv_block_sparse_page_indices_bytes] + .view(torch.int32) + .view(self.q_token_kv_block_sparse_page_indices_shape) + ) + q_token_kv_block_sparse_page_memberships = ( + workspace_bytes[ + self.q_token_kv_block_sparse_page_memberships_byte_offset : ( + self.q_token_kv_block_sparse_page_memberships_byte_offset + + self.q_token_kv_block_sparse_page_memberships_bytes + ) + ] + .view(torch.int32) + .view(self.q_token_kv_block_sparse_page_memberships_shape) + ) + seq_lens = workspace_bytes[ + self.seq_lens_byte_offset : self.seq_lens_byte_offset + self.seq_lens_bytes + ].view(torch.int32) + attention_workspace_buffer = workspace_bytes[ + self.attention_workspace_byte_offset : self.attention_workspace_byte_offset + + self.attention_scratch_bytes + ] + return _PrimsTSQTokenKvBlockSparseWorkspaceViews( + q_token_kv_block_sparse_page_indices=q_token_kv_block_sparse_page_indices, + q_token_kv_block_sparse_page_memberships=q_token_kv_block_sparse_page_memberships, + seq_lens=seq_lens, + attention_workspace_buffer=attention_workspace_buffer, + ) + + +@dataclass(frozen=True) +class _PrimsTSQTokenKvBlockSparseMetadataPlan: + """Unchecked metadata launch state with all geometry pre-resolved.""" + + q_token_kv_block_sparse_page_indices: torch.Tensor + q_token_kv_block_sparse_page_memberships: torch.Tensor + seq_lens: torch.Tensor + qo_indptr: Optional[torch.Tensor] + metadata_run: Any + group_size: int + use_packed_q: bool + sparse_block_size: int + storage_page_size: int + max_seq_len_kv: int + release_attention_pdl: bool + + def run( + self, + block_indices: torch.Tensor, + block_table: torch.Tensor, + token_to_request: torch.Tensor, + query_positions: torch.Tensor, + ) -> None: + """Launch the prepared CUDA metadata kernel without shape resolution.""" + + if self.use_packed_q: + self.metadata_run( + block_indices, + block_table, + token_to_request, + query_positions, + self.qo_indptr, + self.q_token_kv_block_sparse_page_indices, + self.q_token_kv_block_sparse_page_memberships, + self.seq_lens, + self.group_size, + self.storage_page_size, + self.sparse_block_size, + self.max_seq_len_kv, + self.release_attention_pdl, + ) + else: + self.metadata_run( + block_indices, + block_table, + token_to_request, + query_positions, + self.q_token_kv_block_sparse_page_indices, + self.q_token_kv_block_sparse_page_memberships, + self.seq_lens, + self.group_size, + self.storage_page_size, + self.sparse_block_size, + self.max_seq_len_kv, + self.release_attention_pdl, + ) + + +@dataclass(frozen=True) +class _PrimsTSQTokenKvBlockSparsePlan: + """Prepared compact-metadata and PrimTS-attention launch state. + + The plan binds output/workspace storage and freezes input geometry once. + Eager calls may pass new input storage with the same shapes, strides, + devices, and dtypes; CUDA graph replay retains its usual stable-address + requirement. The hot path performs synchronization-free replacement-storage + checks followed by already-prepared launches. G1 resolves logical blocks + directly inside attention, without launching the metadata kernel; grouped + routes first build their union. The + original tensors are validated during preparation. Callers must keep all + replacement tensors alive and disjoint, and must not mutate them + concurrently with a launch or CUDA-graph replay that reads them. Call + :meth:`run` once before capture so both metadata and attention kernels are + compiled and their workspace state is initialized outside the graph. + """ + + _metadata_plan: _PrimsTSQTokenKvBlockSparseMetadataPlan + _bmm1_scale: float + _bmm2_scale: float + _attention_plan: Any + _query: _QTokenKvBlockSparseTensorDescriptor + _block_indices: _QTokenKvBlockSparseTensorDescriptor + _block_table: _QTokenKvBlockSparseTensorDescriptor + _token_to_request: _QTokenKvBlockSparseTensorDescriptor + _query_positions: _QTokenKvBlockSparseTensorDescriptor + _out: _QTokenKvBlockSparseTensorDescriptor + _fixed_query_group_size: Optional[int] = None + + def run( + self, + query: torch.Tensor, + block_indices: torch.Tensor, + block_table: torch.Tensor, + token_to_request: torch.Tensor, + query_positions: torch.Tensor, + *, + out: torch.Tensor, + sm_scale: Optional[float] = None, + v_scale: Optional[float] = None, + ) -> torch.Tensor: + """Launch prepared QToken-KvBlock-Sparse-Attention metadata and attention on the current stream.""" + + _validate_q_token_kv_block_sparse_plan_tensor(query, "query", self._query) + _validate_q_token_kv_block_sparse_plan_tensor( + block_indices, "block_indices", self._block_indices + ) + _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, + ) + _validate_q_token_kv_block_sparse_plan_tensor( + query_positions, + "query_positions", + self._query_positions, + ) + _validate_q_token_kv_block_sparse_plan_tensor(out, "out", self._out) + + from .decode import _validate_16byte_alignment + + _validate_16byte_alignment(query, "query") + _validate_16byte_alignment(out, "out") + + direct_q1_inputs = None + if self._attention_plan._direct_q1_inputs: + direct_q1_inputs = ( + block_indices, + block_table, + token_to_request, + query_positions, + ) + else: + self._metadata_plan.run( + block_indices, + block_table, + token_to_request, + query_positions, + ) + attention_query = _flatten_fixed_q_token_kv_block_sparse_groups( + query, self._fixed_query_group_size + ) + attention_out = _flatten_fixed_q_token_kv_block_sparse_groups( + out, self._fixed_query_group_size + ) + from .decode import _validate_scale + + scale_qk = ( + self._bmm1_scale + if sm_scale is None + else _validate_scale(sm_scale, "sm_scale") + ) + scale_v = ( + self._bmm2_scale if v_scale is None else _validate_scale(v_scale, "v_scale") + ) + self._attention_plan._run_unchecked( + attention_query, + attention_out, + scale_qk, + scale_v, + direct_q1_inputs, + ) + return out + + +@dataclass(frozen=True) +class _QTokenKvBlockSparsePlanConfig: + """Capacity and dtype contract published by the public wrapper.""" + + batch_size: int + seq_len_q: int + num_qo_heads: int + num_kv_heads: int + head_dim: int + kv_block_size: int + page_size: int + block_topk: int + max_seq_len_kv: int + device: torch.device + workspace_buffer: torch.Tensor + use_packed_q: 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 + + +def _q_token_tensor_abi_key( + tensor: Optional[torch.Tensor], + *, + include_storage: bool = False, +) -> object: + """Return the structural ABI, optionally including retained storage.""" + + if tensor is None: + return None + key: tuple[object, ...] = ( + tuple(tensor.shape), + tuple(tensor.stride()), + tensor.device, + tensor.dtype, + ) + return key + (tensor.data_ptr(),) if include_storage else key + + +class QTokenKvBlockSparsePagedTSWrapper: + """Plan QToken-KvBlock-Sparse-Attention capacity and run live routes. + + :meth:`plan` fixes query grouping, K/V capacity, dtypes, and the single + caller-owned workspace. It must run outside CUDA Graph capture. The first + eager :meth:`run` binds the live tensor ABI, compiles the selected PrimTS + kernels, and initializes split-KV scratch. Later runs rebuild compact + routes from the current indexer output and are CUDA-graph capturable. + + One wrapper revision owns mutable route and split-KV workspace. Runs must + therefore be ordered on one stream or externally synchronized; unordered + concurrent runs require distinct wrappers. + + Output and workspace must be disjoint from each other and from all live + inputs, including the original indexer and request metadata. Storage + overlap is not checked; this precondition also applies after rebinding + inputs and across CUDA-graph replays. + """ + + def __init__(self) -> None: + self._config: Optional[_QTokenKvBlockSparsePlanConfig] = None + self._prepared_plan: Optional[_PrimsTSQTokenKvBlockSparsePlan] = None + self._prepared_key: Optional[tuple[object, ...]] = None + self._default_out: Optional[torch.Tensor] = None + self._plan_lock = threading.Lock() + + def plan( + self, + batch_size: int, + seq_len_q: int, + num_qo_heads: int, + num_kv_heads: int, + head_dim: int, + kv_block_size: int, + page_size: int, + block_topk: int, + max_seq_len_kv: int, + *, + device: torch.device | str | int, + workspace_buffer: torch.Tensor, + use_packed_q: bool = False, + split_kv: bool = True, + share_pattern_across_kv_heads: 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, + ) -> None: + """Publish a capacity-only QToken-KvBlock-Sparse-Attention plan. + + ``batch_size`` is the number of query routes: ``len(qo_indptr)-1`` for + packed Q, or ``B * Nq`` for fixed ``[B, Nq, G, Hq, D]`` Q. The planned + ``seq_len_q`` is the maximum packed route length or fixed group size + ``G``. ``kv_block_size`` is the indexer's sparse K/V atom and currently + supports 4/8/16/32/64/128 tokens; ``page_size`` is the physical cache + page. Blocks may cross cache pages. The loader uses fragments of size + gcd(kv_block_size, page_size), resolving every fragment independently. + + ``share_pattern_across_kv_heads=True`` preserves the existing shared + 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. + + ``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 + must use the same value. Change it only by replanning outside capture. + """ + + from .decode import ( + _resolve_cuda_device, + _validate_head_dim, + _validate_head_geometry, + _validate_positive_int, + _validate_runtime_device, + ) + + batch_size = _validate_positive_int(batch_size, "batch_size") + seq_len_q = _validate_positive_int(seq_len_q, "seq_len_q") + block_topk = _validate_positive_int(block_topk, "block_topk") + max_seq_len_kv = _validate_positive_int(max_seq_len_kv, "max_seq_len_kv") + page_size = _validate_positive_int(page_size, "page_size") + head_dim = _validate_head_dim(head_dim) + _validate_head_geometry(num_qo_heads, num_kv_heads) + kv_block_size = _validate_sparse_block_size(kv_block_size) + _validate_storage_page_size(page_size) + _validate_shape_parameters(batch_size * seq_len_q, block_topk, seq_len_q) + if mask_type != "causal": + 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") + for dtype, name in ( + (q_data_type, "q_data_type"), + (kv_data_type, "kv_data_type"), + (o_data_type, "o_data_type"), + ): + 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 + planned_device, device_index = _resolve_cuda_device(device) + _validate_runtime_device(planned_device) + if not isinstance(workspace_buffer, torch.Tensor): + raise TypeError("workspace_buffer must be a torch.Tensor") + if workspace_buffer.device != planned_device: + raise ValueError("workspace_buffer must be on the planned device") + + max_num_storage_pages = (max_seq_len_kv + page_size - 1) // page_size + layout = _get_prims_ts_q_token_kv_block_sparse_workspace_layout( + batch_size * seq_len_q, + block_topk, + max_num_storage_pages, + page_size, + seq_len_q, + max_seq_len_kv=max_seq_len_kv, + 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, + device=planned_device, + num_query_groups=batch_size, + use_packed_q=use_packed_q, + sparse_block_size=kv_block_size, + split_kv=split_kv, + share_pattern_across_kv_heads=share_pattern_across_kv_heads, + ) + _validate_q_token_kv_block_sparse_attention_workspace( + workspace_buffer, layout.total_bytes + ) + with torch.cuda.device(device_index): + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "QToken-KvBlock-Sparse-Attention planning is unsupported during " + "CUDA Graph capture" + ) + candidate = _QTokenKvBlockSparsePlanConfig( + batch_size=batch_size, + seq_len_q=seq_len_q, + num_qo_heads=num_qo_heads, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + kv_block_size=kv_block_size, + page_size=page_size, + block_topk=block_topk, + max_seq_len_kv=max_seq_len_kv, + device=planned_device, + workspace_buffer=workspace_buffer, + use_packed_q=use_packed_q, + 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, + ) + with self._plan_lock: + self._config = candidate + self._prepared_plan = None + self._prepared_key = None + self._default_out = None + + def _require_config(self) -> _QTokenKvBlockSparsePlanConfig: + config = self._config + if config is None: + raise RuntimeError("plan() must be called before run()") + return config + + @flashinfer_experimental_api + def run( + self, + q: torch.Tensor, + 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, + *, + qo_indptr: Optional[torch.Tensor] = None, + sm_scale: Optional[float] = None, + v_scale: Optional[float] = None, + out: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Build live routes and launch the current QToken attention plan. + + Packed Q is ``[total_q, Hq, D]`` with ``qo_indptr``. Fixed Q is + ``[B, Nq, G, Hq, D]`` without ``qo_indptr``. ``indexer_block_ids`` is + Int32 ``[total_q, block_topk]`` for shared patterns or + ``[total_q, Hkv, block_topk]`` for independent patterns, with a + 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)``. + 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. + + The KV-cache writer must zero-fill unused token slots in allocated + pages before attention. TMA loads complete fragments, including the + partial causal tail; score masking does not sanitize NaN/Inf padding + in V. Attention does not clear cache padding itself. + + Run once eagerly after :meth:`plan`; capture only subsequent calls with + stable tensor storage. A changed packed extent or retained K/V/offset + storage triggers a new preparation outside capture. ``sm_scale`` and + ``v_scale`` are live launch scalars, so frameworks may update Q/K and V + dequantization scales without rebuilding the capacity plan. + + Parameters + ---------- + q : torch.Tensor + Packed ``[total_q, Hq, D]`` or fixed ``[B, Nq, G, Hq, D]`` query, + with the planned device, dtype and head geometry. + paged_kv_cache : tuple[torch.Tensor, torch.Tensor] + Separate HND K/V caches shaped + ``[num_pages, Hkv, page_size, D]``. + block_table : torch.Tensor + Dense CUDA Int32 physical-page table + ``[num_requests, max_storage_pages]``. + indexer_block_ids : torch.Tensor + CUDA Int32 logical block IDs ``[T, block_topk]`` when the plan + 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 + CUDA Int32 or Int64 zero-based causal position for each query token. + qo_indptr : torch.Tensor, optional + CUDA Int32 packed-route offsets. Required for packed Q and omitted + for fixed Q; each route must stay within one request and contain + at most the planned ``seq_len_q`` tokens. + sm_scale : float, optional + Softmax scale, defaulting to ``head_dim**-0.5``. Include Q/K + dequantization scales here when using FP8 inputs. + v_scale : float, optional + Value-cache dequantization scale applied to the output, + defaulting to one. + out : torch.Tensor, optional + Caller-owned output with Q's logical shape and the planned output + dtype. When omitted, the wrapper retains and reuses its output. + Must not overlap any input or workspace storage; this is not checked. + + Returns + ------- + torch.Tensor + Attention output in Q's logical layout; returns ``out`` when given. + """ + + 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.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( + qo_indptr, expected_device=config.device + ) + if qo_indptr.numel() != config.batch_size + 1: + raise ValueError("qo_indptr length must match the planned route count") + if q.shape[0] <= 0 or q.shape[0] > config.batch_size * config.seq_len_q: + raise ValueError("packed q exceeds the planned query-token capacity") + else: + if qo_indptr is not None: + raise ValueError("fixed Q does not accept qo_indptr") + if ( + q.ndim != 5 + or q.shape[0] * q.shape[1] != config.batch_size + or q.shape[2] != config.seq_len_q + ): + raise ValueError( + "fixed q must be [B,Nq,G,Hq,D] with B*Nq and G matching plan" + ) + if tuple(q.shape[-2:]) != (config.num_qo_heads, config.head_dim): + raise ValueError("q must use the planned query-head and head dimensions") + num_query_tokens = ( + int(q.shape[0]) + if config.use_packed_q + else (config.batch_size * config.seq_len_q) + ) + expected_ids_shape = ( + (num_query_tokens, config.block_topk) + if config.share_pattern_across_kv_heads + else (num_query_tokens, config.num_kv_heads, config.block_topk) + ) + if ( + not isinstance(indexer_block_ids, torch.Tensor) + or indexer_block_ids.shape != expected_ids_shape + ): + raise ValueError(f"indexer_block_ids must have shape {expected_ids_shape}") + for tensor, name in ( + (token_to_request, "token_to_request"), + (query_positions, "query_positions"), + (block_table, "block_table"), + ): + if not isinstance(tensor, torch.Tensor): + raise TypeError(f"{name} must be a torch.Tensor") + if 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") + + 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 tuple(k_cache.shape[1:]) + != (config.num_kv_heads, config.page_size, config.head_dim) + ): + raise ValueError("paged_kv_cache does not match the planned HND cache ABI") + required_columns = ( + config.max_seq_len_kv + config.page_size - 1 + ) // config.page_size + if ( + block_table.ndim != 2 + or block_table.dtype != torch.int32 + or block_table.device != config.device + or block_table.shape[1] < required_columns + ): + raise ValueError( + "block_table must be CUDA Int32 [requests, pages] with planned capacity" + ) + if out is None: + if self._default_out is None or self._default_out.shape != q.shape: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "run once eagerly before capture when out is not supplied" + ) + self._default_out = torch.empty_like(q, dtype=config.o_data_type) + 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 + ): + raise ValueError("out must match q shape and the planned output dtype") + + prepared_key = ( + _q_token_tensor_abi_key(q), + _q_token_tensor_abi_key(indexer_block_ids), + _q_token_tensor_abi_key(block_table), + _q_token_tensor_abi_key(token_to_request), + _q_token_tensor_abi_key(query_positions), + _q_token_tensor_abi_key(out), + _q_token_tensor_abi_key(k_cache, include_storage=True), + _q_token_tensor_abi_key(v_cache, include_storage=True), + _q_token_tensor_abi_key(qo_indptr, include_storage=True), + ) + if self._prepared_plan is None or prepared_key != self._prepared_key: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "run once eagerly after changing the QToken tensor ABI" + ) + with self._plan_lock: + if self._prepared_plan is None or prepared_key != self._prepared_key: + self._prepared_plan = _prepare_q_token_kv_block_sparse_attention( + q, + paged_kv_cache, + indexer_block_ids, + block_table, + token_to_request, + query_positions, + config.workspace_buffer, + out=out, + max_seq_len_kv=config.max_seq_len_kv, + bmm1_scale=sm_scale, + qo_indptr=qo_indptr, + max_seq_len_q=( + config.seq_len_q if config.use_packed_q else None + ), + sparse_block_size=config.kv_block_size, + split_kv=config.split_kv, + share_pattern_across_kv_heads=config.share_pattern_across_kv_heads, + ) + self._prepared_key = prepared_key + assert self._prepared_plan is not None + return self._prepared_plan.run( + q, + indexer_block_ids, + block_table, + token_to_request, + query_positions, + out=out, + sm_scale=sm_scale, + v_scale=v_scale, + ) + + +def _flatten_fixed_q_token_kv_block_sparse_groups( + tensor: torch.Tensor, + group_size: Optional[int], +) -> torch.Tensor: + """Return the lower-level decode view for a canonical fixed QToken-KvBlock-Sparse-Attention tensor.""" + + if group_size is None: + return tensor + flattened = tensor.flatten(0, 1) + return flattened.squeeze(1) if group_size == 1 else flattened + + +def _get_q_token_kv_block_sparse_metadata_output_shapes( + num_query_tokens: int, + block_topk: int, + group_size: int, + *, + num_query_groups: Optional[int] = None, + sparse_block_size: int = 4, + storage_page_size: Optional[int] = None, + num_pattern_heads: int = 1, +) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]: + """Return ``(q_token_kv_block_sparse_page_indices, q_token_kv_block_sparse_page_memberships, seq_lens)`` shapes. + + Physical fragments have size gcd(sparse_block_size, storage_page_size). + Metadata has one row per query group and pattern head. Grouped masks use + four packed bytes per Int32 word. Q1 does not consume membership metadata + and returns a zero-width membership shape. + + Parameters + ---------- + num_query_tokens : int + Number of flattened query rows represented by the metadata. + block_topk : int + Number of selected logical sparse blocks supplied for each query row. + group_size : int + Maximum number of query rows represented by one QToken-KvBlock-Sparse-Attention route. + num_query_groups : int, optional + Explicit route count for a packed query. If omitted, query rows are + partitioned into fixed groups and ``num_query_tokens`` must be + divisible by ``group_size``. + sparse_block_size : int + Logical sparse-block size in tokens: 4, 8, 16, 32, 64, or 128. + storage_page_size : int, optional + Physical storage-page size; defaults to sparse_block_size. A block + contains sparse_block_size / gcd(sparse_block_size, storage_page_size) + independently addressable fragments. + num_pattern_heads : int + One for shared patterns, otherwise the number of KV heads. + + Returns + ------- + tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]] + Shapes for ``q_token_kv_block_sparse_page_indices``, ``q_token_kv_block_sparse_page_memberships``, and + ``seq_lens``, in that order. + """ + + _validate_sparse_block_size(sparse_block_size) + _validate_shape_parameters(num_query_tokens, block_topk, group_size) + groups = _resolve_num_query_groups( + num_query_tokens, + group_size, + num_query_groups, + ) + if storage_page_size is None: + storage_page_size = sparse_block_size + _validate_storage_page_size(storage_page_size) + fragments_per_block = sparse_block_size // math.gcd( + sparse_block_size, storage_page_size + ) + page_capacity = group_size * (block_topk + 1) * fragments_per_block + groups *= num_pattern_heads + membership_words = ( + 0 + if group_size == 1 + else (page_capacity + _Q_TOKEN_KV_BLOCK_SPARSE_MEMBERSHIPS_PER_WORD - 1) + // _Q_TOKEN_KV_BLOCK_SPARSE_MEMBERSHIPS_PER_WORD + ) + return (groups, page_capacity), (groups, membership_words), (groups,) + + +def _get_prims_ts_q_token_kv_block_sparse_workspace_layout( + num_query_tokens: int, + block_topk: int, + max_num_storage_pages: int, + storage_page_size: int, + group_size: int, + *, + max_seq_len_kv: int, + num_qo_heads: int, + num_kv_heads: int, + head_dim: int, + q_dtype: torch.dtype = torch.float16, + kv_dtype: Optional[torch.dtype] = None, + out_dtype: Optional[torch.dtype] = None, + device: Optional[torch.device | str | int] = None, + num_query_groups: Optional[int] = None, + use_packed_q: bool = False, + sparse_block_size: int = 4, + split_kv: bool = True, + share_pattern_across_kv_heads: bool = True, +) -> _PrimsTSQTokenKvBlockSparseWorkspaceLayout: + """Return the unified allocation layout for QToken-KvBlock-Sparse-Attention metadata and attention. + + The workspace owns the dense QToken-KvBlock-Sparse-Attention page-index table, grouped membership + words, compact sequence lengths, and kernel scratch. The caller continues + to provide the model block table, request mappings, and query positions as + explicit semantic inputs. + + Metadata and attention use disjoint regions. The caller controls whether + split-KV may be selected, independently of query layout. Split launches + own dedicated partial-output, statistics and completion-counter storage. + """ + + sparse_block_size = _validate_sparse_block_size(sparse_block_size) + _validate_shape_parameters(num_query_tokens, block_topk, group_size) + if num_query_tokens == 0: + raise ValueError( + "num_query_tokens must be positive for QToken-KvBlock-Sparse-Attention attention" + ) + if not isinstance(split_kv, bool): + raise TypeError("split_kv must be a bool") + if not isinstance(max_num_storage_pages, int) or isinstance( + max_num_storage_pages, bool + ): + raise TypeError("max_num_storage_pages must be an integer") + if max_num_storage_pages <= 0: + raise ValueError("max_num_storage_pages must be positive") + _validate_storage_page_size(storage_page_size) + max_seq_len_kv = _validate_q_token_kv_block_sparse_max_seq_len_kv( + max_seq_len_kv, + block_table_token_capacity=max_num_storage_pages * storage_page_size, + ) + + from .decode import ( + _resolve_decode_workspace_layout, + _validate_prims_ts_q_token_kv_block_sparse_group_capacity, + ) + + group_size = _validate_prims_ts_q_token_kv_block_sparse_group_capacity( + group_size, + num_qo_heads, + num_kv_heads, + ) + + groups = _resolve_num_query_groups( + num_query_tokens, + group_size, + num_query_groups, + ) + 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( + num_kv_heads, share_pattern_across_kv_heads + ) + metadata_rows = groups * pattern_heads + fragment_size = math.gcd(sparse_block_size, storage_page_size) + block_capacity = min( + group_size * (block_topk + 1), + (max_seq_len_kv + sparse_block_size - 1) // sparse_block_size, + ) + page_capacity = block_capacity * (sparse_block_size // fragment_size) + membership_words = ( + 0 + if group_size == 1 + else (page_capacity + _Q_TOKEN_KV_BLOCK_SPARSE_MEMBERSHIPS_PER_WORD - 1) + // _Q_TOKEN_KV_BLOCK_SPARSE_MEMBERSHIPS_PER_WORD + ) + q_token_kv_block_sparse_page_indices_shape = (metadata_rows, page_capacity) + q_token_kv_block_sparse_page_indices_bytes = metadata_rows * page_capacity * 4 + q_token_kv_block_sparse_page_memberships_shape = (metadata_rows, membership_words) + q_token_kv_block_sparse_page_memberships_byte_offset = ( + _align_up_q_token_kv_block_sparse_workspace( + q_token_kv_block_sparse_page_indices_bytes + ) + ) + q_token_kv_block_sparse_page_memberships_bytes = ( + metadata_rows * membership_words * 4 + ) + seq_lens_byte_offset = _align_up_q_token_kv_block_sparse_workspace( + q_token_kv_block_sparse_page_memberships_byte_offset + + q_token_kv_block_sparse_page_memberships_bytes + ) + seq_lens_numel = metadata_rows + seq_lens_bytes = seq_lens_numel * 4 + max_seq_len = ( + block_topk * sparse_block_size + (sparse_block_size - 1) + if group_size == 1 + else page_capacity * fragment_size + ) + max_seq_len = min(max_seq_len, max_seq_len_kv) + if kv_dtype is None: + kv_dtype = q_dtype + if out_dtype is None: + out_dtype = q_dtype + attention_layout = _resolve_decode_workspace_layout( + groups, + num_qo_heads, + num_kv_heads, + head_dim, + fragment_size, + max_seq_len, + group_size, + q_dtype, + kv_dtype, + kv_dtype, + out_dtype, + "HND", + "causal", + use_packed_q, + -1, + storage_page_size, + device, + use_q_token_kv_block_sparse_route=True, + use_pdl=True, + split_kv=split_kv, + share_pattern_across_kv_heads=share_pattern_across_kv_heads, + ) + attention_scratch_bytes = attention_layout.total_bytes + attention_workspace_byte_offset = _align_up_q_token_kv_block_sparse_workspace( + seq_lens_byte_offset + seq_lens_bytes + ) + return _PrimsTSQTokenKvBlockSparseWorkspaceLayout( + q_token_kv_block_sparse_page_indices_shape=q_token_kv_block_sparse_page_indices_shape, + q_token_kv_block_sparse_page_indices_bytes=q_token_kv_block_sparse_page_indices_bytes, + q_token_kv_block_sparse_page_memberships_shape=q_token_kv_block_sparse_page_memberships_shape, + q_token_kv_block_sparse_page_memberships_byte_offset=q_token_kv_block_sparse_page_memberships_byte_offset, + q_token_kv_block_sparse_page_memberships_bytes=q_token_kv_block_sparse_page_memberships_bytes, + seq_lens_byte_offset=seq_lens_byte_offset, + seq_lens_bytes=seq_lens_bytes, + attention_workspace_byte_offset=attention_workspace_byte_offset, + attention_scratch_bytes=attention_scratch_bytes, + uses_split_kv=attention_layout.uses_split_kv, + max_seq_len=max_seq_len, + total_bytes=attention_workspace_byte_offset + attention_scratch_bytes, + ) + + +@flashinfer_experimental_api +def get_q_token_kv_block_sparse_workspace_size( + q: torch.Tensor, + k_cache: torch.Tensor, + block_table: torch.Tensor, + *, + block_topk: int, + max_seq_len_kv: int, + o_data_type: Optional[torch.dtype] = None, + qo_indptr: Optional[torch.Tensor] = None, + seq_len_q: Optional[int] = None, + kv_block_size: int = 4, + split_kv: bool = True, + share_pattern_across_kv_heads: bool = True, +) -> int: + """Return bytes for QToken-KvBlock-Sparse-Attention workspace storage. + + Packed queries use ``[num_query_tokens, num_qo_heads, head_dim]`` and + Int32 ``qo_indptr`` partitions their rows into request-safe routes of + length at most ``seq_len_q``. Fixed queries use + ``[batch, num_query_groups, group_size, num_qo_heads, head_dim]`` and do + not require query-offset metadata. The first two fixed axes are flattened + into the attention route axis without copying tensor storage. + + The dense QToken-KvBlock-Sparse-Attention page indices, grouped membership words, and compact sequence + lengths are owned by the workspace. These persistent metadata outputs and + the attention scratch occupy disjoint regions, and callers do not allocate + any metadata output separately. + The attention region contains split-KV partials, statistics, and counters + only when the resolved policy uses split-KV; direct prefill retains only + small uniform-ABI placeholders. + ``kv_block_size`` supports 4/8/16/32/64/128 tokens. Shared patterns use + one metadata row per query group; independent patterns use one per group + and KV head. Storage-page crossings increase the number of physical + fragments, not the number of logical candidates sorted. + + Parameters + ---------- + q : torch.Tensor + Packed ``[total_q, Hq, D]`` query when ``qo_indptr`` is supplied, or + contiguous fixed ``[B, Nq, G, Hq, D]`` query otherwise. + k_cache : torch.Tensor + HND key cache shaped ``[num_pages, Hkv, storage_page_size, D]``. The + matching value cache does not affect workspace sizing. + block_table : torch.Tensor + Dense CUDA Int32 physical-page table shaped + ``[num_requests, max_storage_pages]``. + block_topk : int + Number of selected logical sparse blocks supplied for each query row. + max_seq_len_kv : int + 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 + Output dtype used to size attention partials. It defaults to the query + dtype. + qo_indptr : torch.Tensor, optional + Contiguous CUDA Int32 cumulative offsets for packed query routes. + seq_len_q : int, optional + Maximum packed-route length and QToken-KvBlock-Sparse-Attention group size. It is required when + ``qo_indptr`` is supplied. + kv_block_size : int + Logical sparse-block size in tokens: 4, 8, 16, 32, 64, or 128. + split_kv : bool + Allow automatic split-KV selection (default True), independently of + query layout. Use False for prefill and True for decode. Pass the same + value to attention planning; False omits split-KV scratch. + + share_pattern_across_kv_heads : bool + 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. + + Returns + ------- + int + Required unified workspace size in bytes, including persistent QToken-KvBlock-Sparse-Attention + metadata outputs and a disjoint attention-scratch region. + """ + + return _get_prims_ts_q_token_kv_block_sparse_workspace_layout_from_tensors( + q, + k_cache, + block_table, + block_topk, + max_seq_len_kv=max_seq_len_kv, + out_dtype=o_data_type, + qo_indptr=qo_indptr, + max_seq_len_q=seq_len_q, + sparse_block_size=kv_block_size, + split_kv=split_kv, + share_pattern_across_kv_heads=share_pattern_across_kv_heads, + ).total_bytes + + +def _get_prims_ts_q_token_kv_block_sparse_workspace_layout_from_tensors( + query: torch.Tensor, + k_cache: torch.Tensor, + block_table: torch.Tensor, + block_topk: int, + *, + max_seq_len_kv: int, + out_dtype: Optional[torch.dtype], + qo_indptr: Optional[torch.Tensor] = None, + max_seq_len_q: Optional[int] = None, + sparse_block_size: int = 4, + split_kv: bool = True, + share_pattern_across_kv_heads: bool = True, +) -> _PrimsTSQTokenKvBlockSparseWorkspaceLayout: + sparse_block_size = _validate_sparse_block_size(sparse_block_size) + if not isinstance(query, torch.Tensor): + raise TypeError("query must be a torch.Tensor") + if not query.is_cuda: + raise ValueError("query must be a CUDA tensor") + if not isinstance(k_cache, torch.Tensor): + 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 not use_packed_q and max_seq_len_q is not None: + raise ValueError( + "max_seq_len_q is only valid with packed QToken-KvBlock-Sparse-Attention qo_indptr" + ) + if use_packed_q: + if query.ndim != 3: + raise ValueError( + "packed QToken-KvBlock-Sparse-Attention query must have shape [total_q,Hq,D]" + ) + if max_seq_len_q is None: + raise ValueError( + "max_seq_len_q is required with QToken-KvBlock-Sparse-Attention qo_indptr" + ) + _validate_q_token_kv_block_sparse_qo_indptr_layout_tensor(qo_indptr) + group_size = int(max_seq_len_q) + groups = int(qo_indptr.numel()) - 1 + num_query_tokens, num_qo_heads, head_dim = query.shape + elif query.ndim == 5: + batch_size, groups_per_request, group_size, num_qo_heads, head_dim = query.shape + if batch_size <= 0 or groups_per_request <= 0: + raise ValueError( + "fixed QToken-KvBlock-Sparse-Attention batch and query-group counts must be positive" + ) + groups = int(batch_size) * int(groups_per_request) + else: + raise ValueError( + "query must be packed [total_q,Hq,D] with qo_indptr or fixed " + "[B,Nq,G,Hq,D] without qo_indptr" + ) + if k_cache.ndim != 4: + raise ValueError("k_cache must have shape [pages, Hkv, storage_page_size, D]") + if k_cache.device != query.device or k_cache.shape[3] != head_dim: + raise ValueError("query and k_cache must share device and head dimension") + if block_table.ndim != 2 or block_table.dtype != torch.int32: + raise ValueError("block_table must be a rank-two int32 tensor") + if block_table.device != query.device: + raise ValueError("block_table must be on the query device") + _validate_q_token_kv_block_sparse_locator_capacity( + k_cache, + sparse_block_size=sparse_block_size, + ) + 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, + block_topk, + block_table.shape[1], + k_cache.shape[2], + group_size, + max_seq_len_kv=max_seq_len_kv, + num_qo_heads=num_qo_heads, + num_kv_heads=k_cache.shape[1], + head_dim=head_dim, + q_dtype=query.dtype, + kv_dtype=k_cache.dtype, + out_dtype=out_dtype, + device=query.device, + num_query_groups=groups, + use_packed_q=use_packed_q, + sparse_block_size=sparse_block_size, + split_kv=split_kv, + share_pattern_across_kv_heads=share_pattern_across_kv_heads, + ) + + +def _build_prims_ts_q_token_kv_block_sparse_metadata( + block_indices: torch.Tensor, + block_table: torch.Tensor, + token_to_request: torch.Tensor, + query_positions: torch.Tensor, + *, + group_size: int, + storage_page_size: int, + max_seq_len_kv: int, + sparse_block_size: int = 4, + qo_indptr: Optional[torch.Tensor] = None, + out: Optional[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = None, + release_attention_pdl: bool = False, + share_pattern_across_kv_heads: bool = True, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build dense QToken-KvBlock-Sparse-Attention indices and memberships from compact sparse-block IDs. + + ``block_indices`` contains logical sparse-block IDs with shape + ``[num_query_tokens, block_topk]`` when shared, or ``[num_query_tokens, + Hkv, block_topk]`` for per-head selections. Adjacent Q2--Q8 rows must belong to + the same request and have consecutive positions. Q1 output locators are + ordinary encoded subpage locators. Grouped routes return the same plain + locators plus a separate Int32 table containing four 8-bit membership masks + per word; membership bit ``i`` marks visibility for query ``i`` in the + group. Valid locators must fit in nonnegative signed Int32. The combined + attention APIs prove this from the cache capacity, while callers of this + advanced metadata-only API own that input-value contract. + + Pass preallocated ``out`` tensors for CUDA graph capture. The page-index + table has a fixed row width of + ``group_size * (block_topk + 1) * fragments_per_block``; ``seq_lens`` + selects the live token prefix. Fragments divide both logical blocks and + storage pages, so physical page IDs never imply logical adjacency. + The membership output has one Int32 word per four page slots, or zero + columns for Q1. + Q1 uses one CUDA C++ CTA per route to map its selected blocks and causal + tail directly. Q2--Q8 use one CUDA C++ CTA per route to radix-sort the + bounded ``group_size * (block_topk + 1)`` candidates, unique them while + OR-reducing membership bits, and map the resulting logical pages through + the dense block table. The terminal metadata grid releases a following + PDL-capable QToken-KvBlock-Sparse-Attention attention launch when requested by the combined API. + Packed ``qo_indptr`` must be an Int32 device copy of CPU-validated route + offsets; this builder checks only its structural tensor contract and does + not read route values back to the host. + """ + + _validate_inputs( + block_indices, + block_table, + token_to_request, + query_positions, + group_size, + storage_page_size, + max_seq_len_kv, + sparse_block_size, + share_pattern_across_kv_heads, + ) + 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: + _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 + expected_shapes = _get_q_token_kv_block_sparse_metadata_output_shapes( + rows, + block_topk, + group_size, + num_query_groups=groups, + sparse_block_size=sparse_block_size, + storage_page_size=storage_page_size, + num_pattern_heads=pattern_heads, + ) + if out is None: + outputs = tuple( + torch.empty(shape, dtype=torch.int32, device=block_indices.device) + for shape in expected_shapes + ) + else: + if not isinstance(out, tuple) or len(out) != 3: + raise TypeError( + "out must be a (q_token_kv_block_sparse_page_indices, q_token_kv_block_sparse_page_memberships, seq_lens) tuple" + ) + outputs = out + for tensor, shape in zip(outputs, expected_shapes, strict=True): + if ( + tensor.shape != shape + or tensor.dtype != torch.int32 + or tensor.device != block_indices.device + or not tensor.is_contiguous() + ): + raise ValueError( + "QToken-KvBlock-Sparse-Attention metadata outputs must be contiguous int32 tensors with " + f"shapes {expected_shapes} on {block_indices.device}" + ) + ( + q_token_kv_block_sparse_page_indices, + q_token_kv_block_sparse_page_memberships, + seq_lens, + ) = outputs + groups = _resolve_num_query_groups(rows, group_size, groups) + if groups == 0: + return outputs + + 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: + metadata_module.run_packed( + block_indices, + block_table, + token_to_request, + query_positions, + qo_indptr, + q_token_kv_block_sparse_page_indices, + q_token_kv_block_sparse_page_memberships, + seq_lens, + group_size, + storage_page_size, + sparse_block_size, + max_seq_len_kv, + release_attention_pdl, + ) + else: + metadata_module.run_fixed( + block_indices, + block_table, + token_to_request, + query_positions, + q_token_kv_block_sparse_page_indices, + q_token_kv_block_sparse_page_memberships, + seq_lens, + group_size, + storage_page_size, + sparse_block_size, + max_seq_len_kv, + release_attention_pdl, + ) + return outputs + + +def _build_q_token_kv_block_sparse_metadata( + block_indices: torch.Tensor, + block_table: torch.Tensor, + token_to_request: torch.Tensor, + query_positions: torch.Tensor, + *, + group_size: int, + storage_page_size: int, + max_seq_len_kv: int, + sparse_block_size: int = 4, + qo_indptr: Optional[torch.Tensor] = None, + out: Optional[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = None, + share_pattern_across_kv_heads: bool = True, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Build dense QToken-KvBlock-Sparse-Attention page indices and memberships from sparse-block IDs. + + ``block_indices`` is ``[T,K]`` when shared, otherwise ``[T,Hkv,K]``. + ``block_table`` remains the dense ``[requests, max_storage_pages]`` physical + table. Let F = gcd(sparse_block_size, storage_page_size). Output rows number + groups * pattern_heads, with width G * (K+1) * (sparse_block_size/F). + ``seq_lens`` gives each row's live token prefix. Memberships pack four + fragment bytes per Int32; G1 has zero membership columns. Every locator is + physical_page * (storage_page_size/F) + subpage. No CSR conversion is needed. + + Q1 uses one CUDA C++ direct-mapping kernel. Q2--Q8 use one CUDA C++ + radix-sort and union kernel per route. Neither path requires + caller-provided scratch. Advanced callers that capture this raw path must + preallocate ``out``, warm the same metadata geometry once before capture, + and retain every tensor at a stable address through replay. The three + ``out`` storage ranges must be mutually disjoint and must not alias inputs. + Packed ``qo_indptr`` must be an Int32 device copy of CPU-validated route + offsets and is checked structurally without a host-side value read. + + ``sparse_block_size`` defaults to four. The API validates power-of-two + values in {4,8,16,32,64,128}. Logical blocks may cross physical cache + pages; metadata emits only the physical fragments the loader needs. + + Parameters + ---------- + block_indices : torch.Tensor + CUDA Int32 selected logical sparse-block IDs shaped ``[T,K]`` or + ``[T,Hkv,K]``, with a contiguous K dimension. + block_table : torch.Tensor + Dense CUDA Int32 physical-page table shaped + ``[num_requests, max_storage_pages]`` with contiguous rows. + token_to_request : torch.Tensor + Contiguous CUDA Int32 request index for each flattened query row, + shaped ``[num_query_tokens]``. + query_positions : torch.Tensor + Contiguous CUDA Int32 or Int64 absolute position for each flattened + query row, shaped ``[num_query_tokens]``. + group_size : int + Maximum number of query rows represented by one route. + storage_page_size : int + Number of tokens in each physical K/V storage page. + max_seq_len_kv : int + Static maximum visible logical K/V length in tokens, including current + query or MTP tokens. It must fit within each dense block-table row. + sparse_block_size : int + Logical sparse-block size in tokens: 4, 8, 16, 32, 64, or 128. + qo_indptr : torch.Tensor, optional + Contiguous CUDA Int32 cumulative route offsets for packed queries. If + omitted, adjacent rows form fixed groups of exactly ``group_size``. + out : tuple[torch.Tensor, torch.Tensor, torch.Tensor], optional + Preallocated contiguous CUDA Int32 tensors for page indices, packed + memberships, and live sequence lengths. Their shapes must match + :func:`_get_q_token_kv_block_sparse_metadata_output_shapes`. If omitted, the three + outputs are allocated internally. + + share_pattern_across_kv_heads : bool + 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. + + Returns + ------- + tuple[torch.Tensor, torch.Tensor, torch.Tensor] + ``(q_token_kv_block_sparse_page_indices, q_token_kv_block_sparse_page_memberships, seq_lens)``. Memberships + pack four per-page bytes into each Int32 word and have zero columns for + Q1. + """ + + return _build_prims_ts_q_token_kv_block_sparse_metadata( + block_indices, + block_table, + token_to_request, + query_positions, + group_size=group_size, + storage_page_size=storage_page_size, + max_seq_len_kv=max_seq_len_kv, + sparse_block_size=sparse_block_size, + qo_indptr=qo_indptr, + out=out, + share_pattern_across_kv_heads=share_pattern_across_kv_heads, + ) + + +def _prepare_prims_ts_q_token_kv_block_sparse_metadata_plan( + block_indices: torch.Tensor, + block_table: torch.Tensor, + token_to_request: torch.Tensor, + query_positions: torch.Tensor, + q_token_kv_block_sparse_page_indices: torch.Tensor, + q_token_kv_block_sparse_page_memberships: torch.Tensor, + seq_lens: torch.Tensor, + *, + group_size: int, + storage_page_size: int, + max_seq_len_kv: int, + sparse_block_size: int = 4, + qo_indptr: Optional[torch.Tensor] = None, +) -> _PrimsTSQTokenKvBlockSparseMetadataPlan: + """Freeze validated metadata tensors and launch constants.""" + + sparse_block_size = _validate_sparse_block_size(sparse_block_size) + rows = block_indices.shape[0] + use_packed_q = qo_indptr is not None + _resolve_num_query_groups( + rows, + group_size, + int(qo_indptr.numel()) - 1 if use_packed_q else None, + ) + max_seq_len_kv = _validate_q_token_kv_block_sparse_max_seq_len_kv( + max_seq_len_kv, + block_table_token_capacity=block_table.shape[1] * storage_page_size, + ) + release_attention_pdl = device_support_pdl(block_indices.device) + metadata_module = _get_prims_ts_q_token_kv_block_sparse_metadata_module() + metadata_run = ( + metadata_module.run_packed if use_packed_q else metadata_module.run_fixed + ) + return _PrimsTSQTokenKvBlockSparseMetadataPlan( + q_token_kv_block_sparse_page_indices=q_token_kv_block_sparse_page_indices, + q_token_kv_block_sparse_page_memberships=q_token_kv_block_sparse_page_memberships, + seq_lens=seq_lens, + qo_indptr=qo_indptr, + metadata_run=metadata_run, + group_size=group_size, + use_packed_q=use_packed_q, + 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, + ) + + +def _prepare_q_token_kv_block_sparse_attention( + query: torch.Tensor, + paged_kv_cache: tuple[torch.Tensor, torch.Tensor], + block_indices: torch.Tensor, + block_table: torch.Tensor, + token_to_request: torch.Tensor, + query_positions: torch.Tensor, + workspace_buffer: torch.Tensor, + *, + out: torch.Tensor, + max_seq_len_kv: int, + bmm1_scale: Optional[float] = None, + bmm2_scale: float = 1.0, + qo_indptr: Optional[torch.Tensor] = None, + max_seq_len_q: Optional[int] = None, + sparse_block_size: int = 4, + split_kv: bool = True, + share_pattern_across_kv_heads: bool = True, +) -> _PrimsTSQTokenKvBlockSparsePlan: + """Prepare one graph-stable metadata-plus-attention QToken-KvBlock-Sparse-Attention launch. + + Packed Q/output use ``[num_query_tokens, Hq, D]`` and describe Q1--Q8 + routes with ``qo_indptr`` and ``max_seq_len_q``. A request's final route + may be shorter than the maximum group size. Fixed Q/output use + ``[B, Nq, G, Hq, D]`` and omit ``qo_indptr``; ``B * Nq`` becomes the + internal attention route count through a zero-copy view. + + Frameworks call :meth:`_PrimsTSQTokenKvBlockSparsePlan.run` with current inputs that preserve + the prepared geometry. The plan owns its dense page-index table, grouped + membership words, and compact sequence lengths inside ``workspace_buffer`` + and resolves metadata geometry and the PrimTS attention launch only once. + Packed ``qo_indptr`` must be an + Int32 device copy of CPU-validated route offsets, such as those from + :func:`make_q_token_kv_block_sparse_qo_indptr`. Preparation checks only the + structural tensor contract and does not materialize values on the host. + + Call ``run`` once outside CUDA graph capture to compile and initialize the + plan, then capture with all semantic inputs, output, and workspace storage + kept at stable addresses. + + G1 reads the raw indexer inputs directly inside attention and does not + materialize the workspace's page-index or sequence-length views. The + workspace size remains a conservative capacity, shared with grouped routes. + + ``sparse_block_size`` defaults to four and selects the semantic block + specialization in {4, 8, 16, 32, 64, 128}. + + Parameters + ---------- + query : torch.Tensor + Packed ``[total_q, Hq, D]`` query when ``qo_indptr`` is supplied, or + contiguous fixed ``[B, Nq, G, Hq, D]`` query otherwise. + paged_kv_cache : tuple[torch.Tensor, torch.Tensor] + Separate HND key and value caches, each shaped + ``[num_pages, Hkv, storage_page_size, D]`` with matching shape, dtype, + and device. + block_indices : torch.Tensor + CUDA Int32 logical sparse-block IDs shaped ``[T, block_topk]`` when + shared, otherwise ``[T, Hkv, block_topk]``, with a contiguous last dimension. + block_table : torch.Tensor + Dense CUDA Int32 physical-page table shaped + ``[num_requests, max_storage_pages]``. + token_to_request : torch.Tensor + Contiguous CUDA Int32 request index for each flattened query row. + query_positions : torch.Tensor + Contiguous CUDA Int32 or Int64 absolute position for each flattened + query row. + workspace_buffer : torch.Tensor + Caller-owned contiguous CUDA byte workspace with at least the size + returned by :func:`get_q_token_kv_block_sparse_workspace_size`. Its storage must + remain stable and disjoint from all inputs and ``out``. + out : torch.Tensor + Caller-owned output with the same logical shape as ``query``. + max_seq_len_kv : int + Static maximum visible logical K/V length in tokens, including current + query or MTP tokens. It must fit within each dense block-table row and + remain unchanged across CUDA graph replay. + bmm1_scale : float, optional + Scale applied to QK scores. It defaults to ``head_dim**-0.5``. + bmm2_scale : float + Scale applied to the attention output. + qo_indptr : torch.Tensor, optional + Contiguous CUDA Int32 cumulative route offsets selecting packed-query + mode. + max_seq_len_q : int, optional + Maximum packed-route length and QToken-KvBlock-Sparse-Attention group size. It is required when + ``qo_indptr`` is supplied. + sparse_block_size : int + Logical sparse-block size in tokens: 4, 8, 16, 32, 64, or 128. + + share_pattern_across_kv_heads : bool + 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. + + Returns + ------- + _PrimsTSQTokenKvBlockSparsePlan + Reusable prepared metadata-plus-attention plan. Call + :meth:`_PrimsTSQTokenKvBlockSparsePlan.run` once before CUDA graph capture. + """ + + sparse_block_size = _validate_sparse_block_size(sparse_block_size) + k_cache, v_cache = _validate_q_token_kv_block_sparse_paged_kv_cache(paged_kv_cache) + if not isinstance(query, torch.Tensor): + raise TypeError("query must be a torch.Tensor") + if not isinstance(out, torch.Tensor): + raise TypeError("out must be a caller-owned torch.Tensor") + + use_packed_q = qo_indptr is not None + fixed_query_group_size = None + if use_packed_q: + if query.ndim != 3: + raise ValueError( + "packed QToken-KvBlock-Sparse-Attention query must have shape [total_q,Hq,D]" + ) + if max_seq_len_q is None: + raise ValueError( + "max_seq_len_q is required with QToken-KvBlock-Sparse-Attention qo_indptr" + ) + group_size = int(max_seq_len_q) + num_query_tokens = int(query.shape[0]) + _validate_q_token_kv_block_sparse_qo_indptr_tensor( + qo_indptr, + expected_device=query.device, + ) + attention_query = query + attention_out = out + elif query.ndim == 5: + if out.shape != query.shape: + raise ValueError( + "fixed QToken-KvBlock-Sparse-Attention output must have the same shape as query" + ) + if not query.is_contiguous() or not out.is_contiguous(): + raise ValueError( + "fixed QToken-KvBlock-Sparse-Attention query and output must be contiguous" + ) + group_size = int(query.shape[2]) + num_query_tokens = int(query.shape[0] * query.shape[1]) * group_size + attention_query = _flatten_fixed_q_token_kv_block_sparse_groups( + query, group_size + ) + attention_out = _flatten_fixed_q_token_kv_block_sparse_groups(out, group_size) + fixed_query_group_size = group_size + else: + raise ValueError( + "query must be packed [total_q,Hq,D] with qo_indptr or fixed " + "[B,Nq,G,Hq,D] without qo_indptr" + ) + + _validate_inputs( + block_indices, + block_table, + token_to_request, + query_positions, + group_size, + int(k_cache.shape[2]), + max_seq_len_kv, + sparse_block_size, + share_pattern_across_kv_heads, + ) + 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") + if block_indices.shape[0] != num_query_tokens: + raise ValueError("block_indices must have one row per flattened query token") + + layout = _get_prims_ts_q_token_kv_block_sparse_workspace_layout_from_tensors( + query, + k_cache, + block_table, + int(block_indices.shape[-1]), + max_seq_len_kv=max_seq_len_kv, + out_dtype=out.dtype, + qo_indptr=qo_indptr, + max_seq_len_q=max_seq_len_q, + sparse_block_size=sparse_block_size, + split_kv=split_kv, + share_pattern_across_kv_heads=share_pattern_across_kv_heads, + ) + views = layout.bind(workspace_buffer) + metadata_plan = _prepare_prims_ts_q_token_kv_block_sparse_metadata_plan( + block_indices, + block_table, + token_to_request, + query_positions, + views.q_token_kv_block_sparse_page_indices, + views.q_token_kv_block_sparse_page_memberships, + views.seq_lens, + group_size=group_size, + storage_page_size=int(k_cache.shape[2]), + max_seq_len_kv=max_seq_len_kv, + sparse_block_size=sparse_block_size, + qo_indptr=qo_indptr, + ) + from .decode import _prepare_prims_ts_batch_decode_plan, _validate_scale + + scale_qk = _validate_scale( + query.shape[-1] ** -0.5 if bmm1_scale is None else bmm1_scale, + "bmm1_scale", + ) + scale_v = _validate_scale(bmm2_scale, "bmm2_scale") + attention_plan, _ = _prepare_prims_ts_batch_decode_plan( + attention_query, + paged_kv_cache, + views.attention_workspace_buffer, + views.q_token_kv_block_sparse_page_indices, + views.seq_lens, + layout.max_seq_len, + out=attention_out, + seq_len_q=group_size, + qo_indptr=qo_indptr, + max_seq_len_q=group_size if use_packed_q else None, + out_dtype=out.dtype, + mask_type="causal", + window_left=-1, + kv_layout="HND", + page_size=math.gcd(sparse_block_size, int(k_cache.shape[2])), + direct_q1_sparse_block_size=sparse_block_size, + use_q_token_kv_block_sparse_route=True, + use_pdl=True, + split_kv=split_kv, + share_pattern_across_kv_heads=share_pattern_across_kv_heads, + q_token_kv_block_sparse_page_memberships=( + views.q_token_kv_block_sparse_page_memberships if group_size > 1 else None + ), + direct_q1_inputs=( + (block_indices, block_table, token_to_request, query_positions) + if group_size == 1 + else () + ), + direct_q1_max_model_len=max_seq_len_kv if group_size == 1 else None, + ) + return _PrimsTSQTokenKvBlockSparsePlan( + _metadata_plan=metadata_plan, + _bmm1_scale=scale_qk, + _bmm2_scale=scale_v, + _attention_plan=attention_plan, + _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), + _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, + ) + + +@flashinfer_experimental_api +def q_token_kv_block_sparse_attention_with_paged_kv_cache( + q: torch.Tensor, + 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, + workspace_buffer: torch.Tensor, + *, + max_seq_len_kv: int, + seq_len_q: Optional[int] = None, + kv_block_size: int = 4, + mask_type: Literal["causal"] = "causal", + sm_scale: Optional[float] = None, + v_scale: Optional[float] = None, + out: Optional[torch.Tensor] = None, + o_data_type: Optional[torch.dtype] = None, + qo_indptr: Optional[torch.Tensor] = None, + split_kv: bool = True, + share_pattern_across_kv_heads: bool = True, +) -> torch.Tensor: + """Plan and run one QToken-KvBlock-Sparse-Attention launch. + + This eager convenience API creates a temporary + :class:`QTokenKvBlockSparsePagedTSWrapper`, plans capacity, and launches + once. It is not CUDA-graph capturable. Repeated launches should retain a + wrapper, call :meth:`QTokenKvBlockSparsePagedTSWrapper.plan` outside + capture, run once eagerly, and capture only :meth:`~QTokenKvBlockSparsePagedTSWrapper.run`. + + Packed Q is ``[total_q, Hq, D]`` with ``qo_indptr``; ``seq_len_q`` + is its maximum route length. Fixed Q is ``[B, Nq, G, Hq, D]`` without + ``qo_indptr``, and ``G`` determines ``seq_len_q``. The indexer + output is named ``indexer_block_ids`` to distinguish it from BSR + ``block_indices``. ``kv_block_size`` is the semantic indexer atom + (4/8/16/32/64/128), while the physical cache extent is inferred from + ``paged_kv_cache``. + + Parameters + ---------- + q : torch.Tensor + Packed ``[total_q, Hq, D]`` or fixed ``[B, Nq, G, Hq, D]`` query. + paged_kv_cache : tuple[torch.Tensor, torch.Tensor] + Separate HND K/V cache tensors shaped + ``[num_pages, Hkv, page_size, D]``. The cache writer must zero-fill + unused token slots in allocated pages, as required by + :meth:`QTokenKvBlockSparsePagedTSWrapper.run`. + block_table : torch.Tensor + Dense CUDA Int32 physical-page table + ``[num_requests, max_storage_pages]``. + indexer_block_ids : torch.Tensor + CUDA Int32 indexer-selected logical K/V blocks ``[T, block_topk]`` + when shared, otherwise ``[T, Hkv, block_topk]``, with a contiguous last + 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 + CUDA Int32 or Int64 causal position for each query token. + workspace_buffer : torch.Tensor + Caller-owned byte workspace sized by + :func:`get_q_token_kv_block_sparse_workspace_size`. + max_seq_len_kv : int + Static maximum logical K/V length in tokens. + seq_len_q : int, optional + Maximum packed route length. Fixed Q derives it from ``G``. + kv_block_size : int + Semantic sparse K/V block size: 4/8/16/32/64/128, default four. + mask_type : {"causal"} + QToken-KvBlock-Sparse-Attention currently supports causal masking only. + sm_scale : float, optional + Softmax scale, defaulting to ``head_dim**-0.5``. + v_scale : float, optional + Value-cache dequantization scale applied to the attention output, + defaulting to one. Frameworks using an FP8 K/V cache should pass the + cache's live V scale. + out : torch.Tensor, optional + 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 + Output dtype, defaulting to ``out.dtype`` or ``q.dtype``. + qo_indptr : torch.Tensor, optional + CUDA Int32 packed-route offsets. + split_kv : bool + Allow automatic split-KV selection (default True), independently of + query layout. Pass False for prefill and True for decode, matching + the value used to size ``workspace_buffer``. + + share_pattern_across_kv_heads : bool + 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. + + Returns + ------- + torch.Tensor + Attention output in the same logical layout as ``q``. + """ + + k_cache, v_cache = _validate_q_token_kv_block_sparse_paged_kv_cache(paged_kv_cache) + if not isinstance(q, torch.Tensor): + raise TypeError("q must be a torch.Tensor") + pattern_rank = 2 if share_pattern_across_kv_heads else 3 + if ( + not isinstance(indexer_block_ids, torch.Tensor) + or indexer_block_ids.ndim != pattern_rank + ): + raise ValueError( + "indexer_block_ids must be [T,K] when shared or [T,Hkv,K] otherwise" + ) + 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 qo_indptr is None: + if q.ndim != 5: + raise ValueError( + "fixed q must be [B,Nq,G,Hq,D], or supply qo_indptr for packed q" + ) + fixed_seq_len_q = int(q.shape[2]) + if seq_len_q is not None and seq_len_q != fixed_seq_len_q: + raise ValueError("fixed seq_len_q must match q.shape[2]") + seq_len_q = fixed_seq_len_q + batch_size = int(q.shape[0] * q.shape[1]) + use_packed_q = False + else: + if q.ndim != 3: + raise ValueError("packed q must have shape [total_q,Hq,D]") + if seq_len_q is None: + raise ValueError("seq_len_q is required with packed qo_indptr") + batch_size = int(qo_indptr.numel()) - 1 + use_packed_q = True + + wrapper = QTokenKvBlockSparsePagedTSWrapper() + wrapper.plan( + batch_size, + seq_len_q, + int(q.shape[-2]), + int(k_cache.shape[1]), + int(q.shape[-1]), + kv_block_size, + int(k_cache.shape[2]), + int(indexer_block_ids.shape[-1]), + max_seq_len_kv, + device=q.device, + workspace_buffer=workspace_buffer, + use_packed_q=use_packed_q, + 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, + ) + return wrapper.run( + q, + (k_cache, v_cache), + block_table, + indexer_block_ids, + token_to_request, + query_positions, + qo_indptr=qo_indptr, + sm_scale=sm_scale, + v_scale=v_scale, + out=out, + ) + + +def _validate_shape_parameters( + num_query_tokens: int, + block_topk: int, + group_size: int, +) -> None: + for value, name in ( + (num_query_tokens, "num_query_tokens"), + (block_topk, "block_topk"), + (group_size, "group_size"), + ): + if not isinstance(value, int) or isinstance(value, bool): + raise TypeError(f"{name} must be an integer") + if num_query_tokens < 0: + raise ValueError("num_query_tokens must be nonnegative") + if block_topk <= 0: + raise ValueError("block_topk must be positive") + from .decode import _validate_prims_ts_q_token_kv_block_sparse_group_value + + _validate_prims_ts_q_token_kv_block_sparse_group_value(group_size) + + +def _resolve_num_query_groups( + num_query_tokens: int, + group_size: int, + num_query_groups: Optional[int], +) -> int: + """Resolve fixed or request-partitioned route count.""" + + if num_query_groups is None: + if num_query_tokens % group_size: + raise ValueError( + "num_query_tokens must be divisible by group_size without " + "packed QToken-KvBlock-Sparse-Attention qo_indptr" + ) + return num_query_tokens // group_size + if not isinstance(num_query_groups, int) or isinstance(num_query_groups, bool): + raise TypeError("num_query_groups must be an integer") + if num_query_groups <= 0: + raise ValueError("num_query_groups must be positive") + min_groups = (num_query_tokens + group_size - 1) // group_size + if num_query_groups < min_groups or num_query_groups > num_query_tokens: + raise ValueError( + "num_query_groups cannot partition num_query_tokens into nonempty " + f"routes of size at most {group_size}" + ) + return num_query_groups + + +def _validate_q_token_kv_block_sparse_qo_indptr_layout_tensor( + qo_indptr: torch.Tensor, +) -> None: + """Validate the packed-Q offsets required to size QToken-KvBlock-Sparse-Attention workspace.""" + + if ( + not isinstance(qo_indptr, torch.Tensor) + or qo_indptr.ndim != 1 + or qo_indptr.numel() < 2 + or qo_indptr.dtype != torch.int32 + ): + raise ValueError( + "QToken-KvBlock-Sparse-Attention qo_indptr must be a rank-one int32 tensor" + ) + + +def _validate_q_token_kv_block_sparse_qo_indptr_tensor( + qo_indptr: torch.Tensor, + *, + expected_device: torch.device, +) -> None: + """Validate the synchronization-free packed-Q offset tensor contract.""" + + _validate_q_token_kv_block_sparse_qo_indptr_layout_tensor(qo_indptr) + if qo_indptr.device != expected_device or not qo_indptr.is_contiguous(): + raise ValueError( + "QToken-KvBlock-Sparse-Attention qo_indptr must be a contiguous CUDA int32 tensor on the query device" + ) + + +def _validate_sparse_block_size(sparse_block_size: int) -> int: + """Validate the generic sparse-block API and current specialization.""" + + if not isinstance(sparse_block_size, int) or isinstance(sparse_block_size, bool): + raise TypeError("sparse_block_size must be an integer") + if sparse_block_size <= 0 or sparse_block_size & (sparse_block_size - 1): + raise ValueError("sparse_block_size must be a positive power of two") + if sparse_block_size not in _Q_TOKEN_KV_BLOCK_SPARSE_BLOCK_SIZES: + raise NotImplementedError("sparse_block_size must be 4, 8, 16, 32, 64, or 128") + return sparse_block_size + + +def _validate_storage_page_size(storage_page_size: int) -> None: + if not isinstance(storage_page_size, int) or isinstance(storage_page_size, bool): + raise TypeError("storage_page_size must be an integer") + if storage_page_size < 4 or storage_page_size % 4: + raise ValueError("storage_page_size must be a positive multiple of four") + + +def _validate_q_token_kv_block_sparse_max_seq_len_kv( + max_seq_len_kv: int, + *, + block_table_token_capacity: Optional[int] = None, +) -> int: + """Validate the static logical K/V bound used by metadata kernels.""" + + if not isinstance(max_seq_len_kv, int) or isinstance(max_seq_len_kv, bool): + raise TypeError("max_seq_len_kv must be an integer") + if max_seq_len_kv <= 0: + raise ValueError("max_seq_len_kv must be positive") + if max_seq_len_kv > _Q_TOKEN_KV_BLOCK_SPARSE_MAX_SEQ_LEN_KV: + raise ValueError( + "max_seq_len_kv exceeds the current grouped QToken-KvBlock-Sparse-Attention radix-key bound " + f"({_Q_TOKEN_KV_BLOCK_SPARSE_MAX_SEQ_LEN_KV})" + ) + if ( + block_table_token_capacity is not None + and max_seq_len_kv > block_table_token_capacity + ): + raise ValueError( + "max_seq_len_kv must not exceed the dense block-table token " + f"capacity ({block_table_token_capacity})" + ) + return max_seq_len_kv + + +def _align_up_q_token_kv_block_sparse_workspace(value: int) -> int: + return ( + (value + _Q_TOKEN_KV_BLOCK_SPARSE_WORKSPACE_ALIGNMENT - 1) + // _Q_TOKEN_KV_BLOCK_SPARSE_WORKSPACE_ALIGNMENT + * _Q_TOKEN_KV_BLOCK_SPARSE_WORKSPACE_ALIGNMENT + ) + + +def _validate_inputs( + block_indices: torch.Tensor, + block_table: torch.Tensor, + token_to_request: 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, +) -> None: + sparse_block_size = _validate_sparse_block_size(sparse_block_size) + for name, tensor in ( + ("block_indices", block_indices), + ("block_table", block_table), + ("token_to_request", token_to_request), + ("query_positions", query_positions), + ): + 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(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 + if block_indices.ndim != rank or block_indices.dtype != torch.int32: + shape = "[T,K]" if share_pattern_across_kv_heads else "[T,Hkv,K]" + raise ValueError(f"block_indices must be a rank-{rank} int32 tensor {shape}") + if rank == 3 and block_indices.shape[1] <= 0: + raise ValueError("block_indices must have a positive pattern head count") + rows, block_topk = block_indices.shape[0], block_indices.shape[-1] + _validate_shape_parameters(rows, block_topk, group_size) + if block_topk > _Q_TOKEN_KV_BLOCK_SPARSE_MAX_BLOCK_TOPK: + raise NotImplementedError( + f"PrimTS QToken-KvBlock-Sparse-Attention currently supports block_topk <= {_Q_TOKEN_KV_BLOCK_SPARSE_MAX_BLOCK_TOPK}" + ) + _validate_storage_page_size(storage_page_size) + if block_table.ndim != 2 or block_table.dtype != torch.int32: + raise ValueError("block_table must be a nonempty rank-two int32 tensor") + if not all(block_table.shape): + raise ValueError("block_table must be nonempty") + _validate_q_token_kv_block_sparse_max_seq_len_kv( + 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 ( + 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) + if any(tensor.device != block_indices.device for tensor in tensors): + raise ValueError( + "QToken-KvBlock-Sparse-Attention metadata inputs must share one CUDA device" + ) + if block_indices.stride(-1) != 1 or block_table.stride(1) != 1: + raise ValueError("block_indices and block_table rows must be contiguous") + if block_table.stride(0) < block_table.shape[1]: + 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: + raise ValueError( + "per-row QToken-KvBlock-Sparse-Attention metadata must be contiguous" + ) + + +def _validate_q_token_kv_block_sparse_attention_workspace( + workspace_buffer: torch.Tensor, + required_bytes: int, +) -> None: + if not isinstance(workspace_buffer, torch.Tensor): + raise TypeError("workspace_buffer must be a torch.Tensor") + if workspace_buffer.dtype not in (torch.int8, torch.uint8): + raise TypeError("workspace_buffer must have dtype torch.int8 or torch.uint8") + if not workspace_buffer.is_cuda: + raise ValueError("workspace_buffer must be a CUDA tensor") + if not workspace_buffer.is_contiguous(): + raise ValueError("workspace_buffer must be contiguous") + available_bytes = workspace_buffer.numel() * workspace_buffer.element_size() + if available_bytes < required_bytes: + raise ValueError( + "workspace_buffer is too small: requires at least " + f"{required_bytes} bytes, got {available_bytes}" + ) + if workspace_buffer.data_ptr() % 32: + raise ValueError("workspace_buffer data pointer must be 32-byte aligned") + + +__all__ = [ + "QTokenKvBlockSparsePagedTSWrapper", + "get_q_token_kv_block_sparse_workspace_size", + "q_token_kv_block_sparse_attention_with_paged_kv_cache", +] diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/sage.py b/tensorrt_llm/_torch/attention/backends/prims_ts/sage.py new file mode 100644 index 000000000000..dcf0d4ae05a0 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/sage.py @@ -0,0 +1,220 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 by FlashInfer team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Sage attention inputs for the PrimTS decode kernels. + +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 + +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 ``[Hkv, D]``. ``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, 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, 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": (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"] = (num_kv_heads, head_dim) + return shapes + + +def validate_sage_params( + params: SageAttentionParams, + expected_shapes: Mapping[str, tuple[int, 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 + ) + + +__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/tests/unittest/_torch/attention/test_prims_ts_fmha.py b/tests/unittest/_torch/attention/test_prims_ts_fmha.py index a30f05d25c2e..8902890021ce 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_fmha.py +++ b/tests/unittest/_torch/attention/test_prims_ts_fmha.py @@ -16,7 +16,7 @@ import math from dataclasses import replace from types import SimpleNamespace -from unittest.mock import Mock +from unittest.mock import Mock, create_autospec import pytest import torch @@ -947,7 +947,7 @@ def test_context_wrapper_plans_once_and_reads_live_fixed_metadata( ) ) context_postprocess = Mock() - wrapper = Mock() + wrapper = create_autospec(prims_context_module.BatchPrefillPagedTSWrapper, instance=True) wrapper_factory = Mock(return_value=wrapper) monkeypatch.setattr( prims_ts_module.thop, @@ -1034,7 +1034,8 @@ def test_context_wrapper_plans_once_and_reads_live_fixed_metadata( "num_kv_heads": attn.num_kv_heads, "head_dim": attn.head_dim, "q_dtype": torch.bfloat16, - "kv_dtype": torch.bfloat16, + "k_dtype": torch.bfloat16, + "v_dtype": torch.bfloat16, "out_dtype": torch.bfloat16, "page_size": 32, "mask_type": "causal", From e7abf19645f9e21d2d9e39029eaa8f3d6260d616 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Mon, 12 Oct 2026 10:40:52 +0000 Subject: [PATCH 2/3] [None][chore] Minimize PrimTS refresh compatibility patch Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../vendor_patches/flashinfer-prims-ts.patch | 99 ++----------------- 3rdparty/vendor_sources.lock.yaml | 4 +- .../_q_token_kv_block_sparse_policy.py | 1 - .../attention/backends/prims_ts/context.py | 1 - .../fmha_decode/direct_sparse_metadata.py | 1 - .../fmha_decode_resources/sage_scales.py | 1 - .../smem_transformed_kv.py | 1 - .../tmem_transformed_kv.py | 1 - .../backends/prims_ts/kernels/task_compat.py | 1 - .../attention/backends/prims_ts/mla_decode.py | 5 +- .../attention/backends/prims_ts/sage.py | 1 - 11 files changed, 16 insertions(+), 100 deletions(-) diff --git a/3rdparty/vendor_patches/flashinfer-prims-ts.patch b/3rdparty/vendor_patches/flashinfer-prims-ts.patch index a4319901a3d5..d6f80daccc16 100644 --- a/3rdparty/vendor_patches/flashinfer-prims-ts.patch +++ b/3rdparty/vendor_patches/flashinfer-prims-ts.patch @@ -1,12 +1,3 @@ -diff --git a/_q_token_kv_block_sparse_policy.py b/_q_token_kv_block_sparse_policy.py -index a0ea3425aab7a23e26b7292eaaf7f5fb772621e8..45f495c955d878625fd88e13cd86212fd85f3e1b 100644 ---- a/_q_token_kv_block_sparse_policy.py -+++ b/_q_token_kv_block_sparse_policy.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"); - # you may not use this file except in compliance with the License. diff --git a/block_sparse.py b/block_sparse.py index d4652b8b68958ee3dab29d062a7cd082ea604c47..bb56724e986d302854722bf33962bea063f9e6e4 100644 --- a/block_sparse.py @@ -70,16 +61,10 @@ index d4652b8b68958ee3dab29d062a7cd082ea604c47..bb56724e986d302854722bf33962bea0 q: torch.Tensor, paged_kv_cache: PagedKVCache, diff --git a/context.py b/context.py -index dd9a29c3e36111a0de9a57c487068b6b5caa4ea2..a8a337373ad1607ca12ac12b2bb3117e485704d8 100644 +index dd9a29c3e36111a0de9a57c487068b6b5caa4ea2..1163324ec632f2a0938f78a3636fb415450369df 100644 --- a/context.py +++ b/context.py -@@ -1,4 +1,5 @@ - # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - # Copyright (c) 2026 by FlashInfer team. - # - # Licensed under the Apache License, Version 2.0 (the "License"); -@@ -46,7 +47,7 @@ from typing import TYPE_CHECKING, Callable, Literal, Optional +@@ -46,7 +46,7 @@ from typing import TYPE_CHECKING, Callable, Literal, Optional import torch @@ -88,7 +73,7 @@ index dd9a29c3e36111a0de9a57c487068b6b5caa4ea2..a8a337373ad1607ca12ac12b2bb3117e if TYPE_CHECKING: -@@ -488,7 +489,7 @@ def _validate_device(device: torch.device) -> int: +@@ -488,7 +488,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): @@ -146,53 +131,8 @@ index 41479a386d72c27eb0df76107318dc09256aff06..3fcec5efdc2a3825be0dc1d9a238aaaa def batch_decode_with_paged_kv_cache( q: torch.Tensor, paged_kv_cache: PagedKVCache, -diff --git a/kernels/fmha_decode/direct_sparse_metadata.py b/kernels/fmha_decode/direct_sparse_metadata.py -index 8bffb076c4a3ef56bc9e463733066abcaf0a7725..15d6f271029540de1e76633e1a81cd78bc553729 100644 ---- a/kernels/fmha_decode/direct_sparse_metadata.py -+++ b/kernels/fmha_decode/direct_sparse_metadata.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"); - # you may not use this file except in compliance with the License. -diff --git a/kernels/fmha_decode/fmha_decode_resources/sage_scales.py b/kernels/fmha_decode/fmha_decode_resources/sage_scales.py -index f2df9bbeafcbdf22e67f9a754d06144b8e3a1dab..4d862f9fa1f1be3e8830a7bbc502e0369237f0d4 100644 ---- a/kernels/fmha_decode/fmha_decode_resources/sage_scales.py -+++ b/kernels/fmha_decode/fmha_decode_resources/sage_scales.py -@@ -1,3 +1,4 @@ -+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - # Copyright (c) 2026 by FlashInfer team. - # - # Licensed under the Apache License, Version 2.0 (the "License"); -diff --git a/kernels/fmha_decode/fmha_decode_resources/smem_transformed_kv.py b/kernels/fmha_decode/fmha_decode_resources/smem_transformed_kv.py -index 184326c482f21ecb1a81f8ba337a996f9ab2bea8..148ebcae6b7431985bb91c02e7fa6d5ddfe69253 100644 ---- a/kernels/fmha_decode/fmha_decode_resources/smem_transformed_kv.py -+++ b/kernels/fmha_decode/fmha_decode_resources/smem_transformed_kv.py -@@ -1,3 +1,4 @@ -+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - # Copyright (c) 2026 by FlashInfer team. - # - # Licensed under the Apache License, Version 2.0 (the "License"); -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 18f2c272cf5c1e089193a8a5437560f45cbbbb8e..1d607a92e63df03c04465a9c7836422088c13d49 100644 ---- a/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py -+++ b/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py -@@ -1,3 +1,4 @@ -+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - # Copyright (c) 2026 by FlashInfer team. - # - # Licensed under the Apache License, Version 2.0 (the "License"); -diff --git a/kernels/task_compat.py b/kernels/task_compat.py -index 80a80102c3ecf6a856d441c0dac2ca5c8215ad82..69f60317b172457e7fda347a3b935ae67dbdd892 100644 ---- a/kernels/task_compat.py -+++ b/kernels/task_compat.py -@@ -1,3 +1,4 @@ -+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - # Copyright (c) 2026 by FlashInfer team. - # - # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/mla_decode.py b/mla_decode.py -index 41ba3b9055d229d227ac10cfc6f7d0f403eae69f..6d69ae7b554afe1502f4c369e5ddc805847b876c 100644 +index 41ba3b9055d229d227ac10cfc6f7d0f403eae69f..3d5fb159a3298d1edf72935b43d9a2eef74a1a4b 100644 --- a/mla_decode.py +++ b/mla_decode.py @@ -1,3 +1,4 @@ @@ -200,20 +140,19 @@ index 41ba3b9055d229d227ac10cfc6f7d0f403eae69f..6d69ae7b554afe1502f4c369e5ddc805 # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); -@@ -21,11 +22,7 @@ from typing import Any, Literal, Optional, cast +@@ -21,10 +22,9 @@ from typing import Any, Literal, Optional, cast import torch -from flashinfer.api_logging import flashinfer_experimental_api --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 flashinfer.trace.templates.attention import ( + prims_ts_decode_mla_one_shot_trace_dispatch, +- prims_ts_decode_mla_wrapper_trace_dispatch, + ) from .decode import ( - _WorkspaceSection, -@@ -1576,7 +1573,7 @@ class BatchMLADecodePagedTSWrapper: +@@ -1576,7 +1576,7 @@ class BatchMLADecodePagedTSWrapper: split_kv=int(dict(policy)["split_kv"]), ) @@ -222,15 +161,6 @@ index 41ba3b9055d229d227ac10cfc6f7d0f403eae69f..6d69ae7b554afe1502f4c369e5ddc805 def run( self, query: torch.Tensor, -@@ -1681,7 +1678,7 @@ class BatchMLADecodePagedTSWrapper: - ) - - --@flashinfer_experimental_api(trace=prims_ts_decode_mla_one_shot_trace_dispatch) -+@flashinfer_experimental_api - def batch_mla_decode_with_paged_kv_cache( - query: torch.Tensor, - kv_cache: torch.Tensor, diff --git a/q_token_kv_block_sparse_metadata.py b/q_token_kv_block_sparse_metadata.py index 35efc60b7ff3dac6af1c5c5e46631712a171fb72..0c5cf40dbd4e4708b8f74bfa698f3c0788f9f3a4 100644 --- a/q_token_kv_block_sparse_metadata.py @@ -261,12 +191,3 @@ index 35efc60b7ff3dac6af1c5c5e46631712a171fb72..0c5cf40dbd4e4708b8f74bfa698f3c07 return gen_prims_ts_q_token_kv_block_sparse_metadata_module().build_and_load() -diff --git a/sage.py b/sage.py -index 97708c79525ba840e4ebfd11876084eae952254d..dcf0d4ae05a0d1358cfa796a3e23dbe098fbab10 100644 ---- a/sage.py -+++ b/sage.py -@@ -1,3 +1,4 @@ -+# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. - # Copyright (c) 2026 by FlashInfer team. - # - # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 9f5f510adefd..2e914da8f052 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -9,8 +9,8 @@ vendors: include: - '**/*.py' patch: 3rdparty/vendor_patches/flashinfer-prims-ts.patch - patch_digest: sha256:a9caa1909814f67df34fff35f21043d95c659ccbd86975924c9d8784adcdaf39 - digest: sha256-tree-v1:1f2f960dfe7b4f875b2e8d6a8b119fec9aae18a527e175b2cd4bde66393e16b0 + patch_digest: sha256:65cc08e20a9cc7bdc190eb6f3e79124986dced74394c1f3250ae2d100a8ee340 + digest: sha256-tree-v1:e411b6c6679474d362b166e586df4276f52322d75a35c59c782d6266143e560e openengine-proto: url: https://github.com/ai-dynamo/openengine.git tag: v0.1.0 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 45f495c955d8..a0ea3425aab7 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 @@ -1,4 +1,3 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2026 by FlashInfer team. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/context.py b/tensorrt_llm/_torch/attention/backends/prims_ts/context.py index a8a337373ad1..1163324ec632 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/context.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/context.py @@ -1,5 +1,4 @@ # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/direct_sparse_metadata.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/direct_sparse_metadata.py index 15d6f2710295..8bffb076c4a3 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 @@ -1,4 +1,3 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2026 by FlashInfer team. # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. 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 4d862f9fa1f1..f2df9bbeafcb 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 @@ -1,4 +1,3 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_transformed_kv.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_transformed_kv.py index 148ebcae6b74..184326c482f2 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_transformed_kv.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_transformed_kv.py @@ -1,4 +1,3 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py index 1d607a92e63d..18f2c272cf5c 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_transformed_kv.py @@ -1,4 +1,3 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/task_compat.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/task_compat.py index 69f60317b172..80a80102c3ec 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/task_compat.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/task_compat.py @@ -1,4 +1,3 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py b/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py index 6d69ae7b554a..3d5fb159a329 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/mla_decode.py @@ -23,6 +23,9 @@ import torch from flashinfer.api_logging import flashinfer_api as flashinfer_experimental_api +from flashinfer.trace.templates.attention import ( + prims_ts_decode_mla_one_shot_trace_dispatch, +) from .decode import ( _WorkspaceSection, @@ -1678,7 +1681,7 @@ def run( ) -@flashinfer_experimental_api +@flashinfer_experimental_api(trace=prims_ts_decode_mla_one_shot_trace_dispatch) def batch_mla_decode_with_paged_kv_cache( query: torch.Tensor, kv_cache: torch.Tensor, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/sage.py b/tensorrt_llm/_torch/attention/backends/prims_ts/sage.py index dcf0d4ae05a0..97708c79525b 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/sage.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/sage.py @@ -1,4 +1,3 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2026 by FlashInfer team. # # Licensed under the Apache License, Version 2.0 (the "License"); From 79188af438a2bbf079fa93e47b654861ee1140f5 Mon Sep 17 00:00:00 2001 From: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> Date: Tue, 13 Oct 2026 04:02:44 +0000 Subject: [PATCH 3/3] [None][fix] Align PrimTS refresh consumers with upstream contracts Signed-off-by: Yuxian Qiu <142763828+yuxianq@users.noreply.github.com> --- .../attention/backends/sparse/params.py | 3 +- .../attention_backend/sparse/sol/kernels.py | 3 +- .../sparse/test_prims_ts_block_sparse.py | 4 +- .../test_prims_ts_attention_backend.py | 74 +++++++------------ .../sparse_attention/sol/test_sol_parity.py | 9 +-- .../sol/test_sol_predictor.py | 6 +- 6 files changed, 38 insertions(+), 61 deletions(-) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/params.py b/tensorrt_llm/_torch/attention/backends/sparse/params.py index 2c87709bdc5a..065ac68a503e 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/params.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/params.py @@ -114,7 +114,8 @@ class BlockSparseForwardInputs: Exactly one routing representation is present. Canonical BSR uses ``block_indptr`` and ``block_indices``; packed bitmask routing uses ``exact_block_bits``. Paired K/V summaries enable proxy routes without - encoding an algorithm name in this shared carrier. + encoding an algorithm name in this shared carrier. Both summaries are + block means over the tokens present, including a partial final block. """ q_block_size: int diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/kernels.py b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/kernels.py index 5bf870459233..98d9de32ea22 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/kernels.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/kernels.py @@ -655,6 +655,7 @@ def visual_gen_sol_predictor( ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Return ``(exact_block_bits, k_summary, v_summary)`` for compact BSHD ``q``, ``k`` and ``v``. + K/V summaries are block means over the tokens present, including a partial final block. The outputs are fresh tensors of this call. Inside CUDA Graph capture they come from the graph pool and stay valid for every replay; under torch.compile the operator is opaque and shapes come from the fake implementation. @@ -668,7 +669,7 @@ def visual_gen_sol_predictor( v_summary = torch.empty(summary_shape, dtype=v.dtype, device=v.device) block_pool(q, q_centroid, block_size=block_size, reduce="mean") block_pool(k, k_summary, block_size=block_size, reduce="mean") - block_pool(v, v_summary, block_size=block_size, reduce="sum") + block_pool(v, v_summary, block_size=block_size, reduce="mean") threshold = block_thresholds( q_centroid, k_summary, tau=tau, sm_scale=sm_scale, thresh_type=thresh_type ) diff --git a/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py b/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py index 44562722fce1..090385c73dcb 100644 --- a/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py +++ b/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py @@ -113,7 +113,7 @@ def _proxy_reference( weights = torch.exp(logits - logits.amax(dim=1, keepdim=True)) exact_weights, proxy_weights = weights.split((block_size, len(proxy_blocks)), dim=1) numerator = exact_weights @ v[0, exact_tokens, 0].float() - numerator += proxy_weights @ v_summary[0, proxy_blocks, 0].float() + numerator += (proxy_weights * block_size) @ v_summary[0, proxy_blocks, 0].float() denominator = exact_weights.sum(dim=1, keepdim=True) denominator += proxy_weights.sum(dim=1, keepdim=True) * block_size return (numerator / denominator).to(q.dtype)[None, :, None] @@ -632,7 +632,7 @@ def test_real_gpu_proxy_adapter_replays_live_routes_and_summaries() -> None: k_blocks = k.float().view(1, 3, 64, 1, 128) v_blocks = v.float().view(1, 3, 64, 1, 128) initial_k_summary = k_blocks.mean(dim=2).to(k.dtype) - initial_v_summary = v_blocks.sum(dim=2).to(v.dtype) + initial_v_summary = v_blocks.mean(dim=2).to(v.dtype) live_k_summary = initial_k_summary.clone() live_v_summary = initial_v_summary.clone() live_exact_bits = torch.tensor([[[[1]]]], device="cuda", dtype=torch.uint32) 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 7e0b31bf97be..2adfac26c11c 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py +++ b/tests/unittest/_torch/attention/test_prims_ts_attention_backend.py @@ -656,7 +656,8 @@ def compile_with_record(*args): num_kv_heads=num_kv_heads, head_dim=head_dim, q_dtype=dtype, - kv_dtype=dtype, + k_dtype=dtype, + v_dtype=dtype, out_dtype=dtype, page_size=page_size, mask_type="causal", @@ -727,8 +728,8 @@ def compile_with_record(*args): def test_prims_ts_decode_live_wrapper_cuda_graph_replay() -> None: from tensorrt_llm._torch.attention.backends.prims_ts import ( BatchDecodePagedTSWrapper, + batch_decode_with_paged_kv_cache, get_prims_ts_batch_decode_workspace_size, - prims_ts_batch_decode_with_kv_cache, ) batch_size = 2 @@ -779,7 +780,7 @@ def test_prims_ts_decode_live_wrapper_cuda_graph_replay() -> None: device=device, ) external_workspace = torch.zeros( - max(workspace_bytes, query.numel() * query.element_size()), + workspace_bytes, device=device, dtype=torch.uint8, ) @@ -804,18 +805,6 @@ def test_prims_ts_decode_live_wrapper_cuda_graph_replay() -> None: assert plan_state is not None compiled_main = plan_state.compiled_main - aliased_query = ( - external_workspace[: query.numel() * query.element_size()].view(dtype).view_as(query) - ) - with pytest.raises(ValueError, match="workspace_buffer must not overlap query storage"): - wrapper.run( - aliased_query, - kv_cache, - seq_lens, - block_tables, - out=output, - ) - external_workspace.zero_() wrapper.run( query, @@ -840,13 +829,13 @@ def test_prims_ts_decode_live_wrapper_cuda_graph_replay() -> None: block_tables.copy_(torch.tensor([[2, 3], [0, 1]], device=device, dtype=torch.int32)) seq_lens.copy_(torch.tensor([32, 64], device=device, dtype=torch.int32)) reference_workspace = torch.zeros_like(external_workspace) - reference = prims_ts_batch_decode_with_kv_cache( + reference = batch_decode_with_paged_kv_cache( query, kv_cache, - reference_workspace, block_tables, seq_lens, - max_seq_len, + workspace_buffer=reference_workspace, + max_kv_len=max_seq_len, out_dtype=dtype, out=torch.empty_like(query), mask_type="causal", @@ -868,8 +857,8 @@ def test_prims_ts_decode_live_wrapper_cuda_graph_replay() -> None: def test_prims_ts_mla_live_wrapper_cuda_graph_replay(max_seq_len: int) -> None: from tensorrt_llm._torch.attention.backends.prims_ts import ( BatchMLADecodePagedTSWrapper, + batch_mla_decode_with_paged_kv_cache, get_prims_ts_batch_mla_decode_workspace_size, - prims_ts_batch_mla_decode_with_kv_cache, ) batch_size = 2 @@ -927,7 +916,7 @@ def test_prims_ts_mla_live_wrapper_cuda_graph_replay(max_seq_len: int) -> None: device=device, ) external_workspace = torch.empty( - max(workspace_bytes, query.numel() * query.element_size()), + workspace_bytes, device=device, dtype=torch.uint8, ) @@ -955,19 +944,6 @@ def test_prims_ts_mla_live_wrapper_cuda_graph_replay(max_seq_len: int) -> None: assert dict(plan_state.policy)["kernel"] == "throughput_2cta" bmm1_scale = (128 + qk_rope_head_dim) ** -0.5 - aliased_query = ( - external_workspace[: query.numel() * query.element_size()].view(dtype).view_as(query) - ) - with pytest.raises(ValueError, match="workspace_buffer must not overlap query storage"): - wrapper.run( - aliased_query, - kv_cache, - block_tables, - seq_lens, - bmm1_scale=bmm1_scale, - out=output, - ) - wrapper.run( query, kv_cache, @@ -992,15 +968,15 @@ def test_prims_ts_mla_live_wrapper_cuda_graph_replay(max_seq_len: int) -> None: block_tables.copy_(block_tables.flip(0).clone()) seq_lens.copy_(seq_lens.flip(0).clone()) reference_workspace = torch.empty_like(external_workspace) - reference = prims_ts_batch_mla_decode_with_kv_cache( + reference = batch_mla_decode_with_paged_kv_cache( query, kv_cache, - reference_workspace, - kv_lora_rank, - qk_rope_head_dim, block_tables, seq_lens, - max_seq_len, + workspace_buffer=reference_workspace, + 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, out_dtype=dtype, @@ -1020,8 +996,8 @@ def test_prims_ts_mla_live_wrapper_cuda_graph_replay(max_seq_len: int) -> None: def test_prims_ts_decode_graph_profiles_reset_shared_workspace_a_b_a() -> None: from tensorrt_llm._torch.attention.backends.prims_ts import ( BatchDecodePagedTSWrapper, + batch_decode_with_paged_kv_cache, get_prims_ts_batch_decode_workspace_size, - prims_ts_batch_decode_with_kv_cache, ) num_qo_heads = 8 @@ -1141,26 +1117,26 @@ def capture( plan(1, max_seq_len_a) reference_workspace = torch.zeros_like(shared_workspace) - reference_a = prims_ts_batch_decode_with_kv_cache( + reference_a = batch_decode_with_paged_kv_cache( query_a, kv_cache, - reference_workspace, block_tables_a, seq_lens_a, - max_seq_len_a, + workspace_buffer=reference_workspace, + max_kv_len=max_seq_len_a, out_dtype=dtype, out=torch.empty_like(output_a), mask_type="causal", kv_layout="HND", ).clone() reference_workspace.zero_() - reference_b = prims_ts_batch_decode_with_kv_cache( + reference_b = batch_decode_with_paged_kv_cache( query_b, kv_cache, - reference_workspace, block_tables_b, seq_lens_b, - max_seq_len_b, + workspace_buffer=reference_workspace, + max_kv_len=max_seq_len_b, out_dtype=dtype, out=torch.empty_like(output_b), mask_type="causal", @@ -1196,8 +1172,8 @@ def capture( def test_prims_ts_decode_wrappers_share_workspace_across_serialized_layers() -> None: from tensorrt_llm._torch.attention.backends.prims_ts import ( BatchDecodePagedTSWrapper, + batch_decode_with_paged_kv_cache, get_prims_ts_batch_decode_workspace_size, - prims_ts_batch_decode_with_kv_cache, ) batch_size = 2 @@ -1282,13 +1258,13 @@ def test_prims_ts_decode_wrappers_share_workspace_across_serialized_layers() -> out=output, ) reference_workspace = torch.zeros_like(shared_workspace) - reference = prims_ts_batch_decode_with_kv_cache( + reference = batch_decode_with_paged_kv_cache( query, kv_cache, - reference_workspace, block_tables, seq_lens, - max_seq_len, + workspace_buffer=reference_workspace, + max_kv_len=max_seq_len, out_dtype=dtype, out=torch.empty_like(query), mask_type="causal", diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_parity.py b/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_parity.py index 34c240070148..df6735fdc696 100644 --- a/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_parity.py +++ b/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_parity.py @@ -106,17 +106,16 @@ def _mixed_proxy_reference( weights = torch.exp(logits - logits.amax(dim=1, keepdim=True)) exact_weights = weights[:, : exact_tokens.numel()] proxy_weights = weights[:, exact_tokens.numel() :] + for proxy_offset, block_idx in enumerate(proxy_blocks): + tokens_in_block = min(_BLOCK, k.shape[1] - block_idx * _BLOCK) + proxy_weights[:, proxy_offset] *= tokens_in_block numerator = exact_weights @ v[batch_idx, exact_tokens, head_idx].float() if proxy_blocks: numerator += ( proxy_weights @ outputs.v_summary[batch_idx, proxy_blocks, head_idx].float() ) denominator = exact_weights.sum(dim=1, keepdim=True) - for proxy_offset, block_idx in enumerate(proxy_blocks): - tokens_in_block = min(_BLOCK, k.shape[1] - block_idx * _BLOCK) - denominator += proxy_weights[:, proxy_offset : proxy_offset + 1] * ( - tokens_in_block - ) + denominator += proxy_weights.sum(dim=1, keepdim=True) reference[batch_idx, q_begin:q_end, head_idx] = (numerator / denominator).to( q.dtype ) diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor.py b/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor.py index 2b5d51a7fef9..f5ce73cc777b 100644 --- a/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor.py +++ b/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor.py @@ -55,7 +55,7 @@ def _summary_oracle(k: torch.Tensor, v: torch.Tensor) -> tuple[torch.Tensor, tor begin = block_idx * BLOCK_SIZE end = min(begin + BLOCK_SIZE, k.shape[1]) k_summary[:, block_idx] = k[:, begin:end].float().mean(dim=1).to(torch.bfloat16) - v_summary[:, block_idx] = v[:, begin:end].float().sum(dim=1).to(torch.bfloat16) + v_summary[:, block_idx] = v[:, begin:end].float().mean(dim=1).to(torch.bfloat16) return k_summary, v_summary @@ -213,7 +213,7 @@ def test_sol_predictor_runtime_scale_and_tau_extremes() -> None: @_REQUIRES_CUDA -def test_sol_predictor_long_proxy_group_keeps_tail_mass_and_clears_padding_bits() -> None: +def test_sol_predictor_long_proxy_group_keeps_tail_mean_and_clears_padding_bits() -> None: tokens = 16_451 q, k, _ = _inputs((1, tokens, 1, 128), 31) v = torch.ones_like(q) @@ -225,7 +225,7 @@ def test_sol_predictor_long_proxy_group_keeps_tail_mass_and_clears_padding_bits( assert torch.equal(outputs.exact_block_bits, expected) torch.testing.assert_close(outputs.k_summary[:, -1], expected_k[:, -1], rtol=1e-2, atol=1e-2) assert torch.equal(outputs.v_summary[:, -1], expected_v[:, -1]) - assert torch.all(outputs.v_summary[:, -1] == 3) + assert torch.all(outputs.v_summary == 1) assert int(outputs.exact_block_bits[..., -1].to(torch.int64).max()) < 4