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

2026-08-22 12:56:45 +08:00
parent 3e096629cf
commit 83e9ece672
9 changed files with 516 additions and 2 deletions
@@ -98,6 +98,7 @@ These fusion families mount under `quality="high"`:
 | --- | --- | --- | --- |
 | `fused_inplace_qknorm_rope` | JIT CUDA | one bf16 rounding step vs the split baseline; exact with `round_norm_before_rope=True` | separate QK-norm kernel + RoPE |
 | `rope_rotate_half` | Triton | bit-exact | `chunk` → `cat(-x2, x1)` → two muls + add → `cat(tail)`, about 7 kernels per projection |
+| `interleaved_rope_fp64` | JIT CUDA | bit-exact | paired SANA-Video Q/K RoPE with fp64 tables, about 14 eager kernels |
 | `ltx2_qknorm_split_rope` | JIT CUDA | close (validated on B200) | LTX-2 QK-norm + split RoPE |
 | `ltx25_decoder_rope` | JIT CUDA | bit-exact | paired LTX-2.5 decoder 3D RoPE from cached compact axis tables |
 | `hunyuan_qkv_rope_pack` | Triton | bit-exact | QKV pack and RoPE in one pass |
@@ -146,6 +147,7 @@ Kernels are written against a specific eager chain in a specific model, so cover
 | LTX-2 | QK-norm + split RoPE, ada-values split, RMSNorm+modulate, modulate, residual-gate add, linear+GELU |
 | LTX-2.5 decoder | paired 3D RoPE with shared axis-table cache |
 | HunyuanVideo | QKV+RoPE pack, strided QK RMSNorm, linear+GELU |
+| SANA-Video | paired fp64 interleaved RoPE |
 | 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 |
@@ -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
@@ -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 |
@@ -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",
@@ -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",
+]
@@ -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)
@@ -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))
@@ -0,0 +1,87 @@
+from dataclasses import dataclass
+
+import torch
+
+from sglang.kernels.jit.benchmark import marker
+from sglang.kernels.ops.diffusion import fused_interleaved_rope_fp64
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(
+    est_time=8, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)
+
+
+@dataclass(frozen=True)
+class Case:
+    name: str
+    batch: int
+    seq_len: int
+    num_heads: int
+    head_dim: int
+
+
+CASES = {
+    case.name: case
+    for case in (
+        Case("sana_video_480p", 2, 7800, 20, 112),
+        Case("sana_video_small", 2, 1920, 20, 112),
+    )
+}
+
+
+def eager_interleaved_rope(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    def apply(hidden_states: torch.Tensor) -> torch.Tensor:
+        x1, x2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1)
+        output = torch.empty_like(hidden_states)
+        output[..., 0::2] = x1 * cos[..., 0::2] - x2 * sin[..., 1::2]
+        output[..., 1::2] = x1 * sin[..., 1::2] + x2 * cos[..., 0::2]
+        return output
+
+    return apply(q), apply(k)
+
+
+FN_MAP = {
+    "eager": eager_interleaved_rope,
+    "jit": fused_interleaved_rope_fp64,
+}
+
+
+@marker.parametrize("case_name", list(CASES), ci_vals=["sana_video_480p"])
+@marker.benchmark("impl", ["eager", "jit"], unit="ms")
+def benchmark(case_name: str, impl: str) -> marker.BenchResult:
+    case = CASES[case_name]
+    generator = torch.Generator(device="cuda").manual_seed(42)
+    q = torch.randn(
+        case.batch,
+        case.seq_len,
+        case.num_heads,
+        case.head_dim,
+        dtype=torch.bfloat16,
+        device="cuda",
+        generator=generator,
+    )
+    k = torch.randn_like(q)
+    cos = torch.randn(
+        1,
+        case.seq_len,
+        1,
+        case.head_dim,
+        dtype=torch.float64,
+        device="cuda",
+        generator=generator,
+    )
+    sin = torch.randn_like(cos)
+    return marker.do_bench(
+        FN_MAP[impl],
+        input_args=(q, k, cos, sin),
+        disable_log_bandwidth=True,
+    )
+
+
+if __name__ == "__main__":
+    benchmark.run()
@@ -0,0 +1,98 @@
+import sys
+
+import pytest
+import torch
+
+from sglang.kernels.jit.utils import get_ci_test_range
+from sglang.kernels.ops.diffusion import (
+    can_use_interleaved_rope_fp64,
+    fused_interleaved_rope_fp64,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=35, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+def eager_interleaved_rope(
+    hidden_states: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> torch.Tensor:
+    x1, x2 = hidden_states.unflatten(-1, (-1, 2)).unbind(-1)
+    output = torch.empty_like(hidden_states)
+    output[..., 0::2] = x1 * cos[..., 0::2] - x2 * sin[..., 1::2]
+    output[..., 1::2] = x1 * sin[..., 1::2] + x2 * cos[..., 0::2]
+    return output
+
+
+CASES = get_ci_test_range(
+    [
+        (1, 1, 1, 2),
+        (1, 17, 3, 12),
+        (2, 129, 5, 112),
+        (2, 7800, 20, 112),
+    ],
+    [(1, 17, 3, 12), (2, 7800, 20, 112)],
+)
+
+
+@pytest.mark.parametrize("batch,seq_len,num_heads,head_dim", CASES)
+def test_interleaved_rope_fp64_is_bit_exact(
+    batch: int,
+    seq_len: int,
+    num_heads: int,
+    head_dim: int,
+) -> None:
+    generator = torch.Generator(device="cuda").manual_seed(42)
+    q = torch.randn(
+        batch,
+        seq_len,
+        num_heads,
+        head_dim,
+        dtype=torch.bfloat16,
+        device="cuda",
+        generator=generator,
+    )
+    k = torch.randn_like(q)
+    cos = torch.randn(
+        1,
+        seq_len,
+        1,
+        head_dim,
+        dtype=torch.float64,
+        device="cuda",
+        generator=generator,
+    )
+    sin = torch.randn_like(cos)
+
+    q_out, k_out = fused_interleaved_rope_fp64(q, k, cos, sin)
+
+    assert torch.equal(q_out, eager_interleaved_rope(q, cos, sin))
+    assert torch.equal(k_out, eager_interleaved_rope(k, cos, sin))
+    assert q_out.data_ptr() not in (q.data_ptr(), k.data_ptr())
+    assert k_out.data_ptr() not in (q.data_ptr(), k.data_ptr())
+
+
+def test_interleaved_rope_fp64_predicate_rejects_unsupported_inputs() -> None:
+    q = torch.empty(1, 17, 3, 12, dtype=torch.bfloat16, device="cuda")
+    k = torch.empty_like(q)
+    cos = torch.empty(1, 17, 1, 12, dtype=torch.float64, device="cuda")
+    sin = torch.empty_like(cos)
+
+    assert can_use_interleaved_rope_fp64(q, k, cos, sin)
+    assert not can_use_interleaved_rope_fp64(q.flatten(), k, cos, sin)
+    assert not can_use_interleaved_rope_fp64(q.float(), k, cos, sin)
+    assert not can_use_interleaved_rope_fp64(q, k[..., ::2], cos, sin)
+    assert not can_use_interleaved_rope_fp64(q, k, cos.float(), sin)
+    assert not can_use_interleaved_rope_fp64(q, k, cos[..., :-2], sin[..., :-2])
+    unaligned = torch.empty(q.numel() + 1, dtype=torch.bfloat16, device="cuda")[
+        1:
+    ].view_as(q)
+    assert unaligned.is_contiguous()
+    assert not can_use_interleaved_rope_fp64(unaligned, k, cos, sin)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))