feat: MLA RoPE Triton kernels (port from llm-training)

Ports Motif3 GDLAttention RoPE optimization kernels from llm-training
to the activation package. Python/Triton-only addition — no CUDA
recompilation needed.

## Kernels added (fused_rope.py)

### _rope_fwd_kernel / _rope_bwd_kernel (Phase 1)
Generic RoPE on [B*S, H, rope_dim] with sin/cos direct multiplication,
interleaved → contiguous layout change inside kernel. Eliminates
view_as_complex + complex_mul + reorder_headdim pattern. Used for k_pe
(1 head, standalone) after Phase 4 landed (Q path uses the in-place
variant below).

### _kv_rope_fwd_kernel (Phase 2)
Fused KV split + k_pe RoPE broadcast + cat. Reads kv_latent (k_nope+v)
and k_pe (1 head, already RoPE'd), outputs key=[k_nope||k_pe_broadcast]
and value. k_pe broadcast across heads happens in register — no
memory copy.

### _q_rope_inplace_fwd_kernel / _q_rope_bwd_kernel (Phase 4)
In-place Q RoPE. Modifies q[..., nope_dim:] from interleaved to
contiguous format without touching the nope section. Eliminates the
subsequent cat(q_nope, q_pe) shuffler. Uses torch.library.custom_op
with mutates_args=("q",) + ctx.mark_dirty(q) for autograd safety.

Distinguishes from the reverted Phase 3 (out-of-place) which copied the
nope section into a new buffer — that cost more than the cat it
replaced. In-place avoids that by never touching nope.

### Autotune (Phase 5)
All kernels autotune over BLOCK_H × num_warps × num_stages = 128
configs via _rope_autotune_configs(). Backward kernels gained ~30%
from this vs BLOCK_H-only search; forward ~5-10%. _kv_rope_fwd_kernel
was already at ~73% HBM BW so showed no improvement.

## Public API

from activation import (
fused_apply_rope, # generic RoPE (Phase 1)
fused_q_rope_inplace, # Q in-place RoPE (Phase 4)
fused_kv_split_rope_cat, # KV split+RoPE+cat (Phase 2)
)

## Measured perf (motif3_seq, B=4, S=4096, n_layers=6, B200)

Attention block per layer (fwd+bwd), vs baseline (PyTorch native MLA):

Phase 0 (baseline) 14,730 µs (reference)
Phase 1 (Triton RoPE) 13,046 (-11.4%)
Phase 2 (KV fused) 12,956 (-12.0%)
Phase 4 (in-place Q) ~12,371 (-16.0%)
Phase 5 (wide autotune) ~12,125 (-17.7%)

(Phase 6 — FA4 dense wrapper — stays in llm-training since it's an
external flash-attn wrapper, not a Motif custom kernel.)

Loss matches baseline to 4 decimal places at step 10 (0.0001). No NaN.

torch.library op namespace 'motif::*' retained (no conflict with
activation's existing TORCH_LIBRARY_EXPAND-registered CUDA ops).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

torch-ext/activation/__init__.py

@@ -2,6 +2,11 @@ import torch
 
 from . import layers, parallel_style
 from ._ops import ops
+from .fused_rope import (
+    fused_apply_rope,
+    fused_kv_split_rope_cat,
+    fused_q_rope_inplace,
+)
 from .grouped_poly_norm import fused_mul_grouped_poly_norm
 from .poly_norm import FusedMulPolyNormFunction, PolyNormFunction
 from .rms_norm import FusedAddRMSNormFunction, RMSNormFunction
@@ -49,6 +54,9 @@ __all__ = [
     "fused_mul_grouped_poly_norm",
     "rms_norm",
     "fused_add_rms_norm",
+    "fused_apply_rope",
+    "fused_q_rope_inplace",
+    "fused_kv_split_rope_cat",
     "layers",
     "parallel_style",
     "ops",

torch-ext/activation/fused_rope.py

@@ -0,0 +1,723 @@
+"""Fused MLA RoPE kernels for Motif3 GDLAttention.
+
+Applies RoPE to the input tensor and outputs in contiguous format
+[real..., imag...] so no reorder_headdim_elements_rope is needed.
+
+Registered as torch custom_op for torch.compile compatibility (no graph break).
+
+Reference: Megatron-LM fused_mla_yarn_rope_apply.py
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+# ---------------------------------------------------------------------------
+# Autotune search space (BLOCK_H × num_warps × num_stages = 128 configs/kernel)
+# Triton caches per-(key) best config — one-time compile cost per distinct shape.
+# ---------------------------------------------------------------------------
+
+def _rope_autotune_configs():
+    configs = []
+    for bh in [1, 2, 4, 8, 16, 32, 64, 128]:
+        for nw in [1, 2, 4, 8]:
+            for ns in [1, 2, 3, 4]:
+                configs.append(triton.Config({"BLOCK_H": bh}, num_warps=nw, num_stages=ns))
+    return configs
+
+
+# ---------------------------------------------------------------------------
+# Triton kernels
+# ---------------------------------------------------------------------------
+
+@triton.autotune(
+    configs=_rope_autotune_configs(),
+    key=["rope_dim", "head_num"],
+)
+@triton.jit
+def _rope_fwd_kernel(
+    X,  # [B*S, H, rope_dim] interleaved
+    O,  # [B*S, H, rope_dim] contiguous
+    COS,  # [S, rope_dim // 2]
+    SIN,  # [S, rope_dim // 2]
+    rope_dim: tl.constexpr,
+    head_num: tl.constexpr,
+    seq_len,
+    stride_x_token,
+    stride_x_head,
+    stride_o_token,
+    stride_o_head,
+    BLOCK_H: tl.constexpr,
+):
+    HALF: tl.constexpr = rope_dim // 2
+
+    pid_token = tl.program_id(0)
+    pid_hblock = tl.program_id(1)
+
+    pos = pid_token % seq_len
+    cos = tl.load(COS + pos * HALF + tl.arange(0, HALF))
+    sin = tl.load(SIN + pos * HALF + tl.arange(0, HALF))
+    cos = cos.expand_dims(0).broadcast_to(BLOCK_H, HALF)
+    sin = sin.expand_dims(0).broadcast_to(BLOCK_H, HALF)
+
+    X_ptr = X + pid_token * stride_x_token + pid_hblock * BLOCK_H * stride_x_head
+    O_ptr = O + pid_token * stride_o_token + pid_hblock * BLOCK_H * stride_o_head
+
+    h_off = tl.arange(0, BLOCK_H)[:, None]
+    mask = (pid_hblock * BLOCK_H + h_off) < head_num
+
+    x1_off = h_off * stride_x_head + tl.arange(0, HALF)[None, :] * 2
+    x2_off = x1_off + 1
+    x1 = tl.load(X_ptr + x1_off, mask=mask).to(tl.float32)
+    x2 = tl.load(X_ptr + x2_off, mask=mask).to(tl.float32)
+
+    out_real = x1 * cos - x2 * sin
+    out_imag = x1 * sin + x2 * cos
+
+    real_off = h_off * stride_o_head + tl.arange(0, HALF)[None, :]
+    imag_off = real_off + HALF
+    tl.store(O_ptr + real_off, out_real, mask=mask)
+    tl.store(O_ptr + imag_off, out_imag, mask=mask)
+
+
+@triton.autotune(
+    configs=_rope_autotune_configs(),
+    key=["rope_dim", "head_num"],
+)
+@triton.jit
+def _rope_bwd_kernel(
+    DO,  # [B*S, H, rope_dim] contiguous
+    DX,  # [B*S, H, rope_dim] interleaved
+    COS,  # [S, rope_dim // 2]
+    SIN,  # [S, rope_dim // 2]
+    rope_dim: tl.constexpr,
+    head_num: tl.constexpr,
+    seq_len,
+    stride_do_token,
+    stride_do_head,
+    stride_dx_token,
+    stride_dx_head,
+    BLOCK_H: tl.constexpr,
+):
+    HALF: tl.constexpr = rope_dim // 2
+
+    pid_token = tl.program_id(0)
+    pid_hblock = tl.program_id(1)
+
+    pos = pid_token % seq_len
+    cos = tl.load(COS + pos * HALF + tl.arange(0, HALF))
+    sin = tl.load(SIN + pos * HALF + tl.arange(0, HALF))
+    cos = cos.expand_dims(0).broadcast_to(BLOCK_H, HALF)
+    sin = sin.expand_dims(0).broadcast_to(BLOCK_H, HALF)
+
+    DO_ptr = DO + pid_token * stride_do_token + pid_hblock * BLOCK_H * stride_do_head
+    DX_ptr = DX + pid_token * stride_dx_token + pid_hblock * BLOCK_H * stride_dx_head
+
+    h_off = tl.arange(0, BLOCK_H)[:, None]
+    mask = (pid_hblock * BLOCK_H + h_off) < head_num
+
+    real_off = h_off * stride_do_head + tl.arange(0, HALF)[None, :]
+    imag_off = real_off + HALF
+    d_real = tl.load(DO_ptr + real_off, mask=mask).to(tl.float32)
+    d_imag = tl.load(DO_ptr + imag_off, mask=mask).to(tl.float32)
+
+    dx1 = d_real * cos + d_imag * sin
+    dx2 = -d_real * sin + d_imag * cos
+
+    x1_off = h_off * stride_dx_head + tl.arange(0, HALF)[None, :] * 2
+    x2_off = x1_off + 1
+    tl.store(DX_ptr + x1_off, dx1, mask=mask)
+    tl.store(DX_ptr + x2_off, dx2, mask=mask)
+
+
+# ---------------------------------------------------------------------------
+# custom_op registration (no graph break under torch.compile)
+# ---------------------------------------------------------------------------
+
+@torch.library.custom_op("motif::rope_fwd", mutates_args=())
+def _rope_fwd(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    seq_len: int,
+) -> torch.Tensor:
+    B_S, H, D = x.shape
+    out = torch.empty_like(x)
+    grid = lambda META: (B_S, triton.cdiv(H, META["BLOCK_H"]))
+    _rope_fwd_kernel[grid](
+        x, out, cos, sin,
+        D, H, seq_len,
+        x.stride(0), x.stride(1),
+        out.stride(0), out.stride(1),
+    )
+    return out
+
+
+@_rope_fwd.register_fake
+def _rope_fwd_fake(x, cos, sin, seq_len):
+    return torch.empty_like(x)
+
+
+@torch.library.custom_op("motif::rope_bwd", mutates_args=())
+def _rope_bwd(
+    grad_out: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    seq_len: int,
+) -> torch.Tensor:
+    B_S, H, D = grad_out.shape
+    grad_in = torch.empty_like(grad_out)
+    grid = lambda META: (B_S, triton.cdiv(H, META["BLOCK_H"]))
+    _rope_bwd_kernel[grid](
+        grad_out, grad_in, cos, sin,
+        D, H, seq_len,
+        grad_out.stride(0), grad_out.stride(1),
+        grad_in.stride(0), grad_in.stride(1),
+    )
+    return grad_in
+
+
+@_rope_bwd.register_fake
+def _rope_bwd_fake(grad_out, cos, sin, seq_len):
+    return torch.empty_like(grad_out)
+
+
+# ---------------------------------------------------------------------------
+# Autograd wrapper
+# ---------------------------------------------------------------------------
+
+class FusedRoPE(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, cos, sin, seq_len):
+        # x: [B, S, H, rope_dim] interleaved
+        B, S, H, D = x.shape
+        out_3d = _rope_fwd(x.reshape(B * S, H, D), cos, sin, seq_len)
+        ctx.save_for_backward(cos, sin)
+        ctx.seq_len = seq_len
+        ctx.shape = (B, S, H, D)
+        return out_3d.view(B, S, H, D)
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        cos, sin = ctx.saved_tensors
+        B, S, H, D = ctx.shape
+        grad_in_3d = _rope_bwd(
+            grad_out.contiguous().reshape(B * S, H, D),
+            cos, sin, ctx.seq_len,
+        )
+        return grad_in_3d.view(B, S, H, D), None, None, None
+
+
+def fused_apply_rope(x, freqs_cis):
+    """Apply fused RoPE, outputting contiguous format. No graph break.
+
+    Args:
+        x: [B, S, H, rope_dim] interleaved format
+        freqs_cis: [max_seq_len, rope_dim//2] complex64
+
+    Returns:
+        [B, S, H, rope_dim] contiguous format [real..., imag...]
+    """
+    S = x.shape[1]
+    cos = freqs_cis[:S].real.contiguous()
+    sin = freqs_cis[:S].imag.contiguous()
+    return FusedRoPE.apply(x, cos, sin, S)
+
+
+# ---------------------------------------------------------------------------
+# Phase 2: Fused KV split + RoPE + expand + cat kernel
+# Fuses: split(kv_latent, [k_dim, v_dim]) + RoPE(k_pe) + expand + cat → key, value
+# Reference: Megatron rotary_fwd_kv_kernel
+# ---------------------------------------------------------------------------
+
+@triton.autotune(
+    configs=_rope_autotune_configs(),
+    key=["rope_dim", "k_nope_dim", "v_dim", "head_num"],
+)
+@triton.jit
+def _kv_rope_fwd_kernel(
+    KV,       # [B*S, H, k_nope_dim + v_dim]
+    K_PE,     # [B*S, 1, rope_dim] (already RoPE'd, contiguous format)
+    O_KEY,    # [B*S, H, k_nope_dim + rope_dim]  output
+    O_VALUE,  # [B*S, H, v_dim]                  output
+    rope_dim: tl.constexpr,
+    k_nope_dim: tl.constexpr,
+    v_dim: tl.constexpr,
+    head_num: tl.constexpr,
+    stride_kv_token,
+    stride_kv_head,
+    stride_pe_token,
+    stride_k_token,
+    stride_k_head,
+    stride_v_token,
+    stride_v_head,
+    BLOCK_H: tl.constexpr,
+):
+    pid_token = tl.program_id(0)
+    pid_hblock = tl.program_id(1)
+
+    KV_ptr = KV + pid_token * stride_kv_token + pid_hblock * BLOCK_H * stride_kv_head
+    K_ptr = O_KEY + pid_token * stride_k_token + pid_hblock * BLOCK_H * stride_k_head
+    V_ptr = O_VALUE + pid_token * stride_v_token + pid_hblock * BLOCK_H * stride_v_head
+
+    h_off = tl.arange(0, BLOCK_H)[:, None]
+    mask = (pid_hblock * BLOCK_H + h_off) < head_num
+
+    # Read k_nope from KV
+    k_nope_off = h_off * stride_kv_head + tl.arange(0, k_nope_dim)[None, :]
+    k_nope = tl.load(KV_ptr + k_nope_off, mask=mask)
+
+    # Read v from KV
+    v_off = h_off * stride_kv_head + k_nope_dim + tl.arange(0, v_dim)[None, :]
+    v = tl.load(KV_ptr + v_off, mask=mask)
+
+    # Read k_pe (shared across all heads, already RoPE'd)
+    # K_PE is [B*S, 1, rope_dim], broadcast to all heads
+    pe_ptr = K_PE + pid_token * stride_pe_token
+    k_pe = tl.load(pe_ptr + tl.arange(0, rope_dim)[None, :])
+    k_pe = k_pe.broadcast_to(BLOCK_H, rope_dim)
+
+    # Write key = [k_nope | k_pe]
+    k_nope_out = h_off * stride_k_head + tl.arange(0, k_nope_dim)[None, :]
+    tl.store(K_ptr + k_nope_out, k_nope, mask=mask)
+    k_pe_out = h_off * stride_k_head + k_nope_dim + tl.arange(0, rope_dim)[None, :]
+    tl.store(K_ptr + k_pe_out, k_pe, mask=mask)
+
+    # Write value
+    v_out = h_off * stride_v_head + tl.arange(0, v_dim)[None, :]
+    tl.store(V_ptr + v_out, v, mask=mask)
+
+
+@torch.library.custom_op("motif::kv_rope_fwd", mutates_args=())
+def _kv_rope_fwd(
+    kv_latent: torch.Tensor,  # [B*S, H, k_nope_dim + v_dim]
+    k_pe: torch.Tensor,       # [B*S, 1, rope_dim] (already RoPE'd)
+    k_nope_dim: int,
+    v_dim: int,
+    rope_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    B_S, H, _ = kv_latent.shape
+    key = kv_latent.new_empty(B_S, H, k_nope_dim + rope_dim)
+    value = kv_latent.new_empty(B_S, H, v_dim)
+    grid = lambda META: (B_S, triton.cdiv(H, META["BLOCK_H"]))
+    _kv_rope_fwd_kernel[grid](
+        kv_latent, k_pe, key, value,
+        rope_dim, k_nope_dim, v_dim, H,
+        kv_latent.stride(0), kv_latent.stride(1),
+        k_pe.stride(0),
+        key.stride(0), key.stride(1),
+        value.stride(0), value.stride(1),
+    )
+    return key, value
+
+
+@_kv_rope_fwd.register_fake
+def _kv_rope_fwd_fake(kv_latent, k_pe, k_nope_dim, v_dim, rope_dim):
+    B_S, H, _ = kv_latent.shape
+    key = kv_latent.new_empty(B_S, H, k_nope_dim + rope_dim)
+    value = kv_latent.new_empty(B_S, H, v_dim)
+    return key, value
+
+
+class FusedKVRoPE(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, kv_latent, k_pe, k_nope_dim, v_dim, rope_dim):
+        # kv_latent: [B, S, H, k_nope_dim + v_dim]
+        # k_pe: [B, S, 1, rope_dim] (already RoPE'd)
+        B, S, H, D = kv_latent.shape
+        key_3d, value_3d = _kv_rope_fwd(
+            kv_latent.reshape(B * S, H, D),
+            k_pe.reshape(B * S, 1, rope_dim),
+            k_nope_dim, v_dim, rope_dim,
+        )
+        ctx.save_for_backward(k_pe)
+        ctx.k_nope_dim = k_nope_dim
+        ctx.v_dim = v_dim
+        ctx.rope_dim = rope_dim
+        ctx.shape = (B, S, H)
+        return key_3d.view(B, S, H, k_nope_dim + rope_dim), value_3d.view(B, S, H, v_dim)
+
+    @staticmethod
+    def backward(ctx, grad_key, grad_value):
+        k_pe, = ctx.saved_tensors
+        B, S, H = ctx.shape
+        k_nope_dim = ctx.k_nope_dim
+        v_dim = ctx.v_dim
+        rope_dim = ctx.rope_dim
+
+        # grad_key: [B, S, H, k_nope_dim + rope_dim]
+        # grad_value: [B, S, H, v_dim]
+        # Reconstruct grad_kv_latent = [grad_k_nope | grad_value]
+        grad_k_nope = grad_key[..., :k_nope_dim]
+        grad_kv_latent = torch.cat([grad_k_nope, grad_value], dim=-1)
+
+        # grad_k_pe: sum across heads (k_pe was broadcast)
+        grad_k_pe = grad_key[..., k_nope_dim:].sum(dim=2, keepdim=True)
+
+        return grad_kv_latent, grad_k_pe, None, None, None
+
+
+# ---------------------------------------------------------------------------
+# Phase 3: Fused Q nope passthrough + RoPE kernel (eliminates cat)
+# Q [B*S, H, nope_dim + rope_dim] → Q_out [B*S, H, nope_dim + rope_dim]
+# First nope_dim elements: copy as-is
+# Last rope_dim elements: apply RoPE (interleaved → contiguous)
+# ---------------------------------------------------------------------------
+
+@triton.autotune(
+    configs=_rope_autotune_configs(),
+    key=["nope_dim", "rope_dim", "head_num"],
+)
+@triton.jit
+def _q_rope_fwd_kernel(
+    Q,      # [B*S, H, nope_dim + rope_dim] input
+    O,      # [B*S, H, nope_dim + rope_dim] output
+    COS,    # [S, rope_dim // 2]
+    SIN,    # [S, rope_dim // 2]
+    nope_dim: tl.constexpr,
+    rope_dim: tl.constexpr,
+    head_num: tl.constexpr,
+    seq_len,
+    stride_q_token,
+    stride_q_head,
+    stride_o_token,
+    stride_o_head,
+    BLOCK_H: tl.constexpr,
+):
+    HALF: tl.constexpr = rope_dim // 2
+
+    pid_token = tl.program_id(0)
+    pid_hblock = tl.program_id(1)
+
+    pos = pid_token % seq_len
+
+    Q_ptr = Q + pid_token * stride_q_token + pid_hblock * BLOCK_H * stride_q_head
+    O_ptr = O + pid_token * stride_o_token + pid_hblock * BLOCK_H * stride_o_head
+
+    h_off = tl.arange(0, BLOCK_H)[:, None]
+    mask = (pid_hblock * BLOCK_H + h_off) < head_num
+
+    # Copy nope portion as-is
+    nope_off = h_off * stride_q_head + tl.arange(0, nope_dim)[None, :]
+    nope_data = tl.load(Q_ptr + nope_off, mask=mask)
+    nope_out = h_off * stride_o_head + tl.arange(0, nope_dim)[None, :]
+    tl.store(O_ptr + nope_out, nope_data, mask=mask)
+
+    # RoPE on rope portion: read interleaved, write contiguous
+    cos = tl.load(COS + pos * HALF + tl.arange(0, HALF))
+    sin = tl.load(SIN + pos * HALF + tl.arange(0, HALF))
+    cos = cos.expand_dims(0).broadcast_to(BLOCK_H, HALF)
+    sin = sin.expand_dims(0).broadcast_to(BLOCK_H, HALF)
+
+    x1_off = h_off * stride_q_head + nope_dim + tl.arange(0, HALF)[None, :] * 2
+    x2_off = x1_off + 1
+    x1 = tl.load(Q_ptr + x1_off, mask=mask).to(tl.float32)
+    x2 = tl.load(Q_ptr + x2_off, mask=mask).to(tl.float32)
+
+    out_real = x1 * cos - x2 * sin
+    out_imag = x1 * sin + x2 * cos
+
+    real_off = h_off * stride_o_head + nope_dim + tl.arange(0, HALF)[None, :]
+    imag_off = real_off + HALF
+    tl.store(O_ptr + real_off, out_real, mask=mask)
+    tl.store(O_ptr + imag_off, out_imag, mask=mask)
+
+
+@triton.autotune(
+    configs=_rope_autotune_configs(),
+    key=["nope_dim", "rope_dim", "head_num"],
+)
+@triton.jit
+def _q_rope_bwd_kernel(
+    DO,     # [B*S, H, nope_dim + rope_dim] grad (contiguous rope fmt)
+    DQ,     # [B*S, H, nope_dim + rope_dim] grad output (interleaved rope fmt)
+    COS,    # [S, rope_dim // 2]
+    SIN,    # [S, rope_dim // 2]
+    nope_dim: tl.constexpr,
+    rope_dim: tl.constexpr,
+    head_num: tl.constexpr,
+    seq_len,
+    stride_do_token,
+    stride_do_head,
+    stride_dq_token,
+    stride_dq_head,
+    BLOCK_H: tl.constexpr,
+):
+    HALF: tl.constexpr = rope_dim // 2
+
+    pid_token = tl.program_id(0)
+    pid_hblock = tl.program_id(1)
+
+    pos = pid_token % seq_len
+
+    DO_ptr = DO + pid_token * stride_do_token + pid_hblock * BLOCK_H * stride_do_head
+    DQ_ptr = DQ + pid_token * stride_dq_token + pid_hblock * BLOCK_H * stride_dq_head
+
+    h_off = tl.arange(0, BLOCK_H)[:, None]
+    mask = (pid_hblock * BLOCK_H + h_off) < head_num
+
+    # Copy nope grad as-is
+    nope_off_in = h_off * stride_do_head + tl.arange(0, nope_dim)[None, :]
+    nope_grad = tl.load(DO_ptr + nope_off_in, mask=mask)
+    nope_off_out = h_off * stride_dq_head + tl.arange(0, nope_dim)[None, :]
+    tl.store(DQ_ptr + nope_off_out, nope_grad, mask=mask)
+
+    # Inverse RoPE: contiguous → interleaved
+    cos = tl.load(COS + pos * HALF + tl.arange(0, HALF))
+    sin = tl.load(SIN + pos * HALF + tl.arange(0, HALF))
+    cos = cos.expand_dims(0).broadcast_to(BLOCK_H, HALF)
+    sin = sin.expand_dims(0).broadcast_to(BLOCK_H, HALF)
+
+    real_off = h_off * stride_do_head + nope_dim + tl.arange(0, HALF)[None, :]
+    imag_off = real_off + HALF
+    d_real = tl.load(DO_ptr + real_off, mask=mask).to(tl.float32)
+    d_imag = tl.load(DO_ptr + imag_off, mask=mask).to(tl.float32)
+
+    dx1 = d_real * cos + d_imag * sin
+    dx2 = -d_real * sin + d_imag * cos
+
+    x1_off = h_off * stride_dq_head + nope_dim + tl.arange(0, HALF)[None, :] * 2
+    x2_off = x1_off + 1
+    tl.store(DQ_ptr + x1_off, dx1, mask=mask)
+    tl.store(DQ_ptr + x2_off, dx2, mask=mask)
+
+
+@torch.library.custom_op("motif::q_rope_fwd", mutates_args=())
+def _q_rope_fwd(
+    q: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    nope_dim: int,
+    rope_dim: int,
+    seq_len: int,
+) -> torch.Tensor:
+    B_S, H, D = q.shape
+    out = torch.empty_like(q)
+    grid = lambda META: (B_S, triton.cdiv(H, META["BLOCK_H"]))
+    _q_rope_fwd_kernel[grid](
+        q, out, cos, sin,
+        nope_dim, rope_dim, H, seq_len,
+        q.stride(0), q.stride(1),
+        out.stride(0), out.stride(1),
+    )
+    return out
+
+
+@_q_rope_fwd.register_fake
+def _q_rope_fwd_fake(q, cos, sin, nope_dim, rope_dim, seq_len):
+    return torch.empty_like(q)
+
+
+@torch.library.custom_op("motif::q_rope_bwd", mutates_args=())
+def _q_rope_bwd(
+    grad_out: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    nope_dim: int,
+    rope_dim: int,
+    seq_len: int,
+) -> torch.Tensor:
+    B_S, H, D = grad_out.shape
+    grad_in = torch.empty_like(grad_out)
+    grid = lambda META: (B_S, triton.cdiv(H, META["BLOCK_H"]))
+    _q_rope_bwd_kernel[grid](
+        grad_out, grad_in, cos, sin,
+        nope_dim, rope_dim, H, seq_len,
+        grad_out.stride(0), grad_out.stride(1),
+        grad_in.stride(0), grad_in.stride(1),
+    )
+    return grad_in
+
+
+@_q_rope_bwd.register_fake
+def _q_rope_bwd_fake(grad_out, cos, sin, nope_dim, rope_dim, seq_len):
+    return torch.empty_like(grad_out)
+
+
+class FusedQRoPE(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, q, cos, sin, nope_dim, rope_dim, seq_len):
+        B, S, H, D = q.shape
+        out_3d = _q_rope_fwd(q.reshape(B * S, H, D), cos, sin, nope_dim, rope_dim, seq_len)
+        ctx.save_for_backward(cos, sin)
+        ctx.nope_dim = nope_dim
+        ctx.rope_dim = rope_dim
+        ctx.seq_len = seq_len
+        ctx.shape = (B, S, H, D)
+        return out_3d.view(B, S, H, D)
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        cos, sin = ctx.saved_tensors
+        B, S, H, D = ctx.shape
+        grad_in_3d = _q_rope_bwd(
+            grad_out.contiguous().reshape(B * S, H, D),
+            cos, sin, ctx.nope_dim, ctx.rope_dim, ctx.seq_len,
+        )
+        return grad_in_3d.view(B, S, H, D), None, None, None, None, None
+
+
+def fused_q_rope(q, freqs_cis, nope_dim, rope_dim):
+    """Fused Q nope passthrough + RoPE. Eliminates split + cat.
+
+    Replaces:
+        q_nope, q_pe = split(q, [nope_dim, rope_dim])
+        q_pe = fused_apply_rope(q_pe, freqs_cis)
+        q_total = cat([q_nope, q_pe])
+
+    Args:
+        q: [B, S, H, nope_dim + rope_dim] from wq_b output
+        freqs_cis: [max_seq_len, rope_dim//2] complex64
+
+    Returns:
+        [B, S, H, nope_dim + rope_dim] with nope unchanged, rope in contiguous format
+    """
+    S = q.shape[1]
+    cos = freqs_cis[:S].real.contiguous()
+    sin = freqs_cis[:S].imag.contiguous()
+    return FusedQRoPE.apply(q, cos, sin, nope_dim, rope_dim, S)
+
+
+# ---------------------------------------------------------------------------
+# Phase 4: In-place Q RoPE kernel (eliminates cat + nope copy)
+# Modifies q[..., nope_dim:] from interleaved → contiguous format IN-PLACE.
+# nope section [..., :nope_dim] is untouched.
+# Forward: in-place on Q (mutates)
+# Backward: out-of-place via existing _q_rope_bwd_kernel (nope grad copy + inverse rope)
+# ---------------------------------------------------------------------------
+
+@triton.autotune(
+    configs=_rope_autotune_configs(),
+    key=["nope_dim", "rope_dim", "head_num"],
+)
+@triton.jit
+def _q_rope_inplace_fwd_kernel(
+    Q,  # [B*S, H, nope_dim + rope_dim]  modified in-place on [..., nope_dim:]
+    COS,  # [S, rope_dim // 2]
+    SIN,  # [S, rope_dim // 2]
+    nope_dim: tl.constexpr,
+    rope_dim: tl.constexpr,
+    head_num: tl.constexpr,
+    seq_len,
+    stride_q_token,
+    stride_q_head,
+    BLOCK_H: tl.constexpr,
+):
+    HALF: tl.constexpr = rope_dim // 2
+
+    pid_token = tl.program_id(0)
+    pid_hblock = tl.program_id(1)
+
+    pos = pid_token % seq_len
+    cos = tl.load(COS + pos * HALF + tl.arange(0, HALF))
+    sin = tl.load(SIN + pos * HALF + tl.arange(0, HALF))
+    cos = cos.expand_dims(0).broadcast_to(BLOCK_H, HALF)
+    sin = sin.expand_dims(0).broadcast_to(BLOCK_H, HALF)
+
+    Q_ptr = Q + pid_token * stride_q_token + pid_hblock * BLOCK_H * stride_q_head
+
+    h_off = tl.arange(0, BLOCK_H)[:, None]
+    mask = (pid_hblock * BLOCK_H + h_off) < head_num
+
+    # Read rope section interleaved: [r0,i0,r1,i1,...]
+    x1_off = h_off * stride_q_head + nope_dim + tl.arange(0, HALF)[None, :] * 2
+    x2_off = x1_off + 1
+    x1 = tl.load(Q_ptr + x1_off, mask=mask).to(tl.float32)
+    x2 = tl.load(Q_ptr + x2_off, mask=mask).to(tl.float32)
+
+    out_real = x1 * cos - x2 * sin
+    out_imag = x1 * sin + x2 * cos
+
+    # Write back to SAME rope section in contiguous format: [r0..r31, i0..i31]
+    real_off = h_off * stride_q_head + nope_dim + tl.arange(0, HALF)[None, :]
+    imag_off = real_off + HALF
+    tl.store(Q_ptr + real_off, out_real, mask=mask)
+    tl.store(Q_ptr + imag_off, out_imag, mask=mask)
+
+
+@torch.library.custom_op("motif::q_rope_inplace_fwd", mutates_args=("q",))
+def _q_rope_inplace_fwd(
+    q: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    nope_dim: int,
+    rope_dim: int,
+    seq_len: int,
+) -> None:
+    B_S, H, _ = q.shape
+    grid = lambda META: (B_S, triton.cdiv(H, META["BLOCK_H"]))
+    _q_rope_inplace_fwd_kernel[grid](
+        q, cos, sin,
+        nope_dim, rope_dim, H, seq_len,
+        q.stride(0), q.stride(1),
+    )
+
+
+class FusedQRoPEInplace(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, q, cos, sin, nope_dim, rope_dim, seq_len):
+        # q: [B, S, H, nope_dim + rope_dim]  mutated in-place on [..., nope_dim:]
+        B, S, H, D = q.shape
+        assert D == nope_dim + rope_dim
+        _q_rope_inplace_fwd(q.reshape(B * S, H, D), cos, sin, nope_dim, rope_dim, seq_len)
+        ctx.mark_dirty(q)
+        ctx.save_for_backward(cos, sin)
+        ctx.nope_dim = nope_dim
+        ctx.rope_dim = rope_dim
+        ctx.seq_len = seq_len
+        ctx.shape = (B, S, H, D)
+        return q
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        # grad_out: [B, S, H, nope_dim + rope_dim]  rope section in contiguous grad fmt
+        # Produce grad_in: same shape, rope section in interleaved grad fmt (matches Q input layout)
+        cos, sin = ctx.saved_tensors
+        B, S, H, D = ctx.shape
+        # Reuse existing Q backward kernel (copies nope grad + inverse-ropes pe section)
+        grad_in_3d = _q_rope_bwd(
+            grad_out.contiguous().reshape(B * S, H, D),
+            cos, sin, ctx.nope_dim, ctx.rope_dim, ctx.seq_len,
+        )
+        return grad_in_3d.view(B, S, H, D), None, None, None, None, None
+
+
+def fused_q_rope_inplace(q, freqs_cis, nope_dim, rope_dim):
+    """In-place fused Q RoPE. Modifies q[..., nope_dim:] from interleaved → contiguous format.
+
+    Replaces:
+        q_nope, q_pe = split(q, [nope_dim, rope_dim])
+        q_pe = fused_apply_rope(q_pe, freqs_cis)
+        q_total = cat([q_nope, q_pe])
+
+    Saves the cat copy (~415 µs/layer for Motif3) compared to out-of-place variants.
+
+    Args:
+        q: [B, S, H, nope_dim + rope_dim] from wq_b output. Will be mutated.
+        freqs_cis: [max_seq_len, rope_dim//2] complex64
+
+    Returns:
+        Same tensor `q`, now with rope section in contiguous format.
+    """
+    S = q.shape[1]
+    cos = freqs_cis[:S].real.contiguous()
+    sin = freqs_cis[:S].imag.contiguous()
+    return FusedQRoPEInplace.apply(q, cos, sin, nope_dim, rope_dim, S)
+
+
+def fused_kv_split_rope_cat(kv_latent, k_pe, k_nope_dim, v_dim, rope_dim):
+    """Fused KV split + k_pe expand + cat. No graph break.
+
+    Replaces:
+        k_nope, v = split(kv_latent, [k_nope_dim, v_dim])
+        k_full = cat([k_nope, k_pe.expand(-1,-1,H,-1)])
+
+    Args:
+        kv_latent: [B, S, H, k_nope_dim + v_dim]
+        k_pe: [B, S, 1, rope_dim] (already RoPE'd, contiguous format)
+
+    Returns:
+        key:   [B, S, H, k_nope_dim + rope_dim]
+        value: [B, S, H, v_dim]
+    """
+    return FusedKVRoPE.apply(kv_latent, k_pe, k_nope_dim, v_dim, rope_dim)