From 39d4d65a518b007465c02e5916d40eb80a9d7a14 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Fri, 21 Aug 2026 18:05:43 +0800 Subject: [PATCH] [diffusion] Accelerate SANA-Video linear attention in quality=high (#35728) --- docs/docs/sglang-diffusion/fused_kernels.mdx | 2 + python/sglang/kernels/ops/diffusion/README.md | 4 + .../sglang/kernels/ops/diffusion/__init__.py | 5 ++ .../sites/sana_video_linear_attention_site.py | 83 +++++++++++++++++++ .../SKILL.md | 3 + .../benchmark-and-profile.md | 2 +- .../runtime/models/dits/sana_video.py | 19 +++-- .../pipelines_core/stages/denoising.py | 7 ++ .../kernels/ops/diffusion/test_sites.py | 44 ++++++++++ 9 files changed, 163 insertions(+), 6 deletions(-) create mode 100644 python/sglang/kernels/ops/diffusion/sites/sana_video_linear_attention_site.py diff --git a/docs/docs/sglang-diffusion/fused_kernels.mdx b/docs/docs/sglang-diffusion/fused_kernels.mdx index 3f1209e50..ee3a9ec3f 100644 --- a/docs/docs/sglang-diffusion/fused_kernels.mdx +++ b/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) | diff --git a/python/sglang/kernels/ops/diffusion/README.md b/python/sglang/kernels/ops/diffusion/README.md index 6de12b39b..ba8f053a2 100644 --- a/python/sglang/kernels/ops/diffusion/README.md +++ b/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: diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py index 3310fad5f..59e1a9113 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/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", diff --git a/python/sglang/kernels/ops/diffusion/sites/sana_video_linear_attention_site.py b/python/sglang/kernels/ops/diffusion/sites/sana_video_linear_attention_site.py new file mode 100644 index 000000000..fb4966bf0 --- /dev/null +++ b/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 diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/SKILL.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/SKILL.md index 6a16fd4f6..d35121d2c 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/SKILL.md +++ b/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, diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/benchmark-and-profile.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/benchmark-and-profile.md index ea03164d3..5d9b8ee2a 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/benchmark-and-profile.md +++ b/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. | diff --git a/python/sglang/multimodal_gen/runtime/models/dits/sana_video.py b/python/sglang/multimodal_gen/runtime/models/dits/sana_video.py index f8f1765a3..8bdd63d27 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/sana_video.py +++ b/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) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py index 3fa92758c..8bfb544ed 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising.py +++ b/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, + ), ) diff --git a/test/registered/kernels/ops/diffusion/test_sites.py b/test/registered/kernels/ops/diffusion/test_sites.py index 75973a0fe..866185a49 100644 --- a/test/registered/kernels/ops/diffusion/test_sites.py +++ b/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"]))