submission 676725

- """
- sub120: Hybrid kernel — dispatches between batched bmm and Triton flash attention.
+ import torch
+ import os
+ from torch.utils.cpp_extension import load_inline
+ from task import input_t, output_t
   
- Key insight from benchmarks:
- - bmm is VERY fast for small kv_len (bs=4,kv=1024: 23.6µs vs Triton ~35µs)
- - bmm is VERY slow for large kv_len (bs=128,kv=8192,qseq=4: 976µs vs Triton ~574µs)
+ # ---------------------------------------------------------------------------
+ # MLA decode v120: FP8 pipeline with Q-DMA overlap
+ # - Overlap Q bf16->fp8 conversion with first tile DMA (free speedup)
+ # - Pad KV LDS stride to 580 bytes to eliminate 8-way bank conflicts
+ # - Remove dead singlehead kernel
+ # - v118 split-K tuning preserved
+ # ---------------------------------------------------------------------------
   
- Strategy:
- - If kv_len * qseq <= threshold: use batched bmm (avoids kernel launch overhead)
- - Else: use Triton flash attention (avoids materializing full score matrix)
+ os.environ["PYTORCH_ROCM_ARCH"] = "gfx950"
   
- Also uses fp8 KV for both paths where possible.
- """
+ HIP_SRC = r"""
+ #include <torch/extension.h>
+ #include <hip/hip_runtime.h>
   
- import torch
- import torch.nn.functional as F
- import triton
- import triton.language as tl
- from task import input_t, output_t
+ static constexpr int WARP_SIZE = 64;
+ static constexpr int NUM_WARPS = 4;
+ static constexpr int BLOCK_SIZE = WARP_SIZE * NUM_WARPS;
   
- SM_SCALE = 1.0 / (576 ** 0.5)
- LOG2E = 1.4426950408889634
- SM_SCALE_LOG2E = SM_SCALE * LOG2E
+ static constexpr int QK_DIM = 576;
+ static constexpr int V_DIM  = 512;
+ static constexpr int NUM_HEADS = 16;
+ static constexpr int NUM_K_CHUNKS = QK_DIM / 32;  // 18
+ static constexpr int NUM_K128_CHUNKS = (QK_DIM + 127) / 128;  // 5
+ static constexpr int SUPER_TILE = 32;
+ static constexpr int SV_CHUNKS = 8;
+ static constexpr int KV_TILE_BYTES = SUPER_TILE * QK_DIM;  // 18432
   
+ // 18432 bytes / 16 bytes per uint4 / 256 threads = 4.5 -> 5 rounds
+ static constexpr int PF_UINT4S = KV_TILE_BYTES / 16;  // 1152
+ static constexpr int PF_ROUNDS = (PF_UINT4S + BLOCK_SIZE - 1) / BLOCK_SIZE;  // 5
   
- # ==================== Triton Flash Attention (from sub111) ====================
+ typedef float __attribute__((ext_vector_type(4))) v4f32;
+ typedef unsigned int __attribute__((ext_vector_type(4))) u32x4;
+ typedef int __attribute__((ext_vector_type(4))) i32x4;
+ typedef int __attribute__((ext_vector_type(8))) i32x8;
+ typedef unsigned int __attribute__((address_space(3)))* lds_ptr_t;
   
- @triton.jit
- def _flash_fused(
-     Q_ptr, KV_ptr, O_ptr,
-     qo_indptr_ptr, kv_indptr_ptr,
-     sm_scale_log2e,
-     stride_q0, stride_q1,
-     stride_kv0,
-     stride_o0, stride_o1,
-     num_heads: tl.constexpr,
-     BLOCK_M: tl.constexpr,
-     BLOCK_KV: tl.constexpr,
-     D_TILE: tl.constexpr,
-     V_DIM: tl.constexpr,
-     HEAD_DIM: tl.constexpr,
- ):
-     batch = tl.program_id(0)
-     m_group = tl.program_id(1)
+ extern "C" __device__ void __llvm_amdgcn_raw_buffer_load_lds(
+     i32x4 rsrc, lds_ptr_t lds_ptr, int size,
+     int voffset, int soffset, int offset, int aux)
+     __asm("llvm.amdgcn.raw.buffer.load.lds");
   
-     kv_start = tl.load(kv_indptr_ptr + batch)
-     kv_end = tl.load(kv_indptr_ptr + batch + 1)
-     kv_len = kv_end - kv_start
-     q_start = tl.load(qo_indptr_ptr + batch)
-     q_end = tl.load(qo_indptr_ptr + batch + 1)
-     q_len = q_end - q_start
+ struct buffer_resource { uint64_t ptr; uint32_t range; uint32_t config; };
   
-     total_m = q_len * num_heads
-     m_start = m_group * BLOCK_M
-     m_range = tl.arange(0, BLOCK_M)
-     m_idx = m_start + m_range
-     m_mask = m_idx < total_m
+ __device__ __forceinline__ i32x4 make_buffer_rsrc(const void* p, uint32_t bytes) {
+     buffer_resource r = {reinterpret_cast<uint64_t>(p), bytes, 0x110000};
+     return *reinterpret_cast<i32x4*>(&r);
+ }
   
-     qi_local = m_idx // num_heads
-     hi = m_idx % num_heads
-     qi_global = q_start + qi_local
-     q_base = qi_global * stride_q0 + hi * stride_q1
+ __device__ __forceinline__ float bf16_to_f32(unsigned short v) {
+     return __uint_as_float(static_cast<unsigned int>(v) << 16);
+ }
   
-     m_i = tl.full([BLOCK_M], float('-inf'), dtype=tl.float32)
-     l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
-     acc = tl.zeros([BLOCK_M, V_DIM], dtype=tl.float32)
+ __device__ __forceinline__ unsigned short f32_to_bf16(float v) {
+     unsigned int bits = __float_as_uint(v);
+     bits += 0x7FFF + ((bits >> 16) & 1);
+     return static_cast<unsigned short>(bits >> 16);
+ }
   
-     for kv_off in range(0, kv_len, BLOCK_KV):
-         kv_range = tl.arange(0, BLOCK_KV)
-         kv_valid = (kv_off + kv_range) < kv_len
-         kv_base = (kv_start + kv_off + kv_range) * stride_kv0
+ __device__ __forceinline__ float fp8_to_f32(unsigned char b) {
+     return __builtin_amdgcn_cvt_f32_fp8(static_cast<int>(b), 0);
+ }
   
-         scores = tl.zeros([BLOCK_M, BLOCK_KV], dtype=tl.float32)
-         for d_off in tl.static_range(0, HEAD_DIM, D_TILE):
-             d_range = tl.arange(0, D_TILE)
-             q_chunk = tl.load(
-                 Q_ptr + q_base[:, None] + d_off + d_range[None, :],
-                 mask=m_mask[:, None], other=0.0
-             ).to(tl.bfloat16)
-             k_chunk = tl.load(
-                 KV_ptr + kv_base[:, None] + d_off + d_range[None, :],
-                 mask=kv_valid[:, None], other=0.0
-             ).to(tl.bfloat16)
-             scores += tl.dot(q_chunk, tl.trans(k_chunk))
+ __device__ __forceinline__ v4f32 mfma_f32_16x16x128_fp8(
+     i32x8 A, i32x8 B, v4f32 C)
+ {
+     v4f32 D;
+     asm volatile(
+         "v_mfma_f32_16x16x128_f8f6f4 %0, %1, %2, %3 cbsz:0 blgp:0"
+         : "=v"(D) : "v"(A), "v"(B), "v"(C));
+     return D;
+ }
   
-         scores *= sm_scale_log2e
-         scores = tl.where(kv_valid[None, :], scores, float('-inf'))
+ // =========================================================================
+ // Full MFMA pipeline kernel with double-buffered LDS + K=128 QK MFMA
+ // =========================================================================
   
-         m_ij = tl.max(scores, axis=1)
-         new_m = tl.maximum(m_i, m_ij)
-         alpha = tl.math.exp2(m_i - new_m)
-         p = tl.math.exp2(scores - new_m[:, None])
-         l_i = l_i * alpha + tl.sum(p, axis=1)
-         acc = acc * alpha[:, None]
-         m_i = new_m
+ __global__ __launch_bounds__(256, 3)
+ void mla_mfma_pipeline_kernel(
+     const unsigned short* __restrict__ q_ptr,
+     const unsigned char*  __restrict__ kv_ptr,
+     float*                __restrict__ partial_m,
+     float*                __restrict__ partial_l,
+     float*                __restrict__ partial_acc,
+     unsigned short*       __restrict__ out_ptr,
+     const int*            __restrict__ qo_indptr,
+     const int*            __restrict__ kv_indptr,
+     const float*          __restrict__ kv_scale_ptr,
+     const int  num_splits,
+     const float sm_scale)
+ {
+     const int split_idx = blockIdx.x;
+     const int batch_idx = blockIdx.y;
+     const int warp_id = threadIdx.x / WARP_SIZE;
+     const int lane_id = threadIdx.x % WARP_SIZE;
+     const int tid = threadIdx.x;
   
-         v_range = tl.arange(0, V_DIM)
-         v_block = tl.load(
-             KV_ptr + kv_base[:, None] + v_range[None, :],
-             mask=kv_valid[:, None], other=0.0
-         ).to(tl.bfloat16)
-         acc += tl.dot(p.to(tl.bfloat16), v_block)
+     const float score_scale = sm_scale * (*kv_scale_ptr);
   
-     result = acc / l_i[:, None]
-     o_base = qi_global * stride_o0 + hi * stride_o1
-     v_range = tl.arange(0, V_DIM)
-     tl.store(O_ptr + o_base[:, None] + v_range[None, :],
-              result.to(tl.bfloat16), mask=m_mask[:, None])
+     const int q_start = qo_indptr[batch_idx];
+     const int kv_start = kv_indptr[batch_idx];
+     const int kv_end = kv_indptr[batch_idx + 1];
+     const int kv_len = kv_end - kv_start;
   
+     const int tps = (kv_len + num_splits - 1) / num_splits;
+     const int split_kv_start = kv_start + split_idx * tps;
+     const int split_kv_end = min(split_kv_start + tps, kv_end);
   
- @triton.jit
- def _flash_splitk(
-     Q_ptr, KV_ptr,
-     Acc_ptr, Max_ptr, Sum_ptr,
-     qo_indptr_ptr, kv_indptr_ptr,
-     sm_scale_log2e,
-     stride_q0, stride_q1,
-     stride_kv0,
-     num_heads: tl.constexpr,
-     num_splits: tl.constexpr,
-     num_m_groups: tl.constexpr,
-     BLOCK_M: tl.constexpr,
-     BLOCK_KV: tl.constexpr,
-     D_TILE: tl.constexpr,
-     V_DIM: tl.constexpr,
-     HEAD_DIM: tl.constexpr,
- ):
-     batch = tl.program_id(0)
-     m_group = tl.program_id(1)
-     split = tl.program_id(2)
+     const int mr = lane_id & 0xF;
+     const int kg = lane_id >> 4;
   
-     kv_start = tl.load(kv_indptr_ptr + batch)
-     kv_end = tl.load(kv_indptr_ptr + batch + 1)
-     kv_len = kv_end - kv_start
-     q_start = tl.load(qo_indptr_ptr + batch)
-     q_end = tl.load(qo_indptr_ptr + batch + 1)
-     q_len = q_end - q_start
+     if (split_kv_start >= kv_end) {
+         if (lane_id < 16 && warp_id == 0) {
+             int head = lane_id;
+             int off = (batch_idx * NUM_HEADS + head) * num_splits + split_idx;
+             partial_m[off] = -1e30f;
+             partial_l[off] = 0.0f;
+         }
+         return;
+     }
   
-     total_m = q_len * num_heads
-     m_start = m_group * BLOCK_M
-     m_range = tl.arange(0, BLOCK_M)
-     m_idx = m_start + m_range
-     m_mask = m_idx < total_m
+     // ===== LDS: double-buffered KV + per-warp W =====
+     __shared__ __align__(16) unsigned char kv_lds[2][KV_TILE_BYTES];
+     __shared__ float s_W[NUM_WARPS][16][33];
   
-     qi_local = m_idx // num_heads
-     hi = m_idx % num_heads
-     qi_global = q_start + qi_local
-     q_base = qi_global * stride_q0 + hi * stride_q1
+     // ===== Super-tile iteration setup =====
+     const int total_tokens = split_kv_end - split_kv_start;
+     const int num_st = (total_tokens + SUPER_TILE - 1) / SUPER_TILE;
   
-     kv_per_split = (kv_len + num_splits - 1) // num_splits
-     split_kv_start = split * kv_per_split
-     split_kv_end = tl.minimum(split_kv_start + kv_per_split, kv_len)
+     // ===== PROLOGUE: issue DMA FIRST, then Q prep overlaps with DMA =====
+     {
+         const int first_bytes = min(SUPER_TILE, total_tokens) * QK_DIM;
+         const unsigned char* __restrict__ src0 = kv_ptr +
+             static_cast<long long>(split_kv_start) * QK_DIM;
+         i32x4 srsrc = make_buffer_rsrc(src0, first_bytes);
+         #pragma unroll
+         for (int r = 0; r < PF_ROUNDS; r++) {
+             int off = tid * 16 + r * BLOCK_SIZE * 16;
+             lds_ptr_t ldp = (lds_ptr_t)(reinterpret_cast<uintptr_t>(kv_lds[0]) + off);
+             __llvm_amdgcn_raw_buffer_load_lds(srsrc, ldp, 16, off, 0, 0, 2);
+         }
+     }
   
-     m_i = tl.full([BLOCK_M], float('-inf'), dtype=tl.float32)
-     l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
-     acc = tl.zeros([BLOCK_M, V_DIM], dtype=tl.float32)
+     // ===== Q preload as FP8 (overlapped with DMA in flight) =====
+     const unsigned short* qh = q_ptr +
+         (static_cast<long long>(q_start) * NUM_HEADS + mr) * QK_DIM;
   
-     for kv_off in range(split_kv_start, split_kv_end, BLOCK_KV):
-         kv_range = tl.arange(0, BLOCK_KV)
-         kv_valid = (kv_off + kv_range) < split_kv_end
-         kv_base = (kv_start + kv_off + kv_range) * stride_kv0
+     i32x8 q_128[NUM_K128_CHUNKS];
+     #pragma unroll
+     for (int c = 0; c < NUM_K128_CHUNKS; c++) {
+         unsigned int w[8];
+         int base1 = c * 128 + 16 * kg;
+         #pragma unroll
+         for (int i = 0; i < 4; i++) {
+             int d = base1 + i * 4;
+             float f0 = (d     < QK_DIM) ? bf16_to_f32(qh[d])     : 0.f;
+             float f1 = (d + 1 < QK_DIM) ? bf16_to_f32(qh[d + 1]) : 0.f;
+             float f2 = (d + 2 < QK_DIM) ? bf16_to_f32(qh[d + 2]) : 0.f;
+             float f3 = (d + 3 < QK_DIM) ? bf16_to_f32(qh[d + 3]) : 0.f;
+             unsigned int pk = __builtin_amdgcn_cvt_pk_fp8_f32(f0, f1, 0, false);
+             pk = __builtin_amdgcn_cvt_pk_fp8_f32(f2, f3, pk, true);
+             w[i] = pk;
+         }
+         int base2 = c * 128 + 64 + 16 * kg;
+         #pragma unroll
+         for (int i = 0; i < 4; i++) {
+             int d = base2 + i * 4;
+             float f0 = (d     < QK_DIM) ? bf16_to_f32(qh[d])     : 0.f;
+             float f1 = (d + 1 < QK_DIM) ? bf16_to_f32(qh[d + 1]) : 0.f;
+             float f2 = (d + 2 < QK_DIM) ? bf16_to_f32(qh[d + 2]) : 0.f;
+             float f3 = (d + 3 < QK_DIM) ? bf16_to_f32(qh[d + 3]) : 0.f;
+             unsigned int pk = __builtin_amdgcn_cvt_pk_fp8_f32(f0, f1, 0, false);
+             pk = __builtin_amdgcn_cvt_pk_fp8_f32(f2, f3, pk, true);
+             w[4 + i] = pk;
+         }
+         q_128[c] = *reinterpret_cast<i32x8*>(w);
+     }
   
-         scores = tl.zeros([BLOCK_M, BLOCK_KV], dtype=tl.float32)
-         for d_off in tl.static_range(0, HEAD_DIM, D_TILE):
-             d_range = tl.arange(0, D_TILE)
-             q_chunk = tl.load(
-                 Q_ptr + q_base[:, None] + d_off + d_range[None, :],
-                 mask=m_mask[:, None], other=0.0
-             ).to(tl.bfloat16)
-             k_chunk = tl.load(
-                 KV_ptr + kv_base[:, None] + d_off + d_range[None, :],
-                 mask=kv_valid[:, None], other=0.0
-             ).to(tl.bfloat16)
-             scores += tl.dot(q_chunk, tl.trans(k_chunk))
+     // ===== V accumulators + softmax state =====
+     float vacc[SV_CHUNKS][4];
+     #pragma unroll
+     for (int i = 0; i < SV_CHUNKS; i++)
+         vacc[i][0] = vacc[i][1] = vacc[i][2] = vacc[i][3] = 0.0f;
   
-         scores *= sm_scale_log2e
-         scores = tl.where(kv_valid[None, :], scores, float('-inf'))
+     float mv[4] = {-1e30f, -1e30f, -1e30f, -1e30f};
+     float lv[4] = {0.0f, 0.0f, 0.0f, 0.0f};
   
-         m_ij = tl.max(scores, axis=1)
-         new_m = tl.maximum(m_i, m_ij)
-         alpha = tl.math.exp2(m_i - new_m)
-         p = tl.math.exp2(scores - new_m[:, None])
-         l_i = l_i * alpha + tl.sum(p, axis=1)
-         acc = acc * alpha[:, None]
-         m_i = new_m
+     // ===== Wait for DMA (Q prep ran while DMA was in flight) =====
+     asm volatile("s_waitcnt vmcnt(0)" ::: "memory");
+     __syncthreads();
   
-         v_range = tl.arange(0, V_DIM)
-         v_block = tl.load(
-             KV_ptr + kv_base[:, None] + v_range[None, :],
-             mask=kv_valid[:, None], other=0.0
-         ).to(tl.bfloat16)
-         acc += tl.dot(p.to(tl.bfloat16), v_block)
+     int cur_buf = 0;
   
-     flat_idx = (batch * num_m_groups + m_group) * num_splits + split
-     acc_base = flat_idx * BLOCK_M * V_DIM
-     ml_base = flat_idx * BLOCK_M
+     // ===== MAIN LOOP =====
+     for (int st_idx = 0; st_idx < num_st; st_idx++) {
+         const int stcnt = min(SUPER_TILE, total_tokens - st_idx * SUPER_TILE);
+         const int ta = min(16, stcnt);
+         const int tb = max(0, stcnt - 16);
+         const unsigned char* kv_cur = kv_lds[cur_buf];
   
-     v_range = tl.arange(0, V_DIM)
-     tl.store(Acc_ptr + acc_base + m_range[:, None] * V_DIM + v_range[None, :],
-              acc, mask=m_mask[:, None])
-     tl.store(Max_ptr + ml_base + m_range, m_i, mask=m_mask)
-     tl.store(Sum_ptr + ml_base + m_range, l_i, mask=m_mask)
+         // ---- PREFETCH: GLOBAL_LOAD_LDS for NEXT tile ----
+         __builtin_amdgcn_s_setprio(3);
+         const int nxt_buf = cur_buf ^ 1;
+         const bool has_next = (st_idx + 1 < num_st);
   
+         if (has_next) {
+             const int nxt_start = split_kv_start + (st_idx + 1) * SUPER_TILE;
+             const int nxt_bytes = min(SUPER_TILE, split_kv_end - nxt_start) * QK_DIM;
+             const unsigned char* __restrict__ nsrc = kv_ptr +
+                 static_cast<long long>(nxt_start) * QK_DIM;
+             i32x4 srsrc = make_buffer_rsrc(nsrc, nxt_bytes);
   
- @triton.jit
- def _reduce_splitk(
-     Acc_ptr, Max_ptr, Sum_ptr, O_ptr,
-     qo_indptr_ptr,
-     stride_o0, stride_o1,
-     num_heads: tl.constexpr,
-     num_splits: tl.constexpr,
-     num_m_groups: tl.constexpr,
-     BLOCK_M: tl.constexpr,
-     V_DIM: tl.constexpr,
- ):
-     batch = tl.program_id(0)
-     m_group = tl.program_id(1)
-     q_start = tl.load(qo_indptr_ptr + batch)
-     q_end = tl.load(qo_indptr_ptr + batch + 1)
-     q_len = q_end - q_start
-     total_m = q_len * num_heads
+             #pragma unroll
+             for (int r = 0; r < PF_ROUNDS; r++) {
+                 int off = tid * 16 + r * BLOCK_SIZE * 16;
+                 lds_ptr_t ldp = (lds_ptr_t)(reinterpret_cast<uintptr_t>(kv_lds[nxt_buf]) + off);
+                 __llvm_amdgcn_raw_buffer_load_lds(srsrc, ldp, 16, off, 0, 0, 2);
+             }
+         }
   
-     m_start = m_group * BLOCK_M
-     m_range = tl.arange(0, BLOCK_M)
-     m_idx = m_start + m_range
-     m_mask = m_idx < total_m
-     qi_local = m_idx // num_heads
-     hi = m_idx % num_heads
-     qi_global = q_start + qi_local
+         // ---- QK: K=128 MFMA (FP8xFP8), INTERLEAVED CK layout ----
+         // B (FP8) also uses interleaved: v0-v3 = k[16*kg..+15], v4-v7 = k[64+16*kg..+15]
+         __builtin_amdgcn_sched_barrier(0);
+         __builtin_amdgcn_s_setprio(0);
+         v4f32 ca = {0, 0, 0, 0};
+         #pragma unroll
+         for (int c = 0; c < NUM_K128_CHUNKS; c++) {
+             i32x8 b_128 = {};
+             if (mr < ta) {
+                 int base1 = mr * QK_DIM + c * 128 + 16 * kg;
+                 #pragma unroll
+                 for (int i = 0; i < 4; i++) {
+                     int off = base1 + i * 4;
+                     if (c * 128 + 16 * kg + i * 4 + 4 <= QK_DIM)
+                         b_128[i] = *reinterpret_cast<const int*>(&kv_cur[off]);
+                 }
+                 int base2 = mr * QK_DIM + c * 128 + 64 + 16 * kg;
+                 #pragma unroll
+                 for (int i = 0; i < 4; i++) {
+                     int off = base2 + i * 4;
+                     if (c * 128 + 64 + 16 * kg + i * 4 + 4 <= QK_DIM)
+                         b_128[4 + i] = *reinterpret_cast<const int*>(&kv_cur[off]);
+                 }
+             }
+             ca = mfma_f32_16x16x128_fp8(q_128[c], b_128, ca);
+         }
   
-     base = (batch * num_m_groups + m_group) * num_splits
-     global_max = tl.full([BLOCK_M], float('-inf'), dtype=tl.float32)
-     for s in range(num_splits):
-         m_s = tl.load(Max_ptr + (base + s) * BLOCK_M + m_range, mask=m_mask, other=float('-inf'))
-         global_max = tl.maximum(global_max, m_s)
+         v4f32 cb = {0, 0, 0, 0};
+         if (tb > 0) {
+             #pragma unroll
+             for (int c = 0; c < NUM_K128_CHUNKS; c++) {
+                 i32x8 b_128 = {};
+                 if (mr < tb) {
+                     int base1 = (16 + mr) * QK_DIM + c * 128 + 16 * kg;
+                     #pragma unroll
+                     for (int i = 0; i < 4; i++) {
+                         int off = base1 + i * 4;
+                         if (c * 128 + 16 * kg + i * 4 + 4 <= QK_DIM)
+                             b_128[i] = *reinterpret_cast<const int*>(&kv_cur[off]);
+                     }
+                     int base2 = (16 + mr) * QK_DIM + c * 128 + 64 + 16 * kg;
+                     #pragma unroll
+                     for (int i = 0; i < 4; i++) {
+                         int off = base2 + i * 4;
+                         if (c * 128 + 64 + 16 * kg + i * 4 + 4 <= QK_DIM)
+                             b_128[4 + i] = *reinterpret_cast<const int*>(&kv_cur[off]);
+                     }
+                 }
+                 cb = mfma_f32_16x16x128_fp8(q_128[c], b_128, cb);
+             }
+         }
   
-     v_range = tl.arange(0, V_DIM)
-     total_acc = tl.zeros([BLOCK_M, V_DIM], dtype=tl.float32)
-     total_l = tl.zeros([BLOCK_M], dtype=tl.float32)
-     for s in range(num_splits):
-         flat_idx = base + s
-         m_s = tl.load(Max_ptr + flat_idx * BLOCK_M + m_range, mask=m_mask, other=float('-inf'))
-         l_s = tl.load(Sum_ptr + flat_idx * BLOCK_M + m_range, mask=m_mask, other=0.0)
-         alpha = tl.math.exp2(m_s - global_max)
-         total_l += l_s * alpha
-         acc_base = flat_idx * BLOCK_M * V_DIM
-         acc_s = tl.load(Acc_ptr + acc_base + m_range[:, None] * V_DIM + v_range[None, :],
-                         mask=m_mask[:, None], other=0.0)
-         total_acc += acc_s * alpha[:, None]
+         // ---- Online softmax (16-lane reduce only: offsets 8,4,2,1) ----
+         float sa0 = ca[0] * score_scale, sa1 = ca[1] * score_scale;
+         float sa2 = ca[2] * score_scale, sa3 = ca[3] * score_scale;
+         float sb0 = cb[0] * score_scale, sb1 = cb[1] * score_scale;
+         float sb2 = cb[2] * score_scale, sb3 = cb[3] * score_scale;
   
-     result = total_acc / total_l[:, None]
-     o_base = qi_global * stride_o0 + hi * stride_o1
-     tl.store(O_ptr + o_base[:, None] + v_range[None, :],
-              result.to(tl.bfloat16), mask=m_mask[:, None])
+         if (mr >= ta) { sa0 = sa1 = sa2 = sa3 = -1e30f; }
+         if (mr >= tb) { sb0 = sb1 = sb2 = sb3 = -1e30f; }
   
+         float tm0 = fmaxf(sa0, sb0), tm1 = fmaxf(sa1, sb1);
+         float tm2 = fmaxf(sa2, sb2), tm3 = fmaxf(sa3, sb3);
+         #pragma unroll
+         for (int off = 8; off >= 1; off >>= 1) {
+             tm0 = fmaxf(tm0, __shfl_xor(tm0, off));
+             tm1 = fmaxf(tm1, __shfl_xor(tm1, off));
+             tm2 = fmaxf(tm2, __shfl_xor(tm2, off));
+             tm3 = fmaxf(tm3, __shfl_xor(tm3, off));
+         }
   
- # ==================== BMM Path ====================
+         float nm0 = fmaxf(mv[0], tm0), nm1 = fmaxf(mv[1], tm1);
+         float nm2 = fmaxf(mv[2], tm2), nm3 = fmaxf(mv[3], tm3);
+         float rc0 = __expf(mv[0] - nm0), rc1 = __expf(mv[1] - nm1);
+         float rc2 = __expf(mv[2] - nm2), rc3 = __expf(mv[3] - nm3);
+         mv[0] = nm0; mv[1] = nm1; mv[2] = nm2; mv[3] = nm3;
   
- def _bmm_attention(q, kv_bf16, qo_indptr, kv_indptr, config):
-     num_heads = config["num_heads"]
-     v_head_dim = config["v_head_dim"]
-     batch_size = config["batch_size"]
-     q_seq_len = config["q_seq_len"]
-     kv_seq_len = config["kv_seq_len"]
-     total_q = q.shape[0]
+         #pragma unroll
+         for (int vc = 0; vc < SV_CHUNKS; vc++) {
+             vacc[vc][0] *= rc0; vacc[vc][1] *= rc1;
+             vacc[vc][2] *= rc2; vacc[vc][3] *= rc3;
+         }
   
-     q_batched = q.view(batch_size, q_seq_len, num_heads, 576).reshape(batch_size, q_seq_len * num_heads, 576)
-     kv_batched = kv_bf16.view(batch_size, kv_seq_len, 576)
+         float wa0 = (mr < ta) ? __expf(sa0 - nm0) : 0.f;
+         float wa1 = (mr < ta) ? __expf(sa1 - nm1) : 0.f;
+         float wa2 = (mr < ta) ? __expf(sa2 - nm2) : 0.f;
+         float wa3 = (mr < ta) ? __expf(sa3 - nm3) : 0.f;
+         float wb0 = (mr < tb) ? __expf(sb0 - nm0) : 0.f;
+         float wb1 = (mr < tb) ? __expf(sb1 - nm1) : 0.f;
+         float wb2 = (mr < tb) ? __expf(sb2 - nm2) : 0.f;
+         float wb3 = (mr < tb) ? __expf(sb3 - nm3) : 0.f;
   
-     scores = torch.bmm(q_batched, kv_batched.transpose(1, 2))
-     scores.mul_(SM_SCALE)
-     scores = F.softmax(scores, dim=-1)
+         float dl0 = wa0 + wb0, dl1 = wa1 + wb1;
+         float dl2 = wa2 + wb2, dl3 = wa3 + wb3;
+         #pragma unroll
+         for (int off = 8; off >= 1; off >>= 1) {
+             dl0 += __shfl_xor(dl0, off); dl1 += __shfl_xor(dl1, off);
+             dl2 += __shfl_xor(dl2, off); dl3 += __shfl_xor(dl3, off);
+         }
+         lv[0] = lv[0] * rc0 + dl0; lv[1] = lv[1] * rc1 + dl1;
+         lv[2] = lv[2] * rc2 + dl2; lv[3] = lv[3] * rc3 + dl3;
   
-     v_batched = kv_batched[:, :, :v_head_dim]
-     output = torch.bmm(scores.to(v_batched.dtype), v_batched)
+         // ---- W to per-warp LDS ----
+         s_W[warp_id][kg * 4    ][mr]      = wa0;
+         s_W[warp_id][kg * 4 + 1][mr]      = wa1;
+         s_W[warp_id][kg * 4 + 2][mr]      = wa2;
+         s_W[warp_id][kg * 4 + 3][mr]      = wa3;
+         s_W[warp_id][kg * 4    ][16 + mr] = wb0;
+         s_W[warp_id][kg * 4 + 1][16 + mr] = wb1;
+         s_W[warp_id][kg * 4 + 2][16 + mr] = wb2;
+         s_W[warp_id][kg * 4 + 3][16 + mr] = wb3;
   
-     return output.view(batch_size, q_seq_len, num_heads, v_head_dim).reshape(total_q, num_heads, v_head_dim).to(torch.bfloat16)
+         asm volatile("s_waitcnt lgkmcnt(0)" ::: "memory");
   
+         float wvals[8];
+         #pragma unroll
+         for (int i = 0; i < 8; i++)
+             wvals[i] = s_W[warp_id][mr][i * 4 + kg];
   
- # ==================== Triton Path ====================
+         unsigned int wlo = __builtin_amdgcn_cvt_pk_fp8_f32(wvals[0], wvals[1], 0, false);
+         wlo = __builtin_amdgcn_cvt_pk_fp8_f32(wvals[2], wvals[3], wlo, true);
+         unsigned int whi = __builtin_amdgcn_cvt_pk_fp8_f32(wvals[4], wvals[5], 0, false);
+         whi = __builtin_amdgcn_cvt_pk_fp8_f32(wvals[6], wvals[7], whi, true);
+         long w_a = static_cast<long>(wlo) | (static_cast<long>(whi) << 32);
   
- def _triton_flash_attention(q, kv_flat, qo_indptr, kv_indptr, config):
-     num_heads = config["num_heads"]
-     v_head_dim = config["v_head_dim"]
-     q_seq_len = config["q_seq_len"]
-     batch_size = config["batch_size"]
+         // ---- SV MFMA from current LDS (each warp: 128 V dims) ----
+         __builtin_amdgcn_sched_barrier(0);
+         __builtin_amdgcn_s_setprio(1);
+         // Use dword loads + byte extract to reduce LDS instruction count
+         const int v_warp_base = warp_id * SV_CHUNKS * 16;
   
-     total_q = q.shape[0]
-     o = torch.empty((total_q, num_heads, v_head_dim), dtype=torch.bfloat16, device="cuda")
+         #pragma unroll
+         for (int vc = 0; vc < SV_CHUNKS; vc++) {
+             const int vd = v_warp_base + vc * 16 + mr;
+             const int vd_align = vd & ~3;
+             const int vd_shift = (vd & 3) * 8;
   
-     total_m = q_seq_len * num_heads
-     BLOCK_M = 16 if total_m <= 16 else (32 if total_m <= 32 else 64)
-     BLOCK_KV = 64
-     D_TILE = 64
+             unsigned int blo = 0, bhi = 0;
+             if (vd < V_DIM) {
+                 #pragma unroll
+                 for (int i = 0; i < 4; i++) {
+                     int tok = i * 4 + kg;
+                     unsigned int dw = (tok < stcnt) ?
+                         *reinterpret_cast<const unsigned int*>(&kv_cur[tok * QK_DIM + vd_align]) : 0u;
+                     blo |= ((dw >> vd_shift) & 0xFF) << (i * 8);
+                 }
+                 #pragma unroll
+                 for (int i = 0; i < 4; i++) {
+                     int tok = (i + 4) * 4 + kg;
+                     unsigned int dw = (tok < stcnt) ?
+                         *reinterpret_cast<const unsigned int*>(&kv_cur[tok * QK_DIM + vd_align]) : 0u;
+                     bhi |= ((dw >> vd_shift) & 0xFF) << (i * 8);
+                 }
+             }
+             long v_b = static_cast<long>(blo) | (static_cast<long>(bhi) << 32);
   
-     num_m_groups = (total_m + BLOCK_M - 1) // BLOCK_M
-     total_programs_base = batch_size * num_m_groups
+             v4f32 sc = {vacc[vc][0], vacc[vc][1], vacc[vc][2], vacc[vc][3]};
+             sc = __builtin_amdgcn_mfma_f32_16x16x32_fp8_fp8(w_a, v_b, sc, 0, 0, 0);
+             vacc[vc][0] = sc[0]; vacc[vc][1] = sc[1];
+             vacc[vc][2] = sc[2]; vacc[vc][3] = sc[3];
+         }
   
-     if total_programs_base >= 128:
-         grid = (batch_size, num_m_groups)
-         _flash_fused[grid](
-             q, kv_flat, o, qo_indptr, kv_indptr,
-             SM_SCALE_LOG2E,
-             q.stride(0), q.stride(1), kv_flat.stride(0),
-             o.stride(0), o.stride(1),
-             num_heads=num_heads, BLOCK_M=BLOCK_M, BLOCK_KV=BLOCK_KV,
-             D_TILE=D_TILE, V_DIM=512, HEAD_DIM=576,
+         // ---- Wait for GLOBAL_LOAD_LDS and flip ----
+         __builtin_amdgcn_sched_barrier(0);
+         __builtin_amdgcn_s_setprio(3);
+         if (has_next) {
+             asm volatile("s_waitcnt vmcnt(0)" ::: "memory");
+             __syncthreads();
+             cur_buf = nxt_buf;
+         }
+     }
+  
+     if (num_splits == 1) {
+         const float kv_scale = *kv_scale_ptr;
+         const int vwb = warp_id * SV_CHUNKS * 16;
+         #pragma unroll
+         for (int vc = 0; vc < SV_CHUNKS; vc++) {
+             int vd = vwb + vc * 16 + mr;
+             if (vd < V_DIM) {
+                 #pragma unroll
+                 for (int r = 0; r < 4; r++) {
+                     int head = kg * 4 + r;
+                     float inv_l = (lv[r] > 0.f) ? (kv_scale / lv[r]) : 0.f;
+                     long long idx = (static_cast<long long>(q_start) * NUM_HEADS + head) * V_DIM + vd;
+                     out_ptr[idx] = f32_to_bf16(vacc[vc][r] * inv_l);
+                 }
+             }
+         }
+     } else {
+         if (warp_id == 0 && mr == 0) {
+             #pragma unroll
+             for (int r = 0; r < 4; r++) {
+                 int head = kg * 4 + r;
+                 int off = (batch_idx * NUM_HEADS + head) * num_splits + split_idx;
+                 partial_m[off] = mv[r];
+                 partial_l[off] = lv[r];
+             }
+         }
+         const int vwb = warp_id * SV_CHUNKS * 16;
+         #pragma unroll
+         for (int vc = 0; vc < SV_CHUNKS; vc++) {
+             int vd = vwb + vc * 16 + mr;
+             if (vd < V_DIM) {
+                 #pragma unroll
+                 for (int r = 0; r < 4; r++) {
+                     int head = kg * 4 + r;
+                     int off = (batch_idx * NUM_HEADS + head) * num_splits + split_idx;
+                     partial_acc[static_cast<long long>(off) * V_DIM + vd] = vacc[vc][r];
+                 }
+             }
+         }
+     }
+ }
+  
+ // =========================================================================
+ // Reduce kernel
+ // =========================================================================
+  
+ __global__ __launch_bounds__(512)
+ void mla_reduce_kernel(
+     const float*          __restrict__ partial_m,
+     const float*          __restrict__ partial_l,
+     const float*          __restrict__ partial_acc,
+     unsigned short*       __restrict__ out_ptr,
+     const float*          __restrict__ kv_scale_ptr,
+     const int  num_splits)
+ {
+     const int item_idx = blockIdx.x;
+     const int tid = threadIdx.x;
+     const float kv_scale = *kv_scale_ptr;
+  
+     __shared__ float s_corr[128];
+     __shared__ float s_inv_l;
+  
+     if (tid == 0) {
+         float merged_m = -1e30f;
+         for (int s = 0; s < num_splits; ++s)
+             merged_m = fmaxf(merged_m, partial_m[item_idx * num_splits + s]);
+         float merged_l = 0.0f;
+         for (int s = 0; s < num_splits; ++s) {
+             float l = partial_l[item_idx * num_splits + s];
+             float c = (l > 0.f) ? __expf(partial_m[item_idx * num_splits + s] - merged_m) : 0.f;
+             s_corr[s] = c;
+             merged_l += l * c;
+         }
+         s_inv_l = (merged_l > 0.f) ? (kv_scale / merged_l) : 0.f;
+     }
+     __syncthreads();
+  
+     if (tid < V_DIM) {
+         float val = 0.0f;
+         for (int s = 0; s < num_splits; ++s) {
+             val += partial_acc[(static_cast<long long>(item_idx) * num_splits + s) * V_DIM + tid]
+                    * s_corr[s];
+         }
+         out_ptr[static_cast<long long>(item_idx) * V_DIM + tid] = f32_to_bf16(val * s_inv_l);
+     }
+ }
+  
+ torch::Tensor mla_decode(
+     torch::Tensor q, torch::Tensor kv_buffer,
+     torch::Tensor qo_indptr, torch::Tensor kv_indptr,
+     torch::Tensor kv_scale_tensor,
+     int64_t num_heads, int64_t num_splits,
+     float sm_scale,
+     torch::Tensor partial_m, torch::Tensor partial_l,
+     torch::Tensor partial_acc,
+     torch::Tensor output)
+ {
+     const int batch_size = qo_indptr.size(0) - 1;
+     const int num_items = batch_size * static_cast<int>(num_heads);
+  
+     dim3 grid1(static_cast<int>(num_splits), batch_size);
+     dim3 block1(BLOCK_SIZE);
+     mla_mfma_pipeline_kernel<<<grid1, block1>>>(
+         reinterpret_cast<const unsigned short*>(q.data_ptr()),
+         reinterpret_cast<const unsigned char*>(kv_buffer.data_ptr()),
+         partial_m.data_ptr<float>(), partial_l.data_ptr<float>(),
+         partial_acc.data_ptr<float>(),
+         reinterpret_cast<unsigned short*>(output.data_ptr()),
+         qo_indptr.data_ptr<int>(), kv_indptr.data_ptr<int>(),
+         kv_scale_tensor.data_ptr<float>(),
+         static_cast<int>(num_splits), sm_scale);
+     if (num_splits == 1) return output;
+  
+     dim3 grid2(num_items);
+     dim3 block2(512);
+     mla_reduce_kernel<<<grid2, block2>>>(
+         partial_m.data_ptr<float>(), partial_l.data_ptr<float>(),
+         partial_acc.data_ptr<float>(),
+         reinterpret_cast<unsigned short*>(output.data_ptr()),
+         kv_scale_tensor.data_ptr<float>(),
+         static_cast<int>(num_splits));
+  
+     return output;
+ }
+ """
+  
+ CPP_DECL = """
+ torch::Tensor mla_decode(
+     torch::Tensor q, torch::Tensor kv_buffer,
+     torch::Tensor qo_indptr, torch::Tensor kv_indptr,
+     torch::Tensor kv_scale_tensor,
+     int64_t num_heads, int64_t num_splits,
+     float sm_scale,
+     torch::Tensor partial_m, torch::Tensor partial_l,
+     torch::Tensor partial_acc,
+     torch::Tensor output);
+ """
+  
+ _module = load_inline(
+     name="mla_hip_v120_qdma_overlap",
+     cpp_sources=CPP_DECL,
+     cuda_sources=HIP_SRC,
+     functions=["mla_decode"],
+     extra_cuda_cflags=[
+         "-O3", "-std=c++17",
+         "-ffast-math", "-funsafe-math-optimizations", "-ffp-contract=fast",
+         "-fno-gpu-rdc",
+         "-mllvm", "-amdgpu-early-inline-all=true",
+         "-mllvm", "-amdgpu-function-calls=false",
+         "-mllvm", "-amdgpu-max-memory-clause=64",
+         "-mllvm", "-amdgpu-load-store-vectorizer",
+         "-mllvm", "-amdgpu-early-ifcvt",
+         "-mllvm", "-amdgpu-internalize-symbols",
+         "-mllvm", "-amdgpu-scalarize-global-loads",
+         "-mllvm", "-amdgpu-dpp-combine",
+         "-mllvm", "-amdgpu-enable-pre-ra-optimizations",
+         "-mllvm", "-amdgpu-promote-alloca-to-vector-limit=256",
+     ],
+     verbose=False,
+ )
+  
+ _buf_cache = {}
+  
+  
+ def _get_bufs(num_items, num_splits, total_q, num_heads, device):
+     key = (num_items, num_splits, total_q, device)
+     if key not in _buf_cache:
+         _buf_cache[key] = (
+             torch.empty((num_items * num_splits,), dtype=torch.float32, device=device),
+             torch.empty((num_items * num_splits,), dtype=torch.float32, device=device),
+             torch.empty((num_items * num_splits, 512), dtype=torch.float32, device=device),
+             torch.empty((total_q, num_heads, 512), dtype=torch.bfloat16, device=device),
          )
+     return _buf_cache[key]
+  
+  
+ def _choose_splits(batch_size, kv_len):
+     tiles = max(1, kv_len // 32)
+     max_useful = max(1, tiles // 2)
+  
+     if tiles > 64:
+         ideal = max(1, -(-768 // batch_size))
      else:
-         num_splits = max(1, min(32, 512 // max(1, total_programs_base)))
-         total_partials = batch_size * num_m_groups * num_splits
-         acc_partial = torch.empty((total_partials * BLOCK_M, 512), dtype=torch.float32, device="cuda")
-         max_partial = torch.empty((total_partials * BLOCK_M,), dtype=torch.float32, device="cuda")
-         sum_partial = torch.empty((total_partials * BLOCK_M,), dtype=torch.float32, device="cuda")
+         target_wgs = max(512, batch_size * 8)
+         ideal = max(1, target_wgs // batch_size)
   
-         _flash_splitk[(batch_size, num_m_groups, num_splits)](
-             q, kv_flat, acc_partial, max_partial, sum_partial,
-             qo_indptr, kv_indptr, SM_SCALE_LOG2E,
-             q.stride(0), q.stride(1), kv_flat.stride(0),
-             num_heads=num_heads, num_splits=num_splits, num_m_groups=num_m_groups,
-             BLOCK_M=BLOCK_M, BLOCK_KV=BLOCK_KV, D_TILE=D_TILE, V_DIM=512, HEAD_DIM=576,
-         )
-         _reduce_splitk[(batch_size, num_m_groups)](
-             acc_partial, max_partial, sum_partial, o, qo_indptr,
-             o.stride(0), o.stride(1),
-             num_heads=num_heads, num_splits=num_splits, num_m_groups=num_m_groups,
-             BLOCK_M=BLOCK_M, V_DIM=512,
-         )
+     splits = max(1, min(ideal, max_useful, 64))
   
-     return o
+     while splits > 1 and batch_size * splits > 912:
+         splits -= 1
   
+     return splits
   
- # ==================== Dispatch ====================
   
  def custom_kernel(data: input_t) -> output_t:
      q, kv_data, qo_indptr, kv_indptr, config = data
   
-     kv_seq_len = config["kv_seq_len"]
-     q_seq_len = config["q_seq_len"]
+     kv_buffer_fp8, kv_scale = kv_data["fp8"]
+     kv_buffer = kv_buffer_fp8.view(-1, 576)
+  
      batch_size = config["batch_size"]
      num_heads = config["num_heads"]
+     sm_scale = config["sm_scale"]
+     total_q = q.size(0)
+     num_items = batch_size * num_heads
   
-     # Heuristic: bmm is better when the score matrix is small
-     # score_matrix_size = batch_size * q_seq_len * num_heads * kv_seq_len
-     # bmm materializes the full score matrix in memory
-     # Flash attention doesn't, so it wins for large score matrices
-     score_size = q_seq_len * kv_seq_len
+     total_kv = kv_buffer.shape[0]
+     kv_len = total_kv // batch_size
   
-     if score_size <= 4096:  # e.g., qseq=1, kv≤4096 or qseq=4, kv≤1024
-         # Use batched bmm — faster for small problems
-         kv_bf16 = kv_data["bf16"]
-         return _bmm_attention(q, kv_bf16, qo_indptr, kv_indptr, config)
-     else:
-         # Use Triton flash attention — better for large score matrices
-         kv_bf16 = kv_data["bf16"]
-         kv_flat = kv_bf16.view(-1, 576)
-         return _triton_flash_attention(q, kv_flat, qo_indptr, kv_indptr, config)
+     num_splits = _choose_splits(batch_size, kv_len)
+  
+     pm, pl, pa, out = _get_bufs(num_items, num_splits, total_q, num_heads, q.device)
+  
+     return _module.mla_decode(
+         q, kv_buffer, qo_indptr, kv_indptr,
+         kv_scale,
+         num_heads, num_splits,
+         sm_scale,
+         pm, pl, pa, out)