From 83e9ece672d9080d7e785c5d090bbed50ccf08d6 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Sat, 22 Aug 2026 12:56:45 +0800 Subject: [PATCH] [diffusion] Fuse SANA-Video interleaved RoPE (#35695) --- docs/docs/sglang-diffusion/fused_kernels.mdx | 2 + .../csrc/diffusion/interleaved_rope_fp64.cuh | 148 ++++++++++++++++++ python/sglang/kernels/ops/diffusion/README.md | 1 + .../sglang/kernels/ops/diffusion/__init__.py | 9 ++ .../rope/interleaved_rope_fp64_jit.py | 105 +++++++++++++ .../runtime/models/dits/sana_video.py | 54 ++++++- .../test/unit/test_sana_video.py | 14 ++ .../diffusion/bench_interleaved_rope_fp64.py | 87 ++++++++++ .../diffusion/test_interleaved_rope_fp64.py | 98 ++++++++++++ 9 files changed, 516 insertions(+), 2 deletions(-) create mode 100644 python/sglang/kernels/jit/csrc/diffusion/interleaved_rope_fp64.cuh create mode 100644 python/sglang/kernels/ops/diffusion/rope/interleaved_rope_fp64_jit.py create mode 100644 test/registered/kernels/benchmark/diffusion/bench_interleaved_rope_fp64.py create mode 100644 test/registered/kernels/ops/diffusion/test_interleaved_rope_fp64.py diff --git a/docs/docs/sglang-diffusion/fused_kernels.mdx b/docs/docs/sglang-diffusion/fused_kernels.mdx index 07a0facce..bee260769 100644 --- a/docs/docs/sglang-diffusion/fused_kernels.mdx +++ b/docs/docs/sglang-diffusion/fused_kernels.mdx @@ -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 | diff --git a/python/sglang/kernels/jit/csrc/diffusion/interleaved_rope_fp64.cuh b/python/sglang/kernels/jit/csrc/diffusion/interleaved_rope_fp64.cuh new file mode 100644 index 000000000..bcdd599b2 --- /dev/null +++ b/python/sglang/kernels/jit/csrc/diffusion/interleaved_rope_fp64.cuh @@ -0,0 +1,148 @@ +#pragma once + +#include +#include + +#include +#include +#include + +#include +#include +#include + +namespace sglang { + +namespace interleaved_rope_fp64 { + +namespace { + +constexpr uint32_t kBlockSize = 256; +constexpr uint32_t kMaxGrid = 65535; + +template +SGL_DEVICE device::AlignedVector rotate_pair(device::AlignedVector input, double cos, double sin) { + static_assert(std::is_same_v); + const double x1 = static_cast(input[0]); + const double x2 = static_cast(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 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(static_cast(__dsub_rn(x1_cos, x2_sin))); + output[1] = static_cast(static_cast(__dadd_rn(x1_sin, x2_cos))); + return output; +} + +template +__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; + const int64_t stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t pair_index = static_cast(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 +struct InterleavedRopeFP64Kernel { + static_assert(std::is_same_v); + + 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; + + auto N = SymbolicSize{"activation_elements"}; + auto R = SymbolicSize{"table_elements"}; + auto device = SymbolicDevice{}; + device.set_options(); + TensorMatcher({N}).with_dtype().with_device(device).verify(q_out).verify(k_out).verify(q).verify(k); + TensorMatcher({R}).with_dtype().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(q_out.data_ptr()) % alignof(Pair) == 0 && + reinterpret_cast(k_out.data_ptr()) % alignof(Pair) == 0 && + reinterpret_cast(q.data_ptr()) % alignof(Pair) == 0 && + reinterpret_cast(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(std::min(div_ceil(num_pairs, static_cast(kBlockSize)), kMaxGrid)); + LaunchKernel(blocks, kBlockSize, device.unwrap())( + interleaved_rope_fp64_kernel, + static_cast(q_out.data_ptr()), + static_cast(k_out.data_ptr()), + static_cast(q.data_ptr()), + static_cast(k.data_ptr()), + static_cast(cos.data_ptr()), + static_cast(sin.data_ptr()), + num_pairs, + seq_len, + num_heads, + pairs_per_head, + head_dim); + } +}; + +} // namespace interleaved_rope_fp64 + +} // namespace sglang diff --git a/python/sglang/kernels/ops/diffusion/README.md b/python/sglang/kernels/ops/diffusion/README.md index 8d017789e..985972af4 100644 --- a/python/sglang/kernels/ops/diffusion/README.md +++ b/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 | diff --git a/python/sglang/kernels/ops/diffusion/__init__.py b/python/sglang/kernels/ops/diffusion/__init__.py index 279f8aa70..b57589fde 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/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", diff --git a/python/sglang/kernels/ops/diffusion/rope/interleaved_rope_fp64_jit.py b/python/sglang/kernels/ops/diffusion/rope/interleaved_rope_fp64_jit.py new file mode 100644 index 000000000..ebf55a8df --- /dev/null +++ b/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", +] 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 8bdd63d27..ac4235ccc 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/sana_video.py +++ b/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) diff --git a/python/sglang/multimodal_gen/test/unit/test_sana_video.py b/python/sglang/multimodal_gen/test/unit/test_sana_video.py index 73a3219ff..5b99165ef 100644 --- a/python/sglang/multimodal_gen/test/unit/test_sana_video.py +++ b/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)) diff --git a/test/registered/kernels/benchmark/diffusion/bench_interleaved_rope_fp64.py b/test/registered/kernels/benchmark/diffusion/bench_interleaved_rope_fp64.py new file mode 100644 index 000000000..807a8176f --- /dev/null +++ b/test/registered/kernels/benchmark/diffusion/bench_interleaved_rope_fp64.py @@ -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() diff --git a/test/registered/kernels/ops/diffusion/test_interleaved_rope_fp64.py b/test/registered/kernels/ops/diffusion/test_interleaved_rope_fp64.py new file mode 100644 index 000000000..5c77c3485 --- /dev/null +++ b/test/registered/kernels/ops/diffusion/test_interleaved_rope_fp64.py @@ -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"]))