Use portable Sol-Attn backend

README.md

@@ -28,7 +28,7 @@ Blackwell-native ComfyUI optimization path:
   high-signal change probe before deciding whether blocks 2–50 can reuse their previous joint residual. An optional
   Ultra Fast mode forecasts between exact anchors, while Exact evaluates every requested step. This is not
   prompt-to-video output caching.
-- NVIDIA Sol-Attn's native SM120 kernel sparsifies target-video attention after a dense warmup while retaining text,
+- NVIDIA Sol-Attn's released Triton kernel sparsifies target-video attention after a dense warmup while retaining text,
   conditioning-video and generated-audio rows as an exact KV sink with dense prefix queries.
 - Segment-wise in-place AdaLN modulation and gated residual accumulation.
 - Each block converts its complete AdaLN table in one launch instead of six, and the final RMSNorm runs only on
@@ -103,7 +103,7 @@ does not install or launch the ComfyUI application. The small adapter uses only
 1. Dynamic NVFP4 activation quantization and native FP4 matrix multiplication.
 2. Fused in-place Q/K RMSNorm and three-axis split-half rotary embedding.
 
-Attention uses NVIDIA Sol-Attn's SM120 CuTe kernel on eligible target-video calls and keeps diffusers'
+Attention uses NVIDIA Sol-Attn's portable Triton kernel on eligible target-video calls and keeps diffusers'
 `_native_cudnn` backend for the first ten denoising steps, the first two blocks, Exact mode, short sequences and any
 safe fallback. The MLP uses one fused QKV-style gate/up matrix, in-place SiLU×up, and the NVFP4 down projection.
 
@@ -156,6 +156,7 @@ exact original denoiser, set `H3_ENGINE=bf16`; this restores the 61.7 GiB unquan
 | `H3_FORECAST_BLEND` | `0.65` | Ultra Fast linear-trend strength; lower values stay closer to last-residual reuse. |
 | `H3_FUSED_ADALN` | `0` | Opt in to fused AdaLN kernels; compilation is expensive on fresh ZeroGPU workers. |
 | `H3_SOL_ATTN` | `1` | Use NVIDIA Sol-Attn outside Exact mode, with automatic dense fallback. |
+| `H3_SOL_ATTN_BACKEND` | `triton` | Released portable backend; `auto` tries CuTe when its rapidly moving DSL/FFI ABI matches. |
 | `H3_SOL_ATTN_TAU` | `1.0` | Official H3 sparse-routing threshold. |
 | `H3_SOL_ATTN_DENSE_STEPS` | `10` | Initial denoising steps kept on exact dense attention. |
 | `H3_SOL_ATTN_DENSE_LAYERS` | `2` | Initial transformer blocks kept on exact dense attention. |

h3_nvfp4.py

@@ -60,6 +60,7 @@ FIRST_BLOCK_DENSE_END = max(1, int(os.environ.get("H3_FIRST_BLOCK_DENSE_END", "2
 FORECAST_BLEND = min(1.0, max(0.0, float(os.environ.get("H3_FORECAST_BLEND", "0.65"))))
 FUSED_ADALN = os.environ.get("H3_FUSED_ADALN", "0") == "1" and triton is not None
 SOL_ATTN = os.environ.get("H3_SOL_ATTN", "1") == "1"
+SOL_ATTN_BACKEND = os.environ.get("H3_SOL_ATTN_BACKEND", "triton").lower()
 SOL_ATTN_TAU = float(os.environ.get("H3_SOL_ATTN_TAU", "1.0"))
 SOL_ATTN_DENSE_STEPS = max(0, int(os.environ.get("H3_SOL_ATTN_DENSE_STEPS", "10")))
 SOL_ATTN_DENSE_LAYERS = max(0, int(os.environ.get("H3_SOL_ATTN_DENSE_LAYERS", "2")))
@@ -360,19 +361,21 @@ class H3SolAttention:
             self.dense_calls += 1
             return None
         try:
-            from sol_attn import sol_attn
+            if SOL_ATTN_BACKEND == "triton":
+                from sol_attn.triton_ref import sol_attn
+            else:
+                from sol_attn import sol_attn
 
             q, k, v = (tensor.contiguous() for tensor in (query, key, value))
-            attended = sol_attn(
-                q,
-                k,
-                v,
-                tau=SOL_ATTN_TAU,
-                thresh_type="diag",
-                kv_splits=1,
-                sink_start=0,
-                sink_tokens=self.video_start,
-            )
+            kwargs = {
+                "tau": SOL_ATTN_TAU,
+                "thresh_type": "diag",
+                "sink_start": 0,
+                "sink_tokens": self.video_start,
+            }
+            if SOL_ATTN_BACKEND != "triton":
+                kwargs["kv_splits"] = 1
+            attended = sol_attn(q, k, v, **kwargs)
             # An exact KV sink does not make the prefix's own queries dense. H3 jointly generates audio in that
             # prefix, so reproduce those rows with exact attention as NVIDIA's H3 integration does.
             prefix = self.video_start

requirements.txt

@@ -19,15 +19,9 @@ spaces==0.51.1
 # Blackwell-native NVFP4 GEMMs and the fused Q/K RMSNorm + split-half RoPE kernel used by h3_nvfp4.py.
 # CUDA 13 is mandatory: older builds emulate this path and are slower than BF16.
 comfy-kitchen==0.2.26
-# NVIDIA Sol-Engine's released sparse-attention package, pinned to the H3 release. The native SM120 path uses
-# CuTe DSL; the package retains its portable Triton fallback if that optional backend cannot initialize.
+# NVIDIA Sol-Engine's released sparse-attention package, pinned to the H3 release. The Space selects its portable
+# Triton backend because the separately released CuTe DSL/TVM-FFI packages currently have an incompatible ABI.
 sol-attn @ git+https://github.com/NVlabs/Sana.git@46031940ba8af5d18054217e571149579424c0b1#subdirectory=techniques/sparse_backends
-# Sol-Engine's Aug-03 H3 kernel targets the 4.5 `cute.math.fmax` signature; 4.6 changed that API.
-nvidia-cutlass-dsl==4.5.3
-cuda-python>=13.0
-# Sol-Attn compiles its CuTe callable with `--enable-tvm-ffi`, but its package currently leaves this transitive
-# runtime undeclared. Without it the SM120 path compiles up to dispatch and then falls back to dense attention.
-apache-tvm-ffi>=0.1.12
 # No `kernels` pin on purpose: the Hub attention backends want `kernels>=0.12.3`, and that version breaks
 # transformers 5.8.0 at import.
 # PyAV muxes the generated soundtrack onto the frames (`encode_video`).