diff --git a/docs/source/features/visualgen-sparse-attention.md b/docs/source/features/visualgen-sparse-attention.md index cb734cb9c888..01940b73d59f 100644 --- a/docs/source/features/visualgen-sparse-attention.md +++ b/docs/source/features/visualgen-sparse-attention.md @@ -6,56 +6,23 @@ 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 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. 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 | `algorithm` | Config class | Status | |---|---|---| | `skip_softmax` | `SkipSoftmaxAttentionConfig` | Supported | -| `vsa` | `VideoSparseAttentionConfig` | Supported (CUTEDSL) | -| `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. +| `vsa` | `VideoSparseAttentionConfig` | Supported (`CUTEDSL`, `TRTLLM`) | +| `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,145 @@ 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 # block threshold policy: "diag" or "exact" +``` + +- `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 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 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 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 +dense until every token is below the cutoff. A dense capture therefore cannot +be reused for the sparse phase. + ## Video Sparse Attention (VSA) -TODO +VSA reduces the cost of self-attention in video diffusion models by attending only to the most relevant spatial-temporal blocks. It uses a two-branch design: a lightweight coarse branch mean-pools tokens into (4, 4, 4) cubes and computes cube-level attention scores to select the top-K most relevant cubes per query, then a fine branch runs the selected backend's block-sparse kernel over those cubes only. The two outputs are blended with learned gates, which is why VSA needs a fine-tuned checkpoint. Select either `CUTEDSL` for the CuTe DSL fine-stage kernel or `TRTLLM` for PrimTS block-sparse attention. + +Requirements: + +- A VSA-fine-tuned checkpoint such as [`FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers`](https://huggingface.co/FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers). Standard Wan checkpoints do not carry the learned VSA gates. +- The `CUTEDSL` or `TRTLLM` attention backend. `CUTEDSL` uses the CuTe DSL fine-stage kernel; `TRTLLM` lowers the selected cubes through the generic PrimTS block-sparse FMHA contract. +- A supported CUDA device and tensor shape for the selected block-sparse kernel. When that kernel is unavailable or the input is outside its envelope, the fine branch runs the compact Q/K/V through that backend's dense path (SDPA for `CUTEDSL`, TRTLLM attention for `TRTLLM`) and the VSA post-processing still applies. +- VSA cannot be combined with `quant_attention_config`. +- VSA is not compatible with Ring attention or Attention2D (it does not produce per-split LSE). Ulysses is supported. + +Configure VSA with `VideoSparseAttentionConfig` and either backend: + +```python +from tensorrt_llm import VisualGenArgs +from tensorrt_llm.visual_gen.args import AttentionConfig, VideoSparseAttentionConfig + +args = VisualGenArgs( + model="FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers", + attention_config=AttentionConfig( + backend="TRTLLM", # or "CUTEDSL" + sparse_attention_config=VideoSparseAttentionConfig(vsa_sparsity=0.9), + ), +) +``` + +The equivalent YAML, for `--visual_gen_args` or `trtllm-serve`: + +```yaml +attention_config: + backend: TRTLLM # or CUTEDSL + sparse_attention_config: + algorithm: vsa + vsa_sparsity: 0.90 # fraction of K/V cubes skipped by the fine branch +``` + +`vsa_sparsity` controls the fraction of K/V cubes the fine branch skips (0.0 keeps every cube and reproduces dense attention, 0.9 skips 90% of them). Higher sparsity gives more speedup at the cost of some quality. + +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 +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..03656d075055 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -56,15 +56,15 @@ Models are auto-detected from the checkpoint directory. Diffusers-format models ### Feature Matrix -| Model | FP8 blockwise | NVFP4 | TeaCache | Cache-DiT | CPU Offloading | CFG Parallelism | Ulysses Parallelism | Parallel VAE | CUDA Graph | torch.compile | trtllm-serve | Attention2D | Ring Attention | Tensor Parallelism | VSA | +| Model | FP8 blockwise | NVFP4 | TeaCache | Cache-DiT | CPU Offloading | CFG Parallelism | Ulysses Parallelism | Parallel VAE | CUDA Graph | torch.compile | trtllm-serve | Attention2D | Ring Attention | Tensor Parallelism | SOL-Attn | |---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---| | **FLUX.1** | Yes | Yes | Yes | Yes | No | No | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | No | | **FLUX.2** | Yes | Yes | Yes | Yes | No | No | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | No | -| **Wan 2.1** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | +| **Wan 2.1** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | | **Wan 2.1 VSA** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes | Yes | -| **Wan 2.2** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | No | +| **Wan 2.2** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | | **FastWan 2.2** | Yes | Yes | No | No | No | No | No | No | Yes | Yes | Yes | No | No | No | No | -| **LTX-2** | Yes | Yes | Yes | Yes | No | Yes | Yes | No | No | Yes | Yes | Yes | Yes | No | No | +| **LTX-2** | Yes | Yes | Yes | Yes | No | Yes | Yes | No | No | Yes | Yes | Yes | Yes | No | Yes | | **MiniMax-H3** | Yes | Yes | No | No | No | No | No | No | No | Yes | Yes | No | No | No | No | | **Qwen-Image** | Yes | Yes | Yes | Yes | No | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | No | No | | **Qwen-Image-Layered** | No | No | No | No | No | No | No | No | Yes | Yes | Yes | No | No | No | No | @@ -73,6 +73,8 @@ Models are auto-detected from the checkpoint directory. Diffusers-format models | **HunyuanVideo 1.5** | Yes | Yes | No | No | No | No | No | No | No | No | Yes | No | No | No | No | | **GlmImage** | Yes | Yes | No | No | No | No | No | No | No | No | Yes | No | No | No | No | +SOL-Attn is a runtime sparse-attention feature of the `TRTLLM` and `CUTEDSL` backends (bfloat16 self-attention with `head_dim=128` on SM100/SM103); it needs no dedicated checkpoint. Video Sparse Attention (VSA) requires a VSA-fine-tuned checkpoint and is therefore listed as its own model, `Wan 2.1 VSA`. See [VisualGen Sparse Attention](../features/visualgen-sparse-attention.md) for both. + ## Quick Start Here is a simple example to generate a video with Wan 2.1: @@ -358,42 +360,15 @@ args = VisualGenArgs( **Wan 2.2 dual-transformer note:** Wan 2.2 uses two expert transformers (high-noise and low-noise stacks). All `CacheDiTConfig` parameters apply to both stacks, except `max_warmup_steps` and `max_cached_steps`: the low-noise stack always uses fixed internal caps (`max_warmup_steps=2`, `max_cached_steps=20`) regardless of user config. -### 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. - -**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. -- 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. - -Python API: - -```python -from tensorrt_llm import VisualGenArgs -from tensorrt_llm.visual_gen.args import AttentionConfig, VideoSparseAttentionConfig +### Sparse Attention -args = VisualGenArgs( - model="FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers", - attention_config=AttentionConfig( - backend="CUTEDSL", - sparse_attention_config=VideoSparseAttentionConfig(vsa_sparsity=0.9), - ), -) -``` +Attention dominates a denoising step at high resolution or long video length. VisualGen ships three sparse attention algorithms, all configured through `attention_config.sparse_attention_config` and all served by the `TRTLLM` and `CUTEDSL` attention backends: -YAML (for use with `--visual_gen_args` or `trtllm-serve`): +- **Skip Softmax Attention** (`algorithm: skip_softmax`) skips negligible softmax blocks inside the FlashAttention-style kernel. It is plug-and-play for existing checkpoints; a ModelOpt-calibrated checkpoint can additionally map a `target_sparsity` to the kernel threshold. +- **SOL Attention** (`algorithm: sol_attn`) routes attention blocks at runtime: blocks whose scores stand out are computed exactly and the rest are folded in from compact K/V proxies. It needs no dedicated checkpoint and applies to bfloat16 self-attention with `head_dim=128` on SM100/SM103 GPUs. +- **Video Sparse Attention (VSA)** (`algorithm: vsa`) combines a coarse mean-pooled branch with a top-K block-sparse fine branch and requires a VSA-fine-tuned checkpoint such as [`FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers`](https://huggingface.co/FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers). -```yaml -attention_config: - backend: CUTEDSL - sparse_attention_config: - algorithm: vsa - vsa_sparsity: 0.90 -``` +Skip Softmax and SOL share the `disabled_until_timestep` schedule that keeps the high-noise prefix dense, and every algorithm participates in the CUDA Graph key so dense and sparse phases never share a graph. Configuration, requirements and Python/YAML examples for each algorithm are in [VisualGen Sparse Attention](../features/visualgen-sparse-attention.md). ### CPU Offloading diff --git a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index a44632837d7c..df7bbf1decf6 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md @@ -224,11 +224,15 @@ 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 planned wrappers in the attention metadata +(`TrtllmAttentionMetadata.fmha_plan_caches`), like the FlashInfer wrappers and +the MSA plans. Every layer that runs with one metadata object plans each static +profile once and reuses its graph-stable route workspace, so the sharing +granularity is the metadata: one per batch in the LLM path, where each +CUDA-graph batch size gets its own copy, and one per component and sequence +shape in VisualGen, whose layers share a shape-keyed metadata cache. Layers +that share a metadata object execute serially, so they never run a planned +wrapper concurrently. `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/fallback.py b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py index 5631950e9e06..a3b9379fce52 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py @@ -45,6 +45,7 @@ "out_scale_sf", # promoted into ``out_scale`` in ``TrtllmAttention.forward`` for NVFP4 path "skip_mla_rope_generation", # handled in ``TrtllmAttention.forward`` for the test-only MLA path "timestep", # consumed by sparse prediction before FMHA dispatch + "sparse_attn_phase", # host-resolved sparse phase, consumed before FMHA dispatch } ) 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..6d754d7a8747 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 @@ -47,10 +47,7 @@ BlockSparsePagedTSWrapper, BlockSparseTSWrapper, ) - from tensorrt_llm._torch.attention.backends.trtllm import ( - TrtllmAttention, - TrtllmAttentionMetadata, - ) + from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata from tensorrt_llm._torch.attention.backends.prims_ts import ( BlockSparsePagedTSWrapper as _BlockSparsePagedTSWrapper, @@ -190,30 +187,9 @@ def _uniform_seq_len_q( class PrimsTSBlockSparseFmha(PrimsTSFmha): """Contiguous context and fixed-Q paged generation block-sparse FMHA.""" - 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``. + PLAN_CACHE_KEY = "prims_ts_block_sparse" - 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. - """ - - self._contiguous_wrappers = cast( - dict[_BlockSparsePlanKey, "BlockSparseTSWrapper"], - cache_state.setdefault("contiguous_wrappers", {}), - ) - self._paged_wrappers = cast( - dict[_BlockSparsePlanKey, "BlockSparsePagedTSWrapper"], - cache_state.setdefault("paged_wrappers", {}), - ) + supports_block_sparse_inputs = True def _is_supported( self, @@ -331,8 +307,15 @@ def _make_plan_key( def _get_or_plan_wrapper( self, key: _BlockSparsePlanKey, + metadata: "TrtllmAttentionMetadata", ) -> "BlockSparseTSWrapper | BlockSparsePagedTSWrapper": - cache = self._paged_wrappers if key.page_size is not None else self._contiguous_wrappers + """Return the planned wrapper for ``key`` from the metadata's plan cache. + + The cache lives on the attention metadata, like the FlashInfer wrappers + and the MSA plans, so every layer that runs with one metadata object + reuses a plan and its graph-stable route workspace. + """ + cache = metadata.fmha_plan_caches.setdefault(self.PLAN_CACHE_KEY, {}) wrapper = cache.get(key) if wrapper is None: wrapper = key.plan() @@ -516,7 +499,7 @@ def run_generation(self, params: FmhaParams) -> None: page_size=page_size, mask_type=self._get_prims_mask_type(forward_args), ) - wrapper = cast("BlockSparsePagedTSWrapper", self._get_or_plan_wrapper(key)) + wrapper = cast("BlockSparsePagedTSWrapper", self._get_or_plan_wrapper(key, metadata)) wrapper.run( query, (k_cache, v_cache), @@ -534,6 +517,7 @@ def _forward_contiguous( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + metadata: "TrtllmAttentionMetadata", forward_args: AttentionForwardArgs, ) -> None: inputs = _get_block_sparse_inputs(forward_args) @@ -552,7 +536,7 @@ def _forward_contiguous( page_size=None, mask_type=self._get_prims_mask_type(forward_args), ) - wrapper = cast("BlockSparseTSWrapper", self._get_or_plan_wrapper(key)) + wrapper = cast("BlockSparseTSWrapper", self._get_or_plan_wrapper(key, metadata)) wrapper.run( query, key_states, @@ -577,6 +561,6 @@ def forward( ) -> None: if metadata.kv_cache_manager is None: assert k is not None and v is not None - self._forward_contiguous(q, k, v, forward_args) + self._forward_contiguous(q, k, v, metadata, forward_args) return super().forward(q, k, v, metadata, forward_args) diff --git a/tensorrt_llm/_torch/attention/backends/interface.py b/tensorrt_llm/_torch/attention/backends/interface.py index 55b88face9d4..f638265d81ce 100644 --- a/tensorrt_llm/_torch/attention/backends/interface.py +++ b/tensorrt_llm/_torch/attention/backends/interface.py @@ -964,6 +964,10 @@ class AttentionForwardArgs: update_kv_cache: bool = True # Optional normalized diffusion timestep for timestep-varying sparse attention. timestep: Optional[torch.Tensor] = None + # Dense-prefix (0) or sparse (1) phase of a timestep-scheduled sparse + # algorithm, resolved on the host by the caller. When set, backends use it + # instead of reading ``timestep``, which CUDA Graph capture cannot do. + sparse_attn_phase: Optional[int] = None sparse_backend_args: Optional[SparseBackendForwardArgs] = None sparse_runtime_params: SparseRuntimeParams = field( diff --git a/tensorrt_llm/_torch/attention/backends/sparse/hooks.py b/tensorrt_llm/_torch/attention/backends/sparse/hooks.py index abddc8ad0b2d..ea7917178137 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/hooks.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/hooks.py @@ -257,5 +257,6 @@ def prepare_sparse_runtime_params( runtime_params = sparse_params.scheduler.get_runtime_params( runtime_params=runtime_params, timestep=forward_args.timestep, + graph_phase=forward_args.sparse_attn_phase, ) return runtime_params 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/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 44f4ef79e8a0..6f792d64dd5a 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -18,7 +18,7 @@ import os import weakref from dataclasses import dataclass, field -from typing import TYPE_CHECKING, List, Optional, Tuple +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple import torch @@ -128,6 +128,13 @@ class TrtllmAttentionMetadata(AttentionMetadata): # when beam search is enabled. beam_width: int = 1 + # Plan caches of the FMHA libraries, keyed by library. A planned wrapper owns + # workspaces that CUDA graphs capture, so it lives exactly as long as this + # metadata; every layer that runs with this metadata reuses it. + fmha_plan_caches: Dict[str, dict] = field(default_factory=dict, + init=False, + repr=False) + @property def effective_beam_width(self) -> int: # Only use this for the fallback kernel's beam_width argument. @@ -346,6 +353,9 @@ def __post_init__(self) -> None: or self.runtime_features.cache_reuse or self.runtime_features.has_speculative_draft_tokens ) if self.runtime_features is not None else False + # CUDA-graph metadata is a shallow copy that re-runs this method; give it + # its own plan caches so each captured batch size plans its own wrappers. + self.fmha_plan_caches = {} self._post_init_with_buffers(self.cuda_graph_buffers) def update_position_offsets_for_cpp(self, query_len: int) -> None: diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py b/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py index bdc9fe9f33c1..6800d9a629bd 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 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", @@ -49,17 +39,11 @@ "FlashAttn4Attention", "FlashInferAttention", "RingAttention", - "SolAttention", "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..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,34 +15,12 @@ """ CuTe DSL attention backend family for visual generation models. - fmha.py — CuTeDSLAttention (dense and blockscaled JIT FMHA) - vsa.py — VSAAttention (Video Sparse Attention, CuTe JIT + SDPA fallback) - 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 -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..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 @@ -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 @@ -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/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/sol/__init__.py b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/__init__.py new file mode 100644 index 000000000000..c93e7392ffc3 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/__init__.py @@ -0,0 +1,17 @@ +# 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 SolPredictorOutputs, predict, support_reason + +__all__ = [ + "SOLCuTeDSLAttention", + "SOLTrtllmAttention", + "SolParams", + "SolPredictorOutputs", + "predict", + "support_reason", +] 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 60% 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 5d0fea9a1645..5872ccbbcb6b 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,55 @@ # 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: 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 + 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.sparse.skip_softmax import SkipSoftmaxScheduler +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 . import predictor as sol_predictor +from .params import SolParams +from .predictor import BLOCK_SIZE _sol_attn_import_error = None try: @@ -69,13 +75,13 @@ 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: - 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 +90,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 +107,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 @@ -118,23 +128,22 @@ 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." ) 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 @@ -173,9 +182,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, ) @@ -187,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", @@ -248,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 @@ -313,3 +322,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) + + @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 dense or + sparse phase, resolved by the wrapper forward, from ``forward_args``. + """ + + if not self.should_use_sparse(forward_args.sparse_attn_phase): + return None + + if self.quant_attention_config is not None: + raise ValueError("SOL sparse execution does not support quant_attention_config") + if not any( + isinstance(fmha, PrimsTSBlockSparseFmha) for fmha in self._fmha_manager.fmha_libs + ): + 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 = sol_predictor.support_reason(q, k, v) + if unsupported_reason is not None: + raise ValueError(unsupported_reason) + + 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, + 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..5bf870459233 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/kernels.py @@ -0,0 +1,702 @@ +# 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 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 + +import math +from typing import Literal + +import torch +import triton +import triton.language as tl + +_POOL_MAX_WIDTH = 1024 +_POOL_TOKENS_PER_LOAD = 8 +_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]: + """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): + 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=((offsets < BLOCK) & (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 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" + ) + 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 thresholds +@triton.jit +def _block_statistics_kernel( + x_ptr, + mean_ptr, + var_ptr, + moment_ptr, + num_blocks, + num_heads, + stride_x_batch, + stride_x_block, + stride_x_head, + stride_s_batch, + stride_s_head, + stride_m_batch, + stride_m_head, + EXACT: tl.constexpr, + ROWS: tl.constexpr, + COLS: tl.constexpr, + HEAD_DIM: tl.constexpr, +): + """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) + row_valid = rows < num_blocks + tile = tl.load( + base + rows[:, None] * stride_x_block + cols[None, :], + mask=row_valid[:, None], + other=0.0, + ) + 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 + 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_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_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: + 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. + """ + 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 +@triton.jit +def _select_exact_blocks_kernel( + centroid_ptr, + keys_ptr, + threshold_ptr, + bits_ptr, + num_q_blocks, + num_kv_blocks, + num_words, + num_heads, + local_radius, + log2_scale, + stride_c_batch, + stride_c_block, + stride_c_head, + stride_k_batch, + stride_k_block, + stride_k_head, + stride_t_batch, + stride_t_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, + ) + 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) + 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, + threshold: torch.Tensor, + exact_block_bits: torch.Tensor, + *, + sm_scale: float, +) -> None: + log2_scale = float(sm_scale) * _LOG2_E + q = centroid.to(torch.float64) + k = k_summary.to(torch.float64) + scores = torch.einsum("bqhd,bkhd->bhqk", q, k) * log2_scale + 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] + 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, + threshold: torch.Tensor, + exact_block_bits: torch.Tensor, + *, + 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'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). + 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)]``. + 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") + 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() + 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, threshold, exact_block_bits, 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, + threshold, + bits, + q_blocks, + kv_blocks, + num_words(kv_blocks), + num_heads, + _LOCAL_RADIUS, + float(sm_scale) * _LOG2_E, + centroid.stride(0), + centroid.stride(1), + centroid.stride(2), + k_summary.stride(0), + k_summary.stride(1), + k_summary.stride(2), + threshold.stride(0), + threshold.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=(), device_types="cuda") +def visual_gen_sol_predictor( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + block_size: int, + tau: float, + sm_scale: float, + 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") + 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") +def _( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + block_size: int, + tau: float, + sm_scale: float, + 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_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 new file mode 100644 index 000000000000..f8a2c7e79c28 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/params.py @@ -0,0 +1,75 @@ +# 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 block routing threshold policy of both. + """ + + 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..231826d659e9 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/predictor.py @@ -0,0 +1,102 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""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 +from dataclasses import dataclass +from typing import Literal + +import torch + +from . import kernels as _kernels # noqa: F401 (registers trtllm::visual_gen_sol_predictor) + +BLOCK_SIZE = 64 +HEAD_DIM = 128 +ThreshType = Literal["diag", "exact"] + + +@dataclass(frozen=True) +class SolPredictorOutputs: + """Predictor tensors consumed by block-sparse attention.""" + + exact_block_bits: torch.Tensor + k_summary: torch.Tensor + v_summary: torch.Tensor + + +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 + + +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/_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..3971961794ed --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/__init__.py @@ -0,0 +1,35 @@ +# 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 + +__all__ = [ + "VSA_TILE_SIZE", + "VSAForwardInputs", + "VSAMetadata", + "VSAMetadataBuilder", + "VSAPredictor", + "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/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 new file mode 100644 index 000000000000..83abf9cff8f1 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/metadata.py @@ -0,0 +1,236 @@ +# 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] +_BITS_PER_WORD = 32 + + +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) + + +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. + + ``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 + 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 + tile_source_index: torch.LongTensor + untile_idx: torch.LongTensor + kv_valid_words: torch.Tensor + + +class VSAMetadataBuilder: + """Build VSA metadata while caching shape-dependent index tensors. + + 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( + 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) + ] + + 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 + + return _VSAShapeMetadata( + num_cubes=num_cubes, + padded_seq_length=padded_seq_length, + variable_block_sizes=variable_block_sizes, + tile_source_index=tile_source_index, + untile_idx=untile_idx, + kv_valid_words=_pack_valid_slots(tile_source_index >= 0), + ) + + 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: + 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..6fa244039981 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/predictor.py @@ -0,0 +1,284 @@ +# 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 math import ceil +from typing import Optional + +import torch + +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 VSA_BLOCK_SIZE, VSAMetadata, get_vsa_forward_context + +_SIGNED_INT32_MAX = torch.iinfo(torch.int32).max + + +@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`` 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: torch.LongTensor = field(repr=False) + fine_is_tiled: bool + + +@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) -> None: + self._indptr_cache: dict[tuple[torch.device, int, int, int, int], torch.Tensor] = {} + + def from_selected_blocks( + self, + selected_blocks: 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) + block_indptr = self._indptr_cache.get(key) + if block_indptr is None: + 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=sort_last_dim(selected_blocks).reshape(-1), + kv_valid_bits=kv_valid_words.unsqueeze(0).expand(batch_size, -1).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 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: + """Produce the complete per-call VSA block-attention input envelope.""" + + 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( + "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() + + @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() + num_cubes = metadata.num_cubes + cur_topk = max(1, ceil((1.0 - metadata.vsa_sparsity) * 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) + # 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: + block_sparse_inputs = self._route_builder.from_selected_blocks( + topk_indices, + metadata.kv_valid_words, + ) + + return VSAForwardInputs( + 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=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, + gate_compress=gate_compress, + gate_fine=gate_fine, + untile_idx=metadata.untile_idx, + fine_is_tiled=use_sparse_fine, + ), + ) + + +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.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, + ) + + +__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..53d29a2adf35 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py @@ -27,10 +27,16 @@ 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.sparse.timestep_phase import graph_phase_for_timestep from ...attention.backends.trtllm import TrtllmAttention as BaseTrtllmAttention from ...attention.backends.trtllm import TrtllmAttentionMetadata as BaseTrtllmAttentionMetadata +from ..cuda_graph_runner import resolved_extra_key from .interface import AttentionBackend, AttentionTensorLayout @@ -185,6 +191,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 +206,14 @@ 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." + ) super().__init__( layer_idx=layer_idx, @@ -222,6 +234,48 @@ def __init__( self.quant_attention_config = quant_attention_config + @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) + + @property + def dense_layers(self) -> frozenset[int]: + """Layer indices that always run dense attention.""" + + return getattr(self.sparse_params, "dense_layers", frozenset()) + + @torch.compiler.disable + def _sparse_attn_phase(self, timestep: object) -> Optional[int]: + """Return the dense (0) or sparse (1) phase of this call, or ``None``. + + The phase belongs to the transformer call, not to the layer: every layer + of one call sees the same timestep and the same cutoff. The CUDA graph + runner resolves it once per call on the host and publishes it during + warmup and capture, so capture never reads the timestep tensor. Eager + calls outside the runner derive it from the timestep. Layers without a + timestep cutoff have no phase. + """ + + if self.timestep_cutoff is None: + return None + phase = resolved_extra_key("sparse_attn_phase") + if phase is None: + phase = graph_phase_for_timestep(timestep, disabled_until_timestep=self.timestep_cutoff) + return phase + + def should_use_sparse(self, sparse_attn_phase: Optional[int]) -> bool: + """Return whether this layer runs its sparse path in ``sparse_attn_phase``. + + Dense layers never do. Otherwise only the dense prefix (phase 0) is + dense; a call without a phase (no cutoff or no timestep) is sparse. + """ + + if self.layer_idx in self.dense_layers: + return False + return sparse_attn_phase != 0 + # Needed to work with torch compile cause of attention metadata # make attn metadata as input for it to work @torch.compiler.disable @@ -245,6 +299,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, + ) -> 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. + 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, @@ -254,6 +326,8 @@ def forward( seq_len: int, 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: """ @@ -263,10 +337,13 @@ 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 the sparse calls of + backends that reject fused QKV, will be passed to the core as + separate tensors; the dense-layer and dense-phase calls of those + backends are fused internally like regular attention Args: q: Query tensor [B, S, H, D] or fused QKV [B, S, H_qkv, D] @@ -276,49 +353,73 @@ 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. + 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] """ + sparse_attn_phase = self._sparse_attn_phase(timestep) + block_sparse_inputs = ( + sparse_backend_args.block_sparse_inputs if sparse_backend_args is not None else None + ) + # A backend that predicts routes inside the core needs separate Q, K and + # V only on the calls that run its sparse path. Its dense-layer and + # dense-phase calls take the fused path: the TRTLLM kernel serves dense + # self-attention from fused QKV only. + use_separate_qkv = ( + block_sparse_inputs is not None + or self.quant_attention_config is not None + or (not self.support_fused_qkv() and self.should_use_sparse(sparse_attn_phase)) + ) + if use_separate_qkv and (k is None or v is None): + raise ValueError("This TRTLLM attention call requires separate q, k, and v tensors.") + if block_sparse_inputs is not None and self.quant_attention_config is not None: + raise ValueError( + "Generic block-sparse attention does not support quant_attention_config." + ) + kv_seq_len = seq_len_kv if seq_len_kv is not None else seq_len prepared_metadata = self._prepare_metadata(batch_size, seq_len) - timestep = kwargs.pop("timestep", None) - - 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" - ) + sage_kwargs = {} + if use_separate_qkv: + q, k, v = self._compact_qkv(q, k, v, batch_size, seq_len, kv_seq_len) quant_cfg = self.quant_attention_config - 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"), - ) + 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_attn_phase=sparse_attn_phase, + 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..e3109e354e50 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py @@ -32,12 +32,16 @@ def get_visual_gen_attention_backend( backend_name: str, + sparse_algorithm: Optional[str] = None, ) -> Type[AttentionBackend]: """ - Get diffusion attention backend class by name. + Get diffusion attention backend class by name and sparse algorithm. Args: backend_name: Backend identifier ("VANILLA", "TRTLLM", "CUDNN", "FLASHINFER", "CUTEDSL", "FA4") + sparse_algorithm: ``algorithm`` of the sparse attention config, if any. VSA + ("vsa") and SOL ("sol_attn") have their own TRTLLM and CUTEDSL classes; + other algorithms, such as skip-softmax, run in the dense class. Returns: Diffusion attention backend class @@ -45,41 +49,63 @@ 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) Requires flash-attn package with cute interface - - "CUTEDSL": CuTe DSL kernels. create_attention selects dense/SkipSoftmax FMHA or VSA - from AttentionConfig.sparse_attention_config. + - "CUTEDSL": CuTe DSL kernels: dense/SkipSoftmax FMHA, VSA or SOL depending on + ``sparse_algorithm``. - "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": + if sparse_algorithm == "vsa" and backend_name == "TRTLLM": + from .sparse.vsa.backend import VSATrtllmAttention + + return VSATrtllmAttention + elif sparse_algorithm == "vsa" and backend_name == "CUTEDSL": + from .sparse.vsa.backend import VSACuTeDSLAttention + + return VSACuTeDSLAttention + elif sparse_algorithm == "sol_attn" and backend_name == "TRTLLM": + from .sparse.sol.backend import SOLTrtllmAttention + + return SOLTrtllmAttention + elif sparse_algorithm == "sol_attn" and backend_name == "CUTEDSL": + from .sparse.sol.backend import SOLCuTeDSLAttention + + return SOLCuTeDSLAttention + elif 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 +143,51 @@ 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_config = ( + attention_config.sparse_attention_config if attention_config is not None else 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" + + backend_name = backend.upper() + attn_cls = get_visual_gen_attention_backend(backend_name, sparse_algorithm) + + 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.") + 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_config.to_sparse_params() # 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..bc06b1836772 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, whose metadata objects also carry the FMHA + plan caches, so one static block-sparse profile is planned once per component + and shape 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/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/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/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/_torch/visual_gen/models/wan/pipeline_wan.py b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py index 3a4bc10b33c0..5f78ed1d5be0 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, ) @@ -148,6 +148,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 +618,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..91dc235a4dde 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py @@ -1,11 +1,26 @@ +# 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 +from typing import Any, Optional, Tuple import torch 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 @@ -96,19 +111,26 @@ 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 = base_backend == "CUTEDSL" and _sa_algo == "vsa" - _is_sol_attn = base_backend == "CUTEDSL" and _sa_algo == "sol_attn" - separate_qkv_cross_attention = ( + 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 ) - # Cross-attention fallback: TRTLLM and CUTEDSL VSA 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): + # 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( @@ -121,10 +143,10 @@ 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: + 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"{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." ) @@ -225,18 +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(sparse_config, SolAttentionConfig) and backend_name in ( + "TRTLLM", + "CUTEDSL", + ): + sparse_params = sparse_config.to_sparse_params() self.sparse_params = sparse_params # Create compute backend @@ -253,12 +279,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( @@ -533,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. @@ -600,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. @@ -650,6 +671,7 @@ def forward_async( hidden_states: torch.Tensor, freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, timestep: Optional[torch.Tensor] = None, + **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 @@ -689,8 +711,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 +768,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..b08d322f95a0 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -126,8 +126,8 @@ class AttentionConfig(StrictBaseModel): status="prototype", description=( "Sparse attention recipe. Discriminated by algorithm: " - "skip_softmax (TRTLLM / CUTEDSL backends), vsa (CUTEDSL backend), " - "or sol_attn (CUTEDSL backend)." + "skip_softmax (TRTLLM / CUTEDSL backends), vsa (CUTEDSL / TRTLLM backends), " + "or sol_attn (TRTLLM / CUTEDSL backends)." ), ) @@ -224,8 +224,8 @@ def _validate_sparse_attention_config(self) -> "AttentionConfig": algo = self.sparse_attention_config.algorithm supported_backends = { "skip_softmax": ("TRTLLM", "CUTEDSL"), - "vsa": ("CUTEDSL",), - "sol_attn": ("CUTEDSL",), + "vsa": ("CUTEDSL", "TRTLLM"), + "sol_attn": ("TRTLLM", "CUTEDSL"), }.get(algo) if supported_backends is None: return self @@ -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 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 and quant_attention_config are mutually exclusive.") return self @@ -848,6 +842,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 bd4c22e22395..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. @@ -43,6 +46,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. @@ -223,29 +241,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=( + "Block routing threshold policy: 'diag' models each key channel " + "independently, 'exact' uses the full key covariance." + ), ) disabled_until_timestep: Optional[float] = PydanticField( None, @@ -254,15 +283,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( @@ -284,20 +311,27 @@ 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): - # 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 + 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 + + 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): - """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 456697cc3676..2ca039299f9c 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -302,12 +302,18 @@ 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_cute_dsl_sol_attn.py + - unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_attention.py + - unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_kernels.py + - unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_parity.py + - unittest/_torch/visual_gen/sparse_attention/sol/test_sol_attention.py + - unittest/_torch/visual_gen/sparse_attention/sol/test_sol_parity.py + - unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor.py + - unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor_kernels.py - unittest/_torch/visual_gen/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/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 14307de1b9dc..74553b291ec5 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -35,7 +35,7 @@ l0_b300: - unittest/_torch/visual_gen/kernels/parallel - unittest/_torch/visual_gen/test_attention_flashinfer.py # ------------- Visual Gen sparse attention (sm103) --------------- - - unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py + - unittest/_torch/visual_gen/sparse_attention/sol/test_sol_attention.py - unittest/_torch/thop/serial - unittest/_torch/executor # 250s - unittest/_torch/disaggregation diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index c1da5c4035a6..c484a49e71c9 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -90,7 +90,10 @@ l0_cpu: - unittest/_torch/visual_gen/multi_gpu/test_attn2d_attention.py - 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/skip_softmax/test_skip_softmax.py + - unittest/_torch/visual_gen/sparse_attention/sol/test_sol_attention.py + - unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor.py + - unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor_kernels.py - unittest/_torch/visual_gen/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_prims_ts_block_sparse.py b/tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py index 40d74b28ede3..44562722fce1 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 @@ -149,6 +149,7 @@ def _contiguous_case(): helix_position_offsets=None, kv_cache_manager=None, seq_lens=torch.tensor([64, 64], dtype=torch.int32), + fmha_plan_caches={}, ) args = AttentionForwardArgs( output=torch.empty_like(q), @@ -187,6 +188,7 @@ def _paged_metadata(): host_kv_cache_pool_pointers=torch.tensor([[1234, 5678]], dtype=torch.int64), host_kv_cache_pool_mapping=torch.tensor([[0, 0]], dtype=torch.int32), kv_cache_manager=manager, + fmha_plan_caches={}, ) @@ -320,7 +322,7 @@ def test_block_sparse_support_rejects_invalid_static_kernel_profile( def test_contiguous_wrappers_cache_static_profile_and_keep_routes_live(monkeypatch) -> None: - _attention, fmha, q, k, v, _metadata, args = _contiguous_case() + _attention, fmha, q, k, v, metadata, args = _contiguous_case() wrapper = Mock() factory = Mock(return_value=wrapper) monkeypatch.setattr(block_sparse_fmha, "_BlockSparseTSWrapper", factory) @@ -337,12 +339,12 @@ def test_contiguous_wrappers_cache_static_profile_and_keep_routes_live(monkeypat ] for inputs in bsr_inputs: _set_block_sparse_inputs(args, inputs) - fmha._forward_contiguous(q, k, v, args) + fmha._forward_contiguous(q, k, v, metadata, args) proxy_inputs = [_bitmask_inputs(proxy=True), _bitmask_inputs(proxy=True)] for inputs in proxy_inputs: _set_block_sparse_inputs(args, inputs) - fmha._forward_contiguous(q, k, v, args) + fmha._forward_contiguous(q, k, v, metadata, args) assert factory.call_count == 2 assert wrapper.plan.call_count == 2 @@ -385,23 +387,25 @@ def _key(fmha): assert _key(first) != _key(second) -def test_block_sparse_plan_cache_is_shared_only_when_explicitly_bound() -> 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 +def test_block_sparse_plan_cache_lives_on_the_metadata(monkeypatch) -> None: + """Layers that run with one metadata object share its planned wrappers.""" + wrapper = Mock() + factory = Mock(return_value=wrapper) + monkeypatch.setattr(block_sparse_fmha, "_BlockSparseTSWrapper", factory) + _attention, first, q, k, v, metadata, args = _contiguous_case() + second_attention = _Attention() + second = block_sparse_fmha.PrimsTSBlockSparseFmha(second_attention) - cache_state = {} - first.bind_plan_cache(cache_state) - second.bind_plan_cache(cache_state) + first._forward_contiguous(q, k, v, metadata, args) + second._forward_contiguous(q, k, v, metadata, args) + assert factory.call_count == 1 + assert list(metadata.fmha_plan_caches) == [ + block_sparse_fmha.PrimsTSBlockSparseFmha.PLAN_CACHE_KEY + ] - assert first._contiguous_wrappers is second._contiguous_wrappers - assert first._paged_wrappers is second._paged_wrappers - assert cache_state == { - "contiguous_wrappers": {}, - "paged_wrappers": {}, - } + other_metadata = SimpleNamespace(**{**vars(metadata), "fmha_plan_caches": {}}) + second._forward_contiguous(q, k, v, other_metadata, args) + assert factory.call_count == 2 def test_paged_wrapper_uses_zero_copy_padded_row_stride_block_tables(monkeypatch) -> None: @@ -541,6 +545,7 @@ def test_prepare_workspace_checks_capture_before_resize(monkeypatch) -> None: max_num_requests=2, tokens_per_block=64, num_generations=2, + fmha_plan_caches={}, ) workspace = torch.empty(0, dtype=torch.uint8) monkeypatch.setattr( @@ -654,6 +659,7 @@ def test_real_gpu_proxy_adapter_replays_live_routes_and_summaries() -> None: is_cross=False, kv_cache_manager=None, seq_lens=torch.tensor([64], dtype=torch.int32), + fmha_plan_caches={}, ) flat_q, flat_k, flat_v = (tensor.flatten(0, 2) for tensor in (q, k, v)) diff --git a/tests/unittest/_torch/attention/sparse/test_sparse_attention.py b/tests/unittest/_torch/attention/sparse/test_sparse_attention.py index 6cf7b15fcc38..d36c02ad903c 100644 --- a/tests/unittest/_torch/attention/sparse/test_sparse_attention.py +++ b/tests/unittest/_torch/attention/sparse/test_sparse_attention.py @@ -216,9 +216,19 @@ def test_prepare_sparse_runtime_params_schedules_skip_softmax_thresholds() -> No attention, torch.empty(0), None, None, None, forward_args ) - schedule.assert_called_once_with(runtime_params=SparseRuntimeParams(), timestep=timestep) + schedule.assert_called_once_with( + runtime_params=SparseRuntimeParams(), timestep=timestep, graph_phase=None + ) assert runtime_params == scheduler.get_runtime_params(timestep=timestep) + # A caller-resolved phase reaches the scheduler, which then never reads the timestep. + forward_args = AttentionForwardArgs(timestep=timestep, sparse_attn_phase=1) + with patch.object( + scheduler, "get_runtime_params", wraps=scheduler.get_runtime_params + ) as schedule: + prepare_sparse_runtime_params(attention, torch.empty(0), None, None, None, forward_args) + assert schedule.call_args.kwargs["graph_phase"] == 1 + def _make_block_sparse_inputs() -> BlockSparseForwardInputs: return BlockSparseForwardInputs( 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..1769854c7841 --- /dev/null +++ b/tests/unittest/_torch/attention/sparse/test_timestep_phase.py @@ -0,0 +1,47 @@ +# 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) -> None: + assert timestep_to_float(timestep) == ( + expected if expected is None else pytest.approx(expected) + ) + + +@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 +) -> 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 8415bd6bc52a..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 @@ -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,55 @@ def run_test_in_distributed(world_size: int, test_fn: Callable, use_cuda: bool = ) +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: torch.Tensor, k: torch.Tensor, v: torch.Tensor, **kwargs: object + ) -> torch.Tensor: + 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: torch.Tensor, **kwargs: object) -> torch.Tensor: + 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..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,5 +317,66 @@ 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() -> None: + from tensorrt_llm._torch.visual_gen.modules.attention import Attention + + class _CaptureAsyncAttention(torch.nn.Module): + 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() + self.kwargs = kwargs + return self.q + + class _CaptureOutputProjection(torch.nn.Module): + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + 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 + 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) + + 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..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 @@ -521,6 +521,55 @@ 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: 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 + 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/sparse_attention/__init__.py b/tests/unittest/_torch/visual_gen/sparse_attention/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/sparse_attention/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/skip_softmax/__init__.py b/tests/unittest/_torch/visual_gen/sparse_attention/skip_softmax/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/sparse_attention/skip_softmax/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py b/tests/unittest/_torch/visual_gen/sparse_attention/skip_softmax/test_skip_softmax.py similarity index 98% rename from tests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py rename to tests/unittest/_torch/visual_gen/sparse_attention/skip_softmax/test_skip_softmax.py index 3347c9240232..ddd1cb0ad219 100644 --- a/tests/unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py +++ b/tests/unittest/_torch/visual_gen/sparse_attention/skip_softmax/test_skip_softmax.py @@ -13,10 +13,8 @@ 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 ( CuTeDSLAttention, @@ -393,6 +391,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 +403,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 +486,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 diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/sol/__init__.py b/tests/unittest/_torch/visual_gen/sparse_attention/sol/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/sparse_attention/sol/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_attention.py b/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_attention.py new file mode 100644 index 000000000000..d0ef18211375 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_attention.py @@ -0,0 +1,1252 @@ +# 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. + +"""Behavior tests for the VisualGen SOL backends. + +The TRTLLM backend predicts SOL routes inside the core prediction hook and +executes them through the PrimTS block-sparse FMHA; the CuTeDSL backend runs +the vendored fused kernel. This module covers dispatch, the public config, the +timestep schedule, CUDA Graph phase keys and the guards of both backends. +Numerical checks against dense attention, the fp32 SOL reference and each +other live in ``test_sol_parity.py``. +""" + +from __future__ import annotations + +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.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 TrtllmAttention as CoreTrtllmAttention +from tensorrt_llm._torch.visual_gen.attention_backend import CuTeDSLAttention +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 ( + SOLCuTeDSLAttention, + 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, +) +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, + create_attention_metadata_state, +) +from tensorrt_llm._torch.visual_gen.cuda_graph_runner import ( + CUDAGraphRunner, + CUDAGraphRunnerConfig, + resolved_extra_key, + resolved_extra_keys_scope, +) +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 + +_CPU_ONLY = pytest.mark.cpu_only +_REQUIRES_CUDA = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +_BACKENDS = pytest.mark.parametrize("backend", ["TRTLLM", "CUTEDSL"]) + + +# --------------------------------------------------------------------------- TRTLLM helpers +def _make_backend(params: SolParams, *, 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.metadata = TrtllmAttentionMetadata( + device=torch.device("cpu"), attention_metadata_state={} + ) + 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, +) -> BlockSparseForwardInputs | 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, + sparse_attn_phase=backend._sparse_attn_phase(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( + monkeypatch, + params: SolParams | None = None, + *, + seq_len: int = 64, + unsupported_reason: str | None = None, +) -> 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 + + +# --------------------------------------------------------------------------- TRTLLM backend +@_CPU_ONLY +def test_sol_backend_uses_the_runner_resolved_phase_during_capture(monkeypatch) -> None: + q = _bshd() + backend, predictor = _stub_backend(monkeypatch, SolParams(tau=1.0, disabled_until_timestep=0.6)) + + # The runner resolved the sparse phase for this call; the dense-looking + # timestep tensor is not read. + with resolved_extra_keys_scope({"sparse_attn_phase": 1}): + assert _predict(backend, q, q, q, timestep=torch.tensor(0.8)) is not None + predictor.predict.assert_called_once() + + +@_CPU_ONLY +def test_sol_backend_without_timestep_runs_sparse_like_skip_softmax(monkeypatch) -> None: + q = _bshd() + backend, predictor = _stub_backend(monkeypatch, SolParams(tau=1.0, disabled_until_timestep=0.6)) + + assert _predict(backend, q, q, q, timestep=None) is not None + predictor.predict.assert_called_once() + + +@_CPU_ONLY +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._sparse_attn_phase(torch.tensor([0.0, 0.8]))) + assert backend.should_use_sparse(backend._sparse_attn_phase(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 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__ + 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(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, + "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) + 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(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) + + 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( + monkeypatch, + k: torch.Tensor | None, + v: torch.Tensor | None, + attention_mask: PredefinedAttentionMask, + message: str, +) -> None: + q = _bshd() + backend, predictor = _stub_backend(monkeypatch) + + with pytest.raises(ValueError, match=message): + _predict(backend, q, k, v, attention_mask=attention_mask) + + 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( + monkeypatch, + params: SolParams, + layer_idx: int, + timestep: float | None, +) -> None: + q = _bshd() + backend, predictor = _stub_backend(monkeypatch, 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(monkeypatch) -> None: + q = _bshd() + backend, predictor = _stub_backend(monkeypatch) + 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() + + +# --------------------------------------------------------------------------- CuTeDSL helpers +def _kernel_wrapper(): + from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell import sol_attn_backend + + return sol_attn_backend + + +def _cutedsl_backend( + *, layer_idx: int = 1, num_heads: int = 2, head_dim: int = 128 +) -> SOLCuTeDSLAttention: + """CuTeDSL SOL backend with the dense prefix disabled unless a test sets it.""" + + backend = SOLCuTeDSLAttention(layer_idx=layer_idx, num_heads=num_heads, head_dim=head_dim) + backend.disabled_until_timestep = None + return backend + + +def _cutedsl_module_config( + hidden_size: int, num_heads: int, head_dim: int, *, tau: float +) -> DiffusionModelConfig: + """Minimal DiffusionModelConfig for one CuTeDSL SOL attention module.""" + + return DiffusionModelConfig( + pretrained_config=SimpleNamespace( + hidden_size=hidden_size, + num_attention_heads=num_heads, + attention_head_dim=head_dim, + eps=1e-6, + ), + attention=AttentionConfig( + backend="CUTEDSL", sparse_attention_config=SolAttentionConfig(tau=tau) + ), + skip_create_weights_in_init=False, + ) + + +def _fake_cuda_q(shape, dtype=torch.bfloat16) -> SimpleNamespace: + """A tensor-like that reports ``is_cuda=True`` without needing a GPU. + + ``sol_attn_ineligible_reason`` checks ``is_cuda`` first and returns early, so + a CPU tensor can never reach the rank, head_dim, or dtype branches. These + stubs pass that first gate so each later reason is actually exercised; the + architecture check comes after them and is not reached. + """ + + return SimpleNamespace(is_cuda=True, ndim=len(shape), shape=shape, dtype=dtype) + + +def _masked_sdpa_reference(q, k, v, *, key_padding_mask=None, is_causal=False) -> torch.Tensor: + attn_mask = ( + None if key_padding_mask is None else key_padding_mask.to(torch.bool)[:, None, None, :] + ) + return torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2).float(), + k.transpose(1, 2).float(), + v.transpose(1, 2).float(), + attn_mask=attn_mask, + is_causal=is_causal and attn_mask is None, + ).transpose(1, 2) + + +# --------------------------------------------------------------------------- CuTeDSL backend +@_CPU_ONLY +def test_cutedsl_factory_dispatches_dense_and_sol() -> None: + dense_attention = create_attention( + backend="CUTEDSL", + layer_idx=0, + num_heads=8, + head_dim=128, + attention_config=AttentionConfig(backend="CUTEDSL"), + ) + sol_attention = create_attention( + backend="CUTEDSL", + layer_idx=0, + num_heads=8, + head_dim=128, + attention_config=AttentionConfig( + backend="CUTEDSL", + sparse_attention_config=SolAttentionConfig(tau=2.0, disabled_until_timestep=0.9545), + ), + ) + + assert isinstance(dense_attention, CuTeDSLAttention) + assert isinstance(sol_attention, SOLCuTeDSLAttention) + assert sol_attention.tau == 2.0 + assert sol_attention.disabled_until_timestep == 0.9545 + + +@_REQUIRES_CUDA +def test_cutedsl_sol_cross_attention_delegates_per_call() -> None: + """Cross-attention is delegated to the dense backend, decided per call. + + SOL is self-attention only. It is still the module's backend for a + cross-attention module and delegates at ``forward`` because ``_can_serve`` + sees ``k.shape[1] != q.shape[1]``. Deciding this per call rather than at + construction matters: ``qkv_mode`` describes how Q/K/V are projected, not + whether K/V come from another sequence, so a construction-time rule keyed + on SEPARATE_QKV would strip the configured backend from self-attention + modules that use that mode for unrelated reasons (Qwen-Image; WAN attn1 + under async Ulysses). + """ + + device = torch.device("cuda") + cross_attn = ( + Attention( + 64, + 4, + qkv_mode=QKVMode.SEPARATE_QKV, + config=_cutedsl_module_config(64, 4, 16, tau=1.0), + ) + .to(device=device, dtype=torch.bfloat16) + .eval() + ) + assert cross_attn.attn_backend == "CUTEDSL" + assert isinstance(cross_attn.attn, SOLCuTeDSLAttention) + + q = torch.randn(1, 32, 4, 16, device=device, dtype=torch.bfloat16) + k = torch.randn(1, 77, 4, 16, device=device, dtype=torch.bfloat16) + assert not cross_attn.attn._can_serve(q, k), ( + "differing q/k sequence lengths must be delegated, not routed to the sparse kernel" + ) + + +@_CPU_ONLY +def test_cutedsl_sol_serves_separate_qkv_self_attention() -> None: + """SEPARATE_QKV self-attention keeps SOL: WAN attn1 under async Ulysses and Qwen-Image.""" + + backend = _cutedsl_backend() + # 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. + q = k = torch.randn(1, 64, 2, 128, dtype=torch.bfloat16) + assert backend._can_serve(q, k), "equal q/k sequence lengths must reach the sparse kernel" + + +@_CPU_ONLY +def test_cutedsl_sol_rejects_gqa_mqa() -> None: + with pytest.raises(ValueError, match="MHA-only"): + SOLCuTeDSLAttention(layer_idx=0, num_heads=8, head_dim=128, num_kv_heads=2) + + +@_CPU_ONLY +def test_cutedsl_sol_dense_prefix_skips_kernel(monkeypatch) -> None: + """Inside the dense prefix the sparse kernel must not be invoked at all. + + CPU tensors, so ``_delegate`` takes its SDPA branch here; that the dense + path routes to the CuTe kernel on CUDA is covered by + ``test_cutedsl_sol_dense_paths_use_cutedsl_dense_kernel``. + """ + + def _fail_if_called(*args, **kwargs): + raise AssertionError("kernel must not run inside the dense prefix") + + monkeypatch.setattr(sol_backend, "_sol_attn_run", _fail_if_called) + backend = SOLCuTeDSLAttention(layer_idx=0, num_heads=2, head_dim=16) + backend.disabled_until_timestep = 0.9 + q = k = v = torch.randn(1, 4, 2, 16) + + out = backend.forward(q, k, v, timestep=torch.tensor(0.95)) + + assert out.shape == q.shape + assert torch.isfinite(out).all() + + +@_CPU_ONLY +def test_cutedsl_sol_missing_timestep_fails_open_to_sparse(monkeypatch) -> None: + """Without a timestep the prefix cannot be applied; run sparse, do not raise.""" + + called = {"n": 0} + + def _record(*args, **kwargs): + called["n"] += 1 + return args[0] + + monkeypatch.setattr(sol_backend, "_sol_attn_run", _record) + backend = SOLCuTeDSLAttention(layer_idx=0, num_heads=2, head_dim=16) + backend.disabled_until_timestep = 0.9 + q = k = v = torch.randn(1, 4, 2, 16) + + backend.forward(q, k, v) + + assert called["n"] == 1, "expected the sparse kernel, not a silent dense fallback" + + +@_CPU_ONLY +def test_cutedsl_sol_dense_layers_guard_skips_kernel(monkeypatch) -> None: + def _fail_if_called(*args, **kwargs): + raise AssertionError("kernel must not be invoked for a dense_layers-forced layer") + + monkeypatch.setattr(sol_backend, "_sol_attn_run", _fail_if_called) + backend = SOLCuTeDSLAttention(layer_idx=3, num_heads=2, head_dim=16) + backend.dense_layers = frozenset({3}) + q = k = v = torch.randn(1, 4, 2, 16) + + out = backend.forward(q, k, v) + + assert out.shape == q.shape + assert torch.isfinite(out).all() + + +@_CPU_ONLY +def test_cutedsl_sol_dense_by_step_prefers_runner_resolved_phase(monkeypatch) -> None: + """Inside a runner-driven call the phase comes from the runner, not the tensor. + + CUDA Graph capture forbids device-to-host syncs, and reading the timestep is + a ``.item()``. The runner resolves the phase host-side to build the graph + key and republishes it during capture; ``_dense_by_step`` must use that and + never touch the tensor. + """ + + import tensorrt_llm._torch.attention.backends.sparse.timestep_phase as timestep_phase + + reads = {"n": 0} + real = timestep_phase.timestep_to_float + + def spy(value): + reads["n"] += 1 + return real(value) + + monkeypatch.setattr(timestep_phase, "timestep_to_float", spy) + backend = _cutedsl_backend() + backend.disabled_until_timestep = 0.9 + + # The 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({"sparse_attn_phase": 0}): + assert backend._dense_by_step(torch.tensor(0.1)) is True + with resolved_extra_keys_scope({"sparse_attn_phase": 1}): + assert backend._dense_by_step(torch.tensor(0.99)) is False + assert reads["n"] == 0, "timestep tensor was read despite a runner-resolved phase" + + # Outside a runner-driven call the tensor is the only source of truth. + assert backend._dense_by_step(torch.tensor(0.99)) is True + assert reads["n"] == 1 + + +@_REQUIRES_CUDA +def test_cutedsl_sol_dense_paths_use_cutedsl_dense_kernel(monkeypatch) -> None: + """Every path SOL cannot serve must reach the configured dense kernel. + + SOL runs dense attention on the ``dense_layers`` guard, the + ``disabled_until_timestep`` prefix and the kernel-ineligibility fallback. If + those ran torch SDPA instead of the CuTe DSL FMHA, a ``backend: CUTEDSL`` + run would differ from a ``backend: CUTEDSL`` dense baseline on those steps + and an A/B against that baseline would measure a backend swap rather than + sparsity. + """ + + if not sol_backend._cute_dense_available(): + pytest.skip("no CuTe DSL dense kernel for this device") + + device = torch.device("cuda") + q = k = v = torch.randn(1, 64, 2, 128, device=device, dtype=torch.bfloat16) + + def _make() -> tuple[SOLCuTeDSLAttention, dict]: + backend = _cutedsl_backend() + calls = {"n": 0} + real = backend._inner.forward + + def _spy(*args, **kwargs): + calls["n"] += 1 + return real(*args, **kwargs) + + monkeypatch.setattr(backend._inner, "forward", _spy) + return backend, calls + + backend, calls = _make() + backend.disabled_until_timestep = 0.9 + backend.forward(q, k, v, timestep=torch.tensor(0.95)) + assert calls["n"] == 1, "dense prefix was not delegated to the CuTeDSL dense kernel" + + backend, calls = _make() + backend.dense_layers = frozenset({backend.layer_idx}) + backend.forward(q, k, v) + assert calls["n"] == 1, "dense_layers guard was not delegated to the CuTeDSL dense kernel" + + backend, calls = _make() + monkeypatch.setattr( + sol_backend, + "_sol_attn_run", + lambda *args, **kw: kw["dense_fn"](*args[:3]), + ) + backend.forward(q, k, v) + assert calls["n"] == 1, "dense_fn did not route the fallback to the CuTeDSL dense kernel" + + +@_CPU_ONLY +@pytest.mark.parametrize( + ("make", "expect"), + [ + (lambda: torch.randn(1, 4, 2, 128), "not a CUDA tensor"), + (lambda: _fake_cuda_q((1, 4, 128)), "must be 4-D"), + (lambda: _fake_cuda_q((1, 4, 2, 64)), "head_dim must be 128"), + (lambda: _fake_cuda_q((1, 4, 2, 128), dtype=torch.float16), "dtype must be bfloat16"), + ], + ids=["cpu-tensor", "wrong-rank", "wrong-head-dim", "wrong-dtype"], +) +def test_cutedsl_ineligible_reason_is_reported(make, expect) -> None: + """Ineligibility must name the specific reason, never fail silently.""" + + wrapper = _kernel_wrapper() + reason = wrapper.sol_attn_ineligible_reason(make()) + assert reason is not None and expect in reason + assert not wrapper.sol_attn_supported(make()) + + +@_CPU_ONLY +def test_cutedsl_strict_raises_on_ineligible_input(monkeypatch) -> None: + """TRTLLM_SOL_ATTN_STRICT=1 turns an unservable input into an error.""" + + wrapper = _kernel_wrapper() + monkeypatch.setenv("TRTLLM_SOL_ATTN_STRICT", "1") + q = k = v = torch.randn(1, 4, 2, 128) # CPU -> ineligible + with pytest.raises(RuntimeError, match="cannot run the CuTe kernel"): + wrapper._run_sol_attn_bthd(q, k, v) + + +@_CPU_ONLY +def test_cutedsl_ineligible_falls_back_to_dense_and_counts(monkeypatch) -> None: + """Without STRICT the same input degrades to dense and increments the counter.""" + + wrapper = _kernel_wrapper() + monkeypatch.delenv("TRTLLM_SOL_ATTN_STRICT", raising=False) + wrapper.reset_sol_attn_stats() + q = k = v = torch.randn(1, 4, 2, 128) + + out = wrapper._run_sol_attn_bthd(q, k, v) + + ref = torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) + ).transpose(1, 2) + assert torch.allclose(out, ref), "dense fallback must be plain SDPA" + assert wrapper.get_sol_attn_stats()["dense_fallback_calls"] == 1 + assert wrapper.get_sol_attn_stats()["kernel_calls"] == 0 + + +@_CPU_ONLY +def test_cutedsl_supported_archs_match_kernel_dispatch_map() -> None: + """SUPPORTED_ARCHS mirrors the kernel dispatch map; drift would reject a served arch.""" + + from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.sol_attn import interface + + assert _kernel_wrapper().SUPPORTED_ARCHS == frozenset( + major * 10 + minor for major, minor in interface._CUTE_BACKENDS + ) + + +@_CPU_ONLY +def test_cutedsl_datacenter_blackwell_archs_are_supported() -> None: + """Both datacenter Blackwell steppings dispatch to the vendored kernel.""" + + from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.sol_attn import interface + + # Resolve, don't just look up: a malformed value would satisfy a key check + # and then fail at dispatch. + for arch in ((10, 0), (10, 3)): + kernel = interface._backend_for_arch(arch, cute_available=True) + assert kernel in interface._CUTE_BACKENDS.values() + assert isinstance(kernel, str) and kernel + + +@_CPU_ONLY +def test_cutedsl_sol_key_padding_mask_routes_to_vanilla_and_is_honored(monkeypatch) -> None: + """Masked self-attention must never reach the mask-blind sparse kernel. + + Equal Q/K lengths do not mean unmasked: HunyuanVideo 1.5 and GLM-Image pass + a ``[B, S]`` ``key_padding_mask`` on self-attention. The sparse kernel takes + no mask and CuTeDSL's dense ``_fwd(**kwargs)`` swallows it, so the only + correct destination is VANILLA. The output is checked against a masked SDPA + reference so a mask that was routed but then dropped still fails. + """ + + monkeypatch.setattr( + sol_backend, + "_sol_attn_run", + lambda *a, **k: pytest.fail("sparse kernel ran on a masked call"), + ) + backend = _cutedsl_backend() + calls = {"vanilla": 0} + real = backend._vanilla.forward + + def spy(*a, **k): + calls["vanilla"] += 1 + return real(*a, **k) + + monkeypatch.setattr(backend._vanilla, "forward", spy) + + torch.manual_seed(0) + q, k, v = (torch.randn(2, 64, 2, 128, dtype=torch.float32) for _ in range(3)) + mask = torch.ones(2, 64, dtype=torch.bool) + mask[0, 48:] = False # pad the tail of sample 0 + mask[1, :16] = False # pad the head of sample 1 + + assert not backend._can_serve(q, k, key_padding_mask=mask) + out = backend.forward(q, k, v, key_padding_mask=mask) + assert calls["vanilla"] == 1 + ref = _masked_sdpa_reference(q, k, v, key_padding_mask=mask) + torch.testing.assert_close(out.float(), ref, rtol=1e-4, atol=1e-4) + # And the mask actually mattered: unmasked attention gives a different answer. + assert (ref - _masked_sdpa_reference(q, k, v)).abs().max() > 1e-3 + + +@_CPU_ONLY +def test_cutedsl_sol_causal_mask_routes_to_dense_and_is_honored(monkeypatch) -> None: + """CAUSAL disqualifies the noncausal sparse kernel and is honored downstream.""" + + monkeypatch.setattr( + sol_backend, + "_sol_attn_run", + lambda *a, **k: pytest.fail("sparse kernel ran on a causal call"), + ) + backend = _cutedsl_backend() + torch.manual_seed(0) + q, k, v = (torch.randn(1, 64, 2, 128, dtype=torch.float32) for _ in range(3)) + + assert not backend._can_serve(q, k, attention_mask=PredefinedAttentionMask.CAUSAL) + out = backend.forward(q, k, v, attention_mask=PredefinedAttentionMask.CAUSAL) + ref = _masked_sdpa_reference(q, k, v, is_causal=True) + torch.testing.assert_close(out.float(), ref, rtol=1e-4, atol=1e-4) + assert (ref - _masked_sdpa_reference(q, k, v)).abs().max() > 1e-3 + + +@_CPU_ONLY +def test_cutedsl_sol_unmasked_self_attention_is_served() -> None: + """The mask routing must not touch the measured path: no mask, sparse.""" + + backend = _cutedsl_backend() + q = k = torch.randn(1, 64, 2, 128, dtype=torch.bfloat16) + assert backend._can_serve(q, k) + assert backend._can_serve( + q, k, attention_mask=PredefinedAttentionMask.FULL, key_padding_mask=None + ) + + +# --------------------------------------------------------------------------- public config +@_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_defaults_to_no_dense_prefix() -> None: + """``None``, not ``0.0``, means no dense prefix; ``0.0`` would run dense on every step.""" + + assert SolAttentionConfig(tau=2.0).disabled_until_timestep is None + + +@_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 +@pytest.mark.parametrize( + ("backend", "quant_recipe"), + [ + ("TRTLLM", {"qk_dtype": "fp8", "q_block_size": 1, "k_block_size": 1, "v_block_size": 1}), + ("CUTEDSL", {"qk_dtype": "bf16", "v_dtype": "fp8"}), + ], + ids=["TRTLLM", "CUTEDSL"], +) +def test_sol_and_attention_quantization_are_mutually_exclusive( + backend: str, quant_recipe: dict +) -> None: + """A recipe the backend accepts on its own is still rejected next to SOL.""" + + AttentionConfig(backend=backend, quant_attention_config=QuantAttentionConfig(**quant_recipe)) + with pytest.raises(ValidationError, match="SOL and quant_attention_config"): + AttentionConfig( + backend=backend, + quant_attention_config=QuantAttentionConfig(**quant_recipe), + sparse_attention_config=SolAttentionConfig(), + ) + + +@_CPU_ONLY +@pytest.mark.parametrize( + "layers", + [[-1], [0, -2], ["x"], "0,2", [1.5]], + ids=["negative", "negative_in_list", "non_numeric", "string_spec", "float"], +) +def test_sol_dense_layers_rejects_invalid_values(layers) -> None: + """``dense_layers`` is a list of non-negative layer indices; anything else fails at + config time rather than at attention construction.""" + + with pytest.raises(ValueError): + SolAttentionConfig(tau=2.0, dense_layers=layers) + + +@_CPU_ONLY +@pytest.mark.parametrize( + ("layers", "expected"), + [(None, None), ([0], [0]), ([0, 2, 4], [0, 2, 4]), ([4, 2, 2, 0], [0, 2, 4])], + ids=["none", "single", "list", "unsorted_with_duplicates"], +) +def test_sol_dense_layers_normalizes_valid_values(layers, expected) -> None: + assert SolAttentionConfig(tau=2.0, dense_layers=layers).dense_layers == expected + + +@_CPU_ONLY +def test_sol_config_has_no_kv_splits_knob() -> None: + """Only one KV split exists on the shipped kernels, so the config rejects the knob.""" + + assert "kv_splits" not in SolAttentionConfig.model_fields + with pytest.raises(ValueError): + SolAttentionConfig(tau=2.0, kv_splits="1") + + +# --------------------------------------------------------------------------- attention module +def _sol_model_config(*, backend: str = "TRTLLM", cp_size: int = 1) -> DiffusionModelConfig: + config = DiffusionModelConfig( + pretrained_config=SimpleNamespace(), + attention=AttentionConfig( + backend=backend, + 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 +@_BACKENDS +def test_sol_attention_rejects_context_parallelism(backend: str) -> None: + """SOL needs the full sequence per rank, so Ring/Attention2D fail at construction.""" + + 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(backend=backend, cp_size=2), + ) + + +# --------------------------------------------------------------------------- CUDA Graph phase keys +def _sol_model(*, disabled_until_timestep=None, dense_layers=None) -> BaseDiffusionModel: + """Minimal BaseDiffusionModel carrying a SOL sparse config.""" + + config = DiffusionModelConfig( + pretrained_config=SimpleNamespace( + hidden_size=64, num_attention_heads=4, attention_head_dim=16, eps=1e-6 + ), + attention=AttentionConfig( + backend="CUTEDSL", + sparse_attention_config=SolAttentionConfig( + tau=2.0, + disabled_until_timestep=disabled_until_timestep, + dense_layers=dense_layers, + ), + ), + skip_create_weights_in_init=False, + ) + return BaseDiffusionModel(config) + + +def _graph_runner() -> CUDAGraphRunner: + return CUDAGraphRunner(CUDAGraphRunnerConfig(use_cuda_graph=True)) + + +@_CPU_ONLY +def test_sol_cuda_graph_phase_is_keyed_without_model_scope(monkeypatch) -> None: + q = _bshd() + backend, predictor = _stub_backend(monkeypatch, SolParams(tau=1.0, disabled_until_timestep=0.6)) + model = _SolModel((backend,)) + runner = _graph_runner() + 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) + + # The stubbed core echoes its q: fused QKV in the dense prefix, Q alone in + # the sparse phase. + assert model(q, timestep=torch.tensor(0.8)).shape == (1, 64, 3 * 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 + + +@_CPU_ONLY +def test_sol_cuda_graph_key_separates_dense_prefix_from_sparse_phase() -> None: + """The prefix swaps kernels without changing any tensor shape, so a graph + captured in the dense prefix must not be replayed for the sparse phase.""" + + model = _sol_model(disabled_until_timestep=0.9) + runner = _graph_runner() + model.register_cuda_graph_extra_key_fns(runner) + + base = {"hidden_states": torch.empty(1, 8, 64)} + key_dense = runner.get_graph_key(**base, timestep=torch.empty(1).fill_(0.95)) + key_sparse = runner.get_graph_key(**base, timestep=torch.empty(1).fill_(0.10)) + + assert key_dense != key_sparse, ( + "dense-prefix and sparse phases share a CUDA graph key despite running " + "different kernels; a graph captured in one phase would be replayed in " + "the other" + ) + + +@_CPU_ONLY +def test_sol_cuda_graph_key_unregistered_without_prefix() -> None: + """``dense_layers`` alone is fixed per layer, so it needs no graph key.""" + + model = _sol_model(disabled_until_timestep=None, dense_layers=[0, 2]) + runner = _graph_runner() + model.register_cuda_graph_extra_key_fns(runner) + + base = {"hidden_states": torch.empty(1, 8, 64)} + key_a = runner.get_graph_key(**base, timestep=torch.empty(1).fill_(0.95)) + key_b = runner.get_graph_key(**base, timestep=torch.empty(1).fill_(0.10)) + assert key_a == key_b, "no phase key should be registered without a dense prefix" + + +# --------------------------------------------------------------------------- CUDA Graph capture +def _base_runner() -> CUDAGraphRunner: + return CUDAGraphRunner(CUDAGraphRunnerConfig(use_cuda_graph=True)) + + +def _ltx2_runner(): + from tensorrt_llm._torch.visual_gen.models.ltx2.pipeline_ltx2 import _LTX2CUDAGraphRunner + + return _LTX2CUDAGraphRunner(CUDAGraphRunnerConfig(use_cuda_graph=True)) + + +def _ltx2_two_stage_runner(): + from tensorrt_llm._torch.visual_gen.models.ltx2.pipeline_ltx2_two_stages import ( + _LTX2TwoStageCUDAGraphRunner, + ) + + return _LTX2TwoStageCUDAGraphRunner( + CUDAGraphRunnerConfig(use_cuda_graph=True), lambda: "original", lambda: "default" + ) + + +# Every runner that overrides ``capture`` must keep the phase-scope contract; +# LTX-2's runners re-implement capture for ``Modality`` inputs. +_RUNNER_FACTORIES = pytest.mark.parametrize( + "make_runner", + [_base_runner, _ltx2_runner, _ltx2_two_stage_runner], + ids=["base", "ltx2", "ltx2_two_stage"], +) + + +@_REQUIRES_CUDA +@_RUNNER_FACTORIES +def test_cuda_graph_runner_publishes_resolved_extra_keys_during_capture(make_runner) -> None: + """The runner exposes its host-resolved extra keys to the captured callee.""" + + seen = [] + + def fn(x): + seen.append(resolved_extra_key("probe")) + return x * 2 + + runner = make_runner() + runner.register_extra_key_fn("probe", lambda *args, **kwargs: 7) + wrapped = runner.wrap(fn) + x = torch.ones(4, device="cuda") + out = wrapped(x) + torch.cuda.synchronize() + assert torch.equal(out, x * 2) + assert seen and all(v == 7 for v in seen), seen # warmup + capture passes + assert resolved_extra_key("probe") is None, "scope must close after capture" + + +@_REQUIRES_CUDA +@_RUNNER_FACTORIES +def test_cutedsl_sol_dense_prefix_survives_cuda_graph_capture(make_runner) -> None: + """Capture and replay a graph on each side of the cutoff without a sync. + + The dense-prefix decision must be baked into each captured graph via the + runner's phase key, and no ``.item()`` may run inside capture. + ``kernel_calls`` proves the phase-0 graph never launched the sparse kernel + and the phase-1 graph did. + """ + + wrapper = _kernel_wrapper() + if not wrapper.sol_attn_supported( + torch.empty(1, 8, 8, 128, device="cuda", dtype=torch.bfloat16) + ): + pytest.skip("no SOL kernel for this device") + + cutoff = 0.9 + backend = _cutedsl_backend() + backend.disabled_until_timestep = cutoff + runner = make_runner() + runner.register_extra_key_fn( + "sparse_attn_phase", + lambda *args, **kwargs: graph_phase_for_timestep( + kwargs.get("timestep"), disabled_until_timestep=cutoff + ), + ) + + def fwd(q, k, v, *, timestep): + return backend.forward(q, k, v, timestep=timestep) + + wrapped = runner.wrap(fwd) + torch.manual_seed(0) + q = k = v = torch.randn(1, 64, 2, 128, device="cuda", dtype=torch.bfloat16) + t_dense = torch.tensor(0.95, device="cuda") + t_sparse = torch.tensor(0.5, device="cuda") + + before = wrapper._SOL_STATS["kernel_calls"] + wrapped(q, k, v, timestep=t_dense) + torch.cuda.synchronize() + after_dense = wrapper._SOL_STATS["kernel_calls"] + wrapped(q, k, v, timestep=t_sparse) + torch.cuda.synchronize() + after_sparse = wrapper._SOL_STATS["kernel_calls"] + + assert len(runner.graphs) == 2, "one graph per phase" + assert after_dense == before, "dense prefix must not launch the sparse kernel" + assert after_sparse > after_dense, "sparse phase must launch the kernel" + # Replays must not error either. + wrapped(q, k, v, timestep=t_dense) + wrapped(q, k, v, timestep=t_sparse) + torch.cuda.synchronize() diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_parity.py b/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_parity.py new file mode 100644 index 000000000000..34c240070148 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_parity.py @@ -0,0 +1,603 @@ +# 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. + +"""Numerical parity of the VisualGen SOL backends. + +The TRTLLM backend (Triton predictor plus the PrimTS block-sparse FMHA) and the +CuTeDSL backend (vendored fused kernel) are checked against dense attention, +against the exact-token/proxy-summary contract of the routes the predictor +emits, against an independent fp32 re-derivation of the SOL routing and +block-mean approximation, and against each other. Every test here needs an +SM100-family GPU. +""" + +from __future__ import annotations + +import math + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.fmha.prims_ts_block_sparse import PrimsTSBlockSparseFmha +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 ( + SOLCuTeDSLAttention, + SOLTrtllmAttention, +) +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.predictor import ( + SolPredictorOutputs, +) +from tensorrt_llm._torch.visual_gen.attention_backend.utils import create_attention +from tensorrt_llm._torch.visual_gen.config import create_attention_metadata_state +from tensorrt_llm._torch.visual_gen.cuda_graph_runner import resolved_extra_keys_scope +from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell import sol_attn_backend +from tensorrt_llm.visual_gen import SolAttentionConfig +from tensorrt_llm.visual_gen.args import AttentionConfig + +_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", +) +_BLOCK = 64 +_HEAD_DIM = 128 +_THRESH_TYPES = pytest.mark.parametrize("thresh_type", ["diag", "exact"]) + + +# --------------------------------------------------------------------------- references +def _dense_reference(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + scores = torch.einsum("bqhd,bkhd->bhqk", q.float(), k.float()) * (_HEAD_DIM**-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 of a predicted route.""" + + num_blocks = math.ceil(k.shape[1] / _BLOCK) + 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)): + q_begin = q_block_idx * _BLOCK + q_end = min(q_begin + _BLOCK, 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, + min((block_idx + 1) * _BLOCK, 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, k.shape[1] - block_idx * _BLOCK) + 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 + + +def _sol_reference(q, k, v, *, tau, scale, block=_BLOCK): + """fp32 reference of the SOL ``diag`` routing and block-mean approximation. + + Mirrors the vendored kernel (``preprocess.py``, ``common/selector.py``, + ``sm100/mainloop.py``): + + * ``kc`` is the per-block mean of K, ``vc`` the per-block sum of V. + * Per query block the threshold is ``mean + tau * std`` in the log2 domain, + where mean and variance come from the query centroid against the global + K mean and the per-channel K variance. + * A KV block is exact for a query block when the column-mean route score + exceeds that threshold or when it lies within one block of the diagonal. + * Non-exact blocks contribute ``exp(q . kc) * vc`` to the numerator and + ``exp(q . kc) * block_len`` to the denominator: every key in the block is + treated as its mean. + + Returns the output and the exact-routing mask ``[B, nb_q, nb_kv, H]``. + """ + + q, k, v = (t.float() for t in (q, k, v)) + B, S, H, D = q.shape + nb = -(-S // block) + log2e = math.log2(math.e) + log2_scale = scale * log2e + dev = q.device + idx = torch.arange(S, device=dev) // block + blen = torch.bincount(idx, minlength=nb).float() + + def _block_sum(t): + return torch.zeros(B, nb, H, D, device=dev).index_add_(1, idx, t) + + kc = _block_sum(k) / blen[None, :, None, None] + vc = _block_sum(v) + qc = _block_sum(q) / blen[None, :, None, None] + kmean = k.mean(dim=1) + kvar = ((k * k).mean(dim=1) - kmean * kmean).clamp(min=0) + mean = (qc * kmean[:, None]).sum(-1) * log2_scale + var = (qc * qc * kvar[:, None]).sum(-1) * log2_scale**2 + thr = mean + tau * torch.sqrt(var.clamp(min=0) + 1e-6) # [B, nb_q, H] + + s2 = torch.einsum("bshd,bjhd->bsjh", q, kc) * log2_scale # [B, S, nb_kv, H] + col_mean = ( + torch.zeros(B, nb, nb, H, device=dev).index_add_(1, idx, s2) / blen[None, :, None, None] + ) + ar = torch.arange(nb, device=dev) + near_diag = (ar[:, None] - ar[None, :]).abs() <= 1 + exact = (col_mean > thr[:, :, None, :]) | near_diag[None, :, :, None] # [B, nb_q, nb_kv, H] + + exact_rows = exact[:, idx] # [B, S, nb_kv, H] + exact_pair = exact_rows[:, :, idx, :] # [B, S, T, H] + scores = torch.einsum("bshd,bthd->bsth", q, k) * scale + approx = s2 / log2e + neg = torch.finfo(torch.float32).min + m = torch.maximum( + scores.masked_fill(~exact_pair, neg).amax(2), + approx.masked_fill(exact_rows, neg).amax(2), + ) + pe = torch.exp(scores - m[:, :, None]) * exact_pair + pa = torch.exp(approx - m[:, :, None]) * (~exact_rows) + num = torch.einsum("bsth,bthd->bshd", pe, v) + torch.einsum("bsjh,bjhd->bshd", pa, vc) + den = pe.sum(2) + (pa * blen[None, None, :, None]).sum(2) + return num / den[..., None], exact + + +# --------------------------------------------------------------------------- backends +def _sol_config(*, tau: float, thresh_type: str = "diag", **kwargs) -> SolAttentionConfig: + return SolAttentionConfig(tau=tau, thresh_type=thresh_type, **kwargs) + + +def _trtllm_backend(sparse_config: SolAttentionConfig, *, num_heads: int) -> SOLTrtllmAttention: + attention_config = AttentionConfig(backend="TRTLLM", sparse_attention_config=sparse_config) + backend = create_attention( + backend="TRTLLM", + layer_idx=1, + num_heads=num_heads, + head_dim=_HEAD_DIM, + dtype=torch.bfloat16, + attention_config=attention_config, + attention_metadata_state=create_attention_metadata_state(), + ) + assert isinstance(backend, SOLTrtllmAttention) + assert any(isinstance(fmha, PrimsTSBlockSparseFmha) for fmha in backend._fmha_manager.fmha_libs) + return backend + + +def _cutedsl_backend(sparse_config: SolAttentionConfig, *, num_heads: int) -> SOLCuTeDSLAttention: + if not sol_attn_backend.sol_attn_supported( + torch.empty(1, 8, num_heads, _HEAD_DIM, device="cuda", dtype=torch.bfloat16) + ): + pytest.skip("no SOL CuTeDSL kernel for this device") + attention_config = AttentionConfig(backend="CUTEDSL", sparse_attention_config=sparse_config) + backend = create_attention( + backend="CUTEDSL", + layer_idx=1, + num_heads=num_heads, + head_dim=_HEAD_DIM, + dtype=torch.bfloat16, + attention_config=attention_config, + ) + assert isinstance(backend, SOLCuTeDSLAttention) + return backend + + +def _forward_trtllm(backend: SOLTrtllmAttention, q, k, v, *, timestep=None) -> torch.Tensor: + batch_size, seq_len = q.shape[:2] + return backend.forward( + q=q, + k=k, + v=v, + batch_size=batch_size, + seq_len=seq_len, + seq_len_kv=seq_len, + timestep=timestep, + ).view_as(q) + + +def _forward_cutedsl(backend: SOLCuTeDSLAttention, q, k, v) -> torch.Tensor: + before = sol_attn_backend._SOL_STATS["kernel_calls"] + out = backend.forward(q, k, v) + torch.cuda.synchronize() + assert sol_attn_backend._SOL_STATS["kernel_calls"] == before + 1, ( + "the CuTeDSL SOL kernel did not run; this would compare dense against dense" + ) + return out + + +def _spy_primts_forward(monkeypatch) -> dict: + """Record every PrimTS block-sparse FMHA launch; a dense fallback would leave it at zero.""" + + calls = {"n": 0} + original = PrimsTSBlockSparseFmha.forward + + def counted(*args, **kwargs): + calls["n"] += 1 + return original(*args, **kwargs) + + monkeypatch.setattr(PrimsTSBlockSparseFmha, "forward", counted) + return calls + + +def _packed_strided_inputs(generator, shape): + """Q/K/V as non-contiguous split views of one packed projection, like the module.""" + + 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)) + + +def _gaussian_inputs(seed: int, shape): + generator = torch.Generator(device="cuda").manual_seed(seed) + return tuple( + torch.randn(shape, generator=generator, device="cuda", dtype=torch.bfloat16) + for _ in range(3) + ) + + +# --------------------------------------------------------------------------- TRTLLM backend +@_REQUIRES_SM100 +@torch.no_grad() +def test_trtllm_sol_all_exact_routes_match_dense_under_cuda_graph() -> None: + """A threshold no block can miss routes everything exact, so SOL equals dense attention.""" + + backend = _trtllm_backend(_sol_config(tau=-1.0e6, disabled_until_timestep=0.6), num_heads=2) + generator = torch.Generator(device="cuda").manual_seed(20260901) + shape = (1, 257, 2, _HEAD_DIM) + + q, k, v = _packed_strided_inputs(generator, shape) + timestep = torch.tensor(0.2, device="cuda") + assert not any(tensor.is_contiguous() for tensor in (q, k, v)) + eager = _forward_trtllm(backend, q, k, v, timestep=timestep) + torch.cuda.synchronize() + torch.testing.assert_close(eager, _dense_reference(q, k, v), rtol=2e-2, atol=2e-2) + + # Capture sees the sparse phase the CUDA graph runner resolved for the call. + graph = torch.cuda.CUDAGraph() + with resolved_extra_keys_scope({"sparse_attn_phase": 1}), torch.cuda.graph(graph): + captured = _forward_trtllm(backend, q, k, v, timestep=timestep) + + next_q, next_k, next_v = _packed_strided_inputs(generator, shape) + 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) + + +@_REQUIRES_SM100 +@torch.no_grad() +@pytest.mark.parametrize( + ("config_kwargs", "dense_timestep"), + ( + ({"dense_layers": [1]}, None), + ({"disabled_until_timestep": 0.6}, 0.9), + ), + ids=("dense_layer", "dense_phase"), +) +def test_trtllm_sol_dense_calls_run_the_fused_dense_kernel( + monkeypatch, config_kwargs, dense_timestep +) -> None: + """A dense layer and a dense-phase step match dense attention without a + block-sparse launch; a later sparse-phase step of the same backend still + routes through the block-sparse FMHA.""" + + backend = _trtllm_backend(_sol_config(tau=1.0, **config_kwargs), num_heads=2) + calls = _spy_primts_forward(monkeypatch) + q, k, v = _gaussian_inputs(20260922, (2, 320, 2, _HEAD_DIM)) + timestep = None if dense_timestep is None else torch.tensor(dense_timestep, device="cuda") + + dense = _forward_trtllm(backend, q, k, v, timestep=timestep) + torch.cuda.synchronize() + + torch.testing.assert_close(dense, _dense_reference(q, k, v), rtol=2e-2, atol=2e-2) + assert calls["n"] == 0 + if dense_timestep is not None: + _forward_trtllm(backend, q, k, v, timestep=torch.tensor(0.2, device="cuda")) + torch.cuda.synchronize() + assert calls["n"] == 1 + + +@_REQUIRES_SM100 +@torch.no_grad() +@pytest.mark.parametrize("seq_len", [256, 257]) +def test_trtllm_sol_mixed_proxy_routes_match_reference_under_cuda_graph(seq_len: int) -> None: + """A threshold most blocks miss exercises the proxy path; the output follows the + exact-token/proxy-summary contract of the predicted route, also after replay.""" + + backend = _trtllm_backend(_sol_config(tau=1.0e6), num_heads=2) + generator = torch.Generator(device="cuda").manual_seed(20260903) + shape = (1, seq_len, 2, _HEAD_DIM) + + q, k, v = _packed_strided_inputs(generator, shape) + eager = _forward_trtllm(backend, q, k, v) + predictor_outputs = sol_predictor.predict( + q.contiguous(), + k.contiguous(), + v.contiguous(), + tau=1.0e6, + sm_scale=_HEAD_DIM**-0.5, + ) + torch.cuda.synchronize() + exact_bits = predictor_outputs.exact_block_bits + num_blocks = math.ceil(seq_len / _BLOCK) + 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, _mixed_proxy_reference(q, k, v, predictor_outputs), rtol=2e-2, atol=2e-2 + ) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = _forward_trtllm(backend, q, k, v) + + next_q, next_k, next_v = _packed_strided_inputs(generator, shape) + q.copy_(next_q) + k.copy_(next_k) + v.copy_(next_v) + 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=_HEAD_DIM**-0.5, + ) + torch.testing.assert_close( + captured, _mixed_proxy_reference(q, k, v, replayed_outputs), rtol=2e-2, atol=2e-2 + ) + + +# --------------------------------------------------------------------------- CuTeDSL kernel +@_REQUIRES_SM100 +def test_cutedsl_sol_kernel_matches_dense_on_a_single_block() -> None: + """With ``tokens <= BLOCK`` there is exactly one KV block, so nothing can be routed + away whatever ``tau`` says and the kernel must reproduce dense attention. + + ``kernel_calls`` is asserted to have advanced: without that, a fallback to + dense would make this pass by comparing dense against itself. + """ + + if not sol_attn_backend.sol_attn_supported( + torch.empty(1, 8, 8, _HEAD_DIM, device="cuda", dtype=torch.bfloat16) + ): + pytest.skip("no SOL CuTeDSL kernel for this device") + + torch.manual_seed(0) + shape = (1, _BLOCK, 4, _HEAD_DIM) + q, k, v = (torch.randn(shape, device="cuda", dtype=torch.bfloat16) for _ in range(3)) + reference = torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2).float(), k.transpose(1, 2).float(), v.transpose(1, 2).float() + ).transpose(1, 2) + + before = sol_attn_backend._SOL_STATS["kernel_calls"] + out = sol_attn_backend._run_sol_attn_bthd(q, k, v, tau=2.0, thresh_type="diag", kv_splits=1) + torch.cuda.synchronize() + assert sol_attn_backend._SOL_STATS["kernel_calls"] == before + 1, ( + "the kernel did not run; this comparison would be dense against dense" + ) + + torch.testing.assert_close(out.float(), reference, rtol=2e-2, atol=2e-2) + + +@_REQUIRES_SM100 +def test_cutedsl_sol_kernel_matches_reference_with_approximated_blocks() -> None: + """The kernel agrees with the fp32 reference where blocks are approximated. + + Plain Gaussian inputs are the right stimulus: every KV block carries + comparable softmax mass and no block's column-mean score clears the + ``mean + tau*std`` threshold, so every block two or more away from the + diagonal is approximated, and because keys vary within a block, treating 64 + keys as their mean is measurably different from dense attention. The test + asserts that some blocks were approximated and that the result differs + from dense, so neither an accidental all-exact run nor a mass-free + approximation can pass. + """ + + if not sol_attn_backend.sol_attn_supported( + torch.empty(1, 8, 8, _HEAD_DIM, device="cuda", dtype=torch.bfloat16) + ): + pytest.skip("no SOL CuTeDSL kernel for this device") + + torch.manual_seed(0) + shape = (1, 256, 2, _HEAD_DIM) # 4 KV blocks + q, k, v = (torch.randn(shape, device="cuda", dtype=torch.bfloat16) for _ in range(3)) + scale = _HEAD_DIM**-0.5 + tau = 2.0 + + ref, exact = _sol_reference(q, k, v, tau=tau, scale=scale) + assert (~exact).any(), "inputs did not force any block onto the approximation path" + dense = torch.nn.functional.scaled_dot_product_attention( + q.transpose(1, 2).float(), k.transpose(1, 2).float(), v.transpose(1, 2).float() + ).transpose(1, 2) + # Calibrated on B200: max |ref - dense| = 0.45 for this stimulus. A tight + # margin here would let a mass-free approximation pass vacuously. + assert (ref - dense).abs().max() > 1e-1, "approximation is not measurably different from dense" + + before = sol_attn_backend._SOL_STATS["kernel_calls"] + out = sol_attn_backend._run_sol_attn_bthd(q, k, v, tau=tau, thresh_type="diag", kv_splits=1) + torch.cuda.synchronize() + assert sol_attn_backend._SOL_STATS["kernel_calls"] == before + 1, "kernel did not run" + + # Calibrated on B200: max |kernel - ref| = 1.7e-3 (mean 2e-4); 5e-3 leaves + # ~3x headroom for bf16 rounding while still catching a routing or + # approximation-formula mismatch, which shows up at 1e-1 or worse. + torch.testing.assert_close(out.float(), ref, rtol=5e-3, atol=5e-3) + + +# --------------------------------------------------------------------------- TRTLLM vs CuTeDSL vs dense +@_REQUIRES_SM100 +@_THRESH_TYPES +@torch.no_grad() +def test_sol_backends_match_dense_when_every_block_is_exact(monkeypatch, thresh_type) -> None: + """Both backends reproduce dense attention when the threshold routes every block exact.""" + + num_heads = 2 + q, k, v = _gaussian_inputs(20260921, (2, 320, num_heads, _HEAD_DIM)) + dense = _dense_reference(q, k, v) + primts_calls = _spy_primts_forward(monkeypatch) + config = _sol_config(tau=-1.0e6, thresh_type=thresh_type) + + trtllm = _forward_trtllm(_trtllm_backend(config, num_heads=num_heads), q, k, v) + cutedsl = _forward_cutedsl(_cutedsl_backend(config, num_heads=num_heads), q, k, v) + torch.cuda.synchronize() + + assert primts_calls["n"] == 1 + torch.testing.assert_close(trtllm, dense, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(cutedsl, dense, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(trtllm, cutedsl, rtol=2e-2, atol=2e-2) + + +def _predicted_routes(q, k, v, *, tau: float, thresh_type: str) -> SolPredictorOutputs: + outputs = sol_predictor.predict( + q.contiguous(), + k.contiguous(), + v.contiguous(), + tau=tau, + sm_scale=_HEAD_DIM**-0.5, + thresh_type=thresh_type, + ) + torch.cuda.synchronize() + return outputs + + +def _exact_block_ratio(outputs: SolPredictorOutputs, num_blocks: int) -> float: + bits = outputs.exact_block_bits.to(torch.int64) + exact = sum( + int(((bits[..., block // 32] >> (block % 32)) & 1).sum().item()) + for block in range(num_blocks) + ) + return exact / (math.prod(bits.shape[:3]) * num_blocks) + + +@_REQUIRES_SM100 +@_THRESH_TYPES +# From a threshold few blocks reach, through one inside the score distribution of this +# stimulus, to one most blocks clear. +@pytest.mark.parametrize("tau", [2.0, 1.0, 0.0]) +@torch.no_grad() +def test_sol_backends_follow_the_same_routes(monkeypatch, thresh_type, tau) -> None: + """Both backends execute the routes the predictor emits, at every threshold. + + The oracle is the exact-token/proxy-summary contract evaluated on the predicted + routes, so a backend that routed a borderline block differently from the + predictor fails with a localized error even though a global similarity to + dense attention would barely move. A threshold sitting inside the score + distribution (``tau=1`` here) is exactly the case an independent fp32 + re-derivation of the routing cannot arbitrate, and this test covers it. + """ + + num_heads, num_blocks = 2, 5 + q, k, v = _gaussian_inputs(20260922, (2, num_blocks * _BLOCK, num_heads, _HEAD_DIM)) + dense = _dense_reference(q, k, v).float() + routes = _predicted_routes(q, k, v, tau=tau, thresh_type=thresh_type) + ratio = _exact_block_ratio(routes, num_blocks) + assert 0.0 < ratio < 1.0, f"tau={tau} left no block on the proxy path ({ratio:.2f} exact)" + reference = _mixed_proxy_reference(q, k, v, routes).float() + primts_calls = _spy_primts_forward(monkeypatch) + config = _sol_config(tau=tau, thresh_type=thresh_type) + + trtllm = _forward_trtllm(_trtllm_backend(config, num_heads=num_heads), q, k, v).float() + cutedsl = _forward_cutedsl(_cutedsl_backend(config, num_heads=num_heads), q, k, v).float() + torch.cuda.synchronize() + + assert primts_calls["n"] == 1 + torch.testing.assert_close(trtllm, cutedsl, rtol=1e-2, atol=1e-2) + torch.testing.assert_close(trtllm, reference, rtol=2e-2, atol=2e-2) + torch.testing.assert_close(cutedsl, reference, rtol=2e-2, atol=2e-2) + # The approximation is active: the shared output is measurably not dense attention. + assert (trtllm - dense).abs().max() > 1e-1 + + +@_REQUIRES_SM100 +@torch.no_grad() +def test_sol_backends_match_independent_reference_away_from_the_threshold(monkeypatch) -> None: + """Both backends reproduce an independent fp32 re-derivation of the ``diag`` routing + and the block-mean approximation. + + ``tau=2`` on Gaussian inputs leaves only the near-diagonal blocks exact, far from + the threshold. Right at the threshold the fp32 re-derivation and the bf16 kernels + can legitimately route a block differently, so borderline thresholds are covered + by ``test_sol_backends_follow_the_same_routes`` against the predicted routes. + """ + + num_heads = 2 + q, k, v = _gaussian_inputs(20260922, (2, 320, num_heads, _HEAD_DIM)) + reference, exact = _sol_reference(q, k, v, tau=2.0, scale=_HEAD_DIM**-0.5) + assert (~exact).any() + dense = _dense_reference(q, k, v).float() + assert (reference - dense).abs().max() > 1e-1 + primts_calls = _spy_primts_forward(monkeypatch) + config = _sol_config(tau=2.0, thresh_type="diag") + + trtllm = _forward_trtllm(_trtllm_backend(config, num_heads=num_heads), q, k, v).float() + cutedsl = _forward_cutedsl(_cutedsl_backend(config, num_heads=num_heads), q, k, v).float() + torch.cuda.synchronize() + + assert primts_calls["n"] == 1 + torch.testing.assert_close(trtllm, reference, rtol=1e-2, atol=1e-2) + torch.testing.assert_close(cutedsl, reference, rtol=1e-2, atol=1e-2) diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor.py b/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor.py new file mode 100644 index 000000000000..2b5d51a7fef9 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor.py @@ -0,0 +1,305 @@ +# 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 math + +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 ( + BLOCK_SIZE, + SolPredictorOutputs, + 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_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"), + ) + 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) + + +def _summary_oracle(k: torch.Tensor, v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + 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, + device=k.device, + ) + v_summary = torch.empty_like(k_summary) + for block_idx in range(blocks): + 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 + + +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 _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, + v: torch.Tensor, + *, + tau: float, + sm_scale: float, + thresh_type: str = "diag", +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + k_summary, v_summary = _summary_oracle(k, v) + q_centroids = _centroid_oracle(q) + threshold = _threshold_oracle( + q_centroids, k_summary, tau=tau, sm_scale=sm_scale, thresh_type=thresh_type + ) + scores = torch.einsum("bqhd,bkhd->bhqk", q_centroids, k_summary.double()) * ( + float(sm_scale) * _LOG2_E + ) + 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 + + +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 _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) + + +@_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 +@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) + 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), + ) + + +@_REQUIRES_CUDA +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) + + +@_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) + 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 = 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 = 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 +@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) + 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 = 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) + k.copy_(next_k) + v.copy_(next_v) + graph.replay() + torch.cuda.synchronize() + + _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_operator_is_functional_and_compiles_fullgraph() -> None: + op = torch.ops.trtllm.visual_gen_sol_predictor.default + schema = str(op._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, 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) + 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 + ) + + 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 + ) + + 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/sol/test_sol_predictor_kernels.py b/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor_kernels.py new file mode 100644 index 000000000000..330b23a2fd05 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/sparse_attention/sol/test_sol_predictor_kernels.py @@ -0,0 +1,251 @@ +# 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 thresholds, exact-block selection).""" + +from __future__ import annotations + +import pytest +import torch + +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.kernels import ( + _LOG2_E, + _block_pool_torch, + _block_thresholds_torch, + _select_exact_blocks_torch, + block_pool, + block_thresholds, + 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_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 + 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, threshold, bits, sm_scale=1.0) + exact = _unpack(bits, blocks) + # 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) + # 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])) + + +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]) +# 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_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_size, 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 +@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 +@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 + ) + 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, 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 + ) + 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)) + 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) + 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") + 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() + 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_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_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) + assert torch.equal(bits, expected_bits) diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/vsa/__init__.py b/tests/unittest/_torch/visual_gen/sparse_attention/vsa/__init__.py new file mode 100644 index 000000000000..52a7a9daf028 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/sparse_attention/vsa/__init__.py @@ -0,0 +1,2 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_attention.py b/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_attention.py new file mode 100644 index 000000000000..ba7195246c58 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_attention.py @@ -0,0 +1,685 @@ +# 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. +"""Behavior tests for the VisualGen VSA backends. + +Both the CuTeDSL and the TRTLLM backend share the VisualGen-owned predictor and +post-processing; this module covers dispatch, the predictor envelope, metadata +caching and the CUDA Graph contract. Numerical parity of the two backends with +dense attention lives in ``test_vsa_parity.py``, and the kernels are covered by +``test_vsa_kernels.py``. +""" + +from types import SimpleNamespace + +import pytest +import torch + +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 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, +) +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 _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_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 not inputs.post_context.fine_is_tiled + + +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) + + 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) + + +@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 _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() -> None: + """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.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.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_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"), + } + 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) + routes = [ + route_builder.from_selected_blocks( + 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") +@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: 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 + 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, x_pooled = tile_and_pool_cubes( + x, + meta.tile_source_index, + meta.variable_block_sizes, + cube_size=VSA_BLOCK_SIZE, + ) + + pad_mask = meta.tile_source_index < 0 + if pad_mask.any(): + assert x_tiled[:, pad_mask, :, :].abs().max().item() == 0.0, ( + "tiling must zero-fill padded positions" + ) + + 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}" + ) + 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}" + ) + + 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) diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_kernels.py b/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_kernels.py new file mode 100644 index 000000000000..657beb19b670 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_kernels.py @@ -0,0 +1,411 @@ +# 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 kernels against references. + +The predictor kernels (tile + cube mean, row sort, coarse/fine blend) are +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. The CuTe +DSL block-sparse fine kernel is checked against dense and masked SDPA +references; the TRTLLM fine stage runs the PrimTS block-sparse FMHA covered by +``tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py``. +""" + +import pytest +import torch +import torch.nn.functional as F + +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 []) +_REQUIRES_CUDA = pytest.mark.skipif(not torch.cuda.is_available(), reason="kernel test needs CUDA") + +_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"), +] + + +# --------------------------------------------------------------------------- predictor kernels +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) + + +@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) + + +@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) + + +# --------------------------------------------------------------------------- CuTe DSL fine kernel +def _cute_fine_kernel(q: torch.Tensor): + """Return the CuTe block-sparse fine kernel or skip when this device cannot run it.""" + + 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") + if not is_cute_supported(q): + pytest.skip("CuTe path needs sm_100+ Blackwell (current device unsupported)") + return block_sparse_attn_from_indices_cute + + +@_REQUIRES_CUDA +def test_cute_kernel_matches_dense_at_full_topk() -> None: + """CuTe block-sparse kernel matches dense SDPA when every cube is selected.""" + 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) + kernel = _cute_fine_kernel(q) + + 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 = kernel(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})" + ) + + +@_REQUIRES_CUDA +def test_cute_kernel_matches_ref_with_independent_indices() -> None: + """CuTe kernel: paired Q-blocks (2i, 2i+1) attend to independent KV index lists.""" + 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) + kernel = _cute_fine_kernel(q) + + 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 = kernel(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})" + ) + + +@_REQUIRES_CUDA +def test_cute_kernel_50pct_sparsity_quality_vs_dense() -> None: + """50% sparse CuTe kernel with score-based topk stays close to dense SDPA.""" + 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) + kernel = _cute_fine_kernel(q) + + 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 = kernel(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" + ) + + +@_REQUIRES_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: int) -> None: + """CuTe kernel supports a final Q block that has no paired neighbor.""" + 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) + kernel = _cute_fine_kernel(q) + + 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 = kernel(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/sparse_attention/vsa/test_vsa_parity.py b/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_parity.py new file mode 100644 index 000000000000..0fe6b494a212 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_parity.py @@ -0,0 +1,240 @@ +# 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. +"""Numerical parity of the VisualGen VSA backends. + +A Wan-style self-attention module (fused QKV projection, QK RMSNorm, RoPE) is run +with the CuTeDSL and TRTLLM VSA backends on a ragged latent and compared with +dense attention through the module's own projections, and with each other. +""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F +from utils.util import isSM100Family + +from tensorrt_llm._torch.attention.backends.fmha.prims_ts_block_sparse import PrimsTSBlockSparseFmha +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import ( + VSAMetadataBuilder, + set_vsa_forward_context, +) +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa.backend import VSACuTeDSLAttention +from tensorrt_llm._torch.visual_gen.config import ( + DiffusionModelConfig, + create_attention_metadata_state, +) +from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode, apply_rotary_emb +from tensorrt_llm.visual_gen.args import AttentionConfig, VideoSparseAttentionConfig + +_REQUIRES_CUDA = pytest.mark.skipif(not torch.cuda.is_available(), reason="VSA needs CUDA") +_REQUIRES_SM100 = pytest.mark.skipif( + not isSM100Family(), + reason="CuTe DSL and PrimTS block-sparse parity requires SM100 or SM103", +) +_BACKENDS = pytest.mark.parametrize("backend", ["CUTEDSL", "TRTLLM"]) + +# A ragged latent exercises VSA padding and token-mask lowering on both fine stages. +_LATENT_SHAPE = (9, 9, 9) +_SEQ_LEN = _LATENT_SHAPE[0] * _LATENT_SHAPE[1] * _LATENT_SHAPE[2] +_NUM_HEADS = 4 +_HEAD_DIM = 128 +_HIDDEN_SIZE = _NUM_HEADS * _HEAD_DIM + + +def _make_attention(backend: str, sparsity: float, *, seed: int) -> Attention: + """A Wan self-attention module on CUDA whose weights come from ``seed``.""" + + config = DiffusionModelConfig( + pretrained_config=SimpleNamespace( + hidden_size=_HIDDEN_SIZE, + num_attention_heads=_NUM_HEADS, + attention_head_dim=_HEAD_DIM, + eps=1e-6, + ), + attention=AttentionConfig( + backend=backend, + sparse_attention_config=VideoSparseAttentionConfig(vsa_sparsity=sparsity), + ), + skip_create_weights_in_init=False, + ) + config.attention_metadata_state = ( + create_attention_metadata_state() if backend == "TRTLLM" else None + ) + attention = Attention(_HIDDEN_SIZE, _NUM_HEADS, qkv_mode=QKVMode.FUSE_QKV, config=config) + attention = attention.to(device=torch.device("cuda"), dtype=torch.bfloat16).eval() + # Fail loudly if the VSA path silently fell back to the VANILLA backend. + assert attention.attn_backend == backend, ( + f"Expected {backend} VSA backend, got {attention.attn_backend!r}" + ) + generator = torch.Generator(device="cuda").manual_seed(seed) + for name, parameter in sorted(attention.named_parameters()): + values = torch.randn(parameter.shape, generator=generator, device="cuda") + if "norm" in name: + values = 1.0 + 0.05 * values + else: + values = 0.05 * values + parameter.data.copy_(values.to(parameter.dtype)) + return attention + + +def _vsa_metadata(sparsity: float): + return VSAMetadataBuilder().build( + current_timestep=0, + raw_latent_shape=_LATENT_SHAPE, + patch_size=(1, 1, 1), + vsa_sparsity=sparsity, + device=torch.device("cuda"), + ) + + +def _rope(seq_len: int, head_dim: int) -> tuple[torch.Tensor, torch.Tensor]: + """RoPE tables with the full head_dim in the ``[1, S, 1, D]`` layout the module uses.""" + + device = torch.device("cuda") + position = torch.arange(seq_len, device=device).unsqueeze(1) + div_term = torch.exp( + torch.arange(0, head_dim, device=device) * (-torch.log(torch.tensor(10000.0)) / head_dim) + ) + freqs_cos = torch.cos(position * div_term).unsqueeze(0).unsqueeze(2) + freqs_sin = torch.sin(position * div_term).unsqueeze(0).unsqueeze(2) + return freqs_cos, freqs_sin + + +def _dense_reference(attention: Attention, hidden_states: torch.Tensor, freqs) -> torch.Tensor: + """Dense self-attention through the module's own projections, norms and RoPE.""" + + batch_size, seq_len = hidden_states.shape[:2] + q, k, v = attention.get_qkv(hidden_states, None) + q, k = attention.apply_qk_norm(q, k) + q = apply_rotary_emb(q.view(batch_size, seq_len, _NUM_HEADS, _HEAD_DIM), *freqs) + k = apply_rotary_emb(k.view(batch_size, seq_len, _NUM_HEADS, _HEAD_DIM), *freqs) + v = v.view(batch_size, seq_len, _NUM_HEADS, _HEAD_DIM) + out = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)) + return attention.to_out[0](out.transpose(1, 2).flatten(2)) + + +def _run(attention: Attention, hidden_states: torch.Tensor, freqs, sparsity: float) -> torch.Tensor: + """One VSA forward with the coarse gate closed, so the output is the fine branch only.""" + + with torch.no_grad(), set_vsa_forward_context(_vsa_metadata(sparsity)): + return attention(hidden_states, freqs=freqs, gate_compress=torch.zeros_like(hidden_states)) + + +def _inputs(seed: int, batch_size: int) -> torch.Tensor: + generator = torch.Generator(device="cuda").manual_seed(seed) + return torch.randn( + batch_size, _SEQ_LEN, _HIDDEN_SIZE, generator=generator, device="cuda", dtype=torch.bfloat16 + ) + + +@_REQUIRES_CUDA +@_BACKENDS +def test_vsa_backends_match_dense_at_sparsity_zero(backend: str) -> None: + """VSA at sparsity 0 with the coarse gate closed selects every cube, so the fine branch + is dense attention and must match the naive SDPA reference modulo bf16 rounding.""" + + attention = _make_attention(backend, sparsity=0.0, seed=42) + hidden_states = _inputs(seed=42, batch_size=2) + freqs = _rope(_SEQ_LEN, _HEAD_DIM) + + out_vsa = _run(attention, hidden_states, freqs, sparsity=0.0) + with torch.no_grad(): + out_dense = _dense_reference(attention, hidden_states, freqs) + + assert out_vsa.shape == out_dense.shape + torch.testing.assert_close(out_vsa, out_dense, rtol=1e-2, atol=1e-2) + + +@_REQUIRES_SM100 +def test_vsa_backends_match_each_other_on_ragged_input(monkeypatch: pytest.MonkeyPatch) -> None: + """CuTeDSL and TRTLLM implement the same sparse VSA fine-stage semantics. + + Both modules carry the same weights, both must run their sparse fine stage + (the CuTe kernel and the PrimTS block-sparse FMHA respectively), and both + outputs stay in the same neighbourhood of dense attention. + """ + + sparsity = 0.5 + attentions = { + backend: _make_attention(backend, sparsity=sparsity, seed=0) + for backend in ("CUTEDSL", "TRTLLM") + } + hidden_states = _inputs(seed=0, batch_size=1) + freqs = _rope(_SEQ_LEN, _HEAD_DIM) + + sparse_fine_executed = {"CUTEDSL": 0, "TRTLLM": 0} + cute_execute = VSACuTeDSLAttention._execute_sparse_fine + primts_forward = PrimsTSBlockSparseFmha.forward + + def counted_cute(self, inputs): + sparse_fine_executed["CUTEDSL"] += 1 + return cute_execute(self, inputs) + + def counted_primts(*args, **kwargs): + sparse_fine_executed["TRTLLM"] += 1 + return primts_forward(*args, **kwargs) + + monkeypatch.setattr(VSACuTeDSLAttention, "_execute_sparse_fine", counted_cute) + monkeypatch.setattr(PrimsTSBlockSparseFmha, "forward", counted_primts) + + outputs = { + backend: _run(attention, hidden_states, freqs, sparsity) + for backend, attention in attentions.items() + } + with torch.no_grad(): + dense = _dense_reference(attentions["TRTLLM"], hidden_states, freqs) + + assert sparse_fine_executed == {"CUTEDSL": 1, "TRTLLM": 1} + for backend, output in outputs.items(): + assert torch.isfinite(output).all(), backend + # Random weights spread attention mass evenly, so dropping half of the cubes + # moves the output well away from dense (cosine similarity 0.85 measured on + # B200); the bound only guards against a broken fine stage. + cos_sim = F.cosine_similarity( + output.float().reshape(-1), dense.float().reshape(-1), dim=0 + ).item() + assert cos_sim >= 0.65, f"{backend} VSA at 50% sparsity drifted from dense: {cos_sim:.4f}" + torch.testing.assert_close(outputs["CUTEDSL"], outputs["TRTLLM"], rtol=1e-2, atol=1e-2) + + +@_REQUIRES_SM100 +def test_vsa_trtllm_cuda_graph_replays_live_routes() -> None: + """Captured VSA recomputes routes when graph-stable input storage changes.""" + + attention = _make_attention("TRTLLM", sparsity=0.5, seed=17) + freqs = _rope(_SEQ_LEN, _HEAD_DIM) + static_hidden = _inputs(seed=17, batch_size=1) + static_gate = torch.zeros_like(static_hidden) + metadata = _vsa_metadata(0.5) + + for _ in range(2): + with torch.no_grad(), set_vsa_forward_context(metadata): + attention(static_hidden, freqs=freqs, gate_compress=static_gate) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph), torch.no_grad(), set_vsa_forward_context(metadata): + graph_output = attention(static_hidden, freqs=freqs, 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(metadata): + eager_output = attention(live_hidden, freqs=freqs, 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) 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 deleted file mode 100644 index 0da640d23d1e..000000000000 --- a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_sol_attn.py +++ /dev/null @@ -1,1032 +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. -"""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 -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 -VSA's full-top-k does, but a single KV block makes sparsity structurally -impossible and gives the same guarantee. -""" - -import math -from types import SimpleNamespace - -import pytest -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.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 -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, - resolved_extra_key, - resolved_extra_keys_scope, -) -from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode -from tensorrt_llm.visual_gen.args import AttentionConfig, SolAttentionConfig - - -def test_cute_dsl_factory_dispatches_dense_and_sol_attn() -> None: - dense_config = AttentionConfig(backend="CUTEDSL") - dense_attention = create_attention( - backend="CUTEDSL", - layer_idx=0, - num_heads=8, - head_dim=128, - attention_config=dense_config, - ) - - sparse_config = SolAttentionConfig(tau=2.0, disabled_until_timestep=0.9545) - sol_attn_config = AttentionConfig(backend="CUTEDSL", sparse_attention_config=sparse_config) - sol_attn_attention = create_attention( - backend="CUTEDSL", - layer_idx=0, - num_heads=8, - head_dim=128, - attention_config=sol_attn_config, - ) - - assert isinstance(dense_attention, CuTeDSLAttention) - assert isinstance(sol_attn_attention, SolAttention) - assert sol_attn_attention.tau == 2.0 - assert sol_attn_attention.disabled_until_timestep == 0.9545 - - -def _make_config( - hidden_size: int, - num_heads: int, - head_dim: int, - backend: str, - sol_attn_tau: "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 = ( - SolAttentionConfig(tau=sol_attn_tau) if sol_attn_tau 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="Sol-Attn needs CUDA") -def test_sol_attn_cross_attention_delegates_to_dense_cutedsl(): - """Cross-attention is delegated to the dense backend, decided per call. - - Sol-Attn is self-attention only. It is still the module's backend for a - cross-attention module -- `create_attention` has no rule excluding it -- and - delegates at `forward` because `_can_serve` sees `k.shape[1] != q.shape[1]`. - - Deciding this per call rather than at construction is the point: `qkv_mode` - describes how Q/K/V are projected, not whether K/V come from another - sequence, so a construction-time rule keyed on SEPARATE_QKV silently - stripped the configured backend from self-attention modules that use that - mode for unrelated reasons (Qwen-Image; WAN attn1 under async Ulysses). - """ - device = torch.device("cuda") - dtype = torch.bfloat16 - cfg = _make_config( - hidden_size=64, num_heads=4, head_dim=16, backend="CUTEDSL", sol_attn_tau=1.0 - ) - cross_attn = ( - Attention(64, 4, qkv_mode=QKVMode.SEPARATE_QKV, config=cfg) - .to(device=device, dtype=dtype) - .eval() - ) - assert cross_attn.attn_backend == "CUTEDSL", ( - f"expected CUTEDSL, got {cross_attn.attn_backend!r}" - ) - assert isinstance(cross_attn.attn, SolAttention), ( - "Sol-Attn should remain the backend and delegate per call, not be " - f"swapped out at construction; got {type(cross_attn.attn).__name__}" - ) - # q and k with different sequence lengths -> not self-attention -> delegate - q = torch.randn(1, 32, 4, 16, device=device, dtype=dtype) - k = torch.randn(1, 77, 4, 16, device=device, dtype=dtype) - assert not cross_attn.attn._can_serve(q, k), ( - "differing q/k sequence lengths must be delegated, not routed to the sparse kernel" - ) - - -def test_sol_attn_self_attention_is_served_under_separate_qkv(): - """SEPARATE_QKV self-attention keeps Sol-Attn -- the async-Ulysses case. - - WAN's attn1 switches to SEPARATE_QKV when async Ulysses is active, and - 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.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. - q = k = torch.randn(1, 64, 2, 128, dtype=torch.bfloat16) - assert attn._can_serve(q, k), "equal q/k sequence lengths must reach the sparse kernel" - - -def test_sol_attn_with_context_parallelism_raises(): - """Sol-Attn + Attention2D/Ring must error at construction (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=SolAttentionConfig(tau=1.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_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) - - -@pytest.mark.parametrize( - "timestep,expected", - [ - (0.99, 0), # early/noisy -> dense prefix - (0.9545, 0), # exactly at the cutoff -> still dense - (0.95, 1), # past the cutoff -> sparse - (0.0, 1), # final step -> sparse - (None, None), # no timestep -> no phase to distinguish - ], - ids=["early", "at-cutoff", "past-cutoff", "final", "missing"], -) -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. - """ - assert ( - SkipSoftmaxScheduler.get_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 - ) - - -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 - ) - - -def test_dense_prefix_skips_kernel(monkeypatch): - """Inside the dense prefix the sparse kernel must not be invoked at all. - - CPU tensors, so `_dense` takes its SDPA branch here; that the dense path - 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 - - 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.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)) - assert out.shape == q.shape - assert torch.isfinite(out).all() - - -def test_missing_timestep_fails_open_to_sparse(monkeypatch): - """Without a timestep the prefix cannot be applied; run sparse, do not raise. - - 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 - - called = {"n": 0} - - def _record(*args, **kwargs): - called["n"] += 1 - return args[0] - - monkeypatch.setattr(sol_attn_mod, "_sol_attn_run", _record) - - attn = SolAttention(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 - assert called["n"] == 1, "expected the sparse kernel, not a silent dense fallback" - - -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 - - 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.dense_layers = frozenset({3}) - q = k = v = torch.randn(1, 4, 2, 16) - out = attn.forward(q, k, v) - assert out.shape == q.shape - assert torch.isfinite(out).all() - - -def _make_solattn_model(disabled_until_timestep=None, dense_layers=None): - """Minimal BaseDiffusionModel carrying a Sol-Attn sparse config.""" - from tensorrt_llm._torch.visual_gen.models.modeling import BaseDiffusionModel - - pretrained_config = SimpleNamespace( - hidden_size=64, num_attention_heads=4, attention_head_dim=16, eps=1e-6 - ) - config = DiffusionModelConfig( - pretrained_config=pretrained_config, - attention=AttentionConfig( - backend="CUTEDSL", - sparse_attention_config=SolAttentionConfig( - tau=2.0, - disabled_until_timestep=disabled_until_timestep, - dense_layers=dense_layers, - ), - ), - skip_create_weights_in_init=False, - ) - return BaseDiffusionModel(config) - - -def _graph_runner(): - from tensorrt_llm._torch.visual_gen.cuda_graph_runner import ( - CUDAGraphRunner, - CUDAGraphRunnerConfig, - ) - - return CUDAGraphRunner(CUDAGraphRunnerConfig(use_cuda_graph=True)) - - -def test_cuda_graph_key_separates_dense_prefix_from_sparse_phase(): - """The prefix swaps kernels without changing any tensor shape, so a graph - captured in the dense prefix must not be replayed for the sparse phase.""" - model = _make_solattn_model(disabled_until_timestep=0.9) - runner = _graph_runner() - model.register_cuda_graph_extra_key_fns(runner) - - base = {"hidden_states": torch.empty(1, 8, 64)} - key_dense = runner.get_graph_key(**base, timestep=torch.empty(1).fill_(0.95)) - key_sparse = runner.get_graph_key(**base, timestep=torch.empty(1).fill_(0.10)) - - assert key_dense != key_sparse, ( - "dense-prefix and sparse phases share a CUDA graph key despite running " - "different kernels; a graph captured in one phase would be replayed in " - "the other" - ) - - -def test_cuda_graph_key_unregistered_without_prefix(): - """dense_layers alone is fixed per layer, so it needs no graph key.""" - model = _make_solattn_model(disabled_until_timestep=None, dense_layers=[0, 2]) - runner = _graph_runner() - model.register_cuda_graph_extra_key_fns(runner) - - base = {"hidden_states": torch.empty(1, 8, 64)} - key_a = runner.get_graph_key(**base, timestep=torch.empty(1).fill_(0.95)) - key_b = runner.get_graph_key(**base, timestep=torch.empty(1).fill_(0.10)) - assert key_a == key_b, "no phase key should be registered without a dense prefix" - - -# --- kernel-wrapper eligibility / strictness (sol_attn_backend.py) ----------- -# These run on CPU: every path here is pure Python guard logic, and the CPU -# tensor is itself one of the ineligible cases. - - -def _backend_mod(): - from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell import sol_attn_backend - - return sol_attn_backend - - -def _fake_cuda_q(shape, dtype=torch.bfloat16): - """A tensor-like that reports `is_cuda=True` without needing a GPU. - - `sol_attn_ineligible_reason` checks `is_cuda` first and returns early, so a - CPU tensor can never reach the rank, head_dim, or dtype branches. These - stubs pass that first gate so each later reason is actually exercised; the - architecture check comes after them and is not reached. - """ - return SimpleNamespace(is_cuda=True, ndim=len(shape), shape=shape, dtype=dtype) - - -@pytest.mark.parametrize( - "make,expect", - [ - (lambda: torch.randn(1, 4, 2, 128), "not a CUDA tensor"), - (lambda: _fake_cuda_q((1, 4, 128)), "must be 4-D"), - (lambda: _fake_cuda_q((1, 4, 2, 64)), "head_dim must be 128"), - (lambda: _fake_cuda_q((1, 4, 2, 128), dtype=torch.float16), "dtype must be bfloat16"), - ], - ids=["cpu-tensor", "wrong-rank", "wrong-head-dim", "wrong-dtype"], -) -def test_ineligible_reason_is_reported(make, expect): - """Ineligibility must name the specific reason, never fail silently. - - Each case is built to fail exactly one check, in the order the function - applies them, so the id names the branch that actually fires. - """ - reason = _backend_mod().sol_attn_ineligible_reason(make()) - assert reason is not None and expect in reason - assert not _backend_mod().sol_attn_supported(make()) - - -def test_strict_raises_on_ineligible_input(monkeypatch): - """TRTLLM_SOL_ATTN_STRICT=1 must turn an unservable input into an error. - Without it an unsupported arch degrades to dense silently and the counters - the PR relies on cannot be trusted.""" - sab = _backend_mod() - monkeypatch.setenv("TRTLLM_SOL_ATTN_STRICT", "1") - q = k = v = torch.randn(1, 4, 2, 128) # CPU -> ineligible - with pytest.raises(RuntimeError, match="cannot run the CuTe kernel"): - sab._run_sol_attn_bthd(q, k, v) - - -def test_ineligible_falls_back_to_dense_and_counts(monkeypatch): - """Without STRICT the same input degrades to dense and increments the counter.""" - sab = _backend_mod() - monkeypatch.delenv("TRTLLM_SOL_ATTN_STRICT", raising=False) - sab.reset_sol_attn_stats() - q = k = v = torch.randn(1, 4, 2, 128) - out = sab._run_sol_attn_bthd(q, k, v) - ref = torch.nn.functional.scaled_dot_product_attention( - q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) - ).transpose(1, 2) - assert torch.allclose(out, ref), "dense fallback must be plain SDPA" - assert sab.get_sol_attn_stats()["dense_fallback_calls"] == 1 - assert sab.get_sol_attn_stats()["kernel_calls"] == 0 - - -def test_supported_archs_matches_kernel_dispatch_map(): - """SUPPORTED_ARCHS is a hand-copy of interface.py's _CUTE_BACKENDS. If they - drift, eligibility silently rejects an arch the kernel actually supports.""" - from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.sol_attn import interface - - assert _backend_mod().SUPPORTED_ARCHS == frozenset( - major * 10 + minor for major, minor in interface._CUTE_BACKENDS - ) - - -def test_datacenter_blackwell_archs_are_supported(): - """Both datacenter Blackwell steppings dispatch to the vendored kernel. - - sm103 (B300/GB300) runs the same `sm100/` kernel as sm100: the CuTe DSL JIT - targets the device it compiles on and the kernel body uses no sm100-only - construct. Verified on hardware against a B200 control; see the kernel's - THIRD_PARTY_NOTICES.md. - """ - from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.sol_attn import interface - - # Resolve, don't just look up: a malformed value would satisfy a key check - # and then fail at dispatch. - for arch in ((10, 0), (10, 3)): - backend = interface._backend_for_arch(arch, cute_available=True) - assert backend in interface._CUTE_BACKENDS.values() - assert isinstance(backend, str) and backend - - -def test_no_arch_literal_outside_the_dispatch_map(): - """`_sol_attn_cute`'s guard must key off _CUTE_BACKENDS, not a literal. - - That guard is a second, deeper check than `_backend_for_arch`, and it is - the one no other test reaches: `test_supported_archs_matches_kernel_dispatch_map` - keeps SUPPORTED_ARCHS and _CUTE_BACKENDS in step, so widening both leaves a - hardcoded literal here as the only thing still rejecting the new arch -- - with every test green. Assert on the source so the coupling cannot regress. - """ - import ast - import inspect - import textwrap - - from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.sol_attn import interface - - src = textwrap.dedent(inspect.getsource(interface._sol_attn_cute)) - assert "arch not in _CUTE_BACKENDS" in src, ( - "the guard in _sol_attn_cute must be keyed off _CUTE_BACKENDS" - ) - - # Reject *any* architecture tuple literal, not just (10, 0): pinning one - # value would let a future `arch != (10, 3)` reintroduce the same trap. - offenders = [] - for node in ast.walk(ast.parse(src)): - if not isinstance(node, ast.Compare): - continue - operands = [node.left, *node.comparators] - names = {n.id for n in operands if isinstance(n, ast.Name)} - if "arch" not in names: - continue - for operand in operands: - if isinstance(operand, ast.Tuple) and all( - isinstance(e, ast.Constant) and isinstance(e.value, int) for e in operand.elts - ): - offenders.append(ast.unparse(node)) - assert not offenders, ( - "hardcoded architecture literal(s) compared against `arch` in " - f"_sol_attn_cute: {offenders}; key the guard off _CUTE_BACKENDS instead" - ) - - -def test_quant_attention_config_rejected_with_sol_attn(): - """Sol-Attn replaces the dense CuTeDSL path, so quantized attention cannot - compose with it; accepting the pair would silently ignore the quant request.""" - from tensorrt_llm.visual_gen.args import QuantAttentionConfig - - with pytest.raises(ValueError, match="mutually exclusive"): - AttentionConfig( - backend="CUTEDSL", - quant_attention_config=QuantAttentionConfig(), - sparse_attention_config=SolAttentionConfig(tau=2.0), - ) - - -def test_zero_cutoff_rejected(): - """0.0 is the natural thing to type for 'no prefix', but it would run dense - on every step and turn Sol-Attn off entirely. Must be rejected, not silent.""" - with pytest.raises(ValueError): - SolAttentionConfig(tau=2.0, disabled_until_timestep=0.0) - assert SolAttentionConfig(tau=2.0).disabled_until_timestep is None - - -@pytest.mark.parametrize( - "layers", - [[-1], [0, -2], ["x"], "0,2", [1.5]], - ids=["negative", "negative_in_list", "non_numeric", "string_spec", "float"], -) -def test_dense_layers_rejects_invalid_values(layers): - """dense_layers is a list of non-negative layer indices; anything else fails at - config time rather than at attention construction.""" - with pytest.raises(ValueError): - SolAttentionConfig(tau=2.0, dense_layers=layers) - - -@pytest.mark.parametrize( - "layers,expected", - [(None, None), ([0], [0]), ([0, 2, 4], [0, 2, 4]), ([4, 2, 2, 0], [0, 2, 4])], - ids=["none", "single", "list", "unsorted_with_duplicates"], -) -def test_dense_layers_normalizes_valid_values(layers, expected): - assert SolAttentionConfig(tau=2.0, dense_layers=layers).dense_layers == expected - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Sol-Attn needs CUDA") -def test_cute_kernel_matches_dense_on_a_single_block(): - """Kernel-level numerical check against dense attention. - - The obvious form of this test -- the analogue of VSA's - `test_cute_kernel_matches_dense_at_full_topk` -- needs a tau that provably - forces full non-sparse routing. Sol-Attn has no such value: its routing is - score-derived rather than a plain top-k, so tau=0 is not the dense limit. - - A single KV block sidesteps that entirely. With `tokens <= BLOCK_SIZE` - there is exactly one block, so there is nothing to route away and sparsity - is structurally impossible whatever tau says. The kernel must therefore - reproduce dense attention here, which makes this a genuine numerical test - of the kernel rather than of the routing policy. - - `kernel_calls` is asserted to have advanced: without that, a fallback to - dense would make this pass by comparing dense against itself. - """ - sab = _backend_mod() - if not sab.sol_attn_supported(torch.empty(1, 8, 8, 128, device="cuda", dtype=torch.bfloat16)): - pytest.skip("no Sol-Attn kernel for this device") - - tokens = 64 # == BLOCK_SIZE: exactly one KV block - torch.manual_seed(0) - shape = (1, tokens, 4, 128) - q, k, v = (torch.randn(shape, device="cuda", dtype=torch.bfloat16) for _ in range(3)) - - reference = torch.nn.functional.scaled_dot_product_attention( - q.transpose(1, 2).float(), k.transpose(1, 2).float(), v.transpose(1, 2).float() - ).transpose(1, 2) - - before = sab._SOL_STATS["kernel_calls"] - out = sab._run_sol_attn_bthd(q, k, v, tau=2.0, thresh_type="diag", kv_splits=1) - torch.cuda.synchronize() - assert sab._SOL_STATS["kernel_calls"] == before + 1, ( - "the kernel did not run; this comparison would be dense against dense" - ) - - torch.testing.assert_close(out.float(), reference, rtol=2e-2, atol=2e-2) - - -def test_kv_splits_is_not_a_config_knob(): - """Only one KV split exists on the shipped kernels, so the config exposes - no `kv_splits` and rejects it rather than accepting a no-op choice.""" - assert "kv_splits" not in SolAttentionConfig.model_fields - with pytest.raises(ValueError): - SolAttentionConfig(tau=2.0, kv_splits="1") - - -def _is_dynamo_disabled(fn) -> bool: - """True if `fn` is wrapped by torch.compiler.disable / torch._dynamo.disable.""" - target = getattr(fn, "__func__", fn) - return bool(getattr(target, "_torchdynamo_disable", False)) - - -def test_kernel_launch_is_opaque_to_dynamo(): - """The CuTe DSL launch boundary must be @torch.compiler.disable'd. - - Without it Dynamo traces into the CuTe DSL JIT builder and retraces on every - call: near two orders of magnitude slower on B200 (2496.9 s mean denoise - without it), and silent -- it looks like torch.compile simply not paying off. - """ - assert _is_dynamo_disabled(_backend_mod()._run_sol_attn_bthd), ( - "_run_sol_attn_bthd must be decorated with @torch.compiler.disable" - ) - - -def test_timestep_scalar_read_is_opaque_to_dynamo(): - """The dense-prefix `.item()` must stay in eager. - - 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" - ) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs a CUDA device") -def test_dense_paths_use_cutedsl_backend(monkeypatch): - """Every path Sol-Attn cannot serve must reach the configured dense kernel. - - Sol-Attn does dense attention on the `dense_layers` guard, the - `disabled_until_timestep` prefix, and kernel-ineligibility fallback. If those - call torch SDPA instead of `cute_dsl_fmha_fwd`, a `backend: CUTEDSL` run - differs from a `backend: CUTEDSL` dense baseline on those steps, and any A/B - against that baseline measures a backend swap rather than sparsity. Measured - at LPIPS 0.214 on Wan2.2-T2V-A14B before this was fixed, against a 0.25 gate. - - 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 - - if not sol_attn_mod._cute_dense_available(): - # A CUDA device is not enough: the premise here is that the dense paths - # reach `cute_dsl_fmha_fwd`, which only exists on sm100/sm103. Without - # this the test fails rather than skips on any other CUDA runner. - pytest.skip("no CuTe DSL dense kernel for this device") - - device = torch.device("cuda") - 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) - calls = {"n": 0} - real = a._inner.forward - - def _spy(*args, **kwargs): - calls["n"] += 1 - return real(*args, **kwargs) - - monkeypatch.setattr(a._inner, "forward", _spy) - return a, calls - - # 1. dense prefix: timestep at/above the cutoff - a, calls = _make() - a.disabled_until_timestep = 0.9 - a.dense_layers = frozenset() - a.forward(q, k, v, timestep=torch.tensor(0.95)) - assert calls["n"] == 1, "dense prefix was not delegated to the CuTeDSL dense kernel" - - # 2. dense_layers guard - a, calls = _make() - a.disabled_until_timestep = None - a.dense_layers = frozenset({0}) - a.forward(q, k, v) - assert calls["n"] == 1, "dense_layers guard was not delegated to the CuTeDSL dense kernel" - - # 3. ineligibility fallback, reached through `dense_fn` - a, calls = _make() - a.disabled_until_timestep = None - a.dense_layers = frozenset() - monkeypatch.setattr( - sol_attn_mod, - "_sol_attn_run", - lambda *args, **kw: kw["dense_fn"](*args[:3]), - ) - a.forward(q, k, v) - assert calls["n"] == 1, "dense_fn did not route the fallback to the CuTeDSL dense kernel" - - -# --------------------------------------------------------------------------- -# Dense-prefix phase under CUDA-graph capture -# --------------------------------------------------------------------------- - - -def test_dense_by_step_prefers_runner_resolved_phase(monkeypatch): - """Inside a runner-driven call the phase comes from the runner, not the tensor. - - CUDA-graph capture forbids device-to-host syncs, and reading the timestep - is a `.item()`. The runner resolves the phase host-side to build the graph - key and republishes it during capture; `_dense_by_step` must use that and - never touch the tensor. `torch.compiler.disable` does not help here -- it - only excludes Dynamo, not stream capture. - """ - reads = {"n": 0} - real = SkipSoftmaxScheduler._as_float - - def spy(value): - reads["n"] += 1 - return real(value) - - monkeypatch.setattr(SkipSoftmaxScheduler, "_as_float", staticmethod(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}): - assert attn._dense_by_step(torch.tensor(0.1)) is True - with resolved_extra_keys_scope({"sol_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" - - # Outside a runner-driven call the tensor is the only source of truth. - assert attn._dense_by_step(torch.tensor(0.99)) is True - assert reads["n"] == 1 - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA graphs need a GPU") -def _base_runner(): - return CUDAGraphRunner(CUDAGraphRunnerConfig(use_cuda_graph=True)) - - -def _ltx2_runner(): - from tensorrt_llm._torch.visual_gen.models.ltx2.pipeline_ltx2 import _LTX2CUDAGraphRunner - - return _LTX2CUDAGraphRunner(CUDAGraphRunnerConfig(use_cuda_graph=True)) - - -def _ltx2_two_stage_runner(): - from tensorrt_llm._torch.visual_gen.models.ltx2.pipeline_ltx2_two_stages import ( - _LTX2TwoStageCUDAGraphRunner, - ) - - return _LTX2TwoStageCUDAGraphRunner( - CUDAGraphRunnerConfig(use_cuda_graph=True), lambda: "original", lambda: "default" - ) - - -# Every runner that overrides ``capture`` must keep the phase-scope contract; -# LTX-2's runners re-implement capture for ``Modality`` inputs. -_RUNNER_FACTORIES = pytest.mark.parametrize( - "make_runner", - [_base_runner, _ltx2_runner, _ltx2_two_stage_runner], - ids=["base", "ltx2", "ltx2_two_stage"], -) - - -@_RUNNER_FACTORIES -def test_cuda_graph_runner_publishes_resolved_extra_keys_during_capture(make_runner): - """The runner exposes its host-resolved extra keys to the captured callee.""" - seen = [] - - def fn(x): - seen.append(resolved_extra_key("probe")) - return x * 2 - - runner = make_runner() - runner.register_extra_key_fn("probe", lambda *args, **kwargs: 7) - wrapped = runner.wrap(fn) - x = torch.ones(4, device="cuda") - out = wrapped(x) - torch.cuda.synchronize() - assert torch.equal(out, x * 2) - assert seen and all(v == 7 for v in seen), seen # warmup + capture passes - assert resolved_extra_key("probe") is None, "scope must close after capture" - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Sol-Attn needs CUDA") -@_RUNNER_FACTORIES -def test_sol_attn_dense_prefix_survives_cuda_graph_capture(make_runner): - """Capture and replay a graph on each side of the cutoff without a sync. - - The dense-prefix decision must be baked into each captured graph via the - runner's phase key, and no `.item()` may run inside capture. `kernel_calls` - proves the phase-0 graph never launched the sparse kernel and the phase-1 - graph did. - """ - sab = _backend_mod() - if not sab.sol_attn_supported(torch.empty(1, 8, 8, 128, device="cuda", dtype=torch.bfloat16)): - pytest.skip("no Sol-Attn kernel for this device") - - cutoff = 0.9 - attn = SolAttention(layer_idx=1, num_heads=2, head_dim=128) - 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( - kwargs.get("timestep"), disabled_until_timestep=cutoff - ), - ) - - def fwd(q, k, v, *, timestep): - return attn.forward(q, k, v, timestep=timestep) - - wrapped = runner.wrap(fwd) - torch.manual_seed(0) - q = k = v = torch.randn(1, 64, 2, 128, device="cuda", dtype=torch.bfloat16) - t_dense = torch.tensor(0.95, device="cuda") - t_sparse = torch.tensor(0.5, device="cuda") - - before = sab._SOL_STATS["kernel_calls"] - wrapped(q, k, v, timestep=t_dense) - torch.cuda.synchronize() - after_dense = sab._SOL_STATS["kernel_calls"] - wrapped(q, k, v, timestep=t_sparse) - torch.cuda.synchronize() - after_sparse = sab._SOL_STATS["kernel_calls"] - - assert len(runner.graphs) == 2, "one graph per phase" - assert after_dense == before, "dense prefix must not launch the sparse kernel" - assert after_sparse > after_dense, "sparse phase must launch the kernel" - # Replays must not error either. - wrapped(q, k, v, timestep=t_dense) - wrapped(q, k, v, timestep=t_sparse) - torch.cuda.synchronize() - - -# --------------------------------------------------------------------------- -# Kernel vs a PyTorch reference of the block-mean approximation -# --------------------------------------------------------------------------- - - -def _sol_attn_reference(q, k, v, *, tau, scale, block=64): - """PyTorch fp32 reference of Sol-Attn's routing + block-mean approximation. - - Mirrors the vendored kernel (`preprocess.py`, `common/selector.py`, - `sm100/mainloop.py`): - - * `kc` = per-block mean of K, `vc` = per-block *sum* of V. - * Per query block: threshold = mean + tau * std, in the log2 domain, where - mean/var come from the query centroid against the global K mean and - per-dim K variance (`_diag_threshold_kernel`). - * A KV block is *exact* for a query block if the column-mean route score - exceeds that threshold, or it is within one block of the diagonal - (`sol_attn_route_is_exact`). - * Non-exact blocks contribute `exp(q . kc) * vc` to the numerator and - `exp(q . kc) * block_len` to the denominator -- every key in the block - is treated as its mean. - - Returns the output and the exact-routing mask `[B, nb_q, nb_kv, H]`. - """ - q, k, v = (t.float() for t in (q, k, v)) - B, S, H, D = q.shape - nb = -(-S // block) - log2e = math.log2(math.e) - log2_scale = scale * log2e - dev = q.device - idx = torch.arange(S, device=dev) // block - blen = torch.bincount(idx, minlength=nb).float() - - def _block_sum(t): - return torch.zeros(B, nb, H, D, device=dev).index_add_(1, idx, t) - - kc = _block_sum(k) / blen[None, :, None, None] - vc = _block_sum(v) - qc = _block_sum(q) / blen[None, :, None, None] - kmean = k.mean(dim=1) - kvar = ((k * k).mean(dim=1) - kmean * kmean).clamp(min=0) - mean = (qc * kmean[:, None]).sum(-1) * log2_scale - var = (qc * qc * kvar[:, None]).sum(-1) * log2_scale**2 - thr = mean + tau * torch.sqrt(var.clamp(min=0) + 1e-6) # [B, nb_q, H] - - s2 = torch.einsum("bshd,bjhd->bsjh", q, kc) * log2_scale # [B, S, nb_kv, H] - col_mean = ( - torch.zeros(B, nb, nb, H, device=dev).index_add_(1, idx, s2) / blen[None, :, None, None] - ) - ar = torch.arange(nb, device=dev) - near_diag = (ar[:, None] - ar[None, :]).abs() <= 1 - exact = (col_mean > thr[:, :, None, :]) | near_diag[None, :, :, None] # [B, nb_q, nb_kv, H] - - exact_rows = exact[:, idx] # [B, S, nb_kv, H] - exact_pair = exact_rows[:, :, idx, :] # [B, S, T, H] - scores = torch.einsum("bshd,bthd->bsth", q, k) * scale - approx = s2 / log2e - neg = torch.finfo(torch.float32).min - m = torch.maximum( - scores.masked_fill(~exact_pair, neg).amax(2), - approx.masked_fill(exact_rows, neg).amax(2), - ) - pe = torch.exp(scores - m[:, :, None]) * exact_pair - pa = torch.exp(approx - m[:, :, None]) * (~exact_rows) - num = torch.einsum("bsth,bthd->bshd", pe, v) + torch.einsum("bsjh,bjhd->bshd", pa, vc) - den = pe.sum(2) + (pa * blen[None, None, :, None]).sum(2) - return num / den[..., None], exact - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="Sol-Attn needs CUDA") -def test_cute_kernel_matches_reference_on_multi_block_with_approximation(): - """The real kernel agrees with the reference where blocks are approximated. - - Complements the single-block test, which can only reach the all-exact - path. Plain Gaussian inputs are the right stimulus: every KV block carries - comparable softmax mass and no block's column-mean score clears the - `mean + tau*std` threshold, so every block two or more away from the - diagonal is approximated -- and because keys vary within a block, treating - 64 keys as their mean is measurably different from dense attention (a - Jensen gap). Clustered inputs are the wrong stimulus: distant blocks get - routed to the approximation but carry no mass, so approximating them - changes nothing and the test would pass vacuously. The reference computes - the same routing, and the test asserts that some blocks were approximated - *and* that the result differs from dense, so neither an accidental - all-exact run nor a mass-free approximation can pass. It also asserts - `kernel_calls` advanced, so a dense fallback cannot pass either. - """ - sab = _backend_mod() - if not sab.sol_attn_supported(torch.empty(1, 8, 8, 128, device="cuda", dtype=torch.bfloat16)): - pytest.skip("no Sol-Attn kernel for this device") - - torch.manual_seed(0) - S, H, D, block = 256, 2, 128, 64 # 4 KV blocks; batch of 1 - dev = torch.device("cuda") - q, k, v = (torch.randn(1, S, H, D, device=dev, dtype=torch.bfloat16) for _ in range(3)) - scale = D**-0.5 - tau = 2.0 - - ref, exact = _sol_attn_reference(q, k, v, tau=tau, scale=scale, block=block) - assert (~exact).any(), "inputs did not force any block onto the approximation path" - dense = torch.nn.functional.scaled_dot_product_attention( - q.transpose(1, 2).float(), k.transpose(1, 2).float(), v.transpose(1, 2).float() - ).transpose(1, 2) - # Calibrated on B200: max |ref - dense| = 0.45 for this stimulus. A tight - # margin here would let a mass-free approximation pass vacuously. - assert (ref - dense).abs().max() > 1e-1, "approximation is not measurably different from dense" - - before = sab._SOL_STATS["kernel_calls"] - out = sab._run_sol_attn_bthd(q, k, v, tau=tau, thresh_type="diag", kv_splits=1) - torch.cuda.synchronize() - assert sab._SOL_STATS["kernel_calls"] == before + 1, "kernel did not run" - - # Calibrated on B200: max |kernel - ref| = 1.7e-3 (mean 2e-4); 5e-3 leaves - # ~3x headroom for bf16 rounding while still catching a routing or - # approximation-formula mismatch, which shows up at 1e-1 or worse. - torch.testing.assert_close(out.float(), ref, rtol=5e-3, atol=5e-3) - - -# --------------------------------------------------------------------------- -# Masks: routed to a mask-aware backend, and honored -- not just redirected -# --------------------------------------------------------------------------- - - -def _masked_sdpa_reference(q, k, v, *, key_padding_mask=None, is_causal=False): - attn_mask = ( - None if key_padding_mask is None else key_padding_mask.to(torch.bool)[:, None, None, :] - ) - return torch.nn.functional.scaled_dot_product_attention( - q.transpose(1, 2).float(), - k.transpose(1, 2).float(), - v.transpose(1, 2).float(), - attn_mask=attn_mask, - is_causal=is_causal and attn_mask is None, - ).transpose(1, 2) - - -def test_key_padding_mask_routes_to_vanilla_and_is_honored(monkeypatch): - """Masked self-attention must never reach the mask-blind sparse kernel. - - Equal Q/K lengths do not mean "unmasked": HunyuanVideo1.5 and GLM-Image - pass a `[B, S]` `key_padding_mask` on self-attention. The sparse kernel - takes no mask and CuTeDSL's dense `_fwd(**kwargs)` swallows it, so the only - 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 - - 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.disabled_until_timestep = None - calls = {"vanilla": 0} - real = attn._vanilla.forward - - def spy(*a, **k): - calls["vanilla"] += 1 - return real(*a, **k) - - monkeypatch.setattr(attn._vanilla, "forward", spy) - - torch.manual_seed(0) - q, k, v = (torch.randn(2, 64, 2, 128, dtype=torch.float32) for _ in range(3)) - mask = torch.ones(2, 64, dtype=torch.bool) - mask[0, 48:] = False # pad the tail of sample 0 - mask[1, :16] = False # pad the head of sample 1 - - assert not attn._can_serve(q, k, key_padding_mask=mask) - out = attn.forward(q, k, v, key_padding_mask=mask) - assert calls["vanilla"] == 1 - ref = _masked_sdpa_reference(q, k, v, key_padding_mask=mask) - torch.testing.assert_close(out.float(), ref, rtol=1e-4, atol=1e-4) - # And the mask actually mattered: unmasked attention gives a different answer. - assert (ref - _masked_sdpa_reference(q, k, v)).abs().max() > 1e-3 - - -def test_causal_mask_routes_to_dense_and_is_honored(monkeypatch): - """CAUSAL disqualifies the noncausal sparse kernel and is honored downstream. - - On sm100 the delegate is dense CuTeDSL, which supports CAUSAL; on any other - 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 - - 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.disabled_until_timestep = None - torch.manual_seed(0) - q, k, v = (torch.randn(1, 64, 2, 128, dtype=torch.float32) for _ in range(3)) - - assert not attn._can_serve(q, k, attention_mask=PredefinedAttentionMask.CAUSAL) - out = attn.forward(q, k, v, attention_mask=PredefinedAttentionMask.CAUSAL) - ref = _masked_sdpa_reference(q, k, v, is_causal=True) - torch.testing.assert_close(out.float(), ref, rtol=1e-4, atol=1e-4) - assert (ref - _masked_sdpa_reference(q, k, v)).abs().max() > 1e-3 - - -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.disabled_until_timestep = None - q = k = torch.randn(1, 64, 2, 128, dtype=torch.bfloat16) - assert attn._can_serve(q, k) - assert attn._can_serve(q, k, attention_mask=PredefinedAttentionMask.FULL, key_padding_mask=None) 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..4c7c87b94993 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_integration.py +++ b/tests/unittest/_torch/visual_gen/test_attention_integration.py @@ -722,107 +722,6 @@ def test_fast_cross_attention_wan_shapes( assert is_close, f"{attn_backend} cross-attn mismatch at Wan shapes: max_diff={max_diff:.2e}" -# ============================================================================ -# VSA self-attention (CUTEDSL backend, sparse_attention_config.algorithm='vsa') -# ============================================================================ - - -def _build_vsa_setup(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. - """ - from tensorrt_llm._torch.visual_gen.attention_backend import VSAMetadataBuilder - - latent_shape = (8, 8, 8) - seq_len = latent_shape[0] * latent_shape[1] * latent_shape[2] - num_heads = 4 - head_dim = 128 - hidden_size = num_heads * head_dim - device = torch.device("cuda") - dtype = torch.bfloat16 - - 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 - ) - 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}" - ) - copy_weights_self_attention(naive, integrated) - naive.eval() - integrated.eval() - - metadata = VSAMetadataBuilder().build( - current_timestep=0, - raw_latent_shape=latent_shape, - patch_size=(1, 1, 1), - vsa_sparsity=sparsity, - device=device, - ) - torch.manual_seed(seed) - hidden_states = torch.randn(batch_size, seq_len, hidden_size, device=device, dtype=dtype) - return SimpleNamespace( - naive=naive, - integrated=integrated, - metadata=metadata, - hidden_states=hidden_states, - gate_compress_zero=torch.zeros_like(hidden_states), - freqs_HSD=generate_rope_embeddings(seq_len, head_dim, device, is_HSD=True), - freqs_SHD=generate_rope_embeddings(seq_len, head_dim, device, is_HSD=False), - ) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="VSA needs CUDA") -def test_vsa_self_attention_equivalence_at_sparsity_zero(): - """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 - - s = _build_vsa_setup(sparsity=0.0, batch_size=2, seed=42) - - with torch.no_grad(): - out_naive = s.naive(s.hidden_states, *s.freqs_HSD) - with torch.no_grad(), set_vsa_forward_context(s.metadata): - out_vsa = s.integrated( - s.hidden_states, freqs=s.freqs_SHD, gate_compress=s.gate_compress_zero - ) - - assert out_naive.shape == out_vsa.shape, ( - f"shape mismatch: naive={out_naive.shape}, vsa={out_vsa.shape}" - ) - max_diff = (out_naive - out_vsa).abs().max().item() - mean_diff = (out_naive - out_vsa).abs().mean().item() - assert torch.allclose(out_naive, out_vsa, rtol=1e-2, atol=1e-2), ( - f"VSA(sparsity=0, G_c=0) deviates from naive dense SDPA: " - f"max_diff={max_diff:.2e}, mean_diff={mean_diff:.2e}" - ) - - -@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 - - s = _build_vsa_setup(sparsity=sparsity, batch_size=1, seed=0) - - 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}" - ) - - def test_trtllm_cached_prepare(): """Test that TRTLLM attention cached prepare works correctly. 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_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_trtllm_attention_metadata.py b/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py index fb0b16ba40eb..37796447dd8a 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,29 @@ # 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.sparse.params import ( + BlockSparseForwardInputs, + 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 +from tensorrt_llm._torch.visual_gen.cuda_graph_runner import resolved_extra_keys_scope + +_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: @@ -19,6 +39,89 @@ 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 + attention.layer_idx = 0 + attention.sparse_params = None + attention.metadata = visual_trtllm.TrtllmAttentionMetadata( + device=torch.device("cpu"), attention_metadata_state={} + ) + 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 +165,213 @@ 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_metadata_and_its_plan_cache(monkeypatch): + """Planned FMHA wrappers live on the core metadata, which the layers of one + component share per shape; another component gets its own metadata.""" + monkeypatch.setattr( + visual_trtllm, "BaseTrtllmAttentionMetadata", _FakeBaseTrtllmAttentionMetadata + ) + state = create_attention_metadata_state() + first = visual_trtllm.TrtllmAttentionMetadata( + device=torch.device("cpu"), attention_metadata_state=state + ) + second = visual_trtllm.TrtllmAttentionMetadata( + device=torch.device("cpu"), attention_metadata_state=state + ) + other = visual_trtllm.TrtllmAttentionMetadata( + device=torch.device("cpu"), attention_metadata_state=create_attention_metadata_state() + ) + seq_lens = torch.tensor([64], dtype=torch.int32) + + shared = first.prepare(1, seq_lens) + assert second.prepare(1, seq_lens) is shared + assert other.prepare(1, seq_lens) is not shared + assert set(state) == {"metadata_cache"} + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_cuda_graph_metadata_plans_its_own_wrappers(): + metadata = visual_trtllm.TrtllmAttentionMetadata( + attention_metadata_state=create_attention_metadata_state() + ).prepare(1, 64) + metadata.fmha_plan_caches["library"] = {"profile": object()} + + graph_metadata = metadata.create_cuda_graph_metadata(1) + + assert graph_metadata.fmha_plan_caches == {} + assert set(metadata.fmha_plan_caches) == {"library"} + + +def test_wrapper_uses_the_runner_resolved_phase_without_reading_the_timestep(monkeypatch): + attention = _make_wrapper() + attention.sparse_params = SimpleNamespace(disabled_until_timestep=0.6) + reads = Mock(side_effect=AssertionError("timestep read despite a runner-resolved phase")) + monkeypatch.setattr(visual_trtllm, "graph_phase_for_timestep", reads) + + with resolved_extra_keys_scope({"sparse_attn_phase": 1}): + assert attention._sparse_attn_phase(torch.tensor([0.8])) == 1 + reads.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_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 + + +@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 + + +@_REQUIRES_SM100 +def test_trtllm_skip_softmax_cutoff_is_capturable_with_a_device_timestep() -> None: + """Under the phase the CUDA graph runner resolves for the call, eager warmup + and capture use the host phase and never read the device timestep.""" + 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) + + # The runner publishes the phase it keyed the graph on (0: dense prefix). + with resolved_extra_keys_scope({"sparse_attn_phase": 0}): + 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_visual_gen_args.py b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py index 25e55c4ada7a..006cb38fd73e 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,52 @@ 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): + # 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" + @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..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, @@ -120,7 +147,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, @@ -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()