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

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,
+    ),
 )