[diffusion] Fuse SANA-Video interleaved RoPE (#35695)

python/sglang/kernels/jit/csrc/diffusion/interleaved_rope_fp64.cuh

@@ -0,0 +1,148 @@
+#pragma once
+
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+
+#include <algorithm>
+#include <cstdint>
+#include <type_traits>
+
+namespace sglang {
+
+namespace interleaved_rope_fp64 {
+
+namespace {
+
+constexpr uint32_t kBlockSize = 256;
+constexpr uint32_t kMaxGrid = 65535;
+
+template <typename T>
+SGL_DEVICE device::AlignedVector<T, 2> rotate_pair(device::AlignedVector<T, 2> input, double cos, double sin) {
+  static_assert(std::is_same_v<T, bf16_t>);
+  const double x1 = static_cast<double>(input[0]);
+  const double x2 = static_cast<double>(input[1]);
+  const double x1_cos = __dmul_rn(x1, cos);
+  const double x2_sin = __dmul_rn(x2, sin);
+  const double x1_sin = __dmul_rn(x1, sin);
+  const double x2_cos = __dmul_rn(x2, cos);
+  device::AlignedVector<T, 2> output;
+  // TensorIterator casts the fp64 expression through fp32 before its bf16
+  // store. Preserving both conversions is required at bf16 tie boundaries.
+  output[0] = static_cast<T>(static_cast<float>(__dsub_rn(x1_cos, x2_sin)));
+  output[1] = static_cast<T>(static_cast<float>(__dadd_rn(x1_sin, x2_cos)));
+  return output;
+}
+
+template <typename T>
+__global__ void interleaved_rope_fp64_kernel(
+    T* __restrict__ q_out,
+    T* __restrict__ k_out,
+    const T* __restrict__ q,
+    const T* __restrict__ k,
+    const double* __restrict__ cos,
+    const double* __restrict__ sin,
+    int64_t num_pairs,
+    int64_t seq_len,
+    int64_t num_heads,
+    int64_t pairs_per_head,
+    int64_t head_dim) {
+  using Pair = device::AlignedVector<T, 2>;
+  const int64_t stride = static_cast<int64_t>(gridDim.x) * blockDim.x;
+  for (int64_t pair_index = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x; pair_index < num_pairs;
+       pair_index += stride) {
+    const int64_t pair_in_head = pair_index % pairs_per_head;
+    const int64_t row = pair_index / (num_heads * pairs_per_head);
+    const int64_t token = row % seq_len;
+    const int64_t table_index = token * head_dim + 2 * pair_in_head;
+    const double cos_value = SGLANG_LDG(cos + table_index);
+    const double sin_value = SGLANG_LDG(sin + table_index + 1);
+
+    Pair q_pair;
+    q_pair.load(q, pair_index);
+    rotate_pair(q_pair, cos_value, sin_value).store(q_out, pair_index);
+
+    Pair k_pair;
+    k_pair.load(k, pair_index);
+    rotate_pair(k_pair, cos_value, sin_value).store(k_out, pair_index);
+  }
+}
+
+}  // namespace
+
+/**
+ * \brief Apply Diffusers-compatible interleaved RoPE to Q and K.
+ *
+ * The fp64 multiply and add/subtract roundings deliberately mirror the eager
+ * PyTorch expression used by SANA-Video.
+ *
+ * \tparam T Activation type; currently bf16 only.
+ */
+template <typename T>
+struct InterleavedRopeFP64Kernel {
+  static_assert(std::is_same_v<T, bf16_t>);
+
+  static void
+  run(tvm::ffi::TensorView q_out,
+      tvm::ffi::TensorView k_out,
+      tvm::ffi::TensorView q,
+      tvm::ffi::TensorView k,
+      tvm::ffi::TensorView cos,
+      tvm::ffi::TensorView sin,
+      int64_t batch_size,
+      int64_t seq_len,
+      int64_t num_heads,
+      int64_t head_dim) {
+    using namespace host;
+    using Pair = device::AlignedVector<T, 2>;
+
+    auto N = SymbolicSize{"activation_elements"};
+    auto R = SymbolicSize{"table_elements"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+    TensorMatcher({N}).with_dtype<T>().with_device(device).verify(q_out).verify(k_out).verify(q).verify(k);
+    TensorMatcher({R}).with_dtype<double>().with_device(device).verify(cos).verify(sin);
+
+    CHECK_HOST(batch_size > 0 && seq_len > 0 && num_heads > 0 && head_dim > 0)
+        << "interleaved_rope_fp64 dimensions must be positive";
+    CHECK_HOST(head_dim % 2 == 0) << "interleaved_rope_fp64 head_dim must be even";
+    CHECK_HOST(N.unwrap() == batch_size * seq_len * num_heads * head_dim)
+        << "interleaved_rope_fp64 activation shape does not match dimensions";
+    CHECK_HOST(R.unwrap() == seq_len * head_dim) << "interleaved_rope_fp64 table shape does not match dimensions";
+    CHECK_HOST(
+        q_out.data_ptr() != k_out.data_ptr() && q_out.data_ptr() != q.data_ptr() && q_out.data_ptr() != k.data_ptr() &&
+        k_out.data_ptr() != q.data_ptr() && k_out.data_ptr() != k.data_ptr())
+        << "interleaved_rope_fp64 outputs must not alias inputs";
+    CHECK_HOST(
+        reinterpret_cast<uintptr_t>(q_out.data_ptr()) % alignof(Pair) == 0 &&
+        reinterpret_cast<uintptr_t>(k_out.data_ptr()) % alignof(Pair) == 0 &&
+        reinterpret_cast<uintptr_t>(q.data_ptr()) % alignof(Pair) == 0 &&
+        reinterpret_cast<uintptr_t>(k.data_ptr()) % alignof(Pair) == 0)
+        << "interleaved_rope_fp64 activations must be aligned to rotation pairs";
+
+    const int64_t num_pairs = N.unwrap() / 2;
+    const int64_t pairs_per_head = head_dim / 2;
+    const auto blocks =
+        static_cast<uint32_t>(std::min<int64_t>(div_ceil(num_pairs, static_cast<int64_t>(kBlockSize)), kMaxGrid));
+    LaunchKernel(blocks, kBlockSize, device.unwrap())(
+        interleaved_rope_fp64_kernel<T>,
+        static_cast<T*>(q_out.data_ptr()),
+        static_cast<T*>(k_out.data_ptr()),
+        static_cast<const T*>(q.data_ptr()),
+        static_cast<const T*>(k.data_ptr()),
+        static_cast<const double*>(cos.data_ptr()),
+        static_cast<const double*>(sin.data_ptr()),
+        num_pairs,
+        seq_len,
+        num_heads,
+        pairs_per_head,
+        head_dim);
+  }
+};
+
+}  // namespace interleaved_rope_fp64
+
+}  // namespace sglang

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

@@ -117,6 +117,7 @@ Several norms look interchangeable and are not. Start here.
 | `fused_inplace_qknorm_rope` | JIT CUDA | one bf16 rounding step vs split baseline; `round_norm_before_rope=True` makes it exact |
 | `fused_qknorm_rope_pack_kv` | JIT CUDA | as above, also packs prefix K/V |
 | `fused_rope_rotate_half_bitexact` | Triton | bit-exact (elementwise only) |
+| `fused_interleaved_rope_fp64` | JIT CUDA | bit-exact vs paired SANA-Video fp64 RoPE |
 | `ltx2_qknorm_split_rope_cuda` | JIT CUDA | close; **validated on B200** |
 | `fused_ltx25_decoder_rope` | JIT CUDA | bit-exact paired 3D RoPE from cached compact axis tables |
 | `apply_rotary_embedding` | Triton (+fallbacks) | close; the generic entry point |

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

@@ -214,6 +214,13 @@ _SPECS: tuple[tuple[str, KernelBackend, str, frozenset, str], ...] = (
         _CUDA,
         "Bit-exact rotate-half RoPE.",
     ),
+    (
+        "diffusion.interleaved_rope_fp64",
+        KernelBackend.JIT,
+        "rope.interleaved_rope_fp64_jit:fused_interleaved_rope_fp64",
+        _CUDA,
+        "Paired interleaved RoPE with fp64 Diffusers semantics.",
+    ),
     (
         "diffusion.hunyuan_qkv_rope_pack",
         KernelBackend.TRITON,
@@ -400,6 +407,8 @@ _EXPORTS: dict[str, str] = {
     "fused_qknorm_rope_pack_kv": "rope.qknorm_rope_jit",
     "can_use_fused_rope_rotate_half": "rope.rope_rotate_half_bitexact",
     "fused_rope_rotate_half_bitexact": "rope.rope_rotate_half_bitexact",
+    "can_use_interleaved_rope_fp64": "rope.interleaved_rope_fp64_jit",
+    "fused_interleaved_rope_fp64": "rope.interleaved_rope_fp64_jit",
     "apply_rotary_embedding": "rope.rotary_triton",
     # Activation-function fusions
     "can_use_fused_bias_glu": "activation.sana_conv_post_triton",

python/sglang/kernels/ops/diffusion/rope/interleaved_rope_fp64_jit.py

@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
+from sglang.srt.utils.custom_op import register_custom_op
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_interleaved_rope_fp64_module(dtype: torch.dtype) -> Module:
+    if dtype is not torch.bfloat16:
+        raise RuntimeError(f"Unsupported interleaved_rope_fp64 dtype: {dtype}")
+    args = make_cpp_args(dtype)
+    return load_jit(
+        "diffusion_interleaved_rope_fp64",
+        *args,
+        cuda_files=["diffusion/interleaved_rope_fp64.cuh"],
+        cuda_wrappers=[
+            (
+                "interleaved_rope_fp64",
+                f"interleaved_rope_fp64::InterleavedRopeFP64Kernel<{args}>::run",
+            ),
+        ],
+    )
+
+
+def _fake_impl(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    del cos, sin
+    return torch.empty_like(q), torch.empty_like(k)
+
+
+@register_custom_op(
+    op_name="diffusion_interleaved_rope_fp64",
+    mutates_args=[],
+    fake_impl=_fake_impl,
+)
+def fused_interleaved_rope_fp64(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Apply paired interleaved RoPE with fp64 Diffusers semantics."""
+    q_out = torch.empty_like(q)
+    k_out = torch.empty_like(k)
+    module = _jit_interleaved_rope_fp64_module(q.dtype)
+    module.interleaved_rope_fp64(
+        q_out.view(-1),
+        k_out.view(-1),
+        q.view(-1),
+        k.view(-1),
+        cos.view(-1),
+        sin.view(-1),
+        q.shape[0],
+        q.shape[1],
+        q.shape[2],
+        q.shape[3],
+    )
+    return q_out, k_out
+
+
+def can_use_interleaved_rope_fp64(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> bool:
+    if q.dim() != 4:
+        return False
+    expected_table_shape = (1, q.shape[1], 1, q.shape[3])
+    return (
+        q.dtype is torch.bfloat16
+        and k.dtype is q.dtype
+        and q.is_cuda
+        and k.is_cuda
+        and q.device == k.device == cos.device == sin.device
+        and k.shape == q.shape
+        and q.shape[-1] % 2 == 0
+        and q.is_contiguous()
+        and k.is_contiguous()
+        and q.data_ptr() % 4 == 0
+        and k.data_ptr() % 4 == 0
+        and cos.dtype is torch.float64
+        and sin.dtype is torch.float64
+        and cos.shape == expected_table_shape
+        and sin.shape == expected_table_shape
+        and cos.is_contiguous()
+        and sin.is_contiguous()
+    )
+
+
+__all__ = [
+    "can_use_interleaved_rope_fp64",
+    "fused_interleaved_rope_fp64",
+]

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

@@ -10,7 +10,11 @@ import torch.nn.functional as F
 from diffusers.models.embeddings import PixArtAlphaTextProjection
 
 from sglang.kernels.ops.diffusion import (
+    BitExactFusionGate,
+    can_use_interleaved_rope_fp64,
+    fused_interleaved_rope_fp64,
     mark_sana_video_linear_attention_site,
+    tensors_equal,
     try_sana_video_linear_attention,
 )
 from sglang.multimodal_gen.configs.models.dits.sana_video import SanaVideoConfig
@@ -24,6 +28,11 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload im
 )
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
 from sglang.multimodal_gen.runtime.models.dits.sana import SanaAdaLayerNormSingle
+from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
+
+logger = init_logger(__name__)
+
+_SANA_VIDEO_ROPE = BitExactFusionGate("SANA-Video fused RoPE")
 
 
 def apply_interleaved_rotary_emb(
@@ -41,6 +50,46 @@ def apply_interleaved_rotary_emb(
     return output
 
 
+def apply_interleaved_rotary_emb_pair(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    freqs_cos: torch.Tensor,
+    freqs_sin: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Apply paired interleaved RoPE, with a bit-exact eager fallback."""
+    verified = _SANA_VIDEO_ROPE.verified
+    if (
+        not _SANA_VIDEO_ROPE.disabled
+        and can_use_interleaved_rope_fp64(query, key, freqs_cos, freqs_sin)
+        and (verified or _SANA_VIDEO_ROPE.can_attempt_once())
+    ):
+        try:
+            fused = fused_interleaved_rope_fp64(query, key, freqs_cos, freqs_sin)
+        except Exception as exc:
+            _SANA_VIDEO_ROPE.on_exception(exc, logger=logger)
+        else:
+            if verified:
+                return fused
+            eager = (
+                apply_interleaved_rotary_emb(query, freqs_cos, freqs_sin),
+                apply_interleaved_rotary_emb(key, freqs_cos, freqs_sin),
+            )
+            return _SANA_VIDEO_ROPE.accept_or_fallback(
+                fused,
+                eager,
+                equal=tensors_equal,
+                logger=logger,
+                mismatch_msg=(
+                    "SANA-Video fused RoPE is not bit-exact on this platform; "
+                    "falling back to eager"
+                ),
+            )
+    return (
+        apply_interleaved_rotary_emb(query, freqs_cos, freqs_sin),
+        apply_interleaved_rotary_emb(key, freqs_cos, freqs_sin),
+    )
+
+
 class SanaVideoRotaryPosEmbed(nn.Module):
     """3D RoPE split across temporal, height, and width head dimensions."""
 
@@ -211,8 +260,9 @@ class SanaVideoLinearAttention(nn.Module):
 
         query = F.relu(query)
         key = F.relu(key)
-        query_rotate = apply_interleaved_rotary_emb(query, *rotary_emb)
-        key_rotate = apply_interleaved_rotary_emb(key, *rotary_emb)
+        query_rotate, key_rotate = apply_interleaved_rotary_emb_pair(
+            query, key, *rotary_emb
+        )
 
         query = query.permute(0, 2, 3, 1)
         key = key.permute(0, 2, 3, 1)

python/sglang/multimodal_gen/test/unit/test_sana_video.py

@@ -9,6 +9,8 @@ from sglang.multimodal_gen.configs.sample.sana_video import SanaVideoSamplingPar
 from sglang.multimodal_gen.registry import get_model_info
 from sglang.multimodal_gen.runtime.models.dits.sana_video import (
     SanaVideoRotaryPosEmbed,
+    apply_interleaved_rotary_emb,
+    apply_interleaved_rotary_emb_pair,
 )
 from sglang.multimodal_gen.runtime.pipelines.sana_video import (
     select_sana_video_prompt_window,
@@ -85,3 +87,15 @@ def test_sana_video_rotary_embeddings_follow_video_token_order():
     assert sin.shape == (1, 12, 1, 12)
     assert torch.isfinite(cos).all()
     assert torch.isfinite(sin).all()
+
+
+def test_sana_video_paired_rope_falls_back_to_eager_on_cpu():
+    query = torch.randn(2, 7, 3, 12, dtype=torch.bfloat16)
+    key = torch.randn_like(query)
+    cos = torch.randn(1, 7, 1, 12, dtype=torch.float64)
+    sin = torch.randn_like(cos)
+
+    query_out, key_out = apply_interleaved_rotary_emb_pair(query, key, cos, sin)
+
+    assert torch.equal(query_out, apply_interleaved_rotary_emb(query, cos, sin))
+    assert torch.equal(key_out, apply_interleaved_rotary_emb(key, cos, sin))