--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/09_fmha_preattn_mrope/solution.py
@@ -0,0 +1,492 @@
+"""Fused multi-axis RoPE pre-attention prep.
+
+Fuses: gather per-axis cos/sin from cache, mrope_section mixing, rotate-half,
+and (B,T,H,D)->(B,H,T,D) transpose into a single stream of q and k.
+
+The reference does: gather(6xB,T,D) + mix + broadcast + rotate + transpose(2x).
+This kernel does the same but reads q/k only once and writes q_rot/k_rot directly
+in attention layout, with per-axis cache values computed once per (B,T) and
+reused across all heads.
+"""
+import os
+import torch
+import torch.utils.cpp_extension
+
+_CC = """#include <torch/extension.h>
+#include <cuda/bfloat16.h>
+#include <cuda_fp16.h>
+#include <mma.h>
+#include <stdio.h>
+
+namespace {
+
+#define WARP_SIZE 32
+#define MAX_D 128
+
+// Device lambda to check compute capability
+__device__ int get_sm_count() { return 0; }
+
+struct HalvedDims {
+    uint32_t s0, s1, s2, h;
+};
+
+__device__ uint2 gather_bf16x2(const __half* cache, int64_t idx, int max_pos) {
+    idx = (idx < 0) ? 0 : (idx >= max_pos ? max_pos - 1 : idx);
+    uint2 val;
+    val.x = ((const uint32_t*)cache)[idx * 2];
+    val.y = ((const uint32_t*)cache)[idx * 2 + 1];
+    return val;
+}
+
+__device__ float2 bfi_load(const __half* cache, int64_t idx, int max_pos) {
+    idx = (idx < 0) ? 0 : (idx >= max_pos ? max_pos - 1 : idx);
+    float2 val;
+    val.x = __low2float(((const uint32_t*)cache)[idx * 2]);
+    val.y = __high2float(((const uint32_t*)cache)[idx * 2 + 1]);
+    return val;
+}
+
+template <int D, int MAX_DV>
+__global__ void mrope_kernel(
+    const uint16_t* __restrict__ q_ptr,
+    const uint16_t* __restrict__ k_ptr,
+    const int64_t* __restrict__ pos_t,
+    const int64_t* __restrict__ pos_h,
+    const int64_t* __restrict__ pos_w,
+    const uint16_t* __restrict__ cos_ptr,
+    const uint16_t* __restrict__ sin_ptr,
+    uint16_t* __restrict__ q_out,
+    uint16_t* __restrict__ k_out,
+    int B, int T, int Hq, int Hkv, int D_,
+    int s0, int s1, int s2,
+    int max_pos,
+    int q_stride_B, int q_stride_T, int q_stride_H,
+    int k_stride_B, int k_stride_T, int k_stride_H,
+    int q_out_stride_B, int q_out_stride_H, int q_out_stride_T,
+    int k_out_stride_B, int k_out_stride_H, int k_out_stride_T
+) {
+    // Grid: (B * Hq, T, 1) for q kernel; (B * Hkv, T, 1) for k kernel
+    // Each block handles one head of one batch element, all T positions
+    int b = blockIdx.x / Hq;
+    int h = blockIdx.x % Hq;
+    int t = blockIdx.y;
+
+    if (b >= B) return;
+
+    // Determine if this is a q or k block based on q_ptr vs k_ptr pointer values
+    // Actually we launch two separate kernels, so we only use one pointer
+    // This kernel handles both q and k; we pass a flag via blockIdx.z
+    // blockIdx.z: 0 = q, 1 = k
+    bool is_q = (blockIdx.z == 0);
+    int H = is_q ? Hq : Hkv;
+
+    // Check if this head is valid for k
+    if (!is_q && h >= Hkv) return;
+
+    // Shared mem for position indices (T positions, 3 axes)
+    extern __shared__ uint64_t shared_pos[];
+    int64_t* s_pos_t = (int64_t*)shared_pos;
+    int64_t* s_pos_h = s_pos_t + T;
+    int64_t* s_pos_w = s_pos_h + T;
+
+    // Load position indices for all T values in this B
+    if (threadIdx.x < T) {
+        s_pos_t[threadIdx.x] = pos_t[b * T + threadIdx.x];
+        s_pos_h[threadIdx.x] = pos_h[b * T + threadIdx.x];
+        s_pos_w[threadIdx.x] = pos_w[b * T + threadIdx.x];
+    }
+    __syncthreads();
+
+    int64_t pt = s_pos_t[t];
+    int64_t ph = s_pos_h[t];
+    int64_t pw = s_pos_w[t];
+
+    // Compute cos/sin mix for this temporal position
+    // cos/sin cache is (max_pos, D) bf16
+    // We'll load pairs of bf16 as float2
+
+    float2 cos_val[MAX_D/2];
+    float2 sin_val[MAX_D/2];
+
+    // Gather and mix per-axis cos/sin
+    const int halfD = D_ / 2;
+
+    // cos/sin cache layout: (max_pos, D) stored as (max_pos, D/2) pairs of uint32
+    // idx for halfD elements = position * halfD * 2 uint32's worth, but we access as pairs
+    // Actually cos_ptr is (max_pos, D) bf16 = (max_pos, 2*halfD) uint16_t
+    // We access as uint32 pairs (2 bf16 = 1 uint32)
+    // So index for position p, pair i is: p * D + i*2 (as uint16) or p * halfD + i (as uint32 pairs)
+
+    // For pair i in [0, halfD):
+    // cos_val[i] = mix(cos_t[i], cos_h[i], cos_w[i]) according to mrope_section
+
+    // mrope_section: first s0 pairs use t, next s1 use h, last s2 use w
+    // Mirrored in second half starting at halfD
+
+    // Load all cos/sin pairs and mix
+    #pragma unroll
+    for (int i = 0; i < MAX_D/2; ++i) {
+        if (i >= halfD) {
+            // Second half - zero out, handled below
+            cos_val[i].x = 0.0f; cos_val[i].y = 0.0f;
+            sin_val[i].x = 0.0f; sin_val[i].y = 0.0f;
+            continue;
+        }
+
+        // Gather from cache
+        // cos/sin is stored as (max_pos, D) bf16 = (max_pos, 2*halfD) uint16
+        // We read pairs as uint32 (2 consecutive bf16)
+        const uint32_t* cos32 = (const uint32_t*)cos_ptr;
+        const uint32_t* sin32 = (const uint32_t*)sin_ptr;
+
+        // But we need to be careful: D_ is the actual D, MAX_D is template param
+        uint32_t cos_t_lo = 0, cos_t_hi = 0, cos_h_lo = 0, cos_h_hi = 0, cos_w_lo = 0, cos_w_hi = 0;
+        uint32_t sin_t_lo = 0, sin_t_hi = 0, sin_h_lo = 0, sin_h_hi = 0, sin_w_lo = 0, sin_w_hi = 0;
+
+        // Load pairs
+        int p_t = (int)(pt < 0 ? 0 : (pt >= max_pos ? max_pos - 1 : pt));
+        int p_h = (int)(ph < 0 ? 0 : (ph >= max_pos ? max_pos - 1 : ph));
+        int p_w = (int)(pw < 0 ? 0 : (pw >= max_pos ? max_pos - 1 : pw));
+
+        // Position in uint32 pairs: position * halfD + i
+        uint32_t base_t = p_t * halfD + i;
+        uint32_t base_h = p_h * halfD + i;
+        uint32_t base_w = p_w * halfD + i;
+
+        cos_t_lo = cos32[base_t];
+        sin_t_lo = sin32[base_t];
+        cos_h_lo = cos32[base_

... (11243 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,94 +2,206 @@
 
 Fuses: gather per-axis cos/sin from cache, mrope_section mixing, rotate-half,
 and (B,T,H,D)->(B,H,T,D) transpose into a single stream of q and k.
-
-The reference does: gather(6xB,T,D) + mix + broadcast + rotate + transpose(2x).
-This kernel does the same but reads q/k only once and writes q_rot/k_rot directly
-in attention layout, with per-axis cache values computed once per (B,T) and
-reused across all heads.
 """
 import os
 import torch
 import torch.utils.cpp_extension
 
-_CC = """#include <torch/extension.h>
-#include <cuda/bfloat16.h>
+_CC = """
+#include <torch/extension.h>
+#include <stdio.h>
 #include <cuda_fp16.h>
-#include <mma.h>
-#include <stdio.h>
 
 namespace {
 
 #define WARP_SIZE 32
-#define MAX_D 128
-
-// Device lambda to check compute capability
-__device__ int get_sm_count() { return 0; }
-
-struct HalvedDims {
-    uint32_t s0, s1, s2, h;
-};
-
-__device__ uint2 gather_bf16x2(const __half* cache, int64_t idx, int max_pos) {
-    idx = (idx < 0) ? 0 : (idx >= max_pos ? max_pos - 1 : idx);
-    uint2 val;
-    val.x = ((const uint32_t*)cache)[idx * 2];
-    val.y = ((const uint32_t*)cache)[idx * 2 + 1];
-    return val;
-}
-
-__device__ float2 bfi_load(const __half* cache, int64_t idx, int max_pos) {
-    idx = (idx < 0) ? 0 : (idx >= max_pos ? max_pos - 1 : idx);
-    float2 val;
-    val.x = __low2float(((const uint32_t*)cache)[idx * 2]);
-    val.y = __high2float(((const uint32_t*)cache)[idx * 2 + 1]);
-    return val;
-}
-
-template <int D, int MAX_DV>
-__global__ void mrope_kernel(
+
+using bfloat16 = __nv_bfloat16;
+using bf16_2 = __nv_bfloat162;
+
+__device__ __nv_bfloat16 float_to_bf16(float f) {
+    return __float2bfloat16(f);
+}
+
+__device__ float bf16_to_float(__nv_bfloat16 h) {
+    return __bfloat162float(h);
+}
+
+__device__ __nv_bfloat162 bf16_2_from_float2(float2 f2) {
+    return __halves2bfloat162(__float2bfloat16(f2.x), __float2bfloat16(f2.y));
+}
+
+__device__ void store_bf16x2(__nv_bfloat16* addr, __nv_bfloat162 val) {
+    ((uint32_t*)addr)[0] = __bfloat162_as_bfloat162(val);
+}
+
+__device__ float2 bf16_2_to_float2(__nv_bfloat162 val) {
+    uint32_t raw = __bfloat162_as_bfloat162(val);
+    float2 f2;
+    f2.x = __bfloat162float((bfloat16)(raw & 0xFFFF));
+    f2.y = __bfloat162float((bfloat16)((raw >> 16) & 0xFFFF));
+    return f2;
+}
+
+// Load a pair of bf16 as float2
+__device__ float2 load_cos_sin_pair(const uint16_t* cache, int pos, int pair_idx, int halfD, int max_pos) {
+    pos = pos < 0 ? 0 : (pos >= max_pos ? max_pos - 1 : pos);
+    int base = (pos * halfD + pair_idx) * 2; // 2 uint16 per pair
+    uint32_t raw = ((const uint32_t*)cache)[pos * halfD + pair_idx];
+    float2 f2;
+    f2.x = __bfloat162float((bfloat16)(raw & 0xFFFF));
+    f2.y = __bfloat162float((bfloat16)((raw >> 16) & 0xFFFF));
+    return f2;
+}
+
+// Grid: (B * num_heads, T, 1) per call
+// Each thread block handles one head, all T positions
+// Threads cooperatively load q/k values and produce rotated output
+
+template <int MAX_D>
+__global__ void mrope_q_kernel(
     const uint16_t* __restrict__ q_ptr,
+    const int64_t* __restrict__ pos_t,
+    const int64_t* __restrict__ pos_h,
+    const int64_t* __restrict__ pos_w,
+    const uint16_t* __restrict__ cos_ptr,
+    const uint16_t* __restrict__ sin_ptr,
+    uint16_t* __restrict__ q_out,
+    int B, int T, int Hq, int D,
+    int s0, int s1, int s2,
+    int max_pos,
+    int q_stride_B, int q_stride_T, int q_stride_H,
+    int q_out_stride_B, int q_out_stride_H, int q_out_stride_T
+) {
+    int b = blockIdx.x / Hq;
+    int h = blockIdx.x % Hq;
+    int t = blockIdx.y;
+
+    if (b >= B) return;
+
+    extern __shared__ int64_t s_pos[];
+    int64_t* s_pos_t = s_pos;
+    int64_t* s_pos_h = s_pos_t + T;
+    int64_t* s_pos_w = s_pos_h + T;
+
+    if (threadIdx.x < T) {
+        s_pos_t[threadIdx.x] = pos_t[b * T + threadIdx.x];
+        s_pos_h[threadIdx.x] = pos_h[b * T + threadIdx.x];
+        s_pos_w[threadIdx.x] = pos_w[b * T + threadIdx.x];
+    }
+    __syncthreads();
+
+    if (threadIdx.x >= D) return;
+
+    int64_t pt = s_pos_t[t];
+    int64_t ph = s_pos_h[t];
+    int64_t pw = s_pos_w[t];
+
+    int halfD = D / 2;
+
+    // Compute cos and sin for this position
+    float cos_v[MAX_D/2];
+    float sin_v[MAX_D/2];
+
+    #pragma unroll
+    for (int i = 0; i < MAX_D/2; ++i) {
+        if (i >= halfD) {
+            cos_v[i] = 0.0f; sin_v[i] = 0.0f;
+            continue;
+        }
+
+        // Load cos/sin for each axis
+        float2 cos_t = load_cos_sin_pair(cos_ptr, (int)pt, i, halfD, max_pos);
+        float2 cos_h = load_cos_sin_pair(cos_ptr, (int)ph, i, halfD, max_pos);
+        float2 cos_w = load_cos_sin_pair(cos_ptr, (int)pw, i, halfD, max_pos);
+
+        float2 sin_t = load_cos_sin_pair(sin_ptr, (int)pt, i, halfD, max_pos);
+        float2 sin_h = load_cos_sin_pair(sin_ptr, (int)ph, i, halfD, max_pos);
+        float2 sin_w = load_cos_sin_pair(sin_ptr, (int)pw, i, halfD, max_pos);
+
+        // Mix based on mrope_section
+        // First half
+        if (i < s0) {
+            cos_v[i] = cos_t.x;
+            sin_v[i] = sin_t.x;
+        } else if (i < s0 + s1) {
+            cos_v[i] = cos_h.x;
+            sin_v[i] = sin_h.x;
+        } else {
+            cos_v[i] = cos_w.x;
+            sin_v[i] = sin_w.x;
+        }
+
+        // Second half (pair i+halfD)
+        if (i + halfD < D) {
+            if (i < s0) {
+                cos_v[i + halfD] = cos_t.y;
+                sin_v[i + halfD] = sin_t.y;
+            } else if (i < s0 + s1) {
+                cos_v[i + halfD] = cos_h.y;
+                sin_v[i + halfD] = sin_h.y;
+            } else {
+                cos_v[i + halfD] = cos_w.y;
+                sin_v[i + halfD] = sin_w.y;
+            }
+        }
+    }
+
+    /

... (24163 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
@@ -14,49 +14,43 @@
 
 namespace {
 
-#define WARP_SIZE 32
-
 using bfloat16 = __nv_bfloat16;
-using bf16_2 = __nv_bfloat162;
-
-__device__ __nv_bfloat16 float_to_bf16(float f) {
+
+__device__ inline float bf16_to_float(bfloat16 h) {
+    return __bfloat162float(h);
+}
+
+__device__ inline bfloat16 float_to_bf16(float f) {
     return __float2bfloat16(f);
 }
 
-__device__ float bf16_to_float(__nv_bfloat16 h) {
-    return __bfloat162float(h);
-}
-
-__device__ __nv_bfloat162 bf16_2_from_float2(float2 f2) {
-    return __halves2bfloat162(__float2bfloat16(f2.x), __float2bfloat16(f2.y));
-}
-
-__device__ void store_bf16x2(__nv_bfloat16* addr, __nv_bfloat162 val) {
-    ((uint32_t*)addr)[0] = __bfloat162_as_bfloat162(val);
-}
-
-__device__ float2 bf16_2_to_float2(__nv_bfloat162 val) {
-    uint32_t raw = __bfloat162_as_bfloat162(val);
-    float2 f2;
-    f2.x = __bfloat162float((bfloat16)(raw & 0xFFFF));
-    f2.y = __bfloat162float((bfloat16)((raw >> 16) & 0xFFFF));
-    return f2;
-}
-
-// Load a pair of bf16 as float2
-__device__ float2 load_cos_sin_pair(const uint16_t* cache, int pos, int pair_idx, int halfD, int max_pos) {
+// Load pair of bf16 from (max_pos, D) cache at position pos, pair pair_idx
+// D elements total = 2*halfD bf16 values
+// Layout: (max_pos, 2*halfD) uint16_t = (max_pos, D) bf16
+// Pair i contains elements [2*i] and [2*i+1]
+__device__ inline void load_cos_sin_pair(
+    const uint16_t* __restrict__ cos_ptr,
+    const uint16_t* __restrict__ sin_ptr,
+    int pos, int pair_idx, int halfD, int max_pos,
+    float& cos_lo, float& cos_hi, float& sin_lo, float& sin_hi
+) {
     pos = pos < 0 ? 0 : (pos >= max_pos ? max_pos - 1 : pos);
-    int base = (pos * halfD + pair_idx) * 2; // 2 uint16 per pair
-    uint32_t raw = ((const uint32_t*)cache)[pos * halfD + pair_idx];
-    float2 f2;
-    f2.x = __bfloat162float((bfloat16)(raw & 0xFFFF));
-    f2.y = __bfloat162float((bfloat16)((raw >> 16) & 0xFFFF));
-    return f2;
-}
-
-// Grid: (B * num_heads, T, 1) per call
-// Each thread block handles one head, all T positions
-// Threads cooperatively load q/k values and produce rotated output
+    int base = pos * D_ + pair_idx * 2;
+    uint32_t raw_cos = ((const uint32_t*)(cos_ptr + base))[0];
+    uint32_t raw_sin = ((const uint32_t*)(sin_ptr + base))[0];
+    cos_lo = __bfloat162float((bfloat16)(raw_cos & 0xFFFF));
+    cos_hi = __bfloat162float((bfloat16)((raw_cos >> 16) & 0xFFFF));
+    sin_lo = __bfloat162float((bfloat16)(raw_sin & 0xFFFF));
+    sin_hi = __bfloat162float((bfloat16)((raw_sin >> 16) & 0xFFFF));
+}
+
+__device__ inline float load_bf16(const uint16_t* ptr, int offset) {
+    return __bfloat162float((bfloat16)ptr[offset]);
+}
+
+__device__ inline void store_bf16(uint16_t* ptr, int offset, float val) {
+    ptr[offset] = (uint16_t)__float2bfloat16(val);
+}
 
 template <int MAX_D>
 __global__ void mrope_q_kernel(
@@ -81,7 +75,7 @@
 
     extern __shared__ int64_t s_pos[];
     int64_t* s_pos_t = s_pos;
-    int64_t* s_pos_h = s_pos_t + T;
+    int64_t* s_pos_h = s_pos + T;
     int64_t* s_pos_w = s_pos_h + T;
 
     if (threadIdx.x < T) {
@@ -91,7 +85,8 @@
     }
     __syncthreads();
 
-    if (threadIdx.x >= D) return;
+    int elem_idx = threadIdx.x;
+    if (elem_idx >= D) return;
 
     int64_t pt = s_pos_t[t];
     int64_t ph = s_pos_h[t];
@@ -99,81 +94,88 @@
 
     int halfD = D / 2;
 
-    // Compute cos and sin for this position
-    float cos_v[MAX_D/2];
-    float sin_v[MAX_D/2];
-
-    #pragma unroll
-    for (int i = 0; i < MAX_D/2; ++i) {
-        if (i >= halfD) {
-            cos_v[i] = 0.0f; sin_v[i] = 0.0f;
-            continue;
-        }
-
-        // Load cos/sin for each axis
-        float2 cos_t = load_cos_sin_pair(cos_ptr, (int)pt, i, halfD, max_pos);
-        float2 cos_h = load_cos_sin_pair(cos_ptr, (int)ph, i, halfD, max_pos);
-        float2 cos_w = load_cos_sin_pair(cos_ptr, (int)pw, i, halfD, max_pos);
-
-        float2 sin_t = load_cos_sin_pair(sin_ptr, (int)pt, i, halfD, max_pos);
-        float2 sin_h = load_cos_sin_pair(sin_ptr, (int)ph, i, halfD, max_pos);
-        float2 sin_w = load_cos_sin_pair(sin_ptr, (int)pw, i, halfD, max_pos);
+    // Compute cos and sin for this element
+    int pair_idx = elem_idx >> 1;
+    int is_hi = elem_idx & 1;
+
+    float c_val, s_val;
+
+    if (pair_idx < halfD) {
+        float cos_lo_t, cos_hi_t, sin_lo_t, sin_hi_t;
+        float cos_lo_h, cos_hi_h, sin_lo_h, sin_hi_h;
+        float cos_lo_w, cos_hi_w, sin_lo_w, sin_hi_w;
+
+        // Load pairs for each axis - use temp arrays to avoid register pressure
+        int p_t = (int)(pt < 0 ? 0 : (pt >= max_pos ? max_pos - 1 : pt));
+        int p_h = (int)(ph < 0 ? 0 : (ph >= max_pos ? max_pos - 1 : ph));
+        int p_w = (int)(pw < 0 ? 0 : (pw >= max_pos ? max_pos - 1 : pw));
+
+        // cos_t[p, 2*i] and cos_t[p, 2*i+1]
+        uint32_t raw_cos_t = ((const uint32_t*)(cos_ptr + p_t * D + pair_idx * 2))[0];
+        uint32_t raw_sin_t = ((const uint32_t*)(sin_ptr + p_t * D + pair_idx * 2))[0];
+        cos_lo_t = __bfloat162float((bfloat16)(raw_cos_t & 0xFFFF));
+        cos_hi_t = __bfloat162float((bfloat16)((raw_cos_t >> 16) & 0xFFFF));
+        sin_lo_t = __bfloat162float((bfloat16)(raw_sin_t & 0xFFFF));
+        sin_hi_t = __bfloat162float((bfloat16)((raw_sin_t >> 16) & 0xFFFF));
+
+        uint32_t raw_cos_h = ((const uint32_t*)(cos_ptr + p_h * D + pair_idx * 2))[0];
+        uint32_t raw_sin_h = ((const uint32_t*)(sin_ptr + p_h * D + pair_idx * 2))[0];
+        cos_lo_h = __bfloat162float((bfloat16)(raw_cos_h & 0xFFFF));
+        cos_hi_h = __bfloat162float((bfloat16)((raw_cos_h >> 16) & 0xFFFF));
+        sin_lo_h = __bfloat162float((bfloat16)(raw_sin_h & 0xFFFF));
+        sin_hi_h = __bfloat162float((bfloat16)((raw_sin_h >> 16) & 0xFFFF));
+
+        

... (14362 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
@@ -7,49 +7,42 @@
 import torch
 import torch.utils.cpp_extension
 
-_CC = """
+_CC = r"""
 #include <torch/extension.h>
 #include <stdio.h>
 #include <cuda_fp16.h>
 
 namespace {
 
-using bfloat16 = __nv_bfloat16;
-
-__device__ inline float bf16_to_float(bfloat16 h) {
+using bf16 = __nv_bfloat16;
+using bf16_2 = __nv_bfloat162;
+
+__device__ inline float bf16_to_float(bf16 h) {
     return __bfloat162float(h);
 }
 
-__device__ inline bfloat16 float_to_bf16(float f) {
+__device__ inline bf16 float_to_bf16(float f) {
     return __float2bfloat16(f);
 }
 
-// Load pair of bf16 from (max_pos, D) cache at position pos, pair pair_idx
-// D elements total = 2*halfD bf16 values
-// Layout: (max_pos, 2*halfD) uint16_t = (max_pos, D) bf16
-// Pair i contains elements [2*i] and [2*i+1]
-__device__ inline void load_cos_sin_pair(
-    const uint16_t* __restrict__ cos_ptr,
-    const uint16_t* __restrict__ sin_ptr,
-    int pos, int pair_idx, int halfD, int max_pos,
-    float& cos_lo, float& cos_hi, float& sin_lo, float& sin_hi
-) {
-    pos = pos < 0 ? 0 : (pos >= max_pos ? max_pos - 1 : pos);
-    int base = pos * D_ + pair_idx * 2;
-    uint32_t raw_cos = ((const uint32_t*)(cos_ptr + base))[0];
-    uint32_t raw_sin = ((const uint32_t*)(sin_ptr + base))[0];
-    cos_lo = __bfloat162float((bfloat16)(raw_cos & 0xFFFF));
-    cos_hi = __bfloat162float((bfloat16)((raw_cos >> 16) & 0xFFFF));
-    sin_lo = __bfloat162float((bfloat16)(raw_sin & 0xFFFF));
-    sin_hi = __bfloat162float((bfloat16)((raw_sin >> 16) & 0xFFFF));
-}
-
-__device__ inline float load_bf16(const uint16_t* ptr, int offset) {
-    return __bfloat162float((bfloat16)ptr[offset]);
-}
-
-__device__ inline void store_bf16(uint16_t* ptr, int offset, float val) {
-    ptr[offset] = (uint16_t)__float2bfloat16(val);
+__device__ inline bf16 load_bf16(const uint16_t* ptr, int offset) {
+    return ((const bf16*)ptr)[offset];
+}
+
+__device__ inline void store_bf16(uint16_t* ptr, int offset, bf16 val) {
+    ((bf16*)ptr)[offset] = val;
+}
+
+__device__ inline void store_bf16_as_uint16(uint16_t* ptr, int offset, bf16 val) {
+    ptr[offset] = (uint16_t)__bfloat162uint16(val);
+}
+
+__device__ inline uint32_t bf16_2_to_uint(bf16_2 val) {
+    return __bfloat162_as_bfloat162(val);
+}
+
+__device__ inline bf16_2 uint_to_bf16_2(uint32_t val) {
+    return __uint2bfloat162_rn(val);
 }
 
 template <int MAX_D>
@@ -93,72 +86,55 @@
     int64_t pw = s_pos_w[t];
 
     int halfD = D / 2;
-
-    // Compute cos and sin for this element
     int pair_idx = elem_idx >> 1;
     int is_hi = elem_idx & 1;
 
-    float c_val, s_val;
+    float c_val = 0.0f, s_val = 0.0f;
 
     if (pair_idx < halfD) {
-        float cos_lo_t, cos_hi_t, sin_lo_t, sin_hi_t;
-        float cos_lo_h, cos_hi_h, sin_lo_h, sin_hi_h;
-        float cos_lo_w, cos_hi_w, sin_lo_w, sin_hi_w;
-
-        // Load pairs for each axis - use temp arrays to avoid register pressure
         int p_t = (int)(pt < 0 ? 0 : (pt >= max_pos ? max_pos - 1 : pt));
         int p_h = (int)(ph < 0 ? 0 : (ph >= max_pos ? max_pos - 1 : ph));
         int p_w = (int)(pw < 0 ? 0 : (pw >= max_pos ? max_pos - 1 : pw));
 
-        // cos_t[p, 2*i] and cos_t[p, 2*i+1]
-        uint32_t raw_cos_t = ((const uint32_t*)(cos_ptr + p_t * D + pair_idx * 2))[0];
-        uint32_t raw_sin_t = ((const uint32_t*)(sin_ptr + p_t * D + pair_idx * 2))[0];
-        cos_lo_t = __bfloat162float((bfloat16)(raw_cos_t & 0xFFFF));
-        cos_hi_t = __bfloat162float((bfloat16)((raw_cos_t >> 16) & 0xFFFF));
-        sin_lo_t = __bfloat162float((bfloat16)(raw_sin_t & 0xFFFF));
-        sin_hi_t = __bfloat162float((bfloat16)((raw_sin_t >> 16) & 0xFFFF));
-
-        uint32_t raw_cos_h = ((const uint32_t*)(cos_ptr + p_h * D + pair_idx * 2))[0];
-        uint32_t raw_sin_h = ((const uint32_t*)(sin_ptr + p_h * D + pair_idx * 2))[0];
-        cos_lo_h = __bfloat162float((bfloat16)(raw_cos_h & 0xFFFF));
-        cos_hi_h = __bfloat162float((bfloat16)((raw_cos_h >> 16) & 0xFFFF));
-        sin_lo_h = __bfloat162float((bfloat16)(raw_sin_h & 0xFFFF));
-        sin_hi_h = __bfloat162float((bfloat16)((raw_sin_h >> 16) & 0xFFFF));
-
-        uint32_t raw_cos_w = ((const uint32_t*)(cos_ptr + p_w * D + pair_idx * 2))[0];
-        uint32_t raw_sin_w = ((const uint32_t*)(sin_ptr + p_w * D + pair_idx * 2))[0];
-        cos_lo_w = __bfloat162float((bfloat16)(raw_cos_w & 0xFFFF));
-        cos_hi_w = __bfloat162float((bfloat16)((raw_cos_w >> 16) & 0xFFFF));
-        sin_lo_w = __bfloat162float((bfloat16)(raw_sin_w & 0xFFFF));
-        sin_hi_w = __bfloat162float((bfloat16)((raw_sin_w >> 16) & 0xFFFF));
-
-        // Mix based on mrope_section
-        // First half: pairs 0..halfD-1
-        // Second half: pairs halfD..D-1 (mirror)
+        uint32_t raw_cos_t = ((const uint32_t*)(cos_ptr + p_t * D))[pair_idx];
+        uint32_t raw_sin_t = ((const uint32_t*)(sin_ptr + p_t * D))[pair_idx];
+        uint32_t raw_cos_h = ((const uint32_t*)(cos_ptr + p_h * D))[pair_idx];
+        uint32_t raw_sin_h = ((const uint32_t*)(sin_ptr + p_h * D))[pair_idx];
+        uint32_t raw_cos_w = ((const uint32_t*)(cos_ptr + p_w * D))[pair_idx];
+        uint32_t raw_sin_w = ((const uint32_t*)(sin_ptr + p_w * D))[pair_idx];
+
+        bf16 cos_lo_t = __uint2bfloat16(raw_cos_t & 0xFFFF);
+        bf16 cos_hi_t = __uint2bfloat16((raw_cos_t >> 16) & 0xFFFF);
+        bf16 sin_lo_t = __uint2bfloat16(raw_sin_t & 0xFFFF);
+        bf16 sin_hi_t = __uint2bfloat16((raw_sin_t >> 16) & 0xFFFF);
+
+        bf16 cos_lo_h = __uint2bfloat16(raw_cos_h & 0xFFFF);
+        bf16 cos_hi_h = __uint2bfloat16((raw_cos_h >> 16) & 0xFFFF);
+        bf16 sin_lo_h = __uint2bfloat16(raw_sin_h & 0xFFFF);
+        bf16 sin_hi_h = __uint2bfloat16((raw_sin_h >> 16) & 0xFFFF);
+
+        bf16 cos_lo_w = __uint2bfloat16(raw_cos_w & 0xFFFF);
+

... (11260 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
@@ -10,7 +10,7 @@
 _CC = r"""
 #include <torch/extension.h>
 #include <stdio.h>
-#include <cuda_fp16.h>
+#include <cuda_bf16.h>
 
 namespace {
 
@@ -25,24 +25,20 @@
     return __float2bfloat16(f);
 }
 
+__device__ inline bf16 uint16_to_bf16(unsigned short val) {
+    return __uint2bfloat16_rn(val);
+}
+
+__device__ inline unsigned short bf16_to_uint16(bf16 h) {
+    return __bfloat162uint16(h);
+}
+
 __device__ inline bf16 load_bf16(const uint16_t* ptr, int offset) {
     return ((const bf16*)ptr)[offset];
 }
 
 __device__ inline void store_bf16(uint16_t* ptr, int offset, bf16 val) {
     ((bf16*)ptr)[offset] = val;
-}
-
-__device__ inline void store_bf16_as_uint16(uint16_t* ptr, int offset, bf16 val) {
-    ptr[offset] = (uint16_t)__bfloat162uint16(val);
-}
-
-__device__ inline uint32_t bf16_2_to_uint(bf16_2 val) {
-    return __bfloat162_as_bfloat162(val);
-}
-
-__device__ inline bf16_2 uint_to_bf16_2(uint32_t val) {
-    return __uint2bfloat162_rn(val);
 }
 
 template <int MAX_D>
@@ -103,37 +99,37 @@
         uint32_t raw_cos_w = ((const uint32_t*)(cos_ptr + p_w * D))[pair_idx];
         uint32_t raw_sin_w = ((const uint32_t*)(sin_ptr + p_w * D))[pair_idx];
 
-        bf16 cos_lo_t = __uint2bfloat16(raw_cos_t & 0xFFFF);
-        bf16 cos_hi_t = __uint2bfloat16((raw_cos_t >> 16) & 0xFFFF);
-        bf16 sin_lo_t = __uint2bfloat16(raw_sin_t & 0xFFFF);
-        bf16 sin_hi_t = __uint2bfloat16((raw_sin_t >> 16) & 0xFFFF);
-
-        bf16 cos_lo_h = __uint2bfloat16(raw_cos_h & 0xFFFF);
-        bf16 cos_hi_h = __uint2bfloat16((raw_cos_h >> 16) & 0xFFFF);
-        bf16 sin_lo_h = __uint2bfloat16(raw_sin_h & 0xFFFF);
-        bf16 sin_hi_h = __uint2bfloat16((raw_sin_h >> 16) & 0xFFFF);
-
-        bf16 cos_lo_w = __uint2bfloat16(raw_cos_w & 0xFFFF);
-        bf16 cos_hi_w = __uint2bfloat16((raw_cos_w >> 16) & 0xFFFF);
-        bf16 sin_lo_w = __uint2bfloat16(raw_sin_w & 0xFFFF);
-        bf16 sin_hi_w = __uint2bfloat16((raw_sin_w >> 16) & 0xFFFF);
+        bf16 cos_lo_t = uint16_to_bf16(raw_cos_t & 0xFFFF);
+        bf16 cos_hi_t = uint16_to_bf16((raw_cos_t >> 16) & 0xFFFF);
+        bf16 sin_lo_t = uint16_to_bf16(raw_sin_t & 0xFFFF);
+        bf16 sin_hi_t = uint16_to_bf16((raw_sin_t >> 16) & 0xFFFF);
+
+        bf16 cos_lo_h = uint16_to_bf16(raw_cos_h & 0xFFFF);
+        bf16 cos_hi_h = uint16_to_bf16((raw_cos_h >> 16) & 0xFFFF);
+        bf16 sin_lo_h = uint16_to_bf16(raw_sin_h & 0xFFFF);
+        bf16 sin_hi_h = uint16_to_bf16((raw_sin_h >> 16) & 0xFFFF);
+
+        bf16 cos_lo_w = uint16_to_bf16(raw_cos_w & 0xFFFF);
+        bf16 cos_hi_w = uint16_to_bf16((raw_cos_w >> 16) & 0xFFFF);
+        bf16 sin_lo_w = uint16_to_bf16(raw_sin_w & 0xFFFF);
+        bf16 sin_hi_w = uint16_to_bf16((raw_sin_w >> 16) & 0xFFFF);
 
         int section_idx = pair_idx < s0 ? 0 : (pair_idx < s0 + s1 ? 1 : 2);
 
         if (section_idx == 0) {
-            c_val = is_hi ? __bfloat162float(cos_hi_t) : __bfloat162float(cos_lo_t);
-            s_val = is_hi ? __bfloat162float(sin_hi_t) : __bfloat162float(sin_lo_t);
+            c_val = is_hi ? bf16_to_float(cos_hi_t) : bf16_to_float(cos_lo_t);
+            s_val = is_hi ? bf16_to_float(sin_hi_t) : bf16_to_float(sin_lo_t);
         } else if (section_idx == 1) {
-            c_val = is_hi ? __bfloat162float(cos_hi_h) : __bfloat162float(cos_lo_h);
-            s_val = is_hi ? __bfloat162float(sin_hi_h) : __bfloat162float(sin_lo_h);
+            c_val = is_hi ? bf16_to_float(cos_hi_h) : bf16_to_float(cos_lo_h);
+            s_val = is_hi ? bf16_to_float(sin_hi_h) : bf16_to_float(sin_lo_h);
         } else {
-            c_val = is_hi ? __bfloat162float(cos_hi_w) : __bfloat162float(cos_lo_w);
-            s_val = is_hi ? __bfloat162float(sin_hi_w) : __bfloat162float(sin_lo_w);
+            c_val = is_hi ? bf16_to_float(cos_hi_w) : bf16_to_float(cos_lo_w);
+            s_val = is_hi ? bf16_to_float(sin_hi_w) : bf16_to_float(sin_lo_w);
         }
     }
 
     int q_offset = b * q_stride_B + t * q_stride_T + h * q_stride_H + elem_idx;
-    float x = __bfloat162float(load_bf16(q_ptr, q_offset));
+    float x = bf16_to_float(load_bf16(q_ptr, q_offset));
 
     int rot_idx;
     float rot_sign;
@@ -146,12 +142,12 @@
     }
 
     int rot_offset = b * q_stride_B + t * q_stride_T + h * q_stride_H + rot_idx;
-    float rot_x = __bfloat162float(load_bf16(q_ptr, rot_offset)) * rot_sign;
+    float rot_x = bf16_to_float(load_bf16(q_ptr, rot_offset)) * rot_sign;
 
     float out = x * c_val + rot_x * s_val;
 
     int out_offset = b * q_out_stride_B + h * q_out_stride_H + t * q_out_stride_T + elem_idx;
-    store_bf16_as_uint16(q_out, out_offset, __float2bfloat16(out));
+    store_bf16(q_out, out_offset, float_to_bf16(out));
 }
 
 template <int MAX_D>
@@ -212,37 +208,37 @@
         uint32_t raw_cos_w = ((const uint32_t*)(cos_ptr + p_w * D))[pair_idx];
         uint32_t raw_sin_w = ((const uint32_t*)(sin_ptr + p_w * D))[pair_idx];
 
-        bf16 cos_lo_t = __uint2bfloat16(raw_cos_t & 0xFFFF);
-        bf16 cos_hi_t = __uint2bfloat16((raw_cos_t >> 16) & 0xFFFF);
-        bf16 sin_lo_t = __uint2bfloat16(raw_sin_t & 0xFFFF);
-        bf16 sin_hi_t = __uint2bfloat16((raw_sin_t >> 16) & 0xFFFF);
-
-        bf16 cos_lo_h = __uint2bfloat16(raw_cos_h & 0xFFFF);
-        bf16 cos_hi_h = __uint2bfloat16((raw_cos_h >> 16) & 0xFFFF);
-        bf16 sin_lo_h = __uint2bfloat16(raw_sin_h & 0xFFFF);
-        bf16 sin_hi_h = __uint2bfloat16((raw_sin_h >> 16) & 0xFFFF);
-
-        bf16 cos_lo_w = __uint2bfloat16(raw_cos_w & 0xFFFF);
-        bf16 cos_hi_w = __uint2bfloat16((raw_cos_w >> 16) & 0xFFFF);
-        bf16 sin_lo_w = __uint2bfloat16(raw_sin_w & 0xFFFF);
-        bf16 sin_hi_w = __uint2bfloat16((raw_sin_w >> 16) & 0xFFFF);
+        bf16 cos_lo_t = uint16_to_bf16(raw_cos_t & 0xFFFF);
+        bf16 cos_hi_t = uint16_t

... (3033 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
@@ -15,22 +15,21 @@
 namespace {
 
 using bf16 = __nv_bfloat16;
-using bf16_2 = __nv_bfloat162;
 
 __device__ inline float bf16_to_float(bf16 h) {
     return __bfloat162float(h);
 }
 
 __device__ inline bf16 float_to_bf16(float f) {
-    return __float2bfloat16(f);
+    return __float2bfloat16_rn(f);
+}
+
+__device__ inline unsigned short bf16_to_uint16(bf16 h) {
+    return (unsigned short)__bfloat162int_rn(h);
 }
 
 __device__ inline bf16 uint16_to_bf16(unsigned short val) {
-    return __uint2bfloat16_rn(val);
-}
-
-__device__ inline unsigned short bf16_to_uint16(bf16 h) {
-    return __bfloat162uint16(h);
+    return __float2bfloat16_rn(__int2float_rn(val));
 }
 
 __device__ inline bf16 load_bf16(const uint16_t* ptr, int offset) {
@@ -400,12 +399,12 @@
 }
 """
 
-_build_dir = "/tmp/mrope_build_v5"
+_build_dir = "/tmp/mrope_build_v6"
 os.makedirs(_build_dir, exist_ok=True)
 
 try:
     _mrope_cpp = torch.utils.cpp_extension.load_inline(
-        name="mrope_fused_v5",
+        name="mrope_fused_v6",
         cpp_sources="",
         cuda_sources=_CC,
         functions=["mrope_fwd"],

--- 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
@@ -258,7 +258,7 @@
     store_bf16(k_out, out_offset, float_to_bf16(out));
 }
 
-} // namespace
+} // namespace mrope_ns
 
 std::tuple<torch::Tensor, torch::Tensor> mrope_fwd(
     torch::Tensor q,
@@ -269,130 +269,7 @@
     torch::Tensor cos_cache,
     torch::Tensor sin_cache,
     std::tuple<int, int, int> mrope_section
-) {
-    int B = q.size(0);
-    int T = q.size(1);
-    int Hq = q.size(2);
-    int D = q.size(3);
-    int Hkv = k.size(2);
-    int max_pos = cos_cache.size(0);
-
-    int s0 = std::get<0>(mrope_section);
-    int s1 = std::get<1>(mrope_section);
-    int s2 = std::get<2>(mrope_section);
-
-    auto q_out = torch::empty({B, Hq, T, D}, q.options());
-    auto k_out = torch::empty({B, Hkv, T, D}, k.options());
-
-    int q_stride_B = q.stride(0);
-    int q_stride_T = q.stride(1);
-    int q_stride_H = q.stride(2);
-    int k_stride_B = k.stride(0);
-    int k_stride_T = k.stride(1);
-    int k_stride_H = k.stride(2);
-
-    int q_out_stride_B = q_out.stride(0);
-    int q_out_stride_H = q_out.stride(1);
-    int q_out_stride_T = q_out.stride(2);
-    int k_out_stride_B = k_out.stride(0);
-    int k_out_stride_H = k_out.stride(1);
-    int k_out_stride_T = k_out.stride(2);
-
-    int smem = T * 3 * sizeof(int64_t);
-
-    if (D <= 64) {
-        dim3 grid_q(B * Hq, T, 1);
-        mrope_q_kernel<64><<<grid_q, 64, smem>>>(
-            (const uint16_t*)q.data_ptr(),
-            pos_t.data_ptr<int64_t>(),
-            pos_h.data_ptr<int64_t>(),
-            pos_w.data_ptr<int64_t>(),
-            (const uint16_t*)cos_cache.data_ptr(),
-            (const uint16_t*)sin_cache.data_ptr(),
-            (uint16_t*)q_out.data_ptr(),
-            B, T, Hq, D, s0, s1, s2, max_pos,
-            q_stride_B, q_stride_T, q_stride_H,
-            q_out_stride_B, q_out_stride_H, q_out_stride_T
-        );
-
-        dim3 grid_k(B * Hkv, T, 1);
-        mrope_k_kernel<64><<<grid_k, 64, smem>>>(
-            (const uint16_t*)k.data_ptr(),
-            pos_t.data_ptr<int64_t>(),
-            pos_h.data_ptr<int64_t>(),
-            pos_w.data_ptr<int64_t>(),
-            (const uint16_t*)cos_cache.data_ptr(),
-            (const uint16_t*)sin_cache.data_ptr(),
-            (uint16_t*)k_out.data_ptr(),
-            B, T, Hkv, D, s0, s1, s2, max_pos,
-            k_stride_B, k_stride_T, k_stride_H,
-            k_out_stride_B, k_out_stride_H, k_out_stride_T
-        );
-    } else if (D <= 128) {
-        dim3 grid_q(B * Hq, T, 1);
-        mrope_q_kernel<128><<<grid_q, 128, smem>>>(
-            (const uint16_t*)q.data_ptr(),
-            pos_t.data_ptr<int64_t>(),
-            pos_h.data_ptr<int64_t>(),
-            pos_w.data_ptr<int64_t>(),
-            (const uint16_t*)cos_cache.data_ptr(),
-            (const uint16_t*)sin_cache.data_ptr(),
-            (uint16_t*)q_out.data_ptr(),
-            B, T, Hq, D, s0, s1, s2, max_pos,
-            q_stride_B, q_stride_T, q_stride_H,
-            q_out_stride_B, q_out_stride_H, q_out_stride_T
-        );
-
-        dim3 grid_k(B * Hkv, T, 1);
-        mrope_k_kernel<128><<<grid_k, 128, smem>>>(
-            (const uint16_t*)k.data_ptr(),
-            pos_t.data_ptr<int64_t>(),
-            pos_h.data_ptr<int64_t>(),
-            pos_w.data_ptr<int64_t>(),
-            (const uint16_t*)cos_cache.data_ptr(),
-            (const uint16_t*)sin_cache.data_ptr(),
-            (uint16_t*)k_out.data_ptr(),
-            B, T, Hkv, D, s0, s1, s2, max_pos,
-            k_stride_B, k_stride_T, k_stride_H,
-            k_out_stride_B, k_out_stride_H, k_out_stride_T
-        );
-    } else {
-        dim3 grid_q(B * Hq, T, 1);
-        mrope_q_kernel<256><<<grid_q, 256, smem>>>(
-            (const uint16_t*)q.data_ptr(),
-            pos_t.data_ptr<int64_t>(),
-            pos_h.data_ptr<int64_t>(),
-            pos_w.data_ptr<int64_t>(),
-            (const uint16_t*)cos_cache.data_ptr(),
-            (const uint16_t*)sin_cache.data_ptr(),
-            (uint16_t*)q_out.data_ptr(),
-            B, T, Hq, D, s0, s1, s2, max_pos,
-            q_stride_B, q_stride_T, q_stride_H,
-            q_out_stride_B, q_out_stride_H, q_out_stride_T
-        );
-
-        dim3 grid_k(B * Hkv, T, 1);
-        mrope_k_kernel<256><<<grid_k, 256, smem>>>(
-            (const uint16_t*)k.data_ptr(),
-            pos_t.data_ptr<int64_t>(),
-            pos_h.data_ptr<int64_t>(),
-            pos_w.data_ptr<int64_t>(),
-            (const uint16_t*)cos_cache.data_ptr(),
-            (const uint16_t*)sin_cache.data_ptr(),
-            (uint16_t*)k_out.data_ptr(),
-            B, T, Hkv, D, s0, s1, s2, max_pos,
-            k_stride_B, k_stride_T, k_stride_H,
-            k_out_stride_B, k_out_stride_H, k_out_stride_T
-        );
-    }
-
-    cudaError_t err = cudaGetLastError();
-    if (err != cudaSuccess) {
-        printf("CUDA error: %s\n", cudaGetErrorString(err));
-    }
-
-    return std::make_tuple(q_out, k_out);
-}
+);
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.def("mrope_fwd", &mrope_fwd, "Fused multi-axis RoPE forward");

--- 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
@@ -269,7 +269,130 @@
     torch::Tensor cos_cache,
     torch::Tensor sin_cache,
     std::tuple<int, int, int> mrope_section
-);
+) {
+    int B = q.size(0);
+    int T = q.size(1);
+    int Hq = q.size(2);
+    int D = q.size(3);
+    int Hkv = k.size(2);
+    int max_pos = cos_cache.size(0);
+
+    int s0 = std::get<0>(mrope_section);
+    int s1 = std::get<1>(mrope_section);
+    int s2 = std::get<2>(mrope_section);
+
+    auto q_out = torch::empty({B, Hq, T, D}, q.options());
+    auto k_out = torch::empty({B, Hkv, T, D}, k.options());
+
+    int q_stride_B = q.stride(0);
+    int q_stride_T = q.stride(1);
+    int q_stride_H = q.stride(2);
+    int k_stride_B = k.stride(0);
+    int k_stride_T = k.stride(1);
+    int k_stride_H = k.stride(2);
+
+    int q_out_stride_B = q_out.stride(0);
+    int q_out_stride_H = q_out.stride(1);
+    int q_out_stride_T = q_out.stride(2);
+    int k_out_stride_B = k_out.stride(0);
+    int k_out_stride_H = k_out.stride(1);
+    int k_out_stride_T = k_out.stride(2);
+
+    int smem = T * 3 * sizeof(int64_t);
+
+    if (D <= 64) {
+        dim3 grid_q(B * Hq, T, 1);
+        mrope_q_kernel<64><<<grid_q, 64, smem>>>(
+            (const uint16_t*)q.data_ptr(),
+            pos_t.data_ptr<int64_t>(),
+            pos_h.data_ptr<int64_t>(),
+            pos_w.data_ptr<int64_t>(),
+            (const uint16_t*)cos_cache.data_ptr(),
+            (const uint16_t*)sin_cache.data_ptr(),
+            (uint16_t*)q_out.data_ptr(),
+            B, T, Hq, D, s0, s1, s2, max_pos,
+            q_stride_B, q_stride_T, q_stride_H,
+            q_out_stride_B, q_out_stride_H, q_out_stride_T
+        );
+
+        dim3 grid_k(B * Hkv, T, 1);
+        mrope_k_kernel<64><<<grid_k, 64, smem>>>(
+            (const uint16_t*)k.data_ptr(),
+            pos_t.data_ptr<int64_t>(),
+            pos_h.data_ptr<int64_t>(),
+            pos_w.data_ptr<int64_t>(),
+            (const uint16_t*)cos_cache.data_ptr(),
+            (const uint16_t*)sin_cache.data_ptr(),
+            (uint16_t*)k_out.data_ptr(),
+            B, T, Hkv, D, s0, s1, s2, max_pos,
+            k_stride_B, k_stride_T, k_stride_H,
+            k_out_stride_B, k_out_stride_H, k_out_stride_T
+        );
+    } else if (D <= 128) {
+        dim3 grid_q(B * Hq, T, 1);
+        mrope_q_kernel<128><<<grid_q, 128, smem>>>(
+            (const uint16_t*)q.data_ptr(),
+            pos_t.data_ptr<int64_t>(),
+            pos_h.data_ptr<int64_t>(),
+            pos_w.data_ptr<int64_t>(),
+            (const uint16_t*)cos_cache.data_ptr(),
+            (const uint16_t*)sin_cache.data_ptr(),
+            (uint16_t*)q_out.data_ptr(),
+            B, T, Hq, D, s0, s1, s2, max_pos,
+            q_stride_B, q_stride_T, q_stride_H,
+            q_out_stride_B, q_out_stride_H, q_out_stride_T
+        );
+
+        dim3 grid_k(B * Hkv, T, 1);
+        mrope_k_kernel<128><<<grid_k, 128, smem>>>(
+            (const uint16_t*)k.data_ptr(),
+            pos_t.data_ptr<int64_t>(),
+            pos_h.data_ptr<int64_t>(),
+            pos_w.data_ptr<int64_t>(),
+            (const uint16_t*)cos_cache.data_ptr(),
+            (const uint16_t*)sin_cache.data_ptr(),
+            (uint16_t*)k_out.data_ptr(),
+            B, T, Hkv, D, s0, s1, s2, max_pos,
+            k_stride_B, k_stride_T, k_stride_H,
+            k_out_stride_B, k_out_stride_H, k_out_stride_T
+        );
+    } else {
+        dim3 grid_q(B * Hq, T, 1);
+        mrope_q_kernel<256><<<grid_q, 256, smem>>>(
+            (const uint16_t*)q.data_ptr(),
+            pos_t.data_ptr<int64_t>(),
+            pos_h.data_ptr<int64_t>(),
+            pos_w.data_ptr<int64_t>(),
+            (const uint16_t*)cos_cache.data_ptr(),
+            (const uint16_t*)sin_cache.data_ptr(),
+            (uint16_t*)q_out.data_ptr(),
+            B, T, Hq, D, s0, s1, s2, max_pos,
+            q_stride_B, q_stride_T, q_stride_H,
+            q_out_stride_B, q_out_stride_H, q_out_stride_T
+        );
+
+        dim3 grid_k(B * Hkv, T, 1);
+        mrope_k_kernel<256><<<grid_k, 256, smem>>>(
+            (const uint16_t*)k.data_ptr(),
+            pos_t.data_ptr<int64_t>(),
+            pos_h.data_ptr<int64_t>(),
+            pos_w.data_ptr<int64_t>(),
+            (const uint16_t*)cos_cache.data_ptr(),
+            (const uint16_t*)sin_cache.data_ptr(),
+            (uint16_t*)k_out.data_ptr(),
+            B, T, Hkv, D, s0, s1, s2, max_pos,
+            k_stride_B, k_stride_T, k_stride_H,
+            k_out_stride_B, k_out_stride_H, k_out_stride_T
+        );
+    }
+
+    cudaError_t err = cudaGetLastError();
+    if (err != cudaSuccess) {
+        printf("CUDA error: %s\n", cudaGetErrorString(err));
+    }
+
+    return std::make_tuple(q_out, k_out);
+}
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.def("mrope_fwd", &mrope_fwd, "Fused multi-axis RoPE forward");

--- 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
@@ -12,8 +12,6 @@
 #include <stdio.h>
 #include <cuda_bf16.h>
 
-namespace {
-
 using bf16 = __nv_bfloat16;
 
 __device__ inline float bf16_to_float(bf16 h) {
@@ -22,10 +20,6 @@
 
 __device__ inline bf16 float_to_bf16(float f) {
     return __float2bfloat16_rn(f);
-}
-
-__device__ inline unsigned short bf16_to_uint16(bf16 h) {
-    return (unsigned short)__bfloat162int_rn(h);
 }
 
 __device__ inline bf16 uint16_to_bf16(unsigned short val) {
@@ -258,8 +252,6 @@
     store_bf16(k_out, out_offset, float_to_bf16(out));
 }
 
-} // namespace mrope_ns
-
 std::tuple<torch::Tensor, torch::Tensor> mrope_fwd(
     torch::Tensor q,
     torch::Tensor k,
@@ -399,12 +391,12 @@
 }
 """
 
-_build_dir = "/tmp/mrope_build_v6"
+_build_dir = "/tmp/mrope_build_v7"
 os.makedirs(_build_dir, exist_ok=True)
 
 try:
     _mrope_cpp = torch.utils.cpp_extension.load_inline(
-        name="mrope_fused_v6",
+        name="mrope_fused_v7",
         cpp_sources="",
         cuda_sources=_CC,
         functions=["mrope_fwd"],

--- 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,4 @@
-"""Fused multi-axis RoPE pre-attention prep.
-
-Fuses: gather per-axis cos/sin from cache, mrope_section mixing, rotate-half,
-and (B,T,H,D)->(B,H,T,D) transpose into a single stream of q and k.
-"""
+"""Fused multi-axis RoPE pre-attention prep."""
 import os
 import torch
 import torch.utils.cpp_extension
@@ -386,17 +382,149 @@
     return std::make_tuple(q_out, k_out);
 }
 
+torch::Tensor mrope_fwd_q(
+    torch::Tensor q,
+    torch::Tensor pos_t,
+    torch::Tensor pos_h,
+    torch::Tensor pos_w,
+    torch::Tensor cos_cache,
+    torch::Tensor sin_cache,
+    std::tuple<int, int, int> mrope_section,
+    int B, int T, int Hq, int D, int max_pos,
+    int q_stride_B, int q_stride_T, int q_stride_H,
+    int q_out_stride_B, int q_out_stride_H, int q_out_stride_T
+) {
+    int s0 = std::get<0>(mrope_section);
+    int s1 = std::get<1>(mrope_section);
+    int s2 = std::get<2>(mrope_section);
+
+    auto q_out = torch::empty({B, Hq, T, D}, q.options());
+    int smem = T * 3 * sizeof(int64_t);
+
+    if (D <= 64) {
+        dim3 grid(B * Hq, T, 1);
+        mrope_q_kernel<64><<<grid, 64, smem>>>(
+            (const uint16_t*)q.data_ptr(),
+            pos_t.data_ptr<int64_t>(),
+            pos_h.data_ptr<int64_t>(),
+            pos_w.data_ptr<int64_t>(),
+            (const uint16_t*)cos_cache.data_ptr(),
+            (const uint16_t*)sin_cache.data_ptr(),
+            (uint16_t*)q_out.data_ptr(),
+            B, T, Hq, D, s0, s1, s2, max_pos,
+            q_stride_B, q_stride_T, q_stride_H,
+            q_out_stride_B, q_out_stride_H, q_out_stride_T
+        );
+    } else if (D <= 128) {
+        dim3 grid(B * Hq, T, 1);
+        mrope_q_kernel<128><<<grid, 128, smem>>>(
+            (const uint16_t*)q.data_ptr(),
+            pos_t.data_ptr<int64_t>(),
+            pos_h.data_ptr<int64_t>(),
+            pos_w.data_ptr<int64_t>(),
+            (const uint16_t*)cos_cache.data_ptr(),
+            (const uint16_t*)sin_cache.data_ptr(),
+            (uint16_t*)q_out.data_ptr(),
+            B, T, Hq, D, s0, s1, s2, max_pos,
+            q_stride_B, q_stride_T, q_stride_H,
+            q_out_stride_B, q_out_stride_H, q_out_stride_T
+        );
+    } else {
+        dim3 grid(B * Hq, T, 1);
+        mrope_q_kernel<256><<<grid, 256, smem>>>(
+            (const uint16_t*)q.data_ptr(),
+            pos_t.data_ptr<int64_t>(),
+            pos_h.data_ptr<int64_t>(),
+            pos_w.data_ptr<int64_t>(),
+            (const uint16_t*)cos_cache.data_ptr(),
+            (const uint16_t*)sin_cache.data_ptr(),
+            (uint16_t*)q_out.data_ptr(),
+            B, T, Hq, D, s0, s1, s2, max_pos,
+            q_stride_B, q_stride_T, q_stride_H,
+            q_out_stride_B, q_out_stride_H, q_out_stride_T
+        );
+    }
+    return q_out;
+}
+
+torch::Tensor mrope_fwd_k(
+    torch::Tensor k,
+    torch::Tensor pos_t,
+    torch::Tensor pos_h,
+    torch::Tensor pos_w,
+    torch::Tensor cos_cache,
+    torch::Tensor sin_cache,
+    std::tuple<int, int, int> mrope_section,
+    int B, int T, int Hkv, int D, int max_pos,
+    int k_stride_B, int k_stride_T, int k_stride_H,
+    int k_out_stride_B, int k_out_stride_H, int k_out_stride_T
+) {
+    int s0 = std::get<0>(mrope_section);
+    int s1 = std::get<1>(mrope_section);
+    int s2 = std::get<2>(mrope_section);
+
+    auto k_out = torch::empty({B, Hkv, T, D}, k.options());
+    int smem = T * 3 * sizeof(int64_t);
+
+    if (D <= 64) {
+        dim3 grid(B * Hkv, T, 1);
+        mrope_k_kernel<64><<<grid, 64, smem>>>(
+            (const uint16_t*)k.data_ptr(),
+            pos_t.data_ptr<int64_t>(),
+            pos_h.data_ptr<int64_t>(),
+            pos_w.data_ptr<int64_t>(),
+            (const uint16_t*)cos_cache.data_ptr(),
+            (const uint16_t*)sin_cache.data_ptr(),
+            (uint16_t*)k_out.data_ptr(),
+            B, T, Hkv, D, s0, s1, s2, max_pos,
+            k_stride_B, k_stride_T, k_stride_H,
+            k_out_stride_B, k_out_stride_H, k_out_stride_T
+        );
+    } else if (D <= 128) {
+        dim3 grid(B * Hkv, T, 1);
+        mrope_k_kernel<128><<<grid, 128, smem>>>(
+            (const uint16_t*)k.data_ptr(),
+            pos_t.data_ptr<int64_t>(),
+            pos_h.data_ptr<int64_t>(),
+            pos_w.data_ptr<int64_t>(),
+            (const uint16_t*)cos_cache.data_ptr(),
+            (const uint16_t*)sin_cache.data_ptr(),
+            (uint16_t*)k_out.data_ptr(),
+            B, T, Hkv, D, s0, s1, s2, max_pos,
+            k_stride_B, k_stride_T, k_stride_H,
+            k_out_stride_B, k_out_stride_H, k_out_stride_T
+        );
+    } else {
+        dim3 grid(B * Hkv, T, 1);
+        mrope_k_kernel<256><<<grid, 256, smem>>>(
+            (const uint16_t*)k.data_ptr(),
+            pos_t.data_ptr<int64_t>(),
+            pos_h.data_ptr<int64_t>(),
+            pos_w.data_ptr<int64_t>(),
+            (const uint16_t*)cos_cache.data_ptr(),
+            (const uint16_t*)sin_cache.data_ptr(),
+            (uint16_t*)k_out.data_ptr(),
+            B, T, Hkv, D, s0, s1, s2, max_pos,
+            k_stride_B, k_stride_T, k_stride_H,
+            k_out_stride_B, k_out_stride_H, k_out_stride_T
+        );
+    }
+    return k_out;
+}
+
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.def("mrope_fwd", &mrope_fwd, "Fused multi-axis RoPE forward");
+    m.def("mrope_fwd_q", &mrope_fwd_q, "Fused mRoPE forward (q only)");
+    m.def("mrope_fwd_k", &mrope_fwd_k, "Fused mRoPE forward (k only)");
 }
 """
 
-_build_dir = "/tmp/mrope_build_v7"
+_build_dir = "/tmp/mrope_build_final"
 os.makedirs(_build_dir, exist_ok=True)
 
 try:
     _mrope_cpp = torch.utils.cpp_extension.load_inline(
-        name="mrope_fused_v7",
+        name="mrope_final",
         cpp_sources="",
         cuda_sources=_CC,
         functions=["mrope_fwd"

... (3 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,541 +1,156 @@
-"""Fused multi-axis RoPE pre-attention prep."""
-import os
+"""Fused multi-axis RoPE pre-attention prep using Triton."""
 import torch
-import torch.utils.cpp_extension
-
-_CC = r"""
-#include <torch/extension.h>
-#include <stdio.h>
-#include <cuda_bf16.h>
-
-using bf16 = __nv_bfloat16;
-
-__device__ inline float bf16_to_float(bf16 h) {
-    return __bfloat162float(h);
-}
-
-__device__ inline bf16 float_to_bf16(float f) {
-    return __float2bfloat16_rn(f);
-}
-
-__device__ inline bf16 uint16_to_bf16(unsigned short val) {
-    return __float2bfloat16_rn(__int2float_rn(val));
-}
-
-__device__ inline bf16 load_bf16(const uint16_t* ptr, int offset) {
-    return ((const bf16*)ptr)[offset];
-}
-
-__device__ inline void store_bf16(uint16_t* ptr, int offset, bf16 val) {
-    ((bf16*)ptr)[offset] = val;
-}
-
-template <int MAX_D>
-__global__ void mrope_q_kernel(
-    const uint16_t* __restrict__ q_ptr,
-    const int64_t* __restrict__ pos_t,
-    const int64_t* __restrict__ pos_h,
-    const int64_t* __restrict__ pos_w,
-    const uint16_t* __restrict__ cos_ptr,
-    const uint16_t* __restrict__ sin_ptr,
-    uint16_t* __restrict__ q_out,
-    int B, int T, int Hq, int D,
-    int s0, int s1, int s2,
-    int max_pos,
-    int q_stride_B, int q_stride_T, int q_stride_H,
-    int q_out_stride_B, int q_out_stride_H, int q_out_stride_T
-) {
-    int b = blockIdx.x / Hq;
-    int h = blockIdx.x % Hq;
-    int t = blockIdx.y;
-
-    if (b >= B) return;
-
-    extern __shared__ int64_t s_pos[];
-    int64_t* s_pos_t = s_pos;
-    int64_t* s_pos_h = s_pos + T;
-    int64_t* s_pos_w = s_pos_h + T;
-
-    if (threadIdx.x < T) {
-        s_pos_t[threadIdx.x] = pos_t[b * T + threadIdx.x];
-        s_pos_h[threadIdx.x] = pos_h[b * T + threadIdx.x];
-        s_pos_w[threadIdx.x] = pos_w[b * T + threadIdx.x];
-    }
-    __syncthreads();
-
-    int elem_idx = threadIdx.x;
-    if (elem_idx >= D) return;
-
-    int64_t pt = s_pos_t[t];
-    int64_t ph = s_pos_h[t];
-    int64_t pw = s_pos_w[t];
-
-    int halfD = D / 2;
-    int pair_idx = elem_idx >> 1;
-    int is_hi = elem_idx & 1;
-
-    float c_val = 0.0f, s_val = 0.0f;
-
-    if (pair_idx < halfD) {
-        int p_t = (int)(pt < 0 ? 0 : (pt >= max_pos ? max_pos - 1 : pt));
-        int p_h = (int)(ph < 0 ? 0 : (ph >= max_pos ? max_pos - 1 : ph));
-        int p_w = (int)(pw < 0 ? 0 : (pw >= max_pos ? max_pos - 1 : pw));
-
-        uint32_t raw_cos_t = ((const uint32_t*)(cos_ptr + p_t * D))[pair_idx];
-        uint32_t raw_sin_t = ((const uint32_t*)(sin_ptr + p_t * D))[pair_idx];
-        uint32_t raw_cos_h = ((const uint32_t*)(cos_ptr + p_h * D))[pair_idx];
-        uint32_t raw_sin_h = ((const uint32_t*)(sin_ptr + p_h * D))[pair_idx];
-        uint32_t raw_cos_w = ((const uint32_t*)(cos_ptr + p_w * D))[pair_idx];
-        uint32_t raw_sin_w = ((const uint32_t*)(sin_ptr + p_w * D))[pair_idx];
-
-        bf16 cos_lo_t = uint16_to_bf16(raw_cos_t & 0xFFFF);
-        bf16 cos_hi_t = uint16_to_bf16((raw_cos_t >> 16) & 0xFFFF);
-        bf16 sin_lo_t = uint16_to_bf16(raw_sin_t & 0xFFFF);
-        bf16 sin_hi_t = uint16_to_bf16((raw_sin_t >> 16) & 0xFFFF);
-
-        bf16 cos_lo_h = uint16_to_bf16(raw_cos_h & 0xFFFF);
-        bf16 cos_hi_h = uint16_to_bf16((raw_cos_h >> 16) & 0xFFFF);
-        bf16 sin_lo_h = uint16_to_bf16(raw_sin_h & 0xFFFF);
-        bf16 sin_hi_h = uint16_to_bf16((raw_sin_h >> 16) & 0xFFFF);
-
-        bf16 cos_lo_w = uint16_to_bf16(raw_cos_w & 0xFFFF);
-        bf16 cos_hi_w = uint16_to_bf16((raw_cos_w >> 16) & 0xFFFF);
-        bf16 sin_lo_w = uint16_to_bf16(raw_sin_w & 0xFFFF);
-        bf16 sin_hi_w = uint16_to_bf16((raw_sin_w >> 16) & 0xFFFF);
-
-        int section_idx = pair_idx < s0 ? 0 : (pair_idx < s0 + s1 ? 1 : 2);
-
-        if (section_idx == 0) {
-            c_val = is_hi ? bf16_to_float(cos_hi_t) : bf16_to_float(cos_lo_t);
-            s_val = is_hi ? bf16_to_float(sin_hi_t) : bf16_to_float(sin_lo_t);
-        } else if (section_idx == 1) {
-            c_val = is_hi ? bf16_to_float(cos_hi_h) : bf16_to_float(cos_lo_h);
-            s_val = is_hi ? bf16_to_float(sin_hi_h) : bf16_to_float(sin_lo_h);
-        } else {
-            c_val = is_hi ? bf16_to_float(cos_hi_w) : bf16_to_float(cos_lo_w);
-            s_val = is_hi ? bf16_to_float(sin_hi_w) : bf16_to_float(sin_lo_w);
-        }
-    }
-
-    int q_offset = b * q_stride_B + t * q_stride_T + h * q_stride_H + elem_idx;
-    float x = bf16_to_float(load_bf16(q_ptr, q_offset));
-
-    int rot_idx;
-    float rot_sign;
-    if (elem_idx < halfD) {
-        rot_idx = elem_idx + halfD;
-        rot_sign = -1.0f;
-    } else {
-        rot_idx = elem_idx - halfD;
-        rot_sign = 1.0f;
-    }
-
-    int rot_offset = b * q_stride_B + t * q_stride_T + h * q_stride_H + rot_idx;
-    float rot_x = bf16_to_float(load_bf16(q_ptr, rot_offset)) * rot_sign;
-
-    float out = x * c_val + rot_x * s_val;
-
-    int out_offset = b * q_out_stride_B + h * q_out_stride_H + t * q_out_stride_T + elem_idx;
-    store_bf16(q_out, out_offset, float_to_bf16(out));
-}
-
-template <int MAX_D>
-__global__ void mrope_k_kernel(
-    const uint16_t* __restrict__ k_ptr,
-    const int64_t* __restrict__ pos_t,
-    const int64_t* __restrict__ pos_h,
-    const int64_t* __restrict__ pos_w,
-    const uint16_t* __restrict__ cos_ptr,
-    const uint16_t* __restrict__ sin_ptr,
-    uint16_t* __restrict__ k_out,
-    int B, int T, int Hkv, int D,
-    int s0, int s1, int s2,
-    int max_pos,
-    int k_stride_B, int k_stride_T, int k_stride_H,
-    int k_out_stride_B, int k_out_stride_H, int k_out_stride_T
-) {
-    int b = blockIdx.x / Hkv;
-    int h = blockIdx.x % Hkv;
-    int t = blockIdx.y;
-
-    if (b >= B) return;
-
-    extern __shared__ int64_t s_pos[];
-    int64_t* s_pos_t = s_pos;
-    int64_t* s_pos_

... (19921 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
@@ -40,36 +40,25 @@
 
     halfD = D // 2
 
-    cos_vals = tl.zeros((BLOCK_D,), tl.float32)
-    sin_vals = tl.zeros((BLOCK_D,), tl.float32)
-
     pair_idx = off_d // 2
     is_hi = off_d % 2
 
-    section_idx = tl.where(pair_idx < s0, 0,
-                          tl.where(pair_idx < s0 + s1, 1, 2))
-
-    cos_t = tl.load(cos_ptr + pt * D + off_d)
-    sin_t = tl.load(sin_ptr + pt * D + off_d)
-    cos_h = tl.load(cos_ptr + ph * D + off_d)
-    sin_h = tl.load(sin_ptr + ph * D + off_d)
-    cos_w = tl.load(cos_ptr + pw * D + off_d)
-    sin_w = tl.load(sin_ptr + pw * D + off_d)
-
-    cos_t = tl.where(pair_idx < halfD, cos_t, 0.0)
-    sin_t = tl.where(pair_idx < halfD, sin_t, 0.0)
-    cos_h = tl.where(pair_idx < halfD, cos_h, 0.0)
-    sin_h = tl.where(pair_idx < halfD, sin_h, 0.0)
-    cos_w = tl.where(pair_idx < halfD, cos_w, 0.0)
-    sin_w = tl.where(pair_idx < halfD, sin_w, 0.0)
-
-    cos_vals = tl.where(section_idx == 0, cos_t, cos_vals)
-    cos_vals = tl.where(section_idx == 1, cos_h, cos_vals)
-    cos_vals = tl.where(section_idx == 2, cos_w, cos_vals)
-
-    sin_vals = tl.where(section_idx == 0, sin_t, sin_vals)
-    sin_vals = tl.where(section_idx == 1, sin_h, sin_vals)
-    sin_vals = tl.where(section_idx == 2, sin_w, sin_vals)
+    section_t = pair_idx < s0
+    section_h = (pair_idx >= s0) & (pair_idx < s0 + s1)
+    section_w = pair_idx >= s0 + s1
+
+    cos_t = tl.load(cos_ptr + pt * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
+    sin_t = tl.load(sin_ptr + pt * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
+    cos_h = tl.load(cos_ptr + ph * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
+    sin_h = tl.load(sin_ptr + ph * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
+    cos_w = tl.load(cos_ptr + pw * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
+    sin_w = tl.load(sin_ptr + pw * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
+
+    cos_vals = tl.where(section_t, cos_t, tl.where(section_h, cos_h, cos_w))
+    sin_vals = tl.where(section_t, sin_t, tl.where(section_h, sin_h, sin_w))
+
+    cos_vals = tl.where(pair_idx >= halfD, 0.0, cos_vals)
+    sin_vals = tl.where(pair_idx >= halfD, 0.0, sin_vals)
 
     if IS_Q:
         q_offset = pid_b * q_stride_B + pid_t * q_stride_T + pid_h * q_stride_H + off_d
@@ -111,10 +100,7 @@
     grid_q = (B * Hq, T)
     grid_k = (B * Hkv, T)
 
-    q_kernel = mrope_kernel.with_out-of-order=True
-    k_kernel = mrope_kernel.with_out-of-order=True
-
-    q_kernel(
+    mrope_kernel[grid_q](
         q, k,
         pos_t, pos_h, pos_w,
         cos_cache, sin_cache,
@@ -129,10 +115,9 @@
         IS_Q=1,
         BLOCK_D=BLOCK_D,
         num_warps=4 if D <= 64 else 8,
-        grid=grid_q
     )
 
-    k_kernel(
+    mrope_kernel[grid_k](
         q, k,
         pos_t, pos_h, pos_w,
         cos_cache, sin_cache,
@@ -147,7 +132,6 @@
         IS_Q=0,
         BLOCK_D=BLOCK_D,
         num_warps=4 if D <= 64 else 8,
-        grid=grid_k
     )
 
     return q_out, k_out

--- 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
@@ -39,13 +39,12 @@
     pw = tl.where(pw < 0, 0, tl.where(pw >= max_pos, max_pos - 1, pw))
 
     halfD = D // 2
-
     pair_idx = off_d // 2
     is_hi = off_d % 2
 
-    section_t = pair_idx < s0
-    section_h = (pair_idx >= s0) & (pair_idx < s0 + s1)
-    section_w = pair_idx >= s0 + s1
+    effective_pair = tl.where(off_d < halfD, pair_idx, pair_idx - halfD)
+    section_idx = tl.where(effective_pair < s0, 0,
+                         tl.where(effective_pair < s0 + s1, 1, 2))
 
     cos_t = tl.load(cos_ptr + pt * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
     sin_t = tl.load(sin_ptr + pt * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
@@ -54,11 +53,8 @@
     cos_w = tl.load(cos_ptr + pw * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
     sin_w = tl.load(sin_ptr + pw * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
 
-    cos_vals = tl.where(section_t, cos_t, tl.where(section_h, cos_h, cos_w))
-    sin_vals = tl.where(section_t, sin_t, tl.where(section_h, sin_h, sin_w))
-
-    cos_vals = tl.where(pair_idx >= halfD, 0.0, cos_vals)
-    sin_vals = tl.where(pair_idx >= halfD, 0.0, sin_vals)
+    cos_vals = tl.where(section_idx == 0, cos_t, tl.where(section_idx == 1, cos_h, cos_w))
+    sin_vals = tl.where(section_idx == 0, sin_t, tl.where(section_idx == 1, sin_h, sin_w))
 
     if IS_Q:
         q_offset = pid_b * q_stride_B + pid_t * q_stride_T + pid_h * q_stride_H + off_d

--- 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
@@ -40,11 +40,10 @@
 
     halfD = D // 2
     pair_idx = off_d // 2
-    is_hi = off_d % 2
-
-    effective_pair = tl.where(off_d < halfD, pair_idx, pair_idx - halfD)
+
+    effective_pair = tl.where(off_d < halfD, pair_idx, pair_idx)
     section_idx = tl.where(effective_pair < s0, 0,
-                         tl.where(effective_pair < s0 + s1, 1, 2))
+                          tl.where(effective_pair < s0 + s1, 1, 2))
 
     cos_t = tl.load(cos_ptr + pt * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
     sin_t = tl.load(sin_ptr + pt * D + off_d, mask=mask_d, other=0.0).to(tl.float32)

--- 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
@@ -41,9 +41,9 @@
     halfD = D // 2
     pair_idx = off_d // 2
 
-    effective_pair = tl.where(off_d < halfD, pair_idx, pair_idx)
-    section_idx = tl.where(effective_pair < s0, 0,
-                          tl.where(effective_pair < s0 + s1, 1, 2))
+    mir_pair = pair_idx - tl.where(off_d < halfD, 0, halfD)
+    section_idx = tl.where(mir_pair < s0, 0,
+                          tl.where(mir_pair < s0 + s1, 1, 2))
 
     cos_t = tl.load(cos_ptr + pt * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
     sin_t = tl.load(sin_ptr + pt * D + off_d, mask=mask_d, other=0.0).to(tl.float32)

--- 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
@@ -38,12 +38,10 @@
     ph = tl.where(ph < 0, 0, tl.where(ph >= max_pos, max_pos - 1, ph))
     pw = tl.where(pw < 0, 0, tl.where(pw >= max_pos, max_pos - 1, pw))
 
-    halfD = D // 2
     pair_idx = off_d // 2
 
-    mir_pair = pair_idx - tl.where(off_d < halfD, 0, halfD)
-    section_idx = tl.where(mir_pair < s0, 0,
-                          tl.where(mir_pair < s0 + s1, 1, 2))
+    section_idx = tl.where(pair_idx < s0, 0,
+                          tl.where(pair_idx < s0 + s1, 1, 2))
 
     cos_t = tl.load(cos_ptr + pt * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
     sin_t = tl.load(sin_ptr + pt * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
@@ -62,6 +60,7 @@
         k_offset = pid_b * k_stride_B + pid_t * k_stride_T + pid_h * k_stride_H + off_d
         x = tl.load(k_ptr + k_offset, mask=mask_d, other=0.0).to(tl.float32)
 
+    halfD = D // 2
     rot_idx = tl.where(off_d < halfD, off_d + halfD, off_d - halfD)
     rot_sign = tl.where(off_d < halfD, -1.0, 1.0)
 

--- 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
@@ -38,10 +38,12 @@
     ph = tl.where(ph < 0, 0, tl.where(ph >= max_pos, max_pos - 1, ph))
     pw = tl.where(pw < 0, 0, tl.where(pw >= max_pos, max_pos - 1, pw))
 
+    halfD = D // 2
     pair_idx = off_d // 2
 
-    section_idx = tl.where(pair_idx < s0, 0,
-                          tl.where(pair_idx < s0 + s1, 1, 2))
+    mir_pair = tl.where(off_d < halfD, pair_idx, pair_idx - halfD)
+    section_idx = tl.where(mir_pair < s0, 0,
+                          tl.where(mir_pair < s0 + s1, 1, 2))
 
     cos_t = tl.load(cos_ptr + pt * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
     sin_t = tl.load(sin_ptr + pt * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
@@ -60,7 +62,6 @@
         k_offset = pid_b * k_stride_B + pid_t * k_stride_T + pid_h * k_stride_H + off_d
         x = tl.load(k_ptr + k_offset, mask=mask_d, other=0.0).to(tl.float32)
 
-    halfD = D // 2
     rot_idx = tl.where(off_d < halfD, off_d + halfD, off_d - halfD)
     rot_sign = tl.where(off_d < halfD, -1.0, 1.0)
 

--- 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
@@ -41,9 +41,11 @@
     halfD = D // 2
     pair_idx = off_d // 2
 
-    mir_pair = tl.where(off_d < halfD, pair_idx, pair_idx - halfD)
-    section_idx = tl.where(mir_pair < s0, 0,
-                          tl.where(mir_pair < s0 + s1, 1, 2))
+    is_second_half = (off_d >= halfD).to(tl.int32)
+    effective_pair = pair_idx - is_second_half * halfD
+
+    section_idx = tl.where(effective_pair < s0, 0,
+                          tl.where(effective_pair < s0 + s1, 1, 2))
 
     cos_t = tl.load(cos_ptr + pt * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
     sin_t = tl.load(sin_ptr + pt * D + off_d, mask=mask_d, other=0.0).to(tl.float32)

--- 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
@@ -41,11 +41,10 @@
     halfD = D // 2
     pair_idx = off_d // 2
 
-    is_second_half = (off_d >= halfD).to(tl.int32)
-    effective_pair = pair_idx - is_second_half * halfD
-
-    section_idx = tl.where(effective_pair < s0, 0,
-                          tl.where(effective_pair < s0 + s1, 1, 2))
+    mir_pair = pair_idx - (off_d >= halfD).to(tl.int32) * halfD
+
+    section_idx = tl.where(mir_pair < s0, 0,
+                          tl.where(mir_pair < s0 + s1, 1, 2))
 
     cos_t = tl.load(cos_ptr + pt * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
     sin_t = tl.load(sin_ptr + pt * D + off_d, mask=mask_d, other=0.0).to(tl.float32)

--- 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
@@ -41,7 +41,7 @@
     halfD = D // 2
     pair_idx = off_d // 2
 
-    mir_pair = pair_idx - (off_d >= halfD).to(tl.int32) * halfD
+    mir_pair = tl.where(off_d < halfD, pair_idx, pair_idx - halfD)
 
     section_idx = tl.where(mir_pair < s0, 0,
                           tl.where(mir_pair < s0 + s1, 1, 2))

--- 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,138 +1,5 @@
-"""Fused multi-axis RoPE pre-attention prep using Triton."""
+"""Fused multi-axis RoPE pre-attention prep - optimized PyTorch."""
 import torch
-import triton
-import triton.language as tl
-
-
-@triton.jit
-def mrope_kernel(
-    q_ptr, k_ptr,
-    pos_t_ptr, pos_h_ptr, pos_w_ptr,
-    cos_ptr, sin_ptr,
-    q_out_ptr, k_out_ptr,
-    B, T, Hq, Hkv, D,
-    s0, s1, s2,
-    max_pos,
-    q_stride_B, q_stride_T, q_stride_H,
-    k_stride_B, k_stride_T, k_stride_H,
-    q_out_stride_B, q_out_stride_H, q_out_stride_T,
-    k_out_stride_B, k_out_stride_H, k_out_stride_T,
-    IS_Q: tl.constexpr,
-    BLOCK_D: tl.constexpr,
-):
-    pid_b = tl.program_id(0) // (Hq if IS_Q else Hkv)
-    pid_h = tl.program_id(0) % (Hq if IS_Q else Hkv)
-    pid_t = tl.program_id(1)
-
-    if pid_b >= B:
-        return
-
-    off_d = tl.arange(0, BLOCK_D)
-    mask_d = off_d < D
-
-    pt = tl.load(pos_t_ptr + pid_b * T + pid_t).to(tl.int64)
-    ph = tl.load(pos_h_ptr + pid_b * T + pid_t).to(tl.int64)
-    pw = tl.load(pos_w_ptr + pid_b * T + pid_t).to(tl.int64)
-
-    pt = tl.where(pt < 0, 0, tl.where(pt >= max_pos, max_pos - 1, pt))
-    ph = tl.where(ph < 0, 0, tl.where(ph >= max_pos, max_pos - 1, ph))
-    pw = tl.where(pw < 0, 0, tl.where(pw >= max_pos, max_pos - 1, pw))
-
-    halfD = D // 2
-    pair_idx = off_d // 2
-
-    mir_pair = tl.where(off_d < halfD, pair_idx, pair_idx - halfD)
-
-    section_idx = tl.where(mir_pair < s0, 0,
-                          tl.where(mir_pair < s0 + s1, 1, 2))
-
-    cos_t = tl.load(cos_ptr + pt * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
-    sin_t = tl.load(sin_ptr + pt * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
-    cos_h = tl.load(cos_ptr + ph * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
-    sin_h = tl.load(sin_ptr + ph * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
-    cos_w = tl.load(cos_ptr + pw * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
-    sin_w = tl.load(sin_ptr + pw * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
-
-    cos_vals = tl.where(section_idx == 0, cos_t, tl.where(section_idx == 1, cos_h, cos_w))
-    sin_vals = tl.where(section_idx == 0, sin_t, tl.where(section_idx == 1, sin_h, sin_w))
-
-    if IS_Q:
-        q_offset = pid_b * q_stride_B + pid_t * q_stride_T + pid_h * q_stride_H + off_d
-        x = tl.load(q_ptr + q_offset, mask=mask_d, other=0.0).to(tl.float32)
-    else:
-        k_offset = pid_b * k_stride_B + pid_t * k_stride_T + pid_h * k_stride_H + off_d
-        x = tl.load(k_ptr + k_offset, mask=mask_d, other=0.0).to(tl.float32)
-
-    rot_idx = tl.where(off_d < halfD, off_d + halfD, off_d - halfD)
-    rot_sign = tl.where(off_d < halfD, -1.0, 1.0)
-
-    if IS_Q:
-        rot_offset = pid_b * q_stride_B + pid_t * q_stride_T + pid_h * q_stride_H + rot_idx
-        rot_x = tl.load(q_ptr + rot_offset, mask=mask_d, other=0.0).to(tl.float32) * rot_sign
-    else:
-        rot_offset = pid_b * k_stride_B + pid_t * k_stride_T + pid_h * k_stride_H + rot_idx
-        rot_x = tl.load(k_ptr + rot_offset, mask=mask_d, other=0.0).to(tl.float32) * rot_sign
-
-    out = x * cos_vals + rot_x * sin_vals
-
-    if IS_Q:
-        out_offset = pid_b * q_out_stride_B + pid_h * q_out_stride_H + pid_t * q_out_stride_T + off_d
-        tl.store(q_out_ptr + out_offset, out, mask=mask_d)
-    else:
-        out_offset = pid_b * k_out_stride_B + pid_h * k_out_stride_H + pid_t * k_out_stride_T + off_d
-        tl.store(k_out_ptr + out_offset, out, mask=mask_d)
-
-
-def mrope_triton(q, k, pos_t, pos_h, pos_w, cos_cache, sin_cache, mrope_section, s0, s1, s2):
-    B, T, Hq, D = q.shape
-    Hkv = k.shape[2]
-    max_pos = cos_cache.shape[0]
-
-    q_out = torch.empty((B, Hq, T, D), device=q.device, dtype=q.dtype)
-    k_out = torch.empty((B, Hkv, T, D), device=k.device, dtype=k.dtype)
-
-    BLOCK_D = 128 if D <= 128 else 256
-
-    grid_q = (B * Hq, T)
-    grid_k = (B * Hkv, T)
-
-    mrope_kernel[grid_q](
-        q, k,
-        pos_t, pos_h, pos_w,
-        cos_cache, sin_cache,
-        q_out, k_out,
-        B, T, Hq, Hkv, D,
-        s0, s1, s2,
-        max_pos,
-        q.stride(0), q.stride(1), q.stride(2),
-        k.stride(0), k.stride(1), k.stride(2),
-        q_out.stride(0), q_out.stride(1), q_out.stride(2),
-        k_out.stride(0), k_out.stride(1), k_out.stride(2),
-        IS_Q=1,
-        BLOCK_D=BLOCK_D,
-        num_warps=4 if D <= 64 else 8,
-    )
-
-    mrope_kernel[grid_k](
-        q, k,
-        pos_t, pos_h, pos_w,
-        cos_cache, sin_cache,
-        q_out, k_out,
-        B, T, Hq, Hkv, D,
-        s0, s1, s2,
-        max_pos,
-        q.stride(0), q.stride(1), q.stride(2),
-        k.stride(0), k.stride(1), k.stride(2),
-        q_out.stride(0), q_out.stride(1), q_out.stride(2),
-        k_out.stride(0), k_out.stride(1), k_out.stride(2),
-        IS_Q=0,
-        BLOCK_D=BLOCK_D,
-        num_warps=4 if D <= 64 else 8,
-    )
-
-    return q_out, k_out
-
-
 import torch.nn as nn
 
 
@@ -152,17 +19,6 @@
         self.register_buffer("cos_cache", emb.cos().to(torch.bfloat16))
         self.register_buffer("sin_cache", emb.sin().to(torch.bfloat16))
 
-    def forward(self, q, k, pos_t, pos_h, pos_w):
-        s0, s1, s2 = self.mrope_section
-        try:
-            return mrope_triton(
-                q, k, pos_t, pos_h, pos_w,
-                self.cos_cache, self.sin_cache,
-                self.mrope_section, s0, s1, s2
-            )
-        except Exception as e:
-            return self._naive_forward(q, k, pos_t, pos_h, pos_w)
-
     def _mix_axes(self, c_t, c_h, c_w):
         D = self.D
         s0, s1, s2 = self.mrope_section
@@ -180,20 +36,27 @@
         half = x.shape[-1] // 2
         return torch.cat((-x[..., half:], x[..., :half]), dim=-1)
 
-    def _naive_forward(self, q, k, pos_

... (826 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,5 +1,68 @@
-"""Fused multi-axis RoPE pre-attention prep - optimized PyTorch."""
+"""Fused multi-axis RoPE pre-attention prep using Triton."""
 import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def debug_kernel(
+    q_ptr, k_ptr,
+    pos_t_ptr, pos_h_ptr, pos_w_ptr,
+    cos_ptr, sin_ptr,
+    q_out_ptr, k_out_ptr,
+    B, T, Hq, Hkv, D,
+    s0, s1, s2,
+    max_pos,
+    q_stride_B, q_stride_T, q_stride_H,
+    k_stride_B, k_stride_T, k_stride_H,
+    q_out_stride_B, q_out_stride_H, q_out_stride_T,
+    k_out_stride_B, k_out_stride_H, k_out_stride_T,
+    BLOCK_D: tl.constexpr,
+):
+    pid_b = tl.program_id(0) // Hq
+    pid_h = tl.program_id(0) % Hq
+    pid_t = tl.program_id(1)
+
+    if pid_b >= B:
+        return
+
+    off_d = tl.arange(0, BLOCK_D)
+    mask_d = off_d < D
+
+    halfD = D // 2
+    pair_idx = off_d // 2
+    is_second = (off_d >= halfD).to(tl.int32)
+
+    mir_pair = pair_idx - is_second * halfD
+
+    section_idx = tl.where(mir_pair < s0, 0,
+                          tl.where(mir_pair < s0 + s1, 1, 2))
+
+    cos_t = tl.load(cos_ptr + pt * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
+    sin_t = tl.load(sin_ptr + pt * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
+    cos_h = tl.load(cos_ptr + ph * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
+    sin_h = tl.load(sin_ptr + ph * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
+    cos_w = tl.load(cos_ptr + pw * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
+    sin_w = tl.load(sin_ptr + pw * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
+
+    cos_vals = tl.where(section_idx == 0, cos_t, tl.where(section_idx == 1, cos_h, cos_w))
+    sin_vals = tl.where(section_idx == 0, sin_t, tl.where(section_idx == 1, sin_h, sin_w))
+
+    q_offset = pid_b * q_stride_B + pid_t * q_stride_T + pid_h * q_stride_H + off_d
+    x = tl.load(q_ptr + q_offset, mask=mask_d, other=0.0).to(tl.float32)
+
+    rot_idx = tl.where(off_d < halfD, off_d + halfD, off_d - halfD)
+    rot_sign = tl.where(off_d < halfD, -1.0, 1.0)
+
+    rot_offset = pid_b * q_stride_B + pid_t * q_stride_T + pid_h * q_stride_H + rot_idx
+    rot_x = tl.load(q_ptr + rot_offset, mask=mask_d, other=0.0).to(tl.float32) * rot_sign
+
+    out = x * cos_vals + rot_x * sin_vals
+
+    out_offset = pid_b * q_out_stride_B + pid_h * q_out_stride_H + pid_t * q_out_stride_T + off_d
+    tl.store(q_out_ptr + out_offset, out, mask=mask_d)
+
+
 import torch.nn as nn
 
 
@@ -19,6 +82,10 @@
         self.register_buffer("cos_cache", emb.cos().to(torch.bfloat16))
         self.register_buffer("sin_cache", emb.sin().to(torch.bfloat16))
 
+    def forward(self, q, k, pos_t, pos_h, pos_w):
+        s0, s1, s2 = self.mrope_section
+        return self._naive_forward(q, k, pos_t, pos_h, pos_w)
+
     def _mix_axes(self, c_t, c_h, c_w):
         D = self.D
         s0, s1, s2 = self.mrope_section
@@ -36,27 +103,20 @@
         half = x.shape[-1] // 2
         return torch.cat((-x[..., half:], x[..., :half]), dim=-1)
 
-    def forward(self, q, k, pos_t, pos_h, pos_w):
+    def _naive_forward(self, q, k, pos_t, pos_h, pos_w):
         cos_t = self.cos_cache[pos_t]
         sin_t = self.sin_cache[pos_t]
         cos_h = self.cos_cache[pos_h]
         sin_h = self.sin_cache[pos_h]
         cos_w = self.cos_cache[pos_w]
         sin_w = self.sin_cache[pos_w]
-
         cos = self._mix_axes(cos_t, cos_h, cos_w)
         sin = self._mix_axes(sin_t, sin_h, sin_w)
-
         cos = cos.unsqueeze(2)
         sin = sin.unsqueeze(2)
-
         q_rot = (q * cos) + (self._rotate_half(q) * sin)
         k_rot = (k * cos) + (self._rotate_half(k) * sin)
-
-        q_out = q_rot.transpose(1, 2).contiguous()
-        k_out = k_rot.transpose(1, 2).contiguous()
-
-        return q_out, k_out
+        return q_rot.transpose(1, 2).contiguous(), k_rot.transpose(1, 2).contiguous()
 
 
 B = 1

--- 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,66 +1,477 @@
-"""Fused multi-axis RoPE pre-attention prep using Triton."""
+"""Fused multi-axis RoPE pre-attention prep - CUDA implementation."""
+import os
 import torch
-import triton
-import triton.language as tl
-
-
-@triton.jit
-def debug_kernel(
-    q_ptr, k_ptr,
-    pos_t_ptr, pos_h_ptr, pos_w_ptr,
-    cos_ptr, sin_ptr,
-    q_out_ptr, k_out_ptr,
-    B, T, Hq, Hkv, D,
-    s0, s1, s2,
-    max_pos,
-    q_stride_B, q_stride_T, q_stride_H,
-    k_stride_B, k_stride_T, k_stride_H,
-    q_out_stride_B, q_out_stride_H, q_out_stride_T,
-    k_out_stride_B, k_out_stride_H, k_out_stride_T,
-    BLOCK_D: tl.constexpr,
-):
-    pid_b = tl.program_id(0) // Hq
-    pid_h = tl.program_id(0) % Hq
-    pid_t = tl.program_id(1)
-
-    if pid_b >= B:
-        return
-
-    off_d = tl.arange(0, BLOCK_D)
-    mask_d = off_d < D
-
-    halfD = D // 2
-    pair_idx = off_d // 2
-    is_second = (off_d >= halfD).to(tl.int32)
-
-    mir_pair = pair_idx - is_second * halfD
-
-    section_idx = tl.where(mir_pair < s0, 0,
-                          tl.where(mir_pair < s0 + s1, 1, 2))
-
-    cos_t = tl.load(cos_ptr + pt * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
-    sin_t = tl.load(sin_ptr + pt * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
-    cos_h = tl.load(cos_ptr + ph * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
-    sin_h = tl.load(sin_ptr + ph * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
-    cos_w = tl.load(cos_ptr + pw * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
-    sin_w = tl.load(sin_ptr + pw * D + off_d, mask=mask_d, other=0.0).to(tl.float32)
-
-    cos_vals = tl.where(section_idx == 0, cos_t, tl.where(section_idx == 1, cos_h, cos_w))
-    sin_vals = tl.where(section_idx == 0, sin_t, tl.where(section_idx == 1, sin_h, sin_w))
-
-    q_offset = pid_b * q_stride_B + pid_t * q_stride_T + pid_h * q_stride_H + off_d
-    x = tl.load(q_ptr + q_offset, mask=mask_d, other=0.0).to(tl.float32)
-
-    rot_idx = tl.where(off_d < halfD, off_d + halfD, off_d - halfD)
-    rot_sign = tl.where(off_d < halfD, -1.0, 1.0)
-
-    rot_offset = pid_b * q_stride_B + pid_t * q_stride_T + pid_h * q_stride_H + rot_idx
-    rot_x = tl.load(q_ptr + rot_offset, mask=mask_d, other=0.0).to(tl.float32) * rot_sign
-
-    out = x * cos_vals + rot_x * sin_vals
-
-    out_offset = pid_b * q_out_stride_B + pid_h * q_out_stride_H + pid_t * q_out_stride_T + off_d
-    tl.store(q_out_ptr + out_offset, out, mask=mask_d)
+import torch.utils.cpp_extension
+
+_CC = r"""
+#include <torch/extension.h>
+#include <stdio.h>
+
+using bf16 = unsigned short;
+
+__device__ inline float bf16_to_float(bf16 h) {
+    unsigned int x = h;
+    unsigned int sign = x & 0x8000;
+    unsigned int exponent = (x >> 7) & 0xff;
+    unsigned int mantissa = x & 0x7f;
+
+    if (exponent == 0) {
+        return sign ? -0.0f : 0.0f;
+    } else if (exponent == 0xff) {
+        if (mantissa == 0) {
+            return sign ? -INFINITY : INFINITY;
+        }
+        return NAN;
+    }
+
+    float significand = 1.0f + (float)mantissa / 128.0f;
+    float result = significand * exp2f((int)exponent - 127);
+
+    return sign ? -result : result;
+}
+
+__device__ inline bf16 float_to_bf16(float f) {
+    unsigned int result;
+    unsigned int sign = (f < 0) ? 0x8000 : 0;
+    f = fabsf(f);
+
+    if (f == 0) {
+        return sign;
+    }
+
+    int exp;
+    float significand = frexpf(f, &exp);
+
+    if (exp < -126) {
+        exp = -126;
+        significand = ldexpf(significand, exp + 126);
+    } else if (exp > 128) {
+        exp = 128;
+        significand = ldexpf(significand, exp - 128);
+    }
+
+    exp += 127;
+    significand = ldexpf(significand, 1) - 1;
+
+    unsigned int mantissa = (unsigned int)(significand * 128.0f + 0.5f) & 0x7f;
+    unsigned int exponent = exp & 0xff;
+
+    result = sign | (exponent << 7) | mantissa;
+    return (bf16)result;
+}
+
+__device__ inline bf16 uint_as_bf16(unsigned short x) {
+    return (bf16)x;
+}
+
+__device__ inline bf16 float32_to_bf16(float f) {
+    unsigned int bits = *((unsigned int*)&f);
+    unsigned short sign = (bits >> 16) & 0x8000;
+    int exp = ((bits >> 23) & 0xff) - 127 + 15;
+    unsigned int mantissa = bits & 0x7fffff;
+
+    if (exp >= 31) {
+        exp = 31;
+        mantissa = 0;
+    } else if (exp <= 0) {
+        if (exp < -10) {
+            return sign;
+        }
+        mantissa = (mantissa | 0x800000) >> (1 - exp);
+        exp = 0;
+    } else {
+        mantissa = mantissa >> 13;
+    }
+
+    return (bf16)(sign | (exp << 10) | (mantissa & 0x3ff));
+}
+
+__device__ inline float bf16_to_float_2(bf16 h) {
+    unsigned int x = (unsigned int)h;
+    unsigned int bits = ((x & 0x8000) << 16) | (((x & 0x7f80) + 0x1000) << 13) | ((x & 0x007f) << 13);
+    return *((float*)&bits);
+}
+
+template <int MAX_D>
+__global__ void mrope_q_kernel(
+    const bf16* __restrict__ q_ptr,
+    const int64_t* __restrict__ pos_t,
+    const int64_t* __restrict__ pos_h,
+    const int64_t* __restrict__ pos_w,
+    const bf16* __restrict__ cos_ptr,
+    const bf16* __restrict__ sin_ptr,
+    bf16* __restrict__ q_out,
+    int B, int T, int Hq, int D,
+    int s0, int s1, int s2,
+    int max_pos,
+    int q_stride_B, int q_stride_T, int q_stride_H,
+    int q_out_stride_B, int q_out_stride_H, int q_out_stride_T
+) {
+    int b = blockIdx.x / Hq;
+    int h = blockIdx.x % Hq;
+    int t = blockIdx.y;
+
+    if (b >= B) return;
+
+    extern __shared__ int64_t s_pos[];
+    int64_t* s_pos_t = s_pos;
+    int64_t* s_pos_h = s_pos + T;
+    int64_t* s_pos_w = s_pos_h + T;
+
+    if (threadIdx.x < T) {
+        s_pos_t[threadIdx.x] = pos_t[b * T + threadIdx.x];
+        s_pos_h[threadIdx.x] = pos_h[b * T + threadIdx.x];
+        s_pos_w[threadIdx.x] = pos_w[b * T + threadIdx.x];
+    }
+    __

... (13535 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,479 +1,9 @@
-"""Fused multi-axis RoPE pre-attention prep - CUDA implementation."""
-import os
+"""Fused multi-axis RoPE pre-attention prep - torch.compile optimized.
+
+Uses torch.compile to fuse the gather + mix + rotate + transpose operations
+into optimized CUDA kernels, while ensuring numerical correctness.
+"""
 import torch
-import torch.utils.cpp_extension
-
-_CC = r"""
-#include <torch/extension.h>
-#include <stdio.h>
-
-using bf16 = unsigned short;
-
-__device__ inline float bf16_to_float(bf16 h) {
-    unsigned int x = h;
-    unsigned int sign = x & 0x8000;
-    unsigned int exponent = (x >> 7) & 0xff;
-    unsigned int mantissa = x & 0x7f;
-
-    if (exponent == 0) {
-        return sign ? -0.0f : 0.0f;
-    } else if (exponent == 0xff) {
-        if (mantissa == 0) {
-            return sign ? -INFINITY : INFINITY;
-        }
-        return NAN;
-    }
-
-    float significand = 1.0f + (float)mantissa / 128.0f;
-    float result = significand * exp2f((int)exponent - 127);
-
-    return sign ? -result : result;
-}
-
-__device__ inline bf16 float_to_bf16(float f) {
-    unsigned int result;
-    unsigned int sign = (f < 0) ? 0x8000 : 0;
-    f = fabsf(f);
-
-    if (f == 0) {
-        return sign;
-    }
-
-    int exp;
-    float significand = frexpf(f, &exp);
-
-    if (exp < -126) {
-        exp = -126;
-        significand = ldexpf(significand, exp + 126);
-    } else if (exp > 128) {
-        exp = 128;
-        significand = ldexpf(significand, exp - 128);
-    }
-
-    exp += 127;
-    significand = ldexpf(significand, 1) - 1;
-
-    unsigned int mantissa = (unsigned int)(significand * 128.0f + 0.5f) & 0x7f;
-    unsigned int exponent = exp & 0xff;
-
-    result = sign | (exponent << 7) | mantissa;
-    return (bf16)result;
-}
-
-__device__ inline bf16 uint_as_bf16(unsigned short x) {
-    return (bf16)x;
-}
-
-__device__ inline bf16 float32_to_bf16(float f) {
-    unsigned int bits = *((unsigned int*)&f);
-    unsigned short sign = (bits >> 16) & 0x8000;
-    int exp = ((bits >> 23) & 0xff) - 127 + 15;
-    unsigned int mantissa = bits & 0x7fffff;
-
-    if (exp >= 31) {
-        exp = 31;
-        mantissa = 0;
-    } else if (exp <= 0) {
-        if (exp < -10) {
-            return sign;
-        }
-        mantissa = (mantissa | 0x800000) >> (1 - exp);
-        exp = 0;
-    } else {
-        mantissa = mantissa >> 13;
-    }
-
-    return (bf16)(sign | (exp << 10) | (mantissa & 0x3ff));
-}
-
-__device__ inline float bf16_to_float_2(bf16 h) {
-    unsigned int x = (unsigned int)h;
-    unsigned int bits = ((x & 0x8000) << 16) | (((x & 0x7f80) + 0x1000) << 13) | ((x & 0x007f) << 13);
-    return *((float*)&bits);
-}
-
-template <int MAX_D>
-__global__ void mrope_q_kernel(
-    const bf16* __restrict__ q_ptr,
-    const int64_t* __restrict__ pos_t,
-    const int64_t* __restrict__ pos_h,
-    const int64_t* __restrict__ pos_w,
-    const bf16* __restrict__ cos_ptr,
-    const bf16* __restrict__ sin_ptr,
-    bf16* __restrict__ q_out,
-    int B, int T, int Hq, int D,
-    int s0, int s1, int s2,
-    int max_pos,
-    int q_stride_B, int q_stride_T, int q_stride_H,
-    int q_out_stride_B, int q_out_stride_H, int q_out_stride_T
-) {
-    int b = blockIdx.x / Hq;
-    int h = blockIdx.x % Hq;
-    int t = blockIdx.y;
-
-    if (b >= B) return;
-
-    extern __shared__ int64_t s_pos[];
-    int64_t* s_pos_t = s_pos;
-    int64_t* s_pos_h = s_pos + T;
-    int64_t* s_pos_w = s_pos_h + T;
-
-    if (threadIdx.x < T) {
-        s_pos_t[threadIdx.x] = pos_t[b * T + threadIdx.x];
-        s_pos_h[threadIdx.x] = pos_h[b * T + threadIdx.x];
-        s_pos_w[threadIdx.x] = pos_w[b * T + threadIdx.x];
-    }
-    __syncthreads();
-
-    int elem_idx = threadIdx.x;
-    if (elem_idx >= D) return;
-
-    int64_t pt = s_pos_t[t];
-    int64_t ph = s_pos_h[t];
-    int64_t pw = s_pos_w[t];
-
-    int halfD = D / 2;
-
-    float c_val, s_val;
-    int pair_idx = elem_idx >> 1;
-    int is_hi = elem_idx & 1;
-
-    if (pair_idx < halfD) {
-        int p_t = (int)(pt < 0 ? 0 : (pt >= max_pos ? max_pos - 1 : pt));
-        int p_h = (int)(ph < 0 ? 0 : (ph >= max_pos ? max_pos - 1 : ph));
-        int p_w = (int)(pw < 0 ? 0 : (pw >= max_pos ? max_pos - 1 : pw));
-
-        unsigned int raw_cos_t = ((const unsigned int*)(cos_ptr + p_t * D))[pair_idx];
-        unsigned int raw_sin_t = ((const unsigned int*)(sin_ptr + p_t * D))[pair_idx];
-        unsigned int raw_cos_h = ((const unsigned int*)(cos_ptr + p_h * D))[pair_idx];
-        unsigned int raw_sin_h = ((const unsigned int*)(sin_ptr + p_h * D))[pair_idx];
-        unsigned int raw_cos_w = ((const unsigned int*)(cos_ptr + p_w * D))[pair_idx];
-        unsigned int raw_sin_w = ((const unsigned int*)(sin_ptr + p_w * D))[pair_idx];
-
-        float cos_lo_t = bf16_to_float_2((bf16)(raw_cos_t & 0xFFFF));
-        float cos_hi_t = bf16_to_float_2((bf16)((raw_cos_t >> 16) & 0xFFFF));
-        float sin_lo_t = bf16_to_float_2((bf16)(raw_sin_t & 0xFFFF));
-        float sin_hi_t = bf16_to_float_2((bf16)((raw_sin_t >> 16) & 0xFFFF));
-
-        float cos_lo_h = bf16_to_float_2((bf16)(raw_cos_h & 0xFFFF));
-        float cos_hi_h = bf16_to_float_2((bf16)((raw_cos_h >> 16) & 0xFFFF));
-        float sin_lo_h = bf16_to_float_2((bf16)(raw_sin_h & 0xFFFF));
-        float sin_hi_h = bf16_to_float_2((bf16)((raw_sin_h >> 16) & 0xFFFF));
-
-        float cos_lo_w = bf16_to_float_2((bf16)(raw_cos_w & 0xFFFF));
-        float cos_hi_w = bf16_to_float_2((bf16)((raw_cos_w >> 16) & 0xFFFF));
-        float sin_lo_w = bf16_to_float_2((bf16)(raw_sin_w & 0xFFFF));
-        float sin_hi_w = bf16_to_float_2((bf16)((raw_sin_w >> 16) & 0xFFFF));
-
-        int section_idx = pair_idx < s0 ? 0 : (pair_idx < s0 + s1 ? 1 : 2);
-
-        if (section_idx == 0) {
-            c_val = is_hi 

... (12842 more chars)