[diffusion] fuse Helios paired transposed RoPE (#36502)

2026-08-28 08:57:37 +08:00
parent 26fd7fdaa2
commit 96b31770f9
8 changed files with 487 additions and 2 deletions
@@ -0,0 +1,118 @@
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+
+#include <tvm/ffi/container/tensor.h>
+
+#include <cstdint>
+#include <limits>
+#include <type_traits>
+
+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 <typename T>
+__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<T, fp16_t> || std::is_same_v<T, bf16_t>);
+  using Packed = packed_t<T>;
+
+  auto* q_pairs = reinterpret_cast<Packed*>(q);
+  auto* k_pairs = reinterpret_cast<Packed*>(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<fp32x2_t, Packed>(q_pairs[pair_index]);
+    const auto k_value = device::cast<fp32x2_t, Packed>(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<Packed, fp32x2_t>(make_float2(q_even, q_odd));
+    k_pairs[pair_index] = device::cast<Packed, fp32x2_t>(make_float2(k_even, k_odd));
+  }
+}
+
+/** \brief Validate and launch the paired Helios Q/K RoPE kernel. */
+template <typename DType>
+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<kDLCUDA>();
+
+    TensorMatcher({N, H, D}).with_dtype<DType>().with_device(device).verify(q).verify(k);
+    TensorMatcher({N, F}).with_dtype<fp32_t>().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<uintptr_t>(q.data_ptr()) % alignof(packed_t<DType>) == 0)
+        << "Helios QK RoPE query pointer is not pair aligned";
+    CHECK_HOST(reinterpret_cast<uintptr_t>(k.data_ptr()) % alignof(packed_t<DType>) == 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<uint32_t>::max())
+        << "Helios QK RoPE pair count exceeds uint32: " << num_pairs_i64;
+
+    const uint32_t num_pairs = static_cast<uint32_t>(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<DType>,
+        static_cast<DType*>(q.data_ptr()),
+        static_cast<DType*>(k.data_ptr()),
+        static_cast<const float*>(freqs.data_ptr()),
+        num_pairs,
+        static_cast<uint32_t>(head_dim / 2),
+        static_cast<uint32_t>(heads),
+        static_cast<uint32_t>(freq_dim));
+  }
+};
+
+}  // namespace sglang
@@ -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 |
@@ -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",
@@ -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"]
@@ -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
@@ -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