diff --git a/3rdparty/vendor_patches/flashinfer-prims-ts.patch b/3rdparty/vendor_patches/flashinfer-prims-ts.patch index 44f3e7a4581d..b5d4bc75155c 100644 --- a/3rdparty/vendor_patches/flashinfer-prims-ts.patch +++ b/3rdparty/vendor_patches/flashinfer-prims-ts.patch @@ -1,5 +1,5 @@ diff --git a/block_sparse.py b/block_sparse.py -index 77a9fc4542ab5fb82aa0df6f7c73e56b7d5fbe89..af132c6a69ce0e4a35903b8e4428f481e949f6b8 100644 +index c7e881ca73dfb95314da532912db39fe7b05eb94..e19dc7b1de7fca03c44b8e98217f8d10fb17e26b 100644 --- a/block_sparse.py +++ b/block_sparse.py @@ -1,3 +1,4 @@ @@ -12,15 +12,15 @@ index 77a9fc4542ab5fb82aa0df6f7c73e56b7d5fbe89..af132c6a69ce0e4a35903b8e4428f481 from flashinfer.api_logging import flashinfer_api -from flashinfer.trace.templates.attention import ( -- prims_ts_block_sparse_trace, +- prims_ts_block_sparse_trace_dispatch, - prims_ts_block_sparse_wrapper_trace_dispatch, - prims_ts_paged_block_sparse_trace_dispatch, - prims_ts_paged_block_sparse_wrapper_trace_dispatch, -) + from ._block_sparse.common import _validate_contiguous_route_mode from ._block_sparse.config import _validate_block_sparse_static_profile - from ._block_sparse.inspection import ( -@@ -220,7 +215,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase): +@@ -239,7 +234,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase): # previously published revision intact and runnable. self._plan_state = candidate @@ -29,16 +29,16 @@ index 77a9fc4542ab5fb82aa0df6f7c73e56b7d5fbe89..af132c6a69ce0e4a35903b8e4428f481 def run( self, q: torch.Tensor, -@@ -302,7 +297,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase): +@@ -361,7 +356,7 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase): return self._launch_validated_run(state, run_args, run_stream) --@flashinfer_api(trace=prims_ts_block_sparse_trace) +-@flashinfer_api(trace=prims_ts_block_sparse_trace_dispatch) +@flashinfer_api def block_sparse_attention( q: torch.Tensor, k: torch.Tensor, -@@ -512,7 +507,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase): +@@ -607,7 +602,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase): ) self._plan_state = candidate @@ -47,7 +47,7 @@ index 77a9fc4542ab5fb82aa0df6f7c73e56b7d5fbe89..af132c6a69ce0e4a35903b8e4428f481 def run( self, q: torch.Tensor, -@@ -618,7 +613,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase): +@@ -727,7 +722,7 @@ class BlockSparsePagedTSWrapper(_BlockSparseWrapperBase): return self._launch_validated_run(state, run_args, run_stream) @@ -57,10 +57,10 @@ index 77a9fc4542ab5fb82aa0df6f7c73e56b7d5fbe89..af132c6a69ce0e4a35903b8e4428f481 q: torch.Tensor, paged_kv_cache: PagedKVCache, diff --git a/context.py b/context.py -index 7245bee1a0f725086171c9c5002115757e425d84..47996c867a962a3684c78e07d3a14eebf34b8452 100644 +index cea5a41438d9d72e152175999453881cc9c3e5e6..3a93019ff4a5995968c5efa9000e8ebd52c8b424 100644 --- a/context.py +++ b/context.py -@@ -29,8 +29,7 @@ position is ``q + (S_kv - S_q)`` and ``window_left`` is measured from that +@@ -31,8 +31,7 @@ position is ``q + (S_kv - S_q)`` and ``window_left`` is measured from that position. PrimTS context entry points are intentionally excluded from ``fi_trace`` for @@ -70,7 +70,7 @@ index 7245bee1a0f725086171c9c5002115757e425d84..47996c867a962a3684c78e07d3a14eeb """ from dataclasses import dataclass -@@ -424,7 +423,7 @@ def _validate_device(device: torch.device) -> int: +@@ -426,7 +425,7 @@ def _validate_device(device: torch.device) -> int: # Rubin runs through the sm_100f family target; a CuTe DSL older than 4.8 # cannot emit for it unless CUTE_DSL_ARCH=sm_100f is set before import. if capability == (10, 7): diff --git a/3rdparty/vendor_sources.lock.yaml b/3rdparty/vendor_sources.lock.yaml index 705d8af9091d..97a88dadbe93 100644 --- a/3rdparty/vendor_sources.lock.yaml +++ b/3rdparty/vendor_sources.lock.yaml @@ -1,13 +1,13 @@ schema_version: 1 vendors: flashinfer-prims-ts: - url: https://github.com/yuxianq/flashinfer.git - branch: trtllm-prims-ts-dev - commit: e500966b575ab83db7c0e84e5a0f8fde6a4f3505 + url: https://github.com/heyuhhh/flashinfer.git + branch: yuhangh/tmp-sol-attn-trtllm-dev + commit: 61454c5ce6e9f020a2158f24059d0b91322e3e77 source: flashinfer/attention/prims_ts destination: tensorrt_llm/_torch/attention/backends/prims_ts include: - '**/*.py' patch: 3rdparty/vendor_patches/flashinfer-prims-ts.patch - patch_digest: sha256:0e2f58c6633f57fee03df42049bc78d4d038063b810ad3a2f0ba6d62f8183887 - digest: sha256-tree-v1:e9af5482f6406af3128e711c3fb2d044359fb6e1d1bc9907743dc86d006d23ac + patch_digest: sha256:b590a3e86c8a2a54a8da5a2268aa9e67f8425f401f6b98972841bfe5484b5732 + digest: sha256-tree-v1:e89b89471aac59e2689ed0fb8312cef77b08e1b46115e7668d751b9f402a29ba diff --git a/docs/source/developer-guide/sparse-attention-development-guide.md b/docs/source/developer-guide/sparse-attention-development-guide.md index 32d103b50058..20f2b771374c 100644 --- a/docs/source/developer-guide/sparse-attention-development-guide.md +++ b/docs/source/developer-guide/sparse-attention-development-guide.md @@ -136,19 +136,48 @@ Dynamic generation-phase KV eviction is tracked as future work. ### Prediction hooks -`TrtllmAttention`-based sparse backends expose two prediction methods that +`TrtllmAttention`-based sparse backends expose three prediction methods that algorithm-specific subclasses override: ```python sparse_kv_indices, sparse_kv_offsets = self.sparse_kv_predict(q, k, metadata, forward_args) sparse_attn_indices, sparse_attn_offsets = self.sparse_attn_predict(q, k, metadata, forward_args) +block_sparse_inputs = self.block_sparse_attn_predict(q, k, v, metadata, forward_args) ``` -`hooks.py` writes these results to `SparseRuntimeParams`. SkipSoftmax writes -its thresholds to the same runtime interface consumed by `AttentionOp`. -`AttentionForwardArgs.sparse_backend_args` carries algorithm inputs from the -module to the backend, while `sparse_runtime_params` carries lowered inputs -from the backend to `AttentionOp`. +`prepare_sparse_runtime_params` in `sparse/hooks.py` runs all three hooks once +per call regardless of whether the backend carries `SparseParams`, applies the +SkipSoftmax threshold schedule when the backend carries `SkipSoftmaxParams`, +and returns a new per-call `SparseRuntimeParams` built from the caller's +`AttentionForwardArgs.sparse_runtime_params` plus the hook results. The core +forward never assigns that field; it dispatches with the returned carrier. +Backends that need runtime state outside the three hooks (DSA's auxiliary pool +pointer, DeepSeek-V4's per-token KV lengths) write it into the caller's carrier +before or inside their hooks, and `prepare_sparse_runtime_params` carries those +fields over. +`AttentionForwardArgs.sparse_backend_args` carries +algorithm inputs from the module to the backend, while +`AttentionForwardArgs.sparse_runtime_params` carries the complete lowered state +from the backend through FMHA dispatch to `AttentionOp`. + +`SparseRuntimeParams.block_sparse_inputs` is the nested carrier for optional, +algorithm-neutral `BlockSparseForwardInputs`. The selected general +block-sparse FMHA validates and consumes that field; dense FMHA libraries reject +it instead of silently ignoring its routes. `AttentionForwardArgs` defaults +the field to an empty `SparseRuntimeParams()`; the core forward always +dispatches with the carrier prepared for the current call. + +`block_sparse_attn_predict` runs even when the backend has no `SparseParams`. +Its default implementation hands through +`SparseBackendForwardArgs.block_sparse_inputs`, so an attention module that +predicts routes before the core forward only needs to place the complete +payload in `sparse_backend_args`. Algorithms that predict inside the backend +override the hook, read `metadata` for the batch layout and `forward_args` for +per-call state such as `timestep`, and return `None` for dense phases. + +The core contract owns this runtime transport and general block-sparse FMHA +execution. Algorithm integrations own their prediction policy, effective Q/K/V +preparation, and any post-processing around the normal core forward. Different KV heads are allowed to emit different sparse index sets; Q heads that map to the same KV head share the KV head's sparse pattern. @@ -288,6 +317,19 @@ prediction methods. A `VanillaAttention` implementation instead overrides different index layouts. Match the selected kernel contract; do not pass request-local block indices to the physical-token path. +**`block_sparse_attn_predict(self, q, k, v, metadata, forward_args)`** + +- **Behavior**: return the `BlockSparseForwardInputs` consumed by the + general block-sparse FMHA, or `None` for a dense call. +- **Outputs**: block geometry plus exactly one route representation + (BSR `block_indptr`/`block_indices` or a packed `exact_block_bits` + bitmask), optional K/V summaries for proxy routes, and optional + `kv_valid_bits` masking ragged KV tails. +- **Default**: hands through `SparseBackendForwardArgs.block_sparse_inputs`, + so modules that predict before the core forward do not override it. + Override it to predict inside the backend from the flattened Q/K/V, + the batch layout in `metadata`, and per-call state in `forward_args`. + Prediction is on the critical path and can dominate latency in low-latency scenarios. Plan for custom kernels (Triton or CUDA) rather than relying on generic PyTorch ops. diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index f358de023608..f5edc0f49810 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -73,7 +73,7 @@ Models are auto-detected from the checkpoint directory. Diffusers-format models [^1]: FLUX models use embedded guidance and do not have a separate negative prompt path, so CFG parallelism is not applicable. -[^2]: `FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers` — VSA-fine-tuned checkpoint with learned sparse-attention gates. Requires `CUTEDSL` on Blackwell sm_100+ (falls back to dense SDPA on older hardware). Ring and Attention2D not supported (no LSE output); Ulysses supported. +[^2]: `FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers` — VSA-fine-tuned checkpoint with learned sparse-attention gates. Supports the `CUTEDSL` and `TRTLLM` attention backends; each uses its block-sparse fine stage when supported and otherwise uses its dense path with compact Q/K/V. Ring and Attention2D are not supported (no LSE output); Ulysses is supported. [^3]: Wan 2.2 has two stage transformers; TeaCache requires explicit `teacache.coefficients` (high-noise) and `teacache.coefficients_2` (low-noise). There is no built-in coefficient table for Wan 2.2. @@ -408,12 +408,15 @@ args = VisualGenArgs( ### Video Sparse Attention (VSA) -VSA reduces the compute cost of self-attention in video diffusion models by selectively attending to only the most relevant spatial-temporal blocks. It uses a two-branch design: a lightweight coarse mean-pool branch computes block-level attention scores to identify the top-K most relevant token blocks, then a fine branch runs a block-sparse CuTe kernel over only those blocks. The two outputs are blended with learned gates. +VSA reduces the compute cost of self-attention in video diffusion models by selectively attending to only the most relevant spatial-temporal blocks. It uses a two-branch design: a lightweight coarse mean-pool branch computes block-level attention scores to identify the top-K most relevant token blocks, then a fine branch runs the selected backend's block-sparse kernel over only those blocks. The two outputs are blended with learned gates. + +VisualGen owns VSA route prediction and coarse/fine post-processing. With the `TRTLLM` backend, it nests the predicted routes in `SparseRuntimeParams.block_sparse_inputs` and passes those precomputed runtime parameters through the normal core attention forward. The core `PrimsTSBlockSparseFmha` owns the general block-sparse execution contract; it does not own VSA-specific prediction or blending. **Requirements:** - VSA-fine-tuned checkpoint: [`FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers`](https://huggingface.co/FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers). Standard Wan checkpoints do not have the learned VSA gates. -- Blackwell GPU (sm_100+) for the CuTe JIT kernel. Falls back to dense SDPA on older hardware with no accuracy loss. -- `CUTEDSL` attention backend. +- `CUTEDSL` or `TRTLLM` attention backend. `CUTEDSL` uses the CuTe DSL fine-stage kernel; `TRTLLM` lowers the selected blocks through the generic PrimTS block-sparse FMHA contract. +- A supported CUDA device and tensor shape for the selected block-sparse kernel. When that kernel is unavailable or the input is outside its supported envelope, the fine branch uses the selected backend's compact dense path (`SDPA` for `CUTEDSL`, TRTLLM attention for `TRTLLM`). +- VSA cannot be combined with `quant_attention_config`. - Not compatible with Ring attention or Attention2D (VSA does not produce per-split LSE). Ulysses is supported. **`vsa_sparsity`** controls the fraction of K/V blocks skipped in the fine branch (0.0 = dense, 0.9 = 90% blocks skipped). Higher sparsity gives more speedup at the cost of some quality. @@ -427,7 +430,7 @@ from tensorrt_llm.visual_gen.args import AttentionConfig, VideoSparseAttentionCo args = VisualGenArgs( model="FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers", attention_config=AttentionConfig( - backend="CUTEDSL", + backend="TRTLLM", # Use "CUTEDSL" for the CuTe DSL fine-stage kernel. sparse_attention_config=VideoSparseAttentionConfig(vsa_sparsity=0.9), ), ) @@ -437,7 +440,7 @@ YAML (for use with `--visual_gen_args` or `trtllm-serve`): ```yaml attention_config: - backend: CUTEDSL + backend: TRTLLM # CUTEDSL is also supported. sparse_attention_config: algorithm: vsa vsa_sparsity: 0.90 diff --git a/docs/source/visual-gen/features/sparse-attention.md b/docs/source/visual-gen/features/sparse-attention.md index 7edeff8fb63c..1b1919f8dd8e 100644 --- a/docs/source/visual-gen/features/sparse-attention.md +++ b/docs/source/visual-gen/features/sparse-attention.md @@ -7,20 +7,22 @@ This page is an unindexed draft until the VisualGen documentation hub is introdu - [Overview](#overview) - [Algorithms](#algorithms) - [Skip Softmax Attention](#skip-softmax-attention) +- [SOL Attention](#sol-attention) - [Video Sparse Attention (VSA)](#video-sparse-attention-vsa) ## Overview Visual generation models naturally operate on long image or video token sequences. Each denoising step is closer to a full-context prefill pass than to autoregressive decoding, and attention can dominate runtime for high-resolution image generation or long video generation. -Sparse attention in VisualGen is configured through `VisualGenArgs.attention_config.sparse_attention_config`. The user-facing config stays in VisualGen args or model config. Checkpoint calibration metadata remains internal and is lowered into per-attention-backend `SparseParams` when each attention module is constructed. +Sparse attention in VisualGen is configured through `VisualGenArgs.attention_config.sparse_attention_config`. The user-facing config stays in VisualGen args or model config, while `attention_config.backend` selects the kernel family. Algorithms produce their block-sparse routes through the core `block_sparse_attn_predict` hook: a backend either predicts inside that hook from the flattened Q/K/V, or predicts before the core forward and hands the complete `BlockSparseForwardInputs` through `AttentionForwardArgs.sparse_backend_args`, which the default hook passes through. `SparseRuntimeParams` is the single lowered runtime carrier passed as `AttentionForwardArgs.sparse_runtime_params`; its optional `block_sparse_inputs` field nests the algorithm-neutral routes for the general block-sparse FMHA. `None` means prediction has not run, while an empty `SparseRuntimeParams()` records that prediction ran without a sparse payload. ### Algorithms | `algorithm` | Config class | Status | |---|---|---| | `skip_softmax` | `SkipSoftmaxAttentionConfig` | Supported | -| VSA | TBD | TODO | +| `vsa` | `VideoSparseAttentionConfig` | Supported (`CUTEDSL`, `TRTLLM`) | +| `sol_attn` | `SolAttentionConfig` | Experimental (`TRTLLM`) | ## Skip Softmax Attention @@ -214,6 +216,80 @@ attention_config: Graphs are captured lazily. The first denoising step seen for a given tensor shape and sparse-attention phase captures a graph; later steps with the same shape and phase replay that graph. When denoising crosses the cutoff, the phase key changes, so VisualGen captures a second graph for the enabled phase instead of replaying the graph from the disabled phase. +## SOL Attention + +SOL is a two-stage self-attention algorithm. A TRT-LLM-owned predictor first +produces an exact block bitmask and K/V proxy summaries from Q/K/V. The shared +`PrimsTSBlockSparseFmha` library then executes that route from the shared +`SparseRuntimeParams`. `SOLTrtllmAttention` is only the VisualGen bridge; +SOL does not add an algorithm-specific core attention backend or FMHA library. + +Configure SOL with `SolAttentionConfig` and the `TRTLLM` backend: + +```python +from tensorrt_llm.visual_gen import AttentionConfig, SolAttentionConfig + +attention_config = AttentionConfig( + backend="TRTLLM", + sparse_attention_config=SolAttentionConfig( + tau=1.0, + disabled_until_timestep=0.6, + dense_layers="0,2-4", + ), +) +``` + +The equivalent YAML is: + +```yaml +attention_config: + backend: TRTLLM + sparse_attention_config: + algorithm: sol_attn + tau: 1.0 + disabled_until_timestep: 0.6 + dense_layers: "0,2-4" +``` + +- `tau` is the finite float32 routing threshold consumed by the predictor. +- `disabled_until_timestep` is an optional normalized cutoff in `(0, 1]`. + Attention stays dense while the current timestep is greater than or equal to + the cutoff and switches to SOL below it. +- `dense_layers` is an optional comma-separated list of zero-based layer indices + and inclusive ranges that always use dense attention. + +The initial SOL envelope is full-mask BF16 self-attention on SM100 or SM103, +with 4-D BSHD Q/K/V tensors, equal Q/K/V shapes, and head dimension 128. +`SOLTrtllmAttention` overrides the core `block_sparse_attn_predict` hook, so +prediction runs inside the core forward from the flattened Q/K/V, the batch +layout in the attention metadata, and the `timestep` in the forward arguments; +dense layers and dense timestep phases return no routes. The VisualGen wrapper +compacts fused projection split views once and shares those tensors between +prediction and the block-sparse FMHA. Cross-attention, context parallelism, +attention quantization, and unsupported tensor envelopes raise an error instead +of silently falling back to dense attention. +SOL uses a host-side graph break to prepare and own predictor plans, so +`torch_compile_config.enable_fullgraph=True` is not supported; keep the default +`False` setting. + +When a cutoff is configured, VisualGen includes the dense-or-sparse phase in +the CUDA Graph key. Each SOL backend prepares that phase during graph warmup +and reuses it during capture, while predictor route buffers remain stable for +replay. A dense capture therefore cannot be reused for the sparse phase. + ## Video Sparse Attention (VSA) -TODO +VSA combines a coarse mean-pooled branch with a top-K block-sparse fine branch. Select either `CUTEDSL` for the CuTe DSL kernel or `TRTLLM` for PrimTS block-sparse attention. If the selected sparse kernel is unavailable or the known VSA tensor envelope is not met, the fine branch uses the compact Q/K/V tensors with that backend's dense path. VSA cannot be combined with `quant_attention_config`. + +VSA retains shape-dependent metadata and route tensors so CUDA Graph replay can reuse stable addresses. A pipeline instance accepts up to 16 distinct VSA shape profiles; reuse configured resolution/frame profiles or restart the pipeline before serving additional shapes. + +Both VSA backends share one VisualGen-owned predictor and identical +post-processing. The `TRTLLM` path runs the coarse stage before the core +forward, hands the predicted `BlockSparseForwardInputs` (including the +tile-padding validity bits only the VSA predictor knows) through +`sparse_backend_args`, lets the default core prediction hook pass them to the +general block-sparse FMHA, and then blends the fine and coarse outputs. Its +compact dense fallback passes no sparse inputs, so the core runs dense attention +and VSA post-processing still runs. `CUTEDSL` retains only its backend-specific +fine-attention execution. The core FMHA registry owns the reusable block-sparse +implementation rather than a VSA-specific lifecycle. diff --git a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md index e56bbb5fa24a..f8e438cabab6 100644 --- a/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/attention/ATTENTION_DEVELOPER_GUIDE.md @@ -155,9 +155,12 @@ their module-to-backend inputs in a `SparseBackendForwardArgs` subclass and pass it through the registered `AttentionForwardArgs.sparse_backend_args` field. For example, DSA owns `DSABackendForwardArgs`, whose indexer intermediates are consumed by `DSATrtllmAttention.sparse_attn_predict`. -Shared sparse carriers, including `SparseBackendForwardArgs.topk_indices` and -the backend-to-AttentionOp `SparseRuntimeParams`, live in -`attention/backends/sparse/params.py`. +Shared sparse carriers, including `SparseBackendForwardArgs.topk_indices`, +`SparseBackendForwardArgs.block_sparse_inputs`, and the +backend-to-FMHA/`AttentionOp` `SparseRuntimeParams`, live in +`attention/backends/sparse/params.py`. The latter is carried by +`AttentionForwardArgs.sparse_runtime_params` and nests optional general +block-sparse inputs in `SparseRuntimeParams.block_sparse_inputs`. For MLA-related tasks, first check whether the work fits the current projection structure, can stay on an existing backend and metadata family, and @@ -211,6 +214,35 @@ that file for the current config/backend combinations. Consult the for the supported attention shapes; do not infer support from algorithm registration alone. +Block-sparse FMHA is a kernel-library contract rather than a sparse algorithm. +Algorithms lower their live routing state to an algorithm-neutral +`BlockSparseForwardInputs`, nested at +`SparseRuntimeParams.block_sparse_inputs`: block geometry plus either canonical +BSR routes or an exact packed bitmask. Optional K/V summaries enable proxy +routes, and optional token-validity bits mask ragged KV tails. Plans contain +only static format, proxy, geometry, and capacity choices; every run receives +the live routes, summaries, validity bits, page tables, and sequence lengths. + +`PrimsTSBlockSparseFmha` owns its wrapper-plan cache by default. Integrations +whose attention layers execute serially may explicitly bind a model-scoped +cache to reuse graph-stable route workspaces across compatible layers. The +cache must not be shared by concurrent forwards; each independent model +component must own separate state. + +`TrtllmAttention.block_sparse_attn_predict(q, k, v, metadata, forward_args)` +is the backend hook that produces this payload; `prepare_sparse_runtime_params` +calls it even when the backend has no `SparseParams`. The default hands through +`SparseBackendForwardArgs.block_sparse_inputs`, which lets an attention module +predict routes before the core forward and pass the complete payload in +`AttentionForwardArgs.sparse_backend_args`. Algorithms that predict inside the +backend override the hook and return `None` for dense phases. + +The core library owns this general planning, validation, and execution +contract. Algorithm integrations own the surrounding lifecycle: prediction +policy, effective Q/K/V preparation before the core forward, plus any +algorithm-specific post-processing afterward. They route their payload through +these hooks instead of adding algorithm-specific FMHA libraries. + ### 2.3 Backend contract All backends implement the `AttentionBackend` interface. @@ -343,13 +375,23 @@ starting with an empty selection cache. `TrtllmAttention` prepares the complete per-forward state, passes itself to the manager for selection, and then executes the selected library. -`TLLM_FMHA_LIBS` controls the ordered selection. PrimTS is opt-in because it may -add host overhead; use `TLLM_FMHA_LIBS=+prims_ts` to add it to the defaults or -`TLLM_FMHA_LIBS=fallback` to force the fallback path. Delta entries update the +`TLLM_FMHA_LIBS` controls the ordered selection. Dense PrimTS is opt-in because +it may add host overhead; use `TLLM_FMHA_LIBS=+prims_ts` to add it to the +defaults or `TLLM_FMHA_LIBS=fallback` to force the fallback path. Generic +block-sparse PrimTS remains enabled by default because a dense fallback cannot +preserve its routing semantics. Delta entries update the default membership and follow canonical registry order, while an exact list preserves the user-specified order. Each FMHA library exposes `is_available()` for module/static environment checks and `is_supported()` for per-forward -request checks. For mixed non-MLA batches, the manager checks each active phase +request checks. `AttentionForwardArgs.sparse_runtime_params` is the sole +per-call lowered sparse runtime carrier and defaults to an empty +`SparseRuntimeParams()`. The core forward never assigns that field; it +dispatches with a new carrier that `prepare_sparse_runtime_params` builds from +the caller's carrier plus the hook results. The carrier holds both flat `AttentionOp` parameters and optional +`BlockSparseForwardInputs` in its nested `block_sparse_inputs` field. The +block-sparse implementation validates and consumes the latter, while every +other library rejects it rather than silently dropping sparse routing semantics. +For mixed non-MLA batches, the manager checks each active phase independently with `is_supported(..., phase=...)`; a phased library accepts only phases backed by its corresponding `run_*()` entry point. @@ -367,6 +409,10 @@ The FMHA package is split by role: `TrtllmAttention` can pair it with a later causal-generation provider through `CombinedFmha`. - `fmha/cute_dsl_mla.py` implements the CuTe DSL MLA decode FMHA library. +- `fmha/prims_ts_block_sparse.py` adapts generic block-sparse requests to the + vendored PrimTS contiguous and paged wrappers. Paged generation passes a + live, zero-copy 2D K-page-table view with its TRT-LLM padded row stride; it + does not stage page tables through CSR metadata. - `fmha/prims_ts.py` adapts TRT-LLM inputs and paged-cache metadata to the vendored PrimTS kernels. Before changing the managed source under `backends/prims_ts`, read the diff --git a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py index 9584d131e245..12342a23ea2e 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py @@ -285,6 +285,8 @@ def _is_supported_with_reason( meta: "TrtllmAttentionMetadata", fwd: AttentionForwardArgs, ) -> tuple[bool, str]: + if fwd.sparse_runtime_params.block_sparse_inputs is not None: + return False, "block_sparse_inputs are not supported." if fwd.attention_input_type != AttentionInputType.generation_only: return False, "CuTe DSL MLA FMHA only supports generation-only attention." # Disable mixed context/generation batches until the CuTe DSL host diff --git a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py index f03ae9a01dc1..033ba3b4aa03 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/fallback.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/fallback.py @@ -40,10 +40,11 @@ _THOP_EXCLUDED_FIELDS: frozenset = frozenset( { "sparse_backend_args", # consumed by sparse prediction before the attention op + "block_sparse_inputs", # consumed by the selected block-sparse FMHA "attention_mask_data", # custom-mask code path "out_scale_sf", # promoted into ``out_scale`` in ``TrtllmAttention.forward`` for NVFP4 path "skip_mla_rope_generation", # handled in ``TrtllmAttention.forward`` for the test-only MLA path - "timestep", # used to populate skip-softmax params in ``TrtllmAttention.forward`` + "timestep", # consumed by sparse prediction before FMHA dispatch } ) @@ -81,8 +82,9 @@ def is_supported( del k, v, phase if q is not None and q.dtype == torch.float8_e4m3fn: return False - return forward_args.attention_mask != CustomAttentionMask.CUSTOM and ( - forward_args.update_kv_cache or metadata.is_cross + return forward_args.sparse_runtime_params.block_sparse_inputs is None and ( + forward_args.attention_mask != CustomAttentionMask.CUSTOM + and (forward_args.update_kv_cache or metadata.is_cross) ) def forward( @@ -218,7 +220,7 @@ def forward( forward_args.sparse_runtime_params.sparse_attn_indices_block_size ), sparse_attn_kv_lens=forward_args.sparse_runtime_params.sparse_attn_kv_lens, - aux_kv_cache_pool_ptr=(forward_args.sparse_runtime_params.aux_kv_cache_pool_ptr), + aux_kv_cache_pool_ptr=forward_args.sparse_runtime_params.aux_kv_cache_pool_ptr, skip_softmax_threshold_scale_factor_prefill=( forward_args.sparse_runtime_params.threshold_scale_factor_prefill ), diff --git a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_sparse_mla.py b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_sparse_mla.py index 2924cc440fc0..1605e9e45151 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_sparse_mla.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_sparse_mla.py @@ -13,7 +13,7 @@ from tensorrt_llm._utils import get_sm_version from tensorrt_llm.logger import logger -from .interface import Fmha +from .interface import Fmha, FmhaPhase if TYPE_CHECKING: from tensorrt_llm._torch.attention.backends.trtllm import ( @@ -70,6 +70,19 @@ def is_available(cls, attn: "TrtllmAttention") -> bool: return False return is_flashinfer_sparse_mla_enabled(getattr(attn.sparse_params, "algorithm", None)) + def is_supported( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + *, + phase: Optional[FmhaPhase] = None, + ) -> bool: + del q, k, v, metadata, phase + return forward_args.sparse_runtime_params.block_sparse_inputs is None + def forward( self, q: torch.Tensor, diff --git a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py index dfc5d2e4af45..cc805ff9d724 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py @@ -571,6 +571,9 @@ def _is_supported_with_reason( *, phase: Optional[FmhaPhase] = None, ) -> Tuple[bool, str]: + sparse_runtime_params = fwd.sparse_runtime_params + if sparse_runtime_params.block_sparse_inputs is not None: + return False, "block_sparse_inputs are not supported." is_mla_enable = attn.is_mla_enable if phase is None: has_context_phase = fwd.attention_input_type != AttentionInputType.generation_only @@ -652,8 +655,8 @@ def _is_supported_with_reason( return False, "sage attention." if meta.helix_position_offsets is not None: return False, "helix parallelism." - sparse_kv_indices = fwd.sparse_runtime_params.sparse_kv_indices - sparse_attn_indices = fwd.sparse_runtime_params.sparse_attn_indices + sparse_kv_indices = sparse_runtime_params.sparse_kv_indices + sparse_attn_indices = sparse_runtime_params.sparse_attn_indices if ( (sparse_kv_indices is not None and sparse_kv_indices.numel() > 0) or (sparse_attn_indices is not None and sparse_attn_indices.numel() > 0) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/manager.py b/tensorrt_llm/_torch/attention/backends/fmha/manager.py index 07837c39eadd..c9afff66e1d8 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/manager.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/manager.py @@ -134,6 +134,7 @@ class _FmhaCacheKey(NamedTuple): generation_seq_len_q: int attention_mask_type: AttentionMaskType use_spec_decoding: bool + has_block_sparse_inputs: bool # LoRA can change the effective output from packed NVFP4 to unpacked BF16 # without changing the request shape. Keep those selection regimes apart. output_dtype: torch.dtype | None @@ -366,12 +367,14 @@ def _make_cache_key( generation_seq_len_q, _FMHA_CACHE_SEQ_LEN_Q_GRID ) + block_sparse_inputs = forward_args.sparse_runtime_params.block_sparse_inputs return _FmhaCacheKey( context_batch_size=context_batch_size, generation_batch_size=generation_batch_size, generation_seq_len_q=generation_seq_len_q, attention_mask_type=attention_mask_type, use_spec_decoding=metadata.use_spec_decoding, + has_block_sparse_inputs=block_sparse_inputs is not None, output_dtype=output_dtype, output_sf_dtype=output_sf_dtype, ) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/msa_sparse_gqa.py b/tensorrt_llm/_torch/attention/backends/fmha/msa_sparse_gqa.py index ee85aaf02e0d..6f6ce04b576c 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/msa_sparse_gqa.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/msa_sparse_gqa.py @@ -5,7 +5,7 @@ MsaSparseGqaFmha wraps the fmha_sm100 paged sparse GQA kernel and participates in the standard TrtllmAttention.forward dispatch loop. The owning MiniMax-M3 MSA attention layer runs an MsaIndexer to select the -per-query KV blocks and publishes them on forward_args.sparse_runtime_params; +per-query KV blocks and publishes them through the aggregate sparse prediction; this class attends over them. """ @@ -17,7 +17,7 @@ from tensorrt_llm._utils import is_sm_100f -from .interface import Fmha +from .interface import Fmha, FmhaPhase if TYPE_CHECKING: from tensorrt_llm._torch.attention.backends.interface import AttentionForwardArgs @@ -175,7 +175,7 @@ class MsaSparseGqaFmha(Fmha): """SM100 paged GQA FMHA powered by MSA's fmha_sm100 kernel. Handles every MiniMax-M3 MSA layer. Sparse layers pass the indexer's - selected KV block indices on forward_args.sparse_runtime_params.sparse_attn_indices + selected KV block indices on the aggregate sparse prediction and attend those blocks; dense layers leave the indices None and attend the full page table. @@ -210,6 +210,19 @@ def is_available(cls, attn: Optional["TrtllmAttention"] = None) -> bool: # alone, so no custom library discovery is needed. return attn.sparse_params is not None and attn.sparse_params.algorithm == "minimax_m3" + def is_supported( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: "TrtllmAttentionMetadata", + forward_args: "AttentionForwardArgs", + *, + phase: Optional[FmhaPhase] = None, + ) -> bool: + del q, k, v, metadata, phase + return forward_args.sparse_runtime_params.block_sparse_inputs is None + def forward( self, q: torch.Tensor, diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py index afad6479067d..3b5be3d9bcff 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py @@ -20,7 +20,7 @@ import math from importlib import import_module from importlib.metadata import PackageNotFoundError, version -from typing import TYPE_CHECKING, Optional +from typing import TYPE_CHECKING, Any, Optional import torch from packaging.version import InvalidVersion, Version @@ -59,6 +59,96 @@ _WORKSPACE_ALIGNMENT = 32 +def get_paged_kv_storage_unsupported_reason( + attn: "TrtllmAttention", + metadata: "TrtllmAttentionMetadata", +) -> Optional[str]: + """Return why the TRT-LLM paged KV storage cannot feed a fixed page-table kernel.""" + if metadata.kv_cache_manager is None: + return "a KV cache manager is required." + if metadata.kv_cache_block_offsets is None: + return "paged KV-cache block offsets are required." + if metadata.host_kv_cache_pool_pointers is None: + return "KV-cache pool pointers are required." + pool_mapping = metadata.host_kv_cache_pool_mapping + if pool_mapping is None: + return "KV-cache pool mapping is required." + if metadata.kv_layout != "HND": + return "only HND KV-cache layout is supported." + manager = metadata.kv_cache_manager + if isinstance(manager, KVCacheManagerV2): + if manager.enable_swa_scratch_reuse: + return "KVCacheManagerV2 SWA scratch reuse is not supported." + elif isinstance(manager, KVCacheManager): + if manager.num_pools != 1: + return "KVCacheManagerV1 with multiple memory pools is not supported." + local_layer_idx = attn.local_layer_idx + if ( + pool_mapping.ndim != 2 + or pool_mapping.shape[1] < 2 + or local_layer_idx is None + or not 0 <= local_layer_idx < pool_mapping.shape[0] + ): + return "KVCacheManagerV1 has an invalid layer-to-pool mapping." + pool_index = int(pool_mapping[local_layer_idx, 0]) + layer_idx_in_pool = int(pool_mapping[local_layer_idx, 1]) + if pool_index != 0 or not 0 <= layer_idx_in_pool < manager.num_local_layers: + return "KVCacheManagerV1 has an invalid layer-to-pool mapping." + else: + return f"unsupported KV cache manager {type(manager).__name__}." + return None + + +def get_paged_kv_policy_unsupported_reason( + attn: "TrtllmAttention", + metadata: "TrtllmAttentionMetadata", +) -> Optional[str]: + """Return why the request's decoding policy is outside the fixed page-table envelope.""" + if metadata.beam_width != 1: + return "beam search is not supported." + if ( + metadata.is_spec_decoding_enabled + or metadata.use_spec_decoding + or metadata.is_spec_dec_tree + or metadata.is_spec_dec_dynamic_tree + ): + return "speculative decoding is not supported by the initial adapter." + position_embedding_type = int(attn.position_embedding_type) + if position_embedding_type in (4, 5, 6, 7, 10): + return f"position embedding type {position_embedding_type} is not supported." + try: + quant_mode = QuantMode(attn.quant_mode) + except (TypeError, ValueError): + return "invalid KV-cache quantization mode." + if quant_mode.has_kv_cache_quant(): + return "quantized KV cache is not supported by the initial adapter." + return None + + +def get_attention_feature_unsupported_reason( + metadata: "TrtllmAttentionMetadata", + forward_args: "AttentionForwardArgs", +) -> Optional[str]: + """Return which optional attention feature the fused-kernel adapters do not implement.""" + if metadata.helix_position_offsets is not None: + return "Helix parallelism is not supported." + if forward_args.relative_attention_bias is not None: + return "relative attention bias is not supported." + if forward_args.attention_sinks is not None: + return "attention sinks are not supported." + if forward_args.attention_mask_data is not None: + return "custom attention masks are not supported." + if forward_args.enable_dsv4_epilogue_fusion: + return "DSv4 epilogue fusion is not supported." + if ( + forward_args.sage_attn_num_elts_per_blk_q > 0 + or forward_args.sage_attn_num_elts_per_blk_k > 0 + or forward_args.sage_attn_num_elts_per_blk_v > 0 + ): + return "SageAttention is not supported." + return None + + class PrimsTSFmha(PhasedFmha): """Blackwell task-scheduled paged context and decode FMHA library.""" @@ -180,6 +270,9 @@ def _is_supported_with_reason( phase: Optional[FmhaPhase] = None, ) -> tuple[bool, str]: """Return a conservative, side-effect-free whole-request support decision.""" + sparse_runtime_params = fwd.sparse_runtime_params + if sparse_runtime_params.block_sparse_inputs is not None: + return False, "block_sparse_inputs are not supported." # PrimTS prepares workspace for every active request phase before # dispatch. Accept the phased dispatcher keyword, but do not narrow # support until that preparation is phase-aware too. @@ -194,40 +287,9 @@ def _is_supported_with_reason( return False, "only fused QKV input is supported." if meta.is_cross: return False, "cross attention is not supported." - if meta.kv_cache_manager is None: - return False, "a KV cache manager is required." - if meta.kv_cache_block_offsets is None: - return False, "paged KV-cache block offsets are required." - if meta.host_kv_cache_pool_pointers is None: - return False, "KV-cache pool pointers are required." - if meta.host_kv_cache_pool_mapping is None: - return False, "KV-cache pool mapping is required." - if meta.kv_layout != "HND": - return False, "only HND KV-cache layout is supported." - kv_cache_manager = meta.kv_cache_manager - if isinstance(kv_cache_manager, KVCacheManagerV2): - if kv_cache_manager.enable_swa_scratch_reuse: - return False, "KVCacheManagerV2 SWA scratch reuse is not supported." - elif isinstance(kv_cache_manager, KVCacheManager): - if kv_cache_manager.num_pools != 1: - return False, "KVCacheManagerV1 with multiple memory pools is not supported." - pool_mapping = meta.host_kv_cache_pool_mapping - local_layer_idx = attn.local_layer_idx - num_local_layers = kv_cache_manager.num_local_layers - if ( - pool_mapping.ndim != 2 - or pool_mapping.shape[1] < 2 - or local_layer_idx is None - or local_layer_idx < 0 - or local_layer_idx >= pool_mapping.shape[0] - ): - return False, "KVCacheManagerV1 has an invalid layer-to-pool mapping." - pool_index = int(pool_mapping[local_layer_idx, 0]) - layer_idx_in_pool = int(pool_mapping[local_layer_idx, 1]) - if pool_index != 0 or not 0 <= layer_idx_in_pool < num_local_layers: - return False, "KVCacheManagerV1 has an invalid layer-to-pool mapping." - else: - return False, f"unsupported KV cache manager {type(kv_cache_manager).__name__}." + storage_reason = get_paged_kv_storage_unsupported_reason(attn, meta) + if storage_reason is not None: + return False, storage_reason output = fwd.output if output is None: @@ -240,38 +302,15 @@ def _is_supported_with_reason( if attn.sparse_params is not None: return False, "sparse attention is not supported." if ( - fwd.sparse_runtime_params.sparse_kv_indices is not None - or fwd.sparse_runtime_params.sparse_attn_indices is not None + sparse_runtime_params.sparse_kv_indices is not None + or sparse_runtime_params.sparse_attn_indices is not None ): return False, "sparse attention metadata is not supported." if meta.num_sparse_topk > 0: return False, "sparse attention metadata is not supported." - if meta.helix_position_offsets is not None: - return False, "Helix parallelism is not supported." - if fwd.relative_attention_bias is not None: - return False, "relative attention bias is not supported." - if fwd.attention_sinks is not None: - return False, "attention sinks are not supported." - if fwd.attention_mask_data is not None: - return False, "custom attention masks are not supported." - if fwd.enable_dsv4_epilogue_fusion: - return False, "DSv4 epilogue fusion is not supported." - if ( - fwd.sage_attn_num_elts_per_blk_q > 0 - or fwd.sage_attn_num_elts_per_blk_k > 0 - or fwd.sage_attn_num_elts_per_blk_v > 0 - ): - return False, "SageAttention is not supported." - - if meta.beam_width != 1: - return False, "beam search is not supported." - if ( - meta.is_spec_decoding_enabled - or meta.use_spec_decoding - or meta.is_spec_dec_tree - or meta.is_spec_dec_dynamic_tree - ): - return False, "speculative decoding is not supported by the initial adapter." + feature_reason = get_attention_feature_unsupported_reason(meta, fwd) + if feature_reason is not None: + return False, feature_reason try: mask_type = AttentionMaskType(fwd.mask_type) @@ -279,17 +318,9 @@ def _is_supported_with_reason( return False, "the attention mask is not causal or dense." if mask_type not in (AttentionMaskType.causal, AttentionMaskType.padding): return False, f"attention mask type {mask_type} is not supported." - - position_embedding_type = int(attn.position_embedding_type) - if position_embedding_type in (4, 5, 6, 7, 10): - return False, f"position embedding type {position_embedding_type} is not supported." - - try: - quant_mode = QuantMode(attn.quant_mode) - except (TypeError, ValueError): - return False, "invalid KV-cache quantization mode." - if quant_mode.has_kv_cache_quant(): - return False, "quantized KV cache is not supported by the initial adapter." + policy_reason = get_paged_kv_policy_unsupported_reason(attn, meta) + if policy_reason is not None: + return False, policy_reason input_type = fwd.attention_input_type if input_type not in ( @@ -432,6 +463,27 @@ def _get_fixed_block_tables( ) return block_tables[:batch_size, 0, :] + def _get_generation_workspace_layout( + self, + dtype: torch.dtype, + num_requests: int, + num_tokens: int, + ) -> dict[str, int]: + """Return the shared TRT-LLM generation preprocessing layout.""" + + return thop.get_trtllm_gen_generation_workspace_layout( + dtype, + num_requests, + num_tokens, + self.attn.num_heads, + self.attn.head_dim, + self.attn.rope_dim, + self.attn.num_kv_heads, + 0, + False, + skip_fmha_workspace=True, + ) + @staticmethod def _get_sequence_lengths( sequence_lengths: torch.Tensor, @@ -652,17 +704,10 @@ def prepare_workspace( if input_type == AttentionInputType.generation_only else q.shape[0] - int(metadata.num_ctx_tokens) ) - generation_layout = thop.get_trtllm_gen_generation_workspace_layout( + generation_layout = self._get_generation_workspace_layout( q.dtype, int(metadata.num_generations), num_gen_tokens_for_layout, - self.attn.num_heads, - self.attn.head_dim, - self.attn.rope_dim, - self.attn.num_kv_heads, - 0, - False, - skip_fmha_workspace=True, ) required_preprocess_bytes = max( required_preprocess_bytes, int(generation_layout["total_size"]) @@ -956,34 +1001,17 @@ def run_context(self, params: FmhaParams) -> None: skip_fmha_workspace=True, ) - def run_generation(self, params: FmhaParams) -> None: - if params.qkv_input is None or params.context_buf is None: - raise RuntimeError("PrimTS decode requires QKV input and an output buffer.") - if params.sequence_lengths is None: - raise RuntimeError("PrimTS decode requires sequence lengths.") - if self._multi_processor_count is None: - raise RuntimeError("PrimTS decode workspace was not prepared.") + def _run_generation_preprocess(self, params: FmhaParams) -> tuple[Any, ...]: + """Run the shared TRT-LLM generation QKV and cache preprocessing.""" + if self._multi_processor_count is None: + raise RuntimeError("PrimTS generation workspace was not prepared.") attn = params.attn meta = params.meta fwd = params.fwd rope_params = attn.rope_params - batch_size = params.batch_size attention_chunk_size = attn.attention_chunk_size or 0 - ( - q_processed, - kv_pool, - block_tables, - _kv_scale_pool, - _bmm1_scale, - _bmm2_scale, - fmha_workspace, - _cu_seqlens, - _max_q_len, - _max_kv_len, - window_left, - is_multi_token_gen, - ) = thop.trtllm_gen_generation_preprocess( + return thop.trtllm_gen_generation_preprocess( params.qkv_input, params.workspace, params.sequence_lengths, @@ -1008,7 +1036,7 @@ def run_generation(self, params: FmhaParams) -> None: params.max_attention_window_size, params.cyclic_attention_window_size, params.num_tokens, - batch_size, + params.batch_size, params.input_seq_length, params.max_past_kv_length, rope_params.dim, @@ -1029,6 +1057,33 @@ def run_generation(self, params: FmhaParams) -> None: False, skip_fmha_workspace=True, ) + + def run_generation(self, params: FmhaParams) -> None: + if params.qkv_input is None or params.context_buf is None: + raise RuntimeError("PrimTS decode requires QKV input and an output buffer.") + if params.sequence_lengths is None: + raise RuntimeError("PrimTS decode requires sequence lengths.") + if self._multi_processor_count is None: + raise RuntimeError("PrimTS decode workspace was not prepared.") + + attn = params.attn + meta = params.meta + fwd = params.fwd + batch_size = params.batch_size + ( + q_processed, + kv_pool, + block_tables, + _kv_scale_pool, + _bmm1_scale, + _bmm2_scale, + fmha_workspace, + _cu_seqlens, + _max_q_len, + _max_kv_len, + window_left, + is_multi_token_gen, + ) = self._run_generation_preprocess(params) if fmha_workspace.numel() != 0: raise RuntimeError("PrimTS generation preprocessing returned an FMHA workspace.") if is_multi_token_gen: diff --git a/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py new file mode 100644 index 000000000000..b8596041572a --- /dev/null +++ b/tensorrt_llm/_torch/attention/backends/fmha/prims_ts_block_sparse.py @@ -0,0 +1,584 @@ +# 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. + +"""TRT-LLM FMHA adapter for the vendored PrimTS block-sparse kernels.""" + +import math +from dataclasses import dataclass, fields +from typing import TYPE_CHECKING, Literal, cast + +import torch + +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + AttentionInputType, + PredefinedAttentionMask, +) +from tensorrt_llm._torch.attention.backends.prims_ts._block_sparse.config import ( + _validate_block_sparse_static_profile, +) +from tensorrt_llm._torch.attention.backends.sparse.params import BlockSparseForwardInputs +from tensorrt_llm.logger import logger + +from .interface import FmhaPhase +from .phased import FmhaParams +from .prims_ts import ( + PrimsTSFmha, + get_attention_feature_unsupported_reason, + get_paged_kv_policy_unsupported_reason, + get_paged_kv_storage_unsupported_reason, +) +from .utils import get_kv_page_offset + +if TYPE_CHECKING: + from tensorrt_llm._torch.attention.backends.prims_ts import ( + BlockSparsePagedTSWrapper, + BlockSparseTSWrapper, + ) + from tensorrt_llm._torch.attention.backends.trtllm import ( + TrtllmAttention, + TrtllmAttentionMetadata, + ) + +try: + from tensorrt_llm._torch.attention.backends.prims_ts import ( + BlockSparsePagedTSWrapper as _BlockSparsePagedTSWrapper, + ) + from tensorrt_llm._torch.attention.backends.prims_ts import ( + BlockSparseTSWrapper as _BlockSparseTSWrapper, + ) +except (ImportError, OSError): + _BlockSparseTSWrapper = _BlockSparsePagedTSWrapper = None + + +@dataclass(frozen=True, slots=True) +class _BlockSparsePlanKey: + """Static wrapper profile shared by compatible attention layers. + + The key is the single description of a plan: support checks validate it + against the kernel library and the wrapper cache plans from it. + """ + + device: torch.device + batch_size: int + seq_len_q: int + kv_capacity: int + num_heads: int + num_kv_heads: int + head_dim: int + page_size: int | None + q_block_size: int + kv_block_size: int + max_blocks_per_row: int + mask_type: Literal["dense", "causal"] + dtype: torch.dtype + use_kv_valid_bits: bool + sparse_format: Literal["bsr", "bitmask"] + use_proxy_routes: bool + + def unsupported_reason(self) -> str | None: + try: + _validate_block_sparse_static_profile( + batch_size=self.batch_size, + seq_len_q=self.seq_len_q, + seq_len_kv=self.kv_capacity, + num_qo_heads=self.num_heads, + num_kv_heads=self.num_kv_heads, + head_dim=self.head_dim, + q_block_size=self.q_block_size, + kv_block_size=self.kv_block_size, + use_kv_valid_bits=self.use_kv_valid_bits, + mask_type=self.mask_type, + q_dtype=self.dtype, + kv_dtype=self.dtype, + output_dtype=self.dtype, + max_blocks_per_row=self.max_blocks_per_row, + page_size=self.page_size, + ) + except (ValueError, NotImplementedError, OverflowError) as error: + return str(error) + return None + + def plan(self) -> "BlockSparseTSWrapper | BlockSparsePagedTSWrapper": + paged = self.page_size is not None + wrapper_type = _BlockSparsePagedTSWrapper if paged else _BlockSparseTSWrapper + assert wrapper_type is not None + wrapper = wrapper_type() + plan_args = ( + self.batch_size, + self.seq_len_q, + self.kv_capacity, + self.num_heads, + self.num_kv_heads, + self.head_dim, + self.q_block_size, + self.kv_block_size, + ) + plan_kwargs = { + "device": self.device, + "max_blocks_per_row": self.max_blocks_per_row, + "use_kv_valid_bits": self.use_kv_valid_bits, + "mask_type": self.mask_type, + "q_data_type": self.dtype, + "kv_data_type": self.dtype, + "o_data_type": self.dtype, + } + if paged: + plan_args += (self.page_size,) + else: + plan_kwargs.update( + sparse_format=self.sparse_format, + use_proxy_routes=self.use_proxy_routes, + ) + wrapper.plan(*plan_args, **plan_kwargs) + return wrapper + + +def _get_block_sparse_inputs( + forward_args: AttentionForwardArgs, +) -> BlockSparseForwardInputs | None: + return forward_args.sparse_runtime_params.block_sparse_inputs + + +def _has_other_sparse_runtime(forward_args: AttentionForwardArgs) -> bool: + """Whether the runtime carrier holds any sparse state besides block-sparse routes.""" + params = forward_args.sparse_runtime_params + for field in fields(params): + if field.name == "block_sparse_inputs": + continue + value = getattr(params, field.name) + if isinstance(value, torch.Tensor) or (value is not None and value != 0): + return True + return False + + +def _route_batch_size(inputs: BlockSparseForwardInputs) -> int: + routes = inputs.block_indptr if inputs.sparse_format == "bsr" else inputs.exact_block_bits + return int(routes.shape[0]) + + +def _uniform_seq_len_q( + q: torch.Tensor, + metadata: "TrtllmAttentionMetadata", + batch_size: int, +) -> int | None: + """Return the fixed per-request query length, or ``None`` if the batch is ragged.""" + seq_lens = metadata.seq_lens + if batch_size <= 0 or q.shape[0] % batch_size: + return None + if seq_lens is None or seq_lens.numel() < batch_size: + return None + seq_len_q = int(q.shape[0]) // batch_size + if not bool(seq_lens[:batch_size].eq(seq_len_q).all()): + return None + return seq_len_q + + +class PrimsTSBlockSparseFmha(PrimsTSFmha): + """Contiguous context and fixed-Q paged generation block-sparse FMHA.""" + + def __init__(self, attn: "TrtllmAttention") -> None: + super().__init__(attn) + self.bind_plan_cache({}) + + def bind_plan_cache(self, cache_state: dict[str, object]) -> None: + """Bind wrapper plans to an explicitly owned cache.""" + + self._contiguous_wrappers = cast( + dict[_BlockSparsePlanKey, "BlockSparseTSWrapper"], + cache_state.setdefault("contiguous_wrappers", {}), + ) + self._paged_wrappers = cast( + dict[_BlockSparsePlanKey, "BlockSparsePagedTSWrapper"], + cache_state.setdefault("paged_wrappers", {}), + ) + + @classmethod + def is_available(cls, attn: "TrtllmAttention") -> bool: + return ( + _BlockSparseTSWrapper is not None + and _BlockSparsePagedTSWrapper is not None + and super().is_available(attn) + ) + + def is_supported( + self, + q: torch.Tensor, + k: torch.Tensor | None, + v: torch.Tensor | None, + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + *, + phase: FmhaPhase | None = None, + ) -> bool: + supported, reason = self._is_supported_with_reason( + q, + k, + v, + metadata, + forward_args, + phase=phase, + ) + if not supported: + logger.debug(f"PrimTS block-sparse FMHA does not support request: {reason}") + return supported + + def _is_supported_with_reason( + self, + q: torch.Tensor, + k: torch.Tensor | None, + v: torch.Tensor | None, + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + *, + phase: FmhaPhase | None = None, + ) -> tuple[bool, str]: + reason = self._common_unsupported_reason(metadata, forward_args) + if reason is None: + paged = metadata.kv_cache_manager is not None + expected_phase = FmhaPhase.GENERATION if paged else FmhaPhase.CONTEXT + if phase not in (None, expected_phase): + storage = "paged" if paged else "contiguous" + reason = ( + f"{storage} block-sparse attention only supports the " + f"{expected_phase.name.lower()} phase" + ) + elif paged: + reason = self._paged_unsupported_reason(q, metadata, forward_args) + else: + reason = self._contiguous_unsupported_reason(q, k, v, metadata, forward_args) + return reason is None, reason or "" + + def _common_unsupported_reason( + self, + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + ) -> str | None: + """Gates shared by the contiguous and paged block-sparse paths.""" + if _get_block_sparse_inputs(forward_args) is None: + return "block-sparse forward inputs are required" + if metadata.is_cross: + return "cross attention is not supported" + if self.attn.is_mla_enable: + return "MLA is not supported" + if metadata.num_sparse_topk > 0 or _has_other_sparse_runtime(forward_args): + return "legacy sparse attention cannot be combined with block-sparse inputs" + feature_reason = get_attention_feature_unsupported_reason(metadata, forward_args) + if feature_reason is not None: + return feature_reason + if forward_args.softmax_stats_tensor is not None: + return "softmax statistics output is not supported" + if ( + forward_args.output_sf is not None + or forward_args.out_scale is not None + or forward_args.out_scale_sf is not None + ): + return "quantized output is not supported" + if forward_args.attention_mask not in ( + PredefinedAttentionMask.FULL, + PredefinedAttentionMask.CAUSAL, + ): + return "only full and causal masks are supported" + return None + + def _make_plan_key( + self, + q: torch.Tensor, + inputs: BlockSparseForwardInputs, + *, + batch_size: int, + seq_len_q: int, + kv_capacity: int, + page_size: int | None, + mask_type: Literal["dense", "causal"], + ) -> _BlockSparsePlanKey: + max_blocks_per_row = inputs.max_blocks_per_row + if max_blocks_per_row is None: + max_blocks_per_row = math.ceil(kv_capacity / inputs.kv_block_size) + return _BlockSparsePlanKey( + device=q.device, + batch_size=batch_size, + seq_len_q=seq_len_q, + kv_capacity=kv_capacity, + num_heads=self.attn.num_heads, + num_kv_heads=self.attn.num_kv_heads, + head_dim=self.attn.head_dim, + page_size=page_size, + q_block_size=inputs.q_block_size, + kv_block_size=inputs.kv_block_size, + max_blocks_per_row=max_blocks_per_row, + mask_type=mask_type, + dtype=q.dtype, + use_kv_valid_bits=inputs.kv_valid_bits is not None, + sparse_format=inputs.sparse_format, + use_proxy_routes=inputs.use_proxy_routes, + ) + + def _get_or_plan_wrapper( + self, + key: _BlockSparsePlanKey, + ) -> "BlockSparseTSWrapper | BlockSparsePagedTSWrapper": + cache = self._paged_wrappers if key.page_size is not None else self._contiguous_wrappers + wrapper = cache.get(key) + if wrapper is None: + wrapper = key.plan() + cache[key] = wrapper + return wrapper + + def _contiguous_unsupported_reason( + self, + q: torch.Tensor, + k: torch.Tensor | None, + v: torch.Tensor | None, + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + ) -> str | None: + if forward_args.is_fused_qkv or k is None or v is None: + return "contiguous block-sparse attention requires separate Q, K, and V" + if self.attn.position_embedding_type != 0 or forward_args.mrope_position_deltas is not None: + return "contiguous Q/K/V must have position embedding applied before attention" + if forward_args.cu_q_seqlens is not None or forward_args.cu_kv_seqlens is not None: + return "packed variable-length Q/KV inputs are not supported" + inputs = _get_block_sparse_inputs(forward_args) + assert inputs is not None + mask_type = self._get_prims_mask_type(forward_args) + if inputs.use_proxy_routes and mask_type != "dense": + return "block-sparse proxy routes require mask_type='dense'" + batch_size = _route_batch_size(inputs) + seq_len_q = _uniform_seq_len_q(q, metadata, batch_size) + if seq_len_q is None or k.shape[0] % batch_size: + return "query and KV token counts must be batch-uniform over the route batch size" + key = self._make_plan_key( + q, + inputs, + batch_size=batch_size, + seq_len_q=seq_len_q, + kv_capacity=int(k.shape[0]) // batch_size, + page_size=None, + mask_type=mask_type, + ) + return key.unsupported_reason() + + def _paged_unsupported_reason( + self, + q: torch.Tensor, + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + ) -> str | None: + inputs = _get_block_sparse_inputs(forward_args) + assert inputs is not None + if inputs.sparse_format != "bsr" or inputs.use_proxy_routes: + return "paged block-sparse attention only supports BSR exact routes" + if not forward_args.is_fused_qkv: + return "paged block-sparse attention requires fused QKV input" + if ( + forward_args.attention_input_type != AttentionInputType.generation_only + or metadata.num_contexts != 0 + ): + return "only generation-only paged requests are supported" + reason = get_paged_kv_storage_unsupported_reason( + self.attn, metadata + ) or get_paged_kv_policy_unsupported_reason(self.attn, metadata) + if reason is not None: + return reason + if metadata.tokens_per_block not in self.SUPPORTED_PAGE_SIZES: + return f"page size {metadata.tokens_per_block} is unsupported" + if self.attn.attention_chunk_size: + return "chunked attention is not supported" + if get_kv_page_offset(self.attn, metadata, 0, cache=self._kv_page_offset_cache) is None: + return "the K-to-V page displacement could not be resolved" + + batch_size = int(metadata.num_generations) + seq_len_q = _uniform_seq_len_q(q, metadata, batch_size) + if seq_len_q is None: + return "query lengths must be batch-uniform and match the fixed query shape" + block_tables = metadata.kv_cache_block_offsets + if block_tables.shape[1] < batch_size: + return "paged KV-cache block offsets must cover the generation batch" + page_size = int(metadata.tokens_per_block) + kv_capacity = int(block_tables.shape[-1]) * page_size + logical_max_seq_len = int(metadata.max_seq_len) + if logical_max_seq_len > kv_capacity: + return "logical maximum sequence length must fit the page-table capacity" + attention_window_size = forward_args.attention_window_size + if ( + attention_window_size is None + or attention_window_size < logical_max_seq_len + or attention_window_size > kv_capacity + ): + return "attention window must fit the non-cyclic page-table capacity" + host_seq_lens = metadata.kv_lens_runtime[:batch_size] + min_seq_len_kv = int(host_seq_lens.min()) + if min_seq_len_kv <= 0: + return "every active request must contain at least one KV token" + mask_type = self._get_prims_mask_type(forward_args) + if mask_type == "causal" and min_seq_len_kv < seq_len_q: + return "causal KV lengths must be at least the fixed query length" + if int(host_seq_lens.max()) > logical_max_seq_len: + return "an active KV length exceeds the logical maximum sequence length" + key = self._make_plan_key( + q, + inputs, + batch_size=batch_size, + seq_len_q=seq_len_q, + kv_capacity=kv_capacity, + page_size=page_size, + mask_type=mask_type, + ) + return key.unsupported_reason() + + def prepare_workspace( + self, + q: torch.Tensor, + k: torch.Tensor | None, + v: torch.Tensor | None, + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + workspace: torch.Tensor, + ) -> None: + del k, v, forward_args + with torch.cuda.device(q.device): + # Contiguous requests run without a KV cache and never touch the + # generation preprocessing workspace. + if metadata.kv_cache_manager is not None: + layout = self._get_generation_workspace_layout( + q.dtype, + int(metadata.num_generations), + int(q.shape[0]), + ) + required_bytes = int(layout["total_size"]) + if workspace.numel() * workspace.element_size() < required_bytes: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "TRT-LLM QKV preprocessing workspace must be sized before " + "CUDA Graph capture" + ) + workspace.resize_((math.ceil(required_bytes / workspace.element_size()),)) + if self._multi_processor_count is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("GPU properties must be prepared before CUDA Graph capture") + self._multi_processor_count = torch.cuda.get_device_properties( + q.device + ).multi_processor_count + + def run_generation(self, params: FmhaParams) -> None: + q = params.qkv_input + output_buffer = params.context_buf + sequence_lengths = params.sequence_lengths + assert q is not None and output_buffer is not None and sequence_lengths is not None + metadata = params.meta + forward_args = params.fwd + inputs = _get_block_sparse_inputs(forward_args) + assert inputs is not None + batch_size = params.num_requests + seq_len_q = params.input_seq_length + page_size = params.tokens_per_block + block_offsets = metadata.kv_cache_block_offsets + assert block_offsets is not None + preprocess = self._run_generation_preprocess(params) + q_processed, kv_pool, block_tables = preprocess[:3] + fmha_workspace = preprocess[6] + if fmha_workspace.numel() != 0: + raise RuntimeError("PrimTS block-sparse preprocessing returned an FMHA workspace.") + if q_processed is None or kv_pool is None or block_tables is None: + raise RuntimeError("TRT-LLM preprocessing did not return paged PrimTS metadata") + kv_page_offset = get_kv_page_offset( + params.attn, + metadata, + params.seq_offset, + cache=self._kv_page_offset_cache, + ) + if kv_page_offset is None: + raise RuntimeError("PrimTS could not resolve the K-to-V page displacement") + k_cache, v_cache = self._standard_kv_views(kv_pool, kv_page_offset) + query = q_processed.view( + batch_size, + seq_len_q, + self.attn.num_heads, + self.attn.head_dim, + ) + key = self._make_plan_key( + query, + inputs, + batch_size=batch_size, + seq_len_q=seq_len_q, + kv_capacity=int(block_offsets.shape[-1]) * page_size, + page_size=page_size, + mask_type=self._get_prims_mask_type(forward_args), + ) + wrapper = cast("BlockSparsePagedTSWrapper", self._get_or_plan_wrapper(key)) + wrapper.run( + query, + (k_cache, v_cache), + block_tables=self._get_fixed_block_tables(block_tables, batch_size), + seq_lens_kv=self._get_sequence_lengths(sequence_lengths, batch_size), + block_indptr=inputs.block_indptr, + block_indices=inputs.block_indices, + kv_valid_bits=inputs.kv_valid_bits, + sm_scale=self._get_bmm1_scale(self.attn), + out=output_buffer.view_as(query), + ) + + def _forward_contiguous( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + forward_args: AttentionForwardArgs, + ) -> None: + inputs = _get_block_sparse_inputs(forward_args) + assert inputs is not None + assert forward_args.output is not None + batch_size = _route_batch_size(inputs) + query = q.view(batch_size, -1, self.attn.num_heads, self.attn.head_dim) + key_states = k.view(batch_size, -1, self.attn.num_kv_heads, self.attn.head_dim) + value_states = v.view_as(key_states) + key = self._make_plan_key( + query, + inputs, + batch_size=batch_size, + seq_len_q=int(query.shape[1]), + kv_capacity=int(key_states.shape[1]), + page_size=None, + mask_type=self._get_prims_mask_type(forward_args), + ) + wrapper = cast("BlockSparseTSWrapper", self._get_or_plan_wrapper(key)) + wrapper.run( + query, + key_states, + value_states, + block_indptr=inputs.block_indptr, + block_indices=inputs.block_indices, + exact_block_bits=inputs.exact_block_bits, + k_summary=inputs.k_summary, + v_summary=inputs.v_summary, + kv_valid_bits=inputs.kv_valid_bits, + sm_scale=self._get_bmm1_scale(self.attn), + out=forward_args.output.view_as(query), + ) + + def forward( + self, + q: torch.Tensor, + k: torch.Tensor | None, + v: torch.Tensor | None, + metadata: "TrtllmAttentionMetadata", + forward_args: AttentionForwardArgs, + ) -> None: + if metadata.kv_cache_manager is None: + assert k is not None and v is not None + self._forward_contiguous(q, k, v, forward_args) + return + super().forward(q, k, v, metadata, forward_args) diff --git a/tensorrt_llm/_torch/attention/backends/fmha/registry.py b/tensorrt_llm/_torch/attention/backends/fmha/registry.py index 4e545924e63e..240e5242feb0 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/registry.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/registry.py @@ -35,6 +35,7 @@ def init_fmha_libs() -> dict[str, "FmhaCls"]: """ from .flashinfer_sparse_mla import FlashInferSparseMlaFmha from .msa_sparse_gqa import MsaSparseGqaFmha + from .prims_ts_block_sparse import PrimsTSBlockSparseFmha return { "triton_custom_mask": TritonCustomMaskFmha, @@ -42,6 +43,7 @@ def init_fmha_libs() -> dict[str, "FmhaCls"]: "msa_sparse_gqa": MsaSparseGqaFmha, "flashinfer_sparse_mla": FlashInferSparseMlaFmha, "prims_ts": PrimsTSFmha, + "prims_ts_block_sparse": PrimsTSBlockSparseFmha, "flashinfer_trtllm_gen": FlashInferTrtllmGenFmha, "fallback": FallbackFmha, } diff --git a/tensorrt_llm/_torch/attention/backends/fmha/triton_custom_mask.py b/tensorrt_llm/_torch/attention/backends/fmha/triton_custom_mask.py index 4391ef94bb22..61aef90d9d61 100644 --- a/tensorrt_llm/_torch/attention/backends/fmha/triton_custom_mask.py +++ b/tensorrt_llm/_torch/attention/backends/fmha/triton_custom_mask.py @@ -105,6 +105,8 @@ def _check_support_with_reason( *, phase: Optional[FmhaPhase] = None, ) -> tuple[bool, str]: + if forward_args.sparse_runtime_params.block_sparse_inputs is not None: + return False, "block_sparse_inputs are not supported." if phase != FmhaPhase.CONTEXT: return False, "Only context attention is supported." if forward_args.attention_mask != CustomAttentionMask.CUSTOM: diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/common.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/common.py index a82deafc7a3f..13cf6038efd9 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/common.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/common.py @@ -22,6 +22,20 @@ _BLOCK_SPARSE_MAX_HEADS_Q_PER_KV = 32 +def _validate_contiguous_route_mode( + sparse_format: object, + use_proxy_routes: object, +) -> None: + """Validate the two public continuous-route axes before device work.""" + + if not isinstance(sparse_format, str): + raise TypeError("sparse_format must be 'bsr' or 'bitmask'") + if sparse_format not in ("bsr", "bitmask"): + raise ValueError("sparse_format must be 'bsr' or 'bitmask'") + if type(use_proxy_routes) is not bool: + raise TypeError("use_proxy_routes must be a bool") + + def _validate_sparse_q_block_size(value: object) -> int: """Return a positive semantic Q block size representable by the ABI.""" @@ -110,6 +124,17 @@ def _block_sparse_kv_atom_size(kv_block_size: int) -> int: ) +def _block_sparse_proxy_summary_geometry( + seq_len_kv: int, + kv_block_size: int, +) -> tuple[int, int]: + """Return the summary count and final summary's represented-token mass.""" + + num_summaries = (seq_len_kv + kv_block_size - 1) // kv_block_size + tail_mass = seq_len_kv - (num_summaries - 1) * kv_block_size + return num_summaries, tail_mass + + def _prepared_kv_routes_are_block_aligned( kv_block_size: int, kv_route_size: int, @@ -117,3 +142,24 @@ def _prepared_kv_routes_are_block_aligned( """Return whether each prepared route stays within one semantic BSR block.""" return _validate_sparse_kv_block_size(kv_block_size) % kv_route_size == 0 + + +def _block_sparse_contiguous_kv_copy_geometry( + *, + kv_block_size: int, + kv_route_size: int, +) -> tuple[int, int, bool]: + """Return source-independent primary/atom TensorMap geometry. + + Exact and proxy routes address different logical matrices, but a route's + physical copy shape depends only on its semantic block and physical route + sizes. Coarse routes prefer KV128 copies and keep a KV64 descriptor only + when KV256 staging or runtime adjacency requires it. + """ + + atom_size = _block_sparse_kv_atom_size(kv_block_size) + primary_box_size = 2 * atom_size if atom_size == 64 else atom_size + needs_aux_atom = atom_size == 64 and ( + kv_route_size == 256 or kv_block_size % kv_route_size != 0 + ) + return primary_box_size, atom_size, needs_aux_atom diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/compiler.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/compiler.py index a73f26060753..fcaa0735436f 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/compiler.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/compiler.py @@ -36,31 +36,46 @@ def _compile_block_sparse(key: _BlockSparseCompileKey) -> Callable[..., object]: from ..kernels.fmha_decode.fmha_decode_config import FmhaDecodeConfig from ..kernels.fmha_decode.block_sparse_prepare import ( - _PrepareBlockSparseRoutes, + _PrepareBitmaskRoutes, + _PrepareBsrRoutes, ) from ..kernels.fmha_decode.fmha_decode_kernel import ( fmha_block_sparse_launch, ) config = _make_block_sparse_config(key) - prepare_routes = _PrepareBlockSparseRoutes( - batch_size=key.batch_size, - num_kv_heads=key.num_kv_heads, - seq_len_q=key.seq_len_q, - seq_len_kv=key.seq_len_kv, - q_block_size=key.q_block_size, - kv_block_size=key.kv_block_size, - kv_route_size=key.kv_route_size, - has_token_bits=key.use_kv_valid_bits, - page_size=key.page_size, - mask_type=key.mask_type, - ) + prepare_kwargs = { + "batch_size": key.batch_size, + "num_kv_heads": key.num_kv_heads, + "seq_len_q": key.seq_len_q, + "seq_len_kv": key.seq_len_kv, + "q_block_size": key.q_block_size, + "kv_block_size": key.kv_block_size, + "kv_route_size": key.kv_route_size, + "use_proxy_routes": key.use_proxy_routes, + "use_causal_mask": key.mask_type == "causal", + "apply_token_mask": key.use_kv_valid_bits, + "store_score_words": config.uses_prepared_score_keep_words, + } + if key.page_size is not None: + if key.sparse_format != "bsr" or key.use_proxy_routes: + raise AssertionError("paged block-sparse supports exact BSR routes only") + prepare_kwargs["page_size"] = key.page_size + + if key.sparse_format == "bsr": + prepare_routes = _PrepareBsrRoutes(**prepare_kwargs) + elif key.sparse_format == "bitmask": + prepare_routes = _PrepareBitmaskRoutes(**prepare_kwargs) + else: + raise AssertionError("sparse_format must be 'bsr' or 'bitmask'") + + route_metadata_base = prepare_routes.route_metadata_base_word_offset Int32 = cutlass.Int32 Int64 = cutlass.Int64 Float32 = cutlass.Float32 @cute.jit - def contiguous_tensor_adapter( + def exact_bsr_adapter( q: cute.Tensor, k: cute.Tensor, v: cute.Tensor, @@ -86,7 +101,7 @@ def contiguous_tensor_adapter( kv_valid_bits, None, None, - None, + Int64(0), Int64(0), row_route_offsets, route_workspace, @@ -95,9 +110,7 @@ def contiguous_tensor_adapter( ) # Live per-row route counts occupy the first words of run scratch. row_route_counts = route_workspace.iterator - route_metadata = route_workspace.iterator + Int32( - prepare_routes.route_metadata_base_word_offset - ) + route_metadata = route_workspace.iterator + Int32(route_metadata_base) fmha_block_sparse_launch( ( Int32(static_batch_size), @@ -109,6 +122,8 @@ def contiguous_tensor_adapter( q.iterator, k.iterator, v.iterator, + k.iterator, + v.iterator, out.iterator, row_route_offsets.iterator, row_route_counts, @@ -119,6 +134,169 @@ def contiguous_tensor_adapter( static_seq_len_kv, ) + @cute.jit + def exact_bitmask_adapter( + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + out: cute.Tensor, + exact_block_bits: cute.Tensor, + kv_valid_bits: cute.Tensor, + row_route_offsets: cute.Tensor, + route_workspace: cute.Tensor, + max_blocks_per_row: cutlass.Int32, + sm_scale: cutlass.Float32, + stream: cuda_drv.CUstream, + static_config: cutlass.Constexpr[FmhaDecodeConfig], + static_batch_size: cutlass.Constexpr[int], + static_seq_len_kv: cutlass.Constexpr[int], + static_num_qo_heads: cutlass.Constexpr[int], + static_num_kv_heads: cutlass.Constexpr[int], + static_head_dim: cutlass.Constexpr[int], + ) -> None: + prepare_routes( + exact_block_bits, + kv_valid_bits, + row_route_offsets, + route_workspace, + max_blocks_per_row, + stream, + ) + fmha_block_sparse_launch( + ( + Int32(static_batch_size), + Int32(static_num_qo_heads), + Int32(static_num_kv_heads), + Int32(static_seq_len_kv), + Int32(static_head_dim), + ), + q.iterator, + k.iterator, + v.iterator, + k.iterator, + v.iterator, + out.iterator, + row_route_offsets.iterator, + route_workspace.iterator, + route_workspace.iterator + Int32(route_metadata_base), + sm_scale, + stream, + static_config, + static_seq_len_kv, + ) + + @cute.jit + def proxy_bsr_adapter( + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + k_summary: cute.Tensor, + v_summary: cute.Tensor, + out: cute.Tensor, + block_indptr: cute.Tensor, + block_indices: cute.Tensor, + kv_valid_bits: cute.Tensor, + row_route_offsets: cute.Tensor, + route_workspace: cute.Tensor, + max_blocks_per_row: cutlass.Int32, + sm_scale: cutlass.Float32, + stream: cuda_drv.CUstream, + static_config: cutlass.Constexpr[FmhaDecodeConfig], + static_batch_size: cutlass.Constexpr[int], + static_seq_len_kv: cutlass.Constexpr[int], + static_num_qo_heads: cutlass.Constexpr[int], + static_num_kv_heads: cutlass.Constexpr[int], + static_head_dim: cutlass.Constexpr[int], + ) -> None: + prepare_routes( + block_indptr, + block_indices, + kv_valid_bits, + None, + None, + Int64(0), + Int64(0), + row_route_offsets, + route_workspace, + max_blocks_per_row, + stream, + ) + fmha_block_sparse_launch( + ( + Int32(static_batch_size), + Int32(static_num_qo_heads), + Int32(static_num_kv_heads), + Int32(static_seq_len_kv), + Int32(static_head_dim), + ), + q.iterator, + k.iterator, + v.iterator, + k_summary.iterator, + v_summary.iterator, + out.iterator, + row_route_offsets.iterator, + route_workspace.iterator, + route_workspace.iterator + Int32(route_metadata_base), + sm_scale, + stream, + static_config, + static_seq_len_kv, + ) + + @cute.jit + def proxy_bitmask_adapter( + q: cute.Tensor, + k: cute.Tensor, + v: cute.Tensor, + k_summary: cute.Tensor, + v_summary: cute.Tensor, + out: cute.Tensor, + exact_block_bits: cute.Tensor, + kv_valid_bits: cute.Tensor, + row_route_offsets: cute.Tensor, + route_workspace: cute.Tensor, + max_blocks_per_row: cutlass.Int32, + sm_scale: cutlass.Float32, + stream: cuda_drv.CUstream, + static_config: cutlass.Constexpr[FmhaDecodeConfig], + static_batch_size: cutlass.Constexpr[int], + static_seq_len_kv: cutlass.Constexpr[int], + static_num_qo_heads: cutlass.Constexpr[int], + static_num_kv_heads: cutlass.Constexpr[int], + static_head_dim: cutlass.Constexpr[int], + ) -> None: + prepare_routes( + exact_block_bits, + kv_valid_bits, + row_route_offsets, + route_workspace, + max_blocks_per_row, + stream, + ) + fmha_block_sparse_launch( + ( + Int32(static_batch_size), + Int32(static_num_qo_heads), + Int32(static_num_kv_heads), + Int32(static_seq_len_kv), + Int32(static_head_dim), + ), + q.iterator, + k.iterator, + v.iterator, + k_summary.iterator, + v_summary.iterator, + out.iterator, + row_route_offsets.iterator, + route_workspace.iterator, + route_workspace.iterator + Int32(route_metadata_base), + sm_scale, + stream, + static_config, + static_seq_len_kv, + ) + @cute.jit def paged_tensor_adapter( q: cute.Tensor, @@ -128,13 +306,13 @@ def paged_tensor_adapter( block_indptr: cute.Tensor, block_indices: cute.Tensor, kv_valid_bits: cute.Tensor, - paged_kv_indptr: cute.Tensor, - paged_kv_indices: cute.Tensor, + block_tables: cute.Tensor, seq_lens_kv: cute.Tensor, row_route_offsets: cute.Tensor, route_workspace: cute.Tensor, max_blocks_per_row: cutlass.Int32, num_physical_kv_pages: cutlass.Int64, + block_table_row_stride: cutlass.Int64, k_page_stride: cutlass.Int64, v_page_stride: cutlass.Int64, sm_scale: cutlass.Float32, @@ -151,18 +329,16 @@ def paged_tensor_adapter( block_indices, kv_valid_bits, seq_lens_kv, - paged_kv_indptr, - paged_kv_indices, + block_tables, num_physical_kv_pages, + block_table_row_stride, row_route_offsets, route_workspace, max_blocks_per_row, stream, ) row_route_counts = route_workspace.iterator - route_metadata = route_workspace.iterator + Int32( - prepare_routes.route_metadata_base_word_offset - ) + route_metadata = route_workspace.iterator + Int32(route_metadata_base) fmha_block_sparse_launch( ( Int32(static_batch_size), @@ -174,6 +350,8 @@ def paged_tensor_adapter( q.iterator, k_cache.iterator, v_cache.iterator, + k_cache.iterator, + v_cache.iterator, out.iterator, row_route_offsets.iterator, row_route_counts, @@ -203,6 +381,7 @@ def fake_compact( logical_workspace_words = cute.sym_int() q_shape = (key.batch_size, key.seq_len_q, key.num_qo_heads, key.head_dim) num_q_blocks = ceil_div(key.seq_len_q, key.q_block_size) + num_kv_blocks = ceil_div(key.seq_len_kv, key.kv_block_size) indptr_fake = fake_compact( Int32, (key.batch_size, key.num_kv_heads, num_q_blocks + 1), @@ -238,24 +417,96 @@ def fake_compact( ) k_fake = fake_compact(config.kv_dtype, kv_shape, 16) v_fake = fake_compact(config.kv_dtype, kv_shape, 16) - tensor_adapter = contiguous_tensor_adapter - dynamic_args = ( - q_fake, - k_fake, - v_fake, - out_fake, - indptr_fake, - indices_fake, + exact_bits_fake = fake_compact( + cutlass.Uint32, + ( + key.batch_size, + key.num_kv_heads, + num_q_blocks, + ceil_div(num_kv_blocks, 32), + ), + 4, + ) + common_tail = ( valid_bits_fake, row_route_offsets_fake, route_workspace_fake, Int32(0), Float32(1.0), ) + if key.sparse_format == "bsr" and not key.use_proxy_routes: + tensor_adapter = exact_bsr_adapter + dynamic_args = ( + q_fake, + k_fake, + v_fake, + out_fake, + indptr_fake, + indices_fake, + *common_tail, + ) + elif key.sparse_format == "bitmask" and not key.use_proxy_routes: + tensor_adapter = exact_bitmask_adapter + dynamic_args = ( + q_fake, + k_fake, + v_fake, + out_fake, + exact_bits_fake, + *common_tail, + ) + elif key.sparse_format == "bsr" and key.use_proxy_routes: + summary_shape = ( + key.batch_size, + num_kv_blocks, + key.num_kv_heads, + key.head_dim, + ) + k_summary_fake = fake_compact(config.kv_dtype, summary_shape, 16) + v_summary_fake = fake_compact(config.kv_dtype, summary_shape, 16) + proxy_prefix = ( + q_fake, + k_fake, + v_fake, + k_summary_fake, + v_summary_fake, + out_fake, + ) + tensor_adapter = proxy_bsr_adapter + dynamic_args = ( + *proxy_prefix, + indptr_fake, + indices_fake, + *common_tail, + ) + elif key.sparse_format == "bitmask" and key.use_proxy_routes: + summary_shape = ( + key.batch_size, + num_kv_blocks, + key.num_kv_heads, + key.head_dim, + ) + k_summary_fake = fake_compact(config.kv_dtype, summary_shape, 16) + v_summary_fake = fake_compact(config.kv_dtype, summary_shape, 16) + tensor_adapter = proxy_bitmask_adapter + dynamic_args = ( + q_fake, + k_fake, + v_fake, + k_summary_fake, + v_summary_fake, + out_fake, + exact_bits_fake, + *common_tail, + ) + else: + raise AssertionError("continuous sparse_format must be 'bsr' or 'bitmask'") else: page_size = key.page_size + assert page_size is not None physical_pages = cute.sym_int() - logical_pages = cute.sym_int() + runtime_page_columns = cute.sym_int() + runtime_page_row_stride = cute.sym_int64(divisibility=1) k_outer_stride = cute.sym_int64(divisibility=1) v_outer_stride = cute.sym_int64(divisibility=1) kv_shape = ( @@ -286,12 +537,12 @@ def fake_compact( ), assumed_align=16, ) - paged_kv_indptr_fake = fake_compact( + block_tables_fake = cute.runtime.make_fake_tensor( Int32, - (key.batch_size + 1,), - 4, + (key.batch_size, runtime_page_columns), + stride=(runtime_page_row_stride, 1), + assumed_align=4, ) - paged_kv_indices_fake = fake_compact(Int32, (logical_pages,), 4) seq_lens_kv_fake = fake_compact(Int32, (key.batch_size,), 4) tensor_adapter = paged_tensor_adapter dynamic_args = ( @@ -302,8 +553,7 @@ def fake_compact( indptr_fake, indices_fake, valid_bits_fake, - paged_kv_indptr_fake, - paged_kv_indices_fake, + block_tables_fake, seq_lens_kv_fake, row_route_offsets_fake, route_workspace_fake, @@ -311,6 +561,7 @@ def fake_compact( Int64(1), Int64(1), Int64(1), + Int64(1), Float32(1.0), ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/config.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/config.py index 47d7c60340b5..cc4458a4298a 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/config.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/config.py @@ -64,7 +64,7 @@ @dataclass(frozen=True) class _BlockSparseCompileKey: - """Named, hashable inputs that determine one compiled adapter.""" + """Named, hashable inputs that determine one compiled sparse adapter.""" device_index: int batch_size: int @@ -81,6 +81,8 @@ class _BlockSparseCompileKey: use_kv_valid_bits: bool use_persistent_scheduler: bool use_parallel_sparse_kv_loads: bool + sparse_format: Literal["bsr", "bitmask"] = "bsr" + use_proxy_routes: bool = False page_size: int | None = None @@ -90,6 +92,9 @@ class _BlockSparseLaunchSpec: policy: tuple[tuple[str, object], ...] compile_key: _BlockSparseCompileKey + # Whether prepared routes carry K32 score-validity words; the decode + # config owns this rule and the plan sizes its route storage from it. + prepares_score_words: bool _CAPACITY_UNSET = object() @@ -194,12 +199,12 @@ def _select_block_sparse_scheduler( use_kv_valid_bits: bool, max_row_route_capacity: int, ) -> tuple[int, bool]: - """Select the Q tile and scheduler without depending on KV storage.""" + """Select the Q tile and scheduler without depending on KV storage. - from ..kernels.fmha_decode.fmha_decode_config import ( - _select_auto_launch_mode, - make_q_tile_geometry, - ) + Proxy routes add one summary route per row on top of the exact routes and + see the same per-tile fixed cost the persistent scheduler amortizes, so + the selection does not depend on the route kind. + """ heads_q_per_kv = num_qo_heads // num_kv_heads q_tile_size = _select_block_sparse_q_tile_size( @@ -216,6 +221,11 @@ def _select_block_sparse_scheduler( ): return q_tile_size, False + from ..kernels.fmha_decode.fmha_decode_config import ( + _select_auto_launch_mode, + make_q_tile_geometry, + ) + q_geometry = make_q_tile_geometry( rows_per_cta=q_tile_size, heads_q_per_kv=heads_q_per_kv, @@ -352,11 +362,13 @@ def _validate_block_sparse_static_profile( kv_block_size=kv_block_size, ) if page_size is not None: + # Validate the paged route geometry with a capacity-free layout; the + # score-word slots do not take part in the page/atom checks. _BlockSparseRouteLayout.create( kv_route_size=kv_route_size, kv_block_size=kv_block_size, page_size=page_size, - has_token_bits=use_kv_valid_bits, + has_token_bits=False, route_metadata_capacity=0, num_rows=1, ) @@ -413,6 +425,8 @@ def _make_block_sparse_config(key: _BlockSparseCompileKey) -> "FmhaDecodeConfig" } if key.use_persistent_scheduler: config_args["use_persistent_scheduler"] = True + if key.use_proxy_routes: + config_args["use_block_sparse_proxy_routes"] = True layout_args: dict[str, object] if key.page_size is None: layout_args = {"qkv_layout": "contiguousKv"} @@ -455,14 +469,17 @@ def _resolve_block_sparse_launch_spec( mask_type: Literal["dense", "causal"], use_kv_valid_bits: bool, max_row_route_capacity: int, + sparse_format: Literal["bsr", "bitmask"] = "bsr", + use_proxy_routes: bool = False, page_size: int | None = None, ) -> _BlockSparseLaunchSpec: """Resolve and cache one validated static or CLC launch. ``max_row_route_capacity`` is a conservative prepared-route bound. Live index values and physical-tail morphology never specialize this cache - entry. If the selected persistent profile is unsupported, retain the valid - static profile instead. + entry. Proxy and exact routes share one scheduler selection. An + unsupported persistent profile falls back to its valid static + counterpart. """ q_tile_size, use_persistent_scheduler = _select_block_sparse_scheduler( @@ -499,10 +516,12 @@ def _resolve_block_sparse_launch_spec( max_row_route_capacity=max_row_route_capacity, use_persistent_scheduler=use_persistent_scheduler, ), + sparse_format=sparse_format, + use_proxy_routes=use_proxy_routes, page_size=page_size, ) try: - _make_block_sparse_config(compile_key) + config = _make_block_sparse_config(compile_key) except ValueError: if not compile_key.use_persistent_scheduler: raise @@ -516,11 +535,15 @@ def _resolve_block_sparse_launch_spec( use_persistent_scheduler=False, ), ) - _make_block_sparse_config(compile_key) + config = _make_block_sparse_config(compile_key) policy_entries: list[tuple[str, object]] = [ ("tile_size_q", q_tile_size), ("tile_size_kv", kv_route_size), + ( + "scheduler", + "persistent" if compile_key.use_persistent_scheduler else "static", + ), ] if page_size is not None: policy_entries.append(("page_size", page_size)) @@ -538,6 +561,7 @@ def _resolve_block_sparse_launch_spec( return _BlockSparseLaunchSpec( policy=tuple(policy_entries), compile_key=compile_key, + prepares_score_words=config.uses_prepared_score_keep_words, ) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/inspection.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/inspection.py index 39c0c22d6f64..7abf22e52792 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/inspection.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/inspection.py @@ -64,12 +64,7 @@ def _raise_for_invalid_paged_metadata( return reason = { 4: (f"seq_lens_kv values must lie in [{minimum_seq_len_kv}, {max_seq_len_kv}]"), - 5: ( - "paged_kv_indptr must start at zero and each row must be " - "bounded and monotone" - ), - 6: "paged_kv_indptr rows must contain enough pages for seq_lens_kv", - 7: "paged_kv_indices must contain an in-range physical page ID", + 5: "block_tables must contain an in-range physical page ID for every live page", }.get(error_code) if reason is None: reason = ( @@ -152,8 +147,7 @@ def launch(summary: torch.Tensor, device_index: int) -> None: def _inspect_paged_block_sparse_metadata( block_indptr: torch.Tensor, block_indices: torch.Tensor, - paged_kv_indptr: torch.Tensor, - paged_kv_indices: torch.Tensor, + block_tables: torch.Tensor, seq_lens_kv: torch.Tensor, *, static: _BlockSparseStaticProfile, @@ -186,8 +180,8 @@ def launch(summary: torch.Tensor, device_index: int) -> None: inspect_metadata( block_indptr, block_indices, - paged_kv_indptr, - paged_kv_indices, + block_tables, + block_tables.stride(0), seq_lens_kv, num_physical_kv_pages, summary, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/plan.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/plan.py index e70ed415fc97..453e8e35134a 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/plan.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/plan.py @@ -18,13 +18,13 @@ from dataclasses import dataclass import functools import _thread -from typing import Concatenate, ParamSpec, Protocol, TypeVar, cast +from typing import Concatenate, Literal, ParamSpec, Protocol, TypeVar, cast import torch from flashinfer.utils import ceil_div -from .common import _SIGNED_INT32_MAX +from .common import _SIGNED_INT32_MAX, _block_sparse_proxy_summary_geometry from .compiler import _get_compiled_block_sparse from .config import ( _BlockSparseStaticProfile, @@ -85,6 +85,8 @@ class _BlockSparsePlanState: cannot prevent in-place modification or replace graph ownership. """ + sparse_format: Literal["bsr", "bitmask"] + use_proxy_routes: bool device: torch.device batch_size: int seq_len_q: int @@ -93,6 +95,7 @@ class _BlockSparsePlanState: num_kv_heads: int head_dim: int q_block_size: int + kv_block_size: int q_dtype: torch.dtype kv_dtype: torch.dtype output_dtype: torch.dtype @@ -179,27 +182,31 @@ def _build_block_sparse_plan_state( device: torch.device, device_index: int, plan_stream: torch.cuda.Stream, + sparse_format: Literal["bsr", "bitmask"] = "bsr", + use_proxy_routes: bool = False, ) -> _BlockSparsePlanState: - """Build and close one complete state after storage validation.""" + """Build one format- and route-specialized plan atomically.""" assert static.max_blocks_per_row is not None - max_row_route_capacity = ceil_div( - static.max_blocks_per_row * static.kv_block_size, - static.kv_route_size, - ) + if static.page_size is not None: + assert sparse_format == "bsr" and not use_proxy_routes num_rows = ( static.batch_size * static.num_kv_heads * ceil_div(static.seq_len_q, static.q_block_size) ) - route_layout = _BlockSparseRouteLayout.create( - kv_route_size=static.kv_route_size, - kv_block_size=static.kv_block_size, - page_size=static.page_size, - has_token_bits=static.use_kv_valid_bits, - route_metadata_capacity=num_rows * max_row_route_capacity, - num_rows=num_rows, + if num_rows > _SIGNED_INT32_MAX: + raise OverflowError("row_count must fit in signed int32") + max_row_route_capacity = ceil_div( + static.max_blocks_per_row * static.kv_block_size, + static.kv_route_size, ) + if use_proxy_routes: + num_summaries, _ = _block_sparse_proxy_summary_geometry( + static.seq_len_kv, + static.kv_block_size, + ) + max_row_route_capacity += ceil_div(num_summaries, static.kv_route_size) with torch.cuda.device(device_index), torch.cuda.stream(plan_stream): spec = _resolve_block_sparse_launch_spec( device_index=device_index, @@ -217,11 +224,21 @@ def _build_block_sparse_plan_state( mask_type=static.mask_type, use_kv_valid_bits=static.use_kv_valid_bits, max_row_route_capacity=max_row_route_capacity, + sparse_format=sparse_format, + use_proxy_routes=use_proxy_routes, ) policy = ( *spec.policy, ("max_blocks_per_row", static.max_blocks_per_row), ) + route_layout = _BlockSparseRouteLayout.create( + kv_route_size=static.kv_route_size, + kv_block_size=static.kv_block_size, + page_size=static.page_size, + has_token_bits=spec.prepares_score_words, + route_metadata_capacity=num_rows * max_row_route_capacity, + num_rows=num_rows, + ) compiled = _get_compiled_block_sparse(spec.compile_key) dummy_kv_valid_bits = ( None @@ -240,6 +257,8 @@ def _build_block_sparse_plan_state( ready_event = _record_block_sparse_plan_ready_event(plan_stream) return _BlockSparsePlanState( + sparse_format=sparse_format, + use_proxy_routes=use_proxy_routes, device=device, batch_size=static.batch_size, seq_len_q=static.seq_len_q, @@ -248,6 +267,7 @@ def _build_block_sparse_plan_state( num_kv_heads=static.num_kv_heads, head_dim=static.head_dim, q_block_size=static.q_block_size, + kv_block_size=static.kv_block_size, q_dtype=static.q_dtype, kv_dtype=static.kv_dtype, output_dtype=static.output_dtype, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/prepared.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/prepared.py index 091becbd33b8..66815f890a55 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/prepared.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/prepared.py @@ -22,6 +22,7 @@ _SECTION_ALIGNMENT_WORDS = 4 _PREPARED_ROUTE_IS_FULL_FLAG = 1 << 0 +_PREPARED_ROUTE_IS_PROXY_FLAG = 1 << 1 _SUPPORTED_KV_ROUTE_SIZES = (128, 256) _SUPPORTED_PAGED_KV_PAGE_SIZES = (16, 32, 64, 128) @@ -91,13 +92,16 @@ class _BlockSparseRouteLayout: Each route's metadata stores logical KV-token atom origins, optional physical page IDs, one atom-valid-mask word, one route-flags word, and - optional token-valid words. ``page_size is None`` selects the contiguous + optional token-mask words. ``page_size is None`` selects the contiguous record; otherwise the paged record adds one page-ID word per logical origin. Logical origins remain independent of the K/V storage locator used - by the attention load path. An invalid logical origin is encoded as - ``-1``. Bit ``i`` of the atom-valid mask corresponds to logical origin - ``i``. ``_PREPARED_ROUTE_IS_FULL_FLAG`` (bit 0) states that the route is - both structurally full and, when token bits are present, token-full. + by the attention load path. Exact routes address raw-token origins, while + proxy routes address summary-token origins and set + ``_PREPARED_ROUTE_IS_PROXY_FLAG`` (bit 1). An invalid logical origin is + encoded as ``-1``. Bit ``i`` of the atom-valid mask corresponds to logical + origin ``i``. ``_PREPARED_ROUTE_IS_FULL_FLAG`` (bit 0) states that the + route is both structurally full and, when token-mask bits are present, + mask-full. """ # Store semantic inputs plus the three validated allocation values. All @@ -238,6 +242,17 @@ def token_words_word_offset(self) -> int | None: return self.route_flags_word_offset + 1 if self.has_token_bits else None + @property + def uses_one_warp_transport(self) -> bool: + """Whether this layout uses the continuous one-warp transport.""" + + token_words_word_offset = self.token_words_word_offset + return ( + not self.is_paged + and token_words_word_offset is not None + and token_words_word_offset + self.token_words_per_route <= 32 + ) + @property def route_metadata_capacity(self) -> int: """Number of routes whose metadata fits in the mutable workspace.""" @@ -249,5 +264,6 @@ def route_metadata_capacity(self) -> int: __all__ = [ "_PREPARED_ROUTE_IS_FULL_FLAG", + "_PREPARED_ROUTE_IS_PROXY_FLAG", "_BlockSparseRouteLayout", ] diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/runtime.py b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/runtime.py index 43e8ff439d31..4f24524e3d4a 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/runtime.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/_block_sparse/runtime.py @@ -16,7 +16,7 @@ from dataclasses import dataclass import math -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Literal import torch @@ -25,6 +25,7 @@ PagedKVCache, _normalize_paged_kv_cache, _validate_16byte_alignment, + _validate_block_table_metadata, _validate_exact_compact_strides, _validate_scale, ) @@ -47,8 +48,7 @@ class _PagedKVStorage: """Paged K/V storage and request metadata consumed by one live run.""" paged_kv_cache: PagedKVCache - paged_kv_indptr: torch.Tensor - paged_kv_indices: torch.Tensor + block_tables: torch.Tensor seq_lens_kv: torch.Tensor @@ -56,8 +56,8 @@ class _PagedKVStorage: class _PagedKVLaunchPayload: """Launch-only live paged metadata derived during shared validation.""" - paged_kv_indptr: torch.Tensor - paged_kv_indices: torch.Tensor + block_tables: torch.Tensor + block_table_row_stride: int seq_lens_kv: torch.Tensor num_physical_kv_pages: int k_page_stride: int @@ -72,12 +72,15 @@ class _BlockSparseRunArgs: k: torch.Tensor v: torch.Tensor out: torch.Tensor - block_indptr: torch.Tensor - block_indices: torch.Tensor + block_indptr: torch.Tensor | None + block_indices: torch.Tensor | None kv_valid_bits: torch.Tensor kv_valid_bits_is_live: bool sm_scale: float paged_kv: _PagedKVLaunchPayload | None + exact_block_bits: torch.Tensor | None = None + k_summary: torch.Tensor | None = None + v_summary: torch.Tensor | None = None def _validate_metadata_tensor( @@ -113,36 +116,71 @@ def _validate_metadata_tensor( def validate_block_sparse_metadata( - block_indptr: torch.Tensor, - block_indices: torch.Tensor, - kv_valid_bits: torch.Tensor | None, *, + sparse_format: Literal["bsr", "bitmask"], + block_indptr: torch.Tensor | None, + block_indices: torch.Tensor | None, + exact_block_bits: torch.Tensor | None, + kv_valid_bits: torch.Tensor | None, device: torch.device, batch_size: int, seq_len_q: int, seq_len_kv: int, num_kv_heads: int, q_block_size: int, + kv_block_size: int, use_kv_valid_bits: bool, ) -> None: - """Validate raw runtime routing without reading device-side values.""" + """Validate the planned route frontend without reading tensor values.""" num_q_blocks = (seq_len_q + q_block_size - 1) // q_block_size - _validate_metadata_tensor( - block_indptr, - "block_indptr", - ndim=3, - dtype=torch.int32, - expected_device=device, - expected_shape=(batch_size, num_kv_heads, num_q_blocks + 1), - ) - _validate_metadata_tensor( - block_indices, - "block_indices", - ndim=1, - dtype=torch.int32, - expected_device=device, - ) + if sparse_format == "bsr": + if block_indptr is None or block_indices is None: + raise ValueError( + "block_indptr and block_indices are required by a BSR plan" + ) + if exact_block_bits is not None: + raise ValueError("exact_block_bits is valid only for a bitmask plan") + _validate_metadata_tensor( + block_indptr, + "block_indptr", + ndim=3, + dtype=torch.int32, + expected_device=device, + expected_shape=(batch_size, num_kv_heads, num_q_blocks + 1), + ) + _validate_metadata_tensor( + block_indices, + "block_indices", + ndim=1, + dtype=torch.int32, + expected_device=device, + ) + elif sparse_format == "bitmask": + if ( + exact_block_bits is None + or block_indptr is not None + or block_indices is not None + ): + raise ValueError( + "runtime route inputs must match planned sparse_format='bitmask'" + ) + num_kv_blocks = (seq_len_kv + kv_block_size - 1) // kv_block_size + _validate_metadata_tensor( + exact_block_bits, + "exact_block_bits", + ndim=4, + dtype=torch.uint32, + expected_device=device, + expected_shape=( + batch_size, + num_kv_heads, + num_q_blocks, + (num_kv_blocks + 31) // 32, + ), + ) + else: + raise AssertionError(f"unsupported sparse format {sparse_format!r}") if use_kv_valid_bits: if kv_valid_bits is None: @@ -159,31 +197,6 @@ def validate_block_sparse_metadata( raise ValueError("kv_valid_bits must be None when use_kv_valid_bits=False") -def validate_paged_kv_metadata( - paged_kv_indptr: torch.Tensor, - paged_kv_indices: torch.Tensor, - seq_lens_kv: torch.Tensor, - *, - device: torch.device, - batch_size: int, -) -> None: - """Validate the shared structural ABI for live paged request metadata.""" - - for tensor, name, shape in ( - (paged_kv_indptr, "paged_kv_indptr", (batch_size + 1,)), - (paged_kv_indices, "paged_kv_indices", None), - (seq_lens_kv, "seq_lens_kv", (batch_size,)), - ): - _validate_metadata_tensor( - tensor, - name, - ndim=1, - dtype=torch.int32, - expected_device=device, - expected_shape=shape, - ) - - def _validate_bshd_tensor( tensor: torch.Tensor, name: str, @@ -214,8 +227,11 @@ def validate_block_sparse_run( kv_storage: _ContiguousKVStorage | _PagedKVStorage, *, state: "_BlockSparsePlanState", - block_indptr: torch.Tensor, - block_indices: torch.Tensor, + block_indptr: torch.Tensor | None, + block_indices: torch.Tensor | None, + exact_block_bits: torch.Tensor | None = None, + k_summary: torch.Tensor | None = None, + v_summary: torch.Tensor | None = None, kv_valid_bits: torch.Tensor | None, sm_scale: float | None, out: torch.Tensor | None, @@ -229,18 +245,44 @@ def validate_block_sparse_run( input. ``sm_scale=None`` is materialized as ``1 / sqrt(D)``. """ + use_proxy_routes = state.use_proxy_routes + num_kv_blocks = (state.seq_len_kv + state.kv_block_size - 1) // state.kv_block_size validate_block_sparse_metadata( - block_indptr, - block_indices, - kv_valid_bits, + sparse_format=state.sparse_format, + block_indptr=block_indptr, + block_indices=block_indices, + exact_block_bits=exact_block_bits, + kv_valid_bits=kv_valid_bits, device=state.device, batch_size=state.batch_size, seq_len_q=state.seq_len_q, seq_len_kv=state.seq_len_kv, num_kv_heads=state.num_kv_heads, q_block_size=state.q_block_size, + kv_block_size=state.kv_block_size, use_kv_valid_bits=state.use_kv_valid_bits, ) + + if use_proxy_routes: + if k_summary is None or v_summary is None: + raise ValueError("K/V summaries are required when proxy routes are enabled") + summary_shape = ( + state.batch_size, + num_kv_blocks, + state.num_kv_heads, + state.head_dim, + ) + for tensor, name in ((k_summary, "k_summary"), (v_summary, "v_summary")): + _validate_bshd_tensor( + tensor, + name, + expected_shape=summary_shape, + expected_dtype=state.kv_dtype, + expected_device=state.device, + ) + elif k_summary is not None or v_summary is not None: + raise ValueError("summaries are valid only when proxy routes are enabled") + if state.use_kv_valid_bits: assert kv_valid_bits is not None effective_kv_valid_bits = kv_valid_bits @@ -258,7 +300,7 @@ def validate_block_sparse_run( expected_device=state.device, ) paged_kv: _PagedKVLaunchPayload | None = None - overlap_inputs: tuple[tuple[str, torch.Tensor], ...] + overlap_inputs: list[tuple[str, torch.Tensor]] if isinstance(kv_storage, _ContiguousKVStorage): if state.page_size is not None: raise TypeError("contiguous K/V storage requires a contiguous plan state") @@ -278,16 +320,11 @@ def validate_block_sparse_run( ) k = kv_storage.k v = kv_storage.v - overlap_inputs = ( + overlap_inputs = [ ("q", q), ("k", k), ("v", v), - ("block_indptr", block_indptr), - ("block_indices", block_indices), - ("kv_valid_bits", effective_kv_valid_bits), - ("row_route_offsets", state.row_route_offsets), - ("route_workspace", state.route_workspace), - ) + ] elif isinstance(kv_storage, _PagedKVStorage): page_size = state.page_size if page_size is None: @@ -324,36 +361,60 @@ def validate_block_sparse_run( raise ValueError( f"K/V dtype must match the plan ({state.kv_dtype}), got {k.dtype}" ) - validate_paged_kv_metadata( - kv_storage.paged_kv_indptr, - kv_storage.paged_kv_indices, - kv_storage.seq_lens_kv, - device=state.device, - batch_size=state.batch_size, + metadata_device, metadata_batch_size, table_capacity = ( + _validate_block_table_metadata( + kv_storage.block_tables, + kv_storage.seq_lens_kv, + ) ) + if metadata_device != state.device: + raise ValueError( + f"per-run metadata must be on {state.device}, got {metadata_device}" + ) + if metadata_batch_size != state.batch_size: + raise ValueError( + "per-run metadata batch size must match the plan " + f"({state.batch_size}), got {metadata_batch_size}" + ) + if table_capacity * page_size < state.seq_len_kv: + raise ValueError( + "block_tables must cover the planned K/V capacity: expected at " + f"least {(state.seq_len_kv + page_size - 1) // page_size} columns, " + f"got {table_capacity}" + ) paged_kv = _PagedKVLaunchPayload( - paged_kv_indptr=kv_storage.paged_kv_indptr, - paged_kv_indices=kv_storage.paged_kv_indices, + block_tables=kv_storage.block_tables, + block_table_row_stride=kv_storage.block_tables.stride(0), seq_lens_kv=kv_storage.seq_lens_kv, num_physical_kv_pages=num_physical_kv_pages, k_page_stride=k_page_stride, v_page_stride=v_page_stride, ) - overlap_inputs = ( + overlap_inputs = [ ("q", q), ("k_cache", k), ("v_cache", v), - ("block_indptr", block_indptr), - ("block_indices", block_indices), - ("kv_valid_bits", effective_kv_valid_bits), - ("paged_kv_indptr", kv_storage.paged_kv_indptr), - ("paged_kv_indices", kv_storage.paged_kv_indices), + ("block_tables", kv_storage.block_tables), ("seq_lens_kv", kv_storage.seq_lens_kv), + ] + else: + raise TypeError("kv_storage must be _ContiguousKVStorage or _PagedKVStorage") + + if block_indptr is not None and block_indices is not None: + overlap_inputs.extend( + (("block_indptr", block_indptr), ("block_indices", block_indices)) + ) + if exact_block_bits is not None: + overlap_inputs.append(("exact_block_bits", exact_block_bits)) + if k_summary is not None and v_summary is not None: + overlap_inputs.extend((("k_summary", k_summary), ("v_summary", v_summary))) + overlap_inputs.extend( + ( + ("kv_valid_bits", effective_kv_valid_bits), ("row_route_offsets", state.row_route_offsets), ("route_workspace", state.route_workspace), ) - else: - raise TypeError("kv_storage must be _ContiguousKVStorage or _PagedKVStorage") + ) effective_scale = _validate_scale( 1.0 / math.sqrt(state.head_dim) if sm_scale is None else sm_scale, @@ -377,6 +438,9 @@ def validate_block_sparse_run( out=out, block_indptr=block_indptr, block_indices=block_indices, + exact_block_bits=exact_block_bits, + k_summary=k_summary, + v_summary=v_summary, kv_valid_bits=effective_kv_valid_bits, kv_valid_bits_is_live=state.use_kv_valid_bits, sm_scale=effective_scale, @@ -384,23 +448,100 @@ def validate_block_sparse_run( ) +def prepare_block_sparse_run_unchecked( + q: torch.Tensor, + kv_storage: _ContiguousKVStorage | _PagedKVStorage, + *, + state: "_BlockSparsePlanState", + block_indptr: torch.Tensor | None, + block_indices: torch.Tensor | None, + exact_block_bits: torch.Tensor | None = None, + k_summary: torch.Tensor | None = None, + v_summary: torch.Tensor | None = None, + kv_valid_bits: torch.Tensor | None, + sm_scale: float | None, + out: torch.Tensor | None, +) -> _BlockSparseRunArgs: + """Canonicalize one trusted run without invoking explicit validators. + + Only the work every launch needs happens here: K/V view selection, the + plan-owned dummy token mask when the plan disabled token bits, the default + softmax scale, and allocation of an omitted output tensor. + """ + + paged_kv: _PagedKVLaunchPayload | None = None + if isinstance(kv_storage, _ContiguousKVStorage): + k = kv_storage.k + v = kv_storage.v + else: + paged_kv_cache = kv_storage.paged_kv_cache + if isinstance(paged_kv_cache, torch.Tensor): + k = paged_kv_cache[:, 0] + v = paged_kv_cache[:, 1] + else: + k, v = paged_kv_cache + paged_kv = _PagedKVLaunchPayload( + block_tables=kv_storage.block_tables, + block_table_row_stride=kv_storage.block_tables.stride(0), + seq_lens_kv=kv_storage.seq_lens_kv, + num_physical_kv_pages=int(k.shape[0]), + k_page_stride=int(k.stride(0)), + v_page_stride=int(v.stride(0)), + ) + if state.use_kv_valid_bits: + effective_kv_valid_bits = kv_valid_bits + else: + effective_kv_valid_bits = state.dummy_kv_valid_bits + assert effective_kv_valid_bits is not None + if out is None: + out = torch.empty( + (state.batch_size, state.seq_len_q, state.num_qo_heads, state.head_dim), + device=state.device, + dtype=state.output_dtype, + ) + return _BlockSparseRunArgs( + q=q, + k=k, + v=v, + out=out, + block_indptr=block_indptr, + block_indices=block_indices, + exact_block_bits=exact_block_bits, + k_summary=k_summary, + v_summary=v_summary, + kv_valid_bits=effective_kv_valid_bits, + kv_valid_bits_is_live=state.use_kv_valid_bits, + sm_scale=1.0 / math.sqrt(state.head_dim) + if sm_scale is None + else float(sm_scale), + paged_kv=paged_kv, + ) + + def record_block_sparse_run_args( run_args: _BlockSparseRunArgs, stream: torch.cuda.Stream, ) -> None: """Extend tensor lifetimes for the asynchronous launch currently in flight.""" - run_args.q.record_stream(stream) - run_args.k.record_stream(stream) - run_args.v.record_stream(stream) + for tensor in (run_args.q, run_args.k, run_args.v): + tensor.record_stream(stream) + if run_args.k_summary is not None: + run_args.k_summary.record_stream(stream) + assert run_args.v_summary is not None + run_args.v_summary.record_stream(stream) run_args.out.record_stream(stream) - run_args.block_indptr.record_stream(stream) - run_args.block_indices.record_stream(stream) + if run_args.block_indptr is not None: + run_args.block_indptr.record_stream(stream) + assert run_args.block_indices is not None + run_args.block_indices.record_stream(stream) + else: + assert run_args.exact_block_bits is not None + run_args.exact_block_bits.record_stream(stream) if run_args.kv_valid_bits_is_live: run_args.kv_valid_bits.record_stream(stream) if run_args.paged_kv is not None: - run_args.paged_kv.paged_kv_indptr.record_stream(stream) - run_args.paged_kv.paged_kv_indices.record_stream(stream) + run_args.paged_kv.block_tables.record_stream(stream) run_args.paged_kv.seq_lens_kv.record_stream(stream) @@ -409,9 +550,13 @@ def launch_block_sparse( *, state: "_BlockSparsePlanState", ) -> torch.Tensor: - """Invoke the exact contiguous or paged ABI chosen by validated payload.""" + """Invoke the layout- and route-specific ABI chosen by the frozen plan.""" - if run_args.paged_kv is None: + sparse_format = state.sparse_format + use_proxy_routes = state.use_proxy_routes + if run_args.paged_kv is not None: + assert run_args.block_indptr is not None + assert run_args.block_indices is not None state.compiled( run_args.q, run_args.k, @@ -420,12 +565,20 @@ def launch_block_sparse( run_args.block_indptr, run_args.block_indices, run_args.kv_valid_bits, + run_args.paged_kv.block_tables, + run_args.paged_kv.seq_lens_kv, state.row_route_offsets, state.route_workspace, state.max_blocks_per_row, + run_args.paged_kv.num_physical_kv_pages, + run_args.paged_kv.block_table_row_stride, + run_args.paged_kv.k_page_stride, + run_args.paged_kv.v_page_stride, run_args.sm_scale, ) - else: + elif sparse_format == "bsr" and not use_proxy_routes: + assert run_args.block_indptr is not None + assert run_args.block_indices is not None state.compiled( run_args.q, run_args.k, @@ -434,17 +587,65 @@ def launch_block_sparse( run_args.block_indptr, run_args.block_indices, run_args.kv_valid_bits, - run_args.paged_kv.paged_kv_indptr, - run_args.paged_kv.paged_kv_indices, - run_args.paged_kv.seq_lens_kv, state.row_route_offsets, state.route_workspace, state.max_blocks_per_row, - run_args.paged_kv.num_physical_kv_pages, - run_args.paged_kv.k_page_stride, - run_args.paged_kv.v_page_stride, run_args.sm_scale, ) + elif sparse_format == "bitmask" and not use_proxy_routes: + assert run_args.exact_block_bits is not None + state.compiled( + run_args.q, + run_args.k, + run_args.v, + run_args.out, + run_args.exact_block_bits, + run_args.kv_valid_bits, + state.row_route_offsets, + state.route_workspace, + state.max_blocks_per_row, + run_args.sm_scale, + ) + elif sparse_format == "bsr" and use_proxy_routes: + assert run_args.block_indptr is not None + assert run_args.block_indices is not None + assert run_args.k_summary is not None + assert run_args.v_summary is not None + state.compiled( + run_args.q, + run_args.k, + run_args.v, + run_args.k_summary, + run_args.v_summary, + run_args.out, + run_args.block_indptr, + run_args.block_indices, + run_args.kv_valid_bits, + state.row_route_offsets, + state.route_workspace, + state.max_blocks_per_row, + run_args.sm_scale, + ) + elif sparse_format == "bitmask" and use_proxy_routes: + assert run_args.exact_block_bits is not None + assert run_args.k_summary is not None + assert run_args.v_summary is not None + state.compiled( + run_args.q, + run_args.k, + run_args.v, + run_args.k_summary, + run_args.v_summary, + run_args.out, + run_args.exact_block_bits, + run_args.kv_valid_bits, + state.row_route_offsets, + state.route_workspace, + state.max_blocks_per_row, + run_args.sm_scale, + ) + else: + raise AssertionError("frozen block-sparse plan has an unsupported route mode") return run_args.out @@ -454,8 +655,8 @@ def launch_block_sparse( "_PagedKVLaunchPayload", "_PagedKVStorage", "launch_block_sparse", + "prepare_block_sparse_run_unchecked", "record_block_sparse_run_args", "validate_block_sparse_metadata", "validate_block_sparse_run", - "validate_paged_kv_metadata", ] diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/block_sparse.py b/tensorrt_llm/_torch/attention/backends/prims_ts/block_sparse.py index af132c6a69ce..e19dc7b1de7f 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/block_sparse.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/block_sparse.py @@ -37,6 +37,7 @@ from flashinfer.api_logging import flashinfer_api +from ._block_sparse.common import _validate_contiguous_route_mode from ._block_sparse.config import _validate_block_sparse_static_profile from ._block_sparse.inspection import ( _inspect_block_sparse_bsr, @@ -53,15 +54,16 @@ _ContiguousKVStorage, _PagedKVStorage, launch_block_sparse as _launch_block_sparse, + prepare_block_sparse_run_unchecked as _prepare_block_sparse_run_unchecked, record_block_sparse_run_args as _record_block_sparse_run_args, validate_block_sparse_metadata as _validate_block_sparse_metadata, validate_block_sparse_run as _validate_block_sparse_run, - validate_paged_kv_metadata as _validate_paged_kv_metadata, ) from .decode import ( PagedKVCache, _normalize_paged_kv_cache, _resolve_cuda_device, + _validate_block_table_metadata, ) @@ -111,10 +113,12 @@ class BlockSparseTSWrapper(_BlockSparseWrapperBase): Q is ``[B, Sq, Hq, D]`` and K/V are ``[B, Skv, Hkv, D]``. Sparse rows are owned per batch, KV head, and query block, so every Q head in one grouped KV head consumes the same sparse row. A plan fixes geometry and a per-row - capacity; every run supplies its own BSR and optional token mask. - Callers must keep those tensors alive and immutable until the queued run or - captured graph finishes using them. CUDA Graph capture pins plan-owned - state only, so captured routing storage remains the caller's responsibility. + capacity; every run supplies either BSR or a packed exact-block bitmask. + Proxy-enabled plans additionally consume caller-owned K/V summaries, while + an optional token mask applies only to exact routes. Callers must keep those + tensors alive and immutable until the queued run or captured graph finishes + using them. CUDA Graph capture pins plan-owned state only, so captured + routing storage remains the caller's responsibility. One plan revision owns one mutable route workspace. Its runs must be ordered on one stream or externally synchronized; unordered concurrent runs require @@ -136,6 +140,8 @@ def plan( device: torch.device | str | int, max_blocks_per_row: int, use_kv_valid_bits: bool, + sparse_format: Literal["bsr", "bitmask"] = "bsr", + use_proxy_routes: bool = False, mask_type: Literal["dense", "causal"] = "dense", q_data_type: torch.dtype = torch.float16, kv_data_type: torch.dtype | None = None, @@ -144,11 +150,16 @@ def plan( """Choose a legal profile and allocate reusable routing capacity. The plan owns immutable geometry and a uniform route workspace, not a - sparse pattern. ``max_blocks_per_row`` bounds each runtime BSR row in - semantic ``kv_block_size`` blocks. ``use_kv_valid_bits`` selects whether - every :meth:`run` must supply the shared batch token mask. Callers may - pass different routing tensor identities and index extents to each run - as long as they fit this declared capacity. + sparse pattern. ``max_blocks_per_row`` bounds each runtime sparse row in + semantic ``kv_block_size`` blocks. ``sparse_format="bsr"`` consumes + canonical CSR-style rows, while ``"bitmask"`` consumes packed exact- + block bits. Enabling proxy routes represents unselected blocks through + caller-provided K/V summaries and currently requires + ``mask_type="dense"``. Exact-only plans continue to support causal + masking. ``use_kv_valid_bits`` selects whether every :meth:`run` must + supply the shared batch token mask. Callers may pass different routing + tensor identities and index extents to each run as long as they fit + this declared capacity. MHA, GQA, and MQA are supported with ``Hq / Hkv`` a power of two no greater than 32 and ``D=128``. Q, K, V, and O use one matching @@ -162,10 +173,13 @@ def plan( respectively. ``kv_block_size`` may be 8, 16, 32, or a positive multiple of 64. The Q tile groups complete Q-head groups and as many Q tokens as fit without crossing a semantic Q-block row, up to Q128; - fine KV blocks cap this at a SWAPAB Q32 tile. Every run prepares - per-KV-head canonical BSR into compact, profile-selected fixed-width - route metadata, and the attention core consumes only that metadata. - This remains true when every KV block is selected; + fine KV blocks cap this at a SWAPAB Q32 tile. Proxy routes reuse the + same Q-tile, KV-route, and MMA geometry as exact routes, but currently + use the direct scheduler because reusable planning cannot observe live + exact-route work. Every run prepares its selected BSR or bitmask into + compact, profile-selected fixed-width route metadata, and the attention + core consumes only that metadata. This remains true when every KV block + is selected; callers that know a pattern is dense should choose the dense FMHA API explicitly. @@ -182,6 +196,7 @@ def plan( runs require distinct wrappers. """ + _validate_contiguous_route_mode(sparse_format, use_proxy_routes) static = _validate_block_sparse_static_profile( batch_size=batch_size, seq_len_q=seq_len_q, @@ -198,6 +213,8 @@ def plan( output_dtype=o_data_type, max_blocks_per_row=max_blocks_per_row, ) + if use_proxy_routes and static.mask_type != "dense": + raise ValueError("block-sparse proxy routes require mask_type='dense'") device, device_index = _resolve_cuda_device(device) plan_stream = torch.cuda.current_stream(device) with torch.cuda.device(device_index), torch.cuda.stream(plan_stream): @@ -210,6 +227,8 @@ def plan( device=device, device_index=device_index, plan_stream=plan_stream, + sparse_format=sparse_format, + use_proxy_routes=use_proxy_routes, ) # This is the only wrapper mutation. Every failure above leaves the # previously published revision intact and runnable. @@ -221,12 +240,16 @@ def run( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - block_indptr: torch.Tensor, - block_indices: torch.Tensor, + block_indptr: torch.Tensor | None = None, + block_indices: torch.Tensor | None = None, *, + exact_block_bits: torch.Tensor | None = None, + k_summary: torch.Tensor | None = None, + v_summary: torch.Tensor | None = None, kv_valid_bits: torch.Tensor | None = None, sm_scale: float | None = None, out: torch.Tensor | None = None, + validate: bool = True, ) -> torch.Tensor: """Launch the current plan on the caller's current CUDA stream. @@ -238,17 +261,34 @@ def run( Only O is returned; this PrimTS API does not return LSE. The launch is enqueued asynchronously on the caller's current CUDA stream. - ``block_indptr`` is compact Int32 - ``[B, Hkv, ceil(Sq / q_block_size) + 1]`` and indexes compact - ``block_indices``. Every row must fit the planned semantic-block - capacity; referenced block IDs must be strictly increasing, unique, - and in range. Reusable runs trust these device-side values. CuTe DSL - assertions can diagnose violations when enabled before compilation; - otherwise invalid values have undefined behavior and may access out of - bounds. A masked plan requires + ``validate=True`` performs structural, plan-geometry, and alias + validation without reading tensor values; it is the safe public + default. ``validate=False`` treats every run argument as a trusted + binding and performs no explicit wrapper validation. K/V view + selection, scale forwarding, and optional output allocation are + unavoidable in both modes. + + A BSR plan consumes compact Int32 ``block_indptr`` with shape + ``[B, Hkv, ceil(Sq / q_block_size) + 1]`` and compact Int32 + ``block_indices``. A bitmask plan instead requires both BSR arguments + to be ``None`` and consumes packed UInt32 ``exact_block_bits`` with + shape ``[B, Hkv, ceil(Sq / q_block_size), ceil(num_kv_blocks / 32)]``. + Bit ``r`` of word ``w`` selects block ``32 * w + r``; final-word + padding bits are ignored. A proxy plan additionally consumes compact + ``k_summary`` and ``v_summary`` with shape + ``[B, num_kv_blocks, Hkv, D]``. K summaries are block means and V + summaries are block sums; the final partial block covers only its + structural tokens. + + Every row must fit the planned semantic-block capacity. Reusable runs + trust routing values. CuTe DSL assertions can diagnose violations when + enabled before compilation; otherwise invalid values have undefined + behavior and may access out of bounds. A masked plan requires ``kv_valid_bits`` with shape ``[B, ceil(Skv / 32)]`` and dtype UInt32; an unmasked plan requires - ``None``. Routing tensors may have different identities on every run. + ``None``. The mask applies only to raw exact routes; proxy summaries and + their represented-token mass remain caller-defined. Routing tensors may + have different identities on every run. Keep this wrapper alive until every captured CUDA Graph is destroyed. @@ -261,12 +301,18 @@ def run( v : torch.Tensor Compact value tensor with the same shape, dtype, and strides as ``k``. - block_indptr : torch.Tensor + block_indptr : torch.Tensor, optional Contiguous Int32 BSR row offsets with shape - ``[B, Hkv, ceil(Sq / q_block_size) + 1]``. - block_indices : torch.Tensor + ``[B, Hkv, ceil(Sq / q_block_size) + 1]``. Required by BSR plans. + block_indices : torch.Tensor, optional Contiguous Int32 semantic KV-block IDs referenced by - ``block_indptr``. + ``block_indptr``. Required by BSR plans. + exact_block_bits : torch.Tensor, optional + Compact packed UInt32 exact-block bitmap required by bitmask plans. + k_summary : torch.Tensor, optional + Per-block mean K tensor required by proxy plans. + v_summary : torch.Tensor, optional + Per-block summed V tensor required by proxy plans. kv_valid_bits : torch.Tensor, optional Contiguous UInt32 token-validity bitmap ``[B, ceil(Skv / 32)]``. Supply it exactly when the plan enabled token validity bits. @@ -275,6 +321,9 @@ def run( out : torch.Tensor, optional Caller-owned compact output buffer ``[B, Sq, Hq, D]`` with the planned output dtype. + validate : bool + Whether to validate tensor structure, plan geometry, and aliasing + before launching. Defaults to ``True``. Returns ------- @@ -283,17 +332,27 @@ def run( """ state = self._require_run_state() - run_stream = torch.cuda.current_stream(state.device) - run_args = _validate_block_sparse_run( + if not isinstance(validate, bool): + raise TypeError("validate must be a bool") + prepare_run = ( + _validate_block_sparse_run + if validate + else _prepare_block_sparse_run_unchecked + ) + run_args = prepare_run( q, _ContiguousKVStorage(k=k, v=v), state=state, block_indptr=block_indptr, block_indices=block_indices, + exact_block_bits=exact_block_bits, + k_summary=k_summary, + v_summary=v_summary, kv_valid_bits=kv_valid_bits, sm_scale=sm_scale, out=out, ) + run_stream = torch.cuda.current_stream(state.device) return self._launch_validated_run(state, run_args, run_stream) @@ -302,23 +361,29 @@ def block_sparse_attention( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - block_indptr: torch.Tensor, - block_indices: torch.Tensor, + block_indptr: torch.Tensor | None, + block_indices: torch.Tensor | None, q_block_size: int, kv_block_size: int, *, + exact_block_bits: torch.Tensor | None = None, + k_summary: torch.Tensor | None = None, + v_summary: torch.Tensor | None = None, kv_valid_bits: torch.Tensor | None = None, + sparse_format: Literal["bsr", "bitmask"] = "bsr", + use_proxy_routes: bool = False, mask_type: Literal["dense", "causal"] = "dense", sm_scale: float | None = None, out: torch.Tensor | None = None, ) -> torch.Tensor: """Plan and run one compact-BSHD block-sparse attention launch. - This one-shot form synchronously inspects canonical BSR, derives its largest - semantic row, creates a capacity-only plan, and passes the original routing - tensors to :meth:`BlockSparseTSWrapper.run`. It therefore cannot be invoked - inside CUDA Graph capture; plan a wrapper outside capture and capture only - ``run()`` instead. + This one-shot form creates a capacity-only plan and passes the original + routing tensors to :meth:`BlockSparseTSWrapper.run`. BSR inputs are + synchronously inspected to validate canonical rows and derive their maximum + width. Bitmask inputs use the structural KV-block count as a conservative + capacity bound. It therefore cannot be invoked inside CUDA Graph capture; + plan a wrapper outside capture and capture only ``run()`` instead. Parameters ---------- @@ -328,11 +393,12 @@ def block_sparse_attention( Compact key tensor ``[B, Skv, Hkv, D]``. v : torch.Tensor Compact value tensor with the same shape, dtype, and strides as ``k``. - block_indptr : torch.Tensor + block_indptr : torch.Tensor, optional Contiguous Int32 BSR row offsets with shape - ``[B, Hkv, ceil(Sq / q_block_size) + 1]``. - block_indices : torch.Tensor + ``[B, Hkv, ceil(Sq / q_block_size) + 1]``. Required in BSR mode. + block_indices : torch.Tensor, optional Contiguous Int32 semantic KV-block IDs referenced by ``block_indptr``. + Required in BSR mode. q_block_size : int Positive number of logical query tokens represented by one BSR row. The product with ``Hq / Hkv`` must be divisible by 8 so a physical Q @@ -340,8 +406,19 @@ def block_sparse_attention( kv_block_size : int Number of logical KV tokens represented by one BSR block ID; it must be 8, 16, 32, or a positive multiple of 64. + exact_block_bits : torch.Tensor, optional + Compact UInt32 exact-block bitmap required in bitmask mode. + k_summary : torch.Tensor, optional + Per-block mean K tensor required when proxy routes are enabled. + v_summary : torch.Tensor, optional + Per-block summed V tensor required when proxy routes are enabled. kv_valid_bits : torch.Tensor, optional Contiguous UInt32 token-validity bitmap ``[B, ceil(Skv / 32)]``. + sparse_format : {"bsr", "bitmask"}, optional + Runtime sparse representation. Defaults to ``"bsr"``. + use_proxy_routes : bool, optional + Represent unselected blocks through K/V summaries. Proxy routes + currently require dense masking. mask_type : {"dense", "causal"}, optional Attention mask applied inside each selected sparse block. sm_scale : float, optional @@ -370,6 +447,7 @@ def block_sparse_attention( raise ValueError("K and V must have identical shapes") use_kv_valid_bits = kv_valid_bits is not None + _validate_contiguous_route_mode(sparse_format, use_proxy_routes) static = _validate_block_sparse_static_profile( batch_size=batch_size, seq_len_q=seq_len_q, @@ -385,25 +463,37 @@ def block_sparse_attention( kv_dtype=k.dtype, output_dtype=q.dtype if out is None else out.dtype, ) + if use_proxy_routes and static.mask_type != "dense": + raise ValueError("block-sparse proxy routes require mask_type='dense'") device, _ = _resolve_cuda_device(q.device) _validate_block_sparse_metadata( - block_indptr, - block_indices, - kv_valid_bits, + sparse_format=sparse_format, + block_indptr=block_indptr, + block_indices=block_indices, + exact_block_bits=exact_block_bits, + kv_valid_bits=kv_valid_bits, device=device, batch_size=static.batch_size, seq_len_q=static.seq_len_q, seq_len_kv=static.seq_len_kv, num_kv_heads=static.num_kv_heads, q_block_size=static.q_block_size, + kv_block_size=static.kv_block_size, use_kv_valid_bits=static.use_kv_valid_bits, ) - max_blocks_per_row = _inspect_block_sparse_bsr( - block_indptr, - block_indices, - static=static, - stream=torch.cuda.current_stream(device), - ) + if sparse_format == "bsr": + max_blocks_per_row = _inspect_block_sparse_bsr( + block_indptr, + block_indices, + static=static, + stream=torch.cuda.current_stream(device), + ) + elif sparse_format == "bitmask": + max_blocks_per_row = ( + static.seq_len_kv + static.kv_block_size - 1 + ) // static.kv_block_size + else: + raise AssertionError(f"unsupported sparse format {sparse_format!r}") wrapper = BlockSparseTSWrapper() wrapper.plan( @@ -418,6 +508,8 @@ def block_sparse_attention( device=device, max_blocks_per_row=max_blocks_per_row, use_kv_valid_bits=static.use_kv_valid_bits, + sparse_format=sparse_format, + use_proxy_routes=use_proxy_routes, mask_type=static.mask_type, q_data_type=static.q_dtype, kv_data_type=static.kv_dtype, @@ -429,6 +521,9 @@ def block_sparse_attention( v, block_indptr, block_indices, + exact_block_bits=exact_block_bits, + k_summary=k_summary, + v_summary=v_summary, kv_valid_bits=kv_valid_bits, sm_scale=sm_scale, out=out, @@ -512,8 +607,7 @@ def run( self, q: torch.Tensor, paged_kv_cache: PagedKVCache, - paged_kv_indptr: torch.Tensor, - paged_kv_indices: torch.Tensor, + block_tables: torch.Tensor, seq_lens_kv: torch.Tensor, block_indptr: torch.Tensor, block_indices: torch.Tensor, @@ -521,6 +615,7 @@ def run( kv_valid_bits: torch.Tensor | None = None, sm_scale: float | None = None, out: torch.Tensor | None = None, + validate: bool = True, ) -> torch.Tensor: """Launch with live lengths, page tables, and sparse routes. @@ -529,15 +624,20 @@ def run( tuple whose members are ``[P, Hkv, page, D]`` with compact inner HND strides and arbitrary non-overlapping outer page strides. - ``paged_kv_indptr`` is compact Int32 ``[B + 1]``; - ``paged_kv_indices`` is compact Int32 with capacity at least its live - final offset; and ``seq_lens_kv`` is compact Int32 ``[B]``. All values - are read on device. The caller must keep every dense length in - ``[1, max_seq_len_kv]`` and every causal length in - ``[Sq, max_seq_len_kv]``. ``paged_kv_indptr`` must start at zero and - contain bounded, monotone rows with at least - ``ceil(seq_lens_kv[b] / page_size)`` entries. Every page ID in its live - prefix must lie in ``[0, P)``. Each BSR row must contain strictly + ``validate=True`` performs structural, plan-geometry, and alias + validation without reading tensor values; it is the safe public + default. ``validate=False`` treats every run argument as a trusted + binding and performs no explicit wrapper validation. K/V view + selection, scale forwarding, and optional output allocation are + unavoidable in both modes. + + ``block_tables`` is Int32 ``[B, C]``, contiguous within each row but + permitted to use a padded outer row stride; ``seq_lens_kv`` is compact + Int32 ``[B]``. All values are read on device. The caller must keep + every dense length in ``[1, max_seq_len_kv]`` and every causal length + in ``[Sq, max_seq_len_kv]``. Every page-table row must contain at least + ``ceil(seq_lens_kv[b] / page_size)`` live entries. Every page ID in its + live prefix must lie in ``[0, P)``. Each BSR row must contain strictly increasing, unique block IDs whose final block starts before that request's live K/V length, and its width must not exceed the planned ``max_blocks_per_row``. Reusable runs trust all of these device-side @@ -564,10 +664,10 @@ def run( Either a combined cache ``[P, 2, Hkv, page_size, D]`` or a ``(K, V)`` tuple whose tensors are ``[P, Hkv, page_size, D]``. - paged_kv_indptr : torch.Tensor - Contiguous Int32 live request offsets with shape ``[B + 1]``. - paged_kv_indices : torch.Tensor - Contiguous Int32 physical-page ID capacity. + block_tables : torch.Tensor + Live Int32 physical page IDs with shape ``[B, C]``. Entries are + contiguous within each row; padded, non-overlapping row strides are + supported and inactive tail entries are ignored. seq_lens_kv : torch.Tensor Contiguous Int32 live logical K/V lengths with shape ``[B]``. Values must satisfy the dense or causal bounds above. @@ -586,6 +686,9 @@ def run( out : torch.Tensor, optional Caller-owned compact output buffer ``[B, Sq, Hq, D]`` with the planned output dtype. + validate : bool + Whether to validate tensor structure, plan geometry, and aliasing + before launching. Defaults to ``True``. Returns ------- @@ -594,13 +697,19 @@ def run( """ state = self._require_run_state() + if not isinstance(validate, bool): + raise TypeError("validate must be a bool") run_stream = torch.cuda.current_stream(state.device) - run_args = _validate_block_sparse_run( + prepare_run = ( + _validate_block_sparse_run + if validate + else _prepare_block_sparse_run_unchecked + ) + run_args = prepare_run( q, _PagedKVStorage( paged_kv_cache=paged_kv_cache, - paged_kv_indptr=paged_kv_indptr, - paged_kv_indices=paged_kv_indices, + block_tables=block_tables, seq_lens_kv=seq_lens_kv, ), state=state, @@ -617,15 +726,14 @@ def run( def block_sparse_attention_with_paged_kv_cache( q: torch.Tensor, paged_kv_cache: PagedKVCache, - paged_kv_indptr: torch.Tensor, - paged_kv_indices: torch.Tensor, + block_tables: torch.Tensor, + seq_lens_kv: torch.Tensor, block_indptr: torch.Tensor, block_indices: torch.Tensor, q_block_size: int, kv_block_size: int, *, max_seq_len_kv: int, - seq_lens_kv: torch.Tensor, kv_valid_bits: torch.Tensor | None = None, mask_type: Literal["dense", "causal"] = "dense", sm_scale: float | None = None, @@ -633,12 +741,11 @@ def block_sparse_attention_with_paged_kv_cache( ) -> torch.Tensor: """Plan and run one fixed-Q paged block-sparse attention launch. - This convenience entry point synchronously validates live page and sparse - metadata, including the complete live physical-page-ID prefix, creates a - capacity-only temporary plan, then forwards the inspected tensors through - the trusted live run API. It cannot run during CUDA Graph capture; plan a - wrapper outside capture and capture only - :meth:`BlockSparsePagedTSWrapper.run` instead. + This convenience entry point synchronously validates the live page tables, + K/V lengths, and sparse metadata, creates a capacity-only temporary plan, + then forwards the inspected tensors through the trusted live run API. It + cannot run during CUDA Graph capture; plan a wrapper outside capture and + capture only :meth:`BlockSparsePagedTSWrapper.run` instead. Parameters ---------- @@ -647,11 +754,13 @@ def block_sparse_attention_with_paged_kv_cache( paged_kv_cache : PagedKVCache Either a combined cache ``[P, 2, Hkv, page_size, D]`` or a ``(K, V)`` tuple whose tensors are ``[P, Hkv, page_size, D]``. - paged_kv_indptr : torch.Tensor - Contiguous Int32 request offsets into ``paged_kv_indices``, with shape - ``[B + 1]``. - paged_kv_indices : torch.Tensor - Contiguous Int32 physical page IDs referenced by ``paged_kv_indptr``. + block_tables : torch.Tensor + Int32 physical page IDs ``[B, C]``, contiguous within each row and free + to use a padded outer row stride. ``C * page_size`` must cover + ``max_seq_len_kv``; only the first ``ceil(seq_lens_kv[b] / page_size)`` + entries of each row are read. + seq_lens_kv : torch.Tensor + Contiguous Int32 per-request logical KV lengths with shape ``[B]``. block_indptr : torch.Tensor Contiguous Int32 BSR row offsets with shape ``[B, Hkv, ceil(Sq / q_block_size) + 1]``. @@ -666,8 +775,6 @@ def block_sparse_attention_with_paged_kv_cache( be 8, 16, 32, or a positive multiple of 64. max_seq_len_kv : int Static maximum logical K/V length used for planning. - seq_lens_kv : torch.Tensor - Contiguous Int32 per-request logical KV lengths with shape ``[B]``. kv_valid_bits : torch.Tensor, optional Contiguous UInt32 logical-token validity bitmap ``[B, ceil(max_seq_len_kv / 32)]``. @@ -693,13 +800,16 @@ def block_sparse_attention_with_paged_kv_cache( batch_size, seq_len_q, num_qo_heads, head_dim = map(int, q.shape) metadata_device, _ = _resolve_cuda_device(q.device) - _validate_paged_kv_metadata( - paged_kv_indptr, - paged_kv_indices, - seq_lens_kv, - device=metadata_device, - batch_size=batch_size, + table_device, table_batch_size, table_capacity = _validate_block_table_metadata( + block_tables, seq_lens_kv ) + if table_device != q.device: + raise ValueError(f"paged-KV metadata must be on {q.device}, got {table_device}") + if table_batch_size != batch_size: + raise ValueError( + "seq_lens_kv must have one entry per request: " + f"expected {batch_size}, got {table_batch_size}" + ) ( k_cache, @@ -734,23 +844,31 @@ def block_sparse_attention_with_paged_kv_cache( kv_dtype=k_cache.dtype, output_dtype=q.dtype if out is None else out.dtype, ) + if table_capacity * page_size < static.seq_len_kv: + raise ValueError( + "block_tables must cover the planned K/V capacity: expected at " + f"least {(static.seq_len_kv + page_size - 1) // page_size} columns, " + f"got {table_capacity}" + ) _validate_block_sparse_metadata( - block_indptr, - block_indices, - kv_valid_bits, + sparse_format="bsr", + block_indptr=block_indptr, + block_indices=block_indices, + exact_block_bits=None, + kv_valid_bits=kv_valid_bits, device=metadata_device, batch_size=static.batch_size, seq_len_q=static.seq_len_q, seq_len_kv=static.seq_len_kv, num_kv_heads=static.num_kv_heads, q_block_size=static.q_block_size, + kv_block_size=static.kv_block_size, use_kv_valid_bits=static.use_kv_valid_bits, ) max_blocks_per_row = _inspect_paged_block_sparse_metadata( block_indptr, block_indices, - paged_kv_indptr, - paged_kv_indices, + block_tables, seq_lens_kv, static=static, num_physical_kv_pages=num_physical_kv_pages, @@ -780,8 +898,7 @@ def block_sparse_attention_with_paged_kv_cache( return wrapper.run( q, paged_kv_cache, - paged_kv_indptr, - paged_kv_indices, + block_tables, seq_lens_kv, block_indptr, block_indices, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_inspect.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_inspect.py index de53ccb9c7de..5728a7c6d373 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_inspect.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_inspect.py @@ -20,10 +20,10 @@ their caller. Token-mask contents belong to the run-time prepare kernel and are not read here. -Paged inspection first validates live sequence lengths and page rows, then -four warps validate four BSR Q-block rows per CTA. Both publish one validation -status plus the maximum row width in one Int64 summary; no route payload is -constructed. +Paged inspection first validates live sequence lengths and the live prefix of +every page-table row, then four warps validate four BSR Q-block rows per CTA. +Both publish one validation status plus the maximum row width in one Int64 +summary; no route payload is constructed. """ import functools @@ -51,9 +51,7 @@ _BSR_ERROR_INDEX_OUT_OF_RANGE = 2 _BSR_ERROR_INVALID_INDPTR = 3 _ERROR_INVALID_SEQ_LEN = 4 -_ERROR_INVALID_PAGE_INDPTR = 5 -_ERROR_INSUFFICIENT_PAGE_CAPACITY = 6 -_ERROR_INVALID_PHYSICAL_PAGE_ID = 7 +_ERROR_INVALID_PHYSICAL_PAGE_ID = 5 @cute.jit @@ -296,7 +294,7 @@ def __call__( class _InspectPagedKvMetadata: - """Validate live lengths and page rows with one warp per request.""" + """Validate live lengths and page-table rows with one warp per request.""" def __init__( self, @@ -314,16 +312,16 @@ def __init__( @cute.jit def __call__( self, - paged_kv_indptr: cute.Tensor, - paged_kv_indices: cute.Tensor, + block_tables: cute.Tensor, + block_table_row_stride: cutlass.Int64, seq_lens_kv: cute.Tensor, num_physical_kv_pages: cutlass.Int64, summary: cute.Tensor, stream: cuda_drv.CUstream, ) -> None: self.kernel( - paged_kv_indptr, - paged_kv_indices, + block_tables, + block_table_row_stride, seq_lens_kv, num_physical_kv_pages, summary, @@ -340,8 +338,8 @@ def __call__( @cute.kernel def kernel( self, - paged_kv_indptr: cute.Tensor, - paged_kv_indices: cute.Tensor, + block_tables: cute.Tensor, + block_table_row_stride: cutlass.Int64, seq_lens_kv: cute.Tensor, num_physical_kv_pages: cutlass.Int64, summary: cute.Tensor, @@ -353,9 +351,7 @@ def kernel( batch_idx = block_idx * _WARPS_PER_CTA + warp_idx request_is_valid = batch_idx < self.batch_size - request_begin = cutlass.Int32(0) - request_end = cutlass.Int32(0) - request_range_is_valid = cutlass.Int32(0) + live_pages = cutlass.Int32(0) error_code = cutlass.Int32(_BSR_ERROR_NONE) if lane_idx == 0 and request_is_valid: seq_len_kv = cutlass.Int32(seq_lens_kv[batch_idx]) @@ -363,38 +359,20 @@ def kernel( seq_len_kv >= cutlass.Int32(self.minimum_seq_len_kv) and seq_len_kv <= cutlass.Int32(self.max_seq_len_kv) ) - if not seq_len_is_valid: - error_code = cutlass.Int32(_ERROR_INVALID_SEQ_LEN) - - request_begin = cutlass.Int32(paged_kv_indptr[batch_idx]) - request_end = cutlass.Int32(paged_kv_indptr[batch_idx + 1]) - num_page_indices = cutlass.Int32(cute.size(paged_kv_indices)) - request_range_is_valid = cutlass.Int32( - paged_kv_indptr[cutlass.Int32(0)] == cutlass.Int32(0) - and request_begin >= cutlass.Int32(0) - and request_begin <= request_end - and request_end <= num_page_indices - ) - if request_range_is_valid == cutlass.Int32(0): - error_code = cutlass.Int32(_ERROR_INVALID_PAGE_INDPTR) - elif seq_len_is_valid: - required_pages = (seq_len_kv - cutlass.Int32(1)) // cutlass.Int32( + if seq_len_is_valid: + live_pages = (seq_len_kv - cutlass.Int32(1)) // cutlass.Int32( self.page_size ) + cutlass.Int32(1) - if request_end - request_begin < required_pages: - error_code = cutlass.Int32(_ERROR_INSUFFICIENT_PAGE_CAPACITY) + else: + error_code = cutlass.Int32(_ERROR_INVALID_SEQ_LEN) - request_begin = _warp_broadcast_i32(request_begin, 0) - request_end = _warp_broadcast_i32(request_end, 0) - request_range_is_valid = _warp_broadcast_i32(request_range_is_valid, 0) - if request_is_valid and request_range_is_valid != cutlass.Int32(0): + live_pages = _warp_broadcast_i32(live_pages, 0) + if request_is_valid: + row_begin = cutlass.Int64(batch_idx) * block_table_row_stride page_offset = cutlass.Int64(lane_idx) - request_page_count = cutlass.Int64(request_end) - cutlass.Int64( - request_begin - ) - while page_offset < request_page_count: - page_position = cutlass.Int64(request_begin) + page_offset - physical_page_id = cutlass.Int32(paged_kv_indices[page_position]) + while page_offset < cutlass.Int64(live_pages): + page_position = cutlass.Int64(row_begin + page_offset) + physical_page_id = cutlass.Int32(block_tables.iterator[page_position]) if ( physical_page_id < cutlass.Int32(0) or cutlass.Int64(physical_page_id) >= num_physical_kv_pages @@ -433,16 +411,16 @@ def __call__( self, block_indptr: cute.Tensor, block_indices: cute.Tensor, - paged_kv_indptr: cute.Tensor, - paged_kv_indices: cute.Tensor, + block_tables: cute.Tensor, + block_table_row_stride: cutlass.Int64, seq_lens_kv: cute.Tensor, num_physical_kv_pages: cutlass.Int64, summary: cute.Tensor, stream: cuda_drv.CUstream, ) -> None: self.inspect_requests( - paged_kv_indptr, - paged_kv_indices, + block_tables, + block_table_row_stride, seq_lens_kv, num_physical_kv_pages, summary, @@ -532,8 +510,9 @@ def compile_paged_block_sparse_metadata_inspection( """Compile one paged metadata entry that launches request then live-BSR.""" num_q_block_rows = (seq_len_q + q_block_size - 1) // q_block_size - logical_page_capacity = cute.sym_int() logical_nnz = cute.sym_int() + runtime_page_columns = cute.sym_int() + runtime_page_row_stride = cute.sym_int64(divisibility=1) stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True) inspect_requests = _InspectPagedKvMetadata( @@ -564,8 +543,13 @@ def compile_paged_block_sparse_metadata_inspection( alignment=4, ), _fake_compact(cutlass.Int32, (logical_nnz,), alignment=4), - _fake_compact(cutlass.Int32, (batch_size + 1,), alignment=4), - _fake_compact(cutlass.Int32, (logical_page_capacity,), alignment=4), + cute.runtime.make_fake_tensor( + cutlass.Int32, + (batch_size, runtime_page_columns), + stride=(runtime_page_row_stride, 1), + assumed_align=4, + ), + cutlass.Int64(1), _fake_compact(cutlass.Int32, (batch_size,), alignment=4), cutlass.Int64(1), _fake_compact(cutlass.Int64, (_SUMMARY_FIELDS,), alignment=8), diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_prepare.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_prepare.py index 6a61d0be86cd..08f2872809cb 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_prepare.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/block_sparse_prepare.py @@ -12,13 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Prepare live canonical BSR rows for the PrimTS FMHA route consumer. +"""Prepare exact-first sparse routes for the PrimTS FMHA route consumer. -The kernel converts caller-owned semantic KV blocks into fixed-stride route -metadata on every run. Route origins are logical KV-token coordinates; -the paged specialization also resolves each origin to a physical page ID for -the attention load path. One warp handles one BSR row and iterates only that -row's live routes, while four warps share a CTA. +The BSR frontend is shared by continuous exact/proxy and paged exact storage. +It validates each canonical row once, emits the same logical exact records, +then either resolves paged locators or appends continuous proxy records. The +bitmask frontend shares the record geometry and emitters but remains limited +to continuous storage. Proxy suffixes contain one stable record per summary +group; a fully exact group remains present with zero score words. One warp owns +one sparse row and four warps share a CTA. ``row_route_offsets`` is a separate plan-owned immutable Int32 tensor. ``route_workspace`` contains only mutable row counts and route metadata @@ -27,21 +29,20 @@ semantic BSR-block limit, which remains distinct from packed-route capacity. """ -import math from dataclasses import dataclass import cutlass import cutlass.cute as cute from cuda.bindings import driver as cuda_drv from cutlass.cute.testing import assert_ as runtime_assert -from cutlass.experimental import primitives as prims from ..._block_sparse.prepared import ( _PREPARED_ROUTE_IS_FULL_FLAG, + _PREPARED_ROUTE_IS_PROXY_FLAG, _BlockSparseRouteLayout, ) -from .block_sparse_inspect import _validate_bsr_row_lane from .fmha_decode_resources.helpers_common import _warp_broadcast_i32 +from .block_sparse_inspect import _validate_bsr_row_lane _WARPS_PER_CTA = 4 @@ -50,22 +51,27 @@ @dataclass(frozen=True) -class _PreparedRouteConfig: - """Compile-time geometry shared by contiguous and paged route packing.""" +class _RouteConfig: + """Compile-time route geometry shared across sparse input/storage modes.""" num_kv_heads: int - num_q_block_rows: int + num_q_blocks: int num_kv_blocks: int + num_exact_words: int + num_proxy_groups: int num_rows: int seq_len_kv: int kv_block_size: int atom_size: int + atoms_per_block: int logical_origins_per_route: int token_words_per_route: int atom_valid_mask_word_offset: int route_flags_word_offset: int token_words_word_offset: int - has_token_bits: bool + stores_score_words: bool + apply_token_mask: bool + use_proxy_routes: bool route_metadata_stride_words: int route_metadata_base_word_offset: int @@ -78,18 +84,30 @@ def create( seq_len_kv: int, q_block_size: int, kv_block_size: int, - ) -> "_PreparedRouteConfig": - """Build shared prepare geometry without adding a storage-mode flag.""" + apply_token_mask: bool, + use_proxy_routes: bool, + ) -> "_RouteConfig": + """Build storage-independent route geometry and policy flags.""" - num_q_block_rows = (seq_len_q + q_block_size - 1) // q_block_size - return _PreparedRouteConfig( + stores_score_words = layout.token_words_word_offset is not None + if apply_token_mask and not stores_score_words: + raise ValueError("token masking requires prepared score words") + if use_proxy_routes and not stores_score_words: + raise ValueError("proxy routes require prepared score words") + num_q_blocks = (seq_len_q + q_block_size - 1) // q_block_size + num_kv_blocks = (seq_len_kv + kv_block_size - 1) // kv_block_size + return _RouteConfig( num_kv_heads=num_kv_heads, - num_q_block_rows=num_q_block_rows, - num_kv_blocks=(seq_len_kv + kv_block_size - 1) // kv_block_size, + num_q_blocks=num_q_blocks, + num_kv_blocks=num_kv_blocks, + num_exact_words=(num_kv_blocks + _WARP_SIZE - 1) // _WARP_SIZE, + num_proxy_groups=(num_kv_blocks + layout.kv_route_size - 1) + // layout.kv_route_size, num_rows=layout.num_rows, seq_len_kv=seq_len_kv, kv_block_size=kv_block_size, atom_size=layout.atom_size, + atoms_per_block=kv_block_size // layout.atom_size, logical_origins_per_route=layout.logical_origins_per_route, token_words_per_route=layout.token_words_per_route, atom_valid_mask_word_offset=layout.atom_valid_mask_word_offset, @@ -99,7 +117,9 @@ def create( if layout.token_words_word_offset is not None else 0 ), - has_token_bits=layout.has_token_bits, + stores_score_words=stores_score_words, + apply_token_mask=apply_token_mask, + use_proxy_routes=use_proxy_routes, route_metadata_stride_words=layout.route_metadata_stride_words, route_metadata_base_word_offset=layout.route_metadata_base_word_offset, ) @@ -115,35 +135,44 @@ def _positive_i32_ceil_div( @cute.jit -def _retained_atom_count( - block_indices: cute.Tensor, - row_begin: cutlass.Int32, - row_end: cutlass.Int32, - kv_block_size: cutlass.Constexpr[int], - atom_size: cutlass.Constexpr[int], - seq_len_kv: cutlass.Int32, +def _prepared_route_counts( + selected_block_count: cutlass.Int32, + cfg: cutlass.Constexpr[_RouteConfig], +) -> tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32]: + """Return exact atoms, exact routes, and total prepared routes for one row.""" + + exact_atom_count = selected_block_count * cutlass.Int32(cfg.atoms_per_block) + exact_route_count = ( + exact_atom_count + cutlass.Int32(cfg.logical_origins_per_route - 1) + ) // cutlass.Int32(cfg.logical_origins_per_route) + total_route_count = exact_route_count + if cutlass.const_expr(cfg.use_proxy_routes): + total_route_count += cutlass.Int32(cfg.num_proxy_groups) + return exact_atom_count, exact_route_count, total_route_count + + +@cute.jit +def _prepared_row_route_begin( + row_route_offsets: cute.Tensor, + linear_row_idx: cutlass.Int32, + lane_idx: cutlass.Int32, + row_is_valid: cutlass.Boolean, + total_route_count: cutlass.Int32, ) -> cutlass.Int32: - """Count selected atoms whose logical origin precedes ``seq_len_kv``.""" - - row_nnz = row_end - row_begin - retained_atoms = cutlass.Int32(0) - if row_nnz > cutlass.Int32(0): - atoms_per_block = kv_block_size // atom_size - retained_atoms = (row_nnz - cutlass.Int32(1)) * cutlass.Int32(atoms_per_block) - last_block_idx = cutlass.Int32(block_indices[row_end - cutlass.Int32(1)]) - last_block_origin = last_block_idx * cutlass.Int32(kv_block_size) - remaining_tokens = cutlass.Int32(seq_len_kv) - last_block_origin + """Load and validate one row's plan-owned prepared-route span.""" + + row_route_begin = cutlass.Int32(0) + if lane_idx == cutlass.Int32(0) and row_is_valid: + row_route_begin = cutlass.Int32(row_route_offsets[linear_row_idx]) + row_route_end = cutlass.Int32(row_route_offsets[linear_row_idx + 1]) + row_capacity = row_route_end - row_route_begin runtime_assert( - remaining_tokens > cutlass.Int32(0), - "block_indices row exceeds the live KV block range", + row_route_begin >= cutlass.Int32(0) + and row_capacity >= cutlass.Int32(0) + and total_route_count <= row_capacity, + "prepared routes exceed planned row capacity", ) - retained_last_atoms = (remaining_tokens - cutlass.Int32(1)) // cutlass.Int32( - atom_size - ) + cutlass.Int32(1) - if retained_last_atoms > cutlass.Int32(atoms_per_block): - retained_last_atoms = cutlass.Int32(atoms_per_block) - retained_atoms = retained_atoms + retained_last_atoms - return retained_atoms + return _warp_broadcast_i32(row_route_begin, 0) @cute.jit @@ -177,71 +206,72 @@ def _resolve_route_logical_atom_origin( @cute.jit -def _load_coarse_token_word( - block_indices: cute.Tensor, - kv_valid_bits: cute.Tensor, - row_begin: cutlass.Int32, - row_end: cutlass.Int32, - route_idx: cutlass.Int32, - logical_word_idx: cutlass.Int32, - batch_idx: cutlass.Int32, - kv_block_size: cutlass.Constexpr[int], - atom_size: cutlass.Constexpr[int], - logical_origins_per_route: cutlass.Constexpr[int], - seq_len_kv: cutlass.Int32, -) -> cutlass.Uint32: - """Load one logical K32 word from a coarse atom larger than K32.""" - - logical_word = cutlass.Uint32(0) - words_per_atom = atom_size // 32 - atom_in_route = logical_word_idx // cutlass.Int32(words_per_atom) - word_in_atom = logical_word_idx % cutlass.Int32(words_per_atom) - logical_origin, valid = _resolve_route_logical_atom_origin( - block_indices, - row_begin, - row_end, - route_idx, - atom_in_route, - kv_block_size, - atom_size, - logical_origins_per_route, - seq_len_kv, - ) - logical_word_origin = logical_origin + word_in_atom * cutlass.Int32(32) - if valid and logical_word_origin < cutlass.Int32(seq_len_kv): - valid_bits_word_idx = logical_word_origin >> cutlass.Int32(5) - logical_word = cutlass.Uint32(kv_valid_bits[batch_idx, valid_bits_word_idx]) - remaining_tokens = cutlass.Int32(seq_len_kv) - logical_word_origin - if remaining_tokens < cutlass.Int32(32): - logical_word = logical_word & ( - (cutlass.Uint32(1) << remaining_tokens) - cutlass.Uint32(1) - ) - return logical_word +def _low_bits_mask(valid_bits: cutlass.Int32) -> cutlass.Uint32: + """Return a Uint32 mask with its lowest clamped bit count set.""" + + mask = cutlass.Uint32(0) + if valid_bits >= cutlass.Int32(_WARP_SIZE): + mask = cutlass.Uint32(0xFFFFFFFF) + elif valid_bits > cutlass.Int32(0): + mask = (cutlass.Uint32(1) << valid_bits) - cutlass.Uint32(1) + return mask @cute.jit -def _load_atom_token_chunk( +def _load_exact_score_word( + route_workspace: cute.Tensor, kv_valid_bits: cute.Tensor, + route_metadata_word_index: cutlass.Int32, + logical_word_idx: cutlass.Int32, batch_idx: cutlass.Int32, - logical_origin: cutlass.Int32, - origin_is_valid: cutlass.Boolean, - atom_size: cutlass.Constexpr[int], seq_len_kv: cutlass.Int32, + cfg: cutlass.Constexpr[_RouteConfig], ) -> cutlass.Uint32: - """Load the <=K32 mask chunk owned by one resolved-origin lane.""" - - token_chunk = cutlass.Uint32(0) - if origin_is_valid: - valid_bits_word_idx = logical_origin >> cutlass.Int32(5) - source_word = cutlass.Uint32(kv_valid_bits[batch_idx, valid_bits_word_idx]) - token_chunk = source_word >> (logical_origin & cutlass.Int32(31)) - token_chunk = token_chunk & cutlass.Uint32((1 << atom_size) - 1) - remaining_tokens = cutlass.Int32(seq_len_kv) - logical_origin - if remaining_tokens < cutlass.Int32(atom_size): - token_chunk = token_chunk & ( - (cutlass.Uint32(1) << remaining_tokens) - cutlass.Uint32(1) - ) - return token_chunk + """Build one exact score word with optional caller-token masking.""" + + token_word = cutlass.Uint32(0) + if cutlass.const_expr(cfg.atom_size <= _WARP_SIZE): + atoms_per_word = _WARP_SIZE // cfg.atom_size + first_atom_idx = logical_word_idx * cutlass.Int32(atoms_per_word) + for atom_in_word in cutlass.range_constexpr(atoms_per_word): + atom_idx = first_atom_idx + cutlass.Int32(atom_in_word) + if atom_idx < cutlass.Int32(cfg.logical_origins_per_route): + origin = cutlass.Int32( + route_workspace[route_metadata_word_index + atom_idx] + ) + atom_word = cutlass.Uint32(0) + if origin >= cutlass.Int32(0): + if cutlass.const_expr(cfg.apply_token_mask): + source_word_idx = origin >> cutlass.Int32(5) + atom_word = cutlass.Uint32( + kv_valid_bits[batch_idx, source_word_idx] + ) + atom_word = atom_word >> (origin & cutlass.Int32(31)) + atom_word = atom_word & cutlass.Uint32((1 << cfg.atom_size) - 1) + atom_word = atom_word & _low_bits_mask(seq_len_kv - origin) + else: + atom_word = _low_bits_mask( + seq_len_kv - origin + ) & cutlass.Uint32((1 << cfg.atom_size) - 1) + token_word = token_word | ( + atom_word << cutlass.Int32(atom_in_word * cfg.atom_size) + ) + else: + words_per_atom = cfg.atom_size // _WARP_SIZE + atom_idx = logical_word_idx // cutlass.Int32(words_per_atom) + word_in_atom = logical_word_idx % cutlass.Int32(words_per_atom) + origin = cutlass.Int32(route_workspace[route_metadata_word_index + atom_idx]) + word_origin = origin + word_in_atom * cutlass.Int32(_WARP_SIZE) + if origin >= cutlass.Int32(0): + if cutlass.const_expr(cfg.apply_token_mask): + if word_origin < seq_len_kv: + source_word_idx = word_origin >> cutlass.Int32(5) + token_word = cutlass.Uint32( + kv_valid_bits[batch_idx, source_word_idx] + ) & _low_bits_mask(seq_len_kv - word_origin) + else: + token_word = _low_bits_mask(seq_len_kv - word_origin) + return token_word @cute.jit @@ -251,7 +281,7 @@ def _resolve_prepared_bsr_row( linear_row_idx: cutlass.Int32, lane_idx: cutlass.Int32, row_is_valid: cutlass.Boolean, - cfg: cutlass.Constexpr[_PreparedRouteConfig], + cfg: cutlass.Constexpr[_RouteConfig], ) -> tuple[cutlass.Int32, cutlass.Int32, cutlass.Int32]: """Resolve one trusted canonical runtime BSR row.""" @@ -259,8 +289,8 @@ def _resolve_prepared_bsr_row( row_end = cutlass.Int32(0) batch_idx = cutlass.Int32(0) if lane_idx == cutlass.Int32(0) and row_is_valid: - q_block_row_idx = linear_row_idx % cfg.num_q_block_rows - linear_batch_head_idx = linear_row_idx // cfg.num_q_block_rows + q_block_row_idx = linear_row_idx % cfg.num_q_blocks + linear_batch_head_idx = linear_row_idx // cfg.num_q_blocks kv_head_idx = linear_batch_head_idx % cfg.num_kv_heads batch_idx = linear_batch_head_idx // cfg.num_kv_heads row_begin = cutlass.Int32(block_indptr[batch_idx, kv_head_idx, q_block_row_idx]) @@ -297,209 +327,256 @@ def _resolve_prepared_bsr_row( @cute.jit -def _publish_prepared_route_count( - block_indices: cute.Tensor, - row_route_offsets: cute.Tensor, +def _finalize_exact_route( route_workspace: cute.Tensor, - row_begin: cutlass.Int32, - row_end: cutlass.Int32, - linear_row_idx: cutlass.Int32, - lane_idx: cutlass.Int32, - row_is_valid: cutlass.Boolean, - max_blocks_per_row: cutlass.Int32, - seq_len_kv: cutlass.Int32, - cfg: cutlass.Constexpr[_PreparedRouteConfig], -) -> tuple[cutlass.Int32, cutlass.Int32]: - """Assert semantic capacity, publish the header, and return its live span.""" - - row_route_begin = cutlass.Int32(0) - required_route_count = cutlass.Int32(0) - if lane_idx == cutlass.Int32(0) and row_is_valid: - row_route_begin = cutlass.Int32(row_route_offsets[linear_row_idx]) - selected_block_count = row_end - row_begin - runtime_assert( - selected_block_count <= max_blocks_per_row, - "selected BSR blocks exceed planned semantic capacity", - ) - retained_atom_count = _retained_atom_count( - block_indices, - row_begin, - row_end, - cfg.kv_block_size, - cfg.atom_size, - seq_len_kv, - ) - required_route_count = ( - retained_atom_count + cutlass.Int32(cfg.logical_origins_per_route - 1) - ) // cutlass.Int32(cfg.logical_origins_per_route) - route_workspace[linear_row_idx] = required_route_count - row_route_begin = _warp_broadcast_i32(row_route_begin, 0) - required_route_count = _warp_broadcast_i32(required_route_count, 0) - return required_route_count, row_route_begin - - -@cute.jit -def _store_prepared_route_validity( - block_indices: cute.Tensor, kv_valid_bits: cute.Tensor, - route_workspace: cute.Tensor, - row_begin: cutlass.Int32, - row_end: cutlass.Int32, - route_idx: cutlass.Int32, + route_metadata_word_index: cutlass.Int32, batch_idx: cutlass.Int32, lane_idx: cutlass.Int32, - logical_origin: cutlass.Int32, - logical_origin_is_valid: cutlass.Boolean, - stored_atom_is_full: cutlass.Boolean, - route_metadata_word_index: cutlass.Int32, + atom_is_valid: cutlass.Boolean, seq_len_kv: cutlass.Int32, - cfg: cutlass.Constexpr[_PreparedRouteConfig], + cfg: cutlass.Constexpr[_RouteConfig], ) -> None: - """Store storage-independent atom, token, and route validity metadata.""" + """Finalize an exact record after its logical origins are stored.""" - stored_atom_valid_mask = cutlass.Int32( - cute.arch.vote_ballot_sync(logical_origin_is_valid) - ) - structural_route_is_full = cute.arch.vote_all_sync( - lane_idx >= cutlass.Int32(cfg.logical_origins_per_route) or stored_atom_is_full + atom_is_full = cutlass.Boolean(False) + if lane_idx < cutlass.Int32(cfg.logical_origins_per_route): + origin = cutlass.Int32(route_workspace[route_metadata_word_index + lane_idx]) + atom_is_full = cutlass.Boolean( + atom_is_valid and origin <= seq_len_kv - cutlass.Int32(cfg.atom_size) + ) + atom_valid_mask = cutlass.Int32(cute.arch.vote_ballot_sync(atom_is_valid)) + structural_full = cute.arch.vote_all_sync( + lane_idx >= cutlass.Int32(cfg.logical_origins_per_route) or atom_is_full ) - route_is_full = structural_route_is_full - if cutlass.const_expr(cfg.has_token_bits): - token_word = cutlass.Uint32(0) - if cutlass.const_expr(cfg.atom_size <= 32): - token_chunk = _load_atom_token_chunk( + + score_words_are_full = cutlass.Boolean(True) + if cutlass.const_expr(cfg.stores_score_words): + score_word = cutlass.Uint32(0) + if lane_idx < cutlass.Int32(cfg.token_words_per_route): + score_word = _load_exact_score_word( + route_workspace, kv_valid_bits, + route_metadata_word_index, + lane_idx, batch_idx, - logical_origin, - logical_origin_is_valid, - cfg.atom_size, seq_len_kv, + cfg, ) - atoms_per_word = 32 // cfg.atom_size - if lane_idx < cutlass.Int32(cfg.logical_origins_per_route): - atom_in_word = lane_idx % cutlass.Int32(atoms_per_word) - token_word = token_chunk << ( - atom_in_word * cutlass.Int32(cfg.atom_size) - ) - active_origin_lanes = (1 << cfg.logical_origins_per_route) - 1 - for shuffle_step in cutlass.range_constexpr( - int(math.log2(atoms_per_word)) - ): - peer_word = cutlass.Uint32( - prims.shfl_sync( - thread_mask=active_origin_lanes, - val=token_word, - offset=1 << shuffle_step, - mask_and_clamp=0x1F, - kind=prims.Shfl.BFLY, - ) - ) - token_word = token_word | peer_word - if atom_in_word == cutlass.Int32(0): - logical_word_idx = lane_idx // cutlass.Int32(atoms_per_word) - route_workspace[ - route_metadata_word_index - + cutlass.Int32(cfg.token_words_word_offset) - + logical_word_idx - ] = cutlass.Int32(token_word) - full_atom_mask = cutlass.Uint32((1 << cfg.atom_size) - 1) - token_route_is_full = cute.arch.vote_all_sync( - lane_idx >= cutlass.Int32(cfg.logical_origins_per_route) - or token_chunk == full_atom_mask - ) - else: - if lane_idx < cutlass.Int32(cfg.token_words_per_route): - token_word = _load_coarse_token_word( - block_indices, - kv_valid_bits, - row_begin, - row_end, - route_idx, - lane_idx, - batch_idx, - cfg.kv_block_size, - cfg.atom_size, - cfg.logical_origins_per_route, - seq_len_kv, - ) - route_workspace[ - route_metadata_word_index - + cutlass.Int32(cfg.token_words_word_offset) - + lane_idx - ] = cutlass.Int32(token_word) - token_route_is_full = cute.arch.vote_all_sync( - lane_idx >= cutlass.Int32(cfg.token_words_per_route) - or token_word == cutlass.Uint32(0xFFFFFFFF) - ) - route_is_full = cutlass.Boolean( - structural_route_is_full and token_route_is_full + route_workspace[ + route_metadata_word_index + + cutlass.Int32(cfg.token_words_word_offset) + + lane_idx + ] = cutlass.Int32(score_word) + score_words_are_full = cute.arch.vote_all_sync( + lane_idx >= cutlass.Int32(cfg.token_words_per_route) + or score_word == cutlass.Uint32(0xFFFFFFFF) ) - if lane_idx == cutlass.Int32(0): route_workspace[ route_metadata_word_index + cutlass.Int32(cfg.atom_valid_mask_word_offset) - ] = stored_atom_valid_mask + ] = atom_valid_mask route_workspace[ route_metadata_word_index + cutlass.Int32(cfg.route_flags_word_offset) ] = ( cutlass.Int32(_PREPARED_ROUTE_IS_FULL_FLAG) - if route_is_full + if structural_full and score_words_are_full else cutlass.Int32(0) ) -@cute.jit -def _paged_request_page_range_is_valid( - request_begin: cutlass.Int32, - request_end: cutlass.Int32, - num_indices: cutlass.Int32, - required_pages: cutlass.Int32, -) -> cutlass.Boolean: - """Validate one request's page-table range before any index load.""" - - return cutlass.Boolean( - request_begin >= cutlass.Int32(0) - and request_begin <= request_end - and request_end <= num_indices - and request_end - request_begin >= required_pages - ) - - @cute.jit def _resolve_paged_route_atom_page_id( - paged_kv_indices: cute.Tensor, - request_begin: cutlass.Int32, + block_tables: cute.Tensor, + batch_idx: cutlass.Int32, + block_table_row_stride: cutlass.Int64, logical_origin: cutlass.Int32, logical_origin_is_valid: cutlass.Boolean, lane_idx: cutlass.Int32, page_size: cutlass.Constexpr[int], num_physical_kv_pages: cutlass.Int64, ) -> cutlass.Int32: - """Resolve one trusted selected logical atom to its raw physical page ID.""" + """Resolve one trusted selected logical atom to its physical page ID.""" physical_page_id = cutlass.Int32(-1) page_id_is_valid = cutlass.Boolean(True) if logical_origin_is_valid: logical_page_idx = logical_origin // cutlass.Int32(page_size) - candidate_page_id = cutlass.Int32( - paged_kv_indices[request_begin + logical_page_idx] + physical_page_id = cutlass.Int32( + block_tables.iterator[ + cutlass.Int64(batch_idx) * block_table_row_stride + + cutlass.Int64(logical_page_idx) + ] ) - physical_page_id = candidate_page_id page_id_is_valid = cutlass.Boolean( - candidate_page_id >= cutlass.Int32(0) - and cutlass.Int64(candidate_page_id) < num_physical_kv_pages + physical_page_id >= cutlass.Int32(0) + and cutlass.Int64(physical_page_id) < num_physical_kv_pages ) page_ids_are_valid = cute.arch.vote_all_sync(page_id_is_valid) if lane_idx == cutlass.Int32(0): runtime_assert( page_ids_are_valid, - "paged_kv_indices contains an out-of-range physical page ID", + "block_tables contains an out-of-range physical page ID", ) return physical_page_id -class _PrepareBlockSparseRoutes: - """Prepare contiguous or paged sparse routes for one static geometry.""" +@cute.jit +def _exact_lane_rank( + exact_ballot: cutlass.Uint32, + lane_idx: cutlass.Int32, + exact_prefix: cutlass.Int32, +) -> cutlass.Int32: + """Return one exact lane's global semantic-block rank.""" + + lower_lane_mask = (cutlass.Uint32(1) << lane_idx) - cutlass.Uint32(1) + return exact_prefix + cutlass.Int32(cute.arch.popc(exact_ballot & lower_lane_mask)) + + +@cute.jit +def _emit_exact_block_atoms( + route_workspace: cute.Tensor, + row_route_begin: cutlass.Int32, + semantic_block_idx: cutlass.Int32, + exact_block_rank: cutlass.Int32, + cfg: cutlass.Constexpr[_RouteConfig], +) -> None: + """Expand one bitmask-selected block into fixed row-global atom slots.""" + + first_atom_rank = exact_block_rank * cutlass.Int32(cfg.atoms_per_block) + atom_in_block = cutlass.Int32(0) + while atom_in_block < cutlass.Int32(cfg.atoms_per_block): + atom_rank = first_atom_rank + atom_in_block + route_idx = atom_rank // cutlass.Int32(cfg.logical_origins_per_route) + atom_in_route = atom_rank % cutlass.Int32(cfg.logical_origins_per_route) + route_word_index = cutlass.Int32(cfg.route_metadata_base_word_offset) + ( + (row_route_begin + route_idx) + * cutlass.Int32(cfg.route_metadata_stride_words) + ) + logical_origin = semantic_block_idx * cutlass.Int32( + cfg.kv_block_size + ) + atom_in_block * cutlass.Int32(cfg.atom_size) + stored_origin = cutlass.Int32(-1) + if logical_origin < cutlass.Int32(cfg.seq_len_kv): + stored_origin = logical_origin + route_workspace[route_word_index + atom_in_route] = stored_origin + atom_in_block += cutlass.Int32(1) + + +@cute.jit +def _load_bitmask_word( + exact_block_bits: cute.Tensor, + batch_idx: cutlass.Int32, + kv_head_idx: cutlass.Int32, + q_block_idx: cutlass.Int32, + logical_word_idx: cutlass.Int32, + cfg: cutlass.Constexpr[_RouteConfig], + for_proxy: cutlass.Constexpr[bool], +) -> cutlass.Uint32: + """Load one in-range exact or proxy semantic-block word.""" + + valid_word = _low_bits_mask( + cutlass.Int32(cfg.num_kv_blocks) - logical_word_idx * cutlass.Int32(_WARP_SIZE) + ) + selected_word = cutlass.Uint32( + exact_block_bits[batch_idx, kv_head_idx, q_block_idx, logical_word_idx] + ) + if cutlass.const_expr(for_proxy): + selected_word = ~selected_word + return valid_word & selected_word + + +@cute.jit +def _load_bsr_proxy_word( + block_indices: cute.Tensor, + row_begin: cutlass.Int32, + row_end: cutlass.Int32, + logical_word_idx: cutlass.Int32, + cfg: cutlass.Constexpr[_RouteConfig], +) -> cutlass.Uint32: + """Build one proxy word from a canonical sorted-BSR interval.""" + + word_begin = logical_word_idx * cutlass.Int32(_WARP_SIZE) + valid_word = _low_bits_mask(cutlass.Int32(cfg.num_kv_blocks) - word_begin) + selected_word = cutlass.Uint32(0) + lower = row_begin + upper = row_end + while lower < upper: + middle = lower + (upper - lower) // cutlass.Int32(2) + if cutlass.Int32(block_indices[middle]) < word_begin: + lower = middle + cutlass.Int32(1) + else: + upper = middle + cursor = lower + word_end = word_begin + cutlass.Int32(_WARP_SIZE) + scanning = cutlass.Boolean(True) + while cursor < row_end and scanning: + block_idx = cutlass.Int32(block_indices[cursor]) + if block_idx < word_end: + selected_word = selected_word | ( + cutlass.Uint32(1) << (block_idx - word_begin) + ) + cursor += cutlass.Int32(1) + else: + scanning = cutlass.Boolean(False) + return valid_word & ~selected_word + + +@cute.jit +def _emit_proxy_route( + route_workspace: cute.Tensor, + row_route_begin: cutlass.Int32, + exact_route_count: cutlass.Int32, + group_idx: cutlass.Int32, + proxy_word: cutlass.Uint32, + lane_idx: cutlass.Int32, + cfg: cutlass.Constexpr[_RouteConfig], +) -> None: + """Emit one fixed summary-group proxy record, including an empty mask.""" + + route_metadata_word_index = cutlass.Int32(cfg.route_metadata_base_word_offset) + ( + row_route_begin + exact_route_count + group_idx + ) * cutlass.Int32(cfg.route_metadata_stride_words) + group_start = group_idx * cutlass.Int32(cfg.token_words_per_route * _WARP_SIZE) + group_size = cutlass.Int32(cfg.num_kv_blocks) - group_start + if group_size > cutlass.Int32(cfg.token_words_per_route * _WARP_SIZE): + group_size = cutlass.Int32(cfg.token_words_per_route * _WARP_SIZE) + origin_is_valid = cutlass.Boolean(False) + if lane_idx < cutlass.Int32(cfg.logical_origins_per_route): + summary_origin = group_start + lane_idx * cutlass.Int32(cfg.atom_size) + origin_is_valid = cutlass.Boolean(summary_origin < cfg.num_kv_blocks) + stored_origin = cutlass.Int32(-1) + if origin_is_valid: + stored_origin = summary_origin + route_workspace[route_metadata_word_index + lane_idx] = stored_origin + atom_valid_mask = cutlass.Int32(cute.arch.vote_ballot_sync(origin_is_valid)) + if lane_idx < cutlass.Int32(cfg.token_words_per_route): + route_workspace[ + route_metadata_word_index + + cutlass.Int32(cfg.token_words_word_offset) + + lane_idx + ] = cutlass.Int32(proxy_word) + score_full = cute.arch.vote_all_sync( + lane_idx >= cutlass.Int32(cfg.token_words_per_route) + or proxy_word == cutlass.Uint32(0xFFFFFFFF) + ) + if lane_idx == cutlass.Int32(0): + route_workspace[ + route_metadata_word_index + cutlass.Int32(cfg.atom_valid_mask_word_offset) + ] = atom_valid_mask + proxy_is_full = cutlass.Boolean( + group_size == cutlass.Int32(cfg.token_words_per_route * _WARP_SIZE) + and score_full + ) + route_workspace[ + route_metadata_word_index + cutlass.Int32(cfg.route_flags_word_offset) + ] = cutlass.Int32(_PREPARED_ROUTE_IS_PROXY_FLAG) | ( + cutlass.Int32(proxy_is_full) * cutlass.Int32(_PREPARED_ROUTE_IS_FULL_FLAG) + ) + + +class _PrepareRoutesBase: + """Own shared route geometry and compile-time storage/policy flags.""" def __init__( self, @@ -511,38 +588,58 @@ def __init__( q_block_size: int, kv_block_size: int, kv_route_size: int, - has_token_bits: bool, + use_proxy_routes: bool, + use_causal_mask: bool = False, + apply_token_mask: bool = False, + store_score_words: bool = False, page_size: int | None = None, - mask_type: str, ) -> None: - if mask_type not in ("dense", "causal"): - raise ValueError(f"unsupported mask_type: {mask_type}") - num_q_block_rows = (seq_len_q + q_block_size - 1) // q_block_size - num_rows = batch_size * num_kv_heads * num_q_block_rows + if not isinstance(use_proxy_routes, bool): + raise TypeError("use_proxy_routes must be a bool") + if not isinstance(apply_token_mask, bool): + raise TypeError("apply_token_mask must be a bool") + if not isinstance(store_score_words, bool): + raise TypeError("store_score_words must be a bool") + if not isinstance(use_causal_mask, bool): + raise TypeError("use_causal_mask must be a bool") + if use_proxy_routes and page_size is not None: + raise ValueError("paged KV does not support proxy routes") + + num_q_blocks = (seq_len_q + q_block_size - 1) // q_block_size + num_rows = batch_size * num_kv_heads * num_q_blocks + # Structural score words (sequence tail, invalid atoms) can be stored + # without a caller token mask; proxy routes and token masks require them. + stores_score_words = use_proxy_routes or apply_token_mask or store_score_words layout = _BlockSparseRouteLayout.create( kv_route_size=kv_route_size, kv_block_size=kv_block_size, page_size=page_size, - has_token_bits=has_token_bits, + has_token_bits=stores_score_words, route_metadata_capacity=0, num_rows=num_rows, ) - self.cfg = _PreparedRouteConfig.create( + self.route_layout = layout + self.cfg = _RouteConfig.create( layout=layout, num_kv_heads=num_kv_heads, seq_len_q=seq_len_q, seq_len_kv=seq_len_kv, q_block_size=q_block_size, kv_block_size=kv_block_size, + apply_token_mask=apply_token_mask, + use_proxy_routes=use_proxy_routes, ) - self.route_layout = layout self.page_size = page_size if page_size is not None else 1 - self.minimum_seq_len_kv = seq_len_q if mask_type == "causal" else 1 + self.minimum_seq_len_kv = seq_len_q if use_causal_mask else 1 self.physical_page_ids_word_offset = ( layout.physical_page_ids_word_offset if layout.is_paged else 0 ) self.route_metadata_base_word_offset = layout.route_metadata_base_word_offset + +class _PrepareBsrRoutes(_PrepareRoutesBase): + """Prepare continuous exact/proxy or paged exact routes from one BSR flow.""" + @cute.jit def __call__( self, @@ -550,24 +647,24 @@ def __call__( block_indices: cute.Tensor, kv_valid_bits: cute.Tensor, seq_lens_kv: cute.Tensor | None, - paged_kv_indptr: cute.Tensor | None, - paged_kv_indices: cute.Tensor | None, + block_tables: cute.Tensor | None, num_physical_kv_pages: cutlass.Int64, + block_table_row_stride: cutlass.Int64, row_route_offsets: cute.Tensor, route_workspace: cute.Tensor, max_blocks_per_row: cutlass.Int32, stream: cuda_drv.CUstream, ) -> None: - """Launch four independent row preparers per CTA.""" + """Launch four independent BSR row preparers per CTA.""" self.kernel( block_indptr, block_indices, kv_valid_bits, seq_lens_kv, - paged_kv_indptr, - paged_kv_indices, + block_tables, num_physical_kv_pages, + block_table_row_stride, row_route_offsets, route_workspace, max_blocks_per_row, @@ -588,14 +685,14 @@ def kernel( block_indices: cute.Tensor, kv_valid_bits: cute.Tensor, seq_lens_kv: cute.Tensor | None, - paged_kv_indptr: cute.Tensor | None, - paged_kv_indices: cute.Tensor | None, + block_tables: cute.Tensor | None, num_physical_kv_pages: cutlass.Int64, + block_table_row_stride: cutlass.Int64, row_route_offsets: cute.Tensor, route_workspace: cute.Tensor, max_blocks_per_row: cutlass.Int32, ) -> None: - """Pack logical routes and, when paged, translate physical locators.""" + """Assert trusted inputs, emit routes, resolve storage, then publish.""" thread_idx, _, _ = cute.arch.thread_idx() block_idx, _, _ = cute.arch.block_idx() @@ -613,8 +710,8 @@ def kernel( self.cfg, ) - request_begin = cutlass.Int32(0) live_seq_len_kv = cutlass.Int32(self.cfg.seq_len_kv) + selected_block_count = row_end - row_begin if cutlass.const_expr(self.route_layout.is_paged): raw_seq_len_kv = cutlass.Int32(self.cfg.seq_len_kv) if lane_idx == cutlass.Int32(0) and row_is_valid: @@ -624,110 +721,366 @@ def kernel( and raw_seq_len_kv <= cutlass.Int32(self.cfg.seq_len_kv), "seq_lens_kv is outside the planned live-length range", ) - raw_seq_len_kv = _warp_broadcast_i32(raw_seq_len_kv, 0) - live_seq_len_kv = raw_seq_len_kv + live_seq_len_kv = _warp_broadcast_i32(raw_seq_len_kv, 0) + + if ( + lane_idx == cutlass.Int32(0) + and row_is_valid + and selected_block_count > cutlass.Int32(0) + ): + last_block_idx = cutlass.Int32( + block_indices[row_end - cutlass.Int32(1)] + ) + runtime_assert( + last_block_idx * cutlass.Int32(self.cfg.kv_block_size) + < live_seq_len_kv, + "block_indices row exceeds the live KV block range", + ) if lane_idx == cutlass.Int32(0) and row_is_valid: required_pages = _positive_i32_ceil_div( live_seq_len_kv, self.page_size, ) - request_begin = cutlass.Int32(paged_kv_indptr[batch_idx]) - request_end = cutlass.Int32( - paged_kv_indptr[batch_idx + cutlass.Int32(1)] - ) - metadata_starts_at_zero = cutlass.Boolean( - paged_kv_indptr[cutlass.Int32(0)] == cutlass.Int32(0) - ) runtime_assert( - metadata_starts_at_zero - and _paged_request_page_range_is_valid( - request_begin, - request_end, - cutlass.Int32(cute.size(paged_kv_indices)), - required_pages, - ), - "paged_kv_indptr row lacks the required live page capacity", + required_pages <= cutlass.Int32(block_tables.shape[1]), + "block_tables row lacks the required live page capacity", ) - request_begin = _warp_broadcast_i32(request_begin, 0) - route_count, row_route_begin = _publish_prepared_route_count( - block_indices, + _, exact_route_count, total_route_count = _prepared_route_counts( + selected_block_count, + self.cfg, + ) + if lane_idx == cutlass.Int32(0) and row_is_valid: + runtime_assert( + selected_block_count <= max_blocks_per_row, + "selected BSR blocks exceed planned semantic capacity", + ) + row_route_begin = _prepared_row_route_begin( row_route_offsets, - route_workspace, - row_begin, - row_end, linear_row_idx, lane_idx, row_is_valid, - max_blocks_per_row, - live_seq_len_kv, - self.cfg, + total_route_count, ) - route_idx = cutlass.Int32(0) - while route_idx < route_count: - route_ordinal = row_route_begin + route_idx - route_metadata_word_index = cutlass.Int32( - self.cfg.route_metadata_base_word_offset - ) + route_ordinal * cutlass.Int32(self.cfg.route_metadata_stride_words) - logical_origin = cutlass.Int32(-1) - logical_origin_is_valid = cutlass.Boolean(False) - physical_page_id = cutlass.Int32(-1) - atom_is_full = cutlass.Boolean(False) - if lane_idx < cutlass.Int32(self.cfg.logical_origins_per_route): - ( - logical_origin, - logical_origin_is_valid, - ) = _resolve_route_logical_atom_origin( - block_indices, - row_begin, - row_end, - route_idx, + if row_is_valid: + route_idx = cutlass.Int32(0) + while route_idx < exact_route_count: + route_word_index = cutlass.Int32( + self.cfg.route_metadata_base_word_offset + ) + (row_route_begin + route_idx) * cutlass.Int32( + self.cfg.route_metadata_stride_words + ) + logical_origin = cutlass.Int32(-1) + logical_origin_is_valid = cutlass.Boolean(False) + physical_page_id = cutlass.Int32(-1) + if lane_idx < cutlass.Int32(self.cfg.logical_origins_per_route): + ( + logical_origin, + logical_origin_is_valid, + ) = _resolve_route_logical_atom_origin( + block_indices, + row_begin, + row_end, + route_idx, + lane_idx, + self.cfg.kv_block_size, + self.cfg.atom_size, + self.cfg.logical_origins_per_route, + live_seq_len_kv, + ) + if cutlass.const_expr(self.route_layout.is_paged): + physical_page_id = _resolve_paged_route_atom_page_id( + block_tables, + batch_idx, + block_table_row_stride, + logical_origin, + logical_origin_is_valid, + lane_idx, + self.page_size, + num_physical_kv_pages, + ) + if lane_idx < cutlass.Int32(self.cfg.logical_origins_per_route): + route_workspace[route_word_index + lane_idx] = logical_origin + if cutlass.const_expr(self.route_layout.is_paged): + route_workspace[ + route_word_index + + cutlass.Int32(self.physical_page_ids_word_offset) + + lane_idx + ] = physical_page_id + cute.arch.sync_warp() + + _finalize_exact_route( + route_workspace, + kv_valid_bits, + route_word_index, + batch_idx, lane_idx, - self.cfg.kv_block_size, - self.cfg.atom_size, - self.cfg.logical_origins_per_route, + logical_origin_is_valid, live_seq_len_kv, + self.cfg, ) - if cutlass.const_expr(self.route_layout.is_paged): - physical_page_id = _resolve_paged_route_atom_page_id( - paged_kv_indices, - request_begin, - logical_origin, - logical_origin_is_valid, - lane_idx, - self.page_size, - num_physical_kv_pages, + route_idx += cutlass.Int32(1) + + if cutlass.const_expr(self.cfg.use_proxy_routes): + group_idx = cutlass.Int32(0) + while group_idx < cutlass.Int32(self.cfg.num_proxy_groups): + proxy_word = cutlass.Uint32(0) + logical_word_idx = ( + group_idx * cutlass.Int32(self.cfg.token_words_per_route) + + lane_idx + ) + if lane_idx < cutlass.Int32(self.cfg.token_words_per_route): + if logical_word_idx < cutlass.Int32(self.cfg.num_exact_words): + proxy_word = _load_bsr_proxy_word( + block_indices, + row_begin, + row_end, + logical_word_idx, + self.cfg, + ) + _emit_proxy_route( + route_workspace, + row_route_begin, + exact_route_count, + group_idx, + proxy_word, + lane_idx, + self.cfg, + ) + group_idx += cutlass.Int32(1) + + if lane_idx == cutlass.Int32(0) and row_is_valid: + route_workspace[linear_row_idx] = total_route_count + + +class _PrepareBitmaskRoutes(_PrepareRoutesBase): + """Lower packed exact-block bits to continuous exact-first routes.""" + + def __init__( + self, + *, + batch_size: int, + num_kv_heads: int, + seq_len_q: int, + seq_len_kv: int, + q_block_size: int, + kv_block_size: int, + kv_route_size: int, + use_proxy_routes: bool, + use_causal_mask: bool = False, + apply_token_mask: bool = False, + store_score_words: bool = False, + ) -> None: + super().__init__( + batch_size=batch_size, + num_kv_heads=num_kv_heads, + seq_len_q=seq_len_q, + seq_len_kv=seq_len_kv, + q_block_size=q_block_size, + kv_block_size=kv_block_size, + kv_route_size=kv_route_size, + use_proxy_routes=use_proxy_routes, + use_causal_mask=use_causal_mask, + apply_token_mask=apply_token_mask, + store_score_words=store_score_words, + ) + + @cute.jit + def __call__( + self, + exact_block_bits: cute.Tensor, + kv_valid_bits: cute.Tensor, + row_route_offsets: cute.Tensor, + route_workspace: cute.Tensor, + max_blocks_per_row: cutlass.Int32, + stream: cuda_drv.CUstream, + ) -> None: + self.kernel( + exact_block_bits, + kv_valid_bits, + row_route_offsets, + route_workspace, + max_blocks_per_row, + ).launch( + grid=[ + (self.cfg.num_rows + _WARPS_PER_CTA - 1) // _WARPS_PER_CTA, + 1, + 1, + ], + block=[_THREADS_PER_CTA, 1, 1], + stream=stream, + ) + + @cute.kernel + def kernel( + self, + exact_block_bits: cute.Tensor, + kv_valid_bits: cute.Tensor, + row_route_offsets: cute.Tensor, + route_workspace: cute.Tensor, + max_blocks_per_row: cutlass.Int32, + ) -> None: + """Pack one bitmask row after proving its complete payload fits.""" + + thread_idx, _, _ = cute.arch.thread_idx() + block_idx, _, _ = cute.arch.block_idx() + warp_idx = thread_idx // _WARP_SIZE + lane_idx = thread_idx % _WARP_SIZE + linear_row_idx = block_idx * _WARPS_PER_CTA + warp_idx + row_is_valid = linear_row_idx < self.cfg.num_rows + q_block_idx = linear_row_idx % self.cfg.num_q_blocks + linear_batch_head_idx = linear_row_idx // self.cfg.num_q_blocks + kv_head_idx = linear_batch_head_idx % self.cfg.num_kv_heads + batch_idx = linear_batch_head_idx // self.cfg.num_kv_heads + + lane_exact_count = cutlass.Int32(0) + word_idx = lane_idx + while word_idx < cutlass.Int32(self.cfg.num_exact_words): + if row_is_valid: + exact_word = _load_bitmask_word( + exact_block_bits, + batch_idx, + kv_head_idx, + q_block_idx, + word_idx, + self.cfg, + for_proxy=False, + ) + lane_exact_count += cutlass.Int32(cute.arch.popc(exact_word)) + word_idx += cutlass.Int32(_WARP_SIZE) + exact_block_count = cutlass.Int32( + cute.arch.warp_redux_sync(lane_exact_count, "add") + ) + + exact_atom_count, exact_route_count, total_route_count = _prepared_route_counts( + exact_block_count, + self.cfg, + ) + if lane_idx == cutlass.Int32(0) and row_is_valid: + runtime_assert( + exact_block_count <= max_blocks_per_row, + "selected bitmask blocks exceed planned semantic capacity", + ) + row_route_begin = _prepared_row_route_begin( + row_route_offsets, + linear_row_idx, + lane_idx, + row_is_valid, + total_route_count, + ) + + if row_is_valid: + exact_prefix = cutlass.Int32(0) + word_idx = cutlass.Int32(0) + while word_idx < cutlass.Int32(self.cfg.num_exact_words): + exact_word_i32 = cutlass.Int32(0) + if lane_idx == cutlass.Int32(0): + exact_word_i32 = _load_bitmask_word( + exact_block_bits, + batch_idx, + kv_head_idx, + q_block_idx, + word_idx, + self.cfg, + for_proxy=False, + ).bitcast(cutlass.Int32) + exact_word = _warp_broadcast_i32(exact_word_i32, 0).bitcast( + cutlass.Uint32 ) - if logical_origin_is_valid: - atom_is_full = cutlass.Boolean( - logical_origin - <= live_seq_len_kv - cutlass.Int32(self.cfg.atom_size) + is_exact = cutlass.Boolean( + (exact_word & (cutlass.Uint32(1) << lane_idx)) != cutlass.Uint32(0) ) - if lane_idx < cutlass.Int32(self.cfg.logical_origins_per_route): - route_workspace[route_metadata_word_index + lane_idx] = logical_origin - if cutlass.const_expr(self.route_layout.is_paged): - route_workspace[ - route_metadata_word_index - + cutlass.Int32(self.physical_page_ids_word_offset) - + lane_idx - ] = physical_page_id + exact_ballot = cute.arch.vote_ballot_sync(is_exact).bitcast( + cutlass.Uint32 + ) + exact_rank = _exact_lane_rank(exact_ballot, lane_idx, exact_prefix) + if is_exact: + _emit_exact_block_atoms( + route_workspace, + row_route_begin, + word_idx * cutlass.Int32(_WARP_SIZE) + lane_idx, + exact_rank, + self.cfg, + ) + exact_prefix += cutlass.Int32(cute.arch.popc(exact_ballot)) + word_idx += cutlass.Int32(1) - _store_prepared_route_validity( - block_indices, - kv_valid_bits, - route_workspace, - row_begin, - row_end, - route_idx, - batch_idx, - lane_idx, - logical_origin, - logical_origin_is_valid, - atom_is_full, - route_metadata_word_index, - live_seq_len_kv, - self.cfg, + final_route_atom_count = exact_atom_count % cutlass.Int32( + self.cfg.logical_origins_per_route ) - route_idx = route_idx + cutlass.Int32(1) + if final_route_atom_count != cutlass.Int32(0): + if lane_idx >= final_route_atom_count and lane_idx < cutlass.Int32( + self.cfg.logical_origins_per_route + ): + final_route_word_index = cutlass.Int32( + self.cfg.route_metadata_base_word_offset + ) + (row_route_begin + exact_route_count - cutlass.Int32(1)) * ( + cutlass.Int32(self.cfg.route_metadata_stride_words) + ) + route_workspace[final_route_word_index + lane_idx] = cutlass.Int32( + -1 + ) + cute.arch.sync_warp() + + route_idx = cutlass.Int32(0) + while route_idx < exact_route_count: + route_word_index = cutlass.Int32( + self.cfg.route_metadata_base_word_offset + ) + (row_route_begin + route_idx) * cutlass.Int32( + self.cfg.route_metadata_stride_words + ) + atom_is_valid = cutlass.Boolean(False) + if lane_idx < cutlass.Int32(self.cfg.logical_origins_per_route): + atom_is_valid = cutlass.Boolean( + cutlass.Int32(route_workspace[route_word_index + lane_idx]) + >= cutlass.Int32(0) + ) + _finalize_exact_route( + route_workspace, + kv_valid_bits, + route_word_index, + batch_idx, + lane_idx, + atom_is_valid, + cutlass.Int32(self.cfg.seq_len_kv), + self.cfg, + ) + route_idx += cutlass.Int32(1) + + if cutlass.const_expr(self.cfg.use_proxy_routes): + group_idx = cutlass.Int32(0) + while group_idx < cutlass.Int32(self.cfg.num_proxy_groups): + proxy_word = cutlass.Uint32(0) + logical_word_idx = ( + group_idx * cutlass.Int32(self.cfg.token_words_per_route) + + lane_idx + ) + if lane_idx < cutlass.Int32(self.cfg.token_words_per_route): + if logical_word_idx < cutlass.Int32(self.cfg.num_exact_words): + proxy_word = _load_bitmask_word( + exact_block_bits, + batch_idx, + kv_head_idx, + q_block_idx, + logical_word_idx, + self.cfg, + for_proxy=True, + ) + _emit_proxy_route( + route_workspace, + row_route_begin, + exact_route_count, + group_idx, + proxy_word, + lane_idx, + self.cfg, + ) + group_idx += cutlass.Int32(1) + + if lane_idx == cutlass.Int32(0) and row_is_valid: + route_workspace[linear_row_idx] = total_route_count + + +__all__ = ["_PrepareBitmaskRoutes", "_PrepareBsrRoutes"] diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_config.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_config.py index c67e4c83bf10..0626410d1f10 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_config.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_config.py @@ -102,6 +102,16 @@ (Float16, Float16, Float16, 256, 128, 1, 1), } +# Per-thread register budgets for the Q64/KV256 warp groups once the launch +# bound enables ``setmaxnreg``. The MMA, load, and scheduler warps keep 56, so +# the two softmax groups and the correction group share the remainder: +# 8 * softmax + 4 * correction = 65536 / 32 - 4 * 56. An even split measured +# fastest on B200 for both the static grid and the persistent scheduler; the +# correction group needs the extra room for its persistent bookkeeping and +# the KV256 tail merge, while the rolled softmax fragment loop needs less. +KV_TILE_256_SOFTMAX_TASK_REGISTERS = 152 +KV_TILE_256_CORRECTION_TASK_REGISTERS = 152 + _KV_TILE_256_PHYSICAL_DEFAULTS: Mapping[str, ConfigValue] = { "tmem_s_cols": 128, "tmem_stats_cols": 32, @@ -505,6 +515,10 @@ class FmhaDecodeConfig: # restricted to 8/16/32 or positive multiples of 64 and are assembled into # a profile-selected fixed KV128 or KV256 route. use_block_sparse: bool = False + # Interpret prepared records as a typed proxy/exact stream. Proxy records + # source semantic-block summaries while exact records retain the K/V + # atom path. Source selection is orthogonal to the physical Q/KV profile. + use_block_sparse_proxy_routes: bool = False q_block_size: int = 0 kv_block_size: int = 0 # Optional batch-wide physical-token validity metadata shared by every head @@ -644,13 +658,13 @@ def uses_task_register_reallocation(self) -> bool: def softmax_task_num_registers(self) -> int | None: if not self.uses_task_register_reallocation: return None - return 176 if self.tile_size_kv == 256 else 184 + return KV_TILE_256_SOFTMAX_TASK_REGISTERS if self.tile_size_kv == 256 else 184 @property def correction_task_num_registers(self) -> int | None: if not self.uses_task_register_reallocation: return None - return 104 if self.tile_size_kv == 256 else 88 + return KV_TILE_256_CORRECTION_TASK_REGISTERS if self.tile_size_kv == 256 else 88 @property def mma_load_task_num_registers(self) -> int | None: @@ -1066,12 +1080,13 @@ def num_s_regs_per_thread(self) -> int: def softmax_score_fragment_regs(self) -> int: """Return the maximum score fragment kept live in registers. - KV256 owns 128 score values per lane but streams them as four native - 32-register LDTM atoms. Other profiles retain their complete score - fragment, so this property is intentionally distinct from - ``num_s_regs_per_thread`` (the total logical ownership). + Streamed profiles own 128 score values per lane but process them as + four native 32-register LDTM atoms. Other profiles retain their + complete score fragment, so this property is intentionally distinct + from ``num_s_regs_per_thread`` (the total logical ownership). The + selection of streamed profiles lives in ``streams_tmem_p_fragments``. """ - if self.tile_size_kv == 256: + if self.streams_tmem_p_fragments: return 32 return self.num_s_regs_per_thread @@ -1080,6 +1095,30 @@ def num_softmax_score_fragments(self) -> int: """Return score fragments used to cover one logical KV tile.""" return self.num_s_regs_per_thread // self.softmax_score_fragment_regs + @property + def block_sparse_kv_atom_size(self) -> int: + """Return the K token span of one block-sparse route origin.""" + assert self.use_block_sparse + return _block_sparse_kv_atom_size(self.kv_block_size) + + @property + def softmax_fragments_per_route_atom(self) -> int: + """Return the streamed score fragments that share one route origin. + + Route origins are staged per K64 atom, so a 128-token KV block spans + two origins; the fragment-to-origin mapping follows the atom. + """ + return self.block_sparse_kv_atom_size // self.softmax_score_fragment_regs + + @property + def uses_ws_2x2_datapath(self) -> bool: + """Whether QK and PV issue the WS 2x2 instruction over two lane halves. + + KV256 exposes two spatial KV128 partials per logical Q row; every other + Keeps profile issues the plain CTA-local instruction. + """ + return self.tile_size_kv == 256 + @property def num_packed_p_regs(self) -> int: """Return packed P registers stored by each softmax producer lane.""" @@ -1246,6 +1285,10 @@ def validate_paged_kv_staging_config(self) -> None: def validate_block_sparse_profile(self, *, heads_q_per_kv: int) -> None: """Validate the qualified host profile for block-sparse.""" + if self.use_block_sparse_proxy_routes and not self.use_block_sparse: + raise ValueError("proxy routes require block-sparse attention") + if self.use_block_sparse_proxy_routes and self.mask_type != DENSE: + raise ValueError("block-sparse proxy routes require mask_type='dense'") if not self.use_block_sparse: if self.use_parallel_sparse_kv_loads: raise ValueError( @@ -1283,6 +1326,13 @@ def validate_block_sparse_profile(self, *, heads_q_per_kv: int) -> None: "block-sparse tile_size_kv=256 requires the Q64 16-bit Keeps " "profile with coarse KV blocks and one load task" ) + if self.use_keeps_mma_ab and not self.streams_tmem_p_fragments: + # The block-sparse Keeps softmax and P passes exist only in their + # streamed K32-fragment form. + raise ValueError( + "block-sparse KeepsMmaAb requires a streamed TMEM-P profile " + "(Q64/KV256 or 16-bit Q128/KV128)" + ) if self.tile_size_q != selected_q_tile: raise ValueError( "block-sparse tile_size_q must match its grouped-Q geometry" @@ -1332,6 +1382,42 @@ def compile_signature(self) -> tuple[tuple[str, object], ...]: for config_field in fields(self) ) + @property + def uses_prepared_score_keep_words(self) -> bool: + """Whether prepared routes carry BMM1 score-column validity words. + + Dense block-sparse Keeps plans prepare them even without a caller + token mask: the streamed max pass trusts the words directly, which is + cheaper than deriving each fragment's visible range in the softmax + warps. The plan sizes its route storage and the prepare kernels store + the words from this same property, via the resolved launch spec. + """ + + return ( + self.use_kv_valid_bits + or self.use_block_sparse_proxy_routes + or ( + self.use_block_sparse + and self.use_keeps_mma_ab + and self.mask_type == DENSE + ) + ) + + @property + def trusts_prepared_score_words(self) -> bool: + """Whether prepared words fully describe dense score-column validity. + + Dense prepared routes have already combined structural tail validity + with any caller-provided exact-token bits. Their K32 words therefore + apply to exact and proxy sources alike. + """ + + return ( + self.use_block_sparse + and self.uses_prepared_score_keep_words + and self.mask_type == DENSE + ) + @property def uses_q_desc_ref(self) -> bool: """Whether QK derives Q's descriptor from shared resource state.""" @@ -1435,10 +1521,10 @@ def has_static_dense_full_kv_tiles(self) -> bool: @property def uses_ordered_softmax_barrier(self) -> bool: """Whether this profile selects the ordered P0/P1 softmax barrier.""" - if self.tile_size_kv == 256: - # KV256 uses independent four-stage P-fragment pipelines. Ordering - # the two softmax groups would serialize fragment production and - # defeat the intended P/PV overlap. + if self.streams_tmem_p_fragments: + # Streamed profiles use independent per-fragment P pipelines. + # Ordering the two softmax groups would serialize fragment + # production and defeat the intended P/PV overlap. return False if self.ordered_softmax_barrier_mode == 2: return True @@ -1554,10 +1640,11 @@ def uses_staged_one_inst_tmem_p(self) -> bool: def uses_two_inst_tmem_p(self) -> bool: """Whether a two-instance Keeps profile uses the TMEM-P overlay. - Q128/KV128 and sparse Q64/KV128 publish a complete packed-P row per - pipeline token. Q64/KV256 uses the same S-to-P aliasing contract but - streams four independently ready K32 fragments. Dense Q64/KV128 keeps - the base kernel's faster SMEM-P cadence. + FP8 Q128/KV128 and dense 16-bit Q128/KV128 publish a complete + packed-P row per pipeline token. Q64/KV256 and block-sparse 16-bit + Q128/KV128 use the same S-to-P aliasing contract but stream four + independently ready K32 fragments (see ``streams_tmem_p_fragments``). + Q64/KV128 keeps the base kernel's faster SMEM-P cadence. """ # Two-instance Keeps keeps stats outside S, so both static and persistent # work tiles can overlay P on the consumed S instance. The split K/V @@ -1568,11 +1655,6 @@ def uses_two_inst_tmem_p(self) -> bool: and ( (self.tile_size_q == 128 and self.tile_size_kv == 128) or (self.tile_size_q == 64 and self.tile_size_kv == 256) - or ( - self.use_block_sparse - and self.tile_size_q == 64 - and self.tile_size_kv == 128 - ) ) and self.head_dim_per_stage_kv == 0 and self.num_insts_kv == 2 @@ -1582,42 +1664,51 @@ def uses_two_inst_tmem_p(self) -> bool: @property def streams_tmem_p_fragments(self) -> bool: - """Whether P is published as independently ready TMEM fragments.""" - return self.uses_two_inst_tmem_p and self.num_softmax_score_fragments > 1 - - @property - def matches_kv256_task_topology(self) -> bool: - """Whether task roles match KV256's validated 16-warp layout.""" - return all( - getattr(self, field) == expected - for field, expected in _KV_TILE_256_TASK_TOPOLOGY_DEFAULTS.items() + """Whether P is published as independently ready TMEM fragments. + + Streamed profiles produce their K32 fragments from one rolled runtime + loop: the max pass writes masked scores back to TMEM, so the P pass + reloads each fragment without mask logic and the exponentiation body + exists once in the instruction stream. Each published fragment lets + the MMA warp start its PV k-slice before the row is complete, at the + cost of one barrier round per fragment. + + Streaming is limited to the 16-bit two-instance profiles whose route + loop waits on the K/V loads, where the earlier PV start hides load + latency: Q64/KV256 and block-sparse Q128/KV128. Dense Q128/KV128 + keeps the complete row because its route loop is not load-bound, so + the per-fragment barriers are not compensated. FP8 Q128 keeps the + complete row because its P publication packs four values per column + into one store. + """ + return ( + self.uses_two_inst_tmem_p + and not self.use_fp8_qkv + and (self.tile_size_kv == 256 or self.use_block_sparse) ) @property - def uses_rotating_kv256_exchange(self) -> bool: - """Whether this profile selects KV-ring scratch for correction. + def defers_softmax_anchor_updates(self) -> bool: + """Whether small row-max increases keep the previous exponent anchor. - Persistent direct output can overlap the next work tile's first two - K loads with correction by placing its exchange in the third, drained - KV stage. Split-KV and attention sinks retain the fixed exchange because - their tail storage and lifetime differ from direct output. + Keeps correction skips the in-place O rescale whenever the anchor is + unchanged, so keeping the prior anchor within + ``SOFTMAX_RESCALE_THRESHOLD_LOG2`` trades a bounded 16-bit P range + (2**8) for fewer TMEM rescales. The profiles listed here are the ones + where that trade was measured to pay: KV256 tiles and block-sparse + routes, whose row maximum moves often but rarely by much. """ - selects_persistent_kv256 = ( - self.streams_tmem_p_fragments - and self.tile_size_q == 64 - and self.tile_size_kv == 256 - and self.use_persistent_scheduler + return self.use_keeps_mma_ab and ( + self.tile_size_kv == 256 or self.use_block_sparse ) - if not selects_persistent_kv256: - return False - has_rotating_kv_ring = ( - self.num_head_dim_stages_kv == 1 - and self.kv_stages == KV_TILE_256_SHARED_FIFO_STAGES - and self.load_num_warps == 1 + @property + def matches_kv256_task_topology(self) -> bool: + """Whether task roles match KV256's validated 16-warp layout.""" + return all( + getattr(self, field) == expected + for field, expected in _KV_TILE_256_TASK_TOPOLOGY_DEFAULTS.items() ) - has_direct_output_lifetime = not (self.use_split_kv or self.use_attention_sinks) - return has_rotating_kv_ring and has_direct_output_lifetime @property def keeps_separates_tmem_s_and_stats(self) -> bool: @@ -2196,17 +2287,6 @@ def _require_python_int(field_name: str) -> int: f"{pipeline_smem_bytes} bytes, limit is " f"{pipeline_smem_budget_bytes} bytes" ) - if ( - cfg.use_persistent_scheduler - and not cfg.use_split_kv - and not cfg.use_attention_sinks - and cfg.kv_stages != KV_TILE_256_SHARED_FIFO_STAGES - ): - raise ValueError( - "persistent KV256 requires kv_stages=" - f"{KV_TILE_256_SHARED_FIFO_STAGES} for the rotating shared-KV " - f"exchange, got {cfg.kv_stages}" - ) if not cfg.supports_grouped_keeps: raise ValueError( "KV256 currently supports only the qualified Q64 FP16/BF16/D128 " diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_constants.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_constants.py index 68686febfbbb..8bf94760cf31 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_constants.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_constants.py @@ -45,8 +45,12 @@ # Keep the old maximum as the exponent reference while a new maximum is at # most eight log2 units larger. This avoids an output-correction round without # letting an intermediate probability exceed 2**8; the softmax identity is -# unchanged apart from normal finite-precision rounding. -KV_TILE_256_RESCALE_THRESHOLD_LOG2 = 8.0 +# unchanged apart from normal finite-precision rounding. As in the +# FlashInfer/TRT-LLM policy, this assumes normal model logits rather than +# adversarial values outside the qualified probability bound. Streamed KV256 +# and block-sparse Keeps profiles apply it; see +# ``FmhaDecodeConfig.defers_softmax_anchor_updates``. +SOFTMAX_RESCALE_THRESHOLD_LOG2 = 8.0 # A launch bound makes ptxas honor warpgroup ``setmaxnreg`` allocations, but # the resulting register hand-off has a fixed cost. Paired B200 measurements diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_kernel.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_kernel.py index 929696c30d3e..d2540ec61920 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_kernel.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_kernel.py @@ -56,8 +56,8 @@ from cutlass.experimental.task_scheduling.task_manager import TaskManager from ..._block_sparse.common import ( - _block_sparse_kv_atom_size, - _prepared_kv_routes_are_block_aligned, + _block_sparse_contiguous_kv_copy_geometry, + _block_sparse_proxy_summary_geometry, ) from ..._block_sparse.prepared import _BlockSparseRouteLayout from ..tensor_map import ( @@ -95,7 +95,7 @@ from .fmha_decode_tasks import ( PackedDecodeWorkQueue, ScheduleTokenThrottleResource, - SmemKvReuseCreditResource, + _prefetch_prepared_sparse_row, create_block_sparse_load_tasks_per_inst, create_correction_task, create_correction_task_one_inst_qkv, @@ -113,7 +113,6 @@ create_softmax0_task, create_softmax1_task, ) - from .reduction import ( # noqa: F401 decode_gen_separate_reduction_kernel, fmha_decode_separate_reduction_launch, @@ -322,6 +321,10 @@ def _build_decode_gen_schedule( tma_desc_v: cutlass.Pointer | None = None, tma_desc_k_atom: cutlass.Pointer | None = None, tma_desc_v_atom: cutlass.Pointer | None = None, + tma_desc_k_summary: cutlass.Pointer | None = None, + tma_desc_v_summary: cutlass.Pointer | None = None, + tma_desc_k_summary_atom: cutlass.Pointer | None = None, + tma_desc_v_summary_atom: cutlass.Pointer | None = None, page_idx_kv: cute.Pointer | None = None, h_k_idx: Int32 | None = None, b_idx: Int32 | None = None, @@ -341,6 +344,8 @@ def _build_decode_gen_schedule( sparse_row_route_offsets: cute.Pointer | None = None, sparse_row_route_counts: cute.Pointer | None = None, sparse_route_metadata: cute.Pointer | None = None, + sparse_row_route_begin: Int32 | None = None, + sparse_route_count: Int32 | None = None, ) -> tuple[ list[Task], dict[MemoryResource, list[MemoryResource]], @@ -408,6 +413,15 @@ def _build_decode_gen_schedule( "tma_desc_k_atom": tma_desc_k_atom, "tma_desc_v_atom": tma_desc_v_atom, } + if cfg.use_block_sparse_proxy_routes: + segment_tensormaps.update( + { + "tma_desc_k_summary": tma_desc_k_summary, + "tma_desc_v_summary": tma_desc_v_summary, + "tma_desc_k_summary_atom": tma_desc_k_summary_atom, + "tma_desc_v_summary_atom": tma_desc_v_summary_atom, + } + ) for name, descriptor in segment_tensormaps.items(): if descriptor is None: raise ValueError( @@ -741,7 +755,6 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: # ------------------------------------------------------------------ work_queue = None schedule_token_throttle = None - smem_kv_reuse_credit = None # CLC remains the single persistent policy for every supported topology. # The stock static WorkQueue advances and decodes coordinates separately # in every task, which regresses multi-wave decode workloads. CLC computes @@ -799,17 +812,6 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: ), name="schedule_token_throttle", ) - if cfg.uses_rotating_kv256_exchange: - smem_kv_reuse_credit = SmemKvReuseCreditResource( - cfg=cfg, - pipeline_config=PipelineConfig.create_async_async_pipeline_cfg( - num_stages=1, - producer_group=load_grp, - consumer_group=correction_grp, - cta_layout_vmnk=cta_layout, - ), - name="smem_kv_reuse_credit", - ) smem_q = SmemQResource( pipeline_config=smem_q_cfg, cfg=cfg, @@ -873,10 +875,12 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: sparse_softmax_metadata0 = None sparse_softmax_metadata1 = None if cfg.use_block_sparse: + # This selects the prepared-record storage ABI. Causal consumers still + # intersect these column-validity words with each Q row's causal mask. prepared_route_layout = _BlockSparseRouteLayout.create( kv_route_size=cfg.tile_size_kv, kv_block_size=cfg.kv_block_size, - has_token_bits=cfg.use_kv_valid_bits, + has_token_bits=cfg.uses_prepared_score_keep_words, route_metadata_capacity=0, num_rows=1, page_size=cfg.num_tokens_per_page if cfg.use_paged_kv else None, @@ -928,6 +932,10 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: tma_desc_v=tma_desc_v, tma_desc_k_atom=tma_desc_k_atom, tma_desc_v_atom=tma_desc_v_atom, + tma_desc_k_summary=tma_desc_k_summary, + tma_desc_v_summary=tma_desc_v_summary, + tma_desc_k_summary_atom=tma_desc_k_summary_atom, + tma_desc_v_summary_atom=tma_desc_v_summary_atom, sparse_kv_metadata=sparse_kv_metadata0, page_offsets_kv=smem_page_offsets, seqlens_kv=kv_seqlens, @@ -947,6 +955,10 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: tma_desc_v=tma_desc_v, tma_desc_k_atom=tma_desc_k_atom, tma_desc_v_atom=tma_desc_v_atom, + tma_desc_k_summary=tma_desc_k_summary, + tma_desc_v_summary=tma_desc_v_summary, + tma_desc_k_summary_atom=tma_desc_k_summary_atom, + tma_desc_v_summary_atom=tma_desc_v_summary_atom, sparse_kv_metadata=sparse_kv_metadata1, page_offsets_kv=smem_page_offsets, seqlens_kv=kv_seqlens, @@ -966,6 +978,10 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: tma_desc_v=tma_desc_v, tma_desc_k_atom=tma_desc_k_atom, tma_desc_v_atom=tma_desc_v_atom, + tma_desc_k_summary=tma_desc_k_summary, + tma_desc_v_summary=tma_desc_v_summary, + tma_desc_k_summary_atom=tma_desc_k_summary_atom, + tma_desc_v_summary_atom=tma_desc_v_summary_atom, sparse_kv_metadata=sparse_kv_metadata0, page_offsets_kv=smem_page_offsets_v or smem_page_offsets, seqlens_kv=kv_seqlens, @@ -985,6 +1001,10 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: tma_desc_v=tma_desc_v, tma_desc_k_atom=tma_desc_k_atom, tma_desc_v_atom=tma_desc_v_atom, + tma_desc_k_summary=tma_desc_k_summary, + tma_desc_v_summary=tma_desc_v_summary, + tma_desc_k_summary_atom=tma_desc_k_summary_atom, + tma_desc_v_summary_atom=tma_desc_v_summary_atom, sparse_kv_metadata=sparse_kv_metadata1, page_offsets_kv=smem_page_offsets_v or smem_page_offsets, seqlens_kv=kv_seqlens, @@ -1005,6 +1025,10 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: tma_desc_v=tma_desc_v, tma_desc_k_atom=tma_desc_k_atom, tma_desc_v_atom=tma_desc_v_atom, + tma_desc_k_summary=tma_desc_k_summary, + tma_desc_v_summary=tma_desc_v_summary, + tma_desc_k_summary_atom=tma_desc_k_summary_atom, + tma_desc_v_summary_atom=tma_desc_v_summary_atom, sparse_kv_metadata0=sparse_kv_metadata0, sparse_kv_metadata1=sparse_kv_metadata1, page_offsets_kv=smem_page_offsets, @@ -1213,6 +1237,8 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: "seq_len_q": seq_len_q, "sparse_row_route_offsets": sparse_row_route_offsets, "sparse_row_route_counts": sparse_row_route_counts, + "sparse_row_route_begin": sparse_row_route_begin, + "sparse_route_count": sparse_route_count, "num_heads_kv": num_heads_kv, } if use_one_inst_qkv: @@ -1283,7 +1309,6 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: smem_kv, work_queue, schedule_token_throttle, - smem_kv_reuse_credit, cfg, domain=load_domain, domain_bias=0, @@ -1431,7 +1456,6 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: tmem_corr0, tmem_corr1, work_queue, - smem_kv_reuse_credit, cfg, domain=corr_domain, tmem_stats_done0=tmem_stats_done0, @@ -1607,10 +1631,6 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: ) if schedule_token_throttle is not None: resource_dependency_graph[schedule_token_throttle] = [work_queue] - if smem_kv_reuse_credit is not None: - # A self-edge models the one-slot ownership token: Load produces it - # for the current tile and Correction consumes it before the next Load. - resource_dependency_graph[smem_kv_reuse_credit] = [smem_kv_reuse_credit] dma_consumer_release_labels: dict[ tuple[MemoryResource, MemoryResource], set[str] ] = {} @@ -1662,8 +1682,6 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: smem_allocator.add_resource(work_queue) if schedule_token_throttle is not None: smem_allocator.add_resource(schedule_token_throttle) - if smem_kv_reuse_credit is not None: - smem_allocator.add_resource(smem_kv_reuse_credit) smem_allocator.add_resource(smem_q) if smem_page_offsets is not None: smem_allocator.add_resource(smem_page_offsets) @@ -1701,21 +1719,10 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: smem_allocator.add_resource(tmem_corr0) if not use_one_inst_qkv: smem_allocator.add_resource(tmem_corr1) - if cfg.tile_size_kv == 256: - # KV256 direct-output correction rotates one compact 35,840-byte - # payload through the shared 192-KiB K/V ring. Split-KV retains the - # fixed full exchange. Neither path increases the CTA SMEM footprint. - smem_allocator.add_alias_group( - [ - smem_kv.get_smem_requirements(), - tmem_corr1.get_smem_requirements(), - ] - ) smem_allocator.add_tmem_ptr( SmemAllocation("fmha_tmem_ptr_i32", dtype=cutlass.Int32, alignment=4) ) smem_allocator.compute_layout() - tmem_allocator = TmemAllocator() if cfg.use_keeps_mma_ab: if use_one_inst_qkv: @@ -1774,13 +1781,15 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: p0_alloc = smem_p0.get_tmem_requirements()[0] p1_alloc = smem_p1.get_tmem_requirements()[0] o_alloc = tmem_o.get_tmem_requirements()[0] - if cfg.tile_size_kv == 256: - # KV256 keeps O in the low 256 columns and overlays packed P on - # each S region from its first column. Softmax streams K32 - # fragments in order, so every 16-column P store only overwrites - # scores that have already been consumed. Starting P after the - # nominal stats columns would instead clobber the next unread S - # fragment; KV256 keeps its softmax stats in SMEM. + if cfg.streams_tmem_p_fragments: + # Streamed profiles keep O in the low 256 columns and overlay + # packed P on each S region from its first column. Softmax streams + # K32 fragments in order, so every 16-column P store only + # overwrites scores that have already been consumed. Starting P + # after the nominal stats columns would instead clobber the next + # unread S fragment; streamed profiles keep their softmax stats in + # SMEM. + assert cfg.keeps_stats_via_smem o_alloc.offset = 0 s0_alloc.offset = 2 * cfg.tmem_o_stage_cols s1_alloc.offset = s0_alloc.offset + cfg.tmem_s_cols @@ -1835,12 +1844,8 @@ def _make_page_offsets_cfg(num_stages: int | None = None) -> PipelineConfig: eager_init_resources = ( [tmem_corr0] if use_one_inst_qkv else [tmem_corr0, tmem_corr1] ) - if smem_kv_reuse_credit is not None: - # Initialize the persistent ring cursor under the same CTA-wide fence - # and barrier used by other manually managed SMEM control state. - eager_init_resources.append(smem_kv_reuse_credit) - if cfg.tile_size_kv == 256: - # KV256's TMEM P operands use one-way per-fragment ready barriers. + if cfg.streams_tmem_p_fragments: + # Streamed TMEM P operands use one-way per-fragment ready barriers. # Initialize them beside correction's manually managed SMEM state. eager_init_resources.extend([smem_p0, smem_p1]) @@ -1863,10 +1868,10 @@ def _has_unmodeled_tmem_p_alias_protocol(cfg: FmhaDecodeConfig) -> bool: """Whether exhaustive TS checking would report a known false P/S race. The staged D256 path selects one of two physical P/S stages at runtime. - Static KV256 instead orders streamed P fragments with private mbarriers and + Static streamed profiles instead order P fragments with private mbarriers and reuses the matching TmemO-full barrier as the next-QK overwrite credit. Those intra-work protocols are below TaskManager's resource transitions, - so its allocation-level checker cannot prove them. Persistent KV256 has + so its allocation-level checker cannot prove them. Persistent streaming has enough task-level ordering for the checker and remains covered. """ return cfg.uses_staged_one_inst_tmem_p or ( @@ -1982,6 +1987,10 @@ def _run_decode_gen_active( g_sparse_row_route_offsets: cute.Pointer | None = None, g_sparse_row_route_counts: cute.Pointer | None = None, g_sparse_route_metadata: cute.Pointer | None = None, + tma_desc_k_summary: cutlass.GridConstant[cuda.TensorMap] | None = None, + tma_desc_v_summary: cutlass.GridConstant[cuda.TensorMap] | None = None, + tma_desc_k_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, + tma_desc_v_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, ) -> None: """Run the complete decode body for one runtime-valid Q tile. @@ -2021,28 +2030,43 @@ def _run_decode_gen_active( else Int32(cfg.static_seq_len_kv) ) use_clc_dynamic_scheduler = cfg.use_persistent_scheduler + tma_desc_k_summary_ptr = None + tma_desc_v_summary_ptr = None + tma_desc_k_summary_atom_ptr = None + tma_desc_v_summary_atom_ptr = None + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): + assert tma_desc_k_summary is not None + assert tma_desc_v_summary is not None + assert tma_desc_k_summary_atom is not None + assert tma_desc_v_summary_atom is not None + tma_desc_k_summary_ptr = tma_desc_k_summary.get_ptr() + tma_desc_v_summary_ptr = tma_desc_v_summary.get_ptr() + tma_desc_k_summary_atom_ptr = tma_desc_k_summary_atom.get_ptr() + tma_desc_v_summary_atom_ptr = tma_desc_v_summary_atom.get_ptr() # Prefetch TMA + uses_atom_desc = False + if cutlass.const_expr(cfg.use_block_sparse): + _, _, uses_atom_desc = _block_sparse_contiguous_kv_copy_geometry( + kv_block_size=cfg.kv_block_size, + kv_route_size=cfg.tile_size_kv, + ) init_warp = 1 if warp_idx == init_warp: prims.prefetch_tensormap(tma_desc_q.get_ptr()) prims.prefetch_tensormap(tma_desc_k.get_ptr()) prims.prefetch_tensormap(tma_desc_v.get_ptr()) - if cutlass.const_expr( - cfg.use_block_sparse - and _block_sparse_kv_atom_size(cfg.kv_block_size) == 64 - and ( - cfg.tile_size_kv == 256 - or not _prepared_kv_routes_are_block_aligned( - cfg.kv_block_size, - cfg.tile_size_kv, - ) - ) - ): - # KV256 always issues semantic KV64 atoms. KV128 needs this second - # descriptor only for non-aligned coarse routes. + if cutlass.const_expr(cfg.use_block_sparse and uses_atom_desc): + # KV256 and non-aligned coarse KV128 may select the exact atom maps. prims.prefetch_tensormap(tma_desc_k_atom.get_ptr()) prims.prefetch_tensormap(tma_desc_v_atom.get_ptr()) + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): + if cutlass.const_expr(cfg.tile_size_kv != 256): + prims.prefetch_tensormap(tma_desc_k_summary_ptr) + prims.prefetch_tensormap(tma_desc_v_summary_ptr) + if cutlass.const_expr(uses_atom_desc): + prims.prefetch_tensormap(tma_desc_k_summary_atom_ptr) + prims.prefetch_tensormap(tma_desc_v_summary_atom_ptr) init_warp += 1 clc_response_ptr = None @@ -2055,6 +2079,27 @@ def _run_decode_gen_active( if cutlass.const_expr(cfg.max_seq_len_q > 1): q_output_rows = g_h_r * Int32(cfg.max_seq_len_q) + # Static block-sparse tiles read their prepared row header here, before + # TMEM allocation and barrier setup, so that global-memory round trip is + # hidden instead of stalling every task at its first schedule step. + sparse_row_route_begin = None + sparse_route_count = None + if cutlass.const_expr( + cfg.use_block_sparse + and not cfg.use_persistent_scheduler + and g_sparse_row_route_offsets is not None + and g_sparse_row_route_counts is not None + ): + sparse_row_route_begin, sparse_route_count = _prefetch_prepared_sparse_row( + cfg, + g_sparse_row_route_offsets, + g_sparse_row_route_counts, + q_group_idx, + h_k_idx, + b_idx, + g_h_k, + ) + ( task_list, dep_graph, @@ -2083,6 +2128,10 @@ def _run_decode_gen_active( tma_desc_v=tma_desc_v.get_ptr(), tma_desc_k_atom=tma_desc_k_atom.get_ptr(), tma_desc_v_atom=tma_desc_v_atom.get_ptr(), + tma_desc_k_summary=tma_desc_k_summary_ptr, + tma_desc_v_summary=tma_desc_v_summary_ptr, + tma_desc_k_summary_atom=tma_desc_k_summary_atom_ptr, + tma_desc_v_summary_atom=tma_desc_v_summary_atom_ptr, page_idx_kv=g_page_idx_kv, h_k_idx=h_k_idx, b_idx=b_idx, @@ -2102,6 +2151,8 @@ def _run_decode_gen_active( sparse_row_route_offsets=g_sparse_row_route_offsets, sparse_row_route_counts=g_sparse_row_route_counts, sparse_route_metadata=g_sparse_route_metadata, + sparse_row_route_begin=sparse_row_route_begin, + sparse_route_count=sparse_route_count, ) smem_allocator.allocate() @@ -2264,6 +2315,10 @@ def _run_decode_gen_runtime_prefix( g_sparse_row_route_offsets: cute.Pointer | None, g_sparse_row_route_counts: cute.Pointer | None, g_sparse_route_metadata: cute.Pointer | None, + tma_desc_k_summary: cutlass.GridConstant[cuda.TensorMap] | None = None, + tma_desc_v_summary: cutlass.GridConstant[cuda.TensorMap] | None = None, + tma_desc_k_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, + tma_desc_v_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, ) -> None: """Run the general runtime split-prefix producer or retire its suffix.""" @@ -2328,6 +2383,10 @@ def _run_decode_gen_runtime_prefix( g_sparse_row_route_offsets, g_sparse_row_route_counts, g_sparse_route_metadata, + tma_desc_k_summary=tma_desc_k_summary, + tma_desc_v_summary=tma_desc_v_summary, + tma_desc_k_summary_atom=tma_desc_k_summary_atom, + tma_desc_v_summary_atom=tma_desc_v_summary_atom, ) else: _run_decode_gen_inactive_cluster_rank() @@ -2371,6 +2430,10 @@ def _run_decode_gen_runtime_prefix( g_sparse_row_route_offsets, g_sparse_row_route_counts, g_sparse_route_metadata, + tma_desc_k_summary=tma_desc_k_summary, + tma_desc_v_summary=tma_desc_v_summary, + tma_desc_k_summary_atom=tma_desc_k_summary_atom, + tma_desc_v_summary_atom=tma_desc_v_summary_atom, ) else: _signal_padded_pdl_producer(cfg) @@ -2409,6 +2472,10 @@ def decode_gen_kernel( g_sparse_row_route_counts: cute.Pointer | None = None, g_sparse_route_metadata: cute.Pointer | None = None, static_full_split_prefix: cutlass.Constexpr[bool] = False, + tma_desc_k_summary: cutlass.GridConstant[cuda.TensorMap] | None = None, + tma_desc_v_summary: cutlass.GridConstant[cuda.TensorMap] | None = None, + tma_desc_k_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, + tma_desc_v_summary_atom: cutlass.GridConstant[cuda.TensorMap] | None = None, ) -> None: """Dispatch one static Q/split tile and drain padded launch slots safely.""" q_group_cta_idx, h_k_idx, b_idx = cute.arch.block_idx() @@ -2470,6 +2537,10 @@ def decode_gen_kernel( g_sparse_row_route_offsets, g_sparse_row_route_counts, g_sparse_route_metadata, + tma_desc_k_summary=tma_desc_k_summary, + tma_desc_v_summary=tma_desc_v_summary, + tma_desc_k_summary_atom=tma_desc_k_summary_atom, + tma_desc_v_summary_atom=tma_desc_v_summary_atom, ) else: _run_decode_gen_runtime_prefix( @@ -2510,6 +2581,10 @@ def decode_gen_kernel( g_sparse_row_route_offsets, g_sparse_row_route_counts, g_sparse_route_metadata, + tma_desc_k_summary=tma_desc_k_summary, + tma_desc_v_summary=tma_desc_v_summary, + tma_desc_k_summary_atom=tma_desc_k_summary_atom, + tma_desc_v_summary_atom=tma_desc_v_summary_atom, ) else: # Packed-Q grids use a batch-wide maximum envelope. These Q CTAs own no @@ -2754,7 +2829,6 @@ def fmha_decode_launch( tma_desc_q, tma_desc_k, tma_desc_v, - # Dense/paged profiles never inspect the 64-token descriptor slots. tma_desc_k, tma_desc_v, o_iter, @@ -2799,6 +2873,8 @@ def fmha_block_sparse_launch( q_iter: cute.Pointer, k_iter: cute.Pointer, v_iter: cute.Pointer, + k_summary_iter: cute.Pointer, + v_summary_iter: cute.Pointer, o_iter: cute.Pointer, row_route_offsets_iter: cute.Pointer, row_route_counts_iter: cute.Pointer, @@ -2813,14 +2889,18 @@ def fmha_block_sparse_launch( k_page_stride: Int64 = 0, v_page_stride: Int64 = 0, ) -> None: - """Launch attention over contiguous or paged prepared KV routes. + """Launch attention over exact and typed exact/proxy prepared KV routes. A preceding prepare kernel has already resolved each BSR row into compact logical atom origins, storage locators, validity flags, and optional token - words. Both layouts execute the same ``decode_gen_kernel`` schedule. + words. Exact routes address K/V; proxy routes address summary K/V. Both + layouts execute the same ``decode_gen_kernel`` schedule and + physical copy policy. Exact builds constexpr-elide summary TensorMaps. """ if cutlass.const_expr(not cfg.use_block_sparse): - raise ValueError("fmha_block_sparse_launch requires cfg.use_block_sparse=True") + raise ValueError("fmha_block_sparse_launch requires block-sparse config") + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes and cfg.use_paged_kv): + raise ValueError("block-sparse proxy routes require contiguous K/V") log2_e = math.log2(math.e) b, h_q, h_k, s_k, d = problem_shape @@ -2855,7 +2935,18 @@ def fmha_block_sparse_launch( swizzle=tma_swizzle, ) - kv_atom_size = _block_sparse_kv_atom_size(cfg.kv_block_size) + ( + primary_kv_box_size, + kv_atom_size, + uses_atom_desc, + ) = _block_sparse_contiguous_kv_copy_geometry( + kv_block_size=cfg.kv_block_size, + kv_route_size=cfg.tile_size_kv, + ) + k_desc_summary_primary = None + v_desc_summary_primary = None + k_desc_summary_atom = None + v_desc_summary_atom = None if cutlass.const_expr(cfg.use_paged_kv): # Paged HND storage is addressed as (D, token-in-page, Hkv, page). # Prepared routes already contain each atom's physical page ID, so no @@ -2899,9 +2990,10 @@ def fmha_block_sparse_launch( k_desc_primary = k_desc_atom v_desc_primary = v_desc_atom else: - # Contiguous sparse coordinates retain the logical (D, S, H, B) - # order and the established primary/atom descriptor split. - primary_kv_box_size = 2 * kv_atom_size if kv_atom_size == 64 else kv_atom_size + # Exact and summary tensors form one logical segmented KV coordinate + # space. Each physical source owns the same primary/atom descriptor + # pair; the prepared route kind selects the pair, while the loader + # retains the existing KV128/fine/KV256 copy policy. kv_dims = (d, s_k, h_k, b) k_desc_primary = create_tensor_map_tiled( global_address=k_iter.toint(), @@ -2921,16 +3013,7 @@ def fmha_block_sparse_launch( ) k_desc_atom = k_desc_primary v_desc_atom = v_desc_primary - if cutlass.const_expr( - kv_atom_size == 64 - and ( - cfg.tile_size_kv == 256 - or not _prepared_kv_routes_are_block_aligned( - cfg.kv_block_size, - cfg.tile_size_kv, - ) - ) - ): + if cutlass.const_expr(uses_atom_desc): # KV256 always stages four semantic KV64 atoms. KV128 needs this # map only when a route may join unrelated BSR entries. k_desc_atom = create_tensor_map_tiled( @@ -2950,6 +3033,55 @@ def fmha_block_sparse_launch( swizzle=tma_swizzle, ) + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): + num_kv_blocks, _ = _block_sparse_proxy_summary_geometry( + seq_len_kv, + cfg.kv_block_size, + ) + _, summary_kv_strides = _block_sparse_bshd_tma_strides( + q_seq=q_seq, + h_q=h_q, + h_k=h_k, + s_k=num_kv_blocks, + d=d, + ) + summary_dims = (d, num_kv_blocks, h_k, b) + k_desc_summary_primary = create_tensor_map_tiled( + global_address=k_summary_iter.toint(), + dtype=cfg.kv_dtype, + global_dims=summary_dims, + global_strides=summary_kv_strides, + box_dims=(tma_box0, primary_kv_box_size, 1, 1), + swizzle=tma_swizzle, + ) + v_desc_summary_primary = create_tensor_map_tiled( + global_address=v_summary_iter.toint(), + dtype=cfg.kv_dtype, + global_dims=summary_dims, + global_strides=summary_kv_strides, + box_dims=(tma_box0, primary_kv_box_size, 1, 1), + swizzle=tma_swizzle, + ) + k_desc_summary_atom = k_desc_summary_primary + v_desc_summary_atom = v_desc_summary_primary + if cutlass.const_expr(uses_atom_desc): + k_desc_summary_atom = create_tensor_map_tiled( + global_address=k_summary_iter.toint(), + dtype=cfg.kv_dtype, + global_dims=summary_dims, + global_strides=summary_kv_strides, + box_dims=(tma_box0, kv_atom_size, 1, 1), + swizzle=tma_swizzle, + ) + v_desc_summary_atom = create_tensor_map_tiled( + global_address=v_summary_iter.toint(), + dtype=cfg.kv_dtype, + global_dims=summary_dims, + global_strides=summary_kv_strides, + box_dims=(tma_box0, kv_atom_size, 1, 1), + swizzle=tma_swizzle, + ) + q_groups = Int32( (cfg.max_seq_len_q + cfg.q_tokens_per_cta - 1) // cfg.q_tokens_per_cta ) @@ -3008,6 +3140,10 @@ def fmha_block_sparse_launch( row_route_counts_iter, route_metadata_iter, False, # static_full_split_prefix + tma_desc_k_summary=k_desc_summary_primary, + tma_desc_v_summary=v_desc_summary_primary, + tma_desc_k_summary_atom=k_desc_summary_atom, + tma_desc_v_summary_atom=v_desc_summary_atom, ).launch( grid=grid, block=[cfg.threads_per_cta, 1, 1], diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_common.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_common.py index e3190f09097c..77412cddca7b 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_common.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_common.py @@ -25,6 +25,8 @@ import cutlass import cutlass.cute as cute from cutlass import BFloat16, Float16, Float32, Int32, Int64, Uint32 +from cutlass._mlir.dialects import llvm +from cutlass.cutlass_dsl import dsl_user_op from cutlass.experimental import primitives as prims from cutlass.experimental.task_scheduling.resources import ( @@ -59,6 +61,22 @@ ) ResourceVars = dict[str, ResourceVarValue] + +@dsl_user_op +def _assume_nonnegative_i32(value: Int32, *, loc=None, ip=None) -> Int32: + """Express a caller-guaranteed nonnegative Int32 contract to codegen.""" + + condition = cutlass.Boolean(value >= Int32(0)) + llvm.intr_assume( + condition.ir_value(loc=loc, ip=ip), + [], + [], + loc=loc, + ip=ip, + ) + return value + + # Offsets into DecodeGenTask.make_task_cache(). Keeping these symbolic makes # resource code explicit about which task-local lane or address value it needs. _TASK_CACHE_TMEM_BASE_OFFSET = 0 @@ -109,6 +127,33 @@ def _warp_broadcast_i32(value: Int32, source_lane: Constexpr[int]) -> Int32: ) +@cute.jit +def _swaps_routed_coordinate( + cfg: Constexpr[FmhaDecodeConfig], + lane_k_offset: Int32, + origin0: Int32, + origin1: Int32, + origin2: Int32, + origin3: Int32, + *, + token_group_idx: Constexpr[int], +) -> tuple[Int32, Int32]: + """Map one SWAP register group to its staged atom and logical coordinate.""" + + atom_size = min(cfg.kv_block_size, 32) + groups_per_atom = atom_size // 8 + origin_idx = token_group_idx // groups_per_atom + atom_origin = origin0 + if cutlass.const_expr(origin_idx == 1): + atom_origin = origin1 + elif cutlass.const_expr(origin_idx == 2): + atom_origin = origin2 + elif cutlass.const_expr(origin_idx == 3): + atom_origin = origin3 + token_offset = (token_group_idx % groups_per_atom) * 8 + return atom_origin, atom_origin + Int32(token_offset) + lane_k_offset + + def _mma_kind_for_qkv(cfg: FmhaDecodeConfig) -> prims.Tcgen05MMAKind: """Select the tcgen05 MMA opcode family used for Q/K/V operands.""" return prims.Tcgen05MMAKind.F8F6F4 if cfg.use_fp8_qkv else prims.Tcgen05MMAKind.F16 diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_softmax.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_softmax.py index 4585213cd80b..2dbe3c7e7ff8 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_softmax.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/helpers_softmax.py @@ -71,6 +71,70 @@ def _pack_float4_to_fp8_e4m3_inline( ) +@cute.jit +def _combine_int_frac_ex2( + x_rounded: Float32, frac_ex2: Float32, *, loc=None, ip=None +) -> Float32: + """Scale ``frac_ex2`` by ``2**floor(x)`` through the FP32 exponent bits. + + ``x_rounded`` still carries the magic rounding constant, so its low + mantissa bits hold ``floor(x)``; shifting them into the exponent field and + adding the bits of the polynomial result multiplies by the integer power. + """ + return cute.arch.inline_ptx( + "{\n" + " .reg .b32 xi;\n" + " .reg .b32 fi;\n" + " .reg .b32 xe;\n" + " .reg .b32 oi;\n" + " mov.b32 xi, {$r0};\n" + " mov.b32 fi, {$r1};\n" + " shl.b32 xe, xi, 23;\n" + " add.s32 oi, xe, fi;\n" + " mov.b32 {$w0}, oi;\n" + "}", + write_only_types=[Float32], + read_only_args=[x_rounded, frac_ex2], + loc=loc, + ip=ip, + ) + + +@cute.jit +def _ex2_emulation_packed_f32x2(x: Float32, y: Float32) -> tuple[Float32, Float32]: + """Evaluate ``2**x`` and ``2**y`` on the FMA pipe instead of MUFU. + + Inputs are non-positive scaled scores minus the row maximum. The integer + part is split off with a magic-constant rounding add, the fraction in + [0, 1) goes through a degree-3 minimax polynomial, and the two parts are + recombined through the exponent bits. The relative error stays below the + BF16 rounding of the P operand, matching the emulation used by the dense + Blackwell FMHA kernels. + """ + fp32_round_int = float(2**23 + 2**22) + xy_clamped = (cute.arch.fmax(x, -127.0), cute.arch.fmax(y, -127.0)) + xy_rounded = cute.arch.add_packed_f32x2( + xy_clamped, (fp32_round_int, fp32_round_int), rnd="rm" + ) + xy_rounded_back = cute.arch.sub_packed_f32x2( + xy_rounded, (fp32_round_int, fp32_round_int) + ) + xy_frac = cute.arch.sub_packed_f32x2(xy_clamped, xy_rounded_back) + coeff = ( + 1.0, + 0.695146143436431884765625, + 0.227564394474029541015625, + 0.077119089663028717041015625, + ) + out = (coeff[3], coeff[3]) + for degree in cutlass.range_constexpr(2, -1, -1): + out = cute.arch.fma_packed_f32x2(out, xy_frac, (coeff[degree], coeff[degree])) + return ( + _combine_int_frac_ex2(xy_rounded[0], out[0]), + _combine_int_frac_ex2(xy_rounded[1], out[1]), + ) + + @cute.jit def _compute_fp8_p_regs_and_local_sums( scale_softmax_log2: Float32, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_block_sparse_metadata.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_block_sparse_metadata.py index 3254de0e79e7..f79f3a1e0566 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_block_sparse_metadata.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_block_sparse_metadata.py @@ -43,6 +43,7 @@ from ...._block_sparse.prepared import ( _PREPARED_ROUTE_IS_FULL_FLAG, + _PREPARED_ROUTE_IS_PROXY_FLAG, _BlockSparseRouteLayout, ) from ...placeholder_helpers import _placeholder_smem_array @@ -56,7 +57,6 @@ DecodeGenResourceBase, ResourceVars, _decode_gen_task_cache, - _keeps_col_base, _sparse_task_cache_route_begin, _sparse_task_cache_route_count, _warp_broadcast_i32, @@ -65,13 +65,20 @@ # Keeps staging uses the low four bits for structural KV64 validity. Bit 4 # carries the conservative prepared summary that token masking can be skipped; -# structural, tail, and causal masking remain independent. +# structural, tail, and causal masking remain independent. The streamed Keeps +# max pass derives its keep words from the token words directly, so the bit is +# currently staged for the consumer but not read. _SOFTMAX_TOKEN_MASK_IS_FULL_FLAG = 1 << 4 - -# B8 SWAP origins are eight-token aligned, so bit 0 is free while the route is -# in Softmax's private staging payload. Reusing it avoids adding a word to every -# pipeline stage merely to forward prepare's route-full summary. -_SWAPS_PACKED_ROUTE_FULL_CLEAR_MASK = ~_PREPARED_ROUTE_IS_FULL_FLAG +# Keeps reserves bit 5 for the prepared route kind. The low four structural +# validity bits and bit 4 keep their existing meaning. +_SOFTMAX_ROUTE_IS_PROXY_FLAG = 1 << 5 + +# SWAP origins are at least eight-token aligned, so their low two bits are free +# while the route is in Softmax's private staging payload. Reusing them avoids +# adding a word to every pipeline stage for prepared FULL/PROXY route flags. +_SWAPS_PACKED_ROUTE_FLAGS_CLEAR_MASK = ~( + _PREPARED_ROUTE_IS_FULL_FLAG | _PREPARED_ROUTE_IS_PROXY_FLAG +) @cute.jit @@ -108,7 +115,7 @@ def _swaps_forwards_packed_route_full(cfg: FmhaDecodeConfig) -> bool: return ( cfg.tile_size_q == 8 and cfg.kv_block_size == 8 - and not cfg.use_kv_valid_bits + and not cfg.uses_prepared_score_keep_words and not cfg.uses_uniform_causal_mask and not cfg.uses_per_row_causal_mask ) @@ -116,12 +123,16 @@ def _swaps_forwards_packed_route_full(cfg: FmhaDecodeConfig) -> bool: def _kv_retained_route_words( route_layout: _BlockSparseRouteLayout, + *, + retain_proxy_kind: bool = False, ) -> int: """Return the aligned SMEM words retained from K issue through V. - Contiguous routes retain their existing load-origin payload. Paged routes - retain parallel logical-origin and physical-page-ID arrays so every atom - has an independent storage locator; invalid entries use ``(-1, -1)``. + Contiguous routes retain their load-origin payload. Paged routes retain + parallel logical-origin and physical-page-ID arrays so every atom has an + independent storage locator. A two-origin contiguous route additionally + keeps its atom-valid mask. Proxy-capable exact/proxy routes reserve the + final aligned word for an explicit source kind. """ payload_words = route_layout.logical_origins_per_route @@ -129,6 +140,8 @@ def _kv_retained_route_words( payload_words *= 2 elif route_layout.logical_origins_per_route == 2: payload_words += 1 + if retain_proxy_kind: + payload_words += 1 return ((payload_words + 3) // 4) * 4 @@ -139,9 +152,9 @@ class _BlockSparseSoftmaxStagingLayout: Keeps retains all route origins, a flags word, alignment padding, and the optional K32 token words. KV256 consumers then select the four words owned by their spatial half. SWAP stores execution-ordered origins followed by - optional logical K32 token words, one for each consumer warp. Its - noncausal Q8/B8 profile without token bits packs route-full into the - otherwise-zero low bit of each warp's first aligned origin. + optional logical K32 token words, one for each consumer warp. Selected + prepared route flags travel in the otherwise-zero low bits of each warp's + first aligned origin. """ # Logical-origin scalars staged for one complete KV route. @@ -206,14 +219,23 @@ class SmemBlockSparseKvMetadataResource(DecodeGenResourceBase): """Pipeline-free route metadata retained from one K issue through V. ``route_metadata`` points at the first prepared GMEM record. Resolution - returns logical origins to masking consumers. The private SMEM copy keeps - contiguous load origins or paged ``(logical origin, physical page ID)`` - pairs through the matching V issue. Invalid atoms are retained as safe - storage-specific OOB coordinates. + returns logical/source-domain origins to masking consumers. The private + SMEM copy keeps contiguous load origins or paged ``(logical origin, + physical page ID)`` pairs through the matching V issue. Invalid atoms are + retained as safe storage-specific OOB coordinates. The proxy-capable + contiguous specialization interprets origins in the selected source domain + (summary tokens for proxy routes, K/V tokens for exact routes) and + retains the prepared route kind in a separate aligned word. Exact-only + specializations keep their original allocation. """ _task_local_specs: ClassVar[tuple[tuple, ...]] = ( - ("resolved_origin0_slot", Int32, Int32(0), "First logical origin."), + ( + "resolved_record_word_slot", + Int32, + Int32(0), + "Lane-owned record word; locator lanes carry route origins.", + ), ("resolved_origin1_slot", Int32, Int32(0), "Second logical origin."), ( "resolved_atom_validity_slot", @@ -227,6 +249,18 @@ class SmemBlockSparseKvMetadataResource(DecodeGenResourceBase): Int32(-1), "Metadata-relative record offset, or -1 for a dummy route.", ), + ( + "prefetched_record_word_slot", + Int32, + Int32(0), + "Lane-owned record word loaded one resolution ahead.", + ), + ( + "prefetched_record_offset_slot", + Int32, + Int32(-1), + "Record offset of the prefetched route, or -1 for a dummy route.", + ), ) cfg: Constexpr[FmhaDecodeConfig] = None inst_id: Constexpr[int] = 0 @@ -236,7 +270,7 @@ class SmemBlockSparseKvMetadataResource(DecodeGenResourceBase): _retained_route_words: Constexpr[int] = 0 _alloc: Constexpr[SmemAllocation | None] = None _smem_words: cutlass.Array = None - resolved_origin0_slot: Constexpr[TaskLocalVariable] = ( + resolved_record_word_slot: Constexpr[TaskLocalVariable] = ( TaskLocalVariable.uninitialized() ) resolved_origin1_slot: Constexpr[TaskLocalVariable] = ( @@ -248,13 +282,25 @@ class SmemBlockSparseKvMetadataResource(DecodeGenResourceBase): route_record_word_offset_slot: Constexpr[TaskLocalVariable] = ( TaskLocalVariable.uninitialized() ) + prefetched_record_word_slot: Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) + prefetched_record_offset_slot: Constexpr[TaskLocalVariable] = ( + TaskLocalVariable.uninitialized() + ) def __post_init__(self) -> None: """Derive the retained K/V payload from the prepared route layout.""" assert self.route_layout is not None assert self.route_layout.is_paged == self.cfg.use_paged_kv - self._retained_route_words = _kv_retained_route_words(self.route_layout) + if self.cfg.use_block_sparse_proxy_routes: + assert not self.route_layout.is_paged + assert self.route_layout.uses_one_warp_transport + self._retained_route_words = _kv_retained_route_words( + self.route_layout, + retain_proxy_kind=self.cfg.use_block_sparse_proxy_routes, + ) super().__post_init__() def _init_placeholder_state(self) -> None: @@ -332,9 +378,81 @@ def _prepared_route_physical_page_id_if_valid( ) return physical_page_id + @cute.jit + def _route_record_word_offset( + self, stage_info: StageInfo, route_idx: Int32 + ) -> Int32: + """Return the record offset of one route index, or -1 past the row.""" + + task_cache = _decode_gen_task_cache(stage_info) + row_route_begin = _sparse_task_cache_route_begin(task_cache) + route_count = _sparse_task_cache_route_count(task_cache) + route_record_word_offset = Int32(-1) + if route_idx < route_count: + route_record_word_offset = (row_route_begin + route_idx) * Int32( + self.route_layout.route_metadata_stride_words + ) + return cute.arch.make_warp_uniform(route_record_word_offset) + + @consumer_work( + returns=( + prefetched_record_word_slot, + prefetched_record_offset_slot, + ) + ) + @cute.jit + def prefetch_route( + self, stage_info: StageInfo, *, target: Constexpr[str] + ) -> tuple[Int32, Int32]: + """Issue the record load for a route that ``resolve_route`` uses later. + + ``target`` selects the route relative to the calling section: + ``"head"`` is this instance's HEAD route, ``"first_loop"`` the route of + LOOP iteration 0 (called from HEAD), ``"current_loop"`` the route of + the calling LOOP iteration (no pipelining), and ``"next_loop"`` the + route of the following LOOP iteration. Only the lane-distributed load is issued + here; the warp broadcasts happen in ``resolve_route`` so the global + memory latency overlaps the TMA issue of the current route instead of + stalling the load warp. Layouts without one-warp transport keep their + loads in ``resolve_route`` and get placeholder values here. + """ + + assert self.route_metadata is not None + num_insts = Int32(self.cfg.num_insts_kv) + if cutlass.const_expr(target == "head"): + route_idx = Int32(self.inst_id) + elif cutlass.const_expr(target == "first_loop"): + route_idx = num_insts + Int32(self.inst_id) + elif cutlass.const_expr(target == "current_loop"): + route_idx = (stage_info.loop_offset + Int32(1)) * num_insts + Int32( + self.inst_id + ) + else: + route_idx = (stage_info.loop_offset + Int32(2)) * num_insts + Int32( + self.inst_id + ) + route_record_word_offset = self._route_record_word_offset(stage_info, route_idx) + record_word = Int32(0) + if cutlass.const_expr(self.route_layout.uses_one_warp_transport): + assert self.route_layout.token_words_word_offset is not None + meaningful_words = ( + self.route_layout.token_words_word_offset + + self.route_layout.token_words_per_route + ) + lane_idx = cute.arch.thread_idx()[0] & Int32(0x1F) + if lane_idx < Int32(self.route_layout.logical_origins_per_route): + record_word = Int32(-1) + if route_record_word_offset >= Int32(0) and lane_idx < Int32( + meaningful_words + ): + record_word = Int32( + self.route_metadata[route_record_word_offset + lane_idx] + ) + return record_word, route_record_word_offset + @consumer_work( returns=( - resolved_origin0_slot, + resolved_record_word_slot, resolved_origin1_slot, resolved_atom_validity_slot, route_record_word_offset_slot, @@ -342,14 +460,23 @@ def _prepared_route_physical_page_id_if_valid( ) @cute.jit def resolve_route( - self, stage_info: StageInfo, *, section: Constexpr[FmhaStage] + self, + stage_info: StageInfo, + *, + section: Constexpr[FmhaStage], + prefetched_record_word_slot: Int32, + prefetched_record_offset_slot: Int32, ) -> tuple[Int32, Int32, Int32, Int32]: - """Load this resource instance's real or dummy prepared KV route.""" + """Resolve this instance's real or dummy prepared KV route. + + One-warp-transport layouts consume the words that ``prefetch_route`` + loaded earlier; other layouts load their record here. The routed + inputs carry the task-local slot names so that every ``prefetch_route`` + call, including the one at the end of the previous LOOP iteration, + updates the value read here. + """ assert self.route_metadata is not None - task_cache = _decode_gen_task_cache(stage_info) - row_route_begin = _sparse_task_cache_route_begin(task_cache) - route_count = _sparse_task_cache_route_count(task_cache) # HEAD publishes one route per instruction. LOOP starts after those # two publications, hence the one-based loop offset below. Keeping the # constexpr branch local lets the task scheduler specialize each work @@ -361,12 +488,12 @@ def resolve_route( self.cfg.num_insts_kv ) + Int32(self.inst_id) lane_idx = cute.arch.thread_idx()[0] & Int32(0x1F) - route_record_word_offset = Int32(-1) - if route_idx < route_count: - route_record_word_offset = (row_route_begin + route_idx) * Int32( - self.route_layout.route_metadata_stride_words + if cutlass.const_expr(self.route_layout.uses_one_warp_transport): + route_record_word_offset = prefetched_record_offset_slot + else: + route_record_word_offset = self._route_record_word_offset( + stage_info, route_idx ) - route_record_word_offset = cute.arch.make_warp_uniform(route_record_word_offset) num_logical_origins = self.route_layout.logical_origins_per_route uses_two_fragment_route = num_logical_origins == 2 @@ -378,13 +505,37 @@ def resolve_route( atom_valid_mask = Int32(0) route_record_is_valid = route_record_word_offset >= Int32(0) + if cutlass.const_expr(self.route_layout.uses_one_warp_transport): + resolved_record_word = prefetched_record_word_slot + atom_valid_mask = _warp_broadcast_i32( + resolved_record_word, + self.route_layout.atom_valid_mask_word_offset, + ) + if cutlass.const_expr(uses_two_fragment_route): + return ( + resolved_record_word, + _warp_broadcast_i32(resolved_record_word, 1), + atom_valid_mask, + route_record_word_offset, + ) + atom_is_valid = cutlass.Boolean( + lane_idx < Int32(num_logical_origins) + and (atom_valid_mask & (Int32(1) << lane_idx)) != Int32(0) + ) + return ( + resolved_record_word, + Int32(0), + Int32(atom_is_valid), + route_record_word_offset, + ) + if route_record_is_valid: if lane_idx < Int32(num_logical_origins): logical_origin = self._prepared_route_logical_origin( route_record_word_offset, lane_idx, ) - if cutlass.const_expr(self.cfg.use_kv_valid_bits): + if cutlass.const_expr(self.cfg.uses_prepared_score_keep_words): if lane_idx == Int32(valid_mask_lane): atom_valid_mask = Int32( self.route_metadata[ @@ -396,7 +547,7 @@ def resolve_route( if cutlass.const_expr(uses_two_fragment_route): origin0 = _warp_broadcast_i32(logical_origin, 0) origin1 = _warp_broadcast_i32(logical_origin, 1) - if cutlass.const_expr(self.cfg.use_kv_valid_bits): + if cutlass.const_expr(self.cfg.uses_prepared_score_keep_words): # The validity word shares the prepared record's cache line # with fields consumed shortly afterward by Softmax. atom_valid_mask = _warp_broadcast_i32(atom_valid_mask, valid_mask_lane) @@ -412,7 +563,7 @@ def resolve_route( # Wider routes stay lane-distributed: each active lane carries only # its origin and validity through the existing three-scalar K/V ABI. valid = cutlass.Boolean(False) - if cutlass.const_expr(self.cfg.use_kv_valid_bits): + if cutlass.const_expr(self.cfg.uses_prepared_score_keep_words): atom_valid_mask = _warp_broadcast_i32(atom_valid_mask, valid_mask_lane) if lane_idx < Int32(num_logical_origins): valid = (atom_valid_mask & (Int32(1) << lane_idx)) != Int32(0) @@ -427,7 +578,7 @@ def store_route( self, stage_info: StageInfo, *, - resolved_origin0: Int32, + resolved_record_word: Int32, resolved_origin1: Int32, resolved_atom_validity: Int32, route_record_word_offset: Int32, @@ -439,11 +590,11 @@ def store_route( num_origins = self.route_layout.logical_origins_per_route if cutlass.const_expr(self.route_layout.is_paged): if lane_idx < Int32(num_origins): - logical_origin = Int32(resolved_origin0) + logical_origin = Int32(resolved_record_word) atom_is_valid = resolved_atom_validity != Int32(0) if cutlass.const_expr(num_origins == 2): if lane_idx == Int32(0): - logical_origin = resolved_origin0 + logical_origin = resolved_record_word else: logical_origin = resolved_origin1 atom_is_valid = ( @@ -463,14 +614,14 @@ def store_route( ] = physical_page_id elif cutlass.const_expr(num_origins == 2): if lane_idx == Int32(0): - self._smem_words[Int32(0)] = resolved_origin0 + self._smem_words[Int32(0)] = resolved_record_word self._smem_words[Int32(1)] = resolved_origin1 self._smem_words[ Int32(self.route_layout.atom_valid_mask_word_offset) ] = resolved_atom_validity else: - if lane_idx < Int32(self.route_layout.logical_origins_per_route): - load_origin = Int32(resolved_origin0) + if lane_idx < Int32(num_origins): + load_origin = Int32(resolved_record_word) if resolved_atom_validity == Int32(0): # Fine-route K and V both consume this retained value. # Materialize their TensorMap OOB coordinate once here @@ -478,6 +629,12 @@ def store_route( # every atom copy in both producer passes. load_origin = Int32(self.tma_oob_origin) self._smem_words[lane_idx] = load_origin + if cutlass.const_expr(self.cfg.use_block_sparse_proxy_routes): + if lane_idx == Int32(self.route_layout.route_flags_word_offset): + prepared_route_flags = Int32(resolved_record_word) + self._smem_words[Int32(self._retained_route_words - 1)] = Int32( + prepared_route_flags + ) & Int32(_PREPARED_ROUTE_IS_PROXY_FLAG) # K consumes this slot immediately, while V consumes it at the start # of the next cadence. Both execute in this warp, so a warp fence is # sufficient; no cross-warp mbarrier belongs here. @@ -523,6 +680,20 @@ def route_atom_valid_mask(self) -> Int32: self._smem_words[Int32(self.route_layout.atom_valid_mask_word_offset)] ) + @cute.jit + def route_is_proxy(self) -> cutlass.Boolean: + """Return the retained prepared route kind for the current K/V pair.""" + + if cutlass.const_expr(not self.cfg.use_block_sparse_proxy_routes): + return cutlass.Boolean(False) + return cutlass.Boolean( + ( + Int32(self._smem_words[Int32(self._retained_route_words - 1)]) + & Int32(_PREPARED_ROUTE_IS_PROXY_FLAG) + ) + != Int32(0) + ) + @dataclass(kw_only=True) class SmemBlockSparseSoftmaxMetadataResource(DecodeGenResourceBase): @@ -533,8 +704,8 @@ class SmemBlockSparseSoftmaxMetadataResource(DecodeGenResourceBase): producer passes the resolved payload explicitly instead of recomputing it. For Keeps, every route token word moves through SMEM without a data-dependent branch; each consumer receives at most four words through - the stable task-local ABI. A runtime route-full bit can skip per-score token - predicates while leaving structural masking independent. + the stable task-local ABI. Runtime route flags carry the conservative FULL + summary and, for proxy-capable builds, the route source kind. """ _task_local_specs: ClassVar[tuple[tuple, ...]] = ( @@ -562,7 +733,7 @@ class SmemBlockSparseSoftmaxMetadataResource(DecodeGenResourceBase): "softmax_token_word2_slot", Uint32, Uint32(0xFFFFFFFF), - "Loaded third Keeps token word or SWAP's B8 route-full summary.", + "Loaded third Keeps token word or SWAP's packed route flags.", ), ( "softmax_token_word3_slot", @@ -605,6 +776,8 @@ def __post_init__(self) -> None: assert self.route_layout is not None assert self.route_layout.is_paged == self.cfg.use_paged_kv + if self.cfg.use_block_sparse_proxy_routes: + assert self.route_layout.uses_one_warp_transport self.staging_layout = _BlockSparseSoftmaxStagingLayout.create( use_keeps_mma_ab=self.cfg.use_keeps_mma_ab, route_layout=self.route_layout, @@ -680,57 +853,82 @@ def _consumer_stage_base(self) -> Int32: def _store_route_swaps( self, stage_info: StageInfo, - resolved_origin0: Int32, + resolved_record_word: Int32, resolved_origin1: Int32, resolved_atom_validity: Int32, route_record_word_offset: Int32, ) -> None: """Stage SWAP origins and optional logical-K32 token metadata. - The noncausal Q8/B8 profile without token bits also packs prepare's - route-full summary into bit 0 of each warp's first aligned origin. + Selected prepared route flags use the free low bits of each warp's + first aligned origin. """ lane_idx = cute.arch.thread_idx()[0] & Int32(0x1F) stage_base = self._producer_stage_base(stage_info) task_cache = _decode_gen_task_cache(stage_info) seq_len_kv = Int32(task_cache[_TASK_CACHE_SEQ_LEN_KV]) - route_record_is_valid = route_record_word_offset >= Int32(0) - - packed_route_full = Int32(0) - if cutlass.const_expr(_swaps_forwards_packed_route_full(self.cfg)): - if lane_idx == Int32(0) and route_record_is_valid: - packed_route_full = Int32( - self.route_metadata[ - route_record_word_offset - + Int32(self.route_layout.route_flags_word_offset) - ] - ) & Int32(_PREPARED_ROUTE_IS_FULL_FLAG) - packed_route_full = _warp_broadcast_i32(packed_route_full, 0) + uses_one_warp_transport = self.route_layout.uses_one_warp_transport + if cutlass.const_expr( + not uses_one_warp_transport + and ( + _swaps_forwards_packed_route_full(self.cfg) + or self.cfg.uses_prepared_score_keep_words + ) + ): + route_record_is_valid = route_record_word_offset >= Int32(0) + + packed_route_flags = Int32(0) + if cutlass.const_expr( + _swaps_forwards_packed_route_full(self.cfg) + or self.cfg.use_block_sparse_proxy_routes + ): + if cutlass.const_expr(uses_one_warp_transport): + packed_route_flags = _warp_broadcast_i32( + resolved_record_word, + self.route_layout.route_flags_word_offset, + ) + else: + if lane_idx == Int32(0) and route_record_is_valid: + packed_route_flags = Int32( + self.route_metadata[ + route_record_word_offset + + Int32(self.route_layout.route_flags_word_offset) + ] + ) & Int32(_PREPARED_ROUTE_IS_FULL_FLAG) + packed_route_flags = _warp_broadcast_i32(packed_route_flags, 0) softmax_origin = Int32(-1) if cutlass.const_expr(self.cfg.kv_block_size < 64): if lane_idx < Int32(self.staging_layout.num_origin_words): - softmax_origin = Int32(resolved_origin0) + softmax_origin = Int32(resolved_record_word) if resolved_atom_validity == Int32(0): softmax_origin = Int32(-1) - if cutlass.const_expr(_swaps_forwards_packed_route_full(self.cfg)): - # Replicate route-full in each K32 slice's first origin; - # B8 alignment leaves bit 0 free for the summary. + if cutlass.const_expr( + _swaps_forwards_packed_route_full(self.cfg) + or self.cfg.use_block_sparse_proxy_routes + ): + # Every SWAP atom is at least B8 aligned. Replicate the + # route flags in each K32 slice's first origin so the + # established seven-slot Softmax ABI also carries source + # kind without growing the staged payload. if lane_idx % Int32(self.staging_layout.origins_per_warp) == Int32( 0 ): softmax_origin = ( - softmax_origin & Int32(_SWAPS_PACKED_ROUTE_FULL_CLEAR_MASK) - ) | packed_route_full + softmax_origin & Int32(_SWAPS_PACKED_ROUTE_FLAGS_CLEAR_MASK) + ) | packed_route_flags self._smem_words[stage_base + lane_idx] = softmax_origin else: # SWAP with a coarse KV atom expands the two resolved KV64 # fragments into the four logical K32 origins consumed by its # four softmax warps. + coarse_origin0 = Int32(resolved_record_word) + if cutlass.const_expr(uses_one_warp_transport): + coarse_origin0 = _warp_broadcast_i32(resolved_record_word, 0) if lane_idx < Int32(4): fragment_idx = lane_idx >> Int32(1) - softmax_origin = Int32(resolved_origin0) + softmax_origin = coarse_origin0 valid = (resolved_atom_validity & Int32(1)) != Int32(0) if fragment_idx == Int32(1): softmax_origin = Int32(resolved_origin1) @@ -738,13 +936,30 @@ def _store_route_swaps( softmax_origin = softmax_origin + (lane_idx & Int32(1)) * Int32(32) if not valid or softmax_origin >= seq_len_kv: softmax_origin = Int32(-1) + if cutlass.const_expr(self.cfg.use_block_sparse_proxy_routes): + # Coarse SWAP expands KV64 atoms to K32-aligned origins; + # their low bits carry the same typed-route flags as the + # fine-route representation above. + softmax_origin = ( + softmax_origin & Int32(_SWAPS_PACKED_ROUTE_FLAGS_CLEAR_MASK) + ) | packed_route_flags self._smem_words[stage_base + lane_idx] = softmax_origin - if cutlass.const_expr(self.cfg.use_kv_valid_bits): + if cutlass.const_expr(self.cfg.uses_prepared_score_keep_words): assert self.route_metadata is not None assert self.route_layout.token_words_word_offset is not None assert self.staging_layout.token_words_word_offset is not None - if lane_idx < Int32(self.route_layout.token_words_per_route): + if cutlass.const_expr(uses_one_warp_transport): + token_begin = Int32(self.route_layout.token_words_word_offset) + token_end = token_begin + Int32(self.route_layout.token_words_per_route) + if lane_idx >= token_begin and lane_idx < token_end: + self._smem_words[ + stage_base + + Int32(self.staging_layout.token_words_word_offset) + + lane_idx + - token_begin + ] = Int32(resolved_record_word) + elif lane_idx < Int32(self.route_layout.token_words_per_route): logical_word = Uint32(0) if route_record_is_valid: logical_word = Uint32( @@ -765,7 +980,7 @@ def _store_route_swaps( def _store_route_keeps( self, stage_info: StageInfo, - resolved_origin0: Int32, + resolved_record_word: Int32, resolved_origin1: Int32, resolved_atom_validity: Int32, route_record_word_offset: Int32, @@ -779,8 +994,12 @@ def _store_route_keeps( assert self.staging_layout.route_flags_word_offset is not None lane_idx = cute.arch.thread_idx()[0] & Int32(0x1F) - route_record_is_valid = route_record_word_offset >= Int32(0) num_origins = self.route_layout.logical_origins_per_route + uses_one_warp_transport = self.route_layout.uses_one_warp_transport + if cutlass.const_expr( + not uses_one_warp_transport and self.cfg.uses_prepared_score_keep_words + ): + route_record_is_valid = route_record_word_offset >= Int32(0) route_flags = Int32(resolved_atom_validity) if cutlass.const_expr(num_origins > 2): route_flags = Int32( @@ -788,47 +1007,73 @@ def _store_route_keeps( lane_idx < Int32(num_origins) and resolved_atom_validity != Int32(0) ) ) + token_word = Uint32(0) route_token_mask_is_full = cutlass.Boolean(False) - if cutlass.const_expr(self.cfg.use_kv_valid_bits): - assert self.route_metadata is not None + if cutlass.const_expr(self.cfg.uses_prepared_score_keep_words): assert self.route_layout.token_words_word_offset is not None - gmem_route_flags = Int32(0) - if lane_idx == Int32(0) and route_record_is_valid: - gmem_route_flags = Int32( - self.route_metadata[ - route_record_word_offset - + Int32(self.route_layout.route_flags_word_offset) - ] - ) - gmem_route_flags = _warp_broadcast_i32(gmem_route_flags, 0) - # Prepared bit 0 summarizes the whole route. Staged low bits are - # already fragment validity, so remap the summary above them. - route_token_mask_is_full = cutlass.Boolean( - (gmem_route_flags & Int32(_PREPARED_ROUTE_IS_FULL_FLAG)) != Int32(0) - ) - if ( - lane_idx < Int32(self.route_layout.token_words_per_route) - and route_record_is_valid - ): - token_word = Uint32( - self.route_metadata[ - route_record_word_offset - + Int32(self.route_layout.token_words_word_offset) - + lane_idx - ] + assert self.route_metadata is not None + if cutlass.const_expr(not uses_one_warp_transport): + gmem_route_flags = Int32(0) + if lane_idx == Int32(0) and route_record_is_valid: + gmem_route_flags = Int32( + self.route_metadata[ + route_record_word_offset + + Int32(self.route_layout.route_flags_word_offset) + ] + ) + gmem_route_flags = _warp_broadcast_i32(gmem_route_flags, 0) + # Prepared bit 0 summarizes the whole route. Staged low bits + # already hold fragment validity, so remap it above them. + route_token_mask_is_full = cutlass.Boolean( + (gmem_route_flags & Int32(_PREPARED_ROUTE_IS_FULL_FLAG)) != Int32(0) ) + if ( + lane_idx < Int32(self.route_layout.token_words_per_route) + and route_record_is_valid + ): + token_word = Uint32( + self.route_metadata[ + route_record_word_offset + + Int32(self.route_layout.token_words_word_offset) + + lane_idx + ] + ) stage_base = self._producer_stage_base(stage_info) if cutlass.const_expr(num_origins == 2): if lane_idx == Int32(0): - self._smem_words[stage_base] = Int32(resolved_origin0) + self._smem_words[stage_base] = Int32(resolved_record_word) self._smem_words[stage_base + Int32(1)] = Int32(resolved_origin1) - else: - if lane_idx < Int32(num_origins): - self._smem_words[stage_base + lane_idx] = Int32(resolved_origin0) - if lane_idx == Int32(0): - if cutlass.const_expr(self.cfg.use_kv_valid_bits): + elif lane_idx < Int32(num_origins): + self._smem_words[stage_base + lane_idx] = Int32(resolved_record_word) + + if cutlass.const_expr(uses_one_warp_transport): + if lane_idx == Int32(self.route_layout.route_flags_word_offset): + prepared_route_flags = Int32(resolved_record_word) + route_flags = route_flags | ( + Int32( + (prepared_route_flags & Int32(_PREPARED_ROUTE_IS_FULL_FLAG)) + != Int32(0) + ) + * Int32(_SOFTMAX_TOKEN_MASK_IS_FULL_FLAG) + ) + if cutlass.const_expr(self.cfg.use_block_sparse_proxy_routes): + route_flags = route_flags | ( + Int32( + ( + prepared_route_flags + & Int32(_PREPARED_ROUTE_IS_PROXY_FLAG) + ) + != Int32(0) + ) + * Int32(_SOFTMAX_ROUTE_IS_PROXY_FLAG) + ) + self._smem_words[ + stage_base + Int32(self.staging_layout.route_flags_word_offset) + ] = route_flags + elif lane_idx == Int32(0): + if cutlass.const_expr(self.cfg.uses_prepared_score_keep_words): route_flags = route_flags | ( Int32(route_token_mask_is_full) * Int32(_SOFTMAX_TOKEN_MASK_IS_FULL_FLAG) @@ -836,9 +1081,20 @@ def _store_route_keeps( self._smem_words[ stage_base + Int32(self.staging_layout.route_flags_word_offset) ] = route_flags - if cutlass.const_expr(self.cfg.use_kv_valid_bits): + + if cutlass.const_expr(self.cfg.uses_prepared_score_keep_words): assert self.staging_layout.token_words_word_offset is not None - if lane_idx < Int32(self.route_layout.token_words_per_route): + if cutlass.const_expr(uses_one_warp_transport): + token_begin = Int32(self.route_layout.token_words_word_offset) + token_end = token_begin + Int32(self.route_layout.token_words_per_route) + if lane_idx >= token_begin and lane_idx < token_end: + self._smem_words[ + stage_base + + Int32(self.staging_layout.token_words_word_offset) + + lane_idx + - token_begin + ] = Int32(resolved_record_word) + elif lane_idx < Int32(self.route_layout.token_words_per_route): self._smem_words[ stage_base + Int32(self.staging_layout.token_words_word_offset) @@ -852,7 +1108,7 @@ def store_route( self, stage_info: StageInfo, *, - resolved_origin0: Int32, + resolved_record_word: Int32, resolved_origin1: Int32, resolved_atom_validity: Int32, route_record_word_offset: Int32, @@ -862,7 +1118,7 @@ def store_route( if cutlass.const_expr(self.cfg.use_keeps_mma_ab): self._store_route_keeps( stage_info, - resolved_origin0, + resolved_record_word, resolved_origin1, resolved_atom_validity, route_record_word_offset, @@ -870,7 +1126,7 @@ def store_route( else: self._store_route_swaps( stage_info, - resolved_origin0, + resolved_record_word, resolved_origin1, resolved_atom_validity, route_record_word_offset, @@ -886,8 +1142,8 @@ def _load_route_swaps_values( this Softmax warp's logical K32 slice; unused or invalid origins are negative. To preserve the shared seven-slot task ABI, origin 2/3 subsequently travel through the shared route-flags/token-word-0 slots. - Token-word 1 carries the logical K32 mask, token-word 2 carries - route-full, and token-word 3 is unused. + Token-word 1 carries the logical K32 mask, token-word 2 carries packed + route flags, and token-word 3 is unused. """ stage_base = self._consumer_stage_base() @@ -907,7 +1163,7 @@ def _load_route_swaps_values( origin3 = Int32(self._smem_words[warp_origin_base + Int32(3)]) token_word = Uint32(0xFFFFFFFF) - if cutlass.const_expr(self.cfg.use_kv_valid_bits): + if cutlass.const_expr(self.cfg.uses_prepared_score_keep_words): assert self.staging_layout.token_words_word_offset is not None token_word = Uint32( self._smem_words[ @@ -917,9 +1173,25 @@ def _load_route_swaps_values( ] ) route_flags = Uint32(0) - if cutlass.const_expr(_swaps_forwards_packed_route_full(self.cfg)): - route_flags = Uint32(origin0 & Int32(1)) - origin0 = origin0 & Int32(_SWAPS_PACKED_ROUTE_FULL_CLEAR_MASK) + if cutlass.const_expr( + _swaps_forwards_packed_route_full(self.cfg) + or self.cfg.use_block_sparse_proxy_routes + ): + packed_route_flags = origin0 & Int32( + _PREPARED_ROUTE_IS_FULL_FLAG | _PREPARED_ROUTE_IS_PROXY_FLAG + ) + route_flags = Uint32( + packed_route_flags & Int32(_PREPARED_ROUTE_IS_FULL_FLAG) + ) + if cutlass.const_expr(self.cfg.use_block_sparse_proxy_routes): + route_flags = route_flags | Uint32( + Int32( + (packed_route_flags & Int32(_PREPARED_ROUTE_IS_PROXY_FLAG)) + != Int32(0) + ) + * Int32(_SOFTMAX_ROUTE_IS_PROXY_FLAG) + ) + origin0 = origin0 & Int32(_SWAPS_PACKED_ROUTE_FLAGS_CLEAR_MASK) return ( origin0, origin1, @@ -986,17 +1258,21 @@ def load_route( valid0 = (stored_route_flags >> origin0_idx) & Int32(1) valid1 = (stored_route_flags >> origin1_idx) & Int32(1) route_flags = valid0 | (valid1 << Int32(1)) - if cutlass.const_expr(self.cfg.use_kv_valid_bits): + if cutlass.const_expr(self.cfg.uses_prepared_score_keep_words): route_flags = route_flags | ( stored_route_flags & Int32(_SOFTMAX_TOKEN_MASK_IS_FULL_FLAG) ) + if cutlass.const_expr(self.cfg.use_block_sparse_proxy_routes): + route_flags = route_flags | ( + stored_route_flags & Int32(_SOFTMAX_ROUTE_IS_PROXY_FLAG) + ) origin0 = Int32(self._smem_words[stage_base + origin0_idx]) origin1 = Int32(self._smem_words[stage_base + origin1_idx]) token_word0 = Uint32(0xFFFFFFFF) token_word1 = Uint32(0xFFFFFFFF) token_word2 = Uint32(0xFFFFFFFF) token_word3 = Uint32(0xFFFFFFFF) - if cutlass.const_expr(self.cfg.use_kv_valid_bits): + if cutlass.const_expr(self.cfg.uses_prepared_score_keep_words): assert self.staging_layout.token_words_word_offset is not None if cutlass.const_expr(self.route_layout.kv_route_size == 256): token_base = Int32(self.staging_layout.token_words_word_offset) @@ -1014,28 +1290,6 @@ def load_route( token_word3 = Uint32( self._smem_words[stage_base + token_base + word1_idx + Int32(1)] ) - elif cutlass.const_expr(self.cfg.tile_size_q == 64): - lane_idx = cute.arch.thread_idx()[0] & Int32(0x1F) - local_word_base = _keeps_col_base( - self.cfg, - lane_idx, - self.cfg.num_s_regs_per_thread, - ) >> Int32(5) - token_word0 = Uint32( - self._smem_words[ - stage_base - + Int32(self.staging_layout.token_words_word_offset) - + local_word_base - ] - ) - token_word1 = Uint32( - self._smem_words[ - stage_base - + Int32(self.staging_layout.token_words_word_offset) - + local_word_base - + Int32(1) - ] - ) else: token_word0 = Uint32( self._smem_words[ diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_p.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_p.py index bb27c3c553f0..330f15b05691 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_p.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_p.py @@ -25,7 +25,7 @@ import cutlass import cutlass.cute as cute -from cutlass import Float32, Int32, Int64 +from cutlass import Float32, Int32, Int64, Uint32 from cutlass.experimental import primitives as prims from cutlass.experimental.task_scheduling.memory import ( @@ -42,6 +42,7 @@ ) from ..fmha_decode_config import FmhaDecodeConfig +from ...._block_sparse.common import _block_sparse_proxy_summary_geometry from ...placeholder_helpers import _placeholder_smem_array from .helpers_common import ( Constexpr, @@ -58,11 +59,13 @@ _is_last_loop_iteration, _keeps_col_base, _keeps_row_idx, + _keeps_tcgen05_ld, _keeps_tcgen05_st, _named_barrier_arrive, _neg_max_f32, _pack_float2_to_bf16, _pack_float2_to_fp16, + _swaps_routed_coordinate, _wait_for_mbarrier_phase, ) from .helpers_output import ( @@ -76,11 +79,27 @@ _compute_fp8_p_regs_and_local_sums, _compute_fp8_p_regs_and_local_sums_dense, _compute_p_values_and_local_sums_dense, + _ex2_emulation_packed_f32x2, _pack_float4_to_fp8_e4m3, _pack_float4_to_fp8_e4m3_inline, ) +from .smem_block_sparse_metadata import _SOFTMAX_ROUTE_IS_PROXY_FLAG from .tmem_s import TmemSResource +# Tunable: number of score pairs per streamed fragment whose exponentials run +# as FMA polynomials instead of MUFU. The MUFU issue rate bounds the fragment +# otherwise, while the FMA pipe is nearly idle in the softmax warps. Larger +# shares grow the fragment body and the softmax warps become instruction-fetch +# bound again, so one quarter of the 16 pairs is the measured optimum. +KV_TILE_256_EX2_EMULATED_PAIRS = 4 + + +def _pair_uses_ex2_emulation(pair_idx: int, pairs_per_fragment: int) -> bool: + """Spread the emulated pairs evenly across a fragment's score pairs.""" + count = KV_TILE_256_EX2_EMULATED_PAIRS + pairs = pairs_per_fragment + return ((pair_idx + 1) * count) // pairs != (pair_idx * count) // pairs + @dataclass(kw_only=True) class SmemPResource(DecodeGenResourceBase): @@ -88,7 +107,7 @@ class SmemPResource(DecodeGenResourceBase): Softmax producers convert S to P, store it in the profile's TMEM or SMEM layout, and publish local sums back to TmemS. Most profiles use the generic - full/empty P pipeline. KV256 instead publishes four independently ready + full/empty P pipeline. Streamed profiles instead publish four independently ready K32 TMEM fragments; BMM2 consumes those fragments in order, while the matching TmemO full barrier prevents the next QK from overwriting aliased P. """ @@ -152,7 +171,7 @@ def _init_placeholder_state(self) -> None: ) def get_smem_requirements(self) -> list[SmemAllocation]: - """Allocate P storage or the KV256 fragment-ready barriers.""" + """Allocate P storage or the streamed fragment-ready barriers.""" if self.cfg.streams_tmem_p_fragments: if self._fragment_ready_alloc is None: self._fragment_ready_alloc = SmemAllocation( @@ -173,7 +192,7 @@ def get_smem_requirements(self) -> list[SmemAllocation]: @cute.jit def _bind_fragment_ready(self, context: ResourceContext | None = None) -> None: - """Bind the one-way KV256 P-ready barriers from the SMEM context.""" + """Bind the one-way streamed P-ready barriers from the SMEM context.""" if cutlass.const_expr( self.cfg.streams_tmem_p_fragments and context is not None @@ -191,7 +210,7 @@ def _bind_fragment_ready(self, context: ResourceContext | None = None) -> None: def create_function_variables( self, context: ResourceContext | None = None ) -> ResourceVars: - """Bind and initialize KV256's per-fragment ready barriers.""" + """Bind and initialize the streamed per-fragment ready barriers.""" self._bind_fragment_ready(context) if cutlass.const_expr(self.cfg.streams_tmem_p_fragments): tidx, _, _ = cute.arch.thread_idx() @@ -299,43 +318,203 @@ def init_descriptor_state(self, stage_info: StageInfo) -> None: # work can publish a valid descriptor or TMEM address for this tile. self._create_initial_task_locals(stage_info.context) + @cute.jit + def _apply_proxy_route_denominator_mass( + self, + local_sum: Float32, + tail_p: Float32, + route_is_proxy: Int32, + ) -> Float32: + """Weight only a proxy route's softmax denominator by block mass.""" + + if cutlass.const_expr(not self.cfg.use_block_sparse_proxy_routes): + return local_sum + if route_is_proxy != Int32(0): + _, tail_len = _block_sparse_proxy_summary_geometry( + self.cfg.static_seq_len_kv, + self.cfg.kv_block_size, + ) + local_sum *= Float32(self.cfg.kv_block_size) + tail_delta = tail_len - self.cfg.kv_block_size + if cutlass.const_expr(tail_delta != 0): + local_sum += Float32(tail_delta) * tail_p + return local_sum + @producer_work @cute.jit - def compute_p_fragment( + def compute_p_fragments( self, stage_info: StageInfo, *, - fragment_idx: Constexpr[int], new_max_arr: cutlass.Array, - s_arr: cutlass.Array, ) -> None: - """Convert one KV256 K32 score fragment and publish its TMEM P slice.""" + """Stream every ordinary K32 fragment from one rolled loop.""" + self._compute_p_fragments_impl( + stage_info, + new_max_arr=new_max_arr, + route_is_proxy=Int32(0), + route_origin0=Int32(0), + route_origin1=Int32(0), + ) + + @producer_work + @cute.jit + def compute_proxy_route_p_fragments( + self, + stage_info: StageInfo, + *, + new_max_arr: cutlass.Array, + route_flags: Int32, + route_origin0: Int32, + route_origin1: Int32, + ) -> None: + """Stream every proxy-capable KV256 K32 fragment from one rolled loop.""" + assert self.cfg.use_block_sparse_proxy_routes + route_is_proxy = Int32( + (route_flags & Int32(_SOFTMAX_ROUTE_IS_PROXY_FLAG)) != Int32(0) + ) + self._compute_p_fragments_impl( + stage_info, + new_max_arr=new_max_arr, + route_is_proxy=route_is_proxy, + route_origin0=route_origin0, + route_origin1=route_origin1, + ) + + @cute.jit + def _compute_p_fragments_impl( + self, + stage_info: StageInfo, + *, + new_max_arr: cutlass.Array, + route_is_proxy: Int32, + route_origin0: Int32, + route_origin1: Int32, + ) -> None: + """Reload, exponentiate, and publish all K32 fragments in a rolled loop. + + The fragment index is a runtime loop variable, so the exponentiation + body exists once in the instruction stream and only the TMEM column + offset, the fragment barrier, and the proxy tail bookkeeping depend on + it. Unrolling the fragments would replicate that body for every + fragment and both softmax instances and leave the softmax warps + instruction-fetch bound. The max pass has already written masked + scores back to TMEM, so the reload needs no mask logic of its own. + """ + _ = stage_info cfg = self.cfg assert cfg.streams_tmem_p_fragments - assert not cfg.use_fp8_qkv and cfg.uses_two_inst_tmem_p - assert cfg.softmax_score_fragment_regs == 32 + assert self._tmem_alloc.offset == self.tmem_s_ref._alloc.offset + # One FP32 score per column, two packed 16-bit probabilities per column. + fragment_regs = cfg.softmax_score_fragment_regs + fragment_cols = fragment_regs // 2 new_max = new_max_arr[0] safe_new_max = new_max if safe_new_max == _neg_max_f32(): safe_new_max = Float32(0.0) minus_max_scale = Float32(-self.scale_softmax_log2 * safe_new_max) + tmem_base = self._tmem_base_addr + Int32(self._tmem_alloc.offset) + tidx, _, _ = cute.arch.thread_idx() + publishes_fragment = (tidx & Int32(31)) == Int32(0) + + total_sum = Float32(0.0) + for fragment_idx in cutlass.range(cfg.num_softmax_score_fragments, unroll=1): + fragment = Int32(fragment_idx) + loaded = _keeps_tcgen05_ld( + cfg, + prims.make_tmem_ptr( + tmem_base + fragment * Int32(fragment_regs), Float32 + ), + num=fragment_regs, + offset=cfg.tile_size_kv // 2, + ) + prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) + s_arr = cutlass.Array( + Float32, fragment_regs, space=cutlass.AddressSpace.rmem + ) + for score_idx in cutlass.range_constexpr(fragment_regs): + s_arr[score_idx] = loaded[score_idx] + + local_sum = self._exponentiate_fragment_pairs(s_arr, minus_max_scale) + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): + if route_is_proxy != Int32(0): + local_sum = self._proxy_fragment_sum( + local_sum, + s_arr, + fragment_origin=self._runtime_fragment_origin( + fragment, route_origin0, route_origin1 + ), + ) + + packed_p = ( + s_arr.data_ptr() + .load(count=fragment_regs, alignment=4) + .to(cfg.q_dtype) + .bitcast(Int32) + ) + _keeps_tcgen05_st( + cfg, + prims.make_tmem_ptr(tmem_base + fragment * Int32(fragment_cols), Int32), + packed_p, + offset=cfg.tmem_p_cols_per_inst, + ) + cute.arch.fence_view_async_tmem_store() + prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) + if publishes_fragment: + prims.mbarrier_arrive(self._fragment_ready.data_ptr() + fragment) + total_sum += local_sum + self.tmem_s_ref.store_p_local_sum(0, total_sum) + + @cute.jit + def _runtime_fragment_origin( + self, fragment: Int32, route_origin0: Int32, route_origin1: Int32 + ) -> Int32: + """Return the token origin of a fragment selected at runtime. + + Each lane's fragments cover two K64 route atoms in order: the first + atom's fragments start at ``route_origin0``, the second atom's at + ``route_origin1``, and consecutive fragments within an atom advance by + one fragment width. + """ + cfg = self.cfg + fragment_regs = cfg.softmax_score_fragment_regs + fragments_per_origin = cfg.softmax_fragments_per_route_atom + fragment_origin = Int32(route_origin0) + if fragment >= Int32(fragments_per_origin): + fragment_origin = Int32(route_origin1) + return fragment_origin + (fragment % Int32(fragments_per_origin)) * Int32( + fragment_regs + ) + + @cute.jit + def _exponentiate_fragment_pairs( + self, s_arr: cutlass.Array, minus_max_scale: Float32 + ) -> Float32: + """Turn one fragment of scaled scores into probabilities in place. - # Eight independent chains keep the denominator update off one long - # dependency chain. Reuse s_arr for probabilities so only one K32 score - # fragment remains live while P is packed. + Returns the fragment's probability sum. Eight independent chains keep + the denominator update off one long dependency chain, and a configurable + subset of pairs runs its exponentials on the FMA pipe. + """ + pairs_per_fragment = self.cfg.softmax_score_fragment_regs // 2 sum_chains = cutlass.Array(Float32, 8, space=cutlass.AddressSpace.rmem) for chain_idx in cutlass.range_constexpr(8): sum_chains[chain_idx] = Float32(0.0) - for pair_idx in cutlass.range_constexpr(16): + for pair_idx in cutlass.range_constexpr(pairs_per_fragment): value_idx = pair_idx * 2 p0, p1 = cute.arch.fma_packed_f32x2( (Float32(s_arr[value_idx]), Float32(s_arr[value_idx + 1])), (self.scale_softmax_log2, self.scale_softmax_log2), (minus_max_scale, minus_max_scale), ) - p0 = Float32(cute.math.exp2(p0, fastmath=True)) - p1 = Float32(cute.math.exp2(p1, fastmath=True)) + if cutlass.const_expr( + _pair_uses_ex2_emulation(pair_idx, pairs_per_fragment) + ): + p0, p1 = _ex2_emulation_packed_f32x2(p0, p1) + else: + p0 = Float32(cute.math.exp2(p0, fastmath=True)) + p1 = Float32(cute.math.exp2(p1, fastmath=True)) s_arr[value_idx] = p0 s_arr[value_idx + 1] = p1 chain_idx = (pair_idx & 3) * 2 @@ -345,10 +524,6 @@ def compute_p_fragment( (p0, p1), ) ) - - # Collapse the eight reduction chains before packing P and publishing - # its barrier. This keeps only one sum scalar live across STTM instead - # of overlapping the full reduction state with packed P and addresses. sum01 = cute.arch.add_packed_f32x2( (sum_chains[0], sum_chains[1]), (sum_chains[2], sum_chains[3]), @@ -358,40 +533,43 @@ def compute_p_fragment( (sum_chains[6], sum_chains[7]), ) total_pair = cute.arch.add_packed_f32x2(sum01, sum23) - local_sum = Float32(total_pair[0] + total_pair[1]) + return Float32(total_pair[0] + total_pair[1]) - packed_p = ( - s_arr.data_ptr().load(count=32, alignment=4).to(cfg.q_dtype).bitcast(Int32) - ) + @cute.jit + def _proxy_fragment_sum( + self, + local_sum: Float32, + s_arr: cutlass.Array, + *, + fragment_origin: Int32, + ) -> Float32: + """Weight a proxy fragment's sum by the token mass each summary stands for. - fragment_cols = cfg.softmax_score_fragment_regs // 2 - p_tmem_addr = ( - self._tmem_base_addr - + Int32(self._tmem_alloc.offset) - + Int32(fragment_idx * fragment_cols) - ) - _keeps_tcgen05_st( - cfg, - prims.make_tmem_ptr(p_tmem_addr, Int32), - packed_p, - offset=cfg.tmem_p_cols_per_inst, + KC stores one mean K vector per semantic KV block while VC stores its V + sum. P itself stays unweighted for PV; only the denominator accounts + for the represented token count, with the final summary covering the + shorter tail block. + """ + cfg = self.cfg + fragment_regs = cfg.softmax_score_fragment_regs + num_summaries, tail_len = _block_sparse_proxy_summary_geometry( + cfg.static_seq_len_kv, + cfg.kv_block_size, ) - # This lowers to the warp-collective tcgen05.wait::st. The explicit - # proxy fence then makes every lane's completed STTM visible through - # the lane-0 mbarrier publication consumed by the MMA warp. - cute.arch.fence_view_async_tmem_store() - prims.tcgen05_fence(prims.Tcgen05Fence.BEFORE_THREAD_SYNC) - - # KV256 aliases P with the score tile that produced it. Each softmax - # warp publishes its own rows after the TMEM store drains; BMM2 waits - # for all producer warps before consuming the fragment. - tidx, _, _ = cute.arch.thread_idx() - if (tidx & Int32(31)) == Int32(0): - prims.mbarrier_arrive(self._fragment_ready.data_ptr() + Int32(fragment_idx)) - - if cutlass.const_expr(fragment_idx != 0): - local_sum += self.tmem_s_ref.load_p_local_sum(0) - self.tmem_s_ref.store_p_local_sum(0, local_sum) + local_sum *= Float32(cfg.kv_block_size) + tail_delta = tail_len - cfg.kv_block_size + if cutlass.const_expr(tail_delta != 0): + final_summary_idx = num_summaries - 1 + final_summary_offset = Int32(final_summary_idx) - fragment_origin + if final_summary_offset >= Int32(0) and final_summary_offset < Int32( + fragment_regs + ): + # Proxy fragment origins are fragment-aligned in summary + # coordinates, so the tail's in-fragment lane is a compile-time + # constant even though route ownership is decided at runtime. + tail_lane = final_summary_idx % fragment_regs + local_sum += Float32(tail_delta) * Float32(s_arr[tail_lane]) + return local_sum @cute.jit def _compute_keeps_p( @@ -401,16 +579,18 @@ def _compute_keeps_p( new_max_arr: cutlass.Array, s_arr: cutlass.Array, ) -> None: - """Materialize one non-KV256 row-major Keeps probability tile. + """Materialize one complete-row Keeps probability tile. TQ128 gives each warp-group thread a complete 128-column row. TQ64 gives paired lanes the low/high 64-column halves of one row. Each lane writes disjoint packed blocks into the TMEM or SMEM layout consumed by - BMM2. + BMM2. Streamed profiles, including every block-sparse Keeps profile, + produce P through the rolled fragment loop instead. """ cfg = self.cfg - # KV256 uses compute_p_fragment so only one K32 score fragment is live. - assert not cfg.streams_tmem_p_fragments + # Every block-sparse Keeps profile streams P; only dense complete rows + # reach this path. + assert not cfg.streams_tmem_p_fragments and not cfg.use_block_sparse task_cache = _decode_gen_task_cache(stage_info) warp_grp_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] lane_idx = task_cache[_TASK_CACHE_LANE_IDX] @@ -439,7 +619,6 @@ def _compute_keeps_p( # without keeping a second 16-value P array live beside the S row. local_sum_pair_01 = (Float32(0.0), Float32(0.0)) local_sum_pair_23 = (Float32(0.0), Float32(0.0)) - # Each vector block is exactly 16 bytes after conversion. Compute and # pack adjacent pairs directly into their final register payload. packed_p_regs = cfg.num_packed_p_regs if cfg.uses_two_inst_tmem_p else 4 @@ -538,10 +717,10 @@ def _compute_keeps_p( packed_p.data_ptr().load(count=4, alignment=4), alignment=16 ) if cutlass.const_expr(cfg.uses_two_inst_tmem_p): - # FP8 publishes a complete row with one x16/x32 STTM. FP16/BF16 - # uses x16 slices to limit Softmax register pressure. This is the - # complete-row Q128/KV128 path; KV256 publishes K32 fragments. - assert cfg.num_packed_p_regs in (16, 32, 64) + # FP8 publishes the complete row with one x32 STTM. Dense 16-bit + # Q128/KV128 uses x16 slices to limit Softmax register pressure. + # Block-sparse two-instance profiles stream K32 fragments instead. + assert cfg.num_packed_p_regs in (32, 64) regs_per_store = cfg.num_packed_p_regs if cfg.use_fp8_qkv else 16 assert cfg.num_packed_p_regs % regs_per_store == 0 for store_idx in cutlass.range_constexpr( @@ -591,14 +770,18 @@ def _compute_keeps_p( # point and no extra named barrier is needed here. cute.arch.fence_view_async_shared() - @producer_work @cute.jit - def compute_p( + def _compute_p_impl( self, stage_info: StageInfo, *, new_max_arr: cutlass.Array, s_arr: cutlass.Array, + route_is_proxy: Int32, + route_origin0: Int32, + route_origin1: Int32, + route_origin2: Int32, + route_origin3: Int32, ) -> None: """Compute P from S, stage its BMM2 operand, and publish local sums.""" cfg = self.cfg @@ -616,6 +799,7 @@ def compute_p( # warp/lane ownership for SMEM offsets and STSM swizzles. task_cache = _decode_gen_task_cache(stage_info) warp_grp_thread_idx = task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX] + lane_idx = Int32(task_cache[_TASK_CACHE_LANE_IDX]) if cutlass.const_expr(cfg.tile_size_q == 32 and cfg.use_fp8_qkv): # Tile-Q=32 FP8 fast path: compute E4M3 P registers in the # same order consumed by the STSM helper, while also capturing @@ -796,10 +980,14 @@ def compute_p( local_sums = cutlass.Array( Float32, num_scale_groups, space=cutlass.AddressSpace.rmem ) + proxy_tail_p = cutlass.Array( + Float32, num_scale_groups, space=cutlass.AddressSpace.rmem + ) for idx in cutlass.range_constexpr(num_s_regs): p_vals[idx] = Float32(0.0) for idx in cutlass.range_constexpr(num_scale_groups): local_sums[idx] = Float32(0.0) + proxy_tail_p[idx] = Float32(0.0) for scale_idx in cutlass.range_constexpr(num_scale_groups): # Convert each softmax scale group from S to P. Masked rows have @@ -828,9 +1016,32 @@ def compute_p( ) p_vals[s_idx] = p_val local_sums[scale_idx] += p_val + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): + num_summaries, _ = _block_sparse_proxy_summary_geometry( + cfg.static_seq_len_kv, + cfg.kv_block_size, + ) + atom_origin, logical_summary = _swaps_routed_coordinate( + cfg, + lane_idx >> Int32(2), + route_origin0, + route_origin1, + route_origin2, + route_origin3, + token_group_idx=k_pair_idx, + ) + if atom_origin >= Int32(0) and logical_summary == Int32( + num_summaries - 1 + ): + proxy_tail_p[scale_idx] = p_val # Hand off denominator contributions through TmemS. P remains a pure # MMA operand in SMEM; sums are not reloaded from the P tile. for scale_idx in cutlass.range_constexpr(num_scale_groups): + local_sums[scale_idx] = self._apply_proxy_route_denominator_mass( + local_sums[scale_idx], + proxy_tail_p[scale_idx], + route_is_proxy, + ) self.tmem_s_ref.store_p_local_sum(scale_idx, local_sums[scale_idx]) if cutlass.const_expr(cfg.use_fp8_qkv): @@ -1057,6 +1268,33 @@ def compute_p( p_vals[p_base + 4] = p_pair[1] local_sum[scale_idx] += p_pair[0] local_sum[scale_idx] += p_pair[1] + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): + num_summaries, _ = _block_sparse_proxy_summary_geometry( + cfg.static_seq_len_kv, + cfg.kv_block_size, + ) + for scale_idx in cutlass.range_constexpr(cfg.num_softmax_scale_groups): + proxy_tail_p = Float32(0.0) + for token_group_idx in cutlass.range_constexpr(4): + atom_origin, logical_summary = _swaps_routed_coordinate( + cfg, + lane_idx >> Int32(2), + route_origin0, + route_origin1, + route_origin2, + route_origin3, + token_group_idx=token_group_idx, + ) + if atom_origin >= Int32(0) and logical_summary == Int32( + num_summaries - 1 + ): + p_idx = scale_idx + token_group_idx * 2 + proxy_tail_p = Float32(p_vals[p_idx]) + local_sum[scale_idx] = self._apply_proxy_route_denominator_mass( + local_sum[scale_idx], + proxy_tail_p, + route_is_proxy, + ) # Pack the P scalars to match the dtype consumed by BMM2. regs_p = cutlass.Array( Int32, cfg.num_packed_p_regs, space=cutlass.AddressSpace.rmem @@ -1107,6 +1345,77 @@ def compute_p( # BMM2 cannot observe a partially written P tile. prims.barrier_cta_sync(4 + self.inst_id, thread_count=128) + @producer_work + @cute.jit + def compute_p( + self, + stage_info: StageInfo, + *, + new_max_arr: cutlass.Array, + s_arr: cutlass.Array, + ) -> None: + """Compute an exact/dense P tile without typed-route metadata.""" + + self._compute_p_impl( + stage_info, + new_max_arr=new_max_arr, + s_arr=s_arr, + route_is_proxy=Int32(0), + route_origin0=Int32(0), + route_origin1=Int32(0), + route_origin2=Int32(0), + route_origin3=Int32(0), + ) + + @producer_work + @cute.jit + def compute_proxy_route_p( + self, + stage_info: StageInfo, + *, + new_max_arr: cutlass.Array, + s_arr: cutlass.Array, + route_origin0: Int32, + route_origin1: Int32, + keeps_route_flags_or_swaps_origin2: Int32, + swaps_route_origin3_bits: Uint32, + swaps_route_flags: Uint32, + ) -> None: + """Normalize the active Keeps/SWAP metadata view and compute P. + + The shared Int32 input is Keeps route flags or SWAP origin2. SWAP's + origin3 and flags stay bit-preserving Uint32 values until this work + boundary because schedule-level dataflow tokens cannot be cast. + """ + + assert self.cfg.use_block_sparse_proxy_routes + route_origin2 = Int32(0) + route_origin3 = Int32(0) + if cutlass.const_expr(self.cfg.use_keeps_mma_ab): + route_is_proxy = Int32( + ( + keeps_route_flags_or_swaps_origin2 + & Int32(_SOFTMAX_ROUTE_IS_PROXY_FLAG) + ) + != Int32(0) + ) + else: + route_is_proxy = Int32( + (swaps_route_flags & Uint32(_SOFTMAX_ROUTE_IS_PROXY_FLAG)) != Uint32(0) + ) + route_origin2 = keeps_route_flags_or_swaps_origin2 + route_origin3 = swaps_route_origin3_bits.bitcast(Int32) + self._compute_p_impl( + stage_info, + new_max_arr=new_max_arr, + s_arr=s_arr, + route_is_proxy=route_is_proxy, + route_origin0=route_origin0, + route_origin1=route_origin1, + route_origin2=route_origin2, + route_origin3=route_origin3, + ) + @consumer_work( returns=( p_desc_0_slot, @@ -1136,7 +1445,7 @@ def p_operands( # stats-free columns of the corresponding S stage. p_stage_cols = cfg.tmem_s_cols if cutlass.const_expr(cfg.streams_tmem_p_fragments): - # KV256's four pipeline stages are K32 fragments of one P + # A streamed profile's four pipeline stages are K32 fragments of one P # operand, not four independent full S/P stages. p_stage_cols = cfg.softmax_score_fragment_regs // 2 p_tmem_addr = self._tmem_base_addr + Int32( @@ -1170,7 +1479,7 @@ def wait_p_fragment( *, fragment_idx: Constexpr[int], ) -> Int32: - """Wait for and return the next KV256 P-fragment TMEM address.""" + """Wait for and return the next streamed P-fragment TMEM address.""" cfg = self.cfg _ = stage_info assert cfg.streams_tmem_p_fragments @@ -1185,25 +1494,3 @@ def wait_p_fragment( self._tmem_alloc.offset + fragment_idx * fragment_cols ) return p_tmem_addr - - @consumer_work(work_attrs=WorkAttr.AUXILIARY) - @cute.jit - def wait_until_reusable_before_qk(self, stage_info: StageInfo) -> None: - """Wait until the previous same-instance PV has stopped reading P. - - KV256 aliases each streamed P instance with its next S accumulator. - The existing two-stage O pipeline commits stage ``inst_id`` only when - the matching PV completes, so its full barrier is also the P-reuse - credit. The S producer phase supplies the generation: the first QK - waits on the initially complete opposite parity, and every later QK - waits for the preceding PV without another commit or barrier. - """ - _ = stage_info - cfg = self.cfg - assert cfg.streams_tmem_p_fragments - assert cfg.o_stages == cfg.num_insts_kv == 2 - barrier = self.tmem_o_ref.pipeline.sync_object_full.get_barrier( - Int32(self.inst_id) - ) - _wait_for_mbarrier_phase(barrier, self.tmem_s_ref.producer_state.phase) - prims.tcgen05_fence(prims.Tcgen05Fence.AFTER_THREAD_SYNC) diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_resources.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_resources.py index 20bb78fb810c..173a13cb5447 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_resources.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/smem_resources.py @@ -497,6 +497,10 @@ class SmemKvTileResource(DecodeGenResourceBase): tma_desc_v: cutlass.Pointer | None = None tma_desc_k_atom: cutlass.Pointer | None = None tma_desc_v_atom: cutlass.Pointer | None = None + tma_desc_k_summary: cutlass.Pointer | None = None + tma_desc_v_summary: cutlass.Pointer | None = None + tma_desc_k_summary_atom: cutlass.Pointer | None = None + tma_desc_v_summary_atom: cutlass.Pointer | None = None sparse_kv_metadata: "SmemBlockSparseKvMetadataResource | None" = None page_offsets_kv: "SmemPageOffsetsKvResource | None" = None seqlens_kv: cute.Pointer | None = None @@ -703,15 +707,24 @@ def _producer_load( assert self.sparse_kv_metadata is not None assert self.tma_desc_k_atom is not None assert self.tma_desc_v_atom is not None - # The positional TensorMaps keep the decode ABI stable. The - # primary K/V descriptors are KV128 for coarse routes and one atom - # for fine routes. The auxiliary slots always expose the atom - # descriptor and alias the primary descriptor for fine routes. + # K/V and summary sources expose the same primary/atom descriptor + # pair. Route kind selects the source; the geometry below alone + # selects the physical copy policy. tma_desc_atom = ( self.tma_desc_v_atom if cutlass.const_expr(self.kv_kind == KV_KIND_V) else self.tma_desc_k_atom ) + tma_desc_summary = ( + self.tma_desc_v_summary + if cutlass.const_expr(self.kv_kind == KV_KIND_V) + else self.tma_desc_k_summary + ) + tma_desc_summary_atom = ( + self.tma_desc_v_summary_atom + if cutlass.const_expr(self.kv_kind == KV_KIND_V) + else self.tma_desc_k_summary_atom + ) kv_atom_size = _block_sparse_kv_atom_size(cfg.kv_block_size) head_dim_stage = cfg.head_dim_kv_stage head_dim_stage_offset = head_dim_stage_idx * head_dim_stage @@ -775,6 +788,12 @@ def _producer_load( # join unrelated entries and must prove physical adjacency. fragment_chunk_elems = chunk_hd * 64 if prims.elect_sync(): + route_tma_desc = tma_desc + route_tma_desc_atom = tma_desc_atom + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): + if self.sparse_kv_metadata.route_is_proxy(): + route_tma_desc = tma_desc_summary + route_tma_desc_atom = tma_desc_summary_atom origin0, _ = self.sparse_kv_metadata.route_tma_coordinate( Int32(0), logical_b_idx, @@ -800,7 +819,7 @@ def _producer_load( local_tile_offset = chunk_idx * tile_chunk_elems prims.cp_async_bulk_tensor_shared_cta_global( stage_base.subview(local_tile_offset), - tma_desc, + route_tma_desc, ( Int32(global_head_dim_offset), origin0, @@ -832,7 +851,7 @@ def _producer_load( if adjacent: prims.cp_async_bulk_tensor_shared_cta_global( stage_base.subview(local_tile_offset), - tma_desc, + route_tma_desc, ( Int32(global_head_dim_offset), origin0, @@ -844,7 +863,7 @@ def _producer_load( else: prims.cp_async_bulk_tensor_shared_cta_global( stage_base.subview(local_tile_offset), - tma_desc_atom, + route_tma_desc_atom, ( Int32(global_head_dim_offset), origin0, @@ -857,7 +876,7 @@ def _producer_load( stage_base.subview( local_tile_offset + fragment_chunk_elems ), - tma_desc_atom, + route_tma_desc_atom, ( Int32(global_head_dim_offset), origin1, @@ -875,6 +894,10 @@ def _producer_load( atom_chunk_elems = chunk_hd * kv_atom_size atoms_per_route = cfg.tile_size_kv // kv_atom_size if prims.elect_sync(): + route_tma_desc_atom = tma_desc_atom + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): + if self.sparse_kv_metadata.route_is_proxy(): + route_tma_desc_atom = tma_desc_summary_atom stage_base = self._stage_base(stage_info) # Reuse each retained origin across all head-dimension # chunks. The copies still target disjoint SMEM regions @@ -913,7 +936,7 @@ def _producer_load( stage_base.subview( local_tile_offset + atom_idx * atom_chunk_elems ), - tma_desc_atom, + route_tma_desc_atom, ( Int32(global_head_dim_offset), origin, @@ -1487,6 +1510,10 @@ class SmemKvResource(DecodeGenResourceBase): tma_desc_v: cutlass.Pointer | None = None tma_desc_k_atom: cutlass.Pointer | None = None tma_desc_v_atom: cutlass.Pointer | None = None + tma_desc_k_summary: cutlass.Pointer | None = None + tma_desc_v_summary: cutlass.Pointer | None = None + tma_desc_k_summary_atom: cutlass.Pointer | None = None + tma_desc_v_summary_atom: cutlass.Pointer | None = None sparse_kv_metadata0: "SmemBlockSparseKvMetadataResource | None" = None sparse_kv_metadata1: "SmemBlockSparseKvMetadataResource | None" = None page_offsets_kv: SmemPageOffsetsKvResource | None = None @@ -1922,6 +1949,29 @@ def _producer_load_kv_tile_256( ): assert self.page_offsets_kv is not None dense_page_ids = self.page_offsets_kv.page_ids(grouped_tile_idx) + # Select the logical source before the constexpr 4 x 2 loop so K/V + # and summary routes retain one physical KV256 staging body. + route_tma_desc = tma_desc + if cutlass.const_expr(cfg.use_block_sparse): + sparse_tma_desc = ( + self.tma_desc_v_atom + if cutlass.const_expr(kv_kind == KV_KIND_V) + else self.tma_desc_k_atom + ) + assert sparse_tma_desc is not None + route_tma_desc = sparse_tma_desc + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): + assert sparse_kv_metadata is not None + summary_tma_desc = ( + self.tma_desc_v_summary_atom + if cutlass.const_expr(kv_kind == KV_KIND_V) + else self.tma_desc_k_summary_atom + ) + assert summary_tma_desc is not None + route_is_proxy = sparse_kv_metadata.route_is_proxy() + route_tma_desc = ( + summary_tma_desc if route_is_proxy else sparse_tma_desc + ) for semantic_block in cutlass.range_constexpr(4): token_coord = Int32(0) storage_coord = logical_b_idx @@ -1950,15 +2000,9 @@ def _producer_load_kv_tile_256( ) if cutlass.const_expr(cfg.use_block_sparse): - sparse_tma_desc = ( - self.tma_desc_v_atom - if cutlass.const_expr(kv_kind == KV_KIND_V) - else self.tma_desc_k_atom - ) - assert sparse_tma_desc is not None prims.cp_async_bulk_tensor_shared_cta_global( stage_base.subview(block_base), - sparse_tma_desc, + route_tma_desc, ( Int32(dim_half * 64), token_coord, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_corr.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_corr.py index c59712aa8f3c..11b384a103bd 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_corr.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_corr.py @@ -91,7 +91,9 @@ _KV_TILE_256_CORRECTION_THREADS = 128 _KV_TILE_256_LOGICAL_OUTPUT_ROWS = 64 -_KV_TILE_256_EXCHANGE_ROW_STRIDE = 132 +# One D32 fragment per logical output row, padded by four floats so adjacent +# rows fall on different bank groups. +_KV_TILE_256_EXCHANGE_FRAGMENT_STRIDE = 36 _KV_TILE_256_STATS_PER_THREAD = 4 @@ -157,10 +159,10 @@ def get_o_stage_dtype_bytes(self) -> int: ) def _kv_tile_256_exchange_entries(self) -> int: - """Return 128 lane-local stats plus 64 logical output rows.""" + """Return 128 lane-local stats plus one D32 fragment per output row.""" return ( _KV_TILE_256_CORRECTION_THREADS * _KV_TILE_256_STATS_PER_THREAD - + _KV_TILE_256_LOGICAL_OUTPUT_ROWS * _KV_TILE_256_EXCHANGE_ROW_STRIDE + + _KV_TILE_256_LOGICAL_OUTPUT_ROWS * _KV_TILE_256_EXCHANGE_FRAGMENT_STRIDE ) def _init_placeholder_state(self) -> None: @@ -306,22 +308,14 @@ def get_smem_requirements(self) -> list[SmemAllocation]: ) if self.cfg.tile_size_kv == 256 and self._kv_tile_256_exchange_alloc is None: # Tail correction exchanges all lane-local stats, then pipelines - # D32 fragments through 64 logical output rows. Upper lanes publish - # one spatial half while lower lanes retain the matching fragment - # in registers. The dependency graph places this scratch after the - # shared KV ring so it can reuse the dead storage. - payload_bytes = self._kv_tile_256_exchange_entries() * 4 - exchange_bytes = payload_bytes - if self.cfg.uses_rotating_kv256_exchange: - assert payload_bytes <= self.cfg.smem_kv_tile_bytes - # Runtime selects one compact payload inside this explicit - # full-ring alias envelope. The envelope keeps every dynamic - # pointer within a declared allocation while the actual live - # exchange remains only 35,840 B in one 64-KiB stage. - exchange_bytes = self.cfg.smem_kv_tile_bytes * self.cfg.kv_stages + # D32 fragments through 64 logical output rows one fragment at a + # time. Upper lanes publish one spatial half while lower lanes + # retain the matching fragment in registers. The buffer is + # dedicated, so the shared KV ring keeps streaming the next tile's + # routes while the tail runs. self._kv_tile_256_exchange_alloc = SmemAllocation( name=f"{self.name}_kvTile256Exchange", - size_bytes=exchange_bytes, + size_bytes=self._kv_tile_256_exchange_entries() * 4, alignment=16, ) allocs = [] @@ -874,60 +868,145 @@ def _fold_split_o_vec8( return output_vals, sum_val, new_max, new_max @cute.jit - def _store_final_o_vec8( + def _merge_kv_tile_256_peer_fragment( + self, + own_vals: cutlass.Array, + peer_vals, + first_col: Constexpr[int], + count: Constexpr[int], + ) -> cutlass.Array: + """Add the peer spatial half to ``count`` own columns from ``first_col``.""" + merged_vals = cutlass.Array( + Float32, + count, + space=cutlass.AddressSpace.rmem, + ) + for elem in cutlass.range_constexpr(0, count, 2): + value_idx = first_col + elem + merged = fadd2( + (own_vals[value_idx], own_vals[value_idx + 1]), + ( + Float32(peer_vals[value_idx]), + Float32(peer_vals[value_idx + 1]), + ), + ) + merged_vals[elem] = merged[0] + merged_vals[elem + 1] = merged[1] + return merged_vals + + @cute.jit + def _store_final_o_columns( self, final_o_dst, output_vals: cutlass.Array, norm_scale: Float32, + *, + count: Constexpr[int], + sector_aligned: cutlass.Boolean, ) -> None: - """Pack one contiguous 8-element output fragment to the final O dtype.""" + """Scale, pack, and store ``count`` contiguous final output columns. + + FP8 output packs four values per register and writes eight columns per + 8-byte store. 16-bit output packs pairs; sixteen columns fill one + 32-byte sector and go out as a single 256-bit store when the + destination is sector aligned, otherwise every eight columns use one + 16-byte store. Callers choose ``count`` per path: the KV256 tail owns + whole rows per lane and pays for half-written sectors, while the + split-KV reducers write eight-column fragments. + """ cfg = self.cfg + assert count % 8 == 0 if cutlass.const_expr(cfg.use_fp8_output): - final_pairs = cutlass.Array(Float32, 8, space=cutlass.AddressSpace.rmem) - for pair_idx in cutlass.range_constexpr(4): - val_base = pair_idx * 2 - pair = fmul2( - (norm_scale, norm_scale), - (output_vals[val_base], output_vals[val_base + 1]), + for chunk_idx in cutlass.range_constexpr(count // 8): + fp8_regs = self._pack_fp8_output_quads( + output_vals, norm_scale, chunk_idx * 8 ) - final_pairs[val_base] = pair[0] - final_pairs[val_base + 1] = pair[1] - final_fp8_regs = cutlass.Array(Int32, 2, space=cutlass.AddressSpace.rmem) - final_fp8_regs[0] = _pack_float4_to_fp8_e4m3( - final_pairs[0], - final_pairs[1], - final_pairs[2], - final_pairs[3], - ) - final_fp8_regs[1] = _pack_float4_to_fp8_e4m3( - final_pairs[4], - final_pairs[5], - final_pairs[6], - final_pairs[7], - ) - final_o_dst.store( - final_fp8_regs.data_ptr().load(count=2, alignment=4), - alignment=8, - ) - else: - final_regs = cutlass.Array(Int32, 4, space=cutlass.AddressSpace.rmem) - for reg_idx in cutlass.range_constexpr(4): - pair = fmul2( - (norm_scale, norm_scale), - ( - output_vals[reg_idx * 2], - output_vals[reg_idx * 2 + 1], - ), + (final_o_dst + Int32(chunk_idx * 2)).store( + fp8_regs.data_ptr().load(count=2, alignment=4), + alignment=8, ) - if cutlass.const_expr(cfg.use_bf16_output): - final_regs[reg_idx] = _pack_float2_to_bf16(pair[0], pair[1]) + else: + final_regs = self._pack_final_o_regs(output_vals, norm_scale, count) + if cutlass.const_expr(count == 16): + if sector_aligned: + final_o_dst.store( + final_regs.data_ptr().load(count=8, alignment=4), + alignment=32, + ) else: - final_regs[reg_idx] = _pack_float2_to_fp16(pair[0], pair[1]) - final_o_dst.store( - final_regs.data_ptr().load(count=4, alignment=4), + self._store_16bit_output_chunks(final_o_dst, final_regs, count) + else: + self._store_16bit_output_chunks(final_o_dst, final_regs, count) + + @cute.jit + def _store_16bit_output_chunks( + self, + final_o_dst, + final_regs: cutlass.Array, + count: Constexpr[int], + ) -> None: + """Store packed 16-bit output columns as 16-byte chunks of eight columns.""" + for chunk_idx in cutlass.range_constexpr(count // 8): + (final_o_dst + Int32(chunk_idx * 4)).store( + (final_regs.data_ptr() + Int32(chunk_idx * 4)).load( + count=4, alignment=4 + ), alignment=16, ) + @cute.jit + def _pack_fp8_output_quads( + self, + output_vals: cutlass.Array, + norm_scale: Float32, + first_col: Constexpr[int], + ) -> cutlass.Array: + """Scale eight output columns from ``first_col`` into two FP8 registers.""" + final_pairs = cutlass.Array(Float32, 8, space=cutlass.AddressSpace.rmem) + for pair_idx in cutlass.range_constexpr(4): + val_base = pair_idx * 2 + pair = fmul2( + (norm_scale, norm_scale), + ( + output_vals[first_col + val_base], + output_vals[first_col + val_base + 1], + ), + ) + final_pairs[val_base] = pair[0] + final_pairs[val_base + 1] = pair[1] + fp8_regs = cutlass.Array(Int32, 2, space=cutlass.AddressSpace.rmem) + fp8_regs[0] = _pack_float4_to_fp8_e4m3( + final_pairs[0], final_pairs[1], final_pairs[2], final_pairs[3] + ) + fp8_regs[1] = _pack_float4_to_fp8_e4m3( + final_pairs[4], final_pairs[5], final_pairs[6], final_pairs[7] + ) + return fp8_regs + + @cute.jit + def _pack_final_o_regs( + self, + output_vals: cutlass.Array, + norm_scale: Float32, + count: Constexpr[int], + ) -> cutlass.Array: + """Scale ``count`` output columns and pack them as 16-bit pairs.""" + cfg = self.cfg + final_regs = cutlass.Array(Int32, count // 2, space=cutlass.AddressSpace.rmem) + for reg_idx in cutlass.range_constexpr(count // 2): + pair = fmul2( + (norm_scale, norm_scale), + ( + output_vals[reg_idx * 2], + output_vals[reg_idx * 2 + 1], + ), + ) + if cutlass.const_expr(cfg.use_bf16_output): + final_regs[reg_idx] = _pack_float2_to_bf16(pair[0], pair[1]) + else: + final_regs[reg_idx] = _pack_float2_to_fp16(pair[0], pair[1]) + return final_regs + @cute.jit def _store_softmax_normalized_o_vec8( self, @@ -953,7 +1032,13 @@ def _store_softmax_normalized_o_vec8( mem_space=1, dtype=Int32, ) - self._store_final_o_vec8(final_o_dst, output_vals, norm_scale) + self._store_final_o_columns( + final_o_dst, + output_vals, + norm_scale, + count=8, + sector_aligned=cutlass.Boolean(False), + ) @cute.jit def _softmax_output_row_state( @@ -2274,24 +2359,6 @@ def _publish_and_reduce_cluster_swaps_partials( full_prefix=False, ) - @cute.jit - def _kv_tile_256_exchange_for_stage( - self, - stage_info: StageInfo, - scratch_stage: Int32 | None, - ) -> cutlass.Array: - """Return the fixed exchange or its dynamically selected KV stage.""" - if cutlass.const_expr(scratch_stage is None): - return self._kv_tile_256_exchange - return cutlass.Array( - stage_info.context.smem_base.data_ptr() - + self._kv_tile_256_exchange_alloc.offset - + scratch_stage * Int32(self.cfg.smem_kv_tile_bytes), - dtype=Float32, - shape=(self._kv_tile_256_exchange_entries(),), - addrspace=3, - ) - @cute.jit def _kv_tile_256_temporal_fragment( self, @@ -2329,6 +2396,99 @@ def _kv_tile_256_temporal_fragment( ) return cutlass.Vector.from_elements(combined, Float32) + @cute.jit + def _store_kv_tile_256_direct_fragment( + self, + own_vals: cutlass.Array, + peer_vals, + *, + fragment_col: Constexpr[int], + dst_row_base: Int64, + norm_scale: Float32, + valid_output_row: cutlass.Boolean, + o_is_32b_aligned: cutlass.Boolean, + ) -> None: + """Merge one D32 fragment with its peer half and write the final output. + + Sixteen columns go out per store so each lane writes one full 32-byte + sector. Adjacent lanes own adjacent rows, so 16-byte stores would leave + every sector half-written twice. + """ + cfg = self.cfg + for vector_pair in cutlass.range_constexpr(2): + pair_col = vector_pair * 16 + merged_vals = self._merge_kv_tile_256_peer_fragment( + own_vals, + peer_vals, + pair_col, + 16, + ) + if valid_output_row: + output_col = fragment_col + pair_col + dst_offset = dst_row_base + Int32(output_col * cfg.o_dtype_bytes) + final_o_dst = cutlass.inttoptr( + self.o_ptr.toint() + cutlass.Int64(dst_offset), + mem_space=1, + dtype=Int32, + ) + self._store_final_o_columns( + final_o_dst, + merged_vals, + norm_scale, + count=16, + sector_aligned=o_is_32b_aligned, + ) + + @cute.jit + def _store_kv_tile_256_partial_fragment( + self, + own_vals: cutlass.Array, + peer_vals, + *, + fragment_col: Constexpr[int], + partial_row_base: Int64, + partial_scale: Float32, + valid_output_row: cutlass.Boolean, + ) -> None: + """Merge one D32 fragment with its peer half and write the split-KV partial.""" + cfg = self.cfg + partial_o_uses_bf16 = ( + cfg.use_bf16_separate_partial_o + if cfg.use_separate_reduction_kernel + else cfg.use_bf16_output + ) + for vector_idx in cutlass.range_constexpr(4): + vector_col = vector_idx * 8 + output_vals = self._merge_kv_tile_256_peer_fragment( + own_vals, + peer_vals, + vector_col, + 8, + ) + if valid_output_row: + output_col = fragment_col + vector_col + scaled_values: tuple = () + for elem in cutlass.range_constexpr(0, 8, 2): + scaled_values += fmul2( + (partial_scale, partial_scale), + (output_vals[elem], output_vals[elem + 1]), + ) + scaled_vector = cutlass.Vector.from_elements(scaled_values, Float32) + if cutlass.const_expr(partial_o_uses_bf16): + packed = scaled_vector.to(cutlass.BFloat16).bitcast(Int32) + else: + packed = scaled_vector.to(cutlass.Float16).bitcast(Int32) + # Split-KV partials are 16-bit, so the column offset follows + # the partial element width + partial_o_dst = cutlass.inttoptr( + self.partial_o_ptr.toint() + + partial_row_base + + Int64(output_col * 2), + mem_space=1, + dtype=Int32, + ) + partial_o_dst.store(packed, alignment=16) + @cute.jit def _kv_tile_256_merge_spatial_output( self, @@ -2357,18 +2517,13 @@ def _kv_tile_256_merge_spatial_output( serializing the complete D128 upper and lower halves. """ cfg = self.cfg - partial_o_uses_bf16 = ( - cfg.use_bf16_separate_partial_o - if cfg.use_separate_reduction_kernel - else cfg.use_bf16_output - ) output_exchange_base = Int32( _KV_TILE_256_CORRECTION_THREADS * _KV_TILE_256_STATS_PER_THREAD ) output_lane = exchange_idx < Int32(_KV_TILE_256_LOGICAL_OUTPUT_ROWS) exchange_row_idx = exchange_idx & Int32(_KV_TILE_256_LOGICAL_OUTPUT_ROWS - 1) output_exchange_row_base = output_exchange_base + exchange_row_idx * Int32( - _KV_TILE_256_EXCHANGE_ROW_STRIDE + _KV_TILE_256_EXCHANGE_FRAGMENT_STRIDE ) logical_output_row_idx = q_row_offset + exchange_row_idx valid_output_row = cutlass.Boolean(False) @@ -2394,6 +2549,9 @@ def _kv_tile_256_merge_spatial_output( else: dst_row_base = Int64(0) norm_scale = Float32(1.0) + # Row strides are multiples of 32 bytes for D128, so sector-wide + # stores are legal exactly when the output base pointer is. + o_is_32b_aligned = (self.o_ptr.toint() & Int64(31)) == Int64(0) if output_lane: valid_output_row = _q_row_is_valid_for_seq( cfg, @@ -2419,10 +2577,17 @@ def _kv_tile_256_merge_spatial_output( weight00=weight00, weight10=weight10, ) + if cutlass.const_expr(fragment != 0): + # The single fragment buffer is reused: lower lanes must have + # consumed the previous peer fragment before it is overwritten. + prims.barrier_cta_sync( + self.store_barrier_id, + thread_count=cfg.correction_barrier_threads, + ) if exchange_idx >= Int32(_KV_TILE_256_LOGICAL_OUTPUT_ROWS): - ( - exchange.data_ptr() + output_exchange_row_base + Int32(fragment_col) - ).store(own_vals, alignment=16) + (exchange.data_ptr() + output_exchange_row_base).store( + own_vals, alignment=16 + ) # Lower lanes keep ``own_vals`` live across the barrier. Once every # lane arrives, upper lanes may prepare the next fragment while @@ -2433,69 +2598,28 @@ def _kv_tile_256_merge_spatial_output( ) if output_lane: - peer_vals = ( - exchange.data_ptr() + output_exchange_row_base + Int32(fragment_col) - ).load(count=32, alignment=16) - for vector_idx in cutlass.range_constexpr(4): - vector_col = vector_idx * 8 - output_vals = cutlass.Array( - Float32, - 8, - space=cutlass.AddressSpace.rmem, + peer_vals = (exchange.data_ptr() + output_exchange_row_base).load( + count=32, alignment=16 + ) + if cutlass.const_expr(not cfg.use_split_kv): + self._store_kv_tile_256_direct_fragment( + own_vals, + peer_vals, + fragment_col=fragment_col, + dst_row_base=dst_row_base, + norm_scale=norm_scale, + valid_output_row=valid_output_row, + o_is_32b_aligned=o_is_32b_aligned, + ) + else: + self._store_kv_tile_256_partial_fragment( + own_vals, + peer_vals, + fragment_col=fragment_col, + partial_row_base=partial_row_base, + partial_scale=partial_scale, + valid_output_row=valid_output_row, ) - for elem in cutlass.range_constexpr(0, 8, 2): - value_idx = vector_col + elem - merged = fadd2( - (own_vals[value_idx], own_vals[value_idx + 1]), - ( - Float32(peer_vals[value_idx]), - Float32(peer_vals[value_idx + 1]), - ), - ) - output_vals[elem] = merged[0] - output_vals[elem + 1] = merged[1] - if valid_output_row: - output_col = fragment_col + vector_col - if cutlass.const_expr(cfg.use_split_kv): - scaled_values: tuple = () - for elem in cutlass.range_constexpr(0, 8, 2): - scaled_values += fmul2( - (partial_scale, partial_scale), - (output_vals[elem], output_vals[elem + 1]), - ) - scaled_vector = cutlass.Vector.from_elements( - scaled_values, Float32 - ) - if cutlass.const_expr(partial_o_uses_bf16): - packed = scaled_vector.to(cutlass.BFloat16).bitcast( - Int32 - ) - else: - packed = scaled_vector.to(cutlass.Float16).bitcast( - Int32 - ) - partial_o_dst = cutlass.inttoptr( - self.partial_o_ptr.toint() - + partial_row_base - + Int64(output_col * cfg.o_dtype_bytes), - mem_space=1, - dtype=Int32, - ) - partial_o_dst.store(packed, alignment=16) - else: - dst_offset = dst_row_base + Int32( - output_col * cfg.o_dtype_bytes - ) - final_o_dst = cutlass.inttoptr( - self.o_ptr.toint() + cutlass.Int64(dst_offset), - mem_space=1, - dtype=Int32, - ) - self._store_final_o_vec8( - final_o_dst, - output_vals, - norm_scale, - ) if cutlass.const_expr(cfg.use_split_kv): if valid_output_row: @@ -2538,7 +2662,6 @@ def _kv_tile_256_tail_epilogue( self, stage_info: StageInfo, *, - scratch_stage: Int32 | None, tail_o_stage_idx_0: Int32, tail_o_stage_idx_1: Int32, inst0_new_max_arr: cutlass.Array, @@ -2553,15 +2676,13 @@ def _kv_tile_256_tail_epilogue( The standard decode schedule still owns the two temporal instances. KV256 adds one physical spatial split per instance. Correction exchanges - their stats, stages one spatial half in SMEM after the shared KV ring is - dead, then publishes the ordinary logical Q64xD128 output. + their stats, stages one spatial half through its dedicated SMEM + exchange one D32 fragment at a time, then publishes the ordinary + logical Q64xD128 output. """ cfg = self.cfg assert cfg.headdim == 128 - exchange = self._kv_tile_256_exchange_for_stage( - stage_info, - scratch_stage, - ) + exchange = self._kv_tile_256_exchange exchange_idx = warp_grp_thread_idx peer_idx = exchange_idx ^ Int32(_KV_TILE_256_LOGICAL_OUTPUT_ROWS) @@ -4371,7 +4492,6 @@ def _correction_tail_epilogue_impl( self, stage_info: StageInfo, *, - scratch_stage: Int32 | None, o_stage_idx: Int32, tail_o_stage_idx_0: Int32, tail_o_stage_idx_1: Int32, @@ -4417,7 +4537,6 @@ def _correction_tail_epilogue_impl( if cutlass.const_expr(cfg.tile_size_kv == 256): self._kv_tile_256_tail_epilogue( stage_info, - scratch_stage=scratch_stage, tail_o_stage_idx_0=tail_o_stage_idx_0, tail_o_stage_idx_1=tail_o_stage_idx_1, inst0_new_max_arr=inst0_new_max_arr, @@ -4493,9 +4612,6 @@ def _correction_tail_epilogue_impl( ) return - # Task Scheduling routes every non-constexpr work argument as a required - # data-flow token. Keep separate fixed/rotating entry points so only the - # latter consumes ``scratch_stage``; both still share the implementation. @producer_work @cute.jit def correction_tail_epilogue( @@ -4512,42 +4628,9 @@ def correction_tail_epilogue( inst1_new_max_arr: cutlass.Array, inst1_sum_arr: cutlass.Array, ) -> None: - """Run the ordinary fixed-exchange tail epilogue.""" - self._correction_tail_epilogue_impl( - stage_info, - scratch_stage=None, - o_stage_idx=o_stage_idx, - tail_o_stage_idx_0=tail_o_stage_idx_0, - tail_o_stage_idx_1=tail_o_stage_idx_1, - old_max_arr=old_max_arr, - new_max_arr=new_max_arr, - inst0_new_max_arr=inst0_new_max_arr, - inst0_sum_arr=inst0_sum_arr, - inst1_new_max_arr=inst1_new_max_arr, - inst1_sum_arr=inst1_sum_arr, - ) - - @producer_work - @cute.jit - def correction_tail_epilogue_rotating_exchange( - self, - stage_info: StageInfo, - *, - scratch_stage: Int32, - o_stage_idx: Int32, - tail_o_stage_idx_0: Int32, - tail_o_stage_idx_1: Int32, - old_max_arr: cutlass.Array, - new_max_arr: cutlass.Array, - inst0_new_max_arr: cutlass.Array, - inst0_sum_arr: cutlass.Array, - inst1_new_max_arr: cutlass.Array, - inst1_sum_arr: cutlass.Array, - ) -> None: - """Run persistent direct output in the stage named by its credit.""" + """Normalize the final O stages and publish the output tile.""" self._correction_tail_epilogue_impl( stage_info, - scratch_stage=scratch_stage, o_stage_idx=o_stage_idx, tail_o_stage_idx_0=tail_o_stage_idx_0, tail_o_stage_idx_1=tail_o_stage_idx_1, diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_o.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_o.py index fcf2cbe71730..fa8928695a98 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_o.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_o.py @@ -59,7 +59,7 @@ def _pv_mma_operand_contract_for_config( cfg.headdim if cfg.head_dim_per_stage_kv == 0 else cfg.head_dim_kv_stage ) if cfg.use_keeps_mma_ab: - if cfg.tile_size_kv == 256: + if cfg.uses_ws_2x2_datapath: # The WS 2x2 PV instruction exposes two spatial D128 partials as # one physical KV256 operation. Correction merges those spatial # halves after the two temporal decode streams are complete. @@ -196,7 +196,7 @@ def vp_mma_loop_fragment( p_tmem_addr: Int32, fragment_idx: Constexpr[int], ) -> None: - """Issue one K32 fragment of a KV256 loop PV tile.""" + """Issue one K32 fragment of a streamed loop PV tile.""" self._vp_mma_fragment( stage_info, v_desc=v_desc, @@ -215,7 +215,7 @@ def vp_mma_tail_fragment( p_tmem_addr: Int32, fragment_idx: Constexpr[int], ) -> None: - """Issue one K32 fragment of the final KV256 PV tile.""" + """Issue one K32 fragment of the final streamed PV tile.""" self._vp_mma_fragment( stage_info, v_desc=v_desc, @@ -234,14 +234,16 @@ def _vp_mma_fragment( fragment_idx: Constexpr[int], initial_scale_d, ) -> None: - """Issue the two WS MMA steps covered by one KV256 P fragment. + """Issue the two MMA K-steps covered by one streamed P fragment. ``p_tmem_addr`` is already the base of the fragment selected by ``wait_p_fragment``. Only the two local K-step offsets are added here; - ``fragment_idx`` must not be applied to the TMEM address again. + ``fragment_idx`` must not be applied to the TMEM address again. KV256 + issues the WS 2x2 instruction over its two spatial halves; KV128 issues + the plain M=128 instruction and advances V by one K16 slice per step. """ cfg = self.cfg - assert cfg.tile_size_kv == 256 and cfg.uses_two_inst_tmem_p + assert cfg.streams_tmem_p_fragments v_desc = _freeze_smem_descriptor(v_desc) task_cache = _decode_gen_task_cache(stage_info) @@ -268,17 +270,32 @@ def _vp_mma_fragment( p_operand = prims.make_tmem_ptr( p_tmem_addr + Int32(local_k_step * 8), Int32 ) - iter_v_desc = v_desc + Int32( - (k_step // 4) * cfg.headdim * 16 + (k_step % 4) * 128 - ) - tcgen05_mma_ws( - _mma_kind_for_qkv(cfg), - tmem_col, - p_operand, - iter_v_desc, - idesc, - initial_scale_d or fragment_idx != 0 or local_k_step != 0, - ) + scale_d = initial_scale_d or fragment_idx != 0 or local_k_step != 0 + if cutlass.const_expr(cfg.uses_ws_2x2_datapath): + # V holds four K64 atoms; jump between atoms every four + # K16 steps. + iter_v_desc = v_desc + Int32( + (k_step // 4) * cfg.headdim * 16 + (k_step % 4) * 128 + ) + tcgen05_mma_ws( + _mma_kind_for_qkv(cfg), + tmem_col, + p_operand, + iter_v_desc, + idesc, + scale_d, + ) + else: + iter_v_desc = v_desc + Int32(k_step * 128) + prims.tcgen05_mma( + _mma_kind_for_qkv(cfg), + prims.CTAGroup.CTA_1, + tmem_col, + p_operand, + iter_v_desc, + idesc, + scale_d, + ) @cute.jit def _vp_mma( diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_s.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_s.py index 703134a5b61d..b010dfdb9d21 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_s.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_resources/tmem_s.py @@ -41,10 +41,9 @@ producer_work, ) -from ...._block_sparse.common import _MAX_KV_ATOM_SIZE from ...._block_sparse.prepared import _PREPARED_ROUTE_IS_FULL_FLAG from ..fmha_decode_config import CAUSAL, FmhaDecodeConfig -from ..fmha_decode_constants import KV_TILE_256_RESCALE_THRESHOLD_LOG2 +from ..fmha_decode_constants import SOFTMAX_RESCALE_THRESHOLD_LOG2 from ...tcgen05_compat import tcgen05_mma_ws from ...placeholder_helpers import ( _placeholder_local_array, @@ -76,13 +75,14 @@ _mma_kind_for_qkv, _neg_max_f32, _softmax_scale_pair_width, + _swaps_routed_coordinate, _q_row_is_valid_for_seq, _q_row_token_and_local_head, _q_group_token_base, _softmax_tile_idx, ) from .smem_block_sparse_metadata import ( - _SOFTMAX_TOKEN_MASK_IS_FULL_FLAG, + _SOFTMAX_ROUTE_IS_PROXY_FLAG, _swaps_forwards_packed_route_full, ) from .helpers_kv_tile_idx import ( @@ -103,20 +103,13 @@ _wspro_reduce_max4, ) -# A block-sparse route often changes the exact row maximum without changing it -# enough to justify rescaling the live O tile. Keeping the prior anchor within -# this bound makes the correction scale exactly one and bounds FP16/BF16 P by -# 2**8. As in the FlashInfer/TRT-LLM policy, this assumes normal model logits -# rather than adversarial values outside the qualified probability bound. -_BLOCK_SPARSE_RESCALE_THRESHOLD_LOG2 = 8.0 - def _swaps_uses_origin0_k32_full_guard(cfg: FmhaDecodeConfig) -> bool: """Whether one staged origin can prove this warp's K32 slice valid.""" return ( cfg.kv_block_size >= 32 - and not cfg.use_kv_valid_bits + and not cfg.uses_prepared_score_keep_words and not cfg.uses_uniform_causal_mask and not cfg.uses_per_row_causal_mask ) @@ -126,7 +119,7 @@ def _swaps_token_word_covers_kv_tail(cfg: FmhaDecodeConfig) -> bool: """Whether SWAP's prepared token word covers the logical KV tail.""" return ( - cfg.use_kv_valid_bits + cfg.uses_prepared_score_keep_words and not cfg.uses_uniform_causal_mask and not cfg.uses_per_row_causal_mask ) @@ -145,42 +138,34 @@ def _swaps_uses_token_only_score_validity(cfg: FmhaDecodeConfig) -> bool: @cute.jit -def _can_skip_sparse_keeps_structural_mask( - q_row_is_valid: Boolean, - origin0: Int32, - origin1: Int32, - valid0: Int32, - valid1: Int32, - seq_len_kv: Int32, - causal_end: Int32, +def _dense_fragment_keep_word( + rows_are_active: Boolean, + visible_start: Int32, + visible_end: Int32, *, - apply_causal_mask: cutlass.Constexpr[bool], -) -> Boolean: - """Return whether one Keeps row needs no Q/tail/causal predicate. + fragment_regs: cutlass.Constexpr[int], +) -> Uint32: + """Return the keep word of one dense K32 fragment. - Token-bit masking is independent. Comparing against the last complete - KV64 origin avoids overflowing an origin near the Int32 upper bound. + ``visible_start`` and ``visible_end`` are the visible token range relative + to the fragment's first column. Columns outside ``[start, end)`` are + masked; an inactive tile or Q row masks the whole fragment. """ - - fragment_size = Int32(_MAX_KV_ATOM_SIZE) - last_complete_origin = seq_len_kv - fragment_size - can_skip = Boolean( - q_row_is_valid - and valid0 != Int32(0) - and valid1 != Int32(0) - and origin0 <= last_complete_origin - and origin1 <= last_complete_origin - ) - if cutlass.const_expr(apply_causal_mask): - last_causal_origin = causal_end - fragment_size - can_skip = Boolean( - can_skip and origin0 <= last_causal_origin and origin1 <= last_causal_origin - ) - return can_skip + keep_word = Uint32(0) + if rows_are_active: + first_kept = cute.math.max(visible_start, Int32(0)) + end_kept = cute.math.min(visible_end, Int32(fragment_regs)) + if first_kept < end_kept: + # Both shift amounts stay strictly below the register width: + # 1 <= end_kept <= fragment_regs and 0 <= first_kept < end_kept. + keep_word = (Uint32(0xFFFFFFFF) >> (Int32(fragment_regs) - end_kept)) & ( + Uint32(0xFFFFFFFF) << first_kept + ) + return keep_word @cute.jit -def _sparse_k32_effective_keep_word( +def _sparse_effective_keep_word( q_row_is_valid: Boolean, fragment_origin: Int32, fragment_valid: Int32, @@ -667,7 +652,7 @@ def _qk_mma( a_desc, b_desc = q_desc, k_desc else: a_desc, b_desc = k_desc, q_desc - if cutlass.const_expr(cfg.tile_size_kv == 256): + if cutlass.const_expr(cfg.uses_ws_2x2_datapath): tcgen05_mma_ws( _mma_kind_for_qkv(cfg), tmem_col, @@ -842,6 +827,7 @@ def _resolve_keeps_tile_context(self, stage_info: StageInfo): is_valid_effective_tile, is_masked_final_wave, tile_is_unmasked, + tile_has_valid_scores, ) @cute.jit @@ -896,49 +882,32 @@ def _publish_keeps_softmax_state( ) -> None: """Publish a masked Keeps row and its updated softmax anchor.""" - new_anchor = cute.math.max(old_max, tile_max, ftz=True) - if cutlass.const_expr( - self.cfg.use_block_sparse and _BLOCK_SPARSE_RESCALE_THRESHOLD_LOG2 > 0.0 - ): - # Online softmax only requires a common finite reference for P, - # sum, and O; it does not require the exact row maximum. Defer a - # small anchor increase so correction can skip a TMEM O rescale. - rescale_log2 = (old_max - new_anchor) * self.scale_softmax_log2 - if (old_max != _neg_max_f32()) and ( - rescale_log2 >= Float32(-_BLOCK_SPARSE_RESCALE_THRESHOLD_LOG2) - ): - new_anchor = old_max old_max_arr[0] = old_max sum_arr[0] = running_sum - new_max_arr[0] = new_anchor + new_max_arr[0] = self._softmax_anchor(old_max, tile_max) for reg_idx in cutlass.range_constexpr(self.cfg.num_s_regs_per_thread): s_arr[reg_idx] = s_vals[reg_idx] @cute.jit - def _mask_and_store_sparse_keeps_atom( - self, - s_vals: cutlass.Array, - loaded: cutlass.Vector, - token_word: Uint32, - *, - atom_col: Constexpr[int], - token_mask_is_required: cutlass.Boolean, - ) -> None: - """Store one 32-score atom, applying its token word when required.""" - - if token_mask_is_required: - for atom_reg_idx in cutlass.range_constexpr(32): - score_idx = atom_col + atom_reg_idx - s_vals[score_idx] = loaded[atom_reg_idx] - token_bit_is_valid = ( - (token_word >> Int32(atom_reg_idx)) & Uint32(1) - ) != Uint32(0) - if not token_bit_is_valid: - s_vals[score_idx] = _neg_max_f32() - else: - for atom_reg_idx in cutlass.range_constexpr(32): - score_idx = atom_col + atom_reg_idx - s_vals[score_idx] = loaded[atom_reg_idx] + def _softmax_anchor(self, old_max: Float32, tile_max: Float32) -> Float32: + """Return the exponent reference max for the tile's P pass. + + Online softmax only requires a common finite reference for P, the + running sum, and O; it does not require the exact row maximum. + Profiles that defer anchor updates keep the previous reference while + the tile raises it by less than ``SOFTMAX_RESCALE_THRESHOLD_LOG2`` + log2 units, so correction can skip the in-place TMEM O rescale. The + 16-bit P path represents the bounded values above one, and the + numerator and denominator stay in the same scale frame. Larger jumps + still rebase to keep P comfortably in range. + """ + new_max = cute.math.max(old_max, tile_max, ftz=True) + if cutlass.const_expr(self.cfg.defers_softmax_anchor_updates): + if old_max != _neg_max_f32(): + max_delta_log2 = self.scale_softmax_log2 * (old_max - new_max) + if max_delta_log2 >= Float32(-SOFTMAX_RESCALE_THRESHOLD_LOG2): + new_max = old_max + return new_max @cute.jit def _load_keeps_fragment_impl( @@ -954,9 +923,8 @@ def _load_keeps_fragment_impl( is_masked_final_wave: cutlass.Boolean, *, apply_boundary_mask: Constexpr[bool], - fragment_idx: Constexpr[int] = 0, ) -> None: - """Load one Keeps score fragment with a compile-time mask policy. + """Load one complete-row Keeps score tile with a compile-time mask policy. The caller chooses the masked/unmasked path before TMEM load. Keeping the score fragment out of the branch condition avoids carrying 64/128 live @@ -966,8 +934,7 @@ def _load_keeps_fragment_impl( """ cfg = self.cfg task_cache = _decode_gen_task_cache(stage_info) - num_s_regs = cfg.softmax_score_fragment_regs - fragment_reg_base = fragment_idx * num_s_regs + num_s_regs = cfg.num_s_regs_per_thread base_addr = ( task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + Int32(self._alloc.offset) @@ -977,9 +944,7 @@ def _load_keeps_fragment_impl( atom_col = load_atom_idx * 32 loaded = _keeps_tcgen05_ld( cfg, - prims.make_tmem_ptr( - base_addr + Int32(fragment_reg_base + atom_col), Float32 - ), + prims.make_tmem_ptr(base_addr + Int32(atom_col), Float32), num=32, offset=cfg.tile_size_kv // 2, ) @@ -1023,7 +988,7 @@ def _load_keeps_fragment_impl( token_idx = tile_offset_k + _keeps_score_col( cfg, warp_grp_thread_idx, - fragment_reg_base + reg_idx, + reg_idx, col_base, ) if token_idx >= element_mask_end_idx: @@ -1051,7 +1016,7 @@ def _load_keeps_fragment_impl( score_col = _keeps_score_col( cfg, warp_grp_thread_idx, - fragment_reg_base + reg_idx, + reg_idx, col_base, ) if cutlass.const_expr(cfg.use_sliding_window_causal): @@ -1084,8 +1049,6 @@ def _load_keeps_fragment( is_valid_effective_tile: cutlass.Boolean, is_masked_final_wave: cutlass.Boolean, tile_is_unmasked: cutlass.Boolean, - *, - fragment_idx: Constexpr[int] = 0, ) -> None: """Select the masked or unmasked fragment loader before LDTM. @@ -1106,7 +1069,6 @@ def _load_keeps_fragment( is_valid_effective_tile, is_masked_final_wave, apply_boundary_mask=False, - fragment_idx=fragment_idx, ) else: self._load_keeps_fragment_impl( @@ -1120,275 +1082,8 @@ def _load_keeps_fragment( is_valid_effective_tile, is_masked_final_wave, apply_boundary_mask=True, - fragment_idx=fragment_idx, ) - @cute.jit - def _reduce_keeps_fragment_max(self, s_vals: cutlass.Array) -> Float32: - """Reduce the row maximum of a previously loaded Keeps fragment.""" - cfg = self.cfg - num_s_regs = cfg.softmax_score_fragment_regs - - max_chains = cutlass.Array(Float32, 4, space=cutlass.AddressSpace.rmem) - for chain_idx in cutlass.range_constexpr(4): - max_chains[chain_idx] = _neg_max_f32() - for reg_base in cutlass.range_constexpr(0, num_s_regs, 4): - for chain_idx in cutlass.range_constexpr(4): - max_chains[chain_idx] = cute.math.max( - max_chains[chain_idx], - s_vals[reg_base + chain_idx], - ftz=True, - ) - tile_max = cute.math.max( - cute.math.max(max_chains[0], max_chains[1], ftz=True), - cute.math.max(max_chains[2], max_chains[3], ftz=True), - ftz=True, - ) - if cutlass.const_expr(cfg.tile_size_q == 64 and cfg.tile_size_kv != 256): - return cute.math.max( - tile_max, - Float32( - prims.shfl_sync( - thread_mask=0xFFFFFFFF, - val=tile_max, - offset=16, - mask_and_clamp=0x1F, - kind=prims.Shfl.BFLY, - ) - ), - ftz=True, - ) - return tile_max - - @cute.jit - def _decode_sparse_mask_metadata( - self, - routed_origin0: Int32, - routed_origin1: Int32, - routed_route_flags: Int32, - routed_token_word0: Uint32, - routed_token_word1: Uint32, - routed_token_word2: Uint32, - routed_token_word3: Uint32, - ) -> tuple[Int32, Int32, Int32, Int32, cutlass.Array, cutlass.Boolean]: - """Decode one prepared, register-routed mask payload.""" - - origin0 = Int32(routed_origin0) - origin1 = Int32(routed_origin1) - route_flags = Int32(routed_route_flags) - valid0 = route_flags & Int32(1) - valid1 = (route_flags >> Int32(1)) & Int32(1) - route_token_mask_is_full = cutlass.Boolean(False) - if cutlass.const_expr(self.cfg.use_kv_valid_bits): - route_token_mask_is_full = cutlass.Boolean( - (route_flags & Int32(_SOFTMAX_TOKEN_MASK_IS_FULL_FLAG)) != Int32(0) - ) - - num_local_words = 4 if self.cfg.tile_size_q == 128 else 2 - local_token_words = cutlass.Array( - Uint32, - num_local_words, - space=cutlass.AddressSpace.rmem, - ) - for word_idx in cutlass.range_constexpr(num_local_words): - local_token_words[word_idx] = Uint32(0xFFFFFFFF) - if cutlass.const_expr(self.cfg.use_kv_valid_bits): - if not route_token_mask_is_full: - if cutlass.const_expr(self.cfg.tile_size_q == 128): - local_token_words[0] = Uint32(routed_token_word0) - local_token_words[1] = Uint32(routed_token_word1) - local_token_words[2] = Uint32(routed_token_word2) - local_token_words[3] = Uint32(routed_token_word3) - else: - local_word0 = Uint32(routed_token_word0) - local_word1 = Uint32(routed_token_word1) - local_token_words[0] = local_word0 - local_token_words[1] = local_word1 - return ( - origin0, - origin1, - valid0, - valid1, - local_token_words, - route_token_mask_is_full, - ) - - @cute.jit - def _compute_softmax_loop_sparse_keeps( - self, - stage_info: StageInfo, - *, - old_max_arr: cutlass.Array, - sum_arr: cutlass.Array, - new_max_arr: cutlass.Array, - s_arr: cutlass.Array, - routed_origin0: Int32, - routed_origin1: Int32, - routed_route_flags: Int32, - routed_token_word0: Uint32, - routed_token_word1: Uint32, - routed_token_word2: Uint32, - routed_token_word3: Uint32, - ) -> tuple[object, object, object, object]: - """Load Keeps scores and mask them in logical KV coordinates.""" - cfg = self.cfg - num_s_regs = cfg.num_s_regs_per_thread - old_max = new_max_arr[0] - running_sum = sum_arr[0] - s_vals = cutlass.Array(Float32, num_s_regs, space=cutlass.AddressSpace.rmem) - task_cache = _decode_gen_task_cache(stage_info) - seq_len_kv = _load_runtime_seq_len_kv( - self.seqlens_kv, - self.max_seq_len_kv, - stage_info, - Int32(0), - Int32(0), - ) - warp_grp_thread_idx = Int32(task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX]) - lane_idx = Int32(task_cache[_TASK_CACHE_LANE_IDX]) - tile_row_idx = _keeps_row_idx(cfg, warp_grp_thread_idx) - col_base = _keeps_col_base(cfg, lane_idx, num_s_regs) - ( - origin0, - origin1, - valid0, - valid1, - local_token_words, - route_token_mask_is_full, - ) = self._decode_sparse_mask_metadata( - routed_origin0=routed_origin0, - routed_origin1=routed_origin1, - routed_route_flags=routed_route_flags, - routed_token_word0=routed_token_word0, - routed_token_word1=routed_token_word1, - routed_token_word2=routed_token_word2, - routed_token_word3=routed_token_word3, - ) - - base_addr = ( - task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] - + Int32(self._alloc.offset) - + self._softmax_loop_stage_slot_offset(stage_info) - ) - num_load_atoms = num_s_regs // 32 - if cutlass.const_expr(cfg.tile_size_q == 64 and cfg.use_kv_valid_bits): - token_mask_is_required = not route_token_mask_is_full - - # Keep each Q64 atom's load, wait, and mask together. A/B testing - # showed that hoisting both loads extends live fragment ranges and - # regresses the Q64 code generated by ptxas. - for load_atom_idx in cutlass.range_constexpr(2): - atom_col = load_atom_idx * 32 - loaded = _keeps_tcgen05_ld( - cfg, - prims.make_tmem_ptr(base_addr + Int32(atom_col), Float32), - num=32, - offset=cfg.tile_size_kv // 2, - ) - prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) - self._mask_and_store_sparse_keeps_atom( - s_vals, - loaded, - local_token_words[load_atom_idx], - atom_col=atom_col, - token_mask_is_required=token_mask_is_required, - ) - else: - for load_atom_idx in cutlass.range_constexpr(num_load_atoms): - atom_col = load_atom_idx * 32 - loaded = _keeps_tcgen05_ld( - cfg, - prims.make_tmem_ptr(base_addr + Int32(atom_col), Float32), - num=32, - offset=cfg.tile_size_kv // 2, - ) - prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) - for atom_reg_idx in cutlass.range_constexpr(32): - score_idx = atom_col + atom_reg_idx - s_vals[score_idx] = loaded[atom_reg_idx] - - logical_q_group_idx = _logical_q_group_idx(cfg, stage_info, self.q_group_idx) - q_token_idx, _ = _q_row_token_and_local_head( - cfg, - self.h_r, - logical_q_group_idx, - tile_row_idx, - ) - q_row_is_valid = _q_row_is_valid_for_seq( - cfg, - self.h_r, - logical_q_group_idx, - tile_row_idx, - self.seq_len_q, - ) - causal_end = seq_len_kv - self.seq_len_q + q_token_idx + Int32(1) - can_skip_structural_mask = _can_skip_sparse_keeps_structural_mask( - q_row_is_valid, - origin0, - origin1, - valid0, - valid1, - seq_len_kv, - causal_end, - apply_causal_mask=cfg.mask_type == CAUSAL, - ) - # This guard covers only route/Q/tail/causal structure. Q64 token - # holes were applied while materializing its two LDTM atoms; Q128 - # applies them in the post-pass below. - if not can_skip_structural_mask: - for reg_idx in cutlass.range_constexpr(num_s_regs): - fragment_offset = Int32(reg_idx) - logical_k = origin0 + fragment_offset - fragment_valid = valid0 - if cutlass.const_expr(cfg.tile_size_q == 128 and reg_idx >= 64): - fragment_offset = Int32(reg_idx - 64) - logical_k = origin1 + fragment_offset - fragment_valid = valid1 - elif cutlass.const_expr(cfg.tile_size_q == 64): - if col_base >= Int32(64): - logical_k = origin1 + fragment_offset - fragment_valid = valid1 - - score_is_valid = ( - q_row_is_valid - and fragment_valid != Int32(0) - and logical_k < seq_len_kv - ) - if cutlass.const_expr(cfg.mask_type == CAUSAL): - score_is_valid = score_is_valid and logical_k < causal_end - if not score_is_valid: - s_vals[reg_idx] = _neg_max_f32() - - # Q128 deliberately keeps all four LDTM atoms adjacent: unlike Q64, - # interleaving each load with mask control flow regresses its codegen. - # The post-pass follows structural masking; the producer's runtime - # route flag skips it only when all four current token words are full. - if cutlass.const_expr(cfg.tile_size_q == 128 and cfg.use_kv_valid_bits): - token_mask_is_required = not route_token_mask_is_full - if token_mask_is_required: - for word_idx in cutlass.range_constexpr(4): - token_word = local_token_words[word_idx] - for bit_idx in cutlass.range_constexpr(32): - reg_idx = word_idx * 32 + bit_idx - token_bit_is_valid = ( - (token_word >> Int32(bit_idx)) & Uint32(1) - ) != Uint32(0) - if not token_bit_is_valid: - s_vals[reg_idx] = _neg_max_f32() - - tile_max = self._reduce_keeps_row_max(s_vals) - self._publish_keeps_softmax_state( - s_vals, - tile_max, - old_max, - running_sum, - old_max_arr, - sum_arr, - new_max_arr, - s_arr, - ) - return old_max_arr, sum_arr, new_max_arr, s_arr - @cute.jit def _compute_softmax_loop_keeps( self, @@ -1407,8 +1102,26 @@ def _compute_softmax_loop_keeps( reduction, whose 16x256b register mapping is unrelated to Keeps. """ cfg = self.cfg + if cutlass.const_expr(cfg.streams_tmem_p_fragments): + # Streamed profiles share the fragment max pass with block-sparse + # routes; dense tiles describe their visible range as keep words. + return self._compute_softmax_loop_keeps_fragments( + stage_info, + old_max_arr=old_max_arr, + sum_arr=sum_arr, + new_max_arr=new_max_arr, + s_arr=s_arr, + use_sparse=False, + sparse_origin0=Int32(0), + sparse_origin1=Int32(0), + sparse_route_flags=Int32(0), + sparse_token_word0=Uint32(0xFFFFFFFF), + sparse_token_word1=Uint32(0xFFFFFFFF), + sparse_token_word2=Uint32(0xFFFFFFFF), + sparse_token_word3=Uint32(0xFFFFFFFF), + ) task_cache = _decode_gen_task_cache(stage_info) - num_s_regs = cfg.softmax_score_fragment_regs + num_s_regs = cfg.num_s_regs_per_thread old_max = new_max_arr[0] running_sum = sum_arr[0] s_vals = cutlass.Array(Float32, num_s_regs, space=cutlass.AddressSpace.rmem) @@ -1434,50 +1147,9 @@ def _compute_softmax_loop_keeps( is_valid_effective_tile, is_masked_final_wave, tile_is_unmasked, + tile_has_valid_scores, ) = self._resolve_keeps_tile_context(stage_info) - if cutlass.const_expr(cfg.tile_size_kv == 256): - # KV256 owns four physical K32 fragments per lane. Reduce the max - # one fragment at a time so only one native LDTM atom is live; the - # P pass reloads the same fragments after the reference max is - # known. - tile_max = _neg_max_f32() - for fragment_idx in cutlass.range_constexpr( - cfg.num_softmax_score_fragments - ): - self._load_keeps_fragment( - stage_info, - s_vals, - tile_offset_k, - element_mask_end_idx, - window_start_idx, - seq_len_kv, - logical_q_group_idx, - is_valid_effective_tile, - is_masked_final_wave, - tile_is_unmasked, - fragment_idx=fragment_idx, - ) - fragment_max = self._reduce_keeps_fragment_max(s_vals) - tile_max = cute.math.max(tile_max, fragment_max, ftz=True) - - new_max = cute.math.max(old_max, tile_max, ftz=True) - if old_max != _neg_max_f32(): - # Keeping the previous reference max avoids an in-place O - # rescale when the new tile raises it only modestly. The - # 16-bit P path can represent the bounded values above one; the - # numerator and denominator remain in the same scale frame. - # Large jumps still rebase to keep P comfortably in range. - max_delta_log2 = self.scale_softmax_log2 * (old_max - new_max) - if max_delta_log2 >= Float32(-KV_TILE_256_RESCALE_THRESHOLD_LOG2): - new_max = old_max - old_max_arr[0] = old_max - sum_arr[0] = running_sum - new_max_arr[0] = new_max - for reg_idx in cutlass.range_constexpr(num_s_regs): - s_arr[reg_idx] = s_vals[reg_idx] - return old_max_arr, sum_arr, new_max_arr, s_arr - if cutlass.const_expr(use_preload_mask_split): # Select the complete unmasked/masked TMEM load+max path before any S # registers are materialized. The shared predicate covers the @@ -1495,7 +1167,7 @@ def _compute_softmax_loop_keeps( is_masked_final_wave, tile_is_unmasked, ) - tile_max = self._reduce_keeps_fragment_max(s_vals) + tile_max = self._reduce_keeps_row_max(s_vals) self._publish_keeps_softmax_state( s_vals, @@ -1509,12 +1181,7 @@ def _compute_softmax_loop_keeps( ) return old_max_arr, sum_arr, new_max_arr, s_arr - should_load_s = ( - is_valid_effective_tile - and (tile_offset_k < seq_len_kv) - and not is_masked_final_wave - ) - if should_load_s: + if tile_has_valid_scores: base_addr = ( task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] + Int32(self._alloc.offset) @@ -1600,73 +1267,6 @@ def _compute_softmax_loop_keeps( ) return old_max_arr, sum_arr, new_max_arr, s_arr - @consumer_work(returns=s_arr, work_attrs=WorkAttr.AUXILIARY) - @cute.jit - def load_softmax_p_fragment( - self, - stage_info: StageInfo, - *, - fragment_idx: Constexpr[int], - s_arr: cutlass.Array, - ) -> cutlass.Array: - """Reload and mask one KV256 K32 fragment for P materialization.""" - if cutlass.const_expr(self.cfg.use_block_sparse): - return self._load_block_sparse_softmax_p_fragment( - stage_info, - fragment_idx=fragment_idx, - s_arr=s_arr, - ) - ( - seq_len_kv, - logical_q_group_idx, - element_mask_end_idx, - tile_offset_k, - window_start_idx, - is_valid_effective_tile, - is_masked_final_wave, - tile_is_unmasked, - ) = self._resolve_keeps_tile_context(stage_info) - self._load_keeps_fragment( - stage_info, - s_arr, - tile_offset_k, - element_mask_end_idx, - window_start_idx, - seq_len_kv, - logical_q_group_idx, - is_valid_effective_tile, - is_masked_final_wave, - tile_is_unmasked, - fragment_idx=fragment_idx, - ) - return s_arr - - @cute.jit - def _sparse_swaps_logical_k( - self, - lane_k_offset: Int32, - sparse_origin0: Int32, - sparse_origin1: Int32, - sparse_origin2: Int32, - sparse_origin3: Int32, - *, - token_group_idx: Constexpr[int], - ) -> tuple[Int32, Int32]: - """Map one SWAP register group to its routed logical K position.""" - - atom_size = min(self.cfg.kv_block_size, 32) - groups_per_atom = atom_size // 8 - origin_idx = token_group_idx // groups_per_atom - atom_origin = sparse_origin0 - if cutlass.const_expr(origin_idx == 1): - atom_origin = sparse_origin1 - elif cutlass.const_expr(origin_idx == 2): - atom_origin = sparse_origin2 - elif cutlass.const_expr(origin_idx == 3): - atom_origin = sparse_origin3 - token_offset = (token_group_idx % groups_per_atom) * 8 - return atom_origin, atom_origin + Int32(token_offset) + lane_k_offset - @cute.jit def _compute_softmax_loop_swaps( self, @@ -1860,6 +1460,11 @@ def _compute_softmax_loop_swaps( s_vals[q_repeats * 4 + ld_base + 2] = loaded1[ld_base + 2] s_vals[q_repeats * 4 + ld_base + 3] = loaded1[ld_base + 3] + route_is_proxy = cutlass.Boolean(False) + if cutlass.const_expr(use_sparse and cfg.use_block_sparse_proxy_routes): + route_is_proxy = cutlass.Boolean( + (sparse_route_flags & Uint32(_SOFTMAX_ROUTE_IS_PROXY_FLAG)) != Uint32(0) + ) if cutlass.const_expr(use_sparse): # Route, KV-tail, uniform-causal, and token validity depend only on # K, so one predicate masks the adjacent pair of Q-row registers. @@ -1894,7 +1499,8 @@ def _compute_softmax_loop_swaps( lane_k_offset = Int32(task_cache[_TASK_CACHE_LANE_IDX]) >> Int32(2) token_word_covers_kv_tail = _swaps_token_word_covers_kv_tail(cfg) for token_group_idx in cutlass.range_constexpr(4): - atom_origin, logical_k = self._sparse_swaps_logical_k( + atom_origin, logical_k = _swaps_routed_coordinate( + cfg, lane_k_offset, sparse_origin0, sparse_origin1, @@ -1906,19 +1512,20 @@ def _compute_softmax_loop_swaps( # tail. Qualified profiles can therefore omit the local # atom-origin guard, independently of the K/V issuer warp. score_is_valid = cutlass.Boolean(True) - if cutlass.const_expr( - not _swaps_uses_token_only_score_validity(cfg) - ): - score_is_valid = cutlass.Boolean(atom_origin >= Int32(0)) - if cutlass.const_expr(not token_word_covers_kv_tail): - score_is_valid = cutlass.Boolean( - score_is_valid and logical_k < seq_len_kv - ) - if cutlass.const_expr(cfg.uses_uniform_causal_mask): - score_is_valid = cutlass.Boolean( - score_is_valid and logical_k < element_mask_end_idx - ) - if cutlass.const_expr(cfg.use_kv_valid_bits): + if not route_is_proxy: + if cutlass.const_expr( + not _swaps_uses_token_only_score_validity(cfg) + ): + score_is_valid = cutlass.Boolean(atom_origin >= Int32(0)) + if cutlass.const_expr(not token_word_covers_kv_tail): + score_is_valid = cutlass.Boolean( + score_is_valid and logical_k < seq_len_kv + ) + if cutlass.const_expr(cfg.uses_uniform_causal_mask): + score_is_valid = cutlass.Boolean( + score_is_valid and logical_k < element_mask_end_idx + ) + if cutlass.const_expr(cfg.uses_prepared_score_keep_words): token_bit_idx = Int32(token_group_idx * 8) + lane_k_offset token_is_valid = ( (sparse_token_word >> token_bit_idx) & Uint32(1) @@ -2061,7 +1668,8 @@ def _compute_softmax_loop_swaps( tile_offset_k + local_idx_k0 + Int32(token_group_idx * 8) ) if cutlass.const_expr(use_sparse): - _, token_idx = self._sparse_swaps_logical_k( + _, token_idx = _swaps_routed_coordinate( + cfg, lane_idx >> Int32(2), sparse_origin0, sparse_origin1, @@ -2351,7 +1959,7 @@ def reduce_sums( return sum_arr @cute.jit - def _compute_softmax_loop_sparse_keeps_kv256( + def _compute_softmax_loop_keeps_fragments( self, stage_info: StageInfo, *, @@ -2359,6 +1967,7 @@ def _compute_softmax_loop_sparse_keeps_kv256( sum_arr: cutlass.Array, new_max_arr: cutlass.Array, s_arr: cutlass.Array, + use_sparse: Constexpr[bool], sparse_origin0: Int32, sparse_origin1: Int32, sparse_route_flags: Int32, @@ -2367,77 +1976,141 @@ def _compute_softmax_loop_sparse_keeps_kv256( sparse_token_word2: Uint32, sparse_token_word3: Uint32, ) -> tuple[object, object, object, object]: - """Reduce one sparse KV256 route as four bounded K32 fragments. - - The full route path only loads and reduces scores. A partial route - predicates one native 32-score fragment at a time and writes it back - to TMEM, so the later P pass can replay masked scores without keeping - the logical 128-score tile live in registers. + """Mask streamed K32 score fragments in place and reduce their max. + + Every fragment gets one keep word. Block-sparse routes derive it from + their two K64 atom origins, validity flags and prepared token words; + dense tiles derive it from the tile's visible token range (sequence + end, uniform or per-row causal end, sliding-window start) and the Q + row's validity. Masked fragments are written back to TMEM so the P + pass can reload them without any mask logic. """ - cfg = self.cfg - assert cfg.tile_size_kv == 256 + assert cfg.streams_tmem_p_fragments + num_fragments = cfg.num_softmax_score_fragments + fragment_regs = cfg.softmax_score_fragment_regs + # The seven-slot softmax metadata ABI carries exactly four token words. + assert num_fragments == 4 and fragment_regs == 32 task_cache = _decode_gen_task_cache(stage_info) - token_words = ( - sparse_token_word0, - sparse_token_word1, - sparse_token_word2, - sparse_token_word3, + keep_words = cutlass.Array( + Uint32, num_fragments, space=cutlass.AddressSpace.rmem ) - keep_words = cutlass.Array(Uint32, 4, space=cutlass.AddressSpace.rmem) warp_group_thread_idx = Int32(task_cache[_TASK_CACHE_WARP_GRP_THREAD_IDX]) tile_row_idx = _keeps_row_idx(cfg, warp_group_thread_idx) - logical_q_group_idx = _logical_q_group_idx(cfg, stage_info, self.q_group_idx) - q_token_idx, _ = _q_row_token_and_local_head( - cfg, - self.h_r, - logical_q_group_idx, - tile_row_idx, - ) - q_row_is_valid = _q_row_is_valid_for_seq( - cfg, - self.h_r, - logical_q_group_idx, - tile_row_idx, - self.seq_len_q, - ) - seq_len_kv = _load_runtime_seq_len_kv( - self.seqlens_kv, - self.max_seq_len_kv, - stage_info, - Int32(0), - Int32(0), - ) - causal_end = seq_len_kv - self.seq_len_q + q_token_idx + Int32(1) - origin0 = Int32(sparse_origin0) - origin1 = Int32(sparse_origin1) - valid0 = sparse_route_flags & Int32(1) - valid1 = (sparse_route_flags >> Int32(1)) & Int32(1) - for fragment_idx in cutlass.range_constexpr(4): - fragment_origin = origin0 + Int32((fragment_idx % 2) * 32) - fragment_valid = valid0 - if cutlass.const_expr(fragment_idx >= 2): - fragment_origin = origin1 + Int32((fragment_idx % 2) * 32) - fragment_valid = valid1 - keep_words[fragment_idx] = _sparse_k32_effective_keep_word( - q_row_is_valid, - fragment_origin, - fragment_valid, - Uint32(token_words[fragment_idx]), - seq_len_kv, - causal_end, - apply_causal_mask=cfg.mask_type == CAUSAL, - apply_token_mask=cfg.use_kv_valid_bits, + if cutlass.const_expr(use_sparse): + token_words = ( + sparse_token_word0, + sparse_token_word1, + sparse_token_word2, + sparse_token_word3, + ) + logical_q_group_idx = _logical_q_group_idx( + cfg, stage_info, self.q_group_idx + ) + q_token_idx, _ = _q_row_token_and_local_head( + cfg, + self.h_r, + logical_q_group_idx, + tile_row_idx, + ) + q_row_is_valid = _q_row_is_valid_for_seq( + cfg, + self.h_r, + logical_q_group_idx, + tile_row_idx, + self.seq_len_q, + ) + seq_len_kv = _load_runtime_seq_len_kv( + self.seqlens_kv, + self.max_seq_len_kv, + stage_info, + Int32(0), + Int32(0), ) + causal_end = seq_len_kv - self.seq_len_q + q_token_idx + Int32(1) + origin0 = Int32(sparse_origin0) + origin1 = Int32(sparse_origin1) + valid0 = sparse_route_flags & Int32(1) + valid1 = (sparse_route_flags >> Int32(1)) & Int32(1) + fragments_per_origin = cfg.softmax_fragments_per_route_atom + for fragment_idx in cutlass.range_constexpr(num_fragments): + atom_offset = Int32( + (fragment_idx % fragments_per_origin) * fragment_regs + ) + fragment_origin = origin0 + atom_offset + fragment_valid = valid0 + if cutlass.const_expr(fragment_idx >= fragments_per_origin): + fragment_origin = origin1 + atom_offset + fragment_valid = valid1 + if cutlass.const_expr(cfg.trusts_prepared_score_words): + prepared_keep_word = Uint32(0) + if q_row_is_valid: + prepared_keep_word = Uint32(token_words[fragment_idx]) + keep_words[fragment_idx] = prepared_keep_word + else: + keep_words[fragment_idx] = _sparse_effective_keep_word( + q_row_is_valid, + fragment_origin, + fragment_valid, + Uint32(token_words[fragment_idx]), + seq_len_kv, + causal_end, + apply_causal_mask=cfg.mask_type == CAUSAL, + apply_token_mask=cfg.uses_prepared_score_keep_words, + ) - warp_scores_are_unmasked = cutlass.Boolean(True) - for fragment_idx in cutlass.range_constexpr(4): - warp_scores_are_unmasked = cutlass.Boolean( - warp_scores_are_unmasked - and keep_words[fragment_idx] == Uint32(0xFFFFFFFF) + else: + ( + seq_len_kv, + logical_q_group_idx, + element_mask_end_idx, + tile_offset_k, + window_start_idx, + _is_valid_effective_tile, + _is_masked_final_wave, + tile_is_unmasked, + rows_are_active, + ) = self._resolve_keeps_tile_context(stage_info) + if cutlass.const_expr(cfg.q_score_rows_need_mask): + rows_are_active = cutlass.Boolean( + rows_are_active + and _q_row_is_valid_for_seq( + cfg, + self.h_r, + logical_q_group_idx, + tile_row_idx, + self.seq_len_q, + ) + ) + visible_start = Int32(0) + visible_end = element_mask_end_idx + if cutlass.const_expr(cfg.uses_per_row_causal_mask): + q_token_idx, _ = _q_row_token_and_local_head( + cfg, + self.h_r, + logical_q_group_idx, + tile_row_idx, + ) + visible_end = seq_len_kv - self.seq_len_q + q_token_idx + Int32(1) + visible_start = _sliding_window_start_idx( + cfg, seq_len_kv, self.seq_len_q, q_token_idx + ) + elif cutlass.const_expr(cfg.use_sliding_window_causal): + visible_start = window_start_idx + # A tile that is unmasked for the whole Q group has all-ones keep + # words on every active row, so only masked tiles build them. + warp_scores_are_unmasked = cute.arch.vote_all_sync( + cutlass.Boolean(tile_is_unmasked and rows_are_active) ) - # The load/store branch must be uniform for each participating warp. - warp_scores_are_unmasked = cute.arch.vote_all_sync(warp_scores_are_unmasked) + if cutlass.const_expr(use_sparse): + warp_scores_are_unmasked = cutlass.Boolean(True) + for fragment_idx in cutlass.range_constexpr(num_fragments): + warp_scores_are_unmasked = cutlass.Boolean( + warp_scores_are_unmasked + and keep_words[fragment_idx] == Uint32(0xFFFFFFFF) + ) + # The load/store branch must be uniform for each participating warp. + warp_scores_are_unmasked = cute.arch.vote_all_sync(warp_scores_are_unmasked) score_tmem_addr = ( task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] @@ -2449,17 +2122,17 @@ def _compute_softmax_loop_sparse_keeps_kv256( max_chains[chain_idx] = _neg_max_f32() if warp_scores_are_unmasked: - for fragment_idx in cutlass.range_constexpr(4): + for fragment_idx in cutlass.range_constexpr(num_fragments): loaded = _keeps_tcgen05_ld( cfg, prims.make_tmem_ptr( - score_tmem_addr + Int32(fragment_idx * 32), Float32 + score_tmem_addr + Int32(fragment_idx * fragment_regs), Float32 ), - num=32, + num=fragment_regs, offset=cfg.tile_size_kv // 2, ) prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) - for score_idx in cutlass.range_constexpr(32): + for score_idx in cutlass.range_constexpr(fragment_regs): chain_idx: Constexpr[int] = score_idx % 4 max_chains[chain_idx] = cute.math.max( max_chains[chain_idx], @@ -2467,19 +2140,35 @@ def _compute_softmax_loop_sparse_keeps_kv256( ftz=True, ) else: - for fragment_idx in cutlass.range_constexpr(4): - fragment_addr = score_tmem_addr + Int32(fragment_idx * 32) + if cutlass.const_expr(not use_sparse): + lane_idx = Int32(task_cache[_TASK_CACHE_LANE_IDX]) + col_base = _keeps_col_base(cfg, lane_idx, num_fragments * fragment_regs) + for fragment_idx in cutlass.range_constexpr(num_fragments): + fragment_token_base = tile_offset_k + _keeps_score_col( + cfg, + warp_group_thread_idx, + fragment_idx * fragment_regs, + col_base, + ) + keep_words[fragment_idx] = _dense_fragment_keep_word( + rows_are_active, + visible_start - fragment_token_base, + visible_end - fragment_token_base, + fragment_regs=fragment_regs, + ) + for fragment_idx in cutlass.range_constexpr(num_fragments): + fragment_addr = score_tmem_addr + Int32(fragment_idx * fragment_regs) loaded = _keeps_tcgen05_ld( cfg, prims.make_tmem_ptr(fragment_addr, Float32), - num=32, + num=fragment_regs, offset=cfg.tile_size_kv // 2, ) prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) masked_scores = cutlass.Array( - Float32, 32, space=cutlass.AddressSpace.rmem + Float32, fragment_regs, space=cutlass.AddressSpace.rmem ) - for score_idx in cutlass.range_constexpr(32): + for score_idx in cutlass.range_constexpr(fragment_regs): score = Float32(loaded[score_idx]) score_is_kept = ( (keep_words[fragment_idx] >> Int32(score_idx)) & Uint32(1) @@ -2494,7 +2183,7 @@ def _compute_softmax_loop_sparse_keeps_kv256( _keeps_tcgen05_st( cfg, prims.make_tmem_ptr(fragment_addr, Float32), - masked_scores.data_ptr().load(count=32, alignment=4), + masked_scores.data_ptr().load(count=fragment_regs, alignment=4), offset=cfg.tile_size_kv // 2, ) prims.tcgen05_wait(kind=prims.Tcgen05Wait.STORE) @@ -2506,44 +2195,11 @@ def _compute_softmax_loop_sparse_keeps_kv256( ftz=True, ) old_max = new_max_arr[0] - new_max = cute.math.max(old_max, tile_max, ftz=True) - if old_max != _neg_max_f32(): - max_delta_log2 = self.scale_softmax_log2 * (old_max - new_max) - if max_delta_log2 >= Float32(-KV_TILE_256_RESCALE_THRESHOLD_LOG2): - new_max = old_max + new_max = self._softmax_anchor(old_max, tile_max) old_max_arr[0] = old_max new_max_arr[0] = new_max return old_max_arr, sum_arr, new_max_arr, s_arr - @cute.jit - def _load_block_sparse_softmax_p_fragment( - self, - stage_info: StageInfo, - *, - fragment_idx: Constexpr[int], - s_arr: cutlass.Array, - ) -> cutlass.Array: - """Reload one full or already-predicated sparse KV256 fragment for P.""" - - assert self.cfg.tile_size_kv == 256 - task_cache = _decode_gen_task_cache(stage_info) - score_tmem_addr = ( - task_cache[_TASK_CACHE_TMEM_BASE_OFFSET] - + Int32(self._alloc.offset) - + self._softmax_loop_stage_slot_offset(stage_info) - + Int32(fragment_idx * 32) - ) - loaded = _keeps_tcgen05_ld( - self.cfg, - prims.make_tmem_ptr(score_tmem_addr, Float32), - num=32, - offset=self.cfg.tile_size_kv // 2, - ) - prims.tcgen05_wait(kind=prims.Tcgen05Wait.LOAD) - for score_idx in cutlass.range_constexpr(32): - s_arr[score_idx] = loaded[score_idx] - return s_arr - @consumer_work(returns=("old_max_arr", "sum_arr", "new_max_arr", "s_arr")) @cute.jit def compute_block_sparse_softmax_loop( @@ -2566,34 +2222,21 @@ def compute_block_sparse_softmax_loop( assert self.cfg.use_block_sparse if cutlass.const_expr(self.cfg.use_keeps_mma_ab): - if cutlass.const_expr(self.cfg.tile_size_kv == 256): - return self._compute_softmax_loop_sparse_keeps_kv256( - stage_info, - old_max_arr=old_max_arr, - sum_arr=sum_arr, - new_max_arr=new_max_arr, - s_arr=s_arr, - sparse_origin0=sparse_origin0, - sparse_origin1=sparse_origin1, - sparse_route_flags=sparse_route_flags, - sparse_token_word0=sparse_token_word0, - sparse_token_word1=sparse_token_word1, - sparse_token_word2=sparse_token_word2, - sparse_token_word3=sparse_token_word3, - ) - return self._compute_softmax_loop_sparse_keeps( + # Every block-sparse Keeps profile streams K32 fragments. + return self._compute_softmax_loop_keeps_fragments( stage_info, old_max_arr=old_max_arr, sum_arr=sum_arr, new_max_arr=new_max_arr, s_arr=s_arr, - routed_origin0=sparse_origin0, - routed_origin1=sparse_origin1, - routed_route_flags=sparse_route_flags, - routed_token_word0=sparse_token_word0, - routed_token_word1=sparse_token_word1, - routed_token_word2=sparse_token_word2, - routed_token_word3=sparse_token_word3, + use_sparse=True, + sparse_origin0=sparse_origin0, + sparse_origin1=sparse_origin1, + sparse_route_flags=sparse_route_flags, + sparse_token_word0=sparse_token_word0, + sparse_token_word1=sparse_token_word1, + sparse_token_word2=sparse_token_word2, + sparse_token_word3=sparse_token_word3, ) # SWAP reuses the Keeps seven-slot task ABI: all four origins remain # logical KV atom bases, but origin2 occupies the flags slot and diff --git a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_tasks.py b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_tasks.py index 4321a0d0d46e..b75de72ba008 100644 --- a/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_tasks.py +++ b/tensorrt_llm/_torch/attention/backends/prims_ts/kernels/fmha_decode/fmha_decode_tasks.py @@ -32,16 +32,13 @@ import cutlass import cutlass.cute as cute -from cutlass import Int32 from cutlass.experimental import primitives as prims from cutlass.experimental.task_scheduling.memory import ( ResourceContext, - SmemAllocation, ) from cutlass.experimental.task_scheduling.resources import ( MemoryResource, StageInfo, - TaskLocalVariable, WorkQueue, consumer_work, producer_work, @@ -61,10 +58,9 @@ KV_INST1, KV_KIND_K, KV_KIND_V, - KV_TILE_256_SHARED_FIFO_STAGES, ) from .fmha_decode_resources.helpers_common import ( - ResourceVars, + _assume_nonnegative_i32, _q_group_token_base, _q_seq_bounds, _warp_broadcast_i32, @@ -122,14 +118,18 @@ def restore_slots(*resource_proxies: object) -> None: return traced -def _block_sparse_route_loop_domain( - route_count: cutlass.Int32, +def _loop_domain_after_head( + total_kv_tiles: cutlass.Int32, *, num_insts_kv: int, ) -> cutlass.Int32: - """Return LOOP iterations after HEAD reserves one candidate per instance.""" + """Return LOOP iterations after HEAD reserves one KV tile per instance. - remaining = route_count - cutlass.Int32(num_insts_kv) + Dense tiles and block-sparse routes share this recurrence; only the tile + count's source differs. + """ + + remaining = total_kv_tiles - cutlass.Int32(num_insts_kv) remaining = cute.math.max(remaining, cutlass.Int32(0)) insts = cutlass.Int32(num_insts_kv) return (remaining + insts - cutlass.Int32(1)) // insts @@ -152,102 +152,6 @@ def consume_schedule_token(self, stage_info: StageInfo) -> None: del stage_info -@dataclass(kw_only=True) -class SmemKvReuseCreditResource(MemoryResource): - """One-slot credit carrying the rotating KV256 exchange-stage index. - - Load publishes which drained 64-KiB physical K/V stage Correction may use - as tail scratch. The one-stage pipeline couples that payload to the same - ownership epoch: the following Load may use the other two physical stages, - but cannot publish a new alias until Correction releases this credit after - all output work completes. - """ - - cfg: cutlass.Constexpr[FmhaDecodeConfig] = None - _alloc: cutlass.Constexpr[SmemAllocation | None] = None - scratch_stage_slot: cutlass.Constexpr[TaskLocalVariable] = ( - TaskLocalVariable.uninitialized() - ) - - def __post_init__(self) -> None: - """Create the routed consumer slot for one physical K/V stage.""" - assert KV_TILE_256_SHARED_FIFO_STAGES == 3, ( - "KV256 reuse-credit rotation requires exactly three shared FIFO stages" - ) - if not self.cfg.uses_rotating_kv256_exchange: - raise ValueError( - "rotating KV scratch requires persistent direct Q64/KV256 " - "with two KV instructions, one head-dimension stage, and " - "one load warp" - ) - self.scratch_stage_slot = TaskLocalVariable( - dtype=Int32, - default=Int32(0), - docs="Physical shared-K/V stage reserved for KV256 tail exchange.", - ) - - def get_smem_requirements(self) -> list[SmemAllocation]: - """Allocate the one-word stage payload guarded by this pipeline.""" - if self._alloc is None: - self._alloc = SmemAllocation( - name=f"{self.name}_scratchStage", - size_bytes=4, - alignment=4, - ) - return [self._alloc] - - @cute.jit - def _payload(self, stage_info: StageInfo) -> cutlass.Array: - """Return the natural next-stage cursor owned by this credit.""" - return cutlass.Array( - stage_info.context.smem_base.data_ptr() + self._alloc.offset, - dtype=Int32, - shape=(1,), - addrspace=3, - ) - - @cute.jit - def create_function_variables( - self, - context: ResourceContext | None = None, - ) -> ResourceVars: - """Initialize the persistent ring cursor before TS tasks start.""" - if cutlass.const_expr(context is not None and context.smem_base is not None): - payload = cutlass.Array( - context.smem_base.data_ptr() + self._alloc.offset, - dtype=Int32, - shape=(1,), - addrspace=3, - ) - thread_idx, _, _ = cute.arch.thread_idx() - if thread_idx == Int32(0): - payload[0] = Int32(0) - return {} - - @producer_work - @cute.jit - def publish_scratch_stage(self, stage_info: StageInfo) -> None: - """Advance the persistent ring cursor and publish the drained stage.""" - num_stages = Int32(KV_TILE_256_SHARED_FIFO_STAGES) - if prims.elect_sync(): - payload = self._payload(stage_info) - # Each work commits T = 4 * (loop_end + 1) K/V transactions. - # Since 4 == 1 (mod 3), the cursor advances by loop_end + 1. - # loop_end is the resolved per-work domain, so heterogeneous - # runtime sequence lengths do not inherit a captured host bound. - payload[0] = ( - Int32(payload[0]) + stage_info.loop_end + Int32(1) - ) % num_stages - - @consumer_work(returns=scratch_stage_slot) - @cute.jit - def read_scratch_stage(self, stage_info: StageInfo) -> Int32: - """Read the alias only after the matching credit wait completes.""" - num_stages = Int32(KV_TILE_256_SHARED_FIFO_STAGES) - next_stage = Int32(self._payload(stage_info)[0]) - return (next_stage + num_stages - Int32(1)) % num_stages - - @dataclass(kw_only=True) class PackedDecodeWorkQueue(WorkQueue): """CLC work queue that drops packed-Q tiles beyond a batch's Q length.""" @@ -414,28 +318,23 @@ def _produce_staged_page_offsets( def _consume_staged_qk_mma( smem_kv: MemoryResource, tmem_s: MemoryResource, - aliased_p: MemoryResource, q_desc: Any, k_desc_label: str, qk_mma_label: str, section: FmhaStage, cfg: FmhaDecodeConfig, ) -> None: - """Consume all K head-dim stages for one QK MMA wave.""" + """Consume all K head-dim stages for one QK MMA wave. + + Streamed KV256 aliases P with the S columns this QK overwrites. The + preceding same-instance PV reads P as its TMEM A operand from the same + issuing thread, and the tensor core interlocks that read against a later + MMA's accumulator write, so no completion wait is needed before QK. + """ tmem_s.acquire() for head_dim_stage_idx in range(cfg.num_head_dim_stages_kv): smem_kv.wait() kv_desc = getattr(smem_kv, k_desc_label)() - if cutlass.const_expr( - cfg.streams_tmem_p_fragments - and head_dim_stage_idx == 0 - and (section == FmhaStage.Loop or cfg.use_persistent_scheduler) - ): - # Wait as late as possible: K staging overlaps the previous PV, - # but QK cannot overwrite the matching S/P alias until PV is done. - # Static HEAD has no previous tile; persistent HEAD may follow the - # same CTA's tail from another logical work tile and must wait. - aliased_p.wait_until_reusable_before_qk() if cfg.uses_q_desc_ref: getattr(tmem_s, f"{qk_mma_label}_from_q_ref")( kv_desc=kv_desc, @@ -451,6 +350,44 @@ def _consume_staged_qk_mma( tmem_s.commit() +def _consume_streamed_pv_fragments( + smem_kv: MemoryResource, + tmem_p: MemoryResource, + tmem_o: MemoryResource, + v_desc_label: str, + vp_mma_label: str, + cfg: FmhaDecodeConfig, +) -> None: + """Issue one PV wave as its K32 P fragments become ready. + + P fragment 0 is the earliest dependency: wait for it and for the + correction credit before holding the V stage. Later fragments may become + ready while the previous PV fragment is already executing; every slot + stays live through the complete async UMMA wave so the producer cannot + overwrite an operand prematurely. + """ + assert cfg.num_head_dim_stages_kv == 1 + fragment_label = f"{vp_mma_label}_fragment" + p_tmem_addr = tmem_p.wait_p_fragment(fragment_idx=0) + tmem_o.acquire() + smem_kv.wait() + v_desc = getattr(smem_kv, v_desc_label)() + getattr(tmem_o, fragment_label)( + v_desc=v_desc, + p_tmem_addr=p_tmem_addr, + fragment_idx=0, + ) + for fragment_idx in range(1, cfg.num_softmax_score_fragments): + p_tmem_addr = tmem_p.wait_p_fragment(fragment_idx=fragment_idx) + getattr(tmem_o, fragment_label)( + v_desc=v_desc, + p_tmem_addr=p_tmem_addr, + fragment_idx=fragment_idx, + ) + smem_kv.release() + tmem_o.commit() + + def _consume_staged_pv_mma( smem_kv: MemoryResource, tmem_p: MemoryResource, @@ -464,33 +401,9 @@ def _consume_staged_pv_mma( """Consume all V head-dim stages for one PV MMA wave.""" _ = section if cutlass.const_expr(cfg.streams_tmem_p_fragments): - assert cfg.num_head_dim_stages_kv == 1 - fragment_label = f"{vp_mma_label}_fragment" - - # P fragment 0 is the earliest dependency. Wait for it and for the - # correction credit before holding the shared V FIFO stage. - p_tmem_addr = tmem_p.wait_p_fragment(fragment_idx=0) - tmem_o.acquire() - smem_kv.wait() - v_desc = getattr(smem_kv, v_desc_label)() - getattr(tmem_o, fragment_label)( - v_desc=v_desc, - p_tmem_addr=p_tmem_addr, - fragment_idx=0, - ) - - # Later P fragments may become ready while the previous PV fragment is - # already executing. Keep every slot live through the complete async - # UMMA wave so the producer cannot overwrite an operand prematurely. - for fragment_idx in range(1, cfg.num_softmax_score_fragments): - p_tmem_addr = tmem_p.wait_p_fragment(fragment_idx=fragment_idx) - getattr(tmem_o, fragment_label)( - v_desc=v_desc, - p_tmem_addr=p_tmem_addr, - fragment_idx=fragment_idx, - ) - smem_kv.release() - tmem_o.commit() + _consume_streamed_pv_fragments( + smem_kv, tmem_p, tmem_o, v_desc_label, vp_mma_label, cfg + ) return tmem_p.wait() @@ -634,6 +547,47 @@ def _decode_work_tile_schedule_with_invariant_bridge( _work_queue_tail(work_queue) +@cute.jit +def _prepared_sparse_row_address( + cfg: cutlass.Constexpr[FmhaDecodeConfig], + q_group_idx: cutlass.Int32, + h_idx: cutlass.Int32, + b_idx: cutlass.Int32, + num_heads_kv: cutlass.Int32, +) -> cutlass.Int32: + """Map a (q_group, head, batch) tile to its prepared row header index.""" + + q_token_base = _q_group_token_base(cfg, q_group_idx) + q_block = q_token_base // cutlass.Int32(cfg.q_block_size) + num_q_blocks = (cfg.max_seq_len_q + cfg.q_block_size - 1) // cfg.q_block_size + return (b_idx * num_heads_kv + h_idx) * cutlass.Int32(num_q_blocks) + q_block + + +@cute.jit +def _prefetch_prepared_sparse_row( + cfg: cutlass.Constexpr[FmhaDecodeConfig], + row_route_offsets: cute.Pointer, + row_route_counts: cute.Pointer, + q_group_idx: cutlass.Int32, + h_idx: cutlass.Int32, + b_idx: cutlass.Int32, + num_heads_kv: cutlass.Int32, +) -> tuple[cutlass.Int32, cutlass.Int32]: + """Load one static tile's prepared row header before its tasks start. + + Every thread loads the same two words, so each warp issues one request + and the global-memory latency overlaps the CTA prologue (TMEM allocation + and barrier setup) instead of stalling every task at its first step. + """ + + row_address = _prepared_sparse_row_address( + cfg, q_group_idx, h_idx, b_idx, num_heads_kv + ) + row_route_begin = cutlass.Int32(row_route_offsets[row_address]) + route_count = _assume_nonnegative_i32(cutlass.Int32(row_route_counts[row_address])) + return row_route_begin, route_count + + @cute.jit def _load_prepared_sparse_row_warp( row_route_offsets: cute.Pointer, @@ -653,7 +607,7 @@ def _load_prepared_sparse_row_warp( loaded_row_route_begin = cutlass.Int32(row_route_offsets[row_address]) loaded_route_count = cutlass.Int32(row_route_counts[row_address]) row_route_begin = _warp_broadcast_i32(loaded_row_route_begin, 0) - route_count = _warp_broadcast_i32(loaded_route_count, 0) + route_count = _assume_nonnegative_i32(_warp_broadcast_i32(loaded_route_count, 0)) return row_route_begin, route_count @@ -667,6 +621,9 @@ def __init__(self, **kwargs: TaskKwarg) -> None: self.block_table_capacity = kwargs.pop("block_table_capacity", None) self.sparse_row_route_offsets = kwargs.pop("sparse_row_route_offsets", None) self.sparse_row_route_counts = kwargs.pop("sparse_row_route_counts", None) + # Static tiles may pass the already loaded row header instead. + self.sparse_row_route_begin = kwargs.pop("sparse_row_route_begin", None) + self.sparse_route_count = kwargs.pop("sparse_route_count", None) self.num_heads_kv = kwargs.pop("num_heads_kv", None) self.max_seq_len_kv = kwargs.pop("max_seq_len_kv", cutlass.Int32(0)) self.seq_len_q = kwargs.pop("seq_len_q", None) @@ -923,20 +880,20 @@ def get_domain(self, tile_coord: cute.Coord) -> cutlass.Int32 | int: q_group_idx = cutlass.Int32(tile_coord[0]) h_idx = cutlass.Int32(tile_coord[1]) b_idx = cutlass.Int32(tile_coord[2]) - q_token_base = _q_group_token_base(self.cfg, q_group_idx) - - q_block = q_token_base // self.cfg.q_block_size - num_q_blocks = ( - self.cfg.max_seq_len_q + self.cfg.q_block_size - 1 - ) // self.cfg.q_block_size - row_address = (b_idx * self.num_heads_kv + h_idx) * num_q_blocks + q_block - - row_route_begin, route_count = _load_prepared_sparse_row_warp( - row_route_offsets, - row_route_counts, - cutlass.Int32(row_address), - self._lane_idx, - ) + if self.sparse_row_route_begin is not None: + # The static kernel prologue already loaded this tile's header. + row_route_begin = self.sparse_row_route_begin + route_count = self.sparse_route_count + else: + row_address = _prepared_sparse_row_address( + self.cfg, q_group_idx, h_idx, b_idx, self.num_heads_kv + ) + row_route_begin, route_count = _load_prepared_sparse_row_warp( + row_route_offsets, + row_route_counts, + row_address, + self._lane_idx, + ) # Sparse route-span accessors share two underlying cache words # with paged KV. Clear dense/paged-only coordinates on every @@ -951,7 +908,7 @@ def get_domain(self, tile_coord: cute.Coord) -> cutlass.Int32 | int: self._kv_valid_tile_end = route_count self._kv_window_start = cutlass.Int32(0) - loop_domain = _block_sparse_route_loop_domain( + loop_domain = _loop_domain_after_head( route_count, num_insts_kv=self.cfg.num_insts_kv, ) @@ -997,14 +954,10 @@ def get_domain(self, tile_coord: cute.Coord) -> cutlass.Int32 | int: self._kv_window_start = cutlass.Int32(0) self._kv_valid_tile_end = total_kv_tiles self._kv_raw_tile_base = cutlass.Int32(0) - remaining_kv_tiles = cute.math.max( - total_kv_tiles - cutlass.Int32(self.cfg.num_insts_kv), - cutlass.Int32(0), + loop_domain = _loop_domain_after_head( + total_kv_tiles, + num_insts_kv=self.cfg.num_insts_kv, ) - num_insts_kv = cutlass.Int32(self.cfg.num_insts_kv) - loop_domain = ( - remaining_kv_tiles + num_insts_kv - cutlass.Int32(1) - ) // num_insts_kv return loop_domain + cutlass.Int32(self.domain_bias) # Decode the logical Q tile with the configured physical split fanout, @@ -1058,13 +1011,10 @@ def get_domain(self, tile_coord: cute.Coord) -> cutlass.Int32 | int: self._kv_raw_tile_base = skipped_tiles + split_idx * total_kv_tiles else: self._kv_raw_tile_base = skipped_tiles - remaining_kv_tiles = cute.math.max( - total_kv_tiles - cutlass.Int32(self.cfg.num_insts_kv), cutlass.Int32(0) + loop_domain = _loop_domain_after_head( + total_kv_tiles, + num_insts_kv=self.cfg.num_insts_kv, ) - num_insts_kv = cutlass.Int32(self.cfg.num_insts_kv) - loop_domain = ( - remaining_kv_tiles + num_insts_kv - cutlass.Int32(1) - ) // num_insts_kv # All tasks share the MMA-loop domain; tail-only tasks add a bias. return loop_domain + cutlass.Int32(self.domain_bias) @@ -1081,29 +1031,63 @@ def get_domain(self, tile_coord: cute.Coord) -> cutlass.Int32 | int: def _resolve_and_store_sparse_route( sparse_kv_metadata: MemoryResource | None, section: FmhaStage, -) -> tuple[Any, Any, Any, Any] | None: - """Resolve one prepared route and retain it for the matching K/V pair.""" + prefetch: tuple[Any, Any] | None = None, + *, + pipeline: bool = True, +) -> tuple[tuple[Any, Any, Any, Any] | None, tuple[Any, Any] | None]: + """Resolve one prepared route and retain it for the matching K/V pair. + + Returns ``(route, prefetch)``. With ``pipeline`` the record load is issued + one resolution ahead: HEAD loads its own record immediately, and every + resolution issues the load for the next one (LOOP iteration 0 from HEAD, + iteration i + 1 from iteration i) before the caller's K TMA burst, so the + global-memory latency overlaps that issue instead of stalling the load + warp. Callers pass the returned ``prefetch`` back into the next resolution + of the same instance, the way ``_staged_kv_load`` threads its cached page + IDs. Without ``pipeline`` the record is loaded where it is resolved and no + state is returned; the split-ring load variants use this because the + pipelined form measured slower for them. Dense profiles pass ``None`` and + get ``(None, None)``. + """ if sparse_kv_metadata is None: - return None + return None, None + if not pipeline: + prefetch = sparse_kv_metadata.prefetch_route( + target="head" if section == FmhaStage.Head else "current_loop" + ) + elif section == FmhaStage.Head: + prefetch = sparse_kv_metadata.prefetch_route(target="head") + assert prefetch is not None + prefetched_record_word, prefetched_record_offset = prefetch ( - resolved_origin0, + resolved_record_word, resolved_origin1, resolved_atom_validity, route_record_word_offset, - ) = sparse_kv_metadata.resolve_route(section=section) + ) = sparse_kv_metadata.resolve_route( + section=section, + prefetched_record_word_slot=prefetched_record_word, + prefetched_record_offset_slot=prefetched_record_offset, + ) sparse_kv_metadata.store_route( - resolved_origin0=resolved_origin0, + resolved_record_word=resolved_record_word, resolved_origin1=resolved_origin1, resolved_atom_validity=resolved_atom_validity, route_record_word_offset=route_record_word_offset, ) - return ( - resolved_origin0, + next_prefetch = None + if pipeline: + next_prefetch = sparse_kv_metadata.prefetch_route( + target="first_loop" if section == FmhaStage.Head else "next_loop" + ) + route = ( + resolved_record_word, resolved_origin1, resolved_atom_validity, route_record_word_offset, ) + return route, next_prefetch def _publish_sparse_softmax_route( @@ -1116,14 +1100,14 @@ def _publish_sparse_softmax_route( return assert route is not None ( - resolved_origin0, + resolved_record_word, resolved_origin1, resolved_atom_validity, route_record_word_offset, ) = route sparse_softmax_metadata.acquire() sparse_softmax_metadata.store_route( - resolved_origin0=resolved_origin0, + resolved_record_word=resolved_record_word, resolved_origin1=resolved_origin1, resolved_atom_validity=resolved_atom_validity, route_record_word_offset=route_record_word_offset, @@ -1136,7 +1120,6 @@ def create_load_task( smem_kv: MemoryResource, work_queue: WorkQueue | None, schedule_token_throttle: MemoryResource | None, - smem_kv_reuse_credit: MemoryResource | None, cfg: FmhaDecodeConfig, *, domain: int | cutlass.Int32, @@ -1158,7 +1141,6 @@ def load_schedule_body( smem_kv: MemoryResource, smem_page_offsets: MemoryResource | None, schedule_token_throttle: MemoryResource | None, - smem_kv_reuse_credit: MemoryResource | None, sparse_kv_metadata0: MemoryResource | None = None, sparse_kv_metadata1: MemoryResource | None = None, sparse_softmax_metadata0: MemoryResource | None = None, @@ -1208,26 +1190,22 @@ def _kv_load(label: str, section: FmhaStage) -> None: smem_page_offsets.wait() else: _page_offsets_consume(smem_page_offsets) - if sparse_kv_metadata0 is None: - for label in ("load_k0", "load_k1"): - _kv_load(label, FmhaStage.Head) - else: - route0 = _resolve_and_store_sparse_route( - sparse_kv_metadata0, FmhaStage.Head - ) - _kv_load("load_k0", FmhaStage.Head) - route1 = _resolve_and_store_sparse_route( - sparse_kv_metadata1, FmhaStage.Head - ) - _kv_load("load_k1", FmhaStage.Head) - # Issue both K tiles before either metadata FIFO can backpressure - # the load warp, matching the split-resource sparse cadence. - _publish_sparse_softmax_route(sparse_softmax_metadata0, route0) - _publish_sparse_softmax_route(sparse_softmax_metadata1, route1) - if smem_kv_reuse_credit is not None: - # K0/K1 occupy the two stages disjoint from the previous work's - # scratch. Acquire only before issuing the third K/V transaction. - smem_kv_reuse_credit.acquire() + # Dense profiles have no route metadata: the resolve and publish + # helpers are no-ops for ``None`` resources, so one cadence serves + # both dense and block-sparse loads. + route0, prefetch0 = _resolve_and_store_sparse_route( + sparse_kv_metadata0, FmhaStage.Head + ) + _kv_load("load_k0", FmhaStage.Head) + route1, prefetch1 = _resolve_and_store_sparse_route( + sparse_kv_metadata1, FmhaStage.Head + ) + _kv_load("load_k1", FmhaStage.Head) + # Issue both K tiles before either metadata FIFO can backpressure + # the load warp, matching the split-resource sparse cadence. + _publish_sparse_softmax_route(sparse_softmax_metadata0, route0) + _publish_sparse_softmax_route(sparse_softmax_metadata1, route1) + prefetch_by_label = {"load_k0": prefetch0, "load_k1": prefetch1} # LOOP: each iter prefetches the full ``num_insts_kv`` K/V pair set. # When P aliases the consumed S columns, MMA must consume each V/P pair @@ -1239,41 +1217,30 @@ def _kv_load(label: str, section: FmhaStage) -> None: else ("load_k0", "load_v0", "load_k1", "load_v1") ) with domain_loop(0, domain, 1, unroll=1): - if sparse_kv_metadata0 is None: - for label in loop_labels: - _kv_load(label, FmhaStage.Loop) - else: - # Follow the dense KV256 stage order exactly. Each V consumes - # its retained route before the matching K label replaces it. - loop_routes = [] - for label in loop_labels: - route = None - sparse_softmax_metadata = None - if label == "load_k0": - route = _resolve_and_store_sparse_route( - sparse_kv_metadata0, FmhaStage.Loop - ) - sparse_softmax_metadata = sparse_softmax_metadata0 - elif label == "load_k1": - route = _resolve_and_store_sparse_route( - sparse_kv_metadata1, FmhaStage.Loop - ) - sparse_softmax_metadata = sparse_softmax_metadata1 - _kv_load(label, FmhaStage.Loop) - if route is not None: - loop_routes.append((sparse_softmax_metadata, route)) - for sparse_softmax_metadata, route in loop_routes: - _publish_sparse_softmax_route(sparse_softmax_metadata, route) + # Generic V-first profiles consume their retained route before the + # matching K label replaces it. + route_metadata_by_label = { + "load_k0": (sparse_kv_metadata0, sparse_softmax_metadata0), + "load_k1": (sparse_kv_metadata1, sparse_softmax_metadata1), + } + loop_routes = [] + for label in loop_labels: + kv_metadata, softmax_metadata = route_metadata_by_label.get( + label, (None, None) + ) + route, prefetch_by_label[label] = _resolve_and_store_sparse_route( + kv_metadata, FmhaStage.Loop, prefetch_by_label.get(label) + ) + _kv_load(label, FmhaStage.Loop) + if route is not None: + loop_routes.append((softmax_metadata, route)) + for sparse_softmax_metadata, route in loop_routes: + _publish_sparse_softmax_route(sparse_softmax_metadata, route) # TAIL: after no more future K tiles are needed, load the final two V # tiles consumed by the final BMM2 calls. for label in ("load_v0", "load_v1"): _kv_load(label, FmhaStage.Tail) - if smem_kv_reuse_credit is not None: - # Publish the physical stage drained by this work together with - # the ownership token consumed by the correction tail. - smem_kv_reuse_credit.publish_scratch_stage() - smem_kv_reuse_credit.commit() if hold_page_window: _page_offsets_release(smem_page_offsets) @@ -1288,7 +1255,6 @@ def load_schedule( sparse_softmax_metadata1: MemoryResource | None, work_queue: WorkQueue | None, schedule_token_throttle: MemoryResource | None, - smem_kv_reuse_credit: MemoryResource | None, ) -> None: """Schedule shared-KV loads with only the resources in this profile.""" @@ -1300,7 +1266,6 @@ def load_schedule( smem_kv, smem_page_offsets, schedule_token_throttle, - smem_kv_reuse_credit, sparse_kv_metadata0, sparse_kv_metadata1, sparse_softmax_metadata0, @@ -1334,7 +1299,6 @@ def load_schedule( sparse_softmax_metadata1, work_queue, schedule_token_throttle, - smem_kv_reuse_credit, ) src = [] for sparse_kv_metadata in (sparse_kv_metadata0, sparse_kv_metadata1): @@ -1350,8 +1314,6 @@ def load_schedule( dst.append(sparse_resource) if schedule_token_throttle is not None: dst.append(schedule_token_throttle) - if smem_kv_reuse_credit is not None: - dst.append(smem_kv_reuse_credit) return task_class( src_resources=src, dst_resources=dst, @@ -1720,7 +1682,9 @@ def load_tile( ) in active_instances: if smem_k is None: continue - route = _resolve_and_store_sparse_route(sparse_kv_metadata, FmhaStage.Head) + route, _ = _resolve_and_store_sparse_route( + sparse_kv_metadata, FmhaStage.Head, pipeline=False + ) load_tile(smem_k, load_k, smem_page_offsets_k, FmhaStage.Head) head_routes.append((sparse_softmax_metadata, route)) # In the combined task, preserve both K issues ahead of Softmax @@ -1748,8 +1712,8 @@ def load_tile( smem_page_offsets_v_local, FmhaStage.Loop, ) - route = _resolve_and_store_sparse_route( - sparse_kv_metadata, FmhaStage.Loop + route, _ = _resolve_and_store_sparse_route( + sparse_kv_metadata, FmhaStage.Loop, pipeline=False ) load_tile(smem_k, load_k, smem_page_offsets_k, FmhaStage.Loop) loop_routes.append((sparse_softmax_metadata, route)) @@ -2175,26 +2139,9 @@ def pv_mma( section: FmhaStage, ) -> None: """Issue one scheduled PV wave using the selected phase work.""" - _ = section - tmem_p.wait() - p_desc_0, p_desc_1, p_tmem_addr_0, p_tmem_addr_1 = tmem_p.p_operands() - tmem_o.acquire() - for head_dim_stage_idx in range(cfg.num_head_dim_stages_kv): - smem_kv.wait() - v_desc = smem_kv.v_desc() - getattr(tmem_o, vp_mma_label)( - v_desc_0=v_desc, - v_desc_1=v_desc, - p_desc_0=p_desc_0, - p_desc_1=p_desc_1, - p_tmem_addr_0=p_tmem_addr_0, - p_tmem_addr_1=p_tmem_addr_1, - inst_idx=inst_idx, - head_dim_stage_idx=head_dim_stage_idx, - ) - smem_kv.release() - tmem_o.commit() - tmem_p.release() + _consume_staged_pv_mma( + smem_kv, tmem_p, tmem_o, "v_desc", vp_mma_label, inst_idx, section, cfg + ) qk_mma( smem_k0, @@ -2652,7 +2599,6 @@ def mma_schedule_body( _consume_staged_qk_mma( smem_kv, tmem_s0, - smem_p0, q_desc, "k_desc_0", "qk_mma_head", @@ -2662,7 +2608,6 @@ def mma_schedule_body( _consume_staged_qk_mma( smem_kv, tmem_s1, - smem_p1, q_desc, "k_desc_1", "qk_mma_head", @@ -2671,8 +2616,9 @@ def mma_schedule_body( ) # LOOP: consume aliased TMEM P before the next same-instance QK - # overwrites its S columns. SMEM-P profiles retain their established - # K-before-V cadence because P no longer depends on S lifetime. + # overwrites its S columns. Full-SMEM P remains score-dependent during + # Softmax replay, but owns independent storage once the replay commits + # and releases S; that completed handoff enables the QK-before-PV cadence. with domain_loop(0, domain, 1, unroll=1): if cfg.uses_two_inst_tmem_p: _consume_staged_pv_mma( @@ -2688,7 +2634,6 @@ def mma_schedule_body( _consume_staged_qk_mma( smem_kv, tmem_s0, - smem_p0, q_desc, "k_desc_0", "qk_mma_loop", @@ -2720,7 +2665,6 @@ def mma_schedule_body( _consume_staged_qk_mma( smem_kv, tmem_s1, - smem_p1, q_desc, "k_desc_1", "qk_mma_loop", @@ -2739,6 +2683,12 @@ def mma_schedule_body( cfg, ) + # Q is live for every BMM1 call, and the last BMM1 has been issued once + # the loop ends. Releasing here commits after those MMAs complete, so + # the next tile's Q load overlaps the final softmax and BMM2 waves + # instead of waiting for them. + smem_q.release() + # TAIL: no future K tiles remain, so only the final two BMM2 waves run. _consume_staged_pv_mma( smem_kv, @@ -2760,8 +2710,6 @@ def mma_schedule_body( FmhaStage.Tail, cfg, ) - # Q is live for every BMM1 call and can be released only after the loop. - smem_q.release() def mma_schedule_prelude( smem_q: MemoryResource, @@ -2884,9 +2832,9 @@ def softmax0_schedule_body( sparse_softmax_metadata.init_read_state() with domain_loop(0, domain, 1, unroll=1) as d: - # ConsWait/ConsWork: load S from TMEM and compute the tile max. - tmem_s0.wait() if sparse_softmax_metadata is not None: + # Consume the independent metadata stream first so its SMEM + # loads and release can overlap the subsequent score wait. sparse_softmax_metadata.wait() # Copy the complete payload to registers before release, so # masking cannot race the producer's next SMEM-stage reuse. @@ -2900,6 +2848,9 @@ def softmax0_schedule_body( sparse_token_word3, ) = sparse_softmax_metadata.load_route() sparse_softmax_metadata.release() + # ConsWait/ConsWork: load S from TMEM and compute the tile max. + tmem_s0.wait() + if sparse_softmax_metadata is not None: old_max_arr, sum_arr, new_max_arr, s_arr = ( tmem_s0.compute_block_sparse_softmax_loop( old_max_arr=old_max_arr, @@ -2937,18 +2888,17 @@ def softmax0_schedule_body( ) tmem_softmax_local0.commit() if cutlass.const_expr(cfg.streams_tmem_p_fragments): - # Publish one K32 probability fragment at a time so PV can - # consume early fragments while later scores are processed. - for fragment_idx in range(cfg.num_softmax_score_fragments): - s_arr = tmem_s0.load_softmax_p_fragment( - fragment_idx=fragment_idx, - s_arr=s_arr, - ) - smem_p0.compute_p_fragment( - fragment_idx=fragment_idx, + # One rolled loop streams every K32 probability fragment; the + # fragment body exists once in the instruction stream. + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): + smem_p0.compute_proxy_route_p_fragments( new_max_arr=new_max_arr, - s_arr=s_arr, + route_flags=sparse_route_flags, + route_origin0=sparse_origin0, + route_origin1=sparse_origin1, ) + else: + smem_p0.compute_p_fragments(new_max_arr=new_max_arr) else: # Wait for a free P stage before entering the ordered window so # BMM2 backpressure on this group's P pipeline cannot extend the @@ -2959,16 +2909,27 @@ def softmax0_schedule_body( # ProdWork: compute P=exp(S-new_max), store it in the profile's # SMEM or staged-TMEM operand, and record local sums for the # running softmax sum update. - smem_p0.compute_p( - new_max_arr=new_max_arr, - s_arr=s_arr, - ) # publishes the local denominator through tmem_s0 + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): + smem_p0.compute_proxy_route_p( + new_max_arr=new_max_arr, + s_arr=s_arr, + route_origin0=sparse_origin0, + route_origin1=sparse_origin1, + keeps_route_flags_or_swaps_origin2=sparse_route_flags, + swaps_route_origin3_bits=sparse_token_word0, + swaps_route_flags=sparse_token_word2, + ) + else: + smem_p0.compute_p( + new_max_arr=new_max_arr, + s_arr=s_arr, + ) # publishes the local denominator through tmem_s0 smem_p0.commit() if tmem_softmax_order is not None: tmem_softmax_order.release_softmax1() if cutlass.const_expr(cfg.use_keeps_mma_ab and cfg.uses_tmem_p): - # The TMEM-P store has consumed the aliased S columns, so the - # next QK wave can now overwrite them. + # TMEM-P has consumed the aliased S columns, so the next QK + # wave can now overwrite S. tmem_s0.release() # ProdWork: FP8 path applies the cross-resource sum correction # before TmemS.reduce_sums publishes the new running sums. @@ -3106,9 +3067,9 @@ def softmax1_schedule_body( sparse_softmax_metadata.init_read_state() with domain_loop(0, domain, 1, unroll=1) as d: - # ConsWait/ConsWork: load the second S instance and compute max. - tmem_s1.wait() if sparse_softmax_metadata is not None: + # Consume the independent metadata stream first so its SMEM + # loads and release can overlap the subsequent score wait. sparse_softmax_metadata.wait() # Copy to registers before release so the producer can reuse # the SMEM stage while this warp group applies the masks. @@ -3122,6 +3083,9 @@ def softmax1_schedule_body( sparse_token_word3, ) = sparse_softmax_metadata.load_route() sparse_softmax_metadata.release() + # ConsWait/ConsWork: load the second S instance and compute max. + tmem_s1.wait() + if sparse_softmax_metadata is not None: old_max_arr, sum_arr, new_max_arr, s_arr = ( tmem_s1.compute_block_sparse_softmax_loop( old_max_arr=old_max_arr, @@ -3157,16 +3121,15 @@ def softmax1_schedule_body( ) tmem_softmax_local1.commit() if cutlass.const_expr(cfg.streams_tmem_p_fragments): - for fragment_idx in range(cfg.num_softmax_score_fragments): - s_arr = tmem_s1.load_softmax_p_fragment( - fragment_idx=fragment_idx, - s_arr=s_arr, - ) - smem_p1.compute_p_fragment( - fragment_idx=fragment_idx, + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): + smem_p1.compute_proxy_route_p_fragments( new_max_arr=new_max_arr, - s_arr=s_arr, + route_flags=sparse_route_flags, + route_origin0=sparse_origin0, + route_origin1=sparse_origin1, ) + else: + smem_p1.compute_p_fragments(new_max_arr=new_max_arr) else: # Wait for a free P stage before entering the ordered window so # BMM2 backpressure on this group's P pipeline cannot extend the @@ -3175,7 +3138,18 @@ def softmax1_schedule_body( if tmem_softmax_order is not None: tmem_softmax_order.wait_softmax1() # ProdWork: compute and publish P1 for BMM2. - smem_p1.compute_p(new_max_arr=new_max_arr, s_arr=s_arr) + if cutlass.const_expr(cfg.use_block_sparse_proxy_routes): + smem_p1.compute_proxy_route_p( + new_max_arr=new_max_arr, + s_arr=s_arr, + route_origin0=sparse_origin0, + route_origin1=sparse_origin1, + keeps_route_flags_or_swaps_origin2=sparse_route_flags, + swaps_route_origin3_bits=sparse_token_word0, + swaps_route_flags=sparse_token_word2, + ) + else: + smem_p1.compute_p(new_max_arr=new_max_arr, s_arr=s_arr) smem_p1.commit() if tmem_softmax_order is not None: tmem_softmax_order.release_softmax0() @@ -3290,7 +3264,6 @@ def create_correction_task( tmem_corr0: MemoryResource, tmem_corr1: MemoryResource, work_queue: WorkQueue | None, - smem_kv_reuse_credit: MemoryResource | None, cfg: FmhaDecodeConfig, *, domain: int | cutlass.Int32, @@ -3301,9 +3274,6 @@ def create_correction_task( ) -> Task: """Create the two-instance correction and output task.""" - if smem_kv_reuse_credit is not None and work_queue is None: - raise ValueError("KV reuse credit requires a work queue") - def correction_schedule_body( tmem_softmax_local0: MemoryResource, tmem_softmax_local1: MemoryResource, @@ -3312,7 +3282,6 @@ def correction_schedule_body( tmem_corr1: MemoryResource, tmem_stats_done0: MemoryResource | None, tmem_stats_done1: MemoryResource | None, - smem_kv_reuse_credit: MemoryResource | None, ) -> None: """Schedule two-instance O correction and final output normalization.""" @@ -3506,37 +3475,17 @@ def correct_o( tail_o_stage_idx_1=tail_1, inst_idx=KV_INST1, ) - if smem_kv_reuse_credit is None: - tmem_corr1.correction_tail_epilogue( - o_stage_idx=o_stage_idx, - tail_o_stage_idx_0=tail_0, - tail_o_stage_idx_1=tail_1, - old_max_arr=old_max_arr, - new_max_arr=new_max_arr, - inst0_new_max_arr=inst0_new_max_arr, - inst0_sum_arr=inst0_sum_arr, - inst1_new_max_arr=inst1_new_max_arr, - inst1_sum_arr=inst1_sum_arr, - ) - else: - # The stage selector and ownership token share one pipeline epoch. - # Wait before the first aliased access and release immediately - # after correction stops touching the selected KV-ring stage. - smem_kv_reuse_credit.wait() - scratch_stage = smem_kv_reuse_credit.read_scratch_stage() - tmem_corr1.correction_tail_epilogue_rotating_exchange( - scratch_stage=scratch_stage, - o_stage_idx=o_stage_idx, - tail_o_stage_idx_0=tail_0, - tail_o_stage_idx_1=tail_1, - old_max_arr=old_max_arr, - new_max_arr=new_max_arr, - inst0_new_max_arr=inst0_new_max_arr, - inst0_sum_arr=inst0_sum_arr, - inst1_new_max_arr=inst1_new_max_arr, - inst1_sum_arr=inst1_sum_arr, - ) - smem_kv_reuse_credit.release() + tmem_corr1.correction_tail_epilogue( + o_stage_idx=o_stage_idx, + tail_o_stage_idx_0=tail_0, + tail_o_stage_idx_1=tail_1, + old_max_arr=old_max_arr, + new_max_arr=new_max_arr, + inst0_new_max_arr=inst0_new_max_arr, + inst0_sum_arr=inst0_sum_arr, + inst1_new_max_arr=inst1_new_max_arr, + inst1_sum_arr=inst1_sum_arr, + ) # Inst1 final reduction consumes both O0 and O1, so defer O0 release # until after inst1 has finished reading it. tmem_o.release() @@ -3550,7 +3499,6 @@ def run_correction_schedule( tmem_corr1: MemoryResource, tmem_stats_done0: MemoryResource | None, tmem_stats_done1: MemoryResource | None, - smem_kv_reuse_credit: MemoryResource | None, work_queue: WorkQueue | None, ) -> None: """Wrap correction with optional stats lifetime gates.""" @@ -3565,7 +3513,6 @@ def run_correction_schedule( tmem_corr1, tmem_stats_done0, tmem_stats_done1, - smem_kv_reuse_credit, ), ) @@ -3577,7 +3524,6 @@ def correction_schedule( tmem_corr0: MemoryResource, tmem_corr1: MemoryResource, work_queue: WorkQueue | None = None, - smem_kv_reuse_credit: MemoryResource | None = None, ) -> None: """Capture the Swaps correction schedule.""" run_correction_schedule( @@ -3588,7 +3534,6 @@ def correction_schedule( tmem_corr1, None, None, - smem_kv_reuse_credit, work_queue, ) @@ -3602,7 +3547,6 @@ def correction_keeps_schedule( tmem_stats_done0: MemoryResource, tmem_stats_done1: MemoryResource, work_queue: WorkQueue | None = None, - smem_kv_reuse_credit: MemoryResource | None = None, ) -> None: """Capture Keeps correction with explicit stats lifetime gates.""" run_correction_schedule( @@ -3613,7 +3557,6 @@ def correction_keeps_schedule( tmem_corr1, tmem_stats_done0, tmem_stats_done1, - smem_kv_reuse_credit, work_queue, ) @@ -3626,15 +3569,6 @@ def correction_keeps_schedule( tmem_corr0, tmem_corr1, ) - elif smem_kv_reuse_credit is None: - captured_schedule = correction_schedule( - tmem_softmax_local0, - tmem_softmax_local1, - tmem_o, - tmem_corr0, - tmem_corr1, - work_queue, - ) else: captured_schedule = correction_schedule( tmem_softmax_local0, @@ -3643,7 +3577,6 @@ def correction_keeps_schedule( tmem_corr0, tmem_corr1, work_queue, - smem_kv_reuse_credit, ) src = [tmem_softmax_local0, tmem_softmax_local1, tmem_o] else: @@ -3657,17 +3590,6 @@ def correction_keeps_schedule( tmem_stats_done0, tmem_stats_done1, ) - elif smem_kv_reuse_credit is None: - captured_schedule = correction_keeps_schedule( - tmem_softmax_local0, - tmem_softmax_local1, - tmem_o, - tmem_corr0, - tmem_corr1, - tmem_stats_done0, - tmem_stats_done1, - work_queue, - ) else: captured_schedule = correction_keeps_schedule( tmem_softmax_local0, @@ -3678,7 +3600,6 @@ def correction_keeps_schedule( tmem_stats_done0, tmem_stats_done1, work_queue, - smem_kv_reuse_credit, ) src = [ tmem_softmax_local0, @@ -3689,8 +3610,6 @@ def correction_keeps_schedule( ] if work_queue is not None: src.append(work_queue) - if smem_kv_reuse_credit is not None: - src.append(smem_kv_reuse_credit) return task_class( src_resources=src, dst_resources=[tmem_corr0, tmem_corr1], diff --git a/tensorrt_llm/_torch/attention/backends/sparse/hooks.py b/tensorrt_llm/_torch/attention/backends/sparse/hooks.py index 53cabb9101b1..44437c0a9037 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/hooks.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/hooks.py @@ -12,6 +12,8 @@ from importlib import import_module from typing import TYPE_CHECKING, Optional +from .skip_softmax import SkipSoftmaxParams + if TYPE_CHECKING: import torch @@ -220,26 +222,36 @@ def prepare_sparse_runtime_params( backend: "TrtllmAttention", q: "torch.Tensor", k: Optional["torch.Tensor"], + v: Optional["torch.Tensor"], metadata: "AttentionMetadata", forward_args: "AttentionForwardArgs", ) -> "SparseRuntimeParams": - """Run backend prediction hooks and update attention-op parameters.""" - runtime_params = forward_args.sparse_runtime_params - if backend.sparse_params is None: - return runtime_params - + """Predict all sparse inputs for one attention call. + + Runs the ``sparse_kv_predict``, ``sparse_attn_predict`` and + ``block_sparse_attn_predict`` hooks once each and returns a new + ``SparseRuntimeParams`` built from ``forward_args.sparse_runtime_params`` + plus the hook results. Fields a backend writes into that carrier outside + the hooks, such as an auxiliary pool pointer, are carried over. SkipSoftmax + backends receive their threshold schedule last. + """ kv_indices, kv_offsets = backend.sparse_kv_predict(q, k, metadata, forward_args) attn_indices, attn_offsets = backend.sparse_attn_predict(q, k, metadata, forward_args) - block_size = ( - backend.sparse_params.indices_block_size - if attn_indices is not None or attn_offsets is not None - else runtime_params.sparse_attn_indices_block_size - ) - return replace( - runtime_params, + block_sparse_inputs = backend.block_sparse_attn_predict(q, k, v, metadata, forward_args) + has_attn_indices = attn_indices is not None or attn_offsets is not None + sparse_params = backend.sparse_params + runtime_params = replace( + forward_args.sparse_runtime_params, sparse_kv_indices=kv_indices, sparse_kv_offsets=kv_offsets, sparse_attn_indices=attn_indices, sparse_attn_offsets=attn_offsets, - sparse_attn_indices_block_size=block_size, + sparse_attn_indices_block_size=sparse_params.indices_block_size if has_attn_indices else 0, + block_sparse_inputs=block_sparse_inputs, ) + if isinstance(sparse_params, SkipSoftmaxParams): + runtime_params = sparse_params.scheduler.get_runtime_params( + runtime_params=runtime_params, + timestep=forward_args.timestep, + ) + return runtime_params diff --git a/tensorrt_llm/_torch/attention/backends/sparse/params.py b/tensorrt_llm/_torch/attention/backends/sparse/params.py index 956fc8c6fdbd..8e6829f8c879 100644 --- a/tensorrt_llm/_torch/attention/backends/sparse/params.py +++ b/tensorrt_llm/_torch/attention/backends/sparse/params.py @@ -15,7 +15,7 @@ """Shared sparse attention parameter types.""" from dataclasses import dataclass -from typing import Optional +from typing import Literal, Optional import torch @@ -36,11 +36,56 @@ class SparseBackendForwardArgs: # Shared by algorithms that accept precomputed top-k indices. topk_indices: Optional[torch.Tensor] = None + # Complete block-sparse routing payload predicted by the module before the + # core forward; the default backend hook hands it through unchanged. + block_sparse_inputs: Optional["BlockSparseForwardInputs"] = None + + +@dataclass(frozen=True, slots=True) +class BlockSparseForwardInputs: + """Block geometry and live routing payload for one attention call. + + Exactly one routing representation is present. Canonical BSR uses + ``block_indptr`` and ``block_indices``; packed bitmask routing uses + ``exact_block_bits``. Paired K/V summaries enable proxy routes without + encoding an algorithm name in this shared carrier. + """ + + q_block_size: int + kv_block_size: int + max_blocks_per_row: Optional[int] = None + block_indptr: Optional[torch.Tensor] = None + block_indices: Optional[torch.Tensor] = None + exact_block_bits: Optional[torch.Tensor] = None + k_summary: Optional[torch.Tensor] = None + v_summary: Optional[torch.Tensor] = None + kv_valid_bits: Optional[torch.Tensor] = None + + def __post_init__(self) -> None: + has_bsr = self.block_indptr is not None + if has_bsr != (self.block_indices is not None): + raise ValueError("block_indptr and block_indices must be provided together") + if has_bsr == (self.exact_block_bits is not None): + raise ValueError("exactly one route representation must be provided") + if has_bsr and self.max_blocks_per_row is None: + raise ValueError("BSR routes require max_blocks_per_row") + if (self.k_summary is None) != (self.v_summary is None): + raise ValueError("k_summary and v_summary must be provided together") + + @property + def sparse_format(self) -> Literal["bsr", "bitmask"]: + """Routing representation selected by the live payload.""" + return "bitmask" if self.exact_block_bits is not None else "bsr" + + @property + def use_proxy_routes(self) -> bool: + """Whether unselected blocks are represented by K/V summaries.""" + return self.k_summary is not None @dataclass(kw_only=True, slots=True) class SparseRuntimeParams: - """Flat optional sparse inputs passed from a backend to ``AttentionOp``.""" + """Complete per-attention sparse runtime state consumed by FMHA/``AttentionOp``.""" # Sparse index inputs shared by multiple algorithms. sparse_kv_indices: Optional[torch.Tensor] = None @@ -57,3 +102,13 @@ class SparseRuntimeParams: threshold_scale_factor_prefill: float = 0.0 # SkipSoftmax decode threshold; diffusion models leave it at zero. threshold_scale_factor_decode: float = 0.0 + block_sparse_inputs: Optional[BlockSparseForwardInputs] = None + + +__all__ = [ + "BlockSparseForwardInputs", + "SparseBackendForwardArgs", + "SparseMetadataParams", + "SparseParams", + "SparseRuntimeParams", +] diff --git a/tensorrt_llm/_torch/attention/backends/trtllm.py b/tensorrt_llm/_torch/attention/backends/trtllm.py index 914f6f3a654e..d5e273b06483 100644 --- a/tensorrt_llm/_torch/attention/backends/trtllm.py +++ b/tensorrt_llm/_torch/attention/backends/trtllm.py @@ -17,7 +17,7 @@ import math import os import weakref -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from typing import TYPE_CHECKING, List, Optional, Tuple import torch @@ -46,8 +46,7 @@ PredefinedAttentionMask, RopeParams, merge_attention_forward_args) from .sparse.hooks import prepare_sparse_runtime_params -from .sparse.params import SparseParams -from .sparse.skip_softmax import SkipSoftmaxParams +from .sparse.params import BlockSparseForwardInputs, SparseParams _SKIP_CORRECTION_SUPPORTED_SMS = frozenset((100, 103)) @@ -1928,8 +1927,13 @@ def forward( seq_start=num_ctx, ) - forward_args.sparse_runtime_params = prepare_sparse_runtime_params( - self, q, k, metadata, forward_args) + sparse_runtime_params = prepare_sparse_runtime_params( + self, q, k, v, metadata, forward_args) + forward_args = replace( + forward_args, + sparse_runtime_params=sparse_runtime_params, + ) + has_block_sparse_inputs = sparse_runtime_params.block_sparse_inputs is not None # Compute FlashMLA tile-scheduler metadata once per forward pass. # The flag is invalidated whenever FlashMLA inputs change. The metadata @@ -1982,11 +1986,11 @@ def forward( assert k.shape[1] == kv_hidden_size assert v.shape[1] == kv_hidden_size num_tokens = q.shape[0] - if k is not None and not metadata.is_cross: + if k is not None and not metadata.is_cross and not has_block_sparse_inputs: assert k.shape[0] == num_tokens assert v.shape[0] == num_tokens else: - sparse_attn_indices = forward_args.sparse_runtime_params.sparse_attn_indices + sparse_attn_indices = sparse_runtime_params.sparse_attn_indices is_sparse_attn = sparse_attn_indices is not None and sparse_attn_indices.numel( ) > 0 if attention_input_type == AttentionInputType.context_only and is_sparse_attn: @@ -2059,14 +2063,6 @@ def forward( if forward_args.kv_scale_quant_orig is None: forward_args.kv_scale_quant_orig = self.kv_scale_quant_orig - sparse_params = self.sparse_params - if isinstance(sparse_params, SkipSoftmaxParams): - forward_args.sparse_runtime_params = ( - sparse_params.scheduler.get_runtime_params( - runtime_params=forward_args.sparse_runtime_params, - timestep=forward_args.timestep, - )) - # max_context_q_len_override is only set when encoder CUDA graphs are enabled. if metadata.max_context_q_len_override is not None: assert metadata.is_cuda_graph @@ -2286,6 +2282,26 @@ def sparse_kv_predict( """Predict sparse KV indices when required by an algorithm.""" return None, None + def block_sparse_attn_predict( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: TrtllmAttentionMetadata, + forward_args: AttentionForwardArgs, + ) -> Optional[BlockSparseForwardInputs]: + """Predict the block-sparse routing payload for one attention call. + + The default hands through routes that the attention module predicted + before the core forward via ``sparse_backend_args``. Algorithms that + predict inside the backend override this method and return ``None`` + for dense phases. + """ + backend_args = forward_args.sparse_backend_args + if backend_args is None: + return None + return backend_args.block_sparse_inputs + def sparse_attn_predict( self, q: torch.Tensor, diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py b/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py index b919f17c0a0c..f6efaffa171a 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/__init__.py @@ -20,15 +20,7 @@ simplified metadata that doesn't require KV caching. """ -from .cute_dsl import ( - VSA_TILE_SIZE, - CuTeDSLAttention, - VSAAttention, - VSAMetadata, - VSAMetadataBuilder, - get_vsa_forward_context, - set_vsa_forward_context, -) +from .cute_dsl import CuTeDSLAttention from .flash_attn4 import FlashAttn4Attention from .flashinfer import FlashInferAttention from .interface import AttentionBackend, AttentionTensorLayout @@ -38,7 +30,6 @@ from .vanilla import VanillaAttention __all__ = [ - "VSA_TILE_SIZE", "Attention2DAttention", "AttentionBackend", "AttentionTensorLayout", @@ -49,13 +40,8 @@ "TrtllmAttention", "TrtllmAttentionMetadata", "UlyssesAttention", - "VSAAttention", - "VSAMetadata", - "VSAMetadataBuilder", "VanillaAttention", "create_attention", "get_visual_gen_attention_backend", - "get_vsa_forward_context", - "set_vsa_forward_context", "wrap_parallel_attention", ] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/__init__.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/__init__.py index 9b70421c3b81..1dde6e1f9aab 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/__init__.py @@ -15,31 +15,12 @@ """ CuTe DSL attention backend family for visual generation models. - fmha.py — CuTeDSLAttention (dense and blockscaled JIT FMHA) - vsa.py — VSAAttention (Video Sparse Attention, CuTe JIT + SDPA fallback) + fmha.py — CuTeDSLAttention (dense and blockscaled JIT FMHA) """ from .fmha import CuTeDSLAttention, _cute_dsl_import_error -from .vsa import ( - VSA_KERNEL_MAX_CUBES, - VSA_TILE_SIZE, - VSAAttention, - VSAMetadata, - VSAMetadataBuilder, - VSAPreprocessor, - get_vsa_forward_context, - set_vsa_forward_context, -) __all__ = [ "CuTeDSLAttention", - "VSAAttention", - "VSAMetadata", - "VSAMetadataBuilder", - "VSAPreprocessor", - "VSA_TILE_SIZE", - "VSA_KERNEL_MAX_CUBES", - "set_vsa_forward_context", - "get_vsa_forward_context", "_cute_dsl_import_error", ] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py index e9638ef9075c..a7bbc7baa2cf 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/fmha.py @@ -17,7 +17,7 @@ JIT-compiles dense or SkipSoftmax FMHA and caches the compiled artifact for each kernel configuration. Expects NHD layout ([B, S, H, D]) and supports float16/bfloat16 inputs. The VSA -sparse path uses VSAAttention from vsa.py instead. +sparse backend uses `VSACuTeDSLAttention` in `attention_backend.sparse.vsa` instead. """ import math @@ -72,20 +72,6 @@ def _resolve_skip_softmax_threshold_scale_factor( ) -> float | None: """Resolve the active CuTeDSL threshold for the current denoising phase.""" if sparse_params is not None: - if timestep is None and sparse_params.scheduler.disabled_until_timestep is not None: - # Fail-open: a missing timestep resolves to the full (unthrottled) - # threshold, i.e. skip-softmax runs during the high-noise steps - # `disabled_until_timestep` exists to protect. This is silent - # elsewhere (a quality regression, not an error), so surface it - # once per process instead of only in this function's return value. - logger.warning_once( - "SkipSoftmax scheduler has disabled_until_timestep=" - f"{sparse_params.scheduler.disabled_until_timestep} configured, but no " - "`timestep` was passed to the CuTeDSL attention forward call. Skip-softmax " - "will run unthrottled (as if past the cutoff) until `timestep` is threaded " - "through.", - key="cute_dsl_skip_softmax_missing_timestep", - ) runtime_params = sparse_params.scheduler.get_runtime_params(timestep=timestep) threshold_scale_factor = runtime_params.threshold_scale_factor_prefill if threshold_scale_factor is None or threshold_scale_factor <= 0.0: diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/vsa.py b/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/vsa.py deleted file mode 100644 index 741f82097ca2..000000000000 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/cute_dsl/vsa.py +++ /dev/null @@ -1,412 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -""" -Video Sparse Attention (VSA) backend for visual generation models. - -VSAAttention implements hierarchical sparse attention: - - Coarse branch: mean-pooled cube attention (always dense) - - Fine branch: block-sparse top-K attention via CuTe JIT kernel (sm100+) - or dense SDPA fallback when CuTe is unavailable / head_dim != 128. -""" - -import contextvars -from contextlib import contextmanager -from dataclasses import dataclass -from math import ceil -from typing import Dict, Optional, Tuple - -import torch -import torch.nn.functional as F - -from ..interface import AttentionBackend, AttentionTensorLayout - -_vsa_import_error = None -try: - from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.video_sparse_attention import ( - block_sparse_attn_from_indices_cute, - is_cute_supported, - ) -except (ImportError, OSError) as e: - block_sparse_attn_from_indices_cute = None - is_cute_supported = None - _vsa_import_error = e - - -# Must match the Blackwell kernel's block_size expectation. -VSA_TILE_SIZE: Tuple[int, int, int] = (4, 4, 4) - -# Kernel's SMEM buffer for variable_block_sizes is fixed-size and unchecked, -# so num_cubes must stay <= this. -VSA_KERNEL_MAX_CUBES: int = 4 * 1024 - - -def _get_tile_partition_indices( - dit_seq_shape: Tuple[int, int, int], - tile_size: Tuple[int, int, int], - device: torch.device, -) -> torch.LongTensor: - T, H, W = dit_seq_shape - tT, tH, tW = tile_size - nT, nH, nW = ceil(T / tT), ceil(H / tH), ceil(W / tW) - - bt = torch.arange(nT, device=device).view(nT, 1, 1, 1, 1, 1) - bh = torch.arange(nH, device=device).view(1, nH, 1, 1, 1, 1) - bw = torch.arange(nW, device=device).view(1, 1, nW, 1, 1, 1) - lt = torch.arange(tT, device=device).view(1, 1, 1, tT, 1, 1) - lh = torch.arange(tH, device=device).view(1, 1, 1, 1, tH, 1) - lw = torch.arange(tW, device=device).view(1, 1, 1, 1, 1, tW) - - gt = bt * tT + lt - gh = bh * tH + lh - gw = bw * tW + lw - valid = (gt < T) & (gh < H) & (gw < W) - flat = gt * (H * W) + gh * W + gw - out = torch.where(valid, flat, torch.full_like(flat, -1)) - return out.reshape(-1).to(torch.long) - - -def _construct_variable_block_sizes( - dit_seq_shape: Tuple[int, int, int], - num_tiles: Tuple[int, int, int], - tile_size: Tuple[int, int, int], - device: torch.device, -) -> torch.LongTensor: - T, H, W = dit_seq_shape - tT, tH, tW = tile_size - nT, nH, nW = num_tiles - - bt = torch.arange(nT, device=device) - bh = torch.arange(nH, device=device) - bw = torch.arange(nW, device=device) - valid_t = (T - bt * tT).clamp(max=tT) - valid_h = (H - bh * tH).clamp(max=tH) - valid_w = (W - bw * tW).clamp(max=tW) - sizes = valid_t.view(nT, 1, 1) * valid_h.view(1, nH, 1) * valid_w.view(1, 1, nW) - return sizes.reshape(-1).to(torch.long) - - -@dataclass -class VSAMetadata: - """Per-timestep metadata required by the VSA sparse path.""" - - current_timestep: int - dit_seq_shape: Tuple[int, int, int] - vsa_sparsity: float - num_tiles: Tuple[int, int, int] - total_seq_length: int - padded_seq_length: int - tile_partition_indices: torch.LongTensor - reverse_tile_partition_indices: torch.LongTensor - variable_block_sizes: torch.LongTensor - non_pad_index: torch.LongTensor - gather_idx: torch.LongTensor - - -class VSAMetadataBuilder: - """Builds VSAMetadata; caches per-shape index tensors so torch.compile - guards stay stable across denoising steps.""" - - def __init__(self) -> None: - self._cache: Dict[Tuple[Tuple[int, int, int], str], Dict[str, object]] = {} - - def _build_shape_payload( - self, - dit_seq_shape: Tuple[int, int, int], - device: torch.device, - ) -> Dict[str, object]: - T, H, W = dit_seq_shape - tT, tH, tW = VSA_TILE_SIZE - num_tiles = (ceil(T / tT), ceil(H / tH), ceil(W / tW)) - total_seq_length = T * H * W - padded_seq_length = num_tiles[0] * num_tiles[1] * num_tiles[2] * tT * tH * tW - - tile_partition_indices = _get_tile_partition_indices(dit_seq_shape, VSA_TILE_SIZE, device) - non_pad_index = (tile_partition_indices >= 0).nonzero(as_tuple=True)[0] - gather_idx = tile_partition_indices[non_pad_index] - - reverse = torch.zeros(total_seq_length, dtype=torch.long, device=device) - reverse[gather_idx] = torch.arange(len(non_pad_index), dtype=torch.long, device=device) - - variable_block_sizes = _construct_variable_block_sizes( - dit_seq_shape, num_tiles, VSA_TILE_SIZE, device - ) - - return { - "dit_seq_shape": dit_seq_shape, - "num_tiles": num_tiles, - "total_seq_length": total_seq_length, - "padded_seq_length": padded_seq_length, - "tile_partition_indices": tile_partition_indices, - "reverse_tile_partition_indices": reverse, - "variable_block_sizes": variable_block_sizes, - "non_pad_index": non_pad_index, - "gather_idx": gather_idx, - } - - def build( - self, - current_timestep: int, - raw_latent_shape: Tuple[int, int, int], - patch_size: Tuple[int, int, int], - vsa_sparsity: float, - device: torch.device, - ) -> VSAMetadata: - dit_seq_shape = ( - raw_latent_shape[0] // patch_size[0], - raw_latent_shape[1] // patch_size[1], - raw_latent_shape[2] // patch_size[2], - ) - cache_key = (dit_seq_shape, str(device)) - payload = self._cache.get(cache_key) - if payload is None: - payload = self._build_shape_payload(dit_seq_shape, device) - self._cache[cache_key] = payload - - return VSAMetadata( - current_timestep=current_timestep, - vsa_sparsity=vsa_sparsity, - **payload, # type: ignore[arg-type] - ) - - -_vsa_forward_context_var: contextvars.ContextVar[Optional[VSAMetadata]] = contextvars.ContextVar( - "_vsa_forward_context", default=None -) - - -@contextmanager -def set_vsa_forward_context(metadata: VSAMetadata): - token = _vsa_forward_context_var.set(metadata) - try: - yield - finally: - _vsa_forward_context_var.reset(token) - - -def get_vsa_forward_context() -> Optional[VSAMetadata]: - return _vsa_forward_context_var.get(None) - - -def _mean_pool_cubes( - x_tiled: torch.Tensor, - variable_block_sizes: torch.LongTensor, - prod_tile: int, - num_cubes: int, -) -> torch.Tensor: - B, _padded, H, D = x_tiled.shape - x_cubes = x_tiled.view(B, num_cubes, prod_tile, H, D) - # fp32 accumulation: bf16 sum over 64 tokens perturbs the coarse softmax. - x_sum = x_cubes.float().sum(dim=2) - valid_counts = variable_block_sizes.float().clamp(min=1).view(1, num_cubes, 1, 1) - return (x_sum / valid_counts).to(x_tiled.dtype) - - -class VSAPreprocessor: - """Reorders NHD tokens into tile-major layout and zero-pads to tile boundaries.""" - - @staticmethod - def tile( - x: torch.Tensor, - non_pad_index: torch.LongTensor, - gather_idx: torch.LongTensor, - padded_seq_len: int, - ) -> torch.Tensor: - # index_select + index_copy_ instead of chained advanced indexing so - # torch.compile can trace this without a graph break. - B, _S, H, D = x.shape - x_valid = x.index_select(1, gather_idx) - x_padded = x.new_zeros(B, padded_seq_len, H, D) - x_padded.index_copy_(1, non_pad_index, x_valid) - return x_padded - - @staticmethod - def untile( - x: torch.Tensor, - reverse_tile_partition_indices: torch.LongTensor, - non_pad_index: torch.LongTensor, - ) -> torch.Tensor: - return x.index_select(1, non_pad_index).index_select(1, reverse_tile_partition_indices) - - -class VSAAttention(AttentionBackend): - """ - Video Sparse Attention (VSA) backend for diffusion models. - - Implements coarse mean-pool + fine block-sparse top-K attention. - The fine branch uses a JIT-compiled CuTe kernel on sm100+ for - head_dim=128 / fp16-bf16; otherwise falls back to dense SDPA. - - Requires an active VSA forward context (set_vsa_forward_context) during - each forward call. Does not support LSE output. - """ - - def __init__( - self, - layer_idx: int = 0, - num_heads: int = 8, - head_dim: int = 128, - num_kv_heads: Optional[int] = None, - dtype: Optional[torch.dtype] = None, - sparse_attention_config=None, - **kwargs, - ): - self.layer_idx = layer_idx - self.num_heads = num_heads - self.head_dim = head_dim - self.num_kv_heads = num_kv_heads or num_heads - assert self.num_kv_heads == self.num_heads, ( - f"VSA coarse mean-pool assumes MHA (num_kv_heads == num_heads), " - f"got num_kv_heads={self.num_kv_heads}, num_heads={self.num_heads}. " - f"GQA/MQA is not supported." - ) - self.dtype = dtype - self.sparse_attention_config = sparse_attention_config - - # Dynamo can't guard on the module-level mutable global, so this read - # runs in eager. - @torch.compiler.disable - def _get_vsa_inputs(self): - ctx: Optional[VSAMetadata] = get_vsa_forward_context() - if ctx is None: - raise RuntimeError( - "VSAAttention.forward called without an active VSA forward context. " - "Wrap each transformer call with set_vsa_forward_context()." - ) - return ( - ctx.non_pad_index, - ctx.gather_idx, - ctx.reverse_tile_partition_indices, - ctx.variable_block_sizes, - ctx.padded_seq_length, - ctx.num_tiles[0] * ctx.num_tiles[1] * ctx.num_tiles[2], - ctx.vsa_sparsity, - ) - - def forward( - self, - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - *, - gate_compress: Optional[torch.Tensor] = None, - gate_fine: Optional[torch.Tensor] = None, - **kwargs, - ) -> torch.Tensor: - """ - VSA forward: coarse mean-pool + fine block-sparse top-K. - - Args: - q, k, v: [B, S, H, D] in original (un-tiled) token order. - gate_compress: [B, S, H, D] G_c gate weighting the coarse branch O_c. - gate_fine: Optional [B, S, H, D] G_f gate weighting the fine branch - O_f. None means constant 1 (dense behavior preserved). - - Returns: - [B, S, H, D] in the same original token order. - """ - if gate_compress is None: - raise ValueError( - "VSAAttention requires gate_compress. " - "Ensure to_gate_compress is wired in the transformer block." - ) - - ( - non_pad_index, - gather_idx, - reverse_tile_partition_indices, - variable_block_sizes, - padded_len, - num_cubes, - vsa_sparsity, - ) = self._get_vsa_inputs() - - B, S, H, D = q.shape - prod_tile = VSA_TILE_SIZE[0] * VSA_TILE_SIZE[1] * VSA_TILE_SIZE[2] - cur_topk = max(1, ceil((1.0 - vsa_sparsity) * num_cubes)) - - q_t = VSAPreprocessor.tile(q, non_pad_index, gather_idx, padded_len) - k_t = VSAPreprocessor.tile(k, non_pad_index, gather_idx, padded_len) - v_t = VSAPreprocessor.tile(v, non_pad_index, gather_idx, padded_len) - - q_c = _mean_pool_cubes(q_t, variable_block_sizes, prod_tile, num_cubes) - k_c = _mean_pool_cubes(k_t, variable_block_sizes, prod_tile, num_cubes) - v_c = _mean_pool_cubes(v_t, variable_block_sizes, prod_tile, num_cubes) - - scale = D**-0.5 - scores_c = torch.einsum("bnhd,bmhd->bhnm", q_c, k_c) * scale - attn_probs_c = scores_c.softmax(dim=-1) - o_c = torch.einsum("bhnm,bmhd->bnhd", attn_probs_c, v_c) - - use_cute = ( - _vsa_import_error is None - and is_cute_supported(q) - and (q.dtype == k.dtype == v.dtype) - and num_cubes <= VSA_KERNEL_MAX_CUBES - ) - topk_indices = attn_probs_c.topk(cur_topk, dim=-1).indices.to(torch.int32) - - o_c_tiled = ( - o_c.unsqueeze(2).expand(B, num_cubes, prod_tile, H, D).reshape(B, padded_len, H, D) - ) - - if use_cute: - q_hnd = q_t.transpose(1, 2).contiguous() - k_hnd = k_t.transpose(1, 2).contiguous() - v_hnd = v_t.transpose(1, 2).contiguous() - q2k_num = torch.full((B, H, num_cubes), cur_topk, dtype=torch.int32, device=q.device) - o_hnd, _lse = block_sparse_attn_from_indices_cute( - q_hnd, - k_hnd, - v_hnd, - q2k_idx=topk_indices.contiguous(), - q2k_num=q2k_num, - variable_block_sizes=variable_block_sizes.to(torch.int32), - ) - o_f_tiled = o_hnd.transpose(1, 2) - - # Padded rows hold kernel garbage; zero-padded gates mask the coarse - # term and untile discards padded positions from both branches. - gate_c_t = VSAPreprocessor.tile(gate_compress, non_pad_index, gather_idx, padded_len) - if gate_fine is not None: - gate_f_t = VSAPreprocessor.tile(gate_fine, non_pad_index, gather_idx, padded_len) - combined_tiled = gate_c_t * o_c_tiled + gate_f_t * o_f_tiled - else: - combined_tiled = gate_c_t * o_c_tiled + o_f_tiled - return VSAPreprocessor.untile( - combined_tiled, reverse_tile_partition_indices, non_pad_index - ) - - # SDPA must run on the un-tiled Q/K/V — padded zero K/V slots would - # otherwise absorb softmax mass and pollute the output. Untile o_c so - # both branches combine in original-flat order. - o_c_full = VSAPreprocessor.untile(o_c_tiled, reverse_tile_partition_indices, non_pad_index) - o_f = F.scaled_dot_product_attention( - q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) - ).transpose(1, 2) - if gate_fine is not None: - return gate_compress * o_c_full + gate_fine * o_f - return gate_compress * o_c_full + o_f - - @classmethod - def support_lse(cls) -> bool: - return False - - @property - def preferred_layout(self) -> AttentionTensorLayout: - return AttentionTensorLayout.NHD - - @classmethod - def support_fused_qkv(cls) -> bool: - return False diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py b/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py index 5024a43b192a..a441ce7edac7 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/flash_attn4.py @@ -39,10 +39,64 @@ def _install_cutlass_dsl_compatibility() -> None: cute.make_fragment = cute.make_rmem_tensor +def _install_flash_attn_tile_scheduler_compatibility() -> None: + """Keep FA4's four-axis ``WorkTileInfo`` independent of CUTLASS task scheduling. + + Importing ``cutlass.experimental.task_scheduling`` rewrites the shared + ``cutlass.utils.WorkTileInfo`` class in place: its constructor unpacks + ``tile_idx`` into exactly three scalars and ``tile_idx`` / ``is_valid_tile`` + become properties over those scalars. The vendored PrimTS kernels import + that package, so once any PrimTS FMHA has been probed or planned in a + process, every later FA4 kernel trace fails with ``ValueError: too many + values to unpack (expected 3)``: FA4 subclasses the same CUTLASS class with + a (block, head, batch, split) coordinate but does not define its own + constructor. Installing the upstream tuple semantics directly on the FA4 + subclass makes it immune to the parent rewrite regardless of import order. + Remove once CUTLASS stops patching the shared class or FA4 owns these + members itself. + """ + try: + from flash_attn.cute import tile_scheduler + except (ImportError, OSError): + return + import cutlass.cute as cute + from cutlass.cutlass_dsl import Boolean, extract_mlir_values + + work_tile_info = tile_scheduler.WorkTileInfo + if "__init__" in vars(work_tile_info): + return + + def __init__(self, tile_idx: cute.Coord, is_valid_tile: Boolean) -> None: + self._tile_idx = tile_idx + self._is_valid_tile = Boolean(is_valid_tile) + self._tile_idx_num_values = None + + def __extract_mlir_values__(self) -> list: + tile_idx_values = extract_mlir_values(self._tile_idx) + valid_values = extract_mlir_values(self._is_valid_tile) + self._tile_idx_num_values = len(tile_idx_values) + return tile_idx_values + valid_values + + @cute.jit + def tile_idx(self) -> cute.Coord: + return self._tile_idx + + @cute.jit + def is_valid_tile(self) -> Boolean: + return self._is_valid_tile + + work_tile_info.__init__ = __init__ + work_tile_info.__extract_mlir_values__ = __extract_mlir_values__ + work_tile_info.tile_idx = property(tile_idx) + work_tile_info.is_valid_tile = property(is_valid_tile) + + _flash_attn_fwd_import_error = None try: _install_cutlass_dsl_compatibility() from flash_attn.cute.interface import _flash_attn_fwd + + _install_flash_attn_tile_scheduler_compatibility() except (ImportError, OSError) as e: _flash_attn_fwd = None _flash_attn_fwd_import_error = e diff --git a/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..e4aa9431c1e9 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/__init__.py @@ -0,0 +1,19 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""TRT-LLM-owned predictor for two-stage SOL attention.""" + +from .predictor import ( + SolPredictorGeometry, + SolPredictorOutputs, + SolPredictorPlan, + SolPredictorPlanKey, + SOLSparsePredictor, +) + +__all__ = [ + "SOLSparsePredictor", + "SolPredictorGeometry", + "SolPredictorOutputs", + "SolPredictorPlan", + "SolPredictorPlanKey", +] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/backend.py b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/backend.py new file mode 100644 index 000000000000..270151266d02 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/backend.py @@ -0,0 +1,150 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""VisualGen SOL attention using the generic TRTLLM sparse lifecycle.""" + +from __future__ import annotations + +from typing import Optional + +import torch + +from tensorrt_llm._torch.attention.backends.fmha.prims_ts_block_sparse import PrimsTSBlockSparseFmha +from tensorrt_llm._torch.attention.backends.fmha.utils import get_bmm1_scale +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + PredefinedAttentionMask, +) +from tensorrt_llm._torch.attention.backends.sparse.params import BlockSparseForwardInputs +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata + +from ...trtllm import TrtllmAttention +from .params import SolParams +from .predictor import BLOCK_SIZE, SOLSparsePredictor + + +class SOLTrtllmAttention(TrtllmAttention): + """Predict SOL routes inside the core prediction hook, then execute them + through the generic block-sparse FMHA.""" + + def __init__(self, *, sparse_params: SolParams | None = None, **kwargs) -> None: + if not isinstance(sparse_params, SolParams): + raise TypeError("SOLTrtllmAttention requires SolParams") + self.sol_params = sparse_params + self._prepared_graph_phase: int | None = None + super().__init__(sparse_params=None, **kwargs) + attention_metadata_state = kwargs["attention_metadata_state"] + predictor_cache = attention_metadata_state.setdefault("sparse_predictors", {}) + predictor = predictor_cache.get("sol_attn") + if predictor is None: + predictor = SOLSparsePredictor() + predictor_cache["sol_attn"] = predictor + elif not isinstance(predictor, SOLSparsePredictor): + raise TypeError("model-scoped SOL predictor cache contains an invalid value") + self.predictor = predictor + + def _resolve_graph_phase(self, timestep: object) -> int | None: + """Resolve the dense-or-sparse phase, reusing the warmup value under capture.""" + + if self.sol_params.disabled_until_timestep is None: + return None + if torch.cuda.is_current_stream_capturing(): + if self._prepared_graph_phase is None: + raise RuntimeError("SOL graph phase must be prepared before CUDA Graph capture") + return self._prepared_graph_phase + graph_phase = self.sol_params.get_graph_phase_for_timestep( + timestep, + disabled_until_timestep=self.sol_params.disabled_until_timestep, + ) + self._prepared_graph_phase = graph_phase + return graph_phase + + def block_sparse_attn_predict( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + metadata: TrtllmAttentionMetadata, + forward_args: AttentionForwardArgs, + ) -> BlockSparseForwardInputs | None: + """Return SOL routes for sparse calls and ``None`` for dense calls. + + ``q``, ``k``, and ``v`` arrive in the flattened ``[B*S, H*D]`` core + layout; the batch layout comes from ``metadata`` and the timestep from + ``forward_args``. + """ + + timestep = forward_args.timestep + graph_phase = self._resolve_graph_phase(timestep) + if not self.sol_params.should_use_sparse( + layer_idx=self.layer_idx, + timestep=timestep, + graph_phase=graph_phase, + ): + return None + + if self.quant_attention_config is not None: + raise ValueError("SOL sparse execution does not support quant_attention_config") + if not any( + isinstance(fmha, PrimsTSBlockSparseFmha) for fmha in self._fmha_manager.fmha_libs + ): + raise RuntimeError("SOL sparse execution requires PrimTS block-sparse FMHA") + if forward_args.attention_mask != PredefinedAttentionMask.FULL: + raise ValueError("SOL sparse execution requires a full attention mask") + if k is None or v is None: + raise ValueError("SOL sparse execution requires separate q, k, and v tensors") + + batch_size = metadata.num_seqs + seq_len = metadata.max_seq_len + num_tokens = batch_size * seq_len + if q.shape[0] != num_tokens or k.shape[0] != num_tokens or v.shape[0] != num_tokens: + raise ValueError( + "SOL sparse execution supports only uniform-length self-attention; " + f"got {q.shape[0]} query and {k.shape[0]} key tokens for " + f"{batch_size} sequences of length {seq_len}" + ) + + # The VisualGen wrapper compacts the flattened tensors once; these views + # are shared between prediction and the generic block-sparse FMHA. + q = q.view(batch_size, seq_len, self.num_heads, self.head_dim) + k = k.view(batch_size, seq_len, self.num_kv_heads, self.head_dim) + v = v.view(batch_size, seq_len, self.num_kv_heads, self.head_dim) + unsupported_reason = self.predictor.support_reason(q, k, v) + if unsupported_reason is not None: + raise ValueError(unsupported_reason) + + outputs = self.predictor.predict( + q, + k, + v, + tau=self.sol_params.tau, + sm_scale=get_bmm1_scale(self), + ) + return BlockSparseForwardInputs( + q_block_size=BLOCK_SIZE, + kv_block_size=BLOCK_SIZE, + exact_block_bits=outputs.exact_block_bits, + k_summary=outputs.k_summary, + v_summary=outputs.v_summary, + ) + + @classmethod + def support_fused_qkv(cls) -> bool: + """SOL prediction requires separate Q, K, and V tensors.""" + + return False + + +__all__ = ["SOLTrtllmAttention"] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/kernels.py b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/kernels.py new file mode 100644 index 000000000000..139a034b4aeb --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/kernels.py @@ -0,0 +1,507 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Memory-bound kernels of the two-stage SOL predictor. + +The predictor summarises ``[batch, tokens, heads, head_dim]`` activations per token block, derives per +channel key statistics, and thresholds centroid scores into packed exact-block words. CUDA tensors run +Triton kernels; other tensors use PyTorch implementations of the same rule. Launch shapes are derived +from tensor shapes, so no autotuning happens at call time and every launch is CUDA Graph safe. Every +helper writes into caller-owned storage so a plan can keep graph-stable outputs. +""" + +from __future__ import annotations + +import math +from typing import Literal + +import torch +import triton +import triton.language as tl + +_POOL_MAX_WIDTH = 1024 +_POOL_TOKENS_PER_LOAD = 8 +_STATS_WIDTH = 128 +_STATS_ROWS_PER_LOAD = 32 +_SELECT_Q_BLOCKS = 64 +_WORD_BITS = 32 +_LOG2_E = math.log2(math.e) +_THRESHOLD_EPSILON = 1.0e-6 +_LOCAL_RADIUS = 1 + +_Reduce = Literal["mean", "sum"] + + +def _column_launch(row_width: int, max_width: int) -> tuple[int, int]: + """Columns per program and number of column chunks for a row of ``row_width`` channels.""" + width = min(max_width, triton.next_power_of_2(row_width)) + return width, triton.cdiv(row_width, width) + + +def num_blocks(seq_len: int, block_size: int) -> int: + return (seq_len + block_size - 1) // block_size + + +def num_words(num_kv_blocks: int) -> int: + return (num_kv_blocks + _WORD_BITS - 1) // _WORD_BITS + + +# --------------------------------------------------------------------------- block pooling +@triton.jit +def _block_pool_kernel( + x_ptr, + out_ptr, + seq_len, + num_blocks, + num_chunks, + row_width, + stride_x_batch, + stride_x_token, + stride_out_batch, + stride_out_block, + MEAN: tl.constexpr, + BLOCK: tl.constexpr, + TOKENS: tl.constexpr, + WIDTH: tl.constexpr, +): + """One program per (batch, block, column chunk): fp32 sum over the block's valid tokens.""" + pid = tl.program_id(0).to(tl.int64) + chunk = pid % num_chunks + batch_block = pid // num_chunks + block = batch_block % num_blocks + batch = batch_block // num_blocks + columns = chunk * WIDTH + tl.arange(0, WIDTH) + in_row = columns < row_width + first_token = block * BLOCK + total = tl.zeros([WIDTH], dtype=tl.float32) + for start in range(0, BLOCK, TOKENS): + tokens = first_token + start + tl.arange(0, TOKENS) + values = tl.load( + x_ptr + batch * stride_x_batch + tokens[:, None] * stride_x_token + columns[None, :], + mask=(tokens < seq_len)[:, None] & in_row[None, :], + other=0.0, + ) + total += tl.sum(values.to(tl.float32), axis=0) + if MEAN: + total = total / tl.minimum(seq_len - first_token, BLOCK).to(tl.float32) + tl.store( + out_ptr + batch * stride_out_batch + block * stride_out_block + columns, + total.to(out_ptr.dtype.element_ty), + mask=in_row, + ) + + +def _block_pool_torch( + x: torch.Tensor, out: torch.Tensor, *, block_size: int, reduce: _Reduce +) -> None: + batch_size, seq_len, num_heads, head_dim = x.shape + blocks = num_blocks(seq_len, block_size) + padded = torch.nn.functional.pad(x, (0, 0, 0, 0, 0, blocks * block_size - seq_len)) + total = padded.view(batch_size, blocks, block_size, num_heads, head_dim).sum( + dim=2, dtype=torch.float32 + ) + if reduce == "mean": + valid = torch.clamp( + seq_len - torch.arange(blocks, device=x.device) * block_size, max=block_size + ) + total = total / valid.to(torch.float32).view(1, -1, 1, 1) + out.copy_(total) + + +def block_pool(x: torch.Tensor, out: torch.Tensor, *, block_size: int, reduce: _Reduce) -> None: + """Reduce every run of ``block_size`` tokens of ``x`` into ``out`` with an fp32 accumulator. + + Args: + x: ``[batch, seq_len, heads, head_dim]`` activations; heads and head_dim must be contiguous, + the batch and token strides are arbitrary. + out: Contiguous ``[batch, ceil(seq_len / block_size), heads, head_dim]`` buffer of any float + dtype; it receives the rounded fp32 result. + block_size: Tokens per block; the final block may be shorter. + reduce: ``"mean"`` divides by the number of valid tokens of the block, ``"sum"`` does not. + """ + if reduce not in ("mean", "sum"): + raise ValueError(f"reduce must be 'mean' or 'sum'; got {reduce!r}") + if x.ndim != 4 or x.stride(3) != 1 or x.stride(2) != x.shape[3]: + raise ValueError( + "x must be [batch, seq_len, heads, head_dim] with contiguous heads and head_dim" + ) + batch_size, seq_len, num_heads, head_dim = x.shape + blocks = num_blocks(seq_len, block_size) + expected = (batch_size, blocks, num_heads, head_dim) + if tuple(out.shape) != expected or not out.is_contiguous(): + raise ValueError( + f"out must be a contiguous tensor of shape {expected}; got {tuple(out.shape)}" + ) + if x.device.type != "cuda": + _block_pool_torch(x, out, block_size=block_size, reduce=reduce) + return + row_width = num_heads * head_dim + width, chunks = _column_launch(row_width, _POOL_MAX_WIDTH) + _block_pool_kernel[(batch_size * blocks * chunks,)]( + x, + out, + seq_len, + blocks, + chunks, + row_width, + x.stride(0), + x.stride(1), + out.stride(0), + out.stride(1), + MEAN=reduce == "mean", + BLOCK=block_size, + TOKENS=min(_POOL_TOKENS_PER_LOAD, block_size), + WIDTH=width, + num_warps=4, + ) + + +# --------------------------------------------------------------------------- block statistics +@triton.jit +def _block_statistics_kernel( + x_ptr, + mean_ptr, + var_ptr, + num_blocks, + num_chunks, + row_width, + stride_x_batch, + stride_x_block, + stride_out_batch, + ROWS: tl.constexpr, + WIDTH: tl.constexpr, +): + """One program per (batch, column chunk): mean and clamped variance over the block axis.""" + pid = tl.program_id(0).to(tl.int64) + chunk = pid % num_chunks + batch = pid // num_chunks + columns = chunk * WIDTH + tl.arange(0, WIDTH) + in_row = columns < row_width + total = tl.zeros([WIDTH], dtype=tl.float32) + total_sq = tl.zeros([WIDTH], dtype=tl.float32) + for start in range(0, num_blocks, ROWS): + rows = start + tl.arange(0, ROWS) + values = tl.load( + x_ptr + batch * stride_x_batch + rows[:, None] * stride_x_block + columns[None, :], + mask=(rows < num_blocks)[:, None] & in_row[None, :], + other=0.0, + ).to(tl.float32) + total += tl.sum(values, axis=0) + total_sq += tl.sum(values * values, axis=0) + count = num_blocks.to(tl.float32) + mean = total / count + variance = tl.maximum(total_sq / count - mean * mean, 0.0) + tl.store(mean_ptr + batch * stride_out_batch + columns, mean, mask=in_row) + tl.store(var_ptr + batch * stride_out_batch + columns, variance, mask=in_row) + + +def _block_statistics_torch(x: torch.Tensor, out_mean: torch.Tensor, out_var: torch.Tensor) -> None: + values = x.to(torch.float32) + mean = values.mean(dim=1) + out_mean.copy_(mean) + out_var.copy_(torch.clamp(values.square().mean(dim=1) - mean.square(), min=0.0)) + + +def block_statistics(x: torch.Tensor, out_mean: torch.Tensor, out_var: torch.Tensor) -> None: + """Per-channel mean and biased variance of ``x`` over its block axis. + + Args: + x: Contiguous ``[batch, num_blocks, heads, head_dim]`` block summaries. + out_mean: Contiguous fp32 ``[batch, heads, head_dim]`` buffer. + out_var: Contiguous fp32 ``[batch, heads, head_dim]`` buffer; negative rounding is clamped to zero. + """ + batch_size, blocks, num_heads, head_dim = x.shape + expected = (batch_size, num_heads, head_dim) + for name, tensor in (("out_mean", out_mean), ("out_var", out_var)): + if ( + tuple(tensor.shape) != expected + or tensor.dtype != torch.float32 + or not tensor.is_contiguous() + ): + raise ValueError(f"{name} must be a contiguous fp32 tensor of shape {expected}") + if not x.is_contiguous(): + raise ValueError("x must be contiguous") + if x.device.type != "cuda": + _block_statistics_torch(x, out_mean, out_var) + return + row_width = num_heads * head_dim + width, chunks = _column_launch(row_width, _STATS_WIDTH) + _block_statistics_kernel[(batch_size * chunks,)]( + x, + out_mean, + out_var, + blocks, + chunks, + row_width, + x.stride(0), + x.stride(1), + out_mean.stride(0), + ROWS=_STATS_ROWS_PER_LOAD, + WIDTH=width, + num_warps=4, + ) + + +# --------------------------------------------------------------------------- exact-block selection +@triton.jit +def _select_exact_blocks_kernel( + centroid_ptr, + keys_ptr, + mean_ptr, + var_ptr, + bits_ptr, + num_q_blocks, + num_kv_blocks, + num_words, + num_heads, + local_radius, + tau, + log2_scale, + epsilon, + stride_c_batch, + stride_c_block, + stride_c_head, + stride_k_batch, + stride_k_block, + stride_k_head, + stride_s_batch, + stride_s_head, + stride_b_batch, + stride_b_head, + stride_b_block, + Q_BLOCKS: tl.constexpr, + HEAD_DIM: tl.constexpr, +): + """One program per (batch, head, tile of Q_BLOCKS query blocks); emits every word of the tile. + + Scores are ``log2_scale * ``. The fp32 centroid is split into three terms + of the key dtype so every tensor-core product is exact and only the fp32 accumulation rounds. + """ + tile = tl.program_id(0) + batch_head = tl.program_id(1).to(tl.int64) + batch = batch_head // num_heads + head = batch_head % num_heads + q_blocks = tile * Q_BLOCKS + tl.arange(0, Q_BLOCKS) + q_valid = q_blocks < num_q_blocks + dims = tl.arange(0, HEAD_DIM) + + centroid = tl.load( + centroid_ptr + + batch * stride_c_batch + + q_blocks[:, None] * stride_c_block + + head * stride_c_head + + dims[None, :], + mask=q_valid[:, None], + other=0.0, + ) + key_mean = tl.load(mean_ptr + batch * stride_s_batch + head * stride_s_head + dims) + key_var = tl.load(var_ptr + batch * stride_s_batch + head * stride_s_head + dims) + projected_mean = tl.sum(centroid * key_mean[None, :], axis=1) * log2_scale + projected_var = tl.sum(centroid * centroid * key_var[None, :], axis=1) * log2_scale * log2_scale + threshold = projected_mean + tau * tl.sqrt(tl.maximum(projected_var, 0.0) + epsilon) + + high = centroid.to(keys_ptr.dtype.element_ty) + rest = centroid - high.to(tl.float32) + mid = rest.to(keys_ptr.dtype.element_ty) + low = (rest - mid.to(tl.float32)).to(keys_ptr.dtype.element_ty) + + lanes = tl.arange(0, 32) + lane_bits = 1 << lanes.to(tl.int64) + for word in range(num_words): + kv_blocks = word * 32 + lanes + kv_valid = kv_blocks < num_kv_blocks + keys = tl.load( + keys_ptr + + batch * stride_k_batch + + kv_blocks[:, None] * stride_k_block + + head * stride_k_head + + dims[None, :], + mask=kv_valid[:, None], + other=0.0, + ) + keys_t = tl.trans(keys) + scores = (tl.dot(high, keys_t) + tl.dot(mid, keys_t) + tl.dot(low, keys_t)) * log2_scale + distance = q_blocks[:, None] - kv_blocks[None, :] + is_local = (distance >= -local_radius) & (distance <= local_radius) + exact = kv_valid[None, :] & ((scores > threshold[:, None]) | is_local) + packed = tl.sum(tl.where(exact, lane_bits[None, :], 0), axis=1) + tl.store( + bits_ptr + + batch * stride_b_batch + + head * stride_b_head + + q_blocks * stride_b_block + + word, + packed.to(tl.int32), + mask=q_valid, + ) + + +def _select_exact_blocks_torch( + centroid: torch.Tensor, + k_summary: torch.Tensor, + k_mean: torch.Tensor, + k_var: torch.Tensor, + exact_block_bits: torch.Tensor, + *, + tau: float, + sm_scale: float, +) -> None: + log2_scale = float(sm_scale) * _LOG2_E + q = centroid.to(torch.float64) + k = k_summary.to(torch.float64) + projected_mean = torch.einsum("bqhd,bhd->bhq", q, k_mean.to(torch.float64)) * log2_scale + projected_var = ( + torch.einsum("bqhd,bhd->bhq", q.square(), k_var.to(torch.float64)) * log2_scale * log2_scale + ) + threshold = projected_mean + float(tau) * torch.sqrt( + torch.clamp(projected_var, min=0.0) + _THRESHOLD_EPSILON + ) + scores = torch.einsum("bqhd,bkhd->bhqk", q, k) * log2_scale + exact = scores > threshold.unsqueeze(-1) + num_kv_blocks = k_summary.shape[1] + ids = torch.arange(num_kv_blocks, device=centroid.device) + exact |= ((ids[:, None] - ids[None, :]).abs() <= _LOCAL_RADIUS)[None, None] + words = num_words(num_kv_blocks) + padded = torch.nn.functional.pad(exact, (0, words * _WORD_BITS - num_kv_blocks)) + weights = 1 << torch.arange(_WORD_BITS, dtype=torch.int64, device=centroid.device) + packed = (padded.view(*exact.shape[:-1], words, _WORD_BITS).to(torch.int64) * weights).sum( + dim=-1 + ) + exact_block_bits.copy_(packed.to(torch.uint32)) + + +def select_exact_blocks( + centroid: torch.Tensor, + k_summary: torch.Tensor, + k_mean: torch.Tensor, + k_var: torch.Tensor, + exact_block_bits: torch.Tensor, + *, + tau: float, + sm_scale: float, +) -> None: + """Pack the SOL exact-block decision of every (query block, key block) pair into ``exact_block_bits``. + + A key block is exact when ``sm_scale * log2(e) * `` exceeds the row threshold + ``mean + tau * sqrt(var + 1e-6)`` projected from the key statistics, or when it lies within one + block of the query block. Bit ``r`` of word ``w`` selects key block ``32 * w + r``; padding bits of + the final word are zero. + + Args: + centroid: Contiguous fp32 ``[batch, num_q_blocks, heads, head_dim]`` query block means. + k_summary: Contiguous ``[batch, num_kv_blocks, heads, head_dim]`` key block means (bf16 or fp16). + k_mean: fp32 ``[batch, heads, head_dim]`` mean of ``k_summary`` over its block axis. + k_var: fp32 ``[batch, heads, head_dim]`` variance of ``k_summary`` over its block axis. + exact_block_bits: Contiguous uint32 ``[batch, heads, num_q_blocks, ceil(num_kv_blocks / 32)]``. + tau: Threshold slope in standard deviations. + sm_scale: Softmax scale of the attention call. + """ + batch_size, q_blocks, num_heads, head_dim = centroid.shape + kv_blocks = k_summary.shape[1] + expected_bits = (batch_size, num_heads, q_blocks, num_words(kv_blocks)) + if tuple(exact_block_bits.shape) != expected_bits or exact_block_bits.dtype != torch.uint32: + raise ValueError(f"exact_block_bits must be uint32 of shape {expected_bits}") + if tuple(k_summary.shape) != (batch_size, kv_blocks, num_heads, head_dim): + raise ValueError("k_summary must match centroid in batch, heads, and head_dim") + if ( + centroid.dtype != torch.float32 + or not centroid.is_contiguous() + or not k_summary.is_contiguous() + ): + raise ValueError("centroid must be contiguous fp32 and k_summary contiguous") + if not exact_block_bits.is_contiguous(): + raise ValueError("exact_block_bits must be contiguous") + if centroid.device.type != "cuda": + _select_exact_blocks_torch( + centroid, k_summary, k_mean, k_var, exact_block_bits, tau=tau, sm_scale=sm_scale + ) + return + bits = exact_block_bits.view(torch.int32) + grid = (triton.cdiv(q_blocks, _SELECT_Q_BLOCKS), batch_size * num_heads) + _select_exact_blocks_kernel[grid]( + centroid, + k_summary, + k_mean, + k_var, + bits, + q_blocks, + kv_blocks, + num_words(kv_blocks), + num_heads, + _LOCAL_RADIUS, + float(tau), + float(sm_scale) * _LOG2_E, + _THRESHOLD_EPSILON, + centroid.stride(0), + centroid.stride(1), + centroid.stride(2), + k_summary.stride(0), + k_summary.stride(1), + k_summary.stride(2), + k_mean.stride(0), + k_mean.stride(1), + bits.stride(0), + bits.stride(1), + bits.stride(2), + Q_BLOCKS=_SELECT_Q_BLOCKS, + HEAD_DIM=head_dim, + num_warps=4, + ) + + +# --------------------------------------------------------------------------- graph-visible operator +@torch.library.custom_op( + "trtllm::visual_gen_sol_predictor", + mutates_args=( + "exact_block_bits", + "k_summary", + "v_summary", + "k_mean", + "k_var_diag", + "q_centroid", + ), + device_types="cuda", +) +def visual_gen_sol_predictor( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + exact_block_bits: torch.Tensor, + k_summary: torch.Tensor, + v_summary: torch.Tensor, + k_mean: torch.Tensor, + k_var_diag: torch.Tensor, + q_centroid: torch.Tensor, + block_size: int, + tau: float, + sm_scale: float, +) -> None: + """Update caller-owned SOL route and proxy tensors in place.""" + + block_pool(q, q_centroid, block_size=block_size, reduce="mean") + block_pool(k, k_summary, block_size=block_size, reduce="mean") + block_pool(v, v_summary, block_size=block_size, reduce="sum") + block_statistics(k_summary, k_mean, k_var_diag) + select_exact_blocks( + q_centroid, k_summary, k_mean, k_var_diag, exact_block_bits, tau=tau, sm_scale=sm_scale + ) + + +@torch.library.register_fake("trtllm::visual_gen_sol_predictor") +def _( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + exact_block_bits: torch.Tensor, + k_summary: torch.Tensor, + v_summary: torch.Tensor, + k_mean: torch.Tensor, + k_var_diag: torch.Tensor, + q_centroid: torch.Tensor, + block_size: int, + tau: float, + sm_scale: float, +) -> None: + return None + + +__all__ = ["block_pool", "block_statistics", "num_blocks", "num_words", "select_exact_blocks"] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/params.py b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/params.py new file mode 100644 index 000000000000..ac27bd90ea39 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/params.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Lowered parameters for two-stage VisualGen SOL attention.""" + +from __future__ import annotations + +import math +import numbers +import struct +from dataclasses import dataclass, field +from typing import Literal + +import torch + +from tensorrt_llm._torch.attention.backends.sparse.params import SparseParams + + +def _as_timestep_float(timestep: object) -> float | None: + if timestep is None: + return None + if isinstance(timestep, torch.Tensor): + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + raise RuntimeError("SOL graph phase must be precomputed before CUDA Graph capture") + if timestep.numel() == 0: + return None + # WAN I2V can carry one timestep per token, with reference tokens fixed + # at zero. Stay dense until every live token is below the cutoff. + timestep = timestep.amax().item() + if isinstance(timestep, bool) or not isinstance(timestep, numbers.Real): + raise TypeError("timestep must be a real scalar or tensor") + value = float(timestep) + if not math.isfinite(value): + raise ValueError("timestep must be finite") + return value + + +@dataclass(frozen=True, slots=True) +class SolParams(SparseParams): + """Static SOL policy lowered from the user-facing VisualGen config.""" + + algorithm: Literal["sol_attn"] = field(init=False, default="sol_attn") + tau: float = 1.0 + disabled_until_timestep: float | None = None + dense_layers: frozenset[int] = field(default_factory=frozenset) + + def __post_init__(self) -> None: + if isinstance(self.tau, bool) or not isinstance(self.tau, numbers.Real): + raise TypeError("tau must be a finite real number") + try: + tau = struct.unpack("=f", struct.pack("=f", float(self.tau)))[0] + except (OverflowError, ValueError, struct.error) as error: + raise ValueError("tau must be representable as float32") from error + if not math.isfinite(tau): + raise ValueError("tau must be finite") + object.__setattr__(self, "tau", tau) + + cutoff = self.disabled_until_timestep + if cutoff is not None: + if isinstance(cutoff, bool) or not isinstance(cutoff, numbers.Real): + raise TypeError("disabled_until_timestep must be a real number or None") + cutoff = float(cutoff) + if not math.isfinite(cutoff) or not 0.0 < cutoff <= 1.0: + raise ValueError("disabled_until_timestep must be in (0, 1]") + object.__setattr__(self, "disabled_until_timestep", cutoff) + + dense_layers = frozenset(self.dense_layers) + if any( + isinstance(layer, bool) or not isinstance(layer, int) or layer < 0 + for layer in dense_layers + ): + raise ValueError("dense_layers must contain only non-negative integers") + object.__setattr__(self, "dense_layers", dense_layers) + + @staticmethod + def get_graph_phase_for_timestep( + timestep: object, + *, + disabled_until_timestep: float | None, + ) -> int | None: + """Return 0 for the dense prefix and 1 for the sparse suffix.""" + + if disabled_until_timestep is None: + return None + value = _as_timestep_float(timestep) + if value is None: + return None + return int(value < disabled_until_timestep) + + def should_use_sparse( + self, + *, + layer_idx: int, + timestep: object, + graph_phase: int | None = None, + ) -> bool: + """Return whether this layer should execute the SOL sparse path.""" + + if layer_idx in self.dense_layers: + return False + if graph_phase is not None: + if graph_phase not in (0, 1): + raise ValueError("SOL graph_phase must be 0 or 1") + phase = graph_phase + else: + phase = self.get_graph_phase_for_timestep( + timestep, + disabled_until_timestep=self.disabled_until_timestep, + ) + if phase is None: + if self.disabled_until_timestep is not None: + raise ValueError( + "timestep is required when SOL disabled_until_timestep is configured" + ) + return True + return phase == 1 + + +__all__ = ["SolParams"] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/predictor.py b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/predictor.py new file mode 100644 index 000000000000..4fa91d52e09e --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/sol/predictor.py @@ -0,0 +1,282 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Plan-owned runtime for the two-stage VisualGen SOL predictor.""" + +from __future__ import annotations + +import numbers +import struct +from dataclasses import dataclass + +import torch + +from . import kernels as _kernels + +BLOCK_SIZE = 64 +HEAD_DIM = 128 + + +def _positive_int(value: object, name: str) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise TypeError(f"{name} must be a Python integer") + if value <= 0: + raise ValueError(f"{name} must be positive") + return value + + +def _float32_scalar(value: object, name: str) -> float: + if isinstance(value, bool) or not isinstance(value, numbers.Real): + raise TypeError(f"{name} must be a finite Python real") + try: + result = struct.unpack("=f", struct.pack("=f", float(value)))[0] + except (OverflowError, TypeError, ValueError, struct.error) as error: + raise ValueError(f"{name} must be representable as float32") from error + if not -float("inf") < result < float("inf"): + raise ValueError(f"{name} must be finite") + return result + + +def _normalize_runtime_scalars(*, tau: object, sm_scale: object) -> tuple[float, float]: + """Validate and round the two dynamic selector scalars to binary32.""" + + effective_tau = _float32_scalar(tau, "tau") + effective_sm_scale = _float32_scalar(sm_scale, "sm_scale") + if effective_sm_scale <= 0.0: + raise ValueError("sm_scale must be positive") + return effective_tau, effective_sm_scale + + +@dataclass(frozen=True) +class SolPredictorGeometry: + """Static shape specialization for compact BF16 self-MHA.""" + + batch_size: int + seq_len: int + num_heads: int + head_dim: int + num_q_blocks: int + num_kv_blocks: int + exact_words: int + tail_tokens: int + + @classmethod + def create( + cls, + *, + batch_size: object, + seq_len: object, + num_heads: object, + head_dim: object = HEAD_DIM, + ) -> "SolPredictorGeometry": + batch = _positive_int(batch_size, "batch_size") + tokens = _positive_int(seq_len, "seq_len") + heads = _positive_int(num_heads, "num_heads") + dim = _positive_int(head_dim, "head_dim") + if dim != HEAD_DIM: + raise ValueError(f"SOL predictor only supports head_dim={HEAD_DIM}; got {dim}") + blocks = _kernels.num_blocks(tokens, BLOCK_SIZE) + tail = tokens - (blocks - 1) * BLOCK_SIZE + return cls( + batch_size=batch, + seq_len=tokens, + num_heads=heads, + head_dim=dim, + num_q_blocks=blocks, + num_kv_blocks=blocks, + exact_words=_kernels.num_words(blocks), + tail_tokens=tail, + ) + + @property + def tensor_shape(self) -> tuple[int, int, int, int]: + return (self.batch_size, self.seq_len, self.num_heads, self.head_dim) + + @property + def summary_shape(self) -> tuple[int, int, int, int]: + return (self.batch_size, self.num_kv_blocks, self.num_heads, self.head_dim) + + @property + def stats_shape(self) -> tuple[int, int, int]: + return (self.batch_size, self.num_heads, self.head_dim) + + @property + def exact_block_bits_shape(self) -> tuple[int, int, int, int]: + return (self.batch_size, self.num_heads, self.num_q_blocks, self.exact_words) + + +@dataclass(frozen=True) +class SolPredictorPlanKey: + """Cache key containing only static kernel specialization state.""" + + geometry: SolPredictorGeometry + device_index: int + dtype: torch.dtype + + +@dataclass(frozen=True) +class SolPredictorOutputs: + """Live predictor tensors consumed by block-sparse attention.""" + + exact_block_bits: torch.Tensor + k_summary: torch.Tensor + v_summary: torch.Tensor + + +@dataclass(frozen=True) +class SolPredictorPlan: + """One published shape specialization and its stable live storage.""" + + key: SolPredictorPlanKey + outputs: SolPredictorOutputs + k_mean: torch.Tensor + k_var_diag: torch.Tensor + q_centroid: torch.Tensor + + +class SOLSparsePredictor: + """Cache of shape-specialized, graph-stable SOL predictor plans.""" + + def __init__(self) -> None: + self._plans: dict[SolPredictorPlanKey, SolPredictorPlan] = {} + + @property + def num_plans(self) -> int: + return len(self._plans) + + @staticmethod + def support_reason(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> str | None: + """Return why compact two-stage SOL cannot serve these tensors.""" + + if not all(isinstance(tensor, torch.Tensor) for tensor in (q, k, v)): + return "q, k, and v must be torch tensors" + if q.ndim != 4: + return f"q must use compact BSHD layout; got rank {q.ndim}" + if k.shape != q.shape or v.shape != q.shape: + return "SOL predictor requires uniform self-attention q/k/v shapes" + if q.dtype != torch.bfloat16 or k.dtype != q.dtype or v.dtype != q.dtype: + return "SOL predictor requires matching BF16 q/k/v" + if not q.is_cuda or not k.is_cuda or not v.is_cuda: + return "SOL predictor requires CUDA q/k/v" + if k.device != q.device or v.device != q.device: + return "SOL predictor requires q/k/v on one CUDA device" + if not q.is_contiguous() or not k.is_contiguous() or not v.is_contiguous(): + return "SOL predictor requires contiguous BSHD q/k/v" + if q.shape[-1] != HEAD_DIM: + return f"SOL predictor requires head_dim={HEAD_DIM}; got {q.shape[-1]}" + if q.shape[0] <= 0 or q.shape[1] <= 0 or q.shape[2] <= 0: + return "SOL predictor requires positive B, S, and H" + return None + + @classmethod + def _key_from_inputs( + cls, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor + ) -> SolPredictorPlanKey: + reason = cls.support_reason(q, k, v) + if reason is not None: + raise ValueError(reason) + device_index = q.device.index + if device_index is None: + device_index = torch.cuda.current_device() + geometry = SolPredictorGeometry.create( + batch_size=q.shape[0], + seq_len=q.shape[1], + num_heads=q.shape[2], + head_dim=q.shape[3], + ) + return SolPredictorPlanKey( + geometry=geometry, + device_index=device_index, + dtype=q.dtype, + ) + + @staticmethod + def _launch( + plan: SolPredictorPlan, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + tau: float, + sm_scale: float, + ) -> None: + torch.ops.trtllm.visual_gen_sol_predictor( + q, + k, + v, + plan.outputs.exact_block_bits, + plan.outputs.k_summary, + plan.outputs.v_summary, + plan.k_mean, + plan.k_var_diag, + plan.q_centroid, + BLOCK_SIZE, + tau, + sm_scale, + ) + + @torch.compiler.disable + def prepare(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> SolPredictorPlan: + """Allocate one geometry and warm its kernels outside compiled or captured regions. + + The host-only boundary preserves per-instance plan ownership and keeps + kernel compilation and allocation out of Dynamo and CUDA Graph capture. + It requires the VisualGen default ``torch.compile(fullgraph=False)``. + """ + + key = self._key_from_inputs(q, k, v) + existing = self._plans.get(key) + if existing is not None: + return existing + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("SOL predictor plan must be prepared before CUDA graph capture") + + geometry = key.geometry + with torch.cuda.device(key.device_index): + k_summary = torch.empty(geometry.summary_shape, dtype=key.dtype, device=q.device) + v_summary = torch.empty_like(k_summary) + exact_block_bits = torch.empty( + geometry.exact_block_bits_shape, dtype=torch.uint32, device=q.device + ) + k_mean = torch.empty(geometry.stats_shape, dtype=torch.float32, device=q.device) + k_var_diag = torch.empty_like(k_mean) + q_centroid = torch.empty(geometry.summary_shape, dtype=torch.float32, device=q.device) + plan = SolPredictorPlan( + key=key, + outputs=SolPredictorOutputs( + exact_block_bits=exact_block_bits, + k_summary=k_summary, + v_summary=v_summary, + ), + k_mean=k_mean, + k_var_diag=k_var_diag, + q_centroid=q_centroid, + ) + # The warm-up launch compiles every kernel specialization of this geometry. + self._launch(plan, q, k, v, tau=0.0, sm_scale=1.0) + self._plans[key] = plan + return plan + + def predict( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + tau: object, + sm_scale: object, + ) -> SolPredictorOutputs: + """Update and return graph-stable SOL routes and proxy summaries.""" + + effective_tau, effective_sm_scale = _normalize_runtime_scalars(tau=tau, sm_scale=sm_scale) + plan = self.prepare(q, k, v) + self._launch(plan, q, k, v, tau=effective_tau, sm_scale=effective_sm_scale) + return plan.outputs + + +__all__ = [ + "SOLSparsePredictor", + "SolPredictorGeometry", + "SolPredictorOutputs", + "SolPredictorPlan", + "SolPredictorPlanKey", +] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/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..1f6af9b72ecc --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/backend.py @@ -0,0 +1,343 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""VisualGen TRTLLM-first attention backends for Video Sparse Attention.""" + +from typing import Optional + +import torch +import torch.nn.functional as F + +from tensorrt_llm.logger import logger +from tensorrt_llm.models.modeling_utils import QuantConfig +from tensorrt_llm.visual_gen.args import QuantAttentionConfig + +from .....attention.backends.fmha.prims_ts_block_sparse import PrimsTSBlockSparseFmha +from .....attention.backends.interface import PredefinedAttentionMask +from .....attention.backends.sparse.params import SparseBackendForwardArgs +from ...cute_dsl import CuTeDSLAttention +from ...trtllm import TrtllmAttention +from .metadata import VSA_BLOCK_SIZE, VSAMetadata +from .predictor import VSAForwardInputs, VSAPredictor, vsa_post_process + +_vsa_import_error = None +try: + from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.video_sparse_attention import ( + block_sparse_attn_from_indices_cute, + is_cute_supported, + ) +except (ImportError, OSError) as error: + block_sparse_attn_from_indices_cute = None + is_cute_supported = None + _vsa_import_error = error + + +VSA_KERNEL_MAX_CUBES: int = 4 * 1024 + + +def _normalize_qkv_inputs( + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Normalize separate BSHD or Ulysses-packed BSH3HD inputs.""" + + if k is not None and v is not None: + return q, k, v + if k is not None or v is not None: + raise ValueError("VSA requires complete separate Q/K/V or one packed QKV tensor.") + if q.ndim != 5 or q.shape[2] != 3: + raise ValueError("VSA packed QKV must have shape [B, S, 3, H, D].") + return q.unbind(dim=2) + + +def _get_unsupported_primts_reason( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + metadata: VSAMetadata, +) -> str | None: + if q.shape != k.shape or q.shape != v.shape: + return "VSA PrimTS requires matching MHA Q/K/V shapes" + if q.device.type != "cuda": + return f"VSA PrimTS requires CUDA tensors, got {q.device}" + if q.dtype not in (torch.float16, torch.bfloat16): + return f"VSA PrimTS requires FP16 or BF16 tensors, got {q.dtype}" + batch_size, seq_len, num_heads, head_dim = map(int, q.shape) + if min(batch_size, seq_len, num_heads, metadata.num_cubes) <= 0: + return "VSA PrimTS requires positive batch, sequence, head, and cube extents" + if head_dim != 128: + return f"VSA PrimTS requires head_dim=128, got {head_dim}" + if metadata.padded_seq_length != metadata.num_cubes * VSA_BLOCK_SIZE: + return "VSA tiled sequence length must match its 64-token cube count" + if batch_size > 65535 or num_heads > 65535: + return "VSA PrimTS batch and head dimensions must fit the CUDA grid" + return None + + +class VSATrtllmAttention(TrtllmAttention): + """TRTLLM VSA backend using the generic block-sparse forward lifecycle.""" + + def __init__( + self, + layer_idx: int = 0, + num_heads: int = 8, + head_dim: int = 64, + num_kv_heads: Optional[int] = None, + quant_config: Optional[QuantConfig] = None, + dtype: Optional[torch.dtype] = None, + max_batch_size: int = 16, + max_seq_len: int = 4096, + quant_attention_config: Optional[QuantAttentionConfig] = None, + attention_metadata_state: Optional[dict] = None, + ) -> None: + num_kv_heads = num_kv_heads or num_heads + super().__init__( + layer_idx=layer_idx, + num_heads=num_heads, + head_dim=head_dim, + num_kv_heads=num_kv_heads, + quant_config=quant_config, + dtype=dtype, + max_batch_size=max_batch_size, + max_seq_len=max_seq_len, + quant_attention_config=quant_attention_config, + attention_metadata_state=attention_metadata_state, + sparse_params=None, + ) + assert attention_metadata_state is not None + predictor_cache = attention_metadata_state.setdefault("sparse_predictors", {}) + predictor_key = ("vsa", num_heads, num_kv_heads) + predictor = predictor_cache.get(predictor_key) + if predictor is None: + predictor = VSAPredictor( + num_heads=num_heads, + num_kv_heads=num_kv_heads, + ) + predictor_cache[predictor_key] = predictor + elif not isinstance(predictor, VSAPredictor): + raise TypeError("model-scoped VSA predictor cache contains an invalid value") + self.predictor = predictor + + def forward( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + batch_size: int, + seq_len: int, + attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL, + seq_len_kv: Optional[int] = None, + **kwargs, + ) -> torch.Tensor: + """Run the VSA coarse stage, then the fine stage through the core forward. + + The coarse stage predicts the complete block-sparse payload, which is + handed to the core prediction hook via ``sparse_backend_args``. The fine + output is blended with the coarse output afterward. + """ + + q, k, v = _normalize_qkv_inputs(q, k, v) + metadata = self.predictor.get_metadata() + use_primts = any( + isinstance(fmha, PrimsTSBlockSparseFmha) for fmha in self._fmha_manager.fmha_libs + ) + unsupported_reason = _get_unsupported_primts_reason(q, k, v, metadata) + if self.quant_attention_config is not None: + unsupported_reason = "VSA PrimTS does not support quant_attention_config" + if not use_primts: + logger.warning_once( + "TRTLLM VSA cannot use PrimTS block-sparse attention because the " + "prims_ts_block_sparse FMHA library is unavailable; using the compact " + "dense TRTLLM fine stage.", + key="trtllm_vsa_primts_unavailable", + ) + elif unsupported_reason is not None: + logger.warning_once( + "TRTLLM VSA cannot use PrimTS block-sparse attention: " + f"{unsupported_reason}; using the compact dense TRTLLM fine stage.", + key=("trtllm_vsa_primts_unsupported_envelope", unsupported_reason), + ) + use_sparse_fine = use_primts and unsupported_reason is None + + inputs = self.predictor.predict( + q, + k, + v, + batch_size=batch_size, + seq_len=seq_len, + seq_len_kv=seq_len if seq_len_kv is None else seq_len_kv, + attention_mask=attention_mask, + gate_compress=kwargs.pop("gate_compress", None), + gate_fine=kwargs.pop("gate_fine", None), + use_sparse_fine=use_sparse_fine, + produce_block_sparse_inputs=use_sparse_fine, + metadata=metadata, + ) + sparse_backend_args = None + if inputs.block_sparse_inputs is not None: + sparse_backend_args = SparseBackendForwardArgs( + block_sparse_inputs=inputs.block_sparse_inputs, + ) + fine_output = super().forward( + inputs.q, + inputs.k, + inputs.v, + inputs.batch_size, + inputs.seq_len, + attention_mask=attention_mask, + seq_len_kv=inputs.seq_len, + sparse_backend_args=sparse_backend_args, + **kwargs, + ) + combined = vsa_post_process(fine_output, inputs) + return combined.reshape(combined.shape[0], combined.shape[1], -1) + + @classmethod + def support_fused_qkv(cls) -> bool: + return True + + +class VSACuTeDSLAttention(CuTeDSLAttention): + """CuTe DSL VSA backend reusing TRTLLM's predictor and post-processing.""" + + def __init__( + self, + layer_idx: int = 0, + num_heads: int = 8, + head_dim: int = 128, + num_kv_heads: Optional[int] = None, + dtype: Optional[torch.dtype] = None, + **kwargs, + ) -> None: + super().__init__( + layer_idx=layer_idx, + num_heads=num_heads, + head_dim=head_dim, + num_kv_heads=num_kv_heads, + dtype=dtype, + **kwargs, + ) + self.predictor = VSAPredictor( + num_heads=num_heads, + num_kv_heads=num_kv_heads, + ) + + def forward( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + *, + attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL, + **kwargs, + ) -> torch.Tensor: + q, k, v = _normalize_qkv_inputs(q, k, v) + gate_compress = kwargs.pop("gate_compress", None) + gate_fine = kwargs.pop("gate_fine", None) + expected_extents = { + "batch_size": int(q.shape[0]), + "seq_len": int(q.shape[1]), + "seq_len_kv": int(k.shape[1]), + } + for name, expected in expected_extents.items(): + actual = kwargs.pop(name, expected) + if not isinstance(actual, int) or isinstance(actual, bool) or actual != expected: + raise ValueError(f"VSA {name}={actual!r} does not match Q/K/V extent {expected}") + kwargs.pop("timestep", None) + if kwargs: + unexpected_names = ", ".join(sorted(kwargs)) + raise TypeError(f"Unexpected CuTeDSL VSA forward keyword arguments: {unexpected_names}") + + metadata = self.predictor.get_metadata() + # The CuTe kernel's fixed launch topology is bounded by the number of + # VSA cubes; larger shapes retain identical VSA math via dense SDPA. + use_cute = ( + _vsa_import_error is None + and is_cute_supported is not None + and is_cute_supported(q) + and q.dtype == k.dtype == v.dtype + and metadata.num_cubes <= VSA_KERNEL_MAX_CUBES + ) + inputs = self.predictor.predict( + q, + k, + v, + batch_size=int(q.shape[0]), + seq_len=int(q.shape[1]), + seq_len_kv=int(k.shape[1]), + attention_mask=attention_mask, + gate_compress=gate_compress, + gate_fine=gate_fine, + use_sparse_fine=use_cute, + produce_block_sparse_inputs=False, + metadata=metadata, + ) + if use_cute: + fine_output = self._execute_sparse_fine(inputs) + else: + fine_output = F.scaled_dot_product_attention( + inputs.q.transpose(1, 2), + inputs.k.transpose(1, 2), + inputs.v.transpose(1, 2), + ).transpose(1, 2) + return vsa_post_process(fine_output, inputs) + + def _execute_sparse_fine(self, inputs: VSAForwardInputs) -> torch.Tensor: + """Execute only the CuTe-specific VSA fine kernel.""" + + q_hnd = inputs.q.transpose(1, 2).contiguous() + k_hnd = inputs.k.transpose(1, 2).contiguous() + v_hnd = inputs.v.transpose(1, 2).contiguous() + q2k_num = torch.full( + (inputs.batch_size, q_hnd.shape[1], inputs.num_cubes), + inputs.cur_topk, + dtype=torch.int32, + device=inputs.q.device, + ) + output_hnd, _lse = block_sparse_attn_from_indices_cute( + q_hnd, + k_hnd, + v_hnd, + q2k_idx=inputs.topk_indices.contiguous(), + q2k_num=q2k_num, + variable_block_sizes=inputs.variable_block_sizes.to(torch.int32), + ) + return output_hnd.transpose(1, 2) + + def forward_with_lse( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor]: + raise NotImplementedError("CuTe DSL VSA does not support LSE output.") + + @classmethod + def support_fused_qkv(cls) -> bool: + return True + + @classmethod + def support_lse(cls) -> bool: + return False + + +__all__ = [ + "VSACuTeDSLAttention", + "VSATrtllmAttention", + "VSA_KERNEL_MAX_CUBES", +] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/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..31916e397436 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/metadata.py @@ -0,0 +1,241 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Shape, policy, and forward-context metadata for Video Sparse Attention.""" + +import contextvars +from contextlib import contextmanager +from dataclasses import dataclass +from math import ceil +from typing import Iterator, Optional, Tuple, TypedDict + +import torch + +# A 4x4x4 cube is one 64-token sparse block for every VSA fine-stage backend. +VSA_TILE_SIZE: Tuple[int, int, int] = (4, 4, 4) +VSA_BLOCK_SIZE = VSA_TILE_SIZE[0] * VSA_TILE_SIZE[1] * VSA_TILE_SIZE[2] +_DEFAULT_MAX_CACHED_SHAPES = 16 +_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.""" + + def __init__(self, max_cached_shapes: int = _DEFAULT_MAX_CACHED_SHAPES) -> None: + if max_cached_shapes <= 0: + raise ValueError("max_cached_shapes must be positive") + self._max_cached_shapes = max_cached_shapes + self._cache: dict[Tuple[Tuple[int, int, int], torch.device], _VSAShapeMetadata] = {} + + def _build_metadata( + self, + dit_seq_shape: Tuple[int, int, int], + device: torch.device, + ) -> _VSAShapeMetadata: + time, height, width = dit_seq_shape + tile_time, tile_height, tile_width = VSA_TILE_SIZE + num_tiles = ( + ceil(time / tile_time), + ceil(height / tile_height), + ceil(width / tile_width), + ) + total_seq_length = time * height * width + padded_seq_length = ( + num_tiles[0] * num_tiles[1] * num_tiles[2] * tile_time * tile_height * tile_width + ) + num_cubes = num_tiles[0] * num_tiles[1] * num_tiles[2] + tokens_per_cube = VSA_BLOCK_SIZE + + tile_partition_indices = _get_tile_partition_indices(dit_seq_shape, VSA_TILE_SIZE, device) + gather_idx = tile_partition_indices[tile_partition_indices >= 0] + + variable_block_sizes = _construct_variable_block_sizes( + dit_seq_shape, num_tiles, VSA_TILE_SIZE, device + ) + local_offsets = torch.arange(tokens_per_cube, device=device).expand( + num_cubes, tokens_per_cube + ) + cube_offsets = torch.arange(num_cubes, device=device).unsqueeze(1) * tokens_per_cube + non_pad_index = (cube_offsets + local_offsets)[ + local_offsets < variable_block_sizes.unsqueeze(1) + ] + + 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: + if len(self._cache) >= self._max_cached_shapes: + raise RuntimeError( + "VSA metadata cache reached its " + f"{self._max_cached_shapes}-shape limit; restart the pipeline or " + "reuse a configured resolution/frame profile" + ) + shape_metadata = self._build_metadata(dit_seq_shape, device) + self._cache[cache_key] = shape_metadata + + return VSAMetadata( + current_timestep=current_timestep, + vsa_sparsity=vsa_sparsity, + **shape_metadata, + ) + + def clear(self) -> None: + """Release cached tensors after CUDA Graphs that reference them are cleared.""" + + self._cache.clear() + + +_vsa_forward_context_var: contextvars.ContextVar[Optional[VSAMetadata]] = contextvars.ContextVar( + "_vsa_forward_context", default=None +) + + +@contextmanager +def set_vsa_forward_context(metadata: VSAMetadata) -> Iterator[None]: + """Make VSA metadata visible to attention layers for one model forward.""" + + token = _vsa_forward_context_var.set(metadata) + try: + yield + finally: + _vsa_forward_context_var.reset(token) + + +def get_vsa_forward_context() -> Optional[VSAMetadata]: + """Return the metadata for the active VSA model forward, if any.""" + + return _vsa_forward_context_var.get(None) + + +__all__ = [ + "VSA_TILE_SIZE", + "VSAMetadata", + "VSAMetadataBuilder", + "get_vsa_forward_context", + "set_vsa_forward_context", +] diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/predictor.py b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/predictor.py new file mode 100644 index 000000000000..97db41c5cf9b --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/sparse/vsa/predictor.py @@ -0,0 +1,303 @@ +# 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 ( + _DEFAULT_MAX_CACHED_SHAPES, + 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, max_cached_shapes: int = _DEFAULT_MAX_CACHED_SHAPES) -> None: + if max_cached_shapes <= 0: + raise ValueError("max_cached_shapes must be positive") + self._max_cached_shapes = max_cached_shapes + self._indptr_cache: dict[tuple[torch.device, int, int, int, int], torch.Tensor] = {} + + def from_selected_blocks( + self, + selected_blocks: torch.Tensor, + kv_valid_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 len(self._indptr_cache) >= self._max_cached_shapes: + raise RuntimeError( + "VSA route cache reached its " + f"{self._max_cached_shapes}-shape limit; restart the pipeline or " + "reuse a configured resolution/frame profile" + ) + if device.type == "cuda" and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "VSA route cache miss during CUDA Graph capture; " + "run an eager warmup with the same selected-block shape first" + ) + if batch_size * num_kv_heads * num_q_blocks * blocks_per_row > _SIGNED_INT32_MAX: + raise OverflowError("VSA route offsets must fit in signed int32") + row_offsets = torch.arange(num_q_blocks + 1, dtype=torch.int32, device=device) + head_offsets = torch.arange(batch_size * num_kv_heads, dtype=torch.int32, device=device) + return ( + head_offsets.reshape(batch_size, num_kv_heads, 1) * (num_q_blocks * blocks_per_row) + + row_offsets.reshape(1, 1, -1) * blocks_per_row + ).contiguous() + + +class VSAPredictor: + """Produce the complete per-call VSA block-attention input envelope.""" + + def __init__( + self, + num_heads: int, + num_kv_heads: Optional[int] = None, + max_cached_shapes: int = _DEFAULT_MAX_CACHED_SHAPES, + ) -> None: + resolved_num_kv_heads = num_kv_heads or num_heads + if resolved_num_kv_heads != num_heads: + raise ValueError( + "VSA coarse mean-pool assumes MHA (num_kv_heads == num_heads), " + f"got num_kv_heads={resolved_num_kv_heads}, num_heads={num_heads}. " + "GQA/MQA is not supported." + ) + self._route_builder = _VSARouteBuilder(max_cached_shapes=max_cached_shapes) + + @torch.compiler.disable + def get_metadata(self) -> VSAMetadata: + metadata = get_vsa_forward_context() + if metadata is None: + raise RuntimeError( + "VSA attention called without an active VSA forward context. " + "Wrap each transformer call with set_vsa_forward_context()." + ) + return metadata + + @staticmethod + def _validate_inputs( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + gate_compress: Optional[torch.Tensor], + gate_fine: Optional[torch.Tensor], + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + if gate_compress is None: + raise ValueError( + "VSA requires gate_compress. " + "Ensure to_gate_compress is wired in the transformer block." + ) + if q.ndim != 4 or q.shape != k.shape or q.shape != v.shape: + raise ValueError("VSA requires Q, K, and V with the same BSHD shape.") + if any(tensor.device != q.device or tensor.dtype != q.dtype for tensor in (k, v)): + raise ValueError("VSA requires Q, K, and V to share device and dtype.") + if not isinstance(gate_compress, torch.Tensor): + raise TypeError("VSA gate_compress must be a torch.Tensor.") + if ( + gate_compress.shape != q.shape + or gate_compress.device != q.device + or gate_compress.dtype != q.dtype + ): + raise ValueError("VSA gate_compress must share Q's shape, device, and dtype.") + if gate_fine is not None and ( + not isinstance(gate_fine, torch.Tensor) + or gate_fine.shape != q.shape + or gate_fine.device != q.device + or gate_fine.dtype != q.dtype + ): + raise ValueError("VSA gate_fine must share Q's shape, device, and dtype.") + return gate_compress, gate_fine + + def predict( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + batch_size: int, + seq_len: int, + seq_len_kv: int, + attention_mask: PredefinedAttentionMask, + gate_compress: Optional[torch.Tensor], + gate_fine: Optional[torch.Tensor], + use_sparse_fine: bool, + produce_block_sparse_inputs: bool, + metadata: Optional[VSAMetadata] = None, + ) -> VSAForwardInputs: + """Predict routes, effective QKV, and the shared post-process context.""" + + gate_compress, gate_fine = self._validate_inputs(q, k, v, gate_compress, gate_fine) + if attention_mask != PredefinedAttentionMask.FULL: + raise ValueError("VSA supports only full self-attention.") + if seq_len_kv != seq_len: + raise ValueError("VSA requires self-attention with matching Q and KV sequence lengths.") + if tuple(q.shape[:2]) != (batch_size, seq_len): + raise ValueError("VSA batch_size and seq_len must match the compact QKV tensors.") + + metadata = metadata or self.get_metadata() + 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..ff1af1e97d8a 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/trtllm.py @@ -27,8 +27,12 @@ from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.visual_gen.args import QuantAttentionConfig -from ...attention.backends.interface import AttentionRuntimeFeatures, PredefinedAttentionMask -from ...attention.backends.sparse.skip_softmax import SkipSoftmaxParams +from ...attention.backends.interface import ( + AttentionForwardArgs, + AttentionRuntimeFeatures, + PredefinedAttentionMask, +) +from ...attention.backends.sparse.params import SparseBackendForwardArgs, SparseParams from ...attention.backends.trtllm import TrtllmAttention as BaseTrtllmAttention from ...attention.backends.trtllm import TrtllmAttentionMetadata as BaseTrtllmAttentionMetadata from .interface import AttentionBackend, AttentionTensorLayout @@ -73,6 +77,12 @@ def __init__( self._cached_seq_lens: Optional[torch.Tensor] = None self._prepared = False + def get_fmha_cache_state(self, name: str) -> dict[str, object]: + """Return one model-scoped cache owned by this metadata adapter.""" + + fmha_caches = self._metadata_state.setdefault("fmha_caches", {}) + return fmha_caches.setdefault(name, {}) + def _needs_prepare(self, batch_size: int, seq_lens: torch.Tensor) -> bool: """Check if we need to call prepare() (current request seq_lens or shared metadata object seq_lens changed). @@ -185,6 +195,7 @@ class TrtllmAttention(BaseTrtllmAttention, AttentionBackend): - Metadata creation and preparation - No KV cache operation - SageAttention per-block QKV quantization (when a quant_attention_config is provided. requires unfused QKV) + - Separate-QKV forwarding for generic block-sparse attention and backends that reject fused QKV """ def __init__( @@ -199,9 +210,17 @@ def __init__( max_seq_len: int = 4096, quant_attention_config: Optional[QuantAttentionConfig] = None, attention_metadata_state: Optional[dict] = None, - sparse_params: Optional[SkipSoftmaxParams] = None, + sparse_params: Optional[SparseParams] = None, ): num_kv_heads = num_kv_heads or num_heads + if attention_metadata_state is None: + raise ValueError( + "TRTLLM attention requires `attention_metadata_state` to be provided " + "by visual-gen config for model-scoped metadata and plan sharing." + ) + self.metadata = TrtllmAttentionMetadata( + attention_metadata_state=attention_metadata_state, + ) super().__init__( layer_idx=layer_idx, @@ -216,12 +235,19 @@ def __init__( # TRTLLM expects flat [B*S, H*D] format self._preferred_layout = AttentionTensorLayout.NHD - self.metadata = TrtllmAttentionMetadata( - attention_metadata_state=attention_metadata_state, - ) - self.quant_attention_config = quant_attention_config + def update_quant_config(self, new_quant_config: Optional[QuantConfig]) -> None: + """Rebuild FMHA libraries and bind VisualGen-owned shared plan caches.""" + + super().update_quant_config(new_quant_config) + from ...attention.backends.fmha.prims_ts_block_sparse import PrimsTSBlockSparseFmha + + cache_state = self.metadata.get_fmha_cache_state("prims_ts_block_sparse") + for fmha in self._fmha_manager.fmha_libs: + if isinstance(fmha, PrimsTSBlockSparseFmha): + fmha.bind_plan_cache(cache_state) + # Needed to work with torch compile cause of attention metadata # make attn metadata as input for it to work @torch.compiler.disable @@ -245,6 +271,24 @@ def _concat_qkv( qkv = torch.cat([q, k, v], dim=-1) return qkv + @torch.compile + def _compact_qkv( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + batch_size: int, + seq_len: int, + kv_seq_len: int, + ): + # Separate Q, K, V stay separate - compact each into a contiguous token-major matrix. + # Slices of a fused QKV projection are strided; the compiled copy keeps them on a + # vectorized kernel, while already contiguous inputs pass through without a copy. + q = q.reshape(batch_size * seq_len, -1).contiguous() + k = k.reshape(batch_size * kv_seq_len, -1).contiguous() + v = v.reshape(batch_size * kv_seq_len, -1).contiguous() + return q, k, v + def forward( self, q: torch.Tensor, @@ -254,6 +298,7 @@ def forward( seq_len: int, attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL, seq_len_kv: Optional[int] = None, + sparse_backend_args: Optional[SparseBackendForwardArgs] = None, **kwargs, ) -> torch.Tensor: """ @@ -263,10 +308,11 @@ def forward( For diffusion models, expects: - Fused QKV: q contains [Q, K, V] concatenated, k and v are None - - does not support SageAttention + - does not support SageAttention or block-sparse routes - OR separate Q, K, V which: - for regular TRTLLM attention, will be fused internally - - for SageAttention, will be used directly + - for SageAttention, block-sparse routes, and backends that reject + fused QKV, will be passed to the core as separate tensors Args: q: Query tensor [B, S, H, D] or fused QKV [B, S, H_qkv, D] @@ -276,49 +322,69 @@ def forward( seq_len: Sequence length for Q attention_mask: Attention mask type seq_len_kv: Sequence length for K/V (for cross-attention, defaults to seq_len) + sparse_backend_args: Module-predicted sparse inputs handed to the core + prediction hooks. A ``block_sparse_inputs`` payload selects the + generic block-sparse FMHA. + **kwargs: ``timestep`` only; other names are rejected. Returns: Output tensor [B, S, H*D] """ - kv_seq_len = seq_len_kv if seq_len_kv is not None else seq_len - prepared_metadata = self._prepare_metadata(batch_size, seq_len) timestep = kwargs.pop("timestep", None) + if kwargs: + unexpected_names = ", ".join(sorted(kwargs)) + raise TypeError( + f"Unexpected TRTLLM attention forward keyword arguments: {unexpected_names}" + ) - if self.quant_attention_config is not None: - assert k is not None and v is not None, ( - "SageAttention requires separate Q, K, V tensors" + block_sparse_inputs = ( + sparse_backend_args.block_sparse_inputs if sparse_backend_args is not None else None + ) + use_separate_qkv = ( + block_sparse_inputs is not None + or self.quant_attention_config is not None + or not self.support_fused_qkv() + ) + if use_separate_qkv and (k is None or v is None): + raise ValueError("This TRTLLM attention call requires separate q, k, and v tensors.") + if block_sparse_inputs is not None and self.quant_attention_config is not None: + raise ValueError( + "Generic block-sparse attention does not support quant_attention_config." ) + + kv_seq_len = seq_len_kv if seq_len_kv is not None else seq_len + prepared_metadata = self._prepare_metadata(batch_size, seq_len) + sage_kwargs = {} + if use_separate_qkv: + q, 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_backend_args=sparse_backend_args, + **sage_kwargs, + ), + ) + return output.view(batch_size, seq_len, -1) @property def preferred_layout(self) -> AttentionTensorLayout: diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py index d6530132b267..8eaea491890e 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/utils.py @@ -45,8 +45,7 @@ def get_visual_gen_attention_backend( Backend Selection Guide: - "VANILLA": Full support for cross-attention (different Q/KV seq lengths) Uses torch SDPA backend - - "TRTLLM": Optimized for self-attention (requires same Q/KV seq lengths) - Better performance but requires fused QKV + - "TRTLLM": Optimized for self-attention (requires same Q/KV seq lengths). - "FLASHINFER": Dense prefill attention without a KV cache. Supports FP16/BF16 and architecture-specific NVFP4 attention recipes. - "FA4": Flash Attention 4; provides higher speedup on Blackwell GPUs (sm100) @@ -54,27 +53,33 @@ def get_visual_gen_attention_backend( - "CUTEDSL": CuTe DSL kernels. create_attention selects dense/SkipSoftmax FMHA or VSA from AttentionConfig.sparse_attention_config. """ - # Lazy imports to avoid circular dependency - from .cute_dsl import CuTeDSLAttention - from .flash_attn4 import FlashAttn4Attention - from .flashinfer import FlashInferAttention - from .trtllm import TrtllmAttention - from .vanilla import VanillaAttention backend_name = backend_name.upper() if backend_name == "VANILLA": + from .vanilla import VanillaAttention + return VanillaAttention elif backend_name == "TRTLLM": + from .trtllm import TrtllmAttention + return TrtllmAttention elif backend_name == "FLASHINFER": + from .flashinfer import FlashInferAttention + return FlashInferAttention elif backend_name == "FA4": + from .flash_attn4 import FlashAttn4Attention + return FlashAttn4Attention elif backend_name == "CUTEDSL": + from .cute_dsl import CuTeDSLAttention + return CuTeDSLAttention else: # Default to VANILLA for maximum compatibility + from .vanilla import VanillaAttention + return VanillaAttention @@ -112,42 +117,64 @@ def create_attention( will automatically reallocate if longer sequences are encountered. attention_config: Optional AttentionConfig used to select the attention algorithm and forward its quantization or sparsity configuration. - attention_metadata_state: Optional model-scoped metadata state from - visual-gen config. Required for TRTLLM and shared by FlashInfer layers. + attention_metadata_state: Optional per-component VisualGen attention state. + It keeps shape-stable attention metadata alive across layers and + CUDA Graph captures. Required for TRTLLM and shared by FlashInfer layers. **kwargs: Additional backend-specific arguments Returns: AttentionBackend instance """ - attn_cls = get_visual_gen_attention_backend(backend) + sparse_attention_config = ( + attention_config.sparse_attention_config if attention_config is not None else None + ) + is_vsa = ( + sparse_attention_config is not None + and getattr(sparse_attention_config, "algorithm", None) == "vsa" + ) + is_sol = ( + sparse_attention_config is not None + and getattr(sparse_attention_config, "algorithm", None) == "sol_attn" + ) + + backend_name = backend.upper() + if is_vsa and backend_name == "CUTEDSL": + from .sparse.vsa.backend import VSACuTeDSLAttention + + attn_cls = VSACuTeDSLAttention + elif is_vsa and backend_name == "TRTLLM": + from .sparse.vsa.backend import VSATrtllmAttention + + attn_cls = VSATrtllmAttention + elif is_sol and backend_name == "TRTLLM": + from .sparse.sol.backend import SOLTrtllmAttention + + attn_cls = SOLTrtllmAttention + else: + attn_cls = get_visual_gen_attention_backend(backend) + + if is_vsa: + sparse_params = kwargs.pop("sparse_params", None) + if sparse_params is not None: + raise ValueError("VSA does not lower through core SparseParams.") # Forward the validated quantization recipe to TRTLLM, FlashInfer, or the dense CuTe DSL FMHA backend. if attention_config is not None and attention_config.quant_attention_config is not None: kwargs["quant_attention_config"] = attention_config.quant_attention_config - if backend.upper() == "TRTLLM": + if backend_name == "TRTLLM": if attention_metadata_state is None: raise ValueError( "TRTLLM backend requires `attention_metadata_state` from " "DiffusionModelConfig; creation path must not allocate metadata implicitly." ) kwargs["attention_metadata_state"] = attention_metadata_state - elif backend.upper() == "FLASHINFER": + elif backend_name == "FLASHINFER": if attention_metadata_state is None: raise ValueError( "FLASHINFER backend requires `attention_metadata_state` from " "DiffusionModelConfig for shared workspace allocation." ) kwargs["attention_metadata_state"] = attention_metadata_state - if backend.upper() == "CUTEDSL" and attention_config is not None: - if ( - attention_config.sparse_attention_config is not None - and getattr(attention_config.sparse_attention_config, "algorithm", None) == "vsa" - ): - from .cute_dsl.vsa import VSAAttention - - attn_cls = VSAAttention - kwargs["sparse_attention_config"] = attention_config.sparse_attention_config - return attn_cls( layer_idx=layer_idx, num_heads=num_heads, diff --git a/tensorrt_llm/_torch/visual_gen/config.py b/tensorrt_llm/_torch/visual_gen/config.py index f3d45e0d651d..ba19e57be664 100644 --- a/tensorrt_llm/_torch/visual_gen/config.py +++ b/tensorrt_llm/_torch/visual_gen/config.py @@ -74,8 +74,24 @@ def discover_pipeline_components(checkpoint_path: Path) -> Dict[str, Path]: def create_attention_metadata_state() -> Dict[str, Any]: - """Create model-scoped state shared by visual-gen attention layers.""" - return {"metadata_cache": {}} + """Create state shared by attention layers in one model component. + + The state outlives individual forwards and CUDA Graph captures and owns the + shape-keyed TRTLLM metadata and PrimTS plan caches. VisualGen attention + layers execute serially within one component, so sharing graph-stable route + workspaces avoids retaining one worst-case allocation per layer. Each model + component receives a distinct state and must not execute concurrent forwards. + """ + return { + "metadata_cache": {}, + "fmha_caches": { + "prims_ts_block_sparse": { + "contiguous_wrappers": {}, + "paged_wrappers": {}, + }, + }, + "sparse_predictors": {}, + } def _model_config_value(value: Any, *, deep_copy: bool = True) -> Any: @@ -123,6 +139,7 @@ class DiffusionModelConfig(_VisualGenConfigBase): cuda_graph: CudaGraphConfig = PydanticField(default_factory=CudaGraphConfig) cpu_offload_config: CpuOffloadConfig = PydanticField(default_factory=CpuOffloadConfig) attention: AttentionConfig = PydanticField(default_factory=AttentionConfig) + # Per-component metadata cache shared by VisualGen TRTLLM attention layers. attention_metadata_state: Optional[Dict[str, Any]] = None parallel: ParallelConfig = PydanticField(default_factory=ParallelConfig) cache: Optional[CacheConfig] = None @@ -196,6 +213,7 @@ class DiffusionPipelineConfig(_VisualGenConfigBase): cuda_graph: CudaGraphConfig = PydanticField(default_factory=CudaGraphConfig) cpu_offload_config: CpuOffloadConfig = PydanticField(default_factory=CpuOffloadConfig) attention: AttentionConfig = PydanticField(default_factory=AttentionConfig) + # Seed state copied into each model component before attention metadata is created. attention_metadata_state: Optional[Dict[str, Any]] = None parallel: ParallelConfig = PydanticField(default_factory=ParallelConfig) cache: Optional[CacheConfig] = None diff --git a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/interface.py b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/interface.py index 4bee4dad352a..b0438b3969bd 100644 --- a/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/interface.py +++ b/tensorrt_llm/_torch/visual_gen/cute_dsl_kernels/blackwell/video_sparse_attention/interface.py @@ -17,7 +17,7 @@ Blackwell (sm_100) fast path for VSA's fine stage. The kernel JIT-compiles on first call and is cached per process; the caller -(CuTeDSLAttention._forward_vsa) falls back to dense SDPA when the +(VSACuTeDSLAttention) falls back to dense SDPA when the device/dtype/head_dim envelope is not met. """ diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py index d7223d91b413..ede5329319f7 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py @@ -1106,6 +1106,7 @@ def _setup_cuda_graphs(self): self._current_lora_cuda_graph_state, lambda: self.transformer.active_topology, ) + self.transformer.register_cuda_graph_extra_key_fns(runner) compile_note = " (with torch.compile)" if self.pipeline_config.torch_compile.enable else "" logger.info( "CUDA graph runner: wrapping LTX-2 two-stage transformer.forward " diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py index 7cba785b1bf2..84338be4cc32 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py @@ -36,6 +36,7 @@ UlyssesAttention, wrap_parallel_attention, ) +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.params import SolParams from tensorrt_llm._torch.visual_gen.attention_backend.utils import create_attention from tensorrt_llm._torch.visual_gen.models.modeling import BaseDiffusionModel from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode @@ -928,6 +929,8 @@ def forward( text_kv_video: tuple[torch.Tensor, torch.Tensor] | None = None, text_kv_audio: tuple[torch.Tensor, torch.Tensor] | None = None, step_index=None, + video_sol_timestep: torch.Tensor | None = None, + audio_sol_timestep: torch.Tensor | None = None, ) -> tuple[TransformerArgs | None, TransformerArgs | None]: """Forward with optional perturbation masking for STG. @@ -938,6 +941,8 @@ def forward( Required when the video stream runs cross-attn — built by ``LTXModel.prepare_text_cache``. text_kv_audio: Pre-projected (K, V) for audio text cross-attention. + video_sol_timestep: Raw video modality timestep for SOL phase preparation. + audio_sol_timestep: Raw audio modality timestep for SOL phase preparation. """ if video is None and audio is None: raise ValueError("At least one of video or audio must be provided") @@ -991,8 +996,14 @@ def forward( self._fuse_adaln, fp4_input_scale=get_nvfp4_self_attn_input_scale(self.attn1), ) + v_attn_timestep = ( + video_sol_timestep + if isinstance(self.attn1.sparse_params, SolParams) + and video_sol_timestep is not None + else video.timesteps + ) v_attn_raw = self.attn1( - norm_vx, pe=video.positional_embeddings, timestep=video.timesteps + norm_vx, pe=video.positional_embeddings, timestep=v_attn_timestep ) if has_perturbations and perturbations.any_in_batch( PerturbationType.SKIP_VIDEO_SELF_ATTN, self.idx @@ -1045,11 +1056,17 @@ def forward( self._fuse_adaln, fp4_input_scale=get_nvfp4_self_attn_input_scale(self.audio_attn1), ) + a_attn_timestep = ( + audio_sol_timestep + if isinstance(self.audio_attn1.sparse_params, SolParams) + and audio_sol_timestep is not None + else audio.timesteps + ) a_attn_raw = self.audio_attn1( norm_ax, pe=audio.positional_embeddings, key_padding_mask=audio.audio_padding_mask, - timestep=audio.timesteps, + timestep=a_attn_timestep, ) if has_perturbations and perturbations.any_in_batch( PerturbationType.SKIP_AUDIO_SELF_ATTN, self.idx @@ -2255,6 +2272,9 @@ def forward( if not self.model_type.is_audio_enabled() and audio is not None: raise ValueError("Audio is not enabled for this model") + video_sol_timestep = video.timesteps if video is not None else None + audio_sol_timestep = audio.timesteps if audio is not None else None + # Audio padding for Ulysses: when self._audio_pad > 0 (set once by # configure_audio_ulysses to make T_a divisible by ulysses_size), pad # audio on entry to make it shardable. Build a [B, T_a_padded] bool mask @@ -2333,6 +2353,8 @@ def forward( ax, perturbations=perturbations, step_index=step_index, + video_sol_timestep=video_sol_timestep, + audio_sol_timestep=audio_sol_timestep, ) if video_args is not None and vx is not None: video_args = replace(video_args, x=vx) @@ -2347,6 +2369,8 @@ def forward( text_kv_video=v_kv[i] if v_kv else None, text_kv_audio=a_kv[i] if a_kv else None, step_index=step_index, + video_sol_timestep=video_sol_timestep, + audio_sol_timestep=audio_sol_timestep, ) # Gather sequences back to full length for output processing. diff --git a/tensorrt_llm/_torch/visual_gen/models/modeling.py b/tensorrt_llm/_torch/visual_gen/models/modeling.py index 9fd8a02187ea..30ea64ec7a08 100644 --- a/tensorrt_llm/_torch/visual_gen/models/modeling.py +++ b/tensorrt_llm/_torch/visual_gen/models/modeling.py @@ -20,8 +20,9 @@ import torch.nn as nn from tensorrt_llm._torch.attention.backends.sparse.skip_softmax import SkipSoftmaxScheduler +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.params import SolParams from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig -from tensorrt_llm.visual_gen.sparse_attention import SkipSoftmaxAttentionConfig +from tensorrt_llm.visual_gen.sparse_attention import SkipSoftmaxAttentionConfig, SolAttentionConfig if TYPE_CHECKING: from tensorrt_llm._torch.visual_gen.cuda_graph_runner import CUDAGraphRunner @@ -74,6 +75,38 @@ def register_cuda_graph_extra_key_fns(self, runner: "CUDAGraphRunner") -> None: the shared registrations. """ sparse_config = self.model_config.attention.sparse_attention_config + if isinstance(sparse_config, SolAttentionConfig): + disabled_until_timestep = sparse_config.disabled_until_timestep + if disabled_until_timestep is None: + return + + def sol_phase_key(*args, **kwargs): + del args + modality_phases = tuple( + (name, phase) + for name in ("video", "audio") + if (modality := kwargs.get(name)) is not None + and hasattr(modality, "timesteps") + and ( + phase := SolParams.get_graph_phase_for_timestep( + modality.timesteps, + disabled_until_timestep=disabled_until_timestep, + ) + ) + is not None + ) + if modality_phases: + return modality_phases + return SolParams.get_graph_phase_for_timestep( + kwargs.get("timestep"), + disabled_until_timestep=disabled_until_timestep, + ) + + # SOL backend warmup prepares the matching phase before capture. Key + # each modality's live timestep so dense and sparse graphs stay separate. + runner.register_extra_key_fn("sol_attn_phase", sol_phase_key) + return + if not isinstance(sparse_config, SkipSoftmaxAttentionConfig): return diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py index ee730932b5cb..2400a3d6728b 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py @@ -26,7 +26,7 @@ from diffusers.video_processor import VideoProcessor from transformers import AutoTokenizer, UMT5EncoderModel -from tensorrt_llm._torch.visual_gen.attention_backend import ( +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import ( VSAMetadataBuilder, set_vsa_forward_context, ) @@ -115,6 +115,7 @@ hf_ids=[ "Wan-AI/Wan2.1-T2V-1.3B-Diffusers", "Wan-AI/Wan2.1-T2V-14B-Diffusers", + "FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers", "Wan-AI/Wan2.2-T2V-A14B-Diffusers", "Wan-AI/Wan2.2-TI2V-5B-Diffusers", "nvidia/Wan2.2-T2V-A14B-Diffusers-FP8", @@ -148,6 +149,16 @@ def __init__(self, pipeline_config): ) super().__init__(pipeline_config) + # CUDA graphs capture the VSA partition/mask tensor addresses. Keep + # their shape cache alive across requests so replay never references + # metadata owned by a completed forward call. + self._vsa_metadata_builder = VSAMetadataBuilder() + + def cleanup(self): + """Release CUDA graphs before clearing captured attention state.""" + + super().cleanup() + self._vsa_metadata_builder.clear() def _compute_wan_timestep_embedding(self, module, timestep=None, **kwargs): """Compute timestep embedding for WAN transformer. @@ -608,15 +619,15 @@ def forward( f"guidance_scale={guidance_scale}, guidance_scale_2={guidance_scale_2}" ) - # VSA: build metadata builder once per forward() call; reused across timesteps. + # VSA metadata is cached at pipeline scope and reused across requests. _attn_cfg = self.pipeline_config.primary_model_config.attention _sparse_cfg = getattr(_attn_cfg, "sparse_attention_config", None) _vsa_active = ( - getattr(_attn_cfg, "backend", "VANILLA") == "CUTEDSL" + getattr(_attn_cfg, "backend", "VANILLA") in ("CUTEDSL", "TRTLLM") and _sparse_cfg is not None and getattr(_sparse_cfg, "algorithm", None) == "vsa" ) - _vsa_builder = VSAMetadataBuilder() if _vsa_active else None + _vsa_builder = self._vsa_metadata_builder if _vsa_active else None _vsa_patch_size = tuple(getattr(self.config, "patch_size", [1, 2, 2])) # (pT, pH, pW) _vsa_sparsity = _sparse_cfg.vsa_sparsity if _vsa_active else 0.0 diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py b/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py index b717ff98abb5..4f3359c8c468 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py @@ -351,6 +351,7 @@ def __init__( config=model_config, layer_idx=_layer_idx, async_ulysses=self._use_async_ulysses, + separate_qkv_is_self_attention=True, module_name=f"blocks.{_layer_idx}.attn1", ) @@ -406,7 +407,7 @@ def __init__( reduce_output=(tp_size != 1), ) - # VSA gates (CUTEDSL backend, sparse_attention_config.algorithm == "vsa"). + # VSA gates are shared by the backend-specific fine-stage implementations. # G_c weights the coarse branch; G_f weights the fine branch. self.to_gate_compress = None self.to_gate_fine = None @@ -414,7 +415,6 @@ def __init__( _sa_cfg = getattr(_attn_cfg, "sparse_attention_config", None) if _attn_cfg else None _is_vsa = ( _attn_cfg is not None - and getattr(_attn_cfg, "backend", "VANILLA") == "CUTEDSL" and _sa_cfg is not None and getattr(_sa_cfg, "algorithm", None) == "vsa" ) @@ -432,6 +432,10 @@ def __init__( force_dynamic_quantization=force_dynamic_quant, tensor_parallel_mode=gate_tp_mode, reduce_output=False, + override_tp_sharding=( + self.attn1.local_q_dim_start, + self.attn1.local_q_dim_end, + ), ) self.to_gate_fine = Linear( hidden_size, @@ -444,6 +448,10 @@ def __init__( force_dynamic_quantization=force_dynamic_quant, tensor_parallel_mode=gate_tp_mode, reduce_output=False, + override_tp_sharding=( + self.attn1.local_q_dim_start, + self.attn1.local_q_dim_end, + ), ) # I2V: Additional K/V projections for image embeddings. @@ -588,7 +596,12 @@ def forward( # so each V/Q/K GEMM + norm + RoPE overlaps with the peer push on the # side stream; both paths return 3D [B, S, H*D]. if self._use_async_ulysses: - attn1_out = self.attn1.forward_async(normed, freqs=freqs, timestep=timestep) + attn1_out = self.attn1.forward_async( + normed, + freqs=freqs, + timestep=timestep, + **attn1_kwargs, + ) else: attn1_out = self.attn1(normed, freqs=freqs, timestep=timestep, **attn1_kwargs) diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/modules/attention.py index 8c19484229c7..b77bea91af3a 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py @@ -1,3 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from enum import Enum from typing import Optional, Tuple @@ -5,7 +20,7 @@ import torch.nn as nn from tensorrt_llm.logger import logger -from tensorrt_llm.visual_gen.sparse_attention import SkipSoftmaxAttentionConfig +from tensorrt_llm.visual_gen.sparse_attention import SkipSoftmaxAttentionConfig, SolAttentionConfig from ...modules.linear import Linear, TensorParallelMode, WeightMode, WeightsLoadingConfig from ...utils import Fp4QuantizedTensor @@ -97,16 +112,27 @@ def __init__( cp_size = vgm.cp_size if vgm else 1 base_backend = config.attention.backend _sa_cfg = config.attention.sparse_attention_config - _is_vsa = ( - base_backend == "CUTEDSL" - and _sa_cfg is not None - and getattr(_sa_cfg, "algorithm", None) == "vsa" + is_vsa = _sa_cfg is not None and getattr(_sa_cfg, "algorithm", None) == "vsa" + is_sol = _sa_cfg is not None and getattr(_sa_cfg, "algorithm", None) == "sol_attn" + is_sparse_self_attention = is_vsa or is_sol + + is_separate_qkv = self.qkv_mode == QKVMode.SEPARATE_QKV + is_separate_qkv_cross_attention = is_separate_qkv and not separate_qkv_is_self_attention + use_vanilla_cross_attention = is_separate_qkv and ( + (base_backend == "TRTLLM" and not is_sparse_self_attention) + or (is_sparse_self_attention and not separate_qkv_is_self_attention) ) - # Cross-attention fallback: TRTLLM and CUTEDSL VSA are self-attn only. - if self.qkv_mode == QKVMode.SEPARATE_QKV and (base_backend == "TRTLLM" or _is_vsa): + # Cross-attention fallback: dense TRTLLM, VSA, and SOL are self-attn only. + if use_vanilla_cross_attention: backend_name = "VANILLA" - requested = f"{base_backend} (VSA)" if _is_vsa else base_backend + requested = ( + f"{base_backend} (VSA)" + if is_vsa + else f"{base_backend} (SOL)" + if is_sol + else base_backend + ) # Warn once per (module class, requested, resolved) triple so the # fallback is visible without per-module-instance log spam. logger.warning_once( @@ -117,12 +143,18 @@ def __init__( else: backend_name = base_backend - if _is_vsa and cp_size > 1: + if is_vsa and cp_size > 1: raise ValueError( f"VSA needs the full token sequence per rank, so it is incompatible " f"with context parallelism (Attention2D/Ring, cp_size={cp_size}). Use " f"ulysses or cfg parallelism instead." ) + if is_sol and cp_size > 1: + raise ValueError( + f"SOL needs the full token sequence per rank, so it is incompatible " + f"with context parallelism (Attention2D/Ring, cp_size={cp_size}). Use " + f"ulysses or cfg parallelism instead." + ) self.attn_backend = backend_name self.qk_norm = qk_norm self.qk_norm_mode = qk_norm_mode @@ -232,6 +264,8 @@ def __init__( module_name=self.module_name, pretrained_config=config.pretrained_config, ) + elif isinstance(ss_cfg, SolAttentionConfig) and backend_name == "TRTLLM": + sparse_params = ss_cfg.to_sparse_params() self.sparse_params = sparse_params # Create compute backend @@ -248,12 +282,7 @@ def __init__( sparse_params=sparse_params, ) - if ( - enable_sequence_parallel - and self.qkv_mode == QKVMode.SEPARATE_QKV - and not separate_qkv_is_self_attention - and vgm is not None - ): + if enable_sequence_parallel and is_separate_qkv_cross_attention and vgm is not None: ring_size = vgm.ring_size if ring_size > 1: raise ValueError( @@ -645,6 +674,7 @@ def forward_async( hidden_states: torch.Tensor, freqs: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, timestep: Optional[torch.Tensor] = None, + **kwargs, ) -> torch.Tensor: """Async-Ulysses self-attn driver. Structurally mirrors ``forward``: each closure does ``to_{q,k,v}`` + (optional) fused norm+RoPE on the @@ -684,8 +714,8 @@ def forward_async( ) B, S = hidden_states.shape[:2] - H = self.num_attention_heads - KV = self.num_key_value_heads + H = self.local_num_attention_heads + KV = self.local_num_key_value_heads D = self.head_dim # Mirrors forward()'s fused gate. qkv_mode is implicitly SEPARATE_QKV # under async (caller-enforced), so the FUSE_QKV check in forward() @@ -741,6 +771,17 @@ def compute_k(): def compute_v(): return self.to_v(qkv_input).view(B, S, KV, D) - out_4d = self.attn.forward_async(compute_q, compute_k, compute_v, timestep=timestep) + for gate_key in ("gate_compress", "gate_fine"): + gate = kwargs.get(gate_key) + if gate is not None: + kwargs[gate_key] = gate.view(B, S, self.local_num_attention_heads, D) + + out_4d = self.attn.forward_async( + compute_q, + compute_k, + compute_v, + timestep=timestep, + **kwargs, + ) b, t = out_4d.shape[:2] return self.to_out[0](out_4d.reshape(b, t, H * D)) diff --git a/tensorrt_llm/_torch/visual_gen/pipeline_loader.py b/tensorrt_llm/_torch/visual_gen/pipeline_loader.py index b1d95d284105..d63d7d087a7b 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline_loader.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline_loader.py @@ -1,3 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """ Model loader for diffusion pipelines. @@ -228,13 +243,22 @@ def load( _attn_backend = config.attention.backend _sa_cfg = config.attention.sparse_attention_config if ( - _attn_backend == "CUTEDSL" + _attn_backend in ("CUTEDSL", "TRTLLM") and _sa_cfg is not None and getattr(_sa_cfg, "algorithm", None) == "vsa" ): - kernel_path = "CuTe DSL block-sparse" if CUTE_AVAILABLE else "dense SDPA fallback" + if _attn_backend == "CUTEDSL": + kernel_path = ( + "CuTe DSL block-sparse when supported; dense SDPA fallback otherwise" + if CUTE_AVAILABLE + else "dense SDPA fallback" + ) + else: + kernel_path = ( + "PrimTS block-sparse when supported; compact dense TRTLLM fallback otherwise" + ) logger.info( - f"Attention backend: CUTEDSL (algorithm=vsa, " + f"Attention backend: {_attn_backend} (algorithm=vsa, " f"sparsity={_sa_cfg.vsa_sparsity}, fine-stage={kernel_path})" ) else: diff --git a/tensorrt_llm/visual_gen/__init__.py b/tensorrt_llm/visual_gen/__init__.py index ec714e3d2bfe..213a8e2e47b4 100644 --- a/tensorrt_llm/visual_gen/__init__.py +++ b/tensorrt_llm/visual_gen/__init__.py @@ -49,6 +49,7 @@ QuantAttentionConfig, RuntimeLoRAConfig, SkipSoftmaxAttentionConfig, + SolAttentionConfig, SparseAttentionConfig, TeaCacheConfig, TorchCompileConfig, @@ -76,6 +77,7 @@ "QuantAttentionConfig": "tensorrt_llm.visual_gen.args", "RuntimeLoRAConfig": "tensorrt_llm.visual_gen.args", "SkipSoftmaxAttentionConfig": "tensorrt_llm.visual_gen.args", + "SolAttentionConfig": "tensorrt_llm.visual_gen.args", "SparseAttentionConfig": "tensorrt_llm.visual_gen.args", "TeaCacheConfig": "tensorrt_llm.visual_gen.args", "TorchCompileConfig": "tensorrt_llm.visual_gen.args", @@ -133,6 +135,7 @@ def __dir__(): "RuntimeLoRAConfig", "SparseAttentionConfig", "SkipSoftmaxAttentionConfig", + "SolAttentionConfig", "VideoSparseAttentionConfig", "VAEConfig", "CacheConfig", diff --git a/tensorrt_llm/visual_gen/args.py b/tensorrt_llm/visual_gen/args.py index 6b2ee30ccdca..e4f3c3e9982d 100644 --- a/tensorrt_llm/visual_gen/args.py +++ b/tensorrt_llm/visual_gen/args.py @@ -31,7 +31,11 @@ from tensorrt_llm.llmapi.utils import StrictBaseModel, set_api_status from tensorrt_llm.models.modeling_utils import QuantConfig -from .sparse_attention import SkipSoftmaxAttentionConfig, VideoSparseAttentionConfig +from .sparse_attention import ( + SkipSoftmaxAttentionConfig, + SolAttentionConfig, + VideoSparseAttentionConfig, +) # ============================================================================= # Type aliases @@ -95,7 +99,7 @@ class QuantAttentionConfig(StrictBaseModel): # Discriminated union of sparse attention configs. SparseAttentionConfig = Annotated[ - Union[SkipSoftmaxAttentionConfig, VideoSparseAttentionConfig], + Union[SkipSoftmaxAttentionConfig, VideoSparseAttentionConfig, SolAttentionConfig], Field(discriminator="algorithm"), ] @@ -122,7 +126,8 @@ class AttentionConfig(StrictBaseModel): status="prototype", description=( "Sparse attention recipe. Discriminated by algorithm: " - "skip_softmax (TRTLLM / CUTEDSL backends) or VSA (CUTEDSL backend)." + "skip_softmax (TRTLLM / CUTEDSL backends), VSA (CUTEDSL / TRTLLM backends), " + "or SOL (TRTLLM backend)." ), ) @@ -205,7 +210,8 @@ def _validate_sparse_attention_config(self) -> "AttentionConfig": algo = self.sparse_attention_config.algorithm supported_backends = { "skip_softmax": ("TRTLLM", "CUTEDSL"), - "vsa": ("CUTEDSL",), + "vsa": ("CUTEDSL", "TRTLLM"), + "sol_attn": ("TRTLLM",), }.get(algo) if supported_backends is None: return self @@ -220,21 +226,16 @@ def _validate_sparse_attention_config(self) -> "AttentionConfig": return self @model_validator(mode="after") - def _validate_cutedsl_quant_sparse_mutex(self) -> "AttentionConfig": - # VSA replaces the dense CuTeDSL path and cannot compose with quantized - # attention. SkipSoftmax is part of that dense path and can compose. - if ( - self.backend == "CUTEDSL" - and self.quant_attention_config is not None - and self.sparse_attention_config is not None - and self.sparse_attention_config.algorithm == "vsa" - ): - raise ValueError( - "CUTEDSL backend: quant_attention_config and VSA " - "sparse_attention_config are mutually exclusive (the " - "CuTeDSLAttention dispatcher selects either the dense path " - "or the sparse VSA path, not both)." - ) + def _validate_quant_sparse_mutex(self) -> "AttentionConfig": + if self.quant_attention_config is None or self.sparse_attention_config is None: + return self + + if self.sparse_attention_config.algorithm == "vsa": + # VSA consumes the unquantized Q/K/V path, so accepting an attention + # quantization recipe would silently ignore user configuration. + raise ValueError("VSA and quant_attention_config are mutually exclusive.") + if self.sparse_attention_config.algorithm == "sol_attn": + raise ValueError("SOL and quant_attention_config are mutually exclusive.") return self @@ -773,6 +774,21 @@ def _normalize_quant_config(cls, data: Any) -> Any: data = {**data, "quant_config": QuantConfig()} return data + @model_validator(mode="after") + def _validate_sol_fullgraph(self) -> "VisualGenArgs": + sparse_config = self.attention_config.sparse_attention_config + if ( + isinstance(sparse_config, SolAttentionConfig) + and self.torch_compile_config.enable + and self.torch_compile_config.enable_fullgraph + ): + raise ValueError( + "SOL sparse attention does not support torch.compile fullgraph; " + "set torch_compile_config.enable_fullgraph=False or disable " + "torch.compile." + ) + return self + @property def cache_backend(self) -> Optional[CacheBackendName]: return self.cache_config.cache_backend if self.cache_config is not None else None # type: ignore[return-value] @@ -824,6 +840,7 @@ def from_yaml(cls, yaml_path: Union[str, Path], **overrides: Any) -> "VisualGenA "QuantAttentionConfig", "SparseAttentionConfig", "SkipSoftmaxAttentionConfig", + "SolAttentionConfig", "VideoSparseAttentionConfig", "AttentionConfig", "VAEConfig", diff --git a/tensorrt_llm/visual_gen/sparse_attention.py b/tensorrt_llm/visual_gen/sparse_attention.py index 2bccb10e2652..a06c8d9c70e7 100644 --- a/tensorrt_llm/visual_gen/sparse_attention.py +++ b/tensorrt_llm/visual_gen/sparse_attention.py @@ -19,10 +19,33 @@ from typing import Any, Dict, Literal, Optional from pydantic import Field as PydanticField +from pydantic import field_validator from tensorrt_llm.llmapi.utils import StrictBaseModel +def _parse_dense_layers(spec: Optional[str]) -> frozenset[int]: + """Parse comma-separated layer indices and inclusive ranges.""" + layers: set[int] = set() + for raw_item in (spec or "").split(","): + item = raw_item.strip() + if not item: + if spec and spec.strip(): + raise ValueError("dense_layers must not contain empty entries") + continue + bounds = item.split("-") + if len(bounds) not in (1, 2) or not all(bounds): + raise ValueError(f"invalid dense_layers entry: {item!r}") + try: + start, end = int(bounds[0]), int(bounds[-1]) + except ValueError as error: + raise ValueError(f"invalid dense_layers entry: {item!r}") from error + if start < 0 or end < start: + raise ValueError(f"invalid dense_layers range: {item!r}") + layers.update(range(start, end + 1)) + return frozenset(layers) + + class BaseSparseAttentionConfig(StrictBaseModel): """Base for visual-generation sparse attention configs. @@ -221,12 +244,66 @@ def _ckpt_sparse_attention_config_from_kwargs( return None +class SolAttentionConfig(BaseSparseAttentionConfig): + """Two-stage SOL sparse attention backed by TRTLLM PrimTS. + + TRTLLM predicts an exact block mask from compact Q/K/V tensors, then the + generic block-sparse FMHA executes that mask. Unsupported runtime tensor + envelopes raise instead of silently falling back to dense attention. + """ + + algorithm: Literal["sol_attn"] = PydanticField( + "sol_attn", + description="Sparse attention algorithm discriminator.", + ) + tau: float = PydanticField( + 1.0, + allow_inf_nan=False, + description="Standard-deviation multiplier used by the SOL block selector.", + ) + disabled_until_timestep: Optional[float] = PydanticField( + None, + gt=0.0, + le=1.0, + description=( + "Dense-prefix cutoff on the normalized denoising timestep. The " + "attention switches to SOL when timestep is below this value." + ), + ) + dense_layers: Optional[str] = PydanticField( + None, + description=( + "Comma-separated layer indices and inclusive ranges (for example, " + "'0,2-4') that always use dense attention." + ), + ) + + @field_validator("dense_layers") + @classmethod + def _validate_dense_layers(cls, value: Optional[str]) -> Optional[str]: + _parse_dense_layers(value) + return value + + def to_sparse_params(self, **kwargs): + """Lower the public recipe into immutable SOL runtime parameters.""" + del kwargs + from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.params import SolParams + + return SolParams( + tau=self.tau, + disabled_until_timestep=self.disabled_until_timestep, + dense_layers=_parse_dense_layers(self.dense_layers), + ) + + class VideoSparseAttentionConfig(StrictBaseModel): - """Video Sparse Attention (VSA) sparse-attention recipe (CUTEDSL backend only). + """Video Sparse Attention (VSA) sparse-attention recipe. Two-stage hybrid attention: a coarse mean-pooled stage over (4,4,4) cubes and a block-sparse fine stage over the top-K cubes selected per head. vsa_sparsity controls the fraction of cubes dropped on the fine stage. + The fine stage may run on either the CuTeDSL backend or the TRTLLM PrimTS + backend, while the user-facing sparsity semantics stay the same. """ algorithm: Literal["vsa"] = PydanticField( diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 7e10e8baaf9f..fe4de3e39fd5 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -280,10 +280,14 @@ l0_b200: - unittest/_torch/visual_gen/test_pertoken_adaln.py - unittest/_torch/visual_gen/test_attention_cute_dsl.py - unittest/_torch/visual_gen/test_fa4_cutlass_compatibility.py - - unittest/_torch/visual_gen/test_attention_cute_dsl_vsa.py + - unittest/_torch/visual_gen/test_attention_vsa.py - unittest/_torch/visual_gen/test_attention_flashinfer.py + - unittest/_torch/visual_gen/sparse_attention/test_sol_predictor.py + - unittest/_torch/visual_gen/sparse_attention/test_sol_attention.py::test_real_b200_sol_backend_cuda_graph_matches_dense_reference + - unittest/_torch/visual_gen/sparse_attention/test_sol_attention.py::test_real_b200_sol_backend_mixed_proxy_cuda_graph_matches_reference - unittest/_torch/visual_gen/test_attention_trtllm_sage.py - unittest/_torch/visual_gen/test_attention_integration.py + - unittest/_torch/visual_gen/test_trtllm_attention_metadata.py - 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_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 0327b4e498f8..11345983f054 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -59,6 +59,8 @@ l0_cpu: - unittest/_torch/visual_gen/multi_gpu/test_visual_gen_mapping.py - unittest/_torch/visual_gen/multi_gpu/test_visual_gen_multinode.py - unittest/_torch/visual_gen/sparse_attention/test_skip_softmax.py + - unittest/_torch/visual_gen/sparse_attention/test_sol_attention.py + - unittest/_torch/visual_gen/sparse_attention/test_sol_predictor.py - unittest/_torch/visual_gen/test_attention_flashinfer.py::test_flashinfer_backend_is_registered - unittest/_torch/visual_gen/test_attention_integration.py - unittest/_torch/visual_gen/test_cache_dit.py diff --git a/tests/unittest/_torch/attention/sparse/test_sparse_attention.py b/tests/unittest/_torch/attention/sparse/test_sparse_attention.py index 40a151daa252..8b012cac879b 100644 --- a/tests/unittest/_torch/attention/sparse/test_sparse_attention.py +++ b/tests/unittest/_torch/attention/sparse/test_sparse_attention.py @@ -21,11 +21,16 @@ """ from types import ModuleType -from unittest.mock import Mock +from unittest.mock import Mock, patch +import pytest import torch -from tensorrt_llm._torch.attention.backends.interface import AttentionForwardArgs +from tensorrt_llm._torch.attention.backends import trtllm as trtllm_backend +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + AttentionInputType, +) from tensorrt_llm._torch.attention.backends.sparse.hooks import ( AttentionSparseHooks, MLASparseHooks, @@ -35,8 +40,13 @@ register_attention_sparse_hooks, register_mla_sparse_hooks, ) -from tensorrt_llm._torch.attention.backends.sparse.params import SparseParams, SparseRuntimeParams -from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention +from tensorrt_llm._torch.attention.backends.sparse.params import ( + BlockSparseForwardInputs, + SparseBackendForwardArgs, + SparseParams, + SparseRuntimeParams, +) +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention, TrtllmAttentionMetadata from tensorrt_llm._torch.attention.mla import MLA @@ -72,7 +82,7 @@ def test_prepare_sparse_runtime_params_from_predictions() -> None: ) runtime_params = prepare_sparse_runtime_params( - attention, torch.empty(0), None, None, forward_args + attention, torch.empty(0), None, None, None, forward_args ) assert runtime_params.sparse_kv_indices is attention._sparse_kv_indices @@ -144,12 +154,250 @@ def test_mla_backend_only_forward_uses_default_path() -> None: ) -def test_prepare_sparse_runtime_params_without_predictions() -> None: +@pytest.mark.parametrize( + "sparse_params", [None, _StubSparseParams()], ids=["dense_backend", "sparse_backend"] +) +def test_prepare_sparse_runtime_params_without_predictions(sparse_params) -> None: attention = TrtllmAttention.__new__(TrtllmAttention) - attention.sparse_params = _StubSparseParams() + attention.sparse_params = sparse_params runtime_params = prepare_sparse_runtime_params( - attention, torch.empty(0), None, None, AttentionForwardArgs() + attention, torch.empty(0), None, None, None, AttentionForwardArgs() ) assert runtime_params == SparseRuntimeParams() + + +def test_prepare_sparse_runtime_params_runs_index_hooks_once() -> None: + attention = _StaticPredictionAttention.__new__(_StaticPredictionAttention) + attention.sparse_params = _StubSparseParams() + sparse_kv_indices = torch.tensor([1], dtype=torch.int32) + sparse_kv_offsets = torch.tensor([0, 1], dtype=torch.int32) + sparse_attn_indices = torch.tensor([2], dtype=torch.int32) + sparse_attn_offsets = torch.tensor([0, 1], dtype=torch.int32) + attention.sparse_kv_predict = Mock(return_value=(sparse_kv_indices, sparse_kv_offsets)) + attention.sparse_attn_predict = Mock(return_value=(sparse_attn_indices, sparse_attn_offsets)) + q = torch.empty((1, 4)) + k = torch.empty((1, 4)) + v = torch.empty((1, 4)) + metadata = Mock() + caller_kv_lens = torch.tensor([3]) + forward_args = AttentionForwardArgs( + sparse_runtime_params=SparseRuntimeParams(sparse_attn_kv_lens=caller_kv_lens) + ) + + runtime_params = prepare_sparse_runtime_params(attention, q, k, v, metadata, forward_args) + + assert isinstance(runtime_params, SparseRuntimeParams) + assert runtime_params.block_sparse_inputs is None + assert runtime_params.sparse_kv_indices is sparse_kv_indices + assert runtime_params.sparse_kv_offsets is sparse_kv_offsets + assert runtime_params.sparse_attn_indices is sparse_attn_indices + assert runtime_params.sparse_attn_offsets is sparse_attn_offsets + assert runtime_params.sparse_attn_indices_block_size == 1 + assert runtime_params.sparse_attn_kv_lens is caller_kv_lens + attention.sparse_kv_predict.assert_called_once_with(q, k, metadata, forward_args) + attention.sparse_attn_predict.assert_called_once_with(q, k, metadata, forward_args) + + +def test_prepare_sparse_runtime_params_schedules_skip_softmax_thresholds() -> None: + from tensorrt_llm._torch.attention.backends.sparse.skip_softmax import SkipSoftmaxParams + + attention = TrtllmAttention.__new__(TrtllmAttention) + attention.sparse_params = SkipSoftmaxParams() + scheduler = attention.sparse_params.scheduler + timestep = torch.tensor(0.5) + forward_args = AttentionForwardArgs(timestep=timestep) + + with patch.object( + scheduler, "get_runtime_params", wraps=scheduler.get_runtime_params + ) as schedule: + runtime_params = prepare_sparse_runtime_params( + attention, torch.empty(0), None, None, None, forward_args + ) + + schedule.assert_called_once_with(runtime_params=SparseRuntimeParams(), timestep=timestep) + assert runtime_params == scheduler.get_runtime_params(timestep=timestep) + + +def _make_block_sparse_inputs() -> BlockSparseForwardInputs: + return BlockSparseForwardInputs( + q_block_size=64, + kv_block_size=64, + exact_block_bits=torch.zeros((1, 1), dtype=torch.int32), + ) + + +def test_block_sparse_attn_predict_hands_through_backend_args() -> None: + attention = TrtllmAttention.__new__(TrtllmAttention) + attention.sparse_params = None + block_sparse_inputs = _make_block_sparse_inputs() + forward_args = AttentionForwardArgs( + sparse_backend_args=SparseBackendForwardArgs(block_sparse_inputs=block_sparse_inputs) + ) + + runtime_params = prepare_sparse_runtime_params( + attention, torch.empty(0), None, None, None, forward_args + ) + + assert runtime_params.block_sparse_inputs is block_sparse_inputs + assert runtime_params == SparseRuntimeParams(block_sparse_inputs=block_sparse_inputs) + + +def test_block_sparse_attn_predict_override_composes_with_index_predictors() -> None: + attention = _StaticPredictionAttention.__new__(_StaticPredictionAttention) + attention.sparse_params = _StubSparseParams() + sparse_attn_indices = torch.tensor([2], dtype=torch.int32) + sparse_attn_offsets = torch.tensor([0, 1], dtype=torch.int32) + attention.sparse_kv_predict = Mock(return_value=(None, None)) + attention.sparse_attn_predict = Mock(return_value=(sparse_attn_indices, sparse_attn_offsets)) + block_sparse_inputs = _make_block_sparse_inputs() + attention.block_sparse_attn_predict = Mock(return_value=block_sparse_inputs) + q = torch.empty((1, 4)) + k = torch.empty((1, 4)) + v = torch.empty((1, 4)) + metadata = Mock() + forward_args = AttentionForwardArgs() + + runtime_params = prepare_sparse_runtime_params(attention, q, k, v, metadata, forward_args) + + assert runtime_params.block_sparse_inputs is block_sparse_inputs + assert runtime_params.sparse_attn_indices is sparse_attn_indices + assert runtime_params.sparse_attn_indices_block_size == 1 + attention.block_sparse_attn_predict.assert_called_once_with(q, k, v, metadata, forward_args) + + +def test_attention_forward_args_default_to_empty_sparse_runtime_params() -> None: + assert AttentionForwardArgs().sparse_runtime_params == SparseRuntimeParams() + + +class _StopAfterShapeValidation(Exception): + pass + + +def _make_sparse_prediction_forward_backend() -> TrtllmAttention: + attention = _StaticPredictionAttention.__new__(_StaticPredictionAttention) + attention.sparse_params = None + attention.is_mla_enable = False + attention.num_heads = 1 + attention.num_kv_heads = 1 + attention.head_dim = 4 + attention.get_local_layer_idx = Mock(return_value=1) + attention._ensure_rope_table_size = Mock(side_effect=_StopAfterShapeValidation) + return attention + + +def _make_sparse_prediction_forward_metadata() -> TrtllmAttentionMetadata: + metadata = object.__new__(TrtllmAttentionMetadata) + seq_lens = torch.tensor([2], dtype=torch.int32) + metadata._seq_lens = seq_lens + metadata._seq_lens_kv = seq_lens + metadata._seq_lens_cuda = None + metadata.kv_cache_manager = None + metadata._max_seq_len_storage = 4 + metadata.use_paged_context_fmha = False + metadata.cu_q_seqlens = None + metadata.cu_kv_seqlens = None + metadata.enable_flash_mla = False + metadata.spec_bl_tree_first_sparse_mask_offset_kv = None + metadata.kv_lens_cuda_runtime = torch.tensor([2], dtype=torch.int32) + metadata.kv_lens_runtime = torch.tensor([2], dtype=torch.int32) + metadata.prompt_lens_cuda_runtime = torch.tensor([2], dtype=torch.int32) + metadata.prompt_lens_cpu_runtime = torch.tensor([2], dtype=torch.int32) + metadata.host_request_types_runtime = torch.tensor([0], dtype=torch.int32) + return metadata + + +def test_forward_materializes_dynamic_block_sparse_prediction_before_shape_validation() -> None: + attention = _make_sparse_prediction_forward_backend() + block_sparse_inputs = _make_block_sparse_inputs() + prediction = SparseRuntimeParams( + sparse_attn_kv_lens=torch.tensor([4]), + block_sparse_inputs=block_sparse_inputs, + ) + q = torch.empty((2, 4)) + k = torch.empty((4, 4)) + v = torch.empty((4, 4)) + metadata = _make_sparse_prediction_forward_metadata() + forward_args = AttentionForwardArgs( + output=torch.empty_like(q), + attention_input_type=AttentionInputType.context_only, + ) + + with patch.object( + trtllm_backend, "prepare_sparse_runtime_params", return_value=prediction + ) as prepare: + for _ in range(2): + with pytest.raises(_StopAfterShapeValidation): + attention.forward(q, k, v, metadata, forward_args) + + assert prepare.call_count == 2 + prepare.assert_called_with(attention, q, k, v, metadata, forward_args) + assert attention._ensure_rope_table_size.call_count == 2 + assert forward_args.sparse_runtime_params == SparseRuntimeParams() + + +@pytest.mark.parametrize("has_block_sparse_inputs", [False, True]) +def test_forward_replaces_caller_sparse_runtime_params_with_prediction( + has_block_sparse_inputs: bool, +) -> None: + attention = _make_sparse_prediction_forward_backend() + block_sparse_inputs = _make_block_sparse_inputs() if has_block_sparse_inputs else None + prediction = SparseRuntimeParams( + sparse_attn_kv_lens=torch.tensor([2]), + block_sparse_inputs=block_sparse_inputs, + ) + caller_params = SparseRuntimeParams(sparse_attn_kv_lens=torch.tensor([7])) + q = torch.empty((2, 4)) + num_kv_tokens = 4 if has_block_sparse_inputs else 2 + k = torch.empty((num_kv_tokens, 4)) + v = torch.empty((num_kv_tokens, 4)) + metadata = _make_sparse_prediction_forward_metadata() + forward_args = AttentionForwardArgs( + output=torch.empty_like(q), + attention_input_type=AttentionInputType.context_only, + sparse_runtime_params=caller_params, + ) + + with patch.object( + trtllm_backend, "prepare_sparse_runtime_params", return_value=prediction + ) as prepare: + for _ in range(2): + with pytest.raises(_StopAfterShapeValidation): + attention.forward(q, k, v, metadata, forward_args) + + assert prepare.call_count == 2 + assert attention._ensure_rope_table_size.call_count == 2 + assert forward_args.sparse_runtime_params is caller_params + + +@pytest.mark.parametrize("backend_name", ["FLASHINFER", "unknown"]) +def test_sparse_attention_backend_fallback_does_not_redispatch( + backend_name: str, monkeypatch: pytest.MonkeyPatch +) -> None: + from tensorrt_llm._torch.attention.backends import utils as attention_backend_utils + from tensorrt_llm._torch.attention.backends.sparse.skip_softmax import SkipSoftmaxParams + + monkeypatch.setattr(attention_backend_utils, "IS_FLASHINFER_AVAILABLE", False) + + with patch.object( + attention_backend_utils, + "get_trtllm_sparse_attn_attention_backend", + ) as trtllm_sparse_resolver: + backend = attention_backend_utils.get_attention_backend( + backend_name, sparse_params=SkipSoftmaxParams() + ) + + assert backend is TrtllmAttention + trtllm_sparse_resolver.assert_not_called() + + +@pytest.mark.parametrize("backend_name", ["FLASHINFER", "unknown"]) +def test_trtllm_fallback_without_sparse_params_remains_dense( + backend_name: str, monkeypatch: pytest.MonkeyPatch +) -> None: + from tensorrt_llm._torch.attention.backends import utils as attention_backend_utils + + monkeypatch.setattr(attention_backend_utils, "IS_FLASHINFER_AVAILABLE", False) + + assert attention_backend_utils.get_attention_backend(backend_name) is TrtllmAttention diff --git a/tests/unittest/_torch/attention/test_attention_op_sync.py b/tests/unittest/_torch/attention/test_attention_op_sync.py index ba5c75f3d7b6..ccf11544be53 100644 --- a/tests/unittest/_torch/attention/test_attention_op_sync.py +++ b/tests/unittest/_torch/attention/test_attention_op_sync.py @@ -41,7 +41,7 @@ import textwrap import typing from dataclasses import fields -from types import SimpleNamespace +from types import SimpleNamespace, UnionType import pytest import torch @@ -382,7 +382,19 @@ def _dataclass_field_type(cls, name: str): return None if f is None: return None - return f.type if not isinstance(f.type, str) else None + if isinstance(f.type, str): + return None + return _unwrap_optional(f.type) + + +def _unwrap_optional(py_type): + """Return the payload type for ``Optional[T]`` annotations.""" + origin = typing.get_origin(py_type) + if origin in (typing.Union, UnionType): + args = [arg for arg in typing.get_args(py_type) if arg is not type(None)] + if len(args) == 1: + return args[0] + return py_type def _resolve_path(root_cls, path: tuple[str, ...]): @@ -404,7 +416,7 @@ def _python_category(py_type) -> str: confidently (the type check is then skipped for that kwarg).""" # Unwrap Optional[X] / Union[X, None]. origin = typing.get_origin(py_type) - if origin is typing.Union: + if origin in (typing.Union, UnionType): args = [a for a in typing.get_args(py_type) if a is not type(None)] if len(args) == 1: return _python_category(args[0]) @@ -558,7 +570,7 @@ def _verify_consumed(cls, chains: set[tuple[str, ...]], excluded=frozenset()): for f in fields(cls): if f.name in excluded: continue - ftype = f.type if not isinstance(f.type, str) else None + ftype = _dataclass_field_type(cls, f.name) if ftype is not None and dataclasses.is_dataclass(ftype): sub = {p[1:] for p in chains if len(p) >= 2 and p[0] == f.name} assert sub, ( @@ -566,7 +578,7 @@ def _verify_consumed(cls, chains: set[tuple[str, ...]], excluded=frozenset()): f"declared but `{f.name}.` is never read at the " f"call site." ) - _verify_consumed(ftype, sub) + _verify_consumed(ftype, sub, excluded=excluded) else: assert f.name in consumed, ( f"Field `{f.name}` on {cls.__name__} not consumed by the " @@ -609,7 +621,7 @@ def _all_forward_args_field_names() -> set[str]: def _walk(cls) -> None: for f in fields(cls): seen.add(f.name) - ftype = f.type if not isinstance(f.type, str) else None + ftype = _dataclass_field_type(cls, f.name) if ftype is not None and dataclasses.is_dataclass(ftype): _walk(ftype) diff --git a/tests/unittest/_torch/attention/test_fmha_manager.py b/tests/unittest/_torch/attention/test_fmha_manager.py index 17dd1beda5ce..b59a5c80595b 100644 --- a/tests/unittest/_torch/attention/test_fmha_manager.py +++ b/tests/unittest/_torch/attention/test_fmha_manager.py @@ -31,6 +31,10 @@ AttentionInputType, PredefinedAttentionMask, ) +from tensorrt_llm._torch.attention.backends.sparse.params import ( + BlockSparseForwardInputs, + SparseRuntimeParams, +) from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization.mode import QuantAlgo @@ -50,12 +54,14 @@ def _make_metadata( num_generations: int, num_ctx_tokens: int = 0, use_spec_decoding: bool = False, + num_sparse_topk: int = 0, ) -> SimpleNamespace: return SimpleNamespace( num_contexts=num_contexts, num_generations=num_generations, num_ctx_tokens=num_ctx_tokens, use_spec_decoding=use_spec_decoding, + num_sparse_topk=num_sparse_topk, ) @@ -538,6 +544,53 @@ def test_fmha_cache_tracks_attention_mask_data() -> None: assert len(manager._cache) == 2 +@pytest.mark.parametrize("block_sparse_first", [False, True]) +def test_fmha_cache_separates_block_sparse_mode(block_sparse_first: bool) -> None: + events: list[tuple] = [] + attn, manager = _make_manager() + block_sparse_fmha = FakeFmha( + attn, + "block-sparse", + events, + support_predicate=lambda forward_args: ( + forward_args.sparse_runtime_params.block_sparse_inputs is not None + ), + ) + dense_fmha = FakeFmha( + attn, + "dense", + events, + support_predicate=lambda forward_args: ( + forward_args.sparse_runtime_params.block_sparse_inputs is None + ), + ) + manager.fmha_libs = [block_sparse_fmha, dense_fmha] + metadata = _make_metadata(num_contexts=1, num_generations=0, num_ctx_tokens=1) + q = torch.empty((1, 4)) + by_mode = { + False: AttentionForwardArgs(attention_input_type=AttentionInputType.context_only), + True: AttentionForwardArgs( + attention_input_type=AttentionInputType.context_only, + sparse_runtime_params=SparseRuntimeParams( + block_sparse_inputs=BlockSparseForwardInputs( + q_block_size=64, + kv_block_size=64, + exact_block_bits=torch.zeros((1, 1), dtype=torch.uint32), + ), + ), + ), + } + order = (True, False) if block_sparse_first else (False, True) + + with patch.object(fmha_manager, "_is_fmha_cache_enabled", return_value=True): + selected = { + mode: manager.select(attn, q, None, None, metadata, by_mode[mode]) for mode in order + } + + assert selected == {False: dense_fmha, True: block_sparse_fmha} + assert len(manager._cache) == 2 + + @pytest.mark.parametrize("speculative_first", [False, True]) def test_fmha_cache_separates_speculative_decoding(speculative_first: bool) -> None: events: list[tuple] = [] diff --git a/tests/unittest/_torch/attention/test_fmha_registry.py b/tests/unittest/_torch/attention/test_fmha_registry.py index cd4cdb4fd7a2..285f2aa075ba 100644 --- a/tests/unittest/_torch/attention/test_fmha_registry.py +++ b/tests/unittest/_torch/attention/test_fmha_registry.py @@ -14,10 +14,14 @@ # limitations under the License. import pytest +import torch from tensorrt_llm._torch.attention.backends.fmha import registry +from tensorrt_llm._torch.attention.backends.fmha.interface import Fmha +from tensorrt_llm._torch.attention.backends.sparse.params import SparseRuntimeParams PRIMS_TS = "prims_ts" +PRIMS_TS_BLOCK_SPARSE = "prims_ts_block_sparse" def _canonical_names() -> tuple[str, ...]: @@ -39,10 +43,31 @@ def test_default_fmha_libs_exclude_prims_ts(monkeypatch: pytest.MonkeyPatch) -> monkeypatch.delenv("TLLM_FMHA_LIBS", raising=False) assert PRIMS_TS not in registry.DEFAULT_FMHA_LIBS + assert PRIMS_TS_BLOCK_SPARSE in registry.DEFAULT_FMHA_LIBS assert set(registry.DEFAULT_FMHA_LIBS) <= set(registry.FMHA_LIBS) assert _enabled_names() == registry.DEFAULT_FMHA_LIBS +@pytest.mark.parametrize("name", [PRIMS_TS, "fallback"]) +def test_dense_fmhas_reject_unconsumed_block_sparse_inputs(name: str) -> None: + attention = type("Attention", (), {})() + fmha = object.__new__(registry.FMHA_LIBS[name]) + Fmha.__init__(fmha, attention) + forward_args = type( + "ForwardArgs", + (), + {"sparse_runtime_params": SparseRuntimeParams(block_sparse_inputs=object())}, + )() + + assert not fmha.is_supported( + torch.empty((1, 4), dtype=torch.bfloat16), + None, + None, + object(), + forward_args, + ) + + @pytest.mark.parametrize("value", ["", " ", ", ,"]) def test_empty_fmha_lib_env_uses_default( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unittest/_torch/attention/test_prims_ts_block_sparse.py b/tests/unittest/_torch/attention/test_prims_ts_block_sparse.py new file mode 100644 index 000000000000..3849b9481077 --- /dev/null +++ b/tests/unittest/_torch/attention/test_prims_ts_block_sparse.py @@ -0,0 +1,720 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from contextlib import nullcontext +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from utils.util import isSM100Family + +from tensorrt_llm._torch.attention.backends import prims_ts +from tensorrt_llm._torch.attention.backends.fmha import prims_ts_block_sparse as block_sparse_fmha +from tensorrt_llm._torch.attention.backends.fmha.interface import FmhaPhase +from tensorrt_llm._torch.attention.backends.fmha.phased import FmhaParams +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + AttentionInputType, + PredefinedAttentionMask, +) +from tensorrt_llm._torch.attention.backends.sparse.params import ( + BlockSparseForwardInputs, + SparseRuntimeParams, +) +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from tensorrt_llm.functional import PositionEmbeddingType + +pytestmark = pytest.mark.cpu_only + +_REQUIRES_PRIMTS_GPU = pytest.mark.skipif( + not isSM100Family(), + reason="PrimTS block-sparse attention requires SM100 or SM103", +) + + +def _bsr_inputs(*, kv_valid_bits: torch.Tensor | None = None): + return BlockSparseForwardInputs( + q_block_size=64, + kv_block_size=64, + max_blocks_per_row=2, + block_indptr=torch.tensor([[[0, 2]], [[2, 4]]], dtype=torch.int32), + block_indices=torch.tensor([0, 1, 2, 3], dtype=torch.int32), + kv_valid_bits=kv_valid_bits, + ) + + +def _bitmask_inputs(*, proxy: bool): + summaries = { + "k_summary": torch.zeros((2, 4, 1, 128), dtype=torch.bfloat16), + "v_summary": torch.zeros((2, 4, 1, 128), dtype=torch.bfloat16), + } + return BlockSparseForwardInputs( + q_block_size=64, + kv_block_size=64, + exact_block_bits=torch.ones((2, 1, 1, 1), dtype=torch.uint32), + **(summaries if proxy else {}), + ) + + +def _set_block_sparse_inputs( + forward_args: AttentionForwardArgs, + block_sparse_inputs, +) -> None: + forward_args.sparse_runtime_params = SparseRuntimeParams( + block_sparse_inputs=block_sparse_inputs + ) + + +def _get_block_sparse_inputs(forward_args: AttentionForwardArgs): + block_sparse_inputs = forward_args.sparse_runtime_params.block_sparse_inputs + assert block_sparse_inputs is not None + return block_sparse_inputs + + +def _pack_token_mask(mask: torch.Tensor) -> torch.Tensor: + shifts = torch.arange(32, dtype=torch.int64, device=mask.device) + weights = torch.ones_like(shifts).bitwise_left_shift_(shifts) + return (mask.view(1, -1, 32).to(torch.int64) * weights).sum(dim=-1).to(torch.uint32) + + +def _proxy_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + k_summary: torch.Tensor, + v_summary: torch.Tensor, + exact_block: int, +) -> torch.Tensor: + block_size = 64 + exact_tokens = torch.arange( + exact_block * block_size, + (exact_block + 1) * block_size, + device=q.device, + ) + proxy_blocks = [block for block in range(k_summary.shape[1]) if block != exact_block] + q_rows = q[0, :, 0].float() + exact_logits = q_rows @ k[0, exact_tokens, 0].float().T + proxy_logits = q_rows @ k_summary[0, proxy_blocks, 0].float().T + logits = torch.cat((exact_logits, proxy_logits), dim=1) / math.sqrt(q.shape[-1]) + weights = torch.exp(logits - logits.amax(dim=1, keepdim=True)) + exact_weights, proxy_weights = weights.split((block_size, len(proxy_blocks)), dim=1) + numerator = exact_weights @ v[0, exact_tokens, 0].float() + numerator += proxy_weights @ v_summary[0, proxy_blocks, 0].float() + denominator = exact_weights.sum(dim=1, keepdim=True) + denominator += proxy_weights.sum(dim=1, keepdim=True) * block_size + return (numerator / denominator).to(q.dtype)[None, :, None] + + +class _Attention: + def __init__(self) -> None: + self.sparse_params = None + self.num_heads = 2 + self.num_kv_heads = 1 + self.head_dim = 128 + self.is_mla_enable = False + self.kv_lora_rank = None + self.qk_rope_head_dim = None + self.qk_nope_head_dim = None + self.v_head_dim = None + self.q_scaling = 1.0 + self.quant_mode = 0 + self.local_layer_idx = 0 + self.position_embedding_type = PositionEmbeddingType.learned_absolute + self.attention_chunk_size = 0 + + +def _contiguous_case(): + attention = _Attention() + fmha = block_sparse_fmha.PrimsTSBlockSparseFmha(attention) + q = torch.zeros((128, 256), dtype=torch.bfloat16) + k = torch.zeros((512, 128), dtype=torch.bfloat16) + v = torch.zeros_like(k) + metadata = SimpleNamespace( + is_cross=False, + num_sparse_topk=0, + helix_position_offsets=None, + kv_cache_manager=None, + seq_lens=torch.tensor([64, 64], dtype=torch.int32), + ) + args = AttentionForwardArgs( + output=torch.empty_like(q), + attention_input_type=AttentionInputType.context_only, + attention_mask=PredefinedAttentionMask.FULL, + sparse_runtime_params=SparseRuntimeParams(block_sparse_inputs=_bsr_inputs()), + ) + return attention, fmha, q, k, v, metadata, args + + +def _paged_metadata(): + batch_size, max_pages, page_size = 2, 4, 64 + key_pages = torch.arange(batch_size * max_pages, dtype=torch.int32).view(batch_size, max_pages) + block_offsets = torch.stack((key_pages, key_pages + 8), dim=1).unsqueeze(0) + manager = Mock(spec=KVCacheManager) + manager.dtype = torch.bfloat16 + manager.num_pools = manager.num_local_layers = 1 + manager.host_kv_cache_block_offsets = block_offsets + return SimpleNamespace( + is_cross=False, + num_sparse_topk=0, + helix_position_offsets=None, + num_contexts=0, + num_generations=batch_size, + seq_lens=torch.ones(batch_size, dtype=torch.int32), + beam_width=1, + is_spec_decoding_enabled=False, + use_spec_decoding=False, + is_spec_dec_tree=False, + is_spec_dec_dynamic_tree=False, + tokens_per_block=page_size, + max_seq_len=max_pages * page_size, + kv_layout="HND", + kv_lens_runtime=torch.tensor([129, 193], dtype=torch.int32), + kv_cache_block_offsets=block_offsets, + host_kv_cache_pool_pointers=torch.tensor([[1234, 5678]], dtype=torch.int64), + host_kv_cache_pool_mapping=torch.tensor([[0, 0]], dtype=torch.int32), + kv_cache_manager=manager, + ) + + +def _paged_case(): + attention = _Attention() + fmha = block_sparse_fmha.PrimsTSBlockSparseFmha(attention) + fmha._multi_processor_count = 1 + metadata = _paged_metadata() + q = torch.zeros((2, 512), dtype=torch.bfloat16) + args = AttentionForwardArgs( + output=torch.empty((2, 256), dtype=q.dtype), + attention_input_type=AttentionInputType.generation_only, + attention_mask=PredefinedAttentionMask.CAUSAL, + attention_window_size=metadata.max_seq_len, + is_fused_qkv=True, + sparse_runtime_params=SparseRuntimeParams(block_sparse_inputs=_bsr_inputs()), + ) + return attention, fmha, q, metadata, args + + +def test_block_sparse_route_mode_is_derived_from_payload() -> None: + bsr = _bsr_inputs() + exact = _bitmask_inputs(proxy=False) + proxy = _bitmask_inputs(proxy=True) + + assert (bsr.sparse_format, bsr.use_proxy_routes) == ("bsr", False) + assert (exact.sparse_format, exact.use_proxy_routes) == ("bitmask", False) + assert (proxy.sparse_format, proxy.use_proxy_routes) == ("bitmask", True) + + +@pytest.mark.parametrize( + ("overrides", "message"), + [ + ({"block_indices": None}, "block_indptr and block_indices"), + ({"max_blocks_per_row": None}, "max_blocks_per_row"), + ( + {"exact_block_bits": torch.ones((1, 1, 1, 1), dtype=torch.uint32)}, + "exactly one route representation", + ), + ({"k_summary": torch.empty(0)}, "k_summary and v_summary"), + ], +) +def test_block_sparse_payload_rejects_ambiguous_combinations(overrides, message) -> None: + kwargs = { + "q_block_size": 64, + "kv_block_size": 64, + "max_blocks_per_row": 1, + "block_indptr": torch.tensor([[[0, 1]]], dtype=torch.int32), + "block_indices": torch.tensor([0], dtype=torch.int32), + } + kwargs.update(overrides) + + with pytest.raises((TypeError, ValueError), match=message): + BlockSparseForwardInputs(**kwargs) + + +def test_block_sparse_support_is_phase_specific_and_paged_proxy_is_rejected( + monkeypatch, +) -> None: + _attention, contiguous, q, k, v, metadata, args = _contiguous_case() + monkeypatch.setattr(contiguous, "_common_unsupported_reason", Mock(return_value=None)) + assert contiguous.is_supported(q, k, v, metadata, args, phase=FmhaPhase.CONTEXT) + assert not contiguous.is_supported(q, k, v, metadata, args, phase=FmhaPhase.GENERATION) + + _attention, paged, q, metadata, args = _paged_case() + _set_block_sparse_inputs(args, _bitmask_inputs(proxy=True)) + _supported, reason = paged._is_supported_with_reason( + q, None, None, metadata, args, phase=FmhaPhase.GENERATION + ) + assert not _supported + assert reason == "paged block-sparse attention only supports BSR exact routes" + + +def test_contiguous_proxy_routes_reject_causal_mask_before_planning(monkeypatch) -> None: + _attention, fmha, q, k, v, metadata, args = _contiguous_case() + _set_block_sparse_inputs(args, _bitmask_inputs(proxy=True)) + args.attention_mask = PredefinedAttentionMask.CAUSAL + monkeypatch.setattr(fmha, "_common_unsupported_reason", Mock(return_value=None)) + + supported, reason = fmha._is_supported_with_reason( + q, + k, + v, + metadata, + args, + phase=FmhaPhase.CONTEXT, + ) + + assert not supported + assert reason == "block-sparse proxy routes require mask_type='dense'" + + +@pytest.mark.parametrize("paged", [False, True]) +def test_block_sparse_support_rejects_invalid_static_kernel_profile( + monkeypatch, + paged, +) -> None: + if paged: + attention, fmha, q, metadata, args = _paged_case() + attention.head_dim = 64 + q = torch.zeros((2, 256), dtype=torch.bfloat16) + args.output = torch.empty((2, 128), dtype=q.dtype) + monkeypatch.setattr(fmha, "_common_unsupported_reason", Mock(return_value=None)) + supported, reason = fmha._is_supported_with_reason( + q, + None, + None, + metadata, + args, + phase=FmhaPhase.GENERATION, + ) + else: + attention, fmha, _q, _k, _v, metadata, args = _contiguous_case() + attention.head_dim = 64 + q = torch.zeros((128, 128), dtype=torch.bfloat16) + k = torch.zeros((512, 64), dtype=torch.bfloat16) + v = torch.zeros_like(k) + args.output = torch.empty_like(q) + monkeypatch.setattr(fmha, "_common_unsupported_reason", Mock(return_value=None)) + supported, reason = fmha._is_supported_with_reason( + q, + k, + v, + metadata, + args, + phase=FmhaPhase.CONTEXT, + ) + + assert not supported + assert reason == "block-sparse requires head_dim=128" + + +def test_contiguous_wrappers_cache_static_profile_and_keep_routes_live(monkeypatch) -> None: + _attention, fmha, q, k, v, _metadata, args = _contiguous_case() + wrapper = Mock() + factory = Mock(return_value=wrapper) + monkeypatch.setattr(block_sparse_fmha, "_BlockSparseTSWrapper", factory) + + bsr_inputs = [ + _bsr_inputs(), + BlockSparseForwardInputs( + q_block_size=64, + kv_block_size=64, + max_blocks_per_row=2, + block_indptr=torch.tensor([[[0, 1]], [[1, 4]]], dtype=torch.int32), + block_indices=torch.tensor([3, 1, 0, 2], dtype=torch.int32), + ), + ] + for inputs in bsr_inputs: + _set_block_sparse_inputs(args, inputs) + fmha._forward_contiguous(q, k, v, args) + + proxy_inputs = [_bitmask_inputs(proxy=True), _bitmask_inputs(proxy=True)] + for inputs in proxy_inputs: + _set_block_sparse_inputs(args, inputs) + fmha._forward_contiguous(q, k, v, args) + + assert factory.call_count == 2 + assert wrapper.plan.call_count == 2 + assert wrapper.plan.call_args_list[0].kwargs["sparse_format"] == "bsr" + assert wrapper.plan.call_args_list[0].kwargs["use_proxy_routes"] is False + assert wrapper.plan.call_args_list[1].kwargs["sparse_format"] == "bitmask" + assert wrapper.plan.call_args_list[1].kwargs["use_proxy_routes"] is True + assert wrapper.plan.call_args_list[1].kwargs["max_blocks_per_row"] == 4 + assert wrapper.run.call_count == 4 + + for call, inputs in zip(wrapper.run.call_args_list[:2], bsr_inputs): + assert call.kwargs["block_indptr"] is inputs.block_indptr + assert call.kwargs["block_indices"] is inputs.block_indices + for call, inputs in zip(wrapper.run.call_args_list[2:], proxy_inputs): + assert call.kwargs["exact_block_bits"] is inputs.exact_block_bits + assert call.kwargs["k_summary"] is inputs.k_summary + assert call.kwargs["v_summary"] is inputs.v_summary + + +def test_block_sparse_plan_key_includes_attention_head_topology() -> None: + inputs = _bitmask_inputs(proxy=True) + q = torch.empty((128, 256), dtype=torch.bfloat16) + first_attention = _Attention() + second_attention = _Attention() + second_attention.num_heads = 4 + first = block_sparse_fmha.PrimsTSBlockSparseFmha(first_attention) + second = block_sparse_fmha.PrimsTSBlockSparseFmha(second_attention) + + def _key(fmha): + return fmha._make_plan_key( + q, + inputs, + batch_size=1, + seq_len_q=128, + kv_capacity=256, + page_size=None, + mask_type="dense", + ) + + assert _key(first) != _key(second) + + +def test_block_sparse_plan_cache_is_shared_only_when_explicitly_bound() -> None: + first = block_sparse_fmha.PrimsTSBlockSparseFmha(_Attention()) + second = block_sparse_fmha.PrimsTSBlockSparseFmha(_Attention()) + + assert first._contiguous_wrappers is not second._contiguous_wrappers + assert first._paged_wrappers is not second._paged_wrappers + + cache_state = {} + first.bind_plan_cache(cache_state) + second.bind_plan_cache(cache_state) + + assert first._contiguous_wrappers is second._contiguous_wrappers + assert first._paged_wrappers is second._paged_wrappers + assert cache_state == { + "contiguous_wrappers": {}, + "paged_wrappers": {}, + } + + +def test_paged_wrapper_uses_zero_copy_padded_row_stride_block_tables(monkeypatch) -> None: + attention, fmha, q, metadata, args = _paged_case() + wrapper = Mock() + monkeypatch.setattr(block_sparse_fmha, "_BlockSparsePagedTSWrapper", Mock(return_value=wrapper)) + monkeypatch.setattr(block_sparse_fmha, "get_kv_page_offset", Mock(return_value=8)) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", Mock(return_value=False)) + q_processed = torch.zeros((2, 2, 128), dtype=torch.bfloat16) + kv_pool = torch.empty((16, 1, 64, 128), dtype=torch.bfloat16) + block_tables = metadata.kv_cache_block_offsets[0] + empty = torch.empty(0, dtype=torch.uint8) + preprocessed = (q_processed, kv_pool, block_tables, None, 1.0, 1.0) + ( + empty, + None, + 1, + 256, + -1, + False, + ) + monkeypatch.setattr(fmha, "_run_generation_preprocess", Mock(return_value=preprocessed)) + params = FmhaParams( + attn=attention, + meta=metadata, + fwd=args, + workspace=torch.empty(0, dtype=torch.uint8), + qkv_input=q, + context_buf=args.output, + sequence_lengths=torch.tensor([129, 193], dtype=torch.int32), + input_seq_length=1, + tokens_per_block=64, + num_requests=2, + ) + expected_block_tables = block_tables[:2, 0, :] + snapshots = [] + + def snapshot(*_args, **kwargs): + snapshots.append( + ( + kwargs["seq_lens_kv"].clone(), + kwargs["block_tables"], + kwargs["block_tables"].clone(), + kwargs["block_indptr"], + kwargs["block_indices"], + ) + ) + + wrapper.run.side_effect = snapshot + first_inputs = _get_block_sparse_inputs(args) + fmha.run_generation(params) + block_tables[:, 0].add_(10) + params.sequence_lengths = torch.tensor([130, 194], dtype=torch.int32) + _set_block_sparse_inputs( + args, + BlockSparseForwardInputs( + q_block_size=64, + kv_block_size=64, + max_blocks_per_row=2, + block_indptr=torch.tensor([[[0, 1]], [[1, 4]]], dtype=torch.int32), + block_indices=torch.tensor([3, 2, 1, 0], dtype=torch.int32), + ), + ) + fmha.run_generation(params) + + wrapper.plan.assert_called_once() + assert wrapper.run.call_count == 2 + torch.testing.assert_close(snapshots[0][0], torch.tensor([129, 193], dtype=torch.int32)) + torch.testing.assert_close(snapshots[1][0], torch.tensor([130, 194], dtype=torch.int32)) + assert snapshots[0][1].data_ptr() == expected_block_tables.data_ptr() + assert snapshots[1][1].data_ptr() == expected_block_tables.data_ptr() + assert snapshots[0][1].shape == (2, 4) + assert snapshots[0][1].stride() == (8, 1) + torch.testing.assert_close(snapshots[0][2], torch.arange(8, dtype=torch.int32).view(2, 4)) + torch.testing.assert_close(snapshots[1][2], torch.arange(8, dtype=torch.int32).view(2, 4) + 10) + assert snapshots[0][3] is first_inputs.block_indptr + assert snapshots[1][3] is _get_block_sparse_inputs(args).block_indptr + + +def test_paged_block_tables_remain_live_across_graph_replay(monkeypatch) -> None: + attention, fmha, q, metadata, args = _paged_case() + wrapper = Mock() + monkeypatch.setattr(block_sparse_fmha, "_BlockSparsePagedTSWrapper", Mock(return_value=wrapper)) + monkeypatch.setattr(block_sparse_fmha, "get_kv_page_offset", Mock(return_value=8)) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", Mock(return_value=True)) + q_processed = torch.zeros((2, 2, 128), dtype=torch.bfloat16) + kv_pool = torch.empty((16, 1, 64, 128), dtype=torch.bfloat16) + block_tables = metadata.kv_cache_block_offsets[0] + empty = torch.empty(0, dtype=torch.uint8) + preprocessed = (q_processed, kv_pool, block_tables, None, 1.0, 1.0) + ( + empty, + None, + 1, + 256, + -1, + False, + ) + monkeypatch.setattr(fmha, "_run_generation_preprocess", Mock(return_value=preprocessed)) + params = FmhaParams( + attn=attention, + meta=metadata, + fwd=args, + workspace=torch.empty(0, dtype=torch.uint8), + qkv_input=q, + context_buf=args.output, + sequence_lengths=torch.tensor([129, 193], dtype=torch.int32), + input_seq_length=1, + tokens_per_block=64, + num_requests=2, + ) + seen = [] + + def snapshot(*_args, **kwargs): + seen.append((kwargs["block_tables"].data_ptr(), kwargs["block_tables"].clone())) + + wrapper.run.side_effect = snapshot + fmha.run_generation(params) + block_tables[:, 0, :].add_(10) + block_tables[:, 1, :].fill_(-1) + fmha.run_generation(params) + + assert seen[0][0] == seen[1][0] == block_tables.data_ptr() + torch.testing.assert_close(seen[0][1], torch.arange(8, dtype=torch.int32).view(2, 4)) + torch.testing.assert_close(seen[1][1], torch.arange(8, dtype=torch.int32).view(2, 4) + 10) + + +def test_prepare_workspace_checks_capture_before_resize(monkeypatch) -> None: + _attention, fmha, _q, _metadata, _args = _paged_case() + query_device = torch.device("cuda:1") + q = SimpleNamespace( + device=query_device, + dtype=torch.bfloat16, + shape=(2, 512), + ) + metadata = SimpleNamespace( + kv_cache_manager=object(), + kv_cache_block_offsets=SimpleNamespace(device=query_device, shape=(1, 2, 4)), + max_num_requests=2, + tokens_per_block=64, + num_generations=2, + ) + workspace = torch.empty(0, dtype=torch.uint8) + monkeypatch.setattr( + fmha, + "_get_generation_workspace_layout", + Mock(return_value={"total_size": 16}), + ) + fmha._multi_processor_count = 1 + device_scope = Mock(return_value=nullcontext()) + monkeypatch.setattr(torch.cuda, "device", device_scope) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", Mock(return_value=True)) + + with pytest.raises(RuntimeError, match="workspace must be sized"): + fmha.prepare_workspace(q, None, None, metadata, _args, workspace) + + device_scope.assert_called_once_with(query_device) + assert workspace.numel() == 0 + + +def test_prepare_workspace_skips_generation_layout_for_contiguous_requests(monkeypatch) -> None: + _attention, fmha, q, _k, _v, metadata, args = _contiguous_case() + layout = Mock() + monkeypatch.setattr(fmha, "_get_generation_workspace_layout", layout) + monkeypatch.setattr(torch.cuda, "device", Mock(return_value=nullcontext())) + fmha._multi_processor_count = 1 + workspace = torch.empty(0, dtype=torch.uint8) + + fmha.prepare_workspace(q, None, None, metadata, args, workspace) + + layout.assert_not_called() + assert workspace.numel() == 0 + + +@_REQUIRES_PRIMTS_GPU +@torch.no_grad() +def test_real_gpu_raw_routes_and_token_mask_match_reference() -> None: + torch.manual_seed(1234) + q = torch.randn((1, 128, 1, 128), device="cuda", dtype=torch.float16) + k = torch.randn((1, 256, 1, 128), device="cuda", dtype=torch.float16) + v = torch.randn_like(k) + token_mask = torch.ones(256, device="cuda", dtype=torch.bool) + token_mask[[1, 63, 64, 95, 129, 190, 255]] = False + inputs = BlockSparseForwardInputs( + q_block_size=64, + kv_block_size=64, + max_blocks_per_row=3, + block_indptr=torch.tensor([[[0, 2, 5]]], device="cuda", dtype=torch.int32), + block_indices=torch.tensor([0, 2, 0, 1, 3], device="cuda", dtype=torch.int32), + kv_valid_bits=_pack_token_mask(token_mask), + ) + sm_scale = 128**-0.5 + + key_blocks = torch.arange(256, device="cuda") // 64 + allowed = torch.zeros((128, 256), device="cuda", dtype=torch.bool) + for row, selected_blocks in enumerate(((0, 2), (0, 1, 3))): + selected = torch.tensor(selected_blocks, device="cuda") + allowed[row * 64 : (row + 1) * 64] = torch.isin(key_blocks, selected) & token_mask + scores = (q[0, :, 0].float() @ k[0, :, 0].float().T) * sm_scale + expected = ( + torch.softmax(scores.masked_fill(~allowed, float("-inf")), dim=-1) @ v[0, :, 0].float() + ).to(q.dtype)[None, :, None, :] + + actual = prims_ts.block_sparse_attention( + q, + k, + v, + block_indptr=inputs.block_indptr, + block_indices=inputs.block_indices, + q_block_size=inputs.q_block_size, + kv_block_size=inputs.kv_block_size, + kv_valid_bits=inputs.kv_valid_bits, + sm_scale=sm_scale, + ) + torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2) + + +@_REQUIRES_PRIMTS_GPU +@torch.no_grad() +def test_real_gpu_proxy_adapter_replays_live_routes_and_summaries() -> None: + torch.manual_seed(20260901) + q = torch.randn((1, 64, 1, 128), device="cuda", dtype=torch.bfloat16) + k = torch.randn((1, 192, 1, 128), device="cuda", dtype=torch.bfloat16) + v = torch.randn_like(k) + k_blocks = k.float().view(1, 3, 64, 1, 128) + v_blocks = v.float().view(1, 3, 64, 1, 128) + initial_k_summary = k_blocks.mean(dim=2).to(k.dtype) + initial_v_summary = v_blocks.sum(dim=2).to(v.dtype) + live_k_summary = initial_k_summary.clone() + live_v_summary = initial_v_summary.clone() + live_exact_bits = torch.tensor([[[[1]]]], device="cuda", dtype=torch.uint32) + + attention = _Attention() + attention.num_heads = attention.num_kv_heads = 1 + fmha = block_sparse_fmha.PrimsTSBlockSparseFmha(attention) + output = torch.empty_like(q).view(64, 128) + args = AttentionForwardArgs( + output=output, + attention_input_type=AttentionInputType.context_only, + attention_mask=PredefinedAttentionMask.FULL, + sparse_runtime_params=SparseRuntimeParams( + block_sparse_inputs=BlockSparseForwardInputs( + q_block_size=64, + kv_block_size=64, + exact_block_bits=live_exact_bits, + k_summary=live_k_summary, + v_summary=live_v_summary, + ), + ), + ) + metadata = SimpleNamespace( + is_cross=False, + kv_cache_manager=None, + seq_lens=torch.tensor([64], dtype=torch.int32), + ) + flat_q, flat_k, flat_v = (tensor.flatten(0, 2) for tensor in (q, k, v)) + + fmha.forward(flat_q, flat_k, flat_v, metadata, args) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + fmha.forward(flat_q, flat_k, flat_v, metadata, args) + + graph.replay() + torch.cuda.synchronize() + expected = _proxy_reference(q, k, v, live_k_summary, live_v_summary, exact_block=0) + torch.testing.assert_close(output.view_as(q), expected, rtol=2e-2, atol=2e-2) + + live_exact_bits.fill_(1 << 2) + live_k_summary.copy_((initial_k_summary.float() * 0.5 + 0.125).to(k.dtype)) + live_v_summary.copy_((initial_v_summary.float() * -0.25).to(v.dtype)) + graph.replay() + torch.cuda.synchronize() + expected = _proxy_reference(q, k, v, live_k_summary, live_v_summary, exact_block=2) + torch.testing.assert_close(output.view_as(q), expected, rtol=2e-2, atol=2e-2) + + +@_REQUIRES_PRIMTS_GPU +@torch.no_grad() +def test_real_gpu_paged_routes_use_live_length_below_capacity() -> None: + torch.manual_seed(7) + q = torch.randn((1, 64, 1, 128), device="cuda", dtype=torch.float16) + k_cache = torch.randn((4, 1, 64, 128), device="cuda", dtype=torch.float16) + v_cache = torch.randn_like(k_cache) + page_indices = torch.tensor([2, 0, 3, 1], device="cuda", dtype=torch.int32) + seq_lens_kv = torch.tensor([160], device="cuda", dtype=torch.int32) + inputs = BlockSparseForwardInputs( + q_block_size=64, + kv_block_size=64, + max_blocks_per_row=2, + block_indptr=torch.tensor([[[0, 2]]], device="cuda", dtype=torch.int32), + block_indices=torch.tensor([0, 2], device="cuda", dtype=torch.int32), + ) + sm_scale = 128**-0.5 + + actual = prims_ts.block_sparse_attention_with_paged_kv_cache( + q, + (k_cache, v_cache), + block_tables=page_indices.view(1, 4), + seq_lens_kv=seq_lens_kv, + block_indptr=inputs.block_indptr, + block_indices=inputs.block_indices, + max_seq_len_kv=256, + q_block_size=inputs.q_block_size, + kv_block_size=inputs.kv_block_size, + sm_scale=sm_scale, + ) + + logical_k = k_cache.index_select(0, page_indices.long()).reshape(256, 1, 128) + logical_v = v_cache.index_select(0, page_indices.long()).reshape(256, 1, 128) + allowed = torch.zeros(256, device="cuda", dtype=torch.bool) + allowed[:64] = True + allowed[128:160] = True + scores = (q[0, :, 0].float() @ logical_k[:, 0].float().T) * sm_scale + expected = ( + torch.softmax(scores.masked_fill(~allowed, float("-inf")), dim=-1) @ logical_v[:, 0].float() + ).to(q.dtype)[None, :, None, :] + + torch.testing.assert_close(actual, expected, rtol=1e-2, atol=1e-2) diff --git a/tests/unittest/_torch/attention/test_prims_ts_fmha.py b/tests/unittest/_torch/attention/test_prims_ts_fmha.py index 87dbd7364f4b..a18f41c00bf7 100644 --- a/tests/unittest/_torch/attention/test_prims_ts_fmha.py +++ b/tests/unittest/_torch/attention/test_prims_ts_fmha.py @@ -38,6 +38,7 @@ AttentionInputType, PredefinedAttentionMask, ) +from tensorrt_llm._torch.attention.backends.sparse.params import SparseRuntimeParams from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm.bindings import DataType @@ -185,7 +186,7 @@ def _support_result( is_fused_qkv=is_fused_qkv, ) if has_sparse_runtime_metadata: - forward_args.sparse_runtime_params.sparse_kv_indices = torch.empty(1) + forward_args.sparse_runtime_params = SparseRuntimeParams(sparse_kv_indices=torch.empty(1)) if attention_input_type == AttentionInputType.context_only: num_contexts, num_generations, num_ctx_tokens = 1, 0, 4 kv_lens = [4] diff --git a/tests/unittest/_torch/attention/test_skip_softmax_sm120.py b/tests/unittest/_torch/attention/test_skip_softmax_sm120.py index 979b0d3c6908..af8a067f7129 100644 --- a/tests/unittest/_torch/attention/test_skip_softmax_sm120.py +++ b/tests/unittest/_torch/attention/test_skip_softmax_sm120.py @@ -75,13 +75,16 @@ def _run_context( ) -> tuple: """Build a TRTLLM attention layer + no-cache context metadata and run a packed-QKV causal prefill. Mirrors ``test_attention_no_cache``.""" - AttentionCls = get_attention_backend("TRTLLM") + sparse_params = ( + sparse_attention_config.to_sparse_params() if sparse_attention_config is not None else None + ) + AttentionCls = get_attention_backend("TRTLLM", sparse_params=sparse_params) layer = AttentionCls( layer_idx=0, num_heads=num_heads, head_dim=head_dim, num_kv_heads=num_kv_heads, - sparse_attention_config=sparse_attention_config, + sparse_params=sparse_params, ) metadata = AttentionCls.Metadata( diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_attention.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_attention.py index 2b900c603ce6..48f46ddee190 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_attention.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_ulysses_attention.py @@ -1,3 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + """Multi-GPU tests for Ulysses Attention. These tests use torch.multiprocessing.spawn to launch multiple processes internally. @@ -105,6 +120,56 @@ def run_test_in_distributed(world_size: int, test_fn: Callable, use_cuda: bool = ) +def test_forward_async_redistributes_vsa_gates(monkeypatch): + if not MODULES_AVAILABLE: + pytest.skip("Required modules not available") + + import tensorrt_llm._torch.visual_gen.attention_backend.parallel as parallel_backend + + class _CaptureBackend: + preferred_layout = AttentionTensorLayout.NHD + + def forward(self, q, k, v, **kwargs): + self.kwargs = kwargs + return q + + inner_backend = _CaptureBackend() + attention = object.__new__(UlyssesAttention) + attention.world_size = 1 + attention.process_group = None + attention.inner_backend = inner_backend + attention._issue_async = lambda tensor: tensor.unsqueeze(0) + attention._join_async = lambda: None + attention._output_a2a = lambda output, batch_size, seq_len: output + redistributed = [] + + def _fake_all_to_all(tensor, **kwargs): + redistributed.append((tensor, kwargs)) + return tensor + 1 + + monkeypatch.setattr(parallel_backend, "all_to_all_4d", _fake_all_to_all) + q = torch.randn(1, 3, 2, 4) + gate_compress = torch.randn_like(q) + gate_fine = torch.randn_like(q) + + output = attention.forward_async( + lambda: q, + lambda: q, + lambda: q, + gate_compress=gate_compress, + gate_fine=gate_fine, + ) + + assert output.shape == q.shape + assert redistributed[0][0] is gate_compress + assert redistributed[1][0] is gate_fine + assert all(entry[1]["scatter_dim"] == 2 for entry in redistributed) + assert all(entry[1]["gather_dim"] == 1 for entry in redistributed) + assert inner_backend.kwargs["batch_size"] == q.shape[0] + torch.testing.assert_close(inner_backend.kwargs["gate_compress"], gate_compress + 1) + torch.testing.assert_close(inner_backend.kwargs["gate_fine"], gate_fine + 1) + + # ============================================================================= # Test logic functions (module-level so they can be pickled by mp.spawn) # ============================================================================= diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_async_ulysses.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_async_ulysses.py index 042525f82e30..64580eab2f00 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_async_ulysses.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_async_ulysses.py @@ -324,5 +324,56 @@ def test_async_vs_sync_parity(self, backend): run_test_in_distributed(2, _logic_async_vs_sync_parity, backend) +def test_forward_async_uses_tp_local_heads_for_qkv_gates_and_output(): + from tensorrt_llm._torch.visual_gen.modules.attention import Attention + + class _CaptureAsyncAttention(torch.nn.Module): + def forward_async(self, compute_q, compute_k, compute_v, **kwargs): + self.q = compute_q() + self.k = compute_k() + self.v = compute_v() + self.kwargs = kwargs + return self.q + + class _CaptureOutputProjection(torch.nn.Module): + def forward(self, hidden_states): + self.input = hidden_states + return hidden_states + + attention = Attention.__new__(Attention) + torch.nn.Module.__init__(attention) + attention.num_attention_heads = 4 + attention.num_key_value_heads = 4 + attention.local_num_attention_heads = 2 + attention.local_num_key_value_heads = 2 + attention.head_dim = 4 + attention.fuse_qk_norm_rope = False + attention.qk_norm = False + attention._maybe_share_qkv_quantize = False + attention.to_q = torch.nn.Linear(12, 8, bias=False) + attention.to_k = torch.nn.Linear(12, 8, bias=False) + attention.to_v = torch.nn.Linear(12, 8, bias=False) + attention.attn = _CaptureAsyncAttention() + output_projection = _CaptureOutputProjection() + attention.to_out = torch.nn.ModuleList([output_projection]) + gate_compress = torch.randn(1, 3, 8) + gate_fine = torch.randn_like(gate_compress) + + output = attention.forward_async( + torch.randn(1, 3, 12), + gate_compress=gate_compress, + gate_fine=gate_fine, + ) + + expected_shape = (1, 3, 2, 4) + assert attention.attn.q.shape == expected_shape + assert attention.attn.k.shape == expected_shape + assert attention.attn.v.shape == expected_shape + assert attention.attn.kwargs["gate_compress"].shape == expected_shape + assert attention.attn.kwargs["gate_fine"].shape == expected_shape + assert output_projection.input.shape == (1, 3, 8) + assert output.shape == (1, 3, 8) + + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_tp.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_tp.py index fbedc1aa1410..032508230298 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_tp.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_tp.py @@ -529,6 +529,53 @@ def _logic_wan_i2v_tp_vs_single_gpu_with_config(rank, world_size, config_dict): # ============================================================================= +@pytest.mark.parametrize( + ("tp_rank", "expected_head_range"), + [(0, (0, 8)), (1, (8, 12))], +) +def test_wan_vsa_gates_follow_ulysses_aligned_tp_q_shard(monkeypatch, tp_rank, expected_head_range): + """VSA gates must select the same TP-local heads as the Q projection.""" + from tensorrt_llm._torch.visual_gen.models.wan.transformer_wan import WanBlock + from tensorrt_llm._torch.visual_gen.modules import attention as attention_module + from tensorrt_llm.mapping import Mapping + from tensorrt_llm.visual_gen.args import VideoSparseAttentionConfig + + monkeypatch.setattr( + attention_module, + "wrap_parallel_attention", + lambda attention, **_kwargs: attention, + ) + mapping = Mapping(world_size=2, rank=tp_rank, tp_size=2) + monkeypatch.setattr(type(mapping), "tp_rank", property(lambda self: self.rank)) + head_dim = 128 + model_config = DiffusionModelConfig( + pretrained_config=SimpleNamespace( + hidden_size=12 * head_dim, + num_attention_heads=12, + attention_head_dim=head_dim, + ffn_dim=512, + eps=1e-6, + cross_attn_norm=True, + ), + mapping=mapping, + visual_gen_mapping=SimpleNamespace(ulysses_size=4, cp_size=1), + attention=AttentionConfig( + backend="CUTEDSL", + sparse_attention_config=VideoSparseAttentionConfig(vsa_sparsity=0.5), + ), + skip_create_weights_in_init=True, + ) + + block = WanBlock(model_config, _layer_idx=0) + + expected_shard = tuple(index * head_dim for index in expected_head_range) + assert (block.attn1.local_q_dim_start, block.attn1.local_q_dim_end) == expected_shard + assert block.to_gate_compress.tp_sharding == expected_shard + assert block.to_gate_fine.tp_sharding == expected_shard + assert block.to_gate_compress.out_features == block.attn1.local_q_dim + assert block.to_gate_fine.out_features == block.attn1.local_q_dim + + class TestWanT2VTP: """Tensor parallelism tests for WAN T2V transformer.""" diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_vsa_ulysses.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_vsa_ulysses.py index 00dd752d902e..a60d12a01a30 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_vsa_ulysses.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_wan_vsa_ulysses.py @@ -36,9 +36,9 @@ import torch.multiprocessing as mp try: - from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl import ( + from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl import _cute_dsl_import_error + from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import ( VSAMetadataBuilder, - _cute_dsl_import_error, set_vsa_forward_context, ) from tensorrt_llm._torch.visual_gen.config import ( diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/test_sol_attention.py b/tests/unittest/_torch/visual_gen/sparse_attention/test_sol_attention.py new file mode 100644 index 000000000000..0a4a1f1a521e --- /dev/null +++ b/tests/unittest/_torch/visual_gen/sparse_attention/test_sol_attention.py @@ -0,0 +1,898 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the VisualGen SOL TRTLLM backend.""" + +from __future__ import annotations + +import math +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch +from pydantic import ValidationError + +from tensorrt_llm._torch.attention.backends.fmha.prims_ts_block_sparse import PrimsTSBlockSparseFmha +from tensorrt_llm._torch.attention.backends.interface import ( + AttentionForwardArgs, + PredefinedAttentionMask, +) +from tensorrt_llm._torch.attention.backends.sparse.hooks import prepare_sparse_runtime_params +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttention as CoreTrtllmAttention +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol import backend as sol_backend +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.backend import SOLTrtllmAttention +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.params import SolParams +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.predictor import ( + SolPredictorOutputs, + SOLSparsePredictor, +) +from tensorrt_llm._torch.visual_gen.attention_backend.trtllm import TrtllmAttention +from tensorrt_llm._torch.visual_gen.attention_backend.utils import create_attention +from tensorrt_llm._torch.visual_gen.config import ( + DiffusionModelConfig, + create_attention_metadata_state, +) +from tensorrt_llm._torch.visual_gen.cuda_graph_runner import CUDAGraphRunner, CUDAGraphRunnerConfig +from tensorrt_llm._torch.visual_gen.models.modeling import BaseDiffusionModel +from tensorrt_llm._torch.visual_gen.modules import attention as attention_module +from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode +from tensorrt_llm.visual_gen import SolAttentionConfig +from tensorrt_llm.visual_gen.args import AttentionConfig, QuantAttentionConfig + +_REQUIRES_SM100 = pytest.mark.skipif( + not torch.cuda.is_available() or torch.cuda.get_device_capability() not in ((10, 0), (10, 3)), + reason="SOL requires SM100 or SM103", +) +_CPU_ONLY = pytest.mark.cpu_only + + +def _make_backend( + params: SolParams, + predictor: Mock, + *, + layer_idx: int = 1, +) -> SOLTrtllmAttention: + backend = object.__new__(SOLTrtllmAttention) + backend.layer_idx = layer_idx + backend.num_heads = 2 + backend.num_kv_heads = 2 + backend.head_dim = 128 + backend.q_scaling = 1.0 + backend.quant_attention_config = None + backend.sparse_params = None + backend._fmha_manager = SimpleNamespace(fmha_libs=[object.__new__(PrimsTSBlockSparseFmha)]) + backend.sol_params = params + backend._prepared_graph_phase = None + backend.predictor = predictor + return backend + + +def _flatten(tensor: torch.Tensor | None) -> torch.Tensor | None: + """Convert a BSHD tensor into the flattened ``[B*S, H*D]`` core layout.""" + + if tensor is None: + return None + return tensor.reshape(tensor.shape[0] * tensor.shape[1], -1) + + +def _core_metadata(batch_size: int, seq_len: int) -> SimpleNamespace: + return SimpleNamespace(num_seqs=batch_size, max_seq_len=seq_len) + + +def _predict( + backend: SOLTrtllmAttention, + q: torch.Tensor, + k: torch.Tensor | None, + v: torch.Tensor | None, + *, + attention_mask: PredefinedAttentionMask = PredefinedAttentionMask.FULL, + timestep: object = None, +): + """Invoke the core prediction hook the way the core forward does.""" + + return backend.block_sparse_attn_predict( + _flatten(q), + _flatten(k), + _flatten(v), + _core_metadata(q.shape[0], q.shape[1]), + AttentionForwardArgs(attention_mask=attention_mask, timestep=timestep), + ) + + +def _forward( + backend: SOLTrtllmAttention, + q: torch.Tensor, + k: torch.Tensor | None, + v: torch.Tensor | None, + **kwargs, +) -> torch.Tensor: + seq_len_kv = kwargs.pop("seq_len_kv", q.shape[1]) + return backend.forward( + q=q, + k=k, + v=v, + batch_size=q.shape[0], + seq_len=q.shape[1], + seq_len_kv=seq_len_kv, + **kwargs, + ) + + +def _stub_core_forward(monkeypatch) -> dict: + """Replace metadata preparation and the core forward with a recorder that + still runs the backend's sparse prediction.""" + + captured = {} + monkeypatch.setattr( + TrtllmAttention, + "_prepare_metadata", + lambda self, batch_size, seq_len: _core_metadata(batch_size, seq_len), + ) + + def _core_forward(self, q, k, v, metadata, forward_args=None, **kwargs): + forward_args.sparse_runtime_params = prepare_sparse_runtime_params( + self, q, k, v, metadata, forward_args + ) + captured.update(q=q, k=k, v=v, metadata=metadata, forward_args=forward_args) + return q + + monkeypatch.setattr(CoreTrtllmAttention, "forward", _core_forward) + return captured + + +def _predictor_outputs(*, batch_size: int, seq_len: int, num_heads: int) -> SolPredictorOutputs: + num_blocks = (seq_len + 63) // 64 + return SolPredictorOutputs( + exact_block_bits=torch.zeros( + batch_size, + num_heads, + num_blocks, + (num_blocks + 31) // 32, + dtype=torch.uint32, + ), + k_summary=torch.zeros(batch_size, num_blocks, num_heads, 128, dtype=torch.bfloat16), + v_summary=torch.zeros(batch_size, num_blocks, num_heads, 128, dtype=torch.bfloat16), + ) + + +def _bshd(seq_len: int = 64, num_heads: int = 2) -> torch.Tensor: + return torch.zeros(1, seq_len, num_heads, 128, dtype=torch.bfloat16) + + +def _stub_backend( + params: SolParams | None = None, + *, + seq_len: int = 64, + unsupported_reason: str | None = None, +) -> tuple[SOLTrtllmAttention, Mock]: + predictor = Mock(spec=SOLSparsePredictor) + predictor.support_reason.return_value = unsupported_reason + predictor.predict.return_value = _predictor_outputs(batch_size=1, seq_len=seq_len, num_heads=2) + return _make_backend(params or SolParams(tau=1.0), predictor), predictor + + +def _dense_reference(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: + scores = torch.einsum("bqhd,bkhd->bhqk", q.float(), k.float()) * (128**-0.5) + return torch.einsum("bhqk,bkhd->bqhd", scores.softmax(dim=-1), v.float()).to(q.dtype) + + +def _mixed_proxy_reference( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + outputs: SolPredictorOutputs, +) -> torch.Tensor: + """Evaluate the exact-token/proxy-summary attention contract.""" + + block_size = 64 + num_blocks = math.ceil(k.shape[1] / block_size) + exact_words = outputs.exact_block_bits.detach().cpu().to(torch.int64) + reference = torch.empty_like(q) + scale = q.shape[-1] ** -0.5 + for batch_idx in range(q.shape[0]): + for head_idx in range(q.shape[2]): + for q_block_idx in range(math.ceil(q.shape[1] / block_size)): + q_begin = q_block_idx * block_size + q_end = min(q_begin + block_size, q.shape[1]) + exact_blocks = [ + block_idx + for block_idx in range(num_blocks) + if int(exact_words[batch_idx, head_idx, q_block_idx, block_idx // 32]) + & (1 << (block_idx % 32)) + ] + proxy_blocks = [ + block_idx for block_idx in range(num_blocks) if block_idx not in exact_blocks + ] + exact_tokens = torch.cat( + [ + torch.arange( + block_idx * block_size, + min((block_idx + 1) * block_size, k.shape[1]), + device=q.device, + ) + for block_idx in exact_blocks + ] + ) + q_rows = q[batch_idx, q_begin:q_end, head_idx].float() + exact_logits = (q_rows @ k[batch_idx, exact_tokens, head_idx].float().T) * scale + proxy_logits = ( + q_rows @ outputs.k_summary[batch_idx, proxy_blocks, head_idx].float().T + ) * scale + logits = torch.cat((exact_logits, proxy_logits), dim=1) + weights = torch.exp(logits - logits.amax(dim=1, keepdim=True)) + exact_weights = weights[:, : exact_tokens.numel()] + proxy_weights = weights[:, exact_tokens.numel() :] + numerator = exact_weights @ v[batch_idx, exact_tokens, head_idx].float() + if proxy_blocks: + numerator += ( + proxy_weights @ outputs.v_summary[batch_idx, proxy_blocks, head_idx].float() + ) + denominator = exact_weights.sum(dim=1, keepdim=True) + for proxy_offset, block_idx in enumerate(proxy_blocks): + tokens_in_block = min(block_size, k.shape[1] - block_idx * block_size) + denominator += proxy_weights[:, proxy_offset : proxy_offset + 1] * ( + tokens_in_block + ) + reference[batch_idx, q_begin:q_end, head_idx] = (numerator / denominator).to( + q.dtype + ) + return reference + + +@_CPU_ONLY +def test_sol_params_requires_precomputed_phase_during_cuda_graph_capture(monkeypatch) -> None: + params = SolParams(tau=1.0, disabled_until_timestep=0.6) + timestep = torch.tensor(0.2) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + + with pytest.raises(RuntimeError, match="precomputed"): + params.should_use_sparse(layer_idx=1, timestep=timestep) + + assert params.should_use_sparse(layer_idx=1, timestep=timestep, graph_phase=1) + assert not params.should_use_sparse(layer_idx=1, timestep=timestep, graph_phase=0) + + +@_CPU_ONLY +def test_sol_backend_warmup_prepares_phase_for_capture(monkeypatch) -> None: + params = SolParams(tau=1.0, disabled_until_timestep=0.6) + q = _bshd() + backend, predictor = _stub_backend(params) + + assert _predict(backend, q, q, q, timestep=torch.tensor(0.8)) is None + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + assert _predict(backend, q, q, q, timestep=torch.tensor(0.8)) is None + predictor.predict.assert_not_called() + + +@_CPU_ONLY +def test_sol_backend_rejects_cutoff_capture_without_warmup(monkeypatch) -> None: + params = SolParams(tau=1.0, disabled_until_timestep=0.6) + q = _bshd() + backend, predictor = _stub_backend(params) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + + with pytest.raises(RuntimeError, match="prepared before CUDA Graph capture"): + _predict(backend, q, q, q, timestep=torch.tensor(0.2)) + + predictor.predict.assert_not_called() + + +@_CPU_ONLY +def test_sol_backend_cutoff_requires_timestep_during_eager_forward() -> None: + params = SolParams(tau=1.0, disabled_until_timestep=0.6) + q = _bshd() + backend, predictor = _stub_backend(params) + + with pytest.raises(ValueError, match="timestep is required"): + _predict(backend, q, q, q) + + predictor.predict.assert_not_called() + + +@_CPU_ONLY +def test_sol_phase_waits_until_all_token_timesteps_are_below_cutoff() -> None: + params = SolParams(tau=1.0, disabled_until_timestep=0.6) + + assert not params.should_use_sparse(layer_idx=1, timestep=torch.tensor([0.0, 0.8])) + assert params.should_use_sparse(layer_idx=1, timestep=torch.tensor([0.0, 0.2])) + + +@_CPU_ONLY +def test_sol_config_lowers_and_factory_initializes_backend(monkeypatch) -> None: + base_kwargs = {} + + def _base_init(self, **kwargs) -> None: + base_kwargs.update(kwargs) + self.layer_idx = kwargs["layer_idx"] + self.head_dim = kwargs["head_dim"] + self.q_scaling = 1.0 + + monkeypatch.setattr(TrtllmAttention, "__init__", _base_init) + attention_config = AttentionConfig( + backend="TRTLLM", + sparse_attention_config={ + "algorithm": "sol_attn", + "tau": -0.25, + "disabled_until_timestep": 0.6, + "dense_layers": "0,2-4,3", + }, + ) + config = attention_config.sparse_attention_config + assert isinstance(config, SolAttentionConfig) + params = config.to_sparse_params() + + backend = create_attention( + backend="TRTLLM", + layer_idx=3, + num_heads=4, + head_dim=128, + attention_config=attention_config, + sparse_params=params, + attention_metadata_state=create_attention_metadata_state(), + ) + + assert config.algorithm == "sol_attn" + assert params.tau == -0.25 + assert params.disabled_until_timestep == 0.6 + assert params.dense_layers == frozenset({0, 2, 3, 4}) + assert isinstance(backend, SOLTrtllmAttention) + assert backend.sol_params is params + assert isinstance(backend.predictor, SOLSparsePredictor) + assert base_kwargs["sparse_params"] is None + assert "_enable_sparse_workflow" not in SOLTrtllmAttention.__dict__ + assert "_should_use_sparse_workflow" not in SOLTrtllmAttention.__dict__ + assert not backend.support_fused_qkv() + assert "forward" not in SOLTrtllmAttention.__dict__ + assert "block_sparse_attn_predict" in SOLTrtllmAttention.__dict__ + + +@_CPU_ONLY +def test_sol_layers_share_one_model_scoped_predictor(monkeypatch) -> None: + monkeypatch.setattr(TrtllmAttention, "__init__", lambda self, **kwargs: None) + attention_metadata_state = create_attention_metadata_state() + + first = SOLTrtllmAttention( + sparse_params=SolParams(), + attention_metadata_state=attention_metadata_state, + ) + second = SOLTrtllmAttention( + sparse_params=SolParams(), + attention_metadata_state=attention_metadata_state, + ) + + assert first.predictor is second.predictor + assert attention_metadata_state["sparse_predictors"]["sol_attn"] is first.predictor + + +@_CPU_ONLY +def test_sol_backend_sparse_phase_emits_proxy_bitmask_carrier(monkeypatch) -> None: + batch_size, seq_len, num_heads = 1, 65, 2 + q, k, v = (_bshd(seq_len, num_heads) for _ in range(3)) + predictor_outputs = _predictor_outputs( + batch_size=batch_size, + seq_len=seq_len, + num_heads=num_heads, + ) + backend, predictor = _stub_backend(SolParams(tau=0.75), seq_len=seq_len) + predictor.predict.return_value = predictor_outputs + monkeypatch.setattr(sol_backend, "get_bmm1_scale", lambda attn: 0.375) + + carrier = _predict(backend, q, k, v, timestep=0.2) + + predicted_q, predicted_k, predicted_v = predictor.predict.call_args.args + assert predictor.predict.call_args.kwargs == {"tau": 0.75, "sm_scale": 0.375} + predictor.support_reason.assert_called_once_with(predicted_q, predicted_k, predicted_v) + for predicted, source in zip((predicted_q, predicted_k, predicted_v), (q, k, v), strict=True): + assert predicted.shape == (batch_size, seq_len, num_heads, 128) + assert predicted.is_contiguous() + assert predicted.data_ptr() == source.data_ptr() + assert ( + carrier.q_block_size, + carrier.kv_block_size, + carrier.max_blocks_per_row, + carrier.block_indptr, + carrier.block_indices, + carrier.kv_valid_bits, + ) == (64, 64, None, None, None, None) + assert carrier.exact_block_bits is predictor_outputs.exact_block_bits + assert carrier.k_summary is predictor_outputs.k_summary + assert carrier.v_summary is predictor_outputs.v_summary + assert carrier.sparse_format == "bitmask" + assert carrier.use_proxy_routes + + +@_CPU_ONLY +def test_sol_wrapper_compacts_separate_qkv_and_predicts_inside_core(monkeypatch) -> None: + batch_size, seq_len, num_heads = 1, 65, 2 + packed_qkv = torch.zeros(batch_size, seq_len, 3 * num_heads * 128, dtype=torch.bfloat16) + q, k, v = ( + tensor.view(batch_size, seq_len, num_heads, 128) + for tensor in packed_qkv.split(num_heads * 128, dim=-1) + ) + predictor_outputs = _predictor_outputs( + batch_size=batch_size, + seq_len=seq_len, + num_heads=num_heads, + ) + backend, predictor = _stub_backend(SolParams(tau=0.75), seq_len=seq_len) + predictor.predict.return_value = predictor_outputs + monkeypatch.setattr(sol_backend, "get_bmm1_scale", lambda attn: 0.375) + captured = _stub_core_forward(monkeypatch) + + output = _forward(backend, q, k, v, attention_mask=PredefinedAttentionMask.FULL, timestep=0.2) + + assert output.shape == (batch_size, seq_len, num_heads * 128) + assert all( + tensor.is_contiguous() and tensor.shape == (batch_size * seq_len, num_heads * 128) + for tensor in (captured["q"], captured["k"], captured["v"]) + ) + forward_args = captured["forward_args"] + assert forward_args.timestep == 0.2 + assert forward_args.sparse_backend_args is None + carrier = forward_args.sparse_runtime_params.block_sparse_inputs + assert carrier.exact_block_bits is predictor_outputs.exact_block_bits + predicted_q = predictor.predict.call_args.args[0] + assert predicted_q.data_ptr() == captured["q"].data_ptr() + + +@_CPU_ONLY +@pytest.mark.parametrize( + ("k", "v", "attention_mask", "message"), + ( + (None, None, PredefinedAttentionMask.FULL, "separate q, k, and v"), + (_bshd(32), _bshd(32), PredefinedAttentionMask.FULL, "self-attention"), + (_bshd(), _bshd(), PredefinedAttentionMask.CAUSAL, "full attention mask"), + ), +) +def test_sol_backend_rejects_non_sol_sparse_calls( + k: torch.Tensor | None, + v: torch.Tensor | None, + attention_mask: PredefinedAttentionMask, + message: str, +) -> None: + q = _bshd() + backend, predictor = _stub_backend() + + with pytest.raises(ValueError, match=message): + _predict(backend, q, k, v, attention_mask=attention_mask) + + predictor.predict.assert_not_called() + + +@_CPU_ONLY +def test_sol_wrapper_rejects_fused_qkv_before_core(monkeypatch) -> None: + q = _bshd() + backend, predictor = _stub_backend() + prepare_metadata = Mock(return_value=object()) + monkeypatch.setattr(TrtllmAttention, "_prepare_metadata", prepare_metadata) + + with pytest.raises(ValueError, match="separate q, k, and v"): + _forward(backend, q, None, None) + + prepare_metadata.assert_not_called() + predictor.predict.assert_not_called() + + +@_CPU_ONLY +def test_sol_backend_surfaces_predictor_support_reason_before_execution() -> None: + q = _bshd() + reason = "SOL predictor requires compact BSHD q/k/v" + backend, predictor = _stub_backend(unsupported_reason=reason) + + with pytest.raises(ValueError, match=reason): + _predict(backend, q, q, q) + + predictor.predict.assert_not_called() + + +@_CPU_ONLY +@pytest.mark.parametrize( + ("params", "layer_idx", "timestep"), + ( + (SolParams(dense_layers=frozenset({1})), 1, None), + (SolParams(disabled_until_timestep=0.6), 1, 0.8), + ), +) +def test_sol_dense_policy_returns_no_routes_without_predicting( + params: SolParams, + layer_idx: int, + timestep: float | None, +) -> None: + q = _bshd() + backend, predictor = _stub_backend(params) + backend.layer_idx = layer_idx + + assert _predict(backend, q, q, q, timestep=timestep) is None + predictor.support_reason.assert_not_called() + predictor.predict.assert_not_called() + + +@_CPU_ONLY +def test_sol_sparse_phase_without_primts_fails_closed() -> None: + q = _bshd() + backend, predictor = _stub_backend() + backend._fmha_manager = SimpleNamespace(fmha_libs=[]) + + with pytest.raises(RuntimeError, match="requires PrimTS block-sparse FMHA"): + _predict(backend, q, q, q) + + predictor.support_reason.assert_not_called() + predictor.predict.assert_not_called() + + +@_CPU_ONLY +def test_sol_sparse_phase_with_quantization_fails_closed() -> None: + q = _bshd() + backend, predictor = _stub_backend() + backend.quant_attention_config = object() + + with pytest.raises(ValueError, match="does not support quant_attention_config"): + _predict(backend, q, q, q) + + predictor.support_reason.assert_not_called() + predictor.predict.assert_not_called() + + +@_CPU_ONLY +@pytest.mark.parametrize( + "config_kwargs", + ( + {"tau": 1.0e100}, + {"disabled_until_timestep": 0.0}, + {"dense_layers": "2-1"}, + ), +) +def test_sol_public_config_rejects_invalid_policy(config_kwargs) -> None: + with pytest.raises((ValidationError, ValueError)): + SolAttentionConfig(**config_kwargs).to_sparse_params() + + +@_CPU_ONLY +def test_sol_public_config_requires_trtllm_backend() -> None: + with pytest.raises(ValidationError, match="requires backend"): + AttentionConfig( + backend="VANILLA", + sparse_attention_config=SolAttentionConfig(), + ) + + +@_CPU_ONLY +def test_sol_and_attention_quantization_are_mutually_exclusive() -> None: + with pytest.raises(ValidationError, match="SOL and quant_attention_config"): + AttentionConfig( + backend="TRTLLM", + quant_attention_config=QuantAttentionConfig( + qk_dtype="fp8", + q_block_size=1, + k_block_size=1, + v_block_size=1, + ), + sparse_attention_config=SolAttentionConfig(), + ) + + +def _sol_model_config(*, cp_size: int = 1) -> DiffusionModelConfig: + config = DiffusionModelConfig( + pretrained_config=SimpleNamespace(), + attention=AttentionConfig( + backend="TRTLLM", + sparse_attention_config=SolAttentionConfig( + tau=0.75, + disabled_until_timestep=0.6, + dense_layers="0,2-3", + ), + ), + skip_create_weights_in_init=True, + attention_metadata_state=create_attention_metadata_state(), + ) + if cp_size > 1: + config.visual_gen_mapping = SimpleNamespace( + ring_size=cp_size, + ring_group=None, + ulysses_size=1, + ulysses_group=None, + attn2d_row_size=1, + attn2d_col_size=1, + attn2d_row_group=None, + attn2d_col_group=None, + cp_size=cp_size, + ) + return config + + +class _SolModel(BaseDiffusionModel): + def __init__(self, backends: tuple[SOLTrtllmAttention, ...]) -> None: + super().__init__(_sol_model_config()) + self.backends = backends + + def forward(self, q: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor: + for backend in self.backends: + q = backend.forward( + q=q, + k=q, + v=q, + batch_size=q.shape[0], + seq_len=q.shape[1], + seq_len_kv=q.shape[1], + timestep=timestep, + ) + return q + + +@_CPU_ONLY +@pytest.mark.parametrize( + ("is_self_attention", "expected_backend", "expects_sol_params"), + ((True, "TRTLLM", True), (False, "VANILLA", False)), + ids=("self", "cross"), +) +def test_sol_attention_module_dispatches_by_attention_role( + monkeypatch, + is_self_attention: bool, + expected_backend: str, + expects_sol_params: bool, +) -> None: + captured = {} + + def _create_attention(*, backend, **kwargs): + captured.update(backend=backend, **kwargs) + return SimpleNamespace(preferred_layout=None) + + monkeypatch.setattr(attention_module, "create_attention", _create_attention) + + attention = Attention( + hidden_size=256, + num_attention_heads=2, + head_dim=128, + qkv_mode=QKVMode.SEPARATE_QKV, + qk_norm=False, + config=_sol_model_config(), + separate_qkv_is_self_attention=is_self_attention, + ) + + assert attention.attn_backend == expected_backend + if expects_sol_params: + assert isinstance(attention.sparse_params, SolParams) + assert captured["sparse_params"] is attention.sparse_params + else: + assert attention.sparse_params is None + assert captured["sparse_params"] is None + + +@_CPU_ONLY +def test_sol_attention_rejects_context_parallelism() -> None: + with pytest.raises(ValueError, match="SOL.*incompatible with context parallelism"): + Attention( + hidden_size=256, + num_attention_heads=2, + head_dim=128, + qk_norm=False, + config=_sol_model_config(cp_size=2), + ) + + +@_CPU_ONLY +def test_sol_cuda_graph_phase_is_keyed_without_model_scope(monkeypatch) -> None: + q = _bshd() + backend, predictor = _stub_backend(SolParams(tau=1.0, disabled_until_timestep=0.6)) + model = _SolModel((backend,)) + runner = CUDAGraphRunner(CUDAGraphRunnerConfig(use_cuda_graph=True)) + model.register_cuda_graph_extra_key_fns(runner) + _stub_core_forward(monkeypatch) + monkeypatch.setattr(sol_backend, "get_bmm1_scale", lambda attn: 0.125) + capturing = False + captured_outputs = {} + captured_keys = [] + + def _capture(key, fn, args, kwargs): + nonlocal capturing + captured_outputs[key] = fn(*args, **kwargs) + capturing = True + try: + captured_outputs[key] = fn(*args, **kwargs) + captured_keys.append(key) + finally: + capturing = False + + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: capturing) + monkeypatch.setattr(runner, "capture", _capture) + monkeypatch.setattr(runner, "replay", lambda key, args, kwargs: captured_outputs[key]) + model.forward = runner.wrap(model.forward) + + assert model(q, timestep=torch.tensor(0.8)).shape == (1, 64, 256) + assert model(q, timestep=torch.tensor(0.2)).shape == (1, 64, 256) + assert ("sol_attn_phase", 0) in captured_keys[0] + assert ("sol_attn_phase", 1) in captured_keys[1] + assert captured_keys[0] != captured_keys[1] + assert predictor.predict.call_count == 2 + + +@_REQUIRES_SM100 +@torch.no_grad() +def test_real_b200_sol_backend_cuda_graph_matches_dense_reference() -> None: + attention_config = AttentionConfig( + backend="TRTLLM", + sparse_attention_config=SolAttentionConfig( + tau=-1.0e6, + disabled_until_timestep=0.6, + ), + ) + sparse_config = attention_config.sparse_attention_config + assert isinstance(sparse_config, SolAttentionConfig) + backend = create_attention( + backend="TRTLLM", + layer_idx=1, + num_heads=2, + head_dim=128, + dtype=torch.bfloat16, + attention_config=attention_config, + attention_metadata_state=create_attention_metadata_state(), + sparse_params=sparse_config.to_sparse_params(), + ) + assert isinstance(backend, SOLTrtllmAttention) + assert any(isinstance(fmha, PrimsTSBlockSparseFmha) for fmha in backend._fmha_manager.fmha_libs) + + generator = torch.Generator(device="cuda").manual_seed(20260901) + shape = (1, 257, 2, 128) + + def _inputs() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + packed = torch.randint( + -2, + 3, + (shape[0], shape[1], 3 * shape[2] * shape[3]), + generator=generator, + device="cuda", + dtype=torch.bfloat16, + ) + return tuple(tensor.view(shape) for tensor in packed.split(shape[2] * shape[3], dim=-1)) + + q, k, v = _inputs() + timestep = torch.tensor(0.2, device="cuda") + assert not any(tensor.is_contiguous() for tensor in (q, k, v)) + eager = backend.forward( + q=q, + k=k, + v=v, + batch_size=1, + seq_len=257, + seq_len_kv=257, + timestep=timestep, + ) + torch.cuda.synchronize() + torch.testing.assert_close( + eager.view_as(q), + _dense_reference(q, k, v), + rtol=2e-2, + atol=2e-2, + ) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = backend.forward( + q=q, + k=k, + v=v, + batch_size=1, + seq_len=257, + seq_len_kv=257, + timestep=timestep, + ).view_as(q) + + next_q, next_k, next_v = _inputs() + q.copy_(next_q) + k.copy_(next_k) + v.copy_(next_v) + graph.replay() + torch.cuda.synchronize() + + torch.testing.assert_close( + captured, + _dense_reference(q, k, v), + rtol=2e-2, + atol=2e-2, + ) + assert backend.predictor.num_plans == 1 + + +@_REQUIRES_SM100 +@torch.no_grad() +@pytest.mark.parametrize("seq_len", [256, 257]) +def test_real_b200_sol_backend_mixed_proxy_cuda_graph_matches_reference(seq_len: int) -> None: + attention_config = AttentionConfig( + backend="TRTLLM", + sparse_attention_config=SolAttentionConfig(tau=1.0e6), + ) + sparse_config = attention_config.sparse_attention_config + assert isinstance(sparse_config, SolAttentionConfig) + backend = create_attention( + backend="TRTLLM", + layer_idx=1, + num_heads=2, + head_dim=128, + dtype=torch.bfloat16, + attention_config=attention_config, + attention_metadata_state=create_attention_metadata_state(), + sparse_params=sparse_config.to_sparse_params(), + ) + assert isinstance(backend, SOLTrtllmAttention) + + generator = torch.Generator(device="cuda").manual_seed(20260903) + shape = (1, seq_len, 2, 128) + + def _inputs() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + packed = torch.randint( + -2, + 3, + (shape[0], shape[1], 3 * shape[2] * shape[3]), + generator=generator, + device="cuda", + dtype=torch.bfloat16, + ) + return tuple(tensor.view(shape) for tensor in packed.split(shape[2] * shape[3], dim=-1)) + + q, k, v = _inputs() + eager = backend.forward(q=q, k=k, v=v, batch_size=1, seq_len=seq_len, seq_len_kv=seq_len) + predictor_outputs = backend.predictor.predict( + q.contiguous(), + k.contiguous(), + v.contiguous(), + tau=1.0e6, + sm_scale=128**-0.5, + ) + torch.cuda.synchronize() + exact_bits = predictor_outputs.exact_block_bits + num_blocks = math.ceil(seq_len / 64) + num_exact = sum( + int( + (exact_bits[..., block_idx // 32].to(torch.int64) >> (block_idx % 32)) + .bitwise_and(1) + .sum() + .item() + ) + for block_idx in range(num_blocks) + ) + assert 0 < num_exact < math.prod(exact_bits.shape[:3]) * num_blocks + torch.testing.assert_close( + eager.view_as(q), + _mixed_proxy_reference(q, k, v, predictor_outputs), + rtol=2e-2, + atol=2e-2, + ) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = backend.forward( + q=q, + k=k, + v=v, + batch_size=1, + seq_len=seq_len, + seq_len_kv=seq_len, + ).view_as(q) + + next_q, next_k, next_v = _inputs() + q.copy_(next_q) + k.copy_(next_k) + v.copy_(next_v) + graph.replay() + torch.cuda.synchronize() + + torch.testing.assert_close( + captured, + _mixed_proxy_reference(q, k, v, predictor_outputs), + rtol=2e-2, + atol=2e-2, + ) diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/test_sol_predictor.py b/tests/unittest/_torch/visual_gen/sparse_attention/test_sol_predictor.py new file mode 100644 index 000000000000..15b31deb210b --- /dev/null +++ b/tests/unittest/_torch/visual_gen/sparse_attention/test_sol_predictor.py @@ -0,0 +1,374 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Qualification tests for the two-stage VisualGen SOL predictor.""" + +from __future__ import annotations + +import dataclasses +import math +import struct + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.predictor import ( + SolPredictorGeometry, + SolPredictorOutputs, + SolPredictorPlanKey, + SOLSparsePredictor, + _normalize_runtime_scalars, +) + +_CPU_ONLY = pytest.mark.cpu_only + + +@_CPU_ONLY +def test_sol_predictor_geometry_and_static_plan_contract() -> None: + geometry = SolPredictorGeometry.create(batch_size=2, seq_len=257, num_heads=3) + assert ( + geometry.tensor_shape, + geometry.summary_shape, + geometry.stats_shape, + geometry.exact_block_bits_shape, + ) == ((2, 257, 3, 128), (2, 5, 3, 128), (2, 3, 128), (2, 3, 5, 1)) + assert (geometry.num_q_blocks, geometry.num_kv_blocks, geometry.tail_tokens) == (5, 5, 1) + + boundary_cases = ( + (64, 1, 1, 64), + (65, 2, 1, 1), + (64 * 32, 32, 1, 64), + (64 * 32 + 1, 33, 2, 1), + ) + for seq_len, blocks, words, tail in boundary_cases: + current = SolPredictorGeometry.create(batch_size=1, seq_len=seq_len, num_heads=1) + assert (current.num_q_blocks, current.exact_words, current.tail_tokens) == ( + blocks, + words, + tail, + ) + + key = SolPredictorPlanKey(geometry=geometry, device_index=1, dtype=torch.bfloat16) + assert tuple(field.name for field in dataclasses.fields(key)) == ( + "geometry", + "device_index", + "dtype", + ) + assert "tau" not in repr(key) and "sm_scale" not in repr(key) + assert SOLSparsePredictor().num_plans == 0 + + +@_CPU_ONLY +def test_sol_predictor_validates_geometry_and_runtime_scalars() -> None: + invalid_geometry = ( + ({"batch_size": 0, "seq_len": 64, "num_heads": 1}, "batch_size"), + ({"batch_size": 1, "seq_len": 0, "num_heads": 1}, "seq_len"), + ({"batch_size": 1, "seq_len": 64, "num_heads": 0}, "num_heads"), + ({"batch_size": True, "seq_len": 64, "num_heads": 1}, "batch_size"), + ({"batch_size": 1, "seq_len": 64, "num_heads": 1, "head_dim": 64}, "head_dim=128"), + ) + for kwargs, message in invalid_geometry: + with pytest.raises((TypeError, ValueError), match=message): + SolPredictorGeometry.create(**kwargs) + + tau, sm_scale = _normalize_runtime_scalars(tau=0.1, sm_scale=math.sqrt(0.5)) + expected_tau = struct.unpack("=f", struct.pack("=f", 0.1))[0] + expected_scale = struct.unpack("=f", struct.pack("=f", math.sqrt(0.5)))[0] + assert (tau, sm_scale) == (expected_tau, expected_scale) + + invalid_scalars = ( + (True, 0.125, "tau"), + (math.nan, 0.125, "tau"), + (0.0, True, "sm_scale"), + (0.0, math.inf, "sm_scale"), + (0.0, 0.0, "sm_scale"), + (0.0, -0.125, "sm_scale"), + ) + for invalid_tau, invalid_scale, message in invalid_scalars: + with pytest.raises((TypeError, ValueError), match=message): + _normalize_runtime_scalars(tau=invalid_tau, sm_scale=invalid_scale) + + +_REQUIRES_CUDA = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +_LOG2_E = math.log2(math.e) + + +def _summary_oracle(k: torch.Tensor, v: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + blocks = (k.shape[1] + 63) // 64 + k_summary = torch.empty( + (k.shape[0], blocks, k.shape[2], k.shape[3]), + dtype=torch.bfloat16, + device=k.device, + ) + v_summary = torch.empty_like(k_summary) + for block_idx in range(blocks): + begin = block_idx * 64 + end = min(begin + 64, k.shape[1]) + k_summary[:, block_idx] = k[:, begin:end].float().mean(dim=1).to(torch.bfloat16) + v_summary[:, block_idx] = v[:, begin:end].float().sum(dim=1).to(torch.bfloat16) + return k_summary, v_summary + + +def _pack_bits(exact: torch.Tensor) -> torch.Tensor: + words = (exact.shape[-1] + 31) // 32 + padded = F.pad(exact, (0, words * 32 - exact.shape[-1])).view(*exact.shape[:-1], words, 32) + powers = 1 << torch.arange(32, dtype=torch.int64, device=exact.device) + return (padded.to(torch.int64) * powers).sum(dim=-1).to(torch.uint32) + + +def _predictor_oracle( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + tau: float, + sm_scale: float, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + k_summary, v_summary = _summary_oracle(k, v) + blocks = k_summary.shape[1] + padded_q = F.pad(q, (0, 0, 0, 0, 0, blocks * 64 - q.shape[1])) + q_blocks = padded_q.view(q.shape[0], blocks, 64, q.shape[2], q.shape[3]) + q_lengths = torch.clamp( + q.shape[1] - torch.arange(blocks, device=q.device) * 64, + min=1, + max=64, + ) + q_centroids = q_blocks.float().sum(dim=2) / q_lengths[None, :, None, None] + k_float = k_summary.float() + k_mean = k_float.mean(dim=1) + k_var = torch.clamp(k_float.square().mean(dim=1) - k_mean.square(), min=0.0) + log2_scale = float(sm_scale) * _LOG2_E + projected_mean = torch.einsum("bqhd,bhd->bqh", q_centroids, k_mean) * log2_scale + projected_var = ( + torch.einsum("bqhd,bhd->bqh", q_centroids.square(), k_var) * log2_scale * log2_scale + ) + threshold = projected_mean + float(tau) * torch.sqrt(projected_var + 1.0e-6) + scores = torch.einsum("bqhd,bkhd->bhqk", q_centroids, k_float) * log2_scale + exact = scores > threshold.permute(0, 2, 1).unsqueeze(-1) + block_ids = torch.arange(blocks, device=q.device) + exact |= (block_ids[:, None] - block_ids[None, :]).abs()[None, None] <= 1 + return _pack_bits(exact), k_summary, v_summary + + +def _small_integer_bf16(shape: tuple[int, ...], *, seed: int) -> torch.Tensor: + generator = torch.Generator(device="cuda").manual_seed(seed) + return torch.randint(-2, 3, shape, generator=generator, device="cuda", dtype=torch.bfloat16) + + +def _inputs( + shape: tuple[int, int, int, int], seed: int +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return ( + _small_integer_bf16(shape, seed=seed), + _small_integer_bf16(shape, seed=seed + 1), + _small_integer_bf16(shape, seed=seed + 2), + ) + + +def _output_tensors( + outputs: SolPredictorOutputs, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return outputs.exact_block_bits, outputs.k_summary, outputs.v_summary + + +def _assert_outputs_match( + outputs: SolPredictorOutputs, + expected: tuple[torch.Tensor, torch.Tensor, torch.Tensor], +) -> None: + assert torch.equal(outputs.exact_block_bits, expected[0]) + torch.testing.assert_close(outputs.k_summary, expected[1], rtol=1e-2, atol=1e-2) + torch.testing.assert_close(outputs.v_summary, expected[2], rtol=1e-2, atol=2e-2) + + +def _run_custom_op( + q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, *buffers: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + torch.ops.trtllm.visual_gen_sol_predictor(q, k, v, *buffers, 64, 0.5, 0.125) + return buffers[0], buffers[1], buffers[2] + + +@_REQUIRES_CUDA +def test_sol_predictor_s257_bh_gt_one_matches_oracle_and_reuses_storage() -> None: + q, k, v = _inputs((2, 257, 3, 128), 11) + predictor = SOLSparsePredictor() + plan = predictor.prepare(q, k, v) + output_ids = tuple(map(id, _output_tensors(plan.outputs))) + scratch_ids = (id(plan.k_mean), id(plan.k_var_diag)) + + outputs = predictor.predict(q, k, v, tau=0.75, sm_scale=0.125) + reference = _predictor_oracle(q, k, v, tau=0.75, sm_scale=0.125) + + assert outputs is plan.outputs + assert tuple(map(id, _output_tensors(outputs))) == output_ids + assert (id(plan.k_mean), id(plan.k_var_diag)) == scratch_ids + _assert_outputs_match(outputs, reference) + expected_k_mean = reference[1].float().mean(dim=1) + expected_k_var = torch.clamp( + reference[1].float().square().mean(dim=1) - expected_k_mean.square(), + min=0.0, + ) + torch.testing.assert_close(plan.k_mean, expected_k_mean, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(plan.k_var_diag, expected_k_var, rtol=1e-5, atol=1e-5) + + second = predictor.predict(q, k, v, tau=-0.25, sm_scale=0.0625) + assert second is outputs + assert tuple(map(id, _output_tensors(second))) == output_ids + assert predictor.num_plans == 1 + + +@_REQUIRES_CUDA +def test_sol_predictor_s257_runtime_scale_and_tau_extremes() -> None: + q, k, v = _inputs((1, 257, 2, 128), 61) + predictor = SOLSparsePredictor() + + normal = predictor.predict(q, k, v, tau=0.5, sm_scale=0.125) + normal_bits = normal.exact_block_bits.clone() + expected_normal = _predictor_oracle(q, k, v, tau=0.5, sm_scale=0.125)[0] + tiny = predictor.predict(q, k, v, tau=0.5, sm_scale=1.0e-5) + expected_tiny = _predictor_oracle(q, k, v, tau=0.5, sm_scale=1.0e-5)[0] + + assert tiny is normal + assert torch.equal(normal_bits, expected_normal) + assert torch.equal(tiny.exact_block_bits, expected_tiny) + assert not torch.equal(expected_normal, expected_tiny) + + blocks = 5 + block_ids = torch.arange(blocks, device=q.device) + local = (block_ids[:, None] - block_ids[None, :]).abs() <= 1 + expected_extremes = ( + _pack_bits(local[None, None].expand(1, 2, -1, -1)), + _pack_bits(torch.ones((1, 2, blocks, blocks), device=q.device, dtype=torch.bool)), + ) + for tau, expected in zip((1.0e6, -1.0e6), expected_extremes, strict=True): + outputs = predictor.predict(q, k, v, tau=tau, sm_scale=128**-0.5) + assert torch.equal(outputs.exact_block_bits, expected) + + +@_REQUIRES_CUDA +def test_sol_predictor_long_proxy_group_keeps_tail_mass_and_clears_padding_bits() -> None: + tokens = 16_451 + q, k, _ = _inputs((1, tokens, 1, 128), 31) + v = torch.ones_like(q) + outputs = SOLSparsePredictor().predict(q, k, v, tau=1.0e6, sm_scale=0.125) + expected, expected_k, expected_v = _predictor_oracle(q, k, v, tau=1.0e6, sm_scale=0.125) + + assert outputs.k_summary.shape[1] == 258 + assert outputs.exact_block_bits.shape[-1] == 9 + assert torch.equal(outputs.exact_block_bits, expected) + torch.testing.assert_close(outputs.k_summary[:, -1], expected_k[:, -1], rtol=1e-2, atol=1e-2) + assert torch.equal(outputs.v_summary[:, -1], expected_v[:, -1]) + assert torch.all(outputs.v_summary[:, -1] == 3) + assert int(outputs.exact_block_bits[..., -1].to(torch.int64).max()) < 4 + + +@_REQUIRES_CUDA +def test_sol_predictor_cuda_graph_replay_updates_live_outputs() -> None: + q, k, v = _inputs((1, 257, 2, 128), 41) + predictor = SOLSparsePredictor() + plan = predictor.prepare(q, k, v) + predictor.predict(q, k, v, tau=0.5, sm_scale=0.125) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = predictor.predict(q, k, v, tau=0.5, sm_scale=0.125) + assert captured is plan.outputs + + next_q, next_k, next_v = _inputs(q.shape, 51) + q.copy_(next_q) + k.copy_(next_k) + v.copy_(next_v) + graph.replay() + torch.cuda.synchronize() + expected = _predictor_oracle(q, k, v, tau=0.5, sm_scale=0.125) + + _assert_outputs_match(captured, expected) + + +@_REQUIRES_CUDA +def test_sol_predictor_rejects_plan_miss_during_capture_and_reuses_prepared_plan( + monkeypatch, +) -> None: + q, k, v = _inputs((1, 193, 1, 128), 71) + + with monkeypatch.context() as capture: + capture.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + with pytest.raises(RuntimeError, match="plan must be prepared"): + SOLSparsePredictor().predict(q, k, v, tau=0.5, sm_scale=0.125) + + predictor = SOLSparsePredictor() + plan = predictor.prepare(q, k, v) + with monkeypatch.context() as capture: + capture.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + assert predictor.predict(q, k, v, tau=0.5, sm_scale=0.125) is plan.outputs + torch.cuda.synchronize() + _assert_outputs_match(plan.outputs, _predictor_oracle(q, k, v, tau=0.5, sm_scale=0.125)) + + +@_REQUIRES_CUDA +def test_sol_predictor_compiled_public_predict_owns_each_instance_plan(recwarn) -> None: + output_ptrs = [] + for seed in (81, 91): + q, k, v = _inputs((1, 257, 2, 128), seed) + predictor = SOLSparsePredictor() + compiled_predict = torch.compile(predictor.predict, backend="eager", fullgraph=False) + + outputs = compiled_predict(q, k, v, tau=0.5, sm_scale=0.125) + expected = _predictor_oracle(q, k, v, tau=0.5, sm_scale=0.125) + + assert predictor.num_plans == 1 + _assert_outputs_match(outputs, expected) + output_ptrs.append(tuple(tensor.data_ptr() for tensor in _output_tensors(outputs))) + + assert output_ptrs[0] != output_ptrs[1] + assert not any("recompile_limit" in str(warning.message) for warning in recwarn) + + +@_REQUIRES_CUDA +def test_sol_predictor_custom_op_fake_schema_and_fullgraph_compile() -> None: + from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol import kernels # noqa: F401 + + op = torch.ops.trtllm.visual_gen_sol_predictor.default + schema = str(op._schema) + assert "exact_block_bits" in schema + assert "k_summary" in schema + assert "v_summary" in schema + assert "!" in schema + assert torch._C._dispatch_has_kernel_for_dispatch_key( + "trtllm::visual_gen_sol_predictor", "Meta" + ) + + meta_q = torch.empty((1, 65, 1, 128), device="meta", dtype=torch.bfloat16) + meta_summary = torch.empty((1, 2, 1, 128), device="meta", dtype=torch.bfloat16) + meta_stats = torch.empty((1, 1, 128), device="meta", dtype=torch.float32) + meta_args = ( + meta_q, + torch.empty_like(meta_q), + torch.empty_like(meta_q), + torch.empty((1, 1, 2, 1), device="meta", dtype=torch.uint32), + meta_summary, + torch.empty_like(meta_summary), + meta_stats, + torch.empty_like(meta_stats), + torch.empty((1, 2, 1, 128), device="meta", dtype=torch.float32), + ) + assert op(*meta_args, 64, 0.5, 0.125) is None + + q, k, v = _inputs((1, 257, 2, 128), 81) + plan = SOLSparsePredictor().prepare(q, k, v) + buffers = ( + plan.outputs.exact_block_bits, + plan.outputs.k_summary, + plan.outputs.v_summary, + plan.k_mean, + plan.k_var_diag, + plan.q_centroid, + ) + actual = torch.compile(_run_custom_op, backend="eager", fullgraph=True)(q, k, v, *buffers) + expected = _predictor_oracle(q, k, v, tau=0.5, sm_scale=0.125) + + assert all( + actual_tensor is plan_tensor + for actual_tensor, plan_tensor in zip(actual, buffers[:3], strict=True) + ) + _assert_outputs_match(plan.outputs, expected) diff --git a/tests/unittest/_torch/visual_gen/sparse_attention/test_sol_predictor_kernels.py b/tests/unittest/_torch/visual_gen/sparse_attention/test_sol_predictor_kernels.py new file mode 100644 index 000000000000..79556dbcf31d --- /dev/null +++ b/tests/unittest/_torch/visual_gen/sparse_attention/test_sol_predictor_kernels.py @@ -0,0 +1,206 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for the SOL predictor kernels (block pooling, block statistics, exact-block selection).""" + +from __future__ import annotations + +import pytest +import torch + +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.kernels import ( + _block_pool_torch, + _block_statistics_torch, + _select_exact_blocks_torch, + block_pool, + block_statistics, + select_exact_blocks, +) + +_CPU_ONLY = pytest.mark.cpu_only +_REQUIRES_CUDA = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +BLOCK = 64 + + +def _pooled_shape(x: torch.Tensor, block_size: int) -> tuple[int, int, int, int]: + batch, seq_len, heads, head_dim = x.shape + return batch, (seq_len + block_size - 1) // block_size, heads, head_dim + + +def _unpack(bits: torch.Tensor, num_kv_blocks: int) -> torch.Tensor: + words = bits.view(torch.int32).to(torch.int64) & 0xFFFFFFFF + shifts = torch.arange(32, device=bits.device, dtype=torch.int64) + return ( + ((words.unsqueeze(-1) >> shifts) & 1) + .bool() + .reshape(*bits.shape[:-1], -1)[..., :num_kv_blocks] + ) + + +@_CPU_ONLY +def test_block_pool_torch_fallback_means_valid_tokens_only() -> None: + x = torch.zeros((1, 70, 1, 4), dtype=torch.bfloat16) + x[0, :64] = 2.0 + x[0, 64:70] = 3.0 + out = torch.empty(_pooled_shape(x, BLOCK), dtype=torch.float32) + block_pool(x, out, block_size=BLOCK, reduce="mean") + assert torch.equal(out[0, 0], torch.full((1, 4), 2.0)) + assert torch.equal(out[0, 1], torch.full((1, 4), 3.0)) + total = torch.empty_like(out) + block_pool(x, total, block_size=BLOCK, reduce="sum") + assert torch.equal(total[0, 1], torch.full((1, 4), 18.0)) + + +@_CPU_ONLY +def test_block_pool_rejects_mismatched_output() -> None: + x = torch.zeros((1, 70, 1, 4), dtype=torch.bfloat16) + with pytest.raises(ValueError, match="out"): + block_pool(x, torch.empty((1, 3, 1, 4)), block_size=BLOCK, reduce="mean") + with pytest.raises(ValueError, match="reduce"): + block_pool(x, torch.empty(_pooled_shape(x, BLOCK)), block_size=BLOCK, reduce="max") + + +@_CPU_ONLY +def test_select_exact_blocks_torch_fallback_packs_bit_r_of_word_w() -> None: + blocks = 35 + centroid = torch.zeros((1, blocks, 1, 8), dtype=torch.float32) + k_summary = torch.zeros((1, blocks, 1, 8), dtype=torch.bfloat16) + k_summary[0, 33, 0, 0] = 1.0 + centroid[0, 0, 0, 0] = 1.0 + k_mean = torch.zeros((1, 1, 8)) + k_var = torch.zeros((1, 1, 8)) + bits = torch.empty((1, 1, blocks, 2), dtype=torch.uint32) + select_exact_blocks(centroid, k_summary, k_mean, k_var, bits, tau=0.5, sm_scale=1.0) + exact = _unpack(bits, blocks) + # Row 0 scores 1.0 against block 33 only (threshold 0.5 * sqrt(1e-6)) plus its local band. + expected = torch.zeros(blocks, dtype=torch.bool) + expected[[0, 1, 33]] = True + assert torch.equal(exact[0, 0, 0], expected) + # Bit 1 of word 1 is block 33. + assert int(bits[0, 0, 0, 1]) == 2 + # Rows without scores keep only the local band. + assert torch.equal(exact[0, 0, 17].nonzero().flatten(), torch.tensor([16, 17, 18])) + + +@_REQUIRES_CUDA +@pytest.mark.parametrize("seq_len", [64, 257, 4097]) +@pytest.mark.parametrize("out_dtype", [torch.float32, torch.bfloat16]) +def test_block_pool_matches_torch_fallback(seq_len: int, out_dtype: torch.dtype) -> None: + torch.manual_seed(0) + x = torch.randn((2, seq_len, 3, 128), device="cuda", dtype=torch.bfloat16) + for reduce in ("mean", "sum"): + out = torch.empty(_pooled_shape(x, BLOCK), dtype=out_dtype, device="cuda") + block_pool(x, out, block_size=BLOCK, reduce=reduce) + expected = torch.empty_like(out) + _block_pool_torch(x, expected, block_size=BLOCK, reduce=reduce) + tolerance = ( + {"rtol": 1e-5, "atol": 1e-5} + if out_dtype == torch.float32 + else {"rtol": 1e-2, "atol": 1e-2} + ) + torch.testing.assert_close(out, expected, **tolerance) + + +@_REQUIRES_CUDA +def test_block_pool_accepts_strided_batch_and_token_dims() -> None: + torch.manual_seed(1) + full = torch.randn((2, 130, 2, 3, 128), device="cuda", dtype=torch.bfloat16) + x = full[:, :, 1] # heads/head_dim contiguous, token stride wider than a row + out = torch.empty(_pooled_shape(x, BLOCK), dtype=torch.float32, device="cuda") + block_pool(x, out, block_size=BLOCK, reduce="mean") + expected = torch.empty_like(out) + _block_pool_torch(x.contiguous(), expected, block_size=BLOCK, reduce="mean") + torch.testing.assert_close(out, expected, rtol=1e-5, atol=1e-5) + + +@_REQUIRES_CUDA +def test_block_statistics_matches_torch_fallback() -> None: + torch.manual_seed(2) + k_summary = torch.randn((2, 1182, 3, 128), device="cuda", dtype=torch.bfloat16) + mean = torch.empty((2, 3, 128), device="cuda", dtype=torch.float32) + var = torch.empty_like(mean) + block_statistics(k_summary, mean, var) + expected_mean = torch.empty_like(mean) + expected_var = torch.empty_like(var) + _block_statistics_torch(k_summary, expected_mean, expected_var) + torch.testing.assert_close(mean, expected_mean, rtol=1e-5, atol=1e-6) + torch.testing.assert_close(var, expected_var, rtol=1e-4, atol=1e-6) + assert bool((var >= 0).all()) + + +@_REQUIRES_CUDA +@pytest.mark.parametrize("num_blocks", [5, 258]) +def test_select_exact_blocks_matches_fallback_and_clears_padding_bits(num_blocks: int) -> None: + torch.manual_seed(3) + batch, heads, dim = 2, 3, 128 + centroid = torch.randn((batch, num_blocks, heads, dim), device="cuda") * 0.125 + k_summary = (torch.randn((batch, num_blocks, heads, dim), device="cuda") * 0.125).to( + torch.bfloat16 + ) + k_mean = torch.empty((batch, heads, dim), device="cuda") + k_var = torch.empty_like(k_mean) + block_statistics(k_summary, k_mean, k_var) + words = (num_blocks + 31) // 32 + bits = torch.empty((batch, heads, num_blocks, words), device="cuda", dtype=torch.uint32) + expected = torch.empty_like(bits) + for tau in (0.75, -1.0e6, 1.0e6): + select_exact_blocks(centroid, k_summary, k_mean, k_var, bits, tau=tau, sm_scale=0.125) + _select_exact_blocks_torch( + centroid, k_summary, k_mean, k_var, expected, tau=tau, sm_scale=0.125 + ) + assert torch.equal(bits, expected), f"tau={tau}" + padding = words * 32 - num_blocks + if padding: + assert int(bits[..., -1].to(torch.int64).max()) < (1 << (32 - padding)) + exact = _unpack(bits, num_blocks) + ids = torch.arange(num_blocks, device="cuda") + assert bool(exact[..., (ids[:, None] - ids[None, :]).abs() <= 1].all()) + + +@_REQUIRES_CUDA +def test_kernels_replay_inside_cuda_graph() -> None: + torch.manual_seed(4) + q = torch.randn((1, 257, 2, 128), device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + shape = _pooled_shape(q, BLOCK) + centroid = torch.empty(shape, device="cuda", dtype=torch.float32) + k_summary = torch.empty(shape, device="cuda", dtype=torch.bfloat16) + mean = torch.empty((1, 2, 128), device="cuda") + var = torch.empty_like(mean) + bits = torch.empty((1, 2, shape[1], 1), device="cuda", dtype=torch.uint32) + + def run() -> None: + block_pool(q, centroid, block_size=BLOCK, reduce="mean") + block_pool(k, k_summary, block_size=BLOCK, reduce="mean") + block_statistics(k_summary, mean, var) + select_exact_blocks(centroid, k_summary, mean, var, bits, tau=0.5, sm_scale=0.125) + + run() + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + run() + q.copy_(torch.randn_like(q)) + k.copy_(torch.randn_like(k)) + graph.replay() + torch.cuda.synchronize() + + expected_centroid = torch.empty_like(centroid) + expected_summary = torch.empty_like(k_summary) + _block_pool_torch(q, expected_centroid, block_size=BLOCK, reduce="mean") + _block_pool_torch(k, expected_summary, block_size=BLOCK, reduce="mean") + expected_mean = torch.empty_like(mean) + expected_var = torch.empty_like(var) + _block_statistics_torch(expected_summary, expected_mean, expected_var) + expected_bits = torch.empty_like(bits) + _select_exact_blocks_torch( + expected_centroid, + expected_summary, + expected_mean, + expected_var, + expected_bits, + tau=0.5, + sm_scale=0.125, + ) + torch.testing.assert_close(centroid, expected_centroid, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(k_summary, expected_summary, rtol=1e-2, atol=1e-2) + assert torch.equal(bits, expected_bits) diff --git a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_vsa.py b/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_vsa.py deleted file mode 100644 index 63684cd0519a..000000000000 --- a/tests/unittest/_torch/visual_gen/test_attention_cute_dsl_vsa.py +++ /dev/null @@ -1,478 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""VSA correctness tests: CuTe kernel, tile/untile roundtrip, top-k math, backend guards. - -Module-level dense-equivalence and finite-output checks live in -test_attention_integration.py. -""" - -from types import SimpleNamespace - -import pytest -import torch -import torch.nn.functional as F - -from tensorrt_llm._torch.visual_gen.attention_backend import ( - CuTeDSLAttention, - VSAAttention, - VSAMetadataBuilder, -) -from tensorrt_llm._torch.visual_gen.attention_backend.utils import create_attention -from tensorrt_llm._torch.visual_gen.config import ( - DiffusionModelConfig, - create_attention_metadata_state, -) -from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode -from tensorrt_llm.visual_gen.args import ( - AttentionConfig, - QuantAttentionConfig, - VideoSparseAttentionConfig, -) - - -def test_cute_dsl_factory_dispatches_quantized_fmha_and_vsa() -> None: - quant_config = QuantAttentionConfig(qk_dtype="mxfp8", v_dtype="fp8", v_block_size=1) - dense_config = AttentionConfig(backend="CUTEDSL", quant_attention_config=quant_config) - dense_attention = create_attention( - backend="CUTEDSL", - layer_idx=0, - num_heads=8, - head_dim=128, - attention_config=dense_config, - ) - - sparse_config = VideoSparseAttentionConfig(vsa_sparsity=0.9) - vsa_config = AttentionConfig(backend="CUTEDSL", sparse_attention_config=sparse_config) - vsa_attention = create_attention( - backend="CUTEDSL", - layer_idx=0, - num_heads=8, - head_dim=128, - attention_config=vsa_config, - ) - - assert isinstance(dense_attention, CuTeDSLAttention) - assert dense_attention.quant_attention_config is quant_config - assert isinstance(vsa_attention, VSAAttention) - assert vsa_attention.sparse_attention_config is sparse_config - - -def _make_config( - hidden_size: int, - num_heads: int, - head_dim: int, - backend: str, - vsa_sparsity: "float | None" = None, -) -> DiffusionModelConfig: - """Minimal DiffusionModelConfig for one Attention module.""" - pretrained_config = SimpleNamespace( - hidden_size=hidden_size, - num_attention_heads=num_heads, - attention_head_dim=head_dim, - eps=1e-6, - ) - sparse_attention_config = ( - VideoSparseAttentionConfig(vsa_sparsity=vsa_sparsity) if vsa_sparsity is not None else None - ) - config = DiffusionModelConfig( - pretrained_config=pretrained_config, - attention=AttentionConfig(backend=backend, sparse_attention_config=sparse_attention_config), - skip_create_weights_in_init=False, - ) - config.attention_metadata_state = ( - create_attention_metadata_state() if backend == "TRTLLM" else None - ) - return config - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="VSA needs CUDA") -def test_vsa_falls_back_to_vanilla_for_cross_attention(): - """Cross-attention (SEPARATE_QKV) falls back to VANILLA — it has no cube structure.""" - device = torch.device("cuda") - dtype = torch.bfloat16 - cfg = _make_config( - hidden_size=64, num_heads=4, head_dim=16, backend="CUTEDSL", vsa_sparsity=0.5 - ) - cross_attn = ( - Attention(64, 4, qkv_mode=QKVMode.SEPARATE_QKV, config=cfg) - .to(device=device, dtype=dtype) - .eval() - ) - assert cross_attn.attn_backend == "VANILLA", ( - f"VSA on cross-attention should fall back to VANILLA, got {cross_attn.attn_backend!r}" - ) - - -def test_vsa_with_attn2d_raises(): - """VSA + Attention2D must error at construction (VSA needs the full sequence per rank).""" - pretrained_config = SimpleNamespace( - hidden_size=64, - num_attention_heads=4, - attention_head_dim=16, - eps=1e-6, - ) - cfg = DiffusionModelConfig( - pretrained_config=pretrained_config, - attention=AttentionConfig( - backend="CUTEDSL", - sparse_attention_config=VideoSparseAttentionConfig(vsa_sparsity=0.0), - ), - skip_create_weights_in_init=False, - ) - cfg.visual_gen_mapping = SimpleNamespace( - ring_size=1, - ring_group=None, - ulysses_size=1, - ulysses_group=None, - attn2d_row_size=2, - attn2d_col_size=2, - attn2d_row_group=None, - attn2d_col_group=None, - cp_size=4, - ) - with pytest.raises(ValueError, match="incompatible with context parallelism"): - Attention(64, 4, qkv_mode=QKVMode.FUSE_QKV, config=cfg) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="VSA needs CUDA") -def test_vsa_topk_collapses_to_dense_at_sparsity_zero(): - """At sparsity=0, top_k equals num_cubes (dense connectivity).""" - from math import ceil - - device = torch.device("cuda") - builder = VSAMetadataBuilder() - metadata = builder.build( - current_timestep=0, - raw_latent_shape=(8, 8, 8), - patch_size=(1, 1, 1), - vsa_sparsity=0.0, - device=device, - ) - num_cubes = metadata.num_tiles[0] * metadata.num_tiles[1] * metadata.num_tiles[2] - cur_topk = max(1, ceil((1.0 - metadata.vsa_sparsity) * num_cubes)) - assert cur_topk == num_cubes, ( - f"sparsity=0 should select all {num_cubes} cubes, got top_k={cur_topk}" - ) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="VSA needs CUDA") -@pytest.mark.parametrize( - "latent_shape", - [ - (8, 8, 8), - (9, 9, 9), - (21, 45, 80), - ], - ids=["clean_8x8x8", "ragged_9x9x9", "wan720p_21x45x80"], -) -def test_vsa_tile_untile_roundtrip(latent_shape): - """VSAPreprocessor.tile then .untile must losslessly reproduce the input.""" - from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl.vsa import VSAPreprocessor - - device = torch.device("cuda") - dtype = torch.bfloat16 - torch.manual_seed(0) - - B, H, D = 2, 4, 32 - seq_len = latent_shape[0] * latent_shape[1] * latent_shape[2] - - builder = VSAMetadataBuilder() - meta = builder.build( - current_timestep=0, - raw_latent_shape=latent_shape, - patch_size=(1, 1, 1), - vsa_sparsity=0.0, - device=device, - ) - - x = torch.randn(B, seq_len, H, D, device=device, dtype=dtype) - - x_tiled = VSAPreprocessor.tile( - x, - meta.non_pad_index, - meta.gather_idx, - meta.padded_seq_length, - ) - - pad_mask = torch.ones(meta.padded_seq_length, dtype=torch.bool, device=device) - pad_mask[meta.non_pad_index] = False - if pad_mask.any(): - assert x_tiled[:, pad_mask, :, :].abs().max().item() == 0.0, ( - "tile() must zero-fill padded positions" - ) - - x_roundtrip = VSAPreprocessor.untile( - x_tiled, - meta.reverse_tile_partition_indices, - meta.non_pad_index, - ) - - assert x_roundtrip.shape == x.shape, ( - f"shape mismatch after tile/untile: {x_roundtrip.shape} vs {x.shape}" - ) - assert torch.equal(x_roundtrip, x), ( - f"tile/untile round-trip is not lossless for latent_shape={latent_shape}: " - f"max_diff={(x_roundtrip - x).abs().max().item():.3e}" - ) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="kernel test needs CUDA") -def test_cute_kernel_matches_dense_at_full_topk(): - """CuTe block-sparse kernel matches dense SDPA when every cube is selected.""" - from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.video_sparse_attention import ( - CUTE_AVAILABLE, - block_sparse_attn_from_indices_cute, - is_cute_supported, - ) - - if not CUTE_AVAILABLE: - pytest.skip("cuda-bindings or cutlass-dsl not importable") - - device = torch.device("cuda") - dtype = torch.bfloat16 - torch.manual_seed(0) - - B, H, num_cubes, D = 1, 4, 4, 128 - block_size = 64 - seq_len = num_cubes * block_size - - q = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) - k = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) - v = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) - - if not is_cute_supported(q): - pytest.skip("CuTe path needs sm_100+ Blackwell (current device unsupported)") - - topk = num_cubes - q2k_idx = ( - torch.arange(num_cubes, device=device, dtype=torch.int32) - .view(1, 1, 1, num_cubes) - .expand(B, H, num_cubes, topk) - .contiguous() - ) - q2k_num = torch.full((B, H, num_cubes), topk, dtype=torch.int32, device=device) - variable_block_sizes = torch.full((num_cubes,), block_size, dtype=torch.int32, device=device) - - out_kernel, _lse = block_sparse_attn_from_indices_cute( - q, k, v, q2k_idx, q2k_num, variable_block_sizes - ) - out_ref = F.scaled_dot_product_attention(q, k, v) - - max_diff = (out_kernel - out_ref).abs().max().item() - mean_diff = (out_kernel - out_ref).abs().mean().item() - - rtol, atol = 1e-2, 1e-2 - assert torch.allclose(out_kernel, out_ref, rtol=rtol, atol=atol), ( - f"CuTe block-sparse kernel deviates from dense SDPA at full top-K: " - f"max_diff={max_diff:.3e}, mean_diff={mean_diff:.3e} (rtol={rtol}, atol={atol})" - ) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="kernel test needs CUDA") -def test_cute_kernel_matches_ref_with_independent_indices(): - """CuTe kernel: paired Q-blocks (2i, 2i+1) attend to independent KV index lists.""" - from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.video_sparse_attention import ( - CUTE_AVAILABLE, - block_sparse_attn_from_indices_cute, - is_cute_supported, - ) - - if not CUTE_AVAILABLE: - pytest.skip("cuda-bindings or cutlass-dsl not importable") - - device = torch.device("cuda") - dtype = torch.bfloat16 - torch.manual_seed(42) - - B, H, num_cubes, D = 2, 4, 16, 128 - block_size = 64 - topk = num_cubes // 2 - seq_len = num_cubes * block_size - - q = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) - k = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) - v = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) - - if not is_cute_supported(q): - pytest.skip("CuTe path needs sm_100+ Blackwell (current device unsupported)") - - q2k_idx = ( - torch.stack( - [ - torch.randperm(num_cubes, device=device, dtype=torch.int32)[:topk] - for _ in range(B * H * num_cubes) - ] - ) - .view(B, H, num_cubes, topk) - .contiguous() - ) - - paired = q2k_idx.view(B, H, num_cubes // 2, 2, topk).sort(dim=-1).values - pair_mismatch = (paired[..., 0, :] != paired[..., 1, :]).sum().item() - assert pair_mismatch > 0, ( - "Pre-condition failed: random permutations matched across every pair; " - "re-seed or raise num_cubes." - ) - - q2k_num = torch.full((B, H, num_cubes), topk, dtype=torch.int32, device=device) - variable_block_sizes = torch.full((num_cubes,), block_size, dtype=torch.int32, device=device) - - attn_mask = torch.full( - (B, H, seq_len, seq_len), float("-inf"), device=device, dtype=torch.float32 - ) - for b in range(B): - for h in range(H): - for q_blk in range(num_cubes): - for ki in range(topk): - k_blk = q2k_idx[b, h, q_blk, ki].item() - qs = q_blk * block_size - ks = k_blk * block_size - attn_mask[b, h, qs : qs + block_size, ks : ks + block_size] = 0.0 - - out_kernel, _lse = block_sparse_attn_from_indices_cute( - q, k, v, q2k_idx, q2k_num, variable_block_sizes - ) - - scale = 1.0 / (D**0.5) - scores = (q.float() @ k.float().transpose(-2, -1)) * scale - scores = scores + attn_mask - probs = torch.softmax(scores, dim=-1) - out_ref = (probs @ v.float()).to(dtype) - - abs_diff = (out_kernel.float() - out_ref.float()).abs() - max_diff = abs_diff.max().item() - mean_diff = abs_diff.mean().item() - - rtol, atol = 1e-2, 1e-2 - assert torch.allclose(out_kernel, out_ref, rtol=rtol, atol=atol), ( - f"CuTe kernel with independent per-Q-block indices deviated from masked fp32 " - f"reference: max_diff={max_diff:.3e}, mean_diff={mean_diff:.3e} " - f"(rtol={rtol}, atol={atol}, pair_mismatch={pair_mismatch})" - ) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="kernel test needs CUDA") -def test_cute_kernel_50pct_sparsity_quality_vs_dense(): - """50% sparse CuTe kernel with score-based topk should stay close to dense SDPA.""" - - from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.video_sparse_attention import ( - CUTE_AVAILABLE, - block_sparse_attn_from_indices_cute, - is_cute_supported, - ) - - if not CUTE_AVAILABLE: - pytest.skip("cuda-bindings or cutlass-dsl not importable") - - device = torch.device("cuda") - dtype = torch.bfloat16 - torch.manual_seed(0) - - B, H, num_cubes, D = 1, 4, 16, 128 - block_size = 64 - topk = num_cubes // 2 - seq_len = num_cubes * block_size - - q = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) - k = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) - v = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) - - if not is_cute_supported(q): - pytest.skip("CuTe path needs sm_100+ Blackwell (current device unsupported)") - - q_blocks = q.reshape(B, H, num_cubes, block_size, D).mean(dim=3) - k_blocks = k.reshape(B, H, num_cubes, block_size, D).mean(dim=3) - scale = D**-0.5 - block_scores = torch.einsum("bhqd,bhkd->bhqk", q_blocks.float(), k_blocks.float()) * scale - q2k_idx = block_scores.topk(topk, dim=-1).indices.to(torch.int32).contiguous() - - q2k_num = torch.full((B, H, num_cubes), topk, dtype=torch.int32, device=device) - variable_block_sizes = torch.full((num_cubes,), block_size, dtype=torch.int32, device=device) - - out_sparse, _lse = block_sparse_attn_from_indices_cute( - q, k, v, q2k_idx, q2k_num, variable_block_sizes - ) - out_dense = F.scaled_dot_product_attention(q, k, v) - - cos_sim = F.cosine_similarity( - out_sparse.float().reshape(-1), out_dense.float().reshape(-1), dim=0 - ).item() - print(f"\n 50% sparse (score-based topk) vs dense SDPA cos_sim: {cos_sim:.4f}") - - assert cos_sim >= 0.65, ( - f"50% sparse CuTe kernel deviated too far from dense SDPA: cos_sim={cos_sim:.4f} < 0.65" - ) - - -@pytest.mark.skipif(not torch.cuda.is_available(), reason="kernel test needs CUDA") -@pytest.mark.parametrize( - "num_cubes", - [1, 3, 9], - ids=["1cube_odd", "3cubes_odd", "9cubes_odd"], -) -def test_cute_kernel_odd_num_cubes_correctness(num_cubes): - """CuTe kernel with odd num_cubes must match dense SDPA (last Q-block has no pair).""" - from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.video_sparse_attention import ( - CUTE_AVAILABLE, - block_sparse_attn_from_indices_cute, - is_cute_supported, - ) - - if not CUTE_AVAILABLE: - pytest.skip("cuda-bindings or cutlass-dsl not importable") - - assert num_cubes % 2 == 1, f"pre-condition: num_cubes={num_cubes} must be odd" - - device = torch.device("cuda") - dtype = torch.bfloat16 - torch.manual_seed(0) - - B, H, D = 1, 4, 128 - block_size = 64 - seq_len = num_cubes * block_size - - q = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) - k = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) - v = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) - - if not is_cute_supported(q): - pytest.skip("CuTe path needs sm_100+ Blackwell (current device unsupported)") - - topk = num_cubes - q2k_idx = ( - torch.arange(num_cubes, device=device, dtype=torch.int32) - .view(1, 1, 1, num_cubes) - .expand(B, H, num_cubes, topk) - .contiguous() - ) - q2k_num = torch.full((B, H, num_cubes), topk, dtype=torch.int32, device=device) - variable_block_sizes = torch.full((num_cubes,), block_size, dtype=torch.int32, device=device) - - out_kernel, _lse = block_sparse_attn_from_indices_cute( - q, k, v, q2k_idx, q2k_num, variable_block_sizes - ) - out_ref = F.scaled_dot_product_attention(q, k, v) - - assert torch.isfinite(out_kernel).all(), ( - f"CuTe kernel produced non-finite output for odd num_cubes={num_cubes}" - ) - - max_diff = (out_kernel - out_ref).abs().max().item() - mean_diff = (out_kernel - out_ref).abs().mean().item() - rtol, atol = 1e-2, 1e-2 - assert torch.allclose(out_kernel, out_ref, rtol=rtol, atol=atol), ( - f"CuTe kernel deviated from dense SDPA for odd num_cubes={num_cubes}: " - f"max_diff={max_diff:.3e}, mean_diff={mean_diff:.3e} (rtol={rtol}, atol={atol})" - ) diff --git a/tests/unittest/_torch/visual_gen/test_attention_integration.py b/tests/unittest/_torch/visual_gen/test_attention_integration.py index a372db9641f2..a4c62657495a 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_integration.py +++ b/tests/unittest/_torch/visual_gen/test_attention_integration.py @@ -12,6 +12,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from utils.util import isSM100Family from tensorrt_llm._torch.modules.rms_norm import RMSNorm @@ -698,18 +699,19 @@ def test_fast_cross_attention_wan_shapes( # ============================================================================ -# VSA self-attention (CUTEDSL backend, sparse_attention_config.algorithm='vsa') +# VSA self-attention (CUTEDSL/TRTLLM backends) # ============================================================================ -def _build_vsa_setup(sparsity: float, batch_size: int, seed: int): +def _build_vsa_setup(backend: str, sparsity: float, batch_size: int, seed: int): """Build naive + integrated models, VSA metadata, and inputs for a VSA test. - latent (8,8,8) -> 512 tokens (divisible by block_size=64), head_dim=128. + A ragged latent exercises VSA padding and token-mask lowering on both + fine-stage implementations. """ - from tensorrt_llm._torch.visual_gen.attention_backend import VSAMetadataBuilder + from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import VSAMetadataBuilder - latent_shape = (8, 8, 8) + latent_shape = (9, 9, 9) seq_len = latent_shape[0] * latent_shape[1] * latent_shape[2] num_heads = 4 head_dim = 128 @@ -717,17 +719,21 @@ def _build_vsa_setup(sparsity: float, batch_size: int, seed: int): device = torch.device("cuda") dtype = torch.bfloat16 + torch.manual_seed(seed) naive = NaiveWanSelfAttention(hidden_size, num_heads, head_dim, dtype=dtype).to(device) cfg_vsa = create_model_config( - hidden_size, num_heads, head_dim, attn_backend="CUTEDSL", vsa_sparsity=sparsity + hidden_size, + num_heads, + head_dim, + attn_backend=backend, + vsa_sparsity=sparsity, ) integrated = Attention(hidden_size, num_heads, qkv_mode=QKVMode.FUSE_QKV, config=cfg_vsa).to( device ) - # Fail loudly if the VSA path silently fell back to dense (which would set - # attn_backend to "VANILLA") instead of selecting the CUTEDSL/VSA backend. - assert integrated.attn_backend == "CUTEDSL", ( - f"Expected CUTEDSL (VSA) backend, got {integrated.attn_backend!r}" + # Fail loudly if the VSA path silently fell back to the VANILLA backend. + assert integrated.attn_backend == backend, ( + f"Expected {backend} VSA backend, got {integrated.attn_backend!r}" ) copy_weights_self_attention(naive, integrated) naive.eval() @@ -754,12 +760,13 @@ def _build_vsa_setup(sparsity: float, batch_size: int, seed: int): @pytest.mark.skipif(not torch.cuda.is_available(), reason="VSA needs CUDA") -def test_vsa_self_attention_equivalence_at_sparsity_zero(): +@pytest.mark.parametrize("backend", ["CUTEDSL", "TRTLLM"]) +def test_vsa_self_attention_equivalence_at_sparsity_zero(backend: str): """VSA at sparsity=0 with G_c=0 reduces to dense attention (top_k=num_cubes, output=O_f); must match the naive SDPA reference modulo bf16 rounding.""" - from tensorrt_llm._torch.visual_gen.attention_backend import set_vsa_forward_context + from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import set_vsa_forward_context - s = _build_vsa_setup(sparsity=0.0, batch_size=2, seed=42) + s = _build_vsa_setup(backend=backend, sparsity=0.0, batch_size=2, seed=42) with torch.no_grad(): out_naive = s.naive(s.hidden_states, *s.freqs_HSD) @@ -780,22 +787,103 @@ def test_vsa_self_attention_equivalence_at_sparsity_zero(): @pytest.mark.skipif(not torch.cuda.is_available(), reason="VSA needs CUDA") -@pytest.mark.parametrize("sparsity", [0.0, 0.5], ids=["s0", "s0p5"]) -def test_vsa_self_attention_finite(sparsity: float): - """VSA forward must produce finite output (no NaN/Inf) at any supported sparsity.""" - from tensorrt_llm._torch.visual_gen.attention_backend import set_vsa_forward_context +@pytest.mark.skipif( + not isSM100Family(), + reason="CuTe DSL and PrimTS block-sparse parity requires SM100 or SM103", +) +def test_vsa_sparse_backends_match_on_ragged_input(): + """CuTeDSL and TRTLLM implement the same sparse VSA fine-stage semantics.""" + from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import set_vsa_forward_context + + sparsity = 0.5 + setups = { + backend: _build_vsa_setup(backend=backend, sparsity=sparsity, batch_size=1, seed=0) + for backend in ("CUTEDSL", "TRTLLM") + } + sparse_fine_executed = {} + for backend, setup in setups.items(): + if backend == "CUTEDSL": + original_execute = setup.integrated.attn._execute_sparse_fine + + def checked_execute(*args, _original=original_execute, **kwargs): + result = _original(*args, **kwargs) + sparse_fine_executed["CUTEDSL"] = True + return result + + setup.integrated.attn._execute_sparse_fine = checked_execute + else: + original_predict = setup.integrated.attn.block_sparse_attn_predict + + def checked_predict(*args, _original=original_predict, **kwargs): + result = _original(*args, **kwargs) + sparse_fine_executed["TRTLLM"] = result is not None + return result + + setup.integrated.attn.block_sparse_attn_predict = checked_predict + + outputs = {} + for backend, setup in setups.items(): + with torch.no_grad(), set_vsa_forward_context(setup.metadata): + outputs[backend] = setup.integrated( + setup.hidden_states, + freqs=setup.freqs_SHD, + gate_compress=setup.gate_compress_zero, + ) - s = _build_vsa_setup(sparsity=sparsity, batch_size=1, seed=0) + assert sparse_fine_executed == {"CUTEDSL": True, "TRTLLM": True} + assert torch.isfinite(outputs["CUTEDSL"]).all() + assert torch.isfinite(outputs["TRTLLM"]).all() + torch.testing.assert_close( + outputs["CUTEDSL"], + outputs["TRTLLM"], + rtol=1e-2, + atol=1e-2, + ) - with torch.no_grad(), set_vsa_forward_context(s.metadata): - out = s.integrated(s.hidden_states, freqs=s.freqs_SHD, gate_compress=s.gate_compress_zero) - assert out.shape == s.hidden_states.shape - nan_count = torch.isnan(out).sum().item() - inf_count = torch.isinf(out).sum().item() - assert nan_count == 0 and inf_count == 0, ( - f"VSA produced non-finite output at sparsity={sparsity}: NaN={nan_count}, Inf={inf_count}" - ) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="VSA needs CUDA") +@pytest.mark.skipif( + not isSM100Family(), + reason="PrimTS block-sparse CUDA Graph replay requires SM100 or SM103", +) +def test_vsa_trtllm_cuda_graph_replays_live_routes(): + """Captured VSA recomputes routes when graph-stable Q/K/V storage changes.""" + from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import set_vsa_forward_context + + setup = _build_vsa_setup(backend="TRTLLM", sparsity=0.5, batch_size=1, seed=17) + static_hidden = setup.hidden_states.clone() + static_gate = setup.gate_compress_zero.clone() + + for _ in range(2): + with torch.no_grad(), set_vsa_forward_context(setup.metadata): + setup.integrated( + static_hidden, + freqs=setup.freqs_SHD, + gate_compress=static_gate, + ) + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph), torch.no_grad(), set_vsa_forward_context(setup.metadata): + graph_output = setup.integrated( + static_hidden, + freqs=setup.freqs_SHD, + gate_compress=static_gate, + ) + + initial_output = graph_output.clone() + live_hidden = torch.randn_like(static_hidden) + static_hidden.copy_(live_hidden) + graph.replay() + replay_output = graph_output.clone() + with torch.no_grad(), set_vsa_forward_context(setup.metadata): + eager_output = setup.integrated( + live_hidden, + freqs=setup.freqs_SHD, + gate_compress=static_gate, + ) + + assert not torch.equal(initial_output, replay_output) + torch.testing.assert_close(replay_output, eager_output, rtol=1e-2, atol=1e-2) def test_trtllm_cached_prepare(): diff --git a/tests/unittest/_torch/visual_gen/test_attention_perf.py b/tests/unittest/_torch/visual_gen/test_attention_perf.py index a662b788cf0d..855366914932 100644 --- a/tests/unittest/_torch/visual_gen/test_attention_perf.py +++ b/tests/unittest/_torch/visual_gen/test_attention_perf.py @@ -36,11 +36,6 @@ import pytest import torch -from tensorrt_llm._torch.visual_gen.attention_backend import ( - VSAMetadataBuilder, - set_vsa_forward_context, -) - # ============================================================================ # Flash Attention 4 availability # ============================================================================ @@ -49,6 +44,10 @@ from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import ( _flash_attn_fwd_import_error as _fa4_import_error, ) +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import ( + VSAMetadataBuilder, + set_vsa_forward_context, +) from tensorrt_llm._torch.visual_gen.config import ( DiffusionModelConfig, create_attention_metadata_state, @@ -1080,7 +1079,9 @@ def test_vsa_kernel_vs_fa4( block_size: int, sparsity: float, ): - from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl import VSA_KERNEL_MAX_CUBES + from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa.backend import ( + VSA_KERNEL_MAX_CUBES, + ) assert seq_len % block_size == 0, "seq_len must be a multiple of block_size" num_cubes = seq_len // block_size diff --git a/tests/unittest/_torch/visual_gen/test_attention_vsa.py b/tests/unittest/_torch/visual_gen/test_attention_vsa.py new file mode 100644 index 000000000000..d2cb2c663390 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_attention_vsa.py @@ -0,0 +1,1023 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""VSA correctness tests: backend dispatch, preprocessing, and kernel behavior. + +Module-level dense-equivalence and finite-output checks live in +test_attention_integration.py. +""" + +from types import SimpleNamespace + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask +from tensorrt_llm._torch.attention.backends.sparse.params import BlockSparseForwardInputs +from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl import CuTeDSLAttention +from tensorrt_llm._torch.visual_gen.attention_backend.interface import AttentionTensorLayout +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import backend as vsa_backend +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa 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 test_vsa_trtllm_overrides_only_forward_around_the_core() -> None: + assert "forward" in VSATrtllmAttention.__dict__ + for name in ( + "block_sparse_attn_predict", + "sparse_predict", + "sparse_post_process", + "_enable_sparse_workflow", + ): + assert name not in VSATrtllmAttention.__dict__ + + +def _capture_wrapper_forward(monkeypatch: pytest.MonkeyPatch) -> dict: + """Replace the VisualGen wrapper forward with a recorder returning the fine input.""" + + captured = {} + + def _forward( + self, + q, + k, + v, + batch_size, + seq_len, + attention_mask=PredefinedAttentionMask.FULL, + seq_len_kv=None, + sparse_backend_args=None, + **kwargs, + ): + captured.update( + q=q, + k=k, + v=v, + batch_size=batch_size, + seq_len=seq_len, + attention_mask=attention_mask, + seq_len_kv=seq_len_kv, + sparse_backend_args=sparse_backend_args, + kwargs=kwargs, + ) + return q.reshape(batch_size, seq_len, -1) + + monkeypatch.setattr(TrtllmAttention, "forward", _forward) + return captured + + +def test_vsa_backends_share_one_predictor_implementation() -> None: + trtllm_attention = object.__new__(VSATrtllmAttention) + cute_attention = object.__new__(VSACuTeDSLAttention) + trtllm_attention.predictor = VSAPredictor(num_heads=1) + cute_attention.predictor = VSAPredictor(num_heads=1) + + assert type(trtllm_attention.predictor) is type(cute_attention.predictor) is VSAPredictor + assert set(vsa_backend.__all__) >= {"VSATrtllmAttention", "VSACuTeDSLAttention"} + + +def test_vsa_trtllm_layers_share_model_scoped_predictor() -> None: + attention_metadata_state = create_attention_metadata_state() + first = VSATrtllmAttention( + layer_idx=0, + num_heads=2, + num_kv_heads=2, + head_dim=128, + attention_metadata_state=attention_metadata_state, + ) + second = VSATrtllmAttention( + layer_idx=1, + num_heads=2, + num_kv_heads=2, + head_dim=128, + attention_metadata_state=attention_metadata_state, + ) + + assert first.predictor is second.predictor + + +def test_vsa_predictor_produces_sorted_block_inputs_and_effective_tiled_qkv() -> None: + predictor = VSAPredictor(num_heads=1) + metadata = _make_vsa_metadata() + q = torch.randn(1, 80, 1, 8) + + inputs = predictor.predict( + q, + q, + q, + batch_size=1, + seq_len=80, + seq_len_kv=80, + attention_mask=PredefinedAttentionMask.FULL, + gate_compress=torch.zeros_like(q), + gate_fine=None, + use_sparse_fine=True, + produce_block_sparse_inputs=True, + metadata=metadata, + ) + + assert isinstance(inputs, VSAForwardInputs) + assert inputs.q.shape == inputs.k.shape == inputs.v.shape == (1, 128, 1, 8) + assert inputs.seq_len == 128 + block_sparse_inputs = inputs.block_sparse_inputs + assert isinstance(block_sparse_inputs, BlockSparseForwardInputs) + assert block_sparse_inputs.block_indptr.tolist() == [[[0, 2, 4]]] + assert block_sparse_inputs.block_indices.tolist() == [0, 1, 0, 1] + assert block_sparse_inputs.kv_valid_bits.dtype == torch.uint32 + assert block_sparse_inputs.kv_valid_bits.tolist() == [[0xFFFFFFFF, 0xFFFFFFFF, 0xFFFF, 0]] + + +def test_vsa_predictor_dense_fallback_keeps_compact_qkv_and_no_block_inputs() -> None: + predictor = VSAPredictor(num_heads=1) + metadata = _make_vsa_metadata(sparsity=0.5) + q = torch.randn(1, 80, 1, 8) + k = torch.randn_like(q) + v = torch.randn_like(q) + + inputs = predictor.predict( + q, + k, + v, + batch_size=1, + seq_len=80, + seq_len_kv=80, + attention_mask=PredefinedAttentionMask.FULL, + gate_compress=torch.zeros_like(q), + gate_fine=None, + use_sparse_fine=False, + produce_block_sparse_inputs=False, + metadata=metadata, + ) + + assert inputs.q is q + assert inputs.k is k + assert inputs.v is v + assert inputs.seq_len == 80 + assert inputs.block_sparse_inputs is None + assert 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 test_trtllm_vsa_consumes_gates_and_forwards_only_timestep( + monkeypatch: pytest.MonkeyPatch, +) -> None: + captured = _capture_wrapper_forward(monkeypatch) + attention = _make_dense_fallback_vsa_attention() + q = torch.randn(1, 80, 1, 8) + gate_compress = torch.full_like(q, 2.0) + gate_fine = torch.full_like(q, 0.5) + timestep = torch.tensor([12]) + + with set_vsa_forward_context(_make_vsa_metadata(sparsity=0.5)): + output = attention.forward( + q, + q, + q, + batch_size=1, + seq_len=80, + gate_compress=gate_compress, + gate_fine=gate_fine, + timestep=timestep, + ) + + assert captured["kwargs"] == {"timestep": timestep} + assert output.shape == (1, 80, 8) + assert torch.isfinite(output).all() + + +def test_cutedsl_vsa_rejects_unexpected_forward_kwargs( + monkeypatch: pytest.MonkeyPatch, +) -> None: + attention = object.__new__(VSACuTeDSLAttention) + attention.predictor = VSAPredictor(num_heads=1) + q = torch.randn(1, 80, 1, 8) + monkeypatch.setattr(vsa_backend, "_vsa_import_error", RuntimeError("disabled for test")) + + with set_vsa_forward_context(_make_vsa_metadata(sparsity=0.5)): + with pytest.raises(TypeError, match="gate_fnne"): + attention.forward( + q, + q, + q, + gate_compress=torch.zeros_like(q), + gate_fnne=torch.zeros_like(q), + ) + + +def _make_config( + hidden_size: int, + num_heads: int, + head_dim: int, + backend: str, + vsa_sparsity: "float | None" = None, +) -> DiffusionModelConfig: + """Minimal DiffusionModelConfig for one Attention module.""" + pretrained_config = SimpleNamespace( + hidden_size=hidden_size, + num_attention_heads=num_heads, + attention_head_dim=head_dim, + eps=1e-6, + ) + sparse_attention_config = ( + VideoSparseAttentionConfig(vsa_sparsity=vsa_sparsity) if vsa_sparsity is not None else None + ) + config = DiffusionModelConfig( + pretrained_config=pretrained_config, + attention=AttentionConfig(backend=backend, sparse_attention_config=sparse_attention_config), + skip_create_weights_in_init=False, + ) + config.attention_metadata_state = ( + create_attention_metadata_state() if backend == "TRTLLM" else None + ) + return config + + +@pytest.mark.parametrize("backend", ["CUTEDSL", "TRTLLM"]) +@pytest.mark.parametrize( + ("is_self_attention", "expected_backend"), + [(False, "VANILLA"), (True, None)], + ids=["cross", "self"], +) +def test_vsa_separate_qkv_dispatches_by_attention_role( + monkeypatch: pytest.MonkeyPatch, + backend: str, + is_self_attention: bool, + expected_backend: str | None, +) -> None: + monkeypatch.setattr( + attention_module, + "create_attention", + lambda *, backend, **kwargs: SimpleNamespace(backend=backend, kwargs=kwargs), + ) + cfg = _make_config( + hidden_size=64, + num_heads=4, + head_dim=16, + backend=backend, + vsa_sparsity=0.5, + ) + attention = Attention( + 64, + 4, + qkv_mode=QKVMode.SEPARATE_QKV, + config=cfg, + separate_qkv_is_self_attention=is_self_attention, + ) + + assert attention.attn_backend == (expected_backend or backend) + + +def test_plain_trtllm_separate_qkv_self_attention_keeps_vanilla_fallback( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + attention_module, + "create_attention", + lambda *, backend, **kwargs: SimpleNamespace(backend=backend, kwargs=kwargs), + ) + cfg = _make_config( + hidden_size=64, + num_heads=4, + head_dim=16, + backend="TRTLLM", + ) + + attention = Attention( + 64, + 4, + qkv_mode=QKVMode.SEPARATE_QKV, + config=cfg, + separate_qkv_is_self_attention=True, + ) + + assert attention.attn_backend == "VANILLA" + + +def test_vsa_with_attn2d_raises(): + """VSA + Attention2D must error at construction (VSA needs the full sequence per rank).""" + pretrained_config = SimpleNamespace( + hidden_size=64, + num_attention_heads=4, + attention_head_dim=16, + eps=1e-6, + ) + cfg = DiffusionModelConfig( + pretrained_config=pretrained_config, + attention=AttentionConfig( + backend="CUTEDSL", + sparse_attention_config=VideoSparseAttentionConfig(vsa_sparsity=0.0), + ), + skip_create_weights_in_init=False, + ) + cfg.visual_gen_mapping = SimpleNamespace( + ring_size=1, + ring_group=None, + ulysses_size=1, + ulysses_group=None, + attn2d_row_size=2, + attn2d_col_size=2, + attn2d_row_group=None, + attn2d_col_group=None, + cp_size=4, + ) + with pytest.raises(ValueError, match="incompatible with context parallelism"): + Attention(64, 4, qkv_mode=QKVMode.FUSE_QKV, config=cfg) + + +def test_vsa_metadata_builder_reuses_shape_tensors_with_live_step_policy() -> None: + builder = VSAMetadataBuilder() + build_args = { + "raw_latent_shape": (9, 9, 9), + "patch_size": (1, 1, 1), + "device": torch.device("cpu"), + } + + first = builder.build(current_timestep=3, vsa_sparsity=0.25, **build_args) + second = builder.build(current_timestep=4, vsa_sparsity=0.75, **build_args) + + assert first is not second + assert (first.current_timestep, first.vsa_sparsity) == (3, 0.25) + assert (second.current_timestep, second.vsa_sparsity) == (4, 0.75) + assert second.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_bound_shape_profiles() -> None: + builder = VSAMetadataBuilder(max_cached_shapes=1) + build_args = { + "current_timestep": 0, + "patch_size": (1, 1, 1), + "vsa_sparsity": 0.5, + "device": torch.device("cpu"), + } + builder.build(raw_latent_shape=(4, 4, 4), **build_args) + with pytest.raises(RuntimeError, match="metadata cache reached its 1-shape limit"): + builder.build(raw_latent_shape=(8, 4, 4), **build_args) + + route_builder = VSAPredictor(num_heads=1, max_cached_shapes=1)._route_builder + kv_valid_words = torch.ones((1,), dtype=torch.uint32) + route_builder.from_selected_blocks(torch.zeros((1, 1, 1, 1), dtype=torch.int32), kv_valid_words) + with pytest.raises(RuntimeError, match="route cache reached its 1-shape limit"): + route_builder.from_selected_blocks( + torch.zeros((1, 1, 2, 1), dtype=torch.int32), kv_valid_words + ) + + +@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): + """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) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="kernel test needs CUDA") +def test_cute_kernel_matches_dense_at_full_topk(): + """CuTe block-sparse kernel matches dense SDPA when every cube is selected.""" + from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.video_sparse_attention import ( + CUTE_AVAILABLE, + block_sparse_attn_from_indices_cute, + is_cute_supported, + ) + + if not CUTE_AVAILABLE: + pytest.skip("cuda-bindings or cutlass-dsl not importable") + + device = torch.device("cuda") + dtype = torch.bfloat16 + torch.manual_seed(0) + + B, H, num_cubes, D = 1, 4, 4, 128 + block_size = 64 + seq_len = num_cubes * block_size + + q = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) + k = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) + v = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) + + if not is_cute_supported(q): + pytest.skip("CuTe path needs sm_100+ Blackwell (current device unsupported)") + + topk = num_cubes + q2k_idx = ( + torch.arange(num_cubes, device=device, dtype=torch.int32) + .view(1, 1, 1, num_cubes) + .expand(B, H, num_cubes, topk) + .contiguous() + ) + q2k_num = torch.full((B, H, num_cubes), topk, dtype=torch.int32, device=device) + variable_block_sizes = torch.full((num_cubes,), block_size, dtype=torch.int32, device=device) + + out_kernel, _lse = block_sparse_attn_from_indices_cute( + q, k, v, q2k_idx, q2k_num, variable_block_sizes + ) + out_ref = F.scaled_dot_product_attention(q, k, v) + + max_diff = (out_kernel - out_ref).abs().max().item() + mean_diff = (out_kernel - out_ref).abs().mean().item() + + rtol, atol = 1e-2, 1e-2 + assert torch.allclose(out_kernel, out_ref, rtol=rtol, atol=atol), ( + f"CuTe block-sparse kernel deviates from dense SDPA at full top-K: " + f"max_diff={max_diff:.3e}, mean_diff={mean_diff:.3e} (rtol={rtol}, atol={atol})" + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="kernel test needs CUDA") +def test_cute_kernel_matches_ref_with_independent_indices(): + """CuTe kernel: paired Q-blocks (2i, 2i+1) attend to independent KV index lists.""" + from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.video_sparse_attention import ( + CUTE_AVAILABLE, + block_sparse_attn_from_indices_cute, + is_cute_supported, + ) + + if not CUTE_AVAILABLE: + pytest.skip("cuda-bindings or cutlass-dsl not importable") + + device = torch.device("cuda") + dtype = torch.bfloat16 + torch.manual_seed(42) + + B, H, num_cubes, D = 2, 4, 16, 128 + block_size = 64 + topk = num_cubes // 2 + seq_len = num_cubes * block_size + + q = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) + k = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) + v = torch.randn(B, H, seq_len, D, device=device, dtype=dtype) + + if not is_cute_supported(q): + pytest.skip("CuTe path needs sm_100+ Blackwell (current device unsupported)") + + q2k_idx = ( + torch.stack( + [ + torch.randperm(num_cubes, device=device, dtype=torch.int32)[:topk] + for _ in range(B * H * num_cubes) + ] + ) + .view(B, H, num_cubes, topk) + .contiguous() + ) + + paired = q2k_idx.view(B, H, num_cubes // 2, 2, topk).sort(dim=-1).values + pair_mismatch = (paired[..., 0, :] != paired[..., 1, :]).sum().item() + assert pair_mismatch > 0, ( + "Pre-condition failed: random permutations matched across every pair; " + "re-seed or raise num_cubes." + ) + + q2k_num = torch.full((B, H, num_cubes), topk, dtype=torch.int32, device=device) + variable_block_sizes = torch.full((num_cubes,), block_size, dtype=torch.int32, device=device) + + attn_mask = torch.full( + (B, H, seq_len, seq_len), float("-inf"), device=device, dtype=torch.float32 + ) + for b in range(B): + for h in range(H): + for q_blk in range(num_cubes): + for ki in range(topk): + k_blk = q2k_idx[b, h, q_blk, ki].item() + qs = q_blk * block_size + ks = k_blk * block_size + attn_mask[b, h, qs : qs + block_size, ks : ks + block_size] = 0.0 + + out_kernel, _lse = block_sparse_attn_from_indices_cute( + q, k, v, q2k_idx, q2k_num, variable_block_sizes + ) + + scale = 1.0 / (D**0.5) + scores = (q.float() @ k.float().transpose(-2, -1)) * scale + scores = scores + attn_mask + probs = torch.softmax(scores, dim=-1) + out_ref = (probs @ v.float()).to(dtype) + + abs_diff = (out_kernel.float() - out_ref.float()).abs() + max_diff = abs_diff.max().item() + mean_diff = abs_diff.mean().item() + + rtol, atol = 1e-2, 1e-2 + assert torch.allclose(out_kernel, out_ref, rtol=rtol, atol=atol), ( + f"CuTe kernel with independent per-Q-block indices deviated from masked fp32 " + f"reference: max_diff={max_diff:.3e}, mean_diff={mean_diff:.3e} " + f"(rtol={rtol}, atol={atol}, pair_mismatch={pair_mismatch})" + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="kernel test needs CUDA") +def test_cute_kernel_50pct_sparsity_quality_vs_dense(): + """50% sparse CuTe kernel with score-based topk stays close to dense SDPA.""" + from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.video_sparse_attention import ( + CUTE_AVAILABLE, + block_sparse_attn_from_indices_cute, + is_cute_supported, + ) + + if not CUTE_AVAILABLE: + pytest.skip("cuda-bindings or cutlass-dsl not importable") + + device = torch.device("cuda") + dtype = torch.bfloat16 + torch.manual_seed(0) + + batch_size, num_heads, num_cubes, head_dim = 1, 4, 16, 128 + block_size = 64 + topk = num_cubes // 2 + seq_len = num_cubes * block_size + + q = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype) + k = torch.randn_like(q) + v = torch.randn_like(q) + + if not is_cute_supported(q): + pytest.skip("CuTe path needs sm_100+ Blackwell (current device unsupported)") + + q_blocks = q.reshape(batch_size, num_heads, num_cubes, block_size, head_dim).mean(dim=3) + k_blocks = k.reshape(batch_size, num_heads, num_cubes, block_size, head_dim).mean(dim=3) + block_scores = torch.einsum( + "bhqd,bhkd->bhqk", + q_blocks.float(), + k_blocks.float(), + ) * (head_dim**-0.5) + q2k_idx = block_scores.topk(topk, dim=-1).indices.to(torch.int32).contiguous() + q2k_num = torch.full( + (batch_size, num_heads, num_cubes), + topk, + dtype=torch.int32, + device=device, + ) + variable_block_sizes = torch.full( + (num_cubes,), + block_size, + dtype=torch.int32, + device=device, + ) + + out_sparse, _lse = block_sparse_attn_from_indices_cute( + q, + k, + v, + q2k_idx, + q2k_num, + variable_block_sizes, + ) + out_dense = F.scaled_dot_product_attention(q, k, v) + + cos_sim = F.cosine_similarity( + out_sparse.float().reshape(-1), + out_dense.float().reshape(-1), + dim=0, + ).item() + assert cos_sim >= 0.65, ( + f"50% sparse CuTe kernel deviated too far from dense SDPA: cos_sim={cos_sim:.4f} < 0.65" + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="kernel test needs CUDA") +@pytest.mark.parametrize( + "num_cubes", + [1, 3, 9], + ids=["1cube_odd", "3cubes_odd", "9cubes_odd"], +) +def test_cute_kernel_odd_num_cubes_correctness(num_cubes): + """CuTe kernel supports a final Q block that has no paired neighbor.""" + from tensorrt_llm._torch.visual_gen.cute_dsl_kernels.blackwell.video_sparse_attention import ( + CUTE_AVAILABLE, + block_sparse_attn_from_indices_cute, + is_cute_supported, + ) + + if not CUTE_AVAILABLE: + pytest.skip("cuda-bindings or cutlass-dsl not importable") + + assert num_cubes % 2 == 1 + device = torch.device("cuda") + dtype = torch.bfloat16 + torch.manual_seed(0) + + batch_size, num_heads, head_dim = 1, 4, 128 + block_size = 64 + seq_len = num_cubes * block_size + q = torch.randn(batch_size, num_heads, seq_len, head_dim, device=device, dtype=dtype) + k = torch.randn_like(q) + v = torch.randn_like(q) + + if not is_cute_supported(q): + pytest.skip("CuTe path needs sm_100+ Blackwell (current device unsupported)") + + q2k_idx = ( + torch.arange(num_cubes, device=device, dtype=torch.int32) + .view(1, 1, 1, num_cubes) + .expand(batch_size, num_heads, num_cubes, num_cubes) + .contiguous() + ) + q2k_num = torch.full( + (batch_size, num_heads, num_cubes), + num_cubes, + dtype=torch.int32, + device=device, + ) + variable_block_sizes = torch.full( + (num_cubes,), + block_size, + dtype=torch.int32, + device=device, + ) + + out_kernel, _lse = block_sparse_attn_from_indices_cute( + q, + k, + v, + q2k_idx, + q2k_num, + variable_block_sizes, + ) + out_ref = F.scaled_dot_product_attention(q, k, v) + + assert torch.isfinite(out_kernel).all() + torch.testing.assert_close(out_kernel, out_ref, rtol=1e-2, atol=1e-2) diff --git a/tests/unittest/_torch/visual_gen/test_attention_vsa_kernels.py b/tests/unittest/_torch/visual_gen/test_attention_vsa_kernels.py new file mode 100644 index 000000000000..c63df0587b19 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_attention_vsa_kernels.py @@ -0,0 +1,195 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""VSA predictor kernels (tile + cube mean, row sort, coarse/fine blend) against references. + +Every kernel is checked on CUDA (Triton path) and CPU (PyTorch fallback) over cube layouts +with ragged fills and head layouts whose row width is not a power of two. +""" + +import pytest +import torch + +from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa.kernels import ( + blend_coarse_fine, + sort_last_dim, + tile_and_pool_cubes, +) + +CUBE_SIZE = 64 + +_DEVICES = ["cpu"] + (["cuda"] if torch.cuda.is_available() else []) + +_HEAD_LAYOUTS = [ + pytest.param(6, 4, 32, id="row128"), + pytest.param(5, 5, 21, id="row105_odd"), + pytest.param(30, 40, 128, id="wan14b_row5120"), + pytest.param(3, 64, 128, id="row8192"), +] + + +def _random_cube_layout(num_cubes: int, device: str, generator: torch.Generator): + """Assign consecutive compact tokens to cubes with a random fill per cube. + + Returns the padded-slot source index (-1 for pad slots), the valid count per cube, the + padded slot of every compact token, and the compact sequence length. + """ + counts = torch.randint(1, CUBE_SIZE + 1, (num_cubes,), generator=generator) + counts[0] = CUBE_SIZE + counts[-1] = CUBE_SIZE // 2 + seq_len = int(counts.sum()) + tile_source_index = torch.full((num_cubes * CUBE_SIZE,), -1, dtype=torch.long) + untile_index = torch.empty(seq_len, dtype=torch.long) + token = 0 + for cube, count in enumerate(counts.tolist()): + slots = torch.arange(cube * CUBE_SIZE, cube * CUBE_SIZE + count) + tile_source_index[slots] = torch.arange(token, token + count) + untile_index[token : token + count] = slots + token += count + return tile_source_index.to(device), counts.to(device), untile_index.to(device), seq_len + + +def _reference_tile_and_pool(x, tile_source_index, counts, num_cubes): + batch, _, heads, head_dim = x.shape + tiled = x.index_select(1, tile_source_index.clamp(min=0)) + valid = (tile_source_index >= 0).view(1, -1, 1, 1) + tiled = torch.where(valid, tiled, torch.zeros((), dtype=x.dtype, device=x.device)) + pooled = tiled.view(batch, num_cubes, CUBE_SIZE, heads, head_dim).float().sum(dim=2) + pooled = pooled / counts.view(1, -1, 1, 1).float() + return tiled, pooled.to(x.dtype) + + +def _reference_blend(fine, coarse, gate_compress, gate_fine, untile_index, fine_is_tiled): + coarse_per_token = coarse.index_select(1, untile_index // CUBE_SIZE) + fine_compact = fine.index_select(1, untile_index) if fine_is_tiled else fine + if gate_fine is not None: + fine_compact = gate_fine * fine_compact + return gate_compress * coarse_per_token + fine_compact + + +@pytest.mark.parametrize("device", _DEVICES) +@pytest.mark.parametrize(("num_cubes", "heads", "head_dim"), _HEAD_LAYOUTS) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32], ids=["bf16", "fp32"]) +def test_tile_and_pool_cubes_matches_reference(device, num_cubes, heads, head_dim, dtype): + generator = torch.Generator().manual_seed(num_cubes * 31 + heads) + source, counts, _untile, seq_len = _random_cube_layout(num_cubes, device, generator) + x = torch.randn(2, seq_len, heads, head_dim, device=device, dtype=dtype) + + tiled, pooled = tile_and_pool_cubes(x, source, counts, cube_size=CUBE_SIZE) + + ref_tiled, ref_pooled = _reference_tile_and_pool(x, source, counts, num_cubes) + assert tiled.shape == (2, num_cubes * CUBE_SIZE, heads, head_dim) + assert pooled.shape == (2, num_cubes, heads, head_dim) + assert torch.equal(tiled, ref_tiled) + tolerance = 1e-5 if dtype == torch.float32 else 1e-2 + torch.testing.assert_close(pooled, ref_pooled, rtol=tolerance, atol=tolerance) + + +@pytest.mark.parametrize("device", _DEVICES) +def test_tile_and_pool_cubes_accepts_strided_sequence_layout(device): + """Q/K/V unbound from a packed [B, S, 3, H, D] tensor must work without a copy.""" + source, counts, _untile, seq_len = _random_cube_layout( + 8, device, torch.Generator().manual_seed(7) + ) + packed = torch.randn(2, seq_len, 3, 4, 32, device=device, dtype=torch.float16) + q = packed[:, :, 1] + assert not q.is_contiguous() + + tiled, pooled = tile_and_pool_cubes(q, source, counts, cube_size=CUBE_SIZE) + + ref_tiled, ref_pooled = _reference_tile_and_pool(q.contiguous(), source, counts, 8) + assert torch.equal(tiled, ref_tiled) + torch.testing.assert_close(pooled, ref_pooled, rtol=1e-3, atol=1e-3) + + +def test_tile_and_pool_cubes_rejects_split_head_dims(): + source, counts, _untile, seq_len = _random_cube_layout( + 2, "cpu", torch.Generator().manual_seed(1) + ) + x = torch.randn(1, seq_len, 32, 4).transpose(2, 3) + + with pytest.raises(ValueError, match="contiguous"): + tile_and_pool_cubes(x, source, counts, cube_size=CUBE_SIZE) + + +@pytest.mark.parametrize("device", _DEVICES) +@pytest.mark.parametrize("row_length", [1, 30, 144, 257, 2048, 5000]) +def test_sort_last_dim_matches_torch_sort(device, row_length): + generator = torch.Generator().manual_seed(row_length) + values = torch.randint(0, 4096, (2, 3, 5, row_length), generator=generator, dtype=torch.int32) + values = values.to(device) + + assert torch.equal(sort_last_dim(values), torch.sort(values, dim=-1).values) + + +@pytest.mark.parametrize("device", _DEVICES) +def test_blend_coarse_fine_reads_head_major_fine_output(device): + """Fine output stored as [B, H, S, D] (the CuTe layout) is consumed through its strides.""" + num_cubes, batch, heads, head_dim = 6, 2, 4, 32 + source, _counts, untile, seq_len = _random_cube_layout( + num_cubes, device, torch.Generator().manual_seed(11) + ) + fine = torch.randn(batch, heads, num_cubes * CUBE_SIZE, head_dim, device=device).transpose(1, 2) + assert not fine.is_contiguous() + coarse = torch.randn(batch, num_cubes, heads, head_dim, device=device) + gate_compress = torch.randn(batch, seq_len, heads, head_dim, device=device) + + out = blend_coarse_fine( + fine, coarse, gate_compress, None, untile, cube_size=CUBE_SIZE, fine_is_tiled=True + ) + + ref = _reference_blend(fine, coarse, gate_compress, None, untile, True) + torch.testing.assert_close(out, ref, rtol=1e-6, atol=1e-6) + + +def test_sort_last_dim_requires_int32(): + with pytest.raises(TypeError, match="int32"): + sort_last_dim(torch.zeros(2, 4, dtype=torch.int64)) + + +@pytest.mark.parametrize("device", _DEVICES) +@pytest.mark.parametrize("fine_is_tiled", [True, False], ids=["tiled_fine", "compact_fine"]) +@pytest.mark.parametrize("with_gate_fine", [True, False], ids=["gate_fine", "no_gate_fine"]) +@pytest.mark.parametrize( + ("heads", "head_dim", "dtype"), + [(5, 21, torch.float32), (40, 128, torch.bfloat16)], + ids=["row105_fp32", "wan14b_bf16"], +) +def test_blend_coarse_fine_matches_reference( + device, fine_is_tiled, with_gate_fine, heads, head_dim, dtype +): + num_cubes, batch = 6, 2 + source, _counts, untile, seq_len = _random_cube_layout( + num_cubes, device, torch.Generator().manual_seed(3) + ) + fine_len = num_cubes * CUBE_SIZE if fine_is_tiled else seq_len + fine = torch.randn(batch, fine_len, heads, head_dim, device=device, dtype=dtype) + coarse = torch.randn(batch, num_cubes, heads, head_dim, device=device, dtype=dtype) + gate_compress = torch.randn(batch, seq_len, heads, head_dim, device=device, dtype=dtype) + gate_fine = torch.randn_like(gate_compress) if with_gate_fine else None + + out = blend_coarse_fine( + fine, + coarse, + gate_compress, + gate_fine, + untile, + cube_size=CUBE_SIZE, + fine_is_tiled=fine_is_tiled, + ) + + ref = _reference_blend(fine, coarse, gate_compress, gate_fine, untile, fine_is_tiled) + assert out.shape == gate_compress.shape + tolerance = 1e-6 if dtype == torch.float32 else 2e-2 + torch.testing.assert_close(out, ref, rtol=tolerance, atol=tolerance) diff --git a/tests/unittest/_torch/visual_gen/test_fa4_cutlass_compatibility.py b/tests/unittest/_torch/visual_gen/test_fa4_cutlass_compatibility.py index eb9ea6941792..e1203891be26 100644 --- a/tests/unittest/_torch/visual_gen/test_fa4_cutlass_compatibility.py +++ b/tests/unittest/_torch/visual_gen/test_fa4_cutlass_compatibility.py @@ -1,6 +1,9 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import subprocess +import sys +import textwrap from importlib import import_module import pytest @@ -10,6 +13,7 @@ from tensorrt_llm._torch.visual_gen.attention_backend import flash_attn4, parallel # noqa: E402 from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import ( # noqa: E402 _install_cutlass_dsl_compatibility, + _install_flash_attn_tile_scheduler_compatibility, ) @@ -48,3 +52,54 @@ def test_cutlass_dsl_47_aliases_allow_fa4_interface_import() -> None: assert callable(interface.flash_attn_combine) assert callable(flash_attn4._flash_attn_fwd) assert callable(parallel._flash_attn_combine) + + +def test_fa4_work_tile_info_survives_cutlass_task_scheduling_import() -> None: + task_scheduling = pytest.importorskip("cutlass.experimental.task_scheduling") + tile_scheduler = pytest.importorskip("flash_attn.cute.tile_scheduler") + import cutlass + from cutlass.cutlass_dsl import Boolean + from cutlass.utils.static_persistent_tile_scheduler import WorkTileInfo as CutlassWorkTileInfo + + del task_scheduling + _install_flash_attn_tile_scheduler_compatibility() + tile_idx = (cutlass.Int32(1), cutlass.Int32(2), cutlass.Int32(3), cutlass.Int32(0)) + + # The task-scheduling import rewrote the shared CUTLASS class to three scalars. + with pytest.raises(ValueError, match="too many values to unpack"): + CutlassWorkTileInfo(tile_idx, Boolean(True)) + + fa4_tile = tile_scheduler.WorkTileInfo(tile_idx, Boolean(True)) + + assert "__init__" in vars(tile_scheduler.WorkTileInfo) + assert fa4_tile.tile_idx == tile_idx + assert bool(fa4_tile.is_valid_tile) + assert issubclass(tile_scheduler.WorkTileInfo, CutlassWorkTileInfo) + + +def test_fa4_work_tile_info_survives_task_scheduling_imported_first() -> None: + pytest.importorskip("cutlass.experimental.task_scheduling") + pytest.importorskip("flash_attn.cute.tile_scheduler") + script = textwrap.dedent( + """ + import cutlass + import cutlass.experimental.task_scheduling # noqa: F401 + from cutlass.cutlass_dsl import Boolean + + import tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 # noqa: F401 + from flash_attn.cute.tile_scheduler import WorkTileInfo + + tile_idx = (cutlass.Int32(1), cutlass.Int32(2), cutlass.Int32(3), cutlass.Int32(0)) + tile = WorkTileInfo(tile_idx, Boolean(True)) + assert tile.tile_idx == tile_idx + assert bool(tile.is_valid_tile) + print("fa4-work-tile-info-ok") + """ + ) + + result = subprocess.run( + [sys.executable, "-c", script], capture_output=True, text=True, timeout=900, check=False + ) + + assert result.returncode == 0, result.stderr[-4000:] + assert result.stdout.strip().endswith("fa4-work-tile-info-ok") diff --git a/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py b/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py index f1064f780fa1..057fafc91520 100644 --- a/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py @@ -1284,6 +1284,9 @@ def __init__(self): self.active_topology = "default" self.device = "cuda" + def register_cuda_graph_extra_key_fns(self, runner): + del runner + def forward(self, x): return self.lin(x) @@ -1344,6 +1347,9 @@ def test_two_stage_cuda_graph_setup_uses_pipeline_config(self): """CUDA graph setup runs before the two-stage model_config is assigned.""" class TinyTransformer: + def register_cuda_graph_extra_key_fns(self, runner): + del runner + def forward(self, *args, **kwargs): return args, kwargs @@ -1364,6 +1370,180 @@ def forward(self, *args, **kwargs): assert runner._lora_state_getter() == "original" assert pipeline.transformer.forward.__wrapped__.__self__ is pipeline.transformer + def test_two_stage_cuda_graph_setup_registers_sol_phase_key(self): + """Dense and sparse SOL phases must never reuse one CUDA graph.""" + from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig + from tensorrt_llm._torch.visual_gen.models.ltx2.ltx2_core.modality import Modality + from tensorrt_llm._torch.visual_gen.models.modeling import BaseDiffusionModel + from tensorrt_llm.visual_gen.args import SolAttentionConfig + + class TinySolTransformer(BaseDiffusionModel): + def __init__(self): + super().__init__( + DiffusionModelConfig( + attention=AttentionConfig( + backend="TRTLLM", + sparse_attention_config=SolAttentionConfig(disabled_until_timestep=0.6), + ) + ) + ) + self.active_topology = "default" + + def forward(self, video, audio, *, text_cache, timestep=None, step_index=None): + del audio, text_cache, timestep, step_index + return video.latent, None + + pipeline = object.__new__(ltx2_two_stages.LTX2TwoStagesPipeline) + torch.nn.Module.__init__(pipeline) + pipeline.pipeline_config = DiffusionPipelineConfig( + cuda_graph=CudaGraphConfig(enable=True), + torch_compile=TorchCompileConfig(enable=False), + ) + pipeline.transformer = TinySolTransformer() + pipeline._cuda_graph_runners = {} + pipeline._setup_cuda_graphs() + runner = pipeline._cuda_graph_runners["transformer"] + + captured_keys = [] + + def fake_capture(key, fn, args, kwargs): + del fn, args, kwargs + captured_keys.append(key) + runner.graphs[key] = object() + + runner.capture = fake_capture + runner.replay = lambda key, args, kwargs: key + + def modality(timestep): + return Modality( + latent=torch.empty(1, 2, 4), + timesteps=torch.tensor([timestep]), + positions=torch.empty(1, 3, 2), + context=torch.empty(1, 3, 4), + ) + + dense_key = pipeline.transformer( + video=modality(0.8), audio=None, text_cache=None, step_index=0 + ) + sparse_key = pipeline.transformer( + video=modality(0.2), audio=None, text_cache=None, step_index=1 + ) + dense_replay_key = pipeline.transformer( + video=modality(0.8), audio=None, text_cache=None, step_index=2 + ) + + assert "sol_attn_phase" in runner._extra_key_fns + assert ("sol_attn_phase", (("video", 0),)) in dense_key + assert ("sol_attn_phase", (("video", 1),)) in sparse_key + assert dense_key != sparse_key + assert dense_replay_key == dense_key + assert captured_keys == [dense_key, sparse_key] + + def test_ltx2_threads_raw_modality_timestep_for_sol_phase(self): + """SOL phase preparation must not use the AdaLN-transformed timestep.""" + from tensorrt_llm._torch.visual_gen.attention_backend.sparse.sol.params import SolParams + from tensorrt_llm._torch.visual_gen.models.ltx2.ltx2_core.modality import Modality + from tensorrt_llm._torch.visual_gen.models.ltx2.ltx2_core.transformer_args import ( + TransformerArgs, + ) + from tensorrt_llm._torch.visual_gen.models.ltx2.transformer_ltx2 import LTXModel + + def modality(timestep): + return Modality( + latent=torch.empty(1, 1, 1), + timesteps=torch.tensor([timestep]), + positions=torch.empty(1, 1, 1), + context=torch.empty(1, 1, 1), + ) + + def transformer_args(timestep): + return TransformerArgs( + x=torch.empty(1, 1, 1), + context=torch.empty(1, 1, 1), + context_mask=None, + timesteps=torch.tensor([timestep]), + embedded_timestep=torch.empty(1, 1, 1), + positional_embeddings=(torch.empty(1), torch.empty(1)), + cross_positional_embeddings=None, + cross_scale_shift_timestep=None, + cross_gate_timestep=None, + enabled=True, + ) + + class Preprocessor: + def __init__(self, value): + self.value = value + + def prepare(self, *args, **kwargs): + del args, kwargs + return self.value + + observed_phases = [] + + class Block: + def __call__( + self, + *, + video, + audio, + video_sol_timestep=None, + audio_sol_timestep=None, + **kwargs, + ): + del kwargs + video_timestep = ( + video.timesteps if video_sol_timestep is None else video_sol_timestep + ) + audio_timestep = ( + audio.timesteps if audio_sol_timestep is None else audio_sol_timestep + ) + observed_phases.append( + ( + SolParams.get_graph_phase_for_timestep( + video_timestep, + disabled_until_timestep=0.6, + ), + SolParams.get_graph_phase_for_timestep( + audio_timestep, + disabled_until_timestep=0.6, + ), + ) + ) + return video, audio + + model = object.__new__(LTXModel) + torch.nn.Module.__init__(model) + model._active_seq_size = 1 + model.model_type = SimpleNamespace( + is_video_enabled=lambda: True, + is_audio_enabled=lambda: True, + ) + model._audio_pad = 0 + model.video_args_preprocessor = Preprocessor(transformer_args(0.8)) + model.audio_args_preprocessor = Preprocessor(transformer_args(0.2)) + model._active_sharder = SimpleNamespace(is_active=False) + model.model_config = SimpleNamespace(cache_backend=None) + model.transformer_blocks = [Block()] + model.scale_shift_table = model.norm_out = model.proj_out = None + model.audio_scale_shift_table = model.audio_norm_out = model.audio_proj_out = None + model._process_output = lambda table, norm, proj, x, embedded: x + text_cache = SimpleNamespace( + video_context=None, + video_mask=None, + video_pe=None, + video_cross_pe=None, + video_kv=None, + audio_context=None, + audio_mask=None, + audio_pe=None, + audio_cross_pe=None, + audio_kv=None, + ) + + model(video=modality(0.2), audio=modality(0.8), text_cache=text_cache) + + assert observed_phases == [(1, 0)] + def test_cuda_graph_rejects_nonpersistent_lora_bindings(self): """CUDA graph is valid only when distilled LoRA uses persistent bindings.""" pipeline = object.__new__(ltx2_two_stages.LTX2TwoStagesPipeline) diff --git a/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py b/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py index fb0b16ba40eb..23293bfe21f8 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_attention_metadata.py @@ -1,9 +1,20 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest import torch +from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask +from tensorrt_llm._torch.attention.backends.sparse.params import ( + BlockSparseForwardInputs, + SparseBackendForwardArgs, + SparseRuntimeParams, +) from tensorrt_llm._torch.visual_gen.attention_backend import trtllm as visual_trtllm +from tensorrt_llm._torch.visual_gen.config import create_attention_metadata_state class _FakeBaseTrtllmAttentionMetadata: @@ -19,6 +30,84 @@ def prepare(self): self.prepare_calls += 1 +def _make_block_sparse_inputs(): + return BlockSparseForwardInputs( + q_block_size=64, + kv_block_size=64, + max_blocks_per_row=1, + block_indptr=torch.tensor([[[0, 1]]], dtype=torch.int32), + block_indices=torch.tensor([0], dtype=torch.int32), + ) + + +class _StopAtFmhaDispatch(Exception): + pass + + +def _make_core_forward_metadata(): + metadata = object.__new__(visual_trtllm.BaseTrtllmAttentionMetadata) + seq_lens = torch.tensor([4], dtype=torch.int32) + metadata._seq_lens = seq_lens + metadata._seq_lens_kv = seq_lens + metadata._seq_lens_cuda = None + metadata.kv_cache_manager = None + metadata._max_seq_len_storage = 4 + metadata.use_paged_context_fmha = False + metadata.cu_q_seqlens = None + metadata.cu_kv_seqlens = None + metadata.enable_flash_mla = False + metadata.spec_bl_tree_first_sparse_mask_offset_kv = None + metadata.spec_decoding_bl_tree_mask = None + metadata.kv_lens_cuda_runtime = torch.tensor([4], dtype=torch.int32) + metadata.kv_lens_runtime = torch.tensor([4], dtype=torch.int32) + metadata.prompt_lens_cuda_runtime = torch.tensor([4], dtype=torch.int32) + metadata.prompt_lens_cpu_runtime = torch.tensor([4], dtype=torch.int32) + metadata.host_request_types_runtime = torch.tensor([0], dtype=torch.int32) + metadata.max_context_q_len_override = None + return metadata + + +def _make_wrapper(cls=visual_trtllm.TrtllmAttention, *, quant_attention_config=None): + attention = object.__new__(cls) + attention.quant_attention_config = quant_attention_config + return attention + + +def _capture_core_forward(monkeypatch, captured: dict): + prepared_metadata = object() + monkeypatch.setattr( + visual_trtllm.TrtllmAttention, + "_prepare_metadata", + lambda self, batch_size, seq_len: prepared_metadata, + ) + monkeypatch.setattr( + visual_trtllm.TrtllmAttention, + "_concat_qkv", + lambda self, q, k, v, batch_size, seq_len, kv_seq_len: torch.cat( + [ + q.reshape(batch_size * seq_len, -1), + k.reshape(batch_size * kv_seq_len, -1), + v.reshape(batch_size * kv_seq_len, -1), + ], + dim=-1, + ), + ) + + def _capture_base_forward(self, q, k, v, metadata, forward_args=None, **kwargs): + captured.update( + q=q, + k=k, + v=v, + metadata=metadata, + forward_args=forward_args, + kwargs=kwargs, + ) + return q[:, :16] + + monkeypatch.setattr(visual_trtllm.BaseTrtllmAttention, "forward", _capture_base_forward) + return prepared_metadata + + def test_trtllm_attention_metadata_caches_distinct_seq_lens(monkeypatch): monkeypatch.setattr( visual_trtllm, @@ -62,3 +151,278 @@ def test_trtllm_attention_metadata_caches_distinct_seq_lens(monkeypatch): assert first_cached_seq_lens.data_ptr() != second_cached_seq_lens.data_ptr() assert first_metadata.seq_lens is first_cached_seq_lens assert second_metadata.seq_lens is second_cached_seq_lens + + +def test_trtllm_attention_layers_share_block_sparse_plan_cache(monkeypatch): + from tensorrt_llm._torch.attention.backends.fmha import prims_ts_block_sparse + + def _base_update_quant_config(self, new_quant_config): + del new_quant_config + self._fmha_manager = SimpleNamespace( + fmha_libs=[prims_ts_block_sparse.PrimsTSBlockSparseFmha(self)] + ) + + def _base_init(self, **kwargs): + del kwargs + self.is_mla_enable = False + self.kv_lora_rank = None + self.v_head_dim = None + self.head_dim = 64 + self.update_quant_config(None) + + monkeypatch.setattr( + visual_trtllm.BaseTrtllmAttention, + "update_quant_config", + _base_update_quant_config, + ) + monkeypatch.setattr(visual_trtllm.BaseTrtllmAttention, "__init__", _base_init) + attention_metadata_state = create_attention_metadata_state() + assert "block_sparse_fmha_cache" not in attention_metadata_state + + first = visual_trtllm.TrtllmAttention( + attention_metadata_state=attention_metadata_state, + ) + second = visual_trtllm.TrtllmAttention( + attention_metadata_state=attention_metadata_state, + ) + + assert not hasattr(first, "_block_sparse_fmha_cache_state") + assert not hasattr(second, "_block_sparse_fmha_cache_state") + first_fmha = first._fmha_manager.fmha_libs[0] + second_fmha = second._fmha_manager.fmha_libs[0] + assert first_fmha._contiguous_wrappers is second_fmha._contiguous_wrappers + assert first_fmha._paged_wrappers is second_fmha._paged_wrappers + + first.update_quant_config(None) + first_fmha = first._fmha_manager.fmha_libs[0] + assert first_fmha._contiguous_wrappers is second_fmha._contiguous_wrappers + assert first_fmha._paged_wrappers is second_fmha._paged_wrappers + assert attention_metadata_state["fmha_caches"]["prims_ts_block_sparse"] == { + "contiguous_wrappers": {}, + "paged_wrappers": {}, + } + + other = visual_trtllm.TrtllmAttention( + attention_metadata_state=create_attention_metadata_state(), + ) + other_fmha = other._fmha_manager.fmha_libs[0] + assert first_fmha._contiguous_wrappers is not other_fmha._contiguous_wrappers + assert first_fmha._paged_wrappers is not other_fmha._paged_wrappers + + +def test_visual_gen_wrapper_does_not_define_its_own_prediction_lifecycle(): + assert not hasattr(visual_trtllm, "SparseForwardInputs") + for name in ( + "block_sparse_attn_predict", + "sparse_post_process", + "_forward_impl", + ): + assert name not in visual_trtllm.TrtllmAttention.__dict__ + assert getattr(visual_trtllm.TrtllmAttention, "__parameters__", ()) == () + + +def test_forward_rejects_unexpected_kwargs_before_metadata_or_core(monkeypatch): + prepare_metadata = Mock(return_value=object()) + core_forward = Mock(return_value=torch.empty(4, 16)) + monkeypatch.setattr(visual_trtllm.TrtllmAttention, "_prepare_metadata", prepare_metadata) + monkeypatch.setattr(visual_trtllm.BaseTrtllmAttention, "forward", core_forward) + attention = _make_wrapper() + + with pytest.raises(TypeError) as exc_info: + attention.forward( + torch.randn(1, 4, 6, 8), + None, + None, + batch_size=1, + seq_len=4, + attention_maks=PredefinedAttentionMask.FULL, + timstep=torch.tensor([12]), + ) + + assert str(exc_info.value) == ( + "Unexpected TRTLLM attention forward keyword arguments: attention_maks, timstep" + ) + prepare_metadata.assert_not_called() + core_forward.assert_not_called() + + +def test_forward_flattens_fused_qkv_without_copy(monkeypatch): + captured = {} + prepared_metadata = _capture_core_forward(monkeypatch, captured) + attention = _make_wrapper() + qkv = torch.randn(1, 4, 6, 8) + timestep = torch.tensor([12]) + + output = attention.forward(qkv, None, None, batch_size=1, seq_len=4, timestep=timestep) + + assert output.shape == (1, 4, 16) + assert captured["q"].shape == (4, 48) + assert captured["q"].data_ptr() == qkv.data_ptr() + assert captured["k"] is None and captured["v"] is None + assert captured["metadata"] is prepared_metadata + assert captured["forward_args"].timestep is timestep + assert captured["forward_args"].sparse_backend_args is None + assert captured["forward_args"].sparse_runtime_params == SparseRuntimeParams() + assert captured["kwargs"] == {} + + +def test_forward_fuses_separate_qkv_without_sparse_backend_args(monkeypatch): + captured = {} + _capture_core_forward(monkeypatch, captured) + attention = _make_wrapper() + q = torch.randn(1, 4, 2, 8) + k = torch.randn_like(q) + v = torch.randn_like(q) + + attention.forward(q, k, v, batch_size=1, seq_len=4) + + assert captured["q"].shape == (4, 48) + torch.testing.assert_close(captured["q"][:, :16], q.reshape(4, 16)) + assert captured["k"] is None and captured["v"] is None + assert captured["forward_args"].sparse_backend_args is None + + +def test_forward_hands_separate_qkv_and_backend_args_to_core_for_block_sparse_routes( + monkeypatch, +): + captured = {} + _capture_core_forward(monkeypatch, captured) + attention = _make_wrapper() + q = torch.randn(1, 4, 2, 8) + k = torch.randn_like(q) + v = torch.randn_like(q) + backend_args = SparseBackendForwardArgs(block_sparse_inputs=_make_block_sparse_inputs()) + + output = attention.forward( + q, + k, + v, + batch_size=1, + seq_len=4, + sparse_backend_args=backend_args, + ) + + assert output.shape == (1, 4, 16) + assert captured["q"].data_ptr() == q.data_ptr() + assert captured["k"].data_ptr() == k.data_ptr() + assert captured["v"].data_ptr() == v.data_ptr() + assert captured["q"].shape == captured["k"].shape == captured["v"].shape == (4, 16) + assert captured["forward_args"].sparse_backend_args is backend_args + assert captured["forward_args"].sparse_runtime_params == SparseRuntimeParams() + + +def test_forward_hands_separate_qkv_to_core_when_backend_rejects_fused_qkv(monkeypatch): + class _SeparateQkvAttention(visual_trtllm.TrtllmAttention): + @classmethod + def support_fused_qkv(cls) -> bool: + return False + + captured = {} + _capture_core_forward(monkeypatch, captured) + attention = _make_wrapper(_SeparateQkvAttention) + q = torch.randn(1, 4, 2, 8) + + attention.forward(q, q, q, batch_size=1, seq_len=4) + + assert captured["k"] is not None and captured["v"] is not None + assert captured["q"].shape == (4, 16) + assert captured["forward_args"].sparse_backend_args is None + + +def test_forward_applies_sage_quantization_to_separate_qkv(monkeypatch): + captured = {} + _capture_core_forward(monkeypatch, captured) + quant_cfg = SimpleNamespace(q_block_size=1, k_block_size=2, v_block_size=3, qk_dtype="int8") + attention = _make_wrapper(quant_attention_config=quant_cfg) + q = torch.randn(1, 4, 2, 8) + + attention.forward(q, q, q, batch_size=1, seq_len=4) + + forward_args = captured["forward_args"] + assert captured["k"] is not None and captured["v"] is not None + assert forward_args.sage_attn_num_elts_per_blk_q == 1 + assert forward_args.sage_attn_num_elts_per_blk_k == 2 + assert forward_args.sage_attn_num_elts_per_blk_v == 3 + assert forward_args.sage_attn_qk_int8 is True + + +def test_forward_requires_separate_qkv_for_block_sparse_routes(monkeypatch): + prepare_metadata = Mock(return_value=object()) + monkeypatch.setattr(visual_trtllm.TrtllmAttention, "_prepare_metadata", prepare_metadata) + attention = _make_wrapper() + backend_args = SparseBackendForwardArgs(block_sparse_inputs=_make_block_sparse_inputs()) + + with pytest.raises(ValueError, match="separate q, k, and v"): + attention.forward( + torch.randn(1, 4, 6, 8), + None, + None, + batch_size=1, + seq_len=4, + sparse_backend_args=backend_args, + ) + + prepare_metadata.assert_not_called() + + +def test_forward_rejects_block_sparse_routes_with_quant_config(monkeypatch): + prepare_metadata = Mock(return_value=object()) + monkeypatch.setattr(visual_trtllm.TrtllmAttention, "_prepare_metadata", prepare_metadata) + attention = _make_wrapper(quant_attention_config=object()) + q = torch.randn(1, 4, 2, 8) + backend_args = SparseBackendForwardArgs(block_sparse_inputs=_make_block_sparse_inputs()) + + with pytest.raises(ValueError, match="quant_attention_config"): + attention.forward(q, q, q, batch_size=1, seq_len=4, sparse_backend_args=backend_args) + + prepare_metadata.assert_not_called() + + +@pytest.mark.parametrize("has_block_sparse_inputs", [False, True]) +def test_forward_reaches_core_fmha_with_module_predicted_routes( + monkeypatch, + has_block_sparse_inputs, +): + metadata = _make_core_forward_metadata() + monkeypatch.setattr( + visual_trtllm.TrtllmAttention, + "_prepare_metadata", + lambda self, batch_size, seq_len: metadata, + ) + + attention = _make_wrapper() + attention.sparse_params = None + attention.is_mla_enable = False + attention.num_heads = 2 + attention.num_kv_heads = 2 + attention.head_dim = 8 + attention.get_local_layer_idx = Mock(return_value=0) + attention._ensure_rope_table_size = Mock() + attention.print_skip_softmax_stat = False + attention.kv_scale_orig_quant = None + attention.kv_scale_quant_orig = None + attention.sparse_kv_predict = Mock(return_value=(None, None)) + attention.sparse_attn_predict = Mock(return_value=(None, None)) + select_fmha = Mock(side_effect=_StopAtFmhaDispatch) + attention._fmha_manager = SimpleNamespace( + fmha_libs=[object()], + select=select_fmha, + ) + carrier = _make_block_sparse_inputs() if has_block_sparse_inputs else None + backend_args = SparseBackendForwardArgs(block_sparse_inputs=carrier) + q = torch.randn(1, 4, 2, 8) + k = torch.randn_like(q) + v = torch.randn_like(q) + + with pytest.raises(_StopAtFmhaDispatch): + attention.forward(q, k, v, batch_size=1, seq_len=4, sparse_backend_args=backend_args) + + select_fmha.assert_called_once() + attention.sparse_kv_predict.assert_called_once() + attention.sparse_attn_predict.assert_called_once() + core_forward_args = select_fmha.call_args.args[5] + assert core_forward_args.sparse_backend_args is backend_args + runtime_params = core_forward_args.sparse_runtime_params + assert isinstance(runtime_params, SparseRuntimeParams) + assert runtime_params.block_sparse_inputs is carrier + assert runtime_params.sparse_attn_indices_block_size == 0 diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py index f1f29e18bd75..14a0d7c272ba 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,12 @@ ParallelConfig, QuantAttentionConfig, RuntimeLoRAConfig, + SkipSoftmaxAttentionConfig, + SolAttentionConfig, TeaCacheConfig, TorchCompileConfig, VAEConfig, + VideoSparseAttentionConfig, VisualGenArgs, ) @@ -101,6 +104,49 @@ def test_quant_config_rejected_when_unsupported(self): ), ) + @pytest.mark.parametrize( + ("backend", "quant_config"), + [ + ( + "TRTLLM", + QuantAttentionConfig( + qk_dtype="fp8", + q_block_size=1, + k_block_size=1, + v_block_size=1, + ), + ), + ( + "CUTEDSL", + QuantAttentionConfig(qk_dtype="bf16", v_dtype="fp8"), + ), + ], + ) + def test_vsa_and_quantization_are_mutually_exclusive(self, backend, quant_config): + with pytest.raises( + ValidationError, match="VSA and quant_attention_config are mutually exclusive" + ): + AttentionConfig( + backend=backend, + quant_attention_config=quant_config, + sparse_attention_config=VideoSparseAttentionConfig(vsa_sparsity=0.9), + ) + + def test_skip_softmax_and_sage_quantization_can_be_combined(self): + attention = AttentionConfig( + backend="TRTLLM", + quant_attention_config=QuantAttentionConfig( + qk_dtype="int8", + q_block_size=1, + k_block_size=4, + v_block_size=1, + ), + sparse_attention_config=SkipSoftmaxAttentionConfig(threshold_scale_factor=0.3), + ) + + assert attention.sparse_attention_config is not None + assert attention.sparse_attention_config.algorithm == "skip_softmax" + @pytest.mark.parametrize( ("qk_dtype", "q_block_size", "k_block_size", "v_block_size"), [ @@ -448,6 +494,54 @@ def test_from_yaml_unknown_field_raises(self, tmp_path): with pytest.raises(ValidationError, match="Extra inputs are not permitted"): VisualGenArgs.from_yaml(yaml_path) + def test_from_yaml_rejects_sol_with_enabled_fullgraph(self, tmp_path): + yaml_path = tmp_path / "sol_fullgraph.yml" + yaml_path.write_text( + "model: /tmp/model\n" + "attention_config:\n" + " backend: TRTLLM\n" + " sparse_attention_config:\n" + " algorithm: sol_attn\n" + "torch_compile_config:\n" + " enable: true\n" + " enable_fullgraph: true\n" + ) + + with pytest.raises(ValidationError, match="SOL.*fullgraph"): + VisualGenArgs.from_yaml(yaml_path) + + +class TestVisualGenArgsCrossFieldValidation: + def test_rejects_sol_with_enabled_fullgraph(self): + with pytest.raises(ValidationError, match="SOL.*fullgraph"): + VisualGenArgs( + model="/tmp/model", + attention_config=AttentionConfig( + backend="TRTLLM", + sparse_attention_config=SolAttentionConfig(), + ), + torch_compile_config=TorchCompileConfig( + enable=True, + enable_fullgraph=True, + ), + ) + + def test_allows_sol_fullgraph_field_when_torch_compile_disabled(self): + args = VisualGenArgs( + model="/tmp/model", + attention_config=AttentionConfig( + backend="TRTLLM", + sparse_attention_config=SolAttentionConfig(), + ), + torch_compile_config=TorchCompileConfig( + enable=False, + enable_fullgraph=True, + ), + ) + + assert args.torch_compile_config.enable is False + assert args.torch_compile_config.enable_fullgraph is True + class TestParallelConfigValidation: """ParallelConfig no longer checks WORLD_SIZE at construction time.""" diff --git a/tests/unittest/_torch/visual_gen/test_wan_vsa_pipeline.py b/tests/unittest/_torch/visual_gen/test_wan_vsa_pipeline.py index 979619e8e433..7c2eddf63e7f 100644 --- a/tests/unittest/_torch/visual_gen/test_wan_vsa_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_wan_vsa_pipeline.py @@ -148,7 +148,7 @@ def _assert_vsa_matches_dense( """Compare CuTe-DSL VSA against SDPA-fallback VSA (same gated formulation, different fine kernel).""" from unittest.mock import patch - from tensorrt_llm._torch.visual_gen.attention_backend.cute_dsl import vsa as _vsa_module + from tensorrt_llm._torch.visual_gen.attention_backend.sparse.vsa import backend as _vsa_module common_kwargs = dict( prompt=PROMPT,