submission 745757

#!POPCORN leaderboard amd-mixed-mla
#!POPCORN gpu MI355X

import os, math
os.environ['PYTORCH_ROCM_ARCH'] = 'gfx950'

import torch
from torch.utils.cpp_extension import load_inline

# MLA Decode Attention — MFMA-accelerated kernel
# Uses mfma_f32_16x16x128_f8f6f4 on gfx950 for score computation:
#   Q(16heads, 576dims) × K(576dims, 16tokens) → Scores(16, 16)
# Then scalar V accumulation from LDS.
# This eliminates warp_reduce_sum (the main bottleneck), replacing 16 serial
# butterfly reductions with 5 MFMA instructions per 16-token tile.

FP8_DTYPE = torch.float8_e4m3fn

mla_source = r'''
#undef __HIP_NO_HALF_CONVERSIONS__
#undef __HIP_NO_HALF_OPERATORS__
#include <torch/extension.h>
#include <hip/hip_runtime.h>
#include <hip/hip_bf16.h>
#include <float.h>
#include <algorithm>

#define NUM_HEADS 16
#define QK_HEAD_DIM 576
#define V_HEAD_DIM 512
#define WAVEFRONT_SIZE 64
#define HEADS_PER_WARP 4
#define WARPS_PER_BLOCK 4
#define BLOCK_SIZE (WARPS_PER_BLOCK * WAVEFRONT_SIZE)
#define LOG2E 1.44269504089f
#define TILE_TOKENS 32
#define SM_SCALE_INV_SQRT (1.0f / 24.0f)  // 1/sqrt(576) = 1/24

typedef int v8i __attribute__((ext_vector_type(8)));
typedef float v4f __attribute__((ext_vector_type(4)));

__device__ __forceinline__ void unpack_fp8x4(int packed,
    float& f0, float& f1, float& f2, float& f3) {
    f0 = __builtin_amdgcn_cvt_f32_fp8(packed, 0);
    f1 = __builtin_amdgcn_cvt_f32_fp8(packed, 1);
    f2 = __builtin_amdgcn_cvt_f32_fp8(packed, 2);
    f3 = __builtin_amdgcn_cvt_f32_fp8(packed, 3);
}

__device__ __forceinline__ float bpermute_f32(int byte_offset, float val) {
    int tmp = __builtin_bit_cast(int, val);
    tmp = __builtin_amdgcn_ds_bpermute(byte_offset, tmp);
    return __builtin_bit_cast(float, tmp);
}

__device__ __forceinline__ float warp_reduce_max(float val) {
    int lane = threadIdx.x & 63;
    #pragma unroll
    for (int offset = 1; offset < WAVEFRONT_SIZE; offset <<= 1) {
        float other = bpermute_f32(((lane ^ offset)) << 2, val);
        val = fmaxf(val, other);
    }
    return val;
}

__device__ __forceinline__ float warp_reduce_sum(float val) {
    int lane = threadIdx.x & 63;
    #pragma unroll
    for (int offset = 1; offset < WAVEFRONT_SIZE; offset <<= 1) {
        float other = bpermute_f32(((lane ^ offset)) << 2, val);
        val += other;
    }
    return val;
}

__device__ __forceinline__ float read_lane(float val, int src_lane) {
    return bpermute_f32(src_lane << 2, val);
}


// ============================================================================
// Phase 1: MFMA-accelerated partial attention
// Grid: (batch_size, num_kv_splits)
// Block: 256 threads (4 wavefronts)
// Each wavefront handles 4 heads using MFMA for score computation.
// KV tiles loaded cooperatively into LDS, shared across all wavefronts.
// ============================================================================
__global__
__launch_bounds__(BLOCK_SIZE)
__attribute__((amdgpu_num_vgpr(96)))
void mla_mfma_partial_kernel(
    const uint8_t* __restrict__ q,
    const uint8_t* __restrict__ kv,
    const float* __restrict__ q_scale,
    const float* __restrict__ kv_scale,
    const int* __restrict__ kv_indptr,
    float* __restrict__ partial_o,
    float* __restrict__ partial_max,
    float* __restrict__ partial_sum,
    int num_kv_splits
) {
    // LDS: KV tile (32 tokens × 576 bytes) + score buffer (4 warps × 4 heads × 32 tokens)
    __shared__ uint8_t kv_lds[TILE_TOKENS * QK_HEAD_DIM];    // 18432 bytes
    __shared__ float score_lds[WARPS_PER_BLOCK * HEADS_PER_WARP * TILE_TOKENS]; // 2048 bytes
    // Total LDS: ~20 KB (within 64 KB limit)

    const int batch_item = blockIdx.x;
    const int split_idx  = blockIdx.y;
    const int tid        = threadIdx.x;
    const int warp_id    = tid / WAVEFRONT_SIZE;    // 0..3
    const int lane       = tid & 63;
    const int g          = lane >> 4;               // group 0..3
    const int l          = lane & 15;               // lane-in-group 0..15
    const int head_base  = warp_id * HEADS_PER_WARP;

    const float qs  = q_scale[0];
    const float kvs = kv_scale[0];
    const float scale_log2 = qs * kvs * SM_SCALE_INV_SQRT * LOG2E;

    // KV range for this split
    const int kv_start = kv_indptr[batch_item];
    const int kv_end   = kv_indptr[batch_item + 1];
    const int kv_len   = kv_end - kv_start;
    const int split_size  = (kv_len + num_kv_splits - 1) / num_kv_splits;
    const int local_start = min(split_idx * split_size, kv_len);
    const int local_end   = min(local_start + split_size, kv_len);
    const int count       = local_end - local_start;

    // ---- Pre-load Q MFMA operands for 4 heads ----
    // MFMA A[16×128]: lane (g, l) holds A[row=l][cols=g*32..g*32+31]
    // We fill rows 0-3 with Q data (4 heads), rows 4-15 with zeros.
    // 5 MFMA calls for 5 K-dim slices of 128 (total 640, padding 576 to 640).
    v8i q_a[5];
    {
        const uint8_t* q_ptr = nullptr;
        if (l < HEADS_PER_WARP) {
            q_ptr = q + ((long)(batch_item * NUM_HEADS + head_base + l)) * QK_HEAD_DIM;
        }
        #pragma unroll
        for (int d = 0; d < 5; d++) {
            int dim_off = d * 128 + g * 32;
            if (l < HEADS_PER_WARP && dim_off + 32 <= QK_HEAD_DIM) {
                // Full 32-byte load
                const int* src = (const int*)(q_ptr + dim_off);
                q_a[d][0]=src[0]; q_a[d][1]=src[1]; q_a[d][2]=src[2]; q_a[d][3]=src[3];
                q_a[d][4]=src[4]; q_a[d][5]=src[5]; q_a[d][6]=src[6]; q_a[d][7]=src[7];
            } else if (l < HEADS_PER_WARP && dim_off < QK_HEAD_DIM) {
                // Partial load (last slice, groups 0-1 have 32 valid bytes, groups 2-3 zero)
                int valid_ints = (QK_HEAD_DIM - dim_off) / 4;
                const int* src = (const int*)(q_ptr + dim_off);
                q_a[d][0]=0; q_a[d][1]=0; q_a[d][2]=0; q_a[d][3]=0;
                q_a[d][4]=0; q_a[d][5]=0; q_a[d][6]=0; q_a[d][7]=0;
                for (int i = 0; i < valid_ints && i < 8; i++) q_a[d][i] = src[i];
            } else {
                // Zero (padded heads l>=4 or out-of-range dims)
                q_a[d][0]=0; q_a[d][1]=0; q_a[d][2]=0; q_a[d][3]=0;
                q_a[d][4]=0; q_a[d][5]=0; q_a[d][6]=0; q_a[d][7]=0;
            }
        }
    }

    // V accumulators and online softmax state (4 heads × 8 V dims)
    float va0[8]={0}, va1[8]={0}, va2[8]={0}, va3[8]={0};
    float rm0=-FLT_MAX, rm1=-FLT_MAX, rm2=-FLT_MAX, rm3=-FLT_MAX;
    float rs0=0.f, rs1=0.f, rs2=0.f, rs3=0.f;

    if (count > 0) {
        const uint8_t* kv_global = kv + (long)(kv_start + local_start) * QK_HEAD_DIM;

        for (int tile_start = 0; tile_start < count; tile_start += TILE_TOKENS) {
            const int tile_count = min(TILE_TOKENS, count - tile_start);

            // ===== Cooperative LDS load: 256 threads load KV tile =====
            {
                const int total_int4 = tile_count * (QK_HEAD_DIM / 16); // 576/16=36
                const int4* src = (const int4*)(kv_global + (long)tile_start * QK_HEAD_DIM);
                int4* dst = (int4*)kv_lds;
                for (int i = tid; i < total_int4; i += BLOCK_SIZE) {
                    dst[i] = src[i];
                }
            }
            __syncthreads();

            // ===== MFMA score + inline tile max (avoid re-reading score_lds) =====
            float tm_lane = -FLT_MAX;
            for (int r = 0; r < 2; r++) {
                const int token_off = r * 16;
                const int r_tile_count = min(16, tile_count - token_off);
                if (r_tile_count <= 0) break;

                v4f c = {0.f, 0.f, 0.f, 0.f};
                #pragma unroll
                for (int d = 0; d < 5; d++) {
                    int dim_off = d * 128 + g * 32;
                    v8i b;
                    if (l < r_tile_count && dim_off + 32 <= QK_HEAD_DIM) {
                        const int* src = (const int*)(kv_lds + (token_off + l) * QK_HEAD_DIM + dim_off);
                        b[0]=src[0]; b[1]=src[1]; b[2]=src[2]; b[3]=src[3];
                        b[4]=src[4]; b[5]=src[5]; b[6]=src[6]; b[7]=src[7];
                    } else if (l < r_tile_count && dim_off < QK_HEAD_DIM) {
                        int vi = (QK_HEAD_DIM - dim_off) / 4;
                        const int* src = (const int*)(kv_lds + (token_off + l) * QK_HEAD_DIM + dim_off);
                        b[0]=0;b[1]=0;b[2]=0;b[3]=0;b[4]=0;b[5]=0;b[6]=0;b[7]=0;
                        for (int i = 0; i < vi && i < 8; i++) b[i] = src[i];
                    } else {
                        b[0]=0;b[1]=0;b[2]=0;b[3]=0;b[4]=0;b[5]=0;b[6]=0;b[7]=0;
                    }
                    c = __builtin_amdgcn_mfma_scale_f32_16x16x128_f8f6f4(
                        b, q_a[d], c, 0, 0, 0, 0, 0, 0);
                }

                float s0=c[0]*scale_log2, s1=c[1]*scale_log2, s2=c[2]*scale_log2, s3=c[3]*scale_log2;
                if (g*4+0 >= r_tile_count) s0=-1e30f;
                if (g*4+1 >= r_tile_count) s1=-1e30f;
                if (g*4+2 >= r_tile_count) s2=-1e30f;
                if (g*4+3 >= r_tile_count) s3=-1e30f;
                if (l >= HEADS_PER_WARP) { s0=-1e30f; s1=-1e30f; s2=-1e30f; s3=-1e30f; }

                if (l < HEADS_PER_WARP) {
                    int base = warp_id * (HEADS_PER_WARP * TILE_TOKENS) + l * TILE_TOKENS + token_off + g * 4;
                    score_lds[base+0]=s0; score_lds[base+1]=s1; score_lds[base+2]=s2; score_lds[base+3]=s3;
                }

                // Compute per-head max from MFMA output directly (avoid score_lds re-read)
                float lm = fmaxf(fmaxf(s0, s1), fmaxf(s2, s3));
                lm = fmaxf(lm, read_lane(lm, ((g^1)*16)+l));
                lm = fmaxf(lm, read_lane(lm, ((g^2)*16)+l));
                lm = fmaxf(lm, read_lane(lm, ((g^3)*16)+l));
                tm_lane = fmaxf(tm_lane, lm);
            }

            // ===== Online softmax + V accumulation (token-first) =====
            {
                const int sb0 = warp_id * (HEADS_PER_WARP * TILE_TOKENS) + 0 * TILE_TOKENS;
                const int sb1 = warp_id * (HEADS_PER_WARP * TILE_TOKENS) + 1 * TILE_TOKENS;
                const int sb2 = warp_id * (HEADS_PER_WARP * TILE_TOKENS) + 2 * TILE_TOKENS;
                const int sb3 = warp_id * (HEADS_PER_WARP * TILE_TOKENS) + 3 * TILE_TOKENS;

                // Broadcast tile max from lanes 0-3 (one per head)
                float tm0 = read_lane(tm_lane, 0);
                float tm1 = read_lane(tm_lane, 1);
                float tm2 = read_lane(tm_lane, 2);
                float tm3 = read_lane(tm_lane, 3);

                // Online softmax correction
                float nm0=fmaxf(rm0,tm0), nm1=fmaxf(rm1,tm1), nm2=fmaxf(rm2,tm2), nm3=fmaxf(rm3,tm3);
                float c0=__builtin_amdgcn_exp2f(rm0-nm0), c1=__builtin_amdgcn_exp2f(rm1-nm1);
                float c2=__builtin_amdgcn_exp2f(rm2-nm2), c3=__builtin_amdgcn_exp2f(rm3-nm3);
                for (int i=0;i<8;i++) { va0[i]*=c0; va1[i]*=c1; va2[i]*=c2; va3[i]*=c3; }
                rs0*=c0; rs1*=c1; rs2*=c2; rs3*=c3;
                rm0=nm0; rm1=nm1; rm2=nm2; rm3=nm3;

                // Token-first V accumulation: kvs folded out to output write
                for (int t = 0; t < tile_count; t++) {
                    int2 vv = *(const int2*)(kv_lds + t * QK_HEAD_DIM + lane * 8);
                    float v0,v1,v2,v3,v4,v5,v6,v7;
                    unpack_fp8x4(vv.x, v0,v1,v2,v3);
                    unpack_fp8x4(vv.y, v4,v5,v6,v7);

                    float ew;
                    ew=__builtin_amdgcn_exp2f(score_lds[sb0+t]-nm0); rs0+=ew;
                    va0[0]+=ew*v0; va0[1]+=ew*v1; va0[2]+=ew*v2; va0[3]+=ew*v3;
                    va0[4]+=ew*v4; va0[5]+=ew*v5; va0[6]+=ew*v6; va0[7]+=ew*v7;

                    ew=__builtin_amdgcn_exp2f(score_lds[sb1+t]-nm1); rs1+=ew;
                    va1[0]+=ew*v0; va1[1]+=ew*v1; va1[2]+=ew*v2; va1[3]+=ew*v3;
                    va1[4]+=ew*v4; va1[5]+=ew*v5; va1[6]+=ew*v6; va1[7]+=ew*v7;

                    ew=__builtin_amdgcn_exp2f(score_lds[sb2+t]-nm2); rs2+=ew;
                    va2[0]+=ew*v0; va2[1]+=ew*v1; va2[2]+=ew*v2; va2[3]+=ew*v3;
                    va2[4]+=ew*v4; va2[5]+=ew*v5; va2[6]+=ew*v6; va2[7]+=ew*v7;

                    ew=__builtin_amdgcn_exp2f(score_lds[sb3+t]-nm3); rs3+=ew;
                    va3[0]+=ew*v0; va3[1]+=ew*v1; va3[2]+=ew*v2; va3[3]+=ew*v3;
                    va3[4]+=ew*v4; va3[5]+=ew*v5; va3[6]+=ew*v6; va3[7]+=ew*v7;
                }
            }

            __syncthreads();  // Before next tile's cooperative LDS load
        }
    }

    // ===== Write partial results (apply kvs here instead of inner loop) =====
    #define WRITE_HEAD(hh, rmH, rsH, vaH) {                                        \
        int head = head_base + hh;                                                  \
        int out_idx = (batch_item * NUM_HEADS + head) * num_kv_splits + split_idx;  \
        if (lane == 0) {                                                            \
            partial_max[out_idx] = rmH;                                             \
            partial_sum[out_idx] = rsH;                                             \
        }                                                                           \
        float* o_ptr = partial_o + (long)out_idx * V_HEAD_DIM + lane * 8;           \
        float4* o4 = reinterpret_cast<float4*>(o_ptr);                              \
        o4[0] = make_float4(vaH[0]*kvs, vaH[1]*kvs, vaH[2]*kvs, vaH[3]*kvs);      \
        o4[1] = make_float4(vaH[4]*kvs, vaH[5]*kvs, vaH[6]*kvs, vaH[7]*kvs);      \
    }
    WRITE_HEAD(0, rm0, rs0, va0)
    WRITE_HEAD(1, rm1, rs1, va1)
    WRITE_HEAD(2, rm2, rs2, va2)
    WRITE_HEAD(3, rm3, rs3, va3)
    #undef WRITE_HEAD
}


// ============================================================================
// Phase 2: Reduce partial results → final BF16 output
// ============================================================================
__global__
__launch_bounds__(WAVEFRONT_SIZE)
void mla_reduce_kernel(
    const float* __restrict__ partial_o,
    const float* __restrict__ partial_max,
    const float* __restrict__ partial_sum,
    __hip_bfloat16* __restrict__ output,
    int num_kv_splits
) {
    const int batch_head = blockIdx.x;
    const int lane       = threadIdx.x;
    const int meta_base  = batch_head * num_kv_splits;

    float my_max = (lane < num_kv_splits) ? partial_max[meta_base + lane] : -FLT_MAX;
    float gmax = warp_reduce_max(my_max);

    float my_rescale = 0.0f, my_psum = 0.0f;
    if (lane < num_kv_splits) {
        my_rescale = __builtin_amdgcn_exp2f(my_max - gmax);
        my_psum    = partial_sum[meta_base + lane] * my_rescale;
    }
    float total_sum = warp_reduce_sum(my_psum);
    const float inv_sum = (total_sum > 0.0f) ? 1.0f / total_sum : 0.0f;

    const int v_base = lane * 8;
    float a0=0,a1=0,a2=0,a3=0,a4=0,a5=0,a6=0,a7=0;
    for (int s = 0; s < num_kv_splits; s++) {
        const float r = read_lane(my_rescale, s);
        const float4* src = reinterpret_cast<const float4*>(
            partial_o + ((long)meta_base + s) * V_HEAD_DIM + v_base);
        float4 v0 = src[0], v1 = src[1];
        a0+=v0.x*r; a1+=v0.y*r; a2+=v0.z*r; a3+=v0.w*r;
        a4+=v1.x*r; a5+=v1.y*r; a6+=v1.z*r; a7+=v1.w*r;
    }
    const int ob = batch_head * V_HEAD_DIM + v_base;
    __hip_bfloat162* dst = reinterpret_cast<__hip_bfloat162*>(output + ob);
    dst[0]=__halves2bfloat162(__float2bfloat16(a0*inv_sum),__float2bfloat16(a1*inv_sum));
    dst[1]=__halves2bfloat162(__float2bfloat16(a2*inv_sum),__float2bfloat16(a3*inv_sum));
    dst[2]=__halves2bfloat162(__float2bfloat16(a4*inv_sum),__float2bfloat16(a5*inv_sum));
    dst[3]=__halves2bfloat162(__float2bfloat16(a6*inv_sum),__float2bfloat16(a7*inv_sum));
}


torch::Tensor mla_decode_hip(
    torch::Tensor q_fp8,
    torch::Tensor kv_fp8,
    torch::Tensor q_scale,
    torch::Tensor kv_scale,
    torch::Tensor kv_indptr,
    int batch_size
) {
    const int total_q = batch_size;
    auto kv_flat = kv_fp8.reshape({-1, QK_HEAD_DIM});

    int num_kv_splits = std::max(4, std::min(64, 2048 / std::max(batch_size, 1)));

    auto f32_opts  = torch::TensorOptions().dtype(torch::kFloat32).device(q_fp8.device());
    auto bf16_opts = torch::TensorOptions().dtype(torch::kBFloat16).device(q_fp8.device());

    auto partial_o   = torch::empty({(long)total_q * NUM_HEADS * num_kv_splits * V_HEAD_DIM}, f32_opts);
    auto partial_max = torch::empty({total_q * NUM_HEADS * num_kv_splits}, f32_opts);
    auto partial_sum = torch::empty({total_q * NUM_HEADS * num_kv_splits}, f32_opts);

    dim3 grid1(batch_size, num_kv_splits);
    dim3 block1(BLOCK_SIZE);
    mla_mfma_partial_kernel<<<grid1, block1>>>(
        (const uint8_t*)q_fp8.data_ptr(),
        (const uint8_t*)kv_flat.data_ptr(),
        q_scale.data_ptr<float>(),
        kv_scale.data_ptr<float>(),
        kv_indptr.data_ptr<int>(),
        partial_o.data_ptr<float>(),
        partial_max.data_ptr<float>(),
        partial_sum.data_ptr<float>(),
        num_kv_splits
    );

    auto output = torch::empty({total_q, NUM_HEADS, V_HEAD_DIM}, bf16_opts);
    dim3 grid2(batch_size * NUM_HEADS);
    dim3 block2(WAVEFRONT_SIZE);
    mla_reduce_kernel<<<grid2, block2>>>(
        partial_o.data_ptr<float>(),
        partial_max.data_ptr<float>(),
        partial_sum.data_ptr<float>(),
        (__hip_bfloat16*)output.data_ptr(),
        num_kv_splits
    );
    return output;
}
'''

mla_cpp_source = '''
torch::Tensor mla_decode_hip(
    torch::Tensor q_fp8, torch::Tensor kv_fp8,
    torch::Tensor q_scale, torch::Tensor kv_scale,
    torch::Tensor kv_indptr, int batch_size
);
'''

mla_module = load_inline(
    name="mla_decode_v97_hybridq2",
    cpp_sources=mla_cpp_source,
    cuda_sources=mla_source,
    functions=["mla_decode_hip"],
    verbose=True,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--offload-arch=gfx950", "-mllvm", "-amdgpu-early-inline-all=true", "-mllvm", "-amdgpu-function-calls=false", "-mllvm", "-amdgpu-promote-alloca-to-vector-limit=128", "-ffast-math"],
)


from aiter.ops.quant import dynamic_per_tensor_quant
from aiter.mla import mla_decode_fwd
from aiter import get_mla_metadata_info_v1, get_mla_metadata_v1
from task import input_t, output_t

FP8_DT = torch.float8_e4m3fn
_FP8_MAX = torch.finfo(FP8_DT).max
_FP8_MIN = torch.finfo(FP8_DT).min
_qc = {}  # quant buffer cache

_mla = mla_module.mla_decode_hip  # avoid repeated attribute lookup
_meta_cache = {}
_SM_SCALE = 1.0 / (576 ** 0.5)


_aiter_cache = {}

def _aiter_decode(q, kv_fp8, q_scale, kv_scale, qo_indptr, kv_indptr, config):
    """Use aiter's optimized ASM kernel."""
    bs = config["batch_size"]
    total_kv = kv_fp8.shape[0]
    key = (bs, total_kv, q.dtype)
    c = _aiter_cache.get(key)
    if c is None:
        total_kv_len = int(kv_indptr[-1].item())
        kv_idx = torch.arange(total_kv_len, dtype=torch.int32, device="cuda")
        kv_lpl = (kv_indptr[1:] - kv_indptr[:-1]).to(torch.int32)
        o = torch.empty((bs, 16, 512), dtype=torch.bfloat16, device="cuda")
        NKS = 32
        info = get_mla_metadata_info_v1(
            bs, 1, 16, q.dtype, kv_fp8.dtype,
            is_sparse=False, fast_mode=False,
            num_kv_splits=NKS, intra_batch_mode=True)
        work = [torch.empty(s, dtype=t, device="cuda") for s, t in info]
        get_mla_metadata_v1(
            qo_indptr, kv_indptr, kv_lpl, 16, 1, True,
            work[0], work[2], work[1], work[3], work[4], work[5],
            page_size=1, kv_granularity=16, max_seqlen_qo=1, uni_seqlen_qo=1,
            fast_mode=False, max_split_per_batch=NKS,
            intra_batch_mode=True, dtype_q=q.dtype, dtype_kv=kv_fp8.dtype)
        c = (kv_idx, kv_lpl, o, work[0], work[1], work[2], work[3], work[4], work[5], NKS)
        _aiter_cache[key] = c
    kv_idx, kv_lpl, o, wm, wi, wis, ri, rfm, rpm, nks = c
    mla_decode_fwd(
        q, kv_fp8.view(total_kv, 1, 1, 576), o,
        qo_indptr, kv_indptr, kv_idx, kv_lpl,
        1, page_size=1, nhead_kv=1,
        sm_scale=_SM_SCALE, logit_cap=0.0, num_kv_splits=nks,
        q_scale=q_scale, kv_scale=kv_scale, intra_batch_mode=True,
        work_meta_data=wm, work_indptr=wi, work_info_set=wis,
        reduce_indptr=ri, reduce_final_map=rfm, reduce_partial_map=rpm)
    return o


def custom_kernel(data: input_t) -> output_t:
    q, kv_data, qo_indptr, kv_indptr, config = data
    bs = config["batch_size"]
    kv_fp8, kv_scale = kv_data["fp8"]
    total_kv = kv_fp8.shape[0]
    kv_per_batch = total_kv // max(bs, 1)
    use_fp8_q = (kv_per_batch >= 4096 or bs <= 4)
    if use_fp8_q:
        s = q.shape
        qb = _qc.get(s)
        if qb is None:
            qb = (torch.empty(s, dtype=FP8_DT, device=q.device),
                  torch.empty(1, dtype=torch.float32, device=q.device))
            _qc[s] = qb
        dynamic_per_tensor_quant(qb[0], q, qb[1])
        return _aiter_decode(qb[0], kv_fp8, qb[1], kv_scale, qo_indptr, kv_indptr, config)
    else:
        return _aiter_decode(q, kv_fp8, None, kv_scale, qo_indptr, kv_indptr, config)