submission 672263

import torch
import numpy as np
from aiter.mla import mla_decode_fwd
from aiter import get_mla_metadata_v1, get_mla_metadata_info_v1

# 常量保持不变
PAGE_SIZE = 1
NUM_KV_SPLITS = 32
QK_HEAD_DIM = 576
SM_SCALE = 1.0 / (QK_HEAD_DIM ** 0.5)

def _make_mla_decode_metadata(batch_size, max_q_len, nhead, nhead_kv, q_dtype, kv_dtype, qo_indptr, kv_indptr, kv_last_page_len):
    info = get_mla_metadata_info_v1(
        batch_size, max_q_len, nhead, q_dtype, kv_dtype,
        is_sparse=False, fast_mode=False,
        num_kv_splits=NUM_KV_SPLITS, intra_batch_mode=True,
    )
    work = [torch.empty(s, dtype=t, device="cuda") for s, t in info]
    (work_metadata, work_indptr, work_info_set, reduce_indptr, reduce_final_map, reduce_partial_map) = work
    get_mla_metadata_v1(
        qo_indptr, kv_indptr, kv_last_page_len,
        nhead // nhead_kv, nhead_kv, True,
        work_metadata, work_info_set, work_indptr,
        reduce_indptr, reduce_final_map, reduce_partial_map,
        page_size=PAGE_SIZE, kv_granularity=max(PAGE_SIZE, 16),
        max_seqlen_qo=max_q_len, uni_seqlen_qo=max_q_len,
        fast_mode=False, max_split_per_batch=NUM_KV_SPLITS,
        intra_batch_mode=True, dtype_q=q_dtype, dtype_kv=kv_dtype,
    )
    return work_metadata, work_indptr, work_info_set, reduce_indptr, reduce_final_map, reduce_partial_map

def custom_kernel(data):
    q, kv_data, qo_indptr, kv_indptr, config = data
    
    # --- 核心改进：切换到 FP8 模式以换取极致速度 ---
    # 我们优先尝试使用 fp8 格式，它比 bf16 快得多
    if 'fp8' in kv_data:
        kv_buffer, kv_scale = kv_data['fp8']
        q_dtype = torch.float8_e4m3fn # 使用 FP8 精度
    else:
        # 兜底方案：如果没提供 fp8，还用你刚才成功的 bf16
        kv_buffer = kv_data['bf16']
        kv_scale = None
        q_dtype = q.dtype

    batch_size = config["batch_size"]
    nq = config["num_heads"]
    nkv = config["num_kv_heads"]
    dv = config["v_head_dim"]
    q_seq_len = config["q_seq_len"]

    kv_buffer_4d = kv_buffer.view(kv_buffer.shape[0], PAGE_SIZE, nkv, kv_buffer.shape[-1])
    kv_last_page_len = (kv_indptr[1:] - kv_indptr[:-1]).to(torch.int32)
    kv_indices = torch.arange(kv_buffer.shape[0], dtype=torch.int32, device="cuda")

    work_meta = _make_mla_decode_metadata(
        batch_size, q_seq_len, nq, nkv,
        q_dtype, kv_buffer.dtype,
        qo_indptr, kv_indptr, kv_last_page_len
    )

    output = torch.empty((q.shape[0], nq, dv), dtype=torch.bfloat16, device="cuda")

    # 执行内核
    mla_decode_fwd(
        q.view(-1, nq, QK_HEAD_DIM),
        kv_buffer_4d,
        output,
        qo_indptr,
        kv_indptr,
        kv_indices,
        kv_last_page_len,
        q_seq_len,
        page_size=PAGE_SIZE,
        nhead_kv=nkv,
        sm_scale=SM_SCALE,
        logit_cap=0.0,
        num_kv_splits=NUM_KV_SPLITS,
        intra_batch_mode=True,
        kv_scale=kv_scale, # 传入 FP8 缩放系数
        work_meta_data=work_meta[0],
        work_info_set=work_meta[2],
        work_indptr=work_meta[1],
        reduce_indptr=work_meta[3],
        reduce_final_map=work_meta[4],
        reduce_partial_map=work_meta[5]
    )
    return output