feat(attention): integrate Sol-Attn for Wan2.1

docs/en/attention.md

@@ -31,6 +31,7 @@ class AttnImplType(Enum):
     # Sparse attention
     RADIAL_ATTN = auto()
     LOCAL_SPARSE_ATTN = auto()
+    SOL_ATTN = auto()
 ```
 
 ### AttentionConfig
@@ -51,6 +52,7 @@ Factory methods:
 - `AttentionConfig.dense_attention(attn_impl)` - Create dense attention config
 - `AttentionConfig.radial_attention(**kwargs)` - Create radial sparse attention config
 - `AttentionConfig.local_sparse_attention(**kwargs)` - Create local sparse attention config
+- `AttentionConfig.sol_attention(**kwargs)` - Create dynamic Sol-Attn config
 
 ### SparseAttentionConfig
 
@@ -59,13 +61,16 @@ Configuration for sparse attention:
 ```python
 @dataclass
 class SparseAttentionConfig:
-    sparse_impl: str | None = None           # "radial", "local", etc.
+    sparse_impl: str | None = None           # "radial", "local", "sol", etc.
     dense_timesteps: int = 40               # Use dense attention for initial timesteps
     dense_layers: int = 0                   # Use dense attention for initial layers
     decay_factor: float = 1.0               # Decay factor for attention window
     local_window_size: int = 6              # Window size for local sparse attention
     block_size: int = 128                   # Block size for sparse computation
     use_sage_attention: bool = False        # Use sage attention backend
+    sol_tau: float = 1.0                    # Sol-Attn routing threshold
+    sol_threshold_type: str = "diag"        # "diag" or "exact"
+    sol_kv_splits: int | str = "auto"       # "auto", 1, 2, or 4
 ```
 
 ## Calling Flow
@@ -155,12 +160,12 @@ else:
 
 ## Pipeline Support Status
 
-| Pipeline | Dense Attention | Sparse (Radial) | Notes |
-|----------|-----------------|-----------------|-------|
-| `Wan21VideoPipeline` | ✅ | ✅ | Full support for video generation |
-| `Wan22VideoPipeline` | ✅ | ✅ | Full support for video generation |
-| `QwenImagePipeline` | ✅ | ❌ | Image generation doesn't need temporal sparse attention |
-| `ZImagePipeline` | ✅ | ❌ | Image generation doesn't need temporal sparse attention |
+| Pipeline | Dense Attention | Radial | Sol-Attn | Notes |
+|----------|-----------------|--------|----------|-------|
+| `Wan21VideoPipeline` | Yes | Yes | Experimental | Sol-Attn covers eligible self-attention calls |
+| `Wan22VideoPipeline` | Yes | Yes | No | Sol-Attn is not wired into Wan2.2 yet |
+| `QwenImagePipeline` | Yes | No | No | Image generation doesn't need temporal sparse attention |
+| `ZImagePipeline` | Yes | No | No | Image generation doesn't need temporal sparse attention |
 
 ### Wan21VideoPipeline / Wan22VideoPipeline
 
@@ -185,6 +190,18 @@ When using radial attention:
 3. Updates state per timestep/layer in denoising loop
 4. Automatically falls back to dense for early timesteps/layers
 
+Wan2.1 can select Sol-Attn through the same pipeline configuration surface:
+
+```python
+config = AttentionConfig.sol_attention()
+pipe_config.dit_config.attention_config = config
+```
+
+Sol-Attn is used only for contiguous, noncausal BF16 self-attention with equal Q/K/V
+shapes and head dimension 128. Unsupported calls, dense warmup layers or timesteps,
+and kernel runtime failures fall back to the existing dense attention path. Ring/USP
+also remains dense because its online merge requires log-sum-exp output.
+
 ### QwenImagePipeline / ZImagePipeline
 
 Supports only dense attention (image generation doesn't have temporal dimension):
@@ -218,9 +235,14 @@ pipe_config.dit_config.attention_config = config
 |---------|-------------|--------------|
 | `RADIAL_ATTN` | Radial attention for video | `flashinfer` or `sageattention` (tf-kernel prioritized) |
 | `LOCAL_SPARSE_ATTN` | Local window sparse attention | `block_sparse_attn` |
+| `SOL_ATTN` | Dynamic block-sparse video attention | Built in; BF16, head dimension 128, SM80+ |
 
 **Note on SageAttention Priority**: When `use_sage_attention=True` is set, the system will prioritize tf-kernel's sageattention implementation over the standalone `sageattention` package if both are available. This provides better performance and integration with the TeleFuser kernel library.
 
+**Sol-Attn packaging**: Sol-Attn ships with TeleFuser under `telefuser.kernel.sol_attn`; it does not require
+`tf-kernel`. The upstream runtime targets PyTorch 2.10+, CUDA 12.8+, and Triton 3.6+. Specialized CuTe DSL
+kernels are selected when that optional runtime is available, otherwise SM80+ uses the Triton implementation.
+
 ### Installing Sparge Attention
 
 To use `SPARGE_ATTN` backend or sparse sage attention in radial attention, you need to install `spas_sage_attn` from source:
@@ -428,6 +450,7 @@ print(f"FlashInfer: {FLASHINFER_AVAILABLE}")
 | Flash Attention 4 | Build from source (cute interface) | SM90+ (H100, B100/B200) |
 | SageAttention | tf-kernel or [official source](https://github.com/thu-ml/SageAttention) | SM80+ |
 | Radial Attention | tf-kernel or [FlashInfer source](https://github.com/flashinfer-ai/flashinfer) | SM80+ |
+| Sol-Attn | Built into TeleFuser | SM80+; CuTe on supported SM90/100, Triton fallback |
 | Block Sparse | tf-kernel or [official source](https://github.com/mit-han-lab/Block-Sparse-Attention) | SM80+ |
 | Sparge Attention | Install from source (see above) | SM80, SM86, SM89, SM90 |
 

docs/zh/attention.md

@@ -31,6 +31,7 @@ class AttnImplType(Enum):
     # 稀疏注意力
     RADIAL_ATTN = auto()
     LOCAL_SPARSE_ATTN = auto()
+    SOL_ATTN = auto()
 ```
 
 ### AttentionConfig
@@ -51,6 +52,7 @@ class AttentionConfig:
 - `AttentionConfig.dense_attention(attn_impl)` - 创建密集注意力配置
 - `AttentionConfig.radial_attention(**kwargs)` - 创建径向稀疏注意力配置
 - `AttentionConfig.local_sparse_attention(**kwargs)` - 创建局部稀疏注意力配置
+- `AttentionConfig.sol_attention(**kwargs)` - 创建动态 Sol-Attn 配置
 
 ### SparseAttentionConfig
 
@@ -59,13 +61,16 @@ class AttentionConfig:
 ```python
 @dataclass
 class SparseAttentionConfig:
-    sparse_impl: str | None = None           # "radial", "local" 等
+    sparse_impl: str | None = None           # "radial", "local", "sol" 等
     dense_timesteps: int = 40               # 初始时间步使用密集注意力
     dense_layers: int = 0                   # 初始层使用密集注意力
     decay_factor: float = 1.0               # 注意力窗口衰减因子
     local_window_size: int = 6              # 局部稀疏注意力窗口大小
     block_size: int = 128                   # 稀疏计算块大小
     use_sage_attention: bool = False        # 使用 sage attention 后端
+    sol_tau: float = 1.0                    # Sol-Attn 路由阈值
+    sol_threshold_type: str = "diag"        # "diag" 或 "exact"
+    sol_kv_splits: int | str = "auto"       # "auto"、1、2 或 4
 ```
 
 ## 调用流程
@@ -155,12 +160,12 @@ else:
 
 ## Pipeline 支持情况
 
-| Pipeline | 密集注意力 | 稀疏 (径向) | 说明 |
-|----------|-----------|------------|------|
-| `Wan21VideoPipeline` | ✅ | ✅ | 视频生成完整支持 |
-| `Wan22VideoPipeline` | ✅ | ✅ | 视频生成完整支持 |
-| `QwenImagePipeline` | ✅ | ❌ | 图像生成不需要时序稀疏注意力 |
-| `ZImagePipeline` | ✅ | ❌ | 图像生成不需要时序稀疏注意力 |
+| Pipeline | 密集注意力 | Radial | Sol-Attn | 说明 |
+|----------|-----------|--------|----------|------|
+| `Wan21VideoPipeline` | 支持 | 支持 | 实验性 | Sol-Attn 用于满足约束的 self-attention |
+| `Wan22VideoPipeline` | 支持 | 支持 | 不支持 | 尚未接入 Wan2.2 |
+| `QwenImagePipeline` | 支持 | 不支持 | 不支持 | 图像生成不需要时序稀疏注意力 |
+| `ZImagePipeline` | 支持 | 不支持 | 不支持 | 图像生成不需要时序稀疏注意力 |
 
 ### Wan21VideoPipeline / Wan22VideoPipeline
 
@@ -185,6 +190,17 @@ pipe_config.dit_config.attention_config = config
 3. 在去噪循环中每时间步/层更新状态
 4. 早期时间步/层自动回退到密集注意力
 
+Wan2.1 可以通过同一配置入口启用 Sol-Attn：
+
+```python
+config = AttentionConfig.sol_attention()
+pipe_config.dit_config.attention_config = config
+```
+
+Sol-Attn 仅用于连续、非因果、BF16、Q/K/V 形状相同且 head dimension 为 128 的
+self-attention。其他调用、dense 预热层/时间步以及内核运行失败都会回退到现有密集路径。
+Ring/USP 需要 LSE 做在线合并，因此仍使用支持 LSE 的密集后端。
+
 ### QwenImagePipeline / ZImagePipeline
 
 仅支持密集注意力（图像生成没有时序维度）：
@@ -218,9 +234,14 @@ pipe_config.dit_config.attention_config = config
 |------|------|------|
 | `RADIAL_ATTN` | 视频径向注意力 | `flashinfer` 或 `sageattention` (优先使用 tf-kernel) |
 | `LOCAL_SPARSE_ATTN` | 局部窗口稀疏注意力 | `block_sparse_attn` |
+| `SOL_ATTN` | 动态块稀疏视频注意力 | 内置；BF16、head dimension 128、SM80+ |
 
 **SageAttention 优先级说明**: 当设置 `use_sage_attention=True` 时，如果 tf-kernel 和独立的 `sageattention` 包都可用，系统将优先使用 tf-kernel 的 sageattention 实现。这提供了更好的性能和与 TeleFuser 内核库的集成。
 
+**Sol-Attn 打包方式**：Sol-Attn 随 TeleFuser 发布，位于 `telefuser.kernel.sol_attn`，不依赖
+`tf-kernel`。上游运行时要求 PyTorch 2.10+、CUDA 12.8+ 和 Triton 3.6+。可选 CuTe DSL runtime
+可用时选择专用内核，否则 SM80+ 使用 Triton 实现。
+
 ### 安装 Sparge Attention
 
 要使用 `SPARGE_ATTN` 后端或径向注意力中的稀疏 sage attention，需要从源码安装 `spas_sage_attn`：
@@ -426,6 +447,7 @@ print(f"FlashInfer: {FLASHINFER_AVAILABLE}")
 | Flash Attention 4 | 从源码编译（cute 接口） | SM90+ (H100, B100/B200) |
 | SageAttention | tf-kernel 或 [官方源码](https://github.com/thu-ml/SageAttention) | SM80+ |
 | Radial Attention | tf-kernel 或 [FlashInfer 源码](https://github.com/flashinfer-ai/flashinfer) | SM80+ |
+| Sol-Attn | TeleFuser 内置 | SM80+；支持的 SM90/100 使用 CuTe，否则回退 Triton |
 | Block Sparse | tf-kernel 或 [官方源码](https://github.com/mit-han-lab/Block-Sparse-Attention) | SM80+ |
 | Sparge Attention | 从源码安装（见上文） | SM80, SM86, SM89, SM90 |
 

examples/wan_video/README.md

@@ -163,6 +163,17 @@ python examples/wan_video/wan21_1_3b_text_to_video_radial.py \
 - Reduced memory usage for long videos
 - Requires flashinfer or sageattention backend
 
+Wan2.1 also supports Sol-Attn through the same attention configuration:
+
+```python
+from telefuser.core.config import AttentionConfig
+
+pipe_config.dit_config.attention_config = AttentionConfig.sol_attention()
+```
+
+Sol-Attn is built into TeleFuser. Eligible BF16 self-attention calls use the sparse kernel; unsupported calls
+automatically use the existing dense fallback. The defaults follow the official Wan2.1 profile: Morton3D token ordering, dense layer 0, and 10 dense warm-up steps for the standard 50-step schedule.
+
 #### wan21_1_3b_text_to_video_cache_calibrate.py
 
 Calibration tool for AdaTaylorCache.

pyproject.toml

@@ -89,6 +89,7 @@ dev = [
     "pytest-cov>=4.0.0",
     "pre-commit==4.0.1",
     "scikit-image>=0.19.0",
+    "tomli>=2.0.0; python_version < '3.11'",
     "uv"
 ]
 
@@ -128,6 +129,12 @@ telefuser = "telefuser.entrypoints.cli.main:main"
 where = ["."]
 include = ["telefuser*"]
 
+[tool.setuptools.package-data]
+"telefuser.kernel.sol_attn" = [
+    "THIRD_PARTY_NOTICES.md",
+    "sm100/LICENSE.flash-attention",
+]
+
 [tool.pytest.ini_options]
 minversion = "7.0"
 testpaths = ["tests"]
@@ -204,7 +211,8 @@ exclude = [
     "dist",
     "tf-kernel",
     "benchmarks",
-    "telefuser/_version.py"
+    "telefuser/_version.py",
+    "telefuser/kernel/sol_attn",
 ]
 
 [tool.ruff.lint]

telefuser/core/config.py

@@ -130,23 +130,35 @@ class AttnImplType(Enum):
     # Sparse attention implementations
     RADIAL_ATTN = auto()  # Radial attention for video generation
     LOCAL_SPARSE_ATTN = auto()  # Local window sparse attention
+    SOL_ATTN = auto()  # Dynamic on-the-fly sparse attention for video generation
 
 
 @dataclass
 class SparseAttentionConfig:
     """Configuration for sparse attention implementations.
 
-    Used with radial or local sparse attention to reduce memory usage
+    Used with radial, local, or Sol sparse attention to reduce memory usage
     for long sequences like videos.
     """
 
-    sparse_impl: str | None = None  # "radial", "local", or None
+    sparse_impl: str | None = None  # "radial", "local", "sol", or None
     dense_timesteps: int = 40  # Initial timesteps to use dense attention
     dense_layers: int = 0  # Initial layers to use dense attention
     decay_factor: float = 1.0  # Decay for radial attention window
     local_window_size: int = 6  # Window size for local attention
     block_size: int = 128  # Block size for sparse computation
     use_sage_attention: bool = False  # Use sage attention backend
+    sol_tau: float = 1.0  # Sol-Attn routing threshold multiplier
+    sol_threshold_type: str = "diag"  # Sol-Attn threshold estimator: "diag" or "exact"
+    sol_kv_splits: int | str = "auto"  # Auto selects split 4 for long SM90 sequences
+
+    def __post_init__(self) -> None:
+        if self.sparse_impl != "sol":
+            return
+        if self.sol_threshold_type not in ("diag", "exact"):
+            raise ValueError("Sol-Attn threshold type must be 'diag' or 'exact'")
+        if self.sol_kv_splits not in ("auto", 1, 2, 4):
+            raise ValueError("Sol-Attn KV splits must be 'auto', 1, 2, or 4")
 
     def should_use_dense(self, numeral_timestep: int, layer_idx: int) -> bool:
         """Check if dense attention should be used for current step/layer.
@@ -209,14 +221,42 @@ def local_sparse_attention(
             **kwargs,
         )
 
+    @classmethod
+    def sol_attention(
+        cls,
+        dense_timesteps: int = 10,
+        dense_layers: int = 1,
+        tau: float = 1.0,
+        threshold_type: str = "diag",
+        kv_splits: int | str = "auto",
+        **kwargs: any,
+    ) -> AttentionConfig:
+        """Create a Sol-Attn config for dynamic sparse video self-attention."""
+        return cls(
+            attn_impl=AttnImplType.SOL_ATTN,
+            sparse_config=SparseAttentionConfig(
+                sparse_impl="sol",
+                dense_timesteps=dense_timesteps,
+                dense_layers=dense_layers,
+                sol_tau=tau,
+                sol_threshold_type=threshold_type,
+                sol_kv_splits=kv_splits,
+            ),
+            **kwargs,
+        )
+
     @classmethod
     def dense_attention(cls, attn_impl: AttnImplType = AttnImplType.FLASH_ATTN_2, **kwargs: any) -> AttentionConfig:
         """Create config for dense attention."""
         return cls(attn_impl=attn_impl, sparse_config=None, **kwargs)
 
     def is_sparse(self) -> bool:
-        """Check if using sparse attention (radial or local)."""
-        return self.attn_impl in (AttnImplType.RADIAL_ATTN, AttnImplType.LOCAL_SPARSE_ATTN)
+        """Check if using a sparse attention implementation."""
+        return self.attn_impl in (
+            AttnImplType.RADIAL_ATTN,
+            AttnImplType.LOCAL_SPARSE_ATTN,
+            AttnImplType.SOL_ATTN,
+        )
 
     def should_use_dense(self, numeral_timestep: int, layer_idx: int) -> bool:
         """Check if dense attention should be used for current step/layer."""

telefuser/kernel/sol_attn/THIRD_PARTY_NOTICES.md

@@ -0,0 +1,23 @@
+# Third-party notices
+
+Sol-Attn was vendored from NVIDIA's `NVlabs/Sana` `sol-engine` branch at
+commit `8a26fb0ec9e353125ead798cb2e312d5ce48cded`. The upstream repository
+declares its code under the Apache License 2.0. Local changes move the
+implementation under the internal `telefuser.kernel.sol_attn` namespace and
+rewrite its absolute imports.
+
+Upstream source: https://github.com/NVlabs/Sana/tree/sol-engine/techniques/sparse_backends/sol_attn
+
+The files under `telefuser/kernel/sol_attn/_vendor/flash_attn/cute/` and portions of the SM90
+and SM100 design scaffold derive from the FlashAttention project. Its
+BSD-3-Clause license is included at
+`telefuser/kernel/sol_attn/sm100/LICENSE.flash-attention`.
+
+The runtime also depends on NVIDIA CUTLASS / CuTe DSL, cuda-python, PyTorch,
+and Triton. Those dependencies are not redistributed by this repository and
+remain subject to their respective licenses.
+
+The SM120 warp-MMA/TMA execution skeleton and online-softmax helpers are
+adapted from NVIDIA cuDNN Frontend's block-sparse-attention reference at commit
+`74785165de2da954a2c879a5e3e6f95411c2292d`. That source is licensed under the
+Apache License 2.0; adapted files retain the corresponding SPDX header.

telefuser/kernel/sol_attn/__init__.py

@@ -0,0 +1,5 @@
+"""Sol-Attn."""
+
+from .interface import sol_attn
+
+__all__ = ["sol_attn"]

telefuser/kernel/sol_attn/_vendor/__init__.py

@@ -0,0 +1 @@
+"""Private source dependencies bundled with Sol-Attn."""

telefuser/kernel/sol_attn/_vendor/flash_attn/__init__.py

@@ -0,0 +1 @@
+"""Local FlashAttention Cute shim for the SOL_ATTN SM90 release."""

telefuser/kernel/sol_attn/_vendor/flash_attn/cute/__init__.py

@@ -0,0 +1 @@
+"""Vendored FlashAttention Cute Python helpers used by SOL_ATTN SM90."""