perf: remove autotune, hard-code per-kernel configs from live dump

Autotune isn't buying anything for these kernels: at runtime the shapes
are fixed by the model config (H ∈ {1, 80}, rope_dim=64), and live-dump
(TRITON_PRINT_AUTOTUNING=1) on motif3_seq B=8 S=4096 n_layers=4 showed
every kernel's best config sitting in a tight cluster that doesn't
change across reasonable workload sizes.

Replace @triton .autotune + config helpers with a small dict of static
(BLOCK_H, num_warps, num_stages) per kernel, passed at each kernel
launch. Values are the plurality-pick from the live dump:

_rope_fwd_kernel BLOCK_H=1, warps=1, stages=3 # k_pe, H=1
_rope_bwd_kernel BLOCK_H=2, warps=1, stages=1 # k_pe bwd
_kv_rope_fwd_kernel BLOCK_H=32, warps=8, stages=2 # H=80
_q_rope_inplace_fwd_kernel BLOCK_H=8, warps=1, stages=2 # H=80
_q_rope_bwd_kernel BLOCK_H=4, warps=2, stages=2 # H=80

Validation (same shape, autotune-full-sweep vs no-autotune, both B=8
n_layers=4): step 10 loss matches (0.0001), per-kernel time within
±15% (normalized, dominated by run-to-run noise). Warm-up cost drops
by ~213 s (full 640-trial sweep → 5 single-config compiles).

If model config changes materially (new H, different D), rerun the
dump and update the constants. The key=(rope_dim, head_num) invariant
that autotune used to key on is now enforced implicitly: any caller
that exercises a meaningfully different shape must update the configs.

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

torch-ext/activation/fused_rope.py

@@ -13,49 +13,32 @@ import triton
 import triton.language as tl
 
 # ---------------------------------------------------------------------------
-# Autotune search space — hand-picked from a full-sweep dump on motif3_seq
-# (B=4, S=4096, H=80, rope_dim=64, B200) with TRITON_PRINT_AUTOTUNING=1.
-# Best configs clustered around BLOCK_H∈{2,4,8}, warps∈{2,4}, stages=2 for
-# every RoPE kernel; keeping a handful of neighbours for shape robustness.
-# Triton caches per-(key) best config — one-time compile cost per distinct shape.
-# autotune key is (rope_dim, head_num) only; stride is not included because all
-# kernels assume last-dim unit stride + head-major layout (enforced by the
-# `stride(-1) == 1` assert at each wrapper entry), so layout never branches.
+# Static configs per kernel — no autotune.
+#
+# Picked from a full-sweep dump (TRITON_PRINT_AUTOTUNING=1) on motif3_seq,
+# B=8, S=4096, n_layers=4, 8×B200, EP=8 (plurality across 8 ranks; see
+# PR description for details). These are the shapes the kernels see at
+# the current prod workload; hard-coding avoids autotune warm-up cost and
+# tie-break jitter between near-equivalent configs.
+#
+# If model config changes materially (different H, D, batch-size/seq-len
+# regime), re-run the dump and update these. Autotune key used to be
+# (rope_dim, head_num); same shape invariants still apply at runtime.
 # ---------------------------------------------------------------------------
 
-def _rope_autotune_configs_fwd():
-    # Covers best configs observed across the 3 fwd kernels:
-    #   _rope_fwd (k_pe, H=1 tail) → BLOCK_H=2, warps=2
-    #   _q_rope_inplace_fwd       → BLOCK_H=8, warps=2
-    #   _kv_rope_fwd              → BLOCK_H=8, warps=4
-    return [
-        triton.Config({"BLOCK_H": 2}, num_warps=2, num_stages=2),
-        triton.Config({"BLOCK_H": 4}, num_warps=2, num_stages=2),
-        triton.Config({"BLOCK_H": 8}, num_warps=2, num_stages=2),
-        triton.Config({"BLOCK_H": 8}, num_warps=4, num_stages=2),
-    ]
-
-
-def _rope_autotune_configs_bwd():
-    # Covers best configs for both bwd kernels:
-    #   _q_rope_bwd   → BLOCK_H=4, warps=2
-    #   _rope_bwd     → BLOCK_H=2, warps∈{1,2} (flaky; tiebreak jitter on 88µs
-    #                   total budget — any of these is within noise)
-    return [
-        triton.Config({"BLOCK_H": 2}, num_warps=2, num_stages=2),
-        triton.Config({"BLOCK_H": 4}, num_warps=2, num_stages=2),
-        triton.Config({"BLOCK_H": 8}, num_warps=2, num_stages=2),
-    ]
+# 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
+_CFG_KV_ROPE_FWD         = dict(BLOCK_H=32, num_warps=8, num_stages=2)  # kv fused,   H=80
+_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
 
 
 # ---------------------------------------------------------------------------
 # Triton kernels
 # ---------------------------------------------------------------------------
 
-@triton.autotune(
-    configs=_rope_autotune_configs_fwd(),
-    key=["rope_dim", "head_num"],
-)
 @triton.jit
 def _rope_fwd_kernel(
     X,  # [B*S, H, rope_dim] interleaved
@@ -102,10 +85,6 @@ def _rope_fwd_kernel(
     tl.store(O_ptr + imag_off, out_imag, mask=mask)
 
 
-@triton.autotune(
-    configs=_rope_autotune_configs_bwd(),
-    key=["rope_dim", "head_num"],
-)
 @triton.jit
 def _rope_bwd_kernel(
     DO,  # [B*S, H, rope_dim] contiguous
@@ -172,6 +151,7 @@ def _rope_fwd(
         D, H, seq_len,
         x.stride(0), x.stride(1),
         out.stride(0), out.stride(1),
+        **_CFG_ROPE_FWD,
     )
     return out
 
@@ -197,6 +177,7 @@ def _rope_bwd(
         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
 
@@ -254,10 +235,6 @@ def fused_apply_rope(x, freqs_cis):
 # Reference: Megatron rotary_fwd_kv_kernel
 # ---------------------------------------------------------------------------
 
-@triton.autotune(
-    configs=_rope_autotune_configs_fwd(),
-    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]
@@ -338,6 +315,7 @@ def _kv_rope_fwd(
         k_pe.stride(0),
         key.stride(0), key.stride(1),
         value.stride(0), value.stride(1),
+        **_CFG_KV_ROPE_FWD,
     )
     return key, value
 
@@ -394,10 +372,6 @@ class FusedKVRoPE(torch.autograd.Function):
 # gradient is copied through unchanged.
 # ---------------------------------------------------------------------------
 
-@triton.autotune(
-    configs=_rope_autotune_configs_bwd(),
-    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)
@@ -471,6 +445,7 @@ def _q_rope_bwd(
         nope_dim, rope_dim, H, seq_len,
         grad_out.stride(0), grad_out.stride(1),
         grad_in.stride(0), grad_in.stride(1),
+        **_CFG_Q_ROPE_BWD,
     )
     return grad_in
 
@@ -488,10 +463,6 @@ def _q_rope_bwd_fake(grad_out, cos, sin, nope_dim, rope_dim, seq_len):
 # Backward: out-of-place via existing _q_rope_bwd_kernel (nope grad copy + inverse rope)
 # ---------------------------------------------------------------------------
 
-@triton.autotune(
-    configs=_rope_autotune_configs_fwd(),
-    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:]
@@ -555,6 +526,7 @@ def _q_rope_inplace_fwd(
         q, cos, sin,
         nope_dim, rope_dim, H, seq_len,
         q.stride(0), q.stride(1),
+        **_CFG_Q_ROPE_INPLACE_FWD,
     )