From 96b31770f9e58d23a3b5b52c66f32b9208a987dd Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Fri, 28 Aug 2026 08:57:37 +0800 Subject: [PATCH] [diffusion] fuse Helios paired transposed RoPE (#36502) --- .../jit/csrc/diffusion/helios_qk_rope.cuh | 118 ++++++++++++ python/sglang/kernels/ops/diffusion/README.md | 1 + .../sglang/kernels/ops/diffusion/__init__.py | 9 + .../ops/diffusion/rope/helios_qk_rope_jit.py | 81 +++++++++ .../existing-fast-paths.md | 25 +++ .../runtime/models/dits/helios.py | 25 ++- .../diffusion/bench_helios_qk_rope.py | 60 +++++++ .../ops/diffusion/test_helios_qk_rope.py | 170 ++++++++++++++++++ 8 files changed, 487 insertions(+), 2 deletions(-) create mode 100644 python/sglang/kernels/jit/csrc/diffusion/helios_qk_rope.cuh create mode 100644 python/sglang/kernels/ops/diffusion/rope/helios_qk_rope_jit.py create mode 100644 test/registered/kernels/benchmark/diffusion/bench_helios_qk_rope.py create mode 100644 test/registered/kernels/ops/diffusion/test_helios_qk_rope.py diff --git a/python/sglang/kernels/jit/csrc/diffusion/helios_qk_rope.cuh b/python/sglang/kernels/jit/csrc/diffusion/helios_qk_rope.cuh new file mode 100644 index 000000000..ca601a7e4 --- /dev/null +++ b/python/sglang/kernels/jit/csrc/diffusion/helios_qk_rope.cuh @@ -0,0 +1,118 @@ +#include +#include + +#include +#include + +#include + +#include +#include +#include + +namespace sglang { + +/** + * \brief Apply Helios transposed RoPE to normalized Q/K in place. + * + * One thread owns one adjacent rotary pair. The explicit round-to-nearest + * multiply and add/subtract operations preserve the separate eager FP32 + * intermediates before the result is rounded back to fp16/bf16. + * + * \tparam T Activation type: fp16_t or bf16_t + * \param q Normalized query tensor, contiguous [tokens, heads, head_dim] + * \param k Normalized key tensor, contiguous [tokens, heads, head_dim] + * \param freqs Transposed Helios frequency tensor, contiguous + * [tokens, 2 * head_dim] + * \param num_pairs Total adjacent Q/K pairs across tokens and heads + * \param pairs_per_head Number of adjacent pairs in one attention head + * \param num_heads Number of attention heads + * \param freq_stride Last dimension of freqs, equal to 2 * head_dim + */ +template +__global__ void helios_qk_rope_kernel( + T* __restrict__ q, + T* __restrict__ k, + const float* __restrict__ freqs, + uint32_t num_pairs, + uint32_t pairs_per_head, + uint32_t num_heads, + uint32_t freq_stride) { + static_assert(std::is_same_v || std::is_same_v); + using Packed = packed_t; + + auto* q_pairs = reinterpret_cast(q); + auto* k_pairs = reinterpret_cast(k); + const uint32_t stride = blockDim.x * gridDim.x; + for (uint32_t pair_index = blockIdx.x * blockDim.x + threadIdx.x; pair_index < num_pairs; pair_index += stride) { + const uint32_t pair_in_head = pair_index % pairs_per_head; + const uint32_t token_head = pair_index / pairs_per_head; + const uint32_t token_index = token_head / num_heads; + const uint32_t head_dim = pairs_per_head * 2; + const uint32_t freq_base = token_index * freq_stride; + const float cos = freqs[freq_base + pair_in_head * 2]; + const float sin = freqs[freq_base + head_dim + pair_in_head * 2 + 1]; + + const auto q_value = device::cast(q_pairs[pair_index]); + const auto k_value = device::cast(k_pairs[pair_index]); + + const float q_even = __fsub_rn(__fmul_rn(q_value.x, cos), __fmul_rn(q_value.y, sin)); + const float q_odd = __fadd_rn(__fmul_rn(q_value.x, sin), __fmul_rn(q_value.y, cos)); + const float k_even = __fsub_rn(__fmul_rn(k_value.x, cos), __fmul_rn(k_value.y, sin)); + const float k_odd = __fadd_rn(__fmul_rn(k_value.x, sin), __fmul_rn(k_value.y, cos)); + + q_pairs[pair_index] = device::cast(make_float2(q_even, q_odd)); + k_pairs[pair_index] = device::cast(make_float2(k_even, k_odd)); + } +} + +/** \brief Validate and launch the paired Helios Q/K RoPE kernel. */ +template +struct HeliosQKRoPEKernel { + static void run(const tvm::ffi::TensorView q, const tvm::ffi::TensorView k, const tvm::ffi::TensorView freqs) { + using namespace host; + + auto N = SymbolicSize{"tokens"}; + auto H = SymbolicSize{"heads"}; + auto D = SymbolicSize{"head_dim"}; + auto F = SymbolicSize{"freq_dim"}; + auto device = SymbolicDevice{}; + device.set_options(); + + TensorMatcher({N, H, D}).with_dtype().with_device(device).verify(q).verify(k); + TensorMatcher({N, F}).with_dtype().with_device(device).verify(freqs); + + const int64_t tokens = N.unwrap(); + const int64_t heads = H.unwrap(); + const int64_t head_dim = D.unwrap(); + const int64_t freq_dim = F.unwrap(); + CHECK_HOST(tokens > 0 && heads > 0 && head_dim > 0) + << "Helios QK RoPE expects positive dimensions, got tokens=" << tokens << ", heads=" << heads + << ", head_dim=" << head_dim; + CHECK_HOST(head_dim % 2 == 0) << "Helios QK RoPE head_dim must be even, got " << head_dim; + CHECK_HOST(freq_dim == 2 * head_dim) << "Helios QK RoPE expects freq_dim=" << 2 * head_dim << ", got " << freq_dim; + CHECK_HOST(reinterpret_cast(q.data_ptr()) % alignof(packed_t) == 0) + << "Helios QK RoPE query pointer is not pair aligned"; + CHECK_HOST(reinterpret_cast(k.data_ptr()) % alignof(packed_t) == 0) + << "Helios QK RoPE key pointer is not pair aligned"; + + const int64_t num_pairs_i64 = tokens * heads * (head_dim / 2); + CHECK_HOST(num_pairs_i64 <= std::numeric_limits::max()) + << "Helios QK RoPE pair count exceeds uint32: " << num_pairs_i64; + + const uint32_t num_pairs = static_cast(num_pairs_i64); + constexpr uint32_t kBlockSize = 256; + const uint32_t grid = div_ceil(num_pairs, kBlockSize); + LaunchKernel(grid, kBlockSize, device.unwrap())( + helios_qk_rope_kernel, + static_cast(q.data_ptr()), + static_cast(k.data_ptr()), + static_cast(freqs.data_ptr()), + num_pairs, + static_cast(head_dim / 2), + static_cast(heads), + static_cast(freq_dim)); + } +}; + +} // namespace sglang diff --git a/python/sglang/kernels/ops/diffusion/README.md b/python/sglang/kernels/ops/diffusion/README.md index e93bd831a..12e4d44f9 100644 --- a/python/sglang/kernels/ops/diffusion/README.md +++ b/python/sglang/kernels/ops/diffusion/README.md @@ -118,6 +118,7 @@ Several norms look interchangeable and are not. Start here. | `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 | +| `fused_inplace_helios_qk_rope` | JIT CUDA | bit-exact paired in-place RoPE for Helios' transposed frequency layout | | `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 2cba1cbb2..67793a51d 100644 --- a/python/sglang/kernels/ops/diffusion/__init__.py +++ b/python/sglang/kernels/ops/diffusion/__init__.py @@ -221,6 +221,13 @@ _SPECS: tuple[tuple[str, KernelBackend, str, frozenset, str], ...] = ( _CUDA, "Paired interleaved RoPE with fp64 Diffusers semantics.", ), + ( + "diffusion.helios_qk_rope", + KernelBackend.JIT, + "rope.helios_qk_rope_jit:fused_inplace_helios_qk_rope", + _CUDA, + "Paired in-place Helios transposed Q/K RoPE.", + ), ( "diffusion.hunyuan_qkv_rope_pack", KernelBackend.TRITON, @@ -410,6 +417,8 @@ _EXPORTS: dict[str, str] = { "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", + "can_use_helios_qk_rope": "rope.helios_qk_rope_jit", + "fused_inplace_helios_qk_rope": "rope.helios_qk_rope_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/helios_qk_rope_jit.py b/python/sglang/kernels/ops/diffusion/rope/helios_qk_rope_jit.py new file mode 100644 index 000000000..e5f4b72d6 --- /dev/null +++ b/python/sglang/kernels/ops/diffusion/rope/helios_qk_rope_jit.py @@ -0,0 +1,81 @@ +"""Bit-exact paired RoPE for Helios' transposed frequency layout. + +The JIT kernel preserves the eager path's separate fp32 multiply and add/sub +rounding boundaries, then rounds each adjacent pair back to fp16/bf16 in +place. It is verified on head dimensions 64, 128, and 256, including Helios' +production ``[8640, 40, 128]`` Q/K shape. Unsupported layouts retain the eager +model path through :func:`can_use_helios_qk_rope`. +""" + +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_helios_qk_rope_module(dtype: torch.dtype) -> Module: + if dtype not in (torch.float16, torch.bfloat16): + raise RuntimeError( + f"Unsupported Helios QK RoPE dtype {dtype}; expected float16 or bfloat16" + ) + args = make_cpp_args(dtype) + return load_jit( + "helios_qk_rope", + *args, + cuda_files=["diffusion/helios_qk_rope.cuh"], + cuda_wrappers=[("helios_qk_rope", f"HeliosQKRoPEKernel<{args}>::run")], + ) + + +@register_custom_op(mutates_args=["q", "k"]) +def fused_inplace_helios_qk_rope( + q: torch.Tensor, + k: torch.Tensor, + freqs: torch.Tensor, +) -> None: + """Apply Helios' transposed RoPE to contiguous normalized Q/K in place.""" + module = _jit_helios_qk_rope_module(q.dtype) + module.helios_qk_rope(q, k, freqs) + + +def can_use_helios_qk_rope( + q: torch.Tensor, + k: torch.Tensor, + freqs: torch.Tensor, +) -> bool: + """Return whether tensors match the native Helios paired-RoPE contract.""" + if q.dim() != 4 or freqs.dim() != 3: + return False + # Dynamo cannot trace pointer or storage-offset queries. Compiled Helios Q/K + # come directly from aligned linear outputs; eager callers retain the guard. + pair_aligned = True + if not torch.compiler.is_compiling(): + pair_aligned = q.storage_offset() % 2 == 0 and k.storage_offset() % 2 == 0 + return ( + q.is_cuda + and k.is_cuda + and freqs.is_cuda + and q.dtype in (torch.float16, torch.bfloat16) + and k.dtype == q.dtype + and freqs.dtype is torch.float32 + and q.device == k.device == freqs.device + and k.shape == q.shape + and all(size > 0 for size in q.shape) + and freqs.shape == (*q.shape[:2], 2 * q.shape[-1]) + and q.shape[-1] % 2 == 0 + and q.is_contiguous() + and k.is_contiguous() + and freqs.is_contiguous() + and pair_aligned + ) + + +__all__ = ["can_use_helios_qk_rope", "fused_inplace_helios_qk_rope"] diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md index cfaebe848..94fbe4ac5 100644 --- a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md +++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-benchmark-profile/existing-fast-paths.md @@ -24,6 +24,7 @@ framework-specific optimization workflow. - `python/sglang/kernels/ops/diffusion/norm/native_bf16_rmsnorm_triton.py` - `python/sglang/kernels/ops/diffusion/norm/zimage_qk_rmsnorm_triton.py` - `python/sglang/kernels/ops/diffusion/rope/rotary_triton.py` +- `python/sglang/kernels/ops/diffusion/rope/helios_qk_rope_jit.py` - `python/sglang/kernels/ops/diffusion/rope/ltx2_rotary_triton.py` - `python/sglang/kernels/ops/diffusion/rope/ltx2_qknorm_split_rope_jit.py` - `python/sglang/kernels/ops/diffusion/sites/ltx2_rmsnorm_modulate_site.py` @@ -217,6 +218,30 @@ framework-specific optimization workflow. path. Unsupported layouts or padding fall back to the aten chain. - Validation: `test/registered/kernels/ops/diffusion/test_wan_causal_cache.py`. +15. Helios paired transposed RoPE +- Kernel: `fused_inplace_helios_qk_rope`. +- Locations: `rope/helios_qk_rope_jit.py`, + `csrc/diffusion/helios_qk_rope.cuh`, and + `runtime/models/dits/helios.py`. +- Use case: apply Helios' transposed fp32 frequency table to already-normalized + contiguous Q/K together, in place, instead of launching the eager + unflatten/chunk/multiply/add/stack chain twice per attention block. +- Constraints: CUDA fp16/bf16 Q/K with matching contiguous `[B, S, H, D]` + layouts, contiguous fp32 frequencies shaped `[B, S, 2 * D]`, even `D`, and + pair-aligned Q/K pointers. Tensor-parallel RMSNorm keeps the eager path. + Current real-model validation covers one H100; it is not a multi-GPU scaling + claim. +- Numerical contract: explicit round-to-nearest fp32 operations reproduce the + eager elementwise rounding boundaries before the result is cast back to the + activation dtype. Correctness tests require `torch.equal`, including the + production `[8640, 40, 128]` shape. +- Validation: `test/registered/kernels/ops/diffusion/test_helios_qk_rope.py`. +- Microbench: + `test/registered/kernels/benchmark/diffusion/bench_helios_qk_rope.py`. +- Workflow rule: if a Helios trace still shows two transposed-RoPE elementwise + ladders per block, check TP mode, dtype, shape, contiguity, and pointer + alignment before proposing another RoPE kernel. + **Faster CUDA Kernel Usage Points** 1. sgl-kernel RMSNorm and fused add RMSNorm diff --git a/python/sglang/multimodal_gen/runtime/models/dits/helios.py b/python/sglang/multimodal_gen/runtime/models/dits/helios.py index 4be5d334f..048956af5 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/helios.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/helios.py @@ -16,6 +16,10 @@ import torch import torch.nn as nn import torch.nn.functional as F +from sglang.kernels.ops.diffusion import ( + can_use_helios_qk_rope, + fused_inplace_helios_qk_rope, +) from sglang.multimodal_gen.configs.models.dits.helios import HeliosConfig from sglang.multimodal_gen.configs.models.fsdp import is_block from sglang.multimodal_gen.runtime.distributed import ( @@ -285,6 +289,24 @@ class HeliosSelfAttention(nn.Module): self.history_scale_mode = history_scale_mode self.max_scale = 10.0 + def _apply_rotary_qk( + self, + q: torch.Tensor, + k: torch.Tensor, + rotary_emb: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + if not self.tp_rmsnorm and can_use_helios_qk_rope(q, k, rotary_emb): + fused_inplace_helios_qk_rope( + q.view(-1, q.shape[-2], q.shape[-1]), + k.view(-1, k.shape[-2], k.shape[-1]), + rotary_emb.view(-1, rotary_emb.shape[-1]), + ) + return q, k + return ( + apply_rotary_emb_transposed(q, rotary_emb), + apply_rotary_emb_transposed(k, rotary_emb), + ) + def forward(self, hidden_states, rotary_emb=None, original_context_length=None): q, _ = self.to_q(hidden_states) k, _ = self.to_k(hidden_states) @@ -302,8 +324,7 @@ class HeliosSelfAttention(nn.Module): v = v.unflatten(2, (self.local_num_heads, self.head_dim)) if rotary_emb is not None: - q = apply_rotary_emb_transposed(q, rotary_emb) - k = apply_rotary_emb_transposed(k, rotary_emb) + q, k = self._apply_rotary_qk(q, k, rotary_emb) history_seq_len = ( hidden_states.shape[1] - original_context_length diff --git a/test/registered/kernels/benchmark/diffusion/bench_helios_qk_rope.py b/test/registered/kernels/benchmark/diffusion/bench_helios_qk_rope.py new file mode 100644 index 000000000..9108e499e --- /dev/null +++ b/test/registered/kernels/benchmark/diffusion/bench_helios_qk_rope.py @@ -0,0 +1,60 @@ +import torch + +from sglang.kernels.jit.benchmark import marker +from sglang.kernels.ops.diffusion import fused_inplace_helios_qk_rope +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci( + est_time=12, + stage="base-b-kernel-benchmark", + runner_config="1-gpu-large", +) + + +def _split(q: torch.Tensor, k: torch.Tensor, freqs: torch.Tensor) -> None: + def apply(value: torch.Tensor) -> torch.Tensor: + x_1, x_2 = value.unflatten(-1, (-1, 2)).unbind(-1) + cos, sin = freqs.unsqueeze(-2).chunk(2, dim=-1) + out = torch.empty_like(value) + out[..., 0::2] = x_1 * cos[..., 0::2] - x_2 * sin[..., 1::2] + out[..., 1::2] = x_1 * sin[..., 1::2] + x_2 * cos[..., 0::2] + return out.type_as(value) + + apply(q) + apply(k) + + +FN_MAP = { + "eager": _split, + "jit": fused_inplace_helios_qk_rope, +} + + +@marker.parametrize("tokens", [2160, 8640], [8640]) +@marker.benchmark("impl", ["eager", "jit"]) +def benchmark(tokens: int, impl: str): + generator = torch.Generator(device="cuda").manual_seed(20260826) + q = torch.randn( + tokens, + 40, + 128, + device="cuda", + dtype=torch.bfloat16, + generator=generator, + ) + k = torch.randn_like(q) + freqs = torch.randn( + tokens, 256, device="cuda", dtype=torch.float32, generator=generator + ) + return marker.do_bench( + FN_MAP[impl], + input_args=(q, k, freqs), + memory_args=(q, k, freqs), + memory_output=None, + use_cuda_graph=False, + disable_log_bandwidth=True, + ) + + +if __name__ == "__main__": + benchmark.run() diff --git a/test/registered/kernels/ops/diffusion/test_helios_qk_rope.py b/test/registered/kernels/ops/diffusion/test_helios_qk_rope.py new file mode 100644 index 000000000..271719d33 --- /dev/null +++ b/test/registered/kernels/ops/diffusion/test_helios_qk_rope.py @@ -0,0 +1,170 @@ +import sys +from unittest.mock import patch + +import pytest +import torch +import torch.nn as nn + +from sglang.kernels.ops.diffusion import ( + can_use_helios_qk_rope, + fused_inplace_helios_qk_rope, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=45, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _reference(value: torch.Tensor, freqs: torch.Tensor) -> torch.Tensor: + x_1, x_2 = value.unflatten(-1, (-1, 2)).unbind(-1) + cos, sin = freqs.unsqueeze(-2).chunk(2, dim=-1) + out = torch.empty_like(value) + out[..., 0::2] = x_1 * cos[..., 0::2] - x_2 * sin[..., 1::2] + out[..., 1::2] = x_1 * sin[..., 1::2] + x_2 * cos[..., 0::2] + return out.type_as(value) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +@pytest.mark.parametrize( + "tokens,heads,head_dim", + [ + (1, 1, 64), + (17, 8, 128), + (129, 4, 256), + (8640, 40, 128), + ], +) +def test_helios_qk_rope_matches_eager_transposed_path( + dtype: torch.dtype, + tokens: int, + heads: int, + head_dim: int, +) -> None: + generator = torch.Generator(device="cuda").manual_seed(20260826) + q = torch.randn( + tokens, heads, head_dim, device="cuda", dtype=dtype, generator=generator + ) + k = torch.randn_like(q) + freqs = torch.randn( + tokens, 2 * head_dim, device="cuda", dtype=torch.float32, generator=generator + ) + q_ref = _reference(q, freqs) + k_ref = _reference(k, freqs) + q_out, k_out = q.clone(), k.clone() + q_ptr, k_ptr = q_out.data_ptr(), k_out.data_ptr() + + fused_inplace_helios_qk_rope(q_out, k_out, freqs) + torch.cuda.synchronize() + + assert q_out.data_ptr() == q_ptr + assert k_out.data_ptr() == k_ptr + assert torch.equal(q_out, q_ref) + assert torch.equal(k_out, k_ref) + + +def test_helios_qk_rope_runtime_guards() -> None: + q = torch.randn(1, 17, 8, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + freqs = torch.randn(1, 17, 256, device="cuda", dtype=torch.float32) + assert can_use_helios_qk_rope(q, k, freqs) + assert not can_use_helios_qk_rope(q.float(), k, freqs) + assert not can_use_helios_qk_rope(q, k.float(), freqs) + assert not can_use_helios_qk_rope(q, k, freqs.bfloat16()) + assert not can_use_helios_qk_rope(q.cpu(), k.cpu(), freqs.cpu()) + assert not can_use_helios_qk_rope(q, k, freqs[..., :-2]) + assert not can_use_helios_qk_rope(q[:, :, :, ::2], k, freqs) + assert not can_use_helios_qk_rope(q[:, :0], k[:, :0], freqs[:, :0]) + + q_unaligned = torch.empty(q.numel() + 1, device=q.device, dtype=q.dtype)[ + 1: + ].view_as(q) + k_unaligned = torch.empty(k.numel() + 1, device=k.device, dtype=k.dtype)[ + 1: + ].view_as(k) + assert q_unaligned.is_contiguous() and q_unaligned.storage_offset() == 1 + assert k_unaligned.is_contiguous() and k_unaligned.storage_offset() == 1 + assert not can_use_helios_qk_rope(q_unaligned, k, freqs) + assert not can_use_helios_qk_rope(q, k_unaligned, freqs) + + +def test_helios_attention_dispatch_and_tp_fallback() -> None: + import sglang.multimodal_gen.runtime.models.dits.helios as helios + + attention = helios.HeliosSelfAttention.__new__(helios.HeliosSelfAttention) + nn.Module.__init__(attention) + attention.tp_rmsnorm = False + q = torch.randn(1, 17, 8, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + freqs = torch.randn(1, 17, 256, device="cuda", dtype=torch.float32) + + with ( + patch.object(helios, "can_use_helios_qk_rope", return_value=True), + patch.object(helios, "fused_inplace_helios_qk_rope") as fused, + ): + q_out, k_out = attention._apply_rotary_qk(q, k, freqs) + assert q_out is q + assert k_out is k + fused.assert_called_once() + assert fused.call_args.args[0].shape == (17, 8, 128) + assert fused.call_args.args[1].shape == (17, 8, 128) + assert fused.call_args.args[2].shape == (17, 256) + + attention.tp_rmsnorm = True + with patch.object(helios, "fused_inplace_helios_qk_rope") as fused: + q_out, k_out = attention._apply_rotary_qk(q, k, freqs) + fused.assert_not_called() + assert torch.equal(q_out, _reference(q, freqs)) + assert torch.equal(k_out, _reference(k, freqs)) + + +def test_helios_qk_rope_fullgraph_custom_op() -> None: + q = torch.randn(1, 17, 8, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + freqs = torch.randn(1, 17, 256, device="cuda", dtype=torch.float32) + q_ref, k_ref = _reference(q, freqs), _reference(k, freqs) + + @torch.compile(fullgraph=True) + def compiled( + q_arg: torch.Tensor, k_arg: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + fused_inplace_helios_qk_rope( + q_arg.view(-1, q_arg.shape[-2], q_arg.shape[-1]), + k_arg.view(-1, k_arg.shape[-2], k_arg.shape[-1]), + freqs.view(-1, freqs.shape[-1]), + ) + return q_arg, k_arg + + q_out, k_out = compiled(q.clone(), k.clone()) + assert torch.equal(q_out, q_ref) + assert torch.equal(k_out, k_ref) + + +def test_helios_attention_fullgraph_dispatch() -> None: + import sglang.multimodal_gen.runtime.models.dits.helios as helios + + attention = helios.HeliosSelfAttention.__new__(helios.HeliosSelfAttention) + nn.Module.__init__(attention) + attention.tp_rmsnorm = False + q = torch.randn(1, 17, 8, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + freqs = torch.randn(1, 17, 256, device="cuda", dtype=torch.float32) + q_ref, k_ref = _reference(q, freqs), _reference(k, freqs) + + compiled = torch.compile(attention._apply_rotary_qk, fullgraph=True) + q_out, k_out = compiled(q.clone(), k.clone(), freqs) + + assert torch.equal(q_out, q_ref) + assert torch.equal(k_out, k_ref) + + +def test_helios_qk_rope_rejects_bad_frequency_shape() -> None: + q = torch.randn(3, 2, 128, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + freqs = torch.randn(3, 128, device="cuda", dtype=torch.float32) + with pytest.raises(RuntimeError, match="freq_dim"): + fused_inplace_helios_qk_rope(q, k, freqs) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__, "-v"]))