submission 718375

# Auto-generated by submit-single-shape.py
# Target shape: s2 = {"batchsize": 4, "kvseqlen": 8192, "qseqlen": 1}

import importlib.util
import sys
import os
from pathlib import Path
import tempfile

_TARGET_SOURCE = '"""\ntest_s2_v58_parreduce_w1: parallel reduce num_warps 2→1\nBase: test_s2_v57_parreduce_w2.py\nDirection: CONTINUING from v57 (attempt 4) — further reduce warp count\nTarget: s2 (batch=4, kv_seq_len=8192)\nChange: num_warps 2→1 for the parallel reduce kernel. On s1, w2→w1 gave -2.7%\n        (iter 502 KEPT). s2\'s [16, 256] tensor is larger — may or may not work.\nRationale: s2 reduce ~5us est. (22% of 22.3us). w1 eliminates sync overhead.\nScale: INCREMENTAL\n"""\nimport torch\nimport triton\nimport triton.language as tl\nfrom task import input_t, output_t\n\nNUM_HEADS: tl.constexpr = 16\nKV_LORA_RANK = 512\nQK_ROPE_HEAD_DIM = 64\nQK_HEAD_DIM = KV_LORA_RANK + QK_ROPE_HEAD_DIM\nV_HEAD_DIM = KV_LORA_RANK\nSM_SCALE = 1.0 / (QK_HEAD_DIM ** 0.5)\n\nNUM_SPLITS = 16\nSPLIT_LEN = 512\nKV_SEQ_LEN = 8192\nV_BLOCK = 128\n\n_cache = {}\n\n\n@triton.jit\ndef _flash_decode_fp8_s2_exact_vtile(\n    Q_ptr,\n    KV_ptr,\n    Mid_O,\n    Mid_lse,\n    stride_kv: tl.int64,\n    kv_scale,\n    sm_scale: tl.constexpr,\n    QK_DIM: tl.constexpr,\n    V_DIM: tl.constexpr,\n    BLOCK_N: tl.constexpr,\n    BLOCK_K: tl.constexpr,\n    NH: tl.constexpr,\n    V_BLOCK_C: tl.constexpr,\n    NUM_SPLITS_C: tl.constexpr,\n    SPLIT_LEN_C: tl.constexpr,\n    KV_SEQ_LEN_C: tl.constexpr,\n):\n    batch_id = tl.program_id(0)\n    split_id = tl.program_id(1)\n    v_block_id = tl.program_id(2)\n\n    out_idx = batch_id * NUM_SPLITS_C + split_id\n    kv_start = batch_id * KV_SEQ_LEN_C + split_id * SPLIT_LEN_C\n    tl.multiple_of(kv_start, 256)\n    h_offs = tl.arange(0, NH)\n    v_start = v_block_id * V_BLOCK_C\n    v_offs = v_start + tl.arange(0, V_BLOCK_C)\n    v_mask = v_offs < V_DIM\n    q_base = batch_id * NH * QK_DIM\n    tl.multiple_of(q_base, 128)\n    tl.multiple_of(stride_kv, 32)\n    tl.assume(stride_kv > 0)\n\n    acc = tl.zeros([NH, V_BLOCK_C], dtype=tl.float32)\n    m_i = tl.full([NH], float("-inf"), dtype=tl.float32)\n    l_i = tl.zeros([NH], dtype=tl.float32)\n    score_scale = sm_scale * kv_scale\n\n    for kv_offset in range(0, SPLIT_LEN_C, BLOCK_N):\n        kv_pos = kv_start + kv_offset\n        tl.multiple_of(kv_pos, 128)\n        n_offs = tl.arange(0, BLOCK_N)\n\n        scores = tl.zeros([NH, BLOCK_N], dtype=tl.float32)\n        for k_start in range(0, QK_DIM, BLOCK_K):\n            k_offs = tl.arange(0, BLOCK_K)\n            k_mask = (k_start + k_offs) < QK_DIM\n            q_tile = tl.load(\n                Q_ptr + q_base + h_offs[:, None] * QK_DIM + k_start + k_offs[None, :],\n                mask=k_mask[None, :],\n                other=0.0,\n            )\n            k_tile = tl.load(\n                KV_ptr + (kv_pos + n_offs[:, None]) * stride_kv + k_start + k_offs[None, :],\n                mask=k_mask[None, :],\n                other=0.0,\n            )\n            scores += tl.dot(q_tile.to(k_tile.dtype), tl.trans(k_tile))\n\n        scores *= score_scale\n\n        row_max = tl.max(scores, axis=1)\n        m_new = tl.maximum(m_i, row_max)\n        alpha = tl.exp(m_i - m_new)\n        l_i = l_i * alpha\n        exp_scores = tl.exp(scores - m_new[:, None])\n        l_i += tl.sum(exp_scores, axis=1)\n        acc = acc * alpha[:, None]\n\n        v_tile = tl.load(\n            KV_ptr + (kv_pos + n_offs[:, None]) * stride_kv + v_offs[None, :],\n            mask=v_mask[None, :],\n            other=0.0,\n        )\n        acc += tl.dot(exp_scores.to(v_tile.dtype), v_tile)\n        m_i = m_new\n\n    acc = (acc * kv_scale) / l_i[:, None]\n    lse_vals = m_i + tl.log(l_i)\n    tl.store(\n        Mid_O + out_idx * NH * V_DIM + h_offs[:, None] * V_DIM + v_offs[None, :],\n        acc.to(tl.bfloat16),\n        mask=v_mask[None, :],\n    )\n    tl.store(Mid_lse + out_idx * NH + h_offs, lse_vals)\n\n\n@triton.jit\ndef _reduce_splitk_parallel(\n    Mid_O,\n    Mid_lse,\n    O_ptr,\n    NUM_SPLITS_C: tl.constexpr,\n    V_DIM: tl.constexpr,\n    NH: tl.constexpr,\n    BLOCK_V: tl.constexpr,\n):\n    """Parallel vectorized reduce: load all splits at once, reduce in one pass."""\n    batch_id = tl.program_id(0)\n    head_id = tl.program_id(1)\n    v_block = tl.program_id(2)\n\n    v_offs = v_block * BLOCK_V + tl.arange(0, BLOCK_V)\n    v_mask = v_offs < V_DIM\n    s_offs = tl.arange(0, NUM_SPLITS_C)  # [16]\n\n    # Load ALL 16 lse values at once: [16]\n    all_lse = tl.load(\n        Mid_lse + (batch_id * NUM_SPLITS_C + s_offs) * NH + head_id,\n    )\n\n    # Find global max across all splits\n    m_global = tl.max(all_lse, axis=0)\n\n    # Compute all alphas: [16]\n    alphas = tl.exp(all_lse - m_global)\n\n    # Sum of alphas for normalization\n    l_total = tl.sum(alphas, axis=0)\n\n    # Normalize alphas: [16]\n    weights = alphas / l_total\n\n    # Load ALL 16 split partials at once: [16, BLOCK_V]\n    all_partials = tl.load(\n        Mid_O + (batch_id * NUM_SPLITS_C + s_offs[:, None]) * NH * V_DIM\n        + head_id * V_DIM + v_offs[None, :],\n        mask=v_mask[None, :],\n        other=0.0,\n    ).to(tl.float32)\n\n    # Weighted sum: [BLOCK_V] = sum over splits of weights * partials\n    result = tl.sum(weights[:, None] * all_partials, axis=0)\n\n    # Store final output\n    out_base = batch_id * NH * V_DIM + head_id * V_DIM\n    tl.store(O_ptr + out_base + v_offs, result.to(tl.bfloat16), mask=v_mask)\n\n\ndef _ensure_cache(batch_size, total_q, kv_scale_tensor):\n    key = ("s2_v58_par_reduce_w1", batch_size, total_q)\n    if key in _cache:\n        return _cache[key]\n\n    _cache[key] = {\n        "mid_o": torch.empty(\n            (batch_size * NUM_SPLITS, NUM_HEADS, V_HEAD_DIM),\n            dtype=torch.bfloat16,\n            device="cuda",\n        ),\n        "mid_lse": torch.empty(\n            (batch_size * NUM_SPLITS, NUM_HEADS),\n            dtype=torch.float32,\n            device="cuda",\n        ),\n        "o": torch.empty(\n            (total_q, NUM_HEADS, V_HEAD_DIM),\n            dtype=torch.bfloat16,\n            device="cuda",\n        ),\n        "kv_scale_val": kv_scale_tensor.item(),\n    }\n    return _cache[key]\n\n\ndef custom_kernel(data: input_t) -> output_t:\n    q, kv_data, qo_indptr, kv_indptr, config = data\n    del qo_indptr, kv_indptr\n\n    batch_size = config["batch_size"]\n    total_q = q.shape[0]\n    total_kv = batch_size * KV_SEQ_LEN\n\n    kv_fp8, kv_scale = kv_data["fp8"]\n    kv_flat = kv_fp8.view(total_kv, QK_HEAD_DIM)\n    c = _ensure_cache(batch_size, total_q, kv_scale)\n\n    # Stage1: UNCHANGED from v42\n    grid1 = (batch_size, NUM_SPLITS, triton.cdiv(V_HEAD_DIM, V_BLOCK))\n    _flash_decode_fp8_s2_exact_vtile[grid1](\n        q,\n        kv_flat,\n        c["mid_o"],\n        c["mid_lse"],\n        QK_HEAD_DIM,\n        c["kv_scale_val"],\n        SM_SCALE,\n        QK_HEAD_DIM,\n        V_HEAD_DIM,\n        256,\n        256,\n        NUM_HEADS,\n        V_BLOCK,\n        NUM_SPLITS,\n        SPLIT_LEN,\n        KV_SEQ_LEN,\n        num_warps=4,\n        num_stages=3,\n    )\n\n    # Reduce: NEW parallel vectorized reduce\n    REDUCE_BLOCK_V = 256\n    n_v_blocks = triton.cdiv(V_HEAD_DIM, REDUCE_BLOCK_V)\n    _reduce_splitk_parallel[(batch_size, NUM_HEADS, n_v_blocks)](\n        c["mid_o"],\n        c["mid_lse"],\n        c["o"],\n        NUM_SPLITS,\n        V_HEAD_DIM,\n        NUM_HEADS,\n        REDUCE_BLOCK_V,\n        num_warps=1,  # was 2 — minimal warp for parallel reduce (KEPT on s1 w1)\n    )\n    return c["o"]\n'
_REF_SOURCE = '"""\ntest_v143_s6_splits4: s6 splits 8→4 (continue reduce optimization pattern)\nBase: test.py (v142)\nDirection: NEW — s6 splits tuning\nTarget: s6 (64,8192) — reduce overhead with kv=8192\nChange: s6 splits 8→4. batch=64 × splits=4 = 256 programs (100% CU fill).\n        Follows s5 splits reduction pattern (v142 +8.5%).\nRationale: v140 profile s8 reduce=3.3us. s6 with splits=8 has more reduce overhead.\nScale: INCREMENTAL\n"""\nimport torch\nimport aiter\nimport triton\nfrom task import input_t, output_t\n\nfrom aiter import dtypes as aiter_dtypes\nfrom aiter import get_mla_metadata_info_v1, get_mla_metadata_v1\nfrom aiter.mla import get_meta_param, _fwd_kernel_stage2_asm\n\nNUM_HEADS = 16\nNUM_KV_HEADS = 1\nKV_LORA_RANK = 512\nQK_ROPE_HEAD_DIM = 64\nQK_HEAD_DIM = KV_LORA_RANK + QK_ROPE_HEAD_DIM\nV_HEAD_DIM = KV_LORA_RANK\nSM_SCALE = 1.0 / (QK_HEAD_DIM ** 0.5)\nPAGE_SIZE = 1\nFP8_DTYPE = aiter_dtypes.fp8\n\n_cache = {}\n\n\ndef _ensure_cache_nonpers_bf16(batch_size, kv_seq_len, total_q):\n    key = ("npbf16", batch_size, kv_seq_len)\n    if key in _cache:\n        return _cache[key]\n\n    nq = NUM_HEADS\n    total_kv = batch_size * kv_seq_len\n\n    kv_last_page_len = torch.full((batch_size,), kv_seq_len, dtype=torch.int32, device="cuda")\n    kv_indices = torch.arange(total_kv, dtype=torch.int32, device="cuda")\n    num_kv_splits, num_kv_splits_indptr = get_meta_param(None, batch_size, total_kv, nq, 1, torch.bfloat16)\n    o = torch.empty((total_q, nq, V_HEAD_DIM), dtype=torch.bfloat16, device="cuda")\n    logits = torch.empty((total_q, num_kv_splits, nq, V_HEAD_DIM), dtype=torch.float32, device="cuda")\n    attn_lse = torch.empty((total_q, num_kv_splits, nq, 1), dtype=torch.float32, device="cuda")\n\n    _cache[key] = {\n        "kv_indices": kv_indices, "kv_last_page_len": kv_last_page_len,\n        "num_kv_splits": num_kv_splits, "num_kv_splits_indptr": num_kv_splits_indptr,\n        "logits": logits, "attn_lse": attn_lse, "o": o,\n    }\n    return _cache[key]\n\n\ndef _ensure_cache_persistent_fp8(batch_size, kv_seq_len, total_q, qo_indptr, kv_indptr, persistent_splits, fast_mode, kv_gran=16):\n    key = ("pfp8", batch_size, kv_seq_len, persistent_splits, fast_mode, kv_gran)\n    if key in _cache:\n        return _cache[key]\n\n    max_q_len = 1\n    nq, nkv = NUM_HEADS, NUM_KV_HEADS\n    total_kv = batch_size * kv_seq_len\n\n    kv_last_page_len = torch.full((batch_size,), kv_seq_len, dtype=torch.int32, device="cuda")\n    kv_indices = torch.arange(total_kv, dtype=torch.int32, device="cuda")\n\n    info = get_mla_metadata_info_v1(\n        batch_size, max_q_len, nq, FP8_DTYPE, FP8_DTYPE,\n        is_sparse=False, fast_mode=fast_mode,\n        num_kv_splits=persistent_splits, intra_batch_mode=True,\n    )\n    work = [torch.empty(s, dtype=t, device="cuda") for s, t in info]\n    (work_metadata, work_indptr, work_info_set,\n     reduce_indptr, reduce_final_map, reduce_partial_map) = work\n\n    get_mla_metadata_v1(\n        qo_indptr, kv_indptr, kv_last_page_len,\n        nq // nkv, nkv, True,\n        work_metadata, work_info_set, work_indptr,\n        reduce_indptr, reduce_final_map, reduce_partial_map,\n        page_size=PAGE_SIZE,\n        kv_granularity=max(PAGE_SIZE, kv_gran),\n        max_seqlen_qo=max_q_len,\n        uni_seqlen_qo=max_q_len,\n        fast_mode=fast_mode,\n        max_split_per_batch=persistent_splits,\n        intra_batch_mode=True,\n        dtype_q=FP8_DTYPE,\n        dtype_kv=FP8_DTYPE,\n    )\n\n    num_partials = reduce_partial_map.size(0)\n    logits = torch.empty((num_partials, 1, nq, V_HEAD_DIM), dtype=torch.float32, device="cuda")\n    attn_lse = torch.empty((num_partials, 1, nq, 1), dtype=torch.float32, device="cuda")\n    o = torch.empty((total_q, nq, V_HEAD_DIM), dtype=torch.bfloat16, device="cuda")\n    q_fp8 = torch.empty((total_q, nq * QK_HEAD_DIM), dtype=FP8_DTYPE, device="cuda")\n    q_scale = torch.ones(1, dtype=torch.float32, device="cuda")\n\n    _cache[key] = {\n        "kv_indices": kv_indices, "kv_last_page_len": kv_last_page_len,\n        "work_metadata": work_metadata, "work_indptr": work_indptr,\n        "work_info_set": work_info_set, "reduce_indptr": reduce_indptr,\n        "reduce_final_map": reduce_final_map, "reduce_partial_map": reduce_partial_map,\n        "logits": logits, "attn_lse": attn_lse, "o": o,\n        "q_fp8": q_fp8, "q_scale": q_scale,\n        "num_partials": num_partials,\n    }\n    return _cache[key]\n\n\ndef custom_kernel(data: input_t) -> output_t:\n    q, kv_data, qo_indptr, kv_indptr, config = data\n\n    batch_size = config["batch_size"]\n    kv_seq_len = config["kv_seq_len"]\n    total_q = q.shape[0]\n    total_kv = batch_size * kv_seq_len\n\n    # ---- Tier 1: batch<=4 -> bf16/bf16 non-persistent (s1, s2) ----\n    if batch_size <= 4:\n        kv_bf16 = kv_data["bf16"]\n        kv_4d = kv_bf16.view(total_kv, PAGE_SIZE, NUM_KV_HEADS, QK_HEAD_DIM)\n        c = _ensure_cache_nonpers_bf16(batch_size, kv_seq_len, total_q)\n\n        aiter.mla_decode_stage1_asm_fwd(\n            q.view(-1, NUM_HEADS, QK_HEAD_DIM), kv_4d,\n            qo_indptr, kv_indptr, c["kv_indices"], c["kv_last_page_len"],\n            c["num_kv_splits_indptr"],\n            None, None, None,\n            1, PAGE_SIZE, NUM_KV_HEADS, SM_SCALE,\n            c["logits"], c["attn_lse"], c["o"],\n            None, None,\n        )\n\n        Lv = V_HEAD_DIM\n        BLOCK_DV = triton.next_power_of_2(Lv)\n        _fwd_kernel_stage2_asm[(batch_size, NUM_HEADS)](\n            c["logits"], c["attn_lse"], c["o"],\n            qo_indptr, kv_indptr, c["num_kv_splits_indptr"],\n            c["attn_lse"].stride(0), c["attn_lse"].stride(2), c["attn_lse"].stride(1),\n            c["o"].stride(0), c["o"].stride(1),\n            MAYBE_FINAL_OUT=True,\n            BATCH_NUM=batch_size,\n            BLOCK_DV=BLOCK_DV,\n            Lv=Lv,\n            mgc=64,\n            num_warps=4,\n            num_stages=2,\n            waves_per_eu=4,\n        )\n        return c["o"]\n\n    # ---- Tier 2: ALL fp8 shapes -> persistent (s3-s8) ----\n    # Non-persistent fp8 was faster but fails leaderboard correctness (v77, v78).\n    # Persistent + mla_reduce_v1 is the only leaderboard-safe fp8 path.\n    else:\n        kv_buffer_fp8, kv_scale = kv_data["fp8"]\n        kv_buffer_4d = kv_buffer_fp8.view(total_kv, PAGE_SIZE, NUM_KV_HEADS, QK_HEAD_DIM)\n\n        # Per-shape split tuning\n        if total_kv >= 1000000:\n            # s8 (256, 8192) -- splits=4 (from v77)\n            splits, fast_mode = 4, False\n        elif total_kv >= 300000:\n            # s6 (64, 8192) -- splits=4 (from 8, 64*4=256 programs = 100% CU fill)\n            splits, fast_mode = 4, False\n        elif batch_size >= 256:\n            # s7 (256, 1024) -- splits=4 with kv_gran=64 (v131 LB-safe config)\n            # splits=1+kv_gran=64 FAILED LB in v136. splits=4 gives 1024 programs.\n            splits, fast_mode = 4, False\n        elif batch_size >= 64:\n            # s5 (64, 1024) -- splits=2 (from 4, reduce=7us → ~3.5us, 128 programs = 50% CU)\n            splits, fast_mode = 2, False\n        else:\n            # s3 (32, 1024) and s4 (32, 8192)\n            if kv_seq_len <= 1024:\n                splits, fast_mode = 4, True   # s3: reduced from 8 to 4\n            else:\n                splits, fast_mode = 32, True  # s4\n\n        # Use kv_granularity=64 for ALL persistent shapes (matches v131 LB-safe config)\n        # v131 (kv_gran=64 all) PASSED LB at 58.0us. v137/v138 (kv_gran=16 for s3/s5)\n        # FAILED LB on s3. kv_gran=64 is required for LB correctness.\n        kv_gran = 64\n        c = _ensure_cache_persistent_fp8(batch_size, kv_seq_len, total_q, qo_indptr, kv_indptr, splits, fast_mode, kv_gran)\n\n        # Fast FP8 quant: copy_ cast (scale=1.0) -- from v63\n        q_2d = q.view(total_q, NUM_HEADS * QK_HEAD_DIM)\n        c["q_fp8"].copy_(q_2d)\n\n        aiter.mla_decode_stage1_asm_fwd(\n            c["q_fp8"].view(-1, NUM_HEADS, QK_HEAD_DIM), kv_buffer_4d,\n            qo_indptr, kv_indptr, c["kv_indices"], c["kv_last_page_len"],\n            None, c["work_metadata"], c["work_indptr"], c["work_info_set"],\n            1, PAGE_SIZE, NUM_KV_HEADS, SM_SCALE,\n            c["logits"], c["attn_lse"], c["o"],\n            c["q_scale"], kv_scale,\n        )\n\n        aiter.mla_reduce_v1(\n            c["logits"], c["attn_lse"],\n            c["reduce_indptr"], c["reduce_final_map"], c["reduce_partial_map"],\n            1, c["o"], None,\n        )\n        return c["o"]\n'

def _load_module(name, source):
    tmpdir = tempfile.mkdtemp()
    path = os.path.join(tmpdir, name + '.py')
    open(path, 'w').write(source)
    spec = importlib.util.spec_from_file_location(name, path)
    mod = importlib.util.module_from_spec(spec)
    sys.modules[name] = mod
    spec.loader.exec_module(mod)
    return mod

# LAZY loading: modules are loaded on first use, not at import time.
# This prevents aiter (reference) from polluting Triton (target) state.
_target_mod = None
_ref_mod = None

from task import input_t, output_t

def custom_kernel(data: input_t) -> output_t:
    global _target_mod, _ref_mod
    _q, _cfg = data[0], data[4]
    if (_q.shape[0] == 4 and _cfg["kv_seq_len"] == 8192):
        if _target_mod is None:
            _target_mod = _load_module('_bss_target', _TARGET_SOURCE)
        return _target_mod.custom_kernel(data)
    else:
        if _ref_mod is None:
            _ref_mod = _load_module('_bss_ref', _REF_SOURCE)
        return _ref_mod.custom_kernel(data)