From 889b3b0ee81d4d64121f663d3fc72ca3146eb7c8 Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Sat, 19 Sep 2026 05:22:07 +0000 Subject: [PATCH 01/10] refactor: align VisualGen sparse attention workflow Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- .../features/visualgen-sparse-attention.md | 19 +- docs/source/models/visual-generation.md | 13 +- .../visual_gen/attention_backend/__init__.py | 17 +- .../attention_backend/cute_dsl/__init__.py | 21 +- .../attention_backend/cute_dsl/fmha.py | 2 +- .../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 | 331 +++++++ .../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 | 74 +- tensorrt_llm/_torch/visual_gen/config.py | 11 +- .../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 | 54 +- .../_torch/visual_gen/pipeline_loader.py | 15 +- tensorrt_llm/visual_gen/args.py | 36 +- tensorrt_llm/visual_gen/sparse_attention.py | 4 +- .../test_lists/test-db/l0_b200.yml | 3 +- .../multi_gpu/test_ulysses_attention.py | 62 ++ .../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 | 892 ++++++++++++++++++ .../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, 2869 insertions(+), 1088 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/features/visualgen-sparse-attention.md b/docs/source/features/visualgen-sparse-attention.md index cb734cb9c888..289b7be9d17d 100644 --- a/docs/source/features/visualgen-sparse-attention.md +++ b/docs/source/features/visualgen-sparse-attention.md @@ -14,14 +14,14 @@ This feature is in **beta** stage. APIs, supported models, and optimization opti 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` | `VideoSparseAttentionConfig` | Supported (CUTEDSL) | +| `vsa` | `VideoSparseAttentionConfig` | Supported (`CUTEDSL`, `TRTLLM`) | | `sol_attn` | `SolAttentionConfig` | Supported (CUTEDSL, sm100/sm103) | ### Sol-Attn @@ -253,4 +253,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 use the same VisualGen-owned predictor implementation, one +instance per attention layer, 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/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index 399366579bff..866e68dfb5af 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -360,12 +360,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. @@ -379,7 +382,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), ), ) @@ -389,7 +392,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/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py b/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py index bdc9fe9f33c1..2c577f0e95bc 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py @@ -21,16 +21,7 @@ """ from .cudnn import CuDNNAttention -from .cute_dsl import ( - VSA_TILE_SIZE, - CuTeDSLAttention, - SolAttention, - VSAAttention, - VSAMetadata, - VSAMetadataBuilder, - get_vsa_forward_context, - set_vsa_forward_context, -) +from .cute_dsl import CuTeDSLAttention, SolAttention from .flash_attn4 import FlashAttn4Attention from .flashinfer import FlashInferAttention from .interface import AttentionBackend, AttentionTensorLayout @@ -40,7 +31,6 @@ from .vanilla import VanillaAttention __all__ = [ - "VSA_TILE_SIZE", "Attention2DAttention", "AttentionBackend", "AttentionTensorLayout", @@ -53,13 +43,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 e5dccb71a7ba..8c4ca507810d 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 @@ -16,33 +16,14 @@ 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) - sol_attn.py — SolAttention (Sol-Attn dynamic block routing, CuTe JIT + SDPA fallback) + sol_attn.py — SolAttention (Sol-Attn dynamic block routing, CuTe JIT + SDPA fallback) """ from .fmha import CuTeDSLAttention, _cute_dsl_import_error from .sol_attn import SolAttention -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", "SolAttention", ] 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 1d7edbbaa948..a5c96202a6d5 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 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..6968424dac95 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/backend.py @@ -0,0 +1,331 @@ +# 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, + ) + 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], + 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 80cedc79fa9c..b2087a4819d8 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) @@ -56,30 +55,37 @@ def get_visual_gen_attention_backend( - "CUDNN": cuDNN fused SDPA. Unquantized by default; quant_attention_config selects per-tensor FP8 or block-scaled MXFP8 (Blackwell). """ - # Lazy imports to avoid circular dependency - from .cudnn import CuDNNAttention - 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 elif backend_name == "CUDNN": + from .cudnn import CuDNNAttention + return CuDNNAttention else: # Default to VANILLA for maximum compatibility + from .vanilla import VanillaAttention + return VanillaAttention @@ -117,46 +123,60 @@ 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 + ) + sparse_algorithm = getattr(sparse_attention_config, "algorithm", None) + is_vsa = sparse_algorithm == "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 + elif sparse_algorithm == "sol_attn" and backend_name == "CUTEDSL": + from .cute_dsl.sol_attn import SolAttention + + attn_cls = SolAttention + kwargs["sparse_attention_config"] = sparse_attention_config + 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, cuDNN, 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: - sparse_algo = getattr(attention_config.sparse_attention_config, "algorithm", None) - if sparse_algo == "vsa": - from .cute_dsl.vsa import VSAAttention - - attn_cls = VSAAttention - kwargs["sparse_attention_config"] = attention_config.sparse_attention_config - elif sparse_algo == "sol_attn": - from .cute_dsl.sol_attn import SolAttention - - attn_cls = SolAttention - 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 b474090e6c19..6ea10f3f19d5 100644 --- a/tensorrt_llm/_torch/visual_gen/config.py +++ b/tensorrt_llm/_torch/visual_gen/config.py @@ -80,7 +80,14 @@ 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.""" + """Create state shared by attention layers in one model component. + + The state outlives individual forwards and CUDA Graph captures. It owns the + shape-keyed TRTLLM metadata cache, and the TRTLLM attention wrapper adds the + shared PrimTS block-sparse plan caches on demand so one static profile is + planned once per component instead of once per layer. Each model component + receives a distinct state and must not execute concurrent forwards. + """ return {"metadata_cache": {}} @@ -129,6 +136,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 @@ -202,6 +210,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 3a4bc10b33c0..358724667ef6 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 09ef502873f3..1c8ae1c2e6da 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 196289d8c18f..44b9834c7879 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 @@ -98,17 +113,17 @@ def __init__( base_backend = config.attention.backend _sa_cfg = config.attention.sparse_attention_config _sa_algo = getattr(_sa_cfg, "algorithm", None) if _sa_cfg is not None else None - _is_vsa = base_backend == "CUTEDSL" and _sa_algo == "vsa" + is_vsa = _sa_algo == "vsa" _is_sol_attn = base_backend == "CUTEDSL" and _sa_algo == "sol_attn" - separate_qkv_cross_attention = ( + is_separate_qkv_cross_attention = ( self.qkv_mode == QKVMode.SEPARATE_QKV and not separate_qkv_is_self_attention ) - # Cross-attention fallback: TRTLLM and CUTEDSL VSA are self-attn only. + # Cross-attention fallback: dense TRTLLM and every VSA backend are self-attn only. # Sol-Attn is absent by design; see SolAttention._can_serve. - if separate_qkv_cross_attention and (base_backend == "TRTLLM" or _is_vsa): + if is_separate_qkv_cross_attention and (base_backend == "TRTLLM" or is_vsa): 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( @@ -121,8 +136,8 @@ def __init__( # Every sparse algorithm here routes over the whole token sequence, so # none of them can be split across context-parallel ranks. - if (_is_vsa or _is_sol_attn) and cp_size > 1: - _algo_name = "VSA" if _is_vsa else "Sol-Attn" + if (is_vsa or _is_sol_attn) and cp_size > 1: + _algo_name = "VSA" if is_vsa else "Sol-Attn" raise ValueError( f"{_algo_name} needs the full token sequence per rank, so it is incompatible " f"with context parallelism (Attention2D/Ring, cp_size={cp_size}). Use " @@ -253,12 +268,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( @@ -650,6 +660,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 @@ -689,8 +700,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() @@ -746,6 +757,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 db2e11b86145..83b2b8a5dcda 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline_loader.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline_loader.py @@ -259,13 +259,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 1ccf02e1a90f..7540ec2d548e 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -126,7 +126,7 @@ class AttentionConfig(StrictBaseModel): status="prototype", description=( "Sparse attention recipe. Discriminated by algorithm: " - "skip_softmax (TRTLLM / CUTEDSL backends), vsa (CUTEDSL backend), " + "skip_softmax (TRTLLM / CUTEDSL backends), vsa (CUTEDSL / TRTLLM backends), " "or sol_attn (CUTEDSL backend)." ), ) @@ -224,7 +224,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"), "sol_attn": ("CUTEDSL",), }.get(algo) if supported_backends is None: @@ -240,25 +240,19 @@ def _validate_sparse_attention_config(self) -> "AttentionConfig": return self @model_validator(mode="after") - def _validate_cutedsl_quant_sparse_mutex(self) -> "AttentionConfig": - # VSA and Sol-Attn each replace the dense CuTeDSL path and cannot - # compose with quantized attention: create_attention swaps in their own - # backend class, which never consumes quant_attention_config, so the - # request would be silently ignored. SkipSoftmax is part of the dense - # path itself and can compose. - _replaces_dense_path = ("vsa", "sol_attn") - 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 in _replaces_dense_path - ): - raise ValueError( - f"CUTEDSL backend: quant_attention_config and " - f"'{self.sparse_attention_config.algorithm}' sparse_attention_config " - "are mutually exclusive (the CuTeDSLAttention dispatcher selects " - "either the dense path or that sparse 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 + + # VSA and Sol-Attn replace the dense attention path on every backend that + # serves them and never consume quant_attention_config, so accepting a + # quantization recipe would silently ignore user configuration. + # SkipSoftmax is part of the dense path itself and can compose. + algorithm = self.sparse_attention_config.algorithm + if algorithm == "vsa": + raise ValueError("VSA and quant_attention_config are mutually exclusive.") + if algorithm == "sol_attn": + raise ValueError("Sol-Attn 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 bd4c22e22395..929797363298 100644 --- a/tensorrt_llm/visual_gen/sparse_attention.py +++ b/tensorrt_llm/visual_gen/sparse_attention.py @@ -293,11 +293,13 @@ def to_sparse_params(self, **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 30da5481037a..a193a49a3ee5 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -300,12 +300,13 @@ 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_cute_dsl_sol_attn.py - unittest/_torch/visual_gen/test_attention_flashinfer.py - unittest/_torch/visual_gen/test_attention_trtllm_sage.py - unittest/_torch/visual_gen/test_attention_cudnn.py - 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 8415bd6bc52a..b36d780418f5 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. @@ -96,6 +111,53 @@ def run_test_in_distributed(world_size: int, test_fn: Callable, use_cuda: bool = ) +def test_forward_async_redistributes_vsa_gates(monkeypatch): + 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 96fba95677f5..aca1c29f961f 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 @@ -317,5 +317,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 6f3aef886476..0860209d1d3d 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 @@ -521,6 +521,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 a04eca96a0c1..dcd331423f2e 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 @@ -35,9 +35,9 @@ import torch.multiprocessing as mp from utils.llm_data import get_checkpoint -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 faa8d488dbbc..c26efea39fdb 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 @@ -723,18 +724,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 @@ -742,17 +744,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() @@ -779,12 +785,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) @@ -805,22 +812,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..be3576d85e98 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_attention_vsa.py @@ -0,0 +1,892 @@ +# 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_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) + + +@pytest.mark.parametrize( + ("is_self_attention", "expected_backend"), + ((True, "TRTLLM"), (False, "VANILLA")), +) +def test_plain_trtllm_separate_qkv_dispatches_by_attention_role( + monkeypatch: pytest.MonkeyPatch, + is_self_attention: bool, + expected_backend: str, +) -> 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=is_self_attention, + ) + + assert attention.attn_backend == expected_backend + + +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 25e55c4ada7a..19d5e0675440 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 c1c5e08aa397..e1232a9f2dbb 100644 --- a/tests/unittest/_torch/visual_gen/test_wan_vsa_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_wan_vsa_pipeline.py @@ -120,7 +120,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 325b263bd72a1cdd732bc0f52299f49f7e8bea23 Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Sat, 19 Sep 2026 05:22:40 +0000 Subject: [PATCH 02/10] 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 be3576d85e98..0c5d9f174110 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 @@ -190,7 +193,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: @@ -217,7 +220,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) @@ -535,13 +542,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: @@ -557,11 +575,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 ) @@ -576,7 +594,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) @@ -595,24 +613,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}" @@ -622,6 +636,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) From 76cd2f332b6bafa847f852a631074b6f9752f97c Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Sat, 19 Sep 2026 05:25:08 +0000 Subject: [PATCH 03/10] fix: tolerate backend-specific kwargs in VisualGen TRTLLM attention Restore the shared VisualGen backend contract in the TRTLLM wrapper: the attention module forwards a superset of keyword arguments and every backend ignores the ones it does not consume. Rejecting unknown names made TRTLLM the only backend that raised, and HunyuanVideo 1.5 and GLM-Image always pass key_padding_mask, so their TRTLLM paths failed. timestep is now an explicit parameter; other keyword arguments are accepted and ignored as before. Guard the SOL graph-phase resolution with torch.cuda.is_available() so the cpu_only SOL tests run on hosts without a GPU, and pin the SM version in the SkipSoftmax and SAGE combination test like the neighbouring int8 SAGE tests. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- .../visual_gen/attention_backend/trtllm.py | 15 ++++---- .../test_trtllm_attention_metadata.py | 38 +++++++++---------- .../_torch/visual_gen/test_visual_gen_args.py | 23 ++++++----- 3 files changed, 37 insertions(+), 39 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py b/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py index dde2e14094c3..1f85598d787a 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py @@ -281,6 +281,7 @@ def forward( attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL, seq_len_kv: Optional[int] = None, sparse_backend_args: Optional[SparseBackendForwardArgs] = None, + timestep: Optional[torch.Tensor] = None, **kwargs, ) -> torch.Tensor: """ @@ -307,18 +308,16 @@ def forward( 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. + timestep: Denoising timestep consumed by the sparse prediction hooks. + **kwargs: Backend-specific keyword arguments forwarded by the attention + module, such as ``key_padding_mask``; ignored here, as by the other + backends. The TRTLLM kernels have no key-padding input, so callers + that need padded keys must guard at the model level or select the + ``VANILLA`` backend. Returns: Output tensor [B, S, H*D] """ - timestep = kwargs.pop("timestep", None) - if kwargs: - unexpected_names = ", ".join(sorted(kwargs)) - raise TypeError( - f"Unexpected TRTLLM attention forward keyword arguments: {unexpected_names}" - ) - block_sparse_inputs = ( sparse_backend_args.block_sparse_inputs if sparse_backend_args is not None else None ) 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 23293bfe21f8..f901967bc417 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py @@ -7,7 +7,6 @@ import pytest import torch -from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask from tensorrt_llm._torch.attention.backends.sparse.params import ( BlockSparseForwardInputs, SparseBackendForwardArgs, @@ -221,29 +220,26 @@ def test_visual_gen_wrapper_does_not_define_its_own_prediction_lifecycle(): 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) +def test_forward_ignores_backend_specific_kwargs_like_other_backends(monkeypatch): + captured: dict = {} + _capture_core_forward(monkeypatch, captured) attention = _make_wrapper() + monkeypatch.setattr(attention, "support_fused_qkv", lambda: True, raising=False) + q = torch.randn(1, 4, 2, 8) - 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" + output = attention.forward( + q, + q, + q, + batch_size=1, + seq_len=4, + key_padding_mask=torch.ones(1, 4, dtype=torch.bool), + gate_compress=torch.ones(1, 4, 2, 8), ) - prepare_metadata.assert_not_called() - core_forward.assert_not_called() + + assert output.shape == (1, 4, 16) + assert captured["kwargs"] == {} + assert not hasattr(captured["forward_args"], "key_padding_mask") def test_forward_flattens_fused_qkv_without_copy(monkeypatch): 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 19d5e0675440..006cb38fd73e 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py @@ -132,16 +132,19 @@ def test_vsa_and_quantization_are_mutually_exclusive(self, backend, quant_config ) 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), - ) + # int8 Q/K SAGE has a compiled cubin only on SM100; pin the SM so this + # combination check is host-independent (CI CPU stages have no GPU). + with patch("tensorrt_llm.visual_gen.args.get_sm_version", return_value=100): + 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" From 49be991e749085298fd875307082619978aa7652 Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Sat, 19 Sep 2026 05:30:02 +0000 Subject: [PATCH 04/10] refactor: share the sparse timestep schedule across VisualGen algorithms Reduce a denoising timestep to a dense-or-sparse phase in one place, tensorrt_llm/_torch/attention/backends/sparse/timestep_phase.py, and use it from the SkipSoftmax scheduler, the CuTeDSL Sol-Attn backend and the CUDA Graph key. Per-token timesteps reduce to their largest live value, so a Wan I2V conditioning frame at timestep zero no longer disables the configured dense prefix. The VisualGen TRTLLM attention wrapper owns the timestep schedule: it prepares the timestep as a host value once per eager call and reuses it during CUDA Graph capture, so SkipSoftmax with a timestep cutoff can be captured with a device timestep tensor, and it answers whether a layer runs sparse from the cutoff and dense_layers of its sparse parameters. Models register a single sparse_attn_phase CUDA Graph key for every sparse attention config with a timestep cutoff, resolved through the new BaseSparseAttentionConfig.resolve_disabled_until_timestep. The CuTeDSL Skip Softmax and Sol-Attn backends read that key from the CUDA Graph runner during capture, and the Sol-Attn dense-prefix decision uses the shared reduction instead of the scheduler classmethod, which is removed. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- .../features/visualgen-sparse-attention.md | 2 +- .../backends/sparse/skip_softmax/params.py | 34 ++-------- .../backends/sparse/timestep_phase.py | 56 ++++++++++++++++ .../attention_backend/cute_dsl/fmha.py | 2 +- .../attention_backend/cute_dsl/sol_attn.py | 6 +- .../visual_gen/attention_backend/trtllm.py | 53 +++++++++++++++ .../_torch/visual_gen/models/modeling.py | 59 ++++++---------- tensorrt_llm/visual_gen/sparse_attention.py | 15 +++++ .../attention/sparse/test_timestep_phase.py | 52 ++++++++++++++ .../sparse_attention/test_skip_softmax.py | 59 +++++++++++++--- .../test_attention_cute_dsl_sol_attn.py | 41 ++++-------- .../test_trtllm_attention_metadata.py | 67 +++++++++++++++++-- 12 files changed, 334 insertions(+), 112 deletions(-) create mode 100644 tensorrt_llm/_torch/attention/backends/sparse/timestep_phase.py create mode 100644 tests/unittest/_torch/attention/sparse/test_timestep_phase.py diff --git a/docs/source/features/visualgen-sparse-attention.md b/docs/source/features/visualgen-sparse-attention.md index 289b7be9d17d..a3e4d73cb1dd 100644 --- a/docs/source/features/visualgen-sparse-attention.md +++ b/docs/source/features/visualgen-sparse-attention.md @@ -14,7 +14,7 @@ This feature is in **beta** stage. APIs, supported models, and optimization opti 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, 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. +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. The VisualGen TRTLLM wrapper owns the timestep schedule those backends consult: it prepares the denoising timestep once per eager call and answers whether a layer runs sparse for it (dense layers and dense timestep phases do not). `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 diff --git a/tensorrt_llm/_torch/attention/backends/sparse/skip_softmax/params.py b/tensorrt_llm/_torch/attention/backends/sparse/skip_softmax/params.py index f3b54151d604..520a248898d3 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/skip_softmax/params.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/skip_softmax/params.py @@ -17,13 +17,13 @@ from dataclasses import dataclass, field, replace from typing import Any, Dict, Literal, Optional, Union -import torch from pydantic import ConfigDict, model_validator from pydantic import Field as PydanticField from tensorrt_llm.llmapi.utils import StrictBaseModel from ..params import SparseParams, SparseRuntimeParams +from ..timestep_phase import graph_phase_for_timestep _RESERVED_FORMULA_KEYS = frozenset({"formula", "target_sparsity"}) _SKIP_SOFTMAX_ALGORITHMS = frozenset({"skip_softmax", "softmax_skip"}) @@ -352,31 +352,6 @@ def _compute(phase: str, sparsity: Optional[float]) -> Optional[float]: disabled_until_timestep=disabled_until_timestep, ) - @staticmethod - def _as_float(value: Any) -> Optional[float]: - if value is None: - return None - if isinstance(value, torch.Tensor): - if value.numel() == 0: - return None - return float(value.flatten()[0].item()) - return float(value) - - @classmethod - def get_graph_phase_for_timestep( - cls, - timestep: Any, - *, - disabled_until_timestep: Optional[float], - ) -> Optional[int]: - """Return 1 after descending timesteps cross the cutoff, otherwise 0.""" - if disabled_until_timestep is None: - return None - timestep_value = cls._as_float(timestep) - if timestep_value is None: - return None - return int(timestep_value < disabled_until_timestep) - def get_runtime_params( self, *, @@ -394,7 +369,7 @@ def get_runtime_params( if runtime_params is None: runtime_params = SparseRuntimeParams() if graph_phase is None: - graph_phase = self.get_graph_phase_for_timestep( + graph_phase = graph_phase_for_timestep( timestep, disabled_until_timestep=self.disabled_until_timestep, ) @@ -418,3 +393,8 @@ class SkipSoftmaxParams(SparseParams): algorithm: Literal["skip_softmax"] = field(init=False, default="skip_softmax") scheduler: SkipSoftmaxScheduler = field(default_factory=SkipSoftmaxScheduler) uses_spcompress: bool = False + + @property + def disabled_until_timestep(self) -> Optional[float]: + """Normalized timestep cutoff owned by the scheduler.""" + return self.scheduler.disabled_until_timestep diff --git a/tensorrt_llm/_torch/attention/backends/sparse/timestep_phase.py b/tensorrt_llm/_torch/attention/backends/sparse/timestep_phase.py new file mode 100644 index 000000000000..a1bee88c418f --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/sparse/timestep_phase.py @@ -0,0 +1,56 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Timestep-driven dense-or-sparse phase shared by sparse attention algorithms.""" + +from __future__ import annotations + +import math +import numbers +from typing import Optional + +import torch + + +def timestep_to_float(timestep: object) -> Optional[float]: + """Return the denoising timestep as a finite float, or ``None`` when absent. + + Per-token timesteps, such as Wan I2V where the conditioning frame keeps + timestep zero, reduce to their largest value so the schedule stays dense + until every live token is below the cutoff. + """ + + if timestep is None: + return None + if isinstance(timestep, torch.Tensor): + if timestep.numel() == 0: + return None + 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 + + +def graph_phase_for_timestep( + timestep: object, + *, + disabled_until_timestep: Optional[float], +) -> Optional[int]: + """Return 0 for the dense prefix and 1 for the sparse suffix. + + ``None`` when the algorithm has no cutoff or the call carries no timestep; + CUDA Graph runners omit their phase key part in that case. + """ + + if disabled_until_timestep is None: + return None + value = timestep_to_float(timestep) + if value is None: + return None + return int(value < disabled_until_timestep) + + +__all__ = ["graph_phase_for_timestep", "timestep_to_float"] 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 a5c96202a6d5..edebb9da7078 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 @@ -91,7 +91,7 @@ def _resolve_skip_softmax_threshold_scale_factor( # read is a `.item()`, which is illegal while a graph is being captured. runtime_params = sparse_params.scheduler.get_runtime_params( timestep=timestep, - graph_phase=resolved_extra_key("skip_softmax_phase"), + graph_phase=resolved_extra_key("sparse_attn_phase"), ) 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/sol_attn.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py index 5d0fea9a1645..f69a278d9414 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py @@ -48,7 +48,7 @@ import torch -from tensorrt_llm._torch.attention.backends.sparse.skip_softmax import SkipSoftmaxScheduler +from tensorrt_llm._torch.attention.backends.sparse.timestep_phase import graph_phase_for_timestep from tensorrt_llm._torch.visual_gen.cuda_graph_runner import resolved_extra_key from tensorrt_llm.logger import logger @@ -173,9 +173,9 @@ def _dense_by_step(self, timestep: Any) -> bool: # Under CUDA-graph capture the runner has already resolved the phase # host-side (it is part of the graph key); reading the tensor here # would `.item()` inside capture, which CUDA forbids. - phase = resolved_extra_key("sol_attn_phase") + phase = resolved_extra_key("sparse_attn_phase") if phase is None: - phase = SkipSoftmaxScheduler.get_graph_phase_for_timestep( + phase = graph_phase_for_timestep( timestep, disabled_until_timestep=self.disabled_until_timestep, ) diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py b/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py index 1f85598d787a..076e677babd4 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py @@ -33,6 +33,7 @@ PredefinedAttentionMask, ) from ...attention.backends.sparse.params import SparseBackendForwardArgs, SparseParams +from ...attention.backends.sparse.timestep_phase import graph_phase_for_timestep, timestep_to_float from ...attention.backends.trtllm import TrtllmAttention as BaseTrtllmAttention from ...attention.backends.trtllm import TrtllmAttentionMetadata as BaseTrtllmAttentionMetadata from .interface import AttentionBackend, AttentionTensorLayout @@ -236,6 +237,8 @@ def __init__( self._preferred_layout = AttentionTensorLayout.NHD self.quant_attention_config = quant_attention_config + self._prepared_timestep: Optional[float] = None + self._timestep_prepared = False def update_quant_config(self, new_quant_config: Optional[QuantConfig]) -> None: """Rebuild FMHA libraries and bind VisualGen-owned shared plan caches.""" @@ -248,6 +251,55 @@ def update_quant_config(self, new_quant_config: Optional[QuantConfig]) -> None: if isinstance(fmha, PrimsTSBlockSparseFmha): fmha.bind_plan_cache(cache_state) + @property + def timestep_cutoff(self) -> Optional[float]: + """Normalized timestep below which the sparse algorithm is enabled, if any.""" + + return getattr(self.sparse_params, "disabled_until_timestep", None) + + def resolve_timestep(self, timestep: object) -> object: + """Reduce ``timestep`` to a host scalar once per eager call. + + Timestep-scheduled sparse algorithms read the timestep on the host, + which CUDA Graph capture cannot do for a device tensor. Eager calls, + including the warmup that precedes capture, remember the reduced value + and capture reuses it. Without a cutoff the timestep passes through + untouched. + """ + + if self.timestep_cutoff is None: + return timestep + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + if not self._timestep_prepared: + raise RuntimeError( + "sparse attention timestep must be prepared before CUDA Graph capture" + ) + return self._prepared_timestep + self._prepared_timestep = timestep_to_float(timestep) + self._timestep_prepared = True + return self._prepared_timestep + + @property + def dense_layers(self) -> frozenset[int]: + """Layer indices that always run dense attention.""" + + return getattr(self.sparse_params, "dense_layers", frozenset()) + + def should_use_sparse(self, timestep: object) -> bool: + """Return whether this layer runs its sparse path for the prepared ``timestep``. + + Dense layers never do. Otherwise the layer is sparse unless the timestep + schedule places the call in the dense prefix; without a cutoff or a + timestep the call is sparse. + """ + + if self.layer_idx in self.dense_layers: + return False + graph_phase = graph_phase_for_timestep( + timestep, disabled_until_timestep=self.timestep_cutoff + ) + return graph_phase is None or graph_phase == 1 + # Needed to work with torch compile cause of attention metadata # make attn metadata as input for it to work @torch.compiler.disable @@ -318,6 +370,7 @@ def forward( Returns: Output tensor [B, S, H*D] """ + timestep = self.resolve_timestep(timestep) block_sparse_inputs = ( sparse_backend_args.block_sparse_inputs if sparse_backend_args is not None else None ) diff --git a/tensorrt_llm/_torch/visual_gen/models/modeling.py b/tensorrt_llm/_torch/visual_gen/models/modeling.py index cd0f931eb3d6..6f61b91668d0 100644 --- a/tensorrt_llm/_torch/visual_gen/models/modeling.py +++ b/tensorrt_llm/_torch/visual_gen/models/modeling.py @@ -19,9 +19,9 @@ import torch import torch.nn as nn -from tensorrt_llm._torch.attention.backends.sparse.skip_softmax import SkipSoftmaxScheduler +from tensorrt_llm._torch.attention.backends.sparse.timestep_phase import graph_phase_for_timestep from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig -from tensorrt_llm.visual_gen.sparse_attention import SkipSoftmaxAttentionConfig, SolAttentionConfig +from tensorrt_llm.visual_gen.sparse_attention import BaseSparseAttentionConfig if TYPE_CHECKING: from tensorrt_llm._torch.visual_gen.cuda_graph_runner import CUDAGraphRunner @@ -74,42 +74,23 @@ 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, SkipSoftmaxAttentionConfig): - disabled_until_timestep = sparse_config.resolve_disabled_until_timestep( - pretrained_config=self.model_config.pretrained_config, - ) - if disabled_until_timestep is None: - return - - # Skip Softmax switches graph-visible attention behavior at the - # timestep boundary while tensor shapes stay unchanged. Key the dense - # and sparse phases separately; if timestep is absent or None, the - # scheduler returns None and the runner omits this key part. - runner.register_extra_key_fn( - "skip_softmax_phase", - lambda *args, **kwargs: SkipSoftmaxScheduler.get_graph_phase_for_timestep( - kwargs.get("timestep"), - disabled_until_timestep=disabled_until_timestep, - ), - ) + if not isinstance(sparse_config, BaseSparseAttentionConfig): return - - if isinstance(sparse_config, SolAttentionConfig): - disabled_until_timestep = sparse_config.disabled_until_timestep - if disabled_until_timestep is None: - # dense_layers is fixed per layer at construction, so it is - # already baked into each captured graph and needs no key. - return - - # Sol-Attn switches between dense and sparse attention at the - # dense-prefix boundary, again without changing tensor shapes, so - # the two phases must not share a captured graph. - runner.register_extra_key_fn( - "sol_attn_phase", - lambda *args, **kwargs: SkipSoftmaxScheduler.get_graph_phase_for_timestep( - kwargs.get("timestep"), - disabled_until_timestep=disabled_until_timestep, - ), - ) + disabled_until_timestep = sparse_config.resolve_disabled_until_timestep( + pretrained_config=self.model_config.pretrained_config, + ) + if disabled_until_timestep is None: return + + # Timestep-scheduled sparse attention switches graph-visible behavior at + # the cutoff while tensor shapes stay unchanged, and the backends prepare + # the matching phase during warmup. Key each live timestep so dense and + # sparse graphs stay separate; when no timestep is present the callback + # returns None and the runner omits this key part. + runner.register_extra_key_fn( + "sparse_attn_phase", + lambda *args, **kwargs: graph_phase_for_timestep( + kwargs.get("timestep"), + disabled_until_timestep=disabled_until_timestep, + ), + ) diff --git a/tensorrt_llm/visual_gen/sparse_attention.py b/tensorrt_llm/visual_gen/sparse_attention.py index 929797363298..5d2e552defd1 100644 --- a/tensorrt_llm/visual_gen/sparse_attention.py +++ b/tensorrt_llm/visual_gen/sparse_attention.py @@ -43,6 +43,21 @@ def to_sparse_metadata_params(self, **kwargs): """Lower user-facing config into SparseMetadataParams.""" return None + def resolve_disabled_until_timestep( + self, + *, + checkpoint_config: Optional[Dict[str, Any]] = None, + pretrained_config: Any = None, + ) -> Optional[float]: + """Return the normalized timestep below which the algorithm runs sparse. + + ``None`` means the algorithm has no dense prefix. Algorithms with a + ``disabled_until_timestep`` field return it; algorithms that also read a + checkpoint-provided cutoff override this method. + """ + del checkpoint_config, pretrained_config + return getattr(self, "disabled_until_timestep", None) + class SkipSoftmaxAttentionConfig(BaseSparseAttentionConfig): """SkipSoftmax sparse attention configuration for visual generation. diff --git a/tests/unittest/_torch/attention/sparse/test_timestep_phase.py b/tests/unittest/_torch/attention/sparse/test_timestep_phase.py new file mode 100644 index 000000000000..507268cd22d1 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/test_timestep_phase.py @@ -0,0 +1,52 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.sparse.timestep_phase import ( + graph_phase_for_timestep, + timestep_to_float, +) + +pytestmark = pytest.mark.cpu_only + + +@pytest.mark.parametrize( + ("timestep", "expected"), + [ + (None, None), + (torch.empty(0), None), + (0.25, 0.25), + (torch.tensor(0.5), 0.5), + # Per-token timesteps reduce to the largest live value. + (torch.tensor([0.0, 0.8]), 0.8), + ], +) +def test_timestep_to_float_reduces_to_the_largest_live_value(timestep, expected): + assert timestep_to_float(timestep) == ( + expected if expected is None else pytest.approx(expected) + ) + + +def test_timestep_to_float_rejects_non_real_and_non_finite_values(): + with pytest.raises(TypeError, match="real scalar or tensor"): + timestep_to_float(True) + with pytest.raises(ValueError, match="finite"): + timestep_to_float(float("nan")) + + +@pytest.mark.parametrize( + ("timestep", "cutoff", "expected"), + [ + (0.8, 0.6, 0), + (0.6, 0.6, 0), + (0.2, 0.6, 1), + (None, 0.6, None), + (0.2, None, None), + (torch.tensor([0.0, 0.8]), 0.6, 0), + (torch.tensor([0.0, 0.2]), 0.6, 1), + ], +) +def test_graph_phase_for_timestep_marks_dense_prefix_and_sparse_suffix(timestep, cutoff, expected): + assert graph_phase_for_timestep(timestep, disabled_until_timestep=cutoff) == expected diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py b/tests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py index 3347c9240232..35edb7da0671 100644 --- a/tests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py +++ b/tests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py @@ -17,6 +17,7 @@ SkipSoftmaxParams, SkipSoftmaxScheduler, ) +from tensorrt_llm._torch.attention.backends.sparse.timestep_phase import graph_phase_for_timestep from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl import fmha as cute_dsl_fmha from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.fmha import ( CuTeDSLAttention, @@ -393,6 +394,9 @@ def test_user_disabled_until_timestep_overrides_checkpoint_default(self): (0.6, 0), (0.59, 1), (None, None), + # Per-token timesteps stay dense until every live token is below the cutoff. + (torch.tensor([0.0, 0.8]), 0), + (torch.tensor([0.0, 0.2]), 1), ], ) def test_graph_phase_tracks_disabled_until_timestep_boundary( @@ -402,13 +406,7 @@ def test_graph_phase_tracks_disabled_until_timestep_boundary( ): # CUDA graph keys need a stable sparse-attention phase so captured # graphs are not reused across disabled and enabled skip-softmax states. - assert ( - SkipSoftmaxScheduler.get_graph_phase_for_timestep( - timestep, - disabled_until_timestep=0.6, - ) - == expected - ) + assert graph_phase_for_timestep(timestep, disabled_until_timestep=0.6) == expected class TestVisualGenSkipSoftmaxCuTeDSL: @@ -491,7 +489,7 @@ def register_extra_key_fn(self, name, fn): runner = _Runner() model.register_cuda_graph_extra_key_fns(runner) - phase_fn = runner.extra_key_fns["skip_softmax_phase"] + phase_fn = runner.extra_key_fns["sparse_attn_phase"] assert phase_fn(timestep=0.6) == 0 assert phase_fn(timestep=0.59) == 1 @@ -631,3 +629,48 @@ def test_pipeline_config_keeps_checkpoint_metadata_per_model(self, tmp_path): _expected_threshold(-20.0, 4.0, 0.5) ) assert transformer_disabled_params is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA Graph capture requires CUDA") +def test_trtllm_skip_softmax_cutoff_is_capturable_with_a_device_timestep(): + """The wrapper prepares the timestep during warmup so capture never reads the tensor.""" + from tensorrt_llm._torch.visual_gen.attention_backend.trtllm import TrtllmAttention + from tensorrt_llm._torch.visual_gen.config import create_attention_metadata_state + + params = SkipSoftmaxParams( + scheduler=SkipSoftmaxScheduler( + threshold_scale_factor_prefill=5000.0, disabled_until_timestep=0.6 + ) + ) + attention = TrtllmAttention( + layer_idx=0, + num_heads=2, + head_dim=128, + dtype=torch.bfloat16, + attention_metadata_state=create_attention_metadata_state(), + sparse_params=params, + ) + batch_size, seq_len = 1, 256 + q = torch.randn(batch_size, seq_len, 2, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn_like(q) + timestep = torch.tensor([0.8], device="cuda") + + def forward(): + return attention.forward(q, k, v, batch_size=batch_size, seq_len=seq_len, timestep=timestep) + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(2): + eager = forward() + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = forward() + graph.replay() + torch.cuda.synchronize() + + torch.testing.assert_close(captured, eager) diff --git a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py index 0da640d23d1e..f28d556c9720 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py +++ b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py @@ -30,7 +30,7 @@ import torch from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask -from tensorrt_llm._torch.attention.backends.sparse.skip_softmax import SkipSoftmaxScheduler +from tensorrt_llm._torch.attention.backends.sparse.timestep_phase import graph_phase_for_timestep from tensorrt_llm._torch.visual_gen.attention_backend import CuTeDSLAttention from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn import SolAttention from tensorrt_llm._torch.visual_gen.attention_backend.utils import create_attention @@ -207,34 +207,19 @@ def test_sol_attn_rejects_gqa_mqa(): def test_graph_phase_matches_skip_softmax_sense(timestep, expected): """Phase 0 is the dense prefix, 1 the sparse phase, None when undecidable. - Same contract as SkipSoftmaxScheduler.get_graph_phase_for_timestep. + Same contract as the shared graph_phase_for_timestep reduction. """ - assert ( - SkipSoftmaxScheduler.get_graph_phase_for_timestep(timestep, disabled_until_timestep=0.9545) - == expected - ) + assert graph_phase_for_timestep(timestep, disabled_until_timestep=0.9545) == expected def test_graph_phase_none_when_prefix_unset(): - assert ( - SkipSoftmaxScheduler.get_graph_phase_for_timestep(0.5, disabled_until_timestep=None) is None - ) + assert graph_phase_for_timestep(0.5, disabled_until_timestep=None) is None def test_graph_phase_accepts_tensor_timestep(): """Pipelines pass a tensor; a 0-d or 1-element tensor must work.""" - assert ( - SkipSoftmaxScheduler.get_graph_phase_for_timestep( - torch.tensor(0.99), disabled_until_timestep=0.95 - ) - == 0 - ) - assert ( - SkipSoftmaxScheduler.get_graph_phase_for_timestep( - torch.tensor([0.10]), disabled_until_timestep=0.95 - ) - == 1 - ) + assert graph_phase_for_timestep(torch.tensor(0.99), disabled_until_timestep=0.95) == 0 + assert graph_phase_for_timestep(torch.tensor([0.10]), disabled_until_timestep=0.95) == 1 def test_dense_prefix_skips_kernel(monkeypatch): @@ -694,22 +679,24 @@ def test_dense_by_step_prefers_runner_resolved_phase(monkeypatch): never touch the tensor. `torch.compiler.disable` does not help here -- it only excludes Dynamo, not stream capture. """ + import tensorrt_llm._torch.attention.backends.sparse.timestep_phase as timestep_phase + reads = {"n": 0} - real = SkipSoftmaxScheduler._as_float + real = timestep_phase.timestep_to_float def spy(value): reads["n"] += 1 return real(value) - monkeypatch.setattr(SkipSoftmaxScheduler, "_as_float", staticmethod(spy)) + monkeypatch.setattr(timestep_phase, "timestep_to_float", spy) attn = SolAttention(layer_idx=1, num_heads=2, head_dim=128) attn.disabled_until_timestep = 0.9 # Resolved phase wins even when the tensor says otherwise: phase 0 is the # dense prefix, phase 1 the sparse phase. - with resolved_extra_keys_scope({"sol_attn_phase": 0}): + with resolved_extra_keys_scope({"sparse_attn_phase": 0}): assert attn._dense_by_step(torch.tensor(0.1)) is True - with resolved_extra_keys_scope({"sol_attn_phase": 1}): + with resolved_extra_keys_scope({"sparse_attn_phase": 1}): assert attn._dense_by_step(torch.tensor(0.99)) is False assert reads["n"] == 0, "timestep tensor was read despite a runner-resolved phase" @@ -787,8 +774,8 @@ def test_sol_attn_dense_prefix_survives_cuda_graph_capture(make_runner): attn.disabled_until_timestep = cutoff runner = make_runner() runner.register_extra_key_fn( - "sol_attn_phase", - lambda *args, **kwargs: SkipSoftmaxScheduler.get_graph_phase_for_timestep( + "sparse_attn_phase", + lambda *args, **kwargs: graph_phase_for_timestep( kwargs.get("timestep"), disabled_until_timestep=cutoff ), ) 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 f901967bc417..af2618afdf1c 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py @@ -69,6 +69,10 @@ def _make_core_forward_metadata(): def _make_wrapper(cls=visual_trtllm.TrtllmAttention, *, quant_attention_config=None): attention = object.__new__(cls) attention.quant_attention_config = quant_attention_config + attention.layer_idx = 0 + attention.sparse_params = None + attention._prepared_timestep = None + attention._timestep_prepared = False return attention @@ -209,17 +213,68 @@ def _base_init(self, **kwargs): assert first_fmha._paged_wrappers is not other_fmha._paged_wrappers -def test_visual_gen_wrapper_does_not_define_its_own_prediction_lifecycle(): +def test_visual_gen_wrapper_owns_only_the_timestep_schedule(): assert not hasattr(visual_trtllm, "SparseForwardInputs") - for name in ( - "block_sparse_attn_predict", - "sparse_post_process", - "_forward_impl", - ): + for name in ("block_sparse_attn_predict", "sparse_post_process", "_forward_impl"): assert name not in visual_trtllm.TrtllmAttention.__dict__ + for name in ("resolve_timestep", "should_use_sparse"): + assert name in visual_trtllm.TrtllmAttention.__dict__ assert getattr(visual_trtllm.TrtllmAttention, "__parameters__", ()) == () +def test_wrapper_schedule_follows_cutoff_and_dense_layers(): + attention = _make_wrapper() + assert attention.should_use_sparse(0.8) + assert attention.should_use_sparse(None) + + attention.layer_idx = 1 + attention.sparse_params = SimpleNamespace( + disabled_until_timestep=0.6, dense_layers=frozenset({3}) + ) + assert not attention.should_use_sparse(0.8) + assert attention.should_use_sparse(0.2) + assert attention.should_use_sparse(None) + attention.layer_idx = 3 + assert not attention.should_use_sparse(0.2) + + +def test_wrapper_passes_timestep_through_without_a_cutoff(): + attention = _make_wrapper() + timestep = torch.tensor([0.2]) + + assert attention.resolve_timestep(timestep) is timestep + assert attention.resolve_timestep(None) is None + + +def test_wrapper_prepares_timestep_for_cuda_graph_capture(monkeypatch): + attention = _make_wrapper() + attention.sparse_params = SimpleNamespace(disabled_until_timestep=0.6) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + + with pytest.raises(RuntimeError, match="prepared before CUDA Graph capture"): + attention.resolve_timestep(torch.tensor([0.8])) + + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + # Per-token timesteps reduce to the largest live value. + assert attention.resolve_timestep(torch.tensor([0.0, 0.8])) == pytest.approx(0.8) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + assert attention.resolve_timestep(torch.tensor([0.2])) == pytest.approx(0.8) + + +def test_forward_hands_the_prepared_timestep_to_the_core(monkeypatch): + captured: dict = {} + _capture_core_forward(monkeypatch, captured) + attention = _make_wrapper() + attention.sparse_params = SimpleNamespace(disabled_until_timestep=0.6) + monkeypatch.setattr(attention, "support_fused_qkv", lambda: True, raising=False) + q = torch.randn(1, 4, 2, 8) + + attention.forward(q, q, q, batch_size=1, seq_len=4, timestep=torch.tensor([0.0, 0.8])) + + assert captured["forward_args"].timestep == pytest.approx(0.8) + + def test_forward_ignores_backend_specific_kwargs_like_other_backends(monkeypatch): captured: dict = {} _capture_core_forward(monkeypatch, captured) From 34ef975b0f12420b758d3451095adc1ac038fb03 Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Sat, 19 Sep 2026 06:04:13 +0000 Subject: [PATCH 05/10] feat: add the TRTLLM SOL backend and unify VisualGen SOL under one config SOL is served by two backends behind SolAttentionConfig (algorithm "sol_attn"): - SOLTrtllmAttention runs SOL in two stages through the generic TRTLLM sparse lifecycle. A TRT-LLM-owned predictor derives the exact block bitmask and K/V proxy summaries inside the core block_sparse_attn_predict hook, from the flattened Q/K/V, the batch layout in the attention metadata and the prepared timestep, and the shared PrimTS block-sparse FMHA executes that route. Each layer owns its predictor; the predictor keeps one shape-specialized plan with graph-stable route and summary buffers per shape. - SOLCuTeDSLAttention is the fused CuTeDSL kernel backend, moved from attention_backend/cute_dsl/sol_attn.py into attention_backend/sparse/sol next to the TRTLLM backend, the way the VSA backends share attention_backend/sparse/vsa. It consumes the same lowered SolParams and keeps its per-call dense delegation for cross-attention, masks and inputs the kernel cannot serve. SolAttentionConfig lowers into SolParams for both backends: tau, disabled_until_timestep and dense_layers (a list of layer indices) drive both, while thresh_type stays a CUTEDSL kernel knob; the TRTLLM predictor implements the diag policy only, so AttentionConfig rejects "exact" with the TRTLLM backend. Quantized attention and torch.compile fullgraph stay rejected for SOL on either backend, and TRTLLM SEPARATE_QKV cross-attention falls back to VANILLA like dense TRTLLM. The SOL predictor kernels, the TRTLLM SOL tests and their test-list entries come with the backend; the CuTeDSL SOL tests follow the moved module and the shared sparse_attn_phase key. The feature guide describes both backends in one SOL section. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- .../features/visualgen-sparse-attention.md | 128 ++- .../visual_gen/attention_backend/__init__.py | 3 +- .../attention_backend/cute_dsl/__init__.py | 5 +- .../attention_backend/sparse/sol/__init__.py | 25 + .../sol_attn.py => sparse/sol/backend.py} | 190 +++- .../attention_backend/sparse/sol/kernels.py | 507 ++++++++++ .../attention_backend/sparse/sol/params.py | 76 ++ .../attention_backend/sparse/sol/predictor.py | 282 ++++++ .../visual_gen/attention_backend/trtllm.py | 22 +- .../visual_gen/attention_backend/utils.py | 18 +- .../blackwell/sol_attn_backend.py | 4 +- .../_torch/visual_gen/modules/attention.py | 29 +- tensorrt_llm/visual_gen/args.py | 40 +- tensorrt_llm/visual_gen/sparse_attention.py | 72 +- .../test_lists/test-db/l0_b200.yml | 3 + .../integration/test_lists/test-db/l0_cpu.yml | 2 + .../sparse_attention/test_sol_attention.py | 891 ++++++++++++++++++ .../sparse_attention/test_sol_predictor.py | 374 ++++++++ .../test_sol_predictor_kernels.py | 206 ++++ .../test_attention_cute_dsl_sol_attn.py | 46 +- .../_torch/visual_gen/test_ltx2_pipeline.py | 71 ++ .../_torch/visual_gen/test_visual_gen_args.py | 49 + 22 files changed, 2884 insertions(+), 159 deletions(-) create mode 100644 tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/__init__.py rename tensorrt_llm/_torch/visual_gen/attention_backend/{cute_dsl/sol_attn.py => sparse/sol/backend.py} (64%) 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/features/visualgen-sparse-attention.md b/docs/source/features/visualgen-sparse-attention.md index a3e4d73cb1dd..30b86868be2a 100644 --- a/docs/source/features/visualgen-sparse-attention.md +++ b/docs/source/features/visualgen-sparse-attention.md @@ -6,8 +6,8 @@ This feature is in **beta** stage. APIs, supported models, and optimization opti - [Overview](#overview) - [Algorithms](#algorithms) - - [Sol-Attn](#sol-attn) - [Skip Softmax Attention](#skip-softmax-attention) +- [SOL Attention](#sol-attention) - [Video Sparse Attention (VSA)](#video-sparse-attention-vsa) ## Overview @@ -22,40 +22,7 @@ Sparse attention in VisualGen is configured through `VisualGenArgs.attention_con |---|---|---| | `skip_softmax` | `SkipSoftmaxAttentionConfig` | Supported | | `vsa` | `VideoSparseAttentionConfig` | Supported (`CUTEDSL`, `TRTLLM`) | -| `sol_attn` | `SolAttentionConfig` | Supported (CUTEDSL, sm100/sm103) | - -### Sol-Attn - -Sol-Attn ([arXiv:2607.24027](https://arxiv.org/abs/2607.24027)) folds dynamic block -routing, sparse computation, and an approximation-correction term into one -online-softmax pass. It runs on the **CUTEDSL** backend only, on datacenter -Blackwell -- sm100 (B200/GB200) and sm103 (B300/GB300) -- and requires -`head_dim=128`, bfloat16, and MHA -(`num_kv_heads == num_heads`). - -```yaml -attention_config: - backend: CUTEDSL - sparse_attention_config: - algorithm: sol_attn - tau: 2.0 # routing threshold; higher routes more blocks sparse - thresh_type: diag # or "exact" - disabled_until_timestep: 0.9090 # dense while normalized timestep >= cutoff - dense_layers: [0] # optional: layer indices forced dense -``` - -`disabled_until_timestep` has the same meaning as it does for Skip Softmax: -attention runs dense while the normalized denoising timestep is at or above the -cutoff, protecting the high-noise prefix, and switches to the sparse kernel -below it. Use `None` rather than `0.0` to disable the prefix. - -On an input the kernel is known not to serve — an unsupported architecture, a -`head_dim` other than 128, a non-bfloat16 dtype — Sol-Attn runs dense -attention instead (the configured backend's dense kernel where available, -torch SDPA otherwise), logs the specific reason once, and counts the fallback. -Set `TRTLLM_SOL_ATTN_STRICT=1` to raise instead, which is useful when -benchmarking to confirm the kernel actually ran. Errors raised by the kernel -itself are not caught. +| `sol_attn` | `SolAttentionConfig` | Supported (`TRTLLM`, `CUTEDSL`) | ## Skip Softmax Attention @@ -127,7 +94,7 @@ User configuration is supplied through Python or YAML and controls how the check `threshold_scale_factor` and `target_sparsity` are alternatives: if both are present, `threshold_scale_factor` takes precedence and the calibration formula is not used. User-provided `target_sparsity` and `disabled_until_timestep` override checkpoint defaults. Checkpoint `ignore` patterns always disable Skip Softmax Attention for matching layers. -Skip Softmax Attention works with both the **TRTLLM** and **CUTEDSL** attention backends in VisualGen. Set `attention_config.backend` to either when enabling it. On CUTEDSL, Skip Softmax Attention can also be combined with `quant_attention_config`'s block-scaled Q/K recipes (MXFP8, NVFP4); VSA and Sol-Attn are currently the two supported sparse-attention algorithms, both available only through the CuTeDSL attention backend; quantized attention is not yet supported or enabled with either mode. +Skip Softmax Attention works with both the **TRTLLM** and **CUTEDSL** attention backends in VisualGen. Set `attention_config.backend` to either when enabling it. On CUTEDSL, Skip Softmax Attention can also be combined with `quant_attention_config`'s block-scaled Q/K recipes (MXFP8, NVFP4); VSA and SOL cannot be combined with quantized attention on any backend. #### Mapping `disabled_until_timestep` to Actual Denoising Steps @@ -251,6 +218,95 @@ 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 ([arXiv:2607.24027](https://arxiv.org/abs/2607.24027)) routes attention +blocks dynamically: blocks whose scores stand out are computed exactly, the +rest are approximated from compact K/V proxies and folded back into the online +softmax. VisualGen serves it on two backends behind one `SolAttentionConfig`: + +- **TRTLLM** runs SOL in two stages. 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: it + 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. +- **CUTEDSL** runs the fused kernel vendored from the reference implementation + (`SOLCuTeDSLAttention`), which folds routing, sparse computation and the + approximation correction into one online-softmax pass. It runs on datacenter + Blackwell, sm100 (B200/GB200) and sm103 (B300/GB300), and requires + `head_dim=128`, bfloat16 and MHA (`num_kv_heads == num_heads`). On an input + the kernel is known not to serve, such as an unsupported architecture, a + `head_dim` other than 128 or a non-bfloat16 dtype, it runs dense attention + instead (the CuTe DSL FMHA where available, torch SDPA otherwise), logs the + reason once and counts the fallback; set `TRTLLM_SOL_ATTN_STRICT=1` to raise + instead, which is useful when benchmarking to confirm the kernel actually + ran. Errors raised by the kernel itself are not caught. Cross-attention and + masked calls are delegated to dense attention per call. + +Configure SOL with `SolAttentionConfig` and either backend: + +```python +from tensorrt_llm.visual_gen import AttentionConfig, SolAttentionConfig + +attention_config = AttentionConfig( + backend="TRTLLM", # or "CUTEDSL" + sparse_attention_config=SolAttentionConfig( + tau=1.0, + disabled_until_timestep=0.6, + dense_layers=[0, 2, 3, 4], + ), +) +``` + +The equivalent YAML is: + +```yaml +attention_config: + backend: TRTLLM # or CUTEDSL + sparse_attention_config: + algorithm: sol_attn + tau: 1.0 # routing threshold; higher tau routes more blocks sparse + disabled_until_timestep: 0.6 # dense while the normalized timestep >= cutoff + dense_layers: [0, 2, 3, 4] # optional: layer indices forced dense + thresh_type: diag # CUTEDSL kernel threshold policy; "exact" needs CUTEDSL +``` + +- `tau` is the routing threshold in standard deviations above the mean block + score; higher values route more blocks sparse. +- `disabled_until_timestep` has the same meaning as it does for Skip Softmax + Attention: attention runs dense while the normalized denoising timestep is at + or above the cutoff, protecting the high-noise prefix, and switches to SOL + below it. Use `None` rather than `0.0` to disable the prefix. +- `dense_layers` lists zero-based layer indices that always use dense + attention. +- `thresh_type` selects the threshold policy of the CUTEDSL kernel. The TRTLLM + predictor implements `diag` only, and `AttentionConfig` rejects `exact` + together with the TRTLLM backend. + +The TRTLLM 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. The TRTLLM +backend uses a host-side graph break to prepare and own predictor plans, so +`torch_compile_config.enable_fullgraph=True` is rejected for SOL; keep the +default `False` setting. + +When a cutoff is configured, VisualGen includes the dense-or-sparse phase in +the CUDA Graph key. The TRTLLM attention wrapper reduces the timestep to a host +value during graph warmup and reuses it during capture for every +timestep-scheduled algorithm (Skip Softmax Attention and SOL), while SOL +predictor route buffers remain stable for replay; the CuTeDSL backends read the +phase the CUDA Graph runner resolved for the graph key instead of the device +tensor. Per-token timesteps, such as Wan I2V where the conditioning frame stays +at timestep zero, reduce to their largest live value, so the schedule stays +dense until every token is below the cutoff. 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/__init__.py b/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py index 2c577f0e95bc..6800d9a629bd 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py @@ -21,7 +21,7 @@ """ from .cudnn import CuDNNAttention -from .cute_dsl import CuTeDSLAttention, SolAttention +from .cute_dsl import CuTeDSLAttention from .flash_attn4 import FlashAttn4Attention from .flashinfer import FlashInferAttention from .interface import AttentionBackend, AttentionTensorLayout @@ -39,7 +39,6 @@ "FlashAttn4Attention", "FlashInferAttention", "RingAttention", - "SolAttention", "TrtllmAttention", "TrtllmAttentionMetadata", "UlyssesAttention", 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 8c4ca507810d..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,15 +15,12 @@ """ CuTe DSL attention backend family for visual generation models. - fmha.py — CuTeDSLAttention (dense and blockscaled JIT FMHA) - sol_attn.py — SolAttention (Sol-Attn dynamic block routing, CuTe JIT + SDPA fallback) + fmha.py — CuTeDSLAttention (dense and blockscaled JIT FMHA) """ from .fmha import CuTeDSLAttention, _cute_dsl_import_error -from .sol_attn import SolAttention __all__ = [ "CuTeDSLAttention", "_cute_dsl_import_error", - "SolAttention", ] 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..844254c75a7f --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/__init__.py @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""VisualGen SOL sparse attention: shared parameters, the TRT-LLM-owned predictor, +and the TRTLLM and CuTeDSL backends.""" + +from .backend import SOLCuTeDSLAttention, SOLTrtllmAttention +from .params import SolParams +from .predictor import ( + SolPredictorGeometry, + SolPredictorOutputs, + SolPredictorPlan, + SolPredictorPlanKey, + SOLSparsePredictor, +) + +__all__ = [ + "SOLCuTeDSLAttention", + "SOLSparsePredictor", + "SOLTrtllmAttention", + "SolParams", + "SolPredictorGeometry", + "SolPredictorOutputs", + "SolPredictorPlan", + "SolPredictorPlanKey", +] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/backend.py similarity index 64% rename from tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py rename to tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/backend.py index f69a278d9414..3c851f31e381 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/sol_attn.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/backend.py @@ -12,49 +12,54 @@ # 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. -""" -Sol-Attn backend for visual generation models. - -Sol-Attn (https://arxiv.org/abs/2607.24027) is dynamic block routing + -sparse computation + approximation correction folded into one online-softmax -pass. The kernel is vendored from its reference implementation -(https://github.com/NVlabs/Sana, branch -https://github.com/NVlabs/Sana/tree/sol-engine, pinned at commit -https://github.com/NVlabs/Sana/commit/5fe5feb -- see -``cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md`` for the pin -and its currency-check note) under ``..cute_dsl_kernels.blackwell.sol_attn`` -/ ``sol_attn_backend.py``. Only the sm100 (B200/GB200) -Blackwell) kernels are carried; the upstream sm89/sm90 kernels and the Triton -reference path are not, and the FlashAttention CuTe helpers they needed come -from the ``flash-attn-4`` dependency rather than a vendored copy. - -This file is only the TRT-LLM AttentionBackend adapter around that kernel's -public BTHD entry point, plus the dense_layers layer-skip guard. - -``disabled_until_timestep`` is the dense-prefix control, and mirrors -skip_softmax's field of the same name: sparse attention stays disabled (that -is, the layer runs the backend's dense kernel) while the normalized timestep is at -or above the cutoff, and switches to the sparse kernel once it drops below. - -The timestep arrives as a forward kwarg -- ``modules/attention.py`` already -threads it to every backend, and all VisualGen pipelines normalize it to -``[0, 1]`` by ``num_train_timesteps`` per the ``BaseDiffusionModel.forward`` -contract. Nothing has to be wired per pipeline, and there is no process-wide -state to keep in sync. +"""VisualGen SOL attention backends. + +SOL (Sol-Attn, https://arxiv.org/abs/2607.24027) routes attention blocks +dynamically and corrects the approximation of the blocks it skips. Two backends +serve one ``SolParams``: + +* ``SOLTrtllmAttention`` runs SOL in two stages through the generic TRTLLM + sparse lifecycle: a TRT-LLM-owned predictor derives the exact block bitmask + and K/V proxy summaries, then the shared PrimTS block-sparse FMHA executes + that route. +* ``SOLCuTeDSLAttention`` runs the fused kernel vendored from the reference + implementation (https://github.com/NVlabs/Sana, branch ``sol-engine``, pinned + in ``cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md``), which + folds routing, sparse computation and correction into one online-softmax + pass. Only the sm100 kernels are carried; see + ``cute_dsl_kernels/blackwell/sol_attn_backend.py`` for the shape and dtype + guard around them. +``disabled_until_timestep`` is the dense-prefix control shared with skip +softmax: the layer runs dense while the normalized timestep is at or above the +cutoff and switches to SOL below it. The TRTLLM wrapper prepares that timestep +as a host value before CUDA Graph capture; the CuTeDSL backend reads the phase +the CUDA Graph runner resolved for the graph key. """ +from __future__ import annotations + from typing import Any, 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.sparse.timestep_phase import graph_phase_for_timestep +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.visual_gen.cuda_graph_runner import resolved_extra_key from tensorrt_llm.logger import logger -from ....attention.backends.interface import PredefinedAttentionMask -from ..interface import AttentionBackend, AttentionTensorLayout -from ..vanilla import VanillaAttention +from ...interface import AttentionBackend, AttentionTensorLayout +from ...trtllm import TrtllmAttention +from ...vanilla import VanillaAttention +from .params import SolParams +from .predictor import BLOCK_SIZE, SOLSparsePredictor _sol_attn_import_error = None try: @@ -75,7 +80,7 @@ def _cute_dense_available() -> bool: degrades to SDPA instead of raising. """ try: - from .fmha import _check_cute_runtime_available, _get_gpu_arch + from ...cute_dsl.fmha import _check_cute_runtime_available, _get_gpu_arch _check_cute_runtime_available() _get_gpu_arch() @@ -84,8 +89,8 @@ def _cute_dense_available() -> bool: return True -class SolAttention(AttentionBackend): - """Sol-Attn dynamic block-routing sparse attention (CuTeDSL, sm100/sm103). +class SOLCuTeDSLAttention(AttentionBackend): + """Fused SOL sparse attention (CuTeDSL, sm100/sm103). Inputs the kernel cannot serve (shape, dtype, architecture) are delegated to dense attention up front by ``_run_sol_attn_bthd``; a kernel error is not @@ -101,14 +106,18 @@ def __init__( head_dim: int = 128, num_kv_heads: Optional[int] = None, dtype: Optional[torch.dtype] = None, - sparse_attention_config=None, + sparse_params: SolParams | None = None, **kwargs, ): if _sol_attn_run is None: raise ImportError( - "SolAttention requires the vendored sol_attn kernel " + "SOLCuTeDSLAttention requires the vendored sol_attn kernel " f"package; import failed: {_sol_attn_import_error}" ) + if sparse_params is None: + sparse_params = SolParams() + if not isinstance(sparse_params, SolParams): + raise TypeError("SOLCuTeDSLAttention requires SolParams") self.layer_idx = layer_idx self.num_heads = num_heads self.head_dim = head_dim @@ -123,18 +132,17 @@ def __init__( f"GQA/MQA is not supported." ) self.dtype = dtype - cfg = sparse_attention_config - self.tau = getattr(cfg, "tau", 1.0) - self.thresh_type = getattr(cfg, "thresh_type", "diag") - self.disabled_until_timestep = getattr(cfg, "disabled_until_timestep", None) - self.dense_layers = frozenset(getattr(cfg, "dense_layers", None) or ()) + self.tau = sparse_params.tau + self.thresh_type = sparse_params.thresh_type + self.disabled_until_timestep = sparse_params.disabled_until_timestep + self.dense_layers = sparse_params.dense_layers - # Sol-Attn's dense steps must run the backend the user selected. Without + # SOL's dense steps must run the backend the user selected. Without # this they ran torch SDPA while a `backend: CUTEDSL` baseline ran # cute_dsl_fmha_fwd, so candidate and reference differed on the dense # steps too -- measured at LPIPS 0.214 on Wan2.2-T2V-A14B with sparsity # switched off entirely, against a 0.25 gate. - from .fmha import CuTeDSLAttention + from ...cute_dsl.fmha import CuTeDSLAttention # The only backend that consumes `key_padding_mask` (CuTeDSL's `_fwd` # swallows it via **kwargs, silently). Masked self-attention is routed @@ -313,3 +321,97 @@ def preferred_layout(self) -> AttentionTensorLayout: @classmethod def support_fused_qkv(cls) -> bool: return False + + +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 + super().__init__(sparse_params=None, **kwargs) + self.predictor = SOLSparsePredictor() + + @property + def timestep_cutoff(self) -> Optional[float]: + """SOL keeps its parameters outside the core ``sparse_params`` slot.""" + + return self.sol_params.disabled_until_timestep + + @property + def dense_layers(self) -> frozenset[int]: + return self.sol_params.dense_layers + + 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, + already prepared by the wrapper forward, from ``forward_args``. + """ + + if not self.should_use_sparse(forward_args.timestep): + 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__ = ["SOLCuTeDSLAttention", "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..1c5defb10995 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/params.py @@ -0,0 +1,76 @@ +# 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 + +from tensorrt_llm._torch.attention.backends.sparse.params import SparseParams + + +@dataclass(frozen=True, slots=True) +class SolParams(SparseParams): + """Static SOL policy lowered from the user-facing VisualGen config. + + Shared by the TRTLLM (predictor + block-sparse FMHA) and CuTeDSL (fused kernel) + backends; ``thresh_type`` selects the CuTeDSL kernel threshold policy, the TRTLLM + predictor implements ``diag`` only. + """ + + algorithm: Literal["sol_attn"] = field(init=False, default="sol_attn") + tau: float = 1.0 + thresh_type: Literal["diag", "exact"] = "diag" + 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) + + if self.thresh_type not in ("diag", "exact"): + raise ValueError("thresh_type must be 'diag' or 'exact'") + + 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) + + +__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 076e677babd4..9051e53e11c1 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py @@ -323,6 +323,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, @@ -390,9 +408,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 b2087a4819d8..9d2cd4d7a425 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py @@ -50,7 +50,7 @@ def get_visual_gen_attention_backend( and architecture-specific NVFP4 attention recipes. - "FA4": Flash Attention 4; provides higher speedup on Blackwell GPUs (sm100) Requires flash-attn package with cute interface - - "CUTEDSL": CuTe DSL kernels. create_attention selects dense/SkipSoftmax FMHA or VSA + - "CUTEDSL": CuTe DSL kernels. create_attention selects dense/SkipSoftmax FMHA, VSA or SOL from AttentionConfig.sparse_attention_config. - "CUDNN": cuDNN fused SDPA. Unquantized by default; quant_attention_config selects per-tensor FP8 or block-scaled MXFP8 (Blackwell). @@ -136,6 +136,7 @@ def create_attention( ) sparse_algorithm = getattr(sparse_attention_config, "algorithm", None) is_vsa = sparse_algorithm == "vsa" + is_sol = sparse_algorithm == "sol_attn" backend_name = backend.upper() if is_vsa and backend_name == "CUTEDSL": @@ -146,11 +147,14 @@ def create_attention( from .sparse.vsa.backend import VSATrtllmAttention attn_cls = VSATrtllmAttention - elif sparse_algorithm == "sol_attn" and backend_name == "CUTEDSL": - from .cute_dsl.sol_attn import SolAttention + elif is_sol and backend_name == "CUTEDSL": + from .sparse.sol.backend import SOLCuTeDSLAttention - attn_cls = SolAttention - kwargs["sparse_attention_config"] = sparse_attention_config + attn_cls = SOLCuTeDSLAttention + 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) @@ -158,6 +162,10 @@ def create_attention( sparse_params = kwargs.pop("sparse_params", None) if sparse_params is not None: raise ValueError("VSA does not lower through core SparseParams.") + elif is_sol and kwargs.get("sparse_params") is None: + # The attention module lowers the config once per layer; callers that + # construct a backend directly get the same lowering here. + kwargs["sparse_params"] = sparse_attention_config.to_sparse_params() # Forward the validated quantization recipe to TRTLLM, cuDNN, FlashInfer, or the dense CuTe DSL # FMHA backend. diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py index 829a3ce462f5..926ca6e5f957 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/sol_attn_backend.py @@ -44,8 +44,8 @@ ``[batch, tokens, heads, 128]`` Q/K/V, ``tau``, ``thresh_type``, ``kv_splits``, and an optional exact KV sink range. -TRT-LLM's dispatch path (``attention_backend/cute_dsl/sol_attn.py``, -``SolAttention``) consumes exactly two names from this module: +TRT-LLM's dispatch path (``attention_backend/sparse/sol/backend.py``, +``SOLCuTeDSLAttention``) consumes exactly two names from this module: ``_run_sol_attn_bthd`` and ``sol_attn_supported``. The dense-prefix decision lives there too, keyed off the normalized timestep forward kwarg. diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/modules/attention.py index 44b9834c7879..cd767504b735 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 @@ -114,16 +114,23 @@ def __init__( _sa_cfg = config.attention.sparse_attention_config _sa_algo = getattr(_sa_cfg, "algorithm", None) if _sa_cfg is not None else None is_vsa = _sa_algo == "vsa" - _is_sol_attn = base_backend == "CUTEDSL" and _sa_algo == "sol_attn" + is_sol = _sa_algo == "sol_attn" is_separate_qkv_cross_attention = ( self.qkv_mode == QKVMode.SEPARATE_QKV and not separate_qkv_is_self_attention ) - # Cross-attention fallback: dense TRTLLM and every VSA backend are self-attn only. - # Sol-Attn is absent by design; see SolAttention._can_serve. + # Cross-attention fallback: every TRTLLM backend and every VSA backend are + # self-attention only. The CuTeDSL SOL backend delegates cross-attention per + # call instead; see SOLCuTeDSLAttention._can_serve. if is_separate_qkv_cross_attention and (base_backend == "TRTLLM" or is_vsa): 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( @@ -136,13 +143,19 @@ def __init__( # Every sparse algorithm here routes over the whole token sequence, so # none of them can be split across context-parallel ranks. - if (is_vsa or _is_sol_attn) and cp_size > 1: - _algo_name = "VSA" if is_vsa else "Sol-Attn" + if (is_vsa or is_sol) and cp_size > 1: + _algo_name = "VSA" if is_vsa else "SOL" raise ValueError( f"{_algo_name} 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." ) + 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 @@ -252,6 +265,8 @@ def __init__( module_name=self.module_name, pretrained_config=config.pretrained_config, ) + elif isinstance(ss_cfg, SolAttentionConfig) and backend_name in ("TRTLLM", "CUTEDSL"): + sparse_params = ss_cfg.to_sparse_params() self.sparse_params = sparse_params # Create compute backend diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index 7540ec2d548e..33b70f0b2153 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -127,7 +127,7 @@ class AttentionConfig(StrictBaseModel): description=( "Sparse attention recipe. Discriminated by algorithm: " "skip_softmax (TRTLLM / CUTEDSL backends), vsa (CUTEDSL / TRTLLM backends), " - "or sol_attn (CUTEDSL backend)." + "or sol_attn (TRTLLM / CUTEDSL backends)." ), ) @@ -225,7 +225,7 @@ def _validate_sparse_attention_config(self) -> "AttentionConfig": supported_backends = { "skip_softmax": ("TRTLLM", "CUTEDSL"), "vsa": ("CUTEDSL", "TRTLLM"), - "sol_attn": ("CUTEDSL",), + "sol_attn": ("TRTLLM", "CUTEDSL"), }.get(algo) if supported_backends is None: return self @@ -244,7 +244,7 @@ def _validate_quant_sparse_mutex(self) -> "AttentionConfig": if self.quant_attention_config is None or self.sparse_attention_config is None: return self - # VSA and Sol-Attn replace the dense attention path on every backend that + # VSA and SOL replace the dense attention path on every backend that # serves them and never consume quant_attention_config, so accepting a # quantization recipe would silently ignore user configuration. # SkipSoftmax is part of the dense path itself and can compose. @@ -252,7 +252,23 @@ def _validate_quant_sparse_mutex(self) -> "AttentionConfig": if algorithm == "vsa": raise ValueError("VSA and quant_attention_config are mutually exclusive.") if algorithm == "sol_attn": - raise ValueError("Sol-Attn and quant_attention_config are mutually exclusive.") + raise ValueError("SOL and quant_attention_config are mutually exclusive.") + return self + + @model_validator(mode="after") + def _validate_sol_thresh_type(self) -> "AttentionConfig": + # The TRTLLM predictor implements the diag threshold only; the exact + # policy exists in the fused CUTEDSL kernel. + sparse_config = self.sparse_attention_config + if ( + isinstance(sparse_config, SolAttentionConfig) + and self.backend == "TRTLLM" + and sparse_config.thresh_type != "diag" + ): + raise ValueError( + f"TRTLLM SOL supports thresh_type='diag' only, got " + f"thresh_type={sparse_config.thresh_type!r}; use backend='CUTEDSL'." + ) return self @@ -791,6 +807,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] @@ -842,6 +873,7 @@ def from_yaml(cls, yaml_path: Union[str, Path], **overrides: Any) -> "VisualGenA "QuantAttentionConfig", "SparseAttentionConfig", "SkipSoftmaxAttentionConfig", + "SolAttentionConfig", "VideoSparseAttentionConfig", "SolAttentionConfig", "AttentionConfig", diff --git a/tensorrt_llm/visual_gen/sparse_attention.py b/tensorrt_llm/visual_gen/sparse_attention.py index 5d2e552defd1..66f74639d9ba 100644 --- a/tensorrt_llm/visual_gen/sparse_attention.py +++ b/tensorrt_llm/visual_gen/sparse_attention.py @@ -238,29 +238,40 @@ def _ckpt_sparse_attention_config_from_kwargs( class SolAttentionConfig(BaseSparseAttentionConfig): - """Sol-Attn sparse attention configuration for visual generation. - - Dynamic block routing + sparse computation + approximation correction in - one online-softmax pass (arXiv:2607.24027). Kernel is CuTeDSL on - datacenter Blackwell -- sm100 (B200/GB200) and sm103 (B300/GB300) -- - head_dim=128, bf16, MHA. - - On an unsupported *shape, dtype, or architecture* the kernel falls back to - dense attention -- the configured backend's dense kernel where available, - torch SDPA otherwise -- and counts the fallback, so setting this config on the wrong GPU - degrades rather than fails. Two cases are not covered by that fallback and do raise: GQA/MQA - (num_kv_heads != num_heads) here at construction, and context parallelism - (cp_size > 1), rejected in visual_gen/modules/attention.py. + """SOL sparse attention configuration for visual generation. + + SOL (Sol-Attn, arXiv:2607.24027) routes attention blocks dynamically and + corrects the approximation of the blocks it skips. Two backends serve it: + + - ``TRTLLM`` runs SOL in two stages. A TRT-LLM-owned predictor derives an + exact block bitmask and K/V proxy summaries from Q/K/V, then the generic + PrimTS block-sparse FMHA executes that route. Unsupported runtime tensor + envelopes raise instead of silently falling back to dense attention. + - ``CUTEDSL`` runs the fused kernel vendored from the reference + implementation, which folds routing, sparse computation and correction + into one online-softmax pass on datacenter Blackwell (sm100/sm103) with + head_dim=128, bf16 and MHA. Inputs the kernel cannot serve fall back to + dense attention and are counted; GQA/MQA raises at construction. + + Context parallelism (cp_size > 1) is rejected for both backends in + visual_gen/modules/attention.py. """ algorithm: Literal["sol_attn"] = "sol_attn" tau: float = PydanticField( 1.0, - description="Per-block routing threshold; higher tau routes more blocks sparse.", + allow_inf_nan=False, + description=( + "Routing threshold in standard deviations above the mean block score; " + "higher tau routes more blocks sparse." + ), ) thresh_type: Literal["diag", "exact"] = PydanticField( "diag", - description="Threshold policy forwarded to the kernel (kernel default: 'diag').", + description=( + "Threshold policy of the CUTEDSL kernel (kernel default: 'diag'). " + "The TRTLLM predictor implements 'diag' only." + ), ) disabled_until_timestep: Optional[float] = PydanticField( None, @@ -269,15 +280,13 @@ class SolAttentionConfig(BaseSparseAttentionConfig): description=( "Dense-prefix cutoff on the normalized denoising timestep, with the " "same sense as skip_softmax's field of the same name: the layer runs " - "dense while timestep >= this value and switches to the sparse kernel " - "below it. Larger timesteps are earlier, noisier steps, so this " - "protects the high-noise prefix. Use None (not 0.0) to disable the " - "prefix; 0.0 is rejected because it would run dense on every step " - "and silently turn Sol-Attn off entirely. " - "Read from the `timestep` forward kwarg, which must be the normalized " - "scheduler time (larger = noisier). WAN and LTX-2 pass it; a pipeline " - "that does not, or that passes something else, gets a one-time warning " - "and runs sparse on every step (fail-open)." + "dense while timestep >= this value and switches to SOL below it. " + "Larger timesteps are earlier, noisier steps, so this protects the " + "high-noise prefix. Use None (not 0.0) to disable the prefix; 0.0 is " + "rejected because it would run dense on every step and silently turn " + "SOL off entirely. Read from the `timestep` forward kwarg, which must " + "be the normalized scheduler time (larger = noisier); a pipeline that " + "does not pass it runs sparse on every step (fail-open)." ), ) dense_layers: Optional[list[int]] = PydanticField( @@ -300,11 +309,16 @@ def _validate_dense_layers(cls, layers: Optional[list[int]]) -> Optional[list[in return sorted(set(layers)) def to_sparse_params(self, **kwargs): - # Sol-Attn's knobs are consumed directly by SolAttention.__init__ - # (constructed via CUTEDSL backend dispatch in create_attention), not - # lowered into a shared SparseParams -- the vendored kernel has no - # checkpoint-calibration step to resolve here, unlike skip_softmax. - return None + """Lower the public recipe into the SOL parameters shared by both backends.""" + del kwargs + from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.params import SolParams + + return SolParams( + tau=self.tau, + thresh_type=self.thresh_type, + disabled_until_timestep=self.disabled_until_timestep, + dense_layers=frozenset(self.dense_layers or ()), + ) class VideoSparseAttentionConfig(StrictBaseModel): diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index a193a49a3ee5..268383a3684d 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -303,6 +303,9 @@ l0_b200: - unittest/_torch/visual_gen/test_attention_vsa.py - unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.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_cudnn.py - unittest/_torch/visual_gen/test_attention_integration.py diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 38247c9ad825..56b1833c6cbb 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -77,6 +77,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..709d4dc17f83 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/sparse_attention/test_sol_attention.py @@ -0,0 +1,891 @@ +# 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_timestep = None + backend._timestep_prepared = False + 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_backend_reuses_prepared_timestep_during_cuda_graph_capture(monkeypatch) -> None: + backend, _predictor = _stub_backend(SolParams(tau=1.0, disabled_until_timestep=0.6)) + # Per-token timesteps reduce to the largest live value. + assert backend.resolve_timestep(torch.tensor([0.0, 0.2])) == pytest.approx(0.2) + + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + prepared = backend.resolve_timestep(torch.tensor(0.8)) + + assert prepared == pytest.approx(0.2) + assert backend.should_use_sparse(prepared) + + +@_CPU_ONLY +def test_sol_backend_warmup_prepares_dense_phase_for_capture(monkeypatch) -> None: + q = _bshd() + backend, predictor = _stub_backend(SolParams(tau=1.0, disabled_until_timestep=0.6)) + + assert _predict(backend, q, q, q, timestep=backend.resolve_timestep(torch.tensor(0.8))) is None + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + assert _predict(backend, q, q, q, timestep=backend.resolve_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: + backend, predictor = _stub_backend(SolParams(tau=1.0, disabled_until_timestep=0.6)) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + + with pytest.raises(RuntimeError, match="prepared before CUDA Graph capture"): + backend.resolve_timestep(torch.tensor(0.2)) + + predictor.predict.assert_not_called() + + +@_CPU_ONLY +def test_sol_backend_without_timestep_runs_sparse_like_skip_softmax() -> None: + q = _bshd() + backend, predictor = _stub_backend(SolParams(tau=1.0, disabled_until_timestep=0.6)) + + assert _predict(backend, q, q, q, timestep=backend.resolve_timestep(None)) is not None + predictor.predict.assert_called_once() + + +@_CPU_ONLY +def test_sol_phase_waits_until_all_token_timesteps_are_below_cutoff() -> None: + backend, _predictor = _stub_backend(SolParams(tau=1.0, disabled_until_timestep=0.6)) + + assert not backend.should_use_sparse(backend.resolve_timestep(torch.tensor([0.0, 0.8]))) + assert backend.should_use_sparse(backend.resolve_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_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": [-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_supported_backend() -> None: + with pytest.raises(ValidationError, match="requires backend"): + AttentionConfig( + backend="VANILLA", + sparse_attention_config=SolAttentionConfig(), + ) + + +@_CPU_ONLY +def test_sol_exact_threshold_needs_cutedsl_backend() -> None: + with pytest.raises(ValidationError, match="thresh_type"): + AttentionConfig( + backend="TRTLLM", + sparse_attention_config=SolAttentionConfig(thresh_type="exact"), + ) + config = AttentionConfig( + backend="CUTEDSL", + sparse_attention_config=SolAttentionConfig(thresh_type="exact"), + ) + assert config.sparse_attention_config.to_sparse_params().thresh_type == "exact" + + +@_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 ("sparse_attn_phase", 0) in captured_keys[0] + assert ("sparse_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_attention_cute_dsl_sol_attn.py b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py index f28d556c9720..fe4aa7f68d21 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py +++ b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py @@ -15,7 +15,7 @@ """Sol-Attn correctness tests: backend dispatch, config guards, step-context dense_layers/disabled_until_timestep guards. -Mirrors test_attention_cute_dsl_vsa.py's structure and scope for its sibling +Mirrors test_attention_vsa.py's structure and scope for its sibling sparse-attention algorithm. GPU kernel-vs-dense numerical equivalence is covered by test_cute_kernel_matches_dense_on_a_single_block: Sol-Attn's routing is score-derived, so no tau provably forces dense routing the way @@ -32,7 +32,7 @@ from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask from tensorrt_llm._torch.attention.backends.sparse.timestep_phase import graph_phase_for_timestep from tensorrt_llm._torch.visual_gen.attention_backend import CuTeDSLAttention -from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn import SolAttention +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.backend import SOLCuTeDSLAttention from tensorrt_llm._torch.visual_gen.attention_backend.utils import create_attention from tensorrt_llm._torch.visual_gen.config import ( DiffusionModelConfig, @@ -69,7 +69,7 @@ def test_cute_dsl_factory_dispatches_dense_and_sol_attn() -> None: ) assert isinstance(dense_attention, CuTeDSLAttention) - assert isinstance(sol_attn_attention, SolAttention) + assert isinstance(sol_attn_attention, SOLCuTeDSLAttention) assert sol_attn_attention.tau == 2.0 assert sol_attn_attention.disabled_until_timestep == 0.9545 @@ -129,7 +129,7 @@ def test_sol_attn_cross_attention_delegates_to_dense_cutedsl(): assert cross_attn.attn_backend == "CUTEDSL", ( f"expected CUTEDSL, got {cross_attn.attn_backend!r}" ) - assert isinstance(cross_attn.attn, SolAttention), ( + assert isinstance(cross_attn.attn, SOLCuTeDSLAttention), ( "Sol-Attn should remain the backend and delegate per call, not be " f"swapped out at construction; got {type(cross_attn.attn).__name__}" ) @@ -148,7 +148,7 @@ def test_sol_attn_self_attention_is_served_under_separate_qkv(): Qwen-Image uses it unconditionally. Both are self-attention; both must still get the sparse kernel. """ - attn = SolAttention(layer_idx=1, num_heads=2, head_dim=128) + attn = SOLCuTeDSLAttention(layer_idx=1, num_heads=2, head_dim=128) attn.disabled_until_timestep = None # CPU tensors on purpose: `_can_serve` compares shapes and the layer index # and never touches the device, so this runs on CPU-only hosts too. @@ -190,7 +190,7 @@ def test_sol_attn_with_context_parallelism_raises(): def test_sol_attn_rejects_gqa_mqa(): """Sol-Attn is MHA-only; num_kv_heads != num_heads must fail fast at construction.""" with pytest.raises(ValueError, match="MHA-only"): - SolAttention(layer_idx=0, num_heads=8, head_dim=128, num_kv_heads=2) + SOLCuTeDSLAttention(layer_idx=0, num_heads=8, head_dim=128, num_kv_heads=2) @pytest.mark.parametrize( @@ -229,14 +229,14 @@ def test_dense_prefix_skips_kernel(monkeypatch): routes to the CuTe kernel on CUDA is covered by `test_dense_paths_use_cutedsl_backend`. """ - import tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn as sol_attn_mod + import tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.backend as sol_attn_mod def _fail_if_called(*args, **kwargs): raise AssertionError("kernel must not run inside the dense prefix") monkeypatch.setattr(sol_attn_mod, "_sol_attn_run", _fail_if_called) - attn = SolAttention(layer_idx=0, num_heads=2, head_dim=16) + attn = SOLCuTeDSLAttention(layer_idx=0, num_heads=2, head_dim=16) attn.disabled_until_timestep = 0.9 q = k = v = torch.randn(1, 4, 2, 16) out = attn.forward(q, k, v, timestep=torch.tensor(0.95)) @@ -249,7 +249,7 @@ def test_missing_timestep_fails_open_to_sparse(monkeypatch): Matches the CuTeDSL skip-softmax path's fail-open choice. """ - import tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn as sol_attn_mod + import tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.backend as sol_attn_mod called = {"n": 0} @@ -259,7 +259,7 @@ def _record(*args, **kwargs): monkeypatch.setattr(sol_attn_mod, "_sol_attn_run", _record) - attn = SolAttention(layer_idx=0, num_heads=2, head_dim=16) + attn = SOLCuTeDSLAttention(layer_idx=0, num_heads=2, head_dim=16) attn.disabled_until_timestep = 0.9 q = k = v = torch.randn(1, 4, 2, 16) attn.forward(q, k, v) # no timestep kwarg @@ -268,14 +268,14 @@ def _record(*args, **kwargs): def test_dense_layers_guard_skips_kernel(monkeypatch): """A layer_idx in dense_layers must use the dense SDPA path and never invoke the kernel.""" - import tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn as sol_attn_mod + import tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.backend as sol_attn_mod def _fail_if_called(*args, **kwargs): raise AssertionError("kernel must not be invoked for a dense_layers-forced layer") monkeypatch.setattr(sol_attn_mod, "_sol_attn_run", _fail_if_called) - attn = SolAttention(layer_idx=3, num_heads=2, head_dim=16) + attn = SOLCuTeDSLAttention(layer_idx=3, num_heads=2, head_dim=16) attn.dense_layers = frozenset({3}) q = k = v = torch.randn(1, 4, 2, 16) out = attn.forward(q, k, v) @@ -596,8 +596,8 @@ def test_timestep_scalar_read_is_opaque_to_dynamo(): Otherwise it graph-breaks the enclosing block once per attention layer. """ - assert _is_dynamo_disabled(SolAttention._dense_by_step), ( - "SolAttention._dense_by_step must be decorated with @torch.compiler.disable" + assert _is_dynamo_disabled(SOLCuTeDSLAttention._dense_by_step), ( + "SOLCuTeDSLAttention._dense_by_step must be decorated with @torch.compiler.disable" ) @@ -615,7 +615,7 @@ def test_dense_paths_use_cutedsl_backend(monkeypatch): The CPU-tensor tests above cannot see this: `_dense` falls back to SDPA when `q.is_cuda` is false, so they exercise the wrong branch by construction. """ - import tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn as sol_attn_mod + import tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.backend as sol_attn_mod if not sol_attn_mod._cute_dense_available(): # A CUDA device is not enough: the premise here is that the dense paths @@ -627,7 +627,7 @@ def test_dense_paths_use_cutedsl_backend(monkeypatch): q = k = v = torch.randn(1, 64, 2, 128, device=device, dtype=torch.bfloat16) def _make(): - a = SolAttention(layer_idx=0, num_heads=2, head_dim=128) + a = SOLCuTeDSLAttention(layer_idx=0, num_heads=2, head_dim=128) calls = {"n": 0} real = a._inner.forward @@ -689,7 +689,7 @@ def spy(value): return real(value) monkeypatch.setattr(timestep_phase, "timestep_to_float", spy) - attn = SolAttention(layer_idx=1, num_heads=2, head_dim=128) + attn = SOLCuTeDSLAttention(layer_idx=1, num_heads=2, head_dim=128) attn.disabled_until_timestep = 0.9 # Resolved phase wins even when the tensor says otherwise: phase 0 is the @@ -770,7 +770,7 @@ def test_sol_attn_dense_prefix_survives_cuda_graph_capture(make_runner): pytest.skip("no Sol-Attn kernel for this device") cutoff = 0.9 - attn = SolAttention(layer_idx=1, num_heads=2, head_dim=128) + attn = SOLCuTeDSLAttention(layer_idx=1, num_heads=2, head_dim=128) attn.disabled_until_timestep = cutoff runner = make_runner() runner.register_extra_key_fn( @@ -951,14 +951,14 @@ def test_key_padding_mask_routes_to_vanilla_and_is_honored(monkeypatch): correct destination is VANILLA. The output is checked against a masked SDPA reference so a mask that was routed but then dropped still fails. """ - import tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn as sol_attn_mod + import tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.backend as sol_attn_mod monkeypatch.setattr( sol_attn_mod, "_sol_attn_run", lambda *a, **k: pytest.fail("sparse kernel ran on a masked call"), ) - attn = SolAttention(layer_idx=1, num_heads=2, head_dim=128) + attn = SOLCuTeDSLAttention(layer_idx=1, num_heads=2, head_dim=128) attn.disabled_until_timestep = None calls = {"vanilla": 0} real = attn._vanilla.forward @@ -991,14 +991,14 @@ def test_causal_mask_routes_to_dense_and_is_honored(monkeypatch): device it is `_sdpa`, which previously dropped it. Either way the result must match a causal reference. """ - import tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.sol_attn as sol_attn_mod + import tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.backend as sol_attn_mod monkeypatch.setattr( sol_attn_mod, "_sol_attn_run", lambda *a, **k: pytest.fail("sparse kernel ran on a causal call"), ) - attn = SolAttention(layer_idx=1, num_heads=2, head_dim=128) + attn = SOLCuTeDSLAttention(layer_idx=1, num_heads=2, head_dim=128) attn.disabled_until_timestep = None torch.manual_seed(0) q, k, v = (torch.randn(1, 64, 2, 128, dtype=torch.float32) for _ in range(3)) @@ -1012,7 +1012,7 @@ def test_causal_mask_routes_to_dense_and_is_honored(monkeypatch): def test_unmasked_self_attention_is_still_served(): """The routing change must not touch the measured path: no mask, sparse.""" - attn = SolAttention(layer_idx=1, num_heads=2, head_dim=128) + attn = SOLCuTeDSLAttention(layer_idx=1, num_heads=2, head_dim=128) attn.disabled_until_timestep = None q = k = torch.randn(1, 64, 2, 128, dtype=torch.bfloat16) assert attn._can_serve(q, k) diff --git a/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py b/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py index ce4614969841..616d143bab91 100644 --- a/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py @@ -1341,6 +1341,77 @@ def register_cuda_graph_extra_key_fns(self, runner): assert pipeline.transformer.registered_runner is runner assert pipeline.transformer.forward.__wrapped__.__self__ is pipeline.transformer + def test_two_stage_cuda_graph_setup_keys_sparse_phase(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 + + video = Modality( + latent=torch.empty(1, 2, 4), + timesteps=torch.tensor([0.5]), + positions=torch.empty(1, 3, 2), + context=torch.empty(1, 3, 4), + ) + + def run(timestep, step_index): + return pipeline.transformer( + video=video, + audio=None, + text_cache=None, + timestep=torch.tensor([timestep]), + step_index=step_index, + ) + + dense_key = run(0.8, 0) + sparse_key = run(0.2, 1) + dense_replay_key = run(0.8, 2) + + assert "sparse_attn_phase" in runner._extra_key_fns + assert ("sparse_attn_phase", 0) in dense_key + assert ("sparse_attn_phase", 1) in sparse_key + assert dense_key != sparse_key + assert dense_replay_key == dense_key + assert captured_keys == [dense_key, sparse_key] + 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 006cb38fd73e..4684da2950ed 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, @@ -513,6 +514,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 4fe3be5ffe8b4f9060fd0bb9e8945c124a7cc167 Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Sun, 20 Sep 2026 04:17:20 +0000 Subject: [PATCH 06/10] refactor: inject shared FMHA state through the TRTLLM attention constructor TrtllmAttention takes an optional fmha_state dict, defaulting to a private dict per instance, and PrimsTSBlockSparseFmha keeps its wrapper-plan caches under its own key in that dict. Layers constructed with one shared dict plan each static block-sparse profile once and reuse its route workspace; the sharing granularity is the lifetime of the dict the constructor caller hands in. Libraries rebuilt by update_quant_config read the same dict, so the caches survive a rebuild without any rebinding. The VisualGen wrapper passes the fmha_caches entry of its component-scoped attention_metadata_state to the core, which removes bind_plan_cache, the wrapper's update_quant_config override and the metadata adapter's get_fmha_cache_state, and lets the wrapper build its metadata adapter after the core constructor again. The developer guide and the state docstring describe the constructor contract. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- .../attention/ATTENTION_DEVELOPER_GUIDE.md | 12 +++++---- .../backends/fmha/prims_ts_block_sparse.py | 23 +++++++--------- .../_torch/attention/backends/trtllm.py | 6 +++++ .../visual_gen/attention_backend/trtllm.py | 26 +++++-------------- tensorrt_llm/_torch/visual_gen/config.py | 8 +++--- .../sparse/test_prims_ts_block_sparse.py | 25 ++++++++++-------- .../test_trtllm_attention_metadata.py | 10 ++++--- 7 files changed, 52 insertions(+), 58 deletions(-) diff --git a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index aca96dcd2d47..9ed1ab460e8a 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md @@ -224,11 +224,13 @@ 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. +`PrimsTSBlockSparseFmha` keeps its wrapper-plan cache in the attention's +`fmha_state` dict (`TrtllmAttention(fmha_state=...)`), which defaults to a +private dict per layer. Integrations whose attention layers execute serially +hand one dict to every layer of a component to plan each static profile once +and reuse its graph-stable route workspace across compatible layers; the +sharing granularity is the lifetime of that dict. A shared dict must not serve +concurrent forwards, so each independent model component owns 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` 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 a93174f26aa5..c8eb08cef8a7 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 @@ -190,29 +190,24 @@ def _uniform_seq_len_q( class PrimsTSBlockSparseFmha(PrimsTSFmha): """Contiguous context and fixed-Q paged generation block-sparse FMHA.""" + PLAN_CACHE_KEY = "prims_ts_block_sparse" + supports_block_sparse_inputs = True def __init__(self, attn: "TrtllmAttention") -> None: super().__init__(attn) - self._contiguous_wrappers: dict[_BlockSparsePlanKey, "BlockSparseTSWrapper"] = {} - self._paged_wrappers: dict[_BlockSparsePlanKey, "BlockSparsePagedTSWrapper"] = {} - - def bind_plan_cache(self, cache_state: dict[str, object]) -> None: - """Share planned wrappers with every adapter bound to ``cache_state``. - - Attention layers that execute serially, such as the blocks of one - diffusion transformer, see identical static profiles. Binding them to - one model-scoped container plans each profile once and allocates its - route workspace once. Call before the first forward. - """ - + # Planned wrappers live in the attention's FMHA state. Layers constructed + # with one shared state, such as the blocks of one diffusion transformer, + # plan each static profile once and allocate its route workspace once; + # the default per-instance state keeps the caches private to this layer. + caches = attn.fmha_state.setdefault(self.PLAN_CACHE_KEY, {}) self._contiguous_wrappers = cast( dict[_BlockSparsePlanKey, "BlockSparseTSWrapper"], - cache_state.setdefault("contiguous_wrappers", {}), + caches.setdefault("contiguous_wrappers", {}), ) self._paged_wrappers = cast( dict[_BlockSparsePlanKey, "BlockSparsePagedTSWrapper"], - cache_state.setdefault("paged_wrappers", {}), + caches.setdefault("paged_wrappers", {}), ) def _is_supported( diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 77e465b9e171..3a227ad16208 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -1396,6 +1396,7 @@ def __init__( sparse_params: Optional[SparseParams] = None, kv_cache_dtype: str = "auto", skip_correction_threshold: float = 0.0, + fmha_state: Optional[dict] = None, **kwargs, ) -> None: """ @@ -1420,6 +1421,10 @@ def __init__( used by DeepSeek-V4 and DSA on SM120/SM121. skip_correction_threshold (float): Runtime MLA threshold. Zero disables skip-correction. + fmha_state (dict): Optional state shared with the FMHA libraries of this + backend, for example plan caches. Libraries read the entries they own + from it, so handing one dict to several layers shares those entries + across them; the default is a private dict per instance. """ super().__init__(layer_idx, num_heads, head_dim, num_kv_heads, quant_config, **kwargs) @@ -1491,6 +1496,7 @@ def __init__( self.kv_scale_orig_quant = 1.0 / self.kv_cache_scaling_factor self.local_layer_idx: Optional[int] = None + self.fmha_state: dict = {} if fmha_state is None else fmha_state self._fmha_manager: FmhaManager if not skip_create_weights_in_init: self.update_quant_config(self.quant_config) diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py b/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py index 9051e53e11c1..99008b4633cf 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py @@ -78,12 +78,6 @@ 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). @@ -219,9 +213,6 @@ def __init__( "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, @@ -231,26 +222,21 @@ def __init__( quant_config=quant_config, sparse_params=sparse_params, dtype=dtype, + # Every layer of one model component shares its FMHA plan caches. + fmha_state=attention_metadata_state.setdefault("fmha_caches", {}), ) # 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 self._prepared_timestep: Optional[float] = None self._timestep_prepared = False - 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) - @property def timestep_cutoff(self) -> Optional[float]: """Normalized timestep below which the sparse algorithm is enabled, if any.""" diff --git a/tensorrt_llm/_torch/visual_gen/config.py b/tensorrt_llm/_torch/visual_gen/config.py index 6ea10f3f19d5..28fe9324df32 100644 --- a/tensorrt_llm/_torch/visual_gen/config.py +++ b/tensorrt_llm/_torch/visual_gen/config.py @@ -83,10 +83,10 @@ def create_attention_metadata_state() -> Dict[str, Any]: """Create state shared by attention layers in one model component. The state outlives individual forwards and CUDA Graph captures. It owns the - shape-keyed TRTLLM metadata cache, and the TRTLLM attention wrapper adds the - shared PrimTS block-sparse plan caches on demand so one static profile is - planned once per component instead of once per layer. Each model component - receives a distinct state and must not execute concurrent forwards. + shape-keyed TRTLLM metadata cache and the FMHA state the TRTLLM attention + layers of the component share, so one static block-sparse profile is planned + once per component instead of once per layer. Each model component receives + a distinct state and must not execute concurrent forwards. """ return {"metadata_cache": {}} diff --git a/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py b/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py index 40d74b28ede3..df77f68f358f 100644 --- a/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py +++ b/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py @@ -120,7 +120,8 @@ def _proxy_reference( class _Attention: - def __init__(self) -> None: + def __init__(self, fmha_state: dict | None = None) -> None: + self.fmha_state = {} if fmha_state is None else fmha_state self.sparse_params = None self.num_heads = 2 self.num_kv_heads = 1 @@ -385,22 +386,24 @@ def _key(fmha): assert _key(first) != _key(second) -def test_block_sparse_plan_cache_is_shared_only_when_explicitly_bound() -> None: +def test_block_sparse_plan_cache_follows_the_attention_fmha_state() -> None: 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 - 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": {}, + fmha_state = {} + shared_first = block_sparse_fmha.PrimsTSBlockSparseFmha(_Attention(fmha_state)) + shared_second = block_sparse_fmha.PrimsTSBlockSparseFmha(_Attention(fmha_state)) + + assert shared_first._contiguous_wrappers is shared_second._contiguous_wrappers + assert shared_first._paged_wrappers is shared_second._paged_wrappers + assert fmha_state == { + "prims_ts_block_sparse": { + "contiguous_wrappers": {}, + "paged_wrappers": {}, + } } 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 af2618afdf1c..db43faf2bd09 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py @@ -166,7 +166,8 @@ def _base_update_quant_config(self, new_quant_config): ) def _base_init(self, **kwargs): - del kwargs + fmha_state = kwargs.get("fmha_state") + self.fmha_state = {} if fmha_state is None else fmha_state self.is_mla_enable = False self.kv_lora_rank = None self.v_head_dim = None @@ -180,7 +181,7 @@ def _base_init(self, **kwargs): ) 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 + assert "fmha_caches" not in attention_metadata_state first = visual_trtllm.TrtllmAttention( attention_metadata_state=attention_metadata_state, @@ -189,8 +190,9 @@ def _base_init(self, **kwargs): attention_metadata_state=attention_metadata_state, ) - assert not hasattr(first, "_block_sparse_fmha_cache_state") - assert not hasattr(second, "_block_sparse_fmha_cache_state") + assert first.fmha_state is attention_metadata_state["fmha_caches"] + assert second.fmha_state is first.fmha_state + assert "update_quant_config" not in visual_trtllm.TrtllmAttention.__dict__ 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 From 0c00f38792761054a3d93d81e55fc1b21bb2e9d9 Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Sun, 20 Sep 2026 05:58:01 +0000 Subject: [PATCH 07/10] refactor: keep the prepared sparse timestep in the component attention state The host value of the denoising timestep is per-step state of a whole model component, not of one attention layer, so the VisualGen TRTLLM metadata adapter now owns it: prepare_timestep reduces the timestep tensor during eager calls, including the warmup that precedes CUDA Graph capture, stores the value in the component-scoped attention_metadata_state next to the metadata cache and the FMHA state, and returns the stored value under capture, where the tensor cannot be read. The wrapper keeps only the schedule policy: it asks the adapter when a cutoff is configured and answers should_use_sparse from the cutoff and dense_layers of its sparse parameters. The per-layer _prepared_timestep and _timestep_prepared fields are gone; the state is the value itself, and its presence marks the timestep as prepared. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- .../features/visualgen-sparse-attention.md | 7 ++-- .../visual_gen/attention_backend/trtllm.py | 41 +++++++++++-------- tensorrt_llm/_torch/visual_gen/config.py | 7 ++-- .../sparse_attention/test_sol_attention.py | 10 +++-- .../test_trtllm_attention_metadata.py | 26 +++++++++++- 5 files changed, 64 insertions(+), 27 deletions(-) diff --git a/docs/source/features/visualgen-sparse-attention.md b/docs/source/features/visualgen-sparse-attention.md index 30b86868be2a..cd00ddfeee85 100644 --- a/docs/source/features/visualgen-sparse-attention.md +++ b/docs/source/features/visualgen-sparse-attention.md @@ -297,9 +297,10 @@ backend uses a host-side graph break to prepare and own predictor plans, so default `False` setting. When a cutoff is configured, VisualGen includes the dense-or-sparse phase in -the CUDA Graph key. The TRTLLM attention wrapper reduces the timestep to a host -value during graph warmup and reuses it during capture for every -timestep-scheduled algorithm (Skip Softmax Attention and SOL), while SOL +the CUDA Graph key. The TRTLLM attention metadata reduces the timestep to a host +value during graph warmup, keeps it in the component attention state and reuses +it during capture for every timestep-scheduled algorithm (Skip Softmax Attention +and SOL), while SOL predictor route buffers remain stable for replay; the CuTeDSL backends read the phase the CUDA Graph runner resolved for the graph key instead of the device tensor. Per-token timesteps, such as Wan I2V where the conditioning frame stays diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py b/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py index 99008b4633cf..9cacb5742b00 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py @@ -132,6 +132,26 @@ def _select_cached_metadata(self, cached) -> None: self._prepared = cached["prepared"] self._cached_seq_lens = cached["seq_lens"] + def prepare_timestep(self, timestep: object) -> Optional[float]: + """Reduce ``timestep`` to a host scalar and keep it for CUDA Graph capture. + + Timestep-scheduled sparse algorithms read the timestep on the host, + which CUDA Graph capture cannot do for a device tensor. Eager calls, + including the warmup that precedes capture, reduce the tensor and store + the value in the component state; capture returns the stored value. + """ + + state = self._metadata_state + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + if "timestep" not in state: + raise RuntimeError( + "sparse attention timestep must be prepared before CUDA Graph capture" + ) + return state["timestep"] + value = timestep_to_float(timestep) + state["timestep"] = value + return value + def prepare( self, batch_size: int, @@ -234,8 +254,6 @@ def __init__( ) self.quant_attention_config = quant_attention_config - self._prepared_timestep: Optional[float] = None - self._timestep_prepared = False @property def timestep_cutoff(self) -> Optional[float]: @@ -244,26 +262,17 @@ def timestep_cutoff(self) -> Optional[float]: return getattr(self.sparse_params, "disabled_until_timestep", None) def resolve_timestep(self, timestep: object) -> object: - """Reduce ``timestep`` to a host scalar once per eager call. + """Return the host timestep the sparse schedule consumes. - Timestep-scheduled sparse algorithms read the timestep on the host, - which CUDA Graph capture cannot do for a device tensor. Eager calls, - including the warmup that precedes capture, remember the reduced value - and capture reuses it. Without a cutoff the timestep passes through + Layers with a timestep cutoff hand the tensor to the metadata adapter, + which reduces it during eager calls and reuses the prepared value under + CUDA Graph capture. Without a cutoff the timestep passes through untouched. """ if self.timestep_cutoff is None: return timestep - if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): - if not self._timestep_prepared: - raise RuntimeError( - "sparse attention timestep must be prepared before CUDA Graph capture" - ) - return self._prepared_timestep - self._prepared_timestep = timestep_to_float(timestep) - self._timestep_prepared = True - return self._prepared_timestep + return self.metadata.prepare_timestep(timestep) @property def dense_layers(self) -> frozenset[int]: diff --git a/tensorrt_llm/_torch/visual_gen/config.py b/tensorrt_llm/_torch/visual_gen/config.py index 28fe9324df32..c9f305dae5f5 100644 --- a/tensorrt_llm/_torch/visual_gen/config.py +++ b/tensorrt_llm/_torch/visual_gen/config.py @@ -83,9 +83,10 @@ def create_attention_metadata_state() -> Dict[str, Any]: """Create state shared by attention layers in one model component. The state outlives individual forwards and CUDA Graph captures. It owns the - shape-keyed TRTLLM metadata cache and the FMHA state the TRTLLM attention - layers of the component share, so one static block-sparse profile is planned - once per component instead of once per layer. Each model component receives + shape-keyed TRTLLM metadata cache, the FMHA state the TRTLLM attention layers + of the component share, so one static block-sparse profile is planned once per + component instead of once per layer, and the host value of the denoising + timestep prepared for CUDA Graph capture. Each model component receives a distinct state and must not execute concurrent forwards. """ return {"metadata_cache": {}} 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 index 709d4dc17f83..3d0cce51fd95 100644 --- a/tests/unittest/_torch/visual_gen/sparse_attention/test_sol_attention.py +++ b/tests/unittest/_torch/visual_gen/sparse_attention/test_sol_attention.py @@ -39,7 +39,10 @@ SolPredictorOutputs, SOLSparsePredictor, ) -from tensorrt_llm._torch.visual_gen.attention_backend.trtllm import TrtllmAttention +from tensorrt_llm._torch.visual_gen.attention_backend.trtllm import ( + TrtllmAttention, + TrtllmAttentionMetadata, +) from tensorrt_llm._torch.visual_gen.attention_backend.utils import create_attention from tensorrt_llm._torch.visual_gen.config import ( DiffusionModelConfig, @@ -75,8 +78,9 @@ def _make_backend( backend.sparse_params = None backend._fmha_manager = SimpleNamespace(fmha_libs=[object.__new__(PrimsTSBlockSparseFmha)]) backend.sol_params = params - backend._prepared_timestep = None - backend._timestep_prepared = False + backend.metadata = TrtllmAttentionMetadata( + device=torch.device("cpu"), attention_metadata_state={} + ) backend.predictor = predictor return backend 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 db43faf2bd09..4018fd5fad9c 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py @@ -71,8 +71,9 @@ def _make_wrapper(cls=visual_trtllm.TrtllmAttention, *, quant_attention_config=N attention.quant_attention_config = quant_attention_config attention.layer_idx = 0 attention.sparse_params = None - attention._prepared_timestep = None - attention._timestep_prepared = False + attention.metadata = visual_trtllm.TrtllmAttentionMetadata( + device=torch.device("cpu"), attention_metadata_state={} + ) return attention @@ -264,6 +265,27 @@ def test_wrapper_prepares_timestep_for_cuda_graph_capture(monkeypatch): assert attention.resolve_timestep(torch.tensor([0.2])) == pytest.approx(0.8) +def test_trtllm_attention_metadata_prepares_timestep_for_capture(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + state = {} + metadata = visual_trtllm.TrtllmAttentionMetadata( + device=torch.device("cpu"), attention_metadata_state=state + ) + with pytest.raises(RuntimeError, match="prepared before CUDA Graph capture"): + metadata.prepare_timestep(torch.tensor([0.8])) + + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + assert metadata.prepare_timestep(torch.tensor([0.0, 0.8])) == pytest.approx(0.8) + assert state["timestep"] == pytest.approx(0.8) + assert metadata.prepare_timestep(None) is None + assert state["timestep"] is None + + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + state["timestep"] = 0.4 + assert metadata.prepare_timestep(torch.tensor([0.9])) == pytest.approx(0.4) + + def test_forward_hands_the_prepared_timestep_to_the_core(monkeypatch): captured: dict = {} _capture_core_forward(monkeypatch, captured) From 9c666fbc10dfcaebfadd3702d04c15b4b376aa4a Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Sun, 20 Sep 2026 06:09:47 +0000 Subject: [PATCH 08/10] refactor: make the SOL predictor a functional operator and add the exact threshold The TRTLLM SOL predictor is now one functional custom operator, trtllm::visual_gen_sol_predictor(q, k, v, block_size, tau, sm_scale, thresh_type) -> (exact_block_bits, k_summary, v_summary), that allocates its outputs per call. Every output is produced and consumed inside the same transformer forward, so under CUDA Graph capture the tensors come from the graph pool and are replayed with the graph, and the runner's eager warmup already compiles the Triton kernels before capture. That removes the plan and geometry classes, the per-instance plan cache, the capture guard and the host-side graph break, so SOL no longer rejects torch.compile fullgraph and the corresponding VisualGenArgs validator is gone. The routing threshold of every query block comes from two Triton kernels that mirror the fused kernel's preprocess: a key-statistics kernel reduces the key block summaries to their per-channel mean and variance and, for the exact policy, their full second moment with tensor-core dots on the 16-bit summaries, and a threshold kernel projects the query block centroids onto those statistics with an IEEE fp32 dot. diag models each key channel independently and exact uses the full key covariance. The selection kernel consumes the threshold, and the PyTorch implementation of the same rule stays as the CPU path and the reference, computed in float64 because the TF32 matmuls TensorRT-LLM enables would perturb the fp32 second moment. SolParams and SolAttentionConfig lower thresh_type to both backends, so the rule that rejected exact with the TRTLLM backend is removed. Tests: the predictor tests cover both policies against an fp64 oracle, a correlated-key case where the policies differ, fresh outputs under CUDA Graph replay, a fullgraph compile, threshold parity with the fused kernel's preprocess and Triton-versus-reference thresholds; the SOL attention tests stub the predictor functions instead of a predictor object. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- .../features/visualgen-sparse-attention.md | 21 +- .../attention_backend/sparse/sol/__init__.py | 14 +- .../attention_backend/sparse/sol/backend.py | 15 +- .../attention_backend/sparse/sol/kernels.py | 462 ++++++++++++------ .../attention_backend/sparse/sol/params.py | 3 +- .../attention_backend/sparse/sol/predictor.py | 334 +++---------- tensorrt_llm/visual_gen/args.py | 31 -- tensorrt_llm/visual_gen/sparse_attention.py | 4 +- .../sparse_attention/test_sol_attention.py | 105 ++-- .../sparse_attention/test_sol_predictor.py | 393 +++++++-------- .../test_sol_predictor_kernels.py | 132 +++-- .../_torch/visual_gen/test_visual_gen_args.py | 49 -- 12 files changed, 749 insertions(+), 814 deletions(-) diff --git a/docs/source/features/visualgen-sparse-attention.md b/docs/source/features/visualgen-sparse-attention.md index cd00ddfeee85..3f5ec48906a1 100644 --- a/docs/source/features/visualgen-sparse-attention.md +++ b/docs/source/features/visualgen-sparse-attention.md @@ -275,7 +275,7 @@ attention_config: tau: 1.0 # routing threshold; higher tau routes more blocks sparse disabled_until_timestep: 0.6 # dense while the normalized timestep >= cutoff dense_layers: [0, 2, 3, 4] # optional: layer indices forced dense - thresh_type: diag # CUTEDSL kernel threshold policy; "exact" needs CUTEDSL + thresh_type: diag # block threshold policy: "diag" or "exact" ``` - `tau` is the routing threshold in standard deviations above the mean block @@ -286,22 +286,23 @@ attention_config: below it. Use `None` rather than `0.0` to disable the prefix. - `dense_layers` lists zero-based layer indices that always use dense attention. -- `thresh_type` selects the threshold policy of the CUTEDSL kernel. The TRTLLM - predictor implements `diag` only, and `AttentionConfig` rejects `exact` - together with the TRTLLM backend. +- `thresh_type` selects how the routing threshold models the key blocks: + `diag` treats every key channel independently, `exact` uses the full key + covariance. Both backends implement both policies from the same per-block + statistics. The TRTLLM 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. The TRTLLM -backend uses a host-side graph break to prepare and own predictor plans, so -`torch_compile_config.enable_fullgraph=True` is rejected for SOL; keep the -default `False` setting. +BSHD Q/K/V tensors, equal Q/K/V shapes and head dimension 128. The predictor is +one graph-visible operator that allocates its route and summary tensors per +call, so it composes with CUDA Graph capture and with torch.compile, including +`fullgraph`, without keeping any state between calls. When a cutoff is configured, VisualGen includes the dense-or-sparse phase in the CUDA Graph key. The TRTLLM attention metadata reduces the timestep to a host value during graph warmup, keeps it in the component attention state and reuses it during capture for every timestep-scheduled algorithm (Skip Softmax Attention -and SOL), while SOL -predictor route buffers remain stable for replay; the CuTeDSL backends read the +and SOL), while the SOL predictor's outputs are allocated inside the captured +graph and replayed with it; the CuTeDSL backends read the phase the CUDA Graph runner resolved for the graph key instead of the device tensor. Per-token timesteps, such as Wan I2V where the conditioning frame stays at timestep zero, reduce to their largest live value, so the schedule stays 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 index 844254c75a7f..c93e7392ffc3 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/__init__.py @@ -5,21 +5,13 @@ from .backend import SOLCuTeDSLAttention, SOLTrtllmAttention from .params import SolParams -from .predictor import ( - SolPredictorGeometry, - SolPredictorOutputs, - SolPredictorPlan, - SolPredictorPlanKey, - SOLSparsePredictor, -) +from .predictor import SolPredictorOutputs, predict, support_reason __all__ = [ "SOLCuTeDSLAttention", - "SOLSparsePredictor", "SOLTrtllmAttention", "SolParams", - "SolPredictorGeometry", "SolPredictorOutputs", - "SolPredictorPlan", - "SolPredictorPlanKey", + "predict", + "support_reason", ] 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 index 3c851f31e381..6cbc44d5db59 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/backend.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/backend.py @@ -19,9 +19,9 @@ serve one ``SolParams``: * ``SOLTrtllmAttention`` runs SOL in two stages through the generic TRTLLM - sparse lifecycle: a TRT-LLM-owned predictor derives the exact block bitmask - and K/V proxy summaries, then the shared PrimTS block-sparse FMHA executes - that route. + sparse lifecycle: one graph-visible predictor operator derives the exact + block bitmask and K/V proxy summaries, then the shared PrimTS block-sparse + FMHA executes that route. * ``SOLCuTeDSLAttention`` runs the fused kernel vendored from the reference implementation (https://github.com/NVlabs/Sana, branch ``sol-engine``, pinned in ``cute_dsl_kernels/blackwell/sol_attn/THIRD_PARTY_NOTICES.md``), which @@ -58,8 +58,9 @@ from ...interface import AttentionBackend, AttentionTensorLayout from ...trtllm import TrtllmAttention from ...vanilla import VanillaAttention +from . import predictor as sol_predictor from .params import SolParams -from .predictor import BLOCK_SIZE, SOLSparsePredictor +from .predictor import BLOCK_SIZE _sol_attn_import_error = None try: @@ -332,7 +333,6 @@ def __init__(self, *, sparse_params: SolParams | None = None, **kwargs) -> None: raise TypeError("SOLTrtllmAttention requires SolParams") self.sol_params = sparse_params super().__init__(sparse_params=None, **kwargs) - self.predictor = SOLSparsePredictor() @property def timestep_cutoff(self) -> Optional[float]: @@ -388,16 +388,17 @@ def block_sparse_attn_predict( 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) + unsupported_reason = sol_predictor.support_reason(q, k, v) if unsupported_reason is not None: raise ValueError(unsupported_reason) - outputs = self.predictor.predict( + outputs = sol_predictor.predict( q, k, v, tau=self.sol_params.tau, sm_scale=get_bmm1_scale(self), + thresh_type=self.sol_params.thresh_type, ) return BlockSparseForwardInputs( q_block_size=BLOCK_SIZE, diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/kernels.py b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/kernels.py index 139a034b4aeb..b241f79aca33 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/kernels.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/kernels.py @@ -2,11 +2,11 @@ # 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. +The predictor summarises ``[batch, tokens, heads, head_dim]`` activations per token block, derives a +routing threshold per query block from the key block statistics, and thresholds centroid scores into +packed exact-block words. CUDA tensors run Triton kernels (block pooling, key statistics, thresholds and +selection); 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. """ from __future__ import annotations @@ -20,15 +20,17 @@ _POOL_MAX_WIDTH = 1024 _POOL_TOKENS_PER_LOAD = 8 -_STATS_WIDTH = 128 -_STATS_ROWS_PER_LOAD = 32 _SELECT_Q_BLOCKS = 64 +_STATS_ROWS = 32 +_STATS_COLS = 64 +_THRESHOLD_Q_BLOCKS = 32 _WORD_BITS = 32 _LOG2_E = math.log2(math.e) _THRESHOLD_EPSILON = 1.0e-6 _LOCAL_RADIUS = 1 _Reduce = Literal["mean", "sum"] +_ThreshType = Literal["diag", "exact"] def _column_launch(row_width: int, max_width: int) -> tuple[int, int]: @@ -155,90 +157,290 @@ def block_pool(x: torch.Tensor, out: torch.Tensor, *, block_size: int, reduce: _ ) -# --------------------------------------------------------------------------- block statistics +# --------------------------------------------------------------------------- block thresholds @triton.jit def _block_statistics_kernel( x_ptr, mean_ptr, var_ptr, + moment_ptr, num_blocks, - num_chunks, - row_width, + num_heads, stride_x_batch, stride_x_block, - stride_out_batch, + stride_x_head, + stride_s_batch, + stride_s_head, + stride_m_batch, + stride_m_head, + EXACT: tl.constexpr, ROWS: tl.constexpr, - WIDTH: tl.constexpr, + COLS: tl.constexpr, + HEAD_DIM: 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) + """One program per (batch, head, column chunk): key block statistics over the block axis. + + Every program writes the per-channel mean and clamped variance of its columns. With ``EXACT`` + it also accumulates its columns of the raw second moment ``E[k k^T]`` with tensor-core dots on + the 16-bit summaries, whose products are exact in the fp32 accumulator. + """ + chunk = tl.program_id(0) + batch_head = tl.program_id(1).to(tl.int64) + batch = batch_head // num_heads + head = batch_head % num_heads + cols = chunk * COLS + tl.arange(0, COLS) + dims = tl.arange(0, HEAD_DIM) + base = x_ptr + batch * stride_x_batch + head * stride_x_head + total = tl.zeros([COLS], dtype=tl.float32) + total_sq = tl.zeros([COLS], dtype=tl.float32) + if EXACT: + moment = tl.zeros([HEAD_DIM, COLS], 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, :], + row_valid = rows < num_blocks + tile = tl.load( + base + rows[:, None] * stride_x_block + cols[None, :], + mask=row_valid[:, None], other=0.0, - ).to(tl.float32) + ) + values = tile.to(tl.float32) total += tl.sum(values, axis=0) total_sq += tl.sum(values * values, axis=0) + if EXACT: + full = tl.load( + base + rows[:, None] * stride_x_block + dims[None, :], + mask=row_valid[:, None], + other=0.0, + ) + moment += tl.dot(tl.trans(full), tile) 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) + stats = batch * stride_s_batch + head * stride_s_head + cols + tl.store(mean_ptr + stats, mean) + tl.store(var_ptr + stats, tl.maximum(total_sq / count - mean * mean, 0.0)) + if EXACT: + tl.store( + moment_ptr + + batch * stride_m_batch + + head * stride_m_head + + dims[:, None] * HEAD_DIM + + cols[None, :], + moment / count, + ) + + +@triton.jit +def _block_thresholds_kernel( + centroid_ptr, + mean_ptr, + var_ptr, + moment_ptr, + threshold_ptr, + num_q_blocks, + num_heads, + tau, + log2_scale, + epsilon, + stride_c_batch, + stride_c_block, + stride_c_head, + stride_s_batch, + stride_s_head, + stride_m_batch, + stride_m_head, + stride_t_batch, + stride_t_head, + EXACT: tl.constexpr, + Q_BLOCKS: tl.constexpr, + HEAD_DIM: tl.constexpr, +): + """One program per (batch, head, tile of Q_BLOCKS query blocks): the routing threshold per row. + ``mean + tau * sqrt(var + epsilon)`` in the log2 score domain, where the variance of the + centroid score is either the diagonal projection of the per-channel key variance or the full + quadratic form with the key second moment (``EXACT``); the fp32 dot keeps IEEE precision. + """ + 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, + ) + stats = batch * stride_s_batch + head * stride_s_head + dims + mean = tl.load(mean_ptr + stats) + projected_mean = tl.sum(centroid * mean[None, :], axis=1) + if EXACT: + moment = tl.load( + moment_ptr + + batch * stride_m_batch + + head * stride_m_head + + dims[:, None] * HEAD_DIM + + dims[None, :] + ) + projected = tl.dot(centroid, moment, input_precision="ieee") + projected_var = tl.maximum( + tl.sum(projected * centroid, axis=1) - projected_mean * projected_mean, 0.0 + ) + else: + var = tl.load(var_ptr + stats) + projected_var = tl.sum(centroid * centroid * var[None, :], axis=1) + threshold = projected_mean * log2_scale + tau * tl.sqrt( + projected_var * (log2_scale * log2_scale) + epsilon + ) + tl.store( + threshold_ptr + batch * stride_t_batch + head * stride_t_head + q_blocks, + threshold, + mask=q_valid, + ) -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_thresholds_torch( + centroid: torch.Tensor, + k_summary: torch.Tensor, + *, + tau: float, + sm_scale: float, + thresh_type: _ThreshType, +) -> torch.Tensor: + # float64 keeps this path a precise reference for the kernels; TensorRT-LLM + # enables TF32 matmuls, which would perturb the fp32 second moment. + log2_scale = float(sm_scale) * _LOG2_E + c = centroid.to(torch.float64).permute(0, 2, 1, 3) + keys = k_summary.to(torch.float64).permute(0, 2, 1, 3) + k_mean = keys.mean(dim=2) + mean = torch.einsum("bhqd,bhd->bhq", c, k_mean) + if thresh_type == "diag": + k_var = torch.clamp(keys.square().mean(dim=2) - k_mean.square(), min=0.0) + var = torch.einsum("bhqd,bhd->bhq", c.square(), k_var) + else: + second_moment = torch.matmul(keys.transpose(-1, -2), keys) / keys.shape[2] + var = torch.clamp( + torch.einsum("bhqd,bhqd->bhq", torch.matmul(c, second_moment), c) - mean.square(), + min=0.0, + ) + threshold = mean * log2_scale + float(tau) * torch.sqrt( + var * (log2_scale * log2_scale) + _THRESHOLD_EPSILON + ) + return threshold.to(torch.float32).contiguous() -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. + +def block_thresholds( + centroid: torch.Tensor, + k_summary: torch.Tensor, + *, + tau: float, + sm_scale: float, + thresh_type: _ThreshType, +) -> torch.Tensor: + """Routing threshold ``mean + tau * sqrt(var + 1e-6)`` of every query block in the log2 score domain. + + The threshold models the score of the query block centroid ``c`` against a key block drawn from the + observed key block summaries: ``mean = `` with either the diagonal variance + ``sum_d c_d^2 Var[k_d]`` (``"diag"``) or the full covariance ``c^T Cov[k] c`` (``"exact"``), both in + the ``sm_scale * log2(e)`` scaled domain the selection kernel scores in. CUDA tensors with 16-bit + key summaries and a power-of-two head dimension run two Triton kernels, key statistics and then + thresholds; other tensors use the float64 PyTorch implementation of the same rule. 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. + 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. + tau: Threshold slope in standard deviations. + sm_scale: Softmax scale of the attention call. + thresh_type: ``"diag"`` or ``"exact"``. + + Returns: + Contiguous fp32 ``[batch, heads, num_q_blocks]`` thresholds. """ - 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, + if thresh_type not in ("diag", "exact"): + raise ValueError(f"thresh_type must be 'diag' or 'exact'; got {thresh_type!r}") + batch_size, q_blocks, num_heads, head_dim = centroid.shape + kv_blocks = k_summary.shape[1] + 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") + triton_ready = ( + centroid.device.type == "cuda" + and centroid.dtype == torch.float32 + and centroid.is_contiguous() + and k_summary.is_contiguous() + and k_summary.dtype in (torch.bfloat16, torch.float16) + and head_dim >= _STATS_COLS + and head_dim & (head_dim - 1) == 0 + ) + if not triton_ready: + return _block_thresholds_torch( + centroid, k_summary, tau=tau, sm_scale=sm_scale, thresh_type=thresh_type + ) + exact = thresh_type == "exact" + mean = torch.empty( + (batch_size, num_heads, head_dim), dtype=torch.float32, device=centroid.device + ) + var = torch.empty_like(mean) + moment = ( + torch.empty( + (batch_size, num_heads, head_dim, head_dim), dtype=torch.float32, device=centroid.device + ) + if exact + else mean + ) + _block_statistics_kernel[(head_dim // _STATS_COLS, batch_size * num_heads)]( + k_summary, + mean, + var, + moment, + kv_blocks, + num_heads, + k_summary.stride(0), + k_summary.stride(1), + k_summary.stride(2), + mean.stride(0), + mean.stride(1), + moment.stride(0), + moment.stride(1), + EXACT=exact, + ROWS=_STATS_ROWS, + COLS=_STATS_COLS, + HEAD_DIM=head_dim, num_warps=4, ) + threshold = torch.empty( + (batch_size, num_heads, q_blocks), dtype=torch.float32, device=centroid.device + ) + _block_thresholds_kernel[(triton.cdiv(q_blocks, _THRESHOLD_Q_BLOCKS), batch_size * num_heads)]( + centroid, + mean, + var, + moment, + threshold, + q_blocks, + num_heads, + float(tau), + float(sm_scale) * _LOG2_E, + _THRESHOLD_EPSILON, + centroid.stride(0), + centroid.stride(1), + centroid.stride(2), + mean.stride(0), + mean.stride(1), + moment.stride(0), + moment.stride(1), + threshold.stride(0), + threshold.stride(1), + EXACT=exact, + Q_BLOCKS=_THRESHOLD_Q_BLOCKS, + HEAD_DIM=head_dim, + num_warps=4, + ) + return threshold # --------------------------------------------------------------------------- exact-block selection @@ -246,25 +448,22 @@ def block_statistics(x: torch.Tensor, out_mean: torch.Tensor, out_var: torch.Ten def _select_exact_blocks_kernel( centroid_ptr, keys_ptr, - mean_ptr, - var_ptr, + threshold_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_t_batch, + stride_t_head, stride_b_batch, stride_b_head, stride_b_block, @@ -293,11 +492,11 @@ def _select_exact_blocks_kernel( 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) + threshold = tl.load( + threshold_ptr + batch * stride_t_batch + head * stride_t_head + q_blocks, + mask=q_valid, + other=0.0, + ) high = centroid.to(keys_ptr.dtype.element_ty) rest = centroid - high.to(tl.float32) @@ -338,25 +537,16 @@ def _select_exact_blocks_kernel( def _select_exact_blocks_torch( centroid: torch.Tensor, k_summary: torch.Tensor, - k_mean: torch.Tensor, - k_var: torch.Tensor, + threshold: 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) + exact = scores > threshold.to(torch.float64).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] @@ -372,27 +562,22 @@ def _select_exact_blocks_torch( def select_exact_blocks( centroid: torch.Tensor, k_summary: torch.Tensor, - k_mean: torch.Tensor, - k_var: torch.Tensor, + threshold: 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. + A key block is exact when ``sm_scale * log2(e) * `` exceeds the row's + ``threshold`` (see ``block_thresholds``) 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. + threshold: Contiguous fp32 ``[batch, heads, num_q_blocks]`` row thresholds. 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 @@ -402,6 +587,13 @@ def select_exact_blocks( 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") + expected_threshold = (batch_size, num_heads, q_blocks) + if ( + tuple(threshold.shape) != expected_threshold + or threshold.dtype != torch.float32 + or not threshold.is_contiguous() + ): + raise ValueError(f"threshold must be contiguous fp32 of shape {expected_threshold}") if ( centroid.dtype != torch.float32 or not centroid.is_contiguous() @@ -412,7 +604,7 @@ def select_exact_blocks( 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 + centroid, k_summary, threshold, exact_block_bits, sm_scale=sm_scale ) return bits = exact_block_bits.view(torch.int32) @@ -420,25 +612,22 @@ def select_exact_blocks( _select_exact_blocks_kernel[grid]( centroid, k_summary, - k_mean, - k_var, + threshold, 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), + threshold.stride(0), + threshold.stride(1), bits.stride(0), bits.stride(1), bits.stride(2), @@ -449,41 +638,40 @@ def select_exact_blocks( # --------------------------------------------------------------------------- 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", -) +@torch.library.custom_op("trtllm::visual_gen_sol_predictor", mutates_args=(), 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.""" + thresh_type: str, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Return ``(exact_block_bits, k_summary, v_summary)`` for compact BSHD ``q``, ``k`` and ``v``. + The outputs are fresh tensors of this call. Inside CUDA Graph capture they come from the graph + pool and stay valid for every replay; under torch.compile the operator is opaque and shapes come + from the fake implementation. + """ + + batch_size, seq_len, num_heads, head_dim = q.shape + blocks = num_blocks(seq_len, block_size) + summary_shape = (batch_size, blocks, num_heads, head_dim) + q_centroid = torch.empty(summary_shape, dtype=torch.float32, device=q.device) + k_summary = torch.empty(summary_shape, dtype=k.dtype, device=k.device) + v_summary = torch.empty(summary_shape, dtype=v.dtype, device=v.device) block_pool(q, q_centroid, block_size=block_size, reduce="mean") block_pool(k, k_summary, block_size=block_size, reduce="mean") block_pool(v, v_summary, block_size=block_size, reduce="sum") - block_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 + threshold = block_thresholds( + q_centroid, k_summary, tau=tau, sm_scale=sm_scale, thresh_type=thresh_type ) + exact_block_bits = torch.empty( + (batch_size, num_heads, blocks, num_words(blocks)), dtype=torch.uint32, device=q.device + ) + select_exact_blocks(q_centroid, k_summary, threshold, exact_block_bits, sm_scale=sm_scale) + return exact_block_bits, k_summary, v_summary @torch.library.register_fake("trtllm::visual_gen_sol_predictor") @@ -491,17 +679,19 @@ 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 + thresh_type: str, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + batch_size, seq_len, num_heads, head_dim = q.shape + blocks = num_blocks(seq_len, block_size) + summary_shape = (batch_size, blocks, num_heads, head_dim) + return ( + q.new_empty((batch_size, num_heads, blocks, num_words(blocks)), dtype=torch.uint32), + k.new_empty(summary_shape), + v.new_empty(summary_shape), + ) -__all__ = ["block_pool", "block_statistics", "num_blocks", "num_words", "select_exact_blocks"] +__all__ = ["block_pool", "block_thresholds", "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 index 1c5defb10995..f8a2c7e79c28 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/params.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/params.py @@ -31,8 +31,7 @@ class SolParams(SparseParams): """Static SOL policy lowered from the user-facing VisualGen config. Shared by the TRTLLM (predictor + block-sparse FMHA) and CuTeDSL (fused kernel) - backends; ``thresh_type`` selects the CuTeDSL kernel threshold policy, the TRTLLM - predictor implements ``diag`` only. + backends; ``thresh_type`` selects the block routing threshold policy of both. """ algorithm: Literal["sol_attn"] = field(init=False, default="sol_attn") 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 index 4fa91d52e09e..231826d659e9 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/predictor.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/predictor.py @@ -1,282 +1,102 @@ # 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.""" +"""Two-stage VisualGen SOL predictor: exact-block routes and K/V proxy summaries from Q/K/V.""" from __future__ import annotations +import math import numbers -import struct from dataclasses import dataclass +from typing import Literal import torch -from . import kernels as _kernels +from . import kernels as _kernels # noqa: F401 (registers trtllm::visual_gen_sol_predictor) 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 +ThreshType = Literal["diag", "exact"] @dataclass(frozen=True) class SolPredictorOutputs: - """Live predictor tensors consumed by block-sparse attention.""" + """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 +def support_reason(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> str | None: + """Return why compact two-stage SOL cannot serve these tensors, or ``None``.""" + + 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 + + +def _runtime_scalar(value: object, name: str, *, positive: bool = False) -> float: + if isinstance(value, bool) or not isinstance(value, numbers.Real): + raise TypeError(f"{name} must be a finite Python real") + result = float(value) + if not math.isfinite(result): + raise ValueError(f"{name} must be finite") + if positive and result <= 0.0: + raise ValueError(f"{name} must be positive") + return result -__all__ = [ - "SOLSparsePredictor", - "SolPredictorGeometry", - "SolPredictorOutputs", - "SolPredictorPlan", - "SolPredictorPlanKey", -] +def predict( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + tau: object, + sm_scale: object, + thresh_type: ThreshType = "diag", +) -> SolPredictorOutputs: + """Predict SOL routes and proxy summaries for compact BF16 self-attention tensors. + + One graph-visible operator pools ``q``, ``k`` and ``v`` per ``BLOCK_SIZE`` tokens, derives the + routing threshold of every query block from the key block statistics (``thresh_type`` selects + the diagonal or the full key covariance), and packs the exact-block decisions into + ``exact_block_bits``. The outputs are fresh tensors: inside CUDA Graph capture they come from the + graph pool and stay valid for replay, and under torch.compile the operator is opaque. + """ + + reason = support_reason(q, k, v) + if reason is not None: + raise ValueError(reason) + if thresh_type not in ("diag", "exact"): + raise ValueError(f"thresh_type must be 'diag' or 'exact'; got {thresh_type!r}") + exact_block_bits, k_summary, v_summary = torch.ops.trtllm.visual_gen_sol_predictor( + q, + k, + v, + BLOCK_SIZE, + _runtime_scalar(tau, "tau"), + _runtime_scalar(sm_scale, "sm_scale", positive=True), + thresh_type, + ) + return SolPredictorOutputs( + exact_block_bits=exact_block_bits, k_summary=k_summary, v_summary=v_summary + ) + + +__all__ = ["BLOCK_SIZE", "HEAD_DIM", "SolPredictorOutputs", "predict", "support_reason"] diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index 33b70f0b2153..b08d322f95a0 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -255,22 +255,6 @@ def _validate_quant_sparse_mutex(self) -> "AttentionConfig": raise ValueError("SOL and quant_attention_config are mutually exclusive.") return self - @model_validator(mode="after") - def _validate_sol_thresh_type(self) -> "AttentionConfig": - # The TRTLLM predictor implements the diag threshold only; the exact - # policy exists in the fused CUTEDSL kernel. - sparse_config = self.sparse_attention_config - if ( - isinstance(sparse_config, SolAttentionConfig) - and self.backend == "TRTLLM" - and sparse_config.thresh_type != "diag" - ): - raise ValueError( - f"TRTLLM SOL supports thresh_type='diag' only, got " - f"thresh_type={sparse_config.thresh_type!r}; use backend='CUTEDSL'." - ) - return self - class VAEConfig(StrictBaseModel): """Configuration for the variational autoencoder.""" @@ -807,21 +791,6 @@ 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] diff --git a/tensorrt_llm/visual_gen/sparse_attention.py b/tensorrt_llm/visual_gen/sparse_attention.py index 66f74639d9ba..82f5864fbf86 100644 --- a/tensorrt_llm/visual_gen/sparse_attention.py +++ b/tensorrt_llm/visual_gen/sparse_attention.py @@ -269,8 +269,8 @@ class SolAttentionConfig(BaseSparseAttentionConfig): thresh_type: Literal["diag", "exact"] = PydanticField( "diag", description=( - "Threshold policy of the CUTEDSL kernel (kernel default: 'diag'). " - "The TRTLLM predictor implements 'diag' only." + "Block routing threshold policy: 'diag' models each key channel " + "independently, 'exact' uses the full key covariance." ), ) disabled_until_timestep: Optional[float] = PydanticField( 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 index 3d0cce51fd95..5dc8679c05f0 100644 --- a/tests/unittest/_torch/visual_gen/sparse_attention/test_sol_attention.py +++ b/tests/unittest/_torch/visual_gen/sparse_attention/test_sol_attention.py @@ -33,11 +33,11 @@ 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 import predictor as sol_predictor 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, @@ -62,12 +62,7 @@ _CPU_ONLY = pytest.mark.cpu_only -def _make_backend( - params: SolParams, - predictor: Mock, - *, - layer_idx: int = 1, -) -> SOLTrtllmAttention: +def _make_backend(params: SolParams, *, layer_idx: int = 1) -> SOLTrtllmAttention: backend = object.__new__(SOLTrtllmAttention) backend.layer_idx = layer_idx backend.num_heads = 2 @@ -81,7 +76,6 @@ def _make_backend( backend.metadata = TrtllmAttentionMetadata( device=torch.device("cpu"), attention_metadata_state={} ) - backend.predictor = predictor return backend @@ -178,15 +172,20 @@ def _bshd(seq_len: int = 64, num_heads: int = 2) -> torch.Tensor: def _stub_backend( + monkeypatch, 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 +) -> tuple[SOLTrtllmAttention, SimpleNamespace]: + """Backend whose predictor functions are recorded mocks.""" + predictor = SimpleNamespace( + support_reason=Mock(return_value=unsupported_reason), + predict=Mock(return_value=_predictor_outputs(batch_size=1, seq_len=seq_len, num_heads=2)), + ) + monkeypatch.setattr(sol_predictor, "support_reason", predictor.support_reason) + monkeypatch.setattr(sol_predictor, "predict", predictor.predict) + return _make_backend(params or SolParams(tau=1.0)), predictor def _dense_reference(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: @@ -259,7 +258,9 @@ def _mixed_proxy_reference( @_CPU_ONLY def test_sol_backend_reuses_prepared_timestep_during_cuda_graph_capture(monkeypatch) -> None: - backend, _predictor = _stub_backend(SolParams(tau=1.0, disabled_until_timestep=0.6)) + backend, _predictor = _stub_backend( + monkeypatch, SolParams(tau=1.0, disabled_until_timestep=0.6) + ) # Per-token timesteps reduce to the largest live value. assert backend.resolve_timestep(torch.tensor([0.0, 0.2])) == pytest.approx(0.2) @@ -274,7 +275,7 @@ def test_sol_backend_reuses_prepared_timestep_during_cuda_graph_capture(monkeypa @_CPU_ONLY def test_sol_backend_warmup_prepares_dense_phase_for_capture(monkeypatch) -> None: q = _bshd() - backend, predictor = _stub_backend(SolParams(tau=1.0, disabled_until_timestep=0.6)) + backend, predictor = _stub_backend(monkeypatch, SolParams(tau=1.0, disabled_until_timestep=0.6)) assert _predict(backend, q, q, q, timestep=backend.resolve_timestep(torch.tensor(0.8))) is None monkeypatch.setattr(torch.cuda, "is_available", lambda: True) @@ -285,7 +286,7 @@ def test_sol_backend_warmup_prepares_dense_phase_for_capture(monkeypatch) -> Non @_CPU_ONLY def test_sol_backend_rejects_cutoff_capture_without_warmup(monkeypatch) -> None: - backend, predictor = _stub_backend(SolParams(tau=1.0, disabled_until_timestep=0.6)) + backend, predictor = _stub_backend(monkeypatch, SolParams(tau=1.0, disabled_until_timestep=0.6)) monkeypatch.setattr(torch.cuda, "is_available", lambda: True) monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) @@ -296,17 +297,19 @@ def test_sol_backend_rejects_cutoff_capture_without_warmup(monkeypatch) -> None: @_CPU_ONLY -def test_sol_backend_without_timestep_runs_sparse_like_skip_softmax() -> None: +def test_sol_backend_without_timestep_runs_sparse_like_skip_softmax(monkeypatch) -> None: q = _bshd() - backend, predictor = _stub_backend(SolParams(tau=1.0, disabled_until_timestep=0.6)) + backend, predictor = _stub_backend(monkeypatch, SolParams(tau=1.0, disabled_until_timestep=0.6)) assert _predict(backend, q, q, q, timestep=backend.resolve_timestep(None)) is not None predictor.predict.assert_called_once() @_CPU_ONLY -def test_sol_phase_waits_until_all_token_timesteps_are_below_cutoff() -> None: - backend, _predictor = _stub_backend(SolParams(tau=1.0, disabled_until_timestep=0.6)) +def test_sol_phase_waits_until_all_token_timesteps_are_below_cutoff(monkeypatch) -> None: + backend, _predictor = _stub_backend( + monkeypatch, SolParams(tau=1.0, disabled_until_timestep=0.6) + ) assert not backend.should_use_sparse(backend.resolve_timestep(torch.tensor([0.0, 0.8]))) assert backend.should_use_sparse(backend.resolve_timestep(torch.tensor([0.0, 0.2]))) @@ -352,7 +355,7 @@ def _base_init(self, **kwargs) -> None: 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 not hasattr(backend, "predictor") assert base_kwargs["sparse_params"] is None assert "_enable_sparse_workflow" not in SOLTrtllmAttention.__dict__ assert "_should_use_sparse_workflow" not in SOLTrtllmAttention.__dict__ @@ -370,14 +373,18 @@ def test_sol_backend_sparse_phase_emits_proxy_bitmask_carrier(monkeypatch) -> No seq_len=seq_len, num_heads=num_heads, ) - backend, predictor = _stub_backend(SolParams(tau=0.75), seq_len=seq_len) + backend, predictor = _stub_backend(monkeypatch, 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} + assert predictor.predict.call_args.kwargs == { + "tau": 0.75, + "sm_scale": 0.375, + "thresh_type": "diag", + } 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) @@ -411,7 +418,7 @@ def test_sol_wrapper_compacts_separate_qkv_and_predicts_inside_core(monkeypatch) seq_len=seq_len, num_heads=num_heads, ) - backend, predictor = _stub_backend(SolParams(tau=0.75), seq_len=seq_len) + backend, predictor = _stub_backend(monkeypatch, 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) @@ -442,13 +449,14 @@ def test_sol_wrapper_compacts_separate_qkv_and_predicts_inside_core(monkeypatch) ), ) def test_sol_backend_rejects_non_sol_sparse_calls( + monkeypatch, k: torch.Tensor | None, v: torch.Tensor | None, attention_mask: PredefinedAttentionMask, message: str, ) -> None: q = _bshd() - backend, predictor = _stub_backend() + backend, predictor = _stub_backend(monkeypatch) with pytest.raises(ValueError, match=message): _predict(backend, q, k, v, attention_mask=attention_mask) @@ -459,7 +467,7 @@ def test_sol_backend_rejects_non_sol_sparse_calls( @_CPU_ONLY def test_sol_wrapper_rejects_fused_qkv_before_core(monkeypatch) -> None: q = _bshd() - backend, predictor = _stub_backend() + backend, predictor = _stub_backend(monkeypatch) prepare_metadata = Mock(return_value=object()) monkeypatch.setattr(TrtllmAttention, "_prepare_metadata", prepare_metadata) @@ -471,10 +479,10 @@ def test_sol_wrapper_rejects_fused_qkv_before_core(monkeypatch) -> None: @_CPU_ONLY -def test_sol_backend_surfaces_predictor_support_reason_before_execution() -> None: +def test_sol_backend_surfaces_predictor_support_reason_before_execution(monkeypatch) -> None: q = _bshd() reason = "SOL predictor requires compact BSHD q/k/v" - backend, predictor = _stub_backend(unsupported_reason=reason) + backend, predictor = _stub_backend(monkeypatch, unsupported_reason=reason) with pytest.raises(ValueError, match=reason): _predict(backend, q, q, q) @@ -491,12 +499,13 @@ def test_sol_backend_surfaces_predictor_support_reason_before_execution() -> Non ), ) def test_sol_dense_policy_returns_no_routes_without_predicting( + monkeypatch, params: SolParams, layer_idx: int, timestep: float | None, ) -> None: q = _bshd() - backend, predictor = _stub_backend(params) + backend, predictor = _stub_backend(monkeypatch, params) backend.layer_idx = layer_idx assert _predict(backend, q, q, q, timestep=timestep) is None @@ -505,9 +514,9 @@ def test_sol_dense_policy_returns_no_routes_without_predicting( @_CPU_ONLY -def test_sol_sparse_phase_without_primts_fails_closed() -> None: +def test_sol_sparse_phase_without_primts_fails_closed(monkeypatch) -> None: q = _bshd() - backend, predictor = _stub_backend() + backend, predictor = _stub_backend(monkeypatch) backend._fmha_manager = SimpleNamespace(fmha_libs=[]) with pytest.raises(RuntimeError, match="requires PrimTS block-sparse FMHA"): @@ -518,9 +527,9 @@ def test_sol_sparse_phase_without_primts_fails_closed() -> None: @_CPU_ONLY -def test_sol_sparse_phase_with_quantization_fails_closed() -> None: +def test_sol_sparse_phase_with_quantization_fails_closed(monkeypatch) -> None: q = _bshd() - backend, predictor = _stub_backend() + backend, predictor = _stub_backend(monkeypatch) backend.quant_attention_config = object() with pytest.raises(ValueError, match="does not support quant_attention_config"): @@ -554,17 +563,13 @@ def test_sol_public_config_requires_supported_backend() -> None: @_CPU_ONLY -def test_sol_exact_threshold_needs_cutedsl_backend() -> None: - with pytest.raises(ValidationError, match="thresh_type"): - AttentionConfig( - backend="TRTLLM", +def test_sol_exact_threshold_lowers_for_both_backends() -> None: + for backend in ("TRTLLM", "CUTEDSL"): + config = AttentionConfig( + backend=backend, sparse_attention_config=SolAttentionConfig(thresh_type="exact"), ) - config = AttentionConfig( - backend="CUTEDSL", - sparse_attention_config=SolAttentionConfig(thresh_type="exact"), - ) - assert config.sparse_attention_config.to_sparse_params().thresh_type == "exact" + assert config.sparse_attention_config.to_sparse_params().thresh_type == "exact" @_CPU_ONLY @@ -684,7 +689,7 @@ def test_sol_attention_rejects_context_parallelism() -> None: @_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)) + backend, predictor = _stub_backend(monkeypatch, 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) @@ -801,7 +806,6 @@ def _inputs() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: rtol=2e-2, atol=2e-2, ) - assert backend.predictor.num_plans == 1 @_REQUIRES_SM100 @@ -842,7 +846,7 @@ def _inputs() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: 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( + predictor_outputs = sol_predictor.predict( q.contiguous(), k.contiguous(), v.contiguous(), @@ -887,9 +891,18 @@ def _inputs() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: graph.replay() torch.cuda.synchronize() + # The captured graph re-predicts from the refreshed inputs; the reference + # needs the summaries of those inputs too. + replayed_outputs = sol_predictor.predict( + q.contiguous(), + k.contiguous(), + v.contiguous(), + tau=1.0e6, + sm_scale=128**-0.5, + ) torch.testing.assert_close( captured, - _mixed_proxy_reference(q, k, v, predictor_outputs), + _mixed_proxy_reference(q, k, v, replayed_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 index 15b31deb210b..d0f597643990 100644 --- a/tests/unittest/_torch/visual_gen/sparse_attention/test_sol_predictor.py +++ b/tests/unittest/_torch/visual_gen/sparse_attention/test_sol_predictor.py @@ -4,97 +4,65 @@ 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.kernels import ( + block_pool, + block_thresholds, +) from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.predictor import ( - SolPredictorGeometry, + BLOCK_SIZE, SolPredictorOutputs, - SolPredictorPlanKey, - SOLSparsePredictor, - _normalize_runtime_scalars, + _runtime_scalar, + predict, + support_reason, ) _CPU_ONLY = pytest.mark.cpu_only +_REQUIRES_CUDA = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +_LOG2_E = math.log2(math.e) +_POLICIES = ("diag", "exact") @_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", +def test_sol_predictor_support_reason_covers_layout_shape_dtype_and_device() -> None: + good = torch.zeros(1, 64, 1, 128, dtype=torch.bfloat16) + cases = ( + ((good.view(64, 128), good, good), "compact BSHD"), + ((good, torch.zeros(1, 65, 1, 128, dtype=torch.bfloat16), good), "uniform self-attention"), + ((good, good.float(), good), "matching BF16"), + ((good, good, good), "requires CUDA"), + ((None, good, good), "torch tensors"), ) - assert "tau" not in repr(key) and "sm_scale" not in repr(key) - assert SOLSparsePredictor().num_plans == 0 + for tensors, message in cases: + assert message in support_reason(*tensors) + with pytest.raises(ValueError, match="requires CUDA"): + predict(good, good, good, tau=0.5, sm_scale=0.125) @_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"), +def test_sol_predictor_validates_runtime_scalars() -> None: + assert _runtime_scalar(0.1, "tau") == 0.1 + assert _runtime_scalar(2, "sm_scale", positive=True) == 2.0 + invalid = ( + (True, "tau", {}), + (math.nan, "tau", {}), + (math.inf, "sm_scale", {"positive": True}), + (0.0, "sm_scale", {"positive": True}), + (-0.125, "sm_scale", {"positive": True}), + ("0.5", "tau", {}), ) - 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) + for value, name, kwargs in invalid: + with pytest.raises((TypeError, ValueError), match=name): + _runtime_scalar(value, name, **kwargs) def _summary_oracle(k: torch.Tensor, v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - blocks = (k.shape[1] + 63) // 64 + blocks = (k.shape[1] + BLOCK_SIZE - 1) // BLOCK_SIZE k_summary = torch.empty( (k.shape[0], blocks, k.shape[2], k.shape[3]), dtype=torch.bfloat16, @@ -102,8 +70,8 @@ def _summary_oracle(k: torch.Tensor, v: torch.Tensor) -> tuple[torch.Tensor, tor ) v_summary = torch.empty_like(k_summary) for block_idx in range(blocks): - begin = block_idx * 64 - end = min(begin + 64, k.shape[1]) + begin = block_idx * BLOCK_SIZE + end = min(begin + BLOCK_SIZE, k.shape[1]) k_summary[:, block_idx] = k[:, begin:end].float().mean(dim=1).to(torch.bfloat16) v_summary[:, block_idx] = v[:, begin:end].float().sum(dim=1).to(torch.bfloat16) return k_summary, v_summary @@ -116,6 +84,40 @@ def _pack_bits(exact: torch.Tensor) -> torch.Tensor: return (padded.to(torch.int64) * powers).sum(dim=-1).to(torch.uint32) +def _centroid_oracle(q: torch.Tensor) -> torch.Tensor: + blocks = (q.shape[1] + BLOCK_SIZE - 1) // BLOCK_SIZE + padded_q = F.pad(q, (0, 0, 0, 0, 0, blocks * BLOCK_SIZE - q.shape[1])) + q_blocks = padded_q.view(q.shape[0], blocks, BLOCK_SIZE, q.shape[2], q.shape[3]) + q_lengths = torch.clamp( + q.shape[1] - torch.arange(blocks, device=q.device) * BLOCK_SIZE, min=1, max=BLOCK_SIZE + ) + return q_blocks.double().sum(dim=2) / q_lengths[None, :, None, None] + + +def _threshold_oracle( + q_centroids: torch.Tensor, + k_summary: torch.Tensor, + *, + tau: float, + sm_scale: float, + thresh_type: str, +) -> torch.Tensor: + """fp64 ``[batch, heads, num_q_blocks]`` thresholds of the selected policy.""" + c = q_centroids.double().permute(0, 2, 1, 3) + keys = k_summary.double().permute(0, 2, 1, 3) + k_mean = keys.mean(dim=2) + mean = torch.einsum("bhqd,bhd->bhq", c, k_mean) + if thresh_type == "diag": + k_var = torch.clamp(keys.square().mean(dim=2) - k_mean.square(), min=0.0) + var = torch.einsum("bhqd,bhd->bhq", c.square(), k_var) + else: + centered = keys - k_mean.unsqueeze(2) + covariance = torch.matmul(centered.transpose(-1, -2), centered) / keys.shape[2] + var = torch.einsum("bhqd,bhqd->bhq", torch.matmul(c, covariance), c) + log2_scale = float(sm_scale) * _LOG2_E + return mean * log2_scale + float(tau) * torch.sqrt(var * log2_scale * log2_scale + 1.0e-6) + + def _predictor_oracle( q: torch.Tensor, k: torch.Tensor, @@ -123,28 +125,18 @@ def _predictor_oracle( *, tau: float, sm_scale: float, + thresh_type: str = "diag", ) -> 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 = _centroid_oracle(q) + threshold = _threshold_oracle( + q_centroids, k_summary, tau=tau, sm_scale=sm_scale, thresh_type=thresh_type ) - 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 + scores = torch.einsum("bqhd,bkhd->bhqk", q_centroids, k_summary.double()) * ( + float(sm_scale) * _LOG2_E ) - 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) + exact = scores > threshold.unsqueeze(-1) + blocks = k_summary.shape[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 @@ -165,12 +157,6 @@ def _inputs( ) -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], @@ -180,57 +166,57 @@ def _assert_outputs_match( 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_support_reason_on_cuda_tensors() -> None: + q, k, v = _inputs((1, 64, 2, 128), 5) + assert support_reason(q, k, v) is None + strided = q.transpose(1, 2).contiguous().transpose(1, 2) + assert strided.shape == q.shape and not strided.is_contiguous() + assert "contiguous" in support_reason(strided, k, v) + narrow = torch.zeros((1, 64, 2, 64), device="cuda", dtype=torch.bfloat16) + assert "head_dim=128" in support_reason(narrow, narrow, narrow) + empty = torch.zeros((1, 0, 2, 128), device="cuda", dtype=torch.bfloat16) + assert "positive" in support_reason(empty, empty, empty) + with pytest.raises(ValueError, match="thresh_type"): + predict(q, k, v, tau=0.5, sm_scale=0.125, thresh_type="diagonal") @_REQUIRES_CUDA -def test_sol_predictor_s257_bh_gt_one_matches_oracle_and_reuses_storage() -> None: +@pytest.mark.parametrize("thresh_type", _POLICIES) +def test_sol_predictor_matches_oracle(thresh_type: str) -> 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, + outputs = predict(q, k, v, tau=0.75, sm_scale=0.125, thresh_type=thresh_type) + assert outputs.exact_block_bits.shape == (2, 3, 5, 1) + assert outputs.k_summary.shape == outputs.v_summary.shape == (2, 5, 3, 128) + _assert_outputs_match( + outputs, + _predictor_oracle(q, k, v, tau=0.75, sm_scale=0.125, thresh_type=thresh_type), ) - 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() +def test_sol_predictor_exact_policy_sees_key_channel_correlation() -> None: + """Keys whose channels move together have a full-covariance variance the diagonal policy misses.""" + q, _, v = _inputs((1, 320, 2, 128), 23) + shared = _small_integer_bf16((1, 320, 2, 1), seed=29) + k = shared.expand(-1, -1, -1, 128).contiguous() + diag = predict(q, k, v, tau=0.5, sm_scale=0.125, thresh_type="diag") + exact = predict(q, k, v, tau=0.5, sm_scale=0.125, thresh_type="exact") + _assert_outputs_match(diag, _predictor_oracle(q, k, v, tau=0.5, sm_scale=0.125)) + _assert_outputs_match( + exact, _predictor_oracle(q, k, v, tau=0.5, sm_scale=0.125, thresh_type="exact") + ) + assert not torch.equal(diag.exact_block_bits, exact.exact_block_bits) - 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) +@_REQUIRES_CUDA +def test_sol_predictor_runtime_scale_and_tau_extremes() -> None: + q, k, v = _inputs((1, 257, 2, 128), 61) + normal = predict(q, k, v, tau=0.5, sm_scale=0.125).exact_block_bits + tiny = predict(q, k, v, tau=0.5, sm_scale=1.0e-5).exact_block_bits + assert torch.equal(normal, _predictor_oracle(q, k, v, tau=0.5, sm_scale=0.125)[0]) + assert torch.equal(tiny, _predictor_oracle(q, k, v, tau=0.5, sm_scale=1.0e-5)[0]) + assert not torch.equal(normal, tiny) blocks = 5 block_ids = torch.arange(blocks, device=q.device) @@ -240,7 +226,7 @@ def test_sol_predictor_s257_runtime_scale_and_tau_extremes() -> None: _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) + outputs = predict(q, k, v, tau=tau, sm_scale=128**-0.5) assert torch.equal(outputs.exact_block_bits, expected) @@ -249,7 +235,7 @@ def test_sol_predictor_long_proxy_group_keeps_tail_mass_and_clears_padding_bits( 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) + outputs = 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 @@ -262,17 +248,15 @@ def test_sol_predictor_long_proxy_group_keeps_tail_mass_and_clears_padding_bits( @_REQUIRES_CUDA -def test_sol_predictor_cuda_graph_replay_updates_live_outputs() -> None: +@pytest.mark.parametrize("thresh_type", _POLICIES) +def test_sol_predictor_cuda_graph_replay_refreshes_captured_outputs(thresh_type: str) -> 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) + predict(q, k, v, tau=0.5, sm_scale=0.125, thresh_type=thresh_type) 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 + captured = predict(q, k, v, tau=0.5, sm_scale=0.125, thresh_type=thresh_type) next_q, next_k, next_v = _inputs(q.shape, 51) q.copy_(next_q) @@ -280,95 +264,60 @@ def test_sol_predictor_cuda_graph_replay_updates_live_outputs() -> None: 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) + _assert_outputs_match( + captured, _predictor_oracle(q, k, v, tau=0.5, sm_scale=0.125, thresh_type=thresh_type) + ) @_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 - +def test_sol_predictor_operator_is_functional_and_compiles_fullgraph() -> None: 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 "!" not in schema, "the predictor operator must not mutate its inputs" 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 + meta_q = torch.empty((1, 65, 3, 128), device="meta", dtype=torch.bfloat16) + bits, k_summary, v_summary = op(meta_q, meta_q, meta_q, BLOCK_SIZE, 0.5, 0.125, "diag") + assert bits.shape == (1, 3, 2, 1) and bits.dtype == torch.uint32 + assert k_summary.shape == v_summary.shape == (1, 2, 3, 128) 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, + compiled = torch.compile(predict, backend="eager", fullgraph=True) + for thresh_type in _POLICIES: + outputs = compiled(q, k, v, tau=0.5, sm_scale=0.125, thresh_type=thresh_type) + _assert_outputs_match( + outputs, + _predictor_oracle(q, k, v, tau=0.5, sm_scale=0.125, thresh_type=thresh_type), + ) + + +@_REQUIRES_CUDA +@pytest.mark.parametrize("thresh_type", _POLICIES) +def test_sol_predictor_thresholds_match_the_fused_kernel_preprocess(thresh_type: str) -> None: + """Both backends must route from the same per-block threshold.""" + if torch.cuda.get_device_capability()[0] < 9: + pytest.skip("the fused kernel preprocess needs TMA descriptors") + preprocess = pytest.importorskip( + "tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.sol_attn.preprocess" + ) + generator = torch.Generator(device="cuda").manual_seed(97) + shape = (2, 320, 2, 128) + q, k, v = ( + torch.randn(shape, generator=generator, device="cuda").to(torch.bfloat16) for _ in range(3) + ) + _, _, fused_threshold = preprocess.prepare( + q, k, v, tau=0.75, scale=128**-0.5, thresh_type=thresh_type ) - 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) + summary_shape = (shape[0], shape[1] // BLOCK_SIZE, shape[2], shape[3]) + centroid = torch.empty(summary_shape, device="cuda", dtype=torch.float32) + k_summary = torch.empty(summary_shape, device="cuda", dtype=torch.bfloat16) + block_pool(q, centroid, block_size=BLOCK_SIZE, reduce="mean") + block_pool(k, k_summary, block_size=BLOCK_SIZE, reduce="mean") + threshold = block_thresholds( + centroid, k_summary, tau=0.75, sm_scale=128**-0.5, thresh_type=thresh_type ) - _assert_outputs_match(plan.outputs, expected) + + torch.testing.assert_close(threshold.permute(0, 2, 1), fused_threshold, rtol=2e-2, atol=2e-2) 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 index 79556dbcf31d..f92bdc5608a3 100644 --- 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 @@ -1,6 +1,6 @@ # 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).""" +"""Unit tests for the SOL predictor kernels (block pooling, block thresholds, exact-block selection).""" from __future__ import annotations @@ -8,11 +8,12 @@ import torch from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.kernels import ( + _LOG2_E, _block_pool_torch, - _block_statistics_torch, + _block_thresholds_torch, _select_exact_blocks_torch, block_pool, - block_statistics, + block_thresholds, select_exact_blocks, ) @@ -66,12 +67,11 @@ def test_select_exact_blocks_torch_fallback_packs_bit_r_of_word_w() -> None: 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)) + threshold = torch.full((1, 1, blocks), 5.0e-4, dtype=torch.float32) 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) + select_exact_blocks(centroid, k_summary, threshold, bits, 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. + # Row 0 scores log2(e) against block 33 only, above the 5e-4 threshold, plus its local band. expected = torch.zeros(blocks, dtype=torch.bool) expected[[0, 1, 33]] = True assert torch.equal(exact[0, 0, 0], expected) @@ -81,6 +81,58 @@ def test_select_exact_blocks_torch_fallback_packs_bit_r_of_word_w() -> None: assert torch.equal(exact[0, 0, 17].nonzero().flatten(), torch.tensor([16, 17, 18])) +def _threshold_oracle( + centroid: torch.Tensor, + k_summary: torch.Tensor, + *, + tau: float, + sm_scale: float, + thresh_type: str, +) -> torch.Tensor: + c = centroid.double().permute(0, 2, 1, 3) + keys = k_summary.double().permute(0, 2, 1, 3) + k_mean = keys.mean(dim=2) + mean = torch.einsum("bhqd,bhd->bhq", c, k_mean) + if thresh_type == "diag": + var = torch.einsum("bhqd,bhd->bhq", c.square(), keys.var(dim=2, unbiased=False)) + else: + centered = keys - k_mean.unsqueeze(2) + covariance = torch.matmul(centered.transpose(-1, -2), centered) / keys.shape[2] + var = torch.einsum("bhqd,bhqd->bhq", torch.matmul(c, covariance), c) + log2_scale = sm_scale * _LOG2_E + return mean * log2_scale + tau * torch.sqrt(var * log2_scale * log2_scale + 1.0e-6) + + +@_CPU_ONLY +@pytest.mark.parametrize("thresh_type", ["diag", "exact"]) +def test_block_thresholds_match_fp64_oracle(thresh_type: str) -> None: + torch.manual_seed(7) + centroid = torch.randn((2, 9, 3, 16), dtype=torch.float32) + k_summary = torch.randn((2, 9, 3, 16)).to(torch.bfloat16) + threshold = block_thresholds( + centroid, k_summary, tau=0.75, sm_scale=0.25, thresh_type=thresh_type + ) + assert threshold.shape == (2, 3, 9) and threshold.dtype == torch.float32 + assert threshold.is_contiguous() + expected = _threshold_oracle( + centroid, k_summary, tau=0.75, sm_scale=0.25, thresh_type=thresh_type + ) + torch.testing.assert_close(threshold.double(), expected, rtol=1e-4, atol=1e-5) + + +@_CPU_ONLY +def test_block_thresholds_exact_exceeds_diag_for_correlated_keys() -> None: + """A rank-one key distribution has a full-covariance variance the diagonal policy underestimates.""" + torch.manual_seed(8) + centroid = torch.ones((1, 4, 1, 16), dtype=torch.float32) + k_summary = torch.randn((1, 12, 1, 1)).expand(-1, -1, -1, 16).contiguous().to(torch.bfloat16) + diag = block_thresholds(centroid, k_summary, tau=1.0, sm_scale=1.0, thresh_type="diag") + exact = block_thresholds(centroid, k_summary, tau=1.0, sm_scale=1.0, thresh_type="exact") + assert torch.all(exact > diag) + with pytest.raises(ValueError, match="thresh_type"): + block_thresholds(centroid, k_summary, tau=1.0, sm_scale=1.0, thresh_type="full") + + @_REQUIRES_CUDA @pytest.mark.parametrize("seq_len", [64, 257, 4097]) @pytest.mark.parametrize("out_dtype", [torch.float32, torch.bfloat16]) @@ -113,18 +165,24 @@ def test_block_pool_accepts_strided_batch_and_token_dims() -> None: @_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()) +@pytest.mark.parametrize("thresh_type", ["diag", "exact"]) +@pytest.mark.parametrize("num_blocks", [5, 258]) +def test_block_thresholds_kernels_match_torch_fallback(thresh_type: str, num_blocks: int) -> None: + torch.manual_seed(5) + batch, heads, dim = 2, 3, 128 + centroid = torch.randn((batch, 7, heads, dim), device="cuda") * 0.125 + k_summary = (torch.randn((batch, num_blocks, heads, dim), device="cuda") * 0.125).to( + torch.bfloat16 + ) + for tau in (0.75, -0.5): + threshold = block_thresholds( + centroid, k_summary, tau=tau, sm_scale=0.125, thresh_type=thresh_type + ) + expected = _block_thresholds_torch( + centroid, k_summary, tau=tau, sm_scale=0.125, thresh_type=thresh_type + ) + assert threshold.shape == (batch, heads, 7) and threshold.is_contiguous() + torch.testing.assert_close(threshold, expected, rtol=1e-4, atol=1e-5) @_REQUIRES_CUDA @@ -136,18 +194,16 @@ def test_select_exact_blocks_matches_fallback_and_clears_padding_bits(num_blocks 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 + for tau, thresh_type in ((0.75, "diag"), (0.75, "exact"), (-1.0e6, "diag"), (1.0e6, "exact")): + threshold = block_thresholds( + centroid, k_summary, tau=tau, sm_scale=0.125, thresh_type=thresh_type ) - assert torch.equal(bits, expected), f"tau={tau}" + select_exact_blocks(centroid, k_summary, threshold, bits, sm_scale=0.125) + _select_exact_blocks_torch(centroid, k_summary, threshold, expected, sm_scale=0.125) + assert torch.equal(bits, expected), f"tau={tau} thresh_type={thresh_type}" padding = words * 32 - num_blocks if padding: assert int(bits[..., -1].to(torch.int64).max()) < (1 << (32 - padding)) @@ -164,15 +220,15 @@ def test_kernels_replay_inside_cuda_graph() -> None: 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) + threshold = block_thresholds( + centroid, k_summary, tau=0.5, sm_scale=0.125, thresh_type="exact" + ) + select_exact_blocks(centroid, k_summary, threshold, bits, sm_scale=0.125) run() torch.cuda.synchronize() @@ -188,18 +244,12 @@ def run() -> None: 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_threshold = block_thresholds( + expected_centroid, expected_summary, tau=0.5, sm_scale=0.125, thresh_type="exact" + ) 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, + expected_centroid, expected_summary, expected_threshold, expected_bits, 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) 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 4684da2950ed..006cb38fd73e 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py @@ -19,7 +19,6 @@ QuantAttentionConfig, RuntimeLoRAConfig, SkipSoftmaxAttentionConfig, - SolAttentionConfig, TeaCacheConfig, TorchCompileConfig, VAEConfig, @@ -514,54 +513,6 @@ 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 c63a310581255c206f1bbfe559ef59f974ae25de Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Sun, 20 Sep 2026 04:45:32 +0000 Subject: [PATCH 09/10] refactor: unify sparse algorithm naming in the VisualGen attention dispatch The attention module and the backend factory use the same names for the sparse dispatch: sparse_config for the user config, sparse_algorithm for its discriminator, is_vsa and is_sol for the two algorithms, and SOL as the display name in messages and comments, including the CuTeDSL SOL backend moved from the Sol-Attn module. The attention module also drops a duplicate context-parallelism check that the shared one already covered. The Wan pipeline registry no longer lists the FastVideo Wan2.1 VSA checkpoint: the checkpoint resolves through its pipeline class like before, and this change set is about the sparse attention backends, not about adding models. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- .../attention_backend/sparse/sol/backend.py | 12 +++---- .../visual_gen/attention_backend/utils.py | 6 ++-- .../visual_gen/models/wan/pipeline_wan.py | 1 - .../_torch/visual_gen/modules/attention.py | 34 ++++++++----------- 4 files changed, 24 insertions(+), 29 deletions(-) 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 index 6cbc44d5db59..56f4f4745d14 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/backend.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/backend.py @@ -75,9 +75,9 @@ def _cute_dense_available() -> bool: """Whether `cute_dsl_fmha_fwd` can run on the current device. - Checked once at construction. Sol-Attn and the dense CuTe DSL kernel now + Checked once at construction. SOL and the dense CuTe DSL kernel now cover the same set (sm_100a/sm_103a), so in practice this is always true - wherever Sol-Attn runs; the negative branch exists so an unsupported device + wherever SOL runs; the negative branch exists so an unsupported device degrades to SDPA instead of raising. """ try: @@ -128,7 +128,7 @@ def __init__( # would then see unequal Q/K shapes and quietly take its dense # fallback instead of rejecting an unsupported configuration. raise ValueError( - f"Sol-Attn is MHA-only (num_kv_heads == num_heads), got " + f"SOL is MHA-only (num_kv_heads == num_heads), got " f"num_kv_heads={self.num_kv_heads}, num_heads={self.num_heads}. " f"GQA/MQA is not supported." ) @@ -196,7 +196,7 @@ def _dense_by_step(self, timestep: Any) -> bool: logger.warning_once( "SolAttentionConfig.disabled_until_timestep=" f"{self.disabled_until_timestep} is set, but no `timestep` reached " - "the Sol-Attn forward call. The dense prefix it requests will not " + "the SOL forward call. The dense prefix it requests will not " "be applied. Ensure the pipeline passes a normalized timestep, or " "unset disabled_until_timestep.", key="sol_attn_missing_timestep", @@ -257,14 +257,14 @@ def _can_serve(self, q: torch.Tensor, k: torch.Tensor, **kwargs: Any) -> bool: Everything false here is delegated (see ``_delegate``). Deciding it from the tensors, per call, is deliberate, and it is why ``modules/attention.py`` - has no ``SEPARATE_QKV`` rule for Sol-Attn: ``qkv_mode`` describes how + has no ``SEPARATE_QKV`` rule for SOL: ``qkv_mode`` describes how Q/K/V are *projected*, not whether K/V come from another sequence, so a construction-time rule keyed on it mistakes self-attention for cross-attention wherever that mode is chosen for other reasons -- Qwen-Image always, and WAN's ``attn1`` under async Ulysses -- silently costing those modules their configured backend. """ - # Cross-attention: K/V come from another sequence, and Sol-Attn's + # Cross-attention: K/V come from another sequence, and SOL's # routing assumes one self-attending sequence. Unequal Q/K lengths are # a heuristic for that, not a definition: a cross-attention call whose # context happens to match the query length is not caught here. The diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py index 9d2cd4d7a425..536a1e6f30a4 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py @@ -131,10 +131,10 @@ def create_attention( Returns: AttentionBackend instance """ - sparse_attention_config = ( + sparse_config = ( attention_config.sparse_attention_config if attention_config is not None else None ) - sparse_algorithm = getattr(sparse_attention_config, "algorithm", None) + sparse_algorithm = getattr(sparse_config, "algorithm", None) is_vsa = sparse_algorithm == "vsa" is_sol = sparse_algorithm == "sol_attn" @@ -165,7 +165,7 @@ def create_attention( elif is_sol and kwargs.get("sparse_params") is None: # The attention module lowers the config once per layer; callers that # construct a backend directly get the same lowering here. - kwargs["sparse_params"] = sparse_attention_config.to_sparse_params() + kwargs["sparse_params"] = sparse_config.to_sparse_params() # Forward the validated quantization recipe to TRTLLM, cuDNN, FlashInfer, or the dense CuTe DSL # FMHA backend. 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 358724667ef6..5f78ed1d5be0 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py @@ -115,7 +115,6 @@ 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", diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/modules/attention.py index cd767504b735..46577db32d8c 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py @@ -111,10 +111,10 @@ def __init__( ulysses_size = vgm.ulysses_size if vgm else 1 cp_size = vgm.cp_size if vgm else 1 base_backend = config.attention.backend - _sa_cfg = config.attention.sparse_attention_config - _sa_algo = getattr(_sa_cfg, "algorithm", None) if _sa_cfg is not None else None - is_vsa = _sa_algo == "vsa" - is_sol = _sa_algo == "sol_attn" + sparse_config = config.attention.sparse_attention_config + sparse_algorithm = getattr(sparse_config, "algorithm", None) + is_vsa = sparse_algorithm == "vsa" + is_sol = sparse_algorithm == "sol_attn" is_separate_qkv_cross_attention = ( self.qkv_mode == QKVMode.SEPARATE_QKV and not separate_qkv_is_self_attention ) @@ -144,15 +144,9 @@ def __init__( # Every sparse algorithm here routes over the whole token sequence, so # none of them can be split across context-parallel ranks. if (is_vsa or is_sol) and cp_size > 1: - _algo_name = "VSA" if is_vsa else "SOL" + algorithm_name = "VSA" if is_vsa else "SOL" raise ValueError( - f"{_algo_name} 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." - ) - if is_sol and cp_size > 1: - raise ValueError( - f"SOL needs the full token sequence per rank, so it is incompatible " + f"{algorithm_name} 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." ) @@ -253,20 +247,22 @@ def __init__( backend_num_heads = self.local_num_attention_heads backend_num_kv_heads = self.local_num_key_value_heads - # Lower the shared SkipSoftmax user/checkpoint config for each backend - # whose kernel consumes SkipSoftmaxParams. + # Lower the user/checkpoint sparse config for each backend whose kernel + # consumes the lowered SparseParams. sparse_params = None - ss_cfg = config.attention.sparse_attention_config - if isinstance(ss_cfg, SkipSoftmaxAttentionConfig) and backend_name in ( + if isinstance(sparse_config, SkipSoftmaxAttentionConfig) and backend_name in ( "TRTLLM", "CUTEDSL", ): - sparse_params = ss_cfg.to_sparse_params( + sparse_params = sparse_config.to_sparse_params( module_name=self.module_name, pretrained_config=config.pretrained_config, ) - elif isinstance(ss_cfg, SolAttentionConfig) and backend_name in ("TRTLLM", "CUTEDSL"): - sparse_params = ss_cfg.to_sparse_params() + elif isinstance(sparse_config, SolAttentionConfig) and backend_name in ( + "TRTLLM", + "CUTEDSL", + ): + sparse_params = sparse_config.to_sparse_params() self.sparse_params = sparse_params # Create compute backend From 85b6b4d72f7ce6201045de477a8bf7ad82517c93 Mon Sep 17 00:00:00 2001 From: yuhangh <58161490+heyuhhh@users.noreply.github.com> Date: Sun, 20 Sep 2026 07:36:54 +0000 Subject: [PATCH 10/10] fix: resolve review feedback on VisualGen VSA and SOL sparse attention Drop the 16-shape cap on the VSA metadata and route caches. The CUDA Graph runner that captures the cached tensor addresses keeps one graph per shape without a cap, so the limit only turned a new resolution/frame profile into a runtime error. Both caches now keep one entry per distinct shape for as long as the graphs that reference them and are released together with the graphs in cleanup. Make the SOL block pooling kernel mask the block boundary so block sizes that are not a multiple of the load width match the PyTorch fallback, and reject non-positive block sizes. Add the VSA and SOL predictor kernel suites to the CPU and B200 test lists. Move the TRTLLM skip-softmax CUDA Graph capture test out of the cpu_only module into the TRTLLM metadata suite with an SM100 gate so it runs on the GPU stage. Assert that TRTLLM VSA executes the block-sparse FMHA rather than only predicting routes, check the forwarded VSA gate values, verify that the Wan VSA pipeline reference takes the SDPA fallback, and reject non-finite tensor timesteps. Annotate the helpers and tests flagged during review and read the sparse algorithm discriminator directly instead of through getattr. Signed-off-by: yuhangh <58161490+heyuhhh@users.noreply.github.com> --- .../features/visualgen-sparse-attention.md | 2 +- .../attention_backend/sparse/sol/kernels.py | 9 +++- .../attention_backend/sparse/vsa/metadata.py | 19 +++---- .../attention_backend/sparse/vsa/predictor.py | 27 ++-------- .../visual_gen/attention_backend/trtllm.py | 2 +- .../visual_gen/attention_backend/utils.py | 2 +- .../_torch/visual_gen/modules/attention.py | 8 +-- tensorrt_llm/visual_gen/sparse_attention.py | 7 ++- .../test_lists/test-db/l0_b200.yml | 2 + .../integration/test_lists/test-db/l0_cpu.yml | 1 + .../attention/sparse/test_timestep_phase.py | 24 +++++++-- .../multi_gpu/test_ulysses_attention.py | 8 +-- .../multi_gpu/test_wan_async_ulysses.py | 16 ++++-- .../visual_gen/multi_gpu/test_wan_tp.py | 4 +- .../sparse_attention/test_skip_softmax.py | 50 +----------------- .../sparse_attention/test_sol_attention.py | 3 +- .../test_sol_predictor_kernels.py | 14 +++-- .../visual_gen/test_attention_integration.py | 24 +++++---- .../_torch/visual_gen/test_attention_vsa.py | 44 +++++++++------- .../test_trtllm_attention_metadata.py | 51 +++++++++++++++++++ .../visual_gen/test_wan_vsa_pipeline.py | 40 ++++++++++++++- 21 files changed, 217 insertions(+), 140 deletions(-) diff --git a/docs/source/features/visualgen-sparse-attention.md b/docs/source/features/visualgen-sparse-attention.md index 3f5ec48906a1..d76f4dbf7092 100644 --- a/docs/source/features/visualgen-sparse-attention.md +++ b/docs/source/features/visualgen-sparse-attention.md @@ -313,7 +313,7 @@ be reused for the sparse phase. 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. +VSA retains shape-dependent metadata and route tensors so CUDA Graph replay can reuse stable addresses. A pipeline instance keeps one set of these tensors per distinct shape profile for as long as the CUDA Graphs that reference them, so their footprint grows with the number of served resolution/frame profiles exactly like the graphs do. Both VSA backends use the same VisualGen-owned predictor implementation, one instance per attention layer, and identical post-processing. The `TRTLLM` path runs the coarse stage before the core diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/kernels.py b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/kernels.py index b241f79aca33..5bf870459233 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/kernels.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/kernels.py @@ -76,10 +76,13 @@ def _block_pool_kernel( 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) + offsets = start + tl.arange(0, TOKENS) + tokens = first_token + offsets + # The last load of a block whose size is not a multiple of TOKENS stops at the + # block boundary instead of running into the next block. values = tl.load( x_ptr + batch * stride_x_batch + tokens[:, None] * stride_x_token + columns[None, :], - mask=(tokens < seq_len)[:, None] & in_row[None, :], + mask=((offsets < BLOCK) & (tokens < seq_len))[:, None] & in_row[None, :], other=0.0, ) total += tl.sum(values.to(tl.float32), axis=0) @@ -122,6 +125,8 @@ def block_pool(x: torch.Tensor, out: torch.Tensor, *, block_size: int, reduce: _ """ if reduce not in ("mean", "sum"): raise ValueError(f"reduce must be 'mean' or 'sum'; got {reduce!r}") + if block_size <= 0: + raise ValueError(f"block_size must be positive; got {block_size}") 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" 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 31916e397436..83abf9cff8f1 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 @@ -26,7 +26,6 @@ # 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 _BITS_PER_WORD = 32 @@ -118,12 +117,14 @@ class _VSAShapeMetadata(TypedDict): class VSAMetadataBuilder: - """Build VSA metadata while caching shape-dependent index tensors.""" + """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 + CUDA Graphs capture the addresses of the cached tensors, so the cache keeps one + entry per distinct shape for as long as the graphs that reference it, exactly like + the graph set itself; ``clear`` releases both together. + """ + + def __init__(self) -> None: self._cache: dict[Tuple[Tuple[int, int, int], torch.device], _VSAShapeMetadata] = {} def _build_metadata( @@ -189,12 +190,6 @@ def build( 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 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 97db41c5cf9b..6fa244039981 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 @@ -24,12 +24,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, - VSAMetadata, - get_vsa_forward_context, -) +from .metadata import VSA_BLOCK_SIZE, VSAMetadata, get_vsa_forward_context _SIGNED_INT32_MAX = torch.iinfo(torch.int32).max @@ -75,10 +70,7 @@ class VSAForwardInputs: 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 + def __init__(self) -> None: self._indptr_cache: dict[tuple[torch.device, int, int, int, int], torch.Tensor] = {} def from_selected_blocks( @@ -117,12 +109,6 @@ def _build_block_indptr( 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; " @@ -141,12 +127,7 @@ def _build_block_indptr( 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: + def __init__(self, num_heads: int, num_kv_heads: Optional[int] = None) -> None: resolved_num_kv_heads = num_kv_heads or num_heads if resolved_num_kv_heads != num_heads: raise ValueError( @@ -154,7 +135,7 @@ def __init__( 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) + self._route_builder = _VSARouteBuilder() @torch.compiler.disable def get_metadata(self) -> VSAMetadata: diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py b/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py index 9cacb5742b00..af67222ed536 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py @@ -327,7 +327,7 @@ def _compact_qkv( batch_size: int, seq_len: int, kv_seq_len: int, - ): + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: # 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. diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py index 536a1e6f30a4..b9e69cc0d74d 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py @@ -134,7 +134,7 @@ def create_attention( sparse_config = ( attention_config.sparse_attention_config if attention_config is not None else None ) - sparse_algorithm = getattr(sparse_config, "algorithm", None) + sparse_algorithm = sparse_config.algorithm if sparse_config is not None else None is_vsa = sparse_algorithm == "vsa" is_sol = sparse_algorithm == "sol_attn" diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/modules/attention.py index 46577db32d8c..91dc235a4dde 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py @@ -14,7 +14,7 @@ # limitations under the License. from enum import Enum -from typing import Optional, Tuple +from typing import Any, Optional, Tuple import torch import torch.nn as nn @@ -554,7 +554,7 @@ def _attn_impl( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - **kwargs, + **kwargs: Any, ) -> torch.Tensor: """ Call attention backend with appropriate tensor layout. @@ -621,7 +621,7 @@ def forward( encoder_hidden_states: Optional[torch.Tensor] = None, freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, timestep: Optional[torch.Tensor] = None, - **kwargs, + **kwargs: Any, ) -> torch.Tensor: # hidden_states may be [B, S, H] or an Fp4QuantizedTensor from an upstream # fused norm+quant kernel; downstream Linear accepts either. @@ -671,7 +671,7 @@ def forward_async( hidden_states: torch.Tensor, freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, timestep: Optional[torch.Tensor] = None, - **kwargs, + **kwargs: Any, ) -> torch.Tensor: """Async-Ulysses self-attn driver. Structurally mirrors ``forward``: each closure does ``to_{q,k,v}`` + (optional) fused norm+RoPE on the diff --git a/tensorrt_llm/visual_gen/sparse_attention.py b/tensorrt_llm/visual_gen/sparse_attention.py index 82f5864fbf86..24e26316b96c 100644 --- a/tensorrt_llm/visual_gen/sparse_attention.py +++ b/tensorrt_llm/visual_gen/sparse_attention.py @@ -16,13 +16,16 @@ import fnmatch from types import SimpleNamespace -from typing import Any, Dict, Literal, Optional +from typing import TYPE_CHECKING, Any, Dict, Literal, Optional from pydantic import Field as PydanticField from pydantic import field_validator from tensorrt_llm.llmapi.utils import StrictBaseModel +if TYPE_CHECKING: + from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.params import SolParams + class BaseSparseAttentionConfig(StrictBaseModel): """Base for visual-generation sparse attention configs. @@ -308,7 +311,7 @@ def _validate_dense_layers(cls, layers: Optional[list[int]]) -> Optional[list[in raise ValueError(f"dense_layers contains a negative layer index: {index}") return sorted(set(layers)) - def to_sparse_params(self, **kwargs): + def to_sparse_params(self, **kwargs: Any) -> "SolParams": """Lower the public recipe into the SOL parameters shared by both backends.""" del kwargs from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.params import SolParams diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 268383a3684d..a0d641c4dbda 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -301,9 +301,11 @@ l0_b200: - unittest/_torch/visual_gen/test_attention_cute_dsl.py - unittest/_torch/visual_gen/test_fa4_cutlass_compatibility.py - unittest/_torch/visual_gen/test_attention_vsa.py + - unittest/_torch/visual_gen/test_attention_vsa_kernels.py - unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.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_predictor_kernels.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 diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 56b1833c6cbb..29fa03f7f543 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -79,6 +79,7 @@ l0_cpu: - 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/sparse_attention/test_sol_predictor_kernels.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/attention/sparse/test_timestep_phase.py b/tests/unittest/_torch/attention/sparse/test_timestep_phase.py index 507268cd22d1..95241f5e863b 100644 --- a/tests/unittest/_torch/attention/sparse/test_timestep_phase.py +++ b/tests/unittest/_torch/attention/sparse/test_timestep_phase.py @@ -23,17 +23,31 @@ (torch.tensor([0.0, 0.8]), 0.8), ], ) -def test_timestep_to_float_reduces_to_the_largest_live_value(timestep, expected): +def test_timestep_to_float_reduces_to_the_largest_live_value(timestep, expected) -> None: assert timestep_to_float(timestep) == ( expected if expected is None else pytest.approx(expected) ) -def test_timestep_to_float_rejects_non_real_and_non_finite_values(): +def test_timestep_to_float_rejects_non_real_values() -> None: with pytest.raises(TypeError, match="real scalar or tensor"): timestep_to_float(True) + + +@pytest.mark.parametrize( + "timestep", + [ + float("nan"), + float("inf"), + torch.tensor([float("nan")]), + # The reduction keeps the non-finite value. + torch.tensor([0.5, float("inf")]), + ], + ids=["nan", "inf", "tensor_nan", "tensor_inf"], +) +def test_timestep_to_float_rejects_non_finite_values(timestep) -> None: with pytest.raises(ValueError, match="finite"): - timestep_to_float(float("nan")) + timestep_to_float(timestep) @pytest.mark.parametrize( @@ -48,5 +62,7 @@ def test_timestep_to_float_rejects_non_real_and_non_finite_values(): (torch.tensor([0.0, 0.2]), 0.6, 1), ], ) -def test_graph_phase_for_timestep_marks_dense_prefix_and_sparse_suffix(timestep, cutoff, expected): +def test_graph_phase_for_timestep_marks_dense_prefix_and_sparse_suffix( + timestep, cutoff, expected +) -> None: assert graph_phase_for_timestep(timestep, disabled_until_timestep=cutoff) == expected 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 b36d780418f5..69ae41adfba3 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 @@ -111,13 +111,15 @@ def run_test_in_distributed(world_size: int, test_fn: Callable, use_cuda: bool = ) -def test_forward_async_redistributes_vsa_gates(monkeypatch): +def test_forward_async_redistributes_vsa_gates(monkeypatch: pytest.MonkeyPatch) -> None: 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): + def forward( + self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, **kwargs: object + ) -> torch.Tensor: self.kwargs = kwargs return q @@ -131,7 +133,7 @@ def forward(self, q, k, v, **kwargs): attention._output_a2a = lambda output, batch_size, seq_len: output redistributed = [] - def _fake_all_to_all(tensor, **kwargs): + def _fake_all_to_all(tensor: torch.Tensor, **kwargs: object) -> torch.Tensor: redistributed.append((tensor, kwargs)) return tensor + 1 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 aca1c29f961f..1b25a1a93d63 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 @@ -317,11 +317,17 @@ 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(): +def test_forward_async_uses_tp_local_heads_for_qkv_gates_and_output() -> None: 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): + def forward_async( + self, + compute_q: Callable[[], torch.Tensor], + compute_k: Callable[[], torch.Tensor], + compute_v: Callable[[], torch.Tensor], + **kwargs: object, + ) -> torch.Tensor: self.q = compute_q() self.k = compute_k() self.v = compute_v() @@ -329,7 +335,7 @@ def forward_async(self, compute_q, compute_k, compute_v, **kwargs): return self.q class _CaptureOutputProjection(torch.nn.Module): - def forward(self, hidden_states): + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: self.input = hidden_states return hidden_states @@ -364,6 +370,10 @@ def forward(self, hidden_states): 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 + torch.testing.assert_close( + attention.attn.kwargs["gate_compress"], gate_compress.view(expected_shape) + ) + torch.testing.assert_close(attention.attn.kwargs["gate_fine"], gate_fine.view(expected_shape)) assert output_projection.input.shape == (1, 3, 8) assert output.shape == (1, 3, 8) 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 0860209d1d3d..b24ff5186dad 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 @@ -525,7 +525,9 @@ def _logic_wan_i2v_tp_vs_single_gpu_with_config(rank, world_size, config_dict): ("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): +def test_wan_vsa_gates_follow_ulysses_aligned_tp_q_shard( + monkeypatch: pytest.MonkeyPatch, tp_rank: int, expected_head_range: tuple[int, int] +) -> None: """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 diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py b/tests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py index 35edb7da0671..ddd1cb0ad219 100644 --- a/tests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py +++ b/tests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py @@ -13,10 +13,7 @@ import yaml from pydantic import ValidationError -from tensorrt_llm._torch.attention.backends.sparse.skip_softmax import ( - SkipSoftmaxParams, - SkipSoftmaxScheduler, -) +from tensorrt_llm._torch.attention.backends.sparse.skip_softmax import SkipSoftmaxParams from tensorrt_llm._torch.attention.backends.sparse.timestep_phase import graph_phase_for_timestep from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl import fmha as cute_dsl_fmha from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.fmha import ( @@ -629,48 +626,3 @@ def test_pipeline_config_keeps_checkpoint_metadata_per_model(self, tmp_path): _expected_threshold(-20.0, 4.0, 0.5) ) assert transformer_disabled_params is None - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA Graph capture requires CUDA") -def test_trtllm_skip_softmax_cutoff_is_capturable_with_a_device_timestep(): - """The wrapper prepares the timestep during warmup so capture never reads the tensor.""" - from tensorrt_llm._torch.visual_gen.attention_backend.trtllm import TrtllmAttention - from tensorrt_llm._torch.visual_gen.config import create_attention_metadata_state - - params = SkipSoftmaxParams( - scheduler=SkipSoftmaxScheduler( - threshold_scale_factor_prefill=5000.0, disabled_until_timestep=0.6 - ) - ) - attention = TrtllmAttention( - layer_idx=0, - num_heads=2, - head_dim=128, - dtype=torch.bfloat16, - attention_metadata_state=create_attention_metadata_state(), - sparse_params=params, - ) - batch_size, seq_len = 1, 256 - q = torch.randn(batch_size, seq_len, 2, 128, device="cuda", dtype=torch.bfloat16) - k = torch.randn_like(q) - v = torch.randn_like(q) - timestep = torch.tensor([0.8], device="cuda") - - def forward(): - return attention.forward(q, k, v, batch_size=batch_size, seq_len=seq_len, timestep=timestep) - - side = torch.cuda.Stream() - side.wait_stream(torch.cuda.current_stream()) - with torch.cuda.stream(side): - for _ in range(2): - eager = forward() - torch.cuda.current_stream().wait_stream(side) - torch.cuda.synchronize() - - graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(graph): - captured = forward() - graph.replay() - torch.cuda.synchronize() - - torch.testing.assert_close(captured, eager) 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 index 5dc8679c05f0..b30d5f66a8e0 100644 --- a/tests/unittest/_torch/visual_gen/sparse_attention/test_sol_attention.py +++ b/tests/unittest/_torch/visual_gen/sparse_attention/test_sol_attention.py @@ -31,6 +31,7 @@ PredefinedAttentionMask, ) from tensorrt_llm._torch.attention.backends.sparse.hooks import prepare_sparse_runtime_params +from tensorrt_llm._torch.attention.backends.sparse.params import BlockSparseForwardInputs 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 import predictor as sol_predictor @@ -99,7 +100,7 @@ def _predict( *, attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL, timestep: object = None, -): +) -> BlockSparseForwardInputs | None: """Invoke the core prediction hook the way the core forward does.""" return backend.block_sparse_attn_predict( 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 index f92bdc5608a3..d09b490083f7 100644 --- 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 @@ -58,6 +58,8 @@ def test_block_pool_rejects_mismatched_output() -> None: 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") + with pytest.raises(ValueError, match="block_size"): + block_pool(x, torch.empty(_pooled_shape(x, BLOCK)), block_size=0, reduce="mean") @_CPU_ONLY @@ -136,14 +138,18 @@ def test_block_thresholds_exact_exceeds_diag_for_correlated_keys() -> None: @_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: +# Block sizes below, not a multiple of, and equal to the tokens loaded per step. +@pytest.mark.parametrize("block_size", [4, 9, BLOCK]) +def test_block_pool_matches_torch_fallback( + seq_len: int, out_dtype: torch.dtype, block_size: int +) -> 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) + out = torch.empty(_pooled_shape(x, block_size), dtype=out_dtype, device="cuda") + block_pool(x, out, block_size=block_size, reduce=reduce) expected = torch.empty_like(out) - _block_pool_torch(x, expected, block_size=BLOCK, reduce=reduce) + _block_pool_torch(x, expected, block_size=block_size, reduce=reduce) tolerance = ( {"rtol": 1e-5, "atol": 1e-5} if out_dtype == torch.float32 diff --git a/tests/unittest/_torch/visual_gen/test_attention_integration.py b/tests/unittest/_torch/visual_gen/test_attention_integration.py index c26efea39fdb..224dfd94130a 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_integration.py +++ b/tests/unittest/_torch/visual_gen/test_attention_integration.py @@ -728,7 +728,7 @@ def test_fast_cross_attention_wan_shapes( # ============================================================================ -def _build_vsa_setup(backend: str, sparsity: float, batch_size: int, seed: int): +def _build_vsa_setup(backend: str, sparsity: float, batch_size: int, seed: int) -> SimpleNamespace: """Build naive + integrated models, VSA metadata, and inputs for a VSA test. A ragged latent exercises VSA padding and token-mask lowering on both @@ -786,7 +786,7 @@ def _build_vsa_setup(backend: str, sparsity: float, batch_size: int, seed: int): @pytest.mark.skipif(not torch.cuda.is_available(), reason="VSA needs CUDA") @pytest.mark.parametrize("backend", ["CUTEDSL", "TRTLLM"]) -def test_vsa_self_attention_equivalence_at_sparsity_zero(backend: str): +def test_vsa_self_attention_equivalence_at_sparsity_zero(backend: str) -> None: """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.sparse.vsa import set_vsa_forward_context @@ -816,8 +816,11 @@ def test_vsa_self_attention_equivalence_at_sparsity_zero(backend: str): not isSM100Family(), reason="CuTe DSL and PrimTS block-sparse parity requires SM100 or SM103", ) -def test_vsa_sparse_backends_match_on_ragged_input(): +def test_vsa_sparse_backends_match_on_ragged_input(monkeypatch: pytest.MonkeyPatch) -> None: """CuTeDSL and TRTLLM implement the same sparse VSA fine-stage semantics.""" + from tensorrt_llm._torch.attention.backends.fmha.prims_ts_block_sparse import ( + PrimsTSBlockSparseFmha, + ) from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import set_vsa_forward_context sparsity = 0.5 @@ -837,14 +840,15 @@ def checked_execute(*args, _original=original_execute, **kwargs): setup.integrated.attn._execute_sparse_fine = checked_execute else: - original_predict = setup.integrated.attn.block_sparse_attn_predict + # Only the block-sparse FMHA library consumes predicted routes, so its + # forward running proves TRTLLM executed the sparse fine stage. + original_forward = PrimsTSBlockSparseFmha.forward - def checked_predict(*args, _original=original_predict, **kwargs): - result = _original(*args, **kwargs) - sparse_fine_executed["TRTLLM"] = result is not None - return result + def checked_forward(*args, _original=original_forward, **kwargs): + sparse_fine_executed["TRTLLM"] = True + return _original(*args, **kwargs) - setup.integrated.attn.block_sparse_attn_predict = checked_predict + monkeypatch.setattr(PrimsTSBlockSparseFmha, "forward", checked_forward) outputs = {} for backend, setup in setups.items(): @@ -871,7 +875,7 @@ def checked_predict(*args, _original=original_predict, **kwargs): not isSM100Family(), reason="PrimTS block-sparse CUDA Graph replay requires SM100 or SM103", ) -def test_vsa_trtllm_cuda_graph_replays_live_routes(): +def test_vsa_trtllm_cuda_graph_replays_live_routes() -> None: """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 diff --git a/tests/unittest/_torch/visual_gen/test_attention_vsa.py b/tests/unittest/_torch/visual_gen/test_attention_vsa.py index 0c5d9f174110..b26a7dae30b8 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_vsa.py +++ b/tests/unittest/_torch/visual_gen/test_attention_vsa.py @@ -407,7 +407,7 @@ def _make_config( num_heads: int, head_dim: int, backend: str, - vsa_sparsity: "float | None" = None, + vsa_sparsity: float | None = None, ) -> DiffusionModelConfig: """Minimal DiffusionModelConfig for one Attention module.""" pretrained_config = SimpleNamespace( @@ -497,7 +497,7 @@ def test_plain_trtllm_separate_qkv_dispatches_by_attention_role( assert attention.attn_backend == expected_backend -def test_vsa_with_attn2d_raises(): +def test_vsa_with_attn2d_raises() -> None: """VSA + Attention2D must error at construction (VSA needs the full sequence per rank).""" pretrained_config = SimpleNamespace( hidden_size=64, @@ -562,25 +562,35 @@ def test_vsa_metadata_exposes_tile_source_index_and_packed_kv_words() -> None: assert metadata.kv_valid_words.tolist() == [0xFFFFFFFF, 0xFFFFFFFF, 0xFFFF, 0] -def test_vsa_graph_stable_caches_bound_shape_profiles() -> None: - builder = VSAMetadataBuilder(max_cached_shapes=1) +def test_vsa_graph_stable_caches_keep_every_shape_profile() -> None: + """Every distinct shape stays cached, matching the unbounded CUDA Graph set.""" + builder = VSAMetadataBuilder() 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 + shapes = [(4, 4, 4 * (index + 1)) for index in range(20)] + first = [builder.build(raw_latent_shape=shape, **build_args) for shape in shapes] + second = [builder.build(raw_latent_shape=shape, **build_args) for shape in shapes] + assert len(builder._cache) == len(shapes) + for before, after in zip(first, second): + assert after.tile_source_index is before.tile_source_index + + route_builder = VSAPredictor(num_heads=1)._route_builder 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"): + routes = [ route_builder.from_selected_blocks( - torch.zeros((1, 1, 2, 1), dtype=torch.int32), kv_valid_words + torch.zeros((1, 1, num_q_blocks, 1), dtype=torch.int32), kv_valid_words ) + for num_q_blocks in range(1, 21) + ] + assert len(route_builder._indptr_cache) == 20 + again = route_builder.from_selected_blocks( + torch.zeros((1, 1, 7, 1), dtype=torch.int32), kv_valid_words + ) + assert again.block_indptr is routes[6].block_indptr @pytest.mark.skipif(not torch.cuda.is_available(), reason="VSA needs CUDA") @@ -593,7 +603,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): +def test_vsa_tile_untile_roundtrip(latent_shape: tuple[int, int, int]) -> None: """Tiling then untiling must reproduce the input, and pooled cubes must be token means.""" device = torch.device("cuda") dtype = torch.bfloat16 @@ -741,7 +751,7 @@ def run() -> tuple[torch.Tensor, torch.Tensor]: @pytest.mark.skipif(not torch.cuda.is_available(), reason="kernel test needs CUDA") -def test_cute_kernel_matches_dense_at_full_topk(): +def test_cute_kernel_matches_dense_at_full_topk() -> None: """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, @@ -793,7 +803,7 @@ def test_cute_kernel_matches_dense_at_full_topk(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="kernel test needs CUDA") -def test_cute_kernel_matches_ref_with_independent_indices(): +def test_cute_kernel_matches_ref_with_independent_indices() -> None: """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, @@ -876,7 +886,7 @@ def test_cute_kernel_matches_ref_with_independent_indices(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="kernel test needs CUDA") -def test_cute_kernel_50pct_sparsity_quality_vs_dense(): +def test_cute_kernel_50pct_sparsity_quality_vs_dense() -> None: """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, @@ -950,7 +960,7 @@ def test_cute_kernel_50pct_sparsity_quality_vs_dense(): [1, 3, 9], ids=["1cube_odd", "3cubes_odd", "9cubes_odd"], ) -def test_cute_kernel_odd_num_cubes_correctness(num_cubes): +def test_cute_kernel_odd_num_cubes_correctness(num_cubes: int) -> None: """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, 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 4018fd5fad9c..08bfc1a01cdb 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py @@ -12,9 +12,18 @@ SparseBackendForwardArgs, SparseRuntimeParams, ) +from tensorrt_llm._torch.attention.backends.sparse.skip_softmax import ( + SkipSoftmaxParams, + SkipSoftmaxScheduler, +) 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 +_REQUIRES_SM100 = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() not in ((10, 0), (10, 3)), + reason="the VisualGen TRTLLM FMHA path is exercised on SM100 or SM103", +) + class _FakeBaseTrtllmAttentionMetadata: def __init__(self, **kwargs): @@ -501,3 +510,45 @@ def test_forward_reaches_core_fmha_with_module_predicted_routes( assert isinstance(runtime_params, SparseRuntimeParams) assert runtime_params.block_sparse_inputs is carrier assert runtime_params.sparse_attn_indices_block_size == 0 + + +@_REQUIRES_SM100 +def test_trtllm_skip_softmax_cutoff_is_capturable_with_a_device_timestep() -> None: + """The wrapper prepares the timestep during warmup so capture never reads the tensor.""" + params = SkipSoftmaxParams( + scheduler=SkipSoftmaxScheduler( + threshold_scale_factor_prefill=5000.0, disabled_until_timestep=0.6 + ) + ) + attention = visual_trtllm.TrtllmAttention( + layer_idx=0, + num_heads=2, + head_dim=128, + dtype=torch.bfloat16, + attention_metadata_state=create_attention_metadata_state(), + sparse_params=params, + ) + batch_size, seq_len = 1, 256 + q = torch.randn(batch_size, seq_len, 2, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn_like(q) + timestep = torch.tensor([0.8], device="cuda") + + def forward() -> torch.Tensor: + return attention.forward(q, k, v, batch_size=batch_size, seq_len=seq_len, timestep=timestep) + + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + for _ in range(2): + eager = forward() + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = forward() + graph.replay() + torch.cuda.synchronize() + + torch.testing.assert_close(captured, eager) 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 e1232a9f2dbb..330e5a355305 100644 --- a/tests/unittest/_torch/visual_gen/test_wan_vsa_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_wan_vsa_pipeline.py @@ -13,8 +13,11 @@ pytest tests/unittest/_torch/visual_gen/test_wan_vsa_pipeline.py -v -s """ +import contextlib import gc import os +from types import SimpleNamespace +from typing import Iterator os.environ["TLLM_DISABLE_MPI"] = "1" @@ -108,6 +111,30 @@ def _cosine_similarity(a: torch.Tensor, b: torch.Tensor) -> float: return F.cosine_similarity(a_flat.unsqueeze(0), b_flat.unsqueeze(0)).clamp(-1.0, 1.0).item() +@contextlib.contextmanager +def _count_cute_fine_stage_launches() -> Iterator[SimpleNamespace]: + """Count CuTe VSA fine-stage launches without retaining their arguments. + + A recording mock would keep every layer's Q/K/V alive for the whole + denoising loop, so a plain counting wrapper is used instead. + """ + from unittest.mock import patch + + from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa.backend import ( + VSACuTeDSLAttention, + ) + + launches = SimpleNamespace(count=0) + original = VSACuTeDSLAttention._execute_sparse_fine + + def counted(self, inputs): + launches.count += 1 + return original(self, inputs) + + with patch.object(VSACuTeDSLAttention, "_execute_sparse_fine", counted): + yield launches + + def _assert_vsa_matches_dense( checkpoint_subdir: str, height: int, @@ -135,15 +162,24 @@ def _assert_vsa_matches_dense( # --- CuTe-DSL path --- vsa_pipe = _load_vsa_pipeline(checkpoint_subdir, vsa_sparsity=vsa_sparsity) - vsa_video = _capture_trtllm_video(vsa_pipe, **common_kwargs) + with _count_cute_fine_stage_launches() as cute_fine_stage: + vsa_video = _capture_trtllm_video(vsa_pipe, **common_kwargs) + assert cute_fine_stage.count > 0, f"{model_label}: the CuTe-DSL fine stage never ran" del vsa_pipe gc.collect() torch.cuda.empty_cache() # --- SDPA fallback reference (same VSA formulation, fine attn via SDPA) --- sdpa_pipe = _load_vsa_pipeline(checkpoint_subdir, vsa_sparsity=vsa_sparsity) - with patch.object(_vsa_module, "is_cute_supported", return_value=False): + with ( + patch.object(_vsa_module, "is_cute_supported", return_value=False), + _count_cute_fine_stage_launches() as sdpa_fine_stage, + ): sdpa_video = _capture_trtllm_video(sdpa_pipe, **common_kwargs) + # The fine stage has exactly two implementations, so no CuTe launch means SDPA ran. + assert sdpa_fine_stage.count == 0, ( + f"{model_label}: the SDPA reference still ran the CuTe-DSL fine stage" + ) del sdpa_pipe gc.collect() torch.cuda.empty_cache()