Sol Attention Implementation

docs/sol_attention.md

@@ -0,0 +1,146 @@
+# Sol-Attn: sparse self-attention for video/image DiTs
+
+[Sol-Attn](https://arxiv.org/abs/2607.24027) is a training-free sparse
+attention method. For every query block it picks the key/value blocks whose
+block-proxy score `mean(Q_i) · mean(K_j)` exceeds a query-dependent threshold
+`tau_i = mu_i + beta * sigma_i`, computes those blocks exactly, and replaces
+every other block with its zeroth-order approximation (one pooled key with
+weight `log(block size)`, value `mean(V_j)`). Both branches share one
+online-softmax state, so skipped blocks still contribute to the softmax
+numerator and denominator.
+
+The expected exact-block density is about `1 - Phi(beta)`: `beta = 1.0` keeps
+~16% of the blocks (~84% sparsity), close to the paper's Wan 2.1 setting.
+
+## Using it with Wan 2.1 14B
+
+```bash
+python src/maxdiffusion/generate_wan.py src/maxdiffusion/configs/base_wan_14b.yml \
+  attention=ulysses_custom ici_data_parallelism=2 ici_context_parallelism=4 \
+  use_sol_attention=True sol_beta=1.0 sol_block_q=1024 sol_block_kv=256
+```
+
+| Key | Default | Meaning |
+|---|---|---|
+| `use_sol_attention` | `False` | Enable Sol-Attn for self-attention (cross-attention stays dense). Mutually exclusive with `use_svg_attention`. |
+| `sol_beta` | `1.0` | Standardized threshold; higher is sparser. |
+| `sol_block_q` | `1024` | Query routing block and kernel query tile, in tokens (multiple of 128). |
+| `sol_block_kv` | `256` | Key/value block for exact compute and pooling (multiple of 128). |
+| `sol_active_{start,end}_step` | `10, 50` | Sparse steps `[start, end)`; earlier steps are a dense warm-up (paper: first 20%). |
+| `sol_active_{start,end}_layer` | `1, 40` | Sparse layers `[start, end)`; layer 0 stays dense. |
+| `sol_dense_layer_fraction`, `sol_dense_timestep_fraction` | `0.0` | Fraction-based warm-up, used when no explicit interval is set. |
+| `sol_log_density` | `False` | Print the measured exact-block density per call (debug). |
+
+Sol-Attn works with any attention backend whose self-attention either shards
+heads (`flash`) or the sequence over one mesh axis (`ulysses*`): in the latter
+case an all-to-all gives each device full sequences for a subset of heads.
+
+## Results on TPU v6e-8
+
+Wan 2.1 14B T2V, 720p, 81 frames, 50 steps, `attention=ulysses_custom`,
+`ici_data_parallelism=2 ici_context_parallelism=4`, one prompt, seed 12345,
+with the async all-to-all flags below. The dense warm-up follows the paper
+(steps `[0, 10)` and layer 0 dense). PSNR and SSIM compare each video with the
+dense video, as in Table 1 of the paper (the paper reports PSNR 20.59 / SSIM
+0.68 at 85% sparsity on H100).
+
+```bash
+export LIBTPU_INIT_ARGS="--xla_tpu_enable_async_all_to_all=true --xla_tpu_overlap_compute_collective_tc=true \
+  --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true"
+```
+
+| Config | Generation time | Speedup | PSNR | SSIM |
+|---|---|---|---|---|
+| dense | 229.6 s | 1.00x | – | – |
+| `sol_block_q=1024 sol_block_kv=256` (default) | 147.1 s | 1.56x | 20.18 dB | 0.704 |
+| `sol_block_q=2048 sol_block_kv=256` | 138.4 s | 1.66x | 18.49 dB | 0.633 |
+| default + `ulysses_attention_chunks=2` | 146.9 s | 1.56x | 20.28 dB | 0.708 |
+| default + `sol_active_start_step=5` (10% dense warm-up) | 136.6 s | 1.68x | 17.01 dB | 0.587 |
+
+Shortening the dense warm-up below the paper's 20% buys speed at a clear
+fidelity cost.
+
+The measured exact-block density on Wan at `sol_beta=1.0` is about 16-17%.
+
+Per self-attention layer (all 8 chips, synthetic video-like inputs at the same
+shapes), dense `ulysses_custom` takes 82.5 ms and Sol-Attn 22.8 ms at
+`1024x256` (3.6x). Of the 22.8 ms, the kernel is ~14 ms and the Ulysses
+exchange (all-to-alls plus XLA layout copies, the same as for dense) ~8 ms.
+
+## Code layout
+
+```
+kernels/sol_attention/
+  routing.py               pooling, thresholds (Eq. 3-5), index compaction (pure JAX)
+  pooling.py               fused single-pass Pallas pooling of Q/K/V (+ key block radii)
+  sol_attention_kernel.py  Pallas TPU kernel + `sol_attention(q, k, v, ...)` entry point
+  reference.py             dense O(L^2) oracle of Eqs. (9)-(10), for tests
+models/
+  sol_attention.py         config resolution, warm-up dispatch, sharded entry point
+  sparse_schedule.py       step/layer warm-up schedule shared with SVG
+```
+
+`kernels/sol_attention` has no model dependencies: `sol_attention(q, k, v,
+scale=..., beta=...)` takes `[..., L, d]` tensors and is usable on its own.
+
+## TPU adaptation
+
+The paper's GPU kernel fuses routing into the attention kernel to avoid
+materializing the `N x N` block-proxy map. On TPU the design differs in a few
+places, each driven by measurements on v6e:
+
+* **Routing in a small XLA pre-pass.** With TPU-sized blocks the proxy map is
+  tiny (e.g. 148 x 148 per head at 512-token blocks), so `mu_i` and `sigma_i`
+  are taken directly from the materialized row (identical to the moment
+  estimator of Eq. 5). Only compacted per-query-block index rows reach the
+  kernel; the selection mask for the approximate branch is rebuilt on chip
+  from them.
+* **Exact branch via manual DMAs.** Keys/values stay in HBM; the kernel walks
+  the selected blocks through a ring of async copies, so only routed blocks are
+  read. The next grid step's index row is prefetched into SMEM.
+* **Transposed score layout.** Scores are `[keys, queries]` so softmax
+  reductions run across sublanes, as in the dense splash kernel.
+* **Branch-free hot loop.** The ragged tail block (sequence length not a
+  multiple of `sol_block_kv`) runs as a masked epilogue; a per-block branch
+  cost ~40%.
+* **Larger blocks than the paper's 64 x 64.** MXU efficiency and per-block
+  overheads favor 512-2048 query tokens and 256-512 key tokens.
+* **Grouped hot loop.** Selected blocks land in contiguous ring slots, so the
+  loop fuses several of them into one ~1024-key score tile, keeping routing
+  at `sol_block_kv` granularity. This alone is 10-25% faster for 128-256
+  token blocks.
+* **Fixed softmax offset.** Instead of a running max, each query uses
+  `m = max(m_est, ub - 80)` (log2 units): `m_est` (the largest pooled score,
+  a lower bound of the true max) rules out underflow, and
+  `ub = max_j (q . mean(K_j) + ||q|| * radius_j)` (an upper bound of every
+  token score, with per-block key radii from the pooling kernel) rules out
+  f32 overflow. The hot loop then has no max reduction and no accumulator
+  rescale (~20% faster). Query tiles whose bounds are too far apart (e.g. an
+  outlier key inflating a block radius) fall back to the online softmax.
+* **Fused pooling, no padding.** One Pallas pass computes the Q/K/V block
+  means and key radii straight from the unpadded `[H, L, d]` layout (XLA's
+  reshape-and-mean cost ~0.4 ms per tensor). The attention kernel reads
+  unpadded tensors too: the ragged tail key block is fetched as the last
+  `block_kv` rows and masked, and the output is transposed on chip, so no
+  pad/slice/transpose copies remain around the kernel.
+* **Head-chunked exchange.** `ulysses_attention_chunks` splits heads into
+  chunks so XLA can overlap one chunk's all-to-all with another's kernel.
+  On v6e-8 this is neutral (the exchange is dominated by layout copies), but
+  it may help on slower interconnects.
+
+## Extending to other models
+
+1. In the attention module constructor, call
+   `sol_attention.init_sol_config(attention_config, num_layers, is_self_attention)`
+   and `setattr` the result on the module.
+2. In `__call__`, wrap the attention op:
+   `sol_attention.apply_sol_or_dense(self, run_attention, layer_index, step_index, timestep)`.
+   `run_attention(sparse_config_override=...)` must forward the override to
+   `attention_flax._apply_attention(spatiotemporal_config=...)`.
+3. In the pipeline, merge `sol_attention.attention_config_from_pyconfig(config, num_layers)`
+   into the transformer's `attention_config`, and add
+   `aot_cache.extract_sol_meta(config, pipeline)` to the AOT metadata.
+
+For MMDiT-style joint attention (text and video tokens in one sequence), pass
+`always_exact` to the kernel to keep text blocks in the exact branch, as the
+paper does for HunyuanVideo.

src/maxdiffusion/aot_cache.py

@@ -143,6 +143,54 @@ def extract_svg_meta(config: Any, pipeline: Any = None) -> dict[str, Any]:
   return meta
 
 
+def extract_sol_meta(config: Any, pipeline: Any = None) -> dict[str, Any]:
+  """Extracts all graph-affecting Sol-Attn settings for AOT cache identity."""
+  sol_keys = (
+      "use_sol_attention",
+      "sol_beta",
+      "sol_block_q",
+      "sol_block_kv",
+      "sol_active_start_step",
+      "sol_active_end_step",
+      "sol_active_start_layer",
+      "sol_active_end_layer",
+      "sol_dense_layer_fraction",
+      "sol_dense_timestep_fraction",
+      "sol_num_train_timesteps",
+      "sol_num_layers",
+      "sol_vmem_limit_bytes",
+      "sol_log_density",
+  )
+  enabled = config is not None and bool(
+      config.get("use_sol_attention", False) if isinstance(config, dict) else getattr(config, "use_sol_attention", False)
+  )
+  if not enabled:
+    # Dense runs keep their pre-Sol cache identity.
+    return {}
+  meta: dict[str, Any] = {}
+  if config is not None:
+    for k in sol_keys:
+      val = getattr(config, k, None) if not isinstance(config, dict) else config.get(k, None)
+      if val is not None:
+        meta[k] = str(val)
+  if pipeline is not None:
+    for attr in ("transformer", "high_noise_transformer", "low_noise_transformer"):
+      t = getattr(pipeline, attr, None)
+      if t is None:
+        continue
+      t_cfg = getattr(t, "config", None)
+      attn_cfg = getattr(t_cfg, "attention_config", None) or (
+          t_cfg.get("attention_config") if isinstance(t_cfg, dict) else None
+      )
+      if isinstance(attn_cfg, dict):
+        meta[f"{attr}_sol_config"] = json.dumps(
+            {k: v for k, v in attn_cfg.items() if str(k).startswith(("sol_", "use_sol_"))},
+            sort_keys=True,
+            default=str,
+        )
+  return meta
+
+
 def _metadata_fingerprint(meta: dict[str, Any]) -> str:
   """Returns the stable filename fingerprint for install-time metadata."""
   serialized = json.dumps(meta, sort_keys=True, default=str)

src/maxdiffusion/configs/base_wan_14b.yml

@@ -101,6 +101,25 @@ svg_low_noise_density: -1.0
 # compute blocks and the sparse kernel for small ones. Must stay a dict so the
 # command line can override it with JSON.
 svg_flash_block_sizes: {}
+# Sol-Attn (arXiv:2607.24027): training-free sparse self-attention with
+# query-dependent threshold routing and pooled-key approximation of skipped
+# blocks. Mutually exclusive with use_svg_attention. Works with any attention
+# backend whose self-attention shards either heads or (Ulysses) sequence.
+use_sol_attention: False
+# Standardized cutoff: expected exact density ~= 1 - Phi(sol_beta).
+# 1.0 -> ~16% density (~84% sparsity), matching the paper's Wan 2.1 setup.
+sol_beta: 1.0
+# Query routing block / kernel query tile and key/value block, in tokens.
+# Smaller blocks route more precisely but run less efficiently on the MXU.
+sol_block_q: 1024
+sol_block_kv: 256
+# Dense warm-up, [start, end) intervals. The paper keeps the first 20% of
+# denoising steps and the first layer dense: steps [10, 50), layers [1, 40).
+sol_active_start_step: 10
+sol_active_end_step: 50
+sol_active_start_layer: 1
+sol_active_end_layer: 40
+sol_log_density: False
 attention: 'flash' # Supported attention: dot_product, flash, tokamax_flash, cudnn_flash_te, ring, tokamax_ring, ulysses, ulysses_custom, ulysses_ring
 
 use_base2_exp: True

src/maxdiffusion/generate_wan.py

@@ -315,6 +315,7 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None):
           "scan_layers": str(config.scan_layers),
           "jax": jax.__version__,
           **aot_cache.extract_svg_meta(config, pipeline),
+          **aot_cache.extract_sol_meta(config, pipeline),
       },
       mesh=pipeline.mesh,
   )

src/maxdiffusion/kernels/sol_attention/__init__.py

@@ -0,0 +1,20 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Sol-Attn: training-free sparse attention with threshold routing and
+approximate correction (arXiv:2607.24027)."""
+
+from maxdiffusion.kernels.sol_attention.routing import SolBlockSizes
+from maxdiffusion.kernels.sol_attention.sol_attention_kernel import sol_attention
+
+__all__ = ["SolBlockSizes", "sol_attention"]

src/maxdiffusion/kernels/sol_attention/pooling.py

@@ -0,0 +1,116 @@
+"""
+Copyright 2026 Google LLC
+
+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
+
+     https://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.
+
+Fused single-pass block pooling of Q, K and V for Sol-Attn on TPU.
+
+XLA lowers the reshape-and-mean pooling of `routing.pool_blocks` to a relayout
+copy plus a reduction per tensor (~0.4 ms each for Wan 2.1 14B shapes on v6e).
+This kernel reads each tensor once, straight from the unpadded [H, L, d]
+layout, and writes the three block means.
+"""
+
+from __future__ import annotations
+
+import functools
+import math
+
+import jax
+import jax.numpy as jnp
+from jax import lax
+from jax.experimental import pallas as pl
+from jax.experimental.pallas import tpu as pltpu
+
+NUM_SUBLANES = 8
+
+
+def slab_rows(block_q: int, block_kv: int) -> int:
+  """Rows pooled per grid step; each output gets >= 8 pooled rows (sublane tiling)."""
+  return math.lcm(NUM_SUBLANES * block_q, NUM_SUBLANES * block_kv)
+
+
+def _pool_kernel(
+    q_ref, k_ref, v_ref, qo_ref, ko_ref, vo_ref, kr_ref, *, seq_len: int, tile: int, block_q: int, block_kv: int
+):
+  """Block means of one `tile`-row slab of q (block_q) and k/v (block_kv), plus key block radii."""
+  start = pl.program_id(1) * tile
+  rows = start + lax.broadcasted_iota(jnp.int32, (tile, 1), 0)
+  valid = rows < seq_len  # The last slab may run past the sequence end.
+
+  def pool(x_ref, out_ref, block):
+    x = jnp.where(valid, x_ref[...].astype(jnp.float32), 0.0)
+    sums = x.reshape(tile // block, block, x.shape[-1]).sum(axis=1)
+    block_start = start + lax.broadcasted_iota(jnp.int32, (tile // block, 1), 0) * block
+    counts = jnp.clip(seq_len - block_start, 1, block).astype(jnp.float32)
+    mean = sums / counts
+    out_ref[...] = mean
+    return x, mean
+
+  pool(q_ref, qo_ref, block_q)
+  k, k_mean = pool(k_ref, ko_ref, block_kv)
+  pool(v_ref, vo_ref, block_kv)
+  dev = k.reshape(tile // block_kv, block_kv, k.shape[-1]) - k_mean[:, None, :]
+  dist2 = jnp.where(valid.reshape(tile // block_kv, block_kv, 1), dev * dev, 0.0).sum(axis=2, keepdims=True)
+  kr_ref[...] = jnp.sqrt(dist2.max(axis=1))
+
+
+def pool_qkv(
+    q: jax.Array,
+    k: jax.Array,
+    v: jax.Array,
+    *,
+    block_q: int,
+    block_kv: int,
+    interpret: bool = False,
+) -> tuple[jax.Array, jax.Array, jax.Array, jax.Array]:
+  """Means of q over block_q tokens and of k, v over block_kv tokens.
+
+  Args:
+    q, k, v: Unpadded [H, L, d] tensors.
+    block_q, block_kv: Pooling block sizes (multiples of 128).
+    interpret: Run in Pallas interpret mode.
+
+  Returns:
+    float32 (q means [H, ceil(L/block_q), d], k and v means [H, ceil(L/block_kv), d],
+    k radii [H, ceil(L/block_kv)]; see `routing.block_radius`). A ragged last
+    block averages only its real tokens.
+  """
+  num_heads, seq_len, head_dim = q.shape
+  tile = slab_rows(block_q, block_kv)
+  if seq_len < tile:
+    raise ValueError(f"pool_qkv needs at least {tile} tokens, got {seq_len}.")
+  nq, nk = -(-seq_len // block_q), -(-seq_len // block_kv)
+  in_spec = pl.BlockSpec((None, tile, head_dim), lambda h, i: (h, i, 0))
+  kernel = functools.partial(_pool_kernel, seq_len=seq_len, tile=tile, block_q=block_q, block_kv=block_kv)
+  q_pooled, k_pooled, v_pooled, k_radius = pl.pallas_call(
+      kernel,
+      grid=(num_heads, -(-seq_len // tile)),
+      in_specs=[in_spec, in_spec, in_spec],
+      out_specs=[
+          pl.BlockSpec((None, tile // block_q, head_dim), lambda h, i: (h, i, 0)),
+          pl.BlockSpec((None, tile // block_kv, head_dim), lambda h, i: (h, i, 0)),
+          pl.BlockSpec((None, tile // block_kv, head_dim), lambda h, i: (h, i, 0)),
+          pl.BlockSpec((None, tile // block_kv, 1), lambda h, i: (h, i, 0)),
+      ],
+      out_shape=[
+          jax.ShapeDtypeStruct((num_heads, nq, head_dim), jnp.float32),
+          jax.ShapeDtypeStruct((num_heads, nk, head_dim), jnp.float32),
+          jax.ShapeDtypeStruct((num_heads, nk, head_dim), jnp.float32),
+          jax.ShapeDtypeStruct((num_heads, nk, 1), jnp.float32),
+      ],
+      compiler_params=pltpu.CompilerParams(dimension_semantics=("parallel", "parallel")),
+      interpret=interpret,
+      name="sol_pool_qkv",
+  )(q, k, v)
+  return q_pooled, k_pooled, v_pooled, k_radius[..., 0]