From dc4cacaa24d149364e1488599819af1c7f9a0934 Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Tue, 8 Sep 2026 13:43:48 +0000 Subject: [PATCH 1/7] chore: update FlashInfer PrimTS pin Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- .../vendor_patches/flashinfer-prims-ts.patch | 22 +- 3rdparty/vendor_sources.lock.yaml | 10 +- .../backends/prims_ts/_block_sparse/common.py | 46 + .../prims_ts/_block_sparse/compiler.py | 331 ++++- .../backends/prims_ts/_block_sparse/config.py | 46 +- .../prims_ts/_block_sparse/inspection.py | 14 +- .../backends/prims_ts/_block_sparse/plan.py | 48 +- .../prims_ts/_block_sparse/prepared.py | 26 +- .../prims_ts/_block_sparse/runtime.py | 395 ++++-- .../backends/prims_ts/block_sparse.py | 311 +++-- .../fmha_decode/block_sparse_inspect.py | 88 +- .../fmha_decode/block_sparse_prepare.py | 1085 +++++++++++------ .../kernels/fmha_decode/fmha_decode_config.py | 198 ++- .../fmha_decode/fmha_decode_constants.py | 8 +- .../kernels/fmha_decode/fmha_decode_kernel.py | 300 +++-- .../fmha_decode_resources/helpers_common.py | 45 + .../fmha_decode_resources/helpers_softmax.py | 64 + .../smem_block_sparse_metadata.py | 528 +++++--- .../fmha_decode_resources/smem_p.py | 461 +++++-- .../fmha_decode_resources/smem_resources.py | 76 +- .../fmha_decode_resources/tmem_corr.py | 471 ++++--- .../fmha_decode_resources/tmem_o.py | 51 +- .../fmha_decode_resources/tmem_s.py | 885 ++++---------- .../kernels/fmha_decode/fmha_decode_tasks.py | 651 +++++----- 24 files changed, 3871 insertions(+), 2289 deletions(-) diff --git a/3rdparty/vendor_patches/flashinfer-prims-ts.patch b/3rdparty/vendor_patches/flashinfer-prims-ts.patch index 44f3e7a4581d..b5d4bc75155c 100644 --- a/3rdparty/vendor_patches/flashinfer-prims-ts.patch +++ b/3rdparty/vendor_patches/flashinfer-prims-ts.patch @@ -1,5 +1,5 @@ diff --git a/block_sparse.py b/block_sparse.py -index 77a9fc4542ab5fb82aa0df6f7c73e56b7d5fbe89..af132c6a69ce0e4a35903b8e4428f481e949f6b8 100644 +index c7e881ca73dfb95314da532912db39fe7b05eb94..e19dc7b1de7fca03c44b8e98217f8d10fb17e26b 100644 --- a/block_sparse.py +++ b/block_sparse.py @@ -1,3 +1,4 @@ @@ -12,15 +12,15 @@ index 77a9fc4542ab5fb82aa0df6f7c73e56b7d5fbe89..af132c6a69ce0e4a35903b8e4428f481 from flashinfer.api_logging import flashinfer_api -from flashinfer.trace.templates.attention import ( -- prims_ts_block_sparse_trace, +- 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 ._block_sparse.common import _validate_contiguous_route_mode from ._block_sparse.config import _validate_block_sparse_static_profile - from ._block_sparse.inspection import ( -@@ -220,7 +215,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase): +@@ -239,7 +234,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase): # previously published revision intact and runnable. self._plan_state = candidate @@ -29,16 +29,16 @@ index 77a9fc4542ab5fb82aa0df6f7c73e56b7d5fbe89..af132c6a69ce0e4a35903b8e4428f481 def run( self, q: torch.Tensor, -@@ -302,7 +297,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase): +@@ -361,7 +356,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase): return self._launch_validated_run(state, run_args, run_stream) --@flashinfer_api(trace=prims_ts_block_sparse_trace) +-@flashinfer_api(trace=prims_ts_block_sparse_trace_dispatch) +@flashinfer_api def block_sparse_attention( q: torch.Tensor, k: torch.Tensor, -@@ -512,7 +507,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase): +@@ -607,7 +602,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase): ) self._plan_state = candidate @@ -47,7 +47,7 @@ index 77a9fc4542ab5fb82aa0df6f7c73e56b7d5fbe89..af132c6a69ce0e4a35903b8e4428f481 def run( self, q: torch.Tensor, -@@ -618,7 +613,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase): +@@ -727,7 +722,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase): return self._launch_validated_run(state, run_args, run_stream) @@ -57,10 +57,10 @@ index 77a9fc4542ab5fb82aa0df6f7c73e56b7d5fbe89..af132c6a69ce0e4a35903b8e4428f481 q: torch.Tensor, paged_kv_cache: PagedKVCache, diff --git a/context.py b/context.py -index 7245bee1a0f725086171c9c5002115757e425d84..47996c867a962a3684c78e07d3a14eebf34b8452 100644 +index cea5a41438d9d72e152175999453881cc9c3e5e6..3a93019ff4a5995968c5efa9000e8ebd52c8b424 100644 --- a/context.py +++ b/context.py -@@ -29,8 +29,7 @@ position is ``q + (S_kv - S_q)`` and ``window_left`` is measured from that +@@ -31,8 +31,7 @@ position is ``q + (S_kv - S_q)`` and ``window_left`` is measured from that position. PrimTS context entry points are intentionally excluded from ``fi_trace`` for @@ -70,7 +70,7 @@ index 7245bee1a0f725086171c9c5002115757e425d84..47996c867a962a3684c78e07d3a14eeb """ from dataclasses import dataclass -@@ -424,7 +423,7 @@ def _validate_device(device: torch.device) -> int: +@@ -426,7 +425,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): diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 705d8af9091d..97a88dadbe93 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -1,13 +1,13 @@ schema_version: 1 vendors: flashinfer-prims-ts: - url: https://github.com/yuxianq/flashinfer.git - branch: trtllm-prims-ts-dev - commit: e500966b575ab83db7c0e84e5a0f8fde6a4f3505 + url: https://github.com/heyuhhh/flashinfer.git + branch: yuhangh/tmp-sol-attn-trtllm-dev + commit: 61454c5ce6e9f020a2158f24059d0b91322e3e77 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:0e2f58c6633f57fee03df42049bc78d4d038063b810ad3a2f0ba6d62f8183887 - digest: sha256-tree-v1:e9af5482f6406af3128e711c3fb2d044359fb6e1d1bc9907743dc86d006d23ac + patch_digest: sha256:b590a3e86c8a2a54a8da5a2268aa9e67f8425f401f6b98972841bfe5484b5732 + digest: sha256-tree-v1:e89b89471aac59e2689ed0fb8312cef77b08e1b46115e7668d751b9f402a29ba 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 a82deafc7a3f..13cf6038efd9 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 @@ -22,6 +22,20 @@ _BLOCK_SPARSE_MAX_HEADS_Q_PER_KV = 32 +def _validate_contiguous_route_mode( + sparse_format: object, + use_proxy_routes: object, +) -> None: + """Validate the two public continuous-route axes before device work.""" + + if not isinstance(sparse_format, str): + raise TypeError("sparse_format must be 'bsr' or 'bitmask'") + if sparse_format not in ("bsr", "bitmask"): + raise ValueError("sparse_format must be 'bsr' or 'bitmask'") + if type(use_proxy_routes) is not bool: + raise TypeError("use_proxy_routes must be a bool") + + def _validate_sparse_q_block_size(value: object) -> int: """Return a positive semantic Q block size representable by the ABI.""" @@ -110,6 +124,17 @@ def _block_sparse_kv_atom_size(kv_block_size: int) -> int: ) +def _block_sparse_proxy_summary_geometry( + seq_len_kv: int, + kv_block_size: int, +) -> tuple[int, int]: + """Return the summary count and final summary's represented-token mass.""" + + num_summaries = (seq_len_kv + kv_block_size - 1) // kv_block_size + tail_mass = seq_len_kv - (num_summaries - 1) * kv_block_size + return num_summaries, tail_mass + + def _prepared_kv_routes_are_block_aligned( kv_block_size: int, kv_route_size: int, @@ -117,3 +142,24 @@ def _prepared_kv_routes_are_block_aligned( """Return whether each prepared route stays within one semantic BSR block.""" return _validate_sparse_kv_block_size(kv_block_size) % kv_route_size == 0 + + +def _block_sparse_contiguous_kv_copy_geometry( + *, + kv_block_size: int, + kv_route_size: int, +) -> tuple[int, int, bool]: + """Return source-independent primary/atom TensorMap geometry. + + Exact and proxy routes address different logical matrices, but a route's + physical copy shape depends only on its semantic block and physical route + sizes. Coarse routes prefer KV128 copies and keep a KV64 descriptor only + when KV256 staging or runtime adjacency requires it. + """ + + atom_size = _block_sparse_kv_atom_size(kv_block_size) + primary_box_size = 2 * atom_size if atom_size == 64 else atom_size + needs_aux_atom = atom_size == 64 and ( + kv_route_size == 256 or kv_block_size % kv_route_size != 0 + ) + return primary_box_size, atom_size, needs_aux_atom 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 a73f26060753..fcaa0735436f 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/compiler.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/compiler.py @@ -36,31 +36,46 @@ def _compile_block_sparse(key: _BlockSparseCompileKey) -> Callable[..., object]: from ..kernels.fmha_decode.fmha_decode_config import FmhaDecodeConfig from ..kernels.fmha_decode.block_sparse_prepare import ( - _PrepareBlockSparseRoutes, + _PrepareBitmaskRoutes, + _PrepareBsrRoutes, ) from ..kernels.fmha_decode.fmha_decode_kernel import ( fmha_block_sparse_launch, ) config = _make_block_sparse_config(key) - prepare_routes = _PrepareBlockSparseRoutes( - 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, - has_token_bits=key.use_kv_valid_bits, - page_size=key.page_size, - mask_type=key.mask_type, - ) + 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'") + + route_metadata_base = prepare_routes.route_metadata_base_word_offset Int32 = cutlass.Int32 Int64 = cutlass.Int64 Float32 = cutlass.Float32 @cute.jit - def contiguous_tensor_adapter( + def exact_bsr_adapter( q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, @@ -86,7 +101,7 @@ def contiguous_tensor_adapter( kv_valid_bits, None, None, - None, + Int64(0), Int64(0), row_route_offsets, route_workspace, @@ -95,9 +110,7 @@ def contiguous_tensor_adapter( ) # Live per-row route counts occupy the first words of run scratch. row_route_counts = route_workspace.iterator - route_metadata = route_workspace.iterator + Int32( - prepare_routes.route_metadata_base_word_offset - ) + route_metadata = route_workspace.iterator + Int32(route_metadata_base) fmha_block_sparse_launch( ( Int32(static_batch_size), @@ -109,6 +122,8 @@ def contiguous_tensor_adapter( q.iterator, k.iterator, v.iterator, + k.iterator, + v.iterator, out.iterator, row_route_offsets.iterator, row_route_counts, @@ -119,6 +134,169 @@ def contiguous_tensor_adapter( 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( + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + k_summary: cute.Tensor, + v_summary: 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_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 paged_tensor_adapter( q: cute.Tensor, @@ -128,13 +306,13 @@ def paged_tensor_adapter( block_indptr: cute.Tensor, block_indices: cute.Tensor, kv_valid_bits: cute.Tensor, - paged_kv_indptr: cute.Tensor, - paged_kv_indices: cute.Tensor, + block_tables: cute.Tensor, seq_lens_kv: cute.Tensor, row_route_offsets: cute.Tensor, route_workspace: cute.Tensor, max_blocks_per_row: cutlass.Int32, num_physical_kv_pages: cutlass.Int64, + block_table_row_stride: cutlass.Int64, k_page_stride: cutlass.Int64, v_page_stride: cutlass.Int64, sm_scale: cutlass.Float32, @@ -151,18 +329,16 @@ def paged_tensor_adapter( block_indices, kv_valid_bits, seq_lens_kv, - paged_kv_indptr, - paged_kv_indices, + block_tables, num_physical_kv_pages, + block_table_row_stride, row_route_offsets, route_workspace, max_blocks_per_row, stream, ) row_route_counts = route_workspace.iterator - route_metadata = route_workspace.iterator + Int32( - prepare_routes.route_metadata_base_word_offset - ) + route_metadata = route_workspace.iterator + Int32(route_metadata_base) fmha_block_sparse_launch( ( Int32(static_batch_size), @@ -174,6 +350,8 @@ def paged_tensor_adapter( q.iterator, k_cache.iterator, v_cache.iterator, + k_cache.iterator, + v_cache.iterator, out.iterator, row_route_offsets.iterator, row_route_counts, @@ -203,6 +381,7 @@ def fake_compact( 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), @@ -238,24 +417,96 @@ def fake_compact( ) k_fake = fake_compact(config.kv_dtype, kv_shape, 16) v_fake = fake_compact(config.kv_dtype, kv_shape, 16) - tensor_adapter = contiguous_tensor_adapter - dynamic_args = ( - q_fake, - k_fake, - v_fake, - out_fake, - indptr_fake, - indices_fake, + 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: + 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, + ) + tensor_adapter = proxy_bsr_adapter + dynamic_args = ( + *proxy_prefix, + indptr_fake, + indices_fake, + *common_tail, + ) + 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, + ) + 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, + ) + else: + raise AssertionError("continuous sparse_format must be 'bsr' or 'bitmask'") else: page_size = key.page_size + assert page_size is not None physical_pages = cute.sym_int() - logical_pages = cute.sym_int() + runtime_page_columns = cute.sym_int() + runtime_page_row_stride = cute.sym_int64(divisibility=1) k_outer_stride = cute.sym_int64(divisibility=1) v_outer_stride = cute.sym_int64(divisibility=1) kv_shape = ( @@ -286,12 +537,12 @@ def fake_compact( ), assumed_align=16, ) - paged_kv_indptr_fake = fake_compact( + block_tables_fake = cute.runtime.make_fake_tensor( Int32, - (key.batch_size + 1,), - 4, + (key.batch_size, runtime_page_columns), + stride=(runtime_page_row_stride, 1), + assumed_align=4, ) - paged_kv_indices_fake = fake_compact(Int32, (logical_pages,), 4) seq_lens_kv_fake = fake_compact(Int32, (key.batch_size,), 4) tensor_adapter = paged_tensor_adapter dynamic_args = ( @@ -302,8 +553,7 @@ def fake_compact( indptr_fake, indices_fake, valid_bits_fake, - paged_kv_indptr_fake, - paged_kv_indices_fake, + block_tables_fake, seq_lens_kv_fake, row_route_offsets_fake, route_workspace_fake, @@ -311,6 +561,7 @@ def fake_compact( Int64(1), Int64(1), Int64(1), + Int64(1), Float32(1.0), ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/config.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/config.py index 47d7c60340b5..cc4458a4298a 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 @@ -64,7 +64,7 @@ @dataclass(frozen=True) class _BlockSparseCompileKey: - """Named, hashable inputs that determine one compiled adapter.""" + """Named, hashable inputs that determine one compiled sparse adapter.""" device_index: int batch_size: int @@ -81,6 +81,8 @@ class _BlockSparseCompileKey: use_kv_valid_bits: bool use_persistent_scheduler: bool use_parallel_sparse_kv_loads: bool + sparse_format: Literal["bsr", "bitmask"] = "bsr" + use_proxy_routes: bool = False page_size: int | None = None @@ -90,6 +92,9 @@ class _BlockSparseLaunchSpec: policy: tuple[tuple[str, object], ...] compile_key: _BlockSparseCompileKey + # Whether prepared routes carry K32 score-validity words; the decode + # config owns this rule and the plan sizes its route storage from it. + prepares_score_words: bool _CAPACITY_UNSET = object() @@ -194,12 +199,12 @@ def _select_block_sparse_scheduler( use_kv_valid_bits: bool, max_row_route_capacity: int, ) -> tuple[int, bool]: - """Select the Q tile and scheduler without depending on KV storage.""" + """Select the Q tile and scheduler without depending on KV storage. - from ..kernels.fmha_decode.fmha_decode_config import ( - _select_auto_launch_mode, - make_q_tile_geometry, - ) + 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. + """ heads_q_per_kv = num_qo_heads // num_kv_heads q_tile_size = _select_block_sparse_q_tile_size( @@ -216,6 +221,11 @@ def _select_block_sparse_scheduler( ): return q_tile_size, False + from ..kernels.fmha_decode.fmha_decode_config import ( + _select_auto_launch_mode, + make_q_tile_geometry, + ) + q_geometry = make_q_tile_geometry( rows_per_cta=q_tile_size, heads_q_per_kv=heads_q_per_kv, @@ -352,11 +362,13 @@ def _validate_block_sparse_static_profile( kv_block_size=kv_block_size, ) if page_size is not None: + # Validate the paged route geometry with a capacity-free layout; the + # score-word slots do not take part in the page/atom checks. _BlockSparseRouteLayout.create( kv_route_size=kv_route_size, kv_block_size=kv_block_size, page_size=page_size, - has_token_bits=use_kv_valid_bits, + has_token_bits=False, route_metadata_capacity=0, num_rows=1, ) @@ -413,6 +425,8 @@ def _make_block_sparse_config(key: _BlockSparseCompileKey) -> "FmhaDecodeConfig" } if key.use_persistent_scheduler: config_args["use_persistent_scheduler"] = True + if key.use_proxy_routes: + config_args["use_block_sparse_proxy_routes"] = True layout_args: dict[str, object] if key.page_size is None: layout_args = {"qkv_layout": "contiguousKv"} @@ -455,14 +469,17 @@ def _resolve_block_sparse_launch_spec( mask_type: Literal["dense", "causal"], use_kv_valid_bits: bool, max_row_route_capacity: int, + sparse_format: Literal["bsr", "bitmask"] = "bsr", + use_proxy_routes: bool = False, page_size: int | None = None, ) -> _BlockSparseLaunchSpec: """Resolve and cache one validated static or CLC launch. ``max_row_route_capacity`` is a conservative prepared-route bound. Live index values and physical-tail morphology never specialize this cache - entry. If the selected persistent profile is unsupported, retain the valid - static profile instead. + entry. Proxy and exact routes share one scheduler selection. An + unsupported persistent profile falls back to its valid static + counterpart. """ q_tile_size, use_persistent_scheduler = _select_block_sparse_scheduler( @@ -499,10 +516,12 @@ def _resolve_block_sparse_launch_spec( max_row_route_capacity=max_row_route_capacity, use_persistent_scheduler=use_persistent_scheduler, ), + sparse_format=sparse_format, + use_proxy_routes=use_proxy_routes, page_size=page_size, ) try: - _make_block_sparse_config(compile_key) + config = _make_block_sparse_config(compile_key) except ValueError: if not compile_key.use_persistent_scheduler: raise @@ -516,11 +535,15 @@ def _resolve_block_sparse_launch_spec( use_persistent_scheduler=False, ), ) - _make_block_sparse_config(compile_key) + config = _make_block_sparse_config(compile_key) policy_entries: list[tuple[str, object]] = [ ("tile_size_q", q_tile_size), ("tile_size_kv", kv_route_size), + ( + "scheduler", + "persistent" if compile_key.use_persistent_scheduler else "static", + ), ] if page_size is not None: policy_entries.append(("page_size", page_size)) @@ -538,6 +561,7 @@ def _resolve_block_sparse_launch_spec( return _BlockSparseLaunchSpec( policy=tuple(policy_entries), compile_key=compile_key, + prepares_score_words=config.uses_prepared_score_keep_words, ) 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 39c0c22d6f64..7abf22e52792 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 @@ -64,12 +64,7 @@ def _raise_for_invalid_paged_metadata( return reason = { 4: (f"seq_lens_kv values must lie in [{minimum_seq_len_kv}, {max_seq_len_kv}]"), - 5: ( - "paged_kv_indptr must start at zero and each row must be " - "bounded and monotone" - ), - 6: "paged_kv_indptr rows must contain enough pages for seq_lens_kv", - 7: "paged_kv_indices must contain an in-range physical page ID", + 5: "block_tables must contain an in-range physical page ID for every live page", }.get(error_code) if reason is None: reason = ( @@ -152,8 +147,7 @@ def launch(summary: torch.Tensor, device_index: int) -> None: def _inspect_paged_block_sparse_metadata( block_indptr: torch.Tensor, block_indices: torch.Tensor, - paged_kv_indptr: torch.Tensor, - paged_kv_indices: torch.Tensor, + block_tables: torch.Tensor, seq_lens_kv: torch.Tensor, *, static: _BlockSparseStaticProfile, @@ -186,8 +180,8 @@ def launch(summary: torch.Tensor, device_index: int) -> None: inspect_metadata( block_indptr, block_indices, - paged_kv_indptr, - paged_kv_indices, + block_tables, + block_tables.stride(0), seq_lens_kv, num_physical_kv_pages, summary, 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 e70ed415fc97..453e8e35134a 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 @@ -18,13 +18,13 @@ from dataclasses import dataclass import functools import _thread -from typing import Concatenate, ParamSpec, Protocol, TypeVar, cast +from typing import Concatenate, Literal, ParamSpec, Protocol, TypeVar, cast import torch from flashinfer.utils import ceil_div -from .common import _SIGNED_INT32_MAX +from .common import _SIGNED_INT32_MAX, _block_sparse_proxy_summary_geometry from .compiler import _get_compiled_block_sparse from .config import ( _BlockSparseStaticProfile, @@ -85,6 +85,8 @@ class _BlockSparsePlanState: cannot prevent in-place modification or replace graph ownership. """ + sparse_format: Literal["bsr", "bitmask"] + use_proxy_routes: bool device: torch.device batch_size: int seq_len_q: int @@ -93,6 +95,7 @@ class _BlockSparsePlanState: num_kv_heads: int head_dim: int q_block_size: int + kv_block_size: int q_dtype: torch.dtype kv_dtype: torch.dtype output_dtype: torch.dtype @@ -179,27 +182,31 @@ def _build_block_sparse_plan_state( device: torch.device, device_index: int, plan_stream: torch.cuda.Stream, + sparse_format: Literal["bsr", "bitmask"] = "bsr", + use_proxy_routes: bool = False, ) -> _BlockSparsePlanState: - """Build and close one complete state after storage validation.""" + """Build one format- and route-specialized plan atomically.""" assert static.max_blocks_per_row is not None - max_row_route_capacity = ceil_div( - static.max_blocks_per_row * static.kv_block_size, - static.kv_route_size, - ) + 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 * ceil_div(static.seq_len_q, static.q_block_size) ) - 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=static.use_kv_valid_bits, - route_metadata_capacity=num_rows * max_row_route_capacity, - num_rows=num_rows, + if num_rows > _SIGNED_INT32_MAX: + raise OverflowError("row_count must fit in signed int32") + max_row_route_capacity = ceil_div( + static.max_blocks_per_row * static.kv_block_size, + static.kv_route_size, ) + if use_proxy_routes: + num_summaries, _ = _block_sparse_proxy_summary_geometry( + static.seq_len_kv, + static.kv_block_size, + ) + max_row_route_capacity += ceil_div(num_summaries, static.kv_route_size) with torch.cuda.device(device_index), torch.cuda.stream(plan_stream): spec = _resolve_block_sparse_launch_spec( device_index=device_index, @@ -217,11 +224,21 @@ def _build_block_sparse_plan_state( mask_type=static.mask_type, use_kv_valid_bits=static.use_kv_valid_bits, max_row_route_capacity=max_row_route_capacity, + sparse_format=sparse_format, + use_proxy_routes=use_proxy_routes, ) 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 @@ -240,6 +257,8 @@ def _build_block_sparse_plan_state( ready_event = _record_block_sparse_plan_ready_event(plan_stream) return _BlockSparsePlanState( + sparse_format=sparse_format, + use_proxy_routes=use_proxy_routes, device=device, batch_size=static.batch_size, seq_len_q=static.seq_len_q, @@ -248,6 +267,7 @@ def _build_block_sparse_plan_state( num_kv_heads=static.num_kv_heads, head_dim=static.head_dim, q_block_size=static.q_block_size, + kv_block_size=static.kv_block_size, q_dtype=static.q_dtype, kv_dtype=static.kv_dtype, output_dtype=static.output_dtype, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/prepared.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/prepared.py index 091becbd33b8..66815f890a55 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/prepared.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/prepared.py @@ -22,6 +22,7 @@ _SECTION_ALIGNMENT_WORDS = 4 _PREPARED_ROUTE_IS_FULL_FLAG = 1 << 0 +_PREPARED_ROUTE_IS_PROXY_FLAG = 1 << 1 _SUPPORTED_KV_ROUTE_SIZES = (128, 256) _SUPPORTED_PAGED_KV_PAGE_SIZES = (16, 32, 64, 128) @@ -91,13 +92,16 @@ class _BlockSparseRouteLayout: Each route's metadata stores logical KV-token atom origins, optional physical page IDs, one atom-valid-mask word, one route-flags word, and - optional token-valid words. ``page_size is None`` selects the contiguous + optional token-mask words. ``page_size is None`` selects the contiguous record; otherwise the paged record adds one page-ID word per logical origin. Logical origins remain independent of the K/V storage locator used - by the attention load path. An invalid logical origin is encoded as - ``-1``. Bit ``i`` of the atom-valid mask corresponds to logical origin - ``i``. ``_PREPARED_ROUTE_IS_FULL_FLAG`` (bit 0) states that the route is - both structurally full and, when token bits are present, token-full. + by the attention load path. Exact routes address raw-token origins, while + proxy routes address summary-token origins and set + ``_PREPARED_ROUTE_IS_PROXY_FLAG`` (bit 1). An invalid logical origin is + encoded as ``-1``. Bit ``i`` of the atom-valid mask corresponds to logical + origin ``i``. ``_PREPARED_ROUTE_IS_FULL_FLAG`` (bit 0) states that the + route is both structurally full and, when token-mask bits are present, + mask-full. """ # Store semantic inputs plus the three validated allocation values. All @@ -238,6 +242,17 @@ def token_words_word_offset(self) -> int | None: return self.route_flags_word_offset + 1 if self.has_token_bits else None + @property + def uses_one_warp_transport(self) -> bool: + """Whether this layout uses the continuous one-warp transport.""" + + token_words_word_offset = self.token_words_word_offset + return ( + not self.is_paged + and token_words_word_offset is not None + and token_words_word_offset + self.token_words_per_route <= 32 + ) + @property def route_metadata_capacity(self) -> int: """Number of routes whose metadata fits in the mutable workspace.""" @@ -249,5 +264,6 @@ def route_metadata_capacity(self) -> int: __all__ = [ "_PREPARED_ROUTE_IS_FULL_FLAG", + "_PREPARED_ROUTE_IS_PROXY_FLAG", "_BlockSparseRouteLayout", ] 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 43e8ff439d31..4f24524e3d4a 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 @@ -16,7 +16,7 @@ from dataclasses import dataclass import math -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Literal import torch @@ -25,6 +25,7 @@ PagedKVCache, _normalize_paged_kv_cache, _validate_16byte_alignment, + _validate_block_table_metadata, _validate_exact_compact_strides, _validate_scale, ) @@ -47,8 +48,7 @@ class _PagedKVStorage: """Paged K/V storage and request metadata consumed by one live run.""" paged_kv_cache: PagedKVCache - paged_kv_indptr: torch.Tensor - paged_kv_indices: torch.Tensor + block_tables: torch.Tensor seq_lens_kv: torch.Tensor @@ -56,8 +56,8 @@ class _PagedKVStorage: class _PagedKVLaunchPayload: """Launch-only live paged metadata derived during shared validation.""" - paged_kv_indptr: torch.Tensor - paged_kv_indices: torch.Tensor + block_tables: torch.Tensor + block_table_row_stride: int seq_lens_kv: torch.Tensor num_physical_kv_pages: int k_page_stride: int @@ -72,12 +72,15 @@ class _BlockSparseRunArgs: k: torch.Tensor v: torch.Tensor out: torch.Tensor - block_indptr: torch.Tensor - block_indices: torch.Tensor + block_indptr: torch.Tensor | None + block_indices: torch.Tensor | None kv_valid_bits: torch.Tensor 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 def _validate_metadata_tensor( @@ -113,36 +116,71 @@ def _validate_metadata_tensor( def validate_block_sparse_metadata( - block_indptr: torch.Tensor, - block_indices: torch.Tensor, - kv_valid_bits: torch.Tensor | None, *, + sparse_format: Literal["bsr", "bitmask"], + block_indptr: torch.Tensor | None, + block_indices: torch.Tensor | None, + exact_block_bits: torch.Tensor | None, + kv_valid_bits: torch.Tensor | None, device: torch.device, batch_size: int, seq_len_q: int, seq_len_kv: int, num_kv_heads: int, q_block_size: int, + kv_block_size: int, use_kv_valid_bits: bool, ) -> None: - """Validate raw runtime routing without reading device-side values.""" + """Validate the planned route frontend without reading tensor values.""" num_q_blocks = (seq_len_q + q_block_size - 1) // q_block_size - _validate_metadata_tensor( - block_indptr, - "block_indptr", - ndim=3, - dtype=torch.int32, - expected_device=device, - expected_shape=(batch_size, num_kv_heads, num_q_blocks + 1), - ) - _validate_metadata_tensor( - block_indices, - "block_indices", - ndim=1, - dtype=torch.int32, - expected_device=device, - ) + if sparse_format == "bsr": + if block_indptr is None or block_indices is None: + raise ValueError( + "block_indptr and block_indices are required by a BSR plan" + ) + if exact_block_bits is not None: + raise ValueError("exact_block_bits is valid only for a bitmask plan") + _validate_metadata_tensor( + block_indptr, + "block_indptr", + ndim=3, + dtype=torch.int32, + expected_device=device, + expected_shape=(batch_size, num_kv_heads, num_q_blocks + 1), + ) + _validate_metadata_tensor( + block_indices, + "block_indices", + ndim=1, + dtype=torch.int32, + expected_device=device, + ) + elif sparse_format == "bitmask": + if ( + exact_block_bits is None + or block_indptr is not None + or block_indices is not None + ): + raise ValueError( + "runtime route inputs must match planned sparse_format='bitmask'" + ) + num_kv_blocks = (seq_len_kv + kv_block_size - 1) // kv_block_size + _validate_metadata_tensor( + exact_block_bits, + "exact_block_bits", + ndim=4, + dtype=torch.uint32, + expected_device=device, + expected_shape=( + batch_size, + num_kv_heads, + num_q_blocks, + (num_kv_blocks + 31) // 32, + ), + ) + else: + raise AssertionError(f"unsupported sparse format {sparse_format!r}") if use_kv_valid_bits: if kv_valid_bits is None: @@ -159,31 +197,6 @@ def validate_block_sparse_metadata( raise ValueError("kv_valid_bits must be None when use_kv_valid_bits=False") -def validate_paged_kv_metadata( - paged_kv_indptr: torch.Tensor, - paged_kv_indices: torch.Tensor, - seq_lens_kv: torch.Tensor, - *, - device: torch.device, - batch_size: int, -) -> None: - """Validate the shared structural ABI for live paged request metadata.""" - - for tensor, name, shape in ( - (paged_kv_indptr, "paged_kv_indptr", (batch_size + 1,)), - (paged_kv_indices, "paged_kv_indices", None), - (seq_lens_kv, "seq_lens_kv", (batch_size,)), - ): - _validate_metadata_tensor( - tensor, - name, - ndim=1, - dtype=torch.int32, - expected_device=device, - expected_shape=shape, - ) - - def _validate_bshd_tensor( tensor: torch.Tensor, name: str, @@ -214,8 +227,11 @@ def validate_block_sparse_run( kv_storage: _ContiguousKVStorage | _PagedKVStorage, *, state: "_BlockSparsePlanState", - block_indptr: torch.Tensor, - block_indices: torch.Tensor, + block_indptr: torch.Tensor | None, + block_indices: torch.Tensor | None, + exact_block_bits: torch.Tensor | None = None, + k_summary: torch.Tensor | None = None, + v_summary: torch.Tensor | None = None, kv_valid_bits: torch.Tensor | None, sm_scale: float | None, out: torch.Tensor | None, @@ -229,18 +245,44 @@ def validate_block_sparse_run( input. ``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( - block_indptr, - block_indices, - kv_valid_bits, + 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, + ) + 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, + ) + elif k_summary is not None or v_summary is not None: + raise ValueError("summaries are valid only when proxy routes are enabled") + if state.use_kv_valid_bits: assert kv_valid_bits is not None effective_kv_valid_bits = kv_valid_bits @@ -258,7 +300,7 @@ def validate_block_sparse_run( expected_device=state.device, ) paged_kv: _PagedKVLaunchPayload | None = None - overlap_inputs: tuple[tuple[str, torch.Tensor], ...] + 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") @@ -278,16 +320,11 @@ def validate_block_sparse_run( ) k = kv_storage.k v = kv_storage.v - overlap_inputs = ( + overlap_inputs = [ ("q", q), ("k", k), ("v", v), - ("block_indptr", block_indptr), - ("block_indices", block_indices), - ("kv_valid_bits", effective_kv_valid_bits), - ("row_route_offsets", state.row_route_offsets), - ("route_workspace", state.route_workspace), - ) + ] elif isinstance(kv_storage, _PagedKVStorage): page_size = state.page_size if page_size is None: @@ -324,36 +361,60 @@ def validate_block_sparse_run( raise ValueError( f"K/V dtype must match the plan ({state.kv_dtype}), got {k.dtype}" ) - validate_paged_kv_metadata( - kv_storage.paged_kv_indptr, - kv_storage.paged_kv_indices, - kv_storage.seq_lens_kv, - device=state.device, - batch_size=state.batch_size, + 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: + 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}" + ) paged_kv = _PagedKVLaunchPayload( - paged_kv_indptr=kv_storage.paged_kv_indptr, - paged_kv_indices=kv_storage.paged_kv_indices, + 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 = ( + overlap_inputs = [ ("q", q), ("k_cache", k), ("v_cache", v), - ("block_indptr", block_indptr), - ("block_indices", block_indices), - ("kv_valid_bits", effective_kv_valid_bits), - ("paged_kv_indptr", kv_storage.paged_kv_indptr), - ("paged_kv_indices", kv_storage.paged_kv_indices), + ("block_tables", kv_storage.block_tables), ("seq_lens_kv", kv_storage.seq_lens_kv), + ] + 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), ) - else: - raise TypeError("kv_storage must be _ContiguousKVStorage or _PagedKVStorage") + ) effective_scale = _validate_scale( 1.0 / math.sqrt(state.head_dim) if sm_scale is None else sm_scale, @@ -377,6 +438,9 @@ def validate_block_sparse_run( out=out, 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=effective_kv_valid_bits, kv_valid_bits_is_live=state.use_kv_valid_bits, sm_scale=effective_scale, @@ -384,23 +448,100 @@ def validate_block_sparse_run( ) +def prepare_block_sparse_run_unchecked( + q: torch.Tensor, + kv_storage: _ContiguousKVStorage | _PagedKVStorage, + *, + state: "_BlockSparsePlanState", + block_indptr: torch.Tensor | None, + block_indices: torch.Tensor | None, + exact_block_bits: torch.Tensor | None = None, + k_summary: torch.Tensor | None = None, + v_summary: torch.Tensor | None = None, + kv_valid_bits: torch.Tensor | None, + sm_scale: float | None, + out: torch.Tensor | None, +) -> _BlockSparseRunArgs: + """Canonicalize one trusted run without invoking explicit validators. + + Only the work every launch needs happens here: K/V view selection, the + plan-owned dummy token mask when the plan disabled token bits, the default + softmax scale, and allocation of an omitted output tensor. + """ + + paged_kv: _PagedKVLaunchPayload | None = None + if isinstance(kv_storage, _ContiguousKVStorage): + k = kv_storage.k + v = kv_storage.v + else: + paged_kv_cache = kv_storage.paged_kv_cache + if isinstance(paged_kv_cache, torch.Tensor): + k = paged_kv_cache[:, 0] + v = paged_kv_cache[:, 1] + else: + k, v = paged_kv_cache + 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=int(k.shape[0]), + k_page_stride=int(k.stride(0)), + v_page_stride=int(v.stride(0)), + ) + if state.use_kv_valid_bits: + 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 out is None: + out = torch.empty( + (state.batch_size, state.seq_len_q, state.num_qo_heads, state.head_dim), + device=state.device, + dtype=state.output_dtype, + ) + return _BlockSparseRunArgs( + q=q, + k=k, + v=v, + out=out, + 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=effective_kv_valid_bits, + kv_valid_bits_is_live=state.use_kv_valid_bits, + sm_scale=1.0 / math.sqrt(state.head_dim) + if sm_scale is None + else float(sm_scale), + paged_kv=paged_kv, + ) + + def record_block_sparse_run_args( run_args: _BlockSparseRunArgs, stream: torch.cuda.Stream, ) -> None: """Extend tensor lifetimes for the asynchronous launch currently in flight.""" - run_args.q.record_stream(stream) - run_args.k.record_stream(stream) - run_args.v.record_stream(stream) + for tensor in (run_args.q, run_args.k, run_args.v): + tensor.record_stream(stream) + if run_args.k_summary is not None: + run_args.k_summary.record_stream(stream) + assert run_args.v_summary is not None + run_args.v_summary.record_stream(stream) run_args.out.record_stream(stream) - run_args.block_indptr.record_stream(stream) - run_args.block_indices.record_stream(stream) + if run_args.block_indptr is not None: + 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 + run_args.exact_block_bits.record_stream(stream) if run_args.kv_valid_bits_is_live: run_args.kv_valid_bits.record_stream(stream) if run_args.paged_kv is not None: - run_args.paged_kv.paged_kv_indptr.record_stream(stream) - run_args.paged_kv.paged_kv_indices.record_stream(stream) + run_args.paged_kv.block_tables.record_stream(stream) run_args.paged_kv.seq_lens_kv.record_stream(stream) @@ -409,9 +550,13 @@ def launch_block_sparse( *, state: "_BlockSparsePlanState", ) -> torch.Tensor: - """Invoke the exact contiguous or paged ABI chosen by validated payload.""" + """Invoke the layout- and route-specific ABI chosen by the frozen plan.""" - if run_args.paged_kv is 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 state.compiled( run_args.q, run_args.k, @@ -420,12 +565,20 @@ def launch_block_sparse( run_args.block_indptr, run_args.block_indices, run_args.kv_valid_bits, + run_args.paged_kv.block_tables, + run_args.paged_kv.seq_lens_kv, state.row_route_offsets, state.route_workspace, state.max_blocks_per_row, + run_args.paged_kv.num_physical_kv_pages, + run_args.paged_kv.block_table_row_stride, + run_args.paged_kv.k_page_stride, + run_args.paged_kv.v_page_stride, run_args.sm_scale, ) - else: + 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, @@ -434,17 +587,65 @@ def launch_block_sparse( run_args.block_indptr, run_args.block_indices, run_args.kv_valid_bits, - run_args.paged_kv.paged_kv_indptr, - run_args.paged_kv.paged_kv_indices, - run_args.paged_kv.seq_lens_kv, state.row_route_offsets, state.route_workspace, state.max_blocks_per_row, - run_args.paged_kv.num_physical_kv_pages, - run_args.paged_kv.k_page_stride, - run_args.paged_kv.v_page_stride, 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 + state.compiled( + run_args.q, + run_args.k, + run_args.v, + run_args.k_summary, + run_args.v_summary, + 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, + run_args.sm_scale, + ) + else: + raise AssertionError("frozen block-sparse plan has an unsupported route mode") return run_args.out @@ -454,8 +655,8 @@ def launch_block_sparse( "_PagedKVLaunchPayload", "_PagedKVStorage", "launch_block_sparse", + "prepare_block_sparse_run_unchecked", "record_block_sparse_run_args", "validate_block_sparse_metadata", "validate_block_sparse_run", - "validate_paged_kv_metadata", ] 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 af132c6a69ce..e19dc7b1de7f 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/block_sparse.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/block_sparse.py @@ -37,6 +37,7 @@ from flashinfer.api_logging import flashinfer_api +from ._block_sparse.common import _validate_contiguous_route_mode from ._block_sparse.config import _validate_block_sparse_static_profile from ._block_sparse.inspection import ( _inspect_block_sparse_bsr, @@ -53,15 +54,16 @@ _ContiguousKVStorage, _PagedKVStorage, launch_block_sparse as _launch_block_sparse, + prepare_block_sparse_run_unchecked as _prepare_block_sparse_run_unchecked, 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_metadata as _validate_paged_kv_metadata, ) from .decode import ( PagedKVCache, _normalize_paged_kv_cache, _resolve_cuda_device, + _validate_block_table_metadata, ) @@ -111,10 +113,12 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase): Q is ``[B, Sq, Hq, D]`` and K/V are ``[B, Skv, Hkv, D]``. Sparse rows are owned per batch, KV head, and query block, so every Q head in one grouped KV head consumes the same sparse row. A plan fixes geometry and a per-row - capacity; every run supplies its own BSR and optional token mask. - Callers must keep those tensors alive and immutable until the queued run or - captured graph finishes using them. CUDA Graph capture pins plan-owned - state only, so captured routing storage remains the caller's responsibility. + capacity; every run supplies either BSR or a packed exact-block bitmask. + Proxy-enabled plans additionally consume caller-owned K/V summaries, while + an optional token mask applies only to exact routes. Callers must keep those + tensors alive and immutable until the queued run or captured graph finishes + using them. CUDA Graph capture pins plan-owned state only, so captured + routing storage remains the caller's responsibility. One plan revision owns one mutable route workspace. Its runs must be ordered on one stream or externally synchronized; unordered concurrent runs require @@ -136,6 +140,8 @@ def plan( device: torch.device | str | int, max_blocks_per_row: int, use_kv_valid_bits: bool, + sparse_format: Literal["bsr", "bitmask"] = "bsr", + use_proxy_routes: bool = False, mask_type: Literal["dense", "causal"] = "dense", q_data_type: torch.dtype = torch.float16, kv_data_type: torch.dtype | None = None, @@ -144,11 +150,16 @@ def plan( """Choose a legal profile and allocate reusable routing capacity. The plan owns immutable geometry and a uniform route workspace, not a - sparse pattern. ``max_blocks_per_row`` bounds each runtime BSR row in - semantic ``kv_block_size`` blocks. ``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. + sparse pattern. ``max_blocks_per_row`` bounds each runtime sparse row in + 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 + ``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. 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 @@ -162,10 +173,13 @@ def plan( respectively. ``kv_block_size`` may be 8, 16, 32, or a positive 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. Every run prepares - per-KV-head canonical BSR 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; + 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 + 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. @@ -182,6 +196,7 @@ def plan( runs require distinct wrappers. """ + _validate_contiguous_route_mode(sparse_format, use_proxy_routes) static = _validate_block_sparse_static_profile( batch_size=batch_size, seq_len_q=seq_len_q, @@ -198,6 +213,8 @@ def plan( output_dtype=o_data_type, max_blocks_per_row=max_blocks_per_row, ) + 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) plan_stream = torch.cuda.current_stream(device) with torch.cuda.device(device_index), torch.cuda.stream(plan_stream): @@ -210,6 +227,8 @@ def plan( device=device, device_index=device_index, plan_stream=plan_stream, + sparse_format=sparse_format, + use_proxy_routes=use_proxy_routes, ) # This is the only wrapper mutation. Every failure above leaves the # previously published revision intact and runnable. @@ -221,12 +240,16 @@ def run( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - block_indptr: torch.Tensor, - block_indices: torch.Tensor, + block_indptr: torch.Tensor | None = None, + block_indices: torch.Tensor | None = None, *, + exact_block_bits: torch.Tensor | None = None, + k_summary: torch.Tensor | None = None, + v_summary: torch.Tensor | None = None, kv_valid_bits: torch.Tensor | None = None, sm_scale: float | None = None, out: torch.Tensor | None = None, + validate: bool = True, ) -> torch.Tensor: """Launch the current plan on the caller's current CUDA stream. @@ -238,17 +261,34 @@ 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. - ``block_indptr`` is compact Int32 - ``[B, Hkv, ceil(Sq / q_block_size) + 1]`` and indexes compact - ``block_indices``. Every row must fit the planned semantic-block - capacity; referenced block IDs must be strictly increasing, unique, - and in range. Reusable runs trust these device-side values. CuTe DSL - assertions can diagnose violations when enabled before compilation; - otherwise invalid values have undefined behavior and may access out of - bounds. A masked plan requires + ``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 BSR plan consumes compact Int32 ``block_indptr`` with shape + ``[B, Hkv, ceil(Sq / q_block_size) + 1]`` and compact Int32 + ``block_indices``. A bitmask plan instead requires both BSR arguments + to be ``None`` and consumes packed UInt32 ``exact_block_bits`` with + shape ``[B, Hkv, ceil(Sq / q_block_size), ceil(num_kv_blocks / 32)]``. + 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. + + Every row must fit the planned semantic-block capacity. Reusable runs + trust routing values. CuTe DSL assertions can diagnose violations when + enabled before compilation; otherwise invalid values have undefined + behavior and may access out of bounds. A masked plan requires ``kv_valid_bits`` with shape ``[B, ceil(Skv / 32)]`` and dtype UInt32; an unmasked plan requires - ``None``. Routing tensors may have different identities on every run. + ``None``. The mask applies only to raw exact routes; proxy summaries and + their represented-token mass remain caller-defined. Routing tensors may + have different identities on every run. Keep this wrapper alive until every captured CUDA Graph is destroyed. @@ -261,12 +301,18 @@ def run( v : torch.Tensor Compact value tensor with the same shape, dtype, and strides as ``k``. - block_indptr : torch.Tensor + block_indptr : torch.Tensor, optional Contiguous Int32 BSR row offsets with shape - ``[B, Hkv, ceil(Sq / q_block_size) + 1]``. - block_indices : torch.Tensor + ``[B, Hkv, ceil(Sq / q_block_size) + 1]``. Required by BSR plans. + block_indices : torch.Tensor, optional Contiguous Int32 semantic KV-block IDs referenced by - ``block_indptr``. + ``block_indptr``. Required by BSR plans. + exact_block_bits : torch.Tensor, optional + Compact packed UInt32 exact-block bitmap required by bitmask plans. + 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. 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. @@ -275,6 +321,9 @@ def run( out : torch.Tensor, optional Caller-owned compact output buffer ``[B, Sq, Hq, D]`` with the planned output dtype. + validate : bool + Whether to validate tensor structure, plan geometry, and aliasing + before launching. Defaults to ``True``. Returns ------- @@ -283,17 +332,27 @@ def run( """ state = self._require_run_state() - run_stream = torch.cuda.current_stream(state.device) - run_args = _validate_block_sparse_run( + if not isinstance(validate, bool): + raise TypeError("validate must be a bool") + prepare_run = ( + _validate_block_sparse_run + if validate + else _prepare_block_sparse_run_unchecked + ) + run_args = prepare_run( q, _ContiguousKVStorage(k=k, v=v), state=state, 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, sm_scale=sm_scale, out=out, ) + run_stream = torch.cuda.current_stream(state.device) return self._launch_validated_run(state, run_args, run_stream) @@ -302,23 +361,29 @@ def block_sparse_attention( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - block_indptr: torch.Tensor, - block_indices: torch.Tensor, + block_indptr: torch.Tensor | None, + block_indices: torch.Tensor | None, q_block_size: int, kv_block_size: int, *, + exact_block_bits: torch.Tensor | None = None, + k_summary: torch.Tensor | None = None, + v_summary: torch.Tensor | None = None, kv_valid_bits: torch.Tensor | None = None, + sparse_format: Literal["bsr", "bitmask"] = "bsr", + use_proxy_routes: bool = False, mask_type: Literal["dense", "causal"] = "dense", sm_scale: float | None = None, out: torch.Tensor | None = None, ) -> torch.Tensor: """Plan and run one compact-BSHD block-sparse attention launch. - This one-shot form synchronously inspects canonical BSR, derives its largest - semantic row, creates a capacity-only plan, and passes the original routing - tensors to :meth:`BlockSparseTSWrapper.run`. It therefore cannot be invoked - inside CUDA Graph capture; plan a wrapper outside capture and capture only - ``run()`` instead. + This one-shot form creates a capacity-only plan and passes the original + routing tensors to :meth:`BlockSparseTSWrapper.run`. BSR inputs are + synchronously inspected to validate canonical rows and derive their maximum + width. Bitmask inputs use the structural KV-block count as a conservative + capacity bound. It therefore cannot be invoked inside CUDA Graph capture; + plan a wrapper outside capture and capture only ``run()`` instead. Parameters ---------- @@ -328,11 +393,12 @@ def block_sparse_attention( Compact key tensor ``[B, Skv, Hkv, D]``. v : torch.Tensor Compact value tensor with the same shape, dtype, and strides as ``k``. - block_indptr : torch.Tensor + block_indptr : torch.Tensor, optional Contiguous Int32 BSR row offsets with shape - ``[B, Hkv, ceil(Sq / q_block_size) + 1]``. - block_indices : torch.Tensor + ``[B, Hkv, ceil(Sq / q_block_size) + 1]``. Required in BSR mode. + block_indices : torch.Tensor, optional Contiguous Int32 semantic KV-block IDs referenced by ``block_indptr``. + Required in BSR mode. q_block_size : int Positive number of logical query tokens represented by one BSR row. The product with ``Hq / Hkv`` must be divisible by 8 so a physical Q @@ -340,8 +406,19 @@ def block_sparse_attention( kv_block_size : int Number of logical KV tokens represented by one BSR block ID; it must be 8, 16, 32, or a positive multiple of 64. + exact_block_bits : torch.Tensor, optional + Compact UInt32 exact-block bitmap required in bitmask mode. + 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. kv_valid_bits : torch.Tensor, optional Contiguous UInt32 token-validity bitmap ``[B, ceil(Skv / 32)]``. + sparse_format : {"bsr", "bitmask"}, optional + Runtime sparse representation. Defaults to ``"bsr"``. + use_proxy_routes : bool, optional + Represent unselected blocks through K/V summaries. Proxy routes + currently require dense masking. mask_type : {"dense", "causal"}, optional Attention mask applied inside each selected sparse block. sm_scale : float, optional @@ -370,6 +447,7 @@ def block_sparse_attention( raise ValueError("K and V must have identical shapes") use_kv_valid_bits = kv_valid_bits is not None + _validate_contiguous_route_mode(sparse_format, use_proxy_routes) static = _validate_block_sparse_static_profile( batch_size=batch_size, seq_len_q=seq_len_q, @@ -385,25 +463,37 @@ def block_sparse_attention( kv_dtype=k.dtype, output_dtype=q.dtype if out is None else out.dtype, ) + 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( - block_indptr, - block_indices, - kv_valid_bits, + 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, ) - max_blocks_per_row = _inspect_block_sparse_bsr( - block_indptr, - block_indices, - static=static, - stream=torch.cuda.current_stream(device), - ) + 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( @@ -418,6 +508,8 @@ def block_sparse_attention( device=device, 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, @@ -429,6 +521,9 @@ def block_sparse_attention( v, block_indptr, block_indices, + exact_block_bits=exact_block_bits, + k_summary=k_summary, + v_summary=v_summary, kv_valid_bits=kv_valid_bits, sm_scale=sm_scale, out=out, @@ -512,8 +607,7 @@ def run( self, q: torch.Tensor, paged_kv_cache: PagedKVCache, - paged_kv_indptr: torch.Tensor, - paged_kv_indices: torch.Tensor, + block_tables: torch.Tensor, seq_lens_kv: torch.Tensor, block_indptr: torch.Tensor, block_indices: torch.Tensor, @@ -521,6 +615,7 @@ def run( kv_valid_bits: torch.Tensor | None = None, sm_scale: float | None = None, out: torch.Tensor | None = None, + validate: bool = True, ) -> torch.Tensor: """Launch with live lengths, page tables, and sparse routes. @@ -529,15 +624,20 @@ def run( tuple whose members are ``[P, Hkv, page, D]`` with compact inner HND strides and arbitrary non-overlapping outer page strides. - ``paged_kv_indptr`` is compact Int32 ``[B + 1]``; - ``paged_kv_indices`` is compact Int32 with capacity at least its live - final offset; and ``seq_lens_kv`` is compact Int32 ``[B]``. All values - are read on device. The caller must keep every dense length in - ``[1, max_seq_len_kv]`` and every causal length in - ``[Sq, max_seq_len_kv]``. ``paged_kv_indptr`` must start at zero and - contain bounded, monotone rows with at least - ``ceil(seq_lens_kv[b] / page_size)`` entries. Every page ID in its live - prefix must lie in ``[0, P)``. Each BSR row must contain strictly + ``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. + + ``block_tables`` is Int32 ``[B, C]``, contiguous within each row but + permitted to use a padded outer row stride; ``seq_lens_kv`` is compact + Int32 ``[B]``. All values are read on device. The caller must keep + every dense length in ``[1, max_seq_len_kv]`` and every causal length + in ``[Sq, max_seq_len_kv]``. Every page-table row must contain at least + ``ceil(seq_lens_kv[b] / page_size)`` live entries. Every page ID in its + live prefix must lie in ``[0, P)``. Each BSR row must contain strictly increasing, unique block IDs whose final block starts before that request's live K/V length, and its width must not exceed the planned ``max_blocks_per_row``. Reusable runs trust all of these device-side @@ -564,10 +664,10 @@ def run( Either a combined cache ``[P, 2, Hkv, page_size, D]`` or a ``(K, V)`` tuple whose tensors are ``[P, Hkv, page_size, D]``. - paged_kv_indptr : torch.Tensor - Contiguous Int32 live request offsets with shape ``[B + 1]``. - paged_kv_indices : torch.Tensor - Contiguous Int32 physical-page ID capacity. + block_tables : torch.Tensor + Live Int32 physical page IDs with shape ``[B, C]``. Entries are + contiguous within each row; padded, non-overlapping row strides are + supported and inactive tail entries are ignored. seq_lens_kv : torch.Tensor Contiguous Int32 live logical K/V lengths with shape ``[B]``. Values must satisfy the dense or causal bounds above. @@ -586,6 +686,9 @@ def run( out : torch.Tensor, optional Caller-owned compact output buffer ``[B, Sq, Hq, D]`` with the planned output dtype. + validate : bool + Whether to validate tensor structure, plan geometry, and aliasing + before launching. Defaults to ``True``. Returns ------- @@ -594,13 +697,19 @@ def run( """ state = self._require_run_state() + if not isinstance(validate, bool): + raise TypeError("validate must be a bool") run_stream = torch.cuda.current_stream(state.device) - run_args = _validate_block_sparse_run( + prepare_run = ( + _validate_block_sparse_run + if validate + else _prepare_block_sparse_run_unchecked + ) + run_args = prepare_run( q, _PagedKVStorage( paged_kv_cache=paged_kv_cache, - paged_kv_indptr=paged_kv_indptr, - paged_kv_indices=paged_kv_indices, + block_tables=block_tables, seq_lens_kv=seq_lens_kv, ), state=state, @@ -617,15 +726,14 @@ def run( def block_sparse_attention_with_paged_kv_cache( q: torch.Tensor, paged_kv_cache: PagedKVCache, - paged_kv_indptr: torch.Tensor, - paged_kv_indices: torch.Tensor, + block_tables: torch.Tensor, + seq_lens_kv: torch.Tensor, block_indptr: torch.Tensor, block_indices: torch.Tensor, q_block_size: int, kv_block_size: int, *, max_seq_len_kv: int, - seq_lens_kv: torch.Tensor, kv_valid_bits: torch.Tensor | None = None, mask_type: Literal["dense", "causal"] = "dense", sm_scale: float | None = None, @@ -633,12 +741,11 @@ def block_sparse_attention_with_paged_kv_cache( ) -> torch.Tensor: """Plan and run one fixed-Q paged block-sparse attention launch. - This convenience entry point synchronously validates live page and sparse - metadata, including the complete live physical-page-ID prefix, creates a - capacity-only temporary plan, then forwards the inspected tensors through - the trusted live run API. It cannot run during CUDA Graph capture; plan a - wrapper outside capture and capture only - :meth:`BlockSparsePagedTSWrapper.run` instead. + This convenience entry point synchronously validates the live page tables, + K/V lengths, and sparse metadata, creates a capacity-only temporary plan, + then forwards the inspected tensors through the trusted live run API. It + cannot run during CUDA Graph capture; plan a wrapper outside capture and + capture only :meth:`BlockSparsePagedTSWrapper.run` instead. Parameters ---------- @@ -647,11 +754,13 @@ def block_sparse_attention_with_paged_kv_cache( paged_kv_cache : PagedKVCache Either a combined cache ``[P, 2, Hkv, page_size, D]`` or a ``(K, V)`` tuple whose tensors are ``[P, Hkv, page_size, D]``. - paged_kv_indptr : torch.Tensor - Contiguous Int32 request offsets into ``paged_kv_indices``, with shape - ``[B + 1]``. - paged_kv_indices : torch.Tensor - Contiguous Int32 physical page IDs referenced by ``paged_kv_indptr``. + block_tables : torch.Tensor + Int32 physical page IDs ``[B, C]``, contiguous within each row and free + to use a padded outer row stride. ``C * page_size`` must cover + ``max_seq_len_kv``; only the first ``ceil(seq_lens_kv[b] / page_size)`` + entries of each row are read. + seq_lens_kv : torch.Tensor + 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]``. @@ -666,8 +775,6 @@ def block_sparse_attention_with_paged_kv_cache( be 8, 16, 32, or a positive multiple of 64. max_seq_len_kv : int Static maximum logical K/V length used for planning. - seq_lens_kv : torch.Tensor - Contiguous Int32 per-request logical KV lengths with shape ``[B]``. kv_valid_bits : torch.Tensor, optional Contiguous UInt32 logical-token validity bitmap ``[B, ceil(max_seq_len_kv / 32)]``. @@ -693,13 +800,16 @@ 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) - _validate_paged_kv_metadata( - paged_kv_indptr, - paged_kv_indices, - seq_lens_kv, - device=metadata_device, - batch_size=batch_size, + table_device, table_batch_size, table_capacity = _validate_block_table_metadata( + block_tables, seq_lens_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, @@ -734,23 +844,31 @@ def block_sparse_attention_with_paged_kv_cache( kv_dtype=k_cache.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( - block_indptr, - block_indices, - kv_valid_bits, + sparse_format="bsr", + block_indptr=block_indptr, + block_indices=block_indices, + exact_block_bits=None, + kv_valid_bits=kv_valid_bits, device=metadata_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, ) max_blocks_per_row = _inspect_paged_block_sparse_metadata( block_indptr, block_indices, - paged_kv_indptr, - paged_kv_indices, + block_tables, seq_lens_kv, static=static, num_physical_kv_pages=num_physical_kv_pages, @@ -780,8 +898,7 @@ def block_sparse_attention_with_paged_kv_cache( return wrapper.run( q, paged_kv_cache, - paged_kv_indptr, - paged_kv_indices, + block_tables, seq_lens_kv, block_indptr, block_indices, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_inspect.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_inspect.py index de53ccb9c7de..5728a7c6d373 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_inspect.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_inspect.py @@ -20,10 +20,10 @@ their caller. Token-mask contents belong to the run-time prepare kernel and are not read here. -Paged inspection first validates live sequence lengths and page rows, then -four warps validate four BSR Q-block rows per CTA. Both publish one validation -status plus the maximum row width in one Int64 summary; no route payload is -constructed. +Paged inspection first validates live sequence lengths and the live prefix of +every page-table row, then four warps validate four BSR Q-block rows per CTA. +Both publish one validation status plus the maximum row width in one Int64 +summary; no route payload is constructed. """ import functools @@ -51,9 +51,7 @@ _BSR_ERROR_INDEX_OUT_OF_RANGE = 2 _BSR_ERROR_INVALID_INDPTR = 3 _ERROR_INVALID_SEQ_LEN = 4 -_ERROR_INVALID_PAGE_INDPTR = 5 -_ERROR_INSUFFICIENT_PAGE_CAPACITY = 6 -_ERROR_INVALID_PHYSICAL_PAGE_ID = 7 +_ERROR_INVALID_PHYSICAL_PAGE_ID = 5 @cute.jit @@ -296,7 +294,7 @@ def __call__( class _InspectPagedKvMetadata: - """Validate live lengths and page rows with one warp per request.""" + """Validate live lengths and page-table rows with one warp per request.""" def __init__( self, @@ -314,16 +312,16 @@ def __init__( @cute.jit def __call__( self, - paged_kv_indptr: cute.Tensor, - paged_kv_indices: cute.Tensor, + block_tables: cute.Tensor, + block_table_row_stride: cutlass.Int64, seq_lens_kv: cute.Tensor, num_physical_kv_pages: cutlass.Int64, summary: cute.Tensor, stream: cuda_drv.CUstream, ) -> None: self.kernel( - paged_kv_indptr, - paged_kv_indices, + block_tables, + block_table_row_stride, seq_lens_kv, num_physical_kv_pages, summary, @@ -340,8 +338,8 @@ def __call__( @cute.kernel def kernel( self, - paged_kv_indptr: cute.Tensor, - paged_kv_indices: cute.Tensor, + block_tables: cute.Tensor, + block_table_row_stride: cutlass.Int64, seq_lens_kv: cute.Tensor, num_physical_kv_pages: cutlass.Int64, summary: cute.Tensor, @@ -353,9 +351,7 @@ def kernel( batch_idx = block_idx * _WARPS_PER_CTA + warp_idx request_is_valid = batch_idx < self.batch_size - request_begin = cutlass.Int32(0) - request_end = cutlass.Int32(0) - request_range_is_valid = cutlass.Int32(0) + live_pages = cutlass.Int32(0) error_code = cutlass.Int32(_BSR_ERROR_NONE) if lane_idx == 0 and request_is_valid: seq_len_kv = cutlass.Int32(seq_lens_kv[batch_idx]) @@ -363,38 +359,20 @@ def kernel( seq_len_kv >= cutlass.Int32(self.minimum_seq_len_kv) and seq_len_kv <= cutlass.Int32(self.max_seq_len_kv) ) - if not seq_len_is_valid: - error_code = cutlass.Int32(_ERROR_INVALID_SEQ_LEN) - - request_begin = cutlass.Int32(paged_kv_indptr[batch_idx]) - request_end = cutlass.Int32(paged_kv_indptr[batch_idx + 1]) - num_page_indices = cutlass.Int32(cute.size(paged_kv_indices)) - request_range_is_valid = cutlass.Int32( - paged_kv_indptr[cutlass.Int32(0)] == cutlass.Int32(0) - and request_begin >= cutlass.Int32(0) - and request_begin <= request_end - and request_end <= num_page_indices - ) - if request_range_is_valid == cutlass.Int32(0): - error_code = cutlass.Int32(_ERROR_INVALID_PAGE_INDPTR) - elif seq_len_is_valid: - required_pages = (seq_len_kv - cutlass.Int32(1)) // cutlass.Int32( + if seq_len_is_valid: + live_pages = (seq_len_kv - cutlass.Int32(1)) // cutlass.Int32( self.page_size ) + cutlass.Int32(1) - if request_end - request_begin < required_pages: - error_code = cutlass.Int32(_ERROR_INSUFFICIENT_PAGE_CAPACITY) + else: + error_code = cutlass.Int32(_ERROR_INVALID_SEQ_LEN) - request_begin = _warp_broadcast_i32(request_begin, 0) - request_end = _warp_broadcast_i32(request_end, 0) - request_range_is_valid = _warp_broadcast_i32(request_range_is_valid, 0) - if request_is_valid and request_range_is_valid != cutlass.Int32(0): + live_pages = _warp_broadcast_i32(live_pages, 0) + if request_is_valid: + row_begin = cutlass.Int64(batch_idx) * block_table_row_stride page_offset = cutlass.Int64(lane_idx) - request_page_count = cutlass.Int64(request_end) - cutlass.Int64( - request_begin - ) - while page_offset < request_page_count: - page_position = cutlass.Int64(request_begin) + page_offset - physical_page_id = cutlass.Int32(paged_kv_indices[page_position]) + while page_offset < cutlass.Int64(live_pages): + page_position = cutlass.Int64(row_begin + page_offset) + physical_page_id = cutlass.Int32(block_tables.iterator[page_position]) if ( physical_page_id < cutlass.Int32(0) or cutlass.Int64(physical_page_id) >= num_physical_kv_pages @@ -433,16 +411,16 @@ def __call__( self, block_indptr: cute.Tensor, block_indices: cute.Tensor, - paged_kv_indptr: cute.Tensor, - paged_kv_indices: cute.Tensor, + block_tables: cute.Tensor, + block_table_row_stride: cutlass.Int64, seq_lens_kv: cute.Tensor, num_physical_kv_pages: cutlass.Int64, summary: cute.Tensor, stream: cuda_drv.CUstream, ) -> None: self.inspect_requests( - paged_kv_indptr, - paged_kv_indices, + block_tables, + block_table_row_stride, seq_lens_kv, num_physical_kv_pages, summary, @@ -532,8 +510,9 @@ def compile_paged_block_sparse_metadata_inspection( """Compile one paged metadata entry that launches request then live-BSR.""" num_q_block_rows = (seq_len_q + q_block_size - 1) // q_block_size - logical_page_capacity = cute.sym_int() logical_nnz = cute.sym_int() + runtime_page_columns = cute.sym_int() + runtime_page_row_stride = cute.sym_int64(divisibility=1) stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) inspect_requests = _InspectPagedKvMetadata( @@ -564,8 +543,13 @@ def compile_paged_block_sparse_metadata_inspection( alignment=4, ), _fake_compact(cutlass.Int32, (logical_nnz,), alignment=4), - _fake_compact(cutlass.Int32, (batch_size + 1,), alignment=4), - _fake_compact(cutlass.Int32, (logical_page_capacity,), alignment=4), + cute.runtime.make_fake_tensor( + cutlass.Int32, + (batch_size, runtime_page_columns), + stride=(runtime_page_row_stride, 1), + assumed_align=4, + ), + cutlass.Int64(1), _fake_compact(cutlass.Int32, (batch_size,), alignment=4), cutlass.Int64(1), _fake_compact(cutlass.Int64, (_SUMMARY_FIELDS,), alignment=8), diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_prepare.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_prepare.py index 6a61d0be86cd..08f2872809cb 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_prepare.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_prepare.py @@ -12,13 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Prepare live canonical BSR rows for the PrimTS FMHA route consumer. +"""Prepare exact-first sparse routes for the PrimTS FMHA route consumer. -The kernel converts caller-owned semantic KV blocks into fixed-stride route -metadata on every run. Route origins are logical KV-token coordinates; -the paged specialization also resolves each origin to a physical page ID for -the attention load path. One warp handles one BSR row and iterates only that -row's live routes, while four warps share a CTA. +The BSR frontend is shared by continuous exact/proxy and paged exact storage. +It validates each canonical row once, emits the same logical exact records, +then either resolves paged locators or appends continuous proxy records. The +bitmask frontend shares the record geometry and emitters but remains limited +to continuous storage. Proxy suffixes contain one stable record per summary +group; a fully exact group remains present with zero score words. One warp owns +one sparse row and four warps share a CTA. ``row_route_offsets`` is a separate plan-owned immutable Int32 tensor. ``route_workspace`` contains only mutable row counts and route metadata @@ -27,21 +29,20 @@ semantic BSR-block limit, which remains distinct from packed-route capacity. """ -import math from dataclasses import dataclass import cutlass import cutlass.cute as cute from cuda.bindings import driver as cuda_drv from cutlass.cute.testing import assert_ as runtime_assert -from cutlass.experimental import primitives as prims from ..._block_sparse.prepared import ( _PREPARED_ROUTE_IS_FULL_FLAG, + _PREPARED_ROUTE_IS_PROXY_FLAG, _BlockSparseRouteLayout, ) -from .block_sparse_inspect import _validate_bsr_row_lane from .fmha_decode_resources.helpers_common import _warp_broadcast_i32 +from .block_sparse_inspect import _validate_bsr_row_lane _WARPS_PER_CTA = 4 @@ -50,22 +51,27 @@ @dataclass(frozen=True) -class _PreparedRouteConfig: - """Compile-time geometry shared by contiguous and paged route packing.""" +class _RouteConfig: + """Compile-time route geometry shared across sparse input/storage modes.""" num_kv_heads: int - num_q_block_rows: int + num_q_blocks: int num_kv_blocks: int + num_exact_words: int + num_proxy_groups: int num_rows: int seq_len_kv: int kv_block_size: int atom_size: int + atoms_per_block: int logical_origins_per_route: int token_words_per_route: int atom_valid_mask_word_offset: int route_flags_word_offset: int token_words_word_offset: int - has_token_bits: bool + stores_score_words: bool + apply_token_mask: bool + use_proxy_routes: bool route_metadata_stride_words: int route_metadata_base_word_offset: int @@ -78,18 +84,30 @@ def create( seq_len_kv: int, q_block_size: int, kv_block_size: int, - ) -> "_PreparedRouteConfig": - """Build shared prepare geometry without adding a storage-mode flag.""" + apply_token_mask: bool, + use_proxy_routes: bool, + ) -> "_RouteConfig": + """Build storage-independent route geometry and policy flags.""" - num_q_block_rows = (seq_len_q + q_block_size - 1) // q_block_size - return _PreparedRouteConfig( + stores_score_words = layout.token_words_word_offset is not None + if apply_token_mask and not stores_score_words: + raise ValueError("token masking requires prepared score words") + if use_proxy_routes and not stores_score_words: + raise ValueError("proxy routes require prepared score words") + num_q_blocks = (seq_len_q + q_block_size - 1) // q_block_size + num_kv_blocks = (seq_len_kv + kv_block_size - 1) // kv_block_size + return _RouteConfig( num_kv_heads=num_kv_heads, - num_q_block_rows=num_q_block_rows, - num_kv_blocks=(seq_len_kv + kv_block_size - 1) // kv_block_size, + num_q_blocks=num_q_blocks, + num_kv_blocks=num_kv_blocks, + num_exact_words=(num_kv_blocks + _WARP_SIZE - 1) // _WARP_SIZE, + num_proxy_groups=(num_kv_blocks + layout.kv_route_size - 1) + // layout.kv_route_size, num_rows=layout.num_rows, seq_len_kv=seq_len_kv, kv_block_size=kv_block_size, atom_size=layout.atom_size, + atoms_per_block=kv_block_size // layout.atom_size, logical_origins_per_route=layout.logical_origins_per_route, token_words_per_route=layout.token_words_per_route, atom_valid_mask_word_offset=layout.atom_valid_mask_word_offset, @@ -99,7 +117,9 @@ def create( if layout.token_words_word_offset is not None else 0 ), - has_token_bits=layout.has_token_bits, + stores_score_words=stores_score_words, + apply_token_mask=apply_token_mask, + use_proxy_routes=use_proxy_routes, route_metadata_stride_words=layout.route_metadata_stride_words, route_metadata_base_word_offset=layout.route_metadata_base_word_offset, ) @@ -115,35 +135,44 @@ def _positive_i32_ceil_div( @cute.jit -def _retained_atom_count( - block_indices: cute.Tensor, - row_begin: cutlass.Int32, - row_end: cutlass.Int32, - kv_block_size: cutlass.Constexpr[int], - atom_size: cutlass.Constexpr[int], - seq_len_kv: cutlass.Int32, +def _prepared_route_counts( + selected_block_count: cutlass.Int32, + cfg: cutlass.Constexpr[_RouteConfig], +) -> tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32]: + """Return exact atoms, exact routes, and total prepared routes for one row.""" + + exact_atom_count = selected_block_count * cutlass.Int32(cfg.atoms_per_block) + exact_route_count = ( + exact_atom_count + cutlass.Int32(cfg.logical_origins_per_route - 1) + ) // cutlass.Int32(cfg.logical_origins_per_route) + total_route_count = exact_route_count + if cutlass.const_expr(cfg.use_proxy_routes): + total_route_count += cutlass.Int32(cfg.num_proxy_groups) + return exact_atom_count, exact_route_count, total_route_count + + +@cute.jit +def _prepared_row_route_begin( + row_route_offsets: cute.Tensor, + linear_row_idx: cutlass.Int32, + lane_idx: cutlass.Int32, + row_is_valid: cutlass.Boolean, + total_route_count: cutlass.Int32, ) -> cutlass.Int32: - """Count selected atoms whose logical origin precedes ``seq_len_kv``.""" - - row_nnz = row_end - row_begin - retained_atoms = cutlass.Int32(0) - if row_nnz > cutlass.Int32(0): - atoms_per_block = kv_block_size // atom_size - retained_atoms = (row_nnz - cutlass.Int32(1)) * cutlass.Int32(atoms_per_block) - last_block_idx = cutlass.Int32(block_indices[row_end - cutlass.Int32(1)]) - last_block_origin = last_block_idx * cutlass.Int32(kv_block_size) - remaining_tokens = cutlass.Int32(seq_len_kv) - last_block_origin + """Load and validate one row's plan-owned prepared-route span.""" + + row_route_begin = cutlass.Int32(0) + if lane_idx == cutlass.Int32(0) and row_is_valid: + row_route_begin = cutlass.Int32(row_route_offsets[linear_row_idx]) + row_route_end = cutlass.Int32(row_route_offsets[linear_row_idx + 1]) + row_capacity = row_route_end - row_route_begin runtime_assert( - remaining_tokens > cutlass.Int32(0), - "block_indices row exceeds the live KV block range", + row_route_begin >= cutlass.Int32(0) + and row_capacity >= cutlass.Int32(0) + and total_route_count <= row_capacity, + "prepared routes exceed planned row capacity", ) - retained_last_atoms = (remaining_tokens - cutlass.Int32(1)) // cutlass.Int32( - atom_size - ) + cutlass.Int32(1) - if retained_last_atoms > cutlass.Int32(atoms_per_block): - retained_last_atoms = cutlass.Int32(atoms_per_block) - retained_atoms = retained_atoms + retained_last_atoms - return retained_atoms + return _warp_broadcast_i32(row_route_begin, 0) @cute.jit @@ -177,71 +206,72 @@ def _resolve_route_logical_atom_origin( @cute.jit -def _load_coarse_token_word( - block_indices: cute.Tensor, - kv_valid_bits: cute.Tensor, - row_begin: cutlass.Int32, - row_end: cutlass.Int32, - route_idx: cutlass.Int32, - logical_word_idx: cutlass.Int32, - batch_idx: cutlass.Int32, - kv_block_size: cutlass.Constexpr[int], - atom_size: cutlass.Constexpr[int], - logical_origins_per_route: cutlass.Constexpr[int], - seq_len_kv: cutlass.Int32, -) -> cutlass.Uint32: - """Load one logical K32 word from a coarse atom larger than K32.""" - - logical_word = cutlass.Uint32(0) - words_per_atom = atom_size // 32 - atom_in_route = logical_word_idx // cutlass.Int32(words_per_atom) - word_in_atom = logical_word_idx % cutlass.Int32(words_per_atom) - logical_origin, valid = _resolve_route_logical_atom_origin( - block_indices, - row_begin, - row_end, - route_idx, - atom_in_route, - kv_block_size, - atom_size, - logical_origins_per_route, - seq_len_kv, - ) - logical_word_origin = logical_origin + word_in_atom * cutlass.Int32(32) - if valid and logical_word_origin < cutlass.Int32(seq_len_kv): - valid_bits_word_idx = logical_word_origin >> cutlass.Int32(5) - logical_word = cutlass.Uint32(kv_valid_bits[batch_idx, valid_bits_word_idx]) - remaining_tokens = cutlass.Int32(seq_len_kv) - logical_word_origin - if remaining_tokens < cutlass.Int32(32): - logical_word = logical_word & ( - (cutlass.Uint32(1) << remaining_tokens) - cutlass.Uint32(1) - ) - return logical_word +def _low_bits_mask(valid_bits: cutlass.Int32) -> cutlass.Uint32: + """Return a Uint32 mask with its lowest clamped bit count set.""" + + mask = cutlass.Uint32(0) + if valid_bits >= cutlass.Int32(_WARP_SIZE): + mask = cutlass.Uint32(0xFFFFFFFF) + elif valid_bits > cutlass.Int32(0): + mask = (cutlass.Uint32(1) << valid_bits) - cutlass.Uint32(1) + return mask @cute.jit -def _load_atom_token_chunk( +def _load_exact_score_word( + route_workspace: cute.Tensor, kv_valid_bits: cute.Tensor, + route_metadata_word_index: cutlass.Int32, + logical_word_idx: cutlass.Int32, batch_idx: cutlass.Int32, - logical_origin: cutlass.Int32, - origin_is_valid: cutlass.Boolean, - atom_size: cutlass.Constexpr[int], seq_len_kv: cutlass.Int32, + cfg: cutlass.Constexpr[_RouteConfig], ) -> cutlass.Uint32: - """Load the <=K32 mask chunk owned by one resolved-origin lane.""" - - token_chunk = cutlass.Uint32(0) - if origin_is_valid: - valid_bits_word_idx = logical_origin >> cutlass.Int32(5) - source_word = cutlass.Uint32(kv_valid_bits[batch_idx, valid_bits_word_idx]) - token_chunk = source_word >> (logical_origin & cutlass.Int32(31)) - token_chunk = token_chunk & cutlass.Uint32((1 << atom_size) - 1) - remaining_tokens = cutlass.Int32(seq_len_kv) - logical_origin - if remaining_tokens < cutlass.Int32(atom_size): - token_chunk = token_chunk & ( - (cutlass.Uint32(1) << remaining_tokens) - cutlass.Uint32(1) - ) - return token_chunk + """Build one exact score word with optional caller-token masking.""" + + token_word = cutlass.Uint32(0) + if cutlass.const_expr(cfg.atom_size <= _WARP_SIZE): + atoms_per_word = _WARP_SIZE // cfg.atom_size + first_atom_idx = logical_word_idx * cutlass.Int32(atoms_per_word) + for atom_in_word in cutlass.range_constexpr(atoms_per_word): + atom_idx = first_atom_idx + cutlass.Int32(atom_in_word) + if atom_idx < cutlass.Int32(cfg.logical_origins_per_route): + origin = cutlass.Int32( + route_workspace[route_metadata_word_index + atom_idx] + ) + atom_word = cutlass.Uint32(0) + if origin >= cutlass.Int32(0): + if cutlass.const_expr(cfg.apply_token_mask): + source_word_idx = origin >> cutlass.Int32(5) + atom_word = cutlass.Uint32( + kv_valid_bits[batch_idx, source_word_idx] + ) + atom_word = atom_word >> (origin & cutlass.Int32(31)) + atom_word = atom_word & cutlass.Uint32((1 << cfg.atom_size) - 1) + atom_word = atom_word & _low_bits_mask(seq_len_kv - origin) + else: + atom_word = _low_bits_mask( + seq_len_kv - origin + ) & cutlass.Uint32((1 << cfg.atom_size) - 1) + token_word = token_word | ( + atom_word << cutlass.Int32(atom_in_word * cfg.atom_size) + ) + else: + words_per_atom = cfg.atom_size // _WARP_SIZE + atom_idx = logical_word_idx // cutlass.Int32(words_per_atom) + word_in_atom = logical_word_idx % cutlass.Int32(words_per_atom) + origin = cutlass.Int32(route_workspace[route_metadata_word_index + atom_idx]) + word_origin = origin + word_in_atom * cutlass.Int32(_WARP_SIZE) + if origin >= cutlass.Int32(0): + if cutlass.const_expr(cfg.apply_token_mask): + if word_origin < seq_len_kv: + source_word_idx = word_origin >> cutlass.Int32(5) + token_word = cutlass.Uint32( + kv_valid_bits[batch_idx, source_word_idx] + ) & _low_bits_mask(seq_len_kv - word_origin) + else: + token_word = _low_bits_mask(seq_len_kv - word_origin) + return token_word @cute.jit @@ -251,7 +281,7 @@ def _resolve_prepared_bsr_row( linear_row_idx: cutlass.Int32, lane_idx: cutlass.Int32, row_is_valid: cutlass.Boolean, - cfg: cutlass.Constexpr[_PreparedRouteConfig], + cfg: cutlass.Constexpr[_RouteConfig], ) -> tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32]: """Resolve one trusted canonical runtime BSR row.""" @@ -259,8 +289,8 @@ def _resolve_prepared_bsr_row( row_end = cutlass.Int32(0) batch_idx = cutlass.Int32(0) if lane_idx == cutlass.Int32(0) and row_is_valid: - q_block_row_idx = linear_row_idx % cfg.num_q_block_rows - linear_batch_head_idx = linear_row_idx // cfg.num_q_block_rows + q_block_row_idx = linear_row_idx % cfg.num_q_blocks + linear_batch_head_idx = linear_row_idx // cfg.num_q_blocks kv_head_idx = linear_batch_head_idx % cfg.num_kv_heads batch_idx = linear_batch_head_idx // cfg.num_kv_heads row_begin = cutlass.Int32(block_indptr[batch_idx, kv_head_idx, q_block_row_idx]) @@ -297,209 +327,256 @@ def _resolve_prepared_bsr_row( @cute.jit -def _publish_prepared_route_count( - block_indices: cute.Tensor, - row_route_offsets: cute.Tensor, +def _finalize_exact_route( route_workspace: cute.Tensor, - row_begin: cutlass.Int32, - row_end: cutlass.Int32, - linear_row_idx: cutlass.Int32, - lane_idx: cutlass.Int32, - row_is_valid: cutlass.Boolean, - max_blocks_per_row: cutlass.Int32, - seq_len_kv: cutlass.Int32, - cfg: cutlass.Constexpr[_PreparedRouteConfig], -) -> tuple[cutlass.Int32, cutlass.Int32]: - """Assert semantic capacity, publish the header, and return its live span.""" - - row_route_begin = cutlass.Int32(0) - required_route_count = cutlass.Int32(0) - if lane_idx == cutlass.Int32(0) and row_is_valid: - row_route_begin = cutlass.Int32(row_route_offsets[linear_row_idx]) - selected_block_count = row_end - row_begin - runtime_assert( - selected_block_count <= max_blocks_per_row, - "selected BSR blocks exceed planned semantic capacity", - ) - retained_atom_count = _retained_atom_count( - block_indices, - row_begin, - row_end, - cfg.kv_block_size, - cfg.atom_size, - seq_len_kv, - ) - required_route_count = ( - retained_atom_count + cutlass.Int32(cfg.logical_origins_per_route - 1) - ) // cutlass.Int32(cfg.logical_origins_per_route) - route_workspace[linear_row_idx] = required_route_count - row_route_begin = _warp_broadcast_i32(row_route_begin, 0) - required_route_count = _warp_broadcast_i32(required_route_count, 0) - return required_route_count, row_route_begin - - -@cute.jit -def _store_prepared_route_validity( - block_indices: cute.Tensor, kv_valid_bits: cute.Tensor, - route_workspace: cute.Tensor, - row_begin: cutlass.Int32, - row_end: cutlass.Int32, - route_idx: cutlass.Int32, + route_metadata_word_index: cutlass.Int32, batch_idx: cutlass.Int32, lane_idx: cutlass.Int32, - logical_origin: cutlass.Int32, - logical_origin_is_valid: cutlass.Boolean, - stored_atom_is_full: cutlass.Boolean, - route_metadata_word_index: cutlass.Int32, + atom_is_valid: cutlass.Boolean, seq_len_kv: cutlass.Int32, - cfg: cutlass.Constexpr[_PreparedRouteConfig], + cfg: cutlass.Constexpr[_RouteConfig], ) -> None: - """Store storage-independent atom, token, and route validity metadata.""" + """Finalize an exact record after its logical origins are stored.""" - stored_atom_valid_mask = cutlass.Int32( - cute.arch.vote_ballot_sync(logical_origin_is_valid) - ) - structural_route_is_full = cute.arch.vote_all_sync( - lane_idx >= cutlass.Int32(cfg.logical_origins_per_route) or stored_atom_is_full + atom_is_full = cutlass.Boolean(False) + if lane_idx < cutlass.Int32(cfg.logical_origins_per_route): + origin = cutlass.Int32(route_workspace[route_metadata_word_index + lane_idx]) + atom_is_full = cutlass.Boolean( + atom_is_valid and origin <= seq_len_kv - cutlass.Int32(cfg.atom_size) + ) + atom_valid_mask = cutlass.Int32(cute.arch.vote_ballot_sync(atom_is_valid)) + structural_full = cute.arch.vote_all_sync( + lane_idx >= cutlass.Int32(cfg.logical_origins_per_route) or atom_is_full ) - route_is_full = structural_route_is_full - if cutlass.const_expr(cfg.has_token_bits): - token_word = cutlass.Uint32(0) - if cutlass.const_expr(cfg.atom_size <= 32): - token_chunk = _load_atom_token_chunk( + + score_words_are_full = cutlass.Boolean(True) + if cutlass.const_expr(cfg.stores_score_words): + score_word = cutlass.Uint32(0) + if lane_idx < cutlass.Int32(cfg.token_words_per_route): + score_word = _load_exact_score_word( + route_workspace, kv_valid_bits, + route_metadata_word_index, + lane_idx, batch_idx, - logical_origin, - logical_origin_is_valid, - cfg.atom_size, seq_len_kv, + cfg, ) - atoms_per_word = 32 // cfg.atom_size - if lane_idx < cutlass.Int32(cfg.logical_origins_per_route): - atom_in_word = lane_idx % cutlass.Int32(atoms_per_word) - token_word = token_chunk << ( - atom_in_word * cutlass.Int32(cfg.atom_size) - ) - active_origin_lanes = (1 << cfg.logical_origins_per_route) - 1 - for shuffle_step in cutlass.range_constexpr( - int(math.log2(atoms_per_word)) - ): - peer_word = cutlass.Uint32( - prims.shfl_sync( - thread_mask=active_origin_lanes, - val=token_word, - offset=1 << shuffle_step, - mask_and_clamp=0x1F, - kind=prims.Shfl.BFLY, - ) - ) - token_word = token_word | peer_word - if atom_in_word == cutlass.Int32(0): - logical_word_idx = lane_idx // cutlass.Int32(atoms_per_word) - route_workspace[ - route_metadata_word_index - + cutlass.Int32(cfg.token_words_word_offset) - + logical_word_idx - ] = cutlass.Int32(token_word) - full_atom_mask = cutlass.Uint32((1 << cfg.atom_size) - 1) - token_route_is_full = cute.arch.vote_all_sync( - lane_idx >= cutlass.Int32(cfg.logical_origins_per_route) - or token_chunk == full_atom_mask - ) - else: - if lane_idx < cutlass.Int32(cfg.token_words_per_route): - token_word = _load_coarse_token_word( - block_indices, - kv_valid_bits, - row_begin, - row_end, - route_idx, - lane_idx, - batch_idx, - cfg.kv_block_size, - cfg.atom_size, - cfg.logical_origins_per_route, - seq_len_kv, - ) - route_workspace[ - route_metadata_word_index - + cutlass.Int32(cfg.token_words_word_offset) - + lane_idx - ] = cutlass.Int32(token_word) - token_route_is_full = cute.arch.vote_all_sync( - lane_idx >= cutlass.Int32(cfg.token_words_per_route) - or token_word == cutlass.Uint32(0xFFFFFFFF) - ) - route_is_full = cutlass.Boolean( - structural_route_is_full and token_route_is_full + route_workspace[ + route_metadata_word_index + + cutlass.Int32(cfg.token_words_word_offset) + + lane_idx + ] = cutlass.Int32(score_word) + score_words_are_full = cute.arch.vote_all_sync( + lane_idx >= cutlass.Int32(cfg.token_words_per_route) + or score_word == cutlass.Uint32(0xFFFFFFFF) ) - if lane_idx == cutlass.Int32(0): route_workspace[ route_metadata_word_index + cutlass.Int32(cfg.atom_valid_mask_word_offset) - ] = stored_atom_valid_mask + ] = atom_valid_mask route_workspace[ route_metadata_word_index + cutlass.Int32(cfg.route_flags_word_offset) ] = ( cutlass.Int32(_PREPARED_ROUTE_IS_FULL_FLAG) - if route_is_full + if structural_full and score_words_are_full else cutlass.Int32(0) ) -@cute.jit -def _paged_request_page_range_is_valid( - request_begin: cutlass.Int32, - request_end: cutlass.Int32, - num_indices: cutlass.Int32, - required_pages: cutlass.Int32, -) -> cutlass.Boolean: - """Validate one request's page-table range before any index load.""" - - return cutlass.Boolean( - request_begin >= cutlass.Int32(0) - and request_begin <= request_end - and request_end <= num_indices - and request_end - request_begin >= required_pages - ) - - @cute.jit def _resolve_paged_route_atom_page_id( - paged_kv_indices: cute.Tensor, - request_begin: cutlass.Int32, + block_tables: cute.Tensor, + batch_idx: cutlass.Int32, + block_table_row_stride: cutlass.Int64, logical_origin: cutlass.Int32, logical_origin_is_valid: cutlass.Boolean, lane_idx: cutlass.Int32, page_size: cutlass.Constexpr[int], num_physical_kv_pages: cutlass.Int64, ) -> cutlass.Int32: - """Resolve one trusted selected logical atom to its raw physical page ID.""" + """Resolve one trusted selected logical atom to its physical page ID.""" physical_page_id = cutlass.Int32(-1) page_id_is_valid = cutlass.Boolean(True) if logical_origin_is_valid: logical_page_idx = logical_origin // cutlass.Int32(page_size) - candidate_page_id = cutlass.Int32( - paged_kv_indices[request_begin + logical_page_idx] + physical_page_id = cutlass.Int32( + block_tables.iterator[ + cutlass.Int64(batch_idx) * block_table_row_stride + + cutlass.Int64(logical_page_idx) + ] ) - physical_page_id = candidate_page_id page_id_is_valid = cutlass.Boolean( - candidate_page_id >= cutlass.Int32(0) - and cutlass.Int64(candidate_page_id) < num_physical_kv_pages + physical_page_id >= cutlass.Int32(0) + and cutlass.Int64(physical_page_id) < num_physical_kv_pages ) page_ids_are_valid = cute.arch.vote_all_sync(page_id_is_valid) if lane_idx == cutlass.Int32(0): runtime_assert( page_ids_are_valid, - "paged_kv_indices contains an out-of-range physical page ID", + "block_tables contains an out-of-range physical page ID", ) return physical_page_id -class _PrepareBlockSparseRoutes: - """Prepare contiguous or paged sparse routes for one static geometry.""" +@cute.jit +def _exact_lane_rank( + exact_ballot: cutlass.Uint32, + lane_idx: cutlass.Int32, + exact_prefix: cutlass.Int32, +) -> cutlass.Int32: + """Return one exact lane's global semantic-block rank.""" + + lower_lane_mask = (cutlass.Uint32(1) << lane_idx) - cutlass.Uint32(1) + return exact_prefix + cutlass.Int32(cute.arch.popc(exact_ballot & lower_lane_mask)) + + +@cute.jit +def _emit_exact_block_atoms( + route_workspace: cute.Tensor, + row_route_begin: cutlass.Int32, + semantic_block_idx: cutlass.Int32, + exact_block_rank: cutlass.Int32, + cfg: cutlass.Constexpr[_RouteConfig], +) -> None: + """Expand one bitmask-selected block into fixed row-global atom slots.""" + + first_atom_rank = exact_block_rank * cutlass.Int32(cfg.atoms_per_block) + atom_in_block = cutlass.Int32(0) + while atom_in_block < cutlass.Int32(cfg.atoms_per_block): + atom_rank = first_atom_rank + atom_in_block + route_idx = atom_rank // cutlass.Int32(cfg.logical_origins_per_route) + atom_in_route = atom_rank % cutlass.Int32(cfg.logical_origins_per_route) + route_word_index = cutlass.Int32(cfg.route_metadata_base_word_offset) + ( + (row_route_begin + route_idx) + * cutlass.Int32(cfg.route_metadata_stride_words) + ) + logical_origin = semantic_block_idx * cutlass.Int32( + cfg.kv_block_size + ) + atom_in_block * cutlass.Int32(cfg.atom_size) + stored_origin = cutlass.Int32(-1) + if logical_origin < cutlass.Int32(cfg.seq_len_kv): + stored_origin = logical_origin + route_workspace[route_word_index + atom_in_route] = stored_origin + atom_in_block += cutlass.Int32(1) + + +@cute.jit +def _load_bitmask_word( + exact_block_bits: cute.Tensor, + batch_idx: cutlass.Int32, + kv_head_idx: cutlass.Int32, + q_block_idx: cutlass.Int32, + logical_word_idx: cutlass.Int32, + cfg: cutlass.Constexpr[_RouteConfig], + for_proxy: cutlass.Constexpr[bool], +) -> cutlass.Uint32: + """Load one in-range exact or proxy semantic-block word.""" + + valid_word = _low_bits_mask( + cutlass.Int32(cfg.num_kv_blocks) - logical_word_idx * cutlass.Int32(_WARP_SIZE) + ) + selected_word = cutlass.Uint32( + exact_block_bits[batch_idx, kv_head_idx, q_block_idx, logical_word_idx] + ) + if cutlass.const_expr(for_proxy): + selected_word = ~selected_word + return valid_word & selected_word + + +@cute.jit +def _load_bsr_proxy_word( + block_indices: cute.Tensor, + row_begin: cutlass.Int32, + row_end: cutlass.Int32, + logical_word_idx: cutlass.Int32, + cfg: cutlass.Constexpr[_RouteConfig], +) -> cutlass.Uint32: + """Build one proxy word from a canonical sorted-BSR interval.""" + + word_begin = logical_word_idx * cutlass.Int32(_WARP_SIZE) + valid_word = _low_bits_mask(cutlass.Int32(cfg.num_kv_blocks) - word_begin) + selected_word = cutlass.Uint32(0) + lower = row_begin + upper = row_end + while lower < upper: + middle = lower + (upper - lower) // cutlass.Int32(2) + if cutlass.Int32(block_indices[middle]) < word_begin: + lower = middle + cutlass.Int32(1) + else: + upper = middle + cursor = lower + word_end = word_begin + cutlass.Int32(_WARP_SIZE) + scanning = cutlass.Boolean(True) + while cursor < row_end and scanning: + block_idx = cutlass.Int32(block_indices[cursor]) + if block_idx < word_end: + selected_word = selected_word | ( + cutlass.Uint32(1) << (block_idx - word_begin) + ) + cursor += cutlass.Int32(1) + else: + scanning = cutlass.Boolean(False) + return valid_word & ~selected_word + + +@cute.jit +def _emit_proxy_route( + route_workspace: cute.Tensor, + row_route_begin: cutlass.Int32, + exact_route_count: cutlass.Int32, + group_idx: cutlass.Int32, + proxy_word: cutlass.Uint32, + lane_idx: cutlass.Int32, + cfg: cutlass.Constexpr[_RouteConfig], +) -> None: + """Emit one fixed summary-group proxy record, including an empty mask.""" + + route_metadata_word_index = cutlass.Int32(cfg.route_metadata_base_word_offset) + ( + row_route_begin + exact_route_count + group_idx + ) * cutlass.Int32(cfg.route_metadata_stride_words) + group_start = group_idx * cutlass.Int32(cfg.token_words_per_route * _WARP_SIZE) + group_size = cutlass.Int32(cfg.num_kv_blocks) - group_start + if group_size > cutlass.Int32(cfg.token_words_per_route * _WARP_SIZE): + group_size = cutlass.Int32(cfg.token_words_per_route * _WARP_SIZE) + origin_is_valid = cutlass.Boolean(False) + if lane_idx < cutlass.Int32(cfg.logical_origins_per_route): + summary_origin = group_start + lane_idx * cutlass.Int32(cfg.atom_size) + origin_is_valid = cutlass.Boolean(summary_origin < cfg.num_kv_blocks) + stored_origin = cutlass.Int32(-1) + if origin_is_valid: + stored_origin = summary_origin + route_workspace[route_metadata_word_index + lane_idx] = stored_origin + atom_valid_mask = cutlass.Int32(cute.arch.vote_ballot_sync(origin_is_valid)) + if lane_idx < cutlass.Int32(cfg.token_words_per_route): + route_workspace[ + route_metadata_word_index + + cutlass.Int32(cfg.token_words_word_offset) + + lane_idx + ] = cutlass.Int32(proxy_word) + score_full = cute.arch.vote_all_sync( + lane_idx >= cutlass.Int32(cfg.token_words_per_route) + or proxy_word == cutlass.Uint32(0xFFFFFFFF) + ) + if lane_idx == cutlass.Int32(0): + route_workspace[ + route_metadata_word_index + cutlass.Int32(cfg.atom_valid_mask_word_offset) + ] = atom_valid_mask + proxy_is_full = cutlass.Boolean( + group_size == cutlass.Int32(cfg.token_words_per_route * _WARP_SIZE) + and score_full + ) + route_workspace[ + route_metadata_word_index + cutlass.Int32(cfg.route_flags_word_offset) + ] = cutlass.Int32(_PREPARED_ROUTE_IS_PROXY_FLAG) | ( + cutlass.Int32(proxy_is_full) * cutlass.Int32(_PREPARED_ROUTE_IS_FULL_FLAG) + ) + + +class _PrepareRoutesBase: + """Own shared route geometry and compile-time storage/policy flags.""" def __init__( self, @@ -511,38 +588,58 @@ def __init__( q_block_size: int, kv_block_size: int, kv_route_size: int, - has_token_bits: bool, + use_proxy_routes: bool, + use_causal_mask: bool = False, + apply_token_mask: bool = False, + store_score_words: bool = False, page_size: int | None = None, - mask_type: str, ) -> None: - if mask_type not in ("dense", "causal"): - raise ValueError(f"unsupported mask_type: {mask_type}") - num_q_block_rows = (seq_len_q + q_block_size - 1) // q_block_size - num_rows = batch_size * num_kv_heads * num_q_block_rows + if not isinstance(use_proxy_routes, bool): + raise TypeError("use_proxy_routes must be a bool") + if not isinstance(apply_token_mask, bool): + raise TypeError("apply_token_mask must be a bool") + if not isinstance(store_score_words, bool): + raise TypeError("store_score_words must be a bool") + if not isinstance(use_causal_mask, bool): + raise TypeError("use_causal_mask must be a bool") + if use_proxy_routes and page_size is not None: + raise ValueError("paged KV does not support proxy routes") + + num_q_blocks = (seq_len_q + q_block_size - 1) // q_block_size + num_rows = batch_size * num_kv_heads * num_q_blocks + # Structural score words (sequence tail, invalid atoms) can be stored + # without a caller token mask; proxy routes and token masks require them. + stores_score_words = use_proxy_routes or apply_token_mask or store_score_words layout = _BlockSparseRouteLayout.create( kv_route_size=kv_route_size, kv_block_size=kv_block_size, page_size=page_size, - has_token_bits=has_token_bits, + has_token_bits=stores_score_words, route_metadata_capacity=0, num_rows=num_rows, ) - self.cfg = _PreparedRouteConfig.create( + self.route_layout = layout + self.cfg = _RouteConfig.create( layout=layout, num_kv_heads=num_kv_heads, seq_len_q=seq_len_q, seq_len_kv=seq_len_kv, q_block_size=q_block_size, kv_block_size=kv_block_size, + apply_token_mask=apply_token_mask, + use_proxy_routes=use_proxy_routes, ) - self.route_layout = layout self.page_size = page_size if page_size is not None else 1 - self.minimum_seq_len_kv = seq_len_q if mask_type == "causal" else 1 + self.minimum_seq_len_kv = seq_len_q if use_causal_mask else 1 self.physical_page_ids_word_offset = ( layout.physical_page_ids_word_offset if layout.is_paged else 0 ) self.route_metadata_base_word_offset = layout.route_metadata_base_word_offset + +class _PrepareBsrRoutes(_PrepareRoutesBase): + """Prepare continuous exact/proxy or paged exact routes from one BSR flow.""" + @cute.jit def __call__( self, @@ -550,24 +647,24 @@ def __call__( block_indices: cute.Tensor, kv_valid_bits: cute.Tensor, seq_lens_kv: cute.Tensor | None, - paged_kv_indptr: cute.Tensor | None, - paged_kv_indices: cute.Tensor | None, + block_tables: cute.Tensor | None, num_physical_kv_pages: cutlass.Int64, + block_table_row_stride: cutlass.Int64, row_route_offsets: cute.Tensor, route_workspace: cute.Tensor, max_blocks_per_row: cutlass.Int32, stream: cuda_drv.CUstream, ) -> None: - """Launch four independent row preparers per CTA.""" + """Launch four independent BSR row preparers per CTA.""" self.kernel( block_indptr, block_indices, kv_valid_bits, seq_lens_kv, - paged_kv_indptr, - paged_kv_indices, + block_tables, num_physical_kv_pages, + block_table_row_stride, row_route_offsets, route_workspace, max_blocks_per_row, @@ -588,14 +685,14 @@ def kernel( block_indices: cute.Tensor, kv_valid_bits: cute.Tensor, seq_lens_kv: cute.Tensor | None, - paged_kv_indptr: cute.Tensor | None, - paged_kv_indices: cute.Tensor | None, + block_tables: cute.Tensor | None, num_physical_kv_pages: cutlass.Int64, + block_table_row_stride: cutlass.Int64, row_route_offsets: cute.Tensor, route_workspace: cute.Tensor, max_blocks_per_row: cutlass.Int32, ) -> None: - """Pack logical routes and, when paged, translate physical locators.""" + """Assert trusted inputs, emit routes, resolve storage, then publish.""" thread_idx, _, _ = cute.arch.thread_idx() block_idx, _, _ = cute.arch.block_idx() @@ -613,8 +710,8 @@ def kernel( self.cfg, ) - request_begin = cutlass.Int32(0) live_seq_len_kv = cutlass.Int32(self.cfg.seq_len_kv) + selected_block_count = row_end - row_begin if cutlass.const_expr(self.route_layout.is_paged): raw_seq_len_kv = cutlass.Int32(self.cfg.seq_len_kv) if lane_idx == cutlass.Int32(0) and row_is_valid: @@ -624,110 +721,366 @@ def kernel( and raw_seq_len_kv <= cutlass.Int32(self.cfg.seq_len_kv), "seq_lens_kv is outside the planned live-length range", ) - raw_seq_len_kv = _warp_broadcast_i32(raw_seq_len_kv, 0) - live_seq_len_kv = raw_seq_len_kv + live_seq_len_kv = _warp_broadcast_i32(raw_seq_len_kv, 0) + + if ( + lane_idx == cutlass.Int32(0) + and row_is_valid + and selected_block_count > cutlass.Int32(0) + ): + last_block_idx = cutlass.Int32( + block_indices[row_end - cutlass.Int32(1)] + ) + runtime_assert( + last_block_idx * cutlass.Int32(self.cfg.kv_block_size) + < live_seq_len_kv, + "block_indices row exceeds the live KV block range", + ) if lane_idx == cutlass.Int32(0) and row_is_valid: required_pages = _positive_i32_ceil_div( live_seq_len_kv, self.page_size, ) - request_begin = cutlass.Int32(paged_kv_indptr[batch_idx]) - request_end = cutlass.Int32( - paged_kv_indptr[batch_idx + cutlass.Int32(1)] - ) - metadata_starts_at_zero = cutlass.Boolean( - paged_kv_indptr[cutlass.Int32(0)] == cutlass.Int32(0) - ) runtime_assert( - metadata_starts_at_zero - and _paged_request_page_range_is_valid( - request_begin, - request_end, - cutlass.Int32(cute.size(paged_kv_indices)), - required_pages, - ), - "paged_kv_indptr row lacks the required live page capacity", + required_pages <= cutlass.Int32(block_tables.shape[1]), + "block_tables row lacks the required live page capacity", ) - request_begin = _warp_broadcast_i32(request_begin, 0) - route_count, row_route_begin = _publish_prepared_route_count( - block_indices, + _, exact_route_count, total_route_count = _prepared_route_counts( + selected_block_count, + self.cfg, + ) + if lane_idx == cutlass.Int32(0) and row_is_valid: + runtime_assert( + selected_block_count <= max_blocks_per_row, + "selected BSR blocks exceed planned semantic capacity", + ) + row_route_begin = _prepared_row_route_begin( row_route_offsets, - route_workspace, - row_begin, - row_end, linear_row_idx, lane_idx, row_is_valid, - max_blocks_per_row, - live_seq_len_kv, - self.cfg, + total_route_count, ) - route_idx = cutlass.Int32(0) - while route_idx < route_count: - route_ordinal = row_route_begin + route_idx - route_metadata_word_index = cutlass.Int32( - self.cfg.route_metadata_base_word_offset - ) + route_ordinal * cutlass.Int32(self.cfg.route_metadata_stride_words) - logical_origin = cutlass.Int32(-1) - logical_origin_is_valid = cutlass.Boolean(False) - physical_page_id = cutlass.Int32(-1) - atom_is_full = cutlass.Boolean(False) - if lane_idx < cutlass.Int32(self.cfg.logical_origins_per_route): - ( - logical_origin, - logical_origin_is_valid, - ) = _resolve_route_logical_atom_origin( - block_indices, - row_begin, - row_end, - route_idx, + if row_is_valid: + route_idx = cutlass.Int32(0) + while route_idx < exact_route_count: + route_word_index = cutlass.Int32( + self.cfg.route_metadata_base_word_offset + ) + (row_route_begin + route_idx) * cutlass.Int32( + self.cfg.route_metadata_stride_words + ) + logical_origin = cutlass.Int32(-1) + logical_origin_is_valid = cutlass.Boolean(False) + physical_page_id = cutlass.Int32(-1) + if lane_idx < cutlass.Int32(self.cfg.logical_origins_per_route): + ( + logical_origin, + logical_origin_is_valid, + ) = _resolve_route_logical_atom_origin( + block_indices, + row_begin, + row_end, + route_idx, + lane_idx, + self.cfg.kv_block_size, + self.cfg.atom_size, + self.cfg.logical_origins_per_route, + live_seq_len_kv, + ) + if cutlass.const_expr(self.route_layout.is_paged): + physical_page_id = _resolve_paged_route_atom_page_id( + block_tables, + batch_idx, + block_table_row_stride, + logical_origin, + logical_origin_is_valid, + lane_idx, + self.page_size, + num_physical_kv_pages, + ) + if lane_idx < cutlass.Int32(self.cfg.logical_origins_per_route): + route_workspace[route_word_index + lane_idx] = logical_origin + if cutlass.const_expr(self.route_layout.is_paged): + route_workspace[ + route_word_index + + cutlass.Int32(self.physical_page_ids_word_offset) + + lane_idx + ] = physical_page_id + cute.arch.sync_warp() + + _finalize_exact_route( + route_workspace, + kv_valid_bits, + route_word_index, + batch_idx, lane_idx, - self.cfg.kv_block_size, - self.cfg.atom_size, - self.cfg.logical_origins_per_route, + logical_origin_is_valid, live_seq_len_kv, + self.cfg, ) - if cutlass.const_expr(self.route_layout.is_paged): - physical_page_id = _resolve_paged_route_atom_page_id( - paged_kv_indices, - request_begin, - logical_origin, - logical_origin_is_valid, - lane_idx, - self.page_size, - num_physical_kv_pages, + route_idx += cutlass.Int32(1) + + if cutlass.const_expr(self.cfg.use_proxy_routes): + group_idx = cutlass.Int32(0) + while group_idx < cutlass.Int32(self.cfg.num_proxy_groups): + proxy_word = cutlass.Uint32(0) + logical_word_idx = ( + group_idx * cutlass.Int32(self.cfg.token_words_per_route) + + lane_idx + ) + if lane_idx < cutlass.Int32(self.cfg.token_words_per_route): + if logical_word_idx < cutlass.Int32(self.cfg.num_exact_words): + proxy_word = _load_bsr_proxy_word( + block_indices, + row_begin, + row_end, + logical_word_idx, + self.cfg, + ) + _emit_proxy_route( + route_workspace, + row_route_begin, + exact_route_count, + group_idx, + proxy_word, + lane_idx, + self.cfg, + ) + group_idx += cutlass.Int32(1) + + if lane_idx == cutlass.Int32(0) and row_is_valid: + route_workspace[linear_row_idx] = total_route_count + + +class _PrepareBitmaskRoutes(_PrepareRoutesBase): + """Lower packed exact-block bits to continuous exact-first routes.""" + + def __init__( + self, + *, + batch_size: int, + num_kv_heads: int, + seq_len_q: int, + seq_len_kv: int, + q_block_size: int, + kv_block_size: int, + kv_route_size: int, + use_proxy_routes: bool, + use_causal_mask: bool = False, + apply_token_mask: bool = False, + store_score_words: bool = False, + ) -> None: + super().__init__( + batch_size=batch_size, + num_kv_heads=num_kv_heads, + seq_len_q=seq_len_q, + seq_len_kv=seq_len_kv, + q_block_size=q_block_size, + kv_block_size=kv_block_size, + kv_route_size=kv_route_size, + use_proxy_routes=use_proxy_routes, + use_causal_mask=use_causal_mask, + apply_token_mask=apply_token_mask, + store_score_words=store_score_words, + ) + + @cute.jit + def __call__( + self, + 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, + stream: cuda_drv.CUstream, + ) -> None: + self.kernel( + exact_block_bits, + kv_valid_bits, + row_route_offsets, + route_workspace, + max_blocks_per_row, + ).launch( + grid=[ + (self.cfg.num_rows + _WARPS_PER_CTA - 1) // _WARPS_PER_CTA, + 1, + 1, + ], + block=[_THREADS_PER_CTA, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + 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, + ) -> None: + """Pack one bitmask row after proving its complete payload fits.""" + + thread_idx, _, _ = cute.arch.thread_idx() + block_idx, _, _ = cute.arch.block_idx() + warp_idx = thread_idx // _WARP_SIZE + lane_idx = thread_idx % _WARP_SIZE + linear_row_idx = block_idx * _WARPS_PER_CTA + warp_idx + row_is_valid = linear_row_idx < self.cfg.num_rows + q_block_idx = linear_row_idx % self.cfg.num_q_blocks + linear_batch_head_idx = linear_row_idx // self.cfg.num_q_blocks + kv_head_idx = linear_batch_head_idx % self.cfg.num_kv_heads + batch_idx = linear_batch_head_idx // self.cfg.num_kv_heads + + lane_exact_count = cutlass.Int32(0) + word_idx = lane_idx + while word_idx < cutlass.Int32(self.cfg.num_exact_words): + if row_is_valid: + exact_word = _load_bitmask_word( + exact_block_bits, + batch_idx, + kv_head_idx, + q_block_idx, + word_idx, + self.cfg, + for_proxy=False, + ) + lane_exact_count += cutlass.Int32(cute.arch.popc(exact_word)) + word_idx += cutlass.Int32(_WARP_SIZE) + exact_block_count = cutlass.Int32( + cute.arch.warp_redux_sync(lane_exact_count, "add") + ) + + exact_atom_count, exact_route_count, total_route_count = _prepared_route_counts( + exact_block_count, + self.cfg, + ) + if lane_idx == cutlass.Int32(0) and row_is_valid: + runtime_assert( + exact_block_count <= max_blocks_per_row, + "selected bitmask blocks exceed planned semantic capacity", + ) + row_route_begin = _prepared_row_route_begin( + row_route_offsets, + linear_row_idx, + lane_idx, + row_is_valid, + total_route_count, + ) + + if row_is_valid: + exact_prefix = cutlass.Int32(0) + word_idx = cutlass.Int32(0) + while word_idx < cutlass.Int32(self.cfg.num_exact_words): + exact_word_i32 = cutlass.Int32(0) + if lane_idx == cutlass.Int32(0): + exact_word_i32 = _load_bitmask_word( + exact_block_bits, + batch_idx, + kv_head_idx, + q_block_idx, + word_idx, + self.cfg, + for_proxy=False, + ).bitcast(cutlass.Int32) + exact_word = _warp_broadcast_i32(exact_word_i32, 0).bitcast( + cutlass.Uint32 ) - if logical_origin_is_valid: - atom_is_full = cutlass.Boolean( - logical_origin - <= live_seq_len_kv - cutlass.Int32(self.cfg.atom_size) + is_exact = cutlass.Boolean( + (exact_word & (cutlass.Uint32(1) << lane_idx)) != cutlass.Uint32(0) ) - if lane_idx < cutlass.Int32(self.cfg.logical_origins_per_route): - route_workspace[route_metadata_word_index + lane_idx] = logical_origin - if cutlass.const_expr(self.route_layout.is_paged): - route_workspace[ - route_metadata_word_index - + cutlass.Int32(self.physical_page_ids_word_offset) - + lane_idx - ] = physical_page_id + exact_ballot = cute.arch.vote_ballot_sync(is_exact).bitcast( + cutlass.Uint32 + ) + exact_rank = _exact_lane_rank(exact_ballot, lane_idx, exact_prefix) + if is_exact: + _emit_exact_block_atoms( + route_workspace, + row_route_begin, + word_idx * cutlass.Int32(_WARP_SIZE) + lane_idx, + exact_rank, + self.cfg, + ) + exact_prefix += cutlass.Int32(cute.arch.popc(exact_ballot)) + word_idx += cutlass.Int32(1) - _store_prepared_route_validity( - block_indices, - kv_valid_bits, - route_workspace, - row_begin, - row_end, - route_idx, - batch_idx, - lane_idx, - logical_origin, - logical_origin_is_valid, - atom_is_full, - route_metadata_word_index, - live_seq_len_kv, - self.cfg, + final_route_atom_count = exact_atom_count % cutlass.Int32( + self.cfg.logical_origins_per_route ) - route_idx = route_idx + cutlass.Int32(1) + if final_route_atom_count != cutlass.Int32(0): + if lane_idx >= final_route_atom_count and lane_idx < cutlass.Int32( + self.cfg.logical_origins_per_route + ): + final_route_word_index = cutlass.Int32( + self.cfg.route_metadata_base_word_offset + ) + (row_route_begin + exact_route_count - cutlass.Int32(1)) * ( + cutlass.Int32(self.cfg.route_metadata_stride_words) + ) + route_workspace[final_route_word_index + lane_idx] = cutlass.Int32( + -1 + ) + cute.arch.sync_warp() + + route_idx = cutlass.Int32(0) + while route_idx < exact_route_count: + route_word_index = cutlass.Int32( + self.cfg.route_metadata_base_word_offset + ) + (row_route_begin + route_idx) * cutlass.Int32( + self.cfg.route_metadata_stride_words + ) + atom_is_valid = cutlass.Boolean(False) + if lane_idx < cutlass.Int32(self.cfg.logical_origins_per_route): + atom_is_valid = cutlass.Boolean( + cutlass.Int32(route_workspace[route_word_index + lane_idx]) + >= cutlass.Int32(0) + ) + _finalize_exact_route( + route_workspace, + kv_valid_bits, + route_word_index, + batch_idx, + lane_idx, + atom_is_valid, + cutlass.Int32(self.cfg.seq_len_kv), + self.cfg, + ) + route_idx += cutlass.Int32(1) + + if cutlass.const_expr(self.cfg.use_proxy_routes): + group_idx = cutlass.Int32(0) + while group_idx < cutlass.Int32(self.cfg.num_proxy_groups): + proxy_word = cutlass.Uint32(0) + logical_word_idx = ( + group_idx * cutlass.Int32(self.cfg.token_words_per_route) + + lane_idx + ) + if lane_idx < cutlass.Int32(self.cfg.token_words_per_route): + if logical_word_idx < cutlass.Int32(self.cfg.num_exact_words): + proxy_word = _load_bitmask_word( + exact_block_bits, + batch_idx, + kv_head_idx, + q_block_idx, + logical_word_idx, + self.cfg, + for_proxy=True, + ) + _emit_proxy_route( + route_workspace, + row_route_begin, + exact_route_count, + group_idx, + proxy_word, + lane_idx, + self.cfg, + ) + group_idx += cutlass.Int32(1) + + if lane_idx == cutlass.Int32(0) and row_is_valid: + route_workspace[linear_row_idx] = total_route_count + + +__all__ = ["_PrepareBitmaskRoutes", "_PrepareBsrRoutes"] diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_config.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_config.py index c67e4c83bf10..0626410d1f10 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 @@ -102,6 +102,16 @@ (Float16, Float16, Float16, 256, 128, 1, 1), } +# 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 +# the two softmax groups and the correction group share the remainder: +# 8 * softmax + 4 * correction = 65536 / 32 - 4 * 56. An even split measured +# fastest on B200 for both the static grid and the persistent scheduler; the +# correction group needs the extra room for its persistent bookkeeping and +# the KV256 tail merge, while the rolled softmax fragment loop needs less. +KV_TILE_256_SOFTMAX_TASK_REGISTERS = 152 +KV_TILE_256_CORRECTION_TASK_REGISTERS = 152 + _KV_TILE_256_PHYSICAL_DEFAULTS: Mapping[str, ConfigValue] = { "tmem_s_cols": 128, "tmem_stats_cols": 32, @@ -505,6 +515,10 @@ class FmhaDecodeConfig: # restricted to 8/16/32 or positive multiples of 64 and are assembled into # a profile-selected fixed KV128 or KV256 route. use_block_sparse: bool = False + # Interpret prepared records as a typed proxy/exact stream. Proxy records + # source semantic-block summaries while exact records retain the K/V + # atom path. Source selection is orthogonal to the physical Q/KV profile. + use_block_sparse_proxy_routes: bool = False q_block_size: int = 0 kv_block_size: int = 0 # Optional batch-wide physical-token validity metadata shared by every head @@ -644,13 +658,13 @@ def uses_task_register_reallocation(self) -> bool: def softmax_task_num_registers(self) -> int | None: if not self.uses_task_register_reallocation: return None - return 176 if self.tile_size_kv == 256 else 184 + return KV_TILE_256_SOFTMAX_TASK_REGISTERS if self.tile_size_kv == 256 else 184 @property def correction_task_num_registers(self) -> int | None: if not self.uses_task_register_reallocation: return None - return 104 if self.tile_size_kv == 256 else 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: @@ -1066,12 +1080,13 @@ def num_s_regs_per_thread(self) -> int: def softmax_score_fragment_regs(self) -> int: """Return the maximum score fragment kept live in registers. - KV256 owns 128 score values per lane but streams them as four native - 32-register LDTM atoms. Other profiles retain their complete score - fragment, so this property is intentionally distinct from - ``num_s_regs_per_thread`` (the total logical ownership). + Streamed profiles own 128 score values per lane but process them as + four native 32-register LDTM atoms. Other profiles retain their + complete score fragment, so this property is intentionally distinct + from ``num_s_regs_per_thread`` (the total logical ownership). The + selection of streamed profiles lives in ``streams_tmem_p_fragments``. """ - if self.tile_size_kv == 256: + if self.streams_tmem_p_fragments: return 32 return self.num_s_regs_per_thread @@ -1080,6 +1095,30 @@ 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 block_sparse_kv_atom_size(self) -> int: + """Return the K token span of one block-sparse route origin.""" + assert self.use_block_sparse + return _block_sparse_kv_atom_size(self.kv_block_size) + + @property + def softmax_fragments_per_route_atom(self) -> int: + """Return the streamed score fragments that share one route origin. + + Route origins are staged per K64 atom, so a 128-token KV block spans + two origins; the fragment-to-origin mapping follows the atom. + """ + return self.block_sparse_kv_atom_size // self.softmax_score_fragment_regs + + @property + def uses_ws_2x2_datapath(self) -> bool: + """Whether QK and PV issue the WS 2x2 instruction over two lane halves. + + KV256 exposes two spatial KV128 partials per logical Q row; every other + Keeps profile issues the plain CTA-local instruction. + """ + return self.tile_size_kv == 256 + @property def num_packed_p_regs(self) -> int: """Return packed P registers stored by each softmax producer lane.""" @@ -1246,6 +1285,10 @@ def validate_paged_kv_staging_config(self) -> None: def validate_block_sparse_profile(self, *, heads_q_per_kv: int) -> None: """Validate the qualified host profile for block-sparse.""" + if self.use_block_sparse_proxy_routes and not self.use_block_sparse: + raise ValueError("proxy routes require block-sparse attention") + if self.use_block_sparse_proxy_routes and self.mask_type != DENSE: + raise ValueError("block-sparse proxy routes require mask_type='dense'") if not self.use_block_sparse: if self.use_parallel_sparse_kv_loads: raise ValueError( @@ -1283,6 +1326,13 @@ def validate_block_sparse_profile(self, *, heads_q_per_kv: int) -> None: "block-sparse tile_size_kv=256 requires the Q64 16-bit 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)" + ) if self.tile_size_q != selected_q_tile: raise ValueError( "block-sparse tile_size_q must match its grouped-Q geometry" @@ -1332,6 +1382,42 @@ def compile_signature(self) -> tuple[tuple[str, object], ...]: for config_field in fields(self) ) + @property + def uses_prepared_score_keep_words(self) -> bool: + """Whether prepared routes carry BMM1 score-column validity words. + + Dense block-sparse Keeps plans prepare them even without a caller + token mask: the streamed max pass trusts the words directly, which is + cheaper than deriving each fragment's visible range in the softmax + warps. The plan sizes its route storage and the prepare kernels store + the words from this same property, via the resolved launch spec. + """ + + return ( + self.use_kv_valid_bits + or self.use_block_sparse_proxy_routes + or ( + self.use_block_sparse + and self.use_keeps_mma_ab + and self.mask_type == DENSE + ) + ) + + @property + def trusts_prepared_score_words(self) -> bool: + """Whether prepared words fully describe dense score-column validity. + + Dense prepared routes have already combined structural tail validity + with any caller-provided exact-token bits. Their K32 words therefore + apply to exact and proxy sources alike. + """ + + return ( + self.use_block_sparse + and self.uses_prepared_score_keep_words + and self.mask_type == DENSE + ) + @property def uses_q_desc_ref(self) -> bool: """Whether QK derives Q's descriptor from shared resource state.""" @@ -1435,10 +1521,10 @@ def has_static_dense_full_kv_tiles(self) -> bool: @property def uses_ordered_softmax_barrier(self) -> bool: """Whether this profile selects the ordered P0/P1 softmax barrier.""" - if self.tile_size_kv == 256: - # KV256 uses independent four-stage P-fragment pipelines. Ordering - # the two softmax groups would serialize fragment production and - # defeat the intended P/PV overlap. + if self.streams_tmem_p_fragments: + # Streamed profiles use independent per-fragment P pipelines. + # Ordering the two softmax groups would serialize fragment + # production and defeat the intended P/PV overlap. return False if self.ordered_softmax_barrier_mode == 2: return True @@ -1554,10 +1640,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. - Q128/KV128 and sparse Q64/KV128 publish a complete packed-P row per - pipeline token. Q64/KV256 uses the same S-to-P aliasing contract but - streams four independently ready K32 fragments. Dense Q64/KV128 keeps - the base kernel's faster SMEM-P cadence. + 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. """ # 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 @@ -1568,11 +1655,6 @@ def uses_two_inst_tmem_p(self) -> bool: and ( (self.tile_size_q == 128 and self.tile_size_kv == 128) or (self.tile_size_q == 64 and self.tile_size_kv == 256) - or ( - self.use_block_sparse - and self.tile_size_q == 64 - and self.tile_size_kv == 128 - ) ) and self.head_dim_per_stage_kv == 0 and self.num_insts_kv == 2 @@ -1582,42 +1664,51 @@ def uses_two_inst_tmem_p(self) -> bool: @property def streams_tmem_p_fragments(self) -> bool: - """Whether P is published as independently ready TMEM fragments.""" - return self.uses_two_inst_tmem_p and self.num_softmax_score_fragments > 1 - - @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() + """Whether P is published as independently ready TMEM fragments. + + Streamed profiles produce their K32 fragments from one rolled runtime + loop: the max pass writes masked scores back to TMEM, so the P pass + reloads each fragment without mask logic and the exponentiation body + exists once in the instruction stream. Each published fragment lets + 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. + """ + return ( + self.uses_two_inst_tmem_p + and not self.use_fp8_qkv + and (self.tile_size_kv == 256 or self.use_block_sparse) ) @property - def uses_rotating_kv256_exchange(self) -> bool: - """Whether this profile selects KV-ring scratch for correction. + def defers_softmax_anchor_updates(self) -> bool: + """Whether small row-max increases keep the previous exponent anchor. - Persistent direct output can overlap the next work tile's first two - K loads with correction by placing its exchange in the third, drained - KV stage. Split-KV and attention sinks retain the fixed exchange because - their tail storage and lifetime differ from direct output. + Keeps correction skips the in-place O rescale whenever the anchor is + unchanged, so keeping the prior anchor within + ``SOFTMAX_RESCALE_THRESHOLD_LOG2`` trades a bounded 16-bit P range + (2**8) for fewer TMEM rescales. The profiles listed here are the ones + where that trade was measured to pay: KV256 tiles and block-sparse + routes, whose row maximum moves often but rarely by much. """ - selects_persistent_kv256 = ( - self.streams_tmem_p_fragments - and self.tile_size_q == 64 - and self.tile_size_kv == 256 - and self.use_persistent_scheduler + return self.use_keeps_mma_ab and ( + self.tile_size_kv == 256 or self.use_block_sparse ) - if not selects_persistent_kv256: - return False - has_rotating_kv_ring = ( - self.num_head_dim_stages_kv == 1 - and self.kv_stages == KV_TILE_256_SHARED_FIFO_STAGES - and self.load_num_warps == 1 + @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() ) - has_direct_output_lifetime = not (self.use_split_kv or self.use_attention_sinks) - return has_rotating_kv_ring and has_direct_output_lifetime @property def keeps_separates_tmem_s_and_stats(self) -> bool: @@ -2196,17 +2287,6 @@ def _require_python_int(field_name: str) -> int: f"{pipeline_smem_bytes} bytes, limit is " f"{pipeline_smem_budget_bytes} bytes" ) - if ( - cfg.use_persistent_scheduler - and not cfg.use_split_kv - and not cfg.use_attention_sinks - and cfg.kv_stages != KV_TILE_256_SHARED_FIFO_STAGES - ): - raise ValueError( - "persistent KV256 requires kv_stages=" - f"{KV_TILE_256_SHARED_FIFO_STAGES} for the rotating shared-KV " - f"exchange, got {cfg.kv_stages}" - ) if not cfg.supports_grouped_keeps: raise ValueError( "KV256 currently supports only the qualified Q64 FP16/BF16/D128 " 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 68686febfbbb..8bf94760cf31 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 @@ -45,8 +45,12 @@ # Keep the old maximum as the exponent reference while a new maximum is at # most eight log2 units larger. This avoids an output-correction round without # letting an intermediate probability exceed 2**8; the softmax identity is -# unchanged apart from normal finite-precision rounding. -KV_TILE_256_RESCALE_THRESHOLD_LOG2 = 8.0 +# unchanged apart from normal finite-precision rounding. As in the +# FlashInfer/TRT-LLM policy, this assumes normal model logits rather than +# adversarial values outside the qualified probability bound. Streamed KV256 +# and block-sparse Keeps profiles apply it; see +# ``FmhaDecodeConfig.defers_softmax_anchor_updates``. +SOFTMAX_RESCALE_THRESHOLD_LOG2 = 8.0 # A launch bound makes ptxas honor warpgroup ``setmaxnreg`` allocations, but # the resulting register hand-off has a fixed cost. Paired B200 measurements 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 929696c30d3e..d2540ec61920 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 @@ -56,8 +56,8 @@ from cutlass.experimental.task_scheduling.task_manager import TaskManager from ..._block_sparse.common import ( - _block_sparse_kv_atom_size, - _prepared_kv_routes_are_block_aligned, + _block_sparse_contiguous_kv_copy_geometry, + _block_sparse_proxy_summary_geometry, ) from ..._block_sparse.prepared import _BlockSparseRouteLayout from ..tensor_map import ( @@ -95,7 +95,7 @@ from .fmha_decode_tasks import ( PackedDecodeWorkQueue, ScheduleTokenThrottleResource, - SmemKvReuseCreditResource, + _prefetch_prepared_sparse_row, create_block_sparse_load_tasks_per_inst, create_correction_task, create_correction_task_one_inst_qkv, @@ -113,7 +113,6 @@ create_softmax0_task, create_softmax1_task, ) - from .reduction import ( # noqa: F401 decode_gen_separate_reduction_kernel, fmha_decode_separate_reduction_launch, @@ -322,6 +321,10 @@ def _build_decode_gen_schedule( tma_desc_v: 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, + tma_desc_v_summary: cutlass.Pointer | None = None, + 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, h_k_idx: Int32 | None = None, b_idx: Int32 | None = None, @@ -341,6 +344,8 @@ def _build_decode_gen_schedule( 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, ) -> tuple[ list[Task], dict[MemoryResource, list[MemoryResource]], @@ -408,6 +413,15 @@ def _build_decode_gen_schedule( "tma_desc_k_atom": tma_desc_k_atom, "tma_desc_v_atom": tma_desc_v_atom, } + if cfg.use_block_sparse_proxy_routes: + segment_tensormaps.update( + { + "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, + } + ) for name, descriptor in segment_tensormaps.items(): if descriptor is None: raise ValueError( @@ -741,7 +755,6 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: # ------------------------------------------------------------------ work_queue = None schedule_token_throttle = None - smem_kv_reuse_credit = None # CLC remains the single persistent policy for every supported topology. # The stock static WorkQueue advances and decodes coordinates separately # in every task, which regresses multi-wave decode workloads. CLC computes @@ -799,17 +812,6 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: ), name="schedule_token_throttle", ) - if cfg.uses_rotating_kv256_exchange: - smem_kv_reuse_credit = SmemKvReuseCreditResource( - cfg=cfg, - pipeline_config=PipelineConfig.create_async_async_pipeline_cfg( - num_stages=1, - producer_group=load_grp, - consumer_group=correction_grp, - cta_layout_vmnk=cta_layout, - ), - name="smem_kv_reuse_credit", - ) smem_q = SmemQResource( pipeline_config=smem_q_cfg, cfg=cfg, @@ -873,10 +875,12 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: sparse_softmax_metadata0 = None sparse_softmax_metadata1 = None 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. prepared_route_layout = _BlockSparseRouteLayout.create( kv_route_size=cfg.tile_size_kv, kv_block_size=cfg.kv_block_size, - has_token_bits=cfg.use_kv_valid_bits, + has_token_bits=cfg.uses_prepared_score_keep_words, route_metadata_capacity=0, num_rows=1, page_size=cfg.num_tokens_per_page if cfg.use_paged_kv else None, @@ -928,6 +932,10 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: 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_metadata0, page_offsets_kv=smem_page_offsets, seqlens_kv=kv_seqlens, @@ -947,6 +955,10 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: 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, @@ -966,6 +978,10 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: 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_metadata0, page_offsets_kv=smem_page_offsets_v or smem_page_offsets, seqlens_kv=kv_seqlens, @@ -985,6 +1001,10 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: 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_v or smem_page_offsets, seqlens_kv=kv_seqlens, @@ -1005,6 +1025,10 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: 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, @@ -1213,6 +1237,8 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: "seq_len_q": seq_len_q, "sparse_row_route_offsets": sparse_row_route_offsets, "sparse_row_route_counts": sparse_row_route_counts, + "sparse_row_route_begin": sparse_row_route_begin, + "sparse_route_count": sparse_route_count, "num_heads_kv": num_heads_kv, } if use_one_inst_qkv: @@ -1283,7 +1309,6 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: smem_kv, work_queue, schedule_token_throttle, - smem_kv_reuse_credit, cfg, domain=load_domain, domain_bias=0, @@ -1431,7 +1456,6 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: tmem_corr0, tmem_corr1, work_queue, - smem_kv_reuse_credit, cfg, domain=corr_domain, tmem_stats_done0=tmem_stats_done0, @@ -1607,10 +1631,6 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: ) if schedule_token_throttle is not None: resource_dependency_graph[schedule_token_throttle] = [work_queue] - if smem_kv_reuse_credit is not None: - # A self-edge models the one-slot ownership token: Load produces it - # for the current tile and Correction consumes it before the next Load. - resource_dependency_graph[smem_kv_reuse_credit] = [smem_kv_reuse_credit] dma_consumer_release_labels: dict[ tuple[MemoryResource, MemoryResource], set[str] ] = {} @@ -1662,8 +1682,6 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: smem_allocator.add_resource(work_queue) if schedule_token_throttle is not None: smem_allocator.add_resource(schedule_token_throttle) - if smem_kv_reuse_credit is not None: - smem_allocator.add_resource(smem_kv_reuse_credit) smem_allocator.add_resource(smem_q) if smem_page_offsets is not None: smem_allocator.add_resource(smem_page_offsets) @@ -1701,21 +1719,10 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: smem_allocator.add_resource(tmem_corr0) if not use_one_inst_qkv: smem_allocator.add_resource(tmem_corr1) - if cfg.tile_size_kv == 256: - # KV256 direct-output correction rotates one compact 35,840-byte - # payload through the shared 192-KiB K/V ring. Split-KV retains the - # fixed full exchange. Neither path increases the CTA SMEM footprint. - smem_allocator.add_alias_group( - [ - smem_kv.get_smem_requirements(), - tmem_corr1.get_smem_requirements(), - ] - ) 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: @@ -1774,13 +1781,15 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: p0_alloc = smem_p0.get_tmem_requirements()[0] p1_alloc = smem_p1.get_tmem_requirements()[0] o_alloc = tmem_o.get_tmem_requirements()[0] - if cfg.tile_size_kv == 256: - # KV256 keeps O in the low 256 columns and overlays packed P on - # each S region from its first column. Softmax streams K32 - # fragments in order, so every 16-column P store only overwrites - # scores that have already been consumed. Starting P after the - # nominal stats columns would instead clobber the next unread S - # fragment; KV256 keeps its softmax stats in SMEM. + if cfg.streams_tmem_p_fragments: + # Streamed profiles keep O in the low 256 columns and overlay + # packed P on each S region from its first column. Softmax streams + # K32 fragments in order, so every 16-column P store only + # overwrites scores that have already been consumed. Starting P + # after the nominal stats columns would instead clobber the next + # unread S fragment; streamed profiles keep their softmax stats in + # SMEM. + assert cfg.keeps_stats_via_smem o_alloc.offset = 0 s0_alloc.offset = 2 * cfg.tmem_o_stage_cols s1_alloc.offset = s0_alloc.offset + cfg.tmem_s_cols @@ -1835,12 +1844,8 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: eager_init_resources = ( [tmem_corr0] if use_one_inst_qkv else [tmem_corr0, tmem_corr1] ) - if smem_kv_reuse_credit is not None: - # Initialize the persistent ring cursor under the same CTA-wide fence - # and barrier used by other manually managed SMEM control state. - eager_init_resources.append(smem_kv_reuse_credit) - if cfg.tile_size_kv == 256: - # KV256's TMEM P operands use one-way per-fragment ready barriers. + 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]) @@ -1863,10 +1868,10 @@ def _has_unmodeled_tmem_p_alias_protocol(cfg: FmhaDecodeConfig) -> bool: """Whether exhaustive TS checking would report a known false P/S race. The staged D256 path selects one of two physical P/S stages at runtime. - Static KV256 instead orders streamed P fragments with private mbarriers and + Static streamed profiles instead order P fragments with private mbarriers and reuses the matching TmemO-full barrier as the next-QK overwrite credit. Those intra-work protocols are below TaskManager's resource transitions, - so its allocation-level checker cannot prove them. Persistent KV256 has + so its allocation-level checker cannot prove them. Persistent streaming has enough task-level ordering for the checker and remains covered. """ return cfg.uses_staged_one_inst_tmem_p or ( @@ -1982,6 +1987,10 @@ def _run_decode_gen_active( 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, + tma_desc_k_summary: cutlass.GridConstant[cuda.TensorMap] | None = None, + 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, ) -> None: """Run the complete decode body for one runtime-valid Q tile. @@ -2021,28 +2030,43 @@ def _run_decode_gen_active( 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 + tma_desc_v_summary_atom_ptr = None + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): + assert tma_desc_k_summary is not None + assert tma_desc_v_summary is not None + assert tma_desc_k_summary_atom is not None + assert tma_desc_v_summary_atom is not None + tma_desc_k_summary_ptr = tma_desc_k_summary.get_ptr() + tma_desc_v_summary_ptr = tma_desc_v_summary.get_ptr() + tma_desc_k_summary_atom_ptr = tma_desc_k_summary_atom.get_ptr() + tma_desc_v_summary_atom_ptr = tma_desc_v_summary_atom.get_ptr() # Prefetch TMA + uses_atom_desc = False + if cutlass.const_expr(cfg.use_block_sparse): + _, _, uses_atom_desc = _block_sparse_contiguous_kv_copy_geometry( + kv_block_size=cfg.kv_block_size, + kv_route_size=cfg.tile_size_kv, + ) init_warp = 1 if warp_idx == init_warp: 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_block_sparse - and _block_sparse_kv_atom_size(cfg.kv_block_size) == 64 - and ( - cfg.tile_size_kv == 256 - or not _prepared_kv_routes_are_block_aligned( - cfg.kv_block_size, - cfg.tile_size_kv, - ) - ) - ): - # KV256 always issues semantic KV64 atoms. KV128 needs this second - # descriptor only for non-aligned coarse routes. + 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()) prims.prefetch_tensormap(tma_desc_v_atom.get_ptr()) + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): + if cutlass.const_expr(cfg.tile_size_kv != 256): + prims.prefetch_tensormap(tma_desc_k_summary_ptr) + prims.prefetch_tensormap(tma_desc_v_summary_ptr) + if cutlass.const_expr(uses_atom_desc): + prims.prefetch_tensormap(tma_desc_k_summary_atom_ptr) + prims.prefetch_tensormap(tma_desc_v_summary_atom_ptr) init_warp += 1 clc_response_ptr = None @@ -2055,6 +2079,27 @@ def _run_decode_gen_active( if cutlass.const_expr(cfg.max_seq_len_q > 1): q_output_rows = g_h_r * Int32(cfg.max_seq_len_q) + # Static block-sparse tiles read their prepared row header here, before + # TMEM allocation and barrier setup, so that global-memory round trip is + # hidden instead of stalling every task at its first schedule step. + sparse_row_route_begin = None + sparse_route_count = None + if cutlass.const_expr( + cfg.use_block_sparse + and not cfg.use_persistent_scheduler + and g_sparse_row_route_offsets is not None + and g_sparse_row_route_counts is not None + ): + sparse_row_route_begin, sparse_route_count = _prefetch_prepared_sparse_row( + cfg, + g_sparse_row_route_offsets, + g_sparse_row_route_counts, + q_group_idx, + h_k_idx, + b_idx, + g_h_k, + ) + ( task_list, dep_graph, @@ -2083,6 +2128,10 @@ def _run_decode_gen_active( tma_desc_v=tma_desc_v.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, page_idx_kv=g_page_idx_kv, h_k_idx=h_k_idx, b_idx=b_idx, @@ -2102,6 +2151,8 @@ def _run_decode_gen_active( 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, + sparse_row_route_begin=sparse_row_route_begin, + sparse_route_count=sparse_route_count, ) smem_allocator.allocate() @@ -2264,6 +2315,10 @@ def _run_decode_gen_runtime_prefix( g_sparse_row_route_offsets: cute.Pointer | None, g_sparse_row_route_counts: cute.Pointer | None, g_sparse_route_metadata: cute.Pointer | None, + tma_desc_k_summary: cutlass.GridConstant[cuda.TensorMap] | None = None, + 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, ) -> None: """Run the general runtime split-prefix producer or retire its suffix.""" @@ -2328,6 +2383,10 @@ def _run_decode_gen_runtime_prefix( g_sparse_row_route_offsets, g_sparse_row_route_counts, g_sparse_route_metadata, + 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, ) else: _run_decode_gen_inactive_cluster_rank() @@ -2371,6 +2430,10 @@ def _run_decode_gen_runtime_prefix( g_sparse_row_route_offsets, g_sparse_row_route_counts, g_sparse_route_metadata, + 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, ) else: _signal_padded_pdl_producer(cfg) @@ -2409,6 +2472,10 @@ def decode_gen_kernel( g_sparse_row_route_counts: cute.Pointer | None = None, g_sparse_route_metadata: cute.Pointer | None = None, static_full_split_prefix: cutlass.Constexpr[bool] = False, + tma_desc_k_summary: cutlass.GridConstant[cuda.TensorMap] | None = None, + 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, ) -> 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() @@ -2470,6 +2537,10 @@ def decode_gen_kernel( g_sparse_row_route_offsets, g_sparse_row_route_counts, g_sparse_route_metadata, + 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, ) else: _run_decode_gen_runtime_prefix( @@ -2510,6 +2581,10 @@ def decode_gen_kernel( g_sparse_row_route_offsets, g_sparse_row_route_counts, g_sparse_route_metadata, + 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, ) else: # Packed-Q grids use a batch-wide maximum envelope. These Q CTAs own no @@ -2754,7 +2829,6 @@ def fmha_decode_launch( tma_desc_q, tma_desc_k, tma_desc_v, - # Dense/paged profiles never inspect the 64-token descriptor slots. tma_desc_k, tma_desc_v, o_iter, @@ -2799,6 +2873,8 @@ def fmha_block_sparse_launch( q_iter: cute.Pointer, k_iter: cute.Pointer, v_iter: cute.Pointer, + k_summary_iter: cute.Pointer, + v_summary_iter: cute.Pointer, o_iter: cute.Pointer, row_route_offsets_iter: cute.Pointer, row_route_counts_iter: cute.Pointer, @@ -2813,14 +2889,18 @@ def fmha_block_sparse_launch( k_page_stride: Int64 = 0, v_page_stride: Int64 = 0, ) -> None: - """Launch attention over contiguous or paged prepared KV routes. + """Launch attention over exact and typed exact/proxy prepared KV routes. A preceding prepare kernel has already resolved each BSR row into compact logical atom origins, storage locators, validity flags, and optional token - words. Both layouts execute the same ``decode_gen_kernel`` schedule. + 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. """ if cutlass.const_expr(not cfg.use_block_sparse): - raise ValueError("fmha_block_sparse_launch requires cfg.use_block_sparse=True") + raise ValueError("fmha_block_sparse_launch requires block-sparse config") + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes and cfg.use_paged_kv): + raise ValueError("block-sparse proxy routes require contiguous K/V") log2_e = math.log2(math.e) b, h_q, h_k, s_k, d = problem_shape @@ -2855,7 +2935,18 @@ def fmha_block_sparse_launch( swizzle=tma_swizzle, ) - kv_atom_size = _block_sparse_kv_atom_size(cfg.kv_block_size) + ( + primary_kv_box_size, + kv_atom_size, + uses_atom_desc, + ) = _block_sparse_contiguous_kv_copy_geometry( + kv_block_size=cfg.kv_block_size, + kv_route_size=cfg.tile_size_kv, + ) + k_desc_summary_primary = None + v_desc_summary_primary = None + k_desc_summary_atom = None + v_desc_summary_atom = None if cutlass.const_expr(cfg.use_paged_kv): # Paged HND storage is addressed as (D, token-in-page, Hkv, page). # Prepared routes already contain each atom's physical page ID, so no @@ -2899,9 +2990,10 @@ def fmha_block_sparse_launch( k_desc_primary = k_desc_atom v_desc_primary = v_desc_atom else: - # Contiguous sparse coordinates retain the logical (D, S, H, B) - # order and the established primary/atom descriptor split. - primary_kv_box_size = 2 * kv_atom_size if kv_atom_size == 64 else kv_atom_size + # Exact and summary tensors form one logical segmented KV coordinate + # space. Each physical source owns the same primary/atom descriptor + # pair; the prepared route kind selects the pair, while the loader + # retains the existing KV128/fine/KV256 copy policy. kv_dims = (d, s_k, h_k, b) k_desc_primary = create_tensor_map_tiled( global_address=k_iter.toint(), @@ -2921,16 +3013,7 @@ def fmha_block_sparse_launch( ) k_desc_atom = k_desc_primary v_desc_atom = v_desc_primary - if cutlass.const_expr( - kv_atom_size == 64 - and ( - cfg.tile_size_kv == 256 - or not _prepared_kv_routes_are_block_aligned( - cfg.kv_block_size, - cfg.tile_size_kv, - ) - ) - ): + if cutlass.const_expr(uses_atom_desc): # KV256 always stages four semantic KV64 atoms. KV128 needs this # map only when a route may join unrelated BSR entries. k_desc_atom = create_tensor_map_tiled( @@ -2950,6 +3033,55 @@ def fmha_block_sparse_launch( swizzle=tma_swizzle, ) + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): + num_kv_blocks, _ = _block_sparse_proxy_summary_geometry( + seq_len_kv, + cfg.kv_block_size, + ) + _, summary_kv_strides = _block_sparse_bshd_tma_strides( + q_seq=q_seq, + h_q=h_q, + h_k=h_k, + s_k=num_kv_blocks, + d=d, + ) + 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, + global_dims=summary_dims, + global_strides=summary_kv_strides, + box_dims=(tma_box0, primary_kv_box_size, 1, 1), + swizzle=tma_swizzle, + ) + v_desc_summary_primary = create_tensor_map_tiled( + global_address=v_summary_iter.toint(), + dtype=cfg.kv_dtype, + global_dims=summary_dims, + global_strides=summary_kv_strides, + box_dims=(tma_box0, primary_kv_box_size, 1, 1), + swizzle=tma_swizzle, + ) + k_desc_summary_atom = k_desc_summary_primary + v_desc_summary_atom = v_desc_summary_primary + 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, + global_dims=summary_dims, + global_strides=summary_kv_strides, + box_dims=(tma_box0, kv_atom_size, 1, 1), + swizzle=tma_swizzle, + ) + v_desc_summary_atom = create_tensor_map_tiled( + global_address=v_summary_iter.toint(), + dtype=cfg.kv_dtype, + global_dims=summary_dims, + global_strides=summary_kv_strides, + box_dims=(tma_box0, kv_atom_size, 1, 1), + swizzle=tma_swizzle, + ) + q_groups = Int32( (cfg.max_seq_len_q + cfg.q_tokens_per_cta - 1) // cfg.q_tokens_per_cta ) @@ -3008,6 +3140,10 @@ def fmha_block_sparse_launch( row_route_counts_iter, route_metadata_iter, False, # static_full_split_prefix + tma_desc_k_summary=k_desc_summary_primary, + 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, ).launch( grid=grid, block=[cfg.threads_per_cta, 1, 1], diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_common.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_common.py index e3190f09097c..77412cddca7b 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 @@ -25,6 +25,8 @@ import cutlass import cutlass.cute as cute from cutlass import BFloat16, Float16, Float32, Int32, Int64, Uint32 +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import dsl_user_op from cutlass.experimental import primitives as prims from cutlass.experimental.task_scheduling.resources import ( @@ -59,6 +61,22 @@ ) ResourceVars = dict[str, ResourceVarValue] + +@dsl_user_op +def _assume_nonnegative_i32(value: Int32, *, loc=None, ip=None) -> Int32: + """Express a caller-guaranteed nonnegative Int32 contract to codegen.""" + + condition = cutlass.Boolean(value >= Int32(0)) + llvm.intr_assume( + condition.ir_value(loc=loc, ip=ip), + [], + [], + loc=loc, + ip=ip, + ) + return value + + # Offsets into DecodeGenTask.make_task_cache(). Keeping these symbolic makes # resource code explicit about which task-local lane or address value it needs. _TASK_CACHE_TMEM_BASE_OFFSET = 0 @@ -109,6 +127,33 @@ def _warp_broadcast_i32(value: Int32, source_lane: Constexpr[int]) -> Int32: ) +@cute.jit +def _swaps_routed_coordinate( + cfg: Constexpr[FmhaDecodeConfig], + lane_k_offset: Int32, + origin0: Int32, + origin1: Int32, + origin2: Int32, + origin3: Int32, + *, + token_group_idx: Constexpr[int], +) -> tuple[Int32, Int32]: + """Map one SWAP register group to its staged atom and logical coordinate.""" + + atom_size = min(cfg.kv_block_size, 32) + groups_per_atom = atom_size // 8 + origin_idx = token_group_idx // groups_per_atom + atom_origin = origin0 + if cutlass.const_expr(origin_idx == 1): + atom_origin = origin1 + elif cutlass.const_expr(origin_idx == 2): + atom_origin = origin2 + elif cutlass.const_expr(origin_idx == 3): + atom_origin = origin3 + token_offset = (token_group_idx % groups_per_atom) * 8 + 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 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 4585213cd80b..2dbe3c7e7ff8 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 @@ -71,6 +71,70 @@ def _pack_float4_to_fp8_e4m3_inline( ) +@cute.jit +def _combine_int_frac_ex2( + x_rounded: Float32, frac_ex2: Float32, *, loc=None, ip=None +) -> Float32: + """Scale ``frac_ex2`` by ``2**floor(x)`` through the FP32 exponent bits. + + ``x_rounded`` still carries the magic rounding constant, so its low + mantissa bits hold ``floor(x)``; shifting them into the exponent field and + adding the bits of the polynomial result multiplies by the integer power. + """ + return cute.arch.inline_ptx( + "{\n" + " .reg .b32 xi;\n" + " .reg .b32 fi;\n" + " .reg .b32 xe;\n" + " .reg .b32 oi;\n" + " mov.b32 xi, {$r0};\n" + " mov.b32 fi, {$r1};\n" + " shl.b32 xe, xi, 23;\n" + " add.s32 oi, xe, fi;\n" + " mov.b32 {$w0}, oi;\n" + "}", + write_only_types=[Float32], + read_only_args=[x_rounded, frac_ex2], + loc=loc, + ip=ip, + ) + + +@cute.jit +def _ex2_emulation_packed_f32x2(x: Float32, y: Float32) -> tuple[Float32, Float32]: + """Evaluate ``2**x`` and ``2**y`` on the FMA pipe instead of MUFU. + + Inputs are non-positive scaled scores minus the row maximum. The integer + part is split off with a magic-constant rounding add, the fraction in + [0, 1) goes through a degree-3 minimax polynomial, and the two parts are + recombined through the exponent bits. The relative error stays below the + BF16 rounding of the P operand, matching the emulation used by the dense + Blackwell FMHA kernels. + """ + fp32_round_int = float(2**23 + 2**22) + xy_clamped = (cute.arch.fmax(x, -127.0), cute.arch.fmax(y, -127.0)) + xy_rounded = cute.arch.add_packed_f32x2( + xy_clamped, (fp32_round_int, fp32_round_int), rnd="rm" + ) + xy_rounded_back = cute.arch.sub_packed_f32x2( + xy_rounded, (fp32_round_int, fp32_round_int) + ) + xy_frac = cute.arch.sub_packed_f32x2(xy_clamped, xy_rounded_back) + coeff = ( + 1.0, + 0.695146143436431884765625, + 0.227564394474029541015625, + 0.077119089663028717041015625, + ) + out = (coeff[3], coeff[3]) + for degree in cutlass.range_constexpr(2, -1, -1): + out = cute.arch.fma_packed_f32x2(out, xy_frac, (coeff[degree], coeff[degree])) + return ( + _combine_int_frac_ex2(xy_rounded[0], out[0]), + _combine_int_frac_ex2(xy_rounded[1], out[1]), + ) + + @cute.jit def _compute_fp8_p_regs_and_local_sums( scale_softmax_log2: Float32, 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 3254de0e79e7..f79f3a1e0566 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 @@ -43,6 +43,7 @@ from ...._block_sparse.prepared import ( _PREPARED_ROUTE_IS_FULL_FLAG, + _PREPARED_ROUTE_IS_PROXY_FLAG, _BlockSparseRouteLayout, ) from ...placeholder_helpers import _placeholder_smem_array @@ -56,7 +57,6 @@ DecodeGenResourceBase, ResourceVars, _decode_gen_task_cache, - _keeps_col_base, _sparse_task_cache_route_begin, _sparse_task_cache_route_count, _warp_broadcast_i32, @@ -65,13 +65,20 @@ # Keeps staging uses the low four bits for structural KV64 validity. Bit 4 # carries the conservative prepared summary that token masking can be skipped; -# structural, tail, and causal masking remain independent. +# structural, tail, and causal masking remain independent. The streamed Keeps +# max pass derives its keep words from the token words directly, so the bit is +# currently staged for the consumer but not read. _SOFTMAX_TOKEN_MASK_IS_FULL_FLAG = 1 << 4 - -# B8 SWAP origins are eight-token aligned, so bit 0 is free while the route is -# in Softmax's private staging payload. Reusing it avoids adding a word to every -# pipeline stage merely to forward prepare's route-full summary. -_SWAPS_PACKED_ROUTE_FULL_CLEAR_MASK = ~_PREPARED_ROUTE_IS_FULL_FLAG +# 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 +# adding a word to every pipeline stage for prepared FULL/PROXY route flags. +_SWAPS_PACKED_ROUTE_FLAGS_CLEAR_MASK = ~( + _PREPARED_ROUTE_IS_FULL_FLAG | _PREPARED_ROUTE_IS_PROXY_FLAG +) @cute.jit @@ -108,7 +115,7 @@ def _swaps_forwards_packed_route_full(cfg: FmhaDecodeConfig) -> bool: return ( cfg.tile_size_q == 8 and cfg.kv_block_size == 8 - and not cfg.use_kv_valid_bits + and not cfg.uses_prepared_score_keep_words and not cfg.uses_uniform_causal_mask and not cfg.uses_per_row_causal_mask ) @@ -116,12 +123,16 @@ def _swaps_forwards_packed_route_full(cfg: FmhaDecodeConfig) -> bool: def _kv_retained_route_words( route_layout: _BlockSparseRouteLayout, + *, + retain_proxy_kind: bool = False, ) -> int: """Return the aligned SMEM words retained from K issue through V. - Contiguous routes retain their existing load-origin payload. Paged routes - retain parallel logical-origin and physical-page-ID arrays so every atom - has an independent storage locator; invalid entries use ``(-1, -1)``. + Contiguous routes retain their load-origin payload. Paged routes retain + parallel logical-origin and physical-page-ID arrays so every atom has an + independent storage locator. A two-origin contiguous route additionally + keeps its atom-valid mask. Proxy-capable exact/proxy routes reserve the + final aligned word for an explicit source kind. """ payload_words = route_layout.logical_origins_per_route @@ -129,6 +140,8 @@ def _kv_retained_route_words( payload_words *= 2 elif route_layout.logical_origins_per_route == 2: payload_words += 1 + if retain_proxy_kind: + payload_words += 1 return ((payload_words + 3) // 4) * 4 @@ -139,9 +152,9 @@ class _BlockSparseSoftmaxStagingLayout: 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. Its - noncausal Q8/B8 profile without token bits packs route-full into the - otherwise-zero low bit of each warp's first aligned origin. + 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. @@ -206,14 +219,23 @@ class SmemBlockSparseKvMetadataResource(DecodeGenResourceBase): """Pipeline-free route metadata retained from one K issue through V. ``route_metadata`` points at the first prepared GMEM record. Resolution - returns logical origins to masking consumers. The private SMEM copy keeps - contiguous load origins or paged ``(logical origin, physical page ID)`` - pairs through the matching V issue. Invalid atoms are retained as safe - storage-specific OOB coordinates. + returns logical/source-domain origins to masking consumers. The private + SMEM copy keeps contiguous load origins or paged ``(logical origin, + physical page ID)`` pairs through the matching V issue. Invalid atoms are + retained as safe storage-specific OOB coordinates. The proxy-capable + contiguous specialization interprets origins in the selected source domain + (summary tokens for proxy routes, K/V tokens for exact routes) and + retains the prepared route kind in a separate aligned word. Exact-only + specializations keep their original allocation. """ _task_local_specs: ClassVar[tuple[tuple, ...]] = ( - ("resolved_origin0_slot", Int32, Int32(0), "First logical origin."), + ( + "resolved_record_word_slot", + Int32, + Int32(0), + "Lane-owned record word; locator lanes carry route origins.", + ), ("resolved_origin1_slot", Int32, Int32(0), "Second logical origin."), ( "resolved_atom_validity_slot", @@ -227,6 +249,18 @@ class SmemBlockSparseKvMetadataResource(DecodeGenResourceBase): Int32(-1), "Metadata-relative record offset, or -1 for a dummy route.", ), + ( + "prefetched_record_word_slot", + Int32, + Int32(0), + "Lane-owned record word loaded one resolution ahead.", + ), + ( + "prefetched_record_offset_slot", + Int32, + Int32(-1), + "Record offset of the prefetched route, or -1 for a dummy route.", + ), ) cfg: Constexpr[FmhaDecodeConfig] = None inst_id: Constexpr[int] = 0 @@ -236,7 +270,7 @@ class SmemBlockSparseKvMetadataResource(DecodeGenResourceBase): _retained_route_words: Constexpr[int] = 0 _alloc: Constexpr[SmemAllocation | None] = None _smem_words: cutlass.Array = None - resolved_origin0_slot: Constexpr[TaskLocalVariable] = ( + resolved_record_word_slot: Constexpr[TaskLocalVariable] = ( TaskLocalVariable.uninitialized() ) resolved_origin1_slot: Constexpr[TaskLocalVariable] = ( @@ -248,13 +282,25 @@ class SmemBlockSparseKvMetadataResource(DecodeGenResourceBase): route_record_word_offset_slot: Constexpr[TaskLocalVariable] = ( TaskLocalVariable.uninitialized() ) + prefetched_record_word_slot: Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + prefetched_record_offset_slot: Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) def __post_init__(self) -> None: """Derive the retained K/V payload from the prepared route layout.""" assert self.route_layout is not None assert self.route_layout.is_paged == self.cfg.use_paged_kv - self._retained_route_words = _kv_retained_route_words(self.route_layout) + if self.cfg.use_block_sparse_proxy_routes: + assert not self.route_layout.is_paged + assert self.route_layout.uses_one_warp_transport + self._retained_route_words = _kv_retained_route_words( + self.route_layout, + retain_proxy_kind=self.cfg.use_block_sparse_proxy_routes, + ) super().__post_init__() def _init_placeholder_state(self) -> None: @@ -332,9 +378,81 @@ def _prepared_route_physical_page_id_if_valid( ) return physical_page_id + @cute.jit + def _route_record_word_offset( + self, stage_info: StageInfo, route_idx: Int32 + ) -> Int32: + """Return the record offset of one route index, or -1 past the row.""" + + task_cache = _decode_gen_task_cache(stage_info) + row_route_begin = _sparse_task_cache_route_begin(task_cache) + route_count = _sparse_task_cache_route_count(task_cache) + route_record_word_offset = Int32(-1) + if route_idx < route_count: + route_record_word_offset = (row_route_begin + route_idx) * Int32( + self.route_layout.route_metadata_stride_words + ) + return cute.arch.make_warp_uniform(route_record_word_offset) + + @consumer_work( + returns=( + prefetched_record_word_slot, + prefetched_record_offset_slot, + ) + ) + @cute.jit + def prefetch_route( + self, stage_info: StageInfo, *, target: Constexpr[str] + ) -> tuple[Int32, Int32]: + """Issue the record load for a route that ``resolve_route`` uses later. + + ``target`` selects the route relative to the calling section: + ``"head"`` is this instance's HEAD route, ``"first_loop"`` the route of + LOOP iteration 0 (called from HEAD), ``"current_loop"`` the route of + the calling LOOP iteration (no pipelining), and ``"next_loop"`` the + route of the following LOOP iteration. Only the lane-distributed load is issued + here; the warp broadcasts happen in ``resolve_route`` so the global + memory latency overlaps the TMA issue of the current route instead of + stalling the load warp. Layouts without one-warp transport keep their + loads in ``resolve_route`` and get placeholder values here. + """ + + assert self.route_metadata is not None + num_insts = Int32(self.cfg.num_insts_kv) + if cutlass.const_expr(target == "head"): + route_idx = Int32(self.inst_id) + elif cutlass.const_expr(target == "first_loop"): + route_idx = num_insts + Int32(self.inst_id) + elif cutlass.const_expr(target == "current_loop"): + route_idx = (stage_info.loop_offset + Int32(1)) * num_insts + Int32( + self.inst_id + ) + else: + route_idx = (stage_info.loop_offset + Int32(2)) * num_insts + Int32( + self.inst_id + ) + route_record_word_offset = self._route_record_word_offset(stage_info, route_idx) + record_word = Int32(0) + if cutlass.const_expr(self.route_layout.uses_one_warp_transport): + assert self.route_layout.token_words_word_offset is not None + meaningful_words = ( + self.route_layout.token_words_word_offset + + self.route_layout.token_words_per_route + ) + lane_idx = cute.arch.thread_idx()[0] & Int32(0x1F) + if lane_idx < Int32(self.route_layout.logical_origins_per_route): + record_word = Int32(-1) + if route_record_word_offset >= Int32(0) and lane_idx < Int32( + meaningful_words + ): + record_word = Int32( + self.route_metadata[route_record_word_offset + lane_idx] + ) + return record_word, route_record_word_offset + @consumer_work( returns=( - resolved_origin0_slot, + resolved_record_word_slot, resolved_origin1_slot, resolved_atom_validity_slot, route_record_word_offset_slot, @@ -342,14 +460,23 @@ def _prepared_route_physical_page_id_if_valid( ) @cute.jit def resolve_route( - self, stage_info: StageInfo, *, section: Constexpr[FmhaStage] + self, + stage_info: StageInfo, + *, + section: Constexpr[FmhaStage], + prefetched_record_word_slot: Int32, + prefetched_record_offset_slot: Int32, ) -> tuple[Int32, Int32, Int32, Int32]: - """Load this resource instance's real or dummy prepared KV route.""" + """Resolve this instance's real or dummy prepared KV route. + + One-warp-transport layouts consume the words that ``prefetch_route`` + loaded earlier; other layouts load their record here. The routed + inputs carry the task-local slot names so that every ``prefetch_route`` + call, including the one at the end of the previous LOOP iteration, + updates the value read here. + """ assert self.route_metadata is not None - task_cache = _decode_gen_task_cache(stage_info) - row_route_begin = _sparse_task_cache_route_begin(task_cache) - route_count = _sparse_task_cache_route_count(task_cache) # HEAD publishes one route per instruction. LOOP starts after those # two publications, hence the one-based loop offset below. Keeping the # constexpr branch local lets the task scheduler specialize each work @@ -361,12 +488,12 @@ def resolve_route( self.cfg.num_insts_kv ) + Int32(self.inst_id) lane_idx = cute.arch.thread_idx()[0] & Int32(0x1F) - route_record_word_offset = Int32(-1) - if route_idx < route_count: - route_record_word_offset = (row_route_begin + route_idx) * Int32( - self.route_layout.route_metadata_stride_words + if cutlass.const_expr(self.route_layout.uses_one_warp_transport): + route_record_word_offset = prefetched_record_offset_slot + else: + route_record_word_offset = self._route_record_word_offset( + stage_info, route_idx ) - route_record_word_offset = cute.arch.make_warp_uniform(route_record_word_offset) num_logical_origins = self.route_layout.logical_origins_per_route uses_two_fragment_route = num_logical_origins == 2 @@ -378,13 +505,37 @@ def resolve_route( atom_valid_mask = Int32(0) route_record_is_valid = route_record_word_offset >= Int32(0) + if cutlass.const_expr(self.route_layout.uses_one_warp_transport): + resolved_record_word = prefetched_record_word_slot + atom_valid_mask = _warp_broadcast_i32( + resolved_record_word, + self.route_layout.atom_valid_mask_word_offset, + ) + if cutlass.const_expr(uses_two_fragment_route): + return ( + resolved_record_word, + _warp_broadcast_i32(resolved_record_word, 1), + atom_valid_mask, + route_record_word_offset, + ) + atom_is_valid = cutlass.Boolean( + lane_idx < Int32(num_logical_origins) + and (atom_valid_mask & (Int32(1) << lane_idx)) != Int32(0) + ) + return ( + resolved_record_word, + Int32(0), + Int32(atom_is_valid), + route_record_word_offset, + ) + if route_record_is_valid: if lane_idx < Int32(num_logical_origins): logical_origin = self._prepared_route_logical_origin( route_record_word_offset, lane_idx, ) - if cutlass.const_expr(self.cfg.use_kv_valid_bits): + if cutlass.const_expr(self.cfg.uses_prepared_score_keep_words): if lane_idx == Int32(valid_mask_lane): atom_valid_mask = Int32( self.route_metadata[ @@ -396,7 +547,7 @@ def resolve_route( if cutlass.const_expr(uses_two_fragment_route): origin0 = _warp_broadcast_i32(logical_origin, 0) origin1 = _warp_broadcast_i32(logical_origin, 1) - if cutlass.const_expr(self.cfg.use_kv_valid_bits): + if cutlass.const_expr(self.cfg.uses_prepared_score_keep_words): # The validity word shares the prepared record's cache line # with fields consumed shortly afterward by Softmax. atom_valid_mask = _warp_broadcast_i32(atom_valid_mask, valid_mask_lane) @@ -412,7 +563,7 @@ def resolve_route( # Wider routes stay lane-distributed: each active lane carries only # its origin and validity through the existing three-scalar K/V ABI. valid = cutlass.Boolean(False) - if cutlass.const_expr(self.cfg.use_kv_valid_bits): + if cutlass.const_expr(self.cfg.uses_prepared_score_keep_words): atom_valid_mask = _warp_broadcast_i32(atom_valid_mask, valid_mask_lane) if lane_idx < Int32(num_logical_origins): valid = (atom_valid_mask & (Int32(1) << lane_idx)) != Int32(0) @@ -427,7 +578,7 @@ def store_route( self, stage_info: StageInfo, *, - resolved_origin0: Int32, + resolved_record_word: Int32, resolved_origin1: Int32, resolved_atom_validity: Int32, route_record_word_offset: Int32, @@ -439,11 +590,11 @@ def store_route( num_origins = self.route_layout.logical_origins_per_route if cutlass.const_expr(self.route_layout.is_paged): if lane_idx < Int32(num_origins): - logical_origin = Int32(resolved_origin0) + logical_origin = Int32(resolved_record_word) atom_is_valid = resolved_atom_validity != Int32(0) if cutlass.const_expr(num_origins == 2): if lane_idx == Int32(0): - logical_origin = resolved_origin0 + logical_origin = resolved_record_word else: logical_origin = resolved_origin1 atom_is_valid = ( @@ -463,14 +614,14 @@ def store_route( ] = physical_page_id elif cutlass.const_expr(num_origins == 2): if lane_idx == Int32(0): - self._smem_words[Int32(0)] = resolved_origin0 + self._smem_words[Int32(0)] = resolved_record_word self._smem_words[Int32(1)] = resolved_origin1 self._smem_words[ Int32(self.route_layout.atom_valid_mask_word_offset) ] = resolved_atom_validity else: - if lane_idx < Int32(self.route_layout.logical_origins_per_route): - load_origin = Int32(resolved_origin0) + if lane_idx < Int32(num_origins): + load_origin = Int32(resolved_record_word) if resolved_atom_validity == Int32(0): # Fine-route K and V both consume this retained value. # Materialize their TensorMap OOB coordinate once here @@ -478,6 +629,12 @@ def store_route( # every atom copy in both producer passes. load_origin = Int32(self.tma_oob_origin) self._smem_words[lane_idx] = load_origin + if cutlass.const_expr(self.cfg.use_block_sparse_proxy_routes): + if lane_idx == Int32(self.route_layout.route_flags_word_offset): + prepared_route_flags = Int32(resolved_record_word) + self._smem_words[Int32(self._retained_route_words - 1)] = Int32( + prepared_route_flags + ) & Int32(_PREPARED_ROUTE_IS_PROXY_FLAG) # K consumes this slot immediately, while V consumes it at the start # of the next cadence. Both execute in this warp, so a warp fence is # sufficient; no cross-warp mbarrier belongs here. @@ -523,6 +680,20 @@ def route_atom_valid_mask(self) -> Int32: self._smem_words[Int32(self.route_layout.atom_valid_mask_word_offset)] ) + @cute.jit + def route_is_proxy(self) -> cutlass.Boolean: + """Return the retained prepared route kind for the current K/V pair.""" + + if cutlass.const_expr(not self.cfg.use_block_sparse_proxy_routes): + return cutlass.Boolean(False) + return cutlass.Boolean( + ( + Int32(self._smem_words[Int32(self._retained_route_words - 1)]) + & Int32(_PREPARED_ROUTE_IS_PROXY_FLAG) + ) + != Int32(0) + ) + @dataclass(kw_only=True) class SmemBlockSparseSoftmaxMetadataResource(DecodeGenResourceBase): @@ -533,8 +704,8 @@ class SmemBlockSparseSoftmaxMetadataResource(DecodeGenResourceBase): producer passes the resolved payload explicitly instead of recomputing it. For Keeps, every route token word moves through SMEM without a data-dependent branch; each consumer receives at most four words through - the stable task-local ABI. A runtime route-full bit can skip per-score token - predicates while leaving structural masking independent. + the stable task-local ABI. Runtime route flags carry the conservative FULL + summary and, for proxy-capable builds, the route source kind. """ _task_local_specs: ClassVar[tuple[tuple, ...]] = ( @@ -562,7 +733,7 @@ class SmemBlockSparseSoftmaxMetadataResource(DecodeGenResourceBase): "softmax_token_word2_slot", Uint32, Uint32(0xFFFFFFFF), - "Loaded third Keeps token word or SWAP's B8 route-full summary.", + "Loaded third Keeps token word or SWAP's packed route flags.", ), ( "softmax_token_word3_slot", @@ -605,6 +776,8 @@ def __post_init__(self) -> None: assert self.route_layout is not None assert self.route_layout.is_paged == self.cfg.use_paged_kv + if self.cfg.use_block_sparse_proxy_routes: + assert self.route_layout.uses_one_warp_transport self.staging_layout = _BlockSparseSoftmaxStagingLayout.create( use_keeps_mma_ab=self.cfg.use_keeps_mma_ab, route_layout=self.route_layout, @@ -680,57 +853,82 @@ def _consumer_stage_base(self) -> Int32: def _store_route_swaps( self, stage_info: StageInfo, - resolved_origin0: Int32, + resolved_record_word: Int32, resolved_origin1: Int32, resolved_atom_validity: Int32, route_record_word_offset: Int32, ) -> None: """Stage SWAP origins and optional logical-K32 token metadata. - The noncausal Q8/B8 profile without token bits also packs prepare's - route-full summary into bit 0 of each warp's first aligned origin. + Selected prepared route flags use the free low bits of each warp's + first aligned origin. """ lane_idx = cute.arch.thread_idx()[0] & Int32(0x1F) stage_base = self._producer_stage_base(stage_info) task_cache = _decode_gen_task_cache(stage_info) seq_len_kv = Int32(task_cache[_TASK_CACHE_SEQ_LEN_KV]) - route_record_is_valid = route_record_word_offset >= Int32(0) - - packed_route_full = Int32(0) - if cutlass.const_expr(_swaps_forwards_packed_route_full(self.cfg)): - if lane_idx == Int32(0) and route_record_is_valid: - packed_route_full = Int32( - self.route_metadata[ - route_record_word_offset - + Int32(self.route_layout.route_flags_word_offset) - ] - ) & Int32(_PREPARED_ROUTE_IS_FULL_FLAG) - packed_route_full = _warp_broadcast_i32(packed_route_full, 0) + uses_one_warp_transport = self.route_layout.uses_one_warp_transport + if cutlass.const_expr( + not uses_one_warp_transport + and ( + _swaps_forwards_packed_route_full(self.cfg) + or self.cfg.uses_prepared_score_keep_words + ) + ): + route_record_is_valid = route_record_word_offset >= Int32(0) + + packed_route_flags = Int32(0) + if cutlass.const_expr( + _swaps_forwards_packed_route_full(self.cfg) + or self.cfg.use_block_sparse_proxy_routes + ): + if cutlass.const_expr(uses_one_warp_transport): + packed_route_flags = _warp_broadcast_i32( + resolved_record_word, + self.route_layout.route_flags_word_offset, + ) + else: + if lane_idx == Int32(0) and route_record_is_valid: + packed_route_flags = Int32( + self.route_metadata[ + route_record_word_offset + + Int32(self.route_layout.route_flags_word_offset) + ] + ) & Int32(_PREPARED_ROUTE_IS_FULL_FLAG) + packed_route_flags = _warp_broadcast_i32(packed_route_flags, 0) softmax_origin = Int32(-1) if cutlass.const_expr(self.cfg.kv_block_size < 64): if lane_idx < Int32(self.staging_layout.num_origin_words): - softmax_origin = Int32(resolved_origin0) + softmax_origin = Int32(resolved_record_word) if resolved_atom_validity == Int32(0): softmax_origin = Int32(-1) - if cutlass.const_expr(_swaps_forwards_packed_route_full(self.cfg)): - # Replicate route-full in each K32 slice's first origin; - # B8 alignment leaves bit 0 free for the summary. + if cutlass.const_expr( + _swaps_forwards_packed_route_full(self.cfg) + or self.cfg.use_block_sparse_proxy_routes + ): + # Every SWAP atom is at least B8 aligned. Replicate the + # route flags in each K32 slice's first origin so the + # established seven-slot Softmax ABI also carries source + # kind without growing the staged payload. if lane_idx % Int32(self.staging_layout.origins_per_warp) == Int32( 0 ): softmax_origin = ( - softmax_origin & Int32(_SWAPS_PACKED_ROUTE_FULL_CLEAR_MASK) - ) | packed_route_full + softmax_origin & Int32(_SWAPS_PACKED_ROUTE_FLAGS_CLEAR_MASK) + ) | packed_route_flags self._smem_words[stage_base + lane_idx] = softmax_origin else: # SWAP with a coarse KV atom expands the two resolved KV64 # fragments into the four logical K32 origins consumed by its # four softmax warps. + coarse_origin0 = Int32(resolved_record_word) + if cutlass.const_expr(uses_one_warp_transport): + coarse_origin0 = _warp_broadcast_i32(resolved_record_word, 0) if lane_idx < Int32(4): fragment_idx = lane_idx >> Int32(1) - softmax_origin = Int32(resolved_origin0) + softmax_origin = coarse_origin0 valid = (resolved_atom_validity & Int32(1)) != Int32(0) if fragment_idx == Int32(1): softmax_origin = Int32(resolved_origin1) @@ -738,13 +936,30 @@ def _store_route_swaps( softmax_origin = softmax_origin + (lane_idx & Int32(1)) * Int32(32) if not valid or softmax_origin >= seq_len_kv: softmax_origin = Int32(-1) + if cutlass.const_expr(self.cfg.use_block_sparse_proxy_routes): + # Coarse SWAP expands KV64 atoms to K32-aligned origins; + # their low bits carry the same typed-route flags as the + # fine-route representation above. + softmax_origin = ( + softmax_origin & Int32(_SWAPS_PACKED_ROUTE_FLAGS_CLEAR_MASK) + ) | packed_route_flags self._smem_words[stage_base + lane_idx] = softmax_origin - if cutlass.const_expr(self.cfg.use_kv_valid_bits): + if cutlass.const_expr(self.cfg.uses_prepared_score_keep_words): assert self.route_metadata is not None assert self.route_layout.token_words_word_offset is not None assert self.staging_layout.token_words_word_offset is not None - if lane_idx < Int32(self.route_layout.token_words_per_route): + if cutlass.const_expr(uses_one_warp_transport): + token_begin = Int32(self.route_layout.token_words_word_offset) + token_end = token_begin + Int32(self.route_layout.token_words_per_route) + if lane_idx >= token_begin and lane_idx < token_end: + self._smem_words[ + stage_base + + Int32(self.staging_layout.token_words_word_offset) + + lane_idx + - token_begin + ] = Int32(resolved_record_word) + elif lane_idx < Int32(self.route_layout.token_words_per_route): logical_word = Uint32(0) if route_record_is_valid: logical_word = Uint32( @@ -765,7 +980,7 @@ def _store_route_swaps( def _store_route_keeps( self, stage_info: StageInfo, - resolved_origin0: Int32, + resolved_record_word: Int32, resolved_origin1: Int32, resolved_atom_validity: Int32, route_record_word_offset: Int32, @@ -779,8 +994,12 @@ def _store_route_keeps( assert self.staging_layout.route_flags_word_offset is not None lane_idx = cute.arch.thread_idx()[0] & Int32(0x1F) - route_record_is_valid = route_record_word_offset >= Int32(0) num_origins = self.route_layout.logical_origins_per_route + uses_one_warp_transport = self.route_layout.uses_one_warp_transport + if cutlass.const_expr( + not uses_one_warp_transport and self.cfg.uses_prepared_score_keep_words + ): + route_record_is_valid = route_record_word_offset >= Int32(0) route_flags = Int32(resolved_atom_validity) if cutlass.const_expr(num_origins > 2): route_flags = Int32( @@ -788,47 +1007,73 @@ def _store_route_keeps( lane_idx < Int32(num_origins) and resolved_atom_validity != Int32(0) ) ) + token_word = Uint32(0) route_token_mask_is_full = cutlass.Boolean(False) - if cutlass.const_expr(self.cfg.use_kv_valid_bits): - assert self.route_metadata is not None + if cutlass.const_expr(self.cfg.uses_prepared_score_keep_words): assert self.route_layout.token_words_word_offset is not None - gmem_route_flags = Int32(0) - if lane_idx == Int32(0) and route_record_is_valid: - gmem_route_flags = Int32( - self.route_metadata[ - route_record_word_offset - + Int32(self.route_layout.route_flags_word_offset) - ] - ) - gmem_route_flags = _warp_broadcast_i32(gmem_route_flags, 0) - # Prepared bit 0 summarizes the whole route. Staged low bits are - # already fragment validity, so remap the summary above them. - route_token_mask_is_full = cutlass.Boolean( - (gmem_route_flags & Int32(_PREPARED_ROUTE_IS_FULL_FLAG)) != Int32(0) - ) - if ( - lane_idx < Int32(self.route_layout.token_words_per_route) - and route_record_is_valid - ): - token_word = Uint32( - self.route_metadata[ - route_record_word_offset - + Int32(self.route_layout.token_words_word_offset) - + lane_idx - ] + assert self.route_metadata is not None + if cutlass.const_expr(not uses_one_warp_transport): + gmem_route_flags = Int32(0) + if lane_idx == Int32(0) and route_record_is_valid: + gmem_route_flags = Int32( + self.route_metadata[ + route_record_word_offset + + Int32(self.route_layout.route_flags_word_offset) + ] + ) + gmem_route_flags = _warp_broadcast_i32(gmem_route_flags, 0) + # Prepared bit 0 summarizes the whole route. Staged low bits + # already hold fragment validity, so remap it above them. + route_token_mask_is_full = cutlass.Boolean( + (gmem_route_flags & Int32(_PREPARED_ROUTE_IS_FULL_FLAG)) != Int32(0) ) + if ( + lane_idx < Int32(self.route_layout.token_words_per_route) + and route_record_is_valid + ): + token_word = Uint32( + self.route_metadata[ + route_record_word_offset + + Int32(self.route_layout.token_words_word_offset) + + lane_idx + ] + ) stage_base = self._producer_stage_base(stage_info) if cutlass.const_expr(num_origins == 2): if lane_idx == Int32(0): - self._smem_words[stage_base] = Int32(resolved_origin0) + self._smem_words[stage_base] = Int32(resolved_record_word) self._smem_words[stage_base + Int32(1)] = Int32(resolved_origin1) - else: - if lane_idx < Int32(num_origins): - self._smem_words[stage_base + lane_idx] = Int32(resolved_origin0) - if lane_idx == Int32(0): - if cutlass.const_expr(self.cfg.use_kv_valid_bits): + elif lane_idx < Int32(num_origins): + self._smem_words[stage_base + lane_idx] = Int32(resolved_record_word) + + if cutlass.const_expr(uses_one_warp_transport): + if lane_idx == Int32(self.route_layout.route_flags_word_offset): + prepared_route_flags = Int32(resolved_record_word) + route_flags = route_flags | ( + Int32( + (prepared_route_flags & Int32(_PREPARED_ROUTE_IS_FULL_FLAG)) + != Int32(0) + ) + * Int32(_SOFTMAX_TOKEN_MASK_IS_FULL_FLAG) + ) + if cutlass.const_expr(self.cfg.use_block_sparse_proxy_routes): + route_flags = route_flags | ( + Int32( + ( + prepared_route_flags + & Int32(_PREPARED_ROUTE_IS_PROXY_FLAG) + ) + != Int32(0) + ) + * Int32(_SOFTMAX_ROUTE_IS_PROXY_FLAG) + ) + self._smem_words[ + stage_base + Int32(self.staging_layout.route_flags_word_offset) + ] = route_flags + elif lane_idx == Int32(0): + if cutlass.const_expr(self.cfg.uses_prepared_score_keep_words): route_flags = route_flags | ( Int32(route_token_mask_is_full) * Int32(_SOFTMAX_TOKEN_MASK_IS_FULL_FLAG) @@ -836,9 +1081,20 @@ def _store_route_keeps( self._smem_words[ stage_base + Int32(self.staging_layout.route_flags_word_offset) ] = route_flags - if cutlass.const_expr(self.cfg.use_kv_valid_bits): + + if cutlass.const_expr(self.cfg.uses_prepared_score_keep_words): assert self.staging_layout.token_words_word_offset is not None - if lane_idx < Int32(self.route_layout.token_words_per_route): + if cutlass.const_expr(uses_one_warp_transport): + token_begin = Int32(self.route_layout.token_words_word_offset) + token_end = token_begin + Int32(self.route_layout.token_words_per_route) + if lane_idx >= token_begin and lane_idx < token_end: + self._smem_words[ + stage_base + + Int32(self.staging_layout.token_words_word_offset) + + lane_idx + - token_begin + ] = Int32(resolved_record_word) + elif lane_idx < Int32(self.route_layout.token_words_per_route): self._smem_words[ stage_base + Int32(self.staging_layout.token_words_word_offset) @@ -852,7 +1108,7 @@ def store_route( self, stage_info: StageInfo, *, - resolved_origin0: Int32, + resolved_record_word: Int32, resolved_origin1: Int32, resolved_atom_validity: Int32, route_record_word_offset: Int32, @@ -862,7 +1118,7 @@ def store_route( if cutlass.const_expr(self.cfg.use_keeps_mma_ab): self._store_route_keeps( stage_info, - resolved_origin0, + resolved_record_word, resolved_origin1, resolved_atom_validity, route_record_word_offset, @@ -870,7 +1126,7 @@ def store_route( else: self._store_route_swaps( stage_info, - resolved_origin0, + resolved_record_word, resolved_origin1, resolved_atom_validity, route_record_word_offset, @@ -886,8 +1142,8 @@ def _load_route_swaps_values( this Softmax warp's logical K32 slice; unused or invalid origins are negative. To preserve the shared seven-slot task ABI, origin 2/3 subsequently travel through the shared route-flags/token-word-0 slots. - Token-word 1 carries the logical K32 mask, token-word 2 carries - route-full, and token-word 3 is unused. + Token-word 1 carries the logical K32 mask, token-word 2 carries packed + route flags, and token-word 3 is unused. """ stage_base = self._consumer_stage_base() @@ -907,7 +1163,7 @@ def _load_route_swaps_values( origin3 = Int32(self._smem_words[warp_origin_base + Int32(3)]) token_word = Uint32(0xFFFFFFFF) - if cutlass.const_expr(self.cfg.use_kv_valid_bits): + if cutlass.const_expr(self.cfg.uses_prepared_score_keep_words): assert self.staging_layout.token_words_word_offset is not None token_word = Uint32( self._smem_words[ @@ -917,9 +1173,25 @@ def _load_route_swaps_values( ] ) route_flags = Uint32(0) - if cutlass.const_expr(_swaps_forwards_packed_route_full(self.cfg)): - route_flags = Uint32(origin0 & Int32(1)) - origin0 = origin0 & Int32(_SWAPS_PACKED_ROUTE_FULL_CLEAR_MASK) + if cutlass.const_expr( + _swaps_forwards_packed_route_full(self.cfg) + or self.cfg.use_block_sparse_proxy_routes + ): + packed_route_flags = origin0 & Int32( + _PREPARED_ROUTE_IS_FULL_FLAG | _PREPARED_ROUTE_IS_PROXY_FLAG + ) + route_flags = Uint32( + packed_route_flags & Int32(_PREPARED_ROUTE_IS_FULL_FLAG) + ) + if cutlass.const_expr(self.cfg.use_block_sparse_proxy_routes): + route_flags = route_flags | Uint32( + Int32( + (packed_route_flags & Int32(_PREPARED_ROUTE_IS_PROXY_FLAG)) + != Int32(0) + ) + * Int32(_SOFTMAX_ROUTE_IS_PROXY_FLAG) + ) + origin0 = origin0 & Int32(_SWAPS_PACKED_ROUTE_FLAGS_CLEAR_MASK) return ( origin0, origin1, @@ -986,17 +1258,21 @@ def load_route( valid0 = (stored_route_flags >> origin0_idx) & Int32(1) valid1 = (stored_route_flags >> origin1_idx) & Int32(1) route_flags = valid0 | (valid1 << Int32(1)) - if cutlass.const_expr(self.cfg.use_kv_valid_bits): + if cutlass.const_expr(self.cfg.uses_prepared_score_keep_words): route_flags = route_flags | ( stored_route_flags & Int32(_SOFTMAX_TOKEN_MASK_IS_FULL_FLAG) ) + if cutlass.const_expr(self.cfg.use_block_sparse_proxy_routes): + route_flags = route_flags | ( + stored_route_flags & Int32(_SOFTMAX_ROUTE_IS_PROXY_FLAG) + ) origin0 = Int32(self._smem_words[stage_base + origin0_idx]) origin1 = Int32(self._smem_words[stage_base + origin1_idx]) token_word0 = Uint32(0xFFFFFFFF) token_word1 = Uint32(0xFFFFFFFF) token_word2 = Uint32(0xFFFFFFFF) token_word3 = Uint32(0xFFFFFFFF) - if cutlass.const_expr(self.cfg.use_kv_valid_bits): + if cutlass.const_expr(self.cfg.uses_prepared_score_keep_words): assert self.staging_layout.token_words_word_offset is not None if cutlass.const_expr(self.route_layout.kv_route_size == 256): token_base = Int32(self.staging_layout.token_words_word_offset) @@ -1014,28 +1290,6 @@ def load_route( token_word3 = Uint32( self._smem_words[stage_base + token_base + word1_idx + Int32(1)] ) - elif cutlass.const_expr(self.cfg.tile_size_q == 64): - lane_idx = cute.arch.thread_idx()[0] & Int32(0x1F) - local_word_base = _keeps_col_base( - self.cfg, - lane_idx, - self.cfg.num_s_regs_per_thread, - ) >> Int32(5) - token_word0 = Uint32( - self._smem_words[ - stage_base - + Int32(self.staging_layout.token_words_word_offset) - + local_word_base - ] - ) - token_word1 = Uint32( - self._smem_words[ - stage_base - + Int32(self.staging_layout.token_words_word_offset) - + local_word_base - + Int32(1) - ] - ) else: token_word0 = Uint32( self._smem_words[ 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 bb27c3c553f0..330f15b05691 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 @@ -25,7 +25,7 @@ import cutlass import cutlass.cute as cute -from cutlass import Float32, Int32, Int64 +from cutlass import Float32, Int32, Int64, Uint32 from cutlass.experimental import primitives as prims from cutlass.experimental.task_scheduling.memory import ( @@ -42,6 +42,7 @@ ) from ..fmha_decode_config import FmhaDecodeConfig +from ...._block_sparse.common import _block_sparse_proxy_summary_geometry from ...placeholder_helpers import _placeholder_smem_array from .helpers_common import ( Constexpr, @@ -58,11 +59,13 @@ _is_last_loop_iteration, _keeps_col_base, _keeps_row_idx, + _keeps_tcgen05_ld, _keeps_tcgen05_st, _named_barrier_arrive, _neg_max_f32, _pack_float2_to_bf16, _pack_float2_to_fp16, + _swaps_routed_coordinate, _wait_for_mbarrier_phase, ) from .helpers_output import ( @@ -76,11 +79,27 @@ _compute_fp8_p_regs_and_local_sums, _compute_fp8_p_regs_and_local_sums_dense, _compute_p_values_and_local_sums_dense, + _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 .tmem_s import TmemSResource +# Tunable: number of score pairs per streamed fragment whose exponentials run +# as FMA polynomials instead of MUFU. The MUFU issue rate bounds the fragment +# otherwise, while the FMA pipe is nearly idle in the softmax warps. Larger +# shares grow the fragment body and the softmax warps become instruction-fetch +# bound again, so one quarter of the 16 pairs is the measured optimum. +KV_TILE_256_EX2_EMULATED_PAIRS = 4 + + +def _pair_uses_ex2_emulation(pair_idx: int, pairs_per_fragment: int) -> bool: + """Spread the emulated pairs evenly across a fragment's score pairs.""" + count = KV_TILE_256_EX2_EMULATED_PAIRS + pairs = pairs_per_fragment + return ((pair_idx + 1) * count) // pairs != (pair_idx * count) // pairs + @dataclass(kw_only=True) class SmemPResource(DecodeGenResourceBase): @@ -88,7 +107,7 @@ class SmemPResource(DecodeGenResourceBase): Softmax producers convert S to P, store it in the profile's TMEM or SMEM layout, and publish local sums back to TmemS. Most profiles use the generic - full/empty P pipeline. KV256 instead publishes four independently ready + full/empty P pipeline. Streamed profiles instead publish four independently ready K32 TMEM fragments; BMM2 consumes those fragments in order, while the matching TmemO full barrier prevents the next QK from overwriting aliased P. """ @@ -152,7 +171,7 @@ def _init_placeholder_state(self) -> None: ) def get_smem_requirements(self) -> list[SmemAllocation]: - """Allocate P storage or the KV256 fragment-ready barriers.""" + """Allocate P storage or the streamed fragment-ready barriers.""" if self.cfg.streams_tmem_p_fragments: if self._fragment_ready_alloc is None: self._fragment_ready_alloc = SmemAllocation( @@ -173,7 +192,7 @@ def get_smem_requirements(self) -> list[SmemAllocation]: @cute.jit def _bind_fragment_ready(self, context: ResourceContext | None = None) -> None: - """Bind the one-way KV256 P-ready barriers from the SMEM context.""" + """Bind the one-way streamed P-ready barriers from the SMEM context.""" if cutlass.const_expr( self.cfg.streams_tmem_p_fragments and context is not None @@ -191,7 +210,7 @@ def _bind_fragment_ready(self, context: ResourceContext | None = None) -> None: def create_function_variables( self, context: ResourceContext | None = None ) -> ResourceVars: - """Bind and initialize KV256's per-fragment ready barriers.""" + """Bind and initialize the streamed per-fragment ready barriers.""" self._bind_fragment_ready(context) if cutlass.const_expr(self.cfg.streams_tmem_p_fragments): tidx, _, _ = cute.arch.thread_idx() @@ -299,43 +318,203 @@ def init_descriptor_state(self, stage_info: StageInfo) -> None: # work can publish a valid descriptor or TMEM address for this tile. self._create_initial_task_locals(stage_info.context) + @cute.jit + def _apply_proxy_route_denominator_mass( + self, + local_sum: Float32, + tail_p: Float32, + route_is_proxy: Int32, + ) -> 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 + @producer_work @cute.jit - def compute_p_fragment( + def compute_p_fragments( self, stage_info: StageInfo, *, - fragment_idx: Constexpr[int], new_max_arr: cutlass.Array, - s_arr: cutlass.Array, ) -> None: - """Convert one KV256 K32 score fragment and publish its TMEM P slice.""" + """Stream every ordinary K32 fragment from one rolled loop.""" + 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), + ) + + @producer_work + @cute.jit + def compute_proxy_route_p_fragments( + self, + stage_info: StageInfo, + *, + 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.""" + 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_origin0=route_origin0, + route_origin1=route_origin1, + ) + + @cute.jit + def _compute_p_fragments_impl( + self, + stage_info: StageInfo, + *, + new_max_arr: cutlass.Array, + route_is_proxy: Int32, + route_origin0: Int32, + route_origin1: Int32, + ) -> 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. + """ + _ = stage_info cfg = self.cfg assert cfg.streams_tmem_p_fragments - assert not cfg.use_fp8_qkv and cfg.uses_two_inst_tmem_p - assert cfg.softmax_score_fragment_regs == 32 + assert self._tmem_alloc.offset == self.tmem_s_ref._alloc.offset + # One FP32 score per column, two packed 16-bit probabilities per column. + fragment_regs = cfg.softmax_score_fragment_regs + fragment_cols = fragment_regs // 2 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) + 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) + 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 + ), + ) + + packed_p = ( + s_arr.data_ptr() + .load(count=fragment_regs, alignment=4) + .to(cfg.q_dtype) + .bitcast(Int32) + ) + _keeps_tcgen05_st( + cfg, + prims.make_tmem_ptr(tmem_base + fragment * Int32(fragment_cols), Int32), + packed_p, + offset=cfg.tmem_p_cols_per_inst, + ) + cute.arch.fence_view_async_tmem_store() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if publishes_fragment: + prims.mbarrier_arrive(self._fragment_ready.data_ptr() + fragment) + total_sum += local_sum + 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. + + 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. + """ + 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 + ) + + @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. - # Eight independent chains keep the denominator update off one long - # dependency chain. Reuse s_arr for probabilities so only one K32 score - # fragment remains live while P is packed. + 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. + """ + pairs_per_fragment = self.cfg.softmax_score_fragment_regs // 2 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(16): + 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(cute.math.exp2(p0, fastmath=True)) - p1 = Float32(cute.math.exp2(p1, fastmath=True)) + if cutlass.const_expr( + _pair_uses_ex2_emulation(pair_idx, pairs_per_fragment) + ): + p0, p1 = _ex2_emulation_packed_f32x2(p0, p1) + else: + p0 = Float32(cute.math.exp2(p0, fastmath=True)) + p1 = Float32(cute.math.exp2(p1, fastmath=True)) s_arr[value_idx] = p0 s_arr[value_idx + 1] = p1 chain_idx = (pair_idx & 3) * 2 @@ -345,10 +524,6 @@ def compute_p_fragment( (p0, p1), ) ) - - # Collapse the eight reduction chains before packing P and publishing - # its barrier. This keeps only one sum scalar live across STTM instead - # of overlapping the full reduction state with packed P and addresses. sum01 = cute.arch.add_packed_f32x2( (sum_chains[0], sum_chains[1]), (sum_chains[2], sum_chains[3]), @@ -358,40 +533,43 @@ def compute_p_fragment( (sum_chains[6], sum_chains[7]), ) total_pair = cute.arch.add_packed_f32x2(sum01, sum23) - local_sum = Float32(total_pair[0] + total_pair[1]) + return Float32(total_pair[0] + total_pair[1]) - packed_p = ( - s_arr.data_ptr().load(count=32, alignment=4).to(cfg.q_dtype).bitcast(Int32) - ) + @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. - fragment_cols = cfg.softmax_score_fragment_regs // 2 - p_tmem_addr = ( - self._tmem_base_addr - + Int32(self._tmem_alloc.offset) - + Int32(fragment_idx * fragment_cols) - ) - _keeps_tcgen05_st( - cfg, - prims.make_tmem_ptr(p_tmem_addr, Int32), - packed_p, - offset=cfg.tmem_p_cols_per_inst, + 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, ) - # This lowers to the warp-collective tcgen05.wait::st. The explicit - # proxy fence then makes every lane's completed STTM visible through - # the lane-0 mbarrier publication consumed by the MMA warp. - cute.arch.fence_view_async_tmem_store() - prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) - - # KV256 aliases P with the score tile that produced it. Each softmax - # warp publishes its own rows after the TMEM store drains; BMM2 waits - # for all producer warps before consuming the fragment. - tidx, _, _ = cute.arch.thread_idx() - if (tidx & Int32(31)) == Int32(0): - prims.mbarrier_arrive(self._fragment_ready.data_ptr() + Int32(fragment_idx)) - - if cutlass.const_expr(fragment_idx != 0): - local_sum += self.tmem_s_ref.load_p_local_sum(0) - self.tmem_s_ref.store_p_local_sum(0, local_sum) + 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 @cute.jit def _compute_keeps_p( @@ -401,16 +579,18 @@ def _compute_keeps_p( new_max_arr: cutlass.Array, s_arr: cutlass.Array, ) -> None: - """Materialize one non-KV256 row-major Keeps probability tile. + """Materialize one complete-row Keeps probability tile. TQ128 gives each warp-group thread a complete 128-column row. TQ64 gives paired lanes the low/high 64-column halves of one row. Each lane writes disjoint packed blocks into the TMEM or SMEM layout consumed by - BMM2. + BMM2. Streamed profiles, including every block-sparse Keeps profile, + produce P through the rolled fragment loop instead. """ cfg = self.cfg - # KV256 uses compute_p_fragment so only one K32 score fragment is live. - assert not cfg.streams_tmem_p_fragments + # Every block-sparse Keeps profile streams P; only dense complete rows + # reach this path. + assert not cfg.streams_tmem_p_fragments and not cfg.use_block_sparse task_cache = _decode_gen_task_cache(stage_info) warp_grp_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] lane_idx = task_cache[_TASK_CACHE_LANE_IDX] @@ -439,7 +619,6 @@ def _compute_keeps_p( # without keeping a second 16-value P array live beside the S row. local_sum_pair_01 = (Float32(0.0), Float32(0.0)) local_sum_pair_23 = (Float32(0.0), Float32(0.0)) - # Each vector block is exactly 16 bytes after conversion. Compute and # pack adjacent pairs directly into their final register payload. packed_p_regs = cfg.num_packed_p_regs if cfg.uses_two_inst_tmem_p else 4 @@ -538,10 +717,10 @@ 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 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) + # 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 assert cfg.num_packed_p_regs % regs_per_store == 0 for store_idx in cutlass.range_constexpr( @@ -591,14 +770,18 @@ def _compute_keeps_p( # point and no extra named barrier is needed here. cute.arch.fence_view_async_shared() - @producer_work @cute.jit - def compute_p( + def _compute_p_impl( self, stage_info: StageInfo, *, 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, ) -> None: """Compute P from S, stage its BMM2 operand, and publish local sums.""" cfg = self.cfg @@ -616,6 +799,7 @@ def compute_p( # warp/lane ownership for SMEM offsets and STSM swizzles. 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): # Tile-Q=32 FP8 fast path: compute E4M3 P registers in the # same order consumed by the STSM helper, while also capturing @@ -796,10 +980,14 @@ def compute_p( 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 @@ -828,9 +1016,32 @@ def compute_p( ) 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): @@ -1057,6 +1268,33 @@ def compute_p( 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 @@ -1107,6 +1345,77 @@ def compute_p( # BMM2 cannot observe a partially written P tile. prims.barrier_cta_sync(4 + self.inst_id, thread_count=128) + @producer_work + @cute.jit + def compute_p( + self, + stage_info: StageInfo, + *, + new_max_arr: cutlass.Array, + s_arr: cutlass.Array, + ) -> None: + """Compute an exact/dense P tile without typed-route metadata.""" + + self._compute_p_impl( + 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), + ) + + @producer_work + @cute.jit + def compute_proxy_route_p( + self, + stage_info: StageInfo, + *, + 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, + swaps_route_flags: Uint32, + ) -> None: + """Normalize 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. + """ + + 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) + ) + 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) + 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( returns=( p_desc_0_slot, @@ -1136,7 +1445,7 @@ def p_operands( # stats-free columns of the corresponding S stage. p_stage_cols = cfg.tmem_s_cols if cutlass.const_expr(cfg.streams_tmem_p_fragments): - # KV256's four pipeline stages are K32 fragments of one P + # 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_tmem_addr = self._tmem_base_addr + Int32( @@ -1170,7 +1479,7 @@ def wait_p_fragment( *, fragment_idx: Constexpr[int], ) -> Int32: - """Wait for and return the next KV256 P-fragment TMEM address.""" + """Wait for and return the next streamed P-fragment TMEM address.""" cfg = self.cfg _ = stage_info assert cfg.streams_tmem_p_fragments @@ -1185,25 +1494,3 @@ def wait_p_fragment( self._tmem_alloc.offset + fragment_idx * fragment_cols ) return p_tmem_addr - - @consumer_work(work_attrs=WorkAttr.AUXILIARY) - @cute.jit - def wait_until_reusable_before_qk(self, stage_info: StageInfo) -> None: - """Wait until the previous same-instance PV has stopped reading P. - - KV256 aliases each streamed P instance with its next S accumulator. - The existing two-stage O pipeline commits stage ``inst_id`` only when - the matching PV completes, so its full barrier is also the P-reuse - credit. The S producer phase supplies the generation: the first QK - waits on the initially complete opposite parity, and every later QK - waits for the preceding PV without another commit or barrier. - """ - _ = stage_info - cfg = self.cfg - assert cfg.streams_tmem_p_fragments - assert cfg.o_stages == cfg.num_insts_kv == 2 - barrier = self.tmem_o_ref.pipeline.sync_object_full.get_barrier( - Int32(self.inst_id) - ) - _wait_for_mbarrier_phase(barrier, self.tmem_s_ref.producer_state.phase) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) 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 20bb78fb810c..173a13cb5447 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 @@ -497,6 +497,10 @@ class SmemKvTileResource(DecodeGenResourceBase): tma_desc_v: 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 + tma_desc_v_summary: cutlass.Pointer | None = None + tma_desc_k_summary_atom: cutlass.Pointer | None = None + tma_desc_v_summary_atom: cutlass.Pointer | None = None sparse_kv_metadata: "SmemBlockSparseKvMetadataResource | None" = None page_offsets_kv: "SmemPageOffsetsKvResource | None" = None seqlens_kv: cute.Pointer | None = None @@ -703,15 +707,24 @@ def _producer_load( assert self.sparse_kv_metadata is not None assert self.tma_desc_k_atom is not None assert self.tma_desc_v_atom is not None - # The positional TensorMaps keep the decode ABI stable. The - # primary K/V descriptors are KV128 for coarse routes and one atom - # for fine routes. The auxiliary slots always expose the atom - # descriptor and alias the primary descriptor for fine routes. + # K/V and summary sources expose the same primary/atom descriptor + # pair. Route kind selects the source; the geometry below alone + # selects the physical copy policy. tma_desc_atom = ( self.tma_desc_v_atom if cutlass.const_expr(self.kv_kind == KV_KIND_V) else self.tma_desc_k_atom ) + tma_desc_summary = ( + self.tma_desc_v_summary + if cutlass.const_expr(self.kv_kind == KV_KIND_V) + else self.tma_desc_k_summary + ) + tma_desc_summary_atom = ( + self.tma_desc_v_summary_atom + if cutlass.const_expr(self.kv_kind == KV_KIND_V) + else self.tma_desc_k_summary_atom + ) 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 @@ -775,6 +788,12 @@ def _producer_load( # join unrelated entries and must prove physical adjacency. fragment_chunk_elems = chunk_hd * 64 if prims.elect_sync(): + route_tma_desc = tma_desc + route_tma_desc_atom = tma_desc_atom + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): + if self.sparse_kv_metadata.route_is_proxy(): + route_tma_desc = tma_desc_summary + route_tma_desc_atom = tma_desc_summary_atom origin0, _ = self.sparse_kv_metadata.route_tma_coordinate( Int32(0), logical_b_idx, @@ -800,7 +819,7 @@ def _producer_load( local_tile_offset = chunk_idx * tile_chunk_elems prims.cp_async_bulk_tensor_shared_cta_global( stage_base.subview(local_tile_offset), - tma_desc, + route_tma_desc, ( Int32(global_head_dim_offset), origin0, @@ -832,7 +851,7 @@ def _producer_load( if adjacent: prims.cp_async_bulk_tensor_shared_cta_global( stage_base.subview(local_tile_offset), - tma_desc, + route_tma_desc, ( Int32(global_head_dim_offset), origin0, @@ -844,7 +863,7 @@ def _producer_load( else: prims.cp_async_bulk_tensor_shared_cta_global( stage_base.subview(local_tile_offset), - tma_desc_atom, + route_tma_desc_atom, ( Int32(global_head_dim_offset), origin0, @@ -857,7 +876,7 @@ def _producer_load( stage_base.subview( local_tile_offset + fragment_chunk_elems ), - tma_desc_atom, + route_tma_desc_atom, ( Int32(global_head_dim_offset), origin1, @@ -875,6 +894,10 @@ def _producer_load( atom_chunk_elems = chunk_hd * kv_atom_size atoms_per_route = cfg.tile_size_kv // kv_atom_size if prims.elect_sync(): + route_tma_desc_atom = tma_desc_atom + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): + if self.sparse_kv_metadata.route_is_proxy(): + route_tma_desc_atom = tma_desc_summary_atom stage_base = self._stage_base(stage_info) # Reuse each retained origin across all head-dimension # chunks. The copies still target disjoint SMEM regions @@ -913,7 +936,7 @@ def _producer_load( stage_base.subview( local_tile_offset + atom_idx * atom_chunk_elems ), - tma_desc_atom, + route_tma_desc_atom, ( Int32(global_head_dim_offset), origin, @@ -1487,6 +1510,10 @@ class SmemKvResource(DecodeGenResourceBase): tma_desc_v: 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 + tma_desc_v_summary: cutlass.Pointer | None = None + tma_desc_k_summary_atom: cutlass.Pointer | None = None + tma_desc_v_summary_atom: cutlass.Pointer | None = None sparse_kv_metadata0: "SmemBlockSparseKvMetadataResource | None" = None sparse_kv_metadata1: "SmemBlockSparseKvMetadataResource | None" = None page_offsets_kv: SmemPageOffsetsKvResource | None = None @@ -1922,6 +1949,29 @@ def _producer_load_kv_tile_256( ): assert self.page_offsets_kv is not None dense_page_ids = self.page_offsets_kv.page_ids(grouped_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 + if cutlass.const_expr(cfg.use_block_sparse): + sparse_tma_desc = ( + self.tma_desc_v_atom + if cutlass.const_expr(kv_kind == KV_KIND_V) + else self.tma_desc_k_atom + ) + assert sparse_tma_desc is not None + route_tma_desc = sparse_tma_desc + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): + assert sparse_kv_metadata is not None + summary_tma_desc = ( + self.tma_desc_v_summary_atom + if cutlass.const_expr(kv_kind == KV_KIND_V) + else self.tma_desc_k_summary_atom + ) + assert summary_tma_desc is not None + route_is_proxy = sparse_kv_metadata.route_is_proxy() + route_tma_desc = ( + summary_tma_desc if route_is_proxy else sparse_tma_desc + ) for semantic_block in cutlass.range_constexpr(4): token_coord = Int32(0) storage_coord = logical_b_idx @@ -1950,15 +2000,9 @@ def _producer_load_kv_tile_256( ) if cutlass.const_expr(cfg.use_block_sparse): - sparse_tma_desc = ( - self.tma_desc_v_atom - if cutlass.const_expr(kv_kind == KV_KIND_V) - else self.tma_desc_k_atom - ) - assert sparse_tma_desc is not None prims.cp_async_bulk_tensor_shared_cta_global( stage_base.subview(block_base), - sparse_tma_desc, + route_tma_desc, ( Int32(dim_half * 64), token_coord, 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 c59712aa8f3c..11b384a103bd 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 @@ -91,7 +91,9 @@ _KV_TILE_256_CORRECTION_THREADS = 128 _KV_TILE_256_LOGICAL_OUTPUT_ROWS = 64 -_KV_TILE_256_EXCHANGE_ROW_STRIDE = 132 +# One D32 fragment per logical output row, padded by four floats so adjacent +# rows fall on different bank groups. +_KV_TILE_256_EXCHANGE_FRAGMENT_STRIDE = 36 _KV_TILE_256_STATS_PER_THREAD = 4 @@ -157,10 +159,10 @@ def get_o_stage_dtype_bytes(self) -> int: ) def _kv_tile_256_exchange_entries(self) -> int: - """Return 128 lane-local stats plus 64 logical output rows.""" + """Return 128 lane-local stats plus one D32 fragment per output row.""" return ( _KV_TILE_256_CORRECTION_THREADS * _KV_TILE_256_STATS_PER_THREAD - + _KV_TILE_256_LOGICAL_OUTPUT_ROWS * _KV_TILE_256_EXCHANGE_ROW_STRIDE + + _KV_TILE_256_LOGICAL_OUTPUT_ROWS * _KV_TILE_256_EXCHANGE_FRAGMENT_STRIDE ) def _init_placeholder_state(self) -> None: @@ -306,22 +308,14 @@ def get_smem_requirements(self) -> list[SmemAllocation]: ) 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. Upper lanes publish - # one spatial half while lower lanes retain the matching fragment - # in registers. The dependency graph places this scratch after the - # shared KV ring so it can reuse the dead storage. - payload_bytes = self._kv_tile_256_exchange_entries() * 4 - exchange_bytes = payload_bytes - if self.cfg.uses_rotating_kv256_exchange: - assert payload_bytes <= self.cfg.smem_kv_tile_bytes - # Runtime selects one compact payload inside this explicit - # full-ring alias envelope. The envelope keeps every dynamic - # pointer within a declared allocation while the actual live - # exchange remains only 35,840 B in one 64-KiB stage. - exchange_bytes = self.cfg.smem_kv_tile_bytes * self.cfg.kv_stages + # 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. self._kv_tile_256_exchange_alloc = SmemAllocation( name=f"{self.name}_kvTile256Exchange", - size_bytes=exchange_bytes, + size_bytes=self._kv_tile_256_exchange_entries() * 4, alignment=16, ) allocs = [] @@ -874,60 +868,145 @@ def _fold_split_o_vec8( return output_vals, sum_val, new_max, new_max @cute.jit - def _store_final_o_vec8( + def _merge_kv_tile_256_peer_fragment( + self, + own_vals: cutlass.Array, + peer_vals, + first_col: Constexpr[int], + count: Constexpr[int], + ) -> cutlass.Array: + """Add the peer spatial half to ``count`` own columns from ``first_col``.""" + merged_vals = cutlass.Array( + Float32, + count, + space=cutlass.AddressSpace.rmem, + ) + for elem in cutlass.range_constexpr(0, count, 2): + value_idx = first_col + elem + merged = fadd2( + (own_vals[value_idx], own_vals[value_idx + 1]), + ( + Float32(peer_vals[value_idx]), + Float32(peer_vals[value_idx + 1]), + ), + ) + merged_vals[elem] = merged[0] + merged_vals[elem + 1] = merged[1] + return merged_vals + + @cute.jit + def _store_final_o_columns( self, final_o_dst, output_vals: cutlass.Array, norm_scale: Float32, + *, + count: Constexpr[int], + sector_aligned: cutlass.Boolean, ) -> None: - """Pack one contiguous 8-element output fragment to the final O dtype.""" + """Scale, pack, and store ``count`` contiguous final output columns. + + FP8 output packs four values per register and writes eight columns per + 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. + """ cfg = self.cfg + assert count % 8 == 0 if cutlass.const_expr(cfg.use_fp8_output): - 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[val_base], output_vals[val_base + 1]), + for chunk_idx in cutlass.range_constexpr(count // 8): + fp8_regs = self._pack_fp8_output_quads( + output_vals, norm_scale, chunk_idx * 8 ) - final_pairs[val_base] = pair[0] - final_pairs[val_base + 1] = pair[1] - final_fp8_regs = cutlass.Array(Int32, 2, space=cutlass.AddressSpace.rmem) - final_fp8_regs[0] = _pack_float4_to_fp8_e4m3( - final_pairs[0], - final_pairs[1], - final_pairs[2], - final_pairs[3], - ) - final_fp8_regs[1] = _pack_float4_to_fp8_e4m3( - final_pairs[4], - final_pairs[5], - final_pairs[6], - final_pairs[7], - ) - final_o_dst.store( - final_fp8_regs.data_ptr().load(count=2, alignment=4), - alignment=8, - ) - else: - final_regs = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) - for reg_idx in cutlass.range_constexpr(4): - pair = fmul2( - (norm_scale, norm_scale), - ( - output_vals[reg_idx * 2], - output_vals[reg_idx * 2 + 1], - ), + (final_o_dst + Int32(chunk_idx * 2)).store( + fp8_regs.data_ptr().load(count=2, alignment=4), + alignment=8, ) - if cutlass.const_expr(cfg.use_bf16_output): - final_regs[reg_idx] = _pack_float2_to_bf16(pair[0], pair[1]) + else: + final_regs = self._pack_final_o_regs(output_vals, norm_scale, count) + if cutlass.const_expr(count == 16): + if sector_aligned: + final_o_dst.store( + final_regs.data_ptr().load(count=8, alignment=4), + alignment=32, + ) else: - final_regs[reg_idx] = _pack_float2_to_fp16(pair[0], pair[1]) - final_o_dst.store( - final_regs.data_ptr().load(count=4, alignment=4), + self._store_16bit_output_chunks(final_o_dst, final_regs, count) + else: + self._store_16bit_output_chunks(final_o_dst, final_regs, count) + + @cute.jit + def _store_16bit_output_chunks( + self, + final_o_dst, + final_regs: cutlass.Array, + count: Constexpr[int], + ) -> None: + """Store packed 16-bit output columns as 16-byte chunks of eight columns.""" + for chunk_idx in cutlass.range_constexpr(count // 8): + (final_o_dst + Int32(chunk_idx * 4)).store( + (final_regs.data_ptr() + Int32(chunk_idx * 4)).load( + count=4, alignment=4 + ), alignment=16, ) + @cute.jit + def _pack_fp8_output_quads( + self, + output_vals: cutlass.Array, + norm_scale: Float32, + 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], + ), + ) + final_pairs[val_base] = pair[0] + final_pairs[val_base + 1] = pair[1] + fp8_regs = cutlass.Array(Int32, 2, space=cutlass.AddressSpace.rmem) + fp8_regs[0] = _pack_float4_to_fp8_e4m3( + final_pairs[0], final_pairs[1], final_pairs[2], final_pairs[3] + ) + fp8_regs[1] = _pack_float4_to_fp8_e4m3( + final_pairs[4], final_pairs[5], final_pairs[6], final_pairs[7] + ) + return fp8_regs + + @cute.jit + def _pack_final_o_regs( + self, + output_vals: cutlass.Array, + norm_scale: Float32, + 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], + ), + ) + if cutlass.const_expr(cfg.use_bf16_output): + final_regs[reg_idx] = _pack_float2_to_bf16(pair[0], pair[1]) + else: + final_regs[reg_idx] = _pack_float2_to_fp16(pair[0], pair[1]) + return final_regs + @cute.jit def _store_softmax_normalized_o_vec8( self, @@ -953,7 +1032,13 @@ def _store_softmax_normalized_o_vec8( mem_space=1, dtype=Int32, ) - self._store_final_o_vec8(final_o_dst, output_vals, norm_scale) + self._store_final_o_columns( + final_o_dst, + output_vals, + norm_scale, + count=8, + sector_aligned=cutlass.Boolean(False), + ) @cute.jit def _softmax_output_row_state( @@ -2274,24 +2359,6 @@ def _publish_and_reduce_cluster_swaps_partials( full_prefix=False, ) - @cute.jit - def _kv_tile_256_exchange_for_stage( - self, - stage_info: StageInfo, - scratch_stage: Int32 | None, - ) -> cutlass.Array: - """Return the fixed exchange or its dynamically selected KV stage.""" - if cutlass.const_expr(scratch_stage is None): - return self._kv_tile_256_exchange - return cutlass.Array( - stage_info.context.smem_base.data_ptr() - + self._kv_tile_256_exchange_alloc.offset - + scratch_stage * Int32(self.cfg.smem_kv_tile_bytes), - dtype=Float32, - shape=(self._kv_tile_256_exchange_entries(),), - addrspace=3, - ) - @cute.jit def _kv_tile_256_temporal_fragment( self, @@ -2329,6 +2396,99 @@ def _kv_tile_256_temporal_fragment( ) return cutlass.Vector.from_elements(combined, Float32) + @cute.jit + def _store_kv_tile_256_direct_fragment( + self, + own_vals: cutlass.Array, + peer_vals, + *, + fragment_col: Constexpr[int], + dst_row_base: Int64, + norm_scale: Float32, + valid_output_row: cutlass.Boolean, + o_is_32b_aligned: cutlass.Boolean, + ) -> None: + """Merge one D32 fragment with its peer half and write the final output. + + Sixteen columns go out per store so each lane writes one full 32-byte + sector. Adjacent lanes own adjacent rows, so 16-byte stores would leave + every sector half-written twice. + """ + cfg = self.cfg + for vector_pair in cutlass.range_constexpr(2): + pair_col = vector_pair * 16 + merged_vals = self._merge_kv_tile_256_peer_fragment( + own_vals, + peer_vals, + pair_col, + 16, + ) + if valid_output_row: + output_col = fragment_col + pair_col + dst_offset = dst_row_base + Int32(output_col * cfg.o_dtype_bytes) + final_o_dst = cutlass.inttoptr( + self.o_ptr.toint() + cutlass.Int64(dst_offset), + mem_space=1, + dtype=Int32, + ) + self._store_final_o_columns( + final_o_dst, + merged_vals, + norm_scale, + count=16, + sector_aligned=o_is_32b_aligned, + ) + + @cute.jit + def _store_kv_tile_256_partial_fragment( + self, + own_vals: cutlass.Array, + peer_vals, + *, + fragment_col: Constexpr[int], + partial_row_base: Int64, + partial_scale: Float32, + valid_output_row: cutlass.Boolean, + ) -> None: + """Merge one D32 fragment with its peer half and write the split-KV partial.""" + cfg = self.cfg + partial_o_uses_bf16 = ( + cfg.use_bf16_separate_partial_o + if cfg.use_separate_reduction_kernel + else cfg.use_bf16_output + ) + for vector_idx in cutlass.range_constexpr(4): + vector_col = vector_idx * 8 + output_vals = self._merge_kv_tile_256_peer_fragment( + own_vals, + peer_vals, + vector_col, + 8, + ) + if valid_output_row: + output_col = fragment_col + vector_col + scaled_values: tuple = () + for elem in cutlass.range_constexpr(0, 8, 2): + scaled_values += fmul2( + (partial_scale, partial_scale), + (output_vals[elem], output_vals[elem + 1]), + ) + scaled_vector = cutlass.Vector.from_elements(scaled_values, Float32) + if cutlass.const_expr(partial_o_uses_bf16): + packed = scaled_vector.to(cutlass.BFloat16).bitcast(Int32) + else: + packed = scaled_vector.to(cutlass.Float16).bitcast(Int32) + # Split-KV partials are 16-bit, so the column offset follows + # the partial element width + partial_o_dst = cutlass.inttoptr( + self.partial_o_ptr.toint() + + partial_row_base + + Int64(output_col * 2), + mem_space=1, + dtype=Int32, + ) + partial_o_dst.store(packed, alignment=16) + @cute.jit def _kv_tile_256_merge_spatial_output( self, @@ -2357,18 +2517,13 @@ def _kv_tile_256_merge_spatial_output( serializing the complete D128 upper and lower halves. """ cfg = self.cfg - partial_o_uses_bf16 = ( - cfg.use_bf16_separate_partial_o - if cfg.use_separate_reduction_kernel - else cfg.use_bf16_output - ) output_exchange_base = Int32( _KV_TILE_256_CORRECTION_THREADS * _KV_TILE_256_STATS_PER_THREAD ) output_lane = exchange_idx < Int32(_KV_TILE_256_LOGICAL_OUTPUT_ROWS) exchange_row_idx = exchange_idx & Int32(_KV_TILE_256_LOGICAL_OUTPUT_ROWS - 1) output_exchange_row_base = output_exchange_base + exchange_row_idx * Int32( - _KV_TILE_256_EXCHANGE_ROW_STRIDE + _KV_TILE_256_EXCHANGE_FRAGMENT_STRIDE ) logical_output_row_idx = q_row_offset + exchange_row_idx valid_output_row = cutlass.Boolean(False) @@ -2394,6 +2549,9 @@ def _kv_tile_256_merge_spatial_output( 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, @@ -2419,10 +2577,17 @@ def _kv_tile_256_merge_spatial_output( weight00=weight00, weight10=weight10, ) + if cutlass.const_expr(fragment != 0): + # The single fragment buffer is reused: lower lanes must have + # consumed the previous peer fragment before it is overwritten. + prims.barrier_cta_sync( + self.store_barrier_id, + thread_count=cfg.correction_barrier_threads, + ) if exchange_idx >= Int32(_KV_TILE_256_LOGICAL_OUTPUT_ROWS): - ( - exchange.data_ptr() + output_exchange_row_base + Int32(fragment_col) - ).store(own_vals, alignment=16) + (exchange.data_ptr() + output_exchange_row_base).store( + own_vals, alignment=16 + ) # Lower lanes keep ``own_vals`` live across the barrier. Once every # lane arrives, upper lanes may prepare the next fragment while @@ -2433,69 +2598,28 @@ def _kv_tile_256_merge_spatial_output( ) if output_lane: - peer_vals = ( - exchange.data_ptr() + output_exchange_row_base + Int32(fragment_col) - ).load(count=32, alignment=16) - for vector_idx in cutlass.range_constexpr(4): - vector_col = vector_idx * 8 - output_vals = cutlass.Array( - Float32, - 8, - space=cutlass.AddressSpace.rmem, + peer_vals = (exchange.data_ptr() + output_exchange_row_base).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=fragment_col, + dst_row_base=dst_row_base, + norm_scale=norm_scale, + 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=fragment_col, + partial_row_base=partial_row_base, + partial_scale=partial_scale, + valid_output_row=valid_output_row, ) - for elem in cutlass.range_constexpr(0, 8, 2): - value_idx = vector_col + elem - merged = fadd2( - (own_vals[value_idx], own_vals[value_idx + 1]), - ( - Float32(peer_vals[value_idx]), - Float32(peer_vals[value_idx + 1]), - ), - ) - output_vals[elem] = merged[0] - output_vals[elem + 1] = merged[1] - if valid_output_row: - output_col = fragment_col + vector_col - if cutlass.const_expr(cfg.use_split_kv): - scaled_values: tuple = () - for elem in cutlass.range_constexpr(0, 8, 2): - scaled_values += fmul2( - (partial_scale, partial_scale), - (output_vals[elem], output_vals[elem + 1]), - ) - scaled_vector = cutlass.Vector.from_elements( - scaled_values, Float32 - ) - if cutlass.const_expr(partial_o_uses_bf16): - packed = scaled_vector.to(cutlass.BFloat16).bitcast( - Int32 - ) - else: - packed = scaled_vector.to(cutlass.Float16).bitcast( - Int32 - ) - partial_o_dst = cutlass.inttoptr( - self.partial_o_ptr.toint() - + partial_row_base - + Int64(output_col * cfg.o_dtype_bytes), - mem_space=1, - dtype=Int32, - ) - partial_o_dst.store(packed, alignment=16) - else: - dst_offset = dst_row_base + Int32( - output_col * cfg.o_dtype_bytes - ) - final_o_dst = cutlass.inttoptr( - self.o_ptr.toint() + cutlass.Int64(dst_offset), - mem_space=1, - dtype=Int32, - ) - self._store_final_o_vec8( - final_o_dst, - output_vals, - norm_scale, - ) if cutlass.const_expr(cfg.use_split_kv): if valid_output_row: @@ -2538,7 +2662,6 @@ def _kv_tile_256_tail_epilogue( self, stage_info: StageInfo, *, - scratch_stage: Int32 | None, tail_o_stage_idx_0: Int32, tail_o_stage_idx_1: Int32, inst0_new_max_arr: cutlass.Array, @@ -2553,15 +2676,13 @@ 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 in SMEM after the shared KV ring is - dead, then publishes the ordinary logical Q64xD128 output. + their stats, stages one spatial half through its dedicated SMEM + exchange one D32 fragment at a time, then publishes the ordinary + logical Q64xD128 output. """ cfg = self.cfg assert cfg.headdim == 128 - exchange = self._kv_tile_256_exchange_for_stage( - stage_info, - scratch_stage, - ) + exchange = self._kv_tile_256_exchange exchange_idx = warp_grp_thread_idx peer_idx = exchange_idx ^ Int32(_KV_TILE_256_LOGICAL_OUTPUT_ROWS) @@ -4371,7 +4492,6 @@ def _correction_tail_epilogue_impl( self, stage_info: StageInfo, *, - scratch_stage: Int32 | None, o_stage_idx: Int32, tail_o_stage_idx_0: Int32, tail_o_stage_idx_1: Int32, @@ -4417,7 +4537,6 @@ def _correction_tail_epilogue_impl( if cutlass.const_expr(cfg.tile_size_kv == 256): self._kv_tile_256_tail_epilogue( stage_info, - scratch_stage=scratch_stage, tail_o_stage_idx_0=tail_o_stage_idx_0, tail_o_stage_idx_1=tail_o_stage_idx_1, inst0_new_max_arr=inst0_new_max_arr, @@ -4493,9 +4612,6 @@ def _correction_tail_epilogue_impl( ) return - # Task Scheduling routes every non-constexpr work argument as a required - # data-flow token. Keep separate fixed/rotating entry points so only the - # latter consumes ``scratch_stage``; both still share the implementation. @producer_work @cute.jit def correction_tail_epilogue( @@ -4512,42 +4628,9 @@ def correction_tail_epilogue( inst1_new_max_arr: cutlass.Array, inst1_sum_arr: cutlass.Array, ) -> None: - """Run the ordinary fixed-exchange tail epilogue.""" - self._correction_tail_epilogue_impl( - stage_info, - scratch_stage=None, - o_stage_idx=o_stage_idx, - tail_o_stage_idx_0=tail_o_stage_idx_0, - tail_o_stage_idx_1=tail_o_stage_idx_1, - old_max_arr=old_max_arr, - new_max_arr=new_max_arr, - inst0_new_max_arr=inst0_new_max_arr, - inst0_sum_arr=inst0_sum_arr, - inst1_new_max_arr=inst1_new_max_arr, - inst1_sum_arr=inst1_sum_arr, - ) - - @producer_work - @cute.jit - def correction_tail_epilogue_rotating_exchange( - self, - stage_info: StageInfo, - *, - scratch_stage: Int32, - o_stage_idx: Int32, - tail_o_stage_idx_0: Int32, - tail_o_stage_idx_1: Int32, - old_max_arr: cutlass.Array, - new_max_arr: cutlass.Array, - inst0_new_max_arr: cutlass.Array, - inst0_sum_arr: cutlass.Array, - inst1_new_max_arr: cutlass.Array, - inst1_sum_arr: cutlass.Array, - ) -> None: - """Run persistent direct output in the stage named by its credit.""" + """Normalize the final O stages and publish the output tile.""" self._correction_tail_epilogue_impl( stage_info, - scratch_stage=scratch_stage, o_stage_idx=o_stage_idx, tail_o_stage_idx_0=tail_o_stage_idx_0, tail_o_stage_idx_1=tail_o_stage_idx_1, 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 fcf2cbe71730..fa8928695a98 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 @@ -59,7 +59,7 @@ def _pv_mma_operand_contract_for_config( cfg.headdim if cfg.head_dim_per_stage_kv == 0 else cfg.head_dim_kv_stage ) if cfg.use_keeps_mma_ab: - if cfg.tile_size_kv == 256: + if cfg.uses_ws_2x2_datapath: # The WS 2x2 PV instruction exposes two spatial D128 partials as # one physical KV256 operation. Correction merges those spatial # halves after the two temporal decode streams are complete. @@ -196,7 +196,7 @@ def vp_mma_loop_fragment( p_tmem_addr: Int32, fragment_idx: Constexpr[int], ) -> None: - """Issue one K32 fragment of a KV256 loop PV tile.""" + """Issue one K32 fragment of a streamed loop PV tile.""" self._vp_mma_fragment( stage_info, v_desc=v_desc, @@ -215,7 +215,7 @@ def vp_mma_tail_fragment( p_tmem_addr: Int32, fragment_idx: Constexpr[int], ) -> None: - """Issue one K32 fragment of the final KV256 PV tile.""" + """Issue one K32 fragment of the final streamed PV tile.""" self._vp_mma_fragment( stage_info, v_desc=v_desc, @@ -234,14 +234,16 @@ def _vp_mma_fragment( fragment_idx: Constexpr[int], initial_scale_d, ) -> None: - """Issue the two WS MMA steps covered by one KV256 P fragment. + """Issue the two 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. + ``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. """ cfg = self.cfg - assert cfg.tile_size_kv == 256 and cfg.uses_two_inst_tmem_p + assert cfg.streams_tmem_p_fragments v_desc = _freeze_smem_descriptor(v_desc) task_cache = _decode_gen_task_cache(stage_info) @@ -268,17 +270,32 @@ def _vp_mma_fragment( p_operand = prims.make_tmem_ptr( p_tmem_addr + Int32(local_k_step * 8), Int32 ) - iter_v_desc = v_desc + Int32( - (k_step // 4) * cfg.headdim * 16 + (k_step % 4) * 128 - ) - tcgen05_mma_ws( - _mma_kind_for_qkv(cfg), - tmem_col, - p_operand, - iter_v_desc, - idesc, - initial_scale_d or fragment_idx != 0 or local_k_step != 0, - ) + 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. + iter_v_desc = v_desc + Int32( + (k_step // 4) * cfg.headdim * 16 + (k_step % 4) * 128 + ) + tcgen05_mma_ws( + _mma_kind_for_qkv(cfg), + tmem_col, + p_operand, + iter_v_desc, + idesc, + scale_d, + ) + else: + iter_v_desc = v_desc + Int32(k_step * 128) + prims.tcgen05_mma( + _mma_kind_for_qkv(cfg), + prims.CTAGroup.CTA_1, + tmem_col, + p_operand, + iter_v_desc, + idesc, + scale_d, + ) @cute.jit def _vp_mma( 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 703134a5b61d..b010dfdb9d21 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 @@ -41,10 +41,9 @@ producer_work, ) -from ...._block_sparse.common import _MAX_KV_ATOM_SIZE from ...._block_sparse.prepared import _PREPARED_ROUTE_IS_FULL_FLAG from ..fmha_decode_config import CAUSAL, FmhaDecodeConfig -from ..fmha_decode_constants import KV_TILE_256_RESCALE_THRESHOLD_LOG2 +from ..fmha_decode_constants import SOFTMAX_RESCALE_THRESHOLD_LOG2 from ...tcgen05_compat import tcgen05_mma_ws from ...placeholder_helpers import ( _placeholder_local_array, @@ -76,13 +75,14 @@ _mma_kind_for_qkv, _neg_max_f32, _softmax_scale_pair_width, + _swaps_routed_coordinate, _q_row_is_valid_for_seq, _q_row_token_and_local_head, _q_group_token_base, _softmax_tile_idx, ) from .smem_block_sparse_metadata import ( - _SOFTMAX_TOKEN_MASK_IS_FULL_FLAG, + _SOFTMAX_ROUTE_IS_PROXY_FLAG, _swaps_forwards_packed_route_full, ) from .helpers_kv_tile_idx import ( @@ -103,20 +103,13 @@ _wspro_reduce_max4, ) -# A block-sparse route often changes the exact row maximum without changing it -# enough to justify rescaling the live O tile. Keeping the prior anchor within -# this bound makes the correction scale exactly one and bounds FP16/BF16 P by -# 2**8. As in the FlashInfer/TRT-LLM policy, this assumes normal model logits -# rather than adversarial values outside the qualified probability bound. -_BLOCK_SPARSE_RESCALE_THRESHOLD_LOG2 = 8.0 - def _swaps_uses_origin0_k32_full_guard(cfg: FmhaDecodeConfig) -> bool: """Whether one staged origin can prove this warp's K32 slice valid.""" return ( cfg.kv_block_size >= 32 - and not cfg.use_kv_valid_bits + and not cfg.uses_prepared_score_keep_words and not cfg.uses_uniform_causal_mask and not cfg.uses_per_row_causal_mask ) @@ -126,7 +119,7 @@ def _swaps_token_word_covers_kv_tail(cfg: FmhaDecodeConfig) -> bool: """Whether SWAP's prepared token word covers the logical KV tail.""" return ( - cfg.use_kv_valid_bits + cfg.uses_prepared_score_keep_words and not cfg.uses_uniform_causal_mask and not cfg.uses_per_row_causal_mask ) @@ -145,42 +138,34 @@ def _swaps_uses_token_only_score_validity(cfg: FmhaDecodeConfig) -> bool: @cute.jit -def _can_skip_sparse_keeps_structural_mask( - q_row_is_valid: Boolean, - origin0: Int32, - origin1: Int32, - valid0: Int32, - valid1: Int32, - seq_len_kv: Int32, - causal_end: Int32, +def _dense_fragment_keep_word( + rows_are_active: Boolean, + visible_start: Int32, + visible_end: Int32, *, - apply_causal_mask: cutlass.Constexpr[bool], -) -> Boolean: - """Return whether one Keeps row needs no Q/tail/causal predicate. + fragment_regs: cutlass.Constexpr[int], +) -> Uint32: + """Return the keep word of one dense K32 fragment. - Token-bit masking is independent. Comparing against the last complete - KV64 origin avoids overflowing an origin near the Int32 upper bound. + ``visible_start`` and ``visible_end`` are the visible token range relative + to the fragment's first column. Columns outside ``[start, end)`` are + masked; an inactive tile or Q row masks the whole fragment. """ - - fragment_size = Int32(_MAX_KV_ATOM_SIZE) - last_complete_origin = seq_len_kv - fragment_size - can_skip = Boolean( - q_row_is_valid - and valid0 != Int32(0) - and valid1 != Int32(0) - and origin0 <= last_complete_origin - and origin1 <= last_complete_origin - ) - if cutlass.const_expr(apply_causal_mask): - last_causal_origin = causal_end - fragment_size - can_skip = Boolean( - can_skip and origin0 <= last_causal_origin and origin1 <= last_causal_origin - ) - return can_skip + keep_word = Uint32(0) + if rows_are_active: + first_kept = cute.math.max(visible_start, Int32(0)) + end_kept = cute.math.min(visible_end, Int32(fragment_regs)) + if first_kept < end_kept: + # Both shift amounts stay strictly below the register width: + # 1 <= end_kept <= fragment_regs and 0 <= first_kept < end_kept. + keep_word = (Uint32(0xFFFFFFFF) >> (Int32(fragment_regs) - end_kept)) & ( + Uint32(0xFFFFFFFF) << first_kept + ) + return keep_word @cute.jit -def _sparse_k32_effective_keep_word( +def _sparse_effective_keep_word( q_row_is_valid: Boolean, fragment_origin: Int32, fragment_valid: Int32, @@ -667,7 +652,7 @@ def _qk_mma( a_desc, b_desc = q_desc, k_desc else: a_desc, b_desc = k_desc, q_desc - if cutlass.const_expr(cfg.tile_size_kv == 256): + if cutlass.const_expr(cfg.uses_ws_2x2_datapath): tcgen05_mma_ws( _mma_kind_for_qkv(cfg), tmem_col, @@ -842,6 +827,7 @@ def _resolve_keeps_tile_context(self, stage_info: StageInfo): is_valid_effective_tile, is_masked_final_wave, tile_is_unmasked, + tile_has_valid_scores, ) @cute.jit @@ -896,49 +882,32 @@ def _publish_keeps_softmax_state( ) -> None: """Publish a masked Keeps row and its updated softmax anchor.""" - new_anchor = cute.math.max(old_max, tile_max, ftz=True) - if cutlass.const_expr( - self.cfg.use_block_sparse and _BLOCK_SPARSE_RESCALE_THRESHOLD_LOG2 > 0.0 - ): - # Online softmax only requires a common finite reference for P, - # sum, and O; it does not require the exact row maximum. Defer a - # small anchor increase so correction can skip a TMEM O rescale. - rescale_log2 = (old_max - new_anchor) * self.scale_softmax_log2 - if (old_max != _neg_max_f32()) and ( - rescale_log2 >= Float32(-_BLOCK_SPARSE_RESCALE_THRESHOLD_LOG2) - ): - new_anchor = old_max old_max_arr[0] = old_max sum_arr[0] = running_sum - new_max_arr[0] = new_anchor + new_max_arr[0] = self._softmax_anchor(old_max, tile_max) 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 _mask_and_store_sparse_keeps_atom( - self, - s_vals: cutlass.Array, - loaded: cutlass.Vector, - token_word: Uint32, - *, - atom_col: Constexpr[int], - token_mask_is_required: cutlass.Boolean, - ) -> None: - """Store one 32-score atom, applying its token word when required.""" - - if token_mask_is_required: - for atom_reg_idx in cutlass.range_constexpr(32): - score_idx = atom_col + atom_reg_idx - s_vals[score_idx] = loaded[atom_reg_idx] - token_bit_is_valid = ( - (token_word >> Int32(atom_reg_idx)) & Uint32(1) - ) != Uint32(0) - if not token_bit_is_valid: - s_vals[score_idx] = _neg_max_f32() - else: - for atom_reg_idx in cutlass.range_constexpr(32): - score_idx = atom_col + atom_reg_idx - s_vals[score_idx] = loaded[atom_reg_idx] + def _softmax_anchor(self, old_max: Float32, tile_max: Float32) -> Float32: + """Return the exponent reference max for the tile's P pass. + + Online softmax only requires a common finite reference for P, the + running sum, and O; it does not require the exact row maximum. + Profiles that defer anchor updates keep the previous reference while + the tile raises it by less than ``SOFTMAX_RESCALE_THRESHOLD_LOG2`` + log2 units, so correction can skip the in-place TMEM O rescale. The + 16-bit P path represents the bounded values above one, and the + numerator and denominator stay in the same scale frame. Larger jumps + still rebase to keep P comfortably in range. + """ + new_max = cute.math.max(old_max, tile_max, ftz=True) + if cutlass.const_expr(self.cfg.defers_softmax_anchor_updates): + if old_max != _neg_max_f32(): + max_delta_log2 = self.scale_softmax_log2 * (old_max - new_max) + if max_delta_log2 >= Float32(-SOFTMAX_RESCALE_THRESHOLD_LOG2): + new_max = old_max + return new_max @cute.jit def _load_keeps_fragment_impl( @@ -954,9 +923,8 @@ def _load_keeps_fragment_impl( is_masked_final_wave: cutlass.Boolean, *, apply_boundary_mask: Constexpr[bool], - fragment_idx: Constexpr[int] = 0, ) -> None: - """Load one Keeps score fragment with a compile-time mask policy. + """Load one complete-row Keeps score tile with a compile-time mask policy. The caller chooses the masked/unmasked path before TMEM load. Keeping the score fragment out of the branch condition avoids carrying 64/128 live @@ -966,8 +934,7 @@ def _load_keeps_fragment_impl( """ cfg = self.cfg task_cache = _decode_gen_task_cache(stage_info) - num_s_regs = cfg.softmax_score_fragment_regs - fragment_reg_base = fragment_idx * num_s_regs + num_s_regs = cfg.num_s_regs_per_thread base_addr = ( task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + Int32(self._alloc.offset) @@ -977,9 +944,7 @@ def _load_keeps_fragment_impl( atom_col = load_atom_idx * 32 loaded = _keeps_tcgen05_ld( cfg, - prims.make_tmem_ptr( - base_addr + Int32(fragment_reg_base + atom_col), Float32 - ), + prims.make_tmem_ptr(base_addr + Int32(atom_col), Float32), num=32, offset=cfg.tile_size_kv // 2, ) @@ -1023,7 +988,7 @@ def _load_keeps_fragment_impl( token_idx = tile_offset_k + _keeps_score_col( cfg, warp_grp_thread_idx, - fragment_reg_base + reg_idx, + reg_idx, col_base, ) if token_idx >= element_mask_end_idx: @@ -1051,7 +1016,7 @@ def _load_keeps_fragment_impl( score_col = _keeps_score_col( cfg, warp_grp_thread_idx, - fragment_reg_base + reg_idx, + reg_idx, col_base, ) if cutlass.const_expr(cfg.use_sliding_window_causal): @@ -1084,8 +1049,6 @@ def _load_keeps_fragment( is_valid_effective_tile: cutlass.Boolean, is_masked_final_wave: cutlass.Boolean, tile_is_unmasked: cutlass.Boolean, - *, - fragment_idx: Constexpr[int] = 0, ) -> None: """Select the masked or unmasked fragment loader before LDTM. @@ -1106,7 +1069,6 @@ def _load_keeps_fragment( is_valid_effective_tile, is_masked_final_wave, apply_boundary_mask=False, - fragment_idx=fragment_idx, ) else: self._load_keeps_fragment_impl( @@ -1120,275 +1082,8 @@ def _load_keeps_fragment( is_valid_effective_tile, is_masked_final_wave, apply_boundary_mask=True, - fragment_idx=fragment_idx, ) - @cute.jit - def _reduce_keeps_fragment_max(self, s_vals: cutlass.Array) -> Float32: - """Reduce the row maximum of a previously loaded Keeps fragment.""" - cfg = self.cfg - num_s_regs = cfg.softmax_score_fragment_regs - - max_chains = cutlass.Array(Float32, 4, space=cutlass.AddressSpace.rmem) - for chain_idx in cutlass.range_constexpr(4): - max_chains[chain_idx] = _neg_max_f32() - for reg_base in cutlass.range_constexpr(0, num_s_regs, 4): - for chain_idx in cutlass.range_constexpr(4): - max_chains[chain_idx] = cute.math.max( - max_chains[chain_idx], - s_vals[reg_base + chain_idx], - ftz=True, - ) - tile_max = cute.math.max( - cute.math.max(max_chains[0], max_chains[1], ftz=True), - cute.math.max(max_chains[2], max_chains[3], ftz=True), - ftz=True, - ) - if cutlass.const_expr(cfg.tile_size_q == 64 and cfg.tile_size_kv != 256): - return cute.math.max( - tile_max, - Float32( - prims.shfl_sync( - thread_mask=0xFFFFFFFF, - val=tile_max, - offset=16, - mask_and_clamp=0x1F, - kind=prims.Shfl.BFLY, - ) - ), - ftz=True, - ) - return tile_max - - @cute.jit - def _decode_sparse_mask_metadata( - self, - routed_origin0: Int32, - routed_origin1: Int32, - routed_route_flags: Int32, - routed_token_word0: Uint32, - routed_token_word1: Uint32, - routed_token_word2: Uint32, - routed_token_word3: Uint32, - ) -> tuple[Int32, Int32, Int32, Int32, cutlass.Array, cutlass.Boolean]: - """Decode one prepared, register-routed mask payload.""" - - origin0 = Int32(routed_origin0) - origin1 = Int32(routed_origin1) - route_flags = Int32(routed_route_flags) - valid0 = route_flags & Int32(1) - valid1 = (route_flags >> Int32(1)) & Int32(1) - route_token_mask_is_full = cutlass.Boolean(False) - if cutlass.const_expr(self.cfg.use_kv_valid_bits): - route_token_mask_is_full = cutlass.Boolean( - (route_flags & Int32(_SOFTMAX_TOKEN_MASK_IS_FULL_FLAG)) != Int32(0) - ) - - num_local_words = 4 if self.cfg.tile_size_q == 128 else 2 - local_token_words = cutlass.Array( - Uint32, - num_local_words, - space=cutlass.AddressSpace.rmem, - ) - for word_idx in cutlass.range_constexpr(num_local_words): - local_token_words[word_idx] = Uint32(0xFFFFFFFF) - if cutlass.const_expr(self.cfg.use_kv_valid_bits): - if not route_token_mask_is_full: - if cutlass.const_expr(self.cfg.tile_size_q == 128): - local_token_words[0] = Uint32(routed_token_word0) - local_token_words[1] = Uint32(routed_token_word1) - local_token_words[2] = Uint32(routed_token_word2) - local_token_words[3] = Uint32(routed_token_word3) - else: - local_word0 = Uint32(routed_token_word0) - local_word1 = Uint32(routed_token_word1) - local_token_words[0] = local_word0 - local_token_words[1] = local_word1 - return ( - origin0, - origin1, - valid0, - valid1, - local_token_words, - route_token_mask_is_full, - ) - - @cute.jit - def _compute_softmax_loop_sparse_keeps( - self, - stage_info: StageInfo, - *, - old_max_arr: cutlass.Array, - sum_arr: cutlass.Array, - new_max_arr: cutlass.Array, - s_arr: cutlass.Array, - routed_origin0: Int32, - routed_origin1: Int32, - routed_route_flags: Int32, - routed_token_word0: Uint32, - routed_token_word1: Uint32, - routed_token_word2: Uint32, - routed_token_word3: Uint32, - ) -> tuple[object, object, object, object]: - """Load Keeps scores and mask them in logical KV coordinates.""" - cfg = self.cfg - num_s_regs = cfg.num_s_regs_per_thread - old_max = new_max_arr[0] - running_sum = sum_arr[0] - s_vals = cutlass.Array(Float32, num_s_regs, space=cutlass.AddressSpace.rmem) - task_cache = _decode_gen_task_cache(stage_info) - seq_len_kv = _load_runtime_seq_len_kv( - self.seqlens_kv, - self.max_seq_len_kv, - stage_info, - Int32(0), - Int32(0), - ) - warp_grp_thread_idx = Int32(task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX]) - lane_idx = Int32(task_cache[_TASK_CACHE_LANE_IDX]) - tile_row_idx = _keeps_row_idx(cfg, warp_grp_thread_idx) - col_base = _keeps_col_base(cfg, lane_idx, num_s_regs) - ( - origin0, - origin1, - valid0, - valid1, - local_token_words, - route_token_mask_is_full, - ) = self._decode_sparse_mask_metadata( - routed_origin0=routed_origin0, - routed_origin1=routed_origin1, - routed_route_flags=routed_route_flags, - routed_token_word0=routed_token_word0, - routed_token_word1=routed_token_word1, - routed_token_word2=routed_token_word2, - routed_token_word3=routed_token_word3, - ) - - base_addr = ( - task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] - + Int32(self._alloc.offset) - + self._softmax_loop_stage_slot_offset(stage_info) - ) - num_load_atoms = num_s_regs // 32 - if cutlass.const_expr(cfg.tile_size_q == 64 and cfg.use_kv_valid_bits): - token_mask_is_required = not route_token_mask_is_full - - # Keep each Q64 atom's load, wait, and mask together. A/B testing - # showed that hoisting both loads extends live fragment ranges and - # regresses the Q64 code generated by ptxas. - for load_atom_idx in cutlass.range_constexpr(2): - atom_col = load_atom_idx * 32 - loaded = _keeps_tcgen05_ld( - cfg, - prims.make_tmem_ptr(base_addr + Int32(atom_col), Float32), - num=32, - offset=cfg.tile_size_kv // 2, - ) - prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) - self._mask_and_store_sparse_keeps_atom( - s_vals, - loaded, - local_token_words[load_atom_idx], - atom_col=atom_col, - token_mask_is_required=token_mask_is_required, - ) - else: - for load_atom_idx in cutlass.range_constexpr(num_load_atoms): - atom_col = load_atom_idx * 32 - loaded = _keeps_tcgen05_ld( - cfg, - prims.make_tmem_ptr(base_addr + Int32(atom_col), Float32), - num=32, - offset=cfg.tile_size_kv // 2, - ) - prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) - for atom_reg_idx in cutlass.range_constexpr(32): - score_idx = atom_col + atom_reg_idx - s_vals[score_idx] = loaded[atom_reg_idx] - - logical_q_group_idx = _logical_q_group_idx(cfg, stage_info, self.q_group_idx) - q_token_idx, _ = _q_row_token_and_local_head( - cfg, - self.h_r, - logical_q_group_idx, - tile_row_idx, - ) - q_row_is_valid = _q_row_is_valid_for_seq( - cfg, - self.h_r, - logical_q_group_idx, - tile_row_idx, - self.seq_len_q, - ) - causal_end = seq_len_kv - self.seq_len_q + q_token_idx + Int32(1) - can_skip_structural_mask = _can_skip_sparse_keeps_structural_mask( - q_row_is_valid, - origin0, - origin1, - valid0, - valid1, - seq_len_kv, - causal_end, - apply_causal_mask=cfg.mask_type == CAUSAL, - ) - # This guard covers only route/Q/tail/causal structure. Q64 token - # holes were applied while materializing its two LDTM atoms; Q128 - # applies them in the post-pass below. - if not can_skip_structural_mask: - for reg_idx in cutlass.range_constexpr(num_s_regs): - fragment_offset = Int32(reg_idx) - logical_k = origin0 + fragment_offset - fragment_valid = valid0 - if cutlass.const_expr(cfg.tile_size_q == 128 and reg_idx >= 64): - fragment_offset = Int32(reg_idx - 64) - logical_k = origin1 + fragment_offset - fragment_valid = valid1 - elif cutlass.const_expr(cfg.tile_size_q == 64): - if col_base >= Int32(64): - logical_k = origin1 + fragment_offset - fragment_valid = valid1 - - score_is_valid = ( - q_row_is_valid - and fragment_valid != Int32(0) - and logical_k < seq_len_kv - ) - if cutlass.const_expr(cfg.mask_type == CAUSAL): - score_is_valid = score_is_valid and logical_k < causal_end - if not score_is_valid: - s_vals[reg_idx] = _neg_max_f32() - - # Q128 deliberately keeps all four LDTM atoms adjacent: unlike Q64, - # interleaving each load with mask control flow regresses its codegen. - # The post-pass follows structural masking; the producer's runtime - # route flag skips it only when all four current token words are full. - if cutlass.const_expr(cfg.tile_size_q == 128 and cfg.use_kv_valid_bits): - token_mask_is_required = not route_token_mask_is_full - if token_mask_is_required: - for word_idx in cutlass.range_constexpr(4): - token_word = local_token_words[word_idx] - for bit_idx in cutlass.range_constexpr(32): - reg_idx = word_idx * 32 + bit_idx - token_bit_is_valid = ( - (token_word >> Int32(bit_idx)) & Uint32(1) - ) != Uint32(0) - if not token_bit_is_valid: - s_vals[reg_idx] = _neg_max_f32() - - tile_max = self._reduce_keeps_row_max(s_vals) - self._publish_keeps_softmax_state( - s_vals, - tile_max, - old_max, - running_sum, - old_max_arr, - sum_arr, - new_max_arr, - s_arr, - ) - return old_max_arr, sum_arr, new_max_arr, s_arr - @cute.jit def _compute_softmax_loop_keeps( self, @@ -1407,8 +1102,26 @@ def _compute_softmax_loop_keeps( reduction, whose 16x256b register mapping is unrelated to Keeps. """ cfg = self.cfg + if cutlass.const_expr(cfg.streams_tmem_p_fragments): + # Streamed profiles share the fragment max pass with block-sparse + # routes; dense tiles describe their visible range as keep words. + 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, + 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.softmax_score_fragment_regs + num_s_regs = cfg.num_s_regs_per_thread old_max = new_max_arr[0] running_sum = sum_arr[0] s_vals = cutlass.Array(Float32, num_s_regs, space=cutlass.AddressSpace.rmem) @@ -1434,50 +1147,9 @@ def _compute_softmax_loop_keeps( is_valid_effective_tile, is_masked_final_wave, tile_is_unmasked, + tile_has_valid_scores, ) = self._resolve_keeps_tile_context(stage_info) - if cutlass.const_expr(cfg.tile_size_kv == 256): - # KV256 owns four physical K32 fragments per lane. Reduce the max - # one fragment at a time so only one native LDTM atom is live; the - # P pass reloads the same fragments after the reference max is - # known. - tile_max = _neg_max_f32() - for fragment_idx in cutlass.range_constexpr( - cfg.num_softmax_score_fragments - ): - self._load_keeps_fragment( - stage_info, - s_vals, - tile_offset_k, - element_mask_end_idx, - window_start_idx, - seq_len_kv, - logical_q_group_idx, - is_valid_effective_tile, - is_masked_final_wave, - tile_is_unmasked, - fragment_idx=fragment_idx, - ) - fragment_max = self._reduce_keeps_fragment_max(s_vals) - tile_max = cute.math.max(tile_max, fragment_max, ftz=True) - - new_max = cute.math.max(old_max, tile_max, ftz=True) - if old_max != _neg_max_f32(): - # Keeping the previous reference max avoids an in-place O - # rescale when the new tile raises it only modestly. The - # 16-bit P path can represent the bounded values above one; the - # numerator and denominator remain in the same scale frame. - # Large jumps still rebase to keep P comfortably in range. - max_delta_log2 = self.scale_softmax_log2 * (old_max - new_max) - if max_delta_log2 >= Float32(-KV_TILE_256_RESCALE_THRESHOLD_LOG2): - new_max = old_max - old_max_arr[0] = old_max - sum_arr[0] = running_sum - new_max_arr[0] = new_max - for reg_idx in cutlass.range_constexpr(num_s_regs): - s_arr[reg_idx] = s_vals[reg_idx] - return old_max_arr, sum_arr, new_max_arr, s_arr - if cutlass.const_expr(use_preload_mask_split): # Select the complete unmasked/masked TMEM load+max path before any S # registers are materialized. The shared predicate covers the @@ -1495,7 +1167,7 @@ def _compute_softmax_loop_keeps( is_masked_final_wave, tile_is_unmasked, ) - tile_max = self._reduce_keeps_fragment_max(s_vals) + tile_max = self._reduce_keeps_row_max(s_vals) self._publish_keeps_softmax_state( s_vals, @@ -1509,12 +1181,7 @@ def _compute_softmax_loop_keeps( ) return old_max_arr, sum_arr, new_max_arr, s_arr - should_load_s = ( - is_valid_effective_tile - and (tile_offset_k < seq_len_kv) - and not is_masked_final_wave - ) - if should_load_s: + if tile_has_valid_scores: base_addr = ( task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + Int32(self._alloc.offset) @@ -1600,73 +1267,6 @@ def _compute_softmax_loop_keeps( ) return old_max_arr, sum_arr, new_max_arr, s_arr - @consumer_work(returns=s_arr, work_attrs=WorkAttr.AUXILIARY) - @cute.jit - def load_softmax_p_fragment( - self, - stage_info: StageInfo, - *, - fragment_idx: Constexpr[int], - s_arr: cutlass.Array, - ) -> cutlass.Array: - """Reload and mask one KV256 K32 fragment for P materialization.""" - if cutlass.const_expr(self.cfg.use_block_sparse): - return self._load_block_sparse_softmax_p_fragment( - stage_info, - fragment_idx=fragment_idx, - s_arr=s_arr, - ) - ( - 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, - ) = self._resolve_keeps_tile_context(stage_info) - self._load_keeps_fragment( - stage_info, - s_arr, - tile_offset_k, - element_mask_end_idx, - window_start_idx, - seq_len_kv, - logical_q_group_idx, - is_valid_effective_tile, - is_masked_final_wave, - tile_is_unmasked, - fragment_idx=fragment_idx, - ) - return s_arr - - @cute.jit - def _sparse_swaps_logical_k( - self, - lane_k_offset: Int32, - sparse_origin0: Int32, - sparse_origin1: Int32, - sparse_origin2: Int32, - sparse_origin3: Int32, - *, - token_group_idx: Constexpr[int], - ) -> tuple[Int32, Int32]: - """Map one SWAP register group to its routed logical K position.""" - - atom_size = min(self.cfg.kv_block_size, 32) - groups_per_atom = atom_size // 8 - origin_idx = token_group_idx // groups_per_atom - atom_origin = sparse_origin0 - if cutlass.const_expr(origin_idx == 1): - atom_origin = sparse_origin1 - elif cutlass.const_expr(origin_idx == 2): - atom_origin = sparse_origin2 - elif cutlass.const_expr(origin_idx == 3): - atom_origin = sparse_origin3 - token_offset = (token_group_idx % groups_per_atom) * 8 - return atom_origin, atom_origin + Int32(token_offset) + lane_k_offset - @cute.jit def _compute_softmax_loop_swaps( self, @@ -1860,6 +1460,11 @@ def _compute_softmax_loop_swaps( s_vals[q_repeats * 4 + ld_base + 2] = loaded1[ld_base + 2] s_vals[q_repeats * 4 + ld_base + 3] = loaded1[ld_base + 3] + 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) + ) 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. @@ -1894,7 +1499,8 @@ def _compute_softmax_loop_swaps( lane_k_offset = Int32(task_cache[_TASK_CACHE_LANE_IDX]) >> Int32(2) token_word_covers_kv_tail = _swaps_token_word_covers_kv_tail(cfg) for token_group_idx in cutlass.range_constexpr(4): - atom_origin, logical_k = self._sparse_swaps_logical_k( + atom_origin, logical_k = _swaps_routed_coordinate( + cfg, lane_k_offset, sparse_origin0, sparse_origin1, @@ -1906,19 +1512,20 @@ def _compute_softmax_loop_swaps( # tail. Qualified profiles can therefore omit the local # atom-origin guard, independently of the K/V issuer warp. score_is_valid = cutlass.Boolean(True) - if cutlass.const_expr( - not _swaps_uses_token_only_score_validity(cfg) - ): - score_is_valid = cutlass.Boolean(atom_origin >= Int32(0)) - if cutlass.const_expr(not token_word_covers_kv_tail): - score_is_valid = cutlass.Boolean( - score_is_valid and logical_k < seq_len_kv - ) - if cutlass.const_expr(cfg.uses_uniform_causal_mask): - score_is_valid = cutlass.Boolean( - score_is_valid and logical_k < element_mask_end_idx - ) - if cutlass.const_expr(cfg.use_kv_valid_bits): + if not route_is_proxy: + if cutlass.const_expr( + not _swaps_uses_token_only_score_validity(cfg) + ): + score_is_valid = cutlass.Boolean(atom_origin >= Int32(0)) + if cutlass.const_expr(not token_word_covers_kv_tail): + score_is_valid = cutlass.Boolean( + score_is_valid and logical_k < seq_len_kv + ) + if cutlass.const_expr(cfg.uses_uniform_causal_mask): + score_is_valid = cutlass.Boolean( + score_is_valid and logical_k < element_mask_end_idx + ) + if cutlass.const_expr(cfg.uses_prepared_score_keep_words): token_bit_idx = Int32(token_group_idx * 8) + lane_k_offset token_is_valid = ( (sparse_token_word >> token_bit_idx) & Uint32(1) @@ -2061,7 +1668,8 @@ def _compute_softmax_loop_swaps( tile_offset_k + local_idx_k0 + Int32(token_group_idx * 8) ) if cutlass.const_expr(use_sparse): - _, token_idx = self._sparse_swaps_logical_k( + _, token_idx = _swaps_routed_coordinate( + cfg, lane_idx >> Int32(2), sparse_origin0, sparse_origin1, @@ -2351,7 +1959,7 @@ def reduce_sums( return sum_arr @cute.jit - def _compute_softmax_loop_sparse_keeps_kv256( + def _compute_softmax_loop_keeps_fragments( self, stage_info: StageInfo, *, @@ -2359,6 +1967,7 @@ def _compute_softmax_loop_sparse_keeps_kv256( sum_arr: cutlass.Array, new_max_arr: cutlass.Array, s_arr: cutlass.Array, + use_sparse: Constexpr[bool], sparse_origin0: Int32, sparse_origin1: Int32, sparse_route_flags: Int32, @@ -2367,77 +1976,141 @@ def _compute_softmax_loop_sparse_keeps_kv256( sparse_token_word2: Uint32, sparse_token_word3: Uint32, ) -> tuple[object, object, object, object]: - """Reduce one sparse KV256 route as four bounded K32 fragments. - - The full route path only loads and reduces scores. A partial route - predicates one native 32-score fragment at a time and writes it back - to TMEM, so the later P pass can replay masked scores without keeping - the logical 128-score tile live in registers. + """Mask streamed K32 score fragments in place and reduce their max. + + Every fragment gets one keep word. Block-sparse routes derive it from + 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. """ - cfg = self.cfg - assert cfg.tile_size_kv == 256 + assert cfg.streams_tmem_p_fragments + num_fragments = cfg.num_softmax_score_fragments + fragment_regs = cfg.softmax_score_fragment_regs + # The seven-slot softmax metadata ABI carries exactly four token words. + assert num_fragments == 4 and fragment_regs == 32 task_cache = _decode_gen_task_cache(stage_info) - token_words = ( - sparse_token_word0, - sparse_token_word1, - sparse_token_word2, - sparse_token_word3, + keep_words = cutlass.Array( + Uint32, num_fragments, space=cutlass.AddressSpace.rmem ) - keep_words = cutlass.Array(Uint32, 4, 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) - logical_q_group_idx = _logical_q_group_idx(cfg, stage_info, self.q_group_idx) - q_token_idx, _ = _q_row_token_and_local_head( - cfg, - self.h_r, - logical_q_group_idx, - tile_row_idx, - ) - q_row_is_valid = _q_row_is_valid_for_seq( - cfg, - self.h_r, - logical_q_group_idx, - tile_row_idx, - self.seq_len_q, - ) - seq_len_kv = _load_runtime_seq_len_kv( - self.seqlens_kv, - self.max_seq_len_kv, - stage_info, - Int32(0), - Int32(0), - ) - causal_end = seq_len_kv - self.seq_len_q + q_token_idx + Int32(1) - origin0 = Int32(sparse_origin0) - origin1 = Int32(sparse_origin1) - valid0 = sparse_route_flags & Int32(1) - valid1 = (sparse_route_flags >> Int32(1)) & Int32(1) - for fragment_idx in cutlass.range_constexpr(4): - fragment_origin = origin0 + Int32((fragment_idx % 2) * 32) - fragment_valid = valid0 - if cutlass.const_expr(fragment_idx >= 2): - fragment_origin = origin1 + Int32((fragment_idx % 2) * 32) - fragment_valid = valid1 - keep_words[fragment_idx] = _sparse_k32_effective_keep_word( - q_row_is_valid, - fragment_origin, - fragment_valid, - Uint32(token_words[fragment_idx]), - seq_len_kv, - causal_end, - apply_causal_mask=cfg.mask_type == CAUSAL, - apply_token_mask=cfg.use_kv_valid_bits, + if cutlass.const_expr(use_sparse): + token_words = ( + sparse_token_word0, + sparse_token_word1, + sparse_token_word2, + sparse_token_word3, + ) + logical_q_group_idx = _logical_q_group_idx( + cfg, stage_info, self.q_group_idx + ) + q_token_idx, _ = _q_row_token_and_local_head( + cfg, + self.h_r, + logical_q_group_idx, + tile_row_idx, + ) + q_row_is_valid = _q_row_is_valid_for_seq( + cfg, + self.h_r, + logical_q_group_idx, + tile_row_idx, + self.seq_len_q, + ) + seq_len_kv = _load_runtime_seq_len_kv( + self.seqlens_kv, + self.max_seq_len_kv, + stage_info, + Int32(0), + Int32(0), ) + causal_end = seq_len_kv - self.seq_len_q + q_token_idx + Int32(1) + origin0 = Int32(sparse_origin0) + origin1 = Int32(sparse_origin1) + valid0 = sparse_route_flags & Int32(1) + valid1 = (sparse_route_flags >> Int32(1)) & Int32(1) + fragments_per_origin = cfg.softmax_fragments_per_route_atom + for fragment_idx in cutlass.range_constexpr(num_fragments): + atom_offset = Int32( + (fragment_idx % fragments_per_origin) * fragment_regs + ) + fragment_origin = origin0 + atom_offset + fragment_valid = valid0 + if cutlass.const_expr(fragment_idx >= fragments_per_origin): + fragment_origin = origin1 + atom_offset + fragment_valid = valid1 + if cutlass.const_expr(cfg.trusts_prepared_score_words): + prepared_keep_word = Uint32(0) + if q_row_is_valid: + prepared_keep_word = Uint32(token_words[fragment_idx]) + keep_words[fragment_idx] = prepared_keep_word + else: + keep_words[fragment_idx] = _sparse_effective_keep_word( + q_row_is_valid, + fragment_origin, + fragment_valid, + Uint32(token_words[fragment_idx]), + seq_len_kv, + causal_end, + apply_causal_mask=cfg.mask_type == CAUSAL, + apply_token_mask=cfg.uses_prepared_score_keep_words, + ) - warp_scores_are_unmasked = cutlass.Boolean(True) - for fragment_idx in cutlass.range_constexpr(4): - warp_scores_are_unmasked = cutlass.Boolean( - warp_scores_are_unmasked - and keep_words[fragment_idx] == Uint32(0xFFFFFFFF) + else: + ( + 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, + rows_are_active, + ) = self._resolve_keeps_tile_context(stage_info) + if cutlass.const_expr(cfg.q_score_rows_need_mask): + rows_are_active = cutlass.Boolean( + rows_are_active + and _q_row_is_valid_for_seq( + cfg, + self.h_r, + logical_q_group_idx, + tile_row_idx, + self.seq_len_q, + ) + ) + visible_start = Int32(0) + visible_end = element_mask_end_idx + if cutlass.const_expr(cfg.uses_per_row_causal_mask): + q_token_idx, _ = _q_row_token_and_local_head( + cfg, + self.h_r, + logical_q_group_idx, + tile_row_idx, + ) + visible_end = seq_len_kv - self.seq_len_q + q_token_idx + Int32(1) + visible_start = _sliding_window_start_idx( + cfg, seq_len_kv, self.seq_len_q, q_token_idx + ) + elif cutlass.const_expr(cfg.use_sliding_window_causal): + visible_start = window_start_idx + # A tile that is unmasked for the whole Q group has all-ones keep + # words on every active row, so only masked tiles build them. + warp_scores_are_unmasked = cute.arch.vote_all_sync( + cutlass.Boolean(tile_is_unmasked and rows_are_active) ) - # The load/store branch must be uniform for each participating warp. - warp_scores_are_unmasked = cute.arch.vote_all_sync(warp_scores_are_unmasked) + if cutlass.const_expr(use_sparse): + warp_scores_are_unmasked = cutlass.Boolean(True) + for fragment_idx in cutlass.range_constexpr(num_fragments): + warp_scores_are_unmasked = cutlass.Boolean( + warp_scores_are_unmasked + and keep_words[fragment_idx] == Uint32(0xFFFFFFFF) + ) + # The load/store branch must be uniform for each participating warp. + warp_scores_are_unmasked = cute.arch.vote_all_sync(warp_scores_are_unmasked) score_tmem_addr = ( task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] @@ -2449,17 +2122,17 @@ def _compute_softmax_loop_sparse_keeps_kv256( max_chains[chain_idx] = _neg_max_f32() if warp_scores_are_unmasked: - for fragment_idx in cutlass.range_constexpr(4): + for fragment_idx in cutlass.range_constexpr(num_fragments): loaded = _keeps_tcgen05_ld( cfg, prims.make_tmem_ptr( - score_tmem_addr + Int32(fragment_idx * 32), Float32 + score_tmem_addr + Int32(fragment_idx * fragment_regs), Float32 ), - num=32, + num=fragment_regs, offset=cfg.tile_size_kv // 2, ) prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) - for score_idx in cutlass.range_constexpr(32): + 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], @@ -2467,19 +2140,35 @@ def _compute_softmax_loop_sparse_keeps_kv256( ftz=True, ) else: - for fragment_idx in cutlass.range_constexpr(4): - fragment_addr = score_tmem_addr + Int32(fragment_idx * 32) + if cutlass.const_expr(not use_sparse): + lane_idx = Int32(task_cache[_TASK_CACHE_LANE_IDX]) + col_base = _keeps_col_base(cfg, lane_idx, num_fragments * fragment_regs) + for fragment_idx in cutlass.range_constexpr(num_fragments): + fragment_token_base = tile_offset_k + _keeps_score_col( + cfg, + warp_group_thread_idx, + fragment_idx * fragment_regs, + col_base, + ) + keep_words[fragment_idx] = _dense_fragment_keep_word( + rows_are_active, + visible_start - fragment_token_base, + 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=32, + num=fragment_regs, offset=cfg.tile_size_kv // 2, ) prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) masked_scores = cutlass.Array( - Float32, 32, space=cutlass.AddressSpace.rmem + Float32, fragment_regs, space=cutlass.AddressSpace.rmem ) - for score_idx in cutlass.range_constexpr(32): + 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) @@ -2494,7 +2183,7 @@ def _compute_softmax_loop_sparse_keeps_kv256( _keeps_tcgen05_st( cfg, prims.make_tmem_ptr(fragment_addr, Float32), - masked_scores.data_ptr().load(count=32, alignment=4), + masked_scores.data_ptr().load(count=fragment_regs, alignment=4), offset=cfg.tile_size_kv // 2, ) prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) @@ -2506,44 +2195,11 @@ def _compute_softmax_loop_sparse_keeps_kv256( ftz=True, ) old_max = new_max_arr[0] - new_max = cute.math.max(old_max, tile_max, ftz=True) - if old_max != _neg_max_f32(): - max_delta_log2 = self.scale_softmax_log2 * (old_max - new_max) - if max_delta_log2 >= Float32(-KV_TILE_256_RESCALE_THRESHOLD_LOG2): - new_max = old_max + 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 _load_block_sparse_softmax_p_fragment( - self, - stage_info: StageInfo, - *, - fragment_idx: Constexpr[int], - s_arr: cutlass.Array, - ) -> cutlass.Array: - """Reload one full or already-predicated sparse KV256 fragment for P.""" - - assert self.cfg.tile_size_kv == 256 - task_cache = _decode_gen_task_cache(stage_info) - score_tmem_addr = ( - task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] - + Int32(self._alloc.offset) - + self._softmax_loop_stage_slot_offset(stage_info) - + Int32(fragment_idx * 32) - ) - loaded = _keeps_tcgen05_ld( - self.cfg, - prims.make_tmem_ptr(score_tmem_addr, Float32), - num=32, - offset=self.cfg.tile_size_kv // 2, - ) - prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) - for score_idx in cutlass.range_constexpr(32): - s_arr[score_idx] = loaded[score_idx] - return s_arr - @consumer_work(returns=("old_max_arr", "sum_arr", "new_max_arr", "s_arr")) @cute.jit def compute_block_sparse_softmax_loop( @@ -2566,34 +2222,21 @@ def compute_block_sparse_softmax_loop( assert self.cfg.use_block_sparse if cutlass.const_expr(self.cfg.use_keeps_mma_ab): - if cutlass.const_expr(self.cfg.tile_size_kv == 256): - return self._compute_softmax_loop_sparse_keeps_kv256( - stage_info, - 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, - ) - return self._compute_softmax_loop_sparse_keeps( + # Every block-sparse Keeps profile streams K32 fragments. + 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, - routed_origin0=sparse_origin0, - routed_origin1=sparse_origin1, - routed_route_flags=sparse_route_flags, - routed_token_word0=sparse_token_word0, - routed_token_word1=sparse_token_word1, - routed_token_word2=sparse_token_word2, - routed_token_word3=sparse_token_word3, + 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, ) # SWAP reuses the Keeps seven-slot task ABI: all four origins remain # logical KV atom bases, but origin2 occupies the flags slot and 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 4321a0d0d46e..b75de72ba008 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 @@ -32,16 +32,13 @@ import cutlass import cutlass.cute as cute -from cutlass import Int32 from cutlass.experimental import primitives as prims from cutlass.experimental.task_scheduling.memory import ( ResourceContext, - SmemAllocation, ) from cutlass.experimental.task_scheduling.resources import ( MemoryResource, StageInfo, - TaskLocalVariable, WorkQueue, consumer_work, producer_work, @@ -61,10 +58,9 @@ KV_INST1, KV_KIND_K, KV_KIND_V, - KV_TILE_256_SHARED_FIFO_STAGES, ) from .fmha_decode_resources.helpers_common import ( - ResourceVars, + _assume_nonnegative_i32, _q_group_token_base, _q_seq_bounds, _warp_broadcast_i32, @@ -122,14 +118,18 @@ def restore_slots(*resource_proxies: object) -> None: return traced -def _block_sparse_route_loop_domain( - route_count: cutlass.Int32, +def _loop_domain_after_head( + total_kv_tiles: cutlass.Int32, *, num_insts_kv: int, ) -> cutlass.Int32: - """Return LOOP iterations after HEAD reserves one candidate per instance.""" + """Return LOOP iterations after HEAD reserves one KV tile per instance. - remaining = route_count - cutlass.Int32(num_insts_kv) + Dense tiles and block-sparse routes share this recurrence; only the tile + count's source differs. + """ + + remaining = total_kv_tiles - cutlass.Int32(num_insts_kv) remaining = cute.math.max(remaining, cutlass.Int32(0)) insts = cutlass.Int32(num_insts_kv) return (remaining + insts - cutlass.Int32(1)) // insts @@ -152,102 +152,6 @@ def consume_schedule_token(self, stage_info: StageInfo) -> None: del stage_info -@dataclass(kw_only=True) -class SmemKvReuseCreditResource(MemoryResource): - """One-slot credit carrying the rotating KV256 exchange-stage index. - - Load publishes which drained 64-KiB physical K/V stage Correction may use - as tail scratch. The one-stage pipeline couples that payload to the same - ownership epoch: the following Load may use the other two physical stages, - but cannot publish a new alias until Correction releases this credit after - all output work completes. - """ - - cfg: cutlass.Constexpr[FmhaDecodeConfig] = None - _alloc: cutlass.Constexpr[SmemAllocation | None] = None - scratch_stage_slot: cutlass.Constexpr[TaskLocalVariable] = ( - TaskLocalVariable.uninitialized() - ) - - def __post_init__(self) -> None: - """Create the routed consumer slot for one physical K/V stage.""" - assert KV_TILE_256_SHARED_FIFO_STAGES == 3, ( - "KV256 reuse-credit rotation requires exactly three shared FIFO stages" - ) - if not self.cfg.uses_rotating_kv256_exchange: - raise ValueError( - "rotating KV scratch requires persistent direct Q64/KV256 " - "with two KV instructions, one head-dimension stage, and " - "one load warp" - ) - self.scratch_stage_slot = TaskLocalVariable( - dtype=Int32, - default=Int32(0), - docs="Physical shared-K/V stage reserved for KV256 tail exchange.", - ) - - def get_smem_requirements(self) -> list[SmemAllocation]: - """Allocate the one-word stage payload guarded by this pipeline.""" - if self._alloc is None: - self._alloc = SmemAllocation( - name=f"{self.name}_scratchStage", - size_bytes=4, - alignment=4, - ) - return [self._alloc] - - @cute.jit - def _payload(self, stage_info: StageInfo) -> cutlass.Array: - """Return the natural next-stage cursor owned by this credit.""" - return cutlass.Array( - stage_info.context.smem_base.data_ptr() + self._alloc.offset, - dtype=Int32, - shape=(1,), - addrspace=3, - ) - - @cute.jit - def create_function_variables( - self, - context: ResourceContext | None = None, - ) -> ResourceVars: - """Initialize the persistent ring cursor before TS tasks start.""" - if cutlass.const_expr(context is not None and context.smem_base is not None): - payload = cutlass.Array( - context.smem_base.data_ptr() + self._alloc.offset, - dtype=Int32, - shape=(1,), - addrspace=3, - ) - thread_idx, _, _ = cute.arch.thread_idx() - if thread_idx == Int32(0): - payload[0] = Int32(0) - return {} - - @producer_work - @cute.jit - def publish_scratch_stage(self, stage_info: StageInfo) -> None: - """Advance the persistent ring cursor and publish the drained stage.""" - num_stages = Int32(KV_TILE_256_SHARED_FIFO_STAGES) - if prims.elect_sync(): - payload = self._payload(stage_info) - # Each work commits T = 4 * (loop_end + 1) K/V transactions. - # Since 4 == 1 (mod 3), the cursor advances by loop_end + 1. - # loop_end is the resolved per-work domain, so heterogeneous - # runtime sequence lengths do not inherit a captured host bound. - payload[0] = ( - Int32(payload[0]) + stage_info.loop_end + Int32(1) - ) % num_stages - - @consumer_work(returns=scratch_stage_slot) - @cute.jit - def read_scratch_stage(self, stage_info: StageInfo) -> Int32: - """Read the alias only after the matching credit wait completes.""" - num_stages = Int32(KV_TILE_256_SHARED_FIFO_STAGES) - next_stage = Int32(self._payload(stage_info)[0]) - return (next_stage + num_stages - Int32(1)) % num_stages - - @dataclass(kw_only=True) class PackedDecodeWorkQueue(WorkQueue): """CLC work queue that drops packed-Q tiles beyond a batch's Q length.""" @@ -414,28 +318,23 @@ def _produce_staged_page_offsets( def _consume_staged_qk_mma( smem_kv: MemoryResource, tmem_s: MemoryResource, - aliased_p: MemoryResource, q_desc: Any, k_desc_label: str, qk_mma_label: str, section: FmhaStage, cfg: FmhaDecodeConfig, ) -> None: - """Consume all K head-dim stages for one QK MMA wave.""" + """Consume all K head-dim stages for one QK MMA wave. + + Streamed KV256 aliases P with the S columns this QK overwrites. The + preceding same-instance PV reads P as its TMEM A operand from the same + issuing thread, and the tensor core interlocks that read against a later + MMA's accumulator write, so no completion wait is needed before QK. + """ tmem_s.acquire() for head_dim_stage_idx in range(cfg.num_head_dim_stages_kv): smem_kv.wait() kv_desc = getattr(smem_kv, k_desc_label)() - if cutlass.const_expr( - cfg.streams_tmem_p_fragments - and head_dim_stage_idx == 0 - and (section == FmhaStage.Loop or cfg.use_persistent_scheduler) - ): - # Wait as late as possible: K staging overlaps the previous PV, - # but QK cannot overwrite the matching S/P alias until PV is done. - # Static HEAD has no previous tile; persistent HEAD may follow the - # same CTA's tail from another logical work tile and must wait. - aliased_p.wait_until_reusable_before_qk() if cfg.uses_q_desc_ref: getattr(tmem_s, f"{qk_mma_label}_from_q_ref")( kv_desc=kv_desc, @@ -451,6 +350,44 @@ def _consume_staged_qk_mma( tmem_s.commit() +def _consume_streamed_pv_fragments( + smem_kv: MemoryResource, + tmem_p: MemoryResource, + tmem_o: MemoryResource, + v_desc_label: str, + vp_mma_label: str, + cfg: FmhaDecodeConfig, +) -> None: + """Issue one PV wave as its K32 P fragments become ready. + + P fragment 0 is the earliest dependency: wait for it and for the + correction credit before holding the V stage. Later fragments may become + ready while the previous PV fragment is already executing; every slot + stays live through the complete async UMMA wave so the producer cannot + overwrite an operand prematurely. + """ + assert cfg.num_head_dim_stages_kv == 1 + fragment_label = f"{vp_mma_label}_fragment" + p_tmem_addr = tmem_p.wait_p_fragment(fragment_idx=0) + tmem_o.acquire() + smem_kv.wait() + v_desc = getattr(smem_kv, v_desc_label)() + getattr(tmem_o, fragment_label)( + v_desc=v_desc, + p_tmem_addr=p_tmem_addr, + fragment_idx=0, + ) + for fragment_idx in range(1, cfg.num_softmax_score_fragments): + p_tmem_addr = tmem_p.wait_p_fragment(fragment_idx=fragment_idx) + getattr(tmem_o, fragment_label)( + v_desc=v_desc, + p_tmem_addr=p_tmem_addr, + fragment_idx=fragment_idx, + ) + smem_kv.release() + tmem_o.commit() + + def _consume_staged_pv_mma( smem_kv: MemoryResource, tmem_p: MemoryResource, @@ -464,33 +401,9 @@ def _consume_staged_pv_mma( """Consume all V head-dim stages for one PV MMA wave.""" _ = section if cutlass.const_expr(cfg.streams_tmem_p_fragments): - assert cfg.num_head_dim_stages_kv == 1 - fragment_label = f"{vp_mma_label}_fragment" - - # P fragment 0 is the earliest dependency. Wait for it and for the - # correction credit before holding the shared V FIFO stage. - p_tmem_addr = tmem_p.wait_p_fragment(fragment_idx=0) - tmem_o.acquire() - smem_kv.wait() - v_desc = getattr(smem_kv, v_desc_label)() - getattr(tmem_o, fragment_label)( - v_desc=v_desc, - p_tmem_addr=p_tmem_addr, - fragment_idx=0, - ) - - # Later P fragments may become ready while the previous PV fragment is - # already executing. Keep every slot live through the complete async - # UMMA wave so the producer cannot overwrite an operand prematurely. - for fragment_idx in range(1, cfg.num_softmax_score_fragments): - p_tmem_addr = tmem_p.wait_p_fragment(fragment_idx=fragment_idx) - getattr(tmem_o, fragment_label)( - v_desc=v_desc, - p_tmem_addr=p_tmem_addr, - fragment_idx=fragment_idx, - ) - smem_kv.release() - tmem_o.commit() + _consume_streamed_pv_fragments( + smem_kv, tmem_p, tmem_o, v_desc_label, vp_mma_label, cfg + ) return tmem_p.wait() @@ -634,6 +547,47 @@ def _decode_work_tile_schedule_with_invariant_bridge( _work_queue_tail(work_queue) +@cute.jit +def _prepared_sparse_row_address( + cfg: cutlass.Constexpr[FmhaDecodeConfig], + q_group_idx: cutlass.Int32, + h_idx: cutlass.Int32, + b_idx: cutlass.Int32, + num_heads_kv: cutlass.Int32, +) -> cutlass.Int32: + """Map a (q_group, head, batch) tile to its prepared row header index.""" + + 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 + return (b_idx * num_heads_kv + h_idx) * cutlass.Int32(num_q_blocks) + q_block + + +@cute.jit +def _prefetch_prepared_sparse_row( + cfg: cutlass.Constexpr[FmhaDecodeConfig], + row_route_offsets: cute.Pointer, + row_route_counts: cute.Pointer, + q_group_idx: cutlass.Int32, + h_idx: cutlass.Int32, + b_idx: cutlass.Int32, + num_heads_kv: cutlass.Int32, +) -> tuple[cutlass.Int32, cutlass.Int32]: + """Load one static tile's prepared row header before its tasks start. + + Every thread loads the same two words, so each warp issues one request + and the global-memory latency overlaps the CTA prologue (TMEM allocation + and barrier setup) instead of stalling every task at its first step. + """ + + row_address = _prepared_sparse_row_address( + cfg, q_group_idx, h_idx, b_idx, num_heads_kv + ) + row_route_begin = cutlass.Int32(row_route_offsets[row_address]) + route_count = _assume_nonnegative_i32(cutlass.Int32(row_route_counts[row_address])) + return row_route_begin, route_count + + @cute.jit def _load_prepared_sparse_row_warp( row_route_offsets: cute.Pointer, @@ -653,7 +607,7 @@ def _load_prepared_sparse_row_warp( loaded_row_route_begin = cutlass.Int32(row_route_offsets[row_address]) loaded_route_count = cutlass.Int32(row_route_counts[row_address]) row_route_begin = _warp_broadcast_i32(loaded_row_route_begin, 0) - route_count = _warp_broadcast_i32(loaded_route_count, 0) + route_count = _assume_nonnegative_i32(_warp_broadcast_i32(loaded_route_count, 0)) return row_route_begin, route_count @@ -667,6 +621,9 @@ def __init__(self, **kwargs: TaskKwarg) -> None: self.block_table_capacity = kwargs.pop("block_table_capacity", None) self.sparse_row_route_offsets = kwargs.pop("sparse_row_route_offsets", None) self.sparse_row_route_counts = kwargs.pop("sparse_row_route_counts", None) + # Static tiles may pass the already loaded row header instead. + self.sparse_row_route_begin = kwargs.pop("sparse_row_route_begin", None) + self.sparse_route_count = kwargs.pop("sparse_route_count", None) self.num_heads_kv = kwargs.pop("num_heads_kv", None) self.max_seq_len_kv = kwargs.pop("max_seq_len_kv", cutlass.Int32(0)) self.seq_len_q = kwargs.pop("seq_len_q", None) @@ -923,20 +880,20 @@ def get_domain(self, tile_coord: cute.Coord) -> cutlass.Int32 | int: q_group_idx = cutlass.Int32(tile_coord[0]) h_idx = cutlass.Int32(tile_coord[1]) b_idx = cutlass.Int32(tile_coord[2]) - q_token_base = _q_group_token_base(self.cfg, q_group_idx) - - q_block = q_token_base // self.cfg.q_block_size - num_q_blocks = ( - self.cfg.max_seq_len_q + self.cfg.q_block_size - 1 - ) // self.cfg.q_block_size - row_address = (b_idx * self.num_heads_kv + h_idx) * num_q_blocks + q_block - - row_route_begin, route_count = _load_prepared_sparse_row_warp( - row_route_offsets, - row_route_counts, - cutlass.Int32(row_address), - self._lane_idx, - ) + if self.sparse_row_route_begin is not None: + # The static kernel prologue already loaded this tile's header. + row_route_begin = self.sparse_row_route_begin + route_count = self.sparse_route_count + else: + row_address = _prepared_sparse_row_address( + self.cfg, q_group_idx, h_idx, b_idx, self.num_heads_kv + ) + row_route_begin, route_count = _load_prepared_sparse_row_warp( + row_route_offsets, + row_route_counts, + row_address, + self._lane_idx, + ) # Sparse route-span accessors share two underlying cache words # with paged KV. Clear dense/paged-only coordinates on every @@ -951,7 +908,7 @@ def get_domain(self, tile_coord: cute.Coord) -> cutlass.Int32 | int: self._kv_valid_tile_end = route_count self._kv_window_start = cutlass.Int32(0) - loop_domain = _block_sparse_route_loop_domain( + loop_domain = _loop_domain_after_head( route_count, num_insts_kv=self.cfg.num_insts_kv, ) @@ -997,14 +954,10 @@ def get_domain(self, tile_coord: cute.Coord) -> cutlass.Int32 | int: self._kv_window_start = cutlass.Int32(0) self._kv_valid_tile_end = total_kv_tiles self._kv_raw_tile_base = cutlass.Int32(0) - remaining_kv_tiles = cute.math.max( - total_kv_tiles - cutlass.Int32(self.cfg.num_insts_kv), - cutlass.Int32(0), + loop_domain = _loop_domain_after_head( + total_kv_tiles, + num_insts_kv=self.cfg.num_insts_kv, ) - num_insts_kv = cutlass.Int32(self.cfg.num_insts_kv) - loop_domain = ( - remaining_kv_tiles + num_insts_kv - cutlass.Int32(1) - ) // num_insts_kv return loop_domain + cutlass.Int32(self.domain_bias) # Decode the logical Q tile with the configured physical split fanout, @@ -1058,13 +1011,10 @@ def get_domain(self, tile_coord: cute.Coord) -> cutlass.Int32 | int: self._kv_raw_tile_base = skipped_tiles + split_idx * total_kv_tiles else: self._kv_raw_tile_base = skipped_tiles - remaining_kv_tiles = cute.math.max( - total_kv_tiles - cutlass.Int32(self.cfg.num_insts_kv), cutlass.Int32(0) + loop_domain = _loop_domain_after_head( + total_kv_tiles, + num_insts_kv=self.cfg.num_insts_kv, ) - num_insts_kv = cutlass.Int32(self.cfg.num_insts_kv) - loop_domain = ( - remaining_kv_tiles + num_insts_kv - cutlass.Int32(1) - ) // num_insts_kv # All tasks share the MMA-loop domain; tail-only tasks add a bias. return loop_domain + cutlass.Int32(self.domain_bias) @@ -1081,29 +1031,63 @@ def get_domain(self, tile_coord: cute.Coord) -> cutlass.Int32 | int: def _resolve_and_store_sparse_route( sparse_kv_metadata: MemoryResource | None, section: FmhaStage, -) -> tuple[Any, Any, Any, Any] | None: - """Resolve one prepared route and retain it for the matching K/V pair.""" + prefetch: tuple[Any, Any] | None = None, + *, + pipeline: bool = True, +) -> tuple[tuple[Any, Any, Any, Any] | None, tuple[Any, Any] | None]: + """Resolve one prepared route and retain it for the matching K/V pair. + + Returns ``(route, prefetch)``. With ``pipeline`` the record load is issued + one resolution ahead: HEAD loads its own record immediately, and every + resolution issues the load for the next one (LOOP iteration 0 from HEAD, + iteration i + 1 from iteration i) before the caller's K TMA burst, so the + global-memory latency overlaps that issue instead of stalling the load + warp. Callers pass the returned ``prefetch`` back into the next resolution + of the same instance, the way ``_staged_kv_load`` threads its cached page + IDs. Without ``pipeline`` the record is loaded where it is resolved and no + state is returned; the split-ring load variants use this because the + pipelined form measured slower for them. Dense profiles pass ``None`` and + get ``(None, None)``. + """ if sparse_kv_metadata is None: - return None + return None, None + if not pipeline: + prefetch = sparse_kv_metadata.prefetch_route( + target="head" if section == FmhaStage.Head else "current_loop" + ) + elif section == FmhaStage.Head: + prefetch = sparse_kv_metadata.prefetch_route(target="head") + assert prefetch is not None + prefetched_record_word, prefetched_record_offset = prefetch ( - resolved_origin0, + resolved_record_word, resolved_origin1, resolved_atom_validity, route_record_word_offset, - ) = sparse_kv_metadata.resolve_route(section=section) + ) = sparse_kv_metadata.resolve_route( + section=section, + prefetched_record_word_slot=prefetched_record_word, + prefetched_record_offset_slot=prefetched_record_offset, + ) sparse_kv_metadata.store_route( - resolved_origin0=resolved_origin0, + resolved_record_word=resolved_record_word, resolved_origin1=resolved_origin1, resolved_atom_validity=resolved_atom_validity, route_record_word_offset=route_record_word_offset, ) - return ( - resolved_origin0, + next_prefetch = None + if pipeline: + next_prefetch = sparse_kv_metadata.prefetch_route( + target="first_loop" if section == FmhaStage.Head else "next_loop" + ) + route = ( + resolved_record_word, resolved_origin1, resolved_atom_validity, route_record_word_offset, ) + return route, next_prefetch def _publish_sparse_softmax_route( @@ -1116,14 +1100,14 @@ def _publish_sparse_softmax_route( return assert route is not None ( - resolved_origin0, + resolved_record_word, resolved_origin1, resolved_atom_validity, route_record_word_offset, ) = route sparse_softmax_metadata.acquire() sparse_softmax_metadata.store_route( - resolved_origin0=resolved_origin0, + resolved_record_word=resolved_record_word, resolved_origin1=resolved_origin1, resolved_atom_validity=resolved_atom_validity, route_record_word_offset=route_record_word_offset, @@ -1136,7 +1120,6 @@ def create_load_task( smem_kv: MemoryResource, work_queue: WorkQueue | None, schedule_token_throttle: MemoryResource | None, - smem_kv_reuse_credit: MemoryResource | None, cfg: FmhaDecodeConfig, *, domain: int | cutlass.Int32, @@ -1158,7 +1141,6 @@ def load_schedule_body( smem_kv: MemoryResource, smem_page_offsets: MemoryResource | None, schedule_token_throttle: MemoryResource | None, - smem_kv_reuse_credit: MemoryResource | None, sparse_kv_metadata0: MemoryResource | None = None, sparse_kv_metadata1: MemoryResource | None = None, sparse_softmax_metadata0: MemoryResource | None = None, @@ -1208,26 +1190,22 @@ def _kv_load(label: str, section: FmhaStage) -> None: smem_page_offsets.wait() else: _page_offsets_consume(smem_page_offsets) - if sparse_kv_metadata0 is None: - for label in ("load_k0", "load_k1"): - _kv_load(label, FmhaStage.Head) - else: - route0 = _resolve_and_store_sparse_route( - sparse_kv_metadata0, FmhaStage.Head - ) - _kv_load("load_k0", FmhaStage.Head) - route1 = _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) - if smem_kv_reuse_credit is not None: - # K0/K1 occupy the two stages disjoint from the previous work's - # scratch. Acquire only before issuing the third K/V transaction. - smem_kv_reuse_credit.acquire() + # 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} # 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 @@ -1239,41 +1217,30 @@ def _kv_load(label: str, section: FmhaStage) -> None: else ("load_k0", "load_v0", "load_k1", "load_v1") ) with domain_loop(0, domain, 1, unroll=1): - if sparse_kv_metadata0 is None: - for label in loop_labels: - _kv_load(label, FmhaStage.Loop) - else: - # Follow the dense KV256 stage order exactly. Each V consumes - # its retained route before the matching K label replaces it. - loop_routes = [] - for label in loop_labels: - route = None - sparse_softmax_metadata = None - if label == "load_k0": - route = _resolve_and_store_sparse_route( - sparse_kv_metadata0, FmhaStage.Loop - ) - sparse_softmax_metadata = sparse_softmax_metadata0 - elif label == "load_k1": - route = _resolve_and_store_sparse_route( - sparse_kv_metadata1, FmhaStage.Loop - ) - sparse_softmax_metadata = sparse_softmax_metadata1 - _kv_load(label, FmhaStage.Loop) - if route is not None: - loop_routes.append((sparse_softmax_metadata, route)) - for sparse_softmax_metadata, route in loop_routes: - _publish_sparse_softmax_route(sparse_softmax_metadata, route) + # 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( + label, (None, None) + ) + route, prefetch_by_label[label] = _resolve_and_store_sparse_route( + kv_metadata, FmhaStage.Loop, prefetch_by_label.get(label) + ) + _kv_load(label, FmhaStage.Loop) + if route is not None: + loop_routes.append((softmax_metadata, route)) + for sparse_softmax_metadata, route in loop_routes: + _publish_sparse_softmax_route(sparse_softmax_metadata, route) # 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"): _kv_load(label, FmhaStage.Tail) - if smem_kv_reuse_credit is not None: - # Publish the physical stage drained by this work together with - # the ownership token consumed by the correction tail. - smem_kv_reuse_credit.publish_scratch_stage() - smem_kv_reuse_credit.commit() if hold_page_window: _page_offsets_release(smem_page_offsets) @@ -1288,7 +1255,6 @@ def load_schedule( sparse_softmax_metadata1: MemoryResource | None, work_queue: WorkQueue | None, schedule_token_throttle: MemoryResource | None, - smem_kv_reuse_credit: MemoryResource | None, ) -> None: """Schedule shared-KV loads with only the resources in this profile.""" @@ -1300,7 +1266,6 @@ def load_schedule( smem_kv, smem_page_offsets, schedule_token_throttle, - smem_kv_reuse_credit, sparse_kv_metadata0, sparse_kv_metadata1, sparse_softmax_metadata0, @@ -1334,7 +1299,6 @@ def load_schedule( sparse_softmax_metadata1, work_queue, schedule_token_throttle, - smem_kv_reuse_credit, ) src = [] for sparse_kv_metadata in (sparse_kv_metadata0, sparse_kv_metadata1): @@ -1350,8 +1314,6 @@ def load_schedule( dst.append(sparse_resource) if schedule_token_throttle is not None: dst.append(schedule_token_throttle) - if smem_kv_reuse_credit is not None: - dst.append(smem_kv_reuse_credit) return task_class( src_resources=src, dst_resources=dst, @@ -1720,7 +1682,9 @@ def load_tile( ) in active_instances: if smem_k is None: continue - route = _resolve_and_store_sparse_route(sparse_kv_metadata, FmhaStage.Head) + route, _ = _resolve_and_store_sparse_route( + sparse_kv_metadata, FmhaStage.Head, pipeline=False + ) load_tile(smem_k, load_k, smem_page_offsets_k, FmhaStage.Head) head_routes.append((sparse_softmax_metadata, route)) # In the combined task, preserve both K issues ahead of Softmax @@ -1748,8 +1712,8 @@ def load_tile( smem_page_offsets_v_local, FmhaStage.Loop, ) - route = _resolve_and_store_sparse_route( - sparse_kv_metadata, FmhaStage.Loop + route, _ = _resolve_and_store_sparse_route( + sparse_kv_metadata, FmhaStage.Loop, pipeline=False ) load_tile(smem_k, load_k, smem_page_offsets_k, FmhaStage.Loop) loop_routes.append((sparse_softmax_metadata, route)) @@ -2175,26 +2139,9 @@ def pv_mma( section: FmhaStage, ) -> None: """Issue one scheduled PV wave using the selected phase work.""" - _ = section - tmem_p.wait() - p_desc_0, p_desc_1, p_tmem_addr_0, p_tmem_addr_1 = tmem_p.p_operands() - tmem_o.acquire() - for head_dim_stage_idx in range(cfg.num_head_dim_stages_kv): - smem_kv.wait() - v_desc = smem_kv.v_desc() - getattr(tmem_o, vp_mma_label)( - v_desc_0=v_desc, - v_desc_1=v_desc, - p_desc_0=p_desc_0, - p_desc_1=p_desc_1, - p_tmem_addr_0=p_tmem_addr_0, - p_tmem_addr_1=p_tmem_addr_1, - inst_idx=inst_idx, - head_dim_stage_idx=head_dim_stage_idx, - ) - smem_kv.release() - tmem_o.commit() - tmem_p.release() + _consume_staged_pv_mma( + smem_kv, tmem_p, tmem_o, "v_desc", vp_mma_label, inst_idx, section, cfg + ) qk_mma( smem_k0, @@ -2652,7 +2599,6 @@ def mma_schedule_body( _consume_staged_qk_mma( smem_kv, tmem_s0, - smem_p0, q_desc, "k_desc_0", "qk_mma_head", @@ -2662,7 +2608,6 @@ def mma_schedule_body( _consume_staged_qk_mma( smem_kv, tmem_s1, - smem_p1, q_desc, "k_desc_1", "qk_mma_head", @@ -2671,8 +2616,9 @@ def mma_schedule_body( ) # LOOP: consume aliased TMEM P before the next same-instance QK - # overwrites its S columns. SMEM-P profiles retain their established - # K-before-V cadence because P no longer depends on S lifetime. + # overwrites its S columns. Full-SMEM P remains score-dependent during + # Softmax replay, but owns independent storage once the replay commits + # and releases S; that completed handoff enables the QK-before-PV cadence. with domain_loop(0, domain, 1, unroll=1): if cfg.uses_two_inst_tmem_p: _consume_staged_pv_mma( @@ -2688,7 +2634,6 @@ def mma_schedule_body( _consume_staged_qk_mma( smem_kv, tmem_s0, - smem_p0, q_desc, "k_desc_0", "qk_mma_loop", @@ -2720,7 +2665,6 @@ def mma_schedule_body( _consume_staged_qk_mma( smem_kv, tmem_s1, - smem_p1, q_desc, "k_desc_1", "qk_mma_loop", @@ -2739,6 +2683,12 @@ def mma_schedule_body( 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 + # the next tile's Q load overlaps the final softmax and BMM2 waves + # instead of waiting for them. + smem_q.release() + # TAIL: no future K tiles remain, so only the final two BMM2 waves run. _consume_staged_pv_mma( smem_kv, @@ -2760,8 +2710,6 @@ def mma_schedule_body( FmhaStage.Tail, cfg, ) - # Q is live for every BMM1 call and can be released only after the loop. - smem_q.release() def mma_schedule_prelude( smem_q: MemoryResource, @@ -2884,9 +2832,9 @@ def softmax0_schedule_body( sparse_softmax_metadata.init_read_state() with domain_loop(0, domain, 1, unroll=1) as d: - # ConsWait/ConsWork: load S from TMEM and compute the tile max. - tmem_s0.wait() 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. @@ -2900,6 +2848,9 @@ def softmax0_schedule_body( 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, @@ -2937,18 +2888,17 @@ def softmax0_schedule_body( ) tmem_softmax_local0.commit() if cutlass.const_expr(cfg.streams_tmem_p_fragments): - # Publish one K32 probability fragment at a time so PV can - # consume early fragments while later scores are processed. - for fragment_idx in range(cfg.num_softmax_score_fragments): - s_arr = tmem_s0.load_softmax_p_fragment( - fragment_idx=fragment_idx, - s_arr=s_arr, - ) - smem_p0.compute_p_fragment( - fragment_idx=fragment_idx, + # 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, - s_arr=s_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 @@ -2959,16 +2909,27 @@ def softmax0_schedule_body( # 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. - smem_p0.compute_p( - new_max_arr=new_max_arr, - s_arr=s_arr, - ) # publishes the local denominator through tmem_s0 + 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): - # The TMEM-P store has consumed the aliased S columns, so the - # next QK wave can now overwrite them. + # 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. @@ -3106,9 +3067,9 @@ def softmax1_schedule_body( sparse_softmax_metadata.init_read_state() with domain_loop(0, domain, 1, unroll=1) as d: - # ConsWait/ConsWork: load the second S instance and compute max. - tmem_s1.wait() 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. @@ -3122,6 +3083,9 @@ def softmax1_schedule_body( 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, @@ -3157,16 +3121,15 @@ def softmax1_schedule_body( ) tmem_softmax_local1.commit() if cutlass.const_expr(cfg.streams_tmem_p_fragments): - for fragment_idx in range(cfg.num_softmax_score_fragments): - s_arr = tmem_s1.load_softmax_p_fragment( - fragment_idx=fragment_idx, - s_arr=s_arr, - ) - smem_p1.compute_p_fragment( - fragment_idx=fragment_idx, + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): + smem_p1.compute_proxy_route_p_fragments( new_max_arr=new_max_arr, - s_arr=s_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 @@ -3175,7 +3138,18 @@ def softmax1_schedule_body( if tmem_softmax_order is not None: tmem_softmax_order.wait_softmax1() # ProdWork: compute and publish P1 for BMM2. - smem_p1.compute_p(new_max_arr=new_max_arr, s_arr=s_arr) + 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() @@ -3290,7 +3264,6 @@ def create_correction_task( tmem_corr0: MemoryResource, tmem_corr1: MemoryResource, work_queue: WorkQueue | None, - smem_kv_reuse_credit: MemoryResource | None, cfg: FmhaDecodeConfig, *, domain: int | cutlass.Int32, @@ -3301,9 +3274,6 @@ def create_correction_task( ) -> Task: """Create the two-instance correction and output task.""" - if smem_kv_reuse_credit is not None and work_queue is None: - raise ValueError("KV reuse credit requires a work queue") - def correction_schedule_body( tmem_softmax_local0: MemoryResource, tmem_softmax_local1: MemoryResource, @@ -3312,7 +3282,6 @@ def correction_schedule_body( tmem_corr1: MemoryResource, tmem_stats_done0: MemoryResource | None, tmem_stats_done1: MemoryResource | None, - smem_kv_reuse_credit: MemoryResource | None, ) -> None: """Schedule two-instance O correction and final output normalization.""" @@ -3506,37 +3475,17 @@ def correct_o( tail_o_stage_idx_1=tail_1, inst_idx=KV_INST1, ) - if smem_kv_reuse_credit is None: - tmem_corr1.correction_tail_epilogue( - o_stage_idx=o_stage_idx, - tail_o_stage_idx_0=tail_0, - tail_o_stage_idx_1=tail_1, - old_max_arr=old_max_arr, - new_max_arr=new_max_arr, - inst0_new_max_arr=inst0_new_max_arr, - inst0_sum_arr=inst0_sum_arr, - inst1_new_max_arr=inst1_new_max_arr, - inst1_sum_arr=inst1_sum_arr, - ) - else: - # The stage selector and ownership token share one pipeline epoch. - # Wait before the first aliased access and release immediately - # after correction stops touching the selected KV-ring stage. - smem_kv_reuse_credit.wait() - scratch_stage = smem_kv_reuse_credit.read_scratch_stage() - tmem_corr1.correction_tail_epilogue_rotating_exchange( - scratch_stage=scratch_stage, - o_stage_idx=o_stage_idx, - tail_o_stage_idx_0=tail_0, - tail_o_stage_idx_1=tail_1, - old_max_arr=old_max_arr, - new_max_arr=new_max_arr, - inst0_new_max_arr=inst0_new_max_arr, - inst0_sum_arr=inst0_sum_arr, - inst1_new_max_arr=inst1_new_max_arr, - inst1_sum_arr=inst1_sum_arr, - ) - smem_kv_reuse_credit.release() + tmem_corr1.correction_tail_epilogue( + o_stage_idx=o_stage_idx, + tail_o_stage_idx_0=tail_0, + tail_o_stage_idx_1=tail_1, + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + inst0_new_max_arr=inst0_new_max_arr, + inst0_sum_arr=inst0_sum_arr, + inst1_new_max_arr=inst1_new_max_arr, + inst1_sum_arr=inst1_sum_arr, + ) # Inst1 final reduction consumes both O0 and O1, so defer O0 release # until after inst1 has finished reading it. tmem_o.release() @@ -3550,7 +3499,6 @@ def run_correction_schedule( tmem_corr1: MemoryResource, tmem_stats_done0: MemoryResource | None, tmem_stats_done1: MemoryResource | None, - smem_kv_reuse_credit: MemoryResource | None, work_queue: WorkQueue | None, ) -> None: """Wrap correction with optional stats lifetime gates.""" @@ -3565,7 +3513,6 @@ def run_correction_schedule( tmem_corr1, tmem_stats_done0, tmem_stats_done1, - smem_kv_reuse_credit, ), ) @@ -3577,7 +3524,6 @@ def correction_schedule( tmem_corr0: MemoryResource, tmem_corr1: MemoryResource, work_queue: WorkQueue | None = None, - smem_kv_reuse_credit: MemoryResource | None = None, ) -> None: """Capture the Swaps correction schedule.""" run_correction_schedule( @@ -3588,7 +3534,6 @@ def correction_schedule( tmem_corr1, None, None, - smem_kv_reuse_credit, work_queue, ) @@ -3602,7 +3547,6 @@ def correction_keeps_schedule( tmem_stats_done0: MemoryResource, tmem_stats_done1: MemoryResource, work_queue: WorkQueue | None = None, - smem_kv_reuse_credit: MemoryResource | None = None, ) -> None: """Capture Keeps correction with explicit stats lifetime gates.""" run_correction_schedule( @@ -3613,7 +3557,6 @@ def correction_keeps_schedule( tmem_corr1, tmem_stats_done0, tmem_stats_done1, - smem_kv_reuse_credit, work_queue, ) @@ -3626,15 +3569,6 @@ def correction_keeps_schedule( tmem_corr0, tmem_corr1, ) - elif smem_kv_reuse_credit is None: - captured_schedule = correction_schedule( - tmem_softmax_local0, - tmem_softmax_local1, - tmem_o, - tmem_corr0, - tmem_corr1, - work_queue, - ) else: captured_schedule = correction_schedule( tmem_softmax_local0, @@ -3643,7 +3577,6 @@ def correction_keeps_schedule( tmem_corr0, tmem_corr1, work_queue, - smem_kv_reuse_credit, ) src = [tmem_softmax_local0, tmem_softmax_local1, tmem_o] else: @@ -3657,17 +3590,6 @@ def correction_keeps_schedule( tmem_stats_done0, tmem_stats_done1, ) - elif smem_kv_reuse_credit 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, - work_queue, - ) else: captured_schedule = correction_keeps_schedule( tmem_softmax_local0, @@ -3678,7 +3600,6 @@ def correction_keeps_schedule( tmem_stats_done0, tmem_stats_done1, work_queue, - smem_kv_reuse_credit, ) src = [ tmem_softmax_local0, @@ -3689,8 +3610,6 @@ def correction_keeps_schedule( ] if work_queue is not None: src.append(work_queue) - if smem_kv_reuse_credit is not None: - src.append(smem_kv_reuse_credit) return task_class( src_resources=src, dst_resources=[tmem_corr0, tmem_corr1], From 8fdd1b325c3ab36dbced6f1f99967d67ec1d2ede Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Mon, 7 Sep 2026 04:12:25 +0000 Subject: [PATCH 2/7] feat: add generic PrimTS block-sparse support Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- .../attention/ATTENTION_DEVELOPER_GUIDE.md | 31 +- .../attention/backends/fmha/cute_dsl_mla.py | 2 + .../attention/backends/fmha/fallback.py | 6 +- .../backends/fmha/flashinfer_sparse_mla.py | 15 +- .../backends/fmha/flashinfer_trtllm_gen.py | 2 + .../_torch/attention/backends/fmha/manager.py | 3 + .../attention/backends/fmha/msa_sparse_gqa.py | 15 +- .../attention/backends/fmha/prims_ts.py | 91 ++- .../backends/fmha/prims_ts_block_sparse.py | 709 ++++++++++++++++++ .../attention/backends/fmha/registry.py | 2 + .../backends/fmha/triton_custom_mask.py | 2 + .../_torch/attention/backends/interface.py | 5 +- .../attention/backends/sparse/params.py | 53 +- .../_torch/attention/backends/trtllm.py | 8 +- .../_torch/attention/test_fmha_manager.py | 42 ++ .../_torch/attention/test_fmha_registry.py | 20 + .../attention/test_prims_ts_block_sparse.py | 690 +++++++++++++++++ 17 files changed, 1674 insertions(+), 22 deletions(-) create mode 100644 tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py create mode 100644 tests/unittest/_torch/attention/test_prims_ts_block_sparse.py diff --git a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index e56bbb5fa24a..e15d77d2634c 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md @@ -211,6 +211,20 @@ that file for the current config/backend combinations. Consult the for the supported attention shapes; do not infer support from algorithm registration alone. +Block-sparse FMHA is a kernel-library contract rather than a sparse algorithm. +Algorithms lower their live routing state to an algorithm-neutral +`BlockSparseForwardInputs`: block geometry plus either canonical BSR routes or +an exact packed bitmask. Optional K/V summaries enable proxy routes, and +optional token-validity bits mask ragged KV tails. Plans contain only static +format, proxy, geometry, and capacity choices; every run receives the live +routes, summaries, validity bits, page tables, and sequence lengths. + +`PrimsTSBlockSparseFmha` owns its wrapper-plan cache by default. Integrations +whose attention layers execute serially may explicitly bind a model-scoped +cache to reuse graph-stable route workspaces across compatible layers. The +cache must not be shared by concurrent forwards; each independent model +component must own separate state. + ### 2.3 Backend contract All backends implement the `AttentionBackend` interface. @@ -343,13 +357,18 @@ starting with an empty selection cache. `TrtllmAttention` prepares the complete per-forward state, passes itself to the manager for selection, and then executes the selected library. -`TLLM_FMHA_LIBS` controls the ordered selection. PrimTS is opt-in because it may -add host overhead; use `TLLM_FMHA_LIBS=+prims_ts` to add it to the defaults or -`TLLM_FMHA_LIBS=fallback` to force the fallback path. Delta entries update the +`TLLM_FMHA_LIBS` controls the ordered selection. Dense PrimTS is opt-in because +it may add host overhead; use `TLLM_FMHA_LIBS=+prims_ts` to add it to the +defaults or `TLLM_FMHA_LIBS=fallback` to force the fallback path. Generic +block-sparse PrimTS remains enabled by default because a dense fallback cannot +preserve its routing semantics. Delta entries update the default membership and follow canonical registry order, while an exact list preserves the user-specified order. Each FMHA library exposes `is_available()` for module/static environment checks and `is_supported()` for per-forward -request checks. For mixed non-MLA batches, the manager checks each active phase +request checks. `AttentionForwardArgs.block_sparse_inputs` is exclusive: the +block-sparse implementation validates and consumes it, while every other +library rejects it rather than silently dropping sparse routing semantics. +For mixed non-MLA batches, the manager checks each active phase independently with `is_supported(..., phase=...)`; a phased library accepts only phases backed by its corresponding `run_*()` entry point. @@ -367,6 +386,10 @@ The FMHA package is split by role: `TrtllmAttention` can pair it with a later causal-generation provider through `CombinedFmha`. - `fmha/cute_dsl_mla.py` implements the CuTe DSL MLA decode FMHA library. +- `fmha/prims_ts_block_sparse.py` adapts generic block-sparse requests to the + vendored PrimTS contiguous and paged wrappers. Paged generation passes a + live, zero-copy 2D K-page-table view with its TRT-LLM padded row stride; it + does not stage page tables through CSR metadata. - `fmha/prims_ts.py` adapts TRT-LLM inputs and paged-cache metadata to the vendored PrimTS kernels. Before changing the managed source under `backends/prims_ts`, read the diff --git a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py index 9584d131e245..224f301424e2 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py @@ -285,6 +285,8 @@ def _is_supported_with_reason( meta: "TrtllmAttentionMetadata", fwd: AttentionForwardArgs, ) -> tuple[bool, str]: + if fwd.block_sparse_inputs is not None: + return False, "block_sparse_inputs are not supported." if fwd.attention_input_type != AttentionInputType.generation_only: return False, "CuTe DSL MLA FMHA only supports generation-only attention." # Disable mixed context/generation batches until the CuTe DSL host diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py index f03ae9a01dc1..80825703f612 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py @@ -40,6 +40,7 @@ _THOP_EXCLUDED_FIELDS: frozenset = frozenset( { "sparse_backend_args", # consumed by sparse prediction before the attention op + "block_sparse_inputs", # consumed by the selected FMHA library "attention_mask_data", # custom-mask code path "out_scale_sf", # promoted into ``out_scale`` in ``TrtllmAttention.forward`` for NVFP4 path "skip_mla_rope_generation", # handled in ``TrtllmAttention.forward`` for the test-only MLA path @@ -81,8 +82,9 @@ def is_supported( del k, v, phase if q is not None and q.dtype == torch.float8_e4m3fn: return False - return forward_args.attention_mask != CustomAttentionMask.CUSTOM and ( - forward_args.update_kv_cache or metadata.is_cross + return forward_args.block_sparse_inputs is None and ( + forward_args.attention_mask != CustomAttentionMask.CUSTOM + and (forward_args.update_kv_cache or metadata.is_cross) ) def forward( diff --git a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_sparse_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_sparse_mla.py index 2924cc440fc0..80606dfd0006 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_sparse_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_sparse_mla.py @@ -13,7 +13,7 @@ from tensorrt_llm._utils import get_sm_version from tensorrt_llm.logger import logger -from .interface import Fmha +from .interface import Fmha, FmhaPhase if TYPE_CHECKING: from tensorrt_llm._torch.attention.backends.trtllm import ( @@ -70,6 +70,19 @@ def is_available(cls, attn: "TrtllmAttention") -> bool: return False return is_flashinfer_sparse_mla_enabled(getattr(attn.sparse_params, "algorithm", None)) + def is_supported( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + *, + phase: Optional[FmhaPhase] = None, + ) -> bool: + del q, k, v, metadata, phase + return forward_args.block_sparse_inputs is None + def forward( self, q: torch.Tensor, diff --git a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py index dfc5d2e4af45..28aee806faf7 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py @@ -571,6 +571,8 @@ def _is_supported_with_reason( *, phase: Optional[FmhaPhase] = None, ) -> Tuple[bool, str]: + if fwd.block_sparse_inputs is not None: + return False, "block_sparse_inputs are not supported." is_mla_enable = attn.is_mla_enable if phase is None: has_context_phase = fwd.attention_input_type != AttentionInputType.generation_only diff --git a/tensorrt_llm/_torch/attention/backends/fmha/manager.py b/tensorrt_llm/_torch/attention/backends/fmha/manager.py index 07837c39eadd..c18747844cce 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/manager.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/manager.py @@ -134,6 +134,7 @@ class _FmhaCacheKey(NamedTuple): generation_seq_len_q: int attention_mask_type: AttentionMaskType use_spec_decoding: bool + has_block_sparse_inputs: bool # LoRA can change the effective output from packed NVFP4 to unpacked BF16 # without changing the request shape. Keep those selection regimes apart. output_dtype: torch.dtype | None @@ -366,12 +367,14 @@ def _make_cache_key( generation_seq_len_q, _FMHA_CACHE_SEQ_LEN_Q_GRID ) + block_sparse_inputs = forward_args.block_sparse_inputs return _FmhaCacheKey( context_batch_size=context_batch_size, generation_batch_size=generation_batch_size, generation_seq_len_q=generation_seq_len_q, attention_mask_type=attention_mask_type, use_spec_decoding=metadata.use_spec_decoding, + has_block_sparse_inputs=block_sparse_inputs is not None, output_dtype=output_dtype, output_sf_dtype=output_sf_dtype, ) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/msa_sparse_gqa.py b/tensorrt_llm/_torch/attention/backends/fmha/msa_sparse_gqa.py index ee85aaf02e0d..c0b37b651ace 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/msa_sparse_gqa.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/msa_sparse_gqa.py @@ -17,7 +17,7 @@ from tensorrt_llm._utils import is_sm_100f -from .interface import Fmha +from .interface import Fmha, FmhaPhase if TYPE_CHECKING: from tensorrt_llm._torch.attention.backends.interface import AttentionForwardArgs @@ -210,6 +210,19 @@ def is_available(cls, attn: Optional["TrtllmAttention"] = None) -> bool: # alone, so no custom library discovery is needed. return attn.sparse_params is not None and attn.sparse_params.algorithm == "minimax_m3" + def is_supported( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: "AttentionForwardArgs", + *, + phase: Optional[FmhaPhase] = None, + ) -> bool: + del q, k, v, metadata, phase + return forward_args.block_sparse_inputs is None + def forward( self, q: torch.Tensor, diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py index afad6479067d..ab48753bfd8e 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py @@ -20,7 +20,7 @@ import math from importlib import import_module from importlib.metadata import PackageNotFoundError, version -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Any, Optional import torch from packaging.version import InvalidVersion, Version @@ -184,6 +184,8 @@ def _is_supported_with_reason( # dispatch. Accept the phased dispatcher keyword, but do not narrow # support until that preparation is phase-aware too. del phase + if fwd.block_sparse_inputs is not None: + return False, "block_sparse_inputs are not supported." if q.device.type != "cuda": return False, "CUDA tensors are required." if not q.is_contiguous(): @@ -432,6 +434,27 @@ def _get_fixed_block_tables( ) return block_tables[:batch_size, 0, :] + def _get_generation_workspace_layout( + self, + dtype: torch.dtype, + num_requests: int, + num_tokens: int, + ) -> dict[str, int]: + """Return the shared TRT-LLM generation preprocessing layout.""" + + return thop.get_trtllm_gen_generation_workspace_layout( + dtype, + num_requests, + num_tokens, + self.attn.num_heads, + self.attn.head_dim, + self.attn.rope_dim, + self.attn.num_kv_heads, + 0, + False, + skip_fmha_workspace=True, + ) + @staticmethod def _get_sequence_lengths( sequence_lengths: torch.Tensor, @@ -652,17 +675,10 @@ def prepare_workspace( if input_type == AttentionInputType.generation_only else q.shape[0] - int(metadata.num_ctx_tokens) ) - generation_layout = thop.get_trtllm_gen_generation_workspace_layout( + generation_layout = self._get_generation_workspace_layout( q.dtype, int(metadata.num_generations), num_gen_tokens_for_layout, - self.attn.num_heads, - self.attn.head_dim, - self.attn.rope_dim, - self.attn.num_kv_heads, - 0, - False, - skip_fmha_workspace=True, ) required_preprocess_bytes = max( required_preprocess_bytes, int(generation_layout["total_size"]) @@ -956,6 +972,63 @@ def run_context(self, params: FmhaParams) -> None: skip_fmha_workspace=True, ) + def _run_generation_preprocess(self, params: FmhaParams) -> tuple[Any, ...]: + """Run the shared TRT-LLM generation QKV and cache preprocessing.""" + + if self._multi_processor_count is None: + raise RuntimeError("PrimTS generation workspace was not prepared.") + attn = params.attn + meta = params.meta + fwd = params.fwd + rope_params = attn.rope_params + attention_chunk_size = attn.attention_chunk_size or 0 + return thop.trtllm_gen_generation_preprocess( + params.qkv_input, + params.workspace, + params.sequence_lengths, + params.spec_decoding_generation_lengths, + params.spec_decoding_position_offsets, + meta.kv_cache_block_offsets, + meta.host_kv_cache_pool_pointers, + meta.host_kv_cache_pool_mapping, + fwd.kv_scale_orig_quant, + fwd.kv_scale_quant_orig, + fwd.out_scale, + attn.rotary_inv_freq, + attn.rotary_cos_sin, + fwd.mrope_position_deltas, + attn.local_layer_idx, + params.seq_offset, + attn.num_heads, + attn.num_kv_heads, + attn.head_dim, + params.tokens_per_block, + attn.quant_mode, + params.max_attention_window_size, + params.cyclic_attention_window_size, + params.num_tokens, + params.batch_size, + params.input_seq_length, + params.max_past_kv_length, + rope_params.dim, + rope_params.theta, + int(rope_params.scale_type), + rope_params.scale, + rope_params.max_positions, + attn.position_embedding_type, + self._get_bmm1_scale(attn), + 1.0, + False, + attn.predicted_tokens_per_seq, + attention_chunk_size, + self._multi_processor_count, + params.total_num_blocks, + params.kv_factor, + True, + False, + skip_fmha_workspace=True, + ) + def run_generation(self, params: FmhaParams) -> None: if params.qkv_input is None or params.context_buf is None: raise RuntimeError("PrimTS decode requires QKV input and an output buffer.") diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py new file mode 100644 index 000000000000..5b99fdbcb931 --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py @@ -0,0 +1,709 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from dataclasses import dataclass +from typing import TYPE_CHECKING, Literal, cast + +import torch + +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + AttentionInputType, + PredefinedAttentionMask, +) +from tensorrt_llm._torch.attention.backends.prims_ts._block_sparse.config import ( + _validate_block_sparse_static_profile, +) +from tensorrt_llm._torch.attention.backends.sparse.params import BlockSparseForwardInputs +from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from tensorrt_llm.logger import logger +from tensorrt_llm.quantization.mode import QuantMode + +from .interface import FmhaPhase +from .phased import FmhaParams +from .prims_ts import PrimsTSFmha +from .utils import get_kv_page_offset + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention.backends.prims_ts import ( + BlockSparsePagedTSWrapper, + BlockSparseTSWrapper, + ) + from tensorrt_llm._torch.attention.backends.trtllm import ( + TrtllmAttention, + TrtllmAttentionMetadata, + ) + + +@dataclass(frozen=True, slots=True) +class _BlockSparsePlanKey: + """Static wrapper profile shared by compatible attention layers.""" + + device: torch.device + batch_size: int + seq_len_q: int + kv_capacity: int + num_heads: int + num_kv_heads: int + head_dim: int + page_size: int | None + q_block_size: int + kv_block_size: int + max_blocks_per_row: int + mask_type: Literal["dense", "causal"] + dtype: torch.dtype + use_kv_valid_bits: bool + sparse_format: Literal["bsr", "bitmask"] + use_proxy_routes: bool + + +try: + from tensorrt_llm._torch.attention.backends.prims_ts import ( + BlockSparsePagedTSWrapper as _BlockSparsePagedTSWrapper, + ) + from tensorrt_llm._torch.attention.backends.prims_ts import ( + BlockSparseTSWrapper as _BlockSparseTSWrapper, + ) +except (ImportError, OSError): + _BlockSparseTSWrapper = _BlockSparsePagedTSWrapper = None + + +class PrimsTSBlockSparseFmha(PrimsTSFmha): + def __init__(self, attn: "TrtllmAttention") -> None: + super().__init__(attn) + self.bind_plan_cache({}) + + def bind_plan_cache(self, cache_state: dict[str, object]) -> None: + """Bind wrapper plans to an explicitly owned cache.""" + + self._contiguous_wrappers = cast( + dict[_BlockSparsePlanKey, "BlockSparseTSWrapper"], + cache_state.setdefault("contiguous_wrappers", {}), + ) + self._paged_wrappers = cast( + dict[_BlockSparsePlanKey, "BlockSparsePagedTSWrapper"], + cache_state.setdefault("paged_wrappers", {}), + ) + + @classmethod + def is_available(cls, attn: "TrtllmAttention") -> bool: + return ( + _BlockSparseTSWrapper is not None + and _BlockSparsePagedTSWrapper is not None + and super().is_available(attn) + ) + + def is_supported( + self, + q: torch.Tensor, + k: torch.Tensor | None, + v: torch.Tensor | None, + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + *, + phase: FmhaPhase | None = None, + ) -> bool: + supported, reason = self._is_supported_with_reason( + q, + k, + v, + metadata, + forward_args, + phase=phase, + ) + if not supported: + logger.debug(f"PrimTS block-sparse FMHA does not support request: {reason}") + return supported + + def _is_supported_with_reason( + self, + q: torch.Tensor, + k: torch.Tensor | None, + v: torch.Tensor | None, + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + *, + phase: FmhaPhase | None = None, + ) -> tuple[bool, str]: + if getattr(metadata, "kv_cache_manager", None) is None: + reason = ( + "contiguous block-sparse attention only supports the context phase" + if phase not in (None, FmhaPhase.CONTEXT) + else self._contiguous_unsupported_reason(q, k, v, metadata, forward_args) + ) + else: + reason = ( + "paged block-sparse attention only supports the generation phase" + if phase not in (None, FmhaPhase.GENERATION) + else self._paged_unsupported_reason(q, metadata, forward_args) + ) + return reason is None, reason or "" + + def _make_plan_key( + self, + q: torch.Tensor, + inputs: BlockSparseForwardInputs, + *, + batch_size: int, + seq_len_q: int, + kv_capacity: int, + page_size: int | None, + mask_type: Literal["dense", "causal"], + ) -> _BlockSparsePlanKey: + max_blocks_per_row = inputs.max_blocks_per_row + if max_blocks_per_row is None: + max_blocks_per_row = math.ceil(kv_capacity / inputs.kv_block_size) + return _BlockSparsePlanKey( + device=q.device, + batch_size=batch_size, + seq_len_q=seq_len_q, + kv_capacity=kv_capacity, + num_heads=self.attn.num_heads, + num_kv_heads=self.attn.num_kv_heads, + head_dim=self.attn.head_dim, + page_size=page_size, + q_block_size=inputs.q_block_size, + kv_block_size=inputs.kv_block_size, + max_blocks_per_row=max_blocks_per_row, + mask_type=mask_type, + dtype=q.dtype, + use_kv_valid_bits=inputs.kv_valid_bits is not None, + sparse_format=inputs.sparse_format, + use_proxy_routes=inputs.use_proxy_routes, + ) + + def _get_or_plan_wrapper( + self, + key: _BlockSparsePlanKey, + *, + paged: bool, + ) -> "BlockSparseTSWrapper | BlockSparsePagedTSWrapper": + cache = self._paged_wrappers if paged else self._contiguous_wrappers + wrapper = cache.get(key) + if wrapper is not None: + return wrapper + wrapper_type = _BlockSparsePagedTSWrapper if paged else _BlockSparseTSWrapper + assert wrapper_type is not None + wrapper = wrapper_type() + plan_args = ( + key.batch_size, + key.seq_len_q, + key.kv_capacity, + key.num_heads, + key.num_kv_heads, + key.head_dim, + key.q_block_size, + key.kv_block_size, + ) + if paged: + plan_args += (key.page_size,) + plan_kwargs = { + "device": key.device, + "max_blocks_per_row": key.max_blocks_per_row, + "use_kv_valid_bits": key.use_kv_valid_bits, + "mask_type": key.mask_type, + "q_data_type": key.dtype, + "kv_data_type": key.dtype, + "o_data_type": key.dtype, + } + if not paged: + plan_kwargs.update( + sparse_format=key.sparse_format, + use_proxy_routes=key.use_proxy_routes, + ) + wrapper.plan(*plan_args, **plan_kwargs) + cache[key] = wrapper + return wrapper + + def _paged_storage_unsupported_reason( + self, + metadata: "TrtllmAttentionMetadata", + ) -> str | None: + if metadata.kv_layout != "HND": + return "only HND KV-cache layout is supported" + if metadata.host_kv_cache_pool_pointers is None: + return "KV-cache pool pointers are required" + pool_mapping = metadata.host_kv_cache_pool_mapping + if pool_mapping is None: + return "KV-cache pool mapping is required" + + manager = metadata.kv_cache_manager + if isinstance(manager, KVCacheManagerV2): + if manager.enable_swa_scratch_reuse: + return "KVCacheManagerV2 SWA scratch reuse is not supported" + elif isinstance(manager, KVCacheManager): + if manager.num_pools != 1: + return "KVCacheManagerV1 with multiple memory pools is not supported" + local_layer_idx = self.attn.local_layer_idx + if ( + pool_mapping.ndim != 2 + or pool_mapping.shape[1] < 2 + or local_layer_idx is None + or not 0 <= local_layer_idx < pool_mapping.shape[0] + ): + return "KVCacheManagerV1 has an invalid layer-to-pool mapping" + pool_index = int(pool_mapping[local_layer_idx, 0]) + layer_index_in_pool = int(pool_mapping[local_layer_idx, 1]) + if pool_index != 0 or not 0 <= layer_index_in_pool < manager.num_local_layers: + return "KVCacheManagerV1 has an invalid layer-to-pool mapping" + else: + return f"unsupported KV cache manager {type(manager).__name__}" + if metadata.tokens_per_block not in self.SUPPORTED_PAGE_SIZES: + return f"page size {metadata.tokens_per_block} is unsupported" + try: + if ( + get_kv_page_offset( + self.attn, + metadata, + 0, + cache=self._kv_page_offset_cache, + ) + is None + ): + return "the K-to-V page displacement could not be resolved" + except (AttributeError, IndexError, RuntimeError, TypeError, ValueError) as error: + return f"invalid KV-cache storage metadata: {error}" + return None + + @staticmethod + def _has_legacy_sparse_prediction(forward_args: AttentionForwardArgs) -> bool: + prediction = forward_args.sparse_runtime_params + return any( + ( + prediction.sparse_kv_indices is not None, + prediction.sparse_kv_offsets is not None, + prediction.sparse_attn_indices is not None, + prediction.sparse_attn_offsets is not None, + prediction.sparse_attn_kv_lens is not None, + bool(prediction.sparse_attn_indices_block_size), + prediction.aux_kv_cache_pool_ptr is not None, + bool(prediction.threshold_scale_factor_prefill), + bool(prediction.threshold_scale_factor_decode), + ) + ) + + def _common_unsupported_reason( + self, + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + ) -> str | None: + if forward_args.block_sparse_inputs is None: + return "block-sparse forward inputs are required" + if metadata.is_cross: + return "cross attention is not supported" + if self.attn.is_mla_enable: + return "MLA is not supported" + if getattr(metadata, "helix_position_offsets", None) is not None: + return "Helix parallelism is not supported" + if getattr(metadata, "num_sparse_topk", 0) > 0 or self._has_legacy_sparse_prediction( + forward_args + ): + return "legacy sparse attention cannot be combined with block-sparse inputs" + if forward_args.enable_dsv4_epilogue_fusion: + return "DSv4 epilogue fusion is not supported" + if forward_args.sage_attn_qk_int8 or any( + getattr(forward_args, name) > 0 + for name in ( + "sage_attn_num_elts_per_blk_q", + "sage_attn_num_elts_per_blk_k", + "sage_attn_num_elts_per_blk_v", + ) + ): + return "SageAttention is not supported" + if forward_args.softmax_stats_tensor is not None: + return "softmax statistics output is not supported" + if ( + forward_args.output_sf is not None + or forward_args.out_scale is not None + or forward_args.out_scale_sf is not None + ): + return "quantized output is not supported" + if ( + forward_args.attention_mask_data is not None + or forward_args.relative_attention_bias is not None + or forward_args.attention_sinks is not None + ): + return "custom attention masks, bias, and sinks are not supported" + if forward_args.attention_mask not in ( + PredefinedAttentionMask.FULL, + PredefinedAttentionMask.CAUSAL, + ): + return "only full and causal masks are supported" + return None + + def _static_profile_unsupported_reason( + self, + q: torch.Tensor, + inputs: BlockSparseForwardInputs, + *, + batch_size: int, + seq_len_q: int, + seq_len_kv: int, + page_size: int | None, + mask_type: Literal["dense", "causal"], + ) -> str | None: + max_blocks_per_row = inputs.max_blocks_per_row + if max_blocks_per_row is None: + max_blocks_per_row = math.ceil(seq_len_kv / inputs.kv_block_size) + try: + _validate_block_sparse_static_profile( + batch_size=batch_size, + seq_len_q=seq_len_q, + seq_len_kv=seq_len_kv, + num_qo_heads=self.attn.num_heads, + num_kv_heads=self.attn.num_kv_heads, + head_dim=self.attn.head_dim, + q_block_size=inputs.q_block_size, + kv_block_size=inputs.kv_block_size, + use_kv_valid_bits=inputs.kv_valid_bits is not None, + mask_type=mask_type, + q_dtype=q.dtype, + kv_dtype=q.dtype, + output_dtype=q.dtype, + max_blocks_per_row=max_blocks_per_row, + page_size=page_size, + ) + except (ValueError, NotImplementedError, OverflowError) as error: + return str(error) + return None + + def _paged_unsupported_reason( + self, + q: torch.Tensor, + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + ) -> str | None: + inputs = forward_args.block_sparse_inputs + if inputs is None: + return "block-sparse forward inputs are required" + if inputs.sparse_format != "bsr" or inputs.use_proxy_routes: + return "paged block-sparse attention only supports BSR exact routes" + common_reason = self._common_unsupported_reason(metadata, forward_args) + if common_reason is not None: + return common_reason + if not forward_args.is_fused_qkv: + return "paged block-sparse attention requires fused QKV input" + if ( + forward_args.attention_input_type != AttentionInputType.generation_only + or metadata.num_contexts != 0 + ): + return "only generation-only paged requests are supported" + batch_size = int(metadata.num_generations) + if batch_size <= 0 or q.shape[0] % batch_size: + return "query tokens must be uniformly divisible across generation requests" + seq_len_q = int(q.shape[0]) // batch_size + if not bool(metadata.seq_lens[:batch_size].eq(seq_len_q).all()): + return "query lengths must be batch-uniform and match the fixed query shape" + paged_kv_reason = self._paged_storage_unsupported_reason(metadata) + if paged_kv_reason is not None: + return paged_kv_reason + block_tables = metadata.kv_cache_block_offsets + if block_tables is None or block_tables.shape[1] < batch_size: + return "paged KV-cache block offsets must cover the generation batch" + if metadata.beam_width != 1: + return "beam search is not supported" + if any( + getattr(metadata, name, False) + for name in ( + "is_spec_decoding_enabled", + "use_spec_decoding", + "is_spec_dec_tree", + "is_spec_dec_dynamic_tree", + ) + ): + return "speculative decoding is not supported" + if self.attn.attention_chunk_size: + return "chunked attention is not supported" + if self.attn.position_embedding_type in (4, 5, 6, 7, 10): + return f"position embedding type {self.attn.position_embedding_type} is not supported" + page_size = int(metadata.tokens_per_block) + max_seq_len_kv = int(block_tables.shape[-1]) * page_size + logical_max_seq_len = int(metadata.max_seq_len) + if logical_max_seq_len > max_seq_len_kv: + return "logical maximum sequence length must fit the page-table capacity" + attention_window_size = forward_args.attention_window_size + if ( + attention_window_size is None + or attention_window_size < logical_max_seq_len + or attention_window_size > max_seq_len_kv + ): + return "attention window must fit the non-cyclic page-table capacity" + if QuantMode(self.attn.quant_mode).has_kv_cache_quant(): + return "quantized KV cache is not supported" + host_seq_lens = metadata.kv_lens_runtime[:batch_size] + min_seq_len_kv = int(host_seq_lens.min()) + if min_seq_len_kv <= 0: + return "every active request must contain at least one KV token" + mask_type = self._get_prims_mask_type(forward_args) + if mask_type == "causal" and min_seq_len_kv < seq_len_q: + return "causal KV lengths must be at least the fixed query length" + if int(host_seq_lens.max()) > logical_max_seq_len: + return "an active KV length exceeds the logical maximum sequence length" + return self._static_profile_unsupported_reason( + q, + inputs, + batch_size=batch_size, + seq_len_q=seq_len_q, + seq_len_kv=max_seq_len_kv, + page_size=page_size, + mask_type=mask_type, + ) + + def _ensure_preprocess_workspace( + self, + q: torch.Tensor, + workspace: torch.Tensor, + *, + batch_size: int, + ) -> None: + layout = self._get_generation_workspace_layout( + q.dtype, + batch_size, + int(q.shape[0]), + ) + required_bytes = int(layout["total_size"]) + available_bytes = workspace.numel() * workspace.element_size() + if available_bytes < required_bytes: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "TRT-LLM QKV preprocessing workspace must be sized before CUDA Graph capture" + ) + workspace.resize_((math.ceil(required_bytes / workspace.element_size()),)) + + def prepare_workspace( + self, + q: torch.Tensor, + k: torch.Tensor | None, + v: torch.Tensor | None, + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + workspace: torch.Tensor, + ) -> None: + del k, v, forward_args + with torch.cuda.device(q.device): + self._ensure_preprocess_workspace( + q, + workspace, + batch_size=int(metadata.num_generations), + ) + if self._multi_processor_count is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("GPU properties must be prepared before CUDA Graph capture") + self._multi_processor_count = torch.cuda.get_device_properties( + q.device + ).multi_processor_count + + def run_generation(self, params: FmhaParams) -> None: + q = params.qkv_input + output_buffer = params.context_buf + sequence_lengths = params.sequence_lengths + assert q is not None and output_buffer is not None and sequence_lengths is not None + metadata = params.meta + forward_args = params.fwd + inputs = forward_args.block_sparse_inputs + assert inputs is not None + batch_size = params.num_requests + seq_len_q = params.input_seq_length + page_size = params.tokens_per_block + block_offsets = metadata.kv_cache_block_offsets + assert block_offsets is not None + max_seq_len_kv = int(block_offsets.shape[-1]) * page_size + mask_type = self._get_prims_mask_type(forward_args) + assert self._multi_processor_count is not None + ( + q_processed, + kv_pool, + block_tables, + _kv_scale_pool, + _bmm1_scale, + _bmm2_scale, + fmha_workspace, + _cu_seqlens, + _max_q_len, + _max_kv_len, + _window_left, + _is_multi_token_gen, + ) = self._run_generation_preprocess(params) + if fmha_workspace.numel() != 0: + raise RuntimeError("PrimTS block-sparse preprocessing returned an FMHA workspace.") + if q_processed is None or kv_pool is None or block_tables is None: + raise RuntimeError("TRT-LLM preprocessing did not return paged PrimTS metadata") + kv_page_offset = get_kv_page_offset( + params.attn, + metadata, + params.seq_offset, + cache=self._kv_page_offset_cache, + ) + if kv_page_offset is None: + raise RuntimeError("PrimTS could not resolve the K-to-V page displacement") + k_cache, v_cache = self._standard_kv_views(kv_pool, kv_page_offset) + fixed_block_tables = self._get_fixed_block_tables( + block_tables, + batch_size, + ) + seq_lens = self._get_sequence_lengths(sequence_lengths, batch_size) + query = q_processed.view( + batch_size, + seq_len_q, + self.attn.num_heads, + self.attn.head_dim, + ) + output = output_buffer.view_as(query) + key = self._make_plan_key( + query, + inputs, + batch_size=batch_size, + seq_len_q=seq_len_q, + kv_capacity=max_seq_len_kv, + page_size=page_size, + mask_type=mask_type, + ) + wrapper = cast( + "BlockSparsePagedTSWrapper", + self._get_or_plan_wrapper(key, paged=True), + ) + wrapper.run( + query, + (k_cache, v_cache), + block_tables=fixed_block_tables, + seq_lens_kv=seq_lens, + block_indptr=inputs.block_indptr, + block_indices=inputs.block_indices, + kv_valid_bits=inputs.kv_valid_bits, + sm_scale=self._get_bmm1_scale(self.attn), + out=output, + ) + + def _contiguous_views( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + route_tensor: torch.Tensor, + output: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + batch_size = int(route_tensor.shape[0]) + if batch_size <= 0 or q.shape[0] % batch_size or k.shape[0] % batch_size: + raise ValueError("flat Q and K token counts must be divisible by route batch size") + seq_len_q = int(q.shape[0]) // batch_size + seq_len_kv = int(k.shape[0]) // batch_size + return ( + q.view(batch_size, seq_len_q, self.attn.num_heads, self.attn.head_dim), + k.view(batch_size, seq_len_kv, self.attn.num_kv_heads, self.attn.head_dim), + v.view(batch_size, seq_len_kv, self.attn.num_kv_heads, self.attn.head_dim), + output.view(batch_size, seq_len_q, self.attn.num_heads, self.attn.head_dim), + ) + + def _contiguous_unsupported_reason( + self, + q: torch.Tensor, + k: torch.Tensor | None, + v: torch.Tensor | None, + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + ) -> str | None: + common_reason = self._common_unsupported_reason(metadata, forward_args) + if common_reason is not None: + return common_reason + if self.attn.position_embedding_type != 0 or forward_args.mrope_position_deltas is not None: + return "contiguous Q/K/V must have position embedding applied before attention" + if forward_args.is_fused_qkv or k is None or v is None: + return "contiguous block-sparse attention requires separate Q, K, and V" + if forward_args.cu_q_seqlens is not None or forward_args.cu_kv_seqlens is not None: + return "packed variable-length Q/KV inputs are not supported" + inputs = forward_args.block_sparse_inputs + assert inputs is not None + mask_type = self._get_prims_mask_type(forward_args) + if inputs.use_proxy_routes and mask_type != "dense": + return "block-sparse proxy routes require mask_type='dense'" + route_tensor = ( + inputs.block_indptr if inputs.sparse_format == "bsr" else inputs.exact_block_bits + ) + batch_size = int(route_tensor.shape[0]) + if batch_size <= 0 or q.shape[0] % batch_size or k.shape[0] % batch_size: + return "flat Q and K token counts must be divisible by the route batch size" + seq_len_q = int(q.shape[0]) // batch_size + if not bool(metadata.seq_lens[:batch_size].eq(seq_len_q).all()): + return "query lengths must be batch-uniform and match the fixed query shape" + return self._static_profile_unsupported_reason( + q, + inputs, + batch_size=batch_size, + seq_len_q=seq_len_q, + seq_len_kv=int(k.shape[0]) // batch_size, + page_size=None, + mask_type=mask_type, + ) + + def _forward_contiguous( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + forward_args: AttentionForwardArgs, + ) -> None: + inputs = forward_args.block_sparse_inputs + assert inputs is not None + assert forward_args.output is not None + q_view, k_view, v_view, out_view = self._contiguous_views( + q, + k, + v, + inputs.block_indptr if inputs.sparse_format == "bsr" else inputs.exact_block_bits, + forward_args.output, + ) + mask_type = self._get_prims_mask_type(forward_args) + key = self._make_plan_key( + q_view, + inputs, + batch_size=int(q_view.shape[0]), + seq_len_q=int(q_view.shape[1]), + kv_capacity=int(k_view.shape[1]), + page_size=None, + mask_type=mask_type, + ) + wrapper = cast( + "BlockSparseTSWrapper", + self._get_or_plan_wrapper(key, paged=False), + ) + wrapper.run( + q_view, + k_view, + v_view, + block_indptr=inputs.block_indptr, + block_indices=inputs.block_indices, + exact_block_bits=inputs.exact_block_bits, + k_summary=inputs.k_summary, + v_summary=inputs.v_summary, + kv_valid_bits=inputs.kv_valid_bits, + sm_scale=self._get_bmm1_scale(self.attn), + out=out_view, + ) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor | None, + v: torch.Tensor | None, + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + ) -> None: + if metadata.kv_cache_manager is None: + assert k is not None and v is not None + self._forward_contiguous(q, k, v, forward_args) + return + super().forward(q, k, v, metadata, forward_args) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/registry.py b/tensorrt_llm/_torch/attention/backends/fmha/registry.py index 4e545924e63e..240e5242feb0 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/registry.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/registry.py @@ -35,6 +35,7 @@ def init_fmha_libs() -> dict[str, "FmhaCls"]: """ from .flashinfer_sparse_mla import FlashInferSparseMlaFmha from .msa_sparse_gqa import MsaSparseGqaFmha + from .prims_ts_block_sparse import PrimsTSBlockSparseFmha return { "triton_custom_mask": TritonCustomMaskFmha, @@ -42,6 +43,7 @@ def init_fmha_libs() -> dict[str, "FmhaCls"]: "msa_sparse_gqa": MsaSparseGqaFmha, "flashinfer_sparse_mla": FlashInferSparseMlaFmha, "prims_ts": PrimsTSFmha, + "prims_ts_block_sparse": PrimsTSBlockSparseFmha, "flashinfer_trtllm_gen": FlashInferTrtllmGenFmha, "fallback": FallbackFmha, } diff --git a/tensorrt_llm/_torch/attention/backends/fmha/triton_custom_mask.py b/tensorrt_llm/_torch/attention/backends/fmha/triton_custom_mask.py index 4391ef94bb22..1f2e5fbbbec9 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/triton_custom_mask.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/triton_custom_mask.py @@ -105,6 +105,8 @@ def _check_support_with_reason( *, phase: Optional[FmhaPhase] = None, ) -> tuple[bool, str]: + if forward_args.block_sparse_inputs is not None: + return False, "block_sparse_inputs are not supported." if phase != FmhaPhase.CONTEXT: return False, "Only context attention is supported." if forward_args.attention_mask != CustomAttentionMask.CUSTOM: diff --git a/tensorrt_llm/_torch/attention/backends/interface.py b/tensorrt_llm/_torch/attention/backends/interface.py index 7d11cdb28f8d..37144b0e3a96 100644 --- a/tensorrt_llm/_torch/attention/backends/interface.py +++ b/tensorrt_llm/_torch/attention/backends/interface.py @@ -32,8 +32,8 @@ from ...pyexecutor.resource_manager import KVCacheManager from ...pyexecutor.trace_log_utils import log_tensor_size from ...utils import get_model_extra_attrs -from .sparse.params import (SparseBackendForwardArgs, SparseMetadataParams, - SparseRuntimeParams) +from .sparse.params import (BlockSparseForwardInputs, SparseBackendForwardArgs, + SparseMetadataParams, SparseRuntimeParams) try: # Transformers v5 @@ -964,6 +964,7 @@ class AttentionForwardArgs: # Optional normalized diffusion timestep for timestep-varying sparse attention. timestep: Optional[torch.Tensor] = None + block_sparse_inputs: Optional[BlockSparseForwardInputs] = None sparse_backend_args: Optional[SparseBackendForwardArgs] = None sparse_runtime_params: SparseRuntimeParams = field( default_factory=SparseRuntimeParams) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/params.py b/tensorrt_llm/_torch/attention/backends/sparse/params.py index 956fc8c6fdbd..c99061926435 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/params.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/params.py @@ -15,7 +15,7 @@ """Shared sparse attention parameter types.""" from dataclasses import dataclass -from typing import Optional +from typing import Literal, Optional import torch @@ -38,6 +38,48 @@ class SparseBackendForwardArgs: topk_indices: Optional[torch.Tensor] = None +@dataclass(frozen=True, slots=True) +class BlockSparseForwardInputs: + """Block geometry and live routing payload for one attention call. + + 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. + """ + + q_block_size: int + kv_block_size: int + max_blocks_per_row: Optional[int] = None + block_indptr: Optional[torch.Tensor] = None + block_indices: Optional[torch.Tensor] = None + exact_block_bits: Optional[torch.Tensor] = None + k_summary: Optional[torch.Tensor] = None + v_summary: Optional[torch.Tensor] = None + kv_valid_bits: Optional[torch.Tensor] = None + + def __post_init__(self) -> None: + has_bsr = self.block_indptr is not None + if has_bsr != (self.block_indices is not None): + raise ValueError("block_indptr and block_indices must be provided together") + if has_bsr == (self.exact_block_bits is not None): + raise ValueError("exactly one route representation must be provided") + if has_bsr and self.max_blocks_per_row is None: + raise ValueError("BSR routes require max_blocks_per_row") + if (self.k_summary is None) != (self.v_summary is None): + raise ValueError("k_summary and v_summary must be provided together") + + @property + def sparse_format(self) -> Literal["bsr", "bitmask"]: + """Routing representation selected by the live payload.""" + return "bitmask" if self.exact_block_bits is not None else "bsr" + + @property + def use_proxy_routes(self) -> bool: + """Whether unselected blocks are represented by K/V summaries.""" + return self.k_summary is not None + + @dataclass(kw_only=True, slots=True) class SparseRuntimeParams: """Flat optional sparse inputs passed from a backend to ``AttentionOp``.""" @@ -57,3 +99,12 @@ class SparseRuntimeParams: threshold_scale_factor_prefill: float = 0.0 # SkipSoftmax decode threshold; diffusion models leave it at zero. threshold_scale_factor_decode: float = 0.0 + + +__all__ = [ + "BlockSparseForwardInputs", + "SparseBackendForwardArgs", + "SparseMetadataParams", + "SparseParams", + "SparseRuntimeParams", +] diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 914f6f3a654e..5a2d7ea3912b 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -1781,6 +1781,7 @@ def forward( ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: """Execute the TRTLLM attention backend.""" forward_args = merge_attention_forward_args(forward_args, kwargs) + has_block_sparse_inputs = forward_args.block_sparse_inputs is not None assert isinstance( metadata, TrtllmAttentionMetadata, @@ -1983,8 +1984,11 @@ def forward( assert v.shape[1] == kv_hidden_size num_tokens = q.shape[0] if k is not None and not metadata.is_cross: - assert k.shape[0] == num_tokens - assert v.shape[0] == num_tokens + if has_block_sparse_inputs: + assert v is not None and v.shape[0] == k.shape[0] + else: + assert k.shape[0] == num_tokens + assert v.shape[0] == num_tokens else: sparse_attn_indices = forward_args.sparse_runtime_params.sparse_attn_indices is_sparse_attn = sparse_attn_indices is not None and sparse_attn_indices.numel( diff --git a/tests/unittest/_torch/attention/test_fmha_manager.py b/tests/unittest/_torch/attention/test_fmha_manager.py index 17dd1beda5ce..a074aceac960 100644 --- a/tests/unittest/_torch/attention/test_fmha_manager.py +++ b/tests/unittest/_torch/attention/test_fmha_manager.py @@ -31,6 +31,7 @@ AttentionInputType, PredefinedAttentionMask, ) +from tensorrt_llm._torch.attention.backends.sparse.params import BlockSparseForwardInputs from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization.mode import QuantAlgo @@ -538,6 +539,47 @@ def test_fmha_cache_tracks_attention_mask_data() -> None: assert len(manager._cache) == 2 +@pytest.mark.parametrize("block_sparse_first", [False, True]) +def test_fmha_cache_separates_block_sparse_mode(block_sparse_first: bool) -> None: + events: list[tuple] = [] + attn, manager = _make_manager() + block_sparse_fmha = FakeFmha( + attn, + "block-sparse", + events, + support_predicate=lambda forward_args: forward_args.block_sparse_inputs is not None, + ) + dense_fmha = FakeFmha( + attn, + "dense", + events, + support_predicate=lambda forward_args: forward_args.block_sparse_inputs is None, + ) + manager.fmha_libs = [block_sparse_fmha, dense_fmha] + metadata = _make_metadata(num_contexts=1, num_generations=0, num_ctx_tokens=1) + q = torch.empty((1, 4)) + by_mode = { + False: AttentionForwardArgs(attention_input_type=AttentionInputType.context_only), + True: AttentionForwardArgs( + attention_input_type=AttentionInputType.context_only, + block_sparse_inputs=BlockSparseForwardInputs( + q_block_size=64, + kv_block_size=64, + exact_block_bits=torch.zeros((1, 1), dtype=torch.uint32), + ), + ), + } + order = (True, False) if block_sparse_first else (False, True) + + with patch.object(fmha_manager, "_is_fmha_cache_enabled", return_value=True): + selected = { + mode: manager.select(attn, q, None, None, metadata, by_mode[mode]) for mode in order + } + + assert selected == {False: dense_fmha, True: block_sparse_fmha} + assert len(manager._cache) == 2 + + @pytest.mark.parametrize("speculative_first", [False, True]) def test_fmha_cache_separates_speculative_decoding(speculative_first: bool) -> None: events: list[tuple] = [] diff --git a/tests/unittest/_torch/attention/test_fmha_registry.py b/tests/unittest/_torch/attention/test_fmha_registry.py index cd4cdb4fd7a2..64060bdbc935 100644 --- a/tests/unittest/_torch/attention/test_fmha_registry.py +++ b/tests/unittest/_torch/attention/test_fmha_registry.py @@ -14,10 +14,13 @@ # limitations under the License. import pytest +import torch from tensorrt_llm._torch.attention.backends.fmha import registry +from tensorrt_llm._torch.attention.backends.fmha.interface import Fmha PRIMS_TS = "prims_ts" +PRIMS_TS_BLOCK_SPARSE = "prims_ts_block_sparse" def _canonical_names() -> tuple[str, ...]: @@ -39,10 +42,27 @@ def test_default_fmha_libs_exclude_prims_ts(monkeypatch: pytest.MonkeyPatch) -> monkeypatch.delenv("TLLM_FMHA_LIBS", raising=False) assert PRIMS_TS not in registry.DEFAULT_FMHA_LIBS + assert PRIMS_TS_BLOCK_SPARSE in registry.DEFAULT_FMHA_LIBS assert set(registry.DEFAULT_FMHA_LIBS) <= set(registry.FMHA_LIBS) assert _enabled_names() == registry.DEFAULT_FMHA_LIBS +@pytest.mark.parametrize("name", [PRIMS_TS, "fallback"]) +def test_dense_fmhas_reject_unconsumed_block_sparse_inputs(name: str) -> None: + attention = type("Attention", (), {})() + fmha = object.__new__(registry.FMHA_LIBS[name]) + Fmha.__init__(fmha, attention) + forward_args = type("ForwardArgs", (), {"block_sparse_inputs": object()})() + + assert not fmha.is_supported( + torch.empty((1, 4), dtype=torch.bfloat16), + None, + None, + object(), + forward_args, + ) + + @pytest.mark.parametrize("value", ["", " ", ", ,"]) def test_empty_fmha_lib_env_uses_default( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unittest/_torch/attention/test_prims_ts_block_sparse.py b/tests/unittest/_torch/attention/test_prims_ts_block_sparse.py new file mode 100644 index 000000000000..e45f7da1166a --- /dev/null +++ b/tests/unittest/_torch/attention/test_prims_ts_block_sparse.py @@ -0,0 +1,690 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from contextlib import nullcontext +from importlib import import_module +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from utils.util import isSM100Family + +from tensorrt_llm._torch.attention.backends import prims_ts +from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase +from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + AttentionInputType, + PredefinedAttentionMask, +) +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from tensorrt_llm.functional import PositionEmbeddingType + +_REQUIRES_PRIMTS_GPU = pytest.mark.skipif( + not isSM100Family(), + reason="PrimTS block-sparse attention requires SM100 or SM103", +) + + +def _generic_api(): + carrier_module = import_module("tensorrt_llm._torch.attention.backends.sparse.params") + fmha_module = import_module("tensorrt_llm._torch.attention.backends.fmha.prims_ts_block_sparse") + return carrier_module.BlockSparseForwardInputs, fmha_module + + +def _bsr_inputs(*, kv_valid_bits: torch.Tensor | None = None): + inputs_type, _ = _generic_api() + return inputs_type( + q_block_size=64, + kv_block_size=64, + max_blocks_per_row=2, + block_indptr=torch.tensor([[[0, 2]], [[2, 4]]], dtype=torch.int32), + block_indices=torch.tensor([0, 1, 2, 3], dtype=torch.int32), + kv_valid_bits=kv_valid_bits, + ) + + +def _bitmask_inputs(*, proxy: bool): + inputs_type, _ = _generic_api() + summaries = { + "k_summary": torch.zeros((2, 4, 1, 128), dtype=torch.bfloat16), + "v_summary": torch.zeros((2, 4, 1, 128), dtype=torch.bfloat16), + } + return inputs_type( + q_block_size=64, + kv_block_size=64, + exact_block_bits=torch.ones((2, 1, 1, 1), dtype=torch.uint32), + **(summaries if proxy else {}), + ) + + +def _pack_token_mask(mask: torch.Tensor) -> torch.Tensor: + shifts = torch.arange(32, dtype=torch.int64, device=mask.device) + weights = torch.ones_like(shifts).bitwise_left_shift_(shifts) + return (mask.view(1, -1, 32).to(torch.int64) * weights).sum(dim=-1).to(torch.uint32) + + +def _proxy_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_summary: torch.Tensor, + v_summary: torch.Tensor, + exact_block: int, +) -> torch.Tensor: + block_size = 64 + exact_tokens = torch.arange( + exact_block * block_size, + (exact_block + 1) * block_size, + device=q.device, + ) + proxy_blocks = [block for block in range(k_summary.shape[1]) if block != exact_block] + q_rows = q[0, :, 0].float() + exact_logits = q_rows @ k[0, exact_tokens, 0].float().T + proxy_logits = q_rows @ k_summary[0, proxy_blocks, 0].float().T + logits = torch.cat((exact_logits, proxy_logits), dim=1) / math.sqrt(q.shape[-1]) + 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() + 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] + + +class _Attention: + def __init__(self) -> None: + self.sparse_params = None + self.num_heads = 2 + self.num_kv_heads = 1 + self.head_dim = 128 + self.is_mla_enable = False + self.kv_lora_rank = None + self.qk_rope_head_dim = None + self.qk_nope_head_dim = None + self.v_head_dim = None + self.q_scaling = 1.0 + self.quant_mode = 0 + self.local_layer_idx = 0 + self.position_embedding_type = PositionEmbeddingType.learned_absolute + self.attention_chunk_size = 0 + + +def _contiguous_case(): + _inputs_type, fmha_module = _generic_api() + attention = _Attention() + fmha = fmha_module.PrimsTSBlockSparseFmha(attention) + q = torch.zeros((128, 256), dtype=torch.bfloat16) + k = torch.zeros((512, 128), dtype=torch.bfloat16) + v = torch.zeros_like(k) + metadata = SimpleNamespace( + is_cross=False, + kv_cache_manager=None, + seq_lens=torch.tensor([64, 64], dtype=torch.int32), + ) + args = AttentionForwardArgs( + output=torch.empty_like(q), + attention_input_type=AttentionInputType.context_only, + attention_mask=PredefinedAttentionMask.FULL, + block_sparse_inputs=_bsr_inputs(), + ) + return attention, fmha, q, k, v, metadata, args + + +def _paged_metadata(): + batch_size, max_pages, page_size = 2, 4, 64 + key_pages = torch.arange(batch_size * max_pages, dtype=torch.int32).view(batch_size, max_pages) + block_offsets = torch.stack((key_pages, key_pages + 8), dim=1).unsqueeze(0) + manager = Mock(spec=KVCacheManager) + manager.dtype = torch.bfloat16 + manager.num_pools = manager.num_local_layers = 1 + manager.host_kv_cache_block_offsets = block_offsets + return SimpleNamespace( + is_cross=False, + num_contexts=0, + num_generations=batch_size, + seq_lens=torch.ones(batch_size, dtype=torch.int32), + beam_width=1, + tokens_per_block=page_size, + max_seq_len=max_pages * page_size, + kv_layout="HND", + kv_lens_runtime=torch.tensor([129, 193], dtype=torch.int32), + kv_cache_block_offsets=block_offsets, + host_kv_cache_pool_pointers=torch.tensor([[1234, 5678]], dtype=torch.int64), + host_kv_cache_pool_mapping=torch.tensor([[0, 0]], dtype=torch.int32), + kv_cache_manager=manager, + ) + + +def _paged_case(): + _inputs_type, fmha_module = _generic_api() + attention = _Attention() + fmha = fmha_module.PrimsTSBlockSparseFmha(attention) + fmha._multi_processor_count = 1 + metadata = _paged_metadata() + q = torch.zeros((2, 512), dtype=torch.bfloat16) + args = AttentionForwardArgs( + output=torch.empty((2, 256), dtype=q.dtype), + attention_input_type=AttentionInputType.generation_only, + attention_mask=PredefinedAttentionMask.CAUSAL, + attention_window_size=metadata.max_seq_len, + is_fused_qkv=True, + block_sparse_inputs=_bsr_inputs(), + ) + return attention, fmha, q, metadata, args + + +def test_block_sparse_route_mode_is_derived_from_payload() -> None: + bsr = _bsr_inputs() + exact = _bitmask_inputs(proxy=False) + proxy = _bitmask_inputs(proxy=True) + + assert (bsr.sparse_format, bsr.use_proxy_routes) == ("bsr", False) + assert (exact.sparse_format, exact.use_proxy_routes) == ("bitmask", False) + assert (proxy.sparse_format, proxy.use_proxy_routes) == ("bitmask", True) + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"block_indices": None}, "block_indptr and block_indices"), + ({"max_blocks_per_row": None}, "max_blocks_per_row"), + ( + {"exact_block_bits": torch.ones((1, 1, 1, 1), dtype=torch.uint32)}, + "exactly one route representation", + ), + ({"k_summary": torch.empty(0)}, "k_summary and v_summary"), + ], +) +def test_block_sparse_payload_rejects_ambiguous_combinations(overrides, message) -> None: + inputs_type, _ = _generic_api() + kwargs = { + "q_block_size": 64, + "kv_block_size": 64, + "max_blocks_per_row": 1, + "block_indptr": torch.tensor([[[0, 1]]], dtype=torch.int32), + "block_indices": torch.tensor([0], dtype=torch.int32), + } + kwargs.update(overrides) + + with pytest.raises((TypeError, ValueError), match=message): + inputs_type(**kwargs) + + +def test_block_sparse_support_is_phase_specific_and_paged_proxy_is_rejected( + monkeypatch, +) -> None: + _attention, contiguous, q, k, v, metadata, args = _contiguous_case() + monkeypatch.setattr(contiguous, "_common_unsupported_reason", Mock(return_value=None)) + assert contiguous.is_supported(q, k, v, metadata, args, phase=FmhaPhase.CONTEXT) + assert not contiguous.is_supported(q, k, v, metadata, args, phase=FmhaPhase.GENERATION) + + _attention, paged, q, metadata, args = _paged_case() + args.block_sparse_inputs = _bitmask_inputs(proxy=True) + _supported, reason = paged._is_supported_with_reason( + q, None, None, metadata, args, phase=FmhaPhase.GENERATION + ) + assert not _supported + assert reason == "paged block-sparse attention only supports BSR exact routes" + + +def test_contiguous_proxy_routes_reject_causal_mask_before_planning(monkeypatch) -> None: + _attention, fmha, q, k, v, metadata, args = _contiguous_case() + args.block_sparse_inputs = _bitmask_inputs(proxy=True) + args.attention_mask = PredefinedAttentionMask.CAUSAL + monkeypatch.setattr(fmha, "_common_unsupported_reason", Mock(return_value=None)) + + supported, reason = fmha._is_supported_with_reason( + q, + k, + v, + metadata, + args, + phase=FmhaPhase.CONTEXT, + ) + + assert not supported + assert reason == "block-sparse proxy routes require mask_type='dense'" + + +@pytest.mark.parametrize("paged", [False, True]) +def test_block_sparse_support_rejects_invalid_static_kernel_profile( + monkeypatch, + paged, +) -> None: + if paged: + attention, fmha, q, metadata, args = _paged_case() + attention.head_dim = 64 + q = torch.zeros((2, 256), dtype=torch.bfloat16) + args.output = torch.empty((2, 128), dtype=q.dtype) + monkeypatch.setattr(fmha, "_common_unsupported_reason", Mock(return_value=None)) + supported, reason = fmha._is_supported_with_reason( + q, + None, + None, + metadata, + args, + phase=FmhaPhase.GENERATION, + ) + else: + attention, fmha, _q, _k, _v, metadata, args = _contiguous_case() + attention.head_dim = 64 + q = torch.zeros((128, 128), dtype=torch.bfloat16) + k = torch.zeros((512, 64), dtype=torch.bfloat16) + v = torch.zeros_like(k) + args.output = torch.empty_like(q) + monkeypatch.setattr(fmha, "_common_unsupported_reason", Mock(return_value=None)) + supported, reason = fmha._is_supported_with_reason( + q, + k, + v, + metadata, + args, + phase=FmhaPhase.CONTEXT, + ) + + assert not supported + assert reason == "block-sparse requires head_dim=128" + + +def test_contiguous_wrappers_cache_static_profile_and_keep_routes_live(monkeypatch) -> None: + _attention, fmha, q, k, v, _metadata, args = _contiguous_case() + _inputs_type, fmha_module = _generic_api() + wrapper = Mock() + factory = Mock(return_value=wrapper) + monkeypatch.setattr(fmha_module, "_BlockSparseTSWrapper", factory) + + bsr_inputs = [ + _bsr_inputs(), + _inputs_type( + q_block_size=64, + kv_block_size=64, + max_blocks_per_row=2, + block_indptr=torch.tensor([[[0, 1]], [[1, 4]]], dtype=torch.int32), + block_indices=torch.tensor([3, 1, 0, 2], dtype=torch.int32), + ), + ] + for inputs in bsr_inputs: + args.block_sparse_inputs = inputs + fmha._forward_contiguous(q, k, v, args) + + proxy_inputs = [_bitmask_inputs(proxy=True), _bitmask_inputs(proxy=True)] + for inputs in proxy_inputs: + args.block_sparse_inputs = inputs + fmha._forward_contiguous(q, k, v, args) + + assert factory.call_count == 2 + assert wrapper.plan.call_count == 2 + assert wrapper.plan.call_args_list[0].kwargs["sparse_format"] == "bsr" + assert wrapper.plan.call_args_list[0].kwargs["use_proxy_routes"] is False + assert wrapper.plan.call_args_list[1].kwargs["sparse_format"] == "bitmask" + assert wrapper.plan.call_args_list[1].kwargs["use_proxy_routes"] is True + assert wrapper.plan.call_args_list[1].kwargs["max_blocks_per_row"] == 4 + assert wrapper.run.call_count == 4 + + for call, inputs in zip(wrapper.run.call_args_list[:2], bsr_inputs): + assert call.kwargs["block_indptr"] is inputs.block_indptr + assert call.kwargs["block_indices"] is inputs.block_indices + for call, inputs in zip(wrapper.run.call_args_list[2:], proxy_inputs): + assert call.kwargs["exact_block_bits"] is inputs.exact_block_bits + assert call.kwargs["k_summary"] is inputs.k_summary + assert call.kwargs["v_summary"] is inputs.v_summary + + +def test_block_sparse_plan_key_includes_attention_head_topology() -> None: + _, fmha_module = _generic_api() + inputs = _bitmask_inputs(proxy=True) + q = torch.empty((128, 256), dtype=torch.bfloat16) + first_attention = _Attention() + second_attention = _Attention() + second_attention.num_heads = 4 + first = fmha_module.PrimsTSBlockSparseFmha(first_attention) + second = fmha_module.PrimsTSBlockSparseFmha(second_attention) + + def _key(fmha): + return fmha._make_plan_key( + q, + inputs, + batch_size=1, + seq_len_q=128, + kv_capacity=256, + page_size=None, + mask_type="dense", + ) + + assert _key(first) != _key(second) + + +def test_block_sparse_plan_cache_is_shared_only_when_explicitly_bound() -> None: + _, fmha_module = _generic_api() + first = fmha_module.PrimsTSBlockSparseFmha(_Attention()) + second = fmha_module.PrimsTSBlockSparseFmha(_Attention()) + + assert first._contiguous_wrappers is not second._contiguous_wrappers + assert first._paged_wrappers is not second._paged_wrappers + + cache_state = {} + first.bind_plan_cache(cache_state) + second.bind_plan_cache(cache_state) + + assert first._contiguous_wrappers is second._contiguous_wrappers + assert first._paged_wrappers is second._paged_wrappers + assert cache_state == { + "contiguous_wrappers": {}, + "paged_wrappers": {}, + } + + +def test_paged_wrapper_uses_zero_copy_padded_row_stride_block_tables(monkeypatch) -> None: + attention, fmha, q, metadata, args = _paged_case() + _inputs_type, fmha_module = _generic_api() + wrapper = Mock() + monkeypatch.setattr(fmha_module, "_BlockSparsePagedTSWrapper", Mock(return_value=wrapper)) + monkeypatch.setattr(fmha_module, "get_kv_page_offset", Mock(return_value=8)) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", Mock(return_value=False)) + q_processed = torch.zeros((2, 2, 128), dtype=torch.bfloat16) + kv_pool = torch.empty((16, 1, 64, 128), dtype=torch.bfloat16) + block_tables = metadata.kv_cache_block_offsets[0] + empty = torch.empty(0, dtype=torch.uint8) + preprocessed = (q_processed, kv_pool, block_tables, None, 1.0, 1.0) + ( + empty, + None, + 1, + 256, + -1, + False, + ) + monkeypatch.setattr(fmha, "_run_generation_preprocess", Mock(return_value=preprocessed)) + params = FmhaParams( + attn=attention, + meta=metadata, + fwd=args, + workspace=torch.empty(0, dtype=torch.uint8), + qkv_input=q, + context_buf=args.output, + sequence_lengths=torch.tensor([129, 193], dtype=torch.int32), + input_seq_length=1, + tokens_per_block=64, + num_requests=2, + ) + expected_block_tables = block_tables[:2, 0, :] + snapshots = [] + + def snapshot(*_args, **kwargs): + snapshots.append( + ( + kwargs["seq_lens_kv"].clone(), + kwargs["block_tables"], + kwargs["block_tables"].clone(), + kwargs["block_indptr"], + kwargs["block_indices"], + ) + ) + + wrapper.run.side_effect = snapshot + first_inputs = args.block_sparse_inputs + fmha.run_generation(params) + block_tables[:, 0].add_(10) + params.sequence_lengths = torch.tensor([130, 194], dtype=torch.int32) + args.block_sparse_inputs = _inputs_type( + q_block_size=64, + kv_block_size=64, + max_blocks_per_row=2, + block_indptr=torch.tensor([[[0, 1]], [[1, 4]]], dtype=torch.int32), + block_indices=torch.tensor([3, 2, 1, 0], dtype=torch.int32), + ) + fmha.run_generation(params) + + wrapper.plan.assert_called_once() + assert wrapper.run.call_count == 2 + torch.testing.assert_close(snapshots[0][0], torch.tensor([129, 193], dtype=torch.int32)) + torch.testing.assert_close(snapshots[1][0], torch.tensor([130, 194], dtype=torch.int32)) + assert snapshots[0][1].data_ptr() == expected_block_tables.data_ptr() + assert snapshots[1][1].data_ptr() == expected_block_tables.data_ptr() + assert snapshots[0][1].shape == (2, 4) + assert snapshots[0][1].stride() == (8, 1) + torch.testing.assert_close(snapshots[0][2], torch.arange(8, dtype=torch.int32).view(2, 4)) + torch.testing.assert_close(snapshots[1][2], torch.arange(8, dtype=torch.int32).view(2, 4) + 10) + assert snapshots[0][3] is first_inputs.block_indptr + assert snapshots[1][3] is args.block_sparse_inputs.block_indptr + + +def test_paged_block_tables_remain_live_across_graph_replay(monkeypatch) -> None: + attention, fmha, q, metadata, args = _paged_case() + _inputs_type, fmha_module = _generic_api() + wrapper = Mock() + monkeypatch.setattr(fmha_module, "_BlockSparsePagedTSWrapper", Mock(return_value=wrapper)) + monkeypatch.setattr(fmha_module, "get_kv_page_offset", Mock(return_value=8)) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", Mock(return_value=True)) + q_processed = torch.zeros((2, 2, 128), dtype=torch.bfloat16) + kv_pool = torch.empty((16, 1, 64, 128), dtype=torch.bfloat16) + block_tables = metadata.kv_cache_block_offsets[0] + empty = torch.empty(0, dtype=torch.uint8) + preprocessed = (q_processed, kv_pool, block_tables, None, 1.0, 1.0) + ( + empty, + None, + 1, + 256, + -1, + False, + ) + monkeypatch.setattr(fmha, "_run_generation_preprocess", Mock(return_value=preprocessed)) + params = FmhaParams( + attn=attention, + meta=metadata, + fwd=args, + workspace=torch.empty(0, dtype=torch.uint8), + qkv_input=q, + context_buf=args.output, + sequence_lengths=torch.tensor([129, 193], dtype=torch.int32), + input_seq_length=1, + tokens_per_block=64, + num_requests=2, + ) + seen = [] + + def snapshot(*_args, **kwargs): + seen.append((kwargs["block_tables"].data_ptr(), kwargs["block_tables"].clone())) + + wrapper.run.side_effect = snapshot + fmha.run_generation(params) + block_tables[:, 0, :].add_(10) + block_tables[:, 1, :].fill_(-1) + fmha.run_generation(params) + + assert seen[0][0] == seen[1][0] == block_tables.data_ptr() + torch.testing.assert_close(seen[0][1], torch.arange(8, dtype=torch.int32).view(2, 4)) + torch.testing.assert_close(seen[1][1], torch.arange(8, dtype=torch.int32).view(2, 4) + 10) + + +def test_prepare_workspace_checks_capture_before_resize(monkeypatch) -> None: + _attention, fmha, _q, _metadata, _args = _paged_case() + query_device = torch.device("cuda:1") + q = SimpleNamespace( + device=query_device, + dtype=torch.bfloat16, + shape=(2, 512), + ) + metadata = SimpleNamespace( + kv_cache_block_offsets=SimpleNamespace(device=query_device, shape=(1, 2, 4)), + max_num_requests=2, + tokens_per_block=64, + num_generations=2, + ) + workspace = torch.empty(0, dtype=torch.uint8) + monkeypatch.setattr( + fmha, + "_get_generation_workspace_layout", + Mock(return_value={"total_size": 16}), + ) + fmha._multi_processor_count = 1 + device_scope = Mock(return_value=nullcontext()) + monkeypatch.setattr(torch.cuda, "device", device_scope) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", Mock(return_value=True)) + + with pytest.raises(RuntimeError, match="workspace must be sized"): + fmha.prepare_workspace(q, None, None, metadata, _args, workspace) + + device_scope.assert_called_once_with(query_device) + assert workspace.numel() == 0 + + +@_REQUIRES_PRIMTS_GPU +@torch.no_grad() +def test_real_gpu_raw_routes_and_token_mask_match_reference() -> None: + torch.manual_seed(1234) + q = torch.randn((1, 128, 1, 128), device="cuda", dtype=torch.float16) + k = torch.randn((1, 256, 1, 128), device="cuda", dtype=torch.float16) + v = torch.randn_like(k) + token_mask = torch.ones(256, device="cuda", dtype=torch.bool) + token_mask[[1, 63, 64, 95, 129, 190, 255]] = False + inputs_type, _ = _generic_api() + inputs = inputs_type( + q_block_size=64, + kv_block_size=64, + max_blocks_per_row=3, + block_indptr=torch.tensor([[[0, 2, 5]]], device="cuda", dtype=torch.int32), + block_indices=torch.tensor([0, 2, 0, 1, 3], device="cuda", dtype=torch.int32), + kv_valid_bits=_pack_token_mask(token_mask), + ) + sm_scale = 128**-0.5 + + key_blocks = torch.arange(256, device="cuda") // 64 + allowed = torch.zeros((128, 256), device="cuda", dtype=torch.bool) + for row, selected_blocks in enumerate(((0, 2), (0, 1, 3))): + selected = torch.tensor(selected_blocks, device="cuda") + allowed[row * 64 : (row + 1) * 64] = torch.isin(key_blocks, selected) & token_mask + scores = (q[0, :, 0].float() @ k[0, :, 0].float().T) * sm_scale + expected = ( + torch.softmax(scores.masked_fill(~allowed, float("-inf")), dim=-1) @ v[0, :, 0].float() + ).to(q.dtype)[None, :, None, :] + + actual = prims_ts.block_sparse_attention( + q, + k, + v, + block_indptr=inputs.block_indptr, + block_indices=inputs.block_indices, + q_block_size=inputs.q_block_size, + kv_block_size=inputs.kv_block_size, + kv_valid_bits=inputs.kv_valid_bits, + sm_scale=sm_scale, + ) + torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2) + + +@_REQUIRES_PRIMTS_GPU +@torch.no_grad() +def test_real_gpu_proxy_adapter_replays_live_routes_and_summaries() -> None: + torch.manual_seed(20260901) + q = torch.randn((1, 64, 1, 128), device="cuda", dtype=torch.bfloat16) + k = torch.randn((1, 192, 1, 128), device="cuda", dtype=torch.bfloat16) + v = torch.randn_like(k) + 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) + 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) + + inputs_type, fmha_module = _generic_api() + attention = _Attention() + attention.num_heads = attention.num_kv_heads = 1 + fmha = fmha_module.PrimsTSBlockSparseFmha(attention) + output = torch.empty_like(q).view(64, 128) + args = AttentionForwardArgs( + output=output, + attention_input_type=AttentionInputType.context_only, + attention_mask=PredefinedAttentionMask.FULL, + block_sparse_inputs=inputs_type( + q_block_size=64, + kv_block_size=64, + exact_block_bits=live_exact_bits, + k_summary=live_k_summary, + v_summary=live_v_summary, + ), + ) + metadata = SimpleNamespace( + is_cross=False, + kv_cache_manager=None, + seq_lens=torch.tensor([64], dtype=torch.int32), + ) + flat_q, flat_k, flat_v = (tensor.flatten(0, 2) for tensor in (q, k, v)) + + fmha.forward(flat_q, flat_k, flat_v, metadata, args) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + fmha.forward(flat_q, flat_k, flat_v, metadata, args) + + graph.replay() + torch.cuda.synchronize() + expected = _proxy_reference(q, k, v, live_k_summary, live_v_summary, exact_block=0) + torch.testing.assert_close(output.view_as(q), expected, rtol=2e-2, atol=2e-2) + + live_exact_bits.fill_(1 << 2) + live_k_summary.copy_((initial_k_summary.float() * 0.5 + 0.125).to(k.dtype)) + live_v_summary.copy_((initial_v_summary.float() * -0.25).to(v.dtype)) + graph.replay() + torch.cuda.synchronize() + expected = _proxy_reference(q, k, v, live_k_summary, live_v_summary, exact_block=2) + torch.testing.assert_close(output.view_as(q), expected, rtol=2e-2, atol=2e-2) + + +@_REQUIRES_PRIMTS_GPU +@torch.no_grad() +def test_real_gpu_paged_routes_use_live_length_below_capacity() -> None: + torch.manual_seed(7) + q = torch.randn((1, 64, 1, 128), device="cuda", dtype=torch.float16) + k_cache = torch.randn((4, 1, 64, 128), device="cuda", dtype=torch.float16) + v_cache = torch.randn_like(k_cache) + page_indices = torch.tensor([2, 0, 3, 1], device="cuda", dtype=torch.int32) + seq_lens_kv = torch.tensor([160], device="cuda", dtype=torch.int32) + inputs_type, _ = _generic_api() + inputs = inputs_type( + q_block_size=64, + kv_block_size=64, + max_blocks_per_row=2, + block_indptr=torch.tensor([[[0, 2]]], device="cuda", dtype=torch.int32), + block_indices=torch.tensor([0, 2], device="cuda", dtype=torch.int32), + ) + sm_scale = 128**-0.5 + + actual = prims_ts.block_sparse_attention_with_paged_kv_cache( + q, + (k_cache, v_cache), + block_tables=page_indices.view(1, 4), + seq_lens_kv=seq_lens_kv, + block_indptr=inputs.block_indptr, + block_indices=inputs.block_indices, + max_seq_len_kv=256, + q_block_size=inputs.q_block_size, + kv_block_size=inputs.kv_block_size, + sm_scale=sm_scale, + ) + + logical_k = k_cache.index_select(0, page_indices.long()).reshape(256, 1, 128) + logical_v = v_cache.index_select(0, page_indices.long()).reshape(256, 1, 128) + allowed = torch.zeros(256, device="cuda", dtype=torch.bool) + allowed[:64] = True + allowed[128:160] = True + scores = (q[0, :, 0].float() @ logical_k[:, 0].float().T) * sm_scale + expected = ( + torch.softmax(scores.masked_fill(~allowed, float("-inf")), dim=-1) @ logical_v[:, 0].float() + ).to(q.dtype)[None, :, None, :] + + torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2) From ee44e6f816c29f5097f302c20cc32f31beecc0de Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Mon, 7 Sep 2026 10:25:26 +0000 Subject: [PATCH 3/7] refactor: unify sparse attention runtime inputs Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- .../sparse-attention-development-guide.md | 54 +- .../attention/ATTENTION_DEVELOPER_GUIDE.md | 45 +- .../attention/backends/fmha/cute_dsl_mla.py | 2 +- .../attention/backends/fmha/fallback.py | 8 +- .../backends/fmha/flashinfer_sparse_mla.py | 2 +- .../backends/fmha/flashinfer_trtllm_gen.py | 7 +- .../_torch/attention/backends/fmha/manager.py | 2 +- .../attention/backends/fmha/msa_sparse_gqa.py | 6 +- .../attention/backends/fmha/prims_ts.py | 228 +++--- .../backends/fmha/prims_ts_block_sparse.py | 653 +++++++----------- .../backends/fmha/triton_custom_mask.py | 2 +- .../_torch/attention/backends/interface.py | 5 +- .../_torch/attention/backends/sparse/hooks.py | 38 +- .../attention/backends/sparse/params.py | 6 +- .../_torch/attention/backends/trtllm.py | 54 +- .../attention/sparse/test_sparse_attention.py | 264 ++++++- .../attention/test_attention_op_sync.py | 24 +- .../_torch/attention/test_fmha_manager.py | 25 +- .../_torch/attention/test_fmha_registry.py | 7 +- .../attention/test_prims_ts_block_sparse.py | 146 ++-- .../_torch/attention/test_prims_ts_fmha.py | 3 +- .../attention/test_skip_softmax_sm120.py | 7 +- 22 files changed, 924 insertions(+), 664 deletions(-) diff --git a/docs/source/developer-guide/sparse-attention-development-guide.md b/docs/source/developer-guide/sparse-attention-development-guide.md index 32d103b50058..20f2b771374c 100644 --- a/docs/source/developer-guide/sparse-attention-development-guide.md +++ b/docs/source/developer-guide/sparse-attention-development-guide.md @@ -136,19 +136,48 @@ Dynamic generation-phase KV eviction is tracked as future work. ### Prediction hooks -`TrtllmAttention`-based sparse backends expose two prediction methods that +`TrtllmAttention`-based sparse backends expose three prediction methods that algorithm-specific subclasses override: ```python sparse_kv_indices, sparse_kv_offsets = self.sparse_kv_predict(q, k, metadata, forward_args) sparse_attn_indices, sparse_attn_offsets = self.sparse_attn_predict(q, k, metadata, forward_args) +block_sparse_inputs = self.block_sparse_attn_predict(q, k, v, metadata, forward_args) ``` -`hooks.py` writes these results to `SparseRuntimeParams`. SkipSoftmax writes -its thresholds to the same runtime interface consumed by `AttentionOp`. -`AttentionForwardArgs.sparse_backend_args` carries algorithm inputs from the -module to the backend, while `sparse_runtime_params` carries lowered inputs -from the backend to `AttentionOp`. +`prepare_sparse_runtime_params` in `sparse/hooks.py` runs all three hooks once +per call regardless of whether the backend carries `SparseParams`, applies the +SkipSoftmax threshold schedule when the backend carries `SkipSoftmaxParams`, +and returns a new per-call `SparseRuntimeParams` built from the caller's +`AttentionForwardArgs.sparse_runtime_params` plus the hook results. The core +forward never assigns that field; it dispatches with the returned carrier. +Backends that need runtime state outside the three hooks (DSA's auxiliary pool +pointer, DeepSeek-V4's per-token KV lengths) write it into the caller's carrier +before or inside their hooks, and `prepare_sparse_runtime_params` carries those +fields over. +`AttentionForwardArgs.sparse_backend_args` carries +algorithm inputs from the module to the backend, while +`AttentionForwardArgs.sparse_runtime_params` carries the complete lowered state +from the backend through FMHA dispatch to `AttentionOp`. + +`SparseRuntimeParams.block_sparse_inputs` is the nested carrier for optional, +algorithm-neutral `BlockSparseForwardInputs`. The selected general +block-sparse FMHA validates and consumes that field; dense FMHA libraries reject +it instead of silently ignoring its routes. `AttentionForwardArgs` defaults +the field to an empty `SparseRuntimeParams()`; the core forward always +dispatches with the carrier prepared for the current call. + +`block_sparse_attn_predict` runs even when the backend has no `SparseParams`. +Its default implementation hands through +`SparseBackendForwardArgs.block_sparse_inputs`, so an attention module that +predicts routes before the core forward only needs to place the complete +payload in `sparse_backend_args`. Algorithms that predict inside the backend +override the hook, read `metadata` for the batch layout and `forward_args` for +per-call state such as `timestep`, and return `None` for dense phases. + +The core contract owns this runtime transport and general block-sparse FMHA +execution. Algorithm integrations own their prediction policy, effective Q/K/V +preparation, and any post-processing around the normal core forward. Different KV heads are allowed to emit different sparse index sets; Q heads that map to the same KV head share the KV head's sparse pattern. @@ -288,6 +317,19 @@ prediction methods. A `VanillaAttention` implementation instead overrides different index layouts. Match the selected kernel contract; do not pass request-local block indices to the physical-token path. +**`block_sparse_attn_predict(self, q, k, v, metadata, forward_args)`** + +- **Behavior**: return the `BlockSparseForwardInputs` consumed by the + general block-sparse FMHA, or `None` for a dense call. +- **Outputs**: block geometry plus exactly one route representation + (BSR `block_indptr`/`block_indices` or a packed `exact_block_bits` + bitmask), optional K/V summaries for proxy routes, and optional + `kv_valid_bits` masking ragged KV tails. +- **Default**: hands through `SparseBackendForwardArgs.block_sparse_inputs`, + so modules that predict before the core forward do not override it. + Override it to predict inside the backend from the flattened Q/K/V, + the batch layout in `metadata`, and per-call state in `forward_args`. + Prediction is on the critical path and can dominate latency in low-latency scenarios. Plan for custom kernels (Triton or CUDA) rather than relying on generic PyTorch ops. diff --git a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index e15d77d2634c..f8e438cabab6 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md @@ -155,9 +155,12 @@ their module-to-backend inputs in a `SparseBackendForwardArgs` subclass and pass it through the registered `AttentionForwardArgs.sparse_backend_args` field. For example, DSA owns `DSABackendForwardArgs`, whose indexer intermediates are consumed by `DSATrtllmAttention.sparse_attn_predict`. -Shared sparse carriers, including `SparseBackendForwardArgs.topk_indices` and -the backend-to-AttentionOp `SparseRuntimeParams`, live in -`attention/backends/sparse/params.py`. +Shared sparse carriers, including `SparseBackendForwardArgs.topk_indices`, +`SparseBackendForwardArgs.block_sparse_inputs`, and the +backend-to-FMHA/`AttentionOp` `SparseRuntimeParams`, live in +`attention/backends/sparse/params.py`. The latter is carried by +`AttentionForwardArgs.sparse_runtime_params` and nests optional general +block-sparse inputs in `SparseRuntimeParams.block_sparse_inputs`. For MLA-related tasks, first check whether the work fits the current projection structure, can stay on an existing backend and metadata family, and @@ -213,11 +216,12 @@ registration alone. Block-sparse FMHA is a kernel-library contract rather than a sparse algorithm. Algorithms lower their live routing state to an algorithm-neutral -`BlockSparseForwardInputs`: block geometry plus either canonical BSR routes or -an exact packed bitmask. Optional K/V summaries enable proxy routes, and -optional token-validity bits mask ragged KV tails. Plans contain only static -format, proxy, geometry, and capacity choices; every run receives the live -routes, summaries, validity bits, page tables, and sequence lengths. +`BlockSparseForwardInputs`, nested at +`SparseRuntimeParams.block_sparse_inputs`: block geometry plus either canonical +BSR routes or an exact packed bitmask. Optional K/V summaries enable proxy +routes, and optional token-validity bits mask ragged KV tails. Plans contain +only static format, proxy, geometry, and capacity choices; every run receives +the live routes, summaries, validity bits, page tables, and sequence lengths. `PrimsTSBlockSparseFmha` owns its wrapper-plan cache by default. Integrations whose attention layers execute serially may explicitly bind a model-scoped @@ -225,6 +229,20 @@ cache to reuse graph-stable route workspaces across compatible layers. The cache must not be shared by concurrent forwards; each independent model component must own separate state. +`TrtllmAttention.block_sparse_attn_predict(q, k, v, metadata, forward_args)` +is the backend hook that produces this payload; `prepare_sparse_runtime_params` +calls it even when the backend has no `SparseParams`. The default hands through +`SparseBackendForwardArgs.block_sparse_inputs`, which lets an attention module +predict routes before the core forward and pass the complete payload in +`AttentionForwardArgs.sparse_backend_args`. Algorithms that predict inside the +backend override the hook and return `None` for dense phases. + +The core library owns this general planning, validation, and execution +contract. Algorithm integrations own the surrounding lifecycle: prediction +policy, effective Q/K/V preparation before the core forward, plus any +algorithm-specific post-processing afterward. They route their payload through +these hooks instead of adding algorithm-specific FMHA libraries. + ### 2.3 Backend contract All backends implement the `AttentionBackend` interface. @@ -365,9 +383,14 @@ preserve its routing semantics. Delta entries update the default membership and follow canonical registry order, while an exact list preserves the user-specified order. Each FMHA library exposes `is_available()` for module/static environment checks and `is_supported()` for per-forward -request checks. `AttentionForwardArgs.block_sparse_inputs` is exclusive: the -block-sparse implementation validates and consumes it, while every other -library rejects it rather than silently dropping sparse routing semantics. +request checks. `AttentionForwardArgs.sparse_runtime_params` is the sole +per-call lowered sparse runtime carrier and defaults to an empty +`SparseRuntimeParams()`. The core forward never assigns that field; it +dispatches with a new carrier that `prepare_sparse_runtime_params` builds from +the caller's carrier plus the hook results. The carrier holds both flat `AttentionOp` parameters and optional +`BlockSparseForwardInputs` in its nested `block_sparse_inputs` field. The +block-sparse implementation validates and consumes the latter, while every +other library rejects it rather than silently dropping sparse routing semantics. For mixed non-MLA batches, the manager checks each active phase independently with `is_supported(..., phase=...)`; a phased library accepts only phases backed by its corresponding `run_*()` entry point. diff --git a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py index 224f301424e2..12342a23ea2e 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py @@ -285,7 +285,7 @@ def _is_supported_with_reason( meta: "TrtllmAttentionMetadata", fwd: AttentionForwardArgs, ) -> tuple[bool, str]: - if fwd.block_sparse_inputs is not None: + if fwd.sparse_runtime_params.block_sparse_inputs is not None: return False, "block_sparse_inputs are not supported." if fwd.attention_input_type != AttentionInputType.generation_only: return False, "CuTe DSL MLA FMHA only supports generation-only attention." diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py index 80825703f612..033ba3b4aa03 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py @@ -40,11 +40,11 @@ _THOP_EXCLUDED_FIELDS: frozenset = frozenset( { "sparse_backend_args", # consumed by sparse prediction before the attention op - "block_sparse_inputs", # consumed by the selected FMHA library + "block_sparse_inputs", # consumed by the selected block-sparse FMHA "attention_mask_data", # custom-mask code path "out_scale_sf", # promoted into ``out_scale`` in ``TrtllmAttention.forward`` for NVFP4 path "skip_mla_rope_generation", # handled in ``TrtllmAttention.forward`` for the test-only MLA path - "timestep", # used to populate skip-softmax params in ``TrtllmAttention.forward`` + "timestep", # consumed by sparse prediction before FMHA dispatch } ) @@ -82,7 +82,7 @@ def is_supported( del k, v, phase if q is not None and q.dtype == torch.float8_e4m3fn: return False - return forward_args.block_sparse_inputs is None and ( + return forward_args.sparse_runtime_params.block_sparse_inputs is None and ( forward_args.attention_mask != CustomAttentionMask.CUSTOM and (forward_args.update_kv_cache or metadata.is_cross) ) @@ -220,7 +220,7 @@ def forward( forward_args.sparse_runtime_params.sparse_attn_indices_block_size ), sparse_attn_kv_lens=forward_args.sparse_runtime_params.sparse_attn_kv_lens, - aux_kv_cache_pool_ptr=(forward_args.sparse_runtime_params.aux_kv_cache_pool_ptr), + aux_kv_cache_pool_ptr=forward_args.sparse_runtime_params.aux_kv_cache_pool_ptr, skip_softmax_threshold_scale_factor_prefill=( forward_args.sparse_runtime_params.threshold_scale_factor_prefill ), diff --git a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_sparse_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_sparse_mla.py index 80606dfd0006..1605e9e45151 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_sparse_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_sparse_mla.py @@ -81,7 +81,7 @@ def is_supported( phase: Optional[FmhaPhase] = None, ) -> bool: del q, k, v, metadata, phase - return forward_args.block_sparse_inputs is None + return forward_args.sparse_runtime_params.block_sparse_inputs is None def forward( self, diff --git a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py index 28aee806faf7..cc805ff9d724 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py @@ -571,7 +571,8 @@ def _is_supported_with_reason( *, phase: Optional[FmhaPhase] = None, ) -> Tuple[bool, str]: - if fwd.block_sparse_inputs is not None: + sparse_runtime_params = fwd.sparse_runtime_params + if sparse_runtime_params.block_sparse_inputs is not None: return False, "block_sparse_inputs are not supported." is_mla_enable = attn.is_mla_enable if phase is None: @@ -654,8 +655,8 @@ def _is_supported_with_reason( return False, "sage attention." if meta.helix_position_offsets is not None: return False, "helix parallelism." - sparse_kv_indices = fwd.sparse_runtime_params.sparse_kv_indices - sparse_attn_indices = fwd.sparse_runtime_params.sparse_attn_indices + sparse_kv_indices = sparse_runtime_params.sparse_kv_indices + sparse_attn_indices = sparse_runtime_params.sparse_attn_indices if ( (sparse_kv_indices is not None and sparse_kv_indices.numel() > 0) or (sparse_attn_indices is not None and sparse_attn_indices.numel() > 0) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/manager.py b/tensorrt_llm/_torch/attention/backends/fmha/manager.py index c18747844cce..c9afff66e1d8 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/manager.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/manager.py @@ -367,7 +367,7 @@ def _make_cache_key( generation_seq_len_q, _FMHA_CACHE_SEQ_LEN_Q_GRID ) - block_sparse_inputs = forward_args.block_sparse_inputs + block_sparse_inputs = forward_args.sparse_runtime_params.block_sparse_inputs return _FmhaCacheKey( context_batch_size=context_batch_size, generation_batch_size=generation_batch_size, diff --git a/tensorrt_llm/_torch/attention/backends/fmha/msa_sparse_gqa.py b/tensorrt_llm/_torch/attention/backends/fmha/msa_sparse_gqa.py index c0b37b651ace..6f6ce04b576c 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/msa_sparse_gqa.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/msa_sparse_gqa.py @@ -5,7 +5,7 @@ MsaSparseGqaFmha wraps the fmha_sm100 paged sparse GQA kernel and participates in the standard TrtllmAttention.forward dispatch loop. The owning MiniMax-M3 MSA attention layer runs an MsaIndexer to select the -per-query KV blocks and publishes them on forward_args.sparse_runtime_params; +per-query KV blocks and publishes them through the aggregate sparse prediction; this class attends over them. """ @@ -175,7 +175,7 @@ class MsaSparseGqaFmha(Fmha): """SM100 paged GQA FMHA powered by MSA's fmha_sm100 kernel. Handles every MiniMax-M3 MSA layer. Sparse layers pass the indexer's - selected KV block indices on forward_args.sparse_runtime_params.sparse_attn_indices + selected KV block indices on the aggregate sparse prediction and attend those blocks; dense layers leave the indices None and attend the full page table. @@ -221,7 +221,7 @@ def is_supported( phase: Optional[FmhaPhase] = None, ) -> bool: del q, k, v, metadata, phase - return forward_args.block_sparse_inputs is None + return forward_args.sparse_runtime_params.block_sparse_inputs is None def forward( self, diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py index ab48753bfd8e..3b5be3d9bcff 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py @@ -59,6 +59,96 @@ _WORKSPACE_ALIGNMENT = 32 +def get_paged_kv_storage_unsupported_reason( + attn: "TrtllmAttention", + metadata: "TrtllmAttentionMetadata", +) -> Optional[str]: + """Return why the TRT-LLM paged KV storage cannot feed a fixed page-table kernel.""" + if metadata.kv_cache_manager is None: + return "a KV cache manager is required." + if metadata.kv_cache_block_offsets is None: + return "paged KV-cache block offsets are required." + if metadata.host_kv_cache_pool_pointers is None: + return "KV-cache pool pointers are required." + pool_mapping = metadata.host_kv_cache_pool_mapping + if pool_mapping is None: + return "KV-cache pool mapping is required." + if metadata.kv_layout != "HND": + return "only HND KV-cache layout is supported." + manager = metadata.kv_cache_manager + if isinstance(manager, KVCacheManagerV2): + if manager.enable_swa_scratch_reuse: + return "KVCacheManagerV2 SWA scratch reuse is not supported." + elif isinstance(manager, KVCacheManager): + if manager.num_pools != 1: + return "KVCacheManagerV1 with multiple memory pools is not supported." + local_layer_idx = attn.local_layer_idx + if ( + pool_mapping.ndim != 2 + or pool_mapping.shape[1] < 2 + or local_layer_idx is None + or not 0 <= local_layer_idx < pool_mapping.shape[0] + ): + return "KVCacheManagerV1 has an invalid layer-to-pool mapping." + pool_index = int(pool_mapping[local_layer_idx, 0]) + layer_idx_in_pool = int(pool_mapping[local_layer_idx, 1]) + if pool_index != 0 or not 0 <= layer_idx_in_pool < manager.num_local_layers: + return "KVCacheManagerV1 has an invalid layer-to-pool mapping." + else: + return f"unsupported KV cache manager {type(manager).__name__}." + return None + + +def get_paged_kv_policy_unsupported_reason( + attn: "TrtllmAttention", + metadata: "TrtllmAttentionMetadata", +) -> Optional[str]: + """Return why the request's decoding policy is outside the fixed page-table envelope.""" + if metadata.beam_width != 1: + return "beam search is not supported." + if ( + metadata.is_spec_decoding_enabled + or metadata.use_spec_decoding + or metadata.is_spec_dec_tree + or metadata.is_spec_dec_dynamic_tree + ): + return "speculative decoding is not supported by the initial adapter." + position_embedding_type = int(attn.position_embedding_type) + if position_embedding_type in (4, 5, 6, 7, 10): + return f"position embedding type {position_embedding_type} is not supported." + try: + quant_mode = QuantMode(attn.quant_mode) + except (TypeError, ValueError): + return "invalid KV-cache quantization mode." + if quant_mode.has_kv_cache_quant(): + return "quantized KV cache is not supported by the initial adapter." + return None + + +def get_attention_feature_unsupported_reason( + metadata: "TrtllmAttentionMetadata", + forward_args: "AttentionForwardArgs", +) -> Optional[str]: + """Return which optional attention feature the fused-kernel adapters do not implement.""" + if metadata.helix_position_offsets is not None: + return "Helix parallelism is not supported." + if forward_args.relative_attention_bias is not None: + return "relative attention bias is not supported." + if forward_args.attention_sinks is not None: + return "attention sinks are not supported." + if forward_args.attention_mask_data is not None: + return "custom attention masks are not supported." + if forward_args.enable_dsv4_epilogue_fusion: + return "DSv4 epilogue fusion is not supported." + if ( + forward_args.sage_attn_num_elts_per_blk_q > 0 + or forward_args.sage_attn_num_elts_per_blk_k > 0 + or forward_args.sage_attn_num_elts_per_blk_v > 0 + ): + return "SageAttention is not supported." + return None + + class PrimsTSFmha(PhasedFmha): """Blackwell task-scheduled paged context and decode FMHA library.""" @@ -180,12 +270,13 @@ def _is_supported_with_reason( phase: Optional[FmhaPhase] = None, ) -> tuple[bool, str]: """Return a conservative, side-effect-free whole-request support decision.""" + sparse_runtime_params = fwd.sparse_runtime_params + if sparse_runtime_params.block_sparse_inputs is not None: + return False, "block_sparse_inputs are not supported." # PrimTS prepares workspace for every active request phase before # dispatch. Accept the phased dispatcher keyword, but do not narrow # support until that preparation is phase-aware too. del phase - if fwd.block_sparse_inputs is not None: - return False, "block_sparse_inputs are not supported." if q.device.type != "cuda": return False, "CUDA tensors are required." if not q.is_contiguous(): @@ -196,40 +287,9 @@ def _is_supported_with_reason( return False, "only fused QKV input is supported." if meta.is_cross: return False, "cross attention is not supported." - if meta.kv_cache_manager is None: - return False, "a KV cache manager is required." - if meta.kv_cache_block_offsets is None: - return False, "paged KV-cache block offsets are required." - if meta.host_kv_cache_pool_pointers is None: - return False, "KV-cache pool pointers are required." - if meta.host_kv_cache_pool_mapping is None: - return False, "KV-cache pool mapping is required." - if meta.kv_layout != "HND": - return False, "only HND KV-cache layout is supported." - kv_cache_manager = meta.kv_cache_manager - if isinstance(kv_cache_manager, KVCacheManagerV2): - if kv_cache_manager.enable_swa_scratch_reuse: - return False, "KVCacheManagerV2 SWA scratch reuse is not supported." - elif isinstance(kv_cache_manager, KVCacheManager): - if kv_cache_manager.num_pools != 1: - return False, "KVCacheManagerV1 with multiple memory pools is not supported." - pool_mapping = meta.host_kv_cache_pool_mapping - local_layer_idx = attn.local_layer_idx - num_local_layers = kv_cache_manager.num_local_layers - if ( - pool_mapping.ndim != 2 - or pool_mapping.shape[1] < 2 - or local_layer_idx is None - or local_layer_idx < 0 - or local_layer_idx >= pool_mapping.shape[0] - ): - return False, "KVCacheManagerV1 has an invalid layer-to-pool mapping." - pool_index = int(pool_mapping[local_layer_idx, 0]) - layer_idx_in_pool = int(pool_mapping[local_layer_idx, 1]) - if pool_index != 0 or not 0 <= layer_idx_in_pool < num_local_layers: - return False, "KVCacheManagerV1 has an invalid layer-to-pool mapping." - else: - return False, f"unsupported KV cache manager {type(kv_cache_manager).__name__}." + storage_reason = get_paged_kv_storage_unsupported_reason(attn, meta) + if storage_reason is not None: + return False, storage_reason output = fwd.output if output is None: @@ -242,38 +302,15 @@ def _is_supported_with_reason( if attn.sparse_params is not None: return False, "sparse attention is not supported." if ( - fwd.sparse_runtime_params.sparse_kv_indices is not None - or fwd.sparse_runtime_params.sparse_attn_indices is not None + sparse_runtime_params.sparse_kv_indices is not None + or sparse_runtime_params.sparse_attn_indices is not None ): return False, "sparse attention metadata is not supported." if meta.num_sparse_topk > 0: return False, "sparse attention metadata is not supported." - if meta.helix_position_offsets is not None: - return False, "Helix parallelism is not supported." - if fwd.relative_attention_bias is not None: - return False, "relative attention bias is not supported." - if fwd.attention_sinks is not None: - return False, "attention sinks are not supported." - if fwd.attention_mask_data is not None: - return False, "custom attention masks are not supported." - if fwd.enable_dsv4_epilogue_fusion: - return False, "DSv4 epilogue fusion is not supported." - if ( - fwd.sage_attn_num_elts_per_blk_q > 0 - or fwd.sage_attn_num_elts_per_blk_k > 0 - or fwd.sage_attn_num_elts_per_blk_v > 0 - ): - return False, "SageAttention is not supported." - - if meta.beam_width != 1: - return False, "beam search is not supported." - if ( - meta.is_spec_decoding_enabled - or meta.use_spec_decoding - or meta.is_spec_dec_tree - or meta.is_spec_dec_dynamic_tree - ): - return False, "speculative decoding is not supported by the initial adapter." + feature_reason = get_attention_feature_unsupported_reason(meta, fwd) + if feature_reason is not None: + return False, feature_reason try: mask_type = AttentionMaskType(fwd.mask_type) @@ -281,17 +318,9 @@ def _is_supported_with_reason( return False, "the attention mask is not causal or dense." if mask_type not in (AttentionMaskType.causal, AttentionMaskType.padding): return False, f"attention mask type {mask_type} is not supported." - - position_embedding_type = int(attn.position_embedding_type) - if position_embedding_type in (4, 5, 6, 7, 10): - return False, f"position embedding type {position_embedding_type} is not supported." - - try: - quant_mode = QuantMode(attn.quant_mode) - except (TypeError, ValueError): - return False, "invalid KV-cache quantization mode." - if quant_mode.has_kv_cache_quant(): - return False, "quantized KV cache is not supported by the initial adapter." + policy_reason = get_paged_kv_policy_unsupported_reason(attn, meta) + if policy_reason is not None: + return False, policy_reason input_type = fwd.attention_input_type if input_type not in ( @@ -1040,9 +1069,7 @@ def run_generation(self, params: FmhaParams) -> None: attn = params.attn meta = params.meta fwd = params.fwd - rope_params = attn.rope_params batch_size = params.batch_size - attention_chunk_size = attn.attention_chunk_size or 0 ( q_processed, kv_pool, @@ -1056,52 +1083,7 @@ def run_generation(self, params: FmhaParams) -> None: _max_kv_len, window_left, is_multi_token_gen, - ) = thop.trtllm_gen_generation_preprocess( - params.qkv_input, - params.workspace, - params.sequence_lengths, - params.spec_decoding_generation_lengths, - params.spec_decoding_position_offsets, - meta.kv_cache_block_offsets, - meta.host_kv_cache_pool_pointers, - meta.host_kv_cache_pool_mapping, - fwd.kv_scale_orig_quant, - fwd.kv_scale_quant_orig, - fwd.out_scale, - attn.rotary_inv_freq, - attn.rotary_cos_sin, - fwd.mrope_position_deltas, - attn.local_layer_idx, - params.seq_offset, - attn.num_heads, - attn.num_kv_heads, - attn.head_dim, - params.tokens_per_block, - attn.quant_mode, - params.max_attention_window_size, - params.cyclic_attention_window_size, - params.num_tokens, - batch_size, - params.input_seq_length, - params.max_past_kv_length, - rope_params.dim, - rope_params.theta, - int(rope_params.scale_type), - rope_params.scale, - rope_params.max_positions, - attn.position_embedding_type, - self._get_bmm1_scale(attn), - 1.0, - False, - attn.predicted_tokens_per_seq, - attention_chunk_size, - self._multi_processor_count, - params.total_num_blocks, - params.kv_factor, - True, - False, - skip_fmha_workspace=True, - ) + ) = self._run_generation_preprocess(params) if fmha_workspace.numel() != 0: raise RuntimeError("PrimTS generation preprocessing returned an FMHA workspace.") if is_multi_token_gen: diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py index 5b99fdbcb931..b8596041572a 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py @@ -13,8 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""TRT-LLM FMHA adapter for the vendored PrimTS block-sparse kernels.""" + import math -from dataclasses import dataclass +from dataclasses import dataclass, fields from typing import TYPE_CHECKING, Literal, cast import torch @@ -28,14 +30,16 @@ _validate_block_sparse_static_profile, ) from tensorrt_llm._torch.attention.backends.sparse.params import BlockSparseForwardInputs -from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 -from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm.logger import logger -from tensorrt_llm.quantization.mode import QuantMode from .interface import FmhaPhase from .phased import FmhaParams -from .prims_ts import PrimsTSFmha +from .prims_ts import ( + PrimsTSFmha, + get_attention_feature_unsupported_reason, + get_paged_kv_policy_unsupported_reason, + get_paged_kv_storage_unsupported_reason, +) from .utils import get_kv_page_offset if TYPE_CHECKING: @@ -48,10 +52,24 @@ TrtllmAttentionMetadata, ) +try: + from tensorrt_llm._torch.attention.backends.prims_ts import ( + BlockSparsePagedTSWrapper as _BlockSparsePagedTSWrapper, + ) + from tensorrt_llm._torch.attention.backends.prims_ts import ( + BlockSparseTSWrapper as _BlockSparseTSWrapper, + ) +except (ImportError, OSError): + _BlockSparseTSWrapper = _BlockSparsePagedTSWrapper = None + @dataclass(frozen=True, slots=True) class _BlockSparsePlanKey: - """Static wrapper profile shared by compatible attention layers.""" + """Static wrapper profile shared by compatible attention layers. + + The key is the single description of a plan: support checks validate it + against the kernel library and the wrapper cache plans from it. + """ device: torch.device batch_size: int @@ -70,19 +88,107 @@ class _BlockSparsePlanKey: sparse_format: Literal["bsr", "bitmask"] use_proxy_routes: bool + def unsupported_reason(self) -> str | None: + try: + _validate_block_sparse_static_profile( + batch_size=self.batch_size, + seq_len_q=self.seq_len_q, + seq_len_kv=self.kv_capacity, + num_qo_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_dim=self.head_dim, + q_block_size=self.q_block_size, + kv_block_size=self.kv_block_size, + use_kv_valid_bits=self.use_kv_valid_bits, + mask_type=self.mask_type, + q_dtype=self.dtype, + kv_dtype=self.dtype, + output_dtype=self.dtype, + max_blocks_per_row=self.max_blocks_per_row, + page_size=self.page_size, + ) + except (ValueError, NotImplementedError, OverflowError) as error: + return str(error) + return None + + def plan(self) -> "BlockSparseTSWrapper | BlockSparsePagedTSWrapper": + paged = self.page_size is not None + wrapper_type = _BlockSparsePagedTSWrapper if paged else _BlockSparseTSWrapper + assert wrapper_type is not None + wrapper = wrapper_type() + plan_args = ( + self.batch_size, + self.seq_len_q, + self.kv_capacity, + self.num_heads, + self.num_kv_heads, + self.head_dim, + self.q_block_size, + self.kv_block_size, + ) + plan_kwargs = { + "device": self.device, + "max_blocks_per_row": self.max_blocks_per_row, + "use_kv_valid_bits": self.use_kv_valid_bits, + "mask_type": self.mask_type, + "q_data_type": self.dtype, + "kv_data_type": self.dtype, + "o_data_type": self.dtype, + } + if paged: + plan_args += (self.page_size,) + else: + plan_kwargs.update( + sparse_format=self.sparse_format, + use_proxy_routes=self.use_proxy_routes, + ) + wrapper.plan(*plan_args, **plan_kwargs) + return wrapper + + +def _get_block_sparse_inputs( + forward_args: AttentionForwardArgs, +) -> BlockSparseForwardInputs | None: + return forward_args.sparse_runtime_params.block_sparse_inputs -try: - from tensorrt_llm._torch.attention.backends.prims_ts import ( - BlockSparsePagedTSWrapper as _BlockSparsePagedTSWrapper, - ) - from tensorrt_llm._torch.attention.backends.prims_ts import ( - BlockSparseTSWrapper as _BlockSparseTSWrapper, - ) -except (ImportError, OSError): - _BlockSparseTSWrapper = _BlockSparsePagedTSWrapper = None + +def _has_other_sparse_runtime(forward_args: AttentionForwardArgs) -> bool: + """Whether the runtime carrier holds any sparse state besides block-sparse routes.""" + params = forward_args.sparse_runtime_params + for field in fields(params): + if field.name == "block_sparse_inputs": + continue + value = getattr(params, field.name) + if isinstance(value, torch.Tensor) or (value is not None and value != 0): + return True + return False + + +def _route_batch_size(inputs: BlockSparseForwardInputs) -> int: + routes = inputs.block_indptr if inputs.sparse_format == "bsr" else inputs.exact_block_bits + return int(routes.shape[0]) + + +def _uniform_seq_len_q( + q: torch.Tensor, + metadata: "TrtllmAttentionMetadata", + batch_size: int, +) -> int | None: + """Return the fixed per-request query length, or ``None`` if the batch is ragged.""" + seq_lens = metadata.seq_lens + if batch_size <= 0 or q.shape[0] % batch_size: + return None + if seq_lens is None or seq_lens.numel() < batch_size: + return None + seq_len_q = int(q.shape[0]) // batch_size + if not bool(seq_lens[:batch_size].eq(seq_len_q).all()): + return None + return seq_len_q class PrimsTSBlockSparseFmha(PrimsTSFmha): + """Contiguous context and fixed-Q paged generation block-sparse FMHA.""" + def __init__(self, attn: "TrtllmAttention") -> None: super().__init__(attn) self.bind_plan_cache({}) @@ -139,20 +245,54 @@ def _is_supported_with_reason( *, phase: FmhaPhase | None = None, ) -> tuple[bool, str]: - if getattr(metadata, "kv_cache_manager", None) is None: - reason = ( - "contiguous block-sparse attention only supports the context phase" - if phase not in (None, FmhaPhase.CONTEXT) - else self._contiguous_unsupported_reason(q, k, v, metadata, forward_args) - ) - else: - reason = ( - "paged block-sparse attention only supports the generation phase" - if phase not in (None, FmhaPhase.GENERATION) - else self._paged_unsupported_reason(q, metadata, forward_args) - ) + reason = self._common_unsupported_reason(metadata, forward_args) + if reason is None: + paged = metadata.kv_cache_manager is not None + expected_phase = FmhaPhase.GENERATION if paged else FmhaPhase.CONTEXT + if phase not in (None, expected_phase): + storage = "paged" if paged else "contiguous" + reason = ( + f"{storage} block-sparse attention only supports the " + f"{expected_phase.name.lower()} phase" + ) + elif paged: + reason = self._paged_unsupported_reason(q, metadata, forward_args) + else: + reason = self._contiguous_unsupported_reason(q, k, v, metadata, forward_args) return reason is None, reason or "" + def _common_unsupported_reason( + self, + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + ) -> str | None: + """Gates shared by the contiguous and paged block-sparse paths.""" + if _get_block_sparse_inputs(forward_args) is None: + return "block-sparse forward inputs are required" + if metadata.is_cross: + return "cross attention is not supported" + if self.attn.is_mla_enable: + return "MLA is not supported" + if metadata.num_sparse_topk > 0 or _has_other_sparse_runtime(forward_args): + return "legacy sparse attention cannot be combined with block-sparse inputs" + feature_reason = get_attention_feature_unsupported_reason(metadata, forward_args) + if feature_reason is not None: + return feature_reason + if forward_args.softmax_stats_tensor is not None: + return "softmax statistics output is not supported" + if ( + forward_args.output_sf is not None + or forward_args.out_scale is not None + or forward_args.out_scale_sf is not None + ): + return "quantized output is not supported" + if forward_args.attention_mask not in ( + PredefinedAttentionMask.FULL, + PredefinedAttentionMask.CAUSAL, + ): + return "only full and causal masks are supported" + return None + def _make_plan_key( self, q: torch.Tensor, @@ -189,197 +329,47 @@ def _make_plan_key( def _get_or_plan_wrapper( self, key: _BlockSparsePlanKey, - *, - paged: bool, ) -> "BlockSparseTSWrapper | BlockSparsePagedTSWrapper": - cache = self._paged_wrappers if paged else self._contiguous_wrappers + cache = self._paged_wrappers if key.page_size is not None else self._contiguous_wrappers wrapper = cache.get(key) - if wrapper is not None: - return wrapper - wrapper_type = _BlockSparsePagedTSWrapper if paged else _BlockSparseTSWrapper - assert wrapper_type is not None - wrapper = wrapper_type() - plan_args = ( - key.batch_size, - key.seq_len_q, - key.kv_capacity, - key.num_heads, - key.num_kv_heads, - key.head_dim, - key.q_block_size, - key.kv_block_size, - ) - if paged: - plan_args += (key.page_size,) - plan_kwargs = { - "device": key.device, - "max_blocks_per_row": key.max_blocks_per_row, - "use_kv_valid_bits": key.use_kv_valid_bits, - "mask_type": key.mask_type, - "q_data_type": key.dtype, - "kv_data_type": key.dtype, - "o_data_type": key.dtype, - } - if not paged: - plan_kwargs.update( - sparse_format=key.sparse_format, - use_proxy_routes=key.use_proxy_routes, - ) - wrapper.plan(*plan_args, **plan_kwargs) - cache[key] = wrapper + if wrapper is None: + wrapper = key.plan() + cache[key] = wrapper return wrapper - def _paged_storage_unsupported_reason( - self, - metadata: "TrtllmAttentionMetadata", - ) -> str | None: - if metadata.kv_layout != "HND": - return "only HND KV-cache layout is supported" - if metadata.host_kv_cache_pool_pointers is None: - return "KV-cache pool pointers are required" - pool_mapping = metadata.host_kv_cache_pool_mapping - if pool_mapping is None: - return "KV-cache pool mapping is required" - - manager = metadata.kv_cache_manager - if isinstance(manager, KVCacheManagerV2): - if manager.enable_swa_scratch_reuse: - return "KVCacheManagerV2 SWA scratch reuse is not supported" - elif isinstance(manager, KVCacheManager): - if manager.num_pools != 1: - return "KVCacheManagerV1 with multiple memory pools is not supported" - local_layer_idx = self.attn.local_layer_idx - if ( - pool_mapping.ndim != 2 - or pool_mapping.shape[1] < 2 - or local_layer_idx is None - or not 0 <= local_layer_idx < pool_mapping.shape[0] - ): - return "KVCacheManagerV1 has an invalid layer-to-pool mapping" - pool_index = int(pool_mapping[local_layer_idx, 0]) - layer_index_in_pool = int(pool_mapping[local_layer_idx, 1]) - if pool_index != 0 or not 0 <= layer_index_in_pool < manager.num_local_layers: - return "KVCacheManagerV1 has an invalid layer-to-pool mapping" - else: - return f"unsupported KV cache manager {type(manager).__name__}" - if metadata.tokens_per_block not in self.SUPPORTED_PAGE_SIZES: - return f"page size {metadata.tokens_per_block} is unsupported" - try: - if ( - get_kv_page_offset( - self.attn, - metadata, - 0, - cache=self._kv_page_offset_cache, - ) - is None - ): - return "the K-to-V page displacement could not be resolved" - except (AttributeError, IndexError, RuntimeError, TypeError, ValueError) as error: - return f"invalid KV-cache storage metadata: {error}" - return None - - @staticmethod - def _has_legacy_sparse_prediction(forward_args: AttentionForwardArgs) -> bool: - prediction = forward_args.sparse_runtime_params - return any( - ( - prediction.sparse_kv_indices is not None, - prediction.sparse_kv_offsets is not None, - prediction.sparse_attn_indices is not None, - prediction.sparse_attn_offsets is not None, - prediction.sparse_attn_kv_lens is not None, - bool(prediction.sparse_attn_indices_block_size), - prediction.aux_kv_cache_pool_ptr is not None, - bool(prediction.threshold_scale_factor_prefill), - bool(prediction.threshold_scale_factor_decode), - ) - ) - - def _common_unsupported_reason( + def _contiguous_unsupported_reason( self, + q: torch.Tensor, + k: torch.Tensor | None, + v: torch.Tensor | None, metadata: "TrtllmAttentionMetadata", forward_args: AttentionForwardArgs, ) -> str | None: - if forward_args.block_sparse_inputs is None: - return "block-sparse forward inputs are required" - if metadata.is_cross: - return "cross attention is not supported" - if self.attn.is_mla_enable: - return "MLA is not supported" - if getattr(metadata, "helix_position_offsets", None) is not None: - return "Helix parallelism is not supported" - if getattr(metadata, "num_sparse_topk", 0) > 0 or self._has_legacy_sparse_prediction( - forward_args - ): - return "legacy sparse attention cannot be combined with block-sparse inputs" - if forward_args.enable_dsv4_epilogue_fusion: - return "DSv4 epilogue fusion is not supported" - if forward_args.sage_attn_qk_int8 or any( - getattr(forward_args, name) > 0 - for name in ( - "sage_attn_num_elts_per_blk_q", - "sage_attn_num_elts_per_blk_k", - "sage_attn_num_elts_per_blk_v", - ) - ): - return "SageAttention is not supported" - if forward_args.softmax_stats_tensor is not None: - return "softmax statistics output is not supported" - if ( - forward_args.output_sf is not None - or forward_args.out_scale is not None - or forward_args.out_scale_sf is not None - ): - return "quantized output is not supported" - if ( - forward_args.attention_mask_data is not None - or forward_args.relative_attention_bias is not None - or forward_args.attention_sinks is not None - ): - return "custom attention masks, bias, and sinks are not supported" - if forward_args.attention_mask not in ( - PredefinedAttentionMask.FULL, - PredefinedAttentionMask.CAUSAL, - ): - return "only full and causal masks are supported" - return None - - def _static_profile_unsupported_reason( - self, - q: torch.Tensor, - inputs: BlockSparseForwardInputs, - *, - batch_size: int, - seq_len_q: int, - seq_len_kv: int, - page_size: int | None, - mask_type: Literal["dense", "causal"], - ) -> str | None: - max_blocks_per_row = inputs.max_blocks_per_row - if max_blocks_per_row is None: - max_blocks_per_row = math.ceil(seq_len_kv / inputs.kv_block_size) - try: - _validate_block_sparse_static_profile( - batch_size=batch_size, - seq_len_q=seq_len_q, - seq_len_kv=seq_len_kv, - num_qo_heads=self.attn.num_heads, - num_kv_heads=self.attn.num_kv_heads, - head_dim=self.attn.head_dim, - q_block_size=inputs.q_block_size, - kv_block_size=inputs.kv_block_size, - use_kv_valid_bits=inputs.kv_valid_bits is not None, - mask_type=mask_type, - q_dtype=q.dtype, - kv_dtype=q.dtype, - output_dtype=q.dtype, - max_blocks_per_row=max_blocks_per_row, - page_size=page_size, - ) - except (ValueError, NotImplementedError, OverflowError) as error: - return str(error) - return None + if forward_args.is_fused_qkv or k is None or v is None: + return "contiguous block-sparse attention requires separate Q, K, and V" + if self.attn.position_embedding_type != 0 or forward_args.mrope_position_deltas is not None: + return "contiguous Q/K/V must have position embedding applied before attention" + if forward_args.cu_q_seqlens is not None or forward_args.cu_kv_seqlens is not None: + return "packed variable-length Q/KV inputs are not supported" + inputs = _get_block_sparse_inputs(forward_args) + assert inputs is not None + mask_type = self._get_prims_mask_type(forward_args) + if inputs.use_proxy_routes and mask_type != "dense": + return "block-sparse proxy routes require mask_type='dense'" + batch_size = _route_batch_size(inputs) + seq_len_q = _uniform_seq_len_q(q, metadata, batch_size) + if seq_len_q is None or k.shape[0] % batch_size: + return "query and KV token counts must be batch-uniform over the route batch size" + key = self._make_plan_key( + q, + inputs, + batch_size=batch_size, + seq_len_q=seq_len_q, + kv_capacity=int(k.shape[0]) // batch_size, + page_size=None, + mask_type=mask_type, + ) + return key.unsupported_reason() def _paged_unsupported_reason( self, @@ -387,14 +377,10 @@ def _paged_unsupported_reason( metadata: "TrtllmAttentionMetadata", forward_args: AttentionForwardArgs, ) -> str | None: - inputs = forward_args.block_sparse_inputs - if inputs is None: - return "block-sparse forward inputs are required" + inputs = _get_block_sparse_inputs(forward_args) + assert inputs is not None if inputs.sparse_format != "bsr" or inputs.use_proxy_routes: return "paged block-sparse attention only supports BSR exact routes" - common_reason = self._common_unsupported_reason(metadata, forward_args) - if common_reason is not None: - return common_reason if not forward_args.is_fused_qkv: return "paged block-sparse attention requires fused QKV input" if ( @@ -402,48 +388,37 @@ def _paged_unsupported_reason( or metadata.num_contexts != 0 ): return "only generation-only paged requests are supported" + reason = get_paged_kv_storage_unsupported_reason( + self.attn, metadata + ) or get_paged_kv_policy_unsupported_reason(self.attn, metadata) + if reason is not None: + return reason + if metadata.tokens_per_block not in self.SUPPORTED_PAGE_SIZES: + return f"page size {metadata.tokens_per_block} is unsupported" + if self.attn.attention_chunk_size: + return "chunked attention is not supported" + if get_kv_page_offset(self.attn, metadata, 0, cache=self._kv_page_offset_cache) is None: + return "the K-to-V page displacement could not be resolved" + batch_size = int(metadata.num_generations) - if batch_size <= 0 or q.shape[0] % batch_size: - return "query tokens must be uniformly divisible across generation requests" - seq_len_q = int(q.shape[0]) // batch_size - if not bool(metadata.seq_lens[:batch_size].eq(seq_len_q).all()): + seq_len_q = _uniform_seq_len_q(q, metadata, batch_size) + if seq_len_q is None: return "query lengths must be batch-uniform and match the fixed query shape" - paged_kv_reason = self._paged_storage_unsupported_reason(metadata) - if paged_kv_reason is not None: - return paged_kv_reason block_tables = metadata.kv_cache_block_offsets - if block_tables is None or block_tables.shape[1] < batch_size: + if block_tables.shape[1] < batch_size: return "paged KV-cache block offsets must cover the generation batch" - if metadata.beam_width != 1: - return "beam search is not supported" - if any( - getattr(metadata, name, False) - for name in ( - "is_spec_decoding_enabled", - "use_spec_decoding", - "is_spec_dec_tree", - "is_spec_dec_dynamic_tree", - ) - ): - return "speculative decoding is not supported" - if self.attn.attention_chunk_size: - return "chunked attention is not supported" - if self.attn.position_embedding_type in (4, 5, 6, 7, 10): - return f"position embedding type {self.attn.position_embedding_type} is not supported" page_size = int(metadata.tokens_per_block) - max_seq_len_kv = int(block_tables.shape[-1]) * page_size + kv_capacity = int(block_tables.shape[-1]) * page_size logical_max_seq_len = int(metadata.max_seq_len) - if logical_max_seq_len > max_seq_len_kv: + if logical_max_seq_len > kv_capacity: return "logical maximum sequence length must fit the page-table capacity" attention_window_size = forward_args.attention_window_size if ( attention_window_size is None or attention_window_size < logical_max_seq_len - or attention_window_size > max_seq_len_kv + or attention_window_size > kv_capacity ): return "attention window must fit the non-cyclic page-table capacity" - if QuantMode(self.attn.quant_mode).has_kv_cache_quant(): - return "quantized KV cache is not supported" host_seq_lens = metadata.kv_lens_runtime[:batch_size] min_seq_len_kv = int(host_seq_lens.min()) if min_seq_len_kv <= 0: @@ -453,36 +428,16 @@ def _paged_unsupported_reason( return "causal KV lengths must be at least the fixed query length" if int(host_seq_lens.max()) > logical_max_seq_len: return "an active KV length exceeds the logical maximum sequence length" - return self._static_profile_unsupported_reason( + key = self._make_plan_key( q, inputs, batch_size=batch_size, seq_len_q=seq_len_q, - seq_len_kv=max_seq_len_kv, + kv_capacity=kv_capacity, page_size=page_size, mask_type=mask_type, ) - - def _ensure_preprocess_workspace( - self, - q: torch.Tensor, - workspace: torch.Tensor, - *, - batch_size: int, - ) -> None: - layout = self._get_generation_workspace_layout( - q.dtype, - batch_size, - int(q.shape[0]), - ) - required_bytes = int(layout["total_size"]) - available_bytes = workspace.numel() * workspace.element_size() - if available_bytes < required_bytes: - if torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "TRT-LLM QKV preprocessing workspace must be sized before CUDA Graph capture" - ) - workspace.resize_((math.ceil(required_bytes / workspace.element_size()),)) + return key.unsupported_reason() def prepare_workspace( self, @@ -495,11 +450,22 @@ def prepare_workspace( ) -> None: del k, v, forward_args with torch.cuda.device(q.device): - self._ensure_preprocess_workspace( - q, - workspace, - batch_size=int(metadata.num_generations), - ) + # Contiguous requests run without a KV cache and never touch the + # generation preprocessing workspace. + if metadata.kv_cache_manager is not None: + layout = self._get_generation_workspace_layout( + q.dtype, + int(metadata.num_generations), + int(q.shape[0]), + ) + required_bytes = int(layout["total_size"]) + if workspace.numel() * workspace.element_size() < required_bytes: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "TRT-LLM QKV preprocessing workspace must be sized before " + "CUDA Graph capture" + ) + workspace.resize_((math.ceil(required_bytes / workspace.element_size()),)) if self._multi_processor_count is None: if torch.cuda.is_current_stream_capturing(): raise RuntimeError("GPU properties must be prepared before CUDA Graph capture") @@ -514,30 +480,16 @@ def run_generation(self, params: FmhaParams) -> None: assert q is not None and output_buffer is not None and sequence_lengths is not None metadata = params.meta forward_args = params.fwd - inputs = forward_args.block_sparse_inputs + inputs = _get_block_sparse_inputs(forward_args) assert inputs is not None batch_size = params.num_requests seq_len_q = params.input_seq_length page_size = params.tokens_per_block block_offsets = metadata.kv_cache_block_offsets assert block_offsets is not None - max_seq_len_kv = int(block_offsets.shape[-1]) * page_size - mask_type = self._get_prims_mask_type(forward_args) - assert self._multi_processor_count is not None - ( - q_processed, - kv_pool, - block_tables, - _kv_scale_pool, - _bmm1_scale, - _bmm2_scale, - fmha_workspace, - _cu_seqlens, - _max_q_len, - _max_kv_len, - _window_left, - _is_multi_token_gen, - ) = self._run_generation_preprocess(params) + preprocess = self._run_generation_preprocess(params) + q_processed, kv_pool, block_tables = preprocess[:3] + fmha_workspace = preprocess[6] if fmha_workspace.numel() != 0: raise RuntimeError("PrimTS block-sparse preprocessing returned an FMHA workspace.") if q_processed is None or kv_pool is None or block_tables is None: @@ -551,102 +503,32 @@ def run_generation(self, params: FmhaParams) -> None: if kv_page_offset is None: raise RuntimeError("PrimTS could not resolve the K-to-V page displacement") k_cache, v_cache = self._standard_kv_views(kv_pool, kv_page_offset) - fixed_block_tables = self._get_fixed_block_tables( - block_tables, - batch_size, - ) - seq_lens = self._get_sequence_lengths(sequence_lengths, batch_size) query = q_processed.view( batch_size, seq_len_q, self.attn.num_heads, self.attn.head_dim, ) - output = output_buffer.view_as(query) key = self._make_plan_key( query, inputs, batch_size=batch_size, seq_len_q=seq_len_q, - kv_capacity=max_seq_len_kv, + kv_capacity=int(block_offsets.shape[-1]) * page_size, page_size=page_size, - mask_type=mask_type, - ) - wrapper = cast( - "BlockSparsePagedTSWrapper", - self._get_or_plan_wrapper(key, paged=True), + mask_type=self._get_prims_mask_type(forward_args), ) + wrapper = cast("BlockSparsePagedTSWrapper", self._get_or_plan_wrapper(key)) wrapper.run( query, (k_cache, v_cache), - block_tables=fixed_block_tables, - seq_lens_kv=seq_lens, + block_tables=self._get_fixed_block_tables(block_tables, batch_size), + seq_lens_kv=self._get_sequence_lengths(sequence_lengths, batch_size), block_indptr=inputs.block_indptr, block_indices=inputs.block_indices, kv_valid_bits=inputs.kv_valid_bits, sm_scale=self._get_bmm1_scale(self.attn), - out=output, - ) - - def _contiguous_views( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - route_tensor: torch.Tensor, - output: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: - batch_size = int(route_tensor.shape[0]) - if batch_size <= 0 or q.shape[0] % batch_size or k.shape[0] % batch_size: - raise ValueError("flat Q and K token counts must be divisible by route batch size") - seq_len_q = int(q.shape[0]) // batch_size - seq_len_kv = int(k.shape[0]) // batch_size - return ( - q.view(batch_size, seq_len_q, self.attn.num_heads, self.attn.head_dim), - k.view(batch_size, seq_len_kv, self.attn.num_kv_heads, self.attn.head_dim), - v.view(batch_size, seq_len_kv, self.attn.num_kv_heads, self.attn.head_dim), - output.view(batch_size, seq_len_q, self.attn.num_heads, self.attn.head_dim), - ) - - def _contiguous_unsupported_reason( - self, - q: torch.Tensor, - k: torch.Tensor | None, - v: torch.Tensor | None, - metadata: "TrtllmAttentionMetadata", - forward_args: AttentionForwardArgs, - ) -> str | None: - common_reason = self._common_unsupported_reason(metadata, forward_args) - if common_reason is not None: - return common_reason - if self.attn.position_embedding_type != 0 or forward_args.mrope_position_deltas is not None: - return "contiguous Q/K/V must have position embedding applied before attention" - if forward_args.is_fused_qkv or k is None or v is None: - return "contiguous block-sparse attention requires separate Q, K, and V" - if forward_args.cu_q_seqlens is not None or forward_args.cu_kv_seqlens is not None: - return "packed variable-length Q/KV inputs are not supported" - inputs = forward_args.block_sparse_inputs - assert inputs is not None - mask_type = self._get_prims_mask_type(forward_args) - if inputs.use_proxy_routes and mask_type != "dense": - return "block-sparse proxy routes require mask_type='dense'" - route_tensor = ( - inputs.block_indptr if inputs.sparse_format == "bsr" else inputs.exact_block_bits - ) - batch_size = int(route_tensor.shape[0]) - if batch_size <= 0 or q.shape[0] % batch_size or k.shape[0] % batch_size: - return "flat Q and K token counts must be divisible by the route batch size" - seq_len_q = int(q.shape[0]) // batch_size - if not bool(metadata.seq_lens[:batch_size].eq(seq_len_q).all()): - return "query lengths must be batch-uniform and match the fixed query shape" - return self._static_profile_unsupported_reason( - q, - inputs, - batch_size=batch_size, - seq_len_q=seq_len_q, - seq_len_kv=int(k.shape[0]) // batch_size, - page_size=None, - mask_type=mask_type, + out=output_buffer.view_as(query), ) def _forward_contiguous( @@ -656,34 +538,27 @@ def _forward_contiguous( v: torch.Tensor, forward_args: AttentionForwardArgs, ) -> None: - inputs = forward_args.block_sparse_inputs + inputs = _get_block_sparse_inputs(forward_args) assert inputs is not None assert forward_args.output is not None - q_view, k_view, v_view, out_view = self._contiguous_views( - q, - k, - v, - inputs.block_indptr if inputs.sparse_format == "bsr" else inputs.exact_block_bits, - forward_args.output, - ) - mask_type = self._get_prims_mask_type(forward_args) + batch_size = _route_batch_size(inputs) + query = q.view(batch_size, -1, self.attn.num_heads, self.attn.head_dim) + key_states = k.view(batch_size, -1, self.attn.num_kv_heads, self.attn.head_dim) + value_states = v.view_as(key_states) key = self._make_plan_key( - q_view, + query, inputs, - batch_size=int(q_view.shape[0]), - seq_len_q=int(q_view.shape[1]), - kv_capacity=int(k_view.shape[1]), + batch_size=batch_size, + seq_len_q=int(query.shape[1]), + kv_capacity=int(key_states.shape[1]), page_size=None, - mask_type=mask_type, - ) - wrapper = cast( - "BlockSparseTSWrapper", - self._get_or_plan_wrapper(key, paged=False), + mask_type=self._get_prims_mask_type(forward_args), ) + wrapper = cast("BlockSparseTSWrapper", self._get_or_plan_wrapper(key)) wrapper.run( - q_view, - k_view, - v_view, + query, + key_states, + value_states, block_indptr=inputs.block_indptr, block_indices=inputs.block_indices, exact_block_bits=inputs.exact_block_bits, @@ -691,7 +566,7 @@ def _forward_contiguous( v_summary=inputs.v_summary, kv_valid_bits=inputs.kv_valid_bits, sm_scale=self._get_bmm1_scale(self.attn), - out=out_view, + out=forward_args.output.view_as(query), ) def forward( diff --git a/tensorrt_llm/_torch/attention/backends/fmha/triton_custom_mask.py b/tensorrt_llm/_torch/attention/backends/fmha/triton_custom_mask.py index 1f2e5fbbbec9..61aef90d9d61 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/triton_custom_mask.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/triton_custom_mask.py @@ -105,7 +105,7 @@ def _check_support_with_reason( *, phase: Optional[FmhaPhase] = None, ) -> tuple[bool, str]: - if forward_args.block_sparse_inputs is not None: + if forward_args.sparse_runtime_params.block_sparse_inputs is not None: return False, "block_sparse_inputs are not supported." if phase != FmhaPhase.CONTEXT: return False, "Only context attention is supported." diff --git a/tensorrt_llm/_torch/attention/backends/interface.py b/tensorrt_llm/_torch/attention/backends/interface.py index 37144b0e3a96..7d11cdb28f8d 100644 --- a/tensorrt_llm/_torch/attention/backends/interface.py +++ b/tensorrt_llm/_torch/attention/backends/interface.py @@ -32,8 +32,8 @@ from ...pyexecutor.resource_manager import KVCacheManager from ...pyexecutor.trace_log_utils import log_tensor_size from ...utils import get_model_extra_attrs -from .sparse.params import (BlockSparseForwardInputs, SparseBackendForwardArgs, - SparseMetadataParams, SparseRuntimeParams) +from .sparse.params import (SparseBackendForwardArgs, SparseMetadataParams, + SparseRuntimeParams) try: # Transformers v5 @@ -964,7 +964,6 @@ class AttentionForwardArgs: # Optional normalized diffusion timestep for timestep-varying sparse attention. timestep: Optional[torch.Tensor] = None - block_sparse_inputs: Optional[BlockSparseForwardInputs] = None sparse_backend_args: Optional[SparseBackendForwardArgs] = None sparse_runtime_params: SparseRuntimeParams = field( default_factory=SparseRuntimeParams) diff --git a/tensorrt_llm/_torch/attention/backends/sparse/hooks.py b/tensorrt_llm/_torch/attention/backends/sparse/hooks.py index 53cabb9101b1..44437c0a9037 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/hooks.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/hooks.py @@ -12,6 +12,8 @@ from importlib import import_module from typing import TYPE_CHECKING, Optional +from .skip_softmax import SkipSoftmaxParams + if TYPE_CHECKING: import torch @@ -220,26 +222,36 @@ def prepare_sparse_runtime_params( backend: "TrtllmAttention", q: "torch.Tensor", k: Optional["torch.Tensor"], + v: Optional["torch.Tensor"], metadata: "AttentionMetadata", forward_args: "AttentionForwardArgs", ) -> "SparseRuntimeParams": - """Run backend prediction hooks and update attention-op parameters.""" - runtime_params = forward_args.sparse_runtime_params - if backend.sparse_params is None: - return runtime_params - + """Predict all sparse inputs for one attention call. + + Runs the ``sparse_kv_predict``, ``sparse_attn_predict`` and + ``block_sparse_attn_predict`` hooks once each and returns a new + ``SparseRuntimeParams`` built from ``forward_args.sparse_runtime_params`` + plus the hook results. Fields a backend writes into that carrier outside + the hooks, such as an auxiliary pool pointer, are carried over. SkipSoftmax + backends receive their threshold schedule last. + """ kv_indices, kv_offsets = backend.sparse_kv_predict(q, k, metadata, forward_args) attn_indices, attn_offsets = backend.sparse_attn_predict(q, k, metadata, forward_args) - block_size = ( - backend.sparse_params.indices_block_size - if attn_indices is not None or attn_offsets is not None - else runtime_params.sparse_attn_indices_block_size - ) - return replace( - runtime_params, + block_sparse_inputs = backend.block_sparse_attn_predict(q, k, v, metadata, forward_args) + has_attn_indices = attn_indices is not None or attn_offsets is not None + sparse_params = backend.sparse_params + runtime_params = replace( + forward_args.sparse_runtime_params, sparse_kv_indices=kv_indices, sparse_kv_offsets=kv_offsets, sparse_attn_indices=attn_indices, sparse_attn_offsets=attn_offsets, - sparse_attn_indices_block_size=block_size, + sparse_attn_indices_block_size=sparse_params.indices_block_size if has_attn_indices else 0, + block_sparse_inputs=block_sparse_inputs, ) + if isinstance(sparse_params, SkipSoftmaxParams): + runtime_params = sparse_params.scheduler.get_runtime_params( + runtime_params=runtime_params, + timestep=forward_args.timestep, + ) + return runtime_params diff --git a/tensorrt_llm/_torch/attention/backends/sparse/params.py b/tensorrt_llm/_torch/attention/backends/sparse/params.py index c99061926435..8e6829f8c879 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/params.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/params.py @@ -36,6 +36,9 @@ class SparseBackendForwardArgs: # Shared by algorithms that accept precomputed top-k indices. topk_indices: Optional[torch.Tensor] = None + # Complete block-sparse routing payload predicted by the module before the + # core forward; the default backend hook hands it through unchanged. + block_sparse_inputs: Optional["BlockSparseForwardInputs"] = None @dataclass(frozen=True, slots=True) @@ -82,7 +85,7 @@ def use_proxy_routes(self) -> bool: @dataclass(kw_only=True, slots=True) class SparseRuntimeParams: - """Flat optional sparse inputs passed from a backend to ``AttentionOp``.""" + """Complete per-attention sparse runtime state consumed by FMHA/``AttentionOp``.""" # Sparse index inputs shared by multiple algorithms. sparse_kv_indices: Optional[torch.Tensor] = None @@ -99,6 +102,7 @@ class SparseRuntimeParams: threshold_scale_factor_prefill: float = 0.0 # SkipSoftmax decode threshold; diffusion models leave it at zero. threshold_scale_factor_decode: float = 0.0 + block_sparse_inputs: Optional[BlockSparseForwardInputs] = None __all__ = [ diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 5a2d7ea3912b..d5e273b06483 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -17,7 +17,7 @@ import math import os import weakref -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from typing import TYPE_CHECKING, List, Optional, Tuple import torch @@ -46,8 +46,7 @@ PredefinedAttentionMask, RopeParams, merge_attention_forward_args) from .sparse.hooks import prepare_sparse_runtime_params -from .sparse.params import SparseParams -from .sparse.skip_softmax import SkipSoftmaxParams +from .sparse.params import BlockSparseForwardInputs, SparseParams _SKIP_CORRECTION_SUPPORTED_SMS = frozenset((100, 103)) @@ -1781,7 +1780,6 @@ def forward( ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: """Execute the TRTLLM attention backend.""" forward_args = merge_attention_forward_args(forward_args, kwargs) - has_block_sparse_inputs = forward_args.block_sparse_inputs is not None assert isinstance( metadata, TrtllmAttentionMetadata, @@ -1929,8 +1927,13 @@ def forward( seq_start=num_ctx, ) - forward_args.sparse_runtime_params = prepare_sparse_runtime_params( - self, q, k, metadata, forward_args) + sparse_runtime_params = prepare_sparse_runtime_params( + self, q, k, v, metadata, forward_args) + forward_args = replace( + forward_args, + sparse_runtime_params=sparse_runtime_params, + ) + has_block_sparse_inputs = sparse_runtime_params.block_sparse_inputs is not None # Compute FlashMLA tile-scheduler metadata once per forward pass. # The flag is invalidated whenever FlashMLA inputs change. The metadata @@ -1983,14 +1986,11 @@ def forward( assert k.shape[1] == kv_hidden_size assert v.shape[1] == kv_hidden_size num_tokens = q.shape[0] - if k is not None and not metadata.is_cross: - if has_block_sparse_inputs: - assert v is not None and v.shape[0] == k.shape[0] - else: - assert k.shape[0] == num_tokens - assert v.shape[0] == num_tokens + if k is not None and not metadata.is_cross and not has_block_sparse_inputs: + assert k.shape[0] == num_tokens + assert v.shape[0] == num_tokens else: - sparse_attn_indices = forward_args.sparse_runtime_params.sparse_attn_indices + sparse_attn_indices = sparse_runtime_params.sparse_attn_indices is_sparse_attn = sparse_attn_indices is not None and sparse_attn_indices.numel( ) > 0 if attention_input_type == AttentionInputType.context_only and is_sparse_attn: @@ -2063,14 +2063,6 @@ def forward( if forward_args.kv_scale_quant_orig is None: forward_args.kv_scale_quant_orig = self.kv_scale_quant_orig - sparse_params = self.sparse_params - if isinstance(sparse_params, SkipSoftmaxParams): - forward_args.sparse_runtime_params = ( - sparse_params.scheduler.get_runtime_params( - runtime_params=forward_args.sparse_runtime_params, - timestep=forward_args.timestep, - )) - # max_context_q_len_override is only set when encoder CUDA graphs are enabled. if metadata.max_context_q_len_override is not None: assert metadata.is_cuda_graph @@ -2290,6 +2282,26 @@ def sparse_kv_predict( """Predict sparse KV indices when required by an algorithm.""" return None, None + def block_sparse_attn_predict( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: TrtllmAttentionMetadata, + forward_args: AttentionForwardArgs, + ) -> Optional[BlockSparseForwardInputs]: + """Predict the block-sparse routing payload for one attention call. + + The default hands through routes that the attention module predicted + before the core forward via ``sparse_backend_args``. Algorithms that + predict inside the backend override this method and return ``None`` + for dense phases. + """ + backend_args = forward_args.sparse_backend_args + if backend_args is None: + return None + return backend_args.block_sparse_inputs + def sparse_attn_predict( self, q: torch.Tensor, diff --git a/tests/unittest/_torch/attention/sparse/test_sparse_attention.py b/tests/unittest/_torch/attention/sparse/test_sparse_attention.py index 40a151daa252..8b012cac879b 100644 --- a/tests/unittest/_torch/attention/sparse/test_sparse_attention.py +++ b/tests/unittest/_torch/attention/sparse/test_sparse_attention.py @@ -21,11 +21,16 @@ """ from types import ModuleType -from unittest.mock import Mock +from unittest.mock import Mock, patch +import pytest import torch -from tensorrt_llm._torch.attention.backends.interface import AttentionForwardArgs +from tensorrt_llm._torch.attention.backends import trtllm as trtllm_backend +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + AttentionInputType, +) from tensorrt_llm._torch.attention.backends.sparse.hooks import ( AttentionSparseHooks, MLASparseHooks, @@ -35,8 +40,13 @@ register_attention_sparse_hooks, register_mla_sparse_hooks, ) -from tensorrt_llm._torch.attention.backends.sparse.params import SparseParams, SparseRuntimeParams -from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention +from tensorrt_llm._torch.attention.backends.sparse.params import ( + BlockSparseForwardInputs, + SparseBackendForwardArgs, + SparseParams, + SparseRuntimeParams, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata from tensorrt_llm._torch.attention.mla import MLA @@ -72,7 +82,7 @@ def test_prepare_sparse_runtime_params_from_predictions() -> None: ) runtime_params = prepare_sparse_runtime_params( - attention, torch.empty(0), None, None, forward_args + attention, torch.empty(0), None, None, None, forward_args ) assert runtime_params.sparse_kv_indices is attention._sparse_kv_indices @@ -144,12 +154,250 @@ def test_mla_backend_only_forward_uses_default_path() -> None: ) -def test_prepare_sparse_runtime_params_without_predictions() -> None: +@pytest.mark.parametrize( + "sparse_params", [None, _StubSparseParams()], ids=["dense_backend", "sparse_backend"] +) +def test_prepare_sparse_runtime_params_without_predictions(sparse_params) -> None: attention = TrtllmAttention.__new__(TrtllmAttention) - attention.sparse_params = _StubSparseParams() + attention.sparse_params = sparse_params runtime_params = prepare_sparse_runtime_params( - attention, torch.empty(0), None, None, AttentionForwardArgs() + attention, torch.empty(0), None, None, None, AttentionForwardArgs() ) assert runtime_params == SparseRuntimeParams() + + +def test_prepare_sparse_runtime_params_runs_index_hooks_once() -> None: + attention = _StaticPredictionAttention.__new__(_StaticPredictionAttention) + attention.sparse_params = _StubSparseParams() + sparse_kv_indices = torch.tensor([1], dtype=torch.int32) + sparse_kv_offsets = torch.tensor([0, 1], dtype=torch.int32) + sparse_attn_indices = torch.tensor([2], dtype=torch.int32) + sparse_attn_offsets = torch.tensor([0, 1], dtype=torch.int32) + attention.sparse_kv_predict = Mock(return_value=(sparse_kv_indices, sparse_kv_offsets)) + attention.sparse_attn_predict = Mock(return_value=(sparse_attn_indices, sparse_attn_offsets)) + q = torch.empty((1, 4)) + k = torch.empty((1, 4)) + v = torch.empty((1, 4)) + metadata = Mock() + caller_kv_lens = torch.tensor([3]) + forward_args = AttentionForwardArgs( + sparse_runtime_params=SparseRuntimeParams(sparse_attn_kv_lens=caller_kv_lens) + ) + + runtime_params = prepare_sparse_runtime_params(attention, q, k, v, metadata, forward_args) + + assert isinstance(runtime_params, SparseRuntimeParams) + assert runtime_params.block_sparse_inputs is None + assert runtime_params.sparse_kv_indices is sparse_kv_indices + assert runtime_params.sparse_kv_offsets is sparse_kv_offsets + assert runtime_params.sparse_attn_indices is sparse_attn_indices + assert runtime_params.sparse_attn_offsets is sparse_attn_offsets + assert runtime_params.sparse_attn_indices_block_size == 1 + assert runtime_params.sparse_attn_kv_lens is caller_kv_lens + attention.sparse_kv_predict.assert_called_once_with(q, k, metadata, forward_args) + attention.sparse_attn_predict.assert_called_once_with(q, k, metadata, forward_args) + + +def test_prepare_sparse_runtime_params_schedules_skip_softmax_thresholds() -> None: + from tensorrt_llm._torch.attention.backends.sparse.skip_softmax import SkipSoftmaxParams + + attention = TrtllmAttention.__new__(TrtllmAttention) + attention.sparse_params = SkipSoftmaxParams() + scheduler = attention.sparse_params.scheduler + timestep = torch.tensor(0.5) + forward_args = AttentionForwardArgs(timestep=timestep) + + with patch.object( + scheduler, "get_runtime_params", wraps=scheduler.get_runtime_params + ) as schedule: + runtime_params = prepare_sparse_runtime_params( + attention, torch.empty(0), None, None, None, forward_args + ) + + schedule.assert_called_once_with(runtime_params=SparseRuntimeParams(), timestep=timestep) + assert runtime_params == scheduler.get_runtime_params(timestep=timestep) + + +def _make_block_sparse_inputs() -> BlockSparseForwardInputs: + return BlockSparseForwardInputs( + q_block_size=64, + kv_block_size=64, + exact_block_bits=torch.zeros((1, 1), dtype=torch.int32), + ) + + +def test_block_sparse_attn_predict_hands_through_backend_args() -> None: + attention = TrtllmAttention.__new__(TrtllmAttention) + attention.sparse_params = None + block_sparse_inputs = _make_block_sparse_inputs() + forward_args = AttentionForwardArgs( + sparse_backend_args=SparseBackendForwardArgs(block_sparse_inputs=block_sparse_inputs) + ) + + runtime_params = prepare_sparse_runtime_params( + attention, torch.empty(0), None, None, None, forward_args + ) + + assert runtime_params.block_sparse_inputs is block_sparse_inputs + assert runtime_params == SparseRuntimeParams(block_sparse_inputs=block_sparse_inputs) + + +def test_block_sparse_attn_predict_override_composes_with_index_predictors() -> None: + attention = _StaticPredictionAttention.__new__(_StaticPredictionAttention) + attention.sparse_params = _StubSparseParams() + sparse_attn_indices = torch.tensor([2], dtype=torch.int32) + sparse_attn_offsets = torch.tensor([0, 1], dtype=torch.int32) + attention.sparse_kv_predict = Mock(return_value=(None, None)) + attention.sparse_attn_predict = Mock(return_value=(sparse_attn_indices, sparse_attn_offsets)) + block_sparse_inputs = _make_block_sparse_inputs() + attention.block_sparse_attn_predict = Mock(return_value=block_sparse_inputs) + q = torch.empty((1, 4)) + k = torch.empty((1, 4)) + v = torch.empty((1, 4)) + metadata = Mock() + forward_args = AttentionForwardArgs() + + runtime_params = prepare_sparse_runtime_params(attention, q, k, v, metadata, forward_args) + + assert runtime_params.block_sparse_inputs is block_sparse_inputs + assert runtime_params.sparse_attn_indices is sparse_attn_indices + assert runtime_params.sparse_attn_indices_block_size == 1 + attention.block_sparse_attn_predict.assert_called_once_with(q, k, v, metadata, forward_args) + + +def test_attention_forward_args_default_to_empty_sparse_runtime_params() -> None: + assert AttentionForwardArgs().sparse_runtime_params == SparseRuntimeParams() + + +class _StopAfterShapeValidation(Exception): + pass + + +def _make_sparse_prediction_forward_backend() -> TrtllmAttention: + attention = _StaticPredictionAttention.__new__(_StaticPredictionAttention) + attention.sparse_params = None + attention.is_mla_enable = False + attention.num_heads = 1 + attention.num_kv_heads = 1 + attention.head_dim = 4 + attention.get_local_layer_idx = Mock(return_value=1) + attention._ensure_rope_table_size = Mock(side_effect=_StopAfterShapeValidation) + return attention + + +def _make_sparse_prediction_forward_metadata() -> TrtllmAttentionMetadata: + metadata = object.__new__(TrtllmAttentionMetadata) + seq_lens = torch.tensor([2], dtype=torch.int32) + metadata._seq_lens = seq_lens + metadata._seq_lens_kv = seq_lens + metadata._seq_lens_cuda = None + metadata.kv_cache_manager = None + metadata._max_seq_len_storage = 4 + metadata.use_paged_context_fmha = False + metadata.cu_q_seqlens = None + metadata.cu_kv_seqlens = None + metadata.enable_flash_mla = False + metadata.spec_bl_tree_first_sparse_mask_offset_kv = None + metadata.kv_lens_cuda_runtime = torch.tensor([2], dtype=torch.int32) + metadata.kv_lens_runtime = torch.tensor([2], dtype=torch.int32) + metadata.prompt_lens_cuda_runtime = torch.tensor([2], dtype=torch.int32) + metadata.prompt_lens_cpu_runtime = torch.tensor([2], dtype=torch.int32) + metadata.host_request_types_runtime = torch.tensor([0], dtype=torch.int32) + return metadata + + +def test_forward_materializes_dynamic_block_sparse_prediction_before_shape_validation() -> None: + attention = _make_sparse_prediction_forward_backend() + block_sparse_inputs = _make_block_sparse_inputs() + prediction = SparseRuntimeParams( + sparse_attn_kv_lens=torch.tensor([4]), + block_sparse_inputs=block_sparse_inputs, + ) + q = torch.empty((2, 4)) + k = torch.empty((4, 4)) + v = torch.empty((4, 4)) + metadata = _make_sparse_prediction_forward_metadata() + forward_args = AttentionForwardArgs( + output=torch.empty_like(q), + attention_input_type=AttentionInputType.context_only, + ) + + with patch.object( + trtllm_backend, "prepare_sparse_runtime_params", return_value=prediction + ) as prepare: + for _ in range(2): + with pytest.raises(_StopAfterShapeValidation): + attention.forward(q, k, v, metadata, forward_args) + + assert prepare.call_count == 2 + prepare.assert_called_with(attention, q, k, v, metadata, forward_args) + assert attention._ensure_rope_table_size.call_count == 2 + assert forward_args.sparse_runtime_params == SparseRuntimeParams() + + +@pytest.mark.parametrize("has_block_sparse_inputs", [False, True]) +def test_forward_replaces_caller_sparse_runtime_params_with_prediction( + has_block_sparse_inputs: bool, +) -> None: + attention = _make_sparse_prediction_forward_backend() + block_sparse_inputs = _make_block_sparse_inputs() if has_block_sparse_inputs else None + prediction = SparseRuntimeParams( + sparse_attn_kv_lens=torch.tensor([2]), + block_sparse_inputs=block_sparse_inputs, + ) + caller_params = SparseRuntimeParams(sparse_attn_kv_lens=torch.tensor([7])) + q = torch.empty((2, 4)) + num_kv_tokens = 4 if has_block_sparse_inputs else 2 + k = torch.empty((num_kv_tokens, 4)) + v = torch.empty((num_kv_tokens, 4)) + metadata = _make_sparse_prediction_forward_metadata() + forward_args = AttentionForwardArgs( + output=torch.empty_like(q), + attention_input_type=AttentionInputType.context_only, + sparse_runtime_params=caller_params, + ) + + with patch.object( + trtllm_backend, "prepare_sparse_runtime_params", return_value=prediction + ) as prepare: + for _ in range(2): + with pytest.raises(_StopAfterShapeValidation): + attention.forward(q, k, v, metadata, forward_args) + + assert prepare.call_count == 2 + assert attention._ensure_rope_table_size.call_count == 2 + assert forward_args.sparse_runtime_params is caller_params + + +@pytest.mark.parametrize("backend_name", ["FLASHINFER", "unknown"]) +def test_sparse_attention_backend_fallback_does_not_redispatch( + backend_name: str, monkeypatch: pytest.MonkeyPatch +) -> None: + from tensorrt_llm._torch.attention.backends import utils as attention_backend_utils + from tensorrt_llm._torch.attention.backends.sparse.skip_softmax import SkipSoftmaxParams + + monkeypatch.setattr(attention_backend_utils, "IS_FLASHINFER_AVAILABLE", False) + + with patch.object( + attention_backend_utils, + "get_trtllm_sparse_attn_attention_backend", + ) as trtllm_sparse_resolver: + backend = attention_backend_utils.get_attention_backend( + backend_name, sparse_params=SkipSoftmaxParams() + ) + + assert backend is TrtllmAttention + trtllm_sparse_resolver.assert_not_called() + + +@pytest.mark.parametrize("backend_name", ["FLASHINFER", "unknown"]) +def test_trtllm_fallback_without_sparse_params_remains_dense( + backend_name: str, monkeypatch: pytest.MonkeyPatch +) -> None: + from tensorrt_llm._torch.attention.backends import utils as attention_backend_utils + + monkeypatch.setattr(attention_backend_utils, "IS_FLASHINFER_AVAILABLE", False) + + assert attention_backend_utils.get_attention_backend(backend_name) is TrtllmAttention diff --git a/tests/unittest/_torch/attention/test_attention_op_sync.py b/tests/unittest/_torch/attention/test_attention_op_sync.py index ba5c75f3d7b6..ccf11544be53 100644 --- a/tests/unittest/_torch/attention/test_attention_op_sync.py +++ b/tests/unittest/_torch/attention/test_attention_op_sync.py @@ -41,7 +41,7 @@ import textwrap import typing from dataclasses import fields -from types import SimpleNamespace +from types import SimpleNamespace, UnionType import pytest import torch @@ -382,7 +382,19 @@ def _dataclass_field_type(cls, name: str): return None if f is None: return None - return f.type if not isinstance(f.type, str) else None + if isinstance(f.type, str): + return None + return _unwrap_optional(f.type) + + +def _unwrap_optional(py_type): + """Return the payload type for ``Optional[T]`` annotations.""" + origin = typing.get_origin(py_type) + if origin in (typing.Union, UnionType): + args = [arg for arg in typing.get_args(py_type) if arg is not type(None)] + if len(args) == 1: + return args[0] + return py_type def _resolve_path(root_cls, path: tuple[str, ...]): @@ -404,7 +416,7 @@ def _python_category(py_type) -> str: confidently (the type check is then skipped for that kwarg).""" # Unwrap Optional[X] / Union[X, None]. origin = typing.get_origin(py_type) - if origin is typing.Union: + if origin in (typing.Union, UnionType): args = [a for a in typing.get_args(py_type) if a is not type(None)] if len(args) == 1: return _python_category(args[0]) @@ -558,7 +570,7 @@ def _verify_consumed(cls, chains: set[tuple[str, ...]], excluded=frozenset()): for f in fields(cls): if f.name in excluded: continue - ftype = f.type if not isinstance(f.type, str) else None + ftype = _dataclass_field_type(cls, f.name) if ftype is not None and dataclasses.is_dataclass(ftype): sub = {p[1:] for p in chains if len(p) >= 2 and p[0] == f.name} assert sub, ( @@ -566,7 +578,7 @@ def _verify_consumed(cls, chains: set[tuple[str, ...]], excluded=frozenset()): f"declared but `{f.name}.` is never read at the " f"call site." ) - _verify_consumed(ftype, sub) + _verify_consumed(ftype, sub, excluded=excluded) else: assert f.name in consumed, ( f"Field `{f.name}` on {cls.__name__} not consumed by the " @@ -609,7 +621,7 @@ def _all_forward_args_field_names() -> set[str]: def _walk(cls) -> None: for f in fields(cls): seen.add(f.name) - ftype = f.type if not isinstance(f.type, str) else None + ftype = _dataclass_field_type(cls, f.name) if ftype is not None and dataclasses.is_dataclass(ftype): _walk(ftype) diff --git a/tests/unittest/_torch/attention/test_fmha_manager.py b/tests/unittest/_torch/attention/test_fmha_manager.py index a074aceac960..b59a5c80595b 100644 --- a/tests/unittest/_torch/attention/test_fmha_manager.py +++ b/tests/unittest/_torch/attention/test_fmha_manager.py @@ -31,7 +31,10 @@ AttentionInputType, PredefinedAttentionMask, ) -from tensorrt_llm._torch.attention.backends.sparse.params import BlockSparseForwardInputs +from tensorrt_llm._torch.attention.backends.sparse.params import ( + BlockSparseForwardInputs, + SparseRuntimeParams, +) from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization.mode import QuantAlgo @@ -51,12 +54,14 @@ def _make_metadata( num_generations: int, num_ctx_tokens: int = 0, use_spec_decoding: bool = False, + num_sparse_topk: int = 0, ) -> SimpleNamespace: return SimpleNamespace( num_contexts=num_contexts, num_generations=num_generations, num_ctx_tokens=num_ctx_tokens, use_spec_decoding=use_spec_decoding, + num_sparse_topk=num_sparse_topk, ) @@ -547,13 +552,17 @@ def test_fmha_cache_separates_block_sparse_mode(block_sparse_first: bool) -> Non attn, "block-sparse", events, - support_predicate=lambda forward_args: forward_args.block_sparse_inputs is not None, + support_predicate=lambda forward_args: ( + forward_args.sparse_runtime_params.block_sparse_inputs is not None + ), ) dense_fmha = FakeFmha( attn, "dense", events, - support_predicate=lambda forward_args: forward_args.block_sparse_inputs is None, + support_predicate=lambda forward_args: ( + forward_args.sparse_runtime_params.block_sparse_inputs is None + ), ) manager.fmha_libs = [block_sparse_fmha, dense_fmha] metadata = _make_metadata(num_contexts=1, num_generations=0, num_ctx_tokens=1) @@ -562,10 +571,12 @@ def test_fmha_cache_separates_block_sparse_mode(block_sparse_first: bool) -> Non False: AttentionForwardArgs(attention_input_type=AttentionInputType.context_only), True: AttentionForwardArgs( attention_input_type=AttentionInputType.context_only, - block_sparse_inputs=BlockSparseForwardInputs( - q_block_size=64, - kv_block_size=64, - exact_block_bits=torch.zeros((1, 1), dtype=torch.uint32), + sparse_runtime_params=SparseRuntimeParams( + block_sparse_inputs=BlockSparseForwardInputs( + q_block_size=64, + kv_block_size=64, + exact_block_bits=torch.zeros((1, 1), dtype=torch.uint32), + ), ), ), } diff --git a/tests/unittest/_torch/attention/test_fmha_registry.py b/tests/unittest/_torch/attention/test_fmha_registry.py index 64060bdbc935..285f2aa075ba 100644 --- a/tests/unittest/_torch/attention/test_fmha_registry.py +++ b/tests/unittest/_torch/attention/test_fmha_registry.py @@ -18,6 +18,7 @@ from tensorrt_llm._torch.attention.backends.fmha import registry from tensorrt_llm._torch.attention.backends.fmha.interface import Fmha +from tensorrt_llm._torch.attention.backends.sparse.params import SparseRuntimeParams PRIMS_TS = "prims_ts" PRIMS_TS_BLOCK_SPARSE = "prims_ts_block_sparse" @@ -52,7 +53,11 @@ def test_dense_fmhas_reject_unconsumed_block_sparse_inputs(name: str) -> None: attention = type("Attention", (), {})() fmha = object.__new__(registry.FMHA_LIBS[name]) Fmha.__init__(fmha, attention) - forward_args = type("ForwardArgs", (), {"block_sparse_inputs": object()})() + forward_args = type( + "ForwardArgs", + (), + {"sparse_runtime_params": SparseRuntimeParams(block_sparse_inputs=object())}, + )() assert not fmha.is_supported( torch.empty((1, 4), dtype=torch.bfloat16), diff --git a/tests/unittest/_torch/attention/test_prims_ts_block_sparse.py b/tests/unittest/_torch/attention/test_prims_ts_block_sparse.py index e45f7da1166a..3849b9481077 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_block_sparse.py +++ b/tests/unittest/_torch/attention/test_prims_ts_block_sparse.py @@ -15,7 +15,6 @@ import math from contextlib import nullcontext -from importlib import import_module from types import SimpleNamespace from unittest.mock import Mock @@ -24,6 +23,7 @@ from utils.util import isSM100Family from tensorrt_llm._torch.attention.backends import prims_ts +from tensorrt_llm._torch.attention.backends.fmha import prims_ts_block_sparse as block_sparse_fmha from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams from tensorrt_llm._torch.attention.backends.interface import ( @@ -31,24 +31,23 @@ AttentionInputType, PredefinedAttentionMask, ) +from tensorrt_llm._torch.attention.backends.sparse.params import ( + BlockSparseForwardInputs, + SparseRuntimeParams, +) from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm.functional import PositionEmbeddingType +pytestmark = pytest.mark.cpu_only + _REQUIRES_PRIMTS_GPU = pytest.mark.skipif( not isSM100Family(), reason="PrimTS block-sparse attention requires SM100 or SM103", ) -def _generic_api(): - carrier_module = import_module("tensorrt_llm._torch.attention.backends.sparse.params") - fmha_module = import_module("tensorrt_llm._torch.attention.backends.fmha.prims_ts_block_sparse") - return carrier_module.BlockSparseForwardInputs, fmha_module - - def _bsr_inputs(*, kv_valid_bits: torch.Tensor | None = None): - inputs_type, _ = _generic_api() - return inputs_type( + return BlockSparseForwardInputs( q_block_size=64, kv_block_size=64, max_blocks_per_row=2, @@ -59,12 +58,11 @@ def _bsr_inputs(*, kv_valid_bits: torch.Tensor | None = None): def _bitmask_inputs(*, proxy: bool): - inputs_type, _ = _generic_api() summaries = { "k_summary": torch.zeros((2, 4, 1, 128), dtype=torch.bfloat16), "v_summary": torch.zeros((2, 4, 1, 128), dtype=torch.bfloat16), } - return inputs_type( + return BlockSparseForwardInputs( q_block_size=64, kv_block_size=64, exact_block_bits=torch.ones((2, 1, 1, 1), dtype=torch.uint32), @@ -72,6 +70,21 @@ def _bitmask_inputs(*, proxy: bool): ) +def _set_block_sparse_inputs( + forward_args: AttentionForwardArgs, + block_sparse_inputs, +) -> None: + forward_args.sparse_runtime_params = SparseRuntimeParams( + block_sparse_inputs=block_sparse_inputs + ) + + +def _get_block_sparse_inputs(forward_args: AttentionForwardArgs): + block_sparse_inputs = forward_args.sparse_runtime_params.block_sparse_inputs + assert block_sparse_inputs is not None + return block_sparse_inputs + + def _pack_token_mask(mask: torch.Tensor) -> torch.Tensor: shifts = torch.arange(32, dtype=torch.int64, device=mask.device) weights = torch.ones_like(shifts).bitwise_left_shift_(shifts) @@ -125,14 +138,15 @@ def __init__(self) -> None: def _contiguous_case(): - _inputs_type, fmha_module = _generic_api() attention = _Attention() - fmha = fmha_module.PrimsTSBlockSparseFmha(attention) + fmha = block_sparse_fmha.PrimsTSBlockSparseFmha(attention) q = torch.zeros((128, 256), dtype=torch.bfloat16) k = torch.zeros((512, 128), dtype=torch.bfloat16) v = torch.zeros_like(k) metadata = SimpleNamespace( is_cross=False, + num_sparse_topk=0, + helix_position_offsets=None, kv_cache_manager=None, seq_lens=torch.tensor([64, 64], dtype=torch.int32), ) @@ -140,7 +154,7 @@ def _contiguous_case(): output=torch.empty_like(q), attention_input_type=AttentionInputType.context_only, attention_mask=PredefinedAttentionMask.FULL, - block_sparse_inputs=_bsr_inputs(), + sparse_runtime_params=SparseRuntimeParams(block_sparse_inputs=_bsr_inputs()), ) return attention, fmha, q, k, v, metadata, args @@ -155,10 +169,16 @@ def _paged_metadata(): manager.host_kv_cache_block_offsets = block_offsets return SimpleNamespace( is_cross=False, + num_sparse_topk=0, + helix_position_offsets=None, num_contexts=0, num_generations=batch_size, seq_lens=torch.ones(batch_size, dtype=torch.int32), beam_width=1, + is_spec_decoding_enabled=False, + use_spec_decoding=False, + is_spec_dec_tree=False, + is_spec_dec_dynamic_tree=False, tokens_per_block=page_size, max_seq_len=max_pages * page_size, kv_layout="HND", @@ -171,9 +191,8 @@ def _paged_metadata(): def _paged_case(): - _inputs_type, fmha_module = _generic_api() attention = _Attention() - fmha = fmha_module.PrimsTSBlockSparseFmha(attention) + fmha = block_sparse_fmha.PrimsTSBlockSparseFmha(attention) fmha._multi_processor_count = 1 metadata = _paged_metadata() q = torch.zeros((2, 512), dtype=torch.bfloat16) @@ -183,7 +202,7 @@ def _paged_case(): attention_mask=PredefinedAttentionMask.CAUSAL, attention_window_size=metadata.max_seq_len, is_fused_qkv=True, - block_sparse_inputs=_bsr_inputs(), + sparse_runtime_params=SparseRuntimeParams(block_sparse_inputs=_bsr_inputs()), ) return attention, fmha, q, metadata, args @@ -211,7 +230,6 @@ def test_block_sparse_route_mode_is_derived_from_payload() -> None: ], ) def test_block_sparse_payload_rejects_ambiguous_combinations(overrides, message) -> None: - inputs_type, _ = _generic_api() kwargs = { "q_block_size": 64, "kv_block_size": 64, @@ -222,7 +240,7 @@ def test_block_sparse_payload_rejects_ambiguous_combinations(overrides, message) kwargs.update(overrides) with pytest.raises((TypeError, ValueError), match=message): - inputs_type(**kwargs) + BlockSparseForwardInputs(**kwargs) def test_block_sparse_support_is_phase_specific_and_paged_proxy_is_rejected( @@ -234,7 +252,7 @@ def test_block_sparse_support_is_phase_specific_and_paged_proxy_is_rejected( assert not contiguous.is_supported(q, k, v, metadata, args, phase=FmhaPhase.GENERATION) _attention, paged, q, metadata, args = _paged_case() - args.block_sparse_inputs = _bitmask_inputs(proxy=True) + _set_block_sparse_inputs(args, _bitmask_inputs(proxy=True)) _supported, reason = paged._is_supported_with_reason( q, None, None, metadata, args, phase=FmhaPhase.GENERATION ) @@ -244,7 +262,7 @@ def test_block_sparse_support_is_phase_specific_and_paged_proxy_is_rejected( def test_contiguous_proxy_routes_reject_causal_mask_before_planning(monkeypatch) -> None: _attention, fmha, q, k, v, metadata, args = _contiguous_case() - args.block_sparse_inputs = _bitmask_inputs(proxy=True) + _set_block_sparse_inputs(args, _bitmask_inputs(proxy=True)) args.attention_mask = PredefinedAttentionMask.CAUSAL monkeypatch.setattr(fmha, "_common_unsupported_reason", Mock(return_value=None)) @@ -303,14 +321,13 @@ def test_block_sparse_support_rejects_invalid_static_kernel_profile( def test_contiguous_wrappers_cache_static_profile_and_keep_routes_live(monkeypatch) -> None: _attention, fmha, q, k, v, _metadata, args = _contiguous_case() - _inputs_type, fmha_module = _generic_api() wrapper = Mock() factory = Mock(return_value=wrapper) - monkeypatch.setattr(fmha_module, "_BlockSparseTSWrapper", factory) + monkeypatch.setattr(block_sparse_fmha, "_BlockSparseTSWrapper", factory) bsr_inputs = [ _bsr_inputs(), - _inputs_type( + BlockSparseForwardInputs( q_block_size=64, kv_block_size=64, max_blocks_per_row=2, @@ -319,12 +336,12 @@ def test_contiguous_wrappers_cache_static_profile_and_keep_routes_live(monkeypat ), ] for inputs in bsr_inputs: - args.block_sparse_inputs = inputs + _set_block_sparse_inputs(args, inputs) fmha._forward_contiguous(q, k, v, args) proxy_inputs = [_bitmask_inputs(proxy=True), _bitmask_inputs(proxy=True)] for inputs in proxy_inputs: - args.block_sparse_inputs = inputs + _set_block_sparse_inputs(args, inputs) fmha._forward_contiguous(q, k, v, args) assert factory.call_count == 2 @@ -346,14 +363,13 @@ def test_contiguous_wrappers_cache_static_profile_and_keep_routes_live(monkeypat def test_block_sparse_plan_key_includes_attention_head_topology() -> None: - _, fmha_module = _generic_api() inputs = _bitmask_inputs(proxy=True) q = torch.empty((128, 256), dtype=torch.bfloat16) first_attention = _Attention() second_attention = _Attention() second_attention.num_heads = 4 - first = fmha_module.PrimsTSBlockSparseFmha(first_attention) - second = fmha_module.PrimsTSBlockSparseFmha(second_attention) + first = block_sparse_fmha.PrimsTSBlockSparseFmha(first_attention) + second = block_sparse_fmha.PrimsTSBlockSparseFmha(second_attention) def _key(fmha): return fmha._make_plan_key( @@ -370,9 +386,8 @@ def _key(fmha): def test_block_sparse_plan_cache_is_shared_only_when_explicitly_bound() -> None: - _, fmha_module = _generic_api() - first = fmha_module.PrimsTSBlockSparseFmha(_Attention()) - second = fmha_module.PrimsTSBlockSparseFmha(_Attention()) + first = block_sparse_fmha.PrimsTSBlockSparseFmha(_Attention()) + second = block_sparse_fmha.PrimsTSBlockSparseFmha(_Attention()) assert first._contiguous_wrappers is not second._contiguous_wrappers assert first._paged_wrappers is not second._paged_wrappers @@ -391,10 +406,9 @@ def test_block_sparse_plan_cache_is_shared_only_when_explicitly_bound() -> None: def test_paged_wrapper_uses_zero_copy_padded_row_stride_block_tables(monkeypatch) -> None: attention, fmha, q, metadata, args = _paged_case() - _inputs_type, fmha_module = _generic_api() wrapper = Mock() - monkeypatch.setattr(fmha_module, "_BlockSparsePagedTSWrapper", Mock(return_value=wrapper)) - monkeypatch.setattr(fmha_module, "get_kv_page_offset", Mock(return_value=8)) + monkeypatch.setattr(block_sparse_fmha, "_BlockSparsePagedTSWrapper", Mock(return_value=wrapper)) + monkeypatch.setattr(block_sparse_fmha, "get_kv_page_offset", Mock(return_value=8)) monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", Mock(return_value=False)) q_processed = torch.zeros((2, 2, 128), dtype=torch.bfloat16) kv_pool = torch.empty((16, 1, 64, 128), dtype=torch.bfloat16) @@ -436,16 +450,19 @@ def snapshot(*_args, **kwargs): ) wrapper.run.side_effect = snapshot - first_inputs = args.block_sparse_inputs + first_inputs = _get_block_sparse_inputs(args) fmha.run_generation(params) block_tables[:, 0].add_(10) params.sequence_lengths = torch.tensor([130, 194], dtype=torch.int32) - args.block_sparse_inputs = _inputs_type( - q_block_size=64, - kv_block_size=64, - max_blocks_per_row=2, - block_indptr=torch.tensor([[[0, 1]], [[1, 4]]], dtype=torch.int32), - block_indices=torch.tensor([3, 2, 1, 0], dtype=torch.int32), + _set_block_sparse_inputs( + args, + BlockSparseForwardInputs( + q_block_size=64, + kv_block_size=64, + max_blocks_per_row=2, + block_indptr=torch.tensor([[[0, 1]], [[1, 4]]], dtype=torch.int32), + block_indices=torch.tensor([3, 2, 1, 0], dtype=torch.int32), + ), ) fmha.run_generation(params) @@ -460,15 +477,14 @@ def snapshot(*_args, **kwargs): torch.testing.assert_close(snapshots[0][2], torch.arange(8, dtype=torch.int32).view(2, 4)) torch.testing.assert_close(snapshots[1][2], torch.arange(8, dtype=torch.int32).view(2, 4) + 10) assert snapshots[0][3] is first_inputs.block_indptr - assert snapshots[1][3] is args.block_sparse_inputs.block_indptr + assert snapshots[1][3] is _get_block_sparse_inputs(args).block_indptr def test_paged_block_tables_remain_live_across_graph_replay(monkeypatch) -> None: attention, fmha, q, metadata, args = _paged_case() - _inputs_type, fmha_module = _generic_api() wrapper = Mock() - monkeypatch.setattr(fmha_module, "_BlockSparsePagedTSWrapper", Mock(return_value=wrapper)) - monkeypatch.setattr(fmha_module, "get_kv_page_offset", Mock(return_value=8)) + monkeypatch.setattr(block_sparse_fmha, "_BlockSparsePagedTSWrapper", Mock(return_value=wrapper)) + monkeypatch.setattr(block_sparse_fmha, "get_kv_page_offset", Mock(return_value=8)) monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", Mock(return_value=True)) q_processed = torch.zeros((2, 2, 128), dtype=torch.bfloat16) kv_pool = torch.empty((16, 1, 64, 128), dtype=torch.bfloat16) @@ -520,6 +536,7 @@ def test_prepare_workspace_checks_capture_before_resize(monkeypatch) -> None: shape=(2, 512), ) metadata = SimpleNamespace( + kv_cache_manager=object(), kv_cache_block_offsets=SimpleNamespace(device=query_device, shape=(1, 2, 4)), max_num_requests=2, tokens_per_block=64, @@ -543,6 +560,20 @@ def test_prepare_workspace_checks_capture_before_resize(monkeypatch) -> None: assert workspace.numel() == 0 +def test_prepare_workspace_skips_generation_layout_for_contiguous_requests(monkeypatch) -> None: + _attention, fmha, q, _k, _v, metadata, args = _contiguous_case() + layout = Mock() + monkeypatch.setattr(fmha, "_get_generation_workspace_layout", layout) + monkeypatch.setattr(torch.cuda, "device", Mock(return_value=nullcontext())) + fmha._multi_processor_count = 1 + workspace = torch.empty(0, dtype=torch.uint8) + + fmha.prepare_workspace(q, None, None, metadata, args, workspace) + + layout.assert_not_called() + assert workspace.numel() == 0 + + @_REQUIRES_PRIMTS_GPU @torch.no_grad() def test_real_gpu_raw_routes_and_token_mask_match_reference() -> None: @@ -552,8 +583,7 @@ def test_real_gpu_raw_routes_and_token_mask_match_reference() -> None: v = torch.randn_like(k) token_mask = torch.ones(256, device="cuda", dtype=torch.bool) token_mask[[1, 63, 64, 95, 129, 190, 255]] = False - inputs_type, _ = _generic_api() - inputs = inputs_type( + inputs = BlockSparseForwardInputs( q_block_size=64, kv_block_size=64, max_blocks_per_row=3, @@ -602,21 +632,22 @@ def test_real_gpu_proxy_adapter_replays_live_routes_and_summaries() -> None: live_v_summary = initial_v_summary.clone() live_exact_bits = torch.tensor([[[[1]]]], device="cuda", dtype=torch.uint32) - inputs_type, fmha_module = _generic_api() attention = _Attention() attention.num_heads = attention.num_kv_heads = 1 - fmha = fmha_module.PrimsTSBlockSparseFmha(attention) + fmha = block_sparse_fmha.PrimsTSBlockSparseFmha(attention) output = torch.empty_like(q).view(64, 128) args = AttentionForwardArgs( output=output, attention_input_type=AttentionInputType.context_only, attention_mask=PredefinedAttentionMask.FULL, - block_sparse_inputs=inputs_type( - q_block_size=64, - kv_block_size=64, - exact_block_bits=live_exact_bits, - k_summary=live_k_summary, - v_summary=live_v_summary, + sparse_runtime_params=SparseRuntimeParams( + block_sparse_inputs=BlockSparseForwardInputs( + q_block_size=64, + kv_block_size=64, + exact_block_bits=live_exact_bits, + k_summary=live_k_summary, + v_summary=live_v_summary, + ), ), ) metadata = SimpleNamespace( @@ -654,8 +685,7 @@ def test_real_gpu_paged_routes_use_live_length_below_capacity() -> None: v_cache = torch.randn_like(k_cache) page_indices = torch.tensor([2, 0, 3, 1], device="cuda", dtype=torch.int32) seq_lens_kv = torch.tensor([160], device="cuda", dtype=torch.int32) - inputs_type, _ = _generic_api() - inputs = inputs_type( + inputs = BlockSparseForwardInputs( q_block_size=64, kv_block_size=64, max_blocks_per_row=2, diff --git a/tests/unittest/_torch/attention/test_prims_ts_fmha.py b/tests/unittest/_torch/attention/test_prims_ts_fmha.py index 87dbd7364f4b..a18f41c00bf7 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_fmha.py +++ b/tests/unittest/_torch/attention/test_prims_ts_fmha.py @@ -38,6 +38,7 @@ AttentionInputType, PredefinedAttentionMask, ) +from tensorrt_llm._torch.attention.backends.sparse.params import SparseRuntimeParams from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm.bindings import DataType @@ -185,7 +186,7 @@ def _support_result( is_fused_qkv=is_fused_qkv, ) if has_sparse_runtime_metadata: - forward_args.sparse_runtime_params.sparse_kv_indices = torch.empty(1) + forward_args.sparse_runtime_params = SparseRuntimeParams(sparse_kv_indices=torch.empty(1)) if attention_input_type == AttentionInputType.context_only: num_contexts, num_generations, num_ctx_tokens = 1, 0, 4 kv_lens = [4] diff --git a/tests/unittest/_torch/attention/test_skip_softmax_sm120.py b/tests/unittest/_torch/attention/test_skip_softmax_sm120.py index 979b0d3c6908..af8a067f7129 100644 --- a/tests/unittest/_torch/attention/test_skip_softmax_sm120.py +++ b/tests/unittest/_torch/attention/test_skip_softmax_sm120.py @@ -75,13 +75,16 @@ def _run_context( ) -> tuple: """Build a TRTLLM attention layer + no-cache context metadata and run a packed-QKV causal prefill. Mirrors ``test_attention_no_cache``.""" - AttentionCls = get_attention_backend("TRTLLM") + sparse_params = ( + sparse_attention_config.to_sparse_params() if sparse_attention_config is not None else None + ) + AttentionCls = get_attention_backend("TRTLLM", sparse_params=sparse_params) layer = AttentionCls( layer_idx=0, num_heads=num_heads, head_dim=head_dim, num_kv_heads=num_kv_heads, - sparse_attention_config=sparse_attention_config, + sparse_params=sparse_params, ) metadata = AttentionCls.Metadata( From 1b91f9ff271bb2af2a1f42814b1f369dacd8dd81 Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Tue, 8 Sep 2026 13:59:13 +0000 Subject: [PATCH 4/7] fix: keep FA4 WorkTileInfo independent of CUTLASS task scheduling Importing cutlass.experimental.task_scheduling rewrites the shared cutlass.utils.WorkTileInfo class in place so that its constructor unpacks tile_idx into exactly three scalars. FlashAttention 4 subclasses that class with a four-axis coordinate and inherits the constructor, so any process that probes or plans the vendored PrimTS kernels turns every later FA4 kernel trace into a ValueError. Install the upstream tuple semantics on the FA4 subclass from the existing CuTe DSL compatibility layer so the parent rewrite cannot reach it, and cover the worst-case import order in the FA4 compatibility tests. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- .../attention_backend/flash_attn4.py | 54 ++++++++++++++++++ .../test_fa4_cutlass_compatibility.py | 55 +++++++++++++++++++ 2 files changed, 109 insertions(+) diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py b/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py index 5024a43b192a..a441ce7edac7 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py @@ -39,10 +39,64 @@ def _install_cutlass_dsl_compatibility() -> None: cute.make_fragment = cute.make_rmem_tensor +def _install_flash_attn_tile_scheduler_compatibility() -> None: + """Keep FA4's four-axis ``WorkTileInfo`` independent of CUTLASS task scheduling. + + Importing ``cutlass.experimental.task_scheduling`` rewrites the shared + ``cutlass.utils.WorkTileInfo`` class in place: its constructor unpacks + ``tile_idx`` into exactly three scalars and ``tile_idx`` / ``is_valid_tile`` + become properties over those scalars. The vendored PrimTS kernels import + that package, so once any PrimTS FMHA has been probed or planned in a + process, every later FA4 kernel trace fails with ``ValueError: too many + values to unpack (expected 3)``: FA4 subclasses the same CUTLASS class with + a (block, head, batch, split) coordinate but does not define its own + constructor. Installing the upstream tuple semantics directly on the FA4 + subclass makes it immune to the parent rewrite regardless of import order. + Remove once CUTLASS stops patching the shared class or FA4 owns these + members itself. + """ + try: + from flash_attn.cute import tile_scheduler + except (ImportError, OSError): + return + import cutlass.cute as cute + from cutlass.cutlass_dsl import Boolean, extract_mlir_values + + work_tile_info = tile_scheduler.WorkTileInfo + if "__init__" in vars(work_tile_info): + return + + def __init__(self, tile_idx: cute.Coord, is_valid_tile: Boolean) -> None: + self._tile_idx = tile_idx + self._is_valid_tile = Boolean(is_valid_tile) + self._tile_idx_num_values = None + + def __extract_mlir_values__(self) -> list: + tile_idx_values = extract_mlir_values(self._tile_idx) + valid_values = extract_mlir_values(self._is_valid_tile) + self._tile_idx_num_values = len(tile_idx_values) + return tile_idx_values + valid_values + + @cute.jit + def tile_idx(self) -> cute.Coord: + return self._tile_idx + + @cute.jit + def is_valid_tile(self) -> Boolean: + return self._is_valid_tile + + work_tile_info.__init__ = __init__ + work_tile_info.__extract_mlir_values__ = __extract_mlir_values__ + work_tile_info.tile_idx = property(tile_idx) + work_tile_info.is_valid_tile = property(is_valid_tile) + + _flash_attn_fwd_import_error = None try: _install_cutlass_dsl_compatibility() from flash_attn.cute.interface import _flash_attn_fwd + + _install_flash_attn_tile_scheduler_compatibility() except (ImportError, OSError) as e: _flash_attn_fwd = None _flash_attn_fwd_import_error = e diff --git a/tests/unittest/_torch/visual_gen/test_fa4_cutlass_compatibility.py b/tests/unittest/_torch/visual_gen/test_fa4_cutlass_compatibility.py index eb9ea6941792..e1203891be26 100644 --- a/tests/unittest/_torch/visual_gen/test_fa4_cutlass_compatibility.py +++ b/tests/unittest/_torch/visual_gen/test_fa4_cutlass_compatibility.py @@ -1,6 +1,9 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import subprocess +import sys +import textwrap from importlib import import_module import pytest @@ -10,6 +13,7 @@ from tensorrt_llm._torch.visual_gen.attention_backend import flash_attn4, parallel # noqa: E402 from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import ( # noqa: E402 _install_cutlass_dsl_compatibility, + _install_flash_attn_tile_scheduler_compatibility, ) @@ -48,3 +52,54 @@ def test_cutlass_dsl_47_aliases_allow_fa4_interface_import() -> None: assert callable(interface.flash_attn_combine) assert callable(flash_attn4._flash_attn_fwd) assert callable(parallel._flash_attn_combine) + + +def test_fa4_work_tile_info_survives_cutlass_task_scheduling_import() -> None: + task_scheduling = pytest.importorskip("cutlass.experimental.task_scheduling") + tile_scheduler = pytest.importorskip("flash_attn.cute.tile_scheduler") + import cutlass + from cutlass.cutlass_dsl import Boolean + from cutlass.utils.static_persistent_tile_scheduler import WorkTileInfo as CutlassWorkTileInfo + + del task_scheduling + _install_flash_attn_tile_scheduler_compatibility() + tile_idx = (cutlass.Int32(1), cutlass.Int32(2), cutlass.Int32(3), cutlass.Int32(0)) + + # The task-scheduling import rewrote the shared CUTLASS class to three scalars. + with pytest.raises(ValueError, match="too many values to unpack"): + CutlassWorkTileInfo(tile_idx, Boolean(True)) + + fa4_tile = tile_scheduler.WorkTileInfo(tile_idx, Boolean(True)) + + assert "__init__" in vars(tile_scheduler.WorkTileInfo) + assert fa4_tile.tile_idx == tile_idx + assert bool(fa4_tile.is_valid_tile) + assert issubclass(tile_scheduler.WorkTileInfo, CutlassWorkTileInfo) + + +def test_fa4_work_tile_info_survives_task_scheduling_imported_first() -> None: + pytest.importorskip("cutlass.experimental.task_scheduling") + pytest.importorskip("flash_attn.cute.tile_scheduler") + script = textwrap.dedent( + """ + import cutlass + import cutlass.experimental.task_scheduling # noqa: F401 + from cutlass.cutlass_dsl import Boolean + + import tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 # noqa: F401 + from flash_attn.cute.tile_scheduler import WorkTileInfo + + tile_idx = (cutlass.Int32(1), cutlass.Int32(2), cutlass.Int32(3), cutlass.Int32(0)) + tile = WorkTileInfo(tile_idx, Boolean(True)) + assert tile.tile_idx == tile_idx + assert bool(tile.is_valid_tile) + print("fa4-work-tile-info-ok") + """ + ) + + result = subprocess.run( + [sys.executable, "-c", script], capture_output=True, text=True, timeout=900, check=False + ) + + assert result.returncode == 0, result.stderr[-4000:] + assert result.stdout.strip().endswith("fa4-work-tile-info-ok") From afa258f82eb3ef34973dde692ddaa0285d0dafaf Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Fri, 4 Sep 2026 13:42:46 +0000 Subject: [PATCH 5/7] refactor: align VisualGen sparse attention workflow Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- docs/source/models/visual-generation.md | 15 +- .../visual-gen/features/sparse-attention.md | 19 +- .../visual_gen/attention_backend/__init__.py | 16 +- .../attention_backend/cute_dsl/__init__.py | 21 +- .../attention_backend/cute_dsl/fmha.py | 16 +- .../attention_backend/cute_dsl/vsa.py | 412 -------- .../visual_gen/attention_backend/parallel.py | 34 +- .../attention_backend/sparse/__init__.py | 16 + .../attention_backend/sparse/vsa/__init__.py | 36 + .../attention_backend/sparse/vsa/backend.py | 343 +++++++ .../attention_backend/sparse/vsa/metadata.py | 231 +++++ .../attention_backend/sparse/vsa/predictor.py | 366 +++++++ .../visual_gen/attention_backend/trtllm.py | 124 ++- .../visual_gen/attention_backend/utils.py | 65 +- tensorrt_llm/_torch/visual_gen/config.py | 22 +- .../video_sparse_attention/interface.py | 2 +- .../visual_gen/models/wan/pipeline_wan.py | 19 +- .../visual_gen/models/wan/transformer_wan.py | 19 +- .../_torch/visual_gen/modules/attention.py | 59 +- .../_torch/visual_gen/pipeline_loader.py | 30 +- tensorrt_llm/visual_gen/args.py | 27 +- tensorrt_llm/visual_gen/sparse_attention.py | 4 +- .../test_lists/test-db/l0_b200.yml | 3 +- .../multi_gpu/test_ulysses_attention.py | 65 ++ .../multi_gpu/test_wan_async_ulysses.py | 51 + .../visual_gen/multi_gpu/test_wan_tp.py | 47 + .../multi_gpu/test_wan_vsa_ulysses.py | 4 +- .../visual_gen/test_attention_cute_dsl_vsa.py | 478 --------- .../visual_gen/test_attention_integration.py | 140 ++- .../_torch/visual_gen/test_attention_perf.py | 13 +- .../_torch/visual_gen/test_attention_vsa.py | 906 ++++++++++++++++++ .../test_trtllm_attention_metadata.py | 364 +++++++ .../_torch/visual_gen/test_visual_gen_args.py | 45 + .../visual_gen/test_wan_vsa_pipeline.py | 2 +- 34 files changed, 2918 insertions(+), 1096 deletions(-) delete mode 100644 tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/vsa.py create mode 100644 tensorrt_llm/_torch/visual_gen/attention_backend/sparse/__init__.py create mode 100644 tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/__init__.py create mode 100644 tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/backend.py create mode 100644 tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/metadata.py create mode 100644 tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/predictor.py delete mode 100644 tests/unittest/_torch/visual_gen/test_attention_cute_dsl_vsa.py create mode 100644 tests/unittest/_torch/visual_gen/test_attention_vsa.py diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index f358de023608..f5edc0f49810 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -73,7 +73,7 @@ Models are auto-detected from the checkpoint directory. Diffusers-format models [^1]: FLUX models use embedded guidance and do not have a separate negative prompt path, so CFG parallelism is not applicable. -[^2]: `FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers` — VSA-fine-tuned checkpoint with learned sparse-attention gates. Requires `CUTEDSL` on Blackwell sm_100+ (falls back to dense SDPA on older hardware). Ring and Attention2D not supported (no LSE output); Ulysses supported. +[^2]: `FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers` — VSA-fine-tuned checkpoint with learned sparse-attention gates. Supports the `CUTEDSL` and `TRTLLM` attention backends; each uses its block-sparse fine stage when supported and otherwise uses its dense path with compact Q/K/V. Ring and Attention2D are not supported (no LSE output); Ulysses is supported. [^3]: Wan 2.2 has two stage transformers; TeaCache requires explicit `teacache.coefficients` (high-noise) and `teacache.coefficients_2` (low-noise). There is no built-in coefficient table for Wan 2.2. @@ -408,12 +408,15 @@ args = VisualGenArgs( ### Video Sparse Attention (VSA) -VSA reduces the compute cost of self-attention in video diffusion models by selectively attending to only the most relevant spatial-temporal blocks. It uses a two-branch design: a lightweight coarse mean-pool branch computes block-level attention scores to identify the top-K most relevant token blocks, then a fine branch runs a block-sparse CuTe kernel over only those blocks. The two outputs are blended with learned gates. +VSA reduces the compute cost of self-attention in video diffusion models by selectively attending to only the most relevant spatial-temporal blocks. It uses a two-branch design: a lightweight coarse mean-pool branch computes block-level attention scores to identify the top-K most relevant token blocks, then a fine branch runs the selected backend's block-sparse kernel over only those blocks. The two outputs are blended with learned gates. + +VisualGen owns VSA route prediction and coarse/fine post-processing. With the `TRTLLM` backend, it nests the predicted routes in `SparseRuntimeParams.block_sparse_inputs` and passes those precomputed runtime parameters through the normal core attention forward. The core `PrimsTSBlockSparseFmha` owns the general block-sparse execution contract; it does not own VSA-specific prediction or blending. **Requirements:** - VSA-fine-tuned checkpoint: [`FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers`](https://huggingface.co/FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers). Standard Wan checkpoints do not have the learned VSA gates. -- Blackwell GPU (sm_100+) for the CuTe JIT kernel. Falls back to dense SDPA on older hardware with no accuracy loss. -- `CUTEDSL` attention backend. +- `CUTEDSL` or `TRTLLM` attention backend. `CUTEDSL` uses the CuTe DSL fine-stage kernel; `TRTLLM` lowers the selected blocks through the generic PrimTS block-sparse FMHA contract. +- A supported CUDA device and tensor shape for the selected block-sparse kernel. When that kernel is unavailable or the input is outside its supported envelope, the fine branch uses the selected backend's compact dense path (`SDPA` for `CUTEDSL`, TRTLLM attention for `TRTLLM`). +- VSA cannot be combined with `quant_attention_config`. - Not compatible with Ring attention or Attention2D (VSA does not produce per-split LSE). Ulysses is supported. **`vsa_sparsity`** controls the fraction of K/V blocks skipped in the fine branch (0.0 = dense, 0.9 = 90% blocks skipped). Higher sparsity gives more speedup at the cost of some quality. @@ -427,7 +430,7 @@ from tensorrt_llm.visual_gen.args import AttentionConfig, VideoSparseAttentionCo args = VisualGenArgs( model="FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers", attention_config=AttentionConfig( - backend="CUTEDSL", + backend="TRTLLM", # Use "CUTEDSL" for the CuTe DSL fine-stage kernel. sparse_attention_config=VideoSparseAttentionConfig(vsa_sparsity=0.9), ), ) @@ -437,7 +440,7 @@ YAML (for use with `--visual_gen_args` or `trtllm-serve`): ```yaml attention_config: - backend: CUTEDSL + backend: TRTLLM # CUTEDSL is also supported. sparse_attention_config: algorithm: vsa vsa_sparsity: 0.90 diff --git a/docs/source/visual-gen/features/sparse-attention.md b/docs/source/visual-gen/features/sparse-attention.md index 7edeff8fb63c..2d5926255190 100644 --- a/docs/source/visual-gen/features/sparse-attention.md +++ b/docs/source/visual-gen/features/sparse-attention.md @@ -13,14 +13,14 @@ This page is an unindexed draft until the VisualGen documentation hub is introdu Visual generation models naturally operate on long image or video token sequences. Each denoising step is closer to a full-context prefill pass than to autoregressive decoding, and attention can dominate runtime for high-resolution image generation or long video generation. -Sparse attention in VisualGen is configured through `VisualGenArgs.attention_config.sparse_attention_config`. The user-facing config stays in VisualGen args or model config. Checkpoint calibration metadata remains internal and is lowered into per-attention-backend `SparseParams` when each attention module is constructed. +Sparse attention in VisualGen is configured through `VisualGenArgs.attention_config.sparse_attention_config`. The user-facing config stays in VisualGen args or model config, while `attention_config.backend` selects the kernel family. Algorithms produce their block-sparse routes through the core `block_sparse_attn_predict` hook: a backend either predicts inside that hook from the flattened Q/K/V, or predicts before the core forward and hands the complete `BlockSparseForwardInputs` through `AttentionForwardArgs.sparse_backend_args`, which the default hook passes through. `SparseRuntimeParams` is the single lowered runtime carrier passed as `AttentionForwardArgs.sparse_runtime_params`; its optional `block_sparse_inputs` field nests the algorithm-neutral routes for the general block-sparse FMHA. `None` means prediction has not run, while an empty `SparseRuntimeParams()` records that prediction ran without a sparse payload. ### Algorithms | `algorithm` | Config class | Status | |---|---|---| | `skip_softmax` | `SkipSoftmaxAttentionConfig` | Supported | -| VSA | TBD | TODO | +| `vsa` | `VideoSparseAttentionConfig` | Supported (`CUTEDSL`, `TRTLLM`) | ## Skip Softmax Attention @@ -216,4 +216,17 @@ Graphs are captured lazily. The first denoising step seen for a given tensor sha ## Video Sparse Attention (VSA) -TODO +VSA combines a coarse mean-pooled branch with a top-K block-sparse fine branch. Select either `CUTEDSL` for the CuTe DSL kernel or `TRTLLM` for PrimTS block-sparse attention. If the selected sparse kernel is unavailable or the known VSA tensor envelope is not met, the fine branch uses the compact Q/K/V tensors with that backend's dense path. VSA cannot be combined with `quant_attention_config`. + +VSA retains shape-dependent metadata and route tensors so CUDA Graph replay can reuse stable addresses. A pipeline instance accepts up to 16 distinct VSA shape profiles; reuse configured resolution/frame profiles or restart the pipeline before serving additional shapes. + +Both VSA backends share one VisualGen-owned predictor and identical +post-processing. The `TRTLLM` path runs the coarse stage before the core +forward, hands the predicted `BlockSparseForwardInputs` (including the +tile-padding validity bits only the VSA predictor knows) through +`sparse_backend_args`, lets the default core prediction hook pass them to the +general block-sparse FMHA, and then blends the fine and coarse outputs. Its +compact dense fallback passes no sparse inputs, so the core runs dense attention +and VSA post-processing still runs. `CUTEDSL` retains only its backend-specific +fine-attention execution. The core FMHA registry owns the reusable block-sparse +implementation rather than a VSA-specific lifecycle. diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py b/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py index b919f17c0a0c..f6efaffa171a 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py @@ -20,15 +20,7 @@ simplified metadata that doesn't require KV caching. """ -from .cute_dsl import ( - VSA_TILE_SIZE, - CuTeDSLAttention, - VSAAttention, - VSAMetadata, - VSAMetadataBuilder, - get_vsa_forward_context, - set_vsa_forward_context, -) +from .cute_dsl import CuTeDSLAttention from .flash_attn4 import FlashAttn4Attention from .flashinfer import FlashInferAttention from .interface import AttentionBackend, AttentionTensorLayout @@ -38,7 +30,6 @@ from .vanilla import VanillaAttention __all__ = [ - "VSA_TILE_SIZE", "Attention2DAttention", "AttentionBackend", "AttentionTensorLayout", @@ -49,13 +40,8 @@ "TrtllmAttention", "TrtllmAttentionMetadata", "UlyssesAttention", - "VSAAttention", - "VSAMetadata", - "VSAMetadataBuilder", "VanillaAttention", "create_attention", "get_visual_gen_attention_backend", - "get_vsa_forward_context", - "set_vsa_forward_context", "wrap_parallel_attention", ] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/__init__.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/__init__.py index 9b70421c3b81..1dde6e1f9aab 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/__init__.py @@ -15,31 +15,12 @@ """ CuTe DSL attention backend family for visual generation models. - fmha.py — CuTeDSLAttention (dense and blockscaled JIT FMHA) - vsa.py — VSAAttention (Video Sparse Attention, CuTe JIT + SDPA fallback) + fmha.py — CuTeDSLAttention (dense and blockscaled JIT FMHA) """ from .fmha import CuTeDSLAttention, _cute_dsl_import_error -from .vsa import ( - VSA_KERNEL_MAX_CUBES, - VSA_TILE_SIZE, - VSAAttention, - VSAMetadata, - VSAMetadataBuilder, - VSAPreprocessor, - get_vsa_forward_context, - set_vsa_forward_context, -) __all__ = [ "CuTeDSLAttention", - "VSAAttention", - "VSAMetadata", - "VSAMetadataBuilder", - "VSAPreprocessor", - "VSA_TILE_SIZE", - "VSA_KERNEL_MAX_CUBES", - "set_vsa_forward_context", - "get_vsa_forward_context", "_cute_dsl_import_error", ] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py index e9638ef9075c..a7bbc7baa2cf 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py @@ -17,7 +17,7 @@ JIT-compiles dense or SkipSoftmax FMHA and caches the compiled artifact for each kernel configuration. Expects NHD layout ([B, S, H, D]) and supports float16/bfloat16 inputs. The VSA -sparse path uses VSAAttention from vsa.py instead. +sparse backend uses `VSACuTeDSLAttention` in `attention_backend.sparse.vsa` instead. """ import math @@ -72,20 +72,6 @@ def _resolve_skip_softmax_threshold_scale_factor( ) -> float | None: """Resolve the active CuTeDSL threshold for the current denoising phase.""" if sparse_params is not None: - if timestep is None and sparse_params.scheduler.disabled_until_timestep is not None: - # Fail-open: a missing timestep resolves to the full (unthrottled) - # threshold, i.e. skip-softmax runs during the high-noise steps - # `disabled_until_timestep` exists to protect. This is silent - # elsewhere (a quality regression, not an error), so surface it - # once per process instead of only in this function's return value. - logger.warning_once( - "SkipSoftmax scheduler has disabled_until_timestep=" - f"{sparse_params.scheduler.disabled_until_timestep} configured, but no " - "`timestep` was passed to the CuTeDSL attention forward call. Skip-softmax " - "will run unthrottled (as if past the cutoff) until `timestep` is threaded " - "through.", - key="cute_dsl_skip_softmax_missing_timestep", - ) runtime_params = sparse_params.scheduler.get_runtime_params(timestep=timestep) threshold_scale_factor = runtime_params.threshold_scale_factor_prefill if threshold_scale_factor is None or threshold_scale_factor <= 0.0: diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/vsa.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/vsa.py deleted file mode 100644 index 741f82097ca2..000000000000 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/vsa.py +++ /dev/null @@ -1,412 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -""" -Video Sparse Attention (VSA) backend for visual generation models. - -VSAAttention implements hierarchical sparse attention: - - Coarse branch: mean-pooled cube attention (always dense) - - Fine branch: block-sparse top-K attention via CuTe JIT kernel (sm100+) - or dense SDPA fallback when CuTe is unavailable / head_dim != 128. -""" - -import contextvars -from contextlib import contextmanager -from dataclasses import dataclass -from math import ceil -from typing import Dict, Optional, Tuple - -import torch -import torch.nn.functional as F - -from ..interface import AttentionBackend, AttentionTensorLayout - -_vsa_import_error = None -try: - from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.video_sparse_attention import ( - block_sparse_attn_from_indices_cute, - is_cute_supported, - ) -except (ImportError, OSError) as e: - block_sparse_attn_from_indices_cute = None - is_cute_supported = None - _vsa_import_error = e - - -# Must match the Blackwell kernel's block_size expectation. -VSA_TILE_SIZE: Tuple[int, int, int] = (4, 4, 4) - -# Kernel's SMEM buffer for variable_block_sizes is fixed-size and unchecked, -# so num_cubes must stay <= this. -VSA_KERNEL_MAX_CUBES: int = 4 * 1024 - - -def _get_tile_partition_indices( - dit_seq_shape: Tuple[int, int, int], - tile_size: Tuple[int, int, int], - device: torch.device, -) -> torch.LongTensor: - T, H, W = dit_seq_shape - tT, tH, tW = tile_size - nT, nH, nW = ceil(T / tT), ceil(H / tH), ceil(W / tW) - - bt = torch.arange(nT, device=device).view(nT, 1, 1, 1, 1, 1) - bh = torch.arange(nH, device=device).view(1, nH, 1, 1, 1, 1) - bw = torch.arange(nW, device=device).view(1, 1, nW, 1, 1, 1) - lt = torch.arange(tT, device=device).view(1, 1, 1, tT, 1, 1) - lh = torch.arange(tH, device=device).view(1, 1, 1, 1, tH, 1) - lw = torch.arange(tW, device=device).view(1, 1, 1, 1, 1, tW) - - gt = bt * tT + lt - gh = bh * tH + lh - gw = bw * tW + lw - valid = (gt < T) & (gh < H) & (gw < W) - flat = gt * (H * W) + gh * W + gw - out = torch.where(valid, flat, torch.full_like(flat, -1)) - return out.reshape(-1).to(torch.long) - - -def _construct_variable_block_sizes( - dit_seq_shape: Tuple[int, int, int], - num_tiles: Tuple[int, int, int], - tile_size: Tuple[int, int, int], - device: torch.device, -) -> torch.LongTensor: - T, H, W = dit_seq_shape - tT, tH, tW = tile_size - nT, nH, nW = num_tiles - - bt = torch.arange(nT, device=device) - bh = torch.arange(nH, device=device) - bw = torch.arange(nW, device=device) - valid_t = (T - bt * tT).clamp(max=tT) - valid_h = (H - bh * tH).clamp(max=tH) - valid_w = (W - bw * tW).clamp(max=tW) - sizes = valid_t.view(nT, 1, 1) * valid_h.view(1, nH, 1) * valid_w.view(1, 1, nW) - return sizes.reshape(-1).to(torch.long) - - -@dataclass -class VSAMetadata: - """Per-timestep metadata required by the VSA sparse path.""" - - current_timestep: int - dit_seq_shape: Tuple[int, int, int] - vsa_sparsity: float - num_tiles: Tuple[int, int, int] - total_seq_length: int - padded_seq_length: int - tile_partition_indices: torch.LongTensor - reverse_tile_partition_indices: torch.LongTensor - variable_block_sizes: torch.LongTensor - non_pad_index: torch.LongTensor - gather_idx: torch.LongTensor - - -class VSAMetadataBuilder: - """Builds VSAMetadata; caches per-shape index tensors so torch.compile - guards stay stable across denoising steps.""" - - def __init__(self) -> None: - self._cache: Dict[Tuple[Tuple[int, int, int], str], Dict[str, object]] = {} - - def _build_shape_payload( - self, - dit_seq_shape: Tuple[int, int, int], - device: torch.device, - ) -> Dict[str, object]: - T, H, W = dit_seq_shape - tT, tH, tW = VSA_TILE_SIZE - num_tiles = (ceil(T / tT), ceil(H / tH), ceil(W / tW)) - total_seq_length = T * H * W - padded_seq_length = num_tiles[0] * num_tiles[1] * num_tiles[2] * tT * tH * tW - - tile_partition_indices = _get_tile_partition_indices(dit_seq_shape, VSA_TILE_SIZE, device) - non_pad_index = (tile_partition_indices >= 0).nonzero(as_tuple=True)[0] - gather_idx = tile_partition_indices[non_pad_index] - - reverse = torch.zeros(total_seq_length, dtype=torch.long, device=device) - reverse[gather_idx] = torch.arange(len(non_pad_index), dtype=torch.long, device=device) - - variable_block_sizes = _construct_variable_block_sizes( - dit_seq_shape, num_tiles, VSA_TILE_SIZE, device - ) - - return { - "dit_seq_shape": dit_seq_shape, - "num_tiles": num_tiles, - "total_seq_length": total_seq_length, - "padded_seq_length": padded_seq_length, - "tile_partition_indices": tile_partition_indices, - "reverse_tile_partition_indices": reverse, - "variable_block_sizes": variable_block_sizes, - "non_pad_index": non_pad_index, - "gather_idx": gather_idx, - } - - def build( - self, - current_timestep: int, - raw_latent_shape: Tuple[int, int, int], - patch_size: Tuple[int, int, int], - vsa_sparsity: float, - device: torch.device, - ) -> VSAMetadata: - dit_seq_shape = ( - raw_latent_shape[0] // patch_size[0], - raw_latent_shape[1] // patch_size[1], - raw_latent_shape[2] // patch_size[2], - ) - cache_key = (dit_seq_shape, str(device)) - payload = self._cache.get(cache_key) - if payload is None: - payload = self._build_shape_payload(dit_seq_shape, device) - self._cache[cache_key] = payload - - return VSAMetadata( - current_timestep=current_timestep, - vsa_sparsity=vsa_sparsity, - **payload, # type: ignore[arg-type] - ) - - -_vsa_forward_context_var: contextvars.ContextVar[Optional[VSAMetadata]] = contextvars.ContextVar( - "_vsa_forward_context", default=None -) - - -@contextmanager -def set_vsa_forward_context(metadata: VSAMetadata): - token = _vsa_forward_context_var.set(metadata) - try: - yield - finally: - _vsa_forward_context_var.reset(token) - - -def get_vsa_forward_context() -> Optional[VSAMetadata]: - return _vsa_forward_context_var.get(None) - - -def _mean_pool_cubes( - x_tiled: torch.Tensor, - variable_block_sizes: torch.LongTensor, - prod_tile: int, - num_cubes: int, -) -> torch.Tensor: - B, _padded, H, D = x_tiled.shape - x_cubes = x_tiled.view(B, num_cubes, prod_tile, H, D) - # fp32 accumulation: bf16 sum over 64 tokens perturbs the coarse softmax. - x_sum = x_cubes.float().sum(dim=2) - valid_counts = variable_block_sizes.float().clamp(min=1).view(1, num_cubes, 1, 1) - return (x_sum / valid_counts).to(x_tiled.dtype) - - -class VSAPreprocessor: - """Reorders NHD tokens into tile-major layout and zero-pads to tile boundaries.""" - - @staticmethod - def tile( - x: torch.Tensor, - non_pad_index: torch.LongTensor, - gather_idx: torch.LongTensor, - padded_seq_len: int, - ) -> torch.Tensor: - # index_select + index_copy_ instead of chained advanced indexing so - # torch.compile can trace this without a graph break. - B, _S, H, D = x.shape - x_valid = x.index_select(1, gather_idx) - x_padded = x.new_zeros(B, padded_seq_len, H, D) - x_padded.index_copy_(1, non_pad_index, x_valid) - return x_padded - - @staticmethod - def untile( - x: torch.Tensor, - reverse_tile_partition_indices: torch.LongTensor, - non_pad_index: torch.LongTensor, - ) -> torch.Tensor: - return x.index_select(1, non_pad_index).index_select(1, reverse_tile_partition_indices) - - -class VSAAttention(AttentionBackend): - """ - Video Sparse Attention (VSA) backend for diffusion models. - - Implements coarse mean-pool + fine block-sparse top-K attention. - The fine branch uses a JIT-compiled CuTe kernel on sm100+ for - head_dim=128 / fp16-bf16; otherwise falls back to dense SDPA. - - Requires an active VSA forward context (set_vsa_forward_context) during - each forward call. Does not support LSE output. - """ - - def __init__( - self, - layer_idx: int = 0, - num_heads: int = 8, - head_dim: int = 128, - num_kv_heads: Optional[int] = None, - dtype: Optional[torch.dtype] = None, - sparse_attention_config=None, - **kwargs, - ): - self.layer_idx = layer_idx - self.num_heads = num_heads - self.head_dim = head_dim - self.num_kv_heads = num_kv_heads or num_heads - assert self.num_kv_heads == self.num_heads, ( - f"VSA coarse mean-pool assumes MHA (num_kv_heads == num_heads), " - f"got num_kv_heads={self.num_kv_heads}, num_heads={self.num_heads}. " - f"GQA/MQA is not supported." - ) - self.dtype = dtype - self.sparse_attention_config = sparse_attention_config - - # Dynamo can't guard on the module-level mutable global, so this read - # runs in eager. - @torch.compiler.disable - def _get_vsa_inputs(self): - ctx: Optional[VSAMetadata] = get_vsa_forward_context() - if ctx is None: - raise RuntimeError( - "VSAAttention.forward called without an active VSA forward context. " - "Wrap each transformer call with set_vsa_forward_context()." - ) - return ( - ctx.non_pad_index, - ctx.gather_idx, - ctx.reverse_tile_partition_indices, - ctx.variable_block_sizes, - ctx.padded_seq_length, - ctx.num_tiles[0] * ctx.num_tiles[1] * ctx.num_tiles[2], - ctx.vsa_sparsity, - ) - - def forward( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - *, - gate_compress: Optional[torch.Tensor] = None, - gate_fine: Optional[torch.Tensor] = None, - **kwargs, - ) -> torch.Tensor: - """ - VSA forward: coarse mean-pool + fine block-sparse top-K. - - Args: - q, k, v: [B, S, H, D] in original (un-tiled) token order. - gate_compress: [B, S, H, D] G_c gate weighting the coarse branch O_c. - gate_fine: Optional [B, S, H, D] G_f gate weighting the fine branch - O_f. None means constant 1 (dense behavior preserved). - - Returns: - [B, S, H, D] in the same original token order. - """ - if gate_compress is None: - raise ValueError( - "VSAAttention requires gate_compress. " - "Ensure to_gate_compress is wired in the transformer block." - ) - - ( - non_pad_index, - gather_idx, - reverse_tile_partition_indices, - variable_block_sizes, - padded_len, - num_cubes, - vsa_sparsity, - ) = self._get_vsa_inputs() - - B, S, H, D = q.shape - prod_tile = VSA_TILE_SIZE[0] * VSA_TILE_SIZE[1] * VSA_TILE_SIZE[2] - cur_topk = max(1, ceil((1.0 - vsa_sparsity) * num_cubes)) - - q_t = VSAPreprocessor.tile(q, non_pad_index, gather_idx, padded_len) - k_t = VSAPreprocessor.tile(k, non_pad_index, gather_idx, padded_len) - v_t = VSAPreprocessor.tile(v, non_pad_index, gather_idx, padded_len) - - q_c = _mean_pool_cubes(q_t, variable_block_sizes, prod_tile, num_cubes) - k_c = _mean_pool_cubes(k_t, variable_block_sizes, prod_tile, num_cubes) - v_c = _mean_pool_cubes(v_t, variable_block_sizes, prod_tile, num_cubes) - - scale = D**-0.5 - scores_c = torch.einsum("bnhd,bmhd->bhnm", q_c, k_c) * scale - attn_probs_c = scores_c.softmax(dim=-1) - o_c = torch.einsum("bhnm,bmhd->bnhd", attn_probs_c, v_c) - - use_cute = ( - _vsa_import_error is None - and is_cute_supported(q) - and (q.dtype == k.dtype == v.dtype) - and num_cubes <= VSA_KERNEL_MAX_CUBES - ) - topk_indices = attn_probs_c.topk(cur_topk, dim=-1).indices.to(torch.int32) - - o_c_tiled = ( - o_c.unsqueeze(2).expand(B, num_cubes, prod_tile, H, D).reshape(B, padded_len, H, D) - ) - - if use_cute: - q_hnd = q_t.transpose(1, 2).contiguous() - k_hnd = k_t.transpose(1, 2).contiguous() - v_hnd = v_t.transpose(1, 2).contiguous() - q2k_num = torch.full((B, H, num_cubes), cur_topk, dtype=torch.int32, device=q.device) - o_hnd, _lse = block_sparse_attn_from_indices_cute( - q_hnd, - k_hnd, - v_hnd, - q2k_idx=topk_indices.contiguous(), - q2k_num=q2k_num, - variable_block_sizes=variable_block_sizes.to(torch.int32), - ) - o_f_tiled = o_hnd.transpose(1, 2) - - # Padded rows hold kernel garbage; zero-padded gates mask the coarse - # term and untile discards padded positions from both branches. - gate_c_t = VSAPreprocessor.tile(gate_compress, non_pad_index, gather_idx, padded_len) - if gate_fine is not None: - gate_f_t = VSAPreprocessor.tile(gate_fine, non_pad_index, gather_idx, padded_len) - combined_tiled = gate_c_t * o_c_tiled + gate_f_t * o_f_tiled - else: - combined_tiled = gate_c_t * o_c_tiled + o_f_tiled - return VSAPreprocessor.untile( - combined_tiled, reverse_tile_partition_indices, non_pad_index - ) - - # SDPA must run on the un-tiled Q/K/V — padded zero K/V slots would - # otherwise absorb softmax mass and pollute the output. Untile o_c so - # both branches combine in original-flat order. - o_c_full = VSAPreprocessor.untile(o_c_tiled, reverse_tile_partition_indices, non_pad_index) - o_f = F.scaled_dot_product_attention( - q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) - ).transpose(1, 2) - if gate_fine is not None: - return gate_compress * o_c_full + gate_fine * o_f - return gate_compress * o_c_full + o_f - - @classmethod - def support_lse(cls) -> bool: - return False - - @property - def preferred_layout(self) -> AttentionTensorLayout: - return AttentionTensorLayout.NHD - - @classmethod - def support_fused_qkv(cls) -> bool: - return False diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py index a2e1a78d2171..8f9ec0a21afd 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/parallel.py @@ -69,9 +69,9 @@ class UlyssesAttention(AttentionBackend): Wraps any attention backend with sequence parallelism via all-to-all. Not a standalone backend -- compose around a real backend (VANILLA/TRTLLM). - Fully transparent to backend-specific kwargs: everything in ``**kwargs`` - is forwarded to the inner backend unchanged (except ``seq_len`` which is - overridden with the post-all-to-all value). + Backend-specific kwargs are forwarded to the inner backend. Sequence + lengths are updated after all-to-all, and VSA gates are redistributed with + the same sequence/head mapping as Q before they are forwarded. Architecture: Input: [B, S/P, H, D] (sequence sharded across P processes) @@ -340,6 +340,23 @@ def forward_async( self._join_async() q_5d, k_5d, v_5d = recv["q"], recv["k"], recv["v"] + gate_compress = attn_kwargs.pop("gate_compress", None) + gate_fine = attn_kwargs.pop("gate_fine", None) + if gate_compress is not None: + gate_compress = all_to_all_4d( + gate_compress, + scatter_dim=2, + gather_dim=1, + process_group=self.process_group, + ) + if gate_fine is not None: + gate_fine = all_to_all_4d( + gate_fine, + scatter_dim=2, + gather_dim=1, + process_group=self.process_group, + ) + # Fast path: one fused kernel replaces the eager post-A2A chain # (6 ops for HND target: permute+reshape+contig + transpose+contig # per Q/K/V; 3 ops for NHD target). bf16-only because the kernel is @@ -365,8 +382,19 @@ def forward_async( k_out = k_out.transpose(1, 2).contiguous() v_out = v_out.transpose(1, 2).contiguous() + if is_hnd: + if gate_compress is not None: + gate_compress = gate_compress.transpose(1, 2) + if gate_fine is not None: + gate_fine = gate_fine.transpose(1, 2) + + attn_kwargs["batch_size"] = B attn_kwargs["seq_len"] = seq_len_full attn_kwargs["seq_len_kv"] = seq_len_kv_full + if gate_compress is not None: + attn_kwargs["gate_compress"] = gate_compress + if gate_fine is not None: + attn_kwargs["gate_fine"] = gate_fine output = self.inner_backend.forward(q=q_out, k=k_out, v=v_out, **attn_kwargs) return self._output_a2a(output, B, seq_len_full) diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/__init__.py b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/__init__.py new file mode 100644 index 000000000000..43600877b99e --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/__init__.py @@ -0,0 +1,16 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Sparse-attention backend families for VisualGen.""" diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/__init__.py b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/__init__.py new file mode 100644 index 000000000000..2748d3cc0686 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/__init__.py @@ -0,0 +1,36 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Video Sparse Attention backends and shared prediction for VisualGen.""" + +from .metadata import ( + VSA_TILE_SIZE, + VSAMetadata, + VSAMetadataBuilder, + get_vsa_forward_context, + set_vsa_forward_context, +) +from .predictor import VSAForwardInputs, VSAPredictor, VSAPreprocessor + +__all__ = [ + "VSA_TILE_SIZE", + "VSAForwardInputs", + "VSAMetadata", + "VSAMetadataBuilder", + "VSAPredictor", + "VSAPreprocessor", + "get_vsa_forward_context", + "set_vsa_forward_context", +] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/backend.py b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/backend.py new file mode 100644 index 000000000000..1f6af9b72ecc --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/backend.py @@ -0,0 +1,343 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""VisualGen TRTLLM-first attention backends for Video Sparse Attention.""" + +from typing import Optional + +import torch +import torch.nn.functional as F + +from tensorrt_llm.logger import logger +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.visual_gen.args import QuantAttentionConfig + +from .....attention.backends.fmha.prims_ts_block_sparse import PrimsTSBlockSparseFmha +from .....attention.backends.interface import PredefinedAttentionMask +from .....attention.backends.sparse.params import SparseBackendForwardArgs +from ...cute_dsl import CuTeDSLAttention +from ...trtllm import TrtllmAttention +from .metadata import VSA_BLOCK_SIZE, VSAMetadata +from .predictor import VSAForwardInputs, VSAPredictor, vsa_post_process + +_vsa_import_error = None +try: + from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.video_sparse_attention import ( + block_sparse_attn_from_indices_cute, + is_cute_supported, + ) +except (ImportError, OSError) as error: + block_sparse_attn_from_indices_cute = None + is_cute_supported = None + _vsa_import_error = error + + +VSA_KERNEL_MAX_CUBES: int = 4 * 1024 + + +def _normalize_qkv_inputs( + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Normalize separate BSHD or Ulysses-packed BSH3HD inputs.""" + + if k is not None and v is not None: + return q, k, v + if k is not None or v is not None: + raise ValueError("VSA requires complete separate Q/K/V or one packed QKV tensor.") + if q.ndim != 5 or q.shape[2] != 3: + raise ValueError("VSA packed QKV must have shape [B, S, 3, H, D].") + return q.unbind(dim=2) + + +def _get_unsupported_primts_reason( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + metadata: VSAMetadata, +) -> str | None: + if q.shape != k.shape or q.shape != v.shape: + return "VSA PrimTS requires matching MHA Q/K/V shapes" + if q.device.type != "cuda": + return f"VSA PrimTS requires CUDA tensors, got {q.device}" + if q.dtype not in (torch.float16, torch.bfloat16): + return f"VSA PrimTS requires FP16 or BF16 tensors, got {q.dtype}" + batch_size, seq_len, num_heads, head_dim = map(int, q.shape) + if min(batch_size, seq_len, num_heads, metadata.num_cubes) <= 0: + return "VSA PrimTS requires positive batch, sequence, head, and cube extents" + if head_dim != 128: + return f"VSA PrimTS requires head_dim=128, got {head_dim}" + if metadata.padded_seq_length != metadata.num_cubes * VSA_BLOCK_SIZE: + return "VSA tiled sequence length must match its 64-token cube count" + if batch_size > 65535 or num_heads > 65535: + return "VSA PrimTS batch and head dimensions must fit the CUDA grid" + return None + + +class VSATrtllmAttention(TrtllmAttention): + """TRTLLM VSA backend using the generic block-sparse forward lifecycle.""" + + def __init__( + self, + layer_idx: int = 0, + num_heads: int = 8, + head_dim: int = 64, + num_kv_heads: Optional[int] = None, + quant_config: Optional[QuantConfig] = None, + dtype: Optional[torch.dtype] = None, + max_batch_size: int = 16, + max_seq_len: int = 4096, + quant_attention_config: Optional[QuantAttentionConfig] = None, + attention_metadata_state: Optional[dict] = None, + ) -> None: + num_kv_heads = num_kv_heads or num_heads + super().__init__( + layer_idx=layer_idx, + num_heads=num_heads, + head_dim=head_dim, + num_kv_heads=num_kv_heads, + quant_config=quant_config, + dtype=dtype, + max_batch_size=max_batch_size, + max_seq_len=max_seq_len, + quant_attention_config=quant_attention_config, + attention_metadata_state=attention_metadata_state, + sparse_params=None, + ) + assert attention_metadata_state is not None + predictor_cache = attention_metadata_state.setdefault("sparse_predictors", {}) + predictor_key = ("vsa", num_heads, num_kv_heads) + predictor = predictor_cache.get(predictor_key) + if predictor is None: + predictor = VSAPredictor( + num_heads=num_heads, + num_kv_heads=num_kv_heads, + ) + predictor_cache[predictor_key] = predictor + elif not isinstance(predictor, VSAPredictor): + raise TypeError("model-scoped VSA predictor cache contains an invalid value") + self.predictor = predictor + + def forward( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + batch_size: int, + seq_len: int, + attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL, + seq_len_kv: Optional[int] = None, + **kwargs, + ) -> torch.Tensor: + """Run the VSA coarse stage, then the fine stage through the core forward. + + The coarse stage predicts the complete block-sparse payload, which is + handed to the core prediction hook via ``sparse_backend_args``. The fine + output is blended with the coarse output afterward. + """ + + q, k, v = _normalize_qkv_inputs(q, k, v) + metadata = self.predictor.get_metadata() + use_primts = any( + isinstance(fmha, PrimsTSBlockSparseFmha) for fmha in self._fmha_manager.fmha_libs + ) + unsupported_reason = _get_unsupported_primts_reason(q, k, v, metadata) + if self.quant_attention_config is not None: + unsupported_reason = "VSA PrimTS does not support quant_attention_config" + if not use_primts: + logger.warning_once( + "TRTLLM VSA cannot use PrimTS block-sparse attention because the " + "prims_ts_block_sparse FMHA library is unavailable; using the compact " + "dense TRTLLM fine stage.", + key="trtllm_vsa_primts_unavailable", + ) + elif unsupported_reason is not None: + logger.warning_once( + "TRTLLM VSA cannot use PrimTS block-sparse attention: " + f"{unsupported_reason}; using the compact dense TRTLLM fine stage.", + key=("trtllm_vsa_primts_unsupported_envelope", unsupported_reason), + ) + use_sparse_fine = use_primts and unsupported_reason is None + + inputs = self.predictor.predict( + q, + k, + v, + batch_size=batch_size, + seq_len=seq_len, + seq_len_kv=seq_len if seq_len_kv is None else seq_len_kv, + attention_mask=attention_mask, + gate_compress=kwargs.pop("gate_compress", None), + gate_fine=kwargs.pop("gate_fine", None), + use_sparse_fine=use_sparse_fine, + produce_block_sparse_inputs=use_sparse_fine, + metadata=metadata, + ) + sparse_backend_args = None + if inputs.block_sparse_inputs is not None: + sparse_backend_args = SparseBackendForwardArgs( + block_sparse_inputs=inputs.block_sparse_inputs, + ) + fine_output = super().forward( + inputs.q, + inputs.k, + inputs.v, + inputs.batch_size, + inputs.seq_len, + attention_mask=attention_mask, + seq_len_kv=inputs.seq_len, + sparse_backend_args=sparse_backend_args, + **kwargs, + ) + combined = vsa_post_process(fine_output, inputs) + return combined.reshape(combined.shape[0], combined.shape[1], -1) + + @classmethod + def support_fused_qkv(cls) -> bool: + return True + + +class VSACuTeDSLAttention(CuTeDSLAttention): + """CuTe DSL VSA backend reusing TRTLLM's predictor and post-processing.""" + + def __init__( + self, + layer_idx: int = 0, + num_heads: int = 8, + head_dim: int = 128, + num_kv_heads: Optional[int] = None, + dtype: Optional[torch.dtype] = None, + **kwargs, + ) -> None: + super().__init__( + layer_idx=layer_idx, + num_heads=num_heads, + head_dim=head_dim, + num_kv_heads=num_kv_heads, + dtype=dtype, + **kwargs, + ) + self.predictor = VSAPredictor( + num_heads=num_heads, + num_kv_heads=num_kv_heads, + ) + + def forward( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + *, + attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL, + **kwargs, + ) -> torch.Tensor: + q, k, v = _normalize_qkv_inputs(q, k, v) + gate_compress = kwargs.pop("gate_compress", None) + gate_fine = kwargs.pop("gate_fine", None) + expected_extents = { + "batch_size": int(q.shape[0]), + "seq_len": int(q.shape[1]), + "seq_len_kv": int(k.shape[1]), + } + for name, expected in expected_extents.items(): + actual = kwargs.pop(name, expected) + if not isinstance(actual, int) or isinstance(actual, bool) or actual != expected: + raise ValueError(f"VSA {name}={actual!r} does not match Q/K/V extent {expected}") + kwargs.pop("timestep", None) + if kwargs: + unexpected_names = ", ".join(sorted(kwargs)) + raise TypeError(f"Unexpected CuTeDSL VSA forward keyword arguments: {unexpected_names}") + + metadata = self.predictor.get_metadata() + # The CuTe kernel's fixed launch topology is bounded by the number of + # VSA cubes; larger shapes retain identical VSA math via dense SDPA. + use_cute = ( + _vsa_import_error is None + and is_cute_supported is not None + and is_cute_supported(q) + and q.dtype == k.dtype == v.dtype + and metadata.num_cubes <= VSA_KERNEL_MAX_CUBES + ) + inputs = self.predictor.predict( + q, + k, + v, + batch_size=int(q.shape[0]), + seq_len=int(q.shape[1]), + seq_len_kv=int(k.shape[1]), + attention_mask=attention_mask, + gate_compress=gate_compress, + gate_fine=gate_fine, + use_sparse_fine=use_cute, + produce_block_sparse_inputs=False, + metadata=metadata, + ) + if use_cute: + fine_output = self._execute_sparse_fine(inputs) + else: + fine_output = F.scaled_dot_product_attention( + inputs.q.transpose(1, 2), + inputs.k.transpose(1, 2), + inputs.v.transpose(1, 2), + ).transpose(1, 2) + return vsa_post_process(fine_output, inputs) + + def _execute_sparse_fine(self, inputs: VSAForwardInputs) -> torch.Tensor: + """Execute only the CuTe-specific VSA fine kernel.""" + + q_hnd = inputs.q.transpose(1, 2).contiguous() + k_hnd = inputs.k.transpose(1, 2).contiguous() + v_hnd = inputs.v.transpose(1, 2).contiguous() + q2k_num = torch.full( + (inputs.batch_size, q_hnd.shape[1], inputs.num_cubes), + inputs.cur_topk, + dtype=torch.int32, + device=inputs.q.device, + ) + output_hnd, _lse = block_sparse_attn_from_indices_cute( + q_hnd, + k_hnd, + v_hnd, + q2k_idx=inputs.topk_indices.contiguous(), + q2k_num=q2k_num, + variable_block_sizes=inputs.variable_block_sizes.to(torch.int32), + ) + return output_hnd.transpose(1, 2) + + def forward_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor]: + raise NotImplementedError("CuTe DSL VSA does not support LSE output.") + + @classmethod + def support_fused_qkv(cls) -> bool: + return True + + @classmethod + def support_lse(cls) -> bool: + return False + + +__all__ = [ + "VSACuTeDSLAttention", + "VSATrtllmAttention", + "VSA_KERNEL_MAX_CUBES", +] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/metadata.py b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/metadata.py new file mode 100644 index 000000000000..de88e3b6d282 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/metadata.py @@ -0,0 +1,231 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shape, policy, and forward-context metadata for Video Sparse Attention.""" + +import contextvars +from contextlib import contextmanager +from dataclasses import dataclass +from math import ceil +from typing import Iterator, Optional, Tuple, TypedDict + +import torch + +# A 4x4x4 cube is one 64-token sparse block for every VSA fine-stage backend. +VSA_TILE_SIZE: Tuple[int, int, int] = (4, 4, 4) +VSA_BLOCK_SIZE = VSA_TILE_SIZE[0] * VSA_TILE_SIZE[1] * VSA_TILE_SIZE[2] +_DEFAULT_MAX_CACHED_SHAPES = 16 + + +def _get_tile_partition_indices( + dit_seq_shape: Tuple[int, int, int], + tile_size: Tuple[int, int, int], + device: torch.device, +) -> torch.LongTensor: + time, height, width = dit_seq_shape + tile_time, tile_height, tile_width = tile_size + num_time = ceil(time / tile_time) + num_height = ceil(height / tile_height) + num_width = ceil(width / tile_width) + + block_time = torch.arange(num_time, device=device).view(num_time, 1, 1, 1, 1, 1) + block_height = torch.arange(num_height, device=device).view(1, num_height, 1, 1, 1, 1) + block_width = torch.arange(num_width, device=device).view(1, 1, num_width, 1, 1, 1) + local_time = torch.arange(tile_time, device=device).view(1, 1, 1, tile_time, 1, 1) + local_height = torch.arange(tile_height, device=device).view(1, 1, 1, 1, tile_height, 1) + local_width = torch.arange(tile_width, device=device).view(1, 1, 1, 1, 1, tile_width) + + global_time = block_time * tile_time + local_time + global_height = block_height * tile_height + local_height + global_width = block_width * tile_width + local_width + valid = (global_time < time) & (global_height < height) & (global_width < width) + flat = global_time * (height * width) + global_height * width + global_width + indices = torch.where(valid, flat, torch.full_like(flat, -1)) + return indices.reshape(-1).to(torch.long) + + +def _construct_variable_block_sizes( + dit_seq_shape: Tuple[int, int, int], + num_tiles: Tuple[int, int, int], + tile_size: Tuple[int, int, int], + device: torch.device, +) -> torch.LongTensor: + time, height, width = dit_seq_shape + tile_time, tile_height, tile_width = tile_size + num_time, num_height, num_width = num_tiles + + block_time = torch.arange(num_time, device=device) + block_height = torch.arange(num_height, device=device) + block_width = torch.arange(num_width, device=device) + valid_time = (time - block_time * tile_time).clamp(max=tile_time) + valid_height = (height - block_height * tile_height).clamp(max=tile_height) + valid_width = (width - block_width * tile_width).clamp(max=tile_width) + sizes = ( + valid_time.view(num_time, 1, 1) + * valid_height.view(1, num_height, 1) + * valid_width.view(1, 1, num_width) + ) + return sizes.reshape(-1).to(torch.long) + + +@dataclass(frozen=True, slots=True) +class VSAMetadata: + """Per-step policy and shape metadata required by the VSA sparse path.""" + + current_timestep: int + vsa_sparsity: float + num_cubes: int + padded_seq_length: int + variable_block_sizes: torch.LongTensor + kv_token_mask: torch.BoolTensor + non_pad_index: torch.LongTensor + gather_idx: torch.LongTensor + untile_idx: torch.LongTensor + + +class _VSAShapeMetadata(TypedDict): + num_cubes: int + padded_seq_length: int + variable_block_sizes: torch.LongTensor + kv_token_mask: torch.BoolTensor + non_pad_index: torch.LongTensor + gather_idx: torch.LongTensor + untile_idx: torch.LongTensor + + +class VSAMetadataBuilder: + """Build VSA metadata while caching shape-dependent index tensors.""" + + def __init__(self, max_cached_shapes: int = _DEFAULT_MAX_CACHED_SHAPES) -> None: + if max_cached_shapes <= 0: + raise ValueError("max_cached_shapes must be positive") + self._max_cached_shapes = max_cached_shapes + self._cache: dict[Tuple[Tuple[int, int, int], torch.device], _VSAShapeMetadata] = {} + + def _build_metadata( + self, + dit_seq_shape: Tuple[int, int, int], + device: torch.device, + ) -> _VSAShapeMetadata: + time, height, width = dit_seq_shape + tile_time, tile_height, tile_width = VSA_TILE_SIZE + num_tiles = ( + ceil(time / tile_time), + ceil(height / tile_height), + ceil(width / tile_width), + ) + total_seq_length = time * height * width + padded_seq_length = ( + num_tiles[0] * num_tiles[1] * num_tiles[2] * tile_time * tile_height * tile_width + ) + num_cubes = num_tiles[0] * num_tiles[1] * num_tiles[2] + tokens_per_cube = VSA_BLOCK_SIZE + + tile_partition_indices = _get_tile_partition_indices(dit_seq_shape, VSA_TILE_SIZE, device) + gather_idx = tile_partition_indices[tile_partition_indices >= 0] + + variable_block_sizes = _construct_variable_block_sizes( + dit_seq_shape, num_tiles, VSA_TILE_SIZE, device + ) + local_offsets = torch.arange(tokens_per_cube, device=device).expand( + num_cubes, tokens_per_cube + ) + cube_offsets = torch.arange(num_cubes, device=device).unsqueeze(1) * tokens_per_cube + non_pad_index = (cube_offsets + local_offsets)[ + local_offsets < variable_block_sizes.unsqueeze(1) + ] + + untile_idx = torch.empty(total_seq_length, dtype=torch.long, device=device) + untile_idx[gather_idx] = non_pad_index + + kv_token_mask = torch.zeros(padded_seq_length, dtype=torch.bool, device=device) + kv_token_mask[non_pad_index] = True + + return _VSAShapeMetadata( + num_cubes=num_cubes, + padded_seq_length=padded_seq_length, + variable_block_sizes=variable_block_sizes, + kv_token_mask=kv_token_mask, + non_pad_index=non_pad_index, + gather_idx=gather_idx, + untile_idx=untile_idx, + ) + + def build( + self, + current_timestep: int, + raw_latent_shape: Tuple[int, int, int], + patch_size: Tuple[int, int, int], + vsa_sparsity: float, + device: torch.device, + ) -> VSAMetadata: + dit_seq_shape = ( + raw_latent_shape[0] // patch_size[0], + raw_latent_shape[1] // patch_size[1], + raw_latent_shape[2] // patch_size[2], + ) + cache_key = (dit_seq_shape, device) + shape_metadata = self._cache.get(cache_key) + if shape_metadata is None: + if len(self._cache) >= self._max_cached_shapes: + raise RuntimeError( + "VSA metadata cache reached its " + f"{self._max_cached_shapes}-shape limit; restart the pipeline or " + "reuse a configured resolution/frame profile" + ) + shape_metadata = self._build_metadata(dit_seq_shape, device) + self._cache[cache_key] = shape_metadata + + return VSAMetadata( + current_timestep=current_timestep, + vsa_sparsity=vsa_sparsity, + **shape_metadata, + ) + + def clear(self) -> None: + """Release cached tensors after CUDA Graphs that reference them are cleared.""" + + self._cache.clear() + + +_vsa_forward_context_var: contextvars.ContextVar[Optional[VSAMetadata]] = contextvars.ContextVar( + "_vsa_forward_context", default=None +) + + +@contextmanager +def set_vsa_forward_context(metadata: VSAMetadata) -> Iterator[None]: + """Make VSA metadata visible to attention layers for one model forward.""" + + token = _vsa_forward_context_var.set(metadata) + try: + yield + finally: + _vsa_forward_context_var.reset(token) + + +def get_vsa_forward_context() -> Optional[VSAMetadata]: + """Return the metadata for the active VSA model forward, if any.""" + + return _vsa_forward_context_var.get(None) + + +__all__ = [ + "VSA_TILE_SIZE", + "VSAMetadata", + "VSAMetadataBuilder", + "get_vsa_forward_context", + "set_vsa_forward_context", +] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/predictor.py b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/predictor.py new file mode 100644 index 000000000000..800bb7fdaf80 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/predictor.py @@ -0,0 +1,366 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shared Video Sparse Attention prediction and post-processing.""" + +from dataclasses import dataclass, field +from functools import cache +from math import ceil +from typing import Optional + +import torch + +from .....attention.backends.interface import PredefinedAttentionMask +from .....attention.backends.sparse.params import BlockSparseForwardInputs +from .metadata import ( + _DEFAULT_MAX_CACHED_SHAPES, + VSA_BLOCK_SIZE, + VSAMetadata, + get_vsa_forward_context, +) + +_BITS_PER_WORD = 32 +_SIGNED_INT32_MAX = torch.iinfo(torch.int32).max + + +def _mean_pool_cubes( + x_tiled: torch.Tensor, + variable_block_sizes: torch.LongTensor, + prod_tile: int, + num_cubes: int, +) -> torch.Tensor: + batch_size, _padded, num_heads, head_dim = x_tiled.shape + x_cubes = x_tiled.view(batch_size, num_cubes, prod_tile, num_heads, head_dim) + # FP32 accumulation avoids perturbing the coarse softmax when inputs are BF16. + x_sum = x_cubes.float().sum(dim=2) + valid_counts = variable_block_sizes.float().clamp(min=1).view(1, num_cubes, 1, 1) + return (x_sum / valid_counts).to(x_tiled.dtype) + + +class VSAPreprocessor: + """Convert compact BSHD tensors between sequence-major and tile-major order.""" + + @staticmethod + def tile( + x: torch.Tensor, + non_pad_index: torch.LongTensor, + gather_idx: torch.LongTensor, + padded_seq_len: int, + ) -> torch.Tensor: + # index_select + index_copy_ keeps this path traceable by torch.compile. + batch_size, _seq_len, num_heads, head_dim = x.shape + x_valid = x.index_select(1, gather_idx) + x_padded = x.new_zeros(batch_size, padded_seq_len, num_heads, head_dim) + x_padded.index_copy_(1, non_pad_index, x_valid) + return x_padded + + @staticmethod + def untile( + x: torch.Tensor, + untile_idx: torch.LongTensor, + ) -> torch.Tensor: + return torch.index_select(x, 1, untile_idx) + + +@dataclass(frozen=True, slots=True, kw_only=True, eq=False) +class VSAPostProcessContext: + """Per-call tensors needed after the backend executes the fine stage.""" + + coarse_output: torch.Tensor = field(repr=False) + gate_compress: torch.Tensor = field(repr=False) + gate_fine: Optional[torch.Tensor] = field(default=None, repr=False) + untile_idx: Optional[torch.LongTensor] = field(default=None, repr=False) + output_shape: tuple[int, int, int, int] + + +@dataclass(frozen=True, slots=True, kw_only=True, eq=False) +class VSAForwardInputs: + """Typed VSA prediction consumed by TRTLLM or CuTe DSL fine attention. + + The envelope is structurally immutable. Tensor payloads remain live objects + so CUDA Graph-compatible predictors can publish values into stable buffers. + ``q``, ``k``, ``v``, and ``seq_len`` describe the effective fine-stage + inputs: tiled when block-sparse routes are produced, compact otherwise. + """ + + q: torch.Tensor = field(repr=False) + k: torch.Tensor = field(repr=False) + v: torch.Tensor = field(repr=False) + batch_size: int + seq_len: int + block_sparse_inputs: Optional[BlockSparseForwardInputs] = field(repr=False) + topk_indices: torch.IntTensor = field(repr=False) + variable_block_sizes: torch.LongTensor = field(repr=False) + cur_topk: int + num_cubes: int + post_context: VSAPostProcessContext = field(repr=False) + + +class _VSARouteBuilder: + """Lower fixed-width VSA top-K tables into graph-stable BSR routes.""" + + def __init__(self, max_cached_shapes: int = _DEFAULT_MAX_CACHED_SHAPES) -> None: + if max_cached_shapes <= 0: + raise ValueError("max_cached_shapes must be positive") + self._max_cached_shapes = max_cached_shapes + self._indptr_cache: dict[tuple[torch.device, int, int, int, int], torch.Tensor] = {} + + def from_selected_blocks( + self, + selected_blocks: torch.Tensor, + kv_valid_bits: torch.Tensor, + ) -> BlockSparseForwardInputs: + batch_size, num_kv_heads, num_q_blocks, blocks_per_row = map(int, selected_blocks.shape) + key = ( + selected_blocks.device, + batch_size, + num_kv_heads, + num_q_blocks, + blocks_per_row, + ) + block_indptr = self._indptr_cache.get(key) + if block_indptr is None: + if len(self._indptr_cache) >= self._max_cached_shapes: + raise RuntimeError( + "VSA route cache reached its " + f"{self._max_cached_shapes}-shape limit; restart the pipeline or " + "reuse a configured resolution/frame profile" + ) + if selected_blocks.is_cuda and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "VSA route cache miss during CUDA Graph capture; " + "run an eager warmup with the same selected-block shape first" + ) + total_entries = batch_size * num_kv_heads * num_q_blocks * blocks_per_row + if total_entries > _SIGNED_INT32_MAX: + raise OverflowError("VSA route offsets must fit in signed int32") + row_offsets = torch.arange( + num_q_blocks + 1, + dtype=torch.int32, + device=selected_blocks.device, + ).reshape(1, 1, -1) + head_offsets = torch.arange( + batch_size * num_kv_heads, + dtype=torch.int32, + device=selected_blocks.device, + ).reshape(batch_size, num_kv_heads, 1) + block_indptr = ( + head_offsets * (num_q_blocks * blocks_per_row) + row_offsets * blocks_per_row + ).contiguous() + self._indptr_cache[key] = block_indptr + return BlockSparseForwardInputs( + q_block_size=VSA_BLOCK_SIZE, + kv_block_size=VSA_BLOCK_SIZE, + max_blocks_per_row=blocks_per_row, + block_indptr=block_indptr, + block_indices=torch.sort(selected_blocks, dim=-1).values.reshape(-1).contiguous(), + kv_valid_bits=kv_valid_bits, + ) + + +@cache +def _get_bit_weights(device: torch.device) -> torch.Tensor: + bit_positions = torch.arange(_BITS_PER_WORD, dtype=torch.int64, device=device) + return torch.bitwise_left_shift(torch.ones_like(bit_positions), bit_positions) + + +def _pack_kv_token_mask(kv_token_mask: torch.Tensor, batch_size: int) -> torch.Tensor: + if kv_token_mask.ndim == 1: + batched_mask = kv_token_mask.unsqueeze(0).expand(batch_size, -1) + else: + batched_mask = kv_token_mask + seq_len_kv = int(batched_mask.shape[1]) + padded_length = ceil(seq_len_kv / _BITS_PER_WORD) * _BITS_PER_WORD + if padded_length != seq_len_kv: + batched_mask = torch.nn.functional.pad(batched_mask, (0, padded_length - seq_len_kv)) + words = ( + batched_mask.reshape(batch_size, -1, _BITS_PER_WORD).to(torch.int64) + * _get_bit_weights(kv_token_mask.device) + ).sum(dim=-1) + return words.to(torch.uint32).contiguous() + + +class VSAPredictor: + """Produce the complete per-call VSA block-attention input envelope.""" + + def __init__( + self, + num_heads: int, + num_kv_heads: Optional[int] = None, + max_cached_shapes: int = _DEFAULT_MAX_CACHED_SHAPES, + ) -> None: + resolved_num_kv_heads = num_kv_heads or num_heads + if resolved_num_kv_heads != num_heads: + raise ValueError( + "VSA coarse mean-pool assumes MHA (num_kv_heads == num_heads), " + f"got num_kv_heads={resolved_num_kv_heads}, num_heads={num_heads}. " + "GQA/MQA is not supported." + ) + self._route_builder = _VSARouteBuilder(max_cached_shapes=max_cached_shapes) + + @torch.compiler.disable + def get_metadata(self) -> VSAMetadata: + metadata = get_vsa_forward_context() + if metadata is None: + raise RuntimeError( + "VSA attention called without an active VSA forward context. " + "Wrap each transformer call with set_vsa_forward_context()." + ) + return metadata + + @staticmethod + def _validate_inputs( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gate_compress: Optional[torch.Tensor], + gate_fine: Optional[torch.Tensor], + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + if gate_compress is None: + raise ValueError( + "VSA requires gate_compress. " + "Ensure to_gate_compress is wired in the transformer block." + ) + if q.ndim != 4 or q.shape != k.shape or q.shape != v.shape: + raise ValueError("VSA requires Q, K, and V with the same BSHD shape.") + if any(tensor.device != q.device or tensor.dtype != q.dtype for tensor in (k, v)): + raise ValueError("VSA requires Q, K, and V to share device and dtype.") + if not isinstance(gate_compress, torch.Tensor): + raise TypeError("VSA gate_compress must be a torch.Tensor.") + if ( + gate_compress.shape != q.shape + or gate_compress.device != q.device + or gate_compress.dtype != q.dtype + ): + raise ValueError("VSA gate_compress must share Q's shape, device, and dtype.") + if gate_fine is not None and ( + not isinstance(gate_fine, torch.Tensor) + or gate_fine.shape != q.shape + or gate_fine.device != q.device + or gate_fine.dtype != q.dtype + ): + raise ValueError("VSA gate_fine must share Q's shape, device, and dtype.") + return gate_compress, gate_fine + + def predict( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + batch_size: int, + seq_len: int, + seq_len_kv: int, + attention_mask: PredefinedAttentionMask, + gate_compress: Optional[torch.Tensor], + gate_fine: Optional[torch.Tensor], + use_sparse_fine: bool, + produce_block_sparse_inputs: bool, + metadata: Optional[VSAMetadata] = None, + ) -> VSAForwardInputs: + """Predict routes, effective QKV, and the shared post-process context.""" + + gate_compress, gate_fine = self._validate_inputs(q, k, v, gate_compress, gate_fine) + if attention_mask != PredefinedAttentionMask.FULL: + raise ValueError("VSA supports only full self-attention.") + if seq_len_kv != seq_len: + raise ValueError("VSA requires self-attention with matching Q and KV sequence lengths.") + if tuple(q.shape[:2]) != (batch_size, seq_len): + raise ValueError("VSA batch_size and seq_len must match the compact QKV tensors.") + + metadata = metadata or self.get_metadata() + padded_len = metadata.padded_seq_length + num_cubes = metadata.num_cubes + cur_topk = max(1, ceil((1.0 - metadata.vsa_sparsity) * num_cubes)) + q_tiled = VSAPreprocessor.tile(q, metadata.non_pad_index, metadata.gather_idx, padded_len) + k_tiled = VSAPreprocessor.tile(k, metadata.non_pad_index, metadata.gather_idx, padded_len) + v_tiled = VSAPreprocessor.tile(v, metadata.non_pad_index, metadata.gather_idx, padded_len) + + q_coarse = _mean_pool_cubes( + q_tiled, metadata.variable_block_sizes, VSA_BLOCK_SIZE, num_cubes + ) + k_coarse = _mean_pool_cubes( + k_tiled, metadata.variable_block_sizes, VSA_BLOCK_SIZE, num_cubes + ) + v_coarse = _mean_pool_cubes( + v_tiled, metadata.variable_block_sizes, VSA_BLOCK_SIZE, num_cubes + ) + coarse_scores = torch.einsum("bnhd,bmhd->bhnm", q_coarse, k_coarse) * q.shape[-1] ** -0.5 + coarse_probs = coarse_scores.softmax(dim=-1) + coarse_output = torch.einsum("bhnm,bmhd->bnhd", coarse_probs, v_coarse) + topk_indices = coarse_probs.topk(cur_topk, dim=-1).indices.to(torch.int32) + coarse_output_tiled = ( + coarse_output.unsqueeze(2) + .expand(batch_size, num_cubes, VSA_BLOCK_SIZE, q.shape[2], q.shape[3]) + .reshape(batch_size, padded_len, q.shape[2], q.shape[3]) + ) + coarse_output_compact = VSAPreprocessor.untile(coarse_output_tiled, metadata.untile_idx) + + block_sparse_inputs = None + if use_sparse_fine and produce_block_sparse_inputs: + kv_valid_bits = _pack_kv_token_mask(metadata.kv_token_mask, batch_size) + block_sparse_inputs = self._route_builder.from_selected_blocks( + topk_indices, + kv_valid_bits, + ) + + effective_q = q_tiled if use_sparse_fine else q + effective_k = k_tiled if use_sparse_fine else k + effective_v = v_tiled if use_sparse_fine else v + effective_seq_len = padded_len if use_sparse_fine else seq_len + return VSAForwardInputs( + q=effective_q, + k=effective_k, + v=effective_v, + batch_size=batch_size, + seq_len=effective_seq_len, + block_sparse_inputs=block_sparse_inputs, + topk_indices=topk_indices, + variable_block_sizes=metadata.variable_block_sizes, + cur_topk=cur_topk, + num_cubes=num_cubes, + post_context=VSAPostProcessContext( + coarse_output=coarse_output_compact, + gate_compress=gate_compress, + gate_fine=gate_fine, + untile_idx=metadata.untile_idx if use_sparse_fine else None, + output_shape=tuple(q.shape), + ), + ) + + +def vsa_post_process(output: torch.Tensor, inputs: VSAForwardInputs) -> torch.Tensor: + """Combine coarse/fine VSA outputs and restore compact BSHD order.""" + + context = inputs.post_context + fine_output = output.reshape( + inputs.batch_size, + inputs.seq_len, + context.output_shape[2], + context.output_shape[3], + ) + if context.untile_idx is not None: + fine_output = VSAPreprocessor.untile(fine_output, context.untile_idx) + if context.gate_fine is not None: + fine_output = context.gate_fine * fine_output + return context.gate_compress * context.coarse_output + fine_output + + +__all__ = [ + "VSAForwardInputs", + "VSAPostProcessContext", + "VSAPredictor", + "vsa_post_process", +] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py b/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py index 18976929bef1..dde2e14094c3 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py @@ -27,8 +27,12 @@ from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.visual_gen.args import QuantAttentionConfig -from ...attention.backends.interface import AttentionRuntimeFeatures, PredefinedAttentionMask -from ...attention.backends.sparse.skip_softmax import SkipSoftmaxParams +from ...attention.backends.interface import ( + AttentionForwardArgs, + AttentionRuntimeFeatures, + PredefinedAttentionMask, +) +from ...attention.backends.sparse.params import SparseBackendForwardArgs, SparseParams from ...attention.backends.trtllm import TrtllmAttention as BaseTrtllmAttention from ...attention.backends.trtllm import TrtllmAttentionMetadata as BaseTrtllmAttentionMetadata from .interface import AttentionBackend, AttentionTensorLayout @@ -73,6 +77,12 @@ def __init__( self._cached_seq_lens: Optional[torch.Tensor] = None self._prepared = False + def get_fmha_cache_state(self, name: str) -> dict[str, object]: + """Return one model-scoped cache owned by this metadata adapter.""" + + fmha_caches = self._metadata_state.setdefault("fmha_caches", {}) + return fmha_caches.setdefault(name, {}) + def _needs_prepare(self, batch_size: int, seq_lens: torch.Tensor) -> bool: """Check if we need to call prepare() (current request seq_lens or shared metadata object seq_lens changed). @@ -185,6 +195,7 @@ class TrtllmAttention(BaseTrtllmAttention, AttentionBackend): - Metadata creation and preparation - No KV cache operation - SageAttention per-block QKV quantization (when a quant_attention_config is provided. requires unfused QKV) + - Separate-QKV forwarding for generic block-sparse attention and backends that reject fused QKV """ def __init__( @@ -199,9 +210,17 @@ def __init__( max_seq_len: int = 4096, quant_attention_config: Optional[QuantAttentionConfig] = None, attention_metadata_state: Optional[dict] = None, - sparse_params: Optional[SkipSoftmaxParams] = None, + sparse_params: Optional[SparseParams] = None, ): num_kv_heads = num_kv_heads or num_heads + if attention_metadata_state is None: + raise ValueError( + "TRTLLM attention requires `attention_metadata_state` to be provided " + "by visual-gen config for model-scoped metadata and plan sharing." + ) + self.metadata = TrtllmAttentionMetadata( + attention_metadata_state=attention_metadata_state, + ) super().__init__( layer_idx=layer_idx, @@ -216,12 +235,19 @@ def __init__( # TRTLLM expects flat [B*S, H*D] format self._preferred_layout = AttentionTensorLayout.NHD - self.metadata = TrtllmAttentionMetadata( - attention_metadata_state=attention_metadata_state, - ) - self.quant_attention_config = quant_attention_config + def update_quant_config(self, new_quant_config: Optional[QuantConfig]) -> None: + """Rebuild FMHA libraries and bind VisualGen-owned shared plan caches.""" + + super().update_quant_config(new_quant_config) + from ...attention.backends.fmha.prims_ts_block_sparse import PrimsTSBlockSparseFmha + + cache_state = self.metadata.get_fmha_cache_state("prims_ts_block_sparse") + for fmha in self._fmha_manager.fmha_libs: + if isinstance(fmha, PrimsTSBlockSparseFmha): + fmha.bind_plan_cache(cache_state) + # Needed to work with torch compile cause of attention metadata # make attn metadata as input for it to work @torch.compiler.disable @@ -254,6 +280,7 @@ def forward( seq_len: int, attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL, seq_len_kv: Optional[int] = None, + sparse_backend_args: Optional[SparseBackendForwardArgs] = None, **kwargs, ) -> torch.Tensor: """ @@ -263,10 +290,11 @@ def forward( For diffusion models, expects: - Fused QKV: q contains [Q, K, V] concatenated, k and v are None - - does not support SageAttention + - does not support SageAttention or block-sparse routes - OR separate Q, K, V which: - for regular TRTLLM attention, will be fused internally - - for SageAttention, will be used directly + - for SageAttention, block-sparse routes, and backends that reject + fused QKV, will be passed to the core as separate tensors Args: q: Query tensor [B, S, H, D] or fused QKV [B, S, H_qkv, D] @@ -276,49 +304,71 @@ def forward( seq_len: Sequence length for Q attention_mask: Attention mask type seq_len_kv: Sequence length for K/V (for cross-attention, defaults to seq_len) + sparse_backend_args: Module-predicted sparse inputs handed to the core + prediction hooks. A ``block_sparse_inputs`` payload selects the + generic block-sparse FMHA. + **kwargs: ``timestep`` only; other names are rejected. Returns: Output tensor [B, S, H*D] """ - kv_seq_len = seq_len_kv if seq_len_kv is not None else seq_len - prepared_metadata = self._prepare_metadata(batch_size, seq_len) timestep = kwargs.pop("timestep", None) + if kwargs: + unexpected_names = ", ".join(sorted(kwargs)) + raise TypeError( + f"Unexpected TRTLLM attention forward keyword arguments: {unexpected_names}" + ) - if self.quant_attention_config is not None: - assert k is not None and v is not None, ( - "SageAttention requires separate Q, K, V tensors" + block_sparse_inputs = ( + sparse_backend_args.block_sparse_inputs if sparse_backend_args is not None else None + ) + use_separate_qkv = ( + block_sparse_inputs is not None + or self.quant_attention_config is not None + or not self.support_fused_qkv() + ) + if use_separate_qkv and (k is None or v is None): + raise ValueError("This TRTLLM attention call requires separate q, k, and v tensors.") + if block_sparse_inputs is not None and self.quant_attention_config is not None: + raise ValueError( + "Generic block-sparse attention does not support quant_attention_config." ) - quant_cfg = self.quant_attention_config + + kv_seq_len = seq_len_kv if seq_len_kv is not None else seq_len + prepared_metadata = self._prepare_metadata(batch_size, seq_len) + sage_kwargs = {} + if use_separate_qkv: q = q.reshape(batch_size * seq_len, -1).contiguous() k = k.reshape(batch_size * kv_seq_len, -1).contiguous() v = v.reshape(batch_size * kv_seq_len, -1).contiguous() - output = super().forward( - q=q, - k=k, - v=v, - metadata=prepared_metadata, - attention_mask=attention_mask, - timestep=timestep, - sage_attn_num_elts_per_blk_q=quant_cfg.q_block_size, - sage_attn_num_elts_per_blk_k=quant_cfg.k_block_size, - sage_attn_num_elts_per_blk_v=quant_cfg.v_block_size, - sage_attn_qk_int8=(quant_cfg.qk_dtype == "int8"), - ) + quant_cfg = self.quant_attention_config + if quant_cfg is not None: + sage_kwargs = { + "sage_attn_num_elts_per_blk_q": quant_cfg.q_block_size, + "sage_attn_num_elts_per_blk_k": quant_cfg.k_block_size, + "sage_attn_num_elts_per_blk_v": quant_cfg.v_block_size, + "sage_attn_qk_int8": quant_cfg.qk_dtype == "int8", + } else: if k is None and v is None: - qkv = q.reshape(batch_size * seq_len, -1) + q = q.reshape(batch_size * seq_len, -1) else: - qkv = self._concat_qkv(q, k, v, batch_size, seq_len, kv_seq_len) - output = super().forward( - q=qkv, - k=None, - v=None, - metadata=prepared_metadata, + q = self._concat_qkv(q, k, v, batch_size, seq_len, kv_seq_len) + k = None + v = None + output = super().forward( + q=q, + k=k, + v=v, + metadata=prepared_metadata, + forward_args=AttentionForwardArgs( attention_mask=attention_mask, timestep=timestep, - ) - output = output.view(batch_size, seq_len, -1) - return output + sparse_backend_args=sparse_backend_args, + **sage_kwargs, + ), + ) + return output.view(batch_size, seq_len, -1) @property def preferred_layout(self) -> AttentionTensorLayout: diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py index d6530132b267..49a30a41a038 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py @@ -45,8 +45,7 @@ def get_visual_gen_attention_backend( Backend Selection Guide: - "VANILLA": Full support for cross-attention (different Q/KV seq lengths) Uses torch SDPA backend - - "TRTLLM": Optimized for self-attention (requires same Q/KV seq lengths) - Better performance but requires fused QKV + - "TRTLLM": Optimized for self-attention (requires same Q/KV seq lengths). - "FLASHINFER": Dense prefill attention without a KV cache. Supports FP16/BF16 and architecture-specific NVFP4 attention recipes. - "FA4": Flash Attention 4; provides higher speedup on Blackwell GPUs (sm100) @@ -54,27 +53,33 @@ def get_visual_gen_attention_backend( - "CUTEDSL": CuTe DSL kernels. create_attention selects dense/SkipSoftmax FMHA or VSA from AttentionConfig.sparse_attention_config. """ - # Lazy imports to avoid circular dependency - from .cute_dsl import CuTeDSLAttention - from .flash_attn4 import FlashAttn4Attention - from .flashinfer import FlashInferAttention - from .trtllm import TrtllmAttention - from .vanilla import VanillaAttention backend_name = backend_name.upper() if backend_name == "VANILLA": + from .vanilla import VanillaAttention + return VanillaAttention elif backend_name == "TRTLLM": + from .trtllm import TrtllmAttention + return TrtllmAttention elif backend_name == "FLASHINFER": + from .flashinfer import FlashInferAttention + return FlashInferAttention elif backend_name == "FA4": + from .flash_attn4 import FlashAttn4Attention + return FlashAttn4Attention elif backend_name == "CUTEDSL": + from .cute_dsl import CuTeDSLAttention + return CuTeDSLAttention else: # Default to VANILLA for maximum compatibility + from .vanilla import VanillaAttention + return VanillaAttention @@ -112,42 +117,56 @@ def create_attention( will automatically reallocate if longer sequences are encountered. attention_config: Optional AttentionConfig used to select the attention algorithm and forward its quantization or sparsity configuration. - attention_metadata_state: Optional model-scoped metadata state from - visual-gen config. Required for TRTLLM and shared by FlashInfer layers. + attention_metadata_state: Optional per-component VisualGen attention state. + It keeps shape-stable attention metadata alive across layers and + CUDA Graph captures. Required for TRTLLM and shared by FlashInfer layers. **kwargs: Additional backend-specific arguments Returns: AttentionBackend instance """ - attn_cls = get_visual_gen_attention_backend(backend) + sparse_attention_config = ( + attention_config.sparse_attention_config if attention_config is not None else None + ) + is_vsa = ( + sparse_attention_config is not None + and getattr(sparse_attention_config, "algorithm", None) == "vsa" + ) + + backend_name = backend.upper() + if is_vsa and backend_name == "CUTEDSL": + from .sparse.vsa.backend import VSACuTeDSLAttention + + attn_cls = VSACuTeDSLAttention + elif is_vsa and backend_name == "TRTLLM": + from .sparse.vsa.backend import VSATrtllmAttention + + attn_cls = VSATrtllmAttention + else: + attn_cls = get_visual_gen_attention_backend(backend) + + if is_vsa: + sparse_params = kwargs.pop("sparse_params", None) + if sparse_params is not None: + raise ValueError("VSA does not lower through core SparseParams.") # Forward the validated quantization recipe to TRTLLM, FlashInfer, or the dense CuTe DSL FMHA backend. if attention_config is not None and attention_config.quant_attention_config is not None: kwargs["quant_attention_config"] = attention_config.quant_attention_config - if backend.upper() == "TRTLLM": + if backend_name == "TRTLLM": if attention_metadata_state is None: raise ValueError( "TRTLLM backend requires `attention_metadata_state` from " "DiffusionModelConfig; creation path must not allocate metadata implicitly." ) kwargs["attention_metadata_state"] = attention_metadata_state - elif backend.upper() == "FLASHINFER": + elif backend_name == "FLASHINFER": if attention_metadata_state is None: raise ValueError( "FLASHINFER backend requires `attention_metadata_state` from " "DiffusionModelConfig for shared workspace allocation." ) kwargs["attention_metadata_state"] = attention_metadata_state - if backend.upper() == "CUTEDSL" and attention_config is not None: - if ( - attention_config.sparse_attention_config is not None - and getattr(attention_config.sparse_attention_config, "algorithm", None) == "vsa" - ): - from .cute_dsl.vsa import VSAAttention - - attn_cls = VSAAttention - kwargs["sparse_attention_config"] = attention_config.sparse_attention_config - return attn_cls( layer_idx=layer_idx, num_heads=num_heads, diff --git a/tensorrt_llm/_torch/visual_gen/config.py b/tensorrt_llm/_torch/visual_gen/config.py index f3d45e0d651d..ba19e57be664 100644 --- a/tensorrt_llm/_torch/visual_gen/config.py +++ b/tensorrt_llm/_torch/visual_gen/config.py @@ -74,8 +74,24 @@ def discover_pipeline_components(checkpoint_path: Path) -> Dict[str, Path]: def create_attention_metadata_state() -> Dict[str, Any]: - """Create model-scoped state shared by visual-gen attention layers.""" - return {"metadata_cache": {}} + """Create state shared by attention layers in one model component. + + The state outlives individual forwards and CUDA Graph captures and owns the + shape-keyed TRTLLM metadata and PrimTS plan caches. VisualGen attention + layers execute serially within one component, so sharing graph-stable route + workspaces avoids retaining one worst-case allocation per layer. Each model + component receives a distinct state and must not execute concurrent forwards. + """ + return { + "metadata_cache": {}, + "fmha_caches": { + "prims_ts_block_sparse": { + "contiguous_wrappers": {}, + "paged_wrappers": {}, + }, + }, + "sparse_predictors": {}, + } def _model_config_value(value: Any, *, deep_copy: bool = True) -> Any: @@ -123,6 +139,7 @@ class DiffusionModelConfig(_VisualGenConfigBase): cuda_graph: CudaGraphConfig = PydanticField(default_factory=CudaGraphConfig) cpu_offload_config: CpuOffloadConfig = PydanticField(default_factory=CpuOffloadConfig) attention: AttentionConfig = PydanticField(default_factory=AttentionConfig) + # Per-component metadata cache shared by VisualGen TRTLLM attention layers. attention_metadata_state: Optional[Dict[str, Any]] = None parallel: ParallelConfig = PydanticField(default_factory=ParallelConfig) cache: Optional[CacheConfig] = None @@ -196,6 +213,7 @@ class DiffusionPipelineConfig(_VisualGenConfigBase): cuda_graph: CudaGraphConfig = PydanticField(default_factory=CudaGraphConfig) cpu_offload_config: CpuOffloadConfig = PydanticField(default_factory=CpuOffloadConfig) attention: AttentionConfig = PydanticField(default_factory=AttentionConfig) + # Seed state copied into each model component before attention metadata is created. attention_metadata_state: Optional[Dict[str, Any]] = None parallel: ParallelConfig = PydanticField(default_factory=ParallelConfig) cache: Optional[CacheConfig] = None diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/interface.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/interface.py index 4bee4dad352a..b0438b3969bd 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/interface.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/interface.py @@ -17,7 +17,7 @@ Blackwell (sm_100) fast path for VSA's fine stage. The kernel JIT-compiles on first call and is cached per process; the caller -(CuTeDSLAttention._forward_vsa) falls back to dense SDPA when the +(VSACuTeDSLAttention) falls back to dense SDPA when the device/dtype/head_dim envelope is not met. """ diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py index ee730932b5cb..2400a3d6728b 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py @@ -26,7 +26,7 @@ from diffusers.video_processor import VideoProcessor from transformers import AutoTokenizer, UMT5EncoderModel -from tensorrt_llm._torch.visual_gen.attention_backend import ( +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import ( VSAMetadataBuilder, set_vsa_forward_context, ) @@ -115,6 +115,7 @@ hf_ids=[ "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "Wan-AI/Wan2.1-T2V-14B-Diffusers", + "FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers", "Wan-AI/Wan2.2-T2V-A14B-Diffusers", "Wan-AI/Wan2.2-TI2V-5B-Diffusers", "nvidia/Wan2.2-T2V-A14B-Diffusers-FP8", @@ -148,6 +149,16 @@ def __init__(self, pipeline_config): ) super().__init__(pipeline_config) + # CUDA graphs capture the VSA partition/mask tensor addresses. Keep + # their shape cache alive across requests so replay never references + # metadata owned by a completed forward call. + self._vsa_metadata_builder = VSAMetadataBuilder() + + def cleanup(self): + """Release CUDA graphs before clearing captured attention state.""" + + super().cleanup() + self._vsa_metadata_builder.clear() def _compute_wan_timestep_embedding(self, module, timestep=None, **kwargs): """Compute timestep embedding for WAN transformer. @@ -608,15 +619,15 @@ def forward( f"guidance_scale={guidance_scale}, guidance_scale_2={guidance_scale_2}" ) - # VSA: build metadata builder once per forward() call; reused across timesteps. + # VSA metadata is cached at pipeline scope and reused across requests. _attn_cfg = self.pipeline_config.primary_model_config.attention _sparse_cfg = getattr(_attn_cfg, "sparse_attention_config", None) _vsa_active = ( - getattr(_attn_cfg, "backend", "VANILLA") == "CUTEDSL" + getattr(_attn_cfg, "backend", "VANILLA") in ("CUTEDSL", "TRTLLM") and _sparse_cfg is not None and getattr(_sparse_cfg, "algorithm", None) == "vsa" ) - _vsa_builder = VSAMetadataBuilder() if _vsa_active else None + _vsa_builder = self._vsa_metadata_builder if _vsa_active else None _vsa_patch_size = tuple(getattr(self.config, "patch_size", [1, 2, 2])) # (pT, pH, pW) _vsa_sparsity = _sparse_cfg.vsa_sparsity if _vsa_active else 0.0 diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py b/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py index b717ff98abb5..4f3359c8c468 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py @@ -351,6 +351,7 @@ def __init__( config=model_config, layer_idx=_layer_idx, async_ulysses=self._use_async_ulysses, + separate_qkv_is_self_attention=True, module_name=f"blocks.{_layer_idx}.attn1", ) @@ -406,7 +407,7 @@ def __init__( reduce_output=(tp_size != 1), ) - # VSA gates (CUTEDSL backend, sparse_attention_config.algorithm == "vsa"). + # VSA gates are shared by the backend-specific fine-stage implementations. # G_c weights the coarse branch; G_f weights the fine branch. self.to_gate_compress = None self.to_gate_fine = None @@ -414,7 +415,6 @@ def __init__( _sa_cfg = getattr(_attn_cfg, "sparse_attention_config", None) if _attn_cfg else None _is_vsa = ( _attn_cfg is not None - and getattr(_attn_cfg, "backend", "VANILLA") == "CUTEDSL" and _sa_cfg is not None and getattr(_sa_cfg, "algorithm", None) == "vsa" ) @@ -432,6 +432,10 @@ def __init__( force_dynamic_quantization=force_dynamic_quant, tensor_parallel_mode=gate_tp_mode, reduce_output=False, + override_tp_sharding=( + self.attn1.local_q_dim_start, + self.attn1.local_q_dim_end, + ), ) self.to_gate_fine = Linear( hidden_size, @@ -444,6 +448,10 @@ def __init__( force_dynamic_quantization=force_dynamic_quant, tensor_parallel_mode=gate_tp_mode, reduce_output=False, + override_tp_sharding=( + self.attn1.local_q_dim_start, + self.attn1.local_q_dim_end, + ), ) # I2V: Additional K/V projections for image embeddings. @@ -588,7 +596,12 @@ def forward( # so each V/Q/K GEMM + norm + RoPE overlaps with the peer push on the # side stream; both paths return 3D [B, S, H*D]. if self._use_async_ulysses: - attn1_out = self.attn1.forward_async(normed, freqs=freqs, timestep=timestep) + attn1_out = self.attn1.forward_async( + normed, + freqs=freqs, + timestep=timestep, + **attn1_kwargs, + ) else: attn1_out = self.attn1(normed, freqs=freqs, timestep=timestep, **attn1_kwargs) diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/modules/attention.py index 8c19484229c7..8ffd4790728c 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py @@ -1,3 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from enum import Enum from typing import Optional, Tuple @@ -97,16 +112,19 @@ def __init__( cp_size = vgm.cp_size if vgm else 1 base_backend = config.attention.backend _sa_cfg = config.attention.sparse_attention_config - _is_vsa = ( - base_backend == "CUTEDSL" - and _sa_cfg is not None - and getattr(_sa_cfg, "algorithm", None) == "vsa" + is_vsa = _sa_cfg is not None and getattr(_sa_cfg, "algorithm", None) == "vsa" + + is_separate_qkv = self.qkv_mode == QKVMode.SEPARATE_QKV + is_separate_qkv_cross_attention = is_separate_qkv and not separate_qkv_is_self_attention + use_vanilla_cross_attention = is_separate_qkv and ( + (base_backend == "TRTLLM" and not is_vsa) + or (is_vsa and not separate_qkv_is_self_attention) ) - # Cross-attention fallback: TRTLLM and CUTEDSL VSA are self-attn only. - if self.qkv_mode == QKVMode.SEPARATE_QKV and (base_backend == "TRTLLM" or _is_vsa): + # Cross-attention fallback: dense TRTLLM and every VSA backend are self-attn only. + if use_vanilla_cross_attention: backend_name = "VANILLA" - requested = f"{base_backend} (VSA)" if _is_vsa else base_backend + requested = f"{base_backend} (VSA)" if is_vsa else base_backend # Warn once per (module class, requested, resolved) triple so the # fallback is visible without per-module-instance log spam. logger.warning_once( @@ -117,7 +135,7 @@ def __init__( else: backend_name = base_backend - if _is_vsa and cp_size > 1: + if is_vsa and cp_size > 1: raise ValueError( f"VSA needs the full token sequence per rank, so it is incompatible " f"with context parallelism (Attention2D/Ring, cp_size={cp_size}). Use " @@ -248,12 +266,7 @@ def __init__( sparse_params=sparse_params, ) - if ( - enable_sequence_parallel - and self.qkv_mode == QKVMode.SEPARATE_QKV - and not separate_qkv_is_self_attention - and vgm is not None - ): + if enable_sequence_parallel and is_separate_qkv_cross_attention and vgm is not None: ring_size = vgm.ring_size if ring_size > 1: raise ValueError( @@ -645,6 +658,7 @@ def forward_async( hidden_states: torch.Tensor, freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, timestep: Optional[torch.Tensor] = None, + **kwargs, ) -> torch.Tensor: """Async-Ulysses self-attn driver. Structurally mirrors ``forward``: each closure does ``to_{q,k,v}`` + (optional) fused norm+RoPE on the @@ -684,8 +698,8 @@ def forward_async( ) B, S = hidden_states.shape[:2] - H = self.num_attention_heads - KV = self.num_key_value_heads + H = self.local_num_attention_heads + KV = self.local_num_key_value_heads D = self.head_dim # Mirrors forward()'s fused gate. qkv_mode is implicitly SEPARATE_QKV # under async (caller-enforced), so the FUSE_QKV check in forward() @@ -741,6 +755,17 @@ def compute_k(): def compute_v(): return self.to_v(qkv_input).view(B, S, KV, D) - out_4d = self.attn.forward_async(compute_q, compute_k, compute_v, timestep=timestep) + for gate_key in ("gate_compress", "gate_fine"): + gate = kwargs.get(gate_key) + if gate is not None: + kwargs[gate_key] = gate.view(B, S, self.local_num_attention_heads, D) + + out_4d = self.attn.forward_async( + compute_q, + compute_k, + compute_v, + timestep=timestep, + **kwargs, + ) b, t = out_4d.shape[:2] return self.to_out[0](out_4d.reshape(b, t, H * D)) diff --git a/tensorrt_llm/_torch/visual_gen/pipeline_loader.py b/tensorrt_llm/_torch/visual_gen/pipeline_loader.py index b1d95d284105..d63d7d087a7b 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline_loader.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline_loader.py @@ -1,3 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """ Model loader for diffusion pipelines. @@ -228,13 +243,22 @@ def load( _attn_backend = config.attention.backend _sa_cfg = config.attention.sparse_attention_config if ( - _attn_backend == "CUTEDSL" + _attn_backend in ("CUTEDSL", "TRTLLM") and _sa_cfg is not None and getattr(_sa_cfg, "algorithm", None) == "vsa" ): - kernel_path = "CuTe DSL block-sparse" if CUTE_AVAILABLE else "dense SDPA fallback" + if _attn_backend == "CUTEDSL": + kernel_path = ( + "CuTe DSL block-sparse when supported; dense SDPA fallback otherwise" + if CUTE_AVAILABLE + else "dense SDPA fallback" + ) + else: + kernel_path = ( + "PrimTS block-sparse when supported; compact dense TRTLLM fallback otherwise" + ) logger.info( - f"Attention backend: CUTEDSL (algorithm=vsa, " + f"Attention backend: {_attn_backend} (algorithm=vsa, " f"sparsity={_sa_cfg.vsa_sparsity}, fine-stage={kernel_path})" ) else: diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index 6b2ee30ccdca..b3390317af52 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -122,7 +122,7 @@ class AttentionConfig(StrictBaseModel): status="prototype", description=( "Sparse attention recipe. Discriminated by algorithm: " - "skip_softmax (TRTLLM / CUTEDSL backends) or VSA (CUTEDSL backend)." + "skip_softmax (TRTLLM / CUTEDSL backends) or VSA (CUTEDSL / TRTLLM backends)." ), ) @@ -205,7 +205,7 @@ def _validate_sparse_attention_config(self) -> "AttentionConfig": algo = self.sparse_attention_config.algorithm supported_backends = { "skip_softmax": ("TRTLLM", "CUTEDSL"), - "vsa": ("CUTEDSL",), + "vsa": ("CUTEDSL", "TRTLLM"), }.get(algo) if supported_backends is None: return self @@ -220,21 +220,14 @@ def _validate_sparse_attention_config(self) -> "AttentionConfig": return self @model_validator(mode="after") - def _validate_cutedsl_quant_sparse_mutex(self) -> "AttentionConfig": - # VSA replaces the dense CuTeDSL path and cannot compose with quantized - # attention. SkipSoftmax is part of that dense path and can compose. - if ( - self.backend == "CUTEDSL" - and self.quant_attention_config is not None - and self.sparse_attention_config is not None - and self.sparse_attention_config.algorithm == "vsa" - ): - raise ValueError( - "CUTEDSL backend: quant_attention_config and VSA " - "sparse_attention_config are mutually exclusive (the " - "CuTeDSLAttention dispatcher selects either the dense path " - "or the sparse VSA path, not both)." - ) + def _validate_quant_sparse_mutex(self) -> "AttentionConfig": + if self.quant_attention_config is None or self.sparse_attention_config is None: + return self + + if self.sparse_attention_config.algorithm == "vsa": + # VSA consumes the unquantized Q/K/V path, so accepting an attention + # quantization recipe would silently ignore user configuration. + raise ValueError("VSA and quant_attention_config are mutually exclusive.") return self diff --git a/tensorrt_llm/visual_gen/sparse_attention.py b/tensorrt_llm/visual_gen/sparse_attention.py index 2bccb10e2652..21189c249d55 100644 --- a/tensorrt_llm/visual_gen/sparse_attention.py +++ b/tensorrt_llm/visual_gen/sparse_attention.py @@ -222,11 +222,13 @@ def _ckpt_sparse_attention_config_from_kwargs( class VideoSparseAttentionConfig(StrictBaseModel): - """Video Sparse Attention (VSA) sparse-attention recipe (CUTEDSL backend only). + """Video Sparse Attention (VSA) sparse-attention recipe. Two-stage hybrid attention: a coarse mean-pooled stage over (4,4,4) cubes and a block-sparse fine stage over the top-K cubes selected per head. vsa_sparsity controls the fraction of cubes dropped on the fine stage. + The fine stage may run on either the CuTeDSL backend or the TRTLLM PrimTS + backend, while the user-facing sparsity semantics stay the same. """ algorithm: Literal["vsa"] = PydanticField( diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 7e10e8baaf9f..6204cc3e512f 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -280,10 +280,11 @@ l0_b200: - unittest/_torch/visual_gen/test_pertoken_adaln.py - unittest/_torch/visual_gen/test_attention_cute_dsl.py - unittest/_torch/visual_gen/test_fa4_cutlass_compatibility.py - - unittest/_torch/visual_gen/test_attention_cute_dsl_vsa.py + - unittest/_torch/visual_gen/test_attention_vsa.py - unittest/_torch/visual_gen/test_attention_flashinfer.py - unittest/_torch/visual_gen/test_attention_trtllm_sage.py - unittest/_torch/visual_gen/test_attention_integration.py + - unittest/_torch/visual_gen/test_trtllm_attention_metadata.py - unittest/_torch/visual_gen/test_attention_fa4.py - unittest/_torch/visual_gen/test_attention_perf.py - unittest/_torch/visual_gen/test_qwen_image_layered_registry.py diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_attention.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_attention.py index 2b900c603ce6..48f46ddee190 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_attention.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_attention.py @@ -1,3 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Multi-GPU tests for Ulysses Attention. These tests use torch.multiprocessing.spawn to launch multiple processes internally. @@ -105,6 +120,56 @@ def run_test_in_distributed(world_size: int, test_fn: Callable, use_cuda: bool = ) +def test_forward_async_redistributes_vsa_gates(monkeypatch): + if not MODULES_AVAILABLE: + pytest.skip("Required modules not available") + + import tensorrt_llm._torch.visual_gen.attention_backend.parallel as parallel_backend + + class _CaptureBackend: + preferred_layout = AttentionTensorLayout.NHD + + def forward(self, q, k, v, **kwargs): + self.kwargs = kwargs + return q + + inner_backend = _CaptureBackend() + attention = object.__new__(UlyssesAttention) + attention.world_size = 1 + attention.process_group = None + attention.inner_backend = inner_backend + attention._issue_async = lambda tensor: tensor.unsqueeze(0) + attention._join_async = lambda: None + attention._output_a2a = lambda output, batch_size, seq_len: output + redistributed = [] + + def _fake_all_to_all(tensor, **kwargs): + redistributed.append((tensor, kwargs)) + return tensor + 1 + + monkeypatch.setattr(parallel_backend, "all_to_all_4d", _fake_all_to_all) + q = torch.randn(1, 3, 2, 4) + gate_compress = torch.randn_like(q) + gate_fine = torch.randn_like(q) + + output = attention.forward_async( + lambda: q, + lambda: q, + lambda: q, + gate_compress=gate_compress, + gate_fine=gate_fine, + ) + + assert output.shape == q.shape + assert redistributed[0][0] is gate_compress + assert redistributed[1][0] is gate_fine + assert all(entry[1]["scatter_dim"] == 2 for entry in redistributed) + assert all(entry[1]["gather_dim"] == 1 for entry in redistributed) + assert inner_backend.kwargs["batch_size"] == q.shape[0] + torch.testing.assert_close(inner_backend.kwargs["gate_compress"], gate_compress + 1) + torch.testing.assert_close(inner_backend.kwargs["gate_fine"], gate_fine + 1) + + # ============================================================================= # Test logic functions (module-level so they can be pickled by mp.spawn) # ============================================================================= diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_async_ulysses.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_async_ulysses.py index 042525f82e30..64580eab2f00 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_async_ulysses.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_async_ulysses.py @@ -324,5 +324,56 @@ def test_async_vs_sync_parity(self, backend): run_test_in_distributed(2, _logic_async_vs_sync_parity, backend) +def test_forward_async_uses_tp_local_heads_for_qkv_gates_and_output(): + from tensorrt_llm._torch.visual_gen.modules.attention import Attention + + class _CaptureAsyncAttention(torch.nn.Module): + def forward_async(self, compute_q, compute_k, compute_v, **kwargs): + self.q = compute_q() + self.k = compute_k() + self.v = compute_v() + self.kwargs = kwargs + return self.q + + class _CaptureOutputProjection(torch.nn.Module): + def forward(self, hidden_states): + self.input = hidden_states + return hidden_states + + attention = Attention.__new__(Attention) + torch.nn.Module.__init__(attention) + attention.num_attention_heads = 4 + attention.num_key_value_heads = 4 + attention.local_num_attention_heads = 2 + attention.local_num_key_value_heads = 2 + attention.head_dim = 4 + attention.fuse_qk_norm_rope = False + attention.qk_norm = False + attention._maybe_share_qkv_quantize = False + attention.to_q = torch.nn.Linear(12, 8, bias=False) + attention.to_k = torch.nn.Linear(12, 8, bias=False) + attention.to_v = torch.nn.Linear(12, 8, bias=False) + attention.attn = _CaptureAsyncAttention() + output_projection = _CaptureOutputProjection() + attention.to_out = torch.nn.ModuleList([output_projection]) + gate_compress = torch.randn(1, 3, 8) + gate_fine = torch.randn_like(gate_compress) + + output = attention.forward_async( + torch.randn(1, 3, 12), + gate_compress=gate_compress, + gate_fine=gate_fine, + ) + + expected_shape = (1, 3, 2, 4) + assert attention.attn.q.shape == expected_shape + assert attention.attn.k.shape == expected_shape + assert attention.attn.v.shape == expected_shape + assert attention.attn.kwargs["gate_compress"].shape == expected_shape + assert attention.attn.kwargs["gate_fine"].shape == expected_shape + assert output_projection.input.shape == (1, 3, 8) + assert output.shape == (1, 3, 8) + + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_tp.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_tp.py index fbedc1aa1410..032508230298 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_tp.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_tp.py @@ -529,6 +529,53 @@ def _logic_wan_i2v_tp_vs_single_gpu_with_config(rank, world_size, config_dict): # ============================================================================= +@pytest.mark.parametrize( + ("tp_rank", "expected_head_range"), + [(0, (0, 8)), (1, (8, 12))], +) +def test_wan_vsa_gates_follow_ulysses_aligned_tp_q_shard(monkeypatch, tp_rank, expected_head_range): + """VSA gates must select the same TP-local heads as the Q projection.""" + from tensorrt_llm._torch.visual_gen.models.wan.transformer_wan import WanBlock + from tensorrt_llm._torch.visual_gen.modules import attention as attention_module + from tensorrt_llm.mapping import Mapping + from tensorrt_llm.visual_gen.args import VideoSparseAttentionConfig + + monkeypatch.setattr( + attention_module, + "wrap_parallel_attention", + lambda attention, **_kwargs: attention, + ) + mapping = Mapping(world_size=2, rank=tp_rank, tp_size=2) + monkeypatch.setattr(type(mapping), "tp_rank", property(lambda self: self.rank)) + head_dim = 128 + model_config = DiffusionModelConfig( + pretrained_config=SimpleNamespace( + hidden_size=12 * head_dim, + num_attention_heads=12, + attention_head_dim=head_dim, + ffn_dim=512, + eps=1e-6, + cross_attn_norm=True, + ), + mapping=mapping, + visual_gen_mapping=SimpleNamespace(ulysses_size=4, cp_size=1), + attention=AttentionConfig( + backend="CUTEDSL", + sparse_attention_config=VideoSparseAttentionConfig(vsa_sparsity=0.5), + ), + skip_create_weights_in_init=True, + ) + + block = WanBlock(model_config, _layer_idx=0) + + expected_shard = tuple(index * head_dim for index in expected_head_range) + assert (block.attn1.local_q_dim_start, block.attn1.local_q_dim_end) == expected_shard + assert block.to_gate_compress.tp_sharding == expected_shard + assert block.to_gate_fine.tp_sharding == expected_shard + assert block.to_gate_compress.out_features == block.attn1.local_q_dim + assert block.to_gate_fine.out_features == block.attn1.local_q_dim + + class TestWanT2VTP: """Tensor parallelism tests for WAN T2V transformer.""" diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_vsa_ulysses.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_vsa_ulysses.py index 00dd752d902e..a60d12a01a30 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_vsa_ulysses.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_vsa_ulysses.py @@ -36,9 +36,9 @@ import torch.multiprocessing as mp try: - from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl import ( + from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl import _cute_dsl_import_error + from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import ( VSAMetadataBuilder, - _cute_dsl_import_error, set_vsa_forward_context, ) from tensorrt_llm._torch.visual_gen.config import ( diff --git a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_vsa.py b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_vsa.py deleted file mode 100644 index 63684cd0519a..000000000000 --- a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_vsa.py +++ /dev/null @@ -1,478 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""VSA correctness tests: CuTe kernel, tile/untile roundtrip, top-k math, backend guards. - -Module-level dense-equivalence and finite-output checks live in -test_attention_integration.py. -""" - -from types import SimpleNamespace - -import pytest -import torch -import torch.nn.functional as F - -from tensorrt_llm._torch.visual_gen.attention_backend import ( - CuTeDSLAttention, - VSAAttention, - VSAMetadataBuilder, -) -from tensorrt_llm._torch.visual_gen.attention_backend.utils import create_attention -from tensorrt_llm._torch.visual_gen.config import ( - DiffusionModelConfig, - create_attention_metadata_state, -) -from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode -from tensorrt_llm.visual_gen.args import ( - AttentionConfig, - QuantAttentionConfig, - VideoSparseAttentionConfig, -) - - -def test_cute_dsl_factory_dispatches_quantized_fmha_and_vsa() -> None: - quant_config = QuantAttentionConfig(qk_dtype="mxfp8", v_dtype="fp8", v_block_size=1) - dense_config = AttentionConfig(backend="CUTEDSL", quant_attention_config=quant_config) - dense_attention = create_attention( - backend="CUTEDSL", - layer_idx=0, - num_heads=8, - head_dim=128, - attention_config=dense_config, - ) - - sparse_config = VideoSparseAttentionConfig(vsa_sparsity=0.9) - vsa_config = AttentionConfig(backend="CUTEDSL", sparse_attention_config=sparse_config) - vsa_attention = create_attention( - backend="CUTEDSL", - layer_idx=0, - num_heads=8, - head_dim=128, - attention_config=vsa_config, - ) - - assert isinstance(dense_attention, CuTeDSLAttention) - assert dense_attention.quant_attention_config is quant_config - assert isinstance(vsa_attention, VSAAttention) - assert vsa_attention.sparse_attention_config is sparse_config - - -def _make_config( - hidden_size: int, - num_heads: int, - head_dim: int, - backend: str, - vsa_sparsity: "float | None" = None, -) -> DiffusionModelConfig: - """Minimal DiffusionModelConfig for one Attention module.""" - pretrained_config = SimpleNamespace( - hidden_size=hidden_size, - num_attention_heads=num_heads, - attention_head_dim=head_dim, - eps=1e-6, - ) - sparse_attention_config = ( - VideoSparseAttentionConfig(vsa_sparsity=vsa_sparsity) if vsa_sparsity is not None else None - ) - config = DiffusionModelConfig( - pretrained_config=pretrained_config, - attention=AttentionConfig(backend=backend, sparse_attention_config=sparse_attention_config), - skip_create_weights_in_init=False, - ) - config.attention_metadata_state = ( - create_attention_metadata_state() if backend == "TRTLLM" else None - ) - return config - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="VSA needs CUDA") -def test_vsa_falls_back_to_vanilla_for_cross_attention(): - """Cross-attention (SEPARATE_QKV) falls back to VANILLA — it has no cube structure.""" - device = torch.device("cuda") - dtype = torch.bfloat16 - cfg = _make_config( - hidden_size=64, num_heads=4, head_dim=16, backend="CUTEDSL", vsa_sparsity=0.5 - ) - cross_attn = ( - Attention(64, 4, qkv_mode=QKVMode.SEPARATE_QKV, config=cfg) - .to(device=device, dtype=dtype) - .eval() - ) - assert cross_attn.attn_backend == "VANILLA", ( - f"VSA on cross-attention should fall back to VANILLA, got {cross_attn.attn_backend!r}" - ) - - -def test_vsa_with_attn2d_raises(): - """VSA + Attention2D must error at construction (VSA needs the full sequence per rank).""" - pretrained_config = SimpleNamespace( - hidden_size=64, - num_attention_heads=4, - attention_head_dim=16, - eps=1e-6, - ) - cfg = DiffusionModelConfig( - pretrained_config=pretrained_config, - attention=AttentionConfig( - backend="CUTEDSL", - sparse_attention_config=VideoSparseAttentionConfig(vsa_sparsity=0.0), - ), - skip_create_weights_in_init=False, - ) - cfg.visual_gen_mapping = SimpleNamespace( - ring_size=1, - ring_group=None, - ulysses_size=1, - ulysses_group=None, - attn2d_row_size=2, - attn2d_col_size=2, - attn2d_row_group=None, - attn2d_col_group=None, - cp_size=4, - ) - with pytest.raises(ValueError, match="incompatible with context parallelism"): - Attention(64, 4, qkv_mode=QKVMode.FUSE_QKV, config=cfg) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="VSA needs CUDA") -def test_vsa_topk_collapses_to_dense_at_sparsity_zero(): - """At sparsity=0, top_k equals num_cubes (dense connectivity).""" - from math import ceil - - device = torch.device("cuda") - builder = VSAMetadataBuilder() - metadata = builder.build( - current_timestep=0, - raw_latent_shape=(8, 8, 8), - patch_size=(1, 1, 1), - vsa_sparsity=0.0, - device=device, - ) - num_cubes = metadata.num_tiles[0] * metadata.num_tiles[1] * metadata.num_tiles[2] - cur_topk = max(1, ceil((1.0 - metadata.vsa_sparsity) * num_cubes)) - assert cur_topk == num_cubes, ( - f"sparsity=0 should select all {num_cubes} cubes, got top_k={cur_topk}" - ) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="VSA needs CUDA") -@pytest.mark.parametrize( - "latent_shape", - [ - (8, 8, 8), - (9, 9, 9), - (21, 45, 80), - ], - ids=["clean_8x8x8", "ragged_9x9x9", "wan720p_21x45x80"], -) -def test_vsa_tile_untile_roundtrip(latent_shape): - """VSAPreprocessor.tile then .untile must losslessly reproduce the input.""" - from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.vsa import VSAPreprocessor - - device = torch.device("cuda") - dtype = torch.bfloat16 - torch.manual_seed(0) - - B, H, D = 2, 4, 32 - seq_len = latent_shape[0] * latent_shape[1] * latent_shape[2] - - builder = VSAMetadataBuilder() - meta = builder.build( - current_timestep=0, - raw_latent_shape=latent_shape, - patch_size=(1, 1, 1), - vsa_sparsity=0.0, - device=device, - ) - - x = torch.randn(B, seq_len, H, D, device=device, dtype=dtype) - - x_tiled = VSAPreprocessor.tile( - x, - meta.non_pad_index, - meta.gather_idx, - meta.padded_seq_length, - ) - - pad_mask = torch.ones(meta.padded_seq_length, dtype=torch.bool, device=device) - pad_mask[meta.non_pad_index] = False - if pad_mask.any(): - assert x_tiled[:, pad_mask, :, :].abs().max().item() == 0.0, ( - "tile() must zero-fill padded positions" - ) - - x_roundtrip = VSAPreprocessor.untile( - x_tiled, - meta.reverse_tile_partition_indices, - meta.non_pad_index, - ) - - assert x_roundtrip.shape == x.shape, ( - f"shape mismatch after tile/untile: {x_roundtrip.shape} vs {x.shape}" - ) - assert torch.equal(x_roundtrip, x), ( - f"tile/untile round-trip is not lossless for latent_shape={latent_shape}: " - f"max_diff={(x_roundtrip - x).abs().max().item():.3e}" - ) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="kernel test needs CUDA") -def test_cute_kernel_matches_dense_at_full_topk(): - """CuTe block-sparse kernel matches dense SDPA when every cube is selected.""" - from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.video_sparse_attention import ( - CUTE_AVAILABLE, - block_sparse_attn_from_indices_cute, - is_cute_supported, - ) - - if not CUTE_AVAILABLE: - pytest.skip("cuda-bindings or cutlass-dsl not importable") - - device = torch.device("cuda") - dtype = torch.bfloat16 - torch.manual_seed(0) - - B, H, num_cubes, D = 1, 4, 4, 128 - block_size = 64 - seq_len = num_cubes * block_size - - q = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) - k = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) - v = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) - - if not is_cute_supported(q): - pytest.skip("CuTe path needs sm_100+ Blackwell (current device unsupported)") - - topk = num_cubes - q2k_idx = ( - torch.arange(num_cubes, device=device, dtype=torch.int32) - .view(1, 1, 1, num_cubes) - .expand(B, H, num_cubes, topk) - .contiguous() - ) - q2k_num = torch.full((B, H, num_cubes), topk, dtype=torch.int32, device=device) - variable_block_sizes = torch.full((num_cubes,), block_size, dtype=torch.int32, device=device) - - out_kernel, _lse = block_sparse_attn_from_indices_cute( - q, k, v, q2k_idx, q2k_num, variable_block_sizes - ) - out_ref = F.scaled_dot_product_attention(q, k, v) - - max_diff = (out_kernel - out_ref).abs().max().item() - mean_diff = (out_kernel - out_ref).abs().mean().item() - - rtol, atol = 1e-2, 1e-2 - assert torch.allclose(out_kernel, out_ref, rtol=rtol, atol=atol), ( - f"CuTe block-sparse kernel deviates from dense SDPA at full top-K: " - f"max_diff={max_diff:.3e}, mean_diff={mean_diff:.3e} (rtol={rtol}, atol={atol})" - ) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="kernel test needs CUDA") -def test_cute_kernel_matches_ref_with_independent_indices(): - """CuTe kernel: paired Q-blocks (2i, 2i+1) attend to independent KV index lists.""" - from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.video_sparse_attention import ( - CUTE_AVAILABLE, - block_sparse_attn_from_indices_cute, - is_cute_supported, - ) - - if not CUTE_AVAILABLE: - pytest.skip("cuda-bindings or cutlass-dsl not importable") - - device = torch.device("cuda") - dtype = torch.bfloat16 - torch.manual_seed(42) - - B, H, num_cubes, D = 2, 4, 16, 128 - block_size = 64 - topk = num_cubes // 2 - seq_len = num_cubes * block_size - - q = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) - k = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) - v = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) - - if not is_cute_supported(q): - pytest.skip("CuTe path needs sm_100+ Blackwell (current device unsupported)") - - q2k_idx = ( - torch.stack( - [ - torch.randperm(num_cubes, device=device, dtype=torch.int32)[:topk] - for _ in range(B * H * num_cubes) - ] - ) - .view(B, H, num_cubes, topk) - .contiguous() - ) - - paired = q2k_idx.view(B, H, num_cubes // 2, 2, topk).sort(dim=-1).values - pair_mismatch = (paired[..., 0, :] != paired[..., 1, :]).sum().item() - assert pair_mismatch > 0, ( - "Pre-condition failed: random permutations matched across every pair; " - "re-seed or raise num_cubes." - ) - - q2k_num = torch.full((B, H, num_cubes), topk, dtype=torch.int32, device=device) - variable_block_sizes = torch.full((num_cubes,), block_size, dtype=torch.int32, device=device) - - attn_mask = torch.full( - (B, H, seq_len, seq_len), float("-inf"), device=device, dtype=torch.float32 - ) - for b in range(B): - for h in range(H): - for q_blk in range(num_cubes): - for ki in range(topk): - k_blk = q2k_idx[b, h, q_blk, ki].item() - qs = q_blk * block_size - ks = k_blk * block_size - attn_mask[b, h, qs : qs + block_size, ks : ks + block_size] = 0.0 - - out_kernel, _lse = block_sparse_attn_from_indices_cute( - q, k, v, q2k_idx, q2k_num, variable_block_sizes - ) - - scale = 1.0 / (D**0.5) - scores = (q.float() @ k.float().transpose(-2, -1)) * scale - scores = scores + attn_mask - probs = torch.softmax(scores, dim=-1) - out_ref = (probs @ v.float()).to(dtype) - - abs_diff = (out_kernel.float() - out_ref.float()).abs() - max_diff = abs_diff.max().item() - mean_diff = abs_diff.mean().item() - - rtol, atol = 1e-2, 1e-2 - assert torch.allclose(out_kernel, out_ref, rtol=rtol, atol=atol), ( - f"CuTe kernel with independent per-Q-block indices deviated from masked fp32 " - f"reference: max_diff={max_diff:.3e}, mean_diff={mean_diff:.3e} " - f"(rtol={rtol}, atol={atol}, pair_mismatch={pair_mismatch})" - ) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="kernel test needs CUDA") -def test_cute_kernel_50pct_sparsity_quality_vs_dense(): - """50% sparse CuTe kernel with score-based topk should stay close to dense SDPA.""" - - from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.video_sparse_attention import ( - CUTE_AVAILABLE, - block_sparse_attn_from_indices_cute, - is_cute_supported, - ) - - if not CUTE_AVAILABLE: - pytest.skip("cuda-bindings or cutlass-dsl not importable") - - device = torch.device("cuda") - dtype = torch.bfloat16 - torch.manual_seed(0) - - B, H, num_cubes, D = 1, 4, 16, 128 - block_size = 64 - topk = num_cubes // 2 - seq_len = num_cubes * block_size - - q = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) - k = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) - v = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) - - if not is_cute_supported(q): - pytest.skip("CuTe path needs sm_100+ Blackwell (current device unsupported)") - - q_blocks = q.reshape(B, H, num_cubes, block_size, D).mean(dim=3) - k_blocks = k.reshape(B, H, num_cubes, block_size, D).mean(dim=3) - scale = D**-0.5 - block_scores = torch.einsum("bhqd,bhkd->bhqk", q_blocks.float(), k_blocks.float()) * scale - q2k_idx = block_scores.topk(topk, dim=-1).indices.to(torch.int32).contiguous() - - q2k_num = torch.full((B, H, num_cubes), topk, dtype=torch.int32, device=device) - variable_block_sizes = torch.full((num_cubes,), block_size, dtype=torch.int32, device=device) - - out_sparse, _lse = block_sparse_attn_from_indices_cute( - q, k, v, q2k_idx, q2k_num, variable_block_sizes - ) - out_dense = F.scaled_dot_product_attention(q, k, v) - - cos_sim = F.cosine_similarity( - out_sparse.float().reshape(-1), out_dense.float().reshape(-1), dim=0 - ).item() - print(f"\n 50% sparse (score-based topk) vs dense SDPA cos_sim: {cos_sim:.4f}") - - assert cos_sim >= 0.65, ( - f"50% sparse CuTe kernel deviated too far from dense SDPA: cos_sim={cos_sim:.4f} < 0.65" - ) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="kernel test needs CUDA") -@pytest.mark.parametrize( - "num_cubes", - [1, 3, 9], - ids=["1cube_odd", "3cubes_odd", "9cubes_odd"], -) -def test_cute_kernel_odd_num_cubes_correctness(num_cubes): - """CuTe kernel with odd num_cubes must match dense SDPA (last Q-block has no pair).""" - from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.video_sparse_attention import ( - CUTE_AVAILABLE, - block_sparse_attn_from_indices_cute, - is_cute_supported, - ) - - if not CUTE_AVAILABLE: - pytest.skip("cuda-bindings or cutlass-dsl not importable") - - assert num_cubes % 2 == 1, f"pre-condition: num_cubes={num_cubes} must be odd" - - device = torch.device("cuda") - dtype = torch.bfloat16 - torch.manual_seed(0) - - B, H, D = 1, 4, 128 - block_size = 64 - seq_len = num_cubes * block_size - - q = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) - k = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) - v = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) - - if not is_cute_supported(q): - pytest.skip("CuTe path needs sm_100+ Blackwell (current device unsupported)") - - topk = num_cubes - q2k_idx = ( - torch.arange(num_cubes, device=device, dtype=torch.int32) - .view(1, 1, 1, num_cubes) - .expand(B, H, num_cubes, topk) - .contiguous() - ) - q2k_num = torch.full((B, H, num_cubes), topk, dtype=torch.int32, device=device) - variable_block_sizes = torch.full((num_cubes,), block_size, dtype=torch.int32, device=device) - - out_kernel, _lse = block_sparse_attn_from_indices_cute( - q, k, v, q2k_idx, q2k_num, variable_block_sizes - ) - out_ref = F.scaled_dot_product_attention(q, k, v) - - assert torch.isfinite(out_kernel).all(), ( - f"CuTe kernel produced non-finite output for odd num_cubes={num_cubes}" - ) - - max_diff = (out_kernel - out_ref).abs().max().item() - mean_diff = (out_kernel - out_ref).abs().mean().item() - rtol, atol = 1e-2, 1e-2 - assert torch.allclose(out_kernel, out_ref, rtol=rtol, atol=atol), ( - f"CuTe kernel deviated from dense SDPA for odd num_cubes={num_cubes}: " - f"max_diff={max_diff:.3e}, mean_diff={mean_diff:.3e} (rtol={rtol}, atol={atol})" - ) diff --git a/tests/unittest/_torch/visual_gen/test_attention_integration.py b/tests/unittest/_torch/visual_gen/test_attention_integration.py index a372db9641f2..a4c62657495a 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_integration.py +++ b/tests/unittest/_torch/visual_gen/test_attention_integration.py @@ -12,6 +12,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from utils.util import isSM100Family from tensorrt_llm._torch.modules.rms_norm import RMSNorm @@ -698,18 +699,19 @@ def test_fast_cross_attention_wan_shapes( # ============================================================================ -# VSA self-attention (CUTEDSL backend, sparse_attention_config.algorithm='vsa') +# VSA self-attention (CUTEDSL/TRTLLM backends) # ============================================================================ -def _build_vsa_setup(sparsity: float, batch_size: int, seed: int): +def _build_vsa_setup(backend: str, sparsity: float, batch_size: int, seed: int): """Build naive + integrated models, VSA metadata, and inputs for a VSA test. - latent (8,8,8) -> 512 tokens (divisible by block_size=64), head_dim=128. + A ragged latent exercises VSA padding and token-mask lowering on both + fine-stage implementations. """ - from tensorrt_llm._torch.visual_gen.attention_backend import VSAMetadataBuilder + from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import VSAMetadataBuilder - latent_shape = (8, 8, 8) + latent_shape = (9, 9, 9) seq_len = latent_shape[0] * latent_shape[1] * latent_shape[2] num_heads = 4 head_dim = 128 @@ -717,17 +719,21 @@ def _build_vsa_setup(sparsity: float, batch_size: int, seed: int): device = torch.device("cuda") dtype = torch.bfloat16 + torch.manual_seed(seed) naive = NaiveWanSelfAttention(hidden_size, num_heads, head_dim, dtype=dtype).to(device) cfg_vsa = create_model_config( - hidden_size, num_heads, head_dim, attn_backend="CUTEDSL", vsa_sparsity=sparsity + hidden_size, + num_heads, + head_dim, + attn_backend=backend, + vsa_sparsity=sparsity, ) integrated = Attention(hidden_size, num_heads, qkv_mode=QKVMode.FUSE_QKV, config=cfg_vsa).to( device ) - # Fail loudly if the VSA path silently fell back to dense (which would set - # attn_backend to "VANILLA") instead of selecting the CUTEDSL/VSA backend. - assert integrated.attn_backend == "CUTEDSL", ( - f"Expected CUTEDSL (VSA) backend, got {integrated.attn_backend!r}" + # Fail loudly if the VSA path silently fell back to the VANILLA backend. + assert integrated.attn_backend == backend, ( + f"Expected {backend} VSA backend, got {integrated.attn_backend!r}" ) copy_weights_self_attention(naive, integrated) naive.eval() @@ -754,12 +760,13 @@ def _build_vsa_setup(sparsity: float, batch_size: int, seed: int): @pytest.mark.skipif(not torch.cuda.is_available(), reason="VSA needs CUDA") -def test_vsa_self_attention_equivalence_at_sparsity_zero(): +@pytest.mark.parametrize("backend", ["CUTEDSL", "TRTLLM"]) +def test_vsa_self_attention_equivalence_at_sparsity_zero(backend: str): """VSA at sparsity=0 with G_c=0 reduces to dense attention (top_k=num_cubes, output=O_f); must match the naive SDPA reference modulo bf16 rounding.""" - from tensorrt_llm._torch.visual_gen.attention_backend import set_vsa_forward_context + from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import set_vsa_forward_context - s = _build_vsa_setup(sparsity=0.0, batch_size=2, seed=42) + s = _build_vsa_setup(backend=backend, sparsity=0.0, batch_size=2, seed=42) with torch.no_grad(): out_naive = s.naive(s.hidden_states, *s.freqs_HSD) @@ -780,22 +787,103 @@ def test_vsa_self_attention_equivalence_at_sparsity_zero(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="VSA needs CUDA") -@pytest.mark.parametrize("sparsity", [0.0, 0.5], ids=["s0", "s0p5"]) -def test_vsa_self_attention_finite(sparsity: float): - """VSA forward must produce finite output (no NaN/Inf) at any supported sparsity.""" - from tensorrt_llm._torch.visual_gen.attention_backend import set_vsa_forward_context +@pytest.mark.skipif( + not isSM100Family(), + reason="CuTe DSL and PrimTS block-sparse parity requires SM100 or SM103", +) +def test_vsa_sparse_backends_match_on_ragged_input(): + """CuTeDSL and TRTLLM implement the same sparse VSA fine-stage semantics.""" + from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import set_vsa_forward_context + + sparsity = 0.5 + setups = { + backend: _build_vsa_setup(backend=backend, sparsity=sparsity, batch_size=1, seed=0) + for backend in ("CUTEDSL", "TRTLLM") + } + sparse_fine_executed = {} + for backend, setup in setups.items(): + if backend == "CUTEDSL": + original_execute = setup.integrated.attn._execute_sparse_fine + + def checked_execute(*args, _original=original_execute, **kwargs): + result = _original(*args, **kwargs) + sparse_fine_executed["CUTEDSL"] = True + return result + + setup.integrated.attn._execute_sparse_fine = checked_execute + else: + original_predict = setup.integrated.attn.block_sparse_attn_predict + + def checked_predict(*args, _original=original_predict, **kwargs): + result = _original(*args, **kwargs) + sparse_fine_executed["TRTLLM"] = result is not None + return result + + setup.integrated.attn.block_sparse_attn_predict = checked_predict + + outputs = {} + for backend, setup in setups.items(): + with torch.no_grad(), set_vsa_forward_context(setup.metadata): + outputs[backend] = setup.integrated( + setup.hidden_states, + freqs=setup.freqs_SHD, + gate_compress=setup.gate_compress_zero, + ) - s = _build_vsa_setup(sparsity=sparsity, batch_size=1, seed=0) + assert sparse_fine_executed == {"CUTEDSL": True, "TRTLLM": True} + assert torch.isfinite(outputs["CUTEDSL"]).all() + assert torch.isfinite(outputs["TRTLLM"]).all() + torch.testing.assert_close( + outputs["CUTEDSL"], + outputs["TRTLLM"], + rtol=1e-2, + atol=1e-2, + ) - with torch.no_grad(), set_vsa_forward_context(s.metadata): - out = s.integrated(s.hidden_states, freqs=s.freqs_SHD, gate_compress=s.gate_compress_zero) - assert out.shape == s.hidden_states.shape - nan_count = torch.isnan(out).sum().item() - inf_count = torch.isinf(out).sum().item() - assert nan_count == 0 and inf_count == 0, ( - f"VSA produced non-finite output at sparsity={sparsity}: NaN={nan_count}, Inf={inf_count}" - ) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="VSA needs CUDA") +@pytest.mark.skipif( + not isSM100Family(), + reason="PrimTS block-sparse CUDA Graph replay requires SM100 or SM103", +) +def test_vsa_trtllm_cuda_graph_replays_live_routes(): + """Captured VSA recomputes routes when graph-stable Q/K/V storage changes.""" + from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import set_vsa_forward_context + + setup = _build_vsa_setup(backend="TRTLLM", sparsity=0.5, batch_size=1, seed=17) + static_hidden = setup.hidden_states.clone() + static_gate = setup.gate_compress_zero.clone() + + for _ in range(2): + with torch.no_grad(), set_vsa_forward_context(setup.metadata): + setup.integrated( + static_hidden, + freqs=setup.freqs_SHD, + gate_compress=static_gate, + ) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph), torch.no_grad(), set_vsa_forward_context(setup.metadata): + graph_output = setup.integrated( + static_hidden, + freqs=setup.freqs_SHD, + gate_compress=static_gate, + ) + + initial_output = graph_output.clone() + live_hidden = torch.randn_like(static_hidden) + static_hidden.copy_(live_hidden) + graph.replay() + replay_output = graph_output.clone() + with torch.no_grad(), set_vsa_forward_context(setup.metadata): + eager_output = setup.integrated( + live_hidden, + freqs=setup.freqs_SHD, + gate_compress=static_gate, + ) + + assert not torch.equal(initial_output, replay_output) + torch.testing.assert_close(replay_output, eager_output, rtol=1e-2, atol=1e-2) def test_trtllm_cached_prepare(): diff --git a/tests/unittest/_torch/visual_gen/test_attention_perf.py b/tests/unittest/_torch/visual_gen/test_attention_perf.py index a662b788cf0d..855366914932 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_perf.py +++ b/tests/unittest/_torch/visual_gen/test_attention_perf.py @@ -36,11 +36,6 @@ import pytest import torch -from tensorrt_llm._torch.visual_gen.attention_backend import ( - VSAMetadataBuilder, - set_vsa_forward_context, -) - # ============================================================================ # Flash Attention 4 availability # ============================================================================ @@ -49,6 +44,10 @@ from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import ( _flash_attn_fwd_import_error as _fa4_import_error, ) +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import ( + VSAMetadataBuilder, + set_vsa_forward_context, +) from tensorrt_llm._torch.visual_gen.config import ( DiffusionModelConfig, create_attention_metadata_state, @@ -1080,7 +1079,9 @@ def test_vsa_kernel_vs_fa4( block_size: int, sparsity: float, ): - from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl import VSA_KERNEL_MAX_CUBES + from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa.backend import ( + VSA_KERNEL_MAX_CUBES, + ) assert seq_len % block_size == 0, "seq_len must be a multiple of block_size" num_cubes = seq_len // block_size diff --git a/tests/unittest/_torch/visual_gen/test_attention_vsa.py b/tests/unittest/_torch/visual_gen/test_attention_vsa.py new file mode 100644 index 000000000000..82b90fd75dcd --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_attention_vsa.py @@ -0,0 +1,906 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""VSA correctness tests: backend dispatch, preprocessing, and kernel behavior. + +Module-level dense-equivalence and finite-output checks live in +test_attention_integration.py. +""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask +from tensorrt_llm._torch.attention.backends.sparse.params import BlockSparseForwardInputs +from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl import CuTeDSLAttention +from tensorrt_llm._torch.visual_gen.attention_backend.interface import AttentionTensorLayout +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import backend as vsa_backend +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa.backend import ( + VSACuTeDSLAttention, + VSATrtllmAttention, +) +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa.metadata import ( + VSAMetadataBuilder, + set_vsa_forward_context, +) +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa.predictor import ( + VSAForwardInputs, + VSAPredictor, + VSAPreprocessor, +) +from tensorrt_llm._torch.visual_gen.attention_backend.trtllm import TrtllmAttention +from tensorrt_llm._torch.visual_gen.attention_backend.utils import create_attention +from tensorrt_llm._torch.visual_gen.attention_backend.vanilla import VanillaAttention +from tensorrt_llm._torch.visual_gen.config import ( + DiffusionModelConfig, + create_attention_metadata_state, +) +from tensorrt_llm._torch.visual_gen.modules import attention as attention_module +from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode +from tensorrt_llm.visual_gen.args import AttentionConfig, VideoSparseAttentionConfig + + +def test_cute_vsa_backend_preserves_sparse_backend_contract() -> None: + attention = VSACuTeDSLAttention( + num_heads=4, + head_dim=128, + ) + + assert isinstance(attention, CuTeDSLAttention) + assert attention.preferred_layout == AttentionTensorLayout.NHD + assert not attention.support_lse() + with pytest.raises(NotImplementedError, match="VSA does not support LSE"): + attention.forward_with_lse(torch.empty(0), torch.empty(0), torch.empty(0)) + + +def _make_vsa_metadata(*, sparsity: float = 0.0): + return VSAMetadataBuilder().build( + current_timestep=0, + raw_latent_shape=(5, 4, 4), + patch_size=(1, 1, 1), + vsa_sparsity=sparsity, + device=torch.device("cpu"), + ) + + +def test_vsa_trtllm_overrides_only_forward_around_the_core() -> None: + assert "forward" in VSATrtllmAttention.__dict__ + for name in ( + "block_sparse_attn_predict", + "sparse_predict", + "sparse_post_process", + "_enable_sparse_workflow", + ): + assert name not in VSATrtllmAttention.__dict__ + + +def _capture_wrapper_forward(monkeypatch: pytest.MonkeyPatch) -> dict: + """Replace the VisualGen wrapper forward with a recorder returning the fine input.""" + + captured = {} + + def _forward( + self, + q, + k, + v, + batch_size, + seq_len, + attention_mask=PredefinedAttentionMask.FULL, + seq_len_kv=None, + sparse_backend_args=None, + **kwargs, + ): + captured.update( + q=q, + k=k, + v=v, + batch_size=batch_size, + seq_len=seq_len, + attention_mask=attention_mask, + seq_len_kv=seq_len_kv, + sparse_backend_args=sparse_backend_args, + kwargs=kwargs, + ) + return q.reshape(batch_size, seq_len, -1) + + monkeypatch.setattr(TrtllmAttention, "forward", _forward) + return captured + + +def test_vsa_backends_share_one_predictor_implementation() -> None: + trtllm_attention = object.__new__(VSATrtllmAttention) + cute_attention = object.__new__(VSACuTeDSLAttention) + trtllm_attention.predictor = VSAPredictor(num_heads=1) + cute_attention.predictor = VSAPredictor(num_heads=1) + + assert type(trtllm_attention.predictor) is type(cute_attention.predictor) is VSAPredictor + assert set(vsa_backend.__all__) >= {"VSATrtllmAttention", "VSACuTeDSLAttention"} + + +def test_vsa_trtllm_layers_share_model_scoped_predictor() -> None: + attention_metadata_state = create_attention_metadata_state() + first = VSATrtllmAttention( + layer_idx=0, + num_heads=2, + num_kv_heads=2, + head_dim=128, + attention_metadata_state=attention_metadata_state, + ) + second = VSATrtllmAttention( + layer_idx=1, + num_heads=2, + num_kv_heads=2, + head_dim=128, + attention_metadata_state=attention_metadata_state, + ) + + assert first.predictor is second.predictor + + +def test_vsa_predictor_produces_sorted_block_inputs_and_effective_tiled_qkv() -> None: + predictor = VSAPredictor(num_heads=1) + metadata = _make_vsa_metadata() + q = torch.randn(1, 80, 1, 8) + + inputs = predictor.predict( + q, + q, + q, + batch_size=1, + seq_len=80, + seq_len_kv=80, + attention_mask=PredefinedAttentionMask.FULL, + gate_compress=torch.zeros_like(q), + gate_fine=None, + use_sparse_fine=True, + produce_block_sparse_inputs=True, + metadata=metadata, + ) + + assert isinstance(inputs, VSAForwardInputs) + assert inputs.q.shape == inputs.k.shape == inputs.v.shape == (1, 128, 1, 8) + assert inputs.seq_len == 128 + block_sparse_inputs = inputs.block_sparse_inputs + assert isinstance(block_sparse_inputs, BlockSparseForwardInputs) + assert block_sparse_inputs.block_indptr.tolist() == [[[0, 2, 4]]] + assert block_sparse_inputs.block_indices.tolist() == [0, 1, 0, 1] + assert block_sparse_inputs.kv_valid_bits.dtype == torch.uint32 + assert block_sparse_inputs.kv_valid_bits.tolist() == [[0xFFFFFFFF, 0xFFFFFFFF, 0xFFFF, 0]] + + +def test_vsa_predictor_dense_fallback_keeps_compact_qkv_and_no_block_inputs() -> None: + predictor = VSAPredictor(num_heads=1) + metadata = _make_vsa_metadata(sparsity=0.5) + q = torch.randn(1, 80, 1, 8) + k = torch.randn_like(q) + v = torch.randn_like(q) + + inputs = predictor.predict( + q, + k, + v, + batch_size=1, + seq_len=80, + seq_len_kv=80, + attention_mask=PredefinedAttentionMask.FULL, + gate_compress=torch.zeros_like(q), + gate_fine=None, + use_sparse_fine=False, + produce_block_sparse_inputs=False, + metadata=metadata, + ) + + assert inputs.q is q + assert inputs.k is k + assert inputs.v is v + assert inputs.seq_len == 80 + assert inputs.block_sparse_inputs is None + assert inputs.post_context.untile_idx is None + + +def test_vsa_shared_post_process_restores_shape_and_applies_gates() -> None: + predictor = VSAPredictor(num_heads=1) + metadata = _make_vsa_metadata(sparsity=0.5) + q = torch.randn(1, 80, 1, 8) + gate_compress = torch.full_like(q, 2.0) + gate_fine = torch.full_like(q, 0.5) + inputs = predictor.predict( + q, + q, + q, + batch_size=1, + seq_len=80, + seq_len_kv=80, + attention_mask=PredefinedAttentionMask.FULL, + gate_compress=gate_compress, + gate_fine=gate_fine, + use_sparse_fine=False, + produce_block_sparse_inputs=False, + metadata=metadata, + ) + fine_output = torch.randn_like(q) + + output = vsa_backend.vsa_post_process(fine_output, inputs) + + expected = 2.0 * inputs.post_context.coarse_output + 0.5 * fine_output + assert output.shape == q.shape + torch.testing.assert_close(output, expected) + + +@pytest.mark.parametrize("backend", ["CUTEDSL", "TRTLLM"]) +def test_factory_composes_vsa_with_attention_backend( + monkeypatch: pytest.MonkeyPatch, + backend: str, +) -> None: + class _Backend: + def __init__(self, **kwargs) -> None: + self.kwargs = kwargs + + backend_name = "VSACuTeDSLAttention" if backend == "CUTEDSL" else "VSATrtllmAttention" + monkeypatch.setattr(vsa_backend, backend_name, _Backend) + sparse_config = VideoSparseAttentionConfig(vsa_sparsity=0.9) + attention = create_attention( + backend=backend, + layer_idx=0, + num_heads=8, + head_dim=128, + attention_config=AttentionConfig( + backend=backend, + sparse_attention_config=sparse_config, + ), + attention_metadata_state=( + create_attention_metadata_state() if backend == "TRTLLM" else None + ), + ) + + assert isinstance(attention, _Backend) + assert "sparse_params" not in attention.kwargs + + +def test_factory_preserves_local_vanilla_fallback_for_vsa() -> None: + attention = create_attention( + backend="VANILLA", + layer_idx=0, + num_heads=8, + head_dim=128, + attention_config=AttentionConfig( + backend="CUTEDSL", + sparse_attention_config=VideoSparseAttentionConfig(vsa_sparsity=0.9), + ), + ) + + assert isinstance(attention, VanillaAttention) + + +def _make_dense_fallback_vsa_attention() -> VSATrtllmAttention: + attention = object.__new__(VSATrtllmAttention) + attention.predictor = VSAPredictor(num_heads=1) + attention._fmha_manager = SimpleNamespace(fmha_libs=[]) + attention.quant_attention_config = None + return attention + + +def test_trtllm_vsa_dense_fallback_runs_compact_inputs_through_core( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured = _capture_wrapper_forward(monkeypatch) + attention = _make_dense_fallback_vsa_attention() + q = torch.randn(1, 80, 1, 8) + + with set_vsa_forward_context(_make_vsa_metadata(sparsity=0.5)): + output = attention.forward( + q, + q, + q, + batch_size=1, + seq_len=80, + gate_compress=torch.zeros_like(q), + ) + + assert captured["q"] is q + assert captured["k"] is q and captured["v"] is q + assert (captured["batch_size"], captured["seq_len"], captured["seq_len_kv"]) == (1, 80, 80) + assert captured["sparse_backend_args"] is None + assert output.shape == (1, 80, 8) + + +def test_trtllm_vsa_hands_predicted_routes_to_core_via_sparse_backend_args( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured = _capture_wrapper_forward(monkeypatch) + attention = _make_dense_fallback_vsa_attention() + attention._fmha_manager = SimpleNamespace( + fmha_libs=[object.__new__(vsa_backend.PrimsTSBlockSparseFmha)] + ) + monkeypatch.setattr(vsa_backend, "_get_unsupported_primts_reason", lambda *args: None) + q = torch.randn(1, 80, 1, 8) + + with set_vsa_forward_context(_make_vsa_metadata()): + output = attention.forward( + q, + q, + q, + batch_size=1, + seq_len=80, + gate_compress=torch.zeros_like(q), + ) + + assert captured["q"].shape == captured["k"].shape == captured["v"].shape == (1, 128, 1, 8) + assert (captured["seq_len"], captured["seq_len_kv"]) == (128, 128) + block_sparse_inputs = captured["sparse_backend_args"].block_sparse_inputs + assert isinstance(block_sparse_inputs, BlockSparseForwardInputs) + assert block_sparse_inputs.kv_valid_bits is not None + assert output.shape == (1, 80, 8) + + +def test_trtllm_vsa_accepts_packed_qkv_through_shared_predictor( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured = _capture_wrapper_forward(monkeypatch) + attention = _make_dense_fallback_vsa_attention() + qkv = tuple(torch.randn(1, 80, 1, 8) for _ in range(3)) + + with set_vsa_forward_context(_make_vsa_metadata(sparsity=0.5)): + attention.forward( + torch.stack(qkv, dim=2), + None, + None, + batch_size=1, + seq_len=80, + gate_compress=torch.zeros_like(qkv[0]), + ) + + for actual, expected in zip((captured["q"], captured["k"], captured["v"]), qkv): + torch.testing.assert_close(actual, expected) + + +def test_trtllm_vsa_consumes_gates_and_forwards_only_timestep( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured = _capture_wrapper_forward(monkeypatch) + attention = _make_dense_fallback_vsa_attention() + q = torch.randn(1, 80, 1, 8) + gate_compress = torch.full_like(q, 2.0) + gate_fine = torch.full_like(q, 0.5) + timestep = torch.tensor([12]) + + with set_vsa_forward_context(_make_vsa_metadata(sparsity=0.5)): + output = attention.forward( + q, + q, + q, + batch_size=1, + seq_len=80, + gate_compress=gate_compress, + gate_fine=gate_fine, + timestep=timestep, + ) + + assert captured["kwargs"] == {"timestep": timestep} + assert output.shape == (1, 80, 8) + assert torch.isfinite(output).all() + + +def test_cutedsl_vsa_rejects_unexpected_forward_kwargs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + attention = object.__new__(VSACuTeDSLAttention) + attention.predictor = VSAPredictor(num_heads=1) + q = torch.randn(1, 80, 1, 8) + monkeypatch.setattr(vsa_backend, "_vsa_import_error", RuntimeError("disabled for test")) + + with set_vsa_forward_context(_make_vsa_metadata(sparsity=0.5)): + with pytest.raises(TypeError, match="gate_fnne"): + attention.forward( + q, + q, + q, + gate_compress=torch.zeros_like(q), + gate_fnne=torch.zeros_like(q), + ) + + +def _make_config( + hidden_size: int, + num_heads: int, + head_dim: int, + backend: str, + vsa_sparsity: "float | None" = None, +) -> DiffusionModelConfig: + """Minimal DiffusionModelConfig for one Attention module.""" + pretrained_config = SimpleNamespace( + hidden_size=hidden_size, + num_attention_heads=num_heads, + attention_head_dim=head_dim, + eps=1e-6, + ) + sparse_attention_config = ( + VideoSparseAttentionConfig(vsa_sparsity=vsa_sparsity) if vsa_sparsity is not None else None + ) + config = DiffusionModelConfig( + pretrained_config=pretrained_config, + attention=AttentionConfig(backend=backend, sparse_attention_config=sparse_attention_config), + skip_create_weights_in_init=False, + ) + config.attention_metadata_state = ( + create_attention_metadata_state() if backend == "TRTLLM" else None + ) + return config + + +@pytest.mark.parametrize("backend", ["CUTEDSL", "TRTLLM"]) +@pytest.mark.parametrize( + ("is_self_attention", "expected_backend"), + [(False, "VANILLA"), (True, None)], + ids=["cross", "self"], +) +def test_vsa_separate_qkv_dispatches_by_attention_role( + monkeypatch: pytest.MonkeyPatch, + backend: str, + is_self_attention: bool, + expected_backend: str | None, +) -> None: + monkeypatch.setattr( + attention_module, + "create_attention", + lambda *, backend, **kwargs: SimpleNamespace(backend=backend, kwargs=kwargs), + ) + cfg = _make_config( + hidden_size=64, + num_heads=4, + head_dim=16, + backend=backend, + vsa_sparsity=0.5, + ) + attention = Attention( + 64, + 4, + qkv_mode=QKVMode.SEPARATE_QKV, + config=cfg, + separate_qkv_is_self_attention=is_self_attention, + ) + + assert attention.attn_backend == (expected_backend or backend) + + +def test_plain_trtllm_separate_qkv_self_attention_keeps_vanilla_fallback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + attention_module, + "create_attention", + lambda *, backend, **kwargs: SimpleNamespace(backend=backend, kwargs=kwargs), + ) + cfg = _make_config( + hidden_size=64, + num_heads=4, + head_dim=16, + backend="TRTLLM", + ) + + attention = Attention( + 64, + 4, + qkv_mode=QKVMode.SEPARATE_QKV, + config=cfg, + separate_qkv_is_self_attention=True, + ) + + assert attention.attn_backend == "VANILLA" + + +def test_vsa_with_attn2d_raises(): + """VSA + Attention2D must error at construction (VSA needs the full sequence per rank).""" + pretrained_config = SimpleNamespace( + hidden_size=64, + num_attention_heads=4, + attention_head_dim=16, + eps=1e-6, + ) + cfg = DiffusionModelConfig( + pretrained_config=pretrained_config, + attention=AttentionConfig( + backend="CUTEDSL", + sparse_attention_config=VideoSparseAttentionConfig(vsa_sparsity=0.0), + ), + skip_create_weights_in_init=False, + ) + cfg.visual_gen_mapping = SimpleNamespace( + ring_size=1, + ring_group=None, + ulysses_size=1, + ulysses_group=None, + attn2d_row_size=2, + attn2d_col_size=2, + attn2d_row_group=None, + attn2d_col_group=None, + cp_size=4, + ) + with pytest.raises(ValueError, match="incompatible with context parallelism"): + Attention(64, 4, qkv_mode=QKVMode.FUSE_QKV, config=cfg) + + +def test_vsa_metadata_builder_reuses_shape_tensors_with_live_step_policy() -> None: + builder = VSAMetadataBuilder() + build_args = { + "raw_latent_shape": (9, 9, 9), + "patch_size": (1, 1, 1), + "device": torch.device("cpu"), + } + + first = builder.build(current_timestep=3, vsa_sparsity=0.25, **build_args) + second = builder.build(current_timestep=4, vsa_sparsity=0.75, **build_args) + + assert first is not second + assert (first.current_timestep, first.vsa_sparsity) == (3, 0.25) + assert (second.current_timestep, second.vsa_sparsity) == (4, 0.75) + assert second.gather_idx is first.gather_idx + assert first.num_cubes == 27 + + builder.clear() + + rebuilt = builder.build(current_timestep=5, vsa_sparsity=0.5, **build_args) + assert rebuilt.gather_idx is not first.gather_idx + + +def test_vsa_graph_stable_caches_bound_shape_profiles() -> None: + builder = VSAMetadataBuilder(max_cached_shapes=1) + build_args = { + "current_timestep": 0, + "patch_size": (1, 1, 1), + "vsa_sparsity": 0.5, + "device": torch.device("cpu"), + } + builder.build(raw_latent_shape=(4, 4, 4), **build_args) + with pytest.raises(RuntimeError, match="metadata cache reached its 1-shape limit"): + builder.build(raw_latent_shape=(8, 4, 4), **build_args) + + route_builder = VSAPredictor(num_heads=1, max_cached_shapes=1)._route_builder + kv_valid_bits = torch.ones((1, 1), dtype=torch.uint32) + route_builder.from_selected_blocks(torch.zeros((1, 1, 1, 1), dtype=torch.int32), kv_valid_bits) + with pytest.raises(RuntimeError, match="route cache reached its 1-shape limit"): + route_builder.from_selected_blocks( + torch.zeros((1, 1, 2, 1), dtype=torch.int32), kv_valid_bits + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="VSA needs CUDA") +@pytest.mark.parametrize( + "latent_shape", + [ + (8, 8, 8), + (9, 9, 9), + (21, 45, 80), + ], + ids=["clean_8x8x8", "ragged_9x9x9", "wan720p_21x45x80"], +) +def test_vsa_tile_untile_roundtrip(latent_shape): + """VSAPreprocessor.tile then .untile must losslessly reproduce the input.""" + device = torch.device("cuda") + dtype = torch.bfloat16 + torch.manual_seed(0) + + B, H, D = 2, 4, 32 + seq_len = latent_shape[0] * latent_shape[1] * latent_shape[2] + + builder = VSAMetadataBuilder() + meta = builder.build( + current_timestep=0, + raw_latent_shape=latent_shape, + patch_size=(1, 1, 1), + vsa_sparsity=0.0, + device=device, + ) + + x = torch.randn(B, seq_len, H, D, device=device, dtype=dtype) + + x_tiled = VSAPreprocessor.tile( + x, + meta.non_pad_index, + meta.gather_idx, + meta.padded_seq_length, + ) + + pad_mask = torch.ones(meta.padded_seq_length, dtype=torch.bool, device=device) + pad_mask[meta.non_pad_index] = False + if pad_mask.any(): + assert x_tiled[:, pad_mask, :, :].abs().max().item() == 0.0, ( + "tile() must zero-fill padded positions" + ) + + x_roundtrip = VSAPreprocessor.untile( + x_tiled, + meta.untile_idx, + ) + + assert x_roundtrip.shape == x.shape, ( + f"shape mismatch after tile/untile: {x_roundtrip.shape} vs {x.shape}" + ) + assert torch.equal(x_roundtrip, x), ( + f"tile/untile round-trip is not lossless for latent_shape={latent_shape}: " + f"max_diff={(x_roundtrip - x).abs().max().item():.3e}" + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="kernel test needs CUDA") +def test_cute_kernel_matches_dense_at_full_topk(): + """CuTe block-sparse kernel matches dense SDPA when every cube is selected.""" + from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.video_sparse_attention import ( + CUTE_AVAILABLE, + block_sparse_attn_from_indices_cute, + is_cute_supported, + ) + + if not CUTE_AVAILABLE: + pytest.skip("cuda-bindings or cutlass-dsl not importable") + + device = torch.device("cuda") + dtype = torch.bfloat16 + torch.manual_seed(0) + + B, H, num_cubes, D = 1, 4, 4, 128 + block_size = 64 + seq_len = num_cubes * block_size + + q = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) + k = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) + v = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) + + if not is_cute_supported(q): + pytest.skip("CuTe path needs sm_100+ Blackwell (current device unsupported)") + + topk = num_cubes + q2k_idx = ( + torch.arange(num_cubes, device=device, dtype=torch.int32) + .view(1, 1, 1, num_cubes) + .expand(B, H, num_cubes, topk) + .contiguous() + ) + q2k_num = torch.full((B, H, num_cubes), topk, dtype=torch.int32, device=device) + variable_block_sizes = torch.full((num_cubes,), block_size, dtype=torch.int32, device=device) + + out_kernel, _lse = block_sparse_attn_from_indices_cute( + q, k, v, q2k_idx, q2k_num, variable_block_sizes + ) + out_ref = F.scaled_dot_product_attention(q, k, v) + + max_diff = (out_kernel - out_ref).abs().max().item() + mean_diff = (out_kernel - out_ref).abs().mean().item() + + rtol, atol = 1e-2, 1e-2 + assert torch.allclose(out_kernel, out_ref, rtol=rtol, atol=atol), ( + f"CuTe block-sparse kernel deviates from dense SDPA at full top-K: " + f"max_diff={max_diff:.3e}, mean_diff={mean_diff:.3e} (rtol={rtol}, atol={atol})" + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="kernel test needs CUDA") +def test_cute_kernel_matches_ref_with_independent_indices(): + """CuTe kernel: paired Q-blocks (2i, 2i+1) attend to independent KV index lists.""" + from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.video_sparse_attention import ( + CUTE_AVAILABLE, + block_sparse_attn_from_indices_cute, + is_cute_supported, + ) + + if not CUTE_AVAILABLE: + pytest.skip("cuda-bindings or cutlass-dsl not importable") + + device = torch.device("cuda") + dtype = torch.bfloat16 + torch.manual_seed(42) + + B, H, num_cubes, D = 2, 4, 16, 128 + block_size = 64 + topk = num_cubes // 2 + seq_len = num_cubes * block_size + + q = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) + k = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) + v = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) + + if not is_cute_supported(q): + pytest.skip("CuTe path needs sm_100+ Blackwell (current device unsupported)") + + q2k_idx = ( + torch.stack( + [ + torch.randperm(num_cubes, device=device, dtype=torch.int32)[:topk] + for _ in range(B * H * num_cubes) + ] + ) + .view(B, H, num_cubes, topk) + .contiguous() + ) + + paired = q2k_idx.view(B, H, num_cubes // 2, 2, topk).sort(dim=-1).values + pair_mismatch = (paired[..., 0, :] != paired[..., 1, :]).sum().item() + assert pair_mismatch > 0, ( + "Pre-condition failed: random permutations matched across every pair; " + "re-seed or raise num_cubes." + ) + + q2k_num = torch.full((B, H, num_cubes), topk, dtype=torch.int32, device=device) + variable_block_sizes = torch.full((num_cubes,), block_size, dtype=torch.int32, device=device) + + attn_mask = torch.full( + (B, H, seq_len, seq_len), float("-inf"), device=device, dtype=torch.float32 + ) + for b in range(B): + for h in range(H): + for q_blk in range(num_cubes): + for ki in range(topk): + k_blk = q2k_idx[b, h, q_blk, ki].item() + qs = q_blk * block_size + ks = k_blk * block_size + attn_mask[b, h, qs : qs + block_size, ks : ks + block_size] = 0.0 + + out_kernel, _lse = block_sparse_attn_from_indices_cute( + q, k, v, q2k_idx, q2k_num, variable_block_sizes + ) + + scale = 1.0 / (D**0.5) + scores = (q.float() @ k.float().transpose(-2, -1)) * scale + scores = scores + attn_mask + probs = torch.softmax(scores, dim=-1) + out_ref = (probs @ v.float()).to(dtype) + + abs_diff = (out_kernel.float() - out_ref.float()).abs() + max_diff = abs_diff.max().item() + mean_diff = abs_diff.mean().item() + + rtol, atol = 1e-2, 1e-2 + assert torch.allclose(out_kernel, out_ref, rtol=rtol, atol=atol), ( + f"CuTe kernel with independent per-Q-block indices deviated from masked fp32 " + f"reference: max_diff={max_diff:.3e}, mean_diff={mean_diff:.3e} " + f"(rtol={rtol}, atol={atol}, pair_mismatch={pair_mismatch})" + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="kernel test needs CUDA") +def test_cute_kernel_50pct_sparsity_quality_vs_dense(): + """50% sparse CuTe kernel with score-based topk stays close to dense SDPA.""" + from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.video_sparse_attention import ( + CUTE_AVAILABLE, + block_sparse_attn_from_indices_cute, + is_cute_supported, + ) + + if not CUTE_AVAILABLE: + pytest.skip("cuda-bindings or cutlass-dsl not importable") + + device = torch.device("cuda") + dtype = torch.bfloat16 + torch.manual_seed(0) + + batch_size, num_heads, num_cubes, head_dim = 1, 4, 16, 128 + block_size = 64 + topk = num_cubes // 2 + seq_len = num_cubes * block_size + + q = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype) + k = torch.randn_like(q) + v = torch.randn_like(q) + + if not is_cute_supported(q): + pytest.skip("CuTe path needs sm_100+ Blackwell (current device unsupported)") + + q_blocks = q.reshape(batch_size, num_heads, num_cubes, block_size, head_dim).mean(dim=3) + k_blocks = k.reshape(batch_size, num_heads, num_cubes, block_size, head_dim).mean(dim=3) + block_scores = torch.einsum( + "bhqd,bhkd->bhqk", + q_blocks.float(), + k_blocks.float(), + ) * (head_dim**-0.5) + q2k_idx = block_scores.topk(topk, dim=-1).indices.to(torch.int32).contiguous() + q2k_num = torch.full( + (batch_size, num_heads, num_cubes), + topk, + dtype=torch.int32, + device=device, + ) + variable_block_sizes = torch.full( + (num_cubes,), + block_size, + dtype=torch.int32, + device=device, + ) + + out_sparse, _lse = block_sparse_attn_from_indices_cute( + q, + k, + v, + q2k_idx, + q2k_num, + variable_block_sizes, + ) + out_dense = F.scaled_dot_product_attention(q, k, v) + + cos_sim = F.cosine_similarity( + out_sparse.float().reshape(-1), + out_dense.float().reshape(-1), + dim=0, + ).item() + assert cos_sim >= 0.65, ( + f"50% sparse CuTe kernel deviated too far from dense SDPA: cos_sim={cos_sim:.4f} < 0.65" + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="kernel test needs CUDA") +@pytest.mark.parametrize( + "num_cubes", + [1, 3, 9], + ids=["1cube_odd", "3cubes_odd", "9cubes_odd"], +) +def test_cute_kernel_odd_num_cubes_correctness(num_cubes): + """CuTe kernel supports a final Q block that has no paired neighbor.""" + from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.video_sparse_attention import ( + CUTE_AVAILABLE, + block_sparse_attn_from_indices_cute, + is_cute_supported, + ) + + if not CUTE_AVAILABLE: + pytest.skip("cuda-bindings or cutlass-dsl not importable") + + assert num_cubes % 2 == 1 + device = torch.device("cuda") + dtype = torch.bfloat16 + torch.manual_seed(0) + + batch_size, num_heads, head_dim = 1, 4, 128 + block_size = 64 + seq_len = num_cubes * block_size + q = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype) + k = torch.randn_like(q) + v = torch.randn_like(q) + + if not is_cute_supported(q): + pytest.skip("CuTe path needs sm_100+ Blackwell (current device unsupported)") + + q2k_idx = ( + torch.arange(num_cubes, device=device, dtype=torch.int32) + .view(1, 1, 1, num_cubes) + .expand(batch_size, num_heads, num_cubes, num_cubes) + .contiguous() + ) + q2k_num = torch.full( + (batch_size, num_heads, num_cubes), + num_cubes, + dtype=torch.int32, + device=device, + ) + variable_block_sizes = torch.full( + (num_cubes,), + block_size, + dtype=torch.int32, + device=device, + ) + + out_kernel, _lse = block_sparse_attn_from_indices_cute( + q, + k, + v, + q2k_idx, + q2k_num, + variable_block_sizes, + ) + out_ref = F.scaled_dot_product_attention(q, k, v) + + assert torch.isfinite(out_kernel).all() + torch.testing.assert_close(out_kernel, out_ref, rtol=1e-2, atol=1e-2) diff --git a/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py b/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py index fb0b16ba40eb..23293bfe21f8 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py @@ -1,9 +1,20 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest import torch +from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask +from tensorrt_llm._torch.attention.backends.sparse.params import ( + BlockSparseForwardInputs, + SparseBackendForwardArgs, + SparseRuntimeParams, +) from tensorrt_llm._torch.visual_gen.attention_backend import trtllm as visual_trtllm +from tensorrt_llm._torch.visual_gen.config import create_attention_metadata_state class _FakeBaseTrtllmAttentionMetadata: @@ -19,6 +30,84 @@ def prepare(self): self.prepare_calls += 1 +def _make_block_sparse_inputs(): + return BlockSparseForwardInputs( + q_block_size=64, + kv_block_size=64, + max_blocks_per_row=1, + block_indptr=torch.tensor([[[0, 1]]], dtype=torch.int32), + block_indices=torch.tensor([0], dtype=torch.int32), + ) + + +class _StopAtFmhaDispatch(Exception): + pass + + +def _make_core_forward_metadata(): + metadata = object.__new__(visual_trtllm.BaseTrtllmAttentionMetadata) + seq_lens = torch.tensor([4], dtype=torch.int32) + metadata._seq_lens = seq_lens + metadata._seq_lens_kv = seq_lens + metadata._seq_lens_cuda = None + metadata.kv_cache_manager = None + metadata._max_seq_len_storage = 4 + metadata.use_paged_context_fmha = False + metadata.cu_q_seqlens = None + metadata.cu_kv_seqlens = None + metadata.enable_flash_mla = False + metadata.spec_bl_tree_first_sparse_mask_offset_kv = None + metadata.spec_decoding_bl_tree_mask = None + metadata.kv_lens_cuda_runtime = torch.tensor([4], dtype=torch.int32) + metadata.kv_lens_runtime = torch.tensor([4], dtype=torch.int32) + metadata.prompt_lens_cuda_runtime = torch.tensor([4], dtype=torch.int32) + metadata.prompt_lens_cpu_runtime = torch.tensor([4], dtype=torch.int32) + metadata.host_request_types_runtime = torch.tensor([0], dtype=torch.int32) + metadata.max_context_q_len_override = None + return metadata + + +def _make_wrapper(cls=visual_trtllm.TrtllmAttention, *, quant_attention_config=None): + attention = object.__new__(cls) + attention.quant_attention_config = quant_attention_config + return attention + + +def _capture_core_forward(monkeypatch, captured: dict): + prepared_metadata = object() + monkeypatch.setattr( + visual_trtllm.TrtllmAttention, + "_prepare_metadata", + lambda self, batch_size, seq_len: prepared_metadata, + ) + monkeypatch.setattr( + visual_trtllm.TrtllmAttention, + "_concat_qkv", + lambda self, q, k, v, batch_size, seq_len, kv_seq_len: torch.cat( + [ + q.reshape(batch_size * seq_len, -1), + k.reshape(batch_size * kv_seq_len, -1), + v.reshape(batch_size * kv_seq_len, -1), + ], + dim=-1, + ), + ) + + def _capture_base_forward(self, q, k, v, metadata, forward_args=None, **kwargs): + captured.update( + q=q, + k=k, + v=v, + metadata=metadata, + forward_args=forward_args, + kwargs=kwargs, + ) + return q[:, :16] + + monkeypatch.setattr(visual_trtllm.BaseTrtllmAttention, "forward", _capture_base_forward) + return prepared_metadata + + def test_trtllm_attention_metadata_caches_distinct_seq_lens(monkeypatch): monkeypatch.setattr( visual_trtllm, @@ -62,3 +151,278 @@ def test_trtllm_attention_metadata_caches_distinct_seq_lens(monkeypatch): assert first_cached_seq_lens.data_ptr() != second_cached_seq_lens.data_ptr() assert first_metadata.seq_lens is first_cached_seq_lens assert second_metadata.seq_lens is second_cached_seq_lens + + +def test_trtllm_attention_layers_share_block_sparse_plan_cache(monkeypatch): + from tensorrt_llm._torch.attention.backends.fmha import prims_ts_block_sparse + + def _base_update_quant_config(self, new_quant_config): + del new_quant_config + self._fmha_manager = SimpleNamespace( + fmha_libs=[prims_ts_block_sparse.PrimsTSBlockSparseFmha(self)] + ) + + def _base_init(self, **kwargs): + del kwargs + self.is_mla_enable = False + self.kv_lora_rank = None + self.v_head_dim = None + self.head_dim = 64 + self.update_quant_config(None) + + monkeypatch.setattr( + visual_trtllm.BaseTrtllmAttention, + "update_quant_config", + _base_update_quant_config, + ) + monkeypatch.setattr(visual_trtllm.BaseTrtllmAttention, "__init__", _base_init) + attention_metadata_state = create_attention_metadata_state() + assert "block_sparse_fmha_cache" not in attention_metadata_state + + first = visual_trtllm.TrtllmAttention( + attention_metadata_state=attention_metadata_state, + ) + second = visual_trtllm.TrtllmAttention( + attention_metadata_state=attention_metadata_state, + ) + + assert not hasattr(first, "_block_sparse_fmha_cache_state") + assert not hasattr(second, "_block_sparse_fmha_cache_state") + first_fmha = first._fmha_manager.fmha_libs[0] + second_fmha = second._fmha_manager.fmha_libs[0] + assert first_fmha._contiguous_wrappers is second_fmha._contiguous_wrappers + assert first_fmha._paged_wrappers is second_fmha._paged_wrappers + + first.update_quant_config(None) + first_fmha = first._fmha_manager.fmha_libs[0] + assert first_fmha._contiguous_wrappers is second_fmha._contiguous_wrappers + assert first_fmha._paged_wrappers is second_fmha._paged_wrappers + assert attention_metadata_state["fmha_caches"]["prims_ts_block_sparse"] == { + "contiguous_wrappers": {}, + "paged_wrappers": {}, + } + + other = visual_trtllm.TrtllmAttention( + attention_metadata_state=create_attention_metadata_state(), + ) + other_fmha = other._fmha_manager.fmha_libs[0] + assert first_fmha._contiguous_wrappers is not other_fmha._contiguous_wrappers + assert first_fmha._paged_wrappers is not other_fmha._paged_wrappers + + +def test_visual_gen_wrapper_does_not_define_its_own_prediction_lifecycle(): + assert not hasattr(visual_trtllm, "SparseForwardInputs") + for name in ( + "block_sparse_attn_predict", + "sparse_post_process", + "_forward_impl", + ): + assert name not in visual_trtllm.TrtllmAttention.__dict__ + assert getattr(visual_trtllm.TrtllmAttention, "__parameters__", ()) == () + + +def test_forward_rejects_unexpected_kwargs_before_metadata_or_core(monkeypatch): + prepare_metadata = Mock(return_value=object()) + core_forward = Mock(return_value=torch.empty(4, 16)) + monkeypatch.setattr(visual_trtllm.TrtllmAttention, "_prepare_metadata", prepare_metadata) + monkeypatch.setattr(visual_trtllm.BaseTrtllmAttention, "forward", core_forward) + attention = _make_wrapper() + + with pytest.raises(TypeError) as exc_info: + attention.forward( + torch.randn(1, 4, 6, 8), + None, + None, + batch_size=1, + seq_len=4, + attention_maks=PredefinedAttentionMask.FULL, + timstep=torch.tensor([12]), + ) + + assert str(exc_info.value) == ( + "Unexpected TRTLLM attention forward keyword arguments: attention_maks, timstep" + ) + prepare_metadata.assert_not_called() + core_forward.assert_not_called() + + +def test_forward_flattens_fused_qkv_without_copy(monkeypatch): + captured = {} + prepared_metadata = _capture_core_forward(monkeypatch, captured) + attention = _make_wrapper() + qkv = torch.randn(1, 4, 6, 8) + timestep = torch.tensor([12]) + + output = attention.forward(qkv, None, None, batch_size=1, seq_len=4, timestep=timestep) + + assert output.shape == (1, 4, 16) + assert captured["q"].shape == (4, 48) + assert captured["q"].data_ptr() == qkv.data_ptr() + assert captured["k"] is None and captured["v"] is None + assert captured["metadata"] is prepared_metadata + assert captured["forward_args"].timestep is timestep + assert captured["forward_args"].sparse_backend_args is None + assert captured["forward_args"].sparse_runtime_params == SparseRuntimeParams() + assert captured["kwargs"] == {} + + +def test_forward_fuses_separate_qkv_without_sparse_backend_args(monkeypatch): + captured = {} + _capture_core_forward(monkeypatch, captured) + attention = _make_wrapper() + q = torch.randn(1, 4, 2, 8) + k = torch.randn_like(q) + v = torch.randn_like(q) + + attention.forward(q, k, v, batch_size=1, seq_len=4) + + assert captured["q"].shape == (4, 48) + torch.testing.assert_close(captured["q"][:, :16], q.reshape(4, 16)) + assert captured["k"] is None and captured["v"] is None + assert captured["forward_args"].sparse_backend_args is None + + +def test_forward_hands_separate_qkv_and_backend_args_to_core_for_block_sparse_routes( + monkeypatch, +): + captured = {} + _capture_core_forward(monkeypatch, captured) + attention = _make_wrapper() + q = torch.randn(1, 4, 2, 8) + k = torch.randn_like(q) + v = torch.randn_like(q) + backend_args = SparseBackendForwardArgs(block_sparse_inputs=_make_block_sparse_inputs()) + + output = attention.forward( + q, + k, + v, + batch_size=1, + seq_len=4, + sparse_backend_args=backend_args, + ) + + assert output.shape == (1, 4, 16) + assert captured["q"].data_ptr() == q.data_ptr() + assert captured["k"].data_ptr() == k.data_ptr() + assert captured["v"].data_ptr() == v.data_ptr() + assert captured["q"].shape == captured["k"].shape == captured["v"].shape == (4, 16) + assert captured["forward_args"].sparse_backend_args is backend_args + assert captured["forward_args"].sparse_runtime_params == SparseRuntimeParams() + + +def test_forward_hands_separate_qkv_to_core_when_backend_rejects_fused_qkv(monkeypatch): + class _SeparateQkvAttention(visual_trtllm.TrtllmAttention): + @classmethod + def support_fused_qkv(cls) -> bool: + return False + + captured = {} + _capture_core_forward(monkeypatch, captured) + attention = _make_wrapper(_SeparateQkvAttention) + q = torch.randn(1, 4, 2, 8) + + attention.forward(q, q, q, batch_size=1, seq_len=4) + + assert captured["k"] is not None and captured["v"] is not None + assert captured["q"].shape == (4, 16) + assert captured["forward_args"].sparse_backend_args is None + + +def test_forward_applies_sage_quantization_to_separate_qkv(monkeypatch): + captured = {} + _capture_core_forward(monkeypatch, captured) + quant_cfg = SimpleNamespace(q_block_size=1, k_block_size=2, v_block_size=3, qk_dtype="int8") + attention = _make_wrapper(quant_attention_config=quant_cfg) + q = torch.randn(1, 4, 2, 8) + + attention.forward(q, q, q, batch_size=1, seq_len=4) + + forward_args = captured["forward_args"] + assert captured["k"] is not None and captured["v"] is not None + assert forward_args.sage_attn_num_elts_per_blk_q == 1 + assert forward_args.sage_attn_num_elts_per_blk_k == 2 + assert forward_args.sage_attn_num_elts_per_blk_v == 3 + assert forward_args.sage_attn_qk_int8 is True + + +def test_forward_requires_separate_qkv_for_block_sparse_routes(monkeypatch): + prepare_metadata = Mock(return_value=object()) + monkeypatch.setattr(visual_trtllm.TrtllmAttention, "_prepare_metadata", prepare_metadata) + attention = _make_wrapper() + backend_args = SparseBackendForwardArgs(block_sparse_inputs=_make_block_sparse_inputs()) + + with pytest.raises(ValueError, match="separate q, k, and v"): + attention.forward( + torch.randn(1, 4, 6, 8), + None, + None, + batch_size=1, + seq_len=4, + sparse_backend_args=backend_args, + ) + + prepare_metadata.assert_not_called() + + +def test_forward_rejects_block_sparse_routes_with_quant_config(monkeypatch): + prepare_metadata = Mock(return_value=object()) + monkeypatch.setattr(visual_trtllm.TrtllmAttention, "_prepare_metadata", prepare_metadata) + attention = _make_wrapper(quant_attention_config=object()) + q = torch.randn(1, 4, 2, 8) + backend_args = SparseBackendForwardArgs(block_sparse_inputs=_make_block_sparse_inputs()) + + with pytest.raises(ValueError, match="quant_attention_config"): + attention.forward(q, q, q, batch_size=1, seq_len=4, sparse_backend_args=backend_args) + + prepare_metadata.assert_not_called() + + +@pytest.mark.parametrize("has_block_sparse_inputs", [False, True]) +def test_forward_reaches_core_fmha_with_module_predicted_routes( + monkeypatch, + has_block_sparse_inputs, +): + metadata = _make_core_forward_metadata() + monkeypatch.setattr( + visual_trtllm.TrtllmAttention, + "_prepare_metadata", + lambda self, batch_size, seq_len: metadata, + ) + + attention = _make_wrapper() + attention.sparse_params = None + attention.is_mla_enable = False + attention.num_heads = 2 + attention.num_kv_heads = 2 + attention.head_dim = 8 + attention.get_local_layer_idx = Mock(return_value=0) + attention._ensure_rope_table_size = Mock() + attention.print_skip_softmax_stat = False + attention.kv_scale_orig_quant = None + attention.kv_scale_quant_orig = None + attention.sparse_kv_predict = Mock(return_value=(None, None)) + attention.sparse_attn_predict = Mock(return_value=(None, None)) + select_fmha = Mock(side_effect=_StopAtFmhaDispatch) + attention._fmha_manager = SimpleNamespace( + fmha_libs=[object()], + select=select_fmha, + ) + carrier = _make_block_sparse_inputs() if has_block_sparse_inputs else None + backend_args = SparseBackendForwardArgs(block_sparse_inputs=carrier) + q = torch.randn(1, 4, 2, 8) + k = torch.randn_like(q) + v = torch.randn_like(q) + + with pytest.raises(_StopAtFmhaDispatch): + attention.forward(q, k, v, batch_size=1, seq_len=4, sparse_backend_args=backend_args) + + select_fmha.assert_called_once() + attention.sparse_kv_predict.assert_called_once() + attention.sparse_attn_predict.assert_called_once() + core_forward_args = select_fmha.call_args.args[5] + assert core_forward_args.sparse_backend_args is backend_args + runtime_params = core_forward_args.sparse_runtime_params + assert isinstance(runtime_params, SparseRuntimeParams) + assert runtime_params.block_sparse_inputs is carrier + assert runtime_params.sparse_attn_indices_block_size == 0 diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py index f1f29e18bd75..19378f37fb70 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py @@ -18,9 +18,11 @@ ParallelConfig, QuantAttentionConfig, RuntimeLoRAConfig, + SkipSoftmaxAttentionConfig, TeaCacheConfig, TorchCompileConfig, VAEConfig, + VideoSparseAttentionConfig, VisualGenArgs, ) @@ -101,6 +103,49 @@ def test_quant_config_rejected_when_unsupported(self): ), ) + @pytest.mark.parametrize( + ("backend", "quant_config"), + [ + ( + "TRTLLM", + QuantAttentionConfig( + qk_dtype="fp8", + q_block_size=1, + k_block_size=1, + v_block_size=1, + ), + ), + ( + "CUTEDSL", + QuantAttentionConfig(qk_dtype="bf16", v_dtype="fp8"), + ), + ], + ) + def test_vsa_and_quantization_are_mutually_exclusive(self, backend, quant_config): + with pytest.raises( + ValidationError, match="VSA and quant_attention_config are mutually exclusive" + ): + AttentionConfig( + backend=backend, + quant_attention_config=quant_config, + sparse_attention_config=VideoSparseAttentionConfig(vsa_sparsity=0.9), + ) + + def test_skip_softmax_and_sage_quantization_can_be_combined(self): + attention = AttentionConfig( + backend="TRTLLM", + quant_attention_config=QuantAttentionConfig( + qk_dtype="int8", + q_block_size=1, + k_block_size=4, + v_block_size=1, + ), + sparse_attention_config=SkipSoftmaxAttentionConfig(threshold_scale_factor=0.3), + ) + + assert attention.sparse_attention_config is not None + assert attention.sparse_attention_config.algorithm == "skip_softmax" + @pytest.mark.parametrize( ("qk_dtype", "q_block_size", "k_block_size", "v_block_size"), [ diff --git a/tests/unittest/_torch/visual_gen/test_wan_vsa_pipeline.py b/tests/unittest/_torch/visual_gen/test_wan_vsa_pipeline.py index 979619e8e433..7c2eddf63e7f 100644 --- a/tests/unittest/_torch/visual_gen/test_wan_vsa_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_wan_vsa_pipeline.py @@ -148,7 +148,7 @@ def _assert_vsa_matches_dense( """Compare CuTe-DSL VSA against SDPA-fallback VSA (same gated formulation, different fine kernel).""" from unittest.mock import patch - from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl import vsa as _vsa_module + from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import backend as _vsa_module common_kwargs = dict( prompt=PROMPT, From 317e78b9661d34275132776c993d7881d091fdfc Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Fri, 4 Sep 2026 14:52:50 +0000 Subject: [PATCH 6/7] feat: add VisualGen SOL sparse attention The SOL predictor pools Q/K/V per 64-token block, derives key statistics over the block axis, and thresholds centroid scores into packed exact-block words with three Triton stages; the fp32 centroid is split into three bf16 terms so every tensor-core product stays exact. Routes reach the PrimTS block-sparse FMHA as a bitmask with the K/V block summaries as proxy routes. The VisualGen TRTLLM backend compacts separate Q/K/V through a compiled copy, which keeps slices of a fused QKV projection on a vectorized kernel. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- .../visual-gen/features/sparse-attention.md | 63 ++ .../attention_backend/sparse/sol/__init__.py | 19 + .../attention_backend/sparse/sol/backend.py | 150 +++ .../attention_backend/sparse/sol/kernels.py | 507 ++++++++++ .../attention_backend/sparse/sol/params.py | 131 +++ .../attention_backend/sparse/sol/predictor.py | 282 ++++++ .../visual_gen/attention_backend/trtllm.py | 22 +- .../visual_gen/attention_backend/utils.py | 8 + .../models/ltx2/pipeline_ltx2_two_stages.py | 1 + .../models/ltx2/transformer_ltx2.py | 28 +- .../_torch/visual_gen/models/modeling.py | 35 +- .../_torch/visual_gen/modules/attention.py | 26 +- tensorrt_llm/visual_gen/__init__.py | 3 + tensorrt_llm/visual_gen/args.py | 30 +- tensorrt_llm/visual_gen/sparse_attention.py | 75 ++ .../test_lists/test-db/l0_b200.yml | 3 + .../integration/test_lists/test-db/l0_cpu.yml | 2 + .../sparse_attention/test_sol_attention.py | 898 ++++++++++++++++++ .../sparse_attention/test_sol_predictor.py | 374 ++++++++ .../test_sol_predictor_kernels.py | 206 ++++ .../_torch/visual_gen/test_ltx2_pipeline.py | 180 ++++ .../_torch/visual_gen/test_visual_gen_args.py | 49 + 22 files changed, 3078 insertions(+), 14 deletions(-) create mode 100644 tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/__init__.py create mode 100644 tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/backend.py create mode 100644 tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/kernels.py create mode 100644 tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/params.py create mode 100644 tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/predictor.py create mode 100644 tests/unittest/_torch/visual_gen/sparse_attention/test_sol_attention.py create mode 100644 tests/unittest/_torch/visual_gen/sparse_attention/test_sol_predictor.py create mode 100644 tests/unittest/_torch/visual_gen/sparse_attention/test_sol_predictor_kernels.py diff --git a/docs/source/visual-gen/features/sparse-attention.md b/docs/source/visual-gen/features/sparse-attention.md index 2d5926255190..1b1919f8dd8e 100644 --- a/docs/source/visual-gen/features/sparse-attention.md +++ b/docs/source/visual-gen/features/sparse-attention.md @@ -7,6 +7,7 @@ This page is an unindexed draft until the VisualGen documentation hub is introdu - [Overview](#overview) - [Algorithms](#algorithms) - [Skip Softmax Attention](#skip-softmax-attention) +- [SOL Attention](#sol-attention) - [Video Sparse Attention (VSA)](#video-sparse-attention-vsa) ## Overview @@ -21,6 +22,7 @@ Sparse attention in VisualGen is configured through `VisualGenArgs.attention_con |---|---|---| | `skip_softmax` | `SkipSoftmaxAttentionConfig` | Supported | | `vsa` | `VideoSparseAttentionConfig` | Supported (`CUTEDSL`, `TRTLLM`) | +| `sol_attn` | `SolAttentionConfig` | Experimental (`TRTLLM`) | ## Skip Softmax Attention @@ -214,6 +216,67 @@ attention_config: Graphs are captured lazily. The first denoising step seen for a given tensor shape and sparse-attention phase captures a graph; later steps with the same shape and phase replay that graph. When denoising crosses the cutoff, the phase key changes, so VisualGen captures a second graph for the enabled phase instead of replaying the graph from the disabled phase. +## SOL Attention + +SOL is a two-stage self-attention algorithm. A TRT-LLM-owned predictor first +produces an exact block bitmask and K/V proxy summaries from Q/K/V. The shared +`PrimsTSBlockSparseFmha` library then executes that route from the shared +`SparseRuntimeParams`. `SOLTrtllmAttention` is only the VisualGen bridge; +SOL does not add an algorithm-specific core attention backend or FMHA library. + +Configure SOL with `SolAttentionConfig` and the `TRTLLM` backend: + +```python +from tensorrt_llm.visual_gen import AttentionConfig, SolAttentionConfig + +attention_config = AttentionConfig( + backend="TRTLLM", + sparse_attention_config=SolAttentionConfig( + tau=1.0, + disabled_until_timestep=0.6, + dense_layers="0,2-4", + ), +) +``` + +The equivalent YAML is: + +```yaml +attention_config: + backend: TRTLLM + sparse_attention_config: + algorithm: sol_attn + tau: 1.0 + disabled_until_timestep: 0.6 + dense_layers: "0,2-4" +``` + +- `tau` is the finite float32 routing threshold consumed by the predictor. +- `disabled_until_timestep` is an optional normalized cutoff in `(0, 1]`. + Attention stays dense while the current timestep is greater than or equal to + the cutoff and switches to SOL below it. +- `dense_layers` is an optional comma-separated list of zero-based layer indices + and inclusive ranges that always use dense attention. + +The initial SOL envelope is full-mask BF16 self-attention on SM100 or SM103, +with 4-D BSHD Q/K/V tensors, equal Q/K/V shapes, and head dimension 128. +`SOLTrtllmAttention` overrides the core `block_sparse_attn_predict` hook, so +prediction runs inside the core forward from the flattened Q/K/V, the batch +layout in the attention metadata, and the `timestep` in the forward arguments; +dense layers and dense timestep phases return no routes. The VisualGen wrapper +compacts fused projection split views once and shares those tensors between +prediction and the block-sparse FMHA. Cross-attention, context parallelism, +attention quantization, and unsupported tensor envelopes raise an error instead +of silently falling back to dense attention. +SOL uses a host-side graph break to prepare and own predictor plans, so +`torch_compile_config.enable_fullgraph=True` is not supported; keep the default +`False` setting. + +When a cutoff is configured, VisualGen includes the dense-or-sparse phase in +the CUDA Graph key. Each SOL backend prepares that phase during graph warmup +and reuses it during capture, while predictor route buffers remain stable for +replay. A dense capture therefore cannot be reused for the sparse phase. + ## Video Sparse Attention (VSA) VSA combines a coarse mean-pooled branch with a top-K block-sparse fine branch. Select either `CUTEDSL` for the CuTe DSL kernel or `TRTLLM` for PrimTS block-sparse attention. If the selected sparse kernel is unavailable or the known VSA tensor envelope is not met, the fine branch uses the compact Q/K/V tensors with that backend's dense path. VSA cannot be combined with `quant_attention_config`. diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/__init__.py b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/__init__.py new file mode 100644 index 000000000000..e4aa9431c1e9 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/__init__.py @@ -0,0 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""TRT-LLM-owned predictor for two-stage SOL attention.""" + +from .predictor import ( + SolPredictorGeometry, + SolPredictorOutputs, + SolPredictorPlan, + SolPredictorPlanKey, + SOLSparsePredictor, +) + +__all__ = [ + "SOLSparsePredictor", + "SolPredictorGeometry", + "SolPredictorOutputs", + "SolPredictorPlan", + "SolPredictorPlanKey", +] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/backend.py b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/backend.py new file mode 100644 index 000000000000..270151266d02 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/backend.py @@ -0,0 +1,150 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""VisualGen SOL attention using the generic TRTLLM sparse lifecycle.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch.attention.backends.fmha.prims_ts_block_sparse import PrimsTSBlockSparseFmha +from tensorrt_llm._torch.attention.backends.fmha.utils import get_bmm1_scale +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + PredefinedAttentionMask, +) +from tensorrt_llm._torch.attention.backends.sparse.params import BlockSparseForwardInputs +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata + +from ...trtllm import TrtllmAttention +from .params import SolParams +from .predictor import BLOCK_SIZE, SOLSparsePredictor + + +class SOLTrtllmAttention(TrtllmAttention): + """Predict SOL routes inside the core prediction hook, then execute them + through the generic block-sparse FMHA.""" + + def __init__(self, *, sparse_params: SolParams | None = None, **kwargs) -> None: + if not isinstance(sparse_params, SolParams): + raise TypeError("SOLTrtllmAttention requires SolParams") + self.sol_params = sparse_params + self._prepared_graph_phase: int | None = None + super().__init__(sparse_params=None, **kwargs) + attention_metadata_state = kwargs["attention_metadata_state"] + predictor_cache = attention_metadata_state.setdefault("sparse_predictors", {}) + predictor = predictor_cache.get("sol_attn") + if predictor is None: + predictor = SOLSparsePredictor() + predictor_cache["sol_attn"] = predictor + elif not isinstance(predictor, SOLSparsePredictor): + raise TypeError("model-scoped SOL predictor cache contains an invalid value") + self.predictor = predictor + + def _resolve_graph_phase(self, timestep: object) -> int | None: + """Resolve the dense-or-sparse phase, reusing the warmup value under capture.""" + + if self.sol_params.disabled_until_timestep is None: + return None + if torch.cuda.is_current_stream_capturing(): + if self._prepared_graph_phase is None: + raise RuntimeError("SOL graph phase must be prepared before CUDA Graph capture") + return self._prepared_graph_phase + graph_phase = self.sol_params.get_graph_phase_for_timestep( + timestep, + disabled_until_timestep=self.sol_params.disabled_until_timestep, + ) + self._prepared_graph_phase = graph_phase + return graph_phase + + def block_sparse_attn_predict( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: TrtllmAttentionMetadata, + forward_args: AttentionForwardArgs, + ) -> BlockSparseForwardInputs | None: + """Return SOL routes for sparse calls and ``None`` for dense calls. + + ``q``, ``k``, and ``v`` arrive in the flattened ``[B*S, H*D]`` core + layout; the batch layout comes from ``metadata`` and the timestep from + ``forward_args``. + """ + + timestep = forward_args.timestep + graph_phase = self._resolve_graph_phase(timestep) + if not self.sol_params.should_use_sparse( + layer_idx=self.layer_idx, + timestep=timestep, + graph_phase=graph_phase, + ): + return None + + if self.quant_attention_config is not None: + raise ValueError("SOL sparse execution does not support quant_attention_config") + if not any( + isinstance(fmha, PrimsTSBlockSparseFmha) for fmha in self._fmha_manager.fmha_libs + ): + raise RuntimeError("SOL sparse execution requires PrimTS block-sparse FMHA") + if forward_args.attention_mask != PredefinedAttentionMask.FULL: + raise ValueError("SOL sparse execution requires a full attention mask") + if k is None or v is None: + raise ValueError("SOL sparse execution requires separate q, k, and v tensors") + + batch_size = metadata.num_seqs + seq_len = metadata.max_seq_len + num_tokens = batch_size * seq_len + if q.shape[0] != num_tokens or k.shape[0] != num_tokens or v.shape[0] != num_tokens: + raise ValueError( + "SOL sparse execution supports only uniform-length self-attention; " + f"got {q.shape[0]} query and {k.shape[0]} key tokens for " + f"{batch_size} sequences of length {seq_len}" + ) + + # The VisualGen wrapper compacts the flattened tensors once; these views + # are shared between prediction and the generic block-sparse FMHA. + q = q.view(batch_size, seq_len, self.num_heads, self.head_dim) + k = k.view(batch_size, seq_len, self.num_kv_heads, self.head_dim) + v = v.view(batch_size, seq_len, self.num_kv_heads, self.head_dim) + unsupported_reason = self.predictor.support_reason(q, k, v) + if unsupported_reason is not None: + raise ValueError(unsupported_reason) + + outputs = self.predictor.predict( + q, + k, + v, + tau=self.sol_params.tau, + sm_scale=get_bmm1_scale(self), + ) + return BlockSparseForwardInputs( + q_block_size=BLOCK_SIZE, + kv_block_size=BLOCK_SIZE, + exact_block_bits=outputs.exact_block_bits, + k_summary=outputs.k_summary, + v_summary=outputs.v_summary, + ) + + @classmethod + def support_fused_qkv(cls) -> bool: + """SOL prediction requires separate Q, K, and V tensors.""" + + return False + + +__all__ = ["SOLTrtllmAttention"] 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 new file mode 100644 index 000000000000..139a034b4aeb --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/kernels.py @@ -0,0 +1,507 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Memory-bound kernels of the two-stage SOL predictor. + +The predictor summarises ``[batch, tokens, heads, head_dim]`` activations per token block, derives per +channel key statistics, and thresholds centroid scores into packed exact-block words. CUDA tensors run +Triton kernels; other tensors use PyTorch implementations of the same rule. Launch shapes are derived +from tensor shapes, so no autotuning happens at call time and every launch is CUDA Graph safe. Every +helper writes into caller-owned storage so a plan can keep graph-stable outputs. +""" + +from __future__ import annotations + +import math +from typing import Literal + +import torch +import triton +import triton.language as tl + +_POOL_MAX_WIDTH = 1024 +_POOL_TOKENS_PER_LOAD = 8 +_STATS_WIDTH = 128 +_STATS_ROWS_PER_LOAD = 32 +_SELECT_Q_BLOCKS = 64 +_WORD_BITS = 32 +_LOG2_E = math.log2(math.e) +_THRESHOLD_EPSILON = 1.0e-6 +_LOCAL_RADIUS = 1 + +_Reduce = Literal["mean", "sum"] + + +def _column_launch(row_width: int, max_width: int) -> tuple[int, int]: + """Columns per program and number of column chunks for a row of ``row_width`` channels.""" + width = min(max_width, triton.next_power_of_2(row_width)) + return width, triton.cdiv(row_width, width) + + +def num_blocks(seq_len: int, block_size: int) -> int: + return (seq_len + block_size - 1) // block_size + + +def num_words(num_kv_blocks: int) -> int: + return (num_kv_blocks + _WORD_BITS - 1) // _WORD_BITS + + +# --------------------------------------------------------------------------- block pooling +@triton.jit +def _block_pool_kernel( + x_ptr, + out_ptr, + seq_len, + num_blocks, + num_chunks, + row_width, + stride_x_batch, + stride_x_token, + stride_out_batch, + stride_out_block, + MEAN: tl.constexpr, + BLOCK: tl.constexpr, + TOKENS: tl.constexpr, + WIDTH: tl.constexpr, +): + """One program per (batch, block, column chunk): fp32 sum over the block's valid tokens.""" + pid = tl.program_id(0).to(tl.int64) + chunk = pid % num_chunks + batch_block = pid // num_chunks + block = batch_block % num_blocks + batch = batch_block // num_blocks + columns = chunk * WIDTH + tl.arange(0, WIDTH) + in_row = columns < row_width + first_token = block * BLOCK + total = tl.zeros([WIDTH], dtype=tl.float32) + for start in range(0, BLOCK, TOKENS): + tokens = first_token + start + tl.arange(0, TOKENS) + values = tl.load( + x_ptr + batch * stride_x_batch + tokens[:, None] * stride_x_token + columns[None, :], + mask=(tokens < seq_len)[:, None] & in_row[None, :], + other=0.0, + ) + total += tl.sum(values.to(tl.float32), axis=0) + if MEAN: + total = total / tl.minimum(seq_len - first_token, BLOCK).to(tl.float32) + tl.store( + out_ptr + batch * stride_out_batch + block * stride_out_block + columns, + total.to(out_ptr.dtype.element_ty), + mask=in_row, + ) + + +def _block_pool_torch( + x: torch.Tensor, out: torch.Tensor, *, block_size: int, reduce: _Reduce +) -> None: + batch_size, seq_len, num_heads, head_dim = x.shape + blocks = num_blocks(seq_len, block_size) + padded = torch.nn.functional.pad(x, (0, 0, 0, 0, 0, blocks * block_size - seq_len)) + total = padded.view(batch_size, blocks, block_size, num_heads, head_dim).sum( + dim=2, dtype=torch.float32 + ) + if reduce == "mean": + valid = torch.clamp( + seq_len - torch.arange(blocks, device=x.device) * block_size, max=block_size + ) + total = total / valid.to(torch.float32).view(1, -1, 1, 1) + out.copy_(total) + + +def block_pool(x: torch.Tensor, out: torch.Tensor, *, block_size: int, reduce: _Reduce) -> None: + """Reduce every run of ``block_size`` tokens of ``x`` into ``out`` with an fp32 accumulator. + + Args: + x: ``[batch, seq_len, heads, head_dim]`` activations; heads and head_dim must be contiguous, + the batch and token strides are arbitrary. + out: Contiguous ``[batch, ceil(seq_len / block_size), heads, head_dim]`` buffer of any float + dtype; it receives the rounded fp32 result. + block_size: Tokens per block; the final block may be shorter. + reduce: ``"mean"`` divides by the number of valid tokens of the block, ``"sum"`` does not. + """ + if reduce not in ("mean", "sum"): + raise ValueError(f"reduce must be 'mean' or 'sum'; got {reduce!r}") + if x.ndim != 4 or x.stride(3) != 1 or x.stride(2) != x.shape[3]: + raise ValueError( + "x must be [batch, seq_len, heads, head_dim] with contiguous heads and head_dim" + ) + batch_size, seq_len, num_heads, head_dim = x.shape + blocks = num_blocks(seq_len, block_size) + expected = (batch_size, blocks, num_heads, head_dim) + if tuple(out.shape) != expected or not out.is_contiguous(): + raise ValueError( + f"out must be a contiguous tensor of shape {expected}; got {tuple(out.shape)}" + ) + if x.device.type != "cuda": + _block_pool_torch(x, out, block_size=block_size, reduce=reduce) + return + row_width = num_heads * head_dim + width, chunks = _column_launch(row_width, _POOL_MAX_WIDTH) + _block_pool_kernel[(batch_size * blocks * chunks,)]( + x, + out, + seq_len, + blocks, + chunks, + row_width, + x.stride(0), + x.stride(1), + out.stride(0), + out.stride(1), + MEAN=reduce == "mean", + BLOCK=block_size, + TOKENS=min(_POOL_TOKENS_PER_LOAD, block_size), + WIDTH=width, + num_warps=4, + ) + + +# --------------------------------------------------------------------------- block statistics +@triton.jit +def _block_statistics_kernel( + x_ptr, + mean_ptr, + var_ptr, + num_blocks, + num_chunks, + row_width, + stride_x_batch, + stride_x_block, + stride_out_batch, + ROWS: tl.constexpr, + WIDTH: tl.constexpr, +): + """One program per (batch, column chunk): mean and clamped variance over the block axis.""" + pid = tl.program_id(0).to(tl.int64) + chunk = pid % num_chunks + batch = pid // num_chunks + columns = chunk * WIDTH + tl.arange(0, WIDTH) + in_row = columns < row_width + total = tl.zeros([WIDTH], dtype=tl.float32) + total_sq = tl.zeros([WIDTH], dtype=tl.float32) + for start in range(0, num_blocks, ROWS): + rows = start + tl.arange(0, ROWS) + values = tl.load( + x_ptr + batch * stride_x_batch + rows[:, None] * stride_x_block + columns[None, :], + mask=(rows < num_blocks)[:, None] & in_row[None, :], + other=0.0, + ).to(tl.float32) + total += tl.sum(values, axis=0) + total_sq += tl.sum(values * values, axis=0) + count = num_blocks.to(tl.float32) + mean = total / count + variance = tl.maximum(total_sq / count - mean * mean, 0.0) + tl.store(mean_ptr + batch * stride_out_batch + columns, mean, mask=in_row) + tl.store(var_ptr + batch * stride_out_batch + columns, variance, mask=in_row) + + +def _block_statistics_torch(x: torch.Tensor, out_mean: torch.Tensor, out_var: torch.Tensor) -> None: + values = x.to(torch.float32) + mean = values.mean(dim=1) + out_mean.copy_(mean) + out_var.copy_(torch.clamp(values.square().mean(dim=1) - mean.square(), min=0.0)) + + +def block_statistics(x: torch.Tensor, out_mean: torch.Tensor, out_var: torch.Tensor) -> None: + """Per-channel mean and biased variance of ``x`` over its block axis. + + Args: + x: Contiguous ``[batch, num_blocks, heads, head_dim]`` block summaries. + out_mean: Contiguous fp32 ``[batch, heads, head_dim]`` buffer. + out_var: Contiguous fp32 ``[batch, heads, head_dim]`` buffer; negative rounding is clamped to zero. + """ + batch_size, blocks, num_heads, head_dim = x.shape + expected = (batch_size, num_heads, head_dim) + for name, tensor in (("out_mean", out_mean), ("out_var", out_var)): + if ( + tuple(tensor.shape) != expected + or tensor.dtype != torch.float32 + or not tensor.is_contiguous() + ): + raise ValueError(f"{name} must be a contiguous fp32 tensor of shape {expected}") + if not x.is_contiguous(): + raise ValueError("x must be contiguous") + if x.device.type != "cuda": + _block_statistics_torch(x, out_mean, out_var) + return + row_width = num_heads * head_dim + width, chunks = _column_launch(row_width, _STATS_WIDTH) + _block_statistics_kernel[(batch_size * chunks,)]( + x, + out_mean, + out_var, + blocks, + chunks, + row_width, + x.stride(0), + x.stride(1), + out_mean.stride(0), + ROWS=_STATS_ROWS_PER_LOAD, + WIDTH=width, + num_warps=4, + ) + + +# --------------------------------------------------------------------------- exact-block selection +@triton.jit +def _select_exact_blocks_kernel( + centroid_ptr, + keys_ptr, + mean_ptr, + var_ptr, + bits_ptr, + num_q_blocks, + num_kv_blocks, + num_words, + num_heads, + local_radius, + tau, + log2_scale, + epsilon, + stride_c_batch, + stride_c_block, + stride_c_head, + stride_k_batch, + stride_k_block, + stride_k_head, + stride_s_batch, + stride_s_head, + stride_b_batch, + stride_b_head, + stride_b_block, + Q_BLOCKS: tl.constexpr, + HEAD_DIM: tl.constexpr, +): + """One program per (batch, head, tile of Q_BLOCKS query blocks); emits every word of the tile. + + Scores are ``log2_scale * ``. The fp32 centroid is split into three terms + of the key dtype so every tensor-core product is exact and only the fp32 accumulation rounds. + """ + tile = tl.program_id(0) + batch_head = tl.program_id(1).to(tl.int64) + batch = batch_head // num_heads + head = batch_head % num_heads + q_blocks = tile * Q_BLOCKS + tl.arange(0, Q_BLOCKS) + q_valid = q_blocks < num_q_blocks + dims = tl.arange(0, HEAD_DIM) + + centroid = tl.load( + centroid_ptr + + batch * stride_c_batch + + q_blocks[:, None] * stride_c_block + + head * stride_c_head + + dims[None, :], + mask=q_valid[:, None], + other=0.0, + ) + key_mean = tl.load(mean_ptr + batch * stride_s_batch + head * stride_s_head + dims) + key_var = tl.load(var_ptr + batch * stride_s_batch + head * stride_s_head + dims) + projected_mean = tl.sum(centroid * key_mean[None, :], axis=1) * log2_scale + projected_var = tl.sum(centroid * centroid * key_var[None, :], axis=1) * log2_scale * log2_scale + threshold = projected_mean + tau * tl.sqrt(tl.maximum(projected_var, 0.0) + epsilon) + + high = centroid.to(keys_ptr.dtype.element_ty) + rest = centroid - high.to(tl.float32) + mid = rest.to(keys_ptr.dtype.element_ty) + low = (rest - mid.to(tl.float32)).to(keys_ptr.dtype.element_ty) + + lanes = tl.arange(0, 32) + lane_bits = 1 << lanes.to(tl.int64) + for word in range(num_words): + kv_blocks = word * 32 + lanes + kv_valid = kv_blocks < num_kv_blocks + keys = tl.load( + keys_ptr + + batch * stride_k_batch + + kv_blocks[:, None] * stride_k_block + + head * stride_k_head + + dims[None, :], + mask=kv_valid[:, None], + other=0.0, + ) + keys_t = tl.trans(keys) + scores = (tl.dot(high, keys_t) + tl.dot(mid, keys_t) + tl.dot(low, keys_t)) * log2_scale + distance = q_blocks[:, None] - kv_blocks[None, :] + is_local = (distance >= -local_radius) & (distance <= local_radius) + exact = kv_valid[None, :] & ((scores > threshold[:, None]) | is_local) + packed = tl.sum(tl.where(exact, lane_bits[None, :], 0), axis=1) + tl.store( + bits_ptr + + batch * stride_b_batch + + head * stride_b_head + + q_blocks * stride_b_block + + word, + packed.to(tl.int32), + mask=q_valid, + ) + + +def _select_exact_blocks_torch( + centroid: torch.Tensor, + k_summary: torch.Tensor, + k_mean: torch.Tensor, + k_var: torch.Tensor, + exact_block_bits: torch.Tensor, + *, + tau: float, + sm_scale: float, +) -> None: + log2_scale = float(sm_scale) * _LOG2_E + q = centroid.to(torch.float64) + k = k_summary.to(torch.float64) + projected_mean = torch.einsum("bqhd,bhd->bhq", q, k_mean.to(torch.float64)) * log2_scale + projected_var = ( + torch.einsum("bqhd,bhd->bhq", q.square(), k_var.to(torch.float64)) * log2_scale * log2_scale + ) + threshold = projected_mean + float(tau) * torch.sqrt( + torch.clamp(projected_var, min=0.0) + _THRESHOLD_EPSILON + ) + scores = torch.einsum("bqhd,bkhd->bhqk", q, k) * log2_scale + exact = scores > threshold.unsqueeze(-1) + num_kv_blocks = k_summary.shape[1] + ids = torch.arange(num_kv_blocks, device=centroid.device) + exact |= ((ids[:, None] - ids[None, :]).abs() <= _LOCAL_RADIUS)[None, None] + words = num_words(num_kv_blocks) + padded = torch.nn.functional.pad(exact, (0, words * _WORD_BITS - num_kv_blocks)) + weights = 1 << torch.arange(_WORD_BITS, dtype=torch.int64, device=centroid.device) + packed = (padded.view(*exact.shape[:-1], words, _WORD_BITS).to(torch.int64) * weights).sum( + dim=-1 + ) + exact_block_bits.copy_(packed.to(torch.uint32)) + + +def select_exact_blocks( + centroid: torch.Tensor, + k_summary: torch.Tensor, + k_mean: torch.Tensor, + k_var: torch.Tensor, + exact_block_bits: torch.Tensor, + *, + tau: float, + sm_scale: float, +) -> None: + """Pack the SOL exact-block decision of every (query block, key block) pair into ``exact_block_bits``. + + A key block is exact when ``sm_scale * log2(e) * `` exceeds the row threshold + ``mean + tau * sqrt(var + 1e-6)`` projected from the key statistics, or when it lies within one + block of the query block. Bit ``r`` of word ``w`` selects key block ``32 * w + r``; padding bits of + the final word are zero. + + Args: + centroid: Contiguous fp32 ``[batch, num_q_blocks, heads, head_dim]`` query block means. + k_summary: Contiguous ``[batch, num_kv_blocks, heads, head_dim]`` key block means (bf16 or fp16). + k_mean: fp32 ``[batch, heads, head_dim]`` mean of ``k_summary`` over its block axis. + k_var: fp32 ``[batch, heads, head_dim]`` variance of ``k_summary`` over its block axis. + exact_block_bits: Contiguous uint32 ``[batch, heads, num_q_blocks, ceil(num_kv_blocks / 32)]``. + tau: Threshold slope in standard deviations. + sm_scale: Softmax scale of the attention call. + """ + batch_size, q_blocks, num_heads, head_dim = centroid.shape + kv_blocks = k_summary.shape[1] + expected_bits = (batch_size, num_heads, q_blocks, num_words(kv_blocks)) + if tuple(exact_block_bits.shape) != expected_bits or exact_block_bits.dtype != torch.uint32: + raise ValueError(f"exact_block_bits must be uint32 of shape {expected_bits}") + if tuple(k_summary.shape) != (batch_size, kv_blocks, num_heads, head_dim): + raise ValueError("k_summary must match centroid in batch, heads, and head_dim") + if ( + centroid.dtype != torch.float32 + or not centroid.is_contiguous() + or not k_summary.is_contiguous() + ): + raise ValueError("centroid must be contiguous fp32 and k_summary contiguous") + if not exact_block_bits.is_contiguous(): + raise ValueError("exact_block_bits must be contiguous") + if centroid.device.type != "cuda": + _select_exact_blocks_torch( + centroid, k_summary, k_mean, k_var, exact_block_bits, tau=tau, sm_scale=sm_scale + ) + return + bits = exact_block_bits.view(torch.int32) + grid = (triton.cdiv(q_blocks, _SELECT_Q_BLOCKS), batch_size * num_heads) + _select_exact_blocks_kernel[grid]( + centroid, + k_summary, + k_mean, + k_var, + bits, + q_blocks, + kv_blocks, + num_words(kv_blocks), + num_heads, + _LOCAL_RADIUS, + float(tau), + float(sm_scale) * _LOG2_E, + _THRESHOLD_EPSILON, + centroid.stride(0), + centroid.stride(1), + centroid.stride(2), + k_summary.stride(0), + k_summary.stride(1), + k_summary.stride(2), + k_mean.stride(0), + k_mean.stride(1), + bits.stride(0), + bits.stride(1), + bits.stride(2), + Q_BLOCKS=_SELECT_Q_BLOCKS, + HEAD_DIM=head_dim, + num_warps=4, + ) + + +# --------------------------------------------------------------------------- graph-visible operator +@torch.library.custom_op( + "trtllm::visual_gen_sol_predictor", + mutates_args=( + "exact_block_bits", + "k_summary", + "v_summary", + "k_mean", + "k_var_diag", + "q_centroid", + ), + device_types="cuda", +) +def visual_gen_sol_predictor( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + exact_block_bits: torch.Tensor, + k_summary: torch.Tensor, + v_summary: torch.Tensor, + k_mean: torch.Tensor, + k_var_diag: torch.Tensor, + q_centroid: torch.Tensor, + block_size: int, + tau: float, + sm_scale: float, +) -> None: + """Update caller-owned SOL route and proxy tensors in place.""" + + 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_statistics(k_summary, k_mean, k_var_diag) + select_exact_blocks( + q_centroid, k_summary, k_mean, k_var_diag, exact_block_bits, tau=tau, sm_scale=sm_scale + ) + + +@torch.library.register_fake("trtllm::visual_gen_sol_predictor") +def _( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + exact_block_bits: torch.Tensor, + k_summary: torch.Tensor, + v_summary: torch.Tensor, + k_mean: torch.Tensor, + k_var_diag: torch.Tensor, + q_centroid: torch.Tensor, + block_size: int, + tau: float, + sm_scale: float, +) -> None: + return None + + +__all__ = ["block_pool", "block_statistics", "num_blocks", "num_words", "select_exact_blocks"] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/params.py b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/params.py new file mode 100644 index 000000000000..ac27bd90ea39 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/params.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Lowered parameters for two-stage VisualGen SOL attention.""" + +from __future__ import annotations + +import math +import numbers +import struct +from dataclasses import dataclass, field +from typing import Literal + +import torch + +from tensorrt_llm._torch.attention.backends.sparse.params import SparseParams + + +def _as_timestep_float(timestep: object) -> float | None: + if timestep is None: + return None + if isinstance(timestep, torch.Tensor): + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + raise RuntimeError("SOL graph phase must be precomputed before CUDA Graph capture") + if timestep.numel() == 0: + return None + # WAN I2V can carry one timestep per token, with reference tokens fixed + # at zero. Stay dense until every live token is below the cutoff. + timestep = timestep.amax().item() + if isinstance(timestep, bool) or not isinstance(timestep, numbers.Real): + raise TypeError("timestep must be a real scalar or tensor") + value = float(timestep) + if not math.isfinite(value): + raise ValueError("timestep must be finite") + return value + + +@dataclass(frozen=True, slots=True) +class SolParams(SparseParams): + """Static SOL policy lowered from the user-facing VisualGen config.""" + + algorithm: Literal["sol_attn"] = field(init=False, default="sol_attn") + tau: float = 1.0 + disabled_until_timestep: float | None = None + dense_layers: frozenset[int] = field(default_factory=frozenset) + + def __post_init__(self) -> None: + if isinstance(self.tau, bool) or not isinstance(self.tau, numbers.Real): + raise TypeError("tau must be a finite real number") + try: + tau = struct.unpack("=f", struct.pack("=f", float(self.tau)))[0] + except (OverflowError, ValueError, struct.error) as error: + raise ValueError("tau must be representable as float32") from error + if not math.isfinite(tau): + raise ValueError("tau must be finite") + object.__setattr__(self, "tau", tau) + + cutoff = self.disabled_until_timestep + if cutoff is not None: + if isinstance(cutoff, bool) or not isinstance(cutoff, numbers.Real): + raise TypeError("disabled_until_timestep must be a real number or None") + cutoff = float(cutoff) + if not math.isfinite(cutoff) or not 0.0 < cutoff <= 1.0: + raise ValueError("disabled_until_timestep must be in (0, 1]") + object.__setattr__(self, "disabled_until_timestep", cutoff) + + dense_layers = frozenset(self.dense_layers) + if any( + isinstance(layer, bool) or not isinstance(layer, int) or layer < 0 + for layer in dense_layers + ): + raise ValueError("dense_layers must contain only non-negative integers") + object.__setattr__(self, "dense_layers", dense_layers) + + @staticmethod + def get_graph_phase_for_timestep( + timestep: object, + *, + disabled_until_timestep: float | None, + ) -> int | None: + """Return 0 for the dense prefix and 1 for the sparse suffix.""" + + if disabled_until_timestep is None: + return None + value = _as_timestep_float(timestep) + if value is None: + return None + return int(value < disabled_until_timestep) + + def should_use_sparse( + self, + *, + layer_idx: int, + timestep: object, + graph_phase: int | None = None, + ) -> bool: + """Return whether this layer should execute the SOL sparse path.""" + + if layer_idx in self.dense_layers: + return False + if graph_phase is not None: + if graph_phase not in (0, 1): + raise ValueError("SOL graph_phase must be 0 or 1") + phase = graph_phase + else: + phase = self.get_graph_phase_for_timestep( + timestep, + disabled_until_timestep=self.disabled_until_timestep, + ) + if phase is None: + if self.disabled_until_timestep is not None: + raise ValueError( + "timestep is required when SOL disabled_until_timestep is configured" + ) + return True + return phase == 1 + + +__all__ = ["SolParams"] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/predictor.py b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/predictor.py new file mode 100644 index 000000000000..4fa91d52e09e --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/predictor.py @@ -0,0 +1,282 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Plan-owned runtime for the two-stage VisualGen SOL predictor.""" + +from __future__ import annotations + +import numbers +import struct +from dataclasses import dataclass + +import torch + +from . import kernels as _kernels + +BLOCK_SIZE = 64 +HEAD_DIM = 128 + + +def _positive_int(value: object, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be a Python integer") + if value <= 0: + raise ValueError(f"{name} must be positive") + return value + + +def _float32_scalar(value: object, name: str) -> float: + if isinstance(value, bool) or not isinstance(value, numbers.Real): + raise TypeError(f"{name} must be a finite Python real") + try: + result = struct.unpack("=f", struct.pack("=f", float(value)))[0] + except (OverflowError, TypeError, ValueError, struct.error) as error: + raise ValueError(f"{name} must be representable as float32") from error + if not -float("inf") < result < float("inf"): + raise ValueError(f"{name} must be finite") + return result + + +def _normalize_runtime_scalars(*, tau: object, sm_scale: object) -> tuple[float, float]: + """Validate and round the two dynamic selector scalars to binary32.""" + + effective_tau = _float32_scalar(tau, "tau") + effective_sm_scale = _float32_scalar(sm_scale, "sm_scale") + if effective_sm_scale <= 0.0: + raise ValueError("sm_scale must be positive") + return effective_tau, effective_sm_scale + + +@dataclass(frozen=True) +class SolPredictorGeometry: + """Static shape specialization for compact BF16 self-MHA.""" + + batch_size: int + seq_len: int + num_heads: int + head_dim: int + num_q_blocks: int + num_kv_blocks: int + exact_words: int + tail_tokens: int + + @classmethod + def create( + cls, + *, + batch_size: object, + seq_len: object, + num_heads: object, + head_dim: object = HEAD_DIM, + ) -> "SolPredictorGeometry": + batch = _positive_int(batch_size, "batch_size") + tokens = _positive_int(seq_len, "seq_len") + heads = _positive_int(num_heads, "num_heads") + dim = _positive_int(head_dim, "head_dim") + if dim != HEAD_DIM: + raise ValueError(f"SOL predictor only supports head_dim={HEAD_DIM}; got {dim}") + blocks = _kernels.num_blocks(tokens, BLOCK_SIZE) + tail = tokens - (blocks - 1) * BLOCK_SIZE + return cls( + batch_size=batch, + seq_len=tokens, + num_heads=heads, + head_dim=dim, + num_q_blocks=blocks, + num_kv_blocks=blocks, + exact_words=_kernels.num_words(blocks), + tail_tokens=tail, + ) + + @property + def tensor_shape(self) -> tuple[int, int, int, int]: + return (self.batch_size, self.seq_len, self.num_heads, self.head_dim) + + @property + def summary_shape(self) -> tuple[int, int, int, int]: + return (self.batch_size, self.num_kv_blocks, self.num_heads, self.head_dim) + + @property + def stats_shape(self) -> tuple[int, int, int]: + return (self.batch_size, self.num_heads, self.head_dim) + + @property + def exact_block_bits_shape(self) -> tuple[int, int, int, int]: + return (self.batch_size, self.num_heads, self.num_q_blocks, self.exact_words) + + +@dataclass(frozen=True) +class SolPredictorPlanKey: + """Cache key containing only static kernel specialization state.""" + + geometry: SolPredictorGeometry + device_index: int + dtype: torch.dtype + + +@dataclass(frozen=True) +class SolPredictorOutputs: + """Live predictor tensors consumed by block-sparse attention.""" + + exact_block_bits: torch.Tensor + k_summary: torch.Tensor + v_summary: torch.Tensor + + +@dataclass(frozen=True) +class SolPredictorPlan: + """One published shape specialization and its stable live storage.""" + + key: SolPredictorPlanKey + outputs: SolPredictorOutputs + k_mean: torch.Tensor + k_var_diag: torch.Tensor + q_centroid: torch.Tensor + + +class SOLSparsePredictor: + """Cache of shape-specialized, graph-stable SOL predictor plans.""" + + def __init__(self) -> None: + self._plans: dict[SolPredictorPlanKey, SolPredictorPlan] = {} + + @property + def num_plans(self) -> int: + return len(self._plans) + + @staticmethod + def support_reason(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> str | None: + """Return why compact two-stage SOL cannot serve these tensors.""" + + if not all(isinstance(tensor, torch.Tensor) for tensor in (q, k, v)): + return "q, k, and v must be torch tensors" + if q.ndim != 4: + return f"q must use compact BSHD layout; got rank {q.ndim}" + if k.shape != q.shape or v.shape != q.shape: + return "SOL predictor requires uniform self-attention q/k/v shapes" + if q.dtype != torch.bfloat16 or k.dtype != q.dtype or v.dtype != q.dtype: + return "SOL predictor requires matching BF16 q/k/v" + if not q.is_cuda or not k.is_cuda or not v.is_cuda: + return "SOL predictor requires CUDA q/k/v" + if k.device != q.device or v.device != q.device: + return "SOL predictor requires q/k/v on one CUDA device" + if not q.is_contiguous() or not k.is_contiguous() or not v.is_contiguous(): + return "SOL predictor requires contiguous BSHD q/k/v" + if q.shape[-1] != HEAD_DIM: + return f"SOL predictor requires head_dim={HEAD_DIM}; got {q.shape[-1]}" + if q.shape[0] <= 0 or q.shape[1] <= 0 or q.shape[2] <= 0: + return "SOL predictor requires positive B, S, and H" + return None + + @classmethod + def _key_from_inputs( + cls, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor + ) -> SolPredictorPlanKey: + reason = cls.support_reason(q, k, v) + if reason is not None: + raise ValueError(reason) + device_index = q.device.index + if device_index is None: + device_index = torch.cuda.current_device() + geometry = SolPredictorGeometry.create( + batch_size=q.shape[0], + seq_len=q.shape[1], + num_heads=q.shape[2], + head_dim=q.shape[3], + ) + return SolPredictorPlanKey( + geometry=geometry, + device_index=device_index, + dtype=q.dtype, + ) + + @staticmethod + def _launch( + plan: SolPredictorPlan, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + tau: float, + sm_scale: float, + ) -> None: + torch.ops.trtllm.visual_gen_sol_predictor( + q, + k, + v, + plan.outputs.exact_block_bits, + plan.outputs.k_summary, + plan.outputs.v_summary, + plan.k_mean, + plan.k_var_diag, + plan.q_centroid, + BLOCK_SIZE, + tau, + sm_scale, + ) + + @torch.compiler.disable + def prepare(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> SolPredictorPlan: + """Allocate one geometry and warm its kernels outside compiled or captured regions. + + The host-only boundary preserves per-instance plan ownership and keeps + kernel compilation and allocation out of Dynamo and CUDA Graph capture. + It requires the VisualGen default ``torch.compile(fullgraph=False)``. + """ + + key = self._key_from_inputs(q, k, v) + existing = self._plans.get(key) + if existing is not None: + return existing + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("SOL predictor plan must be prepared before CUDA graph capture") + + geometry = key.geometry + with torch.cuda.device(key.device_index): + k_summary = torch.empty(geometry.summary_shape, dtype=key.dtype, device=q.device) + v_summary = torch.empty_like(k_summary) + exact_block_bits = torch.empty( + geometry.exact_block_bits_shape, dtype=torch.uint32, device=q.device + ) + k_mean = torch.empty(geometry.stats_shape, dtype=torch.float32, device=q.device) + k_var_diag = torch.empty_like(k_mean) + q_centroid = torch.empty(geometry.summary_shape, dtype=torch.float32, device=q.device) + plan = SolPredictorPlan( + key=key, + outputs=SolPredictorOutputs( + exact_block_bits=exact_block_bits, + k_summary=k_summary, + v_summary=v_summary, + ), + k_mean=k_mean, + k_var_diag=k_var_diag, + q_centroid=q_centroid, + ) + # The warm-up launch compiles every kernel specialization of this geometry. + self._launch(plan, q, k, v, tau=0.0, sm_scale=1.0) + self._plans[key] = plan + return plan + + def predict( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + tau: object, + sm_scale: object, + ) -> SolPredictorOutputs: + """Update and return graph-stable SOL routes and proxy summaries.""" + + effective_tau, effective_sm_scale = _normalize_runtime_scalars(tau=tau, sm_scale=sm_scale) + plan = self.prepare(q, k, v) + self._launch(plan, q, k, v, tau=effective_tau, sm_scale=effective_sm_scale) + return plan.outputs + + +__all__ = [ + "SOLSparsePredictor", + "SolPredictorGeometry", + "SolPredictorOutputs", + "SolPredictorPlan", + "SolPredictorPlanKey", +] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py b/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py index dde2e14094c3..ff1af1e97d8a 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py @@ -271,6 +271,24 @@ def _concat_qkv( qkv = torch.cat([q, k, v], dim=-1) return qkv + @torch.compile + def _compact_qkv( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + batch_size: int, + seq_len: int, + kv_seq_len: int, + ): + # Separate Q, K, V stay separate - compact each into a contiguous token-major matrix. + # Slices of a fused QKV projection are strided; the compiled copy keeps them on a + # vectorized kernel, while already contiguous inputs pass through without a copy. + q = q.reshape(batch_size * seq_len, -1).contiguous() + k = k.reshape(batch_size * kv_seq_len, -1).contiguous() + v = v.reshape(batch_size * kv_seq_len, -1).contiguous() + return q, k, v + def forward( self, q: torch.Tensor, @@ -338,9 +356,7 @@ def forward( prepared_metadata = self._prepare_metadata(batch_size, seq_len) sage_kwargs = {} if use_separate_qkv: - q = q.reshape(batch_size * seq_len, -1).contiguous() - k = k.reshape(batch_size * kv_seq_len, -1).contiguous() - v = v.reshape(batch_size * kv_seq_len, -1).contiguous() + q, k, v = self._compact_qkv(q, k, v, batch_size, seq_len, kv_seq_len) quant_cfg = self.quant_attention_config if quant_cfg is not None: sage_kwargs = { diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py index 49a30a41a038..8eaea491890e 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py @@ -132,6 +132,10 @@ def create_attention( sparse_attention_config is not None and getattr(sparse_attention_config, "algorithm", None) == "vsa" ) + is_sol = ( + sparse_attention_config is not None + and getattr(sparse_attention_config, "algorithm", None) == "sol_attn" + ) backend_name = backend.upper() if is_vsa and backend_name == "CUTEDSL": @@ -142,6 +146,10 @@ def create_attention( from .sparse.vsa.backend import VSATrtllmAttention attn_cls = VSATrtllmAttention + elif is_sol and backend_name == "TRTLLM": + from .sparse.sol.backend import SOLTrtllmAttention + + attn_cls = SOLTrtllmAttention else: attn_cls = get_visual_gen_attention_backend(backend) diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py index d7223d91b413..ede5329319f7 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py @@ -1106,6 +1106,7 @@ def _setup_cuda_graphs(self): self._current_lora_cuda_graph_state, lambda: self.transformer.active_topology, ) + self.transformer.register_cuda_graph_extra_key_fns(runner) compile_note = " (with torch.compile)" if self.pipeline_config.torch_compile.enable else "" logger.info( "CUDA graph runner: wrapping LTX-2 two-stage transformer.forward " diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py index 7cba785b1bf2..84338be4cc32 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py @@ -36,6 +36,7 @@ UlyssesAttention, wrap_parallel_attention, ) +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.params import SolParams from tensorrt_llm._torch.visual_gen.attention_backend.utils import create_attention from tensorrt_llm._torch.visual_gen.models.modeling import BaseDiffusionModel from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode @@ -928,6 +929,8 @@ def forward( text_kv_video: tuple[torch.Tensor, torch.Tensor] | None = None, text_kv_audio: tuple[torch.Tensor, torch.Tensor] | None = None, step_index=None, + video_sol_timestep: torch.Tensor | None = None, + audio_sol_timestep: torch.Tensor | None = None, ) -> tuple[TransformerArgs | None, TransformerArgs | None]: """Forward with optional perturbation masking for STG. @@ -938,6 +941,8 @@ def forward( Required when the video stream runs cross-attn — built by ``LTXModel.prepare_text_cache``. text_kv_audio: Pre-projected (K, V) for audio text cross-attention. + video_sol_timestep: Raw video modality timestep for SOL phase preparation. + audio_sol_timestep: Raw audio modality timestep for SOL phase preparation. """ if video is None and audio is None: raise ValueError("At least one of video or audio must be provided") @@ -991,8 +996,14 @@ def forward( self._fuse_adaln, fp4_input_scale=get_nvfp4_self_attn_input_scale(self.attn1), ) + v_attn_timestep = ( + video_sol_timestep + if isinstance(self.attn1.sparse_params, SolParams) + and video_sol_timestep is not None + else video.timesteps + ) v_attn_raw = self.attn1( - norm_vx, pe=video.positional_embeddings, timestep=video.timesteps + norm_vx, pe=video.positional_embeddings, timestep=v_attn_timestep ) if has_perturbations and perturbations.any_in_batch( PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx @@ -1045,11 +1056,17 @@ def forward( self._fuse_adaln, fp4_input_scale=get_nvfp4_self_attn_input_scale(self.audio_attn1), ) + a_attn_timestep = ( + audio_sol_timestep + if isinstance(self.audio_attn1.sparse_params, SolParams) + and audio_sol_timestep is not None + else audio.timesteps + ) a_attn_raw = self.audio_attn1( norm_ax, pe=audio.positional_embeddings, key_padding_mask=audio.audio_padding_mask, - timestep=audio.timesteps, + timestep=a_attn_timestep, ) if has_perturbations and perturbations.any_in_batch( PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx @@ -2255,6 +2272,9 @@ def forward( if not self.model_type.is_audio_enabled() and audio is not None: raise ValueError("Audio is not enabled for this model") + video_sol_timestep = video.timesteps if video is not None else None + audio_sol_timestep = audio.timesteps if audio is not None else None + # Audio padding for Ulysses: when self._audio_pad > 0 (set once by # configure_audio_ulysses to make T_a divisible by ulysses_size), pad # audio on entry to make it shardable. Build a [B, T_a_padded] bool mask @@ -2333,6 +2353,8 @@ def forward( ax, perturbations=perturbations, step_index=step_index, + video_sol_timestep=video_sol_timestep, + audio_sol_timestep=audio_sol_timestep, ) if video_args is not None and vx is not None: video_args = replace(video_args, x=vx) @@ -2347,6 +2369,8 @@ def forward( text_kv_video=v_kv[i] if v_kv else None, text_kv_audio=a_kv[i] if a_kv else None, step_index=step_index, + video_sol_timestep=video_sol_timestep, + audio_sol_timestep=audio_sol_timestep, ) # Gather sequences back to full length for output processing. diff --git a/tensorrt_llm/_torch/visual_gen/models/modeling.py b/tensorrt_llm/_torch/visual_gen/models/modeling.py index 9fd8a02187ea..30ea64ec7a08 100644 --- a/tensorrt_llm/_torch/visual_gen/models/modeling.py +++ b/tensorrt_llm/_torch/visual_gen/models/modeling.py @@ -20,8 +20,9 @@ import torch.nn as nn from tensorrt_llm._torch.attention.backends.sparse.skip_softmax import SkipSoftmaxScheduler +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.params import SolParams from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig -from tensorrt_llm.visual_gen.sparse_attention import SkipSoftmaxAttentionConfig +from tensorrt_llm.visual_gen.sparse_attention import SkipSoftmaxAttentionConfig, SolAttentionConfig if TYPE_CHECKING: from tensorrt_llm._torch.visual_gen.cuda_graph_runner import CUDAGraphRunner @@ -74,6 +75,38 @@ def register_cuda_graph_extra_key_fns(self, runner: "CUDAGraphRunner") -> None: the shared registrations. """ sparse_config = self.model_config.attention.sparse_attention_config + if isinstance(sparse_config, SolAttentionConfig): + disabled_until_timestep = sparse_config.disabled_until_timestep + if disabled_until_timestep is None: + return + + def sol_phase_key(*args, **kwargs): + del args + modality_phases = tuple( + (name, phase) + for name in ("video", "audio") + if (modality := kwargs.get(name)) is not None + and hasattr(modality, "timesteps") + and ( + phase := SolParams.get_graph_phase_for_timestep( + modality.timesteps, + disabled_until_timestep=disabled_until_timestep, + ) + ) + is not None + ) + if modality_phases: + return modality_phases + return SolParams.get_graph_phase_for_timestep( + kwargs.get("timestep"), + disabled_until_timestep=disabled_until_timestep, + ) + + # SOL backend warmup prepares the matching phase before capture. Key + # each modality's live timestep so dense and sparse graphs stay separate. + runner.register_extra_key_fn("sol_attn_phase", sol_phase_key) + return + if not isinstance(sparse_config, SkipSoftmaxAttentionConfig): return diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/modules/attention.py index 8ffd4790728c..b77bea91af3a 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py @@ -20,7 +20,7 @@ import torch.nn as nn from tensorrt_llm.logger import logger -from tensorrt_llm.visual_gen.sparse_attention import SkipSoftmaxAttentionConfig +from tensorrt_llm.visual_gen.sparse_attention import SkipSoftmaxAttentionConfig, SolAttentionConfig from ...modules.linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig from ...utils import Fp4QuantizedTensor @@ -113,18 +113,26 @@ def __init__( base_backend = config.attention.backend _sa_cfg = config.attention.sparse_attention_config is_vsa = _sa_cfg is not None and getattr(_sa_cfg, "algorithm", None) == "vsa" + is_sol = _sa_cfg is not None and getattr(_sa_cfg, "algorithm", None) == "sol_attn" + is_sparse_self_attention = is_vsa or is_sol is_separate_qkv = self.qkv_mode == QKVMode.SEPARATE_QKV is_separate_qkv_cross_attention = is_separate_qkv and not separate_qkv_is_self_attention use_vanilla_cross_attention = is_separate_qkv and ( - (base_backend == "TRTLLM" and not is_vsa) - or (is_vsa and not separate_qkv_is_self_attention) + (base_backend == "TRTLLM" and not is_sparse_self_attention) + or (is_sparse_self_attention and not separate_qkv_is_self_attention) ) - # Cross-attention fallback: dense TRTLLM and every VSA backend are self-attn only. + # Cross-attention fallback: dense TRTLLM, VSA, and SOL are self-attn only. if use_vanilla_cross_attention: backend_name = "VANILLA" - requested = f"{base_backend} (VSA)" if is_vsa else base_backend + requested = ( + f"{base_backend} (VSA)" + if is_vsa + else f"{base_backend} (SOL)" + if is_sol + else base_backend + ) # Warn once per (module class, requested, resolved) triple so the # fallback is visible without per-module-instance log spam. logger.warning_once( @@ -141,6 +149,12 @@ def __init__( f"with context parallelism (Attention2D/Ring, cp_size={cp_size}). Use " f"ulysses or cfg parallelism instead." ) + if is_sol and cp_size > 1: + raise ValueError( + f"SOL needs the full token sequence per rank, so it is incompatible " + f"with context parallelism (Attention2D/Ring, cp_size={cp_size}). Use " + f"ulysses or cfg parallelism instead." + ) self.attn_backend = backend_name self.qk_norm = qk_norm self.qk_norm_mode = qk_norm_mode @@ -250,6 +264,8 @@ def __init__( module_name=self.module_name, pretrained_config=config.pretrained_config, ) + elif isinstance(ss_cfg, SolAttentionConfig) and backend_name == "TRTLLM": + sparse_params = ss_cfg.to_sparse_params() self.sparse_params = sparse_params # Create compute backend diff --git a/tensorrt_llm/visual_gen/__init__.py b/tensorrt_llm/visual_gen/__init__.py index ec714e3d2bfe..213a8e2e47b4 100644 --- a/tensorrt_llm/visual_gen/__init__.py +++ b/tensorrt_llm/visual_gen/__init__.py @@ -49,6 +49,7 @@ QuantAttentionConfig, RuntimeLoRAConfig, SkipSoftmaxAttentionConfig, + SolAttentionConfig, SparseAttentionConfig, TeaCacheConfig, TorchCompileConfig, @@ -76,6 +77,7 @@ "QuantAttentionConfig": "tensorrt_llm.visual_gen.args", "RuntimeLoRAConfig": "tensorrt_llm.visual_gen.args", "SkipSoftmaxAttentionConfig": "tensorrt_llm.visual_gen.args", + "SolAttentionConfig": "tensorrt_llm.visual_gen.args", "SparseAttentionConfig": "tensorrt_llm.visual_gen.args", "TeaCacheConfig": "tensorrt_llm.visual_gen.args", "TorchCompileConfig": "tensorrt_llm.visual_gen.args", @@ -133,6 +135,7 @@ def __dir__(): "RuntimeLoRAConfig", "SparseAttentionConfig", "SkipSoftmaxAttentionConfig", + "SolAttentionConfig", "VideoSparseAttentionConfig", "VAEConfig", "CacheConfig", diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index b3390317af52..e4f3c3e9982d 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -31,7 +31,11 @@ from tensorrt_llm.llmapi.utils import StrictBaseModel, set_api_status from tensorrt_llm.models.modeling_utils import QuantConfig -from .sparse_attention import SkipSoftmaxAttentionConfig, VideoSparseAttentionConfig +from .sparse_attention import ( + SkipSoftmaxAttentionConfig, + SolAttentionConfig, + VideoSparseAttentionConfig, +) # ============================================================================= # Type aliases @@ -95,7 +99,7 @@ class QuantAttentionConfig(StrictBaseModel): # Discriminated union of sparse attention configs. SparseAttentionConfig = Annotated[ - Union[SkipSoftmaxAttentionConfig, VideoSparseAttentionConfig], + Union[SkipSoftmaxAttentionConfig, VideoSparseAttentionConfig, SolAttentionConfig], Field(discriminator="algorithm"), ] @@ -122,7 +126,8 @@ class AttentionConfig(StrictBaseModel): status="prototype", description=( "Sparse attention recipe. Discriminated by algorithm: " - "skip_softmax (TRTLLM / CUTEDSL backends) or VSA (CUTEDSL / TRTLLM backends)." + "skip_softmax (TRTLLM / CUTEDSL backends), VSA (CUTEDSL / TRTLLM backends), " + "or SOL (TRTLLM backend)." ), ) @@ -206,6 +211,7 @@ def _validate_sparse_attention_config(self) -> "AttentionConfig": supported_backends = { "skip_softmax": ("TRTLLM", "CUTEDSL"), "vsa": ("CUTEDSL", "TRTLLM"), + "sol_attn": ("TRTLLM",), }.get(algo) if supported_backends is None: return self @@ -228,6 +234,8 @@ def _validate_quant_sparse_mutex(self) -> "AttentionConfig": # VSA consumes the unquantized Q/K/V path, so accepting an attention # quantization recipe would silently ignore user configuration. raise ValueError("VSA and quant_attention_config are mutually exclusive.") + if self.sparse_attention_config.algorithm == "sol_attn": + raise ValueError("SOL and quant_attention_config are mutually exclusive.") return self @@ -766,6 +774,21 @@ def _normalize_quant_config(cls, data: Any) -> Any: data = {**data, "quant_config": QuantConfig()} return data + @model_validator(mode="after") + def _validate_sol_fullgraph(self) -> "VisualGenArgs": + sparse_config = self.attention_config.sparse_attention_config + if ( + isinstance(sparse_config, SolAttentionConfig) + and self.torch_compile_config.enable + and self.torch_compile_config.enable_fullgraph + ): + raise ValueError( + "SOL sparse attention does not support torch.compile fullgraph; " + "set torch_compile_config.enable_fullgraph=False or disable " + "torch.compile." + ) + return self + @property def cache_backend(self) -> Optional[CacheBackendName]: return self.cache_config.cache_backend if self.cache_config is not None else None # type: ignore[return-value] @@ -817,6 +840,7 @@ def from_yaml(cls, yaml_path: Union[str, Path], **overrides: Any) -> "VisualGenA "QuantAttentionConfig", "SparseAttentionConfig", "SkipSoftmaxAttentionConfig", + "SolAttentionConfig", "VideoSparseAttentionConfig", "AttentionConfig", "VAEConfig", diff --git a/tensorrt_llm/visual_gen/sparse_attention.py b/tensorrt_llm/visual_gen/sparse_attention.py index 21189c249d55..a06c8d9c70e7 100644 --- a/tensorrt_llm/visual_gen/sparse_attention.py +++ b/tensorrt_llm/visual_gen/sparse_attention.py @@ -19,10 +19,33 @@ from typing import Any, Dict, Literal, Optional from pydantic import Field as PydanticField +from pydantic import field_validator from tensorrt_llm.llmapi.utils import StrictBaseModel +def _parse_dense_layers(spec: Optional[str]) -> frozenset[int]: + """Parse comma-separated layer indices and inclusive ranges.""" + layers: set[int] = set() + for raw_item in (spec or "").split(","): + item = raw_item.strip() + if not item: + if spec and spec.strip(): + raise ValueError("dense_layers must not contain empty entries") + continue + bounds = item.split("-") + if len(bounds) not in (1, 2) or not all(bounds): + raise ValueError(f"invalid dense_layers entry: {item!r}") + try: + start, end = int(bounds[0]), int(bounds[-1]) + except ValueError as error: + raise ValueError(f"invalid dense_layers entry: {item!r}") from error + if start < 0 or end < start: + raise ValueError(f"invalid dense_layers range: {item!r}") + layers.update(range(start, end + 1)) + return frozenset(layers) + + class BaseSparseAttentionConfig(StrictBaseModel): """Base for visual-generation sparse attention configs. @@ -221,6 +244,58 @@ def _ckpt_sparse_attention_config_from_kwargs( return None +class SolAttentionConfig(BaseSparseAttentionConfig): + """Two-stage SOL sparse attention backed by TRTLLM PrimTS. + + TRTLLM predicts an exact block mask from compact Q/K/V tensors, then the + generic block-sparse FMHA executes that mask. Unsupported runtime tensor + envelopes raise instead of silently falling back to dense attention. + """ + + algorithm: Literal["sol_attn"] = PydanticField( + "sol_attn", + description="Sparse attention algorithm discriminator.", + ) + tau: float = PydanticField( + 1.0, + allow_inf_nan=False, + description="Standard-deviation multiplier used by the SOL block selector.", + ) + disabled_until_timestep: Optional[float] = PydanticField( + None, + gt=0.0, + le=1.0, + description=( + "Dense-prefix cutoff on the normalized denoising timestep. The " + "attention switches to SOL when timestep is below this value." + ), + ) + dense_layers: Optional[str] = PydanticField( + None, + description=( + "Comma-separated layer indices and inclusive ranges (for example, " + "'0,2-4') that always use dense attention." + ), + ) + + @field_validator("dense_layers") + @classmethod + def _validate_dense_layers(cls, value: Optional[str]) -> Optional[str]: + _parse_dense_layers(value) + return value + + def to_sparse_params(self, **kwargs): + """Lower the public recipe into immutable SOL runtime parameters.""" + del kwargs + from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.params import SolParams + + return SolParams( + tau=self.tau, + disabled_until_timestep=self.disabled_until_timestep, + dense_layers=_parse_dense_layers(self.dense_layers), + ) + + class VideoSparseAttentionConfig(StrictBaseModel): """Video Sparse Attention (VSA) sparse-attention recipe. diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 6204cc3e512f..fe4de3e39fd5 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -282,6 +282,9 @@ l0_b200: - unittest/_torch/visual_gen/test_fa4_cutlass_compatibility.py - unittest/_torch/visual_gen/test_attention_vsa.py - unittest/_torch/visual_gen/test_attention_flashinfer.py + - unittest/_torch/visual_gen/sparse_attention/test_sol_predictor.py + - unittest/_torch/visual_gen/sparse_attention/test_sol_attention.py::test_real_b200_sol_backend_cuda_graph_matches_dense_reference + - unittest/_torch/visual_gen/sparse_attention/test_sol_attention.py::test_real_b200_sol_backend_mixed_proxy_cuda_graph_matches_reference - unittest/_torch/visual_gen/test_attention_trtllm_sage.py - unittest/_torch/visual_gen/test_attention_integration.py - unittest/_torch/visual_gen/test_trtllm_attention_metadata.py diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 0327b4e498f8..11345983f054 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -59,6 +59,8 @@ l0_cpu: - unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py - unittest/_torch/visual_gen/multi_gpu/test_visual_gen_multinode.py - unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py + - unittest/_torch/visual_gen/sparse_attention/test_sol_attention.py + - unittest/_torch/visual_gen/sparse_attention/test_sol_predictor.py - unittest/_torch/visual_gen/test_attention_flashinfer.py::test_flashinfer_backend_is_registered - unittest/_torch/visual_gen/test_attention_integration.py - unittest/_torch/visual_gen/test_cache_dit.py diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/test_sol_attention.py b/tests/unittest/_torch/visual_gen/sparse_attention/test_sol_attention.py new file mode 100644 index 000000000000..0a4a1f1a521e --- /dev/null +++ b/tests/unittest/_torch/visual_gen/sparse_attention/test_sol_attention.py @@ -0,0 +1,898 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the VisualGen SOL TRTLLM backend.""" + +from __future__ import annotations + +import math +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from pydantic import ValidationError + +from tensorrt_llm._torch.attention.backends.fmha.prims_ts_block_sparse import PrimsTSBlockSparseFmha +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + PredefinedAttentionMask, +) +from tensorrt_llm._torch.attention.backends.sparse.hooks import prepare_sparse_runtime_params +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention as CoreTrtllmAttention +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol import backend as sol_backend +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.backend import SOLTrtllmAttention +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.params import SolParams +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.predictor import ( + SolPredictorOutputs, + SOLSparsePredictor, +) +from tensorrt_llm._torch.visual_gen.attention_backend.trtllm import TrtllmAttention +from tensorrt_llm._torch.visual_gen.attention_backend.utils import create_attention +from tensorrt_llm._torch.visual_gen.config import ( + DiffusionModelConfig, + create_attention_metadata_state, +) +from tensorrt_llm._torch.visual_gen.cuda_graph_runner import CUDAGraphRunner, CUDAGraphRunnerConfig +from tensorrt_llm._torch.visual_gen.models.modeling import BaseDiffusionModel +from tensorrt_llm._torch.visual_gen.modules import attention as attention_module +from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode +from tensorrt_llm.visual_gen import SolAttentionConfig +from tensorrt_llm.visual_gen.args import AttentionConfig, QuantAttentionConfig + +_REQUIRES_SM100 = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() not in ((10, 0), (10, 3)), + reason="SOL requires SM100 or SM103", +) +_CPU_ONLY = pytest.mark.cpu_only + + +def _make_backend( + params: SolParams, + predictor: Mock, + *, + layer_idx: int = 1, +) -> SOLTrtllmAttention: + backend = object.__new__(SOLTrtllmAttention) + backend.layer_idx = layer_idx + backend.num_heads = 2 + backend.num_kv_heads = 2 + backend.head_dim = 128 + backend.q_scaling = 1.0 + backend.quant_attention_config = None + backend.sparse_params = None + backend._fmha_manager = SimpleNamespace(fmha_libs=[object.__new__(PrimsTSBlockSparseFmha)]) + backend.sol_params = params + backend._prepared_graph_phase = None + backend.predictor = predictor + return backend + + +def _flatten(tensor: torch.Tensor | None) -> torch.Tensor | None: + """Convert a BSHD tensor into the flattened ``[B*S, H*D]`` core layout.""" + + if tensor is None: + return None + return tensor.reshape(tensor.shape[0] * tensor.shape[1], -1) + + +def _core_metadata(batch_size: int, seq_len: int) -> SimpleNamespace: + return SimpleNamespace(num_seqs=batch_size, max_seq_len=seq_len) + + +def _predict( + backend: SOLTrtllmAttention, + q: torch.Tensor, + k: torch.Tensor | None, + v: torch.Tensor | None, + *, + attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL, + timestep: object = None, +): + """Invoke the core prediction hook the way the core forward does.""" + + return backend.block_sparse_attn_predict( + _flatten(q), + _flatten(k), + _flatten(v), + _core_metadata(q.shape[0], q.shape[1]), + AttentionForwardArgs(attention_mask=attention_mask, timestep=timestep), + ) + + +def _forward( + backend: SOLTrtllmAttention, + q: torch.Tensor, + k: torch.Tensor | None, + v: torch.Tensor | None, + **kwargs, +) -> torch.Tensor: + seq_len_kv = kwargs.pop("seq_len_kv", q.shape[1]) + return backend.forward( + q=q, + k=k, + v=v, + batch_size=q.shape[0], + seq_len=q.shape[1], + seq_len_kv=seq_len_kv, + **kwargs, + ) + + +def _stub_core_forward(monkeypatch) -> dict: + """Replace metadata preparation and the core forward with a recorder that + still runs the backend's sparse prediction.""" + + captured = {} + monkeypatch.setattr( + TrtllmAttention, + "_prepare_metadata", + lambda self, batch_size, seq_len: _core_metadata(batch_size, seq_len), + ) + + def _core_forward(self, q, k, v, metadata, forward_args=None, **kwargs): + forward_args.sparse_runtime_params = prepare_sparse_runtime_params( + self, q, k, v, metadata, forward_args + ) + captured.update(q=q, k=k, v=v, metadata=metadata, forward_args=forward_args) + return q + + monkeypatch.setattr(CoreTrtllmAttention, "forward", _core_forward) + return captured + + +def _predictor_outputs(*, batch_size: int, seq_len: int, num_heads: int) -> SolPredictorOutputs: + num_blocks = (seq_len + 63) // 64 + return SolPredictorOutputs( + exact_block_bits=torch.zeros( + batch_size, + num_heads, + num_blocks, + (num_blocks + 31) // 32, + dtype=torch.uint32, + ), + k_summary=torch.zeros(batch_size, num_blocks, num_heads, 128, dtype=torch.bfloat16), + v_summary=torch.zeros(batch_size, num_blocks, num_heads, 128, dtype=torch.bfloat16), + ) + + +def _bshd(seq_len: int = 64, num_heads: int = 2) -> torch.Tensor: + return torch.zeros(1, seq_len, num_heads, 128, dtype=torch.bfloat16) + + +def _stub_backend( + params: SolParams | None = None, + *, + seq_len: int = 64, + unsupported_reason: str | None = None, +) -> tuple[SOLTrtllmAttention, Mock]: + predictor = Mock(spec=SOLSparsePredictor) + predictor.support_reason.return_value = unsupported_reason + predictor.predict.return_value = _predictor_outputs(batch_size=1, seq_len=seq_len, num_heads=2) + return _make_backend(params or SolParams(tau=1.0), predictor), predictor + + +def _dense_reference(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + scores = torch.einsum("bqhd,bkhd->bhqk", q.float(), k.float()) * (128**-0.5) + return torch.einsum("bhqk,bkhd->bqhd", scores.softmax(dim=-1), v.float()).to(q.dtype) + + +def _mixed_proxy_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + outputs: SolPredictorOutputs, +) -> torch.Tensor: + """Evaluate the exact-token/proxy-summary attention contract.""" + + block_size = 64 + num_blocks = math.ceil(k.shape[1] / block_size) + exact_words = outputs.exact_block_bits.detach().cpu().to(torch.int64) + reference = torch.empty_like(q) + scale = q.shape[-1] ** -0.5 + for batch_idx in range(q.shape[0]): + for head_idx in range(q.shape[2]): + for q_block_idx in range(math.ceil(q.shape[1] / block_size)): + q_begin = q_block_idx * block_size + q_end = min(q_begin + block_size, q.shape[1]) + exact_blocks = [ + block_idx + for block_idx in range(num_blocks) + if int(exact_words[batch_idx, head_idx, q_block_idx, block_idx // 32]) + & (1 << (block_idx % 32)) + ] + proxy_blocks = [ + block_idx for block_idx in range(num_blocks) if block_idx not in exact_blocks + ] + exact_tokens = torch.cat( + [ + torch.arange( + block_idx * block_size, + min((block_idx + 1) * block_size, k.shape[1]), + device=q.device, + ) + for block_idx in exact_blocks + ] + ) + q_rows = q[batch_idx, q_begin:q_end, head_idx].float() + exact_logits = (q_rows @ k[batch_idx, exact_tokens, head_idx].float().T) * scale + proxy_logits = ( + q_rows @ outputs.k_summary[batch_idx, proxy_blocks, head_idx].float().T + ) * scale + logits = torch.cat((exact_logits, proxy_logits), dim=1) + weights = torch.exp(logits - logits.amax(dim=1, keepdim=True)) + exact_weights = weights[:, : exact_tokens.numel()] + proxy_weights = weights[:, exact_tokens.numel() :] + 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_size, k.shape[1] - block_idx * block_size) + denominator += proxy_weights[:, proxy_offset : proxy_offset + 1] * ( + tokens_in_block + ) + reference[batch_idx, q_begin:q_end, head_idx] = (numerator / denominator).to( + q.dtype + ) + return reference + + +@_CPU_ONLY +def test_sol_params_requires_precomputed_phase_during_cuda_graph_capture(monkeypatch) -> None: + params = SolParams(tau=1.0, disabled_until_timestep=0.6) + timestep = torch.tensor(0.2) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + + with pytest.raises(RuntimeError, match="precomputed"): + params.should_use_sparse(layer_idx=1, timestep=timestep) + + assert params.should_use_sparse(layer_idx=1, timestep=timestep, graph_phase=1) + assert not params.should_use_sparse(layer_idx=1, timestep=timestep, graph_phase=0) + + +@_CPU_ONLY +def test_sol_backend_warmup_prepares_phase_for_capture(monkeypatch) -> None: + params = SolParams(tau=1.0, disabled_until_timestep=0.6) + q = _bshd() + backend, predictor = _stub_backend(params) + + assert _predict(backend, q, q, q, timestep=torch.tensor(0.8)) is None + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + assert _predict(backend, q, q, q, timestep=torch.tensor(0.8)) is None + predictor.predict.assert_not_called() + + +@_CPU_ONLY +def test_sol_backend_rejects_cutoff_capture_without_warmup(monkeypatch) -> None: + params = SolParams(tau=1.0, disabled_until_timestep=0.6) + q = _bshd() + backend, predictor = _stub_backend(params) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + + with pytest.raises(RuntimeError, match="prepared before CUDA Graph capture"): + _predict(backend, q, q, q, timestep=torch.tensor(0.2)) + + predictor.predict.assert_not_called() + + +@_CPU_ONLY +def test_sol_backend_cutoff_requires_timestep_during_eager_forward() -> None: + params = SolParams(tau=1.0, disabled_until_timestep=0.6) + q = _bshd() + backend, predictor = _stub_backend(params) + + with pytest.raises(ValueError, match="timestep is required"): + _predict(backend, q, q, q) + + predictor.predict.assert_not_called() + + +@_CPU_ONLY +def test_sol_phase_waits_until_all_token_timesteps_are_below_cutoff() -> None: + params = SolParams(tau=1.0, disabled_until_timestep=0.6) + + assert not params.should_use_sparse(layer_idx=1, timestep=torch.tensor([0.0, 0.8])) + assert params.should_use_sparse(layer_idx=1, timestep=torch.tensor([0.0, 0.2])) + + +@_CPU_ONLY +def test_sol_config_lowers_and_factory_initializes_backend(monkeypatch) -> None: + base_kwargs = {} + + def _base_init(self, **kwargs) -> None: + base_kwargs.update(kwargs) + self.layer_idx = kwargs["layer_idx"] + self.head_dim = kwargs["head_dim"] + self.q_scaling = 1.0 + + monkeypatch.setattr(TrtllmAttention, "__init__", _base_init) + attention_config = AttentionConfig( + backend="TRTLLM", + sparse_attention_config={ + "algorithm": "sol_attn", + "tau": -0.25, + "disabled_until_timestep": 0.6, + "dense_layers": "0,2-4,3", + }, + ) + config = attention_config.sparse_attention_config + assert isinstance(config, SolAttentionConfig) + params = config.to_sparse_params() + + backend = create_attention( + backend="TRTLLM", + layer_idx=3, + num_heads=4, + head_dim=128, + attention_config=attention_config, + sparse_params=params, + attention_metadata_state=create_attention_metadata_state(), + ) + + assert config.algorithm == "sol_attn" + assert params.tau == -0.25 + assert params.disabled_until_timestep == 0.6 + assert params.dense_layers == frozenset({0, 2, 3, 4}) + assert isinstance(backend, SOLTrtllmAttention) + assert backend.sol_params is params + assert isinstance(backend.predictor, SOLSparsePredictor) + assert base_kwargs["sparse_params"] is None + assert "_enable_sparse_workflow" not in SOLTrtllmAttention.__dict__ + assert "_should_use_sparse_workflow" not in SOLTrtllmAttention.__dict__ + assert not backend.support_fused_qkv() + assert "forward" not in SOLTrtllmAttention.__dict__ + assert "block_sparse_attn_predict" in SOLTrtllmAttention.__dict__ + + +@_CPU_ONLY +def test_sol_layers_share_one_model_scoped_predictor(monkeypatch) -> None: + monkeypatch.setattr(TrtllmAttention, "__init__", lambda self, **kwargs: None) + attention_metadata_state = create_attention_metadata_state() + + first = SOLTrtllmAttention( + sparse_params=SolParams(), + attention_metadata_state=attention_metadata_state, + ) + second = SOLTrtllmAttention( + sparse_params=SolParams(), + attention_metadata_state=attention_metadata_state, + ) + + assert first.predictor is second.predictor + assert attention_metadata_state["sparse_predictors"]["sol_attn"] is first.predictor + + +@_CPU_ONLY +def test_sol_backend_sparse_phase_emits_proxy_bitmask_carrier(monkeypatch) -> None: + batch_size, seq_len, num_heads = 1, 65, 2 + q, k, v = (_bshd(seq_len, num_heads) for _ in range(3)) + predictor_outputs = _predictor_outputs( + batch_size=batch_size, + seq_len=seq_len, + num_heads=num_heads, + ) + backend, predictor = _stub_backend(SolParams(tau=0.75), seq_len=seq_len) + predictor.predict.return_value = predictor_outputs + monkeypatch.setattr(sol_backend, "get_bmm1_scale", lambda attn: 0.375) + + carrier = _predict(backend, q, k, v, timestep=0.2) + + predicted_q, predicted_k, predicted_v = predictor.predict.call_args.args + assert predictor.predict.call_args.kwargs == {"tau": 0.75, "sm_scale": 0.375} + predictor.support_reason.assert_called_once_with(predicted_q, predicted_k, predicted_v) + for predicted, source in zip((predicted_q, predicted_k, predicted_v), (q, k, v), strict=True): + assert predicted.shape == (batch_size, seq_len, num_heads, 128) + assert predicted.is_contiguous() + assert predicted.data_ptr() == source.data_ptr() + assert ( + carrier.q_block_size, + carrier.kv_block_size, + carrier.max_blocks_per_row, + carrier.block_indptr, + carrier.block_indices, + carrier.kv_valid_bits, + ) == (64, 64, None, None, None, None) + assert carrier.exact_block_bits is predictor_outputs.exact_block_bits + assert carrier.k_summary is predictor_outputs.k_summary + assert carrier.v_summary is predictor_outputs.v_summary + assert carrier.sparse_format == "bitmask" + assert carrier.use_proxy_routes + + +@_CPU_ONLY +def test_sol_wrapper_compacts_separate_qkv_and_predicts_inside_core(monkeypatch) -> None: + batch_size, seq_len, num_heads = 1, 65, 2 + packed_qkv = torch.zeros(batch_size, seq_len, 3 * num_heads * 128, dtype=torch.bfloat16) + q, k, v = ( + tensor.view(batch_size, seq_len, num_heads, 128) + for tensor in packed_qkv.split(num_heads * 128, dim=-1) + ) + predictor_outputs = _predictor_outputs( + batch_size=batch_size, + seq_len=seq_len, + num_heads=num_heads, + ) + backend, predictor = _stub_backend(SolParams(tau=0.75), seq_len=seq_len) + predictor.predict.return_value = predictor_outputs + monkeypatch.setattr(sol_backend, "get_bmm1_scale", lambda attn: 0.375) + captured = _stub_core_forward(monkeypatch) + + output = _forward(backend, q, k, v, attention_mask=PredefinedAttentionMask.FULL, timestep=0.2) + + assert output.shape == (batch_size, seq_len, num_heads * 128) + assert all( + tensor.is_contiguous() and tensor.shape == (batch_size * seq_len, num_heads * 128) + for tensor in (captured["q"], captured["k"], captured["v"]) + ) + forward_args = captured["forward_args"] + assert forward_args.timestep == 0.2 + assert forward_args.sparse_backend_args is None + carrier = forward_args.sparse_runtime_params.block_sparse_inputs + assert carrier.exact_block_bits is predictor_outputs.exact_block_bits + predicted_q = predictor.predict.call_args.args[0] + assert predicted_q.data_ptr() == captured["q"].data_ptr() + + +@_CPU_ONLY +@pytest.mark.parametrize( + ("k", "v", "attention_mask", "message"), + ( + (None, None, PredefinedAttentionMask.FULL, "separate q, k, and v"), + (_bshd(32), _bshd(32), PredefinedAttentionMask.FULL, "self-attention"), + (_bshd(), _bshd(), PredefinedAttentionMask.CAUSAL, "full attention mask"), + ), +) +def test_sol_backend_rejects_non_sol_sparse_calls( + k: torch.Tensor | None, + v: torch.Tensor | None, + attention_mask: PredefinedAttentionMask, + message: str, +) -> None: + q = _bshd() + backend, predictor = _stub_backend() + + with pytest.raises(ValueError, match=message): + _predict(backend, q, k, v, attention_mask=attention_mask) + + predictor.predict.assert_not_called() + + +@_CPU_ONLY +def test_sol_wrapper_rejects_fused_qkv_before_core(monkeypatch) -> None: + q = _bshd() + backend, predictor = _stub_backend() + prepare_metadata = Mock(return_value=object()) + monkeypatch.setattr(TrtllmAttention, "_prepare_metadata", prepare_metadata) + + with pytest.raises(ValueError, match="separate q, k, and v"): + _forward(backend, q, None, None) + + prepare_metadata.assert_not_called() + predictor.predict.assert_not_called() + + +@_CPU_ONLY +def test_sol_backend_surfaces_predictor_support_reason_before_execution() -> None: + q = _bshd() + reason = "SOL predictor requires compact BSHD q/k/v" + backend, predictor = _stub_backend(unsupported_reason=reason) + + with pytest.raises(ValueError, match=reason): + _predict(backend, q, q, q) + + predictor.predict.assert_not_called() + + +@_CPU_ONLY +@pytest.mark.parametrize( + ("params", "layer_idx", "timestep"), + ( + (SolParams(dense_layers=frozenset({1})), 1, None), + (SolParams(disabled_until_timestep=0.6), 1, 0.8), + ), +) +def test_sol_dense_policy_returns_no_routes_without_predicting( + params: SolParams, + layer_idx: int, + timestep: float | None, +) -> None: + q = _bshd() + backend, predictor = _stub_backend(params) + backend.layer_idx = layer_idx + + assert _predict(backend, q, q, q, timestep=timestep) is None + predictor.support_reason.assert_not_called() + predictor.predict.assert_not_called() + + +@_CPU_ONLY +def test_sol_sparse_phase_without_primts_fails_closed() -> None: + q = _bshd() + backend, predictor = _stub_backend() + backend._fmha_manager = SimpleNamespace(fmha_libs=[]) + + with pytest.raises(RuntimeError, match="requires PrimTS block-sparse FMHA"): + _predict(backend, q, q, q) + + predictor.support_reason.assert_not_called() + predictor.predict.assert_not_called() + + +@_CPU_ONLY +def test_sol_sparse_phase_with_quantization_fails_closed() -> None: + q = _bshd() + backend, predictor = _stub_backend() + backend.quant_attention_config = object() + + with pytest.raises(ValueError, match="does not support quant_attention_config"): + _predict(backend, q, q, q) + + predictor.support_reason.assert_not_called() + predictor.predict.assert_not_called() + + +@_CPU_ONLY +@pytest.mark.parametrize( + "config_kwargs", + ( + {"tau": 1.0e100}, + {"disabled_until_timestep": 0.0}, + {"dense_layers": "2-1"}, + ), +) +def test_sol_public_config_rejects_invalid_policy(config_kwargs) -> None: + with pytest.raises((ValidationError, ValueError)): + SolAttentionConfig(**config_kwargs).to_sparse_params() + + +@_CPU_ONLY +def test_sol_public_config_requires_trtllm_backend() -> None: + with pytest.raises(ValidationError, match="requires backend"): + AttentionConfig( + backend="VANILLA", + sparse_attention_config=SolAttentionConfig(), + ) + + +@_CPU_ONLY +def test_sol_and_attention_quantization_are_mutually_exclusive() -> None: + with pytest.raises(ValidationError, match="SOL and quant_attention_config"): + AttentionConfig( + backend="TRTLLM", + quant_attention_config=QuantAttentionConfig( + qk_dtype="fp8", + q_block_size=1, + k_block_size=1, + v_block_size=1, + ), + sparse_attention_config=SolAttentionConfig(), + ) + + +def _sol_model_config(*, cp_size: int = 1) -> DiffusionModelConfig: + config = DiffusionModelConfig( + pretrained_config=SimpleNamespace(), + attention=AttentionConfig( + backend="TRTLLM", + sparse_attention_config=SolAttentionConfig( + tau=0.75, + disabled_until_timestep=0.6, + dense_layers="0,2-3", + ), + ), + skip_create_weights_in_init=True, + attention_metadata_state=create_attention_metadata_state(), + ) + if cp_size > 1: + config.visual_gen_mapping = SimpleNamespace( + ring_size=cp_size, + ring_group=None, + ulysses_size=1, + ulysses_group=None, + attn2d_row_size=1, + attn2d_col_size=1, + attn2d_row_group=None, + attn2d_col_group=None, + cp_size=cp_size, + ) + return config + + +class _SolModel(BaseDiffusionModel): + def __init__(self, backends: tuple[SOLTrtllmAttention, ...]) -> None: + super().__init__(_sol_model_config()) + self.backends = backends + + def forward(self, q: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: + for backend in self.backends: + q = backend.forward( + q=q, + k=q, + v=q, + batch_size=q.shape[0], + seq_len=q.shape[1], + seq_len_kv=q.shape[1], + timestep=timestep, + ) + return q + + +@_CPU_ONLY +@pytest.mark.parametrize( + ("is_self_attention", "expected_backend", "expects_sol_params"), + ((True, "TRTLLM", True), (False, "VANILLA", False)), + ids=("self", "cross"), +) +def test_sol_attention_module_dispatches_by_attention_role( + monkeypatch, + is_self_attention: bool, + expected_backend: str, + expects_sol_params: bool, +) -> None: + captured = {} + + def _create_attention(*, backend, **kwargs): + captured.update(backend=backend, **kwargs) + return SimpleNamespace(preferred_layout=None) + + monkeypatch.setattr(attention_module, "create_attention", _create_attention) + + attention = Attention( + hidden_size=256, + num_attention_heads=2, + head_dim=128, + qkv_mode=QKVMode.SEPARATE_QKV, + qk_norm=False, + config=_sol_model_config(), + separate_qkv_is_self_attention=is_self_attention, + ) + + assert attention.attn_backend == expected_backend + if expects_sol_params: + assert isinstance(attention.sparse_params, SolParams) + assert captured["sparse_params"] is attention.sparse_params + else: + assert attention.sparse_params is None + assert captured["sparse_params"] is None + + +@_CPU_ONLY +def test_sol_attention_rejects_context_parallelism() -> None: + with pytest.raises(ValueError, match="SOL.*incompatible with context parallelism"): + Attention( + hidden_size=256, + num_attention_heads=2, + head_dim=128, + qk_norm=False, + config=_sol_model_config(cp_size=2), + ) + + +@_CPU_ONLY +def test_sol_cuda_graph_phase_is_keyed_without_model_scope(monkeypatch) -> None: + q = _bshd() + backend, predictor = _stub_backend(SolParams(tau=1.0, disabled_until_timestep=0.6)) + model = _SolModel((backend,)) + runner = CUDAGraphRunner(CUDAGraphRunnerConfig(use_cuda_graph=True)) + model.register_cuda_graph_extra_key_fns(runner) + _stub_core_forward(monkeypatch) + monkeypatch.setattr(sol_backend, "get_bmm1_scale", lambda attn: 0.125) + capturing = False + captured_outputs = {} + captured_keys = [] + + def _capture(key, fn, args, kwargs): + nonlocal capturing + captured_outputs[key] = fn(*args, **kwargs) + capturing = True + try: + captured_outputs[key] = fn(*args, **kwargs) + captured_keys.append(key) + finally: + capturing = False + + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: capturing) + monkeypatch.setattr(runner, "capture", _capture) + monkeypatch.setattr(runner, "replay", lambda key, args, kwargs: captured_outputs[key]) + model.forward = runner.wrap(model.forward) + + assert model(q, timestep=torch.tensor(0.8)).shape == (1, 64, 256) + assert model(q, timestep=torch.tensor(0.2)).shape == (1, 64, 256) + assert ("sol_attn_phase", 0) in captured_keys[0] + assert ("sol_attn_phase", 1) in captured_keys[1] + assert captured_keys[0] != captured_keys[1] + assert predictor.predict.call_count == 2 + + +@_REQUIRES_SM100 +@torch.no_grad() +def test_real_b200_sol_backend_cuda_graph_matches_dense_reference() -> None: + attention_config = AttentionConfig( + backend="TRTLLM", + sparse_attention_config=SolAttentionConfig( + tau=-1.0e6, + disabled_until_timestep=0.6, + ), + ) + sparse_config = attention_config.sparse_attention_config + assert isinstance(sparse_config, SolAttentionConfig) + backend = create_attention( + backend="TRTLLM", + layer_idx=1, + num_heads=2, + head_dim=128, + dtype=torch.bfloat16, + attention_config=attention_config, + attention_metadata_state=create_attention_metadata_state(), + sparse_params=sparse_config.to_sparse_params(), + ) + assert isinstance(backend, SOLTrtllmAttention) + assert any(isinstance(fmha, PrimsTSBlockSparseFmha) for fmha in backend._fmha_manager.fmha_libs) + + generator = torch.Generator(device="cuda").manual_seed(20260901) + shape = (1, 257, 2, 128) + + def _inputs() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + packed = torch.randint( + -2, + 3, + (shape[0], shape[1], 3 * shape[2] * shape[3]), + generator=generator, + device="cuda", + dtype=torch.bfloat16, + ) + return tuple(tensor.view(shape) for tensor in packed.split(shape[2] * shape[3], dim=-1)) + + q, k, v = _inputs() + timestep = torch.tensor(0.2, device="cuda") + assert not any(tensor.is_contiguous() for tensor in (q, k, v)) + eager = backend.forward( + q=q, + k=k, + v=v, + batch_size=1, + seq_len=257, + seq_len_kv=257, + timestep=timestep, + ) + torch.cuda.synchronize() + torch.testing.assert_close( + eager.view_as(q), + _dense_reference(q, k, v), + rtol=2e-2, + atol=2e-2, + ) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = backend.forward( + q=q, + k=k, + v=v, + batch_size=1, + seq_len=257, + seq_len_kv=257, + timestep=timestep, + ).view_as(q) + + next_q, next_k, next_v = _inputs() + q.copy_(next_q) + k.copy_(next_k) + v.copy_(next_v) + graph.replay() + torch.cuda.synchronize() + + torch.testing.assert_close( + captured, + _dense_reference(q, k, v), + rtol=2e-2, + atol=2e-2, + ) + assert backend.predictor.num_plans == 1 + + +@_REQUIRES_SM100 +@torch.no_grad() +@pytest.mark.parametrize("seq_len", [256, 257]) +def test_real_b200_sol_backend_mixed_proxy_cuda_graph_matches_reference(seq_len: int) -> None: + attention_config = AttentionConfig( + backend="TRTLLM", + sparse_attention_config=SolAttentionConfig(tau=1.0e6), + ) + sparse_config = attention_config.sparse_attention_config + assert isinstance(sparse_config, SolAttentionConfig) + backend = create_attention( + backend="TRTLLM", + layer_idx=1, + num_heads=2, + head_dim=128, + dtype=torch.bfloat16, + attention_config=attention_config, + attention_metadata_state=create_attention_metadata_state(), + sparse_params=sparse_config.to_sparse_params(), + ) + assert isinstance(backend, SOLTrtllmAttention) + + generator = torch.Generator(device="cuda").manual_seed(20260903) + shape = (1, seq_len, 2, 128) + + def _inputs() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + packed = torch.randint( + -2, + 3, + (shape[0], shape[1], 3 * shape[2] * shape[3]), + generator=generator, + device="cuda", + dtype=torch.bfloat16, + ) + return tuple(tensor.view(shape) for tensor in packed.split(shape[2] * shape[3], dim=-1)) + + q, k, v = _inputs() + eager = backend.forward(q=q, k=k, v=v, batch_size=1, seq_len=seq_len, seq_len_kv=seq_len) + predictor_outputs = backend.predictor.predict( + q.contiguous(), + k.contiguous(), + v.contiguous(), + tau=1.0e6, + sm_scale=128**-0.5, + ) + torch.cuda.synchronize() + exact_bits = predictor_outputs.exact_block_bits + num_blocks = math.ceil(seq_len / 64) + num_exact = sum( + int( + (exact_bits[..., block_idx // 32].to(torch.int64) >> (block_idx % 32)) + .bitwise_and(1) + .sum() + .item() + ) + for block_idx in range(num_blocks) + ) + assert 0 < num_exact < math.prod(exact_bits.shape[:3]) * num_blocks + torch.testing.assert_close( + eager.view_as(q), + _mixed_proxy_reference(q, k, v, predictor_outputs), + rtol=2e-2, + atol=2e-2, + ) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = backend.forward( + q=q, + k=k, + v=v, + batch_size=1, + seq_len=seq_len, + seq_len_kv=seq_len, + ).view_as(q) + + next_q, next_k, next_v = _inputs() + q.copy_(next_q) + k.copy_(next_k) + v.copy_(next_v) + graph.replay() + torch.cuda.synchronize() + + torch.testing.assert_close( + captured, + _mixed_proxy_reference(q, k, v, predictor_outputs), + rtol=2e-2, + atol=2e-2, + ) diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/test_sol_predictor.py b/tests/unittest/_torch/visual_gen/sparse_attention/test_sol_predictor.py new file mode 100644 index 000000000000..15b31deb210b --- /dev/null +++ b/tests/unittest/_torch/visual_gen/sparse_attention/test_sol_predictor.py @@ -0,0 +1,374 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Qualification tests for the two-stage VisualGen SOL predictor.""" + +from __future__ import annotations + +import dataclasses +import math +import struct + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.predictor import ( + SolPredictorGeometry, + SolPredictorOutputs, + SolPredictorPlanKey, + SOLSparsePredictor, + _normalize_runtime_scalars, +) + +_CPU_ONLY = pytest.mark.cpu_only + + +@_CPU_ONLY +def test_sol_predictor_geometry_and_static_plan_contract() -> None: + geometry = SolPredictorGeometry.create(batch_size=2, seq_len=257, num_heads=3) + assert ( + geometry.tensor_shape, + geometry.summary_shape, + geometry.stats_shape, + geometry.exact_block_bits_shape, + ) == ((2, 257, 3, 128), (2, 5, 3, 128), (2, 3, 128), (2, 3, 5, 1)) + assert (geometry.num_q_blocks, geometry.num_kv_blocks, geometry.tail_tokens) == (5, 5, 1) + + boundary_cases = ( + (64, 1, 1, 64), + (65, 2, 1, 1), + (64 * 32, 32, 1, 64), + (64 * 32 + 1, 33, 2, 1), + ) + for seq_len, blocks, words, tail in boundary_cases: + current = SolPredictorGeometry.create(batch_size=1, seq_len=seq_len, num_heads=1) + assert (current.num_q_blocks, current.exact_words, current.tail_tokens) == ( + blocks, + words, + tail, + ) + + key = SolPredictorPlanKey(geometry=geometry, device_index=1, dtype=torch.bfloat16) + assert tuple(field.name for field in dataclasses.fields(key)) == ( + "geometry", + "device_index", + "dtype", + ) + assert "tau" not in repr(key) and "sm_scale" not in repr(key) + assert SOLSparsePredictor().num_plans == 0 + + +@_CPU_ONLY +def test_sol_predictor_validates_geometry_and_runtime_scalars() -> None: + invalid_geometry = ( + ({"batch_size": 0, "seq_len": 64, "num_heads": 1}, "batch_size"), + ({"batch_size": 1, "seq_len": 0, "num_heads": 1}, "seq_len"), + ({"batch_size": 1, "seq_len": 64, "num_heads": 0}, "num_heads"), + ({"batch_size": True, "seq_len": 64, "num_heads": 1}, "batch_size"), + ({"batch_size": 1, "seq_len": 64, "num_heads": 1, "head_dim": 64}, "head_dim=128"), + ) + for kwargs, message in invalid_geometry: + with pytest.raises((TypeError, ValueError), match=message): + SolPredictorGeometry.create(**kwargs) + + tau, sm_scale = _normalize_runtime_scalars(tau=0.1, sm_scale=math.sqrt(0.5)) + expected_tau = struct.unpack("=f", struct.pack("=f", 0.1))[0] + expected_scale = struct.unpack("=f", struct.pack("=f", math.sqrt(0.5)))[0] + assert (tau, sm_scale) == (expected_tau, expected_scale) + + invalid_scalars = ( + (True, 0.125, "tau"), + (math.nan, 0.125, "tau"), + (0.0, True, "sm_scale"), + (0.0, math.inf, "sm_scale"), + (0.0, 0.0, "sm_scale"), + (0.0, -0.125, "sm_scale"), + ) + for invalid_tau, invalid_scale, message in invalid_scalars: + with pytest.raises((TypeError, ValueError), match=message): + _normalize_runtime_scalars(tau=invalid_tau, sm_scale=invalid_scale) + + +_REQUIRES_CUDA = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +_LOG2_E = math.log2(math.e) + + +def _summary_oracle(k: torch.Tensor, v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + blocks = (k.shape[1] + 63) // 64 + k_summary = torch.empty( + (k.shape[0], blocks, k.shape[2], k.shape[3]), + dtype=torch.bfloat16, + device=k.device, + ) + v_summary = torch.empty_like(k_summary) + for block_idx in range(blocks): + begin = block_idx * 64 + end = min(begin + 64, 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) + return k_summary, v_summary + + +def _pack_bits(exact: torch.Tensor) -> torch.Tensor: + words = (exact.shape[-1] + 31) // 32 + padded = F.pad(exact, (0, words * 32 - exact.shape[-1])).view(*exact.shape[:-1], words, 32) + powers = 1 << torch.arange(32, dtype=torch.int64, device=exact.device) + return (padded.to(torch.int64) * powers).sum(dim=-1).to(torch.uint32) + + +def _predictor_oracle( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + tau: float, + sm_scale: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + k_summary, v_summary = _summary_oracle(k, v) + blocks = k_summary.shape[1] + padded_q = F.pad(q, (0, 0, 0, 0, 0, blocks * 64 - q.shape[1])) + q_blocks = padded_q.view(q.shape[0], blocks, 64, q.shape[2], q.shape[3]) + q_lengths = torch.clamp( + q.shape[1] - torch.arange(blocks, device=q.device) * 64, + min=1, + max=64, + ) + q_centroids = q_blocks.float().sum(dim=2) / q_lengths[None, :, None, None] + k_float = k_summary.float() + k_mean = k_float.mean(dim=1) + k_var = torch.clamp(k_float.square().mean(dim=1) - k_mean.square(), min=0.0) + log2_scale = float(sm_scale) * _LOG2_E + projected_mean = torch.einsum("bqhd,bhd->bqh", q_centroids, k_mean) * log2_scale + projected_var = ( + torch.einsum("bqhd,bhd->bqh", q_centroids.square(), k_var) * log2_scale * log2_scale + ) + threshold = projected_mean + float(tau) * torch.sqrt(projected_var + 1.0e-6) + scores = torch.einsum("bqhd,bkhd->bhqk", q_centroids, k_float) * log2_scale + exact = scores > threshold.permute(0, 2, 1).unsqueeze(-1) + block_ids = torch.arange(blocks, device=q.device) + exact |= (block_ids[:, None] - block_ids[None, :]).abs()[None, None] <= 1 + return _pack_bits(exact), k_summary, v_summary + + +def _small_integer_bf16(shape: tuple[int, ...], *, seed: int) -> torch.Tensor: + generator = torch.Generator(device="cuda").manual_seed(seed) + return torch.randint(-2, 3, shape, generator=generator, device="cuda", dtype=torch.bfloat16) + + +def _inputs( + shape: tuple[int, int, int, int], seed: int +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return ( + _small_integer_bf16(shape, seed=seed), + _small_integer_bf16(shape, seed=seed + 1), + _small_integer_bf16(shape, seed=seed + 2), + ) + + +def _output_tensors( + outputs: SolPredictorOutputs, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return outputs.exact_block_bits, outputs.k_summary, outputs.v_summary + + +def _assert_outputs_match( + outputs: SolPredictorOutputs, + expected: tuple[torch.Tensor, torch.Tensor, torch.Tensor], +) -> None: + assert torch.equal(outputs.exact_block_bits, expected[0]) + torch.testing.assert_close(outputs.k_summary, expected[1], rtol=1e-2, atol=1e-2) + torch.testing.assert_close(outputs.v_summary, expected[2], rtol=1e-2, atol=2e-2) + + +def _run_custom_op( + q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, *buffers: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + torch.ops.trtllm.visual_gen_sol_predictor(q, k, v, *buffers, 64, 0.5, 0.125) + return buffers[0], buffers[1], buffers[2] + + +@_REQUIRES_CUDA +def test_sol_predictor_s257_bh_gt_one_matches_oracle_and_reuses_storage() -> None: + q, k, v = _inputs((2, 257, 3, 128), 11) + predictor = SOLSparsePredictor() + plan = predictor.prepare(q, k, v) + output_ids = tuple(map(id, _output_tensors(plan.outputs))) + scratch_ids = (id(plan.k_mean), id(plan.k_var_diag)) + + outputs = predictor.predict(q, k, v, tau=0.75, sm_scale=0.125) + reference = _predictor_oracle(q, k, v, tau=0.75, sm_scale=0.125) + + assert outputs is plan.outputs + assert tuple(map(id, _output_tensors(outputs))) == output_ids + assert (id(plan.k_mean), id(plan.k_var_diag)) == scratch_ids + _assert_outputs_match(outputs, reference) + expected_k_mean = reference[1].float().mean(dim=1) + expected_k_var = torch.clamp( + reference[1].float().square().mean(dim=1) - expected_k_mean.square(), + min=0.0, + ) + torch.testing.assert_close(plan.k_mean, expected_k_mean, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(plan.k_var_diag, expected_k_var, rtol=1e-5, atol=1e-5) + + second = predictor.predict(q, k, v, tau=-0.25, sm_scale=0.0625) + assert second is outputs + assert tuple(map(id, _output_tensors(second))) == output_ids + assert predictor.num_plans == 1 + + +@_REQUIRES_CUDA +def test_sol_predictor_s257_runtime_scale_and_tau_extremes() -> None: + q, k, v = _inputs((1, 257, 2, 128), 61) + predictor = SOLSparsePredictor() + + normal = predictor.predict(q, k, v, tau=0.5, sm_scale=0.125) + normal_bits = normal.exact_block_bits.clone() + expected_normal = _predictor_oracle(q, k, v, tau=0.5, sm_scale=0.125)[0] + tiny = predictor.predict(q, k, v, tau=0.5, sm_scale=1.0e-5) + expected_tiny = _predictor_oracle(q, k, v, tau=0.5, sm_scale=1.0e-5)[0] + + assert tiny is normal + assert torch.equal(normal_bits, expected_normal) + assert torch.equal(tiny.exact_block_bits, expected_tiny) + assert not torch.equal(expected_normal, expected_tiny) + + blocks = 5 + block_ids = torch.arange(blocks, device=q.device) + local = (block_ids[:, None] - block_ids[None, :]).abs() <= 1 + expected_extremes = ( + _pack_bits(local[None, None].expand(1, 2, -1, -1)), + _pack_bits(torch.ones((1, 2, blocks, blocks), device=q.device, dtype=torch.bool)), + ) + for tau, expected in zip((1.0e6, -1.0e6), expected_extremes, strict=True): + outputs = predictor.predict(q, k, v, tau=tau, sm_scale=128**-0.5) + assert torch.equal(outputs.exact_block_bits, expected) + + +@_REQUIRES_CUDA +def test_sol_predictor_long_proxy_group_keeps_tail_mass_and_clears_padding_bits() -> None: + tokens = 16_451 + q, k, _ = _inputs((1, tokens, 1, 128), 31) + v = torch.ones_like(q) + outputs = SOLSparsePredictor().predict(q, k, v, tau=1.0e6, sm_scale=0.125) + expected, expected_k, expected_v = _predictor_oracle(q, k, v, tau=1.0e6, sm_scale=0.125) + + assert outputs.k_summary.shape[1] == 258 + assert outputs.exact_block_bits.shape[-1] == 9 + 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 int(outputs.exact_block_bits[..., -1].to(torch.int64).max()) < 4 + + +@_REQUIRES_CUDA +def test_sol_predictor_cuda_graph_replay_updates_live_outputs() -> None: + q, k, v = _inputs((1, 257, 2, 128), 41) + predictor = SOLSparsePredictor() + plan = predictor.prepare(q, k, v) + predictor.predict(q, k, v, tau=0.5, sm_scale=0.125) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = predictor.predict(q, k, v, tau=0.5, sm_scale=0.125) + assert captured is plan.outputs + + next_q, next_k, next_v = _inputs(q.shape, 51) + q.copy_(next_q) + k.copy_(next_k) + v.copy_(next_v) + graph.replay() + torch.cuda.synchronize() + expected = _predictor_oracle(q, k, v, tau=0.5, sm_scale=0.125) + + _assert_outputs_match(captured, expected) + + +@_REQUIRES_CUDA +def test_sol_predictor_rejects_plan_miss_during_capture_and_reuses_prepared_plan( + monkeypatch, +) -> None: + q, k, v = _inputs((1, 193, 1, 128), 71) + + with monkeypatch.context() as capture: + capture.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + with pytest.raises(RuntimeError, match="plan must be prepared"): + SOLSparsePredictor().predict(q, k, v, tau=0.5, sm_scale=0.125) + + predictor = SOLSparsePredictor() + plan = predictor.prepare(q, k, v) + with monkeypatch.context() as capture: + capture.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + assert predictor.predict(q, k, v, tau=0.5, sm_scale=0.125) is plan.outputs + torch.cuda.synchronize() + _assert_outputs_match(plan.outputs, _predictor_oracle(q, k, v, tau=0.5, sm_scale=0.125)) + + +@_REQUIRES_CUDA +def test_sol_predictor_compiled_public_predict_owns_each_instance_plan(recwarn) -> None: + output_ptrs = [] + for seed in (81, 91): + q, k, v = _inputs((1, 257, 2, 128), seed) + predictor = SOLSparsePredictor() + compiled_predict = torch.compile(predictor.predict, backend="eager", fullgraph=False) + + outputs = compiled_predict(q, k, v, tau=0.5, sm_scale=0.125) + expected = _predictor_oracle(q, k, v, tau=0.5, sm_scale=0.125) + + assert predictor.num_plans == 1 + _assert_outputs_match(outputs, expected) + output_ptrs.append(tuple(tensor.data_ptr() for tensor in _output_tensors(outputs))) + + assert output_ptrs[0] != output_ptrs[1] + assert not any("recompile_limit" in str(warning.message) for warning in recwarn) + + +@_REQUIRES_CUDA +def test_sol_predictor_custom_op_fake_schema_and_fullgraph_compile() -> None: + from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol import kernels # noqa: F401 + + op = torch.ops.trtllm.visual_gen_sol_predictor.default + schema = str(op._schema) + assert "exact_block_bits" in schema + assert "k_summary" in schema + assert "v_summary" in schema + assert "!" in schema + assert torch._C._dispatch_has_kernel_for_dispatch_key( + "trtllm::visual_gen_sol_predictor", "Meta" + ) + + meta_q = torch.empty((1, 65, 1, 128), device="meta", dtype=torch.bfloat16) + meta_summary = torch.empty((1, 2, 1, 128), device="meta", dtype=torch.bfloat16) + meta_stats = torch.empty((1, 1, 128), device="meta", dtype=torch.float32) + meta_args = ( + meta_q, + torch.empty_like(meta_q), + torch.empty_like(meta_q), + torch.empty((1, 1, 2, 1), device="meta", dtype=torch.uint32), + meta_summary, + torch.empty_like(meta_summary), + meta_stats, + torch.empty_like(meta_stats), + torch.empty((1, 2, 1, 128), device="meta", dtype=torch.float32), + ) + assert op(*meta_args, 64, 0.5, 0.125) is None + + q, k, v = _inputs((1, 257, 2, 128), 81) + plan = SOLSparsePredictor().prepare(q, k, v) + buffers = ( + plan.outputs.exact_block_bits, + plan.outputs.k_summary, + plan.outputs.v_summary, + plan.k_mean, + plan.k_var_diag, + plan.q_centroid, + ) + actual = torch.compile(_run_custom_op, backend="eager", fullgraph=True)(q, k, v, *buffers) + expected = _predictor_oracle(q, k, v, tau=0.5, sm_scale=0.125) + + assert all( + actual_tensor is plan_tensor + for actual_tensor, plan_tensor in zip(actual, buffers[:3], strict=True) + ) + _assert_outputs_match(plan.outputs, expected) diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/test_sol_predictor_kernels.py b/tests/unittest/_torch/visual_gen/sparse_attention/test_sol_predictor_kernels.py new file mode 100644 index 000000000000..79556dbcf31d --- /dev/null +++ b/tests/unittest/_torch/visual_gen/sparse_attention/test_sol_predictor_kernels.py @@ -0,0 +1,206 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for the SOL predictor kernels (block pooling, block statistics, exact-block selection).""" + +from __future__ import annotations + +import pytest +import torch + +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.kernels import ( + _block_pool_torch, + _block_statistics_torch, + _select_exact_blocks_torch, + block_pool, + block_statistics, + select_exact_blocks, +) + +_CPU_ONLY = pytest.mark.cpu_only +_REQUIRES_CUDA = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +BLOCK = 64 + + +def _pooled_shape(x: torch.Tensor, block_size: int) -> tuple[int, int, int, int]: + batch, seq_len, heads, head_dim = x.shape + return batch, (seq_len + block_size - 1) // block_size, heads, head_dim + + +def _unpack(bits: torch.Tensor, num_kv_blocks: int) -> torch.Tensor: + words = bits.view(torch.int32).to(torch.int64) & 0xFFFFFFFF + shifts = torch.arange(32, device=bits.device, dtype=torch.int64) + return ( + ((words.unsqueeze(-1) >> shifts) & 1) + .bool() + .reshape(*bits.shape[:-1], -1)[..., :num_kv_blocks] + ) + + +@_CPU_ONLY +def test_block_pool_torch_fallback_means_valid_tokens_only() -> None: + x = torch.zeros((1, 70, 1, 4), dtype=torch.bfloat16) + x[0, :64] = 2.0 + x[0, 64:70] = 3.0 + out = torch.empty(_pooled_shape(x, BLOCK), dtype=torch.float32) + block_pool(x, out, block_size=BLOCK, reduce="mean") + assert torch.equal(out[0, 0], torch.full((1, 4), 2.0)) + assert torch.equal(out[0, 1], torch.full((1, 4), 3.0)) + total = torch.empty_like(out) + block_pool(x, total, block_size=BLOCK, reduce="sum") + assert torch.equal(total[0, 1], torch.full((1, 4), 18.0)) + + +@_CPU_ONLY +def test_block_pool_rejects_mismatched_output() -> None: + x = torch.zeros((1, 70, 1, 4), dtype=torch.bfloat16) + with pytest.raises(ValueError, match="out"): + block_pool(x, torch.empty((1, 3, 1, 4)), block_size=BLOCK, reduce="mean") + with pytest.raises(ValueError, match="reduce"): + block_pool(x, torch.empty(_pooled_shape(x, BLOCK)), block_size=BLOCK, reduce="max") + + +@_CPU_ONLY +def test_select_exact_blocks_torch_fallback_packs_bit_r_of_word_w() -> None: + blocks = 35 + centroid = torch.zeros((1, blocks, 1, 8), dtype=torch.float32) + k_summary = torch.zeros((1, blocks, 1, 8), dtype=torch.bfloat16) + k_summary[0, 33, 0, 0] = 1.0 + centroid[0, 0, 0, 0] = 1.0 + k_mean = torch.zeros((1, 1, 8)) + k_var = torch.zeros((1, 1, 8)) + bits = torch.empty((1, 1, blocks, 2), dtype=torch.uint32) + select_exact_blocks(centroid, k_summary, k_mean, k_var, bits, tau=0.5, sm_scale=1.0) + exact = _unpack(bits, blocks) + # Row 0 scores 1.0 against block 33 only (threshold 0.5 * sqrt(1e-6)) plus its local band. + expected = torch.zeros(blocks, dtype=torch.bool) + expected[[0, 1, 33]] = True + assert torch.equal(exact[0, 0, 0], expected) + # Bit 1 of word 1 is block 33. + assert int(bits[0, 0, 0, 1]) == 2 + # Rows without scores keep only the local band. + assert torch.equal(exact[0, 0, 17].nonzero().flatten(), torch.tensor([16, 17, 18])) + + +@_REQUIRES_CUDA +@pytest.mark.parametrize("seq_len", [64, 257, 4097]) +@pytest.mark.parametrize("out_dtype", [torch.float32, torch.bfloat16]) +def test_block_pool_matches_torch_fallback(seq_len: int, out_dtype: torch.dtype) -> None: + torch.manual_seed(0) + x = torch.randn((2, seq_len, 3, 128), device="cuda", dtype=torch.bfloat16) + for reduce in ("mean", "sum"): + out = torch.empty(_pooled_shape(x, BLOCK), dtype=out_dtype, device="cuda") + block_pool(x, out, block_size=BLOCK, reduce=reduce) + expected = torch.empty_like(out) + _block_pool_torch(x, expected, block_size=BLOCK, reduce=reduce) + tolerance = ( + {"rtol": 1e-5, "atol": 1e-5} + if out_dtype == torch.float32 + else {"rtol": 1e-2, "atol": 1e-2} + ) + torch.testing.assert_close(out, expected, **tolerance) + + +@_REQUIRES_CUDA +def test_block_pool_accepts_strided_batch_and_token_dims() -> None: + torch.manual_seed(1) + full = torch.randn((2, 130, 2, 3, 128), device="cuda", dtype=torch.bfloat16) + x = full[:, :, 1] # heads/head_dim contiguous, token stride wider than a row + out = torch.empty(_pooled_shape(x, BLOCK), dtype=torch.float32, device="cuda") + block_pool(x, out, block_size=BLOCK, reduce="mean") + expected = torch.empty_like(out) + _block_pool_torch(x.contiguous(), expected, block_size=BLOCK, reduce="mean") + torch.testing.assert_close(out, expected, rtol=1e-5, atol=1e-5) + + +@_REQUIRES_CUDA +def test_block_statistics_matches_torch_fallback() -> None: + torch.manual_seed(2) + k_summary = torch.randn((2, 1182, 3, 128), device="cuda", dtype=torch.bfloat16) + mean = torch.empty((2, 3, 128), device="cuda", dtype=torch.float32) + var = torch.empty_like(mean) + block_statistics(k_summary, mean, var) + expected_mean = torch.empty_like(mean) + expected_var = torch.empty_like(var) + _block_statistics_torch(k_summary, expected_mean, expected_var) + torch.testing.assert_close(mean, expected_mean, rtol=1e-5, atol=1e-6) + torch.testing.assert_close(var, expected_var, rtol=1e-4, atol=1e-6) + assert bool((var >= 0).all()) + + +@_REQUIRES_CUDA +@pytest.mark.parametrize("num_blocks", [5, 258]) +def test_select_exact_blocks_matches_fallback_and_clears_padding_bits(num_blocks: int) -> None: + torch.manual_seed(3) + batch, heads, dim = 2, 3, 128 + centroid = torch.randn((batch, num_blocks, heads, dim), device="cuda") * 0.125 + k_summary = (torch.randn((batch, num_blocks, heads, dim), device="cuda") * 0.125).to( + torch.bfloat16 + ) + k_mean = torch.empty((batch, heads, dim), device="cuda") + k_var = torch.empty_like(k_mean) + block_statistics(k_summary, k_mean, k_var) + words = (num_blocks + 31) // 32 + bits = torch.empty((batch, heads, num_blocks, words), device="cuda", dtype=torch.uint32) + expected = torch.empty_like(bits) + for tau in (0.75, -1.0e6, 1.0e6): + select_exact_blocks(centroid, k_summary, k_mean, k_var, bits, tau=tau, sm_scale=0.125) + _select_exact_blocks_torch( + centroid, k_summary, k_mean, k_var, expected, tau=tau, sm_scale=0.125 + ) + assert torch.equal(bits, expected), f"tau={tau}" + padding = words * 32 - num_blocks + if padding: + assert int(bits[..., -1].to(torch.int64).max()) < (1 << (32 - padding)) + exact = _unpack(bits, num_blocks) + ids = torch.arange(num_blocks, device="cuda") + assert bool(exact[..., (ids[:, None] - ids[None, :]).abs() <= 1].all()) + + +@_REQUIRES_CUDA +def test_kernels_replay_inside_cuda_graph() -> None: + torch.manual_seed(4) + q = torch.randn((1, 257, 2, 128), device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + shape = _pooled_shape(q, BLOCK) + centroid = torch.empty(shape, device="cuda", dtype=torch.float32) + k_summary = torch.empty(shape, device="cuda", dtype=torch.bfloat16) + mean = torch.empty((1, 2, 128), device="cuda") + var = torch.empty_like(mean) + bits = torch.empty((1, 2, shape[1], 1), device="cuda", dtype=torch.uint32) + + def run() -> None: + block_pool(q, centroid, block_size=BLOCK, reduce="mean") + block_pool(k, k_summary, block_size=BLOCK, reduce="mean") + block_statistics(k_summary, mean, var) + select_exact_blocks(centroid, k_summary, mean, var, bits, tau=0.5, sm_scale=0.125) + + run() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + q.copy_(torch.randn_like(q)) + k.copy_(torch.randn_like(k)) + graph.replay() + torch.cuda.synchronize() + + expected_centroid = torch.empty_like(centroid) + expected_summary = torch.empty_like(k_summary) + _block_pool_torch(q, expected_centroid, block_size=BLOCK, reduce="mean") + _block_pool_torch(k, expected_summary, block_size=BLOCK, reduce="mean") + expected_mean = torch.empty_like(mean) + expected_var = torch.empty_like(var) + _block_statistics_torch(expected_summary, expected_mean, expected_var) + expected_bits = torch.empty_like(bits) + _select_exact_blocks_torch( + expected_centroid, + expected_summary, + expected_mean, + expected_var, + expected_bits, + tau=0.5, + sm_scale=0.125, + ) + torch.testing.assert_close(centroid, expected_centroid, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(k_summary, expected_summary, rtol=1e-2, atol=1e-2) + assert torch.equal(bits, expected_bits) diff --git a/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py b/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py index f1064f780fa1..057fafc91520 100644 --- a/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py @@ -1284,6 +1284,9 @@ def __init__(self): self.active_topology = "default" self.device = "cuda" + def register_cuda_graph_extra_key_fns(self, runner): + del runner + def forward(self, x): return self.lin(x) @@ -1344,6 +1347,9 @@ def test_two_stage_cuda_graph_setup_uses_pipeline_config(self): """CUDA graph setup runs before the two-stage model_config is assigned.""" class TinyTransformer: + def register_cuda_graph_extra_key_fns(self, runner): + del runner + def forward(self, *args, **kwargs): return args, kwargs @@ -1364,6 +1370,180 @@ def forward(self, *args, **kwargs): assert runner._lora_state_getter() == "original" assert pipeline.transformer.forward.__wrapped__.__self__ is pipeline.transformer + def test_two_stage_cuda_graph_setup_registers_sol_phase_key(self): + """Dense and sparse SOL phases must never reuse one CUDA graph.""" + from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig + from tensorrt_llm._torch.visual_gen.models.ltx2.ltx2_core.modality import Modality + from tensorrt_llm._torch.visual_gen.models.modeling import BaseDiffusionModel + from tensorrt_llm.visual_gen.args import SolAttentionConfig + + class TinySolTransformer(BaseDiffusionModel): + def __init__(self): + super().__init__( + DiffusionModelConfig( + attention=AttentionConfig( + backend="TRTLLM", + sparse_attention_config=SolAttentionConfig(disabled_until_timestep=0.6), + ) + ) + ) + self.active_topology = "default" + + def forward(self, video, audio, *, text_cache, timestep=None, step_index=None): + del audio, text_cache, timestep, step_index + return video.latent, None + + pipeline = object.__new__(ltx2_two_stages.LTX2TwoStagesPipeline) + torch.nn.Module.__init__(pipeline) + pipeline.pipeline_config = DiffusionPipelineConfig( + cuda_graph=CudaGraphConfig(enable=True), + torch_compile=TorchCompileConfig(enable=False), + ) + pipeline.transformer = TinySolTransformer() + pipeline._cuda_graph_runners = {} + pipeline._setup_cuda_graphs() + runner = pipeline._cuda_graph_runners["transformer"] + + captured_keys = [] + + def fake_capture(key, fn, args, kwargs): + del fn, args, kwargs + captured_keys.append(key) + runner.graphs[key] = object() + + runner.capture = fake_capture + runner.replay = lambda key, args, kwargs: key + + def modality(timestep): + return Modality( + latent=torch.empty(1, 2, 4), + timesteps=torch.tensor([timestep]), + positions=torch.empty(1, 3, 2), + context=torch.empty(1, 3, 4), + ) + + dense_key = pipeline.transformer( + video=modality(0.8), audio=None, text_cache=None, step_index=0 + ) + sparse_key = pipeline.transformer( + video=modality(0.2), audio=None, text_cache=None, step_index=1 + ) + dense_replay_key = pipeline.transformer( + video=modality(0.8), audio=None, text_cache=None, step_index=2 + ) + + assert "sol_attn_phase" in runner._extra_key_fns + assert ("sol_attn_phase", (("video", 0),)) in dense_key + assert ("sol_attn_phase", (("video", 1),)) in sparse_key + assert dense_key != sparse_key + assert dense_replay_key == dense_key + assert captured_keys == [dense_key, sparse_key] + + def test_ltx2_threads_raw_modality_timestep_for_sol_phase(self): + """SOL phase preparation must not use the AdaLN-transformed timestep.""" + from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.params import SolParams + from tensorrt_llm._torch.visual_gen.models.ltx2.ltx2_core.modality import Modality + from tensorrt_llm._torch.visual_gen.models.ltx2.ltx2_core.transformer_args import ( + TransformerArgs, + ) + from tensorrt_llm._torch.visual_gen.models.ltx2.transformer_ltx2 import LTXModel + + def modality(timestep): + return Modality( + latent=torch.empty(1, 1, 1), + timesteps=torch.tensor([timestep]), + positions=torch.empty(1, 1, 1), + context=torch.empty(1, 1, 1), + ) + + def transformer_args(timestep): + return TransformerArgs( + x=torch.empty(1, 1, 1), + context=torch.empty(1, 1, 1), + context_mask=None, + timesteps=torch.tensor([timestep]), + embedded_timestep=torch.empty(1, 1, 1), + positional_embeddings=(torch.empty(1), torch.empty(1)), + cross_positional_embeddings=None, + cross_scale_shift_timestep=None, + cross_gate_timestep=None, + enabled=True, + ) + + class Preprocessor: + def __init__(self, value): + self.value = value + + def prepare(self, *args, **kwargs): + del args, kwargs + return self.value + + observed_phases = [] + + class Block: + def __call__( + self, + *, + video, + audio, + video_sol_timestep=None, + audio_sol_timestep=None, + **kwargs, + ): + del kwargs + video_timestep = ( + video.timesteps if video_sol_timestep is None else video_sol_timestep + ) + audio_timestep = ( + audio.timesteps if audio_sol_timestep is None else audio_sol_timestep + ) + observed_phases.append( + ( + SolParams.get_graph_phase_for_timestep( + video_timestep, + disabled_until_timestep=0.6, + ), + SolParams.get_graph_phase_for_timestep( + audio_timestep, + disabled_until_timestep=0.6, + ), + ) + ) + return video, audio + + model = object.__new__(LTXModel) + torch.nn.Module.__init__(model) + model._active_seq_size = 1 + model.model_type = SimpleNamespace( + is_video_enabled=lambda: True, + is_audio_enabled=lambda: True, + ) + model._audio_pad = 0 + model.video_args_preprocessor = Preprocessor(transformer_args(0.8)) + model.audio_args_preprocessor = Preprocessor(transformer_args(0.2)) + model._active_sharder = SimpleNamespace(is_active=False) + model.model_config = SimpleNamespace(cache_backend=None) + model.transformer_blocks = [Block()] + model.scale_shift_table = model.norm_out = model.proj_out = None + model.audio_scale_shift_table = model.audio_norm_out = model.audio_proj_out = None + model._process_output = lambda table, norm, proj, x, embedded: x + text_cache = SimpleNamespace( + video_context=None, + video_mask=None, + video_pe=None, + video_cross_pe=None, + video_kv=None, + audio_context=None, + audio_mask=None, + audio_pe=None, + audio_cross_pe=None, + audio_kv=None, + ) + + model(video=modality(0.2), audio=modality(0.8), text_cache=text_cache) + + assert observed_phases == [(1, 0)] + def test_cuda_graph_rejects_nonpersistent_lora_bindings(self): """CUDA graph is valid only when distilled LoRA uses persistent bindings.""" pipeline = object.__new__(ltx2_two_stages.LTX2TwoStagesPipeline) diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py index 19378f37fb70..14a0d7c272ba 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py @@ -19,6 +19,7 @@ QuantAttentionConfig, RuntimeLoRAConfig, SkipSoftmaxAttentionConfig, + SolAttentionConfig, TeaCacheConfig, TorchCompileConfig, VAEConfig, @@ -493,6 +494,54 @@ def test_from_yaml_unknown_field_raises(self, tmp_path): with pytest.raises(ValidationError, match="Extra inputs are not permitted"): VisualGenArgs.from_yaml(yaml_path) + def test_from_yaml_rejects_sol_with_enabled_fullgraph(self, tmp_path): + yaml_path = tmp_path / "sol_fullgraph.yml" + yaml_path.write_text( + "model: /tmp/model\n" + "attention_config:\n" + " backend: TRTLLM\n" + " sparse_attention_config:\n" + " algorithm: sol_attn\n" + "torch_compile_config:\n" + " enable: true\n" + " enable_fullgraph: true\n" + ) + + with pytest.raises(ValidationError, match="SOL.*fullgraph"): + VisualGenArgs.from_yaml(yaml_path) + + +class TestVisualGenArgsCrossFieldValidation: + def test_rejects_sol_with_enabled_fullgraph(self): + with pytest.raises(ValidationError, match="SOL.*fullgraph"): + VisualGenArgs( + model="/tmp/model", + attention_config=AttentionConfig( + backend="TRTLLM", + sparse_attention_config=SolAttentionConfig(), + ), + torch_compile_config=TorchCompileConfig( + enable=True, + enable_fullgraph=True, + ), + ) + + def test_allows_sol_fullgraph_field_when_torch_compile_disabled(self): + args = VisualGenArgs( + model="/tmp/model", + attention_config=AttentionConfig( + backend="TRTLLM", + sparse_attention_config=SolAttentionConfig(), + ), + torch_compile_config=TorchCompileConfig( + enable=False, + enable_fullgraph=True, + ), + ) + + assert args.torch_compile_config.enable is False + assert args.torch_compile_config.enable_fullgraph is True + class TestParallelConfigValidation: """ParallelConfig no longer checks WORLD_SIZE at construction time.""" From afc6692312d4530294caeced0dc21d15542fd204 Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Wed, 9 Sep 2026 04:11:23 +0000 Subject: [PATCH 7/7] perf: fuse VSA predictor preprocessing into Triton kernels The Video Sparse Attention predictor spent most of its time in generic gather kernels and redundant passes: index_select/index_copy_ tiling of Q/K/V at a fraction of HBM bandwidth, an fp32 copy of the tiled tensors for the cube mean, an expanded copy of the coarse output followed by an untile gather, a value sort inside topk, and torch.sort over int32 rows. Replace them with three memory-bound kernels (Triton on CUDA, PyTorch fallback elsewhere) that keep the numerics of the previous code: - tile_and_pool_cubes gathers tokens into the tile-major layout and emits the fp32-accumulated cube mean in the same pass; - sort_last_dim orders each BSR route row by cube index; - blend_coarse_fine gathers the coarse and fine outputs back to compact order and applies both gates in one kernel, reading head-major fine output through strides. Metadata carries the padded-slot source index and the packed valid-token words per shape; topk skips the value sort when routes are re-sorted; the post-process context keeps the coarse output per cube. Launch shapes come from the row width and grid size, so no autotuning runs at call time and the predictor stays CUDA Graph capturable. On B200 with Wan 14B shapes (B2/H40/D128) the predictor drops from 34.6 ms to 5.1 ms per call for 81 frames and from 4.7 ms to 0.56 ms for 9 frames; the post-process drops from 4.1 ms to 0.9 ms and from 0.61 ms to 0.14 ms. Selected routes and tiled Q/K/V are unchanged. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- .../attention_backend/sparse/vsa/__init__.py | 3 +- .../attention_backend/sparse/vsa/kernels.py | 385 ++++++++++++++++++ .../attention_backend/sparse/vsa/metadata.py | 36 +- .../attention_backend/sparse/vsa/predictor.py | 223 ++++------ .../_torch/visual_gen/test_attention_vsa.py | 157 ++++++- .../visual_gen/test_attention_vsa_kernels.py | 195 +++++++++ 6 files changed, 821 insertions(+), 178 deletions(-) create mode 100644 tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/kernels.py create mode 100644 tests/unittest/_torch/visual_gen/test_attention_vsa_kernels.py diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/__init__.py b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/__init__.py index 2748d3cc0686..3971961794ed 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/__init__.py @@ -22,7 +22,7 @@ get_vsa_forward_context, set_vsa_forward_context, ) -from .predictor import VSAForwardInputs, VSAPredictor, VSAPreprocessor +from .predictor import VSAForwardInputs, VSAPredictor __all__ = [ "VSA_TILE_SIZE", @@ -30,7 +30,6 @@ "VSAMetadata", "VSAMetadataBuilder", "VSAPredictor", - "VSAPreprocessor", "get_vsa_forward_context", "set_vsa_forward_context", ] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/kernels.py b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/kernels.py new file mode 100644 index 000000000000..56e3df83a239 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/kernels.py @@ -0,0 +1,385 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Memory-bound kernels of the Video Sparse Attention predictor and post-process. + +Every helper streams ``[batch, tokens, heads, head_dim]`` activations row by row, where a row +is one token's ``heads * head_dim`` values. CUDA tensors run Triton kernels; other tensors use +PyTorch implementations with the same numerics. Launch shapes are derived from the row width +and the grid size, so no autotuning happens at call time and every launch is CUDA Graph safe. +""" + +import torch +import triton +import triton.language as tl + +_MAX_BLOCK = 1024 +_MIN_ELEMENTS_PER_THREAD = 4 +_SMALL_GRID_PROGRAMS = 2048 +_MAX_SLOTS_PER_ITERATION = 4 +_MAX_TRITON_SORT_LENGTH = 2048 + + +def _row_launch_config(row_width: int, num_row_programs: int) -> tuple[int, int, int]: + """Choose ``(block, num_chunks, num_warps)`` for a kernel that streams rows in chunks. + + A block covers up to 1024 columns; wider rows are split into chunks. Small grids get + more warps per program to expose parallelism, large grids fewer warps so that every + thread keeps several elements in flight. + """ + block = min(_MAX_BLOCK, triton.next_power_of_2(row_width)) + num_chunks = triton.cdiv(row_width, block) + num_warps = 4 if num_row_programs * num_chunks < _SMALL_GRID_PROGRAMS else 2 + num_warps = min(num_warps, max(1, block // (32 * _MIN_ELEMENTS_PER_THREAD))) + return block, num_chunks, num_warps + + +def _slots_per_iteration(cube_size: int) -> int: + """Largest power of two up to four that divides the cube, so slot tiles stay rectangular.""" + slots = _MAX_SLOTS_PER_ITERATION + while cube_size % slots: + slots //= 2 + return slots + + +@triton.jit +def _tile_and_pool_cubes_kernel( + x_ptr, + source_ptr, + count_ptr, + tiled_ptr, + pooled_ptr, + num_cubes, + stride_batch, + stride_token, + ROW: tl.constexpr, + CUBE: tl.constexpr, + SLOTS: tl.constexpr, + BLOCK: tl.constexpr, +): + """One program per (batch, cube, column chunk). + + The cube's CUBE slots are streamed SLOTS at a time: each slot is copied from its compact + source token (zero for a pad slot) into the tiled layout while an fp32 accumulator builds + the cube mean. + """ + batch_cube = tl.program_id(0).to(tl.int64) + chunk = tl.program_id(1) + batch = batch_cube // num_cubes + cube = batch_cube % num_cubes + columns = chunk * BLOCK + tl.arange(0, BLOCK) + in_row = columns < ROW + slot_offsets = tl.arange(0, SLOTS) + + x_batch_ptr = x_ptr + batch * stride_batch + tiled_cube_ptr = tiled_ptr + batch_cube * CUBE * ROW + total = tl.zeros([BLOCK], dtype=tl.float32) + for first_slot in range(0, CUBE, SLOTS): + slots = first_slot + slot_offsets + sources = tl.load(source_ptr + cube * CUBE + slots) + values = tl.load( + x_batch_ptr + sources[:, None] * stride_token + columns[None, :], + mask=(sources >= 0)[:, None] & in_row[None, :], + other=0.0, + ) + tl.store( + tiled_cube_ptr + slots[:, None] * ROW + columns[None, :], + values, + mask=(slots >= 0)[:, None] & in_row[None, :], + ) + total += tl.sum(values.to(tl.float32), axis=0) + + count = tl.load(count_ptr + cube).to(tl.float32) + mean = (total / count).to(pooled_ptr.dtype.element_ty) + tl.store(pooled_ptr + batch_cube * ROW + columns, mean, mask=in_row) + + +def _tile_and_pool_cubes_torch( + x: torch.Tensor, + tile_source_index: torch.Tensor, + cube_valid_counts: torch.Tensor, + *, + cube_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + batch_size, _, num_heads, head_dim = x.shape + num_cubes = cube_valid_counts.shape[0] + tiled = x.index_select(1, tile_source_index.clamp(min=0)) + is_valid_slot = (tile_source_index >= 0).view(1, -1, 1, 1) + tiled = torch.where(is_valid_slot, tiled, torch.zeros((), dtype=x.dtype, device=x.device)) + total = tiled.view(batch_size, num_cubes, cube_size, num_heads, head_dim).sum( + dim=2, dtype=torch.float32 + ) + mean = total / cube_valid_counts.view(1, -1, 1, 1).to(torch.float32) + return tiled, mean.to(x.dtype) + + +def tile_and_pool_cubes( + x: torch.Tensor, + tile_source_index: torch.Tensor, + cube_valid_counts: torch.Tensor, + *, + cube_size: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Gather tokens into the tile-major padded layout and mean-pool every cube. + + Args: + x: ``[batch, seq_len, heads, head_dim]`` activations; the heads and head_dim + dimensions must be contiguous, batch and sequence strides are arbitrary. + tile_source_index: ``[num_cubes * cube_size]`` long tensor with the compact token of + every padded slot, or ``-1`` for a pad slot. + cube_valid_counts: ``[num_cubes]`` number of valid tokens per cube (at least one). + cube_size: Tokens per cube. + + Returns: + ``tiled`` ``[batch, num_cubes * cube_size, heads, head_dim]`` with zeroed pad slots and + ``pooled`` ``[batch, num_cubes, heads, head_dim]`` cube means in ``x.dtype`` computed + with an fp32 accumulator. + """ + batch_size, _, num_heads, head_dim = x.shape + if x.stride(3) != 1 or (num_heads > 1 and x.stride(2) != head_dim): + raise ValueError("heads and head_dim must be contiguous in the VSA input") + num_cubes = cube_valid_counts.shape[0] + if tile_source_index.numel() != num_cubes * cube_size: + raise ValueError( + f"tile_source_index must have {num_cubes * cube_size} slots, " + f"got {tile_source_index.numel()}" + ) + if not x.is_cuda: + return _tile_and_pool_cubes_torch( + x, tile_source_index, cube_valid_counts, cube_size=cube_size + ) + + row_width = num_heads * head_dim + tiled = torch.empty( + (batch_size, num_cubes * cube_size, num_heads, head_dim), dtype=x.dtype, device=x.device + ) + pooled = torch.empty( + (batch_size, num_cubes, num_heads, head_dim), dtype=x.dtype, device=x.device + ) + block, num_chunks, num_warps = _row_launch_config(row_width, batch_size * num_cubes) + _tile_and_pool_cubes_kernel[(batch_size * num_cubes, num_chunks)]( + x, + tile_source_index, + cube_valid_counts, + tiled, + pooled, + num_cubes, + x.stride(0), + x.stride(1), + ROW=row_width, + CUBE=cube_size, + SLOTS=_slots_per_iteration(cube_size), + BLOCK=block, + num_warps=num_warps, + ) + return tiled, pooled + + +@triton.jit +def _sort_rows_kernel(values_ptr, sorted_ptr, ROW: tl.constexpr, BLOCK: tl.constexpr): + """One program per row; pad slots sort to the end and are never stored.""" + row = tl.program_id(0).to(tl.int64) + columns = tl.arange(0, BLOCK) + in_row = columns < ROW + values = tl.load(values_ptr + row * ROW + columns, mask=in_row, other=2147483647) + tl.store(sorted_ptr + row * ROW + columns, tl.sort(values), mask=in_row) + + +def _sort_last_dim_torch(values: torch.Tensor) -> torch.Tensor: + return torch.sort(values, dim=-1).values + + +def sort_last_dim(values: torch.Tensor) -> torch.Tensor: + """Sort an int32 tensor ascending along its last dimension. + + Rows up to 2048 entries are sorted by one Triton program each; longer rows and non-CUDA + tensors use ``torch.sort``. + """ + if values.dtype != torch.int32: + raise TypeError(f"sort_last_dim expects int32 values, got {values.dtype}") + row_length = values.shape[-1] + if not values.is_cuda or row_length > _MAX_TRITON_SORT_LENGTH or values.numel() == 0: + return _sort_last_dim_torch(values) + + rows = values.reshape(-1, row_length).contiguous() + sorted_rows = torch.empty_like(rows) + _sort_rows_kernel[(rows.shape[0],)]( + rows, + sorted_rows, + ROW=row_length, + BLOCK=triton.next_power_of_2(row_length), + num_warps=1, + ) + return sorted_rows.view(values.shape) + + +@triton.jit +def _blend_coarse_fine_kernel( + fine_ptr, + coarse_ptr, + gate_compress_ptr, + gate_fine_ptr, + untile_ptr, + out_ptr, + seq_len, + num_cubes, + stride_fine_batch, + stride_fine_token, + stride_fine_head, + ROW: tl.constexpr, + HEAD_DIM: tl.constexpr, + CUBE: tl.constexpr, + BLOCK: tl.constexpr, + HAS_GATE_FINE: tl.constexpr, + FINE_IS_TILED: tl.constexpr, +): + """One program per (batch, compact token, column chunk). + + The fine output is addressed through explicit strides so that head-major storage + (``[batch, heads, tokens, head_dim]``) is consumed without a copy. Products and the final + sum are rounded to the output dtype after each operation, which matches the PyTorch + expression ``gate_compress * coarse + gate_fine * fine``. + """ + batch_token = tl.program_id(0).to(tl.int64) + chunk = tl.program_id(1) + batch = batch_token // seq_len + token = batch_token % seq_len + slot = tl.load(untile_ptr + token) + cube = slot // CUBE + if FINE_IS_TILED: + fine_row = slot + else: + fine_row = token + columns = chunk * BLOCK + tl.arange(0, BLOCK) + in_row = columns < ROW + head = (columns // HEAD_DIM).to(tl.int64) + dim = columns % HEAD_DIM + out_dtype = out_ptr.dtype.element_ty + + fine_ptrs = ( + fine_ptr + + batch * stride_fine_batch + + fine_row * stride_fine_token + + head * stride_fine_head + + dim + ) + fine = tl.load(fine_ptrs, mask=in_row, other=0.0) + coarse = tl.load( + coarse_ptr + (batch * num_cubes + cube) * ROW + columns, mask=in_row, other=0.0 + ) + gate_compress = tl.load(gate_compress_ptr + batch_token * ROW + columns, mask=in_row, other=0.0) + coarse_term = (gate_compress.to(tl.float32) * coarse.to(tl.float32)).to(out_dtype) + if HAS_GATE_FINE: + gate_fine = tl.load(gate_fine_ptr + batch_token * ROW + columns, mask=in_row, other=0.0) + fine_term = (gate_fine.to(tl.float32) * fine.to(tl.float32)).to(out_dtype) + else: + fine_term = fine.to(out_dtype) + result = (coarse_term.to(tl.float32) + fine_term.to(tl.float32)).to(out_dtype) + tl.store(out_ptr + batch_token * ROW + columns, result, mask=in_row) + + +def _blend_coarse_fine_torch( + fine: torch.Tensor, + coarse: torch.Tensor, + gate_compress: torch.Tensor, + gate_fine: torch.Tensor | None, + untile_index: torch.Tensor, + *, + cube_size: int, + fine_is_tiled: bool, +) -> torch.Tensor: + coarse_per_token = coarse.index_select(1, untile_index // cube_size) + fine_compact = fine.index_select(1, untile_index) if fine_is_tiled else fine + if gate_fine is not None: + fine_compact = gate_fine * fine_compact + return gate_compress * coarse_per_token + fine_compact + + +def blend_coarse_fine( + fine: torch.Tensor, + coarse: torch.Tensor, + gate_compress: torch.Tensor, + gate_fine: torch.Tensor | None, + untile_index: torch.Tensor, + *, + cube_size: int, + fine_is_tiled: bool, +) -> torch.Tensor: + """Restore compact token order and blend the coarse and fine VSA outputs. + + Computes ``gate_compress * coarse[cube(t)] + gate_fine * fine[src(t)]`` for every compact + token ``t``, where ``cube(t)`` is the cube holding the token and ``src(t)`` is its padded + slot when the fine output is tiled, or ``t`` itself otherwise. + + Args: + fine: ``[batch, padded_len or seq_len, heads, head_dim]`` fine-stage output; only + head_dim has to be contiguous, so head-major kernel outputs are accepted as views. + coarse: ``[batch, num_cubes, heads, head_dim]`` coarse-stage output per cube. + gate_compress: ``[batch, seq_len, heads, head_dim]`` gate for the coarse term. + gate_fine: Optional gate for the fine term with the same shape as ``gate_compress``. + untile_index: ``[seq_len]`` long tensor with the padded slot of every compact token. + cube_size: Tokens per cube. + fine_is_tiled: Whether ``fine`` is in the padded tile-major layout. + + Returns: + ``[batch, seq_len, heads, head_dim]`` blended output in the gate dtype. + """ + if not gate_compress.is_cuda: + return _blend_coarse_fine_torch( + fine, + coarse, + gate_compress, + gate_fine, + untile_index, + cube_size=cube_size, + fine_is_tiled=fine_is_tiled, + ) + if fine.stride(3) != 1: + raise ValueError("head_dim must be contiguous in the VSA fine output") + + coarse = coarse.contiguous() + gate_compress = gate_compress.contiguous() + batch_size, seq_len, num_heads, head_dim = gate_compress.shape + row_width = num_heads * head_dim + out = torch.empty_like(gate_compress) + block, num_chunks, num_warps = _row_launch_config(row_width, batch_size * seq_len) + _blend_coarse_fine_kernel[(batch_size * seq_len, num_chunks)]( + fine, + coarse, + gate_compress, + gate_compress if gate_fine is None else gate_fine.contiguous(), + untile_index, + out, + seq_len, + coarse.shape[1], + fine.stride(0), + fine.stride(1), + fine.stride(2), + ROW=row_width, + HEAD_DIM=head_dim, + CUBE=cube_size, + BLOCK=block, + HAS_GATE_FINE=gate_fine is not None, + FINE_IS_TILED=fine_is_tiled, + num_warps=num_warps, + ) + return out + + +__all__ = [ + "blend_coarse_fine", + "sort_last_dim", + "tile_and_pool_cubes", +] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/metadata.py b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/metadata.py index de88e3b6d282..31916e397436 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/metadata.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/metadata.py @@ -27,6 +27,7 @@ VSA_TILE_SIZE: Tuple[int, int, int] = (4, 4, 4) VSA_BLOCK_SIZE = VSA_TILE_SIZE[0] * VSA_TILE_SIZE[1] * VSA_TILE_SIZE[2] _DEFAULT_MAX_CACHED_SHAPES = 16 +_BITS_PER_WORD = 32 def _get_tile_partition_indices( @@ -80,29 +81,40 @@ def _construct_variable_block_sizes( return sizes.reshape(-1).to(torch.long) +def _pack_valid_slots(valid_slots: torch.BoolTensor) -> torch.Tensor: + """Pack a padded-slot validity mask into little-endian 32-bit words.""" + bit_weights = 1 << torch.arange(_BITS_PER_WORD, dtype=torch.int64, device=valid_slots.device) + words = (valid_slots.view(-1, _BITS_PER_WORD).to(torch.int64) * bit_weights).sum(dim=-1) + return words.to(torch.uint32) + + @dataclass(frozen=True, slots=True) class VSAMetadata: - """Per-step policy and shape metadata required by the VSA sparse path.""" + """Per-step policy and shape metadata required by the VSA sparse path. + + ``tile_source_index`` maps every padded slot to its compact token (``-1`` for padding), + ``untile_idx`` maps every compact token back to its padded slot, + ``variable_block_sizes`` counts the valid tokens per cube, and ``kv_valid_words`` packs + the padded valid-token mask into 32-bit words for the block-sparse kernels. + """ current_timestep: int vsa_sparsity: float num_cubes: int padded_seq_length: int variable_block_sizes: torch.LongTensor - kv_token_mask: torch.BoolTensor - non_pad_index: torch.LongTensor - gather_idx: torch.LongTensor + tile_source_index: torch.LongTensor untile_idx: torch.LongTensor + kv_valid_words: torch.Tensor class _VSAShapeMetadata(TypedDict): num_cubes: int padded_seq_length: int variable_block_sizes: torch.LongTensor - kv_token_mask: torch.BoolTensor - non_pad_index: torch.LongTensor - gather_idx: torch.LongTensor + tile_source_index: torch.LongTensor untile_idx: torch.LongTensor + kv_valid_words: torch.Tensor class VSAMetadataBuilder: @@ -147,20 +159,18 @@ def _build_metadata( local_offsets < variable_block_sizes.unsqueeze(1) ] + tile_source_index = torch.full((padded_seq_length,), -1, dtype=torch.long, device=device) + tile_source_index[non_pad_index] = gather_idx untile_idx = torch.empty(total_seq_length, dtype=torch.long, device=device) untile_idx[gather_idx] = non_pad_index - kv_token_mask = torch.zeros(padded_seq_length, dtype=torch.bool, device=device) - kv_token_mask[non_pad_index] = True - return _VSAShapeMetadata( num_cubes=num_cubes, padded_seq_length=padded_seq_length, variable_block_sizes=variable_block_sizes, - kv_token_mask=kv_token_mask, - non_pad_index=non_pad_index, - gather_idx=gather_idx, + tile_source_index=tile_source_index, untile_idx=untile_idx, + kv_valid_words=_pack_valid_slots(tile_source_index >= 0), ) def build( diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/predictor.py b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/predictor.py index 800bb7fdaf80..97db41c5cf9b 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/predictor.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/predictor.py @@ -16,7 +16,6 @@ """Shared Video Sparse Attention prediction and post-processing.""" from dataclasses import dataclass, field -from functools import cache from math import ceil from typing import Optional @@ -24,6 +23,7 @@ from .....attention.backends.interface import PredefinedAttentionMask from .....attention.backends.sparse.params import BlockSparseForwardInputs +from .kernels import blend_coarse_fine, sort_last_dim, tile_and_pool_cubes from .metadata import ( _DEFAULT_MAX_CACHED_SHAPES, VSA_BLOCK_SIZE, @@ -31,58 +31,22 @@ get_vsa_forward_context, ) -_BITS_PER_WORD = 32 _SIGNED_INT32_MAX = torch.iinfo(torch.int32).max -def _mean_pool_cubes( - x_tiled: torch.Tensor, - variable_block_sizes: torch.LongTensor, - prod_tile: int, - num_cubes: int, -) -> torch.Tensor: - batch_size, _padded, num_heads, head_dim = x_tiled.shape - x_cubes = x_tiled.view(batch_size, num_cubes, prod_tile, num_heads, head_dim) - # FP32 accumulation avoids perturbing the coarse softmax when inputs are BF16. - x_sum = x_cubes.float().sum(dim=2) - valid_counts = variable_block_sizes.float().clamp(min=1).view(1, num_cubes, 1, 1) - return (x_sum / valid_counts).to(x_tiled.dtype) - - -class VSAPreprocessor: - """Convert compact BSHD tensors between sequence-major and tile-major order.""" - - @staticmethod - def tile( - x: torch.Tensor, - non_pad_index: torch.LongTensor, - gather_idx: torch.LongTensor, - padded_seq_len: int, - ) -> torch.Tensor: - # index_select + index_copy_ keeps this path traceable by torch.compile. - batch_size, _seq_len, num_heads, head_dim = x.shape - x_valid = x.index_select(1, gather_idx) - x_padded = x.new_zeros(batch_size, padded_seq_len, num_heads, head_dim) - x_padded.index_copy_(1, non_pad_index, x_valid) - return x_padded - - @staticmethod - def untile( - x: torch.Tensor, - untile_idx: torch.LongTensor, - ) -> torch.Tensor: - return torch.index_select(x, 1, untile_idx) - - @dataclass(frozen=True, slots=True, kw_only=True, eq=False) class VSAPostProcessContext: - """Per-call tensors needed after the backend executes the fine stage.""" + """Per-call tensors needed after the backend executes the fine stage. + + ``coarse_output`` stays per cube (``[batch, num_cubes, heads, head_dim]``); the + post-process gathers it to compact token order together with the fine output. + """ coarse_output: torch.Tensor = field(repr=False) gate_compress: torch.Tensor = field(repr=False) gate_fine: Optional[torch.Tensor] = field(default=None, repr=False) - untile_idx: Optional[torch.LongTensor] = field(default=None, repr=False) - output_shape: tuple[int, int, int, int] + untile_idx: torch.LongTensor = field(repr=False) + fine_is_tiled: bool @dataclass(frozen=True, slots=True, kw_only=True, eq=False) @@ -120,76 +84,58 @@ def __init__(self, max_cached_shapes: int = _DEFAULT_MAX_CACHED_SHAPES) -> None: def from_selected_blocks( self, selected_blocks: torch.Tensor, - kv_valid_bits: torch.Tensor, + kv_valid_words: torch.Tensor, ) -> BlockSparseForwardInputs: + """Build the BSR carrier for one prediction. + + Args: + selected_blocks: ``[batch, kv_heads, q_blocks, blocks_per_row]`` int32 selected + KV cube per query cube, in any order within a row. + kv_valid_words: ``[words]`` uint32 packed valid-token mask of the padded + sequence, shared by every batch entry. + """ batch_size, num_kv_heads, num_q_blocks, blocks_per_row = map(int, selected_blocks.shape) - key = ( - selected_blocks.device, - batch_size, - num_kv_heads, - num_q_blocks, - blocks_per_row, - ) + key = (selected_blocks.device, batch_size, num_kv_heads, num_q_blocks, blocks_per_row) block_indptr = self._indptr_cache.get(key) if block_indptr is None: - if len(self._indptr_cache) >= self._max_cached_shapes: - raise RuntimeError( - "VSA route cache reached its " - f"{self._max_cached_shapes}-shape limit; restart the pipeline or " - "reuse a configured resolution/frame profile" - ) - if selected_blocks.is_cuda and torch.cuda.is_current_stream_capturing(): - raise RuntimeError( - "VSA route cache miss during CUDA Graph capture; " - "run an eager warmup with the same selected-block shape first" - ) - total_entries = batch_size * num_kv_heads * num_q_blocks * blocks_per_row - if total_entries > _SIGNED_INT32_MAX: - raise OverflowError("VSA route offsets must fit in signed int32") - row_offsets = torch.arange( - num_q_blocks + 1, - dtype=torch.int32, - device=selected_blocks.device, - ).reshape(1, 1, -1) - head_offsets = torch.arange( - batch_size * num_kv_heads, - dtype=torch.int32, - device=selected_blocks.device, - ).reshape(batch_size, num_kv_heads, 1) - block_indptr = ( - head_offsets * (num_q_blocks * blocks_per_row) + row_offsets * blocks_per_row - ).contiguous() + block_indptr = self._build_block_indptr(*key) self._indptr_cache[key] = block_indptr return BlockSparseForwardInputs( q_block_size=VSA_BLOCK_SIZE, kv_block_size=VSA_BLOCK_SIZE, max_blocks_per_row=blocks_per_row, block_indptr=block_indptr, - block_indices=torch.sort(selected_blocks, dim=-1).values.reshape(-1).contiguous(), - kv_valid_bits=kv_valid_bits, + block_indices=sort_last_dim(selected_blocks).reshape(-1), + kv_valid_bits=kv_valid_words.unsqueeze(0).expand(batch_size, -1).contiguous(), ) - -@cache -def _get_bit_weights(device: torch.device) -> torch.Tensor: - bit_positions = torch.arange(_BITS_PER_WORD, dtype=torch.int64, device=device) - return torch.bitwise_left_shift(torch.ones_like(bit_positions), bit_positions) - - -def _pack_kv_token_mask(kv_token_mask: torch.Tensor, batch_size: int) -> torch.Tensor: - if kv_token_mask.ndim == 1: - batched_mask = kv_token_mask.unsqueeze(0).expand(batch_size, -1) - else: - batched_mask = kv_token_mask - seq_len_kv = int(batched_mask.shape[1]) - padded_length = ceil(seq_len_kv / _BITS_PER_WORD) * _BITS_PER_WORD - if padded_length != seq_len_kv: - batched_mask = torch.nn.functional.pad(batched_mask, (0, padded_length - seq_len_kv)) - words = ( - batched_mask.reshape(batch_size, -1, _BITS_PER_WORD).to(torch.int64) - * _get_bit_weights(kv_token_mask.device) - ).sum(dim=-1) - return words.to(torch.uint32).contiguous() + def _build_block_indptr( + self, + device: torch.device, + batch_size: int, + num_kv_heads: int, + num_q_blocks: int, + blocks_per_row: int, + ) -> torch.Tensor: + if len(self._indptr_cache) >= self._max_cached_shapes: + raise RuntimeError( + "VSA route cache reached its " + f"{self._max_cached_shapes}-shape limit; restart the pipeline or " + "reuse a configured resolution/frame profile" + ) + if device.type == "cuda" and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "VSA route cache miss during CUDA Graph capture; " + "run an eager warmup with the same selected-block shape first" + ) + if batch_size * num_kv_heads * num_q_blocks * blocks_per_row > _SIGNED_INT32_MAX: + raise OverflowError("VSA route offsets must fit in signed int32") + row_offsets = torch.arange(num_q_blocks + 1, dtype=torch.int32, device=device) + head_offsets = torch.arange(batch_size * num_kv_heads, dtype=torch.int32, device=device) + return ( + head_offsets.reshape(batch_size, num_kv_heads, 1) * (num_q_blocks * blocks_per_row) + + row_offsets.reshape(1, 1, -1) * blocks_per_row + ).contiguous() class VSAPredictor: @@ -281,62 +227,52 @@ def predict( raise ValueError("VSA batch_size and seq_len must match the compact QKV tensors.") metadata = metadata or self.get_metadata() - padded_len = metadata.padded_seq_length num_cubes = metadata.num_cubes cur_topk = max(1, ceil((1.0 - metadata.vsa_sparsity) * num_cubes)) - q_tiled = VSAPreprocessor.tile(q, metadata.non_pad_index, metadata.gather_idx, padded_len) - k_tiled = VSAPreprocessor.tile(k, metadata.non_pad_index, metadata.gather_idx, padded_len) - v_tiled = VSAPreprocessor.tile(v, metadata.non_pad_index, metadata.gather_idx, padded_len) - q_coarse = _mean_pool_cubes( - q_tiled, metadata.variable_block_sizes, VSA_BLOCK_SIZE, num_cubes - ) - k_coarse = _mean_pool_cubes( - k_tiled, metadata.variable_block_sizes, VSA_BLOCK_SIZE, num_cubes - ) - v_coarse = _mean_pool_cubes( - v_tiled, metadata.variable_block_sizes, VSA_BLOCK_SIZE, num_cubes + (q_tiled, q_coarse), (k_tiled, k_coarse), (v_tiled, v_coarse) = ( + tile_and_pool_cubes( + x, + metadata.tile_source_index, + metadata.variable_block_sizes, + cube_size=VSA_BLOCK_SIZE, + ) + for x in (q, k, v) ) + coarse_scores = torch.einsum("bnhd,bmhd->bhnm", q_coarse, k_coarse) * q.shape[-1] ** -0.5 coarse_probs = coarse_scores.softmax(dim=-1) coarse_output = torch.einsum("bhnm,bmhd->bnhd", coarse_probs, v_coarse) - topk_indices = coarse_probs.topk(cur_topk, dim=-1).indices.to(torch.int32) - coarse_output_tiled = ( - coarse_output.unsqueeze(2) - .expand(batch_size, num_cubes, VSA_BLOCK_SIZE, q.shape[2], q.shape[3]) - .reshape(batch_size, padded_len, q.shape[2], q.shape[3]) - ) - coarse_output_compact = VSAPreprocessor.untile(coarse_output_tiled, metadata.untile_idx) + # BSR routes are re-sorted by cube index, so their value order is not needed; other + # consumers keep receiving the selected cubes in descending probability order. + topk_indices = coarse_probs.topk( + cur_topk, dim=-1, sorted=not produce_block_sparse_inputs + ).indices.to(torch.int32) block_sparse_inputs = None if use_sparse_fine and produce_block_sparse_inputs: - kv_valid_bits = _pack_kv_token_mask(metadata.kv_token_mask, batch_size) block_sparse_inputs = self._route_builder.from_selected_blocks( topk_indices, - kv_valid_bits, + metadata.kv_valid_words, ) - effective_q = q_tiled if use_sparse_fine else q - effective_k = k_tiled if use_sparse_fine else k - effective_v = v_tiled if use_sparse_fine else v - effective_seq_len = padded_len if use_sparse_fine else seq_len return VSAForwardInputs( - q=effective_q, - k=effective_k, - v=effective_v, + q=q_tiled if use_sparse_fine else q, + k=k_tiled if use_sparse_fine else k, + v=v_tiled if use_sparse_fine else v, batch_size=batch_size, - seq_len=effective_seq_len, + seq_len=metadata.padded_seq_length if use_sparse_fine else seq_len, block_sparse_inputs=block_sparse_inputs, topk_indices=topk_indices, variable_block_sizes=metadata.variable_block_sizes, cur_topk=cur_topk, num_cubes=num_cubes, post_context=VSAPostProcessContext( - coarse_output=coarse_output_compact, + coarse_output=coarse_output, gate_compress=gate_compress, gate_fine=gate_fine, - untile_idx=metadata.untile_idx if use_sparse_fine else None, - output_shape=tuple(q.shape), + untile_idx=metadata.untile_idx, + fine_is_tiled=use_sparse_fine, ), ) @@ -346,16 +282,17 @@ def vsa_post_process(output: torch.Tensor, inputs: VSAForwardInputs) -> torch.Te context = inputs.post_context fine_output = output.reshape( - inputs.batch_size, - inputs.seq_len, - context.output_shape[2], - context.output_shape[3], + inputs.batch_size, inputs.seq_len, *context.gate_compress.shape[2:] + ) + return blend_coarse_fine( + fine_output, + context.coarse_output, + context.gate_compress, + context.gate_fine, + context.untile_idx, + cube_size=VSA_BLOCK_SIZE, + fine_is_tiled=context.fine_is_tiled, ) - if context.untile_idx is not None: - fine_output = VSAPreprocessor.untile(fine_output, context.untile_idx) - if context.gate_fine is not None: - fine_output = context.gate_fine * fine_output - return context.gate_compress * context.coarse_output + fine_output __all__ = [ diff --git a/tests/unittest/_torch/visual_gen/test_attention_vsa.py b/tests/unittest/_torch/visual_gen/test_attention_vsa.py index 82b90fd75dcd..d2cb2c663390 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_vsa.py +++ b/tests/unittest/_torch/visual_gen/test_attention_vsa.py @@ -29,18 +29,21 @@ from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl import CuTeDSLAttention from tensorrt_llm._torch.visual_gen.attention_backend.interface import AttentionTensorLayout from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import backend as vsa_backend +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import kernels as vsa_kernels +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import predictor as vsa_predictor from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa.backend import ( VSACuTeDSLAttention, VSATrtllmAttention, ) +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa.kernels import tile_and_pool_cubes from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa.metadata import ( + VSA_BLOCK_SIZE, VSAMetadataBuilder, set_vsa_forward_context, ) from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa.predictor import ( VSAForwardInputs, VSAPredictor, - VSAPreprocessor, ) from tensorrt_llm._torch.visual_gen.attention_backend.trtllm import TrtllmAttention from tensorrt_llm._torch.visual_gen.attention_backend.utils import create_attention @@ -210,7 +213,7 @@ def test_vsa_predictor_dense_fallback_keeps_compact_qkv_and_no_block_inputs() -> assert inputs.v is v assert inputs.seq_len == 80 assert inputs.block_sparse_inputs is None - assert inputs.post_context.untile_idx is None + assert not inputs.post_context.fine_is_tiled def test_vsa_shared_post_process_restores_shape_and_applies_gates() -> None: @@ -237,7 +240,11 @@ def test_vsa_shared_post_process_restores_shape_and_applies_gates() -> None: output = vsa_backend.vsa_post_process(fine_output, inputs) - expected = 2.0 * inputs.post_context.coarse_output + 0.5 * fine_output + coarse_per_token = inputs.post_context.coarse_output.index_select( + 1, metadata.untile_idx // VSA_BLOCK_SIZE + ) + expected = 2.0 * coarse_per_token + 0.5 * fine_output + assert inputs.post_context.coarse_output.shape == (1, metadata.num_cubes, 1, 8) assert output.shape == q.shape torch.testing.assert_close(output, expected) @@ -549,13 +556,24 @@ def test_vsa_metadata_builder_reuses_shape_tensors_with_live_step_policy() -> No assert first is not second assert (first.current_timestep, first.vsa_sparsity) == (3, 0.25) assert (second.current_timestep, second.vsa_sparsity) == (4, 0.75) - assert second.gather_idx is first.gather_idx + assert second.tile_source_index is first.tile_source_index assert first.num_cubes == 27 builder.clear() rebuilt = builder.build(current_timestep=5, vsa_sparsity=0.5, **build_args) - assert rebuilt.gather_idx is not first.gather_idx + assert rebuilt.tile_source_index is not first.tile_source_index + + +def test_vsa_metadata_exposes_tile_source_index_and_packed_kv_words() -> None: + metadata = _make_vsa_metadata() + + source = metadata.tile_source_index + assert source.shape == (metadata.padded_seq_length,) + assert int((source >= 0).sum()) == 80 + assert torch.equal(source[metadata.untile_idx], torch.arange(80)) + assert metadata.kv_valid_words.dtype == torch.uint32 + assert metadata.kv_valid_words.tolist() == [0xFFFFFFFF, 0xFFFFFFFF, 0xFFFF, 0] def test_vsa_graph_stable_caches_bound_shape_profiles() -> None: @@ -571,11 +589,11 @@ def test_vsa_graph_stable_caches_bound_shape_profiles() -> None: builder.build(raw_latent_shape=(8, 4, 4), **build_args) route_builder = VSAPredictor(num_heads=1, max_cached_shapes=1)._route_builder - kv_valid_bits = torch.ones((1, 1), dtype=torch.uint32) - route_builder.from_selected_blocks(torch.zeros((1, 1, 1, 1), dtype=torch.int32), kv_valid_bits) + kv_valid_words = torch.ones((1,), dtype=torch.uint32) + route_builder.from_selected_blocks(torch.zeros((1, 1, 1, 1), dtype=torch.int32), kv_valid_words) with pytest.raises(RuntimeError, match="route cache reached its 1-shape limit"): route_builder.from_selected_blocks( - torch.zeros((1, 1, 2, 1), dtype=torch.int32), kv_valid_bits + torch.zeros((1, 1, 2, 1), dtype=torch.int32), kv_valid_words ) @@ -590,7 +608,7 @@ def test_vsa_graph_stable_caches_bound_shape_profiles() -> None: ids=["clean_8x8x8", "ragged_9x9x9", "wan720p_21x45x80"], ) def test_vsa_tile_untile_roundtrip(latent_shape): - """VSAPreprocessor.tile then .untile must losslessly reproduce the input.""" + """Tiling then untiling must reproduce the input, and pooled cubes must be token means.""" device = torch.device("cuda") dtype = torch.bfloat16 torch.manual_seed(0) @@ -609,24 +627,20 @@ def test_vsa_tile_untile_roundtrip(latent_shape): x = torch.randn(B, seq_len, H, D, device=device, dtype=dtype) - x_tiled = VSAPreprocessor.tile( + x_tiled, x_pooled = tile_and_pool_cubes( x, - meta.non_pad_index, - meta.gather_idx, - meta.padded_seq_length, + meta.tile_source_index, + meta.variable_block_sizes, + cube_size=VSA_BLOCK_SIZE, ) - pad_mask = torch.ones(meta.padded_seq_length, dtype=torch.bool, device=device) - pad_mask[meta.non_pad_index] = False + pad_mask = meta.tile_source_index < 0 if pad_mask.any(): assert x_tiled[:, pad_mask, :, :].abs().max().item() == 0.0, ( - "tile() must zero-fill padded positions" + "tiling must zero-fill padded positions" ) - x_roundtrip = VSAPreprocessor.untile( - x_tiled, - meta.untile_idx, - ) + x_roundtrip = x_tiled.index_select(1, meta.untile_idx) assert x_roundtrip.shape == x.shape, ( f"shape mismatch after tile/untile: {x_roundtrip.shape} vs {x.shape}" @@ -636,6 +650,109 @@ def test_vsa_tile_untile_roundtrip(latent_shape): f"max_diff={(x_roundtrip - x).abs().max().item():.3e}" ) + expected_pooled = x_tiled.view(B, meta.num_cubes, VSA_BLOCK_SIZE, H, D).float().sum(dim=2) + expected_pooled = expected_pooled / meta.variable_block_sizes.view(1, -1, 1, 1).float() + torch.testing.assert_close(x_pooled, expected_pooled.to(dtype), rtol=1e-2, atol=1e-2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="VSA needs CUDA") +def test_vsa_predictor_kernels_match_torch_fallback(monkeypatch: pytest.MonkeyPatch) -> None: + """The Triton path and the PyTorch fallback must produce the same envelope and output.""" + device = torch.device("cuda") + metadata = VSAMetadataBuilder().build( + current_timestep=0, + raw_latent_shape=(9, 9, 9), + patch_size=(1, 1, 1), + vsa_sparsity=0.75, + device=device, + ) + torch.manual_seed(0) + q, k, v = (torch.randn(2, 729, 4, 32, device=device) for _ in range(3)) + gate = torch.randn_like(q) + fine_output = torch.randn(2, metadata.padded_seq_length, 4, 32, device=device) + call_args = { + "batch_size": 2, + "seq_len": 729, + "seq_len_kv": 729, + "attention_mask": PredefinedAttentionMask.FULL, + "gate_compress": gate, + "gate_fine": gate, + "use_sparse_fine": True, + "produce_block_sparse_inputs": True, + "metadata": metadata, + } + + with_kernels = VSAPredictor(num_heads=4).predict(q, k, v, **call_args) + output_with_kernels = vsa_backend.vsa_post_process(fine_output, with_kernels) + + monkeypatch.setattr( + vsa_predictor, "tile_and_pool_cubes", vsa_kernels._tile_and_pool_cubes_torch + ) + monkeypatch.setattr(vsa_predictor, "sort_last_dim", vsa_kernels._sort_last_dim_torch) + monkeypatch.setattr(vsa_predictor, "blend_coarse_fine", vsa_kernels._blend_coarse_fine_torch) + fallback = VSAPredictor(num_heads=4).predict(q, k, v, **call_args) + output_fallback = vsa_backend.vsa_post_process(fine_output, fallback) + + for name in ("q", "k", "v"): + assert torch.equal(getattr(with_kernels, name), getattr(fallback, name)), name + assert torch.equal( + with_kernels.block_sparse_inputs.block_indices, fallback.block_sparse_inputs.block_indices + ) + torch.testing.assert_close( + with_kernels.post_context.coarse_output, fallback.post_context.coarse_output + ) + torch.testing.assert_close(output_with_kernels, output_fallback) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="VSA needs CUDA") +def test_vsa_predictor_replays_inside_cuda_graph() -> None: + device = torch.device("cuda") + metadata = VSAMetadataBuilder().build( + current_timestep=0, + raw_latent_shape=(8, 8, 8), + patch_size=(1, 1, 1), + vsa_sparsity=0.5, + device=device, + ) + predictor = VSAPredictor(num_heads=2) + torch.manual_seed(0) + q = torch.randn(1, 512, 2, 16, device=device, dtype=torch.bfloat16) + gate = torch.randn_like(q) + + def run() -> tuple[torch.Tensor, torch.Tensor]: + inputs = predictor.predict( + q, + q, + q, + batch_size=1, + seq_len=512, + seq_len_kv=512, + attention_mask=PredefinedAttentionMask.FULL, + gate_compress=gate, + gate_fine=None, + use_sparse_fine=True, + produce_block_sparse_inputs=True, + metadata=metadata, + ) + return vsa_backend.vsa_post_process( + inputs.q, inputs + ), inputs.block_sparse_inputs.block_indices + + eager_output, eager_routes = run() + side_stream = torch.cuda.Stream() + side_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side_stream): + run() + torch.cuda.current_stream().wait_stream(side_stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_output, graph_routes = run() + graph.replay() + torch.cuda.synchronize() + + assert torch.equal(graph_output, eager_output) + assert torch.equal(graph_routes, eager_routes) + @pytest.mark.skipif(not torch.cuda.is_available(), reason="kernel test needs CUDA") def test_cute_kernel_matches_dense_at_full_topk(): diff --git a/tests/unittest/_torch/visual_gen/test_attention_vsa_kernels.py b/tests/unittest/_torch/visual_gen/test_attention_vsa_kernels.py new file mode 100644 index 000000000000..c63df0587b19 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_attention_vsa_kernels.py @@ -0,0 +1,195 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""VSA predictor kernels (tile + cube mean, row sort, coarse/fine blend) against references. + +Every kernel is checked on CUDA (Triton path) and CPU (PyTorch fallback) over cube layouts +with ragged fills and head layouts whose row width is not a power of two. +""" + +import pytest +import torch + +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa.kernels import ( + blend_coarse_fine, + sort_last_dim, + tile_and_pool_cubes, +) + +CUBE_SIZE = 64 + +_DEVICES = ["cpu"] + (["cuda"] if torch.cuda.is_available() else []) + +_HEAD_LAYOUTS = [ + pytest.param(6, 4, 32, id="row128"), + pytest.param(5, 5, 21, id="row105_odd"), + pytest.param(30, 40, 128, id="wan14b_row5120"), + pytest.param(3, 64, 128, id="row8192"), +] + + +def _random_cube_layout(num_cubes: int, device: str, generator: torch.Generator): + """Assign consecutive compact tokens to cubes with a random fill per cube. + + Returns the padded-slot source index (-1 for pad slots), the valid count per cube, the + padded slot of every compact token, and the compact sequence length. + """ + counts = torch.randint(1, CUBE_SIZE + 1, (num_cubes,), generator=generator) + counts[0] = CUBE_SIZE + counts[-1] = CUBE_SIZE // 2 + seq_len = int(counts.sum()) + tile_source_index = torch.full((num_cubes * CUBE_SIZE,), -1, dtype=torch.long) + untile_index = torch.empty(seq_len, dtype=torch.long) + token = 0 + for cube, count in enumerate(counts.tolist()): + slots = torch.arange(cube * CUBE_SIZE, cube * CUBE_SIZE + count) + tile_source_index[slots] = torch.arange(token, token + count) + untile_index[token : token + count] = slots + token += count + return tile_source_index.to(device), counts.to(device), untile_index.to(device), seq_len + + +def _reference_tile_and_pool(x, tile_source_index, counts, num_cubes): + batch, _, heads, head_dim = x.shape + tiled = x.index_select(1, tile_source_index.clamp(min=0)) + valid = (tile_source_index >= 0).view(1, -1, 1, 1) + tiled = torch.where(valid, tiled, torch.zeros((), dtype=x.dtype, device=x.device)) + pooled = tiled.view(batch, num_cubes, CUBE_SIZE, heads, head_dim).float().sum(dim=2) + pooled = pooled / counts.view(1, -1, 1, 1).float() + return tiled, pooled.to(x.dtype) + + +def _reference_blend(fine, coarse, gate_compress, gate_fine, untile_index, fine_is_tiled): + coarse_per_token = coarse.index_select(1, untile_index // CUBE_SIZE) + fine_compact = fine.index_select(1, untile_index) if fine_is_tiled else fine + if gate_fine is not None: + fine_compact = gate_fine * fine_compact + return gate_compress * coarse_per_token + fine_compact + + +@pytest.mark.parametrize("device", _DEVICES) +@pytest.mark.parametrize(("num_cubes", "heads", "head_dim"), _HEAD_LAYOUTS) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32], ids=["bf16", "fp32"]) +def test_tile_and_pool_cubes_matches_reference(device, num_cubes, heads, head_dim, dtype): + generator = torch.Generator().manual_seed(num_cubes * 31 + heads) + source, counts, _untile, seq_len = _random_cube_layout(num_cubes, device, generator) + x = torch.randn(2, seq_len, heads, head_dim, device=device, dtype=dtype) + + tiled, pooled = tile_and_pool_cubes(x, source, counts, cube_size=CUBE_SIZE) + + ref_tiled, ref_pooled = _reference_tile_and_pool(x, source, counts, num_cubes) + assert tiled.shape == (2, num_cubes * CUBE_SIZE, heads, head_dim) + assert pooled.shape == (2, num_cubes, heads, head_dim) + assert torch.equal(tiled, ref_tiled) + tolerance = 1e-5 if dtype == torch.float32 else 1e-2 + torch.testing.assert_close(pooled, ref_pooled, rtol=tolerance, atol=tolerance) + + +@pytest.mark.parametrize("device", _DEVICES) +def test_tile_and_pool_cubes_accepts_strided_sequence_layout(device): + """Q/K/V unbound from a packed [B, S, 3, H, D] tensor must work without a copy.""" + source, counts, _untile, seq_len = _random_cube_layout( + 8, device, torch.Generator().manual_seed(7) + ) + packed = torch.randn(2, seq_len, 3, 4, 32, device=device, dtype=torch.float16) + q = packed[:, :, 1] + assert not q.is_contiguous() + + tiled, pooled = tile_and_pool_cubes(q, source, counts, cube_size=CUBE_SIZE) + + ref_tiled, ref_pooled = _reference_tile_and_pool(q.contiguous(), source, counts, 8) + assert torch.equal(tiled, ref_tiled) + torch.testing.assert_close(pooled, ref_pooled, rtol=1e-3, atol=1e-3) + + +def test_tile_and_pool_cubes_rejects_split_head_dims(): + source, counts, _untile, seq_len = _random_cube_layout( + 2, "cpu", torch.Generator().manual_seed(1) + ) + x = torch.randn(1, seq_len, 32, 4).transpose(2, 3) + + with pytest.raises(ValueError, match="contiguous"): + tile_and_pool_cubes(x, source, counts, cube_size=CUBE_SIZE) + + +@pytest.mark.parametrize("device", _DEVICES) +@pytest.mark.parametrize("row_length", [1, 30, 144, 257, 2048, 5000]) +def test_sort_last_dim_matches_torch_sort(device, row_length): + generator = torch.Generator().manual_seed(row_length) + values = torch.randint(0, 4096, (2, 3, 5, row_length), generator=generator, dtype=torch.int32) + values = values.to(device) + + assert torch.equal(sort_last_dim(values), torch.sort(values, dim=-1).values) + + +@pytest.mark.parametrize("device", _DEVICES) +def test_blend_coarse_fine_reads_head_major_fine_output(device): + """Fine output stored as [B, H, S, D] (the CuTe layout) is consumed through its strides.""" + num_cubes, batch, heads, head_dim = 6, 2, 4, 32 + source, _counts, untile, seq_len = _random_cube_layout( + num_cubes, device, torch.Generator().manual_seed(11) + ) + fine = torch.randn(batch, heads, num_cubes * CUBE_SIZE, head_dim, device=device).transpose(1, 2) + assert not fine.is_contiguous() + coarse = torch.randn(batch, num_cubes, heads, head_dim, device=device) + gate_compress = torch.randn(batch, seq_len, heads, head_dim, device=device) + + out = blend_coarse_fine( + fine, coarse, gate_compress, None, untile, cube_size=CUBE_SIZE, fine_is_tiled=True + ) + + ref = _reference_blend(fine, coarse, gate_compress, None, untile, True) + torch.testing.assert_close(out, ref, rtol=1e-6, atol=1e-6) + + +def test_sort_last_dim_requires_int32(): + with pytest.raises(TypeError, match="int32"): + sort_last_dim(torch.zeros(2, 4, dtype=torch.int64)) + + +@pytest.mark.parametrize("device", _DEVICES) +@pytest.mark.parametrize("fine_is_tiled", [True, False], ids=["tiled_fine", "compact_fine"]) +@pytest.mark.parametrize("with_gate_fine", [True, False], ids=["gate_fine", "no_gate_fine"]) +@pytest.mark.parametrize( + ("heads", "head_dim", "dtype"), + [(5, 21, torch.float32), (40, 128, torch.bfloat16)], + ids=["row105_fp32", "wan14b_bf16"], +) +def test_blend_coarse_fine_matches_reference( + device, fine_is_tiled, with_gate_fine, heads, head_dim, dtype +): + num_cubes, batch = 6, 2 + source, _counts, untile, seq_len = _random_cube_layout( + num_cubes, device, torch.Generator().manual_seed(3) + ) + fine_len = num_cubes * CUBE_SIZE if fine_is_tiled else seq_len + fine = torch.randn(batch, fine_len, heads, head_dim, device=device, dtype=dtype) + coarse = torch.randn(batch, num_cubes, heads, head_dim, device=device, dtype=dtype) + gate_compress = torch.randn(batch, seq_len, heads, head_dim, device=device, dtype=dtype) + gate_fine = torch.randn_like(gate_compress) if with_gate_fine else None + + out = blend_coarse_fine( + fine, + coarse, + gate_compress, + gate_fine, + untile, + cube_size=CUBE_SIZE, + fine_is_tiled=fine_is_tiled, + ) + + ref = _reference_blend(fine, coarse, gate_compress, gate_fine, untile, fine_is_tiled) + assert out.shape == gate_compress.shape + tolerance = 1e-6 if dtype == torch.float32 else 2e-2 + torch.testing.assert_close(out, ref, rtol=tolerance, atol=tolerance)