Uploaded using `kernel-builder`.

build/torch-rocm/__init__.py

@@ -0,0 +1,77 @@
+"""AITER RoPE kernels for AMD ROCm (Triton).
+
+Triton-based RoPE (Rotary Position Embedding) kernels repackaged from the
+[ROCm/aiter](https://github.com/ROCm/aiter) project. Exposes
+``apply_rotary_transformers`` — drop-in replacement for the same-named
+function in ``kernels-community/rotary`` — so transformers can use this as
+the ROCm entry of ``_KERNEL_MAPPING["rotary_pos_emb"]`` with zero model-side
+changes.
+"""
+
+import torch
+
+from .rope import (
+    RotateStyle,
+    rope_cached_fwd,
+    rope_cached_fwd_inplace,
+)
+
+
+__kernel_metadata__ = {
+    "license": "mit",
+}
+
+
+def _to_sbhd(t: torch.Tensor) -> torch.Tensor:
+    # transformers passes attention tensors as (batch, heads, seq, head_dim);
+    # AITER's kernels operate in (seq, batch, heads, head_dim).
+    return t.permute(2, 0, 1, 3).contiguous()
+
+
+def _to_bhsd(t: torch.Tensor) -> torch.Tensor:
+    return t.permute(1, 2, 0, 3)
+
+
+def apply_rotary_transformers(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    """Apply NEOX-style RoPE to ``q`` and ``k``.
+
+    Signature mirrors ``kernels-community/rotary``'s ``apply_rotary_transformers``
+    so this kernel can be a drop-in replacement under the ``"rocm"`` entry of
+    transformers' ``_KERNEL_MAPPING["rotary_pos_emb"]``.
+
+    Args:
+        q, k: ``(batch, heads, seq, head_dim)``.
+        cos, sin: ``(batch, seq, head_dim // 2)`` — pre-``unsqueeze`` form.
+        position_ids, unsqueeze_dim: accepted for API parity; the kernel reads
+            positions from the already-computed ``cos`` / ``sin``.
+
+    Returns:
+        ``(q_embed, k_embed)`` in the same shape as ``q``, ``k``.
+    """
+    q_sbhd = _to_sbhd(q)
+    k_sbhd = _to_sbhd(k)
+
+    # Cached path expects (seq, 1, 1, head_dim/2). Take batch index 0 — for the
+    # standard (unpadded, equal-length) inference case all batch rows share the
+    # same cos/sin since they come from the same position_ids.
+    cos_freqs = cos[0].unsqueeze(1).unsqueeze(1).contiguous()
+    sin_freqs = sin[0].unsqueeze(1).unsqueeze(1).contiguous()
+
+    common = dict(
+        rotate_style=int(RotateStyle.NEOX),
+        reuse_freqs_front_part=True,
+        nope_first=False,
+    )
+    q_out_sbhd = rope_cached_fwd(q_sbhd, cos_freqs, sin_freqs, **common)
+    k_out_sbhd = rope_cached_fwd(k_sbhd, cos_freqs, sin_freqs, **common)
+
+    return _to_bhsd(q_out_sbhd), _to_bhsd(k_out_sbhd)
+
+
+__all__ = [
+    "__kernel_metadata__",
+    "RotateStyle",
+    "apply_rotary_transformers",
+    "rope_cached_fwd",
+    "rope_cached_fwd_inplace",
+]

build/torch-rocm/_kernels/rope.py

@@ -0,0 +1,2040 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _get_neox_rotated_x_1D(
+    x,
+    x_rotated_mask,
+    BLOCK_D: tl.constexpr,
+    BLOCK_D_HALF: tl.constexpr,
+):
+    x_rotated = tl.where(x_rotated_mask, x, -x)
+    x_rotated = tl.reshape(x_rotated, (2, BLOCK_D_HALF))
+    x_rotated = tl.flip(x_rotated, 1)
+    x_rotated = tl.reshape(x_rotated, (BLOCK_D,))
+    x_rotated = tl.flip(x_rotated, 0)
+    return x_rotated
+
+
+@triton.jit
+def _get_gptj_rotated_x_1D(
+    x,
+    x_rotated_mask,
+    BLOCK_D: tl.constexpr,
+    BLOCK_D_HALF: tl.constexpr,
+):
+    x_rotated = tl.where(x_rotated_mask, x, -x)
+    x_rotated = tl.reshape(x_rotated, (BLOCK_D_HALF, 2))
+    x_rotated = tl.flip(x_rotated, 1)
+    x_rotated = tl.reshape(x_rotated, (BLOCK_D,))
+    return x_rotated
+
+
+@triton.jit
+def _get_neox_rotated_x(
+    x,
+    x_rotated_mask,
+    BLOCK_T: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    BLOCK_D_HALF: tl.constexpr,
+    IS_BWD: tl.constexpr = False,
+):
+    if IS_BWD:
+        x_rotated = tl.where(x_rotated_mask, -x, x)
+    else:
+        x_rotated = tl.where(x_rotated_mask, x, -x)
+
+    x_rotated = tl.reshape(x_rotated, (BLOCK_T, 2, BLOCK_D_HALF))
+    x_rotated = tl.flip(x_rotated, 2)
+    x_rotated = tl.reshape(
+        x_rotated,
+        (
+            BLOCK_T,
+            BLOCK_D,
+        ),
+    )
+    x_rotated = tl.flip(x_rotated, 1)
+    return x_rotated
+
+
+@triton.jit
+def _get_gptj_rotated_x(
+    x,
+    x_rotated_mask,
+    BLOCK_T: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    BLOCK_D_HALF: tl.constexpr,
+    IS_BWD: tl.constexpr = False,
+):
+    if IS_BWD:
+        x_rotated = tl.where(x_rotated_mask, -x, x)
+    else:
+        x_rotated = tl.where(x_rotated_mask, x, -x)
+
+    x_rotated = tl.reshape(x_rotated, (BLOCK_T, BLOCK_D_HALF, 2))
+    x_rotated = tl.flip(x_rotated, 2)
+    x_rotated = tl.reshape(
+        x_rotated,
+        (
+            BLOCK_T,
+            BLOCK_D,
+        ),
+    )
+    return x_rotated
+
+
+@triton.jit
+def _rope_kernel_sbhd_fwd(
+    x_ptr,
+    freqs_ptr,
+    out_ptr,
+    stride_x_s,
+    stride_x_b,
+    stride_x_h,
+    stride_x_d,
+    stride_freqs_s,
+    stride_freqs_b,
+    stride_freqs_h,
+    stride_freqs_d,
+    stride_out_s,
+    stride_out_b,
+    stride_out_h,
+    stride_out_d,
+    S,
+    HAVE_NOPE: tl.constexpr,
+    NOPE_FIRST: tl.constexpr,
+    INPLACE: tl.constexpr,
+    REUSE_FREQS_FRONT_PART: tl.constexpr,
+    IS_NEOX: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    BLOCK_D_HALF: tl.constexpr,
+):
+    # Parallelize over batch and head. Handle 1 sequence per program
+    b = tl.program_id(0)
+    h = tl.program_id(1)
+    pid_s = tl.program_id(2)
+
+    s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S)
+    d_offs = tl.arange(0, BLOCK_D)
+    s_mask = s_offs < S
+
+    if REUSE_FREQS_FRONT_PART:
+        if IS_NEOX:
+            d_freqs_offs = tl.where(
+                (d_offs >= BLOCK_D_HALF) & (d_offs < BLOCK_D),
+                d_offs - BLOCK_D_HALF,
+                d_offs,
+            ).to(d_offs.dtype)
+            d_freqs_mask = d_freqs_offs < BLOCK_D
+        else:
+            d_freqs_offs = d_offs // 2
+            d_freqs_mask = d_freqs_offs < BLOCK_D_HALF
+    else:
+        d_freqs_offs = d_offs
+        d_freqs_mask = d_freqs_offs < BLOCK_D
+
+    freqs_mask = s_mask[:, None] & d_freqs_mask[None, :]
+    freqs_offs = (
+        s_offs[:, None] * stride_freqs_s + d_freqs_offs[None, :] * stride_freqs_d
+    )
+
+    freqs = tl.load(freqs_ptr + freqs_offs, mask=freqs_mask)
+    cos = tl.cos(freqs.to(tl.float32))
+    sin = tl.sin(freqs.to(tl.float32))
+
+    nope_offs = 0
+    if HAVE_NOPE and NOPE_FIRST:
+        nope_offs = BLOCK_D
+
+    x_offs = (
+        b * stride_x_b
+        + s_offs[:, None] * stride_x_s
+        + h * stride_x_h
+        + (d_offs + nope_offs)[None, :] * stride_x_d
+    )
+    x_mask = s_mask[:, None] & (d_offs < BLOCK_D)[None, :]
+    x = tl.load(x_ptr + x_offs, mask=x_mask)
+
+    if IS_NEOX:
+        x_rotated_mask = (d_offs < BLOCK_D_HALF)[None, :]
+        x_rotated = _get_neox_rotated_x(
+            x, x_rotated_mask, BLOCK_S, BLOCK_D, BLOCK_D_HALF
+        )
+    else:
+        x_rotated_mask = (d_offs % 2 == 0)[None, :]
+        x_rotated = _get_gptj_rotated_x(
+            x, x_rotated_mask, BLOCK_S, BLOCK_D, BLOCK_D_HALF
+        )
+
+    out_x = x * cos + x_rotated * sin
+    out_x = out_x.to(x_ptr.dtype.element_ty)
+    x_out_offs = (
+        b * stride_out_b
+        + s_offs[:, None] * stride_out_s
+        + h * stride_out_h
+        + (d_offs + nope_offs)[None, :] * stride_out_d
+    )
+
+    tl.store(out_ptr + x_out_offs, out_x, mask=x_mask)
+
+    if HAVE_NOPE and not INPLACE:
+        if NOPE_FIRST:
+            x = tl.load(x_ptr + x_offs - BLOCK_D * stride_x_d, mask=x_mask)
+            tl.store(out_ptr + x_out_offs - BLOCK_D * stride_out_d, x, mask=x_mask)
+        else:
+            x = tl.load(x_ptr + x_offs + BLOCK_D * stride_x_d, mask=x_mask)
+            tl.store(out_ptr + x_out_offs + BLOCK_D * stride_out_d, x, mask=x_mask)
+
+
+@triton.jit
+def _rope_kernel_sbhd_bwd(
+    x_ptr,
+    freqs_ptr,
+    out_ptr,
+    stride_x_s,
+    stride_x_b,
+    stride_x_h,
+    stride_x_d,
+    stride_freqs_s,
+    stride_freqs_b,
+    stride_freqs_h,
+    stride_freqs_d,
+    stride_out_s,
+    stride_out_b,
+    stride_out_h,
+    stride_out_d,
+    S,
+    HAVE_NOPE: tl.constexpr,
+    NOPE_FIRST: tl.constexpr,
+    INPLACE: tl.constexpr,
+    REUSE_FREQS_FRONT_PART: tl.constexpr,
+    IS_NEOX: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    BLOCK_D_HALF: tl.constexpr,
+):
+    # Parallelize over batch and head. Handle 1 sequence per program
+    b = tl.program_id(0)
+    h = tl.program_id(1)
+    pid_s = tl.program_id(2)
+
+    s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S)
+    d_offs = tl.arange(0, BLOCK_D)
+    s_mask = s_offs < S
+
+    if REUSE_FREQS_FRONT_PART:
+        if IS_NEOX:
+            d_freqs_offs = tl.where(
+                (d_offs >= BLOCK_D_HALF) & (d_offs < BLOCK_D),
+                d_offs - BLOCK_D_HALF,
+                d_offs,
+            ).to(d_offs.dtype)
+            d_freqs_mask = d_freqs_offs < BLOCK_D
+        else:
+            d_freqs_offs = d_offs // 2
+            d_freqs_mask = d_freqs_offs < BLOCK_D_HALF
+    else:
+        d_freqs_offs = d_offs
+        d_freqs_mask = d_freqs_offs < BLOCK_D
+
+    freqs_mask = s_mask[:, None] & d_freqs_mask[None, :]
+    freqs_offs = (
+        s_offs[:, None] * stride_freqs_s + d_freqs_offs[None, :] * stride_freqs_d
+    )
+
+    freqs = tl.load(freqs_ptr + freqs_offs, mask=freqs_mask)
+    cos = tl.cos(freqs.to(tl.float32))
+    sin = tl.sin(freqs.to(tl.float32))
+
+    nope_offs = 0
+    if HAVE_NOPE and NOPE_FIRST:
+        nope_offs = BLOCK_D
+
+    x_mask = s_mask[:, None] & (d_offs < BLOCK_D)[None, :]
+
+    if IS_NEOX:
+        x_rotated_mask = (d_offs < BLOCK_D_HALF)[None, :]
+    else:
+        x_rotated_mask = (d_offs % 2 == 0)[None, :]
+
+    d_offs += nope_offs
+    x_offs = (
+        b * stride_x_b
+        + s_offs[:, None] * stride_x_s
+        + h * stride_x_h
+        + d_offs[None, :] * stride_x_d
+    )
+    x = tl.load(x_ptr + x_offs, mask=x_mask)
+
+    if IS_NEOX:
+        x_rotated = _get_neox_rotated_x(
+            x * sin, x_rotated_mask, BLOCK_S, BLOCK_D, BLOCK_D_HALF, True
+        )
+    else:
+        x_rotated = _get_gptj_rotated_x(
+            x * sin, x_rotated_mask, BLOCK_S, BLOCK_D, BLOCK_D_HALF, True
+        )
+
+    out_x = x * cos + x_rotated
+    out_x = out_x.to(x_ptr.dtype.element_ty)
+    x_out_offs = (
+        b * stride_out_b
+        + s_offs[:, None] * stride_out_s
+        + h * stride_out_h
+        + d_offs[None, :] * stride_out_d
+    )
+
+    tl.store(out_ptr + x_out_offs, out_x, mask=x_mask)
+
+    if HAVE_NOPE and not INPLACE:
+        if NOPE_FIRST:
+            x = tl.load(x_ptr + x_offs - BLOCK_D * stride_x_d, mask=x_mask)
+            tl.store(out_ptr + x_out_offs - BLOCK_D * stride_out_d, x, mask=x_mask)
+        else:
+            x = tl.load(x_ptr + x_offs + BLOCK_D * stride_x_d, mask=x_mask)
+            tl.store(out_ptr + x_out_offs + BLOCK_D * stride_out_d, x, mask=x_mask)
+
+
+@triton.jit
+def _rope_kernel_thd_fwd(
+    x_ptr,
+    cu_seqlens_ptr,
+    freqs_ptr,
+    out_ptr,
+    stride_x_t,
+    stride_x_h,
+    stride_x_d,
+    stride_freqs_t,
+    stride_freqs_b,
+    stride_freqs_h,
+    stride_freqs_d,
+    stride_out_t,
+    stride_out_h,
+    stride_out_d,
+    HAVE_NOPE: tl.constexpr,
+    NOPE_FIRST: tl.constexpr,
+    INPLACE: tl.constexpr,
+    REUSE_FREQS_FRONT_PART: tl.constexpr,
+    IS_NEOX: tl.constexpr,
+    BLOCK_T: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    BLOCK_D_HALF: tl.constexpr,
+):
+    b = tl.program_id(0)
+    h = tl.program_id(1)
+    pid_t = tl.program_id(2)
+
+    t_start = tl.load(cu_seqlens_ptr + b)
+    t_end = tl.load(cu_seqlens_ptr + b + 1)
+    T = t_end - t_start
+    if pid_t * BLOCK_T >= T:
+        return
+
+    t_offs = pid_t * BLOCK_T + tl.arange(0, BLOCK_T)
+    d_offs = tl.arange(0, BLOCK_D)
+    t_mask = t_offs < T
+
+    if REUSE_FREQS_FRONT_PART:
+        if IS_NEOX:
+            d_freqs_offs = tl.where(
+                (d_offs >= BLOCK_D_HALF) & (d_offs < BLOCK_D),
+                d_offs - BLOCK_D_HALF,
+                d_offs,
+            ).to(d_offs.dtype)
+            d_freqs_mask = d_freqs_offs < BLOCK_D
+        else:
+            d_freqs_offs = d_offs // 2
+            d_freqs_mask = d_freqs_offs < BLOCK_D_HALF
+    else:
+        d_freqs_offs = d_offs
+        d_freqs_mask = d_freqs_offs < BLOCK_D
+
+    freqs_mask = t_mask[:, None] & d_freqs_mask[None, :]
+    freqs_offs = (
+        t_offs[:, None] * stride_freqs_t + d_freqs_offs[None, :] * stride_freqs_d
+    )
+    freqs = tl.load(freqs_ptr + freqs_offs, mask=freqs_mask)
+    cos = tl.cos(freqs.to(tl.float32))
+    sin = tl.sin(freqs.to(tl.float32))
+
+    nope_offs = 0
+    if HAVE_NOPE and NOPE_FIRST:
+        nope_offs = BLOCK_D
+
+    x_mask = t_mask[:, None] & (d_offs < BLOCK_D)[None, :]
+
+    if IS_NEOX:
+        x_rotated_mask = (d_offs < BLOCK_D_HALF)[None, :]
+    else:
+        x_rotated_mask = (d_offs % 2 == 0)[None, :]
+
+    d_offs += nope_offs
+    x_offs = (
+        (t_start + t_offs)[:, None] * stride_x_t
+        + h * stride_x_h
+        + d_offs[None, :] * stride_x_d
+    )
+    x = tl.load(x_ptr + x_offs, mask=x_mask)
+
+    if IS_NEOX:
+        x_rotated = _get_neox_rotated_x(
+            x, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF
+        )
+    else:
+        x_rotated = _get_gptj_rotated_x(
+            x, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF
+        )
+
+    out_x = x * cos + x_rotated * sin
+    out_x = out_x.to(x_ptr.dtype.element_ty)
+    x_out_offs = (
+        (t_start + t_offs)[:, None] * stride_out_t
+        + h * stride_out_h
+        + d_offs[None, :] * stride_out_d
+    )
+
+    tl.store(out_ptr + x_out_offs, out_x, mask=x_mask)
+
+    if HAVE_NOPE and not INPLACE:
+        if NOPE_FIRST:
+            x = tl.load(x_ptr + x_offs - BLOCK_D * stride_x_d, mask=x_mask)
+            tl.store(out_ptr + x_out_offs - BLOCK_D * stride_out_d, x, mask=x_mask)
+        else:
+            x = tl.load(x_ptr + x_offs + BLOCK_D * stride_x_d, mask=x_mask)
+            tl.store(out_ptr + x_out_offs + BLOCK_D * stride_out_d, x, mask=x_mask)
+
+
+@triton.jit
+def _rope_kernel_thd_bwd(
+    x_ptr,
+    cu_seqlens_ptr,
+    freqs_ptr,
+    out_ptr,
+    stride_x_t,
+    stride_x_h,
+    stride_x_d,
+    stride_freqs_t,
+    stride_freqs_b,
+    stride_freqs_h,
+    stride_freqs_d,
+    stride_out_t,
+    stride_out_h,
+    stride_out_d,
+    HAVE_NOPE: tl.constexpr,
+    NOPE_FIRST: tl.constexpr,
+    INPLACE: tl.constexpr,
+    REUSE_FREQS_FRONT_PART: tl.constexpr,
+    IS_NEOX: tl.constexpr,
+    BLOCK_T: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    BLOCK_D_HALF: tl.constexpr,
+):
+    b = tl.program_id(0)
+    h = tl.program_id(1)
+    pid_t = tl.program_id(2)
+
+    t_start = tl.load(cu_seqlens_ptr + b)
+    t_end = tl.load(cu_seqlens_ptr + b + 1)
+    T = t_end - t_start
+    if pid_t * BLOCK_T >= T:
+        return
+
+    t_offs = pid_t * BLOCK_T + tl.arange(0, BLOCK_T)
+    d_offs = tl.arange(0, BLOCK_D)
+    t_mask = t_offs < T
+
+    if REUSE_FREQS_FRONT_PART:
+        if IS_NEOX:
+            d_freqs_offs = tl.where(
+                (d_offs >= BLOCK_D_HALF) & (d_offs < BLOCK_D),
+                d_offs - BLOCK_D_HALF,
+                d_offs,
+            ).to(d_offs.dtype)
+            d_freqs_mask = d_freqs_offs < BLOCK_D
+        else:
+            d_freqs_offs = d_offs // 2
+            d_freqs_mask = d_freqs_offs < BLOCK_D_HALF
+    else:
+        d_freqs_offs = d_offs
+        d_freqs_mask = d_freqs_offs < BLOCK_D
+
+    freqs_mask = t_mask[:, None] & d_freqs_mask[None, :]
+    freqs_offs = (
+        t_offs[:, None] * stride_freqs_t + d_freqs_offs[None, :] * stride_freqs_d
+    )
+    freqs = tl.load(freqs_ptr + freqs_offs, mask=freqs_mask)
+    cos = tl.cos(freqs.to(tl.float32))
+    sin = tl.sin(freqs.to(tl.float32))
+
+    nope_offs = 0
+    if HAVE_NOPE and NOPE_FIRST:
+        nope_offs = BLOCK_D
+
+    x_mask = t_mask[:, None] & (d_offs < BLOCK_D)[None, :]
+
+    if IS_NEOX:
+        x_rotated_mask = (d_offs < BLOCK_D_HALF)[None, :]
+    else:
+        x_rotated_mask = (d_offs % 2 == 0)[None, :]
+
+    d_offs += nope_offs
+    x_offs = (
+        (t_start + t_offs)[:, None] * stride_x_t
+        + h * stride_x_h
+        + d_offs[None, :] * stride_x_d
+    )
+    x = tl.load(x_ptr + x_offs, mask=x_mask)
+
+    if IS_NEOX:
+        x_rotated = _get_neox_rotated_x(
+            x * sin, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF, True
+        )
+    else:
+        x_rotated = _get_gptj_rotated_x(
+            x * sin, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF, True
+        )
+
+    out_x = x * cos + x_rotated
+    out_x = out_x.to(x_ptr.dtype.element_ty)
+    x_out_offs = (
+        (t_start + t_offs)[:, None] * stride_out_t
+        + h * stride_out_h
+        + d_offs[None, :] * stride_out_d
+    )
+
+    tl.store(out_ptr + x_out_offs, out_x, mask=x_mask)
+
+    if HAVE_NOPE and not INPLACE:
+        if NOPE_FIRST:
+            x = tl.load(x_ptr + x_offs - BLOCK_D * stride_x_d, mask=x_mask)
+            tl.store(out_ptr + x_out_offs - BLOCK_D * stride_out_d, x, mask=x_mask)
+        else:
+            x = tl.load(x_ptr + x_offs + BLOCK_D * stride_x_d, mask=x_mask)
+            tl.store(out_ptr + x_out_offs + BLOCK_D * stride_out_d, x, mask=x_mask)
+
+
+@triton.jit
+def _rope_kernel_sbhd_cached_fwd(
+    x_ptr,
+    cos_ptr,
+    sin_ptr,
+    pos_ptr,
+    off_ptr,
+    out_ptr,
+    stride_x_s,
+    stride_x_b,
+    stride_x_h,
+    stride_x_d,
+    stride_cos_s,
+    stride_cos_b,
+    stride_cos_h,
+    stride_cos_d,
+    stride_pos_s,
+    stride_pos_b,
+    stride_out_s,
+    stride_out_b,
+    stride_out_h,
+    stride_out_d,
+    S,
+    HAVE_NOPE: tl.constexpr,
+    NOPE_FIRST: tl.constexpr,
+    INPLACE: tl.constexpr,
+    REUSE_FREQS_FRONT_PART: tl.constexpr,
+    IS_NEOX: tl.constexpr,
+    HAVE_POS: tl.constexpr,
+    HAVE_OFFS: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    BLOCK_D_HALF: tl.constexpr,
+):
+    b = tl.program_id(0)
+    h = tl.program_id(1)
+    pid_s = tl.program_id(2)
+
+    s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S)
+    d_offs = tl.arange(0, BLOCK_D)
+    s_mask = s_offs < S
+
+    if HAVE_POS:
+        pos_offs = s_offs * stride_pos_s + b * stride_pos_b
+        pos = tl.load(pos_ptr + pos_offs, mask=s_mask)
+        if HAVE_OFFS:
+            offset = tl.load(off_ptr + pos_offs, mask=s_mask)
+            s_cos_offs = pos + offset
+        else:
+            s_cos_offs = pos
+    else:
+        s_cos_offs = s_offs
+
+    if REUSE_FREQS_FRONT_PART:
+        if IS_NEOX:
+            d_cos_offs = d_offs
+            d_cos_offs = tl.where(
+                (d_cos_offs >= BLOCK_D_HALF) & (d_cos_offs < BLOCK_D),
+                d_cos_offs - BLOCK_D_HALF,
+                d_cos_offs,
+            ).to(d_cos_offs.dtype)
+            d_cos_mask = d_cos_offs < BLOCK_D
+        else:
+            d_cos_offs = d_offs // 2
+            d_cos_mask = d_cos_offs < BLOCK_D_HALF
+    else:
+        d_cos_offs = d_offs
+        d_cos_mask = d_cos_offs < BLOCK_D
+
+    cos_mask = s_mask[:, None] & d_cos_mask[None, :]
+    cos_offs = s_cos_offs[:, None] * stride_cos_s + d_cos_offs[None, :] * stride_cos_d
+    cos = tl.load(cos_ptr + cos_offs, mask=cos_mask)
+    sin = tl.load(sin_ptr + cos_offs, mask=cos_mask)
+
+    nope_offs = 0
+    if HAVE_NOPE and NOPE_FIRST:
+        nope_offs = BLOCK_D
+
+    x_mask = s_mask[:, None] & (d_offs < BLOCK_D)[None, :]
+
+    if IS_NEOX:
+        x_rotated_mask = (d_offs < BLOCK_D_HALF)[None, :]
+    else:
+        x_rotated_mask = (d_offs % 2 == 0)[None, :]
+
+    d_offs += nope_offs
+    x_offs = (
+        b * stride_x_b
+        + s_offs[:, None] * stride_x_s
+        + h * stride_x_h
+        + d_offs[None, :] * stride_x_d
+    )
+    x = tl.load(x_ptr + x_offs, mask=x_mask)
+
+    if IS_NEOX:
+        x_rotated = _get_neox_rotated_x(
+            x, x_rotated_mask, BLOCK_S, BLOCK_D, BLOCK_D_HALF
+        )
+    else:
+        x_rotated = _get_gptj_rotated_x(
+            x, x_rotated_mask, BLOCK_S, BLOCK_D, BLOCK_D_HALF
+        )
+
+    out_x = x * cos + x_rotated * sin
+    out_x = out_x.to(x_ptr.dtype.element_ty)
+    x_out_offs = (
+        b * stride_out_b
+        + s_offs[:, None] * stride_out_s
+        + h * stride_out_h
+        + d_offs[None, :] * stride_out_d
+    )
+
+    tl.store(out_ptr + x_out_offs, out_x, mask=x_mask)
+
+    if HAVE_NOPE and not INPLACE:
+        if NOPE_FIRST:
+            x = tl.load(x_ptr + x_offs - BLOCK_D * stride_x_d, mask=x_mask)
+            tl.store(out_ptr + x_out_offs - BLOCK_D * stride_out_d, x, mask=x_mask)
+        else:
+            x = tl.load(x_ptr + x_offs + BLOCK_D * stride_x_d, mask=x_mask)
+            tl.store(out_ptr + x_out_offs + BLOCK_D * stride_out_d, x, mask=x_mask)
+
+
+@triton.jit
+def _rope_kernel_sbhd_cached_bwd(
+    x_ptr,
+    cos_ptr,
+    sin_ptr,
+    pos_ptr,
+    off_ptr,
+    out_ptr,
+    stride_x_s,
+    stride_x_b,
+    stride_x_h,
+    stride_x_d,
+    stride_cos_s,
+    stride_cos_b,
+    stride_cos_h,
+    stride_cos_d,
+    stride_pos_s,
+    stride_pos_b,
+    stride_out_s,
+    stride_out_b,
+    stride_out_h,
+    stride_out_d,
+    S,
+    HAVE_NOPE: tl.constexpr,
+    NOPE_FIRST: tl.constexpr,
+    INPLACE: tl.constexpr,
+    REUSE_FREQS_FRONT_PART: tl.constexpr,
+    IS_NEOX: tl.constexpr,
+    HAVE_POS: tl.constexpr,
+    HAVE_OFFS: tl.constexpr,
+    BLOCK_S: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    BLOCK_D_HALF: tl.constexpr,
+):
+    b = tl.program_id(0)
+    h = tl.program_id(1)
+    pid_s = tl.program_id(2)
+
+    s_offs = pid_s * BLOCK_S + tl.arange(0, BLOCK_S)
+    d_offs = tl.arange(0, BLOCK_D)
+    s_mask = s_offs < S
+
+    if HAVE_POS:
+        pos_offs = s_offs * stride_pos_s + b * stride_pos_b
+        pos = tl.load(pos_ptr + pos_offs, mask=s_mask)
+        if HAVE_OFFS:
+            offset = tl.load(off_ptr + pos_offs, mask=s_mask)
+            s_cos_offs = pos + offset
+        else:
+            s_cos_offs = pos
+    else:
+        s_cos_offs = s_offs
+
+    if REUSE_FREQS_FRONT_PART:
+        if IS_NEOX:
+            d_cos_offs = d_offs
+            d_cos_offs = tl.where(
+                (d_cos_offs >= BLOCK_D_HALF) & (d_cos_offs < BLOCK_D),
+                d_cos_offs - BLOCK_D_HALF,
+                d_cos_offs,
+            ).to(d_cos_offs.dtype)
+            d_cos_mask = d_cos_offs < BLOCK_D
+        else:
+            d_cos_offs = d_offs // 2
+            d_cos_mask = d_cos_offs < BLOCK_D_HALF
+    else:
+        d_cos_offs = d_offs
+        d_cos_mask = d_cos_offs < BLOCK_D
+
+    cos_mask = s_mask[:, None] & d_cos_mask[None, :]
+    cos_offs = s_cos_offs[:, None] * stride_cos_s + d_cos_offs[None, :] * stride_cos_d
+    cos = tl.load(cos_ptr + cos_offs, mask=cos_mask)
+    sin = tl.load(sin_ptr + cos_offs, mask=cos_mask)
+
+    nope_offs = 0
+    if HAVE_NOPE and NOPE_FIRST:
+        nope_offs = BLOCK_D
+
+    x_mask = s_mask[:, None] & (d_offs < BLOCK_D)[None, :]
+
+    if IS_NEOX:
+        x_rotated_mask = (d_offs < BLOCK_D_HALF)[None, :]
+    else:
+        x_rotated_mask = (d_offs % 2 == 0)[None, :]
+
+    d_offs += nope_offs
+    x_offs = (
+        b * stride_x_b
+        + s_offs[:, None] * stride_x_s
+        + h * stride_x_h
+        + d_offs[None, :] * stride_x_d
+    )
+    x = tl.load(x_ptr + x_offs, mask=x_mask)
+
+    if IS_NEOX:
+        x_rotated = _get_neox_rotated_x(
+            x * sin, x_rotated_mask, BLOCK_S, BLOCK_D, BLOCK_D_HALF, True
+        )
+    else:
+        x_rotated = _get_gptj_rotated_x(
+            x * sin, x_rotated_mask, BLOCK_S, BLOCK_D, BLOCK_D_HALF, True
+        )
+
+    out_x = x * cos + x_rotated
+    out_x = out_x.to(x_ptr.dtype.element_ty)
+    x_out_offs = (
+        b * stride_out_b
+        + s_offs[:, None] * stride_out_s
+        + h * stride_out_h
+        + d_offs[None, :] * stride_out_d
+    )
+
+    tl.store(out_ptr + x_out_offs, out_x, mask=x_mask)
+
+    if HAVE_NOPE and not INPLACE:
+        if NOPE_FIRST:
+            x = tl.load(x_ptr + x_offs - BLOCK_D * stride_x_d, mask=x_mask)
+            tl.store(out_ptr + x_out_offs - BLOCK_D * stride_out_d, x, mask=x_mask)
+        else:
+            x = tl.load(x_ptr + x_offs + BLOCK_D * stride_x_d, mask=x_mask)
+            tl.store(out_ptr + x_out_offs + BLOCK_D * stride_out_d, x, mask=x_mask)
+
+
+@triton.jit
+def _rope_kernel_thd_cached_2c_fwd(
+    x_ptr,
+    y_ptr,
+    cos_ptr,
+    sin_ptr,
+    pos_ptr,
+    off_ptr,
+    out_x_ptr,
+    out_y_ptr,
+    stride_x_t,
+    stride_x_h,
+    stride_x_d,
+    stride_y_t,
+    stride_y_h,
+    stride_y_d,
+    stride_cos_t,
+    stride_cos_d,
+    stride_pos_t,
+    stride_out_x_t,
+    stride_out_x_h,
+    stride_out_x_d,
+    stride_out_y_t,
+    stride_out_y_h,
+    stride_out_y_d,
+    T,
+    HAVE_NOPE: tl.constexpr,
+    NOPE_FIRST: tl.constexpr,
+    INPLACE: tl.constexpr,
+    REUSE_FREQS_FRONT_PART: tl.constexpr,
+    IS_NEOX: tl.constexpr,
+    HAVE_POS: tl.constexpr,
+    HAVE_OFFS: tl.constexpr,
+    BLOCK_T: tl.constexpr,
+    SPLIT_H_SIZE: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    BLOCK_D_HALF: tl.constexpr,
+    num_stages: tl.constexpr,
+):
+    h_s = tl.program_id(0)
+    pid_t = tl.program_id(1)
+
+    tl.assume(stride_x_t > 0)
+    tl.assume(stride_x_h > 0)
+    tl.assume(stride_x_d > 0)
+    tl.assume(stride_y_t > 0)
+    tl.assume(stride_y_h > 0)
+    tl.assume(stride_y_d > 0)
+    tl.assume(stride_cos_t > 0)
+    tl.assume(stride_cos_d > 0)
+    tl.assume(stride_pos_t > 0)
+    tl.assume(stride_out_x_t > 0)
+    tl.assume(stride_out_x_h > 0)
+    tl.assume(stride_out_x_d > 0)
+    tl.assume(stride_out_y_t > 0)
+    tl.assume(stride_out_y_h > 0)
+    tl.assume(stride_out_y_d > 0)
+
+    t_offs = pid_t * BLOCK_T + tl.arange(0, BLOCK_T)
+    d_offs = tl.arange(0, BLOCK_D)
+    t_mask = t_offs < T
+
+    if HAVE_POS:
+        pos_offs = t_offs * stride_pos_t
+        pos = tl.load(pos_ptr + pos_offs, mask=t_mask)
+        if HAVE_OFFS:
+            offset = tl.load(off_ptr + pos_offs, mask=t_mask)
+            t_cos_offs = pos + offset
+        else:
+            t_cos_offs = pos
+    else:
+        t_cos_offs = t_offs
+
+    if REUSE_FREQS_FRONT_PART:
+        if IS_NEOX:
+            d_cos_offs = d_offs
+            d_cos_offs = tl.where(
+                (d_cos_offs < BLOCK_D_HALF),
+                d_cos_offs,
+                d_cos_offs - BLOCK_D_HALF,
+            ).to(d_cos_offs.dtype)
+            d_cos_mask = d_cos_offs < BLOCK_D_HALF
+        else:
+            d_cos_offs = tl.arange(0, BLOCK_D) // 2
+            d_cos_mask = d_cos_offs < BLOCK_D_HALF
+    else:
+        d_cos_offs = d_offs
+        d_cos_mask = d_cos_offs < BLOCK_D
+
+    cos_mask = t_mask[:, None] & d_cos_mask[None, :]
+    cos_offs = t_cos_offs[:, None] * stride_cos_t + d_cos_offs[None, :] * stride_cos_d
+    cos = tl.load(cos_ptr + cos_offs, mask=cos_mask)
+    sin = tl.load(sin_ptr + cos_offs, mask=cos_mask)
+
+    nope_offs = 0
+    if HAVE_NOPE and NOPE_FIRST:
+        nope_offs = BLOCK_D
+
+    h_start_idx = h_s * SPLIT_H_SIZE
+
+    x_mask = t_mask[:, None] & (d_offs < BLOCK_D)[None, :]
+
+    if IS_NEOX:
+        x_rotated_mask = (d_offs < BLOCK_D_HALF)[None, :]
+    else:
+        x_rotated_mask = (d_offs % 2 == 0)[None, :]
+
+    d_offs += nope_offs
+    for h in tl.range(0, SPLIT_H_SIZE, 1, num_stages=num_stages):
+        x_offs = (
+            t_offs[:, None] * stride_x_t
+            + d_offs[None, :] * stride_x_d
+            + (h_start_idx + h) * stride_x_h
+        )
+        y_offs = (
+            t_offs[:, None] * stride_y_t
+            + d_offs[None, :] * stride_y_d
+            + (h_start_idx + h) * stride_y_h
+        )
+
+        x = tl.load(x_ptr + x_offs, mask=x_mask)
+        y = tl.load(y_ptr + y_offs, mask=x_mask)
+
+        if IS_NEOX:
+            x_rotated = _get_neox_rotated_x(
+                x, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF
+            )
+            y_rotated = _get_neox_rotated_x(
+                y, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF
+            )
+        else:
+            x_rotated = _get_gptj_rotated_x(
+                x, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF
+            )
+            y_rotated = _get_gptj_rotated_x(
+                y, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF
+            )
+
+        out_x = x * cos + x_rotated * sin
+        out_x = out_x.to(x_ptr.dtype.element_ty)
+        out_y = y * cos + y_rotated * sin
+        out_y = out_y.to(y_ptr.dtype.element_ty)
+
+        out_x_offs = (
+            t_offs[:, None] * stride_out_x_t
+            + d_offs[None, :] * stride_out_x_d
+            + (h_start_idx + h) * stride_out_x_h
+        )
+        out_y_offs = (
+            t_offs[:, None] * stride_out_y_t
+            + d_offs[None, :] * stride_out_y_d
+            + (h_start_idx + h) * stride_out_y_h
+        )
+        tl.store(out_x_ptr + out_x_offs, out_x, mask=x_mask)
+        tl.store(out_y_ptr + out_y_offs, out_y, mask=x_mask)
+
+        if HAVE_NOPE and not INPLACE:
+            if NOPE_FIRST:
+                x = tl.load(x_ptr + x_offs - BLOCK_D * stride_x_d, mask=x_mask)
+                tl.store(
+                    out_x_ptr + out_x_offs - BLOCK_D * stride_out_x_d, x, mask=x_mask
+                )
+                y = tl.load(y_ptr + y_offs - BLOCK_D * stride_y_d, mask=x_mask)
+                tl.store(
+                    out_y_ptr + out_y_offs - BLOCK_D * stride_out_y_d, y, mask=x_mask
+                )
+            else:
+                x = tl.load(x_ptr + x_offs + BLOCK_D * stride_x_d, mask=x_mask)
+                tl.store(
+                    out_x_ptr + out_x_offs + BLOCK_D * stride_out_x_d, x, mask=x_mask
+                )
+                y = tl.load(y_ptr + y_offs + BLOCK_D * stride_y_d, mask=x_mask)
+                tl.store(
+                    out_y_ptr + out_y_offs + BLOCK_D * stride_out_y_d, y, mask=x_mask
+                )
+
+
+@triton.jit
+def _rope_kernel_thd_cached_2c_bwd(
+    x_ptr,
+    y_ptr,
+    cos_ptr,
+    sin_ptr,
+    pos_ptr,
+    off_ptr,
+    out_x_ptr,
+    out_y_ptr,
+    stride_x_t,
+    stride_x_h,
+    stride_x_d,
+    stride_y_t,
+    stride_y_h,
+    stride_y_d,
+    stride_cos_t,
+    stride_cos_d,
+    stride_pos_t,
+    stride_out_x_t,
+    stride_out_x_h,
+    stride_out_x_d,
+    stride_out_y_t,
+    stride_out_y_h,
+    stride_out_y_d,
+    T,
+    HAVE_NOPE: tl.constexpr,
+    NOPE_FIRST: tl.constexpr,
+    INPLACE: tl.constexpr,
+    REUSE_FREQS_FRONT_PART: tl.constexpr,
+    IS_NEOX: tl.constexpr,
+    HAVE_POS: tl.constexpr,
+    HAVE_OFFS: tl.constexpr,
+    BLOCK_T: tl.constexpr,
+    SPLIT_H_SIZE: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    BLOCK_D_HALF: tl.constexpr,
+    num_stages: tl.constexpr,
+):
+    h_s = tl.program_id(0)
+    pid_t = tl.program_id(1)
+
+    tl.assume(stride_x_t > 0)
+    tl.assume(stride_x_h > 0)
+    tl.assume(stride_x_d > 0)
+    tl.assume(stride_y_t > 0)
+    tl.assume(stride_y_h > 0)
+    tl.assume(stride_y_d > 0)
+    tl.assume(stride_cos_t > 0)
+    tl.assume(stride_cos_d > 0)
+    tl.assume(stride_pos_t > 0)
+    tl.assume(stride_out_x_t > 0)
+    tl.assume(stride_out_x_h > 0)
+    tl.assume(stride_out_x_d > 0)
+    tl.assume(stride_out_y_t > 0)
+    tl.assume(stride_out_y_h > 0)
+    tl.assume(stride_out_y_d > 0)
+
+    t_offs = pid_t * BLOCK_T + tl.arange(0, BLOCK_T)
+    d_offs = tl.arange(0, BLOCK_D)
+    t_mask = t_offs < T
+
+    if HAVE_POS:
+        pos_offs = t_offs * stride_pos_t
+        pos = tl.load(pos_ptr + pos_offs, mask=t_mask)
+        if HAVE_OFFS:
+            offset = tl.load(off_ptr + pos_offs, mask=t_mask)
+            t_cos_offs = pos + offset
+        else:
+            t_cos_offs = pos
+    else:
+        t_cos_offs = t_offs
+
+    if REUSE_FREQS_FRONT_PART:
+        if IS_NEOX:
+            d_cos_offs = d_offs
+            d_cos_offs = tl.where(
+                (d_cos_offs < BLOCK_D_HALF),
+                d_cos_offs,
+                d_cos_offs - BLOCK_D_HALF,
+            ).to(d_cos_offs.dtype)
+            d_cos_mask = d_cos_offs < BLOCK_D_HALF
+        else:
+            d_cos_offs = tl.arange(0, BLOCK_D) // 2
+            d_cos_mask = d_cos_offs < BLOCK_D_HALF
+    else:
+        d_cos_offs = d_offs
+        d_cos_mask = d_cos_offs < BLOCK_D
+
+    cos_mask = t_mask[:, None] & d_cos_mask[None, :]
+    cos_offs = t_cos_offs[:, None] * stride_cos_t + d_cos_offs[None, :] * stride_cos_d
+    cos = tl.load(cos_ptr + cos_offs, mask=cos_mask)
+    sin = tl.load(sin_ptr + cos_offs, mask=cos_mask)
+
+    nope_offs = 0
+    if HAVE_NOPE and NOPE_FIRST:
+        nope_offs = BLOCK_D
+
+    h_start_idx = h_s * SPLIT_H_SIZE
+
+    x_mask = t_mask[:, None] & (d_offs < BLOCK_D)[None, :]
+
+    if IS_NEOX:
+        x_rotated_mask = (d_offs < BLOCK_D_HALF)[None, :]
+    else:
+        x_rotated_mask = (d_offs % 2 == 0)[None, :]
+
+    d_offs += nope_offs
+    for h in tl.range(0, SPLIT_H_SIZE, 1, num_stages=num_stages):
+        x_offs = (
+            t_offs[:, None] * stride_x_t
+            + d_offs[None, :] * stride_x_d
+            + (h_start_idx + h) * stride_x_h
+        )
+        y_offs = (
+            t_offs[:, None] * stride_y_t
+            + d_offs[None, :] * stride_y_d
+            + (h_start_idx + h) * stride_y_h
+        )
+
+        x = tl.load(x_ptr + x_offs, mask=x_mask)
+        y = tl.load(y_ptr + y_offs, mask=x_mask)
+
+        if IS_NEOX:
+            x_rotated = _get_neox_rotated_x(
+                x * sin, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF, True
+            )
+            y_rotated = _get_neox_rotated_x(
+                y * sin, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF, True
+            )
+        else:
+            x_rotated = _get_gptj_rotated_x(
+                x * sin, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF, True
+            )
+            y_rotated = _get_gptj_rotated_x(
+                y * sin, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF, True
+            )
+
+        out_x = x * cos + x_rotated
+        out_x = out_x.to(x_ptr.dtype.element_ty)
+        out_y = y * cos + y_rotated
+        out_y = out_y.to(y_ptr.dtype.element_ty)
+
+        out_x_offs = (
+            t_offs[:, None] * stride_out_x_t
+            + d_offs[None, :] * stride_out_x_d
+            + (h_start_idx + h) * stride_out_x_h
+        )
+        out_y_offs = (
+            t_offs[:, None] * stride_out_y_t
+            + d_offs[None, :] * stride_out_y_d
+            + (h_start_idx + h) * stride_out_y_h
+        )
+        tl.store(out_x_ptr + out_x_offs, out_x, mask=x_mask)
+        tl.store(out_y_ptr + out_y_offs, out_y, mask=x_mask)
+
+        if HAVE_NOPE and not INPLACE:
+            # TODO check
+            if NOPE_FIRST:
+                x = tl.load(x_ptr + x_offs - BLOCK_D * stride_x_d, mask=x_mask)
+                tl.store(
+                    out_x_ptr + out_x_offs - BLOCK_D * stride_out_x_d, x, mask=x_mask
+                )
+                y = tl.load(y_ptr + y_offs - BLOCK_D * stride_y_d, mask=x_mask)
+                tl.store(
+                    out_y_ptr + out_y_offs - BLOCK_D * stride_out_y_d, y, mask=x_mask
+                )
+            else:
+                x = tl.load(x_ptr + x_offs + BLOCK_D * stride_x_d, mask=x_mask)
+                tl.store(
+                    out_x_ptr + out_x_offs + BLOCK_D * stride_out_x_d, x, mask=x_mask
+                )
+                y = tl.load(y_ptr + y_offs + BLOCK_D * stride_y_d, mask=x_mask)
+                tl.store(
+                    out_y_ptr + out_y_offs + BLOCK_D * stride_out_y_d, y, mask=x_mask
+                )
+
+
+@triton.jit
+def _rope_kernel_cached_thd_2c_gqa_fwd(
+    x_ptr,
+    y_ptr,
+    cos_ptr,
+    sin_ptr,
+    pos_ptr,
+    off_ptr,
+    out_x_ptr,
+    out_y_ptr,
+    stride_x_t,
+    stride_x_h,
+    stride_x_d,
+    stride_y_t,
+    stride_y_h,
+    stride_y_d,
+    stride_cos_t,
+    stride_cos_d,
+    stride_pos_t,
+    stride_out_x_t,
+    stride_out_x_h,
+    stride_out_x_d,
+    stride_out_y_t,
+    stride_out_y_h,
+    stride_out_y_d,
+    T,
+    HAVE_NOPE: tl.constexpr,
+    NOPE_FIRST: tl.constexpr,
+    INPLACE: tl.constexpr,
+    REUSE_FREQS_FRONT_PART: tl.constexpr,
+    IS_NEOX: tl.constexpr,
+    HAVE_POS: tl.constexpr,
+    HAVE_OFFS: tl.constexpr,
+    BLOCK_T: tl.constexpr,
+    QH_per_G: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    BLOCK_D_HALF: tl.constexpr,
+    num_stages: tl.constexpr,
+):
+    h_s = tl.program_id(0)
+    pid_t = tl.program_id(1)
+
+    tl.assume(stride_x_t > 0)
+    tl.assume(stride_x_h > 0)
+    tl.assume(stride_x_d > 0)
+    tl.assume(stride_y_t > 0)
+    tl.assume(stride_y_h > 0)
+    tl.assume(stride_y_d > 0)
+    tl.assume(stride_cos_t > 0)
+    tl.assume(stride_cos_d > 0)
+    tl.assume(stride_pos_t > 0)
+    tl.assume(stride_out_x_t > 0)
+    tl.assume(stride_out_x_h > 0)
+    tl.assume(stride_out_x_d > 0)
+    tl.assume(stride_out_y_t > 0)
+    tl.assume(stride_out_y_h > 0)
+    tl.assume(stride_out_y_d > 0)
+
+    t_offs = pid_t * BLOCK_T + tl.arange(0, BLOCK_T)
+    d_offs = tl.arange(0, BLOCK_D)
+    t_mask = t_offs < T
+
+    if HAVE_POS:
+        pos_offs = t_offs * stride_pos_t
+        pos = tl.load(pos_ptr + pos_offs, mask=t_mask)
+        if HAVE_OFFS:
+            offset = tl.load(off_ptr + pos_offs, mask=t_mask)
+            t_cos_offs = pos + offset
+        else:
+            t_cos_offs = pos
+    else:
+        t_cos_offs = t_offs
+
+    if REUSE_FREQS_FRONT_PART:
+        if IS_NEOX:
+            d_cos_offs = d_offs
+            d_cos_offs = tl.where(
+                (d_cos_offs < BLOCK_D_HALF),
+                d_cos_offs,
+                d_cos_offs - BLOCK_D_HALF,
+            ).to(d_cos_offs.dtype)
+            d_cos_mask = d_cos_offs < BLOCK_D_HALF
+        else:
+            d_cos_offs = tl.arange(0, BLOCK_D) // 2
+            d_cos_mask = d_cos_offs < BLOCK_D_HALF
+    else:
+        d_cos_offs = d_offs
+        d_cos_mask = d_cos_offs < BLOCK_D
+
+    cos_mask = t_mask[:, None] & d_cos_mask[None, :]
+    cos_offs = t_cos_offs[:, None] * stride_cos_t + d_cos_offs[None, :] * stride_cos_d
+    cos = tl.load(cos_ptr + cos_offs, mask=cos_mask)
+    sin = tl.load(sin_ptr + cos_offs, mask=cos_mask)
+
+    nope_offs = 0
+    if HAVE_NOPE and NOPE_FIRST:
+        nope_offs = BLOCK_D
+
+    h_start_idx = h_s * QH_per_G
+
+    x_mask = t_mask[:, None] & (d_offs < BLOCK_D)[None, :]
+
+    if IS_NEOX:
+        x_rotated_mask = (d_offs < BLOCK_D_HALF)[None, :]
+    else:
+        x_rotated_mask = (d_offs % 2 == 0)[None, :]
+
+    d_offs += nope_offs
+    y_offs = (
+        t_offs[:, None] * stride_y_t + d_offs[None, :] * stride_y_d + h_s * stride_y_h
+    )
+    y = tl.load(y_ptr + y_offs, mask=x_mask)
+
+    if IS_NEOX:
+        y_rotated = _get_neox_rotated_x(
+            y, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF
+        )
+    else:
+        y_rotated = _get_gptj_rotated_x(
+            y, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF
+        )
+
+    out_y_offs = (
+        t_offs[:, None] * stride_out_y_t
+        + d_offs[None, :] * stride_out_y_d
+        + h_s * stride_out_y_h
+    )
+    out_y = y * cos + y_rotated * sin
+    out_y = out_y.to(y_ptr.dtype.element_ty)
+    tl.store(out_y_ptr + out_y_offs, out_y, mask=x_mask)
+
+    if HAVE_NOPE and not INPLACE:
+        if NOPE_FIRST:
+            y = tl.load(y_ptr + y_offs - BLOCK_D * stride_y_d, mask=x_mask)
+            tl.store(out_y_ptr + out_y_offs - BLOCK_D * stride_out_y_d, y, mask=x_mask)
+        else:
+            y = tl.load(y_ptr + y_offs + BLOCK_D * stride_y_d, mask=x_mask)
+            tl.store(out_y_ptr + out_y_offs + BLOCK_D * stride_out_y_d, y, mask=x_mask)
+
+    for h in tl.range(0, QH_per_G, 1, num_stages=num_stages):
+        x_offs = (
+            t_offs[:, None] * stride_x_t
+            + d_offs[None, :] * stride_x_d
+            + (h_start_idx + h) * stride_x_h
+        )
+
+        x = tl.load(x_ptr + x_offs, mask=x_mask)
+
+        if IS_NEOX:
+            x_rotated = _get_neox_rotated_x(
+                x, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF
+            )
+        else:
+            x_rotated = _get_gptj_rotated_x(
+                x, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF
+            )
+
+        out_x_offs = (
+            t_offs[:, None] * stride_out_x_t
+            + d_offs[None, :] * stride_out_x_d
+            + (h_start_idx + h) * stride_out_x_h
+        )
+        out_x = x * cos + x_rotated * sin
+        out_x = out_x.to(x_ptr.dtype.element_ty)
+
+        tl.store(out_x_ptr + out_x_offs, out_x, mask=x_mask)
+
+        if HAVE_NOPE and not INPLACE:
+            if NOPE_FIRST:
+                x = tl.load(x_ptr + x_offs - BLOCK_D * stride_x_d, mask=x_mask)
+                tl.store(
+                    out_x_ptr + out_x_offs - BLOCK_D * stride_out_x_d, x, mask=x_mask
+                )
+            else:
+                x = tl.load(x_ptr + x_offs + BLOCK_D * stride_x_d, mask=x_mask)
+                tl.store(
+                    out_x_ptr + out_x_offs + BLOCK_D * stride_out_x_d, x, mask=x_mask
+                )
+
+
+@triton.jit
+def _rope_kernel_cached_thd_2c_gqa_onehead_fwd(
+    x_ptr,
+    y_ptr,
+    cos_ptr,
+    sin_ptr,
+    pos_ptr,
+    off_ptr,
+    out_x_ptr,
+    out_y_ptr,
+    stride_x_t,
+    stride_x_h,
+    stride_x_d,
+    stride_y_t,
+    stride_y_h,
+    stride_y_d,
+    stride_cos_t,
+    stride_cos_d,
+    stride_pos_t,
+    stride_out_x_t,
+    stride_out_x_h,
+    stride_out_x_d,
+    stride_out_y_t,
+    stride_out_y_h,
+    stride_out_y_d,
+    T,
+    HAVE_NOPE: tl.constexpr,
+    NOPE_FIRST: tl.constexpr,
+    INPLACE: tl.constexpr,
+    REUSE_FREQS_FRONT_PART: tl.constexpr,
+    IS_NEOX: tl.constexpr,
+    HAVE_POS: tl.constexpr,
+    HAVE_OFFS: tl.constexpr,
+    BLOCK_T: tl.constexpr,
+    G: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    BLOCK_D_HALF: tl.constexpr,
+):
+    pid_t = tl.program_id(0)
+    hq = tl.program_id(1)
+
+    tl.assume(stride_x_t > 0)
+    tl.assume(stride_x_h > 0)
+    tl.assume(stride_x_d > 0)
+    tl.assume(stride_y_t > 0)
+    tl.assume(stride_y_h > 0)
+    tl.assume(stride_y_d > 0)
+    tl.assume(stride_cos_t > 0)
+    tl.assume(stride_cos_d > 0)
+    tl.assume(stride_pos_t > 0)
+    tl.assume(stride_out_x_t > 0)
+    tl.assume(stride_out_x_h > 0)
+    tl.assume(stride_out_x_d > 0)
+    tl.assume(stride_out_y_t > 0)
+    tl.assume(stride_out_y_h > 0)
+    tl.assume(stride_out_y_d > 0)
+
+    t_offs = pid_t * BLOCK_T + tl.arange(0, BLOCK_T)
+    d_offs = tl.arange(0, BLOCK_D)
+    t_mask = t_offs < T
+
+    if HAVE_POS:
+        pos_offs = t_offs * stride_pos_t
+        pos = tl.load(pos_ptr + pos_offs, mask=t_mask)
+        if HAVE_OFFS:
+            offset = tl.load(off_ptr + pos_offs, mask=t_mask)
+            t_cos_offs = pos + offset
+        else:
+            t_cos_offs = pos
+    else:
+        t_cos_offs = t_offs
+
+    if REUSE_FREQS_FRONT_PART:
+        if IS_NEOX:
+            d_cos_offs = d_offs
+            d_cos_offs = tl.where(
+                (d_cos_offs < BLOCK_D_HALF),
+                d_cos_offs,
+                d_cos_offs - BLOCK_D_HALF,
+            ).to(d_cos_offs.dtype)
+            d_cos_mask = d_cos_offs < BLOCK_D_HALF
+        else:
+            d_cos_offs = tl.arange(0, BLOCK_D) // 2
+            d_cos_mask = d_cos_offs < BLOCK_D_HALF
+    else:
+        d_cos_offs = d_offs
+        d_cos_mask = d_cos_offs < BLOCK_D
+
+    cos_mask = t_mask[:, None] & d_cos_mask[None, :]
+    cos_offs = t_cos_offs[:, None] * stride_cos_t + d_cos_offs[None, :] * stride_cos_d
+    cos = tl.load(cos_ptr + cos_offs, mask=cos_mask)
+    sin = tl.load(sin_ptr + cos_offs, mask=cos_mask)
+
+    nope_offs = 0
+    if HAVE_NOPE and NOPE_FIRST:
+        nope_offs = BLOCK_D
+
+    x_mask = t_mask[:, None] & (d_offs < BLOCK_D)[None, :]
+
+    if IS_NEOX:
+        x_rotated_mask = (d_offs < BLOCK_D_HALF)[None, :]
+    else:
+        x_rotated_mask = (d_offs % 2 == 0)[None, :]
+
+    d_offs += nope_offs
+    x_offs = (
+        t_offs[:, None] * stride_x_t + d_offs[None, :] * stride_x_d + hq * stride_x_h
+    )
+    x = tl.load(x_ptr + x_offs, mask=x_mask)
+
+    if IS_NEOX:
+        x_rotated = _get_neox_rotated_x(
+            x, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF
+        )
+    else:
+        x_rotated = _get_gptj_rotated_x(
+            x, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF
+        )
+
+    out_x_offs = (
+        t_offs[:, None] * stride_out_x_t
+        + d_offs[None, :] * stride_out_x_d
+        + hq * stride_out_x_h
+    )
+    out_x = x * cos + x_rotated * sin
+    out_x = out_x.to(x_ptr.dtype.element_ty)
+    tl.store(out_x_ptr + out_x_offs, out_x, mask=x_mask)
+
+    if HAVE_NOPE and not INPLACE:
+        if NOPE_FIRST:
+            x = tl.load(x_ptr + x_offs - BLOCK_D * stride_x_d, mask=x_mask)
+            tl.store(out_x_ptr + out_x_offs - BLOCK_D * stride_out_x_d, x, mask=x_mask)
+        else:
+            x = tl.load(x_ptr + x_offs + BLOCK_D * stride_x_d, mask=x_mask)
+            tl.store(out_x_ptr + out_x_offs + BLOCK_D * stride_out_x_d, x, mask=x_mask)
+
+    if hq < G:
+        y_offs = (
+            t_offs[:, None] * stride_y_t
+            + d_offs[None, :] * stride_x_d
+            + hq * stride_y_h
+        )
+        y = tl.load(y_ptr + y_offs, mask=x_mask)
+
+        if IS_NEOX:
+            y_rotated = _get_neox_rotated_x(
+                y, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF
+            )
+        else:
+            y_rotated = _get_gptj_rotated_x(
+                y, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF
+            )
+
+        out_y_offs = (
+            t_offs[:, None] * stride_out_y_t
+            + d_offs[None, :] * stride_out_y_d
+            + hq * stride_out_y_h
+        )
+        out_y = y * cos + y_rotated * sin
+        out_y = out_y.to(y_ptr.dtype.element_ty)
+        tl.store(out_y_ptr + out_y_offs, out_y, mask=x_mask)
+
+        if HAVE_NOPE and not INPLACE:
+            if NOPE_FIRST:
+                y = tl.load(y_ptr + y_offs - BLOCK_D * stride_y_d, mask=x_mask)
+                tl.store(
+                    out_y_ptr + out_y_offs - BLOCK_D * stride_out_y_d, y, mask=x_mask
+                )
+            else:
+                y = tl.load(y_ptr + y_offs + BLOCK_D * stride_y_d, mask=x_mask)
+                tl.store(
+                    out_y_ptr + out_y_offs + BLOCK_D * stride_out_y_d, y, mask=x_mask
+                )
+
+
+@triton.jit
+def _rope_kernel_cached_thd_2c_gqa_bwd(
+    x_ptr,
+    y_ptr,
+    cos_ptr,
+    sin_ptr,
+    pos_ptr,
+    off_ptr,
+    out_x_ptr,
+    out_y_ptr,
+    stride_x_t,
+    stride_x_h,
+    stride_x_d,
+    stride_y_t,
+    stride_y_h,
+    stride_y_d,
+    stride_cos_t,
+    stride_cos_d,
+    stride_pos_t,
+    stride_out_x_t,
+    stride_out_x_h,
+    stride_out_x_d,
+    stride_out_y_t,
+    stride_out_y_h,
+    stride_out_y_d,
+    T,
+    HAVE_NOPE: tl.constexpr,
+    NOPE_FIRST: tl.constexpr,
+    INPLACE: tl.constexpr,
+    REUSE_FREQS_FRONT_PART: tl.constexpr,
+    IS_NEOX: tl.constexpr,
+    HAVE_POS: tl.constexpr,
+    HAVE_OFFS: tl.constexpr,
+    BLOCK_T: tl.constexpr,
+    QH_per_G: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    BLOCK_D_HALF: tl.constexpr,
+    num_stages: tl.constexpr,
+):
+    h_s = tl.program_id(0)
+    pid_t = tl.program_id(1)
+
+    tl.assume(stride_x_t > 0)
+    tl.assume(stride_x_h > 0)
+    tl.assume(stride_x_d > 0)
+    tl.assume(stride_y_t > 0)
+    tl.assume(stride_y_h > 0)
+    tl.assume(stride_y_d > 0)
+    tl.assume(stride_cos_t > 0)
+    tl.assume(stride_cos_d > 0)
+    tl.assume(stride_pos_t > 0)
+    tl.assume(stride_out_x_t > 0)
+    tl.assume(stride_out_x_h > 0)
+    tl.assume(stride_out_x_d > 0)
+    tl.assume(stride_out_y_t > 0)
+    tl.assume(stride_out_y_h > 0)
+    tl.assume(stride_out_y_d > 0)
+
+    t_offs = pid_t * BLOCK_T + tl.arange(0, BLOCK_T)
+    d_offs = tl.arange(0, BLOCK_D)
+    t_mask = t_offs < T
+
+    if HAVE_POS:
+        pos_offs = t_offs * stride_pos_t
+        pos = tl.load(pos_ptr + pos_offs, mask=t_mask)
+        if HAVE_OFFS:
+            offset = tl.load(off_ptr + pos_offs, mask=t_mask)
+            t_cos_offs = pos + offset
+        else:
+            t_cos_offs = pos
+    else:
+        t_cos_offs = t_offs
+
+    if REUSE_FREQS_FRONT_PART:
+        if IS_NEOX:
+            d_cos_offs = d_offs
+            d_cos_offs = tl.where(
+                (d_cos_offs < BLOCK_D_HALF),
+                d_cos_offs,
+                d_cos_offs - BLOCK_D_HALF,
+            ).to(d_cos_offs.dtype)
+            d_cos_mask = d_cos_offs < BLOCK_D_HALF
+        else:
+            d_cos_offs = tl.arange(0, BLOCK_D) // 2
+            d_cos_mask = d_cos_offs < BLOCK_D_HALF
+    else:
+        d_cos_offs = d_offs
+        d_cos_mask = d_cos_offs < BLOCK_D
+
+    cos_mask = t_mask[:, None] & d_cos_mask[None, :]
+    cos_offs = t_cos_offs[:, None] * stride_cos_t + d_cos_offs[None, :] * stride_cos_d
+    cos = tl.load(cos_ptr + cos_offs, mask=cos_mask)
+    sin = tl.load(sin_ptr + cos_offs, mask=cos_mask)
+
+    nope_offs = 0
+    if HAVE_NOPE and NOPE_FIRST:
+        nope_offs = BLOCK_D
+
+    h_start_idx = h_s * QH_per_G
+
+    x_mask = t_mask[:, None] & (d_offs < BLOCK_D)[None, :]
+
+    if IS_NEOX:
+        x_rotated_mask = (d_offs < BLOCK_D_HALF)[None, :]
+    else:
+        x_rotated_mask = (d_offs % 2 == 0)[None, :]
+
+    d_offs += nope_offs
+    y_offs = (
+        t_offs[:, None] * stride_y_t + d_offs[None, :] * stride_y_d + h_s * stride_y_h
+    )
+    y = tl.load(y_ptr + y_offs, mask=x_mask)
+
+    if IS_NEOX:
+        y_rotated = _get_neox_rotated_x(
+            y * sin, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF, True
+        )
+    else:
+        y_rotated = _get_gptj_rotated_x(
+            y * sin, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF, True
+        )
+
+    out_y_offs = (
+        t_offs[:, None] * stride_out_y_t
+        + d_offs[None, :] * stride_out_y_d
+        + h_s * stride_out_y_h
+    )
+    out_y = y * cos + y_rotated
+    out_y = out_y.to(y_ptr.dtype.element_ty)
+    tl.store(out_y_ptr + out_y_offs, out_y, mask=x_mask)
+
+    if HAVE_NOPE and not INPLACE:
+        if NOPE_FIRST:
+            y = tl.load(y_ptr + y_offs - BLOCK_D * stride_y_d, mask=x_mask)
+            tl.store(out_y_ptr + out_y_offs - BLOCK_D * stride_out_y_d, y, mask=x_mask)
+        else:
+            y = tl.load(y_ptr + y_offs + BLOCK_D * stride_y_d, mask=x_mask)
+            tl.store(out_y_ptr + out_y_offs + BLOCK_D * stride_out_y_d, y, mask=x_mask)
+
+    for h in tl.range(0, QH_per_G, 1, num_stages=num_stages):
+        x_offs = (
+            t_offs[:, None] * stride_x_t
+            + d_offs[None, :] * stride_x_d
+            + (h_start_idx + h) * stride_x_h
+        )
+
+        x = tl.load(x_ptr + x_offs, mask=x_mask)
+
+        if IS_NEOX:
+            x_rotated = _get_neox_rotated_x(
+                x * sin, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF, True
+            )
+        else:
+            x_rotated = _get_gptj_rotated_x(
+                x * sin, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF, True
+            )
+
+        out_x_offs = (
+            t_offs[:, None] * stride_out_x_t
+            + d_offs[None, :] * stride_out_x_d
+            + (h_start_idx + h) * stride_out_x_h
+        )
+        out_x = x * cos + x_rotated
+        out_x = out_x.to(x_ptr.dtype.element_ty)
+
+        tl.store(out_x_ptr + out_x_offs, out_x, mask=x_mask)
+
+        if HAVE_NOPE and not INPLACE:
+            if NOPE_FIRST:
+                x = tl.load(x_ptr + x_offs - BLOCK_D * stride_x_d, mask=x_mask)
+                tl.store(
+                    out_x_ptr + out_x_offs - BLOCK_D * stride_out_x_d, x, mask=x_mask
+                )
+            else:
+                x = tl.load(x_ptr + x_offs + BLOCK_D * stride_x_d, mask=x_mask)
+                tl.store(
+                    out_x_ptr + out_x_offs + BLOCK_D * stride_out_x_d, x, mask=x_mask
+                )
+
+
+@triton.jit
+def _rope_kernel_cached_thd_2c_gqa_onehead_bwd(
+    x_ptr,
+    y_ptr,
+    cos_ptr,
+    sin_ptr,
+    pos_ptr,
+    off_ptr,
+    out_x_ptr,
+    out_y_ptr,
+    stride_x_t,
+    stride_x_h,
+    stride_x_d,
+    stride_y_t,
+    stride_y_h,
+    stride_y_d,
+    stride_cos_t,
+    stride_cos_d,
+    stride_pos_t,
+    stride_out_x_t,
+    stride_out_x_h,
+    stride_out_x_d,
+    stride_out_y_t,
+    stride_out_y_h,
+    stride_out_y_d,
+    T,
+    HAVE_NOPE: tl.constexpr,
+    NOPE_FIRST: tl.constexpr,
+    INPLACE: tl.constexpr,
+    REUSE_FREQS_FRONT_PART: tl.constexpr,
+    IS_NEOX: tl.constexpr,
+    HAVE_POS: tl.constexpr,
+    HAVE_OFFS: tl.constexpr,
+    BLOCK_T: tl.constexpr,
+    G: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+    BLOCK_D_HALF: tl.constexpr,
+):
+    pid_t = tl.program_id(0)
+    hq = tl.program_id(1)
+
+    tl.assume(stride_x_t > 0)
+    tl.assume(stride_x_h > 0)
+    tl.assume(stride_x_d > 0)
+    tl.assume(stride_y_t > 0)
+    tl.assume(stride_y_h > 0)
+    tl.assume(stride_y_d > 0)
+    tl.assume(stride_cos_t > 0)
+    tl.assume(stride_cos_d > 0)
+    tl.assume(stride_pos_t > 0)
+    tl.assume(stride_out_x_t > 0)
+    tl.assume(stride_out_x_h > 0)
+    tl.assume(stride_out_x_d > 0)
+    tl.assume(stride_out_y_t > 0)
+    tl.assume(stride_out_y_h > 0)
+    tl.assume(stride_out_y_d > 0)
+
+    t_offs = pid_t * BLOCK_T + tl.arange(0, BLOCK_T)
+    d_offs = tl.arange(0, BLOCK_D)
+    t_mask = t_offs < T
+
+    if HAVE_POS:
+        pos_offs = t_offs * stride_pos_t
+        pos = tl.load(pos_ptr + pos_offs, mask=t_mask)
+        if HAVE_OFFS:
+            offset = tl.load(off_ptr + pos_offs, mask=t_mask)
+            t_cos_offs = pos + offset
+        else:
+            t_cos_offs = pos
+    else:
+        t_cos_offs = t_offs
+
+    if REUSE_FREQS_FRONT_PART:
+        if IS_NEOX:
+            d_cos_offs = d_offs
+            d_cos_offs = tl.where(
+                (d_cos_offs < BLOCK_D_HALF),
+                d_cos_offs,
+                d_cos_offs - BLOCK_D_HALF,
+            ).to(d_cos_offs.dtype)
+            d_cos_mask = d_cos_offs < BLOCK_D_HALF
+        else:
+            d_cos_offs = tl.arange(0, BLOCK_D) // 2
+            d_cos_mask = d_cos_offs < BLOCK_D_HALF
+    else:
+        d_cos_offs = d_offs
+        d_cos_mask = d_cos_offs < BLOCK_D
+
+    cos_mask = t_mask[:, None] & d_cos_mask[None, :]
+    cos_offs = t_cos_offs[:, None] * stride_cos_t + d_cos_offs[None, :] * stride_cos_d
+    cos = tl.load(cos_ptr + cos_offs, mask=cos_mask)
+    sin = tl.load(sin_ptr + cos_offs, mask=cos_mask)
+
+    nope_offs = 0
+    if HAVE_NOPE and NOPE_FIRST:
+        nope_offs = BLOCK_D
+
+    x_mask = t_mask[:, None] & (d_offs < BLOCK_D)[None, :]
+
+    if IS_NEOX:
+        x_rotated_mask = (d_offs < BLOCK_D_HALF)[None, :]
+    else:
+        x_rotated_mask = (d_offs % 2 == 0)[None, :]
+
+    d_offs += nope_offs
+    x_offs = (
+        t_offs[:, None] * stride_x_t + d_offs[None, :] * stride_x_d + hq * stride_x_h
+    )
+    x = tl.load(x_ptr + x_offs, mask=x_mask)
+
+    if IS_NEOX:
+        x_rotated = _get_neox_rotated_x(
+            x * sin, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF, True
+        )
+    else:
+        x_rotated = _get_gptj_rotated_x(
+            x * sin, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF, True
+        )
+
+    out_x_offs = (
+        t_offs[:, None] * stride_out_x_t
+        + d_offs[None, :] * stride_out_x_d
+        + hq * stride_out_x_h
+    )
+    out_x = x * cos + x_rotated
+    out_x = out_x.to(x_ptr.dtype.element_ty)
+    tl.store(out_x_ptr + out_x_offs, out_x, mask=x_mask)
+
+    if HAVE_NOPE and not INPLACE:
+        if NOPE_FIRST:
+            x = tl.load(x_ptr + x_offs - BLOCK_D * stride_x_d, mask=x_mask)
+            tl.store(out_x_ptr + out_x_offs - BLOCK_D * stride_out_x_d, x, mask=x_mask)
+        else:
+            x = tl.load(x_ptr + x_offs + BLOCK_D * stride_x_d, mask=x_mask)
+            tl.store(out_x_ptr + out_x_offs + BLOCK_D * stride_out_x_d, x, mask=x_mask)
+
+    if hq < G:
+        y_offs = (
+            t_offs[:, None] * stride_y_t
+            + d_offs[None, :] * stride_x_d
+            + hq * stride_y_h
+        )
+        y = tl.load(y_ptr + y_offs, mask=x_mask)
+
+        if IS_NEOX:
+            y_rotated = _get_neox_rotated_x(
+                y * sin, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF, True
+            )
+        else:
+            y_rotated = _get_gptj_rotated_x(
+                y * sin, x_rotated_mask, BLOCK_T, BLOCK_D, BLOCK_D_HALF, True
+            )
+
+        out_y_offs = (
+            t_offs[:, None] * stride_out_y_t
+            + d_offs[None, :] * stride_out_y_d
+            + hq * stride_out_y_h
+        )
+        out_y = y * cos + y_rotated
+        out_y = out_y.to(y_ptr.dtype.element_ty)
+        tl.store(out_y_ptr + out_y_offs, out_y, mask=x_mask)
+
+        if HAVE_NOPE and not INPLACE:
+            if NOPE_FIRST:
+                y = tl.load(y_ptr + y_offs - BLOCK_D * stride_y_d, mask=x_mask)
+                tl.store(
+                    out_y_ptr + out_y_offs - BLOCK_D * stride_out_y_d, y, mask=x_mask
+                )
+            else:
+                y = tl.load(y_ptr + y_offs + BLOCK_D * stride_y_d, mask=x_mask)
+                tl.store(
+                    out_y_ptr + out_y_offs + BLOCK_D * stride_out_y_d, y, mask=x_mask
+                )
+
+
+@triton.jit
+def _rope_fwd_2d_kernel_neox(
+    x_ptr,
+    cos_h_ptr,
+    sin_h_ptr,
+    cos_w_ptr,
+    sin_w_ptr,
+    out_ptr,
+    stride_x_b,
+    stride_x_wh,
+    stride_x_h,
+    stride_x_d,
+    stride_cos_h_b,
+    stride_cos_h_ht,
+    stride_cos_h_h,
+    stride_cos_h_d,
+    stride_cos_w_b,
+    stride_cos_w_w,
+    stride_cos_w_h,
+    stride_cos_w_d,
+    WH: tl.constexpr,
+    HEIGHT: tl.constexpr,
+    WEIGHT: tl.constexpr,
+    BLOCK_D: tl.constexpr,
+):
+    b = tl.program_id(0)
+    h = tl.program_id(1)
+
+    # load cos_h [HT, BLOCK_D]
+    offs_wh = tl.arange(0, WH)
+    offs_cos_h_h = offs_wh // WEIGHT
+    offs_d = tl.arange(0, BLOCK_D)
+    offs_cos_h = (
+        stride_cos_h_h * offs_cos_h_h[:, None] + stride_cos_h_d * offs_d[None, :]
+    )
+    mask_cos_h = offs_d < BLOCK_D // 2
+    cos_h = tl.load(cos_h_ptr + offs_cos_h, mask=mask_cos_h[None, :])
+
+    # load sin_h
+    sin_h = tl.load(sin_h_ptr + offs_cos_h, mask=mask_cos_h[None, :])
+
+    # load cos_w
+    offs_cos_w_w = offs_wh % WEIGHT
+    offs_cos_w_d = offs_d - BLOCK_D // 2
+    offs_cos_w = (
+        stride_cos_w_w * offs_cos_w_w[:, None] + stride_cos_w_d * offs_cos_w_d[None, :]
+    )
+    mask_cos_w = (offs_cos_w_d >= 0) & (offs_cos_w_d < BLOCK_D // 2)
+    cos_w = tl.load(cos_w_ptr + offs_cos_w, mask=mask_cos_w[None, :])
+
+    # load sin_w
+    sin_w = tl.load(sin_w_ptr + offs_cos_w, mask=mask_cos_w[None, :])
+
+    # load x
+    offs_wh = tl.arange(0, WH)
+    offs_x = (
+        stride_x_b * b
+        + stride_x_wh * offs_wh[:, None]
+        + stride_x_h * h
+        + stride_x_d * offs_d[None, :]
+    )
+    x = tl.load(x_ptr + offs_x)
+
+    # load x_rotated
+    offs_wh = tl.arange(0, WH)
+    offs_d_rotated = tl.where(offs_d < BLOCK_D // 4, offs_d + BLOCK_D // 4, offs_d)
+    offs_d_rotated = tl.where(
+        (offs_d >= BLOCK_D // 4) & (offs_d < BLOCK_D // 2),
+        offs_d_rotated - BLOCK_D // 4,
+        offs_d_rotated,
+    )
+    offs_d_rotated = tl.where(
+        (offs_d >= BLOCK_D // 2) & (offs_d < 3 * BLOCK_D // 4),
+        offs_d_rotated + BLOCK_D // 4,
+        offs_d_rotated,
+    )
+    offs_d_rotated = tl.where(
+        (offs_d >= 3 * BLOCK_D // 4) & (offs_d < BLOCK_D),
+        offs_d_rotated - BLOCK_D // 4,
+        offs_d_rotated,
+    )
+    offs_x_rotated = (
+        stride_x_b * b
+        + stride_x_wh * offs_wh[:, None]
+        + stride_x_h * h
+        + stride_x_d * offs_d_rotated[None, :]
+    )
+    x_rotated = tl.load(x_ptr + offs_x_rotated)
+    neg_x_rotated = tl.where((offs_d >= BLOCK_D // 4) & (offs_d < BLOCK_D // 2), 1, 0)
+    neg_x_rotated = tl.where(
+        (offs_d >= 3 * BLOCK_D // 4) & (offs_d < BLOCK_D), 1, neg_x_rotated
+    )
+    x_rotated = tl.where(neg_x_rotated, x_rotated, -x_rotated)
+
+    # compute x1
+    x1 = x * cos_h + x_rotated * sin_h
+
+    # compute x2
+    x2 = x * cos_w + x_rotated * sin_w
+
+    # compute output
+    out = x1 + x2
+
+    # store output
+    tl.store(out_ptr + offs_x, out)
+
+
+@triton.jit
+def _rope_fwd_3d(
+    x_ptr,
+    freqs_real_ptr,
+    freqs_imag_ptr,
+    grid_sizes_ptr,
+    out_ptr,
+    stride_x_b,
+    stride_x_l,
+    stride_x_n,
+    stride_x_c,
+    stride_freqs_s,
+    stride_freqs_c,
+    stride_grid_b,
+    stride_grid_d,
+    stride_out_b,
+    stride_out_l,
+    stride_out_n,
+    stride_out_c,
+    L: tl.constexpr,
+    N_HEADS: tl.constexpr,
+    C: tl.constexpr,
+    c_total: tl.constexpr,
+    sp_size: tl.constexpr,
+    sp_rank: tl.constexpr,
+    max_freq_seq_len: tl.constexpr,
+    s_per_rank: tl.constexpr,
+    pad_freq_val_r: tl.constexpr,
+    pad_freq_val_i: tl.constexpr,
+    BLOCK_L: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_C: tl.constexpr,
+    C1: tl.constexpr,
+    C2: tl.constexpr,
+):
+    pid_b = tl.program_id(0)
+    pid_n = tl.program_id(1)
+    pid_l = tl.program_id(2)
+
+    l_start = pid_l * BLOCK_L
+    l_off = l_start + tl.arange(0, BLOCK_L)
+    s_mask = l_off < L
+
+    c_off = tl.arange(0, BLOCK_C)
+    c_mask = c_off < c_total
+
+    # head mask
+    n_mask = pid_n < N_HEADS
+
+    # broadcast to  (BLOCK_L, 1, BLOCK_C)
+    l_b = tl.broadcast_to(l_off[:, None], (BLOCK_L, BLOCK_C))
+    c_b = tl.broadcast_to(c_off[None, :], (BLOCK_L, BLOCK_C))
+
+    # read grid_sizes
+    f_grid = tl.load(
+        grid_sizes_ptr + pid_b * stride_grid_b + 0 * stride_grid_d, mask=n_mask, other=0
+    )
+    h_grid = tl.load(
+        grid_sizes_ptr + pid_b * stride_grid_b + 1 * stride_grid_d, mask=n_mask, other=0
+    )
+    w_grid = tl.load(
+        grid_sizes_ptr + pid_b * stride_grid_b + 2 * stride_grid_d, mask=n_mask, other=0
+    )
+    h_w = h_grid * w_grid
+
+    global_tid = sp_rank * s_per_rank + l_b
+    valid_global_tid = global_tid < f_grid * h_w
+
+    # caculate f h w
+    f_idx = tl.where(valid_global_tid, global_tid // h_w, 0)
+    rem = tl.where(valid_global_tid, global_tid % h_w, 0)
+    h_idx = tl.where(valid_global_tid, rem // w_grid, 0)
+    w_idx = tl.where(valid_global_tid, rem % w_grid, 0)
+
+    freq_row = tl.where(c_b < C1, f_idx, tl.where(c_b < C1 + C2, h_idx, w_idx))
+    freq_row = tl.where(freq_row >= max_freq_seq_len, max_freq_seq_len - 1, freq_row)
+
+    mask_rope = s_mask[:, None] & c_mask[None, :] & n_mask & valid_global_tid[:, :]
+
+    # load freqs_real and freqs_imag
+    off_freq = freq_row * stride_freqs_s + c_b * stride_freqs_c
+    freq_r = tl.load(freqs_real_ptr + off_freq, mask=mask_rope, other=pad_freq_val_r)
+    freq_i = tl.load(freqs_imag_ptr + off_freq, mask=mask_rope, other=pad_freq_val_i)
+
+    off_x_base = pid_b * stride_x_b + pid_n * stride_x_n
+    off_x_r = off_x_base + l_b * stride_x_l + (2 * c_b) * stride_x_c
+    off_x_i = off_x_base + l_b * stride_x_l + (2 * c_b + 1) * stride_x_c
+
+    x_r = tl.load(x_ptr + off_x_r, mask=mask_rope, other=0.0)
+    x_i = tl.load(x_ptr + off_x_i, mask=mask_rope, other=0.0)
+
+    # complex number multiplication
+    out_r = x_r * freq_r - x_i * freq_i
+    out_i = x_r * freq_i + x_i * freq_r
+
+    # write result
+    off_out_base = pid_b * stride_out_b + pid_n * stride_out_n
+    off_out_r = off_out_base + l_b * stride_out_l + (2 * c_b) * stride_out_c
+    off_out_i = off_out_base + l_b * stride_out_l + (2 * c_b + 1) * stride_out_c
+
+    tl.store(out_ptr + off_out_r, out_r, mask=mask_rope)
+    tl.store(out_ptr + off_out_i, out_i, mask=mask_rope)

build/torch-rocm/_ops.py

@@ -0,0 +1,38 @@
+import torch
+
+def get_backend() -> str:
+    """Detect the backend by inspecting torch."""
+    import torch
+
+    if hasattr(torch, "neuron"):
+        # Needs to be sorted before specific Torch builds, since Neuron
+        # extension can be loaded into e.g. CUDA Torch builds.
+        return "neuron"
+    elif torch.version.cuda is not None:
+        return "cuda"
+    elif torch.version.hip is not None:
+        return "rocm"
+    elif torch.backends.mps.is_available():
+        return "metal"
+    elif hasattr(torch.version, "xpu") and torch.version.xpu is not None:
+        return "xpu"
+    else:
+        return "cpu"
+
+
+def _find_ops_name() -> str:
+    kernel_name = "aiter_rope"
+    unique_id = "6c454e0"
+    backend = get_backend()
+    return f"_{kernel_name}_{backend}_{unique_id}"
+
+
+_OPS_NAME = _find_ops_name()
+
+ops = getattr(torch.ops, _OPS_NAME)
+
+def add_op_namespace_prefix(op_name: str) -> str:
+    """
+    Prefix op by namespace.
+    """
+    return f"{_OPS_NAME}::{op_name}"

build/torch-rocm/aiter_rope/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+
+
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+
+
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch-rocm/metadata.json

@@ -0,0 +1,10 @@
+{
+  "name": "aiter-rope",
+  "id": "_aiter_rope_rocm_6c454e0",
+  "version": 1,
+  "license": "MIT",
+  "python-depends": [],
+  "backend": {
+    "type": "rocm"
+  }
+}

build/torch-rocm/rope.py

@@ -0,0 +1,1728 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
+
+import torch
+import triton
+import triton.language as tl
+from torch import autograd
+from enum import IntEnum
+from typing import Tuple, Union
+from ._kernels.rope import (
+    _rope_kernel_sbhd_fwd,
+    _rope_kernel_sbhd_bwd,
+    _rope_kernel_thd_fwd,
+    _rope_kernel_thd_bwd,
+    _rope_kernel_sbhd_cached_fwd,
+    _rope_kernel_sbhd_cached_bwd,
+    _rope_kernel_thd_cached_2c_fwd,
+    _rope_kernel_thd_cached_2c_bwd,
+    _rope_kernel_cached_thd_2c_gqa_fwd,
+    _rope_kernel_cached_thd_2c_gqa_onehead_fwd,
+    _rope_kernel_cached_thd_2c_gqa_bwd,
+    _rope_kernel_cached_thd_2c_gqa_onehead_bwd,
+    _rope_fwd_2d_kernel_neox,
+    _rope_fwd_3d,
+    _get_neox_rotated_x_1D,
+    _get_gptj_rotated_x_1D,
+    _get_neox_rotated_x,
+    _get_gptj_rotated_x,
+)
+from .utils.logger import AiterTritonLogger
+
+__all__ = [
+    "_get_neox_rotated_x_1D",
+    "_get_gptj_rotated_x_1D",
+    "_get_neox_rotated_x",
+    "_get_gptj_rotated_x",
+]
+
+_LOGGER = AiterTritonLogger()
+
+
+class RotateStyle(IntEnum):
+    NEOX = (0,)
+    GPTJ = 1
+
+
+# TODO: For now BLOCK_D is assumed to be power of 2. Expand to handle other value of D.
+def _rope_fwd(
+    x: torch.Tensor,
+    out: torch.Tensor,
+    freqs: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    inplace: bool,
+    transpose_output: bool = False,
+) -> torch.Tensor:
+    s, b, h, d = x.shape
+
+    if freqs.shape[-1] == d // 2:
+        if reuse_freqs_front_part:
+            have_nope = False
+        else:
+            have_nope = True
+    elif freqs.shape[-1] == d // 4:
+        have_nope = True
+    else:
+        have_nope = False
+
+    if have_nope:
+        BLOCK_D = d // 2
+        BLOCK_D_HALF = d // 4
+    else:
+        BLOCK_D = d
+        BLOCK_D_HALF = d // 2
+
+    # TODO: performance optimization
+    BLOCK_S = 32
+    num_warps = 4
+    waves_per_eu = 0
+    grid = (b, h, triton.cdiv(s, BLOCK_S))
+
+    _rope_kernel_sbhd_fwd[grid](
+        x,
+        freqs,
+        out,
+        *x.stride(),
+        *freqs.stride(),
+        *out.stride(),
+        s,
+        HAVE_NOPE=have_nope,
+        NOPE_FIRST=nope_first,
+        INPLACE=inplace,
+        REUSE_FREQS_FRONT_PART=reuse_freqs_front_part,
+        IS_NEOX=(rotate_style == RotateStyle.NEOX),
+        BLOCK_S=BLOCK_S,
+        BLOCK_D=BLOCK_D,
+        BLOCK_D_HALF=BLOCK_D_HALF,
+        num_warps=num_warps,
+        waves_per_eu=waves_per_eu,
+    )
+
+    return out
+
+
+def rope_fwd(
+    x: torch.Tensor,
+    freqs: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    transpose_output: bool = False,
+) -> torch.Tensor:
+    s, b, h, d = x.shape
+    out = torch.empty((s, b, h, d), dtype=x.dtype, device=x.device, requires_grad=False)
+
+    _rope_fwd(
+        x,
+        out,
+        freqs,
+        rotate_style,
+        reuse_freqs_front_part,
+        nope_first,
+        False,
+        transpose_output,
+    )
+
+    return out
+
+
+def rope_fwd_inplace(
+    x: torch.Tensor,
+    freqs: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    transpose_output: bool = False,
+) -> torch.Tensor:
+    out = x
+
+    _rope_fwd(
+        x,
+        out,
+        freqs,
+        rotate_style,
+        reuse_freqs_front_part,
+        nope_first,
+        True,
+        transpose_output,
+    )
+
+    return out
+
+
+def _rope_bwd(
+    x: torch.Tensor,
+    out: torch.Tensor,
+    freqs: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    inplace: bool,
+    transpose_output: bool = False,
+) -> torch.Tensor:
+    s, b, h, d = x.shape
+
+    if freqs.shape[-1] == d // 2:
+        if reuse_freqs_front_part:
+            have_nope = False
+        else:
+            have_nope = True
+    elif freqs.shape[-1] == d // 4:
+        have_nope = True
+    else:
+        have_nope = False
+
+    if have_nope:
+        BLOCK_D = d // 2
+        BLOCK_D_HALF = d // 4
+    else:
+        BLOCK_D = d
+        BLOCK_D_HALF = d // 2
+
+    # TODO: performance optimization
+    BLOCK_S = 32
+    num_warps = 4
+    waves_per_eu = 0
+    grid = (b, h, triton.cdiv(s, BLOCK_S))
+
+    _rope_kernel_sbhd_bwd[grid](
+        x,
+        freqs,
+        out,
+        *x.stride(),
+        *freqs.stride(),
+        *out.stride(),
+        s,
+        HAVE_NOPE=have_nope,
+        NOPE_FIRST=nope_first,
+        INPLACE=inplace,
+        REUSE_FREQS_FRONT_PART=reuse_freqs_front_part,
+        IS_NEOX=(rotate_style == RotateStyle.NEOX),
+        BLOCK_S=BLOCK_S,
+        BLOCK_D=BLOCK_D,
+        BLOCK_D_HALF=BLOCK_D_HALF,
+        num_warps=num_warps,
+        waves_per_eu=waves_per_eu,
+    )
+
+    return out
+
+
+def rope_bwd(
+    x: torch.Tensor,
+    freqs: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    transpose_output: bool = False,
+) -> torch.Tensor:
+    s, b, h, d = x.shape
+    out = torch.empty((s, b, h, d), dtype=x.dtype, device=x.device, requires_grad=False)
+
+    _rope_bwd(
+        x,
+        out,
+        freqs,
+        rotate_style,
+        reuse_freqs_front_part,
+        nope_first,
+        False,
+        transpose_output,
+    )
+
+    return out
+
+
+def _rope_thd_fwd(
+    x: torch.Tensor,
+    out: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    freqs: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    inplace: bool,
+    transpose_output: bool = False,
+) -> torch.Tensor:
+    b = torch.numel(cu_seqlens) - 1
+    t, h, d = x.shape
+
+    if freqs.shape[-1] == d // 2:
+        if reuse_freqs_front_part:
+            have_nope = False
+        else:
+            have_nope = True
+    elif freqs.shape[-1] == d // 4:
+        have_nope = True
+    else:
+        have_nope = False
+
+    if have_nope:
+        BLOCK_D = d // 2
+        BLOCK_D_HALF = d // 4
+    else:
+        BLOCK_D = d
+        BLOCK_D_HALF = d // 2
+
+    # TODO: performance optimization
+    BLOCK_T = 32
+    num_warps = 4
+    waves_per_eu = 0
+    grid = (b, h, triton.cdiv(t, BLOCK_T))
+
+    _rope_kernel_thd_fwd[grid](
+        x,
+        cu_seqlens,
+        freqs,
+        out,
+        *x.stride(),
+        *freqs.stride(),
+        *out.stride(),
+        HAVE_NOPE=have_nope,
+        NOPE_FIRST=nope_first,
+        INPLACE=inplace,
+        REUSE_FREQS_FRONT_PART=reuse_freqs_front_part,
+        IS_NEOX=(rotate_style == RotateStyle.NEOX),
+        BLOCK_T=BLOCK_T,
+        BLOCK_D=BLOCK_D,
+        BLOCK_D_HALF=BLOCK_D_HALF,
+        num_warps=num_warps,
+        waves_per_eu=waves_per_eu,
+    )
+
+    return out
+
+
+def rope_thd_fwd(
+    x: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    freqs: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    transpose_output: bool = False,
+) -> torch.Tensor:
+    t, h, d = x.shape
+    out = torch.empty((t, h, d), dtype=x.dtype, device=x.device, requires_grad=False)
+
+    _rope_thd_fwd(
+        x,
+        out,
+        cu_seqlens,
+        freqs,
+        rotate_style,
+        reuse_freqs_front_part,
+        nope_first,
+        False,
+        transpose_output,
+    )
+
+    return out
+
+
+def rope_thd_fwd_inplace(
+    x: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    freqs: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    transpose_output: bool = False,
+) -> torch.Tensor:
+    out = x
+
+    _rope_thd_fwd(
+        x,
+        out,
+        cu_seqlens,
+        freqs,
+        rotate_style,
+        reuse_freqs_front_part,
+        nope_first,
+        True,
+        transpose_output,
+    )
+
+    return out
+
+
+def _rope_thd_bwd(
+    x: torch.Tensor,
+    out: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    freqs: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    inplace: bool,
+    transpose_output: bool = False,
+) -> torch.Tensor:
+    b = torch.numel(cu_seqlens) - 1
+    t, h, d = x.shape
+
+    if freqs.shape[-1] == d // 2:
+        if reuse_freqs_front_part:
+            have_nope = False
+        else:
+            have_nope = True
+    elif freqs.shape[-1] == d // 4:
+        have_nope = True
+    else:
+        have_nope = False
+
+    if have_nope:
+        BLOCK_D = d // 2
+        BLOCK_D_HALF = d // 4
+    else:
+        BLOCK_D = d
+        BLOCK_D_HALF = d // 2
+
+    # TODO: performance optimization
+    BLOCK_T = 32
+    num_warps = 4
+    waves_per_eu = 0
+    grid = (b, h, triton.cdiv(t, BLOCK_T))
+
+    _rope_kernel_thd_bwd[grid](
+        x,
+        cu_seqlens,
+        freqs,
+        out,
+        *x.stride(),
+        *freqs.stride(),
+        *out.stride(),
+        HAVE_NOPE=have_nope,
+        NOPE_FIRST=nope_first,
+        INPLACE=inplace,
+        REUSE_FREQS_FRONT_PART=reuse_freqs_front_part,
+        IS_NEOX=(rotate_style == RotateStyle.NEOX),
+        BLOCK_T=BLOCK_T,
+        BLOCK_D=BLOCK_D,
+        BLOCK_D_HALF=BLOCK_D_HALF,
+        num_warps=num_warps,
+        waves_per_eu=waves_per_eu,
+    )
+
+    return out
+
+
+def rope_thd_bwd(
+    x: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    freqs: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    transpose_output: bool = False,
+) -> torch.Tensor:
+    t, h, d = x.shape
+    out = torch.empty((t, h, d), dtype=x.dtype, device=x.device, requires_grad=False)
+
+    _rope_thd_bwd(
+        x,
+        out,
+        cu_seqlens,
+        freqs,
+        rotate_style,
+        reuse_freqs_front_part,
+        nope_first,
+        False,
+        transpose_output,
+    )
+
+    return out
+
+
+# TODO: For now BLOCK_D is assumed to be power of 2. Expand to handle other value of D.
+def _rope_cached_fwd(
+    x: torch.Tensor,
+    out: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    positions: torch.Tensor,
+    offsets: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    inplace: bool,
+    transpose_output: bool = False,
+) -> torch.Tensor:
+    s, b, h, d = x.shape
+
+    if cos.shape[-1] == d // 2:
+        if reuse_freqs_front_part:
+            have_nope = False
+        else:
+            have_nope = True
+    elif cos.shape[-1] == d // 4:
+        have_nope = True
+    else:
+        have_nope = False
+
+    if have_nope:
+        BLOCK_D = d // 2
+        BLOCK_D_HALF = d // 4
+    else:
+        BLOCK_D = d
+        BLOCK_D_HALF = d // 2
+
+    # TODO: performance optimization
+    BLOCK_S = 32
+    num_warps = 4
+    waves_per_eu = 0
+    grid = (b, h, triton.cdiv(s, BLOCK_S))
+
+    pos_stride = positions.stride() if positions is not None else (1, 1)
+    _rope_kernel_sbhd_cached_fwd[grid](
+        x,
+        cos,
+        sin,
+        positions,
+        offsets,
+        out,
+        *x.stride(),
+        *cos.stride(),
+        *pos_stride,
+        *out.stride(),
+        s,
+        HAVE_NOPE=have_nope,
+        NOPE_FIRST=nope_first,
+        INPLACE=inplace,
+        REUSE_FREQS_FRONT_PART=reuse_freqs_front_part,
+        IS_NEOX=(rotate_style == RotateStyle.NEOX),
+        HAVE_POS=(positions is not None),
+        HAVE_OFFS=(offsets is not None),
+        BLOCK_S=BLOCK_S,
+        BLOCK_D=BLOCK_D,
+        BLOCK_D_HALF=BLOCK_D_HALF,
+        num_warps=num_warps,
+        waves_per_eu=waves_per_eu,
+    )
+
+    return out
+
+
+def rope_cached_fwd(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    transpose_output: bool = False,
+):
+    s, b, h, d = x.shape
+    out = torch.empty((s, b, h, d), dtype=x.dtype, device=x.device, requires_grad=False)
+
+    _rope_cached_fwd(
+        x,
+        out,
+        cos,
+        sin,
+        None,
+        None,
+        rotate_style,
+        reuse_freqs_front_part,
+        nope_first,
+        False,
+        transpose_output,
+    )
+
+    return out
+
+
+def rope_cached_fwd_inplace(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    transpose_output: bool = False,
+):
+    out = x
+
+    _rope_cached_fwd(
+        x,
+        out,
+        cos,
+        sin,
+        None,
+        None,
+        rotate_style,
+        reuse_freqs_front_part,
+        nope_first,
+        True,
+        transpose_output,
+    )
+
+    return out
+
+
+def rope_cached_positions_fwd(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    positions: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    transpose_output: bool = False,
+) -> torch.Tensor:
+    s, b, h, d = x.shape
+    out = torch.empty((s, b, h, d), dtype=x.dtype, device=x.device, requires_grad=False)
+
+    _rope_cached_fwd(
+        x,
+        out,
+        cos,
+        sin,
+        positions,
+        None,
+        rotate_style,
+        reuse_freqs_front_part,
+        nope_first,
+        False,
+        transpose_output,
+    )
+
+    return out
+
+
+def rope_cached_positions_fwd_inplace(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    positions: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    transpose_output: bool = False,
+) -> torch.Tensor:
+    out = x
+
+    _rope_cached_fwd(
+        x,
+        out,
+        cos,
+        sin,
+        positions,
+        None,
+        rotate_style,
+        reuse_freqs_front_part,
+        nope_first,
+        True,
+        transpose_output,
+    )
+
+    return out
+
+
+def rope_cached_positions_offsets_fwd(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    positions: torch.Tensor,
+    offsets: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    transpose_output: bool = False,
+) -> torch.Tensor:
+    s, b, h, d = x.shape
+    out = torch.empty((s, b, h, d), dtype=x.dtype, device=x.device, requires_grad=False)
+
+    _rope_cached_fwd(
+        x,
+        out,
+        cos,
+        sin,
+        positions,
+        offsets,
+        rotate_style,
+        reuse_freqs_front_part,
+        nope_first,
+        False,
+        transpose_output,
+    )
+
+    return out
+
+
+def rope_cached_positions_offsets_fwd_inplace(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    positions: torch.Tensor,
+    offsets: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    transpose_output: bool = False,
+) -> torch.Tensor:
+    out = x
+
+    _rope_cached_fwd(
+        x,
+        out,
+        cos,
+        sin,
+        positions,
+        offsets,
+        rotate_style,
+        reuse_freqs_front_part,
+        nope_first,
+        True,
+        transpose_output,
+    )
+
+    return out
+
+
+def _rope_cached_bwd(
+    x: torch.Tensor,
+    out: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    positions: torch.Tensor,
+    offsets: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    inplace: bool,
+    transpose_output: bool = False,
+) -> torch.Tensor:
+    s, b, h, d = x.shape
+
+    if cos.shape[-1] == d // 2:
+        if reuse_freqs_front_part:
+            have_nope = False
+        else:
+            have_nope = True
+    elif cos.shape[-1] == d // 4:
+        have_nope = True
+    else:
+        have_nope = False
+
+    if have_nope:
+        BLOCK_D = d // 2
+        BLOCK_D_HALF = d // 4
+    else:
+        BLOCK_D = d
+        BLOCK_D_HALF = d // 2
+
+    # TODO: performance optimization
+    BLOCK_S = 32
+    num_warps = 4
+    waves_per_eu = 0
+    grid = (b, h, triton.cdiv(s, BLOCK_S))
+
+    pos_stride = positions.stride() if positions is not None else (1, 1)
+    _rope_kernel_sbhd_cached_bwd[grid](
+        x,
+        cos,
+        sin,
+        positions,
+        offsets,
+        out,
+        *x.stride(),
+        *cos.stride(),
+        *pos_stride,
+        *out.stride(),
+        s,
+        HAVE_NOPE=have_nope,
+        NOPE_FIRST=nope_first,
+        INPLACE=inplace,
+        REUSE_FREQS_FRONT_PART=reuse_freqs_front_part,
+        IS_NEOX=(rotate_style == RotateStyle.NEOX),
+        HAVE_POS=(positions is not None),
+        HAVE_OFFS=(offsets is not None),
+        BLOCK_S=BLOCK_S,
+        BLOCK_D=BLOCK_D,
+        BLOCK_D_HALF=BLOCK_D_HALF,
+        num_warps=num_warps,
+        waves_per_eu=waves_per_eu,
+    )
+
+    return out
+
+
+def rope_cached_bwd(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    transpose_output: bool = False,
+):
+    s, b, h, d = x.shape
+    out = torch.empty((s, b, h, d), dtype=x.dtype, device=x.device, requires_grad=False)
+
+    _rope_cached_bwd(
+        x,
+        out,
+        cos,
+        sin,
+        None,
+        None,
+        rotate_style,
+        reuse_freqs_front_part,
+        nope_first,
+        False,
+        transpose_output,
+    )
+
+    return out
+
+
+def rope_cached_positions_bwd(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    positions: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    transpose_output: bool = False,
+) -> torch.Tensor:
+    s, b, h, d = x.shape
+    out = torch.empty((s, b, h, d), dtype=x.dtype, device=x.device, requires_grad=False)
+
+    _rope_cached_bwd(
+        x,
+        out,
+        cos,
+        sin,
+        positions,
+        None,
+        rotate_style,
+        reuse_freqs_front_part,
+        nope_first,
+        False,
+        transpose_output,
+    )
+
+    return out
+
+
+def rope_cached_positions_offsets_bwd(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    positions: torch.Tensor,
+    offsets: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    transpose_output: bool = False,
+) -> torch.Tensor:
+    s, b, h, d = x.shape
+    out = torch.empty((s, b, h, d), dtype=x.dtype, device=x.device, requires_grad=False)
+
+    _rope_cached_bwd(
+        x,
+        out,
+        cos,
+        sin,
+        positions,
+        offsets,
+        rotate_style,
+        reuse_freqs_front_part,
+        nope_first,
+        False,
+        transpose_output,
+    )
+
+    return out
+
+
+def _rope_cached_thd_2c_fwd(
+    x: torch.Tensor,
+    y: torch.Tensor,
+    out_x: torch.Tensor,
+    out_y: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    positions: torch.Tensor,
+    offsets: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    inplace: bool,
+    transpose_output: bool = False,
+):
+    t, h, d = x.shape
+    ty, kh, dy = y.shape
+
+    assert (
+        t == ty
+    ), f"The number of tokens should be the same for the two inputs, but got {t} and {ty}"
+    assert (
+        d == dy
+    ), f"The head dimension should be the same for the two inputs, but got {d} and {dy}"
+    assert h % kh == 0, f"QH should be multiple of KH, but got QH={h} and KH={kh}"
+
+    if cos.shape[-1] == d // 2:
+        if reuse_freqs_front_part:
+            have_nope = False
+        else:
+            have_nope = True
+    elif cos.shape[-1] == d // 4:
+        have_nope = True
+    else:
+        have_nope = False
+
+    if have_nope:
+        BLOCK_D = d // 2
+        BLOCK_D_HALF = d // 4
+    else:
+        BLOCK_D = d
+        BLOCK_D_HALF = d // 2
+
+    if h == kh:
+        BLOCK_T = 32
+        SPLIT_T = (triton.next_power_of_2(t) + BLOCK_T - 1) // BLOCK_T
+
+        if t >= 8192:
+            MIN_NUM_WG = 4096
+        elif t >= 1024:
+            MIN_NUM_WG = 1024
+        else:
+            MIN_NUM_WG = 512
+
+        if SPLIT_T < MIN_NUM_WG:
+            SPLIT_H_SIZE = h
+            SPLIT_H = (triton.next_power_of_2(h) + SPLIT_H_SIZE - 1) // SPLIT_H_SIZE
+            while SPLIT_H * SPLIT_T < MIN_NUM_WG and SPLIT_H_SIZE > 1:
+                SPLIT_H_SIZE = SPLIT_H_SIZE // 2
+                SPLIT_H = (triton.next_power_of_2(h) + SPLIT_H_SIZE - 1) // SPLIT_H_SIZE
+        else:
+            SPLIT_H_SIZE = h
+
+        SPLIT_H = (triton.next_power_of_2(h) + SPLIT_H_SIZE - 1) // SPLIT_H_SIZE
+        grid = (SPLIT_H, SPLIT_T, 1)
+        num_warps = 4
+        waves_per_eu = 0
+        num_stages = 2 if SPLIT_H_SIZE > 1 else 1
+
+        _rope_kernel_thd_cached_2c_fwd[grid](
+            x,
+            y,
+            cos,
+            sin,
+            positions,
+            offsets,
+            out_x,
+            out_y,
+            *x.stride(),
+            *y.stride(),
+            *cos.stride(),
+            *positions.stride(),
+            *out_x.stride(),
+            *out_y.stride(),
+            t,
+            HAVE_NOPE=have_nope,
+            NOPE_FIRST=nope_first,
+            INPLACE=inplace,
+            REUSE_FREQS_FRONT_PART=reuse_freqs_front_part,
+            IS_NEOX=(rotate_style == RotateStyle.NEOX),
+            HAVE_POS=(positions is not None),
+            HAVE_OFFS=(offsets is not None),
+            BLOCK_T=BLOCK_T,
+            SPLIT_H_SIZE=SPLIT_H_SIZE,
+            BLOCK_D=BLOCK_D,
+            BLOCK_D_HALF=BLOCK_D_HALF,
+            num_warps=num_warps,
+            waves_per_eu=waves_per_eu,
+            num_stages=num_stages,
+        )
+    else:
+        # TODO check boundary
+        if rotate_style == RotateStyle.GPTJ and t >= 1024:
+            BLOCK_T = 32
+            SPLIT_T = (triton.next_power_of_2(t) + BLOCK_T - 1) // BLOCK_T
+            QH_per_G = h // kh
+            grid = (kh, SPLIT_T, 1)
+            num_warps = 4
+            waves_per_eu = 0
+            num_stages = 2 if QH_per_G > 1 else 1
+
+            _rope_kernel_cached_thd_2c_gqa_fwd[grid](
+                x,
+                y,
+                cos,
+                sin,
+                positions,
+                offsets,
+                out_x,
+                out_y,
+                *x.stride(),
+                *y.stride(),
+                *cos.stride(),
+                *positions.stride(),
+                *out_x.stride(),
+                *out_y.stride(),
+                t,
+                HAVE_NOPE=have_nope,
+                NOPE_FIRST=nope_first,
+                INPLACE=inplace,
+                REUSE_FREQS_FRONT_PART=reuse_freqs_front_part,
+                IS_NEOX=(rotate_style == RotateStyle.NEOX),
+                HAVE_POS=(positions is not None),
+                HAVE_OFFS=(offsets is not None),
+                BLOCK_T=BLOCK_T,
+                QH_per_G=QH_per_G,
+                BLOCK_D=BLOCK_D,
+                BLOCK_D_HALF=BLOCK_D_HALF,
+                num_warps=num_warps,
+                waves_per_eu=waves_per_eu,
+                num_stages=num_stages,
+            )
+        else:
+            BLOCK_T = min(max(triton.next_power_of_2(t), 16), 32)
+            SPLIT_T = (triton.next_power_of_2(t) + BLOCK_T - 1) // BLOCK_T
+            grid = (SPLIT_T, h, 1)
+            num_warps = 4
+            waves_per_eu = 0
+            _rope_kernel_cached_thd_2c_gqa_onehead_fwd[grid](
+                x,
+                y,
+                cos,
+                sin,
+                positions,
+                offsets,
+                out_x,
+                out_y,
+                *x.stride(),
+                *y.stride(),
+                *cos.stride(),
+                *positions.stride(),
+                *out_x.stride(),
+                *out_y.stride(),
+                t,
+                HAVE_NOPE=have_nope,
+                NOPE_FIRST=nope_first,
+                INPLACE=inplace,
+                REUSE_FREQS_FRONT_PART=reuse_freqs_front_part,
+                IS_NEOX=(rotate_style == RotateStyle.NEOX),
+                HAVE_POS=(positions is not None),
+                HAVE_OFFS=(offsets is not None),
+                BLOCK_T=BLOCK_T,
+                G=kh,
+                BLOCK_D=BLOCK_D,
+                BLOCK_D_HALF=BLOCK_D_HALF,
+                num_warps=num_warps,
+                waves_per_eu=waves_per_eu,
+            )
+
+    return out_x, out_y
+
+
+def rope_cached_thd_positions_2c_fwd(
+    x: torch.Tensor,
+    y: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    positions: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    transpose_output: bool = False,
+):
+    out_x = torch.empty(*x.shape, dtype=x.dtype, device=x.device, requires_grad=False)
+    out_y = torch.empty(*y.shape, dtype=x.dtype, device=x.device, requires_grad=False)
+
+    _rope_cached_thd_2c_fwd(
+        x,
+        y,
+        out_x,
+        out_y,
+        cos,
+        sin,
+        positions,
+        None,
+        rotate_style,
+        reuse_freqs_front_part,
+        nope_first,
+        False,
+        transpose_output,
+    )
+
+    return out_x, out_y
+
+
+def rope_cached_thd_positions_2c_fwd_inplace(
+    x: torch.Tensor,
+    y: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    positions: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    transpose_output: bool = False,
+):
+    out_x = x
+    out_y = y
+
+    _rope_cached_thd_2c_fwd(
+        x,
+        y,
+        out_x,
+        out_y,
+        cos,
+        sin,
+        positions,
+        None,
+        rotate_style,
+        reuse_freqs_front_part,
+        nope_first,
+        True,
+        transpose_output,
+    )
+
+    return out_x, out_y
+
+
+def rope_cached_thd_positions_offsets_2c_fwd(
+    x: torch.Tensor,
+    y: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    positions: torch.Tensor,
+    offsets: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    transpose_output: bool = False,
+):
+    out_x = torch.empty(*x.shape, dtype=x.dtype, device=x.device, requires_grad=False)
+    out_y = torch.empty(*y.shape, dtype=x.dtype, device=x.device, requires_grad=False)
+
+    _rope_cached_thd_2c_fwd(
+        x,
+        y,
+        out_x,
+        out_y,
+        cos,
+        sin,
+        positions,
+        offsets,
+        rotate_style,
+        reuse_freqs_front_part,
+        nope_first,
+        False,
+        transpose_output,
+    )
+
+    return out_x, out_y
+
+
+def rope_cached_thd_positions_offsets_2c_fwd_inplace(
+    x: torch.Tensor,
+    y: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    positions: torch.Tensor,
+    offsets: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    transpose_output: bool = False,
+):
+    out_x = x
+    out_y = y
+
+    _rope_cached_thd_2c_fwd(
+        x,
+        y,
+        out_x,
+        out_y,
+        cos,
+        sin,
+        positions,
+        offsets,
+        rotate_style,
+        reuse_freqs_front_part,
+        nope_first,
+        True,
+        transpose_output,
+    )
+
+    return out_x, out_y
+
+
+def _rope_cached_thd_positions_offsets_2c_bwd(
+    x: torch.Tensor,
+    y: torch.Tensor,
+    out_x: torch.Tensor,
+    out_y: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    positions: torch.Tensor,
+    offsets: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    inplace: bool,
+    transpose_output: bool = False,
+):
+    t, h, d = x.shape
+    ty, kh, dy = y.shape
+
+    assert (
+        t == ty
+    ), f"The number of tokens should be the same for the two inputs, but got {t} and {ty}"
+    assert (
+        d == dy
+    ), f"The head dimension should be the same for the two inputs, but got {d} and {dy}"
+    assert h % kh == 0, f"QH should be multiple of KH, but got QH={h} and KH={kh}"
+
+    if cos.shape[-1] == d // 2:
+        if reuse_freqs_front_part:
+            have_nope = False
+        else:
+            have_nope = True
+    elif cos.shape[-1] == d // 4:
+        have_nope = True
+    else:
+        have_nope = False
+
+    if have_nope:
+        BLOCK_D = d // 2
+        BLOCK_D_HALF = d // 4
+    else:
+        BLOCK_D = d
+        BLOCK_D_HALF = d // 2
+
+    if h == kh:
+        BLOCK_T = 32
+        SPLIT_T = (triton.next_power_of_2(t) + BLOCK_T - 1) // BLOCK_T
+
+        if t >= 8192:
+            MIN_NUM_WG = 4096
+        elif t >= 1024:
+            MIN_NUM_WG = 1024
+        else:
+            MIN_NUM_WG = 512
+
+        if SPLIT_T < MIN_NUM_WG:
+            SPLIT_H_SIZE = h
+            SPLIT_H = (triton.next_power_of_2(h) + SPLIT_H_SIZE - 1) // SPLIT_H_SIZE
+            while SPLIT_H * SPLIT_T < MIN_NUM_WG and SPLIT_H_SIZE > 1:
+                SPLIT_H_SIZE = SPLIT_H_SIZE // 2
+                SPLIT_H = (triton.next_power_of_2(h) + SPLIT_H_SIZE - 1) // SPLIT_H_SIZE
+        else:
+            SPLIT_H_SIZE = h
+
+        SPLIT_H = (triton.next_power_of_2(h) + SPLIT_H_SIZE - 1) // SPLIT_H_SIZE
+        grid = (SPLIT_H, SPLIT_T, 1)
+        num_warps = 4
+        waves_per_eu = 0
+        num_stages = 2 if SPLIT_H_SIZE > 1 else 1
+
+        _rope_kernel_thd_cached_2c_bwd[grid](
+            x,
+            y,
+            cos,
+            sin,
+            positions,
+            offsets,
+            out_x,
+            out_y,
+            *x.stride(),
+            *y.stride(),
+            *cos.stride(),
+            *positions.stride(),
+            *out_x.stride(),
+            *out_y.stride(),
+            t,
+            HAVE_NOPE=have_nope,
+            NOPE_FIRST=nope_first,
+            INPLACE=inplace,
+            REUSE_FREQS_FRONT_PART=reuse_freqs_front_part,
+            IS_NEOX=(rotate_style == RotateStyle.NEOX),
+            HAVE_POS=(positions is not None),
+            HAVE_OFFS=(offsets is not None),
+            BLOCK_T=BLOCK_T,
+            SPLIT_H_SIZE=SPLIT_H_SIZE,
+            BLOCK_D=BLOCK_D,
+            BLOCK_D_HALF=BLOCK_D_HALF,
+            num_warps=num_warps,
+            waves_per_eu=waves_per_eu,
+            num_stages=num_stages,
+        )
+    else:
+        # TODO check boundary
+        if rotate_style == RotateStyle.GPTJ and t >= 1024:
+            BLOCK_T = 32
+            SPLIT_T = (triton.next_power_of_2(t) + BLOCK_T - 1) // BLOCK_T
+            QH_per_G = h // kh
+            grid = (kh, SPLIT_T, 1)
+            num_warps = 4
+            waves_per_eu = 0
+            num_stages = 2 if QH_per_G > 1 else 1
+
+            _rope_kernel_cached_thd_2c_gqa_bwd[grid](
+                x,
+                y,
+                cos,
+                sin,
+                positions,
+                offsets,
+                out_x,
+                out_y,
+                *x.stride(),
+                *y.stride(),
+                *cos.stride(),
+                *positions.stride(),
+                *out_x.stride(),
+                *out_y.stride(),
+                t,
+                HAVE_NOPE=have_nope,
+                NOPE_FIRST=nope_first,
+                INPLACE=inplace,
+                REUSE_FREQS_FRONT_PART=reuse_freqs_front_part,
+                IS_NEOX=(rotate_style == RotateStyle.NEOX),
+                HAVE_POS=(positions is not None),
+                HAVE_OFFS=(offsets is not None),
+                BLOCK_T=BLOCK_T,
+                QH_per_G=QH_per_G,
+                BLOCK_D=BLOCK_D,
+                BLOCK_D_HALF=BLOCK_D_HALF,
+                num_warps=num_warps,
+                waves_per_eu=waves_per_eu,
+                num_stages=num_stages,
+            )
+        else:
+            BLOCK_T = min(max(triton.next_power_of_2(t), 16), 32)
+            SPLIT_T = (triton.next_power_of_2(t) + BLOCK_T - 1) // BLOCK_T
+            grid = (SPLIT_T, h, 1)
+            num_warps = 4
+            waves_per_eu = 0
+            _rope_kernel_cached_thd_2c_gqa_onehead_bwd[grid](
+                x,
+                y,
+                cos,
+                sin,
+                positions,
+                offsets,
+                out_x,
+                out_y,
+                *x.stride(),
+                *y.stride(),
+                *cos.stride(),
+                *positions.stride(),
+                *out_x.stride(),
+                *out_y.stride(),
+                t,
+                HAVE_NOPE=have_nope,
+                NOPE_FIRST=nope_first,
+                INPLACE=inplace,
+                REUSE_FREQS_FRONT_PART=reuse_freqs_front_part,
+                IS_NEOX=(rotate_style == RotateStyle.NEOX),
+                HAVE_POS=(positions is not None),
+                HAVE_OFFS=(offsets is not None),
+                BLOCK_T=BLOCK_T,
+                G=kh,
+                BLOCK_D=BLOCK_D,
+                BLOCK_D_HALF=BLOCK_D_HALF,
+                num_warps=num_warps,
+                waves_per_eu=waves_per_eu,
+            )
+
+    return out_x, out_y
+
+
+def rope_cached_thd_positions_2c_bwd(
+    x: torch.Tensor,
+    y: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    positions: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    transpose_output: bool = False,
+):
+    out_x = torch.empty(*x.shape, dtype=x.dtype, device=x.device, requires_grad=False)
+    out_y = torch.empty(*y.shape, dtype=x.dtype, device=x.device, requires_grad=False)
+
+    _rope_cached_thd_positions_offsets_2c_bwd(
+        x,
+        y,
+        out_x,
+        out_y,
+        cos,
+        sin,
+        positions,
+        None,
+        rotate_style,
+        reuse_freqs_front_part,
+        nope_first,
+        False,
+        transpose_output,
+    )
+
+    return out_x, out_y
+
+
+def rope_cached_thd_positions_offsets_2c_bwd(
+    x: torch.Tensor,
+    y: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    positions: torch.Tensor,
+    offsets: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    transpose_output: bool = False,
+):
+    out_x = torch.empty(*x.shape, dtype=x.dtype, device=x.device, requires_grad=False)
+    out_y = torch.empty(*y.shape, dtype=x.dtype, device=x.device, requires_grad=False)
+
+    _rope_cached_thd_positions_offsets_2c_bwd(
+        x,
+        y,
+        out_x,
+        out_y,
+        cos,
+        sin,
+        positions,
+        offsets,
+        rotate_style,
+        reuse_freqs_front_part,
+        nope_first,
+        False,
+        transpose_output,
+    )
+
+    return out_x, out_y
+
+
+def _rope_fwd_2d(
+    x: torch.Tensor,
+    out: torch.Tensor,
+    cos_h: torch.Tensor,
+    sin_h: torch.Tensor,
+    cos_w: torch.Tensor,
+    sin_w: torch.Tensor,
+    img_height: torch.Tensor,
+    img_width: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    transpose_output: bool = False,
+):
+    b, wh, h, d = x.shape
+    # out = torch.empty((b,wh,h,d), dtype=x.dtype, device=x.device, requires_grad=False)
+
+    grid = (b, h, 1)
+    _rope_fwd_2d_kernel_neox[grid](
+        x,
+        cos_h,
+        sin_h,
+        cos_w,
+        sin_w,
+        out,
+        *x.stride(),
+        *cos_h.stride(),
+        *cos_w.stride(),
+        wh,
+        img_height,
+        img_width,
+        BLOCK_D=d,
+    )
+
+    return out
+
+
+def rope_fwd_2d(
+    x: torch.Tensor,
+    cos_h: torch.Tensor,
+    sin_h: torch.Tensor,
+    cos_w: torch.Tensor,
+    sin_w: torch.Tensor,
+    img_height: torch.Tensor,
+    img_width: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    transpose_output: bool = False,
+):
+    b, wh, h, d = x.shape
+    out = torch.empty(
+        (b, wh, h, d), dtype=x.dtype, device=x.device, requires_grad=False
+    )
+
+    # grid = (b,h,1)
+    # _rope_fwd_2d_kernel_neox[grid](x, cos_h, sin_h, cos_w, sin_w, out, *x.stride(), *cos_h.stride(), *cos_w.stride(), wh, img_height, img_width, BLOCK_D=d)
+
+    _rope_fwd_2d(
+        x,
+        out,
+        cos_h,
+        sin_h,
+        cos_w,
+        sin_w,
+        img_height,
+        img_width,
+        rotate_style,
+        reuse_freqs_front_part,
+        nope_first,
+        transpose_output,
+    )
+
+    return out
+
+
+def rope_fwd_2d_inplace(
+    x: torch.Tensor,
+    cos_h: torch.Tensor,
+    sin_h: torch.Tensor,
+    cos_w: torch.Tensor,
+    sin_w: torch.Tensor,
+    img_height: torch.Tensor,
+    img_width: torch.Tensor,
+    rotate_style: int,
+    reuse_freqs_front_part: bool,
+    nope_first: bool,
+    transpose_output: bool = False,
+):
+    out = x
+    _rope_fwd_2d(
+        x,
+        out,
+        cos_h,
+        sin_h,
+        cos_w,
+        sin_w,
+        img_height,
+        img_width,
+        rotate_style,
+        reuse_freqs_front_part,
+        nope_first,
+        transpose_output,
+    )
+
+    return out
+
+
+def rope_fwd_3d(
+    x,
+    grid_sizes: tl.constexpr,
+    freqs: tl.constexpr,
+    sp_size: tl.constexpr,
+    sp_rank: tl.constexpr,
+):
+    B, s, n_heads, C = x.shape
+    c_total = C // 2  # 64
+    c1 = c_total - 2 * (c_total // 3)  # 22
+    c2 = c_total // 3  # 21
+    c3 = c_total // 3  # 21
+    device = x.device
+
+    grid_sizes = grid_sizes.to(device=device, dtype=torch.int32).contiguous()
+
+    freqs_real = freqs.real.to(dtype=torch.float32, device=device).contiguous()
+    freqs_imag = freqs.imag.to(dtype=torch.float32, device=device).contiguous()
+    out = torch.empty_like(x, dtype=torch.float32, device=device)
+
+    BLOCK_L, BLOCK_N, BLOCK_C = 32, 4, 64
+
+    grid = (B, n_heads, triton.cdiv(s, BLOCK_L))
+
+    num_warps = 4
+    waves_per_eu = 1
+
+    _rope_fwd_3d[grid](
+        x,
+        freqs_real,
+        freqs_imag,
+        grid_sizes,
+        out,
+        *x.stride(),
+        freqs_real.stride(0),
+        freqs_real.stride(1),
+        *grid_sizes.stride(),
+        *out.stride(),
+        s,
+        n_heads,
+        C,
+        c_total,
+        sp_size,
+        sp_rank,
+        freqs.shape[0],
+        s,
+        1.0,
+        0.0,
+        BLOCK_L=BLOCK_L,
+        BLOCK_N=BLOCK_N,
+        BLOCK_C=BLOCK_C,
+        C1=c1,
+        C2=c2,
+        num_warps=num_warps,
+        waves_per_eu=waves_per_eu,
+    )
+
+    return out
+
+
+class RoPE(autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        x: torch.Tensor,
+        freqs: torch.Tensor,
+        rotate_style: int,
+        reuse_freqs_front_part: bool,
+        nope_first: bool,
+        transpose_output: bool = False,
+    ) -> torch.Tensor:
+        ctx.rotate_style = rotate_style
+        ctx.reuse_freqs_front_part = reuse_freqs_front_part
+        ctx.nope_first = nope_first
+        ctx.transpose_output = transpose_output
+        ctx.save_for_backward(freqs)
+        return rope_fwd(
+            x, freqs, rotate_style, reuse_freqs_front_part, nope_first, transpose_output
+        )
+
+    @staticmethod
+    def backward(
+        ctx, output_grads: torch.Tensor
+    ) -> Tuple[Union[torch.Tensor, None], ...]:
+        (freqs,) = ctx.saved_tensors
+        return (
+            rope_bwd(
+                output_grads,
+                freqs,
+                ctx.rotate_style,
+                ctx.reuse_freqs_front_part,
+                ctx.nope_first,
+                ctx.transpose_output,
+            ),
+            None,
+            None,
+        )
+
+
+class RoPETHD(autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        x: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        freqs: torch.Tensor,
+        rotate_style: int,
+        reuse_freqs_front_part: bool,
+        nope_first: bool,
+    ):
+        ctx.rotate_style = rotate_style
+        ctx.reuse_freqs_front_part = reuse_freqs_front_part
+        ctx.nope_first = nope_first
+        ctx.save_for_backward(cu_seqlens, freqs)
+        return rope_thd_fwd(
+            x, cu_seqlens, freqs, rotate_style, reuse_freqs_front_part, nope_first
+        )
+
+    @staticmethod
+    def backward(ctx, output_grads) -> Tuple[Union[torch.Tensor, None], ...]:
+        cu_seqlens, freqs = ctx.saved_tensors
+        return (
+            rope_thd_bwd(
+                output_grads,
+                cu_seqlens,
+                freqs,
+                ctx.rotate_style,
+                ctx.reuse_freqs_front_part,
+                ctx.nope_first,
+            ),
+            None,
+            None,
+        )
+
+
+class RoPECached(autograd.Function):
+
+    @staticmethod
+    def forward(
+        ctx,
+        x: torch.Tensor,
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        rotate_style: int,
+        reuse_freqs_front_part: bool,
+        nope_first: bool,
+        transpose_output: bool = False,
+    ) -> torch.Tensor:
+        ctx.rotate_style = rotate_style
+        ctx.reuse_freqs_front_part = reuse_freqs_front_part
+        ctx.nope_first = nope_first
+        ctx.transpose_output = transpose_output
+        ctx.save_for_backward(cos, sin)
+        return rope_cached_fwd(
+            x,
+            cos,
+            sin,
+            rotate_style,
+            reuse_freqs_front_part,
+            nope_first,
+            transpose_output,
+        )
+
+    @staticmethod
+    def backward(ctx, output_grads) -> Tuple[Union[torch.Tensor, None], ...]:
+        cos, sin = ctx.saved_tensors
+        return (
+            rope_cached_bwd(
+                output_grads,
+                cos,
+                sin,
+                ctx.rotate_style,
+                ctx.reuse_freqs_front_part,
+                ctx.nope_first,
+                ctx.transpose_output,
+            ),
+            None,
+            None,
+        )
+
+
+class RoPE2D(autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        x: torch.Tensor,
+        cos_height: torch.Tensor,
+        sin_height: torch.Tensor,
+        cos_width: torch.Tensor,
+        sin_width: torch.Tensor,
+        img_height: int,
+        img_width: int,
+        rotate_style: int,
+        reuse_freqs_front_part: bool,
+        nope_first: bool,
+    ) -> torch.Tensor:
+        ctx.img_height = img_height
+        ctx.img_width = img_width
+        ctx.rotate_style = rotate_style
+        ctx.reuse_freqs_front_part = reuse_freqs_front_part
+        ctx.nope_first = nope_first
+        ctx.save_for_backward(cos_height, sin_height, cos_width, sin_width)
+        return rope_fwd_2d(
+            x,
+            cos_height,
+            sin_height,
+            cos_width,
+            sin_width,
+            img_height,
+            img_width,
+            rotate_style,
+            reuse_freqs_front_part,
+            nope_first,
+        )

build/torch-rocm/utils/logger.py

@@ -0,0 +1,47 @@
+import os
+import logging
+
+
+# AITER Triton Logger which is singleton object around python logging.
+# Note: Python logging is also a singleton object, but we want to read the
+# env var AITER_LOG_LEVEL once at the beginning. Another alternative is to do
+# this in __init__.py. In fact, that's how CK logger is setup. We can look at
+# switching to that at some point
+#
+# AITER_LOG_LEVEL follows python logging levels
+#   DEBUG
+#   INFO
+#   WARNING
+#   ERROR
+#   CRITICAL
+#
+class AiterTritonLogger(object):
+    _instance = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = super(AiterTritonLogger, cls).__new__(cls)
+            log_level_str = os.getenv("AITER_TRITON_LOG_LEVEL", "WARNING").upper()
+            numeric_level = getattr(logging, log_level_str, logging.WARNING)
+            cls._instance._logger = logging.getLogger("AITER_TRITON")
+            cls._instance._logger.setLevel(numeric_level)
+
+        return cls._instance
+
+    def get_logger(self):
+        return self._logger
+
+    def debug(self, msg):
+        self._logger.debug(msg)
+
+    def info(self, msg):
+        self._logger.info(msg)
+
+    def warning(self, msg):
+        self._logger.warning(msg)
+
+    def error(self, msg):
+        self._logger.error(msg)
+
+    def critical(self, msg):
+        self._logger.critical(msg)