cleanup: drop k_pe RoPE custom kernel (caller uses PyTorch native)

Remove _rope_fwd_kernel, _rope_bwd_kernel, their custom_op wrappers,
FusedRoPE autograd class, and the fused_apply_rope user-facing wrapper
(plus the activation.__init__ export).

Rationale: k_pe in MLA is head-shared (H=1), so the kernel handles a
trivially small problem (B*S*1*rope_dim ≈ 8 MB at B=8, S=4096).
At that size the kernel is launch-bound, peak HBM utilization sits
around 5%, and a standalone Triton launch costs more than letting
inductor inline the work into the surrounding graph. The caller now
uses the same view_as_complex + reorder_headdim_elements_rope path
as the upstream PyTorch implementation.

Verified at B=8, S=4096, n_layers=4 on B200: step-10 loss matches
(0.0001), aggregate RoPE-path savings vs vanilla unchanged at -68%.

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

torch-ext/activation/__init__.py

@@ -3,7 +3,6 @@ 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,
 )
@@ -54,7 +53,6 @@ __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",

torch-ext/activation/fused_rope.py

@@ -27,207 +27,12 @@ import triton.language as tl
 # ---------------------------------------------------------------------------
 
 # BLOCK_H, num_warps, num_stages per kernel. H values denote head_num
-# (head-shared k_pe is H=1; per-head q/kv is H=80 in motif3_seq).
-_CFG_ROPE_FWD            = dict(BLOCK_H=1,  num_warps=1, num_stages=3)  # k_pe fwd,   H=1
-_CFG_ROPE_BWD            = dict(BLOCK_H=2,  num_warps=1, num_stages=1)  # k_pe bwd,   H=1
+# (q is per-head H=80, kv fused expands head-shared k_pe to H_kv=16 in motif3_seq).
 _CFG_KV_ROPE_FWD         = dict(BLOCK_H=32, num_warps=8, num_stages=2)  # kv fused,   H=80
+_CFG_KV_ROPE_BWD         = dict(BLOCK_H=16, num_warps=4, num_stages=2)  # kv bwd; BLOCK_H must cover H_kv=16 (single program per token)
 _CFG_Q_ROPE_INPLACE_FWD  = dict(BLOCK_H=8,  num_warps=1, num_stages=2)  # q in-place, H=80
 _CFG_Q_ROPE_BWD          = dict(BLOCK_H=4,  num_warps=2, num_stages=2)  # q bwd,      H=80
-_CFG_KV_ROPE_BWD         = dict(BLOCK_H=16, num_warps=4, num_stages=2)  # kv bwd; BLOCK_H must cover H_kv=16 (single program per token)
-
 
-# ---------------------------------------------------------------------------
-# Triton kernels
-# ---------------------------------------------------------------------------
-
-@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.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:
-    assert x.stride(-1) == 1, "fused_rope kernel requires last-dim unit stride"
-    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),
-        **_CFG_ROPE_FWD,
-    )
-    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:
-    assert grad_out.stride(-1) == 1, "fused_rope kernel requires last-dim unit stride"
-    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),
-        **_CFG_ROPE_BWD,
-    )
-    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)
 
 
 # ---------------------------------------------------------------------------