[diffusion] Accelerate SANA-Video linear attention in quality=high (#35728)

docs/docs/sglang-diffusion/fused_kernels.mdx

@@ -57,6 +57,7 @@ These fusion families mount under `quality="high"`:
 | LTX-2 RMSNorm + modulate | `rms_norm(x) * (1 + scale) + shift` in one launch |
 | Gate RMSNorm (BF16-native) | `RMSNorm + tanh + mul + add` in one pass |
 | HunyuanVideo strided QK RMSNorm | Per-head QK RMSNorm over the packed QKV layout |
+| SANA-Video BF16-input linear attention | Keeps the first linear-attention GEMM's inputs in BF16 with FP32 accumulation/output; the second GEMM remains FP32 |
 
 ## Kernel inventory
 
@@ -146,6 +147,7 @@ Kernels are written against a specific eager chain in a specific model, so cover
 | LTX-2.5 decoder | paired 3D RoPE with shared axis-table cache |
 | HunyuanVideo | QKV+RoPE pack, strided QK RMSNorm, linear+GELU |
 | Sana | LN+modulate, GLUMB bias+SiLU / bias+GLU, residual-gate add |
+| SANA-Video | Packed QKV/KV; BF16-input linear attention at `quality=high` |
 | Sana-WM | bidirectional gated delta-net, fused QK inverse-RMS |
 | Wan | temb table slices; VAE cat+pad and DupUp3D add, `channels_last_3d` RMSNorm+SiLU |
 | Cosmos3 / Krea2 / MiniMax-H3 | QK-norm + RoPE (Krea2 also CuTe-DSL norm+scale/shift; MiniMax-H3 also indexed modulation) |

python/sglang/kernels/ops/diffusion/README.md

@@ -63,6 +63,10 @@ looks harmless and is not: on ERNIE-Image it moved the 50-step trajectory to
 PSNR 18.83 dB at `quality=high`, which is what motivated the bit-exact
 rewrite.
 
+SANA-Video's quality-gated linear-attention site keeps BF16 inputs for the
+first GEMM while requesting FP32 accumulation/output, then runs the second
+GEMM in FP32. The default path still promotes Q/K/V before both GEMMs.
+
 ## Entry-point protocol
 
 Every public kernel is a **predicate + kernel** pair:

python/sglang/kernels/ops/diffusion/__init__.py

@@ -460,6 +460,11 @@ _EXPORTS: dict[str, str] = {
     "mark_ltx2_rms_norm_modulate_site": "sites.ltx2_rmsnorm_modulate_site",
     "mount_ltx2_rms_norm_modulate": "sites.ltx2_rmsnorm_modulate_site",
     "unmount_ltx2_rms_norm_modulate": "sites.ltx2_rmsnorm_modulate_site",
+    "mark_sana_video_linear_attention_site": "sites.sana_video_linear_attention_site",
+    "mount_sana_video_linear_attention": "sites.sana_video_linear_attention_site",
+    "sana_video_linear_attention_active": "sites.sana_video_linear_attention_site",
+    "try_sana_video_linear_attention": "sites.sana_video_linear_attention_site",
+    "unmount_sana_video_linear_attention": "sites.sana_video_linear_attention_site",
     "QualityGatedFusion": "sites.quality_gate",
     # JIT C++/CUDA extensions (not kernels, not in the registry)
     "interpolate": "ext.hunyuan3d_rasterizer",

python/sglang/kernels/ops/diffusion/sites/sana_video_linear_attention_site.py

@@ -0,0 +1,83 @@
+"""SANA-Video BF16-input linear attention, gated by request quality.
+
+The reference path promotes rotated Q/K and V to FP32 before both attention
+GEMMs.  For ``quality="high"``, the first GEMM keeps its BF16 inputs while
+requesting an FP32 output/accumulator from cuBLAS.  The second GEMM stays in
+FP32.  This removes two large dtype-conversion kernels and lets the first GEMM
+use BF16 Tensor Cores, at the cost of half-precision input rounding.
+
+The default ``quality="lossless"`` path remains the original FP32-input chain
+bit-for-bit.  Only the single-batch CUDA layout used by native SANA-Video is
+eligible; unsupported dtypes and layouts fall back to the reference path.
+"""
+
+from __future__ import annotations
+
+import logging
+
+import torch
+import torch.nn as nn
+
+from sglang.kernels.ops.diffusion.sites.quality_gate import QualityGatedFusion
+
+logger = logging.getLogger(__name__)
+
+_FUSION = QualityGatedFusion(
+    name="SANA-Video BF16-input linear attention",
+    marker_attr="_sgl_sana_video_linear_attention_site",
+    enabled_attr="_sgl_sana_video_linear_attention_enabled",
+)
+
+
+def mark_sana_video_linear_attention_site(module: nn.Module) -> None:
+    """Mark a SANA-Video linear-attention module; it starts unmounted."""
+    _FUSION.mark(module)
+
+
+def sana_video_linear_attention_active(module: nn.Module) -> bool:
+    """Whether the request-scoped BF16-input path is mounted on ``module``."""
+    return _FUSION.is_enabled(module)
+
+
+def _site_reject_reason(_site: nn.Module) -> str | None:
+    if torch.version.cuda is None:
+        return "CUDA is unavailable"
+    return None
+
+
+def mount_sana_video_linear_attention(root: nn.Module) -> bool:
+    return _FUSION.mount(root, reject_reason=_site_reject_reason, logger=logger)
+
+
+def unmount_sana_video_linear_attention(root: nn.Module) -> None:
+    _FUSION.unmount(root)
+
+
+def try_sana_video_linear_attention(
+    site: nn.Module,
+    query_rotate: torch.Tensor,
+    key_rotate: torch.Tensor,
+    value: torch.Tensor,
+    normalizer: torch.Tensor,
+) -> torch.Tensor | None:
+    """Return the quality-gated attention result, or ``None`` to fall back."""
+    if not (
+        _FUSION.is_enabled(site)
+        and query_rotate.is_cuda
+        and query_rotate.dtype == torch.bfloat16
+        and key_rotate.dtype == query_rotate.dtype
+        and value.dtype == query_rotate.dtype
+        and query_rotate.dim() == 4
+        and query_rotate.shape == key_rotate.shape == value.shape
+        and query_rotate.shape[0] == 1
+        and normalizer.is_cuda
+    ):
+        return None
+
+    batch_size, num_heads, head_dim, _ = value.shape
+    scores = torch.bmm(
+        value.flatten(0, 1),
+        key_rotate.transpose(-1, -2).flatten(0, 1),
+        out_dtype=torch.float32,
+    ).view(batch_size, num_heads, head_dim, head_dim)
+    return (scores @ query_rotate.float()) * normalizer

python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/SKILL.md

@@ -62,6 +62,9 @@ Always rule out these existing families first:
 - LTX upsampler GroupNorm+SiLU
 - Z-Image bf16-native Triton RMSNorm scale/tanh-residual modulation
 - SANA packed self-attention Q/K/V and cross-attention K/V GEMMs
+- SANA-Video's packed projections and request-scoped BF16-input linear
+  attention at `quality=high`; keep the second attention GEMM in FP32 and
+  compare against `quality=lossless` before changing its precision further
 - SANA-Video reuse of SANA's bit-exact bias/activation, residual-gate, and
   LayerNorm-modulation fast paths before adding video-only kernels
 - MiniMax-H3 indexed modulation, fused QK norm + RoPE, packed Ulysses QKV,

python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/benchmark-and-profile.md

@@ -240,7 +240,7 @@ Use the preset categories this way:
 | `wan-i2v` | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | Yes: `wan22_i2v_a14b_720p` | Nightly cat image and motion prompt, 1280x720, 81 frames, 4 GPUs, CFG parallel, Ulysses degree 2, text encoder CPU offload and pinned CPU memory |
 | `minimax-h3-t2va` | `MiniMaxAI/MiniMax-H3` | Yes: `minimax_h3_t2va_5s` | H3 FL2VA-partition T2VA baseline: 1344x768 resolved canvas, 5 seconds / 124 frames at 24 fps, 50 joint video-audio steps, 4 GPUs, TP2 + Ulysses2, eager BF16/FP32. The helper writes H3's request contract to a generated config. |
 | `longcat-image` | `meituan-longcat/LongCat-Image` | No | Eager DiT baseline at 1024x1024, 50 steps, guidance 4.5; prompt rewrite is disabled so Qwen2.5-VL does not contaminate the DiT A/B. |
-| `sana-video` | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers` | No | CI-sized eager T2V baseline: 832x480, 17 frames, 8 steps, guidance 6.0. |
+| `sana-video` | `Efficient-Large-Model/SANA-Video_2B_480p_diffusers` | No | CI-sized eager T2V baseline: 832x480, 17 frames, 8 steps, guidance 6.0. Compare `quality=lossless` and `quality=high`; high enables the BF16-input first linear-attention GEMM while retaining FP32 output and the FP32 second GEMM. |
 | `lingbot-video-moe` | `robbyant/lingbot-video-moe-30b-a3b` | No | One-GPU eager baseline using the CI structured-JSON caption, 384x640, 17 frames, 12 steps, and text-encoder CPU offload. |
 | `cosmos3-edge-t2i` | `nvidia/Cosmos3-Edge` | No | One-GPU eager T2I baseline at Edge's native 640x640 shape, 35 steps, guidance 7.0. |
 | `cosmos3-super-t2i-distilled` | `nvidia/Cosmos3-Super-Text2Image-4Step` | No | Four-GPU eager distilled T2I baseline. The checkpoint owns its fixed sigma schedule; the preset does not override the step count. |

python/sglang/multimodal_gen/runtime/models/dits/sana_video.py

@@ -9,6 +9,10 @@ import torch.nn as nn
 import torch.nn.functional as F
 from diffusers.models.embeddings import PixArtAlphaTextProjection
 
+from sglang.kernels.ops.diffusion import (
+    mark_sana_video_linear_attention_site,
+    try_sana_video_linear_attention,
+)
 from sglang.multimodal_gen.configs.models.dits.sana_video import SanaVideoConfig
 from sglang.multimodal_gen.runtime.layers.layernorm import RMSNorm
 from sglang.multimodal_gen.runtime.layers.linear import MergedColumnParallelLinear
@@ -186,6 +190,7 @@ class SanaVideoLinearAttention(nn.Module):
         self.to_out = nn.ModuleList(
             [nn.Linear(self.inner_dim, query_dim, bias=True), nn.Identity()]
         )
+        mark_sana_video_linear_attention_site(self)
 
     def forward(
         self,
@@ -211,15 +216,19 @@ class SanaVideoLinearAttention(nn.Module):
 
         query = query.permute(0, 2, 3, 1)
         key = key.permute(0, 2, 3, 1)
-        query_rotate = query_rotate.permute(0, 2, 3, 1).float()
-        key_rotate = key_rotate.permute(0, 2, 3, 1).float()
-        value = value.permute(0, 2, 3, 1).float()
+        query_rotate = query_rotate.permute(0, 2, 3, 1)
+        key_rotate = key_rotate.permute(0, 2, 3, 1)
+        value = value.permute(0, 2, 3, 1)
 
         normalizer = 1.0 / (
             key.sum(dim=-1, keepdim=True).transpose(-2, -1) @ query + 1e-15
         )
-        scores = value @ key_rotate.transpose(-1, -2)
-        hidden_states = (scores @ query_rotate) * normalizer
+        hidden_states = try_sana_video_linear_attention(
+            self, query_rotate, key_rotate, value, normalizer
+        )
+        if hidden_states is None:
+            scores = value.float() @ key_rotate.float().transpose(-1, -2)
+            hidden_states = (scores @ query_rotate.float()) * normalizer
         hidden_states = hidden_states.flatten(1, 2).transpose(1, 2)
         hidden_states = hidden_states.to(original_dtype)
         return self.to_out[0](hidden_states)

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py

@@ -26,11 +26,13 @@ from sglang.kernels.ops.diffusion import (
     mount_fused_ln_modulate,
     mount_hunyuan_qknorm,
     mount_ltx2_rms_norm_modulate,
+    mount_sana_video_linear_attention,
     unmount_fused_gate_rmsnorm,
     unmount_fused_linear_gelu,
     unmount_fused_ln_modulate,
     unmount_hunyuan_qknorm,
     unmount_ltx2_rms_norm_modulate,
+    unmount_sana_video_linear_attention,
 )
 from sglang.multimodal_gen import envs
 from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType, STA_Mode
@@ -182,6 +184,11 @@ _QUALITY_FUSION_HANDLERS: tuple[
         mount_hunyuan_qknorm,
         unmount_hunyuan_qknorm,
     ),
+    (
+        "SANA-Video BF16-input linear attention",
+        mount_sana_video_linear_attention,
+        unmount_sana_video_linear_attention,
+    ),
 )
 
 

test/registered/kernels/ops/diffusion/test_sites.py

@@ -27,6 +27,7 @@ import torch.nn.functional as F
 
 import sglang.kernels.ops.diffusion.sites.fused_gate_rmsnorm_site as gate_rmsnorm
 import sglang.kernels.ops.diffusion.sites.fused_linear_gelu_site as linear_gelu
+import sglang.kernels.ops.diffusion.sites.sana_video_linear_attention_site as sana_video_linear_attention
 from sglang.kernels.ops.diffusion import (
     BitExactFusionGate,
     QualityGatedFusion,
@@ -412,5 +413,48 @@ def test_mounted_gelu_site_compiles_fullgraph():
     )
 
 
+@requires_cuda
+@torch.no_grad()
+def test_sana_video_linear_attention_quality_path_and_guards():
+    torch.manual_seed(0)
+    site = nn.Module()
+    sana_video_linear_attention.mark_sana_video_linear_attention_site(site)
+    shape = (1, 4, 16, 128)
+    query = torch.randn(shape, device="cuda", dtype=torch.bfloat16)
+    key = torch.randn_like(query)
+    value = torch.randn_like(query)
+    normalizer = torch.randn(
+        shape[0], shape[1], 1, shape[-1], device="cuda", dtype=torch.bfloat16
+    )
+
+    assert (
+        sana_video_linear_attention.try_sana_video_linear_attention(
+            site, query, key, value, normalizer
+        )
+        is None
+    )
+    assert sana_video_linear_attention.mount_sana_video_linear_attention(site)
+    output = sana_video_linear_attention.try_sana_video_linear_attention(
+        site, query, key, value, normalizer
+    )
+    reference = ((value.float() @ key.float().transpose(-1, -2)) @ query.float()) * (
+        normalizer
+    )
+    torch.testing.assert_close(output, reference, atol=1e-2, rtol=1e-2)
+
+    assert (
+        sana_video_linear_attention.try_sana_video_linear_attention(
+            site,
+            query.expand(2, -1, -1, -1),
+            key.expand(2, -1, -1, -1),
+            value.expand(2, -1, -1, -1),
+            normalizer.expand(2, -1, -1, -1),
+        )
+        is None
+    )
+    sana_video_linear_attention.unmount_sana_video_linear_attention(site)
+    assert not sana_video_linear_attention.sana_video_linear_attention_active(site)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__, "-v"]))