--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
@@ -0,0 +1,601 @@
+"""Triton kernel for multi-axis RoPE pre-attention prep.
+
+Fuses gather + per-axis mix + rotate-half + transpose into one kernel.
+Streams q/k once and emits the rotated (B, H, T, D) layout directly.
+"""
+import torch
+import torch.nn as nn
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _mrope_kernel(
+    q_ptr, k_ptr,
+    cos_cache, sin_cache,
+    pos_t, pos_h, pos_w,
+    q_rot, k_rot,
+    B: tl.int64, T: tl.int64, Hq: tl.int64, Hkv: tl.int64, D: tl.int64,
+    s0: tl.int64, s1: tl.int64, s2: tl.int64,
+    stride_qb, stride_qt, stride_qh, stride_qd,
+    stride_kb, stride_kt, stride_kh, stride_kd,
+    stride_pb, stride_pt,
+    stride_cocos, stride_cod,
+    stride_qrb, stride_qrr, stride_qrh, stride_qrd,
+    stride_krb, stride_krr, stride_krh, stride_krd,
+    BLOCK_TD: tl.constexpr,
+):
+    """Single kernel that processes one tensor (q or k).
+    
+    Each 1D program over (B*H*R) space. Within each program, a 2D grid
+    of threads handles tiles of (T, D).
+    
+    R = ceil(T / BLOCK_TD)  -- T tiles
+    """
+    # This program handles a specific (b, h, r) triple
+    # where r is the T-tile index
+    pid = tl.program_id(0)
+    
+    # Decompose pid into (b, h, r)
+    R = tl.cdiv(T, BLOCK_TD)
+    H = Hq  # set per call
+    r = pid % R
+    bh = pid // R
+    b = bh // H
+    h = bh % H
+    
+    if b >= B:
+        return
+    
+    # T offset and stride for output layout (B, H, T, D)
+    t_offset = r * BLOCK_TD
+    t_idx = t_offset + tl.arange(0, BLOCK_TD)
+    t_mask = t_idx < T
+    
+    # D indices - all threads in a warp handle different D elements
+    # We use a simple mapping: each thread handles one TD pair
+    d_idx = tl.arange(0, D)
+    
+    # We need threads = BLOCK_TD * D to cover the full tile, but that's too large.
+    # Instead, we'll process D serially in chunks.
+    # Actually let's use BLOCK_D as well.
+    # For now, let's just handle D elements sequentially per thread group.
+    pass  # placeholder
+
+
+@triton.jit
+def _mrope_fused(
+    q_ptr, k_ptr,
+    cos_cache_ptr, sin_cache_ptr,
+    pos_t_ptr, pos_h_ptr, pos_w_ptr,
+    q_rot_ptr, k_rot_ptr,
+    B: tl.int64, T: tl.int64, Hq: tl.int64, Hkv: tl.int64, D: tl.int64,
+    s0: tl.int64, s1: tl.int64, s2: tl.int64,
+    stride_qb, stride_qt, stride_qh, stride_qd,
+    stride_kb, stride_kt, stride_kh, stride_kd,
+    stride_pb, stride_pt,
+    stride_ccd,  # cos/sin cache: (max_pos, D)
+    stride_qrb, stride_qrr, stride_qrh, stride_qrd,
+    stride_krb, stride_krr, stride_krh, stride_krd,
+    BLOCK_DT: tl.constexpr,
+):
+    """
+    2D grid: (B*Hq*T//BLOCK_DT, D). Each program handles D elements for one (b,h,t) group.
+    
+    Actually we'll do TD tiling.
+    """
+    pass  # split below
+
+
+def _apply_mrope_kernel(q, k, cos_cache, sin_cache, pos_t, pos_h, pos_w, 
+                         mrope_section, D):
+    """Launch the fused mRoPE + transpose kernel."""
+    B, T, Hq, _ = q.shape
+    _, _, Hkv, _ = k.shape
+    s0, s1, s2 = mrope_section
+    
+    # Output buffers in (B, H, T, D) layout
+    q_rot = torch.empty(B, Hq, T, D, dtype=q.dtype, device=q.device)
+    k_rot = torch.empty(B, Hkv, T, D, dtype=k.dtype, device=k.device)
+    
+    # Process q and k separately but with same cos/sin
+    # Grid: each program handles (B*T*Hq) * D elements
+    # We'll use a simpler approach: (B * H) programs, each handles T * D
+    
+    BLOCK_D = 64
+    BLOCK_T = 64
+    
+    total_elements = B * T * (Hq + Hkv) * D
+    num_warps = 4
+    
+    # Launch for q
+    if q.numel() > 0:
+        _mrope_single_kernel[(B * Hq * triton.cdiv(T, BLOCK_T),)](
+            q, cos_cache, sin_cache, pos_t, pos_h, pos_w,
+            q_rot,
+            B=B, T=T, H=Hq, D=D,
+            s0=s0, s1=s1, s2=s2,
+            BLOCK_T=BLOCK_T,
+            num_warps=num_warps,
+        )
+    
+    # Launch for k 
+    if k.numel() > 0:
+        _mrope_single_kernel[(B * Hkv * triton.cdiv(T, BLOCK_T),)](
+            k, cos_cache, sin_cache, pos_t, pos_h, pos_w,
+            k_rot,
+            B=B, T=T, H=Hkv, D=D,
+            s0=s0, s1=s1, s2=s2,
+            BLOCK_T=BLOCK_T,
+            num_warps=num_warps,
+        )
+    
+    return q_rot, k_rot
+
+
+@triton.jit
+def _mrope_single_kernel(
+    x_ptr,        # (B, T, H, D) input
+    cos_cache_ptr,  # (max_pos, D)
+    sin_cache_ptr,
+    pos_t_ptr,    # (B, T)
+    pos_h_ptr,
+    pos_w_ptr,
+    out_ptr,      # (B, H, T, D) output
+    B: tl.int64, T: tl.int64, H: tl.int64, D: tl.int64,
+    s0: tl.int64, s1: tl.int64, s2: tl.int64,
+    stride_xb, stride_xt, stride_xh, stride_xd,
+    stride_outb, stride_outh, stride_outt, stride_outd,
+    stride_pb, stride_pt,
+    stride_cachepos, stride_cached,  # cos/sin cache strides
+    BLOCK_T: tl.constexpr,
+):
+    """
+    Each program handles one (b, h, t_tile) group.
+    Within the program, we process all D elements using thread mapping.
+    """
+    pid = tl.program_id(0)
+    R = tl.cdiv(T, BLOCK_T)
+    h_total = H
+    bh = pid // R
+    r = pid % R
+    b = bh // h_total
+    h = bh % h_total
+    
+    if b >= B:
+        return
+    
+    # T indices for this tile
+    t_off = r * BLOCK_T
+    t_idx = t_off + tl.arange(0, BLOCK_T)
+    t_mask = t_idx < T
+    
+    # D indices - we process ALL D across the program's threads
+    # With BLOCK_T=64 and num_warps=4 (128 threads), each thread handles D/128 elements
+    # For D=128, each thread handles 1 element. For D=64, we can have fewer threads.
+    stride = tl.num_programs(0)  # not needed
+    pid_d = tl.program_id(0) % D  # won't work - this is 1D grid
+    
+    # Better approach: make grid 2D or use a different strategy
+    # Let's use a 1D grid over (b, h, t_tile, d_chunk)
+    pass
+
+
+# Let me rewrite with proper tiling strategy
+@triton.jit
+def _mrope_kernel_v2(
+    x_ptr,
+    co

... (16484 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
@@ -11,545 +11,185 @@
 
 @triton.jit
 def _mrope_kernel(
-    q_ptr, k_ptr,
-    cos_cache, sin_cache,
-    pos_t, pos_h, pos_w,
-    q_rot, k_rot,
-    B: tl.int64, T: tl.int64, Hq: tl.int64, Hkv: tl.int64, D: tl.int64,
-    s0: tl.int64, s1: tl.int64, s2: tl.int64,
-    stride_qb, stride_qt, stride_qh, stride_qd,
-    stride_kb, stride_kt, stride_kh, stride_kd,
-    stride_pb, stride_pt,
-    stride_cocos, stride_cod,
-    stride_qrb, stride_qrr, stride_qrh, stride_qrd,
-    stride_krb, stride_krr, stride_krh, stride_krd,
-    BLOCK_TD: tl.constexpr,
-):
-    """Single kernel that processes one tensor (q or k).
-    
-    Each 1D program over (B*H*R) space. Within each program, a 2D grid
-    of threads handles tiles of (T, D).
-    
-    R = ceil(T / BLOCK_TD)  -- T tiles
-    """
-    # This program handles a specific (b, h, r) triple
-    # where r is the T-tile index
-    pid = tl.program_id(0)
-    
-    # Decompose pid into (b, h, r)
-    R = tl.cdiv(T, BLOCK_TD)
-    H = Hq  # set per call
-    r = pid % R
-    bh = pid // R
-    b = bh // H
-    h = bh % H
-    
-    if b >= B:
-        return
-    
-    # T offset and stride for output layout (B, H, T, D)
-    t_offset = r * BLOCK_TD
-    t_idx = t_offset + tl.arange(0, BLOCK_TD)
-    t_mask = t_idx < T
-    
-    # D indices - all threads in a warp handle different D elements
-    # We use a simple mapping: each thread handles one TD pair
-    d_idx = tl.arange(0, D)
-    
-    # We need threads = BLOCK_TD * D to cover the full tile, but that's too large.
-    # Instead, we'll process D serially in chunks.
-    # Actually let's use BLOCK_D as well.
-    # For now, let's just handle D elements sequentially per thread group.
-    pass  # placeholder
-
-
-@triton.jit
-def _mrope_fused(
-    q_ptr, k_ptr,
-    cos_cache_ptr, sin_cache_ptr,
-    pos_t_ptr, pos_h_ptr, pos_w_ptr,
-    q_rot_ptr, k_rot_ptr,
-    B: tl.int64, T: tl.int64, Hq: tl.int64, Hkv: tl.int64, D: tl.int64,
-    s0: tl.int64, s1: tl.int64, s2: tl.int64,
-    stride_qb, stride_qt, stride_qh, stride_qd,
-    stride_kb, stride_kt, stride_kh, stride_kd,
-    stride_pb, stride_pt,
-    stride_ccd,  # cos/sin cache: (max_pos, D)
-    stride_qrb, stride_qrr, stride_qrh, stride_qrd,
-    stride_krb, stride_krr, stride_krh, stride_krd,
-    BLOCK_DT: tl.constexpr,
-):
-    """
-    2D grid: (B*Hq*T//BLOCK_DT, D). Each program handles D elements for one (b,h,t) group.
-    
-    Actually we'll do TD tiling.
-    """
-    pass  # split below
-
-
-def _apply_mrope_kernel(q, k, cos_cache, sin_cache, pos_t, pos_h, pos_w, 
-                         mrope_section, D):
-    """Launch the fused mRoPE + transpose kernel."""
-    B, T, Hq, _ = q.shape
-    _, _, Hkv, _ = k.shape
-    s0, s1, s2 = mrope_section
-    
-    # Output buffers in (B, H, T, D) layout
-    q_rot = torch.empty(B, Hq, T, D, dtype=q.dtype, device=q.device)
-    k_rot = torch.empty(B, Hkv, T, D, dtype=k.dtype, device=k.device)
-    
-    # Process q and k separately but with same cos/sin
-    # Grid: each program handles (B*T*Hq) * D elements
-    # We'll use a simpler approach: (B * H) programs, each handles T * D
-    
-    BLOCK_D = 64
-    BLOCK_T = 64
-    
-    total_elements = B * T * (Hq + Hkv) * D
-    num_warps = 4
-    
-    # Launch for q
-    if q.numel() > 0:
-        _mrope_single_kernel[(B * Hq * triton.cdiv(T, BLOCK_T),)](
-            q, cos_cache, sin_cache, pos_t, pos_h, pos_w,
-            q_rot,
-            B=B, T=T, H=Hq, D=D,
-            s0=s0, s1=s1, s2=s2,
-            BLOCK_T=BLOCK_T,
-            num_warps=num_warps,
-        )
-    
-    # Launch for k 
-    if k.numel() > 0:
-        _mrope_single_kernel[(B * Hkv * triton.cdiv(T, BLOCK_T),)](
-            k, cos_cache, sin_cache, pos_t, pos_h, pos_w,
-            k_rot,
-            B=B, T=T, H=Hkv, D=D,
-            s0=s0, s1=s1, s2=s2,
-            BLOCK_T=BLOCK_T,
-            num_warps=num_warps,
-        )
-    
-    return q_rot, k_rot
-
-
-@triton.jit
-def _mrope_single_kernel(
-    x_ptr,        # (B, T, H, D) input
-    cos_cache_ptr,  # (max_pos, D)
+    x_ptr,
+    cos_cache_ptr,
     sin_cache_ptr,
-    pos_t_ptr,    # (B, T)
+    pos_t_ptr,
     pos_h_ptr,
     pos_w_ptr,
-    out_ptr,      # (B, H, T, D) output
-    B: tl.int64, T: tl.int64, H: tl.int64, D: tl.int64,
-    s0: tl.int64, s1: tl.int64, s2: tl.int64,
-    stride_xb, stride_xt, stride_xh, stride_xd,
-    stride_outb, stride_outh, stride_outt, stride_outd,
-    stride_pb, stride_pt,
-    stride_cachepos, stride_cached,  # cos/sin cache strides
-    BLOCK_T: tl.constexpr,
-):
-    """
-    Each program handles one (b, h, t_tile) group.
-    Within the program, we process all D elements using thread mapping.
-    """
-    pid = tl.program_id(0)
-    R = tl.cdiv(T, BLOCK_T)
-    h_total = H
-    bh = pid // R
-    r = pid % R
-    b = bh // h_total
-    h = bh % h_total
-    
-    if b >= B:
-        return
-    
-    # T indices for this tile
-    t_off = r * BLOCK_T
-    t_idx = t_off + tl.arange(0, BLOCK_T)
-    t_mask = t_idx < T
-    
-    # D indices - we process ALL D across the program's threads
-    # With BLOCK_T=64 and num_warps=4 (128 threads), each thread handles D/128 elements
-    # For D=128, each thread handles 1 element. For D=64, we can have fewer threads.
-    stride = tl.num_programs(0)  # not needed
-    pid_d = tl.program_id(0) % D  # won't work - this is 1D grid
-    
-    # Better approach: make grid 2D or use a different strategy
-    # Let's use a 1D grid over (b, h, t_tile, d_chunk)
-    pass
-
-
-# Let me rewrite with proper tiling strategy
-@triton.jit
-def _mrope_kernel_v2(
-    x_ptr,
-    cos_cache_ptr, sin_cache_ptr,
-    pos_t_ptr, pos_h_ptr, pos_w_ptr,
     out_ptr,
-    B: tl.int64, T: tl.int64, H: tl.int64, D: tl.int64,
-    s0: tl.int64, s1: tl.int

... (20308 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
@@ -2,6 +2,10 @@
 
 Fuses gather + per-axis mix + rotate-half + transpose into one kernel.
 Streams q/k once and emits the rotated (B, H, T, D) layout directly.
+
+Optimization: each d maps to exactly one axis (t/h/w) based on mrope_section.
+Instead of loading cos/sin from all 3 axes, we select the position first
+and load cos/sin from the chosen axis only — 2 cache loads per (t,d) instead of 6.
 """
 import torch
 import torch.nn as nn
@@ -42,8 +46,10 @@
 ):
     """
     Grid: (B * H * ceil(T/BLOCK_T), ceil(D/BLOCK_D))
-    Each program handles a (BLOCK_T, BLOCK_D) tile in the (T, D) plane
-    for a single (b, h) pair.
+    Each program handles a (BLOCK_T, BLOCK_D) tile for one (b, h) pair.
+
+    Key optimization: each d maps to exactly one axis based on mrope_section.
+    Select the position first, then load cos/sin once per (t,d).
     """
     pid = tl.program_id(0)
     pid_d = tl.program_id(1)
@@ -70,68 +76,51 @@
 
     half = D // 2
 
+    # --- Select position index based on which axis each d belongs to ---
+    # d_first = d if d < half else d - half  (axis is determined by first-half index)
+    d_first = d_off + tl.arange(0, BLOCK_D)
+    d_first = tl.where(d_first < half, d_first, d_first - half)
+
+    axis_t = d_first < s0
+    axis_h = (d_first >= s0) & (d_first < s0 + s1)
+    # axis_w = not axis_t and not axis_h (implied)
+
+    pos_base = b * stride_pb
+    pt_vals = tl.load(pos_t_ptr + pos_base + t_rng[:, None] * stride_pt,
+                      mask=t_mask[:, None], other=0)
+    ph_vals = tl.load(pos_h_ptr + pos_base + t_rng[:, None] * stride_pt,
+                      mask=t_mask[:, None], other=0)
+    pw_vals = tl.load(pos_w_ptr + pos_base + t_rng[:, None] * stride_pt,
+                      mask=t_mask[:, None], other=0)
+
+    # Select position per d column
+    pos_mixed = tl.where(axis_t[None, :], pt_vals,
+                 tl.where(axis_h[None, :], ph_vals, pw_vals))
+
+    # --- Load cos/sin once from the selected position ---
+    cos_off = pos_mixed * stride_ccb + d_rng[None, :] * stride_ccd
+    sin_off = cos_off  # same offset, different base pointer
+
+    cos_tile = tl.load(cos_cache_ptr + cos_off, mask=td_mask, other=0.0)
+    sin_tile = tl.load(sin_cache_ptr + sin_off, mask=td_mask, other=0.0)
+
     # --- Load x[b, t, h, d] ---
     x_base_off = b * stride_xb + h * stride_xh
     x_offsets = x_base_off + t_rng[:, None] * stride_xt + d_rng[None, :] * stride_xd
     x_tile = tl.load(x_ptr + x_offsets, mask=td_mask, other=0.0)
 
-    # --- Load position indices (each T loads its position) ---
-    pos_base = b * stride_pb
-    pt_off = pos_base + t_rng * stride_pt
-    ph_off = pos_base + t_rng * stride_pt
-    pw_off = pos_base + t_rng * stride_pt
-
-    pos_t_vals = tl.load(pos_t_ptr + pt_off, mask=t_mask, other=0)
-    pos_h_vals = tl.load(pos_h_ptr + ph_off, mask=t_mask, other=0)
-    pos_w_vals = tl.load(pos_w_ptr + pw_off, mask=t_mask, other=0)
-
-    # --- Load cos for each axis ---
-    # cos_cache is (max_pos, D), contiguous. stride = (D, 1)
-    # offset for (pos, d) = pos * D + d
-    # pos_vals[:, None] * stride_ccb + d_rng[None, :] * stride_ccd
-    # shape: (BLOCK_T, BLOCK_D)
-    cos_t_off = pos_t_vals[:, None] * stride_ccb + d_rng[None, :] * stride_ccd
-    cos_h_off = pos_h_vals[:, None] * stride_ccb + d_rng[None, :] * stride_ccd
-    cos_w_off = pos_w_vals[:, None] * stride_ccb + d_rng[None, :] * stride_ccd
-
-    cos_t_tile = tl.load(cos_cache_ptr + cos_t_off, mask=td_mask, other=0.0)
-    cos_h_tile = tl.load(cos_cache_ptr + cos_h_off, mask=td_mask, other=0.0)
-    cos_w_tile = tl.load(cos_cache_ptr + cos_w_off, mask=td_mask, other=0.0)
-
-    sin_t_tile = tl.load(sin_cache_ptr + cos_t_off, mask=td_mask, other=0.0)
-    sin_h_tile = tl.load(sin_cache_ptr + cos_h_off, mask=td_mask, other=0.0)
-    sin_w_tile = tl.load(sin_cache_ptr + cos_w_off, mask=td_mask, other=0.0)
-
-    # --- Mix axes based on mrope_section ---
-    # d in [0, s0) -> temporal, [s0, s0+s1) -> height, [s0+s1, D//2) -> width
-    # d in [D//2, D) -> same scheme but with d and d-D//2 for cache column
-    # Key: cache column is ALWAYS d, not d_first.
-    # The axis determines WHICH position Row to use (pos_t vs pos_h vs pos_w).
-    # d_first determines axis, but actual cache col = d (mirror has same frequency)
-    d_first = tl.where(d_rng < half, d_rng, d_rng - half)
-    axis_mask_t = d_first < s0
-    axis_mask_h = (d_first >= s0) & (d_first < s0 + s1)
-
-    cos_mixed = tl.where(axis_mask_t, cos_t_tile,
-                 tl.where(axis_mask_h, cos_h_tile, cos_w_tile))
-    sin_mixed = tl.where(axis_mask_t, sin_t_tile,
-                 tl.where(axis_mask_h, sin_h_tile, sin_w_tile))
-
-    # --- Rotate half ---
-    # rotate_half(x)[d] = -x[d+half] if d < half else x[d-half]
-    d_mir = tl.where(d_rng < half, d_rng + half, d_rng - half)
-
-    # x_offsets for mirrored d: x[b, t, h, d_mir]
-    # d_mir might be outside [d_off, d_off+BLOCK_D), so this is a
-    # potentially non-contiguous load. Triton handles it as scalar addressed load.
+    # --- Load x at mirrored position for rotate_half ---
+    # d_mir = d + half if d < half else d - half
+    d_mir = d_rng + half
+    d_mir = tl.where(d_rng < half, d_mir, d_rng - half)
     x_mir_offsets = x_base_off + t_rng[:, None] * stride_xt + d_mir[None, :] * stride_xd
     x_mir_tile = tl.load(x_ptr + x_mir_offsets, mask=td_mask, other=0.0)
 
-    # Sign: negate if d < half (the d that's being moved from second half)
-    x_rot = tl.where(d_rng < half, -x_mir_tile, x_mir_tile)
+    # Sign: rotate_half[x at d] = -x[d+half] if d < half, +x[d-half] if d >= half
+    x_rot = tl.where(d_rng[None, :] < half, -x_mir_tile, x_mir_tile)
 
     # --- Apply RoPE ---
-    out_tile = x_tile * cos_mixed + x_rot * sin_mixed
+    out_tile = x_tile * cos_t

... (897 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
@@ -1,11 +1,8 @@
-"""Triton kernel for multi-axis RoPE pre-attention prep.
-
-Fuses gather + per-axis mix + rotate-half + transpose into one kernel.
-Streams q/k once and emits the rotated (B, H, T, D) layout directly.
-
-Optimization: each d maps to exactly one axis (t/h/w) based on mrope_section.
-Instead of loading cos/sin from all 3 axes, we select the position first
-and load cos/sin from the chosen axis only — 2 cache loads per (t,d) instead of 6.
+"""Optimized Triton kernel for multi-axis RoPE pre-attention prep.
+
+Fuses gather + per-axis mix + rotate-half + transpose for both q and k
+into a single kernel launch. Each thread block handles one (b, h, t_tile, d_tile),
+loading cos/sin once per (t,d) and applying rotation to both q and k heads.
 """
 import torch
 import torch.nn as nn
@@ -14,42 +11,37 @@
 
 
 @triton.jit
-def _mrope_kernel(
-    x_ptr,
+def _mrope_fused_kernel(
+    q_ptr,
+    k_ptr,
     cos_cache_ptr,
     sin_cache_ptr,
     pos_t_ptr,
     pos_h_ptr,
     pos_w_ptr,
-    out_ptr,
+    q_out_ptr,
+    k_out_ptr,
     B: tl.int64,
     T: tl.int64,
-    H: tl.int64,
+    Hq: tl.int64,
+    Hkv: tl.int64,
     D: tl.int64,
     s0: tl.int64,
     s1: tl.int64,
-    s2: tl.int64,
-    stride_xb,
-    stride_xt,
-    stride_xh,
-    stride_xd,
-    stride_pb,
-    stride_pt,
-    stride_ccb,
-    stride_ccd,
-    stride_outb,
-    stride_outh,
-    stride_outt,
-    stride_outd,
+    stride_qb, stride_qt, stride_qh, stride_qd,
+    stride_kb, stride_kt, stride_kh, stride_kd,
+    stride_pb, stride_pt,
+    stride_ccb, stride_ccd,
+    stride_qob, stride_qoh, stride_qot, stride_qod,
+    stride_kob, stride_koh, stride_kot, stride_kod,
     BLOCK_T: tl.constexpr,
     BLOCK_D: tl.constexpr,
+    HAS_K: tl.constexpr,
 ):
     """
-    Grid: (B * H * ceil(T/BLOCK_T), ceil(D/BLOCK_D))
-    Each program handles a (BLOCK_T, BLOCK_D) tile for one (b, h) pair.
-
-    Key optimization: each d maps to exactly one axis based on mrope_section.
-    Select the position first, then load cos/sin once per (t,d).
+    Grid: (B * (Hq + Hkv) * ceil(T/BLOCK_T), ceil(D/BLOCK_D))
+    Each program handles one (b, h, t_tile, d_tile).
+    h_is_q indicates whether this is a q head or k head.
     """
     pid = tl.program_id(0)
     pid_d = tl.program_id(1)
@@ -57,16 +49,34 @@
     R = tl.cdiv(T, BLOCK_T)
     D_tiles = tl.cdiv(D, BLOCK_D)
 
+    total_h = Hq + Hkv if HAS_K else Hq
     bh = pid // R
     r = pid % R
-    b = bh // H
-    h = bh % H
+    b = bh // total_h
+    h = bh % total_h
 
     if b >= B or pid_d >= D_tiles:
         return
 
+    h_is_q = h < Hq
+    h_idx = h if h_is_q else h - Hq
+
+    # Select pointers
+    stride_b = tl.where(h_is_q, stride_qb, tl.where(HAS_K, stride_kb, 0))
+    stride_t = tl.where(h_is_q, stride_qt, tl.where(HAS_K, stride_kt, 0))
+    stride_h = tl.where(h_is_q, stride_qh, tl.where(HAS_K, stride_kh, 0))
+    stride_d = tl.where(h_is_q, stride_qd, tl.where(HAS_K, stride_kd, 0))
+    x_ptr = tl.where(h_is_q, q_ptr, tl.where(HAS_K, k_ptr, 0))
+    
+    stride_ob = tl.where(h_is_q, stride_qob, tl.where(HAS_K, stride_kob, 0))
+    stride_oh = tl.where(h_is_q, stride_qoh, tl.where(HAS_K, stride_koh, 0))
+    stride_ot = tl.where(h_is_q, stride_qot, tl.where(HAS_K, stride_kot, 0))
+    stride_od = tl.where(h_is_q, stride_qod, tl.where(HAS_K, stride_kod, 0))
+    out_ptr = tl.where(h_is_q, q_out_ptr, tl.where(HAS_K, k_out_ptr, 0))
+
     t_off = r * BLOCK_T
     d_off = pid_d * BLOCK_D
+    half = D // 2
 
     t_rng = t_off + tl.arange(0, BLOCK_T)
     d_rng = d_off + tl.arange(0, BLOCK_D)
@@ -74,57 +84,63 @@
     d_mask = d_rng < D
     td_mask = t_mask[:, None] & d_mask[None, :]
 
-    half = D // 2
-
-    # --- Select position index based on which axis each d belongs to ---
-    # d_first = d if d < half else d - half  (axis is determined by first-half index)
-    d_first = d_off + tl.arange(0, BLOCK_D)
-    d_first = tl.where(d_first < half, d_first, d_first - half)
-
+    # Load x[b, t, h_idx, d]
+    x_base = b * stride_b + h_idx * stride_h
+    x_offsets = x_base + t_rng[:, None] * stride_t + d_rng[None, :] * stride_d
+    x_tile = tl.load(x_ptr + x_offsets, mask=td_mask, other=0.0)
+
+    # Load x at mirrored position for rotate_half
+    d_mir = tl.where(d_rng < half, d_rng + half, d_rng - half)
+    x_mir_offsets = x_base + t_rng[:, None] * stride_t + d_mir[None, :] * stride_d
+    x_mir_tile = tl.load(x_ptr + x_mir_offsets, mask=td_mask, other=0.0)
+    x_rot = tl.where(d_rng[None, :] < half, -x_mir_tile, x_mir_tile)
+
+    # Load positions
+    pos_base = b * stride_pb
+    pos_offsets = pos_base + t_rng * stride_pt
+    pos_t_vals = tl.load(pos_t_ptr + pos_offsets, mask=t_mask, other=0)
+    pos_h_vals = tl.load(pos_h_ptr + pos_offsets, mask=t_mask, other=0)
+    pos_w_vals = tl.load(pos_w_ptr + pos_offsets, mask=t_mask, other=0)
+
+    # Axis selection: d_first determines axis
+    d_first = tl.where(d_rng < half, d_rng, d_rng - half)
     axis_t = d_first < s0
     axis_h = (d_first >= s0) & (d_first < s0 + s1)
-    # axis_w = not axis_t and not axis_h (implied)
-
-    pos_base = b * stride_pb
-    pt_vals = tl.load(pos_t_ptr + pos_base + t_rng[:, None] * stride_pt,
-                      mask=t_mask[:, None], other=0)
-    ph_vals = tl.load(pos_h_ptr + pos_base + t_rng[:, None] * stride_pt,
-                      mask=t_mask[:, None], other=0)
-    pw_vals = tl.load(pos_w_ptr + pos_base + t_rng[:, None] * stride_pt,
-                      mask=t_mask[:, None], other=0)
-
-    # Select position per d column
-    pos_mixed = tl.where(axis_t[None, :], pt_vals,
-                 tl.where(axis_h[None, :], ph_vals, pw_vals))
-
-    # --- Load cos/sin once from the selected position ---
+
+    # Broadcast positions to (BLOCK_T, BLOCK_D)
+    pos_t_wide = pos_t_vals[:,

... (4909 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
@@ -1,8 +1,7 @@
 """Optimized Triton kernel for multi-axis RoPE pre-attention prep.
 
-Fuses gather + per-axis mix + rotate-half + transpose for both q and k
-into a single kernel launch. Each thread block handles one (b, h, t_tile, d_tile),
-loading cos/sin once per (t,d) and applying rotation to both q and k heads.
+Strategy: one kernel per tensor (q, k). Each block handles (BLOCK_T x BLOCK_D)
+for one (b, h) pair. Cos/sin loaded once per (t,d) via axis-selected position.
 """
 import torch
 import torch.nn as nn
@@ -11,20 +10,19 @@
 
 
 @triton.jit
-def _mrope_fused_kernel(
+def _mrope_kernel(
     q_ptr,
     k_ptr,
-    cos_cache_ptr,
-    sin_cache_ptr,
-    pos_t_ptr,
-    pos_h_ptr,
-    pos_w_ptr,
-    q_out_ptr,
-    k_out_ptr,
+    cos_cache,
+    sin_cache,
+    pos_t,
+    pos_h,
+    pos_w,
+    q_out,
+    k_out,
     B: tl.int64,
     T: tl.int64,
     Hq: tl.int64,
-    Hkv: tl.int64,
     D: tl.int64,
     s0: tl.int64,
     s1: tl.int64,
@@ -32,47 +30,44 @@
     stride_kb, stride_kt, stride_kh, stride_kd,
     stride_pb, stride_pt,
     stride_ccb, stride_ccd,
-    stride_qob, stride_qoh, stride_qot, stride_qod,
-    stride_kob, stride_koh, stride_kot, stride_kod,
+    stride_qob, stride_qot, stride_qod,
+    stride_kob, stride_kot, stride_kod,
     BLOCK_T: tl.constexpr,
     BLOCK_D: tl.constexpr,
-    HAS_K: tl.constexpr,
 ):
     """
-    Grid: (B * (Hq + Hkv) * ceil(T/BLOCK_T), ceil(D/BLOCK_D))
-    Each program handles one (b, h, t_tile, d_tile).
-    h_is_q indicates whether this is a q head or k head.
+    Grid: (B * H * ceil(T/BLOCK_T), ceil(D/BLOCK_D))
+    where H = Hq for q pass, Hkv for k pass.
+    
+    Actually this is a truly fused kernel: grid_0 over (b, h_or_hkv, t_tile),
+    grid_1 over d_tile. Each program computes RoPE for both q and k if applicable.
+    
+    Wait, q has Hq heads, k has Hkv heads. Different H means we can't simply fuse (b,h).
+    Let's keep things simple: each grid_0 program handles (b, h, t_tile), 
+    and within the kernel we compute for the appropriate tensor.
+    
+    We'll just use separate launches with shared code via constexpr HAS_KV flag.
+    Actually simpler: just two kernel calls, same kernel, different H and pointers.
     """
-    pid = tl.program_id(0)
-    pid_d = tl.program_id(1)
+    pid_t = tl.program_id(0)  # combined (b * H + h) * R + r
+    pid_d = tl.program_id(1)  # D tile
 
     R = tl.cdiv(T, BLOCK_T)
     D_tiles = tl.cdiv(D, BLOCK_D)
 
-    total_h = Hq + Hkv if HAS_K else Hq
-    bh = pid // R
-    r = pid % R
-    b = bh // total_h
-    h = bh % total_h
+    # Total heads depends on which tensor we're processing.
+    # We pass which_h_total to decide decomposition.
+    # Simplify: assume grid_0 = (B * H_total * R,) where H_total comes from
+    # the first argument.
+    H_total = Hq  # set per-caller
+    
+    bh = pid_t // R
+    r = pid_t % R
+    b = bh // H_total
+    h = bh % H_total
 
     if b >= B or pid_d >= D_tiles:
         return
-
-    h_is_q = h < Hq
-    h_idx = h if h_is_q else h - Hq
-
-    # Select pointers
-    stride_b = tl.where(h_is_q, stride_qb, tl.where(HAS_K, stride_kb, 0))
-    stride_t = tl.where(h_is_q, stride_qt, tl.where(HAS_K, stride_kt, 0))
-    stride_h = tl.where(h_is_q, stride_qh, tl.where(HAS_K, stride_kh, 0))
-    stride_d = tl.where(h_is_q, stride_qd, tl.where(HAS_K, stride_kd, 0))
-    x_ptr = tl.where(h_is_q, q_ptr, tl.where(HAS_K, k_ptr, 0))
-    
-    stride_ob = tl.where(h_is_q, stride_qob, tl.where(HAS_K, stride_kob, 0))
-    stride_oh = tl.where(h_is_q, stride_qoh, tl.where(HAS_K, stride_koh, 0))
-    stride_ot = tl.where(h_is_q, stride_qot, tl.where(HAS_K, stride_kot, 0))
-    stride_od = tl.where(h_is_q, stride_qod, tl.where(HAS_K, stride_kod, 0))
-    out_ptr = tl.where(h_is_q, q_out_ptr, tl.where(HAS_K, k_out_ptr, 0))
 
     t_off = r * BLOCK_T
     d_off = pid_d * BLOCK_D
@@ -84,71 +79,168 @@
     d_mask = d_rng < D
     td_mask = t_mask[:, None] & d_mask[None, :]
 
-    # Load x[b, t, h_idx, d]
-    x_base = b * stride_b + h_idx * stride_h
-    x_offsets = x_base + t_rng[:, None] * stride_t + d_rng[None, :] * stride_d
-    x_tile = tl.load(x_ptr + x_offsets, mask=td_mask, other=0.0)
-
-    # Load x at mirrored position for rotate_half
+    # Load position indices (broadcast per t column)
+    pos_off = b * stride_pb + t_rng * stride_pt
+    pos_t_vals = tl.load(pos_t + pos_off, mask=t_mask, other=0)
+    pos_h_vals = tl.load(pos_h + pos_off, mask=t_mask, other=0)
+    pos_w_vals = tl.load(pos_w + pos_off, mask=t_mask, other=0)
+
+    # Axis mask for each d column (independent of t)
+    d_first = tl.where(d_rng < half, d_rng, d_rng - half)
+    # axis_t[d] = d_first < s0
+    # axis_h[d] = s0 <= d_first < s0+s1
+    # axis_w[d] = d_first >= s0+s1
+    axis_t = d_first < s0
+    axis_h = (d_first >= s0) & (d_first < s0 + s1)
+    # axis_w = ~axis_t & ~axis_h (implied)
+
+    # Select position: for each (t, d), which of pos_t/pos_h/pos_w to use
+    # axis_t/axis_h are (BLOCK_D,) vectors, extend to (BLOCK_T, BLOCK_D)
+    pos_sel = tl.where(
+        axis_t[None, :],
+        pos_t_vals[:, None],
+        tl.where(axis_h[None, :], pos_h_vals[:, None], pos_w_vals[:, None])
+    )
+
+    # Load cos/sin from cache using selected position
+    # cache is contiguous (max_pos, D): offset = pos * D + d
+    cos_sin_off = pos_sel * stride_ccb + d_rng[None, :] * stride_ccd
+    cos_tile = tl.load(cos_cache + cos_sin_off, mask=td_mask, other=0.0)
+    sin_tile = tl.load(sin_cache + cos_sin_off, mask=td_mask, other=0.0)
+
+    # --- Q RoPE ---
+    q_base = b * stride_qb + h * stride_qh
+    q_offsets = q_base + t_rng[:, None] * stride_qt + d_rng[None, :] * stride_qd
+    q_tile = tl.load(q_ptr + q_offsets, mask=td_mask, other=0.0)
+
+    # Mirror d for rotate_half
     d_mir = tl.w

... (8826 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
@@ -1,7 +1,10 @@
 """Optimized Triton kernel for multi-axis RoPE pre-attention prep.
 
-Strategy: one kernel per tensor (q, k). Each block handles (BLOCK_T x BLOCK_D)
-for one (b, h) pair. Cos/sin loaded once per (t,d) via axis-selected position.
+Key optimizations:
+1. Each d maps to exactly one axis - load cos/sin ONCE per (t,d)
+2. Large D tiles (128) for coalesced loads
+3. num_warps=8 for better memory throughput
+4. Use TRITON_MAX Sunder for pointer arithmetic
 """
 import torch
 import torch.nn as nn
@@ -9,131 +12,6 @@
 import triton.language as tl
 
 
-@triton.jit
-def _mrope_kernel(
-    q_ptr,
-    k_ptr,
-    cos_cache,
-    sin_cache,
-    pos_t,
-    pos_h,
-    pos_w,
-    q_out,
-    k_out,
-    B: tl.int64,
-    T: tl.int64,
-    Hq: tl.int64,
-    D: tl.int64,
-    s0: tl.int64,
-    s1: tl.int64,
-    stride_qb, stride_qt, stride_qh, stride_qd,
-    stride_kb, stride_kt, stride_kh, stride_kd,
-    stride_pb, stride_pt,
-    stride_ccb, stride_ccd,
-    stride_qob, stride_qot, stride_qod,
-    stride_kob, stride_kot, stride_kod,
-    BLOCK_T: tl.constexpr,
-    BLOCK_D: tl.constexpr,
-):
-    """
-    Grid: (B * H * ceil(T/BLOCK_T), ceil(D/BLOCK_D))
-    where H = Hq for q pass, Hkv for k pass.
-    
-    Actually this is a truly fused kernel: grid_0 over (b, h_or_hkv, t_tile),
-    grid_1 over d_tile. Each program computes RoPE for both q and k if applicable.
-    
-    Wait, q has Hq heads, k has Hkv heads. Different H means we can't simply fuse (b,h).
-    Let's keep things simple: each grid_0 program handles (b, h, t_tile), 
-    and within the kernel we compute for the appropriate tensor.
-    
-    We'll just use separate launches with shared code via constexpr HAS_KV flag.
-    Actually simpler: just two kernel calls, same kernel, different H and pointers.
-    """
-    pid_t = tl.program_id(0)  # combined (b * H + h) * R + r
-    pid_d = tl.program_id(1)  # D tile
-
-    R = tl.cdiv(T, BLOCK_T)
-    D_tiles = tl.cdiv(D, BLOCK_D)
-
-    # Total heads depends on which tensor we're processing.
-    # We pass which_h_total to decide decomposition.
-    # Simplify: assume grid_0 = (B * H_total * R,) where H_total comes from
-    # the first argument.
-    H_total = Hq  # set per-caller
-    
-    bh = pid_t // R
-    r = pid_t % R
-    b = bh // H_total
-    h = bh % H_total
-
-    if b >= B or pid_d >= D_tiles:
-        return
-
-    t_off = r * BLOCK_T
-    d_off = pid_d * BLOCK_D
-    half = D // 2
-
-    t_rng = t_off + tl.arange(0, BLOCK_T)
-    d_rng = d_off + tl.arange(0, BLOCK_D)
-    t_mask = t_rng < T
-    d_mask = d_rng < D
-    td_mask = t_mask[:, None] & d_mask[None, :]
-
-    # Load position indices (broadcast per t column)
-    pos_off = b * stride_pb + t_rng * stride_pt
-    pos_t_vals = tl.load(pos_t + pos_off, mask=t_mask, other=0)
-    pos_h_vals = tl.load(pos_h + pos_off, mask=t_mask, other=0)
-    pos_w_vals = tl.load(pos_w + pos_off, mask=t_mask, other=0)
-
-    # Axis mask for each d column (independent of t)
-    d_first = tl.where(d_rng < half, d_rng, d_rng - half)
-    # axis_t[d] = d_first < s0
-    # axis_h[d] = s0 <= d_first < s0+s1
-    # axis_w[d] = d_first >= s0+s1
-    axis_t = d_first < s0
-    axis_h = (d_first >= s0) & (d_first < s0 + s1)
-    # axis_w = ~axis_t & ~axis_h (implied)
-
-    # Select position: for each (t, d), which of pos_t/pos_h/pos_w to use
-    # axis_t/axis_h are (BLOCK_D,) vectors, extend to (BLOCK_T, BLOCK_D)
-    pos_sel = tl.where(
-        axis_t[None, :],
-        pos_t_vals[:, None],
-        tl.where(axis_h[None, :], pos_h_vals[:, None], pos_w_vals[:, None])
-    )
-
-    # Load cos/sin from cache using selected position
-    # cache is contiguous (max_pos, D): offset = pos * D + d
-    cos_sin_off = pos_sel * stride_ccb + d_rng[None, :] * stride_ccd
-    cos_tile = tl.load(cos_cache + cos_sin_off, mask=td_mask, other=0.0)
-    sin_tile = tl.load(sin_cache + cos_sin_off, mask=td_mask, other=0.0)
-
-    # --- Q RoPE ---
-    q_base = b * stride_qb + h * stride_qh
-    q_offsets = q_base + t_rng[:, None] * stride_qt + d_rng[None, :] * stride_qd
-    q_tile = tl.load(q_ptr + q_offsets, mask=td_mask, other=0.0)
-
-    # Mirror d for rotate_half
-    d_mir = tl.where(d_rng < half, d_rng + half, d_rng - half)
-    q_mir_offsets = q_base + t_rng[:, None] * stride_qt + d_mir[None, :] * stride_qd
-    q_mir = tl.load(q_ptr + q_mir_offsets, mask=td_mask, other=0.0)
-    q_rot_half = tl.where(d_rng[None, :] < half, -q_mir, q_mir)
-
-    q_result = q_tile * cos_tile + q_rot_half * sin_tile
-
-    # Store q output (B, H, T, D)
-    q_out_base = b * stride_qob
-    q_out_off = q_out_base + h * q.stride(1).to(tl.int64) * 0 + t_rng[:, None] * stride_qot + d_rng[None, :] * stride_qod
-    # Wait, q.stride(1) is T stride in OUTPUT layout (B, H, T, D). Output strides are different.
-    tl.store(q_out + q_out_base + h * q_out.stride(1) + t_rng[:, None] * stride_qot + d_rng[None, :] * stride_qod, 
-             q_result, mask=td_mask)
-
-    # --- K RoPE (if this kernel also handles k) ---
-    # We'll vectorize over both by making K a separate launch for now.
-    pass  # handled by separate kernel call
-
-
-# Better to have one simple kernel and call it twice. Let me simplify:
 @triton.jit
 def _rope_kernel(
     x_ptr,
@@ -164,78 +42,77 @@
     BLOCK_T: tl.constexpr,
     BLOCK_D: tl.constexpr,
 ):
-    """
-    Grid: (B * H * ceil(T/BLOCK_T), ceil(D/BLOCK_D))
-    """
+    """2D grid: (B*H*ceil(T/BLOCK_T), ceil(D/BLOCK_D))"""
     pid = tl.program_id(0)
     pid_d = tl.program_id(1)
 
     R = tl.cdiv(T, BLOCK_T)
     D_tiles = tl.cdiv(D, BLOCK_D)
-    total_bh = B * H
 
     bh = pid // R
     r = pid % R
     b = bh // H
     h = bh % H
 
-    if bh >= total_bh or pid_d >= D_tiles:
+    if b >= B or pid_d >= D_tiles:
         return
 
     t_off = r

... (5715 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
@@ -80,8 +80,8 @@
 
     # Axis determination per d (BLOCK_D,)
     d_in_first = tl.where(d_rng < half_d, d_rng, d_rng - half_d)
-    _bt = d_in_first < s0       # (BLOCK_D,)
-    _bh = d_in_first >= s0
+    _bt = d_in_first < s0
+    _bh = (d_in_first >= s0) & (d_in_first < s0 + s1)
     
     # pos mixing: (BLOCK_T, BLOCK_D)
     pos = tl.where(_bt[None, :],

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
@@ -1,10 +1,8 @@
 """Optimized Triton kernel for multi-axis RoPE pre-attention prep.
 
-Key optimizations:
-1. Each d maps to exactly one axis - load cos/sin ONCE per (t,d)
-2. Large D tiles (128) for coalesced loads
-3. num_warps=8 for better memory throughput
-4. Use TRITON_MAX Sunder for pointer arithmetic
+Strategy: each block handles one (b, h, t) — all of D in SRAM.
+Grid = B * H * T programs. Each program has 128 threads for D=128.
+Single cos/sin load per (t,d) via axis-selected position.
 """
 import torch
 import torch.nn as nn
@@ -39,80 +37,130 @@
     stride_outh,
     stride_outt,
     stride_outd,
-    BLOCK_T: tl.constexpr,
+    D_PADDED: tl.constexpr,
+):
+    """
+    Grid: (B * H * T,)
+    Each program loads the full D for one (b, h, t).
+    Thread idx = tid within block maps to D index.
+    D_PADDED should be a power of 2 >= D for efficient thread mapping.
+    """
+    pid = tl.program_id(0)
+    R = B * H * T
+    if pid >= R:
+        return
+
+    t = pid % T
+    bh = (pid // T) % (B * H)
+    b = bh // H
+    h = bh % H
+
+    tid = tl.program_id(0) % D_PADDED
+    d_rng = tid + tl.arange(0, D_PADDED)
+    # Actually in Triton, we use tl.arange for vector operations.
+    # The program handles all D.
+    # With 1D grid, tl.program_id(0) = pid. Each program = one (b,h,t).
+    # Threads handle D via tl.arange.
+    pass
+
+
+# Actually better to use proper Triton patterns. Let me do 1D grid with
+# td-range within each program. Triton handles this via tl.arange(0, N)
+# where N = threads per block / warps.
+
+@triton.jit
+def _rope_per_token(
+    x_ptr,
+    cos_cache,
+    sin_cache,
+    pos_t_ptr,
+    pos_h_ptr,
+    pos_w_ptr,
+    out_ptr,
+    B: tl.int64,
+    T: tl.int64,
+    H: tl.int64,
+    D: tl.int64,
+    s0: tl.int64,
+    s1: tl.int64,
+    stride_xb,
+    stride_xt,
+    stride_xh,
+    stride_xd,
+    stride_pb,
+    stride_pt,
+    stride_cachepos,
+    stride_cached,
+    stride_outb,
+    stride_outh,
+    stride_outt,
+    stride_outd,
     BLOCK_D: tl.constexpr,
 ):
-    """2D grid: (B*H*ceil(T/BLOCK_T), ceil(D/BLOCK_D))"""
+    """
+    Each program handles one (b, h, t) and processes D in tiles of BLOCK_D.
+    Total programs = B * H * T. But that's huge (1*32*4096 = 131072).
+    Each program loops over D tiles.
+    """
     pid = tl.program_id(0)
-    pid_d = tl.program_id(1)
-
-    R = tl.cdiv(T, BLOCK_T)
-    D_tiles = tl.cdiv(D, BLOCK_D)
-
-    bh = pid // R
-    r = pid % R
+
+    t = pid % T
+    bh = (pid // T) % (B * H)
     b = bh // H
     h = bh % H
 
-    if b >= B or pid_d >= D_tiles:
+    if pid >= B * H * T:
         return
 
-    t_off = r * BLOCK_T
-    d_off = pid_d * BLOCK_D
-    half_d = D // 2
-
-    # Thread-local indices
-    _range_t = tl.arange(0, BLOCK_T)
-    _range_d = tl.arange(0, BLOCK_D)
-    
-    t_rng = t_off + _range_t
-    d_rng = d_off + _range_d
-    
-    t_mask = t_rng < T
-    d_mask = d_rng < D
-    td_mask = t_mask[:, None] & d_mask[None, :]
-
-    # --- Position gathering ---
-    pos_off = b * stride_pb + t_rng * stride_pt
-    p_t = tl.load(pos_t_ptr + pos_off, mask=t_mask, other=0)
-    p_h = tl.load(pos_h_ptr + pos_off, mask=t_mask, other=0)
-    p_w = tl.load(pos_w_ptr + pos_off, mask=t_mask, other=0)
-
-    # Axis determination per d (BLOCK_D,)
-    d_in_first = tl.where(d_rng < half_d, d_rng, d_rng - half_d)
-    _bt = d_in_first < s0
-    _bh = (d_in_first >= s0) & (d_in_first < s0 + s1)
-    
-    # pos mixing: (BLOCK_T, BLOCK_D)
-    pos = tl.where(_bt[None, :],
-                   p_t[:, None],
-                   tl.where(_bh[None, :], p_h[:, None], p_w[:, None])
-                   )
-
-    # --- Single cos/sin load per (t,d) ---
-    cache_off = pos * stride_cachepos + d_rng[None, :] * stride_cached
-    cos = tl.load(cos_cache + cache_off, mask=td_mask, other=1.0)
-    sin = tl.load(sin_cache + cache_off, mask=td_mask, other=0.0)
-
-    # --- Load x and x_mirror ---
-    x_base = b * stride_xb + h * stride_xh
-    
-    x_off = x_base + t_rng[:, None] * stride_xt + d_rng[None, :] * stride_xd
-    x = tl.load(x_ptr + x_off, mask=td_mask, other=0.0)
-
-    # Mirror index for rotate_half
-    d_mir = tl.where(d_rng < half_d, d_rng + half_d, d_rng - half_d)
-    x_mir_off = x_base + t_rng[:, None] * stride_xt + d_mir[None, :] * stride_xd
-    x_mir = tl.load(x_ptr + x_mir_off, mask=td_mask, other=0.0)
-    x_rot = tl.where(d_rng[None, :] < half_d, -x_mir, x_mir)
-
-    # --- RoPE ---
-    out = x * cos + x_rot * sin
-
-    # --- Store (B, H, T, D) ---
-    o_base = b * stride_outb + h * stride_outh
-    o_off = o_base + t_rng[:, None] * stride_outt + d_rng[None, :] * stride_outd
-    tl.store(out_ptr + o_off, out, mask=td_mask)
+    half = D // 2
+
+    # Base offsets
+    x_base = b * stride_xb + t * stride_xt + h * stride_xh
+    out_base = b * stride_outb + h * stride_outh + t * stride_outt
+    pos_off = b * stride_pb + t * stride_pt
+
+    # Load positions (scalar or small vector)
+    pos_t_val = tl.load(pos_t_ptr + pos_off)
+    pos_h_val = tl.load(pos_h_ptr + pos_off)
+    pos_w_val = tl.load(pos_w_ptr + pos_off)
+
+    # Process D in blocks
+    d_off = 0
+    while d_off < D:
+        d_rng = d_off + tl.arange(0, BLOCK_D)
+        d_mask = d_rng < D
+
+        # Axis determination
+        d_in_first = tl.where(d_rng < half, d_rng, d_rng - half)
+        is_t = d_in_first < s0
+        is_h = (d_in_first >= s0) & (d_in_first < s0 + s1)
+
+        # Select position
+        pos_sel = tl.where(is_t, pos_t_val, tl.where(is_h, pos_h_val, pos_w_val))
+
+        # Load cos/sin
+        cos_off = cos_cache + pos_sel * stride_cachepos + d_rng * stride_cached
+        sin_off = sin_cache + pos_sel * stride_cachepos + d_rng * stride_cached
+        cos_val = tl.load(cos_off, mask=d_mask, other=1.0)
+      

... (2504 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
@@ -1,8 +1,7 @@
 """Optimized Triton kernel for multi-axis RoPE pre-attention prep.
 
-Strategy: each block handles one (b, h, t) — all of D in SRAM.
-Grid = B * H * T programs. Each program has 128 threads for D=128.
-Single cos/sin load per (t,d) via axis-selected position.
+Winning config from prior iteration: BLOCK_T=16, BLOCK_D=64, num_warps=4.
+Now pushing for better memory throughput with autotuning.
 """
 import torch
 import torch.nn as nn
@@ -10,6 +9,19 @@
 import triton.language as tl
 
 
+@triton.autotune(
+    configs=[
+        triton.Config({"BLOCK_T": 16, "BLOCK_D": 64}, num_warps=4),
+        triton.Config({"BLOCK_T": 16, "BLOCK_D": 64}, num_warps=8),
+        triton.Config({"BLOCK_T": 32, "BLOCK_D": 64}, num_warps=4),
+        triton.Config({"BLOCK_T": 32, "BLOCK_D": 32}, num_warps=4),
+        triton.Config({"BLOCK_T": 64, "BLOCK_D": 32}, num_warps=4),
+        triton.Config({"BLOCK_T": 64, "BLOCK_D": 32}, num_warps=8),
+        triton.Config({"BLOCK_T": 8, "BLOCK_D": 64}, num_warps=2),
+        triton.Config({"BLOCK_T": 8, "BLOCK_D": 64}, num_warps=4),
+    ],
+    key=["B", "T", "H", "D"],
+)
 @triton.jit
 def _rope_kernel(
     x_ptr,
@@ -37,130 +49,70 @@
     stride_outh,
     stride_outt,
     stride_outd,
-    D_PADDED: tl.constexpr,
+    BLOCK_T: tl.constexpr,
+    BLOCK_D: tl.constexpr,
 ):
-    """
-    Grid: (B * H * T,)
-    Each program loads the full D for one (b, h, t).
-    Thread idx = tid within block maps to D index.
-    D_PADDED should be a power of 2 >= D for efficient thread mapping.
-    """
     pid = tl.program_id(0)
-    R = B * H * T
-    if pid >= R:
-        return
-
-    t = pid % T
-    bh = (pid // T) % (B * H)
+    pid_d = tl.program_id(1)
+
+    R = tl.cdiv(T, BLOCK_T)
+    D_tiles = tl.cdiv(D, BLOCK_D)
+
+    bh = pid // R
+    r = pid % R
     b = bh // H
     h = bh % H
 
-    tid = tl.program_id(0) % D_PADDED
-    d_rng = tid + tl.arange(0, D_PADDED)
-    # Actually in Triton, we use tl.arange for vector operations.
-    # The program handles all D.
-    # With 1D grid, tl.program_id(0) = pid. Each program = one (b,h,t).
-    # Threads handle D via tl.arange.
-    pass
-
-
-# Actually better to use proper Triton patterns. Let me do 1D grid with
-# td-range within each program. Triton handles this via tl.arange(0, N)
-# where N = threads per block / warps.
-
-@triton.jit
-def _rope_per_token(
-    x_ptr,
-    cos_cache,
-    sin_cache,
-    pos_t_ptr,
-    pos_h_ptr,
-    pos_w_ptr,
-    out_ptr,
-    B: tl.int64,
-    T: tl.int64,
-    H: tl.int64,
-    D: tl.int64,
-    s0: tl.int64,
-    s1: tl.int64,
-    stride_xb,
-    stride_xt,
-    stride_xh,
-    stride_xd,
-    stride_pb,
-    stride_pt,
-    stride_cachepos,
-    stride_cached,
-    stride_outb,
-    stride_outh,
-    stride_outt,
-    stride_outd,
-    BLOCK_D: tl.constexpr,
-):
-    """
-    Each program handles one (b, h, t) and processes D in tiles of BLOCK_D.
-    Total programs = B * H * T. But that's huge (1*32*4096 = 131072).
-    Each program loops over D tiles.
-    """
-    pid = tl.program_id(0)
-
-    t = pid % T
-    bh = (pid // T) % (B * H)
-    b = bh // H
-    h = bh % H
-
-    if pid >= B * H * T:
+    if b >= B or pid_d >= D_tiles:
         return
 
-    half = D // 2
-
-    # Base offsets
-    x_base = b * stride_xb + t * stride_xt + h * stride_xh
-    out_base = b * stride_outb + h * stride_outh + t * stride_outt
-    pos_off = b * stride_pb + t * stride_pt
-
-    # Load positions (scalar or small vector)
-    pos_t_val = tl.load(pos_t_ptr + pos_off)
-    pos_h_val = tl.load(pos_h_ptr + pos_off)
-    pos_w_val = tl.load(pos_w_ptr + pos_off)
-
-    # Process D in blocks
-    d_off = 0
-    while d_off < D:
-        d_rng = d_off + tl.arange(0, BLOCK_D)
-        d_mask = d_rng < D
-
-        # Axis determination
-        d_in_first = tl.where(d_rng < half, d_rng, d_rng - half)
-        is_t = d_in_first < s0
-        is_h = (d_in_first >= s0) & (d_in_first < s0 + s1)
-
-        # Select position
-        pos_sel = tl.where(is_t, pos_t_val, tl.where(is_h, pos_h_val, pos_w_val))
-
-        # Load cos/sin
-        cos_off = cos_cache + pos_sel * stride_cachepos + d_rng * stride_cached
-        sin_off = sin_cache + pos_sel * stride_cachepos + d_rng * stride_cached
-        cos_val = tl.load(cos_off, mask=d_mask, other=1.0)
-        sin_val = tl.load(sin_off, mask=d_mask, other=0.0)
-
-        # Load x and x_mirror
-        x_off = x_ptr + x_base + d_rng * stride_xd
-        x_val = tl.load(x_off, mask=d_mask, other=0.0)
-
-        d_mir = tl.where(d_rng < half, d_rng + half, d_rng - half)
-        x_mir_off = x_ptr + x_base + d_mir * stride_xd
-        x_mir_val = tl.load(x_mir_off, mask=d_mask, other=0.0)
-        x_rot = tl.where(d_rng < half, -x_mir_val, x_mir_val)
-
-        # RoPE
-        out_val = x_val * cos_val + x_rot * sin_val
-
-        # Store
-        out_off = out_ptr + out_base + d_rng * stride_outd
-        tl.store(out_off, out_val, mask=d_mask)
-
-        d_off += BLOCK_D
+    t_off = r * BLOCK_T
+    d_off = pid_d * BLOCK_D
+    half_d = D // 2
+
+    t_rng = t_off + tl.arange(0, BLOCK_T)
+    d_rng = d_off + tl.arange(0, BLOCK_D)
+    t_mask = t_rng < T
+    d_mask = d_rng < D
+    td_mask = t_mask[:, None] & d_mask[None, :]
+
+    # Position indices
+    pos_off = b * stride_pb + t_rng * stride_pt
+    p_t = tl.load(pos_t_ptr + pos_off, mask=t_mask, other=0)
+    p_h = tl.load(pos_h_ptr + pos_off, mask=t_mask, other=0)
+    p_w = tl.load(pos_w_ptr + pos_off, mask=t_mask, other=0)
+
+    # Axis per d
+    d_in_first = tl.where(d_rng < half_d, d_rng, d_rng - half_d)
+    _bt = d_in_first < s0
+    _bh = (d_in_first >= s0) & (d_in_first < s0 + s1)
+
+    # Select position per (t,d)
+    pos_sel = tl.where(_bt[None, :], p_t[:, None],
+                 tl.where(_bh[N

... (2767 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
@@ -1,7 +1,8 @@
 """Optimized Triton kernel for multi-axis RoPE pre-attention prep.
 
-Winning config from prior iteration: BLOCK_T=16, BLOCK_D=64, num_warps=4.
-Now pushing for better memory throughput with autotuning.
+Strategy: single kernel handles both q and k. Each program covers one (b, h_q_or_k, t_tile, d_tile).
+By processing q-heads first then k-heads in the same kernel, we share code and
+potentially benefit from SM occupancy.
 """
 import torch
 import torch.nn as nn
@@ -9,21 +10,10 @@
 import triton.language as tl
 
 
-@triton.autotune(
-    configs=[
-        triton.Config({"BLOCK_T": 16, "BLOCK_D": 64}, num_warps=4),
-        triton.Config({"BLOCK_T": 16, "BLOCK_D": 64}, num_warps=8),
-        triton.Config({"BLOCK_T": 32, "BLOCK_D": 64}, num_warps=4),
-        triton.Config({"BLOCK_T": 32, "BLOCK_D": 32}, num_warps=4),
-        triton.Config({"BLOCK_T": 64, "BLOCK_D": 32}, num_warps=4),
-        triton.Config({"BLOCK_T": 64, "BLOCK_D": 32}, num_warps=8),
-        triton.Config({"BLOCK_T": 8, "BLOCK_D": 64}, num_warps=2),
-        triton.Config({"BLOCK_T": 8, "BLOCK_D": 64}, num_warps=4),
-    ],
-    key=["B", "T", "H", "D"],
-)
+# Kernel for one tensor: q or k.
+# Uses num_stages for async load pipeline (Blackwell feature).
 @triton.jit
-def _rope_kernel(
+def _rope_tensor(
     x_ptr,
     cos_cache,
     sin_cache,
@@ -37,18 +27,10 @@
     D: tl.int64,
     s0: tl.int64,
     s1: tl.int64,
-    stride_xb,
-    stride_xt,
-    stride_xh,
-    stride_xd,
-    stride_pb,
-    stride_pt,
-    stride_cachepos,
-    stride_cached,
-    stride_outb,
-    stride_outh,
-    stride_outt,
-    stride_outd,
+    stride_xb, stride_xt, stride_xh, stride_xd,
+    stride_pb, stride_pt,
+    stride_ccb, stride_ccd,
+    stride_ob, stride_oh, stride_ot, stride_od,
     BLOCK_T: tl.constexpr,
     BLOCK_D: tl.constexpr,
 ):
@@ -70,83 +52,193 @@
     d_off = pid_d * BLOCK_D
     half_d = D // 2
 
-    t_rng = t_off + tl.arange(0, BLOCK_T)
-    d_rng = d_off + tl.arange(0, BLOCK_D)
+    _rt = tl.arange(0, BLOCK_T)
+    _rd = tl.arange(0, BLOCK_D)
+    t_rng = t_off + _rt
+    d_rng = d_off + _rd
     t_mask = t_rng < T
     d_mask = d_rng < D
     td_mask = t_mask[:, None] & d_mask[None, :]
 
-    # Position indices
-    pos_off = b * stride_pb + t_rng * stride_pt
-    p_t = tl.load(pos_t_ptr + pos_off, mask=t_mask, other=0)
-    p_h = tl.load(pos_h_ptr + pos_off, mask=t_mask, other=0)
-    p_w = tl.load(pos_w_ptr + pos_off, mask=t_mask, other=0)
-
-    # Axis per d
-    d_in_first = tl.where(d_rng < half_d, d_rng, d_rng - half_d)
-    _bt = d_in_first < s0
-    _bh = (d_in_first >= s0) & (d_in_first < s0 + s1)
-
-    # Select position per (t,d)
-    pos_sel = tl.where(_bt[None, :], p_t[:, None],
-                 tl.where(_bh[None, :], p_h[:, None], p_w[:, None]))
-
-    # Cos/sin from cache (ONE load per axis position)
-    cache_off = pos_sel * stride_cachepos + d_rng[None, :] * stride_cached
-    cos = tl.load(cos_cache + cache_off, mask=td_mask, other=1.0)
-    sin = tl.load(sin_cache + cache_off, mask=td_mask, other=0.0)
-
-    # x[b, t, h, d]
-    x_base = b * stride_xb + h * stride_xh
-    x_off = x_base + t_rng[:, None] * stride_xt + d_rng[None, :] * stride_xd
-    x = tl.load(x_ptr + x_off, mask=td_mask, other=0.0)
-
-    # x mirror for rotate_half
+    # --- Stage 1: Load positions ---
+    po = b * stride_pb + t_rng * stride_pt
+    p_t = tl.load(pos_t_ptr + po, mask=t_mask, other=0)
+    p_h = tl.load(pos_h_ptr + po, mask=t_mask, other=0)
+    p_w = tl.load(pos_w_ptr + po, mask=t_mask, other=0)
+
+    # --- Stage 2: Axis select + cos/sin ---
+    d_f = tl.where(d_rng < half_d, d_rng, d_rng - half_d)
+    _bt = d_f < s0
+    _bh = (d_f >= s0) & (d_f < s0 + s1)
+    ps = tl.where(_bt[None, :], p_t[:, None],
+             tl.where(_bh[None, :], p_h[:, None], p_w[:, None]))
+
+    co = ps * stride_ccb + d_rng[None, :] * stride_ccd
+    cos = tl.load(cos_cache + co, mask=td_mask, other=1.0)
+    sin = tl.load(sin_cache + co, mask=td_mask, other=0.0)
+
+    # --- Stage 3: Load x + x_mirror, compute, store ---
+    xb = b * stride_xb + h * stride_xh
+    xo = xb + t_rng[:, None] * stride_xt + d_rng[None, :] * stride_xd
+    x = tl.load(x_ptr + xo, mask=td_mask, other=0.0)
+
     d_mir = tl.where(d_rng < half_d, d_rng + half_d, d_rng - half_d)
-    x_mir_off = x_base + t_rng[:, None] * stride_xt + d_mir[None, :] * stride_xd
-    x_mir = tl.load(x_ptr + x_mir_off, mask=td_mask, other=0.0)
-    x_rot = tl.where(d_rng[None, :] < half_d, -x_mir, x_mir)
-
-    out = x * cos + x_rot * sin
-
-    # Store to (B, H, T, D)
-    out_base = b * stride_outb + h * stride_outh
-    out_off = out_base + t_rng[:, None] * stride_outt + d_rng[None, :] * stride_outd
-    tl.store(out_ptr + out_off, out, mask=td_mask)
-
+    xmo = xb + t_rng[:, None] * stride_xt + d_mir[None, :] * stride_xd
+    xm = tl.load(x_ptr + xmo, mask=td_mask, other=0.0)
+    xr = tl.where(d_rng[None, :] < half_d, -xm, xm)
+
+    out = x * cos + xr * sin
+
+    ob = b * stride_ob + h * stride_oh
+    oo = ob + t_rng[:, None] * stride_ot + d_rng[None, :] * stride_od
+    tl.store(out_ptr + oo, out, mask=td_mask)
+
+
+# Unified kernel: each program knows if it's processing a q-head or k-head.
+# Grid_0 = B * (Hq + Hkv) * ceil(T/BLOCK_T), Grid_1 = ceil(D/BLOCK_D)
+@triton.jit
+def _rope_unified(
+    q_ptr, k_ptr,
+    cos_cache, sin_cache,
+    pos_t_ptr, pos_h_ptr, pos_w_ptr,
+    q_out_ptr, k_out_ptr,
+    B: tl.int64,
+    T: tl.int64,
+    Hq: tl.int64,
+    Hkv: tl.int64,
+    D: tl.int64,
+    s0: tl.int64,
+    s1: tl.int64,
+    stride_qb, stride_qt, stride_qh, stride_qd,
+    stride_kb, stride_kt, stride_kh, stride_kd,
+    stride_pb, stride_pt,
+    stride_ccb, stride_ccd,
+    stride_qob, stride_qot, stride_qod,
+    stride_kob, stride_kot, stride_kod

... (5401 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
@@ -1,8 +1,10 @@
 """Optimized Triton kernel for multi-axis RoPE pre-attention prep.
 
-Strategy: single kernel handles both q and k. Each program covers one (b, h_q_or_k, t_tile, d_tile).
-By processing q-heads first then k-heads in the same kernel, we share code and
-potentially benefit from SM occupancy.
+Unified kernel: grid_0 = B * (Hq + Hkv) * ceil(T/BLOCK_T), grid_1 = ceil(D/BLOCK_D).
+Each program handles one (b, h_combined, t_tile, d_tile).
+head index < Hq -> process q; head index >= Hq -> process k.
+
+No conditional branching on head type — instead we use tl.where for all operations.
 """
 import torch
 import torch.nn as nn
@@ -10,92 +12,6 @@
 import triton.language as tl
 
 
-# Kernel for one tensor: q or k.
-# Uses num_stages for async load pipeline (Blackwell feature).
-@triton.jit
-def _rope_tensor(
-    x_ptr,
-    cos_cache,
-    sin_cache,
-    pos_t_ptr,
-    pos_h_ptr,
-    pos_w_ptr,
-    out_ptr,
-    B: tl.int64,
-    T: tl.int64,
-    H: tl.int64,
-    D: tl.int64,
-    s0: tl.int64,
-    s1: tl.int64,
-    stride_xb, stride_xt, stride_xh, stride_xd,
-    stride_pb, stride_pt,
-    stride_ccb, stride_ccd,
-    stride_ob, stride_oh, stride_ot, stride_od,
-    BLOCK_T: tl.constexpr,
-    BLOCK_D: tl.constexpr,
-):
-    pid = tl.program_id(0)
-    pid_d = tl.program_id(1)
-
-    R = tl.cdiv(T, BLOCK_T)
-    D_tiles = tl.cdiv(D, BLOCK_D)
-
-    bh = pid // R
-    r = pid % R
-    b = bh // H
-    h = bh % H
-
-    if b >= B or pid_d >= D_tiles:
-        return
-
-    t_off = r * BLOCK_T
-    d_off = pid_d * BLOCK_D
-    half_d = D // 2
-
-    _rt = tl.arange(0, BLOCK_T)
-    _rd = tl.arange(0, BLOCK_D)
-    t_rng = t_off + _rt
-    d_rng = d_off + _rd
-    t_mask = t_rng < T
-    d_mask = d_rng < D
-    td_mask = t_mask[:, None] & d_mask[None, :]
-
-    # --- Stage 1: Load positions ---
-    po = b * stride_pb + t_rng * stride_pt
-    p_t = tl.load(pos_t_ptr + po, mask=t_mask, other=0)
-    p_h = tl.load(pos_h_ptr + po, mask=t_mask, other=0)
-    p_w = tl.load(pos_w_ptr + po, mask=t_mask, other=0)
-
-    # --- Stage 2: Axis select + cos/sin ---
-    d_f = tl.where(d_rng < half_d, d_rng, d_rng - half_d)
-    _bt = d_f < s0
-    _bh = (d_f >= s0) & (d_f < s0 + s1)
-    ps = tl.where(_bt[None, :], p_t[:, None],
-             tl.where(_bh[None, :], p_h[:, None], p_w[:, None]))
-
-    co = ps * stride_ccb + d_rng[None, :] * stride_ccd
-    cos = tl.load(cos_cache + co, mask=td_mask, other=1.0)
-    sin = tl.load(sin_cache + co, mask=td_mask, other=0.0)
-
-    # --- Stage 3: Load x + x_mirror, compute, store ---
-    xb = b * stride_xb + h * stride_xh
-    xo = xb + t_rng[:, None] * stride_xt + d_rng[None, :] * stride_xd
-    x = tl.load(x_ptr + xo, mask=td_mask, other=0.0)
-
-    d_mir = tl.where(d_rng < half_d, d_rng + half_d, d_rng - half_d)
-    xmo = xb + t_rng[:, None] * stride_xt + d_mir[None, :] * stride_xd
-    xm = tl.load(x_ptr + xmo, mask=td_mask, other=0.0)
-    xr = tl.where(d_rng[None, :] < half_d, -xm, xm)
-
-    out = x * cos + xr * sin
-
-    ob = b * stride_ob + h * stride_oh
-    oo = ob + t_rng[:, None] * stride_ot + d_rng[None, :] * stride_od
-    tl.store(out_ptr + oo, out, mask=td_mask)
-
-
-# Unified kernel: each program knows if it's processing a q-head or k-head.
-# Grid_0 = B * (Hq + Hkv) * ceil(T/BLOCK_T), Grid_1 = ceil(D/BLOCK_D)
 @triton.jit
 def _rope_unified(
     q_ptr, k_ptr,
@@ -113,8 +29,8 @@
     stride_kb, stride_kt, stride_kh, stride_kd,
     stride_pb, stride_pt,
     stride_ccb, stride_ccd,
-    stride_qob, stride_qot, stride_qod,
-    stride_kob, stride_kot, stride_kod,
+    stride_qob, stride_qoh, stride_qot, stride_qod,
+    stride_kob, stride_koh, stride_kot, stride_kod,
     BLOCK_T: tl.constexpr,
     BLOCK_D: tl.constexpr,
 ):
@@ -128,7 +44,7 @@
     bh = pid // R
     r = pid % R
     b = bh // total_h
-    h = bh % total_h
+    h_combined = bh % total_h
 
     if b >= B or pid_d >= D_tiles:
         return
@@ -145,60 +61,60 @@
     d_mask = d_rng < D
     td_mask = t_mask[:, None] & d_mask[None, :]
 
-    # Positions (broadcast across q and k since same pos)
+    # --- Load positions (shared for q and k since pos is per-token) ---
     po = b * stride_pb + t_rng * stride_pt
     p_t = tl.load(pos_t_ptr + po, mask=t_mask, other=0)
     p_h = tl.load(pos_h_ptr + po, mask=t_mask, other=0)
     p_w = tl.load(pos_w_ptr + po, mask=t_mask, other=0)
 
-    # Axis select
+    # --- Axis select per d ---
     d_f = tl.where(d_rng < half_d, d_rng, d_rng - half_d)
-    _bt = d_f < s0
-    _bh = (d_f >= s0) & (d_f < s0 + s1)
-    ps = tl.where(_bt[None, :], p_t[:, None],
-             tl.where(_bh[None, :], p_h[:, None], p_w[:, None]))
-
-    co = ps * stride_ccb + d_rng[None, :] * stride_ccd
-    cos = tl.load(cos_cache + co, mask=td_mask, other=1.0)
-    sin = tl.load(sin_cache + co, mask=td_mask, other=0.0)
-
-    # Determine if this is a q-head or k-head, and apply RoPE accordingly
-    is_q = h < Hq
-    h_idx = tl.where(is_q, h, h - Hq)
-
-    # Q processing
-    if is_q:
-        qb = b * stride_qb + h * stride_qh
-        qo = qb + t_rng[:, None] * stride_qt + d_rng[None, :] * stride_qd
-        qx = tl.load(q_ptr + qo, mask=td_mask, other=0.0)
-
-        d_mir = tl.where(d_rng < half_d, d_rng + half_d, d_rng - half_d)
-        qmo = qb + t_rng[:, None] * stride_qt + d_mir[None, :] * stride_qd
-        qm = tl.load(q_ptr + qmo, mask=td_mask, other=0.0)
-        qr = tl.where(d_rng[None, :] < half_d, -qm, qm)
-
-        qout = qx * cos + qr * sin
-
-        qob = b * stride_qob
-        qoo = qob + h * stride_qot * 0 + t_rng[:, None] * stride_qot + d_rng[None, :] * stride_qod
-        Wait this is wrong. Let me use actual stride values.
-        # Out stride for (B, H, T, D): stride_qob = T*D*2, stride for h = D*2, stride for t = D*2, stride for d = 2
-    

... (4832 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
@@ -1,10 +1,7 @@
 """Optimized Triton kernel for multi-axis RoPE pre-attention prep.
 
-Unified kernel: grid_0 = B * (Hq + Hkv) * ceil(T/BLOCK_T), grid_1 = ceil(D/BLOCK_D).
-Each program handles one (b, h_combined, t_tile, d_tile).
-head index < Hq -> process q; head index >= Hq -> process k.
-
-No conditional branching on head type — instead we use tl.where for all operations.
+Separate kernels for q and k (pointers can't be tl.where'd in Triton).
+Each kernel: grid_0 = B*H*ceil(T/BLOCK_T), grid_1 = ceil(D/BLOCK_D).
 """
 import torch
 import torch.nn as nn
@@ -13,24 +10,32 @@
 
 
 @triton.jit
-def _rope_unified(
-    q_ptr, k_ptr,
-    cos_cache, sin_cache,
-    pos_t_ptr, pos_h_ptr, pos_w_ptr,
-    q_out_ptr, k_out_ptr,
+def _rope_kernel(
+    x_ptr,
+    cos_cache,
+    sin_cache,
+    pos_t_ptr,
+    pos_h_ptr,
+    pos_w_ptr,
+    out_ptr,
     B: tl.int64,
     T: tl.int64,
-    Hq: tl.int64,
-    Hkv: tl.int64,
+    H: tl.int64,
     D: tl.int64,
     s0: tl.int64,
     s1: tl.int64,
-    stride_qb, stride_qt, stride_qh, stride_qd,
-    stride_kb, stride_kt, stride_kh, stride_kd,
-    stride_pb, stride_pt,
-    stride_ccb, stride_ccd,
-    stride_qob, stride_qoh, stride_qot, stride_qod,
-    stride_kob, stride_koh, stride_kot, stride_kod,
+    stride_xb,
+    stride_xt,
+    stride_xh,
+    stride_xd,
+    stride_pb,
+    stride_pt,
+    stride_ccb,
+    stride_ccd,
+    stride_ob,
+    stride_oh,
+    stride_ot,
+    stride_od,
     BLOCK_T: tl.constexpr,
     BLOCK_D: tl.constexpr,
 ):
@@ -39,12 +44,11 @@
 
     R = tl.cdiv(T, BLOCK_T)
     D_tiles = tl.cdiv(D, BLOCK_D)
-    total_h = Hq + Hkv
 
     bh = pid // R
     r = pid % R
-    b = bh // total_h
-    h_combined = bh % total_h
+    b = bh // H
+    h = bh % H
 
     if b >= B or pid_d >= D_tiles:
         return
@@ -61,59 +65,62 @@
     d_mask = d_rng < D
     td_mask = t_mask[:, None] & d_mask[None, :]
 
-    # --- Load positions (shared for q and k since pos is per-token) ---
+    # Position indices
     po = b * stride_pb + t_rng * stride_pt
     p_t = tl.load(pos_t_ptr + po, mask=t_mask, other=0)
     p_h = tl.load(pos_h_ptr + po, mask=t_mask, other=0)
     p_w = tl.load(pos_w_ptr + po, mask=t_mask, other=0)
 
-    # --- Axis select per d ---
+    # Axis select per d
     d_f = tl.where(d_rng < half_d, d_rng, d_rng - half_d)
-    is_axis_t = d_f < s0
-    is_axis_h = (d_f >= s0) & (d_f < s0 + s1)
-    pos_sel = tl.where(is_axis_t[None, :], p_t[:, None],
-                 tl.where(is_axis_h[None, :], p_h[:, None], p_w[:, None]))
+    is_t = d_f < s0
+    is_h = (d_f >= s0) & (d_f < s0 + s1)
+    pos_sel = tl.where(is_t[None, :], p_t[:, None],
+                 tl.where(is_h[None, :], p_h[:, None], p_w[:, None]))
 
-    # --- Cos/sin cache load (shared) ---
+    # Single cos/sin load
     cache_off = pos_sel * stride_ccb + d_rng[None, :] * stride_ccd
     cos = tl.load(cos_cache + cache_off, mask=td_mask, other=1.0)
     sin = tl.load(sin_cache + cache_off, mask=td_mask, other=0.0)
 
-    # --- Determine if this program handles q or k ---
-    is_q = h_combined < Hq
-    h_idx = tl.where(is_q, h_combined, h_combined - Hq)
+    # x[b, t, h, d]
+    x_base = b * stride_xb + h * stride_xh
+    x_off = x_base + t_rng[:, None] * stride_xt + d_rng[None, :] * stride_xd
+    x = tl.load(x_ptr + x_off, mask=td_mask, other=0.0)
 
-    # Select strides and pointers based on is_q
-    stride_b = tl.where(is_q, stride_qb, stride_kb)
-    stride_t = tl.where(is_q, stride_qt, stride_kt)
-    stride_h = tl.where(is_q, stride_qh, stride_kh)
-    stride_d = tl.where(is_q, stride_qd, stride_kd)
-    in_ptr = tl.where(is_q, q_ptr, k_ptr)
-
-    stride_ob = tl.where(is_q, stride_qob, stride_kob)
-    stride_oh = tl.where(is_q, stride_qoh, stride_koh)
-    stride_ot = tl.where(is_q, stride_qot, stride_kot)
-    stride_od = tl.where(is_q, stride_qod, stride_kod)
-    out_ptr = tl.where(is_q, q_out_ptr, k_out_ptr)
-
-    # Load x
-    x_base = b * stride_b + h_idx * stride_h
-    x_off = x_base + t_rng[:, None] * stride_t + d_rng[None, :] * stride_d
-    x = tl.load(in_ptr + x_off, mask=td_mask, other=0.0)
-
-    # Load x_mirror
+    # Mirror for rotate_half
     d_mir = tl.where(d_rng < half_d, d_rng + half_d, d_rng - half_d)
-    x_mir_off = x_base + t_rng[:, None] * stride_t + d_mir[None, :] * stride_d
-    x_mir = tl.load(in_ptr + x_mir_off, mask=td_mask, other=0.0)
+    x_mir_off = x_base + t_rng[:, None] * stride_xt + d_mir[None, :] * stride_xd
+    x_mir = tl.load(x_ptr + x_mir_off, mask=td_mask, other=0.0)
     x_rot = tl.where(d_rng[None, :] < half_d, -x_mir, x_mir)
 
     # RoPE
     out = x * cos + x_rot * sin
 
-    # Store
-    out_base = b * stride_ob + h_idx * stride_oh
+    # Store to (B, H, T, D)
+    out_base = b * stride_ob + h * stride_oh
     out_off = out_base + t_rng[:, None] * stride_ot + d_rng[None, :] * stride_od
     tl.store(out_ptr + out_off, out, mask=td_mask)
+
+
+def _run_kernel(x, cos_cache, sin_cache, pos_t, pos_h, pos_w, out, H, D,
+                s0, s1, BLOCK_T, BLOCK_D, num_warps):
+    B, T = x.shape[0], x.shape[1]
+    n_t = (T + BLOCK_T - 1) // BLOCK_T
+    n_d = (D + BLOCK_D - 1) // BLOCK_D
+
+    _rope_kernel[(B * H * n_t, n_d)](
+        x, cos_cache, sin_cache, pos_t, pos_h, pos_w, out,
+        B=B, T=T, H=H, D=D, s0=s0, s1=s1,
+        stride_xb=x.stride(0), stride_xt=x.stride(1),
+        stride_xh=x.stride(2), stride_xd=x.stride(3),
+        stride_pb=pos_t.stride(0), stride_pt=pos_t.stride(1),
+        stride_ccb=cos_cache.stride(0), stride_ccd=cos_cache.stride(1),
+        stride_ob=out.stride(0), stride_oh=out.stride(1),
+        stride_ot=out.stride(2), stride_od=out.stride(3),
+        BLOCK_T=BLOCK_T, BLOCK_D=BLOCK_D,
+        num_warps=num_warps,
+    )
 
 
 def apply_mrope_torch(q, k, cos_cache, sin_cache, pos_t, p

... (1504 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
@@ -1,7 +1,9 @@
 """Optimized Triton kernel for multi-axis RoPE pre-attention prep.
 
-Separate kernels for q and k (pointers can't be tl.where'd in Triton).
-Each kernel: grid_0 = B*H*ceil(T/BLOCK_T), grid_1 = ceil(D/BLOCK_D).
+Key optimization: cos/sin depends only on (b, t, d), NOT on h.
+Grid over (B * ceil(T/BLOCK_T), ceil(D/BLOCK_D)).
+Each program loads cos/sin ONCE, then loops over all H (both q and k heads).
+This saves Hq+Hkv factor in random cache reads — critical for bandwidth.
 """
 import torch
 import torch.nn as nn
@@ -11,50 +13,45 @@
 
 @triton.jit
 def _rope_kernel(
-    x_ptr,
-    cos_cache,
-    sin_cache,
-    pos_t_ptr,
-    pos_h_ptr,
-    pos_w_ptr,
-    out_ptr,
+    q_ptr, k_ptr,
+    cos_cache, sin_cache,
+    pos_t_ptr, pos_h_ptr, pos_w_ptr,
+    q_out_ptr, k_out_ptr,
     B: tl.int64,
     T: tl.int64,
-    H: tl.int64,
+    Hq: tl.int64,
+    Hkv: tl.int64,
     D: tl.int64,
     s0: tl.int64,
     s1: tl.int64,
-    stride_xb,
-    stride_xt,
-    stride_xh,
-    stride_xd,
-    stride_pb,
-    stride_pt,
-    stride_ccb,
-    stride_ccd,
-    stride_ob,
-    stride_oh,
-    stride_ot,
-    stride_od,
+    stride_qb, stride_qt, stride_qh, stride_qd,
+    stride_kb, stride_kt, stride_kh, stride_kd,
+    stride_pb, stride_pt,
+    stride_ccb, stride_ccd,
+    stride_qob, stride_qoh, stride_qot, stride_qod,
+    stride_kob, stride_koh, stride_kot, stride_kod,
     BLOCK_T: tl.constexpr,
     BLOCK_D: tl.constexpr,
 ):
+    """
+    Grid: (B * ceil(T/BLOCK_T), ceil(D/BLOCK_D))
+    Each program handles one (b, t_tile, d_tile) and ALL heads.
+    cos/sin loaded once, then RoPE applied to each head.
+    """
     pid = tl.program_id(0)
     pid_d = tl.program_id(1)
 
     R = tl.cdiv(T, BLOCK_T)
     D_tiles = tl.cdiv(D, BLOCK_D)
 
-    bh = pid // R
     r = pid % R
-    b = bh // H
-    h = bh % H
+    b = pid // R
+    d_off = pid_d * BLOCK_D
 
     if b >= B or pid_d >= D_tiles:
         return
 
     t_off = r * BLOCK_T
-    d_off = pid_d * BLOCK_D
     half_d = D // 2
 
     _rt = tl.arange(0, BLOCK_T)
@@ -65,62 +62,65 @@
     d_mask = d_rng < D
     td_mask = t_mask[:, None] & d_mask[None, :]
 
-    # Position indices
+    # --- Load positions ---
     po = b * stride_pb + t_rng * stride_pt
     p_t = tl.load(pos_t_ptr + po, mask=t_mask, other=0)
     p_h = tl.load(pos_h_ptr + po, mask=t_mask, other=0)
     p_w = tl.load(pos_w_ptr + po, mask=t_mask, other=0)
 
-    # Axis select per d
+    # --- Axis select per d ---
     d_f = tl.where(d_rng < half_d, d_rng, d_rng - half_d)
     is_t = d_f < s0
     is_h = (d_f >= s0) & (d_f < s0 + s1)
     pos_sel = tl.where(is_t[None, :], p_t[:, None],
                  tl.where(is_h[None, :], p_h[:, None], p_w[:, None]))
 
-    # Single cos/sin load
+    # --- Cos/sin cache (loaded ONCE per program, reused for all H) ---
     cache_off = pos_sel * stride_ccb + d_rng[None, :] * stride_ccd
     cos = tl.load(cos_cache + cache_off, mask=td_mask, other=1.0)
     sin = tl.load(sin_cache + cache_off, mask=td_mask, other=0.0)
 
-    # x[b, t, h, d]
-    x_base = b * stride_xb + h * stride_xh
-    x_off = x_base + t_rng[:, None] * stride_xt + d_rng[None, :] * stride_xd
-    x = tl.load(x_ptr + x_off, mask=td_mask, other=0.0)
-
-    # Mirror for rotate_half
-    d_mir = tl.where(d_rng < half_d, d_rng + half_d, d_rng - half_d)
-    x_mir_off = x_base + t_rng[:, None] * stride_xt + d_mir[None, :] * stride_xd
-    x_mir = tl.load(x_ptr + x_mir_off, mask=td_mask, other=0.0)
-    x_rot = tl.where(d_rng[None, :] < half_d, -x_mir, x_mir)
-
-    # RoPE
-    out = x * cos + x_rot * sin
-
-    # Store to (B, H, T, D)
-    out_base = b * stride_ob + h * stride_oh
-    out_off = out_base + t_rng[:, None] * stride_ot + d_rng[None, :] * stride_od
-    tl.store(out_ptr + out_off, out, mask=td_mask)
-
-
-def _run_kernel(x, cos_cache, sin_cache, pos_t, pos_h, pos_w, out, H, D,
-                s0, s1, BLOCK_T, BLOCK_D, num_warps):
-    B, T = x.shape[0], x.shape[1]
-    n_t = (T + BLOCK_T - 1) // BLOCK_T
-    n_d = (D + BLOCK_D - 1) // BLOCK_D
-
-    _rope_kernel[(B * H * n_t, n_d)](
-        x, cos_cache, sin_cache, pos_t, pos_h, pos_w, out,
-        B=B, T=T, H=H, D=D, s0=s0, s1=s1,
-        stride_xb=x.stride(0), stride_xt=x.stride(1),
-        stride_xh=x.stride(2), stride_xd=x.stride(3),
-        stride_pb=pos_t.stride(0), stride_pt=pos_t.stride(1),
-        stride_ccb=cos_cache.stride(0), stride_ccd=cos_cache.stride(1),
-        stride_ob=out.stride(0), stride_oh=out.stride(1),
-        stride_ot=out.stride(2), stride_od=out.stride(3),
-        BLOCK_T=BLOCK_T, BLOCK_D=BLOCK_D,
-        num_warps=num_warps,
-    )
+    # --- Process Q heads ---
+    h = tl.arange(0, Hq)
+    # Loop over Hq heads — use scalar loop since Triton can handle it
+    h_idx = 0
+    while h_idx < Hq:
+        qb = b * stride_qb + h_idx * stride_qh
+        q_off = qb + t_rng[:, None] * stride_qt + d_rng[None, :] * stride_qd
+        qx = tl.load(q_ptr + q_off, mask=td_mask, other=0.0)
+
+        d_mir = tl.where(d_rng < half_d, d_rng + half_d, d_rng - half_d)
+        q_mir_off = qb + t_rng[:, None] * stride_qt + d_mir[None, :] * stride_qd
+        qm = tl.load(q_ptr + q_mir_off, mask=td_mask, other=0.0)
+        qr = tl.where(d_rng[None, :] < half_d, -qm, qm)
+
+        qout = qx * cos + qr * sin
+
+        q_out_base = b * stride_qob + h_idx * stride_qoh
+        q_out_off = q_out_base + t_rng[:, None] * stride_qot + d_rng[None, :] * stride_qod
+        tl.store(q_out_ptr + q_out_off, qout, mask=td_mask)
+
+        h_idx += 1
+
+    # --- Process K heads ---
+    h_idx = 0
+    while h_idx < Hkv:
+        kb = b * stride_kb + h_idx * stride_kh
+        k_off = kb + t_rng[:, None] * stride_kt + d_rng[None, :] * stride_kd
+        kx = tl.load(k_ptr + k_off, mask=td_mask, other=0.0)
+
+    

... (2113 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
@@ -81,8 +81,6 @@
     sin = tl.load(sin_cache + cache_off, mask=td_mask, other=0.0)
 
     # --- Process Q heads ---
-    h = tl.arange(0, Hq)
-    # Loop over Hq heads — use scalar loop since Triton can handle it
     h_idx = 0
     while h_idx < Hq:
         qb = b * stride_qb + h_idx * stride_qh

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
@@ -1,9 +1,7 @@
 """Optimized Triton kernel for multi-axis RoPE pre-attention prep.
 
-Key optimization: cos/sin depends only on (b, t, d), NOT on h.
-Grid over (B * ceil(T/BLOCK_T), ceil(D/BLOCK_D)).
-Each program loads cos/sin ONCE, then loops over all H (both q and k heads).
-This saves Hq+Hkv factor in random cache reads — critical for bandwidth.
+Unified kernel: cos/sin loaded once per (b, t_tile, d_tile), reused across all H.
+Tunable BLOCK_T, BLOCK_D, num_warps for each shape.
 """
 import torch
 import torch.nn as nn
@@ -33,11 +31,6 @@
     BLOCK_T: tl.constexpr,
     BLOCK_D: tl.constexpr,
 ):
-    """
-    Grid: (B * ceil(T/BLOCK_T), ceil(D/BLOCK_D))
-    Each program handles one (b, t_tile, d_tile) and ALL heads.
-    cos/sin loaded once, then RoPE applied to each head.
-    """
     pid = tl.program_id(0)
     pid_d = tl.program_id(1)
 
@@ -62,63 +55,53 @@
     d_mask = d_rng < D
     td_mask = t_mask[:, None] & d_mask[None, :]
 
-    # --- Load positions ---
+    # Positions
     po = b * stride_pb + t_rng * stride_pt
     p_t = tl.load(pos_t_ptr + po, mask=t_mask, other=0)
     p_h = tl.load(pos_h_ptr + po, mask=t_mask, other=0)
     p_w = tl.load(pos_w_ptr + po, mask=t_mask, other=0)
 
-    # --- Axis select per d ---
+    # Axis select per d
     d_f = tl.where(d_rng < half_d, d_rng, d_rng - half_d)
     is_t = d_f < s0
     is_h = (d_f >= s0) & (d_f < s0 + s1)
     pos_sel = tl.where(is_t[None, :], p_t[:, None],
                  tl.where(is_h[None, :], p_h[:, None], p_w[:, None]))
 
-    # --- Cos/sin cache (loaded ONCE per program, reused for all H) ---
+    # Cos/sin cache — loaded ONCE, reused for all heads
     cache_off = pos_sel * stride_ccb + d_rng[None, :] * stride_ccd
     cos = tl.load(cos_cache + cache_off, mask=td_mask, other=1.0)
     sin = tl.load(sin_cache + cache_off, mask=td_mask, other=0.0)
 
-    # --- Process Q heads ---
-    h_idx = 0
-    while h_idx < Hq:
-        qb = b * stride_qb + h_idx * stride_qh
-        q_off = qb + t_rng[:, None] * stride_qt + d_rng[None, :] * stride_qd
-        qx = tl.load(q_ptr + q_off, mask=td_mask, other=0.0)
+    # Rotate-half mirror indices (same for all heads)
+    d_mir = tl.where(d_rng < half_d, d_rng + half_d, d_rng - half_d)
+    is_first_half = d_rng[None, :] < half_d
 
-        d_mir = tl.where(d_rng < half_d, d_rng + half_d, d_rng - half_d)
-        q_mir_off = qb + t_rng[:, None] * stride_qt + d_mir[None, :] * stride_qd
-        qm = tl.load(q_ptr + q_mir_off, mask=td_mask, other=0.0)
-        qr = tl.where(d_rng[None, :] < half_d, -qm, qm)
+    # Process Q heads
+    h = 0
+    while h < Hq:
+        q_base = b * stride_qb + h * stride_qh
+        qx = tl.load(q_ptr + q_base + t_rng[:, None] * stride_qt + d_rng[None, :] * stride_qd,
+                      mask=td_mask, other=0.0)
+        qm = tl.load(q_ptr + q_base + t_rng[:, None] * stride_qt + d_mir[None, :] * stride_qd,
+                      mask=td_mask, other=0.0)
+        out = qx * cos + tl.where(is_first_half, -qm, qm) * sin
+        tl.store(q_out_ptr + b * stride_qob + h * stride_qoh + t_rng[:, None] * stride_qot + d_rng[None, :] * stride_qod,
+                 out, mask=td_mask)
+        h += 1
 
-        qout = qx * cos + qr * sin
-
-        q_out_base = b * stride_qob + h_idx * stride_qoh
-        q_out_off = q_out_base + t_rng[:, None] * stride_qot + d_rng[None, :] * stride_qod
-        tl.store(q_out_ptr + q_out_off, qout, mask=td_mask)
-
-        h_idx += 1
-
-    # --- Process K heads ---
-    h_idx = 0
-    while h_idx < Hkv:
-        kb = b * stride_kb + h_idx * stride_kh
-        k_off = kb + t_rng[:, None] * stride_kt + d_rng[None, :] * stride_kd
-        kx = tl.load(k_ptr + k_off, mask=td_mask, other=0.0)
-
-        d_mir = tl.where(d_rng < half_d, d_rng + half_d, d_rng - half_d)
-        k_mir_off = kb + t_rng[:, None] * stride_kt + d_mir[None, :] * stride_kd
-        km = tl.load(k_ptr + k_mir_off, mask=td_mask, other=0.0)
-        kr = tl.where(d_rng[None, :] < half_d, -km, km)
-
-        kout = kx * cos + kr * sin
-
-        k_out_base = b * stride_kob + h_idx * stride_koh
-        k_out_off = k_out_base + t_rng[:, None] * stride_kot + d_rng[None, :] * stride_kod
-        tl.store(k_out_ptr + k_out_off, kout, mask=td_mask)
-
-        h_idx += 1
+    # Process K heads
+    h = 0
+    while h < Hkv:
+        k_base = b * stride_kb + h * stride_kh
+        kx = tl.load(k_ptr + k_base + t_rng[:, None] * stride_kt + d_rng[None, :] * stride_kd,
+                      mask=td_mask, other=0.0)
+        km = tl.load(k_ptr + k_base + t_rng[:, None] * stride_kt + d_mir[None, :] * stride_kd,
+                      mask=td_mask, other=0.0)
+        out = kx * cos + tl.where(is_first_half, -km, km) * sin
+        tl.store(k_out_ptr + b * stride_kob + h * stride_koh + t_rng[:, None] * stride_kot + d_rng[None, :] * stride_kod,
+                 out, mask=td_mask)
+        h += 1
 
 
 def apply_mrope_torch(q, k, cos_cache, sin_cache, pos_t, pos_h, pos_w, mrope_section):
@@ -126,11 +109,18 @@
     _, _, Hkv, _ = k.shape
     s0, s1, s2 = mrope_section
 
-    BLOCK_T = 16
-    BLOCK_D = 64
-
     q_rot = torch.empty(B, Hq, T, D, dtype=q.dtype, device=q.device)
     k_rot = torch.empty(B, Hkv, T, D, dtype=k.dtype, device=k.device)
+
+    # Tunable configs — selected based on shape characteristics
+    # Larger H benefits from smaller grid => larger BLOCK_T
+    # Bandwidth-bound: want maximal memory throughput
+    if D == 64:
+        BLOCK_T, BLOCK_D, num_warps = 32, 64, 4
+    elif T <= 4096:
+        BLOCK_T, BLOCK_D, num_warps = 16, 64, 4
+    else:
+        BLOCK_T, BLOCK_D, num_warps = 16, 64, 8
 
     n_t = (T + BLOCK_T - 1) // BLOCK_T
     n_d = (D + BLOCK_D - 1) // BLOCK_D
@@ -152,7 +142,7 @@
         stride_kob=k_rot.stride(0), stride_koh=k_rot.stride(1),
         

... (185 more chars)